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 12, 17:15 UTC
[Model Runner V2] Limit cudagraph size to max decode batch size (#29221)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Nov 21, 2025 at 20:21 UTC
e9056056fbacecbac4318bd0323745fdd7fe55b6
1 file changed
+9
-4
vllm/v1/worker/gpu/cudagraph_utils.py
+9
-4
@@ -27,9 +27,11 @@ class CudaGraphManager:
27
device: torch.device,
28
):
29
self.vllm_config = vllm_config
30
+ self.scheduler_config = vllm_config.scheduler_config
31
self.device = device
32
33
self.max_model_len = vllm_config.model_config.max_model_len
34
+ self.max_num_reqs = self.scheduler_config.max_num_seqs
35
self.dp_size = vllm_config.parallel_config.data_parallel_size
36
self.compilation_config = vllm_config.compilation_config
37
assert self.compilation_config is not None
@@ -39,9 +41,11 @@ class CudaGraphManager:
41
else:
42
self.cudagraph_mode = self.compilation_config.cudagraph_mode
43
if self.compilation_config.cudagraph_capture_sizes is not None:
42
- self.cudagraph_sizes = sorted(
43
- self.compilation_config.cudagraph_capture_sizes
44
- )
44
+ cudagraph_sizes = sorted(self.compilation_config.cudagraph_capture_sizes)
45
+ # Limit the cudagraph sizes to the max decode batch size.
46
+ self.cudagraph_sizes = [
47
+ x for x in cudagraph_sizes if x <= self.max_num_reqs
48
+ ]
49
else:
50
self.cudagraph_sizes = []
51
self.padded_sizes = self._init_padded_sizes()
@@ -54,9 +58,10 @@ class CudaGraphManager:
58
if not self.cudagraph_mode.has_full_cudagraphs():
59
# Full cuda graphs are not used.
60
return {}
61
+ if not self.cudagraph_sizes:
62
+ return {}
63
64
padded_sizes: dict[int, int] = {}
59
- assert len(self.cudagraph_sizes) > 0
65
for i in range(1, self.cudagraph_sizes[-1] + 1):
66
for x in self.cudagraph_sizes:
67
if i <= x: