[MRV2][Spec] Fuse AR speculator multi-step decodes back into one CUDA graph

[yiz-liu]

Purpose

This PR restores fused multi-step CUDA graph execution for autoregressive speculative decoding in Model Runner V2.

#41162 fixed stale attention metadata by rebuilding it and replaying a separate CUDA graph for every draft step. While correct, that design reintroduces per-step Python dispatch, metadata construction, and CUDA graph launch overhead.

This PR instead captures the post-prefill draft loop into one CUDA graph. Attention metadata is built once, and backend-owned step-dependent metadata is updated in place between draft steps.

Additional Background

Our local tests suggest the historical illegal memory access is not fully explained by stale draft metadata:

1. A pre-[Model Runner V2] Rebuild attn metadata between draft decode steps #41162 version reproduces the illegal memory access, but adding this metadata refresh alone does not fix it.
2. Since v0.23.0, the illegal memory access does not reproduce even when this refresh is disabled.

Design

Design of #41162:

sequenceDiagram
    participant Spec as AutoRegressiveSpeculator
    participant Inputs as Draft Input Buffers
    participant Builder as Attention Metadata Builder
    participant Graph as Draft CUDA graph

    loop draft_step = 0..N-1
        Spec->>Builder: _build_draft_attn_metadata()
        Builder-->>Spec: fresh attn_metadata
        Spec->>Graph: replay per-step draft graph
        Spec->>Inputs: update_draft_inputs()
    end

Loading

Design of this PR:

sequenceDiagram
    participant Spec as AutoRegressiveSpeculator
    participant Input as Draft input buffers
    participant Attn as Attention metadata builder
    participant Graph as Fused draft CUDA graph

    Spec->>Attn: Build metadata once
    Spec->>Graph: Replay fused multi-step graph
    loop Each captured draft step
        Graph->>Graph: Draft model forward and sampling
        Graph->>Input: Advance draft inputs and positions
        Graph->>Attn: `update_draft_decode_metadata()`
        Attn-->>Attn: Update backend-derived metadata
    end

Loading

Backends opt in through supports_draft_decode_metadata_update. The default is disabled, so unadapted backends retain the split-graph path and rebuild metadata between draft steps.

update_draft_decode_metadata() is executed while the fused graph is being captured. Implementations must emit CUDA graph capture-safe operations and keep replayed tensor state in persistent storage; Python code is not executed during replay.

Current backend handling:

Backend	Fused path	Metadata handling
FlashAttention / FA3 without DCP	Enabled	FA3 scheduler metadata is regenerated into persistent storage; other FlashAttention paths are no-op
DeepSeek V4 sparse SWA	Enabled	Recomputes SWA lengths and indices, refreshes tile schedulers, and invalidates the FlashInfer sparse-index cache
Triton Attention	Enabled	Step-dependent fields already reference persistent input buffers
Triton MLA without DCP	Enabled	No additional materialized metadata requires updating
Position-static draft models	Enabled	Positions do not advance, so no metadata update is required
DeepSeek V4 ROCm sparse SWA	Disabled	ROCm-specific ragged SWA indices and indptrs still require adaptation
Other backends	Disabled by default	Fall back to per-step metadata rebuilding

Test Plan

Tested on:

• Device: H100 8 * 80G
• vLLM version: main branch (0406ba2, fix pre-commit broken (fix pre-commit broken #51341)) + this PR

Online serving: device-bound workload

The online serving sweep used:

vllm-bench \
--backend openai-chat \
--base-url http://127.0.0.1:8000 \
--model <model> \
--dataset-name hf \
--dataset-path philschmid/mt-bench \
--ignore-eos \
--request-rate inf \
--temperature <0 or 1> \
--num-warmups 64 \
--hf-output-len 1024 \
--seed 0 \
--disable-shuffle \
--sweep-max-concurrency 1,16,64 \
--sweep-num-prompts-factor 10

VLLM_USE_FLASHINFER_MOE_FP8=0 \
FLASHINFER_DISABLE_VERSION_CHECK=1 \
VLLM_USE_V2_MODEL_RUNNER=1 \
vllm serve /path/to/DeepSeek-V4-Flash \
--tensor-parallel-size 8 \
--enable-expert-parallel \
--max-model-len 32768 \
--max-num-seqs 128 \
--block-size 256 \
--kv-cache-dtype fp8 \
--tokenizer-mode deepseek_v4 \
--tool-call-parser deepseek_v4 \
--reasoning-parser deepseek_v4 \
--no-enable-prefix-caching \
--kernel-config.enable_flashinfer_autotune=False \
--speculative-config '{"method":"mtp","num_speculative_tokens":3}'

# For Qwen3.5-9B, we increased `num_speculative_tokens` to 7 to amplify host-side overhead.
FLASHINFER_DISABLE_VERSION_CHECK=1 \
VLLM_USE_V2_MODEL_RUNNER=1 \
vllm serve /path/to/Qwen3_5-9B \
--tensor-parallel-size 2 \
--max-model-len 32768 \
--max-num-seqs 128 \
--language-model-only \
--enable-auto-tool-choice \
--tool-call-parser qwen3_coder \
--reasoning-parser qwen3 \
--no-enable-prefix-caching \
--speculative-config '{"method":"mtp","num_speculative_tokens":7}'

The baseline forces use_fused_multi_step_decode=False; the fused runs use normal backend capability selection. All 24 runs completed without failed requests.

For DeepSeek-V4-Flash:

Temp	Concurrency	Median ITL baseline (ms)	Median ITL fused (ms)	ITL delta	Acceptance baseline	Acceptance fused	Acceptance delta	Acceptance-normalized ITL delta
0	1	10.75	10.70	-0.4%	46.50%	48.85%	+2.36 pp	+2.5%
0	16	34.44	34.91	+1.4%	53.47%	53.57%	+0.10 pp	+1.5%
0	64	85.44	78.09	-8.6%	54.31%	54.20%	-0.12 pp	-8.7%
1	1	10.88	10.99	+1.0%	38.53%	36.75%	-1.78 pp	-1.5%
1	16	35.28	34.90	-1.1%	43.33%	42.78%	-0.55 pp	-1.8%
1	64	79.16	81.52	+3.0%	43.42%	43.25%	-0.17 pp	+2.7%

For Qwen3.5-9B:

Temp	Concurrency	Median ITL baseline (ms)	Median ITL fused (ms)	ITL delta	Acceptance baseline	Acceptance fused	Acceptance delta	Acceptance-normalized ITL delta
0	1	8.90	9.05	+1.7%	36.70%	36.70%	+0.00 pp	+1.7%
0	16	11.55	11.49	-0.5%	45.07%	44.55%	-0.52 pp	-1.4%
0	64	23.14	22.78	-1.5%	45.75%	45.75%	+0.00 pp	-1.5%
1	1	8.92	9.02	+1.1%	27.40%	26.31%	-1.10 pp	-1.6%
1	16	11.59	11.50	-0.8%	33.52%	33.16%	-0.36 pp	-1.5%
1	64	23.17	23.05	-0.5%	34.01%	33.77%	-0.24 pp	-1.0%

The acceptance-normalized ITL delta compares median ITL * acceptance length, which approximates the time per speculative scheduler step and removes the effect of small acceptance-rate differences between paired runs.

The normalized results fluctuate in both directions and show no systematic regression so we cannot conclude it has performance gain. My guess is that these workloads remain device-bound, so reducing host dispatch helps little.

Host-bound profiling

To expose host dispatch overhead, Qwen3.5-9B was profiled with TP=2 and torch_profiler_with_stack=True.

llm = LLM(
    model=MODEL,
    trust_remote_code=True,
    tensor_parallel_size=2,
    max_model_len=4096,
    max_num_seqs=4,
    enable_prefix_caching=False,
    speculative_config={
        "method": "mtp",
        "num_speculative_tokens": 3,
    },
    disable_log_stats=False,
    profiler_config={
        "profiler": "torch",
        "torch_profiler_dir": "vllm-profile/",
        "torch_profiler_with_stack": True,
        "torch_profiler_use_gzip": False,
        "wait_iterations": 4,
        "max_iterations": 4,
    }
)

The table below averages rank 0 ProfilerStep#5 through ProfilerStep#7, excluding the first profiled iteration.

Metric	Split graph	Fused graph	Change
CUDA graph launches per scheduler step	4	3	-1
Multi-step draft decode CPU span	2.050 ms	0.860 ms	-58.1%
Complete draft propose CPU span	2.971 ms	1.493 ms	-49.8%
Profiler step CPU span	20.385 ms	14.537 ms	-28.7%

With three speculative tokens, the split path launches one target graph, one draft-prefill graph, and two draft-decode graphs per scheduler step. The fused path replaces the two draft-decode launches with one multi-step graph.

Trace screenshots

Split graph: Each post-prefill draft step returns to Python, rebuilds attention metadata, and launches a separate draft graph. In the stable rank 0 iterations, _multi_step_decode takes 1.929–2.149 ms.

Fused graph: Both post-prefill draft steps execute inside one CUDA graph. Input advancement and backend metadata updates are captured as part of that graph. In the same rank 0 interval, _fused_multi_step_decode takes 0.848–0.883 ms.

Raw traces for both TP ranks and the profiler summary outputs are attached here:
fused-graph.zip
split-graph.zip

Per-step Metadata Check

The following values were captured with temporary debug instrumentation. The
final PR does not keep this logging path.

With refresh enabled, SWA metadata advances from draft step 1 to draft step 2:

draft_step=1 num_decodes=4 scheduler=FlashMLASchedMeta:have_initialized=False,tile_scheduler_metadata_shape=None,num_splits_shape=None
field=is_valid_token sample=[True, False, False, False]
field=decode_swa_lens sample=[6, 0, 0, 0]
field=decode_swa_indices sample=[192, 193, 194, 195, 196, 197, -1, -1, -1, -1, -1, -1]

draft_step=2 num_decodes=4 scheduler=FlashMLASchedMeta:have_initialized=False,tile_scheduler_metadata_shape=None,num_splits_shape=None
field=is_valid_token sample=[True, False, False, False]
field=decode_swa_lens sample=[7, 0, 0, 0]
field=decode_swa_indices sample=[192, 193, 194, 195, 196, 197, 198, -1, -1, -1, -1, -1]

With refresh disabled, draft step 2 keeps step-1 SWA metadata:

draft_step=1 num_decodes=4 scheduler=FlashMLASchedMeta:have_initialized=False,tile_scheduler_metadata_shape=None,num_splits_shape=None
field=is_valid_token sample=[True, False, False, False]
field=decode_swa_lens sample=[6, 0, 0, 0]
field=decode_swa_indices sample=[192, 193, 194, 195, 196, 197, -1, -1, -1, -1, -1, -1]

draft_step=2 num_decodes=4 scheduler=FlashMLASchedMeta:have_initialized=True,tile_scheduler_metadata_shape=(132, 8),num_splits_shape=(5,)
field=is_valid_token sample=[True, False, False, False]
field=decode_swa_lens sample=[6, 0, 0, 0]
field=decode_swa_indices sample=[192, 193, 194, 195, 196, 197, -1, -1, -1, -1, -1, -1]

So the refresh path updates the same SWA metadata that the split-graph path would rebuild. is_valid_token stays unchanged, as expected, because the padded request set is unchanged across draft steps.

Other things

AI assistance was used for code review, performance-data aggregation. The human author made and reviewed the changes and ran the evaluations above.

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

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

[TheEpicDolphin]

Thanks for this optimization, it would be really awesome to have this. However, i believe we need to gate this fully-captured decode loop functionality on whether all attention backends for the model support refresh_meta_for_draft_decodes. The reason we can't enable this for all draft models today is because some backends (for example, FlashAttention's scheduler metadata:

vllm/vllm/v1/attention/backends/flash_attn.py

Lines 505 to 521 in 735def4

	return get_scheduler_metadata(
	batch_size=batch_size,
	max_seqlen_q=max_query_len,
	max_seqlen_k=max_seq_len,
	num_heads_q=self.num_heads_q * self.dcp_world_size,
	num_heads_kv=self.num_heads_kv,
	headdim=self.headdim,
	cache_seqlens=seqlens,
	qkv_dtype=qkv_dtype,
	cu_seqlens_q=cu_query_lens,
	page_size=self.block_size,
	causal=causal,
	window_size=_maybe_symmetrize_window(
	self.aot_sliding_window, causal
	),
	num_splits=max_num_splits,
	)

) generate position-dependent metadata. This metadata becomes stale as we draft tokens in the loop, unless we refresh the attention metadata.

Right now only sparse SWA seems to support this refresh. To safely enable this optimization for Deepseek v4 MTP we need to make sure that all of its backends support this refresh.

[mergify]

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

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

[yiz-liu]

Thanks for this optimization, it would be really awesome to have this. However, i believe we need to gate this fully-captured decode loop functionality on whether all attention backends for the model support refresh_meta_for_draft_decodes. The reason we can't enable this for all draft models today is because some backends (for example, FlashAttention's scheduler metadata:

vllm/vllm/v1/attention/backends/flash_attn.py

Lines 505 to 521 in 735def4

	return get_scheduler_metadata(
	batch_size=batch_size,
	max_seqlen_q=max_query_len,
	max_seqlen_k=max_seq_len,
	num_heads_q=self.num_heads_q * self.dcp_world_size,
	num_heads_kv=self.num_heads_kv,
	headdim=self.headdim,
	cache_seqlens=seqlens,
	qkv_dtype=qkv_dtype,
	cu_seqlens_q=cu_query_lens,
	page_size=self.block_size,
	causal=causal,
	window_size=_maybe_symmetrize_window(
	self.aot_sliding_window, causal
	),
	num_splits=max_num_splits,
	)

) generate position-dependent metadata. This metadata becomes stale as we draft tokens in the loop, unless we refresh the attention metadata.
Right now only sparse SWA seems to support this refresh. To safely enable this optimization for Deepseek v4 MTP we need to make sure that all of its backends support this refresh.

@TheEpicDolphin Thanks for the review! Now I understand it's not just about query_start_loc and query_lens for FA3, and I plan to address this incrementally in this PR.

Firstly, a commit will add the configuration and per-backend capability gating with an automatic fallback to the existing per-step rebuild path. The next commit will add the required metadata refresh for FA3, and then I will audit the remaining EAGLE/MTP attention backends for similar sequence- or position-dependent state and see if there are other adjustments required.

[TheEpicDolphin]

/ci run

[github-actions]

✅ Triggered Buildkite CI #83294 for commit fc8ceaf0fc49.

[TheEpicDolphin]

LGTM! cc: @ZJY0516