From 16f73d0da2fdfefde40d704f3af56aa41a6d5654 Mon Sep 17 00:00:00 2001 From: tzhouam Date: Sat, 3 Oct 2026 21:23:18 +0800 Subject: [PATCH] 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 --- .../worker/test_gpu_generation_model_runner.py | 18 ++++++++++++++++++ .../worker/gpu_generation_model_runner.py | 5 ++--- 2 files changed, 20 insertions(+), 3 deletions(-) diff --git a/tests/worker/test_gpu_generation_model_runner.py b/tests/worker/test_gpu_generation_model_runner.py index 0cd0a1ab3..847e16154 100644 --- a/tests/worker/test_gpu_generation_model_runner.py +++ b/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)] diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index c1f4fb8eb..7c2cf2163 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/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