fix: match vLLM 0.31 generation profiling API

Accept the upstream randomize_inputs profiling flag while preserving structured codec inputs during generation-stage dummy runs.

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
This commit is contained in:
tzhouam committed 2026-10-03 21:23:18 +08:00
1 parent 2974f8be59
commit 16f73d0da2
2 files changed
+20 -3

No files matched your search

@@ -869,7 +869,6 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
# in code2wav warmup. The expert-balance rationale behind
# randomize_inputs targets text-MoE models, which generation
# stages are not. The kwarg stays accepted for vLLM API compat.
self.maybe_randomize_inputs(input_ids, inputs_embeds),
set_forward_context(
attn_metadata,
self.vllm_config,
@@ -949,7 +948,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
return hidden_states, None
def profile_run(self) -> None:
def profile_run(self, randomize_inputs: bool = False) -> None:
# Profile with multimodal encoder & encoder cache.
if self.supports_mm_inputs:
mm_config = self.model_config.multimodal_config
@@ -1002,7 +1001,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
self.encoder_cache[f"tmp_{i}"] = output
# Add `is_profile` here to pre-allocate communication buffers
hidden_states, _ = self._dummy_run(self.max_num_tokens, is_profile=True)
hidden_states, _ = self._dummy_run(self.max_num_tokens, is_profile=True, randomize_inputs=randomize_inputs)
output = None
self._sync_device()
del hidden_states