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

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

docs/user_guide/examples/online_serving/aura_omni.md

@@ -0,0 +1,65 @@
+# AURA Omni Pipeline
+
+`aura_omni` wires ASR, AURA, and Qwen3-TTS into one vLLM-Omni pipeline:
+
+```text
+ASR -> AURA -> Qwen3-TTS Talker -> Code2Wav
+```
+
+Qwen3-TTS remains two engine stages so the pipeline reuses the existing native
+Talker and Code2Wav implementation.
+
+```bash
+vllm serve aurateam/AURA \
+  --omni \
+  --deploy-config vllm_omni/deploy/aura_omni.yaml \
+  --served-model-name aurateam/AURA \
+  --trust-remote-code
+```
+
+Configure local checkpoints by editing per-stage `model` values in
+`vllm_omni/deploy/aura_omni.yaml`. The deploy file sets
+`pipeline: aura_omni`, so the four-stage topology is used even if the
+command-line `--model` points at one of the component checkpoints.
+
+Send requests with `"model": "aurateam/AURA"`. The ASR, AURA, and Qwen3-TTS
+checkpoint paths are internal stage models from the deploy YAML, not the
+OpenAI-facing served model name.
+
+The AURA stage can emit `<|silent|>`. Silent outputs are treated as a gate:
+they produce no Qwen3-TTS Talker input, so no audio is synthesized for that
+turn.
+
+## GPU Utilization Recommendation
+
+`gpu_memory_utilization` in `vllm_omni/deploy/aura_omni.yaml` controls how much
+VRAM each stage can reserve. Start with this split for a single GPU:
+
+- Stage 0 (ASR): `0.10`
+- Stage 1 (AURA): `0.40`
+- Stage 2 (Qwen3-TTS Talker): `0.20`
+- Stage 3 (Qwen3-TTS Code2Wav): `0.20`
+
+## TTS Modes
+
+`aura_omni` can pass AURA text to Qwen3-TTS in two task modes:
+
+- `Base`: voice clone from `tts_ref_audio` with ICL enabled in the AURA
+  pipeline. Provide both `tts_ref_audio` and `tts_ref_text`. Set
+  `tts_x_vector_only_mode=true` to disable ICL and use speaker embedding only.
+- `CustomVoice`: predefined speaker mode. Use a Qwen3-TTS CustomVoice
+  checkpoint for stages 2 and 3 in `aura_omni.yaml`, then pass
+  `tts_task_type=CustomVoice` and `tts_speaker`.
+
+By default, AURA responses are passed to Qwen3-TTS as text. Set
+`additional_information.tts_pass_token_ids=true` to pass AURA-generated
+assistant token ids directly instead. Even when token passthrough is disabled,
+the stage processor uses AURA token ids when available to estimate the Talker
+prefill length, so scheduling does not rely on raw character length.
+
+The example client exposes this as:
+
+```bash
+python examples/online_serving/aura_omni/openai_chat_completion_client.py \
+  --tts-pass-token-ids
+```

examples/offline_inference/aura_omni/README.md

@@ -0,0 +1,61 @@
+# AURA Omni
+
+This example runs the native AURA Omni pipeline offline:
+
+```text
+Qwen3-ASR -> AURA/Qwen3-VL -> Qwen3-TTS Talker -> Qwen3-TTS Code2Wav
+```
+
+The first stage consumes speech and produces a transcript. The AURA stage then combines the transcript with the video frames from the original request and returns text or `<|silent|>`. Non-silent responses are passed to Qwen3-TTS as AURA-generated token ids.
+
+## Run
+
+```bash
+cd examples/offline_inference/aura_omni
+bash run_single_prompt.sh
+```
+
+Use local media:
+
+```bash
+python end2end.py \
+  --audio-path /path/to/input.wav \
+  --video-path /path/to/video.mp4 \
+  --modalities text,audio
+```
+
+Base voice clone mode is the default. It uses the AURA reference audio unless
+you override it:
+
+```bash
+python end2end.py \
+  --tts-task-type Base \
+  --tts-ref-audio /data/yrr/rein_test/shuhan.mp3 \
+  --tts-x-vector-only-mode
+```
+
+CustomVoice mode requires stages 2 and 3 in `vllm_omni/deploy/aura_omni.yaml`
+to point at a Qwen3-TTS CustomVoice checkpoint:
+
+```bash
+python end2end.py \
+  --tts-task-type CustomVoice \
+  --tts-speaker Vivian
+```
+
+For local checkpoints, edit the stage `model` entries in
+`vllm_omni/deploy/aura_omni.yaml` or pass a copied deploy config with
+`--deploy-config`.
+
+## GPU Utilization Recommendation
+
+Set `gpu_memory_utilization` in `vllm_omni/deploy/aura_omni.yaml` by stage.
+Suggested starting point:
+
+- Stage 0 (ASR): `0.10`
+- Stage 1 (AURA): `0.4`
+- Stage 2 (Qwen3-TTS Talker): `0.20`
+- Stage 3 (Qwen3-TTS Code2Wav): `0.20`
+
+Generated text and audio are written to `--output-dir`
+(default: `output_aura_omni`).

examples/offline_inference/aura_omni/end2end.py

@@ -0,0 +1,236 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Offline AURA Omni example.
+
+The native AURA pipeline is:
+
+    Qwen3-ASR -> AURA/Qwen3-VL -> Qwen3-TTS Talker -> Code2Wav
+
+Stage 0 consumes audio and emits a transcript. The stage input processor then
+combines that transcript with the original visual payload for AURA.
+"""
+
+from __future__ import annotations
+
+import os
+from typing import NamedTuple
+
+import numpy as np
+import soundfile as sf
+import vllm
+from vllm import SamplingParams
+from vllm.assets.audio import AudioAsset
+from vllm.assets.video import VideoAsset, video_to_ndarrays
+from vllm.multimodal.media.audio import load_audio
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.model_executor.stage_input_processors.aura_omni import (
+    DEFAULT_QWEN3_TTS_REF_TEXT,
+    default_qwen3_tts_ref_audio_path,
+)
+
+SEED = 42
+DEFAULT_MODEL = "aurateam/AURA"
+DEFAULT_DEPLOY_CONFIG = "vllm_omni/deploy/aura_omni.yaml"
+
+
+class QueryResult(NamedTuple):
+    inputs: dict
+    limit_mm_per_prompt: dict[str, int]
+
+
+def _load_audio(audio_path: str | None, sampling_rate: int) -> tuple[np.ndarray, int]:
+    if audio_path:
+        if not os.path.exists(audio_path):
+            raise FileNotFoundError(f"Audio file not found: {audio_path}")
+        audio_signal, sr = load_audio(audio_path, sr=sampling_rate)
+        return audio_signal.astype(np.float32), int(sr)
+    audio_signal, sr = AudioAsset("mary_had_lamb").audio_and_sample_rate
+    return audio_signal.astype(np.float32), int(sr)
+
+
+def _load_video(video_path: str | None, num_frames: int) -> list[np.ndarray]:
+    if video_path:
+        if not os.path.exists(video_path):
+            raise FileNotFoundError(f"Video file not found: {video_path}")
+        return video_to_ndarrays(video_path, num_frames=num_frames)
+    return VideoAsset(name="baby_reading", num_frames=num_frames).np_ndarrays
+
+
+def build_query(args) -> QueryResult:
+    prompt = (
+        "<|im_start|>user\n"
+        "<|audio_start|><|audio_pad|><|audio_end|>"
+        "Please transcribe this speech and use the video context to decide whether a reply is needed.<|im_end|>\n"
+        "<|im_start|>assistant\n"
+    )
+    audio_data = _load_audio(args.audio_path, args.sampling_rate)
+    video_frames = _load_video(args.video_path, args.num_frames)
+    return QueryResult(
+        inputs={
+            "prompt": prompt,
+            "multi_modal_data": {
+                "audio": audio_data,
+                "video": video_frames,
+            },
+            "additional_information": {
+                "aura_system_prompt": args.aura_system_prompt,
+                "tts_task_type": args.tts_task_type,
+                "tts_language": args.tts_language,
+                "tts_speaker": args.tts_speaker,
+                "tts_instruct": args.tts_instruct,
+                "tts_ref_audio": args.tts_ref_audio,
+                "tts_ref_text": args.tts_ref_text,
+                "tts_x_vector_only_mode": args.tts_x_vector_only_mode,
+                "tts_pass_token_ids": args.tts_pass_token_ids,
+            },
+        },
+        limit_mm_per_prompt={"audio": 1, "video": 1},
+    )
+
+
+def build_sampling_params() -> list[SamplingParams]:
+    asr_sampling = SamplingParams(
+        temperature=0.0,
+        top_p=1.0,
+        top_k=-1,
+        max_tokens=256,
+        detokenize=True,
+        seed=SEED,
+    )
+    aura_sampling = SamplingParams(
+        temperature=0.5,
+        top_p=1.0,
+        top_k=-1,
+        max_tokens=256,
+        detokenize=True,
+        repetition_penalty=1.0,
+        seed=SEED,
+    )
+    talker_sampling = SamplingParams(
+        temperature=0.9,
+        top_k=50,
+        max_tokens=4096,
+        detokenize=False,
+        repetition_penalty=1.05,
+        stop_token_ids=[2150],
+        seed=SEED,
+    )
+    code2wav_sampling = SamplingParams(
+        temperature=0.0,
+        top_p=1.0,
+        top_k=-1,
+        max_tokens=65536,
+        detokenize=True,
+        repetition_penalty=1.0,
+        seed=SEED,
+    )
+    return [asr_sampling, aura_sampling, talker_sampling, code2wav_sampling]
+
+
+def save_stage_output(stage_outputs, output_dir: str) -> None:
+    output = stage_outputs.request_output
+    request_id = output.request_id
+    if stage_outputs.final_output_type == "text":
+        text_output = output.outputs[0].text
+        out_txt = os.path.join(output_dir, f"{request_id}.txt")
+        with open(out_txt, "w", encoding="utf-8") as f:
+            f.write(str(text_output).strip() + "\n")
+        print(f"Request ID: {request_id}, text saved to {out_txt}")
+        print(text_output)
+        return
+
+    if stage_outputs.final_output_type == "audio":
+        mm = output.outputs[0].multimodal_output
+        if not mm or "audio" not in mm:
+            print(f"Request ID: {request_id}, no audio output")
+            return
+        audio = mm["audio"]
+        if isinstance(audio, list):
+            import torch
+
+            audio = torch.cat([(item if hasattr(item, "detach") else torch.tensor(item)).flatten() for item in audio])
+        sample_rate = int(mm.get("sr", 24000))
+        audio_numpy = audio.float().detach().cpu().numpy().flatten()
+        out_wav = os.path.join(output_dir, f"{request_id}.wav")
+        sf.write(out_wav, audio_numpy, samplerate=sample_rate, format="WAV")
+        print(f"Request ID: {request_id}, audio saved to {out_wav}")
+
+
+def main(args) -> None:
+    print("=" * 20, "\n", f"vllm version: {vllm.__version__}", "\n", "=" * 20)
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    query = build_query(args)
+    prompts = [query.inputs for _ in range(args.num_prompts)]
+    if args.modalities:
+        modalities = [item.strip() for item in args.modalities.split(",") if item.strip()]
+        for prompt in prompts:
+            prompt["modalities"] = modalities
+
+    omni_kwargs = vars(args).copy()
+    omni_kwargs["model"] = args.model
+    omni = Omni(**omni_kwargs)
+    sampling_params_list = build_sampling_params()[: omni.num_stages]
+
+    try:
+        for stage_outputs in omni.generate(prompts, sampling_params_list):
+            save_stage_output(stage_outputs, args.output_dir)
+    finally:
+        omni.close()
+
+
+def parse_args():
+    parser = FlexibleArgumentParser(description="AURA Omni offline inference example")
+    parser.add_argument("--model", default=DEFAULT_MODEL, help="Top-level model name/path.")
+    parser.add_argument("--deploy-config", default=DEFAULT_DEPLOY_CONFIG, help="AURA Omni deploy config.")
+    parser.add_argument("--audio-path", "-a", default=None, help="Local audio file. Uses a built-in asset by default.")
+    parser.add_argument("--video-path", "-v", default=None, help="Local video file. Uses a built-in asset by default.")
+    parser.add_argument("--num-frames", type=int, default=16, help="Number of video frames to sample.")
+    parser.add_argument("--sampling-rate", type=int, default=16000, help="Audio sampling rate for ASR input.")
+    parser.add_argument("--num-prompts", type=int, default=1, help="Number of repeated prompts.")
+    parser.add_argument("--modalities", default="text,audio", help="Comma-separated final modalities.")
+    parser.add_argument("--output-dir", default="output_aura_omni", help="Directory for generated files.")
+    parser.add_argument(
+        "--aura-system-prompt",
+        default=(
+            "You are receiving a live video stream where the final frame is the present moment. "
+            "Respond only when a response is needed. Otherwise output '<|silent|>'. Respond in Chinese."
+        ),
+        help="System prompt for the AURA stage.",
+    )
+    parser.add_argument("--tts-task-type", default="Base", choices=["Base", "CustomVoice"], help="Qwen3-TTS task.")
+    parser.add_argument("--tts-language", default="Chinese", help="Qwen3-TTS language.")
+    parser.add_argument("--tts-speaker", default="Vivian", help="Qwen3-TTS speaker.")
+    parser.add_argument("--tts-instruct", default="", help="Optional Qwen3-TTS style instruction.")
+    parser.add_argument(
+        "--tts-ref-audio",
+        default=default_qwen3_tts_ref_audio_path(),
+        help="Base-mode reference audio path/URL.",
+    )
+    parser.add_argument(
+        "--tts-ref-text",
+        default=DEFAULT_QWEN3_TTS_REF_TEXT,
+        help="Base-mode reference audio transcript.",
+    )
+    parser.add_argument(
+        "--tts-x-vector-only-mode",
+        action="store_true",
+        help="Use speaker embedding only for Base mode (disable ICL ref_text conditioning).",
+    )
+    parser.add_argument(
+        "--tts-pass-token-ids",
+        action="store_true",
+        help="Pass AURA-generated assistant token ids directly to Qwen3-TTS. Defaults to sending text.",
+    )
+    parser.add_argument("--init-timeout", type=int, default=2000, help="Overall initialization timeout.")
+    parser.add_argument("--stage-init-timeout", type=int, default=2000, help="Per-stage initialization timeout.")
+    parser.add_argument("--log-stats", action="store_true", default=False, help="Enable stats logging.")
+    nullify_stage_engine_defaults(parser)
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    main(parse_args())

examples/offline_inference/aura_omni/run_single_prompt.sh

@@ -0,0 +1,8 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+python end2end.py \
+  --model aurateam/AURA \
+  --deploy-config vllm_omni/deploy/aura_omni.yaml \
+  --modalities text,audio \
+  "$@"

examples/online_serving/aura_omni/README.md

@@ -0,0 +1,146 @@
+# AURA Omni Native Pipeline
+
+`aura_omni` serves AURA as a native multi-stage vLLM-Omni pipeline:
+
+```text
+Qwen3-ASR -> AURA/Qwen3-VL -> Qwen3-TTS Talker -> Qwen3-TTS Code2Wav
+```
+
+The pipeline has three semantic modules, but four engine stages because the
+existing Qwen3-TTS implementation is natively split into Talker and Code2Wav.
+
+Start the server with the deploy profile:
+
+```bash
+vllm serve aurateam/AURA \
+  --omni \
+  --port 8091 \
+  --deploy-config vllm_omni/deploy/aura_omni.yaml \
+  --served-model-name aurateam/AURA \
+  --trust-remote-code
+```
+
+The deploy file sets per-stage model repos:
+
+- Stage 0 ASR: `Qwen/Qwen3-ASR-1.7B`
+- Stage 1 AURA: `aurateam/AURA`
+- Stage 2/3 TTS: `Qwen/Qwen3-TTS-12Hz-1.7B-Base`
+
+For local weights, edit the `model` value on each stage in
+`vllm_omni/deploy/aura_omni.yaml`. The deploy profile includes
+`pipeline: aura_omni`, so the server uses this four-stage topology even when
+the command-line model path points at one component checkpoint.
+
+Expected request shape:
+
+- Send microphone audio as the Stage 0 multimodal audio input.
+- Include video frames in the original request `multi_modal_data`; the
+  `asr2aura` processor carries them forward to AURA.
+- Optional `additional_information` keys:
+  - `aura_system_prompt`
+  - `tts_task_type`
+  - `tts_language`
+  - `tts_speaker`
+  - `tts_instruct`
+  - `tts_ref_audio`
+  - `tts_ref_text`
+  - `tts_x_vector_only_mode`
+  - `tts_pass_token_ids`
+
+If AURA emits `<|silent|>`, the `aura2tts` processor returns no TTS request, so
+the TTS stages are skipped for that turn.
+
+## GPU Utilization Recommendation
+
+Tune `gpu_memory_utilization` per stage in `vllm_omni/deploy/aura_omni.yaml`.
+Recommended baseline on one GPU for H200
+
+- Stage 0 (ASR): `0.10`
+- Stage 1 (AURA): `0.4`
+- Stage 2 (Qwen3-TTS Talker): `0.20`
+- Stage 3 (Qwen3-TTS Code2Wav): `0.20`
+
+## Python Client
+
+```bash
+python examples/online_serving/aura_omni/openai_chat_completion_client.py \
+  --host localhost \
+  --port 8091 \
+  --model aurateam/AURA \
+  --modalities text,audio
+```
+
+Use local media:
+
+```bash
+python examples/online_serving/aura_omni/openai_chat_completion_client.py \
+  --audio-path /path/to/input.wav \
+  --video-path /path/to/video.mp4 \
+  --output-dir output_aura_omni_online
+```
+
+Base voice clone mode (default, recommended as x-vector while debugging ICL):
+
+```bash
+python examples/online_serving/aura_omni/openai_chat_completion_client.py \
+  --tts-task-type Base \
+  --tts-ref-audio vllm-omni/tests/assets/qwen3_tts/clone_2.wav \
+  --tts-ref-text "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
+```
+
+Enable AURA token-id passthrough explicitly:
+
+```bash
+python examples/online_serving/aura_omni/openai_chat_completion_client.py \
+  --tts-pass-token-ids
+```
+
+CustomVoice mode requires stages 2 and 3 in `aura_omni.yaml` to point at a
+Qwen3-TTS CustomVoice checkpoint:
+
+```bash
+python examples/online_serving/aura_omni/openai_chat_completion_client.py \
+  --tts-task-type CustomVoice \
+  --tts-speaker Vivian
+```
+
+By default, AURA responses are passed to Qwen3-TTS as text. Set
+`tts_pass_token_ids=true` to pass AURA-generated assistant token ids directly
+to Qwen3-TTS instead. The processor still uses AURA token ids, when available,
+to estimate the Talker prompt length in the default text path.
+
+## Curl
+
+```bash
+cd examples/online_serving/aura_omni
+bash run_curl_multimodal_generation.sh
+```
+
+Set `PORT`, `MODEL`, or `OUTPUT_DIR` to override defaults:
+
+```bash
+PORT=8666 MODEL=aurateam/AURA bash run_curl_multimodal_generation.sh
+TTS_PASS_TOKEN_IDS=true PORT=8666 MODEL=aurateam/AURA bash run_curl_multimodal_generation.sh
+```
+
+## Gradio
+
+Launch the server and Gradio UI together:
+
+```bash
+cd examples/online_serving/aura_omni
+bash run_gradio_demo.sh
+```
+
+If the server is already running:
+
+```bash
+python examples/online_serving/aura_omni/gradio_demo.py \
+  --model aurateam/AURA \
+  --api-base http://localhost:8091/v1
+```
+
+## Offline
+
+For offline inference, see
+[`examples/offline_inference/aura_omni`](../../offline_inference/aura_omni/).

examples/online_serving/aura_omni/gradio_demo.py

@@ -0,0 +1,305 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Gradio demo for AURA Omni online serving."""
+
+from __future__ import annotations
+
+import argparse
+import base64
+import io
+from pathlib import Path
+from typing import Any
+
+try:
+    import gradio as gr
+except ImportError:
+    raise ImportError("gradio is required to run this demo. Install it with: pip install 'vllm-omni[demo]'") from None
+
+import numpy as np
+import soundfile as sf
+from openai import OpenAI
+
+from vllm_omni.model_executor.stage_input_processors.aura_omni import (
+    DEFAULT_QWEN3_TTS_REF_TEXT,
+    default_qwen3_tts_ref_audio_path,
+)
+
+SEED = 42
+DEFAULT_TTS_REF_AUDIO = default_qwen3_tts_ref_audio_path()
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Gradio demo for AURA Omni online inference.")
+    parser.add_argument("--model", default="aurateam/AURA", help="Model name served by vLLM.")
+    parser.add_argument("--api-base", default="http://localhost:8091/v1", help="OpenAI-compatible API base.")
+    parser.add_argument("--ip", default="127.0.0.1", help="Gradio host.")
+    parser.add_argument("--port", type=int, default=7862, help="Gradio port.")
+    parser.add_argument("--share", action="store_true", help="Share the Gradio demo publicly.")
+    return parser.parse_args()
+
+
+def _audio_to_data_url(audio_file: Any | None) -> str | None:
+    if audio_file is None:
+        return None
+
+    def _path_to_data_url(path_str: str) -> str | None:
+        if not path_str:
+            return None
+        path = Path(path_str)
+        if not path.exists():
+            return None
+        with open(path, "rb") as f:
+            return f"data:audio/wav;base64,{base64.b64encode(f.read()).decode('utf-8')}"
+
+    if isinstance(audio_file, str):
+        return _path_to_data_url(audio_file)
+
+    # Gradio >=5 may return dict payloads with local temp path.
+    if isinstance(audio_file, dict):
+        path_val = audio_file.get("path") or audio_file.get("name")
+        if isinstance(path_val, str):
+            return _path_to_data_url(path_val)
+        return None
+
+    if isinstance(audio_file, tuple):
+        # Common shape: (sample_rate, np.ndarray)
+        if len(audio_file) == 2 and isinstance(audio_file[0], (int, float)):
+            sample_rate, audio_np = audio_file
+            audio_np = np.asarray(audio_np)
+            if audio_np.ndim > 1:
+                audio_np = audio_np[:, 0]
+            if audio_np.dtype != np.int16:
+                audio_np = np.clip(audio_np.astype(np.float32), -1.0, 1.0)
+            buf = io.BytesIO()
+            sf.write(buf, audio_np, int(sample_rate), format="WAV")
+            return f"data:audio/wav;base64,{base64.b64encode(buf.getvalue()).decode('utf-8')}"
+        # Compatibility shape from some gradio builds: (filepath, (sr, np.ndarray))
+        if len(audio_file) >= 1 and isinstance(audio_file[0], str):
+            path_url = _path_to_data_url(audio_file[0])
+            if path_url is not None:
+                return path_url
+    return None
+
+
+def _video_to_data_url(video_file: str | None) -> str | None:
+    if not video_file:
+        return None
+    path = Path(video_file)
+    mime = {
+        ".mp4": "video/mp4",
+        ".webm": "video/webm",
+        ".mov": "video/quicktime",
+        ".avi": "video/x-msvideo",
+        ".mkv": "video/x-matroska",
+    }.get(path.suffix.lower(), "video/mp4")
+    with open(path, "rb") as f:
+        return f"data:{mime};base64,{base64.b64encode(f.read()).decode('utf-8')}"
+
+
+def _sampling_params_list() -> list[dict[str, Any]]:
+    return [
+        {"temperature": 0.0, "top_p": 1.0, "top_k": -1, "max_tokens": 256, "seed": SEED},
+        {
+            "temperature": 0.5,
+            "top_p": 1.0,
+            "top_k": -1,
+            "max_tokens": 256,
+            "seed": SEED,
+            "repetition_penalty": 1.0,
+        },
+        {
+            "temperature": 0.9,
+            "top_k": 50,
+            "max_tokens": 4096,
+            "seed": SEED,
+            "detokenize": False,
+            "repetition_penalty": 1.05,
+            "stop_token_ids": [2150],
+        },
+        {
+            "temperature": 0.0,
+            "top_p": 1.0,
+            "top_k": -1,
+            "max_tokens": 65536,
+            "seed": SEED,
+            "repetition_penalty": 1.0,
+        },
+    ]
+
+
+def _decode_audio_choice(audio_data: str | None) -> tuple[int, np.ndarray] | None:
+    if not audio_data:
+        return None
+    wav_bytes = base64.b64decode(audio_data)
+    audio_np, sample_rate = sf.read(io.BytesIO(wav_bytes))
+    if audio_np.ndim > 1:
+        audio_np = audio_np[:, 0]
+    return int(sample_rate), audio_np.astype(np.float32)
+
+
+def build_interface(client: OpenAI, model: str):
+    def run(
+        audio_file,
+        video_file,
+        prompt: str,
+        aura_system_prompt: str,
+        tts_task_type: str,
+        tts_language: str,
+        tts_speaker: str,
+        tts_instruct: str,
+        tts_ref_audio: str,
+        tts_ref_text: str,
+        tts_x_vector_only_mode: bool,
+        tts_pass_token_ids: bool,
+    ):
+        audio_url = _audio_to_data_url(audio_file)
+        video_url = _video_to_data_url(video_file)
+        if not audio_url:
+            return "Please provide audio input for the ASR stage.", None
+        if tts_task_type == "Base" and (not (tts_ref_audio or "").strip() or not (tts_ref_text or "").strip()):
+            return "Base TTS requires both reference audio and reference transcript.", None
+        content: list[dict[str, Any]] = [{"type": "audio_url", "audio_url": {"url": audio_url}}]
+        if video_url:
+            content.append({"type": "video_url", "video_url": {"url": video_url}})
+        content.append(
+            {"type": "text", "text": prompt or "Use the audio and video together to decide whether a reply is needed."}
+        )
+
+        additional_information: dict[str, Any] = {
+            "aura_system_prompt": aura_system_prompt,
+            "tts_task_type": tts_task_type,
+            "tts_instruct": tts_instruct,
+            "tts_pass_token_ids": bool(tts_pass_token_ids),
+        }
+        if tts_task_type == "CustomVoice":
+            additional_information.update(
+                {
+                    "tts_language": tts_language or "English",
+                    "tts_speaker": tts_speaker or "Vivian",
+                }
+            )
+        else:
+            additional_information.update(
+                {
+                    "tts_ref_audio": tts_ref_audio.strip() if tts_ref_audio else None,
+                    "tts_ref_text": tts_ref_text.strip() if tts_ref_text else None,
+                    "tts_x_vector_only_mode": bool(tts_x_vector_only_mode),
+                }
+            )
+
+        response = client.chat.completions.create(
+            model=model,
+            messages=[{"role": "user", "content": content}],
+            modalities=["text", "audio"],
+            extra_body={
+                "sampling_params_list": _sampling_params_list(),
+                "additional_information": additional_information,
+            },
+            timeout=600.0,
+        )
+
+        text_parts: list[str] = []
+        audio_output = None
+        for choice in response.choices:
+            message = choice.message
+            if message.content:
+                text_parts.append(str(message.content))
+            if getattr(message, "audio", None):
+                audio_output = _decode_audio_choice(message.audio.data)
+        return "\n\n".join(text_parts) or "No text output.", audio_output
+
+    default_system = (
+        "You are receiving a live video stream where the final frame is the present moment. "
+        "Respond only when a response is needed. Otherwise output '<|silent|>'. Respond in English."
+    )
+
+    def _toggle_tts_advanced(task_type: str):
+        is_base = task_type == "Base"
+        return (
+            gr.update(visible=not is_base),
+            gr.update(visible=is_base),
+        )
+
+    with gr.Blocks() as demo:
+        gr.Markdown("# AURA Omni")
+        gr.Markdown("ASR -> AURA/Qwen3-VL -> Qwen3-TTS Talker -> Code2Wav")
+        with gr.Row():
+            audio_input = gr.Audio(label="Audio input", sources=["upload", "microphone"], type="numpy")
+            video_input = gr.Video(label="Video input (optional)", sources=["upload"])
+        prompt = gr.Textbox(
+            label="Instruction",
+            value="Use the audio and video together to decide whether a reply is needed. If needed, respond briefly in English.",
+            lines=2,
+        )
+        with gr.Accordion("Advanced", open=False):
+            aura_system_prompt = gr.Textbox(label="AURA system prompt", value=default_system, lines=4)
+            tts_task_type = gr.Radio(choices=["Base", "CustomVoice"], value="Base", label="TTS task type")
+            tts_instruct = gr.Textbox(label="TTS instruction", value="")
+            tts_pass_token_ids = gr.Checkbox(
+                label="Pass AURA token ids directly to TTS",
+                value=False,
+            )
+
+            with gr.Group(visible=False) as customvoice_config_group:
+                tts_language = gr.Dropdown(
+                    choices=["Chinese", "English", "Japanese", "Korean", "Cantonese"],
+                    value="Chinese",
+                    label="CustomVoice language",
+                )
+                tts_speaker = gr.Textbox(label="CustomVoice speaker", value="Vivian")
+
+            with gr.Group(visible=True) as base_config_group:
+                tts_ref_audio = gr.Textbox(
+                    label="Base reference audio path/URL",
+                    value=DEFAULT_TTS_REF_AUDIO,
+                )
+                tts_ref_text = gr.Textbox(
+                    label="Base reference transcript",
+                    value=DEFAULT_QWEN3_TTS_REF_TEXT,
+                    lines=3,
+                )
+                tts_x_vector_only_mode = gr.Checkbox(
+                    label="Base x-vector only mode (disable ICL)",
+                    value=False,
+                )
+
+            tts_task_type.change(
+                _toggle_tts_advanced,
+                inputs=[tts_task_type],
+                outputs=[customvoice_config_group, base_config_group],
+            )
+        button = gr.Button("Generate", variant="primary")
+        with gr.Row():
+            text_output = gr.Textbox(label="Text output", lines=8)
+            audio_output = gr.Audio(label="Audio output", interactive=False)
+        button.click(
+            run,
+            inputs=[
+                audio_input,
+                video_input,
+                prompt,
+                aura_system_prompt,
+                tts_task_type,
+                tts_language,
+                tts_speaker,
+                tts_instruct,
+                tts_ref_audio,
+                tts_ref_text,
+                tts_x_vector_only_mode,
+                tts_pass_token_ids,
+            ],
+            outputs=[text_output, audio_output],
+        )
+        demo.queue()
+    return demo
+
+
+def main() -> None:
+    args = parse_args()
+    client = OpenAI(api_key="EMPTY", base_url=args.api_base)
+    demo = build_interface(client, args.model)
+    demo.launch(server_name=args.ip, server_port=args.port, share=args.share)
+
+
+if __name__ == "__main__":
+    main()

examples/online_serving/aura_omni/openai_chat_completion_client.py

@@ -0,0 +1,199 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""OpenAI-compatible client for the AURA Omni pipeline."""
+
+from __future__ import annotations
+
+import base64
+import io
+import os
+
+import soundfile as sf
+from openai import OpenAI
+from vllm.assets.audio import AudioAsset
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+from vllm_omni.model_executor.stage_input_processors.aura_omni import (
+    DEFAULT_QWEN3_TTS_REF_TEXT,
+    default_qwen3_tts_ref_audio_path,
+)
+
+SEED = 42
+DEFAULT_MODEL = "aurateam/AURA"
+DEFAULT_VIDEO_URL = "https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4"
+
+
+def _encode_file(path: str) -> str:
+    with open(path, "rb") as f:
+        return base64.b64encode(f.read()).decode("utf-8")
+
+
+def _data_url(path: str, default_mime: str) -> str:
+    suffix = os.path.splitext(path)[1].lower()
+    mime_by_suffix = {
+        ".wav": "audio/wav",
+        ".mp3": "audio/mpeg",
+        ".ogg": "audio/ogg",
+        ".flac": "audio/flac",
+        ".m4a": "audio/mp4",
+        ".mp4": "video/mp4",
+        ".webm": "video/webm",
+        ".mov": "video/quicktime",
+        ".avi": "video/x-msvideo",
+        ".mkv": "video/x-matroska",
+    }
+    return f"data:{mime_by_suffix.get(suffix, default_mime)};base64,{_encode_file(path)}"
+
+
+def media_url(path_or_url: str | None, *, kind: str) -> str:
+    if path_or_url:
+        if path_or_url.startswith(("http://", "https://", "data:")):
+            return path_or_url
+        if not os.path.exists(path_or_url):
+            raise FileNotFoundError(f"{kind} file not found: {path_or_url}")
+        return _data_url(path_or_url, "audio/wav" if kind == "audio" else "video/mp4")
+    if kind == "audio":
+        return AudioAsset("mary_had_lamb").url
+    return DEFAULT_VIDEO_URL
+
+
+def sampling_params_list() -> list[dict]:
+    return [
+        {"temperature": 0.0, "top_p": 1.0, "top_k": -1, "max_tokens": 256, "seed": SEED},
+        {
+            "temperature": 0.5,
+            "top_p": 1.0,
+            "top_k": -1,
+            "max_tokens": 256,
+            "seed": SEED,
+            "repetition_penalty": 1.0,
+        },
+        {
+            "temperature": 0.9,
+            "top_k": 50,
+            "max_tokens": 4096,
+            "seed": SEED,
+            "detokenize": False,
+            "repetition_penalty": 1.05,
+            "stop_token_ids": [2150],
+        },
+        {
+            "temperature": 0.0,
+            "top_p": 1.0,
+            "top_k": -1,
+            "max_tokens": 65536,
+            "seed": SEED,
+            "repetition_penalty": 1.0,
+        },
+    ]
+
+
+def parse_modalities(value: str | None) -> list[str] | None:
+    if not value:
+        return None
+    return [item.strip() for item in value.split(",") if item.strip()]
+
+
+def save_response(response, output_dir: str) -> None:
+    os.makedirs(output_dir, exist_ok=True)
+    for idx, choice in enumerate(response.choices):
+        message = choice.message
+        if message.content:
+            out_txt = os.path.join(output_dir, f"choice_{idx}.txt")
+            with open(out_txt, "w", encoding="utf-8") as f:
+                f.write(str(message.content).strip() + "\n")
+            print(f"Text saved to {out_txt}")
+            print(message.content)
+        if getattr(message, "audio", None):
+            audio_bytes = base64.b64decode(message.audio.data)
+            audio_np, sample_rate = sf.read(io.BytesIO(audio_bytes))
+            out_wav = os.path.join(output_dir, f"choice_{idx}.wav")
+            sf.write(out_wav, audio_np, int(sample_rate), format="WAV")
+            print(f"Audio saved to {out_wav}")
+
+
+def main(args) -> None:
+    client = OpenAI(base_url=f"http://{args.host}:{args.port}/v1", api_key="EMPTY")
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "audio_url", "audio_url": {"url": media_url(args.audio_path, kind="audio")}},
+                {"type": "video_url", "video_url": {"url": media_url(args.video_path, kind="video")}},
+                {"type": "text", "text": args.prompt},
+            ],
+        }
+    ]
+    response = client.chat.completions.create(
+        model=args.model,
+        messages=messages,
+        modalities=parse_modalities(args.modalities),
+        extra_body={
+            "sampling_params_list": sampling_params_list(),
+            "additional_information": {
+                "aura_system_prompt": args.aura_system_prompt,
+                "tts_task_type": args.tts_task_type,
+                "tts_language": args.tts_language,
+                "tts_speaker": args.tts_speaker,
+                "tts_instruct": args.tts_instruct,
+                "tts_ref_audio": args.tts_ref_audio,
+                "tts_ref_text": args.tts_ref_text,
+                "tts_x_vector_only_mode": args.tts_x_vector_only_mode,
+                "tts_pass_token_ids": args.tts_pass_token_ids,
+            },
+        },
+        timeout=args.timeout,
+    )
+    save_response(response, args.output_dir)
+
+
+def parse_args():
+    parser = FlexibleArgumentParser(description="AURA Omni online serving client")
+    parser.add_argument("--host", default="localhost")
+    parser.add_argument("--port", type=int, default=8091)
+    parser.add_argument("--model", default=DEFAULT_MODEL)
+    parser.add_argument("--audio-path", default=None, help="Audio file, URL, or data URL.")
+    parser.add_argument("--video-path", default=None, help="Video file, URL, or data URL.")
+    parser.add_argument(
+        "--prompt",
+        default="Use the audio and video together to decide whether a reply is needed. If needed, respond briefly in English.",
+    )
+    parser.add_argument("--modalities", default="text,audio")
+    parser.add_argument("--output-dir", default="output_aura_omni_online")
+    parser.add_argument(
+        "--aura-system-prompt",
+        default=(
+            "You are receiving a live video stream where the final frame is the present moment. "
+            "Respond only when a response is needed. Otherwise output '<|silent|>'. Respond in English."
+        ),
+    )
+    parser.add_argument("--tts-task-type", default="Base", choices=["Base", "CustomVoice"])
+    parser.add_argument("--tts-language", default="English")
+    parser.add_argument("--tts-speaker", default="Vivian")
+    parser.add_argument("--tts-instruct", default="")
+    parser.add_argument(
+        "--tts-ref-audio",
+        default=default_qwen3_tts_ref_audio_path(),
+        help="Base-mode reference audio path/URL visible to server.",
+    )
+    parser.add_argument(
+        "--tts-ref-text",
+        default=DEFAULT_QWEN3_TTS_REF_TEXT,
+        help="Base-mode reference audio transcript.",
+    )
+    parser.add_argument(
+        "--tts-x-vector-only-mode",
+        action="store_true",
+        help="Use speaker embedding only for Base mode (disable ICL ref_text conditioning).",
+    )
+    parser.add_argument(
+        "--tts-pass-token-ids",
+        action="store_true",
+        help="Pass AURA-generated assistant token ids directly to Qwen3-TTS. Defaults to sending text.",
+    )
+    parser.add_argument("--timeout", type=float, default=600.0)
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    main(parse_args())

examples/online_serving/aura_omni/run_curl_multimodal_generation.sh

@@ -0,0 +1,60 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+PORT="${PORT:-8091}"
+MODEL="${MODEL:-aurateam/AURA}"
+OUTPUT_DIR="${OUTPUT_DIR:-output_aura_omni_online}"
+TTS_PASS_TOKEN_IDS="${TTS_PASS_TOKEN_IDS:-false}"
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+VLLM_OMNI_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
+CLONE_REF_AUDIO="${VLLM_OMNI_ROOT}/tests/assets/qwen3_tts/clone_2.wav"
+CLONE_REF_TEXT="Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
+mkdir -p "$OUTPUT_DIR"
+
+MARY_HAD_LAMB_AUDIO_URL="https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/mary_had_lamb.ogg"
+SAMPLE_VIDEO_URL="https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4"
+
+request_body=$(cat <<EOF
+{
+  "model": "$MODEL",
+  "modalities": ["text", "audio"],
+  "sampling_params_list": [
+    {"temperature": 0.0, "top_p": 1.0, "top_k": -1, "max_tokens": 256, "seed": 42},
+    {"temperature": 0.5, "top_p": 1.0, "top_k": -1, "max_tokens": 256, "seed": 42, "repetition_penalty": 1.0},
+    {"temperature": 0.9, "top_k": 50, "max_tokens": 4096, "seed": 42, "detokenize": false, "repetition_penalty": 1.05, "stop_token_ids": [2150]},
+    {"temperature": 0.0, "top_p": 1.0, "top_k": -1, "max_tokens": 65536, "seed": 42, "repetition_penalty": 1.0}
+  ],
+  "additional_information": {
+    "aura_system_prompt": "You are receiving a live video stream where the final frame is the present moment. Respond only when a response is needed. Otherwise output '<|silent|>'. Respond in English.",
+    "tts_task_type": "Base",
+    "tts_ref_audio": "file://${CLONE_REF_AUDIO}",
+    "tts_ref_text": "${CLONE_REF_TEXT}",
+    "tts_language": "English",
+    "tts_speaker": "Vivian",
+    "tts_instruct": "",
+    "tts_pass_token_ids": ${TTS_PASS_TOKEN_IDS}
+  },
+  "messages": [{
+    "role": "user",
+    "content": [
+      {"type": "audio_url", "audio_url": {"url": "$MARY_HAD_LAMB_AUDIO_URL"}},
+      {"type": "video_url", "video_url": {"url": "$SAMPLE_VIDEO_URL"}},
+      {"type": "text", "text": "Use the audio and video together to decide whether a reply is needed. If needed, respond briefly in English."}
+    ]
+  }]
+}
+EOF
+)
+
+response=$(curl -sS --retry 3 --retry-delay 3 --retry-connrefused \
+  -X POST "http://localhost:${PORT}/v1/chat/completions" \
+  -H "Content-Type: application/json" \
+  -d "$request_body")
+
+echo "$response" | jq '.choices[].message.content'
+
+audio_b64=$(echo "$response" | jq -r '.choices[]?.message.audio.data // empty' | head -n 1)
+if [[ -n "$audio_b64" ]]; then
+  echo "$audio_b64" | base64 -d > "${OUTPUT_DIR}/aura_omni_output.wav"
+  echo "Audio saved to ${OUTPUT_DIR}/aura_omni_output.wav"
+fi

examples/online_serving/aura_omni/run_gradio_demo.sh

@@ -0,0 +1,79 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+MODEL="aurateam/AURA"
+SERVER_MODEL="aurateam/AURA"
+DEPLOY_CONFIG="/data/yrr/vllm-omni/vllm_omni/deploy/aura_omni.yaml"
+SERVER_PORT=8091
+GRADIO_PORT=7862
+SERVER_HOST="0.0.0.0"
+GRADIO_IP="127.0.0.1"
+GRADIO_SHARE=false
+
+
+while [[ $# -gt 0 ]]; do
+  case "$1" in
+    --model) MODEL="$2"; shift 2 ;;
+    --server-model) SERVER_MODEL="$2"; shift 2 ;;
+    --deploy-config) DEPLOY_CONFIG="$2"; shift 2 ;;
+    --server-port) SERVER_PORT="$2"; shift 2 ;;
+    --gradio-port) GRADIO_PORT="$2"; shift 2 ;;
+    --server-host) SERVER_HOST="$2"; shift 2 ;;
+    --gradio-ip) GRADIO_IP="$2"; shift 2 ;;
+    --share) GRADIO_SHARE=true; shift ;;
+    --help)
+      echo "Usage: $0 [--model SERVED_MODEL_NAME] [--server-model MODEL_PATH] [--deploy-config YAML] [--server-port PORT] [--gradio-port PORT] [--share]"
+      exit 0
+      ;;
+    *) echo "Unknown option: $1"; exit 1 ;;
+  esac
+done
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+API_BASE="http://localhost:${SERVER_PORT}/v1"
+LOG_FILE="/tmp/aura_omni_vllm_${SERVER_PORT}.log"
+
+cleanup() {
+  echo "Shutting down..."
+  [[ -n "${SERVER_PID:-}" ]] && kill "$SERVER_PID" 2>/dev/null || true
+  [[ -n "${GRADIO_PID:-}" ]] && kill "$GRADIO_PID" 2>/dev/null || true
+}
+trap cleanup SIGINT SIGTERM EXIT
+
+vllm serve "$SERVER_MODEL" \
+  --omni \
+  --host "$SERVER_HOST" \
+  --port "$SERVER_PORT" \
+  --deploy-config "$DEPLOY_CONFIG" \
+  --served-model-name "$MODEL" \
+  --trust-remote-code 2>&1 | tee "$LOG_FILE" &
+SERVER_PID=$!
+
+echo "Waiting for server startup..."
+for _ in $(seq 1 600); do
+  if grep -q "Application startup complete" "$LOG_FILE" 2>/dev/null; then
+    break
+  fi
+  if ! kill -0 "$SERVER_PID" 2>/dev/null; then
+    echo "vLLM server exited before startup completed"
+    wait "$SERVER_PID" || true
+    exit 1
+  fi
+  sleep 1
+done
+
+# cd "$SCRIPT_DIR"
+GRADIO_CMD=(python gradio_demo.py --model "$MODEL" --api-base "$API_BASE" --ip "$GRADIO_IP" --port "$GRADIO_PORT")
+if [[ "$GRADIO_SHARE" == "true" ]]; then
+  GRADIO_CMD+=(--share)
+fi
+"${GRADIO_CMD[@]}" &
+GRADIO_PID=$!
+
+echo "vLLM server: http://${SERVER_HOST}:${SERVER_PORT}"
+echo "Gradio demo: http://${GRADIO_IP}:${GRADIO_PORT}"
+if [[ -n "${SERVER_PID:-}" ]]; then
+  wait "$SERVER_PID" "$GRADIO_PID"
+else
+  wait "$GRADIO_PID"
+fi

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)

vllm_omni/config/pipeline_registry.py

@@ -44,6 +44,10 @@ _OMNI_PIPELINES: dict[str, tuple[str, str]] = {
         "vllm_omni.model_executor.models.qwen3_omni.pipeline",
         "QWEN3_OMNI_PIPELINE",
     ),
+    "aura_omni": (
+        "vllm_omni.model_executor.models.aura_omni.pipeline",
+        "AURA_OMNI_PIPELINE",
+    ),
     "qwen3_tts": (
         "vllm_omni.model_executor.models.qwen3_tts.pipeline",
         "QWEN3_TTS_PIPELINE",

vllm_omni/deploy/aura_omni.yaml

@@ -0,0 +1,110 @@
+# AURA Omni deploy profile: Qwen3-ASR -> AURA/Qwen3-VL -> Qwen3-TTS Talker -> Code2Wav.
+#
+# Pass `--deploy-config vllm_omni/deploy/aura_omni.yaml` and point the parent
+# `--model` at any local placeholder repo if needed; each stage overrides its
+# own model below.
+pipeline: aura_omni
+# AURA has a semantic ASR -> AURA -> TTS dependency. Do not prewarm downstream
+# stages before AURA text is converted into Qwen3-TTS metadata by aura2tts.
+async_chunk: false
+
+connectors:
+  connector_of_shared_memory:
+    name: SharedMemoryConnector
+    extra:
+      shm_threshold_bytes: 65536
+      codec_streaming: true
+      connector_get_sleep_s: 0.01
+      connector_get_max_wait_first_chunk: 3000
+      connector_get_max_wait: 300
+      codec_chunk_frames: 25
+      codec_left_context_frames: 72
+      initial_codec_chunk_frames: 1
+      decode_cudagraph_capture_sizes: [25, 73, 97, 169, 325]
+
+stages:
+  - stage_id: 0
+    max_num_seqs: 8
+    gpu_memory_utilization: 0.1
+    trust_remote_code: true
+    enable_prefix_caching: true
+    async_scheduling: true
+    max_model_len: 8192
+    max_num_batched_tokens: 4096
+    devices: "0"
+    model: Qwen/Qwen3-ASR-1.7B
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 256
+
+  - stage_id: 1
+    max_num_seqs: 4
+    gpu_memory_utilization: 0.4
+    trust_remote_code: true
+    enable_prefix_caching: true
+    async_scheduling: true
+    max_model_len: 262144
+    max_num_batched_tokens: 15360
+    devices: "0"
+    model: aurateam/AURA
+    default_sampling_params:
+      temperature: 0.5
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 256
+      repetition_penalty: 1.0
+
+  - stage_id: 2
+    max_num_seqs: 8
+    gpu_memory_utilization: 0.2
+    trust_remote_code: true
+    enable_prefix_caching: false
+    async_scheduling: true
+    max_num_batched_tokens: 512
+    max_model_len: 4096
+    devices: "0"
+    output_connectors:
+      to_stage_3: connector_of_shared_memory
+    # model: Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice
+    model: Qwen/Qwen3-TTS-12Hz-1.7B-Base
+    default_sampling_params:
+      temperature: 0.9
+      top_k: 50
+      max_tokens: 4096
+      seed: 42
+      repetition_penalty: 1.05
+    subtalker_sampling_params:
+      do_sample: true
+      temperature: 0.9
+      top_k: 50
+      top_p: 1.0
+
+  - stage_id: 3
+    max_num_seqs: 8
+    gpu_memory_utilization: 0.2
+    enforce_eager: false
+    trust_remote_code: true
+    enable_prefix_caching: false
+    async_scheduling: true
+    max_num_batched_tokens: 65536
+    max_model_len: 65536
+    devices: "0"
+    input_connectors:
+      from_stage_2: connector_of_shared_memory
+    # model: Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice
+    model: Qwen/Qwen3-TTS-12Hz-1.7B-Base
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 65536
+      seed: 42
+      repetition_penalty: 1.0
+
+platforms:
+  npu:
+    stages:
+      - stage_id: 2
+        enforce_eager: true

vllm_omni/engine/orchestrator.py

@@ -20,6 +20,7 @@ import janus
 import torch
 from vllm.config import ModelConfig
 from vllm.logger import init_logger
+from vllm.outputs import CompletionOutput, RequestOutput
 from vllm.pooling_params import PoolingParams
 from vllm.sampling_params import SamplingParams
 from vllm.v1.engine import EngineCoreOutputs
@@ -52,6 +53,68 @@ from vllm_omni.outputs import OmniRequestOutput
 logger = init_logger(__name__)
 
 
+def _build_terminal_empty_output(
+    request_id: str,
+    *,
+    final_output_type: str | None,
+    audio_sample_rate: int = 24000,
+) -> RequestOutput:
+    """Build a terminal empty output when no downstream stage input exists."""
+    completion = CompletionOutput(
+        index=0,
+        text="",
+        token_ids=[],
+        cumulative_logprob=None,
+        logprobs=None,
+        finish_reason="stop",
+        stop_reason=None,
+    )
+    if final_output_type == "audio":
+        completion.multimodal_output = {
+            "audio": torch.zeros((0,), dtype=torch.float32),
+            "sr": audio_sample_rate,
+        }
+    return RequestOutput(
+        request_id=request_id,
+        prompt=None,
+        prompt_token_ids=[],
+        prompt_logprobs=None,
+        outputs=[completion],
+        finished=True,
+    )
+
+
+def _coerce_int_scalar(value: Any) -> int:
+    if value is None:
+        return 0
+    if isinstance(value, (list, tuple)):
+        for item in value:
+            coerced = _coerce_int_scalar(item)
+            if coerced > 0:
+                return coerced
+        return 0
+    if hasattr(value, "item"):
+        value = value.item()
+    try:
+        coerced = int(value)
+    except (TypeError, ValueError):
+        return 0
+    return coerced if coerced > 0 else 0
+
+
+def _infer_stage_audio_sample_rate(stage_pool: StagePool, default: int = 24000) -> int:
+    """Infer the final audio stage sample rate from stage metadata when possible."""
+    sample_rate_attrs = ("audio_sample_rate", "sample_rate", "sampling_rate", "output_sample_rate", "sr")
+    stage_client = getattr(stage_pool, "stage_client", None)
+    stage_config = getattr(stage_pool, "_stage_vllm_config", None)
+    for source in (stage_client, stage_config):
+        for attr in sample_rate_attrs:
+            sample_rate = _coerce_int_scalar(getattr(source, attr, None))
+            if sample_rate > 0:
+                return sample_rate
+    return default
+
+
 def build_engine_core_request_from_tokens(
     request_id: str,
     prompt: dict[str, Any],
@@ -180,6 +243,7 @@ class Orchestrator:
         self._init_metrics_state(stage_pools, running_counter, transfer_emitter, log_stats=log_stats)
 
         self._cfg_tracker = CfgCompanionTracker()
+        self._stage_input_processors: dict[int, Any] = {}
 
         self._shutdown_event = asyncio.Event()
         self._stages_shutdown = False
@@ -884,6 +948,77 @@ class Orchestrator:
     def _next_stage_already_submitted(self, stage_id: int, req_state: OrchestratorRequestState) -> bool:
         return (stage_id + 1) in req_state.stage_submit_ts
 
+    def _get_stage_input_processor(self, stage_id: int) -> Any:
+        processor = self._stage_input_processors.get(stage_id)
+        if processor is None:
+            from vllm_omni.engine.stage_init_utils import build_stage0_input_processor
+
+            processor = build_stage0_input_processor(self.stage_pools[stage_id].stage_vllm_config)
+            self._stage_input_processors[stage_id] = processor
+        return processor
+
+    def _upgrade_processed_stage_request(self, request: Any, raw_prompt: Any) -> Any:
+        prompt_embeds = getattr(request, "prompt_embeds", None)
+        additional_information = None
+
+        if isinstance(raw_prompt, dict):
+            if prompt_embeds is None:
+                raw_prompt_embeds = raw_prompt.get("prompt_embeds")
+                if isinstance(raw_prompt_embeds, torch.Tensor):
+                    prompt_embeds = raw_prompt_embeds
+            additional_information = serialize_additional_information(
+                raw_prompt.get("additional_information"),
+                log_prefix="Orchestrator stage input",
+            )
+
+        if prompt_embeds is None and additional_information is None:
+            return request
+
+        return OmniEngineCoreRequest.from_request(
+            request,
+            prompt_embeds=prompt_embeds,
+            additional_information=additional_information,
+        )
+
+    def _next_stage_input_is_tokens(self, next_input: Any) -> bool:
+        return isinstance(next_input, dict) and "prompt_token_ids" in next_input
+
+    def _build_next_stage_request(
+        self,
+        req_id: str,
+        next_stage_id: int,
+        next_input: Any,
+        params: SamplingParams | PoolingParams,
+        *,
+        mm_features: list | None = None,
+        resumable: bool = False,
+    ) -> Any:
+        next_pool = self.stage_pools[next_stage_id]
+        if self._next_stage_input_is_tokens(next_input):
+            request = build_engine_core_request_from_tokens(
+                request_id=req_id,
+                prompt=next_input,
+                params=params,
+                model_config=next_pool.stage_vllm_config.model_config,
+                mm_features=mm_features,
+                resumable=resumable,
+            )
+            request.external_req_id = request.request_id
+            return request
+
+        processor = self._get_stage_input_processor(next_stage_id)
+        request = processor.process_inputs(
+            request_id=req_id,
+            prompt=next_input,
+            params=params,
+            supported_tasks=("generate",),
+            arrival_time=_time.time(),
+            resumable=resumable,
+        )
+        request = self._upgrade_processed_stage_request(request, next_input)
+        request.external_req_id = req_id
+        return request
+
     async def _handle_cfg_companion_ready(self, req_id: str) -> None:
         """Mark a CFG companion as done; if all companions are done, flush deferred parent."""
         parent_id = self._cfg_tracker.on_companion_completed(req_id)
@@ -1207,22 +1342,64 @@ class Orchestrator:
             )
             raise
 
+        if not next_inputs:
+            if not getattr(output, "finished", False):
+                logger.debug(
+                    "[Orchestrator] req=%s stage-%s produced no inputs for stage-%s; waiting for more outputs",
+                    req_id,
+                    src_stage_id,
+                    next_logical,
+                )
+                return
+
+            final_stage_id = req_state.final_stage_id
+            final_pool = self.stage_pools[final_stage_id]
+            final_output_type = getattr(final_pool.stage_client, "final_output_type", None)
+            terminal_output = _build_terminal_empty_output(
+                req_id,
+                final_output_type=final_output_type,
+                audio_sample_rate=_infer_stage_audio_sample_rate(final_pool),
+            )
+            submit_ts = _time.time()
+            req_state.stage_submit_ts[final_stage_id] = submit_ts
+            logger.info(
+                "[Orchestrator] req=%s stage-%s produced no terminal inputs for stage-%s; "
+                "returning empty %s output from final stage-%s",
+                req_id,
+                src_stage_id,
+                next_logical,
+                final_output_type or "text",
+                final_stage_id,
+            )
+            await self.output_async_queue.put(
+                OutputMessage(
+                    request_id=req_id,
+                    stage_id=final_stage_id,
+                    replica_id=0,
+                    engine_outputs=terminal_output,
+                    metrics=None,
+                    finished=True,
+                    stage_submit_ts=submit_ts,
+                )
+            )
+            await self._cleanup_request_ids([req_id, *self._cfg_tracker.cleanup_parent(req_id)])
+            return
+
         # Build and submit requests for each input
         for next_input in next_inputs:
             # Only AR thinker stages consume encoder mm_features; downstream
             # (talker/code2wav/…) must not see them (avoids encoder-cache misses).
             model_stage = getattr(getattr(next_pool.stage_vllm_config, "model_config", None), "model_stage", None)
             mm_features = req_state.mm_features if model_stage == "thinker" else None
-            request = build_engine_core_request_from_tokens(
-                request_id=req_id,
-                prompt=next_input,
+            request = self._build_next_stage_request(
+                req_id,
+                next_logical,
+                next_input,
                 params=params,
-                model_config=next_pool.stage_vllm_config.model_config,
                 mm_features=mm_features,
                 resumable=next_stage_resumable,
             )
 
-            request.external_req_id = request.request_id
             if already_submitted:
                 await next_pool.submit_update(req_id, req_state, request)
             else:

vllm_omni/engine/stage_init_utils.py

@@ -698,6 +698,7 @@ def build_engine_args_dict(
     stage_id = stage_config.stage_id
 
     engine_args_dict = _to_dict(engine_args)
+    model = engine_args_dict.pop("model", None) or model
     stage_defines_tokenizer = (
         engine_args_dict.get("tokenizer") is not None or engine_args_dict.get("tokenizer_subdir") is not None
     )

vllm_omni/entrypoints/openai/serving_chat.py

@@ -74,6 +74,7 @@ from vllm.entrypoints.serve.utils.api_utils import should_include_usage
 from vllm.entrypoints.serve.utils.tool_calls_utils import maybe_filter_parallel_tool_calls
 from vllm.inputs import PromptType
 from vllm.logger import init_logger
+from vllm.multimodal.media.connector import MediaConnector
 from vllm.outputs import RequestOutput
 from vllm.reasoning import ReasoningParser
 from vllm.renderers import BaseRenderer, merge_kwargs
@@ -104,6 +105,7 @@ from vllm_omni.entrypoints.openai.protocol.images import (
 )
 from vllm_omni.entrypoints.openai.stage_params import (
     build_stage_sampling_params_list,
+    clone_sampling_params,
     get_default_sampling_params_list,
 )
 from vllm_omni.entrypoints.openai.utils import (
@@ -122,6 +124,10 @@ from vllm_omni.utils.audio import audio_chunk_pcm_bytes, audio_chunk_sample_rate
 logger = init_logger(__name__)
 
 
+async def _identity_async(value: Any) -> Any:
+    return value
+
+
 class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
     """OpenAI-compatible chat serving for both LLM and Diffusion models.
 
@@ -654,6 +660,13 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
             default_mm_processor_kwargs=getattr(request, "mm_processor_kwargs", None),
         )
 
+        deferred_multi_modal_data: dict[str, Any] | None = None
+        if self._needs_multistage_multimodal_split():
+            messages, deferred_multi_modal_data = await self._prepare_multistage_multimodal_inputs(
+                messages,
+                request,
+            )
+
         (conversation,), (engine_prompt,) = await renderer.render_chat_async(
             [messages],
             chat_params,
@@ -698,18 +711,25 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
         if hasattr(request, "cache_salt") and request.cache_salt is not None:
             engine_prompt["cache_salt"] = request.cache_salt
 
+        additional_information = getattr(request, "additional_information", None)
+        if isinstance(additional_information, dict):
+            prompt_additional_information = self._ensure_prompt_additional_information(engine_prompt)
+            prompt_additional_information.update(additional_information)
+
+        if deferred_multi_modal_data:
+            prompt_additional_information = self._ensure_prompt_additional_information(engine_prompt)
+            prompt_additional_information["deferred_multi_modal_data"] = deferred_multi_modal_data
+
         speaker = getattr(request, "voice", None) or getattr(request, "speaker", None)
         normalized = validate_requested_speaker(speaker, self._get_supported_speakers())
         if normalized is not None:
-            if "additional_information" not in engine_prompt or engine_prompt["additional_information"] is None:
-                engine_prompt["additional_information"] = {}
-            engine_prompt["additional_information"]["speaker"] = [normalized]
+            prompt_additional_information = self._ensure_prompt_additional_information(engine_prompt)
+            prompt_additional_information["speaker"] = [normalized]
 
         language = getattr(request, "language", None)
         if language is not None and isinstance(language, str) and language.strip():
-            if "additional_information" not in engine_prompt or engine_prompt["additional_information"] is None:
-                engine_prompt["additional_information"] = {}
-            engine_prompt["additional_information"]["language"] = [language.strip()]
+            prompt_additional_information = self._ensure_prompt_additional_information(engine_prompt)
+            prompt_additional_information["language"] = [language.strip()]
 
         # Style instruction — used by Ming-flash-omni instruct TTS path
         # (ming_task="instruct").  For the omni speech path the thinker2talker
@@ -717,20 +737,263 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
         # which hardcodes instruction=None.
         instructions = getattr(request, "instructions", None)
         if instructions is not None and isinstance(instructions, str) and instructions.strip():
-            if "additional_information" not in engine_prompt or engine_prompt["additional_information"] is None:
-                engine_prompt["additional_information"] = {}
-            engine_prompt["additional_information"]["instruction"] = instructions.strip()
+            prompt_additional_information = self._ensure_prompt_additional_information(engine_prompt)
+            prompt_additional_information["instruction"] = instructions.strip()
 
         return conversation, [engine_prompt]
 
+    @staticmethod
+    def _ensure_prompt_additional_information(engine_prompt: dict[str, Any]) -> dict[str, Any]:
+        additional_information = engine_prompt.get("additional_information")
+        if not isinstance(additional_information, dict):
+            additional_information = {}
+            engine_prompt["additional_information"] = additional_information
+        return additional_information
+
+    def _needs_multistage_multimodal_split(self) -> bool:
+        return bool(self._deferred_multimodal_modalities())
+
+    def _deferred_multimodal_modalities(self) -> set[str]:
+        stage_configs = list(getattr(self.engine_client, "stage_configs", []) or [])
+        if len(stage_configs) < 2:
+            return set()
+
+        first_stage_modalities = self._stage_input_modalities(stage_configs[0])
+        if not first_stage_modalities:
+            return set()
+
+        downstream_modalities: set[str] = set()
+        for stage in stage_configs[1:]:
+            downstream_modalities.update(self._stage_input_modalities(stage))
+        return downstream_modalities - first_stage_modalities
+
+    @staticmethod
+    def _stage_input_modalities(stage: Any) -> set[str]:
+        engine_args = getattr(stage, "engine_args", None)
+        explicit = (
+            getattr(stage, "input_modalities", None)
+            or getattr(stage, "modalities", None)
+            or getattr(engine_args, "input_modalities", None)
+            or getattr(engine_args, "modalities", None)
+        )
+        if explicit:
+            return {str(modality) for modality in as_list(explicit)}
+
+        model_stage = str(getattr(engine_args, "model_stage", None) or getattr(stage, "model_stage", "")).lower()
+        if model_stage in {"asr", "stt"} or model_stage.endswith("_asr"):
+            return {"audio"}
+        if any(name in model_stage for name in ("vision", "vl", "aura")):
+            return {"image", "video"}
+        if any(name in model_stage for name in ("tts", "talker", "code2wav")):
+            return set()
+
+        if getattr(stage, "requires_multimodal_data", False):
+            return {"audio", "image", "video"}
+        return set()
+
+    async def _prepare_multistage_multimodal_inputs(
+        self,
+        messages: list[ChatCompletionMessageParam],
+        request: ChatLikeRequest | ResponsesRequest,
+    ) -> tuple[list[ChatCompletionMessageParam], dict[str, Any] | None]:
+        """Hide modalities unsupported by stage 0 and stash them for downstream stages."""
+        deferred_modalities = self._deferred_multimodal_modalities()
+        deferred_parts: dict[str, list[Any]] = {modality: [] for modality in deferred_modalities}
+        stripped_messages: list[ChatCompletionMessageParam] = []
+
+        for message in messages:
+            content = message.get("content") if isinstance(message, dict) else None
+            if not isinstance(content, list):
+                stripped_messages.append(message)
+                continue
+
+            stripped_content: list[Any] = []
+            changed = False
+            for part in content:
+                modality, payload = self._deferred_multimodal_part(part, deferred_modalities)
+                if modality is not None:
+                    if payload is not None:
+                        deferred_parts.setdefault(modality, []).append(payload)
+                    changed = True
+                    continue
+                stripped_content.append(part)
+
+            if changed:
+                stripped_message = dict(message)
+                stripped_message["content"] = stripped_content
+                stripped_messages.append(cast(ChatCompletionMessageParam, stripped_message))
+            else:
+                stripped_messages.append(message)
+
+        deferred_parts = {modality: parts for modality, parts in deferred_parts.items() if parts}
+        if not deferred_parts:
+            return messages, None
+
+        media_connector = MediaConnector(
+            media_io_kwargs=getattr(request, "media_io_kwargs", None),
+            allowed_local_media_path=getattr(self.model_config, "allowed_local_media_path", "") or "",
+            allowed_media_domains=getattr(self.model_config, "allowed_media_domains", None),
+        )
+        multi_modal_data: dict[str, Any] = {}
+        for modality, parts in deferred_parts.items():
+            multi_modal_data[modality] = await self._materialize_deferred_multimodal_parts(
+                media_connector,
+                modality,
+                parts,
+            )
+        return stripped_messages, multi_modal_data
+
+    @staticmethod
+    def _deferred_multimodal_part(part: Any, deferred_modalities: set[str]) -> tuple[str | None, Any | None]:
+        if not isinstance(part, dict):
+            return None, None
+        part_type = part.get("type")
+        if part_type in {"video_url", "video"} and "video" in deferred_modalities:
+            video = part.get("video_url", part.get("video"))
+            if isinstance(video, dict):
+                video = video.get("url")
+            return "video", video
+        if part_type in {"image_url", "image_pil", "image"} and "image" in deferred_modalities:
+            image = part.get("image_pil", part.get("image_url", part.get("image")))
+            if isinstance(image, dict):
+                image = image.get("url")
+            return "image", image
+        if part_type in {"audio_url", "input_audio", "audio"} and "audio" in deferred_modalities:
+            audio = part.get("audio_url", part.get("input_audio", part.get("audio")))
+            if isinstance(audio, dict):
+                audio = audio.get("url") or audio.get("data")
+            return "audio", audio
+        return None, None
+
+    @staticmethod
+    async def _materialize_deferred_multimodal_parts(
+        media_connector: MediaConnector,
+        modality: str,
+        parts: list[Any],
+    ) -> list[Any]:
+        if modality == "video":
+            return list(await asyncio.gather(*(media_connector.fetch_video_async(part) for part in parts)))
+        if modality == "image":
+            fetch_image = getattr(media_connector, "fetch_image_async", None)
+            if fetch_image is None:
+                return parts
+            return list(
+                await asyncio.gather(
+                    *(fetch_image(part) if isinstance(part, str) else _identity_async(part) for part in parts)
+                )
+            )
+        if modality == "audio":
+            fetch_audio = getattr(media_connector, "fetch_audio_async", None)
+            if fetch_audio is None:
+                return parts
+            return list(
+                await asyncio.gather(
+                    *(fetch_audio(part) if isinstance(part, str) else _identity_async(part) for part in parts)
+                )
+            )
+        return parts
+
+    async def _inject_audio_from_video_urls(
+        self,
+        messages: list[ChatCompletionMessageParam],
+    ) -> list[ChatCompletionMessageParam]:
+        """Pre-extract audio from video URLs and inject as audio_url content items.
+
+        When use_audio_in_video=True, the qwen2_5_omni_thinker multimodal
+        processor requires that the number of audio items equals the number of
+        video items (it subtracts mm_counts["video"] from mm_counts["audio"]).
+        The client only sends video_url items; this method adds the matching
+        audio_url items on the server side before the renderer processes them.
+        """
+        import io
+
+        from vllm_omni.entrypoints.chat_utils import extract_audio_from_video_async
+
+        new_messages: list[ChatCompletionMessageParam] = []
+        for msg in messages:
+            content = msg.get("content") if isinstance(msg, dict) else getattr(msg, "content", None)
+            if not isinstance(content, list):
+                new_messages.append(msg)
+                continue
+
+            video_urls = [
+                part.get("video_url", {}).get("url")
+                for part in content
+                if isinstance(part, dict) and part.get("type") == "video_url" and part.get("video_url", {}).get("url")
+            ]
+
+            if not video_urls:
+                new_messages.append(msg)
+                continue
+
+            audios = await asyncio.gather(*(extract_audio_from_video_async(u) for u in video_urls))
+
+            audio_items: list[dict] = []
+            for audio_array, sample_rate in audios:
+                buf = io.BytesIO()
+                if soundfile is not None:
+                    soundfile.write(buf, audio_array, samplerate=int(sample_rate), format="WAV")
+                else:
+                    import struct
+
+                    import numpy as np
+
+                    audio_np = np.asarray(audio_array, dtype=np.float32)
+                    sr = int(sample_rate)
+                    num_channels = 1
+                    bits_per_sample = 32
+                    num_frames = len(audio_np)
+                    data_size = num_frames * num_channels * (bits_per_sample // 8)
+                    # Write minimal RIFF/WAV header
+                    buf.write(b"RIFF")
+                    buf.write(struct.pack("<I", 36 + data_size))
+                    buf.write(b"WAVE")
+                    buf.write(b"fmt ")
+                    buf.write(
+                        struct.pack(
+                            "<IHHIIHH",
+                            16,
+                            3,
+                            num_channels,
+                            sr,
+                            sr * num_channels * (bits_per_sample // 8),
+                            num_channels * (bits_per_sample // 8),
+                            bits_per_sample,
+                        )
+                    )
+                    buf.write(b"data")
+                    buf.write(struct.pack("<I", data_size))
+                    buf.write(audio_np.tobytes())
+
+                audio_b64 = base64.b64encode(buf.getvalue()).decode()
+                audio_items.append(
+                    {
+                        "type": "audio_url",
+                        "audio_url": {"url": f"data:audio/wav;base64,{audio_b64}"},
+                    }
+                )
+
+            new_content = list(content) + audio_items
+            if isinstance(msg, dict):
+                new_msg = {**msg, "content": new_content}
+            else:
+                new_msg = msg.model_copy(update={"content": new_content})
+            new_messages.append(new_msg)
+
+        return new_messages
+
     def _to_sampling_params_list(self, sampling_params_list: list[dict]) -> list[Any]:
         """Convert request dicts to stage-typed sampling params objects.
 
         For diffusion stages, build ``OmniDiffusionSamplingParams`` so
         downstream ``StageDiffusionClient._sampling_params_to_dict`` (which
         requires a dataclass) works. For LLM stages build ``SamplingParams``.
+        If callers provide params for fewer stages than the native pipeline has
+        (for example AURA has three semantic models but four engine stages),
+        append cloned deploy defaults for the omitted tail stages.
         """
         stage_configs = list(getattr(self.engine_client, "stage_configs", []) or [])
+        default_params_list = list(getattr(self.engine_client, "default_sampling_params_list", []) or [])
         final_sampling_params_list: list[Any] = []
         for idx, sampling_params in enumerate(sampling_params_list):
             stage_type = get_stage_type(stage_configs[idx]) if idx < len(stage_configs) else "llm"
@@ -750,6 +1013,11 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 final_sampling_params_list.append(target_cls(**as_dict))
             else:
                 raise ValueError(f"Invalid sampling params: {sampling_params}")
+        for idx in range(len(final_sampling_params_list), len(stage_configs)):
+            if idx < len(default_params_list):
+                final_sampling_params_list.append(clone_sampling_params(default_params_list[idx]))
+            else:
+                final_sampling_params_list.append(SamplingParams())
         return final_sampling_params_list
 
     def _get_comprehension_stage_index(self) -> int:

vllm_omni/model_executor/models/aura_omni/pipeline.py

@@ -0,0 +1,84 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""AURA Omni pipeline topology.
+
+Semantic modules:
+  1. Qwen3-ASR: microphone audio -> transcript text
+  2. AURA/Qwen3-VL: transcript + video -> text or <|silent|>
+  3. Qwen3-TTS: response text -> audio
+
+The Qwen3-TTS module is represented as two native engine stages, Talker and
+Code2Wav, to reuse the existing vLLM-Omni implementation without wrapping or
+duplicating codec streaming behavior.
+"""
+
+from vllm_omni.config.stage_config import (
+    PipelineConfig,
+    StageExecutionType,
+    StagePipelineConfig,
+)
+
+_AURA_PROC = "vllm_omni.model_executor.stage_input_processors.aura_omni"
+_QWEN3_TTS_PROC = "vllm_omni.model_executor.stage_input_processors.qwen3_tts"
+
+
+AURA_OMNI_PIPELINE = PipelineConfig(
+    model_type="aura_omni",
+    model_arch="Qwen3ASRForConditionalGeneration",
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="asr",
+            execution_type=StageExecutionType.LLM_AR,
+            input_sources=(),
+            owns_tokenizer=True,
+            requires_multimodal_data=True,
+            engine_output_type="text",
+            model_arch="Qwen3ASRForConditionalGeneration",
+            sampling_constraints={"detokenize": True},
+        ),
+        StagePipelineConfig(
+            stage_id=1,
+            model_stage="aura",
+            execution_type=StageExecutionType.LLM_AR,
+            input_sources=(0,),
+            final_output=True,
+            final_output_type="text",
+            owns_tokenizer=True,
+            requires_multimodal_data=True,
+            engine_output_type="text",
+            model_arch="AuraQwen3VLForConditionalGeneration",
+            custom_process_input_func=f"{_AURA_PROC}.asr2aura",
+            sampling_constraints={"detokenize": True},
+        ),
+        StagePipelineConfig(
+            stage_id=2,
+            model_stage="qwen3_tts",
+            execution_type=StageExecutionType.LLM_AR,
+            input_sources=(1,),
+            owns_tokenizer=True,
+            engine_output_type="latent",
+            model_arch="Qwen3TTSTalkerForConditionalGeneration",
+            custom_process_input_func=f"{_AURA_PROC}.aura2tts",
+            custom_process_next_stage_input_func=f"{_QWEN3_TTS_PROC}.talker2code2wav_full_payload",
+            async_chunk_process_next_stage_input_func=f"{_QWEN3_TTS_PROC}.talker2code2wav_async_chunk",
+            sampling_constraints={
+                "detokenize": False,
+                "stop_token_ids": [2150],
+            },
+        ),
+        StagePipelineConfig(
+            stage_id=3,
+            model_stage="code2wav",
+            execution_type=StageExecutionType.LLM_GENERATION,
+            input_sources=(2,),
+            final_output=True,
+            final_output_type="audio",
+            engine_output_type="audio",
+            model_arch="Qwen3TTSCode2Wav",
+            sync_process_input_func=f"{_QWEN3_TTS_PROC}.talker2code2wav_token_only",
+            sampling_constraints={"detokenize": True},
+            extras={"tts_args": {"max_instructions_length": 500}},
+        ),
+    ),
+)

vllm_omni/model_executor/models/aura_omni/qwen3_vl.py

@@ -0,0 +1,59 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""AURA-compatible Qwen3-VL wrapper.
+
+AURA ships Qwen3-VL weights with a repository-local ``configuration_qwen3_vl``
+module in ``auto_map``. With ``trust_remote_code=True``, Transformers therefore
+returns a remote ``Qwen3VLConfig`` class. The model structure is still Qwen3-VL,
+but vLLM's stock multimodal processor checks for the upstream Transformers
+``Qwen3VLConfig`` type exactly. This wrapper keeps the stock model and processor
+logic, while accepting structurally-compatible remote Qwen3-VL configs.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from vllm.model_executor.models.qwen3_vl import (
+    Qwen3VLConfig,
+    Qwen3VLDummyInputsBuilder,
+    Qwen3VLForConditionalGeneration,
+    Qwen3VLMultiModalProcessor,
+    Qwen3VLProcessingInfo,
+    Qwen3VLProcessor,
+)
+from vllm.multimodal import MULTIMODAL_REGISTRY
+
+
+class AuraQwen3VLProcessingInfo(Qwen3VLProcessingInfo):
+    def get_hf_config(self) -> Any:
+        try:
+            return self.ctx.get_hf_config(Qwen3VLConfig)
+        except TypeError:
+            config = self.ctx.get_hf_config()
+            if (
+                getattr(config, "model_type", None) == "qwen3_vl"
+                and hasattr(config, "vision_config")
+                and hasattr(config, "text_config")
+            ):
+                return config
+            raise
+
+    def get_hf_processor(self, **kwargs: object) -> Any:
+        # AURA's config advertises `processing_qwen3_vl.py` in auto_map, but
+        # the checkpoint snapshot does not include that file. Force the
+        # upstream processor class while still allowing AURA's remote config.
+        return self.ctx.get_hf_processor(
+            Qwen3VLProcessor,
+            use_fast=kwargs.pop("use_fast", True),
+            **kwargs,
+        )
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+    Qwen3VLMultiModalProcessor,
+    info=AuraQwen3VLProcessingInfo,
+    dummy_inputs=Qwen3VLDummyInputsBuilder,
+)
+class AuraQwen3VLForConditionalGeneration(Qwen3VLForConditionalGeneration):
+    pass

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py

@@ -30,7 +30,7 @@ from vllm_omni.utils.speaker_cache import (
 )
 
 from .configuration_qwen3_tts import Qwen3TTSConfig, Qwen3TTSSpeakerEncoderConfig, Qwen3TTSTalkerConfig
-from .prompt_embeds_builder import Qwen3TTSPromptEmbedsBuilder
+from .prompt_embeds_builder import PRECOMPUTED_TEXT_IDS_KEY, Qwen3TTSPromptEmbedsBuilder
 from .qwen3_tts_code_predictor_vllm import Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM
 from .tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import Qwen3TTSTokenizerV2Config
 from .tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import Qwen3TTSTokenizerV2Encoder
@@ -40,6 +40,19 @@ logger = init_logger(__name__)
 _TRAILING_TEXT_COMPACT_MIN_FRAMES = 64
 
 
+def _has_tts_text_conditioning(info_dict: dict[str, Any], hidden_states: Any | None = None) -> bool:
+    text_list = info_dict.get("text")
+    if isinstance(text_list, list) and bool(text_list) and bool(text_list[0]):
+        return True
+    if PRECOMPUTED_TEXT_IDS_KEY in info_dict:
+        return True
+    if isinstance(hidden_states, dict):
+        tail = hidden_states.get("trailing_text")
+        if isinstance(tail, torch.Tensor):
+            return True
+    return False
+
+
 # ---------------------------------------------------------------------------
 # Components ported from the HuggingFace Qwen3-TTS reference implementation.
 # Only the classes actually needed by the vLLM AR Talker are kept here.
@@ -625,9 +638,8 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             except Exception:
                 is_prefill = span_len > 1
 
-        text_list = info_dict.get("text")
-        if not isinstance(text_list, list) or not text_list or not text_list[0]:
-            raise ValueError("Missing additional_information.text for Qwen3-TTS AR talker.")
+        if not _has_tts_text_conditioning(info_dict, hs):
+            raise ValueError("Missing Qwen3-TTS text conditioning: provide `text` or precomputed text token ids.")
 
         task_type = (info_dict.get("task_type") or ["CustomVoice"])[0]
         codec_streaming_val = meta.get("codec_streaming")
@@ -650,7 +662,6 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         if is_prefill:
             # Prefill (prompt embeddings)
             prompt_embeds_cpu = embed.get("prefill")
-
             # First prefill round: prompt_embeds_cpu is not yet populated.
             # Subsequent prefill rounds (multi-chunk): prompt_embeds_cpu is a Tensor stored by the first round.
             is_first_prefill = not isinstance(prompt_embeds_cpu, torch.Tensor) or prompt_embeds_cpu.ndim != 2
@@ -808,9 +819,8 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             hs = payload.get("hidden_states", {})
             meta = payload.get("meta", {})
 
-            text_list = info_dict.get("text")
-            if not isinstance(text_list, list) or not text_list or not text_list[0]:
-                raise ValueError("Missing additional_information.text for Qwen3-TTS AR talker.")
+            if not _has_tts_text_conditioning(info_dict, hs):
+                raise ValueError("Missing Qwen3-TTS text conditioning: provide `text` or precomputed text token ids.")
 
             task_type = (info_dict.get("task_type") or ["CustomVoice"])[0]
             codec_streaming_val = meta.get("codec_streaming")

vllm_omni/model_executor/models/registry.py

@@ -343,6 +343,11 @@ _OMNI_MODELS = {
         "minicpmo_4_5_omni_tts",
         "MiniCPMO45OmniTTSForConditionalGeneration",
     ),
+    "AuraQwen3VLForConditionalGeneration": (
+        "aura_omni",
+        "qwen3_vl",
+        "AuraQwen3VLForConditionalGeneration",
+    ),
 }
 
 

vllm_omni/model_executor/stage_input_processors/aura_omni.py

@@ -0,0 +1,390 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Stage processors for the AURA Omni pipeline."""
+
+from __future__ import annotations
+
+import math
+import os
+from pathlib import Path
+from typing import Any
+
+import soundfile as sf
+
+from vllm_omni.inputs.data import OmniTokensPrompt
+from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import (
+    PRECOMPUTED_TEXT_IDS_KEY,
+)
+
+DEFAULT_AURA_SYSTEM_PROMPT = (
+    "You are receiving a live video stream where the final frame is the present moment. "
+    "Respond only when a response is needed based on the user's message or the visual context. "
+    "Otherwise, output '<|silent|>' to signify silence. Respond in Chinese."
+)
+
+SILENT_TEXT = "<|silent|>"
+QWEN_IM_START_ID = 151644
+QWEN_IM_END_ID = 151645
+QWEN_ASSISTANT_ID = 77091
+QWEN_NEWLINE_ID = 198
+QWEN_ASSISTANT_PREFIX_IDS = [QWEN_IM_START_ID, QWEN_ASSISTANT_ID, QWEN_NEWLINE_ID]
+QWEN_ASSISTANT_SUFFIX_IDS = [
+    QWEN_IM_END_ID,
+    QWEN_NEWLINE_ID,
+    QWEN_IM_START_ID,
+    QWEN_ASSISTANT_ID,
+    QWEN_NEWLINE_ID,
+]
+DEFAULT_QWEN3_TTS_REF_AUDIO = "vllm-omni/tests/assets/qwen3_tts/clone_2.wav"
+DEFAULT_QWEN3_TTS_REF_TEXT = (
+    "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
+)
+
+
+def default_qwen3_tts_ref_audio_path() -> str:
+    """Return absolute path to the bundled ``clone_2.wav`` reference asset."""
+    bundled = Path(__file__).resolve().parents[3] / "tests" / "assets" / "qwen3_tts" / "clone_2.wav"
+    if bundled.is_file():
+        return str(bundled)
+    return DEFAULT_QWEN3_TTS_REF_AUDIO
+
+
+def _as_list(value: Any) -> list[Any]:
+    if value is None:
+        return []
+    return value if isinstance(value, list) else [value]
+
+
+def _as_prompt_dict(prompt_item: Any) -> dict[str, Any]:
+    return prompt_item if isinstance(prompt_item, dict) else {}
+
+
+def _first_value(value: Any, default: Any = None) -> Any:
+    if isinstance(value, list):
+        return value[0] if value else default
+    return default if value is None else value
+
+
+def _first_bool(value: Any, default: bool = False) -> bool:
+    value = _first_value(value, default)
+    if isinstance(value, str):
+        return value.strip().lower() in {"1", "true", "yes", "on"}
+    return bool(value)
+
+
+def _normalize_qwen3_tts_speaker(speaker: Any) -> Any:
+    if not isinstance(speaker, str):
+        return speaker
+    speaker = speaker.strip()
+    if not speaker:
+        return speaker
+    if "_" in speaker:
+        return speaker
+    return speaker[0].upper() + speaker[1:].lower()
+
+
+def _extract_output(source_output: Any) -> Any:
+    outputs = getattr(source_output, "outputs", None)
+    if isinstance(outputs, list) and outputs:
+        return outputs[0]
+    return source_output
+
+
+def _extract_text(source_output: Any) -> str:
+    output = _extract_output(source_output)
+    cumulative_text = getattr(output, "cumulative_text", None)
+    if isinstance(cumulative_text, str) and cumulative_text:
+        return cumulative_text
+    text = getattr(output, "text", None)
+    if isinstance(text, str):
+        return text
+    mm = getattr(output, "multimodal_output", None)
+    if isinstance(mm, dict):
+        for key in ("text", "transcript", "asr_text"):
+            value = mm.get(key)
+            if isinstance(value, str):
+                return value
+            if isinstance(value, list) and value and isinstance(value[0], str):
+                return value[0]
+    return ""
+
+
+def _extract_token_ids(source_output: Any) -> list[int]:
+    output = _extract_output(source_output)
+    token_ids = getattr(output, "cumulative_token_ids", None)
+    if isinstance(token_ids, list):
+        return [int(token_id) for token_id in token_ids if isinstance(token_id, int)]
+    return []
+
+
+def _trim_aura_response_token_ids(token_ids: list[int]) -> list[int]:
+    ids = list(token_ids)
+    if ids[: len(QWEN_ASSISTANT_PREFIX_IDS)] == QWEN_ASSISTANT_PREFIX_IDS:
+        ids = ids[len(QWEN_ASSISTANT_PREFIX_IDS) :]
+    if QWEN_IM_END_ID in ids:
+        ids = ids[: ids.index(QWEN_IM_END_ID)]
+    while ids and ids[-1] in {QWEN_IM_START_ID, QWEN_IM_END_ID, QWEN_NEWLINE_ID}:
+        ids.pop()
+    return ids
+
+
+def _qwen3_tts_assistant_token_ids_from_aura(source_output: Any) -> list[int]:
+    content_ids = _trim_aura_response_token_ids(_extract_token_ids(source_output))
+    if not content_ids:
+        return []
+    return QWEN_ASSISTANT_PREFIX_IDS + content_ids + QWEN_ASSISTANT_SUFFIX_IDS
+
+
+def _source_prompt_by_request_id(source_outputs: list[Any], prompt: Any) -> dict[str, dict[str, Any]]:
+    prompts = _as_list(prompt)
+    return {
+        str(getattr(source_output, "request_id", idx)): _as_prompt_dict(prompt_item)
+        for idx, (source_output, prompt_item) in enumerate(zip(source_outputs, prompts))
+    }
+
+
+def _vision_placeholder(multi_modal_data: dict[str, Any]) -> str:
+    if "video" in multi_modal_data:
+        return "<|vision_start|><|video_pad|><|vision_end|>"
+    if "image" in multi_modal_data:
+        return "<|vision_start|><|image_pad|><|vision_end|>"
+    return ""
+
+
+def _vision_multimodal_data(multi_modal_data: dict[str, Any]) -> dict[str, Any]:
+    return {key: value for key, value in multi_modal_data.items() if key in {"image", "video"}}
+
+
+def _aura_prompt(system_prompt: str, transcript: str, multi_modal_data: dict[str, Any]) -> str:
+    vision = _vision_placeholder(multi_modal_data)
+    query = transcript.strip()
+    user_body = f"{vision}{query}" if query else vision
+    return (
+        f"<|im_start|>system\n{system_prompt}<|im_end|>\n"
+        f"<|im_start|>user\n{user_body}<|im_end|>\n"
+        "<|im_start|>assistant\n"
+    )
+
+
+def asr2aura(
+    source_outputs: list[Any],
+    prompt: Any = None,
+    requires_multimodal_data: bool = True,
+) -> list[dict[str, Any]]:
+    """Build AURA Qwen3-VL prompts from ASR transcripts and original video payloads."""
+    prompt_by_request_id = _source_prompt_by_request_id(source_outputs, prompt)
+    next_inputs: list[dict[str, Any]] = []
+    for idx, source_output in enumerate(source_outputs):
+        src_prompt = prompt_by_request_id.get(str(getattr(source_output, "request_id", idx)), {})
+        additional_info = src_prompt.get("additional_information") or {}
+        system_prompt = _first_value(additional_info.get("aura_system_prompt"), DEFAULT_AURA_SYSTEM_PROMPT)
+        transcript = _extract_text(source_output)
+        multi_modal_data = {}
+        source_multi_modal_data = src_prompt.get("multi_modal_data") or {}
+        if isinstance(source_multi_modal_data, dict):
+            multi_modal_data.update(source_multi_modal_data)
+        deferred_multi_modal_data = additional_info.get("deferred_multi_modal_data") or {}
+        if isinstance(deferred_multi_modal_data, dict):
+            multi_modal_data.update(deferred_multi_modal_data)
+        multi_modal_data = _vision_multimodal_data(multi_modal_data)
+
+        next_input: dict[str, Any] = {
+            "prompt": _aura_prompt(str(system_prompt), transcript, multi_modal_data),
+        }
+        if requires_multimodal_data:
+            next_input["multi_modal_data"] = multi_modal_data
+        if src_prompt.get("mm_processor_kwargs") is not None:
+            next_input["mm_processor_kwargs"] = src_prompt.get("mm_processor_kwargs")
+        next_inputs.append(next_input)
+    return next_inputs
+
+
+def _estimate_ref_code_len_from_ref_audio(ref_audio: Any) -> int | None:
+    """Estimate Qwen3-TTS ref_code length from a ref-audio payload.
+
+    For Qwen3-TTS 12Hz models, code length is approximately:
+        ceil(duration_seconds * 12.5)
+    i.e. one codec frame per 1920 samples at 24kHz.
+    """
+
+    codec_frame_rate = 24000.0 / 1920.0
+
+    # Unwrap common list wrappers.
+    item = ref_audio
+    while isinstance(item, list) and item:
+        item = item[0]
+
+    # Accept tuple/list like (wav, sr).
+    if isinstance(item, (tuple, list)) and len(item) == 2 and isinstance(item[1], (int, float)):
+        wav, sr = item
+        sr_i = int(sr)
+        if sr_i <= 0:
+            return None
+        if hasattr(wav, "__len__"):
+            n_samples = len(wav)
+        elif hasattr(wav, "shape"):
+            shape = getattr(wav, "shape", None)
+            if not shape:
+                return None
+            n_samples = shape[-1] if len(shape) > 1 else shape[0]
+        else:
+            return None
+        if n_samples <= 0:
+            return None
+        return max(1, int(math.ceil((float(n_samples) / float(sr_i)) * codec_frame_rate)))
+
+    # Accept file path (wav only).
+    if isinstance(item, str) and item:
+        audio_path = item
+        if not os.path.isfile(audio_path) or not audio_path.lower().endswith(".wav"):
+            return None
+        try:
+            info = sf.info(audio_path)
+            n_frames = int(info.frames)
+            sr = int(info.samplerate)
+            if n_frames <= 0 or sr <= 0:
+                return None
+            return max(1, int(math.ceil((float(n_frames) / float(sr)) * codec_frame_rate)))
+        except Exception:
+            return None
+
+    return None
+
+
+def _estimate_tts_prompt_len_from_token_ids(
+    token_ids: list[int],
+    *,
+    task_type: str = "Base",
+    language: str = "Chinese",
+    instruct: str = "",
+    x_vector_only_mode: bool = False,
+    non_streaming_mode: bool | None = None,
+    ref_code_len: int | None = None,
+) -> int:
+    """Estimate Talker prefill length from prompt structure.
+
+    This mirrors Qwen3-TTS prompt assembly at length level:
+      prompt_len = instruct_len + role_len + codec_prefix_len + text/icl term
+    """
+
+    # Official defaults: Base -> streaming, others -> non-streaming.
+    if non_streaming_mode is None:
+        non_streaming_mode = task_type in ("CustomVoice", "VoiceDesign")
+
+    # We do not have tokenizer here; use char length as a monotonic proxy.
+    instruct_len = len(instruct.strip()) if isinstance(instruct, str) else 0
+    assistant_len = max(0, len(token_ids))
+
+    # role_len = 3; codec_prefix_len = (prefill_len + speaker_len + 2) - 1
+    # prefill_len = 4 when language_id exists else 3. Use non-auto language as
+    # the language-id-present proxy.
+    has_language_id = isinstance(language, str) and language.strip().lower() != "auto"
+    prefill_len = 4 if has_language_id else 3
+    speaker_len = 1 if task_type in ("CustomVoice", "Base") else 0
+    base_len = instruct_len + 3 + (prefill_len + speaker_len + 2 - 1)
+    if task_type in ("CustomVoice", "VoiceDesign"):
+        if non_streaming_mode:
+            prompt_len = base_len + max(0, assistant_len - 6)
+        else:
+            prompt_len = base_len + 1
+        return max(32, min(4096, int(prompt_len)))
+
+    if task_type == "Base":
+        in_context_mode = not bool(x_vector_only_mode)
+        if in_context_mode and ref_code_len is not None:
+            codec_lens = 1 + int(ref_code_len)
+            if non_streaming_mode:
+                # Exact non-streaming ICL needs ref_ids token length; unavailable
+                # in this processor. Keep a conservative upper estimate.
+                prompt_len = base_len + codec_lens + max(0, assistant_len - 8) + 1
+            else:
+                # Streaming ICL exact length term: 1 + ref_code_len
+                prompt_len = base_len + codec_lens
+        else:
+            # Base x-vector-only (or missing ref_code length) follows CV shape.
+            if non_streaming_mode:
+                prompt_len = base_len + max(0, assistant_len - 6)
+            else:
+                prompt_len = base_len + 1
+        return max(32, min(4096, int(prompt_len)))
+
+    # Defensive fallback for unknown task types.
+    return max(32, min(4096, int(base_len + max(assistant_len, 1))))
+
+
+def aura2tts(
+    source_outputs: list[Any],
+    prompt: Any = None,
+    requires_multimodal_data: bool = False,
+) -> list[OmniTokensPrompt]:
+    """Convert AURA text output into Qwen3-TTS Talker requests."""
+    del requires_multimodal_data
+    prompt_by_request_id = _source_prompt_by_request_id(source_outputs, prompt)
+    next_inputs: list[OmniTokensPrompt] = []
+    for idx, source_output in enumerate(source_outputs):
+        text = _extract_text(source_output).strip()
+        if not text or text == SILENT_TEXT:
+            continue
+
+        src_prompt = prompt_by_request_id.get(str(getattr(source_output, "request_id", idx)), {})
+        additional_info = src_prompt.get("additional_information") or {}
+        task_type = _first_value(additional_info.get("tts_task_type"), "Base")
+        language = _first_value(additional_info.get("tts_language"), "English")
+        instruct = _first_value(additional_info.get("tts_instruct"), "")
+        x_vector_only_mode = _first_bool(additional_info.get("tts_x_vector_only_mode"), False)
+        non_streaming_mode_raw = _first_value(additional_info.get("tts_non_streaming_mode"), None)
+        non_streaming_mode = non_streaming_mode_raw if isinstance(non_streaming_mode_raw, bool) else None
+        ref_code_len_raw = _first_value(additional_info.get("tts_ref_code_length"), None)
+        ref_code_len = int(ref_code_len_raw) if isinstance(ref_code_len_raw, int) else None
+        ref_audio = None
+        ref_text = None
+        if task_type == "Base" and not x_vector_only_mode and ref_code_len is None:
+            ref_audio = _first_value(additional_info.get("tts_ref_audio"), None)
+            ref_code_len = _estimate_ref_code_len_from_ref_audio(ref_audio)
+
+        assistant_token_ids_for_len = _qwen3_tts_assistant_token_ids_from_aura(source_output)
+        pass_token_ids = _first_bool(additional_info.get("tts_pass_token_ids"), False)
+        tts_info = {
+            "task_type": [task_type],
+            "language": [language],
+            "instruct": [instruct],
+            "max_new_tokens": [int(_first_value(additional_info.get("tts_max_new_tokens"), 2048))],
+        }
+        if pass_token_ids and assistant_token_ids_for_len:
+            tts_info[PRECOMPUTED_TEXT_IDS_KEY] = [assistant_token_ids_for_len]
+        else:
+            tts_info["text"] = [text]
+        if ref_code_len is not None:
+            tts_info["ref_code_length"] = [int(ref_code_len)]
+        prompt_len = _estimate_tts_prompt_len_from_token_ids(
+            assistant_token_ids_for_len if assistant_token_ids_for_len else [0] * max(0, len(text)),
+            task_type=str(task_type),
+            language=str(language),
+            instruct=str(instruct),
+            x_vector_only_mode=x_vector_only_mode,
+            non_streaming_mode=non_streaming_mode,
+            ref_code_len=ref_code_len,
+        )
+
+        if task_type == "Base":
+            ref_audio = ref_audio or _first_value(additional_info.get("tts_ref_audio"), None)
+            ref_text = _first_value(additional_info.get("tts_ref_text"), None)
+            if not ref_audio or not ref_text:
+                raise ValueError("AURA Base TTS requires tts_ref_audio and tts_ref_text.")
+            x_vector_only_mode = _first_bool(additional_info.get("tts_x_vector_only_mode"), False)
+            tts_info["ref_audio"] = [ref_audio]
+            tts_info["ref_text"] = [ref_text]
+            tts_info["x_vector_only_mode"] = [x_vector_only_mode]
+        elif task_type == "CustomVoice":
+            tts_info["speaker"] = [
+                _normalize_qwen3_tts_speaker(_first_value(additional_info.get("tts_speaker"), "Vivian"))
+            ]
+        next_inputs.append(
+            OmniTokensPrompt(
+                prompt_token_ids=[0] * prompt_len,
+                additional_information=tts_info,
+                multi_modal_data=None,
+                mm_processor_kwargs=None,
+            )
+        )
+    return next_inputs

vllm_omni/worker/gpu_model_runner.py

@@ -1,4 +1,5 @@
 import contextlib
+import inspect
 from collections.abc import Callable
 from dataclasses import replace
 from typing import TYPE_CHECKING, Any, cast
@@ -48,6 +49,30 @@ else:
 logger = init_logger(__name__)
 
 
+def _filter_mrope_kwargs_for_model(model: object, kwargs: dict[str, Any]) -> dict[str, Any]:
+    """Return only M-RoPE kwargs accepted by the model implementation."""
+    method = getattr(model, "get_mrope_input_positions")
+    try:
+        signature = inspect.signature(method)
+    except (TypeError, ValueError):
+        return kwargs
+
+    if any(param.kind == inspect.Parameter.VAR_KEYWORD for param in signature.parameters.values()):
+        return kwargs
+
+    accepted = {
+        name
+        for name, param in signature.parameters.items()
+        if name not in {"self", "input_tokens"}
+        and param.kind
+        in {
+            inspect.Parameter.POSITIONAL_OR_KEYWORD,
+            inspect.Parameter.KEYWORD_ONLY,
+        }
+    }
+    return {key: value for key, value in kwargs.items() if key in accepted}
+
+
 class OmniGPUModelRunner(GPUModelRunner):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -338,7 +363,7 @@ class OmniGPUModelRunner(GPUModelRunner):
                 kwargs["target_w"] = target_w
             req_state.mrope_positions, req_state.mrope_position_delta = self.model.get_mrope_input_positions(
                 req_state.prompt_token_ids,
-                **kwargs,
+                **_filter_mrope_kwargs_for_model(self.model, kwargs),
             )
         else:
             req_state.mrope_positions, req_state.mrope_position_delta = MRotaryEmbedding.get_input_positions_tensor(