[Model Runner V2][Spec Decode] Add multi-layer MTP speculator

[TheEpicDolphin]

Context

Multi-Module MTP is a variant of MTP that has a separate model layer for each speculative step (i.e. MTP module #i drafts the ith token). Currently the Inkling model's MTP model supports this, with up to 8 specualtive tokens. This PR adds a new speculator class (MultiModuleMTPSpeculator) to support this functionality.

Currently, only single-module MTP is supported in Model Runner V2, which uses the same, single MTP layer for speculating all N tokens.

NOTE: The changes to vllm/models/inkling/nvidia/model.py were ported over from #48768

NOTE: This feature needs KV cache support (#50062) and multimodal embeddings lookahead (#50306) to function correctly. The benchmark results below were captured with those changes in place. The MultiModuleMTPSpeculator is currently DISABLED until the KV cache support PR is merged.

This PR

The biggest challenge for implementing multi-module MTP was adding support to the scheduler for re-prefilling tokens from previous decode steps. The reason this is needed is because the last N-1 MTP layers receive draft tokens from previous MTP layers, which may or may not be rejected during target validation. If the R draft tokens are rejected, then the last R-1 MTP layers contain stale kv cache values for the rejected tokens, and must be updated. This necessitates re-prefilling tokens from the last decode step. Below is a concrete example, and how it's resolved:

NOTATION: hi is hidden state. ai is also hidden states, but i only use it for representing output hidden states used for sampling draft tokens, not those of prefill tokens. pi is prompt tokens, and di is drafted tokens.

Prefill step: p0,p1,...,p7. 4 MTP modules.
p0,...,p7 => Base => h0,...,h7 (d0 sampled from h7)
p1,...,p7,d0 => MTP #1 => h1,...,h7,a0 (d1 sampled from a0)
p2,...,p7,d0,d1 => MTP #2 => h2,...,h7,a0,a1 (d2 sampled from a1)
p3,...,p7,d0,d1,d2 => MTP #3 => h3,...,h7,a0,a1,a2 (d3 sampled from a2)
p4,...,p7,d0,d1,d2,d3 => MTP #4 => h4,...,h7,a0,a1,a2,a3 (d4 sampled from a3)

Verification:
d0,d1,d2,d3,d4 => Base => a0,a1 (Assume that d3, d4 and d5 were rejected. d3' is sampled from a2).

Drafting (without re-prefills, WRONG):
d1,d2' => MTP #1 => a1,a2 (d3' sampled from a2)
d2,d3' => MTP #2 => a2,a3 (d4' sampled from a3)
d3',d4' => MTP #3 => a3,a4 (d5 sampled from a4)
d4',d5 => MTP #4 => a4,a5 (d6 sampled from a5)

d2 remains stale in MTP #3's KV cache. d2 and d3 remain stale in MTP #4's KV cache.


Drafting (with re-prefills, CORRECT):
p7,d0,d1,d2' => MTP #1 => h7,a0,a1,a2 (d3' sampled from a2)
d0,d1,d2',d3' => MTP #2 => a0,a1,a2,a3 (d4' sampled from a3)
d1,d2',d3',d4' => MTP #3 => a1,a2,a3,a4 (d5 sampled from a4)
d2',d3',d4',d5 => MTP #4 => a2,a3,a4,a5 (d6 sampled from a5)

The stale KVs are corrected in MTP #3 and #4 by re-prefilling tokens p7 and d0 from the last decode step.

Most of the changes to the scheduler, KV cache manager and coordinator were to support re-prefilling tokens in the draft model layers like this. Freeing/caching of KV blocks are delayed by the max number of re-prefillable tokens (num_speculative_tokens - 1) to prevent updating KVs after they've been claimed by another request, or prefix-cached.

The entire N token speculation is cudagraph captured. However, we have to build the attention metadata and slot mappings before the drafting loop, because the sequence lengths can change due to re-prefills.

Remaining TODOs

1. KV cache support for multi-module MTP: [Model Runner V2][Spec Decode] Add KV cache support for multi-layer MTP #50062
2. Multimodal embeddings lookahead for all MTP modules: [Model Runner V2][Spec Decode] Gather MM embeddings for all MTP modules #50306
3. Piecewise cudagraph support - Requires wiring the spec_step_idx into the batch descriptor. Currently disabled.
4. Embed tokens only when any inputs are multimodal. This would allow us to only shift token ids rather than full embeddings. Requires special handling of whether MM inputs are present when capturing/replaying the cudagraph.

Evals

Server Config

export VLLM_USE_V2_MODEL_RUNNER=1
export FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED=1

vllm serve thinkingmachines/Inkling-NVFP4 \
      --tensor-parallel-size 4 \
      --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 \
      --no-enable-prefix-caching \
      --speculative-config '{"method": "mtp", "num_speculative_tokens": 8}'

GSM8K — full test set (1319 prompts, concurrency 64)

metric	Prefix Caching OFF	Prefix Caching ON
Successful / failed requests	1319 / 0	1319 / 0
Accuracy (vs official answers)	95.45% (1259/1319)	95.22% (1256/1319)
Acceptance rate	42.46%	43.03%
Acceptance length	4.40	4.44
Per-position acceptance (%)	75.6 / 60.1 / 49.8 / 41.0 / 35.3 / 29.7 / 26.0 / 22.3	75.7 / 60.3 / 49.9 / 41.5 / 36.0 / 30.6 / 26.9 / 23.3
Prefix cache hit rate	—	84.12% (716,992 / 852,381 tokens)
Mean TTFT	503 ms	440 ms
Mean TPOT	28.0 ms	26.9 ms
Output throughput	2106.6 tok/s	2209.4 tok/s
Benchmark duration	219.4 s	211.4 s

Full MMAU (scored, 1000 questions, mmau.py --task all, concurrency 16)

metric	value
Overall accuracy	77.50% (775/1000, 0 errors)
music	76.65% (256/334)
sound	78.08% (260/333)
speech	77.78% (259/333)
Acceptance rate during eval	34.04% (accept len 3.72)
Per-position acceptance (%)	71.8 / 51.3 / 38.4 / 30.3 / 25.0 / 21.2 / 18.3 / 16.0
Elapsed	272.9 s

Benchmarks

GSM8K, Concurrency = 64, Temperature = 0

Metric	MTP 1	MTP 8	Δ
Benchmark duration (s)	275.18	214.95	-21.9%
Request throughput (req/s)	4.79	6.14	+28.2%
Output token throughput (tok/s)	1646.15	2148.42	+30.5%
Median TTFT (ms)	363.58	426.77	+17.4% (worse)
P99 TTFT (ms)	1437.07	1556.07	+8.3% (worse)
Median TPOT (ms)	37.18	27.48	-26.1% (better)
P99 TPOT (ms)	45.42	38.04	-16.2% (better)
Median ITL (ms)	22.40	159.60	+612.5% (worse)
P99 ITL (ms)	223.44	251.09	+12.4% (worse)
Acceptance rate (%)	83.59	42.53	-41.06 pts
Acceptance length	1.84	4.40	+139.1%
Draft tokens	246,348	839,040	+240.6%
Accepted tokens	205,916	356,839	+73.3%
Per-position acceptance (%) — Pos 0	83.59	75.68	-7.91 pts
Per-position acceptance (%) — Pos 1	—	60.11	—
Per-position acceptance (%) — Pos 2	—	49.76	—
Per-position acceptance (%) — Pos 3	—	41.13	—
Per-position acceptance (%) — Pos 4	—	35.38	—
Per-position acceptance (%) — Pos 5	—	29.80	—
Per-position acceptance (%) — Pos 6	—	26.08	—
Per-position acceptance (%) — Pos 7	—	22.29	—

Speed-Bench (8K prompt / 2K generation), Concurrency = 16, Temperature = 0

Metric	MTP 1	MTP 8	Δ
Benchmark duration (s)	164.94	135.11	-18.1%
Request throughput (req/s)	0.78	0.95	+21.8%
Output token throughput (tok/s)	1589.30	1940.19	+22.1%
Median TTFT (ms)	333.40	294.58	-11.6%
P99 TTFT (ms)	3226.45	3532.19	+9.5% (worse)
Median TPOT (ms)	9.74	7.66	-21.4% (better)
P99 TPOT (ms)	10.47	10.00	-4.5%
Median ITL (ms)	15.40	25.57	+66.0% (worse)
P99 ITL (ms)	169.54	177.97	+5.0%
Median E2EL (ms)	20306.49	16297.28	-19.7%
P99 E2EL (ms)	23469.09	21144.10	-9.9%
Acceptance rate (%)	79.27	37.18	-42.09 pts
Acceptance length	1.79	3.97	+121.8%
Draft tokens	146,188	528,112	+261.2%
Accepted tokens	115,886	196,334	+69.4%
Per-position acceptance (%) — Pos 0	79.27	73.17	-6.10 pts
Per-position acceptance (%) — Pos 1	—	55.13	—
Per-position acceptance (%) — Pos 2	—	42.84	—
Per-position acceptance (%) — Pos 3	—	34.64	—
Per-position acceptance (%) — Pos 4	—	28.59	—
Per-position acceptance (%) — Pos 5	—	24.10	—
Per-position acceptance (%) — Pos 6	—	20.80	—
Per-position acceptance (%) — Pos 7	—	18.14	—

MMAU Audio, Concurrency = 16, Temperature = 0

Metric	MTP 1	MTP 8	Δ
Benchmark duration (s)	37.42	38.35	+2.5%
Request throughput (req/s)	5.34	5.21	-2.4%
Output token throughput (tok/s)	1195.32	1172.70	-1.9%
Median TTFT (ms)	225.82	258.79	+14.6% (worse)
P99 TTFT (ms)	689.84	586.49	-15.0% (better)
Median TPOT (ms)	12.48	12.53	+0.4%
P99 TPOT (ms)	15.62	17.79	+13.9% (worse)
Median ITL (ms)	14.57	24.74	+69.8% (worse)
P99 ITL (ms)	199.25	215.30	+8.1% (worse)
Acceptance rate (%)	78.66	33.79	-44.87 pts
Acceptance length	1.79	3.70	+106.7%
Draft tokens	24,982	97,456	+290.2%
Accepted tokens	19,650	32,928	+67.6%
Per-position acceptance (%) — Pos 0	78.66	73.83	-4.83 pts
Per-position acceptance (%) — Pos 1	—	52.86	—
Per-position acceptance (%) — Pos 2	—	39.44	—
Per-position acceptance (%) — Pos 3	—	29.67	—
Per-position acceptance (%) — Pos 4	—	23.81	—
Per-position acceptance (%) — Pos 5	—	19.64	—
Per-position acceptance (%) — Pos 6	—	16.83	—
Per-position acceptance (%) — Pos 7	—	14.22	—

Speed-Bench (2K prompt / 2K generation / 256 prompts), Concurrency = 16, Temperature = 1)

Metric	MTP 1	MTP 8 (standard)	MTP 8 (block)
Benchmark duration (s)	327.70	278.57	272.24
Request throughput (req/s)	0.78	0.92	0.94
Output token throughput (tok/s)	1599.92	1882.04	1925.82
Median TTFT (ms)	219.03	260.82	246.65
Median TPOT (ms)	9.86	8.29	8.07
Median ITL (ms)	15.24	25.60	25.69
Median E2EL (ms)	20449.51	17242.75	16873.60
Acceptance rate (%)	76.96	33.15	34.79
Acceptance length	1.77	3.65	3.78
Draft tokens	296,198	1,149,512	1,109,600
Accepted tokens	227,942	381,022	386,019
Per-position acceptance (%) — Pos 0	76.96	70.15	70.31
Per-position acceptance (%) — Pos 1	—	51.22	52.24
Per-position acceptance (%) — Pos 2	—	38.47	40.23
Per-position acceptance (%) — Pos 3	—	30.12	32.25
Per-position acceptance (%) — Pos 4	—	24.22	26.45
Per-position acceptance (%) — Pos 5	—	19.91	22.03
Per-position acceptance (%) — Pos 6	—	16.75	18.72
Per-position acceptance (%) — Pos 7	—	14.33	16.09

[TheEpicDolphin]

@benchislett

1. I looked at folding these together, and they can't be unified because they operate on opposite sides of the cache:

• eagle_cache_drop is read-side: On a prefix-cache hit, drop the last matched block and recompute it so that MTP receives hidden states from the target model. It never prevents anything from being cached.
• num_reprefillable_tokens is write-side: Don't cache the last N−1 tokens, because their drafter-module KVs aren't final and the next step can re-prefill and rewrite them in place. Cached blocks must be immutable while referenced, and if a block with unfinalized drafter KVs gets cached, another request can take a reference before the rewrite happens. Dropping a block on the second request's hit doesn't undo the mutation race on the first request's write.
• extra_retained_tokens is freeing-side, and load-bearing for Inkling specifically. It has short-conv layers that are treated as sliding windows of size 4 (smaller than the 7-token re-prefill depth), so without extended retention the re-prefill targets' blocks are freed (and reallocatable) before the repair runs.
  Each of the above targets a different issue, and multi-module MTP needs all of them to work correctly. I pinged @ivanium for a review, because he has more context on the KV cache manager and prefix caching.

2. I don't think that subclassing AutoRegressiveSpeculator will buy us a lot in terms of simplification. The two systems are architecturally very different. AutoRegressiveSpeculator is two-phase: _prefill (reusing the target's attention metadata and buffers) + _multi_step_decode (rebuild attention metadata every step, single token decodes), with two cudagraph managers. Multi-module MTP is a single phase that generates all drafts in a single captured graph. We would end up overriding propose, capture, and init_cudagraph_manager. Additionally, the kernels don't generalize. All the prep kernels have to account for the re-prefilled token gap, as well as the query length adjustments (for continued prefills) that don't happen in the single-module MTP case. The per-step update kernels (update_draft_inputs) are also quite different. In the AutoRegressiveSpeculator case it's just appending the last drafted token. But in the MultiModuleMTPSpeculator case the input token ids need to be shifted + append the last drafted token. The latter doesn't update positions/sequence lengths either, and copies hidden states outside of the kernel (because it is processing multiple query tokens per requests rather than a single one).
   In summary, I think that subclass overriding would make this a worse maintenance story than keeping them as sibling classes.

3. I agree with this. The multi-modal embedding logic can likely be shared. That would potentially allow us to share the _run_model method between the two speculators, perhaps as a helper method. However that will require refactoring AutoRegressiveSpeculator and testing to make sure it doesn't break. I'd like to follow up with that effort in a separate PR to avoid complicating this one further.

[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.

[WoosukKwon]

LGTM. Thanks for the cleanup!

[WoosukKwon]

Please follow up with the investigation on Pos 0 acceptance rate. It looks very mysterious/suspicious to me 🤔

[TheEpicDolphin]

@WoosukKwon I investigated the weird pos 0 acceptance rate discrepancy (observed with gsm8k using temp=0) between MTP 1:

---------------Speculative Decoding---------------
Acceptance rate (%):                     83.39     
Acceptance length:                       1.83      
Drafts:                                  26407     
Draft tokens:                            26407     
Accepted tokens:                         22021     
Per-position acceptance (%):
  Position 0:                            83.39     
==================================================

vs MTP 8:

---------------Speculative Decoding---------------
Acceptance rate (%):                     43.71     
Acceptance length:                       4.50      
Drafts:                                  11043     
Draft tokens:                            88344     
Accepted tokens:                         38611     
Per-position acceptance (%):
  Position 0:                            76.17     
  Position 1:                            61.58     
  Position 2:                            51.25     
  Position 3:                            42.43     
  Position 4:                            36.40     
  Position 5:                            30.94     
  Position 6:                            27.22     
  Position 7:                            23.65     
==================================================

To debug, I bucketed the per-position acceptance rate by the previous step's accepted count (instrumented re-run, same config; headline numbers reproduced at 83.05% / 75.72%). Counts cover the first 25k / 10k drafts of each run.

MTP 1:

prev num accepted	count	position 0 acceptance
0 (draft rejected)	4214	70.43%
1 (draft accepted)	20673	85.66%
all	24887	83.08%

MTP 8:

prev num accepted	count	position 0 acceptance
0	2418	70.10%
1	1529	69.91%
2	955	73.30%
3	846	69.15%
4	597	72.36%
5	543	74.95%
6	368	71.47%
7	361	75.62%
8 (chain never broke)	2268	90.21%
all	9885	75.57%

In both cases the position 0 acceptance rate is much lower when the previous step ended in a rejection: 70.43% vs 85.66% for MTP 1, and ~71% (prev 0–7, weighted) vs 90.21% for MTP 8. A rejection indicates that the current content is less predictable, so acceptance is strongly autocorrelated with local context difficulty.

The discrepancy is a sampling effect, not a regression. A position 0 draft only happens at a chain restart. With MTP 1 that's every token, so only 4214/24887 = 16.9% of position 0 drafts follow a rejection. With MTP 8 it's once per accepted chain, so 7617/9885 = 77.1% of them do — position 0 is measured almost entirely on the hard subset.

Applying MTP 1's conditional rates to MTP 8's restart mix predicts 0.771 × 0.704 + 0.229 × 0.857 = 73.9%, vs 75.57% observed. Bucket-matched, MTP 8's first module is if anything slightly better than MTP 1's (71.2% vs 70.4% post-rejection). The deeper modules look healthy too — within the prev=8 bucket, per-position conditional acceptance is .902 .886 .910 .904 .914 .911 .924 .931, i.e. flat to rising across depths 1–7, which wouldn't survive a KV or depth-indexing bug.