[Model] Aura support - Non async chunk path (#4257)

Signed-off-by: Rein Yang <ruiruyang2@gmail.com>
Signed-off-by: R2-Y <ruiruyang2@gmail.com>

tests/config/test_pipeline_registry.py

@@ -30,6 +30,7 @@ class TestCentralRegistryDeclarations:
         assert "qwen2_5_omni" in _OMNI_PIPELINES
         assert "qwen2_5_omni_thinker_only" in _OMNI_PIPELINES
         assert "qwen3_omni_moe" in _OMNI_PIPELINES
+        assert "aura_omni" in _OMNI_PIPELINES
         assert "qwen3_tts" in _OMNI_PIPELINES
 
 
@@ -45,6 +46,20 @@ class TestLazyLoading:
         assert pipeline.model_type == "qwen3_omni_moe"
         assert pipeline.model_arch == "Qwen3OmniMoeForConditionalGeneration"
 
+    def test_getitem_loads_aura_omni_pipeline(self):
+        pipeline = _PIPELINE_REGISTRY["aura_omni"]
+        assert pipeline.model_type == "aura_omni"
+        assert [stage.model_stage for stage in pipeline.stages] == [
+            "asr",
+            "aura",
+            "qwen3_tts",
+            "code2wav",
+        ]
+        assert pipeline.stages[0].model_arch == "Qwen3ASRForConditionalGeneration"
+        assert pipeline.stages[1].model_arch == "AuraQwen3VLForConditionalGeneration"
+        assert pipeline.stages[2].model_arch == "Qwen3TTSTalkerForConditionalGeneration"
+        assert pipeline.stages[3].model_arch == "Qwen3TTSCode2Wav"
+
     def test_unknown_model_type_returns_none_via_get(self):
         assert _PIPELINE_REGISTRY.get("not_a_real_pipeline") is None
 
@@ -73,6 +88,7 @@ class TestLazyLoading:
                     "attr_fail": ("vllm_omni.config.stage_config", "NON_EXISTENT_VAR"),
                 },
             )
+            monkeypatch.setattr(_PIPELINE_REGISTRY, "_loaded", {})
 
             # Assert that list(_PIPELINE_REGISTRY.values()) and dict(_PIPELINE_REGISTRY.items())
             # both yield only healthy entries.

tests/engine/test_async_omni_engine_stage_init.py

@@ -591,6 +591,24 @@ def test_build_engine_args_cli_tokenizer_overrides_inferred_base_tokenizer(tmp_p
     assert engine_args["tokenizer"] == "/external/tokenizer"
 
 
+def test_build_engine_args_stage_model_overrides_parent_model():
+    from vllm_omni.engine.stage_init_utils import build_engine_args_dict
+
+    stage_cfg = types.SimpleNamespace(
+        stage_id=0,
+        stage_type="llm",
+        engine_args={"model": "/stage/model"},
+        default_sampling_params={},
+    )
+
+    engine_args = build_engine_args_dict(
+        stage_cfg,
+        "/parent/model",
+    )
+
+    assert engine_args["model"] == "/stage/model"
+
+
 def test_build_engine_args_keeps_stage_owned_tokenizer_subdir(tmp_path):
     from vllm_omni.engine.stage_init_utils import build_engine_args_dict
 

tests/engine/test_orchestrator.py

@@ -24,7 +24,12 @@ from vllm_omni.engine.messages import (
     ShutdownRequestMessage,
     StageSubmissionMessage,
 )
-from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState
+from vllm_omni.engine.orchestrator import (
+    Orchestrator,
+    OrchestratorRequestState,
+    _build_terminal_empty_output,
+    _infer_stage_audio_sample_rate,
+)
 from vllm_omni.engine.stage_pool import StagePool
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
@@ -135,6 +140,20 @@ class FakeStageClient:
         self._diffusion_outputs.put_nowait(output)
 
 
+def test_terminal_empty_audio_output_uses_stage_sample_rate() -> None:
+    final_stage = FakeStageClient(final_output=True, final_output_type="audio")
+    final_stage.sample_rate = 44100
+    final_pool = SimpleNamespace(stage_client=final_stage, _stage_vllm_config=None)
+
+    terminal_output = _build_terminal_empty_output(
+        "req-1",
+        final_output_type="audio",
+        audio_sample_rate=_infer_stage_audio_sample_rate(final_pool),
+    )
+
+    assert terminal_output.outputs[0].multimodal_output["sr"] == 44100
+
+
 class FakeCollectiveRpcStageClient(FakeStageClient):
     def __init__(self, *args, rpc_result: Any = None, **kwargs) -> None:
         super().__init__(*args, **kwargs)

tests/engine/test_orchestrator_stage_input_bridge.py

@@ -0,0 +1,156 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import queue
+from types import SimpleNamespace
+from typing import Any
+
+import janus
+import pytest
+from vllm.outputs import CompletionOutput, RequestOutput
+from vllm.sampling_params import SamplingParams
+
+from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState
+from vllm_omni.engine.stage_pool import StagePool
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class FakeStageClient:
+    def __init__(
+        self,
+        *,
+        next_inputs: list[dict[str, Any]] | None = None,
+        final_output: bool = False,
+    ) -> None:
+        self.stage_id = 0
+        self.replica_id = 0
+        self.stage_type = "llm"
+        self.final_output = final_output
+        self.final_output_type = "text"
+        self.default_sampling_params = SamplingParams(max_tokens=1)
+        self.requires_multimodal_data = False
+        self.engine_input_source = [0]
+        self.is_comprehension = False
+        self.model_stage = None
+        self.custom_process_input_func = None
+        self.next_inputs = list(next_inputs or [])
+        self.add_request_calls: list[tuple[Any, ...]] = []
+        self._engine_core_outputs = queue.Queue()
+
+    async def add_request_async(self, *args, **_kwargs) -> None:
+        self.add_request_calls.append(args)
+
+    async def get_output_async(self):
+        try:
+            return self._engine_core_outputs.get_nowait()
+        except queue.Empty:
+            return SimpleNamespace(outputs=[])
+
+    def process_engine_inputs(self, _source_outputs, prompt=None, streaming_context=None):
+        return list(self.next_inputs)
+
+    async def abort_requests_async(self, _request_ids: list[str]) -> None:
+        return None
+
+    def set_engine_outputs(self, _outputs) -> None:
+        return None
+
+    def check_health(self) -> None:
+        return None
+
+    def shutdown(self) -> None:
+        return None
+
+
+class FakeOutputProcessor:
+    def add_request(self, *args, **kwargs) -> None:
+        return None
+
+
+class FakeInputProcessor:
+    def __init__(self) -> None:
+        self.calls: list[dict[str, Any]] = []
+
+    def process_inputs(self, **kwargs):
+        self.calls.append(kwargs)
+        return SimpleNamespace(
+            request_id=kwargs["request_id"],
+            prompt_token_ids=[101, 102],
+            prompt_embeds=None,
+            external_req_id=None,
+        )
+
+
+def _request_output(request_id: str) -> RequestOutput:
+    completion = CompletionOutput(
+        index=0,
+        text="transcript",
+        token_ids=[11, 12],
+        cumulative_logprob=None,
+        logprobs=None,
+        finish_reason="stop",
+        stop_reason=None,
+    )
+    return RequestOutput(
+        request_id=request_id,
+        prompt="prompt",
+        prompt_token_ids=[1, 2],
+        prompt_logprobs=None,
+        outputs=[completion],
+        finished=True,
+        metrics=None,
+        lora_request=None,
+    )
+
+
+@pytest.mark.asyncio
+async def test_forward_text_prompt_uses_target_stage_input_processor() -> None:
+    stage0 = FakeStageClient(final_output=True)
+    stage1 = FakeStageClient(
+        final_output=True,
+        next_inputs=[{"prompt": "hello", "multi_modal_data": {"video": ["frame"]}}],
+    )
+    stage_pools = [
+        StagePool(
+            0,
+            [stage0],
+            output_processor=FakeOutputProcessor(),
+            stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)),
+        ),
+        StagePool(
+            1,
+            [stage1],
+            output_processor=FakeOutputProcessor(),
+            stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)),
+        ),
+    ]
+    request_q = janus.Queue()
+    output_q = janus.Queue()
+    rpc_q = janus.Queue()
+    orchestrator = Orchestrator(
+        request_async_queue=request_q.async_q,
+        output_async_queue=output_q.async_q,
+        rpc_async_queue=rpc_q.async_q,
+        stage_pools=stage_pools,
+        async_chunk=False,
+    )
+    input_processor = FakeInputProcessor()
+    orchestrator._stage_input_processors[1] = input_processor
+    req_state = OrchestratorRequestState(
+        request_id="req-text",
+        prompt={"prompt": "original"},
+        sampling_params_list=[SamplingParams(max_tokens=1), SamplingParams(max_tokens=1)],
+        final_stage_id=1,
+    )
+
+    await orchestrator._forward_to_next_stage("req-text", 0, _request_output("req-text"), req_state)
+
+    assert input_processor.calls
+    assert input_processor.calls[0]["prompt"] == {"prompt": "hello", "multi_modal_data": {"video": ["frame"]}}
+    assert stage1.add_request_calls
+    submitted_request = stage1.add_request_calls[0][0]
+    assert submitted_request.prompt_token_ids == [101, 102]
+    assert submitted_request.external_req_id == "req-text"

tests/entrypoints/openai_api/test_serving_chat_multistage_generation.py

@@ -1,8 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
-"""Regression tests for multistage diffusion generation input construction."""
+"""Regression tests for multistage generation input construction."""
 
 from __future__ import annotations
 
+import asyncio
 from types import SimpleNamespace
 
 import pytest
@@ -93,6 +94,79 @@ def test_build_multistage_generation_inputs_applies_stage_specific_overrides(ser
     assert engine.default_sampling_params_list[2].lora_request is None
 
 
+def test_prepare_multistage_multimodal_inputs_defers_downstream_modalities(serving_chat):
+    from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
+
+    image_payload = object()
+    serving_chat.engine_client = SimpleNamespace(
+        stage_configs=[
+            SimpleNamespace(model_stage="asr", requires_multimodal_data=True),
+            SimpleNamespace(model_stage="aura", requires_multimodal_data=True),
+        ]
+    )
+    serving_chat.model_config = SimpleNamespace(
+        allowed_local_media_path="",
+        allowed_media_domains=None,
+    )
+    request = SimpleNamespace(media_io_kwargs=None)
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "audio_url", "audio_url": {"url": "data:audio/wav;base64,abc"}},
+                {"type": "image_pil", "image_pil": image_payload},
+                {"type": "text", "text": "What changed?"},
+            ],
+        }
+    ]
+
+    stripped_messages, deferred_data = asyncio.run(
+        OmniOpenAIServingChat._prepare_multistage_multimodal_inputs(
+            serving_chat,
+            messages,
+            request,
+        )
+    )
+
+    assert deferred_data == {"image": [image_payload]}
+    stripped_content = stripped_messages[0]["content"]
+    assert {"type": "image_pil", "image_pil": image_payload} not in stripped_content
+    assert {"type": "text", "text": "What changed?"} in stripped_content
+
+
+def test_needs_multistage_multimodal_split_noops_for_single_stage(serving_chat):
+    serving_chat.engine_client = SimpleNamespace(
+        stage_configs=[
+            SimpleNamespace(model_stage="thinker", requires_multimodal_data=True),
+        ]
+    )
+
+    assert serving_chat._needs_multistage_multimodal_split() is False
+
+
+def test_needs_multistage_multimodal_split_noops_for_qwen3_omni(serving_chat):
+    serving_chat.engine_client = SimpleNamespace(
+        stage_configs=[
+            SimpleNamespace(model_stage="thinker", requires_multimodal_data=True),
+            SimpleNamespace(model_stage="talker", requires_multimodal_data=False),
+            SimpleNamespace(model_stage="code2wav", requires_multimodal_data=False),
+        ]
+    )
+
+    assert serving_chat._needs_multistage_multimodal_split() is False
+
+
+def test_needs_multistage_multimodal_split_detects_aura_handoff(serving_chat):
+    serving_chat.engine_client = SimpleNamespace(
+        stage_configs=[
+            SimpleNamespace(model_stage="asr", requires_multimodal_data=True),
+            SimpleNamespace(model_stage="aura", requires_multimodal_data=True),
+        ]
+    )
+
+    assert serving_chat._needs_multistage_multimodal_split() is True
+
+
 def test_build_multistage_generation_inputs_multi_image_emits_n_img_placeholders(serving_chat):
     """N reference images with bot_task set must emit N <img> placeholders.
 

tests/entrypoints/openai_api/test_serving_chat_sampling_params.py

@@ -6,6 +6,8 @@ Tests that standard OpenAI API parameters (max_tokens, temperature, etc.)
 are correctly applied to the comprehension stage while preserving YAML defaults.
 """
 
+from types import SimpleNamespace
+
 import pytest
 from pytest_mock import MockerFixture
 from vllm.sampling_params import SamplingParams
@@ -487,6 +489,35 @@ def test_build_sampling_params_list_empty_stop_preserves_yaml(serving_chat, mock
     assert comprehension_params.stop_token_ids == []
 
 
+def test_to_sampling_params_list_pads_missing_tail_stage_with_defaults():
+    """AURA callers may pass 3 semantic model params for a 4-stage engine pipeline."""
+    from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
+
+    instance = object.__new__(OmniOpenAIServingChat)
+    default_params = [
+        SamplingParams(max_tokens=10),
+        SamplingParams(max_tokens=20),
+        SamplingParams(max_tokens=30),
+        SamplingParams(max_tokens=40),
+    ]
+    instance.engine_client = SimpleNamespace(
+        stage_configs=[SimpleNamespace(stage_type="llm") for _ in range(4)],
+        default_sampling_params_list=default_params,
+    )
+
+    result = instance._to_sampling_params_list(
+        [
+            {"max_tokens": 1},
+            {"max_tokens": 2},
+            {"max_tokens": 3},
+        ]
+    )
+
+    assert len(result) == 4
+    assert [params.max_tokens for params in result] == [1, 2, 3, 40]
+    assert result[3] is not default_params[3]
+
+
 # =============================================================================
 # Tests for _get_comprehension_stage_index
 # =============================================================================

tests/model_executor/models/test_aura_qwen3_vl.py

@@ -0,0 +1,69 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.model_executor.models.aura_omni.qwen3_vl import (
+    AuraQwen3VLForConditionalGeneration,
+    AuraQwen3VLProcessingInfo,
+    Qwen3VLProcessor,
+)
+from vllm_omni.model_executor.models.registry import OmniModelRegistry
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _FakeCtx:
+    def __init__(self, config):
+        self.config = config
+        self.processor_type = None
+
+    def get_hf_config(self, typ=None):
+        if typ is not None:
+            raise TypeError("wrong config class")
+        return self.config
+
+    def get_hf_processor(self, typ=None, **kwargs):
+        self.processor_type = typ
+        return SimpleNamespace(kwargs=kwargs)
+
+
+def test_aura_qwen3_vl_processing_accepts_remote_qwen3_vl_config():
+    info = object.__new__(AuraQwen3VLProcessingInfo)
+    info.ctx = _FakeCtx(
+        SimpleNamespace(
+            model_type="qwen3_vl",
+            vision_config=SimpleNamespace(spatial_merge_size=2),
+            text_config=SimpleNamespace(hidden_size=4096),
+        )
+    )
+
+    assert info.get_hf_config().model_type == "qwen3_vl"
+
+
+def test_aura_qwen3_vl_processing_rejects_unrelated_remote_config():
+    info = object.__new__(AuraQwen3VLProcessingInfo)
+    info.ctx = _FakeCtx(SimpleNamespace(model_type="not_qwen3_vl"))
+
+    with pytest.raises(TypeError, match="wrong config class"):
+        info.get_hf_config()
+
+
+def test_aura_qwen3_vl_model_arch_registered():
+    model_cls = OmniModelRegistry._try_load_model_cls("AuraQwen3VLForConditionalGeneration")
+
+    assert model_cls is AuraQwen3VLForConditionalGeneration
+
+
+def test_aura_qwen3_vl_processing_forces_upstream_processor_class():
+    info = object.__new__(AuraQwen3VLProcessingInfo)
+    ctx = _FakeCtx(SimpleNamespace(model_type="qwen3_vl"))
+    info.ctx = ctx
+
+    processor = info.get_hf_processor(extra="value")
+
+    assert ctx.processor_type is Qwen3VLProcessor
+    assert processor.kwargs["use_fast"] is True
+    assert processor.kwargs["extra"] == "value"

tests/model_executor/stage_input_processors/test_aura_omni.py

@@ -0,0 +1,197 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import (
+    PRECOMPUTED_TEXT_IDS_KEY,
+)
+from vllm_omni.model_executor.stage_input_processors.aura_omni import (
+    SILENT_TEXT,
+    asr2aura,
+    aura2tts,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _source_output(text: str, request_id: str = "req-1", token_ids: list[int] | None = None):
+    output = SimpleNamespace(text=text, cumulative_token_ids=token_ids or [1, 2, 3], multimodal_output={})
+    return SimpleNamespace(request_id=request_id, outputs=[output])
+
+
+def _source_delta_final_output(cumulative_text: str, request_id: str = "req-1"):
+    output = SimpleNamespace(
+        text="",
+        cumulative_text=cumulative_text,
+        cumulative_token_ids=[1, 2, 3],
+        multimodal_output={},
+    )
+    return SimpleNamespace(request_id=request_id, outputs=[output])
+
+
+def test_asr2aura_carries_video_payload_and_transcript():
+    prompt = {
+        "multi_modal_data": {"video": ["frame-0", "frame-1"]},
+        "additional_information": {"aura_system_prompt": ["system"]},
+    }
+
+    [next_input] = asr2aura([_source_output("What is happening now?")], prompt=[prompt])
+
+    assert next_input["multi_modal_data"] == {"video": ["frame-0", "frame-1"]}
+    assert "<|video_pad|>" in next_input["prompt"]
+    assert "What is happening now?" in next_input["prompt"]
+    assert next_input["prompt"].startswith("<|im_start|>system\nsystem")
+
+
+def test_asr2aura_drops_audio_before_qwen3_vl_stage():
+    prompt = {
+        "multi_modal_data": {
+            "audio": ("wave", 16000),
+            "video": ["frame-0", "frame-1"],
+        },
+    }
+
+    [next_input] = asr2aura([_source_output("Check the video")], prompt=[prompt])
+
+    assert next_input["multi_modal_data"] == {"video": ["frame-0", "frame-1"]}
+    assert "<|video_pad|>" in next_input["prompt"]
+
+
+def test_asr2aura_reads_video_stashed_for_downstream_stage():
+    prompt = {
+        "multi_modal_data": {"audio": ("wave", 16000)},
+        "additional_information": {
+            "deferred_multi_modal_data": {"video": ["frame-0", "frame-1"]},
+        },
+    }
+
+    [next_input] = asr2aura([_source_output("Check the video")], prompt=[prompt])
+
+    assert next_input["multi_modal_data"] == {"video": ["frame-0", "frame-1"]}
+    assert "<|video_pad|>" in next_input["prompt"]
+
+
+def test_asr2aura_supports_video_only_observation():
+    prompt = {"multi_modal_data": {"video": ["frame-0", "frame-1"]}}
+
+    [next_input] = asr2aura([_source_output("")], prompt=[prompt])
+
+    assert "<|video_pad|>" in next_input["prompt"]
+    assert "<|im_start|>assistant" in next_input["prompt"]
+
+
+def test_aura2tts_builds_qwen3_tts_prompt_information():
+    prompt = {
+        "additional_information": {
+            "tts_language": ["Chinese"],
+            "tts_instruct": ["Calm voice."],
+            "tts_ref_audio": ["ref.wav"],
+            "tts_ref_text": ["Reference transcript sample."],
+        }
+    }
+
+    [tts_input] = aura2tts([_source_output("Hello.")], prompt=[prompt])
+
+    assert len(tts_input["prompt_token_ids"]) >= 32
+    assert tts_input["additional_information"]["text"] == ["Hello."]
+    assert PRECOMPUTED_TEXT_IDS_KEY not in tts_input["additional_information"]
+    assert tts_input["additional_information"]["task_type"] == ["Base"]
+    assert tts_input["additional_information"]["language"] == ["Chinese"]
+    assert tts_input["additional_information"]["ref_audio"] == ["ref.wav"]
+    assert tts_input["additional_information"]["ref_text"] == ["Reference transcript sample."]
+    assert tts_input["additional_information"]["x_vector_only_mode"] == [False]
+    assert tts_input["additional_information"]["instruct"] == ["Calm voice."]
+
+
+def test_aura2tts_prefers_streaming_cumulative_text():
+    prompt = {
+        "additional_information": {
+            "tts_ref_audio": ["ref.wav"],
+            "tts_ref_text": ["Reference transcript sample."],
+        }
+    }
+
+    [tts_input] = aura2tts(
+        [_source_delta_final_output("The complete AURA reply.")],
+        prompt=[prompt],
+    )
+
+    assert tts_input["additional_information"]["text"] == ["The complete AURA reply."]
+
+
+def test_aura2tts_supports_base_ref_audio_override():
+    prompt = {
+        "additional_information": {
+            "tts_ref_audio": ["custom.wav"],
+            "tts_ref_text": ["custom transcript"],
+        }
+    }
+
+    [tts_input] = aura2tts([_source_output("Hello.")], prompt=[prompt])
+
+    assert tts_input["additional_information"]["task_type"] == ["Base"]
+    assert tts_input["additional_information"]["ref_audio"] == ["custom.wav"]
+    assert tts_input["additional_information"]["ref_text"] == ["custom transcript"]
+    assert tts_input["additional_information"]["x_vector_only_mode"] == [False]
+
+
+def test_aura2tts_supports_x_vector_only_mode_for_base():
+    prompt = {
+        "additional_information": {
+            "tts_task_type": ["Base"],
+            "tts_x_vector_only_mode": [True],
+            "tts_ref_audio": ["ref.wav"],
+            "tts_ref_text": ["Reference transcript sample."],
+        }
+    }
+
+    [tts_input] = aura2tts([_source_output("Hello.")], prompt=[prompt])
+
+    assert tts_input["additional_information"]["x_vector_only_mode"] == [True]
+
+
+def test_aura2tts_supports_custom_voice_mode():
+    prompt = {
+        "additional_information": {
+            "tts_task_type": ["CustomVoice"],
+            "tts_speaker": ["vivian"],
+        }
+    }
+
+    [tts_input] = aura2tts([_source_output("Hello.")], prompt=[prompt])
+
+    assert tts_input["additional_information"]["task_type"] == ["CustomVoice"]
+    assert tts_input["additional_information"]["speaker"] == ["Vivian"]
+    assert "ref_audio" not in tts_input["additional_information"]
+
+
+def test_aura2tts_passes_token_ids_to_qwen3_tts_when_enabled():
+    prompt = {
+        "additional_information": {
+            "tts_ref_audio": ["ref.wav"],
+            "tts_ref_text": ["Reference transcript sample."],
+            "tts_pass_token_ids": [True],
+        }
+    }
+
+    [tts_input] = aura2tts(
+        [
+            _source_output(
+                "Hello.",
+                token_ids=[151644, 77091, 198, 108386, 1773, 151645, 198],
+            )
+        ],
+        prompt=[prompt],
+    )
+
+    assert tts_input["additional_information"][PRECOMPUTED_TEXT_IDS_KEY] == [
+        [151644, 77091, 198, 108386, 1773, 151645, 198, 151644, 77091, 198]
+    ]
+    assert "text" not in tts_input["additional_information"]
+
+
+def test_aura2tts_drops_silent_response():
+    assert aura2tts([_source_output(SILENT_TEXT)]) == []

tests/test_config_factory.py

@@ -1384,6 +1384,34 @@ class TestCLIOverrideFlow:
             assert s.runtime_overrides["enforce_eager"] is True
 
 
+class TestAuraOmniDeploy:
+    def test_aura_omni_deploy_forces_pipeline_override(self):
+        deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "aura_omni.yaml"
+        deploy = load_deploy_config(deploy_path)
+
+        assert deploy.pipeline == "aura_omni"
+
+    def test_aura_omni_deploy_resolves_four_native_stages(self):
+        stages = StageConfigFactory._create_from_registry(
+            "qwen3_tts",
+            cli_overrides={},
+            deploy_config_path=str(Path(__file__).parent.parent / "vllm_omni" / "deploy" / "aura_omni.yaml"),
+        )
+
+        assert [stage.model_stage for stage in stages] == [
+            "asr",
+            "aura",
+            "qwen3_tts",
+            "code2wav",
+        ]
+        assert [stage.final_output for stage in stages] == [False, True, False, True]
+        assert [stage.final_output_type for stage in stages] == [None, "text", None, "audio"]
+        assert stages[0].yaml_engine_args["model_arch"] == "Qwen3ASRForConditionalGeneration"
+        assert stages[1].yaml_engine_args["model_arch"] == "AuraQwen3VLForConditionalGeneration"
+        assert stages[2].yaml_engine_args["model_arch"] == "Qwen3TTSTalkerForConditionalGeneration"
+        assert stages[3].yaml_engine_args["model_arch"] == "Qwen3TTSCode2Wav"
+
+
 class TestSentinelDefaultPrecedence:
     """Caller-typed (non-None) values win over YAML; None values fall through
     to YAML / dataclass defaults (#3035)."""

tests/worker/test_omni_gpu_model_runner.py

@@ -4,7 +4,7 @@ from types import SimpleNamespace
 import pytest
 import torch
 
-from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
+from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner, _filter_mrope_kwargs_for_model
 from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -96,12 +96,44 @@ class CaptureTalkerMTP(torch.nn.Module):
         return req_embeds, codes
 
 
+class StrictMRoPEModel:
+    def get_mrope_input_positions(self, input_tokens, mm_features):
+        raise NotImplementedError
+
+
+class FlexibleMRoPEModel:
+    def get_mrope_input_positions(self, input_tokens, mm_features=None, **kwargs):
+        raise NotImplementedError
+
+
 @contextmanager
 def _noop_forward_context(*args, **kwargs):
     """A no-op context manager to replace vLLM forward context in CPU tests."""
     yield
 
 
+def test_filter_mrope_kwargs_for_strict_model_signature():
+    kwargs = {
+        "mm_features": ["audio"],
+        "hf_config": object(),
+        "image_grid_thw": [],
+    }
+
+    assert _filter_mrope_kwargs_for_model(StrictMRoPEModel(), kwargs) == {
+        "mm_features": ["audio"],
+    }
+
+
+def test_filter_mrope_kwargs_preserves_flexible_model_kwargs():
+    kwargs = {
+        "mm_features": ["video"],
+        "hf_config": object(),
+        "video_grid_thw": [[1, 2, 3]],
+    }
+
+    assert _filter_mrope_kwargs_for_model(FlexibleMRoPEModel(), kwargs) is kwargs
+
+
 def _make_runner(req_ids=("r1", "r2"), hidden_size=4):
     # Create an instance without calling OmniGPUModelRunner.__init__
     runner = object.__new__(OmniGPUModelRunner)