fix: size gemma3/gemma4/exaone_moe prefill mask from live_len under --max-kv-size

[inureyes]

Summary

gemma3, gemma4, and exaone_moe sized their multi-token (l > 1) prefill attention masks from the KV cache's monotonic offset instead of the live window, for both the global (full-attention) causal mask and the sliding-window prefill mask. Under the server --max-kv-size path, enforce_max_kv_size_for trims the global Standard cache via KVCache::trim_front, which advances live_start while offset keeps growing for RoPE, so update_and_fetch returns only live_len = offset - live_start keys. A mask sized from offset is then wider than the live K/V, MLX broadcast_shapes throws, and because the op is a non-Result cxx function the throw becomes std::terminate and aborts the whole server. This completes the offset -> live_len prefill-mask migration started in #419/#420 and #421/#422 for the three Gemma-family / ExaOne dense-cache models they missed.

The fix is mask-only and byte-identical when untrimmed (live_start == 0, so live_len == offset). RoPE, position_ids, and the batched-MTP coordinate math keep reading the monotonic offset.

What changed

• src/models/gemma3.rs: the seq_len > 1 prefill branch sizes the global mask from caches[global_idx].live_len() and the sliding mask from caches[0].live_len(). Added a live_len() accessor to the model's CacheInterface trait (KVCache -> live_len(), RotatingKVCache -> seq_len()).
• src/models/gemma4.rs: added a first_cache_live_len helper (free function plus method form) that mirrors first_cache_offset and returns the live window (live_len() for the full-attention Standard cache, seq_len() for the sliding Rotating cache). It feeds the two genuine prefill-mask sites: the non-ragged l > 1 branch in forward_with_speculative_sinks and the stage-forward execute_hidden path. The batched-MTP sites that use the offset as a per_row_valid_end column coordinate (divergent verify, has_padding, mask_stale_key_gap) intentionally keep the monotonic first_cache_offset; that regime enforces live_len == offset (slot_base == 0) anyway, and the value there is a coordinate, not a pure mask width. The now-unused first_cache_offset method form was removed (the free function remains in use).
• src/models/exaone_moe.rs: the global and sliding prefill masks size from a new AnyKVCache::live_len() (Standard -> live_len(), Rotating -> seq_len()) instead of offset().

Mask sites switched to live_len vs offset uses kept

Switched to live_len/seq_len (prefill mask width):

• gemma3 global mask (gemma3.rs, seq_len > 1 branch) and sliding mask.
• gemma4 non-ragged l > 1 prefill (create_causal_mask + create_sliding_window_prefill_mask) and the execute_hidden stage forward (seq_len > 1).
• exaone_moe global causal mask and sliding-window prefill mask.

Intentionally kept on the monotonic offset:

• All per-layer RoPE / fast_rope offsets (attention reads each layer's own cache.offset()).
• gemma4 batched-MTP coordinate sites: first_present_cache_offset divergence detection, build_divergent_verify_mask, the has_padding left-padding masks, and mask_stale_key_gap (these compare against per_row_valid_end, which lives in monotonic-offset space, and the MTP regime requires slot_base == 0).
• gemma4 speculative_cache_offset drafter RoPE rebind.

gemma4_mtp_target verdict

Verified safe, no change needed. gemma4_mtp_target does not implement LanguageModel and is not registered in model_metadata.rs; its adapters drive forward through Gemma4Wrapper::forward_with_speculative_sinks_explicit_cache on caller-owned Vec<Cache> that never enters the scheduler pool, so enforce_max_kv_size_for cannot reach it. Its only mask builder (create_causal_mask_with_left_padding) sizes from a fixed offset 0 prefill. The Gemma 4 prefill-mask sites it routes into are already covered by the gemma4.rs edits.

Test plan

• cargo fmt
• cargo check --lib --tests (clean, no warnings)
• cargo test --lib gemma3 (22 passed; new global_prefill_mask_sizes_from_live_len_not_offset_after_trim and sliding_cache_live_len_matches_returned_keys pass)
• cargo test --lib gemma4 (new first_cache_live_len_sizes_global_mask_from_live_window_after_trim and first_cache_live_len_sliding_matches_returned_keys pass; one pre-existing unrelated failure load_and_sanitize_weights_dequantizes_nvfp4_gemma4_checkpoint reproduces identically on the clean base commit and is NVFP4 weight-dequant, not mask logic)
• cargo test --lib exaone_moe (2 passed: new global_mask_sizes_from_live_len_not_offset_after_trim and sliding_cache_live_len_matches_returned_keys)
• cargo clippy --lib --tests -- -D warnings (clean)

The new per-model tests trim a Standard cache so live_start > 0, then assert the live_len-sized prefill mask key axis matches the continuation chunk's returned K/V (live_len + m) while the offset-sized mask is strictly wider, plus a Rotating-cache case where live_len equals the returned key axis. The shared cache invariant is already pinned by kv_cache_continuation_mask_sizes_from_live_len_not_offset_after_trim in cache.rs (#419).

Closes #430

[inureyes]

Implementation Review Summary

Intent

Size the multi-token (l > 1) prefill attention mask in gemma3/gemma4/exaone_moe from the cache's live window (live_len()/seq_len()) instead of the monotonic offset, for both the global and sliding masks, completing the #419/#420 + #421/#422 migration. Defensive/consistency hardening (not live-crashable today; the issue's remote-DoS framing was corrected), byte-identical on the untrimmed path.

Findings Addressed

• gemma3 Gemma3StageModel::execute_hidden (pipeline-stage path) still sized both the global and sliding prefill masks from offset() while the gemma4 stage analog and the gemma3 single-process path were switched (HIGH, incomplete migration). Fixed in this branch by switching both sites to as_interface().live_len(), mirroring Gemma3Model::forward_with_caches and Gemma4StageModel::execute_hidden. RoPE kept on the monotonic offset.

Remaining Items

• None. exaone_moe discards/recreates its internal AnyKVCache per forward (only offset round-trips to the external KVCache) is a pre-existing design unrelated to this PR and out of scope.

Regression verdict

• The untrimmed gemma3/gemma4 path is byte-identical. At the cache level KVCache::live_len() == seq_len() == offset - live_start, and untrimmed live_start == 0, so live_len == offset exactly. Confirmed at src/lib/mlxcel-core/src/cache.rs:720-731.
• Both the GLOBAL causal mask and the SLIDING-window prefill mask were switched in every file (gemma3 forward_with_caches 795/799 + stage 1074/1080 after fix; gemma4 non-ragged l > 1 2806/2807 + stage 3833/3834; exaone_moe 702/710).
• Cache accessor usage is correct: Standard -> live_len(), Rotating -> seq_len() (gemma3 CacheInterface, gemma4 first_cache_live_len, exaone_moe AnyKVCache::live_len).
• No RoPE/position/coordinate site was switched. Verified offset is kept at: gemma3 per-layer fast_rope/fast_rope_batched (141/223), gemma4 per-layer fast_rope (1399), exaone_moe per-layer fast_rope (506); gemma4 MTP coordinate sites (first_present_cache_offset 2697, build_divergent_verify_mask 2727/2732/2739, has_padding left-pad masks 2759/2760, mask_stale_key_gap upper bounds 2849/2869, speculative_cache_offset 4252). Each is a column coordinate compared against per_row_valid_end (monotonic-offset space) or a RoPE position, not a bare prefill-mask width, and the MTP regime enforces slot_base == 0 so live_len == offset there anyway.
• gemma4_mtp_target verdict confirmed safe: it does not impl LanguageModel (not registered for scheduling), and routes through Gemma4Wrapper::forward_with_speculative_sinks[_explicit_cache] on caller-owned caches that never enter the scheduler pool. The Gemma 4 prefill-mask sites it reaches are already covered.

Tests

• New per-model tests genuinely exercise the trimmed invariant (construct cache, trim_front so live_start > 0 and live_len < offset, then assert a live_len-sized mask key axis equals the continuation chunk's returned K/V live_len + m while an offset-sized mask is strictly wider), plus an untrimmed Rotating-cache equality case. Not trivially passing.

Verification

• All stated requirements implemented (offset -> live_len for both masks in all three models; one stage site was missed and is now fixed)
• No placeholder/mock code remaining
• Integrated into project code flow (model forward and pipeline-stage paths)
• Project conventions followed (reuses existing live_len()/seq_len()/as_interface() accessors; no new restructuring)
• Existing modules reused where applicable
• No unintended structural changes
• Tests pass

Narrow gate results (this dev host, MLX CPU build cached)

• cargo fmt --check: clean
• cargo check --lib: clean
• cargo test --lib gemma3: 22 passed (incl. 2 new)
• cargo test --lib gemma4: 148 passed (incl. 2 new); 1 pre-existing unrelated failure load_and_sanitize_weights_dequantizes_nvfp4_gemma4_checkpoint (NVFP4 dequant), confirmed failing identically on the clean base 22002d1
• cargo test --lib exaone_moe: 2 passed (2 new)
• cargo clippy --lib --tests -- -D warnings: clean

The fix commit (fix(models): size gemma3 stage prefill mask from live_len) is on the PR branch and not pushed.