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

src/distributed/tensor_parallel/llama_runtime.rs

@@ -2757,11 +2757,17 @@ fn gemma3_masks(
     let sliding_mask = if rank0.sliding_window_pattern > 1 {
         let sliding_offset = gemma3_cache_offset(&mut rank0_caches[0]);
         let max_cache = rank0.sliding_window as i32;
-        let effective_offset = sliding_offset.min((max_cache - seq_len).max(0));
-        Some(mlxcel_core::utils::create_causal_mask_with_window(
+        // Mirror the single-process gemma3 `sliding_prefill_mask` (issue #401/#408):
+        // a fresh single-pass prefill longer than the window keeps every key (the
+        // RotatingKVCache returns all of them) and needs the full-width windowed
+        // mask, otherwise the clamped `(seq_len, window)` mask strands the earliest
+        // query rows with an all-(-inf) row and the broadcast crashes against the
+        // full key axis. The helper's else-branch is byte-identical to the prior
+        // `sliding_offset.min((max_cache - seq_len).max(0))` clamped construction.
+        Some(mlxcel_core::utils::create_sliding_window_prefill_mask(
             seq_len,
-            effective_offset,
-            Some(max_cache),
+            sliding_offset,
+            max_cache,
         ))
     } else {
         None

src/lib/mlxcel-core/src/layers.rs

@@ -3461,7 +3461,14 @@ mod tests {
     }
 
     #[test]
-    fn causal_attention_window_bool_mask_matches_float_reference() {
+    fn causal_attention_multi_token_over_window_uses_full_windowed_mask() {
+        // Multi-token prefill (`q_len = 3`) whose key length (`k_len = 5`)
+        // exceeds the window. Before issue #408 `causal_attention` sliced K/V
+        // to the trailing `window_size` keys and used the clamped
+        // `(q_len, window_size)` mask, which strands the earliest query rows.
+        // The corrected behavior keeps all keys and builds the full
+        // `(q_len, k_len)` windowed-causal mask, so the reference must use the
+        // full K/V and `create_causal_mask_with_window_full`.
         let q = crate::from_slice_f32(
             &[
                 0.1, -0.2, 0.3, 0.4, 0.2, 0.0, -0.1, 0.5, -0.3, 0.1, 0.2, -0.4,
@@ -3483,26 +3490,22 @@ mod tests {
             &[1, 1, 5, 4],
         );
         let scale = 0.5_f32;
-        // Use window_size=3 (not 2) so the post-earlier cap path produces a
-        // non-degenerate mask: with q_len=3, window=3, tril_offset = w - size
-        // = 0, so q=0 attends to k=0. With window=2 the row-0 mask is fully
-        // masked, which produces NaN under the now-correct semantics (a
-        // query whose logical position predates every cache key).
-        let window_size = 3_i32;
+        // window_size=2 is the previously-degenerate case: under the old
+        // clamp+slice path query row 0 (logical position 0) had an all-`-inf`
+        // mask row and softmaxed to NaN. The full mask gives it its own
+        // window, so no row is fully masked.
+        let window_size = 2_i32;
 
         let out = crate::causal_attention(&q, &k, &v, scale, 0.0, window_size);
-        // See note in `causal_attention_single_query_window_respects_mask_when_needed`:
-        // the explicit-mask reference must slice K/V to the last `window_size`
-        // slots to line up with the post-cap mask shape `(q_len, window_size)`.
-        let mask_f = crate::utils::create_causal_mask_with_window(3, 2, Some(window_size));
-        let k_sliced = crate::slice(&k, &[0, 0, 5 - window_size, 0], &[1, 1, 5, 4]);
-        let v_sliced = crate::slice(&v, &[0, 0, 5 - window_size, 0], &[1, 1, 5, 4]);
+        // Reference: full key set + uncapped windowed mask over offset =
+        // k_len - q_len = 2.
+        let mask_full = crate::utils::create_causal_mask_with_window_full(3, 2, Some(window_size));
         let out_ref = attention(
             &q,
-            &k_sliced,
-            &v_sliced,
+            &k,
+            &v,
             scale,
-            Some(mask_f.as_ref().unwrap()),
+            Some(mask_full.as_ref().unwrap()),
             0.0,
             window_size,
         );
@@ -3511,9 +3514,21 @@ mod tests {
         let diff_abs = crate::abs(&diff);
         let diff_sum = crate::sum_all(&diff_abs);
         crate::eval(&diff_sum);
+        let total = crate::item_f32(&diff_sum);
         assert!(
-            crate::item_f32(&diff_sum) < 1e-5,
-            "bool-mask causal window path should match float-mask reference"
+            total.is_finite() && total < 1e-5,
+            "multi-token over-window causal attention should match the full \
+             windowed-mask reference (no NaN, no stranded rows); diff_sum = {total}"
+        );
+
+        // The output itself must be finite (the pre-#408 path produced NaN for
+        // the earliest query row under window_size=2).
+        let out_abs = crate::abs(out.as_ref().unwrap());
+        let out_sum = crate::sum_all(&out_abs);
+        crate::eval(&out_sum);
+        assert!(
+            crate::item_f32(&out_sum).is_finite(),
+            "multi-token over-window causal attention output must be finite"
         );
     }
 

src/lib/mlxcel-core/src/lib.rs

@@ -2570,6 +2570,14 @@ pub fn fast_rope_batched(
 /// explicit mask array. Other hardware keeps using MLX's native causal SDPA
 /// entry point.
 ///
+/// Sliding-window handling splits by query length: a single-query decode
+/// (`q_len == 1`) slices K/V to the trailing `window_size` keys and uses the
+/// clamped `(1, window_size)` mask, while a multi-token prefill that exceeds
+/// the window (`q_len > 1 && k_len > window_size`) keeps all keys and builds a
+/// full-width windowed-causal mask over them. Slicing K/V for a multi-token
+/// prefill would strand the earliest query rows with an all-`-inf` row -> NaN
+/// (issue #401/#408); the windowed correctness instead comes from the mask.
+///
 /// Used by: Llama, Qwen, Mixtral, Gemma, Cohere, Phi, OLMo, Exaone, GLM4,
 /// MiniCPM, DeepSeek, Hunyuan, StarCoder2 and other causal prefill call sites
 pub fn causal_attention(
@@ -2593,13 +2601,25 @@ pub fn causal_attention(
     }
 
     if softcap > 0.0 || window_size > 0 {
-        // When the window is exceeded, slice K/V to the last `window_size`
-        // slots so they line up with the (q_len, window_size)-shaped mask
-        // produced by `create_causal_mask_with_window`. Production callers
-        // backed by `RotatingKVCache` already deliver K/V truncated to
-        // `max_size = window_size`; this branch keeps the wrapper
-        // self-consistent for any other caller (and for our own unit tests).
-        let (k_used, v_used, effective_k_len) = if needs_window_mask {
+        // A multi-token prefill that exceeds the window must NOT slice K/V to
+        // the trailing `window_size` keys: a fresh `RotatingKVCache` (and a
+        // dense `KVCache`) keep every prefill key, and a trailing slice strands
+        // the earliest query rows (logical position `< k_len - window_size`)
+        // with no visible key, producing an all-`-inf` softmax row -> NaN /
+        // `<pad>` (issue #401/#408). For that case build a full-width
+        // windowed-causal mask over ALL keys instead; windowed correctness
+        // comes from the mask, mirroring mlx-lm's `RotatingKVCache`. The
+        // single-query decode path (`q_len == 1`) keeps the trailing-window K/V
+        // slice fast path, byte-for-byte unchanged.
+        let over_window_prefill = needs_window_mask && q_len > 1;
+
+        let (k_used, v_used, effective_k_len) = if needs_window_mask && !over_window_prefill {
+            // Decode (`q_len == 1`) beyond the window: slice K/V to the last
+            // `window_size` keys so they line up with the (1, window_size) mask
+            // produced by `create_causal_mask_with_window`. Production callers
+            // backed by `RotatingKVCache` already deliver K/V truncated to
+            // `max_size = window_size`; this keeps the wrapper self-consistent
+            // for any other single-query caller.
             let k_shape = ffi::array_shape(k);
             let v_shape = ffi::array_shape(v);
             let start = k_len - window_size;
@@ -2621,7 +2641,15 @@ pub fn causal_attention(
         let v_ref: &MlxArray = v_used.as_deref().unwrap_or(v);
 
         let offset = (effective_k_len - q_len).max(0);
-        let mask = if softcap == 0.0 && use_bool_causal_mask_path() {
+        let mask = if over_window_prefill {
+            // Full-width windowed-causal mask over all `k_len` keys (same
+            // builder as the gemma3/gemma4 `sliding_prefill_mask`, issue #401).
+            // Only reached for the previously-degenerate `q_len > 1`,
+            // `k_len > window` case, so there is no decode bit-exactness
+            // constraint here; the experimental bool-mask path is intentionally
+            // bypassed (this case has no all-`-inf` rows left to reproduce).
+            utils::create_causal_mask_with_window_full(q_len, offset, Some(window_size))
+        } else if softcap == 0.0 && use_bool_causal_mask_path() {
             if window_size > 0 {
                 utils::create_causal_bool_mask_with_window(q_len, offset, Some(window_size))
             } else {

src/lib/mlxcel-core/src/utils.rs

@@ -626,6 +626,45 @@ pub fn create_causal_mask_with_window_full(
     ffi::where_cond(&bool_mask, &zeros, &neg_inf)
 }
 
+/// Build the sliding-window attention mask for a multi-token prefill
+/// (`size > 1`), sized to the keys the cache actually returns.
+///
+/// A fresh single-pass prefill that exceeds the window is the degenerate case
+/// behind issue #401/#408: both [`RotatingKVCache`] (on its first prefill
+/// append) and a dense `KVCache` keep *every* prefill key, so the window must
+/// be enforced by the mask over the full key axis, not by capping the mask and
+/// dropping keys. For that case (`size > window && sliding_offset == 0`) this
+/// returns the uncapped [`create_causal_mask_with_window_full`] `[size, size]`
+/// mask. In every other case (within-window prefill, or a rolled-over rotating
+/// cache that already returns at most `window` keys) it returns the clamped
+/// [`create_causal_mask_with_window`] mask with the same
+/// `sliding_offset.min((window - size).max(0))` adjustment the per-model mask
+/// builders used before, so greedy output stays byte-identical there.
+///
+/// The `sliding_offset == 0` gate mirrors the gemma3 prior art: only a fresh
+/// prefill is guaranteed to hold all `size` keys; once a rotating cache has
+/// rolled over (`sliding_offset > 0`) it returns at most `window` keys and the
+/// clamped mask is the matching shape. Models whose attention layer slices K/V
+/// to the trailing window must slice to the mask's key dimension (not blindly
+/// to `window`) so they keep the full key set when this returns the full mask.
+///
+/// Used by: GptOss, Mellum, Exaone4, ExaoneMoE, Ministral3, Step3P5, Cohere2,
+/// Gemma3n, Olmo3 sliding-window prefill mask construction.
+///
+/// [`RotatingKVCache`]: crate::cache::RotatingKVCache
+pub fn create_sliding_window_prefill_mask(
+    size: i32,
+    sliding_offset: i32,
+    window: i32,
+) -> UniquePtr<MlxArray> {
+    if size > window && sliding_offset == 0 {
+        create_causal_mask_with_window_full(size, 0, Some(window))
+    } else {
+        let effective_offset = sliding_offset.min((window - size).max(0));
+        create_causal_mask_with_window(size, effective_offset, Some(window))
+    }
+}
+
 /// Create a boolean causal attention mask with optional sliding window.
 /// Used by: same as `create_causal_mask_with_window` (experimental path)
 ///
@@ -1272,6 +1311,88 @@ mod tests {
         }
     }
 
+    // --- Sliding-window prefill mask selector (issue #408) ----------------
+
+    /// A fresh single-pass prefill (`sliding_offset == 0`) longer than the
+    /// window must select the uncapped full `[size, size]` mask, so no early
+    /// query row is stranded with an all-`-inf` row.
+    #[test]
+    fn sliding_window_prefill_mask_selects_full_when_fresh_over_window() {
+        let prefill = create_sliding_window_prefill_mask(4, 0, 2);
+        assert_eq!(
+            ffi::array_shape(&prefill),
+            vec![4, 4],
+            "fresh over-window prefill must use the full [size, size] mask"
+        );
+        let full = create_causal_mask_with_window_full(4, 0, Some(2));
+        let at =
+            |m: &MlxArray, q: i32, k: i32| ffi::item_f32(&ffi::slice(m, &[q, k], &[q + 1, k + 1]));
+        for q in 0..4 {
+            for k in 0..4 {
+                let a = at(&prefill, q, k);
+                let b = at(&full, q, k);
+                assert_eq!(a.is_finite(), b.is_finite(), "cell ({q},{k}) finiteness");
+                if a.is_finite() {
+                    assert_eq!(a, b, "cell ({q},{k}) value");
+                }
+            }
+            // No row is fully masked: each attends to at least its own key.
+            assert_eq!(at(&prefill, q, q), 0.0, "row {q} must attend to itself");
+        }
+    }
+
+    /// A rolled-over cache (`sliding_offset > 0`) keeps the clamped path,
+    /// byte-identical to the per-model `create_causal_mask_with_window`
+    /// construction it replaced (same `min((window - size).max(0))` clamp).
+    #[test]
+    fn sliding_window_prefill_mask_clamps_when_rolled_over() {
+        let window = 4;
+        let size = 2;
+        let sliding_offset = 9; // well past the window → clamped path
+        let prefill = create_sliding_window_prefill_mask(size, sliding_offset, window);
+        let effective_offset = sliding_offset.min((window - size).max(0));
+        let clamped = create_causal_mask_with_window(size, effective_offset, Some(window));
+        assert_eq!(
+            ffi::array_shape(&prefill),
+            ffi::array_shape(&clamped),
+            "rolled-over prefill must match the clamped mask shape"
+        );
+        let cols = ffi::array_shape(&prefill)[1];
+        let at =
+            |m: &MlxArray, q: i32, k: i32| ffi::item_f32(&ffi::slice(m, &[q, k], &[q + 1, k + 1]));
+        for q in 0..size {
+            for k in 0..cols {
+                let a = at(&prefill, q, k);
+                let b = at(&clamped, q, k);
+                assert_eq!(a.is_finite(), b.is_finite(), "cell ({q},{k}) finiteness");
+                if a.is_finite() {
+                    assert_eq!(a, b, "cell ({q},{k}) value");
+                }
+            }
+        }
+    }
+
+    /// A within-window fresh prefill (`size <= window`) selects the clamped
+    /// (here non-capped) mask, identical to `create_causal_mask_with_window`.
+    #[test]
+    fn sliding_window_prefill_mask_within_window_matches_capped_builder() {
+        let prefill = create_sliding_window_prefill_mask(3, 0, 8);
+        let capped = create_causal_mask_with_window(3, 0, Some(8));
+        assert_eq!(ffi::array_shape(&prefill), ffi::array_shape(&capped));
+        let at =
+            |m: &MlxArray, q: i32, k: i32| ffi::item_f32(&ffi::slice(m, &[q, k], &[q + 1, k + 1]));
+        for q in 0..3 {
+            for k in 0..3 {
+                let a = at(&prefill, q, k);
+                let b = at(&capped, q, k);
+                assert_eq!(a.is_finite(), b.is_finite(), "cell ({q},{k}) finiteness");
+                if a.is_finite() {
+                    assert_eq!(a, b, "cell ({q},{k}) value");
+                }
+            }
+        }
+    }
+
     // --- Windowed left-padding mask (ragged batched MTP prefill) ----------
 
     /// No-padding fast path must be byte-identical to the plain windowed mask.

src/models/cohere2.rs

@@ -23,7 +23,7 @@
 
 use mlxcel_core::generate::LanguageModel;
 use mlxcel_core::layers::{FusedQKVLinear, KVCache, LayerNorm, UnifiedEmbedding, UnifiedLinear};
-use mlxcel_core::utils::{create_causal_mask, create_causal_mask_with_window};
+use mlxcel_core::utils::{create_causal_mask, create_sliding_window_prefill_mask};
 use mlxcel_core::weights::WeightMap;
 use mlxcel_core::{MlxArray, UniquePtr};
 use serde::Deserialize;
@@ -178,16 +178,21 @@ impl Cohere2Attention {
 
         // Scaled dot-product attention
         let attn_out = if l > 1 {
-            // Prefill: use mask. If a windowed mask was supplied for a
-            // sliding-window layer, post-earlier it is shaped
-            // `(l, window_size)` (capped). Plain `KVCache` returns
-            // full-length K/V, so slice K/V to the last `window_size` slots
-            // here, mirroring `causal_attention`'s internal slicing.
+            // Prefill: use mask. A sliding-window layer's mask is either the
+            // clamped `(l, window_size)` mask or the full `(l, k_len)` mask for
+            // a fresh single-pass prefill that exceeds the window (issue #408).
+            // Plain `KVCache` returns full-length K/V, so slice K/V to the
+            // mask's key axis (not blindly to `window_size`): a full mask keeps
+            // every key, a clamped mask drops the oldest. Mirrors
+            // `causal_attention`'s internal handling.
             let k_shape = mlxcel_core::array_shape(&cache_k);
             let k_len = k_shape[2];
-            let (k_used, v_used) = if self.window_size > 0 && k_len > self.window_size {
+            let mask_klen = mask
+                .map(|m| *mlxcel_core::array_shape(m).last().unwrap_or(&k_len))
+                .unwrap_or(k_len);
+            let (k_used, v_used) = if self.window_size > 0 && k_len > mask_klen {
                 let v_shape = mlxcel_core::array_shape(&cache_v);
-                let start = k_len - self.window_size;
+                let start = k_len - mask_klen;
                 (
                     Some(mlxcel_core::slice(
                         &cache_k,
@@ -422,10 +427,14 @@ impl Cohere2Model {
             let swa_offset = caches[self.swa_idx].offset;
 
             let full = Some(create_causal_mask(l as i32, ga_offset));
-            let sliding = Some(create_causal_mask_with_window(
+            // Full-width windowed mask for a fresh single-pass prefill that
+            // exceeds the window; clamped mask otherwise. The attention layer
+            // slices K/V to the mask's key axis, so a full mask keeps all
+            // (dense `KVCache`) keys. See issue #408.
+            let sliding = Some(create_sliding_window_prefill_mask(
                 l as i32,
                 swa_offset,
-                Some(self.config.sliding_window as i32),
+                self.config.sliding_window as i32,
             ));
             (full, sliding)
         } else {
@@ -480,10 +489,14 @@ impl Cohere2Model {
             let swa_offset = caches[self.swa_idx].offset;
 
             let full = Some(create_causal_mask(l as i32, ga_offset));
-            let sliding = Some(create_causal_mask_with_window(
+            // Full-width windowed mask for a fresh single-pass prefill that
+            // exceeds the window; clamped mask otherwise. The attention layer
+            // slices K/V to the mask's key axis, so a full mask keeps all
+            // (dense `KVCache`) keys. See issue #408.
+            let sliding = Some(create_sliding_window_prefill_mask(
                 l as i32,
                 swa_offset,
-                Some(self.config.sliding_window as i32),
+                self.config.sliding_window as i32,
             ));
             (full, sliding)
         } else {