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] Fix DP/EP hang (#25906)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Lucas Wilkinson committed
Sep 30, 2025 at 00:18 UTC
23194d83e8f2a6783b0d8c275f5f8a22faab9aec
1 file changed
+15
-3
vllm/v1/worker/gpu_model_runner.py
+15
-3
@@ -3075,13 +3075,19 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
3075
# We currently only microbatch if the number of tokens is
3076
# over a certain threshold.
3077
if self.parallel_config.enable_dbo and allow_microbatching:
3078
- ubatch_slices, num_tokens_after_padding = ubatch_split(
3078
+ ubatch_slices, ubatch_num_tokens_after_padding = ubatch_split(
3079
num_scheduled_tokens,
3080
total_num_scheduled_tokens,
3081
total_num_scheduled_tokens,
3082
uniform_decode=uniform_decode,
3083
vllm_config=self.vllm_config,
3084
)
3085
+ # Currently when DBO is enabled `ubatch_split` returns
3086
+ # the num_tokens_after_padding for a single ubatch, but we have 2
3087
+ # TODO(sage,lucas): this is cruft that should be addressed in the
3088
+ # padding refactor.
3089
+ if ubatch_num_tokens_after_padding is not None:
3090
+ num_tokens_after_padding = ubatch_num_tokens_after_padding * 2
3091
3092
# If we failed to microbatch, currently need to resynchronize
3093
# TODO(lucas,sage): we should be able to avoid this second sync by
@@ -3198,7 +3204,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
3204
3205
# filter out the valid batch descriptor
3206
_cg_mode, batch_descriptor = self.cudagraph_dispatcher.dispatch(
3201
- BatchDescriptor(num_tokens=num_tokens,
3207
+ BatchDescriptor(num_tokens=num_tokens_after_padding,
3208
uniform_decode=uniform_decode)) \
3209
if not is_profile else (CUDAGraphMode.NONE, None)
3210
if cudagraph_runtime_mode is not None:
@@ -3212,7 +3218,13 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
3218
cudagraph_runtime_mode = _cg_mode
3219
3220
if ubatch_slices is not None:
3215
- num_tokens = num_tokens // 2
3221
+ # Adjust values to reflect a single ubatch.
3222
+ # TODO(sage,lucas): this is cruft that should be addressed in
3223
+ # the padding refactor.
3224
+ num_tokens_after_padding = ubatch_slices[0].num_tokens
3225
+ if num_tokens_across_dp is not None:
3226
+ num_tokens_across_dp[:] = num_tokens_after_padding
3227
+
3228
with self.maybe_randomize_inputs(input_ids), set_forward_context(
3229
attn_metadata,
3230
self.vllm_config,