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 7, 14:00 UTC
[DP] Only use DP padding when cudagraphs are actually used (#34102)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Lucas Wilkinson committed
Feb 27, 2026 at 15:14 UTC
1d532f9d8fb205942035313293af701ee580a7e2
7 files changed
+139
-127
tests/v1/cudagraph/test_cudagraph_dispatch.py
+16
-2
@@ -176,10 +176,14 @@ class TestCudagraphDispatcher:
176
assert rt_mode == CUDAGraphMode.NONE
177
assert key == BatchDescriptor(num_tokens=15)
178
179
- # 4. disable_full should have a fall back mode (e.g., cascade attention)
179
+ # 4. invalid_modes={FULL} should have a fall back mode
180
+ # (e.g., cascade attention)
181
desc_full_exact = BatchDescriptor(num_tokens=8, uniform=False)
182
rt_mode, key = dispatcher.dispatch(
182
- num_tokens=8, uniform_decode=False, has_lora=False, disable_full=True
183
+ num_tokens=8,
184
+ uniform_decode=False,
185
+ has_lora=False,
186
+ invalid_modes={CUDAGraphMode.FULL},
187
)
188
189
if "PIECEWISE" in cudagraph_mode_str: # string contains check
@@ -188,6 +192,16 @@ class TestCudagraphDispatcher:
192
else:
193
assert rt_mode == CUDAGraphMode.NONE
194
195
+ # 5. valid_modes={NONE} always returns NONE even when keys exist
196
+ rt_mode, key = dispatcher.dispatch(
197
+ num_tokens=8,
198
+ uniform_decode=False,
199
+ has_lora=False,
200
+ valid_modes={CUDAGraphMode.NONE},
201
+ )
202
+ assert rt_mode == CUDAGraphMode.NONE
203
+ assert key == BatchDescriptor(num_tokens=8)
204
+
205
@pytest.mark.parametrize(
206
"cudagraph_mode_str,compilation_mode,expected_modes",
207
[
vllm/config/compilation.py
+6
-2
@@ -87,8 +87,12 @@ class CUDAGraphMode(enum.Enum):
87
def separate_routine(self) -> bool:
88
return isinstance(self.value, tuple)
89
90
- def valid_runtime_modes(self) -> bool:
91
- return self in [CUDAGraphMode.NONE, CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL]
90
+ @classmethod
91
+ def valid_runtime_modes(cls) -> frozenset["CUDAGraphMode"]:
92
+ return frozenset({cls.NONE, cls.PIECEWISE, cls.FULL})
93
+
94
+ def is_valid_runtime_mode(self) -> bool:
95
+ return self in CUDAGraphMode.valid_runtime_modes()
96
97
def __str__(self) -> str:
98
return self.name
vllm/forward_context.py
+1
-2
@@ -241,7 +241,7 @@ class ForwardContext:
241
additional_kwargs: dict[str, Any] = field(default_factory=dict)
242
243
def __post_init__(self):
244
- assert self.cudagraph_runtime_mode.valid_runtime_modes(), (
244
+ assert self.cudagraph_runtime_mode.is_valid_runtime_mode(), (
245
f"Invalid cudagraph runtime mode: {self.cudagraph_runtime_mode}"
246
)
247
@@ -347,7 +347,6 @@ def set_forward_context(
347
num_tokens_unpadded=num_tokens,
348
parallel_config=vllm_config.parallel_config,
349
allow_microbatching=False,
350
- allow_dp_padding=False,
350
)
351
assert num_tokens_across_dp is not None
352
dp_metadata = DPMetadata.make(
vllm/v1/cudagraph_dispatcher.py
+40
-22
@@ -1,5 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from collections.abc import Set as AbstractSet
4
from dataclasses import replace
5
from itertools import product
6
@@ -232,8 +233,9 @@ class CudagraphDispatcher:
233
num_tokens: int,
234
uniform_decode: bool = False,
235
has_lora: bool = False,
235
- disable_full: bool = False,
236
num_active_loras: int = 0,
237
+ valid_modes: AbstractSet[CUDAGraphMode] | None = None,
238
+ invalid_modes: AbstractSet[CUDAGraphMode] | None = None,
239
) -> tuple[CUDAGraphMode, BatchDescriptor]:
240
"""
241
Given conditions(e.g.,batch descriptor and if using piecewise only),
@@ -246,15 +248,29 @@ class CudagraphDispatcher:
248
uniform_decode: Whether the batch is uniform decode (i.e. uniform and query
249
length is uniform_decode_query_len).
250
has_lora: Whether LoRA is active.
249
- disable_full: If True, skip FULL cudagraph checks and
250
- return PIECEWISE or NONE only. (can be used for features like
251
- cascade attention that are not supported by full cudagraphs)
251
num_active_loras: Number of distinct active LoRA adapters.
252
+ valid_modes: Set of cudagraph modes that are allowed. None means
253
+ all modes are allowed.
254
+ invalid_modes: Set of cudagraph modes to exclude. Subtracted from
255
+ valid_modes to compute allowed modes. (e.g., {FULL} for
256
+ features like cascade attention not supported by full
257
+ cudagraphs). None means no modes are excluded.
258
"""
259
+ allowed_modes = valid_modes or CUDAGraphMode.valid_runtime_modes()
260
+
261
+ if invalid_modes:
262
+ allowed_modes -= invalid_modes
263
+
264
+ assert len(allowed_modes) >= 1, (
265
+ f"No allowed cudagraph modes: valid_modes={valid_modes}, "
266
+ f"invalid_modes={invalid_modes}"
267
+ )
268
+
269
if (
270
not self.keys_initialized
271
or self.cudagraph_mode == CUDAGraphMode.NONE
272
or num_tokens > self.compilation_config.max_cudagraph_capture_size
273
+ or allowed_modes <= {CUDAGraphMode.NONE}
274
):
275
return CUDAGraphMode.NONE, BatchDescriptor(num_tokens)
276
@@ -281,24 +297,26 @@ class CudagraphDispatcher:
297
num_tokens, uniform_decode, has_lora, effective_num_active_loras
298
)
299
284
- # check if key exists for full cudagraph
285
- # For pure FULL mode, keys are registered with uniform=False.
286
- batch_desc_to_check = batch_desc
287
- if self.cudagraph_mode == CUDAGraphMode.FULL:
288
- batch_desc_to_check = replace(batch_desc, uniform=False)
289
- if (
290
- not disable_full
291
- and batch_desc_to_check in self.cudagraph_keys[CUDAGraphMode.FULL]
292
- ):
293
- return CUDAGraphMode.FULL, batch_desc_to_check
294
-
295
- # also check if the relaxed key exists for more "general"
296
- # piecewise cudagraph
297
- batch_desc_to_check = replace(batch_desc, num_reqs=None, uniform=False)
298
- if batch_desc_to_check in self.cudagraph_keys[CUDAGraphMode.PIECEWISE]:
299
- return CUDAGraphMode.PIECEWISE, batch_desc_to_check
300
-
301
- # finally, just return no cudagraphs and a trivial batch descriptor
300
+ if CUDAGraphMode.FULL in allowed_modes:
301
+ # check if key exists for full cudagraph
302
+ # For pure FULL mode, keys are registered with uniform=False.
303
+ batch_desc_to_check = batch_desc
304
+ if self.cudagraph_mode == CUDAGraphMode.FULL:
305
+ batch_desc_to_check = replace(batch_desc, uniform=False)
306
+ if batch_desc_to_check in self.cudagraph_keys[CUDAGraphMode.FULL]:
307
+ return CUDAGraphMode.FULL, batch_desc_to_check
308
+
309
+ if CUDAGraphMode.PIECEWISE in allowed_modes:
310
+ # also check if the relaxed key exists for more "general"
311
+ # piecewise cudagraph
312
+ batch_desc_to_check = replace(batch_desc, num_reqs=None, uniform=False)
313
+ if batch_desc_to_check in self.cudagraph_keys[CUDAGraphMode.PIECEWISE]:
314
+ return CUDAGraphMode.PIECEWISE, batch_desc_to_check
315
+
316
+ assert CUDAGraphMode.NONE in allowed_modes, (
317
+ f"No matching cudagraph found and NONE is not in "
318
+ f"allowed_modes={allowed_modes}"
319
+ )
320
return CUDAGraphMode.NONE, BatchDescriptor(num_tokens)
321
322
def get_capture_descs(self) -> list[tuple[CUDAGraphMode, list[BatchDescriptor]]]:
vllm/v1/spec_decode/eagle.py
+48
-49
@@ -448,17 +448,10 @@ class SpecDecodeBaseProposer:
448
assert draft_indexer_metadata is not None
449
per_layer_attn_metadata[layer_name] = draft_indexer_metadata
450
451
- num_tokens_dp_padded, num_tokens_across_dp = self._pad_batch_across_dp(
452
- num_tokens_unpadded=num_tokens, num_tokens_padded=num_tokens
451
+ cudagraph_runtime_mode, num_input_tokens, num_tokens_across_dp = (
452
+ self._determine_batch_execution_and_padding(num_tokens)
453
)
454
455
- cudagraph_runtime_mode, batch_desc = self.cudagraph_dispatcher.dispatch(
456
- num_tokens_dp_padded
457
- )
458
- num_input_tokens = batch_desc.num_tokens
459
- if num_tokens_across_dp is not None:
460
- num_tokens_across_dp[self.dp_rank] = num_input_tokens
461
-
455
if self.supports_mm_inputs:
456
mm_embeds, is_mm_embed = mm_embed_inputs or (None, None)
457
@@ -549,16 +542,9 @@ class SpecDecodeBaseProposer:
542
# Generate the remaining draft tokens.
543
draft_token_ids_list = [draft_token_ids]
544
552
- batch_size_dp_padded, batch_size_across_dp = self._pad_batch_across_dp(
553
- num_tokens_unpadded=batch_size, num_tokens_padded=batch_size
554
- )
555
-
556
- cudagraph_runtime_mode, batch_desc = self.cudagraph_dispatcher.dispatch(
557
- batch_size_dp_padded
545
+ cudagraph_runtime_mode, input_batch_size, batch_size_across_dp = (
546
+ self._determine_batch_execution_and_padding(batch_size)
547
)
559
- input_batch_size = batch_desc.num_tokens
560
- if batch_size_across_dp is not None:
561
- batch_size_across_dp[self.dp_rank] = input_batch_size
548
549
common_attn_metadata.num_actual_tokens = batch_size
550
common_attn_metadata.max_query_len = 1
@@ -1568,19 +1554,11 @@ class SpecDecodeBaseProposer:
1554
self.num_speculative_tokens if not is_graph_capturing else 1
1555
):
1556
if fwd_idx <= 1:
1571
- num_tokens_dp_padded, num_tokens_across_dp = self._pad_batch_across_dp(
1572
- num_tokens_unpadded=num_tokens, num_tokens_padded=num_tokens
1573
- )
1574
- if use_cudagraphs:
1575
- cudagraph_runtime_mode, batch_desc = (
1576
- self.cudagraph_dispatcher.dispatch(num_tokens_dp_padded)
1557
+ cudagraph_runtime_mode, num_input_tokens, num_tokens_across_dp = (
1558
+ self._determine_batch_execution_and_padding(
1559
+ num_tokens, use_cudagraphs=use_cudagraphs
1560
)
1578
- num_input_tokens = batch_desc.num_tokens
1579
- else:
1580
- cudagraph_runtime_mode = CUDAGraphMode.NONE
1581
- num_input_tokens = num_tokens_dp_padded
1582
- if num_tokens_across_dp is not None:
1583
- num_tokens_across_dp[self.dp_rank] = num_input_tokens
1561
+ )
1562
1563
# Make sure to use EAGLE's own buffer during cudagraph capture.
1564
if (
@@ -1680,28 +1658,49 @@ class SpecDecodeBaseProposer:
1658
== 1
1659
), "All drafting layers should belong to the same kv cache group"
1660
1683
- def _pad_batch_across_dp(
1661
+ def _determine_batch_execution_and_padding(
1662
self,
1685
- num_tokens_unpadded: int,
1686
- num_tokens_padded: int,
1687
- ) -> tuple[int, torch.Tensor]:
1688
- # TODO(Flechman): support DBO ubatching
1689
- should_ubatch, num_toks_across_dp, _ = coordinate_batch_across_dp(
1690
- num_tokens_unpadded=num_tokens_unpadded,
1691
- parallel_config=self.vllm_config.parallel_config,
1692
- allow_microbatching=False,
1693
- allow_dp_padding=self.cudagraph_dispatcher.cudagraph_mode
1694
- != CUDAGraphMode.NONE,
1695
- num_tokens_padded=num_tokens_padded,
1696
- uniform_decode=None,
1697
- num_scheduled_tokens_per_request=None,
1663
+ num_tokens: int,
1664
+ use_cudagraphs: bool = True,
1665
+ ) -> tuple[CUDAGraphMode, int, torch.Tensor | None]:
1666
+ cudagraph_mode, batch_desc = self.cudagraph_dispatcher.dispatch(
1667
+ num_tokens,
1668
+ valid_modes=({CUDAGraphMode.NONE} if not use_cudagraphs else None),
1669
)
1699
- assert not should_ubatch, "DBO ubatching not implemented for EAGLE"
1670
+ num_tokens_padded = batch_desc.num_tokens
1671
+
1672
+ # Extra coordination when running data-parallel since we need to
1673
+ # coordinate across ranks
1674
+ # TODO(Flechman): support DBO ubatching
1675
+ should_ubatch, num_tokens_across_dp = False, None
1676
+ if self.vllm_config.parallel_config.data_parallel_size > 1:
1677
+ should_ubatch, num_tokens_across_dp, synced_cudagraph_mode = (
1678
+ coordinate_batch_across_dp(
1679
+ num_tokens_unpadded=num_tokens,
1680
+ parallel_config=self.vllm_config.parallel_config,
1681
+ allow_microbatching=False,
1682
+ num_tokens_padded=num_tokens_padded,
1683
+ cudagraph_mode=cudagraph_mode.value,
1684
+ )
1685
+ )
1686
+ assert not should_ubatch, "DBO ubatching not implemented for EAGLE"
1687
+
1688
+ # Extract DP-synced values
1689
+ if num_tokens_across_dp is not None:
1690
+ dp_rank = self.dp_rank
1691
+ num_tokens_padded = int(num_tokens_across_dp[dp_rank].item())
1692
+ # Re-dispatch with DP padding so we have the correct
1693
+ # batch_descriptor
1694
+ cudagraph_mode, batch_desc = self.cudagraph_dispatcher.dispatch(
1695
+ num_tokens_padded,
1696
+ valid_modes={CUDAGraphMode(synced_cudagraph_mode)},
1697
+ )
1698
+ # Assert to make sure the agreed upon token count is correct
1699
+ # otherwise num_tokens_across_dp will no-longer be valid
1700
+ assert batch_desc.num_tokens == num_tokens_padded
1701
+ num_tokens_across_dp[dp_rank] = num_tokens_padded
1702
1701
- num_tokens_dp_padded = num_tokens_padded
1702
- if num_toks_across_dp is not None:
1703
- num_tokens_dp_padded = int(num_toks_across_dp[self.dp_rank].item())
1704
- return num_tokens_dp_padded, num_toks_across_dp
1703
+ return cudagraph_mode, num_tokens_padded, num_tokens_across_dp
1704
1705
1706
class EagleProposer(SpecDecodeBaseProposer):
vllm/v1/worker/dp_utils.py
+18
-30
@@ -37,7 +37,6 @@ def _get_device_and_group(parallel_config: ParallelConfig):
37
38
def _run_ar(
39
should_ubatch: bool,
40
- should_dp_pad: bool,
40
orig_num_tokens_per_ubatch: int,
41
padded_num_tokens_per_ubatch: int,
42
cudagraph_mode: int,
@@ -46,12 +45,11 @@ def _run_ar(
45
dp_size = parallel_config.data_parallel_size
46
dp_rank = parallel_config.data_parallel_rank
47
device, group = _get_device_and_group(parallel_config)
49
- tensor = torch.zeros(5, dp_size, device=device, dtype=torch.int32)
48
+ tensor = torch.zeros(4, dp_size, device=device, dtype=torch.int32)
49
tensor[0][dp_rank] = orig_num_tokens_per_ubatch
50
tensor[1][dp_rank] = padded_num_tokens_per_ubatch
51
tensor[2][dp_rank] = 1 if should_ubatch else 0
53
- tensor[3][dp_rank] = 1 if should_dp_pad else 0
54
- tensor[4][dp_rank] = cudagraph_mode
52
+ tensor[3][dp_rank] = cudagraph_mode
53
dist.all_reduce(tensor, group=group)
54
return tensor
55
@@ -97,14 +95,13 @@ def _post_process_cudagraph_mode(tensor: torch.Tensor) -> int:
95
If any rank has NONE (0), all ranks use NONE.
96
This ensures all ranks send consistent values (all padded or all unpadded).
97
"""
100
- return int(tensor[4, :].min().item())
98
+ return int(tensor[3, :].min().item())
99
100
101
def _synchronize_dp_ranks(
102
num_tokens_unpadded: int,
103
num_tokens_padded: int,
104
should_attempt_ubatching: bool,
107
- should_attempt_dp_padding: bool,
105
cudagraph_mode: int,
106
parallel_config: ParallelConfig,
107
) -> tuple[bool, torch.Tensor | None, int]:
@@ -113,8 +110,8 @@ def _synchronize_dp_ranks(
110
run with microbatching or none of them do.
111
112
2. Determines the total number of tokens that each rank will run.
116
- When running microbatched or if should_attempt_dp_padding is True, all
117
- ranks will be padded out so that the run with the same number of tokens
113
+ When running microbatched or if cudagraph is enabled (synced across ranks),
114
+ all ranks will be padded out so that they run with the same number of tokens.
115
116
3. Synchronizes cudagraph_mode across ranks by taking the minimum.
117
@@ -133,29 +130,26 @@ def _synchronize_dp_ranks(
130
# will run and if we are using ubatching or not.
131
tensor = _run_ar(
132
should_ubatch=should_attempt_ubatching,
136
- should_dp_pad=should_attempt_dp_padding,
133
orig_num_tokens_per_ubatch=num_tokens_unpadded,
134
padded_num_tokens_per_ubatch=num_tokens_padded,
135
cudagraph_mode=cudagraph_mode,
136
parallel_config=parallel_config,
137
)
138
143
- should_dp_pad = bool(torch.all(tensor[3] == 1).item())
144
-
145
- # DP ranks should all have the same value for should_attempt_dp_padding.
146
- assert should_attempt_dp_padding == should_dp_pad
139
+ # Synchronize cudagraph_mode across ranks first (take min).
140
+ # This is needed before DP padding decision since we use the synced
141
+ # cudagraph mode to determine whether DP padding is needed.
142
+ synced_cudagraph_mode = _post_process_cudagraph_mode(tensor)
143
144
# Check conditions for microbatching
145
should_ubatch = _post_process_ubatch(tensor, parallel_config.num_ubatches)
146
151
- if should_ubatch and not should_dp_pad:
152
- logger.debug_once(
153
- "Microbatching has been triggered and requires DP padding. "
154
- "Enabling DP padding even though it has been explicitly "
155
- "disabled.",
156
- scope="global",
157
- )
158
- should_dp_pad = True
147
+ # DP padding is needed when cudagraph is enabled (synced across ranks)
148
+ # or when ubatching/DBO is active (ubatching requires uniform batch
149
+ # sizes across DP ranks currently).
150
+ # Use the synced runtime cudagraph mode rather than the compilation config
151
+ # so we can avoid padding when cudagraph is not enabled for this step.
152
+ should_dp_pad = synced_cudagraph_mode != 0 or should_ubatch
153
154
# Pad all DP ranks up to the maximum token count across ranks if
155
# should_dp_pad is True
@@ -164,16 +158,12 @@ def _synchronize_dp_ranks(
158
should_dp_pad,
159
)
160
167
- # Synchronize cudagraph_mode across ranks (take min)
168
- synced_cudagraph_mode = _post_process_cudagraph_mode(tensor)
169
-
161
return should_ubatch, num_tokens_after_padding, synced_cudagraph_mode
162
163
164
def coordinate_batch_across_dp(
165
num_tokens_unpadded: int,
166
allow_microbatching: bool,
176
- allow_dp_padding: bool,
167
parallel_config: ParallelConfig,
168
num_tokens_padded: int | None = None,
169
uniform_decode: bool | None = None,
@@ -187,7 +177,6 @@ def coordinate_batch_across_dp(
177
Args:
178
num_tokens_unpadded: Number of tokens without accounting for padding
179
allow_microbatching: If microbatching should be attempted
190
- allow_dp_padding: If all DP ranks should be padded up to the same value
180
parallel_config: The parallel config
181
num_tokens_padded: Number of tokens including any non-DP padding (CUDA graphs,
182
TP, etc)
@@ -195,15 +184,15 @@ def coordinate_batch_across_dp(
184
only contains single token decodes
185
num_scheduled_tokens_per_request: Only used if allow_microbatching is True. The
186
number of tokens per request.
198
- cudagraph_mode: The cudagraph mode for this rank (0=NONE, 1=PIECEWISE, 2=FULL)
187
+ cudagraph_mode: The cudagraph mode for this rank (0=NONE, 1=PIECEWISE, 2=FULL).
188
+ DP padding is enabled when synced cudagraph mode across ranks is not NONE.
189
190
Returns: tuple[
191
ubatch_slices: if this is set then all DP ranks have agreed to
192
microbatch
193
num_tokens_after_padding: A tensor containing the total number of
194
tokens per-microbatch for each DP rank including padding. Will be
205
- padded up to the max value across all DP ranks when allow_dp_padding
206
- is True.
195
+ padded up to the max value across all DP ranks when cudagraph is enabled.
196
synced_cudagraph_mode: The synchronized cudagraph mode (min across ranks)
197
]
198
@@ -231,7 +220,6 @@ def coordinate_batch_across_dp(
220
num_tokens_unpadded,
221
num_tokens_padded,
222
should_attempt_ubatching,
234
- allow_dp_padding,
223
cudagraph_mode,
224
parallel_config,
225
)
vllm/v1/worker/gpu_model_runner.py
+10
-20
@@ -2300,7 +2300,7 @@ class GPUModelRunner(
2300
)
2301
# Dispatch for the decoder portion of the model.
2302
_, batch_desc = self.cudagraph_dispatcher.dispatch(
2303
- num_logits, disable_full=True
2303
+ num_logits, invalid_modes={CUDAGraphMode.FULL}
2304
)
2305
num_logits_padded = batch_desc.num_tokens
2306
logits_indices_padded = self.kv_sharing_fast_prefill_logits_indices[
@@ -3174,20 +3174,19 @@ class GPUModelRunner(
3174
has_lora = num_active_loras > 0 if force_has_lora is None else force_has_lora
3175
3176
num_tokens_padded = self._pad_for_sequence_parallelism(num_tokens)
3177
- dispatch_cudagraph = (
3178
- lambda num_tokens, disable_full: self.cudagraph_dispatcher.dispatch(
3177
+
3178
+ def dispatch_cudagraph(num_tokens, disable_full=False, valid_modes=None):
3179
+ return self.cudagraph_dispatcher.dispatch(
3180
num_tokens=num_tokens,
3181
has_lora=has_lora,
3182
uniform_decode=uniform_decode,
3182
- disable_full=disable_full,
3183
num_active_loras=num_active_loras,
3184
+ valid_modes={CUDAGraphMode.NONE} if force_eager else valid_modes,
3185
+ invalid_modes={CUDAGraphMode.FULL} if disable_full else None,
3186
)
3185
- if not force_eager
3186
- else (CUDAGraphMode.NONE, BatchDescriptor(num_tokens_padded))
3187
- )
3187
3188
cudagraph_mode, batch_descriptor = dispatch_cudagraph(
3190
- num_tokens_padded, use_cascade_attn or has_encoder_output
3189
+ num_tokens_padded, disable_full=use_cascade_attn or has_encoder_output
3190
)
3191
num_tokens_padded = batch_descriptor.num_tokens
3192
if self.compilation_config.pass_config.enable_sp:
@@ -3204,20 +3203,11 @@ class GPUModelRunner(
3203
# across ranks
3204
should_ubatch, num_tokens_across_dp = False, None
3205
if self.vllm_config.parallel_config.data_parallel_size > 1:
3207
- # Disable DP padding when running eager to avoid excessive padding when
3208
- # running prefills. This lets us set cudagraph_mode="NONE" on the prefiller
3209
- # in a P/D setup and still use CUDA graphs (enabled by this padding) on the
3210
- # decoder.
3211
- allow_dp_padding = (
3212
- self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE
3213
- )
3214
-
3206
should_ubatch, num_tokens_across_dp, synced_cudagraph_mode = (
3207
coordinate_batch_across_dp(
3208
num_tokens_unpadded=num_tokens,
3209
parallel_config=self.parallel_config,
3210
allow_microbatching=allow_microbatching,
3220
- allow_dp_padding=allow_dp_padding,
3211
num_tokens_padded=num_tokens_padded,
3212
uniform_decode=uniform_decode,
3213
num_scheduled_tokens_per_request=num_scheduled_tokens_np,
@@ -3232,7 +3222,7 @@ class GPUModelRunner(
3222
# Re-dispatch with DP padding so we have the correct batch_descriptor
3223
cudagraph_mode, batch_descriptor = dispatch_cudagraph(
3224
num_tokens_padded,
3235
- disable_full=synced_cudagraph_mode <= CUDAGraphMode.PIECEWISE.value,
3225
+ valid_modes={CUDAGraphMode(synced_cudagraph_mode)},
3226
)
3227
# Assert to make sure the agreed upon token count is correct otherwise
3228
# num_tokens_across_dp will no-longer be valid
@@ -4724,7 +4714,7 @@ class GPUModelRunner(
4714
4715
assert (
4716
cudagraph_runtime_mode is None
4727
- or cudagraph_runtime_mode.valid_runtime_modes()
4717
+ or cudagraph_runtime_mode.is_valid_runtime_mode()
4718
)
4719
4720
# If cudagraph_mode.decode_mode() == FULL and
@@ -5336,7 +5326,7 @@ class GPUModelRunner(
5326
):
5327
assert (
5328
cudagraph_runtime_mode != CUDAGraphMode.NONE
5339
- and cudagraph_runtime_mode.valid_runtime_modes()
5329
+ and cudagraph_runtime_mode.is_valid_runtime_mode()
5330
), f"Invalid cudagraph runtime mode: {cudagraph_runtime_mode}"
5331
5332
if not batch_descriptors: