Revert "[Perf][Qwen3-TTS] Fuse QKV and gate_up projections in code predictor …"

This reverts commit 6fdc4ca62d.

.buildkite/cuda/test-nightly.yml

@@ -427,6 +427,12 @@ steps:
           - pytest -s -v tests/e2e/accuracy/test_diffusers_backend_similarity.py -k '2v_matches_diffusers' -m full_model --run-level full_model
         mirror_hardwares: h100_2
 
+      - label: ":full_moon: Diffusion X2V · MiniMax H3 T2VA Accuracy Test"
+        timeout_in_minutes: 180
+        commands:
+          - pytest -s -v tests/e2e/accuracy/minimax_h3/test_minimax_h3_t2va_similarity.py -m full_model --run-level full_model
+        mirror_hardwares: h100_4
+
       - label: ":full_moon: Diffusion X2V · Other Accuracy Test"
         timeout_in_minutes: 180
         commands:
@@ -477,12 +483,18 @@ steps:
           - pytest -sv tests/diffusion/offloader -m "full_model and cuda"
         mirror_hardwares: l4_1
 
-      - label: ":full_moon: Diffusion · Quantization Test with H100"
+      - label: ":full_moon: Diffusion · Quantization Test with H100 · Single-GPU"
         timeout_in_minutes: 60
         commands:
           - pytest -sv tests/diffusion/quantization -m 'full_model and cuda and H100' --run-level "full_model"
         mirror_hardwares: h100_1
 
+      - label: ":full_moon: Diffusion · MiniMax H3 FP8 Quality Test with H100 · 2-GPU"
+        timeout_in_minutes: 60
+        commands:
+          - pytest -sv tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py -m 'full_model and cuda and H100' --run-level "full_model"
+        mirror_hardwares: h100_2
+
       - label: ":full_moon: Diffusion · Quantization Test with L4"
         timeout_in_minutes: 60
         commands:

README.md

@@ -16,14 +16,10 @@ Easy, fast, and cheap omni-modality model serving for everyone
 ---
 
 *Latest News* 🔥
+- [2026/08] We released [0.26.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.26.0) - aligned with the vLLM 0.26 release line, featuring [MiniMax H3](recipes/MiniMaxAI/MiniMax-H3.md) joint video/audio generation, an experimental full-duplex realtime runtime for [MiniCPM-o 4.5](recipes/OpenBMB/MiniCPM-o-4_5.md), distributed layerwise diffusion offload, and broader model, hardware, streaming, TTS, and quantization support.
 - [2026/07] We released [0.24.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.24.0) - aligned with the vLLM 0.24 release line, expanding production-ready coverage across TTS, speech, diffusion, image/video generation, and robot-policy serving, with major Omni stage runtime refactoring, diffusion request-level batching, async output materialization, quantization/cache/memory improvements, and broad CUDA/ROCm/XPU/NPU support.
-- [2026/06] We released [0.22.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.22.0) - an **omnimodal world-model** release aligned with vLLM 0.22, featuring [Nvidia Cosmos3](recipes/cosmos3/Cosmos3-Nano.md)/DreamZero world model support, expanded quantization coverage across Blackwell/NPU/XPU, TTS production improvements, new models including MiniCPM-o 4.5, MOSS-TTS, and Lance, plus RL integration with [VeRL-Omni](https://github.com/verl-project/verl-omni).
-- [2026/05] We released [0.20.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.20.0) - refreshes the serving/runtime stack for large-scale omni workloads, and improves diffusion model performance, quantization, and hardware readiness across CUDA, ROCm, MUSA, NPU, and XPU backends.
-- [2026/03] We released [0.18.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.18.0) - strengthens the core runtime through a large entrypoint refactor and scheduler/runtime cleanups, expands unified quantization and diffusion execution, broadens multimodal model coverage, and improves production readiness across audio, omni, image, video, RL, and multi-platform deployments.
+- [2026/06] Starting with [0.14.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.14.0), vLLM-Omni publishes a stable release aligned with every even-numbered upstream vLLM minor version. [0.16.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.16.0), [0.18.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.18.0), [0.20.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.20.0), and [0.22.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.22.0) continued this cadence, expanding omni and world-model support with [NVIDIA Cosmos3](recipes/cosmos3/Cosmos3-Nano.md) and DreamZero, adding models such as MiniCPM-o 4.5, MOSS-TTS, and Lance, and advancing TTS, diffusion, distributed execution, quantization, RL integration through [VeRL-Omni](https://github.com/verl-project/verl-omni), and CUDA/ROCm/MUSA/NPU/XPU coverage.
 - [2026/03] Check out our first public [project deepdive](https://youtu.be/sgwNfsNnR9I) at the vLLM Hong Kong Meetup!
-- [2026/03] **[vllm-omni-skills](https://github.com/hsliuustc0106/vllm-omni-skills)** is a community-driven collection of AI assistant skills that help developers work with vLLM-Omni more effectively. These skills can be used with popular agentic AI coding assistants like **Cursor IDE**, **Claude**, **Codex**, and more.
-- [2026/02] We released [0.16.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.16.0) - A major alignment + capability release that rebases onto **upstream vLLM v0.16.0** and significantly expands performance, distributed execution, and production readiness across **Qwen3-Omni / Qwen3-TTS**, **Bagel**, **MiMo-Audio**, **GLM-Image** and the **Diffusion (DiT) image/video stack**—while also improving platform coverage (CUDA / ROCm / NPU / XPU), CI quality, and documentation.
-- [2026/02] We released [0.14.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.14.0) - This is the first **stable release** of vLLM-Omni that expands Omni’s diffusion / image-video generation and audio / TTS stack, improves distributed execution and memory efficiency, and broadens platform/backend coverage (GPU/ROCm/NPU/XPU). It also brings meaningful upgrades to serving APIs, profiling & benchmarking, and overall stability. Please check our latest [paper](https://arxiv.org/abs/2602.02204) for architecture design and performance results.
 - [2025/11] vLLM community officially released [vllm-project/vllm-omni](https://github.com/vllm-project/vllm-omni) in order to support omni-modality models serving.
 
 ---
@@ -55,12 +51,13 @@ vLLM-Omni is flexible and easy to use with:
 - Tensor, pipeline, data and expert parallelism support for distributed inference
 - Streaming outputs
 - OpenAI-compatible API server
+- Full-duplex realtime serving with streaming audio input and output (experimental)
 
 vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, including:
 
-- **Omni-modality models** (e.g. Qwen3-Omni, Cosmos3, HunyuanImage, BAGEL)
+- **Omni-modality models** (e.g. Qwen3-Omni, MiniCPM-o 4.5, Cosmos3, HunyuanImage, BAGEL)
 - **TTS models** (e.g. Qwen3-TTS, VoxCPM2, Ming-Omni-TTS, CosyVoice3)
-- **Diffusion models** — image, video, and audio generation (e.g. Qwen-Image, Wan2.2, FLUX)
+- **Diffusion models** — image, video, and audio generation (e.g. MiniMax H3, Qwen-Image, Wan2.2, FLUX)
 - **Robot-policy and action models** (e.g. GR00T-N1.7, DreamZero-DROID, InternVLA-A1, Cosmos3 action policy)
 
 ## Getting Started

apps/ComfyUI-vLLM-Omni/README.md

@@ -47,7 +47,7 @@ If no, check your shell running the ComfyUI process. There may be some error mes
 This extension offers the following nodes based on the output modalities (at **ComfyUI sidebar -> Node Library**):
 
 - **Generate Image** for text-to-image and image-to-image tasks
-- **Generate Video** for text-to-video and image-to-video tasks
+- **Generate Video** for text-to-video, first-frame/image-to-video, and reference-conditioned video
 - **Multimodality Understanding** for multimodality-to-text and multimodality-to-audio tasks
 - **TTS** and **TTS Voice Clone** for TTS tasks
 
@@ -105,7 +105,7 @@ You can configure per-stage sampling parameters for multi-stage models.
 > [!TIP]
 > The node automatically choose text-to-image or image-to-image API endpoints depending on whether you connect an image input or not.
 
-### Text-to-video and image-to-video generation (e.g., Wan)
+### Text-to-video and image-to-video generation (e.g., Wan, MiniMax-H3)
 
 (Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Video Generation**)
 
@@ -117,7 +117,11 @@ You can configure per-stage sampling parameters for multi-stage models.
 </p>
 
 > [!TIP]
-> The node automatically choose text-to-video or image-to-video API endpoints depending on whether you connect an image input or not.
+> Connect a **frame** image for first-frame / image-to-video (e.g., Wan I2V, MiniMax-H3 FL2VA).
+>
+> For reference-conditioned generation (MiniMax-H3 Ref2VA), connect a **Video References** node instead.
+>
+> Do not use `frame` and `references` together. Task routing is automatic from which inputs you connect.
 
 ### TTS (e.g., Qwen TTS series)
 

apps/ComfyUI-vLLM-Omni/__init__.py

@@ -15,11 +15,13 @@ from .comfyui_vllm_omni.nodes import (
     VLLMOmniDiffusionSampling,
     VLLMOmniGenerateImage,
     VLLMOmniGenerateVideo,
+    VLLMOmniMiniMaxH3Params,
     VLLMOmniQwenTTSParams,
     VLLMOmniRemoteLoRA,
     VLLMOmniSamplingParamsList,
     VLLMOmniTTS,
     VLLMOmniUnderstanding,
+    VLLMOmniVideoReferences,
     VLLMOmniVoiceClone,
     VLLMOmniWanParams,
 )
@@ -32,6 +34,7 @@ NODE_CLASS_MAPPINGS = {
     "VLLMOmniUnderstanding": VLLMOmniUnderstanding,
     "VLLMOmniTTS": VLLMOmniTTS,
     "VLLMOmniVoiceClone": VLLMOmniVoiceClone,
+    "VLLMOmniVideoReferences": VLLMOmniVideoReferences,
     # === Params ===
     "VLLMOmniARSampling": VLLMOmniARSampling,
     "VLLMOmniDiffusionSampling": VLLMOmniDiffusionSampling,
@@ -39,6 +42,7 @@ NODE_CLASS_MAPPINGS = {
     "VLLMOmniRemoteLoRA": VLLMOmniRemoteLoRA,
     "VLLMOmniQwenTTSParams": VLLMOmniQwenTTSParams,
     "VLLMOmniWanParams": VLLMOmniWanParams,
+    "VLLMOmniMiniMaxH3Params": VLLMOmniMiniMaxH3Params,
 }
 
 # A dictionary that contains the friendly/humanly readable titles for the nodes
@@ -49,6 +53,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "VLLMOmniUnderstanding": "Multimodality Understanding",
     "VLLMOmniTTS": "TTS (Text to Speech)",
     "VLLMOmniVoiceClone": "TTS Voice Cloning",
+    "VLLMOmniVideoReferences": "Video References",
     # === Params ===
     "VLLMOmniARSampling": "AR Sampling Params",
     "VLLMOmniDiffusionSampling": "Diffusion Sampling Params",
@@ -56,6 +61,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "VLLMOmniRemoteLoRA": "LoRA",
     "VLLMOmniQwenTTSParams": "Qwen TTS Params",
     "VLLMOmniWanParams": "Wan Video Params",
+    "VLLMOmniMiniMaxH3Params": "MiniMax-H3 Video Params",
 }
 
 WEB_DIRECTORY = "./web"

apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/nodes.py

@@ -10,7 +10,9 @@ from .utils.types import (
     AudioFormat,
     AutoregressionSamplingParams,
     DiffusionSamplingParams,
+    MiniMaxH3ModelSpecificParams,
     QwenTTSModelSpecificParams,
+    VideoReferences,
     WanModelSpecificParams,
 )
 from .utils.validators import (
@@ -163,7 +165,8 @@ class VLLMOmniGenerateVideo(_VLLMOmniGenerateBase):
                 "num_frames": ("INT", {"default": 41, "min": 1}),
             },
             "optional": {
-                "image": ("IMAGE",),
+                "frame": ("IMAGE",),
+                "references": ("VIDEO_REFERENCES",),
                 "sampling_params": ("SAMPLING_PARAMS",),
                 "lora": ("REMOTE_LORA",),
                 "model_params": ("VIDEO_PARAMS",),
@@ -174,6 +177,15 @@ class VLLMOmniGenerateVideo(_VLLMOmniGenerateBase):
     RETURN_NAMES = ("video",)
     FUNCTION = "generate"
 
+    @classmethod
+    def VALIDATE_INPUTS(cls, url, model, frame=None, references=None, **_kwargs) -> str | Literal[True]:
+        base = super().VALIDATE_INPUTS(url, model)
+        if base is not True:
+            return base
+        if frame is not None and references is not None:
+            return "Provide only one of frame or references, not both."
+        return True
+
     async def generate(
         self,
         url: str,
@@ -184,7 +196,8 @@ class VLLMOmniGenerateVideo(_VLLMOmniGenerateBase):
         fps: int,
         num_frames: int,
         negative_prompt: str | None = None,
-        image: torch.Tensor | None = None,
+        frame: torch.Tensor | None = None,
+        references: dict | None = None,
         sampling_params: dict | list[dict] | None = None,
         model_params: dict | None = None,
         lora: dict | None = None,
@@ -207,14 +220,16 @@ class VLLMOmniGenerateVideo(_VLLMOmniGenerateBase):
 
         if sampling_params is not None:
             sampling_params.pop("type", None)  # internal fields
-        if model_params is not None:
-            model_params.pop("type", None)  # internal fields
+        # model_params["type"] is also an internal field, but is used in api_client
+        # if model_params is not None:
+        #     model_params.pop("type", None)  # internal fields
 
         client = VLLMOmniClient(url)
         output = await client.generate_video(
             model=model,
             prompt=prompt,
-            image=image,  # image present => i2v, absent => t2v
+            frame=frame,  # frame present => fl2va / Wan I2V
+            references=references,
             width=width,
             height=height,
             num_frames=num_frames,
@@ -734,3 +749,78 @@ class VLLMOmniWanParams:
 
     def get_params(self, **kwargs):
         return (WanModelSpecificParams(kwargs),)
+
+
+class VLLMOmniMiniMaxH3Params:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "audio_flow_shift": (
+                    "FLOAT",
+                    {"default": 3.0, "min": 0.0, "max": 100.0, "step": 0.1},
+                ),
+                "flow_shift": (
+                    "FLOAT",
+                    {"default": 12.0, "min": 0.0, "max": 100.0, "step": 0.1},
+                ),
+            }
+        }
+
+    RETURN_TYPES = ("VIDEO_PARAMS",)
+    RETURN_NAMES = ("MiniMax-H3 video params",)
+    FUNCTION = "get_params"
+    CATEGORY = "vLLM-Omni/Video Params"
+
+    def get_params(self, **kwargs):
+        params = MiniMaxH3ModelSpecificParams(kwargs)
+        params["type"] = "minimax_h3"
+        return (params,)
+
+
+class VLLMOmniVideoReferences:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {},
+            "optional": {
+                "image_1": ("IMAGE",),
+                "image_2": ("IMAGE",),
+                "audio_1": ("AUDIO",),
+                "audio_2": ("AUDIO",),
+                "video_1": ("VIDEO",),
+                "video_2": ("VIDEO",),
+            },
+        }
+
+    RETURN_TYPES = ("VIDEO_REFERENCES",)
+    RETURN_NAMES = ("references",)
+    FUNCTION = "get_references"
+    CATEGORY = "vLLM-Omni"
+
+    def get_references(
+        self,
+        image_1: torch.Tensor | None = None,
+        image_2: torch.Tensor | None = None,
+        audio_1: AudioInput | None = None,
+        audio_2: AudioInput | None = None,
+        video_1: VideoInput | None = None,
+        video_2: VideoInput | None = None,
+        **kwargs,
+    ):
+        if kwargs:
+            logger.info("Uncaught kwargs: %s", kwargs)
+        refs = VideoReferences()
+        if image_1 is not None:
+            refs["image_1"] = image_1
+        if image_2 is not None:
+            refs["image_2"] = image_2
+        if audio_1 is not None:
+            refs["audio_1"] = audio_1
+        if audio_2 is not None:
+            refs["audio_2"] = audio_2
+        if video_1 is not None:
+            refs["video_1"] = video_1
+        if video_2 is not None:
+            refs["video_2"] = video_2
+        return (refs,)

apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/api_client.py

@@ -24,6 +24,7 @@ from .format import (
     image_tensor_to_base64,
     image_tensor_to_png_bytes,
     video_to_base64,
+    video_to_bytes,
 )
 from .logger import get_logger, pretty_printer
 from .models import lookup_model_spec
@@ -63,7 +64,11 @@ async def url_bytes(session: aiohttp.ClientSession, url: str, verb: str = "get",
 
 class VLLMOmniClient:
     def __init__(
-        self, base_url: str, timeout: float | None = None, poll_interval: float = 5.0, max_poll_duration: float = 60 * 5
+        self,
+        base_url: str,
+        timeout: float | None = None,
+        poll_interval: float = 5.0,
+        max_poll_duration: float = 60 * 30,
     ):
         self.base_url = base_url
         self.timeout = aiohttp.ClientTimeout(total=timeout)
@@ -261,12 +266,17 @@ class VLLMOmniClient:
         num_frames: int,
         fps: int,
         negative_prompt: str | None = None,
-        image: torch.Tensor | None = None,
+        frame: torch.Tensor | None = None,
+        references: dict | None = None,
         sampling_params: dict | None = None,
         model_params: dict | None = None,
         lora: dict | None = None,
-        **extra_body,
+        **extra_params,
     ) -> VideoInput:
+        if frame is not None and references is not None:
+            raise ValueError("Provide only one of frame or references, not both.")
+
+        # === regular payload fields ===
         form = aiohttp.FormData()
         form.add_field("model", model)
         form.add_field("prompt", prompt)
@@ -279,23 +289,91 @@ class VLLMOmniClient:
         if sampling_params is not None:
             for k, v in sampling_params.items():
                 form.add_field(k, str(v))
-        if model_params is not None:
-            for k, v in model_params.items():
-                form.add_field(k, str(v))
         if lora is not None:
             form.add_field("lora", json.dumps(lora, ensure_ascii=False))
-        if extra_body:
-            form.add_field("extra_body", json.dumps(extra_body, ensure_ascii=False))
 
-        if image is not None:
+        # === multimodal inputs (first-last-frames, references, etc.) ===
+        input_reference_image: torch.Tensor | None = None
+        audio_reference: AudioInput | None = None
+        reference_videos: list[VideoInput] = []
+        video_task: str | None = None
+
+        if frame is not None:
+            input_reference_image = frame
+            video_task = "fl2va"
+        elif references is not None:
+            images = [references[k] for k in ("image_1", "image_2") if k in references and references[k] is not None]
+            audios = [references[k] for k in ("audio_1", "audio_2") if k in references and references[k] is not None]
+            videos = [references[k] for k in ("video_1", "video_2") if k in references and references[k] is not None]
+
+            if not images and not audios and not videos:
+                raise ValueError("references is empty; connect at least one image/audio/video.")
+
+            if videos:
+                if images or audios:
+                    raise ValueError(
+                        "Video references cannot be combined with image or audio references. "
+                        "Connect only video_1/video_2 for multi-video Ref2VA."
+                    )
+                reference_videos = videos
+                video_task = "ref2va"
+            elif len(images) == 1 and len(audios) == 1:
+                input_reference_image = images[0]
+                audio_reference = audios[0]
+                video_task = "ref2va"
+            else:
+                raise ValueError(
+                    "Invalid references combination. Supported modes: "
+                    "(1) one or more videos only, or (2) exactly one image and one audio."
+                )
+        else:
+            video_task = "t2va"
+
+        if input_reference_image is not None:
             image_filename = "image.png"  # Required for multipart form
             form.add_field(
                 "input_reference",
-                image_tensor_to_png_bytes(image, image_filename),
+                image_tensor_to_png_bytes(input_reference_image, image_filename),
                 filename=image_filename,
                 content_type="image/png",
             )
 
+        if audio_reference is not None:
+            form.add_field(
+                "audio_reference",
+                json.dumps({"audio_url": audio_to_base64(audio_reference)}, ensure_ascii=False),
+            )
+
+        for idx, video in enumerate(reference_videos, start=1):
+            video_filename = f"reference_{idx}.mp4"
+            form.add_field(
+                "input_references",
+                video_to_bytes(video, video_filename),
+                filename=video_filename,
+                content_type="video/mp4",
+            )
+
+        # === model specific params. Either use a specialized builder, or add flattened fields as-is ===
+        if model_params is not None:
+            model_params = dict(model_params)
+            model_params.pop("type", None)
+
+        spec, _ = lookup_model_spec(model)
+        params_builder = spec.get("params_builder") if spec else None
+        if params_builder is not None:
+            form_fields = params_builder(
+                model_params or {},
+                extra_params={**extra_params, "task": video_task},
+            )
+            for k, v in form_fields.items():
+                form.add_field(k, v if isinstance(v, str) else str(v))
+        else:
+            if model_params is not None:
+                for k, v in model_params.items():
+                    form.add_field(k, str(v))
+            if extra_params:
+                form.add_field("extra_params", json.dumps(extra_params, ensure_ascii=False))
+
         async with aiohttp.ClientSession(timeout=self.timeout) as session:
             # Start the video generation job
             url = f"{self.base_url}/videos"

apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/models.py

@@ -1,7 +1,12 @@
+import json
 import re
+from typing import Any
 
+from .logger import get_logger
 from .types import Modality, ModelMode, Spec
 
+logger = get_logger(__name__)
+
 
 def _bagel_payload_preprocessor(payload: dict) -> dict:
     try:
@@ -29,6 +34,30 @@ def _qwen25_payload_preprocessor(payload: dict) -> dict:
     return payload
 
 
+def _minimaxh3_params_builder(
+    model_params: dict[str, Any],
+    *,
+    extra_params: dict[str, Any],
+) -> dict[str, Any]:
+    """Build multipart form fields for MiniMax-H3 from model_params + routed task."""
+    params = dict(model_params)
+    params.pop("type", None)
+    form_fields: dict[str, Any] = {}
+    merged_extra_params: dict[str, Any] = dict(extra_params or {})
+
+    if "flow_shift" in params:
+        form_fields["flow_shift"] = params.pop("flow_shift")
+    for key in ("audio_flow_shift",):
+        if key in params:
+            merged_extra_params[key] = params.pop(key)
+    if params:
+        logger.warning("Unused MiniMax-H3 model params ignored: %s", sorted(params))
+
+    if merged_extra_params:
+        form_fields["extra_params"] = json.dumps(merged_extra_params, ensure_ascii=False)
+    return form_fields
+
+
 _MODEL_PIPELINE_SPECS: dict[str, Spec] = {
     r"BAGEL-7B-MoT": {
         "stages": [
@@ -71,6 +100,21 @@ _MODEL_PIPELINE_SPECS: dict[str, Spec] = {
             }
         ],
     },
+    r"MiniMax-H3(/FL2VA|/Ref2VA)?": {
+        "stages": ["diffusion"],
+        "modes": [
+            {
+                "mode": ModelMode.VIDEO_GENERATION,
+                "input_modalities": [
+                    Modality.TEXT,
+                    Modality.IMAGE,
+                    Modality.VIDEO,
+                    Modality.AUDIO,
+                ],
+            }
+        ],
+        "params_builder": _minimaxh3_params_builder,
+    },
 }
 # Convert dict keys to regex patterns
 MODEL_PIPELINE_SPECS: dict[re.Pattern, Spec] = {}
@@ -80,12 +124,10 @@ del _MODEL_PIPELINE_SPECS
 
 
 def lookup_model_spec(model: str) -> tuple[Spec | None, str | None]:
-    try:
-        last_component = model.rstrip("/").rsplit("/", 1)[-1]
-    except IndexError:
-        last_component = model
+    normalized = model.rstrip("/").lstrip("/")
+    lookup_key = normalized.split("/", 1)[-1]
     for pattern, spec in MODEL_PIPELINE_SPECS.items():
-        if pattern.search(last_component):
+        if pattern.search(lookup_key):
             return spec, pattern.pattern
     return None, None
 
@@ -101,6 +143,9 @@ if __name__ == "__main__":
         "Qwen/Qwen3-Omni-30B-A3B-Instruct",
         "Custom/Path/UnknownModel-Instruct",
         "Not/Matching/Anything",
+        "/path/to/MiniMax-H3/FL2VA",
+        "/path/to/MiniMax-H3/Ref2VA",
+        "MiniMaxAI/MiniMax-H3",
     ]
 
     test_payload = {"messages": [{"role": "user", "content": "prompt"}]}
@@ -112,7 +157,13 @@ if __name__ == "__main__":
             if preprocessor := spec.get("payload_preprocessor"):
                 result = preprocessor(test_payload)
                 print(f"✓ {path:<40} → {result}")
+            elif params_builder := spec.get("params_builder"):
+                result = params_builder(
+                    {"audio_flow_shift": 3.0, "flow_shift": 12.0, "type": "minimax_h3"},
+                    extra_params={"task": "t2va"},
+                )
+                print(f"✓ {path:<40} → {result}")
             else:
-                print(f"✓ {path:<40} → No preprocessor")
+                print(f"✓ {path:<40} → No preprocessor/params_builder")
         else:
             print(f"✗ {path:<40} → No match")

apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/types.py

@@ -23,6 +23,14 @@ class WanModelSpecificParams(dict):
     pass
 
 
+class MiniMaxH3ModelSpecificParams(dict):
+    pass
+
+
+class VideoReferences(dict):
+    pass
+
+
 class ModelMode(Enum):
     IMAGE_GENERATION = auto()
     VIDEO_GENERATION = auto()
@@ -43,9 +51,11 @@ class ModelModeSpec(TypedDict):
 
 
 PayloadPreprocessor: TypeAlias = Callable[[dict[str, Any]], dict[str, Any]]
+ParamsBuilder: TypeAlias = Callable[..., dict[str, Any]]
 
 
 class Spec(TypedDict):
     stages: list[Literal["diffusion", "autoregression"]]
     modes: list[ModelModeSpec]
     payload_preprocessor: NotRequired[PayloadPreprocessor]
+    params_builder: NotRequired[ParamsBuilder]

apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni Video Generation.json

@@ -1,513 +1 @@
-{
-    "id": "3aea1d08-7de5-4997-ba71-b2f66ad72a21",
-    "revision": 0,
-    "last_node_id": 14,
-    "last_link_id": 28,
-    "nodes": [
-        {
-            "id": 10,
-            "type": "MarkdownNote",
-            "pos": [
-                227.99306819575278,
-                -231.24143306069843
-            ],
-            "size": [
-                240.3326478922474,
-                136.3505820791881
-            ],
-            "flags": {},
-            "order": 0,
-            "mode": 0,
-            "inputs": [],
-            "outputs": [],
-            "title": "Note: Task and Input",
-            "properties": {},
-            "widgets_values": [
-                "vLLM-Omni nodes are categorized based on the output modality. The \"Generate Video\" node supports both text-to-video generation and image-to-video generation. The node will create corresponding payload depending on whether an input image is present or not."
-            ],
-            "color": "#432",
-            "bgcolor": "#000"
-        },
-        {
-            "id": 4,
-            "type": "LoadImage",
-            "pos": [
-                496.31859627609606,
-                -229.71277089860084
-            ],
-            "size": [
-                270,
-                314
-            ],
-            "flags": {},
-            "order": 1,
-            "mode": 0,
-            "inputs": [
-                {
-                    "localized_name": "image",
-                    "name": "image",
-                    "type": "COMBO",
-                    "widget": {
-                        "name": "image"
-                    },
-                    "link": null
-                },
-                {
-                    "localized_name": "choose file to upload",
-                    "name": "upload",
-                    "type": "IMAGEUPLOAD",
-                    "widget": {
-                        "name": "upload"
-                    },
-                    "link": null
-                }
-            ],
-            "outputs": [
-                {
-                    "localized_name": "IMAGE",
-                    "name": "IMAGE",
-                    "type": "IMAGE",
-                    "links": [
-                        25
-                    ]
-                },
-                {
-                    "localized_name": "MASK",
-                    "name": "MASK",
-                    "type": "MASK",
-                    "links": null
-                }
-            ],
-            "properties": {
-                "Node name for S&R": "LoadImage"
-            },
-            "widgets_values": [
-                "cute-cat.jpg",
-                "image"
-            ]
-        },
-        {
-            "id": 8,
-            "type": "VLLMOmniDiffusionSampling",
-            "pos": [
-                465.5140218220927,
-                148.80072646828987
-            ],
-            "size": [
-                284.205078125,
-                226
-            ],
-            "flags": {},
-            "order": 2,
-            "mode": 0,
-            "inputs": [
-                {
-                    "localized_name": "n",
-                    "name": "n",
-                    "type": "INT",
-                    "widget": {
-                        "name": "n"
-                    },
-                    "link": null
-                },
-                {
-                    "localized_name": "num_inference_steps",
-                    "name": "num_inference_steps",
-                    "type": "INT",
-                    "widget": {
-                        "name": "num_inference_steps"
-                    },
-                    "link": null
-                },
-                {
-                    "localized_name": "guidance_scale",
-                    "name": "guidance_scale",
-                    "type": "FLOAT",
-                    "widget": {
-                        "name": "guidance_scale"
-                    },
-                    "link": null
-                },
-                {
-                    "localized_name": "true_cfg_scale",
-                    "name": "true_cfg_scale",
-                    "type": "FLOAT",
-                    "widget": {
-                        "name": "true_cfg_scale"
-                    },
-                    "link": null
-                },
-                {
-                    "localized_name": "vae_use_slicing",
-                    "name": "vae_use_slicing",
-                    "type": "BOOLEAN",
-                    "widget": {
-                        "name": "vae_use_slicing"
-                    },
-                    "link": null
-                },
-                {
-                    "localized_name": "vae_use_tiling",
-                    "name": "vae_use_tiling",
-                    "type": "BOOLEAN",
-                    "widget": {
-                        "name": "vae_use_tiling"
-                    },
-                    "link": null
-                },
-                {
-                    "localized_name": "seed",
-                    "name": "seed",
-                    "type": "INT",
-                    "widget": {
-                        "name": "seed"
-                    },
-                    "link": null
-                }
-            ],
-            "outputs": [
-                {
-                    "localized_name": "diffusion sampling params",
-                    "name": "diffusion sampling params",
-                    "type": "SAMPLING_PARAMS",
-                    "links": [
-                        27
-                    ]
-                }
-            ],
-            "properties": {
-                "Node name for S&R": "VLLMOmniDiffusionSampling"
-            },
-            "widgets_values": [
-                1,
-                50,
-                1,
-                1,
-                false,
-                false,
-                1992,
-                "randomize"
-            ]
-        },
-        {
-            "id": 14,
-            "type": "VLLMOmniWanParams",
-            "pos": [
-                518.6005521845747,
-                428.4455663719664
-            ],
-            "size": [
-                270,
-                106
-            ],
-            "flags": {},
-            "order": 3,
-            "mode": 0,
-            "inputs": [
-                {
-                    "localized_name": "guidance_scale_2",
-                    "name": "guidance_scale_2",
-                    "type": "FLOAT",
-                    "widget": {
-                        "name": "guidance_scale_2"
-                    },
-                    "link": null
-                },
-                {
-                    "localized_name": "boundary_ratio",
-                    "name": "boundary_ratio",
-                    "type": "FLOAT",
-                    "widget": {
-                        "name": "boundary_ratio"
-                    },
-                    "link": null
-                },
-                {
-                    "localized_name": "flow_shift",
-                    "name": "flow_shift",
-                    "type": "FLOAT",
-                    "widget": {
-                        "name": "flow_shift"
-                    },
-                    "link": null
-                }
-            ],
-            "outputs": [
-                {
-                    "localized_name": "Wan video params",
-                    "name": "Wan video params",
-                    "type": "VIDEO_PARAMS",
-                    "links": [
-                        28
-                    ]
-                }
-            ],
-            "properties": {
-                "Node name for S&R": "VLLMOmniWanParams"
-            },
-            "widgets_values": [
-                4,
-                0.875,
-                5
-            ]
-        },
-        {
-            "id": 12,
-            "type": "VLLMOmniGenerateVideo",
-            "pos": [
-                827.2566336087867,
-                -73.77449831827545
-            ],
-            "size": [
-                400,
-                346
-            ],
-            "flags": {},
-            "order": 4,
-            "mode": 0,
-            "inputs": [
-                {
-                    "localized_name": "image",
-                    "name": "image",
-                    "shape": 7,
-                    "type": "IMAGE",
-                    "link": 25
-                },
-                {
-                    "localized_name": "sampling_params",
-                    "name": "sampling_params",
-                    "shape": 7,
-                    "type": "SAMPLING_PARAMS",
-                    "link": 27
-                },
-                {
-                    "localized_name": "lora",
-                    "name": "lora",
-                    "shape": 7,
-                    "type": "REMOTE_LORA",
-                    "link": null
-                },
-                {
-                    "localized_name": "model_params",
-                    "name": "model_params",
-                    "shape": 7,
-                    "type": "VIDEO_PARAMS",
-                    "link": 28
-                },
-                {
-                    "localized_name": "url",
-                    "name": "url",
-                    "type": "STRING",
-                    "widget": {
-                        "name": "url"
-                    },
-                    "link": null
-                },
-                {
-                    "localized_name": "model",
-                    "name": "model",
-                    "type": "STRING",
-                    "widget": {
-                        "name": "model"
-                    },
-                    "link": null
-                },
-                {
-                    "localized_name": "prompt",
-                    "name": "prompt",
-                    "type": "STRING",
-                    "widget": {
-                        "name": "prompt"
-                    },
-                    "link": null
-                },
-                {
-                    "localized_name": "negative_prompt",
-                    "name": "negative_prompt",
-                    "type": "STRING",
-                    "widget": {
-                        "name": "negative_prompt"
-                    },
-                    "link": null
-                },
-                {
-                    "localized_name": "width",
-                    "name": "width",
-                    "type": "INT",
-                    "widget": {
-                        "name": "width"
-                    },
-                    "link": null
-                },
-                {
-                    "localized_name": "height",
-                    "name": "height",
-                    "type": "INT",
-                    "widget": {
-                        "name": "height"
-                    },
-                    "link": null
-                },
-                {
-                    "localized_name": "fps",
-                    "name": "fps",
-                    "type": "INT",
-                    "widget": {
-                        "name": "fps"
-                    },
-                    "link": null
-                },
-                {
-                    "localized_name": "num_frames",
-                    "name": "num_frames",
-                    "type": "INT",
-                    "widget": {
-                        "name": "num_frames"
-                    },
-                    "link": null
-                }
-            ],
-            "outputs": [
-                {
-                    "localized_name": "video",
-                    "name": "video",
-                    "type": "VIDEO",
-                    "links": [
-                        26
-                    ]
-                }
-            ],
-            "properties": {
-                "Node name for S&R": "VLLMOmniGenerateVideo"
-            },
-            "widgets_values": [
-                "http://localhost:8000/v1",
-                "Wan-AI/Wan2.2-TI2V-5B-Diffusers",
-                "Make the cat in the attached video yawning. Eye closed and mouth wide open. Keep the original realistic style.",
-                "Cartoonish",
-                900,
-                1200,
-                12,
-                41
-            ]
-        },
-        {
-            "id": 13,
-            "type": "SaveVideo",
-            "pos": [
-                1257.5441369501657,
-                -44.12956394419863
-            ],
-            "size": [
-                374,
-                372.34896826553796
-            ],
-            "flags": {},
-            "order": 5,
-            "mode": 0,
-            "inputs": [
-                {
-                    "localized_name": "video",
-                    "name": "video",
-                    "type": "VIDEO",
-                    "link": 26
-                },
-                {
-                    "localized_name": "filename_prefix",
-                    "name": "filename_prefix",
-                    "type": "STRING",
-                    "widget": {
-                        "name": "filename_prefix"
-                    },
-                    "link": null
-                },
-                {
-                    "localized_name": "format",
-                    "name": "format",
-                    "type": "COMBO",
-                    "widget": {
-                        "name": "format"
-                    },
-                    "link": null
-                },
-                {
-                    "localized_name": "codec",
-                    "name": "codec",
-                    "type": "COMBO",
-                    "widget": {
-                        "name": "codec"
-                    },
-                    "link": null
-                }
-            ],
-            "outputs": [],
-            "properties": {},
-            "widgets_values": [
-                "video/ComfyUI",
-                "auto",
-                "auto"
-            ]
-        }
-    ],
-    "links": [
-        [
-            25,
-            4,
-            0,
-            12,
-            0,
-            "IMAGE"
-        ],
-        [
-            26,
-            12,
-            0,
-            13,
-            0,
-            "VIDEO"
-        ],
-        [
-            27,
-            8,
-            0,
-            12,
-            1,
-            "SAMPLING_PARAMS"
-        ],
-        [
-            28,
-            14,
-            0,
-            12,
-            3,
-            "VIDEO_PARAMS"
-        ]
-    ],
-    "groups": [
-        {
-            "id": 1,
-            "title": "Input",
-            "bounding": [
-                213.1706010087147,
-                -313.20095750667554,
-                560.8246471437027,
-                407.4246532472182
-            ],
-            "color": "#3f789e",
-            "font_size": 24,
-            "flags": {}
-        }
-    ],
-    "config": {},
-    "extra": {
-        "workflowRendererVersion": "LG",
-        "ds": {
-            "scale": 1.1469075819486894,
-            "offset": [
-                101.7632950847089,
-                492.29121889347806
-            ]
-        }
-    },
-    "version": 0.4
-}
+{"id":"3aea1d08-7de5-4997-ba71-b2f66ad72a21","revision":0,"last_node_id":24,"last_link_id":35,"nodes":[{"id":13,"type":"SaveVideo","pos":[1315.729153459605,-137.92038948394355],"size":[374,313.7692307692308],"flags":{},"order":9,"mode":0,"inputs":[{"localized_name":"video","name":"video","type":"VIDEO","link":32},{"localized_name":"filename_prefix","name":"filename_prefix","type":"STRING","widget":{"name":"filename_prefix"},"link":null},{"localized_name":"format","name":"format","type":"COMBO","widget":{"name":"format"},"link":null},{"localized_name":"codec","name":"codec","type":"COMFY_DYNAMICCOMBO_V3","widget":{"name":"codec"},"link":null}],"outputs":[{"localized_name":"video","name":"video","type":"VIDEO","links":null}],"properties":{},"widgets_values":["video/ComfyUI","auto","auto"]},{"id":10,"type":"MarkdownNote","pos":[183.74636797766794,-571.2529158046766],"size":[240.3326478922474,136.3505820791881],"flags":{},"order":0,"mode":0,"inputs":[],"outputs":[],"title":"Note: Task and Input","properties":{},"widgets_values":["vLLM-Omni nodes are categorized based on the output modality. The \"Generate Video\" node supports both *text-to-video*, *frame-to-video*, and *reference-to-video* generation. The node will create corresponding payload depending on whether an input image is present or not."],"color":"#432","bgcolor":"#000"},{"id":4,"type":"LoadImage","pos":[452.07189605801125,-569.7242536425791],"size":[270,314],"flags":{},"order":1,"mode":0,"inputs":[{"localized_name":"image","name":"image","type":"COMBO","widget":{"name":"image"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"IMAGEUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"IMAGE","name":"IMAGE","type":"IMAGE","links":[29]},{"localized_name":"MASK","name":"MASK","type":"MASK","links":null}],"properties":{"Node name for S&R":"LoadImage"},"widgets_values":["rap.png","image"]},{"id":23,"type":"LoadImage","pos":[9.206110246057657,-7.6955291509170785],"size":[270,314],"flags":{},"order":2,"mode":0,"inputs":[{"localized_name":"image","name":"image","type":"COMBO","widget":{"name":"image"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"IMAGEUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"IMAGE","name":"IMAGE","type":"IMAGE","links":[34]},{"localized_name":"MASK","name":"MASK","type":"MASK","links":null}],"properties":{"Node name for S&R":"LoadImage"},"widgets_values":["white-cat.jpg","image"]},{"id":22,"type":"MarkdownNote","pos":[-9.783113411416666,-148.8790300038605],"size":[297.3012007349086,88],"flags":{},"order":3,"mode":0,"inputs":[],"outputs":[],"title":"Note: Reference to video (r2v)","properties":{},"widgets_values":["For r2v tasks, connect this Video Reference block to the main Generate Video block. Then disconnect the \"frame\" input above."],"color":"#432","bgcolor":"#000"},{"id":24,"type":"LoadAudio","pos":[317.9929742385805,-141.44821500065095],"size":[282.798828125,136],"flags":{},"order":4,"mode":0,"inputs":[{"localized_name":"audio","name":"audio","type":"COMBO","widget":{"name":"audio"},"link":null},{"localized_name":"audioUI","name":"audioUI","type":"AUDIO_UI","widget":{"name":"audioUI"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"AUDIOUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"AUDIO","name":"AUDIO","type":"AUDIO","links":[35]}],"properties":{"Node name for S&R":"LoadAudio"},"widgets_values":["nothing-beats-a-jet2-holiday.mp3",null,""]},{"id":21,"type":"VLLMOmniVideoReferences","pos":[376.6129172897595,105.41557277556292],"size":[156.9380859375,126],"flags":{},"order":7,"mode":0,"inputs":[{"localized_name":"image_1","name":"image_1","shape":7,"type":"IMAGE","link":34},{"localized_name":"image_2","name":"image_2","shape":7,"type":"IMAGE","link":null},{"localized_name":"audio_1","name":"audio_1","shape":7,"type":"AUDIO","link":35},{"localized_name":"audio_2","name":"audio_2","shape":7,"type":"AUDIO","link":null},{"localized_name":"video_1","name":"video_1","shape":7,"type":"VIDEO","link":null},{"localized_name":"video_2","name":"video_2","shape":7,"type":"VIDEO","link":null}],"outputs":[{"localized_name":"references","name":"references","type":"VIDEO_REFERENCES","links":null}],"properties":{"Node name for S&R":"VLLMOmniVideoReferences"},"widgets_values":[]},{"id":8,"type":"VLLMOmniDiffusionSampling","pos":[83.24620208113915,422.08543477542605],"size":[284.205078125,226],"flags":{},"order":5,"mode":0,"inputs":[{"localized_name":"n","name":"n","type":"INT","widget":{"name":"n"},"link":null},{"localized_name":"num_inference_steps","name":"num_inference_steps","type":"INT","widget":{"name":"num_inference_steps"},"link":null},{"localized_name":"guidance_scale","name":"guidance_scale","type":"FLOAT","widget":{"name":"guidance_scale"},"link":null},{"localized_name":"true_cfg_scale","name":"true_cfg_scale","type":"FLOAT","widget":{"name":"true_cfg_scLine truncated

docs/.nav.yml

@@ -118,6 +118,7 @@ nav:
       - design/feature/teacache.md
       - design/feature/skip_softmax.md
       - design/feature/async_chunk.md
+      - design/feature/omni_async_output_materialization.md
       - design/feature/vae_parallel.md
       - design/feature/diffusion_step_execution.md
       - design/feature/diffusion_request_level_batching.md

docs/README.md

@@ -52,12 +52,13 @@ vLLM-Omni is flexible and easy to use with:
 - Tensor, pipeline, data and expert parallelism support for distributed inference
 - Streaming outputs
 - OpenAI-compatible API server
+- Full-duplex realtime serving with streaming audio input and output (experimental)
 
 vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, including:
 
-- Omni-modality models (e.g. Qwen3-Omni, Cosmos3, HunyuanImage, BAGEL)
+- Omni-modality models (e.g. Qwen3-Omni, MiniCPM-o 4.5, Cosmos3, HunyuanImage, BAGEL)
 - TTS models (e.g. Qwen3-TTS, VoxCPM2, Ming-Omni-TTS, CosyVoice3)
-- Diffusion models — image, video, and audio generation (e.g. Qwen-Image, Wan2.2, FLUX)
+- Diffusion models — image, video, and audio generation (e.g. MiniMax H3, Qwen-Image, Wan2.2, FLUX)
 - Robot-policy and action models (e.g. GR00T-N1.7, DreamZero-DROID, InternVLA-A1, Cosmos3 action policy)
 
 For more information, checkout the following:

docs/configuration/README.md

@@ -2,7 +2,7 @@
 
 This section lists the most common options for running vLLM-Omni.
 
-For options within a vLLM Engine. Please refer to [vLLM Configuration](https://docs.vllm.ai/en/v0.16.0/configuration/index.html)
+For options within a vLLM Engine, please refer to the [vLLM 0.26 configuration guide](https://docs.vllm.ai/en/v0.26.0/configuration/index.html).
 
 Currently, the main options are maintained by stage configs for each model.
 

docs/design/feature/omni_async_output_materialization.md

@@ -0,0 +1,386 @@
+# Async Omni Output Materialization
+
+## Table of Contents
+
+1. [Overview](#overview)
+2. [Performance](#performance)
+3. [Architecture](#architecture)
+4. [Enabling the Feature](#enabling-the-feature)
+5. [Compatibility and Fallbacks](#compatibility-and-fallbacks)
+6. [Related Files](#related-files)
+
+## Overview
+
+Async Omni output materialization moves CPU-side construction of
+`OmniModelRunnerOutput` out of the autoregressive (AR) decode critical path.
+It complements [Async Chunk](async_chunk.md): async chunk pipelines partial
+outputs between stages, while async output materialization prevents the
+construction of those partial outputs from blocking the next decode step.
+
+An AR stage must return sampled token IDs quickly because the scheduler needs
+them for the next decode step. However, an Omni stage can also produce hidden
+states, multimodal tensors, streaming inter-stage payloads, and connector
+metadata. Building these payloads can require device-to-host (D2H) copies,
+tensor slicing, flattening, and Python object construction.
+
+Without this feature, all of that work runs inline in `sample_tokens()`:
+
+```text
+sample tokens
+  -> update decode state
+  -> copy hidden and multimodal outputs to CPU
+  -> build per-request payloads
+  -> build the streaming wire payload
+  -> collect connector signals
+  -> return sampled tokens
+  -> launch the next decode step
+```
+
+With async output materialization, only work required by the next decode step
+stays on the critical path:
+
+```text
+sample tokens
+  -> update decode state
+  -> snapshot output state and start asynchronous D2H copies
+  -> register sampled tokens
+  -> launch the next decode step
+
+background output path
+  -> wait for the payload snapshot
+  -> build per-request payloads
+  -> build the streaming wire payload
+  -> collect connector signals
+  -> construct OmniModelRunnerOutput
+```
+
+This allows payload construction for step `N` to overlap with GPU decode work
+for step `N + 1`. It does not change model computation or generated output.
+
+!!! note "Platform and validation scope"
+    The performance and correctness validation in this document covers CUDA
+    and ROCm. This is a validation scope, not a runtime platform guard: XPU
+    inherits `GPUARModelRunner`, and MUSA selects `GPUARWorker`, so those
+    backends may enter the asynchronous path when its other guards pass. XPU
+    and MUSA have not been validated for this optimization. Ascend NPU uses a
+    separate `NPUARModelRunner` path that still constructs
+    `OmniModelRunnerOutput` and drains connector output synchronously before
+    wrapping the result as an asynchronous runner output.
+
+<p align="center">
+  <img
+    alt="Decode step gap before and after async output materialization"
+    src="../../../source/architecture/qwen3-omni-async-output-step-gap.svg"
+    width="100%"
+  >
+</p>
+
+<p align="center">
+  <em>
+    Figure 1: Moving payload construction off the decode path reduces the
+    observed gap between consecutive Talker steps from about 2.8 ms to 41 µs.
+    Adapted from the
+    <a href="https://vllm.ai/blog/2026-07-01-qwen3-omni-optimization">
+      Qwen3-Omni optimization blog
+    </a>.
+  </em>
+</p>
+
+For Qwen3-Omni, the feature is enabled for both AR stages:
+
+- **Thinker**: Defers hidden-state and multimodal payload materialization.
+- **Talker**: Runs its decode-state postprocess eagerly, then defers codec
+  payload materialization. Because Code2Wav consumes codec codes rather than
+  Talker hidden states, the Talker also skips unnecessary hidden-state D2H.
+- **Code2Wav**: Does not use this path because it is a generation stage rather
+  than an AR stage.
+
+Qwen3-TTS uses the same mechanism in its AR Talker stage. The Talker updates
+the decode state it needs for the next step before returning, then constructs
+the codec payload for Code2Wav in the background.
+
+## Performance
+
+In the controlled Qwen3-Omni optimization sweep, enabling async output
+materialization on top of CUDA Graph and async chunk produced the following
+results at concurrency 64:
+
+| Configuration | Request throughput | Mean audio TTFP | Mean audio RTF |
+|---|---:|---:|---:|
+| CUDA Graph + async chunk | 9.3 req/s | 655 ms | 0.63 |
+| + async output materialization | 11.3 req/s | 631 ms | 0.47 |
+| Change | **+22%** | **-4%** | **-25%** |
+
+The sweep used `Qwen/Qwen3-Omni-30B-A3B-Instruct`, 640 Seed-TTS English
+prompts, and one replica each for Thinker, Talker, and Code2Wav. Async output
+materialization recovers throughput by removing CPU payload construction from
+the decode critical path without regressing time to first audio packet (TTFP).
+See the
+[Qwen3-Omni optimization blog](https://vllm.ai/blog/2026-07-01-qwen3-omni-optimization)
+for the complete benchmark methodology and results.
+
+The implementation was also validated against the synchronous safe-copy path.
+The text and audio output hashes matched, confirming that overlap does not
+change the generated result. See
+[PR #4476](https://github.com/vllm-project/vllm-omni/pull/4476) for the
+correctness and profiling details.
+
+## Architecture
+
+### Output Lifecycle
+
+`OmniAsyncGPUModelRunnerOutput` extends vLLM's
+`AsyncGPUModelRunnerOutput`. It preserves the existing asynchronous
+sampled-token feedback while adding a background builder for the complete Omni
+output.
+
+The output lifecycle is:
+
+1. `GPUARModelRunner.sample_tokens()` samples tokens and performs the
+   bookkeeping required by the next decode step.
+2. The runner snapshots step-local metadata such as request IDs, token spans,
+   scheduler output, and query start locations.
+3. CUDA payload tensors are cloned before reusable model or CUDA Graph buffers
+   can be overwritten.
+4. The cloned payload is copied to pinned CPU memory on a dedicated CUDA
+   stream. A CUDA event records when the copy is ready.
+5. `OmniAsyncGPUModelRunnerOutput` starts a background thread that waits for
+   the payload event and builds `OmniModelRunnerOutput`.
+6. The AR runner immediately registers the sampled-token CPU copy with the
+   input batch and returns, allowing async scheduling to advance.
+7. When the engine calls `get_output()`, it joins the background builder,
+   propagates any builder exception, and finalizes sampled tokens and logprobs
+   through the upstream output implementation.
+
+```text
+GPU / decode path                         Background output path
+
+forward + sample
+      |
+      +-- clone output tensors
+      +-- enqueue D2H copy --------------> wait for D2H event
+      +-- register sampled tokens               |
+      +-- return async output                    +-- slice per request
+      |                                         +-- build multimodal payloads
+next decode step                                +-- partition streaming payloads
+                                                +-- build the wire payload
+                                                +-- drain connector signals
+                                                +-- build OmniModelRunnerOutput
+```
+
+### Safe State Snapshots
+
+The next scheduler step can mutate runner state while the background builder is
+still active. Most step-local builder inputs are therefore detached snapshots.
+The snapshotted state includes:
+
+- scheduler token counts and speculative-token metadata
+- request IDs and request-to-batch-index mappings
+- sampled token IDs, logprobs, and prompt logprobs
+- query start locations and scheduled-token spans
+- hidden states and multimodal output tensors
+- KV-connector and encoder-cache outputs captured for the step
+
+CUDA tensor cloning is required because CUDA Graph and model output buffers can
+be reused by the next forward pass. The dedicated copy stream and pinned host
+buffers make the D2H transfer asynchronous; the snapshot retains its cloned
+CUDA sources until the transfer event completes.
+
+### Live Runner State and Ownership
+
+The complete builder input is not a detached snapshot. Connector output state
+remains runner-owned: `get_omni_connector_output()` drains live, per-cycle
+connector signals after the streaming payload is built. Correctness therefore
+depends on that builder being the only consumer that drains connector signals
+for its output cycle. Receive-side connector state written by the background
+receiver is coordinated by the connector mixin's `_lock`.
+`OmniAsyncGPUModelRunnerOutput.get_output()` joins the builder and is the
+completion and exception boundary before the materialized output is consumed.
+
+The shared output builder also contains a full-payload accumulation branch
+that reads `self.requests`. That branch is not reachable when async Omni output
+materialization is enabled: this feature requires `async_chunk`, while
+`should_accumulate_full_payload_output()` returns `False` whenever
+`async_chunk` is enabled. The feature path partitions each step's payload into
+streaming inter-stage and client handoffs instead.
+
+### Output Builder
+
+The background builder performs the work that previously ran inline:
+
+- resolves which requests need downstream payloads
+- converts hidden and multimodal tensors into per-request CPU payloads
+- applies model-specific payload processing
+- partitions per-step payloads into inter-stage and client streams
+- creates the tensor-only `multimodal_outputs` wire payload
+- drains live connector readiness signals after building the streaming payload
+- constructs the final `OmniModelRunnerOutput`
+
+Input-side connector operations remain synchronous with model execution. In
+particular, receiving inputs and flushing previously completed connector
+outputs are not deferred because they affect scheduler-visible request state.
+
+### Qwen3-Omni Stage Behavior
+
+| Stage | Async output behavior | Reason |
+|---|---|---|
+| Thinker | Snapshots hidden states and multimodal outputs; builds the downstream payload in the background | The Talker needs the payload, but the next Thinker decode step only needs sampled-token feedback |
+| Talker | Runs lightweight postprocess eagerly; snapshots codec outputs; omits hidden states from the downstream payload | `hidden_states.last` is needed by the next Talker step, while Code2Wav only needs codec codes |
+| Code2Wav | Uses the normal generation-stage output path | Code2Wav is not executed by `GPUARModelRunner` |
+
+The Qwen3-TTS Talker follows the same pattern as the Qwen3-Omni Talker:
+postprocess runs eagerly because the next decode step needs its updated state,
+while codec payload construction runs asynchronously. Its Code2Wav stage also
+uses the normal generation-stage path.
+
+## Enabling the Feature
+
+There is no separate `async_omni_output` command-line or YAML option. The
+feature is selected automatically for model stages that opt in when the runtime
+conditions are safe. On the validated CUDA and ROCm configurations, using the
+bundled deployment profile is enough to turn it on for the supported
+Qwen3-Omni and Qwen3-TTS pipelines. The selection logic has no CUDA/ROCm-only
+guard; see [Compatibility and Fallbacks](#compatibility-and-fallbacks) for the
+other platform paths. The same profile does not enable background Omni
+materialization on Ascend NPU.
+
+### Example 1: Qwen3-Omni
+
+Start the full Qwen3-Omni pipeline:
+
+```bash
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct \
+  --omni \
+  --port 8091
+```
+
+The model registry automatically loads
+`vllm_omni/deploy/qwen3_omni_moe.yaml`. Its relevant settings are:
+
+```yaml
+async_chunk: true
+
+stages:
+  - stage_id: 0  # Thinker
+    enable_prefix_caching: false
+    # async_scheduling defaults to true for AR stages
+
+  - stage_id: 1  # Talker
+    enable_prefix_caching: false
+    # async_scheduling defaults to true for AR stages
+
+  - stage_id: 2  # Code2Wav
+    enable_prefix_caching: false
+    async_scheduling: false
+```
+
+With this profile, async output materialization activates automatically for:
+
+- **Stage 0, Thinker**: Defers hidden-state and multimodal payload
+  construction.
+- **Stage 1, Talker**: Defers codec payload construction after eagerly
+  updating its decode state.
+
+Stage 2 does not use this feature because Code2Wav is not an AR stage.
+
+### Example 2: Qwen3-TTS
+
+Start Qwen3-TTS with any checkpoint that uses the Qwen3-TTS Talker pipeline.
+For example:
+
+```bash
+vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
+  --omni \
+  --port 8091
+```
+
+The model registry automatically loads `vllm_omni/deploy/qwen3_tts.yaml`. Its
+relevant settings are:
+
+```yaml
+async_chunk: true
+
+stages:
+  - stage_id: 0  # Talker
+    async_scheduling: true
+    enable_prefix_caching: false
+
+  - stage_id: 1  # Code2Wav
+    enable_prefix_caching: false
+```
+
+Async output materialization activates only for Stage 0, the AR Talker. The
+Talker keeps the state required by its next decode step on the GPU, skips an
+unnecessary hidden-state D2H copy, and builds the codec payload for Code2Wav in
+the background. Stage 1 uses the generation-stage output path.
+
+The same behavior applies to the Base and VoiceDesign Qwen3-TTS checkpoints
+because they use the same Talker implementation.
+
+!!! warning
+    Do not pass `--no-async-chunk` or enable prefix caching when you want this
+    optimization. Either change causes the runner to fall back to synchronous
+    output construction. Enabling async chunk alone cannot activate the feature
+    for a model that has not opted into async Omni output.
+
+!!! note "Prefix cache compatibility"
+    Async Omni output materialization and Omni prefix caching cannot currently
+    run together. `_should_use_async_omni_output()` returns `False` whenever
+    `self.omni_prefix_cache` is present, so prefix caching selects synchronous
+    output materialization. Supporting both features requires snapshotting or
+    otherwise synchronizing prefix-cache merge and update state before the
+    background output path can consume it safely.
+
+## Compatibility and Fallbacks
+
+`GPUARModelRunner` uses the async output path only when all of the following
+runtime conditions hold:
+
+| Requirement | Reason |
+|---|---|
+| AR async scheduling is enabled | The optimization relies on the scheduler advancing while the prior output is materialized |
+| `async_chunk` is enabled | The feature targets incremental downstream Omni payloads |
+| The model stage opts in with `use_async_omni_output` | Models must declare that their output lifecycle is safe to defer |
+| Omni prefix cache is disabled | Prefix-cache merge and update ordering currently requires synchronous materialization |
+| Speculative decoding is disabled | Speculative output state is not included in this deferred path |
+| Routed-expert output is disabled | Routed-expert extraction currently requires the synchronous path |
+| Postprocess is absent or explicitly runs eagerly | State needed by the next decode step must be updated before the runner returns |
+
+These checks are evaluated per stage on runners based on `GPUARModelRunner`.
+An unsupported combination does not prevent serving; it only falls back to
+synchronous output construction for that stage. CUDA and ROCm are the validated
+platforms. XPU inherits `GPUARModelRunner`, and MUSA selects `GPUARWorker`, so
+they may select the async path when these checks pass, but this optimization is
+not yet validated on either backend.
+
+Ascend NPU does not use this selection path. `NPUARModelRunner.sample_tokens()`
+fully constructs `OmniModelRunnerOutput`, calls
+`get_omni_connector_output()`, and only then creates the asynchronous wrapper,
+so Omni payload materialization remains synchronous.
+
+!!! note
+    `use_async_omni_output`,
+    `eager_omni_postprocess_before_async_output`, and
+    `omni_pooler_payload_include_hidden` are model implementation contracts,
+    not user-facing configuration fields.
+
+## Related Files
+
+- `vllm_omni/worker/gpu_ar_model_runner.py`: Async output object, safe tensor
+  snapshots, compatibility guards, and deferred Omni output builder.
+- `vllm_omni/platforms/npu/worker/npu_ar_model_runner.py`: Separate Ascend NPU
+  runner, where Omni output materialization remains synchronous.
+- `vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py`: Thinker and
+  Talker opt-in behavior.
+- `vllm_omni/deploy/qwen3_omni_moe.yaml`: Default Qwen3-Omni deployment
+  settings.
+- `vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py`: Qwen3-TTS
+  Talker opt-in and eager postprocess behavior.
+- `vllm_omni/deploy/qwen3_tts.yaml`: Default Qwen3-TTS deployment settings.
+- `tests/worker/test_gpu_ar_model_runner.py`: Snapshot, guard, connector
+  ordering, and background error-propagation tests.
+- [Async Chunk](async_chunk.md): Inter-stage chunking and scheduling design.
+- [Qwen3-Omni optimization blog](https://vllm.ai/blog/2026-07-01-qwen3-omni-optimization):
+  Optimization context and controlled performance results.
+- [PR #4476](https://github.com/vllm-project/vllm-omni/pull/4476): Feature
+  implementation, validation, and review history.

docs/design/index.md

@@ -10,6 +10,7 @@ This section contains design documents and architecture specifications for vLLM-
 
 - [Disaggregated Inference](feature/disaggregated_inference.md)
 - [Ray-based Execution](feature/ray_based_execution.md)
+- [Async Omni Output Materialization](feature/omni_async_output_materialization.md)
 - [Adding Step Execution Support for Diffusion Pipelines](feature/diffusion_step_execution.md)
 - [Request-Level Batching for Diffusion](feature/diffusion_request_level_batching.md)
 - [Continuous Batching for Step-Wise Diffusion](feature/diffusion_continuous_batching.md)

docs/getting_started/installation/README.md

@@ -1,5 +1,8 @@
 # Installation
 
+!!! important
+    vLLM-Omni is released against the matching upstream vLLM major/minor version. vLLM-Omni 0.26.x requires vLLM 0.26.x; the stable v0.26.0 instructions pin vLLM 0.26.0.
+
 vLLM-Omni supports the following hardware platforms:
 
 - [GPU](gpu.md)

docs/getting_started/installation/gpu/cuda.inc.md

@@ -5,7 +5,7 @@
 # --8<-- [end:requirements]
 # --8<-- [start:set-up-using-python]
 
-vLLM-Omni depends vLLM. So please follow instructions below mainly for vLLM.
+vLLM-Omni depends on the matching major/minor release of vLLM. The vLLM-Omni 0.26.x release line uses vLLM 0.26.x.
 
 !!! note
     PyTorch installed via `conda` will statically link `NCCL` library, which can cause issues when vLLM tries to use `NCCL`. See <gh-issue:8420> for more details.
@@ -124,7 +124,7 @@ If you want to specify the base vLLM version:
 ```bash
 DOCKER_BUILDKIT=1 docker build \
   -f docker/Dockerfile.cuda \
-  --build-arg BASE_IMAGE=vllm/vllm-openai:v0.22.1 \
+  --build-arg BASE_IMAGE=vllm/vllm-openai:v0.26.0 \
   -t vllm-omni-cuda .
 ```
 

docs/getting_started/installation/npu/npu.inc.md

@@ -53,16 +53,20 @@ We are keeping [issue #886](https://github.com/vllm-project/vllm-omni/issues/886
 You can also build vLLM-Omni from the latest main branch if you want to use the latest features or bug fixes. (But sometimes it will break for a while. You can check [issue #886](https://github.com/vllm-project/vllm-omni/issues/886) for the status of the latest commit of vLLM-Omni main branch on NPU.)
 
 ```bash
-# Pin vLLM version to 0.18.0
-git clone -b v0.18.0 https://github.com/vllm-project/vllm.git
+# Pin vLLM and vLLM-Ascend to the v0.26 release line
+git clone -b v0.26.0 https://github.com/vllm-project/vllm.git
+cd vllm
 VLLM_TARGET_DEVICE=empty pip install -v -e .
+cd ..
 
-git clone -b v0.18.0rc1 https://github.com/vllm-project/vllm-ascend.git
+git clone -b releases/v0.26.0rc https://github.com/vllm-project/vllm-ascend.git
+cd vllm-ascend
 pip install -v -e .
+cd ..
 
 # Install vLLM-Omni from the latest main branch
 git clone https://github.com/vllm-project/vllm-omni.git
-cd /vllm-workspace/vllm-omni
+cd vllm-omni
 pip install -v -e . --no-build-isolation
 # or VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e .
 export VLLM_WORKER_MULTIPROC_METHOD=spawn

docs/source/architecture/qwen3-omni-async-output-step-gap.svg

@@ -0,0 +1,150 @@
+<svg xmlns="http://www.w3.org/2000/svg" width="1024" height="600" viewBox="0 0 1024 600" role="img" aria-labelledby="title desc">
+  <title id="title">Qwen3-Omni Talker step gap before and after asynchronous output materialization</title>
+  <desc id="desc">Before the optimization, synchronous D2H copy, per-request payload construction, streaming payload handoff, and connector-signal drain separate consecutive Talker forward and sample steps by about 2.8 milliseconds. After the optimization, that Omni output work runs on a background builder while the next Talker step executes, leaving a boundary of about 41 microseconds for snapshot, D2H enqueue, state update, and sampled-token feedback.</desc>
+  <defs>
+    <marker id="time-arrow" viewBox="0 0 10 10" refX="8" refY="5" markerWidth="7" markerHeight="7" orient="auto">
+      <path d="M 0 0 L 10 5 L 0 10 z" fill="#64748b"/>
+    </marker>
+    <marker id="red-end" viewBox="0 0 10 10" refX="8" refY="5" markerWidth="5" markerHeight="5" orient="auto">
+      <path d="M 0 0 L 10 5 L 0 10 z" fill="#b91c1c"/>
+    </marker>
+    <marker id="red-start" viewBox="0 0 10 10" refX="2" refY="5" markerWidth="5" markerHeight="5" orient="auto">
+      <path d="M 10 0 L 0 5 L 10 10 z" fill="#b91c1c"/>
+    </marker>
+    <marker id="blue-end" viewBox="0 0 10 10" refX="8" refY="5" markerWidth="6" markerHeight="6" orient="auto">
+      <path d="M 0 0 L 10 5 L 0 10 z" fill="#2563eb"/>
+    </marker>
+    <style>
+      .canvas { fill: #ffffff; stroke: #e5e7eb; }
+      .title { font: 700 22px system-ui, -apple-system, Segoe UI, sans-serif; fill: #111827; }
+      .subtitle { font: 500 13px system-ui, -apple-system, Segoe UI, sans-serif; fill: #64748b; }
+      .phase { font: 700 16px system-ui, -apple-system, Segoe UI, sans-serif; fill: #1f2937; }
+      .phase-sub { font: 500 12px system-ui, -apple-system, Segoe UI, sans-serif; fill: #64748b; }
+      .lane { font: 600 12px system-ui, -apple-system, Segoe UI, sans-serif; fill: #374151; }
+      .block-title { font: 700 10px system-ui, -apple-system, Segoe UI, sans-serif; letter-spacing: .25px; fill: #ffffff; }
+      .block-op { font: 500 9.5px system-ui, -apple-system, Segoe UI, sans-serif; fill: #ffffff; }
+      .axis { font: 500 11px system-ui, -apple-system, Segoe UI, sans-serif; fill: #94a3b8; }
+      .measure { font: 700 12px system-ui, -apple-system, Segoe UI, sans-serif; fill: #b91c1c; }
+      .inset-title { font: 700 11px system-ui, -apple-system, Segoe UI, sans-serif; fill: #475569; }
+      .inset-step { font: 600 10px system-ui, -apple-system, Segoe UI, sans-serif; fill: #ffffff; }
+      .note-title { font: 700 12px system-ui, -apple-system, Segoe UI, sans-serif; fill: #1e3a8a; }
+      .note-text { font: 500 10.5px system-ui, -apple-system, Segoe UI, sans-serif; fill: #475569; }
+      .stat-cap { font: 600 11px system-ui, -apple-system, Segoe UI, sans-serif; fill: #64748b; }
+      .stat-big { font: 800 16px system-ui, -apple-system, Segoe UI, sans-serif; fill: #111827; }
+      .stat-sub { font: 500 11px system-ui, -apple-system, Segoe UI, sans-serif; fill: #64748b; }
+      .before-card { fill: #eef4ff; stroke: #c7dbf6; }
+      .after-card { fill: #eafaf0; stroke: #bfe7cf; }
+      .talker { fill: #16a34a; }
+      .sync-output { fill: #ea580c; }
+      .async-output { fill: #2563eb; }
+      .zoom-card { fill: #fff7f7; stroke: #f3caca; }
+      .note-card { fill: #eff6ff; stroke: #bfdbfe; }
+      .stat-blue { fill: #eef4ff; stroke: #c7dbf6; }
+      .stat-green { fill: #eafaf0; stroke: #bfe7cf; }
+      .stat-gray { fill: #f4f5f7; stroke: #e2e5ea; }
+      .guide { stroke: #e2e8f0; stroke-width: 1; }
+      .time-line { stroke: #64748b; stroke-width: 1.3; fill: none; marker-end: url(#time-arrow); }
+      .measure-line { stroke: #b91c1c; stroke-width: 1.4; marker-start: url(#red-start); marker-end: url(#red-end); }
+      .zoom-lead { stroke: #dc2626; stroke-width: 1.15; stroke-dasharray: 4 3; fill: none; }
+      .handoff { stroke: #2563eb; stroke-width: 1.3; stroke-dasharray: 3 2; fill: none; marker-end: url(#blue-end); }
+    </style>
+  </defs>
+
+  <rect x="14" y="14" width="996" height="572" rx="14" class="canvas"/>
+
+  <text x="44" y="52" class="title">Async output: what moves off the Talker decode path</text>
+  <text x="44" y="74" class="subtitle">Qwen3-Omni Talker stage &#183; consecutive decode steps in a fixed profiler window</text>
+
+  <!-- ===================== BEFORE ===================== -->
+  <rect x="44" y="94" width="144" height="110" rx="10" class="before-card"/>
+  <text x="60" y="120" class="phase">Before</text>
+  <text x="60" y="139" class="phase-sub">Omni output built</text>
+  <text x="60" y="156" class="phase-sub">synchronously on the</text>
+  <text x="60" y="173" class="phase-sub">AR critical path</text>
+
+  <text x="206" y="122" class="lane">AR path</text>
+  <line x1="300" y1="88" x2="300" y2="210" class="guide"/>
+
+  <rect x="300" y="104" width="190" height="54" rx="5" class="talker"/>
+  <text x="395" y="124" text-anchor="middle" class="block-title">TALKER STAGE &#183; STEP N</text>
+  <text x="395" y="145" text-anchor="middle" class="block-op">forward &#8594; sample</text>
+
+  <rect x="490" y="104" width="250" height="54" rx="5" class="sync-output"/>
+  <text x="615" y="120" text-anchor="middle" class="block-title">SYNC OMNI OUTPUT &#183; STEP N</text>
+  <text x="615" y="138" text-anchor="middle" class="block-op">D2H &#8594; per-request slice / payload build</text>
+  <text x="615" y="152" text-anchor="middle" class="block-op">streaming handoff &#8594; connector-signal drain</text>
+
+  <rect x="740" y="104" width="220" height="54" rx="5" class="talker"/>
+  <text x="850" y="124" text-anchor="middle" class="block-title">TALKER STAGE &#183; STEP N+1</text>
+  <text x="850" y="145" text-anchor="middle" class="block-op">forward &#8594; sample</text>
+
+  <line x1="490" y1="174" x2="740" y2="174" class="measure-line"/>
+  <text x="615" y="194" text-anchor="middle" class="measure">~2.8 ms GPU gap: output work is inline</text>
+
+  <!-- ===================== AFTER ===================== -->
+  <rect x="44" y="218" width="144" height="140" rx="10" class="after-card"/>
+  <text x="60" y="244" class="phase">After</text>
+  <text x="60" y="263" class="phase-sub">Only step feedback</text>
+  <text x="60" y="280" class="phase-sub">stays inline; payload</text>
+  <text x="60" y="297" class="phase-sub">work runs in the</text>
+  <text x="60" y="314" class="phase-sub">background</text>
+
+  <text x="206" y="254" class="lane">Talker GPU</text>
+  <text x="206" y="326" class="lane">Output builder</text>
+  <line x1="300" y1="212" x2="300" y2="370" class="guide"/>
+
+  <rect x="300" y="226" width="235" height="54" rx="5" class="talker"/>
+  <text x="418" y="246" text-anchor="middle" class="block-title">TALKER STAGE &#183; STEP N</text>
+  <text x="418" y="267" text-anchor="middle" class="block-op">forward &#8594; sample</text>
+
+  <rect x="540" y="226" width="235" height="54" rx="5" class="talker"/>
+  <text x="658" y="246" text-anchor="middle" class="block-title">TALKER STAGE &#183; STEP N+1</text>
+  <text x="658" y="267" text-anchor="middle" class="block-op">forward &#8594; sample</text>
+
+  <path d="M 535 280 L 548 298" class="handoff"/>
+  <rect x="540" y="298" width="360" height="58" rx="5" class="async-output"/>
+  <text x="720" y="315" text-anchor="middle" class="block-title">BACKGROUND OMNI OUTPUT &#183; STEP N</text>
+  <text x="720" y="334" text-anchor="middle" class="block-op">wait for D2H &#8594; per-request slice / payload build</text>
+  <text x="720" y="349" text-anchor="middle" class="block-op">streaming handoff &#8594; connector-signal drain</text>
+
+  <line x1="300" y1="366" x2="930" y2="366" class="time-line"/>
+  <text x="922" y="382" text-anchor="end" class="axis">time</text>
+
+  <!-- leads from the real 41 us boundary to the magnified inset -->
+  <line x1="535" y1="280" x2="485" y2="424" class="zoom-lead"/>
+  <line x1="540" y1="280" x2="525" y2="424" class="zoom-lead"/>
+
+  <!-- ===================== 41 µs BOUNDARY ===================== -->
+  <rect x="300" y="386" width="400" height="92" rx="10" class="zoom-card"/>
+  <text x="316" y="404" class="inset-title">Zoom &#183; Talker step boundary (magnified)</text>
+
+  <rect x="320" y="416" width="165" height="26" rx="4" class="talker"/>
+  <rect x="525" y="416" width="155" height="26" rx="4" class="talker"/>
+  <text x="403" y="433" text-anchor="middle" class="inset-step">Talker step N</text>
+  <text x="603" y="433" text-anchor="middle" class="inset-step">Talker step N+1</text>
+
+  <line x1="485" y1="452" x2="525" y2="452" class="measure-line"/>
+  <text x="505" y="469" text-anchor="middle" class="measure">~41 &#181;s</text>
+
+  <rect x="720" y="386" width="240" height="92" rx="10" class="note-card"/>
+  <text x="736" y="405" class="note-title">Inline at the ~41 &#181;s boundary</text>
+  <text x="736" y="424" class="note-text">&#8226; eager Talker state update</text>
+  <text x="736" y="441" class="note-text">&#8226; output snapshot + D2H enqueue</text>
+  <text x="736" y="458" class="note-text">&#8226; sampled-token feedback</text>
+
+  <!-- ===================== SUMMARY ===================== -->
+  <rect x="44" y="494" width="296" height="70" rx="9" class="stat-blue"/>
+  <text x="60" y="514" class="stat-cap">Step-to-step gap</text>
+  <text x="60" y="538" class="stat-big">2.8 ms &#8594; 41 &#181;s</text>
+  <text x="60" y="555" class="stat-sub">~68&#215; tighter</text>
+
+  <rect x="356" y="494" width="296" height="70" rx="9" class="stat-green"/>
+  <text x="372" y="514" class="stat-cap">Talker GPU</text>
+  <text x="372" y="538" class="stat-big">next step starts immediately</text>
+  <text x="372" y="555" class="stat-sub">decode workers stay busy</text>
+
+  <rect x="668" y="494" width="312" height="70" rx="9" class="stat-gray"/>
+  <text x="684" y="514" class="stat-cap">Background builder</text>
+  <text x="684" y="538" class="stat-big">materializes Omni output</text>
+  <text x="684" y="555" class="stat-sub">payload work overlaps the next Talker step</text>
+</svg>

docs/user_guide/quantization/fp8.md

@@ -93,6 +93,7 @@ warmup_quack_fp8([(14040, 2048, 6144), (14040, 2048, 2048)])
 | HunyuanImage-3.0 | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | Yes | Yes | All layers; use the Hunyuan stage config for multi-stage runs | None | |
 | HunyuanVideo-1.5 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `720p_t2v`, `480p_i2v` | Yes | Yes | All layers | None | |
 | Cosmos3 | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | Yes | Not validated | All layers | None | |
+| MiniMax H3 | `MiniMaxAI/MiniMax-H3` | Yes | Not validated | DiT linears except patch, timestep, and final projections; not compatible with layerwise offload | None | |
 
 ### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
 

recipes/MiniMaxAI/MiniMax-H3.md

@@ -164,6 +164,42 @@ command with:
 export MODEL="${MODEL_ROOT}/Ref2VA"
 ```
 
+### Online FP8 quantization
+
+MiniMax H3 supports load-time FP8 quantization of the DiT. The checkpoint
+remains BF16 on disk; vLLM-Omni quantizes eligible weights while loading and
+uses dynamic activation scaling during inference. By default, attention and
+MLP linears in the token refiner and main DiT blocks, the condition
+projection, and all AdaLN projections use FP8. Patch, timestep, and final
+projections remain FP32; the text encoder and VAEs are unchanged.
+
+Add this option to an existing H3 server command:
+
+```bash
+--quantization fp8
+```
+
+Use `ignored_layers` to keep any otherwise eligible linear in BF16. H3
+resolves the `transformer` component before constructing the DiT, so names do
+not start with `transformer.`. Entries are exact runtime linear prefixes; a
+parent name such as `blocks.0.attn` does not exclude its children.
+
+Eligible names are the `attn.qkv_proj`, `attn.out_proj`, `mlp.fc1`, and
+`mlp.fc2` children under `token_refiner.blocks.<0-1>` or `blocks.<0-49>`.
+The other eligible names are `condition_proj`,
+`blocks.<0-49>.adaln_proj.linear`, and `final_layer.adaln_proj.linear`.
+For example, keep the first main block's attention projections in BF16 with:
+
+```bash
+--diffusion-quantization-config \
+  '{"transformer":{"method":"fp8","ignored_layers":["blocks.0.attn.qkv_proj","blocks.0.attn.out_proj"]}}'
+```
+
+The structured option replaces `--quantization fp8`. Online FP8 is currently
+incompatible with H3 layerwise offload because the offload path produces a
+weight stride rejected by the Cutlass FP8 kernel. Use resident FP8 with tensor
+parallelism and VAE tiling instead.
+
 ## HTTP API examples
 
 The following requests use the synchronous endpoint so the returned body can
@@ -310,12 +346,21 @@ throughput guarantee. Multi-video Ref2VA is much slower because the two
 reference videos expand both the Qwen3-VL vision sequence and the packed DiT
 attention sequence.
 
+## Validated FP8 evidence
+
+With eager DiT/text-encoder TP2 and VAE tiling, the 384x672, 107-frame,
+10-step quality case measured LPIPS 0.1156 (limit 0.20), PSNR 23.6316 dB,
+audio spectral cosine 0.9589 (minimum 0.80), and audio RMS ratio 0.9342. The
+resident per-GPU peak was 68.52 GiB for BF16 and 53.51 GiB for FP8, a 22%
+reduction.
+
 ## Known limitations
 
 - Each server process loads only one checkpoint partition.
 - H3 currently executes one generation request per diffusion batch.
 - The first regional-compile request is a warmup and should not be included in
   steady-state performance measurements.
+- Online FP8 is not compatible with layerwise offload.
 - Image+audio Ref2VA accepts exactly one image and one audio reference.
 - Video Ref2VA accepts one or more video files, but not an additional standalone
   audio reference.

tests/core/sched/test_omni_ar_scheduler_logprobs.py

@@ -3,7 +3,8 @@
 from __future__ import annotations
 
 from collections import defaultdict
-from types import SimpleNamespace
+from collections.abc import Callable
+from types import MethodType, SimpleNamespace
 
 import numpy as np
 import pytest
@@ -12,9 +13,22 @@ from vllm.v1.outputs import LogprobsLists
 from vllm.v1.request import RequestStatus
 
 from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler, _slice_sampled_logprobs
+from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
+# Shared update_from_output helpers moved onto OmniSchedulerMixin; bind them on
+# the lightweight stub so direct method calls still exercise the real paths.
+_MIXIN_UPDATE_HELPERS = (
+    "_remove_stopped_requests_from_queues",
+    "_handle_failed_kv_load_outputs",
+    "_aggregate_kv_connector_stats",
+    "_publish_kv_cache_events",
+    "_attach_finished_request_sets",
+    "_attach_scheduler_stats",
+    "_capture_omni_connector_output",
+)
+
 
 def _logprobs(
     token_ids: list[list[int]],
@@ -105,18 +119,14 @@ class _RequestQueue(list):
                 self.remove(request)
 
 
-def test_mid_step_stop_trims_logprob_rows_with_token_ids() -> None:
-    """Spec decode can sample tokens past a stop; the trailing tokens are
-    trimmed by _update_request_with_output, and the emitted logprob rows must
-    be trimmed with them (upstream slices logprobs after the trim)."""
-    request = _Request("req")
-    request.num_computed_tokens = 8
+def _make_scheduler_stub(requests: list[_Request]) -> SimpleNamespace:
+    """Build a minimal stub for OmniARScheduler.update_from_output()."""
     scheduler = SimpleNamespace(
         perf_metrics=None,
         connector=None,
         chunk_transfer_adapter=None,
-        requests={request.request_id: request},
-        running=[request],
+        requests={request.request_id: request for request in requests},
+        running=list(requests),
         waiting=_RequestQueue(),
         skipped_waiting=_RequestQueue(),
         structured_output_manager=SimpleNamespace(should_advance=lambda _request: False),
@@ -129,7 +139,42 @@ def test_mid_step_stop_trims_logprob_rows_with_token_ids() -> None:
         finished_req_ids_dict=defaultdict(set),
         kv_cache_manager=SimpleNamespace(take_events=lambda: None),
         kv_event_publisher=SimpleNamespace(publish=lambda _events: None),
+        recompute_kv_load_failures=False,
     )
+    for name in _MIXIN_UPDATE_HELPERS:
+        setattr(scheduler, name, MethodType(getattr(OmniSchedulerMixin, name), scheduler))
+    scheduler._cleanup_kv_tracking = MethodType(OmniARScheduler._cleanup_kv_tracking, scheduler)
+    scheduler.make_spec_decoding_stats = lambda *args, **kwargs: None
+    scheduler.make_stats = lambda *args, **kwargs: None
+    return scheduler
+
+
+def _bind_request_lifecycle(
+    scheduler: SimpleNamespace,
+    *,
+    update_request: Callable,
+    handle_stopped: Callable | None = None,
+) -> None:
+    def free_request(request):
+        scheduler.requests.pop(request.request_id, None)
+        scheduler.finished_req_ids.add(request.request_id)
+        scheduler.finished_req_ids_dict[request.client_index].add(request.request_id)
+        # vLLM 0.26 contract: (kv_xfer_params, ec_xfer_params)
+        return None, None
+
+    scheduler._update_request_with_output = update_request
+    scheduler._process_kv_transfer_trigger = lambda _request, _tokens: False
+    scheduler._handle_stopped_request = handle_stopped or (lambda _request: True)
+    scheduler._free_request = free_request
+
+
+def test_mid_step_stop_trims_logprob_rows_with_token_ids() -> None:
+    """Spec decode can sample tokens past a stop; the trailing tokens are
+    trimmed by _update_request_with_output, and the emitted logprob rows must
+    be trimmed with them (upstream slices logprobs after the trim)."""
+    request = _Request("req")
+    request.num_computed_tokens = 8
+    scheduler = _make_scheduler_stub([request])
 
     def update_request_trimming(req, token_ids):
         # Mirror vllm Scheduler._update_request_with_output: the first sampled
@@ -139,19 +184,7 @@ def test_mid_step_stop_trims_logprob_rows_with_token_ids() -> None:
         req.status = RequestStatus.FINISHED_STOPPED
         return token_ids, True
 
-    def free_request(req):
-        scheduler.requests.pop(req.request_id, None)
-        scheduler.finished_req_ids.add(req.request_id)
-        scheduler.finished_req_ids_dict[req.client_index].add(req.request_id)
-        return None, None
-
-    scheduler._update_request_with_output = update_request_trimming
-    scheduler._process_kv_transfer_trigger = lambda _request, _tokens: False
-    scheduler._handle_stopped_request = lambda _request: True
-    scheduler._free_request = free_request
-    scheduler.make_spec_decoding_stats = lambda *args, **kwargs: None
-    scheduler.make_stats = lambda *args, **kwargs: None
-    scheduler._capture_omni_connector_output = lambda _output: None
+    _bind_request_lifecycle(scheduler, update_request=update_request_trimming)
 
     scheduler_output = SimpleNamespace(
         num_scheduled_tokens={"req": 3},
@@ -190,45 +223,14 @@ def test_invalid_logprobs_finish_only_the_affected_scheduler_request() -> None:
     """A bad runner response must not bring down unrelated batch requests."""
     bad = _Request("bad")
     good = _Request("good")
-    scheduler = SimpleNamespace(
-        perf_metrics=None,
-        connector=None,
-        chunk_transfer_adapter=None,
-        requests={bad.request_id: bad, good.request_id: good},
-        running=[bad, good],
-        waiting=_RequestQueue(),
-        skipped_waiting=_RequestQueue(),
-        structured_output_manager=SimpleNamespace(should_advance=lambda _request: False),
-        transfer_triggered_requests=set(),
-        active_kv_transfers=set(),
-        pending_stop_after_extraction=set(),
-        waiting_for_transfer_free=set(),
-        _new_prompt_len_snapshot={},
-        finished_req_ids=set(),
-        finished_req_ids_dict=defaultdict(set),
-        kv_cache_manager=SimpleNamespace(take_events=lambda: None),
-        kv_event_publisher=SimpleNamespace(publish=lambda _events: None),
-    )
+    scheduler = _make_scheduler_stub([bad, good])
     update_calls: list[str] = []
 
     def update_request(request, token_ids):
         update_calls.append(request.request_id)
         return token_ids, False
 
-    def free_request(request):
-        scheduler.requests.pop(request.request_id, None)
-        scheduler.finished_req_ids.add(request.request_id)
-        scheduler.finished_req_ids_dict[request.client_index].add(request.request_id)
-        # vLLM 0.26 contract: (kv_xfer_params, ec_xfer_params)
-        return None, None
-
-    scheduler._update_request_with_output = update_request
-    scheduler._process_kv_transfer_trigger = lambda _request, _tokens: False
-    scheduler._handle_stopped_request = lambda _request: True
-    scheduler._free_request = free_request
-    scheduler.make_spec_decoding_stats = lambda *args, **kwargs: None
-    scheduler.make_stats = lambda *args, **kwargs: None
-    scheduler._capture_omni_connector_output = lambda _output: None
+    _bind_request_lifecycle(scheduler, update_request=update_request)
 
     scheduler_output = SimpleNamespace(
         num_scheduled_tokens={"bad": 1, "good": 1},

tests/core/sched/test_omni_scheduler_mixin_shared.py

@@ -0,0 +1,90 @@
+from collections import defaultdict
+from types import SimpleNamespace
+
+import pytest
+from vllm.v1.engine import FinishReason
+
+from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
+from vllm_omni.core.sched.output import OmniChunkRecvHandle
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _Scheduler(OmniSchedulerMixin):
+    pass
+
+
+def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths():
+    calls = []
+    scheduler = _Scheduler()
+    scheduler.waiting = ["waiting"]
+    scheduler.running = ["running"]
+    scheduler.requests = {"request": object()}
+    scheduler._consume_pending_connector_output = lambda mode: calls.append(("consume", mode))
+    scheduler._process_pending_input_timeouts = lambda: calls.append(("timeouts",))
+    scheduler.chunk_transfer_adapter = SimpleNamespace(
+        process_pending_chunks=lambda waiting, running, scheduler_requests: calls.append(
+            ("process", waiting, running, scheduler_requests)
+        ),
+        restore_queues=lambda waiting, running, scheduler_requests: calls.append(
+            ("restore-chunks", waiting, running, scheduler_requests)
+        ),
+    )
+    scheduler.input_coordinator = SimpleNamespace(
+        restore_queues=lambda waiting: calls.append(("restore-full", waiting))
+    )
+
+    scheduler._process_pending_omni_inputs("ar")
+    scheduler._restore_omni_wait_queues()
+
+    assert calls == [
+        ("consume", "ar"),
+        ("timeouts",),
+        ("process", scheduler.waiting, scheduler.running, scheduler.requests),
+        ("restore-chunks", scheduler.waiting, scheduler.running, scheduler.requests),
+        ("restore-full", scheduler.waiting),
+    ]
+
+
+@pytest.mark.parametrize(
+    ("synthesize_abort_outputs", "expected_finish_reason"),
+    [(False, None), (True, FinishReason.ABORT)],
+)
+def test_finished_request_attachment_keeps_ar_abort_policy_explicit(
+    synthesize_abort_outputs,
+    expected_finish_reason,
+):
+    scheduler = _Scheduler()
+    scheduler.finished_req_ids_dict = defaultdict(set, {2: {"req-finished"}})
+    outputs = {}
+
+    scheduler._attach_finished_request_sets(
+        outputs,
+        synthesize_abort_outputs=synthesize_abort_outputs,
+    )
+
+    assert outputs[2].finished_requests == {"req-finished"}
+    if expected_finish_reason is None:
+        assert outputs[2].outputs == []
+    else:
+        assert outputs[2].outputs[0].finish_reason == expected_finish_reason
+    assert scheduler.finished_req_ids_dict == {}
+
+
+def test_chunk_receive_handle_carries_minimal_registration_fields():
+    handle = OmniChunkRecvHandle(request_id="req", external_req_id="external")
+    assert handle.request_id == "req"
+    assert handle.external_req_id == "external"
+
+
+def test_output_helper_preserves_required_nan_counter_default():
+    scheduler = _Scheduler()
+    request = SimpleNamespace(
+        request_id="req-output",
+        trace_headers=None,
+        take_events=lambda: [],
+    )
+
+    output = scheduler._make_omni_engine_output(request, new_token_ids=[])
+
+    assert output.num_nans_in_logits == 0

tests/diffusion/models/minimax_h3/test_minimax_h3_e2e.py

@@ -16,6 +16,17 @@ pytestmark = [
 MODEL_ENV = "VLLM_TEST_MINIMAX_H3_FL2VA_MODEL"
 
 
+def _assert_joint_output(outputs):
+    assert len(outputs) == 1
+    frames = np.asarray(outputs[0].images[0])
+    assert frames.shape == (39, 256, 448, 3)
+    multimodal = outputs[0].multimodal_output
+    assert multimodal is not None
+    assert np.asarray(multimodal["audio"]).shape[1] == 2
+    assert multimodal["audio_sample_rate"] == 32000
+    assert multimodal["fps"] == 24
+
+
 @pytest.mark.skipif(
     not os.environ.get(MODEL_ENV),
     reason=f"set {MODEL_ENV} to an authorized FL2VA checkpoint path",
@@ -59,11 +70,56 @@ def test_minimax_h3_t2va_ulysses8_smoke():
     finally:
         engine.close()
 
-    assert len(outputs) == 1
-    frames = np.asarray(outputs[0].images[0])
-    assert frames.shape == (39, 256, 448, 3)
-    multimodal = outputs[0].multimodal_output
-    assert multimodal is not None
-    assert np.asarray(multimodal["audio"]).shape[1] == 2
-    assert multimodal["audio_sample_rate"] == 32000
-    assert multimodal["fps"] == 24
+    _assert_joint_output(outputs)
+
+
+@pytest.mark.skipif(
+    not os.environ.get(MODEL_ENV),
+    reason=f"set {MODEL_ENV} to an authorized FL2VA checkpoint path",
+)
+def test_minimax_h3_t2va_fp8_single_gpu_smoke():
+    import torch
+
+    from vllm_omni.diffusion.data import DiffusionParallelConfig
+    from vllm_omni.entrypoints.omni import Omni
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+    if torch.accelerator.device_count() < 1:
+        pytest.skip("MiniMax H3 FP8 smoke requires a CUDA device")
+
+    engine = Omni(
+        model=os.environ[MODEL_ENV],
+        quantization_config={
+            "transformer": {"method": "fp8"},
+            "text_encoder": None,
+            "video_vae": None,
+            "audio_vae": None,
+        },
+        parallel_config=DiffusionParallelConfig(ulysses_degree=1),
+        trust_remote_code=True,
+        enable_cpu_offload=True,
+        enforce_eager=True,
+    )
+    try:
+        outputs = engine.generate(
+            "A quiet cinematic night scene with matching ambient sound.",
+            OmniDiffusionSamplingParams(
+                height=256,
+                width=448,
+                num_frames=29,
+                fps=24,
+                num_inference_steps=2,
+                seed=42,
+                output_type="np",
+                extra_args={
+                    "task": "t2va",
+                    "flow_shift": 12.0,
+                    "audio_flow_shift": 3.0,
+                },
+            ),
+            use_tqdm=False,
+        )
+    finally:
+        engine.close()
+
+    _assert_joint_output(outputs)

tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py

@@ -0,0 +1,213 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+import torch.nn as nn
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+class _FakeLinear(nn.Module):
+    def __init__(
+        self,
+        *args,
+        bias=False,
+        params_dtype=None,
+        quant_config=None,
+        prefix="",
+        total_num_heads=None,
+        total_num_kv_heads=None,
+        **kwargs,
+    ):
+        del args, kwargs
+        super().__init__()
+        self.prefix = prefix
+        self.quant_config = quant_config
+        self.weight = nn.Parameter(torch.empty(1, dtype=params_dtype))
+        if bias:
+            self.bias = nn.Parameter(torch.empty(1, dtype=params_dtype))
+        else:
+            self.register_parameter("bias", None)
+        self.num_heads = total_num_heads
+        self.num_kv_heads = total_num_kv_heads
+
+
+class _FakeAttention(nn.Module):
+    def __init__(self, **kwargs):
+        del kwargs
+        super().__init__()
+
+
+def _small_od_config():
+    arch = {
+        "num_layers": 1,
+        "token_refiner_num_layers": 1,
+        "hidden_size": 8,
+        "num_attention_heads": 2,
+        "attention_head_dim": 4,
+        "ffn_hidden_size": 16,
+        "latents_dim": 2,
+        "audio_latents_dim": 2,
+        "patch_size": (1, 2, 2),
+        "text_dim": 6,
+        "timestep_input_dim": 4,
+        "time_embed_hidden_size": 8,
+        "time_embed_dim": 4,
+        "adaln_out_features": 18 * 8,
+        "final_adaln_out_features": 2 * 8,
+        "rope_inv_freq_len": 2,
+    }
+    return SimpleNamespace(
+        tf_model_config=arch,
+        parallel_config=SimpleNamespace(ulysses_degree=1),
+    )
+
+
+def test_fp8_scope_and_prefix_propagation(monkeypatch):
+    from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
+    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
+    from vllm.model_executor.layers.quantization.utils.quant_utils import (
+        is_layer_skipped,
+    )
+
+    from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
+
+    monkeypatch.setattr(h3, "ColumnParallelLinear", _FakeLinear)
+    monkeypatch.setattr(h3, "MergedColumnParallelLinear", _FakeLinear)
+    monkeypatch.setattr(h3, "QKVParallelLinear", _FakeLinear)
+    monkeypatch.setattr(h3, "RowParallelLinear", _FakeLinear)
+    monkeypatch.setattr(h3, "Attention", _FakeAttention)
+    monkeypatch.setattr(h3, "get_tensor_model_parallel_world_size", lambda: 1)
+
+    ignored_layers = {
+        "token_refiner.blocks.0.mlp.fc2",
+        "blocks.0.attn.qkv_proj",
+        "condition_proj",
+        "blocks.0.adaln_proj.linear",
+        "final_layer.adaln_proj.linear",
+    }
+    fp8_config = Fp8Config(ignored_layers=sorted(ignored_layers))
+    model = h3.MiniMaxH3DiTModel(
+        _small_od_config(),
+        quant_config=fp8_config,
+    )
+    linears = {module.prefix: module.quant_config for module in model.modules() if isinstance(module, _FakeLinear)}
+
+    quantized = {
+        "token_refiner.blocks.0.attn.qkv_proj",
+        "token_refiner.blocks.0.attn.out_proj",
+        "token_refiner.blocks.0.mlp.fc1",
+        "token_refiner.blocks.0.mlp.fc2",
+        "blocks.0.attn.qkv_proj",
+        "blocks.0.attn.out_proj",
+        "blocks.0.mlp.fc1",
+        "blocks.0.mlp.fc2",
+        "condition_proj",
+        "blocks.0.adaln_proj.linear",
+        "final_layer.adaln_proj.linear",
+    }
+    full_precision = {
+        "video_patch_proj",
+        "audio_patch_proj",
+        "time_embedder.proj_in",
+        "time_embedder.proj_out",
+        "final_layer.video_out",
+        "final_layer.audio_out",
+    }
+
+    assert quantized <= linears.keys()
+    assert full_precision <= linears.keys()
+    assert all(linears[prefix] is fp8_config for prefix in quantized)
+    assert all(linears[prefix] is None for prefix in full_precision)
+    assert {prefix for prefix in quantized if is_layer_skipped(prefix, fp8_config.ignored_layers)} == ignored_layers
+    linear = Mock(spec=LinearBase)
+    assert all(
+        isinstance(
+            fp8_config.get_quant_method(linear, prefix),
+            UnquantizedLinearMethod,
+        )
+        for prefix in ignored_layers
+    )
+
+
+class _WeightTarget(nn.Module):
+    def __init__(self, loader):
+        super().__init__()
+        self.weight = nn.Parameter(torch.empty(1))
+        self.weight.weight_loader = loader
+
+
+def test_model_load_weights_transforms_before_calling_vllm_loader():
+    from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
+        MiniMaxH3DiTArchConfig,
+        MiniMaxH3DiTModel,
+    )
+
+    qkv_calls = []
+    fc1_calls = []
+
+    def qkv_loader(param, loaded_weight):
+        del param
+        qkv_calls.append(loaded_weight.clone())
+
+    def fc1_loader(param, loaded_weight, shard_id):
+        del param
+        fc1_calls.append((shard_id, loaded_weight.clone()))
+
+    model = object.__new__(MiniMaxH3DiTModel)
+    nn.Module.__init__(model)
+    model.arch = MiniMaxH3DiTArchConfig(
+        hidden_size=1,
+        num_attention_heads=2,
+        attention_head_dim=1,
+        ffn_hidden_size=2,
+    )
+    model.blocks = nn.ModuleList([nn.Module()])
+    model.blocks[0].attn = nn.Module()
+    model.blocks[0].attn.qkv_proj = _WeightTarget(qkv_loader)
+    model.blocks[0].mlp = nn.Module()
+    model.blocks[0].mlp.fc1 = _WeightTarget(fc1_loader)
+
+    qkv = torch.arange(6, dtype=torch.float32).reshape(6, 1)
+    fc1 = torch.arange(4, dtype=torch.float32).reshape(4, 1)
+    loaded = model.load_weights(
+        [
+            ("blocks.0.attn.qkv_proj.weight", qkv),
+            ("blocks.0.mlp.fc1.weight", fc1),
+        ]
+    )
+
+    assert loaded == {
+        "blocks.0.attn.qkv_proj.weight",
+        "blocks.0.mlp.fc1.weight",
+    }
+    assert qkv_calls[0][:, 0].tolist() == [0, 3, 1, 4, 2, 5]
+    assert [(shard_id, tensor[:, 0].tolist()) for shard_id, tensor in fc1_calls] == [
+        (0, [0, 1]),
+        (1, [2, 3]),
+    ]
+
+
+def test_pipeline_resolves_transformer_component_quant_config():
+    from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
+        _resolve_component_quant_config,
+    )
+    from vllm_omni.quantization import build_quant_config
+
+    ignored_layers = ["blocks.0.attn.qkv_proj"]
+    component_config = build_quant_config(
+        {
+            "transformer": {
+                "method": "fp8",
+                "ignored_layers": ignored_layers,
+            }
+        }
+    )
+    transformer_config = component_config.resolve("transformer")
+
+    assert transformer_config.ignored_layers == ignored_layers
+    assert _resolve_component_quant_config(component_config, "transformer") is transformer_config
+    assert _resolve_component_quant_config(transformer_config, "transformer") is transformer_config

tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py

@@ -0,0 +1,297 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+
+from tests.diffusion.quantization.test_quantization_quality import (
+    _OUTPUT_DIR,
+    QualityTestConfig,
+    _compute_lpips,
+    _compute_psnr_and_mae,
+    _free_gpu_memory,
+    _maybe_save_output,
+)
+from tests.helpers.mark import hardware_marks
+
+_MINIMAX_H3_REPO = "MiniMaxAI/MiniMax-H3"
+_MINIMAX_H3_REVISION = "48d93ede732756e404a3b1b2f3b3a9b5a22f6cfc"
+_MINIMAX_H3_AUDIO_SAMPLE_RATE = 32000
+_MIN_AUDIO_SPECTRAL_COSINE = 0.80
+_MIN_AUDIO_RMS_RATIO = 0.50
+_MAX_AUDIO_RMS_RATIO = 2.00
+
+_QUALITY_CONFIG = QualityTestConfig(
+    id="fp8_minimax_h3_fl2va",
+    model=_MINIMAX_H3_REPO,
+    quantization={"transformer": {"method": "fp8"}},
+    task="t2v",
+    prompt=(
+        "A bright daytime cinematic tracking shot of a golden retriever running through a sunlit meadow "
+        "filled with colorful wildflowers, vivid blue sky, crisp natural colors."
+    ),
+    max_lpips=0.20,
+    height=384,
+    width=672,
+    num_frames=107,
+    num_inference_steps=10,
+)
+
+
+def _resolve_fl2va_model_ref() -> str:
+    from huggingface_hub import snapshot_download
+
+    repo_root = Path(
+        snapshot_download(
+            repo_id=_MINIMAX_H3_REPO,
+            revision=_MINIMAX_H3_REVISION,
+            allow_patterns=["FL2VA/**"],
+        )
+    )
+    fl2va_root = repo_root / "FL2VA"
+    if not (fl2va_root / "model_index.json").is_file():
+        raise FileNotFoundError(f"MiniMax H3 FL2VA model_index.json not found under {repo_root}")
+    return str(fl2va_root)
+
+
+def _canonical_audio(audio: object) -> np.ndarray:
+    """Return H3 audio as float32 [channels, samples]."""
+    array = np.asarray(audio, dtype=np.float32)
+    while array.ndim > 2 and array.shape[0] == 1:
+        array = array[0]
+    if array.ndim == 1:
+        array = array[None, :]
+    if array.ndim != 2:
+        raise ValueError(f"Expected mono/stereo audio, got shape {array.shape}")
+    if array.shape[0] > array.shape[1] and array.shape[1] <= 8:
+        array = array.T
+    if not 1 <= array.shape[0] <= 8:
+        raise ValueError(f"Expected channel-first audio, got shape {array.shape}")
+    if not np.isfinite(array).all():
+        raise ValueError("Audio contains non-finite samples")
+    return np.ascontiguousarray(array)
+
+
+def _audio_quality_metrics(baseline: object, quantized: object) -> tuple[float, float]:
+    """Compare phase-tolerant spectral structure and overall signal energy."""
+    baseline_array = _canonical_audio(baseline)
+    quantized_array = _canonical_audio(quantized)
+    if baseline_array.shape != quantized_array.shape:
+        raise ValueError(
+            "Audio shapes do not match for metric computation: "
+            f"baseline={baseline_array.shape}, quantized={quantized_array.shape}"
+        )
+
+    baseline_tensor = torch.from_numpy(baseline_array)
+    quantized_tensor = torch.from_numpy(quantized_array)
+    similarities: list[float] = []
+    for n_fft in (512, 1024, 2048):
+        window = torch.hann_window(n_fft)
+        baseline_magnitude = torch.stft(
+            baseline_tensor,
+            n_fft=n_fft,
+            hop_length=n_fft // 4,
+            window=window,
+            return_complex=True,
+        ).abs()
+        quantized_magnitude = torch.stft(
+            quantized_tensor,
+            n_fft=n_fft,
+            hop_length=n_fft // 4,
+            window=window,
+            return_complex=True,
+        ).abs()
+        cosine = torch.nn.functional.cosine_similarity(
+            baseline_magnitude.flatten(1),
+            quantized_magnitude.flatten(1),
+            dim=1,
+            eps=1e-12,
+        )
+        similarities.extend(float(value) for value in cosine)
+
+    baseline_rms = float(torch.sqrt(torch.mean(baseline_tensor.square())))
+    quantized_rms = float(torch.sqrt(torch.mean(quantized_tensor.square())))
+    rms_ratio = quantized_rms / max(baseline_rms, 1e-12)
+    return float(np.mean(similarities)), rms_ratio
+
+
+def _generate_joint_output(omni, config: QualityTestConfig):
+    """Generate H3 video/audio, returning both payloads and peak GPU memory."""
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+    from vllm_omni.platforms import current_omni_platform
+
+    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(config.seed)
+    torch.accelerator.reset_peak_memory_stats()
+    outputs = omni.generate(
+        {"prompt": config.prompt, "negative_prompt": config.negative_prompt},
+        OmniDiffusionSamplingParams(
+            height=config.height,
+            width=config.width,
+            generator=generator,
+            num_inference_steps=config.num_inference_steps,
+            num_frames=config.num_frames,
+            guidance_scale=config.guidance_scale,
+            sigmas=config.sigmas,
+        ),
+    )
+    main_process_peak_mem = torch.accelerator.max_memory_allocated() / (1024**3)
+
+    first = outputs[0]
+    worker_peak_memory_mb = float(getattr(first, "peak_memory_mb", 0.0) or 0.0)
+    if hasattr(first, "request_output") and isinstance(first.request_output, list):
+        first = first.request_output[0]
+        worker_peak_memory_mb = max(
+            worker_peak_memory_mb,
+            float(getattr(first, "peak_memory_mb", 0.0) or 0.0),
+        )
+    if not hasattr(first, "images") or not first.images:
+        raise ValueError("Could not extract video frames from H3 output")
+    frames = first.images[0]
+    if isinstance(frames, torch.Tensor):
+        frames = frames.detach().cpu()
+        if frames.dim() == 5:
+            frames = frames[0]
+        if frames.dim() == 4 and frames.shape[0] in (3, 4):
+            frames = frames.permute(1, 2, 3, 0)
+        if frames.is_floating_point():
+            frames = frames.clamp(-1, 1) * 0.5 + 0.5
+        video = frames.float().numpy()
+    else:
+        video = np.asarray(frames)
+        if video.ndim == 5:
+            video = video[0]
+
+    multimodal = getattr(first, "multimodal_output", None)
+    if not multimodal or "audio" not in multimodal:
+        raise ValueError("Could not extract audio from H3 output")
+    sample_rate = int(multimodal.get("audio_sample_rate", 0))
+    peak_mem = worker_peak_memory_mb / 1024 if worker_peak_memory_mb > 0 else main_process_peak_mem
+    return video, _canonical_audio(multimodal["audio"]), sample_rate, peak_mem
+
+
+@pytest.mark.full_model
+@pytest.mark.diffusion
+@pytest.mark.parametrize(
+    "config",
+    [
+        pytest.param(
+            _QUALITY_CONFIG,
+            id=_QUALITY_CONFIG.id,
+            marks=hardware_marks(res={"cuda": _QUALITY_CONFIG.gpu}, num_cards=2),
+        )
+    ],
+)
+def test_minimax_h3_quantization_quality(config: QualityTestConfig):
+    from vllm_omni.entrypoints.omni import Omni
+
+    model_ref = _resolve_fl2va_model_ref()
+    common_kwargs = {
+        "model": model_ref,
+        "enforce_eager": True,
+        "tensor_parallel_size": 2,
+        "text_encoder_tp_size": 2,
+        "vae_use_tiling": True,
+    }
+
+    omni_bl = Omni(**common_kwargs)
+    baseline_out, baseline_audio, baseline_sample_rate, bl_mem = _generate_joint_output(omni_bl, config)
+    omni_bl.shutdown()
+    del omni_bl
+    _free_gpu_memory()
+    _maybe_save_output(_OUTPUT_DIR, config, "baseline", baseline_out)
+
+    quantization = config.quantization_ref()
+    omni_qt = Omni(**common_kwargs, quantization_config=quantization)
+    quant_out, quant_audio, quant_sample_rate, qt_mem = _generate_joint_output(omni_qt, config)
+    omni_qt.shutdown()
+    del omni_qt
+    _free_gpu_memory()
+    _maybe_save_output(_OUTPUT_DIR, config, "quantized", quant_out)
+
+    lpips_score = _compute_lpips(baseline_out, quant_out, config.task)
+    psnr_score, mae_score = _compute_psnr_and_mae(baseline_out, quant_out, config.task)
+    if baseline_sample_rate != _MINIMAX_H3_AUDIO_SAMPLE_RATE or quant_sample_rate != baseline_sample_rate:
+        raise AssertionError(
+            f"Unexpected H3 audio sample rate: baseline={baseline_sample_rate}, quantized={quant_sample_rate}"
+        )
+    audio_spectral_cosine, audio_rms_ratio = _audio_quality_metrics(baseline_audio, quant_audio)
+    assert lpips_score <= config.max_lpips, (
+        f"LPIPS {lpips_score:.4f} exceeds threshold {config.max_lpips} "
+        f"for {config.quantization_ref()} on {config.quantized_ref()}"
+    )
+    assert audio_spectral_cosine >= _MIN_AUDIO_SPECTRAL_COSINE, (
+        f"Audio spectral cosine {audio_spectral_cosine:.4f} is below threshold "
+        f"{_MIN_AUDIO_SPECTRAL_COSINE} for {config.id}"
+    )
+    assert _MIN_AUDIO_RMS_RATIO <= audio_rms_ratio <= _MAX_AUDIO_RMS_RATIO, (
+        f"Audio RMS ratio {audio_rms_ratio:.4f} is outside "
+        f"[{_MIN_AUDIO_RMS_RATIO}, {_MAX_AUDIO_RMS_RATIO}] for {config.id}"
+    )
+
+    mem_reduction = (bl_mem - qt_mem) / bl_mem * 100 if bl_mem > 0 else 0
+    print(f"\n{'=' * 60}")
+    print(f"Quantization Quality: {config.id}")
+    print(f"{'=' * 60}")
+    print(f"  Baseline:      {config.baseline_ref()}")
+    print(f"  Quantized:     {config.quantized_ref()}")
+    print(f"  Method:        {config.quantization_ref()}")
+    print(f"  LPIPS:         {lpips_score:.4f}  (threshold: {config.max_lpips})")
+    print(f"  PSNR:          {psnr_score:.4f} dB  (higher is better)")
+    print(f"  MAE:           {mae_score:.6f}  (lower is better)")
+    print(f"  Audio cosine:  {audio_spectral_cosine:.4f}  (threshold: {_MIN_AUDIO_SPECTRAL_COSINE})")
+    print(f"  Audio RMS:     {audio_rms_ratio:.4f}x  (range: {_MIN_AUDIO_RMS_RATIO}-{_MAX_AUDIO_RMS_RATIO})")
+    print(f"  BF16 memory:   {bl_mem:.2f} GiB")
+    print(f"  Quant memory:  {qt_mem:.2f} GiB  ({mem_reduction:.0f}% reduction)")
+    print("  Result:        PASS")
+    print(f"{'=' * 60}\n")
+
+    assert np.isfinite(psnr_score) or np.isinf(psnr_score), f"PSNR is invalid for {config.id}: {psnr_score}"
+    assert np.isfinite(mae_score), f"MAE is not finite for {config.id}: {mae_score}"
+
+
+def test_resolve_fl2va_model_ref(tmp_path, monkeypatch):
+    fl2va_root = tmp_path / "FL2VA"
+    fl2va_root.mkdir()
+    (fl2va_root / "model_index.json").write_text("{}", encoding="utf-8")
+
+    def fake_snapshot_download(*, repo_id, revision, allow_patterns):
+        assert repo_id == _MINIMAX_H3_REPO
+        assert revision == _MINIMAX_H3_REVISION
+        assert allow_patterns == ["FL2VA/**"]
+        return str(tmp_path)
+
+    monkeypatch.setattr("huggingface_hub.snapshot_download", fake_snapshot_download)
+    assert _resolve_fl2va_model_ref() == str(fl2va_root)
+
+
+def test_audio_quality_metrics_ignore_phase_but_reject_spectral_drift():
+    sample_rate = _MINIMAX_H3_AUDIO_SAMPLE_RATE
+    time = np.arange(sample_rate, dtype=np.float32) / sample_rate
+    baseline = np.stack(
+        [
+            np.sin(2 * np.pi * 440 * time),
+            np.sin(2 * np.pi * 880 * time),
+        ]
+    )
+    phase_shifted = np.stack(
+        [
+            np.sin(2 * np.pi * 440 * time + np.pi / 2),
+            np.sin(2 * np.pi * 880 * time + np.pi / 2),
+        ]
+    )
+    different = np.stack(
+        [
+            np.sin(2 * np.pi * 3000 * time),
+            np.sin(2 * np.pi * 4000 * time),
+        ]
+    )
+
+    phase_cosine, phase_rms_ratio = _audio_quality_metrics(baseline, phase_shifted)
+    different_cosine, _ = _audio_quality_metrics(baseline, different)
+
+    assert phase_cosine >= _MIN_AUDIO_SPECTRAL_COSINE
+    assert phase_rms_ratio == pytest.approx(1.0, abs=1e-5)
+    assert different_cosine < _MIN_AUDIO_SPECTRAL_COSINE

tests/e2e/accuracy/minimax_h3/__init__.py

@@ -0,0 +1 @@
+"""MiniMax H3 accuracy tests."""

tests/e2e/accuracy/minimax_h3/test_minimax_h3_t2va_similarity.py

@@ -0,0 +1,204 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import json
+import os
+import subprocess
+from pathlib import Path
+
+import pytest
+import requests
+import torch
+from huggingface_hub import snapshot_download
+
+from tests.e2e.accuracy.helpers import (
+    assert_video_metadata,
+    assert_video_similarity_metrics,
+    probe_binary,
+    probe_video,
+    reset_artifact_dir,
+)
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServer
+
+pytestmark = [pytest.mark.benchmark, pytest.mark.diffusion, pytest.mark.full_model]
+
+MODEL_REPO_ID = "MiniMaxAI/MiniMax-H3"
+MODEL_REVISION = "73372e6cf53e414edd3ab03e357717fb0602e758"
+MODEL_ENV_VAR = "VLLM_TEST_MINIMAX_H3_FL2VA_MODEL"
+REFERENCE_VIDEO_URL = f"https://huggingface.co/{MODEL_REPO_ID}/resolve/main/assets/t2va.mp4"
+
+# This is the official reproducible 768p T2VA prompt associated with
+# ``assets/t2va.mp4`` in the MiniMax H3 repository.
+PROMPT = """integrated_multimodal_description: [Shot 1] Cinematic, medium wide shot, pushing in slowly. In the cavernous, dimly lit bridge of a starship, sleek metallic consoles with glowing amber displays flank a massive, curved observation window. A female captain, in her late 40s with an athletic build and short silver-streaked black hair, stands in the center midground. She wears a structured, high-collared dark navy military tunic with silver chest insignias. Her back is to the camera, silhouetted against the cool, ambient starlight pouring through the thick glass. She stands perfectly still with her hands clasped tightly behind her back. Outside the window, a massive armada of jagged, dark grey dreadnoughts hovers in tight formation against a deep purple space nebula. The fleet's massive rear thrusters begin to glow with an intense, escalating bright blue light. [Shot 2] At 00:04.500, the camera cuts to a close-up of the captain's face and shakes strongly. The brilliant blue-white light from the fleet's gathering energy reflects vividly in her dark eyes. Suddenly, a blinding white flash floods through the window, completely washing out the background as the fleet jumps to hyperspace. The sheer spatial force violently jolts the bridge, causing the captain from Shot 1 to stagger slightly forward, her shoulders tensing as she visibly braces herself against the physical tremors. As the intense white light fades abruptly, leaving only the dim, empty expanse of the purple nebula reflected on her starkly lit skin, her jaw clenches, and she slowly closes her eyes in the newly emptied space.
+overall_soundscape: A low, resonant hum of the ship's ambient life support systems serves as the baseline, soon drowned out by an audible, escalating, high-pitched electronic whine as the fleet outside charges its hyperdrives. A massive, deafening, bass-heavy boom and sharp crackle erupts during the blinding flash, accompanied by the loud metallic creaking, rattling, and deep thuds of the bridge's bulkheads vibrating under immense physical stress. The intense roaring impact then cuts abruptly back to a hollow, echoing room tone, leaving only the faint, steady hum of the isolated bridge.
+non_diegetic_music: Cinematic space-opera orchestral score, slow tempo, featuring a solitary, mournful French horn melody over deep, sustained string dissonances that build rapidly in volume and intensity, swelling to a massive orchestral peak before snapping immediately into silence right after the jump."""
+
+WIDTH = 1344
+HEIGHT = 768
+FPS = 24
+NUM_FRAMES = 243
+NUM_INFERENCE_STEPS = 50
+FLOW_SHIFT = 12.0
+AUDIO_FLOW_SHIFT = 3.0
+DURATION_SECONDS = 10.0
+SEED = 0
+SSIM_THRESHOLD = 0.82
+PSNR_THRESHOLD = 20.0
+REQUEST_TIMEOUT_SECONDS = 60 * 60
+
+
+def _model_name() -> str:
+    configured = os.environ.get(MODEL_ENV_VAR)
+    if configured:
+        return configured
+
+    snapshot_root = snapshot_download(
+        repo_id=MODEL_REPO_ID,
+        revision=MODEL_REVISION,
+        allow_patterns=["FL2VA/**"],
+    )
+    return str(Path(snapshot_root) / "FL2VA")
+
+
+def _server_args() -> list[str]:
+    return [
+        "--trust-remote-code",
+        "--num-gpus",
+        "4",
+        "--usp",
+        "4",
+        "--ring",
+        "1",
+        "--use-hsdp",
+        "--hsdp-shard-size",
+        "4",
+        "--text-encoder-tp-size",
+        "4",
+        "--vae-patch-parallel-size",
+        "4",
+        "--vae-parallel-mode",
+        "tile",
+        "--vae-use-tiling",
+        "--diffusion-attention-backend",
+        "FLASH_ATTN",
+        "--stage-init-timeout",
+        "1800",
+        "--init-timeout",
+        "1800",
+    ]
+
+
+def _download_reference_video(output_dir: Path) -> Path:
+    reference_path = output_dir / "reference.mp4"
+    response = requests.get(REFERENCE_VIDEO_URL, timeout=120)
+    response.raise_for_status()
+    reference_path.write_bytes(response.content)
+    return reference_path
+
+
+def _probe_audio(path: Path) -> dict[str, str | int]:
+    probe_binary("ffprobe")
+    result = subprocess.run(
+        [
+            "ffprobe",
+            "-v",
+            "error",
+            "-select_streams",
+            "a:0",
+            "-show_entries",
+            "stream=codec_name,sample_rate,channels",
+            "-of",
+            "json",
+            str(path),
+        ],
+        capture_output=True,
+        text=True,
+        check=True,
+    )
+    streams = json.loads(result.stdout).get("streams", [])
+    assert len(streams) == 1, f"Expected one audio stream in {path}, got {streams}"
+    stream = streams[0]
+    return {
+        "codec_name": str(stream["codec_name"]),
+        "sample_rate": int(stream["sample_rate"]),
+        "channels": int(stream["channels"]),
+    }
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=4)
+def test_minimax_h3_t2va_matches_official_reference(
+    accuracy_artifact_root: Path,
+) -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("MiniMax H3 T2VA accuracy test requires CUDA.")
+
+    probe_binary("ffmpeg")
+    probe_binary("ffprobe")
+    output_dir = reset_artifact_dir(accuracy_artifact_root / "minimax_h3_t2va")
+    reference_path = _download_reference_video(output_dir)
+    generated_path = output_dir / "vllm_omni.mp4"
+
+    server_env = {
+        "FLASHINFER_DISABLE_VERSION_CHECK": "1",
+        "VLLM_WORKER_MULTIPROC_METHOD": "spawn",
+        "VLLM_OMNI_VIDEO_SYNC_TIMEOUT": str(REQUEST_TIMEOUT_SECONDS),
+        "VLLM_OMNI_STORAGE_PATH": str(output_dir / "storage"),
+    }
+    request_data = {
+        "prompt": PROMPT,
+        "width": str(WIDTH),
+        "height": str(HEIGHT),
+        "fps": str(FPS),
+        "num_inference_steps": str(NUM_INFERENCE_STEPS),
+        "flow_shift": str(FLOW_SHIFT),
+        "seed": str(SEED),
+        "extra_params": json.dumps(
+            {
+                "task": "t2va",
+                "duration": DURATION_SECONDS,
+                "audio_flow_shift": AUDIO_FLOW_SHIFT,
+            }
+        ),
+    }
+
+    with OmniServer(
+        _model_name(),
+        _server_args(),
+        env_dict=server_env,
+        use_omni=True,
+    ) as server:
+        response = requests.post(
+            f"http://{server.host}:{server.port}/v1/videos/sync",
+            data=request_data,
+            timeout=REQUEST_TIMEOUT_SECONDS,
+        )
+
+    response.raise_for_status()
+    assert response.headers["content-type"].startswith("video/mp4")
+    generated_path.write_bytes(response.content)
+
+    reference_metadata = probe_video(reference_path)
+    generated_metadata = probe_video(generated_path)
+    assert reference_metadata == generated_metadata
+    assert_video_metadata(
+        generated_metadata,
+        width=WIDTH,
+        height=HEIGHT,
+        fps=FPS,
+        frame_count=NUM_FRAMES,
+    )
+
+    expected_audio = {"codec_name": "aac", "sample_rate": 32000, "channels": 2}
+    assert _probe_audio(reference_path) == expected_audio
+    assert _probe_audio(generated_path) == expected_audio
+
+    assert_video_similarity_metrics(
+        label="minimax_h3_t2va_official_reference",
+        online_path=generated_path,
+        offline_path=reference_path,
+        ssim_threshold=SSIM_THRESHOLD,
+        psnr_threshold=PSNR_THRESHOLD,
+    )

tests/e2e/features/comfyui/test_comfyui_integration.py

@@ -25,9 +25,15 @@ from comfyui_vllm_omni.nodes import (
     VLLMOmniGenerateVideo,
     VLLMOmniTTS,
     VLLMOmniUnderstanding,
+    VLLMOmniVideoReferences,
     VLLMOmniVoiceClone,
 )
-from comfyui_vllm_omni.utils.types import AutoregressionSamplingParams, DiffusionSamplingParams, WanModelSpecificParams
+from comfyui_vllm_omni.utils.types import (
+    AutoregressionSamplingParams,
+    DiffusionSamplingParams,
+    MiniMaxH3ModelSpecificParams,
+    WanModelSpecificParams,
+)
 from PIL import Image
 from pytest_mock import MockerFixture
 from vllm import SamplingParams
@@ -68,6 +74,8 @@ class SamplingKind(str, Enum):
     TTS_DIFFUSION_SINGLE = "tts_diffusion_single"
     VIDEO_NONE = "video_none"
     VIDEO_DIFFUSION_SINGLE = "video_diffusion_single"
+    VIDEO_REF2VA_IMAGE_AUDIO = "video_ref2va_image_audio"
+    VIDEO_REF2VA_MULTI_VIDEO = "video_ref2va_multi_video"
 
 
 # Pre-defined arguments to be used in function calls during the tests
@@ -127,6 +135,15 @@ AR_LIST_SAMPLING_PARAMS = [
 
 VIDEO_MODEL_PARAMS = WanModelSpecificParams({"guidance_scale_2": 5.0, "boundary_ratio": 0.98, "flow_shift": 12.0})
 
+# Matches MiniMax-H3 recipe Ref2VA defaults (task is auto-routed by the client).
+H3_MODEL_PARAMS = MiniMaxH3ModelSpecificParams(
+    {
+        "type": "minimax_h3",
+        "audio_flow_shift": 3.0,
+        "flow_shift": 12.0,
+    }
+)
+
 LORA_PARAMS = {"local_path": "test_lora_path", "name": "test_name", "scale": 0.7, "int_id": 10}
 
 
@@ -306,6 +323,7 @@ def _build_mock_outputs(outputs: Iterable[OmniRequestOutput], sampling_case: Sam
         received_sampling_params_list: Sequence[OmniSamplingParams] | None = (
             args[2] if len(args) > 2 else kwargs.get("sampling_params_list")
         )
+        prompt = args[0] if len(args) > 0 else kwargs.get("prompt")
 
         assert received_sampling_params_list is not None, (
             "In the current codebase, the API layer always provides not-None sampling parameter list when calling `AsyncOmni.generate`"
@@ -372,6 +390,59 @@ def _build_mock_outputs(outputs: Iterable[OmniRequestOutput], sampling_case: Sam
                 LORA_PARAMS,
             )
             _assert_model_param_values(received_sampling_params_list[0], VIDEO_MODEL_PARAMS)
+        elif sampling_case.kind is SamplingKind.VIDEO_REF2VA_IMAGE_AUDIO:
+            assert len(received_sampling_params_list) == 1
+            assert isinstance(prompt, dict)
+            multi_modal_data = prompt.get("multi_modal_data")
+            assert isinstance(multi_modal_data, dict)
+            assert isinstance(multi_modal_data.get("image"), Image.Image)
+            assert isinstance(multi_modal_data.get("audio"), str)
+            assert "video" not in multi_modal_data
+            _assert_sampling_param_values(
+                received_sampling_params_list[0],
+                {
+                    "width": VIDEO_WIDTH,
+                    "height": VIDEO_HEIGHT,
+                    "num_frames": VIDEO_NUM_FRAMES,
+                    "fps": VIDEO_FPS,
+                },
+            )
+            _assert_model_param_values(
+                received_sampling_params_list[0],
+                {
+                    "flow_shift": 12.0,
+                    "task": "ref2va",
+                    "audio_flow_shift": 3.0,
+                },
+            )
+        elif sampling_case.kind is SamplingKind.VIDEO_REF2VA_MULTI_VIDEO:
+            assert len(received_sampling_params_list) == 1
+            assert isinstance(prompt, dict)
+            multi_modal_data = prompt.get("multi_modal_data")
+            assert isinstance(multi_modal_data, dict)
+            input_videos = multi_modal_data.get("video")
+            assert isinstance(input_videos, list)
+            assert len(input_videos) == 2
+            assert all(isinstance(path, str) for path in input_videos)
+            assert "image" not in multi_modal_data
+            assert "audio" not in multi_modal_data
+            _assert_sampling_param_values(
+                received_sampling_params_list[0],
+                {
+                    "width": VIDEO_WIDTH,
+                    "height": VIDEO_HEIGHT,
+                    "num_frames": VIDEO_NUM_FRAMES,
+                    "fps": VIDEO_FPS,
+                },
+            )
+            _assert_model_param_values(
+                received_sampling_params_list[0],
+                {
+                    "flow_shift": 12.0,
+                    "task": "ref2va",
+                    "audio_flow_shift": 3.0,
+                },
+            )
         else:
             raise AssertionError(f"Unknown sampling case: {sampling_case.kind}")
 
@@ -775,7 +846,7 @@ async def test_video_generation_node(api_server: str, model: str, image_input: b
         "model_params": VIDEO_MODEL_PARAMS,
     }
     if image_input:
-        kwargs["image"] = torch.zeros((1, VIDEO_HEIGHT, VIDEO_WIDTH, 3), dtype=torch.float32)
+        kwargs["frame"] = torch.zeros((1, VIDEO_HEIGHT, VIDEO_WIDTH, 3), dtype=torch.float32)
     if sampling_case.sampling_params is not None:
         kwargs["sampling_params"] = sampling_case.sampling_params
     if sampling_case.lora:
@@ -786,3 +857,79 @@ async def test_video_generation_node(api_server: str, model: str, image_input: b
     assert isinstance(result, tuple)
     assert len(result) == 1
     assert isinstance(result[0], VideoInput)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "server_case,sampling_case,ref_mode",
+    [
+        pytest.param(
+            ServerCase(
+                served_model="MiniMaxAI/MiniMax-H3",
+                stage_list=["diffusion"],
+                stage_configs=[{"stage_type": "diffusion"}],
+                outputs=[_build_diffusion_video_output()],
+            ),
+            SamplingCase(kind=SamplingKind.VIDEO_REF2VA_IMAGE_AUDIO, sampling_params=None),
+            "image_audio",
+            id="ref2va-image-audio",
+        ),
+        pytest.param(
+            ServerCase(
+                served_model="MiniMaxAI/MiniMax-H3",
+                stage_list=["diffusion"],
+                stage_configs=[{"stage_type": "diffusion"}],
+                outputs=[_build_diffusion_video_output()],
+            ),
+            SamplingCase(kind=SamplingKind.VIDEO_REF2VA_MULTI_VIDEO, sampling_params=None),
+            "multi_video",
+            id="ref2va-multi-video",
+        ),
+    ],
+    indirect=["server_case", "sampling_case"],
+)
+async def test_video_generation_node_minimax_h3_ref2va(
+    api_server: str,
+    sampling_case: SamplingCase,
+    ref_mode: str,
+):
+    """MiniMax-H3 Ref2VA paths matching recipes/MiniMaxAI/MiniMax-H3.md."""
+    node = VLLMOmniGenerateVideo()
+    refs_node = VLLMOmniVideoReferences()
+
+    if ref_mode == "image_audio":
+        (references,) = refs_node.get_references(
+            image_1=torch.zeros((1, VIDEO_HEIGHT, VIDEO_WIDTH, 3), dtype=torch.float32),
+            audio_1={"waveform": torch.zeros((1, 1, 24000), dtype=torch.float32), "sample_rate": 24000},
+        )
+        prompt = (
+            "A white cat with black mustache and eyebrow markings sits on a beige couch, "
+            "lip-syncing precisely to the complete reference audio."
+        )
+    elif ref_mode == "multi_video":
+        (references,) = refs_node.get_references(
+            video_1=VideoInput(b"subject-video"),  # type: ignore[reportAbstractUsage]
+            video_2=VideoInput(b"background-video"),  # type: ignore[reportAbstractUsage]
+        )
+        prompt = (
+            "Remove the green screen background of Video 1 and replace it with the fairytale environment from Video 2."
+        )
+    else:
+        raise AssertionError(f"Unknown ref_mode: {ref_mode}")
+
+    result = await node.generate(
+        url=api_server,
+        model="MiniMaxAI/MiniMax-H3",
+        prompt=prompt,
+        negative_prompt="",
+        width=VIDEO_WIDTH,
+        height=VIDEO_HEIGHT,
+        fps=VIDEO_FPS,
+        num_frames=VIDEO_NUM_FRAMES,
+        references=references,
+        model_params=H3_MODEL_PARAMS,
+    )
+
+    assert isinstance(result, tuple)
+    assert len(result) == 1
+    assert isinstance(result[0], VideoInput)

tests/entrypoints/test_omni_new_request_data.py

@@ -2,7 +2,9 @@ from types import SimpleNamespace
 
 import pytest
 import torch
+from vllm.v1.core.sched.output import NewRequestData
 
+from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
 from vllm_omni.core.sched.output import OmniNewRequestData
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -52,3 +54,80 @@ def test_omni_new_request_data_allows_missing_payloads():
 
     assert data.prompt_embeds is None
     assert data.additional_information is None
+
+
+def test_from_base_preserves_every_upstream_field_and_adds_omni_payloads():
+    prompt_embeds = torch.randn(2, 3)
+    base = NewRequestData(
+        req_id="req-base",
+        prompt_token_ids=[1, 2],
+        mm_features=[],
+        sampling_params=None,
+        pooling_params=None,
+        block_ids=([3, 4],),
+        num_computed_tokens=1,
+        lora_request=None,
+        prompt_embeds=prompt_embeds,
+        prompt_is_token_ids=False,
+        prefill_token_ids=[1],
+    )
+    request = SimpleNamespace(
+        external_req_id="external-base",
+        additional_information={"payload": "value"},
+        model_intermediate_buffer={"hidden": "state"},
+    )
+
+    data = OmniNewRequestData.from_base(base, request)
+
+    for field_name in NewRequestData.__dataclass_fields__:
+        assert getattr(data, field_name) is getattr(base, field_name)
+    assert data.external_req_id == "external-base"
+    assert data.additional_information == {"payload": "value"}
+    assert data.model_intermediate_buffer == {"hidden": "state"}
+
+
+def test_rewrap_converts_fallback_entries_and_keeps_fast_path_entries():
+    base = NewRequestData(
+        req_id="req-fallback",
+        prompt_token_ids=[1],
+        mm_features=[],
+        sampling_params=None,
+        pooling_params=None,
+        block_ids=([2],),
+        num_computed_tokens=0,
+        lora_request=None,
+        prompt_embeds=None,
+        prompt_is_token_ids=True,
+        prefill_token_ids=[1],
+    )
+    fast_path = OmniNewRequestData(
+        req_id="req-fast",
+        prompt_token_ids=[3],
+        mm_features=[],
+        sampling_params=None,
+        pooling_params=None,
+        block_ids=([4],),
+        num_computed_tokens=0,
+        lora_request=None,
+        prompt_embeds=None,
+        prompt_is_token_ids=True,
+        prefill_token_ids=[3],
+        additional_information={"fast": True},
+    )
+    scheduler = OmniSchedulerMixin()
+    scheduler.requests = {
+        "req-fallback": SimpleNamespace(
+            external_req_id="external-fallback",
+            additional_information={"fallback": True},
+            model_intermediate_buffer=None,
+        )
+    }
+    output = SimpleNamespace(scheduled_new_reqs=[base, fast_path])
+
+    scheduler._rewrap_scheduled_new_reqs(output)
+
+    fallback_data, fast_path_data = output.scheduled_new_reqs
+    assert isinstance(fallback_data, OmniNewRequestData)
+    assert fallback_data.prefill_token_ids == [1]
+    assert fallback_data.additional_information == {"fallback": True}
+    assert fast_path_data is fast_path

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -2,37 +2,25 @@ from __future__ import annotations
 
 from collections import defaultdict
 from collections.abc import Iterable
-from time import time
 from typing import Any
 
 import numpy as np
 from vllm.compilation.cuda_graph import CUDAGraphStat
-from vllm.distributed.kv_events import KVEventBatch
-from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
 from vllm.logger import init_logger
 from vllm.v1.core.sched.async_scheduler import AsyncScheduler as AsyncVLLMScheduler
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.core.sched.request_queue import create_request_queue
 from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
-from vllm.v1.core.sched.utils import remove_all
-from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs, FinishReason
+from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs
 from vllm.v1.metrics.perf import PerfStats
 from vllm.v1.outputs import ModelRunnerOutput
 from vllm.v1.request import Request, RequestStatus, StreamingUpdate
 from vllm.v1.spec_decode.metrics import SpecDecodingStats
 
 from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
-from vllm_omni.core.sched.omni_scheduling_coordinator import (
-    OmniSchedulingCoordinator,
-    uses_full_payload_input_coordinator,
-)
 from vllm_omni.core.sched.utils import omni_routed_experts_for_request
-from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import (
-    OmniChunkTransferAdapter,
-)
 from vllm_omni.engine import OmniEngineCoreOutput
 from vllm_omni.engine.serialization import deserialize_additional_information
-from vllm_omni.outputs import OmniConnectorOutput
 
 logger = init_logger(__name__)
 
@@ -108,6 +96,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         self.finished_req_ids_dict = defaultdict(set)
 
         # [Omni] Pre-parse KV transfer criteria
+        self._omni_kv_config = getattr(self.vllm_config.model_config, "omni_kv_config", None)
         self.kv_transfer_criteria = self._get_kv_transfer_criteria()
 
         # Track requests that have already triggered prefill transfer to avoid duplicates
@@ -115,16 +104,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
 
         # Cache per-request flag to avoid repeated deserialization of additional_information
         self._omits_kv_transfer_cache: dict[str, bool] = {}
-        model_config = self.vllm_config.model_config
-        self.chunk_transfer_adapter = None
-        if getattr(model_config, "async_chunk", False):
-            self.chunk_transfer_adapter = OmniChunkTransferAdapter(self.vllm_config)
-        self.input_coordinator: OmniSchedulingCoordinator | None = None
-        if uses_full_payload_input_coordinator(model_config):
-            self.input_coordinator = OmniSchedulingCoordinator(
-                stage_id=getattr(model_config, "stage_id", 0),
-            )
-        self._latest_omni_connector_output: OmniConnectorOutput | None = None
+        self._init_omni_io_scheduling_state()
         # Snapshot prompt length for each streaming input update
         self._new_prompt_len_snapshot: dict[str, int] = {}
 
@@ -134,17 +114,15 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         return request.num_computed_tokens - request.num_output_placeholders
 
     def _get_kv_transfer_criteria(self) -> dict | None:
-        # Note: vllm_config is available in Scheduler after super().__init__
-        if not hasattr(self, "vllm_config"):
-            return None
+        return self._get_omni_kv_config_value("kv_transfer_criteria")
 
-        omni_kv_config = getattr(self.vllm_config.model_config, "omni_kv_config", None)
-        if omni_kv_config:
-            if isinstance(omni_kv_config, dict):
-                return omni_kv_config.get("kv_transfer_criteria", None)
-            else:
-                return getattr(omni_kv_config, "kv_transfer_criteria", None)
-        return None
+    def _get_omni_kv_config_value(self, key: str, default: Any = None) -> Any:
+        config = getattr(self, "_omni_kv_config", None)
+        if config is None and hasattr(self, "vllm_config"):
+            config = getattr(self.vllm_config.model_config, "omni_kv_config", None)
+        if isinstance(config, dict):
+            return config.get(key, default)
+        return getattr(config, key, default) if config is not None else default
 
     def _request_omits_kv_transfer_to_next_stage(self, request: Request) -> bool:
         """True when this stage-zero-final request does not need downstream KV.
@@ -207,42 +185,39 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
 
         if criteria_type == "prefill_finished":
             if confirmed_computed >= request.num_prompt_tokens:
-                self.transfer_triggered_requests.add(request.request_id)
-
-                self._mark_request_for_kv_transfer(request.request_id, confirmed_computed)
-                actually_queued = request.request_id in self.requests_needing_kv_transfer
-
-                if stop_decode_on_trigger and actually_queued:
-                    # Defer the stop until KV extraction completes so that
-                    # the kv_ready signal can be emitted while the request
-                    # is still alive.  The request will be stopped on the
-                    # next scheduler step after extraction ack arrives.
-                    self.pending_stop_after_extraction.add(request.request_id)
-
+                self._commit_kv_transfer_trigger(
+                    request.request_id,
+                    confirmed_computed,
+                    stop_decode_on_trigger,
+                )
                 return False
 
         elif criteria_type == "special_token":
             target_token_id = self.kv_transfer_criteria.get("token_id")
             if target_token_id is not None and target_token_id in new_token_ids:
-                self.transfer_triggered_requests.add(request.request_id)
-
-                try:
-                    idx = new_token_ids.index(target_token_id)
-                    tokens_to_exclude = len(new_token_ids) - (idx + 1)
-                    snapshot_len = confirmed_computed - tokens_to_exclude
-                except ValueError:
-                    snapshot_len = confirmed_computed
-
-                self._mark_request_for_kv_transfer(request.request_id, snapshot_len)
-                actually_queued = request.request_id in self.requests_needing_kv_transfer
-
-                if stop_decode_on_trigger and actually_queued:
-                    self.pending_stop_after_extraction.add(request.request_id)
-
+                idx = new_token_ids.index(target_token_id)
+                tokens_to_exclude = len(new_token_ids) - (idx + 1)
+                snapshot_len = confirmed_computed - tokens_to_exclude
+                self._commit_kv_transfer_trigger(
+                    request.request_id,
+                    snapshot_len,
+                    stop_decode_on_trigger,
+                )
                 return False
 
         return False
 
+    def _commit_kv_transfer_trigger(
+        self,
+        req_id: str,
+        seq_len: int,
+        stop_after_transfer: bool,
+    ) -> None:
+        self.transfer_triggered_requests.add(req_id)
+        self._mark_request_for_kv_transfer(req_id, seq_len)
+        if stop_after_transfer and req_id in self.requests_needing_kv_transfer:
+            self.pending_stop_after_extraction.add(req_id)
+
     def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
         # Remove FINISHED_ABORTED requests before the upstream scheduler sees
         # them. Upstream vllm raises RuntimeError on this status; omni allows
@@ -252,12 +227,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
             for req in list(queue):
                 if getattr(req, "status", None) == RequestStatus.FINISHED_ABORTED:
                     queue.remove(req)
-        self._consume_pending_connector_output(model_mode="ar")
-        self._process_pending_input_timeouts()
-        if self.chunk_transfer_adapter:
-            self.chunk_transfer_adapter.process_pending_chunks(
-                self.waiting, self.running, scheduler_requests=self.requests
-            )
+        self._process_pending_omni_inputs(model_mode="ar")
 
         original_waiting = None
         if self._should_defer_waiting_admission():
@@ -272,55 +242,13 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                 if deferred_waiting:
                     original_waiting.prepend_requests(deferred_waiting)
                 self.waiting = original_waiting
-            if self.chunk_transfer_adapter:
-                # Add request waiting for chunk to the waiting and running queue
-                self.chunk_transfer_adapter.restore_queues(
-                    self.waiting,
-                    self.running,
-                    scheduler_requests=self.requests,
-                )
-            if self.input_coordinator:
-                self.input_coordinator.restore_queues(self.waiting)
-        try:
-            # Late import to avoid circulars in some launch modes
-            from .output import OmniNewRequestData
+            self._restore_omni_wait_queues()
 
-            # Rewrap base NewRequestData entries with OmniNewRequestData,
-            # enriching with request-level payloads
-            new_list = []
-            for nr in scheduler_output.scheduled_new_reqs:
-                req_id = getattr(nr, "req_id", None)
-                request = self.requests.get(req_id) if req_id else None
-                # Build omni entry preserving all base fields
-                omni_nr = OmniNewRequestData(
-                    req_id=nr.req_id,
-                    external_req_id=(getattr(request, "external_req_id", None) if request else None),
-                    prompt_token_ids=nr.prompt_token_ids,
-                    mm_features=nr.mm_features,
-                    sampling_params=nr.sampling_params,
-                    pooling_params=nr.pooling_params,
-                    block_ids=nr.block_ids,
-                    num_computed_tokens=nr.num_computed_tokens,
-                    lora_request=nr.lora_request,
-                    # Enrich with omni payloads from the live request object
-                    prompt_embeds=(getattr(request, "prompt_embeds", None) if request else None),
-                    prompt_is_token_ids=nr.prompt_is_token_ids,
-                    additional_information=(getattr(request, "additional_information", None) if request else None),
-                    model_intermediate_buffer=(
-                        getattr(request, "model_intermediate_buffer", None) if request else None
-                    ),
-                )
-                new_list.append(omni_nr)
-
-            scheduler_output.scheduled_new_reqs = new_list  # type: ignore[assignment]
-            if self.chunk_transfer_adapter:
-                self.chunk_transfer_adapter.postprocess_scheduler_output(scheduler_output, self.requests)
-            # Add information about requests needing KV cache transfer
-            finished_reqs = self.get_finished_requests_needing_kv_transfer()
-        except Exception:
-            # If anything goes wrong, leave the original output unchanged
-            init_logger(__name__).exception("Failed to wrap scheduled_new_reqs with OmniNewRequestData")
-            finished_reqs = {}
+        self._postprocess_omni_schedule_output(
+            scheduler_output,
+            include_cached_payloads=True,
+        )
+        finished_reqs = self.get_finished_requests_needing_kv_transfer()
 
         # Wrap in omni scheduler output to carry transfer metadata.
         return self._wrap_omni_scheduler_output(
@@ -550,26 +478,23 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
             # Get prompt logprobs for this request.
             prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
             if new_token_ids or mm_output is not None or pooler_output is not None or kv_transfer_params or stopped:
-                # Add EngineCoreOutput for this Request.
-                outputs[request.client_index].append(
-                    OmniEngineCoreOutput(
-                        request_id=req_id,
-                        new_token_ids=new_token_ids,
-                        finish_reason=finish_reason,
-                        new_logprobs=new_logprobs,
-                        new_prompt_logprobs_tensors=prompt_logprobs_tensors,
-                        pooling_output=pooler_output,
-                        multimodal_output=mm_output,
-                        stop_reason=request.stop_reason,
-                        events=request.take_events(),
-                        prefill_stats=request.take_prefill_stats(),
-                        kv_transfer_params=kv_transfer_params,
-                        trace_headers=request.trace_headers,
-                        routed_experts=routed_experts,
-                        num_nans_in_logits=request.num_nans_in_logits,
-                        is_segment_finished=is_segment_finished,
-                        new_prompt_len_snapshot=self._new_prompt_len_snapshot.get(req_id, None),
-                    )
+                OmniSchedulerMixin._append_request_output(
+                    self,
+                    outputs,
+                    request,
+                    new_token_ids=new_token_ids,
+                    finish_reason=finish_reason,
+                    new_logprobs=new_logprobs,
+                    new_prompt_logprobs_tensors=prompt_logprobs_tensors,
+                    pooling_output=pooler_output,
+                    multimodal_output=mm_output,
+                    stop_reason=request.stop_reason,
+                    prefill_stats=request.take_prefill_stats(),
+                    kv_transfer_params=kv_transfer_params,
+                    routed_experts=routed_experts,
+                    num_nans_in_logits=request.num_nans_in_logits,
+                    is_segment_finished=is_segment_finished,
+                    new_prompt_len_snapshot=self._new_prompt_len_snapshot.get(req_id),
                 )
             else:
                 # Invariant: EngineCore returns no partial prefill outputs.
@@ -584,117 +509,44 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                     is_segment_finished,
                 )
 
-        # Remove the stopped requests from the running and waiting queues.
-        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)
-            self.skipped_waiting.remove_requests(stopped_preempted_reqs)
+        self._remove_stopped_requests_from_queues(
+            stopped_running_reqs,
+            stopped_preempted_reqs,
+        )
 
-        # [Main] Handle failed KV load requests
-        if failed_kv_load_req_ids and not self.recompute_kv_load_failures:
-            requests = [self.requests[req_id] for req_id in failed_kv_load_req_ids]
-            self.finish_requests(failed_kv_load_req_ids, RequestStatus.FINISHED_ERROR)
-            for request in requests:
-                outputs[request.client_index].append(
-                    OmniEngineCoreOutput(
-                        request_id=request.request_id,
-                        new_token_ids=[],
-                        finish_reason=request.get_finished_reason(),
-                        events=request.take_events(),
-                        trace_headers=request.trace_headers,
-                    )
-                )
-                if self.chunk_transfer_adapter is not None:
-                    self.chunk_transfer_adapter.cleanup_receiver(
-                        request.request_id,
-                    )
+        failed_requests = self._handle_failed_kv_load_outputs(
+            failed_kv_load_req_ids,
+            outputs,
+        )
+        if self.chunk_transfer_adapter is not None:
+            for request in failed_requests:
+                self.chunk_transfer_adapter.cleanup_receiver(request.request_id)
 
-        # [Omni] Cleanup state for finished requests
-        for req in stopped_running_reqs:
-            if req.request_id not in self.waiting_for_transfer_free:
-                if req.request_id in self.transfer_triggered_requests:
-                    self.transfer_triggered_requests.remove(req.request_id)
-                if req.request_id in self.active_kv_transfers:
-                    self.active_kv_transfers.remove(req.request_id)
-                self.pending_stop_after_extraction.discard(req.request_id)
-
-        # Same for preempted
-        for req in stopped_preempted_reqs:
-            if req.request_id not in self.waiting_for_transfer_free:
-                if req.request_id in self.transfer_triggered_requests:
-                    self.transfer_triggered_requests.remove(req.request_id)
-                if req.request_id in self.active_kv_transfers:
-                    self.active_kv_transfers.remove(req.request_id)
-                self.pending_stop_after_extraction.discard(req.request_id)
+        self._cleanup_kv_tracking(req.request_id for req in stopped_running_reqs | stopped_preempted_reqs)
 
         # KV Connector: update state for finished KV Transfers.
         if kv_connector_output:
             self._update_from_kv_xfer_finished(kv_connector_output)
 
-        # Worker-side KV connector stats from the model runner output.
-        kv_connector_stats: KVConnectorStats | None = (
-            kv_connector_output.kv_connector_stats if kv_connector_output else None
-        )
-        if self.connector:
-            # Scheduler-side KV connector stats collected after connector update.
-            scheduler_kv_connector_stats = self.connector.get_kv_connector_stats()
-            if scheduler_kv_connector_stats is not None and not scheduler_kv_connector_stats.is_empty():
-                kv_connector_stats = (
-                    kv_connector_stats.aggregate(scheduler_kv_connector_stats)
-                    if kv_connector_stats is not None
-                    else scheduler_kv_connector_stats
-                )
-
-        # collect KV cache events from KV cache manager
-        events = self.kv_cache_manager.take_events()
-
-        # collect KV cache events from connector
-        if self.connector is not None:
-            connector_events = self.connector.take_events()
-            if connector_events:
-                if events is None:
-                    events = list(connector_events)
-                else:
-                    events.extend(connector_events)
-
-        # publish collected KV cache events
-        if events:
-            batch = KVEventBatch(ts=time(), events=events)
-            self.kv_event_publisher.publish(batch)
+        kv_connector_stats = self._aggregate_kv_connector_stats(kv_connector_output)
+        self._publish_kv_cache_events()
 
         # 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()}
 
-        # FIXME: finished_req_ids_dict is unconditionally initialized as
-        # defaultdict(set) in __init__ (not gated by include_finished_set).
-        # This branch is therefore always eligible once any client_index is
-        # populated; revisit when wiring streaming-only / upstream semantics.
-        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():
-                eco = engine_core_outputs.get(client_index)
-                if eco is None:
-                    eco = EngineCoreOutputs()
-                    engine_core_outputs[client_index] = eco
-                emitted = {o.request_id for o in eco.outputs}
-                for req_id in finished_set:
-                    if req_id not in emitted:
-                        eco.outputs.append(EngineCoreOutput(req_id, [], finish_reason=FinishReason.ABORT))
-                eco.finished_requests = finished_set
-            finished_req_ids.clear()
+        self._attach_finished_request_sets(
+            engine_core_outputs,
+            synthesize_abort_outputs=True,
+        )
 
-        if (stats := self.make_stats(spec_decoding_stats, kv_connector_stats, cudagraph_stats, perf_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
+        self._attach_scheduler_stats(
+            engine_core_outputs,
+            spec_decoding_stats,
+            kv_connector_stats,
+            cudagraph_stats,
+            perf_stats,
+        )
 
         self._capture_omni_connector_output(model_runner_output)
 
@@ -720,56 +572,6 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
 
         return engine_core_outputs
 
-    def finish_requests(self, request_ids: str | Iterable[str] | None, finished_status: RequestStatus) -> list[Request]:
-        """Handles the finish signal from outside the scheduler.
-
-        For example, the API server can abort a request when the client
-        disconnects.
-
-        If request_ids is None, all requests will be finished.
-
-        Returns:
-            The Request objects that were aborted. Will not include any that
-            were already finished.
-        """
-        # TODO(yrr): chunk transfer adapter & input_coordinator unified to one
-        if self.chunk_transfer_adapter:
-            self.chunk_transfer_adapter.finish_requests(request_ids, finished_status, self.requests)
-
-        # Realign stale ``request.status`` (chunk-transfer-adapter's
-        # ``requests_origin_status`` table doesn't follow the
-        # ``waiting → running`` admit transition; without this, an abort
-        # arriving between admit and the next deque round-trip leaves
-        # the request in ``self.running`` with ``status=WAITING`` and
-        # upstream ``Scheduler.finish_requests`` silently fails to
-        # release the worker's ``input_batch`` slot -- after
-        # ``max_num_seqs`` such aborts new requests hang at
-        # ``chunks=0``). Only the ``async_chunk`` path triggers the
-        # staleness; with ``async_chunk`` disabled this is a cheap O(n)
-        # no-op over an already-aligned set, kept unconditional so the
-        # abort path stays uniform across configurations. See
-        # ``OmniSchedulerMixin._realign_request_status_to_queues`` and
-        # #3774 discussion.
-        self._realign_request_status_to_queues(request_ids)
-
-        finished = super().finish_requests(request_ids, finished_status)
-
-        # Defensive post-finish purge: belt-and-suspenders to the
-        # realignment above. Even after realign + ``super()``, corner
-        # cases (mid-transition status, connector cleanups that pop
-        # from ``self.requests`` without unwinding ``self.running``)
-        # can leave already-finished or untracked entries in
-        # ``self.running``. Sweep them now so the worker's
-        # ``input_batch`` slot never pins a freed request and starves
-        # new admissions. See ``OmniSchedulerMixin._purge_finished_from_running``.
-        self._purge_finished_from_running()
-
-        input_coordinator = getattr(self, "input_coordinator", None)
-        if input_coordinator is not None:
-            for request in finished:
-                self._free_input_coordinator_request(request.request_id)
-        return finished
-
     def _update_request_as_session(self, session: Request, update: StreamingUpdate) -> None:
         """
         Override: Only extend prompt at stage 0, and replace
@@ -967,44 +769,35 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
 
     def _should_transfer_kv_for_request(self, req_id: str) -> bool:
         """Determine if a request should trigger KV cache transfer."""
-        need_send = False
-        # Try to read from vLLM Config (where YAML config is typically loaded)
-        # Check for omni_kv_config attribute
-        omni_kv_config = getattr(self.vllm_config.model_config, "omni_kv_config", None)
-        if omni_kv_config:
-            # omni_kv_config could be an object or a dict
-            if isinstance(omni_kv_config, dict):
-                need_send = omni_kv_config.get("need_send_cache", False)
-            else:
-                need_send = getattr(omni_kv_config, "need_send_cache", False)
-        if not need_send:
+        if not self._get_omni_kv_config_value("need_send_cache", False):
             return False
         request = self.requests.get(req_id)
         if request is not None and self._request_omits_kv_transfer_to_next_stage(request):
             return False
         return True
 
+    def _cleanup_kv_tracking(self, request_ids: Iterable[str]) -> None:
+        for req_id in request_ids:
+            if req_id in self.waiting_for_transfer_free:
+                continue
+            self.transfer_triggered_requests.discard(req_id)
+            self.active_kv_transfers.discard(req_id)
+            self.pending_stop_after_extraction.discard(req_id)
+
+    def _has_pending_kv_work(self) -> bool:
+        return bool(self.requests_needing_kv_transfer or self.active_kv_transfers or self.waiting_for_transfer_free)
+
     def has_requests(self) -> bool:
         """Check if there are any requests to process, including KV transfers."""
-        # [Omni] Also check for pending KV transfers
-        if self.requests_needing_kv_transfer or self.active_kv_transfers or self.waiting_for_transfer_free:
-            return True
-        return super().has_requests()
+        return self._has_pending_kv_work() or super().has_requests()
 
     def has_finished_requests(self) -> bool:
         """Check if there are any finished requests (including those needing KV transfer)."""
-        if self.requests_needing_kv_transfer or self.active_kv_transfers or self.waiting_for_transfer_free:
-            return True
-        return super().has_finished_requests()
+        return self._has_pending_kv_work() or super().has_finished_requests()
 
     def has_unfinished_requests(self) -> bool:
         """Check if there are any unfinished requests (including those needing KV transfer)."""
-        # [Omni] Also check for pending KV transfers to ensure the engine loop continues
-        # MUST verify waiting_for_transfer_free and active_kv_transfers
-        # Otherwise engine loop might exit before transfer Ack is received.
-        if self.requests_needing_kv_transfer or self.active_kv_transfers or self.waiting_for_transfer_free:
-            return True
-        return super().has_unfinished_requests()
+        return self._has_pending_kv_work() or super().has_unfinished_requests()
 
     def get_finished_requests_needing_kv_transfer(self) -> dict[str, dict]:
         """Get and clear the list of requests needing KV cache transfer.

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -2,13 +2,9 @@ from __future__ import annotations
 
 import time
 from collections import defaultdict
-from collections.abc import Iterable
 from typing import Any
 
 from vllm.compilation.cuda_graph import CUDAGraphStat
-from vllm.distributed.kv_events import KVEventBatch
-from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
-from vllm.logger import init_logger
 from vllm.v1.core.kv_cache_manager import KVCacheBlocks
 from vllm.v1.core.sched.interface import PauseState
 from vllm.v1.core.sched.output import SchedulerOutput
@@ -25,36 +21,18 @@ from vllm.v1.request import Request, RequestStatus, StreamingUpdate
 from vllm.v1.spec_decode.metrics import SpecDecodingStats
 
 from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
-from vllm_omni.core.sched.omni_scheduling_coordinator import (
-    OmniSchedulingCoordinator,
-    uses_full_payload_input_coordinator,
-)
 from vllm_omni.core.sched.output import OmniCachedRequestData, OmniNewRequestData
 from vllm_omni.core.sched.utils import omni_routed_experts_for_request
-from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import (
-    OmniChunkTransferAdapter,
-)
-from vllm_omni.engine import OmniEngineCoreOutput
-from vllm_omni.outputs import OmniConnectorOutput, OmniModelRunnerOutput
-
-logger = init_logger(__name__)
+from vllm_omni.outputs import OmniModelRunnerOutput
 
 
 class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         model_config = self.vllm_config.model_config
-        self.chunk_transfer_adapter = None
-        if getattr(model_config, "async_chunk", False):
-            self.chunk_transfer_adapter = OmniChunkTransferAdapter(self.vllm_config)
+        self._init_omni_io_scheduling_state()
         self._retains_state_across_chunks = bool(getattr(model_config, "retains_state_across_chunks", False))
         self._pending_finish_reqs: list[Request] = []
-        self.input_coordinator: OmniSchedulingCoordinator | None = None
-        if uses_full_payload_input_coordinator(model_config):
-            self.input_coordinator = OmniSchedulingCoordinator(
-                stage_id=getattr(model_config, "stage_id", 0),
-            )
-        self._latest_omni_connector_output: OmniConnectorOutput | None = None
 
     @staticmethod
     def _record_prefill_stats(request: Request) -> None:
@@ -117,12 +95,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
         # Temporary queue: preserve waiting order while requests await input.
         skipped_waiting_requests = create_request_queue(self.policy)
         req_index = 0
-        self._consume_pending_connector_output(model_mode="generation")
-        self._process_pending_input_timeouts()
-        if self.chunk_transfer_adapter:
-            self.chunk_transfer_adapter.process_pending_chunks(
-                self.waiting, self.running, scheduler_requests=self.requests
-            )
+        self._process_pending_omni_inputs(model_mode="generation")
 
         # OMNI: Track requests that are already finished (e.g., marked by connector)
         # These should be removed from running and not scheduled
@@ -250,15 +223,11 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
             if self.chunk_transfer_adapter:
                 # Don't fall back: base scheduler doesn't handle async_chunk
                 # requests with empty prompt_token_ids.
-                self.chunk_transfer_adapter.restore_queues(
-                    self.waiting,
-                    self.running,
-                    scheduler_requests=self.requests,
-                )
+                self._restore_omni_wait_queues()
             else:
                 res = super().schedule(throttle_prefills)
-                if self.input_coordinator:
-                    self.input_coordinator.restore_queues(self.waiting)
+                self._restore_omni_wait_queues()
+                self._postprocess_omni_schedule_output(res)
                 return self._wrap_omni_scheduler_output(res)
 
         # Compute common prefix blocks (aligned with v1)
@@ -342,94 +311,12 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
         self._update_after_schedule(scheduler_output)
 
         try:
-            # Rewrap base NewRequestData entries with OmniNewRequestData,
-            # enriching with request-level payloads
-            new_list = []
-            for nr in scheduler_output.scheduled_new_reqs:
-                req_id = getattr(nr, "req_id", None)
-                request = self.requests.get(req_id) if req_id else None
-                # Build omni entry preserving all base fields
-                omni_nr = OmniNewRequestData(
-                    req_id=nr.req_id,
-                    external_req_id=(getattr(request, "external_req_id", None) if request else None),
-                    prompt_token_ids=nr.prompt_token_ids,
-                    mm_features=nr.mm_features,
-                    sampling_params=nr.sampling_params,
-                    pooling_params=nr.pooling_params,
-                    block_ids=nr.block_ids,
-                    num_computed_tokens=nr.num_computed_tokens,
-                    lora_request=nr.lora_request,
-                    # Enrich with omni payloads from the live request object
-                    prompt_embeds=(getattr(request, "prompt_embeds", None) if request else None),
-                    prompt_is_token_ids=nr.prompt_is_token_ids,
-                    additional_information=(getattr(request, "additional_information", None) if request else None),
-                    model_intermediate_buffer=(
-                        getattr(request, "model_intermediate_buffer", None) if request else None
-                    ),
-                )
-                new_list.append(omni_nr)
-
-            scheduler_output.scheduled_new_reqs = new_list  # type: ignore[assignment]
-
-            if self.chunk_transfer_adapter:
-                self.chunk_transfer_adapter.postprocess_scheduler_output(scheduler_output)
-
-        except Exception:
-            # If anything goes wrong, leave the original output unchanged
-            logger.exception("Failed to wrap scheduled_new_reqs with OmniNewRequestData")
+            self._postprocess_omni_schedule_output(scheduler_output)
         finally:
-            # Ensure chunk-waiting requests are restored even on error,
-            # otherwise they are permanently orphaned in the adapter's
-            # internal deques and never scheduled again.
-            if self.chunk_transfer_adapter:
-                self.chunk_transfer_adapter.restore_queues(
-                    self.waiting,
-                    self.running,
-                    scheduler_requests=self.requests,
-                )
-            if self.input_coordinator:
-                self.input_coordinator.restore_queues(self.waiting)
+            self._restore_omni_wait_queues()
 
         return self._wrap_omni_scheduler_output(scheduler_output)
 
-    def finish_requests(self, request_ids: str | Iterable[str] | None, finished_status: RequestStatus) -> list[Request]:
-        """Handles the finish signal from outside the scheduler.
-
-        For example, the API server can abort a request when the client
-        disconnects.
-
-        If request_ids is None, all requests will be finished.
-
-        Returns:
-            The Request objects that were aborted. Will not include any that
-            were already finished.
-        """
-
-        if self.chunk_transfer_adapter:
-            self.chunk_transfer_adapter.finish_requests(request_ids, finished_status, self.requests)
-
-        # See ``OmniSchedulerMixin._realign_request_status_to_queues`` --
-        # closes the residual hang on the
-        # ``waiting → running → abort-before-next-deque-round-trip`` race
-        # surfaced by #3774's reproduction matrix. Only the
-        # ``async_chunk`` path actually triggers the staleness; with
-        # ``async_chunk`` disabled this is a cheap O(n) no-op, kept
-        # unconditional so the abort path stays uniform.
-        self._realign_request_status_to_queues(request_ids)
-
-        finished = super().finish_requests(request_ids, finished_status)
-
-        # See ``OmniSchedulerMixin._purge_finished_from_running`` --
-        # defensive belt-and-suspenders sweep paired with the realign
-        # above. Closes residual ``self.running`` slot leaks if any
-        # corner case slips past upstream's status-driven removal.
-        self._purge_finished_from_running()
-
-        if self.input_coordinator is not None:
-            for request in finished:
-                self._free_input_coordinator_request(request.request_id)
-        return finished
-
     def _free_request(
         self, request: Request, delay_free_blocks: bool = False
     ) -> tuple[dict[str, Any] | None, dict[str, Any] | None]:
@@ -595,25 +482,22 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
             # Get prompt logprobs for this request.
             prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
             if new_token_ids or mm_output is not None or pooler_output is not None or kv_transfer_params or stopped:
-                # Add EngineCoreOutput for this Request.
-                outputs[request.client_index].append(
-                    OmniEngineCoreOutput(
-                        request_id=req_id,
-                        new_token_ids=new_token_ids,
-                        finish_reason=finish_reason,
-                        new_logprobs=new_logprobs,
-                        new_prompt_logprobs_tensors=prompt_logprobs_tensors,
-                        pooling_output=pooler_output,
-                        multimodal_output=mm_output,
-                        stop_reason=request.stop_reason,
-                        events=request.take_events(),
-                        prefill_stats=request.take_prefill_stats(),
-                        kv_transfer_params=kv_transfer_params,
-                        trace_headers=request.trace_headers,
-                        routed_experts=routed_experts,
-                        num_nans_in_logits=request.num_nans_in_logits,
-                        is_segment_finished=is_segment_finished,
-                    )
+                OmniSchedulerMixin._append_request_output(
+                    self,
+                    outputs,
+                    request,
+                    new_token_ids=new_token_ids,
+                    finish_reason=finish_reason,
+                    new_logprobs=new_logprobs,
+                    new_prompt_logprobs_tensors=prompt_logprobs_tensors,
+                    pooling_output=pooler_output,
+                    multimodal_output=mm_output,
+                    stop_reason=request.stop_reason,
+                    prefill_stats=request.take_prefill_stats(),
+                    kv_transfer_params=kv_transfer_params,
+                    routed_experts=routed_experts,
+                    num_nans_in_logits=request.num_nans_in_logits,
+                    is_segment_finished=is_segment_finished,
                 )
             else:
                 # Invariant: EngineCore returns no partial prefill outputs.
@@ -636,103 +520,58 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                         request.request_id,
                         getattr(request, "external_req_id", None),
                     )
-            outputs[request.client_index].append(
-                OmniEngineCoreOutput(
-                    request_id=request.request_id,
-                    new_token_ids=[],
-                    finish_reason=finish_reason,
-                    stop_reason=request.stop_reason,
-                    events=request.take_events(),
-                    kv_transfer_params=kv_transfer_params,
-                    trace_headers=request.trace_headers,
-                    is_segment_finished=is_segment_finished,
-                )
+            OmniSchedulerMixin._append_request_output(
+                self,
+                outputs,
+                request,
+                new_token_ids=[],
+                finish_reason=finish_reason,
+                stop_reason=request.stop_reason,
+                kv_transfer_params=kv_transfer_params,
+                is_segment_finished=is_segment_finished,
             )
             stopped_running_reqs.add(request)
         self._pending_finish_reqs.clear()
 
-        # Remove the stopped requests from the running and waiting queues.
-        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)
-            self.skipped_waiting.remove_requests(stopped_preempted_reqs)
+        self._remove_stopped_requests_from_queues(
+            stopped_running_reqs,
+            stopped_preempted_reqs,
+        )
 
-        # Handle failed KV load requests
-        if failed_kv_load_req_ids and not self.recompute_kv_load_failures:
-            requests = [self.requests[req_id] for req_id in failed_kv_load_req_ids]
-            self.finish_requests(failed_kv_load_req_ids, RequestStatus.FINISHED_ERROR)
-            for request in requests:
-                outputs[request.client_index].append(
-                    OmniEngineCoreOutput(
-                        request_id=request.request_id,
-                        new_token_ids=[],
-                        finish_reason=request.get_finished_reason(),
-                        events=request.take_events(),
-                        trace_headers=request.trace_headers,
-                    )
+        failed_requests = self._handle_failed_kv_load_outputs(
+            failed_kv_load_req_ids,
+            outputs,
+        )
+        if self.chunk_transfer_adapter is not None:
+            for request in failed_requests:
+                self.chunk_transfer_adapter.cleanup(
+                    request.request_id,
+                    getattr(request, "external_req_id", None),
                 )
-                if self.chunk_transfer_adapter is not None:
-                    self.chunk_transfer_adapter.cleanup(
-                        request.request_id,
-                        getattr(request, "external_req_id", None),
-                    )
 
         # KV Connector: update state for finished KV Transfers.
         if kv_connector_output:
             self._update_from_kv_xfer_finished(kv_connector_output)
 
-        # Worker-side KV connector stats from the model runner output.
-        kv_connector_stats: KVConnectorStats | None = (
-            kv_connector_output.kv_connector_stats if kv_connector_output else None
-        )
-        if self.connector:
-            # Scheduler-side KV connector stats collected after connector update.
-            scheduler_kv_connector_stats = self.connector.get_kv_connector_stats()
-            if scheduler_kv_connector_stats is not None and not scheduler_kv_connector_stats.is_empty():
-                kv_connector_stats = (
-                    kv_connector_stats.aggregate(scheduler_kv_connector_stats)
-                    if kv_connector_stats is not None
-                    else scheduler_kv_connector_stats
-                )
-
-        # Collect and publish KV cache events (align with v0.14.0)
-        events = self.kv_cache_manager.take_events()
-        if self.connector is not None:
-            connector_events = self.connector.take_events()
-            if connector_events:
-                if events is None:
-                    events = list(connector_events)
-                else:
-                    events.extend(connector_events)
-        if events:
-            batch = KVEventBatch(ts=time.time(), events=events)
-            self.kv_event_publisher.publish(batch)
+        kv_connector_stats = self._aggregate_kv_connector_stats(kv_connector_output)
+        self._publish_kv_cache_events()
 
         # 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()
+        self._attach_finished_request_sets(
+            engine_core_outputs,
+            synthesize_abort_outputs=False,
+        )
 
-        if (stats := self.make_stats(spec_decoding_stats, kv_connector_stats, cudagraph_stats, perf_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
+        self._attach_scheduler_stats(
+            engine_core_outputs,
+            spec_decoding_stats,
+            kv_connector_stats,
+            cudagraph_stats,
+            perf_stats,
+        )
 
         self._capture_omni_connector_output(model_runner_output)
 

vllm_omni/core/sched/omni_scheduler_mixin.py

@@ -5,13 +5,36 @@ import time
 from collections.abc import Iterable
 from typing import Any
 
+from vllm.compilation.cuda_graph import CUDAGraphStat
+from vllm.distributed.kv_events import KVEventBatch
+from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
 from vllm.logger import init_logger
 from vllm.v1.core.sched.output import SchedulerOutput
-from vllm.v1.engine import EngineCoreEventType
+from vllm.v1.core.sched.utils import remove_all
+from vllm.v1.engine import (
+    EngineCoreEventType,
+    EngineCoreOutput,
+    EngineCoreOutputs,
+    FinishReason,
+)
+from vllm.v1.metrics.perf import PerfStats
 from vllm.v1.metrics.stats import SchedulerStats
 from vllm.v1.request import Request, RequestStatus, StreamingUpdate
+from vllm.v1.spec_decode.metrics import SpecDecodingStats
 
-from vllm_omni.core.sched.output import OmniChunkRecvHandle, OmniSchedulerOutput
+from vllm_omni.core.sched.omni_scheduling_coordinator import (
+    OmniSchedulingCoordinator,
+    uses_full_payload_input_coordinator,
+)
+from vllm_omni.core.sched.output import (
+    OmniChunkRecvHandle,
+    OmniNewRequestData,
+    OmniSchedulerOutput,
+)
+from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import (
+    OmniChunkTransferAdapter,
+)
+from vllm_omni.engine import OmniEngineCoreOutput
 
 logger = init_logger(__name__)
 
@@ -45,6 +68,19 @@ class OmniSchedulerMixin:
     #  Shared scheduler/output helpers (lift the AR / generation duplicates)
     # ------------------------------------------------------------------ #
 
+    def _init_omni_io_scheduling_state(self) -> None:
+        """Initialize scheduler state shared by AR and generation stages."""
+        model_config = self.vllm_config.model_config
+        self.chunk_transfer_adapter = (
+            OmniChunkTransferAdapter(self.vllm_config) if getattr(model_config, "async_chunk", False) else None
+        )
+        self.input_coordinator = (
+            OmniSchedulingCoordinator(stage_id=getattr(model_config, "stage_id", 0))
+            if uses_full_payload_input_coordinator(model_config)
+            else None
+        )
+        self._latest_omni_connector_output = None
+
     def _free_input_coordinator_request(self, request_id: str) -> None:
         """Prune full-payload coordinator state for a completed request."""
         input_coordinator = getattr(self, "input_coordinator", None)
@@ -102,6 +138,28 @@ class OmniSchedulerMixin:
             connector_output.stage_recv_req_ids if connector_output else set(),
         )
 
+    def _process_pending_omni_inputs(self, model_mode: str) -> None:
+        """Apply pending connector inputs, timeouts, and async chunks."""
+        self._consume_pending_connector_output(model_mode)
+        self._process_pending_input_timeouts()
+        if self.chunk_transfer_adapter:
+            self.chunk_transfer_adapter.process_pending_chunks(
+                self.waiting,
+                self.running,
+                scheduler_requests=self.requests,
+            )
+
+    def _restore_omni_wait_queues(self) -> None:
+        """Restore requests temporarily parked by Omni input gates."""
+        if self.chunk_transfer_adapter:
+            self.chunk_transfer_adapter.restore_queues(
+                self.waiting,
+                self.running,
+                scheduler_requests=self.requests,
+            )
+        if self.input_coordinator:
+            self.input_coordinator.restore_queues(self.waiting)
+
     def _process_pending_input_timeouts(self) -> None:
         """Force-fail requests waiting on the full-payload coordinator too long.
 
@@ -181,6 +239,201 @@ class OmniSchedulerMixin:
             pending_input_registrations=pending_input_registrations,
         )
 
+    def _rewrap_scheduled_new_reqs(self, scheduler_output: SchedulerOutput) -> None:
+        """Attach Omni payloads without reconstructing existing Omni entries."""
+        scheduler_output.scheduled_new_reqs = [  # type: ignore[assignment]
+            data
+            if isinstance(data, OmniNewRequestData)
+            else OmniNewRequestData.from_base(data, self.requests.get(data.req_id))
+            for data in scheduler_output.scheduled_new_reqs
+        ]
+
+    def _postprocess_omni_schedule_output(
+        self,
+        scheduler_output: SchedulerOutput,
+        *,
+        include_cached_payloads: bool = False,
+    ) -> None:
+        """Enrich new requests and apply async-chunk output bookkeeping."""
+        self._rewrap_scheduled_new_reqs(scheduler_output)
+        if not self.chunk_transfer_adapter:
+            return
+        if include_cached_payloads:
+            self.chunk_transfer_adapter.postprocess_scheduler_output(
+                scheduler_output,
+                self.requests,
+            )
+        else:
+            self.chunk_transfer_adapter.postprocess_scheduler_output(scheduler_output)
+
+    def _make_omni_engine_output(
+        self,
+        request: Request,
+        *,
+        new_token_ids: list[int],
+        finish_reason: Any = None,
+        new_logprobs: Any = None,
+        new_prompt_logprobs_tensors: Any = None,
+        pooling_output: Any = None,
+        multimodal_output: Any = None,
+        stop_reason: Any = None,
+        prefill_stats: Any = None,
+        kv_transfer_params: Any = None,
+        routed_experts: Any = None,
+        num_nans_in_logits: int = 0,
+        is_segment_finished: bool | None = False,
+        new_prompt_len_snapshot: int | None = None,
+    ) -> OmniEngineCoreOutput:
+        """Build the common request-output envelope used by LLM schedulers."""
+        return OmniEngineCoreOutput(
+            request_id=request.request_id,
+            new_token_ids=new_token_ids,
+            finish_reason=finish_reason,
+            new_logprobs=new_logprobs,
+            new_prompt_logprobs_tensors=new_prompt_logprobs_tensors,
+            pooling_output=pooling_output,
+            multimodal_output=multimodal_output,
+            stop_reason=stop_reason,
+            events=request.take_events(),
+            prefill_stats=prefill_stats,
+            kv_transfer_params=kv_transfer_params,
+            trace_headers=request.trace_headers,
+            routed_experts=routed_experts,
+            num_nans_in_logits=num_nans_in_logits,
+            is_segment_finished=is_segment_finished,
+            new_prompt_len_snapshot=new_prompt_len_snapshot,
+        )
+
+    def _append_request_output(
+        self,
+        outputs: dict[int, list[EngineCoreOutput]],
+        request: Request,
+        **output_fields: Any,
+    ) -> None:
+        outputs[request.client_index].append(
+            OmniSchedulerMixin._make_omni_engine_output(
+                self,
+                request,
+                **output_fields,
+            )
+        )
+
+    def _handle_failed_kv_load_outputs(
+        self,
+        failed_request_ids: set[str] | None,
+        outputs: dict[int, list[EngineCoreOutput]],
+    ) -> list[Request]:
+        """Finish unrecoverable KV loads and emit their terminal outputs."""
+        if not failed_request_ids or self.recompute_kv_load_failures:
+            return []
+        requests = [self.requests[req_id] for req_id in failed_request_ids]
+        self.finish_requests(failed_request_ids, RequestStatus.FINISHED_ERROR)
+        for request in requests:
+            OmniSchedulerMixin._append_request_output(
+                self,
+                outputs,
+                request,
+                new_token_ids=[],
+                finish_reason=request.get_finished_reason(),
+            )
+        return requests
+
+    def _attach_finished_request_sets(
+        self,
+        engine_core_outputs: dict[int, EngineCoreOutputs],
+        *,
+        synthesize_abort_outputs: bool,
+    ) -> None:
+        """Attach finished IDs while keeping AR's synthetic-abort policy explicit."""
+        finished_req_ids = self.finished_req_ids_dict
+        if not finished_req_ids:
+            return
+        for client_index, finished_set in finished_req_ids.items():
+            output = engine_core_outputs.get(client_index)
+            if output is None:
+                output = EngineCoreOutputs()
+                engine_core_outputs[client_index] = output
+            if synthesize_abort_outputs:
+                emitted = {item.request_id for item in output.outputs}
+                output.outputs.extend(
+                    EngineCoreOutput(req_id, [], finish_reason=FinishReason.ABORT)
+                    for req_id in finished_set
+                    if req_id not in emitted
+                )
+            output.finished_requests = finished_set
+        finished_req_ids.clear()
+
+    def _remove_stopped_requests_from_queues(
+        self,
+        stopped_running_reqs: set[Request],
+        stopped_preempted_reqs: set[Request],
+    ) -> None:
+        if stopped_running_reqs:
+            self.running = remove_all(self.running, stopped_running_reqs)
+        if stopped_preempted_reqs:
+            self.waiting.remove_requests(stopped_preempted_reqs)
+            self.skipped_waiting.remove_requests(stopped_preempted_reqs)
+
+    def _aggregate_kv_connector_stats(
+        self,
+        kv_connector_output: Any,
+    ) -> KVConnectorStats | None:
+        stats = kv_connector_output.kv_connector_stats if kv_connector_output else None
+        if self.connector:
+            scheduler_stats = self.connector.get_kv_connector_stats()
+            if scheduler_stats is not None and not scheduler_stats.is_empty():
+                stats = stats.aggregate(scheduler_stats) if stats is not None else scheduler_stats
+        return stats
+
+    def _publish_kv_cache_events(self) -> None:
+        events = self.kv_cache_manager.take_events()
+        if self.connector is not None:
+            connector_events = self.connector.take_events()
+            if connector_events:
+                if events is None:
+                    events = list(connector_events)
+                else:
+                    events.extend(connector_events)
+        if events:
+            self.kv_event_publisher.publish(KVEventBatch(ts=time.time(), events=events))
+
+    def _attach_scheduler_stats(
+        self,
+        engine_core_outputs: dict[int, EngineCoreOutputs],
+        spec_decoding_stats: SpecDecodingStats | None,
+        kv_connector_stats: KVConnectorStats | None,
+        cudagraph_stats: CUDAGraphStat | None,
+        perf_stats: PerfStats | None,
+    ) -> None:
+        stats = self.make_stats(
+            spec_decoding_stats,
+            kv_connector_stats,
+            cudagraph_stats,
+            perf_stats,
+        )
+        if stats is None:
+            return
+        if (output := next(iter(engine_core_outputs.values()), None)) is None:
+            engine_core_outputs[0] = output = EngineCoreOutputs()
+        output.scheduler_stats = stats
+
+    def finish_requests(
+        self,
+        request_ids: str | Iterable[str] | None,
+        finished_status: RequestStatus,
+    ) -> list[Request]:
+        """Finish requests and clean all Omni-owned queue/coordinator state."""
+        if self.chunk_transfer_adapter:
+            self.chunk_transfer_adapter.finish_requests(request_ids, finished_status, self.requests)
+
+        self._realign_request_status_to_queues(request_ids)
+        finished = super().finish_requests(request_ids, finished_status)
+        self._purge_finished_from_running()
+
+        for request in finished:
+            self._free_input_coordinator_request(request.request_id)
+        return finished
+
     def make_stats(self, *args, **kwargs) -> SchedulerStats | None:
         now = time.monotonic()
         if now - getattr(self, "_last_stats_time", 0.0) < _STATS_INTERVAL_S:

vllm_omni/core/sched/output.py

@@ -1,4 +1,4 @@
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, fields
 
 from vllm.v1.core.sched.output import CachedRequestData, NewRequestData, SchedulerOutput
 from vllm.v1.request import Request
@@ -28,6 +28,21 @@ class OmniNewRequestData(NewRequestData):
     additional_information: AdditionalInformationPayload | None = None
     model_intermediate_buffer: dict[str, object] | None = None
 
+    @classmethod
+    def from_base(
+        cls,
+        data: NewRequestData,
+        request: Request | None,
+    ) -> "OmniNewRequestData":
+        """Preserve upstream request data while attaching Omni payloads."""
+        base_data = {field.name: getattr(data, field.name) for field in fields(NewRequestData)}
+        return cls(
+            **base_data,
+            external_req_id=getattr(request, "external_req_id", None),
+            additional_information=getattr(request, "additional_information", None),
+            model_intermediate_buffer=getattr(request, "model_intermediate_buffer", None),
+        )
+
     @classmethod
     def from_request(
         cls,

vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py

@@ -248,7 +248,8 @@ class MiniMaxH3TimeEmbedder(nn.Module):
     def __init__(
         self,
         arch: MiniMaxH3DiTArchConfig,
-        quant_config: QuantizationConfig | None,
+        *,
+        prefix: str,
     ) -> None:
         super().__init__()
         self.frequency_embedding_size = arch.timestep_input_dim
@@ -258,7 +259,8 @@ class MiniMaxH3TimeEmbedder(nn.Module):
             bias=True,
             gather_output=True,
             params_dtype=_FP32_DTYPE,
-            quant_config=quant_config,
+            quant_config=None,
+            prefix=f"{prefix}.proj_in",
         )
         self.proj_out = RowParallelLinear(
             arch.time_embed_hidden_size,
@@ -266,7 +268,8 @@ class MiniMaxH3TimeEmbedder(nn.Module):
             bias=True,
             input_is_parallel=False,
             params_dtype=_FP32_DTYPE,
-            quant_config=quant_config,
+            quant_config=None,
+            prefix=f"{prefix}.proj_out",
         )
 
     def forward(self, t: torch.Tensor) -> torch.Tensor:
@@ -323,6 +326,7 @@ class MiniMaxH3Attention(nn.Module):
         arch: MiniMaxH3DiTArchConfig,
         quant_config: QuantizationConfig | None,
         *,
+        prefix: str,
         skip_sequence_parallel: bool = False,
     ) -> None:
         super().__init__()
@@ -338,11 +342,11 @@ class MiniMaxH3Attention(nn.Module):
             bias=False,
             params_dtype=_BF16_DTYPE,
             quant_config=quant_config,
+            prefix=f"{prefix}.qkv_proj",
             return_bias=True,
         )
         self.num_heads = self.qkv_proj.num_heads
         self.num_kv_heads = self.qkv_proj.num_kv_heads
-        self._install_qkv_weight_loader(arch)
         self.q_norm = _norm(arch.attention_head_dim, eps=arch.qk_norm_eps)
         self.k_norm = _norm(arch.attention_head_dim, eps=arch.qk_norm_eps)
         self.out_proj = RowParallelLinear(
@@ -352,6 +356,7 @@ class MiniMaxH3Attention(nn.Module):
             input_is_parallel=True,
             params_dtype=_BF16_DTYPE,
             quant_config=quant_config,
+            prefix=f"{prefix}.out_proj",
         )
         self.attention = Attention(
             num_heads=self.num_heads,
@@ -362,24 +367,6 @@ class MiniMaxH3Attention(nn.Module):
             skip_sequence_parallel=skip_sequence_parallel,
         )
 
-    def _install_qkv_weight_loader(self, arch: MiniMaxH3DiTArchConfig) -> None:
-        base_loader = self.qkv_proj.weight.weight_loader
-
-        def _weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
-            # The grouped checkpoint layout is
-            # [num_query_groups, q_per_group + k + v] before splitting.
-            # MiniMax H3 uses MHA, so checkpoint rows are per-head [q, k, v],
-            # while qkv_proj expects [q_all, k_all, v_all].
-            reordered = _reorder_grouped_qkv_to_qkv(
-                loaded_weight,
-                num_query_groups=arch.num_attention_heads,
-                heads_per_group=1,
-                head_dim=arch.attention_head_dim,
-            )
-            base_loader(param, reordered)
-
-        self.qkv_proj.weight.weight_loader = _weight_loader
-
     @torch.compiler.disable
     def _run_packed_attention(
         self,
@@ -472,6 +459,8 @@ class MiniMaxH3MLP(nn.Module):
         self,
         arch: MiniMaxH3DiTArchConfig,
         quant_config: QuantizationConfig | None,
+        *,
+        prefix: str,
     ) -> None:
         super().__init__()
         self.fc1 = MergedColumnParallelLinear(
@@ -481,8 +470,8 @@ class MiniMaxH3MLP(nn.Module):
             gather_output=False,
             params_dtype=_BF16_DTYPE,
             quant_config=quant_config,
+            prefix=f"{prefix}.fc1",
         )
-        self._install_fc1_weight_loader()
         # Chunk the fused fc1 output as [gate, up], then compute
         # silu(gate) * up.
         self.fc2 = RowParallelLinear(
@@ -492,23 +481,9 @@ class MiniMaxH3MLP(nn.Module):
             input_is_parallel=True,
             params_dtype=_BF16_DTYPE,
             quant_config=quant_config,
+            prefix=f"{prefix}.fc2",
         )
 
-    def _install_fc1_weight_loader(self) -> None:
-        base_loader = self.fc1.weight.weight_loader
-
-        def _weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
-            if loaded_weight.shape[0] % 2:
-                raise ValueError(
-                    "MiniMax H3 fc1 checkpoint rows must split evenly into "
-                    f"gate/up matrices, got {tuple(loaded_weight.shape)}"
-                )
-            gate, up = loaded_weight.chunk(2, dim=0)
-            base_loader(param, gate, 0)
-            base_loader(param, up, 1)
-
-        self.fc1.weight.weight_loader = _weight_loader
-
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         hidden, _ = self.fc1(x)
         gate, up = hidden.chunk(2, dim=-1)
@@ -534,6 +509,7 @@ class MiniMaxH3AdalnProj(nn.Module):
         *,
         expand_ratio: int,
         modality_num: int,
+        prefix: str,
     ) -> None:
         super().__init__()
         if out_features != expand_ratio * arch.hidden_size * modality_num:
@@ -550,12 +526,13 @@ class MiniMaxH3AdalnProj(nn.Module):
             gather_output=True,
             params_dtype=_BF16_DTYPE,
             quant_config=quant_config,
+            prefix=f"{prefix}.linear",
         )
 
     def forward(self, t_emb: torch.Tensor) -> tuple[torch.Tensor, ...]:
         """t_emb: [M, t_dim] -> expand_ratio tensors of [M*modality_num, H]."""
         x = nn.functional.silu(t_emb)
-        x, _ = self.linear(x.to(self.linear.weight.dtype))
+        x, _ = self.linear(x.to(_BF16_DTYPE))
         m = x.shape[0]
         x = x.view(m * self.modality_num, self.expand_ratio * self.hidden_size)
         return tuple(x.chunk(self.expand_ratio, dim=-1))
@@ -568,6 +545,8 @@ class MiniMaxH3TokenRefinerBlock(nn.Module):
         self,
         arch: MiniMaxH3DiTArchConfig,
         quant_config: QuantizationConfig | None,
+        *,
+        prefix: str,
     ) -> None:
         super().__init__()
         self.norm1 = _norm(arch.hidden_size, eps=arch.norm_eps)
@@ -578,9 +557,14 @@ class MiniMaxH3TokenRefinerBlock(nn.Module):
         self.attn = MiniMaxH3Attention(
             arch,
             quant_config,
+            prefix=f"{prefix}.attn",
             skip_sequence_parallel=True,
         )
-        self.mlp = MiniMaxH3MLP(arch, quant_config)
+        self.mlp = MiniMaxH3MLP(
+            arch,
+            quant_config,
+            prefix=f"{prefix}.mlp",
+        )
 
     def forward(
         self,
@@ -604,10 +588,19 @@ class MiniMaxH3TokenRefiner(nn.Module):
         self,
         arch: MiniMaxH3DiTArchConfig,
         quant_config: QuantizationConfig | None,
+        *,
+        prefix: str,
     ) -> None:
         super().__init__()
         self.blocks = nn.ModuleList(
-            [MiniMaxH3TokenRefinerBlock(arch, quant_config) for _ in range(arch.token_refiner_num_layers)]
+            [
+                MiniMaxH3TokenRefinerBlock(
+                    arch,
+                    quant_config,
+                    prefix=f"{prefix}.blocks.{i}",
+                )
+                for i in range(arch.token_refiner_num_layers)
+            ]
         )
         self.final_norm = _norm(arch.hidden_size, eps=arch.final_norm_eps)
 
@@ -628,18 +621,29 @@ class MiniMaxH3DiTBlock(nn.Module):
         self,
         arch: MiniMaxH3DiTArchConfig,
         quant_config: QuantizationConfig | None,
+        *,
+        prefix: str,
     ) -> None:
         super().__init__()
         self.norm1 = _norm(arch.hidden_size, eps=arch.norm_eps)
         self.norm2 = _norm(arch.hidden_size, eps=arch.norm_eps)
-        self.attn = MiniMaxH3Attention(arch, quant_config)
-        self.mlp = MiniMaxH3MLP(arch, quant_config)
+        self.attn = MiniMaxH3Attention(
+            arch,
+            quant_config,
+            prefix=f"{prefix}.attn",
+        )
+        self.mlp = MiniMaxH3MLP(
+            arch,
+            quant_config,
+            prefix=f"{prefix}.mlp",
+        )
         self.adaln_proj = MiniMaxH3AdalnProj(
             arch,
             arch.adaln_out_features,
             quant_config,
             expand_ratio=6,
             modality_num=MINIMAX_H3_ADALN_MODALITY_NUM,
+            prefix=f"{prefix}.adaln_proj",
         )
 
     def forward(
@@ -693,6 +697,8 @@ class MiniMaxH3FinalLayer(nn.Module):
         self,
         arch: MiniMaxH3DiTArchConfig,
         quant_config: QuantizationConfig | None,
+        *,
+        prefix: str,
     ) -> None:
         super().__init__()
         video_patch_dim = arch.latents_dim * arch.patch_size[0] * arch.patch_size[1] * arch.patch_size[2]
@@ -703,6 +709,7 @@ class MiniMaxH3FinalLayer(nn.Module):
             quant_config,
             expand_ratio=2,
             modality_num=1,
+            prefix=f"{prefix}.adaln_proj",
         )
         self.video_out = ColumnParallelLinear(
             arch.hidden_size,
@@ -710,7 +717,8 @@ class MiniMaxH3FinalLayer(nn.Module):
             bias=True,
             gather_output=True,
             params_dtype=_FP32_DTYPE,
-            quant_config=quant_config,
+            quant_config=None,
+            prefix=f"{prefix}.video_out",
         )
         self.audio_out = ColumnParallelLinear(
             arch.hidden_size,
@@ -718,7 +726,8 @@ class MiniMaxH3FinalLayer(nn.Module):
             bias=True,
             gather_output=True,
             params_dtype=_FP32_DTYPE,
-            quant_config=quant_config,
+            quant_config=None,
+            prefix=f"{prefix}.audio_out",
         )
 
     def forward(
@@ -864,7 +873,8 @@ class MiniMaxH3DiTModel(nn.Module):
             bias=True,
             gather_output=True,
             params_dtype=_FP32_DTYPE,
-            quant_config=quant_config,
+            quant_config=None,
+            prefix="video_patch_proj",
         )
         self.audio_patch_proj = ColumnParallelLinear(
             arch.audio_latents_dim,
@@ -872,7 +882,8 @@ class MiniMaxH3DiTModel(nn.Module):
             bias=True,
             gather_output=True,
             params_dtype=_FP32_DTYPE,
-            quant_config=quant_config,
+            quant_config=None,
+            prefix="audio_patch_proj",
         )
         self.condition_proj = ColumnParallelLinear(
             arch.text_dim,
@@ -881,14 +892,35 @@ class MiniMaxH3DiTModel(nn.Module):
             gather_output=True,
             params_dtype=_BF16_DTYPE,
             quant_config=quant_config,
+            prefix="condition_proj",
+        )
+        self.time_embedder = MiniMaxH3TimeEmbedder(
+            arch,
+            prefix="time_embedder",
         )
-        self.time_embedder = MiniMaxH3TimeEmbedder(arch, quant_config)
         self.rope = MiniMaxH3Rope(arch.rope_inv_freq_len)
-        self.token_refiner = MiniMaxH3TokenRefiner(arch, quant_config)
-        self.blocks = nn.ModuleList([MiniMaxH3DiTBlock(arch, quant_config) for _ in range(arch.num_layers)])
+        self.token_refiner = MiniMaxH3TokenRefiner(
+            arch,
+            quant_config,
+            prefix="token_refiner",
+        )
+        self.blocks = nn.ModuleList(
+            [
+                MiniMaxH3DiTBlock(
+                    arch,
+                    quant_config,
+                    prefix=f"blocks.{i}",
+                )
+                for i in range(arch.num_layers)
+            ]
+        )
         self.sp_prepare = MiniMaxH3SPPrepare()
         self.sp_gather = MiniMaxH3SPGather()
-        self.final_layer = MiniMaxH3FinalLayer(arch, quant_config)
+        self.final_layer = MiniMaxH3FinalLayer(
+            arch,
+            quant_config,
+            prefix="final_layer",
+        )
         self._mark_missing_params_required()
 
     def _mark_missing_params_required(self) -> None:
@@ -917,7 +949,27 @@ class MiniMaxH3DiTModel(nn.Module):
                 logger.warning("Skipping MiniMax H3 weight not present in model: %s", name)
                 continue
             weight_loader = getattr(param, "weight_loader", default_weight_loader)
-            weight_loader(param, loaded_weight)
+            if name.endswith(".attn.qkv_proj.weight"):
+                # Transform checkpoint layout before entering vLLM's loader so
+                # online FP8 can keep ``online_process_loader`` outermost.
+                loaded_weight = _reorder_grouped_qkv_to_qkv(
+                    loaded_weight,
+                    num_query_groups=self.arch.num_attention_heads,
+                    heads_per_group=1,
+                    head_dim=self.arch.attention_head_dim,
+                )
+                weight_loader(param, loaded_weight)
+            elif name.endswith(".mlp.fc1.weight"):
+                if loaded_weight.shape[0] % 2:
+                    raise ValueError(
+                        "MiniMax H3 fc1 checkpoint rows must split evenly into "
+                        f"gate/up matrices, got {tuple(loaded_weight.shape)}"
+                    )
+                gate, up = loaded_weight.chunk(2, dim=0)
+                weight_loader(param, gate, 0)
+                weight_loader(param, up, 1)
+            else:
+                weight_loader(param, loaded_weight)
             loaded.add(name)
         return loaded
 

vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py

@@ -88,6 +88,12 @@ MINIMAX_H3_REFERENCE_IMAGE_SHORT_EDGE = 2048
 MINIMAX_H3_REFERENCE_IMAGE_MULTIPLE = 32
 
 
+def _resolve_component_quant_config(quant_config, component: str):
+    if hasattr(quant_config, "resolve"):
+        return quant_config.resolve(component)
+    return quant_config
+
+
 def _minimax_h3_post_process(output, output_type: str = "np"):
     """Convert the joint video/audio output without capturing worker state.
 
@@ -293,9 +299,13 @@ class MiniMaxH3Pipeline(
                 fall_back_to_pt=False,
             )
         ]
+        transformer_quant_config = _resolve_component_quant_config(
+            od_config.quantization_config,
+            "transformer",
+        )
         self.transformer = MiniMaxH3DiTModel(
             od_config,
-            quant_config=od_config.quantization_config,
+            quant_config=transformer_quant_config,
         )
 
         self.tokenizer = Qwen2TokenizerFast.from_pretrained(