A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 9, 18:15 UTC
[cudagraphs] Refactor cudagraph capture loop (#32946)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Lucas Wilkinson committed
Jan 23, 2026 at 13:22 UTC
3a4145950176d966558d06ed5b05b29c489555f0
3 files changed
+117
-59
tests/v1/cudagraph/test_cudagraph_dispatch.py
+62
@@ -173,6 +173,68 @@ class TestCudagraphDispatcher:
173
else:
174
assert rt_mode == CUDAGraphMode.NONE
175
176
+ @pytest.mark.parametrize(
177
+ "cudagraph_mode_str,compilation_mode,expected_modes",
178
+ [
179
+ # FULL mode: only FULL keys, no PIECEWISE
180
+ ("FULL", CompilationMode.NONE, [CUDAGraphMode.FULL]),
181
+ # PIECEWISE mode: only PIECEWISE keys
182
+ ("PIECEWISE", CompilationMode.VLLM_COMPILE, [CUDAGraphMode.PIECEWISE]),
183
+ # FULL_DECODE_ONLY: only FULL keys for uniform decode
184
+ ("FULL_DECODE_ONLY", CompilationMode.NONE, [CUDAGraphMode.FULL]),
185
+ # NONE mode: no keys
186
+ ("NONE", CompilationMode.NONE, []),
187
+ ],
188
+ )
189
+ def test_get_capture_descs(
190
+ self, cudagraph_mode_str, compilation_mode, expected_modes
191
+ ):
192
+ """Test get_capture_descs returns correctly grouped and ordered descs."""
193
+ comp_config = CompilationConfig(
194
+ cudagraph_mode=cudagraph_mode_str,
195
+ mode=compilation_mode,
196
+ cudagraph_capture_sizes=[1, 4, 8, 16],
197
+ )
198
+
199
+ config = _create_vllm_config(comp_config, max_num_seqs=16)
200
+ dispatcher = CudagraphDispatcher(config)
201
+ dispatcher.initialize_cudagraph_keys(
202
+ cudagraph_mode=comp_config.cudagraph_mode, uniform_decode_query_len=1
203
+ )
204
+
205
+ capture_descs = dispatcher.get_capture_descs()
206
+
207
+ # Verify we get the expected modes
208
+ actual_modes = [mode for mode, _ in capture_descs]
209
+ assert actual_modes == expected_modes
210
+
211
+ # Verify each group is sorted largest-first
212
+ for mode, descs in capture_descs:
213
+ assert len(descs) > 0, "Each group should have at least one descriptor"
214
+ num_tokens_list = [d.num_tokens for d in descs]
215
+ assert num_tokens_list == sorted(num_tokens_list, reverse=True), (
216
+ f"Descriptors for {mode} should be sorted largest-first"
217
+ )
218
+
219
+ # All descriptors in a group should have same uniform value
220
+ uniform_values = [d.uniform for d in descs]
221
+ assert len(set(uniform_values)) == 1, (
222
+ "All descriptors in a group should have the same uniform value"
223
+ )
224
+
225
+ def test_get_capture_descs_empty_when_not_initialized(self):
226
+ """Test that get_capture_descs returns empty list when keys not initialized."""
227
+ comp_config = CompilationConfig(
228
+ cudagraph_mode="FULL",
229
+ mode=CompilationMode.NONE,
230
+ cudagraph_capture_sizes=[1, 8],
231
+ )
232
+ config = _create_vllm_config(comp_config, max_num_seqs=8)
233
+ dispatcher = CudagraphDispatcher(config)
234
+ # Don't initialize keys
235
+
236
+ assert dispatcher.get_capture_descs() == []
237
+
238
239
@pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda")
240
class TestCUDAGraphWrapper:
vllm/v1/cudagraph_dispatcher.py
+23
@@ -231,3 +231,26 @@ class CudagraphDispatcher:
231
232
# finally, just return no cudagraphs and a trivial batch descriptor
233
return CUDAGraphMode.NONE, BatchDescriptor(num_tokens)
234
+
235
+ def get_capture_descs(self) -> list[tuple[CUDAGraphMode, list[BatchDescriptor]]]:
236
+ """
237
+ Returns capture descriptors for cudagraph capturing.
238
+
239
+ Returns:
240
+ List of (runtime_mode, batch_descriptors) tuples, ordered PIECEWISE
241
+ first then FULL. Batch descriptors are sorted largest-first for
242
+ memory efficiency.
243
+ """
244
+ if not self.keys_initialized or self.cudagraph_mode == CUDAGraphMode.NONE:
245
+ return []
246
+
247
+ result = []
248
+ # Return in order: PIECEWISE first, then FULL
249
+ for mode in [CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL]:
250
+ descs = list(self.cudagraph_keys[mode])
251
+ if descs:
252
+ # Sort by num_tokens descending (largest first)
253
+ descs.sort(key=lambda d: d.num_tokens, reverse=True)
254
+ result.append((mode, descs))
255
+
256
+ return result
vllm/v1/worker/gpu_model_runner.py
+32
-59
@@ -10,7 +10,6 @@ from collections.abc import Iterator, Sequence
10
from contextlib import contextmanager
11
from copy import copy, deepcopy
12
from functools import reduce
13
-from itertools import product
13
from typing import TYPE_CHECKING, Any, NamedTuple, TypeAlias, cast
14
15
import numpy as np
@@ -4839,50 +4838,14 @@ class GPUModelRunner(
4838
set_cudagraph_capturing_enabled(True)
4839
with freeze_gc(), graph_capture(device=self.device):
4840
start_free_gpu_memory = torch.cuda.mem_get_info()[0]
4842
- cudagraph_mode = self.compilation_config.cudagraph_mode
4843
- assert cudagraph_mode is not None
4841
4845
- if self.lora_config:
4846
- if self.compilation_config.cudagraph_specialize_lora:
4847
- lora_cases = [True, False]
4848
- else:
4849
- lora_cases = [True]
4850
- else:
4851
- lora_cases = [False]
4852
-
4853
- if cudagraph_mode.mixed_mode() != CUDAGraphMode.NONE:
4854
- cudagraph_runtime_mode = cudagraph_mode.mixed_mode()
4855
- # make sure we capture the largest batch size first
4856
- compilation_cases = list(
4857
- product(reversed(self.cudagraph_batch_sizes), lora_cases)
4858
- )
4859
- self._capture_cudagraphs(
4860
- compilation_cases,
4861
- cudagraph_runtime_mode=cudagraph_runtime_mode,
4862
- uniform_decode=False,
4863
- )
4864
-
4865
- # Capture full cudagraph for uniform decode batches if we
4866
- # don't already have full mixed prefill-decode cudagraphs.
4867
- if (
4868
- cudagraph_mode.decode_mode() == CUDAGraphMode.FULL
4869
- and cudagraph_mode.separate_routine()
4870
- ):
4871
- max_num_tokens = (
4872
- self.scheduler_config.max_num_seqs * self.uniform_decode_query_len
4873
- )
4874
- decode_cudagraph_batch_sizes = [
4875
- x
4876
- for x in self.cudagraph_batch_sizes
4877
- if max_num_tokens >= x >= self.uniform_decode_query_len
4878
- ]
4879
- compilation_cases_decode = list(
4880
- product(reversed(decode_cudagraph_batch_sizes), lora_cases)
4881
- )
4842
+ for (
4843
+ runtime_mode,
4844
+ batch_descs,
4845
+ ) in self.cudagraph_dispatcher.get_capture_descs():
4846
self._capture_cudagraphs(
4883
- compilation_cases=compilation_cases_decode,
4884
- cudagraph_runtime_mode=CUDAGraphMode.FULL,
4885
- uniform_decode=True,
4847
+ batch_descriptors=batch_descs,
4848
+ cudagraph_runtime_mode=runtime_mode,
4849
)
4850
4851
torch.cuda.synchronize()
@@ -4913,19 +4876,32 @@ class GPUModelRunner(
4876
4877
def _capture_cudagraphs(
4878
self,
4916
- compilation_cases: list[tuple[int, bool]],
4879
+ batch_descriptors: list[BatchDescriptor],
4880
cudagraph_runtime_mode: CUDAGraphMode,
4918
- uniform_decode: bool,
4881
):
4882
assert (
4883
cudagraph_runtime_mode != CUDAGraphMode.NONE
4884
and cudagraph_runtime_mode.valid_runtime_modes()
4885
), f"Invalid cudagraph runtime mode: {cudagraph_runtime_mode}"
4886
4887
+ if not batch_descriptors:
4888
+ return
4889
+
4890
+ uniform_decode = batch_descriptors[0].uniform
4891
+ force_attention = cudagraph_runtime_mode == CUDAGraphMode.FULL
4892
+
4893
+ dummy_run = functools.partial(
4894
+ self._dummy_run,
4895
+ uniform_decode=uniform_decode,
4896
+ skip_eplb=True,
4897
+ remove_lora=False,
4898
+ force_attention=force_attention,
4899
+ )
4900
+
4901
# Only rank 0 should print progress bar during capture
4902
if is_global_first_rank():
4927
- compilation_cases = tqdm(
4928
- compilation_cases,
4903
+ batch_descriptors = tqdm(
4904
+ batch_descriptors,
4905
disable=not self.load_config.use_tqdm_on_load,
4906
desc="Capturing CUDA graphs ({}, {})".format(
4907
"decode" if uniform_decode else "mixed prefill-decode",
@@ -4934,7 +4910,10 @@ class GPUModelRunner(
4910
)
4911
4912
# We skip EPLB here since we don't want to record dummy metrics
4937
- for num_tokens, activate_lora in compilation_cases:
4913
+ for batch_desc in batch_descriptors:
4914
+ num_tokens = batch_desc.num_tokens
4915
+ activate_lora = batch_desc.has_lora
4916
+
4917
# We currently only capture ubatched graphs when its a FULL
4918
# cudagraph, a uniform decode batch, and the number of tokens
4919
# is above the threshold. Otherwise we just capture a non-ubatched
@@ -4952,28 +4931,22 @@ class GPUModelRunner(
4931
4932
for _ in range(self.compilation_config.cudagraph_num_of_warmups):
4933
# Use CUDAGraphRuntimeStyle.NONE (default) for warmup.
4955
- # But be careful, warm up with `NONE`is orthogonal to
4934
+ # But be careful, warm up with `NONE` is orthogonal to
4935
# if we want to warm up attention or not. This is
4936
# different from the case where `FULL` implies capture
4937
# attention while `PIECEWISE` implies no attention.
4959
- force_attention = cudagraph_runtime_mode == CUDAGraphMode.FULL
4960
- self._dummy_run(
4938
+ dummy_run(
4939
num_tokens,
4940
cudagraph_runtime_mode=CUDAGraphMode.NONE,
4963
- force_attention=force_attention,
4964
- uniform_decode=uniform_decode,
4941
allow_microbatching=allow_microbatching,
4966
- skip_eplb=True,
4967
- remove_lora=False,
4942
activate_lora=activate_lora,
4943
)
4970
- self._dummy_run(
4944
+
4945
+ # Capture run
4946
+ dummy_run(
4947
num_tokens,
4948
cudagraph_runtime_mode=cudagraph_runtime_mode,
4973
- uniform_decode=uniform_decode,
4949
allow_microbatching=allow_microbatching,
4975
- skip_eplb=True,
4976
- remove_lora=False,
4950
activate_lora=activate_lora,
4951
is_graph_capturing=True,
4952
)