diff --git a/mistralrs-core/src/kv_cache/mod.rs b/mistralrs-core/src/kv_cache/mod.rs index 46ff55e03..d93385645 100644 --- a/mistralrs-core/src/kv_cache/mod.rs +++ b/mistralrs-core/src/kv_cache/mod.rs @@ -806,6 +806,68 @@ pub(crate) fn ensure_uniform_batch_cache_lens( } } +/// 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 CacheManager for NormalCacheManager { fn clone_in_cache( &self, @@ -840,6 +902,12 @@ impl CacheManager for NormalCa } } + // `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 @@ mod clone_in_cache_invariant_tests { 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)] diff --git a/mistralrs-core/src/kv_cache/single_cache.rs b/mistralrs-core/src/kv_cache/single_cache.rs index 4c017d82a..d8f426a57 100644 --- a/mistralrs-core/src/kv_cache/single_cache.rs +++ b/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 diff --git a/mistralrs-core/src/kv_cache/xs_rolling.rs b/mistralrs-core/src/kv_cache/xs_rolling.rs index b4d84b459..a525f8d69 100644 --- a/mistralrs-core/src/kv_cache/xs_rolling.rs +++ b/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> { @@ -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 \