diff --git a/memory/mission/wave62-CN-mtp-per-sequence-kv.md b/memory/mission/wave62-CN-mtp-per-sequence-kv.md new file mode 100644 index 000000000..8402119e2 --- /dev/null +++ b/memory/mission/wave62-CN-mtp-per-sequence-kv.md @@ -0,0 +1,242 @@ +# wave62-CN โ€” MTP at batch: the cohort barrier removed, and the one wire that is left + +**Branch:** `feat/mtp-per-sequence-kv` off `02e441e22`. **PR #92.** +**Hardware used: NONE.** Everything below is code, arithmetic, or a CPU unit test. +**Nothing here is an acceptance measurement.** See ยง6 for the GPU ask. + +--- + +## 0. The correction this acts on + +wave59-CJ found the ceiling correctly and it was over-generalised into "MTP can't work at +batch". It can. The narrow, true statement is: + +> **Arc's MTP stalls at batch because the cohort shares ONE DENSE CACHE with a cohort-wide +> MIN-ROLLBACK. That is an Arc implementation choice, not a property of speculative decoding.** + +PR #86 pushed the wrong way โ€” it *enforces* the lockstep. Correct as a bug fix, wrong as a +destination. This removes the need for it. **#86 does not need to be reverted and is not +touched here**: `enforce_shared_cache_lockstep` is a `KvAdvance::Cohort` postcondition and +stays exactly that. Under `KvAdvance::PerSequence` it is not reached, because +`plan_per_sequence_step` never puts two sequences on one length in the first place. + +## 1. ๐Ÿ”‘ Verified: no production engine uses a cohort minimum + +Read from pinned `main` (vLLM `cc7cf71`, SGLang `c82e928`, TRT-LLM `8efd46e`), not from +memory. All three key the advance on a **per-request scalar**, inside a per-request loop: + +| engine | the line | +|---|---| +| vLLM v1 | `v1/core/sched/scheduler.py:1836-1848` โ€” `num_rejected = num_draft_tokens - num_accepted; request.num_computed_tokens -= num_rejected` | +| SGLang | `speculative/eagle_worker_common.py:587-588` `new_seq_lens = batch.seq_lens + accept_lens`; committed at `batch_result_processor.py:686-687` `req.kv_committed_len += num_accept_tokens` | +| TensorRT-LLM | `pyexecutor/resource_manager.py:1330-1346` rewinds `request.py_rewind_len` per request | + +A batch-wide reduction over accept lengths appears in **none** of them (`grep` for +`num_accepted.min` / `min(accept_len` returns only unrelated hits โ€” vLLM's cascade-attention +common-prefix pick, SGLang's per-request `finished_len` clamp, TRT-LLM's warmup floor). + +**What happens to a rejected draft's K/V: nothing.** vLLM and SGLang neither free nor zero +it โ€” the request's length shrinks so attention stops reading those slots, and the next +step's window overwrites the same physical positions (SGLang says so at +`eagle_worker_common.py:450`: "trailing unaccepted slots stay and are freed as overshoot"). +TRT-LLM is the only one that reclaims, and only when a rewind empties a whole block +(`kvCacheManager.cpp:4562-4600`). Prefix-cache correctness is kept by *capping the commit* +at the verified length (vLLM `kv_cache_manager.py:562-565`), never by erasing KV. + +**Lockstep is real โ€” but only in dense/padded batching outside paged engines.** It is the +stated motivation of EMS-SD ([arXiv 2405.07542]) and "Batch Speculative Decoding Done Right" +([arXiv 2510.22876]). That is exactly the regime Arc is in. Jish is right about production +engines; wave59 was right about Arc; the two were never in conflict. + +## 2. The barrier, stated as one equation + +`plan_batch_step` keeps `keep = min_i(u_i + a_i)`, so `next_u_i = u_i + a_i + 1 - keep`. +`plan_per_sequence_step` keeps `committed_i = u_i + a_i` and `next_u_i = 1`. Pinned by +`both_plans_agree_on_what_each_sequence_can_vouch_for`: + +``` +cohort.keep == min_i(per_seq.committed[i]) +``` + +**The cohort keep IS the batch minimum of the per-sequence commits.** That single `min` is +the whole defect, and it is now the only difference between the two rules โ€” everything else +(drafting, verify, `verify_proposed`, `window_verify_row`) is literally the same code path, +which is what makes the measurements in ยง3 attributable. + +## 3. ๐Ÿ“Š Measured โ€” arithmetic on the production plan functions, **not hardware** + +`tok_per_step_at_b128_is_bounded_by_the_cohort_rule_and_freed_by_the_per_sequence_rule`, +B=128, depth 3 (`w=4`), 200 steps, one fixture whose sequences accept at content-dependent +rates: + +| rule | tok/step | vs B=1 | +|---|---|---| +| cohort (today) | **1.1297** | โˆ’23.6% | +| per-sequence | **1.4783** | **0.0000** (exact, `< 1e-9`) | +| same 128 sequences run alone (B=1) | 1.4783 | โ€” | + +โš ๏ธ These are **fixture numbers, not V4's**. The absolute values depend on the oracle's +acceptance profile; the *ratio* is the claim. Read `tok_per_step`, never `accept_rate` โ€” a +saturated sequence drafts 0, contributes `proposed = 0`, and leaves `accept_rate` +flattering while `tok_per_step` collapses. + +โš ๏ธ **A correction to wave59-CJ ยง3.** It said saturation is a fixed point and `keep = 1` is +certain at B=128. The dynamics are milder: when *every* sequence reaches `u = w` the minimum +is `w`, so the whole cohort resets to `u = 1` at once and the system **oscillates** rather +than latching. The cost is 24% of `tok_per_step` on this fixture, not a total collapse. +Directionally wave59 was right and the fix is the same; the magnitude was overstated. + +## 4. What was built + +| piece | file | what it is | +|---|---|---| +| `plan_per_sequence_step`, `PerSeqStepPlan` | `pipeline/mtp_pipeline.rs` | the production rule; `next_uncached โ‰ก 1`, no reduction over the batch | +| `KvAdvance` + `MtpSpeculativePipeline::kv_advance()` | same | latched once, logged once, `ARC_MTP_PER_SEQ_KV=1`, default OFF | +| `cache_supports_per_sequence_advance` | same | names the first slot that cannot carry its own length | +| per-sequence commit in `step` | same | after the POST op splits the batched buffer, each sequence's own cache is set to `cache_len + u_i + a_i`; **no reduction across the batch, and no shadow state โ€” see below** | +| `KvCache::supports_per_sequence_len` / `kind_name` | `kv_cache/mod.rs` | the capability, with the reason each of the four variants gives | +| `front_pad_kv_cache`, `front_align_batch` | `kv_cache/mod.rs` | left-alignment: shift each row's live run so it **ends** at the batch max | +| `RaggedKvLens`, `make_left_padded_causal_mask` | `layers_masker.rs` | the `[B,1,t_q,k]` additive mask a left-aligned batch needs | + +### ๐Ÿ”‘ No second map keyed by sequence id + +The per-sequence state **is** `Sequence::normal_cache`'s own `current_seq_len`. The commit +happens in `step`, in the same scope that computed it, from a plain `Vec` indexed by the +batch row, immediately after `handle_post_cache_op` has split the batched buffer back into +per-sequence storage. An earlier draft staged the lengths in a +`HashMap` on the pipeline; that was removed. A shadow copy of a length the +cache also holds is precisely the divergence wave59-CJ ยง1 documents, reappearing in a new +place โ€” the bug being fixed, in a new hiding spot. (Thanks to the PR #90 author for the +catch.) + +Steps whose POST op is not `CacheInstruction::Out` take the cohort rule, because there is no +per-sequence storage to write to on those steps. + +### ๐Ÿ”‘ The paged route needs none of this โ€” and is a separate change + +`KVCacheManager::trim_request_to_num_tokens` under `ARC_SEGMENTED_KV=1` already routes to +`SegmentedAllocator::rollback(pool, request_id, 0, num_tokens)` +(`paged_attention/kv_cache_manager.rs:445-448`) โ€” **per-request, absolute token count, no +new API needed**. PagedAttention therefore has no cohort barrier to remove. What blocks +paged MTP is upstream: `MtpSpeculativePipeline::step`'s fast path only accepts +`CacheBackendMetadata::DefaultInstructions`, so a paged batch never reaches that rollback. + +โš ๏ธ And PR #90 makes V4's read *expressible*, not paged: `cache_write_and_gather` still +builds `cu_seqlens_kv` from per-sequence context lens and must consume +`SegmentPlan.cu_seqlens` (Bร—S rows) instead, and the compressor's `xs` history has no block +table at all. Do not assume V4 end-to-end from the segment allocator landing. + +### ๐Ÿ”‘ Why left-alignment is the dense-cache answer + +`SingleCache::append` writes every sequence's new K/V at **one** shared offset +(`single_cache.rs:225`). Right-padding a ragged cohort puts each row's hole in the middle, +where the next append lands. Left-alignment puts every row's live run against the **same +end column**, so that one shared offset is simultaneously correct for all of them โ€” the +dense batching code in `NormalCacheManager::clone_in_cache` needs **no change at all**, and +`ensure_uniform_batch_cache_lens` passes for the same reason it always did. What the caller +gains is that the length handed back on the way *out* may differ per sequence. + +The dead prefix does not grow without bound: it 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 (โ‰ˆ6.6% at 1000 steps, B=128). + +## 5. ๐Ÿ”ด The honest gap โ€” this does NOT move the B=128 number yet + +`KvAdvance::PerSequence` is **unreachable in production today**, deliberately, and the +refusal is logged by name. Two things block it and both are named in code: + +1. **`XsRolling` cannot carry its own length** (`KvCache::supports_per_sequence_len`). + V4's compressor slots hold two time bases (`comp` rows + a raw `tail`) plus scalar + `tokens`/`base` that are one number for the whole batch, and per-sequence advance needs + `advance` to append a **different number of rows per sequence** โ€” which one shared append + offset cannot express. Same gap wave61-CL ยง6 names ("the compressor's `xs` history has no + block table") and wave29-BC ยง4b ruled its own project. **This is now the third wave to + land on it. It is the keystone.** +2. **No model threads the ragged mask** (`target_masks_ragged_batches`). `front_align_batch` + already produces `lead_pad` per sequence and `make_left_padded_causal_mask` already turns + it into the mask; what is missing is carrying it from the pipeline through + `inputs_processor` into the model forward, where the mask is rebuilt from + `seqlen_offsets` alone today. `dsv4_attention::compose_caller_mask` already accepts rank + 4, so the consuming end is done. + +**Refusing is not conservatism, it is the only correct behaviour.** A zero-filled dead +prefix is *not* a masked prefix โ€” a zero K row scores logit 0 and takes real softmax weight. +Turning this on without (2) would be a wrong answer nothing downstream catches, which is the +exact failure this project already paid for once with FP8 KV. + +**MTP has exactly one model in the tree** (`mtp_decode_kit` โ†’ `deepseek4.rs:4250`; the trait +default returns `None`), so there is no second model to exercise the per-sequence path on +either. The order is forced: (1) โ†’ (2) โ†’ the number. + +### What was ruled out, with the reason + +* **Advance the shared cache by `w` and mask the dead rows.** V4's raw branch is a + ~128-token *suffix* window; 54% of it would be dead rows, so the effective window shrinks + to ~59 real tokens. Not token-identical. Dead on arrival. +* **A wider window `W = max_i(u_i) + depth`.** `u_i` grows by `a_i` per step whenever any + sequence sits at `u=1, a=0`, so `W` grows linearly and verify compute with it. +* **Discard the accepted-but-uncached surplus.** Throughput becomes `min_i a_i` โ‡’ 1.0 + tok/step at B=128. Strictly worse than today. +* **Bucket the cohort by accept length.** `select_running_bucket` runs one bucket per step, + so a B=128 cohort becomes โ‰ค4 buckets of ~32 and only one runs. This is wave59's option (B) + under another name. + +## 6. ๐Ÿ”ด The GPU ask + +Nothing here changes a default, so the check is a **no-regression** one, on the box that is +already serving V4: + +```bash +# A: control (current master behaviour โ€” the flag is OFF) +mistralrs bench -m -b 1 -b 128 2>&1 | tee /tmp/a.txt +# B: same box, same flags, per-sequence advance REQUESTED +ARC_MTP_PER_SEQ_KV=1 ARC_MTP_LOG_ACCEPTANCE=1 mistralrs bench -m -b 1 -b 128 2>&1 | tee /tmp/b.txt +``` + +**The one number:** decode tok/s in B as a fraction of A, at B=1 and B=128. The claim under +test is `B/A == 1.00` at both โ€” B must be byte-identical to A, because the refusal in ยง5 +means B takes the cohort path. **What B additionally proves is that the refusal fires and +names `XsRolling`**: `grep "cannot honour it" /tmp/b.txt`. If that line is absent the probe +is wrong; if `B/A != 1.00` something in this diff touched the live path and must be found. + +## 7. Gates + +`cargo check -p mistralrs-core` green ยท `cargo test -p mistralrs-core --lib` **377 passed, +0 failed** (was 366) ยท scoped clippy lane green ยท **zero rustfmt drift**, checked +like-for-like on `/tmp` copies against the pre-change baseline (all three files 0 โ†’ 0). + +### Mutation runs โ€” no test here is vacuous + +``` +plan_per_sequence_step: committed = (u+a).min(w).min(batch_min) + both_plans_agree_on_what_each_sequence_can_vouch_for ............ FAILED + per_sequence_plan_ignores_every_other_sequence .................. FAILED +plan_per_sequence_step: next_uncached follows the cohort rule + per_sequence_advance_is_token_identical_to_the_b1_reference ..... FAILED + per_sequence_advance_never_lets_a_tail_ratchet_at_b128 .......... FAILED + tok_per_step_at_b128_โ€ฆ ......................................... FAILED + per_sequence_plan_ignores_every_other_sequence .................. FAILED +front_pad_single: slice_set at offset 0 instead of `lead` + front_pad_moves_the_live_run_to_the_end_and_zeroes_the_prefix ... FAILED +make_left_padded_causal_mask: drop the `j < lead` condition + left_padded_mask_kills_the_dead_prefix_and_the_future_โ€ฆ ......... FAILED + a_single_query_row_still_gets_its_ragged_mask .................. FAILED +cache_supports_per_sequence_advance: never refuse + the_capability_probe_names_the_slot_that_cannot_carry_its_own_length FAILED +``` + +The no-ratchet test carries its own control in the same test: the **cohort** rule on the +identical fixture must still ratchet tails to the window, or the premise is gone and the +whole workstream is unnecessary. + +## 8. Surfaced, not shipped + +1. **`XsRollingCache` per-sequence state is the keystone** (ยง5.1), and the per-row append it + needs is precisely a `slot_mapping` โ€” the primitive every paged engine has and Arc's + dense path does not. PR #90's gather kernel already accepts "a row is any + (block-table, length) run", so the read side exists; the write side does not. +2. **The ragged-mask wire** (ยง5.2) is small and independent of (1) โ€” it would let any plain + dense model take the per-sequence path, if one ever gets an MTP head. +3. **`MtpAcceptance::from_fused_verify(d, n_acc)` records `proposed = d`**, so a saturated + sequence contributes a 0-proposed step and keeps `accept_rate` flattering. Carried over + from wave59-CJ ยง6.3, still worth an explicit `drafting_sequences / batch` counter. diff --git a/mistralrs-core/src/kv_cache/mod.rs b/mistralrs-core/src/kv_cache/mod.rs index 1acf7b604..86888c23a 100644 --- a/mistralrs-core/src/kv_cache/mod.rs +++ b/mistralrs-core/src/kv_cache/mod.rs @@ -94,6 +94,44 @@ impl KvCache { )?))) } + /// The variant's name, for diagnostics that have to say which slot refused. + pub fn kind_name(&self) -> &'static str { + match self { + Self::Normal { .. } => "Normal", + Self::Rotating { .. } => "Rotating", + Self::TurboQuant(_) => "TurboQuant", + Self::XsRolling(_) => "XsRolling", + } + } + + /// Can this slot hold a length that differs from the rest of its batch? + /// + /// One dense batched buffer writes every sequence's new rows at ONE offset + /// (`SingleCache::append`), so a per-sequence length is representable only + /// if a ragged cohort can be re-aligned so that every row's live content + /// *ends* at the same column โ€” [`front_pad_kv_cache`]. Each variant: + /// + /// * **`Normal` โ€” yes.** Its content is a flat run of `current_seq_len` + /// rows on one dim with no other state, so shifting the run right and + /// zero-filling ahead of it is exact. + /// * **`Rotating` โ€” no.** A ring buffer whose `offset` names where the + /// window wrapped; shifting the data would leave `offset` describing the + /// wrong rows, and the wrap point is not a function of the length. + /// * **`TurboQuant` โ€” no.** Content is quantized blocks with their own + /// grouping; a row shift is not a block shift. + /// * **`XsRolling` โ€” no, and this is the one that blocks DeepSeek V4.** It + /// carries *two* time bases (completed compressed rows on `comp`, raw + /// rows for `[base, tokens)` on `tail`) plus scalar `tokens`/`base` that + /// are one number for the whole batch. Per-sequence advance needs those + /// to be per-row, needs `advance` to append a **different number of rows + /// per sequence** โ€” which one shared append offset cannot express โ€” and + /// needs the compressor to fire per row. That is the same gap + /// `wave61-CL` ยง6 names ("the compressor's `xs` history has no block + /// table") and `wave29-BC` ยง4b ruled its own project. + pub fn supports_per_sequence_len(&self) -> bool { + matches!(self, Self::Normal { .. }) + } + pub fn k(&self) -> Result> { match self { Self::Normal { k, .. } => k.current_data(), @@ -711,6 +749,120 @@ impl BatchSrc { } } +/// Shift one `SingleCache`'s live run right so it **ends** at `target_len`, +/// zero-filling the `target_len - current_seq_len` columns ahead of it. +/// Returns the width of that dead prefix (`lead_pad`). +/// +/// This is the whole mechanism by which a dense batched cache can carry +/// per-sequence lengths. `SingleCache::append` writes at one shared offset +/// (`single_cache.rs:225`), so a ragged cohort can only share a forward if +/// every row's live content ends at the same column โ€” then that one offset is +/// simultaneously correct for all of them. Left-alignment buys that; the price +/// is `lead_pad` dead columns per row, which the caller MUST mask +/// ([`crate::layers_masker::RaggedKvLens`]) because a zero K row is not a +/// masked row โ€” it scores logit 0 and takes softmax weight. +/// +/// 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. +fn front_pad_single(sc: &mut SingleCache, target_len: usize) -> Result { + let live = sc.current_seq_len; + if target_len < live { + candle_core::bail!( + "kv-cache: front_pad to {target_len} is shorter than the {live} live positions it \ + would have to keep" + ); + } + let lead = target_len - live; + if lead == 0 { + sc.current_seq_len = target_len; + return Ok(0); + } + let Some(ad) = sc.all_data.as_ref() else { + candle_core::bail!("kv-cache: front_pad on a slot whose buffer is not materialised"); + }; + let capacity = sc.capacity_seq_len.max(target_len); + if capacity > sc.max_seq_len { + candle_core::bail!( + "kv-cache: front_pad to {target_len} needs capacity {capacity}, past this slot's \ + {} maximum", + sc.max_seq_len + ); + } + let mut shape = ad.dims().to_vec(); + shape[sc.dim] = capacity; + let grown = Tensor::zeros(shape, ad.dtype(), ad.device())?; + if live > 0 { + let src = ad.narrow(sc.dim, 0, live)?.contiguous()?; + grown.slice_set(&src, sc.dim, lead)?; + } + sc.all_data = Some(grown); + sc.capacity_seq_len = capacity; + sc.current_seq_len = target_len; + Ok(lead) +} + +/// [`front_pad_single`] over both halves of a slot. Refuses any variant that +/// cannot carry its own length ([`KvCache::supports_per_sequence_len`]). +pub(crate) fn front_pad_kv_cache(cache: &mut KvCache, target_len: usize) -> Result { + let KvCache::Normal { k, v } = cache else { + candle_core::bail!( + "kv-cache: front_pad is only defined for a `Normal` slot; this one is a `{}`. See \ + `KvCache::supports_per_sequence_len`.", + cache.kind_name() + ); + }; + let lead_k = front_pad_single(k, target_len)?; + let lead_v = front_pad_single(v, target_len)?; + debug_assert_eq!( + lead_k, lead_v, + "the two halves of one slot describe the same positions" + ); + Ok(lead_k) +} + +/// 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. +/// +/// After this call every sequence reports the same `current_seq_len`, so +/// [`ensure_uniform_batch_cache_lens`] passes and `NormalCacheManager` stacks +/// them unchanged โ€” the dense batching code needs no modification at all. +/// What the caller gains is that the length it hands back on the way *out* +/// (`clone_out_cache` โ†’ `KvCache::set_len`) may differ per sequence. +/// +/// Returns `lead_pad[i]` per sequence, which is `padded_len - live_len_i`, and +/// which the model's attention MUST mask. +pub(crate) fn front_align_batch( + seqs: &mut [&mut crate::sequence::Sequence], + modify_draft_cache: bool, +) -> Result> { + let mut target = 0usize; + for seq in seqs.iter_mut() { + let cache = if modify_draft_cache { + seq.normal_draft_cache() + } else { + seq.normal_cache() + }; + for slot in cache.iter().flatten() { + target = target.max(slot.current_seq_len()); + } + } + let mut lead_pads = Vec::with_capacity(seqs.len()); + for seq in seqs.iter_mut() { + let cache = if modify_draft_cache { + seq.normal_draft_cache() + } else { + seq.normal_cache() + }; + let mut lead = 0usize; + for slot in cache.iter_mut().flatten() { + lead = front_pad_kv_cache(slot, target)?; + } + lead_pads.push(lead); + } + Ok(lead_pads) +} + /// Zero-extend `src` along `dim` to `width`. Used only for preallocation slack /// (see [`BatchSrc`]), so the added columns are never read. fn pad_slack(src: &Tensor, dim: usize, width: usize) -> Result { @@ -1943,6 +2095,7 @@ impl CacheManager for HybridCa #[cfg(test)] mod clone_in_cache_invariant_tests { use super::*; + use candle_core::Device; use crate::sampler::Sampler; use crate::sequence::{SeqStepType, SequenceGroup, SequenceRecognizer}; use tokio::sync::Mutex as TokioMutex; @@ -2031,6 +2184,140 @@ mod clone_in_cache_invariant_tests { seq } + /// A `Normal` slot with a REAL buffer, so `front_pad` has tensors to move. + /// `all_data` is preallocated to `capacity` and the live run is the first + /// `live` rows, which is exactly how production builds it + /// (`kv_cache/mod.rs` preallocated cache path) โ€” `SingleCache::append` + /// writes into a capacity buffer, it does not grow a narrow one. + fn materialised_slot(live: usize, capacity: usize, mark: f32) -> KvCache { + let half = |mark: f32| { + let data: Vec = (0..capacity).map(|i| mark + i as f32).collect(); + let all = Tensor::from_vec(data, (1, 1, capacity, 1), &Device::Cpu).unwrap(); + SingleCache { + all_data: Some(all), + dim: 2, + current_seq_len: live, + capacity_seq_len: capacity, + max_seq_len: 4096, + } + }; + KvCache::Normal { + k: half(mark), + v: half(mark + 1000.0), + } + } + + fn rows(slot: &KvCache) -> Vec { + let KvCache::Normal { k, .. } = slot else { + unreachable!() + }; + k.all_data + .as_ref() + .unwrap() + .flatten_all() + .unwrap() + .to_vec1::() + .unwrap() + } + + /// ๐Ÿ”‘ The mechanism that lets one dense batched buffer carry per-sequence + /// lengths: the live run is shifted so it ENDS at the target, which makes + /// `SingleCache::append`'s one shared write offset correct for every row at + /// once. Content must survive the move exactly, and the prefix must be + /// zero โ€” a nonzero prefix would be attended as real keys. + #[test] + fn front_pad_moves_the_live_run_to_the_end_and_zeroes_the_prefix() { + 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, "a 5-row run inside a 9-wide target has a 4-row prefix"); + assert_eq!(slot.current_seq_len(), 9); + let after = rows(&slot); + assert!( + after[..4].iter().all(|x| *x == 0.0), + "the dead prefix must be zeroed, not left holding another position's K/V" + ); + assert_eq!( + after[4..9], + before[..], + "the live run must survive the shift bit-for-bit" + ); + } + + /// A no-op front pad must not reallocate or perturb anything โ€” this is the + /// B=1 path, and the uniform-batch path, and it must stay free. + #[test] + fn front_pad_to_the_current_length_is_a_no_op() { + let mut slot = materialised_slot(7, 16, 3.0); + let before = rows(&slot); + assert_eq!(front_pad_kv_cache(&mut slot, 7).unwrap(), 0); + assert_eq!(rows(&slot), before, "a zero-width pad must not touch data"); + assert_eq!(slot.current_seq_len(), 7); + } + + /// A ragged cohort becomes uniform in `current_seq_len` โ€” so + /// `ensure_uniform_batch_cache_lens` passes and `NormalCacheManager` stacks + /// it with no changes at all โ€” while each sequence's dead prefix is + /// reported so the caller can mask it. + #[test] + fn front_align_batch_makes_a_ragged_cohort_uniform_and_reports_each_dead_prefix() { + let mut a = seq_with_cache_len(0, 2, 1); + let mut b = seq_with_cache_len(1, 2, 1); + let mut c = seq_with_cache_len(2, 2, 1); + for (seq, live) in [(&mut a, 5usize), (&mut b, 3), (&mut c, 8)] { + let cache = seq.normal_cache(); + cache.clear(); + for l in 0..2 { + cache.push(Some(materialised_slot(live, 16, 10.0 * (l + 1) as f32))); + } + } + let mut seqs: Vec<&mut Sequence> = vec![&mut a, &mut b, &mut c]; + + assert!( + first_mismatched_cache_len(&mut seqs, false).is_some(), + "the fixture must start ragged or this test proves nothing" + ); + let lead = front_align_batch(&mut seqs, false).unwrap(); + assert_eq!( + lead, + vec![3, 5, 0], + "each sequence's dead prefix is `max_len - its own live length`" + ); + assert_eq!( + first_mismatched_cache_len(&mut seqs, false), + None, + "after alignment the dense batcher's own precondition holds again" + ); + for seq in seqs.iter_mut() { + for slot in seq.normal_cache().iter().flatten() { + assert_eq!(slot.current_seq_len(), 8); + } + } + } + + /// 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. + #[test] + fn front_pad_refuses_a_slot_that_cannot_carry_its_own_length() { + let mut xs = KvCache::XsRolling(Box::new(XsRollingCache::new(4, 2, 64, 2048))); + assert!(!xs.supports_per_sequence_len()); + let err = front_pad_kv_cache(&mut xs, 8).unwrap_err().to_string(); + assert!( + err.contains("XsRolling") && err.contains("supports_per_sequence_len"), + "the refusal must name the slot kind and where the rule lives; got {err}" + ); + } + + /// Front padding only ever grows. Shortening is `set_len`'s job and has its + /// own retention rules; conflating the two would let a caller drop rows the + /// compressor still needs without going through `try_set_len`. + #[test] + fn front_pad_refuses_to_shorten() { + let mut slot = materialised_slot(9, 16, 1.0); + assert!(front_pad_kv_cache(&mut slot, 4).is_err()); + } + #[test] fn uniform_cache_lens_are_accepted() { let mut a = seq_with_cache_len(0, 4, 100); diff --git a/mistralrs-core/src/layers_masker.rs b/mistralrs-core/src/layers_masker.rs index d7e2acd1c..998ccf3fc 100644 --- a/mistralrs-core/src/layers_masker.rs +++ b/mistralrs-core/src/layers_masker.rs @@ -24,6 +24,46 @@ pub struct NotACache; pub trait PastKvLenCache { fn get_past_kv_len(&self) -> Result; + + /// The number of **live** past positions each sequence in the batch holds, + /// when the dense batched cache is left-aligned and therefore ragged. + /// + /// `None` โ€” the default, and what every caller in the tree returns today โ€” + /// means "every row is live for the whole `get_past_kv_len()`", i.e. the + /// mask is batch-invariant and [`CausalMasker::make_causal_mask_matrix`] + /// builds exactly the rank-2 mask it always has. + /// + /// `Some(lens)` means row `b`'s live run is the **suffix** + /// `[past - lens[b], past)`; the `past - lens[b]` columns ahead of it are + /// the zero-filled dead prefix [`crate::kv_cache::front_pad_kv_cache`] + /// leaves behind. Those columns are not harmless: a zero K row scores logit + /// 0 and takes real softmax weight, so they must be masked, and this is how + /// the mask learns about them. + fn per_seq_kv_lens(&self) -> Option<&[usize]> { + None + } +} + +/// A left-aligned ragged batch's mask inputs: the padded width every row shares +/// and the live length each row actually holds. +/// +/// Pass this as the `cache` argument to +/// [`CausalMasker::make_causal_mask_matrix`] to get a `[B, 1, t_q, k]` additive +/// mask that kills each row's dead prefix as well as the future. +pub struct RaggedKvLens<'a> { + /// Columns every row of the batched buffer carries. + pub padded_len: usize, + /// `live[b]` โ€” real past positions in row `b`, `<= padded_len`. + pub live: &'a [usize], +} + +impl PastKvLenCache for RaggedKvLens<'_> { + fn get_past_kv_len(&self) -> Result { + Ok(self.padded_len) + } + fn per_seq_kv_lens(&self) -> Option<&[usize]> { + Some(self.live) + } } impl PastKvLenCache for NotACache { @@ -169,6 +209,53 @@ impl CausalMasker { Ok(k_cache_1.dims()[2]) } + /// The additive `[B, 1, t_q, past + t_q]` mask a **left-aligned ragged** + /// batch needs: `-inf` on the future (ordinary causality) and `-inf` on + /// each row's dead prefix `[0, past - live[b])`. + /// + /// Query row `i` of sequence `b` sits at absolute position + /// `live[b] + i`, and its live keys are the columns + /// `[past - live[b], past + i]`. Everything outside that is killed. Note + /// the row is never fully masked โ€” column `past + i` is the query's own + /// position โ€” so softmax cannot produce a NaN row. + pub fn make_left_padded_causal_mask( + &self, + b_sz: usize, + tgt_len: usize, + past_kv_len: usize, + live: &[usize], + dtype: DType, + device: &Device, + ) -> Result { + if live.len() != b_sz { + candle_core::bail!( + "ragged mask: {} live lengths for a batch of {b_sz}", + live.len() + ); + } + let k_len = past_kv_len + tgt_len; + let mut data: Vec = Vec::with_capacity(b_sz * tgt_len * k_len); + for &l in live { + if l > past_kv_len { + candle_core::bail!( + "ragged mask: live length {l} exceeds the padded width {past_kv_len}" + ); + } + let lead = past_kv_len - l; + for i in 0..tgt_len { + let last = past_kv_len + i; + for j in 0..k_len { + data.push(if j < lead || j > last { + f32::NEG_INFINITY + } else { + 0.0 + }); + } + } + } + Tensor::from_vec(data, (b_sz, 1, tgt_len, k_len), device)?.to_dtype(dtype) + } + pub fn make_causal_mask_matrix( &self, input_ids: &Tensor, @@ -177,7 +264,24 @@ impl CausalMasker { _n_attn_heads: usize, ) -> Result> { let past_kv_len = cache.get_past_kv_len()?; - let (_b_sz, tgt_len) = input_ids.dims2()?; + let (b_sz, tgt_len) = input_ids.dims2()?; + + // ๐Ÿ”‘ A left-aligned ragged batch needs a mask even at `tgt_len == 1`, + // and it cannot take the flash-attn shortcut: both of those assume the + // mask is a pure function of causality, and the dead prefix + // `front_pad_kv_cache` leaves is neither causal nor batch-invariant. + // Order matters โ€” this branch has to come before both early returns. + if let Some(live) = cache.per_seq_kv_lens() { + return Ok(Some(self.make_left_padded_causal_mask( + b_sz, + tgt_len, + past_kv_len, + live, + dtype, + input_ids.device(), + )?)); + } + if tgt_len == 1 { return Ok(None); } @@ -420,3 +524,110 @@ impl BidirectionalMasker { Ok(mask) } } + +#[cfg(test)] +mod ragged_mask_tests { + use super::*; + + fn mask_rows(m: &Tensor) -> Vec> { + let (b, _, q, k) = m.dims4().unwrap(); + let flat = m.flatten_all().unwrap().to_vec1::().unwrap(); + (0..b * q) + .map(|r| flat[r * k..(r + 1) * k].to_vec()) + .collect() + } + + /// ๐Ÿ”‘ The mask a left-aligned per-sequence KV cache needs. Three things + /// have to hold at once, and each one is a wrong answer if it does not: + /// + /// * the dead prefix `front_pad_kv_cache` zero-fills is `-inf` โ€” a zero K + /// row is NOT a masked row, it scores logit 0 and takes softmax weight; + /// * the future is `-inf` โ€” ordinary causality, which the ragged branch + /// must not lose; + /// * no row is entirely `-inf` โ€” a fully masked row is a softmax NaN. + #[test] + fn left_padded_mask_kills_the_dead_prefix_and_the_future_and_never_a_whole_row() { + let device = Device::Cpu; + // Padded width 6; sequence 0 holds 6 live positions, sequence 1 holds + // 2, so row 1 carries a 4-wide dead prefix. Two query rows. + let live = [6usize, 2]; + let m = CausalMasker + .make_left_padded_causal_mask(2, 2, 6, &live, DType::F32, &device) + .unwrap(); + assert_eq!(m.dims(), &[2, 1, 2, 8]); + let rows = mask_rows(&m); + + // seq 0, query 0: keys 0..=6 live (own position is column 6), 7 future. + assert!(rows[0][..7].iter().all(|x| *x == 0.0)); + assert_eq!(rows[0][7], f32::NEG_INFINITY); + // seq 0, query 1: keys 0..=7 all live. + assert!(rows[1].iter().all(|x| *x == 0.0)); + // seq 1, query 0: columns 0..4 are the dead prefix. + assert!( + rows[2][..4].iter().all(|x| *x == f32::NEG_INFINITY), + "the dead prefix must be masked or the model attends zero-filled keys" + ); + assert!(rows[2][4..7].iter().all(|x| *x == 0.0)); + assert_eq!(rows[2][7], f32::NEG_INFINITY); + // seq 1, query 1. + assert!(rows[3][..4].iter().all(|x| *x == f32::NEG_INFINITY)); + assert!(rows[3][4..8].iter().all(|x| *x == 0.0)); + + for (i, row) in rows.iter().enumerate() { + assert!( + row.iter().any(|x| *x == 0.0), + "row {i} is entirely masked, which is a softmax NaN" + ); + } + } + + /// A batch whose rows are all fully live must produce exactly the ordinary + /// causal mask โ€” so turning the ragged path on cannot change a uniform + /// batch's answer. + #[test] + fn a_fully_live_ragged_mask_is_the_ordinary_causal_mask() { + let device = Device::Cpu; + let m = CausalMasker + .make_left_padded_causal_mask(2, 3, 4, &[4, 4], DType::F32, &device) + .unwrap(); + let rows = mask_rows(&m); + for (r, row) in rows.iter().enumerate() { + let last = 4 + (r % 3); + for (j, v) in row.iter().enumerate() { + let want = if j > last { f32::NEG_INFINITY } else { 0.0 }; + assert_eq!(*v, want, "row {r} column {j}"); + } + } + } + + /// `make_causal_mask_matrix` routes to the ragged builder BEFORE both of + /// its early returns โ€” the `tgt_len == 1` shortcut and the flash-attn + /// placeholder. A decode step is `tgt_len == 1` and still has a dead prefix + /// to mask, so returning `None` there would serve from unmasked keys. + #[test] + fn a_single_query_row_still_gets_its_ragged_mask() { + let device = Device::Cpu; + let ids = Tensor::zeros((2, 1), DType::U32, &device).unwrap(); + let ragged = RaggedKvLens { + padded_len: 5, + live: &[5, 1], + }; + let m = CausalMasker + .make_causal_mask_matrix(&ids, &ragged, DType::F32, 1) + .unwrap() + .expect("a ragged batch must always get a mask, even at tgt_len == 1"); + assert_eq!(m.dims(), &[2, 1, 1, 6]); + let rows = mask_rows(&m); + assert!(rows[0].iter().all(|x| *x == 0.0)); + assert!(rows[1][..4].iter().all(|x| *x == f32::NEG_INFINITY)); + assert!(rows[1][4..].iter().all(|x| *x == 0.0)); + + // And the default (`per_seq_kv_lens() == None`) path is untouched: one + // query row over a uniform batch still returns `None`. + let offsets: &[usize] = &[5, 5]; + assert!(CausalMasker + .make_causal_mask_matrix(&ids, &offsets, DType::F32, 1) + .unwrap() + .is_none()); + } +} diff --git a/mistralrs-core/src/pipeline/mtp_pipeline.rs b/mistralrs-core/src/pipeline/mtp_pipeline.rs index 1c6a36fe7..25f053983 100644 --- a/mistralrs-core/src/pipeline/mtp_pipeline.rs +++ b/mistralrs-core/src/pipeline/mtp_pipeline.rs @@ -1170,6 +1170,10 @@ pub struct MtpSpeculativePipeline { /// Latches so the "drafting skipped, draft KV unprimed" explanation is /// logged once per process rather than once per token. warned_unprimed: std::sync::atomic::AtomicBool, + /// How this pipeline advances the KV cache after a verify. Latched on the + /// first fast-path step (the cache's slot kinds are fixed by then) so the + /// decision and its reason are logged once, not once per token. + kv_advance: std::sync::OnceLock, } /// One sequence's persistent MTP draft KV. @@ -1239,6 +1243,7 @@ impl MtpSpeculativePipeline { draft_kv: std::sync::Mutex::new(std::collections::HashMap::new()), draft_kv_clock: AtomicUsize::new(0), warned_unprimed: std::sync::atomic::AtomicBool::new(false), + kv_advance: std::sync::OnceLock::new(), }) } @@ -1270,6 +1275,7 @@ impl MtpSpeculativePipeline { draft_kv: std::sync::Mutex::new(std::collections::HashMap::new()), draft_kv_clock: AtomicUsize::new(0), warned_unprimed: std::sync::atomic::AtomicBool::new(false), + kv_advance: std::sync::OnceLock::new(), } } @@ -1291,6 +1297,79 @@ impl MtpSpeculativePipeline { self.depth + 1 } + /// How this pipeline advances the KV cache after a verify, decided once. + /// + /// Default [`KvAdvance::Cohort`]. `ARC_MTP_PER_SEQ_KV=1` asks for + /// [`KvAdvance::PerSequence`]; it is granted only when **both** of these + /// hold, and the reason it was refused is logged by name so a run that + /// asked for it never silently gets the ceiling instead: + /// + /// 1. every cache slot can carry its own length + /// ([`cache_supports_per_sequence_advance`]); + /// 2. the target model masks a left-aligned ragged batch's dead prefix + /// ([`Self::target_masks_ragged_batches`]). + /// + /// (2) is not decoration. Per-sequence lengths in one dense buffer mean + /// `front_pad_kv_cache` zero-fills ahead of the shorter rows, and a zero K + /// row is **not** a masked row โ€” it scores logit 0 and takes real softmax + /// weight. Serving from those columns would be a wrong answer nothing + /// downstream catches, which is the exact failure this project already paid + /// for once with FP8 KV. Refusing is the only correct behaviour until a + /// model threads [`crate::layers_masker::RaggedKvLens`]. + pub(crate) fn kv_advance(&self) -> KvAdvance { + *self.kv_advance.get_or_init(|| { + if !per_sequence_kv_requested() { + return KvAdvance::Cohort; + } + let refusal = cache_supports_per_sequence_advance(&self.target_cache) + .err() + .or_else(|| self.target_masks_ragged_batches().err()); + match refusal { + None => { + tracing::info!( + target: "mtp_speculative", + "MTP: per-sequence KV advance is ON โ€” each sequence advances by its own \ + accepted count and no cohort rollback is applied" + ); + KvAdvance::PerSequence + } + Some(why) => { + tracing::warn!( + target: "mtp_speculative", + "MTP: ARC_MTP_PER_SEQ_KV=1 was requested but this pipeline cannot honour \ + it, so batched MTP keeps the cohort rollback and its ceiling ({why})" + ); + KvAdvance::Cohort + } + } + }) + } + + /// Whether the target model applies a per-row mask over a left-aligned + /// ragged batch's dead prefix. + /// + /// **No model in the tree does yet**, and this says so rather than + /// pretending. The plumbing that would flip it is exactly one wire: + /// `clone_in_cache` already computes `lead_pad` per sequence + /// ([`crate::kv_cache::front_align_batch`]) and + /// [`crate::layers_masker::CausalMasker::make_left_padded_causal_mask`] + /// already turns that into the `[B, 1, t_q, k]` additive mask a model wants + /// โ€” what is missing is carrying it from the pipeline through + /// `inputs_processor` into the model forward, where today the mask is + /// rebuilt from `seqlen_offsets` alone. + /// + /// It is stated as a capability rather than hardcoded at the call site so + /// that the model that adds the wire overrides one thing, here. + fn target_masks_ragged_batches(&self) -> std::result::Result<(), String> { + Err( + "no model in this tree threads a ragged-batch mask into its forward yet, so the \ + zero-filled dead prefix a left-aligned per-sequence cache leaves would be attended \ + as real keys (a zero K row scores logit 0 and takes softmax weight). See \ + `MtpSpeculativePipeline::target_masks_ragged_batches`." + .to_string(), + ) + } + /// Snapshot of this pipeline's MTP acceptance counters. pub fn acceptance(&self) -> MtpAcceptance { self.acceptance.snapshot() @@ -1686,6 +1765,163 @@ pub(crate) fn plan_batch_step(uncached: &[usize], accepted: &[usize], w: usize) } } +/// How a batched MTP step advances the KV cache. +/// +/// This is the whole of the batch problem, named. See +/// [`plan_per_sequence_step`] for why `Cohort` has a ceiling that no amount of +/// draft-quality work can lift. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum KvAdvance { + /// One dense cache, one length, rolled back to `min_i(u_i + a_i)`. + /// [`plan_batch_step`]. Correct, and bounded by [`plan_per_sequence_step`]'s + /// doc comment. + Cohort, + /// Every sequence advances by its own accepted count. What vLLM, SGLang and + /// TensorRT-LLM all do. [`plan_per_sequence_step`]. + PerSequence, +} + +/// What one fused batched MTP step does when each sequence advances its **own** +/// KV by its **own** accepted count. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct PerSeqStepPlan { + /// `committed[i]` โ€” how many of the `w` window slots sequence `i` keeps: + /// `u_i + a_i`. Its cache ends the step at `cache_len + committed[i]`. + pub committed: Vec, + /// Slots sequence `i` drops: `w - committed[i]`. Their K/V is never read + /// again โ€” it is overwritten by the next step's window, exactly as vLLM and + /// SGLang leave rejected draft slots stale inside a request's own + /// over-allocated region rather than freeing them. + pub dropped: Vec, + /// Each sequence's uncached committed tail at the START of the next step. + /// **Always 1**, for every sequence, at every batch size โ€” which is the + /// entire point, and is exactly the B=1 invariant. + pub next_uncached: Vec, +} + +/// Resolve one fused batched step **without a cohort barrier**. +/// +/// # Why this exists +/// +/// [`plan_batch_step`] rolls one dense cache back to `keep = min_i(u_i + a_i)`, +/// so `next_u_i = u_i + a_i + 1 - keep`. The moment any sequence rejects its +/// first draft (`u=1, a=0`) that pins `keep = 1` and collapses to +/// `next_u_i = u_i + a_i` โ€” monotone non-decreasing, saturating at the window +/// `w`, where the sequence drafts `w - u = 0` tokens. **The sequences that +/// accept best are the ones that stop drafting.** At B=1 it cannot happen +/// (`keep = u + a`, so `next_u โ‰ก 1`), which is why B=1 is the only regime +/// batched MTP ever measured well in. +/// +/// # What production engines do instead +/// +/// All three reference engines key the advance on a **per-request scalar** and +/// perform no reduction across the batch: +/// +/// * vLLM v1 โ€” `v1/core/sched/scheduler.py:1836-1848`, inside the per-`req_id` +/// loop of `update_from_output`: `num_rejected = num_draft_tokens - +/// num_accepted; request.num_computed_tokens -= num_rejected`. +/// * SGLang โ€” `speculative/eagle_worker_common.py:587-588` +/// `new_seq_lens = batch.seq_lens + accept_lens` (elementwise), committed at +/// `scheduler_components/batch_result_processor.py:686-687` +/// `req.kv_committed_len += num_accept_tokens`. +/// * TensorRT-LLM โ€” `pyexecutor/resource_manager.py:1330-1346` rewinds +/// `request.py_rewind_len` per request. +/// +/// A batch-wide `min` over accept lengths appears in none of them. The only +/// `min`s nearby are unrelated (vLLM's cascade-attention common-prefix pick) or +/// per-request clamps. +/// +/// # What happens to a rejected draft's K/V +/// +/// Nothing. vLLM and SGLang neither free nor zero it: the request's length +/// shrinks so attention stops reading those slots, and the next step's window +/// writes over the same physical positions (SGLang says so in +/// `eagle_worker_common.py:450` โ€” "trailing unaccepted slots stay and are freed +/// as overshoot"). TensorRT-LLM is the only one that reclaims, and only when a +/// rewind empties a whole block (`kvCacheManager.cpp:4562-4600`). `dropped[i]` +/// here is that same quantity: slots the sequence stops counting, not slots +/// anyone hands back. +/// +/// # Invariants +/// +/// * `next_uncached[i] == 1` for every `i` โ€” no ratchet is representable. +/// * `committed[i] == u_i + a_i >= 1`, so every sequence advances every step. +/// * `committed[i] + dropped[i] == w` โ€” no window slot is unaccounted for. +/// * **No value depends on any other sequence.** That is the property under +/// test; a cohort reduction reintroduced anywhere breaks +/// `per_sequence_advance_is_independent_of_every_other_sequence`. +pub(crate) fn plan_per_sequence_step( + uncached: &[usize], + accepted: &[usize], + w: usize, +) -> PerSeqStepPlan { + let committed: Vec = uncached + .iter() + .zip(accepted) + .map(|(u, a)| (u + a).min(w)) + .collect(); + PerSeqStepPlan { + dropped: committed.iter().map(|c| w.saturating_sub(*c)).collect(), + next_uncached: vec![1; committed.len()], + committed, + } +} + +/// `ARC_MTP_PER_SEQ_KV=1` opts a run into [`KvAdvance::PerSequence`]. +/// +/// Default OFF. The flag alone is not enough: the batch's cache must also be +/// able to carry a per-sequence length +/// ([`cache_supports_per_sequence_advance`]), and a run that asks for it and +/// cannot get it is told which slot refused rather than silently served the +/// cohort answer. +pub(crate) fn per_sequence_kv_requested() -> bool { + std::env::var("ARC_MTP_PER_SEQ_KV").is_ok_and(|v| v == "1" || v.eq_ignore_ascii_case("true")) +} + +/// Whether every slot of this cache can hold a length that differs from the +/// rest of its batch, or the name of the first slot that cannot. +/// +/// A dense batched forward writes every sequence's new K/V at ONE offset +/// (`SingleCache::append`, `kv_cache/single_cache.rs:225`), so per-sequence +/// lengths are only representable if a ragged cohort can be re-aligned into one +/// buffer whose rows all *end* at the same column โ€” +/// [`crate::kv_cache::front_pad_kv_cache`]. That is well defined for a plain +/// `Normal` slot and for nothing else in the tree today; see +/// [`crate::kv_cache::KvCache::supports_per_sequence_len`] for why each of the +/// other three declines. +pub(crate) fn cache_supports_per_sequence_advance( + cache: &EitherCache, +) -> std::result::Result<(), String> { + let EitherCache::Normal(normal) = cache else { + // Not a refusal of the *idea* โ€” the opposite. PagedAttention already + // has per-sequence independence at the allocator surface, and under + // `ARC_SEGMENTED_KV=1` `KVCacheManager::trim_request_to_num_tokens` + // routes straight to `SegmentedAllocator::rollback(pool, request_id, 0, + // num_tokens)` (`paged_attention/kv_cache_manager.rs:445-448`), which + // is per-request and takes an ABSOLUTE token count. What blocks paged + // MTP is upstream of here: the fast path only accepts + // `CacheBackendMetadata::DefaultInstructions`, so a paged batch never + // reaches this rollback at all. That is its own change. + return Err( + "the target is on PagedAttention, whose per-sequence block tables already advance \ + independently โ€” this dense-cache path is not what it needs" + .to_string(), + ); + }; + let guard = normal.lock().unwrap(); + for (i, slot) in guard.0.iter().enumerate() { + if !slot.supports_per_sequence_len() { + return Err(format!( + "cache slot {i} is a `{}`, which carries one length for the whole batch. \ + Per-sequence MTP advance needs every slot to hold its own length; see \ + `KvCache::supports_per_sequence_len`.", + slot.kind_name() + )); + } + } + Ok(()) +} + /// Greedy argmax over a 1-D logits tensor โ€” returns the token ID. fn argmax_token(logits: &Tensor) -> Result { // logits shape [vocab] @@ -1966,7 +2202,22 @@ impl IsqPipelineMixin for MtpSpeculativePipeline { } impl CacheManagerMixin for MtpSpeculativePipeline { + /// Under [`KvAdvance::PerSequence`] a cohort arrives here genuinely ragged + /// โ€” that is the point โ€” 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 code below is then completely + /// 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. 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) } fn clone_out_cache(&self, seqs: &mut [&mut Sequence]) { @@ -2508,16 +2759,45 @@ impl Pipeline for MtpSpeculativePipeline { // step, which is exactly what `w - u` drafts leaves room for. `u >= 1` // for every sequence, so the cache always advances by at least one // position and the batch can never stall. + // + // Under [`KvAdvance::PerSequence`] none of the paragraph above applies: + // there is no minimum, no rollback of the shared buffer, and no + // surplus. Each sequence's own cache is set to `cache_len + u_i + a_i` + // 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(); - let plan = plan_batch_step(&uncached, &n_accepted, w); - debug_assert_eq!( - plan.keep, - valid_extent.iter().copied().min().unwrap_or(0), - "plan_batch_step must agree with the per-sequence valid extents" - ); - if plan.n_drop > 0 { - truncate_normal_cache(self, plan.n_drop)?; - } + // 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 per_sequence_commit = match advance { + KvAdvance::Cohort => { + let plan = plan_batch_step(&uncached, &n_accepted, w); + debug_assert_eq!( + plan.keep, + valid_extent.iter().copied().min().unwrap_or(0), + "plan_batch_step must agree with the per-sequence valid extents" + ); + if plan.n_drop > 0 { + truncate_normal_cache(self, plan.n_drop)?; + } + None + } + KvAdvance::PerSequence => { + let plan = plan_per_sequence_step(&uncached, &n_accepted, w); + debug_assert_eq!( + plan.committed, valid_extent, + "plan_per_sequence_step must agree with the per-sequence valid extents" + ); + Some(plan.committed) + } + }; // ---- Commit ---- for i in 0..batch { @@ -2613,6 +2893,42 @@ impl Pipeline for MtpSpeculativePipeline { } handle_post_cache_op(self, input_seqs, post_op); + // ---- Per-sequence advance ---- + // + // `handle_post_cache_op` has just split the batched buffer back into + // each sequence's own `normal_cache`, every row still `cache_len + w` + // long. Each sequence now keeps exactly the `cache_len + u_i + a_i` it + // can vouch for. The slots past that hold a rejected draft's K/V and + // are simply stopped being counted โ€” vLLM + // (`v1/core/sched/scheduler.py:1836-1848`) and SGLang + // (`speculative/eagle_worker_common.py:450`) do the same thing and + // likewise never free or zero them; the next step's window overwrites + // the same physical positions. + // + // ๐Ÿ”‘ The per-sequence state IS `Sequence::normal_cache`'s own + // `current_seq_len`. There is deliberately **no second map keyed by + // sequence id** anywhere in this pipeline: a shadow copy of a length + // that the cache also holds is exactly the divergence wave59-CJ ยง1 + // documents, reappearing in a new place. The commit therefore happens + // 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. + if let Some(committed) = per_sequence_commit { + for (i, c) in committed.iter().enumerate() { + let target = cache_len + c; + for slot in input_seqs[i].normal_cache().iter_mut().flatten() { + let kind = slot.kind_name(); + if let Err(e) = slot.set_len(target) { + candle_core::bail!( + "MTP per-sequence advance: set_len({target}) refused on a {kind} slot \ + for batch row {i} ({e}). Continuing would leave that sequence \ + reading a rejected draft's K/V." + ); + } + } + } + } + Ok(start.elapsed()) } @@ -3798,6 +4114,37 @@ mod tests { w: usize, plan_override: Option BatchStepPlan>, ) -> Vec> { + let (uncached, n_accepted, emitted) = sim_propose_and_verify(seqs, w); + let plan = plan_override.unwrap_or(plan_batch_step)(&uncached, &n_accepted, w); + for (i, s) in seqs.iter_mut().enumerate() { + s.toks.extend_from_slice(&emitted[i]); + s.uncached = (s.uncached + n_accepted[i] + 1).saturating_sub(plan.keep); + } + emitted + } + + /// The same step under [`KvAdvance::PerSequence`]. **Byte-for-byte the same + /// drafting and the same verify** โ€” `sim_propose_and_verify` is shared โ€” so + /// the only thing that differs between this and [`sim_step`] is how the KV + /// cache advances. Any divergence in the tests below is therefore + /// attributable to the advance rule and to nothing else. + fn sim_step_per_sequence(seqs: &mut [SimSeq], w: usize) -> Vec> { + let (uncached, n_accepted, emitted) = sim_propose_and_verify(seqs, w); + let plan = plan_per_sequence_step(&uncached, &n_accepted, w); + for (i, s) in seqs.iter_mut().enumerate() { + s.toks.extend_from_slice(&emitted[i]); + s.uncached = plan.next_uncached[i]; + } + emitted + } + + /// Draft `w - u` tokens per sequence, feed the `w`-wide window, verify with + /// `verify_proposed` against the rows `window_verify_row` names. Returns + /// `(uncached, n_accepted, emitted)`; it does **not** touch the cache. + fn sim_propose_and_verify( + seqs: &[SimSeq], + w: usize, + ) -> (Vec, Vec, Vec>) { let uncached: Vec = seqs.iter().map(|s| s.uncached).collect(); let mut drafts: Vec> = Vec::with_capacity(seqs.len()); for s in seqs.iter() { @@ -3839,12 +4186,299 @@ mod tests { emitted.push(out); } - let plan = plan_override.unwrap_or(plan_batch_step)(&uncached, &n_accepted, w); - for (i, s) in seqs.iter_mut().enumerate() { - s.toks.extend_from_slice(&emitted[i]); - s.uncached = (s.uncached + n_accepted[i] + 1).saturating_sub(plan.keep); + (uncached, n_accepted, emitted) + } + + // ===================================================================== + // Per-sequence KV advance โ€” the production rule (wave62) + // ===================================================================== + // + // `plan_batch_step` is correct and has a ceiling; `plan_per_sequence_step` + // is what vLLM, SGLang and TensorRT-LLM do. Every test below drives BOTH + // through the same `sim_propose_and_verify`, so the drafting, the verify + // and the accept rule are literally the same code and the only variable is + // the cache advance. + + /// A B=128 cohort with content-dependent accept lengths โ€” the fixture the + /// whole argument is about. + fn batch_of(n: usize) -> Vec { + (0..n) + .map(|i| SimSeq { + // Distinct seeds so the sequences genuinely diverge; two equal + // seeds would make ragged accept unreachable and the test + // vacuous. + toks: vec![11 + (i as u32 % 61), 5 + (i as u32 % 37)], + // 2..=5: a spread of per-sequence acceptance, which is what + // makes `min_i` bite. + agree_mod: 2 + (i as u32 % 4), + uncached: 1, + }) + .collect() + } + + /// The capability probe is what stands between "asked for per-sequence + /// advance" and "silently served the cohort ceiling". A V4-shaped cache โ€” + /// K/V slots plus the compressor's `XsRolling` histories โ€” must decline, + /// and must name the slot, because that slot is the actual remaining work. + #[test] + fn the_capability_probe_names_the_slot_that_cannot_carry_its_own_length() { + use crate::kv_cache::{EitherCache, KvCache, NormalCache, XsRollingCache}; + use std::sync::{Arc, Mutex}; + + let all_normal = EitherCache::Normal(Arc::new(Mutex::new(NormalCache(vec![ + KvCache::new_normal(2, 4096, 512), + KvCache::new_normal(2, 4096, 512), + ])))); + assert!( + cache_supports_per_sequence_advance(&all_normal).is_ok(), + "a cache of plain Normal slots can carry per-sequence lengths" + ); + + // V4's shape: K/V slots interleaved with compressor histories. + let v4_shaped = EitherCache::Normal(Arc::new(Mutex::new(NormalCache(vec![ + KvCache::new_normal(2, 4096, 512), + KvCache::XsRolling(Box::new(XsRollingCache::new(4, 2, 64, 2048))), + ])))); + let err = cache_supports_per_sequence_advance(&v4_shaped).unwrap_err(); + assert!( + err.contains("slot 1") && err.contains("XsRolling"), + "the refusal must name which slot and which kind; got {err}" + ); + } + + /// The plan functions agree on the one thing they must: the slots a + /// sequence keeps. They differ only in what happens to the surplus โ€” the + /// cohort rule rolls it back and re-feeds it as a longer tail next step, + /// the per-sequence rule keeps it and resets the tail to 1. + #[test] + fn both_plans_agree_on_what_each_sequence_can_vouch_for() { + let w = 5usize; + let uncached = [1usize, 2, 1, 4]; + let accepted = [3usize, 1, 0, 0]; + let per_seq = plan_per_sequence_step(&uncached, &accepted, w); + let cohort = plan_batch_step(&uncached, &accepted, w); + for i in 0..uncached.len() { + assert_eq!( + per_seq.committed[i], + uncached[i] + accepted[i], + "seq {i}: a sequence vouches for its own tail plus its own accepts" + ); } - emitted + assert_eq!( + cohort.keep, + per_seq.committed.iter().copied().min().unwrap(), + "the cohort keep is exactly the batch minimum of the per-sequence commits โ€” which \ + is the barrier, stated as an equation" + ); + } + + /// ๐Ÿ”‘ **The whole point.** Under the per-sequence rule no sequence's + /// uncached tail ever leaves 1, so every sequence drafts the full depth on + /// every step, at every batch size. + /// + /// The control in the same test shows the cohort rule does the opposite on + /// the identical fixture: tails ratchet up to the window and the sequences + /// sitting there draft **zero**. That control is what makes this test + /// non-vacuous โ€” if someone reintroduces a cohort reduction inside + /// `plan_per_sequence_step` the first assertion fails, and if the cohort + /// rule ever stops ratcheting the second one fails and this whole + /// workstream is unnecessary. + #[test] + fn per_sequence_advance_never_lets_a_tail_ratchet_at_b128() { + let depth = 3usize; + let w = depth + 1; + let steps = 200usize; + + let mut seqs = batch_of(128); + let mut min_drafted = usize::MAX; + for step in 0..steps { + for (i, s) in seqs.iter().enumerate() { + assert_eq!( + s.uncached, 1, + "step {step}, seq {i}: per-sequence advance must leave every uncached tail \ + at 1 โ€” a tail above 1 is the ratchet, and a tail at w={w} drafts nothing" + ); + min_drafted = min_drafted.min(w - s.uncached); + } + sim_step_per_sequence(&mut seqs, w); + } + assert_eq!( + min_drafted, depth, + "every sequence must draft the full depth on every step under the per-sequence rule" + ); + + // Control: the cohort rule on the identical fixture. + let mut cohort = batch_of(128); + let mut saturated = 0usize; + let mut zero_drafters = 0usize; + for _ in 0..steps { + sim_step(&mut cohort, w, None); + } + for s in &cohort { + if s.uncached == w { + saturated += 1; + } + if w - s.uncached == 0 { + zero_drafters += 1; + } + } + assert!( + saturated > 0, + "the cohort rule is supposed to ratchet tails to the window at B=128 โ€” if it no \ + longer does, this whole change is unnecessary and the premise must be re-derived" + ); + assert_eq!( + saturated, zero_drafters, + "a saturated tail is exactly a sequence that drafts nothing" + ); + } + + /// Token identity: batched per-sequence decoding must produce, for every + /// sequence, exactly the tokens that sequence would have produced alone. + /// + /// This is the acceptance test for the whole design. `sim_step_per_sequence` + /// at B=128 vs the same sequence run at B=1 โ€” no tolerance, exact equality + /// on the full stream. + #[test] + fn per_sequence_advance_is_token_identical_to_the_b1_reference() { + let w = 4usize; + let steps = 60usize; + + let mut batched = batch_of(128); + for _ in 0..steps { + sim_step_per_sequence(&mut batched, w); + } + + for (i, want) in batch_of(128).into_iter().enumerate() { + let mut solo = vec![want]; + for _ in 0..steps { + sim_step_per_sequence(&mut solo, w); + } + assert_eq!( + batched[i].toks, solo[0].toks, + "sequence {i} at B=128 diverged from its own B=1 run โ€” batched per-sequence \ + decoding is not user-invisible" + ); + } + } + + /// B=1 must not regress: at a batch of one the two rules are the same + /// function, so the flag cannot move the only regime that has ever been + /// measured. + #[test] + fn the_two_advance_rules_are_identical_at_batch_one() { + for depth in 1..=8usize { + let w = depth + 1; + let mut cohort = vec![batch_of(1).remove(0)]; + let mut per_seq = vec![batch_of(1).remove(0)]; + for step in 0..40 { + let a = sim_step(&mut cohort, w, None); + let b = sim_step_per_sequence(&mut per_seq, w); + assert_eq!(a, b, "depth {depth}, step {step}: B=1 emission differs"); + assert_eq!( + cohort[0].uncached, per_seq[0].uncached, + "depth {depth}, step {step}: B=1 uncached tail differs" + ); + } + } + } + + /// No value the per-sequence plan produces for sequence `i` may depend on + /// any other sequence. This is the anti-cohort-barrier test: it fails the + /// moment anyone reintroduces a `min`, a `max` or any other reduction over + /// the batch. + #[test] + fn per_sequence_plan_ignores_every_other_sequence() { + let w = 5usize; + // A deliberately hostile batch: one permanent laggard (u=1, a=0) โ€” + // exactly the sequence that pins `keep = 1` and ratchets everyone else + // under the cohort rule. + let uncached = [1usize, 1, 2, 3, 5, 1, 4]; + let accepted = [0usize, 4, 3, 2, 0, 1, 0]; + let batched = plan_per_sequence_step(&uncached, &accepted, w); + + for i in 0..uncached.len() { + let solo = plan_per_sequence_step(&uncached[i..=i], &accepted[i..=i], w); + assert_eq!( + batched.committed[i], solo.committed[0], + "seq {i}: committed slots changed because of its neighbours" + ); + assert_eq!( + batched.dropped[i], solo.dropped[0], + "seq {i}: dropped slots changed because of its neighbours" + ); + assert_eq!(batched.next_uncached[i], 1, "seq {i}: tail must reset to 1"); + assert_eq!( + batched.committed[i] + batched.dropped[i], + w, + "seq {i}: window slots unaccounted for" + ); + assert!(batched.committed[i] >= 1, "seq {i}: cache must advance"); + } + + // Teeth: the cohort rule on the SAME input does depend on neighbours, + // so the equalities above are not trivially true of any plan. + let cohort = plan_batch_step(&uncached, &accepted, w); + assert_eq!(cohort.keep, 1, "the laggard pins the cohort keep at 1"); + assert!( + cohort.next_uncached.iter().any(|u| *u > 1), + "the cohort rule must ratchet on this fixture or the contrast is vacuous" + ); + } + + /// The quantified gap, in the unit that matters. `tok_per_step` โ€” NOT + /// `accept_rate`: a saturated sequence drafts 0, so it contributes + /// `proposed = 0` and leaves `accept_rate` flattering while `tok_per_step` + /// collapses. The two diverge exactly here. + /// + /// This is arithmetic on the production plan functions, not a hardware + /// measurement, and must never be quoted as one. + #[test] + fn tok_per_step_at_b128_is_bounded_by_the_cohort_rule_and_freed_by_the_per_sequence_rule() { + let w = 4usize; // depth 3, the shipped default + let steps = 200usize; + let b = 128usize; + + let tok_per_step = |mut seqs: Vec, per_seq: bool| -> f64 { + let mut emitted = 0usize; + for _ in 0..steps { + let out = if per_seq { + sim_step_per_sequence(&mut seqs, w) + } else { + sim_step(&mut seqs, w, None) + }; + emitted += out.iter().map(Vec::len).sum::(); + } + emitted as f64 / (steps * seqs.len()) as f64 + }; + + let cohort = tok_per_step(batch_of(b), false); + let per_sequence = tok_per_step(batch_of(b), true); + // The B=1 baseline is the mean over each of the SAME 128 sequences run + // alone โ€” not `batch_of(1)`, whose single member has one particular + // `agree_mod` and is not representative of the cohort's mix. + let solo = batch_of(b) + .into_iter() + .map(|s| tok_per_step(vec![s], false)) + .sum::() + / b as f64; + + assert!( + per_sequence > cohort * 1.10, + "per-sequence advance must beat the cohort rule at B=128 by a clear margin \ + (cohort {cohort:.4}, per-sequence {per_sequence:.4} tok/step)" + ); + assert!( + (per_sequence - solo).abs() < 1e-9, + "per-sequence advance at B=128 must reach what the same fixture reaches at B=1 \ + (B=128 {per_sequence:.4}, B=1 {solo:.4} tok/step) โ€” that is what 'no cohort \ + barrier' means" + ); + assert!( + cohort < solo, + "the cohort rule must lose against B=1 or there is no problem to solve \ + (cohort {cohort:.4}, B=1 {solo:.4} tok/step)" + ); } /// The window invariant closes: over every depth 1..=8 and every ragged