@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 284cc92275
Fork 0

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 14, 05:30 UTC

[MISC] `cudagraph_capture_sizes` related improvements (#26016)

Signed-off-by: fhl <2410591650@qq.com> Signed-off-by: fhl2000 <63384265+fhl2000@users.noreply.github.com> Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

fhl2000 committed Oct 24, 2025 at 20:11 UTC 284cc922757b0494a5920f70d4e31feeb4bdfea0
14 files changed +303 -110
tests/compile/test_config.py
+73
@@ -1,6 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import copy
4 +from contextlib import nullcontext
5
6 import pytest
7
@@ -8,6 +9,8 @@ from vllm.compilation.counter import compilation_counter
9 from vllm.compilation.fix_functionalization import FixFunctionalizationPass
10 from vllm.config import CompilationConfig, CUDAGraphMode, VllmConfig
11 from vllm.config.compilation import CompilationMode
12 +from vllm.engine.arg_utils import EngineArgs
13 +from vllm.platforms import current_platform
14 from vllm.utils.torch_utils import _is_torch_equal_or_newer, is_torch_equal_or_newer
15
16
@@ -233,3 +236,73 @@ def test_resolve_operator_overload():
236 assert len(resolved) == 2 # Only 2 valid ops
237 assert resolved[0] is torch.ops.aten.mm.default
238 assert resolved[1] is torch.ops.aten.addmm.default
239 +
240 +
241 +@pytest.mark.skipif(
242 + not current_platform.support_static_graph_mode(),
243 + reason="Skip if not cudagraph mode supported",
244 +)
245 +@pytest.mark.parametrize(
246 + (
247 + "cudagraph_capture_sizes",
248 + "max_cudagraph_capture_size",
249 + "tp_size",
250 + "enable_sequence_parallelism",
251 + "max_num_batched_tokens",
252 + "use_cudagraph",
253 + "expected_max_size",
254 + ),
255 + [
256 + (None, None, 1, False, 2048, True, 512),
257 + ([1, 2, 4], 4, 1, False, 2048, True, 4),
258 + ([1, 2, 4], 8, 1, False, 2048, True, RuntimeError),
259 + ([1, 256], None, 1, False, 2048, 256),
260 + ([], None, 1, False, 2048, False, 0),
261 + (None, 0, 1, False, 2048, False, 0),
262 + # truncated to nearest multiple of 8 or 16
263 + (None, 257, 1, False, 2048, True, 256),
264 + ([1, 2, 4, 15], None, 1, False, 2048, True, 15), # max from list
265 + ([1, 2, 4, 15], None, 2, True, 2048, True, 4), # filtered out 15 due to SP
266 + ([1, 2, 4, 15], None, 1, False, 8, True, 4), # limited by the max_tokens
267 + # the list should contain at least 1 element when use cudagraph
268 + ([], None, 1, False, 2048, True, RuntimeError),
269 + # the max capturing size should be >= 1 when use cudagraph
270 + (None, 0, 1, False, 2048, True, RuntimeError),
271 + ],
272 +)
273 +def test_cudagraph_sizes_post_init(
274 + cudagraph_capture_sizes,
275 + max_cudagraph_capture_size,
276 + tp_size,
277 + enable_sequence_parallelism,
278 + max_num_batched_tokens,
279 + use_cudagraph,
280 + expected_max_size,
281 +):
282 + ctx = nullcontext()
283 + if isinstance(expected_max_size, Exception):
284 + ctx = pytest.raises(expected_max_size)
285 +
286 + cudagraph_mode = CUDAGraphMode.PIECEWISE if use_cudagraph else CUDAGraphMode.NONE
287 + with ctx:
288 + compilation_config = CompilationConfig(
289 + cudagraph_capture_sizes=cudagraph_capture_sizes,
290 + max_cudagraph_capture_size=max_cudagraph_capture_size,
291 + pass_config={
292 + "enable_sequence_parallelism": enable_sequence_parallelism,
293 + "enable_fusion": True,
294 + "enable_noop": True,
295 + },
296 + cudagraph_mode=cudagraph_mode,
297 + )
298 + engine_args = EngineArgs(
299 + model="facebook/opt-125m",
300 + tensor_parallel_size=tp_size,
301 + max_num_batched_tokens=max_num_batched_tokens,
302 + compilation_config=compilation_config,
303 + )
304 + vllm_config = engine_args.create_engine_config()
305 +
306 + assert (
307 + vllm_config.compilation_config.max_cudagraph_capture_size == expected_max_size
308 + )
vllm/config/compilation.py
+40 -38
@@ -154,6 +154,8 @@ class CompilationConfig:
154 - [`cudagraph_mode`][vllm.config.CompilationConfig.cudagraph_mode]
155 - [`cudagraph_capture_sizes`]
156 [vllm.config.CompilationConfig.cudagraph_capture_sizes]
157 + - [`max_cudagraph_capture_size`]
158 + [vllm.config.CompilationConfig.max_cudagraph_capture_size]
159 - [`cudagraph_num_of_warmups`]
160 [vllm.config.CompilationConfig.cudagraph_num_of_warmups]
161 - [`cudagraph_copy_inputs`]
@@ -327,18 +329,16 @@ class CompilationConfig:
329 more modes may be added.
330 """
331 use_cudagraph: bool = True
330 - """Whether to use cudagraph inside compilation.
331 - - False: cudagraph inside compilation is not used.
332 + """Whether to use cudagraph inside compilation:
333 +
334 + - False: cudagraph inside compilation is not used.\n
335 - True: cudagraph inside compilation is used. It requires
336 that all input buffers have fixed addresses, and all
337 splitting ops write their outputs to input buffers.
335 - In the vLLM V1 Engine, this flag only applies for
336 - CompilationMode.VLLM_COMPILE (aka -O3).
337 - Note that this is orthogonal to the cudagraph capture logic
338 - outside of compilation.
338 +
339 Warning: This flag is deprecated and will be removed in the next major or
340 - minor release, i.e. v0.11.0 or v1.0.0. Please use cudagraph_mode=PIECEWISE
341 - instead.
340 + minor release, i.e. v0.11.0 or v1.0.0. Please use cudagraph_mode=FULL_AND
341 + _PIECEWISE instead.
342 """
343 cudagraph_num_of_warmups: int = 0
344 """Number of warmup runs for cudagraph.
@@ -398,8 +398,22 @@ class CompilationConfig:
398 pass_config: PassConfig = field(default_factory=PassConfig)
399 """Custom inductor passes, see PassConfig for more details"""
400
401 - max_capture_size: int = field(default=None, init=False) # type: ignore
402 - """not configurable, computed after init"""
401 + max_cudagraph_capture_size: int | None = field(default=None)
402 + """The maximum cudagraph capture size.
403 +
404 + If cudagraph_capture_sizes is specified, this will be set to the largest
405 + size in that list (or checked for consistency if specified). If
406 + cudagraph_capture_sizes is not specified, the list of sizes is generated
407 + automatically following the pattern:
408 +
409 + [1, 2, 4] + list(range(8, 256, 8)) + list(
410 + range(256, max_cudagraph_capture_size + 1, 16))
411 +
412 + If not specified, max_cudagraph_capture_size is set to min(max_num_seqs*2,
413 + 512) by default. This voids OOM in tight memory scenarios with small
414 + max_num_seqs, and prevents capture of many large graphs (>512) that would
415 + greatly increase startup time with limited performance benefit.
416 + """
417 local_cache_dir: str = field(default=None, init=False) # type: ignore
418 """local cache dir for each rank"""
419 bs_to_padded_graph_size: list[int] = field(
@@ -408,7 +422,7 @@ class CompilationConfig:
422 )
423 """optimization:
424 Intuitively, bs_to_padded_graph_size should be dict[int, int].
411 - since we know all keys are in a range [0, max_capture_size],
425 + since we know all keys are in a range [0, max_cudagraph_capture_size],
426 we can optimize it to list[int] for better lookup performance."""
427
428 # keep track of enabled and disabled custom ops
@@ -672,25 +686,12 @@ class CompilationConfig:
686
687 return VllmBackend(vllm_config)
688
675 - def init_with_cudagraph_sizes(self, cudagraph_capture_sizes: list[int]) -> None:
676 - """To complete the initialization of config,
677 - we need to know the cudagraph sizes."""
678 -
679 - if self.cudagraph_capture_sizes is None:
680 - self.cudagraph_capture_sizes = cudagraph_capture_sizes
681 - else:
682 - # de-duplicate the sizes provided by the config
683 - dedup_sizes = list(set(self.cudagraph_capture_sizes))
684 - if len(dedup_sizes) < len(self.cudagraph_capture_sizes):
685 - logger.info(
686 - (
687 - "cudagraph sizes specified by model runner"
688 - " %s is overridden by config %s"
689 - ),
690 - cudagraph_capture_sizes,
691 - dedup_sizes,
692 - )
693 - self.cudagraph_capture_sizes = dedup_sizes
689 + def post_init_cudagraph_sizes(self) -> None:
690 + """To complete the initialization after cudagraph related
691 + configs are set. This includes:
692 + - initialize compile_sizes
693 + - pre-compute the mapping bs_to_padded_graph_size
694 + """
695
696 computed_compile_sizes = []
697 if self.compile_sizes is not None:
@@ -708,23 +709,24 @@ class CompilationConfig:
709 computed_compile_sizes.append(x)
710 self.compile_sizes = computed_compile_sizes # type: ignore
711
711 - # sort to make sure cudagraph capture sizes are in descending order
712 - self.cudagraph_capture_sizes.sort(reverse=True)
713 - self.max_capture_size = (
714 - self.cudagraph_capture_sizes[0] if self.cudagraph_capture_sizes else 0
715 - )
712 + # make sure the sizes are in ascending order
713 + self.cudagraph_capture_sizes.sort()
714 + if self.cudagraph_capture_sizes:
715 + assert self.cudagraph_capture_sizes[-1] == self.max_cudagraph_capture_size
716
717 # pre-compute the mapping from batch size to padded graph size
718 - self.bs_to_padded_graph_size = [0 for i in range(self.max_capture_size + 1)]
718 + self.bs_to_padded_graph_size = [
719 + 0 for i in range(self.max_cudagraph_capture_size + 1)
720 + ]
721 for end, start in zip(
720 - self.cudagraph_capture_sizes, self.cudagraph_capture_sizes[1:] + [0]
722 + self.cudagraph_capture_sizes + [self.max_cudagraph_capture_size + 1],
723 + [0] + self.cudagraph_capture_sizes,
724 ):
725 for bs in range(start, end):
726 if bs == start:
727 self.bs_to_padded_graph_size[bs] = start
728 else:
729 self.bs_to_padded_graph_size[bs] = end
727 - self.bs_to_padded_graph_size[self.max_capture_size] = self.max_capture_size
730
731 def set_splitting_ops_for_v1(self):
732 # NOTE: this function needs to be called only when mode is
vllm/config/scheduler.py
-15
@@ -71,14 +71,6 @@ class SchedulerConfig:
71 NOTE: This will be replaced by speculative config in the future; it is
72 present to enable correctness tests until then."""
73
74 - cuda_graph_sizes: list[int] = field(default_factory=list)
75 - """Cuda graph capture sizes
76 - 1. if none provided, then default set to [min(max_num_seqs * 2, 512)]
77 - 2. if one value is provided, then the capture list would follow the
78 - pattern: [1, 2, 4] + [i for i in range(8, cuda_graph_sizes + 1, 8)]
79 - 3. more than one value (e.g. 1 2 128) is provided, then the capture list
80 - will follow the provided list."""
81 -
74 enable_chunked_prefill: SkipValidation[bool] = None # type: ignore
75 """If True, prefill requests can be chunked based
76 on the remaining max_num_batched_tokens."""
@@ -235,13 +227,6 @@ class SchedulerConfig:
227 self.long_prefill_token_threshold,
228 )
229
238 - # NOTE: Default set cuda_graph_sizes to [min(max_num_seqs * 2, 512)].
239 - # This avoids OOM in tight memory scenarios with small max_num_seqs,
240 - # and prevents capture of many large graphs (>512) that would greatly
241 - # increase startup time with limited performance benefit.
242 - if not self.cuda_graph_sizes:
243 - self.cuda_graph_sizes = [min(self.max_num_seqs * 2, 512)]
244 -
230 if self.async_scheduling:
231 self.scheduler_cls = "vllm.v1.core.sched.async_scheduler.AsyncScheduler"
232
vllm/config/vllm.py
+107 -26
@@ -197,10 +197,10 @@ class VllmConfig:
197 return hash_str
198
199 def pad_for_cudagraph(self, batch_size: int) -> int:
200 - # if batch_size > self.compilation_config.max_capture_size,
200 + # if batch_size > self.compilation_config.max_cudagraph_capture_size,
201 # it should raise an IndexError.
202 # the caller should make sure the batch_size is within the range,
203 - # i.e., batch_size <= self.compilation_config.max_capture_size
203 + # i.e., batch_size <= self.compilation_config.max_cudagraph_capture_size
204 return self.compilation_config.bs_to_padded_graph_size[batch_size]
205
206 @staticmethod
@@ -396,6 +396,9 @@ class VllmConfig:
396 if self.model_config is not None and self.model_config.enforce_eager:
397 logger.info("Cudagraph is disabled under eager mode")
398 self.compilation_config.cudagraph_mode = CUDAGraphMode.NONE
399 + # override related settings when enforce eager
400 + self.compilation_config.max_cudagraph_capture_size = 0
401 + self.compilation_config.cudagraph_capture_sizes = []
402 elif envs.VLLM_USE_V1:
403 self.compilation_config.cudagraph_num_of_warmups = 1
404
@@ -654,11 +657,13 @@ class VllmConfig:
657
658 ```python
659 max_graph_size = min(max_num_seqs * 2, 512)
657 - # 1, 2, 4, then multiples of 8 up to max_graph_size
658 - cuda_graph_sizes = [1, 2, 4, 8, 16, 24, 32, 40, ..., max_graph_size]
660 + # 1, 2, 4, then multiples of 8 up to 256 and then multiples of 16
661 + # up to max_graph_size
662 + cuda_graph_sizes = [1, 2, 4] + list(range(8, 256, 8)) + list(
663 + range(256, max_graph_size + 1, 16))
664
665 In the end, `vllm_config.compilation_config.cudagraph_capture_sizes`
661 - will be the final sizes to capture cudagraph (in descending order).
666 + will be the final sizes to capture cudagraph (in ascending order).
667
668 These sizes are used to capture and reuse CUDA graphs for
669 performance-critical paths (e.g., decoding). Capturing enables
@@ -685,35 +690,111 @@ class VllmConfig:
690 not be used.
691 """
692
688 - # calculate the default `batch_size_capture_list`
689 - batch_size_capture_list = []
690 - if self.model_config is not None and not self.model_config.enforce_eager:
691 - cuda_graph_sizes = self.scheduler_config.cuda_graph_sizes
692 - if len(cuda_graph_sizes) == 1:
693 - max_graph_size = cuda_graph_sizes[0]
694 - assert max_graph_size >= 1, (
695 - "Maximum cudagraph size should be greater than or equal to 1."
693 + if (
694 + self.model_config is not None
695 + and not self.model_config.enforce_eager
696 + and self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE
697 + ):
698 + # determine the initial max_cudagraph_capture_size
699 + max_cudagraph_capture_size = (
700 + self.compilation_config.max_cudagraph_capture_size
701 + )
702 + if max_cudagraph_capture_size is None:
703 + max_cudagraph_capture_size = min(
704 + self.scheduler_config.max_num_seqs * 2, 512
705 )
697 - batch_size_capture_list = [
698 - i for i in [1, 2, 4] if i <= max_graph_size
699 - ] + list(range(8, max_graph_size + 1, 8))
700 - elif len(cuda_graph_sizes) > 1:
701 - batch_size_capture_list = sorted(cuda_graph_sizes)
706 + max_num_tokens = self.scheduler_config.max_num_batched_tokens
707 + max_cudagraph_capture_size = min(max_num_tokens, max_cudagraph_capture_size)
708 +
709 + assert max_cudagraph_capture_size >= 1, (
710 + "Maximum cudagraph size should be greater than or equal to 1 "
711 + "when using cuda graph."
712 + )
713 +
714 + # determine the cudagraph_capture_sizes
715 + if self.compilation_config.cudagraph_capture_sizes is not None:
716 + assert len(self.compilation_config.cudagraph_capture_sizes) > 0, (
717 + "cudagraph_capture_sizes should contain at least one element "
718 + "when using cuda graph."
719 + )
720 + # de-duplicate the sizes provided by the config
721 + dedup_sizes = list(set(self.compilation_config.cudagraph_capture_sizes))
722 + cudagraph_capture_sizes = dedup_sizes
723 + # sort to make sure the sizes are in ascending order
724 + cudagraph_capture_sizes.sort()
725 else:
703 - raise TypeError(f"Invalid value for {cuda_graph_sizes=}.")
726 + cudagraph_capture_sizes = [
727 + i for i in [1, 2, 4] if i <= max_cudagraph_capture_size
728 + ]
729 + if max_cudagraph_capture_size >= 8:
730 + # Step size 8 for small batch sizes, up to 256(not included)
731 + cudagraph_capture_sizes += list(
732 + range(8, min(max_cudagraph_capture_size + 1, 256), 8)
733 + )
734 + if max_cudagraph_capture_size >= 256:
735 + # Step size 16 for larger batch sizes
736 + cudagraph_capture_sizes += list(
737 + range(256, max_cudagraph_capture_size + 1, 16)
738 + )
739 +
740 if (
741 self.parallel_config.tensor_parallel_size > 1
742 and self.compilation_config.pass_config.enable_sequence_parallelism
743 ):
708 - batch_size_capture_list = self.update_sizes_for_sequence_parallelism(
709 - batch_size_capture_list
744 + cudagraph_capture_sizes = self.update_sizes_for_sequence_parallelism(
745 + cudagraph_capture_sizes
746 )
711 - max_num_tokens = self.scheduler_config.max_num_batched_tokens
712 - batch_size_capture_list = [
713 - size for size in batch_size_capture_list if size <= max_num_tokens
714 - ]
747
716 - self.compilation_config.init_with_cudagraph_sizes(batch_size_capture_list)
748 + # user-specific compilation_config.max_cudagraph_capture_size get
749 + # truncated to valid_max_size when they are inconsistent.
750 + valid_max_size = (
751 + cudagraph_capture_sizes[-1] if cudagraph_capture_sizes else 0
752 + )
753 + if (
754 + self.compilation_config.max_cudagraph_capture_size is not None
755 + and self.compilation_config.max_cudagraph_capture_size != valid_max_size
756 + ):
757 + # raise error only when both two flags are user-specified
758 + # and they are inconsistent with each other
759 + if self.compilation_config.cudagraph_capture_sizes is not None:
760 + raise ValueError(
761 + "customized max_cudagraph_capture_size"
762 + f"(={self.compilation_config.max_cudagraph_capture_size}) "
763 + "should be consistent with the max value of "
764 + f"cudagraph_capture_sizes(={valid_max_size})"
765 + )
766 +
767 + logger.warning(
768 + "Truncating max_cudagraph_capture_size to %d",
769 + valid_max_size,
770 + )
771 + # always set the final max_cudagraph_capture_size
772 + self.compilation_config.max_cudagraph_capture_size = valid_max_size
773 +
774 + if self.compilation_config.cudagraph_capture_sizes is not None and len(
775 + cudagraph_capture_sizes
776 + ) < len(self.compilation_config.cudagraph_capture_sizes):
777 + # If users have specified capture sizes, we only need to
778 + # compare the lens before and after modification since the modified
779 + # list is only the subset of the original list.
780 + logger.warning(
781 + (
782 + "cudagraph_capture_sizes specified in compilation_config"
783 + " %s is overridden by config %s"
784 + ),
785 + self.compilation_config.cudagraph_capture_sizes,
786 + cudagraph_capture_sizes,
787 + )
788 + # always write back the final sizes
789 + self.compilation_config.cudagraph_capture_sizes = cudagraph_capture_sizes
790 +
791 + else:
792 + # no cudagraph in use
793 + self.compilation_config.max_cudagraph_capture_size = 0
794 + self.compilation_config.cudagraph_capture_sizes = []
795 +
796 + # complete the remaining process.
797 + self.compilation_config.post_init_cudagraph_sizes()
798
799 def recalculate_max_model_len(self, max_model_len: int):
800 # Can only be called in try_verify_and_update_config
vllm/engine/arg_utils.py
+62 -5
@@ -364,7 +364,13 @@ class EngineArgs:
364 kv_cache_dtype: CacheDType = CacheConfig.cache_dtype
365 seed: int | None = ModelConfig.seed
366 max_model_len: int | None = ModelConfig.max_model_len
367 - cuda_graph_sizes: list[int] = get_field(SchedulerConfig, "cuda_graph_sizes")
367 + cuda_graph_sizes: list[int] | None = CompilationConfig.cudagraph_capture_sizes
368 + cudagraph_capture_sizes: list[int] | None = (
369 + CompilationConfig.cudagraph_capture_sizes
370 + )
371 + max_cudagraph_capture_size: int | None = get_field(
372 + CompilationConfig, "max_cudagraph_capture_size"
373 + )
374 # Note: Specifying a custom executor backend by passing a class
375 # is intended for expert use only. The API may change without
376 # notice.
@@ -1007,9 +1013,6 @@ class EngineArgs:
1013 "--max-long-partial-prefills",
1014 **scheduler_kwargs["max_long_partial_prefills"],
1015 )
1010 - scheduler_group.add_argument(
1011 - "--cuda-graph-sizes", **scheduler_kwargs["cuda_graph_sizes"]
1012 - )
1016 scheduler_group.add_argument(
1017 "--long-prefill-token-threshold",
1018 **scheduler_kwargs["long_prefill_token_threshold"],
@@ -1039,6 +1042,29 @@ class EngineArgs:
1042 "--async-scheduling", **scheduler_kwargs["async_scheduling"]
1043 )
1044
1045 + # Compilation arguments
1046 + compilation_kwargs = get_kwargs(CompilationConfig)
1047 + compilation_group = parser.add_argument_group(
1048 + title="CompilationConfig",
1049 + description=CompilationConfig.__doc__,
1050 + )
1051 + compilation_group.add_argument(
1052 + "--cudagraph-capture-sizes", **compilation_kwargs["cudagraph_capture_sizes"]
1053 + )
1054 + compilation_kwargs["cudagraph_capture_sizes"]["help"] = (
1055 + "--cuda-graph-sizes is deprecated and will be removed in v0.13.0 or v1.0.0,"
1056 + " whichever is soonest. Please use --cudagraph-capture-sizes instead."
1057 + )
1058 + compilation_group.add_argument(
1059 + "--cuda-graph-sizes",
1060 + **compilation_kwargs["cudagraph_capture_sizes"],
1061 + deprecated=True,
1062 + )
1063 + compilation_group.add_argument(
1064 + "--max-cudagraph-capture-size",
1065 + **compilation_kwargs["max_cudagraph_capture_size"],
1066 + )
1067 +
1068 # vLLM arguments
1069 vllm_kwargs = get_kwargs(VllmConfig)
1070 vllm_group = parser.add_argument_group(
@@ -1548,7 +1574,6 @@ class EngineArgs:
1574 max_num_batched_tokens=self.max_num_batched_tokens,
1575 max_num_seqs=self.max_num_seqs,
1576 max_model_len=model_config.max_model_len,
1551 - cuda_graph_sizes=self.cuda_graph_sizes,
1577 num_lookahead_slots=num_lookahead_slots,
1578 enable_chunked_prefill=self.enable_chunked_prefill,
1579 disable_chunked_mm_input=self.disable_chunked_mm_input,
@@ -1616,6 +1641,38 @@ class EngineArgs:
1641 collect_detailed_traces=self.collect_detailed_traces,
1642 )
1643
1644 + # Compilation config overrides
1645 + if self.cuda_graph_sizes is not None:
1646 + logger.warning(
1647 + "--cuda-graph-sizes is deprecated and will be removed in v0.13.0 or "
1648 + "v1.0.0, whichever is soonest. Please use --cudagraph-capture-sizes "
1649 + "instead."
1650 + )
1651 + if self.compilation_config.cudagraph_capture_sizes is not None:
1652 + raise ValueError(
1653 + "cuda_graph_sizes and compilation_config."
1654 + "cudagraph_capture_sizes are mutually exclusive"
1655 + )
1656 + self.compilation_config.cudagraph_capture_sizes = self.cuda_graph_sizes
1657 + if self.cudagraph_capture_sizes is not None:
1658 + if self.compilation_config.cudagraph_capture_sizes is not None:
1659 + raise ValueError(
1660 + "cudagraph_capture_sizes and compilation_config."
1661 + "cudagraph_capture_sizes are mutually exclusive"
1662 + )
1663 + self.compilation_config.cudagraph_capture_sizes = (
1664 + self.cudagraph_capture_sizes
1665 + )
1666 + if self.max_cudagraph_capture_size is not None:
1667 + if self.compilation_config.max_cudagraph_capture_size is not None:
1668 + raise ValueError(
1669 + "max_cudagraph_capture_size and compilation_config."
1670 + "max_cudagraph_capture_size are mutually exclusive"
1671 + )
1672 + self.compilation_config.max_cudagraph_capture_size = (
1673 + self.max_cudagraph_capture_size
1674 + )
1675 +
1676 config = VllmConfig(
1677 model_config=model_config,
1678 cache_config=cache_config,
vllm/model_executor/layers/quantization/mxfp4.py
+1 -1
@@ -185,7 +185,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
185 self.moe = moe
186 self.mxfp4_backend = get_mxfp4_backend()
187 self.max_capture_size = (
188 - get_current_vllm_config().compilation_config.max_capture_size
188 + get_current_vllm_config().compilation_config.max_cudagraph_capture_size
189 )
190
191 assert self.mxfp4_backend != Mxfp4Backend.NONE, (
vllm/model_executor/models/config.py
+11 -13
@@ -259,21 +259,19 @@ class GptOssForCausalLMConfig(VerifyAndUpdateConfig):
259 # Increase the max capture size from 512 to 992 for performance.
260 # NOTE(woosuk): This will increase the number of CUDA graphs
261 # from 67 to 81.
262 - scheduler_config = vllm_config.scheduler_config
263 - if len(scheduler_config.cuda_graph_sizes) == 1:
264 - max_capture_size = scheduler_config.cuda_graph_sizes[0]
262 + compilation_config = vllm_config.compilation_config
263 + # Only override when the user has not set either of
264 + # cudagraph_capture_sizes or max_cudagraph_capture_size.
265 + if (
266 + compilation_config.cudagraph_capture_sizes is None
267 + and compilation_config.max_cudagraph_capture_size is None
268 + ):
269 # FIXME(woosuk): When using full cuda graph with FA3, the max
270 # supported size is 992.
267 - if max_capture_size < 992:
268 - cuda_graph_sizes = [1, 2, 4]
269 - # Step size 8 for small batch sizes
270 - cuda_graph_sizes += [i for i in range(8, 256, 8)]
271 - # Step size 16 for larger batch sizes
272 - cuda_graph_sizes += [i for i in range(256, 993, 16)]
273 - scheduler_config.cuda_graph_sizes = cuda_graph_sizes
274 - logger.info(
275 - "Overriding max cuda graph capture size to %d for performance.", 992
276 - )
271 + compilation_config.max_cudagraph_capture_size = 992
272 + logger.info(
273 + "Overriding max cuda graph capture size to %d for performance.", 992
274 + )
275
276
277 class MambaModelConfig(VerifyAndUpdateConfig):
vllm/v1/attention/backends/flash_attn.py
+1 -1
@@ -236,7 +236,7 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
236 self.use_full_cuda_graph = (
237 self.compilation_config.cudagraph_mode.has_full_cudagraphs()
238 )
239 - self.max_cudagraph_size = self.compilation_config.max_capture_size
239 + self.max_cudagraph_size = self.compilation_config.max_cudagraph_capture_size
240
241 if self.use_full_cuda_graph and self.aot_schedule:
242 if self.max_cudagraph_size > 992:
vllm/v1/attention/backends/flashinfer.py
+1 -1
@@ -324,7 +324,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
324 ] = {}
325 self._decode_cudagraph_max_bs = min(
326 (1 + num_spec_tokens) * max_num_reqs,
327 - self.compilation_config.max_capture_size,
327 + self.compilation_config.max_cudagraph_capture_size,
328 )
329
330 self.num_qo_heads = self.model_config.get_num_attention_heads(
vllm/v1/attention/backends/gdn_attn.py
+1 -1
@@ -87,7 +87,7 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
87 )
88 self.decode_cudagraph_max_bs = min(
89 self.vllm_config.scheduler_config.max_num_seqs * (self.num_spec + 1),
90 - self.compilation_config.max_capture_size,
90 + self.compilation_config.max_cudagraph_capture_size,
91 )
92
93 self.spec_state_indices_tensor = torch.empty(
vllm/v1/attention/backends/mamba_attn.py
+1 -1
@@ -36,7 +36,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
36 self.compilation_config = vllm_config.compilation_config
37 self.decode_cudagraph_max_bs = min(
38 self.vllm_config.scheduler_config.max_num_seqs,
39 - self.compilation_config.max_capture_size,
39 + self.compilation_config.max_cudagraph_capture_size,
40 )
41 self.state_indices_tensor = torch.empty(
42 (self.decode_cudagraph_max_bs,),
vllm/v1/attention/backends/mla/flashattn_mla.py
+1 -1
@@ -89,7 +89,7 @@ class FlashAttnMLAMetadataBuilder(MLACommonMetadataBuilder[FlashAttnMLAMetadata]
89 self.use_full_cuda_graph = (
90 self.compilation_config.cudagraph_mode.has_full_cudagraphs()
91 )
92 - self.max_cudagraph_size = self.compilation_config.max_capture_size
92 + self.max_cudagraph_size = self.compilation_config.max_cudagraph_capture_size
93
94 if self.use_full_cuda_graph and self.fa_aot_schedule:
95 if self.max_cudagraph_size > 992:
vllm/v1/spec_decode/eagle.py
+1 -1
@@ -104,7 +104,7 @@ class EagleProposer:
104 )
105
106 self.cudagraph_batch_sizes = (
107 - list(reversed(self.vllm_config.compilation_config.cudagraph_capture_sizes))
107 + (sorted(self.vllm_config.compilation_config.cudagraph_capture_sizes))
108 if self.use_cuda_graph
109 else []
110 )
vllm/v1/worker/gpu_model_runner.py
+3 -6
@@ -379,16 +379,13 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
379 self.async_output_copy_stream = torch.cuda.Stream()
380 self.prepare_inputs_event = torch.cuda.Event()
381
382 - # TODO(woosuk): Provide an option to tune the max cudagraph batch size.
383 - # The convention is different.
382 # self.cudagraph_batch_sizes sorts in ascending order.
385 - # The batch sizes in the config are in descending order.
383 if (
384 self.compilation_config.cudagraph_capture_sizes
385 and self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE
386 ):
390 - self.cudagraph_batch_sizes = list(
391 - reversed(self.compilation_config.cudagraph_capture_sizes)
387 + self.cudagraph_batch_sizes = sorted(
388 + self.compilation_config.cudagraph_capture_sizes
389 )
390
391 # Cache the device properties.
@@ -3791,7 +3788,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
3788
3789 if cudagraph_mode.mixed_mode() != CUDAGraphMode.NONE:
3790 cudagraph_runtime_mode = cudagraph_mode.mixed_mode()
3794 -
3791 + # make sure we capture the largest batch size first
3792 compilation_cases = list(
3793 product(reversed(self.cudagraph_batch_sizes), lora_cases)
3794 )