[inureyes]

Summary

The exaone-4.0 CLI decode loss after #2225 was a fixed per-call cost from stale state, not a slower engine step. ExaOne 4 cannot batch and keeps its KV caches inside the wrapper; its model-owned pool entry never calls make_caches(), so DirectEngine never reset them. mlxcel generate's timed run prefilled the prompt on top of the one-token warmup's KV: the output was wrong, and the stale-offset prefill loaded about 290 ms (under nsys) of cuBLAS kernels (gdb on cuLibraryLoadData) inside the timed window.

DirectEngine::open_sequence now calls make_caches() when supports_batching() is false, the rule BatchScheduler applies at prefill time and the retired loop applied through reset_with_model. GPT-OSS, Ministral 3, Step 3.5, Mellum and Laguna hold their caches the same way. Server path unchanged.

Behavior change: mlxcel generate (and loglikelihoods, MlxInferenceSession) on these families now decodes from fresh caches, so its text changes: exaone4 greedy text now matches the d379d2f CxxGenerator build; gpt-oss-20b no longer opens with a different reply; laguna-xs-2.1-nvfp4 greedy text with a warmup now equals its --profile text without one (checked in review). --prompt-lookup refuses models that cannot batch, so it is unaffected.

Measurement

GB10 sm_121, kernel 7.0.0-1019-nvidia, driver 580.178.04, MLX_CUDA_ARCHITECTURES=121, test-fast; ref d379d2f, main 8116e26, fix, null = fix binary again; rotated interleaving.

model	ref	main	fix	null
exaone4-1.2b, -n 384, n=7	202.73	179.28	202.04	201.94
llama-3.2-1b, -n 384, n=5	251.33	254.95	254.37	254.45
granite-4.0-h-tiny, -n 256, n=7	not in this block	99.07	97.38	99.37

exaone fix vs ref paired -0.72% (-1.14..+0.82), null floor +0.11% (-1.48..+1.46). Fit time = fixed + n * per_token: main 302 ms + 4.85 ms/tok vs ref 70 ms + 4.77 ms/tok; fix 57 ms + 4.97 ms/tok vs ref 70 ms + 4.89 ms/tok in the same block. granite main vs fix +2.14% sits inside its null spread (+1.55% median, same binary). Full record: docs/benchmark_results/exaone4-engine-decode-gb10-2026-10-22.md.

Test plan

• New engine::shared_cache_tests (2): fail with the reset disabled, pass with it
• -p mlxcel-core --lib engine:: (70), --lib exaone (3), granitemoehybrid (17), lfm2 (72), --test-threads=1, CUDA
• cargo clippy --profile test-fast --features cuda -p mlxcel-core --lib --tests -- -D warnings, cargo fmt --check
• Greedy --show-reasoning -n 48 text: exaone4 fix == ref; granite-4.0-h and ministral-3b fix == main

Unverified

• Step 3.5, Mellum and the Ministral3 text wrapper on a real checkpoint (none local; ministral-3b-4bit loads as a mistral3 VLM)
• Release-profile numbers; the issue's 198 vs 158 tok/s were release builds, these are test-fast
• No throughput measured for gpt-oss
• VLM --image prefill on a non-batching wrapper not exercised; it also goes through open_sequence, and the server applies the same reset before its embedding prefill

Closes #2494

[inureyes]

Implementation Review Summary

Intent

Reset shared model-owned caches when DirectEngine opens a sequence, so mlxcel generate's timed run stops prefilling on top of its warmup's KV (#2494).

Findings Addressed

No CRITICAL or HIGH findings, so nothing was changed.

Remaining Items

• The family list is incomplete (LOW). Mellum and Laguna also keep RefCell caches that only make_caches() resets, with no prepare_sequence_state. The doc comment, the test header, the PR body and the record name only four families. On laguna-xs-2.1-nvfp4 the warmup-then-timed text matches the --profile text with this PR and differs with the reset removed.
• The PR body says --prompt-lookup changes (LOW). prompt_lookup_unsupported_reason refuses every non-batching model, so it never runs on these families.
• Pre-existing, out of scope: longcat_flash_ngram and nemotron_nas have no reset side effect in make_caches.

Checks

• Non-batching make_caches() either resets a shared RefCell (families with no prepared state and no KV modes), replaces only the fallback slot (the prepare_sequence_state map and its modes stay intact; Engine always passes Some(seq_id)), or only allocates. VLM wrappers delegate to the text model and clear no image state.
• open_sequence callers (generate, loglikelihoods, the speculative entry points) each open one sequence on a capacity-1 engine. No drafter opens a second sequence on the target.
• Server path untouched. The scheduler already does open, then reset at prefill.
• The record's medians and paired deltas recompute from its per-run values. The post-fix fit still has a per-token slope 1.7% above ref; the interleaved -n 384 block (fix vs ref -0.72%, null range -1.48..+1.46%) cannot resolve it, as the record states.

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: engine:: 70, exaone 3, gpt_oss 11, ministral3 10 (CUDA). Without the reset both new tests fail, and exaone4 greedy text changes.

[inureyes]

Security and performance review

No CRITICAL or HIGH findings, so no fix commits were made.

State isolation. Every in-process path that drives a non-batching family now starts from fresh caches. generate (including embedding prefill), loglikelihoods, generate_with_drafter and warm_up_verify_widths all open through open_sequence. MlxInferenceSession builds a new DirectEngine per call, so it inherits the reset, and mlxcel generate, bench_decode, speculative_bench, engine_parity and the engine probe use the same client. The server already resets in execute_full_prefill and start_chunked_prefill. Diffusion, Florence-2 and Nemotron-Parse run on their own batch-1 loops. The Python package is an HTTP client with no in-process binding. Drafter bind runs before the reset but borrows only the embedding and LM head, not caches.

Error handling. make_caches returns Vec<KVCache>, not a Result. For these families the vector holds dummy caches and the reset is a side effect, so let _ = drops nothing that could carry an error. Cache construction is lazy (keys: None, including RotatingKVCache::new), so no allocation can fail at this point.

Performance. The reset runs once per open, never per step. Batching families skip it: their pool entry already called make_caches. Non-batching families get a model-owned pool entry that never calls it, so this is their only call. For ModelOwnedSequenceState families (Mamba, Jamba, NemotronH, LFM2, Kimi, RWKV7, ...) the call replaces only the fallback slot, which the engine path does not use. That costs one cheap object allocation per open, and the per-id state from prepare_sequence_state is left alone.

LOW, not fixed.

• shared_cache_tests.rs does not assert that a batching stub sees exactly one make_caches per open. The gate is a single if, so this is a regression guard only.
• Two DirectEngine clients over the same non-batching model would still share one RefCell cache set. This comes from the shared-cache design, and this PR does not cause it.
• Not in this PR: NemotronNAS reports supports_batching() == false but its forward indexes the caller's caches slice. The model-owned pool entry leaves that slice empty, so the engine path looks like it would panic on the first attention layer. This needs its own issue.

Tests. The mlxcel-core lib tests filtered to engine:: pass (70 of 70, including both new tests), and the session filter passes (13 of 13). Both ran with --profile test-fast --features cuda, --test-threads=1 and MLX_ENABLE_TF32=0 under the GPU lock.