[TheEpicDolphin]

Context

I noticed an issue previously for MRV2 DSV4 where the model would output gibberish for simple prompts. The root cause was that the model was running both the warmup and capture forward passes using the same attention metadata. This was problematic for attention backends (e.g. FlashMLA) that lazily initialize metadata state. The warmup pass (run eagerly) triggers the lazy init and flips an have_initialized flag on the metadata, so the capture pass then sees the flag already set, skips the init, and never records the initialization kernel into the graph. That issue was fixed with this PR by using separate attention metadata for warmup and capture, so the init runs and is recorded during capture.

A similar issue happens with MRV2 DSV4 MTP, and likely also with other speculator attention backends that perform lazy metadata initializations. In the speculator case, PrefillEagleCudaGraphManager does not build its own metadata. It reuses the target model's capture metadata, returning that same object from create_forward_fn both when CudaGraphManager calls it before warmup and again before capture.

The key subtlety is that CUDA graph capture records kernels without executing them. So when the target model captured its graph, the FlashMLA planner's host code set have_initialized=True and bound a freshly-allocated tile_scheduler_metadata / num_splits tensor — but the kernel that fills those buffers was only recorded, never executed. The buffers are valid allocations holding garbage.

The speculator then receives this metadata and runs its prefill warmup eagerly (real execution). Because have_initialized is already True, the planner is skipped and the garbage scheduler buffers are fed straight into the FlashMLA decode kernel, which indexes the KV cache / block table out of bounds. The result is this error during speculator prefill capture:

(Worker_DP3_EP3 pid=2845516) ERROR 06-01 23:50:49 [multiproc_executor.py:966]   File "/home/gdelfin/vllm/vllm/v1/worker/gpu/model_runner.py", line 653, in capture_model
(Worker_DP3_EP3 pid=2845516) ERROR 06-01 23:50:49 [multiproc_executor.py:966]     self.speculator.capture(captured_attn_states)
(Worker_DP3_EP3 pid=2845516) ERROR 06-01 23:50:49 [multiproc_executor.py:966]   File "/home/gdelfin/vllm/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py", line 443, in capture
(Worker_DP3_EP3 pid=2845516) ERROR 06-01 23:50:49 [multiproc_executor.py:966]     self.prefill_cudagraph_manager.capture(
(Worker_DP3_EP3 pid=2845516) ERROR 06-01 23:50:49 [multiproc_executor.py:966]   File "/home/gdelfin/vllm/vllm/v1/worker/gpu/spec_decode/eagle/cudagraph.py", line 53, in capture
(Worker_DP3_EP3 pid=2845516) ERROR 06-01 23:50:49 [multiproc_executor.py:966]     super().capture(create_forward_fn, progress_bar_desc)
(Worker_DP3_EP3 pid=2845516) ERROR 06-01 23:50:49 [multiproc_executor.py:966]   File "/home/gdelfin/vllm/.venv/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 124, in decorate_context
(Worker_DP3_EP3 pid=2845516) ERROR 06-01 23:50:49 [multiproc_executor.py:966]     return func(*args, **kwargs)
(Worker_DP3_EP3 pid=2845516) ERROR 06-01 23:50:49 [multiproc_executor.py:966]            ^^^^^^^^^^^^^^^^^^^^^
(Worker_DP3_EP3 pid=2845516) ERROR 06-01 23:50:49 [multiproc_executor.py:966]   File "/home/gdelfin/vllm/vllm/v1/worker/gpu/cudagraph_utils.py", line 255, in capture
(Worker_DP3_EP3 pid=2845516) ERROR 06-01 23:50:49 [multiproc_executor.py:966]     with torch.cuda.graph(graph, self.pool):
(Worker_DP3_EP3 pid=2845516) ERROR 06-01 23:50:49 [multiproc_executor.py:966]          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(Worker_DP3_EP3 pid=2845516) ERROR 06-01 23:50:49 [multiproc_executor.py:966]   File "/home/gdelfin/vllm/.venv/lib/python3.12/site-packages/torch/cuda/graphs.py", line 239, in __enter__
(Worker_DP3_EP3 pid=2845516) ERROR 06-01 23:50:49 [multiproc_executor.py:966]     torch.cuda.synchronize()
(Worker_DP3_EP3 pid=2845516) ERROR 06-01 23:50:49 [multiproc_executor.py:966]   File "/home/gdelfin/vllm/.venv/lib/python3.12/site-packages/torch/cuda/__init__.py", line 1162, in synchronize
(Worker_DP3_EP3 pid=2845516) ERROR 06-01 23:50:49 [multiproc_executor.py:966]     return torch._C._cuda_synchronize()
(Worker_DP3_EP3 pid=2845516) ERROR 06-01 23:50:49 [multiproc_executor.py:966]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(Worker_DP3_EP3 pid=2845516) ERROR 06-01 23:50:49 [multiproc_executor.py:966] torch.AcceleratorError: CUDA error: unspecified launch failure
(Worker_DP3_EP3 pid=2845516) ERROR 06-01 23:50:49 [multiproc_executor.py:966] Search for `cudaErrorLaunchFailure' in https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__TYPES.html for more information.
(Worker_DP3_EP3 pid=2845516) ERROR 06-01 23:50:49 [multiproc_executor.py:966] CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
(Worker_DP3_EP3 pid=2845516) ERROR 06-01 23:50:49 [multiproc_executor.py:966] For debugging consider passing CUDA_LAUNCH_BLOCKING=1
(Worker_DP3_EP3 pid=2845516) ERROR 06-01 23:50:49 [multiproc_executor.py:966] Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

This PR

Fixes the above issue by modifying CudaGraphManager.capture to return, for each batch descriptor, the attention states used for both warmup and capture. These are passed into PrefillEagleCudaGraphManager and the correct one is used for the speculator forward pass based on the warmup parameter. This ensures that we don't invoke the warmup and capture forward passes on the same attention metadata.

Benchmarks

Server Command

#!/bin/bash

export VLLM_USE_V2_MODEL_RUNNER=1
export VLLM_DEEP_GEMM_WARMUP="skip"

MODEL_CKPT="deepseek-ai/DeepSeek-V4-Flash"

vllm serve $MODEL_CKPT \
    -dp 4 -ep \
    --port 8003 \
    --block-size 256 \
    --kv-cache-dtype fp8 \
    --tokenizer-mode deepseek_v4 \
    --tool-call-parser deepseek_v4 \
    --reasoning-parser deepseek_v4 \
    --moe-backend deep_gemm_mega_moe \
    --speculative-config '{"method": "mtp", "num_speculative_tokens": 3}' \
    --no-enable-prefix-caching \
    --kernel-config.enable_flashinfer_autotune=False

Benchmark Command

vllm-bench --backend openai-chat \
     --base-url http://127.0.0.1:8003 \
     --model deepseek-ai/DeepSeek-V4-Flash\
     --dataset-name speed-bench \
     --speed-bench-config throughput_16k \
     --speed-bench-max-input-len 8192 \
     --speed-bench-category low_entropy \
     --output-len 1024 \
     --save-result \
     --num-warmups 1 \
     --seed 420 \
     --sweep-max-concurrency 1,2,1024 \
     --sweep-num-prompts-factor 5 \
     --temperature 0 \
     --top-p 1

Results

Metric	Concurrency = 1	Concurrency = 2	Concurrency = 1024
Benchmark duration (s)	14.60	17.14	390.02
Total input tokens	40,960	81,920	41,943,040
Total generated tokens	3,688	7,752	2,851,158
Request throughput (req/s)	0.34	0.58	13.13
Output token throughput (tok/s)	252.64	452.38	7,310.21
Peak output token throughput (tok/s)	96.00	187.00	16,893.00
Total token throughput (tok/s)	3,058.57	5,232.93	114,849.87
Mean TTFT (ms)	264.37	296.63	22,956.43
Median TTFT (ms)	264.00	258.02	19,345.86
P90 TTFT (ms)	270.02	444.74	38,323.11
P99 TTFT (ms)	273.17	458.13	61,000.88
Mean TPOT (ms)	3.63	4.04	89.74
Median TPOT (ms)	3.68	3.86	77.74
P90 TPOT (ms)	3.74	4.29	114.70
P99 TPOT (ms)	3.74	5.55	315.51
Mean ITL (ms)	10.59	11.39	250.53
Median ITL (ms)	10.50	10.85	208.79
P90 ITL (ms)	11.02	10.95	240.82
P99 ITL (ms)	15.46	13.35	299.93
Mean E2EL (ms)	2,919.42	3,314.18	70,754.69
Median E2EL (ms)	2,996.37	4,155.82	65,880.41
P90 E2EL (ms)	3,850.75	4,444.94	117,805.97
P99 E2EL (ms)	3,852.71	4,487.04	149,902.46
Acceptance rate (%)	65.22	64.62	64.46
Acceptance length	2.96	2.94	2.93
Drafts	1,248	2,639	971,705
Draft tokens	3,744	7,917	2,915,115
Accepted tokens	2,442	5,116	1,879,132
Position 0 acceptance (%)	94.95	95.60	94.56
Position 1 acceptance (%)	68.83	67.49	68.13
Position 2 acceptance (%)	31.89	30.77	30.69

No more startup failure!