SIGN IN SIGN UP

[Spec Decode] Register draft-model KV as a scheduler-managed cache group (#630)

CC: @ricky-chaoju 

### Problem
`method="draft_model"` speculative decoding re-prefills a request's
entire prompt into the draft model's KV on every new request. The
proposer's private block allocator has no content-based reuse, so even a
resubmitted identical prompt pays a full draft prefill again (#482). PR
#500 fixed this with a proposer-private prefix cache (vLLM's BlockPool
for hashing, refcounting, and eviction), but was closed: the cache lived
in the proposer/worker path while the real cache policy, `cache_salt`,
per-request cache-read rules (`skip_reading_prefix_cache`), and the
actual scheduler admission/allocation limit live in the
scheduler/request layer, which a proposer-private allocator structurally
cannot see.

### Fix
Split the draft model's KV by where cross-request reuse actually
happens:

- **Committed portion (`[0, committed_len)`)**: Registered as a real
second scheduler-owned `KVCacheSpec` group
(`cache_policy.py::_draft_layer_specs`), sized and budgeted alongside
the target's own groups. The scheduler hashes, matches, admits, and
evicts it exactly like the target's, so `cache_salt`,
`skip_reading_prefix_cache`, and the real admission limit apply
automatically. `draft_model_proposer.py`'s `_make_decode_plan` and
`_make_prefill_plan` read block IDs straight off
`state.block_ids[self._committed_group_index]` (the scheduler's own
assignment) rather than a private pool.
- **Speculative lookahead tail**: Positions drafted ahead of
`committed_len` (not yet verified) have no scheduler concept, so it
stays a small proposer-local scratch reservation over-provisioned into
the draft's physical backend beyond the scheduler-visible block count.
- **Validation fail-fast**: `resolve_draft_dims` now rejects draft
models with sliding-window or hybrid attention (e.g., Gemma4) at
config-resolution time, since `_draft_layer_specs` unconditionally emits
`FullAttentionSpec` and a mismatch would silently corrupt KV cache
sizing.
- **Ingest alignment**: `DraftModelProposer` now ingests every active
decode and prefill row each step (chunked or not, greedy-eligible or
not), not just drafting-eligible ones. The scheduler advances
`num_computed_tokens` and marks committed-group blocks "cached"
uniformly per step, trusting whatever it scheduled was actually
computed; skipping ingest for excluded rows would leave the scheduler
believing blocks hold real KV that was never written. Only drafting
(producing returned tokens, running extra lookahead steps) stays
restricted to the greedy-eligible, non-intermediate subset.

### Behavior Changes
The draft cache's memory is now counted against the shared Metal KV
budget rather than allocated after the target's budget was computed and
left untracked. Concretely, each scheduler block now reserves room for
the draft layers as well, so `per_block_bytes` doubles (1,835,008 ->
3,670,016 on a 128 GB Mac) and `num_blocks` halves (30,398 -> 15,189)
for the same `VLLM_METAL_MEMORY_FRACTION`. This is a hard requirement,
not a latent risk: on base at 0.5 the target takes the entire budget
(`kv_budget` 55.78 GB, `num_blocks` 30,398) and drafting dies with a
Metal out-of-memory inside `mx.eval(drafts)` — it only fits at 0.3. With
this PR the same 0.5 fraction works end to end. `RequestState` gains a
`num_computed_tokens` field set at request creation from the scheduler's
own reconciled cache-hit boundary.


### Verification
Same methodology as the original #500 validation thread
(`Qwen/Qwen3-0.6B` draft==target, greedy, resubmitting an identical
prompt under a fresh request ID). The resubmit's first draft plan goes
from full-prompt ingest with zero reuse to ingest 16 / reuse 8176 at 8k
(ingest 16 / reuse 2032 at 2k); first propose drops from 3525 ms to 111
ms and e2e from 66.8 to 31.2 ms/tok at 8k (438 to 64 ms and 21.9 to 19.5
at 2k).

Losslessness is shown by comparing spec-decode output against a
non-speculative reference run on the same tree (rather than cross-tree
SHA comparison, which is unreliable when the base path is broken). The
validation harness is in-tree at
`tools/benchmark/draft_resubmit_benchmark.py`.

### Tests & Documentation
- E2E coverage that the draft group honors `cache_salt` and
`skip_reading_prefix_cache` (the cache-policy gaps named when the
proposer-private draft cache was rejected).
- Stub-level unit test that the scratch reserve covers `max_num_seqs`
concurrent drafters.
- Regression tests verifying sliding-window and hybrid-attention draft
models are rejected at config resolution time.
- Updated `docs/speculative_decoding.md`: documents
scheduler-managed/budgeted draft KV cache, SWA/hybrid draft rejection at
startup, and the block-alignment safety invariant for lookahead writes
into scheduler-owned blocks.

Addresses #482 direction 1 (per-request re-ingest). The steady-state
proposer overhead (#482 Problem 2) is unaffected and left as a
follow-up.

Credits:
@KrxGu  for the first PR #500

---------

Signed-off-by: Mihal Dimo <mihal@kakao.com>
M
Mihal Dimo committed
4d48fbc25c478ebe47a1530bc7ae13ca3930c2c4
Parent: bd32be8
Committed by GitHub <noreply@github.com> on 8/20/2026, 12:15:48 PM