From fc20c3109492b18e3d15c77c921f205f8fcf2f8a Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Thu, 25 Jun 2026 04:45:02 +0900 Subject: [PATCH 1/3] fix(models): size gemma3/gemma4/exaone_moe prefill mask from live_len gemma3, gemma4, and exaone_moe built their multi-token (l > 1) prefill attention masks from the cache's monotonic offset instead of the live window, for both the global (full-attention) causal mask and the sliding-window prefill mask. That is correct only while no trim has happened. Under the server --max-kv-size path, enforce_max_kv_size_for calls KVCache::trim_front, which advances live_start while offset keeps growing for RoPE; update_and_fetch then returns only live_len = offset - live_start keys, so a mask sized from offset is wider than the live K/V, MLX broadcast_shapes throws, and the non-Result cxx boundary turns the throw into std::terminate and aborts the whole server. This completes the offset->live_len prefill-mask migration started in #419/#420 and #421/#422 for the three Gemma-family / ExaOne dense-cache models they missed. Fix is mask-only: size the prefill masks from the cache's live window (live_len() for a Standard KVCache, seq_len() for a Rotating cache, both equal to the keys update_and_fetch returns) instead of offset. RoPE and position bookkeeping keep reading the monotonic offset. Byte-identical when untrimmed (live_start == 0, so live_len == offset). Models changed: - gemma3: the seq_len > 1 prefill branch sizes the global mask from caches[global_idx].live_len() and the sliding mask from caches[0].live_len(); a live_len() accessor was added to the model's CacheInterface trait (KVCache -> live_len(), RotatingKVCache -> seq_len()). - gemma4: a first_cache_live_len helper (free function plus method form) mirrors first_cache_offset and feeds the two genuine prefill-mask sites (the non-ragged l > 1 branch and the stage-forward execute_hidden path). The batched-MTP sites that use the offset as a per_row_valid_end column coordinate (divergent verify, has_padding, mask_stale_key_gap) intentionally keep the monotonic first_cache_offset; that regime enforces live_len == offset (slot_base == 0) anyway. - exaone_moe: the global and sliding prefill masks size from a new AnyKVCache::live_len() (Standard -> live_len(), Rotating -> seq_len()) instead of offset(). gemma4_mtp_target was verified and needs no change: it is not a registered LanguageModel, drives forward through caller-owned caches that never enter the scheduler pool, and its only mask builder sizes from a fixed offset 0 prefill. Adds per-model unit tests (gemma3_mask_tests, gemma4_unified_mask_tests, exaone_moe_mask_tests) that trim a Standard cache so live_start > 0, then assert the live_len-sized prefill mask key axis matches the continuation chunk's returned K/V (live_len + m) while the offset-sized mask is strictly wider, plus a Rotating-cache case where live_len equals the returned key axis. The shared cache invariant is already pinned by kv_cache_continuation_mask_sizes_from_live_len_not_offset_after_trim in cache.rs (#419). --- src/models/exaone_moe.rs | 128 ++++++++++++++++++++++++++++-- src/models/gemma3.rs | 127 ++++++++++++++++++++++++++++-- src/models/gemma4.rs | 166 ++++++++++++++++++++++++++++++++++++--- 3 files changed, 397 insertions(+), 24 deletions(-) diff --git a/src/models/exaone_moe.rs b/src/models/exaone_moe.rs index 81a2156d..bd81da2c 100644 --- a/src/models/exaone_moe.rs +++ b/src/models/exaone_moe.rs @@ -627,6 +627,22 @@ pub enum AnyKVCache { Rotating(RotatingKVCache), } +impl AnyKVCache { + /// Number of live keys the cache will actually return on the next + /// `update_and_fetch`. For a `Standard` (`KVCache`) this is + /// `offset - live_start` (`live_len()`), which shrinks below the + /// monotonic `offset` after a `--max-kv-size` `trim_front`; for a + /// `Rotating` cache `seq_len()` already reports the live window. Prefill + /// masks size from this so the mask key axis matches the returned K/V, + /// while `offset()` stays for RoPE. See issue #430. + fn live_len(&self) -> i32 { + match self { + Self::Standard(c) => c.live_len(), + Self::Rotating(c) => c.seq_len(), + } + } +} + impl KVCacheTrait for AnyKVCache { fn update_and_fetch( &mut self, @@ -668,28 +684,42 @@ impl ExaoneMoeModel { let seq_len = mlxcel_core::array_shape(&h)[1]; - // Create masks if needed + // Create masks if needed. + // + // Size both prefill masks from the cache's live window + // (`live_len()`), not the monotonic `offset`. Under `--max-kv-size`, + // the scheduler trims the global (full-attention) `KVCache` via + // `trim_front`, which advances `live_start` while `offset` keeps + // growing to preserve the RoPE relative positions, so + // `update_and_fetch` returns only `live_len = offset - live_start` + // keys. A mask sized from `offset` would be wider than the returned + // K/V and trip `broadcast_shapes`. With no trim (`live_start == 0`), + // `live_len == offset`, so this is byte-identical to the untrimmed + // path; RoPE keeps reading the monotonic `offset` inside the + // attention layer. See issue #430 (mirrors #419/#420, #421/#422). let (global_mask, swa_mask) = if seq_len > 1 && mask.is_none() { // Create global causal mask - let ga_offset = caches + let ga_live_len = caches .iter() .find(|c| matches!(c, AnyKVCache::Standard(_))) - .map(|c| c.offset()) + .map(|c| c.live_len()) .unwrap_or(0); - let global_mask = Some(utils::create_causal_mask(seq_len, ga_offset)); + let global_mask = Some(utils::create_causal_mask(seq_len, ga_live_len)); // Create sliding window mask if applicable - let swa_offset = caches + let swa_live_len = caches .iter() .find(|c| matches!(c, AnyKVCache::Rotating(_))) - .map(|c| c.offset()) + .map(|c| c.live_len()) .unwrap_or(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 max_cache = self.window_size as i32; let swa_mask = Some(utils::create_sliding_window_prefill_mask( - seq_len, swa_offset, max_cache, + seq_len, + swa_live_len, + max_cache, )); (global_mask, swa_mask) @@ -1057,3 +1087,87 @@ impl LanguageModel for ExaoneMoeModel { vec![2] // } } + +#[cfg(test)] +mod exaone_moe_mask_tests { + // `KVCacheTrait` (defined in this module) brings `update_and_fetch` / + // `offset` into scope for `AnyKVCache`; `live_len` is an inherent method. + use super::*; + + // [1, H, n, D] K/V tensor; only the key axis (n) matters for the shape + // invariant exercised here. + fn make_kv(h: i32, n: i32, d: i32, base: f32) -> UniquePtr { + let count = (h * n * d) as usize; + let data: Vec = (0..count).map(|i| base + i as f32).collect(); + mlxcel_core::from_slice_f32(&data, &[1, h, n, d]) + } + + /// Regression for #430: the global prefill mask must size from + /// `AnyKVCache::live_len()` (the `Standard`/`KVCache` live window), not the + /// monotonic `offset()`. After a `--max-kv-size` `trim_front` + /// (`live_start > 0`), `live_len < offset`; a continuation chunk's + /// `update_and_fetch` returns only `live_len + m` keys, so a `live_len`- + /// sized causal mask matches the returned K/V while an `offset`-sized mask + /// is strictly wider (the crash this fix prevents). + #[test] + fn global_mask_sizes_from_live_len_not_offset_after_trim() { + const H: i32 = 2; + const D: i32 = 4; + let (n1, trimmed, m) = (8, 3, 4); + + let mut cache = AnyKVCache::Standard(KVCache::new()); + let _ = cache.update_and_fetch(make_kv(H, n1, D, 0.0), make_kv(H, n1, D, 100.0)); + let AnyKVCache::Standard(ref mut std_cache) = cache else { + unreachable!("constructed as Standard") + }; + assert_eq!(std_cache.trim_front(trimmed), trimmed); + + let offset = cache.offset(); + let live_len = cache.live_len(); + assert_eq!(offset, n1, "offset stays monotonic after trim"); + assert_eq!(live_len, n1 - trimmed, "live_len == offset - live_start"); + assert!(live_len < offset, "trim must make live_len < offset"); + + let (cont_k, _) = cache.update_and_fetch(make_kv(H, m, D, 200.0), make_kv(H, m, D, 300.0)); + let returned_klen = mlxcel_core::array_shape(&cont_k)[2]; + assert_eq!( + returned_klen, + live_len + m, + "cache returns live_len + m keys" + ); + + let live_mask = utils::create_causal_mask(m, live_len); + mlxcel_core::eval(&live_mask); + assert_eq!( + *mlxcel_core::array_shape(&live_mask).last().unwrap(), + returned_klen, + "live_len-sized mask key axis must match the returned K/V" + ); + let offset_mask = utils::create_causal_mask(m, offset); + mlxcel_core::eval(&offset_mask); + assert!( + *mlxcel_core::array_shape(&offset_mask).last().unwrap() > returned_klen, + "offset-sized mask must be wider than the returned K/V (the bug)" + ); + } + + /// The sliding lookup uses `RotatingKVCache::seq_len()` via + /// `AnyKVCache::live_len()`, which equals the keys `update_and_fetch` + /// returns. + #[test] + fn sliding_cache_live_len_matches_returned_keys() { + const H: i32 = 2; + const D: i32 = 4; + let window = 6; + let m = 4; + + let mut cache = AnyKVCache::Rotating(RotatingKVCache::new(window)); + let (k, _) = cache.update_and_fetch(make_kv(H, m, D, 0.0), make_kv(H, m, D, 100.0)); + let returned_klen = mlxcel_core::array_shape(&k)[2]; + assert_eq!( + cache.live_len(), + returned_klen, + "AnyKVCache::live_len() for Rotating (== seq_len) must equal the returned key axis" + ); + } +} diff --git a/src/models/gemma3.rs b/src/models/gemma3.rs index 05c4d352..41ffe230 100644 --- a/src/models/gemma3.rs +++ b/src/models/gemma3.rs @@ -589,6 +589,12 @@ impl TransformerBlock { // Cache Interface. pub(crate) trait CacheInterface { fn offset(&self) -> i32; + /// Number of live keys the cache will actually return on the next + /// `update_and_fetch` (`offset - live_start` for a trimmed `KVCache`, + /// the live window for a `RotatingKVCache`). Used to size prefill masks + /// so the mask key axis matches the returned K/V; `offset()` stays for + /// RoPE. See issue #430 (mirrors #419/#420, #421/#422). + fn live_len(&self) -> i32; fn update_and_fetch( &mut self, k: UniquePtr, @@ -601,6 +607,10 @@ impl CacheInterface for KVCache { self.offset } + fn live_len(&self) -> i32 { + KVCache::live_len(self) + } + fn update_and_fetch( &mut self, k: UniquePtr, @@ -615,6 +625,12 @@ impl CacheInterface for RotatingKVCache { self.offset } + fn live_len(&self) -> i32 { + // `RotatingKVCache` has no `live_start`; `seq_len()` already reports + // the live window the cache returns, so it is the right mask width. + self.seq_len() + } + fn update_and_fetch( &mut self, k: UniquePtr, @@ -761,16 +777,29 @@ impl Gemma3Model { pipeline_hint(&h, i, n); } } else { - // Prefill path (seq_len > 1): create causal masks + // Prefill path (seq_len > 1): create causal masks. + // + // Size both prefill masks from the cache's live window + // (`live_len() = offset - live_start`), not the monotonic + // `offset`. Under `--max-kv-size`, `trim_front` slices the global + // (full-attention) `KVCache` to its live window and advances + // `live_start` while `offset` keeps growing to preserve the RoPE + // relative positions, so `update_and_fetch` returns only + // `live_len` keys. A mask sized from `offset` would be wider than + // the returned K/V and trip `broadcast_shapes`. With no trim + // (`live_start == 0`), `live_len == offset`, so this is + // byte-identical to the untrimmed path. RoPE keeps reading each + // layer's own monotonic `offset` inside the attention layer. See + // issue #430 (mirrors #419/#420, #421/#422). let global_idx = self.sliding_window_pattern - 1; - let global_offset = caches[global_idx].as_interface().offset(); - let global_mask = Some(create_causal_mask(seq_len, global_offset)); + let global_live_len = caches[global_idx].as_interface().live_len(); + let global_mask = Some(create_causal_mask(seq_len, global_live_len)); let sliding_mask = if self.sliding_window_pattern > 1 { - let sliding_offset = caches[0].as_interface().offset(); + let sliding_live_len = caches[0].as_interface().live_len(); Some(create_sliding_window_prefill_mask( seq_len, - sliding_offset, + sliding_live_len, self.sliding_window as i32, )) } else { @@ -1364,4 +1393,92 @@ mod gemma3_mask_tests { "non-fresh over-window prefill mask must match the window-trimmed key axis" ); } + + // [H, n, D] K/V tensor for an `[1, H, n, D]` cache entry; only the key + // axis (n) matters for these shape invariants. + fn make_kv(h: i32, n: i32, d: i32, base: f32) -> UniquePtr { + let count = (h * n * d) as usize; + let data: Vec = (0..count).map(|i| base + i as f32).collect(); + mlxcel_core::from_slice_f32(&data, &[1, h, n, d]) + } + + /// Regression for #430: the global (full-attention) prefill mask must be + /// sized from the `Cache::Standard` live window via `CacheInterface:: + /// live_len()`, not the monotonic `offset()`. After a `--max-kv-size` + /// `trim_front` (`live_start > 0`), `live_len < offset`, and a continuation + /// chunk's `update_and_fetch` returns only `live_len + m` keys; a mask + /// sized from `live_len` matches that axis while an `offset`-sized mask is + /// strictly wider (the crash). + #[test] + fn global_prefill_mask_sizes_from_live_len_not_offset_after_trim() { + const H: i32 = 2; + const D: i32 = 4; + let (n1, trimmed, m) = (8, 3, 4); + + let mut cache = Cache::Standard(KVCache::new()); + let _ = cache + .as_interface() + .update_and_fetch(make_kv(H, n1, D, 0.0), make_kv(H, n1, D, 100.0)); + let Cache::Standard(ref mut std_cache) = cache else { + unreachable!("constructed as Standard") + }; + assert_eq!(std_cache.trim_front(trimmed), trimmed); + + // RoPE invariant: `offset()` stays monotonic; the live window shrinks. + let offset = cache.as_interface().offset(); + let live_len = cache.as_interface().live_len(); + assert_eq!(offset, n1, "offset stays monotonic after trim"); + assert_eq!(live_len, n1 - trimmed, "live_len == offset - live_start"); + assert!(live_len < offset, "trim must make live_len < offset"); + + // What the continuation chunk's `update_and_fetch` actually returns. + let (cont_k, _) = cache + .as_interface() + .update_and_fetch(make_kv(H, m, D, 200.0), make_kv(H, m, D, 300.0)); + let returned_klen = mlxcel_core::array_shape(&cont_k)[2]; + assert_eq!( + returned_klen, + live_len + m, + "cache returns live_len + m keys" + ); + + // live_len-sized mask matches the returned K/V; offset-sized is wider. + let live_mask = create_causal_mask(m, live_len); + mlxcel_core::eval(&live_mask); + assert_eq!( + *mlxcel_core::array_shape(&live_mask).last().unwrap(), + returned_klen, + "create_causal_mask(m, live_len) key axis must match the returned K/V" + ); + let offset_mask = create_causal_mask(m, offset); + mlxcel_core::eval(&offset_mask); + assert!( + *mlxcel_core::array_shape(&offset_mask).last().unwrap() > returned_klen, + "offset-sized mask must be wider than the returned K/V (the bug)" + ); + } + + /// The sliding (`Cache::Rotating`) prefill mask must size from the same + /// `CacheInterface::live_len()` accessor, which for a `RotatingKVCache` + /// returns `seq_len()` (its live window). Untrimmed, that equals the keys + /// it returns, so `live_len()` is the correct mask width and never exceeds + /// the K/V axis. + #[test] + fn sliding_cache_live_len_matches_returned_keys() { + const H: i32 = 2; + const D: i32 = 4; + let window = 6; + let m = 4; + + let mut cache = Cache::Rotating(RotatingKVCache::new(window)); + let (k, _) = cache + .as_interface() + .update_and_fetch(make_kv(H, m, D, 0.0), make_kv(H, m, D, 100.0)); + let returned_klen = mlxcel_core::array_shape(&k)[2]; + assert_eq!( + cache.as_interface().live_len(), + returned_klen, + "RotatingKVCache live_len() (== seq_len) must equal the returned key axis" + ); + } } diff --git a/src/models/gemma4.rs b/src/models/gemma4.rs index 2f7a1551..b453d187 100644 --- a/src/models/gemma4.rs +++ b/src/models/gemma4.rs @@ -2747,6 +2747,15 @@ impl Gemma4TextModel { } else if has_padding { // Resident prompt padding present (ragged burst). Build per-row // left-padding-aware masks for ANY query width, including `l == 1`. + // + // These offsets stay monotonic (NOT `live_len`): this is the + // batched-MTP regime, which never front-compacts either cache kind + // (`slot_base == 0`, so `live_len == offset` here), and the same + // offset value is the column-coordinate base that the + // `per_row_valid_end` stale-gap masking below compares against + // (`mask_stale_key_gap`). The `--max-kv-size` trim cannot reach + // these model-owned caches, so there is no live/monotonic split to + // reconcile on this path. See issue #430. let global_offset = first_cache_offset(caches, "full_attention"); let sliding_offset = first_cache_offset(caches, "sliding_attention"); let window = self.config.sliding_window as i32; @@ -2786,9 +2795,16 @@ impl Gemma4TextModel { (Some(global_mask), Some(sliding_mask)) } else if l > 1 { // Non-ragged prefill / multi-token verify: derive both masks from - // the cache offsets (byte-identical to the pre-ragged behaviour). - let global_offset = first_cache_offset(caches, "full_attention"); - let sliding_offset = first_cache_offset(caches, "sliding_attention"); + // the cache's live window (`live_len`), not the monotonic + // `offset`. Under `--max-kv-size`, `trim_front` advances the + // full-attention `KVCache`'s `live_start` while `offset` keeps + // growing for RoPE, so `update_and_fetch` returns only `live_len` + // keys; an `offset`-sized mask would be wider than the returned + // K/V and trip `broadcast_shapes`. With no trim, `live_len == + // offset`, so this stays byte-identical to the pre-ragged + // behaviour. See issue #430 (mirrors #419/#420, #421/#422). + let global_live_len = first_cache_live_len(caches, "full_attention"); + let sliding_live_len = first_cache_live_len(caches, "sliding_attention"); let window = self.config.sliding_window as i32; // Shared helper (hoisted in #410) gates the full mask on // `l > window && sliding_offset == 0`; the old local copy gated on @@ -2800,10 +2816,10 @@ impl Gemma4TextModel { // independent of offset). Old-trimmed equals new for every input, // so the migration is behaviour-preserving. ( - Some(create_causal_mask(l, global_offset)), + Some(create_causal_mask(l, global_live_len)), Some(create_sliding_window_prefill_mask( l, - sliding_offset, + sliding_live_len, window, )), ) @@ -3287,6 +3303,31 @@ pub(crate) fn first_cache_offset(caches: &mut [Cache], layer_type: &str) -> i32 0 } +/// Live-window length of the first cache of `layer_type`, the count of keys +/// `update_and_fetch` will actually return. +/// +/// This is the mask-sizing companion to [`first_cache_offset`]: prefill masks +/// must be sized from the live window, not the monotonic `offset`. Under +/// `--max-kv-size`, `trim_front` advances a full-attention `KVCache`'s +/// `live_start` while `offset` keeps growing for RoPE, so the cache returns +/// only `live_len = offset - live_start` keys; a mask sized from `offset` +/// would be wider than the returned K/V and trip `broadcast_shapes`. For a +/// `Rotating` (sliding) cache `seq_len()` already reports the live window. +/// With no trim (`live_start == 0`), this equals [`first_cache_offset`], so +/// the prefill path stays byte-identical. RoPE/position bookkeeping and the +/// batched-MTP `per_row_valid_end` coordinate math keep using +/// [`first_cache_offset`] (the monotonic value). See issue #430. +pub(crate) fn first_cache_live_len(caches: &mut [Cache], layer_type: &str) -> i32 { + for cache in caches.iter_mut() { + match (layer_type, cache) { + ("full_attention", Cache::Standard(c)) => return c.live_len(), + ("sliding_attention", Cache::Rotating(c)) => return c.seq_len(), + _ => {} + } + } + 0 +} + pub(crate) fn slice_layer_input( layer_inputs: &MlxArray, layer_idx: i32, @@ -3783,15 +3824,21 @@ impl Gemma4StageModel { let seq_len = shape[1]; let (global_mask, sliding_mask) = if seq_len > 1 { - let global_offset = self.first_cache_offset(caches, "full_attention"); - let sliding_offset = self.first_cache_offset(caches, "sliding_attention"); + // Size the prefill masks from the cache's live window + // (`live_len`), not the monotonic `offset`, so the mask key axis + // matches the K/V `update_and_fetch` returns after a + // `--max-kv-size` trim (`live_start` advances while `offset` keeps + // growing for RoPE). Byte-identical when untrimmed (`live_len == + // offset`). See issue #430 (mirrors #419/#420, #421/#422). + let global_live_len = self.first_cache_live_len(caches, "full_attention"); + let sliding_live_len = self.first_cache_live_len(caches, "sliding_attention"); ( - Some(create_causal_mask(seq_len, global_offset)), + Some(create_causal_mask(seq_len, global_live_len)), // Shared helper hoisted in #410; behaviour-preserving for Gemma 4 // (see the gate-divergence note in `forward_with_speculative_sinks`). Some(create_sliding_window_prefill_mask( seq_len, - sliding_offset, + sliding_live_len, self.config.sliding_window as i32, )), ) @@ -4014,11 +4061,17 @@ impl Gemma4StageModel { ) } - fn first_cache_offset(&self, caches: &[Cache], layer_type: &str) -> i32 { + /// Live-window length of the first cache of `layer_type` (mask-sizing + /// companion to the free-function [`first_cache_offset`]). See + /// [`first_cache_live_len`] for the full rationale; prefill masks size + /// from this so the mask key axis matches the K/V `update_and_fetch` + /// returns after a `--max-kv-size` trim, while RoPE keeps the monotonic + /// offset. See issue #430. + fn first_cache_live_len(&self, caches: &[Cache], layer_type: &str) -> i32 { for cache in caches { match (layer_type, cache) { - ("full_attention", Cache::Standard(cache)) => return cache.offset, - ("sliding_attention", Cache::Rotating(cache)) => return cache.offset, + ("full_attention", Cache::Standard(cache)) => return cache.live_len(), + ("sliding_attention", Cache::Rotating(cache)) => return cache.seq_len(), _ => {} } } @@ -4893,6 +4946,95 @@ mod gemma4_unified_mask_tests { assert_eq!(mask_at(&out, 3, 1), 0.0); } + // [1, H, n, D] K/V tensor; only the key axis (n) matters here. + fn make_kv(h: i32, n: i32, d: i32, base: f32) -> UniquePtr { + let count = (h * n * d) as usize; + let data: Vec = (0..count).map(|i| base + i as f32).collect(); + mlxcel_core::from_slice_f32(&data, &[1, h, n, d]) + } + + /// Regression for #430: `first_cache_live_len` must return the + /// full-attention `Cache::Standard` LIVE window (`live_len`), not the + /// monotonic `offset`, so the prefill mask matches the K/V a continuation + /// chunk's `update_and_fetch` returns after a `--max-kv-size` trim. + #[test] + fn first_cache_live_len_sizes_global_mask_from_live_window_after_trim() { + const H: i32 = 2; + const D: i32 = 4; + let (n1, trimmed, m) = (8, 3, 4); + let window = 6; + + // Layer 0 sliding (Rotating), layer 1 global (Standard): matches the + // `first_cache_live_len` family lookups. + let mut caches = vec![ + Cache::Rotating(RotatingKVCache::new(window)), + Cache::Standard(KVCache::new()), + ]; + if let Cache::Standard(ref mut c) = caches[1] { + let _ = c.update_and_fetch(make_kv(H, n1, D, 0.0), make_kv(H, n1, D, 100.0)); + assert_eq!(c.trim_front(trimmed), trimmed); + } + + let global_offset = first_cache_offset(&mut caches, "full_attention"); + let global_live_len = first_cache_live_len(&mut caches, "full_attention"); + assert_eq!(global_offset, n1, "first_cache_offset stays monotonic"); + assert_eq!( + global_live_len, + n1 - trimmed, + "first_cache_live_len reports the trimmed live window" + ); + assert!(global_live_len < global_offset); + + let (cont_k, _) = if let Cache::Standard(ref mut c) = caches[1] { + c.update_and_fetch(make_kv(H, m, D, 200.0), make_kv(H, m, D, 300.0)) + } else { + unreachable!() + }; + let returned_klen = mlxcel_core::array_shape(&cont_k)[2]; + assert_eq!(returned_klen, global_live_len + m); + + let live_mask = create_causal_mask(m, global_live_len); + mlxcel_core::eval(&live_mask); + assert_eq!( + *mlxcel_core::array_shape(&live_mask).last().unwrap(), + returned_klen, + "live_len-sized global mask must match the returned K/V" + ); + let offset_mask = create_causal_mask(m, global_offset); + mlxcel_core::eval(&offset_mask); + assert!( + *mlxcel_core::array_shape(&offset_mask).last().unwrap() > returned_klen, + "offset-sized global mask must be wider than the returned K/V (the bug)" + ); + } + + /// The sliding lookup of `first_cache_live_len` returns the + /// `RotatingKVCache` live window (`seq_len`), which equals the keys it + /// returns from `update_and_fetch`. + #[test] + fn first_cache_live_len_sliding_matches_returned_keys() { + const H: i32 = 2; + const D: i32 = 4; + let window = 6; + let m = 4; + + let mut caches = vec![ + Cache::Rotating(RotatingKVCache::new(window)), + Cache::Standard(KVCache::new()), + ]; + let (k, _) = if let Cache::Rotating(ref mut c) = caches[0] { + c.update_and_fetch(make_kv(H, m, D, 0.0), make_kv(H, m, D, 100.0)) + } else { + unreachable!() + }; + let returned_klen = mlxcel_core::array_shape(&k)[2]; + assert_eq!( + first_cache_live_len(&mut caches, "sliding_attention"), + returned_klen, + "sliding first_cache_live_len (== seq_len) must equal the returned key axis" + ); + } + // The Gemma 4 over-window (`sliding_prefill_mask_over_window_*`) and // within-window prefill-mask tests were hoisted to the shared // `mlxcel_core::utils::create_sliding_window_prefill_mask` tests in #410 From 75dac72be7794677395ac02a2c981fa9001510e2 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Thu, 25 Jun 2026 04:58:52 +0900 Subject: [PATCH 2/3] fix(models): size gemma3 stage prefill mask from live_len The Gemma3StageModel::execute_hidden pipeline-stage path still sized both the global causal mask and the sliding-window prefill mask from the cache's monotonic offset, while the gemma4 stage analog (Gemma4StageModel::execute_hidden) and the gemma3 single-process path (Gemma3Model::forward_with_caches) were already switched to live_len in this PR. This completes the offset to live_len prefill-mask migration for gemma3 so both stage paths match. Like the rest of #430, this is byte-identical on the untrimmed path: the stage executor's PointerOwnedCacheStore syncs only offset (never live_start) into the model-owned internal caches, so live_start stays 0 and live_len == offset there today. The change uses the existing as_interface().live_len() accessor (KVCache to live_len, RotatingKVCache to seq_len), and keeps RoPE on the monotonic offset. --- src/models/gemma3.rs | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/src/models/gemma3.rs b/src/models/gemma3.rs index 41ffe230..b993dbdd 100644 --- a/src/models/gemma3.rs +++ b/src/models/gemma3.rs @@ -1071,17 +1071,26 @@ impl Gemma3StageModel { let seq_len = mlxcel_core::array_shape(hidden.as_ref().unwrap())[1]; if seq_len > 1 { - let global_offset = self + // Size both prefill masks from the cache's live window + // (`live_len()`), not the monotonic `offset`, so the mask key axis + // matches the K/V `update_and_fetch` returns after a + // `--max-kv-size` trim advances `live_start` while `offset` keeps + // growing for RoPE. Byte-identical when untrimmed (`live_len == + // offset`); mirrors `Gemma3Model::forward_with_caches` and + // `Gemma4StageModel::execute_hidden`. See issue #430. + let global_live_len = self .first_global_cache_index() - .map(|idx| caches[idx].offset()) + .map(|idx| caches[idx].as_interface().live_len()) .unwrap_or(0); - let global_mask = create_causal_mask(seq_len, global_offset); + let global_mask = create_causal_mask(seq_len, global_live_len); let sliding_mask = if self.first_sliding_cache_index().is_some() { - let sliding_offset = caches[self.first_sliding_cache_index().unwrap()].offset(); + let sliding_live_len = caches[self.first_sliding_cache_index().unwrap()] + .as_interface() + .live_len(); Some(create_sliding_window_prefill_mask( seq_len, - sliding_offset, + sliding_live_len, self.sliding_window as i32, )) } else { From 6751008edf0df8748be44b5fee10ba7bae4d9c83 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Thu, 25 Jun 2026 05:03:06 +0900 Subject: [PATCH 3/3] docs(changelog): record gemma3/gemma4/exaone_moe prefill-mask consistency fix (#430) --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index c726b6d8..4b77395c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). ### Fixed - **Double-transpose crash on mlx-community conv checkpoints (Gemma 4 audio, phi4mm patch-embed, nemotron audio, RT-DETRv2).** Several weight-sanitizer functions transposed conv weights from PyTorch `[out, in, kH, kW]` to MLX channel-last `[out, kH, kW, in]` unconditionally. Pre-converted mlx-community checkpoints already store these weights in channel-last order, so the unconditional transpose double-converted them and produced a corrupted shape. The confirmed crash: loading `mlx-community/gemma-4-e4b-it-qat-4bit` turned the audio subsample conv weight `[128, 3, 3, 1]` into `[128, 3, 1, 3]`, which MLX conv2d rejected because the input C_in=1 did not match the weight C_in=3. All four affected sanitizers now check the tensor shape before transposing: `conv2d_weight_is_channel_last` (already-MLX `[out, kH, kW, in]` skips; PyTorch `[out, in, kH, kW]` transposes) and `conv1d_weight_is_channel_last` (depthwise-only; MLX `[out, kW, 1]` skips; PyTorch `[out, 1, kW]` transposes). Both predicates are idempotent. Resolves #428. +- **Prefill attention mask sized from the monotonic offset instead of the live window (gemma3, gemma4, exaone_moe).** The multi-token prefill causal and sliding-window masks were sized from the cache's monotonic `offset` rather than its live window (`live_len()`/`seq_len()` = `offset - live_start`). After a `--max-kv-size` `trim_front` advances `live_start`, a mask sized from `offset` is wider than the K/V the cache returns and MLX rejects the broadcast. These three models are not reachable through the trimmable scheduler pool today (gemma3/gemma4 use `ModelOwnedSequenceState`; exaone_moe rebuilds offset-0 caches per forward), so this is a defensive consistency fix that hardens the mask-width invariant and matches the offset-to-live_len migration already applied to cohere2/gemma3n/olmo3 (#419/#420) and mistral4/nemotron_nas/qwen-vl (#421/#422). Byte-identical on the untrimmed path; RoPE and position ids keep the monotonic offset. Resolves #430. ## [v0.3.3] - 2026-06-23