diff --git a/mistralrs-core/src/kv_cache/mod.rs b/mistralrs-core/src/kv_cache/mod.rs index cf3d0136d..39d9faf6b 100644 --- a/mistralrs-core/src/kv_cache/mod.rs +++ b/mistralrs-core/src/kv_cache/mod.rs @@ -825,6 +825,17 @@ impl BatchSrc { /// The dead prefix does **not** grow without bound: `lead_pad_i` is /// `max_j L_j - L_i`, and sequences in one batch accept at statistically the /// same rate, so the spread is a `sqrt(steps)` random walk, not a linear one. +/// +/// ⚠️ That last sentence is only true if the prefix is **taken back out** on +/// the way to the per-sequence caches. `current_seq_len` here counts the dead +/// columns as live (it has to — this is where the shared append offset comes +/// from), so a caller that records this length as the sequence's own makes the +/// next `front_align_batch` pad relative to an already-padded length and the +/// prefix accumulates **linearly**. [`drop_dead_prefix`] is the inverse that +/// keeps the claim honest, and +/// `MtpSpeculativePipeline::step`'s per-sequence commit is the caller that +/// applies it. `the_dead_prefix_does_not_accumulate_across_steps` pins the +/// difference, with the un-stripped variant as its negative control. fn front_pad_single(sc: &mut SingleCache, target_len: usize) -> Result { let live = sc.current_seq_len; if target_len < live { @@ -897,6 +908,73 @@ pub(crate) fn front_pad_kv_cache(cache: &mut KvCache, target_len: usize) -> Resu Ok(lead_k) } +/// Drop the `lead` dead columns [`front_pad_single`] put ahead of one +/// `SingleCache`'s live run, so the run is left-anchored at column 0 again and +/// `current_seq_len` counts only positions that hold real K/V. +fn drop_front_single(sc: &mut SingleCache, lead: usize) -> Result<()> { + if lead == 0 { + return Ok(()); + } + if lead > sc.current_seq_len { + candle_core::bail!( + "kv-cache: cannot drop a {lead}-column dead prefix from a slot that holds only {} \ + position(s) — the caller's idea of this row's padding and the slot's own length \ + disagree", + sc.current_seq_len + ); + } + let Some(ad) = sc.all_data.as_ref() else { + candle_core::bail!("kv-cache: drop_dead_prefix on a slot whose buffer is not materialised"); + }; + let have = ad.dims()[sc.dim]; + if have < sc.current_seq_len { + candle_core::bail!( + "kv-cache: slot claims {} live position(s) but its buffer is only {have} wide", + sc.current_seq_len + ); + } + let kept = have - lead; + sc.all_data = Some(ad.narrow(sc.dim, lead, kept)?.contiguous()?); + sc.capacity_seq_len = kept; + sc.current_seq_len -= lead; + Ok(()) +} + +/// The inverse of [`front_pad_kv_cache`]: take a row's dead prefix back out +/// once the dense batched forward that needed it is over. +/// +/// 🔑 This is what stops the prefix from accumulating. Left-alignment is the +/// only way a ragged cohort can share one append offset, but the length it +/// leaves behind (`lead + live`) is not the sequence's own length. Recording +/// *that* as the per-sequence length makes the next `front_align_batch` pad +/// relative to it, so every step adds another `max_j c_j - c_i` columns and the +/// buffer grows linearly in steps rather than tracking the `sqrt(steps)` spread +/// of the sequences themselves. Stripping the prefix restores the invariant +/// every other part of the cache assumes — that a slot's live run starts at +/// column 0 — which is exactly what `front_pad_single`'s `narrow(dim, 0, live)` +/// requires the next time round. +/// +/// `lead == 0` (every B=1 request, every uniform batch) is a no-op that touches +/// no tensor. An `XsRolling` slot is likewise untouched: `front_pad_kv_cache` +/// never gave it a prefix — its compressed rows are start-anchored and its raw +/// window is re-anchored per row by `XsRollingCache::split_row`. +pub(crate) fn drop_dead_prefix(cache: &mut KvCache, lead: usize) -> Result<()> { + if matches!(cache, KvCache::XsRolling(_)) || lead == 0 { + return Ok(()); + } + let KvCache::Normal { k, v } = cache else { + candle_core::bail!( + "kv-cache: drop_dead_prefix is only defined for a `Normal` slot; this one is a `{}`. \ + Nothing front-pads it, so nothing should be stripping it either. See \ + `KvCache::supports_per_sequence_len`.", + cache.kind_name() + ); + }; + drop_front_single(k, lead)?; + drop_front_single(v, lead)?; + Ok(()) +} + /// Left-align a whole ragged cohort so every sequence's live K/V ends at the /// batch maximum, and report each sequence's dead prefix. /// @@ -2455,6 +2533,153 @@ mod clone_in_cache_invariant_tests { } } + /// 🔑 The inverse round-trip. Front padding then stripping must return the + /// slot to exactly what it was — same rows, same length — because that is + /// the only reason a per-sequence length recorded after a dense batched + /// forward means what it says. + #[test] + fn dropping_the_dead_prefix_inverts_front_pad() { + let mut slot = materialised_slot(5, 16, 100.0); + let before = rows(&slot)[..5].to_vec(); + let lead = front_pad_kv_cache(&mut slot, 9).unwrap(); + assert_eq!(lead, 4); + drop_dead_prefix(&mut slot, lead).unwrap(); + assert_eq!( + slot.current_seq_len(), + 5, + "after the strip the slot's length is its OWN live run again, not the batch's width" + ); + assert_eq!( + rows(&slot)[..5], + before[..], + "the live run must survive the round trip bit-for-bit" + ); + } + + /// 🔴 The claim `front_pad_single`'s doc makes — that the dead prefix + /// tracks the `sqrt(steps)` spread of the sequences rather than growing + /// linearly — is a property of the code, not of the batch. It holds only + /// because the prefix is stripped on the way out. + /// + /// The fixture is a two-row cohort that diverges by a fixed 2 positions + /// every step, which is the *worst* case for the spread and still must not + /// compound. The negative control is the same loop with the strip removed: + /// that is the pre-change behaviour, and it grows without bound. + /// + /// ⚠️ Arithmetic on real tensors, not a hardware measurement (D14). + #[test] + fn the_dead_prefix_does_not_accumulate_across_steps() { + // One "step": pad both rows up to the batch max, pretend the forward + // appended `w`, keep `commit` of it, then (optionally) strip. + let run = |strip: bool, steps: usize| -> (Vec, Vec>) { + let w = 4usize; + let mut slots = [ + materialised_slot(4, 512, 1.0), + materialised_slot(4, 512, 2.0), + ]; + let commits = [3usize, 1]; + for _ in 0..steps { + let target = slots + .iter() + .map(|s| s.current_seq_len()) + .max() + .expect("two slots"); + let leads: Vec = slots + .iter_mut() + .map(|s| front_pad_kv_cache(s, target).unwrap()) + .collect(); + for (i, slot) in slots.iter_mut().enumerate() { + // The batched forward writes `w` new positions for every row. + slot.set_len(target + w).unwrap(); + if strip { + drop_dead_prefix(slot, leads[i]).unwrap(); + // The slot's OWN length after the strip, deliberately + // not recomputed here: `target + w - lead`. Deriving it + // in the test instead would let a strip that moves no + // data and updates no length still look right. + let live = slot.current_seq_len() - w; + slot.set_len(live + commits[i]).unwrap(); + } else { + slot.set_len(target + commits[i]).unwrap(); + } + } + } + ( + slots.iter().map(|s| s.current_seq_len()).collect(), + slots.iter().map(|s| rows(s)[..4].to_vec()).collect(), + ) + }; + + // 4 committed to start, then `n` steps of 3 and 1 accepted tokens. The + // slow row's true length is `4 + n`; anything above that is padding + // being counted as content. + for n in [6usize, 12] { + let (lens, heads) = run(true, n); + assert_eq!( + lens, + vec![4 + n * 3, 4 + n], + "with the strip, each row's recorded length is exactly the tokens it committed" + ); + assert_eq!( + heads, + vec![vec![1.0, 2.0, 3.0, 4.0], vec![2.0, 3.0, 4.0, 5.0]], + "and the live run is back at column 0 holding its own K/V — the pad/strip pair \ + has to be an identity on content, or the length is right about the wrong rows" + ); + } + // Negative control — the pre-change behaviour. The error is the slow + // row's recorded length minus its true one, and it must grow with the + // step count rather than settling. + let err_at = |n: usize| run(false, n).0[1] - (4 + n); + assert_eq!( + (err_at(6), err_at(12)), + (10, 22), + "without the strip the slow row is inflated, and doubling the steps must roughly \ + double the inflation — that is the linear accumulation, not a `sqrt(steps)` spread" + ); + assert!( + err_at(12) >= 2 * err_at(6) - 2, + "the inflation must scale with the step count, not saturate" + ); + } + + /// A strip wider than the live run means the caller's per-row padding and + /// the slot disagree. Silently clamping would leave the sequence reading + /// another row's K/V, so it refuses by name (D18). + #[test] + fn dropping_more_dead_prefix_than_the_slot_holds_is_refused() { + let mut slot = materialised_slot(5, 16, 1.0); + let err = drop_dead_prefix(&mut slot, 6).unwrap_err().to_string(); + assert!( + err.contains("dead prefix") && err.contains("disagree"), + "the refusal must say whose bookkeeping disagrees; got {err}" + ); + } + + /// A zero-width strip is the B=1 path and the uniform-batch path. It must + /// not reallocate or perturb anything. + #[test] + fn dropping_a_zero_width_dead_prefix_is_a_no_op() { + let mut slot = materialised_slot(7, 16, 3.0); + let before = rows(&slot); + drop_dead_prefix(&mut slot, 0).unwrap(); + assert_eq!(rows(&slot), before); + assert_eq!(slot.current_seq_len(), 7); + } + + /// The compressor slot is never front-padded, so it must never be + /// stripped either — its per-row token counts are the state the whole + /// path exists to carry. + #[test] + fn dropping_the_dead_prefix_leaves_an_xs_slot_untouched() { + xs_rolling::test_override::with(true, || { + let mut xs = KvCache::XsRolling(Box::new(XsRollingCache::new(4, 2, 64, 2048))); + let before = xs.current_seq_len(); + drop_dead_prefix(&mut xs, 3).unwrap(); + assert_eq!(xs.current_seq_len(), before); + }); + } + /// With `ARC_V4_XS_PER_SEQ` off the slot that blocks DeepSeek V4 refuses /// **by name**. A silent success here would corrupt the compressor's /// distant-context branch, which nothing downstream checks. diff --git a/mistralrs-core/src/pipeline/inputs_processor.rs b/mistralrs-core/src/pipeline/inputs_processor.rs index f5686c93c..7625cd044 100644 --- a/mistralrs-core/src/pipeline/inputs_processor.rs +++ b/mistralrs-core/src/pipeline/inputs_processor.rs @@ -188,9 +188,37 @@ pub mod text_models_inputs_processor { pub seq_indices: Vec, } + /// The per-row absolute offsets a batch needs, or `None` when it shares + /// one — which is every batch except the fused MTP step's left-aligned + /// ragged cohort. + /// + /// 🔑 Returning `None` for a batch that shares an offset is a **dispatch, + /// not an optimisation**: `None` sends [`make_prompt_chunk`] down the + /// pre-change scalar code verbatim, so every prompt, every decode and every + /// batch built with per-sequence KV advance off produces the characters it + /// always did. Split out as its own function for the reason PR #100's + /// `resolve_ragged_rows` was: the two branches agree exactly on their + /// common domain, so no numeric test can tell you which one ran, and a + /// mutation that routes a uniform batch down the per-row path is invisible + /// to everything except a test of the dispatch itself. + pub(crate) fn resolve_row_offsets( + per_seq: &[Option], + shared: usize, + ) -> Option> { + if !per_seq.iter().any(Option::is_some) { + return None; + } + Some(per_seq.iter().map(|o| o.unwrap_or(shared)).collect()) + } + // chunk_offset_toks is the number of tokens by which the tokens are offset, // chunk_offset_toks / prompt_chunksize = number of batches // + // row_offsets: when provided, replaces `chunk_offset_toks` per sequence — + // the batch is a left-aligned ragged cohort with no shared absolute + // position (see `resolve_row_offsets`). `None` is the shared-offset path + // and is bit-for-bit what it always was. + // // prefix_cache_lens: when provided, indicates how many tokens per sequence are already // cached in the paged KV cache. Only new (non-cached) tokens will be included in the // input tensor, and slot_mappings will only cover new token slots. Block tables still @@ -206,7 +234,17 @@ pub mod text_models_inputs_processor { mut paged_attn_metadata: Option<&mut PagedAttentionMeta>, mapper: Option<&dyn DeviceMapper>, prefix_cache_lens: Option<&[usize]>, + row_offsets: Option<&[usize]>, ) -> Result { + if let Some(offsets) = row_offsets { + if offsets.len() != seq_ids.len() { + anyhow::bail!( + "inputs processor: {} per-row absolute offsets for a batch of {}", + offsets.len(), + seq_ids.len() + ); + } + } // Determine effective tokens per sequence after prefix cache trimming let effective_lens: Vec = toks .iter() @@ -234,6 +272,9 @@ pub mod text_models_inputs_processor { let has_any_cache_hit = prefix_cache_lens.is_some_and(|lens| lens.iter().any(|&l| l > 0)); for (seq_idx, (seq_id, ctxt)) in seq_ids.iter().zip(&toks).enumerate() { let cached = prefix_cache_lens.map_or(0, |lens| lens[seq_idx]); + // This row's own absolute start, which is the shared + // `chunk_offset_toks` for every batch that has one. + let chunk_offset_toks = row_offsets.map_or(chunk_offset_toks, |o| o[seq_idx]); let full_prompt_len = ctxt.len(); // The new (non-cached) tokens to process let new_toks = &ctxt[cached..]; @@ -848,6 +889,15 @@ pub mod text_models_inputs_processor { mapper: Option<&dyn DeviceMapper>, ) -> Result { let offset = input_seqs[0].token_offset(); + // A left-aligned ragged cohort has no shared absolute position, so each + // row carries its own. `None` — every other batch — keeps `offset`. + let row_offsets = resolve_row_offsets( + &input_seqs + .iter() + .map(|s| s.prefill_seqlen_offset()) + .collect::>(), + offset, + ); // Collect prefix cache lens when paged attention is in use let prefix_cache_lens: Vec = input_seqs.iter().map(|s| s.prefix_cache_len()).collect(); @@ -866,6 +916,7 @@ pub mod text_models_inputs_processor { } else { None }, + row_offsets.as_deref(), ) .map(|inputs| InnerInputProcessorOutput { inputs, @@ -1146,4 +1197,110 @@ pub mod text_models_inputs_processor { InputsProcessorType::Text } } + + #[cfg(test)] + mod tests { + use super::{make_prompt_chunk, resolve_row_offsets}; + use crate::device_map::DummyDeviceMapper; + use candle_core::Device; + + /// 🔑 The wire itself: what the model actually receives. `positions` is + /// `seqlen_offsets`, which is where RoPE places each row's queries and — + /// on DeepSeek V4 — the per-row `row_q0` that masks a left-aligned + /// cohort. A per-row vector in must produce a per-row vector out. + #[test] + fn per_row_offsets_reach_seqlen_offsets_and_a_shared_one_still_does_not() { + let mapper = DummyDeviceMapper { + nm_device: Device::Cpu, + }; + let window: Vec = vec![7, 8, 9, 10]; + let toks: Vec<&[u32]> = vec![&window, &window, &window]; + let run = |row_offsets: Option<&[usize]>, shared: usize| -> Vec { + make_prompt_chunk( + shared, + toks.clone(), + &[0, 1, 2], + &Device::Cpu, + Some((window.len(), 0)), + false, + None, + Some(&mapper), + None, + row_offsets, + ) + .unwrap() + .positions + }; + + assert_eq!( + run(None, 40), + vec![40, 40, 40], + "with no per-row offsets the batch shares one, exactly as it always did" + ); + assert_eq!( + run(Some(&[40, 33, 37]), 40), + vec![40, 33, 37], + "each row's own absolute position must survive to `seqlen_offsets`" + ); + } + + /// A per-row vector that does not describe this batch is a caller bug + /// that would silently mis-place two rows' queries. It refuses. + #[test] + fn a_per_row_offset_vector_of_the_wrong_width_is_refused() { + let mapper = DummyDeviceMapper { + nm_device: Device::Cpu, + }; + let window: Vec = vec![7, 8]; + let err = make_prompt_chunk( + 0, + vec![&window, &window, &window], + &[0, 1, 2], + &Device::Cpu, + Some((2, 0)), + false, + None, + Some(&mapper), + None, + Some(&[1, 2]), + ) + .err() + .expect("a per-row vector that does not describe this batch must be refused") + .to_string(); + assert!( + err.contains("per-row absolute offsets") && err.contains("batch of 3"), + "the refusal must name both widths; got {err}" + ); + } + + /// 🔑 The dispatch, pinned directly. Every batch that has ever existed + /// carries no per-row override, and must return `None` — because `None` + /// is what runs `make_prompt_chunk`'s pre-change scalar code verbatim. + /// + /// ⚠️ Routing a batch that happens to agree down the per-row path is + /// **numerically invisible** (the vector would be the shared offset + /// repeated), so only a test of the dispatch itself can see that + /// mutation. Same trap PR #100 documented for `resolve_ragged_rows`. + #[test] + fn a_batch_with_no_per_row_offset_takes_the_shared_path() { + assert_eq!(resolve_row_offsets(&[None, None, None], 17), None); + assert_eq!(resolve_row_offsets(&[], 17), None); + } + + /// One row carrying its own position is enough to make the batch + /// per-row; the rest fall back to the shared offset rather than to zero, + /// which is what keeps a partially-overridden batch coherent. + #[test] + fn a_row_with_its_own_offset_makes_the_whole_batch_per_row() { + assert_eq!( + resolve_row_offsets(&[Some(40), Some(33), Some(37)], 17), + Some(vec![40, 33, 37]) + ); + assert_eq!( + resolve_row_offsets(&[Some(40), None], 17), + Some(vec![40, 17]), + "a row without its own offset keeps the batch's, not 0" + ); + } + } } diff --git a/mistralrs-core/src/pipeline/mtp_pipeline.rs b/mistralrs-core/src/pipeline/mtp_pipeline.rs index 4fc90f201..8e1ff95df 100644 --- a/mistralrs-core/src/pipeline/mtp_pipeline.rs +++ b/mistralrs-core/src/pipeline/mtp_pipeline.rs @@ -1309,11 +1309,18 @@ impl MtpSpeculativePipeline { /// wave63-CO under `ARC_V4_XS_PER_SEQ`; /// 2. the target model masks a left-aligned ragged batch's dead prefix AND /// its start-anchored compressed columns - /// ([`Self::target_masks_ragged_batches`]) — cleared for V4 by the - /// per-row `row_q0` in `dsv4_attention`; - /// 3. this pipeline's own step arithmetic can express a batch whose rows - /// hold different lengths ([`Self::step_carries_per_sequence_lengths`]) - /// — **still refused**, and now the only one. + /// ([`model_masks_ragged_batches`]) — cleared for V4 by the per-row + /// `row_q0` in `dsv4_attention`; + /// 3. the MTP **draft** chain can carry rows at different absolute + /// positions ([`draft_chain_carries_per_sequence_positions`]) — **still + /// refused**, and now the only one. + /// + /// This pipeline's own step arithmetic used to be (3). It no longer is: + /// every quantity in [`Self::step`] now comes from `cache_lens[i]` + /// ([`resolve_row_cache_lens`]), the verify window routes + /// `seqlen_offsets[i]` per row, and the per-sequence commit strips the dead + /// prefix ([`crate::kv_cache::drop_dead_prefix`]) so a row's recorded length + /// is its own rather than the batch's padding. /// /// None of these is decoration. Per-sequence lengths in one dense buffer /// mean `front_pad_kv_cache` zero-fills ahead of the shorter rows, and a @@ -1323,11 +1330,11 @@ impl MtpSpeculativePipeline { /// already paid for once with FP8 KV. /// /// (3) is worth stating explicitly because granting the mode on (1) and (2) - /// alone would be **actively harmful**, not merely useless: `window_ok` - /// would fail on every ragged step, the fused path would defer to the - /// target's own decode, and that path builds no mask at all at `t_q == 1` - /// (`layers_masker.rs:285`) — so the dead prefix would be attended as real - /// keys by the very fallback that was supposed to be the safe one. + /// alone would mean every ragged row silently stopped drafting — the group + /// loop drops any row whose seed is not at the group's chain start — so MTP + /// would report a healthy `accept_rate` over chains it never proposed. A + /// mode that is inert is bad; a mode that is inert *and* still reports is + /// D18's "the absence of a signal read as a specific signal". pub(crate) fn kv_advance(&self) -> KvAdvance { *self.kv_advance.get_or_init(|| { if !per_sequence_kv_requested() { @@ -1619,6 +1626,37 @@ impl MtpSpeculativePipeline { Ok(true) } + /// `clone_in_cache`, keeping the one thing it computes and throws away. + /// + /// Under [`KvAdvance::PerSequence`] the cohort arriving here is genuinely + /// ragged, so it is left-aligned first: every sequence's live K/V is shifted + /// to END at the batch maximum, which makes `SingleCache::append`'s one + /// shared write offset simultaneously correct for all of them. The dense + /// batching below is then completely unchanged. + /// + /// 🔑 The returned `lead_pad[i]` is the **only** record that row `i`'s share + /// of the batched buffer starts `lead_pad[i]` columns in. Once + /// `front_pad_single` has run, the slot's own `current_seq_len` is the + /// batch's width for every row — the per-row length is gone. Discarding this + /// vector is therefore what forced every downstream quantity onto one shared + /// `cache_len`; keeping it is what lets [`resolve_row_cache_lens`] hand it + /// straight back. + /// + /// Off the per-sequence path it is `vec![0; seqs.len()]` and no tensor is + /// touched, exactly as before. + fn clone_in_cache_reporting_lead( + &self, + seqs: &mut [&mut Sequence], + ) -> candle_core::Result> { + let lead_pad = if self.kv_advance() == KvAdvance::PerSequence { + crate::kv_cache::front_align_batch(seqs, false)? + } else { + vec![0; seqs.len()] + }; + get_mut_arcmutex!(self.target).clone_in_cache(seqs)?; + Ok(lead_pad) + } + /// Log the "draft KV could not be primed, drafting skipped" explanation /// once per process. fn warn_unprimed_once(&self, reason: &str) { @@ -1923,12 +1961,13 @@ pub(crate) fn cache_supports_per_sequence_advance( /// /// Order matters only for which reason gets reported. Each check is /// independent, and each has been the *last* one at some point: (1) fell to -/// wave63-CO, (2) to the per-row `row_q0`, and (3) is where it stands now. +/// wave63-CO, (2) to the per-row `row_q0`, (3) to this wave's per-row step +/// arithmetic, and (4) is where it stands now. pub(crate) fn per_sequence_refusal(cache: &EitherCache) -> Option { cache_supports_per_sequence_advance(cache) .err() .or_else(|| model_masks_ragged_batches(cache).err()) - .or_else(|| step_carries_per_sequence_lengths().err()) + .or_else(|| draft_chain_carries_per_sequence_positions().err()) } /// Whether the model behind this cache masks a left-aligned ragged batch, or @@ -1966,55 +2005,146 @@ pub(crate) fn model_masks_ragged_batches(cache: &EitherCache) -> std::result::Re ) } -/// Whether the **fused step's own bookkeeping** can express a batch whose -/// rows hold different committed lengths. +/// Per-row committed KV lengths for one fused step, or `None` when the batch +/// shares one — which is every batch that was not built by +/// [`crate::kv_cache::front_align_batch`] under per-sequence advance. +/// +/// `lead_pad[i]` is the dead prefix that call left ahead of row `i`, so +/// `cache_len - lead_pad[i]` is row `i`'s own committed KV length: the batched +/// buffer's width is the batch **maximum**, and every row was padded up to it. +/// No value is invented here — the vector already existed and was being +/// discarded by `clone_in_cache`. +/// +/// 🔑 Returning `None` for a uniform batch is a **dispatch, not an +/// optimisation**. `None` runs the pre-change scalar arithmetic verbatim, so +/// every B=1 request and every batch built with the flag off is byte-identical +/// by control flow rather than by a number some test happened to observe. It is +/// a separate function for the reason PR #100's `resolve_ragged_rows` is: the +/// two branches agree exactly on their common domain, so no numeric test can +/// say which one ran. +/// +/// The two refusals describe the same invariant from opposite sides. +/// `front_align_batch` pads every row **up to** the batch maximum, so no lead +/// can exceed the width and at least one row must have no lead at all. A batch +/// that is ragged in some other way is not the layout this step knows how to +/// unpick, and guessing would be a wrong answer nothing downstream catches. +pub(crate) fn resolve_row_cache_lens( + enabled: bool, + cache_len: usize, + lead_pad: &[usize], +) -> Result>> { + if !enabled || lead_pad.iter().all(|&l| l == 0) { + return Ok(None); + } + if let Some(&over) = lead_pad.iter().find(|&&l| l > cache_len) { + candle_core::bail!( + "MTP per-sequence step: a dead prefix of {over} column(s) is wider than the \ + {cache_len}-column batched cache it was padded into, so the pipeline's idea of this \ + batch's layout and `front_align_batch`'s disagree" + ); + } + if !lead_pad.contains(&0) { + candle_core::bail!( + "MTP per-sequence step: every row of this batch carries a dead prefix ({lead_pad:?}), \ + so none of them sits at the shared end column {cache_len}. `front_align_batch` pads \ + UP TO the batch maximum, so at least one row must have no prefix — this is not the \ + left-aligned cohort the per-row arithmetic is defined for" + ); + } + Ok(Some(lead_pad.iter().map(|&l| cache_len - l).collect())) +} + +/// The window invariant, per row. +/// +/// `uncached[i]` is the committed tail sequence `i` has not yet put in its KV, +/// and `ok` is whether every row is in a state this fast path produced. Both +/// used to be computed from ONE shared `cache_len` — the *batched* buffer's +/// width, which left-alignment makes `max_j L_j` rather than any row's length — +/// which made them wrong in two different ways the moment the rows diverged: +/// +/// * `toks_i.len() - cache_len` saturated to **0** for every row but the +/// longest, so `ok` failed on every ragged step and per-sequence advance +/// would have been *reachable but inert*; +/// * `cache_len` then indexed the verify window (`toks[cache_len..]`), a +/// **panic** site the saturation was the only thing hiding. +/// +/// Both close by taking row `i`'s own length. `ok` also guarantees the slice is +/// in range: `uncached[i] >= 1` means `cache_lens[i] < tok_lens[i]`. +pub(crate) fn plan_step_window(cache_lens: &[usize], tok_lens: &[usize], w: usize) -> StepWindow { + let uncached: Vec = cache_lens + .iter() + .zip(tok_lens) + .map(|(&c, &t)| t.saturating_sub(c)) + .collect(); + let ok = cache_lens + .iter() + .zip(&uncached) + .all(|(&c, &u)| c + u >= 2 && (1..=w).contains(&u)); + StepWindow { uncached, ok } +} + +/// [`plan_step_window`]'s answer. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct StepWindow { + pub uncached: Vec, + pub ok: bool, +} + +/// Which rows of a batch may draft together: those needing the same number of +/// drafts (`w - u`) **and** sitting at the same absolute chain start. +/// +/// Both halves are forced by the draft path, not chosen: one batched MTP-block +/// forward takes ONE `pos` (`MtpBlock::forward_step`, used as the RoPE position +/// and the draft-KV slot) and `batch_draft_caches` refuses rows of differing +/// length. Named so the cost of the second half is something a test can state +/// rather than something a benchmark discovers — see +/// [`draft_chain_carries_per_sequence_positions`]. +pub(crate) const fn draft_group_key(uncached: usize, cache_len: usize) -> (usize, usize) { + (uncached, cache_len + uncached - 2) +} + +/// Whether the **MTP draft chain** can carry a batch whose rows sit at +/// different absolute positions. /// -/// It cannot, and this is the blocker that remains after the mask. It is -/// stated separately, and named, because the two previous waves each -/// believed the mask was the last one. +/// It cannot, and this is the blocker that remains after the fused step's own +/// arithmetic. It is stated separately, and named, because each of the three +/// previous waves believed the one it removed was the last. /// -/// [`MtpSpeculativePipeline::step_full`] keys the entire step on ONE scalar -/// `cache_len = current_normal_cache_len(self)` — the *batched* buffer's -/// width, which under left-alignment is `max_j L_j`, not any row's length. -/// Four of its uses are arithmetically wrong the moment the rows diverge, -/// and they fail in different ways: +/// The target's time base is now per-row end to end: the cache carries its own +/// length per row (wave63-CO), `dsv4_attention` masks a left-aligned cohort +/// (PR #100), and this step derives every quantity from `cache_lens[i]` and +/// routes `seqlen_offsets[i]` per row. The **draft** KV is a second, independent +/// time base, and it is still keyed on one scalar in four places: /// -/// * `uncached[i] = toks_i.len() - cache_len` saturates to **0** for every -/// row but the longest, which fails `window_ok`'s `(1..=w)` check and -/// sends every step down the deferral path — per-sequence advance would -/// be *reachable but inert*, and the deferral is itself unsafe here -/// because the target's own decode builds no mask at `t_q == 1`; -/// * `toks[cache_len..]` (the verify window) is a **panic** site once -/// `cache_len > toks_i.len()`, reachable only because `window_ok` bails -/// first; -/// * `run_target_forward_batch(.., Some((w, cache_len)))` becomes -/// `seqlen_offsets[i] = cache_len` for every row -/// (`inputs_processor.rs:244` — `last_n_context_len.1` is a shared -/// scalar), so RoPE would place a shorter row's queries `lead_i` past -/// their true positions. That same vector is what feeds the model's -/// `row_q0`, so the mask this PR built is correct but would be handed -/// uniform values; -/// * `set_len(cache_len + committed_i)` counts the dead prefix **inside** -/// the recorded length (`front_pad_single` sets `current_seq_len = -/// target_len`, and `clone_out_cache` copies the batched length verbatim -/// into every row). So the next `front_align_batch` pads relative to -/// already-padded lengths and the dead prefix **accumulates linearly**, -/// not as the `sqrt(steps)` random walk `front_pad_single`'s doc claims. -/// Nothing in `Sequence` or `SingleCache` records a row's true live -/// length; the only honest source is `seq.get_toks().len()`. +/// * `MtpHiddenCapture::store(seqlen_offsets.first(), ..)` +/// (`deepseek4.rs`) records row 0's absolute position for the whole capture, +/// so `extend_draft_kv_row`'s `off` describes only the leading row; +/// * `propose_chain_batched(.., start_pos, ..)` takes ONE `start_pos` for the +/// group and hands `start_pos + i` to `MtpBlock::forward_step`, whose `pos` is +/// both the RoPE position and the draft-KV slot; +/// * `batch_draft_caches` refuses rows whose `current_seq_len` differs, which +/// under per-sequence advance is all of them; +/// * so the fused step can only group rows that share an absolute chain start. +/// That grouping is now explicit and correct — see the `(u, chain_start)` key +/// in [`MtpSpeculativePipeline::step`] — but with `u_i ≡ 1` and every row at +/// its own length it degenerates to one group per sequence, i.e. `B` single-row +/// MTP-block forwards per chain step instead of `<= depth + 1` batched ones. /// -/// Fixing it is a change to this pipeline's step arithmetic plus a per-row -/// route through `inputs_processor` — not a model change, which is why it -/// is not in the PR that added the mask. -pub(crate) fn step_carries_per_sequence_lengths() -> std::result::Result<(), String> { +/// Correct and `B`x slower is not a mode worth granting, and granting it would +/// also hide the real fix, which is a per-row `pos` in `MtpBlock::forward_step` +/// (a positions vector for RoPE plus a per-row draft-KV slot) and a per-row +/// capture offset. That is a **model** change, in `deepseek4`, which is why it +/// is not in the PR that fixed the pipeline arithmetic. +pub(crate) fn draft_chain_carries_per_sequence_positions() -> std::result::Result<(), String> { Err( - "`step_full` keys the whole fused step on one shared `cache_len` (the BATCHED buffer \ - width, which left-alignment makes `max_j L_j`): `uncached[i]` saturates to 0 for \ - every row but the longest and trips `window_ok`, `toks[cache_len..]` would panic, \ - `prefill_window` collapses `seqlen_offsets` to one scalar so the model's per-row \ - `row_q0` is fed uniform values, and `set_len(cache_len + committed_i)` counts the \ - dead prefix inside the length so it accumulates. See \ - `MtpSpeculativePipeline::step_carries_per_sequence_lengths`." + "the MTP draft chain keys on ONE absolute position for a whole group: \ + `MtpHiddenCapture::store` records row 0's offset for the entire capture, \ + `propose_chain_batched` takes a single `start_pos` and passes `start_pos + i` to \ + `MtpBlock::forward_step` as both the RoPE position and the draft-KV slot, and \ + `batch_draft_caches` refuses rows of differing length. Rows at different positions can \ + only be grouped one-per-sequence, which is correct but costs B single-row block \ + forwards per chain step. See \ + `MtpSpeculativePipeline::draft_chain_carries_per_sequence_positions`." .to_string(), ) } @@ -2307,15 +2437,13 @@ impl CacheManagerMixin for MtpSpeculativePipeline { /// unchanged, and `ensure_uniform_batch_cache_lens` passes for the same /// reason it always did. /// - /// `front_align_batch` returns each row's dead prefix. It is discarded here - /// only because [`MtpSpeculativePipeline::kv_advance`] refuses to enter - /// this mode until a model consumes that mask; the moment one does, this is - /// where the mask is born. + /// `front_align_batch` returns each row's dead prefix, which is the only + /// record anywhere of how much of the batched buffer is padding rather than + /// this row's own KV. The fused step needs it, so it calls + /// [`MtpSpeculativePipeline::clone_in_cache_reporting_lead`] directly; this + /// trait entry point is for every other caller and discards it. fn clone_in_cache(&self, seqs: &mut [&mut Sequence]) -> candle_core::Result<()> { - if self.kv_advance() == KvAdvance::PerSequence { - let _lead_pad = crate::kv_cache::front_align_batch(seqs, false)?; - } - get_mut_arcmutex!(self.target).clone_in_cache(seqs) + self.clone_in_cache_reporting_lead(seqs).map(|_| ()) } fn clone_out_cache(&self, seqs: &mut [&mut Sequence]) { get_mut_arcmutex!(self.target).clone_out_cache(seqs); @@ -2593,8 +2721,19 @@ impl Pipeline for MtpSpeculativePipeline { // PRE-cache instruction: clone-in / reset, matching what // `Pipeline::step` does internally for the target alone. + // + // `lead_pad[i]` is how many columns of the batched buffer are padding + // ahead of row `i` rather than that row's own K/V — see + // [`Self::clone_in_cache_reporting_lead`]. It is all-zero unless this + // step assembled the batch itself, which is why a ragged cohort can only + // be unpicked on the `In` arm. + let mut lead_pad = vec![0usize; input_seqs.len()]; + let mut assembled_here = false; match pre_op { - CacheInstruction::In => self.clone_in_cache(input_seqs)?, + CacheInstruction::In => { + lead_pad = self.clone_in_cache_reporting_lead(input_seqs)?; + assembled_here = true; + } CacheInstruction::Nothing => (), CacheInstruction::Reset { reset_non_granular, @@ -2622,22 +2761,57 @@ impl Pipeline for MtpSpeculativePipeline { // ---- The window invariant ---- // - // `cache_len` (C) is shared by the whole batch — the scheduler buckets - // on it (`scheduler/default_scheduler.rs`) and `clone_in_cache` builds - // one dense batched cache from it. Each sequence's committed length is - // `C + u`, where `u` is its uncached committed tail; plain decode holds - // `u == 1`. `w = depth + 1` is the fixed window width every sequence - // feeds: `u` real tokens, then `w - u` drafts. + // `cache_len` (C) is the width of the one dense batched buffer + // `clone_in_cache` built. Under the cohort rule it is also every + // sequence's committed length, and `cache_lens` is `C` repeated. Under + // per-sequence advance it is only the batch **maximum**: row `i` was + // padded up to it and owns `C - lead_pad[i]` of it, which is what + // `cache_lens[i]` carries. + // + // Each sequence's committed length is `cache_lens[i] + u_i`, where `u` + // is its uncached committed tail; plain decode holds `u == 1`. + // `w = depth + 1` is the fixed window width every sequence feeds: `u` + // real tokens, then `w - u` drafts. let w = self.window(); let batch = input_seqs.len(); let cache_len = current_normal_cache_len(self); - let uncached: Vec = input_seqs - .iter() - .map(|s| s.get_toks().len().saturating_sub(cache_len)) - .collect(); - let window_ok = cache_len + uncached.iter().copied().min().unwrap_or(0) >= 2 - && uncached.iter().all(|u| (1..=w).contains(u)); + // Per-sequence advance needs the batched buffer split back into each + // sequence's own storage before the lengths can diverge, and that is + // what `CacheInstruction::Out` does. It also needs this step to have + // assembled the batch, because that is the only place `lead_pad` comes + // from. Any other pair leaves the batched cache in the pipeline, where + // there is no per-sequence length to set, so those steps take the cohort + // rule — correctly, and without a mid-run state that no other path knows + // how to read. + let per_seq_advance = assembled_here + && matches!(post_op, CacheInstruction::Out) + && self.kv_advance() == KvAdvance::PerSequence; + let row_cache_lens = resolve_row_cache_lens(per_seq_advance, cache_len, &lead_pad)?; + let cache_lens: Vec = match &row_cache_lens { + Some(rows) => rows.clone(), + None => vec![cache_len; batch], + }; + let tok_lens: Vec = input_seqs.iter().map(|s| s.get_toks().len()).collect(); + let StepWindow { + uncached, + ok: window_ok, + } = plan_step_window(&cache_lens, &tok_lens, w); if !window_ok { + // 🔴 A left-aligned batch cannot be handed to the target's own + // decode: that path builds no mask at all at `t_q == 1` + // (`layers_masker.rs:285`), so every row's dead prefix would be + // attended as real keys — a zero K row scores logit 0 and takes + // softmax weight. Deferring is only safe for a batch that is not + // padded, so a padded one refuses instead. + if row_cache_lens.is_some() { + candle_core::bail!( + "MTP per-sequence step: this batch is left-aligned (dead prefixes \ + {lead_pad:?}) but its window invariant does not hold (uncached tails \ + {uncached:?} against a window of {w}). Deferring to the target's own decode \ + would attend the dead prefix as real keys, because that path builds no mask \ + at `t_q == 1`." + ); + } // Not a state this fast path produced — a cache reset, a // prefix-cache hit, a chunked prefill. The target's own bookkeeping // is authoritative there, so defer (the PRE op is already applied). @@ -2684,26 +2858,43 @@ impl Pipeline for MtpSpeculativePipeline { .map(|id| self.checkout_draft_kv(*id)) .collect(); - // ---- Draft, grouped by uncached-tail length ---- + // ---- Draft, grouped by (uncached-tail length, absolute chain start) ---- // // Within a group every sequence sits at the SAME absolute chain start - // `C + u - 2` (because they share `C`) and needs the SAME number of - // drafts `w - u`, so one batched MTP-block forward per chain step - // drafts for all of them. `u` ranges over `[1, w]`, so a batch splits - // into at most `depth + 1` groups no matter how large it is — the draft - // cost stays a small multiple of one single-layer forward, not a - // multiple of the batch size. + // `cache_lens[i] + u - 2` and needs the SAME number of drafts `w - u`, + // so one batched MTP-block forward per chain step drafts for all of + // them: `MtpBlock::forward_step` takes ONE `pos`, used both as the RoPE + // position and as the draft-KV slot, and `batch_draft_caches` refuses + // rows of differing length. + // + // When the batch shares `cache_len` — every batch today — the chain + // start is a function of `u` alone, so this partitions and iterates + // exactly as the `u`-keyed map it replaces: `u` ranges over `[1, w]`, so + // a batch splits into at most `depth + 1` groups no matter how large it + // is, and the draft cost stays a small multiple of one single-layer + // forward rather than a multiple of the batch size. + // + // 🔴 It is stated as a pair anyway because under per-sequence advance it + // is NOT a function of `u`: rows sit at their own lengths, so the + // partition degenerates to one group per sequence. That is correct and + // `B`x slower, and it is exactly what + // [`draft_chain_carries_per_sequence_positions`] refuses the mode over. + // Keying on `u` alone there would not be slow, it would be silent: a row + // whose seed is not at the group's chain start is skipped, so it would + // stop drafting with no signal but a falling `tok_per_step`. let mut drafts: Vec> = vec![Vec::new(); batch]; - let mut groups: std::collections::BTreeMap> = + let mut groups: std::collections::BTreeMap<(usize, usize), Vec> = std::collections::BTreeMap::new(); for (i, &u) in uncached.iter().enumerate() { if u < w { - groups.entry(u).or_default().push(i); + groups + .entry(draft_group_key(u, cache_lens[i])) + .or_default() + .push(i); } } - for (u, members) in groups { + for ((u, chain_start), members) in groups { let n_draft = w - u; - let chain_start = cache_len + u - 2; let mut rows: Vec = Vec::with_capacity(members.len()); let mut seeds: Vec = Vec::with_capacity(members.len()); let mut last_toks: Vec = Vec::with_capacity(members.len()); @@ -2804,15 +2995,30 @@ impl Pipeline for MtpSpeculativePipeline { // other direction, skipping the final draft forward because the // draft-extend it must run anyway already produced that token // (`eagle_worker.py:871-873`). + // + // 🔑 `seqlen_offsets[i]` is where the model places row `i`'s queries — + // for RoPE, and (on DeepSeek V4) as the per-row `row_q0` that masks a + // left-aligned cohort's dead prefix and its start-anchored compressed + // columns. A shared `prefill_window` collapses it to one scalar, which + // is right only when the rows share a length. When they do not, each row + // carries its own and the shared component is zero — the vector the + // model needs is the one this step already computed. for (i, seq) in input_seqs.iter_mut().enumerate() { let toks = seq.get_toks(); - let mut window: Vec = toks[cache_len..].to_vec(); + let mut window: Vec = toks[cache_lens[i]..].to_vec(); window.extend(drafts[i].iter().copied()); let pad = *window.last().expect("uncached tail is at least 1 token"); window.resize(w, pad); - seq.set_prefill_toks(window); + match row_cache_lens { + Some(_) => seq.set_prefill_toks_at(window, cache_lens[i]), + None => seq.set_prefill_toks(window), + } } - let verify = run_target_forward_batch(self, input_seqs, Some((w, cache_len))).await; + let shared_offset = match row_cache_lens { + Some(_) => 0, + None => cache_len, + }; + let verify = run_target_forward_batch(self, input_seqs, Some((w, shared_offset))).await; for seq in input_seqs.iter_mut() { seq.reset_prefill_toks(); } @@ -2863,15 +3069,10 @@ impl Pipeline for MtpSpeculativePipeline { // on the way out and its next uncached tail is 1 — the B=1 invariant, // at every batch size. let n_accepted: Vec = accepted.iter().map(Vec::len).collect(); - // Per-sequence advance needs the batched buffer split back into each - // sequence's own storage before the lengths can diverge, and that is - // what `CacheInstruction::Out` does. Any other POST op leaves the - // batched cache in the pipeline, where there is no per-sequence length - // to set, so those steps take the cohort rule — correctly, and without - // a mid-run state that no other path knows how to read. - let advance = match post_op { - CacheInstruction::Out => self.kv_advance(), - _ => KvAdvance::Cohort, + let advance = if per_seq_advance { + KvAdvance::PerSequence + } else { + KvAdvance::Cohort }; let per_sequence_commit = match advance { KvAdvance::Cohort => { @@ -3010,11 +3211,31 @@ impl Pipeline for MtpSpeculativePipeline { // here, in the same scope that computed it, from a plain `Vec` indexed // by the batch row — and there is no reduction across `input_seqs` in // it, which is the whole difference from the cohort path. + // + // 🔑 And the length recorded is this row's OWN, not the batch's width. + // `clone_out_cache` hands each row the batched buffer verbatim, dead + // prefix included, so the prefix is stripped first + // (`drop_dead_prefix` — the inverse of `front_pad_kv_cache`). Recording + // `cache_len + c_i` instead would count `lead_pad[i]` padding columns as + // this sequence's content, and the next `front_align_batch` would pad + // relative to that: the prefix would compound every step, growing + // linearly rather than tracking the `sqrt(steps)` spread of the + // sequences themselves, and the rows would silently re-converge onto one + // length — the cohort barrier this whole path exists to remove, + // reappearing as padding. if let Some(committed) = per_sequence_commit { for (i, c) in committed.iter().enumerate() { - let target = cache_len + c; + let lead = lead_pad[i]; + let target = cache_lens[i] + c; for slot in input_seqs[i].normal_cache().iter_mut().flatten() { let kind = slot.kind_name(); + if let Err(e) = crate::kv_cache::drop_dead_prefix(slot, lead) { + candle_core::bail!( + "MTP per-sequence advance: stripping the {lead}-column dead prefix \ + failed on a {kind} slot for batch row {i} ({e}). Continuing would \ + record that padding as the sequence's own length." + ); + } if let Err(e) = slot.set_len(target) { candle_core::bail!( "MTP per-sequence advance: set_len({target}) refused on a {kind} slot \ @@ -4375,20 +4596,18 @@ mod tests { ); } - /// 🔴 **The refusal moved; it did not disappear.** With both of the + /// 🔴 **The refusal moved; it did not disappear.** With all three /// previously-named blockers cleared — the cache carries per-row lengths - /// (wave63-CO) and V4's attention masks a ragged cohort (this wave) — a V4 - /// target must STILL be refused, and the reason must now name the fused - /// step's shared `cache_len`. + /// (wave63-CO), V4's attention masks a ragged cohort (PR #100), and this + /// step's arithmetic is per-row — a V4 target must STILL be refused, and the + /// reason must now name the MTP **draft** chain, which is a second and + /// independent time base. /// - /// This is the test that stops the third check from being quietly dropped. - /// Granting the mode on the first two alone is not merely useless, it is - /// harmful: `window_ok` fails on every ragged step, the fused path defers - /// to the target's own decode, and that path builds no mask at all at - /// `t_q == 1` — so the dead prefix would be attended as real keys by the - /// fallback that was supposed to be the safe one. + /// This is the test that stops the remaining check from being quietly + /// dropped, and the assertions that it no longer names the cleared ones are + /// what make "the refusal moved" a claim a test can fail. #[test] - fn a_fully_capable_v4_cache_is_still_refused_and_now_names_the_step() { + fn a_fully_capable_v4_cache_is_still_refused_and_now_names_the_draft_chain() { use crate::kv_cache::{EitherCache, KvCache, NormalCache, XsRollingCache}; use std::sync::{Arc, Mutex}; @@ -4412,20 +4631,228 @@ mod tests { }); let why = refusal.expect( - "per-sequence advance must STILL be refused: the fused step keys on one shared \ - `cache_len`, and granting the mode without fixing that defers every ragged step \ - into a decode path that builds no mask at all", + "per-sequence advance must STILL be refused: the MTP draft chain keys a whole group \ + on one absolute position, so rows at their own lengths can only be grouped \ + one-per-sequence", ); assert!( - why.contains("cache_len") && why.contains("window_ok"), - "the refusal must name the fused step's shared `cache_len`, not a blocker that is \ - already cleared; got {why}" + why.contains("draft") && why.contains("forward_step"), + "the refusal must name the draft chain and the scalar `pos` it turns on, not a \ + blocker that is already cleared; got {why}" ); assert!( !why.contains("XsRolling"), "the refusal still names the compressor slot, so the flag is not reaching the \ - cache and this wave changed nothing; got {why}" + cache and wave63-CO changed nothing; got {why}" + ); + assert!( + !why.contains("window_ok") && !why.contains("cache_len"), + "the refusal still names the fused step's shared `cache_len`, which THIS wave \ + removed — so it changed nothing; got {why}" + ); + } + + /// 🔑 The dispatch, pinned directly. A uniform batch — every B=1 request, + /// every batch built with the flag off — must return `None`, because `None` + /// is what runs the pre-change scalar arithmetic verbatim. + /// + /// ⚠️ This is the mutation PR #100 warned about and caught the hard way: + /// routing a uniform batch down the per-row path is **numerically + /// invisible**, since `cache_lens` would just be `cache_len` repeated. It is + /// still a real defect — it flips `set_prefill_toks_at` on for every decode, + /// builds a `[B]` offsets vector per forward, and sends `dsv4_attention` a + /// `row_q0` it has to re-check on all 43 layers. Only a test of the dispatch + /// itself can see it. + #[test] + fn a_uniform_batch_never_reaches_the_per_row_step_arithmetic() { + assert_eq!( + resolve_row_cache_lens(true, 40, &[0, 0, 0, 0]).unwrap(), + None, + "no row carries a dead prefix, so the batch is uniform and must take the scalar path" ); + assert_eq!( + resolve_row_cache_lens(true, 40, &[]).unwrap(), + None, + "an empty batch is uniform" + ); + assert_eq!( + resolve_row_cache_lens(true, 40, &[0, 7, 0, 3]).unwrap(), + Some(vec![40, 33, 40, 37]), + "a row's own length is the batch width minus its dead prefix" + ); + } + + /// The per-row route is opt-in on the same flag that authorizes the only + /// producer of a ragged batch. With it off, even a batch that somehow + /// carries leads takes the scalar path — the gate is a guarantee, not a + /// coincidence. + #[test] + fn the_per_row_step_arithmetic_is_opt_in() { + assert_eq!( + resolve_row_cache_lens(false, 40, &[0, 7, 0, 3]).unwrap(), + None, + "with per-sequence advance off, nothing may take the per-row path" + ); + } + + /// `front_align_batch` pads every row UP TO the batch maximum, so no lead + /// can exceed the width and at least one row must have none. A vector that + /// breaks either half is not the layout this arithmetic is defined for, and + /// guessing would be a wrong answer nothing downstream catches. + #[test] + fn a_batch_that_is_not_a_left_aligned_cohort_is_refused_by_name() { + let err = resolve_row_cache_lens(true, 10, &[0, 12]) + .unwrap_err() + .to_string(); + assert!( + err.contains("wider than") && err.contains("disagree"), + "a lead past the batch width must say whose bookkeeping disagrees; got {err}" + ); + + let err = resolve_row_cache_lens(true, 10, &[2, 3]) + .unwrap_err() + .to_string(); + assert!( + err.contains("every row") && err.contains("end column"), + "a batch with no row at the shared end column must say so; got {err}" + ); + } + + /// 🔴 Defect 1, pinned as arithmetic. Fed the batched buffer's width, a + /// ragged cohort's uncached tails saturate to 0 for every row but the + /// longest, `window_ok` fails, and every ragged step defers — per-sequence + /// advance would be reachable but inert. Fed each row's own length, every + /// tail is 1, which is the B=1 invariant at every batch size. + /// + /// ⚠️ Arithmetic on the plan function, not a hardware measurement (D14). + #[test] + fn per_row_lengths_are_what_keep_a_ragged_cohorts_tails_inside_the_window() { + let w = 4usize; + // Three rows at 40/33/37 committed tokens, each one token ahead of its + // own cache — the steady state per-sequence advance produces. + let rows = [40usize, 33, 37]; + let tok_lens: Vec = rows.iter().map(|r| r + 1).collect(); + let cache_len = 40usize; + + let shared = plan_step_window(&[cache_len; 3], &tok_lens, w); + assert_eq!( + shared.uncached, + vec![1, 0, 0], + "the shared width saturates every row but the longest" + ); + assert!( + !shared.ok, + "which trips the window invariant, so the fused step would defer on every ragged step" + ); + + let per_row = plan_step_window(&rows, &tok_lens, w); + assert_eq!(per_row.uncached, vec![1, 1, 1]); + assert!(per_row.ok, "each row's own length puts every tail at 1"); + } + + /// 🔴 Defect 2, pinned as arithmetic: `toks[cache_lens[i]..]` is a panic + /// site unless `ok` guarantees the slice is in range. It does, for every + /// row, and that guarantee is the ONLY thing standing between the verify + /// window and an index past the end. + /// + /// ⚠️ Arithmetic on the plan function, not a hardware measurement (D14). + #[test] + fn the_window_invariant_is_what_keeps_the_verify_slice_in_range() { + let w = 4usize; + for cache_lens in [ + vec![10usize, 10, 10], + vec![10usize, 7, 9], + vec![3usize, 3, 3], + ] { + for extra in [0usize, 1, 2, 4, 5] { + let tok_lens: Vec = cache_lens.iter().map(|c| c + extra).collect(); + let plan = plan_step_window(&cache_lens, &tok_lens, w); + if plan.ok { + for (i, &c) in cache_lens.iter().enumerate() { + assert!( + c < tok_lens[i], + "ok must imply row {i}'s slice start {c} is inside its {} tokens", + tok_lens[i] + ); + } + } + } + } + // A cache longer than the sequence — what the shared width looks like to + // a short row — must never be reported ok. + assert!( + !plan_step_window(&[12, 12], &[13, 9], w).ok, + "a row whose cache claims more positions than it has tokens is not a state this \ + fast path produced" + ); + } + + /// 🔴 The evidence behind the refusal that remains. Drafting groups on + /// `(u, chain_start)` because one MTP-block forward carries one absolute + /// position; on a batch that shares a length that is a function of `u` + /// alone, so the partition is unchanged and bounded by `depth + 1`. Under + /// per-sequence advance `u ≡ 1` and every row is at its own length, so the + /// partition degenerates to one group per sequence — `B` single-row block + /// forwards per chain step. + /// + /// Correct, and `B`x slower. That is why the mode is refused rather than + /// granted, and this is the arithmetic that says so. + /// + /// ⚠️ A count of groups, not a hardware measurement (D14). + #[test] + fn the_draft_group_partition_collapses_when_the_rows_carry_their_own_lengths() { + use std::collections::BTreeSet; + let count = |uncached: &[usize], cache_lens: &[usize]| -> usize { + uncached + .iter() + .zip(cache_lens) + .map(|(&u, &c)| draft_group_key(u, c)) + .collect::>() + .len() + }; + + // The cohort batch: one shared width, tails spread over the window. + let uncached = [1usize, 2, 1, 3, 2, 1, 4, 1]; + let shared = [40usize; 8]; + assert_eq!( + count(&uncached, &shared), + 4, + "sharing a width makes the chain start a function of `u`, so the partition is \ + bounded by the window and one batched forward drafts for the whole group" + ); + + // The per-sequence batch: every tail is 1 (the point of the mode), every + // row is at its own length. + let per_seq_lens = [40usize, 33, 37, 31, 39, 35, 30, 38]; + let ones = [1usize; 8]; + assert_eq!( + count(&ones, &per_seq_lens), + 8, + "with rows at their own lengths the partition is one group per sequence — the draft \ + chain's shared absolute position, which is what `per_sequence_refusal` names" + ); + } + + /// The two window rules coincide on a uniform batch, which is the other + /// half of "flag off is byte-identical": the extracted function must not + /// have changed what the inline expression decided. + #[test] + fn the_window_invariant_is_unchanged_for_a_uniform_batch() { + let w = 4usize; + for cache_len in [0usize, 1, 2, 9] { + for tails in [vec![1usize, 1], vec![1, 4], vec![0, 1], vec![5, 1]] { + let tok_lens: Vec = tails.iter().map(|t| cache_len + t).collect(); + let plan = plan_step_window(&[cache_len; 2], &tok_lens, w); + let legacy = cache_len + tails.iter().copied().min().unwrap_or(0) >= 2 + && tails.iter().all(|u| (1..=w).contains(u)); + assert_eq!( + plan.ok, legacy, + "cache_len {cache_len}, tails {tails:?}: the per-row rule must agree with the \ + scalar one it replaced" + ); + assert_eq!(plan.uncached, tails); + } + } } /// The plan functions agree on the one thing they must: the slots a diff --git a/mistralrs-core/src/pipeline/normal.rs b/mistralrs-core/src/pipeline/normal.rs index 29b2e553c..942083156 100644 --- a/mistralrs-core/src/pipeline/normal.rs +++ b/mistralrs-core/src/pipeline/normal.rs @@ -969,6 +969,7 @@ impl Loader for NormalLoader { None, Some(pipeline_mapper.as_ref()), None, + None, )?; model.forward( diff --git a/mistralrs-core/src/pipeline/vision.rs b/mistralrs-core/src/pipeline/vision.rs index 9ffca047f..b33f86cf5 100644 --- a/mistralrs-core/src/pipeline/vision.rs +++ b/mistralrs-core/src/pipeline/vision.rs @@ -770,6 +770,7 @@ impl Loader for VisionLoader { None, None, None, + None, )?; let _ = model.forward( &inputs.input, diff --git a/mistralrs-core/src/sequence.rs b/mistralrs-core/src/sequence.rs index 03a862aea..734ec2c72 100644 --- a/mistralrs-core/src/sequence.rs +++ b/mistralrs-core/src/sequence.rs @@ -437,6 +437,21 @@ pub struct Sequence { // Prefix caching prefill_prompt_toks: Option>, + /// This sequence's OWN absolute position for the first token of + /// [`Self::prefill_prompt_toks`], when the batch it is in does not share + /// one. + /// + /// `None` — always, unless a caller set it with + /// [`Self::set_prefill_toks_at`] — means the batch shares + /// `input_seqs[0].token_offset()`, which is what every path did before and + /// still does. It exists because a left-aligned ragged cohort + /// ([`crate::kv_cache::front_align_batch`]) has no shared offset to collapse + /// to: row `i`'s queries start at its own committed KV length, and RoPE and + /// DeepSeek V4's per-row `row_q0` both read that from `seqlen_offsets`. + /// + /// Cleared by [`Self::reset_prefill_toks`] together with the window itself, + /// so an offset can never outlive the tokens it describes. + prefill_seqlen_offset: Option, /// Number of tokens at the start of the prompt that are cached (KV already computed). /// These tokens should be skipped during prefill. prefix_cache_len: usize, @@ -575,6 +590,7 @@ impl Sequence { creation_time, recognizer, prefill_prompt_toks: None, + prefill_seqlen_offset: None, prefix_cache_len: 0, suffix, prefix, @@ -847,9 +863,28 @@ impl Sequence { self.prefill_prompt_toks = Some(toks) } + /// [`Self::set_prefill_toks`] for a sequence whose batch has **no shared** + /// absolute position: `offset` is this row's own, and it replaces the + /// batch-wide `token_offset` in `seqlen_offsets`. + /// + /// Only the fused MTP step's left-aligned ragged path calls this. Every + /// other caller uses [`Self::set_prefill_toks`] and leaves the offset + /// `None`, which is the pre-change behaviour verbatim — see + /// `inputs_processor::resolve_row_offsets`, where the dispatch lives. + pub fn set_prefill_toks_at(&mut self, toks: Vec, offset: usize) { + self.prefill_prompt_toks = Some(toks); + self.prefill_seqlen_offset = Some(offset); + } + + /// This row's own absolute position for its prefill window, if it has one. + pub fn prefill_seqlen_offset(&self) -> Option { + self.prefill_seqlen_offset + } + /// Remove the prefill tokens. pub fn reset_prefill_toks(&mut self) { - self.prefill_prompt_toks = None + self.prefill_prompt_toks = None; + self.prefill_seqlen_offset = None; } /// Internal api to add one raw token.