fix: reset model-owned VLM fallback state

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

@@ -237,6 +237,20 @@ pub trait LanguageModel {
     /// so that padding positions do not corrupt subsequent decode steps.
     fn trim_internal_caches(&self, _excess: i32) {}
 
+    /// Reset model-owned fallback runtime state before a fresh single-row
+    /// generation starts.
+    ///
+    /// Most models store all request-local state in the `KVCache` slice owned
+    /// by [`CxxGenerator`], so the default is a no-op. Models that keep a
+    /// fallback cache slot on `&self` (for legacy CLI / benchmark paths that
+    /// do not carry a `SequenceId`) override this to clear that slot without
+    /// touching scheduler-owned per-sequence maps.
+    ///
+    /// Used by: `CxxGenerator::{reset_with_model, generate_streaming,
+    /// generate_with_stats, generate_streaming_with_embeddings,
+    /// generate_with_stats_and_embeddings, evaluate_loglikelihoods}`.
+    fn reset_runtime_state(&self) {}
+
     /// Release any model-owned sequence state associated with the provided
     /// external cache slice before the scheduler drops that cache set.
     ///
@@ -782,8 +796,9 @@ impl CxxGenerator {
     /// Reset generator state including model-internal caches.
     ///
     /// Models with internal RefCell caches (sliding window, SSM, hybrid) reset
-    /// their own state inside `make_caches()`. This method ensures both the
-    /// generator's cache vector and the model's internal caches are cleared.
+    /// their own fallback state inside `reset_runtime_state()`. This method
+    /// ensures both the generator's cache vector and the model's internal
+    /// single-row caches are cleared.
     /// The kv_cache_mode is applied to the freshly created caches.
     ///
     /// Honors the Boundary-V policy (issue #478): when `self.kv_cache_mode`
@@ -792,6 +807,7 @@ impl CxxGenerator {
     /// The boundary count is read from `MLXCEL_KV_BOUNDARY_V_LAYERS` so a
     /// runtime-tuned count is honored on every reset.
     pub fn reset_with_model<M: LanguageModel + ?Sized>(&mut self, model: &M) {
+        model.reset_runtime_state();
         self.caches = model.make_caches();
         // Apply the configured KV cache mode (with Boundary-V upgrade) to
         // all freshly created caches.
@@ -854,8 +870,11 @@ impl CxxGenerator {
         sampling: &SamplingConfig,
         mut on_token: F,
     ) -> Vec<i32> {
-        // Reset state
-        self.reset();
+        // Reset generator-owned caches and model-owned fallback caches. The
+        // latter matters for model-owned cache families (Gemma 3/4, Llama 4,
+        // Qwen 3.5 Next) whose legacy single-row path does not carry a
+        // `SequenceId`.
+        self.reset_with_model(model);
 
         // Axis B: merge any generator-cached language-bias map into the
         // sampling config before seeding/penalty evaluation. Empty cached
@@ -1122,8 +1141,9 @@ impl CxxGenerator {
         sampling: &SamplingConfig,
         mut on_token: F,
     ) -> Vec<i32> {
-        // Reset state
-        self.reset();
+        // Reset generator-owned caches and model-owned fallback caches. See
+        // `generate_streaming` for the model-owned cache rationale.
+        self.reset_with_model(model);
 
         // Axis B: inject generator-cached language-bias into the sampling config.
         let sampling_cow = self.compose_sampling(sampling);
@@ -1269,7 +1289,7 @@ impl CxxGenerator {
     ) -> (Vec<i32>, GenerationStats) {
         use std::time::Instant;
 
-        self.reset();
+        self.reset_with_model(model);
 
         // Axis B: inject generator-cached language-bias into the sampling config.
         let sampling_cow = self.compose_sampling(sampling);
@@ -1417,8 +1437,9 @@ impl CxxGenerator {
     ) -> (Vec<i32>, GenerationStats) {
         use std::time::Instant;
 
-        // Reset state
-        self.reset();
+        // Reset generator-owned caches and model-owned fallback caches. See
+        // `generate_streaming` for the model-owned cache rationale.
+        self.reset_with_model(model);
 
         // Axis B: inject generator-cached language-bias into the sampling config.
         let sampling_cow = self.compose_sampling(sampling);
@@ -2142,6 +2163,61 @@ mod tests {
         model.trim_internal_caches(-1);
     }
 
+    struct TrackingResetModel {
+        reset_call_count: std::cell::Cell<usize>,
+    }
+
+    impl TrackingResetModel {
+        fn new() -> Self {
+            Self {
+                reset_call_count: std::cell::Cell::new(0),
+            }
+        }
+    }
+
+    impl LanguageModel for TrackingResetModel {
+        fn forward(
+            &self,
+            input_ids: &MlxArray,
+            _caches: &mut [KVCache],
+            _mask: Option<&MlxArray>,
+        ) -> UniquePtr<MlxArray> {
+            ffi::eval(input_ids);
+            ffi::zeros(&[1, 1, 4], crate::dtype::FLOAT32)
+        }
+
+        fn make_caches(&self) -> Vec<KVCache> {
+            vec![KVCache::new()]
+        }
+
+        fn num_layers(&self) -> usize {
+            1
+        }
+
+        fn eos_token_ids(&self) -> Vec<i32> {
+            vec![99]
+        }
+
+        fn reset_runtime_state(&self) {
+            self.reset_call_count.set(self.reset_call_count.get() + 1);
+        }
+    }
+
+    #[test]
+    fn reset_with_model_invokes_runtime_state_hook() {
+        let model = TrackingResetModel::new();
+        let mut generator = CxxGenerator::new(model.num_layers());
+
+        assert_eq!(model.reset_call_count.get(), 0);
+        generator.reset_with_model(&model);
+
+        assert_eq!(
+            model.reset_call_count.get(),
+            1,
+            "reset_with_model must reset model-owned fallback state"
+        );
+    }
+
     /// A model that overrides trim_internal_caches records each call so we can
     /// verify the generation machinery actually invokes the method.
     struct TrackingTrimModel {

src/loaded_model.rs

@@ -300,6 +300,10 @@ impl LanguageModel for LoadedModel {
         delegate_language_model!(self, after_prefill())
     }
 
+    fn reset_runtime_state(&self) {
+        delegate_language_model!(self, reset_runtime_state())
+    }
+
     fn release_sequence_state(&self, caches: &mut [mlxcel_core::layers::KVCache]) {
         delegate_language_model!(self, release_sequence_state(caches))
     }

src/models/gemma3.rs

@@ -1204,6 +1204,17 @@ impl mlxcel_core::generate::LanguageModel for Gemma3Wrapper {
             .prepare_sequence_state(seq_id, self.model.make_caches());
     }
 
+    fn reset_runtime_state(&self) {
+        // Used by: CxxGenerator single-row generation paths. Gemma 3 owns
+        // its fallback sliding/global KV cache state inside
+        // `ModelOwnedSequenceState`; the caller-provided `KVCache` slice is
+        // intentionally empty/ignored. Reset the fallback slot before each
+        // fresh CLI / benchmark generation so a second VLM prefill does not
+        // reuse offsets from a previous run with a stale 4D attention mask
+        // (issue #731).
+        self.reset_caches();
+    }
+
     fn release_sequence_state_by_id(&self, seq_id: SequenceId) {
         self.sequence_state.release_sequence_state(seq_id)
     }

src/models/gemma4.rs

@@ -3346,6 +3346,14 @@ impl LanguageModel for Gemma4Wrapper {
             .prepare_sequence_state(seq_id, self.model.make_caches());
     }
 
+    fn reset_runtime_state(&self) {
+        // Used by: CxxGenerator single-row generation paths. Gemma 4 owns
+        // its fallback cache slot inside `ModelOwnedSequenceState`; reset it
+        // for fresh CLI / benchmark runs without touching scheduler-owned
+        // per-sequence entries.
+        self.reset_caches();
+    }
+
     fn release_sequence_state_by_id(&self, seq_id: SequenceId) {
         self.sequence_state.release_sequence_state(seq_id);
     }

src/models/llama4.rs

@@ -1728,6 +1728,14 @@ impl LanguageModel for Llama4Wrapper {
             .prepare_sequence_state(seq_id, self.model.make_llama4_caches());
     }
 
+    fn reset_runtime_state(&self) {
+        // Used by: CxxGenerator single-row generation paths. Llama 4 owns
+        // iGQA cache state inside `ModelOwnedSequenceState`; reset only the
+        // legacy fallback slot, leaving scheduler-owned per-sequence entries
+        // untouched.
+        self.reset_caches();
+    }
+
     fn release_sequence_state_by_id(&self, seq_id: SequenceId) {
         self.sequence_state.release_sequence_state(seq_id)
     }

src/models/qwen3_5.rs

@@ -2575,6 +2575,17 @@ impl LanguageModel for Qwen35Model {
             .prepare_sequence_state(seq_id, self.make_internal_caches());
     }
 
+    fn reset_runtime_state(&self) {
+        // Used by: CxxGenerator single-row generation paths. Qwen 3.5 Next
+        // owns mixed attention / GatedDelta cache state in
+        // `ModelOwnedSequenceState`; reset the fallback cache slot for fresh
+        // CLI / benchmark runs without touching scheduler-owned sequence
+        // entries. Do not clear the MRoPE fallback here: VLM callers populate
+        // it during embedding preparation immediately before generation.
+        self.sequence_state
+            .replace_internal(self.make_internal_caches());
+    }
+
     fn release_sequence_state_by_id(&self, seq_id: SequenceId) {
         self.sequence_state.release_sequence_state(seq_id);
         // Issue #540: drop the per-sequence MRoPE entry alongside the

src/vision/gemma4_vl.rs

@@ -551,6 +551,14 @@ impl LanguageModel for Gemma4VLModel {
         self.text_model.prepare_sequence_state(seq_id);
     }
 
+    fn reset_runtime_state(&self) {
+        // Used by: CxxGenerator single-row generation paths. Reset only the
+        // text backbone's fallback cache slot; the VLM per-layer-inputs
+        // fallback is populated by `get_input_embeddings*` immediately before
+        // `forward_with_embeddings*` consumes it.
+        self.text_model.reset_runtime_state();
+    }
+
     fn release_sequence_state_by_id(&self, seq_id: SequenceId) {
         // Issue #543: drop the per-sequence `per_layer_inputs`
         // alongside the text model's per-sequence cache release so

src/vision/mod.rs

@@ -262,6 +262,10 @@ impl LanguageModel for VisionLanguageModel {
         self.text_model.after_prefill()
     }
 
+    fn reset_runtime_state(&self) {
+        self.text_model.reset_runtime_state()
+    }
+
     fn trim_internal_caches(&self, excess: i32) {
         self.text_model.trim_internal_caches(excess)
     }

src/vision/qwen3_5_vl.rs

@@ -377,6 +377,14 @@ impl LanguageModel for Qwen35VLModel {
         );
     }
 
+    fn reset_runtime_state(&self) {
+        // Used by: CxxGenerator single-row generation paths. Reset the
+        // hybrid text backbone's fallback mixed-cache slot while preserving
+        // the MRoPE fallback written by VLM embedding preparation just before
+        // generation starts.
+        mlxcel_core::generate::LanguageModel::reset_runtime_state(&self.text_model);
+    }
+
     /// Issue #540: forward `seq_ids` to the underlying Qwen3.5 model so
     /// each row's MRoPE state resolves correctly in mixed VL+text
     /// batches. `Qwen35Model::forward_batched_with_context_and_ids`