perf(cache): make the dense KV tail trim a logical rewind instead of a buffer rewrite

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

@@ -925,7 +925,13 @@ impl KVCache {
             let new_v = ffi::zeros(&[b, n_kv_heads, buf_size, v_head_dim], v_dtype);
 
             if self.keys.is_some() {
-                if prev % self.step != 0 {
+                // Normalize the buffer to the live prefix before extending.
+                // The buffer can be longer than `prev`: a logical `trim`
+                // keeps the old capacity (issue #1209), and a partial growth
+                // chunk leaves a non-step-aligned length. Either way the
+                // slots past `prev` are dead and must not survive the
+                // concatenate.
+                if prev != self.buffer_seq_len() {
                     self.keys = Some(ffi::slice(
                         self.keys.as_ref().unwrap(),
                         &[0, 0, 0, 0],
@@ -1011,7 +1017,10 @@ impl KVCache {
             let new_vs_buf = ffi::zeros(&[b, n_kv_heads, buf_size, 1], dtype::FLOAT16);
 
             if self.keys.is_some() {
-                if prev % self.step != 0 {
+                // Normalize to the live prefix before extending — the buffer
+                // can be longer than `prev` after a logical `trim`
+                // (issue #1209) or a partial growth chunk; see `update_fp16`.
+                if prev != self.buffer_seq_len() {
                     self.keys = Some(ffi::slice(
                         self.keys.as_ref().unwrap(),
                         &[0, 0, 0, 0],
@@ -2165,7 +2174,18 @@ impl KVCache {
     /// Trim the last `n` entries from the cache.
     ///
     /// Returns the number of entries actually trimmed.
-    /// In INT8 mode the corresponding scale buffers are also trimmed.
+    ///
+    /// In `Fp16` and `Int8` modes this is a **logical rewind**: only `offset`
+    /// moves, the step-aligned buffers (scale sidecars included) keep their
+    /// allocation, and the next `update_*` overwrites the abandoned slots
+    /// through the same `slice_update` it always uses. Every reader slices to
+    /// `buffer_idx()` rather than trusting the buffer's physical length, which
+    /// is already the normal state between growths. This is the contract
+    /// [`RotatingKVCache::trim`] has always had, and it makes a speculative
+    /// rollback cost O(1) instead of an O(context) buffer rewrite per round
+    /// (issue #1209).
+    ///
+    /// The quantized-sidecar modes keep the physical trim:
     /// In Turbo4Asym mode the V-packed and V-norms buffers are also trimmed.
     /// In Turbo4 (symmetric) mode both the K- and V-side packed sidecars are
     /// trimmed.
@@ -2297,13 +2317,30 @@ impl KVCache {
             }
             // Suppress unused-variable warnings in non-delegated branches.
             let _ = hot_trim;
+        } else if matches!(self.mode, KVCacheMode::Fp16 | KVCacheMode::Int8) {
+            // Dense modes: the `self.offset -= n` above IS the trim. The
+            // buffers (scale sidecars included) keep their step-aligned
+            // capacity; every reader slices to `buffer_idx()`
+            // (`update_and_fetch`, `keys_view`, `trim_front_keep_sink`, the
+            // detached-handle `trim_to`), and the next `slice_update`
+            // overwrites the abandoned slots. Re-slicing here cost an
+            // O(context) copy per speculative round, plus a second one in the
+            // next `update_*` regrow because the capacity was gone
+            // (issue #1209).
         } else {
-            // Non-delegated modes: simple buffer-prefix trim.
-            // The slice upper bound is the post-trim *live* window length
-            // (`live_len_after`), not the monotonic `self.offset`. For
-            // Turbo modes `live_start == 0` so `live_len_after == self.offset`
-            // and behaviour is bit-identical to the earlier implementation.
-            // Trim keys (always present except in Turbo4Asym pre-init).
+            // Turbo modes keep the physical buffer-prefix trim. Their packed
+            // sidecars feed fused kernels and fold/dequant helpers that have
+            // not been audited for a stale tail beyond the live window, and
+            // the speculative rollbacks this rewind serves run on Fp16/Int8
+            // caches; an O(context) slice on an explicit Turbo trim is the
+            // status quo, kept until those consumers are audited
+            // (issue #1209).
+            //
+            // The slice upper bound for K/V is the post-trim *live* window
+            // length (`live_len_after`); `live_start == 0` for Turbo modes so
+            // it equals `self.offset` and the sidecar slices below match.
+            // Trim keys (fp16 K in Turbo4Asym/Turbo3Asym; None in symmetric
+            // Turbo4 pre-init, whose K lives in `k_packed`).
             if let Some(ref k) = self.keys {
                 let k_shape = ffi::array_shape(k);
                 self.keys = Some(ffi::slice(
@@ -2312,7 +2349,6 @@ impl KVCache {
                     &[k_shape[0], k_shape[1], live_len_after, k_shape[3]],
                 ));
             }
-            // Trim values for FP16/INT8 modes (Turbo4Asym keeps values None).
             if let Some(ref v) = self.values {
                 let v_shape = ffi::array_shape(v);
                 self.values = Some(ffi::slice(
@@ -2321,25 +2357,6 @@ impl KVCache {
                     &[v_shape[0], v_shape[1], live_len_after, v_shape[3]],
                 ));
             }
-            // Trim INT8 scale sidecars.
-            if self.mode == KVCacheMode::Int8 {
-                if let Some(ref ks) = self.key_scales {
-                    let ks_shape = ffi::array_shape(ks);
-                    self.key_scales = Some(ffi::slice(
-                        ks,
-                        &[0, 0, 0, 0],
-                        &[ks_shape[0], ks_shape[1], live_len_after, 1],
-                    ));
-                }
-                if let Some(ref vs) = self.val_scales {
-                    let vs_shape = ffi::array_shape(vs);
-                    self.val_scales = Some(ffi::slice(
-                        vs,
-                        &[0, 0, 0, 0],
-                        &[vs_shape[0], vs_shape[1], live_len_after, 1],
-                    ));
-                }
-            }
             // Trim Turbo4* V sidecars (per-token; speculative-decode rewinds
             // <1 block at a time so we never need to re-quantize a partial
             // block — the trimmed tail is already block-aligned in the buffer).

src/lib/mlxcel-core/src/cache/detach.rs

@@ -443,9 +443,9 @@ impl KVCache {
     /// Semantics:
     /// * `new_len == 0` fully rewinds the cache (equivalent to
     ///   `trim(self.offset)`) and drops all backing buffers.
-    /// * `0 < new_len < self.offset` keeps the pre-allocated buffer but
-    ///   re-slices its visible region to `new_len`. This mirrors
-    ///   [`KVCache::trim`] but takes an absolute target instead of a delta.
+    /// * `0 < new_len < self.offset` delegates to [`KVCache::trim`] with the
+    ///   delta: a logical rewind for `Fp16`/`Int8` (the buffer keeps its
+    ///   capacity, issue #1209), a physical re-slice for the Turbo modes.
     /// * `new_len == self.offset` is a no-op.
     /// * `new_len < 0` or `new_len > self.offset` returns `Err`.
     ///

src/lib/mlxcel-core/src/speculative/mod.rs

@@ -950,6 +950,7 @@ mod distribution_tests;
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::cache::KVCacheMode;
     use crate::dtype;
 
     #[test]
@@ -961,17 +962,63 @@ mod tests {
         let values = ffi::ones(&[1, 2, 5, 4], dtype::FLOAT32);
         cache.update(keys, values);
         assert_eq!(cache.offset, 5);
+        let cap_before = ffi::array_shape(cache.keys.as_ref().unwrap())[2];
 
-        // Trim 2
+        // Trim 2. A dense trim is a logical rewind (issue #1209): only the
+        // offset moves, and the buffer keeps its pre-allocated capacity so
+        // the next append does not regrow.
         let trimmed = cache.trim(2);
         assert_eq!(trimmed, 2);
         assert_eq!(cache.offset, 3);
-
-        // Verify shapes
+        assert_eq!(cache.seq_len(), 3);
         let k_shape = ffi::array_shape(cache.keys.as_ref().unwrap());
-        assert_eq!(k_shape, vec![1, 2, 3, 4]);
+        assert_eq!(k_shape, vec![1, 2, cap_before, 4]);
         let v_shape = ffi::array_shape(cache.values.as_ref().unwrap());
-        assert_eq!(v_shape, vec![1, 2, 3, 4]);
+        assert_eq!(v_shape, vec![1, 2, cap_before, 4]);
+
+        // The visible window is what the offset says, not the buffer length.
+        let (k_win, v_win) = cache.update_and_fetch(
+            ffi::ones(&[1, 2, 1, 4], dtype::FLOAT32),
+            ffi::ones(&[1, 2, 1, 4], dtype::FLOAT32),
+        );
+        assert_eq!(ffi::array_shape(&k_win), vec![1, 2, 4, 4]);
+        assert_eq!(ffi::array_shape(&v_win), vec![1, 2, 4, 4]);
+    }
+
+    /// The rollback contract behind the logical trim (issue #1209): a
+    /// trimmed-then-overwritten cache must fetch a window byte-identical to a
+    /// cache that never saw the rejected block, in Fp16 and in Int8 (whose
+    /// scale sidecars must rewind in lockstep).
+    #[test]
+    fn kv_cache_trim_then_append_matches_untrimmed_reference() {
+        let fill = |v: f32, n: i32| {
+            let data: Vec<f32> = (0..(2 * n * 4)).map(|i| v + i as f32 * 0.125).collect();
+            ffi::from_slice_f32(&data, &[1, 2, n, 4])
+        };
+        let bytes = |arr: &crate::ffi::MlxArray| {
+            ffi::eval(arr);
+            ffi::array_to_raw_bytes(arr)
+        };
+        for mode in [KVCacheMode::Fp16, KVCacheMode::Int8] {
+            let mut cache = KVCache::new_with_mode(mode);
+            let mut reference = KVCache::new_with_mode(mode);
+            let _ = cache.update_and_fetch(fill(1.0, 5), fill(100.0, 5));
+            let _ = reference.update_and_fetch(fill(1.0, 5), fill(100.0, 5));
+
+            // A rejected speculative block reaches only the trimmed cache.
+            let _ = cache.update_and_fetch(fill(-7.0, 3), fill(-70.0, 3));
+            assert_eq!(cache.trim(3), 3);
+
+            let (k, v) = cache.update_and_fetch(fill(2.0, 2), fill(200.0, 2));
+            let (k_ref, v_ref) = reference.update_and_fetch(fill(2.0, 2), fill(200.0, 2));
+            assert_eq!(ffi::array_shape(&k), ffi::array_shape(&k_ref));
+            assert_eq!(bytes(&k), bytes(&k_ref), "{mode:?} keys diverge after trim");
+            assert_eq!(
+                bytes(&v),
+                bytes(&v_ref),
+                "{mode:?} values diverge after trim"
+            );
+        }
     }
 
     #[test]