[Model Runner V2][Spec Decode] Add KV cache support for multi-layer MTP

[TheEpicDolphin]

Summary

This PR adds the scheduler and KV-cache-manager support required for multi-module MTP
(one MTP module per speculative step, e.g. Inkling's 8-depth checkpoint). It is the
companion to #48892, which introduced the speculator itself to Model Runner V2.

The core property this PR protects: the multi-module drafter reads ahead of the
computed tokens during prefill. MTP module m computes position p's KV from token
p + m + 1, so at every chunked-prefill boundary the drafter consumes the next
num_speculative_tokens (N) known prompt tokens — the "prefill lookahead" — to write
exact KVs into all modules. That lookahead, plus the decode-time rejection re-prefill
(which rewrites up to N−1 trailing positions), creates three hazards that the scheduler
and KV cache layer must handle: chunk boundaries landing where lookahead tokens don't
exist, caching/freeing KV that can still change, and prefix-cache hits serving KV that
encodes another request's continuation.

Changes

Config (vllm/config/speculative.py)

• Inkling MTP hf-config override now exposes all checkpoint MTP depths
  (n_predict = num_nextn_predict_layers instead of clamping to 1), and the
  "exactly one speculative token" restriction is removed. Module i drafts
  speculative token i.

Speculator selection (vllm/v1/worker/gpu/spec_decode/__init__.py)

• init_speculator routes to MultiModuleMTPSpeculator when
  use_multi_module_mtp() (method "mtp" with >1 usable MTP layer).

Scheduler (vllm/v1/core/sched/scheduler.py)

• New unified field num_prefill_lookahead: how many positions past the computed
  tokens the drafter reads during prefill (N for multi-module MTP, 1 for other
  eagle-family drafters, 0 without spec decode). All three consumers below are
  projections of it.
• _reserve_prefill_lookahead (both scheduling loops, after encoder
  truncation): never end a prefill chunk with 0 < remaining < N tokens before
  the prefill end — either finish the prefill or leave ≥ N for the next chunk.
  Without this, a boundary near the prompt end has no real lookahead tokens; the
  drafter falls back to sampled drafts and the trailing modules' KVs at those
  positions are polluted permanently (they fall outside every future query
  window, so nothing rewrites them). Vacuous no-op for eagle-family (lookahead 1).
• Encoder scheduling shift generalized: shift_computed_tokens to
  _try_schedule_encoder_inputs is now num_prefill_lookahead (previously the
  hardcoded eagle +1). MM spans starting inside the lookahead window get
  encoded one chunk early, so the drafter's future-token embeddings are
  available; when the encoder budget can't cover a span, the existing rollback
  lands the boundary at span_start − N, keeping the lookahead window all-text.
  (The runner-side consumption of these early embeddings lands in a follow-up
  PR; until then MM placeholder lookahead tokens use text embeddings.)
• Passes num_prefill_lookahead to KVCacheManager.

KV cache manager / coordinator

(kv_cache_manager.py, kv_cache_coordinator.py, single_type_kv_cache_manager.py)

• num_prefill_lookahead is threaded into KVCacheCoordinator, which derives
  num_reprefillable_tokens = max(0, lookahead − 1).
• Delayed caching: cache_blocks only hash-registers tokens up to
  num_computed − num_reprefillable_tokens (mirrored in the hybrid
  coordinator's EAGLE lookahead-block eligibility). During decode, rejection
  re-prefill can rewrite the last N−1 tokens' draft KVs; registering them
  earlier would expose unverified KV to other requests and mutate blocks after
  they are shared.
• Prefix-hit soundness assert: for every EAGLE-flagged group,
  block_size >= num_prefill_lookahead. The last N slots of any cached
  prefix hold draft KVs computed from tokens past the block hash (the writer's
  continuation — lookahead tokens at a chunk boundary, or sampled/draft tokens
  later). The pre-existing EAGLE last-block drop already recomputes the trailing
  block on every hit — with the new request's own lookahead tokens, making the
  rewrite exact — but a single dropped block only covers all N polluted slots if
  the block size is large enough. Rather than generalizing the drop to multiple
  blocks, we assert (real configs use block 16/64 with N ≤ 8).

Sliding-window retention

(kv_cache_interface.py, kv_cache_utils.py, single_type_kv_cache_manager.py)

• New SlidingWindowSpec.extra_retained_tokens (tagged as N−1 for multi-module
  MTP in get_kv_cache_configs): the SWA free boundary lags by N−1 tokens, and
  max_memory_usage_bytes accounts for the extra retained blocks. Rejection
  re-prefill recomputes positions up to N−1 behind the tip, and each recomputed
  position needs its full attention window; without the lag, those windows reach
  into already-freed (null) blocks and the corrected KVs would be computed from
  garbage. Single source of truth on the spec keeps pool sizing, the admission
  cap, and block eviction consistent. Prior spec methods never needed this knob
  because no drafter re-processed positions behind the tip (eagle's ±1 is a
  token-id shift, not a position shift).

Why this is not duplicating an existing PR

Multi-module MTP is not supported by any open PR; the eagle/single-module paths this
builds on (drop_eagle_block, encoder shift, use_eagle plumbing) are extended
in place rather than duplicated.

Testing

• tests/v1/core/test_scheduler.py — 137 passed (includes the eagle encoder-shift
  regression test, chunked-prefill, and preemption paths).
• tests/v1/core/test_prefix_caching.py — 89 passed.
• tests/config/test_speculative_draft_hf_overrides.py — updated for the
  all-depths Inkling override.
• Direct construction check of the new assert: accepts block_size=16, N=8 and
  spec-decode-off; rejects block_size=4, N=8 with a clear message.
• End-to-end acceptance-rate evaluation with the Inkling multi-module MTP model
  (chunked prefill + prefix caching enabled): [results to be added].

AI assistance (Claude) was used for implementation and review of this PR; all
changes were human-reviewed.

Evals

GSM8K, 1319 questions, 5-shot, concurrency=64, max_tokens=2048

Metric	Baseline	MTP8, Prefix Caching ENABLED	MTP8, Prefix Caching DISABLED
Accuracy	0.952	0.958	0.953
Invalid responses	0.025	0.023	0.025
Total latency (s)	334.306	206.880	122.817
Questions per second	3.945	6.376	10.740
Total output tokens	463,363	456,388	459,382
Output tokens per second	1386.045	2206.053	3740.375

Benchmarks

Server Config

export VLLM_USE_V2_MODEL_RUNNER=1
export FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED=1
# Needed to prevent hangs during serving. Still need to investigate the root cause.
export LAMPORT_RS_SCONV=0 

vllm serve thinkingmachines/Inkling-NVFP4 \
      --tensor-parallel-size 4 \
      --max-num-batched-tokens 8192 \
      --gpu-memory-utilization 0.9 \
      --max-model-len 16384 \
      --max-num-seqs 64 \
      --tokenizer-mode inkling \
      --tool-call-parser inkling \
      --reasoning-parser inkling \
      --enable-auto-tool-choice \
      --trust-remote-code \
      --kernel-config.enable_flashinfer_autotune=False \
      --speculative-config '{"method": "mtp", "num_speculative_tokens": 8}'

GSM8K, 1319 questions, concurrency=64, max_tokens=2048

Metric	MTP8, No Prefix Caching	MTP8, Prefix Caching
Successful requests	1319	1319
Maximum request concurrency	64	64
Benchmark duration (s)	217.19	124.03
Total input tokens	872,166	872,166
Total generated tokens	487,670	475,072
Request throughput (req/s)	6.07	10.63
Output token throughput (tok/s)	2245.34	3830.33
Peak output token throughput (tok/s)	868.00	925.00
Total token throughput (tok/s)	6260.97	10862.27
Time to First Token
Mean TTFT (ms)	435.87	230.70
Median TTFT (ms)	390.64	211.64
P99 TTFT (ms)	1430.27	595.46
Time per Output Token (excl. 1st)
Mean TPOT (ms)	25.84	14.71
Median TPOT (ms)	25.82	14.61
P99 TPOT (ms)	35.48	19.74
Inter-token Latency
Mean ITL (ms)	138.38	74.91
Median ITL (ms)	122.13	69.06
P99 ITL (ms)	220.47	128.72
Speculative Decoding
Acceptance rate (%)	43.38	43.09
Acceptance length	4.47	4.45
Drafts	109,085	106,821
Draft tokens	872,680	854,568
Accepted tokens	378,539	368,231
Per-position acceptance (%)
Position 0	76.12	75.87
Position 1	60.98	60.33
Position 2	50.58	50.20
Position 3	41.87	41.90
Position 4	36.19	35.90
Position 5	30.78	30.55
Position 6	27.04	26.83
Position 7	23.43	23.14

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @TheEpicDolphin.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[github-actions]

✅ Triggered Buildkite CI #83251 for commit 01b6e182e3eb.

[TheEpicDolphin]

/ci run

[github-actions]

✅ Triggered Buildkite CI #83401 for commit 0841fed141e9.

[WoosukKwon]

LGTM and sorry for all the delays.

I'm probably not the best person to verify the correctness of the PR, but it seems correct to me at the high level, and only affect the behavior on multi-layer MTP. Therefore, I think it is safe to merge. As the last step, can you please run Inkling and add the accuracy & acceptance rates?

End-to-end acceptance-rate evaluation with the Inkling multi-module MTP model
(chunked prefill + prefix caching enabled): [results to be added].

[TheEpicDolphin]

/ci run

[github-actions]

✅ Triggered Buildkite CI #83778 for commit ba0d589e8fe5.

[TheEpicDolphin]

@WoosukKwon added eval and benchmarks to the PR summary

[TheEpicDolphin]

/ci run

[github-actions]

✅ Triggered Buildkite CI #83794 for commit 02f7c1e52ef1.