fix: harden sliding-window over-window prefill across models

[inureyes]

Summary

A single-pass prefill longer than a model's sliding window degenerated to NaN / <pad> (or hard-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 [broadcast_shapes] crash instead. PR #405 fixed only gemma3/gemma4.

Chosen approach: general fix + shared helper (no loud-guard fallback needed)

The general causal_attention fix landed cleanly and is bit-exact for every currently-correct path, so the loud-guard fallback was not required.

1. 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. This is the shared safety net and directly fixes baichuan (which reaches this path with mask.is_none()), and incidentally also fixes the tensor-parallel gemma4 path (which routes through causal_attention).

2. Shared helper create_sliding_window_prefill_mask(size, sliding_offset, window) (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. Its else-branch is byte-identical to the per-model sliding_offset.min((window - size).max(0)) construction it replaces (the capped tril offset window - size is offset-independent, and the helper's non-capped triu is a no-op for size <= window), so greedy temp-0 output is unchanged for every non-degenerate case.

3. Per-model wiring of the helper into the prefill mask builders (see audit table). gemma3/gemma4 keep their fix(models): correct gemma sliding-window prefill beyond the window #405 sliding_prefill_mask helpers as-is; refactor: hoist duplicated sliding_prefill_mask helper from gemma3/gemma4 to a shared module #410 will consolidate them onto the shared helper.

4. Tensor-parallel gemma3 path (src/distributed/tensor_parallel/llama_runtime.rs, gemma3_masks): the TP path independently computed its own clamped offset and called create_causal_mask_with_window, so a TP gemma3 >window prefill hit the same NaN/crash as the single-process path. Replaced with create_sliding_window_prefill_mask, mirroring the single-process fix. TP gemma4 was incidentally fixed by the general causal_attention change in item 1.

Per-model audit (every RotatingKVCache / create_causal_mask_with_window sliding-window caller)

Model	Cache	Pre-fix >window prefill	Fix applied	How validated
gpt_oss	RotatingKVCache, full K/V	CRASH (mask [N,128] vs scores [..,N,N])	helper at both mask sites (forward + pipeline stage)	real model, window 128
mellum	RotatingKVCache, full K/V	CRASH (capped mask vs full K/V)	helper at prefill mask builder	unit test + code analysis (no local ckpt)
exaone4	RotatingKVCache, full K/V (bf16 cast)	CRASH	helper (astype preserved)	unit test + analysis (local ckpt not sliding)
exaone_moe	RotatingKVCache, full K/V	CRASH	helper	unit test + analysis (no local ckpt)
ministral3	RotatingKVCache, full K/V	CRASH	helper	unit test + analysis (local ckpt not sliding)
step3p5	RotatingKVCache, full K/V	CRASH	helper	unit test + analysis (no local ckpt)
cohere2	dense KVCache + model-level K/V slice	NaN (capped mask + trailing slice strands early rows)	helper + slice to mask key-axis (not blindly window)	unit test + analysis (window 4096, not quick to exceed)
gemma3n	dense KVCache + model-level K/V slice	NaN	helper at both mask sites + slice to mask key-axis	unit test + analysis
olmo3	dense KVCache + model-level K/V slice	NaN (capped mask + trailing slice strands early rows)	helper + slice to mask key-axis (not blindly window)	unit test + analysis
baichuan	dense KVCache, prefill via causal_attention(mask=None)	NaN / shape issue	covered by the general causal_attention fix (no model edit)	unit test + analysis
tensor-parallel gemma3	RotatingKVCache via gemma3_masks in llama_runtime.rs	NaN / CRASH (same clamped construction as single-process)	create_sliding_window_prefill_mask in gemma3_masks	logit-match test (tensor_parallel_gemma3_matches_full_model_logits)
tensor-parallel gemma4	routes through causal_attention	NaN / CRASH	incidentally fixed by the general causal_attention fix (item 1 above)	covered by the general causal_attention fix
recurrent_gemma	KVCache, full causal mask, window_size = 0	unaffected (never caps; window not mask-enforced)	none	code analysis
diffusion_gemma	dense KVCache, dense_windowed_causal_mask (uncapped)	unaffected (already full windowed mask)	none	code analysis
gemma3 / gemma4	RotatingKVCache, sliding_prefill_mask	already fixed by #405	untouched	real model (gemma-4-12b) + existing tests

ministral-3b-4bit and exaone4-1.2b-4bit ship with sliding_window: null (no windowed layers), so those local checkpoints do not exercise the windowed path; their fix rests on the unit test, the byte-identical helper proof, and the gpt-oss real-model result for the identical code shape.

Real-model validation (M1 Ultra, --temp 0)

Bug existed on main (HEAD 346018a), gpt-oss-20b-mxfp4, 298-token prompt (window 128):

exit=134
libc++abi: terminating due to uncaught exception of type std::invalid_argument:
[broadcast_shapes] Shapes (298,128) and (1,64,298,298) cannot be broadcast.

Same prompt on this branch:

...<|channel|>final<|message|>The bakery sits on the second hill, and the clock tower rings thirteen times at noon.
[Generated 90 tokens in 1.57s = 57.36 tok/s]

gemma-4-12b-it-4bit, ~1600-token prompt (window 1024), this branch (#405 not regressed):

The bakery is located on the second hill. At noon, the clock tower rings thirteen times as a joke that began two centuries ago.

Regression (byte-identical main vs this branch, generated text diffed):

• gpt-oss-20b short prompt (<128, prefill + decode): BYTE-IDENTICAL
• gemma-4-12b short prompt ("The capital of France is Paris."): BYTE-IDENTICAL
• qwen3-0.6b non-windowed ("Mercury / Venus / Earth"): BYTE-IDENTICAL

This confirms decode (q_len == 1), non-windowed (window_size == 0), and within-window prefill paths are unchanged; only the previously-degenerate >window prefill changed (crash/NaN -> coherent).

Tests (fast, no model loads)

• utils::tests::sliding_window_prefill_mask_selects_full_when_fresh_over_window / _clamps_when_rolled_over / _within_window_matches_capped_builder (the selector's three regimes).
• layers::tests::causal_attention_multi_token_over_window_uses_full_windowed_mask (the q_len > 1, k_len > window case: finite, matches the full windowed-mask reference; the prior clamp+slice path produced NaN). Updated from the old test that asserted the now-removed sliced behavior.
• Existing decode/window tests, all utils::tests (30) and layers::tests (28), and gemma3/gemma4 sliding_prefill_mask tests still pass.
• tensor_parallel_gemma3_matches_full_model_logits validates the TP gemma3 mask path.

Test plan

• cargo clippy --features metal,accelerate -p mlxcel-core -p mlxcel --lib --tests -- -D warnings (clean)
• cargo test --release -p mlxcel-core sliding_window_prefill_mask and causal_attention (all pass)
• cargo test --release -p mlxcel-core "utils::tests" / "layers::tests" (30 / 28 pass)
• cargo test --release -p mlxcel sliding_prefill_mask (gemma3/gemma4, 4 pass)
• cargo test --release -p mlxcel tensor_parallel_gemma3 (logit-match test passes)
• real-model gpt-oss-20b >window (crash on main -> coherent here), gemma-4-12b >window coherent, byte-identical short/decode regression

Known follow-ups

• fix(models): dense-cache sliding-window models still degenerate on rolled-over over-window prefill (offset>0 && seq>window) #413 (pre-existing, out of scope here): dense-cache offset > 0 rolled-over over-window path; in that case the cache already trims to window keys so the clamped mask is the correct shape, but this has not been confirmed safe under every model's K/V trimming logic.
• refactor: hoist duplicated sliding_prefill_mask helper from gemma3/gemma4 to a shared module #410 consolidation: migrate gemma3/gemma4 sliding_prefill_mask onto the shared create_sliding_window_prefill_mask.
• recurrent_gemma applies its local-attention window via a full causal mask with window_size = 0 (window not enforced by the mask); orthogonal to this issue, worth a separate look.

Closes #408

[inureyes]

Implementation Review Summary

Intent

Harden sliding-window >window single-pass prefill across all models by fixing the shared causal_attention, adding create_sliding_window_prefill_mask, and wiring per-model fixes (closes #408).

Findings Addressed

• olmo3 missed (CRITICAL) — src/models/olmo3.rs is a registered, reachable sliding-window model on a dense KVCache with the exact over-window prefill NaN/<pad> bug this PR fixes for cohere2/gemma3n: it built the clamped (l, window) mask via create_causal_mask_with_window and sliced K/V to the trailing window_size keys, stranding the earliest query rows. It was not in the audit table and not covered by the general causal_attention fix (its windowed path goes through the explicit-mask branch, and it passes window_size = 0 to causal_attention). Fixed by mirroring the cohere2/gemma3n pattern exactly: build the mask via create_sliding_window_prefill_mask, and slice K/V to the mask's key axis (mask_klen) instead of blindly to window_size. Used by: on the helper corrected to add Olmo3 and drop Baichuan (baichuan does not call the helper).

Verification

• All stated requirements implemented (now including olmo3)
• No placeholder/mock code remaining
• Integrated into project code flow (mask builders + attention K/V slicing wired per model)
• Project conventions followed (no em dashes, no AI attribution; // Used by: updated on causal_attention and the new helper)
• Existing modules reused (create_sliding_window_prefill_mask / create_causal_mask_with_window_full, no parallel builder)
• No unintended structural changes
• Tests pass (helper selector 3/3, causal_attention_multi_token_over_window_uses_full_windowed_mask)

Review notes (verified correct, no change needed)

• Bit-exactness: over_window_prefill = needs_window_mask && q_len > 1 is the only new branch. For q_len == 1 decode, window_size == 0, and within-window (k_len <= window) prefills, over_window_prefill is always false and both the K/V-slice selection and the mask chain are byte-identical to main. softcap > 0 composes correctly: the guard is independent of softcap, and softcap is still passed through to attention when the full windowed mask is built.
• Helper else-branch: byte-identical to the prior per-model sliding_offset.min((window - size).max(0)) construction. For step3p5 (which previously passed the raw offset) the resulting mask is still byte-identical, because the capped path's tril diagonal window - size is offset-independent and the non-capped triu term is a no-op for size <= window.
• Dense-cache fixes (cohere2, gemma3n): mask_klen correctly tracks the mask key axis (full mask keeps all keys, clamped mask drops the oldest). Both gemma3n sites (forward + forward_with_embeddings) handled.
• baichuan: covered by the general causal_attention fix on the mask.is_none() prefill path; no separate clamped-mask site. The M5 tile-aligned padded-mask path does not crash (full causal mask matches full K/V) and its lack of window enforcement there is pre-existing and orthogonal to fix: harden sliding-window >window prefill across non-Gemma models; guard causal_attention windowed K/V slice #408.

Out of scope (pre-existing, flagged for a separate chore)

• src/lib/mlxcel-core/build.rs:90 trips clippy::needless_borrows_for_generic_args under the pinned toolchain 1.93.1 (bridge.flag_if_supported(&format!(...)) should drop the &). It is byte-identical to main and unrelated to fix: harden sliding-window >window prefill across non-Gemma models; guard causal_attention windowed K/V slice #408, but a mlxcel-core recompile re-emits it under -D warnings. Recommend a one-line chore fix separately rather than folding it into this PR.