Fix LingBot diffusion alignment CI

Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>

.buildkite/test-merge.yml

@@ -627,11 +627,10 @@ steps:
       - label: "Diffusion · FLUX.2 Klein Test"
         source_file_dependencies:
           - tests/e2e/offline_inference/test_flux2_klein.py
-          - tests/e2e/online_serving/test_flux2_klein.py
           - vllm_omni/diffusion/models/flux2_klein/
         timeout_in_minutes: 30
         commands:
-          - pytest -s -v tests/e2e/offline_inference/test_flux2_klein.py tests/e2e/online_serving/test_flux2_klein.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
+          - pytest -s -v tests/e2e/offline_inference/test_flux2_klein.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
         agents:
           queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
         plugins:

.buildkite/test-nightly.yml

@@ -598,6 +598,38 @@ steps:
                       path: /mnt/hf-cache
                       type: DirectoryOrCreate
 
+  - group: ":card_index_dividers: Tiny Model Tests [Multi-GPU]"
+    key: nightly-tiny-model-test-group
+    depends_on: upload-nightly-pipeline
+    if: >-
+      build.env("NIGHTLY") == "1" ||
+      build.pull_request.labels includes "nightly-test"
+    steps:
+      - label: ":full_moon: Tiny Model Tests · Diffusion (Multi-GPU accelerations)"
+        timeout_in_minutes: 30
+        # NOTE: Single-GPU (core_model) tiny model tests run on the ready label,
+        # so nightly only adds the multi-GPU parallelism configurations.
+        #
+        # This is a bit confusing, but `full_model` is used to mark the multigpu tests, and
+        # --run-level `core_model` is used to indicate that we should use tiny model weights
+        # for the CI, since these tests should pass irrespective of whether or not we use
+        # the full weights or not.
+        commands:
+          - timeout 25m pytest -sv tests/model_tests/diffusion -m 'full_model and cuda' --run-level "core_model"
+        agents:
+          queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              shm-size: "8gb"
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+                - "HF_TOKEN"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
+
   # Diffusion X2I suite: x2i / x2a / x2t and related non-video paths; x2v is only in "Diffusion X2V Model Test" below.
   - group: ":card_index_dividers: Diffusion X2I(&A&T) Model Test"
     key: nightly-diffusion-x2iat-group
@@ -1287,6 +1319,44 @@ steps:
                       path: /mnt/hf-cache
                       type: DirectoryOrCreate
 
+      - label: ":full_moon: Diffusion X2V · Doc Test"
+        timeout_in_minutes: 60
+        commands:
+          - pytest -s -v tests/examples/offline_inference/test_image_to_video.py -m "full_model and example and H100" --run-level "full_model"
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 1
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
+
       - label: ":full_moon: Diffusion X2V · Other Function Test"
         timeout_in_minutes: 90
         commands:

.buildkite/test-ready.yml

@@ -157,6 +157,27 @@ steps:
                 - "/fsx/hf_cache:/fsx/hf_cache"
 
 
+  - group: ":card_index_dividers: Tiny Model Tests [Single GPU]"
+    depends_on: upload-ready-pipeline
+    steps:
+      - label: "Tiny Model Tests · Diffusion (base)"
+        commands:
+          - timeout 15m pytest -sv tests/model_tests/diffusion -m 'core_model and cuda' --run-level "core_model"
+        agents:
+          # Ready label only runs the cases that don't need parallelism for tiny
+          # model tests, which is usually just a smoke test for TTI and I2I
+          queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+                - "HF_TOKEN"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
+
   - label: "Engine&Model Executor Test"
     depends_on: upload-ready-pipeline
     commands:

.gitignore

@@ -208,6 +208,7 @@ cache/
 .cache/
 diffusion_cache/
 kv_cache/
+!vllm_omni/experimental/ar_diffusion/kv_cache/
 
 # Logs
 logs/

benchmarks/diffusion/lingbot_video_parity.py

@@ -0,0 +1,217 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import argparse
+import json
+import re
+import statistics
+import subprocess
+import sys
+import time
+from pathlib import Path
+from typing import Any
+
+import imageio.v3 as iio
+import numpy as np
+import torch
+
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.platforms import current_omni_platform
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="LingBot dense parity and timing harness.")
+    parser.add_argument("--model", default="/home/models/lingbot-video-dense-1.3b")
+    parser.add_argument("--official-repo", default="/tmp/lingbot-video")
+    parser.add_argument("--output-dir", default="/tmp/lingbot_dense_parity")
+    parser.add_argument("--prompt", default="a robotic arm picks up a red block")
+    parser.add_argument("--height", type=int, default=192)
+    parser.add_argument("--width", type=int, default=320)
+    parser.add_argument("--num-frames", type=int, default=9)
+    parser.add_argument("--steps", type=int, default=2)
+    parser.add_argument("--guidance-scale", type=float, default=3.0)
+    parser.add_argument("--shift", type=float, default=3.0)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--fps", type=int, default=24)
+    parser.add_argument("--runs", type=int, default=0, help="Optional in-process steady-state native runs.")
+    return parser.parse_args()
+
+
+def _run(cmd: list[str], cwd: Path | None = None) -> tuple[float, str]:
+    start = time.perf_counter()
+    proc = subprocess.run(cmd, cwd=cwd, text=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
+    elapsed = time.perf_counter() - start
+    if proc.returncode != 0:
+        print(proc.stdout)
+        raise subprocess.CalledProcessError(proc.returncode, proc.args, output=proc.stdout)
+    return elapsed, proc.stdout
+
+
+def _read_video(path: Path) -> np.ndarray:
+    frames = iio.imread(path, plugin="pyav")
+    if frames.dtype != np.float32:
+        frames = frames.astype(np.float32) / 255.0
+    return frames
+
+
+def _compare(a: np.ndarray, b: np.ndarray) -> dict[str, Any]:
+    if a.shape != b.shape:
+        return {"shape_a": list(a.shape), "shape_b": list(b.shape), "shape_match": False}
+    diff = a - b
+    mse = float(np.mean(diff * diff))
+    mae = float(np.mean(np.abs(diff)))
+    psnr = float("inf") if mse == 0 else float(20.0 * np.log10(1.0 / np.sqrt(mse)))
+    return {
+        "shape": list(a.shape),
+        "shape_match": True,
+        "mae": mae,
+        "mse": mse,
+        "psnr": psnr,
+    }
+
+
+def _extract_total_generation_seconds(log: str) -> float | None:
+    match = re.search(r"Total generation time:\s*([0-9.]+)\s*seconds", log)
+    return float(match.group(1)) if match else None
+
+
+def _peak_memory_mb(output: Any) -> float:
+    result = output[0] if isinstance(output, list) and output else output
+    return float(getattr(result, "peak_memory_mb", 0.0) or 0.0)
+
+
+def _benchmark_native_requests(args: argparse.Namespace) -> dict[str, Any]:
+    load_start = time.perf_counter()
+    omni = Omni(
+        model=args.model,
+        model_class_name="LingBotVideoPipeline",
+        flow_shift=args.shift,
+        parallel_config=DiffusionParallelConfig(),
+    )
+    load_seconds = time.perf_counter() - load_start
+    prompt = {"prompt": args.prompt}
+    latencies: list[float] = []
+    peak_mb = 0.0
+
+    for _ in range(args.runs):
+        generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
+        sampling_params = OmniDiffusionSamplingParams(
+            height=args.height,
+            width=args.width,
+            num_frames=args.num_frames,
+            num_inference_steps=args.steps,
+            guidance_scale=args.guidance_scale,
+            generator=generator,
+            output_type="pt",
+            extra_args={"shift": args.shift},
+        )
+        start = time.perf_counter()
+        output = omni.generate(prompt, sampling_params)
+        latencies.append(time.perf_counter() - start)
+        peak_mb = max(peak_mb, _peak_memory_mb(output))
+        del output
+    omni.close()
+
+    return {
+        "load_seconds": load_seconds,
+        "run_seconds": latencies,
+        "mean_seconds": statistics.mean(latencies) if latencies else None,
+        "median_seconds": statistics.median(latencies) if latencies else None,
+        "min_seconds": min(latencies) if latencies else None,
+        "max_seconds": max(latencies) if latencies else None,
+        "peak_mb": peak_mb,
+    }
+
+
+def main() -> None:
+    args = parse_args()
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    official_out = output_dir / "official_diffusers.mp4"
+    native_out = output_dir / "vllm_omni_native.mp4"
+
+    official_cmd = [
+        sys.executable,
+        str(Path(args.official_repo) / "scripts" / "inference.py"),
+        "--backend",
+        "diffusers",
+        "--model_dir",
+        args.model,
+        "--mode",
+        "t2v",
+        "--prompt",
+        args.prompt,
+        "--output",
+        str(official_out),
+        "--height",
+        str(args.height),
+        "--width",
+        str(args.width),
+        "--num_frames",
+        str(args.num_frames),
+        "--steps",
+        str(args.steps),
+        "--guidance_scale",
+        str(args.guidance_scale),
+        "--shift",
+        str(args.shift),
+        "--seed",
+        str(args.seed),
+        "--fps",
+        str(args.fps),
+        "--quiet_progress",
+    ]
+    native_cmd = [
+        sys.executable,
+        "examples/offline_inference/text_to_video/text_to_video_lingbot.py",
+        "--model",
+        args.model,
+        "--prompt",
+        args.prompt,
+        "--output",
+        str(native_out),
+        "--height",
+        str(args.height),
+        "--width",
+        str(args.width),
+        "--num-frames",
+        str(args.num_frames),
+        "--num-inference-steps",
+        str(args.steps),
+        "--guidance-scale",
+        str(args.guidance_scale),
+        "--flow-shift",
+        str(args.shift),
+        "--seed",
+        str(args.seed),
+        "--fps",
+        str(args.fps),
+    ]
+
+    official_time, official_log = _run(official_cmd, cwd=Path(args.official_repo))
+    native_time, native_log = _run(native_cmd, cwd=Path.cwd())
+
+    result = {
+        "official_seconds": official_time,
+        "native_process_seconds": native_time,
+        "native_request_seconds": _extract_total_generation_seconds(native_log),
+        "official_output": str(official_out),
+        "native_output": str(native_out),
+        "accuracy": _compare(_read_video(official_out), _read_video(native_out)),
+        "official_log_tail": official_log.splitlines()[-20:],
+        "native_log_tail": native_log.splitlines()[-20:],
+    }
+    if args.runs > 0:
+        result["native_steady_state"] = _benchmark_native_requests(args)
+
+    result_path = output_dir / "comparison.json"
+    result_path.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n")
+    print(json.dumps(result, indent=2, sort_keys=True))
+
+
+if __name__ == "__main__":
+    main()

docs/models/supported_models.md

@@ -46,6 +46,7 @@ th {
 | `LTX2ImageToVideoTwoStagesPipeline` | LTX-2-I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
 | `LTX23Pipeline` | LTX-2.3-T2V | `dg845/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
 | `LTX23ImageToVideoPipeline` | LTX-2.3-I2V | `dg845/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
+| `LingBotVideoPipeline` | LingBot-Video dense T2V | `robbyant/lingbot-video-dense-1.3b` | ✅︎ | | | |
 | `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | |
 | `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | |

docs/user_guide/diffusion/attention_backends.md

@@ -20,6 +20,8 @@ The full set of backends and their platform defaults is in the **Backend Options
 | `TORCH_SDPA` | PyTorch `scaled_dot_product_attention` with the default backend dispatcher. Most conservative; always available. |
 | `SAGE_ATTN` | SageAttention 2.2 — INT8-quantized attention with FP16 accumulation. Lossy but typically visually indistinguishable on diffusion outputs. Requires `sageattention`. |
 | `SAGE_ATTN_3` | Requires `sageattn3` from `SageAttention/sageattention3_blackwell`. CUDA only, intended for Blackwell GPUs, with GQA/MQA requests falling back to PyTorch SDPA. |
+| `FLASH_ATTN_HUB` | FlashAttention 2 from HuggingFace `kernels` library. Useful for train/rollout alignment. |
+| `FLASH_ATTN_3_HUB` | FlashAttention 3 from HuggingFace `kernels` library. CUDA Hopper (sm_90+) only; falls back to `FLASH_ATTN_HUB` on older GPUs. |
 
 
 ## Configuration
@@ -210,6 +212,36 @@ Notes:
 - `SAGE_ATTN_3` is only selected on CUDA when `sageattn3` is importable and the GPU is Blackwell-class.
 - SageAttention3's Blackwell kernel assumes `Hq == Hkv`. In vLLM-Omni, GQA/MQA diffusion requests fall back to PyTorch SDPA for correctness.
 
+## HuggingFace Kernels Hub Backends
+
+To achieve perfect numerical consistency between **training** (typically using HuggingFace Diffusers with Hub kernels) and **serving/rollout** (using vLLM-Omni), you can use the Hub-based attention backends. This eliminates numerical drift / sampling divergence caused by executing different local kernel versions during rollout.
+
+The following backend options are supported:
+- `FLASH_ATTN_HUB` (HuggingFace `kernels-community/flash-attn2`)
+- `FLASH_ATTN_3_HUB` (HuggingFace `kernels-community/flash-attn3`, Hopper sm_90+ only)
+
+### Installation
+
+To use these backends, you must install the `kernels` library:
+
+```bash
+pip install kernels==0.14.1
+```
+
+If the `kernels` library is not available in the environment, vLLM-Omni will log a warning and fall back gracefully to the corresponding local backend implementations (`FLASH_ATTN`). On CUDA GPUs below Hopper (compute capability < 9.0), `FLASH_ATTN_3_HUB` falls back to `FLASH_ATTN_HUB`.
+
+### Usage
+
+Select a Hub backend using the global CLI flag or environment variables:
+
+```bash
+# Environment variable
+export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN_HUB
+
+# CLI flag
+vllm-omni serve <model> --diffusion-attention-backend FLASH_ATTN_HUB
+```
+
 ## Usage Examples
 
 ### Default (auto-route)

docs/user_guide/diffusion_features.md

@@ -126,7 +126,7 @@ The following tables show which models support each feature:
 | **MagiHuman**            |     ❌     |     ❌      |           ❌           |       ❓        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **MammothModa2(T2I)**    |     ❌     |     ❌      |           ❌           |       ❌        |         ❌         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **Nextstep_1(T2I)**      |     ❓     |     ❓      |           ❌           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
-| **OmniGen2**             |     ❌     |     ✅      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
+| **OmniGen2**             |     ❌     |     ✅      |           ✅           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **Ovis-Image**           |     ❌     |     ✅      |           ❌           |       ✅        |         ❌         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **Qwen-Image**           |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ✅        |        ✅         |
 | **Qwen-Image-2512**      |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ✅        |        ✅         |

docs/user_guide/examples/offline_inference/mammothmodal2_preview.md

@@ -22,7 +22,7 @@ Text-to-image now runs through the shared offline image example
 ```bash
 python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py \
   --model ./MammothModa2-Preview \
-  --stage-config ./vllm_omni/model_executor/stage_configs/mammoth_moda2_ar.yaml \
+  --deploy-config ./vllm_omni/deploy/mammoth_moda2_ar.yaml \
   --question "Summarize this image." \
   --image ./image.png
 ```

examples/offline_inference/image_to_image/image_edit.py

@@ -138,6 +138,15 @@ def parse_args() -> argparse.Namespace:
         required=True,
         help="Path(s) to input image file(s) (PNG, JPG, etc.). Can specify multiple images.",
     )
+    parser.add_argument(
+        "--deploy-config",
+        type=str,
+        default=None,
+        help=(
+            "Path to a deploy YAML. Required for multi-stage image-edit pipelines "
+            "whose deploy config is not auto-loaded."
+        ),
+    )
     parser.add_argument(
         "--prompt",
         type=str,
@@ -349,7 +358,12 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--enforce-eager",
         action="store_true",
-        help="Disable torch.compile and force eager execution.",
+        default=None,
+        help=(
+            "Disable torch.compile and force eager execution. Left unset (None) "
+            "so it is only forwarded when explicitly given; "
+            "otherwise the per-stage deploy YAML value wins."
+        ),
     )
     parser.add_argument(
         "--vae-use-slicing",
@@ -468,7 +482,7 @@ def main():
         }
 
     # Initialize Omni with appropriate pipeline
-    omni = Omni(
+    omni_kwargs: dict[str, Any] = dict(
         model=args.model,
         enable_layerwise_offload=args.enable_layerwise_offload,
         vae_use_slicing=args.vae_use_slicing,
@@ -476,11 +490,15 @@ def main():
         cache_backend=args.cache_backend,
         cache_config=cache_config,
         parallel_config=parallel_config,
-        enforce_eager=args.enforce_eager,
         enable_cpu_offload=args.enable_cpu_offload,
         enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
         profiler_config=args.profiler_config,
     )
+    if args.enforce_eager is not None:
+        omni_kwargs["enforce_eager"] = args.enforce_eager
+    if args.deploy_config:
+        omni_kwargs["deploy_config"] = args.deploy_config
+    omni = Omni(**omni_kwargs)
     model_class_name = get_model_class_name(omni)
     declared_extra_body_params = get_extra_body_params(model_class_name)
     print("Pipeline loaded")

examples/offline_inference/image_to_video/README.md

@@ -3,14 +3,140 @@
 This shared example generates videos from images with VACE, Wan2.2,
 LTX-2/LTX-2.3, HunyuanVideo-1.5, Cosmos3, and other compatible pipelines.
 
-## Local CLI Usage
+- `image_to_video.py`: command-line script for single video generation with advanced options.
 
-Download the example image:
+## Table of Contents
+
+- [Overview](#overview)
+- [Prerequisites](#prerequisites)
+- [Quick Start](#quick-start)
+- [Key Arguments](#key-arguments)
+- [More CLI Examples](#more-cli-examples)
+- [Advanced Features](#advanced-features)
+- [FAQ](#faq)
+
+## Overview
+
+This folder provides a unified CLI script for image-to-video generation using vLLM-Omni diffusion/video pipelines. The script selects practical defaults for supported model families while still exposing common sampling, memory, and parallelism options.
+
+### Supported Models
+
+| Model | Default Resolution | Default Frames | Default Steps | Guidance | VRAM Notes |
+| ----- | ------------------ | -------------- | ------------- | -------- | ---------- |
+| `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | 480 x 832 | 81 | 50 | 5.0 | Around 60 GiB BF16 for basic single-card usage |
+| `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | 480 x 832 | 81 | 50 | 4.0 | Around 20–25 GiB BF16, smallest I2V model |
+| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v` | 480 x 832 | 121 | 50 | 6.0 | Around 100 GiB at default settings; the example enables `--enable-cpu-offload` + VAE tiling/slicing to fit an 80 GiB card |
+| LTX2 (local path + `--model-class-name LTX2ImageToVideoPipeline`) | 512 x 768 | 121 | 40 | 4.0 | Memory use depends on frame count and tensor parallelism |
+
+!!! info
+    Peak VRAM: based on basic single-card usage, batch size = 1, without any acceleration/optimization features. Some model weights cannot fit into one card with 80 GiB VRAM, which may need to use CPU offloading.
+
+Default model: `Wan-AI/Wan2.2-I2V-A14B-Diffusers`.
+
+## Prerequisites
+
+Download the example image used in the snippets below:
 
 ```bash
-wget https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg
+curl -L -o cherry_blossom.jpg https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg
 ```
 
+## Quick Start
+
+### Python API
+
+Single-prompt generation using TI2V-5B (lightweight):
+
+```python
+import PIL.Image
+import torch
+
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+if __name__ == "__main__":
+    image = PIL.Image.open("cherry_blossom.jpg").convert("RGB")
+    image = image.resize((576, 320))
+
+    omni = Omni(
+        model="Wan-AI/Wan2.2-TI2V-5B-Diffusers",
+        flow_shift=12.0,
+    )
+    outputs = omni.generate(
+        {
+            "prompt": "Cherry blossoms swaying gently in the breeze, petals falling",
+            "multi_modal_data": {"image": image},
+        },
+        OmniDiffusionSamplingParams(
+            height=320,
+            width=576,
+            num_frames=17,
+            num_inference_steps=20,
+            guidance_scale=4.0,
+            generator=torch.Generator(device="cuda").manual_seed(42),
+        ),
+    )
+    from diffusers.utils import export_to_video
+    frames = outputs[0].request_output.images
+    export_to_video(frames, "quick_test_i2v.mp4", fps=16)
+```
+
+### Local CLI Usage
+
+Quick test using TI2V-5B with a small resolution and few frames:
+
+```bash
+python image_to_video.py \
+  --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \
+  --image cherry_blossom.jpg \
+  --prompt "Cherry blossoms swaying gently in the breeze, petals falling, smooth motion" \
+  --height 320 \
+  --width 576 \
+  --num-frames 17 \
+  --guidance-scale 4.0 \
+  --num-inference-steps 20 \
+  --flow-shift 12.0 \
+  --fps 16 \
+  --output quick_test_i2v.mp4
+```
+
+## Key Arguments
+
+| Argument | Type | Default | Description |
+| -------- | ---- | ------- | ----------- |
+| `--model` | str | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | Diffusers I2V model ID or local path |
+| `--model-class-name` | str | `None` | Override model class name (e.g., `LTX2ImageToVideoPipeline`) |
+| `--image` | str | (required) | Path to input image |
+| `--prompt` | str | `""` | Text description of desired motion/animation |
+| `--negative-prompt` | str | `""` | Optional list of artifacts to suppress |
+| `--seed` | int | `42` | Random seed for deterministic sampling |
+| `--guidance-scale` | float | `5.0` | CFG scale |
+| `--guidance-scale-high` | float | `None` | Separate CFG for high-noise stage (MoE only) |
+| `--height` | int | auto | Video height (auto-calculated from image if not set). Multiples of 16 |
+| `--width` | int | auto | Video width (auto-calculated from image if not set). Multiples of 16 |
+| `--num-frames` | int | `81` | Number of frames |
+| `--num-inference-steps` | int | `50` | Number of denoising steps |
+| `--boundary-ratio` | float | `0.875` | Boundary split ratio for two-stage MoE models |
+| `--flow-shift` | float | `5.0` | Scheduler flow shift (5.0 for 720p, 12.0 for 480p) |
+| `--sample-solver` | str | `unipc` | Wan2.2 sampling solver (`unipc` or `euler` for Lightning/Distill) |
+| `--fps` | int | `None` | Frames per second for the saved MP4 |
+| `--frame-rate` | float | `None` | Generation frame rate for pipelines that require it (e.g., LTX2) |
+| `--output` | str | `i2v_output.mp4` | Path to save the generated video |
+| `--vae-use-slicing` | flag | off | Enable VAE slicing for memory optimization |
+| `--vae-use-tiling` | flag | off | Enable VAE tiling for memory optimization |
+| `--enable-cpu-offload` | flag | off | Enable CPU offloading for diffusion models |
+| `--enable-layerwise-offload` | flag | off | Enable layerwise offloading on DiT modules |
+| `--cfg-parallel-size` | int | `1` | Set to `2` to enable CFG Parallel |
+| `--tensor-parallel-size` | int | `1` | Tensor parallel size (effective for models that support TP, e.g. LTX2) |
+| `--ulysses-degree` | int | `1` | Ulysses sequence parallel degree |
+| `--ring-degree` | int | `1` | Ring sequence parallel degree |
+| `--cache-backend` | str | `None` | Cache backend: `cache_dit` or `tea_cache` |
+| `--use-hsdp` | flag | off | Enable Hybrid Sharded Data Parallel |
+| `--hsdp-shard-size` | int | `-1` | GPUs per shard group (-1 auto-calculates) |
+| `--hsdp-replicate-size` | int | `1` | Number of replica groups for HSDP |
+
+## More CLI Examples
+
 ### Wan2.2-I2V-A14B-Diffusers (MoE)
 
 ```bash
@@ -18,7 +144,7 @@ python image_to_video.py \
   --model Wan-AI/Wan2.2-I2V-A14B-Diffusers \
   --image cherry_blossom.jpg \
   --prompt "Cherry blossoms swaying gently in the breeze, petals falling, smooth motion" \
-  --negative-prompt "<optional quality filter>" \
+  --negative-prompt "low quality, blurry" \
   --height 480 \
   --width 832 \
   --num-frames 48 \
@@ -28,7 +154,7 @@ python image_to_video.py \
   --boundary-ratio 0.875 \
   --flow-shift 12.0 \
   --fps 16 \
-  --output i2v_output.mp4
+  --output i2v_wan_moe.mp4
 ```
 
 ### Wan2.2-TI2V-5B-Diffusers (Unified)
@@ -38,7 +164,7 @@ python image_to_video.py \
   --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \
   --image cherry_blossom.jpg \
   --prompt "Cherry blossoms swaying gently in the breeze, petals falling, smooth motion" \
-  --negative-prompt "<optional quality filter>" \
+  --negative-prompt "low quality, blurry" \
   --height 480 \
   --width 832 \
   --num-frames 48 \
@@ -46,7 +172,65 @@ python image_to_video.py \
   --num-inference-steps 40 \
   --flow-shift 12.0 \
   --fps 16 \
-  --output i2v_output.mp4
+  --output i2v_wan_ti2v.mp4
+```
+
+### HunyuanVideo-1.5 I2V (480p)
+
+```bash
+python image_to_video.py \
+  --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v \
+  --image cherry_blossom.jpg \
+  --prompt "Cherry blossoms swaying gently in the breeze, petals falling, smooth motion" \
+  --height 480 \
+  --width 832 \
+  --num-frames 121 \
+  --guidance-scale 6.0 \
+  --flow-shift 5.0 \
+  --num-inference-steps 50 \
+  --fps 24 \
+  --enable-cpu-offload \
+  --vae-use-tiling \
+  --vae-use-slicing \
+  --output hunyuan_i2v.mp4
+```
+
+### LTX2 Image-to-Video
+
+```bash
+python image_to_video.py \
+  --model /path/to/LTX-2 \
+  --model-class-name LTX2ImageToVideoPipeline \
+  --image cherry_blossom.jpg \
+  --prompt "A cinematic dolly shot of cherry blossoms" \
+  --height 512 \
+  --width 768 \
+  --num-frames 121 \
+  --num-inference-steps 40 \
+  --guidance-scale 4.0 \
+  --frame-rate 24 \
+  --fps 24 \
+  --output ltx2_i2v.mp4
+```
+
+## Advanced Features
+
+### CFG Parallel
+
+Set `--cfg-parallel-size 2` to enable CFG Parallel for faster inference on multi-GPU setups.
+See more examples in the [cfg_parallel user guide](../../../docs/user_guide/diffusion/parallelism/cfg_parallel.md).
+
+### Cache Acceleration
+
+Use `--cache-backend cache_dit` for Cache-DiT acceleration or `--cache-backend tea_cache` for Timestep Embedding Aware Cache:
+
+```bash
+python image_to_video.py \
+  --model Wan-AI/Wan2.2-I2V-A14B-Diffusers \
+  --image cherry_blossom.jpg \
+  --prompt "Cherry blossoms swaying gently in the breeze" \
+  --cache-backend cache_dit \
+  --output i2v_cached.mp4
 ```
 
 ### LTX-2.3
@@ -220,3 +404,11 @@ The VACE T2V command is documented in the shared
 
 For Wan2.2 LightX2V-converted local Diffusers directories and related LoRA
 assets, see the [LoRA guide](../../../docs/user_guide/diffusion/lora.md#wan22-lightx2v-offline-assembly).
+
+## FAQ
+
+**OOM errors**: Try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage. For very large models, add `--enable-cpu-offload` or `--enable-layerwise-offload`.
+
+**Auto-calculated resolution**: If `--height` and `--width` are not provided, the script calculates output dimensions from the input image while maintaining aspect ratio and targeting 480 x 832 area (or 512 x 768 for LTX2).
+
+**Wan2.2 MoE vs unified**: `I2V-A14B` is a larger Mixture-of-Experts model with separate low/high-noise DiT stages (use `--guidance-scale-high` and `--boundary-ratio`). `TI2V-5B` is a smaller unified T2V+I2V model that does not need these extra arguments.

examples/offline_inference/mammothmodal2_preview/README.md

@@ -19,7 +19,7 @@ hf download bytedance-research/MammothModa2-Preview --local-dir ./MammothModa2-P
 ```bash
 python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py \
   --model ./MammothModa2-Preview \
-  --stage-config ./vllm_omni/model_executor/stage_configs/mammoth_moda2_ar.yaml \
+  --deploy-config ./vllm_omni/deploy/mammoth_moda2_ar.yaml \
   --question "Summarize this image." \
   --image ./image.png
 ```

examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py

@@ -4,7 +4,7 @@ Offline inference example: MammothModa2 image summarization (single AR stage).
 Example:
   uv run python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py \
     --model path/to/MammothModa2-Preview \
-    --stage-config vllm_omni/model_executor/stage_configs/mammoth_moda2_ar.yaml \
+    --deploy-config vllm_omni/deploy/mammoth_moda2_ar.yaml \
     --image /path/to/input.jpg \
     --question "Please summarize the content of this image."
 """
@@ -27,9 +27,7 @@ DEFAULT_QUESTION = "Please summarize the content of this image."
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(description="MammothModa2 image summarization (offline, AR only).")
     parser.add_argument("--model", type=str, required=True, help="Path to model directory or model id.")
-    parser.add_argument(
-        "--stage-config", type=str, required=True, help="Path to stage config yaml (single-stage AR->text)."
-    )
+    parser.add_argument("--deploy-config", type=str, required=True, help="Path to deploy config yaml.")
     parser.add_argument("--image", type=str, required=True, help="Path to input image.")
     parser.add_argument("--question", type=str, default=DEFAULT_QUESTION, help="Question/instruction for the model.")
     parser.add_argument("--system", type=str, default=DEFAULT_SYSTEM, help="System prompt.")
@@ -73,7 +71,7 @@ def main() -> None:
 
     omni = Omni(
         model=args.model,
-        stage_configs_path=args.stage_config,
+        deploy_config=args.deploy_config,
         trust_remote_code=args.trust_remote_code,
         enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
     )

examples/offline_inference/ming_flash_omni/README.md

@@ -90,14 +90,40 @@ Pass `--deploy-config /path/to/your_deploy.yaml` to any of the commands above to
 
 ### Image generation (text-to-image / img2img)
 
-The diffusion-stage image-generation path is currently only wired through
-the online OpenAI-compatible chat endpoint (`/v1/chat/completions` with
-`"modalities": ["image"]`). For payloads, optional knobs
-(`steps`/`cfg`/`seed`/`byte5_text`/`negative_prompt`), and the img2img
-reference-image flow, see the
-[image-gen section in the online-serving README](../../online_serving/ming_flash_omni/README.md#image-generation-text-to-image).
-`end2end.py` does not yet exercise the imagegen stage; that is tracked as
-follow-up work.
+Image generation is served through the standard task examples.
+The diffusion-stage knobs are declared centrally in `vllm_omni/model_extras/ming_flash_omni.py` and routed via `--extra-body`. With deploy yaml assigned properly, the model can run through the shared example scripts like any other diffusion model:
+
+Text-to-image (offline):
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+    --model Jonathan1909/Ming-flash-omni-2.0 \
+    --deploy-config vllm_omni/deploy/ming_flash_omni_image.yaml \
+    --prompt "Please draw a cute cat." \
+    --height 1024 \
+    --width 1024 \
+    --extra-body '{"steps": 30, "cfg": 2.0, "seed": 42}' \
+    --output ming_flash_omni_t2i.png
+```
+
+Image-to-image (offline):
+```bash
+# Reference image: figures/cases/person_gen_05.png from the upstream Ming repo
+# https://github.com/inclusionAI/Ming/blob/3954fcb880ff5e61ff128bcf7f1ec344d46a6fe3/examples/vllm_demo.py
+wget https://raw.githubusercontent.com/inclusionAI/Ming/3954fcb880ff5e61ff128bcf7f1ec344d46a6fe3/figures/cases/person_gen_05.png
+
+python examples/offline_inference/image_to_image/image_edit.py \
+    --model Jonathan1909/Ming-flash-omni-2.0 \
+    --deploy-config vllm_omni/deploy/ming_flash_omni_image.yaml \
+    --image person_gen_05.png \
+    --prompt "Put a pair of sunglasses on the person." \
+    --extra-args '{"steps": 30, "cfg": 2.0, "seed": 42}' \
+    --output ming_flash_omni_i2i.png
+```
+
+For the online path and the full knob list
+(`steps`/`cfg`/`seed`/`byte5_text`/`negative_prompt`/`height`/`width`), see the
+[image-generation section in the recipe](../../../recipes/inclusionAI/Ming-flash-omni-2.0.md#image-generation-text-to-image--img2img).
 
 ## Online serving
 

examples/offline_inference/text_to_image/text_to_image.py

@@ -69,7 +69,16 @@ def parse_args() -> argparse.Namespace:
         "--stage-configs-path",
         type=str,
         default=None,
-        help="Path to a YAML file containing stage configurations for Omni.",
+        help="[Deprecated] Path to a legacy stage_args-format YAML. Prefer --deploy-config.",
+    )
+    parser.add_argument(
+        "--deploy-config",
+        type=str,
+        default=None,
+        help=(
+            "Path to a deploy YAML (new pipeline/stages format). Required for multi-stage "
+            "text-to-image pipelines whose deploy config is not auto-loaded."
+        ),
     )
     parser.add_argument("--prompt", default="a cup of coffee on the table", help="Text prompt for image generation.")
     parser.add_argument(
@@ -155,7 +164,12 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--enforce-eager",
         action="store_true",
-        help="Disable torch.compile and force eager execution.",
+        default=None,
+        help=(
+            "Disable torch.compile and force eager execution. Left unset (None) "
+            "so it is only forwarded when explicitly given; "
+            "otherwise the per-stage deploy YAML value wins."
+        ),
     )
     parser.add_argument(
         "--enable-cpu-offload",
@@ -221,8 +235,13 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--tensor-parallel-size",
         type=int,
-        default=1,
-        help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
+        default=None,
+        help=(
+            "Number of GPUs used for tensor parallelism (TP) inside the DiT. "
+            "Left unset so it is only forwarded when explicitly given; "
+            "otherwise the per-stage deploy YAML (or engine default of 1) wins. "
+            "Passing it always overrides the deploy YAML."
+        ),
     )
     parser.add_argument(
         "--enable-expert-parallel",
@@ -410,10 +429,8 @@ def main():
         "ring_degree": args.ring_degree,
         "ulysses_mode": args.ulysses_mode,
         "cfg_parallel_size": args.cfg_parallel_size,
-        "tensor_parallel_size": args.tensor_parallel_size,
         "vae_patch_parallel_size": args.vae_patch_parallel_size,
         "enable_expert_parallel": args.enable_expert_parallel,
-        "enforce_eager": args.enforce_eager,
         "enable_cpu_offload": args.enable_cpu_offload,
         "mode": "text-to-image",
         "log_stats": args.log_stats,
@@ -425,8 +442,14 @@ def main():
         **lora_args,
         **quant_kwargs,
     }
+    if args.tensor_parallel_size is not None:
+        omni_kwargs["tensor_parallel_size"] = args.tensor_parallel_size
+    if args.enforce_eager is not None:
+        omni_kwargs["enforce_eager"] = args.enforce_eager
     if args.stage_configs_path:
         omni_kwargs["stage_configs_path"] = args.stage_configs_path
+    if args.deploy_config:
+        omni_kwargs["deploy_config"] = args.deploy_config
     if use_nextstep:
         # NextStep-1.1 requires explicit pipeline class
         omni_kwargs["model_class_name"] = "NextStep11Pipeline"
@@ -451,8 +474,9 @@ def main():
     print(f"  Quantization: {args.quantization if args.quantization else 'None (BF16)'}")
     if ignored_layers:
         print(f"  Ignored layers: {ignored_layers}")
+    tp_display = args.tensor_parallel_size if args.tensor_parallel_size is not None else "deploy/default"
     print(
-        f"  Parallel configuration: tensor_parallel_size={args.tensor_parallel_size}, "
+        f"  Parallel configuration: tensor_parallel_size={tp_display}, "
         f"ulysses_degree={args.ulysses_degree}, ulysses_mode={args.ulysses_mode}, "
         f"ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}, "
         f"vae_patch_parallel_size={args.vae_patch_parallel_size}, "
@@ -522,8 +546,8 @@ def main():
         diffusion_params.extra_args.update({k: v for k, v in user_extra.items() if v is not None})
 
     if lora_request:
-        diffusion_params.extra_args["lora_request"] = lora_request
-        diffusion_params.extra_args["lora_scale"] = args.lora_scale
+        diffusion_params.lora_request = lora_request
+        diffusion_params.lora_scale = args.lora_scale
 
     # Build per-stage sampling params for multi-stage models (e.g. BAGEL),
     # or wrap single diffusion params for single-stage models.

examples/offline_inference/text_to_speech/omnivoice/end2end.py

@@ -34,7 +34,7 @@ def run_e2e():
         help="Model name or path (HuggingFace or local)",
     )
     parser.add_argument(
-        "--stage-config",
+        "--deploy-config",
         type=str,
         default="vllm_omni/deploy/omnivoice.yaml",
     )
@@ -87,14 +87,14 @@ def run_e2e():
     )
     args = parser.parse_args()
 
-    if not os.path.exists(args.stage_config):
-        raise FileNotFoundError(f"Stage config not found: {args.stage_config}")
+    if not os.path.exists(args.deploy_config):
+        raise FileNotFoundError(f"Deploy config not found: {args.deploy_config}")
 
     print(f"Initializing OmniVoice with model={args.model}")
 
     omni = Omni(
         model=args.model,
-        stage_configs_path=args.stage_config,
+        deploy_config=args.deploy_config,
         log_stats=True,
     )
 

examples/offline_inference/text_to_video/text_to_video_lingbot.py

@@ -0,0 +1,124 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import argparse
+import time
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.platforms import current_omni_platform
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Generate text-to-video with the LingBot dense preset.")
+    parser.add_argument("--model", default="/home/models/lingbot-video-dense-1.3b")
+    parser.add_argument("--prompt", default="a robotic arm picks up a red block")
+    parser.add_argument("--negative-prompt", default="")
+    parser.add_argument("--output", default="lingbot_video_output.mp4")
+    parser.add_argument("--height", type=int, default=192)
+    parser.add_argument("--width", type=int, default=320)
+    parser.add_argument("--num-frames", type=int, default=9)
+    parser.add_argument("--num-inference-steps", type=int, default=2)
+    parser.add_argument("--guidance-scale", type=float, default=3.0)
+    parser.add_argument("--flow-shift", type=float, default=3.0)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--fps", type=int, default=24)
+    parser.add_argument("--output-type", choices=["pt", "np"], default="pt")
+    parser.add_argument("--batch-cfg", action="store_true")
+    return parser.parse_args()
+
+
+def _peak_memory_mb(output: Any) -> float:
+    result = output[0] if isinstance(output, list) and output else output
+    return float(getattr(result, "peak_memory_mb", 0.0) or 0.0)
+
+
+def _extract_video(output: Any) -> Any:
+    if isinstance(output, list):
+        output = output[0] if output else None
+    if isinstance(output, OmniRequestOutput):
+        if output.final_output_type != "image":
+            raise ValueError(f"Unexpected output type {output.final_output_type!r}.")
+        if not output.images:
+            raise ValueError("No video frames found in OmniRequestOutput.")
+        return output.images[0] if len(output.images) == 1 else output.images
+    return output
+
+
+def _to_frame_array(video: Any) -> np.ndarray:
+    if isinstance(video, torch.Tensor):
+        tensor = video.detach().cpu()
+        if tensor.dim() == 5:
+            tensor = tensor[0] if tensor.shape[1] not in (3, 4) else tensor[0].permute(1, 2, 3, 0)
+        elif tensor.dim() == 4 and tensor.shape[0] in (3, 4):
+            tensor = tensor.permute(1, 2, 3, 0)
+        if tensor.is_floating_point():
+            if float(tensor.min()) < 0.0:
+                tensor = tensor.clamp(-1, 1) * 0.5 + 0.5
+            else:
+                tensor = tensor.clamp(0, 1)
+        return tensor.float().numpy()
+    if isinstance(video, np.ndarray):
+        array = video[0] if video.ndim == 5 else video
+        if np.issubdtype(array.dtype, np.integer):
+            return array.astype(np.float32) / 255.0
+        return array.astype(np.float32, copy=False)
+    if isinstance(video, list):
+        return np.stack([_to_frame_array(frame) for frame in video], axis=0)
+    return np.asarray(video, dtype=np.float32)
+
+
+def main() -> None:
+    args = parse_args()
+    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
+    omni = Omni(
+        model=args.model,
+        model_class_name="LingBotVideoPipeline",
+        flow_shift=args.flow_shift,
+        parallel_config=DiffusionParallelConfig(),
+    )
+
+    prompt: dict[str, str] = {"prompt": args.prompt}
+    if args.negative_prompt:
+        prompt["negative_prompt"] = args.negative_prompt
+
+    sampling_params = OmniDiffusionSamplingParams(
+        height=args.height,
+        width=args.width,
+        num_frames=args.num_frames,
+        num_inference_steps=args.num_inference_steps,
+        guidance_scale=args.guidance_scale,
+        generator=generator,
+        output_type=args.output_type,
+        extra_args={"batch_cfg": args.batch_cfg, "shift": args.flow_shift},
+    )
+
+    start = time.perf_counter()
+    output = omni.generate(prompt, sampling_params)
+    elapsed = time.perf_counter() - start
+    print(f"Total generation time: {elapsed:.4f} seconds ({elapsed * 1000:.2f} ms)")
+    peak_mb = _peak_memory_mb(output)
+    if peak_mb:
+        print(f"Worker peak GPU memory (reserved): {peak_mb:.2f} MiB ({peak_mb / 1024:.2f} GiB)")
+
+    video = _to_frame_array(_extract_video(output))
+    output_path = Path(args.output)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    from diffusers.utils import export_to_video
+
+    export_to_video(list(video), str(output_path), fps=args.fps)
+    print(f"Saved generated video to {output_path}")
+    omni.close()
+
+
+if __name__ == "__main__":
+    main()

examples/online_serving/ming_flash_omni/README.md

@@ -49,22 +49,14 @@ python examples/online_serving/openai_chat_completion_client_for_multimodal_gene
     --modalities text
 ```
 
-Parameterized curl wrapper in this directory:
 
-```bash
-bash run_curl_multimodal_generation.sh text
-bash run_curl_multimodal_generation.sh use_image
-bash run_curl_multimodal_generation.sh use_audio
-bash run_curl_multimodal_generation.sh use_video
-bash run_curl_multimodal_generation.sh use_mixed_modalities
-bash run_curl_multimodal_generation.sh use_image_gen
-```
+## Image generation (text-to-image / img2img)
 
-## Image generation (text-to-image)
+Ming-flash-omni-2.0 also exposes an image-generation (diffusion) stage. Launch with the image deploy YAML, which adds an image-gen stage behind the thinker.
 
-Ming-flash-omni-2.0 also exposes an image-generation (diffusion) stage.
-Launch with the image deploy YAML, which adds an image-gen stage behind
-the thinker:
+The image-generation stage is a standard vLLM-Omni diffusion pipeline (`MingImagePipeline`); its request knobs are declared in `vllm_omni/model_extras/ming_flash_omni.py` and routed through `extra_body`, so they no longer need a bespoke `sampling_params_list` recipe (that form is still available for per-stage thinker sampling — see below).
+
+### Launch
 
 ```bash
 vllm serve Jonathan1909/Ming-flash-omni-2.0 --omni \
@@ -74,32 +66,58 @@ vllm serve Jonathan1909/Ming-flash-omni-2.0 --omni \
     --port 8091
 ```
 
-Then request image output by passing `"modalities": ["image"]`:
+
+### Text-to-image
+
+Request image output with `"modalities": ["image"]`:
 
 ```bash
 curl -s http://127.0.0.1:8091/v1/chat/completions \
-    -H "Content-Type: application/json" \
-    -d '{
-      "model": "Jonathan1909/Ming-flash-omni-2.0",
-      "messages": [{"role": "user", "content": "Please draw a cute cat."}],
-      "modalities": ["image"]
-    }' | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' | base64 -d > ming_imagegen.png
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "Jonathan1909/Ming-flash-omni-2.0",
+    "messages": [
+      {
+        "role": "user",
+        "content": "Please draw a cute cat."
+      }
+    ],
+    "modalities": ["image"]
+  }' \
+  | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' \
+  | base64 -d > ming_imagegen.png
 ```
 
-### Optional knobs
+Pass generation knobs under a literal `extra_body` object (the OpenAI Python client's `extra_body=` kwarg produces the same request). Keys are filtered against the declared set and routed into every stage's `extra_args`:
 
-Pass image-gen overrides as flat keys on the diffusion-stage `sampling_params_list[1].extra_args`:
+```bash
+curl http://127.0.0.1:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "Jonathan1909/Ming-flash-omni-2.0",
+    "modalities": ["image"],
+    "messages": [
+      {
+        "role": "user",
+        "content": "Draw a poster."
+      }
+    ],
+    "extra_body": {
+      "steps": 6,
+      "cfg": 1.5,
+      "height": 512,
+      "width": 512,
+      "seed": 123,
+      "byte5_text": ["理解与生成统一"],
+      "negative_prompt": "ugly, blurry, distorted"
+    }
+  }' \
+  | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' \
+  | base64 -d > ming_imagegen_extra_body.png
+```
 
-| Key | Default | Description |
-| --- | --- | --- |
-| `height` / `width` | from config (1024) | Output resolution (multiples of `vae_scale_factor * 2`, currently 16). |
-| `steps` | 30 | Number of FlowMatchEuler denoise steps. |
-| `cfg` | 2.0 | Classifier-free guidance scale. |
-| `seed` | 42 | Per-request RNG seed (deterministic when ByT5 is also seed-stable). |
-| `byte5_text` | (auto) | Override the glyph text for ByT5 enhancement; raw strings are auto-wrapped to Ming's `Text "...". ` format. |
-| `negative_prompt` | empty | Real CFG negative conditioning (set on **stage-0 thinker** `extra_args` so `expand_cfg_prompts` spawns the companion). |
-
-Example with all the knobs:
+NOTE: `extra_body` does **not** set the thinker's own sampling params. To tune the thinker (stage-0) sampling (`temperature` / `top_p` / `top_k` / `max_tokens`) or place knobs explicitly per stage, use `sampling_params_list` (`[thinker, imagegen]`).
+`negative_prompt` must sit on the **stage-0** entry to trigger the real-CFG companion; the imagegen knobs go on **stage-1**:
 
 ```bash
 curl http://127.0.0.1:8091/v1/chat/completions \
@@ -143,27 +161,48 @@ curl http://127.0.0.1:8091/v1/chat/completions \
 
 ### img2img (reference image)
 
-Add an `image_url` content part to the user message; the parser routes it
-into the diffusion stage as `extra[reference_image]`:
+Add an `image_url` content part (a base64 data URL) to the user message; it is routed into the DiT stage as `extra[reference_image]`. base64-encode a local file and stream it through `jq` via stdin - piping `base64 -> jq -> curl` so that avoids the shell `ARG_MAX` limit that inlining a large base64 string in `-d '…'` hits.
 
-```jsonc
-"messages": [{
-  "role": "user",
-  "content": [
-    {"type": "text", "text": "Change the background to a sandy beach at sunset."},
-    {"type": "image_url", "image_url": {"url": "data:image/png;base64,<base64>"}}
-  ]
-}]
+```bash
+# Reference image: figures/cases/person_gen_05.png from the upstream Ming repo
+# Check https://github.com/inclusionAI/Ming/blob/3954fcb880ff5e61ff128bcf7f1ec344d46a6fe3/examples/vllm_demo.py
+wget https://raw.githubusercontent.com/inclusionAI/Ming/3954fcb880ff5e61ff128bcf7f1ec344d46a6fe3/figures/cases/person_gen_05.png
+
+base64 -w0 person_gen_05.png \
+| jq -Rs --arg prompt "Put a pair of sunglasses on the person." '{
+    model: "Jonathan1909/Ming-flash-omni-2.0",
+    modalities: ["image"],
+    messages: [
+      {
+        role: "user",
+        content: [
+          { type: "text", text: $prompt },
+          { type: "image_url", image_url: { url: ("data:image/png;base64," + (. | rtrimstr("\n"))) } }
+        ]
+      }
+    ]
+  }' \
+| curl http://127.0.0.1:8091/v1/chat/completions \
+    -H "Content-Type: application/json" \
+    -d @- \
+| jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' \
+| base64 -d > ming_img2img.png
 ```
 
-### GPU layout
+The reference image can also be a public URL (`"url": "https://…/photo.jpg"`) or the simplified `{"image": "<base64>"}` content-part form — see the [image-to-image request formats](../image_to_image/README.md#request-format).
 
-The shipped `ming_flash_omni_image.yaml` allocates the thinker on GPUs 0–3
-(TP=4) and the diffusion stage on GPU 4 (TP=1). Copy the YAML and edit
-`devices` per stage to relocate; with fewer GPUs available, drop
-the thinker TP to 2 and run the diffusion stage on a free card. Image-gen
-warmup takes roughly an extra 30–60 s on top of the thinker — set
-`--stage-init-timeout 1800` if the default 300 s is too tight.
+### Knobs (declared `extra_body` params)
+
+| Key | Default | Description |
+| --- | --- | --- |
+| `height` / `width` | 1024 | Output resolution (multiples of `vae_scale_factor * 2`, currently 16). |
+| `steps` | 30 | Number of FlowMatchEuler denoise steps. |
+| `cfg` | 2.0 | Classifier-free guidance scale. |
+| `seed` | 42 | Per-request RNG seed. |
+| `byte5_text` | (auto) | Glyph text for ByT5 enhancement; raw strings are auto-wrapped to Ming's `Text "…". ` format. Auto-extracted from quoted spans in the prompt when omitted. |
+| `negative_prompt` | (empty) | Real CFG negative conditioning (text-to-image only). |
+
+For the offline `text_to_image.py` / `image_edit.py` scripts and the full knob reference, see the [image-generation section in the recipe](../../../recipes/inclusionAI/Ming-flash-omni-2.0.md#image-generation-text-to-image--img2img).
 
 ## Modality control
 

examples/online_serving/ming_flash_omni/run_curl_multimodal_generation.sh

@@ -1,169 +0,0 @@
-#!/usr/bin/env bash
-set -euo pipefail
-
-# Server port
-PORT="${PORT:-8091}"
-# Default query type
-QUERY_TYPE="${1:-text}"
-
-# Validate query type
-if [[ ! "$QUERY_TYPE" =~ ^(text|use_audio|use_image|use_video|use_mixed_modalities|use_image_gen)$ ]]; then
-    echo "Error: Invalid query type '$QUERY_TYPE'"
-    echo "Usage: $0 [text|use_audio|use_image|use_video|use_mixed_modalities|use_image_gen]"
-    echo "  text: Text-only query"
-    echo "  use_audio: Audio + Text query"
-    echo "  use_image: Image + Text query"
-    echo "  use_video: Video + Text query"
-    echo "  use_mixed_modalities: Audio + Image + Video + Text query"
-    echo "  use_image_gen: Text-to-image (diffusion stage) — saves the generated PNG locally"
-    exit 1
-fi
-
-# Define URLs for assets
-MARY_HAD_LAMB_AUDIO_URL="https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/mary_had_lamb.ogg"
-CHERRY_BLOSSOM_IMAGE_URL="https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg"
-SAMPLE_VIDEO_URL="https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4"
-
-# Build user content based on query type
-case "$QUERY_TYPE" in
-  text)
-    user_content='[
-      {
-        "type": "text",
-        "text": "请详细介绍鹦鹉的生活习性。"
-      }
-    ]'
-    ;;
-  use_image)
-    user_content='[
-        {
-          "type": "image_url",
-          "image_url": {
-            "url": "'"$CHERRY_BLOSSOM_IMAGE_URL"'"
-          }
-        },
-        {
-          "type": "text",
-          "text": "Describe this image in detail."
-        }
-      ]'
-    ;;
-  use_audio)
-    user_content='[
-        {
-          "type": "audio_url",
-          "audio_url": {
-            "url": "'"$MARY_HAD_LAMB_AUDIO_URL"'"
-          }
-        },
-        {
-          "type": "text",
-          "text": "Please recognize the language of this speech and transcribe it. Format: oral."
-        }
-      ]'
-    ;;
-  use_video)
-    user_content='[
-        {
-          "type": "video_url",
-          "video_url": {
-            "url": "'"$SAMPLE_VIDEO_URL"'"
-          }
-        },
-        {
-          "type": "text",
-          "text": "Describe what is happening in this video."
-        }
-      ]'
-    ;;
-  use_mixed_modalities)
-    user_content='[
-        {
-          "type": "image_url",
-          "image_url": {
-            "url": "'"$CHERRY_BLOSSOM_IMAGE_URL"'"
-          }
-        },
-        {
-          "type": "audio_url",
-          "audio_url": {
-            "url": "'"$MARY_HAD_LAMB_AUDIO_URL"'"
-          }
-        },
-        {
-          "type": "text",
-          "text": "Describe the image, and recognize the language of this speech and transcribe it. Format: oral"
-        }
-      ]'
-    ;;
-esac
-
-echo "Running query type: $QUERY_TYPE"
-echo ""
-
-if [[ "$QUERY_TYPE" == "use_image_gen" ]]; then
-    # Image-gen branch: dual-stage diffusion endpoint, modalities=["image"].
-    # The shipped ming_flash_omni_image.yaml runs the thinker on cards 0-3
-    # (TP=4) and the diffusion stage on card 4 (TP=1). See the image-gen
-    # section in README.md for knobs and img2img.
-    out_path="${OUT_PATH:-/tmp/ming_imagegen.png}"
-    request_body=$(cat <<'EOF'
-{
-  "model": "Jonathan1909/Ming-flash-omni-2.0",
-  "modalities": ["image"],
-  "messages": [
-    {"role": "user", "content": "Draw a beautiful girl with short black hair and red dress."}
-  ]
-}
-EOF
-)
-    output=$(curl -sS --retry 3 --retry-delay 3 --retry-connrefused \
-        -X POST "http://localhost:${PORT}/v1/chat/completions" \
-        -H "Content-Type: application/json" \
-        -d "$request_body")
-    # Server returns base64-encoded PNG either as a data URL in
-    # ``choices[0].message.content[0].image_url.url`` or as a raw string.
-    b64=$(echo "$output" | jq -r '.choices[0].message.content
-        | if type == "array"
-            then (map(select(.image_url? != null)) | .[0].image_url.url // "")
-            elif type == "string" then . else "" end')
-    if [[ -z "$b64" ]]; then
-        echo "Error: no image returned. Raw response:"
-        echo "$output" | jq '.'
-        exit 1
-    fi
-    if [[ "$b64" == data:image* ]]; then b64="${b64#*,}"; fi
-    printf '%s' "$b64" | base64 --decode > "$out_path"
-    echo "Saved generated image: $out_path ($(stat -c%s "$out_path" 2>/dev/null || stat -f%z "$out_path") bytes)"
-    exit 0
-fi
-
-request_body=$(cat <<EOF
-{
-  "model": "Jonathan1909/Ming-flash-omni-2.0",
-  "modalities": ["text"],
-  "messages": [
-    {
-      "role": "system",
-      "content": [
-        {
-          "type": "text",
-          "text": "你是一个友好的AI助手。\n\ndetailed thinking off"
-        }
-      ]
-    },
-    {
-      "role": "user",
-      "content": $user_content
-    }
-  ]
-}
-EOF
-)
-
-output=$(curl -sS --retry 3 --retry-delay 3 --retry-connrefused \
-    -X POST http://localhost:${PORT}/v1/chat/completions \
-    -H "Content-Type: application/json" \
-    -d "$request_body")
-
-echo "Output of request: $(echo "$output" | jq '.choices[0].message.content')"

pyproject.toml

@@ -41,6 +41,7 @@ dev = [
     "pytest-cov>=4.0.0",
     "pytest-mock>=3.10.0",
     "websockets>=15.0.1",
+    "pytest-subtests>=0.11.0",
     "datasets>=2.14.0",
     "mypy==1.11.1",
     "pre-commit==4.0.1",

recipes/MammothModa2/MammothModa2-Preview.md

@@ -32,12 +32,12 @@ flags. Image size uses the standard `--height` / `--width` flags.
   [`examples/offline_inference/text_to_image/text_to_image.py`](../../examples/offline_inference/text_to_image/text_to_image.py)
 - Declared parameters:
   [`vllm_omni/model_extras/mammothmodal2_preview.py`](../../vllm_omni/model_extras/mammothmodal2_preview.py)
-- Stage config:
-  [`vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml`](../../vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml)
+- Deploy config:
+  [`vllm_omni/deploy/mammoth_moda2.yaml`](../../vllm_omni/deploy/mammoth_moda2.yaml)
 
 ## Hardware Support
 
-The default stage config runs both the AR and DiT stages on a single GPU
+The default deploy config runs both the AR and DiT stages on a single GPU
 (`devices: "0"`). The committed `gpu_memory_utilization` split (stage-0 AR `0.5`,
 stage-1 DiT `0.3`) is sized for an ~80 GB GPU. The model also fits on a 48 GB GPU
 after rebalancing the split so the AR weights (~23 GB) leave room for the KV
@@ -48,7 +48,7 @@ cache — see the note under *1x L40S 48GB*.
 ### 1x L40S 48GB
 
 > **48 GB config adjustment:** the committed
-> `vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml` uses
+> `vllm_omni/deploy/mammoth_moda2.yaml` uses
 > `gpu_memory_utilization` 0.5 / 0.3 (sized for ~80 GB). To fit on a 48 GB L40S,
 > set the stage-0 (AR) value to `0.8` and the stage-1 (DiT) value to `0.16`
 > before running. (On an ~80 GB GPU, leave the defaults unchanged.)
@@ -71,13 +71,13 @@ hf download bytedance-research/MammothModa2-Preview --local-dir ./MammothModa2-P
 ```
 
 Run text-to-image with the shared offline example from the repository root. The
-stage config sets `trust_remote_code`, so no extra flag is needed. Forward the
+deploy config sets `trust_remote_code`, so no extra flag is needed. Forward the
 MammothModa2 generation parameters as a JSON object through `--extra-body`:
 
 ```bash
 python examples/offline_inference/text_to_image/text_to_image.py \
   --model ./MammothModa2-Preview \
-  --stage-configs-path vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml \
+  --stage-configs-path vllm_omni/deploy/mammoth_moda2.yaml \
   --prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \
   --height 1024 \
   --width 1024 \

recipes/OpenBMB/MiniCPM-o-4_5.md

@@ -102,8 +102,9 @@ curl http://localhost:8099/v1/chat/completions \
 ```
 
 **Text + speech in one response** (the headline 4.5 feature). The TTS
-path is gated by a Jinja flag on the chat template, passed through
-`extra_body.chat_template_kwargs.use_tts_template=true`:
+path is gated by a Jinja flag on the chat template. Pass
+`use_tts_template=true` via the **top-level** `chat_template_kwargs`
+field (curl does not flatten nested `extra_body`):
 
 ```bash
 curl http://localhost:8099/v1/chat/completions \
@@ -112,12 +113,18 @@ curl http://localhost:8099/v1/chat/completions \
         "model": "openbmb/MiniCPM-o-4_5",
         "messages": [{"role": "user", "content": "Say hello, then introduce vLLM in one sentence."}],
         "modalities": ["text", "audio"],
-        "extra_body": {"chat_template_kwargs": {"use_tts_template": true}}
+        "chat_template_kwargs": {"use_tts_template": true}
     }'
 ```
 
-Response carries text in `choices[0].message.content` and base64 WAV
-in `choices[0].message.audio.data` (24 kHz mono, see Notes).
+When using the OpenAI Python SDK, the same flag can also be sent as
+`extra_body={"chat_template_kwargs": {"use_tts_template": True}}`
+because the client merges `extra_body` into the request root.
+
+Response carries text in one choice's `message.content` and base64 WAV
+in another choice's `message.audio.data` (24 kHz mono, see Notes). With
+`modalities: ["text", "audio"]` you typically get two `choices` entries
+(one text, one audio).
 
 **Gradio demo (text + image + audio + video UI)**:
 
@@ -209,9 +216,13 @@ vllm serve openbmb/MiniCPM-o-4_5 --omni \
   missing dep raises `ImportError` at first request with the same
   install hint instead of silently emitting empty audio.
 
-- **TTS trigger**: speech output is only emitted when the client passes
-  `extra_body.chat_template_kwargs.use_tts_template=true`. Without it,
-  the response is text-only (which is also faster).
+- **TTS trigger**: speech output requires
+  `chat_template_kwargs.use_tts_template=true` so the chat template
+  appends `<|tts_bos|>` before generation. Without it, Stage-1 talker
+  receives no TTS token span and returns silent audio (not text-only).
+  For **curl**, put `chat_template_kwargs` at the request root; nested
+  `extra_body.chat_template_kwargs` is ignored. The OpenAI Python SDK
+  may use `extra_body` because it flattens those fields into the root.
 
 - **Output audio**: 24 kHz mono WAV inside the OpenAI-style
   `message.audio.data` (base64). The Gradio demo's WAV player decodes

recipes/README.md

@@ -48,6 +48,7 @@ recipes/
 | [`LTX/LTX-2.3.md`](./LTX/LTX-2.3.md) | Text-to-video with audio generation (22B) | 1x GPU (96GB VRAM) |
 | [`MammothModa2/MammothModa2-Preview.md`](./MammothModa2/MammothModa2-Preview.md) | Text-to-image with the shared offline image example (AR → DiT) | 1x L40S 48GB / 1x ≥40GB GPU |
 | [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB |
+| [`Robbyant/LingBot-Video.md`](./Robbyant/LingBot-Video.md) | Native dense text-to-video serving | 1x L20X / A100 / H100 |
 | [`cosmos3/Cosmos3-Nano.md`](./cosmos3/Cosmos3-Nano.md) | Text-to-image, text-to-video, image-to-video, video-to-video generation, text to video with sound, action policy | 1x H200 141GB / B300 |
 | [`cosmos3/Cosmos3-Super.md`](./cosmos3/Cosmos3-Super.md) | 64B T2I / T2V / I2V / V2V generation (+ optional audio) / Action policy | 8x H200/H100/A100 / 2x H200 / B300 |
 | [`OpenBMB/MiniCPM-o-4_5.md`](./OpenBMB/MiniCPM-o-4_5.md) | Online serving for omni multimodal chat (text / image / audio / video → text + 24 kHz speech) | 2x A100/H100 80GB / 3x mid-tier GPU / 8x RTX 4090 24GB |

recipes/Robbyant/LingBot-Video.md

@@ -0,0 +1,143 @@
+# LingBot-Video
+
+> Native dense text-to-video serving for `robbyant/lingbot-video-dense-1.3b`
+
+## Summary
+
+- Vendor: Robbyant
+- Model: `robbyant/lingbot-video-dense-1.3b`
+- Task: Text-to-video generation
+- Mode: Offline generation and online serving with the OpenAI-compatible `/v1/videos` API
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe when you want to run the dense LingBot-Video checkpoint with
+vLLM-Omni's native pipeline. The runtime path does not import the upstream
+`lingbot_video` Python package; it loads the checkpoint components directly
+with the in-tree `LingBotVideoPipeline`, dense DiT transformer, shared FlowUniPC
+scheduler, Qwen3-VL text encoder, and Wan VAE.
+
+This recipe is intentionally text-to-video only. Text-to-image, image-conditioned
+video, TI2V, MoE checkpoints, and fused-expert kernels are not covered here.
+
+## References
+
+- Dense checkpoint: <https://huggingface.co/robbyant/lingbot-video-dense-1.3b>
+- Upstream project: <https://github.com/Robbyant/lingbot-video>
+- Related offline example: [`examples/offline_inference/text_to_video/text_to_video_lingbot.py`](../../examples/offline_inference/text_to_video/text_to_video_lingbot.py)
+- Related online video API docs: [`docs/serving/videos_api.md`](../../docs/serving/videos_api.md)
+
+## Hardware Support
+
+This recipe documents the CUDA single-GPU dense checkpoint path. Multi-GPU
+parallelism, Cache-DiT, CPU offload, and the LingBot MoE path are not validated
+for this model in this PR.
+
+## GPU
+
+### 1 x NVIDIA L20X / A100 / H100
+
+The dense 1.3B checkpoint has been smoke-tested on a single NVIDIA L20X at a
+small validation shape (`192x320`, 9 frames, 2 steps). Use a larger GPU or lower
+resolution/frame count if your target workload runs out of memory.
+
+#### Offline T2V
+
+```bash
+CUDA_VISIBLE_DEVICES=0 \
+python examples/offline_inference/text_to_video/text_to_video_lingbot.py \
+  --model robbyant/lingbot-video-dense-1.3b \
+  --prompt "a robotic arm picks up a red block" \
+  --output lingbot_t2v.mp4 \
+  --height 192 \
+  --width 320 \
+  --num-frames 9 \
+  --num-inference-steps 2 \
+  --guidance-scale 3.0 \
+  --flow-shift 3.0 \
+  --seed 42 \
+  --fps 24
+```
+
+#### Online Serving
+
+```bash
+CUDA_VISIBLE_DEVICES=0 \
+vllm serve robbyant/lingbot-video-dense-1.3b \
+  --omni \
+  --model-class-name LingBotVideoPipeline \
+  --port 8091
+```
+
+After the server is ready, submit a text-to-video job:
+
+```bash
+create_response=$(curl -s http://localhost:8091/v1/videos \
+  -F "model=robbyant/lingbot-video-dense-1.3b" \
+  -F "prompt=a robotic arm picks up a red block" \
+  -F "width=320" \
+  -F "height=192" \
+  -F "num_frames=9" \
+  -F "fps=24" \
+  -F "num_inference_steps=2" \
+  -F "guidance_scale=3.0" \
+  -F "flow_shift=3.0" \
+  -F "seed=42")
+
+video_id=$(echo "${create_response}" | jq -r '.id')
+while true; do
+  status=$(curl -s "http://localhost:8091/v1/videos/${video_id}" | jq -r '.status')
+  if [ "${status}" = "completed" ]; then
+    break
+  fi
+  if [ "${status}" = "failed" ]; then
+    curl -s "http://localhost:8091/v1/videos/${video_id}" | jq .
+    exit 1
+  fi
+  sleep 2
+done
+
+curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o lingbot_t2v.mp4
+```
+
+## Key Parameters
+
+| Parameter | Suggested smoke value | Notes |
+|-----------|-----------------------|-------|
+| `height` | `192` | Must be a multiple of 16 |
+| `width` | `320` | Must be a multiple of 16 |
+| `num_frames` | `9` | Must be `1` or `4n + 1`; this PR validates T2V with video outputs |
+| `num_inference_steps` | `2` | Use more steps for quality sweeps |
+| `guidance_scale` | `3.0` | CFG is active when this is greater than `1.0` |
+| `flow_shift` | `3.0` | Scheduler flow-shift; aliases the pipeline's internal `shift` |
+| `negative_prompt` | model default | Optional text describing artifacts to avoid |
+| `fps` | `24` | Output MP4 frame rate |
+
+## Validation
+
+Local dense smoke run:
+
+- Shape: 9 frames at `192x320`
+- Steps: 2
+- Request generation time: `0.2923s`
+- Peak reserved GPU memory: `14548 MiB`
+
+Local parity harness against the upstream repository:
+
+- Shape: `[9, 192, 320, 3]`
+- MAE: `0.0065238`
+- MSE: `0.00006650`
+- PSNR: `41.77 dB`
+- Native request time: `0.2875s`
+
+Do not treat these as production benchmarks; they are draft-PR validation
+numbers for the small dense smoke configuration.
+
+## Known Limitations
+
+- T2V only. T2I, I2V, and TI2V are not claimed by this PR.
+- Dense checkpoint only. MoE/fused-expert support belongs in a separate PR.
+- No Cache-DiT, TeaCache, tensor parallelism, sequence parallelism, CFG
+  parallelism, HSDP, CPU offload, or VAE patch parallelism validation yet.
+- Only one request per LingBot pipeline batch is currently supported.

recipes/inclusionAI/Ming-flash-omni-2.0.md

@@ -18,6 +18,7 @@ Use this recipe when you want a known-good starting point for serving
 
 - **Thinker only** — multimodal understanding with text output.
 - **Thinker + Talker (omni-speech)** — multimodal understanding with text and spoken output.
+- **Thinker + Imagegen (DiT)** — text-to-image / img2img
 - **Talker only (TTS)** — standalone text-to-speech via the OpenAI `/v1/audio/speech` endpoint.
 
 ## References
@@ -160,6 +161,129 @@ deltas appear at `choices[0].delta.content`.
 - Memory usage: size depends on output modalities and multimodal input; leave
   headroom for video frames and audio caches.
 
+## Image generation (text-to-image / img2img)
+
+Ming-flash-omni-2.0 also exposes an image-generation (diffusion) stage. Launch with the image deploy YAML, which adds an image-gen stage behind the thinker.
+
+The image-generation stage is a standard vLLM-Omni diffusion pipeline (`MingImagePipeline`); its request knobs are declared in `vllm_omni/model_extras/ming_flash_omni.py` and routed through `extra_body`, so they no longer need a bespoke `sampling_params_list` recipe (that form is still available for per-stage thinker sampling — see below).
+
+```bash
+vllm serve Jonathan1909/Ming-flash-omni-2.0 --omni \
+    --deploy-config vllm_omni/deploy/ming_flash_omni_image.yaml \
+    --stage-init-timeout 1800 \
+    --init-timeout 1800 \
+    --port 8091
+```
+
+With fewer GPUs, copy the YAML and drop the thinker to TP=2 with the DiT on a
+free card.
+
+### Online (text-to-image)
+
+Request image output with `"modalities": ["image"]`:
+
+```bash
+curl -s http://127.0.0.1:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "Jonathan1909/Ming-flash-omni-2.0",
+    "messages": [
+      {
+        "role": "user",
+        "content": "Please draw a cute cat."
+      }
+    ],
+    "modalities": ["image"]
+  }' \
+  | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' \
+  | base64 -d > ming_imagegen.png
+```
+
+**Quick form — `extra_body`** (keys are filtered against the declared set and routed into every stage's `extra_args`). Convenient, but it does **not** let you set the thinker (stage-0) sampling params. Wrap the knobs in a literal `"extra_body"` object for raw curl; the OpenAI Python client's `extra_body=` kwarg produces the same request:
+
+```bash
+curl http://127.0.0.1:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "Jonathan1909/Ming-flash-omni-2.0",
+    "modalities": ["image"],
+    "messages": [
+      {
+        "role": "user",
+        "content": "Draw a poster."
+      }
+    ],
+    "extra_body": {
+      "steps": 6,
+      "cfg": 1.5,
+      "height": 512,
+      "width": 512,
+      "seed": 123,
+      "byte5_text": ["理解与生成统一"],
+      "negative_prompt": "ugly, blurry, distorted"
+    }
+  }' \
+  | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' \
+  | base64 -d > ming_imagegen_extra_body.png
+```
+
+**Full control — `sampling_params_list`** (one entry per stage: `[thinker, imagegen]`).
+Use this when you need to tune the thinker's own sampling (`temperature` / `top_p` / `top_k` / `max_tokens`), or to place knobs explicitly per stage.
+Note `negative_prompt` must sit on the **stage-0 thinker** `extra_args`; the imagegen-stage knobs (`steps` / `cfg` / `height` / `width` / `seed` / `byte5_text`) go on the **stage-1** entry:
+
+```bash
+curl http://127.0.0.1:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "Jonathan1909/Ming-flash-omni-2.0",
+    "modalities": ["image"],
+    "sampling_params_list": [
+      {
+        "temperature": 0.4,
+        "top_p": 0.9,
+        "top_k": 1,
+        "max_tokens": 1,
+        "seed": 42,
+        "extra_args": {
+          "negative_prompt": "ugly, blurry, distorted"
+        }
+      },
+      {
+        "seed": 42,
+        "extra_args": {
+          "steps": 6,
+          "cfg": 1.5,
+          "height": 512,
+          "width": 512,
+          "seed": 123,
+          "byte5_text": ["理解与生成统一"]
+        }
+      }
+    ],
+    "messages": [
+      {
+        "role": "user",
+        "content": "Draw a poster."
+      }
+    ]
+  }' \
+  | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' \
+  | base64 -d > ming_imagegen_knobs.png
+```
+
+### Knobs (declared `extra_body` params)
+
+| Key | Default | Description |
+| --- | --- | --- |
+| `height` / `width` | 1024 | Output resolution (multiples of `vae_scale_factor * 2`, currently 16). |
+| `steps` | 30 | Number of FlowMatchEuler denoise steps. |
+| `cfg` | 2.0 | Classifier-free guidance scale. |
+| `seed` | 42 | Per-request RNG seed. |
+| `byte5_text` | (auto) | Glyph text for ByT5 enhancement; raw strings are auto-wrapped to Ming's `Text "…". ` format. Auto-extracted from quoted spans in the prompt when omitted. |
+| `negative_prompt` | (empty) | Real CFG negative conditioning. Spawns a CFG-text companion via `expand_cfg_prompts`; **online / text-to-image only** (offline uses Ming's default zero-negative). |
+
+For img2img, add an `image_url` content part to the user message (online) or pass `--image` (offline); the reference image is routed into the DiT stage as `extra[reference_image]`.
+
 ### 1x H100 80GB — standalone TTS (talker only)
 
 The bundled `ming_flash_omni_tts.yaml` runs the talker on a single GPU and exposes the OpenAI `/v1/audio/speech` endpoint.

requirements/common.txt

@@ -13,7 +13,7 @@ transformers >= 5.5.3
 av>=14.0.0
 omegaconf>=2.3.0
 diffusers==0.38.0
-safetensors>=0.8.0rc0
+safetensors>=0.8.0
 accelerate==1.12.0
 soundfile>=0.13.1
 cache-dit==1.3.0

tests/ar_diffusion/__init__.py

@@ -0,0 +1 @@
+# SPDX-License-Identifier: Apache-2.0

tests/ar_diffusion/test_engine_routing.py

@@ -0,0 +1,136 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Tests for the generic engine-backend dispatcher + AR-Diffusion runner wiring.
+
+``DiffusionEngine.resolve_engine_class`` is a generic dispatcher (``"default"`` / a
+``DiffusionEngine`` subclass / an import-path string). DreamZero selects the
+AR-Diffusion engine via its deploy config's ``engine_backend`` qualname — no
+DreamZero/ar_diffusion knowledge lives in the public base, so the routing check
+here is simply that the ``ARDiffusionEngine`` qualname resolves correctly.
+"""
+
+from dataclasses import fields
+from types import SimpleNamespace
+
+import pytest
+from vllm.utils.import_utils import resolve_obj_by_qualname
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
+from vllm_omni.experimental.ar_diffusion.engine import (
+    AR_DIFFUSION_MODEL_RUNNER_CLS,
+    ARDiffusionEngine,
+)
+
+
+def _cfg(backend):
+    return SimpleNamespace(engine_backend=backend)
+
+
+def test_default_resolves_to_diffusion_engine():
+    assert DiffusionEngine.resolve_engine_class(_cfg("default")) is DiffusionEngine
+
+
+def test_subclass_type_is_returned():
+    assert DiffusionEngine.resolve_engine_class(_cfg(ARDiffusionEngine)) is ARDiffusionEngine
+
+
+def test_ar_diffusion_qualname_resolves():
+    # How DreamZero's deploy config selects the engine: a full import-path string.
+    cls = DiffusionEngine.resolve_engine_class(_cfg("vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine"))
+    assert cls is ARDiffusionEngine
+
+
+def test_non_engine_type_raises():
+    with pytest.raises(TypeError):
+        DiffusionEngine.resolve_engine_class(_cfg(dict))
+
+
+def test_qualname_not_an_engine_raises():
+    with pytest.raises(TypeError):
+        DiffusionEngine.resolve_engine_class(_cfg("builtins.dict"))
+
+
+def test_bad_qualname_raises():
+    with pytest.raises(ValueError):
+        DiffusionEngine.resolve_engine_class(_cfg("not.a.real.module.NoSuchEngine"))
+
+
+def test_ar_diffusion_engine_is_diffusion_engine_subclass():
+    assert issubclass(ARDiffusionEngine, DiffusionEngine)
+
+
+def test_config_engine_backend_field_defaults_to_default():
+    field = {f.name: f for f in fields(OmniDiffusionConfig)}["engine_backend"]
+    assert field.default == "default"
+
+
+# --- AR-Diffusion worker/runner wiring -----------------------------------------------
+
+
+def test_ar_diffusion_model_runner_cls_resolves():
+    cls = resolve_obj_by_qualname(AR_DIFFUSION_MODEL_RUNNER_CLS)
+    assert cls.__name__ == "ARDiffusionModelRunner"
+
+
+def test_engine_class_declares_its_runner():
+    """Routing lives on the engine class (review: hsliuustc0106) — engines
+    never mutate od_config; the worker resolves the runner off the class."""
+    assert ARDiffusionEngine.default_diffusion_model_runner_cls == AR_DIFFUSION_MODEL_RUNNER_CLS
+    assert DiffusionEngine.default_diffusion_model_runner_cls is None
+
+
+def _select_runner(od_config, platform_default="PLATFORM"):
+    """Mirror of the worker's runner-selection chain (override > engine > platform)."""
+    runner_override = getattr(od_config, "diffusion_model_runner_cls", None)
+    if isinstance(runner_override, str) and runner_override:
+        return runner_override
+    engine_cls = DiffusionEngine.resolve_engine_class(od_config)
+    engine_runner = getattr(engine_cls, "default_diffusion_model_runner_cls", None)
+    if isinstance(engine_runner, str) and engine_runner:
+        return engine_runner
+    return platform_default
+
+
+def test_runner_selection_chain():
+    ar = "vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine"
+    # engine-declared runner wins over platform default
+    assert (
+        _select_runner(SimpleNamespace(diffusion_model_runner_cls=None, engine_backend=ar))
+        == AR_DIFFUSION_MODEL_RUNNER_CLS
+    )
+    # explicit override wins over the engine's declaration
+    assert (
+        _select_runner(SimpleNamespace(diffusion_model_runner_cls="my.custom.Runner", engine_backend=ar))
+        == "my.custom.Runner"
+    )
+    # default engine -> platform default
+    assert _select_runner(SimpleNamespace(diffusion_model_runner_cls=None, engine_backend="default")) == "PLATFORM"
+
+
+def test_dreamzero_pipeline_rejects_non_ar_diffusion_engine():
+    """Review (hsliuustc0106): a stale config with engine_backend='default'
+    must fail at init, not mid-forward on the first KV access."""
+    from types import SimpleNamespace
+
+    import pytest
+
+    from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
+
+    with pytest.raises(ValueError, match="requires the AR-Diffusion engine"):
+        DreamZeroPipeline(od_config=SimpleNamespace(engine_backend="default"))
+    with pytest.raises(ValueError, match="requires the AR-Diffusion engine"):
+        DreamZeroPipeline(od_config=SimpleNamespace())
+
+
+def test_runner_rejects_multi_sequence_config():
+    """Review (hsliuustc0106): the paged path writes batch index 0 only, so
+    max_num_seqs > 1 must be rejected at init."""
+    from types import SimpleNamespace
+
+    import pytest
+
+    from vllm_omni.experimental.ar_diffusion.runner import ARDiffusionModelRunner
+
+    fake = SimpleNamespace(od_config=SimpleNamespace(max_num_seqs=4))
+    with pytest.raises(ValueError, match="max_num_seqs=1"):
+        ARDiffusionModelRunner._preallocate_kv_cache(fake)

tests/ar_diffusion/test_kv_cache.py

@@ -0,0 +1,262 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Unit tests for the AR-Diffusion KV cache helpers (Phase 1, PR-2).
+
+Covers the request adapter, the chunk-window spec/manager (registration + the
+eviction policy), and the pool builder — exercised against the installed vLLM
+KV stack on CPU (block bookkeeping only, no GPU tensors).
+"""
+
+import pytest
+import torch
+from vllm.v1.kv_cache_interface import KVCacheSpecKind, get_kv_cache_spec_kind
+from vllm.v1.request import RequestStatus
+
+try:
+    from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
+except ModuleNotFoundError:
+    KVCacheSpecRegistry = None
+    from vllm.v1.core.single_type_kv_cache_manager import spec_manager_map
+
+from vllm_omni.experimental.ar_diffusion.kv_cache import (
+    ARDiffusionKVCache,
+    ARDiffusionKVConfig,
+    ARDiffusionRequestAdapter,
+    ChunkWindowManager,
+    ChunkWindowSpec,
+    allocate_kv_pool_with_views,
+    build_kv_manager,
+    compute_num_blocks,
+)
+from vllm_omni.experimental.ar_diffusion.kv_cache.paged import chunk_window_skipped_tokens
+
+BLOCK = 16
+
+
+def make_spec(*, chunk_size=BLOCK, window_chunks=2, sink_chunks=0, reset_at_boundary=False):
+    return ChunkWindowSpec(
+        block_size=BLOCK,
+        num_kv_heads=4,
+        head_size=64,
+        dtype=torch.float16,
+        sliding_window=window_chunks * chunk_size,
+        chunk_size=chunk_size,
+        window_chunks=window_chunks,
+        sink_chunks=sink_chunks,
+        reset_at_boundary=reset_at_boundary,
+    )
+
+
+# --- ChunkWindowSpec registration -------------------------------------------
+
+
+def test_spec_registration_resolves_to_chunk_window_manager():
+    # Without explicit registration the MRO walk would fall back to the parent
+    # SlidingWindowManager; assert the subclass manager wins.
+    spec = make_spec()
+    if KVCacheSpecRegistry is not None:
+        assert KVCacheSpecRegistry.get_manager_class(spec) is ChunkWindowManager
+    else:
+        assert spec_manager_map[type(spec)] is ChunkWindowManager
+
+
+def test_spec_kind_is_sliding_window():
+    assert get_kv_cache_spec_kind(make_spec()) == KVCacheSpecKind.SLIDING_WINDOW
+
+
+def test_spec_rejects_inconsistent_window():
+    with pytest.raises(ValueError):
+        ChunkWindowSpec(
+            block_size=BLOCK,
+            num_kv_heads=4,
+            head_size=64,
+            dtype=torch.float16,
+            sliding_window=99,  # != window_chunks * chunk_size
+            chunk_size=BLOCK,
+            window_chunks=2,
+        )
+
+
+# --- eviction policy (pure) -------------------------------------------------
+
+
+def test_sliding_replace_keeps_window():
+    # window = 2 chunks * 16 = 32. Base sliding formula keeps `window` tokens;
+    # the snap is to chunk boundaries.
+    def skip(n):
+        return chunk_window_skipped_tokens(n, chunk_size=16, sliding_window=32, sink_chunks=0, reset_at_boundary=False)
+
+    assert skip(32) == 0  # nothing past the window yet
+    assert skip(48) == 16  # one chunk fell out of the window
+    assert skip(64) == 32
+
+
+def test_sliding_replace_snaps_to_chunk_boundary():
+    # A non-chunk-aligned overflow must snap down so a chunk is never half-evicted.
+    skip = chunk_window_skipped_tokens(50, chunk_size=16, sliding_window=32, sink_chunks=0, reset_at_boundary=False)
+    assert skip % 16 == 0 and skip == 16
+
+
+def test_sink_chunks_protected():
+    # sink = 1 chunk (16 tokens) is never skipped.
+    skip = chunk_window_skipped_tokens(80, chunk_size=16, sliding_window=32, sink_chunks=1, reset_at_boundary=False)
+    # raw overflow = 80 - 32 - 16 = 32 -> snapped 32
+    assert skip == 32
+
+
+def test_reset_at_boundary_drops_completed_past_sink():
+    skip = chunk_window_skipped_tokens(48, chunk_size=16, sliding_window=32, sink_chunks=1, reset_at_boundary=True)
+    # completed = 48; sink = 16 -> drop 32
+    assert skip == 32
+
+
+# --- ARDiffusionKVConfig ------------------------------------------------------------
+
+
+def test_kv_config_sliding_window_property():
+    assert ARDiffusionKVConfig(chunk_size=16, window_chunks=3).sliding_window == 48
+    assert ARDiffusionKVConfig(chunk_size=16, window_chunks=None).sliding_window is None
+
+
+# --- ARDiffusionRequestAdapter ------------------------------------------------------
+
+
+def test_adapter_advances_per_chunk_not_per_step():
+    a = ARDiffusionRequestAdapter("r0", chunk_size=16)
+    assert a.num_computed_tokens == 0
+    assert a.num_tokens == 16  # in-flight chunk
+    a.on_chunk_committed()
+    assert a.num_computed_tokens == 16
+    assert a.num_tokens == 32
+
+
+def test_adapter_accounts_for_prefill_prefix():
+    a = ARDiffusionRequestAdapter("r0", chunk_size=16, prefill_prefix_tokens=4)
+    assert a.num_computed_tokens == 4
+    assert a.num_prompt_tokens == 4
+    assert a.num_tokens == 20
+
+
+def test_adapter_status_is_vllm_enum():
+    assert isinstance(ARDiffusionRequestAdapter("r0", chunk_size=16).status, RequestStatus)
+
+
+# --- pool / manager ---------------------------------------------------------
+
+
+def test_compute_num_blocks():
+    # 1 MiB budget at fraction 0.5 with 16 KiB pages -> 32 blocks.
+    assert compute_num_blocks(1 << 20, 0.5, 16 << 10) == 32
+
+
+def test_paged_pool_layout_exposes_flat_slot_views():
+    kv_pools, k_pools, v_pools = allocate_kv_pool_with_views(
+        num_blocks=4,
+        block_size=BLOCK,
+        num_layers=1,
+        num_kv_heads=4,
+        head_dim=64,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
+    )
+    assert kv_pools[0].shape == (2, 4, BLOCK, 4, 64)
+    assert k_pools[0].shape == (4 * BLOCK, 4, 64)
+    assert v_pools[0].shape == (4 * BLOCK, 4, 64)
+
+    k_pools[0][BLOCK + 3].fill_(7)
+    v_pools[0][2 * BLOCK + 5].fill_(11)
+    assert torch.equal(kv_pools[0][0, 1, 3], k_pools[0][BLOCK + 3])
+    assert torch.equal(kv_pools[0][1, 2, 5], v_pools[0][2 * BLOCK + 5])
+
+
+def test_build_manager_allocate_free_roundtrip():
+    """End-to-end: a ARDiffusionRequestAdapter drives a real KVCacheManager.
+
+    This is the adapter conformance check — if the adapter were missing an
+    attribute the manager reads, allocate_slots/free would raise here.
+    """
+    spec = make_spec()
+    mgr = build_kv_manager(spec, ["layer0"], num_blocks=16, max_model_len=1024)
+    free_before = mgr.block_pool.get_num_free_blocks()
+
+    adapter = ARDiffusionRequestAdapter("req-0", chunk_size=BLOCK)
+    blocks = mgr.allocate_slots(adapter, num_new_tokens=BLOCK, full_sequence_must_fit=True)
+    assert blocks is not None
+    assert mgr.block_pool.get_num_free_blocks() < free_before
+
+    mgr.free(adapter)
+    assert mgr.block_pool.get_num_free_blocks() == free_before
+
+
+def test_cross_attn_pool_deducted_from_self_attn_budget():
+    """The cross-attn pool is allocated directly; its bytes are subtracted from
+    the self-attn paged-pool budget so the two together stay within the free
+    memory budget (review: zwhzzz0821)."""
+    L = 512
+    avail = 1 << 30  # 1 GiB
+    kv = ARDiffusionKVCache(
+        ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=2, gpu_memory_fraction=0.5),
+        num_layers=2,
+        num_kv_heads=4,
+        head_size=64,
+        dtype=torch.float16,
+        block_size=BLOCK,
+        max_model_len=4096,
+        available_bytes=avail,
+        cross_attn_length=L,
+        device=torch.device("cpu"),
+    )
+    cross_bytes = 2 * 2 * L * 4 * 64 * torch.float16.itemsize * 2  # K+V, pos+neg, layers
+    expected = compute_num_blocks(avail - cross_bytes, 0.5, kv.spec.page_size_bytes * 2)
+    assert kv.num_blocks == expected
+    assert expected > 2 * (2 + 1) + 2  # above the min-blocks floor, so the deduction is what's tested
+
+
+def _make_kv(*, local_branches, num_frame_per_block=2, window_chunks=9):
+    return ARDiffusionKVCache(
+        ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=window_chunks),
+        num_layers=1,
+        num_kv_heads=4,
+        head_size=64,
+        dtype=torch.float32,
+        block_size=BLOCK,
+        max_model_len=4096,
+        available_bytes=1 << 16,  # tiny -> the floor binds
+        device=torch.device("cpu"),
+        local_branches=local_branches,
+        num_frame_per_block=num_frame_per_block,
+    )
+
+
+def test_pool_floor_is_branch_aware():
+    """CFG-parallel rank (one local branch) sizes for one window + in-flight chunk;
+    a single-process run (both branches) sizes for two. Scratch scales the same way."""
+    one = _make_kv(local_branches=1)
+    two = _make_kv(local_branches=2)
+
+    assert one.managed_num_blocks == 1 * (9 + 2) + 2  # 13
+    assert two.managed_num_blocks == 2 * (9 + 2) + 2  # 24
+    assert one.scratch_num_blocks == one.scratch_blocks_per_branch
+    assert two.scratch_num_blocks == 2 * two.scratch_blocks_per_branch
+    assert one.num_blocks_total == 13 + one.scratch_blocks_per_branch
+
+
+def test_scratch_maps_to_slot_zero_with_one_local_branch():
+    """A CFG-parallel rank runs exactly one branch: whichever CFG side it is,
+    its scratch lands in the rank's single slot (no dead second slot)."""
+    one = _make_kv(local_branches=1)
+    assert one.scratch_block_ids(True, 0, 2) == one.scratch_block_ids(False, 0, 2)
+
+    two = _make_kv(local_branches=2)
+    assert two.scratch_block_ids(True, 0, 2) != two.scratch_block_ids(False, 0, 2)
+
+
+def test_scratch_exhaustion_still_raises():
+    one = _make_kv(local_branches=1)
+    cap = one.scratch_blocks_per_branch
+    with pytest.raises(RuntimeError, match="scratch blocks exhausted"):
+        one.scratch_block_ids(False, 0, cap + 1)
+
+
+def test_invalid_local_branches_rejected():
+    with pytest.raises(ValueError, match="local_branches"):
+        _make_kv(local_branches=3)

tests/ar_diffusion/test_kv_orchestration.py

@@ -0,0 +1,93 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Tests for ARDiffusionKVCache — the engine-level KV orchestration body (Phase 1)."""
+
+import pytest
+import torch
+
+from vllm_omni.experimental.ar_diffusion.kv_cache import ARDiffusionKVCache, ARDiffusionKVConfig
+
+DIMS = dict(num_layers=2, num_kv_heads=4, head_size=64, dtype=torch.float16, block_size=16)
+
+
+def make_cache(*, chunk_size=16, window_chunks=2, available_bytes=1 << 24):
+    cfg = ARDiffusionKVConfig(enable=True, chunk_size=chunk_size, window_chunks=window_chunks)
+    return ARDiffusionKVCache(cfg, max_model_len=4096, available_bytes=available_bytes, **DIMS)
+
+
+def test_requires_enabled_config():
+    with pytest.raises(ValueError):
+        ARDiffusionKVCache(ARDiffusionKVConfig(enable=False), max_model_len=256, available_bytes=1 << 20, **DIMS)
+
+
+def test_requires_bounded_window():
+    cfg = ARDiffusionKVConfig(enable=True, chunk_size=16, window_chunks=None)
+    with pytest.raises(ValueError):
+        ARDiffusionKVCache(cfg, max_model_len=256, available_bytes=1 << 20, **DIMS)
+
+
+def test_full_request_lifecycle_and_eviction():
+    """begin -> per-chunk allocate/slots/commit over a long rollout -> free.
+
+    Exercises the orchestrator end-to-end and asserts the chunk window bounds
+    memory (pool plateaus) and frees cleanly.
+    """
+    kv = make_cache(chunk_size=16, window_chunks=2)
+    free_total = kv.manager.block_pool.get_num_free_blocks()
+
+    adapter = kv.begin_request("req-0")
+    free_after = []
+    for k in range(10):
+        block_table = kv.allocate_chunk(adapter)
+        slots = kv.chunk_write_slots(adapter)
+        # Slots target real blocks in this chunk's table; length == chunk_size.
+        assert len(slots) == kv.spec.chunk_size
+        used = {int(s) // kv.block_size for s in slots}
+        assert kv.null_block_id not in used
+        assert used <= set(block_table)
+        # Resident window stays bounded by window + the in-flight chunk.
+        assert len(kv.window_block_ids(adapter)) <= kv.spec.window_chunks + 1
+        free_after.append(kv.manager.block_pool.get_num_free_blocks())
+        kv.commit_chunk(adapter)
+
+    # Pool memory plateaus once the window is full (eviction recycles blocks).
+    assert free_after[-1] == free_after[kv.spec.window_chunks]
+
+    kv.end_request(adapter)
+    assert kv.manager.block_pool.get_num_free_blocks() == free_total
+
+
+def test_num_computed_advances_per_chunk():
+    kv = make_cache(chunk_size=16, window_chunks=2)
+    a = kv.begin_request("r")
+    assert a.num_computed_tokens == 0
+    kv.allocate_chunk(a)
+    kv.commit_chunk(a)
+    assert a.num_computed_tokens == 16
+    kv.end_request(a)
+
+
+def test_state_close_frees_both_branch_blocks():
+    """ARDiffusionKVState.close() returns both CFG branches' pool blocks to the pool.
+
+    This is the primitive the runner's LRU eviction relies on: when a session is
+    evicted, close() must free the blocks both adapters hold, or session churn
+    leaks pool ownership (review P1).
+    """
+    from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState
+
+    kv = make_cache(chunk_size=16, window_chunks=2)
+    free_total = kv.manager.block_pool.get_num_free_blocks()
+
+    pos = kv.begin_request("bde__s")
+    neg = kv.begin_request("bde__s__neg")
+    state = ARDiffusionKVState(kv, pos, neg, num_layers=kv.num_layers)
+    for _ in range(3):
+        for adapter in (pos, neg):
+            kv.allocate_chunk(adapter)
+            kv.commit_chunk(adapter)
+    # Both branches hold resident blocks now.
+    assert kv.manager.block_pool.get_num_free_blocks() < free_total
+
+    state.close()
+    # All blocks returned — no leak across an evicted session.
+    assert kv.manager.block_pool.get_num_free_blocks() == free_total

tests/ar_diffusion/test_paged_attention.py

@@ -0,0 +1,394 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Tests for AR-Diffusion paged self-attention contexts."""
+
+from __future__ import annotations
+
+import subprocess
+from importlib.util import find_spec
+from pathlib import Path
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+
+from vllm_omni.experimental.ar_diffusion.kv_cache import (
+    ARDiffusionKVCache,
+    ARDiffusionKVConfig,
+    ARDiffusionPagedLayerContext,
+    ARDiffusionPagedLayerInputs,
+    ar_diffusion_paged_attention,
+    paged_write_attn,
+)
+from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState
+
+BLOCK = 16
+N_HEADS = 4
+HEAD_DIM = 64
+
+
+def make_state(*, num_layers=1, window_chunks=2, dtype=torch.float32, device=torch.device("cpu")):
+    cfg = ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=window_chunks)
+    kv = ARDiffusionKVCache(
+        cfg,
+        num_layers=num_layers,
+        num_kv_heads=N_HEADS,
+        head_size=HEAD_DIM,
+        dtype=dtype,
+        block_size=BLOCK,
+        max_model_len=4096,
+        available_bytes=1 << 26,
+        device=device,
+    )
+    pos = kv.begin_request("r-pos")
+    neg = kv.begin_request("r-neg")
+    return kv, ARDiffusionKVState(kv, pos, neg, num_layers=num_layers)
+
+
+def _dense_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor:
+    scores = torch.einsum("bqhd,bkhd->bhqk", query.float(), key.float()) * (HEAD_DIM**-0.5)
+    probs = torch.softmax(scores, dim=-1).to(value.dtype)
+    return torch.einsum("bhqk,bkhd->bqhd", probs, value)
+
+
+def _cuda_flash_attn_usable() -> bool:
+    if not torch.cuda.is_available():
+        return False
+    try:
+        spec = find_spec("vllm.vllm_flash_attn")
+        if spec is None or spec.origin is None:
+            return True
+        fa2_so = Path(spec.origin).parent / "_vllm_fa2_C.abi3.so"
+        linked = subprocess.check_output(["ldd", str(fa2_so)], text=True, timeout=5)
+    except Exception:
+        return True
+    if "libcudart.so.13" not in linked:
+        return True
+    try:
+        out = subprocess.check_output(
+            ["nvidia-smi", "--query-gpu=driver_version", "--format=csv,noheader"],
+            text=True,
+            timeout=5,
+        )
+        driver_major = int(out.splitlines()[0].split(".")[0])
+    except Exception:
+        return True
+    return driver_major >= 580
+
+
+def _commit_video_span(
+    kv: ARDiffusionKVCache,
+    st: ARDiffusionKVState,
+    *,
+    is_negative: bool,
+    n_chunks: int,
+    dtype: torch.dtype,
+    device: torch.device,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    ctx = st.get_kv_caches(is_negative, seq_len=n_chunks * BLOCK, commit_current=True)[0].forward_ctx
+    ctx.ensure_video_slots(device)
+    k = torch.randn(1, n_chunks * BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    v = torch.randn(1, n_chunks * BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    kv._k_pools[0][ctx.current_video_slot_mapping] = k[0]
+    kv._v_pools[0][ctx.current_video_slot_mapping] = v[0]
+    st.commit_paged_context(is_negative)
+    return k, v
+
+
+def test_paged_context_allocates_lazily_and_commits_after_forward():
+    _, st = make_state()
+
+    contexts = st.get_kv_caches(False, seq_len=BLOCK, commit_current=True)
+    ctx = contexts[0].forward_ctx
+    assert isinstance(contexts[0], ARDiffusionPagedLayerContext)
+    assert st.pos.completed_chunks == 0
+    assert ctx.current_video_slot_mapping is None
+
+    ctx.ensure_video_slots(torch.device("cpu"))
+    assert st.pos.completed_chunks == 0
+    assert len(ctx.current_video_block_ids) == 1
+
+    st.commit_paged_context(False)
+    assert st.pos.completed_chunks == 1
+    assert st._committed[False] == BLOCK
+
+
+def test_scratch_video_and_action_blocks_do_not_commit():
+    kv, st = make_state()
+
+    ctx = st.get_kv_caches(False, seq_len=2 * BLOCK, commit_current=False)[0].forward_ctx
+    ctx.ensure_video_slots(torch.device("cpu"))
+    ctx.ensure_action_slots(3, torch.device("cpu"))
+
+    assert ctx.current_video_block_ids == kv.scratch_block_ids(False, 0, 2)
+    assert ctx.action_scratch_block_ids == kv.scratch_block_ids(False, 2, 1)
+    st.commit_paged_context(False)
+    assert st.pos.completed_chunks == 0
+    assert st._committed[False] == 0
+
+
+def test_pipeline_kv_get_paged_path_has_no_gather_backend():
+    kv, st = make_state()
+    assert not hasattr(kv, "gather_window_all_layers")
+
+    from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
+
+    pipeline = DreamZeroPipeline.__new__(DreamZeroPipeline)
+    pipeline._ar_diffusion_kv_state = st
+    contexts = pipeline._kv_get(MagicMock(), False, seq_len=BLOCK, update_kv_cache=False)
+
+    assert len(contexts) == 1
+    assert isinstance(contexts[0], ARDiffusionPagedLayerContext)
+
+
+@pytest.mark.parametrize("history_chunks", [0, 1, 3])
+@pytest.mark.parametrize("action_len", [0, 3])
+@pytest.mark.parametrize("commit_current", [False, True])
+def test_paged_attention_matches_dense_reference_cpu(history_chunks, action_len, commit_current):
+    torch.manual_seed(0)
+    device = torch.device("cpu")
+    dtype = torch.float32
+    kv, st = make_state(dtype=dtype, device=device, window_chunks=2)
+
+    history_k_parts: list[torch.Tensor] = []
+    history_v_parts: list[torch.Tensor] = []
+    if history_chunks:
+        k, v = _commit_video_span(
+            kv,
+            st,
+            is_negative=False,
+            n_chunks=history_chunks,
+            dtype=dtype,
+            device=device,
+        )
+        history_k_parts.append(k)
+        history_v_parts.append(v)
+
+    ctx = st.get_kv_caches(False, seq_len=BLOCK, commit_current=commit_current)[0].forward_ctx
+    ctx.ensure_video_slots(device)
+    current_k = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    current_v = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    kv._k_pools[0][ctx.current_video_slot_mapping] = current_k[0]
+    kv._v_pools[0][ctx.current_video_slot_mapping] = current_v[0]
+
+    action_k = action_v = None
+    if action_len:
+        ctx.ensure_action_slots(action_len, device)
+        action_k = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+        action_v = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+        kv._k_pools[0][ctx.action_slot_mapping] = action_k[0]
+        kv._v_pools[0][ctx.action_slot_mapping] = action_v[0]
+
+    query = torch.randn(1, BLOCK + action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    block_table, query_start_loc, seq_lens, max_query_len, max_seq_len = ctx.build_block_table(
+        action_len=action_len,
+        query_len=query.shape[1],
+        device=device,
+    )
+    paged = ar_diffusion_paged_attention(
+        query,
+        kv.key_cache(0),
+        kv.value_cache(0),
+        block_table=block_table,
+        query_start_loc=query_start_loc,
+        seq_lens=seq_lens,
+        max_query_len=max_query_len,
+        max_seq_len=max_seq_len,
+        softmax_scale=HEAD_DIM**-0.5,
+        causal=False,
+    )
+
+    if history_k_parts:
+        history_k = torch.cat(history_k_parts, dim=1)
+        history_v = torch.cat(history_v_parts, dim=1)
+    else:
+        history_k = torch.empty(1, 0, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+        history_v = torch.empty(1, 0, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    new_k = torch.cat([history_k, current_k], dim=1)[:, -kv.spec.sliding_window :]
+    new_v = torch.cat([history_v, current_v], dim=1)[:, -kv.spec.sliding_window :]
+    if action_len:
+        new_k = torch.cat([new_k, action_k], dim=1)
+        new_v = torch.cat([new_v, action_v], dim=1)
+    ref = _dense_attention(query, new_k, new_v)
+
+    torch.testing.assert_close(paged, ref, rtol=1e-5, atol=1e-5)
+
+    before = st.pos.completed_chunks
+    st.commit_paged_context(False)
+    assert st.pos.completed_chunks == before + (1 if commit_current else 0)
+
+
+@pytest.mark.skipif(not _cuda_flash_attn_usable(), reason="usable CUDA FlashAttention is required")
+@pytest.mark.parametrize("history_chunks", [1, 3])
+@pytest.mark.parametrize("action_len", [0, 3])
+@pytest.mark.parametrize("commit_current", [False, True])
+def test_paged_attention_matches_dense_reference_gpu(history_chunks, action_len, commit_current):
+    pytest.importorskip("vllm.vllm_flash_attn")
+    torch.manual_seed(0)
+    device = torch.device("cuda")
+    dtype = torch.float16
+    kv, st = make_state(dtype=dtype, device=device, window_chunks=2)
+
+    history_k, history_v = _commit_video_span(
+        kv,
+        st,
+        is_negative=False,
+        n_chunks=history_chunks,
+        dtype=dtype,
+        device=device,
+    )
+
+    layer_ctx = st.get_kv_caches(False, seq_len=BLOCK, commit_current=commit_current)[0]
+    ctx = layer_ctx.forward_ctx
+    current_k = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    current_v = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    action_k = action_v = None
+    if action_len:
+        action_k = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+        action_v = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+
+    query = torch.randn(1, BLOCK + action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+
+    # The production path: once-per-forward host prep, then the fused
+    # write+attend custom op consuming the NamedTuple payload.
+    ctx.prepare(device=device, action_len=action_len, query_len=query.shape[1])
+    inputs = layer_ctx.to_layer_inputs()
+    assert isinstance(inputs, ARDiffusionPagedLayerInputs)
+    paged = paged_write_attn(
+        inputs,
+        query[0],
+        current_k[0],
+        current_v[0],
+        action_k[0] if action_len else None,
+        action_v[0] if action_len else None,
+        HEAD_DIM**-0.5,
+    ).unsqueeze(0)
+
+    # Direct python-fn call on the same (already written) pools must be
+    # bit-exact: identical kernel, identical inputs.
+    direct = ar_diffusion_paged_attention(
+        query,
+        kv.key_cache(0),
+        kv.value_cache(0),
+        block_table=ctx.block_table,
+        query_start_loc=ctx.query_start_loc,
+        seq_lens=ctx.seq_lens,
+        max_query_len=ctx.max_query_len,
+        max_seq_len=ctx.max_seq_len,
+        softmax_scale=HEAD_DIM**-0.5,
+        causal=False,
+    )
+    assert torch.equal(paged, direct)
+
+    new_k = torch.cat([history_k, current_k], dim=1)[:, -kv.spec.sliding_window :]
+    new_v = torch.cat([history_v, current_v], dim=1)[:, -kv.spec.sliding_window :]
+    if action_len:
+        new_k = torch.cat([new_k, action_k], dim=1)
+        new_v = torch.cat([new_v, action_v], dim=1)
+    ref = _dense_attention(query, new_k, new_v)
+
+    torch.testing.assert_close(paged, ref, rtol=2e-2, atol=2e-2)
+
+
+def test_block_table_padded_to_fixed_width():
+    """Shapes must be constant across window growth: only values change."""
+    device = torch.device("cpu")
+    kv, st = make_state(window_chunks=2)
+
+    ctx1 = st.get_kv_caches(False, seq_len=BLOCK, commit_current=True)[0].forward_ctx
+    ctx1.prepare(device=device, action_len=0, query_len=BLOCK)
+    st.commit_paged_context(False)
+
+    ctx2 = st.get_kv_caches(False, seq_len=BLOCK, commit_current=True)[0].forward_ctx
+    ctx2.prepare(device=device, action_len=0, query_len=BLOCK)
+    st.commit_paged_context(False)
+
+    # 1-block vs 2-block visible history: same table width, same max_seq_len.
+    assert ctx1.block_table.shape == ctx2.block_table.shape
+    assert ctx1.max_seq_len == ctx2.max_seq_len
+    expected_width = kv.spec.sliding_window // kv.block_size + 1
+    assert ctx1.block_table.shape == (1, expected_width)
+    # Real lengths live in seq_lens, not the padded table.
+    assert int(ctx1.seq_lens[0]) == BLOCK
+    assert int(ctx2.seq_lens[0]) == 2 * BLOCK
+
+
+def test_prepare_is_idempotent_and_layers_share_metadata():
+    device = torch.device("cpu")
+    _, st = make_state(num_layers=2)
+    contexts = st.get_kv_caches(False, seq_len=BLOCK, commit_current=False)
+    fctx = contexts[0].forward_ctx
+    fctx.prepare(device=device, action_len=0, query_len=BLOCK)
+    table = fctx.block_table
+    fctx.prepare(device=device, action_len=0, query_len=BLOCK)
+    assert fctx.block_table is table  # memoized, not rebuilt
+
+    i0, i1 = contexts[0].to_layer_inputs(), contexts[1].to_layer_inputs()
+    # 0-dim tensors (NOT python ints) so dynamo doesn't install per-layer
+    # value guards on the shared block code object.
+    assert isinstance(i0.layer_idx, torch.Tensor) and int(i0.layer_idx) == 0
+    assert isinstance(i1.layer_idx, torch.Tensor) and int(i1.layer_idx) == 1
+    # All layers share the same metadata tensor objects.
+    assert i0.block_table is i1.block_table
+    assert i0.seq_lens is i1.seq_lens
+    assert i0.video_slots is i1.video_slots
+
+
+def test_layer_inputs_before_prepare_raises():
+    _, st = make_state()
+    layer_ctx = st.get_kv_caches(False, seq_len=BLOCK, commit_current=False)[0]
+    with pytest.raises(RuntimeError, match="before prepare"):
+        layer_ctx.to_layer_inputs()
+
+
+def test_custom_op_registration_idempotent():
+    import importlib
+    import sys
+
+    assert hasattr(torch.ops.vllm_omni, "ar_diffusion_paged_write_attn")
+    mod = "vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention"
+    saved = sys.modules.pop(mod)
+    try:
+        importlib.import_module(mod)  # re-registration must not raise
+    finally:
+        sys.modules[mod] = saved
+    assert hasattr(torch.ops.vllm_omni, "ar_diffusion_paged_write_attn")
+
+
+def test_custom_op_compiles_fullgraph_without_recompile_on_value_change():
+    """The op must trace as one opaque node: fullgraph OK, and changed tensor
+    VALUES (new slots / block ids) must not trigger recompilation."""
+    import torch._dynamo
+
+    device = torch.device("cpu")
+    kv, st = make_state(num_layers=2, window_chunks=2)
+
+    def run_one_forward(commit):
+        contexts = st.get_kv_caches(False, seq_len=BLOCK, commit_current=commit)
+        fctx = contexts[0].forward_ctx
+        fctx.prepare(device=device, action_len=0, query_len=BLOCK)
+        q = torch.randn(BLOCK, N_HEADS, HEAD_DIM)
+        k = torch.randn(BLOCK, N_HEADS, HEAD_DIM)
+        v = torch.randn(BLOCK, N_HEADS, HEAD_DIM)
+        # Both layers through ONE compiled fn: layer_idx is a tensor, so a
+        # different layer must NOT recompile (all 40 DiT blocks share the
+        # block-forward code object in production).
+        for layer_ctx in contexts:
+            out = compiled(layer_ctx.to_layer_inputs(), q, k, v)
+        st.commit_paged_context(False)
+        return out
+
+    torch._dynamo.reset()
+    from torch._dynamo.testing import CompileCounter
+
+    counter = CompileCounter()
+
+    def fn(inputs, q, k, v):
+        return paged_write_attn(inputs, q, k, v, None, None, HEAD_DIM**-0.5) * 1.0
+
+    compiled = torch.compile(fn, backend=counter, fullgraph=True)
+
+    run_one_forward(commit=True)  # history grows between calls ->
+    run_one_forward(commit=True)  # block-table VALUES change, shapes don't
+    run_one_forward(commit=False)
+
+    assert counter.frame_count == 1, f"recompiled: frame_count={counter.frame_count}"

tests/ar_diffusion/test_pipeline_bridge.py

@@ -0,0 +1,222 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Tests for the ARDiffusionKVState paged-attention pipeline bridge."""
+
+import pytest
+import torch
+
+from vllm_omni.experimental.ar_diffusion.kv_cache import (
+    ARDiffusionKVCache,
+    ARDiffusionKVConfig,
+    ARDiffusionPagedLayerContext,
+)
+from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState
+
+BLOCK = 16
+N_HEADS = 4
+HEAD_DIM = 64
+
+
+def make_state(num_layers=1, window_chunks=4, cross_attn_length=0):
+    cfg = ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=window_chunks)
+    kv = ARDiffusionKVCache(
+        cfg,
+        num_layers=num_layers,
+        num_kv_heads=N_HEADS,
+        head_size=HEAD_DIM,
+        dtype=torch.float32,
+        block_size=BLOCK,
+        max_model_len=4096,
+        available_bytes=1 << 24,
+        cross_attn_length=cross_attn_length,
+        device=torch.device("cpu"),
+    )
+    pos = kv.begin_request("r-pos")
+    neg = kv.begin_request("r-neg")
+    return kv, ARDiffusionKVState(kv, pos, neg, num_layers=num_layers)
+
+
+def _prepare_and_commit(st: ARDiffusionKVState, is_negative: bool, n_chunks: int) -> None:
+    ctx = st.get_kv_caches(is_negative, seq_len=n_chunks * BLOCK, commit_current=True)[0].forward_ctx
+    ctx.ensure_video_slots(torch.device("cpu"))
+    st.commit_paged_context(is_negative)
+
+
+def test_get_returns_paged_layer_contexts_when_nothing_committed():
+    _, st = make_state(num_layers=2)
+    for neg in (False, True):
+        contexts = st.get_kv_caches(neg, seq_len=BLOCK, commit_current=False)
+        assert len(contexts) == 2
+        assert all(isinstance(ctx, ARDiffusionPagedLayerContext) for ctx in contexts)
+        assert contexts[0].history_block_ids == []
+
+
+def test_managed_paged_context_commits_only_after_forward():
+    kv, st = make_state(num_layers=1)
+    contexts = st.get_kv_caches(False, seq_len=2 * BLOCK, commit_current=True)
+    ctx = contexts[0].forward_ctx
+
+    assert st.pos.completed_chunks == 0
+    assert kv.window_block_ids(st.pos) == []
+    ctx.ensure_video_slots(torch.device("cpu"))
+    assert st.pos.completed_chunks == 0
+    assert len(ctx.current_video_block_ids) == 2
+
+    st.commit_paged_context(False)
+
+    assert st.pos.completed_chunks == 2
+    assert st._committed[False] == 2 * BLOCK
+    assert len(kv.window_block_ids(st.pos)) == 2
+
+
+def test_scratch_paged_context_does_not_commit():
+    kv, st = make_state(num_layers=1)
+    contexts = st.get_kv_caches(False, seq_len=BLOCK, commit_current=False)
+    ctx = contexts[0].forward_ctx
+    ctx.ensure_video_slots(torch.device("cpu"))
+
+    assert ctx.current_video_block_ids == kv.scratch_block_ids(False, 0, 1)
+    st.commit_paged_context(False)
+
+    assert st.pos.completed_chunks == 0
+    assert st._committed[False] == 0
+    assert kv.window_block_ids(st.pos) == []
+
+
+def test_eviction_bounds_resident_window():
+    kv, st = make_state(num_layers=1, window_chunks=3)
+    _prepare_and_commit(st, False, 3)
+    assert len(kv.window_block_ids(st.pos)) == 3
+
+    _prepare_and_commit(st, False, 2)
+    ctx = st.get_kv_caches(False, seq_len=BLOCK, commit_current=False)[0].forward_ctx
+    visible_blocks, video_len = ctx.video_block_table(torch.device("cpu"))
+    assert len(visible_blocks) == 3
+    assert video_len == 3 * BLOCK
+
+
+def test_branches_are_independent():
+    kv, st = make_state()
+    _prepare_and_commit(st, False, 1)
+    assert len(kv.window_block_ids(st.pos)) == 1
+    assert kv.window_block_ids(st.neg) == []
+
+
+def test_reset_clears_session_window():
+    kv, st = make_state(num_layers=1, window_chunks=4)
+    _prepare_and_commit(st, False, 2)
+    _prepare_and_commit(st, True, 1)
+    assert st._committed[False] == 2 * BLOCK and st._committed[True] == BLOCK
+    free_before = kv.manager.block_pool.get_num_free_blocks()
+
+    st.reset()
+
+    assert st._committed == {False: 0, True: 0}
+    assert st._paged_pending == {False: None, True: None}
+    assert kv.window_block_ids(st.pos) == []
+    assert kv.window_block_ids(st.neg) == []
+    assert kv.manager.block_pool.get_num_free_blocks() > free_before
+    _prepare_and_commit(st, False, 1)
+    assert len(kv.window_block_ids(st.pos)) == 1
+
+
+def test_get_cross_kv_caches_raises_before_populate():
+    """Engine ownership guard: a cross-attn read before _kv_populate_cross must fail loud."""
+    _, st = make_state()  # cross_attn_length == 0, nothing populated
+    for neg in (False, True):
+        with pytest.raises(RuntimeError, match="cross-attn read before"):
+            st.get_cross_kv_caches(neg)
+    # Guards the full AND: the populated flag alone is not enough without a pool.
+    st._cross_text_populated[False] = True
+    with pytest.raises(RuntimeError, match="cross-attn read before"):
+        st.get_cross_kv_caches(False)
+
+
+def test_get_cross_kv_caches_returns_pool_dicts_when_populated():
+    L = 8
+    kv, st = make_state(num_layers=2, cross_attn_length=L)
+    written = []
+    for i in range(2):
+        k = torch.randn(1, L, N_HEADS, HEAD_DIM)
+        v = torch.randn(1, L, N_HEADS, HEAD_DIM)
+        kv.write_cross_kv(i, False, k, v)
+        written.append((k, v))
+    st._cross_text_populated[False] = True
+
+    out = st.get_cross_kv_caches(False)
+    assert len(out) == 2
+    for i, (k, v) in enumerate(written):
+        assert out[i]["is_init"] is True
+        assert out[i]["k"].shape == (1, L, N_HEADS, HEAD_DIM)
+        assert torch.equal(out[i]["k"], k)
+        assert torch.equal(out[i]["v"], v)
+
+
+def test_kv_create_is_noop_engine_owns_allocation():
+    from unittest.mock import MagicMock
+
+    from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
+
+    p = DreamZeroPipeline.__new__(DreamZeroPipeline)
+    p._ar_diffusion_kv_state = object()
+    state = MagicMock()
+    p._kv_create(state, 1, "float32", "cpu", 24, 4, 64)
+    state.create_kv_caches.assert_not_called()
+
+
+def test_get_kv_cache_requires_frame_aligned_seqlen():
+    _, st = make_state()
+    with pytest.raises(AssertionError, match="frame-aligned"):
+        st.get_kv_caches(False, seq_len=BLOCK + 1, commit_current=True)
+
+
+def test_prepare_one_branch_allocates_nothing_for_the_other():
+    """CFG-parallel laziness: a rank prepares only the branch it runs."""
+    kv, st = make_state(num_layers=1, window_chunks=4)
+    pos_ctx = st.get_kv_caches(False, seq_len=BLOCK, commit_current=True)[0].forward_ctx
+    neg_ctx = st.get_kv_caches(True, seq_len=BLOCK, commit_current=True)[0].forward_ctx
+
+    pos_ctx.prepare(device=torch.device("cpu"), action_len=0, query_len=BLOCK)
+
+    assert pos_ctx._allocated_video
+    assert not neg_ctx._allocated_video
+    assert kv.window_block_ids(st.neg) == []
+
+
+def test_failed_forward_tears_down_session(monkeypatch):
+    """Review follow-up (hsliuustc0106): a forward that dies partway must not
+    leave allocated-but-uncommitted KV behind — the session is torn down (pool
+    blocks freed, model-local state dropped) and the error propagates."""
+    from types import SimpleNamespace
+
+    from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
+    from vllm_omni.experimental.ar_diffusion.runner import ARDiffusionModelRunner
+
+    kv, st = make_state(num_layers=1, window_chunks=4)
+    _prepare_and_commit(st, False, 2)  # session owns pool blocks
+    free_before_failure = kv.manager.block_pool.get_num_free_blocks()
+
+    runner = object.__new__(ARDiffusionModelRunner)
+    runner.kv_cache = kv
+    runner._ar_diffusion_states = __import__("collections").OrderedDict({"s1": st})
+    runner._max_ar_diffusion_states = 4
+    runner.pipeline = SimpleNamespace(_states={"s1": object()}, _ar_diffusion_kv_state=None)
+    runner.device = None
+    runner._perf_e2e_times = []
+
+    def boom(self, req):
+        raise RuntimeError("layer 17 exploded")
+
+    monkeypatch.setattr(DiffusionModelRunner, "execute_model", boom)
+    req = SimpleNamespace(
+        request_id="r0",
+        sampling_params=SimpleNamespace(extra_args={"session_id": "s1"}),
+    )
+
+    with pytest.raises(RuntimeError, match="layer 17 exploded"):
+        ARDiffusionModelRunner.execute_model(runner, req)
+
+    assert "s1" not in runner._ar_diffusion_states
+    assert "s1" not in runner.pipeline._states
+    assert runner.pipeline._ar_diffusion_kv_state is None
+    # close() freed the session's resident blocks.
+    assert kv.manager.block_pool.get_num_free_blocks() > free_before_failure

tests/ar_diffusion/test_slot_mapping.py

@@ -0,0 +1,76 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Slot-mapping / block-table tests for the AR-Diffusion engine (Phase 1, Step 3)."""
+
+import pytest
+import torch
+
+from vllm_omni.experimental.ar_diffusion.kv_cache import (
+    ARDiffusionRequestAdapter,
+    ChunkWindowSpec,
+    build_kv_manager,
+    chunk_slot_mapping,
+    compute_slot_mapping,
+    resident_block_ids,
+)
+
+BLOCK = 16
+
+
+def make_spec(*, window_chunks=2):
+    return ChunkWindowSpec(
+        block_size=BLOCK,
+        num_kv_heads=4,
+        head_size=64,
+        dtype=torch.float16,
+        sliding_window=window_chunks * BLOCK,
+        chunk_size=BLOCK,
+        window_chunks=window_chunks,
+    )
+
+
+def test_slot_mapping_matches_block_offsets():
+    # block table [5, 2, 9], block_size 4. slot = block*4 + (pos % 4).
+    slots = compute_slot_mapping([5, 2, 9], [0, 1, 4, 5, 8, 9], block_size=4)
+    assert slots.tolist() == [20, 21, 8, 9, 36, 37]
+
+
+def test_slot_mapping_rejects_bad_block_size():
+    with pytest.raises(ValueError):
+        compute_slot_mapping([1], [0], block_size=0)
+
+
+def test_chunk_slot_mapping_targets_current_chunk():
+    # After 2 committed chunks (num_computed=32), the in-flight chunk maps to the
+    # block covering positions 32..47 — i.e. block table index 2.
+    block_ids = [7, 4, 3]  # block index 2 -> physical block 3
+    slots = chunk_slot_mapping(block_ids, num_computed_tokens=32, chunk_size=4, block_size=16)
+    # positions 32..35 -> block_index 2 -> physical 3 -> slots 48..51
+    assert slots.tolist() == [48, 49, 50, 51]
+
+
+def test_resident_block_ids_excludes_null():
+    assert resident_block_ids([0, 1, 0, 2, 3], null_block_id=0) == [1, 2, 3]
+
+
+def test_blocktable_build_single_request():
+    """Integration: a real KVCacheManager block table -> a valid slot mapping.
+
+    Every slot for the in-flight chunk must land inside a real (non-null) block
+    that the manager actually allocated for the request.
+    """
+    spec = make_spec()
+    mgr = build_kv_manager(spec, ["l0"], num_blocks=16, max_model_len=1024)
+    null_id = mgr.block_pool.null_block.block_id
+    adapter = ARDiffusionRequestAdapter("req", chunk_size=BLOCK)
+
+    mgr.allocate_slots(adapter, num_new_tokens=BLOCK)
+    block_ids = mgr.get_block_ids(adapter.request_id)[0]
+    real = set(resident_block_ids(block_ids, null_id))
+    assert real, "expected at least one real block after allocation"
+
+    slots = chunk_slot_mapping(block_ids, adapter.num_computed_tokens, spec.chunk_size, spec.block_size)
+    # Each slot resolves to a real allocated block, none to null_block.
+    blocks_used = {int(s) // spec.block_size for s in slots}
+    assert blocks_used <= real
+    assert null_id not in blocks_used
+    assert len(slots) == spec.chunk_size

tests/config/test_endpoint_policy.py

@@ -0,0 +1,87 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for endpoint restrictions logic."""
+
+import pytest
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
+
+from vllm_omni.config.endpoint_policy import (
+    EndpointRestriction,
+    OmniServingCapability,
+    shutdown_unsupported_routes,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+SERVER_ROUTE = "/v1/completions"
+REJECTION_REASON = "This route is banned!"
+
+
+def _make_app_with_server_route():
+    """Create a FastAPI app with a route to reject."""
+    app = FastAPI()
+
+    @app.post(SERVER_ROUTE)
+    async def existing_handler():
+        """Returns a response; we want to disable this."""
+        return {"ok": True}
+
+    return app
+
+
+def test_restricted_completions_returns_400():
+    """Ensure that simple route rejection works."""
+    app = _make_app_with_server_route()
+    restrictions = (EndpointRestriction(OmniServingCapability.COMPLETIONS, REJECTION_REASON),)
+    shutdown_unsupported_routes(app, restrictions)
+
+    client = TestClient(app)
+    resp = client.post(SERVER_ROUTE, json={"prompt": "hello", "model": "m"})
+    assert resp.status_code == 400
+    body = resp.json()
+    assert body["error"]["message"] == REJECTION_REASON
+    assert body["error"]["type"] == "BadRequestError"
+
+
+def test_unrestricted_completions_not_blocked():
+    """Ensure that if we don't shutdown any routes, the route stays."""
+    app = _make_app_with_server_route()
+    shutdown_unsupported_routes(app, ())
+
+    client = TestClient(app)
+    resp = client.post(SERVER_ROUTE)
+    assert resp.status_code == 200
+    assert resp.json() == {"ok": True}
+
+
+def test_route_rejection_is_idempotent():
+    """Ensure that calling rejection twice doesn't break things."""
+    app = _make_app_with_server_route()
+    restrictions = (EndpointRestriction(OmniServingCapability.COMPLETIONS, REJECTION_REASON),)
+    shutdown_unsupported_routes(app, restrictions)
+    shutdown_unsupported_routes(app, restrictions)
+
+    client = TestClient(app)
+    resp = client.post(SERVER_ROUTE, json={"prompt": "hello", "model": "m"})
+    assert resp.status_code == 400
+    body = resp.json()
+    assert body["error"]["message"] == REJECTION_REASON
+    assert body["error"]["type"] == "BadRequestError"
+
+
+def test_invalid_route_rejection():
+    """Ensure that a restriction won't crash if a route doesn't exist."""
+    # NOTE: This test is for safety since the route organization is still messy.
+    # Since the restriction needs a valid OmniServingCapability, its implied the route
+    # is a valid vLLM route, but was already removed from the app by earlier processing.
+    app = FastAPI()
+    restrictions = (EndpointRestriction(OmniServingCapability.COMPLETIONS, REJECTION_REASON),)
+    shutdown_unsupported_routes(app, restrictions)
+
+    client = TestClient(app)
+    resp = client.post(SERVER_ROUTE, json={"prompt": "hello", "model": "m"})
+    assert resp.status_code == 400
+    body = resp.json()
+    assert body["error"]["message"] == REJECTION_REASON
+    assert body["error"]["type"] == "BadRequestError"

tests/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py

@@ -6,6 +6,7 @@ import pytest
 from vllm.v1.request import RequestStatus
 
 import vllm_omni.core.sched.omni_ar_scheduler as scheduler_mod
+import vllm_omni.model_executor.models.voxcpm2.scheduler as voxcpm2_scheduler_mod
 from vllm_omni.model_executor.models.voxcpm2.scheduler import VoxCPM2OmniARAsyncScheduler
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -51,14 +52,30 @@ class _MockRequest:
         return RequestStatus.is_finished(self.status)
 
 
-def _make_scheduler(*, enable_unified_decode_graph: bool = True) -> VoxCPM2OmniARAsyncScheduler:
+@pytest.fixture(autouse=True)
+def _mock_cuda_graph_platform(monkeypatch) -> None:
+    monkeypatch.setattr(voxcpm2_scheduler_mod.current_omni_platform, "is_cuda", lambda: True)
+
+
+def _make_scheduler(
+    *,
+    enable_unified_decode_graph: bool | None = True,
+    deterministic_cfm_noise: bool = False,
+    runtime_config_on_model_config: bool = False,
+) -> VoxCPM2OmniARAsyncScheduler:
     sched = VoxCPM2OmniARAsyncScheduler.__new__(VoxCPM2OmniARAsyncScheduler)
-    runtime_config = SimpleNamespace(enable_unified_decode_graph=enable_unified_decode_graph)
-    sched.vllm_config = SimpleNamespace(
-        model_config=SimpleNamespace(
-            hf_config=SimpleNamespace(voxcpm2_runtime_config=runtime_config),
+    hf_config = SimpleNamespace()
+    model_config = SimpleNamespace(hf_config=hf_config)
+    if enable_unified_decode_graph is not None or deterministic_cfm_noise:
+        runtime_config = SimpleNamespace(
+            enable_unified_decode_graph=True if enable_unified_decode_graph is None else enable_unified_decode_graph,
+            deterministic_cfm_noise=deterministic_cfm_noise,
         )
-    )
+        if runtime_config_on_model_config:
+            model_config.voxcpm2_runtime_config = runtime_config
+        else:
+            hf_config.voxcpm2_runtime_config = runtime_config
+    sched.vllm_config = SimpleNamespace(model_config=model_config)
     return sched
 
 
@@ -86,6 +103,39 @@ def test_voxcpm2_unified_decode_graph_does_not_defer_when_disabled() -> None:
     assert not scheduler._should_defer_waiting_for_unified_decode_graph()
 
 
+def test_voxcpm2_unified_decode_graph_uses_model_runtime_defaults() -> None:
+    scheduler = _make_scheduler(enable_unified_decode_graph=None)
+    scheduler.running = [_MockRequest("decode")]
+    scheduler.waiting = _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)])
+
+    assert scheduler._should_defer_waiting_for_unified_decode_graph()
+
+
+def test_voxcpm2_unified_decode_graph_reads_model_config_runtime_config() -> None:
+    scheduler = _make_scheduler(enable_unified_decode_graph=True, runtime_config_on_model_config=True)
+    scheduler.running = [_MockRequest("decode")]
+    scheduler.waiting = _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)])
+
+    assert scheduler._should_defer_waiting_for_unified_decode_graph()
+
+
+def test_voxcpm2_unified_decode_graph_does_not_defer_with_deterministic_noise() -> None:
+    scheduler = _make_scheduler(deterministic_cfm_noise=True)
+    scheduler.running = [_MockRequest("decode")]
+    scheduler.waiting = _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)])
+
+    assert not scheduler._should_defer_waiting_for_unified_decode_graph()
+
+
+def test_voxcpm2_unified_decode_graph_does_not_defer_without_cuda_graph(monkeypatch) -> None:
+    monkeypatch.setattr(voxcpm2_scheduler_mod.current_omni_platform, "is_cuda", lambda: False)
+    scheduler = _make_scheduler()
+    scheduler.running = [_MockRequest("decode")]
+    scheduler.waiting = _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)])
+
+    assert not scheduler._should_defer_waiting_for_unified_decode_graph()
+
+
 def test_unified_decode_graph_deferral_restores_waiting_queue(monkeypatch) -> None:
     scheduler = _make_scheduler()
     scheduler.running = [_MockRequest("decode")]

tests/dfx/perf/tests/test_lingbot_video_vllm_omni.json

@@ -0,0 +1,34 @@
+[
+    {
+        "test_name": "test_lingbot_video_single_device",
+        "description": "Single-device LingBot dense T2V baseline at 320x192, 9 frames, 2 steps.",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "robbyant/lingbot-video-dense-1.3b",
+            "serve_args": {
+                "model-class-name": "LingBotVideoPipeline",
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "320x192_frames9_steps2",
+                "dataset": "random",
+                "task": "t2v",
+                "backend": "v1/videos",
+                "width": 320,
+                "height": 192,
+                "num-frames": 9,
+                "fps": 24,
+                "num-inference-steps": 2,
+                "num-prompts": 3,
+                "max-concurrency": 1,
+                "warmup-requests": 1,
+                "warmup-concurrency": 1,
+                "warmup-num-inference-steps": 2,
+                "enable-negative-prompt": true,
+                "skip-performance-assertion": true
+            }
+        ]
+    }
+]

tests/diffusion/attention/test_kernels_hub.py

@@ -0,0 +1,104 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import sys
+import types
+
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum
+from vllm_omni.diffusion.attention.backends.sdpa import SDPAImpl
+from vllm_omni.platforms import current_omni_platform
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_kernels_hub_platform_fallback(monkeypatch: pytest.MonkeyPatch):
+    """Test that when kernels package is not available, the platform fallback logic
+
+    routes to local FLASH_ATTN.
+    """
+    from vllm_omni.platforms.cuda.platform import CudaOmniPlatform
+
+    # Temporarily hide/remove kernels module if present
+    monkeypatch.setitem(sys.modules, "kernels", None)
+
+    # Use monkeypatch to mock CudaOmniPlatform capability and package checker to allow FLASH_ATTN
+    from vllm.platforms.interface import DeviceCapability
+
+    from vllm_omni.diffusion.envs import PACKAGES_CHECKER
+
+    monkeypatch.setattr(
+        CudaOmniPlatform,
+        "get_device_capability",
+        classmethod(lambda cls, device_id=0: DeviceCapability(8, 0)),
+    )
+    monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": True})
+
+    # Test FLASH_ATTN_HUB falls back to FLASH_ATTN when kernels is unavailable
+    backend_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_HUB", head_size=64)
+    assert backend_path == DiffusionAttentionBackendEnum.FLASH_ATTN.get_path()
+
+    # Test FLASH_ATTN_3_HUB falls back to FLASH_ATTN when kernels is unavailable
+    backend_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_3_HUB", head_size=64)
+    assert backend_path == DiffusionAttentionBackendEnum.FLASH_ATTN.get_path()
+
+    # Test FLASH_ATTN_3_HUB falls back to FLASH_ATTN_HUB on pre-Hopper GPUs
+    kernels_module = types.ModuleType("kernels")
+    monkeypatch.setitem(sys.modules, "kernels", kernels_module)
+    backend_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_3_HUB", head_size=64)
+    assert backend_path == DiffusionAttentionBackendEnum.FLASH_ATTN_HUB.get_path()
+
+
+@pytest.mark.core_model
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_kernels_hub_execution():
+    """Verify basic forward of flash_attn_hub and flash_attn_3_hub, comparing with SDPA reference."""
+    device = torch.device(current_omni_platform.device_type)
+    dtype = torch.bfloat16
+
+    num_heads = 8
+    head_dim = 64
+    seq_len = 32
+    batch_size = 1
+
+    torch.manual_seed(42)
+    q = torch.randn(batch_size, seq_len, num_heads, head_dim, device=device, dtype=dtype)
+    k = q.clone()
+    v = q.clone()
+
+    # 1. Test PyTorch SDPA reference
+    sdpa_impl = SDPAImpl(num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False)
+    attn_metadata_sdpa = AttentionMetadata(attn_mask=None)
+    output_ref = sdpa_impl.forward(q.clone(), k.clone(), v.clone(), attn_metadata_sdpa)
+
+    # 2. Test FlashAttentionHubBackend (FlashAttention 2)
+    from vllm_omni.diffusion.attention.backends.flash_attn_hub import FlashAttentionHubImpl
+
+    fa_hub_impl = FlashAttentionHubImpl(
+        num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False
+    )
+    output_fa_hub = fa_hub_impl.forward(q.clone(), k.clone(), v.clone(), attn_metadata_sdpa)
+    assert output_fa_hub.shape == q.shape
+    assert not torch.isnan(output_fa_hub).any()
+    max_diff = torch.max(torch.abs(output_ref - output_fa_hub)).item()
+    assert max_diff < 1e-2, f"FlashAttentionHub output differs too much from SDPA reference: {max_diff}"
+
+    # 3. Test FlashAttention3HubBackend (FlashAttention 3, Hopper+ only)
+    major, _minor = torch.cuda.get_device_capability()
+    if major < 9:
+        pytest.skip("FLASH_ATTN_3_HUB execution requires Hopper-class GPU (compute capability >= 9.0)")
+
+    from vllm_omni.diffusion.attention.backends.flash_attn_hub import FlashAttention3HubImpl
+
+    fa3_hub_impl = FlashAttention3HubImpl(
+        num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False
+    )
+    output_fa3_hub = fa3_hub_impl.forward(q.clone(), k.clone(), v.clone(), attn_metadata_sdpa)
+    assert output_fa3_hub.shape == q.shape
+    assert not torch.isnan(output_fa3_hub).any()
+    max_diff = torch.max(torch.abs(output_ref - output_fa3_hub)).item()
+    assert max_diff < 1e-2, f"FlashAttention3Hub output differs too much from SDPA reference: {max_diff}"

tests/diffusion/distributed/test_expert_parallel_layout.py

@@ -19,6 +19,10 @@ class _FakeGroup:
         self.group_ranks = group_ranks
         self.parallel_mode = parallel_mode
         self.device_group = object()
+        self.device_communicator = kwargs.get("device_communicator")
+        reduce_scatter = kwargs.get("reduce_scatter")
+        if reduce_scatter is not None:
+            self.reduce_scatter = reduce_scatter
         self.ulysses_group = kwargs.get("ulysses_group")
         self.ring_group = kwargs.get("ring_group")
         self.local_group = next(group for group in group_ranks if local_rank in group)
@@ -55,6 +59,23 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
         created_groups.append(group)
         return group
 
+    def fake_init_vllm_model_parallel_group(
+        group_ranks,
+        local_rank,
+        backend,
+        group_name,
+    ):
+        del backend
+        group = _FakeGroup(
+            [list(ranks) for ranks in group_ranks],
+            local_rank,
+            f"vllm_{group_name}",
+            device_communicator=object(),
+            reduce_scatter=lambda tensor, **kwargs: tensor,
+        )
+        created_groups.append(group)
+        return group
+
     fake_world_group = SimpleNamespace(
         rank_in_group=local_rank,
         local_rank=local_rank,
@@ -68,6 +89,7 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
     monkeypatch.setattr(parallel_state, "get_world_group", lambda: fake_world_group)
     monkeypatch.setattr(parallel_state, "get_forward_context", lambda: fake_forward_context)
     monkeypatch.setattr(parallel_state, "init_model_parallel_group", fake_init_model_parallel_group)
+    monkeypatch.setattr(parallel_state, "init_vllm_model_parallel_group", fake_init_vllm_model_parallel_group)
     monkeypatch.setattr(parallel_state, "init_dit_group", lambda *_args, **_kwargs: None)
 
     for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_EXPERT_PARALLEL_GROUP_RANKS"):
@@ -87,7 +109,7 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
         backend="gloo",
     )
 
-    assert parallel_state.vllm_parallel_state._PCP is parallel_state._SP
+    assert parallel_state.vllm_parallel_state._PCP is not parallel_state._SP
     assert parallel_state.vllm_parallel_state._PCP.world_size == 2
     assert parallel_state._DP.world_size == 2
     assert parallel_state.vllm_parallel_state._DP is not parallel_state._DP
@@ -96,6 +118,13 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
     assert parallel_state.vllm_parallel_state._TP.world_size == 2
     assert parallel_state._PP.world_size == 2
 
+    assert parallel_state.vllm_parallel_state._PCP.device_communicator is not None
+    assert parallel_state.vllm_parallel_state._DP.device_communicator is not None
+    assert parallel_state.vllm_parallel_state._EP.device_communicator is not None
+    assert hasattr(parallel_state.vllm_parallel_state._PCP, "reduce_scatter")
+    assert hasattr(parallel_state.vllm_parallel_state._DP, "reduce_scatter")
+    assert hasattr(parallel_state.vllm_parallel_state._EP, "reduce_scatter")
+
     assert parallel_state.vllm_parallel_state._PCP.local_group == [0, 2]
     assert parallel_state.vllm_parallel_state._DP.local_group == [0, 8, 16, 24]
     assert parallel_state.vllm_parallel_state._EP.local_group == [
@@ -155,7 +184,9 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
         ],
     ]
 
-    ep_groups = [group.local_group for group in created_groups if group.parallel_mode == "expert"]
+    vllm_group_names = [group.parallel_mode for group in created_groups if group.parallel_mode.startswith("vllm_")]
+    assert vllm_group_names == ["vllm_pcp", "vllm_tp", "vllm_dp", "vllm_ep"]
+    ep_groups = [group.local_group for group in created_groups if group.parallel_mode == "vllm_ep"]
     assert ep_groups == [parallel_state.vllm_parallel_state._EP.local_group]
 
 

tests/diffusion/layers/test_rope_broadcast.py

@@ -0,0 +1,236 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Shape broadcasting regression tests for RotaryEmbedding.forward_native.
+
+Covers the batch-dimension shape mismatch fix: forward_native now strips the
+batch dim from 3D cos/sin before applying RoPE, consistent with forward_cuda,
+forward_hip, and apply_rotary_emb_mindiesd paths.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.layers.rope import RotaryEmbedding
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_inputs(
+    batch: int,
+    seq_len: int,
+    num_heads: int,
+    head_dim: int,
+    cos_batch: int | None = None,
+    dtype: torch.dtype = torch.float32,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Generate x, cos, sin tensors for RoPE testing.
+
+    Args:
+        batch: Batch dimension of x.
+        seq_len: Sequence length.
+        num_heads: Number of attention heads.
+        head_dim: Head dimension (must be even).
+        cos_batch: Batch dimension of cos/sin. Defaults to batch if None.
+        dtype: Data type for the tensors.
+    """
+    if cos_batch is None:
+        cos_batch = batch
+
+    torch.manual_seed(42)
+    x = torch.randn(batch, seq_len, num_heads, head_dim, dtype=dtype)
+    cos = torch.randn(cos_batch, seq_len, head_dim // 2, dtype=dtype)
+    sin = torch.randn(cos_batch, seq_len, head_dim // 2, dtype=dtype)
+    return x, cos, sin
+
+
+class TestRotaryEmbeddingNativeShapeRegression:
+    """Regression tests for forward_native shape handling (#8e297d81)."""
+
+    @pytest.fixture
+    def rope_neox(self) -> RotaryEmbedding:
+        """NeoX-style (non-interleaved) RoPE."""
+        return RotaryEmbedding(is_neox_style=True)
+
+    @pytest.fixture
+    def rope_gptj(self) -> RotaryEmbedding:
+        """GPT-J style (interleaved) RoPE."""
+        return RotaryEmbedding(is_neox_style=False)
+
+    # ── 2D cos/sin (no batch dim) ──────────────────────────────────────
+
+    def test_2d_cos_sin_unchanged(self, rope_neox: RotaryEmbedding) -> None:
+        """2D cos/sin [S, D/2] pass through unchanged (no-op path)."""
+        x, cos, sin = _make_inputs(batch=2, seq_len=16, num_heads=8, head_dim=64, cos_batch=2)
+        # Manually drop batch dim to create 2D cos/sin
+        cos_2d = cos[0]
+        sin_2d = sin[0]
+        assert cos_2d.dim() == 2
+        output = rope_neox.forward_native(x, cos_2d, sin_2d)
+        assert output.shape == x.shape
+        assert not torch.isnan(output).any()
+        assert not torch.isinf(output).any()
+
+    # ── 3D cos/sin — batch-dim stripping ──────────────────────────────
+
+    @pytest.mark.parametrize("cos_batch", [1, 2, 4])
+    def test_3d_cos_sin_no_error(self, rope_neox: RotaryEmbedding, cos_batch: int) -> None:
+        """3D cos/sin with batch dim should not raise errors."""
+        x, cos, sin = _make_inputs(batch=2, seq_len=16, num_heads=8, head_dim=64, cos_batch=cos_batch)
+        output = rope_neox.forward_native(x, cos, sin)
+        assert output.shape == x.shape
+
+    @pytest.mark.parametrize("interleaved", [False, True])
+    def test_3d_cos_identical_to_2d(self, interleaved: bool) -> None:
+        """3D cos/sin with identical batch slices == 2D cos/sin result."""
+        rope = RotaryEmbedding(is_neox_style=not interleaved)
+        B, S, H, D = 2, 16, 8, 64
+        x = torch.randn(B, S, H, D)
+        cos_2d = torch.randn(S, D // 2)
+        sin_2d = torch.randn(S, D // 2)
+
+        # 3D variant: expand the same 2D data across batch
+        cos_3d = cos_2d.unsqueeze(0).expand(B, -1, -1)
+        sin_3d = sin_2d.unsqueeze(0).expand(B, -1, -1)
+
+        out_2d = rope.forward_native(x, cos_2d, sin_2d)
+        out_3d = rope.forward_native(x, cos_3d, sin_3d)
+
+        torch.testing.assert_close(out_3d, out_2d, atol=0, rtol=0)
+
+    # ── Batch size mismatch — the regression scenario ─────────────────
+
+    @pytest.mark.parametrize(
+        "x_batch,cos_batch",
+        [
+            (2, 1),  # cos has fewer batch elements
+            (1, 2),  # cos has more batch elements (takes [0])
+            (4, 1),  # large x batch, single cos
+        ],
+    )
+    def test_batch_size_mismatch_regression(
+        self,
+        rope_neox: RotaryEmbedding,
+        x_batch: int,
+        cos_batch: int,
+    ) -> None:
+        """forward_native works when cos batch dim differs from x batch dim.
+
+        This is the exact regression scenario — before the fix, a 3D cos/sin
+        with a different batch size than x would cause a shape mismatch error
+        inside apply_rotary_emb_torch during broadcast.
+        """
+        x, cos, sin = _make_inputs(batch=x_batch, seq_len=16, num_heads=8, head_dim=64, cos_batch=cos_batch)
+        output = rope_neox.forward_native(x, cos, sin)
+        assert output.shape == x.shape
+        assert output.dtype == x.dtype
+
+    # ── Batch consistency — all elements get same rotation ────────────
+
+    def test_2d_cos_broadcasts_identically_across_batch(self, rope_neox: RotaryEmbedding) -> None:
+        """With 2D cos, identical x slices → identical output slices."""
+        B, S, H, D = 4, 16, 8, 64
+        x_template = torch.randn(1, S, H, D)
+        x = x_template.expand(B, -1, -1, -1)  # all batch elements identical
+        cos = torch.randn(S, D // 2)
+        sin = torch.randn(S, D // 2)
+
+        output = rope_neox.forward_native(x, cos, sin)
+
+        # All output batch slices must be identical
+        for i in range(1, B):
+            torch.testing.assert_close(output[i], output[0], atol=0, rtol=0)
+
+    def test_3d_identical_cos_broadcasts_identically_across_batch(self, rope_neox: RotaryEmbedding) -> None:
+        """With 3D cos (identical slices), identical x → identical output."""
+        B, S, H, D = 4, 16, 8, 64
+        x_template = torch.randn(1, S, H, D)
+        x = x_template.expand(B, -1, -1, -1)
+        cos = torch.randn(1, S, D // 2).expand(B, -1, -1)
+        sin = torch.randn(1, S, D // 2).expand(B, -1, -1)
+
+        output = rope_neox.forward_native(x, cos, sin)
+
+        for i in range(1, B):
+            torch.testing.assert_close(output[i], output[0], atol=0, rtol=0)
+
+    # ── Output shape / dtype invariants ───────────────────────────────
+
+    def test_output_shape_preserved(self, rope_neox: RotaryEmbedding) -> None:
+        """Output shape must match input shape for all cos/sin dimensionalities."""
+        B, S, H, D = 2, 32, 4, 128
+        x = torch.randn(B, S, H, D)
+
+        # 2D cos/sin
+        cos_2d = torch.randn(S, D // 2)
+        sin_2d = torch.randn(S, D // 2)
+        assert rope_neox.forward_native(x, cos_2d, sin_2d).shape == (B, S, H, D)
+
+        # 3D cos/sin, matching batch
+        cos_3d = torch.randn(B, S, D // 2)
+        sin_3d = torch.randn(B, S, D // 2)
+        assert rope_neox.forward_native(x, cos_3d, sin_3d).shape == (B, S, H, D)
+
+        # 3D cos/sin, different batch
+        cos_3d_1 = torch.randn(1, S, D // 2)
+        sin_3d_1 = torch.randn(1, S, D // 2)
+        assert rope_neox.forward_native(x, cos_3d_1, sin_3d_1).shape == (B, S, H, D)
+
+    def test_output_dtype_preserved(self, rope_neox: RotaryEmbedding) -> None:
+        """Output dtype must match input dtype."""
+        for dtype in [torch.float32, torch.float16, torch.bfloat16]:
+            x, cos, sin = _make_inputs(batch=2, seq_len=16, num_heads=8, head_dim=64, dtype=dtype)
+            output = rope_neox.forward_native(x, cos, sin)
+            assert output.dtype == dtype, f"dtype mismatch: {output.dtype} != {dtype}"
+
+    def test_output_no_nan_inf(
+        self,
+        rope_neox: RotaryEmbedding,
+        rope_gptj: RotaryEmbedding,
+    ) -> None:
+        """Output must be finite for both RoPE styles."""
+        for rope in (rope_neox, rope_gptj):
+            x, cos, sin = _make_inputs(batch=2, seq_len=32, num_heads=4, head_dim=64)
+            output = rope.forward_native(x, cos, sin)
+            assert not torch.isnan(output).any(), f"NaN in {rope}"
+            assert not torch.isinf(output).any(), f"Inf in {rope}"
+
+    def test_output_contiguous(self, rope_neox: RotaryEmbedding) -> None:
+        """Output should be contiguous."""
+        x, cos, sin = _make_inputs(batch=2, seq_len=16, num_heads=8, head_dim=64)
+        output = rope_neox.forward_native(x, cos, sin)
+        assert output.is_contiguous()
+
+    # ── CFG-style doubled batch ───────────────────────────────────────
+
+    def test_cfg_style_doubled_batch(self, rope_neox: RotaryEmbedding) -> None:
+        """CFG (classifier-free guidance) doubles batch — cos may have
+        original batch size while x has 2× batch."""
+        B, S, H, D = 2, 16, 8, 64
+        # Simulate CFG: x has 2× batch (cond + uncond concatenated)
+        x = torch.randn(B * 2, S, H, D)
+        # cos/sin only have the original batch
+        cos = torch.randn(B, S, D // 2)
+        sin = torch.randn(B, S, D // 2)
+
+        output = rope_neox.forward_native(x, cos, sin)
+        assert output.shape == (B * 2, S, H, D)
+
+    # ── RotaryEmbeddingWan forward_native is unchanged ────────────────
+
+    def test_wan_forward_native_unchanged(self) -> None:
+        """RotaryEmbeddingWan.forward_native still uses its own implementation
+        (unflatten + stack + flatten) and should not crash with 3D cos/sin."""
+        from vllm_omni.diffusion.layers.rope import RotaryEmbeddingWan
+
+        rope = RotaryEmbeddingWan(is_neox_style=False, half_head_dim=True)
+        B, S, H, D = 2, 16, 8, 64
+        x = torch.randn(B, S, H, D)
+        cos = torch.randn(B, S, 1, D // 2)
+        sin = torch.randn(B, S, 1, D // 2)
+
+        output = rope.forward_native(x, cos, sin)
+        assert output.shape == x.shape
+        assert not torch.isnan(output).any()

tests/diffusion/models/hunyuan_image3/test_hunyuan_fused_moe.py

@@ -50,6 +50,48 @@ class TestSetForwardContextNumTokens:
         mock_get.assert_not_called()
 
 
+class TestSetForwardContextDPMetadata:
+    """Test diffusion MoE DP metadata setup for vLLM 0.24 all2all dispatch."""
+
+    def test_sets_dp_metadata_from_vllm_dp_group(self, mocker):
+        import vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe as hunyuan_moe
+
+        mock_ctx = mocker.MagicMock()
+        mock_ctx.dp_metadata = None
+        fake_group = mocker.MagicMock()
+        fake_group.world_size = 2
+        fake_group.cpu_group = object()
+
+        def _all_gather_object(output, obj, group):
+            assert group is fake_group.cpu_group
+            output[:] = [obj, obj + 1]
+
+        mocker.patch.object(hunyuan_moe._vllm_fc, "is_forward_context_available", return_value=True)
+        mocker.patch.object(hunyuan_moe._vllm_fc, "get_forward_context", return_value=mock_ctx)
+        mocker.patch.object(hunyuan_moe._vllm_ps, "_DP", fake_group, create=True)
+        mocker.patch.object(hunyuan_moe.torch.distributed, "all_gather_object", side_effect=_all_gather_object)
+
+        hunyuan_moe._set_forward_context_dp_metadata(1024)
+
+        assert mock_ctx.dp_metadata is not None
+        assert mock_ctx.dp_metadata.num_tokens_across_dp_cpu.tolist() == [1024, 1025]
+
+    def test_keeps_existing_dp_metadata(self, mocker):
+        import vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe as hunyuan_moe
+
+        existing_metadata = object()
+        mock_ctx = mocker.MagicMock()
+        mock_ctx.dp_metadata = existing_metadata
+        mocker.patch.object(hunyuan_moe._vllm_fc, "is_forward_context_available", return_value=True)
+        mocker.patch.object(hunyuan_moe._vllm_fc, "get_forward_context", return_value=mock_ctx)
+        mock_all_gather = mocker.patch.object(hunyuan_moe.torch.distributed, "all_gather_object")
+
+        hunyuan_moe._set_forward_context_dp_metadata(1024)
+
+        assert mock_ctx.dp_metadata is existing_metadata
+        mock_all_gather.assert_not_called()
+
+
 class TestHunyuanFusedMoEPlatformDispatch:
     """Test platform dispatch via platform qualname hooks."""
 

tests/diffusion/models/lingbot_video/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

tests/diffusion/models/lingbot_video/test_lingbot_video_transformer.py

@@ -0,0 +1,117 @@
+# 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 _tiny_transformer(**overrides):
+    from vllm_omni.diffusion.models.lingbot_video import LingBotVideoTransformer3DModel
+
+    config = {
+        "patch_size": (1, 1, 1),
+        "in_channels": 2,
+        "out_channels": 2,
+        "hidden_size": 16,
+        "num_attention_heads": 1,
+        "depth": 0,
+        "intermediate_size": 32,
+        "text_dim": 8,
+        "freq_dim": 8,
+        "axes_dims": (4, 4, 8),
+        "axes_lens": (32, 32, 32),
+    }
+    config.update(overrides)
+    return LingBotVideoTransformer3DModel(**config)
+
+
+def test_joint_position_ids_video_then_text_order():
+    from vllm_omni.diffusion.models.lingbot_video.lingbot_video_transformer import make_joint_position_ids
+
+    positions = make_joint_position_ids(text_len=3, grid_t=1, grid_h=2, grid_w=2, device=torch.device("cpu"))
+
+    assert positions.shape == (7, 3)
+    assert positions[:4, 0].tolist() == [4, 4, 4, 4]
+    assert positions[:4, 1:].tolist() == [[0, 0], [0, 1], [1, 0], [1, 1]]
+    assert positions[4:].tolist() == [[1, 0, 0], [2, 0, 0], [3, 0, 0]]
+
+
+def test_tiny_transformer_depth_zero_forward_shape():
+    model = _tiny_transformer()
+    hidden_states = torch.randn(1, 2, 1, 2, 2)
+    timestep = torch.tensor([300.0])
+    encoder_hidden_states = torch.randn(1, 3, 8)
+    encoder_attention_mask = torch.ones(1, 3, dtype=torch.long)
+
+    with torch.no_grad():
+        out = model(
+            hidden_states,
+            timestep,
+            encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            return_dict=False,
+        )[0]
+
+    assert out.shape == hidden_states.shape
+    assert torch.isfinite(out).all()
+
+
+def test_packed_attention_uses_sdpa_fallback_without_flash_varlen(monkeypatch):
+    from vllm_omni.diffusion.models.lingbot_video import lingbot_video_transformer as module
+
+    monkeypatch.setattr(module, "flash_attn_varlen_func_v3", None)
+    attn = module.LingBotVideoAttention(
+        hidden_size=8,
+        num_heads=2,
+        norm_eps=1e-6,
+        qkv_bias=False,
+        out_bias=False,
+    )
+    captured = {}
+
+    def fake_sdpa_forward(query, key, value, attn_metadata):
+        captured["mask"] = attn_metadata.attn_mask
+        return torch.zeros_like(query)
+
+    monkeypatch.setattr(attn.attn.sdpa_fallback, "forward", fake_sdpa_forward)
+    x = torch.randn(1, 5, 8)
+    rotary = torch.ones(1, 5, 2, dtype=torch.complex64)
+    packed_indices = {
+        "cu_seqlens_kv": torch.tensor([0, 2, 5], dtype=torch.int32),
+        "max_seqlen_in_batch_kv": 3,
+        "attention_mask": module._packed_block_attention_mask([2, 3], x.device),
+    }
+
+    out = attn(x, rotary, packed_indices=packed_indices)
+
+    assert out.shape == x.shape
+    mask = captured["mask"]
+    assert mask.shape == (1, 1, 5, 5)
+    assert mask[0, 0, :2, :2].all()
+    assert mask[0, 0, 2:, 2:].all()
+    assert not mask[0, 0, :2, 2:].any()
+    assert not mask[0, 0, 2:, :2].any()
+
+
+def test_tiny_transformer_rejects_invalid_rope_dims():
+    from vllm_omni.diffusion.models.lingbot_video import LingBotVideoTransformer3DModel
+
+    with pytest.raises(AssertionError, match="head_dim"):
+        LingBotVideoTransformer3DModel(
+            hidden_size=16,
+            num_attention_heads=1,
+            axes_dims=(4, 4, 4),
+            depth=0,
+        )
+
+
+def test_transformer_to_keeps_sensitive_modules_in_fp32():
+    model = _tiny_transformer()
+
+    model.to(dtype=torch.bfloat16)
+
+    assert model.patch_embedder.weight.dtype == torch.bfloat16
+    assert model.time_embedder.linear_1.weight.dtype == torch.float32
+    assert model.norm_out_modulation[1].weight.dtype == torch.float32

tests/diffusion/models/lingbot_video/test_pipeline_lingbot_video.py

@@ -0,0 +1,195 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def _make_pipeline():
+    from vllm_omni.diffusion.models.lingbot_video.pipeline_lingbot_video import LingBotVideoPipeline
+
+    pipeline = object.__new__(LingBotVideoPipeline)
+    nn.Module.__init__(pipeline)
+    pipeline.device = torch.device("cpu")
+    pipeline.default_negative_prompt = "default negative"
+    pipeline.od_config = SimpleNamespace(flow_shift=None)
+    return pipeline
+
+
+def _make_request_batch(prompt, **sampling_overrides):
+    sampling = OmniDiffusionSamplingParams(**sampling_overrides)
+    request = OmniDiffusionRequest(prompt=prompt, sampling_params=sampling, request_id="lingbot-test")
+    return DiffusionRequestBatch([request])
+
+
+def test_lingbot_video_pipeline_import_and_registry():
+    from vllm_omni.diffusion.models.lingbot_video import (
+        LingBotVideoPipeline,
+        LingBotVideoTransformer3DModel,
+        get_lingbot_video_post_process_func,
+    )
+    from vllm_omni.diffusion.registry import _DIFFUSION_MODELS, _DIFFUSION_POST_PROCESS_FUNCS
+
+    assert LingBotVideoPipeline is not None
+    assert LingBotVideoTransformer3DModel is not None
+    assert get_lingbot_video_post_process_func is not None
+    assert _DIFFUSION_MODELS["LingBotVideoPipeline"] == (
+        "lingbot_video",
+        "pipeline_lingbot_video",
+        "LingBotVideoPipeline",
+    )
+    assert _DIFFUSION_POST_PROCESS_FUNCS["LingBotVideoPipeline"] == "get_lingbot_video_post_process_func"
+
+
+def test_component_discovery_declarations():
+    from vllm_omni.diffusion.models.lingbot_video import LingBotVideoPipeline
+
+    assert LingBotVideoPipeline._dit_modules == ["transformer"]
+    assert LingBotVideoPipeline._encoder_modules == ["text_encoder"]
+    assert LingBotVideoPipeline._vae_modules == ["vae"]
+    assert LingBotVideoPipeline.supports_step_execution is False
+
+
+def test_extra_body_params_include_video_flow_shift_alias():
+    from vllm_omni.model_extras import get_extra_body_params
+
+    params = get_extra_body_params("LingBotVideoPipeline")
+    assert "shift" in params
+    assert "flow_shift" in params
+    assert "batch_cfg" in params
+
+
+def test_extract_prompt_string_and_mapping():
+    from vllm_omni.diffusion.models.lingbot_video.pipeline_lingbot_video import _extract_prompt
+
+    string_req = SimpleNamespace(prompt="a robot waves")
+    assert _extract_prompt(string_req) == ("a robot waves", None)
+
+    mapping_req = SimpleNamespace(prompt={"prompt": "a robot walks", "negative_prompt": "low quality"})
+    assert _extract_prompt(mapping_req) == ("a robot walks", "low quality")
+
+    with pytest.raises(TypeError, match="Unsupported LingBot prompt type"):
+        _extract_prompt(SimpleNamespace(prompt=["not", "supported"]))
+
+
+def test_check_inputs_accepts_t2v_shapes_and_rejects_invalid_values():
+    from vllm_omni.diffusion.models.lingbot_video import LingBotVideoPipeline
+
+    LingBotVideoPipeline.check_inputs(height=192, width=320, num_frames=1)
+    LingBotVideoPipeline.check_inputs(height=192, width=320, num_frames=9)
+
+    with pytest.raises(ValueError, match="num_frames"):
+        LingBotVideoPipeline.check_inputs(height=192, width=320, num_frames=8)
+    with pytest.raises(ValueError, match="height"):
+        LingBotVideoPipeline.check_inputs(height=190, width=320, num_frames=9)
+
+
+def test_postprocess_keeps_torch_by_default_and_converts_np_when_requested():
+    from vllm_omni.diffusion.models.lingbot_video import get_lingbot_video_post_process_func
+
+    frames = torch.ones(2, 4, 4, 3)
+    post = get_lingbot_video_post_process_func(SimpleNamespace())
+
+    assert post(frames, SimpleNamespace(output_type="pt")) is frames
+    array = post(frames, SimpleNamespace(output_type="np"))
+    assert isinstance(array, np.ndarray)
+    assert array.shape == (2, 4, 4, 3)
+
+
+def test_forward_resolves_t2v_sampling_and_flow_shift_alias():
+    pipeline = _make_pipeline()
+    calls = []
+
+    def fake_generate(**kwargs):
+        calls.append(kwargs)
+        return torch.zeros(9, 192, 320, 3)
+
+    pipeline._generate = fake_generate
+    req = _make_request_batch(
+        {"prompt": "a robot arm", "negative_prompt": "bad hands"},
+        height=192,
+        width=320,
+        num_frames=9,
+        num_inference_steps=2,
+        guidance_scale=3.0,
+        seed=42,
+        output_type="pt",
+        extra_args={"flow_shift": 4.0, "batch_cfg": True},
+    )
+
+    out = pipeline.forward(req)
+
+    assert torch.equal(out.output, torch.zeros(9, 192, 320, 3))
+    assert len(calls) == 1
+    call = calls[0]
+    assert call["prompt"] == "a robot arm"
+    assert call["negative_prompt"] == "bad hands"
+    assert call["height"] == 192
+    assert call["width"] == 320
+    assert call["num_frames"] == 9
+    assert call["num_inference_steps"] == 2
+    assert call["guidance_scale"] == 3.0
+    assert call["shift"] == 4.0
+    assert call["batch_cfg"] is True
+    assert call["output_type"] == "pt"
+
+
+def test_forward_prefers_shift_over_flow_shift_and_defaults_negative_prompt():
+    pipeline = _make_pipeline()
+    calls = []
+
+    def fake_generate(**kwargs):
+        calls.append(kwargs)
+        return torch.zeros(5, 192, 192, 3)
+
+    pipeline._generate = fake_generate
+    req = _make_request_batch(
+        "a robot arm",
+        height=192,
+        width=192,
+        num_frames=5,
+        num_inference_steps=2,
+        guidance_scale=1.0,
+        seed=42,
+        extra_args={"shift": 5.0, "flow_shift": 4.0},
+    )
+
+    pipeline.forward(req)
+
+    assert calls[0]["shift"] == 5.0
+    assert calls[0]["negative_prompt"] == "default negative"
+
+
+def test_forward_rejects_multi_request_batches():
+    pipeline = _make_pipeline()
+    first = OmniDiffusionRequest(
+        prompt="first",
+        sampling_params=OmniDiffusionSamplingParams(seed=1),
+        request_id="lingbot-first",
+    )
+    second = OmniDiffusionRequest(
+        prompt="second",
+        sampling_params=OmniDiffusionSamplingParams(seed=2),
+        request_id="lingbot-second",
+    )
+
+    with pytest.raises(ValueError, match="only supports one request"):
+        pipeline.forward(DiffusionRequestBatch([first, second]))
+
+
+def test_load_weights_rejects_external_weight_stream():
+    pipeline = _make_pipeline()
+
+    assert pipeline.load_weights([]) == set()
+    with pytest.raises(RuntimeError, match="components are loaded directly"):
+        pipeline.load_weights([("transformer.weight", torch.zeros(1))])

tests/dreamzero/test_pipeline_state.py

@@ -4,7 +4,6 @@
 from collections import OrderedDict
 
 import pytest
-import torch
 
 from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
 from vllm_omni.diffusion.models.dreamzero.state_dreamzero import DreamZeroState
@@ -46,14 +45,11 @@ def test_dreamzero_pipeline_state_lru_caps_retained_sessions() -> None:
     assert "session-b" not in pipeline._states
 
 
-def test_dreamzero_state_cache_access_requires_initialization() -> None:
+def test_dreamzero_state_owns_no_kv_caches() -> None:
+    """KV (self- and cross-attention) is engine-owned since the AR-Diffusion
+    paged backend: the model-local cache accessors must be gone so nothing can
+    silently bypass the engine pool."""
     state = DreamZeroState()
 
-    with pytest.raises(RuntimeError, match="KV caches not initialized"):
-        state.get_kv_caches()
-
-    with pytest.raises(RuntimeError, match="Cross-attn caches not initialized"):
-        state.get_crossattn_caches()
-
-    with pytest.raises(RuntimeError, match="create_kv_caches first"):
-        state.update_kv_cache(0, torch.empty(0))
+    for removed in ("get_kv_caches", "create_kv_caches", "update_kv_cache", "get_crossattn_caches"):
+        assert not hasattr(state, removed)

tests/e2e/offline_inference/test_flux2_klein.py

@@ -18,6 +18,8 @@ from tests.helpers.runtime import DiffusionResponse, OmniRunnerHandler
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform
 
+# NOTE: Most of the flux2klein tests have migrated to the common diffusion tests.
+# This model will also use the tiny model builder when run at the core model level.
 MODEL = "black-forest-labs/FLUX.2-klein-4B"
 
 _OMNI_RUNNER_PARAM = (MODEL, None)
@@ -70,41 +72,6 @@ def _send_inpaint_with_generator(
     )
 
 
-# Regression test for https://github.com/vllm-project/vllm-omni/issues/3097
-@pytest.mark.advanced_model
-@pytest.mark.diffusion
-def test_flux2_klein_can_accept_text_inputs(omni_runner_handler: OmniRunnerHandler):
-    omni_runner_handler.send_diffusion_request(
-        {
-            "model": MODEL,
-            "prompt": "a cup of coffee on the table",
-            "sampling_params": OmniDiffusionSamplingParams(num_inference_steps=2, seed=42),
-        }
-    )
-
-
-@pytest.mark.advanced_model
-@pytest.mark.diffusion
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
-def test_flux2_klein_inpaint_basic(omni_runner_handler: OmniRunnerHandler):
-    input_image, mask_image = _create_test_inputs()
-    omni_runner_handler.send_diffusion_request(
-        {
-            "model": MODEL,
-            "prompt": "Fill in the masked area with a beautiful garden",
-            "multi_modal_data": {"image": input_image, "mask_image": mask_image},
-            "sampling_params": OmniDiffusionSamplingParams(
-                height=_HEIGHT,
-                width=_WIDTH,
-                num_inference_steps=_NUM_INFERENCE_STEPS,
-                guidance_scale=0.0,
-                generator=torch.Generator(current_omni_platform.device_type).manual_seed(42),
-                num_outputs_per_prompt=1,
-            ),
-        }
-    )
-
-
 @pytest.mark.advanced_model
 @pytest.mark.diffusion
 @hardware_test(res={"cuda": "L4"}, num_cards=1)

tests/e2e/offline_inference/test_magi_human_expansion.py

@@ -22,6 +22,7 @@ _OMNI_RUNNER_PARAM = (
     {
         "init_timeout": 1200,
         "tensor_parallel_size": 2,
+        "vae_use_tiling": True,
     },
 )
 
@@ -32,6 +33,7 @@ _OMNI_RUNNER_PARAM_OFFLOAD = (
         "init_timeout": 1200,
         "tensor_parallel_size": 2,
         "enable_cpu_offload": True,
+        "vae_use_tiling": True,
     },
 )
 
@@ -45,10 +47,6 @@ pytestmark = [
     ),
 ]
 
-_SKIP_ISSUE_3236 = pytest.mark.skip(
-    reason="https://github.com/vllm-project/vllm-omni/issues/3236",
-)
-
 
 def _validate_mp4(video_bytes: bytes, min_frames: int = 10) -> None:
     """Validate that the MP4 contains meaningful video and audio tracks."""
@@ -74,7 +72,6 @@ def _validate_mp4(video_bytes: bytes, min_frames: int = 10) -> None:
     container.close()
 
 
-@_SKIP_ISSUE_3236
 @hardware_test(res={"cuda": "H100"}, num_cards=2)
 def test_magi_human_e2e(omni_runner: OmniRunner):
     """End-to-end test for MagiHuman generating video and audio."""

tests/e2e/offline_inference/test_mammoth_moda2_expansion.py

@@ -5,7 +5,7 @@ Verifies that the AR->DiT pipeline produces an image tensor whose pixel values
 match a golden reference.
 
 Model Hub repo id: ``bytedance-research/MammothModa2-Preview``.
-Stage config: ``get_deploy_config_path("mammoth_moda2.yaml")`` → ``vllm_omni/deploy/mammoth_moda2.yaml``
+Deploy config: ``get_deploy_config_path("mammoth_moda2.yaml")`` -> ``vllm_omni/deploy/mammoth_moda2.yaml``
 
 Golden pixel file: ``tests/e2e/offline_inference/fixtures/mammoth_moda2_t2i_golden.json``
   Regenerate with: ``UPDATE_GOLDEN=1 pytest tests/e2e/offline_inference/test_mammoth_moda2_expansion.py``
@@ -36,9 +36,9 @@ _VISION_END_TOKEN_ID = 151653  # "<|vision_end|>"
 _AR_PATCH_SIZE = 16
 
 MODEL_PATH = "bytedance-research/MammothModa2-Preview"
-T2I_STAGE_CONFIG = get_deploy_config_path("mammoth_moda2.yaml")
+T2I_DEPLOY_CONFIG = get_deploy_config_path("mammoth_moda2.yaml")
 
-_OMNI_RUNNER_PARAM = (MODEL_PATH, T2I_STAGE_CONFIG)
+_OMNI_RUNNER_PARAM = (MODEL_PATH, None, {"deploy_config": T2I_DEPLOY_CONFIG})
 
 # Golden pixel reference file.  Set UPDATE_GOLDEN=1 to regenerate.
 _GOLDEN_T2I_PATH = Path(__file__).parent / "fixtures" / "mammoth_moda2_t2i_golden.json"

tests/e2e/offline_inference/test_omnivoice_expansion.py

@@ -18,13 +18,15 @@ from tests.helpers.runtime import OmniRunner
 from tests.helpers.stage_config import get_deploy_config_path
 
 MODEL = "k2-fsa/OmniVoice"
-STAGE_CONFIG = get_deploy_config_path("omnivoice.yaml")
+DEPLOY_CONFIG = get_deploy_config_path("omnivoice.yaml")
 
-# (model, stage_config_path, extra_omni_kwargs) — see ``omni_runner`` in tests.helpers.fixtures.runtime
+# OmniRunner tuple: model, legacy stage config path, extra Omni kwargs.
+# The migrated test passes deploy_config through extra Omni kwargs.
 _OMNI_RUNNER_PARAM = (
     MODEL,
-    STAGE_CONFIG,
+    None,
     {
+        "deploy_config": DEPLOY_CONFIG,
         "trust_remote_code": True,
         "log_stats": True,
     },

tests/e2e/online_serving/test_flux2_klein.py

@@ -1,62 +0,0 @@
-"""
-Tests for Flux2 Klein; currently Dev is implemented separately,
-but ideally these models will fold together in the future.
-"""
-
-import pytest
-
-from tests.helpers.mark import hardware_marks
-from tests.helpers.runtime import (
-    OmniServer,
-    OmniServerParams,
-    OpenAIClientHandler,
-    dummy_messages_from_mix_data,
-)
-
-FOUR_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"}, num_cards=4)
-POSITIVE_PROMPT = "A cat sitting on a windowsill"
-NEGATIVE_PROMPT = "blurry, low quality"
-
-
-# For now, we only test a good configuration for this model to keep the CI
-# light, but in the near future may explore using random models to check
-# cross-feature compatibility more generally.
-def _get_diffusion_feature_cases(model: str):
-    return [
-        # FP8 / Hybrid sequence parallelism
-        pytest.param(
-            OmniServerParams(
-                model=model,
-            ),
-            id="default",
-            marks=FOUR_CARD_FEATURE_MARKS,
-        ),
-    ]
-
-
-@pytest.mark.advanced_model
-@pytest.mark.core_model
-@pytest.mark.diffusion
-@pytest.mark.parametrize(
-    "omni_server",
-    _get_diffusion_feature_cases(
-        model="black-forest-labs/FLUX.2-klein-4B",
-    ),
-    indirect=True,
-)
-def test_text_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):
-    messages = dummy_messages_from_mix_data(content_text=POSITIVE_PROMPT)
-    request_config = {
-        "model": omni_server.model,
-        "messages": messages,
-        "extra_body": {
-            "height": 512,
-            "width": 512,
-            "num_inference_steps": 2,
-            "negative_prompt": NEGATIVE_PROMPT,
-            "true_cfg_scale": 4.0,
-            "seed": 42,
-        },
-    }
-
-    openai_client.send_diffusion_request(request_config)

tests/e2e/online_serving/test_flux2_klein_expansion.py

@@ -1,163 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""
-End-to-end tests for Flux2 Klein inpainting in online serving mode.
-
-Uses /v1/images/edits endpoint which is the correct API for image inpainting.
-"""
-
-import base64
-from io import BytesIO
-
-import httpx
-import pytest
-from PIL import Image, ImageDraw
-
-from tests.helpers.runtime import OmniServer, OmniServerParams
-
-pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
-
-MODEL = "black-forest-labs/FLUX.2-klein-4B"
-
-_HEIGHT = 512
-_WIDTH = 512
-_NUM_INFERENCE_STEPS = 4
-
-
-def _get_diffusion_feature_cases(model: str):
-    return [
-        pytest.param(
-            OmniServerParams(
-                model=model,
-                server_args=["--tensor-parallel-size", "2", "--enable-cpu-offload"],
-            ),
-            id="tp2_cpu_offload",
-        ),
-        pytest.param(
-            OmniServerParams(
-                model=model,
-                server_args=["--tensor-parallel-size", "2"],
-            ),
-            id="tp2_basic",
-        ),
-    ]
-
-
-def _image_to_base64_jpeg(image: Image.Image) -> str:
-    buffer = BytesIO()
-    image.save(buffer, format="JPEG")
-    buffer.seek(0)
-    return base64.b64encode(buffer.read()).decode("utf-8")
-
-
-def _create_test_mask_base64(width: int = _WIDTH, height: int = _HEIGHT) -> str:
-    mask = Image.new("L", (width, height), 0)
-    draw = ImageDraw.Draw(mask)
-    draw.rectangle([width // 4, height // 4, width * 3 // 4, height * 3 // 4], fill=255)
-    return _image_to_base64_jpeg(mask)
-
-
-def _compare_images(img1: Image.Image, img2: Image.Image) -> bool:
-    return list(img1.getdata()) == list(img2.getdata())
-
-
-def _send_edit_request(host: str, port: int, model: str, image_b64: str, mask_b64: str, prompt: str, **kwargs):
-    url = f"http://{host}:{port}/v1/images/edits"
-    files = {
-        "image": ("image.jpg", base64.b64decode(image_b64), "image/jpeg"),
-        "mask_image": ("mask.jpg", base64.b64decode(mask_b64), "image/jpeg"),
-    }
-    data = {"prompt": prompt, "model": model, **kwargs}
-    with httpx.Client(timeout=60.0) as client:
-        response = client.post(url, files=files, data=data)
-        response.raise_for_status()
-        return response.json()
-
-
-@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
-def test_flux2_klein_inpaint_basic(omni_server: OmniServer):
-    input_image_b64 = _image_to_base64_jpeg(Image.new("RGB", (_WIDTH, _HEIGHT), (128, 128, 128)))
-    mask_b64 = _create_test_mask_base64()
-
-    result = _send_edit_request(
-        host=omni_server.host,
-        port=omni_server.port,
-        model=MODEL,
-        image_b64=input_image_b64,
-        mask_b64=mask_b64,
-        prompt="Fill in the masked area with a beautiful garden",
-        guidance_scale=1.0,
-        num_inference_steps=_NUM_INFERENCE_STEPS,
-        n=1,
-        seed=42,
-    )
-
-    assert "data" in result and len(result["data"]) == 1
-    img_data = result["data"][0].get("b64_json") or result["data"][0].get("url", "").split(",")[-1]
-    img = Image.open(BytesIO(base64.b64decode(img_data)))
-    assert img.size == (_WIDTH, _HEIGHT)
-
-
-@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
-def test_flux2_klein_inpaint_deterministic(omni_server: OmniServer):
-    input_image_b64 = _image_to_base64_jpeg(Image.new("RGB", (_WIDTH, _HEIGHT), (128, 128, 128)))
-    mask_b64 = _create_test_mask_base64()
-    prompt = "A red flower in a field"
-
-    result1 = _send_edit_request(
-        host=omni_server.host,
-        port=omni_server.port,
-        model=MODEL,
-        image_b64=input_image_b64,
-        mask_b64=mask_b64,
-        prompt=prompt,
-        guidance_scale=1.0,
-        num_inference_steps=_NUM_INFERENCE_STEPS,
-        n=1,
-        seed=12345,
-    )
-
-    result2 = _send_edit_request(
-        host=omni_server.host,
-        port=omni_server.port,
-        model=MODEL,
-        image_b64=input_image_b64,
-        mask_b64=mask_b64,
-        prompt=prompt,
-        guidance_scale=1.0,
-        num_inference_steps=_NUM_INFERENCE_STEPS,
-        n=1,
-        seed=12345,
-    )
-
-    img1_data = result1["data"][0].get("b64_json") or result1["data"][0].get("url", "").split(",")[-1]
-    img2_data = result2["data"][0].get("b64_json") or result2["data"][0].get("url", "").split(",")[-1]
-
-    img1 = Image.open(BytesIO(base64.b64decode(img1_data)))
-    img2 = Image.open(BytesIO(base64.b64decode(img2_data)))
-
-    assert _compare_images(img1, img2), (
-        "Same input with same seed should produce identical output. This is critical for offline/online consistency."
-    )
-
-
-@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
-def test_flux2_klein_inpaint_multiple_outputs(omni_server: OmniServer):
-    input_image_b64 = _image_to_base64_jpeg(Image.new("RGB", (_WIDTH, _HEIGHT), (128, 128, 128)))
-    mask_b64 = _create_test_mask_base64()
-
-    result = _send_edit_request(
-        host=omni_server.host,
-        port=omni_server.port,
-        model=MODEL,
-        image_b64=input_image_b64,
-        mask_b64=mask_b64,
-        prompt="A beautiful landscape",
-        guidance_scale=1.0,
-        num_inference_steps=_NUM_INFERENCE_STEPS,
-        n=2,
-        seed=42,
-    )
-
-    assert "data" in result and len(result["data"]) == 2

tests/e2e/online_serving/test_higgs_audio_v3.py

@@ -127,7 +127,7 @@ class TestHiggsAudioV3OnlineHappyPath:
         windows stitch into a coherent PCM stream. The byte-count gate is the same as
         the sync paths; per-chunk audio content is verified offline against Whisper.
 
-        NOTE: ``min_hnr_db=0.0`` sits below the typical speech-noise floor so the
+        NOTE: ``min_hnr_db=0.0`` sit below the typical speech-noise floor so the
         check still catches catastrophic codec failure (silence, white noise, sample
         scramble all give HNR << 0) while allowing for the sliding-window codec's
         slightly-noisier-than-sync output. The default 1.0 dB threshold has only

tests/e2e/online_serving/test_lingbot_video.py

@@ -0,0 +1,61 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Online serving smoke for ``robbyant/lingbot-video-dense-1.3b``.
+
+The LingBot dense PR supports text-to-video only, so this module exercises the
+native ``/v1/videos`` path with a small 9-frame request.
+"""
+
+import os
+
+import pytest
+
+from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+MODEL = "robbyant/lingbot-video-dense-1.3b"
+PROMPT = "a robotic arm picks up a red block"
+NEGATIVE_PROMPT = "low quality, blurry, watermark, text"
+
+SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
+
+
+def _get_diffusion_feature_cases(model: str):
+    return [
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=["--model-class-name", "LingBotVideoPipeline"],
+            ),
+            id="default",
+            marks=SINGLE_CARD_FEATURE_MARKS,
+        ),
+    ]
+
+
+@pytest.mark.core_model
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
+def test_text_to_video_001(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+    request_config = {
+        "model": omni_server.model,
+        "form_data": {
+            "model": omni_server.model,
+            "prompt": PROMPT,
+            "negative_prompt": NEGATIVE_PROMPT,
+            "height": 192,
+            "width": 320,
+            "num_frames": 9,
+            "fps": 24,
+            "num_inference_steps": 2,
+            "guidance_scale": 3.0,
+            "flow_shift": 3.0,
+            "seed": 42,
+        },
+    }
+    openai_client.send_video_diffusion_request(request_config)

tests/e2e/online_serving/test_lingbot_video_expansion.py

@@ -0,0 +1,82 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+L4 expansion coverage for ``robbyant/lingbot-video-dense-1.3b``.
+
+This is intentionally T2V-only. T2I, I2V, TI2V, MoE, and multi-GPU feature
+rows should be added in follow-up PRs when those paths are implemented.
+"""
+
+import json
+import os
+
+import pytest
+
+from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
+
+MODEL = "robbyant/lingbot-video-dense-1.3b"
+PROMPT = "a robotic arm picks up a red block"
+NEGATIVE_PROMPT = "low quality, blurry, watermark, text"
+
+SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
+
+
+def _get_diffusion_feature_cases(model: str):
+    return [
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=["--model-class-name", "LingBotVideoPipeline"],
+            ),
+            id="default",
+            marks=SINGLE_CARD_FEATURE_MARKS,
+        ),
+    ]
+
+
+@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
+def test_cfg_off(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+    request_config = {
+        "model": omni_server.model,
+        "form_data": {
+            "model": omni_server.model,
+            "prompt": PROMPT,
+            "height": 192,
+            "width": 192,
+            "num_frames": 9,
+            "fps": 24,
+            "num_inference_steps": 2,
+            "guidance_scale": 1.0,
+            "flow_shift": 3.0,
+            "seed": 42,
+        },
+    }
+    openai_client.send_video_diffusion_request(request_config)
+
+
+@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
+def test_batch_cfg_extra_params(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+    request_config = {
+        "model": omni_server.model,
+        "form_data": {
+            "model": omni_server.model,
+            "prompt": PROMPT,
+            "negative_prompt": NEGATIVE_PROMPT,
+            "height": 192,
+            "width": 320,
+            "num_frames": 9,
+            "fps": 24,
+            "num_inference_steps": 2,
+            "guidance_scale": 3.0,
+            "flow_shift": 3.0,
+            "seed": 42,
+            "extra_params": json.dumps({"batch_cfg": True}, separators=(",", ":")),
+        },
+    }
+    openai_client.send_video_diffusion_request(request_config)

tests/e2e/online_serving/test_qwen3_omni.py

@@ -226,3 +226,27 @@ def test_thinker_prefix_caching_audio_output(omni_server, openai_client) -> None
     }
 
     _run_prefix_cache_check(openai_client, request_config)
+
+
+@pytest.mark.advanced_model
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_completions_rejected_for_thinker_talker(omni_server, openai_client) -> None:
+    """Ensure Thinker-talker models reject /v1/completions; we do this because the
+    thinker-talker handoff implementations currently use ChatML <|im_start|> and
+    <|im_end|> markers to segment the input sequence; when we don't have them,
+    the talker does not get any embeddings, which breaks the server.
+    """
+    responses = openai_client.send_completions_http_request(
+        {
+            "json": {
+                "model": omni_server.model,
+                "prompt": "Hello, how are you?",
+                "max_tokens": 10,
+            },
+        },
+        err_code=400,
+    )
+    assert not responses[0].success

tests/entrypoints/openai_api/test_serving_speech.py

@@ -3450,6 +3450,48 @@ class TestTTSAsyncOffloading:
 
         mock_coerce.assert_called_once_with(qwen3_tts_server.engine_client.default_sampling_params_list, True)
 
+    def test_prepare_speech_generation_no_async_chunk_stream_uses_final_only(
+        self, qwen3_tts_server, mocker: MockerFixture
+    ):
+        """Full-payload TTS streaming should not request delta multimodal outputs."""
+        qwen3_tts_server.engine_client.model_config.async_chunk = False
+        qwen3_tts_server._validate_tts_request = mocker.MagicMock(return_value=None)
+        qwen3_tts_server._build_tts_params = mocker.MagicMock(
+            return_value={"text": ["hello"], "task_type": ["CustomVoice"], "speaker": ["Vivian"]}
+        )
+        qwen3_tts_server._estimate_prompt_len_async = mocker.AsyncMock(return_value=512)
+        mock_coerce = mocker.patch(
+            "vllm_omni.entrypoints.openai.serving_speech.coerce_param_message_types",
+            return_value=qwen3_tts_server.engine_client.default_sampling_params_list,
+        )
+        request = OpenAICreateSpeechRequest(input="hello", stream_format="audio", response_format="pcm")
+
+        asyncio.run(qwen3_tts_server._prepare_speech_generation(request))
+
+        mock_coerce.assert_called_once_with(qwen3_tts_server.engine_client.default_sampling_params_list, False)
+
+    def test_prepare_speech_generation_no_async_chunk_stream_keeps_delta_for_non_qwen3(
+        self, voxtral_server, mocker: MockerFixture
+    ):
+        """FINAL_ONLY streaming for async_chunk=False is scoped to qwen3_tts only."""
+        voxtral_server.engine_client.model_config.async_chunk = False
+        mocker.patch.object(voxtral_server._get_tts_adapter(), "validate", return_value=None)
+        voxtral_server._build_voxtral_prompt_async = mocker.AsyncMock(
+            return_value={
+                "prompt_token_ids": [1, 2, 3],
+                "additional_information": {"voice": ["test"]},
+            }
+        )
+        mock_coerce = mocker.patch(
+            "vllm_omni.entrypoints.openai.serving_speech.coerce_param_message_types",
+            return_value=voxtral_server.engine_client.default_sampling_params_list,
+        )
+        request = OpenAICreateSpeechRequest(input="hello", voice="test", stream_format="audio", response_format="pcm")
+
+        asyncio.run(voxtral_server._prepare_speech_generation(request))
+
+        mock_coerce.assert_called_once_with(voxtral_server.engine_client.default_sampling_params_list, True)
+
     def test_prepare_speech_generation_qwen3_voicedesign_non_streaming_mode_false(
         self, qwen3_tts_server, mocker: MockerFixture
     ):

tests/entrypoints/test_omni_entrypoints.py

@@ -107,6 +107,7 @@ class FakeAsyncOmniEngine:
         self.stage_vllm_configs = [None for _ in range(self.num_stages)]
         self.output_processors = [SimpleNamespace(tokenizer=None) for _ in range(self.num_stages)]
         self.input_processor = None
+        self.endpoint_restrictions = ()
 
         self.output_q: queue.Queue[Any] = queue.Queue()
         self.submitted: list[dict[str, Any]] = []

tests/entrypoints/test_omni_sleep_mode.py

@@ -179,7 +179,7 @@ class TestOmniSleepMode:
         try:
             device_id = 0
             used_before = get_device_global_memory_used_gib(device_id)
-            acks = await llm_engine.sleep(stage_ids=[0], level=2)
+            acks = await llm_engine.sleep(stage_ids=[0], level=1)
             await asyncio.sleep(1.5)
             used_after = get_device_global_memory_used_gib(device_id)
             drop_gib = used_before - used_after
@@ -209,7 +209,7 @@ class TestOmniSleepMode:
         """Diffusion Worker stage signal loop"""
         try:
             logger.info("Starting Diffusion Worker Handshake Test")
-            acks = await diffusion_engine.sleep(stage_ids=[0], level=2)
+            acks = await diffusion_engine.sleep(stage_ids=[0], level=1)
 
             def _get_status(ack):
                 return ack.status if hasattr(ack, "status") else ack.get("status")
@@ -264,9 +264,9 @@ class TestOmniSleepMode:
                 base_output = output
             assert base_output is not None and len(base_output.images) > 0
             logger.info("Baseline Generation successful.")
-            # Sleep Level 2
+            # Sleep Level 1
             logger.info("Entering Deep Sleep (VRAM Scavenging)...")
-            await diffusion_engine.sleep(stage_ids=[0], level=2)
+            await diffusion_engine.sleep(stage_ids=[0], level=1)
             # Wake-up
             logger.info("Waking up (Reloading Weights)...")
             await diffusion_engine.wake_up(stage_ids=[0])
@@ -353,8 +353,8 @@ class TestOmniSleepMode:
             logger.info(f"Diffusion Initial VRAM: {vram_initial:.2f} GiB")
 
             # Sleep
-            logger.info("Triggering Level 2 Deep Sleep (Full Weight Offloading)...")
-            acks = await diffusion_engine.sleep(stage_ids=[0], level=2)
+            logger.info("Triggering Level 1 Deep Sleep (Partial Weight Offloading)...")
+            acks = await diffusion_engine.sleep(stage_ids=[0], level=1)
 
             reported_freed_bytes = sum(getattr(ack, "freed_bytes", 0) for ack in acks)
             reported_freed_gib = reported_freed_bytes / 1024**3
@@ -407,6 +407,47 @@ class TestOmniSleepMode:
             diffusion_engine.shutdown()
             await asyncio.sleep(1)
 
+    @pytest.mark.asyncio
+    @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+    async def test_level2_sleep_wake_raises(self, llm_engine: AsyncOmni):
+        """Regression for #4473 Repro A: wake_up() after sleep(level=2) must raise
+        NotImplementedError instead of silently producing corrupted output."""
+        try:
+            await llm_engine.sleep(stage_ids=[0], level=2)
+            with pytest.raises(NotImplementedError, match="sleep\\(level=2\\)"):
+                await llm_engine.wake_up(stage_ids=[0])
+        finally:
+            llm_engine.shutdown()
+
+    @pytest.mark.asyncio
+    @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+    async def test_partial_wake_blocks_generate(self, llm_engine: AsyncOmni):
+        """Regression for #4473 Repro B: generate() must be rejected if kv_cache
+        is still asleep after wake_up(tags=["weights"]), instead of crashing with
+        CUDA illegal memory access."""
+        try:
+            await llm_engine.sleep(stage_ids=[0], level=1)
+            await llm_engine.wake_up(stage_ids=[0], tags=["weights"])
+            with pytest.raises(RuntimeError, match="partially or fully asleep"):
+                async for _ in llm_engine.generate("test", sampling_params=SamplingParams(max_tokens=4)):
+                    pass
+        finally:
+            llm_engine.shutdown()
+
+    @pytest.mark.asyncio
+    @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+    async def test_duplicate_wake_is_idempotent(self, llm_engine: AsyncOmni):
+        """Regression for #4473 Repro C: duplicate wake_up(tags=None) must be a
+        safe no-op instead of raising a cumem CUDA invalid argument error."""
+        try:
+            await llm_engine.sleep(stage_ids=[0], level=1)
+            first_acks = await llm_engine.wake_up(stage_ids=[0])
+            assert len(first_acks) > 0, "First wake_up() should return ACKs"
+            second_acks = await llm_engine.wake_up(stage_ids=[0])
+            assert second_acks == [], f"Duplicate wake_up() should return [] but got {second_acks}"
+        finally:
+            llm_engine.shutdown()
+
 
 # ---------------------------------------------------------------------------
 # H100 (or MI325) e2e: BAGEL / pure diffusion, multi-TP
@@ -444,7 +485,7 @@ async def test_diffusion_model_sleep_tp(tp_size: int):
         async for _ in engine.generate("test", sampling_params_list=[llm_sp, diff_sp]):
             pass
 
-        acks = await engine.sleep(level=2)
+        acks = await engine.sleep(level=1)
         statuses = [get_ack_info(ack, "status") for ack in acks]
         assert all(s == "SUCCESS" for s in statuses), f"Sleep failed. Statuses: {statuses}"
 
@@ -500,7 +541,7 @@ async def test_multistage_sleep_h100(tp_size: int):
         async for _ in engine.generate("warmup", sampling_params_list=[SamplingParams(), sp]):
             pass
 
-        acks = await engine.sleep(stage_ids=[0, 1], level=2)
+        acks = await engine.sleep(stage_ids=[0, 1], level=1)
         assert len(acks) == 2
 
         await engine.wake_up(stage_ids=[0, 1])

tests/entrypoints/test_resolve_dreamzero_config.py

@@ -24,8 +24,8 @@ def test_dreamzero_vla_resolves_to_dreamzero_config(monkeypatch):
 
 def test_dreamzero_config_sets_model_class_and_policy_config(monkeypatch):
     monkeypatch.setattr(
-        "vllm_omni.config.config_factory.StageConfigFactory._auto_detect_model_type",
-        classmethod(lambda _cls, _model, trust_remote_code=True: ("vla", None)),
+        "vllm_omni.config.config_factory.StageConfigFactory._try_infer_model_type",
+        classmethod(lambda _cls, model, trust_remote_code=True: "vla"),
     )
     monkeypatch.setattr(
         "vllm_omni.diffusion.utils.hf_utils._looks_like_dreamzero",

tests/examples/offline_inference/test_image_to_video.py

@@ -0,0 +1,81 @@
+"""
+Offline inference tests: image-to-video.
+See examples/offline_inference/image_to_video/README.md
+"""
+
+import shutil
+from pathlib import Path
+
+import pytest
+import requests
+
+from tests.examples.helpers import EXAMPLES, ExampleRunner, ReadmeSnippet
+from tests.helpers.assertions import assert_video_valid
+from tests.helpers.mark import hardware_marks
+
+pytestmark = [
+    pytest.mark.usefixtures("clean_gpu_memory_between_tests"),
+    pytest.mark.full_model,
+    pytest.mark.example,
+    *hardware_marks(res={"cuda": "H100"}),
+]
+
+I2V_SCRIPT = EXAMPLES / "offline_inference" / "image_to_video" / "image_to_video.py"
+README_PATH = I2V_SCRIPT.with_name("README.md")
+EXAMPLE_OUTPUT_SUBFOLDER = "example_offline_i2v"
+
+_IMAGE_URL = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg"
+_IMAGE_NAME = "cherry_blossom.jpg"
+
+_SKIP_SECTIONS = {
+    "Prerequisites",
+    "Advanced Features",
+    "FAQ",
+    # VACE conditional tasks are covered by test_vace_video_generation.py, which
+    # runs the same README snippets with synthetic assets and smoke settings. The
+    # prep block here downloads inputs via wget (not portable, not collected), so
+    # exercising this section from the shared runner would only duplicate and fail.
+    "Wan2.1 VACE Conditional Tasks",
+}
+
+
+def _ensure_test_image(run_dir: Path) -> None:
+    """Download or copy the example image into run_dir so CLI snippets can find it."""
+    dest = run_dir / _IMAGE_NAME
+    if dest.exists():
+        return
+    src = I2V_SCRIPT.parent / _IMAGE_NAME
+    if src.exists():
+        shutil.copy2(src, dest)
+        return
+    response = requests.get(_IMAGE_URL, timeout=60)
+    response.raise_for_status()
+    dest.write_bytes(response.content)
+
+
+def _skip_readme_snippet(language: str, code: str, h2_title: str) -> tuple[bool, str]:
+    if h2_title in _SKIP_SECTIONS:
+        return True, f"README section '{h2_title}' is intentionally excluded for examples tests"
+    if language == "python":
+        return True, "Python API snippets produce video files that ExampleRunner does not auto-collect"
+    if "/path/to/" in code:
+        return True, "Snippet references a placeholder local model path"
+    return False, ""
+
+
+README_SNIPPETS = ReadmeSnippet.extract_readme_snippets(README_PATH, skipif=_skip_readme_snippet)
+
+
+@pytest.mark.parametrize("snippet", README_SNIPPETS, ids=lambda snippet: snippet.test_id)
+def test_image_to_video(snippet: ReadmeSnippet, example_runner: ExampleRunner):
+    should_skip, reason = snippet.skip
+    if should_skip:
+        pytest.skip(reason)
+
+    run_dir = example_runner.output_root / EXAMPLE_OUTPUT_SUBFOLDER / snippet.test_id
+    run_dir.mkdir(parents=True, exist_ok=True)
+    _ensure_test_image(run_dir)
+
+    result = example_runner.run(snippet, output_subfolder=Path(EXAMPLE_OUTPUT_SUBFOLDER))
+    for asset in result.assets:
+        assert_video_valid(asset)

tests/helpers/assertions.py

@@ -19,19 +19,20 @@ from PIL import Image
 from tests.helpers.media import (
     convert_audio_bytes_to_text,
     cosine_similarity_text,
+    preprocess_text,
 )
 
 _GENDER_PIPELINE = None
 _GENDER_PIPELINE_LOCK = threading.Lock()
 # Transcript gates default to whisper ``small`` for speed. ``small`` mishears a
 # short TTS clip ~0.5% of the time (e.g. "Hello"->"fellow", or hallucinating a
-# leading SFX token), which flakes the deterministic similarity gate. A test can
-# opt in to ASR escalation by setting ``transcript_escalation_model`` to a
-# whisper model name (e.g. ``"large-v3"``) in its request_config: on a failed
-# fast pass the clip is re-transcribed with that stronger ASR before the test
-# fails, so a weak-ASR mishear is rescued while a genuine model artifact still
-# fails (the strong ASR mismatches too). Tests that do not opt in keep the
-# strict behaviour.
+# leading SFX token), which flakes the deterministic similarity gate. Short
+# clips first get a conservative containment fallback for minor ASR repeats/noise.
+# A test can also opt in to ASR escalation by setting
+# ``transcript_escalation_model`` to a whisper model name (e.g. ``"large-v3"``)
+# in its request_config: on a failed fast pass the clip is re-transcribed with
+# that stronger ASR before the test fails, so a weak-ASR mishear is rescued while
+# a genuine model artifact still fails (the strong ASR mismatches too).
 _PCM_SPEECH_SAMPLE_RATE_HZ = 24_000
 _MIN_PCM_SPEECH_HNR_DB = 1.0
 _PRESET_VOICE_GENDER_MAP: dict[str, str] = {
@@ -43,6 +44,23 @@ _PRESET_VOICE_GENDER_MAP: dict[str, str] = {
 }
 
 
+def _short_transcript_contains_expected(transcript: str, expected: str) -> bool:
+    """Allow minor ASR repeats/noise for very short speech clips."""
+    transcript_clean = preprocess_text(transcript)
+    expected_clean = preprocess_text(expected)
+    if not transcript_clean or not expected_clean:
+        return False
+
+    transcript_words = transcript_clean.split()
+    expected_words = expected_clean.split()
+    if not transcript_words or not expected_words:
+        return False
+
+    short_text = min(len(transcript_clean), len(expected_clean)) <= 15
+    small_word_delta = len(transcript_words) <= len(expected_words) + 2
+    return short_text and small_word_delta and expected_clean in transcript_clean
+
+
 def assert_image_diffusion_response(
     response: "DiffusionResponse",
     request_config: dict[str, Any],
@@ -620,13 +638,16 @@ def _assert_transcript_matches(
     ``transcript`` is the fast whisper-``small`` result. If it clears
     ``threshold`` the check passes immediately.
 
+    If the cosine check fails, very short clips get a conservative containment
+    fallback that accepts minor ASR repeats/noise only when the expected text is
+    still present and the transcript has few extra words.
+
     When ``escalation_model`` is set (opt-in via the ``transcript_escalation_model``
-    request_config key) and the fast check fails, the clip is re-transcribed with
-    that stronger ASR and the assertion is decided on its verdict -- so a weak
-    whisper-``small`` mishear on a short clip does not flake the gate, while a
-    genuine model artifact still fails (the strong ASR mismatches too). When
-    ``escalation_model`` is ``None`` the original strict behaviour is preserved,
-    so other tests are unaffected.
+    request_config key) and the fast check plus containment fallback fail, the
+    clip is re-transcribed with that stronger ASR and the assertion is decided on
+    its verdict -- so a weak whisper-``small`` mishear on a short clip does not
+    flake the gate, while a genuine model artifact still fails (the strong ASR
+    mismatches too).
     """
     expected = str(expected_text).strip().lower()
     similarity = cosine_similarity_text(transcript.strip().lower(), expected)
@@ -634,6 +655,10 @@ def _assert_transcript_matches(
     if similarity > threshold:
         return
 
+    if _short_transcript_contains_expected(transcript, expected):
+        print("short speech containment check passed")
+        return
+
     if escalation_model and audio_bytes:
         print(
             f"whisper-small below threshold ({similarity:.2f} <= {threshold}); "

tests/helpers/runtime.py

@@ -47,6 +47,7 @@ from tests.helpers.media import (
     _merge_base64_audio_to_segment,
     decode_b64_image,
 )
+from tests.model_tests.diffusion.utils import resolve_tiny_model_path
 from vllm_omni.config.stage_config import resolve_deploy_yaml
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt
 from vllm_omni.outputs import OmniRequestOutput
@@ -1064,6 +1065,29 @@ class OpenAIClientHandler:
         )
         return [resp]
 
+    def send_completions_http_request(
+        self,
+        request_config: dict[str, Any],
+        *,
+        err_code: int | tuple[int, ...] | list[int] | None = None,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+    ) -> list[HttpResponse]:
+        """POST ``/v1/completions`` with ``json`` or ``raw_body``."""
+        # TODO (Alex): A lot of these helpers should be consolidated as they differ only by endpoint
+        cfg = _merge_http_expectation_kwargs(
+            request_config,
+            err_code=err_code,
+            err_message=err_message,
+        )
+        r = self._post_json_endpoint("/v1/completions", cfg, default_timeout=120.0)
+        resp = self._http_response_from_requests(r)
+        assert_http_error(
+            resp,
+            err_code=cfg.get("err_code"),
+            err_message=cfg.get("err_message"),
+        )
+        return [resp]
+
     def send_omni_sleep_http_request(
         self,
         request_config: dict[str, Any],
@@ -2792,11 +2816,24 @@ def iter_omni_server(
 
     with omni_fixture_lock:
         params: OmniServerParams = request.param
-        model = model_prefix + params.model
+        # For now, when a tiny model is substituted, we preserve the original model
+        # name via --served-model-name (so that the server still accepts requests with
+        # the original name). We also do the same for server.model so that tests reading
+        # server.model send the correct name in requests.
+        #
+        # TODO: core models on this path currently do not clean up tiny models, although
+        # tiny model paths are deterministic, so it's not a huge footprint. Still, it would
+        # be ideal to cleanup consistently everywhere.
+        original_model = model_prefix + params.model
+        model = original_model
+        if run_level == "core_model":
+            model = resolve_tiny_model_path(model)
         port = params.port
         stage_config_path = stage_config_path_for_run_level(params.stage_config_path, run_level)
 
         server_args = params.server_args or []
+        if model != original_model:
+            server_args = [*server_args, "--served-model-name", original_model]
         if params.use_omni and params.stage_init_timeout is not None:
             server_args = [*server_args, "--stage-init-timeout", str(params.stage_init_timeout)]
         else:
@@ -2822,6 +2859,8 @@ def iter_omni_server(
                 port=port,
                 env_dict=params.env_dict,
             ) as server:
+                if model != original_model:
+                    server.model = original_model
                 print("OmniServer started successfully")
                 yield server
                 print("OmniServer stopping...")
@@ -2845,6 +2884,8 @@ def iter_omni_server(
                     use_omni=params.use_omni,
                 )
             ) as server:
+                if model != original_model:
+                    server.model = original_model
                 print("OmniServer started successfully")
                 yield server
                 print("OmniServer stopping...")
@@ -2876,6 +2917,8 @@ def iter_omni_runner(
             extra_omni_kwargs = dict(extra) if extra is not None else {}
         stage_config_path = stage_config_path_for_run_level(stage_config_path, run_level)
         model = model_prefix + model
+        if run_level == "core_model":
+            model = resolve_tiny_model_path(model)
         with OmniRunner(model, seed=42, stage_configs_path=stage_config_path, **extra_omni_kwargs) as runner:
             print("OmniRunner started successfully")
             yield runner

tests/helpers/test_assertions.py

@@ -0,0 +1,25 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+from tests.helpers.assertions import _assert_transcript_matches
+
+
+def test_short_transcript_repeat_passes_containment_fallback():
+    _assert_transcript_matches(
+        " How... how are you?",
+        audio_bytes=None,
+        expected_text="how are you",
+        threshold=0.9,
+    )
+
+
+def test_short_transcript_unrelated_text_still_fails():
+    with pytest.raises(AssertionError, match="Transcript doesn't match input"):
+        _assert_transcript_matches(
+            " I don't know, sorry.",
+            audio_bytes=None,
+            expected_text="how are you",
+            threshold=0.9,
+        )

tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py

@@ -339,6 +339,91 @@ class TestCodePredictorDtypeAlignment:
         assert not torch.equal(first[:, 1:], different[:, 1:])
 
 
+class TestCodePredictorPerRowGenerators:
+    """Seeded requests must stay deterministic inside a multi-row batch (#4883)."""
+
+    def _make_predictor(self, mocker: MockerFixture, loaded_target_classes):
+        _, _, code_predictor_wrapper, _, _ = loaded_target_classes
+        common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
+        mocker.patch.object(common_mod.current_omni_platform, "is_npu", return_value=False)
+        cp_config, talker_config = _make_tiny_config(loaded_target_classes)
+        vllm_config = _make_vllm_config(mocker, max_num_seqs=4)
+        predictor = code_predictor_wrapper(
+            vllm_config=vllm_config,
+            config=cp_config,
+            talker_config=talker_config,
+        )
+        predictor._wrapper_config.use_cuda_graphs = False
+        return predictor, talker_config
+
+    def test_multinomial_per_row_matches_single_row_draws(self, mocker: MockerFixture, loaded_target_classes) -> None:
+        common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
+        multinomial = common_mod.CodePredictorWrapper._multinomial
+        torch.manual_seed(7)
+        probs = torch.softmax(torch.randn(3, 64), dim=-1)
+
+        def seeded(seed: int) -> torch.Generator:
+            generator = torch.Generator(device=probs.device)
+            generator.manual_seed(seed)
+            return generator
+
+        batched = multinomial(probs, None, [seeded(11), None, seeded(22)])
+        assert batched.shape == (3, 1)
+        # Each seeded row must reproduce a standalone draw from the same seed,
+        # independent of what the other rows in the batch consume.
+        assert torch.equal(batched[0:1], torch.multinomial(probs[0:1], num_samples=1, generator=seeded(11)))
+        assert torch.equal(batched[2:3], torch.multinomial(probs[2:3], num_samples=1, generator=seeded(22)))
+
+    def test_forward_per_row_generators_are_row_independent(self, mocker: MockerFixture, loaded_target_classes) -> None:
+        predictor, talker_config = self._make_predictor(mocker, loaded_target_classes)
+        hidden = talker_config.hidden_size
+        torch.manual_seed(123)
+        row_embed = torch.randn(1, hidden)
+        row_hidden = torch.randn(1, hidden)
+        # Three identical rows: two share a seed, one differs.
+        layer0_code = torch.zeros(3, dtype=torch.long)
+        layer0_embed = row_embed.expand(3, hidden).contiguous()
+        last_talker_hidden = row_hidden.expand(3, hidden).contiguous()
+
+        def seeded(seed: int) -> torch.Generator:
+            generator = torch.Generator(device=layer0_code.device)
+            generator.manual_seed(seed)
+            return generator
+
+        def run(generators):
+            return predictor(
+                layer0_code=layer0_code,
+                layer0_embed=layer0_embed,
+                last_talker_hidden=last_talker_hidden,
+                do_sample=True,
+                temperature=0.9,
+                top_k=50,
+                top_p=1.0,
+                generators=generators,
+            )
+
+        first = run([seeded(1234), seeded(1234), seeded(4321)])
+        second = run([seeded(1234), seeded(1234), seeded(4321)])
+
+        # Same per-row seeds -> the whole batch reproduces across calls.
+        assert torch.equal(first, second)
+        # Identical inputs + identical seeds -> identical rows within a call.
+        assert torch.equal(first[0], first[1])
+        # A different seed on the same inputs must diverge in the residual layers.
+        assert not torch.equal(first[0, 1:], first[2, 1:])
+
+    def test_forward_rejects_mismatched_generators_length(self, mocker: MockerFixture, loaded_target_classes) -> None:
+        predictor, talker_config = self._make_predictor(mocker, loaded_target_classes)
+        hidden = talker_config.hidden_size
+        with pytest.raises(ValueError, match="one entry per row"):
+            predictor(
+                layer0_code=torch.zeros(2, dtype=torch.long),
+                layer0_embed=torch.randn(2, hidden),
+                last_talker_hidden=torch.randn(2, hidden),
+                generators=[torch.Generator()],
+            )
+
+
 class TestCodePredictorModelDtype:
     """Test the inner model forward with different dtypes."""
 

tests/model_executor/models/voxcpm2/test_runner_assisted_unified_graph_static.py

@@ -12,6 +12,7 @@ REPO_ROOT = Path(__file__).resolve().parents[4]
 RUNNER = REPO_ROOT / "vllm_omni" / "worker" / "gpu_ar_model_runner.py"
 RUNNER_ASSISTED_METADATA = REPO_ROOT / "vllm_omni" / "worker" / "runner_assisted_metadata.py"
 TALKER = REPO_ROOT / "vllm_omni" / "model_executor" / "models" / "voxcpm2" / "voxcpm2_talker.py"
+VOXCPM2_RUNTIME_CONFIG = REPO_ROOT / "vllm_omni" / "model_executor" / "models" / "voxcpm2" / "runtime_config.py"
 VOXCPM2_PIPELINE = REPO_ROOT / "vllm_omni" / "model_executor" / "models" / "voxcpm2" / "pipeline.py"
 VOXCPM2_SCHEDULER = REPO_ROOT / "vllm_omni" / "model_executor" / "models" / "voxcpm2" / "scheduler.py"
 OMNI_AR_SCHEDULER = REPO_ROOT / "vllm_omni" / "core" / "sched" / "omni_ar_scheduler.py"
@@ -74,6 +75,7 @@ def test_ar_runner_without_model_hook_stays_on_normal_path():
 
 def test_voxcpm2_graph_paths_fail_closed_and_preserve_deterministic_noise():
     source = TALKER.read_text()
+    runtime_config_source = VOXCPM2_RUNTIME_CONFIG.read_text()
     compact_source = "".join(source.split())
 
     assert "_voxcpm2_compile_without_inductor_cudagraphs" in source
@@ -86,8 +88,8 @@ def test_voxcpm2_graph_paths_fail_closed_and_preserve_deterministic_noise():
     assert '"triton.cudagraphs": False' in compile_source
     assert '"triton.cudagraph_trees": False' in compile_source
 
-    assert "self._enable_unified_decode_graph=(use_cuda_graph" in compact_source
-    assert "andnotself._deterministic_cfm_noise" in compact_source
+    assert "self._enable_unified_decode_graph=self._runtime_config.unified_decode_graph_available(" in compact_source
+    assert "andnotself.deterministic_cfm_noise" in "".join(runtime_config_source.split())
     assert "decode_tail_graph" not in source
 
     unified_source = source[source.index("def _forward_unified_decode") :]
@@ -183,7 +185,9 @@ def test_voxcpm2_scheduler_policy_stays_model_local():
     assert "def _should_defer_waiting_admission(self) -> bool:" in common_source
 
     assert "class VoxCPM2OmniARAsyncScheduler(OmniARAsyncScheduler)" in voxcpm2_scheduler_source
-    assert "voxcpm2_runtime_config" in voxcpm2_scheduler_source
+    assert "from .runtime_config import _VoxCPM2RuntimeConfig" in voxcpm2_scheduler_source
+    assert "_VoxCPM2RuntimeConfig.from_vllm_config(self.vllm_config)" in voxcpm2_scheduler_source
+    assert "unified_decode_graph_available(use_cuda_graph=current_omni_platform.is_cuda())" in voxcpm2_scheduler_source
     assert "_should_defer_waiting_for_unified_decode_graph" in voxcpm2_scheduler_source
     assert "def schedule(" not in voxcpm2_scheduler_source
     assert "create_request_queue" not in voxcpm2_scheduler_source

tests/model_executor/models/voxtral_tts/test_tts_preprocess.py

@@ -0,0 +1,78 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import functools
+
+import pytest
+import torch
+import torch.nn as nn
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@functools.lru_cache(maxsize=1)
+def _voxtral_tts_model_cls():
+    from tests.model_executor.helpers import bootstrap_vllm_layer_custom_op_modules
+
+    bootstrap_vllm_layer_custom_op_modules()
+    import vllm.model_executor.models.utils  # noqa: F401
+
+    from vllm_omni.model_executor.models.voxtral_tts.voxtral_tts import (
+        VoxtralTTSForConditionalGeneration,
+    )
+
+    return VoxtralTTSForConditionalGeneration
+
+
+class FakeAudioGeneration(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.embed_multimodal_calls = []
+
+    def embed_multimodal(self, **kwargs):
+        self.embed_multimodal_calls.append(kwargs)
+        return [torch.full((1, 4), 7.0)]
+
+
+def _make_voxtral_tts_model():
+    model_cls = _voxtral_tts_model_cls()
+    model = model_cls.__new__(model_cls)
+    nn.Module.__init__(model)
+    model.model_stage = "audio_generation"
+    model.model = FakeAudioGeneration()
+    model._audio_token_id = 42
+    return model
+
+
+def test_tts_preprocess_consumes_nested_codes_audio_feedback():
+    model = _make_voxtral_tts_model()
+    input_ids = torch.tensor([model._audio_token_id])
+    input_embeds = torch.zeros((1, 4))
+    audio_tokens = torch.tensor([[1, 2, 3, 4]])
+
+    _, output_embeds, _ = model.tts_preprocess(
+        input_ids=input_ids,
+        input_embeds=input_embeds,
+        codes={"audio": audio_tokens},
+    )
+
+    assert len(model.model.embed_multimodal_calls) == 1
+    torch.testing.assert_close(model.model.embed_multimodal_calls[0]["audio_tokens"], audio_tokens)
+    torch.testing.assert_close(output_embeds, torch.full((1, 4), 7.0))
+
+
+def test_tts_preprocess_keeps_legacy_top_level_audio_feedback():
+    model = _make_voxtral_tts_model()
+    input_ids = torch.tensor([model._audio_token_id])
+    input_embeds = torch.zeros((1, 4))
+    audio_tokens = torch.tensor([[5, 6, 7, 8]])
+
+    _, output_embeds, _ = model.tts_preprocess(
+        input_ids=input_ids,
+        input_embeds=input_embeds,
+        audio=audio_tokens,
+    )
+
+    assert len(model.model.embed_multimodal_calls) == 1
+    torch.testing.assert_close(model.model.embed_multimodal_calls[0]["audio_tokens"], audio_tokens)
+    torch.testing.assert_close(output_embeds, torch.full((1, 4), 7.0))

tests/model_extras/test_model_extras.py

@@ -126,6 +126,102 @@ def test_magi_human_extra_registry_declares_request_and_response_params() -> Non
     assert should_init_extra_args_for_non_diffusion_stages("MagiHumanPipeline") is False
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_ming_flash_omni_extra_registry_declares_request_and_response_params() -> None:
+    assert get_extra_body_params("MingImagePipeline") == frozenset(
+        {
+            "height",
+            "width",
+            "steps",
+            "cfg",
+            "seed",
+            "byte5_text",
+            "negative_prompt",
+        }
+    )
+    assert get_extra_output_params("MingImagePipeline") == frozenset()
+    assert should_init_extra_args_for_non_diffusion_stages("MingImagePipeline") is True
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_ming_flash_omni_declared_extra_args_route_into_sampling_params() -> None:
+    params = OmniDiffusionSamplingParams()
+    declared = get_extra_body_params("MingImagePipeline")
+    apply_declared_extra_args(
+        params,
+        declared,
+        {
+            "steps": 6,
+            "cfg": 1.5,
+            "byte5_text": ["理解与生成统一"],
+            "negative_prompt": "ugly, blurry",
+            "unknown": "ignored",
+        },
+    )
+    assert params.extra_args == {
+        "steps": 6,
+        "cfg": 1.5,
+        "byte5_text": ["理解与生成统一"],
+        "negative_prompt": "ugly, blurry",
+    }
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_ming_flash_omni_text_to_image_prompt_builder() -> None:
+    assert build_text_to_image_prompt(
+        "MingImagePipeline",
+        prompt="Please draw a cute cat.",
+        negative_prompt="ugly, blurry",
+        height=512,
+        width=768,
+    ) == {
+        "prompt": "Please draw a cute cat.",
+        "modalities": ["image"],
+        "mm_processor_kwargs": {
+            "modalities": ["image"],
+            "target_h": 512,
+            "target_w": 768,
+        },
+        "negative_prompt": "ugly, blurry",
+    }
+    # target_h/w are omitted when height/width are not supplied so the
+    # pipeline's 1024x1024 default still applies.
+    assert build_text_to_image_prompt(
+        "MingImagePipeline",
+        prompt="Draw a poster.",
+        negative_prompt=None,
+    ) == {
+        "prompt": "Draw a poster.",
+        "modalities": ["image"],
+        "mm_processor_kwargs": {"modalities": ["image"]},
+    }
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_ming_flash_omni_image_to_image_prompt_builder() -> None:
+    dummy_image = Image.new("RGB", (64, 64))
+    result = build_image_to_image_prompt(
+        "MingImagePipeline",
+        prompt="Change the background to a sandy beach at sunset.",
+        negative_prompt=None,
+        input_image=dummy_image,
+        height=256,
+        width=256,
+    )
+    assert result["modalities"] == ["img2img"]
+    assert result["multi_modal_data"]["img2img"] is dummy_image
+    assert result["mm_processor_kwargs"] == {
+        "modalities": ["img2img"],
+        "target_h": 256,
+        "target_w": 256,
+    }
+    assert "negative_prompt" not in result
+
+
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_cosmos3_text_to_image_prompt_builder_selects_image_modality() -> None:

tests/model_tests/conftest.py

@@ -0,0 +1,41 @@
+from shutil import rmtree
+
+import pytest
+
+from tests.model_tests.diffusion.model_settings import DIFFUSION_TEST_SETTINGS
+
+
+@pytest.fixture(scope="session")
+def tiny_model_paths(request, run_level):
+    """Build or download the tiny models for the selected tests.
+
+    At core_model level, builds tiny models via the builder function.
+    At advanced_model / full_model level, uses the real HF model path.
+
+    NOTE: this is session scoped to avoid churn in tiny model creation,
+    but will ensure all the tiny models you need are created for the selected tests
+    before it starts to execute them."""
+    model_paths = {}
+    built_paths = []
+    print("Initializing tiny models...")
+    for item in request.session.items:
+        if not hasattr(item, "callspec"):
+            continue
+        model_name = item.callspec.params.get("model_name")
+        if model_name is None:
+            continue
+        if model_name not in model_paths:
+            settings = DIFFUSION_TEST_SETTINGS[model_name]
+            if run_level == "core_model":
+                print(f"Calling tiny model builder for: {model_name}")
+                path = settings.builder()
+                built_paths.append(path)
+            else:
+                print(f"Run level is {run_level}; {model_name} will use the full model")
+                path = settings.model
+            model_paths[model_name] = path
+
+    yield model_paths
+    for path in built_paths:
+        print(f"Removing tiny model: {path}")
+        rmtree(path, ignore_errors=True)

tests/model_tests/diffusion/case_filtering.py

@@ -0,0 +1,104 @@
+"""
+Analogous to: https://github.com/vllm-project/vllm/blob/v0.23.0/tests/models/multimodal/generation/vlm_utils/case_filtering.py
+
+This follows a similar pattern so that we can use `get_parametrized_options` to expand our test
+configurations out, similar to the way in which vLLM expands supported input type tests. For now,
+this mostly just means parametrizing over the test groups, which may be unique per model, but
+doing so in this way lets us cleanly define test groups that may be heterogeneous + avoid having
+excessive skip marks for features that may not be supported in every pipeline.
+"""
+
+# Engine initialization is currently quite expensive. For efficiency,
+# when we run the tests, we should try to collapse compatible accelerations,
+# because even with tiny models, the overhead of spinning up a new instance
+# stacks up fast.
+import itertools
+
+import pytest
+
+from tests.helpers.mark import hardware_marks
+from tests.model_tests.diffusion.config_types import DiffusionAccs, DiffusionModelTestOpts, get_required_device_count
+from vllm_omni.platforms import current_omni_platform
+
+# These tests are intended to run on L4 GPUs since they test parallelism (i.e., need GPUs),
+# but are intended to be smoke tests that we can properly validate without real model
+# weights in the CI.
+#
+# As such, no DiffusionAccs may require more than 4 GPUs, since we would otherwise skip
+# them in the CI.
+MAX_CI_DEVICES = 4
+
+
+def get_test_group_marks(model_name: str, test_group: list[DiffusionAccs] | None, model_marks: list | None) -> list:
+    """Build the full set of pytest marks for a test group.
+
+    For now, single device groups default to core model, and multi device groups
+    default to advanced_model marks. We also validate and throw if a test group exceeds
+    the max number of GPUs for the CI to prevent accidentally adding configurations that
+    would be skipped in the CI."""
+    marks = list(model_marks) if model_marks is not None else []
+
+    required_devices = get_required_device_count(test_group)
+    if required_devices > MAX_CI_DEVICES:
+        raise ValueError(
+            f"Test group {test_group} for {model_name} requires {required_devices} devices, "
+            f"but the max CI device count is {MAX_CI_DEVICES}. "
+        )
+
+    marks.extend(hardware_marks(res={"cuda": "L4"}, num_cards=required_devices))
+    # hardware_marks only adds skipif for num_cards > 1, so we currently handle the single-device case
+    # directly here. This should probably be handled in a more common way later on.
+    assert current_omni_platform is not None and current_omni_platform.device_count is not None
+    device_count = current_omni_platform.device_count()
+    if current_omni_platform.is_cuda() and device_count < required_devices:
+        marks.append(pytest.mark.skip(reason=f"Need {required_devices} devices, got {device_count}"))
+    if required_devices > 1:
+        marks.append(pytest.mark.full_model)
+    else:
+        marks.append(pytest.mark.core_model)
+    return marks
+
+
+def get_model_parametrization(model_name: str, test_info: DiffusionModelTestOpts, online: bool):
+    """Given a model & its corresponding test options, build the list of pytest params
+    to be run for this model. The base case (no accelerations) is always included. Extra
+    test groups are always appended for offline tests, but are only added to the online
+    tests if online_base_only is overridden to False to avoid redundant testing.
+    """
+    test_groups: list[list[DiffusionAccs] | None] = [None]
+    if not (online and test_info.online_base_only) and test_info.extra_test_groups:
+        test_groups.extend(test_info.extra_test_groups)
+
+    return [
+        pytest.param(
+            model_name,
+            test_group,
+            test_info.supported_tasks,
+            test_info.check_multi_output and test_group is None,
+            test_info.check_determinism and test_group is None,
+            id=f"{model_name}[{'+'.join(test_group)}]" if test_group else model_name,
+            marks=get_test_group_marks(model_name, test_group, test_info.marks),
+        )
+        for test_group in test_groups
+    ]
+
+
+def get_parametrized_options(
+    test_settings: dict[str, DiffusionModelTestOpts],
+    online: bool = False,
+):
+    """Converts all the DiffusionModelTestOpts into an expanded list of parameters
+    based on which accelerations are available.
+
+    When online=True and a model has online_base_only=True, only the base case
+    (no accelerations) is included for that model's parametrization to minimize
+    redundancy between offline / online tests.
+    """
+    # Get a list per model type, where each entry contains a tuple of all of
+    # that model type's cases, then flatten them into the top level so that
+    # we can consume them in one mark.parametrize call.
+    parametrization = [
+        get_model_parametrization(model_name, test_info, online=online)
+        for model_name, test_info in test_settings.items()
+    ]
+    return list(itertools.chain(*parametrization))

tests/model_tests/diffusion/config_types.py

@@ -0,0 +1,187 @@
+"""
+Common definitions for controlling what tests run where.
+"""
+
+from collections.abc import Callable
+from enum import StrEnum, auto
+from typing import NamedTuple, TypeAlias
+
+from pytest import MarkDecorator
+
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.entrypoints.omni import Omni
+
+# All builder funcs take no params and return a path
+TinyDiffusionBuilder: TypeAlias = Callable[[], str]
+
+
+class DiffusionAccs(StrEnum):
+    """Supported acceleration types / test settings for Diffusion Models."""
+
+    HSDP = auto()
+    TEA_CACHE = auto()
+    CACHE_DIT = auto()
+    SEQUENCE_PARALLEL = auto()
+    CFG_PARALLEL = auto()
+    TENSOR_PARALLEL = auto()
+    CPU_OFFLOAD = auto()
+    LAYERWISE_OFFLOAD = auto()
+    VAE_PATCH_PARALLEL = auto()
+
+
+class DiffusionTasks(StrEnum):
+    """Supported tasks for Diffusion Models."""
+
+    TEXT_TO_IMAGE = auto()
+    IMAGE_TO_IMAGE = auto()
+    # Text to video, text to audio, etc should be added here as needed
+
+
+class DiffusionModelTestOpts(NamedTuple):
+    """Configuration for one Diffusion model's tests."""
+
+    # HF model name for real-weight tests (advanced_model / full_model level).
+    # For now, whether we use the real weights vs tiny weights in the common tests
+    # depends on the run level.
+    model: str
+
+    # Creates a tiny model for the given architecture. We should always use tiny
+    # model weights for tests that do not require us to check the model quality.
+    builder: TinyDiffusionBuilder
+
+    # Actual tasks which controls the tests actually run
+    supported_tasks: list[DiffusionTasks]
+
+    # Additional acceleration groups to run beyond the base case (no acceleration).
+    # The base case is always run for every model in the test settings. None means
+    # we only run the base case.
+    extra_test_groups: list[list[DiffusionAccs]] | None = None
+
+    # Pytest Marks for this model. This may be useful for selecting which models
+    # we want to run where, similar to the way vLLM's multimodal tests mark some
+    # as core models to always run in the CI.
+    # Example: https://github.com/vllm-project/vllm/blob/v0.23.0/tests/models/multimodal/generation/test_common.py#L131
+    marks: list[MarkDecorator] | None = None
+
+    # When True (default), online tests only run the base case (no accelerations).
+    # When False, online tests run all test_groups, same as offline. This should be
+    # True unless there is a good reason for it not to be, because the execution of
+    # the acceleration should be the same on both codepaths, and CLI parsing etc should
+    # be tested by tests adding the acceleration, and not per model.
+    online_base_only: bool = True
+
+    # Additional checks to run for the base case.
+    check_multi_output: bool = True  # Runs multiple generations in one request
+    check_determinism: bool = True  # Runs 2 generations with the same seed and check determinism
+
+
+### Mappings & utils for building offline Omni() instances given a list of enabled accelerations
+ACC_OMNI_KWARGS = {
+    DiffusionAccs.VAE_PATCH_PARALLEL: {"vae_use_tiling": True},
+    DiffusionAccs.CPU_OFFLOAD: {"enable_cpu_offload": True},
+    DiffusionAccs.LAYERWISE_OFFLOAD: {"enable_layerwise_offload": True},
+    DiffusionAccs.CACHE_DIT: {"cache_backend": "cache_dit"},
+    DiffusionAccs.TEA_CACHE: {"cache_backend": "tea_cache"},
+}
+
+ACC_PARALLEL_KWARGS = {
+    DiffusionAccs.HSDP: {"use_hsdp": True, "hsdp_shard_size": 2},
+    DiffusionAccs.TENSOR_PARALLEL: {"tensor_parallel_size": 2},
+    DiffusionAccs.CFG_PARALLEL: {"cfg_parallel_size": 2},
+    DiffusionAccs.VAE_PATCH_PARALLEL: {"vae_patch_parallel_size": 2},
+    # For SP, we don't run ring here to conserve devices, since it is easy
+    # to blow up the number of needed devices fast. Compatibility for ring
+    # and ulysses together should be tested generically outside of these tests.
+    DiffusionAccs.SEQUENCE_PARALLEL: {"ulysses_degree": 2},
+}
+
+ACC_DEVICE_COUNT_KEYS = {
+    DiffusionAccs.TENSOR_PARALLEL: "tensor_parallel_size",
+    DiffusionAccs.CFG_PARALLEL: "cfg_parallel_size",
+    DiffusionAccs.SEQUENCE_PARALLEL: "ulysses_degree",
+    DiffusionAccs.VAE_PATCH_PARALLEL: "vae_patch_parallel_size",
+    DiffusionAccs.HSDP: "hsdp_shard_size",
+}
+
+### CLI args for launching an OmniServer subprocess per acceleration.
+# Follows the same pattern as e2e online serving tests, which hand-code
+# their server_args lists in OmniServerParams.
+ACC_SERVER_ARGS: dict[DiffusionAccs, list[str]] = {
+    DiffusionAccs.HSDP: ["--use-hsdp", "--hsdp-shard-size", "2"],
+    DiffusionAccs.TEA_CACHE: ["--cache-backend", "tea_cache"],
+    DiffusionAccs.CACHE_DIT: ["--cache-backend", "cache_dit"],
+    DiffusionAccs.SEQUENCE_PARALLEL: ["--usp", "2"],
+    DiffusionAccs.CFG_PARALLEL: ["--cfg-parallel-size", "2"],
+    DiffusionAccs.TENSOR_PARALLEL: ["--tensor-parallel-size", "2"],
+    DiffusionAccs.CPU_OFFLOAD: ["--enable-cpu-offload"],
+    DiffusionAccs.LAYERWISE_OFFLOAD: ["--enable-layerwise-offload"],
+    DiffusionAccs.VAE_PATCH_PARALLEL: ["--vae-use-tiling", "--vae-patch-parallel-size", "2"],
+}
+
+## TODO ^ These should be in the same object, it's getting messy.
+
+
+def get_required_device_count(accelerations: list[DiffusionAccs] | None) -> int:
+    """Compute the minimum number of devices needed for a set of accelerations.
+    The total is the product of all parallel dimensions (defaulting to 1).
+
+    If not enough devices are available for a test group's accelerations,
+    that test will be skipped."""
+    count = 1
+    if accelerations is None:
+        return count
+
+    for acc in accelerations:
+        key = ACC_DEVICE_COUNT_KEYS.get(acc)
+        if key is not None:
+            count *= ACC_PARALLEL_KWARGS[acc][key]
+    return count
+
+
+def build_parallel_config_from_diff_accelerations(accelerations: list[DiffusionAccs]) -> DiffusionParallelConfig | None:
+    """Given a list of accelerations pertaining to the current test group,
+    build the parallel config needed for the Omni() object (if any)."""
+    config_kwargs = {}
+    for acc in accelerations:
+        update_dict = ACC_PARALLEL_KWARGS.get(acc, {})
+        config_kwargs.update(update_dict)
+    if config_kwargs:
+        return DiffusionParallelConfig(**config_kwargs)
+    return None
+
+
+### Offline Omni() object builder
+def build_omni_from_diff_accelerations(accelerations: list[DiffusionAccs] | None, **kwargs) -> Omni:
+    """Given one or more acceleration types, build the corresponding Omni() object."""
+    # Coerce to a list and build the parallel config, since that depends on the accelerations
+    if accelerations is None:
+        accelerations = []
+    parallel_config = build_parallel_config_from_diff_accelerations(accelerations)
+
+    # Then add anything else that's a top-level kwarg
+    acc_kwargs = {}
+    if parallel_config is not None:
+        acc_kwargs["parallel_config"] = parallel_config
+    for acc in accelerations:
+        update_dict = ACC_OMNI_KWARGS.get(acc, {})
+        acc_kwargs.update(update_dict)
+
+    # Keys passed through should mostly be things like enforce_eager;
+    # if there's overlap, it's probably due to a misconfiguration
+    shared_keys = acc_kwargs.keys() & kwargs.keys()
+    if shared_keys:
+        raise ValueError(f"Explicit Omni kwargs and inferred Omni kwargs for accelerations overlap: {shared_keys}")
+    omni_kwargs = {**acc_kwargs, **kwargs}
+    return Omni(**omni_kwargs)
+
+
+### Online server flag builder
+def build_server_args_from_diff_accelerations(accelerations: list[DiffusionAccs] | None) -> list[str]:
+    """Given one or more acceleration types, build the corresponding CLI args
+    for launching an OmniServer subprocess."""
+    if accelerations is None:
+        return []
+    args = []
+    for acc in accelerations:
+        args.extend(ACC_SERVER_ARGS.get(acc, []))
+    return args

tests/model_tests/diffusion/diff_model_builders.py

@@ -0,0 +1,63 @@
+import os
+import tempfile
+
+import torch
+from diffusers.models.autoencoders.autoencoder_kl_flux2 import AutoencoderKLFlux2
+from diffusers.models.transformers.transformer_flux2 import Flux2Transformer2DModel
+from diffusers.pipelines.flux2.pipeline_flux2_klein import Flux2KleinPipeline
+from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
+from transformers import AutoTokenizer, Qwen3Config, Qwen3ForCausalLM
+
+TINY_MODEL_DIR = os.path.join(tempfile.gettempdir(), "vllm-omni-tiny-models")
+
+
+def _get_tiny_model_path(name: str) -> str:
+    path = os.path.join(TINY_MODEL_DIR, name)
+    os.makedirs(path, exist_ok=True)
+    return path
+
+
+def tiny_flux2_klein_builder() -> str:
+    """Build a tiny Flux2Klein model."""
+    model_dir = _get_tiny_model_path("Flux2KleinPipeline")
+
+    pipe = Flux2KleinPipeline(
+        scheduler=FlowMatchEulerDiscreteScheduler(),
+        vae=AutoencoderKLFlux2(
+            down_block_types=("DownEncoderBlock2D",),
+            up_block_types=("UpDecoderBlock2D",),
+            block_out_channels=(32,),
+            layers_per_block=1,
+            latent_channels=16,
+            norm_num_groups=16,
+        ),
+        # NOTE: For now we need 28 layers because of hardcoded stuff in the model :(
+        text_encoder=Qwen3ForCausalLM(
+            Qwen3Config(
+                hidden_size=32,
+                intermediate_size=64,
+                num_hidden_layers=28,
+                num_attention_heads=2,
+                num_key_value_heads=2,
+                head_dim=16,
+                vocab_size=151936,
+                max_position_embeddings=512,
+            )
+        ),
+        tokenizer=AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B"),
+        # NOTE: For now we need at least 2 layers for the transformer
+        # due to hardcoded hacks in CacheDiT for Flux2Klein specifically.
+        transformer=Flux2Transformer2DModel(
+            in_channels=64,
+            num_layers=2,
+            num_single_layers=2,
+            attention_head_dim=16,
+            num_attention_heads=2,
+            joint_attention_dim=96,
+            timestep_guidance_channels=32,
+            axes_dims_rope=(4, 4, 4, 4),
+        ),
+    )
+    # Need dtypes to be consistent; for now we just put it on bfloat16
+    pipe.to(torch.bfloat16).save_pretrained(model_dir)
+    return model_dir

tests/model_tests/diffusion/model_settings.py

@@ -0,0 +1,42 @@
+from tests.model_tests.diffusion import diff_model_builders
+from tests.model_tests.diffusion.config_types import (
+    DiffusionAccs,
+    DiffusionModelTestOpts,
+    DiffusionTasks,
+)
+
+# This object defines the (tiny) model configurations for common tests.
+#
+# TO ADD A NEW MODEL:
+# Map the pipeline name (i.e., the name of the pipeline in _DIFFUSION_MODELS)
+# to a DiffusionModelTestOpts and define the supported tasks. A base case with no
+# accelerations is always run. Use extra_test_groups to define additional
+# acceleration configurations to run a basic smoke test for (i.e., simple inference).
+#
+# The model will only be loaded once per test group, and will execute each of its
+# supported tasks as a pytest subtest. The base case also runs additional checks
+# that are independent of accelerations, e.g., for determinism and compatibility with
+# multi-output, unless explicitly disabled in the test settings.
+#
+# TO RUN A SUBSET OF THE TESTS:
+# These tests should be fast, but if you'd like to run a subset of the tests, you can do so
+# with `-k`; the IDs of the tests are the name of the pipeline joined by +.
+#
+# Example:
+# $ pytest test_common_offline.py -k test_pipeline_on_supported_tasks[Flux2KleinPipeline]
+#   ^ Runs only the base (no accelerations) case for Flux2KleinPipeline
+#
+# $ pytest test_common_offline.py -k test_pipeline_on_supported_tasks[Flux2KleinPipeline
+#   ^ Runs all test groups for Flux2KleinPipeline only
+DIFFUSION_TEST_SETTINGS = {
+    "Flux2KleinPipeline": DiffusionModelTestOpts(
+        model="black-forest-labs/FLUX.2-klein-4B",
+        builder=diff_model_builders.tiny_flux2_klein_builder,
+        supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE, DiffusionTasks.IMAGE_TO_IMAGE],
+        extra_test_groups=[
+            [DiffusionAccs.HSDP, DiffusionAccs.TEA_CACHE],
+            [DiffusionAccs.SEQUENCE_PARALLEL, DiffusionAccs.CACHE_DIT, DiffusionAccs.LAYERWISE_OFFLOAD],
+            [DiffusionAccs.CFG_PARALLEL, DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CPU_OFFLOAD],
+        ],
+    )
+}

tests/model_tests/diffusion/task_runners.py

@@ -0,0 +1,172 @@
+"""
+Definitions running and validating individual tasks, e.g., text to image,
+image to image, and so on. These are called by the core test runner.
+"""
+
+import base64
+import io
+from dataclasses import replace
+
+import numpy as np
+from PIL import Image
+
+from tests.helpers.runtime import DiffusionResponse, OmniServer, OpenAIClientHandler, dummy_messages_from_mix_data
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+
+PROMPT = "Dummy prompt"
+IMAGE_DIMS = (512, 512)
+HEIGHT, WIDTH = IMAGE_DIMS
+INPUT_IMAGE = Image.new("RGB", IMAGE_DIMS)
+
+# Offline sampling params
+IMAGE_GEN_SAMPLING_PARAMS = OmniDiffusionSamplingParams(
+    num_inference_steps=4,
+    height=HEIGHT,
+    width=WIDTH,
+    seed=42,
+)
+
+# Online extra_body for diffusion requests
+IMAGE_GEN_EXTRA_BODY = {
+    "height": HEIGHT,
+    "width": WIDTH,
+    "num_inference_steps": 4,
+    "seed": 42,
+}
+
+
+### Shared validation
+def _validate_images(images: list[Image.Image], expected_n: int = 1):
+    """Given a set of images, ensure we got the expected count of images,
+    and that all provided images match the expected dimensions."""
+    assert len(images) == expected_n
+    for img in images:
+        assert isinstance(img, Image.Image)
+        assert img.size == IMAGE_DIMS
+    return images
+
+
+def _validate_image_gen_determinism(images_a: list[Image.Image], images_b: list[Image.Image]):
+    """Ensure that two image sets are valid and that the results match."""
+    _validate_images(images_a)
+    _validate_images(images_b)
+    assert np.array_equal(np.array(images_a[0]), np.array(images_b[0]))
+
+
+### Output extractor utils for offline / online paths respectively
+def _get_offline_images(outputs: list[OmniRequestOutput]) -> list[Image.Image]:
+    """Extract the images from an Omni .generate() call."""
+    assert len(outputs) == 1
+    return outputs[0].images
+
+
+def _get_online_images(responses: list[DiffusionResponse]) -> list[Image.Image]:
+    """Extract the images from a server response."""
+    assert len(responses) == 1
+    images = responses[0].images
+    assert images is not None
+    return images
+
+
+### Offline helpers
+def _run_offline_t2i(omni: Omni, params: OmniDiffusionSamplingParams = IMAGE_GEN_SAMPLING_PARAMS):
+    return omni.generate({"prompt": PROMPT}, params)
+
+
+def _run_offline_i2i(omni: Omni):
+    return omni.generate(
+        {"prompt": PROMPT, "multi_modal_data": {"image": INPUT_IMAGE}},
+        IMAGE_GEN_SAMPLING_PARAMS,
+    )
+
+
+### Online helpers
+def _build_online_image_data_url() -> str:
+    """Get a valid base 64 encoded data URL corresponding to an image."""
+    buf = io.BytesIO()
+    INPUT_IMAGE.save(buf, format="PNG")
+    b64 = base64.b64encode(buf.getvalue()).decode()
+    return f"data:image/png;base64,{b64}"
+
+
+def _run_online_t2i(
+    server: OmniServer, client: OpenAIClientHandler, extra_body: dict | None = None
+) -> list[DiffusionResponse]:
+    """Run a text to image request through the server."""
+    messages = dummy_messages_from_mix_data(content_text=PROMPT)
+    request_config = {
+        "model": server.model,
+        "messages": messages,
+        "extra_body": extra_body or IMAGE_GEN_EXTRA_BODY,
+    }
+    return client.send_diffusion_request(request_config)
+
+
+def _run_online_i2i(server: OmniServer, client: OpenAIClientHandler) -> list[DiffusionResponse]:
+    """Run an image to image request through the server."""
+    image_data_url = _build_online_image_data_url()
+    messages = dummy_messages_from_mix_data(
+        content_text=PROMPT,
+        image_data_url=image_data_url,
+    )
+    request_config = {
+        "model": server.model,
+        "messages": messages,
+        "extra_body": IMAGE_GEN_EXTRA_BODY,
+    }
+    return client.send_diffusion_request(request_config)
+
+
+### Offline task runners
+def run_and_validate_text_to_image_request(omni: Omni):
+    """Run and validate a text to image request."""
+    _validate_images(_get_offline_images(_run_offline_t2i(omni)))
+
+
+def run_and_validate_image_to_image_request(omni: Omni):
+    """Run and validate an image to image request."""
+    _validate_images(_get_offline_images(_run_offline_i2i(omni)))
+
+
+def run_and_validate_text_to_image_determinism(omni: Omni):
+    """Checks for determinism; for now we just keep this for TTI."""
+    _validate_image_gen_determinism(
+        _get_offline_images(_run_offline_t2i(omni)),
+        _get_offline_images(_run_offline_t2i(omni)),
+    )
+
+
+def run_and_validate_text_to_image_multi_output(omni: Omni):
+    """Checks for multi-output; for now we just keep this for TTI."""
+    params = replace(IMAGE_GEN_SAMPLING_PARAMS, num_outputs_per_prompt=2)
+    _validate_images(_get_offline_images(_run_offline_t2i(omni, params)), expected_n=2)
+
+
+### Online task runners
+def run_and_validate_online_text_to_image_request(server: OmniServer, client: OpenAIClientHandler):
+    """Run and validate a text to image request through the server."""
+    _validate_images(_get_online_images(_run_online_t2i(server, client)))
+
+
+def run_and_validate_online_image_to_image_request(server: OmniServer, client: OpenAIClientHandler):
+    """Run and validate an image to image request through the server."""
+    _validate_images(_get_online_images(_run_online_i2i(server, client)))
+
+
+def run_and_validate_online_text_to_image_determinism(server: OmniServer, client: OpenAIClientHandler):
+    """Checks for determinism through the server; for now we just keep this for TTI."""
+    _validate_image_gen_determinism(
+        _get_online_images(_run_online_t2i(server, client)),
+        _get_online_images(_run_online_t2i(server, client)),
+    )
+
+
+def run_and_validate_online_text_to_image_multi_output(server: OmniServer, client: OpenAIClientHandler):
+    """Checks for multi-output through the server; for now we just keep this for TTI."""
+    extra_body = {**IMAGE_GEN_EXTRA_BODY, "num_outputs_per_prompt": 2}
+    _validate_images(
+        _get_online_images(_run_online_t2i(server, client, extra_body=extra_body)),
+        expected_n=2,
+    )

tests/model_tests/diffusion/test_alignment.py

@@ -0,0 +1,99 @@
+"""
+Tests for alignment between _DIFFUSION_MODELS and DIFFUSION_TEST_SETTINGS; if
+tests in this file are failing, you are probably adding a new model, and need
+to account for it in DIFFUSION_TEST_SETTINGS
+"""
+
+import pytest
+
+from tests.model_tests.diffusion.model_settings import DIFFUSION_TEST_SETTINGS
+from vllm_omni.diffusion.registry import _DIFFUSION_MODELS
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+
+# These diffusion models currently do not have any added (tiny) model tests.
+# In general, pipelines should only be in the list below if they have not been
+# migrated yet or there is a compelling reason that they should not be (which for
+# now may just be CI pressure etc due to the volume of models and the cost of
+# initializing the server).
+#
+# The tests here validate that all keys in the diffusion registry are accounted for;
+# if you are adding a new model and see the tests below fail, please follow the pattern
+# for adding a new tiny model builder and corresponding entry in DIFFUSION_TEST_SETTINGS.
+EXCLUDED_MODELS = [
+    "QwenImagePipeline",
+    "QwenImageEditPipeline",
+    "QwenImageEditPlusPipeline",
+    "QwenImageLayeredPipeline",
+    "GlmImagePipeline",
+    "ZImagePipeline",
+    "OvisImagePipeline",
+    "WanPipeline",
+    "WanVACEPipeline",
+    "LTX2Pipeline",
+    "LTX2ImageToVideoPipeline",
+    "LTX2TwoStagesPipeline",
+    "LTX2ImageToVideoTwoStagesPipeline",
+    "LTX2T2VDMD2Pipeline",
+    "LTX2I2VDMD2Pipeline",
+    "LTX23Pipeline",
+    "LTX23ImageToVideoPipeline",
+    "StableAudioPipeline",
+    "WanImageToVideoPipeline",
+    "WanS2VPipeline",
+    "WanT2VDMD2Pipeline",
+    "WanI2VDMD2Pipeline",
+    "LongCatImagePipeline",
+    "BagelPipeline",
+    "LancePipeline",
+    "MingImagePipeline",
+    "InternVLAA1Pipeline",
+    "LongCatImageEditPipeline",
+    "StableDiffusion3Pipeline",
+    "FluxKontextPipeline",
+    "HunyuanImage3ForCausalMM",
+    "ErnieImagePipeline",
+    "NextStep11Pipeline",
+    "FluxPipeline",
+    "FluxDMD2Pipeline",
+    "QwenImageDMD2Pipeline",
+    "OmniGen2Pipeline",
+    "HeliosPipeline",
+    "HeliosPyramidPipeline",
+    "Flux2Pipeline",
+    "DreamIDOmniPipeline",
+    "SenseNovaU1Pipeline",
+    "AudioXPipeline",
+    "HunyuanVideo15Pipeline",
+    "HunyuanVideo15ImageToVideoPipeline",
+    "LingBotVideoPipeline",
+    "MagiHumanPipeline",
+    "OmniVoicePipeline",
+    "OmniVoice",
+    "Cosmos3OmniDiffusersPipeline",
+    "DiffusersAdapterPipeline",
+    "HiDreamImagePipeline",
+    "DreamZeroPipeline",
+    "StableDiffusionXLPipeline",
+    "Gr00tN1d7Pipeline",
+    "SoulXSingerPipeline",
+    "SoulXSingerSVCPipeline",
+]
+
+
+def test_all_non_excluded_pipelines_have_tests():
+    """Ensure that every pipeline in the diffusion registry either has a test
+    configuration in DIFFUSION_TEST_SETTINGS, or is explicitly excluded for a
+    known reason."""
+    non_excluded_diff_pipes = _DIFFUSION_MODELS.keys() - set(EXCLUDED_MODELS)
+    missing_tests = non_excluded_diff_pipes - DIFFUSION_TEST_SETTINGS.keys()
+    assert len(missing_tests) == 0, f"Pipelines missing test settings: {missing_tests}"
+
+
+def test_no_excluded_models_have_test_settings():
+    """Ensure that no models in the exclude list have a test configuration in
+    DIFFUSION_TEST_SETTINGS (i.e., if they have test settings, we should not
+    exclude them)."""
+    non_excluded_diff_pipes = _DIFFUSION_MODELS.keys() - set(EXCLUDED_MODELS)
+    missing_pipes = DIFFUSION_TEST_SETTINGS.keys() - set(non_excluded_diff_pipes)
+    assert len(missing_pipes) == 0

tests/model_tests/diffusion/test_common_offline.py

@@ -0,0 +1,69 @@
+import pytest
+
+from tests.model_tests.diffusion.case_filtering import get_parametrized_options
+from tests.model_tests.diffusion.config_types import (
+    DiffusionAccs,
+    DiffusionTasks,
+    build_omni_from_diff_accelerations,
+)
+from tests.model_tests.diffusion.model_settings import DIFFUSION_TEST_SETTINGS
+from tests.model_tests.diffusion.task_runners import (
+    run_and_validate_image_to_image_request,
+    run_and_validate_text_to_image_determinism,
+    run_and_validate_text_to_image_multi_output,
+    run_and_validate_text_to_image_request,
+)
+
+# NOTE : Hardware marks are added dynamically based on test requirements
+pytestmark = [pytest.mark.diffusion]
+
+
+@pytest.mark.parametrize(
+    "model_name,accelerations,supported_tasks,check_multioutput,check_determinism",
+    get_parametrized_options(DIFFUSION_TEST_SETTINGS),
+)
+def test_pipeline_on_supported_tasks(
+    model_name,
+    accelerations: list[DiffusionAccs] | None,
+    supported_tasks: list[DiffusionTasks],
+    check_multioutput: bool,
+    check_determinism: bool,
+    tiny_model_paths: dict[str, str],
+    subtests,
+):
+    """Run a smoke test on all of the pipelines supported tasks using a set of enabled accelerations."""
+    assert len(supported_tasks) > 0
+    # We initialize the Omni object before running the tasks, then run each task as a pytest subtask.
+    # This lets us init the model once, but display separate failures in pytest, and avoid halting the
+    # checks on other tasks if one fails.
+    #
+    # This allows us to have some degree of test isolation without the cost of redundant initialization,
+    # since starting the server can take 10+ seconds, even for tiny models.
+    #
+    # NOTE: Be sure to install pytest-subtests if you're running on pytest < 9
+    omni = build_omni_from_diff_accelerations(
+        accelerations=accelerations,
+        model=tiny_model_paths[model_name],
+        enforce_eager=True,
+    )
+    try:
+        for task_type in supported_tasks:
+            with subtests.test(msg=task_type):
+                if task_type == DiffusionTasks.TEXT_TO_IMAGE:
+                    run_and_validate_text_to_image_request(omni)
+                elif task_type == DiffusionTasks.IMAGE_TO_IMAGE:
+                    run_and_validate_image_to_image_request(omni)
+                else:
+                    raise ValueError(f"Task type {task_type} is not yet supported")
+
+        # NOTE: For now, we only check determinism + multi output for the base case,
+        # since checking it on every extra acceleration configuration is redundant
+        # (see case_filtering).
+        if check_determinism:
+            with subtests.test(msg="determinism"):
+                run_and_validate_text_to_image_determinism(omni)
+        if check_multioutput:
+            with subtests.test(msg="multi_output"):
+                run_and_validate_text_to_image_multi_output(omni)
+    finally:
+        omni.close()

tests/model_tests/diffusion/test_common_online.py

@@ -0,0 +1,75 @@
+"""
+Analogous to test_common_offline, but for server tests. Validates the full
+online serving stack (CLI arg parsing, subprocess, API routing, response
+encoding) using tiny models.
+"""
+
+import pytest
+
+from tests.helpers.runtime import OmniServer, OpenAIClientHandler
+from tests.model_tests.diffusion.case_filtering import get_parametrized_options
+from tests.model_tests.diffusion.config_types import (
+    DiffusionAccs,
+    DiffusionTasks,
+    build_server_args_from_diff_accelerations,
+)
+from tests.model_tests.diffusion.model_settings import DIFFUSION_TEST_SETTINGS
+from tests.model_tests.diffusion.task_runners import (
+    run_and_validate_online_image_to_image_request,
+    run_and_validate_online_text_to_image_determinism,
+    run_and_validate_online_text_to_image_multi_output,
+    run_and_validate_online_text_to_image_request,
+)
+
+# NOTE : Hardware marks are added dynamically based on test requirements
+pytestmark = [pytest.mark.diffusion]
+
+
+@pytest.mark.parametrize(
+    "model_name,accelerations,supported_tasks,check_multioutput,check_determinism",
+    get_parametrized_options(DIFFUSION_TEST_SETTINGS, online=True),
+)
+def test_online_on_supported_tasks(
+    model_name: str,
+    accelerations: list[DiffusionAccs] | None,
+    supported_tasks: list[DiffusionTasks],
+    check_multioutput: bool,
+    check_determinism: bool,
+    tiny_model_paths: dict[str, str],
+    run_level: str,
+    subtests,
+):
+    """Smoke test: start a tiny model server and run each supported task via the API."""
+    model_path = tiny_model_paths[model_name]
+    server_args = build_server_args_from_diff_accelerations(accelerations)
+    server_args.append("--enforce-eager")
+
+    with OmniServer(model_path, server_args) as server:
+        # TODO: We may want to revisit run_level validation here,
+        # because checks for things like image size etc should not
+        # depend on whether or not the weights are real or random
+        client = OpenAIClientHandler(
+            host=server.host,
+            port=server.port,
+            api_key="EMPTY",
+            run_level=run_level,
+            log_stats=server.log_stats,
+        )
+        for task_type in supported_tasks:
+            with subtests.test(msg=task_type):
+                if task_type == DiffusionTasks.TEXT_TO_IMAGE:
+                    run_and_validate_online_text_to_image_request(server, client)
+                elif task_type == DiffusionTasks.IMAGE_TO_IMAGE:
+                    run_and_validate_online_image_to_image_request(server, client)
+                else:
+                    raise ValueError(f"Task type {task_type} is not yet supported")
+
+        # NOTE: For now, we only check determinism + multi output for the base case,
+        # since checking it on every extra acceleration configuration is redundant
+        # (see case_filtering).
+        if check_determinism:
+            with subtests.test(msg="determinism"):
+                run_and_validate_online_text_to_image_determinism(server, client)
+        if check_multioutput:
+            with subtests.test(msg="multi_output"):
+                run_and_validate_online_text_to_image_multi_output(server, client)

tests/model_tests/diffusion/utils.py

@@ -0,0 +1,33 @@
+"""
+Utilities for resolving real models to their tiny model equivalents.
+"""
+
+import logging
+
+from tests.model_tests.diffusion.model_settings import DIFFUSION_TEST_SETTINGS
+from vllm_omni.diffusion.data import resolve_model_class_name
+
+logger = logging.getLogger(__name__)
+
+
+def resolve_tiny_model_path(model: str) -> str:
+    """Given a real model name/path, resolve it to a tiny model path.
+
+    Raises ValueError if the pipeline class cannot be determined (invalid
+    model). Returns the original model path if no tiny builder exists yet."""
+    pipeline_class = resolve_model_class_name(model)
+    if pipeline_class is None:
+        raise ValueError(
+            f"Cannot resolve pipeline class for model: {model}. The model path may be invalid or its config unreadable."
+        )
+
+    test_opts = DIFFUSION_TEST_SETTINGS.get(pipeline_class)
+    if test_opts is None:
+        logger.warning(
+            "No tiny model builder for pipeline %s (model: %s). Using original model.",
+            pipeline_class,
+            model,
+        )
+        return model
+
+    return test_opts.builder()

tests/test_config_factory.py

@@ -16,6 +16,7 @@ from transformers import PretrainedConfig, Qwen3OmniMoeConfig
 
 from tests.helpers.stage_config import get_deploy_config_path, get_deploy_config_stage
 from vllm_omni.config.config_factory import StageConfigFactory
+from vllm_omni.config.endpoint_policy import EndpointRestriction, OmniServingCapability
 from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, register_pipeline
 from vllm_omni.config.stage_config import (
     DeployConfig,
@@ -39,6 +40,15 @@ from vllm_omni.engine.arg_utils import SHARED_FIELDS, EngineArgs, internal_black
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+@pytest.fixture(autouse=True)
+def clear_config_factory_caches():
+    """Clear cached classmethods from the StageConfigFactory to prevent test pollution."""
+    yield
+    StageConfigFactory.get_hf_config.cache_clear()
+    StageConfigFactory.try_infer_model_type.cache_clear()
+    StageConfigFactory.get_pipeline_config.cache_clear()
+
+
 Q3_OMNI_ALL_STAGES_HF_CONFIG = Qwen3OmniMoeConfig(enable_audio_output=True)
 Q3_OMNI_THINKER_HF_CONFIG = Qwen3OmniMoeConfig(enable_audio_output=False)
 
@@ -666,6 +676,43 @@ class TestPipelineRegistration:
         assert resolved_config is not None
         assert len(resolved_config) > 0
 
+    def test_deploy_override_uses_correct_endpoint_restrictions(self, clean_pipeline_registry, tmp_path):
+        """Ensure endpoint restrictions must come from the final pipeline
+        after deploy config overrides, not the auto-detected pipeline.
+        """
+        # Register two pipeline configs, where one has an endpoint restriction, and one doesn't
+        restriction = EndpointRestriction(
+            OmniServingCapability.COMPLETIONS,
+            "pipeline_a blocks completions",
+        )
+        pipe_a = PipelineConfig(model_type="detect_type", endpoint_restrictions=(restriction,))
+        pipe_b = PipelineConfig(model_type="override_type", endpoint_restrictions=())
+        register_pipeline(pipe_a)
+        register_pipeline(pipe_b)
+
+        # Create a config with the autodetected type, and write the
+        # deploy config specifying the override type to a temp path
+        class FakeConfig(PretrainedConfig):
+            model_type = "detect_type"
+
+        deploy_yaml = tmp_path / "override.yaml"
+        deploy_yaml.write_text("pipeline: override_type\n")
+
+        # Get the endpoint restrictions, passing the deploy config with the override
+        # type + patching the config for the detected type. Ensure that the endpoint
+        # restrictions correspond to the type in the deploy config.
+        with patch(
+            "vllm_omni.config.config_factory.get_config",
+            return_value=FakeConfig(),
+        ):
+            restrictions = StageConfigFactory.get_pipeline_endpoint_restrictions(
+                model="fake/model",
+                trust_remote_code=False,
+                deploy_config_path=str(deploy_yaml),
+            )
+
+        assert restrictions == ()
+
 
 class TestResolveScheduler:
     def test_all_execution_types_handled(self):
@@ -749,6 +796,32 @@ stages:
         runtime_config = deploy.stages[0].engine_extras["hf_overrides"]["voxcpm2_runtime_config"]
         assert runtime_config == expected_runtime_config
 
+    @pytest.mark.parametrize(
+        ("deploy_name", "pipeline_name", "stage_count", "final_output_type"),
+        [
+            ("mammoth_moda2.yaml", "mammoth_moda2", 2, "image"),
+            ("mammoth_moda2_ar.yaml", "mammoth_moda2_ar", 1, "text"),
+            ("omnivoice.yaml", "omnivoice", 1, "audio"),
+        ],
+    )
+    def test_load_new_registry_backed_deploy_configs(
+        self,
+        deploy_name: str,
+        pipeline_name: str,
+        stage_count: int,
+        final_output_type: str,
+    ):
+        deploy_path = Path(get_deploy_config_path(deploy_name))
+        deploy = load_deploy_config(deploy_path)
+        assert deploy.pipeline == pipeline_name
+
+        with patch("vllm_omni.platforms.current_omni_platform") as platform:
+            platform.device_name = "cuda"
+            stages = merge_pipeline_deploy(OMNI_PIPELINES[pipeline_name], deploy)
+        assert len(stages) == stage_count
+        assert stages[-1].final_output is True
+        assert stages[-1].final_output_type == final_output_type
+
     def test_merge_pipeline_deploy(self):
         deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
         if not deploy_path.exists():

tests/unit/higgs_audio_v3/test_higgs_audio_v3.py

@@ -7,6 +7,7 @@ without requiring the actual checkpoint or GPU.
 """
 
 import asyncio
+import json
 import time
 from collections import OrderedDict, defaultdict
 from types import SimpleNamespace
@@ -1309,7 +1310,10 @@ class TestVoiceCloneReferenceCache:
         parent = tmp_path / "OmniVoice"
         subdir = parent / "audio_tokenizer"
         subdir.mkdir(parents=True)
-        (subdir / "config.json").write_text("{}", encoding="utf-8")
+        (subdir / "config.json").write_text(
+            json.dumps({"model_type": "higgs_audio_v2_tokenizer"}),
+            encoding="utf-8",
+        )
 
         monkeypatch.setenv("HIGGS_AUDIO_TOKENIZER_PATH", str(parent))
         assert tok._resolve_audio_tokenizer_dir() == str(subdir)
@@ -1317,6 +1321,32 @@ class TestVoiceCloneReferenceCache:
         monkeypatch.setenv("HIGGS_AUDIO_TOKENIZER_PATH", str(subdir))
         assert tok._resolve_audio_tokenizer_dir() == str(subdir)
 
+    def test_audio_tokenizer_dir_rejects_omnivoice_config(self, tmp_path, monkeypatch):
+        import huggingface_hub
+        import huggingface_hub.constants as hub_constants
+
+        from vllm_omni.model_executor.models.higgs_audio_v2 import higgs_audio_v2_tokenizer as tok
+
+        parent = tmp_path / "OmniVoice"
+        subdir = parent / "audio_tokenizer"
+        subdir.mkdir(parents=True)
+        (subdir / "config.json").write_text(
+            json.dumps({"model_type": "omnivoice"}),
+            encoding="utf-8",
+        )
+
+        monkeypatch.setenv("HIGGS_AUDIO_TOKENIZER_PATH", str(parent))
+        monkeypatch.setenv("HIGGS_AUDIO_V2_TOKENIZER_PATH", "")
+        monkeypatch.setattr(
+            huggingface_hub,
+            "try_to_load_from_cache",
+            lambda **_: str(subdir / "config.json"),
+        )
+        monkeypatch.setattr(hub_constants, "HF_HUB_CACHE", str(tmp_path / "empty_cache"))
+
+        assert tok._normalize_audio_tokenizer_dir(str(parent)) is None
+        assert tok._resolve_audio_tokenizer_dir() is None
+
     def test_higgs_v3_ref_code_cache_returns_clone(self):
         from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
 

tests/worker/test_gpu_ar_model_runner.py

@@ -266,7 +266,6 @@ def test_build_omni_output_uses_snapshots_and_connector_after_accumulation(monke
         ec_connector_output=None,
         cudagraph_stats=None,
         kv_extracted_req_ids=["r2"],
-        seq_len=3,
         num_scheduled_tokens_np=torch.tensor([1, 2], dtype=torch.int32).numpy(),
         query_start_loc_cpu=torch.tensor([0, 1], dtype=torch.long),
     )
@@ -312,7 +311,6 @@ def test_build_omni_output_copies_hidden_for_partial_downstream_batch(monkeypatc
         ec_connector_output=None,
         cudagraph_stats=None,
         kv_extracted_req_ids=None,
-        seq_len=6,
         num_scheduled_tokens_np=np.array([1, 2, 3], dtype=np.int32),
         query_start_loc_cpu=torch.tensor([0, 1, 3], dtype=torch.long),
     )
@@ -366,7 +364,6 @@ def test_process_additional_information_uses_snapshot_request_order(monkeypatch)
         ec_connector_output=None,
         cudagraph_stats=None,
         kv_extracted_req_ids=None,
-        seq_len=3,
         num_scheduled_tokens_np=torch.tensor([1, 2], dtype=torch.int32).numpy(),
         query_start_loc_cpu=torch.tensor([0, 1], dtype=torch.long),
     )
@@ -427,7 +424,6 @@ def test_build_omni_output_skips_hidden_when_model_opts_out(monkeypatch):
         ec_connector_output=None,
         cudagraph_stats=None,
         kv_extracted_req_ids=None,
-        seq_len=2,
         num_scheduled_tokens_np=np.array([2], dtype=np.int32),
         query_start_loc_cpu=torch.tensor([0], dtype=torch.long),
     )
@@ -439,6 +435,97 @@ def test_build_omni_output_skips_hidden_when_model_opts_out(monkeypatch):
     assert output.multimodal_outputs is None
 
 
+def test_build_omni_output_splits_mm_by_scheduled_tokens_when_hidden_is_tail_only(monkeypatch):
+    runner = _make_async_output_runner(engine_output_type="latent")
+    runner.model.omni_pooler_payload_include_hidden = False
+    runner._async_chunk = False
+    runner.requests = {"r1": object(), "r2": object(), "r3": object()}
+
+    monkeypatch.setattr(
+        GPUARModelRunner,
+        "_resolve_pooler_payload_req_ids",
+        lambda self, req_ids: ("latent", req_ids),
+    )
+    monkeypatch.setattr(GPUARModelRunner, "_should_accumulate_full_payload_output", lambda self: False)
+    monkeypatch.setattr(GPUARModelRunner, "get_omni_connector_output", lambda self: None)
+    monkeypatch.setattr(GPUARModelRunner, "_process_additional_information_updates", lambda *args, **kwargs: None)
+
+    codes = torch.arange(48, dtype=torch.long).reshape(3, 16)
+    output = GPUARModelRunner._build_omni_model_runner_output_from_snapshot(
+        runner,
+        scheduler_output=SimpleNamespace(
+            total_num_scheduled_tokens=3,
+            num_scheduled_tokens={"r1": 1, "r2": 1, "r3": 1},
+        ),
+        hidden_states=torch.tensor([[1.0]]),
+        staged_hidden_states_cpu=None,
+        multimodal_outputs={"codes": {"audio": codes}},
+        req_ids_output_copy=["r1", "r2", "r3"],
+        req_id_to_index_output_copy={"r1": 0, "r2": 1, "r3": 2},
+        valid_sampled_token_ids=[[101], [102], [103]],
+        logprobs_lists=None,
+        prompt_logprobs_dict={},
+        num_nans_in_logits=None,
+        kv_connector_output=None,
+        ec_connector_output=None,
+        cudagraph_stats=None,
+        kv_extracted_req_ids=None,
+        num_scheduled_tokens_np=np.array([1, 1, 1], dtype=np.int32),
+        query_start_loc_cpu=torch.tensor([0, 1, 2], dtype=torch.long),
+    )
+
+    assert output.inter_stage_outputs is not None
+    assert torch.equal(output.inter_stage_outputs[0]["codes.audio"], codes[0:1])
+    assert torch.equal(output.inter_stage_outputs[1]["codes.audio"], codes[1:2])
+    assert torch.equal(output.inter_stage_outputs[2]["codes.audio"], codes[2:3])
+
+
+def test_build_omni_output_splits_mm_by_hidden_len_when_scheduled_is_padded(monkeypatch):
+    """Thinker mm rows align to hidden_states.shape[0], not padded scheduled count."""
+    runner = _make_async_output_runner(engine_output_type="latent")
+    runner.model.omni_pooler_payload_include_hidden = False
+    runner._async_chunk = False
+    runner.requests = {"r1": object(), "r2": object(), "r3": object()}
+
+    monkeypatch.setattr(
+        GPUARModelRunner,
+        "_resolve_pooler_payload_req_ids",
+        lambda self, req_ids: ("latent", req_ids),
+    )
+    monkeypatch.setattr(GPUARModelRunner, "_should_accumulate_full_payload_output", lambda self: False)
+    monkeypatch.setattr(GPUARModelRunner, "get_omni_connector_output", lambda self: None)
+    monkeypatch.setattr(GPUARModelRunner, "_process_additional_information_updates", lambda *args, **kwargs: None)
+
+    layers = torch.arange(96, dtype=torch.long).reshape(3, 32)
+    output = GPUARModelRunner._build_omni_model_runner_output_from_snapshot(
+        runner,
+        scheduler_output=SimpleNamespace(
+            total_num_scheduled_tokens=5,
+            num_scheduled_tokens={"r1": 1, "r2": 1, "r3": 1},
+        ),
+        hidden_states=torch.randn(3, 4),
+        staged_hidden_states_cpu=None,
+        multimodal_outputs={"hidden_states": {"layers": {0: layers}}},
+        req_ids_output_copy=["r1", "r2", "r3"],
+        req_id_to_index_output_copy={"r1": 0, "r2": 1, "r3": 2},
+        valid_sampled_token_ids=[[101], [102], [103]],
+        logprobs_lists=None,
+        prompt_logprobs_dict={},
+        num_nans_in_logits=None,
+        kv_connector_output=None,
+        ec_connector_output=None,
+        cudagraph_stats=None,
+        kv_extracted_req_ids=None,
+        num_scheduled_tokens_np=np.array([1, 1, 1], dtype=np.int32),
+        query_start_loc_cpu=torch.tensor([0, 1, 2], dtype=torch.long),
+    )
+
+    assert output.inter_stage_outputs is not None
+    assert torch.equal(output.inter_stage_outputs[0]["hidden_states.layer_0"], layers[0:1])
+    assert torch.equal(output.inter_stage_outputs[1]["hidden_states.layer_0"], layers[1:2])
+    assert torch.equal(output.inter_stage_outputs[2]["hidden_states.layer_0"], layers[2:3])
+
+
 def test_async_snapshot_payload_omits_hidden_when_model_opts_out():
     runner = _make_async_output_runner()
     runner.model.omni_pooler_payload_include_hidden = False
@@ -627,7 +714,6 @@ def test_build_omni_output_falls_back_to_mm_cpu_without_prefix_merge(monkeypatch
         ec_connector_output=None,
         cudagraph_stats=None,
         kv_extracted_req_ids=None,
-        seq_len=2,
         num_scheduled_tokens_np=np.array([1, 1], dtype=np.int32),
         query_start_loc_cpu=torch.tensor([0, 1], dtype=torch.long),
     )

tests/worker/test_min_tokens_stop_ids_guard.py

@@ -0,0 +1,115 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Guard against out-of-vocabulary stop ids in min-tokens masking (#4962).
+
+vLLM's input processor folds the stage tokenizer's EOS id into
+``SamplingParams.all_stop_token_ids``. For AR stages whose lm_head is
+narrower than the tokenizer vocabulary (codec talkers such as Qwen3-TTS:
+3072 logits vs text EOS 151645), any ``min_tokens >= 1`` makes
+``MinTokensLogitsProcessor`` write ``-inf`` at an out-of-range index and
+the engine dies with a CUDA device-side assert.
+"""
+
+import pytest
+import torch
+from vllm import SamplingParams
+from vllm.v1.sample.logits_processor import (
+    BatchUpdate,
+    LogitsProcessors,
+    MinTokensLogitsProcessor,
+)
+
+from vllm_omni.worker.sampling_utils import sanitize_min_tokens_stop_ids
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+TALKER_VOCAB = 3072
+CODEC_EOS = 2150
+TEXT_EOS = 151645
+
+
+def _make_min_tokens_proc(stop_token_ids: list[int], eos_token_id: int | None) -> MinTokensLogitsProcessor:
+    """Build a MinTokensLogitsProcessor with one active request, mirroring
+    the production path (engine folds the tokenizer EOS via
+    ``update_from_generation_config``)."""
+    params = SamplingParams(min_tokens=2, stop_token_ids=stop_token_ids)
+    if eos_token_id is not None:
+        params.update_from_generation_config({}, eos_token_id)
+    proc = MinTokensLogitsProcessor(None, device=torch.device("cpu"), is_pin_memory=False)
+    proc.update_state(
+        BatchUpdate(
+            batch_size=1,
+            removed=[],
+            added=[(0, params, None, [])],
+            moved=[],
+        )
+    )
+    return proc
+
+
+def test_oob_stop_id_crashes_without_guard():
+    """Document the failure mode: the folded text EOS is out of range for a
+    narrow codec head and index_put_ raises (device-side assert on CUDA)."""
+    proc = _make_min_tokens_proc([CODEC_EOS], eos_token_id=TEXT_EOS)
+    logits = torch.zeros(1, TALKER_VOCAB)
+    with pytest.raises((IndexError, RuntimeError)):
+        proc.apply(logits)
+
+
+def test_guard_filters_oob_and_keeps_in_range_mask():
+    proc = _make_min_tokens_proc([CODEC_EOS], eos_token_id=TEXT_EOS)
+    sanitize_min_tokens_stop_ids(LogitsProcessors([proc]), TALKER_VOCAB)
+
+    logits = torch.zeros(1, TALKER_VOCAB)
+    out = proc.apply(logits)
+
+    assert out[0, CODEC_EOS] == -float("inf")
+    masked = (out == -float("inf")).nonzero()
+    assert masked.tolist() == [[0, CODEC_EOS]]
+
+
+def test_guard_noop_when_all_stop_ids_in_range():
+    """Full-width heads (e.g. CosyVoice3 talker) must be untouched: no
+    tensor rebuild, mask unchanged."""
+    proc = _make_min_tokens_proc([CODEC_EOS], eos_token_id=None)
+    slice_before = proc.logits_slice
+
+    sanitize_min_tokens_stop_ids(LogitsProcessors([proc]), TALKER_VOCAB)
+
+    assert proc.logits_slice is slice_before
+    out = proc.apply(torch.zeros(1, TALKER_VOCAB))
+    assert out[0, CODEC_EOS] == -float("inf")
+
+
+def test_guard_sanitizes_params_only_once():
+    """The stop-id set is shared with the request's SamplingParams, so the
+    second call finds nothing to drop (no repeated rebuilds/warnings)."""
+    proc = _make_min_tokens_proc([CODEC_EOS], eos_token_id=TEXT_EOS)
+    sanitize_min_tokens_stop_ids(LogitsProcessors([proc]), TALKER_VOCAB)
+    slice_after_first = proc.logits_slice
+
+    sanitize_min_tokens_stop_ids(LogitsProcessors([proc]), TALKER_VOCAB)
+
+    assert proc.logits_slice is slice_after_first
+
+
+def test_guard_handles_multiple_requests():
+    params_oob = SamplingParams(min_tokens=4, stop_token_ids=[CODEC_EOS])
+    params_oob.update_from_generation_config({}, TEXT_EOS)
+    params_ok = SamplingParams(min_tokens=4, stop_token_ids=[7])
+    proc = MinTokensLogitsProcessor(None, device=torch.device("cpu"), is_pin_memory=False)
+    proc.update_state(
+        BatchUpdate(
+            batch_size=2,
+            removed=[],
+            added=[(0, params_oob, None, []), (1, params_ok, None, [])],
+            moved=[],
+        )
+    )
+
+    sanitize_min_tokens_stop_ids(LogitsProcessors([proc]), TALKER_VOCAB)
+    out = proc.apply(torch.zeros(2, TALKER_VOCAB))
+
+    assert out[0, CODEC_EOS] == -float("inf")
+    assert out[1, 7] == -float("inf")
+    assert int((out == -float("inf")).sum()) == 2

tests/worker/test_omni_gpu_model_runner.py

@@ -81,6 +81,7 @@ class CaptureTalkerMTP(torch.nn.Module):
         top_k=None,
         top_p=None,
         generator=None,
+        generators=None,
     ):
         self.calls.append(
             {
@@ -90,6 +91,7 @@ class CaptureTalkerMTP(torch.nn.Module):
                 "top_k": top_k,
                 "top_p": top_p,
                 "generator": generator,
+                "generators": generators,
             }
         )
         codes = torch.zeros((req_embeds.shape[0], 1), dtype=torch.int64)
@@ -306,6 +308,7 @@ def test_talker_mtp_forward_passes_qwen3_tts_subtalker_sampling_params_to_talker
             "top_k": 9,
             "top_p": 0.55,
             "generator": runner.talker_mtp.calls[0]["generator"],
+            "generators": None,
         }
     ]
     assert runner.talker_mtp.calls[0]["generator"] is not None
@@ -350,6 +353,57 @@ def test_talker_mtp_forward_keeps_explicit_seeded_requests_scalar(monkeypatch):
     assert torch.equal(runner.talker_mtp_inputs_embeds.gpu, saved_embeds)
 
 
+def test_talker_mtp_forward_batches_seeded_requests_for_opted_in_models(monkeypatch):
+    """Models with talker_mtp_accepts_per_row_generators get one batched call (#4883)."""
+    import vllm_omni.worker.gpu_model_runner as mod
+
+    monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context)
+
+    runner = _make_runner(req_ids=("r1", "r2"), hidden_size=4)
+    runner.requests["r1"].sampling_params = SimpleNamespace(
+        seed=11,
+        extra_args={"tts_local_seed": 11},
+    )
+    runner.requests["r2"].sampling_params = SimpleNamespace(
+        seed=22,
+        extra_args={"tts_local_seed": 22},
+    )
+    runner.talker_mtp = CaptureTalkerMTP()
+    runner.model = SimpleNamespace(
+        talker_mtp_output_key=("codes", "audio"),
+        talker_mtp_accepts_per_row_generators=True,
+    )
+    runner.vllm_config = SimpleNamespace(model_config=SimpleNamespace(subtalker_sampling_params={}))
+
+    def fake_determine(self, num_tokens, num_reqs, num_scheduled_tokens_np, max_num_scheduled_tokens, use_cascade_attn):
+        batch_desc = SimpleNamespace(num_tokens=int(num_tokens))
+        return (False, batch_desc, None, None, None)
+
+    monkeypatch.setattr(runner, "_determine_batch_execution_and_padding", fake_determine.__get__(runner, type(runner)))
+
+    inputs_embeds = torch.zeros((6, 4), dtype=torch.float32)
+    OmniGPUModelRunner._talker_mtp_forward(runner, ["r1", "r2"], inputs_embeds)
+
+    # One batched call with distinct per-row generators, not two scalar calls.
+    assert [call["batch_size"] for call in runner.talker_mtp.calls] == [2]
+    row_generators = runner.talker_mtp.calls[0]["generators"]
+    assert runner.talker_mtp.calls[0]["generator"] is None
+    assert len(row_generators) == 2
+    assert all(generator is not None for generator in row_generators)
+    assert row_generators[0] is not row_generators[1]
+
+    # The per-request generator stream persists across steps...
+    OmniGPUModelRunner._talker_mtp_forward(runner, ["r1", "r2"], inputs_embeds)
+    assert runner.talker_mtp.calls[1]["generators"][0] is row_generators[0]
+    assert runner.talker_mtp.calls[1]["generators"][1] is row_generators[1]
+
+    # ...and is evicted once its request finishes.
+    del runner.requests["r2"]
+    OmniGPUModelRunner._talker_mtp_forward(runner, ["r1"], inputs_embeds)
+    assert set(runner._talker_mtp_generators) == {"r1"}
+    assert runner.talker_mtp.calls[2]["generator"] is row_generators[0]
+
+
 def test_update_intermediate_buffer_writes_to_buffer_and_setattr(monkeypatch):
     """Validate that _update_intermediate_buffer writes to model_intermediate_buffer
     (forward path) and mirrors to additional_information_cpu setattr (backward compat)."""

vllm_omni/config/config_factory.py

@@ -5,6 +5,7 @@
 from __future__ import annotations
 
 import dataclasses
+import functools
 from dataclasses import asdict
 from pathlib import Path
 from typing import Any
@@ -14,6 +15,7 @@ from vllm.logger import init_logger
 from vllm.transformers_utils.config import get_config
 from vllm.transformers_utils.repo_utils import get_hf_file_to_dict
 
+from vllm_omni.config.endpoint_policy import EndpointRestriction
 from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
 from vllm_omni.config.stage_config import (
     _DEPLOY_DIR,
@@ -46,83 +48,174 @@ class StageConfigFactory:
     """
 
     @classmethod
-    def create_from_model(
+    def get_pipeline_endpoint_restrictions(
         cls,
         model: str,
-        cli_overrides: dict[str, Any] | None = None,
-        deploy_config_path: str | None = None,
-        **deprecated_kwargs: Any,
-    ) -> list[StageConfig] | None:
-        """Load pipeline + deploy config, merge with CLI overrides.
+        trust_remote_code: bool,
+        deploy_config_path: str | None,
+    ) -> tuple[EndpointRestriction, ...]:
+        """Given a model string, determine the corresponding endpoint restrictions.
 
-        Checks OMNI_PIPELINES first, since supported models should be explicitly
-        registered. If a model is not registered in OMNI_PIPELINES, tries to fall
-        back to using the Transformers config & finding pipelines that have overlapping
-        supported architectures.
+        Args:
+            model: Model name or path.
+            trust_remote_code: Whether to trust remote code for HF config loading.
+            deploy_config_path: Optional path to the deploy config for the pipeline.
+
+        Returns:
+            A tuple of model specific endpoint restrictions.
         """
-        if cli_overrides is None:
-            cli_overrides = {}
+        pipeline_cfg = StageConfigFactory.get_pipeline_config(
+            model=model,
+            trust_remote_code=trust_remote_code,
+            deploy_config_path=deploy_config_path,
+        )
+        return pipeline_cfg.endpoint_restrictions if pipeline_cfg else ()
 
-        trust_remote_code = cli_overrides.get("trust_remote_code", True)
-        if trust_remote_code is None:
-            trust_remote_code = False
+    @classmethod
+    @functools.cache
+    def get_hf_config(cls, model: str, trust_remote_code: bool) -> PretrainedConfig | None:
+        """Fetch the HF config (if it exists) from the model directory.
 
-        model_type, hf_config = cls._auto_detect_model_type(model, trust_remote_code=trust_remote_code)
+        Args:
+            model: Model name or path.
+            trust_remote_code: Whether to trust remote code for HF config loading.
+
+        Returns:
+            the model's config or None.
+        """
+        hf_config = None
+        try:
+            return get_config(model, trust_remote_code=trust_remote_code)
+        except Exception as e:
+            logger.debug(f"`get_config` failed with exception {e}; inferred HF config is None")
+        return hf_config
+
+    @classmethod
+    @functools.cache
+    def try_infer_model_type(cls, model: str, trust_remote_code: bool) -> str | None:
+        """Auto-detect model_type from model directory and apply any model
+        specific patches to get the correct model_type str. If we are unable
+        to infer it from the model directory, we fall back to the PipelineConfig.
+
+        Args:
+            model: Model name or path.
+            trust_remote_code: Whether to trust remote code for HF config loading.
+
+        Returns:
+            model_type as a string; may be None on failure.
+        """
+        model_type = cls._try_infer_model_type(
+            model=model,
+            trust_remote_code=trust_remote_code,
+        )
         if model_type == "vla":
             if _looks_like_dreamzero(model):
                 model_type = "dreamzero"
+        return model_type
 
-        # --- 1. Explicit deploy-config pipeline override (highest precedence) ---
-        # A deploy YAML may set ``pipeline: <model_type>`` to force routing for
-        # models whose HF config has a generic/colliding ``model_type`` or no
-        # matching architectures. Ming-omni-tts, for example, reports
-        # model_type="dense" with arch BailingMMNative…, which matches no
-        # pipeline — without this it falls through to the diffusion default and
-        # dies with "Model class BailingMMNativeForConditionalGeneration not
-        # found in diffusion model registry". Honor the key before auto-detection.
-        explicit_pipeline = None
-        if deploy_config_path:
-            try:
-                explicit_pipeline = load_deploy_config(deploy_config_path).pipeline
-            except Exception:
-                logger.exception("Failed to read 'pipeline' key from deploy config %s", deploy_config_path)
-        if explicit_pipeline:
-            pipeline_cfg = cls.resolve_pipeline_config(explicit_pipeline, hf_config)
-            if pipeline_cfg is not None:
-                return cls._create_from_registry(
-                    explicit_pipeline,
-                    pipeline_cfg,
-                    cli_overrides,
-                    deploy_config_path,
-                )
-            logger.warning(
-                "Deploy config %s requested pipeline %r which is not in OMNI_PIPELINES; "
-                "falling back to auto-detection.",
-                deploy_config_path,
-                explicit_pipeline,
-            )
+    @classmethod
+    def _try_infer_model_type(cls, model: str, trust_remote_code: bool) -> str | None:
+        """Auto-detect model_type from model directory.
 
-        # --- 2. Auto-detected model_type registered in OMNI_PIPELINES ---
+        Args:
+            model: Model name or path.
+            trust_remote_code: Whether to trust remote code for HF config loading.
+
+        Returns:
+            model_type as a string; may be None on failure.
+        """
+        hf_config = cls.get_hf_config(
+            model=model,
+            trust_remote_code=trust_remote_code,
+        )
+        if hf_config is not None:
+            return hf_config.model_type
+
+        # Fallback: read config.json directly for custom model types that
+        # are not registered with transformers (e.g. qwen3_tts).
+        try:
+            config_dict = get_hf_file_to_dict("config.json", model, revision=None)
+            if config_dict:
+                if "model_type" in config_dict:
+                    return config_dict["model_type"]
+                # VoxCPM2-style configs use singular ``architecture`` rather
+                # than HF's standard ``model_type`` / ``architectures``. Accept
+                # it as a fallback so the pipeline registry can still match.
+                if "architecture" in config_dict and isinstance(config_dict["architecture"], str):
+                    return config_dict["architecture"]
+        except Exception as e:
+            logger.debug(f"Failed to auto-detect model type for {model}: {e}")
+
+        # Fallback for diffusers-style models: check model_index.json.
+        # Some models (e.g. GLM-Image) have no root config.json but ship a
+        # model_index.json with _class_name that maps to a pipeline key via
+        # PipelineConfig.diffusers_class_name.
+        try:
+            model_index = get_hf_file_to_dict("model_index.json", model, revision=None)
+            if model_index and "_class_name" in model_index:
+                class_name = model_index["_class_name"]
+                for obj in OMNI_PIPELINES.values():
+                    # If we have a resolver, call it with the optional hf_config
+                    # to get the default pipeline config for this key
+                    pipeline_cfg = obj(hf_config) if callable(obj) else obj
+                    if pipeline_cfg is not None and pipeline_cfg.diffusers_class_name == class_name:
+                        logger.info(
+                            "Detected pipeline %r from model_index.json (_class_name=%r)",
+                            pipeline_cfg.model_type,
+                            class_name,
+                        )
+                        return pipeline_cfg.model_type
+        except Exception as e:
+            logger.debug(f"Failed to detect model type for diffusers-style models: {e}")
+
+        # Final fallback: some models (e.g. CosyVoice3) ship an empty
+        # config.json and rely on naming conventions. Match the model path
+        # basename against registered pipeline keys — longest match wins
+        # so "cosyvoice3" (length 10) beats "cosyvoice" (length 9).
+        model_lower = model.lower().replace("-", "").replace("_", "")
+        best: str | None = None
+        best_len = 0
+        for registered_key in OMNI_PIPELINES.keys():
+            candidate = registered_key.lower().replace("-", "").replace("_", "")
+            if candidate and candidate in model_lower and len(candidate) > best_len:
+                best = registered_key
+                best_len = len(candidate)
+        if best is not None:
+            return best
+
+        return None
+
+    @classmethod
+    @functools.cache
+    def get_pipeline_config(
+        cls,
+        model: str,
+        trust_remote_code: bool,
+        deploy_config_path: str | None = None,
+    ) -> PipelineConfig | None:
+        """Resolve the PipelineConfig for a model path/name."""
+        model_type = cls.try_infer_model_type(model=model, trust_remote_code=trust_remote_code)
+        hf_config = cls.get_hf_config(model=model, trust_remote_code=trust_remote_code)
+
+        # Resolve the deploy config & check if the user set the pipeline;
+        # If the pipeline is explicitly set, it takes highest priority
+        deploy_config_pipe = cls._get_deploy_override_pipe_config(hf_config, deploy_config_path)
+        if deploy_config_pipe is not None:
+            return deploy_config_pipe
+
+        # Pipeline isn't set in the yaml spec, so we need infer it ourselves.
         if model_type and model_type in OMNI_PIPELINES:
             pipeline_cfg = cls.resolve_pipeline_config(model_type, hf_config)
             if pipeline_cfg is not None:
-                return cls._create_from_registry(
-                    model_type,
-                    pipeline_cfg,
-                    cli_overrides,
-                    deploy_config_path,
-                )
+                return pipeline_cfg
 
-        # --- HF architecture fallback: some models report a generic
-        # model_type that collides with another model. Match by the
-        # hf_architectures declared on each registered PipelineConfig.
         if hf_config is not None:
-            logger.warning("Inferred model type %s is not registered to an Omni pipeline", model_type)
+            if model_type is not None:
+                logger.warning("Inferred model type %s is not registered to an Omni pipeline", model_type)
             hf_archs = set(getattr(hf_config, "architectures", []) or [])
             if hf_archs:
                 for registered in OMNI_PIPELINES.values():
                     pipeline_cfg = registered if isinstance(registered, PipelineConfig) else registered(hf_config)
-                    # Resolvers that get configs of the incorrect type should return None
                     if pipeline_cfg is None:
                         continue
                     predicate = pipeline_cfg.hf_config_predicate
@@ -143,36 +236,61 @@ class StageConfigFactory:
                                 pipeline_cfg.model_type,
                             )
                             continue
-
                     if isinstance(pipeline_cfg, PipelineConfig) and hf_archs.intersection(
                         pipeline_cfg.hf_architectures
                     ):
-                        return cls._create_from_registry(
-                            pipeline_cfg.model_type,
-                            pipeline_cfg,
-                            cli_overrides,
-                            deploy_config_path,
-                        )
+                        return pipeline_cfg
+        return None
 
-        # --- Explicit deploy-config pipeline ---
-        # When auto-detection above resolves nothing (generic/missing HF model_type
-        # and no matching architecture), honor an explicit pipeline key in the deploy config
-        if deploy_config_path is not None:
-            deploy_path = Path(deploy_config_path)
-            if deploy_path.exists():
-                deploy_cfg = load_deploy_config(deploy_path)
-                if deploy_cfg.pipeline:
-                    pipeline_cfg = cls.resolve_pipeline_config(deploy_cfg.pipeline, hf_config)
-                    if pipeline_cfg is not None:
-                        return cls._create_from_registry(
-                            pipeline_cfg.model_type,
-                            pipeline_cfg,
-                            cli_overrides,
-                            deploy_config_path,
-                        )
+    @classmethod
+    def _get_deploy_override_pipe_config(
+        cls,
+        hf_config: PretrainedConfig | None,
+        deploy_config_path: str | None,
+    ) -> PipelineConfig | None:
+        """Load the deploy config and extract + resolve its pipeline field."""
+        if deploy_config_path is None:
+            return None
 
-        # Not in the pipeline registry — let the caller fall back to the
-        # legacy ``stage_configs/*.yaml`` path (resolve_model_config_path).
+        deploy_path = Path(deploy_config_path)
+        if deploy_path.exists():
+            deploy_cfg = load_deploy_config(deploy_path)
+            if deploy_cfg.pipeline is not None:
+                return cls.resolve_pipeline_config(deploy_cfg.pipeline, hf_config)
+        return None
+
+    @classmethod
+    def create_from_model(
+        cls,
+        model: str,
+        *,
+        trust_remote_code: bool = False,
+        cli_overrides: dict[str, Any] | None = None,
+        deploy_config_path: str | None = None,
+        **deprecated_kwargs: Any,
+    ) -> list[StageConfig] | None:
+        """Load pipeline + deploy config, merge with CLI overrides.
+
+        Checks OMNI_PIPELINES first, since supported models should be explicitly
+        registered. If a model is not registered in OMNI_PIPELINES, tries to fall
+        back to using the Transformers config & finding pipelines that have overlapping
+        supported architectures.
+        """
+        if cli_overrides is None:
+            cli_overrides = {}
+
+        pipeline_cfg = cls.get_pipeline_config(
+            model=model,
+            trust_remote_code=trust_remote_code,
+            deploy_config_path=deploy_config_path,
+        )
+        if pipeline_cfg is not None:
+            return cls._create_from_registry(
+                pipeline_cfg.model_type,
+                pipeline_cfg,
+                cli_overrides,
+                deploy_config_path,
+            )
         return None
 
     @classmethod
@@ -288,79 +406,6 @@ class StageConfigFactory:
 
         return [config_dict]
 
-    @classmethod
-    def _auto_detect_model_type(cls, model: str, trust_remote_code: bool = True) -> tuple[str | None, Any]:
-        """Auto-detect model_type from model directory.
-
-        Args:
-            model: Model name or path.
-            trust_remote_code: Whether to trust remote code for HF config loading.
-
-        Returns:
-            Tuple of (model_type, hf_config). Both may be None on failure.
-        """
-        hf_config = None
-
-        try:
-            hf_config = get_config(model, trust_remote_code=trust_remote_code)
-            return hf_config.model_type, hf_config
-        except Exception as e:
-            logger.debug(f"`get_config` failed for {e}; Falling back to raw config.json path")
-
-        # Fallback: read config.json directly for custom model types that
-        # are not registered with transformers (e.g. qwen3_tts).
-        try:
-            config_dict = get_hf_file_to_dict("config.json", model, revision=None)
-            if config_dict:
-                if "model_type" in config_dict:
-                    return config_dict["model_type"], None
-                # VoxCPM2-style configs use singular ``architecture`` rather
-                # than HF's standard ``model_type`` / ``architectures``. Accept
-                # it as a fallback so the pipeline registry can still match.
-                if "architecture" in config_dict and isinstance(config_dict["architecture"], str):
-                    return config_dict["architecture"], None
-        except Exception as e:
-            logger.debug(f"Failed to auto-detect model type for {model}: {e}")
-
-        # Fallback for diffusers-style models: check model_index.json.
-        # Some models (e.g. GLM-Image) have no root config.json but ship a
-        # model_index.json with _class_name that maps to a pipeline key via
-        # PipelineConfig.diffusers_class_name.
-        try:
-            model_index = get_hf_file_to_dict("model_index.json", model, revision=None)
-            if model_index and "_class_name" in model_index:
-                class_name = model_index["_class_name"]
-                for obj in OMNI_PIPELINES.values():
-                    # If we have a resolver, call it with the optional hf_config
-                    # to get the default pipeline config for this key
-                    pipeline_cfg = obj(hf_config) if callable(obj) else obj
-                    if pipeline_cfg is not None and pipeline_cfg.diffusers_class_name == class_name:
-                        logger.info(
-                            "Detected pipeline %r from model_index.json (_class_name=%r)",
-                            pipeline_cfg.model_type,
-                            class_name,
-                        )
-                        return pipeline_cfg.model_type, None
-        except Exception as e:
-            logger.debug(f"Failed to detect model type for diffusers-style models: {e}")
-
-        # Final fallback: some models (e.g. CosyVoice3) ship an empty
-        # config.json and rely on naming conventions. Match the model path
-        # basename against registered pipeline keys — longest match wins
-        # so "cosyvoice3" (length 10) beats "cosyvoice" (length 9).
-        model_lower = model.lower().replace("-", "").replace("_", "")
-        best: str | None = None
-        best_len = 0
-        for registered_key in OMNI_PIPELINES.keys():
-            candidate = registered_key.lower().replace("-", "").replace("_", "")
-            if candidate and candidate in model_lower and len(candidate) > best_len:
-                best = registered_key
-                best_len = len(candidate)
-        if best is not None:
-            return best, None
-
-        return None, None
-
     @classmethod
     def _merge_cli_overrides(
         cls,
@@ -381,6 +426,7 @@ class StageConfigFactory:
         """Given a model type, resolve to the pipeline to be used. If the pipeline
         maps to a callable we resolve based on the HF config."""
         if model_type not in OMNI_PIPELINES:
+            logger.warning("Model type %s is not registered to OMNI_PIPELINES", model_type)
             return None
         obj = OMNI_PIPELINES[model_type]
         return obj(hf_config) if callable(obj) else obj

vllm_omni/config/endpoint_policy.py

@@ -0,0 +1,91 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Endpoint restriction policy for omni pipelines."""
+
+from dataclasses import dataclass
+from enum import Enum
+from typing import NamedTuple
+
+from fastapi import FastAPI, Request
+from fastapi.responses import JSONResponse
+from vllm.entrypoints.serve.utils.error_response import create_error_response
+
+
+class RouteTarget(NamedTuple):
+    """A server path & supported methods."""
+
+    path: str
+    methods: frozenset[str]
+
+
+class OmniServingCapability(Enum):
+    """Serving capabilities that pipelines can shut down."""
+
+    CHAT_COMPLETIONS_BATCH = RouteTarget("/v1/chat/completions/batch", frozenset({"POST"}))
+    COMPLETIONS = RouteTarget("/v1/completions", frozenset({"POST"}))
+
+    @property
+    def path(self) -> str:
+        return self.value.path
+
+    @property
+    def methods(self) -> frozenset[str]:
+        return self.value.methods
+
+
+@dataclass(frozen=True)
+class EndpointRestriction:
+    capability: OmniServingCapability
+    reason: str
+
+
+# Routes that are not supported for any model, but are supported in vLLM.
+# This is only temporary to avoid 500s for batched chat completions.
+UNSUPPORTED_ROUTES: tuple[EndpointRestriction, ...] = (
+    EndpointRestriction(
+        OmniServingCapability.CHAT_COMPLETIONS_BATCH,
+        "Batched chat completions are not yet supported by vLLM Omni.",
+    ),
+)
+
+
+def build_rejection_handler(reason: str):
+    """Build a rejection handler for a given endpoint for the provided reason."""
+
+    async def rejection_handler(raw_request: Request):
+        error = create_error_response(message=reason)
+        return JSONResponse(
+            content=error.model_dump(),
+            status_code=error.error.code,
+        )
+
+    return rejection_handler
+
+
+def shutdown_unsupported_routes(
+    app: FastAPI,
+    endpoint_restrictions: tuple[EndpointRestriction, ...],
+):
+    """Given an initialized FastAPI server instance and a set of model specific endpoint
+    restrictions, remove the restricted routes and patch a handler that returns 400.
+    """
+    from vllm_omni.entrypoints.openai.api_server import _remove_route_from_app
+
+    # Generally these should not overlap since there is no point. If they do,
+    # we use the reason message in UNSUPPORTED_ROUTES, for consistent error messages.
+    restricted_endpoints = (*endpoint_restrictions, *UNSUPPORTED_ROUTES)
+
+    for end_restrict in restricted_endpoints:
+        capability = end_restrict.capability
+        # Remove the route from the app
+        _remove_route_from_app(app, capability.path, capability.methods)
+
+        # Patch the bad request error with the model specific
+        # reason for shutting down this endpoint
+        rejection_handler = build_rejection_handler(end_restrict.reason)
+
+        app.add_api_route(
+            capability.path,
+            rejection_handler,
+            methods=list(capability.methods),
+        )

vllm_omni/config/pipeline_registry.py

@@ -56,6 +56,10 @@ from vllm_omni.model_executor.models.hunyuan_image3.pipeline import (
 )
 from vllm_omni.model_executor.models.indextts2.pipeline import INDEXTTS2_PIPELINE
 from vllm_omni.model_executor.models.lance.pipeline import LANCE_PIPELINE
+from vllm_omni.model_executor.models.mammoth_moda2.pipeline import (
+    MAMMOTH_MODA2_AR_PIPELINE,
+    MAMMOTH_MODA2_PIPELINE,
+)
 from vllm_omni.model_executor.models.mimo_audio.pipeline import MIMO_AUDIO_PIPELINE
 from vllm_omni.model_executor.models.ming_flash_omni.pipeline import (
     MING_FLASH_OMNI_IMAGE_PIPELINE,
@@ -74,6 +78,7 @@ from vllm_omni.model_executor.models.moss_tts.pipeline import (
     MOSS_TTS_REALTIME_PIPELINE,
 )
 from vllm_omni.model_executor.models.moss_tts_nano.pipeline import MOSS_TTS_NANO_PIPELINE
+from vllm_omni.model_executor.models.omnivoice.pipeline import OMNIVOICE_PIPELINE
 from vllm_omni.model_executor.models.qwen2_5_omni.pipeline import (
     QWEN2_5_OMNI_PIPELINE,
     QWEN2_5_OMNI_THINKER_ONLY_PIPELINE,
@@ -118,6 +123,9 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
     "ming_flash_omni_thinker_only": MING_FLASH_OMNI_THINKER_ONLY_PIPELINE,
     "ming_flash_omni_image": MING_FLASH_OMNI_IMAGE_PIPELINE,
     "moss_tts_nano": MOSS_TTS_NANO_PIPELINE,
+    "omnivoice": OMNIVOICE_PIPELINE,
+    "mammoth_moda2": MAMMOTH_MODA2_PIPELINE,
+    "mammoth_moda2_ar": MAMMOTH_MODA2_AR_PIPELINE,
     "moss_tts_delay": MOSS_TTS_PIPELINE,
     "moss_tts_realtime": MOSS_TTS_REALTIME_PIPELINE,
     "moss_tts_local": MOSS_TTS_LOCAL_PIPELINE,

vllm_omni/config/stage_config.py

@@ -18,6 +18,7 @@ from transformers import PretrainedConfig
 from vllm.logger import init_logger
 from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
 
+from vllm_omni.config.endpoint_policy import EndpointRestriction
 from vllm_omni.config.yaml_util import create_config, load_yaml_config, to_dict
 from vllm_omni.core.sched.omni_ar_scheduler import OmniARAsyncScheduler, OmniARScheduler
 from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
@@ -268,6 +269,7 @@ class PipelineConfig:
     # this value to auto-detect the pipeline.  Only needed for diffusers-style
     # multi-component repos (e.g. GLM-Image).  ``None`` = not a diffusers model.
     diffusers_class_name: str | None = None
+    endpoint_restrictions: tuple[EndpointRestriction, ...] = ()
 
     def get_stage(self, stage_id: int) -> StagePipelineConfig | None:
         """Look up a stage by its ID."""

vllm_omni/deploy/dreamzero.yaml

@@ -14,6 +14,8 @@ stages:
     max_num_seqs: 1
     enforce_eager: false
     model_class_name: DreamZeroPipeline
+    # DreamZero runs on the AR-Diffusion engine (engine-level paged KV cache).
+    engine_backend: vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine
     model_config:
       default_robot_embodiment: roboarena
       policy_server_config:

vllm_omni/deploy/mammoth_moda2.yaml

@@ -1,42 +1,27 @@
-# MammothModa2 AR→DiT (text-to-image). Canonical copy lives at
-# ``vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml``; keep in sync.
-#
-# Must remain legacy ``stage_args`` layout (no top-level ``stages``): Omni resolves
-# explicit ``--stage-configs-path`` via ``load_stage_configs_from_yaml``.
+# MammothModa2 AR→DiT (text-to-image).
 
-stage_args:
+async_chunk: false
+pipeline: mammoth_moda2
+
+stages:
   - stage_id: 0
-    runtime:
-      devices: "0"
-    engine_args:
-      model_stage: ar
-      max_num_seqs: 100
-      model_arch: MammothModa2ForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
-      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      max_model_len: 8192
-      gpu_memory_utilization: 0.5
-      enforce_eager: true
-      trust_remote_code: true
-      engine_output_type: latent
-      enable_prefix_caching: false
-    final_output: false
+    devices: "0"
+    max_num_seqs: 100
+    max_model_len: 8192
+    gpu_memory_utilization: 0.5
+    enforce_eager: true
+    trust_remote_code: true
+    enable_prefix_caching: false
 
   - stage_id: 1
-    runtime:
-      devices: "0"
-    engine_args:
-      model_stage: dit
-      max_num_seqs: 1
-      model_arch: MammothModa2ForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
-      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
-      gpu_memory_utilization: 0.3
-      enforce_eager: true
-      trust_remote_code: true
-      engine_output_type: image
-      enable_prefix_caching: false
-    engine_input_source: [0]
-    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.mammoth_moda2.ar2dit
-    final_output: true
-    final_output_type: image
+    devices: "0"
+    max_num_seqs: 1
+    gpu_memory_utilization: 0.3
+    enforce_eager: true
+    trust_remote_code: true
+    enable_prefix_caching: false
+    default_sampling_params:
+      extra_args:
+        text_guidance_scale: 9.0
+        cfg_range: [0.0, 1.0]
+        num_inference_steps: 50

vllm_omni/deploy/mammoth_moda2_ar.yaml

@@ -0,0 +1,14 @@
+# MammothModa2 AR-only (image/video understanding and summarization).
+
+async_chunk: false
+pipeline: mammoth_moda2_ar
+
+stages:
+  - stage_id: 0
+    devices: "0"
+    max_num_seqs: 16
+    max_model_len: 8192
+    gpu_memory_utilization: 0.5
+    enforce_eager: true
+    trust_remote_code: true
+    enable_prefix_caching: false

vllm_omni/deploy/omnivoice.yaml

@@ -1,19 +1,13 @@
 # OmniVoice TTS deploy: diffusion engine single stage (dit).
 # Text → iterative unmasking → codebook tokens → DAC decode → 24kHz audio.
-# (Migrated from model_executor/stage_configs/omnivoice.yaml for deploy/ layout.)
 
-stage_args:
+async_chunk: false
+pipeline: omnivoice
+
+stages:
   - stage_id: 0
-    stage_type: diffusion
-    is_comprehension: true
-    runtime:
-      devices: 0
-    engine_args:
-      model_stage: dit
-      model_class_name: "OmniVoicePipeline"
-      enforce_eager: true
-      trust_remote_code: true
-      distributed_executor_backend: "mp"
-      dtype: "float32"
-    final_output: true
-    final_output_type: audio
+    devices: "0"
+    enforce_eager: true
+    trust_remote_code: true
+    distributed_executor_backend: "mp"
+    dtype: "float32"

vllm_omni/deploy/qwen3_tts.yaml

@@ -54,7 +54,10 @@ stages:
       temperature: 0.9
       top_k: 50
       max_tokens: 4096
-      seed: 42
+      # Official qwen-tts generates the talker with min_new_tokens=2; without
+      # it, checkpoints that sample codec EOS on the first frame return empty
+      # audio (HTTP 400). See #4962.
+      min_tokens: 2
       repetition_penalty: 1.05
     subtalker_sampling_params:
       do_sample: true
@@ -85,7 +88,6 @@ stages:
       top_p: 1.0
       top_k: -1
       max_tokens: 65536
-      seed: 42
       repetition_penalty: 1.0
 
 platforms:

vllm_omni/deploy/qwen3_tts_high_concurrency.yaml

@@ -60,6 +60,10 @@ stages:
       temperature: 0.9
       top_k: 50
       max_tokens: 4096
+      # Official qwen-tts generates the talker with min_new_tokens=2; without
+      # it, checkpoints that sample codec EOS on the first frame return empty
+      # audio (HTTP 400). See #4962.
+      min_tokens: 2
       repetition_penalty: 1.05
     subtalker_sampling_params:
       do_sample: true

vllm_omni/diffusion/attention/backends/flash_attn_hub.py

@@ -0,0 +1,406 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import threading
+from functools import partial
+
+import torch
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.attention.backends.abstract import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+)
+from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import (
+    piecewise_attn,
+)
+
+logger = init_logger(__name__)
+
+
+_hub_modules: dict[str, object] = {}
+_hub_lock = threading.Lock()
+
+
+def _load_hub_module(repo_id: str):
+    from kernels import get_kernel
+
+    logger.info("Loading %s kernel from HuggingFace Hub...", repo_id)
+    last_error = None
+    for version in (1, 2, None):
+        try:
+            if version is not None:
+                return get_kernel(repo_id, version=version)
+            return get_kernel(repo_id)
+        except Exception as exc:
+            logger.info("Failed to load %s version %s: %s", repo_id, version, exc)
+            last_error = exc
+    raise RuntimeError(f"Failed to load HuggingFace Hub kernel {repo_id!r}") from last_error
+
+
+def _get_hub_module(repo_id: str):
+    with _hub_lock:
+        if repo_id not in _hub_modules:
+            _hub_modules[repo_id] = _load_hub_module(repo_id)
+        return _hub_modules[repo_id]
+
+
+def _run_varlen_dense(
+    varlen_func,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    *,
+    causal: bool,
+    softmax_scale: float,
+) -> torch.Tensor:
+    batch_size, q_len = query.size()[:2]
+    k_len = key.size(1)
+    cu_seqlens_q = torch.arange(0, (batch_size + 1) * q_len, step=q_len, dtype=torch.int32, device=query.device)
+    cu_seqlens_k = torch.arange(0, (batch_size + 1) * k_len, step=k_len, dtype=torch.int32, device=query.device)
+
+    out = varlen_func(
+        q=query.flatten(0, 1),
+        k=key.flatten(0, 1),
+        v=value.flatten(0, 1),
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        max_seqlen_q=q_len,
+        max_seqlen_k=k_len,
+        causal=causal,
+        softmax_scale=softmax_scale,
+    )
+    if isinstance(out, tuple):
+        out = out[0]
+    return out.reshape(batch_size, q_len, *out.shape[1:])
+
+
+class FlashAttentionHubBackend(AttentionBackend):
+    accept_output_buffer: bool = True
+
+    @classmethod
+    def supports_attention_mask(cls) -> bool:
+        return True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [64, 96, 128, 192, 256]
+
+    @staticmethod
+    def get_name() -> str:
+        return "FLASH_ATTN_HUB"
+
+    @staticmethod
+    def get_impl_cls() -> type["FlashAttentionHubImpl"]:
+        return FlashAttentionHubImpl
+
+
+class FlashAttentionHubImpl(AttentionImpl):
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        softmax_scale: float,
+        causal: bool = False,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        qkv_layout: str | None = None,
+        backend_kwargs: dict | None = None,
+        **extra_impl_args,
+    ) -> None:
+        self.num_heads = num_heads
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+        self.qkv_layout = qkv_layout
+        if backend_kwargs:
+            logger.warning("FlashAttentionHubImpl ignoring backend_kwargs: %s", list(backend_kwargs.keys()))
+
+        hub_module = _get_hub_module("kernels-community/flash-attn2")
+        self.flash_attn_func = getattr(hub_module, "flash_attn_func", None)
+        self.flash_attn_varlen_func = getattr(hub_module, "flash_attn_varlen_func", None)
+        if self.flash_attn_func is None and self.flash_attn_varlen_func is None:
+            raise RuntimeError("Failed to load flash-attn2 kernel from HuggingFace Hub: no functions found")
+
+    @staticmethod
+    def _unwrap_flash_output(out: torch.Tensor | tuple[torch.Tensor, ...]) -> torch.Tensor:
+        # FA3 may return (out, lse), FA2 returns out
+        return out[0] if isinstance(out, tuple) else out
+
+    def _flash_wrapper(self, q, k, v, *, attn_func, **kwargs):
+        if attn_func is not None:
+            return self._unwrap_flash_output(attn_func(q, k, v, **kwargs))
+        return _run_varlen_dense(
+            self.flash_attn_varlen_func,
+            q,
+            k,
+            v,
+            causal=kwargs.get("causal", self.causal),
+            softmax_scale=kwargs.get("softmax_scale", self.softmax_scale),
+        )
+
+    def _forward_varlen_masked(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attention_mask: torch.Tensor,
+    ) -> torch.Tensor:
+        from vllm_omni.diffusion.attention.backends.utils.fa import (
+            _pad_input,
+            _unpad_input,
+            _upad_input,
+        )
+
+        assert attention_mask.ndim == 2, "attention_mask must be 2D, (batch_size, seq_len)"
+        query_length = query.size(1)
+        q, k, v, indices_q, (cu_seq_lens_q, cu_seq_lens_k), (max_length_q, max_length_k) = _upad_input(
+            query, key, value, attention_mask, query_length, _unpad_input
+        )
+
+        out_unpad = self.flash_attn_varlen_func(
+            q,
+            k,
+            v,
+            cu_seqlens_q=cu_seq_lens_q,
+            cu_seqlens_k=cu_seq_lens_k,
+            max_seqlen_q=max_length_q,
+            max_seqlen_k=max_length_k,
+            **{
+                "causal": self.causal,
+                "softmax_scale": self.softmax_scale,
+            },
+        )
+        out_unpad = self._unwrap_flash_output(out_unpad)
+        return _pad_input(out_unpad, indices_q, query.size(0), query_length)
+
+    def _forward_varlen_dense(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+    ) -> torch.Tensor:
+        return _run_varlen_dense(
+            self.flash_attn_varlen_func,
+            query,
+            key,
+            value,
+            causal=self.causal,
+            softmax_scale=self.softmax_scale,
+        )
+
+    def forward_cuda(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata = None,
+    ) -> torch.Tensor:
+        attention_mask = attn_metadata.attn_mask if attn_metadata is not None else None
+        full_attn_spans = attn_metadata.full_attn_spans if attn_metadata is not None else None
+
+        # Try piecewise attention
+        if full_attn_spans is not None:
+            logger.debug("Using piecewise Flash Attention for mixed causal/full mask")
+            attn_func = partial(
+                self._flash_wrapper,
+                attn_func=self.flash_attn_func,
+            )
+
+            return piecewise_attn(
+                query,
+                key,
+                value,
+                full_attn_spans,
+                self.softmax_scale,
+                attn_func,
+            )
+
+        if attention_mask is not None and torch.any(~attention_mask):
+            return self._forward_varlen_masked(
+                query,
+                key,
+                value,
+                attention_mask,
+            )
+
+        if self.flash_attn_func is not None:
+            out = self.flash_attn_func(
+                query,
+                key,
+                value,
+                causal=self.causal,
+                softmax_scale=self.softmax_scale,
+            )
+            return self._unwrap_flash_output(out)
+
+        return self._forward_varlen_dense(
+            query,
+            key,
+            value,
+        )
+
+
+class FlashAttention3HubBackend(AttentionBackend):
+    accept_output_buffer: bool = True
+
+    @classmethod
+    def supports_attention_mask(cls) -> bool:
+        return True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [64, 96, 128, 192, 256]
+
+    @staticmethod
+    def get_name() -> str:
+        return "FLASH_ATTN_3_HUB"
+
+    @staticmethod
+    def get_impl_cls() -> type["FlashAttention3HubImpl"]:
+        return FlashAttention3HubImpl
+
+
+class FlashAttention3HubImpl(AttentionImpl):
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        softmax_scale: float,
+        causal: bool = False,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        qkv_layout: str | None = None,
+        backend_kwargs: dict | None = None,
+        **extra_impl_args,
+    ) -> None:
+        self.num_heads = num_heads
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+        self.qkv_layout = qkv_layout
+        if backend_kwargs:
+            logger.warning("FlashAttention3HubImpl ignoring backend_kwargs: %s", list(backend_kwargs.keys()))
+
+        hub_module = _get_hub_module("kernels-community/flash-attn3")
+        self.flash_attn_func = getattr(hub_module, "flash_attn_func", None)
+        self.flash_attn_varlen_func = getattr(hub_module, "flash_attn_varlen_func", None)
+        if self.flash_attn_func is None and self.flash_attn_varlen_func is None:
+            raise RuntimeError("Failed to load flash-attn3 kernel from HuggingFace Hub: no functions found")
+
+    @staticmethod
+    def _unwrap_flash_output(out: torch.Tensor | tuple[torch.Tensor, ...]) -> torch.Tensor:
+        # FA3 returns (out, lse)
+        return out[0] if isinstance(out, tuple) else out
+
+    def _flash_wrapper(self, q, k, v, *, attn_func, **kwargs):
+        if attn_func is not None:
+            return self._unwrap_flash_output(attn_func(q, k, v, **kwargs))
+        return _run_varlen_dense(
+            self.flash_attn_varlen_func,
+            q,
+            k,
+            v,
+            causal=kwargs.get("causal", self.causal),
+            softmax_scale=kwargs.get("softmax_scale", self.softmax_scale),
+        )
+
+    def _forward_varlen_masked(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attention_mask: torch.Tensor,
+    ) -> torch.Tensor:
+        from vllm_omni.diffusion.attention.backends.utils.fa import (
+            _pad_input,
+            _unpad_input,
+            _upad_input,
+        )
+
+        assert attention_mask.ndim == 2, "attention_mask must be 2D, (batch_size, seq_len)"
+        query_length = query.size(1)
+        q, k, v, indices_q, (cu_seq_lens_q, cu_seq_lens_k), (max_length_q, max_length_k) = _upad_input(
+            query, key, value, attention_mask, query_length, _unpad_input
+        )
+
+        out_unpad = self.flash_attn_varlen_func(
+            q,
+            k,
+            v,
+            cu_seqlens_q=cu_seq_lens_q,
+            cu_seqlens_k=cu_seq_lens_k,
+            max_seqlen_q=max_length_q,
+            max_seqlen_k=max_length_k,
+            **{
+                "causal": self.causal,
+                "softmax_scale": self.softmax_scale,
+            },
+        )
+        out_unpad = self._unwrap_flash_output(out_unpad)
+        return _pad_input(out_unpad, indices_q, query.size(0), query_length)
+
+    def _forward_varlen_dense(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+    ) -> torch.Tensor:
+        return _run_varlen_dense(
+            self.flash_attn_varlen_func,
+            query,
+            key,
+            value,
+            causal=self.causal,
+            softmax_scale=self.softmax_scale,
+        )
+
+    def forward_cuda(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata = None,
+    ) -> torch.Tensor:
+        attention_mask = attn_metadata.attn_mask if attn_metadata is not None else None
+        full_attn_spans = attn_metadata.full_attn_spans if attn_metadata is not None else None
+
+        # Try piecewise attention
+        if full_attn_spans is not None:
+            logger.debug("Using piecewise Flash Attention for mixed causal/full mask")
+            attn_func = partial(
+                self._flash_wrapper,
+                attn_func=self.flash_attn_func,
+            )
+
+            return piecewise_attn(
+                query,
+                key,
+                value,
+                full_attn_spans,
+                self.softmax_scale,
+                attn_func,
+            )
+
+        if attention_mask is not None and torch.any(~attention_mask):
+            return self._forward_varlen_masked(
+                query,
+                key,
+                value,
+                attention_mask,
+            )
+
+        if self.flash_attn_func is not None:
+            out = self.flash_attn_func(
+                query,
+                key,
+                value,
+                causal=self.causal,
+                softmax_scale=self.softmax_scale,
+            )
+            return self._unwrap_flash_output(out)
+
+        return self._forward_varlen_dense(
+            query,
+            key,
+            value,
+        )

vllm_omni/diffusion/attention/backends/registry.py

@@ -62,6 +62,8 @@ class DiffusionAttentionBackendEnum(Enum, metaclass=_DiffusionBackendEnumMeta):
     SAGE_ATTN_3 = "vllm_omni.diffusion.attention.backends.sage_attn3.SageAttention3Backend"
     CUDNN_ATTN = "vllm_omni.diffusion.attention.backends.cudnn_attn.CuDNNAttentionBackend"
     FLASHINFER_ATTN = "vllm_omni.diffusion.attention.backends.flashinfer_attn.FlashInferAttentionBackend"
+    FLASH_ATTN_HUB = "vllm_omni.diffusion.attention.backends.flash_attn_hub.FlashAttentionHubBackend"
+    FLASH_ATTN_3_HUB = "vllm_omni.diffusion.attention.backends.flash_attn_hub.FlashAttention3HubBackend"
 
     def get_path(self, include_classname: bool = True) -> str:
         """Get the class path for this backend (respects overrides).

vllm_omni/diffusion/data.py

@@ -6,6 +6,7 @@ import os
 import random
 from collections.abc import Callable, Mapping
 from dataclasses import dataclass, field, fields
+from enum import Enum
 from typing import TYPE_CHECKING, Any
 
 import diffusers
@@ -613,6 +614,20 @@ class OmniDiffusionConfig:
     distributed_executor_backend: str = "mp"
     nccl_port: int | None = None
 
+    # Engine backend selection, resolved by ``DiffusionEngine.resolve_engine_class``
+    # (mirrors ``DiffusionExecutor.get_class``). Config files use a string:
+    # "default" -> DiffusionEngine, or an import path (set e.g. by a deploy
+    # config). Programmatic callers may pass a DiffusionEngine subclass
+    # directly — hence ``str | type`` (structured-config mirrors should expose
+    # the string form only).
+    engine_backend: str | type = "default"
+
+    # Optional override for the diffusion model runner class (import path).
+    # Precedence in the worker: this override > the runner declared by the
+    # selected engine class (``default_diffusion_model_runner_cls``) > the
+    # platform default. Never mutated by engines.
+    diffusion_model_runner_cls: str | None = None
+
     # HuggingFace specific parameters
     trust_remote_code: bool = False
     revision: str | None = None
@@ -1173,7 +1188,8 @@ class OmniDiffusionConfig:
         valid_fields = {f.name for f in fields(cls)}
         filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_fields}
 
-        return cls(**filtered_kwargs)
+        instance = cls(**filtered_kwargs)
+        return instance
 
 
 @dataclass
@@ -1473,5 +1489,12 @@ class OmniWakeTask:
     tags: list[str] | None = None
 
 
+class CuMemTag(str, Enum):
+    """Tags representing specific CuMem allocations for sleep/wake state tracking."""
+
+    WEIGHTS = "weights"
+    KV_CACHE = "kv_cache"
+
+
 # Special message broadcast via scheduler queues to signal worker shutdown.
 SHUTDOWN_MESSAGE = {"type": "shutdown"}

vllm_omni/diffusion/diffusion_engine.py

@@ -134,6 +134,14 @@ class _RpcTask:
 class DiffusionEngine:
     """The diffusion engine for vLLM-Omni diffusion models."""
 
+    #: Import path of the model runner this engine's workers should build, or
+    #: ``None`` for the platform default. Subclasses declare their runner here
+    #: (e.g. the AR-Diffusion engine), the worker resolves it through
+    #: :meth:`resolve_engine_class` — so engine->runner routing lives on the
+    #: engine class itself, and ``od_config.diffusion_model_runner_cls``
+    #: remains a pure explicit user override (never mutated by engines).
+    default_diffusion_model_runner_cls: str | None = None
+
     def __init__(
         self,
         od_config: OmniDiffusionConfig,
@@ -609,19 +617,55 @@ class DiffusionEngine:
             self._put_streaming_output_with_cv(request_id, out)
 
     @staticmethod
-    def make_engine(
-        config: OmniDiffusionConfig,
-        scheduler: SchedulerInterface | None = None,
-    ) -> DiffusionEngine:
-        """Factory method to create a DiffusionEngine instance.
+    def resolve_engine_class(config: OmniDiffusionConfig) -> type[DiffusionEngine]:
+        """Resolve the engine class selected by ``config.engine_backend``.
+
+        Mirrors ``DiffusionExecutor.get_class``: accepts ``"default"``, a
+        ``DiffusionEngine`` subclass, or an import-path string (e.g. a deploy
+        config's ``engine_backend``). Kept separate from :meth:`make_engine` so the
+        selection is testable without constructing an engine (which runs a dummy forward).
 
         Args:
             config: The configuration for the diffusion engine.
 
         Returns:
-            An instance of DiffusionEngine.
+            The ``DiffusionEngine`` (sub)class to instantiate.
         """
-        return DiffusionEngine(config, scheduler=scheduler)
+        backend = getattr(config, "engine_backend", "default") or "default"
+
+        if isinstance(backend, type):
+            if not issubclass(backend, DiffusionEngine):
+                raise TypeError(f"engine_backend must be a DiffusionEngine subclass. Got {backend}.")
+            return backend
+        if backend == "default":
+            return DiffusionEngine
+        if isinstance(backend, str):
+            try:
+                engine_class = resolve_obj_by_qualname(backend)
+            except (ImportError, ValueError) as e:
+                raise ValueError(
+                    f"Failed to load engine_backend '{backend}'. Ensure it is a valid python path. Error: {e}"
+                ) from e
+            if not issubclass(engine_class, DiffusionEngine):
+                raise TypeError(f"engine_backend must resolve to a DiffusionEngine subclass. Got {engine_class}.")
+            return engine_class
+        raise ValueError(f"Unknown engine_backend: {backend!r}")
+
+    @staticmethod
+    def make_engine(
+        config: OmniDiffusionConfig,
+        scheduler: SchedulerInterface | None = None,
+    ) -> DiffusionEngine:
+        """Factory method to create the engine selected by ``config.engine_backend``.
+
+        Args:
+            config: The configuration for the diffusion engine.
+
+        Returns:
+            An instance of the resolved ``DiffusionEngine`` (sub)class.
+        """
+        engine_class = DiffusionEngine.resolve_engine_class(config)
+        return engine_class(config, scheduler=scheduler)
 
     def add_request(self, request: OmniDiffusionRequest) -> str:
         with self._cv:

vllm_omni/diffusion/distributed/parallel_state.py

@@ -512,6 +512,21 @@ def init_model_parallel_group(
         )
 
 
+def init_vllm_model_parallel_group(
+    group_ranks: list[list[int]],
+    local_rank: int,
+    backend: str,
+    group_name: str,
+) -> vllm_parallel_state.GroupCoordinator:
+    return vllm_parallel_state.init_model_parallel_group(
+        group_ranks=group_ranks,
+        local_rank=local_rank,
+        backend=backend,
+        group_name=group_name,
+        use_device_communicator=True,
+    )
+
+
 def init_dit_group(
     dit_parallel_size: int,
     backend: str,
@@ -838,23 +853,40 @@ def initialize_model_parallel(
     if use_moe_parallel_mapping:
         # Diffusion normally uses its own SP group. Map it to vLLM PCP only for
         # expert-parallel runtimes that rely on vLLM FusedMoE group semantics.
-        vllm_parallel_state._PCP = _SP
+        # vLLM 0.24 MoE kernels require GroupCoordinator.device_communicator
+        # and reduce_scatter(), which the diffusion SP coordinator intentionally
+        # does not own. Keep the rank layout but build a vLLM coordinator.
+        vllm_parallel_state._PCP = init_vllm_model_parallel_group(
+            group_ranks=sp_group_ranks,
+            local_rank=get_world_group().local_rank,
+            backend=backend,
+            group_name="pcp",
+        )
 
     assert vllm_parallel_state._TP is None, "Tensor parallel group is already initialized"
-    vllm_parallel_state._TP = init_model_parallel_group(
-        group_ranks=rank_generator.get_ranks("tp"),
-        local_rank=get_world_group().local_rank,
-        backend=backend,
-        parallel_mode="tensor",
-    )
-    if use_moe_parallel_mapping and cfg_parallel_size > 1:
+    tp_group_ranks = rank_generator.get_ranks("tp")
+    if use_moe_parallel_mapping:
+        vllm_parallel_state._TP = init_vllm_model_parallel_group(
+            group_ranks=tp_group_ranks,
+            local_rank=get_world_group().local_rank,
+            backend=backend,
+            group_name="tp",
+        )
+    else:
+        vllm_parallel_state._TP = init_model_parallel_group(
+            group_ranks=tp_group_ranks,
+            local_rank=get_world_group().local_rank,
+            backend=backend,
+            parallel_mode="tensor",
+        )
+    if use_moe_parallel_mapping:
         # CFG is a diffusion-specific replica dimension. Fold it into vLLM DP
         # only when constructing the vLLM EP layout for expert-parallel paths.
-        vllm_parallel_state._DP = init_model_parallel_group(
+        vllm_parallel_state._DP = init_vllm_model_parallel_group(
             group_ranks=rank_generator.get_ranks("cfg-dp"),
             local_rank=get_world_group().local_rank,
             backend=backend,
-            parallel_mode="data",
+            group_name="dp",
         )
 
     global _FS, _EXPERT_PARALLEL_GROUP_RANKS
@@ -869,11 +901,11 @@ def initialize_model_parallel(
     _EXPERT_PARALLEL_GROUP_RANKS = None
     if use_moe_parallel_mapping:
         ep_group_ranks = rank_generator.get_ranks("tp-sp-cfg-dp")
-        vllm_parallel_state._EP = init_model_parallel_group(
+        vllm_parallel_state._EP = init_vllm_model_parallel_group(
             group_ranks=ep_group_ranks,
             local_rank=get_world_group().local_rank,
             backend=backend,
-            parallel_mode="expert",
+            group_name="ep",
         )
         _EXPERT_PARALLEL_GROUP_RANKS = ep_group_ranks
 

vllm_omni/diffusion/layers/rope.py

@@ -188,6 +188,12 @@ class RotaryEmbedding(CustomOp):
         cos: torch.Tensor,
         sin: torch.Tensor,
     ) -> torch.Tensor:
+        # All batch elements share the same rotary position encoding.
+        # Strip the batch dim so the underlying op broadcasts over the batch,
+        # consistent with forward_cuda / forward_hip / apply_rotary_emb_mindiesd.
+        if cos.dim() == 3:
+            cos = cos[0]
+            sin = sin[0]
         return apply_rotary_emb_torch(
             x,
             cos,