[Test] Cover MiniMax H3 AdaLN offload ownership

Signed-off-by: Sy03 <1370724210@qq.com>

.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:

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/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/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

@@ -115,6 +115,36 @@ steady-state latency. H3 is CFG-distilled, so `--cfg-parallel-size` must remain
 1. The H3 VAE supports its native `tile` mode, not
 `spatial_shard_height` or `spatial_shard_width`.
 
+### Optional exact-schedule AdaLN cache
+
+H3 can precompute every AdaLN projection for the request's known timestep
+schedule and reuse the exact rows in each denoise step:
+
+```bash
+export VLLM_OMNI_H3_ADALN_SCHEDULE_CACHE=1
+```
+
+The cache key uses the exact float32 timestep bit patterns. A request with the
+same schedule reuses the installed tables; a changed schedule rebuilds them.
+This is disabled by default because the tables consume persistent device
+memory proportional to the number of denoise steps and layers.
+
+After the tables are installed, the original projection weights can also move
+to CPU and release their device storage:
+
+```bash
+export VLLM_OMNI_H3_ADALN_SCHEDULE_CACHE=1
+export VLLM_OMNI_H3_ADALN_OFFLOAD_WEIGHTS=1
+```
+
+Weight offload requires the schedule cache. On a schedule change, the weights
+are restored, the replacement tables are built, and the weights are offloaded
+again. The optimization falls back to direct projection when the tables would
+occupy at least as much memory as the projection weights. DTensor/HSDP
+parameters remain under their distributed owner. With layerwise offload, H3's
+AdaLN projection parameters stay outside the per-block transfer set while the
+exact-schedule cache owns them.
+
 ### Text encoder tensor parallelism
 
 The Qwen3-VL text encoder (~51.5 GB in BF16 for the retained 50 layers) is by
@@ -164,6 +194,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 +376,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_adaln_cache.py

@@ -0,0 +1,409 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def _make_branch(layout: str):
+    from vllm_omni.diffusion.models.minimax_h3.denoise_loop import (
+        MiniMaxH3DenoiseBranch,
+    )
+
+    update_mask = {
+        "t2va": [True, True],
+        "fl2va": [False, True],
+        "ref2va": [True, True],
+    }[layout]
+    audio_update_mask = {
+        "t2va": [True, True],
+        "fl2va": [True, True],
+        "ref2va": [False, True],
+    }[layout]
+    packed = {
+        "seq_len": 8,
+        "text_pos": torch.tensor([0, 1]),
+        "img_pos": torch.tensor([2, 3]),
+        "audio_pos": torch.tensor([4, 5]),
+        "update_mask": torch.tensor(update_mask),
+        "audio_update_mask": torch.tensor(audio_update_mask),
+        "img_position_ids": torch.zeros(8, 3, dtype=torch.long),
+        "cu_seqlens": torch.tensor([0, 6, 8], dtype=torch.int32),
+    }
+    return MiniMaxH3DenoiseBranch(
+        packed=packed,
+        text_embeddings=torch.zeros(2, 4),
+        token_tags=torch.zeros(8, dtype=torch.long),
+        device=torch.device("cpu"),
+    )
+
+
+def _reference_full_vector_unique(
+    branch,
+    *,
+    sigma_video: float,
+    sigma_audio: float,
+    imgvid_cond_timestep: float,
+    audio_ref_cond_timestep: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    t_video = 1.0 - sigma_video
+    t_audio = 1.0 - sigma_audio
+    timesteps = torch.full((branch.seq_len,), t_video, dtype=torch.float32)
+    timesteps[branch.img_pos[branch.update_mask]] = t_video
+    timesteps[branch.img_pos[~branch.update_mask]] = max(t_video, imgvid_cond_timestep)
+    timesteps[branch.audio_pos[branch.audio_update_mask]] = t_audio
+    timesteps[branch.audio_pos[~branch.audio_update_mask]] = max(t_audio, audio_ref_cond_timestep)
+    return torch.unique(timesteps, sorted=True, return_inverse=True)
+
+
+@pytest.mark.parametrize("layout", ["t2va", "fl2va", "ref2va"])
+def test_timestep_metadata_matches_full_vector_unique(layout: str):
+    branch = _make_branch(layout)
+    sigmas_video = [1.0, 0.5, 0.0]
+    sigmas_audio = [1.0, 0.25, 0.0]
+
+    cached = branch.prepare_timestep_metadata(
+        sigmas_video=sigmas_video,
+        sigmas_audio=sigmas_audio,
+        imgvid_cond_noise_aug_for_inference=0.999,
+        audio_cond_noise_aug_for_inference=1.0,
+    )
+
+    assert len(cached) == 2
+    for step, metadata in enumerate(cached):
+        expected_unique, expected_inverse = _reference_full_vector_unique(
+            branch,
+            sigma_video=sigmas_video[step],
+            sigma_audio=sigmas_audio[step],
+            imgvid_cond_timestep=0.999,
+            audio_ref_cond_timestep=1.0,
+        )
+        torch.testing.assert_close(metadata.unique_timesteps.cpu(), expected_unique, rtol=0, atol=0)
+        torch.testing.assert_close(metadata.inverse_indices.cpu(), expected_inverse, rtol=0, atol=0)
+
+
+def test_adaln_schedule_key_uses_exact_float32_bits():
+    from vllm_omni.diffusion.models.minimax_h3.adaln_schedule_cache import (
+        minimax_h3_adaln_schedule_key,
+    )
+
+    branch = _make_branch("t2va")
+    first = branch.prepare_timestep_metadata(
+        sigmas_video=[1.0, 0.5, 0.0],
+        sigmas_audio=[1.0, 0.25, 0.0],
+        imgvid_cond_noise_aug_for_inference=0.999,
+        audio_cond_noise_aug_for_inference=1.0,
+    )
+    same_after_float32_rounding = branch.prepare_timestep_metadata(
+        sigmas_video=[1.0, 0.5 + 1e-12, 0.0],
+        sigmas_audio=[1.0, 0.25 + 1e-12, 0.0],
+        imgvid_cond_noise_aug_for_inference=0.999,
+        audio_cond_noise_aug_for_inference=1.0,
+    )
+    changed = branch.prepare_timestep_metadata(
+        sigmas_video=[1.0, 0.5, 0.0],
+        sigmas_audio=[1.0, 0.25000003, 0.0],
+        imgvid_cond_noise_aug_for_inference=0.999,
+        audio_cond_noise_aug_for_inference=1.0,
+    )
+
+    first_key = minimax_h3_adaln_schedule_key(first)
+    assert first_key == minimax_h3_adaln_schedule_key(same_after_float32_rounding)
+    assert first_key != minimax_h3_adaln_schedule_key(changed)
+
+
+class _FakeAdalnLinear(torch.nn.Module):
+    def __init__(self, out_features: int):
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.arange(out_features * 4, dtype=torch.bfloat16).reshape(out_features, 4))
+        self.bias = torch.nn.Parameter(torch.arange(out_features, dtype=torch.bfloat16))
+        self.out_features = out_features
+        self.calls = 0
+
+    def forward(self, x):
+        assert self.weight.numel() > 0
+        assert self.bias.numel() > 0
+        self.calls += 1
+        rows = torch.arange(
+            x.shape[0] * self.out_features,
+            device=x.device,
+            dtype=torch.float32,
+        ).reshape(x.shape[0], self.out_features)
+        rows = rows + x.float().sum(dim=-1, keepdim=True)
+        return rows.to(torch.bfloat16), None
+
+
+def _make_tiny_adaln_projection(*, expand_ratio: int = 2, modality_num: int = 3):
+    from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
+        MiniMaxH3AdalnProj,
+    )
+
+    projection = object.__new__(MiniMaxH3AdalnProj)
+    torch.nn.Module.__init__(projection)
+    projection.expand_ratio = expand_ratio
+    projection.modality_num = modality_num
+    projection.hidden_size = 2
+    projection.linear = _FakeAdalnLinear(out_features=expand_ratio * modality_num * projection.hidden_size)
+    projection.register_buffer("_precomputed_table", None, persistent=False)
+    projection._precomputed_step = None
+    projection._offloaded_projection_state = None
+    return projection
+
+
+def test_adaln_projection_reuses_exact_precomputed_rows_without_linear_calls():
+    projection = _make_tiny_adaln_projection()
+    first_t_emb = torch.tensor([[0.0, 1.0, 2.0, 3.0]])
+    second_t_emb = torch.tensor([[4.0, 5.0, 6.0, 7.0], [8.0, 9.0, 10.0, 11.0]])
+    first_direct = projection.compute_flat(first_t_emb)
+    second_direct = projection.compute_flat(second_t_emb)
+    first_padded = torch.cat(
+        [first_direct, first_direct.new_zeros(second_direct.shape[0] - first_direct.shape[0], first_direct.shape[1])]
+    )
+    table = torch.stack([first_padded, second_direct])
+    cursor = torch.zeros((), dtype=torch.long)
+    direct_calls = projection.linear.calls
+
+    projection.install_precomputed(table, cursor)
+    first_cached = projection(first_t_emb)
+    cursor.fill_(1)
+    second_cached = projection(second_t_emb)
+
+    assert projection.linear.calls == direct_calls
+    for actual, expected in zip(first_cached, first_direct.chunk(2, dim=-1)):
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+    for actual, expected in zip(second_cached, second_direct.chunk(2, dim=-1)):
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+def test_adaln_projection_offload_preserves_cache_and_clear_restores_weights():
+    projection = _make_tiny_adaln_projection()
+    t_emb = torch.tensor([[0.0, 1.0, 2.0, 3.0]])
+    direct = projection.compute_flat(t_emb)
+    expected_state = {name: tensor.clone() for name, tensor in projection.linear.state_dict().items()}
+    expected_bytes = sum(tensor.numel() * tensor.element_size() for tensor in expected_state.values())
+    projection.install_precomputed(direct.unsqueeze(0), torch.zeros((), dtype=torch.long))
+
+    freed_bytes = projection.offload_projection()
+
+    assert freed_bytes == expected_bytes
+    assert all(parameter.numel() == 0 for parameter in projection.linear.parameters())
+    cached = projection(t_emb)
+    for actual, expected in zip(cached, direct.chunk(projection.expand_ratio, dim=-1)):
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+    projection.clear_precomputed()
+    assert not projection.projection_weights_offloaded
+    for name, tensor in projection.linear.state_dict().items():
+        torch.testing.assert_close(tensor, expected_state[name], rtol=0, atol=0)
+    torch.testing.assert_close(projection.compute_flat(t_emb), direct, rtol=0, atol=0)
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for allocator accounting")
+def test_adaln_projection_offload_releases_cuda_allocator_storage():
+    device = torch.device("cuda")
+    torch.accelerator.empty_cache()
+    projection = _make_tiny_adaln_projection().to(device)
+    t_emb = torch.tensor([[0.0, 1.0, 2.0, 3.0]], device=device)
+    direct = projection.compute_flat(t_emb)
+    projection.install_precomputed(
+        direct.unsqueeze(0),
+        torch.zeros((), dtype=torch.long, device=device),
+    )
+    torch.accelerator.synchronize()
+    allocated_before = torch.cuda.memory_allocated(device)
+
+    freed_bytes = projection.offload_projection()
+    torch.accelerator.synchronize()
+    allocated_after = torch.cuda.memory_allocated(device)
+
+    assert freed_bytes > 0
+    assert allocated_before - allocated_after >= freed_bytes
+    assert all(parameter.is_cuda and parameter.numel() == 0 for parameter in projection.linear.parameters())
+
+    projection.clear_precomputed()
+    assert all(parameter.is_cuda and parameter.numel() > 0 for parameter in projection.linear.parameters())
+
+
+class _CountingTimeEmbedder(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.calls = 0
+
+    def forward(self, timesteps):
+        self.calls += 1
+        return timesteps[:, None]
+
+
+class _TinyAdalnBlock(torch.nn.Module):
+    def __init__(self, *, modality_num: int = 3):
+        super().__init__()
+        self.adaln_proj = _make_tiny_adaln_projection(modality_num=modality_num)
+
+
+def _make_adaln_schedule_model():
+    from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
+        MiniMaxH3DiTModel,
+    )
+
+    model = object.__new__(MiniMaxH3DiTModel)
+    torch.nn.Module.__init__(model)
+    model.time_embedder = _CountingTimeEmbedder()
+    model.blocks = torch.nn.ModuleList([_TinyAdalnBlock(), _TinyAdalnBlock()])
+    model.final_layer = _TinyAdalnBlock(modality_num=1)
+    model._adaln_projections = tuple(block.adaln_proj for block in model.blocks) + (model.final_layer.adaln_proj,)
+    model._adaln_schedule_cache = None
+    return model
+
+
+def _schedule(branch, *, middle_video: float, middle_audio: float):
+    return branch.prepare_timestep_metadata(
+        sigmas_video=[1.0, middle_video, 0.0],
+        sigmas_audio=[1.0, middle_audio, 0.0],
+        imgvid_cond_noise_aug_for_inference=0.999,
+        audio_cond_noise_aug_for_inference=1.0,
+    )
+
+
+def test_model_cache_hit_and_changed_schedule_rebuild_with_weight_offload():
+    from vllm_omni.diffusion.models.minimax_h3.adaln_schedule_cache import (
+        minimax_h3_adaln_schedule_key,
+    )
+
+    branch = _make_branch("t2va")
+    first_plan = _schedule(branch, middle_video=0.5, middle_audio=0.25)
+    changed_plan = _schedule(branch, middle_video=0.4, middle_audio=0.2)
+    model = _make_adaln_schedule_model()
+    expected_projection_bytes = sum(
+        parameter.numel() * parameter.element_size()
+        for projection in model._adaln_projections
+        for parameter in projection.linear.parameters()
+    )
+
+    first = model.prepare_adaln_schedule_cache(
+        unique_timestep_plan=tuple(step.unique_timesteps for step in first_plan),
+        schedule_key=minimax_h3_adaln_schedule_key(first_plan),
+        offload_weights=True,
+    )
+    calls_after_first = [projection.linear.calls for projection in model._adaln_projections]
+    hit = model.prepare_adaln_schedule_cache(
+        unique_timestep_plan=tuple(step.unique_timesteps for step in first_plan),
+        schedule_key=minimax_h3_adaln_schedule_key(first_plan),
+        offload_weights=True,
+    )
+    replaced = model.prepare_adaln_schedule_cache(
+        unique_timestep_plan=tuple(step.unique_timesteps for step in changed_plan),
+        schedule_key=minimax_h3_adaln_schedule_key(changed_plan),
+        offload_weights=True,
+    )
+
+    assert first is hit
+    assert replaced is not None and replaced is not first
+    assert first is not None
+    assert first.offloaded_projection_bytes == expected_projection_bytes
+    assert first.net_memory_saved_bytes == expected_projection_bytes - first.table_bytes
+    assert calls_after_first == [2, 2, 2]
+    assert [projection.linear.calls for projection in model._adaln_projections] == [4, 4, 4]
+    assert all(projection.projection_weights_offloaded for projection in model._adaln_projections)
+
+
+def test_model_does_not_manually_offload_dtensor_owned_parameters(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer
+    from vllm_omni.diffusion.models.minimax_h3.adaln_schedule_cache import (
+        minimax_h3_adaln_schedule_key,
+    )
+
+    class FakeDTensor(torch.nn.Parameter):
+        def to_local(self):
+            return self
+
+    monkeypatch.setattr(minimax_h3_transformer, "DTensor", FakeDTensor)
+    branch = _make_branch("t2va")
+    plan = _schedule(branch, middle_video=0.5, middle_audio=0.25)
+    model = _make_adaln_schedule_model()
+    projection = model._adaln_projections[0]
+    projection.linear.weight = FakeDTensor(projection.linear.weight.detach())
+
+    cache = model.prepare_adaln_schedule_cache(
+        unique_timestep_plan=tuple(step.unique_timesteps for step in plan),
+        schedule_key=minimax_h3_adaln_schedule_key(plan),
+        offload_weights=True,
+    )
+
+    assert cache is not None
+    assert cache.offloaded_projection_bytes == 0
+    assert all(not item.projection_weights_offloaded for item in model._adaln_projections)
+    assert projection.linear.weight.numel() > 0
+
+
+def test_model_weight_offload_falls_back_when_tables_are_larger():
+    model = _make_adaln_schedule_model()
+    unique_timestep_plan = tuple(torch.tensor([float(step)]) for step in range(12))
+
+    cache = model.prepare_adaln_schedule_cache(
+        unique_timestep_plan=unique_timestep_plan,
+        schedule_key=tuple((step,) for step in range(12)),
+        offload_weights=True,
+    )
+
+    assert cache is None
+    assert all(projection._precomputed_table is None for projection in model._adaln_projections)
+    assert all(not projection.projection_weights_offloaded for projection in model._adaln_projections)
+
+
+def test_denoise_loop_prepares_cache_once_and_passes_step_indices(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3 import denoise_loop
+
+    monkeypatch.setenv("VLLM_OMNI_H3_ADALN_SCHEDULE_CACHE", "1")
+    monkeypatch.setenv("VLLM_OMNI_H3_ADALN_OFFLOAD_WEIGHTS", "1")
+
+    class FakeModel:
+        def __init__(self):
+            self.prepare_calls = 0
+            self.prepare_for_request_calls = 0
+            self.step_indices = []
+
+        def prepare_for_request_caches(self):
+            assert torch.is_inference_mode_enabled()
+            self.prepare_for_request_calls += 1
+
+        def prepare_adaln_schedule_cache(self, *, unique_timestep_plan, schedule_key, offload_weights):
+            assert torch.is_inference_mode_enabled()
+            assert len(unique_timestep_plan) == len(schedule_key) == 2
+            assert offload_weights
+            self.prepare_calls += 1
+            return object()
+
+        def __call__(self, **kwargs):
+            self.step_indices.append(kwargs["adaln_step_index"])
+            return torch.zeros(2, 96), torch.zeros(2, 32)
+
+    model = FakeModel()
+    denoise_loop.minimax_h3_denoise_loop(
+        model=model,
+        positive=_make_branch("t2va"),
+        initial_video_rows=torch.zeros(2, 96),
+        initial_audio_rows=torch.zeros(2, 32),
+        keyframe_cond_rows=None,
+        sigmas_video=[1.0, 0.5, 0.0],
+        sigmas_audio=[1.0, 0.25, 0.0],
+        device=torch.device("cpu"),
+    )
+
+    assert model.prepare_for_request_calls == 1
+    assert model.prepare_calls == 1
+    assert model.step_indices == [0, 1]
+
+
+def test_weight_offload_env_requires_schedule_cache(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3.adaln_schedule_cache import (
+        minimax_h3_adaln_weight_offload_enabled,
+    )
+
+    monkeypatch.delenv("VLLM_OMNI_H3_ADALN_SCHEDULE_CACHE", raising=False)
+    monkeypatch.setenv("VLLM_OMNI_H3_ADALN_OFFLOAD_WEIGHTS", "1")
+    assert not minimax_h3_adaln_weight_offload_enabled()
+    monkeypatch.setenv("VLLM_OMNI_H3_ADALN_SCHEDULE_CACHE", "1")
+    assert minimax_h3_adaln_weight_offload_enabled()

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_parallel.py

@@ -25,7 +25,7 @@ def test_grouped_qkv_checkpoint_reorder():
     assert reordered[:, 0].tolist() == [0, 3, 1, 4, 2, 5]
 
 
-def test_transformer_declares_cache_sp_layerwise_offload_and_hsdp():
+def test_transformer_declares_cache_sp_layerwise_offload_and_hsdp(monkeypatch):
     from cache_dit import ForwardPattern
 
     from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
@@ -38,6 +38,11 @@ def test_transformer_declares_cache_sp_layerwise_offload_and_hsdp():
     assert not MiniMaxH3DiTModel._cache_dit_adapter_config.has_separate_cfg
     assert set(MiniMaxH3DiTModel._sp_plan) == {"sp_prepare", "sp_gather"}
 
+    monkeypatch.setenv("VLLM_OMNI_H3_ADALN_SCHEDULE_CACHE", "1")
+    assert MiniMaxH3DiTModel.layerwise_offload_excluded_parameter_prefixes() == ("adaln_proj.linear.",)
+    monkeypatch.delenv("VLLM_OMNI_H3_ADALN_SCHEDULE_CACHE")
+    assert MiniMaxH3DiTModel.layerwise_offload_excluded_parameter_prefixes() == ()
+
     model = object.__new__(MiniMaxH3DiTModel)
     nn.Module.__init__(model)
     model.blocks = nn.ModuleList([nn.Linear(4, 4) for _ in range(2)])
@@ -61,6 +66,14 @@ def test_packed_attention_is_a_regional_compile_boundary():
     assert getattr(MiniMaxH3Attention._run_packed_attention, "_torchdynamo_disable", False)
 
 
+def test_adaln_projection_is_a_regional_compile_boundary():
+    from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
+        MiniMaxH3AdalnProj,
+    )
+
+    assert getattr(MiniMaxH3AdalnProj.forward, "_torchdynamo_disable", False)
+
+
 @pytest.mark.parametrize(
     ("tp_size", "message"),
     [

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/diffusion/offloader/test_layerwise_backend.py

@@ -135,6 +135,14 @@ class _DummyBlock(nn.Module):
         self.weight = nn.Parameter(torch.randn(10, 10))
 
 
+class _OwnedProjectionBlock(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.weight = nn.Parameter(torch.arange(4, dtype=torch.float32))
+        self.adaln_proj = nn.Module()
+        self.adaln_proj.linear = nn.Linear(2, 3)
+
+
 class _SingleBlockModel(nn.Module):
     _layerwise_offload_blocks_attrs = ["blocks"]
 
@@ -224,6 +232,49 @@ class TestGetBlocksFromDit:
         assert len(blocks) == 2
 
 
+def test_layerwise_hook_leaves_excluded_parameters_under_model_ownership(patched_offload_runtime):
+    current_block = _OwnedProjectionBlock()
+    next_block = _OwnedProjectionBlock()
+    current_adaln = {name: tensor.clone() for name, tensor in current_block.adaln_proj.linear.state_dict().items()}
+    next_adaln = {name: tensor.clone() for name, tensor in next_block.adaln_proj.linear.state_dict().items()}
+    hook = LayerwiseOffloadHook(
+        next_block=next_block,
+        device=torch.device("cpu"),
+        stream=DummyStream(),
+        pin_memory=False,
+        excluded_parameter_prefixes=("adaln_proj.linear.",),
+    )
+
+    hook.initialize_hook(current_block)
+
+    assert set(hook.block_parameters) == {"weight"}
+    assert set(hook.next_block_parameters) == {"weight"}
+    assert next_block.weight.numel() == 0
+    for name, tensor in next_block.adaln_proj.linear.state_dict().items():
+        torch.testing.assert_close(tensor, next_adaln[name], rtol=0, atol=0)
+
+    hook.offload_layer()
+    assert current_block.weight.numel() == 0
+    for name, tensor in current_block.adaln_proj.linear.state_dict().items():
+        torch.testing.assert_close(tensor, current_adaln[name], rtol=0, atol=0)
+
+
+def test_layerwise_backend_materializes_excluded_parameters_on_target_device():
+    block = _OwnedProjectionBlock()
+    expected_bytes = sum(
+        parameter.numel() * parameter.element_size() for parameter in block.adaln_proj.linear.parameters()
+    )
+
+    moved_bytes = LayerWiseOffloadBackend._move_excluded_parameters_to_device(
+        block,
+        torch.device("cpu"),
+        ("adaln_proj.linear.",),
+    )
+
+    assert moved_bytes == expected_bytes
+    assert all(parameter.device.type == "cpu" for parameter in block.adaln_proj.linear.parameters())
+
+
 class TestGetBlocksAttrNames:
     def test_get_blocks_attr_names_new_format(self):
         model = _MultiBlockModel()

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/openai_api/test_video_server.py

@@ -548,6 +548,23 @@ def test_multi_video_generation_preserves_uploaded_files_until_generation(
     assert engine.captured_sampling_params_list[0].extra_args["duration"] == 15.0
 
 
+def test_decode_video_bytes_can_keep_first_frames():
+    from vllm_omni.entrypoints.openai.video_api_utils import _decode_video_bytes
+
+    frames = _decode_video_bytes(
+        _make_test_video_bytes((32, 24), num_frames=6),
+        source="input_reference",
+        max_frames=2,
+        keep="first",
+    )
+
+    assert len(frames) == 2
+    assert frames.fps == pytest.approx(8.0)
+    red_means = [np.asarray(frame)[:, :, 0].mean() for frame in frames]
+    assert red_means[0] < red_means[1]
+    assert red_means[1] < 100
+
+
 def test_decode_video_bytes_can_keep_last_frames():
     from vllm_omni.entrypoints.openai.video_api_utils import _decode_video_bytes
 

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

tests/worker/test_omni_connector_mixin.py

@@ -105,6 +105,34 @@ class _FakeTPGroup:
 # ------------------------------------------------------------------ #
 
 
+@pytest.mark.parametrize(
+    ("role", "custom_func", "expected"),
+    [
+        ("sender", None, False),
+        ("sender", "package.build_payload", True),
+        ("receiver", None, True),
+    ],
+)
+def test_init_payload_connector_ownership(role, custom_func, expected):
+    model_config = _make_model_config(custom_func=custom_func)
+    model_config.stage_connector_config = {
+        "name": "MooncakeTransferEngineConnector",
+        "extra": {"role": role},
+    }
+
+    host = MixinHost()
+    connector = MockConnector()
+    with (
+        patch.object(host, "_create_connector", return_value=connector) as create,
+        patch.object(host, "_load_custom_func", return_value=(None, None)),
+    ):
+        host.init_omni_connectors(model_config)
+
+    assert (create.call_count == 1) is expected
+    assert (host._omni_connector is connector) is expected
+    host.shutdown_omni_connectors()
+
+
 class TestMixinAsyncChunkSendRecv(unittest.TestCase):
     """Test 2: Async chunk send/recv + bg threads."""
 

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/bagel/bagel_transformer.py

@@ -2426,7 +2426,7 @@ class Bagel(CFGParallelMixin, nn.Module):
 
             batched_sequence = packed_sequence.repeat(num_branches, 1)
             batched_vae_indexes = torch.cat([packed_vae_token_indexes + i * seq_len for i in range(num_branches)])
-            batched_position_ids = torch.cat(cfg_branch_pids)
+            batched_position_ids = torch.cat(cfg_branch_pids, dim=1 if cfg_branch_pids[0].ndim == 2 else 0)
             batched_seqlens = packed_seqlens.repeat(num_branches)
             merged_cache = NaiveCache.merge(cfg_branch_caches)
 

vllm_omni/diffusion/models/minimax_h3/adaln_schedule_cache.py

@@ -0,0 +1,98 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Exact schedule identity and policy for MiniMax-H3 AdaLN reuse."""
+
+from __future__ import annotations
+
+import os
+from collections.abc import Iterable
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+
+MINIMAX_H3_ADALN_SCHEDULE_CACHE_ENV = "VLLM_OMNI_H3_ADALN_SCHEDULE_CACHE"
+MINIMAX_H3_ADALN_OFFLOAD_WEIGHTS_ENV = "VLLM_OMNI_H3_ADALN_OFFLOAD_WEIGHTS"
+_TRUTHY_ENV_VALUES = frozenset({"1", "true", "yes", "on"})
+
+MiniMaxH3AdalnScheduleKey = tuple[tuple[int, ...], ...]
+
+
+@dataclass(frozen=True)
+class MiniMaxH3AdalnScheduleCache:
+    """One model-local exact modulation schedule."""
+
+    key: MiniMaxH3AdalnScheduleKey
+    t_embs: tuple[torch.Tensor, ...]
+    step_cursor: torch.Tensor
+    steps: int
+    max_unique_timesteps: int
+    table_bytes: int
+    projection_bytes: int = 0
+    offloaded_projection_bytes: int = 0
+    net_memory_saved_bytes: int = 0
+
+
+def minimax_h3_adaln_schedule_cache_enabled() -> bool:
+    """Return whether persistent exact-schedule AdaLN reuse is enabled."""
+    return os.getenv(MINIMAX_H3_ADALN_SCHEDULE_CACHE_ENV, "").strip().lower() in _TRUTHY_ENV_VALUES
+
+
+def minimax_h3_adaln_weight_offload_enabled() -> bool:
+    """Return whether exact tables should take ownership of AdaLN weights."""
+    return minimax_h3_adaln_schedule_cache_enabled() and (
+        os.getenv(MINIMAX_H3_ADALN_OFFLOAD_WEIGHTS_ENV, "").strip().lower() in _TRUTHY_ENV_VALUES
+    )
+
+
+def minimax_h3_float32_bits(values: Any) -> tuple[int, ...]:
+    """Return the signed int32 bit patterns of a CPU float32 tensor."""
+    if not isinstance(values, torch.Tensor) or values.device.type != "cpu" or values.dtype != torch.float32:
+        raise ValueError("schedule key values must be a CPU float32 tensor")
+    return tuple(int(value) for value in values.contiguous().view(torch.int32).tolist())
+
+
+def minimax_h3_adaln_schedule_key(
+    metadata: Iterable[Any],
+) -> MiniMaxH3AdalnScheduleKey:
+    """Build an immutable schedule key without reading a device tensor."""
+    key: list[tuple[int, ...]] = []
+    for step in metadata:
+        bits = getattr(step, "timestep_bits", None)
+        if bits is None:
+            raise ValueError("timestep metadata is missing host timestep_bits")
+        key.append(tuple(int(value) for value in bits))
+    if not key:
+        raise ValueError("AdaLN schedule must contain at least one denoise step")
+    return tuple(key)
+
+
+def minimax_h3_build_adaln_table(
+    projection: Any,
+    t_embs: tuple[torch.Tensor, ...],
+) -> torch.Tensor:
+    """Project each step at its original shape and pad only unused rows."""
+    rows_by_step = tuple(projection.compute_flat(t_emb) for t_emb in t_embs)
+    max_rows = max(int(rows.shape[0]) for rows in rows_by_step)
+    padded: list[torch.Tensor] = []
+    for rows in rows_by_step:
+        missing = max_rows - int(rows.shape[0])
+        if missing:
+            rows = torch.cat(
+                [rows, rows.new_zeros((missing, rows.shape[1]))],
+                dim=0,
+            )
+        padded.append(rows)
+    return torch.stack(padded, dim=0)
+
+
+__all__ = [
+    "MINIMAX_H3_ADALN_OFFLOAD_WEIGHTS_ENV",
+    "MINIMAX_H3_ADALN_SCHEDULE_CACHE_ENV",
+    "MiniMaxH3AdalnScheduleCache",
+    "MiniMaxH3AdalnScheduleKey",
+    "minimax_h3_build_adaln_table",
+    "minimax_h3_adaln_schedule_cache_enabled",
+    "minimax_h3_adaln_weight_offload_enabled",
+    "minimax_h3_adaln_schedule_key",
+    "minimax_h3_float32_bits",
+]

vllm_omni/diffusion/models/minimax_h3/denoise_loop.py

@@ -10,10 +10,17 @@ from __future__ import annotations
 
 from collections.abc import Callable
 from contextlib import AbstractContextManager, nullcontext
+from dataclasses import dataclass
 from typing import Any
 
 import torch
 
+from .adaln_schedule_cache import (
+    minimax_h3_adaln_schedule_cache_enabled,
+    minimax_h3_adaln_schedule_key,
+    minimax_h3_adaln_weight_offload_enabled,
+    minimax_h3_float32_bits,
+)
 from .scheduling_minimax_h3_euler_ancestral import (
     minimax_h3_euler_eta0_step,
     minimax_h3_rf_v_to_x0,
@@ -28,6 +35,15 @@ MINIMAX_H3_VIDEO_ROW_WIDTH = 96
 MINIMAX_H3_AUDIO_ROW_WIDTH = 32
 
 
+@dataclass(frozen=True)
+class MiniMaxH3StepTimestepMetadata:
+    """Device-ready timestep routing for one denoise forward."""
+
+    unique_timesteps: torch.Tensor
+    inverse_indices: torch.Tensor
+    timestep_bits: tuple[int, ...]
+
+
 class MiniMaxH3DenoiseBranch:
     """Static per-branch state: packed layout + fixed forward kwargs.
 
@@ -45,6 +61,7 @@ class MiniMaxH3DenoiseBranch:
         device: torch.device,
     ) -> None:
         seq_len = int(packed["seq_len"])
+        self.device = device
         self.seq_len = seq_len
         self.img_pos = packed["img_pos"].view(-1).to(torch.long)
         self.audio_pos = packed["audio_pos"].view(-1).to(torch.long)
@@ -65,6 +82,12 @@ class MiniMaxH3DenoiseBranch:
         self.audio_pos_dev = self.audio_pos.to(device)
         self.update_mask_dev = self.update_mask.to(device)
         self.audio_update_mask_dev = self.audio_update_mask.to(device)
+        timestep_category_ids = torch.zeros(seq_len, dtype=torch.long)
+        timestep_category_ids[self.img_pos[~self.update_mask]] = 1
+        timestep_category_ids[self.audio_pos[self.audio_update_mask]] = 2
+        timestep_category_ids[self.audio_pos[~self.audio_update_mask]] = 3
+        self.timestep_category_ids = timestep_category_ids
+        self.used_timestep_categories = torch.unique(timestep_category_ids, sorted=True)
         self.x_base = torch.zeros(1, seq_len, MINIMAX_H3_VIDEO_ROW_WIDTH, dtype=torch.float32, device=device)
         self.audio_x_base = torch.zeros(1, seq_len, MINIMAX_H3_AUDIO_ROW_WIDTH, dtype=torch.float32, device=device)
         text_pos_dev = packed["text_pos"].view(-1).to(torch.long).to(device)
@@ -87,43 +110,103 @@ class MiniMaxH3DenoiseBranch:
                 "max_seqlen_q": text_len,
             },
         }
+        self._adaln_schedule_cache_active = False
+
+    def prepare_timestep_metadata(
+        self,
+        *,
+        sigmas_video: list[float],
+        sigmas_audio: list[float],
+        imgvid_cond_noise_aug_for_inference: float,
+        audio_cond_noise_aug_for_inference: float,
+    ) -> list[MiniMaxH3StepTimestepMetadata]:
+        """Precompute the full schedule's packed-row timestep routing on CPU."""
+        if len(sigmas_video) != len(sigmas_audio):
+            raise ValueError("video/audio sigma schedules must have equal length")
+        if len(sigmas_video) < 2:
+            raise ValueError("sigma schedules need at least 2 entries")
+
+        result: list[MiniMaxH3StepTimestepMetadata] = []
+        for step in range(len(sigmas_video) - 1):
+            t_video = 1.0 - float(sigmas_video[step])
+            t_audio = 1.0 - float(sigmas_audio[step])
+            category_timesteps = torch.tensor(
+                [
+                    t_video,
+                    max(t_video, float(imgvid_cond_noise_aug_for_inference)),
+                    t_audio,
+                    max(t_audio, float(audio_cond_noise_aug_for_inference)),
+                ],
+                dtype=torch.float32,
+            )
+            active_timesteps = category_timesteps.index_select(0, self.used_timestep_categories)
+            unique_timesteps, active_inverse = torch.unique(
+                active_timesteps,
+                sorted=True,
+                return_inverse=True,
+            )
+            category_inverse = torch.zeros(4, dtype=torch.long)
+            category_inverse.index_copy_(0, self.used_timestep_categories, active_inverse)
+            inverse_indices = category_inverse.index_select(0, self.timestep_category_ids)
+            result.append(
+                MiniMaxH3StepTimestepMetadata(
+                    unique_timesteps=unique_timesteps.to(self.device),
+                    inverse_indices=inverse_indices.to(self.device),
+                    timestep_bits=minimax_h3_float32_bits(unique_timesteps),
+                )
+            )
+        return result
+
+    def prepare_adaln_schedule_cache(
+        self,
+        model: Any,
+        timestep_metadata: list[MiniMaxH3StepTimestepMetadata],
+        *,
+        enabled: bool,
+    ) -> Any | None:
+        """Prepare exact model-local modulation tables for this schedule."""
+        self._adaln_schedule_cache_active = False
+        if not enabled:
+            clear = getattr(model, "clear_adaln_schedule_cache", None)
+            if clear is not None:
+                clear()
+            return None
+        prepare = getattr(model, "prepare_adaln_schedule_cache", None)
+        if prepare is None:
+            return None
+        with torch.inference_mode():
+            cache = prepare(
+                unique_timestep_plan=tuple(step.unique_timesteps for step in timestep_metadata),
+                schedule_key=minimax_h3_adaln_schedule_key(timestep_metadata),
+                offload_weights=minimax_h3_adaln_weight_offload_enabled(),
+            )
+        self._adaln_schedule_cache_active = cache is not None
+        return cache
 
     def forward_kwargs(
         self,
         *,
         video_rows: torch.Tensor,
         audio_rows: torch.Tensor,
-        t_video: float,
-        t_audio: float,
-        imgvid_cond_timestep: float,
-        audio_ref_cond_timestep: float,
+        timestep_metadata: MiniMaxH3StepTimestepMetadata,
+        step_index: int | None = None,
     ) -> dict[str, Any]:
         x = self.x_base.clone()
         x[0].index_copy_(0, self.img_pos_dev, video_rows)
         audio_x = self.audio_x_base.clone()
         audio_x[0].index_copy_(0, self.audio_pos_dev, audio_rows)
-        # Packed-sequence timestep semantics: non-media rows (text and
-        # padding) inherit the current video timestep. Later steps must reuse
-        # the previous step's updated rows; re-initializing from zeros is only
-        # valid at step 0.
-        timesteps = torch.full(
-            (self.seq_len,),
-            float(t_video),
-            dtype=torch.float32,
-            device=x.device,
-        )
-        timesteps[self.img_pos_dev[self.update_mask_dev]] = t_video
-        timesteps[self.img_pos_dev[~self.update_mask_dev]] = imgvid_cond_timestep
-        timesteps[self.audio_pos_dev[self.audio_update_mask_dev]] = t_audio
-        timesteps[self.audio_pos_dev[~self.audio_update_mask_dev]] = audio_ref_cond_timestep
-        unique_timesteps, inverse_indices = torch.unique(timesteps, sorted=True, return_inverse=True)
-        return {
+        dynamic_kwargs = {
             **self.static_kwargs,
             "x": x,
             "audio_x": audio_x,
-            "unique_timesteps": unique_timesteps,
-            "inverse_indices": inverse_indices,
+            "unique_timesteps": timestep_metadata.unique_timesteps,
+            "inverse_indices": timestep_metadata.inverse_indices,
         }
+        if self._adaln_schedule_cache_active:
+            if step_index is None:
+                raise ValueError("step_index is required while the AdaLN schedule cache is active")
+            dynamic_kwargs["adaln_step_index"] = step_index
+        return dynamic_kwargs
 
 
 def minimax_h3_denoise_loop(
@@ -153,6 +236,12 @@ def minimax_h3_denoise_loop(
         raise ValueError("video/audio sigma schedules must have equal length")
     if len(sigmas_video) < 2:
         raise ValueError("sigma schedules need at least 2 entries")
+    timestep_metadata = positive.prepare_timestep_metadata(
+        sigmas_video=sigmas_video,
+        sigmas_audio=sigmas_audio,
+        imgvid_cond_noise_aug_for_inference=imgvid_cond_noise_aug_for_inference,
+        audio_cond_noise_aug_for_inference=audio_cond_noise_aug_for_inference,
+    )
     n_cond = int((~positive.update_mask).sum())
     if keyframe_cond_rows is None:
         if n_cond != 0:
@@ -187,6 +276,16 @@ def minimax_h3_denoise_loop(
     if audio_anchor is not None:
         audio_rows[~audio_update] = audio_anchor
 
+    prepare_for_request_caches = getattr(model, "prepare_for_request_caches", None)
+    if prepare_for_request_caches is not None:
+        with torch.inference_mode():
+            prepare_for_request_caches()
+    positive.prepare_adaln_schedule_cache(
+        model,
+        timestep_metadata,
+        enabled=minimax_h3_adaln_schedule_cache_enabled(),
+    )
+
     num_steps = len(sigmas_video) - 1
     for step in range(num_steps):
         step_cm = step_profiler(step) if step_profiler is not None else nullcontext()
@@ -194,16 +293,12 @@ def minimax_h3_denoise_loop(
             s_v, s_v_next = sigmas_video[step], sigmas_video[step + 1]
             s_a, s_a_next = sigmas_audio[step], sigmas_audio[step + 1]
             t_v, t_a = 1.0 - s_v, 1.0 - s_a
-            imgvid_cond_t = max(t_v, float(imgvid_cond_noise_aug_for_inference))
-            audio_ref_cond_t = max(t_a, float(audio_cond_noise_aug_for_inference))
 
             fk = positive.forward_kwargs(
                 video_rows=video_rows,
                 audio_rows=audio_rows,
-                t_video=t_v,
-                t_audio=t_a,
-                imgvid_cond_timestep=imgvid_cond_t,
-                audio_ref_cond_timestep=audio_ref_cond_t,
+                timestep_metadata=timestep_metadata[step],
+                step_index=step,
             )
             with torch.inference_mode():
                 v_video, v_audio = model(**fk)
@@ -245,5 +340,6 @@ __all__ = [
     "MINIMAX_H3_IMGVID_COND_TIMESTEP",
     "MINIMAX_H3_VIDEO_ROW_WIDTH",
     "MiniMaxH3DenoiseBranch",
+    "MiniMaxH3StepTimestepMetadata",
     "minimax_h3_denoise_loop",
 ]

vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py

@@ -16,6 +16,7 @@ from typing import TYPE_CHECKING, Any
 import torch
 import torch.nn as nn
 from cache_dit import ForwardPattern
+from torch.distributed.tensor import DTensor
 from vllm.distributed import get_tensor_model_parallel_world_size
 from vllm.logger import init_logger
 from vllm.model_executor.layers.linear import (
@@ -33,6 +34,14 @@ from vllm_omni.diffusion.distributed.sp_plan import (
     SequenceParallelInput,
     SequenceParallelOutput,
 )
+from vllm_omni.platforms import current_omni_platform
+
+from .adaln_schedule_cache import (
+    MiniMaxH3AdalnScheduleCache,
+    MiniMaxH3AdalnScheduleKey,
+    minimax_h3_adaln_schedule_cache_enabled,
+    minimax_h3_build_adaln_table,
+)
 
 if TYPE_CHECKING:
     from vllm.model_executor.layers.quantization.base_config import (
@@ -132,6 +141,7 @@ _FORWARD_SUPPORTED_KWARGS = frozenset(
         "img_pos_for_infer_output_info",
         "packed_seq_params",
         "refiner_packed_seq_params",
+        "adaln_step_index",
     }
 )
 
@@ -248,7 +258,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 +269,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 +278,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 +336,7 @@ class MiniMaxH3Attention(nn.Module):
         arch: MiniMaxH3DiTArchConfig,
         quant_config: QuantizationConfig | None,
         *,
+        prefix: str,
         skip_sequence_parallel: bool = False,
     ) -> None:
         super().__init__()
@@ -338,11 +352,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 +366,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 +377,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 +469,8 @@ class MiniMaxH3MLP(nn.Module):
         self,
         arch: MiniMaxH3DiTArchConfig,
         quant_config: QuantizationConfig | None,
+        *,
+        prefix: str,
     ) -> None:
         super().__init__()
         self.fc1 = MergedColumnParallelLinear(
@@ -481,8 +480,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 +491,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 +519,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,15 +536,120 @@ class MiniMaxH3AdalnProj(nn.Module):
             gather_output=True,
             params_dtype=_BF16_DTYPE,
             quant_config=quant_config,
+            prefix=f"{prefix}.linear",
+        )
+        self.register_buffer("_precomputed_table", None, persistent=False)
+        self._precomputed_step: torch.Tensor | None = None
+        self._offloaded_projection_state: dict[str, tuple[torch.Tensor, torch.device, torch.Size]] | None = None
+
+    @property
+    def projection_weights_offloaded(self) -> bool:
+        return self._offloaded_projection_state is not None
+
+    @property
+    def projection_nbytes(self) -> int:
+        state = self._offloaded_projection_state
+        if state is not None:
+            return sum(tensor.numel() * tensor.element_size() for tensor, _, _ in state.values())
+        total = 0
+        for parameter in self.linear.parameters():
+            local = parameter.to_local() if isinstance(parameter, DTensor) else parameter
+            total += local.numel() * local.element_size()
+        return total
+
+    def can_offload_projection(self) -> bool:
+        """Whether this projection can safely own ordinary parameter storage."""
+        if self.projection_weights_offloaded:
+            return True
+        parameters = tuple(self.linear.parameters())
+        return bool(parameters) and all(
+            not isinstance(parameter, DTensor) and not parameter.is_meta and parameter.numel() > 0
+            for parameter in parameters
         )
 
-    def forward(self, t_emb: torch.Tensor) -> tuple[torch.Tensor, ...]:
-        """t_emb: [M, t_dim] -> expand_ratio tensors of [M*modality_num, H]."""
+    def offload_projection(self) -> int:
+        """Move exact projection parameters to CPU and release device storage."""
+        if self.projection_weights_offloaded:
+            return 0
+        if self._precomputed_table is None:
+            raise RuntimeError("AdaLN projection weights require an installed precomputed table before offload")
+        if not self.can_offload_projection():
+            raise RuntimeError("AdaLN projection weight offload does not support DTensor, meta, or empty parameters")
+
+        state: dict[str, tuple[torch.Tensor, torch.device, torch.Size]] = {}
+        try:
+            for name, parameter in self.linear.named_parameters():
+                cpu_tensor = parameter.detach().to(device="cpu", copy=True)
+                state[name] = (cpu_tensor, parameter.device, parameter.shape)
+                parameter.data = torch.empty((0,), device=parameter.device, dtype=parameter.dtype)
+        except Exception:
+            for name, (cpu_tensor, device, shape) in state.items():
+                self.linear.get_parameter(name).data = cpu_tensor.to(device=device).view(shape)
+            raise
+        self._offloaded_projection_state = state
+        return sum(tensor.numel() * tensor.element_size() for tensor, _, _ in state.values())
+
+    def restore_projection(self) -> None:
+        """Restore a previously offloaded projection to its original devices."""
+        state = self._offloaded_projection_state
+        if state is None:
+            return
+        restored = {
+            name: cpu_tensor.to(device=device).view(shape) for name, (cpu_tensor, device, shape) in state.items()
+        }
+        for name, tensor in restored.items():
+            self.linear.get_parameter(name).data = tensor
+        self._offloaded_projection_state = None
+
+    def compute_flat(self, t_emb: torch.Tensor) -> torch.Tensor:
+        """Run the original projection and return `[M*modalities, ratio*H]`."""
         x = nn.functional.silu(t_emb)
         x, _ = self.linear(x.to(self.linear.weight.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))
+        return x.view(m * self.modality_num, self.expand_ratio * self.hidden_size)
+
+    def install_precomputed(
+        self,
+        table: torch.Tensor,
+        step_cursor: torch.Tensor,
+    ) -> None:
+        """Atomically select a non-persistent exact-schedule modulation table."""
+        expected_width = self.expand_ratio * self.hidden_size
+        if table.dim() != 3 or table.shape[-1] != expected_width:
+            raise ValueError(
+                f"precomputed AdaLN table must be [steps, rows, {expected_width}], got {list(table.shape)}"
+            )
+        if table.dtype != self.linear.weight.dtype:
+            raise ValueError(f"precomputed AdaLN table dtype mismatch: {table.dtype} != {self.linear.weight.dtype}")
+        if step_cursor.dtype != torch.long or step_cursor.dim() != 0:
+            raise ValueError("precomputed AdaLN step cursor must be a scalar torch.long tensor")
+        if table.device != step_cursor.device:
+            raise ValueError(f"precomputed AdaLN table/cursor device mismatch: {table.device} != {step_cursor.device}")
+        self._buffers["_precomputed_table"] = table
+        self._precomputed_step = step_cursor
+
+    def clear_precomputed(self) -> None:
+        """Restore direct projection without changing checkpoint parameters."""
+        self.restore_projection()
+        self._buffers["_precomputed_table"] = None
+        self._precomputed_step = None
+
+    @torch.compiler.disable
+    def forward(self, t_emb: torch.Tensor) -> tuple[torch.Tensor, ...]:
+        """t_emb: [M, t_dim] -> expand_ratio tensors of [M*modality_num, H]."""
+        table = self._precomputed_table
+        if table is None:
+            rows = self.compute_flat(t_emb)
+            return tuple(rows.chunk(self.expand_ratio, dim=-1))
+
+        step_cursor = self._precomputed_step
+        if step_cursor is None:
+            raise RuntimeError("precomputed AdaLN table is installed without a step cursor")
+        row_count = t_emb.shape[0] * self.modality_num
+        if row_count > table.shape[1]:
+            raise ValueError(f"precomputed AdaLN table row capacity is too small: {table.shape[1]} < {row_count}")
+        rows = table.index_select(0, step_cursor.reshape(1))[0].narrow(0, 0, row_count)
+        return tuple(rows.chunk(self.expand_ratio, dim=-1))
 
 
 class MiniMaxH3TokenRefinerBlock(nn.Module):
@@ -568,6 +659,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 +671,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 +702,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 +735,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 +811,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 +823,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 +831,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 +840,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(
@@ -772,6 +895,13 @@ class MiniMaxH3DiTModel(nn.Module):
     _repeated_blocks = ["MiniMaxH3DiTBlock"]
     _layerwise_offload_blocks_attrs = ["blocks"]
 
+    @staticmethod
+    def layerwise_offload_excluded_parameter_prefixes() -> tuple[str, ...]:
+        """Parameters owned by the exact-schedule cache instead of block hooks."""
+        if minimax_h3_adaln_schedule_cache_enabled():
+            return ("adaln_proj.linear.",)
+        return ()
+
     @staticmethod
     def _is_transformer_block(name: str, module: nn.Module) -> bool:
         del module
@@ -864,7 +994,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 +1003,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 +1013,37 @@ 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._adaln_projections = tuple(block.adaln_proj for block in self.blocks) + (self.final_layer.adaln_proj,)
+        self._adaln_schedule_cache: MiniMaxH3AdalnScheduleCache | None = None
         self._mark_missing_params_required()
 
     def _mark_missing_params_required(self) -> None:
@@ -903,6 +1058,126 @@ class MiniMaxH3DiTModel(nn.Module):
             if name in MINIMAX_H3_FP32_BUFFER_NAMES and buffer.dtype != _FP32_DTYPE:
                 raise ValueError(f"{name} must stay fp32 after load, got {buffer.dtype}.")
 
+    @torch.inference_mode()
+    def prepare_for_request_caches(self) -> None:
+        """Materialize a sequentially offloaded DiT before auxiliary cache work."""
+        registry = getattr(self, "_hook_registry", None)
+        if registry is None:
+            return
+        hook = registry.get_hook("sequential_offload")
+        if hook is not None:
+            hook.pre_forward(self)
+
+    @torch.inference_mode()
+    def prepare_adaln_schedule_cache(
+        self,
+        *,
+        unique_timestep_plan: tuple[torch.Tensor, ...],
+        schedule_key: MiniMaxH3AdalnScheduleKey,
+        offload_weights: bool = False,
+    ) -> MiniMaxH3AdalnScheduleCache | None:
+        """Build or reuse the single exact AdaLN schedule held by this model."""
+        if not unique_timestep_plan:
+            raise ValueError("AdaLN schedule must contain at least one denoise step")
+        normalized_key = tuple(tuple(int(value) for value in step) for step in schedule_key)
+        if len(normalized_key) != len(unique_timestep_plan):
+            raise ValueError(
+                f"AdaLN schedule key/plan length mismatch: {len(normalized_key)} != {len(unique_timestep_plan)}"
+            )
+        existing = self._adaln_schedule_cache
+        if existing is not None and existing.key == normalized_key:
+            return existing
+
+        for projection in self._adaln_projections:
+            projection.restore_projection()
+
+        t_embs = tuple(self.time_embedder(timestep.view(-1)) for timestep in unique_timestep_plan)
+        device = t_embs[0].device
+        if any(t_emb.device != device for t_emb in t_embs):
+            raise ValueError("AdaLN timestep embeddings must share one device")
+        step_cursor = torch.zeros((), dtype=torch.long, device=device)
+        tables = tuple(minimax_h3_build_adaln_table(projection, t_embs) for projection in self._adaln_projections)
+        table_bytes = sum(table.numel() * table.element_size() for table in tables)
+        projection_bytes = sum(projection.projection_nbytes for projection in self._adaln_projections)
+        offload_supported = offload_weights and all(
+            projection.can_offload_projection() for projection in self._adaln_projections
+        )
+        if offload_supported and table_bytes >= projection_bytes:
+            for projection in self._adaln_projections:
+                projection.clear_precomputed()
+            self._adaln_schedule_cache = None
+            logger.warning(
+                "MiniMax H3 AdaLN weight offload skipped because table bytes (%d) "
+                "would not save memory versus projection bytes (%d)",
+                table_bytes,
+                projection_bytes,
+            )
+            return None
+
+        # Build every table before changing the live projection path.
+        for projection, table in zip(self._adaln_projections, tables):
+            projection.install_precomputed(table, step_cursor)
+
+        offloaded_projection_bytes = 0
+        if offload_supported:
+            try:
+                offloaded_projection_bytes = sum(
+                    projection.offload_projection() for projection in self._adaln_projections
+                )
+            except Exception:
+                for projection in self._adaln_projections:
+                    projection.restore_projection()
+                logger.exception("MiniMax H3 AdaLN projection offload failed; retaining exact tables and weights")
+        elif offload_weights:
+            logger.warning("MiniMax H3 AdaLN projection offload is unsupported for the active parameter layout")
+
+        if offloaded_projection_bytes:
+            current_omni_platform.empty_cache()
+
+        cache = MiniMaxH3AdalnScheduleCache(
+            key=normalized_key,
+            t_embs=t_embs,
+            step_cursor=step_cursor,
+            steps=len(t_embs),
+            max_unique_timesteps=max(int(t_emb.shape[0]) for t_emb in t_embs),
+            table_bytes=table_bytes,
+            projection_bytes=projection_bytes,
+            offloaded_projection_bytes=offloaded_projection_bytes,
+            net_memory_saved_bytes=offloaded_projection_bytes - table_bytes,
+        )
+        self._adaln_schedule_cache = cache
+        if offloaded_projection_bytes:
+            gib = 1024**3
+            logger.info(
+                "MiniMax H3 AdaLN projection weights offloaded: %.3f GiB weights, "
+                "%.3f GiB tables, %.3f GiB net device-memory reduction",
+                offloaded_projection_bytes / gib,
+                table_bytes / gib,
+                cache.net_memory_saved_bytes / gib,
+            )
+        return cache
+
+    def clear_adaln_schedule_cache(self) -> None:
+        """Drop the model-local table and restore direct AdaLN projections."""
+        if self._adaln_schedule_cache is None:
+            return
+        for projection in self._adaln_projections:
+            projection.clear_precomputed()
+        self._adaln_schedule_cache = None
+
+    def _select_adaln_schedule_step(self, step_index: Any) -> torch.Tensor | None:
+        cache = self._adaln_schedule_cache
+        if cache is None:
+            if step_index is not None:
+                raise ValueError("adaln_step_index was supplied without a prepared AdaLN schedule")
+            return None
+        if isinstance(step_index, bool) or not isinstance(step_index, int):
+            raise ValueError("adaln_step_index must be an integer while the AdaLN schedule is active")
+        if step_index < 0 or step_index >= cache.steps:
+            raise ValueError(f"adaln_step_index {step_index} is outside [0, {cache.steps})")
+        cache.step_cursor.fill_(step_index)
+        return cache.t_embs[step_index]
+
     def load_weights(
         self,
         weights: Iterable[tuple[str, torch.Tensor]],
@@ -917,7 +1192,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
 
@@ -955,6 +1250,7 @@ class MiniMaxH3DiTModel(nn.Module):
         refiner_max_seqlen: int,
         seq_len: int,
         device: torch.device,
+        precomputed_t_emb: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """Build packed multimodal embeddings for the TP=1/SP=1 inference path.
 
@@ -980,7 +1276,7 @@ class MiniMaxH3DiTModel(nn.Module):
         embeddings.index_add_(0, img_pos, video_embed.to(_BF16_DTYPE)[: img_pos.shape[0]])
         embeddings.index_add_(0, audio_pos, audio_embed.to(_BF16_DTYPE)[: audio_pos.shape[0]])
 
-        t_emb = self.time_embedder(unique_timesteps)
+        t_emb = self.time_embedder(unique_timesteps) if precomputed_t_emb is None else precomputed_t_emb
         return embeddings, t_emb
 
     def forward(self, **kwargs: Any) -> tuple[torch.Tensor, torch.Tensor]:
@@ -1000,6 +1296,8 @@ class MiniMaxH3DiTModel(nn.Module):
                 f"{sorted(_FORWARD_SUPPORTED_KWARGS)}"
             )
 
+        precomputed_t_emb = self._select_adaln_schedule_step(kwargs.get("adaln_step_index"))
+
         x = _required_kwarg(kwargs, "x")
         audio_x = _required_kwarg(kwargs, "audio_x")
         img_position_ids = _required_kwarg(kwargs, "img_position_ids")
@@ -1052,6 +1350,7 @@ class MiniMaxH3DiTModel(nn.Module):
             refiner_max_seqlen=refiner_max,
             seq_len=seq_len,
             device=device,
+            precomputed_t_emb=precomputed_t_emb,
         )
 
         combined_indices = (inverse_indices * MINIMAX_H3_ADALN_MODALITY_NUM + token_tags.clamp(min=0)).to(device)

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(

vllm_omni/diffusion/offloader/layerwise_backend.py

@@ -41,6 +41,7 @@ class LayerwiseOffloadHook(ModelHook):
         device: torch.device,
         stream: current_omni_platform.Stream | None = None,
         pin_memory: bool = True,
+        excluded_parameter_prefixes: tuple[str, ...] = (),
     ):
         assert isinstance(next_block, nn.Module), "transformer block must be type `torch.nn.Module`"
 
@@ -48,6 +49,7 @@ class LayerwiseOffloadHook(ModelHook):
         self.device = device
         self.copy_stream = stream or current_omni_platform.current_stream()
         self.pin_memory = pin_memory
+        self.excluded_parameter_prefixes = tuple(excluded_parameter_prefixes)
 
         # Per-block synchronization primitive: set after H2D copy completes.
         self._prefetch_done: current_omni_platform.Event | None = None
@@ -90,10 +92,10 @@ class LayerwiseOffloadHook(ModelHook):
         # the input module is kept intact
         module = super().initialize_hook(module)
 
-        self.block_parameters: dict[str, nn.Parameter] = dict(module.named_parameters())
+        self.block_parameters = self._owned_parameters(module)
         self.block_buffers: dict[str, torch.Tensor] = dict(module.named_buffers())
 
-        self.next_block_parameters: dict[str, nn.Parameter] = dict(self.next_block.named_parameters())
+        self.next_block_parameters = self._owned_parameters(self.next_block)
         self.next_block_buffers: dict[str, torch.Tensor] = dict(self.next_block.named_buffers())
 
         # Pre-allocate gpu tensors in a flattened way
@@ -106,6 +108,13 @@ class LayerwiseOffloadHook(ModelHook):
 
         return module
 
+    def _owned_parameters(self, module: nn.Module) -> dict[str, nn.Parameter]:
+        return {
+            name: parameter
+            for name, parameter in module.named_parameters()
+            if not name.startswith(self.excluded_parameter_prefixes)
+        }
+
     @staticmethod
     def _to_cpu(
         params: dict[str, nn.Parameter],
@@ -255,9 +264,10 @@ def apply_block_hook(
     device: torch.device,
     stream: current_omni_platform.Stream | None = None,
     pin_memory: bool = True,
+    excluded_parameter_prefixes: tuple[str, ...] = (),
 ) -> LayerwiseOffloadHook:
     registry = HookRegistry.get_or_create(module)
-    hook = LayerwiseOffloadHook(next_block, device, stream, pin_memory)
+    hook = LayerwiseOffloadHook(next_block, device, stream, pin_memory, excluded_parameter_prefixes)
     registry.register_hook(LayerwiseOffloadHook._HOOK_NAME, hook)
 
     return hook
@@ -284,6 +294,20 @@ class LayerWiseOffloadBackend(OffloadBackend):
         self.copy_stream = current_omni_platform.Stream()
         self._blocks: list[list[nn.Module]] = []
 
+    @staticmethod
+    def _move_excluded_parameters_to_device(
+        module: nn.Module,
+        device: torch.device,
+        excluded_parameter_prefixes: tuple[str, ...],
+    ) -> int:
+        moved_bytes = 0
+        for name, parameter in module.named_parameters():
+            if not name.startswith(excluded_parameter_prefixes):
+                continue
+            moved_bytes += parameter.numel() * parameter.element_size()
+            parameter.data = parameter.data.to(device, non_blocking=True)
+        return moved_bytes
+
     def enable(self, pipeline: nn.Module) -> None:
         if self.enabled:
             logger.warning("LayerWiseOffloadBackend already enabled")
@@ -341,6 +365,30 @@ class LayerWiseOffloadBackend(OffloadBackend):
                 dit_module.to(self.device)
                 continue
 
+            excluded_parameter_prefixes: tuple[str, ...] = ()
+            get_excluded_prefixes = getattr(dit_module, "layerwise_offload_excluded_parameter_prefixes", None)
+            if callable(get_excluded_prefixes):
+                excluded_parameter_prefixes = tuple(get_excluded_prefixes())
+                if excluded_parameter_prefixes:
+                    logger.info(
+                        "Leaving layer-wise parameters matching %s under %s ownership",
+                        excluded_parameter_prefixes,
+                        dit_module.__class__.__name__,
+                    )
+                    moved_bytes = sum(
+                        self._move_excluded_parameters_to_device(
+                            block,
+                            self.device,
+                            excluded_parameter_prefixes,
+                        )
+                        for block in blocks
+                    )
+                    logger.info(
+                        "Materialized %.3f GiB of model-owned layer-wise parameters on %s",
+                        moved_bytes / 1024**3,
+                        self.device,
+                    )
+
             # Move non-block modules to GPU (they stay resident)
             for name, m in dit_module.named_children():
                 if name not in blocks_attr_names:
@@ -368,6 +416,7 @@ class LayerWiseOffloadBackend(OffloadBackend):
                 self.device,
                 self.copy_stream,
                 self.config.pin_cpu_memory,
+                excluded_parameter_prefixes,
             )
             last_hook.prefetch_layer(non_blocking=False)
 
@@ -381,6 +430,7 @@ class LayerWiseOffloadBackend(OffloadBackend):
                     self.device,
                     self.copy_stream,
                     self.config.pin_cpu_memory,
+                    excluded_parameter_prefixes,
                 )
                 block_hooks.append(hook)
 

vllm_omni/entrypoints/openai/video_api_utils.py

@@ -8,7 +8,6 @@ from __future__ import annotations
 
 import base64
 import binascii
-from collections import deque
 from io import BytesIO
 from typing import Any, Literal
 
@@ -16,6 +15,12 @@ import httpx
 import numpy as np
 import torch
 from PIL import Image, UnidentifiedImageError
+from vllm.multimodal.video import (
+    VIDEO_LOADER_REGISTRY,
+    VideoBackend,
+    VideoSourceMetadata,
+    VideoTargetMetadata,
+)
 
 from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError
 from vllm_omni.entrypoints.openai.protocol.videos import (
@@ -58,6 +63,30 @@ def _decode_image_bytes(image_bytes: bytes, *, source: str) -> Image.Image:
         raise InvalidInputReferenceError(f"Invalid {source}: provided content is not a valid image.") from exc
 
 
+@VIDEO_LOADER_REGISTRY.register("omni")
+class OmniVideoBackend(VideoBackend):
+    """Video backend that selects sequential first-N or last-N frames."""
+
+    @classmethod
+    def compute_frames_index_to_sample(
+        cls,
+        source: VideoSourceMetadata,
+        target: VideoTargetMetadata,
+        *,
+        keep: Literal["first", "last"] = "first",
+        **kwargs,
+    ) -> list[int]:
+        num_frames_to_sample = source.total_frames_num
+        if target.num_frames > 0:
+            num_frames_to_sample = min(num_frames_to_sample, target.num_frames)
+        num_frames_to_sample = max(1, num_frames_to_sample)
+        if num_frames_to_sample >= source.total_frames_num:
+            return list(range(source.total_frames_num))
+        if keep == "last":
+            return list(range(source.total_frames_num - num_frames_to_sample, source.total_frames_num))
+        return list(range(num_frames_to_sample))
+
+
 def _decode_video_bytes(
     video_bytes: bytes,
     *,
@@ -65,45 +94,29 @@ def _decode_video_bytes(
     max_frames: int | None = None,
     keep: Literal["first", "last"] = "first",
 ) -> VideoFrames:
-    try:
-        import av
-    except ImportError as exc:  # pragma: no cover - av is a serving dependency via media_utils
-        raise InvalidInputReferenceError(f"Invalid {source}: video decoding requires PyAV.") from exc
-
     if keep not in {"first", "last"}:
         raise InvalidInputReferenceError(f"Invalid {source}: video frame selection must be 'first' or 'last'.")
     if max_frames is not None and max_frames <= 0:
         raise InvalidInputReferenceError(f"Invalid {source}: max video frames must be positive.")
 
-    frames: list[Image.Image] = []
-    tail_frames: deque[Image.Image] | None = (
-        deque(maxlen=max_frames) if keep == "last" and max_frames is not None else None
-    )
-    fps: float | None = None
+    loader = VIDEO_LOADER_REGISTRY.load("omni")
+    num_frames = max_frames if max_frames is not None else -1
     try:
-        with av.open(BytesIO(video_bytes)) as container:
-            video_stream = container.streams.video[0] if container.streams.video else None
-            if video_stream is not None:
-                fps = (
-                    positive_float(getattr(video_stream, "average_rate", None))
-                    or positive_float(getattr(video_stream, "base_rate", None))
-                    or positive_float(getattr(video_stream, "guessed_rate", None))
-                )
-            for frame in container.decode(video=0):
-                image = frame.to_image().convert("RGB")
-                if tail_frames is not None:
-                    tail_frames.append(image)
-                else:
-                    frames.append(image)
-                if keep == "first" and max_frames is not None and len(frames) >= max_frames:
-                    break
+        frames_array, metadata = loader.load_bytes(
+            video_bytes,
+            num_frames=num_frames,
+            backend="pyav",
+            keep=keep,
+        )
     except Exception as exc:
         raise InvalidInputReferenceError(f"Invalid {source}: provided content is not a valid video.") from exc
 
-    if tail_frames is not None:
-        frames = list(tail_frames)
+    fps: float | None = positive_float(metadata.get("fps"))
+
+    frames = [Image.fromarray(f, "RGB") for f in frames_array]
     if not frames:
         raise InvalidInputReferenceError(f"Invalid {source}: provided content is not a valid video.")
+
     return VideoFrames(frames, fps=fps)
 
 

vllm_omni/worker/omni_connector_model_runner_mixin.py

@@ -25,7 +25,10 @@ from vllm.logger import init_logger
 
 from vllm_omni.data_entry_keys import OmniPayload
 from vllm_omni.distributed.omni_connectors.factory import OmniConnectorFactory
-from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec
+from vllm_omni.distributed.omni_connectors.utils.config import (
+    ConnectorSpec,
+    get_stage_connector_role,
+)
 from vllm_omni.outputs import OmniConnectorOutput
 from vllm_omni.worker.payload_span import (
     get_tensor_span,
@@ -46,6 +49,19 @@ if TYPE_CHECKING:
 logger = init_logger(__name__)
 
 
+def _should_create_payload_connector(model_config: Any) -> bool:
+    """Whether this stage owns runner payload transport for its edge.
+
+    Sender edges may instead be owned solely by KV transfer. Receivers still
+    need a connector even though they do not declare a downstream payload hook.
+    """
+    if get_stage_connector_role(model_config) != "sender":
+        return True
+
+    next_stage_func = getattr(model_config, "custom_process_next_stage_input_func", None)
+    return isinstance(next_stage_func, str) and bool(next_stage_func)
+
+
 def should_accumulate_full_payload_output(model_config, custom_process_func) -> bool:
     """Producer-side structural gate.
 
@@ -96,7 +112,9 @@ class OmniConnectorModelRunnerMixin:
             model_config: Stage-level model config with connector settings.
             kv_transfer_manager: Existing KV transfer manager to delegate to.
         """
-        self._omni_connector: OmniConnectorBase | None = self._create_connector(model_config)
+        self._omni_connector: OmniConnectorBase | None = (
+            self._create_connector(model_config) if _should_create_payload_connector(model_config) else None
+        )
         self._kv_transfer_manager = kv_transfer_manager
 
         self._async_chunk: bool = getattr(model_config, "async_chunk", False)