diff --git a/src/distributed/tensor_parallel/llama_runtime.rs b/src/distributed/tensor_parallel/llama_runtime.rs index def84862..f9ca3c46 100644 --- a/src/distributed/tensor_parallel/llama_runtime.rs +++ b/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 diff --git a/src/lib/mlxcel-core/src/layers.rs b/src/lib/mlxcel-core/src/layers.rs index 9431e4d9..ec25969b 100644 --- a/src/lib/mlxcel-core/src/layers.rs +++ b/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" ); } diff --git a/src/lib/mlxcel-core/src/lib.rs b/src/lib/mlxcel-core/src/lib.rs index e67fc251..59473a64 100644 --- a/src/lib/mlxcel-core/src/lib.rs +++ b/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 / + // `` (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 { diff --git a/src/lib/mlxcel-core/src/utils.rs b/src/lib/mlxcel-core/src/utils.rs index c3ebaf07..f1a919f6 100644 --- a/src/lib/mlxcel-core/src/utils.rs +++ b/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 { + 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. diff --git a/src/models/cohere2.rs b/src/models/cohere2.rs index 9d620ebe..bd4e91aa 100644 --- a/src/models/cohere2.rs +++ b/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 { diff --git a/src/models/exaone4.rs b/src/models/exaone4.rs index b8510409..b692b9a9 100644 --- a/src/models/exaone4.rs +++ b/src/models/exaone4.rs @@ -22,7 +22,7 @@ //! - Output-norm architecture: h = x + norm(attn_out), out = h + norm(mlp_out) use mlxcel_core::layers::{KVCache, RMSNorm, RotatingKVCache, 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; @@ -655,8 +655,10 @@ impl ExaOne4Model { // Create masks (cast to bfloat16 to match model weights dtype for SDPA) let mask_local = if self.sliding_window.is_some() { let max_cache = self.sliding_window.map(|w| w as i32).unwrap_or(i32::MAX); - let effective_offset = local_offset.min((max_cache - seq_len).max(0)); - let mask = create_causal_mask_with_window(seq_len, effective_offset, Some(max_cache)); + // Full-width windowed mask for a fresh single-pass prefill that + // exceeds the window (RotatingKVCache keeps all prefill keys), + // clamped mask otherwise. See issue #408. + let mask = create_sliding_window_prefill_mask(seq_len, local_offset, max_cache); Some(mlxcel_core::astype(&mask, mlxcel_core::dtype::BFLOAT16)) } else { None diff --git a/src/models/exaone_moe.rs b/src/models/exaone_moe.rs index ca632a0b..81a2156d 100644 --- a/src/models/exaone_moe.rs +++ b/src/models/exaone_moe.rs @@ -684,13 +684,12 @@ impl ExaoneMoeModel { .find(|c| matches!(c, AnyKVCache::Rotating(_))) .map(|c| c.offset()) .unwrap_or(0); - // Clamp offset so mask shape matches RotatingKVCache output + // Full-width windowed mask for a fresh single-pass prefill that + // exceeds the window (RotatingKVCache keeps all prefill keys), + // clamped mask otherwise. See issue #408. let max_cache = self.window_size as i32; - let effective_swa_offset = swa_offset.min((max_cache - seq_len).max(0)); - let swa_mask = Some(utils::create_causal_mask_with_window( - seq_len, - effective_swa_offset, - Some(max_cache), + let swa_mask = Some(utils::create_sliding_window_prefill_mask( + seq_len, swa_offset, max_cache, )); (global_mask, swa_mask) diff --git a/src/models/gemma3n.rs b/src/models/gemma3n.rs index 37360444..a911b135 100644 --- a/src/models/gemma3n.rs +++ b/src/models/gemma3n.rs @@ -30,7 +30,7 @@ use crate::models::gemma3n_helpers::{ }; use mlxcel_core::generate::LanguageModel; use mlxcel_core::layers::{KVCache, Linear, RMSNorm, 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; @@ -514,16 +514,21 @@ impl Gemma3nAttention { self.window_size, ) } else { - // Single token or explicit mask path. If the caller supplied a - // windowed mask and this layer is a sliding-window layer, the - // mask is shaped `(q_len, window_size)` post-earlier cap. Plain - // `KVCache` returns full-length K/V, so we must slice K/V to - // match the capped mask, mirroring `causal_attention` internals. + // Single token or explicit mask path. A sliding-window layer's + // mask is either the clamped `(q_len, window_size)` mask or the + // full `(q_len, 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(&keys); 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(&values); - let start = k_len - self.window_size; + let start = k_len - mask_klen; ( Some(mlxcel_core::slice( &keys, @@ -1200,10 +1205,14 @@ impl Gemma3nLanguageModel { None }; let sliding_mask = if l > 1 { - 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 every + // (dense `KVCache`) key. See issue #408. + Some(create_sliding_window_prefill_mask( l, sliding_offset, - Some(self.config.sliding_window as i32), + self.config.sliding_window as i32, )) } else { None @@ -1373,10 +1382,14 @@ impl Gemma3nLanguageModel { None }; let sliding_mask = if l > 1 { - 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 every + // (dense `KVCache`) key. See issue #408. + Some(create_sliding_window_prefill_mask( l, sliding_offset, - Some(self.config.sliding_window as i32), + self.config.sliding_window as i32, )) } else { None diff --git a/src/models/gpt_oss.rs b/src/models/gpt_oss.rs index 12aeeac7..20290916 100644 --- a/src/models/gpt_oss.rs +++ b/src/models/gpt_oss.rs @@ -26,7 +26,7 @@ //! Reference: mlx-lm gpt_oss.py use mlxcel_core::layers::{KVCache, RMSNorm, RotatingKVCache, 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; @@ -908,9 +908,11 @@ impl GptOssModel { let sliding_offset = caches[swa_idx].as_interface().offset(); let max_cache = self.sliding_window as i32; - let effective_offset = sliding_offset.min((max_cache - seq_len).max(0)); + // Full-width windowed mask for a fresh single-pass prefill that + // exceeds the window (RotatingKVCache keeps all prefill keys), + // clamped mask otherwise. See issue #408. let sliding_mask = - create_causal_mask_with_window(seq_len, effective_offset, Some(max_cache)); + create_sliding_window_prefill_mask(seq_len, sliding_offset, max_cache); for (i, layer) in self.layers.iter().enumerate() { let mask = if self.layer_types[i] == "full_attention" { @@ -1240,9 +1242,10 @@ impl GptOssStageModel { let sliding_mask = if seq_len > 1 { self.first_layer_offset(caches, "sliding_attention") .map(|offset| { + // See issue #408: full-width windowed mask for a fresh + // >window prefill, clamped otherwise. let max_cache = self.sliding_window as i32; - let effective_offset = offset.min((max_cache - seq_len).max(0)); - create_causal_mask_with_window(seq_len, effective_offset, Some(max_cache)) + create_sliding_window_prefill_mask(seq_len, offset, max_cache) }) } else { None diff --git a/src/models/mellum.rs b/src/models/mellum.rs index 0ffaf605..7e7fd2a3 100644 --- a/src/models/mellum.rs +++ b/src/models/mellum.rs @@ -35,7 +35,7 @@ use crate::models::switch_layers::{SwitchGLU, fused_moe_enabled, moe_weighted_sum}; use mlxcel_core::layers::{KVCache, RMSNorm, RotatingKVCache, 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; @@ -753,10 +753,12 @@ impl MellumModel { let mask_full = Some(create_causal_mask(seq_len, full_offset)); let mask_sliding = sliding_idx.map(|idx| { + // Full-width windowed mask for a fresh single-pass prefill that + // exceeds the window (the RotatingKVCache keeps all prefill keys), + // clamped mask otherwise. See issue #408. let sliding_offset = caches[idx].as_interface().offset(); let max_cache = self.sliding_window as i32; - let effective_offset = sliding_offset.min((max_cache - seq_len).max(0)); - create_causal_mask_with_window(seq_len, effective_offset, Some(max_cache)) + create_sliding_window_prefill_mask(seq_len, sliding_offset, max_cache) }); self.forward( diff --git a/src/models/ministral3.rs b/src/models/ministral3.rs index c900e5cd..0d106b26 100644 --- a/src/models/ministral3.rs +++ b/src/models/ministral3.rs @@ -20,7 +20,7 @@ //! - RotatingKVCache for sliding_attention layers use mlxcel_core::layers::{KVCache, RMSNorm, RotatingKVCache, 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; @@ -462,10 +462,11 @@ impl Ministral3Model { // Create masks let fa_mask = Some(create_causal_mask(seq_len, fa_offset)); let swa_mask = swa_idx.map(|_| { - // Clamp offset so mask shape matches RotatingKVCache output + // Full-width windowed mask for a fresh single-pass prefill that + // exceeds the window (RotatingKVCache keeps all prefill keys), + // clamped mask otherwise. See issue #408. let max_cache = self.sliding_window.map(|w| w as i32).unwrap_or(i32::MAX); - let effective_offset = swa_offset.min((max_cache - seq_len).max(0)); - create_causal_mask_with_window(seq_len, effective_offset, Some(max_cache)) + create_sliding_window_prefill_mask(seq_len, swa_offset, max_cache) }); // Compute Llama 4 attention scale using h.shape[1] (after embedding) diff --git a/src/models/olmo3.rs b/src/models/olmo3.rs index 24be1fe6..1ab8372d 100644 --- a/src/models/olmo3.rs +++ b/src/models/olmo3.rs @@ -24,7 +24,7 @@ use mlxcel_core::generate::LanguageModel; use mlxcel_core::layers::{KVCache, RMSNorm, 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; @@ -205,15 +205,21 @@ impl OLMo3Attention { // Scaled dot-product attention let attn_out = if l > 1 && mask.is_some() { - // Prefill with mask. If this is a sliding layer the windowed - // mask is capped to `(l, window_size)` post-earlier. Plain - // `KVCache` returns full-length K/V, so slice K/V to the last - // `window_size` slots to match the mask. + // 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, @@ -458,10 +464,14 @@ impl OLMo3Model { 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 every + // (dense `KVCache`) key. 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 { diff --git a/src/models/step3p5.rs b/src/models/step3p5.rs index 2dc3e603..95f4ac98 100644 --- a/src/models/step3p5.rs +++ b/src/models/step3p5.rs @@ -25,7 +25,7 @@ use mlxcel_core::generate::LanguageModel; use mlxcel_core::layers::{KVCache, RMSNorm, RotatingKVCache, 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; @@ -964,8 +964,11 @@ impl Step3p5Model { let swa_mask = if seq_len > 1 { self.swa_idx.map(|idx| { + // Full-width windowed mask for a fresh single-pass prefill that + // exceeds the window (RotatingKVCache keeps all prefill keys), + // clamped mask otherwise. See issue #408. let offset = caches[idx].as_interface().offset(); - create_causal_mask_with_window(seq_len, offset, Some(self.config.sliding_window)) + create_sliding_window_prefill_mask(seq_len, offset, self.config.sliding_window) }) } else { None