From 24613e87a52ea40486cf8e432007eb9080e83a17 Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Mon, 17 Aug 2026 11:54:40 +0100 Subject: [PATCH] =?UTF-8?q?feat(mtp):=20per-sequence=20KV=20advance=20?= =?UTF-8?q?=E2=80=94=20remove=20the=20cohort=20min-rollback?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Arc's batched MTP rolls one dense cache back to `keep = min_i(u_i + a_i)`, so `next_u_i = u_i + a_i + 1 - keep`. One sequence rejecting its first draft pins `keep = 1`, and every other sequence's uncached tail ratchets up to the window, where it drafts zero. The sequences that accept best are the ones that stop drafting. No production engine does this. vLLM (`scheduler.py:1836-1848`), SGLang (`eagle_worker_common.py:587-588`) and TensorRT-LLM (`resource_manager.py: 1330-1346`) all key the advance on a per-request scalar with no reduction over the batch, and leave a rejected draft's KV stale to be overwritten next step. This adds that rule — `plan_per_sequence_step`: `committed_i = u_i + a_i`, `next_u_i = 1` for every sequence at every batch size — behind `ARC_MTP_PER_SEQ_KV`, default OFF, plus the dense-cache mechanism that can honour it (left-alignment, so one shared append offset is correct for every row) and the ragged mask that mechanism requires. Measured on the production plan functions at B=128, depth 3, 200 steps (CPU arithmetic, NOT hardware): cohort 1.1297 tok/step, per-sequence 1.4783, which is exactly what the same 128 sequences reach at B=1. Token-identical to each sequence's solo run. The flag is refused today and says why: V4's `XsRolling` compressor slots carry one `tokens`/`base` for the whole batch, and no model threads a ragged-batch mask yet. Both are named in code and in wave62-CM. `enforce_shared_cache_lockstep` (PR #86) is untouched and stays a Cohort-mode postcondition. Co-Authored-By: Claude Opus 5 (1M context) --- .../mission/wave62-CN-mtp-per-sequence-kv.md | 242 +++++++ mistralrs-core/src/kv_cache/mod.rs | 287 ++++++++ mistralrs-core/src/layers_masker.rs | 213 +++++- mistralrs-core/src/pipeline/mtp_pipeline.rs | 662 +++++++++++++++++- 4 files changed, 1389 insertions(+), 15 deletions(-) create mode 100644 memory/mission/wave62-CN-mtp-per-sequence-kv.md 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 0000000000..8402119e21 --- /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 51d3dd1ecd..93025479c0 100644 --- a/mistralrs-core/src/kv_cache/mod.rs +++ b/mistralrs-core/src/kv_cache/mod.rs @@ -84,6 +84,44 @@ impl KvCache { Self::TurboQuant(Box::new(TurboQuantCache::new(config))) } + /// 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(), @@ -499,6 +537,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 { @@ -1639,6 +1791,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; @@ -1727,6 +1880,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 d7e2acd1c4..998ccf3fcf 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 1c6a36fe7f..25f0539837 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