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, 15:00 UTC
[BugFix][MM]support VLLM_RANDOMIZE_DP_DUMMY_INPUTS (#30472)
Signed-off-by: Xingyu Liu <charlotteliu12x@gmail.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Xingyu Liu committed
Dec 11, 2025 at 13:00 UTC
90d6cf921fe623524f618740616a6cf494d4a8df
1 file changed
+23
-6
vllm/v1/worker/gpu_model_runner.py
+23
-6
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import functools
5
import gc
6
import itertools
7
import time
@@ -3892,19 +3893,21 @@ class GPUModelRunner(
3893
return {}
3894
3895
@contextmanager
3895
- def maybe_randomize_inputs(self, input_ids: torch.Tensor):
3896
+ def maybe_randomize_inputs(
3897
+ self, input_ids: torch.Tensor | None, inputs_embeds: torch.Tensor | None
3898
+ ):
3899
"""
3900
Randomize input_ids if VLLM_RANDOMIZE_DP_DUMMY_INPUTS is set.
3901
This is to help balance expert-selection
3902
- during profile_run
3903
- during DP rank dummy run
3904
"""
3905
+
3906
dp_size = self.vllm_config.parallel_config.data_parallel_size
3907
randomize_inputs = envs.VLLM_RANDOMIZE_DP_DUMMY_INPUTS and dp_size > 1
3908
if not randomize_inputs:
3909
yield
3906
- else:
3907
- import functools
3910
+ elif input_ids is not None:
3911
3912
@functools.cache
3913
def rand_input_ids() -> torch.Tensor:
@@ -3912,13 +3915,27 @@ class GPUModelRunner(
3915
self.input_ids.gpu,
3916
low=0,
3917
high=self.model_config.get_vocab_size(),
3915
- dtype=input_ids.dtype,
3918
)
3919
3918
- logger.debug_once("Randomizing dummy data for DP Rank")
3920
+ logger.debug_once("Randomizing dummy input_ids for DP Rank")
3921
input_ids.copy_(rand_input_ids()[: input_ids.size(0)], non_blocking=True)
3922
yield
3923
input_ids.fill_(0)
3924
+ else:
3925
+
3926
+ @functools.cache
3927
+ def rand_inputs_embeds() -> torch.Tensor:
3928
+ return torch.randn_like(
3929
+ self.inputs_embeds.gpu,
3930
+ )
3931
+
3932
+ assert inputs_embeds is not None
3933
+ logger.debug_once("Randomizing dummy inputs_embeds for DP Rank")
3934
+ inputs_embeds.copy_(
3935
+ rand_inputs_embeds()[: inputs_embeds.size(0)], non_blocking=True
3936
+ )
3937
+ yield
3938
+ inputs_embeds.fill_(0)
3939
3940
def _get_mm_dummy_batch(
3941
self,
@@ -4167,7 +4184,7 @@ class GPUModelRunner(
4184
num_tokens_across_dp[:] = num_tokens_padded
4185
4186
with (
4170
- self.maybe_randomize_inputs(input_ids),
4187
+ self.maybe_randomize_inputs(input_ids, inputs_embeds),
4188
set_forward_context(
4189
attn_metadata,
4190
self.vllm_config,