@tduncan / Nvidia-Optimization-Stack-VLLM / commits / e5e38ba9b7

[Model Runner V2] Support randomized dummy inputs (#58411)

Signed-off-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>

Robert Shaw committed Oct 1, 2026 at 03:52 UTC e5e38ba9b7d18f9746d989e389a51a94b0f96f6f

4 files changed +43 -9

vllm/distributed/elastic_ep/elastic_execute.py

+6 -1

@@ -721,7 +721,12 @@ class ElasticEPScalingExecutor:

721 self._suppress_eplb(),

722 serving_state,

723 ):

724 - runner._dummy_run(runner.max_num_tokens, is_profile=True, skip_eplb=True)

724 + runner._dummy_run(

725 + runner.max_num_tokens,

726 + is_profile=True,

727 + skip_eplb=True,

728 + randomize_inputs=self.worker.randomize_dummy_inputs,

729 + )

730 self.worker.compile_or_warm_up_model()

731

732 lock_workspace()

vllm/v1/worker/gpu/model_runner.py

+11 -2

@@ -805,6 +805,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

805 skip_eplb: bool = False,

806 is_profile: bool = False,

807 valid_dummy_state_slots: bool = False,

808 + randomize_inputs: bool = False,

809 **kwargs,

810 ) -> tuple[torch.Tensor | None, torch.Tensor | None]:

811 if skip_attn and not is_profile:

@@ -878,6 +879,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

879 is_profile=is_profile,

880 context_len=context_len,

881 valid_dummy_state_slots=valid_dummy_state_slots,

882 + randomize_inputs=randomize_inputs,

883 )

884 self.kv_connector.set_disabled(False)

885

@@ -968,7 +970,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

970 self.pooling_runner.dummy_pooler_run(hidden_states)

971

972 @torch.inference_mode()

971 - def profile_run(self) -> None:

973 + def profile_run(self, randomize_inputs: bool = False) -> None:

974 if self.supports_mm_inputs and self.is_first_pp_rank:

975 mm_config = self.model_config.multimodal_config

976 if mm_config is not None and not mm_config.skip_mm_profiling:

@@ -981,7 +983,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):

983 )

984

985 hidden_states, sample_hidden_states = self._dummy_run(

984 - self.max_num_tokens, skip_attn=True, is_profile=True

986 + self.max_num_tokens,

987 + skip_attn=True,

988 + is_profile=True,

989 + randomize_inputs=randomize_inputs,

990 )

991

992 # Only run sampler/pooler on last PP rank (non-last ranks return None).

@@ -1696,6 +1701,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1701 is_profile: bool = False,

1702 context_len: int = 0,

1703 valid_dummy_state_slots: bool = False,

1704 + randomize_inputs: bool = False,

1705 ) -> ModelRunnerOutput | IntermediateTensors | None:

1706 if not dummy_run:

1707 # Update the request states.

@@ -1814,6 +1820,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1820 # so MoE memory is measured and MoE kernels are exercised.

1821 is_padding=not is_profile,

1822 )

1823 + if randomize_inputs:

1824 + # All-zero input_ids route every token to the same experts.

1825 + input_batch.input_ids.random_(0, self.vocab_size)

1826 if self.pcp_manager is not None:

1827 input_batch = self.pcp_manager.prepare_inputs_to_capture(input_batch)

1828 if skip_attn_for_dummy_run:

vllm/v1/worker/gpu_model_runner.py

+5 -2

@@ -5735,6 +5735,9 @@ class GPUModelRunner(

5735 - during profile_run

5736 - during DP rank dummy run

5737 """

5738 + # The worker also resolves this flag and passes randomize_inputs, but

5739 + # V1 has internal dummy runs (e.g. CUDA graph capture) that don't go

5740 + # through the worker, so keep the check here too.

5741 dp_size = self.vllm_config.parallel_config.data_parallel_size

5742 randomize_inputs = randomize_inputs or (

5743 envs.VLLM_RANDOMIZE_DP_DUMMY_INPUTS and dp_size > 1

@@ -6420,7 +6423,7 @@ class GPUModelRunner(

6423 max_task = max(output_size.items(), key=lambda x: x[1])[0]

6424 return self._dummy_pooler_run_task(hidden_states, max_task)

6425

6423 - def profile_run(self) -> None:

6426 + def profile_run(self, randomize_inputs: bool = False) -> None:

6427 # Profile with multimodal encoder & encoder cache.

6428 if self.supports_mm_inputs:

6429 mm_config = self.model_config.multimodal_config

@@ -6481,7 +6484,7 @@ class GPUModelRunner(

6484

6485 # Add `is_profile` here to pre-allocate communication buffers

6486 hidden_states, last_hidden_states = self._dummy_run(

6484 - self.max_num_tokens, is_profile=True

6487 + self.max_num_tokens, is_profile=True, randomize_inputs=randomize_inputs

6488 )

6489 if get_pp_group().is_last_rank:

6490 if self.is_pooling_model:

vllm/v1/worker/gpu_worker.py

+21 -4

@@ -613,7 +613,7 @@ class Worker(WorkerBase):

613 if kv_cache_memory_bytes := self.cache_config.kv_cache_memory_bytes:

614 # still need a profile run which compiles the model for

615 # max_num_batched_tokens

616 - self.model_runner.profile_run()

616 + self.model_runner.profile_run(randomize_inputs=self.randomize_dummy_inputs)

617

618 msg = (

619 f"Initial free memory {format_gib(self.init_snapshot.free_memory)} "

@@ -650,7 +650,7 @@ class Worker(WorkerBase):

650 # memory_profiling measures, restoring the original limit.

651 self._scoped_allocator_max_split(max_split_size_mb=20),

652 ):

653 - self.model_runner.profile_run()

653 + self.model_runner.profile_run(randomize_inputs=self.randomize_dummy_inputs)

654

655 # Profile CUDA graph memory if graphs will be captured.

656 # ROCm is included: #44825 moved the profiler to

@@ -885,7 +885,12 @@ class Worker(WorkerBase):

885 # We skip EPLB here since we don't want to record dummy metrics

886 for size in sorted(warmup_sizes, reverse=True):

887 logger.info("Compile and warming up model for size %d", size)

888 - self.model_runner._dummy_run(size, skip_eplb=True, remove_lora=False)

888 + self.model_runner._dummy_run(

889 + size,

890 + skip_eplb=True,

891 + remove_lora=False,

892 + randomize_inputs=self.randomize_dummy_inputs,

893 + )

894 self.model_runner.maybe_remove_all_loras(self.model_runner.lora_config)

895

896 # Warmup and tune the kernels used during model execution before

@@ -1387,9 +1392,21 @@ class Worker(WorkerBase):

1392 # Recreate it so the next profile_prefix is honored.

1393 self.profiler = None

1394

1395 + @property

1396 + def randomize_dummy_inputs(self) -> bool:

1397 + # Not cached: elastic EP rewrites parallel_config in place on reconfigure.

1398 + return (

1399 + envs.VLLM_RANDOMIZE_DP_DUMMY_INPUTS

1400 + and self.parallel_config.data_parallel_size > 1

1401 + )

1402 +

1403 def execute_dummy_batch(self) -> None:

1404 num_tokens = getattr(self.model_runner, "uniform_decode_query_len", 1)

1392 - self.model_runner._dummy_run(num_tokens, uniform_decode=True)

1405 + self.model_runner._dummy_run(

1406 + num_tokens,

1407 + uniform_decode=True,

1408 + randomize_inputs=self.randomize_dummy_inputs,

1409 + )

1410

1411 def add_lora(self, lora_request: LoRARequest) -> bool:

1412 return self.model_runner.add_lora(lora_request)