@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 2261340806

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 9, 23:00 UTC

[Misc] Remove pad_for_cudagraphs from config (#30143)

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>

Lucas Wilkinson committed Jan 20, 2026 at 13:05 UTC 2261340806d4afd9dda90053975747d3453e46f3

9 files changed +157 -136

tests/compile/test_config.py

+34 -2

@@ -2,14 +2,20 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import copy

4 from contextlib import nullcontext

5 -from unittest.mock import patch

5 +from unittest.mock import MagicMock, patch

7 import pytest

8 from pydantic import ValidationError

10 from vllm.compilation.counter import compilation_counter

11 from vllm.compilation.fix_functionalization import FixFunctionalizationPass

12 -from vllm.config import CompilationConfig, CUDAGraphMode, ParallelConfig, VllmConfig

12 +from vllm.config import (

13 + CompilationConfig,

14 + CUDAGraphMode,

15 + ParallelConfig,

16 + SchedulerConfig,

17 + VllmConfig,

18 +)

19 from vllm.config.compilation import CompilationMode, PassConfig

20 from vllm.engine.arg_utils import EngineArgs

21 from vllm.platforms import current_platform

@@ -17,6 +23,7 @@ from vllm.utils.torch_utils import (

23 _is_torch_equal_or_newer,

24 is_torch_equal,

25 )

26 +from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher

28 # This import automatically registers `torch.ops.silly.attention`

29 from . import silly_attention # noqa: F401

@@ -472,6 +479,19 @@ def test_cached_compilation_config(default_vllm_config):

479 assert "torch.ops._C.static_scaled_fp8_quant.default(" in code

480

481

482 +def _create_vllm_config_for_validation(

483 + compilation_config: CompilationConfig,

484 +) -> MagicMock:

485 + """Helper to create a mock VllmConfig for padding validation testing."""

486 + mock_config = MagicMock(spec=VllmConfig)

487 + mock_config.compilation_config = compilation_config

488 + mock_config.scheduler_config = SchedulerConfig.default_factory(max_num_seqs=8)

489 + mock_config.parallel_config = ParallelConfig()

490 + mock_config.speculative_config = None

491 + mock_config.lora_config = None

492 + return mock_config

493 +

494 +

495 def test_compile_sizes_padding_validation():

496 """Test that compile_sizes with values that would be padded raises an error."""

497 # cudagraph_capture_sizes=[1, 2, 4, 8] means:

@@ -488,29 +508,39 @@ def test_compile_sizes_padding_validation():

508 cudagraph_capture_sizes=[1, 2, 4, 8],

509 max_cudagraph_capture_size=8,

510 compile_sizes=[3],

511 + cudagraph_mode=CUDAGraphMode.FULL,

512 )

513 config.post_init_cudagraph_sizes()

514 + dispatcher = CudagraphDispatcher(_create_vllm_config_for_validation(config))

515 + dispatcher.initialize_cudagraph_keys(CUDAGraphMode.FULL)

516

517 with pytest.raises(ValueError, match="would be padded to"):

518 config = CompilationConfig(

519 cudagraph_capture_sizes=[1, 2, 4, 8],

520 max_cudagraph_capture_size=8,

521 compile_sizes=[5],

522 + cudagraph_mode=CUDAGraphMode.FULL,

523 )

524 config.post_init_cudagraph_sizes()

525 + dispatcher = CudagraphDispatcher(_create_vllm_config_for_validation(config))

526 + dispatcher.initialize_cudagraph_keys(CUDAGraphMode.FULL)

527

528 config = CompilationConfig(

529 cudagraph_capture_sizes=[1, 2, 4, 8],

530 max_cudagraph_capture_size=8,

531 compile_sizes=[1, 2, 4, 8],

532 + cudagraph_mode=CUDAGraphMode.FULL,

533 )

534 config.post_init_cudagraph_sizes()

535 assert sorted(config.compile_sizes) == [1, 2, 4, 8]

536 + dispatcher = CudagraphDispatcher(_create_vllm_config_for_validation(config))

537 + dispatcher.initialize_cudagraph_keys(CUDAGraphMode.FULL) # Should not raise

538

539 config = CompilationConfig(

540 cudagraph_capture_sizes=[1, 2, 4, 8],

541 max_cudagraph_capture_size=8,

542 compile_sizes=["cudagraph_capture_sizes"],

543 + cudagraph_mode=CUDAGraphMode.FULL,

544 )

545 config.post_init_cudagraph_sizes()

546 assert sorted(config.compile_sizes) == [1, 2, 4, 8]

@@ -535,3 +565,5 @@ def test_compile_sizes_padding_validation():

565 )

566 config.post_init_cudagraph_sizes()

567 assert sorted(config.compile_sizes) == [3, 5, 7]

568 + dispatcher = CudagraphDispatcher(_create_vllm_config_for_validation(config))

569 + dispatcher.initialize_cudagraph_keys(CUDAGraphMode.NONE) # Should not raise

tests/models/language/generation/test_hybrid.py

+9 -1

@@ -9,6 +9,7 @@ from tests.models.registry import HF_EXAMPLE_MODELS

9 from tests.utils import multi_gpu_test

10 from vllm.engine.arg_utils import EngineArgs

11 from vllm.sampling_params import SamplingParams

12 +from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher

14 from ...utils import check_logprobs_close, check_outputs_equal

@@ -172,7 +173,14 @@ def test_mamba_cache_cg_padding(

173 tensor dimensions aren't compatible.

174 """

175 vllm_config = EngineArgs(model=model, trust_remote_code=True).create_engine_config()

175 - while len(example_prompts) == vllm_config.pad_for_cudagraph(len(example_prompts)):

176 + cudagraph_dispatcher = CudagraphDispatcher(vllm_config)

177 + cudagraph_dispatcher.initialize_cudagraph_keys(

178 + vllm_config.compilation_config.cudagraph_mode

179 + )

180 + while (

181 + len(example_prompts)

182 + == cudagraph_dispatcher.dispatch(len(example_prompts))[1].num_tokens

183 + ):

184 example_prompts.append(example_prompts[0])

185

186 try:

tests/v1/cudagraph/test_cudagraph_dispatch.py

+4 -6

@@ -61,9 +61,6 @@ def _create_vllm_config(

61 )

63 compilation_config.post_init_cudagraph_sizes()

64 - mock_config.pad_for_cudagraph = (

65 - lambda batch_size: compilation_config.bs_to_padded_graph_size[batch_size]

66 - )

65 return mock_config

@@ -169,6 +166,7 @@ class TestCudagraphDispatcher:

166 rt_mode, key = dispatcher.dispatch(

167 num_tokens=8, uniform_decode=False, has_lora=False, disable_full=True

168 )

169 +

170 if "PIECEWISE" in cudagraph_mode_str: # string contains check

171 assert rt_mode == CUDAGraphMode.PIECEWISE

172 assert key == desc_full_exact.relax_for_mixed_batch_cudagraphs()

@@ -360,7 +358,7 @@ class TestCudagraphIntegration:

358 ):

359 full_wrapper(input_1)

360

363 - rt_mode, key = self.dispatcher.dispatch(desc_1)

361 + rt_mode, key = self.dispatcher.dispatch(num_tokens=desc_1.num_tokens)

362 # 1. Capture first shape

363 action = self._run_and_monitor_call(full_wrapper, input_1, rt_mode, key)

364 assert action == "capture_global"

@@ -369,7 +367,7 @@ class TestCudagraphIntegration:

367 action = self._run_and_monitor_call(full_wrapper, input_1, rt_mode, key)

368 assert action == "replay"

369

372 - rt_mode, key = self.dispatcher.dispatch(desc_2)

370 + rt_mode, key = self.dispatcher.dispatch(num_tokens=desc_2.num_tokens)

371 # 3. Capture second shape

372 action = self._run_and_monitor_call(full_wrapper, input_2, rt_mode, key)

373 assert action == "capture_global"

@@ -381,7 +379,7 @@ class TestCudagraphIntegration:

379 assert action == "replay"

380

381 # 5. Bypass if no key match

384 - rt_mode, key = self.dispatcher.dispatch(desc_3_unseen)

382 + rt_mode, key = self.dispatcher.dispatch(num_tokens=desc_3_unseen.num_tokens)

383 assert rt_mode == CUDAGraphMode.NONE

384 action = self._run_and_monitor_call(full_wrapper, input_3, rt_mode, key)

385 assert action == "bypass"

vllm/compilation/piecewise_backend.py

+1 -1

@@ -15,7 +15,7 @@ from torch._inductor.runtime.triton_heuristics import CachingAutotuner

15 from vllm.compilation.backends import VllmBackend

16 from vllm.compilation.monitor import end_monitoring_torch_compile

17 from vllm.config import VllmConfig

18 -from vllm.config.compilation import Range

18 +from vllm.config.utils import Range

19 from vllm.logger import init_logger

21 logger = init_logger(__name__)

vllm/compilation/sequence_parallelism.py

+1 -1

@@ -11,7 +11,7 @@ import torch.fx as fx

11 from torch._inductor.pattern_matcher import PatternMatcherPass

13 from vllm.config import VllmConfig

14 -from vllm.config.compilation import Range

14 +from vllm.config.utils import Range

15 from vllm.distributed import get_tp_group, tensor_model_parallel_all_reduce

16 from vllm.distributed.parallel_state import get_tensor_model_parallel_world_size

17 from vllm.logger import init_logger

vllm/config/compilation.py

-47

@@ -581,15 +581,6 @@ class CompilationConfig:

581 local_cache_dir: str = field(default=None, init=False) # type: ignore

582 """local cache dir for each rank"""

583

584 - bs_to_padded_graph_size: list[int] = field(

585 - default=None, # type: ignore

586 - init=False,

587 - )

588 - """optimization:

589 - Intuitively, bs_to_padded_graph_size should be dict[int, int].

590 - since we know all keys are in a range [0, max_cudagraph_capture_size],

591 - we can optimize it to list[int] for better lookup performance."""

592 -

584 # keep track of enabled and disabled custom ops

585 enabled_custom_ops: Counter[str] = field(default_factory=Counter, init=False)

586 """custom ops that are enabled"""

@@ -639,7 +630,6 @@ class CompilationConfig:

630 "debug_dump_path",

631 "cache_dir",

632 "local_cache_dir",

642 - "bs_to_padded_graph_size",

633 "traced_files",

634 "compilation_time",

635 "static_forward_context",

@@ -661,7 +651,6 @@ class CompilationConfig:

651 "enabled_custom_ops": True,

652 "disabled_custom_ops": True,

653 "compilation_time": True,

664 - "bs_to_padded_graph_size": True,

654 "traced_files": True,

655 "inductor_compile_config": {

656 "post_grad_custom_post_pass": True,

@@ -882,7 +871,6 @@ class CompilationConfig:

871 """To complete the initialization after cudagraph related

872 configs are set. This includes:

873 - initialize compile_sizes

885 - - pre-compute the mapping bs_to_padded_graph_size

874 """

875

876 computed_compile_sizes = []

@@ -906,23 +894,6 @@ class CompilationConfig:

894 if self.cudagraph_capture_sizes:

895 assert self.cudagraph_capture_sizes[-1] == self.max_cudagraph_capture_size

896

909 - # May get recomputed in the model runner if adjustment is needed for spec-decode

910 - self.compute_bs_to_padded_graph_size()

911 -

912 - # Validate that compile_sizes won't be changed by padding.

913 - # Only validate when cudagraphs are actually being used.

914 - if self.compile_sizes and self.cudagraph_mode != CUDAGraphMode.NONE:

915 - for size in self.compile_sizes:

916 - if size <= self.max_cudagraph_capture_size:

917 - padded = self.bs_to_padded_graph_size[size]

918 - if padded != size:

919 - raise ValueError(

920 - f"compile_sizes contains {size} which would be "

921 - f"padded to {padded}. All compile_sizes must be "

922 - "values that won't be changed by cudagraph padding. "

923 - "Use values from cudagraph_capture_sizes."

924 - )

925 -

897 def set_splitting_ops_for_v1(

898 self, all2all_backend: str, data_parallel_size: int = 1

899 ):

@@ -1134,24 +1105,6 @@ class CompilationConfig:

1105 self.max_cudagraph_capture_size = rounded_sizes[-1]

1106 self.cudagraph_capture_sizes = rounded_sizes

1107

1137 - # Recompute after adjusting the cudagraph sizes

1138 - self.compute_bs_to_padded_graph_size()

1139 -

1140 - def compute_bs_to_padded_graph_size(self):

1141 - # pre-compute the mapping from batch size to padded graph size

1142 - self.bs_to_padded_graph_size = [

1143 - 0 for i in range(self.max_cudagraph_capture_size + 1)

1144 - ]

1145 - for end, start in zip(

1146 - self.cudagraph_capture_sizes + [self.max_cudagraph_capture_size + 1],

1147 - [0] + self.cudagraph_capture_sizes,

1148 - ):

1149 - for bs in range(start, end):

1150 - if bs == start:

1151 - self.bs_to_padded_graph_size[bs] = start

1152 - else:

1153 - self.bs_to_padded_graph_size[bs] = end

1154 -

1108 def get_compile_ranges(self) -> list[Range]:

1109 """Get the compile ranges for the compilation config."""

1110 if self.compile_ranges_split_points is None:

vllm/v1/cudagraph_dispatcher.py

+55 -5

@@ -57,13 +57,47 @@ class CudagraphDispatcher:

57 )

59 self.keys_initialized = False

60 + # Default cudagraph_mode to NONE until initialize_cudagraph_keys is called

61 + self.cudagraph_mode = CUDAGraphMode.NONE

62 +

63 + def _compute_bs_to_padded_graph_size(self) -> None:

64 + """Pre-compute the mapping from batch size to padded graph size."""

65 + max_size = self.compilation_config.max_cudagraph_capture_size

66 + capture_sizes = self.compilation_config.cudagraph_capture_sizes

67 + self._bs_to_padded_graph_size: list[int] = [0] * (max_size + 1)

68 + for end, start in zip(

69 + capture_sizes + [max_size + 1],

70 + [0] + capture_sizes,

71 + ):

72 + for bs in range(start, end):

73 + if bs == start:

74 + self._bs_to_padded_graph_size[bs] = start

75 + else:

76 + self._bs_to_padded_graph_size[bs] = end

77 +

78 + # Validate that compile_sizes won't be changed by padding.

79 + # Only validate when cudagraphs are actually being used.

80 + if (

81 + self.compilation_config.compile_sizes

82 + and self.cudagraph_mode != CUDAGraphMode.NONE

83 + ):

84 + for size in self.compilation_config.compile_sizes:

85 + if size <= self.compilation_config.max_cudagraph_capture_size:

86 + padded = self._bs_to_padded_graph_size[size]

87 + if padded != size:

88 + raise ValueError(

89 + f"compile_sizes contains {size} which would be "

90 + f"padded to {padded}. All compile_sizes must be "

91 + "values that won't be changed by cudagraph padding. "

92 + "Use values from cudagraph_capture_sizes."

93 + )

95 def _create_padded_batch_descriptor(

96 self, num_tokens: int, uniform_decode: bool, has_lora: bool

97 ) -> BatchDescriptor:

98 max_num_seqs = self.vllm_config.scheduler_config.max_num_seqs

99 uniform_decode_query_len = self.uniform_decode_query_len

66 - num_tokens_padded = self.vllm_config.pad_for_cudagraph(num_tokens)

100 + num_tokens_padded = self._bs_to_padded_graph_size[num_tokens]

101

102 if uniform_decode and self.cudagraph_mode.has_mode(CUDAGraphMode.FULL):

103 num_reqs = num_tokens_padded // uniform_decode_query_len

@@ -88,12 +122,19 @@ class CudagraphDispatcher:

122 self.cudagraph_keys[runtime_mode].add(batch_descriptor)

123

124 def initialize_cudagraph_keys(

91 - self, cudagraph_mode: CUDAGraphMode, uniform_decode_query_len: int

125 + self, cudagraph_mode: CUDAGraphMode, uniform_decode_query_len: int = 1

126 ):

127 # This should be called only after attention backend is initialized. So we can

128 # get the correct cudagraph mode after backend support is resolved.

129 self.cudagraph_mode = cudagraph_mode

130

131 + # Early exit if cudagraphs are disabled

132 + if cudagraph_mode == CUDAGraphMode.NONE:

133 + self.keys_initialized = True

134 + return

135 +

136 + self._compute_bs_to_padded_graph_size()

137 +

138 # LoRA activation cases to specialize the cuda graphs on

139 if self.vllm_config.lora_config:

140 if self.compilation_config.cudagraph_specialize_lora:

@@ -143,15 +184,24 @@ class CudagraphDispatcher:

184 def dispatch(

185 self,

186 num_tokens: int,

146 - uniform_decode: bool,

147 - has_lora: bool,

187 + uniform_decode: bool = False,

188 + has_lora: bool = False,

189 disable_full: bool = False,

190 ) -> tuple[CUDAGraphMode, BatchDescriptor]:

191 """

151 - Given conditions(e.g.,batch descriptor and if using cascade attention),

192 + Given conditions(e.g.,batch descriptor and if using piecewise only),

193 dispatch to a cudagraph runtime mode and the valid batch descriptor.

194 A new batch descriptor is returned as we might dispatch a uniform batch

195 to a graph that supports a more general batch (uniform to non-uniform).

196 +

197 + Args:

198 + num_tokens: Number of tokens in the batch.

199 + uniform_decode: Whether the batch is uniform decode (i.e. uniform and query

200 + length is uniform_decode_query_len).

201 + has_lora: Whether LoRA is active.

202 + disable_full: If True, skip FULL cudagraph checks and

203 + return PIECEWISE or NONE only. (can be used for features like

204 + cascade attention that are not supported by full cudagraphs)

205 """

206 if (

207 not self.keys_initialized

vllm/v1/spec_decode/eagle.py

+43 -64

@@ -9,7 +9,6 @@ import torch

9 import torch.nn as nn

11 from vllm.config import (

12 - CompilationMode,

12 CUDAGraphMode,

13 VllmConfig,

14 get_layers_from_vllm_config,

@@ -36,6 +35,7 @@ from vllm.v1.attention.backends.tree_attn import (

35 TreeAttentionMetadataBuilder,

36 )

37 from vllm.v1.attention.backends.triton_attn import TritonAttentionMetadata

38 +from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher

39 from vllm.v1.kv_cache_interface import KVCacheConfig

40 from vllm.v1.sample.metadata import SamplingMetadata

41 from vllm.v1.sample.sampler import _SAMPLING_EPS

@@ -100,24 +100,13 @@ class SpecDecodeBaseProposer:

100 self._get_eagle3_use_aux_hidden_state_from_config()

101 )

102

103 - self.use_cuda_graph = False

104 -

103 self.compilation_config = self.vllm_config.compilation_config

106 - if self.compilation_config.mode == CompilationMode.VLLM_COMPILE:

107 - cudagraph_mode = self.compilation_config.cudagraph_mode

108 - if cudagraph_mode != CUDAGraphMode.NONE and not cudagraph_mode.has_mode(

109 - CUDAGraphMode.PIECEWISE

110 - ):

111 - logger.warning(

112 - "Currently the eagle proposer only supports cudagraph_mode "

113 - "PIECEWISE, if you want the drafter to use cuda graphs, "

114 - "please set compilation_config.cudagraph_mode to PIECEWISE "

115 - "or FULL_AND_PIECEWISE"

116 - )

117 - self.use_cuda_graph = (

118 - cudagraph_mode.has_mode(CUDAGraphMode.PIECEWISE)

119 - and not self.speculative_config.enforce_eager

120 - )

104 +

105 + # Cudagraph dispatcher for PIECEWISE-only dispatching in eagle.

106 + # Keys are initialized later via initialize_cudagraph_keys() called from

107 + # gpu_model_runner._check_and_update_cudagraph_mode after

108 + # adjust_cudagraph_sizes_for_spec_decode is called.

109 + self.cudagraph_dispatcher = CudagraphDispatcher(self.vllm_config)

110

111 # persistent buffers for cuda graph

112 self.input_ids = torch.zeros(

@@ -234,6 +223,23 @@ class SpecDecodeBaseProposer:

223 else:

224 self.positions[:num_tokens] = positions

225

226 + def initialize_cudagraph_keys(self, cudagraph_mode: CUDAGraphMode) -> None:

227 + """Initialize cudagraph dispatcher keys for eagle.

228 +

229 + Eagle only supports PIECEWISE cudagraphs (via mixed_mode).

230 + This should be called after adjust_cudagraph_sizes_for_spec_decode.

231 + """

232 + if (

233 + not self.speculative_config.enforce_eager

234 + and cudagraph_mode.mixed_mode()

235 + in [CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL]

236 + ):

237 + eagle_cudagraph_mode = CUDAGraphMode.PIECEWISE

238 + else:

239 + eagle_cudagraph_mode = CUDAGraphMode.NONE

240 +

241 + self.cudagraph_dispatcher.initialize_cudagraph_keys(eagle_cudagraph_mode)

242 +

243 def propose(

244 self,

245 # [num_tokens]

@@ -304,16 +310,10 @@ class SpecDecodeBaseProposer:

310 num_tokens_unpadded=num_tokens, num_tokens_padded=num_tokens

311 )

312

307 - cudagraph_runtime_mode = CUDAGraphMode.NONE

308 - if (

309 - self.use_cuda_graph

310 - and num_tokens_dp_padded

311 - <= self.compilation_config.max_cudagraph_capture_size

312 - ):

313 - num_input_tokens = self.vllm_config.pad_for_cudagraph(num_tokens_dp_padded)

314 - cudagraph_runtime_mode = CUDAGraphMode.PIECEWISE

315 - else:

316 - num_input_tokens = num_tokens_dp_padded

313 + cudagraph_runtime_mode, batch_desc = self.cudagraph_dispatcher.dispatch(

314 + num_tokens_dp_padded

315 + )

316 + num_input_tokens = batch_desc.num_tokens

317 if num_tokens_across_dp is not None:

318 num_tokens_across_dp[self.dp_rank] = num_input_tokens

319

@@ -412,16 +412,10 @@ class SpecDecodeBaseProposer:

412 num_tokens_unpadded=batch_size, num_tokens_padded=batch_size

413 )

414

415 - if (

416 - self.use_cuda_graph

417 - and batch_size_dp_padded

418 - <= self.compilation_config.max_cudagraph_capture_size

419 - ):

420 - input_batch_size = self.vllm_config.pad_for_cudagraph(batch_size_dp_padded)

421 - cudagraph_runtime_mode = CUDAGraphMode.PIECEWISE

422 - else:

423 - input_batch_size = batch_size_dp_padded

424 - cudagraph_runtime_mode = CUDAGraphMode.NONE

415 + cudagraph_runtime_mode, batch_desc = self.cudagraph_dispatcher.dispatch(

416 + batch_size_dp_padded

417 + )

418 + input_batch_size = batch_desc.num_tokens

419 if batch_size_across_dp is not None:

420 batch_size_across_dp[self.dp_rank] = input_batch_size

421

@@ -870,15 +864,10 @@ class SpecDecodeBaseProposer:

864 self.positions[:num_tokens] = tree_positions.view(-1)

865 self.hidden_states[:num_tokens] = tree_hidden_states.view(num_tokens, -1)

866

873 - if (

874 - self.use_cuda_graph

875 - and num_tokens <= self.compilation_config.max_cudagraph_capture_size

876 - ):

877 - num_input_tokens = self.vllm_config.pad_for_cudagraph(num_tokens)

878 - cudagraph_runtime_mode = CUDAGraphMode.PIECEWISE

879 - else:

880 - num_input_tokens = num_tokens

881 - cudagraph_runtime_mode = CUDAGraphMode.NONE

867 + cudagraph_runtime_mode, batch_desc = self.cudagraph_dispatcher.dispatch(

868 + num_tokens

869 + )

870 + num_input_tokens = batch_desc.num_tokens

871 # Run the model.

872 with set_forward_context(

873 per_layer_attn_metadata,

@@ -1216,9 +1205,6 @@ class SpecDecodeBaseProposer:

1205 use_cudagraphs: bool = True,

1206 is_graph_capturing: bool = False,

1207 ) -> None:

1219 - # Determine if CUDA graphs should be used for this run.

1220 - cudagraphs_enabled = use_cudagraphs and self.use_cuda_graph

1221 -

1208 # FIXME: when using tree-based specdec, adjust number of forward-passes

1209 # according to the depth of the tree.

1210 for fwd_idx in range(

@@ -1228,16 +1214,10 @@ class SpecDecodeBaseProposer:

1214 num_tokens_dp_padded, num_tokens_across_dp = self._pad_batch_across_dp(

1215 num_tokens_unpadded=num_tokens, num_tokens_padded=num_tokens

1216 )

1231 - if (

1232 - cudagraphs_enabled

1233 - and num_tokens_dp_padded

1234 - <= self.compilation_config.max_cudagraph_capture_size

1235 - ):

1236 - num_input_tokens = self.vllm_config.pad_for_cudagraph(

1237 - num_tokens_dp_padded

1238 - )

1239 - else:

1240 - num_input_tokens = num_tokens_dp_padded

1217 + cudagraph_runtime_mode, batch_desc = self.cudagraph_dispatcher.dispatch(

1218 + num_tokens_dp_padded

1219 + )

1220 + num_input_tokens = batch_desc.num_tokens

1221 if num_tokens_across_dp is not None:

1222 num_tokens_across_dp[self.dp_rank] = num_input_tokens

1223

@@ -1246,9 +1226,7 @@ class SpecDecodeBaseProposer:

1226 self.vllm_config,

1227 num_tokens=num_input_tokens,

1228 num_tokens_across_dp=num_tokens_across_dp,

1249 - cudagraph_runtime_mode=CUDAGraphMode.PIECEWISE

1250 - if cudagraphs_enabled

1251 - else CUDAGraphMode.NONE,

1229 + cudagraph_runtime_mode=cudagraph_runtime_mode,

1230 ):

1231 if self.supports_mm_inputs:

1232 input_ids = None

@@ -1340,7 +1318,8 @@ class SpecDecodeBaseProposer:

1318 num_tokens_unpadded=num_tokens_unpadded,

1319 parallel_config=self.vllm_config.parallel_config,

1320 allow_microbatching=False,

1343 - allow_dp_padding=self.use_cuda_graph,

1321 + allow_dp_padding=self.cudagraph_dispatcher.cudagraph_mode

1322 + != CUDAGraphMode.NONE,

1323 num_tokens_padded=num_tokens_padded,

1324 uniform_decode=None,

1325 num_scheduled_tokens_per_request=None,

vllm/v1/worker/gpu_model_runner.py

+10 -9

@@ -2139,15 +2139,11 @@ class GPUModelRunner(

2139 self.kv_sharing_fast_prefill_logits_indices[num_logits:].fill_(

2140 logits_indices[-1].item()

2141 )

2142 - if (

2143 - self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE

2144 - and num_logits <= self.cudagraph_batch_sizes[-1]

2145 - ):

2146 - # Use piecewise CUDA graphs.

2147 - # Add padding to the batch size.

2148 - num_logits_padded = self.vllm_config.pad_for_cudagraph(num_logits)

2149 - else:

2150 - num_logits_padded = num_logits

2142 + # Dispatch for the decoder portion of the model.

2143 + _, batch_desc = self.cudagraph_dispatcher.dispatch(

2144 + num_logits, disable_full=True

2145 + )

2146 + num_logits_padded = batch_desc.num_tokens

2147 logits_indices_padded = self.kv_sharing_fast_prefill_logits_indices[

2148 :num_logits_padded

2149 ]

@@ -5212,6 +5208,11 @@ class GPUModelRunner(

5208 cudagraph_mode, self.uniform_decode_query_len

5209 )

5210

5211 + # Initialize eagle's cudagraph dispatcher if using eagle spec decode.

5212 + if self.speculative_config and self.speculative_config.use_eagle():

5213 + assert isinstance(self.drafter, EagleProposer)

5214 + self.drafter.initialize_cudagraph_keys(cudagraph_mode)

5215 +

5216 def calculate_reorder_batch_threshold(self) -> None:

5217 """

5218 Choose the minimum reorder batch threshold from all attention groups.