Skip to content

fix(models): correct gemma sliding-window prefill beyond the window#405

Merged
inureyes merged 3 commits into
mainfrom
fix/issue-401-gemma4-sliding-window-prefill
Jun 23, 2026
Merged

fix(models): correct gemma sliding-window prefill beyond the window#405
inureyes merged 3 commits into
mainfrom
fix/issue-401-gemma4-sliding-window-prefill

Conversation

@inureyes

Copy link
Copy Markdown
Member

Summary

For Gemma 4 sliding-window models (gemma4_unified and gemma4_vl), and identically for Gemma 3, a single-pass prefill whose length exceeds the sliding window (1024 for gemma-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's RotatingKVCache.

Root cause

The RotatingKVCache keeps every key on the first prefill pass (update_concat returns all keys when self.keys.is_none()) and only trims to the window on the subsequent decode step (update_in_place). But the sliding mask was built with create_causal_mask_with_window, which caps the key axis to window whenever seq_len + offset > window, on the assumption that the cache only ever returns the most recent window keys. For a single-pass prefill longer than the window that assumption is false. The capped [seq_len, window] mask is then discarded by gemma4's trim_mask_to_keys (mask shorter than the full key axis), so attend falls back to causal_attention, which slices K/V down to the trailing window keys (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 separate overlay_block_bidirectional reshape panic; this is the remaining numerical degeneration.

The fix and why

New shared helper create_causal_mask_with_window_full(size, offset, window) in mlxcel-core::utils builds 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_keys still crops it on any later capped fetch, so the change is robust for cache rollover too. Both gemma4.rs mask sites (forward_with_speculative_sinks and the pipeline-stage execute_hidden) and both gemma3.rs sites now route through a small sliding_prefill_mask helper that uses the full mask only when seq_len > window and keeps the existing clamped create_causal_mask_with_window path verbatim for seq_len <= window. That preserves byte-identical greedy output at or below the window (including the batched MTP verify rounds, which always have l <= window), and the full mask is mathematically the uncapped form of the same windowed-causal mask, so its last window columns equal the old capped mask. This reuses the existing rotating-cache + trim_mask_to_keys machinery rather than adding a parallel chunked-prefill path, and it preserves the gemma4_unified vision-span bidirectional overlay (the overlay now runs on the clean k == id_len branch).

Validation matrix (M1 Ultra, real checkpoints, --temp 0)

  1. gemma-4-12b-it-4bit, text-only >window prefill (1511 tokens):

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.

  1. gemma-4-12b-it-4bit, the default-fps feat: video input support for Gemma 4 Unified (gemma4_unified) #164 video command generate -m models/gemma-4-12b-it-4bit --video <clip>.mp4 -p "Describe this video." -n 100 (loader reports 20 total frames, 1377 total tokens > 1024, no --fps lowering):

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...

  1. Regression at/below the window (must be unchanged): generated text is byte-identical before vs after (only the non-deterministic model-load and tok/s timing lines differ in the captured stdout).

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).

  1. gemma3-4b-4bit, >window text prefill (1511 tokens):

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-core utils: 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) and full_windowed_mask_within_window_matches_capped (cell-for-cell equality for size + offset <= window).
  • mlxcel gemma4: sliding_prefill_mask_over_window_has_no_fully_masked_row and sliding_prefill_mask_within_window_matches_clamped.
  • Suites green: gemma4 (146), gemma3 (18), diffusion_gemma (33), core cache (435), core utils (27). cargo fmt --check clean; cargo clippy -p mlxcel-core -p mlxcel --lib --tests -- -D warnings clean.

What changed

  • src/lib/mlxcel-core/src/utils.rs: add create_causal_mask_with_window_full + two unit tests.
  • src/models/gemma4.rs: add sliding_prefill_mask helper, route both prefill mask sites through it.
  • src/models/gemma3.rs: add sliding_prefill_mask helper, route both prefill mask sites through it.
  • docs/supported-models.md: replace the gemma4_unified known-limitation note (long clips needing lower --fps) with the fixed behavior.

Not regressed

diffusion_gemma (shares overlay_block_bidirectional) and gemma3 are not regressed: the <= window path 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

@inureyes inureyes added status:review Under review type:bug Bug fixes, error corrections, or issue resolutions priority:medium Medium priority area:models Model architectures, weights, loading, metadata area:inference Generation, sampling, decoding (incl. speculative, DRY) labels Jun 23, 2026
@inureyes

Copy link
Copy Markdown
Member Author

Implementation Review Summary

Intent

Fix Gemma 4 (gemma4_unified / gemma4_vl) and Gemma 3 single-pass prefill degenerating to <pad> when the prompt exceeds the sliding window, by making windowed-prefill correctness come from the attention mask instead of dropping keys.

Verification (all focus areas confirmed)

  • All stated requirements implemented. New create_causal_mask_with_window_full is the uncapped form of the existing windowed-causal mask (tril(offset)triu(offset-w+1), where_cond to 0/-inf). Both Gemma 4 prefill sites (forward_with_speculative_sinks non-ragged l > 1 branch, Gemma4StageModel::execute_hidden) and both Gemma 3 sites (Gemma3Model::forward, Gemma3StageModel) route through sliding_prefill_mask.
  • seq_len > window gating is correct and the <= window path is byte-identical. The gate exactly matches the cap condition in create_causal_mask_with_window (for prefill offset == 0, the old code capped iff l > w). For l <= window the helper reproduces the prior code verbatim (sliding_offset.min((window - l).max(0)) + same create_causal_mask_with_window call), so output is unchanged. The two new unit tests per layer (*_over_window_has_no_fully_masked_row, *_within_window_matches_clamped) lock both halves.
  • Prefill -> first-decode boundary aligns. The mask change does not alter cache behavior (the rotating cache always kept all prefill keys and trims only on the first decode step). Fresh prefill: full [l, l] mask matches the full key axis, so trim_mask_to_keys is a no-op and attend uses the explicit mask instead of the causal_attention K/V-slicing fallback. On any later capped fetch the trailing-window slice of the full mask equals the old capped mask (verified algebraically: in the trailing columns the window lower bound triu is trivially satisfied, leaving only tril), so trim_mask_to_keys keeps it aligned. Decode (l == 1) passes None and is untouched.
  • Vision-span bidirectional overlay runs on the clean k == id_len branch. For unified video/image prefill the VLM passes mask = None + bidirectional_block_ids = Some; the full [l, l] mask (offset 0) has k == id_len == l, so overlay_block_bidirectional takes the clean branch. The PR feat: add video input support for Gemma 4 Unified (gemma4_unified) #400 windowed-cap (k < id_len) branch is retained for cache-rollover safety but is no longer exercised for over-window prefill (no contradiction).
  • diffusion_gemma path unchanged. It builds its own dense_windowed_causal_mask and calls the (unmodified) overlay_block_bidirectional; it does not use create_causal_mask_with_window or sliding_prefill_mask. The PR only adds a new utils function plus two private per-model helpers, so the diffusion path is byte-identical.
  • MTP / speculative routes consistently. Single-sequence MTP prefill flows through the same l > 1 branch (sliding_prefill_mask). Batched ragged MTP is admitted only when max_prompt_len <= sliding_window (gated via sliding_window_value), and verify rounds always have l <= window, so neither hits the over-window cap. This matches the byte-identical <= window claim.
  • Conventions followed. // Used by: present on the new shared helper; no AI attribution; docs/supported-models.md note accurately describes the fixed behavior.
  • Tests pass (per orchestrator: core utils 27, core cache 435, gemma4 146, gemma3 18, diffusion_gemma 33; build + clippy clean) and real-checkpoint validation on M1 Ultra (text/video over-window coherent; <=window byte-identical).

Remaining Items (non-blocking, LOW / informational)

  • (LOW, reuse) sliding_prefill_mask is duplicated near-identically in gemma4.rs and gemma3.rs; it could live once in mlxcel-core::utils beside create_causal_mask_with_window_full. Tiny (5 lines), optional.
  • (INFO, out of scope) Other sliding-window models that use the capped create_causal_mask_with_window (cohere2, gpt_oss, ministral3, exaone4, olmo3, gemma3n, mellum, step3p5) may share the same latent over-window single-pass prefill weakness depending on their attention path. Out of fix(models): gemma4 sliding-window single-pass prefill degenerates to <pad> beyond the window #401 scope; worth a follow-up audit.

No CRITICAL or HIGH findings. The fix is correct, complete, well-tested, and properly integrated; no auto-fix required.

@inureyes

Copy link
Copy Markdown
Member Author

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 pass

My earlier Focus-2 note claimed trim_mask_to_keys keeps the mask aligned on cache rollover. That is true for Gemma 4 but NOT for Gemma 3: Gemma3 Attention::forward (src/models/gemma3.rs:205-216) passes the mask straight to attention_from_ptr with no trim step.

Confirmed mechanism:

  • RotatingKVCache::update_concat (src/lib/mlxcel-core/src/cache.rs:4203-4218) trims the returned K/V to max_size (= window) on a NON-fresh multi-token append (keys.is_some() and total_len > max_size). On a fresh prefill (keys.is_none(), cache.rs:4187-4192) it keeps all seq_len keys.
  • When seq_len > window AND sliding_offset > 0, the old sliding_prefill_mask built a [seq_len, seq_len + sliding_offset] mask while the cache returned only window keys, so it could not broadcast against the [B, H, seq_len, window] scores: shape error / crash.
  • Reachable via server chunked prefill when --batch-size / --ubatch-size / prefill_chunk_size exceeds the sliding window (not clamped to it, see resolve_prefill_chunk_size), or multi-turn cache reuse feeding a >window chunk. The validated default paths (CLI single-pass with offset == 0, server default chunk 512 < 1024) are unaffected, which is why validation passed. Gemma 4 was never exposed because trim_mask_to_keys re-crops the full mask to the window (and the trailing-window slice of the full mask equals the clamped mask).

Fix (commit 4cb2c18)

src/models/gemma3.rs: gate the full mask on a fresh prefill only.

-    if seq_len > window {
+    if seq_len > window && sliding_offset == 0 {

For seq_len > window && sliding_offset > 0 the else branch's effective_offset collapses to 0, producing the clamped [seq_len, window] mask, which both broadcasts against the window-trimmed cache and carries the correct windowed-causal values (the offset cancels in the relative query/key relation). Gemma 4 is intentionally left unchanged (protected by trim_mask_to_keys). Doc comment updated; two regression tests added ([4,0,2] -> [4,4] fresh, [4,3,2] -> [4,2] non-fresh, the assertion that catches the crash).

Re-verification (narrow selectors)

  • cargo build --release --features metal,accelerate: clean.
  • cargo test --release -p mlxcel gemma3: 20 passed (18 prior + 2 new), 0 failed.
  • cargo fmt --check: clean. cargo clippy -p mlxcel --lib --tests --features metal,accelerate -- -D warnings: clean.

Updated verdict

The original issue #401 fix (over-window single-pass prefill, offset == 0) is preserved and correct for both Gemma 3 and Gemma 4. With this follow-up, the Gemma 3 non-fresh over-window path no longer crashes. No remaining CRITICAL/HIGH findings.

@inureyes

Copy link
Copy Markdown
Member Author

PR Finalization Complete

Summary

Docs (docs/supported-models.md): verified. The implementer already replaced the stale "lower --fps / `<pad>`" caveat with the correct description ("over-window single-pass prefill is handled correctly... long clips decode coherently without lowering --fps"). No further change needed.

Doc comments (two LOW items from review):

  • gemma3.rs sliding_prefill_mask: tightened the trailing sentence to state the exact gate (seq_len > window && sliding_offset == 0) rather than only the seq_len <= window case, which omitted the non-fresh over-window (sliding_offset > 0) path.
  • utils.rs and gemma4.rs: removed two em dashes from doc comments (project style prohibition; the reviewer removed one, these were two remaining).

Tests: added mask_at helper and diagonal assertion (mask_at(q, q) == 0.0 for all rows) to gemma3_mask_tests::sliding_prefill_mask_fresh_over_window_spans_full_key_axis, bringing it to parity with gemma4's equivalent test.

Lint/format: cargo fmt -p mlxcel-core -p mlxcel clean; cargo clippy --features metal,accelerate -p mlxcel-core -p mlxcel --lib --tests -- -D warnings passed with no warnings.

No em dashes remain anywhere in the PR diff. Branch is pushed and clean.

@inureyes inureyes added status:done Completed and removed status:review Under review labels Jun 23, 2026
inureyes added 3 commits June 24, 2026 05:31
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.
@inureyes
inureyes force-pushed the fix/issue-401-gemma4-sliding-window-prefill branch from df974b1 to dad9c36 Compare June 23, 2026 20:31
@inureyes
inureyes merged commit 346018a into main Jun 23, 2026
5 checks passed
@inureyes
inureyes deleted the fix/issue-401-gemma4-sliding-window-prefill branch June 23, 2026 20:32
inureyes added a commit that referenced this pull request Jun 23, 2026
* 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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:inference Generation, sampling, decoding (incl. speculative, DRY) area:models Model architectures, weights, loading, metadata priority:medium Medium priority status:done Completed type:bug Bug fixes, error corrections, or issue resolutions

Projects

None yet

Development

Successfully merging this pull request may close these issues.

fix(models): gemma4 sliding-window single-pass prefill degenerates to <pad> beyond the window

1 participant