Honor Omni stage runner selection in shared vLLM config

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

tests/engine/test_arg_utils.py

@@ -20,7 +20,11 @@ import pytest
 from omegaconf import OmegaConf
 from pydantic import ValidationError
 from transformers import PretrainedConfig, Qwen3OmniMoeConfig
+from vllm.config import VllmConfig
 from vllm.engine.arg_utils import EngineArgs
+from vllm.exceptions import VLLMValidationError
+from vllm.v1.engine.input_processor import InputProcessor
+from vllm.v1.worker.gpu.sample.logit_bias import LogitBiasState
 
 from tests.helpers.mock import patch_hf_snapshot_download
 from vllm_omni.config.model import OmniModelConfig
@@ -32,6 +36,36 @@ from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorMode
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+@pytest.mark.parametrize("use_v2", [False, True])
+def test_stage_runner_selects_matching_input_processor_validator(monkeypatch, use_v2):
+    from vllm.sampling_params import SamplingParams
+
+    from vllm_omni.model_executor.models.audex.cfg import AudexCFGLogitsProcessor
+
+    # Upstream's process-wide preference must not override a stage's deploy choice.
+    monkeypatch.setattr("vllm.envs.VLLM_USE_V2_MODEL_RUNNER", not use_v2)
+    config = VllmConfig()
+    model_config = object.__new__(OmniModelConfig)
+    model_config.use_v2_model_runner = use_v2
+    model_config.runner_type = "generate"
+    model_config.logits_processors = [LogitBiasState if use_v2 else AudexCFGLogitsProcessor]
+    config.model_config = model_config
+    assert config.use_v2_model_runner is use_v2
+    processor = object.__new__(InputProcessor)
+    processor.vllm_config, processor.model_config = config, model_config
+    validator = processor._build_logits_processors_params_validator()
+    validator(SamplingParams(extra_args={"cfg_role": "cond", "cfg_scale": 1.5}))
+    if not use_v2:
+        with pytest.raises(VLLMValidationError, match="cfg_role"):
+            validator(SamplingParams(extra_args={"cfg_role": "invalid"}))
+
+
+@pytest.mark.parametrize("use_v2", [False, True])
+def test_plain_vllm_config_retains_upstream_runner_preference(monkeypatch, use_v2):
+    monkeypatch.setattr("vllm.envs.VLLM_USE_V2_MODEL_RUNNER", use_v2)
+    assert VllmConfig().use_v2_model_runner is use_v2
+
+
 def test_sync_config_is_omni():
     """Ensure create_model_config gives the right type."""
     cfg = OmniEngineArgs().create_model_config()

vllm_omni/patch.py

@@ -10,6 +10,7 @@ from functools import cached_property
 import torch
 from aenum import extend_enum
 from vllm.config import ModelConfig as _OriginalModelConfig
+from vllm.config import VllmConfig as _OriginalVllmConfig
 from vllm.inputs import TokensPrompt as _OriginalTokensPrompt
 from vllm.model_executor.layers.rotary_embedding import (
     MRotaryEmbedding as _OriginalMRotaryEmbedding,
@@ -29,6 +30,23 @@ from vllm_omni.request import OmniRequest, OmniStreamingUpdate
 
 _PATCH_LOGGER = logging.getLogger("vllm_omni.patch")
 
+# Omni's deploy config selects a runner per stage. vLLM 0.31's config property
+# otherwise defaults to v2 independently of the runner selected by our workers,
+# making shared consumers such as InputProcessor choose the wrong API.
+_runner_property = _OriginalVllmConfig.use_v2_model_runner
+if not getattr(_runner_property.fget, "_omni_stage_runner", False):
+
+    def _patched_use_v2_model_runner(self, _upstream=_runner_property.fget):
+        # Lazy import keeps the package's patch/config initialization acyclic.
+        from vllm_omni.config.model import OmniModelConfig
+
+        if isinstance(self.model_config, OmniModelConfig):
+            return self.model_config.use_v2_model_runner
+        return _upstream(self)
+
+    _patched_use_v2_model_runner._omni_stage_runner = True
+    _OriginalVllmConfig.use_v2_model_runner = property(_patched_use_v2_model_runner)
+
 # =============================================================================
 # Patch ModelConfig.is_mm_prefix_lm to support omni-specific models
 # =============================================================================
@@ -390,8 +408,8 @@ def _patch_chat_template_registry():
         )
 
         if "qwen3_omni_moe" not in _MODEL_TYPE_TO_CHAT_TEMPLATE_FALLBACK:
-            _MODEL_TYPE_TO_CHAT_TEMPLATE_FALLBACK["qwen3_omni_moe"] = (
-                lambda _: CHAT_TEMPLATES_DIR / "template_chatml.jinja"
+            _MODEL_TYPE_TO_CHAT_TEMPLATE_FALLBACK["qwen3_omni_moe"] = lambda _: (
+                CHAT_TEMPLATES_DIR / "template_chatml.jinja"
             )
     except ImportError:
         pass