fix(models): dense-cache sliding-window models still degenerate on rolled-over over-window prefill (offset>0 && seq>window)

[inureyes]

Problem / Background

PR #412 (#408) hardened sliding-window >window single-pass prefill across models. For dense-KVCache sliding-window models (cohere2, gemma3n, olmo3) it routes mask construction through the shared create_sliding_window_prefill_mask(size, sliding_offset, window) helper and slices K/V to the mask's key axis.

That helper only builds the uncapped full windowed mask when size > window && sliding_offset == 0 (a FRESH prefill). For sliding_offset > 0 && seq_len > window (a rolled-over / continued prefill, e.g. a multi-turn continuation chunk longer than the window with prior cache, or a chunked-prefill chunk larger than the window) it falls through to the clamped [seq_len, window] mask.

This was explicitly left out of #412's scope. It is pre-existing (NOT introduced by #412): the behavior is byte-identical to the pre-#412 code path, so this is a latent degeneration, not a regression.

The bug

A dense KVCache retains ALL keys (unlike a RotatingKVCache, which trims to window). With the clamped [seq_len, window] mask, whose tril diagonal window - seq_len < 0, and the model slicing K/V to the trailing window keys, the earliest query rows (logical position < seq_len - window) are stranded with an all--inf row. That produces NaN, which decodes to <pad>.

So a rolled-over over-window prefill on cohere2 / gemma3n / olmo3 degenerates to <pad> / NaN output for the early positions.

Root cause (the asymmetry)

The shared helper's gate size > window && sliding_offset == 0 is NARROWER than the two other windowed-mask gates in the tree:

• causal_attention's windowed path gates on q_len > 1 && k_len > window (no offset gate; the offset case was fixed in fix: harden sliding-window >window prefill across non-Gemma models; guard causal_attention windowed K/V slice #408).
• gemma4's sliding_prefill_mask gates on l > window (no offset gate; protected by trim_mask_to_keys).

Models that reach causal_attention's windowed path (e.g. baichuan, gemma4) therefore handle the offset>0 over-window case correctly. The dense-cache models that build the helper mask directly do not, because of the extra sliding_offset == 0 condition.

Proposed Solution

Extend create_sliding_window_prefill_mask (or the dense-cache call sites) to build the correct full windowed mask for the sliding_offset > 0 && seq_len > window case as well: the full [seq_len, seq_len + sliding_offset] windowed-causal mask over all retained keys, mirroring create_causal_mask_with_window_full with a non-zero offset. Keep the dense-cache mask-key-axis slicing consistent (a full mask keeps every key).

Coordinate with #410 (the sliding_prefill_mask consolidation), since both touch this helper and its gate. Verify that gemma3 / gemma4 single-process and the RotatingKVCache models are not affected by the helper change (RotatingKVCache trims to window, so its offset>0 path differs and must not be regressed).

Acceptance Criteria

• A rolled-over over-window prefill (offset>0, seq>window) on a dense-cache sliding-window model (cohere2 / gemma3n / olmo3) produces coherent output instead of <pad> / NaN.
• Fresh over-window and within-window paths are unchanged.
• RotatingKVCache models and gemma3 / gemma4 are not regressed.
• A regression test at the helper / mask layer covers the offset>0 over-window case.

Technical Considerations

• Dense KVCache vs RotatingKVCache divergence is the crux: the helper change must keep every retained key for dense caches while leaving the trimming-cache path alone.
• Touches the same helper as refactor: hoist duplicated sliding_prefill_mask helper from gemma3/gemma4 to a shared module #410; sequence the work so the two changes do not conflict on the gate.

[inureyes]

Per maintainer request, the scope of PR #418 was expanded to also fix the adjacent latent case size <= window && size + sliding_offset > window on dense caches (a within-window chunk whose offset pushes the total key span past the window), rather than deferring it as a follow-up.

This simplified 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. create_sliding_window_prefill_mask_dense is now an unconditional call to create_causal_mask_with_window_full (no clamped branch). For size + sliding_offset <= window it is byte-identical to the rotating helper, so within-total greedy output is unchanged.

The acceptance criterion "Fresh over-window and within-window paths are unchanged" is refined to "within-window AND within-total (size + sliding_offset <= window) paths are byte-identical". The previously-deferred within-window over-total case (stranded oldest in-window keys, coherent-but-wrong output) is now fixed in the same PR. RotatingKVCache models and gemma3 / gemma4 / the tensor-parallel path keep create_sliding_window_prefill_mask and its clamped path and are not regressed.