[Rebase] Align vLLM-Omni with vLLM v0.31.0 (#8459)

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()

tests/engine/test_arg_utils_shared_fields.py

@@ -8,6 +8,7 @@ from __future__ import annotations
 from dataclasses import fields
 
 import pytest
+from vllm.config import AuxOutputConfig
 
 from vllm_omni.engine.arg_utils import (
     SHARED_FIELDS,
@@ -49,3 +50,26 @@ def test_internal_blacklist_keys_derived_from_orchestrator():
     """
     blacklist = internal_blacklist_keys()
     assert blacklist == orchestrator_field_names() - SHARED_FIELDS
+
+
+@pytest.mark.parametrize(
+    "aux_args",
+    [
+        {"enable_return_routed_experts": True},
+        {"aux_output_config": {"enable_return_routed_experts": True}},
+        {"aux_output_config": AuxOutputConfig(enable_return_routed_experts=True)},
+    ],
+)
+@pytest.mark.parametrize("use_v2_model_runner", [False, True])
+@pytest.mark.parametrize("is_npu", [False, True])
+def test_aux_outputs_require_a_supported_omni_stage_runner(aux_args, use_v2_model_runner, is_npu, monkeypatch):
+    # Upstream runner selection reads this environment variable; Omni workers
+    # select their runner from the separate per-stage flag instead.
+    monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
+    monkeypatch.setattr("vllm_omni.engine.arg_utils.current_omni_platform.is_npu", lambda: is_npu)
+    if use_v2_model_runner or is_npu:
+        args = OmniEngineArgs(model="unused", use_v2_model_runner=use_v2_model_runner, **aux_args)
+        assert args.aux_output_config.enabled
+    else:
+        with pytest.raises(ValueError, match="use_v2_model_runner=True for this Omni stage"):
+            OmniEngineArgs(model="unused", **aux_args)

tests/engine/test_async_omni_engine_input.py

@@ -168,6 +168,37 @@ def test_build_add_request_message_preserves_additional_information(mocker: Mock
     output_processor.add_request.assert_not_called()
 
 
+def test_build_add_request_message_preserves_kv_hints_after_omni_payload_upgrade(mocker):
+    from vllm.v1.engine.input_processor import InputProcessor
+    from vllm.v1.kv_hints import KvHintsEnvelope
+
+    engine = object.__new__(AsyncOmniEngine)
+    params = SamplingParams(max_tokens=8)
+    engine.default_sampling_params_list = [params]
+    engine.stage_metadata = [StageRuntimeInfo(final_output=False, final_output_type=None, stage_type="llm")]
+    engine.supported_tasks = ("generate",)
+    processor = mocker.create_autospec(InputProcessor, instance=True)
+    hints = KvHintsEnvelope("1", "hint-message", [])
+
+    def process(**kwargs):
+        request = _make_engine_core_request()
+        request.kv_hints = kwargs["kv_hints"]
+        return request
+
+    processor.process_inputs.side_effect = process
+    engine.input_processor = processor
+    msg = engine._build_add_request_message(
+        request_id="req-1",
+        prompt={"prompt_token_ids": [1], "additional_information": {"speaker": "voice"}},
+        sampling_params_list=[params],
+        kv_hints=hints,
+    )
+    assert processor.process_inputs.call_args.kwargs["kv_hints"] is hints
+    assert isinstance(msg.prompt, OmniEngineCoreRequest)
+    assert msg.prompt.kv_hints is hints
+    assert msg.prompt.additional_information is not None
+
+
 def test_build_add_request_message_injects_global_id_before_prompt_transform(mocker: MockerFixture):
     engine = object.__new__(AsyncOmniEngine)
     params = SamplingParams(max_tokens=8)

tests/engine/test_async_omni_engine_stage_init.py

@@ -296,11 +296,21 @@ def test_stage_engine_core_client_module_reload_keeps_forward_refs_deferred():
     """Regression test for forward references in make_async_mp_client."""
     import vllm_omni.engine.stage_engine_core_client as client_mod
 
-    importlib.reload(client_mod)
+    # Snapshot the module namespace so a reload cannot strand other test
+    # modules that captured the pre-reload class objects (isinstance checks
+    # against the stale class would then fail). ``importlib.reload`` mutates
+    # the module dict in place; restoring it leaves the runtime exactly as it
+    # was before this test ran.
+    saved = dict(client_mod.__dict__)
+    try:
+        importlib.reload(client_mod)
 
-    assert client_mod.StageEngineCoreClientBase.make_async_mp_client.__annotations__["return"] == (
-        "StageEngineCoreClient | DPLBStageEngineCoreClient"
-    )
+        assert client_mod.StageEngineCoreClientBase.make_async_mp_client.__annotations__["return"] == (
+            "StageEngineCoreClient | DPLBStageEngineCoreClient"
+        )
+    finally:
+        client_mod.__dict__.clear()
+        client_mod.__dict__.update(saved)
 
 
 def test_async_omni_engine_initialize_stages_passes_log_stats_and_client_config_to_runtime(monkeypatch):

tests/engine/test_omni_request.py

@@ -13,13 +13,16 @@ substitutable with the base class — including base-style calls that pass
 
 import inspect
 
+import msgspec
 import numpy as np
 import pytest
 import torch
 from vllm.sampling_params import SamplingParams
+from vllm.v1.engine import EngineCoreRequest
+from vllm.v1.kv_hints import KvHintAction, KvHintsEnvelope
 from vllm.v1.request import Request
 
-from vllm_omni.engine import PromptEmbedsPayload
+from vllm_omni.engine import OmniEngineCoreRequest, PromptEmbedsPayload
 from vllm_omni.request import OmniRequest, OmniStreamingUpdate
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -111,3 +114,26 @@ def test_model_intermediate_buffer_round_trips_to_streaming_update():
     assert update is not None
     assert request.model_intermediate_buffer is buffer
     assert update.model_intermediate_buffer is buffer
+
+
+def test_kv_hints_survive_omni_upgrade_wire_transport_and_scheduler_conversion():
+    hints = KvHintsEnvelope("1", "hint-message", [KvHintAction("a", "prefetch", "1", {"blocks": [7]})])
+    core = EngineCoreRequest(
+        request_id="hint-request",
+        prompt_token_ids=[1, 2],
+        mm_features=None,
+        sampling_params=SamplingParams(max_tokens=2),
+        pooling_params=None,
+        arrival_time=0.0,
+        lora_request=None,
+        cache_salt=None,
+        data_parallel_rank=None,
+        session_id="session",
+        kv_hints=hints,
+    )
+    upgraded = OmniEngineCoreRequest.from_request(core)
+    decoded = msgspec.msgpack.decode(msgspec.msgpack.encode(upgraded), type=OmniEngineCoreRequest)
+    scheduled = OmniRequest.from_engine_core_request(decoded, block_hasher=None)
+    assert decoded.session_id == "session"
+    assert scheduled.kv_hints == hints
+    assert scheduled.session_id == "session"

tests/engine/test_omni_request_output.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Tests for OmniRequestOutput class."""
 
+import numpy as np
 import pytest
 from PIL import Image
 from vllm.outputs import CompletionOutput, RequestOutput
@@ -157,6 +158,7 @@ class TestMsgpackRoundTrip:
         )
 
         source = _make_text_request_output(text="round trip")
+        source.prompt_token_id_logprobs = np.array([[-0.5, -1.5], [-0.2, -2.2]], dtype=np.float32)
         out = OmniRequestOutput.from_stage_output(
             source,
             request_id="rt-1",
@@ -178,6 +180,7 @@ class TestMsgpackRoundTrip:
         assert decoded.outputs[0].text == "round trip"
         assert decoded.prompt == "test prompt"
         assert decoded.prompt_token_ids == [1, 2, 3]
+        np.testing.assert_array_equal(decoded.prompt_token_id_logprobs, source.prompt_token_id_logprobs)
 
     def test_round_trip_diffusion_output(self):
         """A diffusion output round-trips without data loss."""
@@ -302,6 +305,16 @@ class TestRequestOutputFieldParity:
         assert omni.num_cached_tokens is None
         assert omni.num_cache_creation_tokens is None
 
+    def test_fixed_token_scores_survive_stage_output_wrapping(self) -> None:
+        source = _make_text_request_output()
+        source.prompt_token_id_logprobs = np.array([[-0.5, -1.5], [-0.2, -2.2]], dtype=np.float32)
+
+        omni = OmniRequestOutput.from_stage_output(source, stage_id=0)
+        forwarded = OmniRequestOutput.from_stage_output(omni, stage_id=1)
+
+        np.testing.assert_array_equal(forwarded.prompt_token_id_logprobs, source.prompt_token_id_logprobs)
+        assert OmniRequestOutput(request_id="unscored").prompt_token_id_logprobs is None
+
     def test_transfer_param_fields_are_copied_from_the_stage_output(self) -> None:
         """Declaring the fields is not enough — they must also be copied.
 

tests/engine/test_orchestrator_error_handling.py

@@ -1282,6 +1282,7 @@ async def test_cache_reset_rpc_failure_is_nonfatal_and_visits_all_replicas(metho
     [
         ("reset_mm_cache", {}, True),
         ("sleep", {}, True),
+        ("release_kv_cache_memory", {}, True),
         ("pause_scheduler", {"clear_cache": True}, True),
         ("pause_scheduler", {"clear_cache": False}, False),
         ("reset_prefix_cache", {}, False),

tests/engine/test_output_processor_mrv2_text.py

@@ -7,8 +7,9 @@ from unittest.mock import MagicMock
 
 import pytest
 import torch
-from vllm.sampling_params import RequestOutputKind
-from vllm.v1.engine import FinishReason
+from vllm.sampling_params import RequestOutputKind, SamplingParams
+from vllm.v1.engine import EngineCoreRequest, FinishReason
+from vllm.v1.engine.logprobs import LogprobsProcessor
 
 from vllm_omni.engine import OmniEngineCoreOutput
 from vllm_omni.engine.output_modality import OutputModality
@@ -16,6 +17,7 @@ from vllm_omni.engine.output_processor import (
     MultimodalOutputProcessor,
     OmniRequestState,
 )
+from vllm_omni.outputs import OmniRequestOutput
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -109,6 +111,37 @@ def test_text_tokens_are_detokenized_when_mrv2_ar_output_has_pooling_payload():
     assert not state.mm_accumulated.is_empty
 
 
+def test_fixed_token_scores_survive_output_processing_and_stage_wrapping():
+    processor, state = _make_processor(OutputModality.TEXT, _Detokenizer())
+    request = EngineCoreRequest(
+        request_id="r",
+        prompt_token_ids=[1, 2, 3],
+        mm_features=None,
+        sampling_params=SamplingParams(prompt_logprob_token_ids=[4, 5]),
+        pooling_params=None,
+        arrival_time=0.0,
+        lora_request=None,
+        cache_salt=None,
+        data_parallel_rank=None,
+    )
+    state.logprobs_processor = LogprobsProcessor.from_new_request(None, request)
+    scores = torch.tensor([[-0.5, -1.5], [-0.2, -2.2]])
+    processed = processor.process_outputs(
+        [
+            OmniEngineCoreOutput(
+                request_id="r",
+                new_token_ids=[42],
+                prompt_token_id_logprobs=scores,
+                finish_reason=FinishReason.STOP,
+            )
+        ]
+    )
+
+    (raw_output,) = processed.request_outputs
+    final_output = OmniRequestOutput.from_stage_output(raw_output, stage_id=0)
+    torch.testing.assert_close(torch.from_numpy(final_output.prompt_token_id_logprobs), scores)
+
+
 @pytest.mark.parametrize("streaming", [False, True])
 def test_generation_stage_mm_delivered_only_with_terminal(streaming):
     # Non-terminal chunks accumulate; the terminal output carries the audio.

tests/engine/test_stage_engine_args.py

@@ -97,7 +97,6 @@ _OMNI_ONLY_LLM_STAGE_ENGINE_FIELDS = frozenset(
         "enable_multithread_weight_load",
         "env",
         "has_sampling_extra_args",
-        "log_level",
         "log_stats",
         "model_arch",
         "model_subdir",

tests/engine/test_stage_engine_core_client.py

@@ -5,11 +5,19 @@
 from __future__ import annotations
 
 from types import SimpleNamespace
+from unittest.mock import patch
 
 import pytest
+from vllm import SamplingParams
+from vllm.v1.engine.core_client import AsyncMPClient, DPLBAsyncMPClient
 from vllm.v1.engine.exceptions import EngineDeadError
 
-from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient
+from vllm_omni.engine.stage_engine_core_client import (
+    DPLBStageEngineCoreClient,
+    StageEngineCoreClient,
+    StageEngineCoreClientBase,
+)
+from vllm_omni.engine.stage_init_utils import StageMetadata
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -30,3 +38,46 @@ def test_check_health_raises_when_resources_engine_dead():
     client = _make_client(engine_dead=True)
     with pytest.raises(EngineDeadError, match="engine core is dead"):
         client.check_health()
+
+
+@pytest.mark.parametrize("dp_size, external_lb", [(1, False), (2, False), (2, True)])
+@pytest.mark.parametrize("renderer", [None, object()])
+def test_stage_client_forwards_renderer_to_the_actual_upstream_constructor(dp_size, external_lb, renderer):
+    config = SimpleNamespace(
+        model_config=None,
+        parallel_config=SimpleNamespace(data_parallel_size=dp_size, data_parallel_external_lb=external_lb),
+    )
+    metadata = StageMetadata(
+        stage_id=0,
+        stage_type="llm",
+        engine_output_type="text",
+        is_comprehension=True,
+        requires_multimodal_data=False,
+        engine_input_source=[],
+        final_output=True,
+        final_output_type="text",
+        default_sampling_params=SamplingParams(),
+        custom_process_input_func=None,
+        model_stage=None,
+        runtime_cfg=None,
+    )
+    internal_lb = dp_size > 1 and not external_lb
+    parent = DPLBAsyncMPClient if internal_lb else AsyncMPClient
+    # Keep Omni construction/factory routing real; intercept only upstream's
+    # subprocess/socket boundary, using its current signature via autospec.
+    with patch.object(parent, "__init__", autospec=True, return_value=None) as init:
+        addresses = {"input_address": "tcp://192.0.2.1:1"}
+        client = StageEngineCoreClientBase.make_async_mp_client(
+            config, object, metadata=metadata, client_addresses=addresses, renderer=renderer
+        )
+        assert isinstance(client, DPLBStageEngineCoreClient if internal_lb else StageEngineCoreClient)
+        init.assert_called_once_with(
+            client,
+            config,
+            object,
+            log_stats=False,
+            client_addresses=addresses,
+            client_count=1,
+            client_index=0,
+            renderer=renderer,
+        )

tests/engine/test_stage_pool_collective_rpc.py

@@ -83,6 +83,27 @@ def test_collective_rpc_control_method_reraises_worker_error():
     asyncio.run(run())
 
 
+@pytest.mark.cpu
+def test_kv_memory_release_uses_engine_core_and_propagates_pause_errors():
+    async def run() -> None:
+        release = AsyncMock(return_value=None)
+        pool, client = _make_pool(release_kv_cache_memory_async=release)
+        await pool.collective_rpc(0, "release_kv_cache_memory", timeout=2.0)
+        release.assert_awaited_once_with()
+        client.collective_rpc_async.assert_not_awaited()
+
+        release.side_effect = RuntimeError("requires a completed pause first")
+        with pytest.raises(RuntimeError, match="completed pause"):
+            await pool.collective_rpc(0, "release_kv_cache_memory")
+
+        del client.release_kv_cache_memory_async
+        result = await pool.collective_rpc(0, "release_kv_cache_memory")
+        assert result["supported"] is False
+        client.collective_rpc_async.assert_not_awaited()
+
+    asyncio.run(run())
+
+
 @pytest.mark.cpu
 def test_collective_rpc_reset_caches_use_control_helpers():
     async def run() -> None: