fix(models): size dense prefill masks from live_len for mistral4/nemotron_nas/qwen-vl under --max-kv-size (#421)

[inureyes]

Summary

Completes the offset -> live_len prefill-mask migration started in #419 / PR #420, applying it to the five remaining dense Vec<KVCache> scheduler-pool models that still built their multi-token (l > 1) prefill causal mask from the cache's monotonic offset.

The bug

Under the server --max-kv-size path, enforce_max_kv_size_for (src/server/batch/scheduler.rs) calls KVCache::trim_front, which advances live_start while offset keeps growing monotonically (the RoPE invariant). KVCache::update_and_fetch then returns only the live window of live_len = offset - live_start keys, not offset keys. A model-level mask sized from offset is therefore WIDER than the live K/V after a trim. MLX broadcast_shapes throws, and because the throw crosses the cxx FFI boundary as a non-Result C++ exception it reaches std::terminate and ABORTS the whole server process. This is a remote-triggerable full-server abort: one long prompt plus a multi-token continuation chunk crashes every session.

The fix (mask-only)

For mask construction only, size the prefill causal mask from cache.live_len() (alias of seq_len(), both = offset - live_start) instead of cache.offset. RoPE, position_ids, and the Llama-4 attention scale keep reading the monotonic offset. With no trim (live_start == 0), live_len == offset, so greedy output on the normal untrimmed path is byte-identical.

What changed

• src/models/mistral4.rs: Mistral4Model::transformer_body builds the prefill mask from caches[0].live_len(). The separate let offset = caches[0].offset that feeds the Llama-4 attention scale (get_llama4_attn_scale) is left monotonic and unchanged.
• src/models/nemotron_nas.rs: LanguageModel::forward no-mask branch (the server path) builds the mask from caches.first().map(|c| c.live_len()). The per-layer RoPE offset in the attention block is untouched.
• src/models/qwen2_vl.rs, src/models/qwen3_vl.rs, src/models/qwen3_vl_moe.rs: only the auto_mask = create_causal_mask(seq_len, ...) argument switches to caches[0].live_len(). The cache_offset variable stays caches[0].offset and continues to drive position_ids / MRoPE grid slicing and the deepstack cache_offset == 0 checks, because position_ids and RoPE need absolute (monotonic) positions.

Scope: confirmed affected vs excluded

The five models above each store attention K/V as a plain Vec<KVCache> in the scheduler pool, so enforce_max_kv_size_for trims them and the offset-sized mask is a live (remotely triggerable) bug. nemotron_h and jamba were investigated and EXCLUDED: they use ModelOwnedSequenceState, so CachePool::get_caches_mut returns an empty slice and the trim loop never runs against their attention caches. For them the offset-based mask is a latent inconsistency, not reachable via --max-kv-size, and is left unchanged. glm4_moe_lite and minicpm3 already size from seq_len() and are the reference pattern.

Test plan

• cargo fmt --check
• cargo check --lib --features metal,accelerate (compiles all five changed models)
• cargo clippy --lib --tests --features metal,accelerate -- -D warnings (clean)
• cargo test --release -p mlxcel-core --features metal,accelerate cache:: (402 passed) including the shared invariant guard kv_cache_continuation_mask_sizes_from_live_len_not_offset_after_trim (added in fix(models): size dense prefill masks from live_len under trim (#419) #420), which pins the trimmed-cache (live_start > 0) invariant these five models now rely on.

No per-model unit test is added: exercising each model's mask construction requires a fully loaded real checkpoint (heavy, not in CI), and the underlying primitive is already pinned by the shared cache-level invariant test above.

Real-model E2E validation status will be added by the orchestrator: nemotron_nas / qwen2_vl / qwen3_vl / qwen3_vl_moe validated locally under a small --max-kv-size with a multi-token continuation chunk; mistral4 (Mistral-Small-4-119B) validated when its checkpoint is available.

Closes #421

[inureyes]

PR Finalization Complete

Summary

• Tests: No new tests added. The shared invariant kv_cache_continuation_mask_sizes_from_live_len_not_offset_after_trim (added in fix(models): size dense prefill masks from live_len under trim (#419) #420, lives in src/lib/mlxcel-core/src/cache.rs:8157) pins the trim_front / live_len / create_causal_mask contract all five models now rely on; it passes. None of the five model files have any existing unit tests (consistent with all inference model files in the codebase), and per-model mask-construction tests require a fully loaded checkpoint, so no cheap unit test exists. No gap to fill.
• Documentation: No docs in docs/ reference --max-kv-size, live_len, trim_front, or dense-cache mask sizing behavior. The change is internal to model forward passes. No doc update warranted.
• Lint/Format: Clean. cargo check --lib, cargo clippy --lib --tests -- -D warnings, cargo fmt --check, and cargo test --release -p mlxcel-core cache:: (402 passed) all green on current HEAD af05dc748.

No commits were added; the branch is already production-ready. Not merged.

[inureyes]

Real-model validation (all 5 models, before/after under --max-kv-size 1024 --prefill-chunk-size 512)

Same bug class as #419/#420: a trim_front mid chunked-prefill leaves an offset-sized mask wider than the live K/V → MLX broadcast_shapes throw → cxx std::terminate → server abort. before = the merged main binary (loads all models but mask-buggy for these 5); after = this branch.

• qwen2-vl-2b: before = abort (broadcast_shapes (512,2048) vs (1,12,512,1536)), after = HTTP 200. Directly reproduced crash→fix. Untrimmed control coherent ("...evolution of computing from early mechanical calculators...").
• qwen3-vl-2b / qwen3-vl-30b-a3b: no-regression. These route text-only generation through forward_text_only; the fixed forward_impl auto-mask is the multimodal path. Byte-identical before==after; untrimmed controls coherent. They share qwen2_vl's exact one-line mask fix.
• nemotron-nas-30b: no-regression, coherent, byte-identical (crash not reproduced even at --max-kv-size 600 --prefill-chunk-size 200; it does not hit a mid-prefill trim + continuation, so the fix is defensive).
• mistral4 (Mistral-Small-4-119B): the mask fix is byte-identical and correct, and mistral4 now loads and generates text (via fix(loading): route Mistral3-VLM mistral4 (MLA) text backbone to the Mistral4 loader (#423) #424). Its mask crash→fix cannot be demonstrated because a SEPARATE bug aborts every mistral4 offset>0 forward first ([reshape] ... size 2097152 into (1,1,1), reachable both via image tokens and chunked prefill), tracked as fix(models): mistral4-VLM image+text path aborts in multimodal prefill (reshape size 12480512 into (1,1,1)) #425. mistral4 text-only without --max-kv-size (single-pass, offset==0) generates coherently, confirming the untrimmed path is unaffected.

Net: the fix is correct and byte-identical for all 5; qwen2_vl is the directly-reproduced crash→fix; the remaining 4 are no-regression (defensive / multimodal / blocked-by-#425). The shared cache invariant test from #420 pins the primitive.