fix(kv-cache): reconcile XsRolling base — a plain-path batching bug split out of #86

mistralrs-core/src/kv_cache/mod.rs

@@ -292,7 +292,7 @@
 
 /// Enable TurboQuant for all subsequently-created [`NormalCache`]s, with equal
 /// K and V widths.
 pub fn set_turboquant_head_dim(head_dim: usize) {
     set_turboquant_kv_head_dims(head_dim, head_dim);
 }
 
@@ -806,6 +806,68 @@
     }
 }
 
+/// Reconcile the one quantity [`ensure_uniform_batch_cache_lens`] cannot see:
+/// [`XsRollingCache::base`].
+///
+/// Uniform `current_seq_len` is **not** sufficient for `tail` (which is
+/// `tokens - base` wide) to batch, because `base` tracks the sequence's
+/// *high-water* length rather than its current one: `set_len` narrows the tail
+/// without moving `base`, and `advance` never lowers it. Two sequences at an
+/// identical token count can therefore hold tails of different width — see
+/// [`XsRollingCache::trim_tail_to`] for the ways that happens.
+///
+/// This needs no speculative decoding to reach. `prefix_cacher.rs` calls
+/// `set_len` on every stored layer, so a sequence restored from a prefix-cache
+/// entry that was stored at a greater length holds a narrower tail than one that
+/// arrived directly — and `clone_in_cache` then batches the two with
+/// `slice_set`, which demands an exact dim match. That is a second, independent
+/// route into the same `shape mismatch on dim 1, …` failure that
+/// `ensure_uniform_batch_cache_lens` was added to close, and it is on the
+/// **plain decode path**.
+///
+/// Raising every member to the batch maximum is lossless by construction: the
+/// member already sitting at `base_max` proves no future compressed row needs a
+/// token below it. `trim_tail_to` refuses rather than silently dropping history
+/// if that were ever untrue.
+fn reconcile_xs_bases(
+    seqs: &mut [&mut crate::sequence::Sequence],
+    layer: usize,
+    modify_draft_cache: bool,
+) -> Result<()> {
+    let mut base_max = 0usize;
+    let mut any = false;
+    for seq in seqs.iter_mut() {
+        let cache = if modify_draft_cache {
+            seq.normal_draft_cache()
+        } else {
+            seq.normal_cache()
+        };
+        if let Some(KvCache::XsRolling(xs)) = cache.get(layer).and_then(|s| s.as_ref()) {
+            base_max = base_max.max(xs.base);
+            any = true;
+        }
+    }
+    if !any || base_max == 0 {
+        return Ok(());
+    }
+    for (i, seq) in seqs.iter_mut().enumerate() {
+        let cache = if modify_draft_cache {
+            seq.normal_draft_cache()
+        } else {
+            seq.normal_cache()
+        };
+        if let Some(KvCache::XsRolling(xs)) = cache.get_mut(layer).and_then(|s| s.as_mut()) {
+            xs.trim_tail_to(base_max).map_err(|e| {
+                candle_core::Error::msg(format!(
+                    "kv-cache: cannot reconcile cache slot {layer} of seqs[{i}] to the batch's \
+                     retained-window start {base_max}: {e}"
+                ))
+            })?;
+        }
+    }
+    Ok(())
+}
+
 impl<T: CacheManagerMixin + MetadataMixin + ?Sized> CacheManager<T> for NormalCacheManager {
     fn clone_in_cache(
         &self,
@@ -840,6 +902,12 @@
                 }
             }
 
+            // `tail` is content, not slack, so it cannot be padded — but it
+            // *can* be trimmed to a common start, which is lossless. Do that
+            // before gathering, so the batch's widths agree by construction
+            // rather than by luck.
+            reconcile_xs_bases(seqs, layer, modify_draft_cache)?;
+
             // Gather every sequence's two tensors for this slot up front, so
             // the batch's shape is decided by the whole batch instead of by
             // `seqs[0]` alone. `None` marks a sequence with no slot here; it is
@@ -2378,6 +2446,121 @@
             Some((2, 100, 101, 1))
         );
     }
+
+    /// Equal `current_seq_len` is NOT sufficient for `tail` to batch.
+    ///
+    /// This is the second, independent route into the `shape mismatch on dim 1`
+    /// failure that `ensure_uniform_batch_cache_lens` closed only half of — and
+    /// it needs **no MTP whatsoever**. `prefix_cacher` calls `set_len` on every
+    /// stored layer, so a sequence restored from an entry stored at a greater
+    /// length holds a narrower tail than one that reached the same token count
+    /// directly.
+    ///
+    /// Mutation check: delete the `reconcile_xs_bases` call from
+    /// `clone_in_cache` and this test fails on the `clone_in_cache` line with
+    /// `shape mismatch on dim 1, 4 <> 132`.
+    #[test]
+    fn xs_base_divergence_at_equal_lengths_is_reconciled_not_refused() {
+        // Restored from a prefix-cache entry stored at 300 tokens, truncated to
+        // 260 — `base` stays at canonical(300), which is past canonical(260).
+        let mut restored = xs_state(128, 1);
+        feed_xs(&mut restored, 300);
+        assert_eq!(
+            restored.base, 256,
+            "canonical(300) for ratio 128, margin 16"
+        );
+        restored.set_len(260).unwrap();
+
+        // Reached 260 directly.
+        let mut direct = xs_state(128, 1);
+        feed_xs(&mut direct, 260);
+        assert_eq!(direct.base, 128, "canonical(260) sits a group lower");
+
+        // --- Fixture discrimination (D12) -------------------------------
+        assert_eq!(
+            (restored.current_seq_len(), direct.current_seq_len()),
+            (260, 260),
+            "the LENGTHS must agree, or this test is reproducing the ragged-length \
+             failure instead of the base one — `ensure_uniform_batch_cache_lens` \
+             would catch that and never reach the tensor"
+        );
+        assert_eq!(
+            (
+                restored.tail.as_ref().unwrap().dims()[1],
+                direct.tail.as_ref().unwrap().dims()[1]
+            ),
+            (4, 132),
+            "and the WIDTHS must disagree, or there is nothing to reconcile"
+        );
+        assert_eq!(
+            legacy_batch_error(
+                restored.tail.as_ref().unwrap(),
+                direct.tail.as_ref().unwrap()
+            ),
+            "shape mismatch on dim 1, 4 <> 132",
+            "uniform lengths still panic the legacy batcher — this is a second, \
+             independent way into the slice_set shape mismatch"
+        );
+
+        // --- The fix ----------------------------------------------------
+        let pipeline = StubPipeline::new(2);
+        let mut a = seq_with_slots(
+            0,
+            260,
+            vec![
+                v4_kv_slot(512, 8, 260),
+                KvCache::XsRolling(Box::new(restored)),
+            ],
+        );
+        let mut b = seq_with_slots(
+            1,
+            260,
+            vec![
+                v4_kv_slot(512, 8, 260),
+                KvCache::XsRolling(Box::new(direct)),
+            ],
+        );
+        let mut seqs: Vec<&mut Sequence> = vec![&mut a, &mut b];
+        NormalCacheManager
+            .clone_in_cache(&pipeline, &mut seqs, false)
+            .expect("equal lengths with unequal retained windows must reconcile");
+
+        let batched = pipeline.cache().normal();
+        let KvCache::XsRolling(xs) = &batched.0[1] else {
+            panic!("slot 1 must stay an XsRolling entry")
+        };
+        assert_eq!(
+            xs.base, 256,
+            "trimmed to the batch's largest retained start"
+        );
+        assert_eq!(xs.tail.as_ref().unwrap().dims(), &[2, 4, XS_HIDDEN]);
+    }
+
+    /// `trim_tail_to` is lossless only up to what the compressor still needs,
+    /// and must refuse — by name — past it, rather than silently dropping the
+    /// rows the next compressed row is built from.
+    #[test]
+    fn trimming_the_retained_window_past_what_the_compressor_needs_is_refused() {
+        let mut state = xs_state(4, 2);
+        feed_xs(&mut state, 260);
+        let needs_from = state.compressor_needs_from();
+        assert!(
+            needs_from >= state.base,
+            "an untouched cache always retains what it needs"
+        );
+        assert!(
+            state.trim_tail_to(needs_from).is_ok(),
+            "trimming to exactly the needed start is the lossless boundary"
+        );
+        let err = state
+            .trim_tail_to(needs_from + 1)
+            .expect_err("one token past it drops history the next row is built from")
+            .to_string();
+        assert!(
+            err.contains("next compressed row is built from"),
+            "the refusal must name what it would have destroyed, got: {err}"
+        );
+    }
 }
 
 #[cfg(test)]

mistralrs-core/src/kv_cache/single_cache.rs

@@ -106,7 +106,9 @@ impl SingleCache {
         }
         let seq_len = src.dim(self.dim)?;
         if seq_len != 1 {
-            candle_core::bail!("append_graph is decode-only (src seq len must be 1, got {seq_len})");
+            candle_core::bail!(
+                "append_graph is decode-only (src seq len must be 1, got {seq_len})"
+            );
         }
         // Buffer must be at least `read_capacity` along the seq dim. Reuse the
         // existing (eager-populated) all_data so past K/V is present; only

mistralrs-core/src/kv_cache/xs_rolling.rs

@@ -187,6 +187,70 @@ impl XsRollingCache {
         Ok(())
     }
 
+    /// The first raw token any *future* compressed row can still need.
+    ///
+    /// Row `g` (the next one to complete) is built from group
+    /// `g + 1 - span_groups` onward, so everything below this index is dead
+    /// weight the cache is free to drop. It is the same expression
+    /// [`Self::advance`] uses for `need_start`, named once so the retention
+    /// rule and the trim rule cannot drift apart.
+    pub fn compressor_needs_from(&self) -> usize {
+        (self.comp.current_seq_len() + 1).saturating_sub(self.span_groups) * self.ratio
+    }
+
+    /// Raise `base` to `new_base`, dropping the leading raw rows.
+    ///
+    /// # Why this exists
+    ///
+    /// `tail` is `tokens - base` wide and `base` is **not** a function of
+    /// `tokens`: it is `canonical(high-water tokens)`, because [`Self::set_len`]
+    /// narrows the tail without moving `base` and [`Self::advance`] never lowers
+    /// it. Two sequences at the identical token count can therefore hold
+    /// different-width tails — one that rolled back across a group boundary, one
+    /// that arrived directly, one restored from a prefix-cache entry stored at a
+    /// greater length. `NormalCacheManager::clone_in_cache` batches `tail` with
+    /// `slice_set`, which demands an exact dim match, so that difference is the
+    /// `shape mismatch on dim 1, 18 <> 22` class of failure.
+    ///
+    /// Trimming to the batch's **largest** `base` reconciles them, and it is
+    /// lossless by construction: the member that already sits at `base_max`
+    /// proves no future compressed row needs a token below it. The rows dropped
+    /// here only ever shortened how far a *rollback* could reach, and after this
+    /// call every member of the batch shares the same reach — the one the shared
+    /// batched cache would have had anyway.
+    ///
+    /// Refuses (rather than silently dropping history) if `new_base` would cut
+    /// into rows the next compressed row is built from, which is the same
+    /// condition [`Self::advance`] bails on as a "compressor history gap".
+    pub fn trim_tail_to(&mut self, new_base: usize) -> Result<()> {
+        if new_base <= self.base {
+            return Ok(());
+        }
+        if new_base > self.tokens {
+            candle_core::bail!(
+                "xs rolling cache: cannot trim the raw window to start at {new_base}; the cache \
+                 only holds {} tokens",
+                self.tokens
+            );
+        }
+        let need_start = self.compressor_needs_from();
+        if need_start < new_base {
+            candle_core::bail!(
+                "xs rolling cache: trimming the raw window to {new_base} would drop tokens the \
+                 next compressed row is built from (it needs them from {need_start}). Resuming \
+                 would silently skip history."
+            );
+        }
+        let keep = self.tokens - new_base;
+        let drop = new_base - self.base;
+        self.tail = match self.tail.take() {
+            Some(t) if keep > 0 => Some(t.narrow(1, drop, keep)?.contiguous()?),
+            _ => None,
+        };
+        self.base = new_base;
+        Ok(())
+    }
+
     /// The compressed rows `[B, G, head_dim]`, or `None` when no group has
     /// completed yet (the sliding-window branch alone covers the context).
     pub fn compressed_rows(&self) -> Result<Option<Tensor>> {
@@ -230,7 +294,11 @@ impl XsRollingCache {
         if g_target > g_done {
             // Rows g_done..g_target need raw tokens from group
             // `g_done + 1 - span_groups` onward.
-            let need_start = (g_done + 1).saturating_sub(self.span_groups) * self.ratio;
+            let need_start = self.compressor_needs_from();
+            debug_assert_eq!(
+                need_start,
+                (g_done + 1).saturating_sub(self.span_groups) * self.ratio
+            );
             if need_start < win_start {
                 candle_core::bail!(
                     "xs rolling cache: compressor history gap — row {g_done} needs tokens from \