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 11, 18:45 UTC
[BugFix] Fix `assert batch_descriptor.num_tokens == num_tokens_padded` (#30173)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Lucas Wilkinson committed
Dec 9, 2025 at 10:36 UTC
56037dfa2fc10a5f28e919d2c26e4c2abbd313a6
6 files changed
+65
-33
tests/v1/cudagraph/test_cudagraph_dispatch.py
+2
-2
@@ -161,10 +161,10 @@ class TestCudagraphDispatcher:
161
assert rt_mode == CUDAGraphMode.NONE
162
assert key == BatchDescriptor(num_tokens=15)
163
164
- # 4. Cascade attention should have a fall back mode
164
+ # 4. disable_full should have a fall back mode (e.g., cascade attention)
165
desc_full_exact = BatchDescriptor(num_tokens=8, uniform=False)
166
rt_mode, key = dispatcher.dispatch(
167
- num_tokens=8, uniform_decode=False, has_lora=False, use_cascade_attn=True
167
+ num_tokens=8, uniform_decode=False, has_lora=False, disable_full=True
168
)
169
if "PIECEWISE" in cudagraph_mode_str: # string contains check
170
assert rt_mode == CUDAGraphMode.PIECEWISE
vllm/forward_context.py
+1
-1
@@ -292,7 +292,7 @@ def set_forward_context(
292
if num_tokens_across_dp is None:
293
assert ubatch_slices is None
294
assert num_tokens is not None
295
- _, num_tokens_across_dp = coordinate_batch_across_dp(
295
+ _, num_tokens_across_dp, _ = coordinate_batch_across_dp(
296
num_tokens_unpadded=num_tokens,
297
parallel_config=vllm_config.parallel_config,
298
allow_microbatching=False,
vllm/v1/cudagraph_dispatcher.py
+2
-2
@@ -145,7 +145,7 @@ class CudagraphDispatcher:
145
num_tokens: int,
146
uniform_decode: bool,
147
has_lora: bool,
148
- use_cascade_attn: bool = False,
148
+ disable_full: bool = False,
149
) -> tuple[CUDAGraphMode, BatchDescriptor]:
150
"""
151
Given conditions(e.g.,batch descriptor and if using cascade attention),
@@ -165,7 +165,7 @@ class CudagraphDispatcher:
165
)
166
relaxed_batch_desc = batch_desc.relax_for_mixed_batch_cudagraphs()
167
168
- if not use_cascade_attn:
168
+ if not disable_full:
169
# check if key exists for full cudagraph
170
if batch_desc in self.cudagraph_keys[CUDAGraphMode.FULL]:
171
return CUDAGraphMode.FULL, batch_desc
vllm/v1/spec_decode/eagle.py
+1
-1
@@ -1258,7 +1258,7 @@ class EagleProposer:
1258
num_tokens_padded: int,
1259
) -> tuple[int, torch.Tensor]:
1260
# TODO(Flechman): support DBO ubatching
1261
- should_ubatch, num_toks_across_dp = coordinate_batch_across_dp(
1261
+ should_ubatch, num_toks_across_dp, _ = coordinate_batch_across_dp(
1262
num_tokens_unpadded=num_tokens_unpadded,
1263
parallel_config=self.vllm_config.parallel_config,
1264
allow_microbatching=False,
vllm/v1/worker/dp_utils.py
+37
-12
@@ -40,16 +40,18 @@ def _run_ar(
40
should_dp_pad: bool,
41
orig_num_tokens_per_ubatch: int,
42
padded_num_tokens_per_ubatch: int,
43
+ cudagraph_mode: int,
44
parallel_config: ParallelConfig,
45
) -> torch.Tensor:
46
dp_size = parallel_config.data_parallel_size
47
dp_rank = parallel_config.data_parallel_rank
48
device, group = _get_device_and_group(parallel_config)
48
- tensor = torch.zeros(4, dp_size, device=device, dtype=torch.int32)
49
+ tensor = torch.zeros(5, dp_size, device=device, dtype=torch.int32)
50
tensor[0][dp_rank] = orig_num_tokens_per_ubatch
51
tensor[1][dp_rank] = padded_num_tokens_per_ubatch
52
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
55
dist.all_reduce(tensor, group=group)
56
return tensor
57
@@ -89,13 +91,23 @@ def _post_process_dp_padding(tensor: torch.Tensor, should_dp_pad: bool) -> torch
91
return num_tokens_across_dp.cpu()
92
93
94
+def _post_process_cudagraph_mode(tensor: torch.Tensor) -> int:
95
+ """
96
+ Synchronize cudagraph_mode across DP ranks by taking the minimum.
97
+ If any rank has NONE (0), all ranks use NONE.
98
+ This ensures all ranks send consistent values (all padded or all unpadded).
99
+ """
100
+ return int(tensor[4, :].min().item())
101
+
102
+
103
def _synchronize_dp_ranks(
104
num_tokens_unpadded: int,
105
num_tokens_padded: int,
106
should_attempt_ubatching: bool,
107
should_attempt_dp_padding: bool,
108
+ cudagraph_mode: int,
109
parallel_config: ParallelConfig,
98
-) -> tuple[bool, torch.Tensor | None]:
110
+) -> tuple[bool, torch.Tensor | None, int]:
111
"""
112
1. Decides if each DP rank is going to microbatch. Either all ranks
113
run with microbatching or none of them do.
@@ -104,10 +116,13 @@ def _synchronize_dp_ranks(
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
118
119
+ 3. Synchronizes cudagraph_mode across ranks by taking the minimum.
120
+
121
Returns: tuple[
122
should_ubatch: Are all DP ranks going to microbatch
123
num_tokens_after_padding: A tensor containing the total number of
124
tokens per-microbatch for each DP rank including any DP padding.
125
+ synced_cudagraph_mode: The synchronized cudagraph mode (min across ranks)
126
]
127
128
"""
@@ -121,6 +136,7 @@ def _synchronize_dp_ranks(
136
should_dp_pad=should_attempt_dp_padding,
137
orig_num_tokens_per_ubatch=num_tokens_unpadded,
138
padded_num_tokens_per_ubatch=num_tokens_padded,
139
+ cudagraph_mode=cudagraph_mode,
140
parallel_config=parallel_config,
141
)
142
@@ -148,7 +164,10 @@ def _synchronize_dp_ranks(
164
should_dp_pad,
165
)
166
151
- return should_ubatch, num_tokens_after_padding
167
+ # Synchronize cudagraph_mode across ranks (take min)
168
+ synced_cudagraph_mode = _post_process_cudagraph_mode(tensor)
169
+
170
+ return should_ubatch, num_tokens_after_padding, synced_cudagraph_mode
171
172
173
def coordinate_batch_across_dp(
@@ -159,7 +178,8 @@ def coordinate_batch_across_dp(
178
num_tokens_padded: int | None = None,
179
uniform_decode: bool | None = None,
180
num_scheduled_tokens_per_request: np.ndarray | None = None,
162
-) -> tuple[bool, torch.Tensor | None]:
181
+ cudagraph_mode: int = 0,
182
+) -> tuple[bool, torch.Tensor | None, int]:
183
"""
184
Coordinates amongst all DP ranks to determine if and how the full batch
185
should be split into microbatches.
@@ -175,6 +195,7 @@ def coordinate_batch_across_dp(
195
only contains single token decodes
196
num_scheduled_tokens_per_request: Only used if allow_microbatching is True. The
197
number of tokens per request.
198
+ cudagraph_mode: The cudagraph mode for this rank (0=NONE, 1=PIECEWISE, 2=FULL)
199
200
Returns: tuple[
201
ubatch_slices: if this is set then all DP ranks have agreed to
@@ -183,12 +204,13 @@ def coordinate_batch_across_dp(
204
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.
207
+ synced_cudagraph_mode: The synchronized cudagraph mode (min across ranks)
208
]
209
210
"""
211
if parallel_config.data_parallel_size == 1:
212
# Early exit.
191
- return False, None
213
+ return False, None, cudagraph_mode
214
215
# If the caller has explicitly enabled microbatching.
216
should_attempt_ubatching = False
@@ -204,12 +226,15 @@ def coordinate_batch_across_dp(
226
if num_tokens_padded is None:
227
num_tokens_padded = num_tokens_unpadded
228
207
- (should_ubatch, num_tokens_after_padding) = _synchronize_dp_ranks(
208
- num_tokens_unpadded,
209
- num_tokens_padded,
210
- should_attempt_ubatching,
211
- allow_dp_padding,
212
- parallel_config,
229
+ (should_ubatch, num_tokens_after_padding, synced_cudagraph_mode) = (
230
+ _synchronize_dp_ranks(
231
+ num_tokens_unpadded,
232
+ num_tokens_padded,
233
+ should_attempt_ubatching,
234
+ allow_dp_padding,
235
+ cudagraph_mode,
236
+ parallel_config,
237
+ )
238
)
239
215
- return (should_ubatch, num_tokens_after_padding)
240
+ return (should_ubatch, num_tokens_after_padding, synced_cudagraph_mode)
vllm/v1/worker/gpu_model_runner.py
+22
-15
@@ -2788,17 +2788,19 @@ class GPUModelRunner(
2788
)
2789
2790
dispatch_cudagraph = (
2791
- lambda num_tokens: self.cudagraph_dispatcher.dispatch(
2791
+ lambda num_tokens, disable_full: self.cudagraph_dispatcher.dispatch(
2792
num_tokens=num_tokens,
2793
has_lora=has_lora,
2794
- use_cascade_attn=use_cascade_attn,
2794
uniform_decode=uniform_decode,
2795
+ disable_full=disable_full,
2796
)
2797
if not force_eager
2798
else (CUDAGraphMode.NONE, BatchDescriptor(num_tokens_padded))
2799
)
2800
2801
- cudagraph_mode, batch_descriptor = dispatch_cudagraph(num_tokens_padded)
2801
+ cudagraph_mode, batch_descriptor = dispatch_cudagraph(
2802
+ num_tokens_padded, use_cascade_attn
2803
+ )
2804
num_tokens_padded = batch_descriptor.num_tokens
2805
2806
# Extra coordination when running data-parallel since we need to coordinate
@@ -2813,23 +2815,28 @@ class GPUModelRunner(
2815
self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE
2816
)
2817
2816
- should_ubatch, num_tokens_across_dp = coordinate_batch_across_dp(
2817
- num_tokens_unpadded=num_tokens,
2818
- parallel_config=self.parallel_config,
2819
- allow_microbatching=allow_microbatching,
2820
- allow_dp_padding=allow_dp_padding,
2821
- num_tokens_padded=num_tokens_padded,
2822
- uniform_decode=uniform_decode,
2823
- num_scheduled_tokens_per_request=num_scheduled_tokens_np,
2818
+ should_ubatch, num_tokens_across_dp, synced_cudagraph_mode = (
2819
+ coordinate_batch_across_dp(
2820
+ num_tokens_unpadded=num_tokens,
2821
+ parallel_config=self.parallel_config,
2822
+ allow_microbatching=allow_microbatching,
2823
+ allow_dp_padding=allow_dp_padding,
2824
+ num_tokens_padded=num_tokens_padded,
2825
+ uniform_decode=uniform_decode,
2826
+ num_scheduled_tokens_per_request=num_scheduled_tokens_np,
2827
+ cudagraph_mode=cudagraph_mode.value,
2828
+ )
2829
)
2830
2826
- # Extract DP padding if there is any
2831
+ # Extract DP-synced values
2832
if num_tokens_across_dp is not None:
2833
dp_rank = self.parallel_config.data_parallel_rank
2834
num_tokens_padded = int(num_tokens_across_dp[dp_rank].item())
2830
-
2831
- # Re-dispatch with DP padding
2832
- cudagraph_mode, batch_descriptor = dispatch_cudagraph(num_tokens_padded)
2835
+ # Re-dispatch with DP padding so we have the correct batch_descriptor
2836
+ cudagraph_mode, batch_descriptor = dispatch_cudagraph(
2837
+ num_tokens_padded,
2838
+ disable_full=synced_cudagraph_mode <= CUDAGraphMode.PIECEWISE.value,
2839
+ )
2840
# Assert to make sure the agreed upon token count is correct otherwise
2841
# num_tokens_across_dp will no-longer be valid
2842
assert batch_descriptor.num_tokens == num_tokens_padded