fix: size gemma3/gemma4/exaone_moe prefill mask from live_len under --max-kv-size

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
 

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] // </s>
     }
 }
+
+#[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<MlxArray> {
+        let count = (h * n * d) as usize;
+        let data: Vec<f32> = (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"
+        );
+    }
+}

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<MlxArray>,
@@ -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<MlxArray>,
@@ -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<MlxArray>,
@@ -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 {
@@ -1042,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 {
@@ -1364,4 +1402,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<MlxArray> {
+        let count = (h * n * d) as usize;
+        let data: Vec<f32> = (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"
+        );
+    }
 }