[Bugfix] Capture prefill kernels for mixed FULL graphs (#58275)
Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Andreas Karatzas committed
Sep 24, 2026 at 04:05 UTC
23110a0c9ef8823406f04f919fcd61e030da8b81
4 files changed
+117
-16
tests/compile/fullgraph/test_full_cudagraph.py
+25
-15
@@ -36,7 +36,7 @@ def temporary_environ(env_vars):
36
model_backends_full_cudagraph = []
37
38
# deepseek-ai/DeepSeek-V2-Lite with MLA
39
-MLA_backends = ["FlashMLA", "FlashAttentionMLA", "CutlassMLA"]
39
+MLA_backends = ["FlashMLA", "FlashAttentionMLA", "CutlassMLA", "FlashInferMLA"]
40
for mla_backend in MLA_backends:
41
model_backends_full_cudagraph.append(
42
("deepseek-ai/DeepSeek-V2-Lite", backend_configs[mla_backend])
@@ -53,27 +53,28 @@ for backend_config in other_backend_configs:
53
@pytest.fixture(scope="class")
54
def llm_pair(request):
55
model, backend_config, use_inductor_graph_partition = request.param
56
- backend_config.comp_config["use_inductor_graph_partition"] = (
57
- use_inductor_graph_partition
58
- )
59
-
56
if use_inductor_graph_partition and not is_torch_equal_or_newer("2.9.0.dev"):
57
pytest.skip("Inductor graph partition only supported in torch>=2.9")
58
59
+ backend = AttentionBackendEnum[backend_config.attention_config["backend"]]
60
+ if current_platform.is_rocm() and backend not in (
61
+ AttentionBackendEnum.TRITON_ATTN,
62
+ AttentionBackendEnum.ROCM_ATTN,
63
+ ):
64
+ pytest.skip(f"{backend_config.name} is a CUDA-only attention backend")
65
+ if backend == AttentionBackendEnum.ROCM_ATTN and not current_platform.is_rocm():
66
+ pytest.skip("ROCM_ATTN requires ROCm")
67
+
68
# Dynamically skip test if GPU capability is not met
64
- if (
65
- backend_config.specific_gpu_arch
66
- and backend_config.specific_gpu_arch != current_platform.get_device_capability()
69
+ if backend_config.specific_gpu_arch and (
70
+ not current_platform.is_cuda()
71
+ or backend_config.specific_gpu_arch != current_platform.get_device_capability()
72
):
73
if backend_config.specific_gpu_arch == (9, 0):
74
pytest.skip("Only Hopper GPUs support FA3 and FlashMLA")
75
elif backend_config.specific_gpu_arch == (10, 0):
76
pytest.skip("Only Blackwell GPUs support Cutlass MLA")
77
73
- # FlashInfer is not supported on ROCm
74
- if backend_config == AttentionBackendEnum.FLASHINFER and current_platform.is_rocm():
75
- pytest.skip("FlashInfer is not supported on ROCm")
76
-
78
env_vars = {
79
# Force native sampler to avoid potential nondeterminism in FlashInfer
80
# when per-request generators are not used in V1.
@@ -86,7 +87,11 @@ def llm_pair(request):
87
trust_remote_code=True,
88
max_model_len=1024,
89
max_num_seqs=128,
89
- compilation_config=CompilationConfig(**backend_config.comp_config),
90
+ attention_config=backend_config.attention_config,
91
+ compilation_config=CompilationConfig(
92
+ **backend_config.comp_config,
93
+ use_inductor_graph_partition=use_inductor_graph_partition,
94
+ ),
95
generation_config="vllm",
96
seed=42,
97
)
@@ -96,8 +101,10 @@ def llm_pair(request):
101
trust_remote_code=True,
102
max_model_len=1024,
103
max_num_seqs=128,
104
+ attention_config=backend_config.attention_config,
105
compilation_config=CompilationConfig(
100
- cudagraph_mode=CUDAGraphMode.PIECEWISE
106
+ cudagraph_mode=CUDAGraphMode.PIECEWISE,
107
+ use_inductor_graph_partition=use_inductor_graph_partition,
108
),
109
generation_config="vllm",
110
seed=42,
@@ -117,7 +124,10 @@ def llm_pair(request):
124
@pytest.mark.parametrize(
125
"llm_pair",
126
[
120
- pytest.param((model, backend_config, use_inductor_graph_partition))
127
+ pytest.param(
128
+ (model, backend_config, use_inductor_graph_partition),
129
+ id=f"{backend_config.name}-partition={use_inductor_graph_partition}",
130
+ )
131
for model, backend_config in model_backends_full_cudagraph
132
for use_inductor_graph_partition in [True, False]
133
],
tests/v1/cudagraph/test_cudagraph_manager.py
+86
@@ -22,6 +22,9 @@ from vllm.v1.worker.gpu import cudagraph_utils as gpu_cudagraph_utils
22
from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
23
from vllm.v1.worker.gpu.input_batch import InputBuffers
24
from vllm.v1.worker.gpu.pcp_manager import PCPManager
25
+from vllm.v1.worker.gpu.spec_decode.autoregressive.cudagraph_utils import (
26
+ SpeculatorCudaGraphManager,
27
+)
28
29
pytestmark = pytest.mark.cpu_test
30
@@ -165,6 +168,89 @@ def test_piecewise_capture_uses_pcp_dummy_slot_mappings():
168
block_tables.get_dummy_slot_mappings.assert_not_called()
169
170
171
+@pytest.mark.parametrize(
172
+ "full_cudagraph,max_query_len,expected_max_query_len",
173
+ [(False, None, None), (True, None, 8), (True, 1, 1), (True, 3, 3)],
174
+)
175
+def test_capture_query_length_covers_replay_batches(
176
+ full_cudagraph, max_query_len, expected_max_query_len
177
+):
178
+ """Mixed FULL graphs must capture prefill kernels as well as decode kernels."""
179
+ num_tokens = num_reqs = 8
180
+ buffers = InputBuffers(num_reqs, num_tokens, torch.device("cpu"))
181
+ block_tables = MagicMock()
182
+ block_tables.cp_size = 1
183
+ block_tables.get_dummy_block_tables.return_value = ()
184
+ block_tables.get_dummy_slot_mappings.return_value = torch.empty(
185
+ 0, num_tokens, dtype=torch.int64
186
+ )
187
+ model_state = MagicMock()
188
+ kv_cache_config = KVCacheConfig(
189
+ num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[]
190
+ )
191
+
192
+ gpu_cudagraph_utils.prepare_inputs_to_capture(
193
+ num_reqs,
194
+ num_tokens,
195
+ model_state,
196
+ buffers,
197
+ block_tables,
198
+ [],
199
+ kv_cache_config,
200
+ full_cudagraph=full_cudagraph,
201
+ max_query_len=max_query_len,
202
+ )
203
+
204
+ batch = model_state.prepare_attn.call_args.args[0]
205
+ assert batch.num_scheduled_tokens.max() == 1
206
+ assert batch.max_query_len == expected_max_query_len
207
+
208
+
209
+@pytest.mark.parametrize(
210
+ "uniform_token_count,max_query_len,expected_max_query_len",
211
+ [(1, None, 1), (None, 3, 3), (None, None, 8)],
212
+)
213
+def test_speculator_capture_preserves_decode_query_bounds(
214
+ monkeypatch, uniform_token_count, max_query_len, expected_max_query_len
215
+):
216
+ """Draft decode graphs keep narrow queries while mixed graphs cover prefill."""
217
+ num_tokens = num_reqs = 8
218
+ desc = BatchExecutionDescriptor(
219
+ CUDAGraphMode.FULL,
220
+ num_tokens,
221
+ num_reqs,
222
+ uniform_token_count=uniform_token_count,
223
+ max_query_len=max_query_len,
224
+ )
225
+ manager = SpeculatorCudaGraphManager.__new__(SpeculatorCudaGraphManager)
226
+ manager.max_num_reqs = num_reqs
227
+ manager.dp_size = 1
228
+ buffers = InputBuffers(num_reqs, num_tokens, torch.device("cpu"))
229
+ block_tables = MagicMock()
230
+ block_tables.cp_size = 1
231
+ block_tables.get_dummy_block_tables.return_value = ()
232
+ block_tables.get_dummy_slot_mappings.return_value = torch.empty(
233
+ 0, num_tokens, dtype=torch.int64
234
+ )
235
+ model_state = MagicMock()
236
+ kv_cache_config = KVCacheConfig(
237
+ num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[]
238
+ )
239
+
240
+ def capture(self, create_forward_fn, progress_bar_desc):
241
+ for warmup in (True, False):
242
+ create_forward_fn(desc, warmup)(desc.cg_mode)
243
+
244
+ monkeypatch.setattr(gpu_cudagraph_utils.CudaGraphManager, "capture", capture)
245
+ manager.capture(
246
+ MagicMock(), model_state, buffers, block_tables, [], kv_cache_config
247
+ )
248
+
249
+ assert model_state.prepare_attn.call_count == 2
250
+ for call in model_state.prepare_attn.call_args_list:
251
+ assert call.args[0].max_query_len == expected_max_query_len
252
+
253
+
254
_DECODE_QUERY_LEN = 3
255
256
vllm/v1/worker/gpu/cudagraph_utils.py
+5
-1
@@ -704,7 +704,7 @@ class ModelCudaGraphManager(CudaGraphManager):
704
attn_groups,
705
kv_cache_config,
706
full_cudagraph=desc.cg_mode == CUDAGraphMode.FULL,
707
- max_query_len=desc.max_query_len,
707
+ max_query_len=desc.max_query_len or desc.uniform_token_count,
708
pcp_manager=pcp_manager,
709
)
710
@@ -775,6 +775,10 @@ def prepare_inputs_to_capture(
775
max_query_len: int | None = None,
776
pcp_manager: "PCPManager | None" = None,
777
) -> AttentionState:
778
+ if full_cudagraph and max_query_len is None:
779
+ # Mixed graphs can replay a single prefill spanning the entire batch,
780
+ # even when the dummy batch distributes one token to each request.
781
+ max_query_len = num_tokens
782
input_batch = InputBatch.make_dummy(
783
num_reqs, num_tokens, input_buffers, max_query_len=max_query_len
784
)
vllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.py
+1
@@ -57,6 +57,7 @@ class SpeculatorCudaGraphManager(CudaGraphManager):
57
attn_groups,
58
kv_cache_config,
59
full_cudagraph=desc.cg_mode == CUDAGraphMode.FULL,
60
+ max_query_len=desc.max_query_len or desc.uniform_token_count,
61
)
62
63
return lambda cg_mode: forward_fn(