[Model Runner V2][MTP] Share topk index buffer between draft steps

[TheEpicDolphin]

Context

Deepseek-style topk index selection for MTP has a bug during proposal stage when the topk indices are shared among all MTP draft step forward passes. After the first draft step, set_skip_topk is called on the MTP model to force the remaining draft steps to reuse the same topk indices values (held within topk_indices_buffer of the attention module).

This feature was introduce in #44420 for Model Runner V1.

This PR

Enable this feature for Model Runner V2. I didn't want to leak model-specific optimizations into AutoRegressiveSpeculator, so i introduced a callback pattern. New optimizations/features that need to run at specific points during the propose operation can inherit from this class and override whatever methods are needed.

class AutoRegressiveSpeculatorCallbacks(Protocol):
    def on_prefill_begin(self) -> None: ...
    def on_prefill_end(self, last_token_indices: torch.Tensor) -> None: ...
    def on_multi_step_decode_begin(self) -> None: ...
    def on_multi_step_decode_end(self) -> None: ...

The topk indexer sharing feature registers itself after MTPSpeculator loads its draft model. After determining whether the draft model is a MTP model that can do index sharing AND index sharing is enabled, it is added to the list of callback classes for the auto regressive speculator to call during prefill begin/end and multi step decode begin/end.

Benchmarks

Server

vllm serve nvidia/DeepSeek-V3.2-NVFP4 \
  -tp 4 \
  -O3 \
  --no-enable-prefix-caching \
  --max-num-seqs 64 \
  --kernel-config.enable_flashinfer_autotune=False \
  --speculative_config '{"method":"mtp","num_speculative_tokens":3}' \
  --compilation-config '{"cudagraph_mode": "full_decode_only", "mode": 0}'

Bench

vllm-bench \
    --model $MODEL \
    --backend openai-chat \
    --dataset-name speed-bench \
    --speed-bench-config throughput_16k \
    --speed-bench-max-input-len 10240 \
    --speed-bench-category low_entropy \
    --request-rate inf \
    --max-concurrency 32 \
    --num-prompts 256 \
    --output-len 2048 \
    --ignore-eos \
    --temperature 1.0 \
    --percentile-metrics "ttft,tpot,itl,e2el" \
    --result-dir ~/dev \
    --save-result

Results

Metric	topk index sharing = False	topk index sharing = True
Request throughput (req/s)	0.80	0.80
Output token throughput (tok/s)	1629.90	1640.71
Total token throughput (tok/s)	9779.42	9844.27
Median TTFT (ms)	608.38	601.95
P90 TTFT (ms)	4206.77	3501.30
Median TPOT (ms)	18.41	18.24
P90 TPOT (ms)	22.15	21.66
Acceptance rate (%)	60.96	60.72
Acceptance length	2.83	2.82
Drafts	185,354	185,828
Draft tokens	556,062	557,484
Accepted tokens	338,992	338,495
Pos 0 (%)	87.45	86.91
Pos 1 (%)	60.10	60.18
Pos 2 (%)	35.34	35.06

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

/ci run

[github-actions]

✅ Triggered Buildkite CI #83222 for commit b2cae9c53691.