From fc20c3109492b18e3d15c77c921f205f8fcf2f8a Mon Sep 17 00:00:00 2001
From: Jeongkyu Shin <inureyes@gmail.com>
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] // </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"
+        );
+    }
+}
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<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 {
@@ -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<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"
+        );
+    }
 }
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<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: `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 <inureyes@gmail.com>
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 <inureyes@gmail.com>
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
 
