fix(models): dense-cache sliding-window prefill uses full mask over all retained keys (#413)

[inureyes]

Summary

Dense-KVCache sliding-window models (cohere2, gemma3n, olmo3) degenerated on prefills where the total key span exceeded the window. This makes their prefill mask construction unconditionally use the full windowed-causal mask over all retained keys, fixing both degeneration cases while leaving the rotating-cache models untouched.

Root cause

A dense KVCache retains every key and the attention layer slices K/V to the mask's key axis, but these three models routed mask construction through create_sliding_window_prefill_mask, whose full-mask gate is size > window && sliding_offset == 0 (a FRESH prefill only). Any other over-window or over-total case fell to the clamped [size, window] mask, which assumes the cache physically trimmed to window keys. For a dense cache that assumption is false, so the clamped mask is the wrong shape for the keys actually present. This is a latent degeneration, not a regression introduced by #412 / #408: the behavior is byte-identical to the pre-#408 path.

The fix

A dense KVCache keeps every key, so the correct prefill mask is ALWAYS the full [size, size + sliding_offset] windowed-causal mask over all retained keys, for every size/offset combination. The new create_sliding_window_prefill_mask_dense helper is therefore an unconditional call to create_causal_mask_with_window_full (no clamped branch). The window is enforced by the mask, not by physically dropping keys.

For the common within-window prefill (size + sliding_offset <= window) this is byte-identical to create_sliding_window_prefill_mask: the clamped builder takes its non-capped path there and produces the same tril(offset) intersect triu(offset - window + 1) band over [size, size + sliding_offset], so greedy output is unchanged.

Both degeneration cases fixed (dense caches: cohere2 / gemma3n / olmo3)

1. Over-window (size > window, at any offset): the clamped [size, window] mask's tril diagonal window - size < 0 stranded the earliest query rows (logical position < size - window) with an all--inf row, which softmaxed to NaN and decoded to <pad>.
2. Within-window but over-total (size <= window && size + sliding_offset > window): the clamped path's trailing-window K/V slice dropped the OLDEST in-window keys for the earliest query rows, silently narrowing their attention window. This produced coherent-but-wrong output rather than NaN/<pad>, so it was harder to spot, but it is the same root cause.

Prefills with size + sliding_offset <= window remain byte-identical (greedy output unchanged). RotatingKVCache models (gpt_oss, mellum, exaone4, exaone_moe, ministral3, step3p5, gemma3, gemma4) and the gemma3 tensor-parallel path are untouched: they keep create_sliding_window_prefill_mask and its clamped path, because a rotating cache physically returns at most window keys once it has rolled over, so the clamped mask is the matching shape.

What changed

• create_sliding_window_prefill_mask_dense in src/lib/mlxcel-core/src/utils.rs simplified to an unconditional full mask; doc comment rewritten to explain the dense-vs-rotating divergence and both fixed cases.
• Reroute / keep the five dense call sites on the dense helper: src/models/cohere2.rs (two forward variants), src/models/gemma3n.rs (two forward variants), src/models/olmo3.rs (one); call-site comments updated to describe the unconditional full mask.
• Used by: lines updated: create_sliding_window_prefill_mask no longer lists the dense models; create_causal_mask_with_window_full lists the dense prefill consumers.
• Rotating-cache call sites left untouched: src/distributed/tensor_parallel/llama_runtime.rs (gemma3 TP), gemma3.rs, gemma4.rs, gpt_oss.rs, mellum.rs, exaone4.rs, exaone_moe.rs, ministral3.rs, step3p5.rs.
• Tests in the utils test module: sliding_window_prefill_mask_dense_full_when_rolled_over_window (over-window NaN case) and sliding_window_prefill_mask_dense_matches_fresh_over_window (parity with the rotating helper for a fresh over-window prefill); sliding_window_prefill_mask_dense_within_window_byte_identical now covers within-total inputs (3, 0, 8) and (3, 2, 8); new sliding_window_prefill_mask_dense_full_when_within_window_over_total pins the within-window over-total fix (full [2, 11] mask, not clamped [2, 4]; row 0 at logical position 9 attends to its oldest in-window key at column 6 that the clamped path dropped).

Test plan

• cargo test --release -p mlxcel-core utils:: (34 passed, 0 failed; includes the four sliding_window_prefill_mask_dense_* tests)
• cargo check --lib (root crate compiles with the rerouted call sites)
• cargo clippy --lib --tests -- -D warnings (clean)
• cargo fmt --check (clean)

Closes #413

[inureyes]

Real-model validation (gemma3n-e2b-4bit, before/after)

Validated on a real checkpoint per the inference-change rule, using the single-node server's chunked prefill to drive sliding_offset > 0. gemma3n has a 512-token sliding window; with --prefill-chunk-size 768 and a 2631-token prompt, prefill chunk 2 (tokens[768:1536]) hits the bug case directly: size = 768 > window = 512 at offset = 768.

Identical request (/v1/completions, temperature 0, max_tokens 80, same 2631-token prompt) against both binaries:

• Unfixed (main 4d5f4fb30): output is 80 consecutive <pad> tokens. The clamped [size, window] mask strands the earliest chunk-2 query rows with an all--inf row; the resulting NaN cascades forward through the overlapping sliding windows to the final prompt position, so the first decoded logits are NaN and argmax degenerates to <pad>. This is exactly the degeneration fix(models): dense-cache sliding-window models still degenerate on rolled-over over-window prefill (offset>0 && seq>window) #413 describes.
• Fixed (this PR 60dfe2bc3): coherent, on-topic output, "The most important theme is that the history of computing is a long chain of incremental ideas, each building on the last." (Mild repetition is expected from a 2B model at greedy temperature 0 and is unrelated to the fix.)

The deterministic mask-layer unit tests cover both fixed cases at the exact layer of the bug: the over-window NaN case (sliding_window_prefill_mask_dense_full_when_rolled_over_window) and the latent within-window-but-over-total stranded-oldest-keys case (sliding_window_prefill_mask_dense_full_when_within_window_over_total), plus byte-identical parity for within-total prefills.

[inureyes]

Implementation Review Summary

Intent

Stop dense-KVCache sliding-window models (cohere2, gemma3n, olmo3) from degenerating on over-window prefills by always building the full windowed-causal mask over all retained keys.

Findings Addressed

None. The implementation is correct and complete as submitted; no fixes were required.

Verification

• All stated requirements implemented (both degeneration cases fixed: over-window NaN/<pad>, and within-window over-total silent window narrowing)
• No placeholder/mock code remaining (create_sliding_window_prefill_mask_dense is an unconditional, real call to create_causal_mask_with_window_full)
• Integrated into project code flow (rerouted at all 5 dense call sites: cohere2 x2, gemma3n x2, olmo3 x1)
• Project conventions followed (Used by: updated on both create_sliding_window_prefill_mask and create_causal_mask_with_window_full; new helper carries its own Used by: line)
• Existing modules reused where applicable (delegates to create_causal_mask_with_window_full, no duplicated mask logic)
• No unintended structural changes (4 files touched; adds a sibling helper plus tests, no restructuring)
• Tests pass (4 dense mask-layer tests cover over-window, within-window over-total, fresh-over-window parity, and within-total byte-identical parity)

Correctness notes

• The 8 RotatingKVCache consumers (gpt_oss, mellum, exaone4, exaone_moe, ministral3, step3p5, gemma3, gemma4) and the gemma3 tensor-parallel path keep the original clamped create_sliding_window_prefill_mask; verified all 8 use RotatingKVCache for their sliding layers, so the clamped shape is correct for them.
• The three rerouted models use a plain dense KVCache (zero RotatingKVCache references) and slice K/V to the mask's key axis (k_len > mask_klen), so the full mask makes slicing a no-op and every retained key stays visible.
• Within-total parity (size + sliding_offset <= window) is byte-identical: the old else branch computes effective_offset == sliding_offset and create_causal_mask_with_window takes its non-capped path, producing the same band as create_causal_mask_with_window_full.

Closes #413 is present. No CRITICAL/HIGH/MEDIUM/LOW findings to escalate.

[inureyes]

Security and performance review (#413 dense sliding-window prefill mask)

Reviewed the mask change and all five rerouted call sites (cohere2 x2, gemma3n x2, olmo3 x1) plus the shared create_causal_mask_with_window_full builder. No CRITICAL or HIGH findings. Summary:

Performance: acceptable / inherent (no decode regression)

• The change is confined to the multi-token prefill mask (l > 1). Single-token decode builds no mask (the l > 1 gate yields None) and routes to causal_attention, so the decode hot path is untouched: no extra allocations, no extra graph nodes.
• For dense-cache models the full [size, size + offset] mask is wider than the old clamped [size, window] mask once size + offset > window, so prefill attention now runs over all retained keys instead of the trailing window. This is the correct and necessary behavior: the K/V are already resident (dense KVCache retains every key, an existing design choice this PR does not change), and the clamped path it replaces produced NaN/<pad> or silently narrowed output. The prior fix: harden sliding-window >window prefill across non-Gemma models; guard causal_attention windowed K/V slice #408 fresh-over-window path already used the full mask. The attention K/V slice keys off mask_klen, so the full mask keeps all keys with no slice.
• Optional future optimization (not a defect, out of scope): a multi-token chunk only needs the trailing size + window - 1 keys, so prefill attention width could be bounded to min(size + offset, size + window - 1) rather than size + offset, which would matter for long-context dense sliding-window serving. This adds K/V banding complexity and its own correctness risk, so it belongs in a separate change.

Integer arithmetic: safe for realistic contexts

• size + offset and offset - window + 1 are i32. Overflow needs > 2.1B tokens, far beyond any reachable KV-cache memory limit (mlxcel also caps admission at 1M tokens). No realistic overflow/underflow.

Within-window parity: confirmed

• For size + offset <= window the new helper is byte-identical to the rotating helper (the clamped builder takes its non-capped path and produces the same band over [size, size + offset]), so greedy output for the common prefill is unchanged. The four sliding_window_prefill_mask_dense_* tests cover this and both divergent cases.

Note (preexisting, out of scope): --max-kv-size + chunked prefill

• The dense prefill masks are built from the monotonic cache.offset. After enforce_max_kv_size_for calls trim_front (server path, --max-kv-size below the prompt length), offset keeps growing while update_and_fetch returns only the live_len = offset - live_start window, so a full [l, l + offset] mask would be wider than the live K/V on a later l > 1 continuation chunk.
• This is NOT introduced by this PR and does not change the end-to-end outcome: the untouched global-attention path in the same models builds create_causal_mask(l, ga_offset) from the same monotonic offset with window_size == 0 (never sliced), so a forward under that config already mismatches at a global layer regardless of this change. Suggest a separate follow-up to align dense prefill masks (global and sliding) with live_len() under trim; the byte-identical-when-live_start == 0 property makes that a safe, isolated fix.

No fixes applied: the PR is correct within its stated scope and validated configuration, and the only latent item is preexisting and broader than this diff.

[inureyes]

PR Finalization Complete

Tests

4 new unit tests for create_sliding_window_prefill_mask_dense in src/lib/mlxcel-core/src/utils.rs:

• sliding_window_prefill_mask_dense_full_when_rolled_over_window: the primary fix, over-window prefill at non-zero offset produces a full [size, size + offset] mask with no all-masked rows
• sliding_window_prefill_mask_dense_full_when_within_window_over_total: the adjacent fix, within-window chunk whose offset pushes total past window keeps the oldest in-window keys
• sliding_window_prefill_mask_dense_matches_fresh_over_window: identity with the rotating helper at sliding_offset == 0
• sliding_window_prefill_mask_dense_within_window_byte_identical: byte-identical to rotating helper for all within-total cases (greedy output unchanged)

All 34 utils:: tests pass. No gaps found; no padding added.

Documentation

No existing doc is inaccurate due to this fix. The sliding-window mask selection behavior (dense vs rotating) is not documented at the doc level, and the change is purely internal. The Used by: comments in utils.rs are the correct discovery mechanism per the code-guidelines convention and are already updated in the PR.

Lint/Format

• cargo fmt --check: clean
• cargo clippy --lib --tests -- -D warnings: clean
• cargo check --lib: clean
• No debug artifacts, stray comments, or todo!()/dbg!() in changed files

All checks passing. Ready for merge.