fix(gemma4): size chunked-prefill masks to returned keys per family

[inureyes]

Summary

Gemma 4 sliding-window models served with chunked prefill under concurrent load could collapse into a repeating cycle of reserved <unused...> tokens instead of a real answer. The root cause is a chunk-local prefill mask that is sized to the wrong key axis and then silently discarded, dropping causal / sliding-window alignment during a chunked-prefill continuation. This makes each family's prefill mask match the keys its cache actually returns, so the discard never happens and chunked prefill matches unchunked prefill.

Root cause (confirmed against code)

During a chunked-prefill CONTINUATION the scheduler passes the server's chunk-local create_padded_prefill_mask into the model forward. Gemma 4 is a model-owned-cache batching model (supports_batching() == true, make_caches() returns an empty Vec), so the cache pool hands the scheduler dummy placeholder caches and continue_chunked_prefill computes kv_offset = caches.first().map_or(0, |c| c.offset) == 0. The caller mask is therefore built with offset 0 (key axis 160 in the reporter's mask_len=160 log). Inside the model, forward_with_sequence_id resolves the REAL per-sequence caches (offset 2048), so the full-attention KVCache returns global_live_len + l keys (2208) and the sliding RotatingKVCache returns min(sliding_live_len, window - 1) + l keys (671, the reporter's key_len=671). The else if let Some(mask) = mask branch of forward_with_speculative_sinks (src/models/gemma4.rs:3058) reused that single short mask for BOTH families, so trim_mask_to_keys (src/models/gemma4.rs:2364) took its too-short branch and discarded the mask. Every layer then fell through to the maskless causal_attention path, which cannot reconstruct the correct causal + windowed alignment for queries at [offset, total) against a rotated key buffer, corrupting attention and driving the LM head into the reserved <unused...> region.

The full-attention (global) layers were affected by the same discard: their offset-0 caller mask (key axis 160) is also shorter than the full cache's returned 2208 keys, so their mask was discarded too and they relied on the maskless causal fallback. The fix now gives them an explicit mask sized to the live window as well.

What changed

• src/models/gemma4.rs (forward_with_speculative_sinks, else if let Some(mask) = mask branch): when l > 1 and the caller mask's key axis is shorter than the keys either family returns, rebuild BOTH masks from the caches' live windows exactly as the proven non-caller l > 1 branch does: create_causal_mask(l, global_live_len) for full attention and create_sliding_window_prefill_mask(l, sliding_live_len, window) for the sliding family. A correctly-sized caller mask (cold batched prefill, or a full prefill that adopted a prompt-cache prefix) is preserved unchanged and cropped per family by trim_mask_to_keys.
• src/models/gemma4.rs (trim_mask_to_keys / Attention::attend): the function now takes query_len. A too-short mask on a MULTI-token chunk (query_len > 1) fails loudly (panic referencing issue bug: Gemma 4 sliding-window server output collapses to reserved <unused> tokens under chunked prefill #885) instead of silently discarding, since silent discard is the corruption. Single-token decode (query_len == 1) keeps the safe discard-and-derive-own-mask behaviour. The mask-longer-than-keys crop path is unchanged.
• src/models/gemma4_tests.rs and the gemma4_unified_mask_tests module: deterministic regression coverage (below).

Test plan

• cargo test --lib gemma4_unified_mask_tests (9 passed) including new: chunked_continuation_sliding_mask_matches_returned_keys and chunked_continuation_global_mask_matches_returned_keys (rebuilt mask key axis equals the rotating/standard cache's returned keys, while the offset-0 caller mask is shorter), trim_mask_to_keys_crops_when_mask_longer_than_keys, trim_mask_to_keys_panics_on_too_short_multi_token (#[should_panic]), trim_mask_to_keys_discards_too_short_single_token_decode.
• cargo test --lib models::gemma4_tests::cache_isolation (4 passed, 3 ignored) including new PRIMARY parity test gemma4_chunked_prefill_with_server_caller_mask_matches_single_pass: on a tiny synthetic sliding-window Gemma 4 model, chunked prefill that feeds the server's offset-0 create_padded_prefill_mask on every continuation chunk produces logits within 1e-3 of single-pass prefill for sliding (window 4 and 8, chunk 8, over-window) and full attention. The rebuild is load-bearing: without it the sliding window-4 continuation chunk hits the new trim_mask_to_keys panic.
• cargo test --lib models::gemma4 (88 passed, 8 ignored) confirms the MTP / ragged / divergent-verify and existing chunked-parity paths still pass.
• cargo check --lib and cargo clippy --lib clean (only pre-existing mlxcel-core C++ build notes).

Full CUDA build plus server E2E on the GB10 host (gemma-4-e4b-it-qat-4bit, concurrent chunked prefill) is left to the orchestrator, since the MLX C++ link and a 26B/GPU run exceed the sandbox watchdog.

Closes #885

[inureyes]

Security & performance review: one CRITICAL finding (fixed on-branch)

CRITICAL: the trim_mask_to_keys panic is a reachable server-crash (DoS), via the MTP verify path

The prior reviewers concluded the new panic! is unreachable. That is correct only for the chunked-prefill path: the caller-mask-branch rebuild (forward_with_speculative_sinks, the else if let Some(mask) = mask arm) sizes both attention-family masks to the keys the caches return, so trim_mask_to_keys never sees a too-short multi-token mask there. Confirmed.

But trim_mask_to_keys is shared, and a different caller reaches it with a too-short l>1 mask: the Gemma 4 MTP (speculative) verify path. A B=1 (or equal-length batched) verify round calls forward_with_speculative_sinks with mask=None, left_padding=None, uniform per_row_valid_end, so it takes the plain else if l > 1 branch, which builds create_sliding_window_prefill_mask(l, sliding_live_len, window) (key axis capped at min(sliding_live_len, window-1) + l). In MTP the sliding cache is a buffered RotatingKVCache (enable_speculative_buffer), whose update_and_fetch returns sliding_live_len + l uncompacted keys. Once sliding_live_len >= window the mask is shorter than the returned keys, so the new panic fires and crashes the worker. Pre-#885 this same case silently discarded the mask and fell back to causal_attention(..., window_size), which is correct for the chronological buffered layout (no crash).

Empirically reproduced with a real Gemma4MtpTargetAdapter on the synthetic sliding_window=8 fixture: an 8-token prompt then one K=2 verify block panics with caller mask key axis (9) is shorter than the returned key length (10) ... (issue #885). Reverting only the panic makes the same call return normally. So the panic is a regression introduced by this PR, not a pre-existing crash. This is reachable by any normal MTP request that generates past the sliding window.

Fix applied on this branch (confined to src/models/gemma4.rs and src/models/gemma4_mtp_target_tests.rs): trim_mask_to_keys now discards a too-short mask with a warning for any query_len (the pre-#885 behavior), instead of panicking on query_len > 1. This removes the DoS while keeping the real #885 fix (the caller-mask-branch rebuild) fully intact: the rebuild already guarantees the chunked-prefill path never delivers a too-short multi-token mask to trim_mask_to_keys, so nothing on that path relies on the panic. The #[should_panic] unit test was converted to assert the discard, and a permanent MTP regression test (mtp_multi_token_verify_at_window_does_not_panic) was added.

Validation (narrow selectors, CPU sandbox): cargo check --lib, cargo clippy --lib -- -D warnings clean; models::gemma4 suite 90 passed / 0 failed / 8 ignored, including the chunked-prefill parity test (gemma4_chunked_prefill_with_server_caller_mask_matches_single_pass) and the new MTP regression test. Full CUDA build + real-model MTP greedy-parity on the GB10 host is left to the orchestrator.

Performance: no regression

The per-chunk mask rebuild is built once per forward (before the layer loop); the loop only selects a reference per layer, so it is not per-layer work. Both rebuilt masks use the same create_causal_mask / create_sliding_window_prefill_mask builders as the proven non-chunked l > 1 branch, so cost is O(l * keys), not quadratic-per-chunk. The common paths are untouched: single-token decode (l == 1) keeps the cheap (None, None) fallback, and the rebuild only fires on a genuine chunked continuation (mask.is_some() && l > 1 && caller mask shorter than the returned keys). A correctly-sized caller mask (cold batched prefill, prompt-cache-prefix adoption) still takes the cheap copy path.

Shared-branch other users: no further regressions

has_padding (ragged MTP) builds full-axis windowed masks (mask >= keys, crop-safe, never discarded); divergent_verify and per_row_valid_end require mask.is_none() or coexist only on prefill where per_row_valid_end is None, and use full-axis / gap-masked builders. The only correctness-as-safety regression in the shared branch was the MTP verify panic, now fixed.

Verdict on the panic-DoS reachability question

Reachable. CRITICAL. Not via chunked prefill (the reviewers' scope, where the rebuild makes it unreachable) but via the MTP speculative verify path, which shares trim_mask_to_keys. Fixed on-branch.

[inureyes]

PR Finalization Complete

Verified the branch at HEAD e230d8b (fix/issue-885-gemma4-chunked-prefill-sliding-mask) against the finalization checklist. No further changes were needed.

Formatting / Lint

• cargo fmt --all -- --check: clean across the whole workspace.
• cargo clippy --lib -- -D warnings: clean (only pre-existing mlxcel-core C++ build notes from the vendored MLX headers, unrelated to this change).
• cargo check --lib: clean.

Tests

• cargo test --lib models::gemma4: 90 passed, 8 ignored (serial-MLX-execution tests, expected), 0 failed. This includes the chunked-prefill parity test (gemma4_chunked_prefill_with_server_caller_mask_matches_single_pass), the mask-axis unit tests in gemma4_unified_mask_tests (including the two issue-bug: Gemma 4 sliding-window server output collapses to reserved <unused> tokens under chunked prefill #885 regression tests and the three trim_mask_to_keys crop/discard tests), and the MTP-verify regression test (mtp_multi_token_verify_at_window_does_not_panic) that covers the fix(gemma4): size chunked-prefill masks to returned keys per family #891 follow-on DoS fix.
• Reviewed the coverage for gaps: the one branch without a dedicated unit test is the "correctly-sized caller mask preserved unchanged" else-branch in forward_with_speculative_sinks, which is byte-for-byte the pre-existing (pre-bug: Gemma 4 sliding-window server output collapses to reserved <unused> tokens under chunked prefill #885) behavior and isn't touched by this fix's logic, so no new test was added for it to keep the diff scoped to the actual fix.

Documentation

• Checked docs/environment-variables.md (MLXCEL_PREFILL_CHUNK) and docs/CONTINUOUS_BATCHING.md (--prefill-chunk-size, batched-prefill mask sizing) for accuracy against this fix. Both describe code paths this fix doesn't change (the CLI single-sequence chunked-prefill path always passes None for the mask, and the continuous-batching doc's mask-memory discussion is about transient sizing, not per-family causal/sliding-window correctness), so no doc update was needed.

No commits were added; the branch was already clean and ready. All checks passing.