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>

tests/worker/test_gpu_generation_model_runner.py

@@ -3,6 +3,7 @@
 
 import contextlib
 from types import SimpleNamespace
+from unittest.mock import Mock
 
 import pytest
 import torch
@@ -22,6 +23,23 @@ class _StopExecutionError(Exception):
     """Halts execute_model at the model boundary."""
 
 
+@pytest.mark.parametrize("randomize_inputs", [False, True])
+def test_generation_profile_accepts_upstream_randomize_inputs(randomize_inputs):
+    """Release workers can pass the profiling flag to codec-stage runners."""
+    runner = object.__new__(GPUGenerationModelRunner)
+    runner.supports_mm_inputs = False
+    runner.max_num_tokens = 128
+    runner.encoder_cache = {"temporary": object()}
+    runner._dummy_run = Mock(return_value=(torch.zeros(1), None))
+    runner._sync_device = Mock()
+
+    runner.profile_run(randomize_inputs=randomize_inputs)
+
+    runner._dummy_run.assert_called_once_with(128, is_profile=True, randomize_inputs=randomize_inputs)
+    runner._sync_device.assert_called_once_with()
+    assert runner.encoder_cache == {}
+
+
 class _DummyInputBatch:
     def __init__(self, num_reqs: int = 1):
         self.req_ids = [f"req-{i + 1}" for i in range(num_reqs)]

vllm_omni/worker/gpu_generation_model_runner.py

@@ -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