fix(models): NemotronNAS and LongCat-Flash-Ngram cache handling on the engine path #2544

Problem / Background

Found during the security review of PR #2542 (#2494, #2542 (comment)), not caused by it. On the engine path a model with supports_batching() == false gets a ModelOwned pool entry whose cache vector is empty (SequenceCacheSet::model_owned, src/lib/mlxcel-core/src/cache.rs:6530-6538, selected at cache.rs:6735), and Engine::prefill hands that empty slice to forward (src/lib/mlxcel-core/src/engine/prefill.rs:110-132). Such a model must keep its own caches and reset them in make_caches(), which the scheduler calls at prefill (src/server/batch/scheduler/prefill.rs:754-760, planned_prefill.rs:295-298) and which PR #2542 adds to DirectEngine::open_sequence (engine/direct.rs:322-332 on that branch). Two families break this contract. Implement both in one PR, run the local gate once at the end, and keep pushes to a minimum. PR #2542 is still open: branch from it, or from main after it merges, so the CLI reset is present.

1. NemotronNAS panics on the first attention layer

Current behavior. NemotronNASModel::forward (src/models/nemotron_nas.rs:696-747) indexes caches[cache_idx] at line 731. The mask code at 715 tolerates an empty slice (caches.first()), the layer loop does not, so every prefill through Engine::prefill (mlxcel generate since #2225, and mlxcel-server) panics with index out of bounds on the first layer that has attention. make_caches (749-757) only returns fresh vectors and resets no internal state. #1688 shows the family generating text before #2225, when the CLI loop built and passed its own caches.

Fix. Use the ExaOne4Wrapper pattern (src/models/exaone4.rs:688-736): add caches: RefCell<Vec<KVCache>> holding one KVCache per attention layer, have forward borrow it and ignore the caller's slice, size the causal mask from the internal first cache's live_len() (keeping the #421 behavior), and have make_caches replace the internal vector with fresh caches and return the same dummy vector it returns today. Keep supports_batching() == false and supports_padded_prefill() == false. Rejected: making the family batch-capable. make_caches yields one cache per attention layer (49 on the 49B checkpoint) while num_layers() is num_hidden_layers (80), and num_layers() drives the server's KV-mode table (scheduler/paged_layout.rs:91, checked against detached sets at scheduler/mod.rs:935) and the paged layout (paged_layout.rs:229); flipping it would also enable APC and paged decode for a family with no local checkpoint. That belongs in a follow-up.

Also correct two stale comments: nemotron_nas.rs:764 says padding corrupts "Mamba recurrent state" and src/server/batch/scheduler_tests.rs:1131 lists NemotronNAS among hybrid SSM families. It has no recurrent state, only a plain KVCache per attention layer; say it keeps them model-owned.

Acceptance.

• A tiny NemotronNASConfig (at least one attention block, one no_op attention block, one replace_with_linear block) built through from_weights with synthetic weights (pattern: tiny_config/tiny_weights in src/models/mellum_tests.rs:205,252) runs mlxcel_core::engine::DirectEngine::new(&model, 0).run(&prompt, 8, &greedy) to completion. This test panics on current code.
• Two consecutive run calls on one DirectEngine return identical tokens (the per-sequence reset works).
• Last-position logits of prefill-then-single-token-steps equal a full-sequence forward on fresh caches (f32 weights, max abs diff at most 1e-3).

2. LongCat-Flash keeps no KV between decode steps

Current behavior. LongcatFlashNgramModel::forward (src/models/longcat_flash_ngram.rs:1008-1021) builds fresh [KVCache; 2] per layer and ngram_context = None on every call. On a decode step cache.offset is 0, so RoPE runs at position 0 (longcat_flash_ngram.rs:238-254), update_and_fetch (261) returns only the current token's latent K/V, and the token attends to itself alone. The n-gram embedding degrades the same way: with an empty context, shift_right returns zeros when seq_len <= n (736-749), so the hashes see only the current token. Each step is not a recompute from scratch; it conditions on the last token only, so decode was never correct on any path. A single-piece prefill is correct, but prompts longer than the prefill chunk (DEFAULT_PREFILL_CHUNK = 2048, src/lib/mlxcel-core/src/prefill_plan.rs:61) lose every earlier piece. The struct serves both longcat_flash and longcat_flash_ngram (src/loaded_model.rs:243-244), so both are affected.

Fix. Same ExaOne4Wrapper pattern: hold RefCell<LongcatState> with caches: Vec<[KVCache; 2]> and ngram_context: Option<UniquePtr<MlxArray>>, have forward pass both to the existing internal_forward (970-1004), and have make_caches (1035-1038) reset both. Size the mask in internal_forward (988) from live_len() rather than offset, as NemotronNAS does. Override supports_padded_prefill() to return false: with persistent internal caches, pad positions would land in state no trim reaches (trait contract, src/lib/mlxcel-core/src/generate.rs:601-623, #1755).

Decode cost. Today each decode step attends over one key per sub-layer (constant and wrong). After the fix each step attends over all T cached positions in each of the two sub-layers per layer, and the KV grows by kv_lora_rank + qk_rope_head_dim values per token per sub-layer. Decode tok/s will fall as context grows; that is the expected, correct cost. Chunked prefill pieces now also attend to earlier pieces.

Acceptance.

• On tiny configs for both variants (is_ngram_model() true and false), last-position logits of prefill then single-token steps equal a full-sequence forward on fresh caches (f32 weights, max abs diff at most 1e-3). The helpers at longcat_flash_ngram.rs:1187 (mla_config) and 1225 (affine_kv_b_proj_weights) are a starting point. This test fails on current code.
• Two consecutive DirectEngine::run calls on one client return identical tokens.
• supports_padded_prefill() returns false.

Shared scope

• Add both families to the family lists in the open_sequence doc comment (engine/direct.rs:313-315) and the module doc of src/lib/mlxcel-core/src/engine/shared_cache_tests.rs:15-16.
• Out of scope: NASMLP::forward maps every activation other than silu/relu/gelu (including gelu_new, gelu_fast) to silu via _ => (nemotron_nas.rs:361-366); tracked with investigate: gelu_approx computes erf GELU where upstream models use the tanh approximation #2423's activation work. --kv-cache-mode does not reach these internal caches, the same as ExaOne4.
• Real checkpoints: none of either family exists on the dev hosts (models-101/mlx/nemotron-nas-30b-4bit is model_type: nemotron_h). If the implementer fetches mlx-community/Llama-3_3-Nemotron-Super-49B-v1-mlx-4bit (26 GB, fix(nemotron_nas): greedy output is still degraded after the RoPE fix in b429947c #1688), record a greedy mlxcel generate run and one mlxcel-server completion in the PR; otherwise say it was not run. The synthetic tests are the gate.

Verification

MLX_ENABLE_TF32=0 cargo test --profile test-fast --features cuda --lib nemotron_nas -- --test-threads=1
MLX_ENABLE_TF32=0 cargo test --profile test-fast --features cuda --lib longcat_flash -- --test-threads=1
cargo test -p mlxcel-core --profile test-fast --features cuda --lib engine:: -- --test-threads=1
cargo fmt --all -- --check
cargo clippy --workspace --all-targets --features cuda -- -D warnings
make verify-test-cuda

On macOS use --features metal,accelerate. On GB10 run test binaries under gpu-lock, never the build. A pass is all new tests green and no new failures in the full suite.