feat(mtp): per-sequence KV advance — remove the cohort min-rollback

[heydryft]

The correction

wave59-CJ found the ceiling correctly and it got over-generalised into "MTP can't work at batch". 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. #86 is not reverted and not touched: enforce_shared_cache_lockstep is a KvAdvance::Cohort postcondition and stays exactly that. Under PerSequence it is not reached, because two sequences are never put on one length.

Verified against pinned source, not memory

vLLM cc7cf71, SGLang c82e928, TRT-LLM 8efd46e. 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
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. Rejected draft KV is neither freed nor zeroed — the request's length shrinks and the next step overwrites the same slots (SGLang says so at eagle_worker_common.py:450). Lockstep is real, but only in dense/padded batching outside paged engines — the stated motivation of EMS-SD (arXiv 2405.07542) and "Batch Speculative Decoding Done Right" (arXiv 2510.22876). That is the regime Arc is in.

The barrier, as one equation

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 — drafting, verify, verify_proposed and window_verify_row are literally the same code path.

Measured — arithmetic on the production plan functions, not hardware

B=128, depth 3 (w=4), 200 steps:

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	1.4783	—

Fixture numbers, not V4's — the ratio is the claim. Read tok_per_step, never accept_rate: a saturated sequence drafts 0 and contributes proposed = 0.

⚠️ Correction to wave59-CJ §3. It said saturation is a fixed point. When every sequence reaches u = w the minimum is w, so the cohort resets to u = 1 at once and the system oscillates rather than latching. Cost is 24% of tok_per_step, not a collapse. Same fix, overstated magnitude.

What is in the diff

• plan_per_sequence_step / PerSeqStepPlan — committed_i = u_i + a_i, next_u_i = 1, no reduction over the batch.
• KvAdvance + kv_advance() — latched once, logged once, ARC_MTP_PER_SEQ_KV=1, default OFF.
• cache_supports_per_sequence_advance — names the first slot that cannot carry its own length.
• Per-sequence commit in step, immediately after the POST op splits the batched buffer.
• front_pad_kv_cache / front_align_batch — left-alignment. SingleCache::append writes at one shared offset, so aligning every row's live run to the same end column makes that offset correct for all of them; NormalCacheManager::clone_in_cache needs no change and ensure_uniform_batch_cache_lens still passes.
• RaggedKvLens + make_left_padded_causal_mask — the [B,1,t_q,k] additive mask that layout requires.

🔑 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 scope that computed it, from a plain Vec indexed by the batch row. An earlier draft staged it in a HashMap<seq_id, usize> 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 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

KVCacheManager::trim_request_to_num_tokens under ARC_SEGMENTED_KV=1 already routes to SegmentedAllocator::rollback(pool, request_id, 0, num_tokens) (kv_cache_manager.rs:445-448) — per-request, absolute token count, no new API. PagedAttention has no cohort barrier to remove. What blocks paged MTP is upstream: the fast path only accepts CacheBackendMetadata::DefaultInstructions, so a paged batch never reaches that rollback. Separate change. 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; the xs history has no block table) — do not assume V4 end-to-end from it.

🔴 The honest gap — this does NOT move the B=128 number yet

PerSequence is unreachable in production today, deliberately, and the refusal is logged by name. Two blockers, both named in code:

1. XsRolling cannot carry its own length. V4's compressor slots hold two time bases plus scalar tokens/base for the whole batch, and per-sequence advance needs advance to append a different number of rows per sequence — one shared append offset cannot express that. Same gap wave61-CL §6 and wave29-BC §4b name. Third wave to land on it; it is the keystone.
2. No model threads the ragged mask. front_align_batch produces lead_pad; make_left_padded_causal_mask turns it into the mask; what is missing is carrying it through inputs_processor into the forward. dsv4_attention::compose_caller_mask already accepts rank 4.

Refusing 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) is a wrong answer nothing downstream catches.

Ruled out with reasons in memory/mission/wave62-CN-mtp-per-sequence-kv.md §5: advance-by-w-and-mask (would shrink V4's 128-token suffix window to ~59 real tokens — not token-identical), a growing window, discarding the surplus, and bucketing by accept length.

Gates

cargo check 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, all three files 0 → 0).

Mutation runs (fix stubbed, tests re-run) — 5 mutations, every one caught:

plan_per_sequence_step: committed = (u+a).min(w).min(batch_min)   -> 2 FAILED
plan_per_sequence_step: next_uncached follows the cohort rule     -> 4 FAILED
front_pad_single: slice_set at offset 0 instead of `lead`         -> 1 FAILED
make_left_padded_causal_mask: drop the `j < lead` condition        -> 2 FAILED
cache_supports_per_sequence_advance: never refuse                  -> 1 FAILED

The no-ratchet test carries its own control: the cohort rule on the identical fixture must still ratchet, or the premise is gone.

🤖 Generated with Claude Code

[github-actions]

Code Metrics Report

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 Language              Files        Lines         Code     Comments       Blanks
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 C Header                  5          305          210           52           43
 CSS                       2         1181         1036           34          111
 CUDA                     72        24328        17592         4018         2718
 Dockerfile                1           39           22            8            9
 JavaScript               16         3546         2676          482          388
 Jinja2                    7          694          656            5           33
 JSON                     74         4600         4597            0            3
 Makefile                  1            6            5            0            1
 Metal Shading Lan|       33        12224         9431         1142         1651
 PowerShell                1          300          227           30           43
 Python                  138        13219        10939          658         1622
 Shell                    18         3885         2668          963          254
 Plain Text                4         3801            0         2479         1322
 TOML                     32         1458         1271           40          147
 YAML                      3           25           23            2            0
─────────────────────────────────────────────────────────────────────────────────
 HTML                      3         2500         2461           14           25
 |- CSS                    1          249          235            6            8
 (Total)                             2749         2696           20           33
─────────────────────────────────────────────────────────────────────────────────
 Jupyter Notebooks         4          122           83           23           16
 |- Markdown               1           60           30           22            8
 |- Python                 1          122          113            1            8
 (Total)                              304          226           46           32
─────────────────────────────────────────────────────────────────────────────────
 Markdown                175        33119            0        25276         7843
 |- BASH                  61         1426         1074          256           96
 |- C                      1           10           10            0            0
 |- CUDA                   2           84           56           16           12
 |- JSON                  18          708          708            0            0
 |- PowerShell             1            1            1            0            0
 |- Python                23         1008          787          113          108
 |- Rust                  62         2019         1688           73          258
 |- TOML                   6          207          164            0           43
 |- YAML                   4           38           33            5            0
 (Total)                            38620         4521        25739         8360
─────────────────────────────────────────────────────────────────────────────────
 Rust                    650       300989       260895        13277        26817
 |- Markdown             464        20958          471        18032         2455
 (Total)                           321947       261366        31309        29272
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 Total                  1239       433231       320162        67027        46042
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

[heydryft]

Bypass record. Admin-merged without review — this fork has no second reviewer, and the ruleset grants bypass_mode: always to the admin role. The bypass covered the review requirement only; every required status check was green.

At merge time the required contexts on master were the six Rust lanes, and all six passed on head 882429c3: Check (ubuntu-latest, stable), Test Suite (ubuntu-latest, stable), Clippy, Rustfmt, Typos, MSRV Check (1.90).

All 15 substantive checks were success, 0 pending, 0 failures — including cargo check (cuda, workspace), nvcc compile (sm_80) and nvcc compile (sm_90), which matter here because this PR touches #[cfg(feature = "cuda")] code that a macOS cargo check structurally cannot type-check.

Merged without --delete-branch, deliberately: deleting a branch that another PR is based on closes that PR rather than retargeting it, and GitHub then refuses to reopen it while the base ref is missing. #95 was retargeted to master first, and only then was feat/mtp-per-seq-kv deleted. That sequence is now D20.

Landed default-off: per-sequence KV advance is opt-in behind ARC_MTP_PER_SEQ_KV=1 and kv_advance() returns KvAdvance::Cohort unless it is set, so this changes no existing path.