[Feature] add support for Qwen3-omni (#55)

Signed-off-by: Rein Yang <ruiruyang2@gmail.com>
Co-authored-by: Rein <rein@Reins-MacBook-Air.local>

docs/configuration/stage_configs/qwen2_5_omni.yaml

@@ -30,7 +30,7 @@ stage_args:
     runtime:
       process: true
       devices: "1"
-      max_batch_size: 3
+      max_batch_size: 1
     engine_args:
       model_stage: talker
       model_arch: Qwen2_5OmniForConditionalGeneration

examples/offline_inference/qwen2_5_omni/end2end.py

@@ -16,7 +16,7 @@ from vllm.multimodal.image import convert_image_mode
 from vllm.sampling_params import SamplingParams
 from vllm.utils import FlexibleArgumentParser
 
-from vllm_omni import Omni
+from vllm_omni.entrypoints.omni_llm import OmniLLM
 
 SEED = 42
 
@@ -140,7 +140,7 @@ def main(args):
     model_name = "Qwen/Qwen2.5-Omni-7B"
     query_result = query_map[args.query_type]()
 
-    omni_llm = Omni(
+    omni_llm = OmniLLM(
         model=model_name,
         log_stats=args.enable_stats,
         log_file=("omni_llm_pipeline.log" if args.enable_stats else None),

examples/offline_inference/qwen3_omni/README.md

@@ -0,0 +1,34 @@
+# Offline Example of vLLM-Omni for Qwen2.5-omni
+
+## 🛠️ Installation
+
+Please refer to [README.md](../../../README.md)
+
+## Run examples (Qwen3-omni)
+### Multiple Prompts
+Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). For processing dataset please refer to [Qwen2.5-omni README.md](../qwen2_5_omni/README.md)
+Get into the example folder
+```bash
+cd examples/offline_inference/qwen3_omni
+```
+Then run the command below.
+```bash
+bash run_multiple_prompts.sh
+```
+### Single Prompts
+Get into the example folder
+```bash
+cd examples/offline_inference/qwen3_omni
+```
+Then run the command below.
+```bash
+bash run_single_prompt.sh
+```
+
+### FAQ
+
+If you encounter error about backend of librosa, try to install ffmpeg with command below.
+```
+sudo apt update
+sudo apt install ffmpeg
+```

examples/offline_inference/qwen3_omni/end2end.py

@@ -0,0 +1,289 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+This example shows how to use vLLM for running offline inference
+with the correct prompt format on Qwen3-Omni (thinker only).
+"""
+
+import os
+from typing import NamedTuple
+
+import soundfile as sf
+from vllm import SamplingParams
+from vllm.assets.audio import AudioAsset
+from vllm.assets.image import ImageAsset
+from vllm.assets.video import VideoAsset
+from vllm.multimodal.image import convert_image_mode
+from vllm.utils import FlexibleArgumentParser
+
+from vllm_omni.entrypoints.omni_llm import OmniLLM
+
+SEED = 42
+
+
+class QueryResult(NamedTuple):
+    inputs: dict
+    limit_mm_per_prompt: dict[str, int]
+
+
+# NOTE: The default `max_num_seqs` and `max_model_len` may result in OOM on
+# lower-end GPUs.
+# Unless specified, these settings have been tested to work on a single L4.
+
+default_system = (
+    "You are Qwen, a virtual human developed by the Qwen Team, Alibaba "
+    "Group, capable of perceiving auditory and visual inputs, as well as "
+    "generating text and speech."
+)
+
+
+def get_text_query(question: str = None) -> QueryResult:
+    if question is None:
+        question = "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words."
+    prompt = (
+        f"<|im_start|>system\n{default_system}<|im_end|>\n"
+        "<|im_start|>user\n"
+        f"{question}<|im_end|>\n"
+        f"<|im_start|>assistant\n"
+    )
+    return QueryResult(
+        inputs={
+            "prompt": prompt,
+        },
+        limit_mm_per_prompt={},
+    )
+
+
+def get_video_query(question: str = None) -> QueryResult:
+    if question is None:
+        question = "Why is this video funny?"
+    prompt = (
+        f"<|im_start|>system\n{default_system}<|im_end|>\n"
+        "<|im_start|>user\n<|vision_start|><|video_pad|><|vision_end|>"
+        f"{question}<|im_end|>\n"
+        f"<|im_start|>assistant\n"
+    )
+
+    return QueryResult(
+        inputs={
+            "prompt": prompt,
+            "multi_modal_data": {
+                "video": VideoAsset(name="baby_reading", num_frames=16).np_ndarrays,
+            },
+        },
+        limit_mm_per_prompt={"video": 1},
+    )
+
+
+def get_image_query(question: str = None) -> QueryResult:
+    if question is None:
+        question = "What is the content of this image?"
+    prompt = (
+        f"<|im_start|>system\n{default_system}<|im_end|>\n"
+        "<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>"
+        f"{question}<|im_end|>\n"
+        f"<|im_start|>assistant\n"
+    )
+    return QueryResult(
+        inputs={
+            "prompt": prompt,
+            "multi_modal_data": {
+                "image": convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB"),
+            },
+        },
+        limit_mm_per_prompt={"image": 1},
+    )
+
+
+def get_audio_query(question: str = None) -> QueryResult:
+    if question is None:
+        question = "What is the content of this audio?"
+    prompt = (
+        f"<|im_start|>system\n{default_system}<|im_end|>\n"
+        "<|im_start|>user\n<|audio_start|><|audio_pad|><|audio_end|>"
+        f"{question}<|im_end|>\n"
+        f"<|im_start|>assistant\n"
+    )
+    return QueryResult(
+        inputs={
+            "prompt": prompt,
+            "multi_modal_data": {
+                "audio": AudioAsset("mary_had_lamb").audio_and_sample_rate,
+            },
+        },
+        limit_mm_per_prompt={"audio": 1},
+    )
+
+
+query_map = {
+    "text": get_text_query,
+    "use_audio": get_audio_query,
+    "use_image": get_image_query,
+    "use_video": get_video_query,
+}
+
+
+def main(args):
+    model_name = "/mnt/gh/yrr_vllm/Qwen3-Omni-30B-A3B-Instruct-Full"
+    query_result = query_map[args.query_type]()
+
+    omni_llm = OmniLLM(
+        model=model_name,
+    )
+
+    thinker_sampling_params = SamplingParams(
+        temperature=0.4,
+        top_p=0.9,
+        top_k=-1,
+        max_tokens=1200,
+        repetition_penalty=1.05,
+        logit_bias={},
+        seed=SEED,
+    )
+
+    talker_sampling_params = SamplingParams(
+        temperature=0.9,
+        top_k=50,
+        max_tokens=4096,
+        seed=SEED,
+        detokenize=False,
+        repetition_penalty=1.05,
+        stop_token_ids=[2150],  # TALKER_CODEC_EOS_TOKEN_ID
+    )
+
+    # Sampling parameters for Code2Wav stage (audio generation)
+    code2wav_sampling_params = SamplingParams(
+        temperature=0.0,
+        top_p=1.0,
+        top_k=-1,
+        max_tokens=4096 * 16,
+        seed=SEED,
+        detokenize=True,
+        repetition_penalty=1.1,
+    )
+
+    sampling_params_list = [
+        thinker_sampling_params,
+        talker_sampling_params,  # code predictor is integrated into talker for Qwen3 Omni
+        code2wav_sampling_params,
+    ]
+
+    if args.txt_prompts is None:
+        prompts = [query_result.inputs for _ in range(args.num_prompts)]
+    else:
+        assert args.query_type == "text", "txt-prompts is only supported for text query type"
+        with open(args.txt_prompts, encoding="utf-8") as f:
+            lines = [ln.strip() for ln in f.readlines()]
+            prompts = [get_text_query(ln).inputs for ln in lines if ln != ""]
+            print(f"[Info] Loaded {len(prompts)} prompts from {args.txt_prompts}")
+
+    omni_outputs = omni_llm.generate(prompts, sampling_params_list)
+    # Determine output directory: prefer --output-dir; fallback to --output-wav
+    output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav
+    os.makedirs(output_dir, exist_ok=True)
+
+    for stage_outputs in omni_outputs:
+        if stage_outputs.final_output_type == "text":
+            for output in stage_outputs.request_output:
+                request_id = int(output.request_id)
+                text_output = output.outputs[0].text
+                # Save aligned text file per request
+                prompt_text = prompts[request_id]["prompt"]
+                out_txt = os.path.join(output_dir, f"{request_id:05d}.txt")
+                lines = []
+                lines.append("Prompt:\n")
+                lines.append(str(prompt_text) + "\n")
+                lines.append("vllm_text_output:\n")
+                lines.append(str(text_output).strip() + "\n")
+                try:
+                    with open(out_txt, "w", encoding="utf-8") as f:
+                        f.writelines(lines)
+                except Exception as e:
+                    print(f"[Warn] Failed writing text file {out_txt}: {e}")
+                print(f"Request ID: {request_id}, Text saved to {out_txt}")
+        elif stage_outputs.final_output_type == "audio":
+            for output in stage_outputs.request_output:
+                request_id = int(output.request_id)
+                audio_tensor = output.multimodal_output["audio"]
+                output_wav = os.path.join(output_dir, f"output_{output.request_id}.wav")
+
+                # Convert to numpy array and ensure correct format
+                audio_numpy = audio_tensor.float().detach().cpu().numpy()
+
+                # Ensure audio is 1D (flatten if needed)
+                if audio_numpy.ndim > 1:
+                    audio_numpy = audio_numpy.flatten()
+
+                # Save audio file with explicit WAV format
+                sf.write(output_wav, audio_numpy, samplerate=24000, format="WAV")
+                print(f"Request ID: {request_id}, Saved audio to {output_wav}")
+
+
+def parse_args():
+    parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
+    parser.add_argument(
+        "--query-type",
+        "-q",
+        type=str,
+        default="mixed_modalities",
+        choices=query_map.keys(),
+        help="Query type.",
+    )
+    parser.add_argument(
+        "--enable-stats",
+        action="store_true",
+        default=False,
+        help="Enable writing detailed statistics (default: disabled)",
+    )
+    parser.add_argument(
+        "--init-sleep-seconds",
+        type=int,
+        default=20,
+        help="Sleep seconds after starting each stage process to allow initialization (default: 20)",
+    )
+    parser.add_argument(
+        "--batch-timeout",
+        type=int,
+        default=5,
+        help="Timeout for batching in seconds (default: 5)",
+    )
+    parser.add_argument(
+        "--init-timeout",
+        type=int,
+        default=300,
+        help="Timeout for initializing stages in seconds (default: 300)",
+    )
+    parser.add_argument(
+        "--shm-threshold-bytes",
+        type=int,
+        default=65536,
+        help="Threshold for using shared memory in bytes (default: 65536)",
+    )
+    parser.add_argument(
+        "--output-wav",
+        default="output_audio",
+        help="[Deprecated] Output wav directory (use --output-dir).",
+    )
+    parser.add_argument(
+        "--num-prompts",
+        type=int,
+        default=1,
+        help="Number of prompts to generate.",
+    )
+    parser.add_argument(
+        "--txt-prompts",
+        type=str,
+        default=None,
+        help="Path to a .txt file with one prompt per line (preferred).",
+    )
+
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    main(args)
+
+    # use examples:
+    # python end2end.py --model Qwen/Qwen3-Omni-30B-A3B-Instruct-Full --query-type text
+    # python end2end.py --model /mnt/gh/yrr_vllm/Qwen3-Omni-30B-A3B-Instruct-Full --query-type use_video_only

examples/offline_inference/qwen3_omni/run_multiple_prompts.sh

@@ -0,0 +1,3 @@
+python end2end.py --output-wav output_audio \
+                  --query-type text \
+                  --txt-prompts top10.txt

examples/offline_inference/qwen3_omni/run_single_prompt.sh

@@ -0,0 +1,2 @@
+python end2end.py --output-wav output_audio \
+                  --query-type use_audio

examples/online_serving/qwen3_omni/README.md

@@ -0,0 +1,77 @@
+# Online serving Example of vLLM-omni for Qwen2.5-omni
+
+## 🛠️ Installation
+
+Please refer to [README.md](../../README.md)
+
+## Run examples (Qwen3-Omni)
+
+Launch the server
+```bash
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
+```
+
+If you have custom stage configs file, launch the server with command below
+```bash
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
+```
+
+Get into the example folder
+```bash
+cd examples/online_serving
+```
+
+Send request via python
+```bash
+python openai_chat_completion_client_for_multimodal_generation.py --query-type mixed_modalities
+```
+
+Send request via curl
+```bash
+bash run_curl_multimodal_generation.sh mixed_modalities
+```
+
+### FAQ
+
+If you encounter error about backend of librosa, try to install ffmpeg with command below.
+```
+sudo apt update
+sudo apt install ffmpeg
+```
+
+## Run Local Web UI Demo
+
+This Web UI demo allows users to interact with the model through a web browser.
+
+### Running Gradio Demo
+
+Once vllm and vllm-omni are installed, you can launch the web service built on AsyncOmniLLM by
+
+```bash
+python gradio_demo.py  --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 7861
+```
+
+Then open `http://localhost:7861/` on your local browser to interact with the web UI.
+
+
+### Options
+
+The gradio demo also supports running with an existing API server and can be customized with the following arguments.
+
+
+```bash
+python gradio_demo.py \
+    --model Qwen/Qwen3-Omni-30B-A3B-Instruct \
+    --use-api-server \
+    --api-base http://localhost:8091/v1 \
+    --ip 127.0.0.1 \
+    --port 7861
+```
+
+- `--model`: Model name
+- `--use-api-server`: If set, connect to an existing vLLM HTTP API server instead of running AsyncOmniLLM locally.
+- `--api-base`: Base URL for vllm serve (only used when `use-api-server` is set, default: http://localhost:8091/v1)
+- `--ip`: Host/IP for Gradio server (default: 127.0.0.1)
+- `--port`: Port for Gradio server (default: 7861)
+- `--stage-configs-path`: Path to custom stage configs YAML file (optional)
+- `--share`: Share the Gradio demo publicly (creates a public link)

examples/online_serving/qwen3_omni/gradio_demo.py

@@ -0,0 +1,537 @@
+import argparse
+import os
+import random
+import signal
+import sys
+from pathlib import Path
+from types import SimpleNamespace
+from typing import Any, Optional
+
+import gradio as gr
+import numpy as np
+import soundfile as sf
+import torch
+from PIL import Image
+from vllm.assets.video import video_get_metadata, video_to_ndarrays
+from vllm.sampling_params import SamplingParams
+
+from vllm_omni.entrypoints.async_omni_llm import AsyncOmniLLM
+
+# Import utils from offline inference example
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../offline_inference/qwen3_omni"))
+try:
+    from utils import make_omni_prompt
+except ImportError:
+    # Fallback if utils doesn't exist - we'll build prompts directly
+    make_omni_prompt = None
+
+SEED = 42
+ASYNC_INIT_TIMEOUT = 600
+
+SUPPORTED_MODELS: dict[str, dict[str, Any]] = {
+    "Qwen/Qwen3-Omni-30B-A3B-Instruct": {
+        "sampling_params": {
+            "thinker": {
+                "temperature": 0.4,
+                "top_p": 0.9,
+                "top_k": 1,
+                "max_tokens": 16384,
+                "detokenize": True,
+                "repetition_penalty": 1.05,
+                "stop_token_ids": [151645],
+                "seed": SEED,
+            },
+            "talker": {
+                "temperature": 0.0,
+                "top_p": 1.0,
+                "top_k": -1,
+                "max_tokens": 4096,
+                "detokenize": True,
+                "repetition_penalty": 1.1,
+                "stop_token_ids": [2150],
+            },
+            "code2wav": {
+                "temperature": 0.0,
+                "top_p": 1.0,
+                "top_k": -1,
+                "max_tokens": 2048,
+                "detokenize": True,
+                "repetition_penalty": 1.1,
+            },
+        },
+    },
+}
+# Ensure deterministic behavior across runs.
+random.seed(SEED)
+np.random.seed(SEED)
+torch.manual_seed(SEED)
+torch.cuda.manual_seed(SEED)
+torch.cuda.manual_seed_all(SEED)
+torch.backends.cudnn.deterministic = True
+torch.backends.cudnn.benchmark = False
+os.environ["PYTHONHASHSEED"] = str(SEED)
+os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Gradio demo for Qwen3-Omni online inference.")
+    parser.add_argument(
+        "--model",
+        default="Qwen/Qwen3-Omni-30B-A3B-Instruct",
+        help="Path to local model checkpoint.",
+    )
+    parser.add_argument(
+        "--ip",
+        default="127.0.0.1",
+        help="Host/IP for gradio `launch`.",
+    )
+    parser.add_argument("--port", type=int, default=7861, help="Port for gradio `launch`.")
+    parser.add_argument("--share", action="store_true", help="Share the Gradio demo publicly.")
+    parser.add_argument(
+        "--stage-configs-path",
+        type=str,
+        default=None,
+        help="Path to custom stage configs YAML file (optional).",
+    )
+    return parser.parse_args()
+
+
+def build_sampling_params(seed: int, model_key: str) -> list[SamplingParams]:
+    """Build SamplingParams objects by reusing the dict definitions."""
+    return [SamplingParams(**params_dict) for params_dict in build_sampling_params_dict(seed, model_key)]
+
+
+def build_sampling_params_dict(seed: int, model_key: str) -> list[dict]:
+    """Build sampling params as dict for HTTP API mode."""
+    model_conf = SUPPORTED_MODELS.get(model_key)
+    if model_conf is None:
+        raise ValueError(f"Unsupported model '{model_key}'")
+
+    sampling_templates: dict[str, dict[str, Any]] = model_conf["sampling_params"]
+    sampling_params: list[dict] = []
+    for stage_name, template in sampling_templates.items():
+        params = dict(template)
+        params["seed"] = seed
+        sampling_params.append(params)
+    return sampling_params
+
+
+def create_prompt_args(base_args: argparse.Namespace) -> SimpleNamespace:
+    # The prompt builder expects a minimal namespace with these attributes.
+    return SimpleNamespace(
+        model=base_args.model,
+        prompt_type="text",
+        tokenize=True,
+        use_torchvision=True,
+        legacy_omni_video=False,
+    )
+
+
+def process_audio_file(
+    audio_file: Optional[Any],
+) -> Optional[tuple[np.ndarray, int]]:
+    """Normalize Gradio audio input to (np.ndarray, sample_rate)."""
+    if audio_file is None:
+        return None
+
+    sample_rate: Optional[int] = None
+    audio_np: Optional[np.ndarray] = None
+
+    def _load_from_path(path_str: str) -> Optional[tuple[np.ndarray, int]]:
+        if not path_str:
+            return None
+        path = Path(path_str)
+        if not path.exists():
+            return None
+        data, sr = sf.read(path)
+        if data.ndim > 1:
+            data = data[:, 0]
+        return data.astype(np.float32), int(sr)
+
+    if isinstance(audio_file, tuple):
+        if len(audio_file) == 2:
+            first, second = audio_file
+            # Case 1: (sample_rate, np.ndarray)
+            if isinstance(first, (int, float)) and isinstance(second, np.ndarray):
+                sample_rate = int(first)
+                audio_np = second
+            # Case 2: (filepath, (sample_rate, np.ndarray or list))
+            elif isinstance(first, str):
+                if isinstance(second, tuple) and len(second) == 2:
+                    sr_candidate, data_candidate = second
+                    if isinstance(sr_candidate, (int, float)) and isinstance(data_candidate, np.ndarray):
+                        sample_rate = int(sr_candidate)
+                        audio_np = data_candidate
+                if audio_np is None:
+                    loaded = _load_from_path(first)
+                    if loaded is not None:
+                        audio_np, sample_rate = loaded
+            # Case 3: (None, (sample_rate, np.ndarray))
+            elif first is None and isinstance(second, tuple) and len(second) == 2:
+                sr_candidate, data_candidate = second
+                if isinstance(sr_candidate, (int, float)) and isinstance(data_candidate, np.ndarray):
+                    sample_rate = int(sr_candidate)
+                    audio_np = data_candidate
+        elif len(audio_file) == 1 and isinstance(audio_file[0], str):
+            loaded = _load_from_path(audio_file[0])
+            if loaded is not None:
+                audio_np, sample_rate = loaded
+    elif isinstance(audio_file, str):
+        loaded = _load_from_path(audio_file)
+        if loaded is not None:
+            audio_np, sample_rate = loaded
+
+    if audio_np is None or sample_rate is None:
+        return None
+
+    if audio_np.ndim > 1:
+        audio_np = audio_np[:, 0]
+
+    return audio_np.astype(np.float32), sample_rate
+
+
+def process_image_file(image_file: Optional[Image.Image]) -> Optional[Image.Image]:
+    """Process image file from Gradio input.
+
+    Returns:
+        PIL Image in RGB mode or None if no image provided.
+    """
+    if image_file is None:
+        return None
+    # Convert to RGB if needed
+    if image_file.mode != "RGB":
+        image_file = image_file.convert("RGB")
+    return image_file
+
+
+def process_video_file(
+    video_file: Optional[str],
+    enable_audio_in_video: bool = False,
+    max_frames: int = 32,
+) -> Optional[tuple[np.ndarray, dict[str, Any], Optional[tuple[np.ndarray, int]]]]:
+    """Process video file and optionally extract audio track."""
+    if video_file is None:
+        return None
+
+    video_path = Path(video_file)
+    if not video_path.exists():
+        print(f"Video file not found: {video_path}")
+        return None
+
+    try:
+        frames = video_to_ndarrays(str(video_path), num_frames=max_frames)
+        metadata = video_get_metadata(str(video_path), num_frames=frames.shape[0])
+    except Exception as exc:  # pylint: disable=broad-except
+        print(f"Failed to decode video {video_path}: {exc}")
+        return None
+
+    audio_tuple: Optional[tuple[np.ndarray, int]] = None
+    if enable_audio_in_video:
+        try:
+            import librosa  # type: ignore import
+
+            audio_signal, sampling_rate = librosa.load(str(video_path), sr=16000)
+            audio_tuple = (audio_signal.astype(np.float32), sampling_rate)
+        except Exception as exc:  # pylint: disable=broad-except
+            print(f"Failed to extract audio from video {video_path}: {exc}")
+
+    return frames, metadata, audio_tuple
+
+
+async def run_inference_async_omni_llm(
+    omni_llm: AsyncOmniLLM,
+    sampling_params: list[SamplingParams],
+    prompt_args_template: SimpleNamespace,
+    user_prompt: str,
+    audio_file: Optional[tuple[str, tuple[int, np.ndarray]]] = None,
+    image_file: Optional[Image.Image] = None,
+    video_file: Optional[str] = None,
+    use_audio_in_video: bool = False,
+):
+    """Run inference using AsyncOmniLLM directly with multimodal support."""
+    if not user_prompt.strip() and not audio_file and not image_file and not video_file:
+        return "Please provide at least a text prompt or multimodal input.", None
+
+    try:
+        # Build prompt with multimodal data
+        prompt_args = SimpleNamespace(**prompt_args_template.__dict__)
+
+        # Process multimodal inputs
+        multi_modal_data = {}
+        mm_processor_kwargs = {}
+
+        # Process audio
+        audio_data = process_audio_file(audio_file)
+        if audio_data is not None:
+            multi_modal_data["audio"] = audio_data
+
+        # Process image
+        image_data = process_image_file(image_file)
+        if image_data is not None:
+            multi_modal_data["image"] = image_data
+
+        # Process video
+        video_payload = process_video_file(video_file, enable_audio_in_video=use_audio_in_video)
+        if video_payload is not None:
+            video_frames, video_metadata, extracted_audio = video_payload
+            video_entry: Any
+            if video_metadata:
+                video_entry = (video_frames, video_metadata)
+            else:
+                video_entry = video_frames
+            multi_modal_data["video"] = video_entry
+            if use_audio_in_video and extracted_audio is not None and "audio" not in multi_modal_data:
+                multi_modal_data["audio"] = extracted_audio
+                mm_processor_kwargs["use_audio_in_video"] = True
+
+        # Build the prompt input
+        if make_omni_prompt is not None:
+            omni_prompt = make_omni_prompt(prompt_args, user_prompt)
+            # Add multimodal data if present
+            if multi_modal_data:
+                if isinstance(omni_prompt, dict):
+                    omni_prompt["multi_modal_data"] = multi_modal_data
+                    if mm_processor_kwargs:
+                        omni_prompt["mm_processor_kwargs"] = mm_processor_kwargs
+                else:
+                    # If make_omni_prompt returns a string, we need to wrap it
+                    omni_prompt = {
+                        "prompt": omni_prompt,
+                        "multi_modal_data": multi_modal_data,
+                    }
+                    if mm_processor_kwargs:
+                        omni_prompt["mm_processor_kwargs"] = mm_processor_kwargs
+        else:
+            # Fallback: build prompt directly
+            default_system = (
+                "You are Qwen, a virtual human developed by the Qwen Team, "
+                "Alibaba Group, capable of perceiving auditory and visual inputs, "
+                "as well as generating text and speech."
+            )
+            prompt = f"<|im_start|>system\n{default_system}<|im_end|>\n<|im_start|>user\n"
+            if audio_data:
+                prompt += "<|audio_bos|><|AUDIO|><|audio_eos|>"
+            if image_data:
+                prompt += "<|vision_bos|><|IMAGE|><|vision_eos|>"
+            if video_payload is not None:
+                prompt += "<|vision_bos|><|VIDEO|><|vision_eos|>"
+            if user_prompt.strip():
+                prompt += f"{user_prompt}"
+            prompt += "<|im_end|>\n<|im_start|>assistant\n"
+
+            omni_prompt = {
+                "prompt": prompt,
+                "multi_modal_data": multi_modal_data,
+            }
+            if mm_processor_kwargs:
+                omni_prompt["mm_processor_kwargs"] = mm_processor_kwargs
+
+        request_id = "0"
+        text_outputs: list[str] = []
+        audio_output = None
+
+        async for stage_outputs in omni_llm.generate(
+            prompt=omni_prompt,
+            request_id=request_id,
+            sampling_params_list=sampling_params,
+        ):
+            # stage_outputs.request_output is a RequestOutput object, not a list
+            request_output = stage_outputs.request_output
+            if stage_outputs.final_output_type == "text":
+                if request_output.outputs:
+                    for output in request_output.outputs:
+                        if output.text:
+                            text_outputs.append(output.text)
+            elif stage_outputs.final_output_type == "audio":
+                # multimodal_output is on the RequestOutput object
+                # See vllm_omni/entrypoints/openai/serving_chat.py:680 for reference
+                if hasattr(request_output, "multimodal_output") and request_output.multimodal_output:
+                    if "audio" in request_output.multimodal_output:
+                        audio_tensor = request_output.multimodal_output["audio"]
+                        # Ensure audio is 1D (flatten if needed)
+                        if hasattr(audio_tensor, "ndim") and audio_tensor.ndim > 1:
+                            audio_tensor = audio_tensor.flatten()
+                        audio_np = audio_tensor.detach().cpu().numpy()
+                        audio_output = (
+                            24000,  # sampling rate in Hz
+                            audio_np,
+                        )
+
+        text_response = "\n\n".join(text_outputs) if text_outputs else "No text output."
+        return text_response, audio_output
+    except Exception as exc:  # pylint: disable=broad-except
+        return f"Inference failed: {exc}", None
+
+
+def build_interface(
+    omni_llm: AsyncOmniLLM,
+    sampling_params: list[SamplingParams],
+    prompt_args_template: SimpleNamespace,
+    model: str,
+):
+    """Build Gradio interface for AsyncOmniLLM mode."""
+
+    async def run_inference(
+        user_prompt: str,
+        audio_file: Optional[tuple[str, tuple[int, np.ndarray]]],
+        image_file: Optional[Image.Image],
+        video_file: Optional[str],
+        use_audio_in_video: bool,
+    ):
+        return await run_inference_async_omni_llm(
+            omni_llm,
+            sampling_params,
+            prompt_args_template,
+            user_prompt,
+            audio_file,
+            image_file,
+            video_file,
+            use_audio_in_video,
+        )
+
+    css = """
+    .media-input-container {
+        display: flex;
+        gap: 10px;
+    }
+    .media-input-container > div {
+        flex: 1;
+    }
+    .media-input-container .image-input,
+    .media-input-container .audio-input {
+        height: 300px;
+    }
+    .media-input-container .video-column {
+        height: 300px;
+        display: flex;
+        flex-direction: column;
+    }
+    .media-input-container .video-input {
+        flex: 1;
+        min-height: 0;
+    }
+    #generate-btn button {
+        width: 100%;
+    }
+    """
+
+    with gr.Blocks(css=css) as demo:
+        gr.Markdown("# vLLM-Omni Online Serving Demo")
+        gr.Markdown(f"**Model:** {model} \n\n")
+
+        with gr.Column():
+            with gr.Row():
+                input_box = gr.Textbox(
+                    label="Text Prompt",
+                    placeholder="For example: Describe what happens in the media inputs.",
+                    lines=4,
+                    scale=1,
+                )
+            with gr.Row(elem_classes="media-input-container"):
+                image_input = gr.Image(
+                    label="Image Input (optional)",
+                    type="pil",
+                    sources=["upload"],
+                    scale=1,
+                    elem_classes="image-input",
+                )
+                with gr.Column(scale=1, elem_classes="video-column"):
+                    video_input = gr.Video(
+                        label="Video Input (optional)",
+                        sources=["upload"],
+                        elem_classes="video-input",
+                    )
+                    use_audio_in_video_checkbox = gr.Checkbox(
+                        label="Use audio from video",
+                        value=False,
+                        info="Extract the video's audio track when provided.",
+                    )
+                audio_input = gr.Audio(
+                    label="Audio Input (optional)",
+                    type="numpy",
+                    sources=["upload", "microphone"],
+                    scale=1,
+                    elem_classes="audio-input",
+                )
+
+        with gr.Row():
+            generate_btn = gr.Button(
+                "Generate",
+                variant="primary",
+                size="lg",
+                elem_id="generate-btn",
+            )
+
+        with gr.Row():
+            text_output = gr.Textbox(label="Text Output", lines=10, scale=2)
+            audio_output = gr.Audio(label="Audio Output", interactive=False, scale=1)
+
+        generate_btn.click(
+            fn=run_inference,
+            inputs=[input_box, audio_input, image_input, video_input, use_audio_in_video_checkbox],
+            outputs=[text_output, audio_output],
+        )
+        demo.queue()
+    return demo
+
+
+def main():
+    args = parse_args()
+    omni_llm = None
+
+    model_name = "/".join(args.model.split("/")[-2:])
+    assert model_name in SUPPORTED_MODELS, (
+        f"Unsupported model '{model_name}'. Supported models: {SUPPORTED_MODELS.keys()}"
+    )
+
+    # Register signal handlers for graceful shutdown
+    def signal_handler(sig, frame):
+        print("\nReceived interrupt signal, shutting down...")
+        if omni_llm is not None:
+            try:
+                omni_llm.shutdown()
+            except Exception as e:
+                print(f"Error during shutdown: {e}")
+        sys.exit(0)
+
+    signal.signal(signal.SIGINT, signal_handler)
+    signal.signal(signal.SIGTERM, signal_handler)
+
+    print(f"Initializing AsyncOmniLLM with model: {args.model}")
+    if args.stage_configs_path:
+        print(f"Using custom stage configs: {args.stage_configs_path}")
+
+    sampling_params = build_sampling_params(SEED, model_name)
+    omni_llm = AsyncOmniLLM(
+        model=args.model,
+        stage_configs_path=args.stage_configs_path,
+        init_timeout=ASYNC_INIT_TIMEOUT,
+    )
+    print("✓ AsyncOmniLLM initialized successfully")
+    prompt_args_template = create_prompt_args(args)
+
+    demo = build_interface(
+        omni_llm,
+        sampling_params,
+        prompt_args_template,
+        args.model,
+    )
+    try:
+        demo.launch(
+            server_name=args.ip,
+            server_port=args.port,
+            share=args.share,
+        )
+    except KeyboardInterrupt:
+        print("\nShutting down...")
+    finally:
+        # Cleanup
+        if omni_llm is not None:
+            try:
+                omni_llm.shutdown()
+            except Exception as e:
+                print(f"Error during cleanup: {e}")
+
+
+if __name__ == "__main__":
+    main()

examples/online_serving/qwen3_omni/openai_chat_completion_client_for_multimodal_generation.py

@@ -0,0 +1,235 @@
+import base64
+from typing import NamedTuple
+
+import requests
+from openai import OpenAI
+from vllm.assets.audio import AudioAsset
+from vllm.utils import FlexibleArgumentParser
+
+# Modify OpenAI's API key and API base to use vLLM's API server.
+openai_api_key = "EMPTY"
+openai_api_base = "http://localhost:8091/v1"
+
+client = OpenAI(
+    # defaults to os.environ.get("OPENAI_API_KEY")
+    api_key=openai_api_key,
+    base_url=openai_api_base,
+)
+
+SEED = 42
+
+
+class QueryResult(NamedTuple):
+    inputs: dict
+    limit_mm_per_prompt: dict[str, int]
+
+
+def encode_base64_content_from_url(content_url: str) -> str:
+    """Encode a content retrieved from a remote url to base64 format."""
+
+    with requests.get(content_url) as response:
+        response.raise_for_status()
+        result = base64.b64encode(response.content).decode("utf-8")
+
+    return result
+
+
+def get_system_prompt():
+    return {
+        "role": "system",
+        "content": [
+            {
+                "type": "text",
+                "text": (
+                    "You are Qwen, a virtual human developed by the Qwen Team, "
+                    "Alibaba Group, capable of perceiving auditory and visual inputs, "
+                    "as well as generating text and speech."
+                ),
+            }
+        ],
+    }
+
+
+def get_text_query():
+    question = "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words."
+    prompt = {
+        "role": "user",
+        "content": [
+            {
+                "type": "text",
+                "text": f"{question}",
+            }
+        ],
+    }
+    return prompt
+
+
+default_system = (
+    "You are Qwen, a virtual human developed by the Qwen Team, Alibaba "
+    "Group, capable of perceiving auditory and visual inputs, as well as "
+    "generating text and speech."
+)
+
+
+def get_video_query() -> QueryResult:
+    question = "Why is this video funny?"
+    prompt = {
+        "role": "user",
+        "content": [
+            {
+                "type": "video_url",
+                "video_url": {
+                    "url": "https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4"
+                },
+            },
+            {
+                "type": "text",
+                "text": f"{question}",
+            },
+        ],
+    }
+    return prompt
+
+
+def get_image_query() -> QueryResult:
+    question = "What is the content of this image?"
+    prompt = {
+        "role": "user",
+        "content": [
+            {
+                "type": "image_url",
+                "image_url": {
+                    "url": "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg"
+                },
+            },
+            {
+                "type": "text",
+                "text": f"{question}",
+            },
+        ],
+    }
+    return prompt
+
+
+def get_audio_query() -> QueryResult:
+    question = "What is the content of this audio?"
+    prompt = {
+        "role": "user",
+        "content": [
+            {
+                "type": "audio_url",
+                "audio_url": {"url": AudioAsset("mary_had_lamb").url},
+            },
+            {
+                "type": "text",
+                "text": f"{question}",
+            },
+        ],
+    }
+    return prompt
+
+
+query_map = {
+    "text": get_text_query,
+    "use_audio": get_audio_query,
+    "use_image": get_image_query,
+    "use_video": get_video_query,
+}
+
+
+def run_multimodal_generation(args) -> None:
+    model_name = args.model
+    thinker_sampling_params = {
+        "temperature": 0.4,  # Deterministic
+        "top_p": 0.9,
+        "top_k": 1,
+        "max_tokens": 16384,
+        "repetition_penalty": 1.05,
+        "stop_token_ids": [151645],  # Qwen EOS token <|im_end|>
+        "seed": SEED,
+    }
+
+    # Sampling parameters for Talker stage (codec generation)
+    # Stop at codec EOS token
+    talker_sampling_params = {
+        "temperature": 0.9,
+        "top_k": 50,
+        "max_tokens": 4096,
+        "seed": SEED,
+        "detokenize": False,
+        "repetition_penalty": 1.05,
+        "stop_token_ids": [2150],  # TALKER_CODEC_EOS_TOKEN_ID
+    }
+
+    # # Sampling parameters for Code2Wav stage (audio generation)
+    code2wav_sampling_params = {
+        "temperature": 0.0,
+        "top_p": 1.0,
+        "top_k": -1,
+        "max_tokens": 4096 * 16,
+        "seed": SEED,
+        "detokenize": True,
+        "repetition_penalty": 1.1,
+    }
+
+    sampling_params_list = [
+        thinker_sampling_params,
+        talker_sampling_params,
+        code2wav_sampling_params,
+    ]
+
+    prompt = query_map[args.query_type]()
+    extra_body = {
+        "sampling_params_list": sampling_params_list  # Optional, it has a default setting in stage_configs of the corresponding model.
+    }
+
+    if args.query_type == "use_audio_in_video":
+        extra_body["mm_processor_kwargs"] = {"use_audio_in_video": True}
+
+    chat_completion = client.chat.completions.create(
+        messages=[
+            get_system_prompt(),
+            prompt,
+        ],
+        model=model_name,
+        extra_body=extra_body,
+    )
+
+    count = 0
+    for choice in chat_completion.choices:
+        print(f"choice: {choice}")
+        if choice.message.audio:
+            audio_data = base64.b64decode(choice.message.audio.data)
+            audio_file_path = f"audio_{count}.wav"
+            with open(audio_file_path, "wb") as f:
+                f.write(audio_data)
+            print(f"Audio saved to {audio_file_path}")
+            count += 1
+        elif choice.message.content:
+            print("Chat completion output from text:", choice.message.content)
+
+
+def parse_args():
+    parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
+    parser.add_argument(
+        "--query-type",
+        "-q",
+        type=str,
+        default="use_video",
+        choices=query_map.keys(),
+        help="Query type.",
+    )
+    parser.add_argument(
+        "--model",
+        "-m",
+        type=str,
+        default="Qwen/Qwen3-Omni-30B-A3B-Instruct",
+        help="Model Name / Path",
+    )
+
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    run_multimodal_generation(args)

examples/online_serving/qwen3_omni/run_curl_multimodal_generation.sh

@@ -0,0 +1,187 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# Default query type
+QUERY_TYPE="${1:-mixed_modalities}"
+
+# Validate query type
+if [[ ! "$QUERY_TYPE" =~ ^(mixed_modalities|use_audio_in_video|multi_audios|text)$ ]]; then
+    echo "Error: Invalid query type '$QUERY_TYPE'"
+    echo "Usage: $0 [mixed_modalities|use_audio_in_video|multi_audios|text]"
+    echo "  mixed_modalities: Audio + Image + Video + Text query"
+    echo "  use_audio_in_video: Video + Text query (with audio extraction from video)"
+    echo "  multi_audios: Two audio clips + Text query"
+    echo "  text: Text query"
+    exit 1
+fi
+
+SEED=42
+
+thinker_sampling_params='{
+  "temperature": 0.4,
+  "top_p": 0.9,
+  "top_k": 1,
+  "max_tokens": 16384,
+  "seed": 42,
+  "detokenize": true,
+  "repetition_penalty": 1.05,
+  "stop_token_ids": [151645]
+}'
+
+talker_sampling_params='{
+  "temperature": 0.9,
+  "top_k": 50,
+  "max_tokens": 4096,
+  "seed": 42,
+  "detokenize": true,
+  "repetition_penalty": 1.05,
+  "stop_token_ids": [2150]
+}'
+
+code2wav_sampling_params='{
+  "temperature": 0.0,
+  "top_p": 1.0,
+  "top_k": -1,
+  "max_tokens": 4096 * 16,
+  "seed": 42,
+  "detokenize": true,
+  "repetition_penalty": 1.1
+}'
+# Above is optional, it has a default setting in stage_configs of the corresponding model.
+
+# Define URLs for assets
+MARY_HAD_LAMB_AUDIO_URL="https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/mary_had_lamb.ogg"
+WINNING_CALL_AUDIO_URL="https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/winning_call.ogg"
+CHERRY_BLOSSOM_IMAGE_URL="https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg"
+SAMPLE_VIDEO_URL="https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4"
+
+# Build user content and extra fields based on query type
+case "$QUERY_TYPE" in
+  text)
+    user_content='[
+      {
+        "type": "text",
+        "text": "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words."
+      }
+    ]'
+    sampling_params_list='[
+      '"$thinker_sampling_params"',
+      '"$talker_sampling_params"',
+      '"$code2wav_sampling_params"'
+    ]'
+    mm_processor_kwargs="{}"
+    ;;
+  mixed_modalities)
+    user_content='[
+        {
+          "type": "audio_url",
+          "audio_url": {
+            "url": "'"$MARY_HAD_LAMB_AUDIO_URL"'"
+          }
+        },
+        {
+          "type": "image_url",
+          "image_url": {
+            "url": "'"$CHERRY_BLOSSOM_IMAGE_URL"'"
+          }
+        },
+        {
+          "type": "video_url",
+          "video_url": {
+            "url": "'"$SAMPLE_VIDEO_URL"'"
+          }
+        },
+        {
+          "type": "text",
+          "text": "What is recited in the audio? What is the content of this image? Why is this video funny?"
+        }
+      ]'
+    sampling_params_list='[
+      '"$thinker_sampling_params"',
+      '"$talker_sampling_params"',
+      '"$code2wav_sampling_params"'
+    ]'
+    mm_processor_kwargs="{}"
+    ;;
+  use_audio_in_video)
+    user_content='[
+        {
+          "type": "video_url",
+          "video_url": {
+            "url": "'"$SAMPLE_VIDEO_URL"'"
+          }
+        },
+        {
+          "type": "text",
+          "text": "Describe the content of the video, then convert what the baby say into text."
+        }
+      ]'
+    sampling_params_list='[
+      '"$thinker_sampling_params"',
+      '"$talker_sampling_params"',
+      '"$code2wav_sampling_params"'
+    ]'
+    mm_processor_kwargs='{
+      "use_audio_in_video": true
+    }'
+    ;;
+  multi_audios)
+    user_content='[
+        {
+          "type": "audio_url",
+          "audio_url": {
+            "url": "'"$MARY_HAD_LAMB_AUDIO_URL"'"
+          }
+        },
+        {
+          "type": "audio_url",
+          "audio_url": {
+            "url": "'"$WINNING_CALL_AUDIO_URL"'"
+          }
+        },
+        {
+          "type": "text",
+          "text": "Are these two audio clips the same?"
+        }
+      ]'
+    sampling_params_list='[
+      '"$thinker_sampling_params"',
+      '"$talker_sampling_params"',
+      '"$code2wav_sampling_params"'
+    ]'
+    mm_processor_kwargs="{}"
+    ;;
+esac
+
+echo "Running query type: $QUERY_TYPE"
+echo ""
+
+
+output=$(curl -sS -X POST http://localhost:8091/v1/chat/completions \
+    -H "Content-Type: application/json" \
+    -d @- <<EOF
+{
+  "model": "Qwen/Qwen2.5-Omni-7B",
+  "sampling_params_list": $sampling_params_list,
+  "mm_processor_kwargs": $mm_processor_kwargs,
+  "messages": [
+    {
+      "role": "system",
+      "content": [
+        {
+          "type": "text",
+          "text": "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, capable of perceiving auditory and visual inputs, as well as generating text and speech."
+        }
+      ]
+    },
+    {
+      "role": "user",
+      "content": $user_content
+    }
+  ]
+}
+EOF
+  )
+
+# Here it only shows the text content of the first choice. Audio content has many binaries, so it's not displayed here.
+echo "Output of request: $(echo "$output" | jq '.choices[0].message.content')"

vllm_omni/__init__.py

@@ -19,11 +19,12 @@ except ModuleNotFoundError as exc:  # pragma: no cover - optional dependency
         raise
     # Allow importing vllm_omni without vllm (e.g., documentation builds)
     patch = None  # type: ignore
+
 from .config import OmniModelConfig
 from .entrypoints.async_omni_llm import AsyncOmniLLM
 
 # Main entry points
-from .entrypoints.omni import Omni
+# from .entrypoints.omni import Omni
 
 from .version import __version__, __version_tuple__  # isort:skip
 
@@ -32,7 +33,7 @@ __all__ = [
     "__version__",
     "__version_tuple__",
     # Main components
-    "Omni",
+    # "Omni",
     "AsyncOmniLLM",
     # Configuration
     "OmniModelConfig",

vllm_omni/config/model.py

@@ -1,11 +1,40 @@
-from typing import Optional
+import json
+import warnings
+from importlib.util import find_spec
+from typing import Any, Literal, Optional
 
+import torch
+import vllm.envs as envs
+from omegaconf import DictConfig, OmegaConf
 from pydantic import ConfigDict
 from pydantic.dataclasses import dataclass
 from vllm.config import ModelConfig, config
+from vllm.config.model import (
+    _RUNNER_CONVERTS,
+    _RUNNER_TASKS,
+    ConvertType,
+    RunnerOption,
+    TaskOption,
+    _get_and_verify_dtype,
+    get_served_model_name,
+)
+from vllm.config.multimodal import MMCacheType, MMEncoderTPMode, MultiModalConfig
+from vllm.config.pooler import PoolerConfig
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+from vllm.transformers_utils.config import (
+    get_config,
+    get_hf_image_processor_config,
+    get_hf_text_config,
+    get_pooling_config,
+    is_interleaved,
+)
+from vllm.transformers_utils.utils import maybe_model_redirect
 
 import vllm_omni.model_executor.models as me_models
 
+logger = init_logger(__name__)
+
 
 @config
 @dataclass(config=ConfigDict(arbitrary_types_allowed=True))
@@ -37,6 +66,7 @@ class OmniModelConfig(ModelConfig):
     model_stage: str = "thinker"
     model_arch: str = "Qwen2_5OmniForConditionalGeneration"
     engine_output_type: Optional[str] = None
+    hf_config_name: Optional[str] = None
 
     @property
     def registry(self):
@@ -45,3 +75,313 @@ class OmniModelConfig(ModelConfig):
     @property
     def architectures(self) -> list[str]:
         return [self.model_arch]
+
+    def draw_hf_text_config(self):
+        # transformers' get_text_config method is used to get the text config from thinker_config.
+        # to handle the case that each model stage has their own text config,
+        # we need to draw the text config from the corresponding model stage.
+        if self.hf_config_name is None:
+            return get_hf_text_config(self.hf_config)
+        return getattr(self.hf_config, self.hf_config_name).get_text_config()
+
+    def __post_init__(
+        self,
+        # Multimodal config init vars
+        limit_mm_per_prompt: Optional[dict[str, int]],
+        media_io_kwargs: Optional[dict[str, dict[str, Any]]],
+        mm_processor_kwargs: Optional[dict[str, Any]],
+        mm_processor_cache_gb: Optional[float],
+        mm_processor_cache_type: Optional[MMCacheType],
+        mm_shm_cache_max_object_size_mb: Optional[int],
+        mm_encoder_tp_mode: Optional[MMEncoderTPMode],
+        interleave_mm_strings: Optional[bool],
+        skip_mm_profiling: Optional[bool],
+        video_pruning_rate: Optional[float],
+    ) -> None:
+        # Set the default seed to 0 in V1.
+        if envs.VLLM_USE_V1 and self.seed is None:
+            self.seed = 0
+            if not envs.VLLM_ENABLE_V1_MULTIPROCESSING:
+                logger.warning(
+                    "The global random seed is set to %d. Since "
+                    "VLLM_ENABLE_V1_MULTIPROCESSING is set to False, this may "
+                    "affect the random state of the Python process that "
+                    "launched vLLM.",
+                    self.seed,
+                )
+
+        # Keep set served_model_name before maybe_model_redirect(self.model)
+        self.served_model_name = get_served_model_name(self.model, self.served_model_name)
+        self.model = maybe_model_redirect(self.model)
+        # The tokenizer is consistent with the model by default.
+        if self.tokenizer is None:
+            self.tokenizer = self.model
+        if self.tokenizer_revision is None:
+            self.tokenizer_revision = self.revision
+        self.tokenizer = maybe_model_redirect(self.tokenizer)
+
+        if isinstance(self.hf_config_path, str):
+            self.hf_config_path = maybe_model_redirect(self.hf_config_path)
+
+        if callable(self.hf_overrides):
+            hf_overrides_kw: dict[str, Any] = {}
+            hf_overrides_fn = self.hf_overrides
+            dict_overrides: dict[str, Any] = {}
+        else:
+            # Separate dict overrides from flat ones; convert OmegaConf to dict
+            hf_overrides_kw = {}
+            dict_overrides = {}
+            for key, value in self.hf_overrides.items():
+                if isinstance(value, DictConfig):
+                    value = OmegaConf.to_container(value, resolve=True, structured_config_mode=None)
+                if isinstance(value, dict):
+                    dict_overrides[key] = value
+                else:
+                    hf_overrides_kw[key] = value
+            hf_overrides_fn = None
+
+        if self.rope_scaling:
+            hf_override: dict[str, Any] = {"rope_scaling": self.rope_scaling}
+            hf_overrides_kw.update(hf_override)
+            hf_overrides_str = json.dumps(hf_overrides_kw)
+            msg = (
+                "`--rope-scaling` will be removed in a future release. "
+                f"'Please instead use `--hf-overrides '{hf_overrides_str}'`"
+            )
+            warnings.warn(DeprecationWarning(msg), stacklevel=2)
+        if self.rope_theta is not None:
+            hf_override = {"rope_theta": self.rope_theta}
+            hf_overrides_kw.update(hf_override)
+            hf_overrides_str = json.dumps(hf_overrides_kw)
+            msg = (
+                "`--rope-theta` will be removed in a future release. "
+                f"'Please instead use `--hf-overrides '{hf_overrides_str}'`"
+            )
+            warnings.warn(DeprecationWarning(msg), stacklevel=2)
+
+        self.maybe_pull_model_tokenizer_for_runai(self.model, self.tokenizer)
+
+        if (backend := envs.VLLM_ATTENTION_BACKEND) and backend == "FLASHINFER" and find_spec("flashinfer") is None:
+            raise ValueError(
+                "VLLM_ATTENTION_BACKEND is set to FLASHINFER, but flashinfer "
+                "module was not found. See "
+                "https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile "
+                "for instructions on how to install it."
+            )
+
+        if self.override_attention_dtype is not None and not current_platform.is_rocm():
+            warnings.warn("override-attention-dtype is set but not using ROCm platform", stacklevel=2)
+
+        if self.enable_sleep_mode and not current_platform.is_sleep_mode_available():
+            raise ValueError("Sleep mode is not supported on current platform.")
+
+        hf_config = get_config(
+            self.hf_config_path or self.model,
+            self.trust_remote_code,
+            self.revision,
+            self.code_revision,
+            self.config_format,
+            hf_overrides_kw=hf_overrides_kw,
+            hf_overrides_fn=hf_overrides_fn,
+        )
+
+        if dict_overrides:
+            self._apply_dict_overrides(hf_config, dict_overrides)
+
+        self.hf_config = hf_config
+        self.hf_text_config = self.draw_hf_text_config()
+        self.attention_chunk_size = getattr(self.hf_text_config, "attention_chunk_size", None)
+        self.encoder_config = self._get_encoder_config()
+        self.hf_image_processor_config = get_hf_image_processor_config(
+            self.model, hf_token=self.hf_token, revision=self.revision
+        )
+
+        architectures = self.architectures
+        registry = self.registry
+        is_generative_model = registry.is_text_generation_model(architectures, self)
+        is_pooling_model = registry.is_pooling_model(architectures, self)
+
+        def _task_to_convert(task: TaskOption) -> ConvertType:
+            if task == "embedding" or task == "embed":
+                return "embed"
+            if task == "classify":
+                return "classify"
+            if task == "reward":
+                return "reward"
+            if task == "score":
+                new_task = self._get_default_pooling_task(architectures)
+                return "classify" if new_task == "classify" else "embed"
+
+            return "none"
+
+        if self.task is not None:
+            runner: RunnerOption = "auto"
+            convert: ConvertType | Literal["auto"] = "auto"  # type: ignore[name-defined]
+            msg_prefix = (
+                "The 'task' option has been deprecated and will be removed in v0.13.0 or v1.0, whichever comes first."
+            )
+            msg_hint = "Please remove this option."
+
+            is_generative_task = self.task in _RUNNER_TASKS["generate"]
+            is_pooling_task = self.task in _RUNNER_TASKS["pooling"]
+
+            if is_generative_model and is_pooling_model:
+                if is_generative_task:
+                    runner = "generate"
+                    convert = "auto"  # type: ignore[assignment]
+                    msg_hint = (
+                        "Please replace this option with `--runner "
+                        "generate` to continue using this model "
+                        "as a generative model."
+                    )
+                elif is_pooling_task:
+                    runner = "pooling"
+                    convert = "auto"  # type: ignore[assignment]
+                    msg_hint = (
+                        "Please replace this option with `--runner "
+                        "pooling` to continue using this model "
+                        "as a pooling model."
+                    )
+                else:  # task == "auto"
+                    pass
+            elif is_generative_model or is_pooling_model:
+                if is_generative_task:
+                    runner = "generate"
+                    convert = "auto"  # type: ignore[assignment]
+                    msg_hint = "Please remove this option"
+                elif is_pooling_task:
+                    runner = "pooling"
+                    convert = _task_to_convert(self.task)
+                    msg_hint = (
+                        "Please replace this option with `--convert "
+                        f"{convert}` to continue using this model "
+                        "as a pooling model."
+                    )
+                else:  # task == "auto"
+                    pass
+            else:
+                raise AssertionError(
+                    f"The model should be a generative or pooling model when task is set to {self.task!r}."
+                )
+
+            self.runner = runner
+            self.convert = convert  # type: ignore[assignment]
+
+            msg = f"{msg_prefix} {msg_hint}"
+            warnings.warn(msg, DeprecationWarning, stacklevel=2)
+
+        self.runner_type = self._get_runner_type(architectures, self.runner)
+        self.convert_type = self._get_convert_type(architectures, self.runner_type, self.convert)
+
+        if self.runner_type == "generate" and not is_generative_model:
+            generate_converts = _RUNNER_CONVERTS["generate"]
+            if self.convert_type not in generate_converts:
+                raise ValueError("This model does not support `--runner generate`.")
+        if self.runner_type == "pooling" and not is_pooling_model:
+            pooling_converts = _RUNNER_CONVERTS["pooling"]
+            if self.convert_type not in pooling_converts:
+                convert_option = "<" + "|".join(pooling_converts) + ">"
+                raise ValueError(
+                    "This model does not support `--runner pooling`. "
+                    f"You can pass `--convert {convert_option} to adapt "
+                    "it into a pooling model."
+                )
+
+        # Note: Initialize these attributes early because transformers fallback
+        # may fail to load dynamic modules in child processes
+        model_info, arch = registry.inspect_model_cls(architectures, self)
+        self._model_info = model_info
+        self._architecture = arch
+        logger.info("Resolved architecture: %s", arch)
+
+        # Init pooler config if needed
+        if self.runner_type == "pooling":
+            if self.override_pooler_config is not None:
+                logger.warning_once(
+                    "`override_pooler_config` is deprecated and will be "
+                    "removed in v0.12.0 or v1.0.0, whichever is sooner. "
+                    "Please use `pooler_config` instead."
+                )
+
+                if isinstance(self.override_pooler_config, dict):
+                    self.pooler_config = PoolerConfig(**self.override_pooler_config)
+                else:
+                    self.pooler_config = self.override_pooler_config
+
+            if self.pooler_config is None:
+                self.pooler_config = PoolerConfig()
+
+            base_config = get_pooling_config(self.model, self.revision)
+            if base_config is not None:
+                for k, v in base_config.items():
+                    if getattr(self.pooler_config, k) is None:
+                        setattr(self.pooler_config, k, v)
+
+            default_pooling_type = self._model_info.default_pooling_type
+            if self.pooler_config.pooling_type is None:
+                self.pooler_config.pooling_type = default_pooling_type
+
+        self.dtype: torch.dtype = _get_and_verify_dtype(
+            self.model,
+            self.hf_config,
+            self.dtype,
+            is_pooling_model=self.runner_type == "pooling",
+            revision=self.revision,
+        )
+
+        # Interleaved attention is not supported by some backends in V0
+        if (
+            not self.disable_sliding_window
+            and is_interleaved(self.hf_text_config)
+            and not envs.VLLM_USE_V1
+            and (backend := envs.VLLM_ATTENTION_BACKEND) in ("XFORMERS", "FLASHINFER")
+        ):
+            logger.warning_once(
+                "%s has interleaved attention, which is currently not "
+                "supported by the %s backend. Disabling sliding window and "
+                "capping the max length to the sliding window size (%d).",
+                self.hf_text_config.model_type,
+                backend,
+                self.hf_text_config.sliding_window,
+            )
+            self.disable_sliding_window = True
+
+        self.original_max_model_len = self.max_model_len
+        self.max_model_len = self.get_and_verify_max_len(self.max_model_len)
+        # Init multimodal config if needed
+        if self._model_info.supports_multimodal:
+            if mm_encoder_tp_mode == "data" and not self._model_info.supports_multimodal_encoder_tp_data:
+                logger.warning_once(
+                    "This model does not support `--mm-encoder-tp-mode data`. "
+                    "Falling back to `--mm-encoder-tp-mode weights`."
+                )
+                mm_encoder_tp_mode = "weights"
+
+            mm_config_kwargs = dict(
+                limit_per_prompt=limit_mm_per_prompt,
+                media_io_kwargs=media_io_kwargs,
+                mm_processor_kwargs=mm_processor_kwargs,
+                mm_processor_cache_gb=mm_processor_cache_gb,
+                mm_processor_cache_type=mm_processor_cache_type,
+                mm_shm_cache_max_object_size_mb=mm_shm_cache_max_object_size_mb,
+                mm_encoder_tp_mode=mm_encoder_tp_mode,
+                interleave_mm_strings=interleave_mm_strings,
+                skip_mm_profiling=skip_mm_profiling,
+                video_pruning_rate=video_pruning_rate,
+            )
+
+            mm_config_kwargs = {k: v for k, v in mm_config_kwargs.items() if v is not None}
+
+            self.multimodal_config = MultiModalConfig(**mm_config_kwargs)
+
+        if self.disable_sliding_window:
+            self.hf_text_config.sliding_window = None
+
+        if not self.skip_tokenizer_init:
+            self._verify_tokenizer_mode()
+
+        self.config_updated = False
+
+        self._verify_quantization()
+        self._verify_cuda_graph()
+        self._verify_bnb_config()

vllm_omni/core/sched/diffusion_scheduler.py

@@ -249,6 +249,12 @@ class DiffusionScheduler(OmniScheduler):
 
             # Get prompt logprobs for this request.
             prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
+            # Convert pooler_output tensor to dict format expected by OmniEngineCoreOutput
+            pooling_output_dict = None
+            if pooler_output is not None:
+                # Wrap tensor in dict to match OmniEngineCoreOutput.pooling_output type
+                # which expects Optional[dict[str, torch.Tensor]]
+                pooling_output_dict = {"model_outputs": pooler_output}
             if new_token_ids or pooler_output is not None or kv_transfer_params:
                 # Add EngineCoreOutput for this Request.
                 outputs[request.client_index].append(
@@ -258,7 +264,7 @@ class DiffusionScheduler(OmniScheduler):
                         finish_reason=request.get_finished_reason(),
                         new_logprobs=new_logprobs,
                         new_prompt_logprobs_tensors=prompt_logprobs_tensors,
-                        pooling_output=pooler_output,
+                        pooling_output=pooling_output_dict,
                         stop_reason=request.stop_reason,
                         events=request.take_events(),
                         kv_transfer_params=kv_transfer_params,

vllm_omni/core/sched/generation_scheduler.py

@@ -0,0 +1,313 @@
+import time
+from collections import defaultdict
+from typing import Optional
+
+from vllm.v1.core.kv_cache_manager import KVCacheBlocks
+from vllm.v1.core.sched.request_queue import create_request_queue
+from vllm.v1.core.sched.scheduler import (
+    EngineCoreOutputs,
+    Request,
+    RequestStatus,
+    SchedulerOutput,
+    SpecDecodingStats,
+)
+from vllm.v1.core.sched.utils import remove_all
+from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput
+
+from vllm_omni.core.sched.output import OmniNewRequestData
+from vllm_omni.core.sched.scheduler import OmniScheduler
+from vllm_omni.outputs import OmniModelRunnerOutput
+
+
+class GenerationScheduler(OmniScheduler):
+    def schedule(self) -> SchedulerOutput:
+        """Diffusion fast path:
+        - Feed all input tokens of the request at once
+          (if 0, allocate 1 placeholder token).
+        - If the token budget cannot be satisfied at once, fall back to the
+          default vLLM scheduling.
+        """
+
+        # Select requests with zero prompt and using pooling
+        # (diffusion results returned via pooler_output)
+        token_budget = self.max_num_scheduled_tokens
+        capacity = self.max_num_running_reqs - len(self.running)
+        scheduled_timestamp = time.monotonic()
+
+        scheduled_new_reqs: list[Request] = []
+        scheduled_resumed_reqs: list[Request] = []
+        scheduled_running_reqs: list[Request] = []
+
+        req_to_new_blocks: dict[str, KVCacheBlocks] = {}
+        num_scheduled_tokens: dict[str, int] = {}
+        scheduled_spec_decode_tokens: dict[str, list[int]] = {}
+        scheduled_encoder_inputs: dict[str, list[int]] = {}
+
+        # Temporary queue: preserve waiting order, do not disturb non-diffusion requests
+        skipped_waiting_requests = create_request_queue(self.policy)
+
+        # Fast path selection and scheduling (treat all as diffusion requests,
+        # independent of pooling_params)
+        while self.waiting and token_budget > 0 and capacity > 0:
+            request = self.waiting.peek_request()
+            # Uniformly treat as diffusion. A feature flag can be added later
+            # via config or request tag.
+
+            # Allocate all input tokens for the request in one shot
+            # (allocate 1 placeholder if zero)
+            required_tokens = max(getattr(request, "num_prompt_tokens", 0), 1)
+            if required_tokens > token_budget:
+                # Insufficient budget to process all inputs at once;
+                # stop fast path attempt
+                break
+            num_new_tokens = required_tokens
+            new_blocks = self.kv_cache_manager.allocate_slots(
+                request,
+                num_new_tokens,
+                num_lookahead_tokens=self.num_lookahead_tokens,
+            )
+            if new_blocks is None:
+                # Allocation failed (e.g., VRAM pressure); stop fast path and
+                # fall back to default scheduling
+                # Put the current request back to the head of the waiting queue
+                # Note: the original queue order is preserved
+                break
+
+            # Officially schedule this request
+            request = self.waiting.pop_request()
+            self.running.append(request)
+            request.status = RequestStatus.RUNNING
+            if self.log_stats:
+                request.record_event(EngineCoreEventType.SCHEDULED, scheduled_timestamp)
+
+            req_to_new_blocks[request.request_id] = new_blocks
+            num_scheduled_tokens[request.request_id] = num_new_tokens
+            token_budget -= num_new_tokens
+            capacity -= 1
+            scheduled_new_reqs.append(request)
+
+        # Return skipped waiting requests
+        if skipped_waiting_requests:
+            self.waiting.prepend_requests(skipped_waiting_requests)
+
+        # If fast path scheduled none, fall back to the original scheduling
+        if not num_scheduled_tokens:
+            return super().schedule()
+
+        # Compute common prefix blocks (aligned with v1)
+        num_common_prefix_blocks = [0] * len(self.kv_cache_config.kv_cache_groups)
+        if self.running:
+            any_request = self.running[0]
+            num_common_prefix_blocks = self.kv_cache_manager.get_num_common_prefix_blocks(
+                any_request, len(self.running)
+            )
+
+        # Assemble SchedulerOutput
+        new_reqs_data = [
+            OmniNewRequestData.from_request(req, req_to_new_blocks[req.request_id].get_block_ids())
+            for req in scheduled_new_reqs
+        ]
+        cached_reqs_data = self._make_cached_request_data(
+            scheduled_running_reqs,
+            scheduled_resumed_reqs,
+            num_scheduled_tokens,
+            scheduled_spec_decode_tokens,
+            req_to_new_blocks,
+        )
+        scheduled_requests = scheduled_new_reqs + scheduled_running_reqs + scheduled_resumed_reqs
+        structured_output_request_ids, grammar_bitmask = self.get_grammar_bitmask(
+            scheduled_requests, scheduled_spec_decode_tokens
+        )
+
+        total_num_scheduled_tokens = sum(num_scheduled_tokens.values())
+        scheduler_output = SchedulerOutput(
+            scheduled_new_reqs=new_reqs_data,
+            scheduled_cached_reqs=cached_reqs_data,
+            num_scheduled_tokens=num_scheduled_tokens,
+            total_num_scheduled_tokens=total_num_scheduled_tokens,
+            scheduled_spec_decode_tokens=scheduled_spec_decode_tokens,
+            scheduled_encoder_inputs=scheduled_encoder_inputs,
+            num_common_prefix_blocks=num_common_prefix_blocks,
+            finished_req_ids=self.finished_req_ids,
+            free_encoder_mm_hashes=self.encoder_cache_manager.get_freed_mm_hashes(),
+            structured_output_request_ids=structured_output_request_ids,
+            grammar_bitmask=grammar_bitmask,
+        )
+
+        # KVTransfer: package metadata
+        if self.connector is not None:
+            meta = self.connector.build_connector_meta(scheduler_output)
+            scheduler_output.kv_connector_metadata = meta
+
+        # Update internal state (advance num_computed_tokens, free encoder inputs,
+        # etc.)
+        self._update_after_schedule(scheduler_output)
+        return scheduler_output
+
+    """
+    Scheduler for the diffusion model.
+    This scheduler is modified to stop the request immediately for the diffusion model.
+    This is because the diffusion model can generate the final image/audio in one step.
+    Note: This is just a minimal modification to the original scheduler,
+    and there should be some further efforts to optimize the scheduler.
+    The original scheduler is still used for the AR model.
+    """
+
+    def update_from_output(
+        self,
+        scheduler_output: SchedulerOutput,
+        model_runner_output: OmniModelRunnerOutput,
+    ) -> dict[int, EngineCoreOutputs]:
+        """Update the scheduler state based on the model runner output.
+
+        This method is modified to stop the request immediately for the diffusion model.
+        """
+        sampled_token_ids = model_runner_output.sampled_token_ids
+        logprobs = model_runner_output.logprobs
+        prompt_logprobs_dict = model_runner_output.prompt_logprobs_dict
+        num_scheduled_tokens = scheduler_output.num_scheduled_tokens
+        pooler_outputs = model_runner_output.pooler_output
+        num_nans_in_logits = model_runner_output.num_nans_in_logits
+        kv_connector_output = model_runner_output.kv_connector_output
+
+        outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list)
+        spec_decoding_stats: Optional[SpecDecodingStats] = None
+        kv_connector_stats = kv_connector_output.kv_connector_stats if kv_connector_output else None
+
+        # NOTE(woosuk): As len(num_scheduled_tokens) can be up to 1K or more,
+        # the below loop can be a performance bottleneck. We should do our best
+        # to avoid expensive operations inside the loop.
+        stopped_running_reqs: set[Request] = set()
+        stopped_preempted_reqs: set[Request] = set()
+        for req_id, num_tokens_scheduled in num_scheduled_tokens.items():
+            assert num_tokens_scheduled > 0
+            request = self.requests.get(req_id)
+            if request is None:
+                # The request is already finished. This can happen if the
+                # request is aborted while the model is executing it (e.g.,
+                # in pipeline parallelism).
+                continue
+
+            req_index = model_runner_output.req_id_to_index[req_id]
+            generated_token_ids = sampled_token_ids[req_index] if sampled_token_ids else []
+
+            scheduled_spec_token_ids = scheduler_output.scheduled_spec_decode_tokens.get(req_id)
+            if scheduled_spec_token_ids:
+                num_draft_tokens = len(scheduled_spec_token_ids)
+                num_accepted = len(generated_token_ids) - 1
+                num_rejected = num_draft_tokens - num_accepted
+                # num_computed_tokens represents the number of tokens
+                # processed in the current step, considering scheduled
+                # tokens and rejections. If some tokens are rejected,
+                # num_computed_tokens is decreased by the number of rejected
+                # tokens.
+                request.num_computed_tokens -= num_rejected
+                spec_decoding_stats = self.make_spec_decoding_stats(
+                    spec_decoding_stats,
+                    num_draft_tokens=num_draft_tokens,
+                    num_accepted_tokens=num_accepted,
+                )
+
+            new_logprobs = None
+            new_token_ids = generated_token_ids
+            kv_transfer_params = None
+            status_before_stop = request.status
+            pooler_output = None
+            if pooler_outputs:
+                pooler_output = pooler_outputs[req_index]
+
+            # Diffusion request: completes in one step; mark finished and free resources
+            request.status = RequestStatus.FINISHED_STOPPED
+            # Optional: set a stop_reason for front-end clarity
+            # (does not affect protocol)
+            request.stop_reason = request.stop_reason  # or "generation_done"
+            kv_transfer_params = self._free_request(request)
+            if status_before_stop == RequestStatus.RUNNING:
+                stopped_running_reqs.add(request)
+            else:
+                stopped_preempted_reqs.add(request)
+
+            # Extract sample logprobs if needed.
+            if request.sampling_params is not None and request.sampling_params.logprobs is not None and logprobs:
+                # NOTE: once we support N tokens per step (spec decode),
+                # the outer lists can be of length > 1.
+                new_logprobs = logprobs.slice(req_index, req_index + 1)
+
+            if new_token_ids and self.structured_output_manager.should_advance(request):
+                # NOTE: structured_output_request should not be None if
+                # use_structured_output, we have check above, so safe to ignore
+                # type warning
+                request.structured_output_request.grammar.accept_tokens(  # type: ignore[union-attr]  # noqa: E501
+                    req_id, new_token_ids
+                )
+
+            # spec_token_ids comes from the model runner output
+            if num_nans_in_logits is not None and req_id in num_nans_in_logits:
+                request.num_nans_in_logits = num_nans_in_logits[req_id]
+
+            # Get prompt logprobs for this request.
+            prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
+            # Convert pooler_output tensor to dict format expected by OmniEngineCoreOutput
+            pooling_output_dict = None
+            if pooler_output is not None:
+                # Wrap tensor in dict to match OmniEngineCoreOutput.pooling_output type
+                # which expects Optional[dict[str, torch.Tensor]]
+                pooling_output_dict = {"model_outputs": pooler_output}
+            if new_token_ids or pooler_output is not None or kv_transfer_params:
+                # Add EngineCoreOutput for this Request.
+                outputs[request.client_index].append(
+                    EngineCoreOutput(
+                        request_id=req_id,
+                        new_token_ids=new_token_ids,
+                        finish_reason=request.get_finished_reason(),
+                        new_logprobs=new_logprobs,
+                        new_prompt_logprobs_tensors=prompt_logprobs_tensors,
+                        pooling_output=pooling_output_dict,
+                        stop_reason=request.stop_reason,
+                        events=request.take_events(),
+                        kv_transfer_params=kv_transfer_params,
+                        trace_headers=request.trace_headers,
+                        num_cached_tokens=request.num_cached_tokens,
+                    )
+                )
+            else:
+                # Invariant: EngineCore returns no partial prefill outputs.
+                assert not prompt_logprobs_tensors
+
+        # Remove the stopped requests from the running and waiting queues.
+        if stopped_running_reqs:
+            if stopped_running_reqs:
+                self.running = remove_all(self.running, stopped_running_reqs)
+        if stopped_preempted_reqs:
+            # This is a rare case and unlikely to impact performance.
+            self.waiting.remove_requests(stopped_preempted_reqs)
+
+        # KV Connector: update state for finished KV Transfers.
+        if model_runner_output.kv_connector_output:
+            self._update_from_kv_xfer_finished(model_runner_output.kv_connector_output)
+
+        # Create EngineCoreOutputs for all clients that have requests with
+        # outputs in this step.
+        engine_core_outputs = {client_index: EngineCoreOutputs(outputs=outs) for client_index, outs in outputs.items()}
+
+        finished_req_ids = self.finished_req_ids_dict
+        if finished_req_ids:
+            # Include ids of requests that finished since last outputs
+            # were sent.
+            for client_index, finished_set in finished_req_ids.items():
+                # Set finished request set in EngineCoreOutputs for this client.
+                if (eco := engine_core_outputs.get(client_index)) is not None:
+                    eco.finished_requests = finished_set
+                else:
+                    engine_core_outputs[client_index] = EngineCoreOutputs(finished_requests=finished_set)
+            finished_req_ids.clear()
+
+        if (stats := self.make_stats(spec_decoding_stats, kv_connector_stats)) is not None:
+            # Return stats to only one of the front-ends.
+            if (eco := next(iter(engine_core_outputs.values()), None)) is None:
+                # We must return the stats even if there are no request
+                # outputs this step.
+                engine_core_outputs[0] = eco = EngineCoreOutputs()
+            eco.scheduler_stats = stats
+
+        return engine_core_outputs

vllm_omni/engine/__init__.py

@@ -2,10 +2,21 @@
 Engine components for vLLM-Omni.
 """
 
-from typing import Any, Optional
+import time
+from collections.abc import Mapping
+from typing import Any, Optional, Union
 
 import msgspec
-from vllm.v1.engine import EngineCoreRequest
+import torch
+from vllm.v1.engine import (
+    EngineCoreEvent,
+    EngineCoreRequest,
+    FinishReason,
+    LogprobsLists,
+    LogprobsTensors,
+    SchedulerStats,
+    UtilityOutput,
+)
 
 
 class PromptEmbedsPayload(msgspec.Struct):
@@ -66,3 +77,62 @@ class OmniEngineCoreRequest(EngineCoreRequest):
     prompt_embeds: Optional[PromptEmbedsPayload] = None
     # Optional additional information dictionary (serialized)
     additional_information: Optional[AdditionalInformationPayload] = None
+
+
+class OmniEngineCoreOutput(
+    msgspec.Struct,
+    array_like=True,  # type: ignore[call-arg]
+    omit_defaults=True,  # type: ignore[call-arg]
+    gc=False,
+):  # type: ignore[call-arg]
+    request_id: str
+    new_token_ids: list[int]
+
+    new_logprobs: Optional[LogprobsLists] = None
+    new_prompt_logprobs_tensors: Optional[LogprobsTensors] = None
+
+    pooling_output: Optional[dict[str, torch.Tensor]] = None
+
+    finish_reason: Optional[FinishReason] = None
+    stop_reason: Union[int, str, None] = None
+    events: Optional[list[EngineCoreEvent]] = None
+    kv_transfer_params: Optional[dict[str, Any]] = None
+
+    trace_headers: Optional[Mapping[str, str]] = None
+    # The number of tokens with prefix cache hits.
+    num_cached_tokens: int = 0
+
+    @property
+    def finished(self) -> bool:
+        return self.finish_reason is not None
+
+
+class OmniEngineCoreOutputs(
+    msgspec.Struct,
+    array_like=True,  # type: ignore[call-arg]
+    omit_defaults=True,  # type: ignore[call-arg]
+    gc=False,
+):  # type: ignore[call-arg]
+    # NOTE(Nick): We could consider ways to make this more compact,
+    # e.g. columnwise layout
+
+    engine_index: int = 0
+
+    # [num_reqs]
+    outputs: list[OmniEngineCoreOutput] = []
+    scheduler_stats: Optional[SchedulerStats] = None
+    timestamp: float = 0.0
+
+    utility_output: Optional[UtilityOutput] = None
+    finished_requests: Optional[set[str]] = None
+
+    # In DP case, used to signal that the current wave of requests
+    # has finished and the engines are paused.
+    wave_complete: Optional[int] = None
+    # In DP case, used to signal that a request was received for an
+    # "old" wave, so the next wave needs to be started in other engines.
+    start_wave: Optional[int] = None
+
+    def __post_init__(self):
+        if self.timestamp == 0.0:
+            self.timestamp = time.monotonic()

vllm_omni/engine/arg_utils.py

@@ -1,19 +1,21 @@
 from dataclasses import dataclass
 from typing import Optional
 
+from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import Qwen3OmniMoeTextConfig
 from vllm.engine.arg_utils import EngineArgs
+from vllm.logger import init_logger
 from vllm.v1.engine.async_llm import AsyncEngineArgs
 
 from vllm_omni.config import OmniModelConfig
 
+logger = init_logger(__name__)
+
 
 @dataclass
 class OmniEngineArgs(EngineArgs):
     """Engine arguments for omni models, extending base EngineArgs.
-
     Adds omni-specific configuration fields for multi-stage pipeline
     processing and output type specification.
-
     Args:
         stage_id: Identifier for the stage in a multi-stage pipeline (default: 0)
         model_stage: Stage type identifier, e.g., "thinker" or "talker"
@@ -29,13 +31,20 @@ class OmniEngineArgs(EngineArgs):
     model_stage: str = "thinker"
     model_arch: str = "Qwen2_5OmniForConditionalGeneration"
     engine_output_type: Optional[str] = None
+    hf_config_name: Optional[str] = None
+
+    def draw_hf_text_config(self, config_dict: dict) -> Qwen3OmniMoeTextConfig:
+        # transformers' get_text_config method is used to get the text config from thinker_config.
+        # to handle the case that each model stage has their own text config,
+        # we need to draw the text config from the corresponding model stage.
+        return getattr(config_dict["hf_config"], config_dict["hf_config_name"]).get_text_config()
 
     def create_model_config(self) -> OmniModelConfig:
         """Create an OmniModelConfig from these engine arguments.
-
         Returns:
             OmniModelConfig instance with all configuration fields set
         """
+
         # First, get the base ModelConfig from the parent class
         base_config = super().create_model_config()
 
@@ -48,7 +57,9 @@ class OmniEngineArgs(EngineArgs):
         config_dict["model_stage"] = self.model_stage
         config_dict["model_arch"] = self.model_arch
         config_dict["engine_output_type"] = self.engine_output_type
-
+        config_dict["hf_config_name"] = self.hf_config_name
+        if self.hf_config_name is not None:
+            config_dict["hf_text_config"] = self.draw_hf_text_config(config_dict)
         # Create and return the OmniModelConfig instance
         omni_config = OmniModelConfig(**config_dict)
         omni_config.hf_config.architectures = omni_config.architectures
@@ -59,10 +70,8 @@ class OmniEngineArgs(EngineArgs):
 @dataclass
 class AsyncOmniEngineArgs(AsyncEngineArgs):
     """Async engine arguments for omni models, extending base AsyncEngineArgs.
-
     Adds omni-specific configuration fields for multi-stage pipeline
     processing and output type specification in async contexts.
-
     Args:
         stage_id: Identifier for the stage in a multi-stage pipeline (default: 0)
         model_stage: Stage type identifier, e.g., "thinker" or "talker"
@@ -78,13 +87,15 @@ class AsyncOmniEngineArgs(AsyncEngineArgs):
     model_stage: str = "thinker"
     model_arch: str = "Qwen2_5OmniForConditionalGeneration"
     engine_output_type: Optional[str] = None
+    hf_config_name: Optional[str] = None
+
+    def draw_hf_text_config(self, config_dict: dict) -> Qwen3OmniMoeTextConfig:
+        # transformers' get_text_config method is used to get the text config from thinker_config.
+        # to handle the case that each model stage has their own text config,
+        # we need to draw the text config from the corresponding model stage.
+        return getattr(config_dict["hf_config"], config_dict["hf_config_name"]).get_text_config()
 
     def create_model_config(self) -> OmniModelConfig:
-        """Create an OmniModelConfig from these engine arguments.
-
-        Returns:
-            OmniModelConfig instance with all configuration fields set
-        """
         # First, get the base ModelConfig from the parent class
         base_config = super().create_model_config()
 
@@ -98,6 +109,9 @@ class AsyncOmniEngineArgs(AsyncEngineArgs):
         config_dict["model_arch"] = self.model_arch
         config_dict["engine_output_type"] = self.engine_output_type
 
+        config_dict["hf_config_name"] = self.hf_config_name
+        if self.hf_config_name is not None:
+            config_dict["hf_text_config"] = self.draw_hf_text_config(config_dict)
         # Create and return the OmniModelConfig instance
         omni_config = OmniModelConfig(**config_dict)
         omni_config.hf_config.architectures = omni_config.architectures

vllm_omni/engine/output_processor.py

@@ -1,3 +1,4 @@
+from ast import Dict
 from typing import Any, Callable, Optional, Union
 
 import torch
@@ -33,7 +34,7 @@ class OmniRequestState(RequestState):
     ):
         super().__init__(*args, **kwargs)
         self.mm_type: Optional[str] = None
-        self.mm_accumulated: Optional[torch.Tensor] = None
+        self.mm_accumulated: Optional[Dict[str, Any]] = None
 
     @classmethod
     def from_new_request(
@@ -92,33 +93,71 @@ class OmniRequestState(RequestState):
             log_stats=log_stats,
         )
 
-    def add_multimodal_tensor(self, tensor: Optional[torch.Tensor], mm_type: Optional[str]) -> None:
-        """Add a multimodal tensor to the accumulated outputs.
-
-        Accumulates multimodal tensors (e.g., image latents, audio waveforms)
-        that are produced incrementally during generation. Tensors are
-        concatenated along the first dimension.
-
-        Args:
-            tensor: Optional tensor to add. If None, this is a no-op.
-            mm_type: Optional modality type identifier (e.g., "image", "audio").
-                Used to categorize the output type.
-        """
-        if tensor is None:
+    def add_multimodal_tensor(self, payload: Optional[Any], mm_type: Optional[str]) -> None:
+        if payload is None:
             return
         try:
             if mm_type:
                 self.mm_type = (mm_type or "").lower()
-            t = tensor.detach()
-            try:
-                t = t.to("cpu")
-            except Exception:
-                # Best-effort CPU move; keep original device if conversion fails
-                logger.debug("Failed to move multimodal tensor to CPU", exc_info=True)
-            if self.mm_accumulated is None:
-                self.mm_accumulated = t
+
+            # Normalize incoming payload to dict on CPU
+            def _to_cpu(x):
+                if isinstance(x, torch.Tensor):
+                    try:
+                        return x.detach().to("cpu").contiguous()
+                    except Exception:
+                        return x
+                return x
+
+            if isinstance(payload, dict):
+                incoming: Dict[str, Any] = {}
+                # Optional remap: if producer used "model_outputs" or "hidden", rename to mm_type
+                # to keep a consistent key namespace per engine_core_output_type.
+                remapped_tensor = dict(payload)
+                target_key = self.mm_type or "hidden"
+                if "model_outputs" in remapped_tensor:
+                    remapped_tensor[target_key] = remapped_tensor.pop("model_outputs")
+                elif "hidden" in remapped_tensor and target_key != "hidden":
+                    remapped_tensor[target_key] = remapped_tensor.pop("hidden")
+                for k, v in remapped_tensor.items():
+                    if isinstance(v, dict):
+                        incoming[k] = {str(sk): _to_cpu(sv) for sk, sv in v.items()}
+                    else:
+                        incoming[k] = _to_cpu(v)
             else:
-                self.mm_accumulated = torch.cat([self.mm_accumulated, t], dim=0)
+                key = self.mm_type or "hidden"
+                incoming = {key: _to_cpu(payload)}
+
+            if self.mm_accumulated is None:
+                self.mm_accumulated = incoming
+            else:
+                # Merge keys; concatenate tensors along token dim when possible
+                for k, v in incoming.items():
+                    if k not in self.mm_accumulated:
+                        self.mm_accumulated[k] = v
+                    else:
+                        existing = self.mm_accumulated[k]
+                        if isinstance(v, torch.Tensor) and isinstance(existing, torch.Tensor):
+                            try:
+                                self.mm_accumulated[k] = torch.cat([existing, v], dim=0)  # type: ignore[index]
+                            except Exception:
+                                self.mm_accumulated[k] = v
+                        elif isinstance(v, dict) and isinstance(existing, dict):
+                            # Merge nested dicts by concatenating tensors along token dim when possible
+                            for sk, sv in v.items():
+                                if sk not in existing:
+                                    existing[sk] = sv
+                                    continue
+                                ev = existing[sk]
+                                if isinstance(sv, torch.Tensor) and isinstance(ev, torch.Tensor):
+                                    try:
+                                        existing[sk] = torch.cat([ev, sv], dim=0)
+                                    except Exception:
+                                        existing[sk] = sv
+                                else:
+                                    existing[sk] = sv
+                        else:
+                            self.mm_accumulated[k] = v
         except Exception:
             # Log and continue without crashing the output pipeline
             logger.exception("Error accumulating multimodal tensor")
@@ -178,18 +217,15 @@ class OmniRequestState(RequestState):
         base_output = super()._new_completion_output(token_ids, finish_reason, stop_reason)
         try:
             if self.mm_accumulated is not None:
-                tensor = self.mm_accumulated
-                try:
-                    tensor = tensor.detach().to("cpu")
-                except Exception:
-                    logger.debug(
-                        "Failed to move accumulated multimodal tensor to CPU",
-                        exc_info=True,
-                    )
-                # Attach on the completion output for downstream consumers.
+                # Attach accumulated multimodal dict on the completion output
                 if not hasattr(base_output, "multimodal_output"):
                     setattr(base_output, "multimodal_output", {})
-                setattr(base_output, "multimodal_output", {self.mm_type: tensor})
+                mm_out = getattr(base_output, "multimodal_output")
+                if isinstance(mm_out, dict):
+                    for k, v in self.mm_accumulated.items():
+                        mm_out[k] = v
+                else:
+                    setattr(base_output, "multimodal_output", self.mm_accumulated)
         except Exception:
             logger.exception("Error in _new_completion_output")
         return base_output
@@ -369,10 +405,9 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
                 # Attach accumulated multimodal payload if any
                 try:
                     if isinstance(req_state, OmniRequestState) and req_state.mm_accumulated is not None:
-                        mm_key = req_state.mm_type or "latents"
                         if not hasattr(ro, "multimodal_output"):
                             setattr(ro, "multimodal_output", {})
-                        ro.multimodal_output[mm_key] = req_state.mm_accumulated
+                        ro.multimodal_output = req_state.mm_accumulated
                 except Exception:
                     logger.exception("Error attaching multimodal payload in process_outputs")
                 if req_state.queue is not None:

vllm_omni/entrypoints/async_omni_llm.py

@@ -2,7 +2,6 @@ import asyncio
 import multiprocessing as mp
 import os
 import socket
-import sys
 import time
 from collections.abc import AsyncGenerator, Iterable, Mapping, Sequence
 from concurrent.futures import ThreadPoolExecutor, as_completed
@@ -94,8 +93,8 @@ class AsyncOmniLLM(EngineClient):
         init_sleep_seconds: int = 20,
         shm_threshold_bytes: int = 65536,
         batch_timeout: int = 10,
-        init_timeout: int = 300,
-        **kwargs: Any,
+        init_timeout: int = 60000,
+        **kwargs,
     ):
         self.batch_timeout = batch_timeout
         self._enable_stats: bool = bool(log_stats)
@@ -194,6 +193,8 @@ class AsyncOmniLLM(EngineClient):
                 logger.warning("[Orchestrator] Failed to stop stage worker: %s", e)
 
     def __del__(self) -> None:  # best-effort
+        print("[AsyncOmniLLM] __del__ close()", flush=True)
+        raise Exception("test")
         try:
             self.close()
         except Exception as e:
@@ -460,9 +461,8 @@ class AsyncOmniLLM(EngineClient):
             logger.exception("[Orchestrator] Failed to build/log summary: %s", e)
 
     def _wait_for_stages_ready(self, timeout: int = 120) -> None:
-        deadline = time.time() + max(0, int(timeout))
         num_stages = len(self.stage_list)
-        while len(self._stages_ready) < num_stages and time.time() < deadline:
+        while len(self._stages_ready) < num_stages:
             progressed = False
             for stage_id, stage in enumerate(self.stage_list):
                 if stage_id in self._stages_ready:
@@ -528,20 +528,6 @@ class AsyncOmniLLM(EngineClient):
                         occurred while logging suggestions",
                 )
 
-            # Attempt graceful shutdown of all stages before exiting
-            try:
-                self.close()
-            except Exception:
-                pass
-
-            # Terminate the current process with non-zero exit code
-            try:
-                sys.exit(1)
-            except SystemExit:
-                raise
-            except Exception:
-                os._exit(1)
-
     @property
     def is_running(self) -> bool:
         # Is None before the loop is started.

vllm_omni/entrypoints/omni_llm.py

@@ -1,6 +1,5 @@
 import multiprocessing as mp
 import os
-import sys
 import time
 from collections.abc import Sequence
 from concurrent.futures import ThreadPoolExecutor, as_completed
@@ -11,7 +10,6 @@ from pydantic import ValidationError
 
 # External library imports (vLLM)
 from vllm.config import CompilationConfig, StructuredOutputsConfig, is_init_field
-from vllm.engine.arg_utils import HfOverrides
 from vllm.entrypoints.llm import LLM
 from vllm.inputs import PromptType
 from vllm.logger import init_logger
@@ -96,12 +94,12 @@ class OmniLLM:
 
         # Optional file handler for orchestrator
         self._log_file = log_file
-        if self._log_file:
-            remove_old_logs(self._log_file, len(self.stage_configs))
-            configure_orchestrator_logger(logger, self._log_file)
 
         self._stats_file, self._overall_stats_file = init_stats_paths(self._enable_stats, self._log_file)
         self._initialize_stages(model, init_sleep_seconds, shm_threshold_bytes, init_timeout)
+        if self._log_file:
+            remove_old_logs(self._log_file, len(self.stage_list))
+            configure_orchestrator_logger(logger, self._log_file)
 
     def _initialize_stages(
         self,
@@ -378,7 +376,16 @@ class OmniLLM:
                 next_stage_id = stage_id + 1
                 if next_stage_id < num_stages:
                     next_stage: OmniStage = self.stage_list[next_stage_id]
-                    next_inputs = next_stage.process_engine_inputs(self.stage_list, [request_id_to_prompt[req_id]])
+                    try:
+                        next_inputs = next_stage.process_engine_inputs(self.stage_list, [request_id_to_prompt[req_id]])
+                    except Exception as e:
+                        logger.exception(
+                            "[Orchestrator] Process engine inputs error for req %s at stage %s: %s",
+                            req_id,
+                            next_stage_id,
+                            e,
+                        )
+                        continue
                     sp_next: SamplingParams = sampling_params_list[next_stage_id]  # type: ignore[index]
                     try:
                         # Measure transfer size and time (encode + enqueue)
@@ -467,7 +474,7 @@ class OmniLLM:
                 progressed = True
                 if result.get("type") == "stage_ready":
                     self._stages_ready.add(stage_id)
-                    logger.debug("[Orchestrator] Stage-%s reported ready", stage_id)
+                    logger.info("[Orchestrator] Stage-%s reported ready", stage_id)
                 else:
                     # No user data should arrive before seeding; ignore other messages
                     pass
@@ -500,20 +507,8 @@ class OmniLLM:
                 logger.error(
                     "[Orchestrator] Stage initialization failed and an error occurred while logging suggestions",
                 )
-
-            # Attempt graceful shutdown of all stages before exiting
-            try:
-                self.close()
-            except Exception:
-                pass
-
-            # Terminate the current process with non-zero exit code
-            try:
-                sys.exit(1)
-            except SystemExit:
-                raise
-            except Exception:
-                os._exit(1)
+        elif len(self._stages_ready) == num_stages:
+            logger.info("[Orchestrator] All stages initialized successfully")
 
 
 class OmniStageLLM(LLM):
@@ -539,9 +534,8 @@ class OmniStageLLM(LLM):
         self,
         model: str,
         compilation_config: Optional[Union[int, dict[str, Any], CompilationConfig]] = None,
-        hf_overrides: Optional[HfOverrides] = None,
         structured_outputs_config: Optional[Union[dict[str, Any], StructuredOutputsConfig]] = None,
-        **kwargs: Any,
+        **kwargs,
     ):
         """LLM constructor."""
         if "disable_log_stats" not in kwargs:
@@ -570,9 +564,6 @@ class OmniStageLLM(LLM):
                 # to provide better context to the user.
                 raise ValueError(f"Invalid 'kv_transfer_config' provided: {e}") from e
 
-        if hf_overrides is None:
-            hf_overrides = {}
-
         if compilation_config is not None:
             if isinstance(compilation_config, int):
                 compilation_config_instance = CompilationConfig(level=compilation_config)
@@ -597,7 +588,6 @@ class OmniStageLLM(LLM):
 
         engine_args = OmniEngineArgs(
             model=model,
-            hf_overrides=hf_overrides,
             compilation_config=compilation_config_instance,
             structured_outputs_config=structured_outputs_instance,
             **kwargs,

vllm_omni/entrypoints/omni_stage.py

@@ -14,6 +14,7 @@ import asyncio
 import importlib
 import logging
 import multiprocessing as mp
+import os
 from typing import Any, Optional, Union
 
 from vllm.inputs import TextPrompt
@@ -237,6 +238,14 @@ class OmniStage:
                     self._proc.terminate()
                 except Exception as e:
                     self._logger.warning("[Stage-%s] terminate() failed: %s", self.stage_id, e)
+        # Cleanup temporary stage log if we created one (only when no log_file provided)
+        try:
+            if not self._log_file:
+                tmp_path = f"/tmp/omni_stage{self.stage_id}.log"
+                if os.path.exists(tmp_path):
+                    os.remove(tmp_path)
+        except Exception:
+            pass
 
     def submit(self, payload: dict[str, Any]) -> None:
         """Submit a task to the stage worker.
@@ -286,8 +295,6 @@ class OmniStage:
                 raise ValueError("engine_input_source is empty")
             source_stage_id = self.engine_input_source[0]
             source_outputs = stage_list[source_stage_id].engine_outputs
-            if not isinstance(prompt, list):
-                prompt = [prompt]
             multi_modal_data = {
                 source_output.request_id: p.get("multi_modal_data", None)
                 for source_output, p in zip(source_outputs, prompt)
@@ -322,6 +329,7 @@ def _stage_worker(
 ) -> None:
     """Stage worker entry: device setup, LLM init, batching, SHM IPC."""
     import logging as _logging
+    import os as _os
     import time as _time
 
     from vllm_omni.entrypoints.log_utils import (
@@ -331,7 +339,6 @@ def _stage_worker(
         log_stage_running_avg,
     )
     from vllm_omni.entrypoints.omni_llm import OmniStageLLM
-
     # no inline JSONL/serialization imports; logging handled by utilities
 
     stage_id = stage_payload["stage_id"]
@@ -339,12 +346,31 @@ def _stage_worker(
     runtime_cfg = stage_payload.get("runtime", {})
     shm_threshold_bytes = int(stage_payload.get("shm_threshold_bytes", 65536))
 
-    # Per-stage file logger (optional)
+    # Per-stage logger: clear inherited handlers to avoid broken parent streams
     try:
-        if log_file:
-            stage_log = _logging.getLogger(__name__)
-            stage_log.setLevel(_logging.DEBUG)
-            fh = _logging.FileHandler(f"{log_file}.stage{stage_id}.log")
+        stage_log = _logging.getLogger(__name__)
+        stage_log.setLevel(_logging.DEBUG)
+        for _h in list(stage_log.handlers):
+            try:
+                stage_log.removeHandler(_h)
+            except Exception:
+                pass
+        stage_log.propagate = False
+
+        class _StageFilter(_logging.Filter):
+            def filter(self, record: _logging.LogRecord) -> bool:
+                setattr(record, "stage", stage_id)
+                return True
+
+        # Prefer file logging even when log_file is not provided, to avoid invalid stdio in child procs
+        try:
+            if log_file:
+                _path = f"{log_file}.stage{stage_id}.log"
+            else:
+                _path = f"/tmp/omni_stage{stage_id}.log"
+            # Ensure dir exists
+            _os.makedirs(_os.path.dirname(_path), exist_ok=True)
+            fh = _logging.FileHandler(_path)
             fh.setLevel(_logging.DEBUG)
             fh.setFormatter(
                 _logging.Formatter("%(asctime)s [PID:%(process)d] [Stage-%(stage)s] %(levelname)s: %(message)s")
@@ -357,6 +383,34 @@ def _stage_worker(
 
             fh.addFilter(_StageFilter())
             stage_log.addHandler(fh)
+            try:
+                # Also route vLLM internal logs to the same handler for this stage
+                _vllm_logger = _logging.getLogger("vllm")
+                _vllm_logger.setLevel(_logging.DEBUG)
+                for _vh in list(_vllm_logger.handlers):
+                    try:
+                        _vllm_logger.removeHandler(_vh)
+                    except Exception:
+                        pass
+                _vllm_logger.propagate = False
+                _vllm_logger.addHandler(fh)
+            except Exception:
+                pass
+        except Exception:
+            # Final fallback: attach NullHandler to avoid logging errors
+            stage_log.addHandler(_logging.NullHandler())
+    except Exception:
+        pass
+
+    try:
+        _logging.raiseExceptions = False
+        _root_logger = _logging.getLogger()
+        for _h in list(_root_logger.handlers):
+            try:
+                _root_logger.removeHandler(_h)
+            except Exception:
+                pass
+        _root_logger.addHandler(_logging.NullHandler())
     except Exception:
         pass
 
@@ -377,9 +431,7 @@ def _stage_worker(
 
     # Init LLM
     _logging.getLogger(__name__).debug(
-        "[Stage-%s] Initializing engine with args keys=%s",
-        stage_id,
-        list(engine_args.keys()),
+        "[Stage-%s] Initializing engine with args keys=%s", stage_id, list(engine_args.keys())
     )
     stage_engine = OmniStageLLM(model=model, **engine_args)
     _logging.getLogger(__name__).debug("[Stage-%s] Engine initialized", stage_id)
@@ -394,21 +446,35 @@ def _stage_worker(
         task = in_q.get()
         _recv_dequeue_ts = _time.time()
         if task is None:
-            _logging.getLogger(__name__).debug("[Stage-%s] Received shutdown signal", stage_id)
+            _logging.getLogger(__name__).error("[Stage-%s] Received shutdown signal", stage_id)
             break
 
         max_batch_size = int(runtime_cfg.get("max_batch_size", 1) or 1)
+        print(f"[Stage-{stage_id}] Max batch size: {max_batch_size}")
         batch_tasks: list[dict[str, Any]] = [task]
         if max_batch_size > 1:
             while len(batch_tasks) < max_batch_size:
+                start_time = _time.time()
                 if not in_q.empty():
-                    extra = in_q.get(timeout=batch_timeout)
+                    extra = in_q.get_nowait()
                     if extra is None:
                         in_q.put(None)
                         break
                     batch_tasks.append(extra)
+                    end_time = _time.time()
+                    duration = end_time - start_time
+                    if duration > batch_timeout:
+                        break
+                    else:
+                        continue
                 else:
-                    break
+                    end_time = _time.time()
+                    duration = end_time - start_time
+                    _time.sleep(0.05)
+                    if duration > batch_timeout:
+                        break
+                    else:
+                        continue
 
         batch_request_ids: list[Any] = []
         batch_engine_inputs: list[Any] = []
@@ -434,7 +500,7 @@ def _stage_worker(
             elif isinstance(ein, dict):
                 batch_engine_inputs.append(ein)
             else:
-                _logging.getLogger(__name__).exception("[Stage-%s] Invalid engine input type: %s", stage_id, type(ein))
+                _logging.getLogger(__name__).error("[Stage-%s] Invalid engine input type: %s", stage_id, type(ein))
         sampling_params = batch_tasks[0]["sampling_params"]
         _logging.getLogger(__name__).debug(
             "[Stage-%s] Received batch size=%d, request_ids=%s",
@@ -456,6 +522,14 @@ def _stage_worker(
                 gen_outputs.append(ro)
             _gen_t1 = _time.time()
             _gen_ms = (_gen_t1 - _gen_t0) * 1000.0
+            try:
+                print(
+                    f"[Stage-{stage_id}] Generate done: batch={len(batch_tasks)}, "
+                    f"req_ids={batch_request_ids}, gen_ms={_gen_ms:.1f}",
+                    flush=True,
+                )
+            except Exception:
+                pass
 
             # Group outputs per request id with fallback
             req_to_outputs: dict[Any, list[Any]] = {rid: [] for rid in batch_request_ids}
@@ -783,6 +857,14 @@ async def _stage_worker_async(
                             "metrics": _metrics,
                         }
                     )
+                    try:
+                        print(
+                            f"[Stage-{stage_id}] Enqueued req={rid}, use_shm={use_shm}, "
+                            f"tokens_out={_metrics['num_tokens_out']}",
+                            flush=True,
+                        )
+                    except Exception:
+                        pass
             except Exception as e:
                 _logging.getLogger(__name__).exception(
                     "[Stage-%s] Failed to enqueue result for request %s: %s",
@@ -815,3 +897,6 @@ async def _stage_worker_async(
                     "error": str(e),
                 }
             )
+    print("--------------------------------", flush=True)
+    print(f"[Stage-{stage_id}] Stage worker exiting", flush=True)
+    print("--------------------------------", flush=True)

vllm_omni/entrypoints/openai/serving_chat.py

@@ -730,7 +730,7 @@ class OmniOpenAIServingChat(OpenAIServingChat):
     def _create_audio_choice(self, omni_outputs: OmniRequestOutput, role: str):
         choices: list[ChatCompletionResponseChoice] = []
         final_res = omni_outputs.request_output
-        audio_tensor = final_res.multimodal_output["audio"].detach().cpu().numpy()
+        audio_tensor = final_res.multimodal_output["audio"].float().detach().cpu().numpy()
 
         # Convert numpy array to WAV bytes and encode as base64
         if soundfile is None:

vllm_omni/entrypoints/stage_utils.py

@@ -53,11 +53,7 @@ def set_stage_gpu_devices(stage_id: int, devices: str | int | None) -> None:
                     if 0 <= logical_idx < len(mapping):
                         selected_physical = mapping[logical_idx]
                 except Exception as e:
-                    logger.debug(
-                        "[Stage-%s] Failed to map logical index via CUDA_VISIBLE_DEVICES: %s",
-                        stage_id,
-                        e,
-                    )
+                    logger.debug("[Stage-%s] Failed to map logical index via CUDA_VISIBLE_DEVICES: %s", stage_id, e)
                     selected_physical = None
             if selected_physical is None:
                 selected_physical = int(logical_idx)
@@ -69,18 +65,12 @@ def set_stage_gpu_devices(stage_id: int, devices: str | int | None) -> None:
                 selected_physical,
             )
         elif devices in (None, "cpu"):
-            logger.debug(
-                "[Stage-%s] Using default device visibility (devices=%s)",
-                stage_id,
-                devices,
-            )
+            logger.debug("[Stage-%s] Using default device visibility (devices=%s)", stage_id, devices)
         else:
             selected_physical = int(str(devices))
             os.environ["CUDA_VISIBLE_DEVICES"] = str(selected_physical)
             logger.debug(
-                "[Stage-%s] Set CUDA_VISIBLE_DEVICES to single device %s (fallback)",
-                stage_id,
-                selected_physical,
+                "[Stage-%s] Set CUDA_VISIBLE_DEVICES to single device %s (fallback)", stage_id, selected_physical
             )
 
         try:

vllm_omni/model_executor/layers/mrope.py

@@ -4,8 +4,11 @@ from typing import Optional, Union
 import numpy as np
 import torch
 from transformers import PretrainedConfig
+from vllm.logger import init_logger
 from vllm.model_executor.layers.rotary_embedding.base import RotaryEmbedding
 
+logger = init_logger(__name__)
+
 
 def _apply_rotary_emb(
     x: torch.Tensor,
@@ -442,9 +445,16 @@ class MRotaryEmbedding(RotaryEmbedding):
         #  _vl_get_input_positions_tensor.
 
         thinker_config = hf_config.thinker_config
-        audio_token_id = thinker_config.audio_token_index
-        image_token_id = thinker_config.image_token_index
-        video_token_id = thinker_config.video_token_index
+        try:
+            audio_token_id = thinker_config.audio_token_index
+            image_token_id = thinker_config.image_token_index
+            video_token_id = thinker_config.video_token_index
+        except Exception:
+            logger.info("Multimodal token idx changed!")
+            audio_token_id = thinker_config.audio_token_id
+            image_token_id = thinker_config.image_token_id
+            video_token_id = thinker_config.video_token_id
+
         audio_start_token_id = thinker_config.audio_start_token_id
         audio_end_token_id = thinker_config.audio_end_token_id
         vision_start_token_id = thinker_config.vision_start_token_id

vllm_omni/model_executor/models/__init__.py

@@ -1 +1,4 @@
+from .qwen3_omni import Qwen3OmniMoeForConditionalGeneration
 from .registry import OmniModelRegistry  # noqa: F401
+
+__all__ = ["Qwen3OmniMoeForConditionalGeneration"]

vllm_omni/model_executor/models/output_templates.py

@@ -0,0 +1,13 @@
+from typing import NamedTuple, Optional
+
+import torch
+from vllm.sequence import IntermediateTensors
+
+
+class OmniOutput(NamedTuple):
+    """Output from the merged Omni model containing both text and audio."""
+
+    text_hidden_states: torch.Tensor
+    multimodal_outputs: Optional[dict] = None
+    intermediate_tensors: Optional[IntermediateTensors] = None
+    next_token_id: Optional[torch.Tensor] = None

vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py

@@ -2,7 +2,7 @@ import glob
 import os
 from collections.abc import Iterable
 from functools import cached_property
-from typing import NamedTuple, Optional, Union
+from typing import Optional, Union
 
 import numpy as np
 import torch
@@ -26,6 +26,7 @@ from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.sample.sampler import Sampler
 
 from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
+from vllm_omni.model_executor.models.output_templates import OmniOutput
 from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import (
     Qwen2_5OmniConditionalGenerationMixin,
     Qwen2_5OmniThinkerDummyInputsBuilder,
@@ -39,14 +40,6 @@ TALKER_CODEC_EOS_TOKEN_ID = 8294
 TALKER_CODEC_BOS_TOKEN_ID = 8293
 
 
-class OmniOutput(NamedTuple):
-    """Output from the merged Omni model containing both text and audio."""
-
-    text_hidden_states: torch.Tensor
-    multimodal_outputs: dict = {}
-    intermediate_tensors: Optional[IntermediateTensors] = None
-
-
 logger = init_logger(__name__)
 
 
@@ -439,7 +432,7 @@ class Qwen2_5OmniForConditionalGeneration(
             code = code[1:] if code[0] == TALKER_CODEC_BOS_TOKEN_ID else code
 
             audio_tensor = self.generate_audio(code, voice_type)
-            return OmniOutput(text_hidden_states=None, multimodal_outputs={"audio": audio_tensor})
+            return OmniOutput(text_hidden_states=None, multimodal_outputs={"model_outputs": audio_tensor})
 
         return OmniOutput(
             text_hidden_states=torch.cat(

vllm_omni/model_executor/models/qwen3_omni/__init__.py

@@ -0,0 +1,3 @@
+from .qwen3_omni import Qwen3OmniMoeForConditionalGeneration
+
+__all__ = ["Qwen3OmniMoeForConditionalGeneration"]

vllm_omni/model_executor/models/qwen3_omni/qwen3_moe.py

@@ -0,0 +1,173 @@
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_moe import FusedMoE
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
+from vllm.model_executor.models.qwen3_moe import (
+    Qwen3MoeDecoderLayer,
+    Qwen3MoeMLP,
+    Qwen3MoeModel,  # as _BaseQwen3MoeModel,
+)
+from vllm.model_executor.models.qwen3_moe import (
+    Qwen3MoeForCausalLM as _BaseQwen3MoeForCausalLM,
+)
+from vllm.model_executor.models.utils import (
+    PPMissingLayer,
+    maybe_prefix,
+)
+
+logger = init_logger(__name__)
+
+
+# Individual expert MoE block using Qwen3MoeMLP instead of FusedMoE
+class Qwen3OmniMoeSparseMoeBlock(nn.Module):
+    """Sparse MoE block using individual Qwen3MoeMLP experts instead of FusedMoE."""
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        prefix: str = "",
+    ):
+        super().__init__()
+
+        config = vllm_config.model_config.hf_text_config
+        quant_config = vllm_config.quant_config
+
+        self.num_experts = config.num_experts
+        self.top_k = config.num_experts_per_tok
+        self.norm_topk_prob = config.norm_topk_prob
+        self.hidden_size = config.hidden_size
+
+        # Create individual expert MLPs
+        self.experts = nn.ModuleList(
+            [
+                Qwen3MoeMLP(
+                    hidden_size=config.hidden_size,
+                    intermediate_size=config.moe_intermediate_size,
+                    hidden_act=config.hidden_act,
+                    quant_config=quant_config,
+                    reduce_results=False,
+                    prefix=f"{prefix}.experts.{i}",
+                )
+                for i in range(self.num_experts)
+            ]
+        )
+
+        # Router for expert selection
+        from vllm.model_executor.layers.linear import ReplicatedLinear
+
+        self.gate = ReplicatedLinear(
+            config.hidden_size, config.num_experts, bias=False, quant_config=quant_config, prefix=f"{prefix}.gate"
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """Forward pass using individual experts."""
+        # Handle 3D inputs (batch, seq_len, hidden_size) by reshaping to 2D
+        orig_shape = hidden_states.shape
+        if hidden_states.dim() == 3:
+            batch_size, seq_len, hidden_dim = hidden_states.shape
+            hidden_states = hidden_states.view(-1, hidden_dim)
+        elif hidden_states.dim() == 2:
+            num_tokens, hidden_dim = hidden_states.shape
+        elif hidden_states.dim() == 1:
+            hidden_states = hidden_states.unsqueeze(0)
+            num_tokens, hidden_dim = hidden_states.shape
+        else:
+            raise ValueError(
+                f"Qwen3OmniMoeSparseMoeBlock only supports 1D, 2D, or 3D inputs, got {hidden_states.dim()}D"
+            )
+
+        is_input_1d = len(orig_shape) == 1
+        hidden_states = hidden_states.view(-1, hidden_dim)
+
+        # Get router logits and select experts (matching transformers)
+        router_logits, _ = self.gate(hidden_states)
+        selected_experts, routing_weights = self._route_tokens(router_logits)
+
+        # Forward through individual experts
+        final_hidden_states = self._forward_experts(hidden_states, selected_experts, routing_weights)
+
+        # Reshape back to original shape
+        if is_input_1d:
+            return final_hidden_states.squeeze(0)
+        elif len(orig_shape) == 3:
+            # Reshape back to 3D (batch, seq_len, hidden_dim)
+            return final_hidden_states.view(orig_shape)
+        else:
+            return final_hidden_states
+
+    def _route_tokens(self, router_logits: torch.Tensor):
+        """Route tokens to experts using top-k selection (matching transformers)."""
+        routing_weights = F.softmax(router_logits, dim=-1, dtype=torch.float)
+        routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1)
+        if self.norm_topk_prob:
+            routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
+        routing_weights = routing_weights.to(router_logits.dtype)
+        return selected_experts, routing_weights
+
+    def _forward_experts(
+        self, hidden_states: torch.Tensor, selected_experts: torch.Tensor, routing_weights: torch.Tensor
+    ):
+        """Forward through individual experts (matching transformers implementation)."""
+        final_hidden_states = torch.zeros_like(hidden_states)
+        expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0)
+
+        expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()
+        for expert_idx in expert_hit:
+            idx, top_x = torch.where(expert_mask[expert_idx].squeeze(0))
+            current_state = hidden_states[None, top_x].reshape(-1, hidden_states.shape[-1])
+            current_hidden_states = self.experts[expert_idx](current_state) * routing_weights[top_x, idx, None]
+            final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype))
+
+        return final_hidden_states
+
+
+class Qwen3MoeForCausalLM(_BaseQwen3MoeForCausalLM):
+    """Thin wrapper to swap in the patched `Qwen3MoeModel`."""
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        # Don't call super().__init__() to avoid duplicate layer registration.
+        nn.Module.__init__(self)
+        config = vllm_config.model_config.hf_text_config
+        quant_config = vllm_config.quant_config
+        self.config = config
+        self.quant_config = quant_config
+        self.model = Qwen3MoeModel(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model"))
+        self.lm_head = ParallelLMHead(
+            config.vocab_size, config.hidden_size, quant_config=quant_config, prefix=maybe_prefix(prefix, "lm_head")
+        )
+        if self.config.tie_word_embeddings:
+            self.lm_head.weight = self.model.embed_tokens.weight
+        self.logits_processor = LogitsProcessor(config.vocab_size)
+        self.make_empty_intermediate_tensors = self.model.make_empty_intermediate_tensors
+
+        # Set MoE hyperparameters for individual experts
+        self.expert_weights = []
+
+        self.moe_layers: list[FusedMoE] = []
+        example_layer = None
+        for layer in self.model.layers:
+            if isinstance(layer, PPMissingLayer):
+                continue
+
+            assert isinstance(layer, Qwen3MoeDecoderLayer)
+            if isinstance(layer.mlp, FusedMoE):
+                example_layer = layer.mlp
+                self.moe_layers.append(layer.mlp)
+
+        if example_layer is None:
+            raise RuntimeError("No Qwen3OmniMoe layer found in the model.layers.")
+
+        self.num_moe_layers = len(self.moe_layers)
+        self.num_expert_groups = 1
+        self.num_shared_experts = 0
+        self.num_logical_experts = example_layer.n_logical_experts
+        self.num_physical_experts = example_layer.n_physical_experts
+        self.num_local_physical_experts = example_layer.n_local_physical_experts
+        self.num_routed_experts = example_layer.n_routed_experts
+        self.num_redundant_experts = example_layer.n_redundant_experts

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py

@@ -0,0 +1,224 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Copyright 2025 The Qwen team.
+"""Inference-only Qwen3-Omni-Moe Code2Wav model."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+
+import numpy as np
+import torch
+import torch.nn as nn
+from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import (
+    Qwen3OmniMoeCode2WavConfig,
+)
+from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import (
+    Qwen3OmniMoeCausalConvNet,
+    Qwen3OmniMoeCausalTransConvNet,
+    Qwen3OmniMoeCode2WavDecoderBlock,
+    Qwen3OmniMoeCode2WavTransformerModel,
+    Qwen3OmniMoeConvNeXtBlock,
+    SnakeBeta,
+)
+from vllm.config import VllmConfig  # type: ignore
+from vllm.logger import init_logger  # type: ignore
+from vllm.model_executor.models.utils import (  # type: ignore
+    AutoWeightsLoader,
+    WeightsMapper,
+)
+
+logger = init_logger(__name__)
+
+
+class Qwen3OmniMoeCode2Wav(nn.Module):
+    """
+    Qwen3 Omni MoE Code2Wav - Converts num_quantizers-layer RVQ codec codes to audio waveform.
+
+    Architecture:
+    1. Code Embedding: Embed and average num_quantizers RVQ layers
+    2. Pre-Transformer: Add temporal context via sliding-window attention
+    3. Upsampling: Progressive upsampling with ConvNeXt blocks
+    4. Decoder: Multi-stage upsampling + residual units → waveform
+
+    Input: [batch, num_quantizers, seq_len] - num_quantizers-layer RVQ codes
+    Output: [batch, 1, waveform_len] - Audio waveform [-1, 1]
+
+    Total upsampling factor: ~1280x
+    Example: 100 codec frames → 128,000 audio samples (8 seconds at 16kHz)
+    """
+
+    input_modalities = "audio"
+
+    # Weight mapper
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_prefix={
+            "code2wav.pre_transformer.": "pre_transformer.",
+            "code2wav.code_embedding.": "code_embedding.",
+            "code2wav.upsample.": "upsample.",
+            "code2wav.decoder.": "decoder.",
+            "code2wav.": "",
+        }
+    )
+
+    def __init__(
+        self,
+        *,
+        vllm_config: VllmConfig | None = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+
+        self.config: Qwen3OmniMoeCode2WavConfig = vllm_config.model_config.hf_config
+
+        # Calculate total upsampling factor
+        self.total_upsample = np.prod(self.config.upsample_rates + self.config.upsampling_ratios)
+
+        # Pre-transformer
+        self.pre_transformer = Qwen3OmniMoeCode2WavTransformerModel._from_config(self.config)
+
+        # Code embedding: Single embedding table for all RVQ layers
+        self.code_embedding = nn.Embedding(
+            self.config.codebook_size * self.config.num_quantizers, self.config.hidden_size
+        )
+
+        # Offset for each RVQ layer (layer 0: 0-1023, layer 1: 1024-2047, etc.)
+        self.register_buffer(
+            "code_offset",
+            torch.arange(self.config.num_quantizers).view(1, -1, 1) * self.config.codebook_size,
+            persistent=False,
+        )
+
+        # Upsampling blocks (e.g., 2x, 2x)
+        upsample = []
+        for factor in self.config.upsampling_ratios:
+            upsample.append(
+                nn.ModuleList(
+                    [
+                        Qwen3OmniMoeCausalTransConvNet(
+                            self.config.hidden_size, self.config.hidden_size, factor, factor
+                        ),
+                        Qwen3OmniMoeConvNeXtBlock(self.config.hidden_size),
+                    ]
+                )
+            )
+        self.upsample = nn.ModuleList(upsample)
+
+        # Decoder: Initial projection + progressive upsampling blocks
+        decoder = [Qwen3OmniMoeCausalConvNet(self.config.hidden_size, self.config.decoder_dim, kernel_size=7)]
+
+        # Add decoder blocks (each upsamples and reduces channels)
+        for i in range(len(self.config.upsample_rates)):
+            decoder.append(Qwen3OmniMoeCode2WavDecoderBlock(self.config, i))
+
+        # Final projection to waveform
+        output_dim = self.config.decoder_dim // 2 ** len(self.config.upsample_rates)
+        decoder += [
+            SnakeBeta(output_dim),
+            Qwen3OmniMoeCausalConvNet(output_dim, 1, kernel_size=7),
+        ]
+        self.decoder = nn.ModuleList(decoder)
+
+    def forward(self, codes: torch.Tensor) -> torch.Tensor:
+        """
+        Convert num_quantizers-layer RVQ codes to audio waveform.
+
+        Args:
+            codes: [batch, num_quantizers, seq_len] - num_quantizers-layer RVQ codec codes
+
+        Returns:
+            waveform: [batch, 1, waveform_len] - Audio waveform clipped to [-1, 1]
+        """
+        if codes.shape[1] != self.config.num_quantizers:
+            raise ValueError(f"Expected {self.config.num_quantizers} layers of codes, got {codes.shape[1]}")
+
+        # Stage 1: Code Embedding
+        # Add offset to separate layer vocabularies, then embed and average
+        hidden = self.code_embedding(codes + self.code_offset).mean(1)
+        # Shape: [batch, seq_len, hidden_size]
+
+        # Stage 2: Pre-Transformer (add temporal context)
+        hidden = self.pre_transformer(inputs_embeds=hidden).last_hidden_state
+        # Shape: [batch, seq_len, hidden_size]
+
+        # Stage 3: Upsampling
+        hidden = hidden.permute(0, 2, 1)  # [batch, hidden_size, seq_len]
+        for blocks in self.upsample:
+            for block in blocks:
+                hidden = block(hidden)
+        # Shape: [batch, hidden_size, seq_len * upsample_factor]
+
+        # Stage 4: Decoder (progressive upsampling to waveform)
+        wav = hidden
+        for block in self.decoder:
+            wav = block(wav)
+        # Shape: [batch, 1, waveform_len]
+
+        # Clamp to valid audio range
+        return wav.clamp(min=-1.0, max=1.0)
+
+    def chunked_decode(
+        self,
+        codes: torch.Tensor,
+        chunk_size: int = 300,
+        left_context_size: int = 25,
+    ) -> torch.Tensor:
+        """
+        Decode long sequences in chunks to avoid OOM.
+
+        Uses overlapping chunks with left context to avoid boundary artifacts.
+
+        Args:
+            codes: [batch, num_quantizers, seq_len] - num_quantizers-layer RVQ codes
+            chunk_size: Number of codec frames per chunk
+            left_context_size: Number of overlapping frames for context
+
+        Returns:
+            waveform: [batch, 1, waveform_len] - Complete waveform
+        """
+        wavs = []
+        start_index = 0
+
+        while start_index < codes.shape[-1]:
+            end_index = min(start_index + chunk_size, codes.shape[-1])
+            context_size = left_context_size if start_index >= left_context_size else start_index
+
+            # Extract chunk with left context
+            codes_chunk = codes[..., start_index - context_size : end_index]
+
+            # Decode chunk
+            wav_chunk = self(codes_chunk)
+
+            # Remove context from output (context_size * total_upsample samples)
+            wavs.append(wav_chunk[..., context_size * self.total_upsample :])
+
+            start_index = end_index
+
+        return torch.cat(wavs, dim=-1)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        """Load weights from HuggingFace checkpoint."""
+        loader = AutoWeightsLoader(
+            self,
+            skip_prefixes=["thinker.", "talker."],  # Already loaded above
+        )
+        loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
+
+        # Log load summary
+        try:
+            total_bytes = 0
+            for name, param in self.named_parameters():
+                if param is not None and param.data is not None:
+                    total_bytes += param.data.numel() * param.data.element_size()
+            device = next(self.parameters()).device
+            logger.info(
+                "[Model Loaded] name=%s, success=%s, size=%.2f MB, device=%s",
+                self.__class__.__name__,
+                True,
+                total_bytes / (1024**2),
+                str(device),
+            )
+        except Exception:
+            pass
+
+        return loaded

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_code_predictor_mtp.py

@@ -0,0 +1,491 @@
+"""Qwen3-Omni Code Predictor with MTP (Multi-Token Prediction) support.
+
+This module implements the code predictor component for Qwen3-Omni talker models.
+
+The code predictor generates residual RVQ (Residual Vector Quantization) codes
+autoregressively, predicting layers 1 to N based on layer-0 codes from the talker.
+"""
+
+from typing import Any, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import Cache, PretrainedConfig
+from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import Qwen3OmniMoeRotaryEmbedding
+from vllm.config import CacheConfig, ModelConfig, VllmConfig
+from vllm.logger import init_logger
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
+
+logger = init_logger(__name__)
+
+
+# ============================================================================
+# Rotary Embeddings and Helper Functions
+# ============================================================================
+
+
+class Qwen3OmniCodePredictorRotaryEmbedding(nn.Module):
+    """Rotary positional embeddings for the code predictor."""
+
+    def __init__(self, dim, max_position_embeddings=2048, base=10000):
+        super().__init__()
+        self.dim = dim
+        self.max_position_embeddings = max_position_embeddings
+        self.base = base
+
+        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float() / self.dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+    def forward(self, x: torch.Tensor, position_ids: torch.LongTensor) -> tuple[torch.Tensor, torch.Tensor]:
+        inv_freq = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
+        position = position_ids[:, None, :].float()
+        freqs = (inv_freq @ position).transpose(1, 2)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        cos = emb.cos().to(x.dtype)
+        sin = emb.sin().to(x.dtype)
+        return cos, sin
+
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
+    """Applies Rotary Position Embedding to the query and key tensors.
+
+    Args:
+        q (`torch.Tensor`): The query tensor.
+        k (`torch.Tensor`): The key tensor.
+        cos (`torch.Tensor`): The cosine part of the rotary embedding.
+        sin (`torch.Tensor`): The sine part of the rotary embedding.
+        position_ids (`torch.Tensor`, *optional*):
+            Deprecated and unused.
+        unsqueeze_dim (`int`, *optional*, defaults to 1):
+            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+    Returns:
+        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+    """
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """Repeat key/value heads for grouped query attention."""
+    if n_rep == 1:
+        return hidden_states
+    bsz, num_kv_heads, seq_len, head_dim = hidden_states.shape
+    hidden_states = hidden_states[:, :, None, :, :].expand(bsz, num_kv_heads, n_rep, seq_len, head_dim)
+    return hidden_states.reshape(bsz, num_kv_heads * n_rep, seq_len, head_dim)
+
+
+# ============================================================================
+# Code Predictor Attention Layer
+# ============================================================================
+
+
+class Qwen3OmniCodePredictorAttention(nn.Module):
+    """Multi-head self-attention for code predictor with vLLM optimization."""
+
+    def __init__(self, config, layer_idx: int, vllm_config: VllmConfig = None):
+        super().__init__()
+
+        self.num_heads = config.code_predictor_config.num_attention_heads
+        self.num_key_value_heads = config.code_predictor_config.num_key_value_heads
+        self.head_dim = getattr(
+            config.code_predictor_config,
+            "head_dim",
+            config.code_predictor_config.hidden_size // config.code_predictor_config.num_attention_heads,
+        )
+        self.hidden_size = config.code_predictor_config.hidden_size
+
+        if self.num_heads % self.num_key_value_heads != 0:
+            raise ValueError("num_attention_heads must be divisible by num_key_value_heads")
+
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+
+        # Projection layers
+        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
+
+        # Query/Key normalization
+        self.q_norm = RMSNorm(self.head_dim, eps=config.code_predictor_config.rms_norm_eps)
+        self.k_norm = RMSNorm(self.head_dim, eps=config.code_predictor_config.rms_norm_eps)
+        self.is_causal = True
+        self.config = config
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        causal_mask: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        past_key_values: Optional[Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        use_cache: bool = False,
+        position_ids: Optional[torch.LongTensor] = None,
+    ) -> torch.Tensor:
+        bsz, seq_len, _ = hidden_states.shape
+
+        # Project to Q, K, V
+        q = self.q_proj(hidden_states).view(bsz, seq_len, self.num_heads, self.head_dim)
+        k = self.k_proj(hidden_states).view(bsz, seq_len, self.num_key_value_heads, self.head_dim)
+        v = self.v_proj(hidden_states).view(bsz, seq_len, self.num_key_value_heads, self.head_dim)
+
+        # Apply normalization
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+
+        # Apply RoPE
+        q_heads = q.transpose(1, 2).contiguous()
+        k_heads = k.transpose(1, 2).contiguous()
+        q_heads, k_heads = apply_rotary_pos_emb(q_heads, k_heads, cos, sin)
+        v_heads = v.transpose(1, 2).contiguous()
+
+        if past_key_values is not None:
+            # sin and cos are specific to RoPE models; cache_position needed for the static cache
+            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+            k_heads, v_heads = past_key_values.update(k_heads, v_heads, self.layer_idx, cache_kwargs)
+
+        from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+        attention_interface = ALL_ATTENTION_FUNCTIONS["flash_attention_2"]
+        attn_output, _ = attention_interface(
+            self,
+            q_heads,
+            k_heads,
+            v_heads,
+            None,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.head_dim**-0.5,
+            sliding_window=None,
+            use_cache=use_cache,
+            position_ids=position_ids,
+            output_hidden_states=True,
+            output_attentions=False,
+        )
+        attn_output = attn_output.reshape(*(hidden_states.shape[:-1]), -1).contiguous()
+
+        attn_output = self.o_proj(attn_output)
+        return attn_output
+
+
+# ============================================================================
+# Code Predictor MLP Layer
+# ============================================================================
+
+
+class Qwen3OmniCodePredictorMLP(nn.Module):
+    """Feed-forward network for code predictor."""
+
+    def __init__(self, config):
+        super().__init__()
+        self.gate_proj = nn.Linear(
+            config.code_predictor_config.hidden_size, config.code_predictor_config.intermediate_size, bias=False
+        )
+        self.up_proj = nn.Linear(
+            config.code_predictor_config.hidden_size, config.code_predictor_config.intermediate_size, bias=False
+        )
+        self.down_proj = nn.Linear(
+            config.code_predictor_config.intermediate_size, config.code_predictor_config.hidden_size, bias=False
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        gate = F.silu(self.gate_proj(hidden_states))
+        up = self.up_proj(hidden_states)
+        return self.down_proj(gate * up)
+
+
+# ============================================================================
+# MTP Layer (Multi-Token Prediction Layer)
+# ============================================================================
+
+
+class Qwen3OmniCodePredictorMTPLayer(nn.Module):
+    """MTP layer for speculative decoding - predicts next residual code layer."""
+
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        prefix: str,
+        model_config: ModelConfig,
+        layer_idx: int,
+        cache_config: Optional[CacheConfig] = None,
+        quant_config: Optional[QuantizationConfig] = None,
+    ) -> None:
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.config = config
+
+        # Qwen3OmniCodePredictorDecoderLayer
+        self.self_attn = Qwen3OmniCodePredictorAttention(
+            config,
+            layer_idx,
+            vllm_config=type(
+                "VllmConfig",
+                (),
+                {"cache_config": cache_config, "quant_config": quant_config, "model_config": model_config},
+            )(),
+        )
+        self.mlp = Qwen3OmniCodePredictorMLP(config)
+        self.input_layernorm = RMSNorm(
+            config.code_predictor_config.hidden_size, eps=config.code_predictor_config.rms_norm_eps
+        )
+        self.post_attention_layernorm = RMSNorm(
+            config.code_predictor_config.hidden_size, eps=config.code_predictor_config.rms_norm_eps
+        )
+
+    def mtp_block(
+        self,
+        hidden_states: torch.Tensor,
+        causal_mask: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        past_key_values: Optional[Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        use_cache: bool = False,
+        position_ids: Optional[torch.LongTensor] = None,
+    ) -> torch.Tensor:
+        # Self-attention with residual
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states = self.self_attn(
+            hidden_states, causal_mask, cos, sin, past_key_values, cache_position, use_cache, position_ids
+        )
+        hidden_states = residual + hidden_states
+
+        # MLP with residual
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        return hidden_states
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        previous_hidden_states: torch.Tensor,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        spec_step_index: int = 0,
+    ) -> torch.Tensor:
+        assert inputs_embeds is not None, "inputs_embeds required for MTP"
+
+        # Mask position 0 (not needed for MTP)
+        inputs_embeds[positions == 0] = 0
+
+        hidden_states = torch.cat([inputs_embeds, previous_hidden_states], dim=-1)
+
+        # Get position info for RoPE
+        batch_size, seq_len, _ = hidden_states.shape
+        device = hidden_states.device
+        position_ids = torch.arange(seq_len, dtype=torch.long, device=device).unsqueeze(0).expand(batch_size, -1)
+
+        # Get RoPE embeddings
+        head_dim = self.self_attn.head_dim
+        rotary_emb = Qwen3OmniCodePredictorRotaryEmbedding(
+            head_dim, max_position_embeddings=self.config.code_predictor_config.max_position_embeddings
+        )
+        rotary_emb = Qwen3OmniMoeRotaryEmbedding(self.config)
+        cos, sin = rotary_emb(hidden_states, position_ids)
+
+        # Create causal mask
+        causal_mask = (
+            torch.triu(torch.ones(seq_len, seq_len, dtype=torch.bool, device=device), diagonal=1)
+            .unsqueeze(0)
+            .unsqueeze(0)
+        )
+        causal_mask = causal_mask.masked_fill(causal_mask, float("-inf"))
+
+        # Forward through MTP block
+        hidden_states = self.mtp_block(hidden_states, causal_mask, cos, sin)
+
+        return hidden_states
+
+
+class Qwen3OmniCodePredictorBaseModel(nn.Module):
+    """
+    Base model for code predictor - matches HF Qwen3OmniMoeTalkerCodePredictorModel structure.
+
+    This is a simple transformer that processes inputs_embeds and outputs hidden states.
+    """
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        config = vllm_config.model_config.hf_config.code_predictor_config
+
+        self.config = config
+        self.vocab_size = config.vocab_size
+        self.num_code_groups = config.num_code_groups
+
+        # Codec embeddings (for layers 1-num_code_groups-1)
+        self.codec_embedding = nn.ModuleList(
+            [
+                VocabParallelEmbedding(
+                    config.vocab_size,
+                    config.hidden_size,
+                )
+                for _ in range(config.num_code_groups - 1)
+            ]
+        )
+
+        # Decoder layers
+        self.layers = nn.ModuleList(
+            [
+                Qwen3OmniCodePredictorMTPLayer(
+                    vllm_config.model_config.hf_config,
+                    f"{prefix}.layers.{idx}",
+                    model_config=vllm_config.model_config,
+                    layer_idx=idx,
+                    cache_config=vllm_config.cache_config,
+                    quant_config=vllm_config.quant_config,
+                )
+                for idx in range(config.num_hidden_layers)
+            ]
+        )
+
+        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        # RoPE
+        self.rotary_emb = Qwen3OmniMoeRotaryEmbedding(config=config)
+
+    def forward(
+        self,
+        inputs_embeds: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Any] = None,
+        use_cache: Optional[bool] = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs,
+    ):
+        """
+        Forward pass matching HF structure.
+
+        Args:
+            inputs_embeds: [batch, seq_len, hidden_size]
+
+        Returns:
+            Object with .last_hidden_state attribute
+        """
+        batch_size, seq_len, _ = inputs_embeds.shape
+
+        # Create positions tensor if not provided
+        # positions must be [num_tokens] or [batch_size, seq_len]
+        if position_ids is None:
+            if cache_position is not None:
+                position_ids = cache_position  # [num_tokens]
+            else:
+                position_ids = torch.arange(seq_len, device=inputs_embeds.device).unsqueeze(0)  # [1, seq_len]
+        else:
+            position_ids = position_ids.flatten()  # Ensure [num_tokens]
+
+        # Extract cos/sin from rotary_emb cache
+        # The cos_sin_cache is [max_pos, rotary_dim * 2]
+        cos, sin = self.rotary_emb(inputs_embeds, position_ids)
+
+        # Create causal mask
+        device = inputs_embeds.device
+        causal_mask = (
+            torch.triu(torch.ones(seq_len, seq_len, dtype=torch.bool, device=device), diagonal=1)
+            .unsqueeze(0)
+            .unsqueeze(0)
+        )
+        causal_mask = causal_mask.masked_fill(causal_mask, float("-inf"))
+
+        # Forward through decoder layers
+        hidden_states = inputs_embeds
+
+        for layer in self.layers:
+            hidden_states = layer.mtp_block(
+                hidden_states, causal_mask, cos, sin, past_key_values, cache_position, use_cache, position_ids
+            )
+
+        # Final norm
+        hidden_states = self.norm(hidden_states)
+
+        # Return in HF-compatible format
+        from collections import namedtuple
+
+        Output = namedtuple("Output", ["last_hidden_state", "past_key_values"])
+        return Output(last_hidden_state=hidden_states, past_key_values=None)  # [batch, num_code_groups-1, hidden_size]
+
+    def get_input_embeddings(self):
+        """Return codec embeddings for HF compatibility."""
+        return self.codec_embedding
+
+
+class Qwen3OmniMoeTalkerCodePredictor(nn.Module):
+    """
+    Code predictor wrapper matching HF structure.
+
+    Structure:
+    - self.model: Qwen3OmniCodePredictorBaseModel (transformer)
+    - self.lm_head: ModuleList of output heads
+    """
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+
+        talker_code_predictor_config = vllm_config.model_config.hf_config
+        self.quant_config = vllm_config.quant_config
+        self.prefix = prefix
+
+        self.config = talker_code_predictor_config
+        self.vocab_size = self.config.code_predictor_config.vocab_size
+        self.num_code_groups = self.config.code_predictor_config.num_code_groups
+
+        # Base transformer model (matches HF structure)
+        self.model = Qwen3OmniCodePredictorBaseModel(vllm_config=vllm_config, prefix=prefix)
+
+        # Output heads for each residual layer (1-num_layers-1)
+        self.lm_head = nn.ModuleList(
+            [
+                nn.Linear(
+                    self.config.code_predictor_config.hidden_size,
+                    self.config.code_predictor_config.vocab_size,
+                    bias=False,
+                )
+                for _ in range(self.num_code_groups - 1)
+            ]
+        )
+
+    def forward(
+        self,
+        inputs_embeds: torch.Tensor,  # [batch, seq_len, hidden_size]
+        layer_idx: int,  # Which layer to predict (0-num_layers-2 for layers 1-num_layers-1)
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Forward pass for code predictor.
+
+        Args:
+            inputs_embeds: Input embeddings [batch_size, seq_len, hidden_size]
+            layer_idx: Which residual layer to predict (0-num_layers-2 for layers 1-num_layers-1)
+
+        Returns:
+            logits: Predicted logits [batch_size, seq_len, vocab_size]
+            hidden_states: Output hidden states [batch_size, seq_len, hidden_size]
+        """
+        # Pass through base model
+        hidden_states = self.model(inputs_embeds)
+
+        # Get logits from corresponding head
+        logits = self.lm_head[layer_idx](hidden_states)
+
+        return logits, hidden_states

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py

@@ -0,0 +1,703 @@
+from collections.abc import Iterable
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import (
+    Qwen3OmniMoeTalkerConfig,
+)
+from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import (
+    Qwen3OmniMoeAudioEncoder,
+)
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+from vllm.model_executor.layers.activation import _ACTIVATION_REGISTRY
+from vllm.model_executor.layers.linear import ReplicatedLinear
+from vllm.model_executor.models.interfaces import (
+    MultiModalEmbeddings,
+    SupportsMultiModal,
+    SupportsPP,
+)
+from vllm.model_executor.models.qwen2_5_omni_thinker import (
+    Qwen2_5OmniThinkerDummyInputsBuilder,
+)
+from vllm.model_executor.models.qwen3_moe import Qwen3MoeMLP
+from vllm.model_executor.models.utils import (
+    AutoWeightsLoader,
+    WeightsMapper,
+    maybe_prefix,
+    sequence_parallel_chunk,
+)
+from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.sequence import IntermediateTensors
+from vllm.v1.sample.metadata import SamplingMetadata
+
+from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_code_predictor_mtp import (
+    Qwen3OmniMoeTalkerCodePredictor,
+)
+from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import (
+    Qwen3MoeLLMForCausalLM,
+    Qwen3Omni_VisionTransformer,
+    Qwen3OmniMoeConditionalGenerationMixin,
+    Qwen3OmniMoeThinkerMultiModalProcessor,
+    Qwen3OmniMoeThinkerProcessingInfo,
+)
+
+try:
+    import flash_attn
+except (ImportError, ModuleNotFoundError):
+    flash_attn = None
+
+from vllm_omni.model_executor.models.utils import safe_tensor_reshape
+
+logger = init_logger(__name__)
+
+Qwen3OmniMoeThinkerDummyInputsBuilder = Qwen2_5OmniThinkerDummyInputsBuilder
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+    Qwen3OmniMoeThinkerMultiModalProcessor,
+    info=Qwen3OmniMoeThinkerProcessingInfo,
+    dummy_inputs=Qwen3OmniMoeThinkerDummyInputsBuilder,
+)
+class Qwen3OmniMoeTalkerForConditionalGeneration(
+    nn.Module,
+    SupportsMultiModal,
+    SupportsPP,
+    Qwen3OmniMoeConditionalGenerationMixin,
+):
+    """
+    Qwen3 Omni MoE Talker - Converts text to audio codec codes.
+
+    The talker is the second stage of Qwen3 Omni MoE's TTS pipeline:
+    1. Thinker: Generates text response + hidden states
+    2. Talker: Converts those to 8-layer audio codec codes
+    3. Code2Wav: Converts codes to waveform
+
+    ## Key Components:
+    - text_projection: Projects thinker text embeddings → talker dimension
+    - hidden_projection: Projects thinker hidden states → talker dimension
+    - language_model: Main MoE transformer (generates layer 0)
+    - codec_head: Projects to codec vocabulary (layer 0 logits)
+    - code_predictor: Small transformer for layers 1-num_layers-1
+    """
+
+    logger = init_logger(__name__)
+
+    # Weight mapping from HuggingFace to vLLM naming convention
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_prefix={
+            # Main MoE transformer model
+            "talker.model.": "language_model.model.",
+            # Codec head remains separate (outputs audio codes, not text)
+            "talker.codec_head.": "codec_head.",
+            # Code predictor: Now matches HF structure exactly (has .model sub-module)
+            # e.g., "talker.code_predictor.model.codec_embedding.0" → "code_predictor.model.codec_embedding.0"
+            "talker.code_predictor.": "code_predictor.",
+            # Projection layers
+            "talker.text_projection.": "text_projection.",
+            "talker.hidden_projection.": "hidden_projection.",
+            # Fallback: strip talker prefix
+            "talker.": "",
+        }
+    )
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        talker_config: Qwen3OmniMoeTalkerConfig = vllm_config.model_config.hf_config
+        self.quant_config = vllm_config.quant_config
+        self.prefix = prefix
+
+        self.config = talker_config
+        self.vocab_size = talker_config.text_config.vocab_size
+        self.router_aux_loss_coef = talker_config.text_config.router_aux_loss_coef
+        self.num_experts = talker_config.text_config.num_experts
+        self.num_experts_per_tok = talker_config.text_config.num_experts_per_tok
+        # thinker projection components for talker
+        self.text_projection = Qwen3OmniMoeTalkerResizeMLP(self.config)
+        self.hidden_projection = Qwen3OmniMoeTalkerResizeMLP(self.config)
+        self.codec_head = nn.Linear(self.config.text_config.hidden_size, self.config.text_config.vocab_size, bias=False)
+
+        self.rope_deltas = None
+        self.spatial_merge_size = self.config.spatial_merge_size
+
+        self.vocab_size = self.config.code_predictor_config.vocab_size
+        self.num_code_groups = self.config.code_predictor_config.num_code_groups
+
+        self.language_model = Qwen3OmniMoeModel(
+            vllm_config=vllm_config,
+            talker_config=self.config,
+            prefix=maybe_prefix(prefix, "language_model"),
+        )
+
+        self.code_predictor = Qwen3OmniMoeTalkerCodePredictor(
+            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "code_predictor")
+        )
+
+    def code_predictor_forward(
+        self,
+        input_ids: torch.Tensor,
+        inputs_embeds: torch.Tensor | None = None,
+        *,
+        temperature: float = 1.0,
+        top_k: int = 50,  # Match transformers default
+        top_p: float = 0.8,  # Match transformers default
+        generation_steps: int | None = None,
+        last_talker_hidden: torch.Tensor | None = None,
+        **_: object,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Generate full RVQ codec codes for the provided sequence.
+
+        The code predictor consumes the layer-0 codec codes produced by the talker
+        alongside the talker's hidden states, and autoregressively predicts the remaining
+        residual layers (to num_codec_groups).
+
+        Returns:
+            tuple containing:
+                - residual_codes: A tensor of shape [batch, num_code_groups, seq_len] containing
+                  the complete set of codec codes
+                - summed_embeddings: A tensor of shape [batch, seq_len, hidden_size]
+                  Sum of all layer embeddings at each position (like Transformers)
+        """
+        if input_ids is None:
+            raise ValueError("`input_ids` containing layer-0 codec codes must be provided.")
+        if inputs_embeds is None:
+            raise ValueError("`inputs_embeds` containing talker hidden states must be provided.")
+
+        if inputs_embeds.ndim == 2:
+            inputs_embeds = inputs_embeds.unsqueeze(0)
+        if input_ids.ndim == 1:
+            input_ids = input_ids.unsqueeze(0)
+
+        # Ensure the tensors are contiguous for the autoregressive sampling loop
+        inputs_embeds = inputs_embeds.contiguous()
+        input_ids = input_ids.contiguous()
+
+        # Generate full codec codes using MTP
+        # This will be the parallel prediction implementation
+        batch_size, seq_len = input_ids.shape
+
+        # For now, use sequential generation (TODO: implement parallel)
+        # Result will be [batch, num_code_groups, seq_len]
+        # - all_codes_per_position will collect [batch, num_code_groups, 1] for each position
+        all_codes_per_position = []
+        middle_hidden_states = []  # Collect hidden states for each position
+
+        # Generate residual layers for each position
+        for pos in range(seq_len):
+            layer0_code = input_ids[:, pos : pos + 1]  # [batch, 1]
+
+            # Predict all residual layers (layers 1 to num_code_groups-1) autoregressively
+            pos_codes = [layer0_code]  # Start with layer 0: [batch, 1]
+
+            # Initial input: [last_talker_hidden, layer0_embed]
+            layer0_embed = self.get_input_embeddings(layer0_code)
+            prev_embed = layer0_embed  # Track previous layer embedding
+            try:
+                current_input = torch.cat([last_talker_hidden, prev_embed], dim=1)  # [batch, 2, hidden_size]
+            except Exception as e:
+                print(f"Error in current_input: {e}")
+                print(f"last_talker_hidden shape: {last_talker_hidden.shape}")
+                print(f"prev_embed shape: {prev_embed.shape}")
+                raise e
+
+            for layer_idx in range(self.num_code_groups - 1):
+                # Input for this layer: [last_talker_hidden, prev_layer_embed]
+
+                # Forward through code_predictor model
+                outputs = self.code_predictor.model(
+                    inputs_embeds=current_input,
+                    attention_mask=None,
+                    position_ids=None,
+                    past_key_values=None,
+                    use_cache=False,
+                    cache_position=None,
+                )
+
+                hidden_state = outputs.last_hidden_state  # [batch, 2, hidden_size]
+
+                # Use the corresponding lm_head for this layer
+                logits = self.code_predictor.lm_head[layer_idx](hidden_state[:, -1:, :])  # [batch, 1, vocab_size]
+                from transformers.generation.logits_process import (
+                    LogitsProcessorList,
+                    TopKLogitsWarper,
+                    TopPLogitsWarper,
+                )
+
+                logits_processors = LogitsProcessorList(
+                    [
+                        TopKLogitsWarper(top_k=top_k),
+                        TopPLogitsWarper(top_p=top_p),
+                    ]
+                )
+                input_ids_for_logits_processors = torch.tensor([pos_codes[1:]]).to(logits.device, dtype=torch.long)
+                logits = logits_processors(input_ids_for_logits_processors, logits.squeeze(0)).unsqueeze(0)
+
+                # Sample from the filtered distribution
+                probs = F.softmax(logits, dim=-1)
+                code = torch.multinomial(probs.squeeze(1), num_samples=1)  # [batch, 1]
+                pos_codes.append(code)
+
+                # Update prev_embed for next layer (if not last layer)
+                # layer_idx=0 predicts layer 1, embed with codec_embedding[1]
+                new_embed = self.code_predictor.model.codec_embedding[layer_idx](code)  # [batch, 1, hidden_size]
+                current_input = torch.cat([current_input, new_embed], dim=1)  # [batch, 3~n, hidden_size]
+
+            # Stack all layers for this position: [batch, num_code_groups, 1]
+            pos_all_layers = torch.stack(pos_codes, dim=1)  # [batch, num_code_groups, 1]
+            all_codes_per_position.append(pos_all_layers)
+            middle_hidden_states.append(current_input[:, 2:-1, :])
+
+        # Concatenate across positions: [batch, num_code_groups, seq_len]
+        result_codes = torch.cat(all_codes_per_position, dim=2)
+
+        # Build summed embeddings for each position (like Transformers)
+        # This combines layer-0 embed, mid layers hidden states, and last layer embed
+        all_summed_embeddings = []
+
+        for pos in range(seq_len):
+            # Layer 0 embedding
+            layer0_code = result_codes[:, 0, pos : pos + 1]  # [batch, 1]
+            layer0_embed = self.get_input_embeddings(layer0_code)  # [batch, 1, hidden_size]
+
+            # mid layers hidden states (from CodePredictor)
+            mid_residual_hiddens = middle_hidden_states[pos]  # [batch, num_code_groups-2, hidden_size]
+            mid_list = list(mid_residual_hiddens.split(1, dim=1))
+
+            # last layer embedding
+            last_layer_code = result_codes[:, -1, pos : pos + 1]  # [batch, 1]
+            last_residual_hidden = self.code_predictor.model.codec_embedding[-1](last_layer_code)
+
+            # Concatenate all layers: [batch, num_code_groups, hidden_size]
+            pos_codec_hiddens = torch.cat(
+                [layer0_embed] + mid_list + [last_residual_hidden],
+                dim=1,
+            )
+
+            # Sum across layers: [batch, 1, hidden_size] (like Transformers)
+            pos_summed = pos_codec_hiddens.sum(dim=1, keepdim=True)
+            all_summed_embeddings.append(pos_summed)
+
+        # Concatenate across positions: [batch, seq_len, hidden_size]
+        summed_embeddings = torch.cat(all_summed_embeddings, dim=1)
+
+        return result_codes, summed_embeddings
+
+    def init_multi_modal(self, thinker_config):
+        """
+        Initialize multimodal components from the thinker.
+
+        Unlike Qwen2.5 Omni which creates audio_tower and visual encoders here,
+        Qwen3 Omni MoE has a cleaner separation: the thinker is the ONLY module
+        that processes raw multimodal inputs. The talker only handles text-to-audio
+        conversion using pre-processed embeddings from the thinker.
+
+        This method exists for API compatibility and stores the thinker config
+        for reference. The actual multimodal processing components (audio_tower,
+        visual) are ONLY in the thinker, not duplicated in the talker.
+
+        Args:
+            thinker_config: Configuration from the thinker model (for reference only)
+        """
+        self.audio_tower = Qwen3OmniMoeAudioEncoder(thinker_config.audio_config)
+        self.visual = Qwen3Omni_VisionTransformer(
+            vision_config=thinker_config.vision_config,
+            norm_eps=getattr(thinker_config.text_config, "rms_norm_eps", 1e-6),
+            quant_config=self.quant_config,
+            prefix=maybe_prefix(self.prefix, "visual"),
+            # attn_backend_override=attn_backend_override,
+        )
+
+    def project_thinker_outputs(
+        self,
+        thinker_embeds: torch.Tensor | None = None,
+        thinker_hidden_states: torch.Tensor | None = None,
+        is_multimodal_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """
+        Project thinker outputs to talker's hidden dimension.
+
+        The talker has a different hidden size than the thinker, so we need
+        to project the inputs appropriately:
+        - Text embeddings (from thinker's embedding layer) → text_projection
+        - Hidden states (from thinker's last layer, for multimodal) → hidden_projection
+
+        Args:
+            thinker_embeds: Text embeddings from thinker [batch, seq, thinker_hidden]
+            thinker_hidden_states: Hidden states from thinker's last layer [batch, seq, thinker_hidden]
+            is_multimodal_mask: Boolean mask indicating multimodal positions [batch, seq]
+
+        Returns:
+            projected_embeds: [batch, seq, talker_hidden]
+        """
+        if thinker_embeds is None and thinker_hidden_states is None:
+            raise ValueError("Either thinker_embeds or thinker_hidden_states must be provided")
+
+        # If only embeddings provided, project all as text
+        if thinker_hidden_states is None or is_multimodal_mask is None:
+            return self.text_projection(thinker_embeds)
+
+        # If only hidden states provided, project all as hidden
+        if thinker_embeds is None:
+            return self.hidden_projection(thinker_hidden_states)
+
+        # Mixed case: use mask to decide which projection
+        batch_size, seq_len, _ = thinker_embeds.shape
+        output = torch.empty(
+            (batch_size, seq_len, self.config.text_config.hidden_size),
+            device=thinker_embeds.device,
+            dtype=thinker_embeds.dtype,
+        )
+
+        # Project multimodal regions using hidden states
+        if is_multimodal_mask.any():
+            mm_hidden = thinker_hidden_states[is_multimodal_mask]
+            projected_mm = self.hidden_projection(mm_hidden)
+            output[is_multimodal_mask] = projected_mm
+
+        # Project text regions using embeddings
+        if (~is_multimodal_mask).any():
+            text_embeds = thinker_embeds[~is_multimodal_mask]
+            projected_text = self.text_projection(text_embeds)
+            output[~is_multimodal_mask] = projected_text
+
+        return output
+
+    def forward(
+        self,
+        batched_talker_inputs: dict,
+        intermediate_tensors: IntermediateTensors | None = None,
+        **kwargs: object,
+    ) -> torch.Tensor | IntermediateTensors:
+        """Forward pass through the talker model.
+
+        For inference, the talker receives inputs_embeds that should already be
+        projected to talker's hidden dimension. If receiving raw thinker outputs,
+        use project_thinker_outputs() first.
+        """
+        # If intermediate_tensors is provided (pipeline parallel),
+        # inputs_embeds should be None
+        batched_input_ids = []
+        batched_positions = []
+        batched_input_embeds = []
+        batched_code_predictor_codes = []
+        if intermediate_tensors is not None:
+            inputs_embeds = None
+        for rid in batched_talker_inputs:
+            input_ids = safe_tensor_reshape(batched_talker_inputs[rid]["input_ids"], (1, -1))
+            positions = batched_talker_inputs[rid]["positions"]
+            inputs_embeds = safe_tensor_reshape(
+                batched_talker_inputs[rid]["inputs_embeds"], (-1, self.config.text_config.hidden_size)
+            )
+            text_step = safe_tensor_reshape(batched_talker_inputs[rid]["text_step"], (1, -1))
+            last_talker_hidden = safe_tensor_reshape(
+                batched_talker_inputs[rid]["last_talker_hidden"], (1, 1, self.config.text_config.hidden_size)
+            )
+            # for profiling
+            if inputs_embeds.shape[-1] == 2048:
+                inputs_embeds = self.text_projection(inputs_embeds)
+            if inputs_embeds.shape[0] == 1:
+                code_predictor_codes, summed_embeddings = self.code_predictor_forward(
+                    input_ids, inputs_embeds.clone(), last_talker_hidden=last_talker_hidden
+                )
+                inputs_embeds = summed_embeddings.clone()
+            else:
+                code_predictor_codes = torch.zeros((0, self.num_code_groups), dtype=torch.long)
+            batched_input_ids.append(input_ids.reshape(-1))
+            batched_positions.append(positions)
+            batched_input_embeds.append((inputs_embeds + text_step).reshape(-1, self.config.text_config.hidden_size))
+            batched_code_predictor_codes.append(code_predictor_codes.squeeze(-1).detach().to("cpu").contiguous())
+        try:
+            talker_input_ids = torch.cat(batched_input_ids, dim=0)
+            talker_positions = torch.cat(batched_positions, dim=0)
+            talker_input_embeds = torch.cat(batched_input_embeds, dim=0)
+        except Exception as e:
+            print(f"Error in talker_input_embeds: {e}")
+            print(f"talker_input_embeds shape: {[tensor.shape for tensor in batched_input_embeds]}")
+            print(f"talker_input_ids shape: {[tensor.shape for tensor in batched_input_ids]}")
+            print(f"talker_positions shape: {[tensor.shape for tensor in batched_positions]}")
+            raise e
+        talker_hidden_states, _ = self.language_model.model(
+            talker_input_ids,
+            talker_positions,
+            intermediate_tensors,
+            inputs_embeds=talker_input_embeds,
+            **kwargs,
+        )
+
+        # Pass talker hidden states to code predictor
+        # Returns: (residual_codes, summed_embeddings=sum of all layer embeddings at each position)
+
+        # Return both talker hidden states and code predictor results
+        # code_predictor_codes: [batch, num_code_groups, seq_len]
+        # summed_embeddings: [batch, seq_len, hidden_size]
+        #   - Sum of all layer embeddings at each position (like Transformers)
+
+        return talker_hidden_states, batched_code_predictor_codes
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor,
+        sampling_metadata: SamplingMetadata = None,
+    ) -> torch.Tensor | None:
+        """Compute logits for audio codec codes (layer 0 of RVQ).
+
+        This projects the hidden states to the codec vocabulary space.
+        For full audio generation, layers except 0 would be predicted by
+        the code_predictor after sampling.
+        """
+        logits = self.codec_head(hidden_states)
+        return logits
+
+    def make_empty_intermediate_tensors(
+        self, batch_size: int, dtype: torch.dtype, device: torch.device
+    ) -> IntermediateTensors:
+        """Create empty intermediate tensors for pipeline parallelism."""
+        return self.language_model.make_empty_intermediate_tensors(batch_size, dtype, device)
+
+    def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:
+        mm_input_by_modality = {}
+
+        # Preserve the order of modalities if there are multiple of them
+        # from the order of kwargs.
+        for input_key in kwargs:
+            if input_key in ("pixel_values", "image_embeds") and "image" not in mm_input_by_modality:
+                mm_input_by_modality["image"] = self._parse_and_validate_image_input(**kwargs)
+            if input_key in ("pixel_values_videos", "video_embeds") and "video" not in mm_input_by_modality:
+                mm_input_by_modality["video"] = self._parse_and_validate_video_input(**kwargs)
+            if input_key in ("input_audio_features") and "audio" not in mm_input_by_modality:
+                mm_input_by_modality["audio"] = self._parse_and_validate_audio_input(**kwargs)
+        return mm_input_by_modality
+
+    def get_multimodal_embeddings(self, **kwargs: object) -> MultiModalEmbeddings | None:
+        mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs)
+        if not mm_input_by_modality:
+            return []
+
+        # The result multimodal_embeddings is tuple of tensors, with each
+        # tensor correspoending to a multimodal data item (image or video).
+        multimodal_embeddings: tuple[torch.Tensor, ...] = ()
+
+        # NOTE: It is important to iterate over the keys in this dictionary
+        # to preserve the order of the modalities.
+        # TODO: do projection for all multimodel
+        for modality in mm_input_by_modality:
+            multimodal_input = mm_input_by_modality[modality]
+            if modality == "image":
+                image_embeddings = self._process_image_input(multimodal_input)
+                image_embeddings = self.hidden_projection(image_embeddings)
+                multimodal_embeddings += tuple(image_embeddings)
+            if modality == "video":
+                video_embeddings = self._process_video_input(multimodal_input)
+                video_video_embeddings_project = ()
+                for video_embed in video_embeddings:
+                    proj = nn.Linear(8192, 2048).to(device=video_embed.device, dtype=torch.bfloat16)
+                    video_embed = proj(video_embed)
+                    video_embed_project = self.hidden_projection(video_embed)
+                    video_video_embeddings_project += (video_embed_project,)
+                multimodal_embeddings += tuple(video_video_embeddings_project)
+            if modality == "audio":
+                audio_embeddings = self._process_audio_input(multimodal_input)
+                audio_embeddings = self.hidden_projection(audio_embeddings)
+                multimodal_embeddings += tuple(audio_embeddings)
+        return multimodal_embeddings
+
+    def get_input_embeddings(self, input_ids: torch.Tensor, multimodal_embeddings: MultiModalEmbeddings | None = None):
+        """Get the input embedding layer (for codec tokens)."""
+        return self.language_model.get_input_embeddings(input_ids, multimodal_embeddings)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        """Load weights for the talker model.
+
+        The weight mapping translates from HuggingFace naming convention
+        to vLLM's internal structure. Code predictor weights are routed
+        to its custom loader for vocab extension support.
+        """
+        loader = AutoWeightsLoader(
+            self,
+            skip_prefixes=["thinker.", "code2wav."],
+            # "code_predictor."],
+        )
+        # Don't apply mapper again since we already did it
+        loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
+
+        # Log load summary
+        try:
+            total_bytes = 0
+            for name, param in self.named_parameters():
+                if param is not None and param.data is not None:
+                    total_bytes += param.data.numel() * param.data.element_size()
+            device = next(self.parameters()).device
+            logger.info(
+                "[Model Loaded] name=%s, success=%s, size=%.2f MB, device=%s",
+                self.__class__.__name__,
+                True,
+                total_bytes / (1024**2),
+                str(device),
+            )
+        except Exception:
+            pass
+
+        multi_model_weights = set()
+        for name, param in self.visual.named_parameters():
+            multi_model_weights.add("visual." + name)
+        for name, param in self.audio_tower.named_parameters():
+            multi_model_weights.add("audio_tower." + name)
+        loaded.update(multi_model_weights)
+
+        return loaded
+
+
+class Qwen3OmniMoeTalkerResizeMLP(nn.Module):
+    """
+    MLP for projecting between thinker and talker hidden dimensions.
+
+    The thinker and talker have different hidden sizes:
+    - Thinker: config.thinker_hidden_size (e.g., 3584)
+    - Talker: config.text_config.hidden_size (e.g., 2048)
+
+    This MLP projects from thinker → talker dimension.
+    Two instances are used:
+    - text_projection: For text embeddings from thinker's embedding layer
+    - hidden_projection: For hidden states from thinker's last transformer layer
+    """
+
+    def __init__(self, config: Qwen3OmniMoeTalkerConfig):
+        super().__init__()
+        self.linear_fc1 = nn.Linear(config.thinker_hidden_size, config.text_config.intermediate_size, bias=True)
+        self.linear_fc2 = nn.Linear(config.text_config.intermediate_size, config.text_config.hidden_size, bias=True)
+        self.act_fn = _ACTIVATION_REGISTRY[config.text_config.hidden_act]  # silu
+
+    def forward(self, hidden_state):
+        return self.linear_fc2(self.act_fn(self.linear_fc1(hidden_state)))
+
+
+class Qwen3OmniMoeModel(Qwen3MoeLLMForCausalLM):
+    def __init__(self, vllm_config, talker_config, prefix):
+        talker_vllm_config = vllm_config.with_hf_config(
+            talker_config.text_config, architectures=["Qwen3MoeForCausalLM"]
+        )
+        talker_vllm_config.model_config.hf_text_config = talker_vllm_config.model_config.hf_config
+        super().__init__(
+            vllm_config=talker_vllm_config,
+            prefix=prefix,
+        )
+
+        self.config = talker_config
+
+        # Remove the inherited LM head so the talker only exposes codec outputs.
+        if hasattr(self, "lm_head"):
+            del self.lm_head
+
+        # Replace the base embed tokens with codec embedding (defined below).
+        if hasattr(self.model, "embed_tokens"):
+            del self.model.embed_tokens
+
+        # Codec embedding for RVQ code generation
+        self.model.codec_embedding = nn.Embedding(
+            talker_config.text_config.vocab_size, talker_config.text_config.hidden_size
+        )
+
+        # Add shared expert to each MoE layer and patch the forward method
+        layer_idx = 0
+        for layer in self.model.layers:
+            # add shared expert to Qwen3OmniMoeSparseMoeBlock layers
+            if hasattr(layer.mlp, "experts"):  # Check if it's a SparseMoeBlock
+                # Shared expert is a regular gated MLP (SwiGLU)
+                layer.mlp.shared_expert = Qwen3MoeMLP(
+                    hidden_size=self.config.text_config.hidden_size,
+                    intermediate_size=self.config.text_config.shared_expert_intermediate_size,
+                    hidden_act=self.config.text_config.hidden_act,
+                    quant_config=talker_vllm_config.quant_config,
+                    reduce_results=False,  # Don't reduce since we'll add it manually
+                    prefix=f"{prefix}.layers.{layer_idx}.mlp.shared_expert",
+                )
+
+                # Shared expert gate outputs a single scalar per token
+                layer.mlp.shared_expert_gate = ReplicatedLinear(
+                    self.config.text_config.hidden_size,
+                    1,  # Output single scalar per token
+                    bias=False,
+                    quant_config=None,
+                    prefix=f"{prefix}.layers.{layer_idx}.mlp.shared_expert_gate",
+                )
+
+                # Store MoE config values for router computation
+                layer.mlp.top_k = self.config.text_config.num_experts_per_tok
+                layer.mlp.norm_topk_prob = self.config.text_config.norm_topk_prob
+                layer.mlp.num_experts = self.config.text_config.num_experts
+
+                # Monkey-patch the forward method to use shared expert
+                layer.mlp.forward = self._create_moe_forward_with_shared_expert(layer.mlp)
+
+            layer_idx += 1
+
+    def _create_moe_forward_with_shared_expert(self, moe_layer):
+        """Create a forward method that includes shared expert computation.
+
+        This matches the Transformers implementation where:
+        1. Compute shared expert output (regular MLP)
+        2. Gate it with sigmoid(shared_expert_gate(x))
+        3. Apply softmax BEFORE top-k selection (matches Transformers router)
+        4. Add to routed expert outputs
+        """
+
+        def forward_with_shared_expert(hidden_states: torch.Tensor, layer_idx: int = 0) -> torch.Tensor:
+            # Save original shape
+            orig_shape = hidden_states.shape
+            hidden_dim = hidden_states.shape[-1]
+            hidden_states = hidden_states.view(-1, hidden_dim)
+
+            # handle sequence parallel if needed
+            if hasattr(moe_layer, "is_sequence_parallel") and moe_layer.is_sequence_parallel:
+                hidden_states = sequence_parallel_chunk(hidden_states)
+
+            # Compute shared expert output
+            # The shared expert is a regular MLP, not a routed MoE
+            shared_output = None
+            if hasattr(moe_layer, "shared_expert") and moe_layer.shared_expert is not None:
+                # Forward through shared expert MLP
+                shared_output = moe_layer.shared_expert(hidden_states)
+
+                # Apply gating with sigmoid: sigmoid(gate(x)) * shared_expert(x)
+                if hasattr(moe_layer, "shared_expert_gate") and moe_layer.shared_expert_gate is not None:
+                    gate_logits, _ = moe_layer.shared_expert_gate(hidden_states)
+                    gate_values = F.sigmoid(gate_logits)  # [batch, 1]
+                    shared_output = gate_values * shared_output  # Broadcasting: [batch, 1] * [batch, hidden]
+
+            # Compute experts results
+            # router_logits: (num_tokens, n_experts)
+            router_logits, _ = moe_layer.gate(hidden_states)
+            experts_output = moe_layer.experts(hidden_states=hidden_states, router_logits=router_logits)
+
+            # combine experts and shared expert results
+            if shared_output is not None:
+                final_hidden_states = experts_output + shared_output
+
+            # Handle sequence parallel if needed
+            if hasattr(moe_layer, "is_sequence_parallel") and moe_layer.is_sequence_parallel:
+                from vllm.distributed import tensor_model_parallel_all_gather
+
+                num_tokens = orig_shape[0] if len(orig_shape) > 1 else 1
+                final_hidden_states = tensor_model_parallel_all_gather(final_hidden_states, 0)
+                final_hidden_states = final_hidden_states[:num_tokens]
+            try:
+                final_hidden_states.view(orig_shape)
+            except Exception as e:
+                print(f"Error viewing final hidden states: {e}")
+                print(f"final_hidden_states.shape: {final_hidden_states.shape}")
+                print(f"orig_shape: {orig_shape}")
+                raise e
+            # Return with original shape
+            return final_hidden_states.view(orig_shape)
+
+        return forward_with_shared_expert
+
+    def get_input_embeddings(
+        self,
+        input_ids: torch.Tensor,
+        multimodal_embeddings: MultiModalEmbeddings | None = None,
+        generation_steps=None,
+    ) -> torch.Tensor:
+        return self.model.codec_embedding(input_ids)

vllm_omni/model_executor/models/registry.py

@@ -27,6 +27,27 @@ _OMNI_MODELS = {
         "Qwen2_5OmniToken2WavModel",
     ),
     "Qwen2ForCausalLM_old": ("qwen2_5_omni", "qwen2_old", "Qwen2ForCausalLM"),  # need to discuss
+    # Qwen3 Omni MoE models
+    "Qwen3OmniMoeForConditionalGeneration": (
+        "qwen3_omni",
+        "qwen3_omni",
+        "Qwen3OmniMoeForConditionalGeneration",
+    ),
+    "Qwen3OmniMoeThinkerForConditionalGeneration": (
+        "qwen3_omni",
+        "qwen3_omni_moe_thinker",
+        "Qwen3OmniMoeThinkerForConditionalGeneration",
+    ),
+    "Qwen3OmniMoeTalkerForConditionalGeneration": (
+        "qwen3_omni",
+        "qwen3_omni_moe_talker",
+        "Qwen3OmniMoeTalkerForConditionalGeneration",
+    ),
+    "Qwen3OmniMoeCode2Wav": (
+        "qwen3_omni",
+        "qwen3_omni_code2wav",
+        "Qwen3OmniMoeCode2Wav",
+    ),
 }
 
 _VLLM_OMNI_MODELS = {

vllm_omni/model_executor/models/utils.py

@@ -28,3 +28,12 @@ def split_list_into_ranges(lst: torch.Tensor, interval: int) -> list[list[int]]:
         ranges[index].append(num)
 
     return ranges
+
+
+def safe_tensor_reshape(tensor: torch.Tensor, shape: tuple) -> torch.Tensor:
+    """
+    Reshape a tensor safely.
+    """
+    if tensor is None:
+        return None
+    return tensor.reshape(shape)

vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml

@@ -30,7 +30,7 @@ stage_args:
     runtime:
       process: true
       devices: "1"
-      max_batch_size: 3
+      max_batch_size: 1
     engine_args:
       model_stage: talker
       model_arch: Qwen2_5OmniForConditionalGeneration

vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml

@@ -0,0 +1,94 @@
+# Stage config for running Qwen3-Omni-MoE with 3-stage architecture
+# Stage 0: Thinker (multimodal understanding + text generation)
+# Stage 1: Talker (text embeddings → 8-layer RVQ codec codes)
+# Stage 2: Code2Wav (8-layer RVQ codes → audio waveform)
+
+stage_args:
+  - stage_id: 0
+    runtime:
+      devices: "0"
+      max_batch_size: 1
+    engine_args:
+      model_stage: thinker
+      model_arch: Qwen3OmniMoeForConditionalGeneration
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
+      gpu_memory_utilization: 0.8
+      enforce_eager: true
+      trust_remote_code: true
+      engine_output_type: latent  # Output hidden states for talker
+      distributed_executor_backend: "mp"
+      enable_prefix_caching: false
+      hf_config_name: thinker_config
+      # tensor_parallel_size: 2  # Commented out - using single GPU for thinker
+    final_output: true
+    final_output_type: text
+    is_comprehension: true
+    default_sampling_params:
+      temperature: 0.4
+      top_p: 0.9
+      top_k: 1
+      max_tokens: 2048
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.05
+
+  - stage_id: 1
+    runtime:
+       devices: "1"
+       max_batch_size: 1
+    engine_args:
+       model_stage: talker
+       model_arch: Qwen3OmniMoeForConditionalGeneration
+       worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+       scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
+       gpu_memory_utilization: 0.5
+       enforce_eager: true
+       trust_remote_code: true
+       engine_output_type: latent  # Output codec codes for code2wav
+      #  tensor_parallel_size: 2
+       enable_prefix_caching: false
+       distributed_executor_backend: "mp"
+       hf_config_name: talker_config
+    engine_input_source: [0]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
+    # final_output: true
+    # final_output_type: text
+    default_sampling_params:
+      temperature: 0.9
+      top_k: 50
+      max_tokens: 4096
+      seed: 42
+      detokenize: False
+      repetition_penalty: 1.05
+      stop_token_ids: [2150]
+
+  - stage_id: 2
+    runtime:
+      devices: "1"
+      max_batch_size: 1
+    engine_args:
+      model_stage: code2wav
+      model_arch: Qwen3OmniMoeForConditionalGeneration
+      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      scheduler_cls: vllm_omni.core.sched.generation_scheduler.GenerationScheduler
+      enforce_eager: true
+      trust_remote_code: true
+      enable_prefix_caching: false
+      engine_output_type: audio  # Final output: audio waveform
+      gpu_memory_utilization: 0.1
+      distributed_executor_backend: "mp"
+      max_num_batched_tokens: 1000000
+      hf_config_name: thinker_config
+    engine_input_source: [1]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
+    final_output: true
+    final_output_type: audio
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 65536
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.1

vllm_omni/model_executor/stage_input_processors/qwen3_omni.py

@@ -0,0 +1,171 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Copyright 2025 The Qwen team.
+"""Stage input processor for Qwen3 Omni MoE: Thinker → Talker transition."""
+
+from typing import Union
+
+import torch
+from vllm.inputs import TextPrompt
+
+from vllm_omni.inputs.data import OmniTokensPrompt
+
+
+def _compute_talker_prompt_ids_length(info):
+    im_start_token_id = 151644
+    system_token_id = 8948
+    user_token_id = 872
+    assistant_token_id = 77091
+
+    thinker_sequences = torch.tensor(info["thinker_sequences"]).unsqueeze(0).cuda().long()  # [1, T]
+    input_ids = torch.tensor(info["thinker_input_ids"]).unsqueeze(0).cuda().long()  # [1, T]
+
+    im_start_indexes = torch.cat(
+        [
+            torch.nonzero(input_ids[0] == im_start_token_id).squeeze(1),
+            torch.tensor([thinker_sequences.shape[-1]], device=input_ids.device, dtype=input_ids.dtype),
+        ],
+        dim=0,
+    )
+
+    sum_user_len = 0
+    assistant_len = 0
+    for i in range(len(im_start_indexes) - 1):
+        s = int(im_start_indexes[i].item())
+        e = int(im_start_indexes[i + 1].item())
+        role = int(input_ids[0, s + 1].item())
+        if role == system_token_id:
+            continue
+        elif role == user_token_id:
+            sum_user_len += e - s
+        elif role == assistant_token_id and i == len(im_start_indexes) - 2:
+            assistant_len += 9  # 3 + 4 + 1 + 1
+        else:
+            pass
+
+    return sum_user_len + assistant_len
+
+
+def thinker2talker(
+    stage_list,
+    engine_input_source,
+    prompt: Union[OmniTokensPrompt, TextPrompt] = None,
+    requires_multimodal_data: bool = False,
+):
+    """
+    Process thinker outputs to create talker inputs.
+
+    Workflow:
+    1. Extract thinker's text generation outputs (token IDs + hidden states)
+    2. Split hidden states into: prompt embeddings + generated embeddings
+    3. Package for talker with additional information
+
+    Args:
+        stage_list: List of stage objects
+        engine_input_source: Source stage IDs (typically [0] for thinker)
+        prompt: Original prompt data
+
+    Returns:
+        List of OmniTokensPrompt for talker stage
+    """
+    if not engine_input_source:
+        raise ValueError("engine_input_source cannot be empty")
+
+    source_stage_id = engine_input_source[0]
+    if source_stage_id >= len(stage_list):
+        raise IndexError(f"Invalid stage_id: {source_stage_id}")
+
+    if stage_list[source_stage_id].engine_outputs is None:
+        raise RuntimeError(f"Stage {source_stage_id} has no outputs yet")
+
+    thinker_outputs = stage_list[source_stage_id].engine_outputs
+    talker_inputs = []
+
+    # Process each thinker output
+    for i, thinker_output in enumerate(thinker_outputs):
+        output = thinker_output.outputs[0]
+        thinker_embeddings = output.multimodal_output["0"].float().clone().detach().cuda()
+
+        thinker_hidden_states = output.multimodal_output["24"].float().clone().detach().cuda()
+        info = {
+            "thinker_embeddings": thinker_embeddings,
+            "thinker_hidden_states": thinker_hidden_states,
+            "thinker_sequences": thinker_output.prompt_token_ids
+            + output.token_ids,  # the thinker_sequences is the whole ids
+            "thinker_input_ids": thinker_output.prompt_token_ids,
+            # Provide thinker-side TTS token embeddings for talker projection
+            "tts_bos_embed": output.multimodal_output.get("tts_bos_embed").float().clone().detach().cuda(),
+            "tts_eos_embed": output.multimodal_output.get("tts_eos_embed").float().clone().detach().cuda(),
+            "tts_pad_embed": output.multimodal_output.get("tts_pad_embed").float().clone().detach().cuda(),
+        }
+        talker_inputs.append(
+            OmniTokensPrompt(
+                prompt_token_ids=[0] * _compute_talker_prompt_ids_length(info),
+                additional_information=info,
+                multi_modal_data=None,
+                mm_processor_kwargs=None,
+            )
+        )
+
+    return talker_inputs
+
+
+def talker2code2wav(
+    stage_list,
+    engine_input_source,
+    prompt: Union[OmniTokensPrompt, TextPrompt] = None,
+    requires_multimodal_data: bool = False,
+):
+    """
+    Process talker outputs to create code2wav inputs.
+
+    Workflow:
+    1. Extract talker's codec code outputs (8-layer RVQ codes)
+    2. Flatten codes for code2wav input
+    3. Package for code2wav stage
+
+    Args:
+        stage_list: List of stage objects
+        engine_input_source: Source stage IDs (typically [1] for talker)
+        prompt: Original prompt data
+
+    Returns:
+        List of OmniTokensPrompt for code2wav stage
+    """
+    if not engine_input_source:
+        raise ValueError("engine_input_source cannot be empty")
+
+    source_stage_id = engine_input_source[0]
+    if source_stage_id >= len(stage_list):
+        raise IndexError(f"Invalid stage_id: {source_stage_id}")
+
+    if stage_list[source_stage_id].engine_outputs is None:
+        raise RuntimeError(f"Stage {source_stage_id} has no outputs yet")
+
+    talker_outputs = stage_list[source_stage_id].engine_outputs
+    code2wav_inputs = []
+
+    # Process each talker output
+    for i, talker_output in enumerate(talker_outputs):
+        output = talker_output.outputs[0]
+
+        # Extract codec codes from talker output
+        # Expected shape: [8, seq_len] (8-layer RVQ codes)
+        codec_codes = (
+            output.multimodal_output["code_predictor_codes"]
+            .to(torch.long)
+            .transpose(0, 1)
+            .cpu()
+            .to(torch.long)
+            .reshape(-1)
+            .tolist()
+        )  # 16, seq_len
+        code2wav_inputs.append(
+            OmniTokensPrompt(
+                prompt_token_ids=codec_codes,
+                multi_modal_data=None,
+                mm_processor_kwargs=None,
+            )
+        )
+
+    return code2wav_inputs

vllm_omni/patch.py

@@ -1,16 +1,24 @@
 import sys
 
 from vllm.inputs.data import TokensPrompt as _OriginalTokensPrompt
-from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding as _OriginalMRotaryEmbedding
+from vllm.model_executor.layers.rotary_embedding import (
+    MRotaryEmbedding as _OriginalMRotaryEmbedding,
+)
+from vllm.v1.engine import EngineCoreOutput as _OriginalEngineCoreOutput
+from vllm.v1.engine import EngineCoreOutputs as _OriginalEngineCoreOutputs
 from vllm.v1.engine import EngineCoreRequest as _OriginalEngineCoreRequest
 from vllm.v1.request import Request as _OriginalRequest
 
-from vllm_omni.engine import OmniEngineCoreRequest
+from vllm_omni.engine import OmniEngineCoreOutput, OmniEngineCoreOutputs, OmniEngineCoreRequest
 from vllm_omni.inputs.data import OmniTokensPrompt
 from vllm_omni.model_executor.layers.mrope import MRotaryEmbedding
 from vllm_omni.request import OmniRequest
 
 for module_name, module in sys.modules.items():
+    if hasattr(module, "EngineCoreOutput") and module.EngineCoreOutput == _OriginalEngineCoreOutput:
+        module.EngineCoreOutput = OmniEngineCoreOutput
+    if hasattr(module, "EngineCoreOutputs") and module.EngineCoreOutputs == _OriginalEngineCoreOutputs:
+        module.EngineCoreOutputs = OmniEngineCoreOutputs
     if hasattr(module, "TokensPrompt") and module.TokensPrompt == _OriginalTokensPrompt:
         module.TokensPrompt = OmniTokensPrompt
     if hasattr(module, "MRotaryEmbedding") and module.MRotaryEmbedding == _OriginalMRotaryEmbedding:

vllm_omni/worker/gpu_ar_model_runner.py

@@ -6,9 +6,6 @@ and also outputs sampled tokens.
 
 from __future__ import annotations
 
-from typing import Any
-
-import numpy as np
 import torch
 from vllm.forward_context import BatchDescriptor
 from vllm.logger import init_logger
@@ -25,10 +22,8 @@ from vllm.v1.worker.gpu_model_runner import (
     has_kv_transfer_group,
     set_forward_context,
 )
-from vllm.v1.worker.ubatch_utils import UBatchSlices
 from vllm.v1.worker.utils import is_residual_scattered_for_sp
 
-from vllm_omni.engine import AdditionalInformationPayload, PromptEmbedsPayload
 from vllm_omni.outputs import OmniModelRunnerOutput
 from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
 
@@ -48,169 +43,9 @@ class GPUARModelRunner(OmniGPUModelRunner):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.input_ids = self._make_buffer(self.max_num_tokens, dtype=torch.int32)
-
-    def _preprocess(
-        self,
-        scheduler_output: SchedulerOutput,
-        num_scheduled_tokens_np: np.ndarray,
-        intermediate_tensors: IntermediateTensors | None = None,
-        ubatch_slices: UBatchSlices | None = None,
-        num_tokens_after_padding: torch.Tensor | None = None,
-    ) -> tuple[
-        int,
-        int,
-        torch.Tensor | None,
-        torch.Tensor | None,
-        torch.Tensor | None,
-        torch.Tensor,
-        IntermediateTensors | None,
-        dict[str, Any],
-        dict[str, dict] | None,
-    ]:
-        num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
-        if ubatch_slices:
-            assert num_tokens_after_padding is not None
-            num_input_tokens = int(num_tokens_after_padding[0].item() * 2)
-            self.pad_out_ubatch_slice(ubatch_slices, num_input_tokens)
-        elif ubatch_slices is None:
-            num_input_tokens = self._get_num_input_tokens(num_scheduled_tokens)
-            num_pad, num_tokens_after_padding = self.get_dp_padding(num_input_tokens)
-            num_input_tokens += num_pad
-
-        # _prepare_inputs may reorder the batch, so we must gather multi
-        # modal outputs after that to ensure the correct order
-        per_req_additional_information: dict[str, dict] | None = None
-        if self.supports_mm_inputs and get_pp_group().is_first_rank and not self.model_config.is_encoder_decoder:
-            # Run the multimodal encoder if any.
-            self._execute_mm_encoder(scheduler_output)
-            mm_embeds = self._gather_mm_embeddings(scheduler_output)
-
-            # NOTE(woosuk): To unify token ids and soft tokens (vision
-            # embeddings), we always use embeddings (rather than token ids)
-            # as input to the multimodal model, even when the input is text.
-            inputs_embeds_scheduled = self.model.get_input_embeddings(
-                input_ids=self.input_ids.gpu[:num_scheduled_tokens],
-                multimodal_embeddings=mm_embeds or None,
-            )
-
-            # TODO(woosuk): Avoid the copy. Optimize.
-            self.inputs_embeds.gpu[:num_scheduled_tokens].copy_(inputs_embeds_scheduled)
-
-            # Reset per-step additional information collector (deprecated concat path)
-            if hasattr(self, "_forward_additional_information"):
-                self._forward_additional_information = None
-            # New: per-request additional information for this step
-            per_req_additional_information = {}
-
-            # Overlay custom prompt_embeds per request for the prompt portion;
-            # collect additional_information (tensor/list) for prefill portion only
-            for req_index, req_id in enumerate(self.input_batch.req_ids):
-                req_state = self.requests[req_id]
-                pe_cpu = getattr(req_state, "prompt_embeds_cpu", None)
-                addi_cpu = getattr(req_state, "additional_information_cpu", None)
-                num_computed_tokens = int(self.input_batch.num_computed_tokens_cpu[req_index])
-                prompt_len = len(req_state.prompt_token_ids)
-                prompt_remaining = max(0, prompt_len - num_computed_tokens)
-                sched_tokens = int(num_scheduled_tokens_np[req_index])
-                overlay_len = min(sched_tokens, prompt_remaining)
-                if overlay_len <= 0:
-                    continue
-                if pe_cpu is not None:
-                    src = pe_cpu[num_computed_tokens : num_computed_tokens + overlay_len].to(
-                        dtype=self.dtype, device=self.device, non_blocking=True
-                    )
-                    start_offset = int(self.query_start_loc.cpu[req_index])
-                    self.inputs_embeds[start_offset : start_offset + overlay_len].copy_(src)
-                # Build per-request additional information (no cross-request concat)
-                if addi_cpu is not None and isinstance(addi_cpu, dict):
-                    req_info: dict[str, object] = {}
-                    for k, v in addi_cpu.items():
-                        if isinstance(v, torch.Tensor):
-                            # For prefill tokens, pass only the scheduled slice;
-                            # for decode or no scheduled tokens, pass whole tensor
-                            if overlay_len > 0:
-                                try:
-                                    seg = (
-                                        v[num_computed_tokens : num_computed_tokens + overlay_len]
-                                        .detach()
-                                        .to("cpu")
-                                        .contiguous()
-                                    )
-                                except Exception:
-                                    seg = v.detach().to("cpu").contiguous()
-                                req_info[k] = seg
-                            else:
-                                req_info[k] = v.detach().to("cpu").contiguous()
-                        elif isinstance(v, list):
-                            req_info[k] = v
-                        else:
-                            req_info[k] = v
-                    per_req_additional_information[req_id] = req_info
-
-            input_ids = self.input_ids.gpu[:num_input_tokens]
-            inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
-            model_kwargs = {
-                **self._init_model_kwargs(num_scheduled_tokens),
-                **self._extract_mm_kwargs(scheduler_output),
-            }
-        elif self.enable_prompt_embeds and get_pp_group().is_first_rank:
-            # Get the input embeddings for the tokens that are not input embeds,
-            # then put them into the appropriate positions.
-            # TODO(qthequartermasterman): Since even when prompt embeds are
-            # enabled, (a) not all requests will use prompt embeds, and (b)
-            # after the initial prompt is processed, the rest of the generated
-            # tokens will be token ids, it is not desirable to have the
-            # embedding layer outside of the CUDA graph all the time. The v0
-            # engine avoids this by "double compiling" the CUDA graph, once
-            # with input_ids and again with inputs_embeds, for all num_tokens.
-            # If a batch only has token ids, then including the embedding layer
-            # in the CUDA graph will be more performant (like in the else case
-            # below).
-            token_ids_idx = self.is_token_ids.gpu[:num_scheduled_tokens].nonzero(as_tuple=False).squeeze(1)
-            # Some tokens ids may need to become embeds
-            if token_ids_idx.numel() > 0:
-                token_ids = self.input_ids.gpu[token_ids_idx]
-                tokens_to_embeds = self.model.get_input_embeddings(input_ids=token_ids)
-                self.inputs_embeds.gpu[token_ids_idx] = tokens_to_embeds
-
-            inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
-            model_kwargs = self._init_model_kwargs(num_input_tokens)
-            input_ids = None
-        else:
-            # For text-only models, we use token ids as input.
-            # While it is possible to use embeddings as input just like the
-            # multimodal models, it is not desirable for performance since
-            # then the embedding layer is not included in the CUDA graph.
-            input_ids = self.input_ids.gpu[:num_input_tokens]
-            inputs_embeds = None
-            model_kwargs = self._init_model_kwargs(num_input_tokens)
-        if self.uses_mrope:
-            positions = self.mrope_positions.gpu[:, :num_input_tokens]
-        else:
-            positions = self.positions.gpu[:num_input_tokens]
-
-        if get_pp_group().is_first_rank:
-            intermediate_tensors = None
-        else:
-            intermediate_tensors = self.sync_and_slice_intermediate_tensors(
-                num_input_tokens, intermediate_tensors, True
-            )
-
-        if self.model_config.is_encoder_decoder and scheduler_output.scheduled_encoder_inputs:
-            encoder_inputs = self._extract_encoder_inputs(scheduler_output)
-            model_kwargs.update(encoder_inputs)
-
-        return (
-            num_scheduled_tokens,
-            num_input_tokens,
-            num_tokens_after_padding,
-            input_ids,
-            inputs_embeds,
-            positions,
-            intermediate_tensors,
-            model_kwargs,
-            per_req_additional_information,
-        )
+        # each model stage has their own hidden size
+        self.hidden_size = self.model_config.hf_text_config.hidden_size
+        self.inputs_embeds = self._make_buffer(self.max_num_tokens, self.hidden_size, dtype=self.dtype, numpy=False)
 
     @torch.inference_mode()
     def execute_model(
@@ -223,58 +58,9 @@ class GPUARModelRunner(OmniGPUModelRunner):
                 # Update persistent batch states.
                 self._update_states(scheduler_output)
 
-                # Decode per-request prompt_embeds / additional_hidden_states payloads
+                # Decode per-request prompt_embeds / additional_information payloads
                 # (if present) into CPU tensors
-                try:
-                    new_reqs = getattr(scheduler_output, "scheduled_new_reqs", [])
-                    if new_reqs:
-                        for nr in new_reqs:
-                            req_id = getattr(nr, "req_id", None) or getattr(nr, "request_id", None)
-                            if req_id is None:
-                                continue
-                            # prompt_embeds
-                            payload_pe = getattr(nr, "prompt_embeds", None)
-                            if payload_pe is not None:
-                                if isinstance(payload_pe, torch.Tensor):
-                                    pe_cpu = payload_pe.detach().to("cpu").contiguous()
-                                elif isinstance(payload_pe, PromptEmbedsPayload):
-                                    dt = np.dtype(getattr(payload_pe, "dtype", "float32"))
-                                    arr = np.frombuffer(payload_pe.data, dtype=dt)
-                                    arr = arr.reshape(payload_pe.shape)
-                                    pe_cpu = torch.from_numpy(arr)
-                                else:
-                                    pe_cpu = None
-                                if pe_cpu is not None and req_id in self.requests:
-                                    setattr(
-                                        self.requests[req_id],
-                                        "prompt_embeds_cpu",
-                                        pe_cpu,
-                                    )
-                            # additional_information
-                            payload_info = getattr(nr, "additional_information", None)
-                            if payload_info is not None:
-                                info_dict = {}
-                                if isinstance(payload_info, dict):
-                                    # Already decoded
-                                    info_dict = payload_info
-                                elif isinstance(payload_info, AdditionalInformationPayload):
-                                    for k, entry in payload_info.entries.items():
-                                        if entry.tensor_data is not None:
-                                            dt = np.dtype(getattr(entry, "tensor_dtype", "float32"))
-                                            arr = np.frombuffer(entry.tensor_data, dtype=dt)
-                                            arr = arr.reshape(entry.tensor_shape)
-                                            info_dict[k] = torch.from_numpy(arr)
-                                        else:
-                                            info_dict[k] = entry.list_data
-                                if info_dict and req_id in self.requests:
-                                    setattr(
-                                        self.requests[req_id],
-                                        "additional_information_cpu",
-                                        info_dict,
-                                    )
-                except Exception as e:
-                    logger.error(f"Error decoding prompt_embeds / additional_information: {e}")
-                    pass
+                self._decode_and_store_request_payloads(scheduler_output)
 
                 if not scheduler_output.total_num_scheduled_tokens:
                     if not has_kv_transfer_group():
@@ -344,29 +130,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
             record_function_or_nullcontext("Forward"),
             self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output,
         ):
-            model_kwargs_extra = {}
-            # Pass per-request additional information map for this step (no concat)
-            if per_req_additional_information:
-                model_kwargs_extra["additional_information_by_req_id"] = per_req_additional_information
-            # Always pass per-request runtime additional_information (persisted in request state)
-            try:
-                per_req_runtime_info = []
-                for req_id in self.input_batch.req_ids:
-                    req_state = self.requests.get(req_id)
-                    info = getattr(req_state, "additional_information_cpu", None) if req_state is not None else None
-                    per_req_runtime_info.append(info if isinstance(info, dict) else {})
-                model_kwargs_extra["runtime_additional_information"] = per_req_runtime_info
-                model_kwargs_extra["request_ids"] = self.input_batch.req_ids
-                # Pass each request's token span within the flattened sequence for this step,
-                # enabling the model to map decode/prefill by request
-                req_token_spans = []
-                for req_index in range(len(self.input_batch.req_ids)):
-                    start_offset = int(self.query_start_loc.cpu[req_index])
-                    sched_tokens = int(num_scheduled_tokens_np[req_index])
-                    req_token_spans.append((start_offset, start_offset + sched_tokens))
-                model_kwargs_extra["request_token_spans"] = req_token_spans
-            except Exception:
-                pass
+            model_kwargs_extra = self._build_model_kwargs_extra(per_req_additional_information, num_scheduled_tokens_np)
             model_output = self.model(
                 input_ids=input_ids,
                 positions=positions,
@@ -388,37 +152,15 @@ class GPUARModelRunner(OmniGPUModelRunner):
                 hidden_states = model_output
                 aux_hidden_states = None
 
-            hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states)
+            # hidden_states, multimodal_outputs = self.extract_multimodal_outputs(
+            #     hidden_states
+            # )
+            multimodal_outputs = model_output.multimodal_outputs
+            hidden_states = model_output.text_hidden_states
             # The model side may return per-request additional_information updates (model-agnostic channel).
             # Convention: multimodal_outputs["additional_information_update"] is a list[dict] in batch order;
             # the runner merges it into the corresponding request's additional_information_cpu for subsequent decode.
-            try:
-                if isinstance(multimodal_outputs, dict) and (
-                    "additional_information_update" in multimodal_outputs
-                    or "additional_information_update_by_req_id" in multimodal_outputs
-                ):
-                    # Option A: list[dict] in batch order
-                    updates_list = multimodal_outputs.get("additional_information_update")
-                    if isinstance(updates_list, list):
-                        for idx, upd in enumerate(updates_list):
-                            if not isinstance(upd, dict) or idx >= len(self.input_batch.req_ids):
-                                continue
-                            req_id = self.input_batch.req_ids[idx]
-                            self._merge_additional_information_update(req_id, upd)
-                    # Option B: dict[str, dict] keyed by req_id
-                    updates_map = multimodal_outputs.get("additional_information_update_by_req_id")
-                    if isinstance(updates_map, dict):
-                        for req_id, upd in updates_map.items():
-                            if not isinstance(upd, dict):
-                                continue
-                            if req_id not in self.requests:
-                                continue
-                            self._merge_additional_information_update(req_id, upd)
-            except Exception as e:
-                logger.error(
-                    f"Error merging for requests:{self.input_batch.req_ids} additional \
-                        information update: {e}, with the multimodal_outputs as {multimodal_outputs}"
-                )
+            self._process_additional_information_updates(multimodal_outputs)
             if not self.broadcast_pp_output:
                 # Common case.
                 if not get_pp_group().is_last_rank:
@@ -534,12 +276,42 @@ class GPUARModelRunner(OmniGPUModelRunner):
 
         # Convert to per-request tensors on CPU
         hidden_states_cpu = hidden_states.detach().to("cpu").contiguous()
+        # pooler_output: list[torch.Tensor | None] = []
         pooler_output: list[torch.Tensor | None] = []
         prev_logits_index = 0
-        for logits_index in logits_indices:
-            pooler_output.append(hidden_states_cpu[prev_logits_index : logits_index + 1])
+        for rid, logits_index in zip(req_ids_output_copy, logits_indices):
+            # Base payload: hidden slice for this request in this iteration
+            hidden_slice = hidden_states_cpu[prev_logits_index : logits_index + 1]
+            payload: dict[str, object] = {"hidden": hidden_slice}
+            # Merge multimodal_outputs if present
+            if isinstance(multimodal_outputs, dict) and multimodal_outputs:
+                mm_payload: dict[str, object] = {}
+                for k, v in multimodal_outputs.items():
+                    try:
+                        # Case 1: tensor aligned on token dimension
+                        if isinstance(v, torch.Tensor) and v.shape[0] == hidden_states_cpu.shape[0]:
+                            mm_payload[k] = v.detach().to("cpu")[prev_logits_index : logits_index + 1].contiguous()
+                        # Case 2: nested dict of tensors aligned on token dimension (e.g., selected_hidden_layers)
+                        elif isinstance(v, dict):
+                            sub_dict: dict[str, torch.Tensor] = {}
+                            for sk, sv in v.items():
+                                if isinstance(sv, torch.Tensor) and sv.shape[0] == hidden_states_cpu.shape[0]:
+                                    sub_dict[str(sk)] = (
+                                        sv.detach().to("cpu")[prev_logits_index : logits_index + 1].contiguous()
+                                    )
+                            if sub_dict:
+                                mm_payload[k] = sub_dict
+                        elif isinstance(v, list):
+                            element: torch.Tensor = v[0]
+                            multimodal_outputs[k] = v[1:] if len(v) > 1 else v
+                            mm_payload[k] = element
+                    except Exception as e:
+                        # Best-effort; skip malformed entries
+                        logger.error(f"Error in merge multimodal outputs: {e}")
+                if mm_payload:
+                    payload.update(mm_payload)
+            pooler_output.append(payload)  # type: ignore[arg-type]
             prev_logits_index = logits_index + 1
-
         output = OmniModelRunnerOutput(
             req_ids=req_ids_output_copy,
             req_id_to_index=req_id_to_index_output_copy,
@@ -583,3 +355,5 @@ class GPUARModelRunner(OmniGPUModelRunner):
             else:
                 merged[k] = v
         setattr(req_state, "additional_information_cpu", merged)
+
+    # ===== Helper functions extracted for clarity and reuse =====

vllm_omni/worker/gpu_generation_model_runner.py

@@ -0,0 +1,582 @@
+"""Code2Wav GPU Model Runner for vLLM-omni.
+
+Handles direct conversion from codec codes to audio waveforms for Qwen3 Omni MoE Code2Wav.
+This is a non-autoregressive model that doesn't require sampling or logits computation.
+"""
+
+from __future__ import annotations
+
+import gc
+import logging
+from typing import Any
+
+import numpy as np
+import torch
+from vllm.config import CUDAGraphMode
+from vllm.forward_context import BatchDescriptor
+from vllm.multimodal.inputs import MultiModalKwargs
+from vllm.v1.attention.backends.utils import CommonAttentionMetadata
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.spec_decode.eagle import EagleProposer
+from vllm.v1.utils import record_function_or_nullcontext
+from vllm.v1.worker.gpu_model_runner import (
+    EMPTY_MODEL_RUNNER_OUTPUT,
+    AsyncGPUModelRunnerOutput,
+    IntermediateTensors,
+    PerLayerAttnMetadata,
+    get_pp_group,
+    set_forward_context,
+)
+from vllm.v1.worker.ubatch_utils import UBatchSlices
+from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs
+
+from vllm_omni.outputs import OmniModelRunnerOutput
+from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
+
+logger = logging.getLogger(__name__)
+
+
+class GPUGenerationModelRunner(OmniGPUModelRunner):
+    """Generation model runner for vLLM-omni (non-autoregressive).
+
+    - Reuses GPUModelRunner preparation, multimodal handling, and TP/PP/DP glue.
+    - Does not compute logits or perform token sampling.
+    - Executes generation process and returns tensors via `pooler_output`.
+    """
+
+    def _preprocess(
+        self,
+        scheduler_output: SchedulerOutput,
+        intermediate_tensors: IntermediateTensors | None = None,
+        ubatch_slices: UBatchSlices | None = None,
+        num_tokens_after_padding: torch.Tensor | None = None,
+    ) -> tuple[
+        int,
+        int,
+        torch.Tensor | None,
+        torch.Tensor | None,
+        torch.Tensor | None,
+        torch.Tensor,
+        IntermediateTensors | None,
+        dict[str, Any],
+    ]:
+        # Input token count for this iteration (not used by diffusion, but
+        # retained to keep DP padding/ordering consistent)
+        num_input_tokens = scheduler_output.total_num_scheduled_tokens
+        num_pad, num_tokens_after_padding = self.get_dp_padding(num_input_tokens)
+        num_input_tokens += num_pad
+
+        # _prepare_inputs may reorder the batch, so we must gather multi
+        # modal outputs after that to ensure the correct order
+        if self.supports_mm_inputs and get_pp_group().is_first_rank and not self.model_config.is_encoder_decoder:
+            # Run the multimodal encoder if any.
+            self._execute_mm_encoder(scheduler_output)
+            mm_embeds = self._gather_mm_embeddings(scheduler_output)
+
+            # NOTE(woosuk): To unify token ids and soft tokens (vision
+            # embeddings), we always use embeddings (rather than token ids)
+            # as input to the multimodal model, even when the input is text.
+            inputs_embeds_scheduled = self.model.get_input_embeddings(
+                input_ids=self.input_ids.gpu[:num_input_tokens],
+                multimodal_embeddings=mm_embeds or None,
+            )
+
+            # TODO(woosuk): Avoid the copy. Optimize.
+            self.inputs_embeds.gpu[:num_input_tokens].copy_(inputs_embeds_scheduled)
+
+            input_ids = self.input_ids.gpu[:num_input_tokens]
+            inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
+            model_kwargs = {
+                **self._init_model_kwargs(num_input_tokens),
+                **self._extract_mm_kwargs(scheduler_output),
+            }
+        elif self.enable_prompt_embeds and get_pp_group().is_first_rank:
+            # Get the input embeddings for the tokens that are not input embeds,
+            # then put them into the appropriate positions.
+            # TODO(qthequartermasterman): Since even when prompt embeds are
+            # enabled, (a) not all requests will use prompt embeds, and (b)
+            # after the initial prompt is processed, the rest of the generated
+            # tokens will be token ids, it is not desirable to have the
+            # embedding layer outside of the CUDA graph all the time. The v0
+            # engine avoids this by "double compiling" the CUDA graph, once
+            # with input_ids and again with inputs_embeds, for all num_tokens.
+            # If a batch only has token ids, then including the embedding layer
+            # in the CUDA graph will be more performant (like in the else case
+            # below).
+            token_ids_idx = self.is_token_ids.gpu[:num_input_tokens].nonzero(as_tuple=False).squeeze(1)
+            # Some tokens ids may need to become embeds
+            if token_ids_idx.numel() > 0:
+                token_ids = self.input_ids.gpu[token_ids_idx]
+                tokens_to_embeds = self.model.get_input_embeddings(input_ids=token_ids)
+                self.inputs_embeds.gpu[token_ids_idx] = tokens_to_embeds
+
+            inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
+            model_kwargs = self._init_model_kwargs(num_input_tokens)
+            input_ids = self.input_ids.gpu[:num_input_tokens]
+        else:
+            # For text-only models, we use token ids as input.
+            # While it is possible to use embeddings as input just like the
+            # multimodal models, it is not desirable for performance since
+            # then the embedding layer is not included in the CUDA graph.
+            input_ids = self.input_ids.gpu[:num_input_tokens]
+            inputs_embeds = None
+            model_kwargs = self._init_model_kwargs(num_input_tokens)
+        if self.uses_mrope:
+            positions = self.mrope_positions.gpu[:, :num_input_tokens]
+        else:
+            positions = self.positions.gpu[:num_input_tokens]
+
+        if get_pp_group().is_first_rank:
+            intermediate_tensors = None
+        else:
+            intermediate_tensors = self.sync_and_slice_intermediate_tensors(
+                num_input_tokens, intermediate_tensors, True
+            )
+
+        if self.model_config.is_encoder_decoder and scheduler_output.scheduled_encoder_inputs:
+            encoder_inputs = self._extract_encoder_inputs(scheduler_output)
+            model_kwargs.update(encoder_inputs)
+
+        return (
+            num_input_tokens,
+            num_input_tokens,
+            num_tokens_after_padding,
+            input_ids,
+            inputs_embeds,
+            positions,
+            intermediate_tensors,
+            model_kwargs,
+        )
+
+    @torch.inference_mode()
+    def execute_model(
+        self,
+        scheduler_output: SchedulerOutput,
+        intermediate_tensors: IntermediateTensors | None = None,
+    ) -> OmniModelRunnerOutput | IntermediateTensors:
+        with record_function_or_nullcontext("Preprocess"):
+            with self.synchronize_input_prep():
+                # Update persistent batch states.
+                self._update_states(scheduler_output)
+                if not scheduler_output.total_num_scheduled_tokens:
+                    return EMPTY_MODEL_RUNNER_OUTPUT
+
+                # Prepare the decoder inputs.
+                (
+                    attn_metadata,
+                    logits_indices,
+                    spec_decode_metadata,
+                    num_scheduled_tokens_np,
+                    spec_decode_common_attn_metadata,
+                    max_query_len,
+                    ubatch_slices,
+                    num_tokens_after_padding,
+                ) = self._prepare_inputs(scheduler_output)
+
+            # Input token count for this iteration (not used by diffusion, but
+            # retained to keep DP padding/ordering consistent)
+            (
+                num_scheduled_tokens,
+                num_input_tokens,
+                num_tokens_across_dp,
+                input_ids,
+                inputs_embeds,
+                positions,
+                intermediate_tensors,
+                model_kwargs,
+            ) = self._preprocess(
+                scheduler_output,
+                intermediate_tensors,
+                ubatch_slices,
+                num_tokens_after_padding,
+            )
+
+        cudagraph_runtime_mode = CUDAGraphMode.NONE
+        # Run the model.
+        # Use persistent buffers for CUDA graphs.
+        with (
+            set_forward_context(
+                attn_metadata,
+                self.vllm_config,
+                num_tokens=num_input_tokens,
+                num_tokens_across_dp=num_tokens_across_dp,
+                cudagraph_runtime_mode=cudagraph_runtime_mode,
+                batch_descriptor=None,
+                ubatch_slices=ubatch_slices,
+            ),
+            record_function_or_nullcontext("Forward"),
+            self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output,
+        ):
+            outputs = self._run_generation_model(
+                input_ids=input_ids,
+                positions=positions,
+                intermediate_tensors=intermediate_tensors,
+                inputs_embeds=inputs_embeds,
+                multimodal_kwargs=model_kwargs,
+                logits_indices=logits_indices,
+            )
+
+        _, multimodal_outputs = self.extract_multimodal_outputs(outputs)
+        # Ensure one tensor per request, map to CPU for output struct
+        pooler_output: list[torch.Tensor | None] = []
+        if isinstance(multimodal_outputs, torch.Tensor):
+            # If model returned a single stacked tensor, split by requests
+            assert multimodal_outputs.shape[0] == self.input_batch.num_reqs
+            for i in range(self.input_batch.num_reqs):
+                pooler_output.append(multimodal_outputs[i].detach().to("cpu").contiguous())
+        elif isinstance(multimodal_outputs, list):
+            for out in multimodal_outputs:
+                pooler_output.append(out.detach().to("cpu").contiguous() if out is not None else None)
+        elif isinstance(multimodal_outputs, dict):
+            for out in multimodal_outputs.values():
+                pooler_output.append(out.detach().to("cpu").contiguous() if out is not None else None)
+        else:
+            raise RuntimeError("Unsupported diffusion output type")
+
+        output = OmniModelRunnerOutput(
+            req_ids=self.input_batch.req_ids,
+            req_id_to_index=self.input_batch.req_id_to_index,
+            sampled_token_ids=[],
+            logprobs=None,
+            prompt_logprobs_dict={},
+            pooler_output=pooler_output,
+            kv_connector_output=kv_connector_output,
+            num_nans_in_logits={},
+        )
+
+        if not self.use_async_scheduling:
+            return output
+
+        return AsyncGPUModelRunnerOutput(
+            model_runner_output=output,
+            sampled_token_ids=[],
+            invalid_req_indices=[],
+            async_output_copy_stream=self.async_output_copy_stream,
+        )
+
+    def _run_generation_model(
+        self,
+        *,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None,
+        inputs_embeds: torch.Tensor | None,
+        multimodal_kwargs: dict,
+        logits_indices: torch.Tensor,
+    ) -> torch.Tensor | list[torch.Tensor]:
+        """Run generation from codec codes to waveforms.
+
+        Args:
+            scheduler_output: Contains codec codes in input_ids or additional info
+            intermediate_tensors: PP intermediate tensors if applicable
+
+        Returns:
+            Audio waveforms: [batch, 1, waveform_len] or list of tensors
+        """
+        # Keep inputs identical to AR runner
+        kwargs = dict(
+            input_ids=input_ids,
+            positions=positions,
+            intermediate_tensors=intermediate_tensors,
+            inputs_embeds=inputs_embeds,
+            **MultiModalKwargs.as_kwargs(multimodal_kwargs, device=self.device),
+            sampling_metadata=self.input_batch.sampling_metadata,
+            logits_index=logits_indices,
+            sampler=self.sampler,
+        )
+
+        if hasattr(self.model, "forward"):
+            return self.model.forward(**kwargs)
+
+        raise RuntimeError(
+            "The loaded model does not expose diffusion interfaces 'sample', "
+            "'forward', or 'diffuse'. Please implement one of them or adapt the runner."
+        )
+
+    @torch.inference_mode()
+    def _dummy_sampler_run(self, hidden_states: torch.Tensor) -> None:
+        logger.warning("Dummy sampler run is not implemented for generation model")
+        return None
+
+    @torch.inference_mode()
+    def _dummy_run(
+        self,
+        num_tokens: int,
+        cudagraph_runtime_mode: CUDAGraphMode | None = None,
+        force_attention: bool = False,
+        uniform_decode: bool = False,
+        skip_eplb: bool = False,
+        is_profile: bool = False,
+        create_mixed_batch: bool = False,
+        remove_lora: bool = True,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Run a dummy forward pass to warm up/profile run or capture the
+        CUDA graph for the model.
+
+        Args:
+            num_tokens: Number of tokens to run the dummy forward pass.
+            cudagraph_runtime_mode: used to control the behavior.
+                - if not set will determine the cudagraph mode based on using
+                    the self.cudagraph_dispatcher.
+                - CUDAGraphMode.NONE: No cudagraph, for warm up and profile run
+                - CUDAGraphMode.PIECEWISE: Piecewise cudagraph.
+                - CUDAGraphMode.FULL: Full cudagraph, attention metadata is
+                    needed.
+            force_attention: If True, always create attention metadata. Used to
+                warm up attention backend when mode is NONE.
+            uniform_decode: If True, the batch is a uniform decode batch.
+            skip_eplb: If True, skip EPLB state update.
+            is_profile: If True, this is a profile run.
+            create_mixed_batch: If True, create a mixed batch with both decode
+                (1 token) and prefill (multiple tokens) requests.
+            remove_lora: If False, dummy LoRAs are not destroyed after the run
+        """
+        assert cudagraph_runtime_mode is None or cudagraph_runtime_mode in {
+            CUDAGraphMode.NONE,
+            CUDAGraphMode.PIECEWISE,
+            CUDAGraphMode.FULL,
+        }
+
+        # If cudagraph_mode.decode_mode() == FULL and
+        # cudagraph_mode.separate_routine(). This means that we are using
+        # different graphs and/or modes for mixed prefill-decode batches vs.
+        # uniform decode batches. A uniform decode batch means that all
+        # requests have identical query length, except a potential virtual
+        # request (shorter) in the batch account for padding.
+        # Uniform decode batch could either be common pure decode, where
+        # max_query_len == 1, or speculative decode, where
+        # max_query_len == 1 + num_spec_decode_tokens.
+
+        # When setting max_query_len = 1, we switch to and capture the optimized
+        # routine of FA2 for pure decode, i.e., Flashdecode + an optimization
+        # for GQA/MQA.
+        max_query_len = self.uniform_decode_query_len if uniform_decode else num_tokens
+
+        # Set num_scheduled_tokens based on num_tokens and max_num_seqs
+        # for dummy run with LoRA so that the num_reqs collectively
+        # has num_tokens in total.
+        assert num_tokens <= self.scheduler_config.max_num_batched_tokens
+        max_num_reqs = self.scheduler_config.max_num_seqs
+
+        num_reqs = min(num_tokens, max_num_reqs)
+        min_tokens_per_req = num_tokens // num_reqs
+        num_scheduled_tokens_list = [min_tokens_per_req] * num_reqs
+        num_scheduled_tokens_list[-1] += num_tokens % num_reqs
+
+        assert sum(num_scheduled_tokens_list) == num_tokens
+        assert len(num_scheduled_tokens_list) == num_reqs
+        num_scheduled_tokens = np.array(num_scheduled_tokens_list, dtype=np.int32)
+
+        num_tokens_after_padding = None
+
+        # If we failed to microbatch, currently need to resynchronize
+        # TODO(lucas,sage): we should be able to avoid this second sync by
+        #  refactoring `get_dp_padding_ubatch` and `get_dp_padding` into
+        #  a single `coordinate_batch_across_dp` function.
+        if num_tokens_after_padding is None:
+            num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens)
+            num_tokens_after_padding = num_tokens + num_pad
+        else:
+            num_tokens_across_dp = num_tokens_after_padding
+            num_tokens_after_padding = int(num_tokens_after_padding[0].item())
+
+        attn_metadata: PerLayerAttnMetadata | None = None
+
+        # If force_attention is True, we always capture attention. Otherwise,
+        # it only happens for cudagraph_runtime_mode=FULL.
+        if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL:
+            attn_metadata = {}
+
+            seq_lens = max_query_len
+            self.seq_lens.np[:num_reqs] = seq_lens
+            self.seq_lens.np[num_reqs:] = 0
+            self.seq_lens.copy_to_gpu()
+
+            cum_num_tokens, _ = self._get_cumsum_and_arange(num_scheduled_tokens)
+            self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens
+            self.query_start_loc.copy_to_gpu()
+
+            for kv_cache_group_id, kv_cache_group_spec in enumerate(self.kv_cache_config.kv_cache_groups):
+                common_attn_metadata = CommonAttentionMetadata(
+                    query_start_loc=self.query_start_loc.gpu[: num_reqs + 1],
+                    query_start_loc_cpu=self.query_start_loc.cpu[: num_reqs + 1],
+                    seq_lens=self.seq_lens.gpu[:num_reqs],
+                    seq_lens_cpu=self.seq_lens.cpu[:num_reqs],
+                    num_computed_tokens_cpu=self.input_batch.num_computed_tokens_cpu_tensor[:num_reqs],  # noqa: E501
+                    num_reqs=num_reqs,
+                    num_actual_tokens=num_tokens,
+                    max_query_len=max_query_len,
+                    max_seq_len=self.max_model_len,
+                    block_table_tensor=self.input_batch.block_table[kv_cache_group_id].get_device_tensor(num_reqs),
+                    slot_mapping=self.input_batch.block_table[kv_cache_group_id].slot_mapping.gpu[:num_tokens],
+                    causal=True,
+                )
+                for attn_group in self.attn_groups[kv_cache_group_id]:
+                    assert type(attn_metadata) is dict
+                    attn_metadata_i = attn_group.get_metadata_builder().build_for_cudagraph_capture(
+                        common_attn_metadata
+                    )  # noqa: E501
+                    for layer_name in attn_group.layer_names:
+                        attn_metadata[layer_name] = attn_metadata_i
+
+        with self.maybe_dummy_run_with_lora(self.lora_config, num_scheduled_tokens, remove_lora):
+            model_kwargs = self._init_model_kwargs(num_tokens)
+            if self.supports_mm_inputs and not self.model_config.is_encoder_decoder:
+                input_ids = None
+                inputs_embeds = self.inputs_embeds.gpu[:num_tokens]
+                model_kwargs = {
+                    **model_kwargs,
+                    **self._dummy_mm_kwargs(num_reqs),
+                }
+            elif self.enable_prompt_embeds:
+                input_ids = None
+                inputs_embeds = self.inputs_embeds.gpu[:num_tokens]
+                model_kwargs = self._init_model_kwargs(num_tokens)
+            else:
+                input_ids = self.input_ids.gpu[:num_tokens]
+                inputs_embeds = None
+
+            if self.uses_mrope:
+                positions = self.mrope_positions.gpu[:, :num_tokens]
+            else:
+                positions = self.positions.gpu[:num_tokens]
+
+            if get_pp_group().is_first_rank:
+                intermediate_tensors = None
+            else:
+                if self.intermediate_tensors is None:
+                    self.intermediate_tensors = self.model.make_empty_intermediate_tensors(
+                        batch_size=self.max_num_tokens,
+                        dtype=self.model_config.dtype,
+                        device=self.device,
+                    )
+
+                intermediate_tensors = self.sync_and_slice_intermediate_tensors(num_tokens, None, False)
+
+            # filter out the valid batch descriptor
+            _cg_mode, batch_descriptor = (
+                self.cudagraph_dispatcher.dispatch(
+                    BatchDescriptor(
+                        num_tokens=num_tokens_after_padding,
+                        uniform_decode=uniform_decode,
+                    )
+                )
+                if not is_profile
+                else (CUDAGraphMode.NONE, None)
+            )
+            if cudagraph_runtime_mode is not None:
+                # we allow forcing NONE when the dispatcher disagrees to support
+                # warm ups for cudagraph capture
+                assert cudagraph_runtime_mode == CUDAGraphMode.NONE or cudagraph_runtime_mode == _cg_mode, (
+                    f"Cudagraph runtime mode mismatch at dummy_run. "
+                    f"Expected {_cg_mode}, but got {cudagraph_runtime_mode}."
+                )
+            else:
+                cudagraph_runtime_mode = _cg_mode
+
+            with (
+                self.maybe_randomize_inputs(input_ids),
+                set_forward_context(
+                    attn_metadata,
+                    self.vllm_config,
+                    num_tokens=num_tokens_after_padding,
+                    num_tokens_across_dp=num_tokens_across_dp,
+                    cudagraph_runtime_mode=cudagraph_runtime_mode,
+                    batch_descriptor=batch_descriptor,
+                ),
+            ):
+                outputs = self.model(
+                    input_ids=input_ids,
+                    positions=positions,
+                    intermediate_tensors=intermediate_tensors,
+                    inputs_embeds=inputs_embeds,
+                    **model_kwargs,
+                )
+
+            if self.use_aux_hidden_state_outputs:
+                hidden_states, _ = outputs
+            else:
+                hidden_states = outputs
+
+            if self.speculative_config and self.speculative_config.use_eagle():
+                assert isinstance(self.drafter, EagleProposer)
+                self.drafter.dummy_run(num_tokens)
+
+        # This is necessary to avoid blocking DP.
+        # For dummy runs, we typically skip EPLB since we don't have any real
+        # requests to process.
+        # However, in DP settings, there may be cases when some DP ranks do
+        # not have any requests to process, so they're executing dummy batches.
+        # In such cases, we still have to trigger EPLB to make sure
+        # ranks execute the rearrangement in synchronization.
+
+        # logit_indices = np.cumsum(num_scheduled_tokens) - 1
+        hidden_states, _ = self.extract_multimodal_outputs(hidden_states)
+        return hidden_states, None
+
+    def profile_run(self) -> None:
+        # Profile with multimodal encoder & encoder cache.
+        if self.supports_mm_inputs:
+            if self.model_config.multimodal_config.skip_mm_profiling:
+                logger.info("Skipping memory profiling for multimodal encoder and encoder cache.")
+            else:
+                mm_budget = self.mm_budget
+                assert mm_budget is not None
+
+                if (encoder_budget := mm_budget.get_encoder_budget()) > 0:
+                    # NOTE: Currently model is profiled with a single non-text
+                    # modality with the max possible input tokens even when
+                    # it supports multiple.
+                    dummy_modality = mm_budget.get_modality_with_max_tokens()
+                    max_mm_items_per_batch = mm_budget.max_items_per_batch_by_modality[dummy_modality]
+
+                    logger.info(
+                        "Encoder cache will be initialized with a budget "
+                        "of %s tokens, and profiled with %s %s items of the maximum feature size.",
+                        encoder_budget,
+                        max_mm_items_per_batch,
+                        dummy_modality,
+                    )
+
+                    # Create dummy batch of multimodal inputs.
+                    batched_dummy_mm_inputs = self._get_mm_dummy_batch(
+                        dummy_modality,
+                        max_mm_items_per_batch,
+                    )
+
+                    # Run multimodal encoder.
+                    dummy_encoder_outputs = self.model.get_multimodal_embeddings(**batched_dummy_mm_inputs)
+
+                    sanity_check_mm_encoder_outputs(
+                        dummy_encoder_outputs,
+                        expected_num_items=max_mm_items_per_batch,
+                    )
+
+                    # NOTE: This happens when encoder cache needs to store
+                    # the embeddings that encoder outputs are scattered onto.
+                    # In this case we create dummy embeddings of size
+                    # (encode_budget, hidden_size) and scatter encoder
+                    # output into it.
+                    encoder_output_shape = dummy_encoder_outputs[0].shape
+                    if encoder_output_shape[0] < encoder_budget:
+                        expanded_outputs = []
+                        for output in dummy_encoder_outputs:
+                            expanded = output.new_zeros((encoder_budget, encoder_output_shape[-1]))
+                            num_tokens = output.shape[0]
+                            expanded[:num_tokens].copy_(output)
+                            expanded_outputs.append(expanded)
+
+                        dummy_encoder_outputs = expanded_outputs
+
+                    # Cache the dummy encoder outputs.
+                    self.encoder_cache["tmp"] = dict(enumerate(dummy_encoder_outputs))
+
+        # Add `is_profile` here to pre-allocate communication buffers
+        hidden_states, _ = self._dummy_run(self.max_num_tokens, is_profile=True)
+        if get_pp_group().is_last_rank:
+            pass
+        self._sync_device()
+        del hidden_states
+        self.encoder_cache.clear()
+        gc.collect()

vllm_omni/worker/gpu_generation_worker.py

@@ -0,0 +1,84 @@
+import gc
+import os
+
+import torch
+from vllm.model_executor import set_random_seed
+from vllm.platforms import current_platform
+from vllm.utils import GiB_bytes, MemorySnapshot
+from vllm.v1.utils import report_usage_stats
+from vllm.v1.worker.gpu_worker import Worker as GPUWorker
+from vllm.v1.worker.gpu_worker import init_worker_distributed_environment
+
+from vllm_omni.worker.gpu_generation_model_runner import GPUGenerationModelRunner
+
+
+class GPUGenerationWorker(GPUWorker):
+    """GPU Worker for Generation model (non-autoregressive waveform generation).
+
+    Usage in stage config:
+        worker_cls: "vllm_omni.worker.gpu_generation_model_runner.GPUGenerationModelRunner"
+    """
+
+    def init_device(self):
+        """Initialize CUDA device and distributed environment."""
+        if self.device_config.device.type == "cuda":
+            # torch.distributed.all_reduce does not free the input tensor until
+            # the synchronization point. This causes the memory usage to grow
+            # as the number of all_reduce calls increases. This env var disables
+            # this behavior.
+            # Related issue:
+            # https://discuss.pytorch.org/t/cuda-allocation-lifetime-for-inputs-to-distributed-all-reduce/191573
+            os.environ["TORCH_NCCL_AVOID_RECORD_STREAMS"] = "1"
+
+            # This env var set by Ray causes exceptions with graph building.
+            os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None)
+            self.device = torch.device(f"cuda:{self.local_rank}")
+            current_platform.set_device(self.device)
+
+            current_platform.check_if_supports_dtype(self.model_config.dtype)
+
+            # Initialize the distributed environment BEFORE taking
+            # memory snapshot
+            # This ensures NCCL buffers are allocated before we measure
+            # available memory
+            init_worker_distributed_environment(
+                self.vllm_config,
+                self.rank,
+                self.distributed_init_method,
+                self.local_rank,
+                current_platform.dist_backend,
+            )
+
+            # Set random seed.
+            set_random_seed(self.model_config.seed)
+
+            # Now take memory snapshot after NCCL is initialized
+            gc.collect()
+            torch.cuda.empty_cache()
+
+            # take current memory snapshot
+            self.init_snapshot = MemorySnapshot()
+            self.requested_memory = self.init_snapshot.total_memory * self.cache_config.gpu_memory_utilization
+            if self.init_snapshot.free_memory < self.requested_memory:
+
+                def GiB(b):
+                    return round(b / GiB_bytes, 2)
+
+                raise ValueError(
+                    f"Free memory on device "
+                    f"({GiB(self.init_snapshot.free_memory)}/"
+                    f"{GiB(self.init_snapshot.total_memory)} GiB) on startup "
+                    f"is less than desired GPU memory utilization "
+                    f"({self.cache_config.gpu_memory_utilization}, "
+                    f"{GiB(self.requested_memory)} GiB). Decrease GPU memory "
+                    f"utilization or reduce GPU memory used by other processes."
+                )
+        else:
+            raise RuntimeError(f"Not support device type: {self.device_config.device}")
+
+        # Construct the model runner
+        self.model_runner: GPUGenerationModelRunner = GPUGenerationModelRunner(self.vllm_config, self.device)
+
+        if self.rank == 0:
+            # If usage stat is enabled, collect relevant info.
+            report_usage_stats(self.vllm_config)

vllm_omni/worker/gpu_model_runner.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Optional, Union, cast
+from typing import TYPE_CHECKING, Any, Optional, Union, cast
 
 import numpy as np
 import torch
@@ -6,15 +6,18 @@ from vllm.config import CUDAGraphMode
 from vllm.distributed.parallel_state import get_pp_group
 from vllm.forward_context import BatchDescriptor, set_forward_context
 from vllm.logger import init_logger
-from vllm.model_executor.models.interfaces import supports_mrope
 from vllm.model_executor.models.interfaces_base import VllmModelForPooling
 from vllm.sampling_params import SamplingType
 from vllm.utils import LazyLoader, cdiv
-from vllm.v1.attention.backends.utils import CommonAttentionMetadata, split_attn_metadata
+from vllm.v1.attention.backends.utils import (
+    CommonAttentionMetadata,
+    split_attn_metadata,
+)
 from vllm.v1.spec_decode.eagle import EagleProposer
 from vllm.v1.worker.gpu_input_batch import CachedRequestState
-from vllm.v1.worker.gpu_model_runner import GPUModelRunner, PerLayerAttnMetadata
+from vllm.v1.worker.gpu_model_runner import GPUModelRunner, IntermediateTensors, PerLayerAttnMetadata
 from vllm.v1.worker.ubatch_splitting import ubatch_split
+from vllm.v1.worker.ubatch_utils import UBatchSlices
 
 if TYPE_CHECKING:
     from vllm.v1.core.sched.output import SchedulerOutput
@@ -30,61 +33,6 @@ logger = init_logger(__name__)
 
 
 class OmniGPUModelRunner(GPUModelRunner):
-    """GPU model runner for omni models with M-RoPE support.
-
-    Extends the base GPUModelRunner to handle multimodal rotary position
-    embeddings (M-RoPE) for processing image, video, and audio inputs
-    with proper positional encoding.
-    """
-
-    def _init_mrope_positions(self, req_state: CachedRequestState):
-        """Initialize M-RoPE positions for multimodal inputs.
-
-        Extracts multimodal feature metadata (image grids, video grids,
-        audio features) and computes M-RoPE positions for proper positional
-        encoding of multimodal tokens.
-
-        Args:
-            req_state: Cached request state containing multimodal features
-
-        Raises:
-            AssertionError: If the model does not support M-RoPE
-        """
-        image_grid_thw = []
-        video_grid_thw = []
-        second_per_grid_ts = []
-        audio_feature_lengths = []
-        use_audio_in_video = False
-        for mm_feature in req_state.mm_features:
-            mm_item = mm_feature.data
-            if mm_item is None:
-                continue
-            mm_input = mm_item.get_data()
-            if (t := mm_input.get("image_grid_thw")) is not None:
-                image_grid_thw.append(t.tolist())
-            if (t := mm_input.get("video_grid_thw")) is not None:
-                video_grid_thw.append(t.tolist())
-            if (t := mm_input.get("second_per_grid_ts")) is not None:
-                second_per_grid_ts.append(t)
-            if (t := mm_input.get("audio_feature_lengths")) is not None:
-                audio_feature_lengths.append(t)
-            # Check for use_audio_in_video
-            use_audio_in_video_value = mm_input.get("use_audio_in_video")
-            if use_audio_in_video_value is not None:
-                use_audio_in_video = bool(use_audio_in_video_value.item())
-
-        assert supports_mrope(self.get_model()), "M-RoPE support is not implemented."
-
-        req_state.mrope_positions, req_state.mrope_position_delta = self.model.get_mrope_input_positions(
-            req_state.prompt_token_ids,
-            hf_config=self.model_config.hf_config,
-            image_grid_thw=image_grid_thw,
-            video_grid_thw=video_grid_thw,
-            second_per_grid_ts=second_per_grid_ts,
-            audio_feature_lengths=audio_feature_lengths,
-            use_audio_in_video=use_audio_in_video,
-        )
-
     def _update_states(self, scheduler_output: "SchedulerOutput") -> None:
         """Update the cached states and the persistent batch with the scheduler
         output.
@@ -600,3 +548,291 @@ class OmniGPUModelRunner(GPUModelRunner):
         logit_indices = np.cumsum(num_scheduled_tokens) - 1
         hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states)
         return hidden_states, hidden_states[logit_indices]
+
+    def _preprocess(
+        self,
+        scheduler_output: "SchedulerOutput",
+        num_scheduled_tokens_np: np.ndarray,
+        intermediate_tensors: Optional[IntermediateTensors] = None,
+        ubatch_slices: Optional[UBatchSlices] = None,
+        num_tokens_after_padding: Optional[torch.Tensor] = None,
+    ) -> tuple[
+        int,
+        int,
+        Optional[torch.Tensor],
+        Optional[torch.Tensor],
+        Optional[torch.Tensor],
+        torch.Tensor,
+        Optional[IntermediateTensors],
+        dict[str, Any],
+        Optional[dict[str, dict]],
+    ]:
+        num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
+        if ubatch_slices:
+            assert num_tokens_after_padding is not None
+            num_input_tokens = int(num_tokens_after_padding[0].item() * 2)
+            self.pad_out_ubatch_slice(ubatch_slices, num_input_tokens)
+        elif ubatch_slices is None:
+            num_input_tokens = self._get_num_input_tokens(num_scheduled_tokens)
+            num_pad, num_tokens_after_padding = self.get_dp_padding(num_input_tokens)
+            num_input_tokens += num_pad
+
+        # _prepare_inputs may reorder the batch, so we must gather multi
+        # modal outputs after that to ensure the correct order
+        per_req_additional_information: Optional[dict[str, dict]] = None
+        if self.supports_mm_inputs and get_pp_group().is_first_rank and not self.model_config.is_encoder_decoder:
+            # Build multimodal inputs and overlay prompt embeds; collect per-request info
+            per_req_additional_information = self._build_mm_inputs_and_overlays(
+                scheduler_output, num_scheduled_tokens, num_scheduled_tokens_np
+            )
+
+            input_ids = self.input_ids.gpu[:num_input_tokens]
+            inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
+            model_kwargs = {
+                **self._init_model_kwargs(num_scheduled_tokens),
+                **self._extract_mm_kwargs(scheduler_output),
+            }
+        elif self.enable_prompt_embeds and get_pp_group().is_first_rank:
+            # Get the input embeddings for the tokens that are not input embeds,
+            # then put them into the appropriate positions.
+            # TODO(qthequartermasterman): Since even when prompt embeds are
+            # enabled, (a) not all requests will use prompt embeds, and (b)
+            # after the initial prompt is processed, the rest of the generated
+            # tokens will be token ids, it is not desirable to have the
+            # embedding layer outside of the CUDA graph all the time. The v0
+            # engine avoids this by "double compiling" the CUDA graph, once
+            # with input_ids and again with inputs_embeds, for all num_tokens.
+            # If a batch only has token ids, then including the embedding layer
+            # in the CUDA graph will be more performant (like in the else case
+            # below).
+            token_ids_idx = self.is_token_ids.gpu[:num_scheduled_tokens].nonzero(as_tuple=False).squeeze(1)
+            # Some tokens ids may need to become embeds
+            if token_ids_idx.numel() > 0:
+                token_ids = self.input_ids.gpu[token_ids_idx]
+                tokens_to_embeds = self.model.get_input_embeddings(input_ids=token_ids)
+                self.inputs_embeds.gpu[token_ids_idx] = tokens_to_embeds
+
+            inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
+            model_kwargs = self._init_model_kwargs(num_input_tokens)
+            input_ids = None
+        else:
+            # For text-only models, we use token ids as input.
+            # While it is possible to use embeddings as input just like the
+            # multimodal models, it is not desirable for performance since
+            # then the embedding layer is not included in the CUDA graph.
+            input_ids = self.input_ids.gpu[:num_input_tokens]
+            inputs_embeds = None
+            model_kwargs = self._init_model_kwargs(num_input_tokens)
+        if self.uses_mrope:
+            positions = self.mrope_positions.gpu[:, :num_input_tokens]
+        else:
+            positions = self.positions.gpu[:num_input_tokens]
+
+        if get_pp_group().is_first_rank:
+            intermediate_tensors = None
+        else:
+            intermediate_tensors = self.sync_and_slice_intermediate_tensors(
+                num_input_tokens, intermediate_tensors, True
+            )
+
+        if self.model_config.is_encoder_decoder and scheduler_output.scheduled_encoder_inputs:
+            encoder_inputs = self._extract_encoder_inputs(scheduler_output)
+            model_kwargs.update(encoder_inputs)
+
+        return (
+            num_scheduled_tokens,
+            num_input_tokens,
+            num_tokens_after_padding,
+            input_ids,
+            inputs_embeds,
+            positions,
+            intermediate_tensors,
+            model_kwargs,
+            per_req_additional_information,
+        )
+
+    def _decode_and_store_request_payloads(self, scheduler_output: "SchedulerOutput") -> None:
+        """Decode per-request prompt_embeds and additional_information for newly
+        scheduled requests and store them to CPU in the request state.
+        This version avoids hard dependency on payload classes by duck-typing."""
+        try:
+            new_reqs = getattr(scheduler_output, "scheduled_new_reqs", [])
+            if not new_reqs:
+                return
+            for nr in new_reqs:
+                req_id = getattr(nr, "req_id", None) or getattr(nr, "request_id", None)
+                if req_id is None:
+                    continue
+                # prompt_embeds
+                payload_pe = getattr(nr, "prompt_embeds", None)
+                pe_cpu = None
+                if payload_pe is not None:
+                    if isinstance(payload_pe, torch.Tensor):
+                        pe_cpu = payload_pe.detach().to("cpu").contiguous()
+                    else:
+                        # Try duck-typing a payload with data/shape/dtype
+                        data = getattr(payload_pe, "data", None)
+                        shape = getattr(payload_pe, "shape", None)
+                        if data is not None and shape is not None:
+                            dt = np.dtype(getattr(payload_pe, "dtype", "float32"))
+                            arr = np.frombuffer(data, dtype=dt)
+                            arr = arr.reshape(shape)
+                            pe_cpu = torch.from_numpy(arr)
+                if pe_cpu is not None and req_id in self.requests:
+                    setattr(self.requests[req_id], "prompt_embeds_cpu", pe_cpu)
+                # additional_information
+                payload_info = getattr(nr, "additional_information", None)
+                if payload_info is not None:
+                    info_dict = {}
+                    if isinstance(payload_info, dict):
+                        info_dict = payload_info
+                    else:
+                        # Try duck-typing a payload with entries, each entry may have
+                        # tensor_data/tensor_dtype/tensor_shape or list_data
+                        entries = getattr(payload_info, "entries", None)
+                        if isinstance(entries, dict):
+                            for k, entry in entries.items():
+                                tensor_data = getattr(entry, "tensor_data", None)
+                                if tensor_data is not None:
+                                    dt = np.dtype(getattr(entry, "tensor_dtype", "float32"))
+                                    arr = np.frombuffer(tensor_data, dtype=dt)
+                                    arr = arr.reshape(getattr(entry, "tensor_shape", ()))
+                                    info_dict[k] = torch.from_numpy(arr)
+                                else:
+                                    info_dict[k] = getattr(entry, "list_data", None)
+                    if info_dict and req_id in self.requests:
+                        setattr(self.requests[req_id], "additional_information_cpu", info_dict)
+        except Exception as e:
+            logger.error(f"Error decoding prompt_embeds / additional_information: {e}")
+
+    def _gather_runtime_additional_information(self) -> list[dict]:
+        """Gather per-request additional_information stored in request state in batch order."""
+        per_req_runtime_info = []
+        for req_id in self.input_batch.req_ids:
+            req_state = self.requests.get(req_id)
+            info = getattr(req_state, "additional_information_cpu", None) if req_state is not None else None
+            per_req_runtime_info.append(info if isinstance(info, dict) else {})
+        return per_req_runtime_info
+
+    def _compute_request_token_spans(self, num_scheduled_tokens_np) -> list[tuple[int, int]]:
+        """Compute (start, end) token spans for each request within the flattened step sequence."""
+        req_token_spans: list[tuple[int, int]] = []
+        for req_index in range(len(self.input_batch.req_ids)):
+            start_offset = int(self.query_start_loc.cpu[req_index])
+            sched_tokens = int(num_scheduled_tokens_np[req_index])
+            req_token_spans.append((start_offset, start_offset + sched_tokens))
+        return req_token_spans
+
+    def _build_model_kwargs_extra(
+        self,
+        per_req_additional_information: Optional[dict[str, dict]],
+        num_scheduled_tokens_np,
+    ) -> dict:
+        """Build extra keyword arguments passed to the model for this step, including:
+        - additional_information_by_req_id: per-request additional information provided by preprocess for this step
+        - runtime_additional_information: per-request additional information stored in request state
+        - request_ids: the request id list in batch order
+        - request_token_spans: token spans per request within the flattened sequence
+        """
+        model_kwargs_extra: dict[str, object] = {}
+        if per_req_additional_information:
+            model_kwargs_extra["additional_information_by_req_id"] = per_req_additional_information
+        try:
+            model_kwargs_extra["runtime_additional_information"] = self._gather_runtime_additional_information()
+            model_kwargs_extra["request_ids"] = self.input_batch.req_ids
+            model_kwargs_extra["request_token_spans"] = self._compute_request_token_spans(num_scheduled_tokens_np)
+        except Exception:
+            pass
+        return model_kwargs_extra
+
+    def _process_additional_information_updates(self, multimodal_outputs: object) -> None:
+        """Process model-provided per-request additional_information updates and merge into request state."""
+        try:
+            if not isinstance(multimodal_outputs, dict):
+                return
+            if (
+                "additional_information_update" not in multimodal_outputs
+                and "additional_information_update_by_req_id" not in multimodal_outputs
+            ):
+                return
+            updates_list = multimodal_outputs.get("additional_information_update")
+            if isinstance(updates_list, list):
+                for idx, upd in enumerate(updates_list):
+                    if not isinstance(upd, dict) or idx >= len(self.input_batch.req_ids):
+                        continue
+                    req_id = self.input_batch.req_ids[idx]
+                    self._merge_additional_information_update(req_id, upd)
+            updates_map = multimodal_outputs.get("additional_information_update_by_req_id")
+            if isinstance(updates_map, dict):
+                for req_id, upd in updates_map.items():
+                    if not isinstance(upd, dict):
+                        continue
+                    if req_id not in self.requests:
+                        continue
+                    self._merge_additional_information_update(req_id, upd)
+        except Exception as e:
+            logger.error(
+                f"Error merging for requests:{self.input_batch.req_ids} "
+                f"additional information update: {e}, with the multimodal_outputs "
+                f"as {multimodal_outputs}"
+            )
+
+    def _build_mm_inputs_and_overlays(
+        self,
+        scheduler_output: "SchedulerOutput",
+        num_scheduled_tokens: int,
+        num_scheduled_tokens_np: np.ndarray,
+    ) -> dict[str, dict]:
+        """Run MM encoder, fill scheduled inputs_embeds, and overlay prompt embeds;
+        collect per-request additional_information for the prefill portion."""
+        # Run the multimodal encoder if any.
+        self._execute_mm_encoder(scheduler_output)
+        mm_embeds = self._gather_mm_embeddings(scheduler_output)
+        # Always use embeddings as inputs for MM models
+        inputs_embeds_scheduled = self.model.get_input_embeddings(
+            input_ids=self.input_ids.gpu[:num_scheduled_tokens],
+            multimodal_embeddings=mm_embeds or None,
+        )
+        self.inputs_embeds.gpu[:num_scheduled_tokens].copy_(inputs_embeds_scheduled)
+        # Reset per-step additional information collector (deprecated concat path)
+        if hasattr(self, "_forward_additional_information"):
+            self._forward_additional_information = None
+        # Overlay and collect per-request info
+        return self._collect_additional_information_for_prefill(num_scheduled_tokens_np)
+
+    def _collect_additional_information_for_prefill(
+        self,
+        num_scheduled_tokens_np: np.ndarray,
+    ) -> dict[str, dict]:
+        """Overlay per-request prompt_embeds for the prefill portion and collect
+        additional_information slices for this step. Returns a map req_id -> dict."""
+        per_req_additional_information: dict[str, dict] = {}
+        for req_index, req_id in enumerate(self.input_batch.req_ids):
+            req_state = self.requests[req_id]
+            pe_cpu = getattr(req_state, "prompt_embeds_cpu", None)
+            addi_cpu = getattr(req_state, "additional_information_cpu", None)
+            num_computed_tokens = int(self.input_batch.num_computed_tokens_cpu[req_index])
+            prompt_len = len(req_state.prompt_token_ids)
+            prompt_remaining = max(0, prompt_len - num_computed_tokens)
+            sched_tokens = int(num_scheduled_tokens_np[req_index])
+            overlay_len = min(sched_tokens, prompt_remaining)
+            if overlay_len <= 0 and not (isinstance(addi_cpu, dict) and addi_cpu):
+                continue
+            if overlay_len > 0 and pe_cpu is not None:
+                src = pe_cpu[num_computed_tokens : num_computed_tokens + overlay_len].to(
+                    dtype=self.dtype, device=self.device, non_blocking=True
+                )
+                start_offset = int(self.query_start_loc.cpu[req_index])
+                self.inputs_embeds[start_offset : start_offset + overlay_len].copy_(src)
+            if addi_cpu is not None and isinstance(addi_cpu, dict):
+                req_info: dict[str, object] = {}
+                for k, v in addi_cpu.items():
+                    if isinstance(v, torch.Tensor):
+                        req_info[k] = v.detach().to("cpu").contiguous()
+                    elif isinstance(v, list):
+                        req_info[k] = v
+                    else:
+                        req_info[k] = v
+                if req_info:
+                    per_req_additional_information[req_id] = req_info
+        return per_req_additional_information