[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/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
 # =============================================================================