[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