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
[Bugfix] Make DP padding optional in coordinate_batch_across_dp (#26375)
Signed-off-by: Sage Moore <sage@neuralmagic.com>
Sage Moore committed
Oct 10, 2025 at 07:53 UTC
ae9d0e7da5bfde55c95b6bd2af0e0d3b02754101
3 files changed
+123
-42
vllm/forward_context.py
+14
-1
@@ -12,6 +12,7 @@ import torch
12
import vllm.envs as envs
13
from vllm.config import CUDAGraphMode, ParallelConfig, VllmConfig
14
from vllm.logger import init_logger
15
+from vllm.v1.worker.dp_utils import coordinate_batch_across_dp
16
from vllm.v1.worker.ubatch_utils import UBatchSlices
17
18
if TYPE_CHECKING:
@@ -278,7 +279,19 @@ def set_forward_context(
279
if vllm_config.parallel_config.data_parallel_size > 1 and (
280
attn_metadata is not None or num_tokens is not None
281
):
281
- assert num_tokens_across_dp is not None
282
+ # If num_tokens_across_dp hasn't already been initialized, then
283
+ # initialize it here. Both DP padding and Microbatching will be
284
+ # disabled.
285
+ if num_tokens_across_dp is None:
286
+ assert ubatch_slices is None
287
+ assert num_tokens is not None
288
+ _, num_tokens_across_dp = coordinate_batch_across_dp(
289
+ num_tokens_unpadded=num_tokens,
290
+ parallel_config=vllm_config.parallel_config,
291
+ allow_microbatching=False,
292
+ allow_dp_padding=False,
293
+ )
294
+ assert num_tokens_across_dp is not None
295
dp_metadata = DPMetadata.make(
296
vllm_config.parallel_config, num_tokens or 0, num_tokens_across_dp
297
)
vllm/v1/worker/dp_utils.py
+80
-26
@@ -7,7 +7,7 @@ import torch
7
import torch.distributed as dist
8
9
from vllm.config import ParallelConfig
10
-from vllm.distributed.parallel_state import get_dp_group
10
+from vllm.distributed.parallel_state import get_dp_group, is_global_first_rank
11
from vllm.logger import init_logger
12
from vllm.platforms import current_platform
13
from vllm.v1.worker.ubatch_utils import (
@@ -37,6 +37,7 @@ def _get_device_and_group(parallel_config: ParallelConfig):
37
38
def _run_ar(
39
should_ubatch: bool,
40
+ should_dp_pad: bool,
41
orig_num_tokens_per_ubatch: int,
42
padded_num_tokens_per_ubatch: int,
43
parallel_config: ParallelConfig,
@@ -44,10 +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)
47
- tensor = torch.zeros(3, 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
52
+ tensor[3][dp_rank] = 1 if should_dp_pad else 0
53
dist.all_reduce(tensor, group=group)
54
return tensor
55
@@ -72,10 +74,26 @@ def _post_process_ubatch(tensor: torch.Tensor) -> bool:
74
return should_ubatch
75
76
77
+def _post_process_dp_padding(tensor: torch.Tensor, should_dp_pad: bool) -> torch.Tensor:
78
+ num_tokens_across_dp = tensor[1, :]
79
+ if should_dp_pad:
80
+ # If DP padding is enabled, ensure that each rank is processing the same number
81
+ # of tokens
82
+ max_num_tokens = int(num_tokens_across_dp.max().item())
83
+ return torch.tensor(
84
+ [max_num_tokens] * len(num_tokens_across_dp),
85
+ device="cpu",
86
+ dtype=torch.int32,
87
+ )
88
+ else:
89
+ return num_tokens_across_dp.cpu()
90
+
91
+
92
def _synchronize_dp_ranks(
93
num_tokens_unpadded: int,
94
num_tokens_padded: int,
95
should_attempt_ubatching: bool,
96
+ should_attempt_dp_padding: bool,
97
parallel_config: ParallelConfig,
98
) -> tuple[bool, Optional[torch.Tensor]]:
99
"""
@@ -83,57 +101,88 @@ def _synchronize_dp_ranks(
101
run with microbatching or none of them do.
102
103
2. Determines the total number of tokens that each rank will run.
86
- All ranks will be padded out so that the run with the same number
87
- of tokens
104
+ When running microbatched or if should_attempt_dp_padding is True, all
105
+ ranks will be padded out so that the run with the same number of tokens
106
107
Returns: tuple[
108
should_ubatch: Are all DP ranks going to microbatch
109
num_tokens_after_padding: A tensor containing the total number of
92
- tokens per-microbatch for each DP rank including padding.
110
+ tokens per-microbatch for each DP rank including any DP padding.
111
]
112
113
"""
114
assert num_tokens_padded >= num_tokens_unpadded
115
98
- # First we coordinate between the DP ranks via an All Reduce
116
+ # Coordinate between the DP ranks via an All Reduce
117
# to determine the total number of tokens that each rank
118
# will run and if we are using ubatching or not.
119
tensor = _run_ar(
120
should_ubatch=should_attempt_ubatching,
121
+ should_dp_pad=should_attempt_dp_padding,
122
orig_num_tokens_per_ubatch=num_tokens_unpadded,
123
padded_num_tokens_per_ubatch=num_tokens_padded,
124
parallel_config=parallel_config,
125
)
126
108
- # Ensure that each rank is processing the same nuber of tokens
109
- num_tokens_across_dp = tensor[1, :]
110
- max_num_tokens = int(num_tokens_across_dp.max().item())
111
- num_tokens_after_padding = torch.tensor(
112
- [max_num_tokens] * len(num_tokens_across_dp), device="cpu", dtype=torch.int32
113
- )
127
+ should_dp_pad = bool(torch.all(tensor[3] == 1).item())
128
+
129
+ # DP ranks should all have the same value for should_attempt_dp_padding.
130
+ assert should_attempt_dp_padding == should_dp_pad
131
132
+ # Check conditions for microbatching
133
should_ubatch = _post_process_ubatch(tensor)
134
135
+ if should_ubatch and not should_dp_pad:
136
+ if is_global_first_rank():
137
+ logger.debug(
138
+ "Microbatching has been triggered and requires DP padding. "
139
+ "Enabling DP padding even though it has been explicitly "
140
+ "disabled."
141
+ )
142
+ should_dp_pad = True
143
+
144
+ # Pad all DP ranks up to the maximum token count across ranks if
145
+ # should_dp_pad is True
146
+ num_tokens_after_padding = _post_process_dp_padding(
147
+ tensor,
148
+ should_dp_pad,
149
+ )
150
+
151
return should_ubatch, num_tokens_after_padding
152
153
154
def coordinate_batch_across_dp(
121
- num_scheduled_tokens_per_request: np.ndarray,
155
num_tokens_unpadded: int,
123
- num_tokens_padded: int,
124
- parallel_config: ParallelConfig,
156
allow_microbatching: bool,
126
- uniform_decode: bool,
157
+ allow_dp_padding: bool,
158
+ parallel_config: ParallelConfig,
159
+ num_tokens_padded: Optional[int] = None,
160
+ uniform_decode: Optional[bool] = None,
161
+ num_scheduled_tokens_per_request: Optional[np.ndarray] = None,
162
) -> tuple[Optional[UBatchSlices], Optional[torch.Tensor]]:
163
"""
164
Coordinates amongst all DP ranks to determine if and how the full batch
165
should be split into microbatches.
166
167
+ Args:
168
+ num_tokens_unpadded: Number of tokens without accounting for padding
169
+ allow_microbatching: If microbatching should be attempted
170
+ allow_dp_padding: If all DP ranks should be padded up to the same value
171
+ parallel_config: The parallel config
172
+ num_tokens_padded: Number of tokens including any non-DP padding (CUDA graphs,
173
+ TP, etc)
174
+ uniform_decode: Only used if allow_microbatching is True. True if the batch
175
+ only contains single token decodes
176
+ num_scheduled_tokens_per_request: Only used if allow_microbatching is True. The
177
+ number of tokens per request.
178
+
179
Returns: tuple[
180
ubatch_slices: if this is set then all DP ranks have agreed to
181
microbatch
182
num_tokens_after_padding: A tensor containing the total number of
136
- tokens per-microbatch for each DP rank including padding.
183
+ tokens per-microbatch for each DP rank including padding. Will be
184
+ padded up to the max value across all DP ranks when allow_dp_padding
185
+ is True.
186
]
187
188
"""
@@ -141,21 +190,25 @@ def coordinate_batch_across_dp(
190
# Early exit.
191
return None, None
192
144
- # Check preconditions for microbatching
145
- should_attempt_ubatching = check_ubatch_thresholds(
146
- parallel_config,
147
- num_tokens_unpadded,
148
- uniform_decode=uniform_decode,
149
- )
193
+ # If the caller has explicitly enabled microbatching.
194
+ should_attempt_ubatching = False
195
+ if allow_microbatching:
196
+ # Check preconditions for microbatching
197
+ assert uniform_decode is not None
198
+ should_attempt_ubatching = check_ubatch_thresholds(
199
+ parallel_config,
200
+ num_tokens_unpadded,
201
+ uniform_decode=uniform_decode,
202
+ )
203
151
- # If the caller has explicitly disabled microbatching.
152
- if not allow_microbatching:
153
- should_attempt_ubatching = False
204
+ if num_tokens_padded is None:
205
+ num_tokens_padded = num_tokens_unpadded
206
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,
213
)
214
@@ -170,6 +223,7 @@ def coordinate_batch_across_dp(
223
assert num_tokens_after_padding is not None
224
token_split_point = int(num_tokens_after_padding[0].item()) // 2
225
226
+ assert num_scheduled_tokens_per_request is not None
227
ubatch_slices = create_ubatch_slices(
228
num_scheduled_tokens_per_request, token_split_point
229
)
vllm/v1/worker/gpu_model_runner.py
+29
-15
@@ -1178,13 +1178,21 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1178
uniform_decode = (
1179
max_num_scheduled_tokens == self.uniform_decode_query_len
1180
) and (total_num_scheduled_tokens == num_reqs * max_num_scheduled_tokens)
1181
+
1182
+ # Disable DP padding when running eager to avoid excessive padding when
1183
+ # running prefills. This lets us set enforce_eager on the prefiller in
1184
+ # a P/D setup and still use CUDA graphs (enabled by this padding) on the
1185
+ # decoder.
1186
+ allow_dp_padding = self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE
1187
+
1188
ubatch_slices, num_tokens_across_dp = coordinate_batch_across_dp(
1182
- num_scheduled_tokens,
1183
- num_tokens_unpadded,
1184
- num_tokens_padded,
1185
- self.parallel_config,
1186
- True,
1187
- uniform_decode,
1189
+ num_tokens_unpadded=num_tokens_unpadded,
1190
+ parallel_config=self.parallel_config,
1191
+ allow_microbatching=True,
1192
+ allow_dp_padding=allow_dp_padding,
1193
+ num_tokens_padded=num_tokens_padded,
1194
+ uniform_decode=uniform_decode,
1195
+ num_scheduled_tokens_per_request=num_scheduled_tokens,
1196
)
1197
1198
self.seq_lens.np[:num_reqs] = (
@@ -2436,12 +2444,13 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2444
use_cascade_attn,
2445
) = self._prepare_inputs(scheduler_output)
2446
2447
+ dp_rank = self.parallel_config.data_parallel_rank
2448
if ubatch_slices:
2449
assert num_tokens_across_dp is not None
2441
- num_input_tokens = int(num_tokens_across_dp[0].item())
2450
+ num_input_tokens = int(num_tokens_across_dp[dp_rank].item())
2451
self.pad_out_ubatch_slice(ubatch_slices, num_input_tokens)
2452
elif num_tokens_across_dp is not None:
2444
- num_input_tokens = int(num_tokens_across_dp[0].item())
2453
+ num_input_tokens = int(num_tokens_across_dp[dp_rank].item())
2454
else:
2455
num_input_tokens = self._get_num_input_tokens(
2456
scheduler_output.total_num_scheduled_tokens
@@ -3256,19 +3265,24 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
3265
num_scheduled_tokens = np.array(num_scheduled_tokens_list, dtype=np.int32)
3266
total_num_scheduled_tokens = int(num_scheduled_tokens.sum())
3267
3268
+ # Disable DP padding when running eager
3269
+ allow_dp_padding = self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE
3270
+
3271
# We currently only microbatch if the number of tokens is
3272
# over a certain threshold.
3273
ubatch_slices, num_tokens_across_dp = coordinate_batch_across_dp(
3262
- num_scheduled_tokens,
3263
- total_num_scheduled_tokens,
3264
- total_num_scheduled_tokens,
3265
- self.vllm_config.parallel_config,
3266
- allow_microbatching,
3267
- uniform_decode,
3274
+ num_tokens_unpadded=total_num_scheduled_tokens,
3275
+ parallel_config=self.vllm_config.parallel_config,
3276
+ allow_microbatching=allow_microbatching,
3277
+ allow_dp_padding=allow_dp_padding,
3278
+ num_tokens_padded=total_num_scheduled_tokens,
3279
+ uniform_decode=uniform_decode,
3280
+ num_scheduled_tokens_per_request=num_scheduled_tokens,
3281
)
3282
num_tokens_after_padding = num_tokens
3283
if num_tokens_across_dp is not None:
3271
- num_tokens_after_padding = int(num_tokens_across_dp[0])
3284
+ dp_rank = self.parallel_config.data_parallel_rank
3285
+ num_tokens_after_padding = int(num_tokens_across_dp[dp_rank])
3286
3287
attn_metadata: Optional[PerLayerAttnMetadata] = None
3288