fix(models): correct gemma sliding-window prefill beyond the window#405
Conversation
Implementation Review SummaryIntent
Verification (all focus areas confirmed)
Remaining Items (non-blocking, LOW / informational)
No CRITICAL or HIGH findings. The fix is correct, complete, well-tested, and properly integrated; no auto-fix required. |
Follow-up: HIGH finding found and fixed (Codex-assisted)The supplementary Codex review surfaced a real defect that my first pass missed, and on verification it is a reachable, PR-introduced crash. It is now fixed on this branch. Finding (HIGH): gemma3 over-window mask crashes on a non-fresh multi-token passMy earlier Focus-2 note claimed Confirmed mechanism:
Fix (commit 4cb2c18)
- if seq_len > window {
+ if seq_len > window && sliding_offset == 0 {For Re-verification (narrow selectors)
Updated verdictThe original issue #401 fix (over-window single-pass prefill, |
PR Finalization CompleteSummaryDocs ( Doc comments (two LOW items from review):
Tests: added Lint/format: No em dashes remain anywhere in the PR diff. Branch is pushed and clean. |
A single-pass prefill longer than the sliding window (1024 for gemma-4-12b-it-4bit) degenerated to <pad> for both gemma4_unified and gemma4_vl, and identically for gemma3. The rotating KV cache keeps every prefill key (it only trims to the window on the first decode step), but the sliding mask was built with create_causal_mask_with_window, which caps the key axis to the window when seq_len > window. The capped [seq_len, window] mask is then discarded by gemma4's trim_mask_to_keys (mask shorter than keys), and causal_attention slices K/V down to the trailing window keys, stranding the earliest query rows (logical position < seq_len - window) with no visible key. That all-masked row softmaxes to NaN and the model emits <pad> (or, for gemma3, a shape mismatch that degenerates to immediate EOS). The fix mirrors mlx-lm's RotatingKVCache: during prefill the windowed correctness comes from the MASK, not from physically dropping keys. A new shared helper create_causal_mask_with_window_full builds the uncapped [seq_len, seq_len + offset] windowed-causal mask so every query row attends to its own window over the full key set; trim_mask_to_keys still crops it when a later fetch caps. Both gemma4.rs mask sites (forward_with_speculative_sinks and the pipeline-stage execute_hidden) and both gemma3.rs sites route through a sliding_prefill_mask helper that uses the full mask only when seq_len > window and keeps the existing clamped path verbatim for seq_len <= window, so greedy output at or below the window is byte-identical (including batched MTP verify rounds). Validated on M1 Ultra with real checkpoints: gemma-4-12b-it-4bit >window text prefill (1511 tokens) and the default-fps #164 video command (20 frames, 1377 tokens) go from all-<pad> to coherent grounded output; gemma3-4b-4bit >window text goes from empty/degenerate to a full grounded answer; <=window short-text and single-image greedy output on gemma-4-12b-it-4bit and gemma3-4b-4bit are byte-identical before and after. gemma4 (146), gemma3 (18), diffusion_gemma (33), core cache (435) and utils (27) unit suites pass, plus new mask regression tests.
…ing cache The #401 fix gated the uncapped full mask purely on `seq_len > window`. On a non-fresh multi-token pass (`sliding_offset > 0`: server chunked prefill where the batch/ubatch/prefill-chunk size exceeds the sliding window, or multi-turn cache reuse feeding a >window chunk) the full builder returns a `[seq_len, seq_len + sliding_offset]` mask, but `RotatingKVCache::update_concat` trims the returned K/V to exactly `window` keys on a non-fresh append. Gemma 3's attention passes the mask straight to `attention_from_ptr` with no `trim_mask_to_keys` step (unlike Gemma 4), so the over-wide mask fails to broadcast against the `[B, H, seq_len, window]` scores and crashes. Gate the full mask on `seq_len > window && sliding_offset == 0` so only a fresh prefill (which keeps all `seq_len` keys, the #401 case) uses it. For `sliding_offset > 0` the else branch already collapses `effective_offset` to `0` (since `(window - seq_len).max(0) == 0` when `seq_len > window`) and yields a `[seq_len, window]` clamped mask that both broadcasts against the trimmed key axis and carries the correct windowed-causal values (the offset cancels in the relative query/key relation). Add gemma3 unit tests asserting the fresh path stays `[4, 4]` and the non-fresh over-window path is `[4, 2]` (not `[4, 7]`). Gemma 4 is unaffected: its `trim_mask_to_keys` re-crops the full mask to the window.
…ertion Tighten the trailing sentence of the gemma3 `sliding_prefill_mask` doc comment so it states the exact gate condition (`seq_len > window && sliding_offset == 0`) rather than only the `seq_len <= window` case, matching the else branch that also covers non-fresh over-window appends. Remove two em dashes from doc comments in utils.rs and gemma4.rs (project style forbids them). Add a `mask_at` helper and diagonal assertion to the gemma3 `sliding_prefill_mask_fresh_over_window_spans_full_key_axis` test for parity with gemma4's equivalent, proving no query row is fully masked.
df974b1 to
dad9c36
Compare
* fix: harden sliding-window over-window prefill across models A single-pass prefill longer than a model's sliding window degenerated to NaN or `<pad>` (or crashed on a mask/key shape mismatch) on every sliding-window model except gemma3/gemma4. When `seq_len > window` a fresh RotatingKVCache (and a dense KVCache) keep ALL prefill keys, but the shared `causal_attention` wrapper and the per-model prefill mask builders capped the windowed mask to `[seq_len, window]` and sliced K/V to the trailing `window` keys. That strands the earliest query rows (logical position `< seq_len - window`) with an all-`-inf` mask row that softmaxes to NaN; models that pass full K/V with a capped mask hit a hard shape mismatch instead. PR #405 fixed only gemma3/gemma4. General fix in `causal_attention` (src/lib/mlxcel-core/src/lib.rs): for a multi-token prefill that exceeds the window (`q_len > 1 && k_len > window`), build the full-width windowed-causal mask over ALL keys via `create_causal_mask_with_window_full` instead of slicing K/V. The single-query decode path (`q_len == 1`) keeps the trailing-window K/V slice byte-for-byte; non-windowed callers (`window_size == 0`) and within-window prefills are untouched. Shared helper `create_sliding_window_prefill_mask` (src/lib/mlxcel-core/src/utils.rs) selects the uncapped full mask for a fresh over-window prefill (`size > window && sliding_offset == 0`, the gemma3 prior-art gate) and the clamped mask otherwise, byte-identical to the per-model `sliding_offset.min((window - size).max(0))` construction it replaces. Wired into the prefill mask builders for gpt_oss, mellum, exaone4, exaone_moe, ministral3, step3p5 (RotatingKVCache, full K/V), and cohere2, gemma3n (dense KVCache with a model-level K/V slice, now sliced to the mask's key axis instead of blindly to `window` so a full mask keeps all keys). baichuan reaches the shared `causal_attention` prefill path directly, so the general fix covers it. recurrent_gemma (full causal mask, no window cap) and diffusion_gemma (already uses an uncapped dense windowed mask) are unaffected; gemma3/gemma4 keep their #405 helpers. Validated on M1 Ultra at --temp 0: gpt-oss-20b (window 128) with a ~260-token prompt now answers coherently where it previously broke; gemma-4-12b (window 1024) with a ~1600-token prompt stays coherent (#405 not regressed); short prompts and decode on gemma-4-12b, gpt-oss-20b, and non-windowed qwen3-0.6b are unchanged. Tests cover the `create_sliding_window_prefill_mask` selector and a `q_len > 1, k_len > window` `causal_attention` regression asserting a finite full-mask result. * fix(models): cover olmo3 in sliding-window >window prefill hardening olmo3 is a registered sliding-window model on a dense KVCache that was missed in the initial #408 audit. It built the clamped create_causal_mask_with_window and sliced K/V to the trailing window_size, so a fresh single-pass prefill longer than sliding_window (4096) stranded the earliest query rows with an all-(-inf) mask row -> NaN/<pad>. It is not covered by the general causal_attention fix because its windowed prefill goes through the explicit-mask branch. Mirror the cohere2/gemma3n dense-cache fix: build via create_sliding_window_prefill_mask and slice K/V to the mask key axis (mask_klen) instead of window_size, so a full mask keeps every key and a clamped mask drops the oldest. Correct the helper Used-by list (add Olmo3; remove Baichuan, which relies on the general causal_attention path). * fix(distributed): cover tensor-parallel gemma3 in over-window prefill hardening The tensor-parallel runtime re-implements gemma3 sliding-window mask construction (gemma3_masks) and still used the clamped create_causal_mask_with_window, so a fresh single-pass prefill longer than the sliding window crashed with a [broadcast_shapes] mismatch (the same DoS the single-process fix removed for gpt-oss): gemma3's RotatingKVCache returns all keys and the consumer passes the mask straight to attention_from_ptr with no trim_mask_to_keys fallback. Tensor-parallel gemma4 was already fixed incidentally, since its attend() trims and falls back to the patched causal_attention. Swap gemma3_masks onto the shared create_sliding_window_prefill_mask, mirroring the single-process gemma3 sliding_prefill_mask: full windowed mask for a fresh over-window prefill, clamped mask otherwise (byte-identical to the prior effective_offset construction for every working case). tensor_parallel_gemma3_matches_full_model_logits passes.
Summary
For Gemma 4 sliding-window models (
gemma4_unifiedandgemma4_vl), and identically for Gemma 3, a single-pass prefill whose length exceeds the sliding window (1024 forgemma-4-12b-it-4bit) degenerated to<pad>. This makes the default-fps video path from #164 (20 frames, ~1377 tokens > 1024) and any ~1500+ token text/image/video prompt produce garbage. The fix makes the windowed prefill correctness come from the attention mask rather than from physically dropping keys, mirroring mlx-lm'sRotatingKVCache.Root cause
The
RotatingKVCachekeeps every key on the first prefill pass (update_concatreturns all keys whenself.keys.is_none()) and only trims to the window on the subsequent decode step (update_in_place). But the sliding mask was built withcreate_causal_mask_with_window, which caps the key axis towindowwheneverseq_len + offset > window, on the assumption that the cache only ever returns the most recentwindowkeys. For a single-pass prefill longer than the window that assumption is false. The capped[seq_len, window]mask is then discarded by gemma4'strim_mask_to_keys(mask shorter than the full key axis), soattendfalls back tocausal_attention, which slices K/V down to the trailingwindowkeys (start = k_len - window_size). The earliest query rows (logical position< seq_len - window) then have no visible key at all: an all-masked attention row that softmaxes to NaN and emits<pad>. For Gemma 3 (whose attention passes the mask directly) the same capped mask mismatches the full key axis and degenerates to immediate EOS. PR #400 had already fixed the separateoverlay_block_bidirectionalreshape panic; this is the remaining numerical degeneration.The fix and why
New shared helper
create_causal_mask_with_window_full(size, offset, window)inmlxcel-core::utilsbuilds the uncapped[size, size + offset]windowed-causal mask (causal lower-triangular intersected with the sliding-window lower bound), so every query row attends to its own window over the full key set the rotating cache returns during prefill.trim_mask_to_keysstill crops it on any later capped fetch, so the change is robust for cache rollover too. Bothgemma4.rsmask sites (forward_with_speculative_sinksand the pipeline-stageexecute_hidden) and bothgemma3.rssites now route through a smallsliding_prefill_maskhelper that uses the full mask only whenseq_len > windowand keeps the existing clampedcreate_causal_mask_with_windowpath verbatim forseq_len <= window. That preserves byte-identical greedy output at or below the window (including the batched MTP verify rounds, which always havel <= window), and the full mask is mathematically the uncapped form of the same windowed-causal mask, so its lastwindowcolumns equal the old capped mask. This reuses the existing rotating-cache +trim_mask_to_keysmachinery rather than adding a parallel chunked-prefill path, and it preserves thegemma4_unifiedvision-span bidirectional overlay (the overlay now runs on the cleank == id_lenbranch).Validation matrix (M1 Ultra, real checkpoints,
--temp 0)BEFORE:
...Be specific and accurate.<pad><pad><pad>...(all<pad>).AFTER:
John Smeaton's Eddystone lighthouse was a significant engineering advancement because he modeled its shape on an oak tree trunk, which directed the force of waves downward into the rock rather than sideways against the masonry. He also pioneered the use of hydraulic lime, a mortar that sets underwater, and designed granite blocks that dovetailed into one another and the bedrock to lock the structure together as a single mass.generate -m models/gemma-4-12b-it-4bit --video <clip>.mp4 -p "Describe this video." -n 100(loader reports20 total frames,1377 total tokens> 1024, no--fpslowering):BEFORE:
Describe this video.<pad><pad><pad>...(all<pad>).AFTER:
The video shows a parking lot with several cars. A black car is parked in the foreground, and a silver car is parked next to it. A red car is parked in the background. The black car is moving forward and hits the silver car. The silver car is pushed back slightly. The black car then continues to move forward and comes to a stop...gemma-4-12b-it-4bit short text:
The sky appears blue because Earth's atmosphere scatters shorter wavelengths of sunlight, such as blue and violet, more effectively than longer wavelengths. This phenomenon, known as Rayleigh scattering, occurs as sunlight interacts with gas molecules and small particles in the air.(identical before/after).gemma-4-12b-it-4bit single image (cats.jpg, 288 tokens):
Two tabby cats are lying on a pink couch, with one cat on the left and another on the right.(identical before/after).BEFORE: empty / degenerate (no generated text after the prompt).
AFTER:
John Smeaton's Eddystone lighthouse was a significant advance because he fundamentally altered the structure's design to withstand the powerful forces of the sea. He modeled the tower's shape on a tree trunk, ensuring that wave energy was directed downwards into the bedrock rather than impacting the masonry directly. Furthermore, he pioneered the use of hydraulic lime and precisely cut granite blocks that dovetailed together... Augustin Fresnel's key optical innovation was the development of a series...gemma3-4b-4bit short text (<=window) is byte-identical before vs after.
Tests
mlxcel-coreutils:full_windowed_mask_over_window_has_no_all_masked_row(the [4,4] mask for window 2 has no fully-masked row, vs the capped [4,2] that strands rows 0 and 1) andfull_windowed_mask_within_window_matches_capped(cell-for-cell equality forsize + offset <= window).mlxcelgemma4:sliding_prefill_mask_over_window_has_no_fully_masked_rowandsliding_prefill_mask_within_window_matches_clamped.cargo fmt --checkclean;cargo clippy -p mlxcel-core -p mlxcel --lib --tests -- -D warningsclean.What changed
src/lib/mlxcel-core/src/utils.rs: addcreate_causal_mask_with_window_full+ two unit tests.src/models/gemma4.rs: addsliding_prefill_maskhelper, route both prefill mask sites through it.src/models/gemma3.rs: addsliding_prefill_maskhelper, route both prefill mask sites through it.docs/supported-models.md: replace thegemma4_unifiedknown-limitation note (long clips needing lower--fps) with the fixed behavior.Not regressed
diffusion_gemma(sharesoverlay_block_bidirectional) andgemma3are not regressed: the<= windowpath is unchanged (proven byte-identical by test and by real-model output), and the over-window path is the strict mlx-lm-equivalent fix. Decode (l == 1) is untouched.Closes #401