[inureyes]

Summary

On the engine path a non-batching model gets an empty cache slice from Engine::prefill, so it must own its caches and reset them in make_caches(). NemotronNAS and LongCat-Flash / LongCat-Flash-Ngram broke that contract.

• NemotronNAS: one KVCache per attention block now lives in a RefCell; forward ignores the caller's slice, make_caches and reset_runtime_state reset it. The mask is sized from the first full-attention slot, because a replace_with_linear block holds a slot it never writes (with caches.first() a linear-first checkpoint aborts on its second prefill chunk). Batching stays off; the cache-per-attention-block vs num_layers() layout is a follow-up.
• LongCat: the dual sub-layer caches and n-gram tail live in a RefCell<LongcatState>, the mask uses live_len(), and supports_padded_prefill() is false (feat(scheduler): sequence-aware cache trim hook so model-owned families can use padded prefill #1755).
• Two pre-existing LongCat forward bugs had to be fixed before any forward could run. The router sliced rank-3 scores with two indices, so MLX threw and the process aborted on every forward. The n-gram embedding also cut the joined context to n - 1 tokens before hashing, which dropped the oldest neighbour on every decode step.
• Stale comments fixed: NemotronNAS has no recurrent state. Both families added to the open_sequence and shared_cache_tests lists. The docs/CONTINUOUS_BATCHING.md exclusion note was updated.

Decode cost: LongCat decode now attends over all cached positions in both sub-layers, so tok/s falls as context grows. That is the correct cost; before this fix each step attended to one key.

Test evidence

New tests in src/models/nemotron_nas_engine_tests.rs and src/models/longcat_flash_ngram_engine_tests.rs use tiny f32 models and run through DirectEngine and through forward with an empty slice.

Test	main 8385f40	after
nemotron_nas_direct_engine_run_completes	panic: index out of bounds (nemotron_nas.rs:731)	ok
nemotron_nas_consecutive_direct_engine_runs_match_recompute	same panic	ok
nemotron_nas_incremental_logits_match_full_forward	same panic	ok
nemotron_nas_linear_first_chunked_prefill_matches_full_forward	panic: index out of bounds (nemotron_nas.rs:731)	ok
longcat_flash{,_ngram}_incremental_logits_match_full_forward	abort: [slice] Invalid number of indices	ok
longcat_flash{,_ngram}_consecutive_direct_engine_runs_match_recompute	same abort	ok
longcat_flash_ngram_chunked_engine_prefill_matches_single_piece	same abort	ok
longcat_flash_refuses_padded_prefill	assert fails	ok

With only the router fix applied, the cache bug shows up on its own: incremental vs full max abs diff is 5.65 (flash) and 6.76 (ngram), and engine decode drifts from a recompute-greedy reference. A mutant that reverts only the n-gram truncation fix and the linear-first mask fix fails longcat_flash_ngram_* (diff 5.14) and aborts nemotron_nas_linear_first_* (broadcast_shapes (3,3) vs (1,4,3,6)).

Local gate on GB10 (--profile test-fast --features cuda, test binaries under gpu-lock, MLX_ENABLE_TF32=0): nemotron_nas 6/6, longcat_flash 8/8 (rerun at the final head), hybrid_ssm_workers_never_use_paged_decode_backend, mlxcel-core engine:: 71/71, cargo fmt --all -- --check, and cargo clippy --lib --tests -- -D warnings on both mlxcel and mlxcel-core.

Real checkpoints

• LongCat-Flash: no checkpoint fits a 128 GB host (560B), so the tiny-config equivalence tests above are the gate for both LongCat variants.
• NemotronNAS: mlx-community/Llama-3_3-Nemotron-Super-49B-v1-mlx-4bit (model_type: nemotron-nas, 80 layers, 26.1 GB resident) on GB10, binaries built from this branch, run under gpu-lock. On main this family panics on the first attention layer through Engine::prefill, so neither path produced text before.
  • mlxcel generate -p "Explain in three sentences why the sky is blue." -n 96 --temp 0: coherent answer ("The sky appears blue because of a phenomenon called Rayleigh scattering, in which shorter (blue) wavelengths of sunlight are scattered more efficiently than longer (red) wavelengths ..."), 96 tokens at 7.06 tok/s.
  • mlxcel-server: two identical greedy /v1/chat/completions requests in a row ("Name the four largest planets in the solar system, largest first.", 64 tokens) returned byte-identical, correct text ("1. Jupiter ... 2. Saturn ..."), so the per-request cache reset works on the server path.

Review

pr-reviewer and the security/performance review found no CRITICAL, HIGH or MEDIUM issues. Low notes shared with every model-owned family (ExaOne 4, GPT-OSS, Mellum, ...) stay out of scope: --kv-cache-mode / --max-kv-size do not reach the internal caches, and the last request's KV stays resident until the next prefill resets it. The finalize pass only merged a duplicated impl NemotronNASModel block.

Closes #2544

[inureyes]

Implementation Review Summary

Intent

NemotronNAS and LongCat-Flash / LongCat-Flash-Ngram keep model-owned caches (plus LongCat's n-gram tail), reset in make_caches, so the engine path stops panicking and decode carries context (#2544).

Findings Addressed

None needed. No CRITICAL or HIGH findings.

Remaining Items

• NemotronNAS real-checkpoint smoke (Llama-3_3-Nemotron-Super-49B 4-bit, greedy mlxcel generate plus one mlxcel-server completion) is still pending, as the PR body says. (LOW, tracked outside this review)

Notes from the review

• Cache contract: both forwards ignore the caller slice, make_caches and reset_runtime_state reset the RefCell state, supports_batching and supports_padded_prefill are false for both families. The default sequence_state_layout is therefore ModelOwned, which the server prompt cache already rejects.
• NemotronNAS mask: first_attention_slot indexes slots among blocks with an attention subblock, so no_op blocks (no slot) and replace_with_linear blocks (slot never written) are both skipped correctly. With no full attention block the mask is unused.
• LongCat n-gram: the hash now runs over the saved tail (at most n - 1 tokens) joined with the new tokens. Every shift up to n - 1 reads the true previous token, or zero when the sequence is shorter, which is exactly what a full-sequence forward computes. The router slice_axis(-1, 0, top_k) matches upstream [..., :top_k].
• RefCell: each borrow_mut is scoped to one forward or reset, and nothing inside the forward re-enters model state.
• No LanguageModel signature change in generate.rs; the scheduler_tests.rs edit is a doc comment only.

Verification

• All stated requirements implemented
• No placeholder/mock code remaining
• Integrated into project code flow
• Project conventions followed
• Existing modules reused where applicable
• No unintended structural changes
• Tests pass (nemotron_nas 6/6 and longcat_flash 8/8 on GB10 under gpu-lock with MLX_ENABLE_TF32=0, cargo fmt --check clean; failing-first logs on main 8385f40 show the index-out-of-bounds panic and the MLX [slice] abort)

[inureyes]

Security and Performance Review

Reviewed at 9630310. No CRITICAL, HIGH or MEDIUM findings, so no fixes were pushed.

Checked

• Cross-request isolation: every path that prefills a non-batching model resets the model-owned state first. That covers execute_full_prefill, start_chunked_prefill (which also covers the batched-window fallbacks) and DirectEngine::open_sequence (CLI generate). The worker clamps max_batch_size to 1 for these families, so one sequence owns the state at a time. An aborted or failed request leaves partial KV that the next prefill discards before any forward reads it.
• Paths that would read or keep state without a reset are all declined for this layout: prompt cache adopt/donate and snapshot warm-up (ModelOwned, no supports_snapshot_reuse), disaggregated handoff (dense only), decode lookahead (no supports_decode_lookahead_rewind), prompt-lookup and drafter verify rounds (NotTrimmable, since the model-owned pool entry has no caches), padded prefill (now false for both), and DFlash (no target impl).
• N-gram tail: the saved tail is capped at n - 1 tokens per call, so it does not grow with the sequence. The joined context is part of every step's logits graph, so evaluation cuts the lazy-graph chain each step.
• Panics: the RefCell borrow is held only inside forward, and nothing in it re-enters make_caches or reset_runtime_state. NemotronNAS indexes a vector built from the same has_attention filter, and LongCat sizes its state from num_layers, the same count as layers. The router slice_axis(-1, 0, top_k) removes the rank-3 MLX abort.
• LongCat decode cost now grows with context because each step attends to the full KV. This is the intended behavior.

LOW, not fixed

• The last request's model-owned KV stays allocated until the next prefill resets it. ExaOne 4, GPT-OSS and the other RefCell families do the same. Resetting it on sequence release instead could clear an in-flight sequence when a queued request is cancelled.
• Internal caches do not see --max-kv-size or --kv-cache-mode. Growth is bounded by the request's context window. This is out of scope per the issue.

Tests

nemotron_nas (6) and longcat_flash (8) lib tests pass on CUDA (--profile test-fast --features cuda, --test-threads=1, MLX_ENABLE_TF32=0).