fix(models): complete offset->live_len prefill-mask migration for dense-KVCache models under --max-kv-size (mistral4/nemotron_nas/qwen*-vl)

[inureyes]

Problem / Background

Follow-up to #419 (PR #420), which fixed the dense-KVCache sliding-window trio (cohere2, gemma3n, olmo3). The same crash class exists in other models that build their multi-token (l > 1) prefill attention mask from the cache's monotonic offset.

Mechanism (unchanged from #419): under the server --max-kv-size path, enforce_max_kv_size_for (src/server/batch/scheduler.rs:2082) calls KVCache::trim_front (src/lib/mlxcel-core/src/cache.rs:2408), which advances live_start while offset keeps growing monotonically (RoPE invariant). KVCache::update_and_fetch (cache.rs:2515, live_len = buffer_idx() at :2535) then returns only the live window of live_len = offset - live_start keys, not offset keys. A model-level mask sized from the monotonic offset is therefore WIDER than the live K/V after a trim. MLX broadcast_shapes throws; because the throw crosses the cxx FFI boundary as a non-Result C++ exception it reaches std::terminate and ABORTS the whole server process. A remote client can crash all sessions with one long prompt plus a multi-token continuation chunk.

Scope is gated by cache ownership (important)

enforce_max_kv_size_for trims only the caches returned by CachePool::get_caches_mut, whose signature is Option<&mut [KVCache]> (cache.rs:5928). It iterates that slice and calls cache.trim_front(...). Two consequences:

• Models that store their attention K/V as a plain Vec<KVCache> in the scheduler cache pool (the standard dense path: make_caches() -> Vec<KVCache>, forward(caches: &mut [KVCache]) used directly) ARE trimmed. These are the affected models.
• Models that use ModelOwnedSequenceState<CustomCacheEnum> get SequenceCacheSet::model_owned, which sets caches: Vec::new() (cache.rs:5699-5707). get_caches_mut returns an empty slice, so the trim loop runs zero iterations and trim_front is never called on their attention caches. Likewise, models that keep their real caches in a custom enum or nested array (Cache, AnyKVCache, Qwen3NextCache, KimiLinearCache, [[KVCache; 2]], ...) and ignore the scheduler-provided &mut [KVCache] (the _caches parameter) are never trimmed through this path.

So an offset-based mask is only a live (remotely triggerable) bug for dense [KVCache] scheduler-pool models. For model-owned and enum/rotating models it is a latent inconsistency, not currently reachable via --max-kv-size.

The batched/serving path confirms the trigger: LanguageModel::forward_batched calls self.forward(input_ids, batch_caches[0], None) (generate.rs:591,600), so the model builds its own causal mask from cache.offset rather than receiving a scheduler-built one.

Confirmed affected (verified)

Each builds an l > 1 prefill create_causal_mask(seq_len, cache.offset) from the monotonic offset, stores its attention K/V as a plain Vec<KVCache> in the scheduler pool (so enforce_max_kv_size_for trims it), and passes the mask straight to attention_from_ptr against update_and_fetch keys with no rebuild against the live K/V:

• src/models/mistral4.rs:639-640 (Mistral4Model::transformer_body, offset = caches[0].offset). transformer_body always builds its own mask and ignores any caller-supplied mask, so the trait forward(.., mask) argument does not save it.
• src/models/nemotron_nas.rs:649-651 (LanguageModel::forward, offset = caches.first().map(|c| c.offset); built only when the caller passes mask = None, which is the server path).
• src/models/qwen2_vl.rs:637,693 (cache_offset = caches[0].offset; auto_mask = create_causal_mask(seq_len, cache_offset) when the caller passes no mask).
• src/models/qwen3_vl.rs:829,898 (same pattern).
• src/models/qwen3_vl_moe.rs:1022,1091 (same pattern).

The VLM three build the mask only on the no-mask branch; the text path (forward_batched -> forward(.., None)) takes that branch. End-to-end triggering on a VLM additionally requires the text decode to run under the trimmed scheduler pool with a multi-token continuation chunk, which should be validated per family.

Originally cited but NOT affected via --max-kv-size (corrected)

These build an offset-based prefill mask but use ModelOwnedSequenceState, so their attention caches are never trimmed by enforce_max_kv_size_for (latent inconsistency only, not remotely triggerable today). Dropping them from the confirmed list:

• src/models/nemotron_h.rs:1629 and :1983 (ModelOwnedSequenceState<NemotronLayerCache>).
• src/models/jamba.rs:855 and :988 (ModelOwnedSequenceState<JambaLayerCache>).

Already-correct reference (the migration pattern)

• src/models/glm4_moe_lite.rs:830-831 and src/models/minicpm3.rs:344-345 already size from caches.first().map(|c| c.seq_len()) (the live window). Both are dense [KVCache] scheduler-pool models, so they would be trimmed; using seq_len() is exactly what keeps them correct. Use them as the pattern.
• cohere2/gemma3n/olmo3 (fixed in fix(models): size dense prefill masks from live_len under trim (#419) #420) now size from *_live_len / live_len().

Audited candidates that are NOT affected (model-owned or internal enum/rotating caches; never trimmed by enforce_max_kv_size_for)

• Model-owned (ModelOwnedSequenceState): qwen3_5, falcon_h1, granitemoehybrid, gemma3, gemma4, llama4, lfm2, plamo2, recurrent_gemma.
• Custom enum / nested-array caches, scheduler &mut [KVCache] ignored: qwen3_next (Qwen3NextCache), kimi_linear (KimiLinearCache), longcat_flash_ngram ([[KVCache; 2]]), gpt_oss / exaone4 / ministral3 / mellum / step3p5 (Cache with Standard + Rotating), exaone_moe (AnyKVCache).

For all of these the offset-based mask is latent: correct today because the trim never runs against the cache the model actually reads. If any later gains a dense [KVCache] scheduler-pool trim path, it would inherit the same crash, so migrating them defensively is reasonable but not required to stop the current abort.

Proposed fix

Mirror #420: for MASK construction only, replace cache.offset with cache.live_len() (alias seq_len()); leave RoPE / position inputs reading the monotonic cache.offset. Byte-identical when untrimmed (live_start == 0, so live_len == offset). Apply to the five confirmed models. Optionally migrate the latent model-owned / enum sites in the same pass for consistency.

Acceptance criteria

• Each confirmed model (mistral4, nemotron_nas, qwen2_vl, qwen3_vl, qwen3_vl_moe), under --max-kv-size small enough to trigger a trim_front plus a multi-token continuation chunk, completes without the broadcast_shapes abort, with the mask key axis aligned to the live K/V.
• Untrimmed paths remain byte-identical (greedy output unchanged).
• Per-family real-checkpoint validation (CLAUDE.md requires it for inference changes); list the checkpoint used per family.
• A regression test pins the trimmed (live_start > 0) invariant, mirroring fix(models): size dense prefill masks from live_len under trim (#419) #420's kv_cache_continuation_mask_sizes_from_live_len_not_offset_after_trim.

References

• Follow-up to fix(server): dense sliding-window models build prefill masks from monotonic offset, not live_len, under --max-kv-size #419 / PR fix(models): size dense prefill masks from live_len under trim (#419) #420 (see the PR fix(models): size dense prefill masks from live_len under trim (#419) #420 security-review comment).