[Spec Decode] Skip re-ingesting accepted draft tokens whose KV the lookahead already wrote (#482 direction 2) (#636)
## Problem
Issue #482's Problem 2: steady-state propose costs ~1.7–1.9x an entire
nospec engine step even at 100% acceptance. The issue names one concrete
piece of it: on full acceptance, **K−1 of K+1 ingested tokens recompute
KV that the previous round's lookahead draft steps already wrote**.
Post-#630 the steady ingest is a 4-token window, so on a full acceptance
that is 2 of 4 tokens doing duplicate work.
#502 fixed the routing half of Problem 2 (94 → ~50 ms/step @8k) and
**prototyped exactly this skip on top of it — measured ~3 ms, and
deliberately left it out of scope** ("rows inside one decode dispatch
largely amortize"). This is that follow-on, completed: implemented,
tested, and measured with the same-day A/B the prototype lacked.
## Fix
Track the speculative KV each round's lookahead steps wrote, and skip
re-ingesting it when the next round's committed tokens prove it still
valid:
1. **Record.** After each drafting round, save per request a ledger of
`position → (physical block id, drafted token)` for the positions the
lookahead steps wrote: `[committed_len, committed_len + K−1)`.
2. **Skip.** The next round's ingest starts at
`_speculative_kv_valid_through` — the first position whose committed KV
is not already valid in the draft cache. A position is skippable only
when:
- its committed token equals the recorded drafted token, **and**
- the scheduler's committed-group block table still maps the position to
the same physical block the speculative write landed in.
Anything else — rejected draft, lookahead KV that crossed into a scratch
block (the scheduler allocates its own block there), a re-allocated
block table, or no ledger — stops the walk, because the ingest forward
must be one contiguous range of positions whose KV is valid up to its
start.
3. **Cap.** The walk is capped one short of `committed_len`: the last
committed token's row must still run — its logits predict this round's
first draft token — so the ingest is never empty.
**Result:** on full acceptance, the steady-state 4-token ingest shrinks
to **2 tokens**.
## Measurements
Same-day, back-to-back A/B of this branch vs plain `main` (Qwen3-0.6B
draft == target, K=3, greedy → 100% acceptance, prefix caching on,
steady-state resubmit runs, in-tree harness
`tools/benchmark/draft_resubmit_benchmark.py`):
| prefix | propose_median_ms (main → branch) | tpot_ms (main → branch) |
|---|---|---|
| 2k | 47.5 → 47.3 | 18.96 → 18.61 |
| 8k | 75.4 → 73.3 | 30.14 → 29.62 |
- **The skip fires as designed:** steady-state plans ingest 2 tokens on
the branch vs 4 on main (23/26 plans per run; the rest are the
resubmit-boundary and partial-acceptance plans).
- **Lossless:** generated token sequences are identical branch vs main
on all four runs (96 tokens each).
- **Honest caveats:** n=1 run per config (each propose_median is a
median of ~26 proposes); the @2k deltas are inside run-to-run noise; the
@8k propose delta (−2.1 ms, ~2.8%) is the strongest signal, and tpot is
directionally better in all four comparisons. The scale matches #502's
prototype (~3 ms): the ceiling is modest because the post-#630 steady
ingest was already a 4-token window — this removes its duplicate half,
not the kernel cost that dominates the rest.
## Correctness
The ledger survives rounds where a request doesn't draft; every entry is
re-validated against the current committed tokens and block table before
use, so staleness can only cost performance, never correctness. The
ledger is cleared on request release and in the periodic GC, so a reused
request id gets no skip.
## Tests & Documentation
- `test_full_acceptance_skips_reingest_of_drafted_kv` — steady-state
ingest shrinks to 2 tokens on full acceptance.
- `test_rejected_draft_reingests_the_full_committed_range` — a rejected
draft falls back to a full re-ingest.
- `test_scratch_block_spec_kv_is_not_reused` — lookahead KV that landed
in a scratch block is not skipped.
- `test_reallocated_block_spec_kv_is_not_reused` — a scheduler
re-allocation that swaps the committed block table's physical blocks
(tokens still matching) stops the walk; nothing is skipped.
- `test_partial_acceptance_never_produces_an_empty_ingest` — the skip
cap keeps the ingest non-empty.
- `test_spec_kv_ledger_cleared_on_release` — releasing a request clears
its ledger.
- `docs/speculative_decoding.md` documents the skip conditions and the
fallback cases.
Stacks on **#630** (reads committed-group block IDs from the
scheduler-owned block table) and **#502** (the 4-token decode-path
ingest this trims).
Addresses **#482 direction 2**. Direction 3 (chunked ingest) and the
residual Problem 2 items (per-row KV re-read in attention, per-layer
block-table hoist) remain open in the issue.
Signed-off-by: Mihal Dimo <mihal@kakao.com> M
Mihal Dimo committed
211a1e4bc976cd6c0c961cad8da59d649ea9bd65
Parent: 3c9bbdc
Committed by GitHub <noreply@github.com>
on 8/22/2026, 1:28:48 PM