[Bugfix] Update routed-expert tests for native auxiliary output (#8500)

Signed-off-by: Yueqian Lin <linyueqian@outlook.com>

.buildkite/amd/test-amd-merge.yml

@@ -91,6 +91,25 @@ steps:
 
 - group: ":card_index_dividers: Diffusion Test"
   steps:
+    # AR paged attention is real GPU work, including cold AITER compilation.
+    # Keep its complete numerical matrix out of the CPU-marked Simple shards.
+    - label: "ROCm · AR Diffusion Paged Attention GPU Test"
+      agent_pool: mi300_1
+      depends_on: amd-build
+      mirror_hardwares: [amdproduction]
+      grade: Blocking
+      timeout_in_minutes: 20
+      commands:
+        - export VLLM_ROCM_USE_AITER=1
+        - export VLLM_OMNI_AR_FA_REQUIRED=1
+        - export AITER_JIT_DIR="/tmp/vllm-omni-ar-aiter-$$BUILDKITE_JOB_ID"
+        - export TORCH_EXTENSIONS_DIR="/tmp/vllm-omni-ar-torch-extensions-$$BUILDKITE_JOB_ID"
+        - mkdir -p "$$AITER_JIT_DIR" "$$TORCH_EXTENSIONS_DIR"
+        - >-
+          timeout --signal=TERM --kill-after=2m 15m
+          pytest -sv tests/diffusion/ar_diffusion/test_paged_attention.py
+          -m 'core_model and rocm and MI325 and cards_1' --run-level core_model
+
     - label: "Diffusion · LoRA Test"
       agent_pool: mi300_1
       depends_on: amd-build

.buildkite/amd/test-amd-ready.yml

@@ -129,6 +129,25 @@ steps:
           -m 'core_model and rocm and MI325 and cards_1'
           --run-level "core_model"
 
+    # AR paged attention is real GPU work, including cold AITER compilation.
+    # Keep its complete numerical matrix out of the CPU-marked Simple shards.
+    - label: "ROCm · AR Diffusion Paged Attention GPU Test"
+      agent_pool: mi300_1
+      depends_on: amd-build
+      mirror_hardwares: [amdproduction]
+      grade: Blocking
+      timeout_in_minutes: 20
+      commands:
+        - export VLLM_ROCM_USE_AITER=1
+        - export VLLM_OMNI_AR_FA_REQUIRED=1
+        - export AITER_JIT_DIR="/tmp/vllm-omni-ar-aiter-$$BUILDKITE_JOB_ID"
+        - export TORCH_EXTENSIONS_DIR="/tmp/vllm-omni-ar-torch-extensions-$$BUILDKITE_JOB_ID"
+        - mkdir -p "$$AITER_JIT_DIR" "$$TORCH_EXTENSIONS_DIR"
+        - >-
+          timeout --signal=TERM --kill-after=2m 15m
+          pytest -sv tests/diffusion/ar_diffusion/test_paged_attention.py
+          -m 'core_model and rocm and MI325 and cards_1' --run-level core_model
+
     - label: "Diffusion · Batch Test"
       agent_pool: mi300_1
       depends_on: amd-build

.buildkite/common/scripts/upload_pipeline.py

@@ -102,6 +102,7 @@ CI_SOURCE_FILE_DEPENDENCIES_PATH = ROOT / ".buildkite/common/ci_source_file_depe
 SOURCE_FILTER_FALLBACK_KEY = "source_filter_fallback"
 CUDA_HF_TOKEN_ENV = "VLLM_CI_HF_TOKEN"
 CUDA_HF_TOKEN_EXPORT = f'if [ -n "$${{{CUDA_HF_TOKEN_ENV}:-}}" ]; then export HF_TOKEN="$${{{CUDA_HF_TOKEN_ENV}}}"; fi'
+CUDA_RUNTIME_CHECK = "python3 .buildkite/cuda/scripts/check_vllm_runtime.py"
 
 # Bootstrap Buildkite ``if`` expressions.
 # ``*_MAIN_IF``: main + env schedule. ``*_LABEL_IF``: PR label (and/or composed with MAIN).
@@ -620,11 +621,11 @@ def _expand_mirror_hardwares(step: dict[str, Any]) -> dict[str, Any] | None:
     if any(preset_name == chip or preset_name.startswith(f"{chip}_") for chip in _cuda_mirror_chips()):
         commands = merged.get("commands")
         if isinstance(commands, list):
-            merged["commands"] = [CUDA_HF_TOKEN_EXPORT, *commands]
+            merged["commands"] = [CUDA_HF_TOKEN_EXPORT, CUDA_RUNTIME_CHECK, *commands]
         elif commands is None:
-            merged["commands"] = [CUDA_HF_TOKEN_EXPORT]
+            merged["commands"] = [CUDA_HF_TOKEN_EXPORT, CUDA_RUNTIME_CHECK]
         else:
-            merged["commands"] = [CUDA_HF_TOKEN_EXPORT, commands]
+            merged["commands"] = [CUDA_HF_TOKEN_EXPORT, CUDA_RUNTIME_CHECK, commands]
     # Preset retry (K8S_RETRY on l4_*) must not clobber a step that opted out.
     if "retry" in step:
         merged["retry"] = step["retry"]

.buildkite/cuda/scripts/check_vllm_runtime.py

@@ -0,0 +1,38 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Verify the pinned vLLM and compiled CUDA imports on a GPU CI worker."""
+
+from __future__ import annotations
+
+import importlib
+import re
+from pathlib import Path
+
+
+def main() -> None:
+    dockerfile = Path(__file__).resolve().parents[3] / "docker" / "Dockerfile.ci"
+    match = re.search(
+        r"^ARG VLLM_VERSION=([0-9]+\.[0-9]+\.[0-9]+)$",
+        dockerfile.read_text(),
+        re.MULTILINE,
+    )
+    if match is None:
+        raise RuntimeError(f"Missing vLLM release version in {dockerfile}")
+    expected_version = match.group(1)
+
+    import torch
+    import vllm
+
+    assert vllm.__version__ == expected_version, vllm.__version__
+    assert torch.cuda.is_available(), "vLLM runtime check requires a GPU worker"
+    for module in ("flashinfer", "vllm._C_stable_libtorch"):
+        importlib.import_module(module)
+    print(
+        f"Verified vLLM CUDA runtime: vllm={vllm.__version__} "
+        f"expected={expected_version} torch={torch.__version__} "
+        f"cuda={torch.version.cuda} device={torch.cuda.get_device_name()}"
+    )
+
+
+if __name__ == "__main__":
+    main()

.buildkite/cuda/test-merge.yml

@@ -38,6 +38,17 @@ steps:
   - group: ":card_index_dividers: Diffusion Test"
     depends_on: upload-merge-pipeline
     steps:
+      - label: "Diffusion · AR Paged Attention GPU Test"
+        timeout_in_minutes: 20
+        commands:
+          - export VLLM_OMNI_AR_FA_REQUIRED=1
+          - >-
+            timeout 15m pytest -sv
+            tests/diffusion/ar_diffusion/test_paged_attention.py
+            -m 'core_model and cuda and L4 and cards_1'
+            --run-level "core_model"
+        mirror_hardwares: l4_1
+
       - label: "Diffusion · LoRA Test"
         timeout_in_minutes: 30
         commands:

.buildkite/cuda/test-nightly.yml

@@ -48,6 +48,8 @@ steps:
       - label: ":full_moon: Omni · Doc Test with H100 · 2-GPU"
         source_file_dependencies: omni_qwen3_omni_doc
         timeout_in_minutes: 90
+        artifact_paths:
+          - audio_*.wav
         commands:
           - pytest -s -v tests/examples/ -m "full_model and omni and H100 and B200 and cards_2" --run-level "full_model"
 
@@ -156,6 +158,11 @@ steps:
       - label: ":full_moon: TTS · Function Test with L4"
         source_file_dependencies: tts_qwen3_tts_function
         timeout_in_minutes: 120
+        env:
+          VLLM_OMNI_FAILED_SPEECH_AUDIO_DIR: .
+        artifact_paths:
+          - failed_speech_*.wav
+          - failed_speech_*.json
         commands:
           - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
           - export VLLM_USE_DEEP_GEMM="0"

benchmarks/diffusion/benchmark_mammoth_moda2_qk_norm_rope.py

@@ -0,0 +1,262 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Benchmark MammothModa2's Q/K RMSNorm + adjacent-pair real RoPE.
+
+Example:
+
+    CUDA_VISIBLE_DEVICES=0 python \
+      benchmarks/diffusion/benchmark_mammoth_moda2_qk_norm_rope.py \
+      --tokens 512 4173 8346 --warmup 50 --iters 200
+"""
+
+from __future__ import annotations
+
+import argparse
+import importlib.metadata
+import json
+import os
+import statistics
+from collections.abc import Callable
+
+import torch
+
+from vllm_omni.diffusion.layers.fused_qk_norm_rope import (
+    _fused_cuda_supported,
+    fused_qk_norm_rope,
+)
+from vllm_omni.diffusion.layers.norm import RMSNorm
+from vllm_omni.diffusion.models.mammoth_moda2.rope_real import (
+    apply_real_rotary_emb,
+)
+
+_HEAD_DIM = 120
+_Q_HEADS = 21
+_KV_HEADS = 7
+_EPS = 1e-5
+
+
+def _parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--tokens", nargs="+", type=int, default=[512, 4173, 8346])
+    parser.add_argument("--warmup", type=int, default=50)
+    parser.add_argument("--iters", type=int, default=200)
+    parser.add_argument("--rounds", type=int, default=3)
+    parser.add_argument("--attention-sequences", nargs="*", type=int, default=[])
+    return parser.parse_args()
+
+
+def _version(distribution: str) -> str | None:
+    try:
+        return importlib.metadata.version(distribution)
+    except importlib.metadata.PackageNotFoundError:
+        return None
+
+
+def _measure(
+    fn: Callable[[], object],
+    warmup: int,
+    iters: int,
+) -> dict[str, float]:
+    for _ in range(warmup):
+        fn()
+    torch.accelerator.synchronize()
+
+    torch.accelerator.reset_peak_memory_stats()
+    baseline_bytes = torch.accelerator.memory_allocated()
+    samples = []
+    for _ in range(iters):
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+        start.record()
+        fn()
+        end.record()
+        end.synchronize()
+        samples.append(float(start.elapsed_time(end)))
+
+    ordered = sorted(samples)
+    return {
+        "median_ms": statistics.median(ordered),
+        "mean_ms": statistics.mean(ordered),
+        "stddev_ms": statistics.pstdev(ordered),
+        "min_ms": ordered[0],
+        "max_ms": ordered[-1],
+        "p90_ms": ordered[min(len(ordered) - 1, int(len(ordered) * 0.9))],
+        "extra_peak_mib": (torch.accelerator.max_memory_allocated() - baseline_bytes) / 2**20,
+    }
+
+
+@torch.inference_mode()
+def _run_shape(tokens: int, warmup: int, iters: int, *, fused_first: bool = False) -> dict[str, object]:
+    dtype = torch.bfloat16
+    device = torch.device("cuda")
+    torch.manual_seed(42)
+    query = torch.randn(tokens, _Q_HEADS, _HEAD_DIM, device=device, dtype=dtype)
+    key = torch.randn(tokens, _KV_HEADS, _HEAD_DIM, device=device, dtype=dtype)
+    norm_q = RMSNorm(_HEAD_DIM, eps=_EPS).to(device=device, dtype=dtype)
+    norm_k = RMSNorm(_HEAD_DIM, eps=_EPS).to(device=device, dtype=dtype)
+    norm_q.weight.normal_(mean=1.0, std=0.2)
+    norm_k.weight.normal_(mean=1.0, std=0.2)
+    angles = torch.rand(tokens, _HEAD_DIM // 2, device=device) * (2 * torch.pi)
+    cos = angles.cos().repeat_interleave(2, dim=-1).to(dtype)
+    sin = angles.sin().repeat_interleave(2, dim=-1).to(dtype)
+
+    def native() -> tuple[torch.Tensor, torch.Tensor]:
+        q_out = apply_real_rotary_emb(norm_q(query).unsqueeze(0), cos.unsqueeze(0), sin.unsqueeze(0)).squeeze(0)
+        k_out = apply_real_rotary_emb(norm_k(key).unsqueeze(0), cos.unsqueeze(0), sin.unsqueeze(0)).squeeze(0)
+        return q_out, k_out
+
+    def fused() -> tuple[torch.Tensor, torch.Tensor]:
+        rope_table = torch.cat((cos[..., 0::2], sin[..., 0::2]), dim=-1)
+        return fused_qk_norm_rope(
+            query,
+            key,
+            norm_q.weight,
+            norm_k.weight,
+            rope_table,
+            _EPS,
+            head_dim=_HEAD_DIM,
+            rotary_dim=_HEAD_DIM,
+            interleaved=True,
+        )
+
+    expected_q, expected_k = native()
+    actual_q, actual_k = fused()
+    errors = {
+        "q_max_abs": (actual_q.float() - expected_q.float()).abs().max().item(),
+        "q_mean_abs": (actual_q.float() - expected_q.float()).abs().mean().item(),
+        "k_max_abs": (actual_k.float() - expected_k.float()).abs().max().item(),
+        "k_mean_abs": (actual_k.float() - expected_k.float()).abs().mean().item(),
+    }
+    torch.testing.assert_close(actual_q, expected_q, atol=0.0625, rtol=0.02)
+    torch.testing.assert_close(actual_k, expected_k, atol=0.0625, rtol=0.02)
+
+    functions = {"native": native, "fused": fused}
+    order = ("fused", "native") if fused_first else ("native", "fused")
+    measured = {name: _measure(functions[name], warmup, iters) for name in order}
+    native_stats, fused_stats = measured["native"], measured["fused"]
+    return {
+        "tokens": tokens,
+        "native": native_stats,
+        "fused_including_rope_pack": fused_stats,
+        "median_speedup": native_stats["median_ms"] / fused_stats["median_ms"],
+        "errors": errors,
+    }
+
+
+@torch.inference_mode()
+def _run_attention(seq: int, warmup: int, iters: int, *, fused_first: bool = False) -> dict[str, object]:
+    from vllm_omni.diffusion.models.mammoth_moda2.mammothmoda2_dit_model import TransformerBlock
+
+    torch.manual_seed(42)
+    block = (
+        TransformerBlock(
+            2520,
+            _Q_HEADS,
+            _KV_HEADS,
+            multiple_of=256,
+            ffn_dim_multiplier=1.0,
+            norm_eps=_EPS,
+            rope_repeats_pairs=True,
+        )
+        .to(device="cuda", dtype=torch.bfloat16)
+        .eval()
+    )
+    block.attn.norm_q.weight.normal_(1.0, 0.2)
+    block.attn.norm_k.weight.normal_(1.0, 0.2)
+    hidden = torch.randn(2, seq, 2520, device="cuda", dtype=torch.bfloat16)
+    mask = torch.ones(2, seq, device="cuda", dtype=torch.bool)
+    mask[0, seq - min(300, seq // 2) :] = False
+    mask[1, seq - min(17, seq // 2) :] = False
+    angles = torch.rand(2, seq, _HEAD_DIM // 2, device="cuda") * (2 * torch.pi)
+    # The production producer supplies FP32 tables; each example has its own positions.
+    rotary = angles.cos().repeat_interleave(2, dim=-1), angles.sin().repeat_interleave(2, dim=-1)
+    env = "VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS"
+    previous = os.environ.get(env)
+
+    def run():
+        return block.attn(hidden, hidden, attention_mask=mask, image_rotary_emb=rotary)
+
+    try:
+        outputs = {}
+        timings = {}
+        order = ("fused", "native") if fused_first else ("native", "fused")
+        for name in order:
+            os.environ[env] = "0" if name == "fused" else str(2 * seq + 1)
+            outputs[name] = run()
+            timings[name] = _measure(run, warmup, iters)
+        diff = (outputs["native"].float() - outputs["fused"].float()).abs()[mask]
+        max_diff, mean_diff = diff.max().item(), diff.mean().item()
+        assert max_diff < 2e-2 and mean_diff < 1e-3, (max_diff, mean_diff)
+        assert torch.count_nonzero(outputs["fused"][~mask]) == 0
+        return {
+            "sequence": seq,
+            "batch": 2,
+            "backend": block.attn.omni_attn.attn_backend.__name__,
+            "native": timings["native"],
+            "fused": timings["fused"],
+            "median_speedup": timings["native"]["median_ms"] / timings["fused"]["median_ms"],
+            "max_abs_diff": max_diff,
+            "mean_abs_diff": mean_diff,
+        }
+    finally:
+        if previous is None:
+            os.environ.pop(env, None)
+        else:
+            os.environ[env] = previous
+
+
+def main() -> None:
+    args = _parse_args()
+    if (
+        args.rounds < 1
+        or args.iters < 1
+        or args.warmup < 0
+        or any(value < 1 for value in [*args.tokens, *args.attention_sequences])
+    ):
+        raise ValueError("rounds, iters and sequence lengths must be positive; warmup must be nonnegative")
+    if not torch.cuda.is_available():
+        raise RuntimeError("This benchmark requires CUDA")
+
+    probe_q = torch.empty(1, _Q_HEADS, _HEAD_DIM, device="cuda", dtype=torch.bfloat16)
+    probe_k = torch.empty(1, _KV_HEADS, _HEAD_DIM, device="cuda", dtype=torch.bfloat16)
+    if not _fused_cuda_supported(probe_q, probe_k, _HEAD_DIM, _HEAD_DIM, interleaved=True):
+        raise RuntimeError("The fused CUDA QK norm/RoPE path is unavailable")
+
+    result = {
+        "environment": {
+            "gpu": torch.cuda.get_device_name(),
+            "torch": torch.__version__,
+            "cuda": torch.version.cuda,
+            "triton": _version("triton"),
+            "vllm": _version("vllm"),
+            "vllm_omni": _version("vllm-omni"),
+        },
+        "config": {
+            "dtype": "bfloat16",
+            "q_heads": _Q_HEADS,
+            "kv_heads": _KV_HEADS,
+            "head_dim": _HEAD_DIM,
+            "warmup": args.warmup,
+            "iters": args.iters,
+            "rounds": args.rounds,
+            "inference_mode": True,
+            "alternating_order": True,
+            "norm_weights": "normal(mean=1.0, std=0.2)",
+        },
+        "results": [
+            {"round": round_id, **_run_shape(tokens, args.warmup, args.iters, fused_first=bool(round_id % 2))}
+            for round_id in range(args.rounds)
+            for tokens in args.tokens
+        ],
+        "attention_results": [
+            {"round": round_id, **_run_attention(seq, args.warmup, args.iters, fused_first=bool(round_id % 2))}
+            for round_id in range(args.rounds)
+            for seq in args.attention_sequences
+        ],
+    }
+    print(json.dumps(result, indent=2))
+
+
+if __name__ == "__main__":
+    main()

benchmarks/diffusion/performance_dashboard/mammoth_moda2_serving_performance.md

@@ -0,0 +1,209 @@
+# MammothModa2-Preview Serving Performance Dashboard
+
+This document describes how to deploy and benchmark **bytedance-research/MammothModa2-Preview** (T2I) using vLLM-Omni. It includes service startup configuration, benchmark methodology, dataset settings, and performance results.
+
+MammothModa2 is a two-stage AR→DiT text-to-image pipeline: stage 0 is an autoregressive LLM that encodes the prompt into a visual-token grid, and stage 1 is a generation-LLM DiT that renders the final image (`final_output_type: image`). Unlike single-stage diffusion pipelines, `/v1/images/generations` requests traverse both stages, and the DiT stage ran with `max_num_seqs: 1` in the then-default deploy config, so concurrent requests serialized on the DiT stage while the AR stage batched (request-level DiT batching landed later in #7476).
+
+> **Legacy-topology baseline**: all numbers in this dashboard were measured on the legacy topology, where the MammothModa2 DiT stage runs as an LLM-generation stage (`stage_type: llm`) built from `VllmConfig`, served through the LLM-typed image-stage compatibility path. #7134 migrates this pipeline to the shared diffusion runtime (DiT registered as a standard `DIFFUSION` stage built from `OmniDiffusionConfig`). These results should be treated as the pre-migration baseline and are expected to be rerun after #7134 lands; they are not directly comparable with post-migration numbers.
+
+---
+
+# 1. Overview
+
+This document covers:
+
+* Service launch configuration
+* Benchmark scripts and usage
+* Dataset and workload settings
+* Performance measurement results
+* Reproducibility guidelines
+
+---
+
+# 2. Test Environment
+
+| Component | Specification |
+|------------|----------------|
+| GPU | NVIDIA A800-SXM4-80GB (1x) |
+| CPU | 18 cores |
+| Model weights | `bytedance-research/MammothModa2-Preview` |
+| Deploy config | `vllm_omni/deploy/mammoth_moda2.yaml` (bundled) |
+| Stage 0 (AR) | `gpu_memory_utilization: 0.5`, `enforce_eager: true`, `max_num_seqs: 100` |
+| Stage 1 (DiT) | `gpu_memory_utilization: 0.3`, `enforce_eager: true`, `max_num_seqs: 1` |
+| Sampling params (per request) | `text_guidance_scale: 9.0`, `cfg_range: [0.0, 1.0]`, `num_inference_steps: 50` sent via `--extra-body` / `--num-inference-steps` |
+
+---
+
+# 3. Service Launch Configuration
+
+## 3.1 Basic Serving Command
+
+```bash
+vllm-omni serve bytedance-research/MammothModa2-Preview --omni \
+    --port 8091 \
+    --trust-remote-code
+```
+
+When serving from a local weights directory (no `model_index.json`), pass the bundled deploy config explicitly so the `mammoth_moda2` pipeline is selected:
+
+```bash
+vllm-omni serve /path/to/MammothModa2-Preview --omni \
+    --port 8091 \
+    --trust-remote-code \
+    --deploy-config vllm_omni/deploy/mammoth_moda2.yaml
+```
+
+## 3.2 Key Parameters
+
+| Parameter | Description |
+|-----------|-------------|
+| `--deploy-config` | Path to the deploy YAML; defaults to the bundled `mammoth_moda2.yaml` for the HF checkpoint |
+| `size` (request) | `WxH`, e.g. `1024x1024`; routed to the AR grid and the DiT `target_w/target_h` |
+| `num_inference_steps` (request) | DiT sampling steps (default 50) |
+| `guidance_scale` (request) | Text guidance scale (default 9.0) |
+
+Note: the benchmark script's `--warmup-num-inference-steps` shortcut is not usable for this pipeline — a reduced step count truncates the AR stage's visual-token output and crashes the DiT stage. Use `--warmup-requests 0` and warm up with a real request instead.
+
+---
+
+# 4. Benchmark Script
+
+## 4.1 Benchmark Entry
+
+```bash
+python benchmarks/diffusion/diffusion_benchmark_serving.py \
+    --base-url http://localhost:8091 \
+    --model bytedance-research/MammothModa2-Preview \
+    --endpoint /v1/images/generations \
+    --task t2i \
+    --dataset random \
+    --num-prompts 8 \
+    --width <W> --height <H> \
+    --num-inference-steps <S> \
+    --extra-body '{"text_guidance_scale": 9.0, "cfg_range": [0.0, 1.0]}' \
+    --max-concurrency <C> \
+    --warmup-requests 0 \
+    --seed 142 \
+    --output-file <OUT>.json
+```
+
+## 4.2 Dataset
+
+`random` synthetic T2I prompts, fixed seed 142, 8 prompts per configuration (4 for 2048x2048 due to per-request latency).
+
+---
+
+# 5. Performance Results
+
+All requests succeeded (0 failures across the sweep). Latency is end-to-end per request (AR + DiT + VAE decode + base64 response).
+
+> Note: these are legacy-topology baseline results (see the callout in the overview); rerun after the #7134 diffusion-runtime migration before comparing across topologies.
+
+## 5.1 Latency
+
+| Resolution | Steps | Concurrency | Mean (s) | P50 (s) | P95 (s) | P99 (s) |
+|------------|-------|-------------|----------|---------|---------|---------|
+| 512x512    | 20    | 1           | 21.8     | 21.7    | 22.1    | 22.2    |
+| 512x512    | 20    | 2           | 24.8     | 25.2    | 25.6    | 25.8    |
+| 512x512    | 20    | 4           | 28.5     | 29.2    | 31.1    | 31.8    |
+| 1024x1024  | 50    | 1           | 96.3     | 95.8    | 99.8    | 100.8   |
+| 1024x1024  | 50    | 2           | 121.5    | 125.4   | 127.9   | 128.5   |
+| 1024x1024  | 50    | 4           | 150.9    | 157.6   | 178.6   | 186.3   |
+| 2048x2048  | 50    | 1           | 271.9    | 271.6   | 273.6   | 273.8   |
+
+## 5.2 Throughput
+
+| Resolution | Steps | Concurrency | Throughput (img/s) |
+|------------|-------|-------------|--------------------|
+| 512x512    | 20    | 1           | 0.0459             |
+| 512x512    | 20    | 2           | 0.0798             |
+| 512x512    | 20    | 4           | 0.1329             |
+| 1024x1024  | 50    | 1           | 0.0104             |
+| 1024x1024  | 50    | 2           | 0.0162             |
+| 1024x1024  | 50    | 4           | 0.0238             |
+| 2048x2048  | 50    | 1           | 0.0037             |
+
+## 5.3 Stage Breakdown (concurrency 1, from server logs)
+
+Per-stage wall time from the server's `[OmniTiming]` records across the §5.1/§5.4/§5.5 sweeps (n = 18/9/5 per row; at concurrency 1 stage wall time equals `stage_gen_time_ms` in the engine stats table):
+
+| Resolution | Steps | AR stage (s) | DiT stage (s) | E2E (s) | AR share | AR output tokens | AR ms/token |
+|------------|-------|--------------|---------------|---------|----------|------------------|-------------|
+| 512x512    | 20    | 20.3         | 1.7           | 22.0    | 92%      | 1,057            | 19.2        |
+| 1024x1024  | 50    | 78.9         | 17.2          | 96.1    | 82%      | 4,161            | 19.0        |
+| 2048x2048  | 50    | 156.3        | 114.8         | 271.2   | 58%      | 8,154            | 19.2        |
+
+Stage 0 is the AR prompt encoder (TTFT 82 ms at 1024x1024); stage 1 is the DiT renderer. Three observations:
+
+1. **AR per-token cost is constant (~19 ms/token) across resolutions**, so AR time is fully predictable: the AR stage decodes a fixed visual grid whose length scales with resolution (1,057 → 4,161 → 8,154 tokens), giving `AR ≈ tokens × 19 ms`. This is also why §5.4's step sweep is sub-linear — the AR token count is step-invariant.
+2. **DiT time scales super-linearly with pixels**: a 16x pixel increase (512 → 2048) multiplies DiT time by 67 (1.7 → 114.8 s).
+3. **The dominant stage flips with resolution**: AR is 92% of E2E at 512x512 (DiT optimization is nearly pointless there) but 58% at 2048x2048, where both stages matter.
+
+At concurrency > 1, per-stage wall times additionally include intra-stage queueing (each stage serializes at `max_num_seqs: 1`), so the c2/c4 rows in §5.1 cannot be decomposed additively from this table.
+
+## 5.4 Inference-Step Sensitivity (1024x1024 / 512x512, concurrency 1)
+
+Steps 25 at 1024x1024 (half the default 50) and steps 25 at 512x512, same protocol as §5.1:
+
+| Resolution | Steps | Concurrency | Mean (s) | P50 (s) | P95 (s) | P99 (s) |
+|------------|-------|-------------|----------|---------|---------|---------|
+| 1024x1024  | 25    | 1           | 85.2     | 85.3    | 87.8    | 88.7    |
+| 1024x1024  | 25    | 4           | 113.7    | 115.8   | 126.1   | 129.4   |
+| 512x512    | 25    | 1           | 21.9     | 21.8    | 22.2    | 22.3    |
+
+Halving DiT steps does not halve end-to-end latency (96.3 → 85.2 s at 1024): the AR stage emits a fixed count of visual tokens (4,161 at 1024x1024) regardless of step count, so its ~79 s contribution (§5.3) is step-invariant; only the DiT sampling time scales with steps. At 512x512, steps 20 → 25 adds ~0.1 s per step (21.8 → 21.9 s).
+
+## 5.5 Prompt-Length Sensitivity (long vs short prompts)
+
+A ~430-word long prompt (repeated via `--random-request-config '[{"weight":1,"prompt":"..."}]'`) against the default short synthetic prompt, same seed and settings:
+
+| Resolution | Steps | Concurrency | Prompt | Mean (s) | P99 (s) | Throughput (img/s) |
+|------------|-------|-------------|--------|----------|---------|---------------------|
+| 1024x1024  | 50    | 1           | short  | 96.3     | 100.8   | 0.0104              |
+| 1024x1024  | 50    | 1           | long   | 93.6     | 94.4    | 0.0107              |
+| 1024x1024  | 50    | 4           | short  | 150.9    | 186.3   | 0.0238              |
+| 1024x1024  | 50    | 4           | long   | 141.7    | 174.2   | 0.0255              |
+| 512x512    | 20    | 1           | short  | 21.8     | 22.2    | 0.0459              |
+| 512x512    | 20    | 1           | long   | 22.2     | 22.5    | 0.0450              |
+
+Prompt length has no material effect on latency (the long-prompt rows are within run-to-run variance and even slightly faster with tighter P99 tails). This confirms the AR stage's cost is dominated by decoding the fixed 4,161-token visual grid, not by prompt encoding: the KV prefill of a few hundred extra text tokens is negligible against 4,161 decode steps.
+
+## 5.6 Peak GPU Memory (per-stage process, nvidia-smi polling at 1 Hz)
+
+Sampled during the §5.4/§5.5 sweeps (stage processes identified by PID; both stages share the single A800):
+
+| Configuration (worst case) | stage0_ar (MB) | stage1_dit (MB) |
+|----------------------------|----------------|-----------------|
+| 1024x1024, 50 steps, c4    | 39,290         | 11,104          |
+| all other configs          | 39,194         | 11,102–11,104   |
+
+Memory is essentially flat across the sweep: the AR stage grows by only ~96 MB at concurrency 4 and the DiT stage by 2 MB, because both stages pre-allocate from `gpu_memory_utilization` budgets (0.5 / 0.3) at startup. Combined steady-state usage ≈ 50.4 GiB of the 80 GiB device.
+
+## 5.7 DiT-Stage Component Attribution (pipeline profiler)
+
+With `enable_diffusion_pipeline_profiler: true` in the stage-1 deploy config (`MammothModa2DiTPipeline` targets: `gen_transformer.forward`, `gen_vae.decode`, `gen_image_condition_refiner.forward`), three warm 1024x1024 / 50-step requests (profiler adds a `synchronize` per call, so timings are indicative, not throughput-grade):
+
+| Component | Per request (s) | Share of DiT pipeline |
+|-----------|-----------------|-----------------------|
+| `gen_transformer.forward` (100 calls: 50 steps x 2 passes) | 14.01–14.07 | 98.6% |
+| `gen_vae.decode` (1 call) | 0.148 | 1.0% |
+| `gen_image_condition_refiner.forward` (1 call) | 0.005 | 0.03% |
+| whole `MammothModa2DiTPipeline.forward` | 14.18–14.27 | — |
+
+The DiT latency is ~99% transformer sampling; VAE decode costs only ~0.15 s per image at 1024x1024 and is not a bottleneck. Per-step transformer time is ~140 ms (two forward passes per step: ~2.8 s nominal x CFG), consistent with §5.3's 17.2 s image-TTFT figure once base64 encoding and stage handoff are included.
+
+## 5.8 Endpoint Fix Verification
+
+Before the serving fix, `/v1/images/generations` returned HTTP 503 ("No diffusion stage found in multi-stage pipeline") because the pipeline's image-output stage is a generation-LLM DiT rather than a `diffusion`-typed stage. After the fix all requests return HTTP 200 with a valid base64 image (verified across every sweep configuration above).
+
+---
+
+# 6. Reproducibility
+
+1. Start the server as in §3.1 and wait for `Application startup complete`.
+2. Warm up with one real request (any size) — do not use reduced-step warmups.
+3. Run the sweep commands from §4.1 for each row in §5.
+4. Server-side per-stage timings are printed in the engine stats table (`stage_gen_time_ms`, `output_unit_count`) at `--log-stats` level info.
+5. New rows use the same seed (`--seed 142`), `--num-prompts 8`, and `--warmup-requests 0` protocol as §4.1; when saving results with `--output-file`, name them `<res>_s<steps>_c<conc>[_long].json` to keep runs comparable (steps-25, long-prompt, peak-memory and profiler runs).
+6. To reproduce §5.6 peak-memory numbers, poll `nvidia-smi --query-gpu=memory.used --format=csv -l 1` per stage PID while a sweep is in flight (e.g. `nvidia-smi pmon` / `--query-compute-apps`).
+7. To reproduce §5.7, copy the bundled deploy config, set `enable_diffusion_pipeline_profiler: true` on the stage-1 entry, restart, and grep `[DiffusionPipelineProfiler]` lines from the server log; timings include a `torch.cuda.synchronize` per call.

docker/Dockerfile.ci

@@ -1,10 +1,7 @@
-# Target vLLM version: v0.30.0 (tag @ ced6857afa0e).
-# Use the matching release image and its coherent CUDA/PyTorch/FlashInfer
-# runtime. Do not reinstall a commit wheel or repair dependencies on top of
-# an older base image at a release boundary. For a future unreleased target,
-# restore an explicit wheel pin and validate its complete runtime stack.
+# Use the official release image with its matching CUDA/PyTorch stack.
+ARG VLLM_VERSION=0.31.0
 ARG VLLM_BASE_IMAGE=vllm/vllm-openai
-ARG VLLM_BASE_TAG=v0.30.0
+ARG VLLM_BASE_TAG=v${VLLM_VERSION}
 
 # We use registry-based layer caching for the dependencies,
 # where the tag is determined based on the hash of the contents
@@ -12,6 +9,7 @@ ARG VLLM_BASE_TAG=v0.30.0
 # everything except the actual project dir, since that should not
 # change much across branches.
 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG} AS deps
+ARG VLLM_VERSION
 # Keep APP_DIR outside /workspace so k8s Buildkite workspace mounts cannot
 # shadow the editable install target (#5364). Docker WORKDIR creates /opt if
 # needed.
@@ -36,12 +34,17 @@ RUN apt-get update && \
 # for CI
 COPY pyproject.toml setup.py ./
 COPY requirements/ requirements/
-ENV SETUPTOOLS_SCM_PRETEND_VERSION=0.0.0
+# Scope the override to Omni; dependency source builds need their own versions.
+ENV VLLM_OMNI_VERSION_OVERRIDE=0.0.0
 # Editable install: deps + project metadata are cached in this layer, while
 # the package itself is linked to APP_DIR (populated later by COPY --link).
 # Do NOT use a non-editable install here — that copies an empty stub into
 # site-packages and breaks ``import vllm_omni.entrypoints``.
 RUN uv pip install --system -e ".[dev]"
+# Fail the image build if dependency resolution replaced the selected vLLM.
+# The CPU builder has no NVIDIA driver; CUDA imports are checked at the start
+# of each GPU test job by .buildkite/cuda/scripts/check_vllm_runtime.py.
+RUN python3 -c "import vllm; assert vllm.__version__ == '${VLLM_VERSION}', vllm.__version__"
 
 # The final stage just copies the working dir on top of the
 # dependencies image. This uses --link so that we can skip

docker/Dockerfile.ci.dockerignore

@@ -1,6 +1,6 @@
 # NOTE: this is largely taken from vLLM
 
-# Exclude .git for CI builds — version is set via SETUPTOOLS_SCM_PRETEND_VERSION
+# Exclude .git for CI builds — version is set via VLLM_OMNI_VERSION_OVERRIDE
 .git
 .venv
 /build

docker/Dockerfile.cuda

@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=vllm/vllm-openai:v0.30.0
+ARG BASE_IMAGE=vllm/vllm-openai:v0.31.0
 FROM ${BASE_IMAGE}
 
 ARG COMMON_WORKDIR=/app

docker/Dockerfile.rocm

@@ -1,7 +1,7 @@
 # Keep the ROCm runtime aligned with the vLLM release targeted by
 # docker/Dockerfile.ci. A stale base image can build successfully while failing
 # at runtime when vLLM-Omni imports APIs added by the newer vLLM release.
-ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.30.0
+ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.31.0
 
 # Keep the dependency-heavy portion independent from the source tree. The
 # registry cache is keyed on these inputs, while the final source copy can be
@@ -12,11 +12,11 @@ FROM ${BASE_IMAGE} AS deps
 # BASE_IMAGE pins the default prebuilt runtime when USE_NIGHTLY_BUILD=0.
 # VLLM_VERSION_OR_COMMIT_HASH is consumed only by the opt-in source rebuild when
 # USE_NIGHTLY_BUILD=1. Keep both aligned with docker/Dockerfile.ci's
-# VLLM_BASE_TAG. This is enforced by tests/buildkite/test_rocm_dockerfile.py. An
+# release version. This is enforced by tests/buildkite/test_rocm_dockerfile.py. An
 # intentional custom source pin must update that regression in the same change,
 # and should return to the stable release as soon as one is available.
 ARG USE_NIGHTLY_BUILD=0
-ARG VLLM_VERSION_OR_COMMIT_HASH=v0.30.0
+ARG VLLM_VERSION_OR_COMMIT_HASH=v0.31.0
 ARG ARG_PYTORCH_ROCM_ARCH
 ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}}
 
@@ -75,7 +75,8 @@ WORKDIR ${COMMON_WORKDIR}/vllm-omni
 RUN mkdir -p vllm_omni && touch vllm_omni/__init__.py
 COPY pyproject.toml setup.py ./
 COPY requirements/ requirements/
-ENV SETUPTOOLS_SCM_PRETEND_VERSION=0.0.0
+# Scope the override to Omni; dependency source builds need their own versions.
+ENV VLLM_OMNI_VERSION_OVERRIDE=0.0.0
 
 RUN uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" \
         --no-cache-dir -e ".[dev]" --no-build-isolation && \

docs/configuration/README.md

@@ -2,7 +2,7 @@
 
 This section lists the most common options for running vLLM-Omni.
 
-For options within a vLLM Engine, please refer to the [vLLM 0.30 configuration guide](https://docs.vllm.ai/en/v0.30.0/configuration/index.html).
+For options within a vLLM Engine, please refer to the [vLLM 0.31 configuration guide](https://docs.vllm.ai/en/v0.31.0/configuration/index.html).
 
 Each model defines fixed topology in a registered `PipelineConfig` and runtime
 overrides in a deploy YAML.

docs/configuration/environment_variables.md

@@ -130,7 +130,7 @@ settings.
 
 | Name | Type and default | Applies to and read time | Precedence and invalid values | Lifecycle |
 | --- | --- | --- | --- | --- |
-| `VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS` | Non-negative integer; unset or empty means each consumer's own default (Boogu-Image: `2048`) | Token gate of the shared `fused_qk_norm_rope` op: a rotary table is packed for the fused kernel only when it spans at least this many positions (`B*S`); read at forward time | Environment-only. `0` = always fuse; a very large value disables the fused path. Any other value raises `ValueError` on the first forward. The crossover is host-dependent; re-benchmark before overriding. | Experimental performance control |
+| `VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS` | Non-negative integer; unset or empty means each consumer's own default (Boogu-Image: `2048`; MammothModa2: `0`) | Token gate of the shared `fused_qk_norm_rope` op: a rotary table is packed for the fused kernel only when it spans at least this many positions (`B*S`); read at forward time | Environment-only. `0` = always fuse; a very large value disables the fused path. Any other value raises `ValueError` on the first forward. The crossover is host-dependent; re-benchmark before overriding. | Experimental performance control |
 | `VLLM_OMNI_SKIP_NVFP4_NAN_CLAMP` | Boolean truthy spellings: `1`, `true`, `yes`, `on`; default false | ModelOpt NVFP4 compatibility patch; read when `vllm_omni.patch` imports | Environment-only escape hatch. Any other value means false. Set only to diagnose the upstream NaN-scale issue. | Diagnostic and temporary |
 | `VLLM_OMNI_USE_QUACK_FP8` | Boolean truthy spellings: `1`, `true`, `yes`, `on`; unset means hardware auto-detection | FP8 scaled matrix multiplication; evaluated when quack capability is selected | A set value overrides auto-detection. Any non-truthy value forces quack off. If quack cannot load, vLLM-Omni warns and falls back to FlashInfer. | Experimental performance control |
 
@@ -186,7 +186,7 @@ their keys only.
 ## Inherited vLLM variables
 
 vLLM-Omni also reads variables through its aligned vLLM dependency. Refer to
-the [vLLM 0.30 environment-variable reference](https://docs.vllm.ai/en/v0.30.0/configuration/env_vars.html)
+the [vLLM 0.31 environment-variable reference](https://docs.vllm.ai/en/v0.31.0/configuration/env_vars.html)
 for their definitions. This includes vLLM launch, cache, logging, plugin, ROCm,
 XPU, ModelScope, and FlashInfer workspace settings.
 

docs/configuration/stage_configs.md

@@ -35,7 +35,7 @@ Common `StagePipelineConfig` fields include:
 | `final_output` / `final_output_type` | Whether the stage produces a user-visible output and its modality. |
 | `owns_tokenizer` | Whether this stage owns the pipeline tokenizer. |
 | `model_arch`, `hf_config_name` | Stage-specific model architecture and nested HF config selector. |
-| `engine_output_type` | Runtime output representation such as `text`, `latent`, or `audio`. |
+| `engine_output_type` | Canonical runtime output type: `text`, `image`, `audio`, `latent`, or `token_ids`. |
 | `custom_process_input_func` | Processor applied to this stage's incoming payload. |
 | `custom_process_next_stage_input_func` | Processor used for full-payload handoff to the next stage. |
 | `async_chunk_process_next_stage_input_func` | Processor used for async chunk handoff. |
@@ -45,6 +45,50 @@ To add or change topology, define and register a new pipeline variant. Use
 deploy YAML only for runtime placement, resource sizing, connectors, and other
 deployment overrides.
 
+### Engine output types
+
+`engine_output_type` is parsed when the LLM stage's output processor is built
+during service initialization. Invalid values fail startup before requests
+are accepted; engine replicas may already have been initialized at that point
+and are cleaned up on failure.
+
+Use the exact lowercase names listed above. Combinations such as `text+audio`
+or `text,audio` remain supported, with every component canonical. Omitted or
+`None` values retain the default behavior (`TEXT` in the output processor).
+Empty strings, whitespace, case variants, and aliases are rejected.
+
+For existing custom pipelines or direct callers, replace old aliases explicitly:
+
+| Old value | Canonical value |
+| ------- | ------- |
+| `speech`, `wav`, `waveform` | `audio` |
+| `images`, `pixel_values`, `pixels` | `image` |
+| `latents` | `latent` |
+| `tokens` | `text` for text output, or `token_ids` for discrete IDs |
+
+`token_ids` is a separate type, not a `text` alias. It describes discrete
+vocabulary or codebook indices, distinct from continuous latent representations.
+Existing models using `latent` retain their historical payload contracts.
+GLM-Image's AR stage keeps `engine_output_type="token_ids"`; generated IDs still
+travel through `cumulative_token_ids`, and source-image IDs through
+`ids.prior_image`.
+
+Bare primary tensors and generic producer keys (`model_outputs` or `hidden`)
+use the declared type as their default output key, including `token_ids`.
+The normalizer does not inspect or convert their dtype. Distinct auxiliary
+outputs must use explicit field names: `model_outputs` and `hidden` both map
+to the same default key and must not represent two independent outputs.
+GLM-Image disables auxiliary hidden output for its default GPU execution.
+Some runner paths still attach it; that field is then named `token_ids` even
+though it contains hidden states. The GLM-Image bridge ignores this auxiliary
+field and reads generated IDs from `cumulative_token_ids` instead.
+
+Token-ID tensors use the existing dimension-0 accumulation default, subject to
+registered per-key overrides, and remain available after DELTA emissions.
+This is separate from the existing accumulation of `cumulative_token_ids`.
+`final_output_type` and Diffusion's `output_type` are separate settings and are
+unchanged by these rules.
+
 ## Deploy configuration schema
 
 The new deploy schema lives under `vllm_omni/deploy/` and is paired with a frozen `PipelineConfig` registered by the model's `pipeline.py`. Each deploy YAML has these top-level fields:
@@ -78,7 +122,44 @@ explicit CLI or stage override, or the downstream vLLM engine default.
 Note: for the diffusion path, an omitted `distributed_executor_backend` selects
 `uni` on a single GPU (in-process worker, no MessageQueue / `/dev/shm` output
 segments) and `mp` when `num_gpus > 1`. Set `mp` explicitly to keep a worker
-subprocess on one GPU. `ray` / `external_launcher` are not fully supported yet.
+subprocess on one GPU. Set `ray` for Ray-managed single- or multi-node
+diffusion workers. `external_launcher` is not supported yet.
+
+The diffusion `ray` backend runs one worker actor per GPU and supports workers
+spanning Ray nodes. Enable it in the deploy YAML:
+
+```yaml
+distributed_executor_backend: ray
+```
+
+The executor reuses Ray when already initialized; otherwise it calls
+`ray.init()`, which handles cluster discovery or local startup and respects
+`RAY_ADDRESS`. There is no diffusion-specific address setting. Ray must be
+installed on every node.
+
+Ray diffusion workers receive the driver's `VLLM_*`, `OMNI_*`, `DIFFUSION_*`,
+`NCCL_*`, `TORCH_NCCL_*`, `UCX_*`, and Hugging Face environment settings,
+plus `PYTHONPATH`, `CUDA_LAUNCH_BLOCKING`, and `OMP_NUM_THREADS`.
+Explicit stage `runtime.env` entries are also forwarded and take precedence.
+Device visibility, rank/rendezvous settings, worker host addresses, and `RAY_*`
+variables are excluded so each worker retains its own assignment and identity.
+
+For multiple nodes, start the cluster before launching Omni. For example,
+with head IP `10.0.0.1`:
+
+```bash
+# On the head node:
+ray start --head --node-ip-address=10.0.0.1 --port=6379
+
+# On each additional worker node:
+ray start --address=10.0.0.1:6379
+
+# On the host launching Omni (MODEL and my_deploy.yaml are placeholders):
+RAY_ADDRESS=10.0.0.1:6379 vllm serve MODEL --omni --deploy-config my_deploy.yaml
+```
+
+Omni creates actors on the resources of that cluster; it does not start Ray
+on additional machines.
 
 ### Stage-level runner selection
 

docs/design/feature/attention_execution_contract_poc.md

@@ -131,3 +131,24 @@ and real-kernel schema/fake checks pass. The full suite must be rerun with the
 vLLM version required by upstream. NPU/ROCm device numerics and compilation remain
 unvalidated. Pre-commit passes with the CI hook skips. Upstream PyTorch/CUTLASS
 warnings remain.
+
+## TRTLLM dense execution
+
+`TRTLLM_ATTN` declares `SUPPORTED` / `CUSTOM_OP` for noncausal dense BF16 on
+B200/GB200 (SM100) and B300/GB300 (SM103), head dimension 128, and equal Q/K/V
+head counts, without parallel, paged-KV, piecewise, or HSDP boundaries. Pre-construction, SAGE, skip-softmax, packed,
+and other unverified paths remain `UNMIGRATED`. Resolution and dispatch share
+metadata validation; workspace mutation is explicit in the custom-op schema.
+
+CPU contract tests cover fullgraph replay and schema/fake consistency using a
+substituted dispatcher. Real-kernel validation requires Blackwell with FlashInfer:
+
+```bash
+python -m pytest tests/diffusion/attention/test_trtllm_attn.py \
+  -k dense_contract_fullgraph_matches_sdpa -q -rs
+```
+
+These hardware tests compare eager/compiled output with FP32 SDPA and check
+schema/fake agreement. The TRTLLM and contract suite passed on Blackwell before the final architecture
+restriction was added: 61 passed, no skips (19 dependency deprecation warnings).
+The SM100/SM103 restriction is additionally covered by CPU contract tests.

docs/design/fullduplex.md

@@ -220,6 +220,7 @@ vllm_omni/
 │       │                            DUPLEX_REALTIME_CAPABILITIES, duplex_response_format
 │       ├── events.py                re-export shim (DuplexEvent = RealtimeEvent) + the runner's
 │       │                            epoch-filter sets DOMAIN_TERMINAL_EVENTS / MODEL_OUTPUT_EVENTS
+│       ├── delivery.py              DuplexOutputBuffer (bounded output shared with the session handle)
 │       ├── realtime_events.py       RealtimeProjectionState: internal event -> typed events
 │       ├── messages.py              queue envelopes (Open/Close/Resume/Touch/Command/Result/Event), DuplexSessionError
 │       ├── config.py                DuplexSessionConfig, DuplexCapabilities, ResponseCreateOptions
@@ -284,7 +285,8 @@ command  handle.submit(DuplexCommand)
          -> engine.submit_command_async -> DuplexSessionCommandMessage on the request queue [one-way]
          -> DuplexSessionManager.dispatch: unknown_session / input_backpressure checks, then runner mailbox
 output   DuplexOrchestrator._intercept_stage_output -> runner.on_stage_output -> mailbox -> typed events
-         -> output_sink (DuplexSessionEventMessage) -> DuplexOmni._route_engine_message -> handle.events()
+         -> session's bounded output buffer -> handle.events(); closure and errors without a live session
+            use output_sink (DuplexSessionEventMessage) -> DuplexOmni._route_engine_message
 detach   DuplexOmni.detach_session(expected_lease_generation) -> touch(DETACH): engine-owned disconnect
          grace, refused for a lease newer than the one the caller held; expiry -> SessionExpired
 resume   DuplexOmni.resume_session(expected_lease_generation) -> lease CAS keyed by the control id (a replay
@@ -303,7 +305,9 @@ reap     DuplexSessionManager.reaper_loop: idle TTL / disconnect grace expiry, c
 
 ### Concurrency and ordering model of the runner
 
-Everything below runs on the orchestrator asyncio loop; there is no lock.
+The runner's session state is owned by the orchestrator asyncio loop and
+needs no lock. The output buffer shared with the caller thread uses a lock
+for queue updates and the final validity check before delivery.
 
 - **Single writer per session.** `DuplexEngineSession` is mutated only by its
   runner. Inputs reach the runner through one mailbox in this order: client
@@ -332,13 +336,15 @@ Everything below runs on the orchestrator asyncio loop; there is no lock.
   on one loop in program order, nothing can be emitted for an old epoch after
   its terminal event; `runner.emit()` still drops a terminal carrying a stale
   epoch and stamps `epoch` on every event for clients that filter
-  defensively. This is a statement about order, not about latency: the engine
-  output queue and each `DuplexSessionHandle` outbox are unbounded and FIFO,
-  so a cancellation is delivered behind whatever audio was already emitted for
-  the response it cancels. A client that must stop quickly cancels its own
-  playback on `response.done` / `audio.cancelled` rather than waiting for the
-  stream to drain. Bounding those buffers and letting an accepted invalidation
-  skip undelivered media is left to the follow-up RFC.
+  defensively. In this experimental follow-up, public output goes directly to
+  one bounded buffer shared by the engine and `DuplexSessionHandle`. Accepted
+  cancellation removes queued `AudioDelta` events for that response through
+  the cancelled epoch. Text increments, audio/transcript completion markers,
+  other responses and normal completion retain FIFO order. The consumer also
+  checks an event already taken from the buffer: WebSocket delivery repeats
+  that check under the connection lock, before sequencing and journaling.
+  Already sequenced replay entries are unchanged; already delivered audio
+  still requires the client to stop its own playback.
 - **Backpressure before the mailbox.** `DuplexSessionManager.dispatch`
   checks `max_pending_input_bytes_per_session` and reserves a pending turn for
   `Commit` before the put; a rejected command is answered with
@@ -349,6 +355,46 @@ Everything below runs on the orchestrator asyncio loop; there is no lock.
   `ResponseDone.stage_metrics` and `metadata.vllm_omni.stage_metrics` on the
   wire keep their meaning.
 
+### Experimental output limits
+
+`max_pending_output_bytes_per_session` defaults to 2 MiB and
+`max_pending_output_events_per_session` to 512. The byte limit counts compact
+Realtime JSON, including base64 audio, not decoded PCM or Python object
+overhead. Audio is already base64-encoded by the producer; counting its length
+avoids serializing or decoding that large field again. Neither default has
+been tuned from a performance study.
+
+Errors and response endings have an additional 64 KiB / eight-event reserve
+and retain FIFO order. One final session-closure notification has an independent
+slot, so an exhausted reserve cannot prevent closure or break delivery to
+other sessions. If that reserve also fills, an error or response-ending
+notification may be omitted; the final session-closure notification still
+has its independent slot. Oversized close details (over 4 KiB with generated identity)
+are replaced by a compact reason, preserving terminal type and identity.
+Late writes after closure are ignored; a locally ended iterator is not reopened
+by a late terminal, and a failed open closes the abandoned handle. Events other
+than errors are not forwarded once their session's buffer has been removed.
+One event held by the consumer is outside the queued budget.
+
+The overflow policy fails the active response without committing its partial
+text to history and closes only the affected session. If overflow occurs while
+a completion batch is being delivered, its not-yet-queued `ResponseDone` is
+converted to failure with its identity and output preserved; any history entry
+for that response is removed. An ending already accepted by the buffer is not
+replaced or followed by a second ending. After overflow, newly
+emitted ordinary events are suppressed, including text/audio completion,
+content-part and output-item completion markers. Previously queued non-audio
+events retain FIFO order; error and failed response-ending notifications may
+use the reserve before the independent closure notification. This differs from
+accepted cancellation, which removes only matching audio. See the
+[public error and closure fields](../serving/realtime_duplex_api.md#output-limits-and-slow-consumers).
+Response-only recovery is not implemented. These limits do not bound the runner's raw stage
+output mailbox, replay journal, or client playback queues. The shared buffer
+requires the current same-process, separate-thread engine arrangement; it is
+not a cross-process delivery protocol. Direct handle iteration and WebSocket
+delivery share buffer validity, but end-to-end client behavior still needs
+independent real-model verification.
+
 ## Orchestrator seams
 
 `OrchestratorBase` keeps the generic machinery (request handler skeleton,

docs/design/module/ar_runtime.md

@@ -47,6 +47,29 @@ adapter, not injected through unrelated scheduler state.
 
 **Rule:** Workers and model runners MUST NOT implement cross-stage routing.
 
+## Token-only output payloads
+
+An AR stage whose downstream decoder consumes sampled tokens rather than
+hidden states can set `omni_pooler_payload_include_hidden = False`. Hidden
+states remain available for logits and sampling; prefill conditioning and
+other required multimodal outputs still cross the stage boundary.
+
+A chunk processor that needs IDs on steps without a tensor payload declares
+`requires_token_updates = True`. The scheduler then enqueues nonempty sampled
+IDs while preserving final-stage exclusions and the processor's EOF handling.
+Processors without this declaration retain their tensor-payload trigger.
+
+When hidden and multimodal payloads are absent, the shared runner skips empty
+per-request payload construction. Prefix-cache merges and pending model
+postprocess retain the normal path. This policy is independent of whether
+output materialization runs inline or in the background.
+
+Inline construction consumes query offsets, scheduled-token counts and
+scheduler metadata before the next execution can reuse them. Background
+construction snapshots those inputs. Both paths reuse the request IDs already
+copied by upstream bookkeeping; connector collection stays on the main thread
+because it can run a tensor-parallel collective.
+
 ## Safe-change guide
 
 Test request lifecycle, abort, cache state, and every affected worker execution

docs/design/module/diffusion/diffusion_runtime.md

@@ -36,6 +36,7 @@ validation_paths:
   - tests/diffusion/test_diffusion_worker.py
   - tests/diffusion/test_multiproc_engine_concurrency.py
   - tests/diffusion/test_uniproc_executor.py
+  - tests/diffusion/test_ray_executor.py
   - tests/diffusion/test_result_pump.py
   - tests/diffusion/test_async_output_worker.py
   - tests/diffusion/test_diffusion_ipc.py
@@ -96,8 +97,10 @@ The executor then chooses how the worker runs:
   RPC.
 - **`mp`** (default when `num_gpus > 1`, or when set explicitly):
   `MultiprocDiffusionExecutor` starts one `WorkerProc` per device.
+- **`ray`** (explicit): `RayDiffusionExecutor` starts one GPU actor per worker
+  across one or more Ray nodes.
 
-Either way, each worker owns its device, distributed state,
+In all backends, each worker owns its device, distributed state,
 `DiffusionWorker`, `DiffusionModelRunner`, and pipeline instance.
 
 ```mermaid
@@ -108,7 +111,7 @@ flowchart TB
         stageProc["StageDiffusionProc<br/>stage transport and lifecycle"]
         engine["DiffusionEngine<br/>admission and control loop"]
         scheduler["RequestScheduler or StepScheduler<br/>request state and policy"]
-        executor["DiffusionExecutor<br/>uni or mp backend"]
+        executor["DiffusionExecutor<br/>uni, mp, or ray backend"]
         streams["Per-request output streams"]
 
         stageProc --> engine
@@ -132,6 +135,8 @@ flowchart TB
     client -.->|"inline: direct call"| engine
     executor -->|"mp: broadcast / result queues"| workerProc
     executor -.->|"uni: in-process call"| worker
+    executor -->|"ray: actor RPC"| rayActor["RayDiffusionWorkerWrapper"]
+    rayActor --> worker
     streams -->|"process-backed"| stageProc
     stageProc -->|"ZMQ results"| client
     streams -.->|"inline"| client
@@ -154,6 +159,7 @@ flowchart TB
 | `DiffusionExecutor` | Execution backend contract, worker RPC, health, shutdown | Admission and request-state transitions |
 | `UniProcDiffusionExecutor` | Single in-process worker for `num_gpus == 1`; no IPC or async output pump | Multi-GPU execution |
 | `MultiprocDiffusionExecutor` | Worker processes, shared-memory message queues, result dispatch, worker monitoring | In-process single-GPU path |
+| `RayDiffusionExecutor` | GPU actors, placement groups, remote RPC, actor monitoring | Local worker subprocesses and shared-memory IPC |
 | `WorkerProc` | One `mp` worker process's IPC loop and reply rules | Scheduling policy; unused by `uni` |
 | `DiffusionWorker` | Device/distributed setup, LoRA activation, profiling, sleep/wake, runner delegation | Request admission |
 | `DiffusionModelRunner` | Pipeline loading, request-local model state, cache/compile setup, request or step execution | Queueing and cross-stage routing |
@@ -219,12 +225,12 @@ uses fused `execute_model_batch()` and requires the pipeline to declare
 request-batch support.
 
 **Distributed layerwise offload with AllGather (DLO DP concurrency)** is a
-separate multiproc dispatch path. It activates when
+separate dispatch path supported by `mp` and `ray`. It activates when
 `data_parallel_size > 1`, `enable_distributed_layerwise_offload` is set, and
 `dlo_use_allgather` is true. The engine then sets `dp_concurrent = True` and
 raises `scheduler.max_num_running_reqs` to `dp_size`, overriding the
 `max_num_seqs` value from `initialize()`. The scheduler still emits a
-multi-request `DiffusionSchedulerOutput`; the multiproc executor routes that
+multi-request `DiffusionSchedulerOutput`; these executors route that
 wave through `execute_request()` instead of fused `execute_model_batch()`.
 
 Optional admission coalescing can fill a `dp_size` wave: when
@@ -233,7 +239,7 @@ the first schedule of a wave (under `dp_concurrent`, the stable window is
 `min(0.3s, wait/2)`). With the default `request_batch_max_wait_ms == 0`,
 there is no wait.
 
-Before dispatch, `MultiprocDiffusionExecutor.execute_request()` rejects waves
+Before dispatch, both backends reject waves
 whose requests differ in sampling-parameter compatibility or `extra_args`
 (AllGather requires every DP rank to follow the same forward schedule). Each
 worker rank then picks one envelope from the wave:
@@ -297,8 +303,9 @@ head-of-line blocking.
 | --- | --- | --- |
 | `uni` | Default for `num_gpus == 1`, or set explicitly | One in-process worker |
 | `mp` | Default for `num_gpus > 1`, or set explicitly | One worker process per device |
+| `ray` | Set explicitly | One GPU actor per worker across Ray nodes |
 
-Ray and external-launcher diffusion backends are not implemented. A custom
+The external-launcher diffusion backend is not implemented. A custom
 `DiffusionExecutor` subclass or import path is also accepted.
 
 ### Uniproc backend
@@ -337,11 +344,27 @@ result path and does not start those pumps.
 See [Async diffusion output](../../feature/async_diffusion_output.md) for the
 multiproc request-mode timeline.
 
+### Ray backend
+
+`RayDiffusionExecutor` reuses initialized Ray or calls `ray.init()` (respecting
+`RAY_ADDRESS`). It reuses the current placement group or reserves one GPU
+bundle per worker with `PACK` placement. Ray owns GPU placement; PyTorch
+owns distributed collectives through a shared TCP rendezvous. Global ranks
+are grouped by host, while each actor uses local device 0.
+
+Request batches, step execution, and DLO DP waves use actor RPCs with
+rank-aware replies. Returned tensors move to CPU before Ray serialization;
+this backend does not use the multiproc async output pumps. Driver inference
+settings and explicit stage `runtime.env` are forwarded, excluding worker
+identity and device assignment. See [Stage configuration](../../../configuration/stage_configs.md)
+for cluster setup and environment details.
+
 ## Worker and runner boundary
 
 On the `mp` path, `WorkerProc` receives messages and calls methods through
 `WorkerWrapperBase`. On the `uni` path, the executor owns that wrapper
-directly. In both cases, `DiffusionWorker` handles the parts tied to a device:
+directly. Ray actors also construct `WorkerWrapperBase`. In all cases,
+`DiffusionWorker` handles the parts tied to a device:
 distributed initialization, model-runner construction, LoRA activation,
 profiling, and memory sleep/wake. It delegates actual model work to
 `DiffusionModelRunner`.
@@ -399,6 +422,10 @@ streams with an error, closes scheduler state, and shuts down the executor.
   async-output futures as failed, then clears them. Completed async outputs
   cached for later `wait_output_ready()` calls are dropped with the executor
   object.
+- **`ray`**: the executor gives actors up to five seconds to shut down, then
+  kills them and removes placement groups it owns. Actor death is monitored
+  even while idle; actor errors and RPC timeouts fail the group and notify the
+  stage. The stage client allows subprocess cleanup before signal fallback.
 - **`uni`**: the executor shuts down the in-process worker, drops the final
   model reference, and empties the accelerator cache so a later engine can reuse
   the device.
@@ -407,7 +434,7 @@ streams with an error, closes scheduler state, and shuts down the executor.
 
     Do not continue using a worker group after a fatal collective timeout or
     unexpected worker exit. Distributed state may be incomplete. The
-    multiprocess executor deliberately fails closed so the stage can restart.
+    `mp` and `ray` executors deliberately fail closed so the stage can restart.
     The uniproc executor fails closed when the accelerator context itself is
     poisoned.
 
@@ -417,10 +444,11 @@ streams with an error, closes scheduler state, and shuts down the executor.
   `diffusion_model_runner_cls` configuration still wins.
 - Tests and custom engine integrations may inject a `BaseScheduler` subclass.
   `SchedulerInterface` remains only as a deprecated compatibility name.
-- `distributed_executor_backend` may be `"uni"`, `"mp"`, a custom `DiffusionExecutor`
+- `distributed_executor_backend` may be `"uni"`, `"mp"`, `"ray"`, a custom `DiffusionExecutor`
   subclass, or an import path. A backend must preserve scheduler output,
   request identity, health, and cleanup contracts. Use `"mp"` when you need
-  process isolation or multi-GPU; `"uni"` is the single-GPU default.
+  local process isolation or multi-GPU, or `"ray"` for workers across nodes;
+  `"uni"` is the single-GPU default.
 - Worker extensions go through `WorkerWrapperBase`. They add worker methods;
   they do not become a second scheduler or request lifecycle.
 
@@ -466,7 +494,7 @@ Test the smallest affected slice, then cover its neighboring boundary:
 | --- | --- |
 | Admission or state transitions | `test_diffusion_scheduler.py`, including duplicate IDs, compatibility, abort, and missing output |
 | Engine loop or output delivery | `test_diffusion_engine.py`, `test_diffusion_engine_cleanup.py`, `test_diffusion_engine_rpc_routing.py` |
-| Executor or IPC | `test_multiproc_engine_concurrency.py`, `test_uniproc_executor.py`, `test_result_pump.py`, `test_diffusion_ipc.py`, `test_async_output_worker.py` |
+| Executor or IPC | `test_multiproc_engine_concurrency.py`, `test_uniproc_executor.py`, `test_ray_executor.py`, `test_result_pump.py`, `test_diffusion_ipc.py`, `test_async_output_worker.py` |
 | Worker or runner | `test_diffusion_worker.py`, `test_diffusion_model_runner.py` |
 | Stage boundary | `test_stage_diffusion_proc.py`, `test_inline_stage_diffusion_client.py` |
 | Warmup or streaming | `test_diffusion_engine_dummy_run.py`, `test_diffusion_streaming_output.py` |

docs/getting_started/installation/README.md

@@ -1,7 +1,7 @@
 # Installation
 
 !!! important
-    vLLM-Omni is released against the matching upstream vLLM major/minor version. The vLLM-Omni 0.30.x release line uses vLLM 0.30.x; the source-install instructions pin vLLM 0.30.0.
+    vLLM-Omni is released against the matching upstream vLLM major/minor version. This source checkout targets vLLM 0.31.0. The published vLLM-Omni 0.30.0 wheel instructions use vLLM 0.30.0; choose the installation path that matches your Omni revision.
 
 vLLM-Omni supports the following hardware platforms:
 

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

@@ -1,11 +1,13 @@
 # --8<-- [start:requirements]
 
-- GPU: compute capability 7.0 or higher (e.g., V100, T4, RTX20xx, A100, L4, H100, etc.)
+- GPU: compute capability 7.5 or higher (e.g., T4, RTX20xx, A100, L4, H100, etc.)
 
 # --8<-- [end:requirements]
 # --8<-- [start:set-up-using-python]
 
-vLLM-Omni depends on the matching major/minor release of vLLM. The vLLM-Omni 0.30.x release line uses vLLM 0.30.x.
+vLLM-Omni depends on the matching major/minor release of vLLM. This source
+checkout targets vLLM 0.31.0. The pre-built Omni wheel instructions below install
+the published 0.30.0 release with its matching vLLM dependency.
 
 !!! note
     PyTorch installed via `conda` will statically link `NCCL` library, which can cause issues when vLLM tries to use `NCCL`. See <gh-issue:8420> for more details.
@@ -41,13 +43,17 @@ uv pip install 'vllm-omni[demo]==0.30.0'
 # --8<-- [start:build-wheel-from-source]
 
 #### Installation of vLLM
-If you do not need to modify source code of vLLM, you can directly install the stable 0.30.0 release version of the library
+If you do not need to modify source code of vLLM, install the final 0.31.0
+release wheel from its commit-pinned index:
 
 ```bash
-uv pip install vllm==0.30.0 --torch-backend=auto
+uv pip install vllm==0.31.0 --torch-backend=auto \
+  --extra-index-url https://wheels.vllm.ai/db9527a46873454610df6dbedf79a36d6bf1a7f6
 ```
 
-The 0.30.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
+The default 0.31.0 wheel contains CUDA 12.9 binaries. If you need a different
+CUDA variant or want to reuse an existing PyTorch installation, build vLLM
+from source instead.
 
 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -68,16 +74,16 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.30.0
+git checkout v0.31.0
 ```
 Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file.
 ```bash
-#For CUDA 13.0 (the default for v0.30.0; the wheel filename has no `+cu130` suffix)
-export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.30.0/vllm-0.30.0-cp38-abi3-manylinux_2_28_x86_64.whl
+# Default CUDA 12.9 wheel for v0.31.0 on x86_64
+export VLLM_PRECOMPILED_WHEEL_LOCATION=https://wheels.vllm.ai/db9527a46873454610df6dbedf79a36d6bf1a7f6/vllm-0.31.0-cp38-abi3-manylinux_2_28_x86_64.whl
 ```
 Install vllm with command below (If you have no existing PyTorch).
 ```bash
-uv pip install --editable .
+VLLM_USE_PRECOMPILED=1 uv pip install --editable . --torch-backend=auto
 ```
 Install vllm with command below (If you already have PyTorch).
 ```bash
@@ -128,12 +134,13 @@ docker run --runtime nvidia --gpus 2 \
 DOCKER_BUILDKIT=1 docker build -f docker/Dockerfile.cuda -t vllm-omni-cuda .
 ```
 
-If you want to specify the base vLLM version:
+The source build defaults to `vllm/vllm-openai:v0.31.0`. To select that
+upstream release image explicitly:
 
 ```bash
 DOCKER_BUILDKIT=1 docker build \
   -f docker/Dockerfile.cuda \
-  --build-arg BASE_IMAGE=vllm/vllm-openai:v0.30.0 \
+  --build-arg BASE_IMAGE=vllm/vllm-openai:v0.31.0 \
   -t vllm-omni-cuda .
 ```
 

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

@@ -10,8 +10,9 @@
 
 For ROCm, vLLM-Omni currently recommends the setup steps through Docker Images.
 
-vLLM-Omni depends on the matching major/minor release of vLLM. The vLLM-Omni
-0.30.x release line uses vLLM 0.30.x.
+vLLM-Omni depends on the matching major/minor release of vLLM. This source
+checkout targets vLLM 0.31.0. The pre-built Omni wheel instructions below install
+the published 0.30.0 release with its matching vLLM dependency.
 
 The Dockerfile's `BASE_IMAGE` pin applies only to Docker builds. The
 `vllm-omni` package does not install vLLM as a dependency, so non-Docker source
@@ -49,13 +50,14 @@ uv pip install onnxruntime-rocm
 
 #### Installation of vLLM
 
-If you do not need to modify source code of vLLM, you can directly install the stable 0.30.0 release version of the library
+If you do not need to modify source code of vLLM, install the final 0.31.0
+release wheel for ROCm 7.2.3:
 
 ```bash
-uv pip install vllm==0.30.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.30.0/rocm723
+uv pip install vllm==0.31.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.31.0/rocm723
 ```
 
-The pre-built 0.30.0 vLLM wheel targets ROCm 7.2.3. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
+The pre-built 0.31.0 vLLM wheel targets ROCm 7.2.3. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
 
 #### Installation of vLLM-Omni
 
@@ -75,7 +77,7 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.30.0
+git checkout v0.31.0
 python3 -m pip install -r requirements/rocm.txt
 python3 setup.py develop
 ```
@@ -87,7 +89,7 @@ python3 setup.py develop
 #### Build docker image
 
 The source build defaults to the published upstream base image
-`vllm/vllm-openai-rocm:v0.30.0`, aligned with the vLLM release used by CI.
+`vllm/vllm-openai-rocm:v0.31.0`, aligned with this source checkout.
 Older or custom bases must provide the vLLM APIs checked by the Dockerfile's
 image-build canary. This upstream base is distinct from the prebuilt
 `vllm/vllm-omni-rocm` images discussed below; their published-tag availability
@@ -102,7 +104,7 @@ To select the upstream ROCm base explicitly, pass `BASE_IMAGE`:
 ```bash
 DOCKER_BUILDKIT=1 docker build \
   -f docker/Dockerfile.rocm \
-  --build-arg BASE_IMAGE=vllm/vllm-openai-rocm:v0.30.0 \
+  --build-arg BASE_IMAGE=vllm/vllm-openai-rocm:v0.31.0 \
   -t vllm-omni-rocm .
 ```
 

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

@@ -20,12 +20,16 @@ vLLM-Omni currently recommends using the Docker image setup steps below.
 #### Build docker image
 
 ```bash
-DOCKER_BUILDKIT=1 docker build -f docker/Dockerfile.xpu -t vllm-omni-xpu --shm-size=4g .
+export VLLM_VERSION=v0.31.0
+DOCKER_BUILDKIT=1 docker build -f docker/Dockerfile.xpu -t vllm-omni-xpu --shm-size=4g \
+  --build-arg VLLM_VERSION="$VLLM_VERSION" .
 ```
 
 This layers vLLM-Omni on top of the published `vllm/vllm-openai-xpu:<VLLM_VERSION>`
-base image (v0.30.0 by default), which Docker pulls automatically. To target a different vLLM release,
-pass `--build-arg VLLM_VERSION=<tag>`. If that tag has not been published yet, build
+base image, which Docker pulls automatically. The Dockerfile still defaults to
+v0.30.0; the command above explicitly selects v0.31.0 to match this source
+checkout. To target a different vLLM release, pass
+`--build-arg VLLM_VERSION=<tag>`. If that tag has not been published yet, build
 the base from upstream's own Dockerfile first and point the build at it:
 
 ```bash

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

@@ -71,10 +71,14 @@ We are keeping [issue #886](https://github.com/vllm-project/vllm-omni/issues/886
 
 # --8<-- [start:installation-main]
 
-You can also install vLLM-Omni from the latest main branch if you want to use the latest features or bug fixes. Use the aligned vLLM-Ascend v0.30.0 image described above, and check [issue #886](https://github.com/vllm-project/vllm-omni/issues/886) for the status of the latest vLLM-Omni commit on NPU.
+You can also install vLLM-Omni from source if you want to use the latest features
+or bug fixes. This source checkout targets vLLM 0.31.0 and needs a matching
+vLLM-Ascend runtime. The v0.30.0 images above are for the published Omni v0.30.0
+release. Check [issue #886](https://github.com/vllm-project/vllm-omni/issues/886)
+for aligned versions and NPU validation before selecting a development image.
 
 ```bash
-# Inside the vLLM-Ascend v0.30.0 container, install vLLM-Omni from main
+# Inside a container with matching vLLM and vLLM-Ascend, install Omni from source
 cd /vllm-workspace
 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni

docs/getting_started/quickstart.md

@@ -19,17 +19,20 @@ uv venv --python 3.12 --seed
 source .venv/bin/activate
 
 # On CUDA
-uv pip install vllm==0.30.0 --torch-backend=auto
+uv pip install vllm==0.31.0 --torch-backend=auto \
+  --extra-index-url https://wheels.vllm.ai/db9527a46873454610df6dbedf79a36d6bf1a7f6
 
 # On ROCm
-uv pip install vllm==0.30.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.30.0/rocm723
+uv pip install vllm==0.31.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.31.0/rocm723
 
 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni
 uv pip install -e .
 ```
 
-For additional installation methods — please see the [installation guide](installation/README.md).
+These commands install vLLM 0.31.0 for this source checkout. The CUDA index is
+pinned to the final release commit. For published vLLM-Omni wheels and additional
+installation methods, see the [installation guide](installation/README.md).
 
 !!! note
     It is important to install the same major & minor version of vLLM and vLLM Omni, otherwise things may not work as expected. If the versions are misaligned, you will see a warning when you import vLLM Omni.

docs/models/supported_models.md

@@ -83,7 +83,7 @@ th {
 | `OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` | ✅︎ | ✅︎ | | ✅︎ | — |
 | `StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` | ✅︎ | ✅︎ | | ✅︎ | — |
 | `MiniMaxMusic3ForConditionalGeneration` | MiniMax Music 3 (text-to-music) | `MiniMaxAI/MiniMax-Music3` | ✅︎ | | | | — |
-| `Yue2ForCausalLM` | YuE2-3B (text-to-music with ABC plan) | `m-a-p/YuE2-3B` | ✅︎ | | | | — |
+| `Yue2ForCausalLM` | YuE2-3B (text-to-music with ABC plan) | `m-a-p/YuE2-3B` | ✅︎ | | | | [Recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/m-a-p/YuE2-3B.md) / [Audio-to-score](https://github.com/vllm-project/vllm-omni/blob/main/recipes/m-a-p/SheetSage2-H200.md) |
 | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | | | ✅︎ | [Published](https://recipes.vllm.ai/Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice) |
 | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |

docs/serving/realtime_duplex_api.md

@@ -250,6 +250,40 @@ call; transport failures raise `DuplexConnectionError`; a session that ends
 while you are awaiting it raises `DuplexSessionClosedError`. All three derive
 from `DuplexClientError`.
 
+### Output limits and slow consumers
+
+Deployment settings under `duplex_session` bound each session's pending
+output: `max_pending_output_bytes_per_session` defaults to 2 MiB and
+`max_pending_output_events_per_session` to 512 events. Bytes count compact
+Realtime JSON, including base64 audio, not decoded audio. The limits apply
+to both WebSocket delivery and direct Python handles, not to replay journals
+or client playback queues.
+
+Exceeding a limit fails the active response without committing its partial
+text to history, then closes only that session. If overflow occurs while a
+completion is being queued, its not-yet-queued `response.done` becomes failed
+with the same identity and output, and its response history entry is removed.
+An ending already accepted by the buffer is not replaced or followed by a
+second ending. Open a new session to continue; the closed session cannot be
+resumed. The notifications use these fields:
+
+| Event | Field values |
+| --- | --- |
+| `error` | `error.type: "rate_limit_error"`, `error.code: "output_backpressure"` |
+| `response.done` (active response or not-yet-queued ending) | `response.status: "failed"`, `response.status_details: {"type": "failed", "reason": "output_backpressure"}` |
+| `session.closed` | `reason: "output_backpressure"` |
+
+Queued audio for the affected response is removed. Previously queued non-audio
+events retain their order, but newly emitted ordinary events, including
+audio/text, content-part and output-item completion markers, are suppressed.
+Clients must therefore handle closure without waiting for every completion
+marker. Errors and response endings have an extra 64 KiB / eight-event reserve;
+if that also fills, either notification may be omitted. Closure has its own
+independent slot. A closure notification exceeding 4 KiB is compacted to
+`reason: "close_details_exceed_output_limit"` with empty details, preserving
+the event type and identity. One event held by the consumer is outside the
+queued budget; audio already sent still requires the client to stop playback.
+
 ### Measure latency
 
 `EventCollector` accumulates events for assertions and metrics:
@@ -315,6 +349,8 @@ single-consumer async iterator that ends after `session.closed` /
 `session.expired`; rejected commands come back as `ErrorEvent`s on it.
 Every event renders the wire JSON with `to_realtime()`, so anything written
 against the WebSocket protocol works unchanged on the typed stream.
+If opening fails, the abandoned handle is closed. Late engine events do not
+reopen a locally closed handle.
 
 `vllm_omni.clients.inline_duplex.InlineDuplexClient` wraps a `DuplexOmni`
 behind the `DuplexClient` API, so the same application code runs in-process
@@ -573,7 +609,7 @@ formats (`pcm16`, `pcm_s16le`, `s16le`, `pcm_f32le`, `g711_ulaw`,
 | `session.created` | 2 | Session opened. OpenAI name; adds `attachment_generation`, `resume_token` and vLLM-Omni keys inside `session` (`id` is the server-allocated session id, `epoch`, `turn_id`, `playback`, `capabilities`, ...). |
 | `session.updated` | 2 | Echo of the effective session config after every `session.update` (same extended `session` object). |
 | `session.heartbeat_ack` | 3 | Reply to `session.heartbeat`. |
-| `session.closed` | 3 | Last event on the socket, emitted once the engine released the session; `reason` ∈ `client_close`, `disconnect`, `timeout`, `transport_error`, `shutdown`, ... |
+| `session.closed` | 3 | Last event on the socket, emitted once the engine released the session; `reason` ∈ `client_close`, `disconnect`, `timeout`, `transport_error`, `shutdown`, `output_backpressure`, `close_details_exceed_output_limit`, ... |
 | `session.resumed` | 3 | Resume accepted; carries the new `attachment_generation` and rotated `resume_token`; journaled events are replayed after it. |
 | `session.replaced` | 3 | Sent to the superseded socket when another socket resumes the session. |
 | `session.expired` | 3 | Engine lease reaped (`disconnect_grace_expired`, `idle_ttl_expired`); socket closes after it. |

docs/user_guide/diffusion_features.md

@@ -154,7 +154,7 @@ The following tables show which models support each feature:
 | **Krea 2**               |     ❌     |     ✅      |           ❌           |       ❌        |         ❌         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **LongCat-Image**        |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **LongCat-Image-Edit**   |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
-| **MammothModa2(T2I)**    |     ❌     |     ✅      |           ❌           |       ❌        |         ❌         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
+| **MammothModa2(T2I)**    |     ✅     |     ✅      |           ❌           |       ❌        |         ❌         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **Nextstep_1(T2I)**      |     ❓     |     ❓      |           ❌           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **OmniGen2**             |     ❌     |     ✅      |           ✅           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **Ovis-Image**           |     ❌     |     ✅      |           ❌           |       ✅        |         ❌         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |

docs/user_guide/examples/offline_inference/text_to_speech.md

@@ -130,7 +130,7 @@ python examples/offline_inference/text_to_speech/cosyvoice3/end2end.py \
 ### Notes
 
 - Stage 0 (`talker`) emits speech tokens; stage 1 (`code2wav`) runs flow matching + HiFiGAN to synthesize waveform.
-- Deploy config auto-loads from `vllm_omni/deploy/cosyvoice3.yaml` based on HF `model_type`. Pass `--deploy-config <path>` to override.
+- H200 automatically selects `cosyvoice3_packed_streaming_optimized_standard.yaml` when the CUDA MPS control tool is available; other devices and environments use `cosyvoice3.yaml`. See the [serving defaults and tradeoffs](../online_serving/text_to_speech.md#cosyvoice3). Pass `--deploy-config cosyvoice3.yaml` for the previous profile or `--deploy-config <path>` to override. The example uses the resolved profile's sampling parameters.
 - `async_chunk: true` is the default; pass `--no-async-chunk` to switch to the legacy synchronous path.
 
 ---

docs/user_guide/examples/online_serving/text_to_speech.md

@@ -155,16 +155,22 @@ vllm serve FunAudioLLM/Fun-CosyVoice3-0.5B-2512 --omni --port 8091 --trust-remot
 ./cosyvoice3/run_server.sh
 ```
 
-Streaming is on by default via `async_chunk: true` in `vllm_omni/deploy/cosyvoice3.yaml`. Pass `--no-async-chunk` (or `NO_ASYNC_CHUNK=1 ./cosyvoice3/run_server.sh`) for the legacy synchronous path.
+Streaming is on by default. On CUDA Hopper GPUs with at least 140 GiB of device memory (H200), serving automatically selects `cosyvoice3_packed_streaming_optimized_standard.yaml` when `nvidia-cuda-mps-control` is available on `PATH`: AR32 / codec32, standard sampling, packed Flow, cached ISTFT, HiFT decode graphs, and local CUDA MPS. No `--hf-overrides` is needed. Other devices and environments without the MPS control tool use `cosyvoice3.yaml` with RAS sampling and AR8 / codec8.
 
-Cross-request Stage-1 flow batching is opt-in. Enable it in the server environment when concurrent requests should share a flow-estimator call:
+The optimized profile is not a universal hardware configuration. Its packed kernels are Hopper-only; the talker's Model Runner V2 is unsupported on NPU/XPU, and CUDA MPS requires a local NVIDIA GPU and its control tool. Larger batches, 25-frame chunks, and graph allocations also require per-device validation on smaller or slower GPUs. Automatic selection retains the generic profile outside the validated device class; an explicit deploy configuration takes precedence.
+
+The H200 default targets streaming throughput at C32/C64. It changes sampling from RAS to standard (temperature 0.7, top-p 0.8, top-k 20, repetition penalty 1.21); output lengths and quality can differ. HiFT captures a recurring exact shape on its third use, which adds latency to that request, and its graph pool needs memory headroom beyond the engine reservations. This profile is not a claim of optimal first-audio latency at low concurrency. Pass `--deploy-config cosyvoice3.yaml` to retain the previous configuration, or `--deploy-config <path>` to select another profile.
+
+Pass `--no-async-chunk` (or `NO_ASYNC_CHUNK=1 ./cosyvoice3/run_server.sh`) for the legacy synchronous path; use `--deploy-config cosyvoice3.yaml` alongside it when selecting that path on H200.
+
+Cross-request Stage-1 flow batching is enabled in the H200 profile. For the generic configuration, enable it in the server environment when concurrent requests should share a flow-estimator call:
 
 ```bash
 export COSYVOICE3_BATCH_FLOW=1
 vllm serve FunAudioLLM/Fun-CosyVoice3-0.5B-2512 --omni --port 8091 --trust-remote-code
 ```
 
-Batching preserves output lengths and streaming cache alignment, but the different GEMM shapes can produce small waveform differences compared with processing each request separately. Leave `COSYVOICE3_BATCH_FLOW` unset (or set it to `0`) when request-independent numerical behavior is required. Set `COSYVOICE3_BATCH_FLOW_DEBUG=1` to log the observed group-size distribution and enable detailed Stage-1 profiler scopes; diagnostics are disabled by default to avoid per-step profiling overhead.
+Batching preserves output lengths and streaming cache alignment, but the different GEMM shapes can produce small waveform differences compared with processing each request separately. Use `--deploy-config cosyvoice3.yaml` and leave `COSYVOICE3_BATCH_FLOW` unset (or set it to `0`) when request-independent numerical behavior is required. Set `COSYVOICE3_BATCH_FLOW_DEBUG=1` to log the observed group-size distribution and enable detailed Stage-1 profiler scopes; diagnostics are disabled by default to avoid per-step profiling overhead.
 
 ### CLI client
 
@@ -194,7 +200,7 @@ The client supports `--api-base`, `--model`, `--text`, `--ref-audio`, `--ref-tex
 ### Notes
 
 - Stage 0 (`talker`) emits speech tokens; stage 1 (`code2wav`) runs flow matching + HiFiGAN to synthesize waveform.
-- Deploy config auto-loads from `vllm_omni/deploy/cosyvoice3.yaml` based on HF `model_type`. Pass `--deploy-config <path>` to override.
+- Deploy config auto-loads the device-specific CosyVoice3 default described above. Pass `--deploy-config <path>` to override.
 - For offline inference and the end-to-end script, see the [offline CosyVoice3 section](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_speech/README.md#cosyvoice3).
 
 ---

examples/offline_inference/text_to_image/text_to_image.py

@@ -14,7 +14,7 @@ from diffusers.utils import numpy_to_pil
 
 from vllm_omni.diffusion.data import logger
 from vllm_omni.diffusion.utils.image_output import extract_images_from_outputs
-from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
+from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args, ar_grid_max_tokens
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.entrypoints.openai.stage_params import clone_sampling_params
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -72,6 +72,24 @@ def build_text_to_image_prompt(prompt: str, negative_prompt: str | None) -> dict
     return result
 
 
+def build_parallel_knob_kwargs(args: argparse.Namespace) -> dict[str, Any]:
+    """Forward explicitly-set parallel knobs so they can override a deploy YAML.
+
+    The argparse defaults are ``None``; only explicitly supplied values enter
+    the returned dict, letting a deploy YAML's parallel_config stay in effect
+    when the flags are omitted (and letting explicit values equal to the old
+    argparse defaults, e.g. ``--ulysses-degree 1``, still override the YAML).
+    """
+    knob_names = (
+        "ulysses_degree",
+        "ring_degree",
+        "ulysses_mode",
+        "cfg_parallel_size",
+        "vae_patch_parallel_size",
+    )
+    return {name: getattr(args, name) for name in knob_names if getattr(args, name) is not None}
+
+
 def _normalize_images_for_save(images: list[Any]) -> list[Any]:
     """Convert NumPy diffusion outputs to PIL images before saving."""
     normalized = []
@@ -83,7 +101,8 @@ def _normalize_images_for_save(images: list[Any]) -> list[Any]:
     return normalized
 
 
-def parse_args() -> argparse.Namespace:
+def build_parser() -> argparse.ArgumentParser:
+    """Build the example's CLI parser (kept separate for testability)."""
     parser = argparse.ArgumentParser(description="Generate an image with supported diffusion models.")
     parser.add_argument(
         "--model",
@@ -166,28 +185,31 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--ulysses-degree",
         type=int,
-        default=1,
-        help="Number of GPUs used for ulysses sequence parallelism.",
+        default=None,
+        help="Number of GPUs used for ulysses sequence parallelism. "
+        "Explicitly passing 1 overrides a deploy YAML value.",
     )
     parser.add_argument(
         "--ulysses-mode",
         type=str,
-        default="strict",
+        default=None,
         choices=["strict", "advanced_uaa"],
-        help="Ulysses sequence-parallel mode: 'strict' (divisibility required) or 'advanced_uaa' (UAA).",
+        help="Ulysses sequence-parallel mode: 'strict' (divisibility required) or 'advanced_uaa' (UAA). "
+        "Explicitly passing 'strict' overrides a deploy YAML value.",
     )
     parser.add_argument(
         "--ring-degree",
         type=int,
-        default=1,
-        help="Number of GPUs used for ring sequence parallelism.",
+        default=None,
+        help="Number of GPUs used for ring sequence parallelism. Explicitly passing 1 overrides a deploy YAML value.",
     )
     parser.add_argument(
         "--cfg-parallel-size",
         type=int,
-        default=1,
+        default=None,
         choices=[1, 2],
-        help="Number of GPUs used for classifier free guidance parallel size.",
+        help="Number of GPUs used for classifier free guidance parallel size. "
+        "Explicitly passing 1 overrides a deploy YAML value.",
     )
     parser.add_argument(
         "--enforce-eager",
@@ -296,8 +318,9 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--vae-patch-parallel-size",
         type=int,
-        default=1,
-        help="Number of ranks used for VAE patch/tile parallelism (decode/encode).",
+        default=None,
+        help="Number of ranks used for VAE patch/tile parallelism (decode/encode). "
+        "Explicitly passing 1 overrides a deploy YAML value.",
     )
     # NextStep-1.1 specific arguments
     parser.add_argument(
@@ -413,7 +436,11 @@ def parse_args() -> argparse.Namespace:
         ),
     )
     current_omni_platform.pre_register_and_update(parser)
-    return parser.parse_args()
+    return parser
+
+
+def parse_args() -> argparse.Namespace:
+    return build_parser().parse_args()
 
 
 def _apply_ar_stage_inputs(
@@ -527,14 +554,8 @@ def main():
             "scm_steps_policy": "dynamic",  # SCM steps policy: "dynamic" or "static"
         }
     elif cache_backend == "tea_cache":
-        # TeaCache configuration
-        # All parameters marked with [tea_cache only] in DiffusionCacheConfig
-        cache_config = {
-            # TeaCache parameters [tea_cache only]
-            "rel_l1_thresh": 0.2,  # Threshold for accumulated relative L1 distance
-            # Note: coefficients will use model-specific defaults based on model_type
-            #       (e.g., QwenImagePipeline or FluxPipeline)
-        }
+        # Let TeaCache select the model-specific threshold and coefficients.
+        cache_config = {}
 
     profiler_enabled = args.profiler_config is not None
 
@@ -562,21 +583,11 @@ def main():
 
     omni_kwargs = {
         "model": args.model,
-        "enable_layerwise_offload": args.enable_layerwise_offload,
-        "vae_use_slicing": args.vae_use_slicing,
-        "vae_use_tiling": args.vae_use_tiling,
         "cache_backend": args.cache_backend,
         "cache_config": cache_config,
-        "enable_cache_dit_summary": args.enable_cache_dit_summary,
-        "ulysses_degree": args.ulysses_degree,
-        "ring_degree": args.ring_degree,
-        "ulysses_mode": args.ulysses_mode,
-        "cfg_parallel_size": args.cfg_parallel_size,
-        "vae_patch_parallel_size": args.vae_patch_parallel_size,
         "enable_expert_parallel": args.enable_expert_parallel,
         "enable_cpu_offload": args.enable_cpu_offload,
         "log_stats": args.log_stats,
-        "enable_diffusion_pipeline_profiler": args.enable_diffusion_pipeline_profiler,
         "profiler_config": args.profiler_config,
         "init_timeout": args.init_timeout,
         "stage_init_timeout": args.stage_init_timeout,
@@ -584,6 +595,8 @@ def main():
         **lora_args,
         **quant_kwargs,
     }
+    # Diffusion-only engine args: pass only when explicitly set, so pipelines
+    # without a DIFFUSION stage pass ownership validation.
     if args.tensor_parallel_size is not None:
         omni_kwargs["tensor_parallel_size"] = args.tensor_parallel_size
     if args.enforce_eager is not None:
@@ -603,6 +616,18 @@ def main():
     # gate is an engine-level config (offline analog of the server's --no-guardrails).
     if args.extra_body and "guardrails" in args.extra_body:
         omni_kwargs["model_config"] = {"guardrails": bool(args.extra_body["guardrails"])}
+    if args.enable_layerwise_offload:
+        omni_kwargs["enable_layerwise_offload"] = True
+    if args.vae_use_slicing:
+        omni_kwargs["vae_use_slicing"] = True
+    if args.vae_use_tiling:
+        omni_kwargs["vae_use_tiling"] = True
+    if args.enable_cache_dit_summary:
+        omni_kwargs["enable_cache_dit_summary"] = True
+    if args.enable_diffusion_pipeline_profiler:
+        omni_kwargs["enable_diffusion_pipeline_profiler"] = True
+    omni_kwargs.update(build_parallel_knob_kwargs(args))
+
     omni = Omni(**omni_kwargs)
     model_class_name = get_model_class_name(omni)
     declared_extra_body_params = get_extra_body_params(model_class_name)
@@ -621,11 +646,17 @@ def main():
     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"
+
+    def _knob_display(value: Any, default: Any) -> Any:
+        return value if value is not None else f"deploy/default ({default})"
+
     print(
         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}, "
+        f"ulysses_degree={_knob_display(args.ulysses_degree, 1)}, "
+        f"ulysses_mode={_knob_display(args.ulysses_mode, 'strict')}, "
+        f"ring_degree={_knob_display(args.ring_degree, 1)}, "
+        f"cfg_parallel_size={_knob_display(args.cfg_parallel_size, 1)}, "
+        f"vae_patch_parallel_size={_knob_display(args.vae_patch_parallel_size, 1)}, "
         f"enable_expert_parallel={args.enable_expert_parallel}."
     )
     print(f"  CPU offload: {args.enable_cpu_offload}; CPU Layerwise Offload: {args.enable_layerwise_offload}")
@@ -736,8 +767,9 @@ def main():
             if idx == 0 and prompt_info.get("omni_task") == ["t2i"]:
                 ar_width = int(prompt_info.get("ar_width", [0])[0])
                 ar_height = int(prompt_info.get("ar_height", [0])[0])
-                if ar_width > 0 and ar_height > 0:
-                    params.max_tokens = ar_height * (ar_width + 1) + 1
+                ar_grid_budget = ar_grid_max_tokens(ar_width, ar_height)
+                if ar_grid_budget is not None:
+                    params.max_tokens = ar_grid_budget
 
     if not diffusion_replaced and len(sampling_params_list) == 1:
         sampling_params_list = [diffusion_params]

examples/offline_inference/text_to_speech/cosyvoice3/end2end.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 import argparse
+import copy
 import functools
 import json
 import os
@@ -10,13 +11,11 @@ from typing import Any
 
 import numpy as np
 import soundfile as sf
-from vllm import SamplingParams
 from vllm.multimodal.media.audio import load_audio
 
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer
 from vllm_omni.model_executor.models.cosyvoice3.utils import extract_text_token
-from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config
 
 # Upstream zero-shot reference clip
 ZERO_SHOT_PROMPT_URL = "https://raw.githubusercontent.com/FunAudioLLM/CosyVoice/main/asset/zero_shot_prompt.wav"
@@ -59,8 +58,7 @@ def run_e2e():
         "--deploy-config",
         type=str,
         default=None,
-        help="Override the deploy config path. If unset, auto-loads "
-        "vllm_omni/deploy/cosyvoice3.yaml based on the HF model_type.",
+        help="Override the deploy config path. If unset, selects the bundled CosyVoice3 profile for the device.",
     )
     parser.add_argument("--text", type=str, default="Hello, this is a test of the CosyVoice system capability.")
     parser.add_argument(
@@ -105,8 +103,6 @@ def run_e2e():
         profiler_config=args.profiler_config,
     )
 
-    sampling_cfg = {"top_p": 0.8, "top_k": 25, "eos_token_id": 6561 + 1}
-
     print("Model initialized. Preparing inputs...")
     ref_audio_path = args.ref_audio or _default_ref_audio()
     if not os.path.exists(ref_audio_path):
@@ -138,7 +134,7 @@ def run_e2e():
 
     print(f"Generating for prompt: {args.text}")
 
-    config = CosyVoice3Config()
+    config = omni.engine.stage_vllm_configs[0].model_config.hf_config
     tokenizer = get_qwen_tokenizer(
         token_path=args.tokenizer,
         skip_special_tokens=config.skip_special_tokens,
@@ -149,29 +145,10 @@ def run_e2e():
     min_len = int(base_len * config.min_token_text_ratio)
     max_len = int(base_len * config.max_token_text_ratio)
 
-    # Build SamplingParams for each stage (GPT, S2Mel, Vocoder)
-    gpt_sampling = SamplingParams(
-        temperature=1.0,
-        top_p=sampling_cfg["top_p"],
-        top_k=sampling_cfg["top_k"],
-        repetition_penalty=2.0,
-        min_tokens=min_len,
-        max_tokens=max_len,
-        stop_token_ids=[sampling_cfg["eos_token_id"]],
-        # allowed_token_ids=list(range(6561+3)),
-        detokenize=False,
-    )
-    # Not used
-    s2mel_sampling = SamplingParams(
-        temperature=1.0,
-        top_p=1.0,
-        top_k=-1,
-        repetition_penalty=2.0,
-        max_tokens=256,
-        detokenize=False,
-    )
-
-    sampling_params_list = [gpt_sampling, s2mel_sampling]
+    # Keep the deploy profile's sampling policy, penalties and required stops.
+    sampling_params_list = copy.deepcopy(omni.default_sampling_params_list)
+    sampling_params_list[0].max_tokens = max(1, min(2048, max_len))
+    sampling_params_list[0].min_tokens = min(max(1, min_len), sampling_params_list[0].max_tokens)
 
     profiler_enabled = args.profiler_config is not None
     if profiler_enabled:

examples/online_serving/text_to_speech/cosyvoice3/run_server.sh

@@ -1,11 +1,15 @@
 #!/bin/bash
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 # Launch vLLM-Omni server for CosyVoice3 TTS
 #
 # Usage:
 #   ./run_server.sh
 #   CUDA_VISIBLE_DEVICES=0 ./run_server.sh
 #
-# Streaming (async-chunk) is on by default via vllm_omni/deploy/cosyvoice3.yaml.
+# Streaming (async-chunk) is on by default. H200 selects the optimized profile;
+# other devices use vllm_omni/deploy/cosyvoice3.yaml.
 # Set NO_ASYNC_CHUNK=1 to use the legacy synchronous path.
 
 set -e

recipes/MammothModa2/MammothModa2.md

@@ -18,12 +18,24 @@ The generic example formats the AR prompt, drives the AR → DiT stage pipeline,
 and forwards MammothModa2-specific generation parameters through the
 pipeline-declared `extra_body` contract.
 
-MammothModa2's DiT stage runs in the shared diffusion runtime in request mode.
-The first integration intentionally supports one request and one image per
-forward only (`max_num_seqs: 1`, `num_outputs_per_prompt: 1`). Request-level
-batching, step execution, continuous batching, cache acceleration,
-compilation, quantization, parallelism, and offload are not enabled by this
-recipe.
+MammothModa2's DiT stage runs in the shared diffusion runtime. The default
+`mammoth_moda2.yaml` uses request-level batching with up to eight compatible
+requests (`max_num_seqs: 8`) and one image per request
+(`num_outputs_per_prompt: 1`). TeaCache and Cache-DiT are optional alternative
+cache backends in request mode. Enable TeaCache with `--cache-backend tea_cache`
+as shown below, or enable Cache-DiT using the commented settings in that YAML.
+Compilation, DiT quantization, parallelism, and offload are not enabled by
+these presets.
+
+To use step execution and continuous batching, select
+[`mammoth_moda2_step.yaml`](../../vllm_omni/deploy/mammoth_moda2_step.yaml)
+with `--deploy-config vllm_omni/deploy/mammoth_moda2_step.yaml` in the
+text-to-image commands below. This preset inherits the same stage placement
+and capacity, enables step execution for Stage 1, and disables request-batch
+admission waiting and diffusion cache acceleration. Compatible requests can
+join between denoising steps and finish independently, including requests
+with different inference-step counts. Step mode cannot be combined with
+TeaCache, Cache-DiT, or other diffusion cache backends.
 
 Image size, seed, guidance, and denoising steps use the standard diffusion
 request fields. `cfg_range` remains a MammothModa2-specific `extra_body`
@@ -176,7 +188,30 @@ sets the relative step range `[start, end]` over which CFG is applied (default
 `[0.0, 1.0]`). For compatibility, `text_guidance_scale` and
 `num_inference_steps` remain accepted `extra_body` aliases and, when non-null,
 take precedence over the standard request fields. Model extras are filtered
-against the declared `extra_body_params` (see
+against the declared `extra_body_params`.
+
+TeaCache can be enabled for the DiT stage with the same user-facing sampling
+parameters:
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model ./MammothModa2-Preview \
+  --deploy-config vllm_omni/deploy/mammoth_moda2.yaml \
+  --prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \
+  --height 1024 \
+  --width 1024 \
+  --guidance-scale 4.0 \
+  --num-inference-steps 50 \
+  --cache-backend tea_cache \
+  --extra-body '{"cfg_range": [0.0, 1.0]}' \
+  --output mammoth_t2i_teacache.png
+```
+
+The bundled TeaCache coefficients were fitted from MammothModa2 full-compute
+traces. MammothModa2 uses the model-specific default `rel_l1_thresh=0.075`,
+selected for the evaluated 1024x1024, 50-step configuration.
+
+The model-specific keys are declared in
 [`vllm_omni/model_extras/mammothmodal2_preview.py`](../../vllm_omni/model_extras/mammothmodal2_preview.py)),
 so unknown MammothModa2 extras may be dropped.
 
@@ -287,6 +322,44 @@ python3 examples/offline_inference/text_to_image/text_to_image.py \
 
 The first request took 85.224 seconds. The AR stage generated 4,161 visual tokens in 72.996 seconds, and the DiT stage took 12.163 seconds. AR weight loading used 21.4 GiB and took 8.250 seconds. DiT weight loading used 5.49 GiB and took 1.824 seconds. The largest one second whole device memory sample was 106.57 GiB, including the AR KV cache reserved by the 0.5 memory setting.
 
+### 1x A800 80GB, MammothModa2 Preview (serving)
+
+#### Environment
+
+- OS: Linux, x86_64
+- CPU: 18 cores
+- GPU: one NVIDIA A800-SXM4-80GB
+- Deploy config: `vllm_omni/deploy/mammoth_moda2.yaml` (default split: AR `gpu_memory_utilization` 0.5, DiT 0.3)
+- vLLM-Omni version or commit: use the commit you are deploying from
+
+#### Serving Commands
+
+```bash
+vllm-omni serve bytedance-research/MammothModa2-Preview --omni \
+    --port 8091 \
+    --trust-remote-code
+
+python benchmarks/diffusion/diffusion_benchmark_serving.py \
+    --model bytedance-research/MammothModa2-Preview \
+    --endpoint /v1/images/generations \
+    --host 127.0.0.1 --port 8091 \
+    --height 1024 --width 1024 --num-inference-steps 50 \
+    --extra-body '{"text_guidance_scale": 9.0, "cfg_range": [0.0, 1.0]}' \
+    --num-prompts 8 --seed 142 --warmup-requests 0 \
+    --output-file mm2_1024_s50_c1.json
+```
+
+#### Verification
+
+On 1024x1024 / 50 steps, single-concurrency end-to-end mean latency is ~96 s
+(P99 ~101 s) per image and concurrency-4 mean is ~151 s (~0.024 img/s);
+steady-state combined GPU memory is ~50.4 GiB (AR ~39.2 GiB, DiT ~11.1 GiB).
+The AR stage dominates (~77 s of a ~96 s request) because it decodes a fixed
+4,161-token visual grid per image, so latency is insensitive to prompt length
+and scales only partially with DiT step count. See the
+[serving performance dashboard](../../benchmarks/diffusion/performance_dashboard/mammoth_moda2_serving_performance.md)
+for the full sweep, peak-memory, and component-attribution data.
+
 The output was a valid 1024 by 1024 RGB PNG.
 
 ### Startup benchmark

recipes/MiniMaxAI/MiniMax-H3.md

@@ -988,6 +988,42 @@ settings, and both videos with any comparison; the command alone is not quality
 evidence. A matching seed does not imply matching noise over differently sized
 full and windowed latents.
 
+## Request-scoped sampler
+
+H3 supports `euler` (the default) and `res_multistep`. Select the sampler per
+request; it applies to both video and audio, using their separate sigma schedules.
+For the HTTP examples above, set `num_inference_steps=20` and add
+`"sampler":"res_multistep"` to the existing `extra_params` object. For example,
+the T2VA fields become:
+
+```bash
+-F 'num_inference_steps=20' \
+-F 'extra_params={"task":"t2va","duration":8.7,"audio_flow_shift":3.0,"sampler":"res_multistep"}'
+```
+
+For the offline Python API, use `extra_args`:
+
+```python
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+sampling_params = OmniDiffusionSamplingParams(
+    num_inference_steps=20,
+    seed=42,
+    extra_args={"task": "t2va", "sampler": "res_multistep"},
+)
+```
+
+Both samplers use one DiT evaluation per step in request mode and step execution.
+`res_multistep` retains the previous denoised prediction independently for each
+request and each stream; its first and final updates use Euler. Omitting `sampler`
+or selecting `euler` preserves the existing Euler update. Reducing the step count
+reduces model evaluations, but quality depends on the prompt and configuration;
+20 RES steps are a candidate to compare against an Euler-50 reference.
+
+Requests using a fixed distilled `base_schedule` from a checkpoint or adapter
+must keep `sampler="euler"`. RES sampling with those schedules has not been
+validated and is rejected with `ValueError`.
+
 ## Request-scoped quality
 
 Add one of these fields to any HTTP request above. No Cache-DiT startup option
@@ -1408,6 +1444,7 @@ four-step adapter keeps its original sampling and fixed-top-k behavior.
 | `duration` | Workload-specific | Decimal seconds in `extra_params`; converted to H3-compatible frame count |
 | `fps` | `24` | H3 output FPS is fixed |
 | `num_inference_steps` | `50` | Matches the reference accuracy workloads |
+| `extra_params.sampler` | `euler` | Per-request `euler` or `res_multistep` for both streams; fixed distilled schedules require Euler |
 | `flow_shift` | `12` | Video sigma shift |
 | `audio_flow_shift` | `3` | Audio sigma shift, passed in `extra_params` |
 | `seed` | Task-specific | Use a fixed value for reproducibility |

recipes/README.md

@@ -66,6 +66,8 @@ recipes/
 | [`krea/Krea-2.md`](./krea/Krea-2.md) | Text-to-image (Turbo + Raw), offline + online, with LoRA | 1x H100 80GB / 1x Intel Arc Pro B70 32GB |
 | [`LTX/LTX-2.md`](./LTX/LTX-2.md) | LTX-2/LTX-2.3 text-to-video and image-to-video with synchronized audio | H200 141GB / 96GB-class GPU / 1x Intel Arc Pro B70 32GB |
 | [`LTX/LTX-2.5.md`](./LTX/LTX-2.5.md) | LTX-2.5-Diffusers: Full/SFT one-stage and distilled two-stage T2V/I2V with synchronized audio | NVIDIA B300; cuDNN-qualified / 1x Intel Arc Pro B70 32GB |
+| [`m-a-p/SheetSage2-H200.md`](./m-a-p/SheetSage2-H200.md) | Official audio-to-score preprocessing and YuE2 request export | 1x H200 141GB |
+| [`m-a-p/YuE2-3B.md`](./m-a-p/YuE2-3B.md) | Offline and online text-to-music with optional ABC score | 1x RTX 4090 24GB |
 | [`MammothModa2/MammothModa2.md`](./MammothModa2/MammothModa2.md) | Preview and Dev text-to-image (AR → DiT); Dev text/image understanding | Preview: 1x L40S 48GB / 1x ≥40GB GPU; Dev: 1x NVIDIA GPU with sufficient cache headroom |
 | [`NVIDIA/SANA-Video-2B.md`](./NVIDIA/SANA-Video-2B.md) | Native and Diffusers-adapter 480p/720p T2V+I2V | 1x RTX 5090 32GB |
 | [`meituan-longcat/LongCat-Video-Avatar-1.5.md`](./meituan-longcat/LongCat-Video-Avatar-1.5.md) | Audio-driven avatar video generation (AT2V / AI2V, single- and multi-speaker, AVC continuation) | 1x H100 80GB |

recipes/m-a-p/SheetSage2-H200.md

@@ -0,0 +1,259 @@
+# SheetSage2 audio-to-score preprocessing — H200
+
+## Summary
+
+- Vendor: m-a-p
+- Model: `m-a-p/SheetSage2` with the `m-a-p/MERT-v2-FullSong` parent
+- Task: recording → ABC score, MIDI and timed music annotations
+- Mode: standalone official Transformers inference; optional YuE2 request export
+- Hardware: one NVIDIA H200 (141 GB)
+- Maintainer: [@princepride](https://github.com/princepride)
+
+## When to use this recipe
+
+Use a recording as the musical score for a YuE2 cover, or inspect its melodies,
+chords and structure. The tool calls SheetSage2's official `transcribe()` API.
+It does not register SheetSage2 as a native vLLM model or start an Omni server.
+
+The transcription tool works independently of YuE2. For score-conditioned
+generation, use an Omni checkout containing the merged
+[PR #7886](https://github.com/vllm-project/vllm-omni/pull/7886)
+(`423f34326ed420e5acf0b1fb862a1b5ffb7e0fa7` or later). See the
+[YuE2 recipe](YuE2-3B.md) for the generation model's runtime contract.
+
+## Supported model contract
+
+| Item | Contract |
+| --- | --- |
+| Input | One local recording decodable by FFmpeg; upstream converts it to mono, 24 kHz |
+| Duration | Entire recording by default; `--max-seconds` limits preprocessing to a prefix |
+| Full score | Default: melody voices and chords; exported request uses `cot=full` |
+| Melody score | `--melody-only`: vocal and instrumental melodies, without score/playback chords; request uses `cot=melody` |
+| Outputs | `score.abc`, `transcription.mid`, upstream annotations and `manifest.json` |
+| YuE2 handoff | `--lyrics-file` and `--style` together also write `yue2_request.json` |
+| Output directory | Must be new or empty; failed exports produce no YuE2 request |
+| Deployment profile | One process, one H200, FP32 weight loading and BF16 inference autocast |
+
+`--melody-only` does not isolate or clone a singer. The upstream model still
+decodes its task annotations and retains those raw annotation files. YuE2
+uses the exported score to condition newly generated audio; a transcription
+does not guarantee an exact reproduction of the recording.
+
+## References
+
+- [Official SheetSage2 model card](https://huggingface.co/m-a-p/SheetSage2)
+- [Pinned SheetSage2 source](https://huggingface.co/m-a-p/SheetSage2/tree/cafc0df1021e14f49e928c4b345f5959d414ef64)
+- [Official YuE2 model card and cover documentation](https://huggingface.co/m-a-p/YuE2-3B)
+- Tool: [`tools/sheetsage2_transcribe.py`](../../tools/sheetsage2_transcribe.py)
+- Offline generation: [`examples/offline_inference/yue2/end2end.py`](../../examples/offline_inference/yue2/end2end.py)
+- [Supported generation models](../../docs/models/supported_models.md)
+
+## Hardware
+
+- Accelerator: one NVIDIA H200, 141 GB; no multi-device interconnect required.
+- Qualification: local M4A transcription, melody-only prefix, full-score
+  whole-song export and one offline YuE2 cover. Run transcription and generation
+  sequentially in separate processes. Other devices, CPU inference and
+  performance scaling are not qualified by this recipe.
+
+## Software environment
+
+- Ubuntu 22.04, Python 3.10, NVIDIA driver 590.48.01.
+- PyTorch/torchaudio 2.8.0 (CUDA 12.8 wheels), Transformers 4.45.2.
+- FFmpeg 4.4.2 was used for the local M4A check; upstream recommends FFmpeg 6.1.
+- vLLM: not required in the preprocessing environment.
+- vLLM-Omni: tool developed against main `a038b38179e9c788d3af6a8e73f94652afb247e4`.
+- YuE2 qualification: vLLM 0.30.0, PyTorch 2.13.0+cu130, Transformers 5.14.1;
+  #7886 at `5cc8acc942e76fa6f12f8b2f2558c2663abc2aa5` with the FP32 VAE-loading
+  and scalar `truncated` fixes subsequently included in the merged integration.
+
+## Command
+
+### Set up transcription
+
+Run from this repository's root. Create a **separate environment**: the
+official transcriber's Transformers version differs from Omni's serving
+dependencies. FFmpeg must be on `PATH`.
+
+```bash
+uv venv --python 3.10 .venv-sheetsage2
+uv pip install --python .venv-sheetsage2/bin/python \
+  -r tools/requirements/sheetsage2.txt
+
+.venv-sheetsage2/bin/python tools/sheetsage2_transcribe.py reference.wav \
+  --trust-remote-code --device cuda:0 --melody-only \
+  --output-dir outputs/reference-score
+```
+
+The default Hub checkpoint and its custom Python code are pinned to
+`cafc0df1021e14f49e928c4b345f5959d414ef64`. Its configuration pins the MERT-v2
+parent to `d8ba1c745e733b3908ce6ad16ebeb17ac7600a42`. Review that code before
+passing `--trust-remote-code`. `--revision` can override the default pin;
+custom Hub IDs use their default revision unless one is supplied.
+
+Local snapshots must preserve this pinned source layout: `modeling_sheetsage2.py`
+and its transitive relative-import dependencies are sibling `.py` files in the
+snapshot root. The local module-cache workaround targets that flat layout;
+reorganized snapshots with nested Python packages are not supported by it.
+
+For offline use, download **both** model snapshots, including their Python,
+JSON and safetensors files, before disconnecting:
+
+```bash
+.venv-sheetsage2/bin/hf download m-a-p/SheetSage2 \
+  --revision cafc0df1021e14f49e928c4b345f5959d414ef64 \
+  --local-dir models/SheetSage2
+.venv-sheetsage2/bin/hf download m-a-p/MERT-v2-FullSong \
+  --revision d8ba1c745e733b3908ce6ad16ebeb17ac7600a42 \
+  --local-dir models/MERT-v2-FullSong
+```
+
+### Transcribe a full recording for a cover
+
+Supply target lyrics in `lyrics.txt`, including YuE2 section tags such as
+`[Verse]` and `[Chorus]`. Match the lyric phrasing to the reference melody.
+This example retains the melody and chords of the whole recording:
+
+```bash
+.venv-sheetsage2/bin/python tools/sheetsage2_transcribe.py reference.wav \
+  --model models/SheetSage2 --base-model-path models/MERT-v2-FullSong \
+  --local-files-only --trust-remote-code --device cuda:0 \
+  --lyrics-file lyrics.txt --style 'Chinese folk, guzheng, gentle vocals' \
+  --seed 831001 --output-dir outputs/cover-score
+```
+
+Add `--melody-only` for a new arrangement without the original chords;
+this also changes the exported request from `cot=full` to `cot=melody`.
+For a short transcription smoke test, add `--max-seconds 45` and use lyrics
+appropriate to that excerpt. Use a fresh output directory for each run.
+
+### Generate offline
+
+Switch to your **Omni environment** for generation. Download `m-a-p/YuE2-3B`
+and `m-a-p/YuE2-Vae` to `models/YuE2-3B` and `models/YuE2-Vae`, respectively.
+The offline entrypoint needs a local YuE2 directory containing `qwen.tiktoken`.
+From the repository root, pass the exported ABC to the existing entrypoint:
+
+```bash
+CUDA_VISIBLE_DEVICES=0 python examples/offline_inference/yue2/end2end.py \
+  --model models/YuE2-3B --vae models/YuE2-Vae \
+  --lyrics "$(cat lyrics.txt)" \
+  --style 'Chinese folk, guzheng, gentle vocals' \
+  --abc-file outputs/cover-score/score.abc --cot full \
+  --seed 831001 --max-frames 9000 --gpu-memory-utilization 0.25 \
+  --output cover.wav
+```
+
+Use `--cot melody` if the score was transcribed with `--melody-only`.
+Changing `--cot` alone does not remove chords from an existing score.
+The 9,000-frame budget permits up to 360 seconds at 25 frames/s; it does
+not force that duration. The entrypoint's default 200-frame budget caps
+audio at 8 seconds. Check that the final report says `truncated=False`.
+The memory fraction above was used on one H200; it is not a measured
+minimum memory requirement.
+
+### Submit over HTTP (configuration-only)
+
+The exported JSON can also be submitted to YuE2's speech endpoint. These
+commands have not been qualified with an HTTP synthesis run in this recipe.
+Start YuE2 in its **Omni environment**:
+
+```bash
+vllm serve m-a-p/YuE2-3B --omni --port 8091
+```
+
+Submit the exported file:
+
+```bash
+curl --fail-with-body http://localhost:8091/v1/audio/speech \
+  -H 'Content-Type: application/json' \
+  --data-binary @outputs/cover-score/yue2_request.json \
+  --output cover.wav
+```
+
+The JSON contains the score itself, not a server-local file path. It sets
+`stream=false`, requests WAV, and preserves an optional generation seed.
+Use `--yue2-model` to match a custom served model name. The tool does not
+submit the request or set a generation token budget; the server's own
+duration and context limits apply.
+
+## Verification
+
+After a successful transcription, set `output` below to its output directory
+(`outputs/reference-score` for the transcription-only example). The YuE2
+request is checked when exported with `--lyrics-file` and `--style`:
+
+```bash
+.venv-sheetsage2/bin/python - <<'PY'
+import json
+from pathlib import Path
+import mido
+
+output = Path("outputs/cover-score")
+abc = (output / "score.abc").read_text(encoding="utf-8")
+manifest = json.loads((output / "manifest.json").read_text(encoding="utf-8"))
+assert abc.strip()
+request_path = output / "yue2_request.json"
+if request_path.exists():
+    request = json.loads(request_path.read_text(encoding="utf-8"))
+    assert request["extra_params"]["abc"] == abc
+    assert request["extra_params"]["cot"] == ("melody" if manifest["melody_only"] else "full")
+    print("YuE2 request verified")
+midi = mido.MidiFile(output / "transcription.mid")
+assert any(msg.type == "note_on" and msg.velocity for track in midi.tracks for msg in track)
+print("ABC, MIDI and manifest verified")
+PY
+```
+
+Run the CPU CLI contract tests in the repository's normal test
+environment (with pytest and pytest-mock installed):
+
+```bash
+python -m pytest tests/tools/test_sheetsage2_transcribe.py -m 'core_model and cpu' -q
+```
+
+These tests cover request construction, loading options, validation and
+failure handling with a mocked backend. They do not measure transcription
+accuracy or generated cover similarity.
+
+Local H200 checks with the pinned backend:
+
+| Input / mode | Result |
+| --- | --- |
+| 45 s M4A prefix, melody-only | 579 ABC characters; parseable MIDI with 118 note-on events |
+| 242.219 s M4A, full score | 3,033 ABC characters; parseable MIDI with 1,087 note-on events |
+| Default Hub ID, 10 s prefix | Successful ABC and MIDI export |
+| Local snapshots, fresh Transformers module cache | Successful offline full-score export |
+| Both generated JSON requests, #7886 adapter at `5cc8acc942e76fa6f12f8b2f2558c2663abc2aa5` | Validation and prompt construction passed (919 / 2,821 prompt tokens) |
+| Full score + new Mandarin lyrics, offline YuE2, `cot=full`, seed 831001, 9,000-frame budget | 241.319 s, 48 kHz stereo WAV; 6,034 generated tokens; `truncated=False` |
+
+The offline cover used a 108 BPM Chinese traditional pop style and one H200
+with `--gpu-memory-utilization 0.25`. The saved waveform contained finite
+samples, peak amplitude 0.9501, RMS 0.1232 and no full-scale samples. This is
+one functional run, not an accuracy or performance benchmark. It does not
+establish perceptual equivalence to the recording. HTTP synthesis and
+recordings longer than 300 seconds were not exercised. The local recording,
+lyrics, generated scores and cover audio are not included in the repository.
+
+## Notes
+
+- Weights load in FP32 because upstream merges MERT adapters before inference
+  autocast. `--dtype fp32` disables BF16 autocast.
+- For local snapshots the tool prepares transitive Python dependencies in
+  Transformers' module cache. This avoids the pinned Transformers version's
+  direct-import-only copy behavior without editing upstream files.
+- SheetSage2 manages long recordings with its own overlapping windows. This
+  tool leaves windowing and notation generation to the official implementation.
+- Failed runs may retain upstream diagnostic annotations. Choose a new output
+  directory when retrying; no incomplete request is published for synthesis.
+- Model weights have upstream license terms (CC BY-NC 4.0); this repository
+  does not bundle weights or recordings.
+
+## Supported features
+
+| Feature | Status |
+| --- | --- |
+| Audio-to-score preprocessing | Official backend; local file input |
+| YuE2 score-conditioned generation | Offline ABC handoff qualified on H200; HTTP request export validated, HTTP synthesis configuration-only |
+| Native Omni inference, continuous batching, TP/PP/SP | Not implemented for SheetSage2 |
+| Streaming transcription or cover audio | Not supported by this tool |

recipes/m-a-p/YuE2-3B.md

@@ -19,6 +19,11 @@ renders it with the lyrics. `cot=off` skips the score and generates music
 directly. Weights are **CC BY-NC 4.0** (non-commercial); this recipe loads
 user-downloaded weights and bundles none.
 
+To start from a recording, use the [SheetSage2 H200 recipe](SheetSage2-H200.md)
+to transcribe it into ABC in a separate environment, then pass the score to
+YuE2 with `--abc-file`. It includes full-song frame budgets and offline cover
+validation details.
+
 ## References
 
 - Model card: [m-a-p/YuE2-3B](https://huggingface.co/m-a-p/YuE2-3B)

requirements/common.txt

@@ -2,7 +2,7 @@
 # kernels==0.16.1 requires the LayerRepository version=/revision= integration
 # introduced in transformers 5.13. Keep the tested upper bound: 5.15 has
 # a reported model-construction regression; later versions need validation.
-# This range also satisfies vLLM 0.30.0's transformers>=5.10.4 requirement.
+# This range also satisfies vLLM 0.31.0's transformers>=5.10.4,<5.18.0 requirement.
 # MOSS-TTS-Nano's list/set weight-loading compatibility is handled by
 # transformers_keys_to_ignore_compat in model_executor/models/utils.py.
 transformers >= 5.13.0, < 5.15

tests/buildkite/test_amd_pipeline.py

@@ -13,6 +13,9 @@ AMD_MERGE_PIPELINE = Path(".buildkite/amd/test-amd-merge.yml")
 AMD_NIGHTLY_PIPELINE = Path(".buildkite/amd/test-amd-nightly.yml")
 AMD_READY_PIPELINE = Path(".buildkite/amd/test-amd-ready.yml")
 AMD_TEMPLATE = Path(".buildkite/amd/test-template-amd-omni.j2")
+AR_PAGED_ATTENTION_LABEL = "ROCm · AR Diffusion Paged Attention GPU Test"
+DIFFUSION_GROUP = ":card_index_dividers: Diffusion Test"
+AR_PAGED_ATTENTION_MARKERS = "core_model and rocm and MI325 and cards_1"
 
 
 def _find_step(label: str, pipeline_path: Path = AMD_MERGE_PIPELINE) -> dict:
@@ -31,6 +34,35 @@ def _find_step(label: str, pipeline_path: Path = AMD_MERGE_PIPELINE) -> dict:
     return step
 
 
+def _find_group(label: str, pipeline_path: Path) -> dict:
+    pipeline = yaml.safe_load(pipeline_path.read_text(encoding="utf-8"))
+    group = next((step for step in pipeline.get("steps", []) if step.get("group") == label), None)
+    assert group is not None, f"missing AMD pipeline group: {label}"
+    return group
+
+
+def test_ar_paged_attention_gpu_lane_is_blocking_and_pinned() -> None:
+    lane_definitions = []
+    for pipeline_path in (AMD_READY_PIPELINE, AMD_MERGE_PIPELINE):
+        group = _find_group(DIFFUSION_GROUP, pipeline_path)
+        step = next((item for item in group["steps"] if item.get("label") == AR_PAGED_ATTENTION_LABEL), None)
+        assert step is not None, f"missing {AR_PAGED_ATTENTION_LABEL} from {DIFFUSION_GROUP} in {pipeline_path}"
+        lane_definitions.append(step)
+
+        assert step["grade"] == "Blocking"
+        assert step["timeout_in_minutes"] == 20
+        assert "export VLLM_OMNI_AR_FA_REQUIRED=1" in step["commands"]
+
+        pytest_command = next(command for command in step["commands"] if "pytest" in command)
+        argv = split(pytest_command)
+        assert argv[:4] == ["timeout", "--signal=TERM", "--kill-after=2m", "15m"]
+        assert "tests/diffusion/ar_diffusion/test_paged_attention.py" in argv
+        assert argv[argv.index("-m") + 1] == AR_PAGED_ATTENTION_MARKERS
+        assert argv[argv.index("--run-level") + 1] == "core_model"
+
+    assert lane_definitions[0] == lane_definitions[1]
+
+
 def test_qwen3_tts_base_preserves_advanced_model_arguments() -> None:
     step = _find_step("Qwen3-TTS Base E2E Test")
     commands = step["commands"]

tests/buildkite/test_cuda_ar_paged_attention_lane.py

@@ -0,0 +1,31 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from pathlib import Path
+from shlex import split
+
+import pytest
+import yaml
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+CUDA_MERGE_PIPELINE = Path(".buildkite/cuda/test-merge.yml")
+DIFFUSION_GROUP = ":card_index_dividers: Diffusion Test"
+AR_PAGED_ATTENTION_LABEL = "Diffusion · AR Paged Attention GPU Test"
+
+
+def test_cuda_merge_keeps_ar_paged_attention_gpu_coverage() -> None:
+    pipeline = yaml.safe_load(CUDA_MERGE_PIPELINE.read_text(encoding="utf-8"))
+    group = next(step for step in pipeline["steps"] if step.get("group") == DIFFUSION_GROUP)
+    lane = next(step for step in group["steps"] if step.get("label") == AR_PAGED_ATTENTION_LABEL)
+
+    assert lane["timeout_in_minutes"] == 20
+    assert lane["mirror_hardwares"] == "l4_1"
+    assert "export VLLM_OMNI_AR_FA_REQUIRED=1" in lane["commands"]
+
+    pytest_command = next(command for command in lane["commands"] if "pytest" in command)
+    argv = split(pytest_command)
+    assert argv[:2] == ["timeout", "15m"]
+    assert "tests/diffusion/ar_diffusion/test_paged_attention.py" in argv
+    assert argv[argv.index("-m") + 1] == "core_model and cuda and L4 and cards_1"
+    assert argv[argv.index("--run-level") + 1] == "core_model"

tests/buildkite/test_rocm_dockerfile.py

@@ -10,7 +10,8 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 REPO_ROOT = Path(__file__).resolve().parents[2]
 AMD_TEMPLATE = REPO_ROOT / ".buildkite/amd/test-template-amd-omni.j2"
 AMD_BUILD_SCRIPT = REPO_ROOT / ".buildkite/amd/scripts/build-ci-image.sh"
-CI_DOCKERFILE = REPO_ROOT / "docker/Dockerfile.ci"
+CUDA_RELEASE_DOCKERFILE = REPO_ROOT / "docker/Dockerfile.cuda"
+CUDA_CI_DOCKERFILE = REPO_ROOT / "docker/Dockerfile.ci"
 ROCM_DOCKERFILE = REPO_ROOT / "docker/Dockerfile.rocm"
 ROCM_DOCKERIGNORE = REPO_ROOT / "docker/Dockerfile.rocm.dockerignore"
 
@@ -30,14 +31,18 @@ def _line_index(lines: list[str], prefix: str) -> int:
     return matches[0]
 
 
-def test_rocm_base_tracks_ci_vllm_release() -> None:
-    ci_release = _docker_arg(CI_DOCKERFILE, "VLLM_BASE_TAG")
+def test_rocm_base_tracks_cuda_vllm_release() -> None:
+    cuda_base = _docker_arg(CUDA_RELEASE_DOCKERFILE, "BASE_IMAGE")
+    cuda_image, _, cuda_release = cuda_base.rpartition(":")
     rocm_base = _docker_arg(ROCM_DOCKERFILE, "BASE_IMAGE")
 
     image_ref, separator, image_tag = rocm_base.rpartition(":")
     assert separator, f"expected a tagged ROCm base image, got {rocm_base}"
-    assert image_ref.rsplit("/", 1)[-1] == "vllm-openai-rocm", image_ref
-    assert image_tag == ci_release
+    assert cuda_image == "vllm/vllm-openai"
+    assert image_ref == "vllm/vllm-openai-rocm"
+    assert image_tag == cuda_release == f"v{_docker_arg(CUDA_CI_DOCKERFILE, 'VLLM_VERSION')}"
+    assert _docker_arg(CUDA_CI_DOCKERFILE, "VLLM_BASE_IMAGE") == cuda_image
+    assert _docker_arg(CUDA_CI_DOCKERFILE, "VLLM_BASE_TAG") == "v${VLLM_VERSION}"
 
 
 def test_rocm_defaults_to_prebuilt_base_image() -> None:
@@ -103,11 +108,11 @@ def test_rocm_build_context_excludes_git_history() -> None:
     assert ".git" in patterns
 
 
-def test_rocm_source_ref_tracks_ci_vllm_release() -> None:
-    ci_release = _docker_arg(CI_DOCKERFILE, "VLLM_BASE_TAG")
+def test_rocm_source_ref_tracks_cuda_vllm_release() -> None:
+    cuda_release = _docker_arg(CUDA_CI_DOCKERFILE, "VLLM_VERSION")
     rocm_source_ref = _docker_arg(ROCM_DOCKERFILE, "VLLM_VERSION_OR_COMMIT_HASH")
 
-    assert rocm_source_ref == ci_release
+    assert rocm_source_ref == f"v{cuda_release}"
 
 
 def test_rocm_dockerfile_contains_vllm_api_canary() -> None:

tests/buildkite/test_upload_pipeline.py

@@ -172,7 +172,11 @@ def test_mirror_hardwares_l4_1_expands_to_agents_and_plugins(monkeypatch: pytest
     env_names = {item["name"] for item in container["env"]}
     assert "VLLM_CI_HF_TOKEN" in env_names
     assert "HF_TOKEN" not in env_names
-    assert step["commands"] == [CUDA_HF_TOKEN_EXPORT, "pytest -sv tests/example"]
+    assert step["commands"] == [
+        CUDA_HF_TOKEN_EXPORT,
+        "python3 .buildkite/cuda/scripts/check_vllm_runtime.py",
+        "pytest -sv tests/example",
+    ]
 
 
 def test_mirror_hardwares_l4_preserves_explicit_retry() -> None:
@@ -239,6 +243,10 @@ def test_all_cuda_mirror_hardwares_restore_hf_token_at_runtime() -> None:
         assert "VLLM_CI_HF_TOKEN" in env_names
         assert "HF_TOKEN" not in env_names
         assert step["commands"][0] == CUDA_HF_TOKEN_EXPORT
+        assert step["commands"][1:] == [
+            "python3 .buildkite/cuda/scripts/check_vllm_runtime.py",
+            "pytest -sv tests/example",
+        ]
 
 
 def _gpu_limit(step: dict) -> int:

tests/config/test_config_factory.py

@@ -571,6 +571,56 @@ class TestPipelineDiscovery:
         assert "qwen3_omni_moe_thinker_only" in OMNI_PIPELINES
         assert "qwen3_tts" in OMNI_PIPELINES
 
+    @pytest.mark.parametrize(
+        "cuda,major,memory_gib,mps,optimized",
+        [
+            (False, None, 0, True, False),
+            (True, None, 141, True, False),
+            (True, 8, 141, True, False),
+            (True, 9, 80, True, False),
+            (True, 10, 180, True, False),
+            (True, 9, 141, False, False),
+            (True, 9, 141, True, True),
+        ],
+    )
+    def test_cosyvoice3_device_default_and_explicit_overrides(
+        self, monkeypatch, cuda, major, memory_gib, mps, optimized
+    ):
+        from vllm.platforms import current_platform
+        from vllm.platforms.interface import DeviceCapability
+
+        from vllm_omni.platforms import current_omni_platform
+
+        monkeypatch.setattr(current_platform, "is_cuda", lambda: cuda)
+        monkeypatch.setattr(current_omni_platform, "device_name", "cuda" if cuda else "cpu")
+        capability = DeviceCapability(major, 0) if major is not None else None
+        monkeypatch.setattr(current_platform, "get_device_capability", lambda: capability)
+        monkeypatch.setattr(current_platform, "get_device_total_memory", lambda: memory_gib * 1024**3)
+        monkeypatch.setattr(
+            "vllm_omni.model_executor.models.cosyvoice3.pipeline.shutil.which",
+            lambda command: "/usr/bin/nvidia-cuda-mps-control" if mps else None,
+        )
+        pipeline = resolve_pipeline_config("cosyvoice3")
+        assert pipeline is not None
+        config = VllmOmniConfig.from_pipeline_config(pipeline)
+        talker, codec = config.stage_configs
+        assert talker.scheduler_config.max_num_seqs == codec.scheduler_config.max_num_seqs == (32 if optimized else 8)
+        assert talker.runtime_config.cuda_mps == codec.runtime_config.cuda_mps == optimized
+        assert talker.model_config.hf_overrides == ({"cosyvoice3_sampling_mode": "standard"} if optimized else None)
+        assert pipeline.stages[0].sampling_constraints["stop_token_ids"] == list(range(6561, 6761))
+        if optimized:
+            assert codec.runtime_config.env["COSYVOICE3_CACHED_ISTFT"] == "1"
+            assert codec.runtime_config.env["COSYVOICE3_HIFT_GRAPH"] == "1"
+            override = VllmOmniConfig.from_pipeline_config(
+                pipeline, cli_overrides={"hf_overrides": {"cosyvoice3_sampling_mode": "ras", "seed": 7}}
+            )
+            assert override.stage_by_id(0).model_config.hf_overrides == {"cosyvoice3_sampling_mode": "ras", "seed": 7}
+        explicit = VllmOmniConfig.from_pipeline_config(
+            pipeline, deploy_config_path=get_deploy_config_path("cosyvoice3.yaml")
+        )
+        assert explicit.stage_by_id(0).scheduler_config.max_num_seqs == 8
+        assert explicit.stage_by_id(0).model_config.hf_overrides is None
+
     def test_registry_resolver_qwen3_omni_all_stages(self):
         """Test that providing the HF config for qwen3 omni with audio enabled uses all stages."""
         pipeline = resolve_pipeline_config(

tests/config/test_omni_config.py

@@ -1948,7 +1948,9 @@ def test_diffusion_config_field_classification_covers_current_fields():
 
     assert classified_fields == {f.name for f in fields(omni_config_module._DiffusionConfigProjection)}
     assert {f.name for f in fields(OmniDiffusionConfig)} <= (
-        classified_fields | omni_config_module._DIFFUSION_MOVED_SHARED_FIELDS
+        classified_fields
+        | omni_config_module._DIFFUSION_MOVED_SHARED_FIELDS
+        | omni_config_module._DIFFUSION_INTERNAL_FIELDS
     )
     assert {
         "enable_prompt_embed_cache",

tests/conftest.py

@@ -23,6 +23,7 @@ bootstrap_vllm_layer_custom_op_modules()
 pytest_plugins = (
     "tests.helpers.fixtures.config",
     "tests.helpers.fixtures.clean",
+    "tests.helpers.fixtures.cpu_threads",
     "tests.helpers.fixtures.log",
     "tests.helpers.fixtures.media",
     "tests.helpers.fixtures.pytest_collection",

tests/core/sched/test_generation_batch_coalescing.py

@@ -171,6 +171,7 @@ def test_cancelled_notification_is_filtered_before_coordinator(monkeypatch, mock
     s, _ = scheduler(monkeypatch, [event])
     s._generation_max_wait_s = 0
     s.waiting = []
+    s.kv_holding_waiting = []
     s.input_coordinator._async_chunk = True
     s.input_coordinator.update_request_metadata = mocker.Mock()
     s.input_coordinator.process_pending_chunks = mocker.Mock()
@@ -178,7 +179,11 @@ def test_cancelled_notification_is_filtered_before_coordinator(monkeypatch, mock
     s.input_coordinator.update_request_metadata.assert_called_once_with(
         s.requests, {"0": {"code_predictor_codes": [7]}}, model_mode="generation"
     )
-    s.input_coordinator.process_pending_chunks.assert_called_once_with(s.waiting, s.running, {"0"}, set())
+    s.input_coordinator.process_pending_chunks.assert_called_once()
+    waiting, running, ready, finished = s.input_coordinator.process_pending_chunks.call_args.args
+    assert list(waiting) == []
+    assert running is s.running
+    assert ready == {"0"} and finished == set()
 
 
 @pytest.mark.parametrize("target,wait", [(0, 1), (257, 1), (4, -1), (4, float("nan")), (4, float("inf"))])

tests/core/sched/test_generation_scheduler_restore.py

@@ -11,7 +11,9 @@ from types import SimpleNamespace
 import pytest
 import torch
 from vllm import SamplingParams
+from vllm.distributed.aux_output_connector.connector import AuxOutputSchedulerConnector
 from vllm.v1.core.sched.interface import PauseState
+from vllm.v1.core.sched.output import CachedRequestData
 from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue
 from vllm.v1.request import Request, RequestStatus
 
@@ -53,12 +55,14 @@ class FakeAdapter:
 def _make_generation_scheduler(waiting_request, *, use_v2_model_runner=False):
     scheduler = OmniGenerationScheduler.__new__(OmniGenerationScheduler)
     scheduler.max_num_scheduled_tokens = 8
+    scheduler.max_num_active_reqs = 1
     scheduler.max_num_running_reqs = 1
     scheduler._pause_state = PauseState.UNPAUSED
     scheduler.running = []
     scheduler.waiting = create_request_queue(SchedulingPolicy.FCFS)
     scheduler.waiting.add_request(waiting_request)
-    scheduler.skipped_waiting = create_request_queue(SchedulingPolicy.FCFS)
+    scheduler.kv_holding_waiting = create_request_queue(SchedulingPolicy.FCFS)
+    scheduler.deferred_waiting = set()
     scheduler.requests = {waiting_request.request_id: waiting_request}
     scheduler.policy = SchedulingPolicy.FCFS
     scheduler.chunk_transfer_adapter = FakeAdapter()
@@ -86,6 +90,7 @@ def _make_generation_scheduler(waiting_request, *, use_v2_model_runner=False):
         get_manager_metadata=lambda: None,
     )
     scheduler.connector = None
+    scheduler.aux_output_connector = None
     scheduler.ec_connector = None
     scheduler.prev_step_scheduled_req_ids = set()
     scheduler._pending_finish_reqs = []
@@ -105,6 +110,45 @@ def _make_generation_scheduler(waiting_request, *, use_v2_model_runner=False):
     return scheduler
 
 
+def test_generation_scheduler_preserves_auxiliary_output_metadata() -> None:
+    request = Request("aux-contract", [1, 2], SamplingParams(max_tokens=2), pooling_params=None)
+    request.block_hashes.append(b"block-hash")
+    scheduler = _make_generation_scheduler(request, use_v2_model_runner=True)
+    scheduler.aux_output_connector = AuxOutputSchedulerConnector()
+
+    output = scheduler._build_generation_scheduler_output(
+        new_reqs_data=[],
+        cached_reqs_data=CachedRequestData.make_empty(),
+        num_scheduled_tokens={request.request_id: 1},
+        scheduled_spec_decode_tokens={},
+        scheduled_encoder_inputs={},
+        num_common_prefix_blocks=[0],
+    )
+
+    metadata = output.aux_output_connector_metadata
+    assert metadata is not None
+    assert metadata.requests == {request.request_id: 0}
+    assert list(metadata.block_hashes[request.request_id]) == request.block_hashes
+
+
+@pytest.mark.parametrize(("runner_slots", "active_limit"), [(1, 2), (2, 1)])
+def test_generation_admits_kv_holder_before_fresh_requests(runner_slots, active_limit) -> None:
+    fresh = Request("fresh", [1, 2], SamplingParams(max_tokens=2), pooling_params=None)
+    holding = Request("holding", [3, 4], SamplingParams(max_tokens=2), pooling_params=None)
+    holding.num_computed_tokens = 1
+    scheduler = _make_generation_scheduler(fresh)
+    scheduler.max_num_running_reqs = runner_slots
+    scheduler.max_num_active_reqs = active_limit
+    scheduler.requests[holding.request_id] = holding
+    scheduler.kv_holding_waiting.add_request(holding)
+
+    output = scheduler.schedule()
+
+    assert set(output.num_scheduled_tokens) == {holding.request_id}
+    assert list(scheduler.waiting) == [fresh]
+    assert list(scheduler.kv_holding_waiting) == []
+
+
 def _chunk_request(request_id, **kwargs):
     defaults = dict(
         request_id=request_id,
@@ -123,7 +167,13 @@ def _chunk_request(request_id, **kwargs):
         record_event=lambda *args, **kwargs: None,
     )
     defaults.update(kwargs)
-    return SimpleNamespace(**defaults)
+    return _HashableChunkRequest(**defaults)
+
+
+class _HashableChunkRequest(SimpleNamespace):
+    # Deferred queues require identity hashing; SimpleNamespace types it as None.
+    __hash__ = object.__hash__  # type: ignore[assignment]
+    __eq__ = object.__eq__
 
 
 def test_chunk_lifecycle_no_resubmit_and_state_survives_requeue(monkeypatch) -> None:
@@ -138,7 +188,8 @@ def test_chunk_lifecycle_no_resubmit_and_state_survives_requeue(monkeypatch) ->
     )
     scheduler.running = []
     scheduler.waiting = create_request_queue(scheduler.policy)
-    scheduler.skipped_waiting = create_request_queue(scheduler.policy)
+    scheduler.kv_holding_waiting = create_request_queue(scheduler.policy)
+    scheduler.deferred_waiting = set()
 
     completed = Request("completed", [1, 2], SamplingParams(max_tokens=4), pooling_params=None)
     completed.status = RequestStatus.RUNNING
@@ -292,6 +343,7 @@ def _express_scheduler():
     scheduler = _make_generation_scheduler(continuation, use_v2_model_runner=True)
     scheduler.waiting.add_request(first)
     scheduler.requests[first.request_id] = first
+    scheduler.max_num_active_reqs = 4
     scheduler.max_num_running_reqs = 4
     scheduler._native_data_plane = True
     scheduler.chunk_transfer_adapter = None

tests/core/sched/test_omni_ar_scheduler_aborted_queue_sweep.py

@@ -9,7 +9,7 @@ handing control to upstream, because upstream ``Scheduler.schedule()`` raises
 admits in a finished state -- which kills the stage's engine core, not just
 the request.
 
-Regression for the sweep missing ``skipped_waiting``, the third queue
+Regression for the sweep missing ``kv_holding_waiting``, the third queue
 upstream admits from. An aborted duplex session parked there was re-selected
 by ``_select_waiting_queue_for_scheduling`` on a later tick and crashed the
 stage.
@@ -49,28 +49,30 @@ class _StubRequest:
         return (self.priority, self.arrival_time) < (other.priority, other.arrival_time)
 
 
-def _make_scheduler(policy: SchedulingPolicy, *, waiting=(), skipped_waiting=(), running=()):
+def _make_scheduler(policy: SchedulingPolicy, *, waiting=(), kv_holding_waiting=(), running=()):
     """Bare scheduler with just the surface the sweep reads."""
     scheduler = OmniARScheduler.__new__(OmniARScheduler)
     scheduler.waiting = create_request_queue(policy)
     for req in waiting:
         scheduler.waiting.add_request(req)
-    scheduler.skipped_waiting = create_request_queue(policy)
-    for req in skipped_waiting:
-        scheduler.skipped_waiting.add_request(req)
+    scheduler.kv_holding_waiting = create_request_queue(policy)
+    scheduler.deferred_waiting = set()
+    for req in kv_holding_waiting:
+        scheduler.kv_holding_waiting.add_request(req)
     scheduler.running = list(running)
     return scheduler
 
 
 @pytest.mark.parametrize("policy", _POLICIES)
-def test_schedule_sweeps_skipped_waiting_before_upstream_selection(
+def test_schedule_sweeps_kv_holding_waiting_before_upstream_selection(
     monkeypatch: pytest.MonkeyPatch,
     policy: SchedulingPolicy,
 ) -> None:
     """The public schedule path must remove aborted skipped requests first."""
     aborted = _StubRequest("req-aborted", RequestStatus.FINISHED_ABORTED)
-    scheduler = _make_scheduler(policy, skipped_waiting=[aborted])
+    scheduler = _make_scheduler(policy, kv_holding_waiting=[aborted])
     scheduler.chunk_transfer_adapter = None
+    scheduler.max_num_active_reqs = 8
     scheduler.max_num_running_reqs = 8
     scheduler.input_coordinator = None
     scheduler.num_waiting_for_streaming_input = 0
@@ -85,7 +87,7 @@ def test_schedule_sweeps_skipped_waiting_before_upstream_selection(
     sentinel = object()
 
     def upstream_schedule(_self, _throttle_prefills: bool = False):
-        assert list(scheduler.skipped_waiting) == []
+        assert list(scheduler.kv_holding_waiting) == []
         return sentinel
 
     monkeypatch.setattr(ar_sched_mod.VLLMScheduler, "schedule", upstream_schedule)

tests/core/sched/test_omni_ar_scheduler_kv_transfer.py

@@ -185,10 +185,15 @@ def test_stage_zero_final_skips_inter_stage_output(mocker):
 def test_streaming_session_update_invalidates_omits_kv_transfer_cache():
     scheduler = OmniARScheduler.__new__(OmniARScheduler)
     scheduler._omits_kv_transfer_cache = {}
-    scheduler.skipped_waiting = []
+    scheduler.deferred_waiting = set()
+    scheduler.kv_holding_waiting = []
     scheduler.log_stats = False
     tagged = apply_omni_final_stage_metadata(_make_engine_request(), final_stage_id=0)
-    session = SimpleNamespace(
+
+    class Session(SimpleNamespace):
+        __hash__ = object.__hash__
+
+    session = Session(
         request_id="req",
         additional_information=tagged.additional_information,
         status=RequestStatus.WAITING,

tests/core/sched/test_omni_ar_scheduler_logprobs.py

@@ -11,12 +11,14 @@ from types import MethodType, SimpleNamespace
 
 import numpy as np
 import pytest
+import torch
 from vllm.v1.engine import FinishReason
 from vllm.v1.outputs import LogprobsLists
 from vllm.v1.request import RequestStatus
 
 from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler, _slice_sampled_logprobs
 from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
+from vllm_omni.outputs import OmniModelRunnerOutput
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -132,11 +134,13 @@ def _make_scheduler_stub(requests: list[_Request]) -> SimpleNamespace:
     scheduler = SimpleNamespace(
         perf_metrics=None,
         connector=None,
+        aux_output_connector=None,
         chunk_transfer_adapter=None,
         requests={request.request_id: request for request in requests},
         running=list(requests),
         waiting=_RequestQueue(),
-        skipped_waiting=_RequestQueue(),
+        kv_holding_waiting=_RequestQueue(),
+        deferred_waiting=set(),
         structured_output_manager=SimpleNamespace(accept_tokens=lambda _request, _token_ids: True),
         transfer_triggered_requests=set(),
         active_kv_transfers=set(),
@@ -292,6 +296,36 @@ def test_invalid_logprobs_finish_only_the_affected_scheduler_request() -> None:
     np.testing.assert_array_equal(output_by_id["good"].new_logprobs.logprob_token_ids[:, 0], [8])
 
 
+@pytest.mark.parametrize("scored", [False, True])
+def test_fixed_token_scores_reach_engine_output(scored) -> None:
+    request = _Request("req")
+    request.sampling_params.num_logprobs = None
+    scheduler = _make_scheduler_stub([request])
+    _bind_request_lifecycle(scheduler, update_request=lambda req, tokens: (tokens, False))
+    scheduler_output = SimpleNamespace(
+        num_scheduled_tokens={"req": 3},
+        scheduled_spec_decode_tokens={},
+        num_invalid_spec_tokens=0,
+    )
+    scores = torch.tensor([[-0.5, -1.5], [-0.2, -2.2]])
+    model_runner_output = OmniModelRunnerOutput(
+        req_ids=["req"],
+        req_id_to_index={"req": 0},
+        sampled_token_ids=[[7]],
+        prompt_logprobs_dict={},
+        prompt_token_id_logprobs_dict={"req": scores} if scored else {},
+    )
+
+    outputs = OmniARScheduler.update_from_output(scheduler, scheduler_output, model_runner_output)
+    (output,) = outputs[0].outputs
+
+    assert output.new_token_ids == [7]
+    if scored:
+        torch.testing.assert_close(output.prompt_token_id_logprobs, scores)
+    else:
+        assert output.prompt_token_id_logprobs is None
+
+
 def _pooling_model_runner_output(pooler_tensor) -> SimpleNamespace:
     return SimpleNamespace(
         sampled_token_ids=[[]],

tests/core/sched/test_omni_ar_scheduler_stale_drain.py

@@ -73,7 +73,8 @@ def _replace_streaming_session(session: Request) -> None:
     sched.num_waiting_for_streaming_input = 0
     sched.log_stats = False
     sched.chunk_transfer_adapter = None
-    sched.skipped_waiting = set()
+    sched.kv_holding_waiting = set()
+    sched.deferred_waiting = set()
     session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
     update = StreamingUpdate(
         mm_features=None,

tests/core/sched/test_omni_ar_scheduler_streaming.py

@@ -23,6 +23,7 @@ from vllm.v1.engine import FinishReason
 from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue
 from vllm.v1.request import Request, RequestStatus, StreamingUpdate
 from vllm_omni.core.sched.omni_ar_scheduler import OmniARAsyncScheduler, OmniARScheduler
+from vllm_omni.model_executor.stage_input_processors.cosyvoice3 import talker2code2wav_async_chunk
 from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import (
     OmniChunkTransferAdapter,
 )
@@ -42,7 +43,8 @@ def _make_scheduler(*, stage_id: int = 0, session_mode: str = "turn") -> OmniARS
     sched.num_waiting_for_streaming_input = 0
     sched.log_stats = False
     sched.chunk_transfer_adapter = None
-    sched.skipped_waiting = set()
+    sched.kv_holding_waiting = set()
+    sched.deferred_waiting = set()
     sched._free_request_blocks = MagicMock()
     sched.encoder_cache_manager = MagicMock()
     sched._inflight_prefills = set()
@@ -287,7 +289,12 @@ def test_resumable_segment_boundary_keeps_pre_transition_send_watermark() -> Non
     )
 
 
-def test_running_decode_step_without_inter_stage_payload_does_not_raise() -> None:
+@pytest.mark.parametrize("processor", [None, talker2code2wav_async_chunk], ids=["no_opt_in", "cosyvoice3"])
+@pytest.mark.parametrize("new_token_ids", [[], [42]], ids=["no_tokens", "sampled_token"])
+@pytest.mark.parametrize("omit_chunk_transfer", [False, True], ids=["downstream", "stage0_final"])
+def test_running_decode_step_without_inter_stage_payload_does_not_raise(
+    processor, new_token_ids, omit_chunk_transfer, mocker
+) -> None:
     """A decode step that neither stops nor carries an inter-stage payload.
 
     ``finished`` is only assigned when the request stops, yet the async-chunk
@@ -296,14 +303,16 @@ def test_running_decode_step_without_inter_stage_payload_does_not_raise() -> Non
     """
     session = _make_request()
     session.status = RequestStatus.RUNNING
+    session.additional_information = {"omni_final_stage_id": 0} if omit_chunk_transfer else None
 
-    sched = MagicMock()
+    sched = mocker.MagicMock()
     sched.requests = {session.request_id: session}
     sched.perf_metrics = None
     sched.structured_output_manager.accept_tokens.return_value = True
-    sched._update_request_with_output.return_value = ([42], False)
+    sched._update_request_with_output.return_value = (new_token_ids, False)
     sched._process_kv_transfer_trigger.return_value = False
-    sched.chunk_transfer_adapter = MagicMock()
+    sched.chunk_transfer_adapter = mocker.MagicMock()
+    sched.chunk_transfer_adapter.custom_process_next_stage_input_func = processor
     sched.running = [session]
     sched.waiting_for_transfer_free = set()
     sched.transfer_triggered_requests = set()
@@ -316,13 +325,13 @@ def test_running_decode_step_without_inter_stage_payload_does_not_raise() -> Non
     sched.make_stats.return_value = None
     bind_omits_transfer_helpers(sched)
 
-    scheduler_output = MagicMock(spec=SchedulerOutput)
+    scheduler_output = mocker.MagicMock(spec=SchedulerOutput)
     scheduler_output.num_scheduled_tokens = {session.request_id: 1}
     scheduler_output.scheduled_spec_decode_tokens = {}
     scheduler_output.num_invalid_spec_tokens = 0
 
-    model_runner_output = MagicMock(spec=ModelRunnerOutput)
-    model_runner_output.sampled_token_ids = [[42]]
+    model_runner_output = mocker.MagicMock(spec=ModelRunnerOutput)
+    model_runner_output.sampled_token_ids = [new_token_ids]
     model_runner_output.logprobs = None
     model_runner_output.prompt_logprobs_dict = {}
     model_runner_output.pooler_output = None
@@ -335,8 +344,16 @@ def test_running_decode_step_without_inter_stage_payload_does_not_raise() -> Non
 
     OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output)
 
-    # Nothing to hand downstream: no payload, no segment boundary, not finished.
-    sched.chunk_transfer_adapter.save_async.assert_not_called()
+    if omit_chunk_transfer or processor is None or not new_token_ids:
+        sched.chunk_transfer_adapter.save_async.assert_not_called()
+    else:
+        sched.chunk_transfer_adapter.save_async.assert_called_once_with(
+            None,
+            session,
+            False,
+            new_token_ids=[42],
+            confirmed_num_computed_tokens=None,
+        )
 
 
 def test_queued_streaming_update_on_async_stop_fences_in_flight_once() -> None:
@@ -1169,7 +1186,8 @@ def test_chunk_segment_cleanup_keeps_requeued_resumable_receiver() -> None:
     sched.vllm_config = SimpleNamespace(model_config=SimpleNamespace(session_mode="duplex"))
     sched.running = []
     sched.waiting = queue()
-    sched.skipped_waiting = queue(session)
+    sched.kv_holding_waiting = queue(session)
+    sched.deferred_waiting = set()
     sched.num_waiting_for_streaming_input = 1
     sched.chunk_transfer_adapter = SimpleNamespace(
         receives_chunks=True,
@@ -1181,7 +1199,7 @@ def test_chunk_segment_cleanup_keeps_requeued_resumable_receiver() -> None:
 
     assert session.status == RequestStatus.WAITING
     assert session in sched.waiting.requests
-    assert session not in sched.skipped_waiting.requests
+    assert session not in sched.kv_holding_waiting.requests
     assert sched.num_waiting_for_streaming_input == 0
     assert session.request_id not in sched.chunk_transfer_adapter.segment_finished_requests
 
@@ -1208,7 +1226,8 @@ def test_chunk_segment_cleanup_keeps_explicit_update_stage_parked(
         receives_chunks=receives_chunks,
         segment_finished_requests={session.request_id},
     )
-    sched.skipped_waiting = MagicMock()
+    sched.kv_holding_waiting = MagicMock()
+    sched.deferred_waiting = set()
     sched._enqueue_waiting_request = MagicMock()
 
     sched._resume_downstream_chunk_receiver(session)
@@ -1216,7 +1235,7 @@ def test_chunk_segment_cleanup_keeps_explicit_update_stage_parked(
     assert session.status == RequestStatus.WAITING_FOR_STREAMING_REQ
     assert sched.num_waiting_for_streaming_input == 1
     assert session.request_id not in sched.chunk_transfer_adapter.segment_finished_requests
-    sched.skipped_waiting.remove_requests.assert_not_called()
+    sched.kv_holding_waiting.remove_requests.assert_not_called()
     sched._enqueue_waiting_request.assert_not_called()
 
 
@@ -1272,7 +1291,8 @@ def _park_session(sched: OmniARScheduler, session: Request) -> None:
     session.num_computed_tokens = 6
     session.num_output_placeholders = 0
     session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
-    sched.skipped_waiting.add_request(session)
+    sched.kv_holding_waiting.add_request(session)
+    sched.deferred_waiting.add(session)
     sched.num_waiting_for_streaming_input = 1
 
 
@@ -1295,7 +1315,7 @@ def test_stage0_streaming_update_that_overflows_max_model_len_finishes_the_sessi
     assert session.status == RequestStatus.FINISHED_ERROR
     assert session.request_id not in sched.requests
     assert len(sched.waiting) == 0
-    assert len(sched.skipped_waiting) == 0
+    assert len(sched.kv_holding_waiting) == 0
     assert sched.num_waiting_for_streaming_input == 0
     sched._free_request_blocks.assert_called_once_with(session)
     assert sched.finished_req_ids == {session.request_id}
@@ -1403,7 +1423,8 @@ def _make_admission_scheduler(*, max_model_len: int) -> OmniARScheduler:
     sched.max_model_len = max_model_len
     sched.policy = SchedulingPolicy.FCFS
     sched.waiting = create_request_queue(sched.policy)
-    sched.skipped_waiting = create_request_queue(sched.policy)
+    sched.kv_holding_waiting = create_request_queue(sched.policy)
+    sched.deferred_waiting = set()
     sched.running = []
     sched.requests = {}
     sched._free_request = MagicMock()  # needs a KV manager; not what this covers
@@ -1432,7 +1453,7 @@ def test_queued_streaming_update_that_overflows_does_not_return_to_admission() -
 
     assert finished is True
     assert len(sched.waiting) == 0
-    assert len(sched.skipped_waiting) == 0
+    assert len(sched.kv_holding_waiting) == 0
     assert session.status == RequestStatus.FINISHED_ERROR
     assert session.resumable is False
     assert "context_length_exceeded: " in sched._streaming_context_overflow[session.request_id][1]
@@ -1451,7 +1472,8 @@ def test_queued_streaming_update_that_fits_still_resumes_the_session() -> None:
     finished = sched._handle_stopped_request(session)
 
     assert finished is False
-    assert len(sched.waiting) == 1
+    assert list(sched.kv_holding_waiting) == [session]
+    assert session not in sched.deferred_waiting
     assert session.status == RequestStatus.WAITING
     assert not getattr(sched, "_streaming_context_overflow", {})
     sched._free_request.assert_not_called()
@@ -1473,7 +1495,7 @@ def test_queued_stop_with_a_recorded_overflow_still_leaves_admission() -> None:
 
     assert finished is True
     assert len(sched.waiting) == 0
-    assert len(sched.skipped_waiting) == 0
+    assert len(sched.kv_holding_waiting) == 0
     assert session.status == RequestStatus.FINISHED_ERROR
     assert session.resumable is False
     sched._free_request.assert_not_called()
@@ -1496,7 +1518,7 @@ def test_parked_stop_with_a_recorded_overflow_keeps_the_streaming_counter_balanc
     finished = sched._handle_stopped_request(session)
 
     assert finished is True
-    assert len(sched.skipped_waiting) == 0
+    assert len(sched.kv_holding_waiting) == 0
     assert sched.num_waiting_for_streaming_input == 0
     assert session.status == RequestStatus.FINISHED_ERROR
     sched._free_request.assert_not_called()
@@ -1529,6 +1551,7 @@ def test_async_chunk_reserves_parked_slots_during_ar_admission(monkeypatch, nati
     sched.running = []
     sched._native_data_plane = native
     sched.use_v2_model_runner = native
+    sched.max_num_active_reqs = 8
     sched.max_num_running_reqs = 8
     sched.input_coordinator = (
         SimpleNamespace(_async_chunk=True, _waiting_for_chunk_running=[parked], restore_queues=lambda _w, _r: None)
@@ -1555,7 +1578,7 @@ def test_async_chunk_reserves_parked_slots_during_ar_admission(monkeypatch, nati
     observed_limits: list[int] = []
 
     def fake_schedule(self, _throttle_prefills=False):
-        observed_limits.append(self.max_num_running_reqs)
+        observed_limits.append(self.max_num_active_reqs)
         return SimpleNamespace(scheduled_new_reqs=[])
 
     monkeypatch.setattr(VLLMScheduler, "schedule", fake_schedule)
@@ -1563,4 +1586,4 @@ def test_async_chunk_reserves_parked_slots_during_ar_admission(monkeypatch, nati
     sched.schedule()
 
     assert observed_limits == [7 if native else 8]
-    assert sched.max_num_running_reqs == 8
+    assert sched.max_num_active_reqs == 8

tests/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py

@@ -79,16 +79,19 @@ def _make_scheduler(
         else:
             hf_config.voxcpm2_runtime_config = runtime_config
     sched.vllm_config = SimpleNamespace(model_config=model_config)
-    # Upstream ``Scheduler.__init__`` always creates ``skipped_waiting``; the
+    # Upstream ``Scheduler.__init__`` always creates ``kv_holding_waiting``; the
     # abort sweep at the top of ``schedule()`` reads it.
-    sched.skipped_waiting = _MockQueue()
+    sched.kv_holding_waiting = _MockQueue()
+    sched.deferred_waiting = set()
     return sched
 
 
-def test_voxcpm2_unified_decode_graph_defers_waiting_when_decode_ready() -> None:
+@pytest.mark.parametrize("queue_name", ["waiting", "kv_holding_waiting"])
+def test_voxcpm2_unified_decode_graph_defers_waiting_when_decode_ready(queue_name) -> None:
     scheduler = _make_scheduler()
     scheduler.running = [_MockRequest("decode")]
-    scheduler.waiting = _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)])
+    scheduler.waiting = _MockQueue()
+    setattr(scheduler, queue_name, _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)]))
 
     assert scheduler._should_defer_waiting_for_unified_decode_graph()
 
@@ -144,12 +147,17 @@ def test_voxcpm2_unified_decode_graph_does_not_defer_without_cuda_graph(monkeypa
 
 def test_unified_decode_graph_deferral_restores_waiting_queue(monkeypatch) -> None:
     scheduler = _make_scheduler()
+    scheduler.max_num_active_reqs = 8
     scheduler.max_num_running_reqs = 8
     scheduler.running = [_MockRequest("decode")]
     original_waiting_req = _MockRequest("waiting", status=RequestStatus.WAITING)
     deferred_by_upstream = _MockRequest("deferred-by-upstream", status=RequestStatus.WAITING)
     original_waiting = _MockQueue([original_waiting_req])
     scheduler.waiting = original_waiting
+    kv_holder = _MockRequest("kv-holder", status=RequestStatus.WAITING)
+    original_kv_holding = _MockQueue([kv_holder])
+    scheduler.kv_holding_waiting = original_kv_holding
+    preempted_by_upstream = _MockRequest("preempted-by-upstream", status=RequestStatus.PREEMPTED)
     scheduler.policy = "fcfs"
     scheduler.chunk_transfer_adapter = None
     scheduler.input_coordinator = None
@@ -160,7 +168,10 @@ def test_unified_decode_graph_deferral_restores_waiting_queue(monkeypatch) -> No
     def fake_upstream_schedule(self, throttle_prefills: bool = False):
         assert self.waiting is not original_waiting
         assert not self.waiting
+        assert self.kv_holding_waiting is not original_kv_holding
+        assert not self.kv_holding_waiting
         self.waiting.add_request(deferred_by_upstream)
+        self.kv_holding_waiting.add_request(preempted_by_upstream)
         raise RuntimeError("stop before output wrapping")
 
     monkeypatch.setattr(scheduler_mod.VLLMScheduler, "schedule", fake_upstream_schedule)
@@ -169,3 +180,6 @@ def test_unified_decode_graph_deferral_restores_waiting_queue(monkeypatch) -> No
         scheduler.schedule()
 
     assert scheduler.waiting._items == [deferred_by_upstream, original_waiting_req]
+    assert scheduler.waiting is original_waiting
+    assert scheduler.kv_holding_waiting is original_kv_holding
+    assert scheduler.kv_holding_waiting._items == [preempted_by_upstream, kv_holder]

tests/core/sched/test_omni_cfg_ar_scheduler.py

@@ -23,10 +23,12 @@ def _scheduler(monkeypatch, waiting, running=()):
     scheduler.waiting = create_request_queue(scheduler.policy)
     for request in waiting:
         scheduler.waiting.add_request(request)
-    scheduler.skipped_waiting = create_request_queue(scheduler.policy)
+    scheduler.kv_holding_waiting = create_request_queue(scheduler.policy)
+    scheduler.deferred_waiting = set()
     scheduler.running = list(running)
     scheduler.requests = {request.request_id: request for request in [*waiting, *running]}
     scheduler.max_model_len = 128
+    scheduler.max_num_active_reqs = 2
     scheduler.max_num_running_reqs = 2
     scheduler.max_num_scheduled_tokens = 256
     scheduler.block_size = 16

tests/core/sched/test_omni_generation_scheduler_update_session.py

@@ -52,10 +52,11 @@ class _SchedulerStub(OmniGenerationScheduler):
         self.num_waiting_for_streaming_input = 0
         self.log_stats = log_stats
         self.chunk_transfer_adapter = _ChunkTransferAdapterStub()
-        self.skipped_waiting = _SkippedWaitingStub()
+        self.kv_holding_waiting = _SkippedWaitingStub()
+        self.deferred_waiting: set[Request] = set()
 
     def _enqueue_waiting_request(self, session: Request) -> None:
-        raise AssertionError("unexpected enqueue for skipped_waiting miss")
+        raise AssertionError("unexpected enqueue for kv_holding_waiting miss")
 
 
 class _AsyncChunkStopSchedulerStub(OmniGenerationScheduler):
@@ -134,7 +135,8 @@ def test_resumable_generation_stop_marks_segment_boundary() -> None:
     sched._handle_stopped_request.side_effect = rearm_resumable_request
     sched.running = [session]
     sched.waiting = MagicMock()
-    sched.skipped_waiting = MagicMock()
+    sched.kv_holding_waiting = MagicMock()
+    sched.deferred_waiting = set()
     sched.structured_output_manager.accept_tokens.return_value = True
     # Async scheduling can observe the segment boundary in the next schedule()
     # before this model output is applied and defer the same request for finish.
@@ -206,7 +208,8 @@ def test_rejected_grammar_finishes_and_frees_generation_request(prompt_complete:
         OmniSchedulerMixin._remove_stopped_requests_from_queues(sched, running, preempted)
     )
     sched.waiting = MagicMock()
-    sched.skipped_waiting = MagicMock()
+    sched.kv_holding_waiting = MagicMock()
+    sched.deferred_waiting = set()
     sched.structured_output_manager.accept_tokens.return_value = False
     sched._pending_finish_reqs = []
     sched.recompute_kv_load_failures = False
@@ -350,6 +353,8 @@ class _RealignSchedulerStub(OmniSchedulerMixin):
         self.requests = requests
         self.running = running
         self.waiting = waiting
+        self.kv_holding_waiting: list[Request] = []
+        self.deferred_waiting: set[Request] = set()
 
 
 class TestRealignRequestStatusToQueues:

tests/core/sched/test_omni_sched_deferred_free_fence.py

@@ -123,6 +123,7 @@ def _make_fast_path_sched(*, defer_block_free: bool) -> tuple[MagicMock, Request
 
     sched = MagicMock()
     sched.max_num_scheduled_tokens = 8
+    sched.max_num_active_reqs = 8
     sched.max_num_running_reqs = 8
     sched._native_data_plane = False
     sched._retains_state_across_chunks = False
@@ -135,6 +136,8 @@ def _make_fast_path_sched(*, defer_block_free: bool) -> tuple[MagicMock, Request
     sched.requests = {request.request_id: request}
     sched.running = [request]
     sched.waiting = []  # falsy: the waiting loop never runs
+    sched.kv_holding_waiting = []
+    sched.deferred_waiting = set()
     sched.policy = SchedulingPolicy.FCFS
     sched.chunk_transfer_adapter = None
     sched.scheduler_config = MagicMock(enable_chunked_prefill=True)

tests/core/sched/test_omni_scheduler_finish_requests_purge.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """Integration tests for the post-finish defensive purge of
 ``self.running`` placed in
 ``OmniARScheduler.finish_requests`` /
@@ -52,6 +55,8 @@ def _make_scheduler(scheduler_cls, *, requests, running, waiting):
     scheduler.requests = requests
     scheduler.running = running
     scheduler.waiting = waiting
+    scheduler.kv_holding_waiting = []
+    scheduler.deferred_waiting = set()
     return scheduler
 
 

tests/core/sched/test_omni_scheduler_input_coordinator_cleanup.py

@@ -52,6 +52,8 @@ def test_finish_requests_cleans_input_coordinator_for_finished_ids(
     scheduler.requests = {}
     scheduler.running = []
     scheduler.waiting = []
+    scheduler.kv_holding_waiting = []
+    scheduler.deferred_waiting = set()
 
     def fake_finish_requests(self, request_ids, finished_status):
         assert request_ids == ["req-a", "req-b"]

tests/core/sched/test_omni_scheduler_mixin_shared.py

@@ -7,7 +7,10 @@ from typing import Any
 
 import pytest
 from vllm.config import SchedulerConfig, VllmConfig
+from vllm.sampling_params import SamplingParams
+from vllm.v1.core.sched.scheduler import Scheduler
 from vllm.v1.engine import FinishReason
+from vllm.v1.request import Request
 
 from vllm_omni.config.model import OmniModelConfig
 from vllm_omni.core.sched import omni_scheduler_mixin
@@ -18,7 +21,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 class _Scheduler(OmniSchedulerMixin):
-    pass
+    _is_blocked_waiting_status = staticmethod(Scheduler._is_blocked_waiting_status)
 
 
 def test_async_chunk_adapter_initializes_for_stage_zero_sender_and_stage_one_receiver(monkeypatch):
@@ -120,7 +123,11 @@ def test_native_data_plane_uses_the_selected_input_protocol(async_chunk, stage_i
 def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths():
     calls: list[tuple[Any, ...]] = []
     scheduler = _Scheduler()
-    scheduler.waiting = ["waiting"]
+    waiting = Request("waiting", [1], SamplingParams(max_tokens=1), pooling_params=None)
+    holding = Request("holding", [1, 2], SamplingParams(max_tokens=1), pooling_params=None)
+    holding.num_computed_tokens = 1
+    scheduler.waiting = [waiting]
+    scheduler.kv_holding_waiting = [holding]
     scheduler.running = ["running"]
     scheduler.requests = {"request": object()}
     scheduler._consume_pending_connector_output = lambda mode: calls.append(("consume", mode))
@@ -137,16 +144,16 @@ def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths():
     scheduler.chunk_transfer_adapter = SimpleNamespace(
         receives_chunks=True,
         process_pending_chunks=lambda waiting, running, scheduler_requests: calls.append(
-            ("process", waiting, running, scheduler_requests)
+            ("process", list(waiting), running, scheduler_requests)
         ),
         restore_queues=lambda waiting, running, scheduler_requests: calls.append(
-            ("restore-chunks", waiting, running, scheduler_requests)
+            ("restore-chunks", list(waiting), running, scheduler_requests)
         ),
         collect_timed_out_request_ids=_collect_timed_out,
         collect_failed_send_request_ids=_collect_failed_sends,
     )
     scheduler.input_coordinator = SimpleNamespace(
-        restore_queues=lambda waiting, running: calls.append(("restore-full", waiting, running))
+        restore_queues=lambda waiting, running: calls.append(("restore-full", list(waiting), running))
     )
 
     scheduler._process_pending_omni_inputs("ar")
@@ -155,13 +162,13 @@ def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths():
     assert calls == [
         ("consume", "ar"),
         ("timeouts",),
-        ("process", scheduler.waiting, scheduler.running, scheduler.requests),
+        ("process", [holding, waiting], scheduler.running, scheduler.requests),
         # The chunk deadline runs after chunks are applied, so a chunk that
         # arrived this cycle resets the clock before it is measured (R1.1).
         ("chunk-timeouts", omni_scheduler_mixin.DEFAULT_INPUT_WAIT_TIMEOUT_S),
         ("failed-sends",),
-        ("restore-chunks", scheduler.waiting, scheduler.running, scheduler.requests),
-        ("restore-full", scheduler.waiting, scheduler.running),
+        ("restore-chunks", [holding, waiting], scheduler.running, scheduler.requests),
+        ("restore-full", [holding, waiting], scheduler.running),
     ]
 
 

tests/core/sched/test_omni_scheduler_ready_inbox.py

@@ -5,8 +5,13 @@
 from types import SimpleNamespace
 
 import pytest
+from vllm import SamplingParams
+from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue
+from vllm.v1.request import Request, RequestStatus
 
+from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
 from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
+from vllm_omni.core.sched.omni_scheduling_coordinator import OmniSchedulingCoordinator
 from vllm_omni.outputs import OmniConnectorOutput
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -15,6 +20,8 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 def test_ready_inbox_coalesces_live_events_and_drops_cancelled(mocker):
     scheduler = OmniSchedulerMixin()
     scheduler.requests, scheduler.waiting, scheduler.running = {"r": object()}, [], []
+    scheduler.kv_holding_waiting = []
+    scheduler._is_blocked_waiting_status = OmniARScheduler._is_blocked_waiting_status
     coordinator = SimpleNamespace(
         _async_chunk=True, update_request_metadata=mocker.Mock(), process_pending_chunks=mocker.Mock()
     )
@@ -32,7 +39,92 @@ def test_ready_inbox_coalesces_live_events_and_drops_cancelled(mocker):
     coordinator.update_request_metadata.assert_called_once_with(
         scheduler.requests, {"r": {"decode_token_end": 2}}, model_mode="ar"
     )
-    coordinator.process_pending_chunks.assert_called_once_with([], [], {"r"}, {"r"})
+    coordinator.process_pending_chunks.assert_called_once_with(mocker.ANY, [], {"r"}, {"r"})
+    assert list(coordinator.process_pending_chunks.call_args.args[0]) == []
     scheduler._consume_pending_connector_output(model_mode="ar")
     assert coordinator.update_request_metadata.call_count == 1
-    coordinator.process_pending_chunks.assert_called_with([], [], set(), set())
+    coordinator.process_pending_chunks.assert_called_with(mocker.ANY, [], set(), set())
+    assert list(coordinator.process_pending_chunks.call_args.args[0]) == []
+
+
+@pytest.mark.parametrize("policy", [SchedulingPolicy.FCFS, SchedulingPolicy.PRIORITY])
+@pytest.mark.parametrize("async_chunk", [False, True])
+def test_native_input_gate_parks_and_restores_kv_holders(policy, async_chunk):
+    scheduler = OmniARScheduler.__new__(OmniARScheduler)
+    scheduler.waiting = create_request_queue(policy)
+    scheduler.kv_holding_waiting = create_request_queue(policy)
+    scheduler.deferred_waiting = set()
+    scheduler.running = []
+    scheduler.chunk_transfer_adapter = None
+    scheduler.input_coordinator = OmniSchedulingCoordinator(
+        stage_id=1, scheduler_max_num_seqs=2, async_chunk=async_chunk
+    )
+    scheduler._init_omni_connector_output_inbox()
+    holder = Request("holder", [1, 2, 3], SamplingParams(max_tokens=1), pooling_params=None)
+    holder.num_computed_tokens = 2
+    fresh = Request("fresh", [1], SamplingParams(max_tokens=1), pooling_params=None)
+    scheduler.requests = {request.request_id: request for request in (holder, fresh)}
+    scheduler._enqueue_waiting_request(fresh)
+    scheduler._enqueue_waiting_request(holder)
+
+    scheduler._consume_pending_connector_output(model_mode="ar")
+
+    assert not scheduler.waiting and not scheduler.kv_holding_waiting
+    expected_status = RequestStatus.WAITING_FOR_CHUNK if async_chunk else RequestStatus.WAITING_FOR_INPUT
+    assert holder.status == fresh.status == expected_status
+    coordinator = scheduler.input_coordinator
+    registrations = coordinator.pending_chunk_registrations if async_chunk else coordinator.pending_input_registrations
+    assert [entry.request_id for entry in registrations] == ["holder", "fresh"]
+
+    scheduler._restore_omni_wait_queues()
+    assert list(scheduler.kv_holding_waiting) == [holder]
+    assert list(scheduler.waiting) == [fresh]
+    ready_output = (
+        OmniConnectorOutput(chunk_ready_req_ids={"holder"})
+        if async_chunk
+        else OmniConnectorOutput(stage_recv_req_ids={"holder"})
+    )
+    scheduler.enqueue_omni_connector_output(ready_output)
+    scheduler._consume_pending_connector_output(model_mode="ar")
+
+    assert holder.status == RequestStatus.WAITING
+    assert fresh.status == expected_status
+    assert list(scheduler.kv_holding_waiting) == [holder]
+    assert not scheduler.waiting
+    scheduler._restore_omni_wait_queues()
+    assert list(scheduler.kv_holding_waiting) == [holder]
+    assert list(scheduler.waiting) == [fresh]
+
+
+@pytest.mark.parametrize("async_chunk", [False, True])
+@pytest.mark.parametrize(
+    "status",
+    [
+        RequestStatus.WAITING_FOR_STREAMING_REQ,
+        RequestStatus.WAITING_FOR_REMOTE_KVS,
+        RequestStatus.WAITING_FOR_STRUCTURED_OUTPUT_GRAMMAR,
+    ],
+)
+def test_native_input_gate_preserves_upstream_blocked_waits(async_chunk, status):
+    scheduler = OmniARScheduler.__new__(OmniARScheduler)
+    scheduler.waiting = create_request_queue(SchedulingPolicy.FCFS)
+    scheduler.kv_holding_waiting = create_request_queue(SchedulingPolicy.FCFS)
+    scheduler.deferred_waiting = set()
+    scheduler.running = []
+    request = Request("blocked", [1, 2, 3], SamplingParams(max_tokens=1), pooling_params=None)
+    request.status = status
+    request.num_computed_tokens = 2
+    scheduler.requests = {request.request_id: request}
+    scheduler._enqueue_waiting_request(request)
+    scheduler.input_coordinator = OmniSchedulingCoordinator(
+        stage_id=1, scheduler_max_num_seqs=1, async_chunk=async_chunk
+    )
+    scheduler._init_omni_connector_output_inbox()
+
+    scheduler._consume_pending_connector_output(model_mode="ar")
+
+    assert request.status == status
+    assert list(scheduler.kv_holding_waiting) == [request]
+    assert scheduler.deferred_waiting == {request}
+    assert not scheduler.input_coordinator.pending_chunk_registrations
+    assert not scheduler.input_coordinator.pending_input_registrations

tests/core/sched/test_omni_scheduler_streaming_input_counter.py

@@ -59,7 +59,8 @@ def _make_scheduler(scheduler_cls, *, requests, running, waiting, counter, skipp
     scheduler.requests = requests
     scheduler.running = running
     scheduler.waiting = waiting
-    scheduler.skipped_waiting = skipped if skipped is not None else []
+    scheduler.kv_holding_waiting = skipped if skipped is not None else []
+    scheduler.deferred_waiting = set()
     scheduler.num_waiting_for_streaming_input = counter
     return scheduler
 
@@ -97,8 +98,8 @@ def test_finish_requests_clears_counter_left_by_a_stomped_status(
     assert scheduler.num_waiting_for_streaming_input == 0
 
 
-def test_resync_counts_parked_requests_in_skipped_waiting() -> None:
-    """``get_num_unfinished_requests`` sums ``waiting`` and ``skipped_waiting``,
+def test_resync_counts_parked_requests_in_kv_holding_waiting() -> None:
+    """``get_num_unfinished_requests`` sums ``waiting`` and ``kv_holding_waiting``,
     so the resync has to look at both or it under-counts and reintroduces the
     spin the counter exists to prevent."""
     parked = _StubRequest("req-skipped", RequestStatus.WAITING_FOR_STREAMING_REQ)
@@ -128,9 +129,11 @@ def test_ar_schedule_resyncs_before_delegating_upstream(monkeypatch: pytest.Monk
     scheduler.requests = {parked.request_id: parked}
     scheduler.running = [parked]
     scheduler.waiting = []
-    scheduler.skipped_waiting = []
+    scheduler.kv_holding_waiting = []
+    scheduler.deferred_waiting = set()
     scheduler.num_waiting_for_streaming_input = 1  # leaked by the status rewrite
     scheduler.policy = "fcfs"
+    scheduler.max_num_active_reqs = 8
     scheduler.max_num_running_reqs = 8
 
     seen: dict[str, int] = {}
@@ -161,12 +164,13 @@ def test_generation_scheduler_sweep_removes_aborted_requests_from_all_queues() -
     scheduler = OmniGenerationScheduler.__new__(OmniGenerationScheduler)
     scheduler.waiting = create_request_queue(SchedulingPolicy.FCFS)
     scheduler.waiting.add_request(aborted)
-    scheduler.skipped_waiting = create_request_queue(SchedulingPolicy.FCFS)
-    scheduler.skipped_waiting.add_request(aborted)
+    scheduler.kv_holding_waiting = create_request_queue(SchedulingPolicy.FCFS)
+    scheduler.deferred_waiting = set()
+    scheduler.kv_holding_waiting.add_request(aborted)
     scheduler.running = [aborted]
 
     scheduler._drop_aborted_queued_requests()
 
     assert list(scheduler.waiting) == []
-    assert list(scheduler.skipped_waiting) == []
+    assert list(scheduler.kv_holding_waiting) == []
     assert scheduler.running == []

tests/dfx/reliability/invalid_param_test/test_invalid_omni_chat.py

@@ -216,44 +216,40 @@ def test_video_chat_stream_invalid_requests(
 # WS /v1/realtime
 # ─────────────────────────────────────────────────────────────────────────────
 
-# Resolved in ``test_realtime_invalid_requests`` (skip ``session.created`` handshake).
-_REALTIME_WS_MODEL_MISMATCH = object()
-_REALTIME_WS_INVALID_AUDIO_APPEND = object()
-
 
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
 @pytest.mark.parametrize(
     "send_frames_spec, ws_error_code, err_message",
     [
-        pytest.param("{not-json", "invalid_json", "Invalid JSON", id="invalid_json"),
+        pytest.param("{not-json", "invalid_event", "Invalid or unrecognized client event", id="invalid_json"),
         pytest.param(
             json.dumps({"type": "session.update"}),
             "invalid_event",
-            ("model", "Missing required field"),
-            id="session_update_missing_model",
+            "Invalid or unrecognized client event",
+            id="session_update_missing_session",
         ),
         pytest.param(
-            _REALTIME_WS_MODEL_MISMATCH,
-            "model_not_found",
-            ("model", "does not exist"),
-            id="session_update_model_mismatch",
+            json.dumps({"type": "session.update", "session": []}),
+            "invalid_event",
+            "Invalid or unrecognized client event",
+            id="session_update_invalid_session",
         ),
         pytest.param(
-            json.dumps({"type": "input_audio_buffer.commit", "final": False}),
-            "model_not_validated",
-            ("session.update", "validate the model"),
-            id="commit_without_model_validation",
+            json.dumps({"type": "input_audio_buffer.commit"}),
+            "invalid_request_error",
+            "Input audio buffer is empty",
+            id="commit_without_audio",
         ),
         pytest.param(
             json.dumps({"type": "unknown.event"}),
-            "unknown_event",
-            "Unknown event type",
+            "invalid_event",
+            "Invalid or unrecognized client event",
             id="unknown_event_type",
         ),
         pytest.param(
-            _REALTIME_WS_INVALID_AUDIO_APPEND,
-            "invalid_audio",
-            "Invalid audio data",
+            json.dumps({"type": "input_audio_buffer.append", "audio": "not-valid-base64!!!"}),
+            "invalid_request_error",
+            "Invalid base64 audio data",
             id="invalid_audio_append",
         ),
     ],
@@ -266,27 +262,12 @@ def test_realtime_invalid_requests(
     ws_error_code: str,
     err_message: str | tuple[str, ...],
 ) -> None:
-    send_frames: Any
-    if send_frames_spec is _REALTIME_WS_MODEL_MISMATCH:
-        send_frames = json.dumps(
-            {
-                "type": "session.update",
-                "model": "this-model-is-not-served-on-this-instance",
-            }
-        )
-    elif send_frames_spec is _REALTIME_WS_INVALID_AUDIO_APPEND:
-        send_frames = [
-            json.dumps({"type": "session.update", "model": omni_server.model}),
-            json.dumps({"type": "input_audio_buffer.append", "audio": "not-valid-base64!!!"}),
-        ]
-    else:
-        send_frames = send_frames_spec
     online_client.send_realtime_ws_request(
         {
-            "send_frames": send_frames,
+            "send_frames": send_frames_spec,
             "timeout": 120,
             "ws_max_size": None,
-            "ws_skip_types": ["session.created"],
+            "ws_skip_types": ["session.created", "conversation.created"],
             "ws_json_type": "error",
             "ws_error_code": ws_error_code,
             "err_message": err_message,

tests/diffusion/ar_diffusion/test_history_staging_gpu.py

@@ -11,7 +11,7 @@ from tests.diffusion.ar_diffusion.test_paged_attention import (
     N_HEADS,
     POS,
     _commit_video_span,
-    _gpu_flash_attn_usable,
+    _require_gpu_flash_attn,
     make_state,
 )
 from tests.helpers.mark import hardware_test
@@ -22,14 +22,18 @@ from vllm_omni.platforms import current_omni_platform
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
 
 
-@pytest.mark.skipif(
-    torch.version.hip is not None or not _gpu_flash_attn_usable(), reason="usable CUDA FlashAttention is required"
-)
+def _require_cuda_flash_attn() -> None:
+    if torch.version.hip is not None:
+        pytest.skip("usable CUDA FlashAttention is required")
+    _require_gpu_flash_attn()
+
+
 @hardware_test(res={"cuda": ["L4", "H100"]}, num_cards=1)
 @pytest.mark.parametrize("window_chunks", [2, 4])
 @torch.inference_mode()
 def test_staged_attention_updates_only_the_active_window_gpu(monkeypatch, window_chunks):
     """Full and reused gathers match fresh attention without touching spare capacity or stale history."""
+    _require_cuda_flash_attn()
     monkeypatch.setenv(KV_GATHER_ENV, "1")
     device, dtype = torch.device("cuda"), torch.bfloat16
     kv, st = make_state(device=device, dtype=dtype, window_chunks=window_chunks, reuse_history_staging=True)
@@ -65,15 +69,13 @@ def test_staged_attention_updates_only_the_active_window_gpu(monkeypatch, window
             assert (staged[inputs.max_seq_len :] == -7).all()
 
 
-@pytest.mark.skipif(
-    torch.version.hip is not None or not _gpu_flash_attn_usable(), reason="usable CUDA FlashAttention is required"
-)
 @hardware_test(res={"cuda": ["L4", "H100"]}, num_cards=1)
 @pytest.mark.parametrize("window_chunks", [2, 4])
 @pytest.mark.parametrize("reuse_history", [False, True])
 @torch.inference_mode()
 def test_staged_attention_replays_inductor_cudagraph(monkeypatch, window_chunks, reuse_history):
     """Inductor must capture the real mutable op and replay with fresh current K/V."""
+    _require_cuda_flash_attn()
     monkeypatch.setenv(KV_GATHER_ENV, "1")
     device, dtype = torch.device("cuda"), torch.bfloat16
     kv, st = make_state(device=device, dtype=dtype, window_chunks=window_chunks, reuse_history_staging=True)

tests/diffusion/ar_diffusion/test_paged_attention.py

@@ -4,6 +4,7 @@
 
 from __future__ import annotations
 
+import os
 import subprocess
 from importlib.util import find_spec
 from pathlib import Path
@@ -12,6 +13,7 @@ from unittest.mock import MagicMock
 import pytest
 import torch
 
+from tests.helpers.mark import hardware_test
 from vllm_omni.experimental.ar_diffusion.capability import ARDiffusionKVBranchSpec
 from vllm_omni.experimental.ar_diffusion.kv_cache import (
     ARDiffusionKVCache,
@@ -26,7 +28,7 @@ from vllm_omni.experimental.ar_diffusion.kv_cache.config import KV_GATHER_ENV
 from vllm_omni.experimental.ar_diffusion.kv_cache.paged import ChunkWindowManager
 from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState
 
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
 
 
 BLOCK = 16
@@ -118,6 +120,17 @@ def _gpu_flash_attn_usable() -> bool:
     return driver_major >= 580
 
 
+def _require_gpu_flash_attn() -> None:
+    """Skip optional local runs, but never let a dedicated GPU lane pass vacuously."""
+    if _gpu_flash_attn_usable():
+        return
+
+    reason = "usable GPU FlashAttention is required"
+    if os.environ.get("VLLM_OMNI_AR_FA_REQUIRED"):
+        pytest.fail(f"{reason} when VLLM_OMNI_AR_FA_REQUIRED is set")
+    pytest.skip(reason)
+
+
 def _commit_video_span(
     kv: ARDiffusionKVCache,
     st: ARDiffusionKVState,
@@ -139,6 +152,7 @@ def _commit_video_span(
     return k, v
 
 
+@pytest.mark.cpu
 def test_paged_context_allocates_lazily_and_commits_after_forward():
     _, st = make_state()
 
@@ -157,6 +171,7 @@ def test_paged_context_allocates_lazily_and_commits_after_forward():
     assert st._committed[POS] == BLOCK
 
 
+@pytest.mark.cpu
 def test_scratch_video_and_action_blocks_do_not_commit():
     kv, st = make_state()
 
@@ -171,6 +186,7 @@ def test_scratch_video_and_action_blocks_do_not_commit():
     assert st._committed[POS] == 0
 
 
+@pytest.mark.cpu
 def test_pipeline_kv_get_paged_path_has_no_gather_backend():
     kv, st = make_state()
     assert not hasattr(kv, "gather_window_all_layers")
@@ -188,6 +204,7 @@ def test_pipeline_kv_get_paged_path_has_no_gather_backend():
 @pytest.mark.parametrize("history_chunks", [0, 1, 3])
 @pytest.mark.parametrize("action_len", [0, 3])
 @pytest.mark.parametrize("commit_current", [False, True])
+@pytest.mark.cpu
 def test_paged_attention_matches_dense_reference_cpu(history_chunks, action_len, commit_current):
     torch.manual_seed(0)
     device = torch.device("cpu")
@@ -264,6 +281,7 @@ def test_paged_attention_matches_dense_reference_cpu(history_chunks, action_len,
 
 @pytest.mark.parametrize("history_chunks", [0, 1, 2, 3])
 @pytest.mark.parametrize("window_chunks", [2, 4])
+@pytest.mark.cpu
 def test_staged_reuse_refreshes_the_current_blocks_at_their_live_offset(monkeypatch, history_chunks, window_chunks):
     """A second probe of the same AR block restages its current K/V where the table actually holds it.
 
@@ -330,6 +348,7 @@ def test_staged_reuse_refreshes_the_current_blocks_at_their_live_offset(monkeypa
 
 
 @pytest.mark.parametrize("gather_enabled", [False, True])
+@pytest.mark.cpu
 def test_staging_is_only_allocated_for_the_gather_path(monkeypatch, gather_enabled):
     if gather_enabled:
         monkeypatch.setenv(KV_GATHER_ENV, "1")
@@ -351,6 +370,7 @@ def test_staging_is_only_allocated_for_the_gather_path(monkeypatch, gather_enabl
 
 
 @pytest.mark.parametrize("window_chunks", [2, 4])
+@pytest.mark.cpu
 def test_layer_inputs_preserve_static_staging_tensors(monkeypatch, window_chunks):
     """Compiled inputs must retain the manager's static-address annotations, including with spare capacity."""
     monkeypatch.setenv(KV_GATHER_ENV, "1")
@@ -368,11 +388,12 @@ def test_layer_inputs_preserve_static_staging_tensors(monkeypatch, window_chunks
         assert stage_key.shape[0] == (window_chunks + 1) * BLOCK
 
 
-@pytest.mark.skipif(not _gpu_flash_attn_usable(), reason="usable GPU FlashAttention is required")
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
 @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):
+    _require_gpu_flash_attn()
     torch.manual_seed(0)
     device = torch.device("cuda")
     dtype = torch.float16
@@ -439,6 +460,7 @@ def test_paged_attention_matches_dense_reference_gpu(history_chunks, action_len,
     torch.testing.assert_close(paged, ref, rtol=2e-2, atol=2e-2)
 
 
+@pytest.mark.cpu
 def test_block_table_padded_to_fixed_width():
     """Shapes must be constant across window growth: only values change."""
     device = torch.device("cpu")
@@ -462,6 +484,7 @@ def test_block_table_padded_to_fixed_width():
     assert int(ctx2.seq_lens[0]) == 2 * BLOCK
 
 
+@pytest.mark.cpu
 def test_prepare_is_idempotent_and_layers_share_metadata():
     device = torch.device("cpu")
     kv, st = make_state(num_layers=2)
@@ -487,6 +510,7 @@ def test_prepare_is_idempotent_and_layers_share_metadata():
     assert i1.value_pool is kv._v_pools[1]
 
 
+@pytest.mark.cpu
 def test_layer_inputs_before_prepare_raises():
     _, st = make_state()
     layer_ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)[0]
@@ -494,6 +518,7 @@ def test_layer_inputs_before_prepare_raises():
         layer_ctx.to_layer_inputs()
 
 
+@pytest.mark.cpu
 def test_custom_op_registration_idempotent():
     import importlib
     import sys
@@ -508,6 +533,7 @@ def test_custom_op_registration_idempotent():
     assert hasattr(torch.ops.vllm_omni, "ar_diffusion_paged_write_attn")
 
 
+@pytest.mark.cpu
 def test_custom_op_mutable_arguments_cannot_be_elided_as_defaults():
     # Older PyTorch ADInplaceOrView handlers index positional mutable inputs
     # directly. Default-valued trailing inputs can disappear before that handler.
@@ -518,6 +544,7 @@ def test_custom_op_mutable_arguments_cannot_be_elided_as_defaults():
 
 
 @pytest.mark.parametrize("reuse_history_staging", [False, True])
+@pytest.mark.cpu
 def test_custom_op_compiles_fullgraph_without_recompile_on_value_change(monkeypatch, reuse_history_staging):
     """The op must trace as one opaque node: fullgraph OK, and changed tensor
     VALUES (new slots / block ids) must not trigger recompilation."""
@@ -576,6 +603,7 @@ RAGGED_CHUNK = 24
 
 @pytest.mark.parametrize("commit_current", [False, True])
 @pytest.mark.parametrize("action_len", [0, 3])
+@pytest.mark.cpu
 def test_a_ragged_chunk_still_matches_the_dense_reference(commit_current, action_len):
     """Attention reads a sequence, not a set of blocks.
 
@@ -643,6 +671,7 @@ def test_a_ragged_chunk_still_matches_the_dense_reference(commit_current, action
     torch.testing.assert_close(paged, ref, rtol=1e-5, atol=1e-5)
 
 
+@pytest.mark.cpu
 def test_a_ragged_scratch_chunk_is_physically_continuous_with_its_history():
     """The write targets themselves must leave no hole.
 
@@ -668,6 +697,7 @@ def test_a_ragged_scratch_chunk_is_physically_continuous_with_its_history():
     assert ctx.action_scratch_block_ids[0] not in ctx.current_video_block_ids
 
 
+@pytest.mark.cpu
 def test_a_ragged_window_keeps_the_block_table_a_fixed_shape():
     """The table's shape must not track how full the window is.
 
@@ -769,6 +799,7 @@ def _model_block_indices(kv, adapter) -> dict[int, int]:
     ("sink_chunks", "window_chunks", "reset_at_boundary"),
     [(1, 2, False), (0, 2, False), (1, 3, False), (2, 2, True), (0, 2, True)],
 )
+@pytest.mark.cpu
 def test_attention_after_eviction_reads_every_kept_token_and_nothing_unwritten(
     sink_chunks, window_chunks, reset_at_boundary
 ):
@@ -859,6 +890,7 @@ def test_attention_after_eviction_reads_every_kept_token_and_nothing_unwritten(
     ("sink_chunks", "window_chunks", "reset_at_boundary"),
     [(1, 2, False), (0, 2, False), (1, 3, False), (2, 2, True), (0, 2, True)],
 )
+@pytest.mark.cpu
 def test_compaction_does_not_change_which_tokens_stay_resident(
     monkeypatch, sink_chunks, window_chunks, reset_at_boundary
 ):
@@ -898,6 +930,7 @@ def test_compaction_does_not_change_which_tokens_stay_resident(
         assert got == want, f"tick {tick}: compaction kept blocks {sorted(got - want)} and dropped {sorted(want - got)}"
 
 
+@pytest.mark.cpu
 def test_history_staging_holds_a_ragged_window_and_restages_it_whole(monkeypatch):
     """The staging buffers must fit the table width this module pads to.
 
@@ -925,6 +958,7 @@ def test_history_staging_holds_a_ragged_window_and_restages_it_whole(monkeypatch
 
 
 @pytest.mark.parametrize("commit_current", [False, True])
+@pytest.mark.cpu
 def test_action_tokens_after_a_partly_written_video_block_are_refused(commit_current):
     """Action K/V follows the video blocks in the table, and the kernel reads it as one run.
 
@@ -939,6 +973,7 @@ def test_action_tokens_after_a_partly_written_video_block_are_refused(commit_cur
         ctx.build_block_table(action_len=3, query_len=RAGGED_CHUNK + 3, device=torch.device("cpu"))
 
 
+@pytest.mark.cpu
 def test_the_shipped_geometry_reads_exactly_the_window_on_every_tick():
     """832x480 with the checkpoint's own sink and window, slot by slot.
 
@@ -988,7 +1023,7 @@ def test_the_shipped_geometry_reads_exactly_the_window_on_every_tick():
         kv_len = int(seq_lens[0])
 
         position_of_block = {b: i * block for b, i in _model_block_indices(kv, st.adapter(POS)).items()}
-        read = []
+        read: list[tuple[int, int]] = []
         for b in table[0].tolist()[: -(-kv_len // block)]:
             assert b in position_of_block, f"tick {tick}: read block {b}, which holds no token of this session"
             read.extend((b * block + o, position_of_block[b] + o) for o in range(block))
@@ -1039,6 +1074,7 @@ class _CountedBlock:
         return self._block.block_id
 
 
+@pytest.mark.cpu
 def test_reading_the_window_touches_only_the_blocks_it_can_keep(monkeypatch):
     """However long the block table is, reading the window must not cost more.
 
@@ -1074,6 +1110,7 @@ def test_reading_the_window_touches_only_the_blocks_it_can_keep(monkeypatch):
     )
 
 
+@pytest.mark.cpu
 def test_the_checkpoints_own_default_resolution_can_build_a_cache():
     """832x480 is what LingBot World v2 ships as its default, and it could not run.
 
@@ -1111,6 +1148,7 @@ def test_the_checkpoints_own_default_resolution_can_build_a_cache():
     assert kv.blocks_per_frame == -(-tokens_per_frame // block_size) == 98
 
 
+@pytest.mark.cpu
 def test_the_shipped_resolution_geometry_matches_dense_attention():
     """The real number, not a scaled-down stand-in: 1560 tokens per frame.
 
@@ -1163,7 +1201,7 @@ def test_the_shipped_resolution_geometry_matches_dense_attention():
     torch.testing.assert_close(paged, ref, rtol=1e-5, atol=1e-5)
 
 
-@pytest.mark.skipif(not _gpu_flash_attn_usable(), reason="usable GPU FlashAttention is required")
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
 @pytest.mark.parametrize("commit_current", [False, True])
 def test_the_shipped_resolution_geometry_on_the_real_kernel(commit_current):
     """The shipped 832x480 geometry through FlashAttention's paged kernel.
@@ -1173,6 +1211,7 @@ def test_the_shipped_resolution_geometry_on_the_real_kernel(commit_current):
     tokens per frame, through the production path -- host prep followed by
     the fused write+attend op, which is what a forward actually calls.
     """
+    _require_gpu_flash_attn()
     torch.manual_seed(0)
     device = torch.device("cuda")
     dtype = torch.float16
@@ -1204,7 +1243,7 @@ def test_the_shipped_resolution_geometry_on_the_real_kernel(commit_current):
     torch.testing.assert_close(paged, ref, rtol=2e-2, atol=2e-2)
 
 
-@pytest.mark.skipif(not _gpu_flash_attn_usable(), reason="usable GPU FlashAttention is required")
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
 @pytest.mark.parametrize("history_chunks", [0, 1, 3])
 @pytest.mark.parametrize("action_len", [0, 3])
 def test_contiguous_kv_gather_path_matches_paged_path_gpu(monkeypatch, history_chunks, action_len):
@@ -1213,6 +1252,7 @@ def test_contiguous_kv_gather_path_matches_paged_path_gpu(monkeypatch, history_c
     Covers an empty, partial and full window plus a partially filled action block:
     the gather reads the tail-padding null block, so its (zeroed) rows must be masked.
     """
+    _require_gpu_flash_attn()
 
     torch.manual_seed(0)
     device = torch.device("cuda")

tests/diffusion/attention/test_trtllm_attn.py

@@ -617,3 +617,68 @@ def test_skip_end_to_end_config_path(monkeypatch):
     ctx = fc.ForwardContext(denoise_timestep=0.3)
     monkeypatch.setattr(fc, "_forward_context", ctx)
     assert impl._resolve_skip_factor(4096) == pytest.approx(expected)
+
+
+@requires_trtllm_attn
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.cuda.get_device_capability() not in ((10, 0), (10, 3)),
+    reason="verified TRTLLM contract requires SM100/SM103",
+)
+@pytest.mark.parametrize("batch", [1, 2])
+def test_dense_contract_fullgraph_matches_sdpa(batch):
+    from vllm_omni.diffusion.attention.capabilities import (
+        CompilationMode,
+        ExecutionContext,
+        SupportStatus,
+    )
+
+    torch.compiler.reset()
+    try:
+        impl = _impl()
+        context = ExecutionContext(platform="cuda", require_fullgraph=True)
+        compiled = torch.compile(impl.forward_cuda, fullgraph=True, dynamic=True)
+        for q_len, kv_len in ((64, 64), (128, 192), (256, 128)):
+            # Noncontiguous Q exercises the reshape before the opaque boundary.
+            q = torch.randn(batch, 8, q_len, 128, device="cuda", dtype=torch.bfloat16).transpose(1, 2)
+            k, v = (torch.randn(batch, kv_len, 8, 128, device="cuda", dtype=torch.bfloat16) for _ in range(2))
+            result = impl.resolve_execution_path(context, q, k, v, None)
+            assert result.requested_support(context).status is SupportStatus.SUPPORTED
+            assert result.compilation_mode is CompilationMode.CUSTOM_OP
+            before = [t.clone() for t in (q, k, v)]
+            eager = impl.forward_cuda(q, k, v)
+            for _ in range(2):
+                out = compiled(q, k, v)
+                assert out.dtype == q.dtype and out.device == q.device and out.is_contiguous()
+                torch.testing.assert_close(out, eager, atol=1e-2, rtol=1e-2)
+                torch.testing.assert_close(out.float(), _sdpa_ref(q, k, v, 128**-0.5), atol=1e-2, rtol=1e-2)
+            for actual, original in zip((q, k, v), before):
+                torch.testing.assert_close(actual, original, atol=0, rtol=0)
+        # Check the real kernel's mutation and output metadata against its schema/fake.
+        qf, kf, vf = (t.reshape(-1, 8, 128).contiguous() for t in (q, k, v))
+        args = (
+            qf,
+            kf,
+            vf,
+            impl._get_workspace(q.device),
+            torch.full((batch,), kv_len, device=q.device, dtype=torch.int32),
+            torch.arange(batch + 1, device=q.device, dtype=torch.int32) * q_len,
+            torch.arange(batch + 1, device=q.device, dtype=torch.int32) * kv_len,
+            None,
+            None,
+            None,
+            q_len,
+            kv_len,
+            batch,
+            128**-0.5,
+            1.0,
+            -1.0,
+            0,
+            0,
+            False,
+        )
+        checks = torch.library.opcheck(
+            tg._trtllm_ragged_attention_op, args, test_utils=("test_schema", "test_faketensor")
+        )
+        assert all(value == "SUCCESS" for value in checks.values())
+    finally:
+        torch.compiler.reset()

tests/diffusion/attention/test_trtllm_contract.py

@@ -0,0 +1,264 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import functools
+from dataclasses import replace
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from vllm_omni.diffusion.attention.backends import trtllm_attn as tg
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata, PackedPaddingMetadata
+from vllm_omni.diffusion.attention.capabilities import (
+    CompilationMode,
+    ExecutionContext,
+    OuterBoundary,
+    ParallelStrategy,
+    SupportStatus,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def _impl(**kwargs):
+    return tg.TrtllmAttentionImpl(2, 128, 128**-0.5, **kwargs)
+
+
+def _inputs(q_len=8, kv_len=12, dtype=torch.bfloat16):
+    return (
+        torch.randn(1, q_len, 2, 128, dtype=dtype),
+        torch.randn(1, kv_len, 2, 128, dtype=dtype),
+        torch.randn(1, kv_len, 2, 128, dtype=dtype),
+    )
+
+
+@pytest.fixture
+def dispatcher(monkeypatch, tmp_path):
+    marker = tmp_path / "kernel"
+    marker.write_text("loaded")
+
+    @functools.cache
+    def load_kernel():
+        return marker.read_text()
+
+    def fake_attention(query, key, value, workspace_buffer, batch_size, bmm1_scale, **kwargs):
+        assert load_kernel() == "loaded"
+        assert kwargs["skip_all_rows_active_check"] is True
+        workspace_buffer.add_(1)
+        q, k, v = (t.reshape(batch_size, -1, t.shape[-2], t.shape[-1]).transpose(1, 2) for t in (query, key, value))
+        return (
+            F.scaled_dot_product_attention(q, k, v, scale=bmm1_scale)
+            .transpose(1, 2)
+            .reshape(query.shape[0], query.shape[1], value.shape[-1])
+            .contiguous()
+        )
+
+    def quantize(q, k, v, **kwargs):
+        scale = torch.ones(1)
+        return q, k, v, scale, scale, scale, None
+
+    # CPU tests substitute device selection; the real kernel tests use the GPU.
+    monkeypatch.setattr(tg, "_is_trtllm_blackwell_device", lambda query: True)
+    monkeypatch.setattr(tg, "HAS_FLASHINFER", True)
+    monkeypatch.setattr(tg, "trtllm_ragged_attention_deepseek", fake_attention, raising=False)
+    monkeypatch.setattr(
+        tg, "flashinfer", SimpleNamespace(__version__="0.6.18", trtllm_sage_attention_quantize=quantize), raising=False
+    )
+    workspace = torch.zeros(16, dtype=torch.uint8)
+    monkeypatch.setattr(tg.TrtllmAttentionImpl, "_get_workspace", classmethod(lambda cls, device: workspace))
+    torch.compiler.reset()
+    yield workspace
+    torch.compiler.reset()
+
+
+def test_dense_contract_uses_initialized_state(dispatcher):
+    context = ExecutionContext(platform="cuda", kernel_variant="caller_sage", causal=True, require_fullgraph=True)
+    result = _impl().resolve_execution_path(context, *_inputs(), None)
+    assert result.path == "trtllm_dense" and result.kernel_variant == "trtllm"
+    assert result.support.status is SupportStatus.SUPPORTED
+    assert result.compilation_mode is CompilationMode.CUSTOM_OP
+    assert result.requested_support(context).status is SupportStatus.SUPPORTED
+    pre = tg.TrtllmAttentionBackend.resolve_capabilities(context)
+    assert pre.kernel_variant is None and pre.support.status is SupportStatus.UNMIGRATED
+    assert pre.requested_support(context).status is SupportStatus.UNSUPPORTED
+
+
+@pytest.mark.parametrize(
+    "change",
+    [
+        {"platform": "rocm"},
+        {"paged_kv": True},
+        {"parallel_strategy": ParallelStrategy.ULYSSES},
+        {"parallel_strategy": ParallelStrategy.RING},
+        {"parallel_strategy": ParallelStrategy.HYBRID_ULYSSES_RING},
+        {"parallel_strategy": ParallelStrategy.ALLGATHER_KV},
+        {"outer_boundaries": frozenset({OuterBoundary.HSDP})},
+    ],
+)
+def test_outer_paths_are_unmigrated(dispatcher, change):
+    context = replace(ExecutionContext(platform="cuda", require_fullgraph=True), **change)
+    result = _impl().resolve_execution_path(context, *_inputs(), None)
+    assert result.support.status is SupportStatus.UNMIGRATED
+    assert result.requested_support(context).status is SupportStatus.UNSUPPORTED
+
+
+@pytest.mark.parametrize(
+    "kwargs,path",
+    [
+        ({"causal": True}, "trtllm_dense"),
+        ({"backend_kwargs": {"quant": {"dtype_qk": "int8"}}}, "trtllm_sage"),
+        ({"backend_kwargs": {"skip_softmax_threshold": 0.01}}, "trtllm_skip_softmax"),
+        ({"backend_kwargs": {"target_sparsity": 0.5}}, "trtllm_skip_softmax"),
+    ],
+)
+def test_unverified_backend_settings(dispatcher, kwargs, path):
+    result = _impl(**kwargs).resolve_execution_path(ExecutionContext(platform="cuda"), *_inputs(), None)
+    assert result.path == path and result.support.status is SupportStatus.UNMIGRATED
+
+
+@pytest.mark.parametrize("padding", [False, True])
+def test_packed_paths_are_not_dense(dispatcher, padding):
+    q, k, v = _inputs()
+    cuq, cuk = torch.tensor([0, 8], dtype=torch.int32), torch.tensor([0, 12], dtype=torch.int32)
+    metadata = AttentionMetadata(
+        extra={
+            "cu_seqlens_q": cuq,
+            "cu_seqlens_k": cuk,
+            "max_seqlen_q": 8,
+            "max_seqlen_k": 12,
+        },
+        packed_padding=PackedPaddingMetadata(8, 12, cuq, cuk) if padding else None,
+    )
+    impl = _impl()
+    result = impl.resolve_execution_path(ExecutionContext(platform="cuda"), q, k, v, metadata)
+    assert result.path == ("trtllm_packed_padding" if padding else "trtllm_packed_varlen")
+    assert result.support.status is SupportStatus.UNMIGRATED
+    assert (
+        impl.resolve_execution_path(ExecutionContext(platform="cuda"), q, k, v, None).support.status
+        is SupportStatus.SUPPORTED
+    )
+
+
+@pytest.mark.parametrize(
+    "metadata,message",
+    [
+        (AttentionMetadata(attn_mask=torch.ones(1, 12)), "does not support attn_mask"),
+        (AttentionMetadata(extra={"cu_seqlens_q": torch.tensor([0, 8])}), "Incomplete packed"),
+        (AttentionMetadata(packed_padding="invalid"), "must be PackedPaddingMetadata"),
+    ],
+)
+def test_resolution_and_execution_share_metadata_validation(dispatcher, metadata, message):
+    impl = _impl()
+    inputs = _inputs()
+    with pytest.raises(ValueError, match=message):
+        impl.resolve_execution_path(ExecutionContext(platform="cuda"), *inputs, metadata)
+    with pytest.raises(ValueError, match=message):
+        impl.forward_cuda(*inputs, metadata)
+
+
+@pytest.mark.parametrize(
+    "metadata",
+    [
+        AttentionMetadata(full_attn_spans=[[(0, 4)]]),
+        AttentionMetadata(extra={"kv_cache_dtype": "fp8"}),
+    ],
+)
+def test_special_metadata_is_unmigrated(dispatcher, metadata):
+    result = _impl().resolve_execution_path(ExecutionContext(platform="cuda"), *_inputs(), metadata)
+    assert result.support.status is SupportStatus.UNMIGRATED
+
+
+def test_invalid_geometry_and_missing_dependency(dispatcher, monkeypatch):
+    q, k, v = _inputs()
+    context = ExecutionContext(platform="cuda")
+    for inputs in ((q[0], k, v), (q, k.float(), v), (q, k, v[:, :-1])):
+        result = _impl().resolve_execution_path(context, *inputs, None)
+        assert result.support.status is SupportStatus.UNSUPPORTED and result.support.reason
+    for inputs in (
+        _inputs(dtype=torch.float32),
+        (q[..., :64], k[..., :64], v[..., :64]),
+        (q, k[:, :, :1], v[:, :, :1]),
+    ):
+        assert _impl().resolve_execution_path(context, *inputs, None).support.status is SupportStatus.UNMIGRATED
+    monkeypatch.setattr(tg, "HAS_FLASHINFER", False)
+    result = _impl().resolve_execution_path(context, q, k, v, None)
+    assert result.support.status is SupportStatus.UNSUPPORTED and "FlashInfer" in result.support.reason
+
+
+@pytest.mark.parametrize("use_sage", [False, True], ids=["dense", "sage"])
+def test_trtllm_dispatcher_is_opaque_to_torch_compile(dispatcher, use_sage):
+    impl = _impl(backend_kwargs={"quant": {"dtype_qk": "int8"}} if use_sage else None)
+    compiled = torch.compile(impl.forward_cuda, fullgraph=True, dynamic=True)
+    for q_len, kv_len in ((8, 12), (11, 17), (8, 12)):
+        q, k, v = _inputs(q_len, kv_len)
+        before = [t.clone() for t in (q, k, v)]
+        eager = impl.forward_cuda(q, k, v)
+        out = compiled(q, k, v)
+        torch.testing.assert_close(out, eager)
+        for actual, original in zip((q, k, v), before):
+            torch.testing.assert_close(actual, original, rtol=0, atol=0)
+    assert torch.all(dispatcher == 6)
+
+
+def test_custom_op_schema_fake_and_noncontiguous_output(dispatcher):
+    q = torch.randn(2, 7, 128).transpose(0, 1)
+    k, v = torch.randn(9, 2, 128), torch.randn(9, 2, 64)
+    args = (
+        q,
+        k,
+        v,
+        dispatcher,
+        torch.tensor([9], dtype=torch.int32),
+        torch.tensor([0, 7], dtype=torch.int32),
+        torch.tensor([0, 9], dtype=torch.int32),
+        None,
+        None,
+        None,
+        7,
+        9,
+        1,
+        128**-0.5,
+        1.0,
+        -1.0,
+        0,
+        0,
+        False,
+    )
+    result = torch.library.opcheck(tg._trtllm_ragged_attention_op, args, test_utils=("test_schema", "test_faketensor"))
+    assert all(v == "SUCCESS" for v in result.values())
+    out = tg._trtllm_ragged_attention_op(*args)
+    assert out.shape == (7, 2, 64) and out.is_contiguous()
+
+
+@pytest.mark.parametrize(
+    "capability,expected",
+    [
+        ((9, 0), False),
+        ((10, 0), True),
+        ((10, 3), True),
+        ((10, 7), False),
+        ((12, 0), False),
+        (None, False),
+    ],
+)
+def test_trtllm_device_capability_is_conservative(monkeypatch, capability, expected):
+    monkeypatch.setattr(
+        tg,
+        "current_omni_platform",
+        SimpleNamespace(
+            is_cuda=lambda: True,
+            get_device_capability=lambda device_id: capability,
+        ),
+    )
+    assert tg._is_trtllm_blackwell_device(SimpleNamespace(device=torch.device("cuda:0"))) is expected
+    assert not tg._is_trtllm_blackwell_device(SimpleNamespace(device=torch.device("cpu")))
+
+
+def test_non_blackwell_cannot_claim_fullgraph(dispatcher, monkeypatch):
+    monkeypatch.setattr(tg, "_is_trtllm_blackwell_device", lambda query: False)
+    context = ExecutionContext(platform="cuda", require_fullgraph=True)
+    result = _impl().resolve_execution_path(context, *_inputs(), None)
+    assert result.support.status is SupportStatus.UNMIGRATED
+    assert result.requested_support(context).status is SupportStatus.UNSUPPORTED

tests/diffusion/cache/test_cache_backends.py

@@ -571,6 +571,57 @@ class TestTeaCacheBackend:
         backend.refresh(mock_pipeline, num_inference_steps=50)
         mock_registry.reset_hook.assert_called_once()
 
+    @patch("vllm_omni.diffusion.cache.teacache.backend.apply_teacache_hook")
+    def test_enable_mammoth_moda2_uses_gen_transformer(self, mock_apply_hook):
+        """MammothModa2 exposes its DiT as gen_transformer, not transformer."""
+        mock_pipeline = Mock()
+        mock_pipeline.__class__.__name__ = "MammothModa2DiTPipeline"
+        mock_pipeline.gen_transformer = Mock()
+        del mock_pipeline.transformer
+
+        backend = TeaCacheBackend(DiffusionCacheConfig())
+        backend.enable(mock_pipeline)
+
+        assert backend.enabled is True
+        mock_apply_hook.assert_called_once()
+        transformer_arg, config_arg = mock_apply_hook.call_args.args
+        assert transformer_arg is mock_pipeline.gen_transformer
+        assert config_arg.transformer_type == "MammothModa2Transformer2DModel"
+        assert config_arg.rel_l1_thresh == 0.075
+        assert mock_pipeline.transformer is mock_pipeline.gen_transformer
+
+    @patch("vllm_omni.diffusion.cache.teacache.backend.apply_teacache_hook")
+    def test_enable_mammoth_moda2_preserves_explicit_threshold(self, mock_apply_hook):
+        pipeline = Mock()
+        pipeline.__class__.__name__ = "MammothModa2DiTPipeline"
+        pipeline.gen_transformer = Mock()
+        del pipeline.transformer
+
+        TeaCacheBackend(DiffusionCacheConfig(rel_l1_thresh=0.05)).enable(pipeline)
+
+        assert mock_apply_hook.call_args.args[1].rel_l1_thresh == 0.05
+
+    @patch("vllm_omni.diffusion.cache.teacache.backend.apply_teacache_hook")
+    def test_refresh_mammoth_moda2_resets_gen_transformer_hook(self, mock_apply_hook):
+        """Refresh should find the hook through the transformer alias set by the enabler."""
+        mock_pipeline = Mock()
+        mock_pipeline.__class__.__name__ = "MammothModa2DiTPipeline"
+        mock_pipeline.gen_transformer = Mock()
+        del mock_pipeline.transformer
+
+        backend = TeaCacheBackend(DiffusionCacheConfig())
+        backend.enable(mock_pipeline)
+
+        mock_hook = Mock()
+        mock_registry = Mock()
+        mock_registry.get_hook = Mock(return_value=mock_hook)
+        mock_registry.reset_hook = Mock()
+        mock_pipeline.gen_transformer._hook_registry = mock_registry
+
+        backend.refresh(mock_pipeline, num_inference_steps=12)
+
+        mock_registry.reset_hook.assert_called_once_with("teacache")
+
 
 class TestCacheSelector:
     """Test cache backend selector function."""

tests/diffusion/cache/test_seacache.py

@@ -329,7 +329,9 @@ def _uneven_cfg_sharded_worker(
         init_method=init_method,
         rank=rank,
         world_size=4,
-        timeout=datetime.timedelta(seconds=10),
+        # Spawned workers import torch/vLLM before joining; allow startup
+        # skew without changing the collective deadlock checks below.
+        timeout=datetime.timedelta(seconds=60),
     )
     try:
         fs_groups = [

tests/diffusion/cache/test_teacache_extractors.py

@@ -28,17 +28,28 @@ import torch
 import torch.nn as nn
 
 from tests.helpers.mark import hardware_test
+from vllm_omni.diffusion.cache.teacache.backend import TeaCacheBackend
+from vllm_omni.diffusion.cache.teacache.config import TeaCacheConfig
 from vllm_omni.diffusion.cache.teacache.extractors import (
     extract_flux2_context,
     extract_flux2_klein_context,
     extract_flux_context,
+    extract_mammoth_moda2_context,
     extract_minimax_h3_context,
     extract_zimage_context,
 )
+from vllm_omni.diffusion.cache.teacache.hook import apply_teacache_hook
+from vllm_omni.diffusion.config import set_current_diffusion_config
+from vllm_omni.diffusion.data import AttentionConfig, DiffusionCacheConfig
 from vllm_omni.diffusion.models.flux.flux_transformer import FluxTransformer2DModel
 from vllm_omni.diffusion.models.flux2_klein.flux2_klein_transformer import (
     Flux2Transformer2DModel,
 )
+from vllm_omni.diffusion.models.mammoth_moda2.mammothmoda2_dit_model import (
+    Transformer2DModel as MammothModa2Transformer2DModel,
+)
+from vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit import MammothModa2DiTPipeline
+from vllm_omni.diffusion.models.mammoth_moda2.rope_real import RotaryPosEmbedReal
 
 pytestmark = [pytest.mark.core_model]
 
@@ -135,6 +146,164 @@ class BaseExtractorTest(ABC):
         pass
 
 
+@pytest.mark.cpu
+class TestMammothModa2Extractor(BaseExtractorTest):
+    """Test extract_mammoth_moda2_context function."""
+
+    def get_extractor(self):
+        return extract_mammoth_moda2_context
+
+    @pytest.fixture
+    def mammoth_module(self):
+        # These tests run on CPU even when the host has CUDA available.
+        od_config = SimpleNamespace(
+            diffusion_attention_config=AttentionConfig(default="TORCH_SDPA"),
+            parallel_config=SimpleNamespace(ring_degree=1),
+        )
+        with set_current_diffusion_config(od_config):
+            yield MammothModa2Transformer2DModel(
+                patch_size=2,
+                in_channels=4,
+                hidden_size=16,
+                num_layers=2,
+                num_refiner_layers=1,
+                num_attention_heads=2,
+                num_kv_heads=1,
+                axes_dim_rope=(2, 2, 4),
+                axes_lens=(16, 16, 16),
+                text_feat_dim=8,
+            )
+
+    def get_module(self, mammoth_module):
+        return mammoth_module
+
+    @pytest.fixture
+    def sample_inputs(self):
+        return {
+            "hidden_states": torch.randn(1, 4, 4, 4),
+            "timestep": torch.tensor([500.0]),
+            "text_hidden_states": torch.randn(1, 3, 8),
+            "text_attention_mask": torch.ones(1, 3, dtype=torch.bool),
+            "freqs_cis": RotaryPosEmbedReal.get_freqs_real((2, 2, 4), (16, 16, 16), theta=10000),
+        }
+
+    def get_sample_inputs(self, sample_inputs):
+        return sample_inputs
+
+    def test_context_shapes_and_postprocess(self, mammoth_module, sample_inputs):
+        context = extract_mammoth_moda2_context(mammoth_module, **sample_inputs)
+
+        assert context.modulated_input.shape == context.hidden_states.shape
+        assert context.encoder_hidden_states is None
+        assert context.temb.shape[0] == 1
+
+        outputs = context.run_transformer_blocks()
+        assert len(outputs) == 1
+        assert outputs[0].shape == context.hidden_states.shape
+
+        output = context.postprocess(context.hidden_states)
+        assert output.shape == sample_inputs["hidden_states"].shape
+
+    def test_branch_hint_is_preserved(self, mammoth_module, sample_inputs):
+        context = extract_mammoth_moda2_context(
+            mammoth_module,
+            **sample_inputs,
+            teacache_branch="positive",
+        )
+
+        assert context.extra_states is not None
+        assert context.extra_states["teacache_branch"] == "positive"
+
+    def test_hook_keeps_positive_and_negative_branch_state_separate(self, mammoth_module, sample_inputs, monkeypatch):
+        config = TeaCacheConfig(
+            transformer_type="MammothModa2Transformer2DModel",
+            coefficients=[0, 0, 0, 0, 0],
+            rel_l1_thresh=999.0,
+        )
+        apply_teacache_hook(mammoth_module, config)
+
+        original_apply_layers = mammoth_module._apply_transformer_layers
+        apply_layers_mock = Mock(wraps=original_apply_layers)
+        monkeypatch.setattr(mammoth_module, "_apply_transformer_layers", apply_layers_mock)
+
+        mammoth_module(**sample_inputs, teacache_branch="positive")
+        assert apply_layers_mock.call_count == 1
+
+        mammoth_module(**sample_inputs, teacache_branch="positive")
+        assert apply_layers_mock.call_count == 1
+
+        mammoth_module(**sample_inputs, teacache_branch="negative")
+        assert apply_layers_mock.call_count == 2
+
+    def test_hook_rejects_invalid_branch_hint(self, mammoth_module, sample_inputs):
+        config = TeaCacheConfig(
+            transformer_type="MammothModa2Transformer2DModel",
+            coefficients=[0, 0, 0, 0, 0],
+            rel_l1_thresh=999.0,
+        )
+        apply_teacache_hook(mammoth_module, config)
+
+        with pytest.raises(ValueError, match="teacache_branch"):
+            mammoth_module(**sample_inputs, teacache_branch="bad")
+
+    @pytest.mark.parametrize("change_shape", [False, True], ids=["same_shape", "different_shape"])
+    @torch.no_grad()
+    def test_backend_refresh_clears_both_branches(self, mammoth_module, sample_inputs, monkeypatch, change_shape):
+        pipeline = MammothModa2DiTPipeline.__new__(MammothModa2DiTPipeline)
+        nn.Module.__init__(pipeline)
+        pipeline.gen_transformer = mammoth_module
+        # Zero coefficients guarantee cache hits after each branch's first call.
+        backend = TeaCacheBackend(DiffusionCacheConfig(coefficients=[0, 0, 0, 0, 0]))
+        backend.enable(pipeline)
+        apply_layers = Mock(wraps=mammoth_module._apply_transformer_layers)
+        monkeypatch.setattr(mammoth_module, "_apply_transformer_layers", apply_layers)
+
+        for branch in ("positive", "negative"):
+            mammoth_module(**sample_inputs, teacache_branch=branch)
+            mammoth_module(**sample_inputs, teacache_branch=branch)
+        assert apply_layers.call_count == 2
+
+        next_inputs = dict(sample_inputs)
+        next_inputs["hidden_states"] = torch.randn(1, 4, 4, 6 if change_shape else 4)
+        text_length = 5 if change_shape else 3
+        next_inputs["text_hidden_states"] = torch.randn(1, text_length, 8)
+        next_inputs["text_attention_mask"] = torch.ones(1, text_length, dtype=torch.bool)
+        # Compute an uncached reference through the original model forward.
+        expected = MammothModa2Transformer2DModel.forward(mammoth_module, **next_inputs)
+        apply_layers.reset_mock()
+
+        backend.refresh(pipeline, num_inference_steps=4)
+
+        for count, branch in enumerate(("positive", "negative"), start=1):
+            actual = mammoth_module(**next_inputs, teacache_branch=branch)
+            assert apply_layers.call_count == count
+            torch.testing.assert_close(actual, expected)
+            # Refresh must preserve caching for subsequent steps of request B.
+            cached = mammoth_module(**next_inputs, teacache_branch=branch)
+            assert apply_layers.call_count == count
+            torch.testing.assert_close(cached, expected)
+
+    def test_forced_full_compute_matches_original_forward(self, mammoth_module, sample_inputs, monkeypatch):
+        with torch.no_grad():
+            expected = mammoth_module(**sample_inputs, teacache_branch="positive")
+
+        config = TeaCacheConfig(
+            transformer_type="MammothModa2Transformer2DModel",
+            coefficients=[0, 0, 0, 0, 0],
+            rel_l1_thresh=0.1,
+        )
+        apply_teacache_hook(mammoth_module, config)
+        hook = mammoth_module._hook_registry.get_hook("teacache")
+        should_compute_full = Mock(return_value=True)
+        monkeypatch.setattr(hook, "_should_compute_full_transformer", should_compute_full)
+
+        with torch.no_grad():
+            actual = mammoth_module(**sample_inputs, teacache_branch="positive")
+
+        should_compute_full.assert_called_once()
+        torch.testing.assert_close(actual, expected)
+
+
 class TestFlux2KleinExtractor(BaseExtractorTest):
     """Test extract_flux2_klein_context function."""
 

tests/diffusion/diffusion_kv/test_kv_cache_utils.py

@@ -103,7 +103,7 @@ def test_block_hashing_uses_upstream_extra_keys_directly(monkeypatch):
     # Independent expected native extras: a crossing item has a negative
     # relative offset; a block can include the end of one item and the next.
     first, second = (feature.identifier for feature in row.mm_features)
-    extras = [((first, 2),), ((first, -2), (second, 3)), ((second, -1),)]
+    extras = [(("mm", first, 2),), (("mm", first, -2), ("mm", second, 3)), (("mm", second, -1),)]
     parent = BlockHash(hash_fn(("diffusion-prefix-root-v1", row.cache_namespace)))
     expected = []
     for start, keys in zip(range(0, 12, 4), extras, strict=True):

tests/diffusion/diffusion_kv/test_mooncake_cfg_fanout.py

@@ -31,6 +31,10 @@ def producer(monkeypatch):
     # vLLM 0.30 send_kv_to_decode also reads the PP size (ce08bb5b34,
     # #56033), the MLA/Mamba flags and the head-resharding validator.
     worker.pp_size = 1
+    worker.opaque_packed_storages = set()
+    worker._prepared_transfer_regions = {}
+    worker.region_shared_groups = []
+    worker.region_row_offsets = []
     worker.use_mla = False
     worker.kv_cache_config = SimpleNamespace(has_mamba_layers=False)
     worker._validate_head_resharding_layout = Mock(return_value=None)

tests/diffusion/diffusion_kv/test_paged_attention_adapter.py

@@ -1253,14 +1253,15 @@ def test_layer_adapter_accepts_platform_native_backend_and_uses_rank_local_heads
         speculative_config=None,
     )
 
-    def select_backend(*, backend, attn_selector_config, num_heads):
-        # Exercise the real 0.29 selector, replacing only platform resolution.
+    def select_backend(*, backend, attn_selector_config, num_heads, _run_kv_cache_dtype):
+        # Exercise the real upstream selector, replacing only platform resolution.
         # Ulysses has already gathered tokens and sharded heads before the
         # native kernel; the MoE PCP mapping must not request PCP attention.
         assert not attn_selector_config.use_pcp
         assert not attn_selector_config.use_dcp
         assert attn_selector_config.use_non_causal
         assert num_heads == 4
+        assert _run_kv_cache_dtype == config.cache_config.cache_dtype
         selected_backends.append((backend, config.attention_config.backend_per_kind))
         if selection_fails:
             raise ValueError("test backend unavailable")

tests/diffusion/distributed/test_comm.py

@@ -14,6 +14,7 @@ from typing import Literal
 import pytest
 import torch
 
+from tests.helpers.distributed import configure_cpu_collective_worker, start_collective_workers
 from tests.helpers.mark import hardware_marks
 from vllm_omni.diffusion.distributed.comm import RingComm, SeqAllToAll4D, SeqAllToAll5D
 from vllm_omni.diffusion.distributed.parallel_state import (
@@ -56,6 +57,8 @@ def _init_worker(
             "MASTER_PORT": str(master_port),
         }
     )
+    if device_kind == "cpu":
+        configure_cpu_collective_worker()
     backend = "gloo" if device_kind == "cpu" else None
     init_distributed_environment(backend=backend)
 
@@ -227,6 +230,10 @@ def _run_ring_p2p(
             atol=atol,
             msg=f"[Rank {local_rank}] Ring P2P data mismatch",
         )
+        if device_kind == "cpu":
+            # P2P completion only synchronizes neighbors. Keep fast CPU ranks
+            # alive until every rank has finished constructing and using groups.
+            torch.distributed.barrier()
     finally:
         destroy_distributed_env()
 
@@ -243,7 +250,7 @@ def _spawn_4d_identity(
     device_kind: DeviceKind,
     master_port: int,
 ) -> None:
-    torch.multiprocessing.spawn(
+    start_collective_workers(
         _run_4d_identity,
         args=(
             world_size,
@@ -257,6 +264,7 @@ def _spawn_4d_identity(
             master_port,
         ),
         nprocs=world_size,
+        device_kind=device_kind,
     )
 
 
@@ -272,7 +280,7 @@ def _spawn_5d_identity(
     device_kind: DeviceKind,
     master_port: int,
 ) -> None:
-    torch.multiprocessing.spawn(
+    start_collective_workers(
         _run_5d_identity,
         args=(
             world_size,
@@ -286,6 +294,7 @@ def _spawn_5d_identity(
             master_port,
         ),
         nprocs=world_size,
+        device_kind=device_kind,
     )
 
 
@@ -299,10 +308,11 @@ def _spawn_ring_p2p(
     device_kind: DeviceKind,
     master_port: int,
 ) -> None:
-    torch.multiprocessing.spawn(
+    start_collective_workers(
         _run_ring_p2p,
         args=(world_size, dtype, batch_size, num_heads, head_size, device_kind, master_port),
         nprocs=world_size,
+        device_kind=device_kind,
     )
 
 
@@ -499,8 +509,9 @@ def _run_fused_qkv_matches_three_4d(
 @pytest.mark.cpu
 @pytest.mark.parametrize("world_size", [2, 4])
 def test_fused_qkv_all_to_all_matches_three_4d_exchanges_cpu(world_size: int) -> None:
-    torch.multiprocessing.spawn(
+    start_collective_workers(
         _run_fused_qkv_matches_three_4d,
         args=(world_size, "cpu", 29613),
         nprocs=world_size,
+        device_kind="cpu",
     )

tests/diffusion/distributed/test_pipeline_parallel.py

@@ -19,6 +19,7 @@ from vllm.utils.network_utils import get_file_store_init_method
 from vllm.v1.worker.gpu_worker import AsyncIntermediateTensors
 
 import vllm_omni.diffusion.distributed.pipeline_parallel as pp_module
+from tests.helpers.distributed import collective_context, start_collective_workers
 from tests.helpers.mark import hardware_marks
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.parallel_state import (
@@ -39,6 +40,25 @@ DeviceKind = Literal["cpu", "cuda"]
 _UNIT_MARKS = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 
 
+@pytest.fixture(autouse=True)
+def _restore_collective_platforms():
+    """CPU baselines run in the parent process; retain its original platforms."""
+    import vllm.platforms
+
+    from vllm_omni.diffusion.distributed import comm, group_coordinator, parallel_state
+
+    original_platform = vllm.platforms.current_platform
+    original_omni_platforms = {
+        module: module.current_omni_platform for module in (comm, group_coordinator, parallel_state)
+    }
+    try:
+        yield
+    finally:
+        vllm.platforms.current_platform = original_platform
+        for module, platform in original_omni_platforms.items():
+            module.current_omni_platform = platform
+
+
 def _cleanup_distributed() -> None:
     """Destroy omni distributed state left dangling after a test.
 
@@ -409,6 +429,10 @@ def init_dist(
     device = _worker_device(local_rank, device_kind)
     if device_kind == "cuda":
         current_omni_platform.set_device(device)
+    else:
+        from tests.helpers.distributed import configure_cpu_collective_worker
+
+        configure_cpu_collective_worker()
     backend = "gloo" if device_kind == "cpu" else None
     init_distributed_environment(
         world_size=world_size,
@@ -500,13 +524,14 @@ def isend_irecv_worker(
 
 
 def _run_isend_irecv(pp_size: int, device_kind: DeviceKind, init_method: str) -> None:
-    mp_context = torch.multiprocessing.get_context("spawn")
+    mp_context = collective_context(device_kind)
     manager = mp_context.Manager()
     q = manager.Queue()
-    torch.multiprocessing.spawn(
+    start_collective_workers(
         isend_irecv_worker,
         args=(pp_size, init_method, device_kind, q),
         nprocs=pp_size,
+        device_kind=device_kind,
     )
     results = {label: tensor for label, tensor in [q.get(), q.get()]}
     torch.testing.assert_close(
@@ -656,15 +681,16 @@ def _run_predict_noise(
 
     baseline_out = compute_single_gpu_baseline(test_config, dtype, do_true_cfg, device_kind)
 
-    mp_context = torch.multiprocessing.get_context("spawn")
+    mp_context = collective_context(device_kind)
     manager = mp_context.Manager()
     pp_q = manager.Queue()
 
     world_size = pp_size * cfg_size
-    torch.multiprocessing.spawn(
+    start_collective_workers(
         predict_noise_worker,
         args=(world_size, init_method, pp_size, cfg_size, do_true_cfg, dtype, test_config, device_kind, pp_q),
         nprocs=world_size,
+        device_kind=device_kind,
     )
 
     pp_out = pp_q.get()
@@ -900,15 +926,16 @@ def _run_scheduler_step(
 
     baseline = compute_scheduler_step_baseline(test_config, do_true_cfg, device_kind)
 
-    mp_context = torch.multiprocessing.get_context("spawn")
+    mp_context = collective_context(device_kind)
     manager = mp_context.Manager()
     q = manager.Queue()
 
     world_size = pp_size * cfg_size
-    torch.multiprocessing.spawn(
+    start_collective_workers(
         scheduler_step_worker,
         args=(world_size, init_method, pp_size, cfg_size, do_true_cfg, test_config, device_kind, q),
         nprocs=world_size,
+        device_kind=device_kind,
     )
 
     result = q.get()

tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py

@@ -239,11 +239,11 @@ def test_constructor_rejects_unsupported_mllm_quantization(mock_dependencies, mo
 )
 def test_constructor_routes_mllm_quantization(mock_dependencies, quantization_config, quantize_mllm):
     from transformers import AutoModel
-    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
     from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig
     from vllm.model_executor.layers.quantization.utils.quant_utils import kFp8Static128BlockSym
 
     from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import BooguImagePipeline
+    from vllm_omni.quantization.fp8_config import DiffusionFp8Config
 
     od_config = OmniDiffusionConfig(
         model="dummy-boogu",
@@ -287,15 +287,14 @@ def test_constructor_routes_mllm_quantization(mock_dependencies, quantization_co
         mock_dependencies["mllm_config_loader"].assert_not_called()
         mock_dependencies["mllm_builder"].assert_not_called()
         assert [source.subfolder for source in pipeline.weights_sources] == ["transformer"]
-    assert isinstance(mock_dependencies["transformer_cls"].call_args.kwargs["quant_config"], Fp8Config)
+    assert isinstance(mock_dependencies["transformer_cls"].call_args.kwargs["quant_config"], DiffusionFp8Config)
 
 
 def test_constructor_preserves_mllm_ignored_layers(mock_dependencies):
-    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
-
     from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import BooguImagePipeline
+    from vllm_omni.quantization.fp8_config import DiffusionFp8Config
 
-    quant_config = Fp8Config(
+    quant_config = DiffusionFp8Config(
         ignored_layers=["mllm.language_model.layers.0.self_attn.q_proj", "transformer.blocks.0.attn.to_q"]
     )
     od_config = OmniDiffusionConfig(

tests/diffusion/models/flux2/test_flux2_quantization.py

@@ -16,17 +16,13 @@ class _RecordingBlock(nn.Module):
 
 
 def test_flux2_propagates_online_quant_config_to_all_transformer_blocks(monkeypatch):
-    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
-
     from vllm_omni.diffusion.models.flux2 import flux2_transformer
+    from vllm_omni.quantization import build_quant_config
 
     monkeypatch.setattr(flux2_transformer, "Flux2TransformerBlock", _RecordingBlock)
     monkeypatch.setattr(flux2_transformer, "Flux2SingleTransformerBlock", _RecordingBlock)
 
-    quant_config = Fp8Config(
-        is_checkpoint_fp8_serialized=False,
-        activation_scheme="dynamic",
-    )
+    quant_config = build_quant_config("fp8")
     model = flux2_transformer.Flux2Transformer2DModel(
         in_channels=4,
         out_channels=4,

tests/diffusion/models/mammoth_moda2/test_dit_attention.py

@@ -187,3 +187,69 @@ def test_empty_text_stream_skips_the_kernel(monkeypatch, with_mask):
         image_rotary_emb=(angles.cos(), angles.sin()),
     )
     assert out.shape == (BATCH, 0, DIM)
+
+
+def test_model_owns_repeated_pair_rope_contract():
+    from vllm_omni.diffusion.models.mammoth_moda2.mammothmoda2_dit_model import Transformer2DModel
+    from vllm_omni.diffusion.models.mammoth_moda2.rope_real import RotaryPosEmbedReal
+
+    with set_current_diffusion_config(_SDPA_CONFIG):
+        model = Transformer2DModel(
+            hidden_size=48,
+            num_attention_heads=6,
+            num_kv_heads=2,
+            num_layers=1,
+            num_refiner_layers=1,
+            multiple_of=8,
+            ffn_dim_multiplier=1.0,
+            axes_dim_rope=(2, 2, 4),
+            axes_lens=(8, 8, 8),
+            text_feat_dim=32,
+        )
+    for blocks in (model.layers, model.noise_refiner, model.ref_image_refiner, model.context_refiner):
+        assert all(block.attn.processor.rope_repeats_pairs for block in blocks)
+    assert not _block(2).attn.processor.rope_repeats_pairs
+    freqs = RotaryPosEmbedReal.get_freqs_real((2, 2, 4), (8, 8, 8), 10000)
+    ids = torch.tensor([[[0, 1, 2], [3, 4, 5]], [[1, 0, 3], [2, 5, 4]]])
+    cos, sin = model.rope_embedder._get_freqs_real(freqs, ids)
+    for table in (cos, sin):
+        torch.testing.assert_close(table[..., 0::2], table[..., 1::2], atol=0, rtol=0)
+    assert not torch.equal(cos[0], cos[1])
+
+
+def test_embedding_preparation_preserves_stubbed_rope_tables(monkeypatch):
+    from types import SimpleNamespace
+
+    from torch import nn
+
+    import vllm_omni.diffusion.models.mammoth_moda2.mammothmoda2_dit_model as mammoth_dit
+
+    # Embedding-only tests can stub the RoPE producer without constructing
+    # attention layers or supplying full transformer configuration.
+    model = object.__new__(mammoth_dit.Transformer2DModel)
+    nn.Module.__init__(model)
+    object.__setattr__(model, "_internal_dict", SimpleNamespace(patch_size=2))
+
+    class CaptionEmbedding(nn.Module):
+        def forward(self, timestep, text, dtype):
+            return torch.zeros(2, 8), text
+
+    model.time_caption_embed = CaptionEmbedding()
+    model.x_embedder = nn.Identity()
+    model.rope_embedder = lambda *args: (None, None, None, None, [4, 4], [8, 8])
+
+    def forbid_gate(*args):
+        raise AssertionError("Embedding preparation must not resolve the fusion token gate")
+
+    monkeypatch.setattr(mammoth_dit, "fused_qk_norm_rope_min_tokens", forbid_gate)
+    prepared = model._prepare_embeddings(
+        torch.zeros(2, 4, 4, 4),
+        torch.ones(2),
+        torch.zeros(2, 4, 8),
+        torch.ones(2, 4, dtype=torch.bool),
+        torch.zeros(1),
+        2,
+        4,
+        4,
+    )
+    assert prepared[6:9] == (None, None, None)

tests/diffusion/models/mammoth_moda2/test_dit_attention_cuda.py

@@ -1,14 +1,16 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-"""Real-shape check of the MammothModa2 DiT attention on CUDA: the shared
-backend (FLASH_ATTN by default, head_dim 120) against the pre-change arithmetic
-in bf16, with the padding mask the joint text+image sequence carries."""
+"""Real-shape checks for MammothModa2's fused QK norm/RoPE and shared CUDA
+attention backend against the native shared RMSNorm path."""
 
 import pytest
 import torch
+from diffusers.models.attention_processor import Attention
 
 import vllm_omni.diffusion.attention.backends.sdpa as sdpa_backend
+import vllm_omni.diffusion.models.mammoth_moda2.mammothmoda2_dit_model as mammoth_dit
 from tests.helpers.mark import hardware_marks
+from vllm_omni.diffusion.layers.norm import RMSNorm
 from vllm_omni.diffusion.models.mammoth_moda2.mammothmoda2_dit_model import TransformerBlock
 
 from .test_dit_attention import _reference_attention
@@ -24,11 +26,127 @@ pytestmark = [
 DIM, HEADS, KV_HEADS = 2520, 21, 7
 
 
+def _norm_attention(head_dim, dtype):
+    attn = Attention(query_dim=head_dim * HEADS, heads=HEADS, kv_heads=KV_HEADS, dim_head=head_dim)
+    attn.norm_q = RMSNorm(head_dim)
+    attn.norm_k = RMSNorm(head_dim)
+    return attn.to(device="cuda", dtype=dtype)
+
+
+def test_unpaired_rope_preserves_native_arithmetic(monkeypatch):
+    """A generic block accepts independently valued even and odd RoPE lanes."""
+    torch.manual_seed(42)
+    q = torch.randn(2, 17, HEADS, 120, device="cuda", dtype=torch.bfloat16)
+    k = torch.randn(2, 17, KV_HEADS, 120, device="cuda", dtype=torch.bfloat16)
+    attn = _norm_attention(120, torch.bfloat16)
+    angles = torch.rand(2, 17, 120, device="cuda") * 6.283
+    rotary = angles.cos().to(q.dtype), angles.sin().to(q.dtype)
+    monkeypatch.setenv("VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS", "0")
+    with torch.no_grad():
+        got = mammoth_dit._apply_qk_norm_rope(attn, q, k, rotary)
+        expected = (
+            mammoth_dit.apply_real_rotary_emb(attn.norm_q(q), *rotary),
+            mammoth_dit.apply_real_rotary_emb(attn.norm_k(k), *rotary),
+        )
+    for actual, want in zip(got, expected):
+        torch.testing.assert_close(actual, want, atol=0, rtol=0)
+
+
+def test_qk_norm_rope_benchmark_measures_inference(monkeypatch):
+    from benchmarks.diffusion import benchmark_mammoth_moda2_qk_norm_rope as benchmark
+
+    calls = []
+
+    def measure(fn, warmup, iters):
+        outputs = fn()
+        assert not torch.is_grad_enabled()
+        assert all(not output.requires_grad for output in outputs)
+        calls.append(fn.__name__)
+        return {"median_ms": 1.0}
+
+    monkeypatch.setattr(benchmark, "_measure", measure)
+    benchmark._run_shape(17, warmup=0, iters=1)
+    assert calls == ["native", "fused"]
+
+
+def _production_rotary(batch, seq, head_dim, dtype):
+    """Mammoth's real RoPE repeats one angle across each adjacent pair."""
+    angles = torch.rand(batch, seq, head_dim // 2, device="cuda") * 6.283
+    return (
+        angles.cos().repeat_interleave(2, dim=-1).to(dtype),
+        angles.sin().repeat_interleave(2, dim=-1).to(dtype),
+    )
+
+
+@pytest.mark.parametrize("force_native", [False, True])
+def test_production_rope_uses_fused_qk_norm_rope(monkeypatch, force_native):
+    torch.manual_seed(0)
+    seq = 512
+    block = (
+        TransformerBlock(
+            DIM,
+            HEADS,
+            KV_HEADS,
+            multiple_of=256,
+            ffn_dim_multiplier=1.0,
+            norm_eps=1e-5,
+            rope_repeats_pairs=True,
+        )
+        .cuda()
+        .to(torch.bfloat16)
+        .eval()
+    )
+    hidden = torch.randn(2, seq, DIM, device="cuda", dtype=torch.bfloat16)
+    mask = torch.ones(2, seq, dtype=torch.bool, device="cuda")
+    mask[0, -37:] = False
+    rotary = _production_rotary(1, seq, block.head_dim, torch.bfloat16)
+    calls = []
+    fused_qk_norm_rope = mammoth_dit.fused_qk_norm_rope
+
+    def record_fused(q, k, q_weight, k_weight, rope_table, eps, **kwargs):
+        calls.append((q.shape, k.shape, rope_table.shape, kwargs))
+        return fused_qk_norm_rope(q, k, q_weight, k_weight, rope_table, eps, **kwargs)
+
+    monkeypatch.setenv("VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS", "0")
+    if force_native:
+        monkeypatch.setattr(mammoth_dit, "_fused_cuda_supported", lambda *args, **kwargs: False)
+    monkeypatch.setattr(mammoth_dit, "fused_qk_norm_rope", record_fused)
+    with torch.no_grad():
+        got = block.attn(hidden, hidden, attention_mask=mask, image_rotary_emb=rotary)
+        want = _reference_attention(block.attn, hidden, mask, rotary)
+
+    supported = mammoth_dit._fused_cuda_supported(hidden, hidden, block.head_dim, block.head_dim, interleaved=True)
+    expected_calls = (
+        [
+            (
+                torch.Size((2 * seq, HEADS, block.head_dim)),
+                torch.Size((2 * seq, KV_HEADS, block.head_dim)),
+                torch.Size((2 * seq, block.head_dim)),
+                {"head_dim": block.head_dim, "rotary_dim": block.head_dim, "interleaved": True},
+            )
+        ]
+        if supported
+        else []
+    )
+    assert calls == expected_calls
+    diff = (got.float() - want.float()).abs()[mask]
+    assert diff.max().item() < 2e-2, diff.max().item()
+    assert diff.mean().item() < 1e-3, diff.mean().item()
+
+
 @pytest.mark.parametrize("seq", [77 + 4096, 512], ids=["t2i_1024", "short"])
 def test_real_shape_matches_previous_arithmetic_bf16(seq):
     torch.manual_seed(0)
     block = (
-        TransformerBlock(DIM, HEADS, KV_HEADS, multiple_of=256, ffn_dim_multiplier=1.0, norm_eps=1e-5)
+        TransformerBlock(
+            DIM,
+            HEADS,
+            KV_HEADS,
+            multiple_of=256,
+            ffn_dim_multiplier=1.0,
+            norm_eps=1e-5,
+            rope_repeats_pairs=True,
+        )
         .cuda()
         .to(torch.bfloat16)
         .eval()
@@ -37,8 +155,7 @@ def test_real_shape_matches_previous_arithmetic_bf16(seq):
     mask = torch.ones(2, seq, dtype=torch.bool, device="cuda")
     mask[0, seq - 300 :] = False
     mask[1, seq - 17 :] = False
-    angles = torch.rand(1, seq, block.head_dim, device="cuda") * 6.283
-    rotary = (angles.cos().to(torch.bfloat16), angles.sin().to(torch.bfloat16))
+    rotary = _production_rotary(1, seq, block.head_dim, torch.bfloat16)
     with torch.no_grad():
         got = block.attn(
             hidden_states=hidden, encoder_hidden_states=hidden, attention_mask=mask, image_rotary_emb=rotary
@@ -88,7 +205,15 @@ def test_fp32_falls_back_to_sdpa_and_matches_reference(monkeypatch, seq, force_g
     torch.manual_seed(0)
     dt = torch.float32
     block = (
-        TransformerBlock(DIM, HEADS, KV_HEADS, multiple_of=256, ffn_dim_multiplier=1.0, norm_eps=1e-5)
+        TransformerBlock(
+            DIM,
+            HEADS,
+            KV_HEADS,
+            multiple_of=256,
+            ffn_dim_multiplier=1.0,
+            norm_eps=1e-5,
+            rope_repeats_pairs=True,
+        )
         .cuda()
         .to(dt)
         .eval()
@@ -122,7 +247,11 @@ def test_fp32_falls_back_to_sdpa_and_matches_reference(monkeypatch, seq, force_g
         calls.append((query.shape, key.shape, value.shape, query.dtype, kwargs))
         return sdpa(query, key, value, **kwargs)
 
+    def reject_fused(*args, **kwargs):
+        raise AssertionError("FP32 must keep Mammoth's native QK norm/RoPE path")
+
     with torch.no_grad(), monkeypatch.context() as patch:
+        patch.setattr(mammoth_dit, "fused_qk_norm_rope", reject_fused)
         patch.setattr(sdpa_backend, "can_sdpa_use_fused_gqa", check_gqa)
         patch.setattr(torch.nn.functional, "scaled_dot_product_attention", record_sdpa)
         patch.setattr(strategy, "pre_attention", prepare)
@@ -147,3 +276,206 @@ def test_fp32_falls_back_to_sdpa_and_matches_reference(monkeypatch, seq, force_g
     assert torch.count_nonzero(got[~mask]) == 0
     diff = (got.float() - want.float()).abs()[mask]
     assert diff.max().item() < 1e-3, diff.max().item()
+
+
+@pytest.mark.parametrize("layout", ["unbatched", "broadcast", "per_example", "strided"])
+@pytest.mark.parametrize("table_dtype", [torch.bfloat16, torch.float32])
+@pytest.mark.parametrize("force_native", [False, True])
+def test_paired_rope_with_learned_norm_weights(monkeypatch, layout, table_dtype, force_native):
+    torch.manual_seed(42)
+    batch, seq, head_dim = 2, 19, 120
+    attn = _norm_attention(head_dim, torch.bfloat16)
+    with torch.no_grad():
+        attn.norm_q.weight.normal_(1.0, 0.2)
+        attn.norm_k.weight.normal_(1.0, 0.2)
+    # Mimic views into packed QKV projections, including their nontrivial stride.
+    packed = torch.randn(batch, seq, HEADS + 2 * KV_HEADS, head_dim, device="cuda", dtype=torch.bfloat16)
+    q, k, _ = packed.split((HEADS, KV_HEADS, KV_HEADS), dim=2)
+    rotary = _production_rotary(batch if layout in ("per_example", "strided") else 1, seq, head_dim, table_dtype)
+    if layout == "unbatched":
+        rotary = tuple(t.squeeze(0) for t in rotary)
+    elif layout == "strided":
+        rotary = tuple(t.transpose(0, 1).contiguous().transpose(0, 1) for t in rotary)
+        assert not rotary[0].is_contiguous()
+    calls = []
+    fused = mammoth_dit.fused_qk_norm_rope
+
+    def record(*args, **kwargs):
+        calls.append(args[4].shape)
+        return fused(*args, **kwargs)
+
+    monkeypatch.setattr(mammoth_dit, "fused_qk_norm_rope", record)
+    if force_native:
+        monkeypatch.setattr(mammoth_dit, "_fused_cuda_supported", lambda *args, **kwargs: False)
+    monkeypatch.setenv("VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS", "0")
+    with torch.no_grad():
+        got = mammoth_dit._apply_qk_norm_rope(attn, q, k, rotary, rope_repeats_pairs=True)
+        expected = (
+            mammoth_dit.apply_real_rotary_emb(attn.norm_q(q), *rotary),
+            mammoth_dit.apply_real_rotary_emb(attn.norm_k(k), *rotary),
+        )
+    supported = mammoth_dit._fused_cuda_supported(q, k, head_dim, head_dim, interleaved=True)
+    assert calls == ([torch.Size((batch * seq, head_dim))] if supported else [])
+    for actual, want in zip(got, expected):
+        torch.testing.assert_close(actual, want.to(actual.dtype), atol=0.0625, rtol=0.02)
+
+
+@pytest.mark.parametrize(
+    "reason", ["support_query", "token_gate", "no_q_norm", "no_k_norm", "epsilon", "fp16", "head_dim", "no_rope"]
+)
+def test_qk_norm_rope_unsupported_inputs_keep_native_chain(monkeypatch, reason):
+    torch.manual_seed(42)
+    dtype = torch.float16 if reason == "fp16" else torch.bfloat16
+    dim = 260 if reason == "head_dim" else 120
+    attn = _norm_attention(dim, dtype)
+    q = torch.randn(2, 19, HEADS, dim, device="cuda", dtype=dtype)
+    k = torch.randn(2, 19, KV_HEADS, dim, device="cuda", dtype=dtype)
+    rotary = _production_rotary(2, 19, dim, dtype) if reason != "no_rope" else None
+    if reason == "support_query":
+        monkeypatch.setattr(mammoth_dit, "_fused_cuda_supported", lambda *args, **kwargs: False)
+    elif reason == "no_q_norm":
+        attn.norm_q = None
+    elif reason == "no_k_norm":
+        attn.norm_k = None
+    elif reason == "epsilon":
+        attn.norm_k.variance_epsilon = 1e-3
+    monkeypatch.setenv("VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS", "39" if reason == "token_gate" else "0")
+
+    def reject(*args, **kwargs):
+        raise AssertionError(f"Unsupported {reason} must keep the native chain")
+
+    monkeypatch.setattr(mammoth_dit, "fused_qk_norm_rope", reject)
+    with torch.no_grad():
+        got = mammoth_dit._apply_qk_norm_rope(attn, q, k, rotary, rope_repeats_pairs=True)
+        expected_q = attn.norm_q(q) if attn.norm_q is not None else q
+        expected_k = attn.norm_k(k) if attn.norm_k is not None else k
+        if rotary is not None:
+            expected_q = mammoth_dit.apply_real_rotary_emb(expected_q, *rotary)
+            expected_k = mammoth_dit.apply_real_rotary_emb(expected_k, *rotary)
+    torch.testing.assert_close(got[0], expected_q, atol=0, rtol=0)
+    torch.testing.assert_close(got[1], expected_k, atol=0, rtol=0)
+
+
+def test_attention_benchmark_single_token_and_restores_gate(monkeypatch):
+    from benchmarks.diffusion import benchmark_mammoth_moda2_qk_norm_rope as benchmark
+
+    env = "VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS"
+    monkeypatch.setenv(env, "123")
+    result = benchmark._run_attention(1, warmup=0, iters=1)
+    assert result["sequence"] == 1
+    assert benchmark.os.environ[env] == "123"
+
+
+@pytest.mark.parametrize("packed", [False, True], ids=["gated_native", "reused_packed"])
+def test_prepared_rope_table_skips_per_layer_gate_and_packing(monkeypatch, packed):
+    torch.manual_seed(7)
+    batch, seq, dim = 2, 19, 120
+    attn = _norm_attention(dim, torch.bfloat16)
+    q = torch.randn(batch, seq, HEADS, dim, device="cuda", dtype=torch.bfloat16)
+    k = torch.randn(batch, seq, KV_HEADS, dim, device="cuda", dtype=torch.bfloat16)
+    cos, sin = _production_rotary(batch, seq, dim, torch.float32)
+    table = torch.cat((cos[..., 0::2], sin[..., 0::2]), dim=-1).reshape(batch * seq, dim) if packed else None
+    rotary = cos, sin, table
+    calls = []
+    fused = mammoth_dit.fused_qk_norm_rope
+
+    def record(*args, **kwargs):
+        calls.append(args[4])
+        return fused(*args, **kwargs)
+
+    def forbid_gate(*args, **kwargs):
+        raise AssertionError("Prepared RoPE must not resolve the token gate per attention layer")
+
+    monkeypatch.setattr(mammoth_dit, "fused_qk_norm_rope", record)
+    monkeypatch.setattr(mammoth_dit, "fused_qk_norm_rope_min_tokens", forbid_gate)
+    with torch.no_grad():
+        for _ in range(2):
+            got = mammoth_dit._apply_qk_norm_rope(attn, q, k, rotary, rope_repeats_pairs=True)
+            expected = (
+                mammoth_dit.apply_real_rotary_emb(attn.norm_q(q), cos, sin),
+                mammoth_dit.apply_real_rotary_emb(attn.norm_k(k), cos, sin),
+            )
+            for actual, want in zip(got, expected):
+                torch.testing.assert_close(
+                    actual, want.to(actual.dtype), atol=0.0625 if packed else 0, rtol=0.02 if packed else 0
+                )
+    supported = packed and mammoth_dit._fused_cuda_supported(q, k, dim, dim, interleaved=True)
+    assert calls == ([table, table] if supported else [])
+
+
+@pytest.mark.parametrize("force_native", [False, True])
+def test_model_prepares_three_rope_tables_with_one_threshold_lookup(monkeypatch, force_native):
+    from vllm_omni.diffusion.models.mammoth_moda2.mammothmoda2_dit_model import Transformer2DModel
+    from vllm_omni.diffusion.models.mammoth_moda2.rope_real import RotaryPosEmbedReal
+
+    torch.manual_seed(7)
+    model = (
+        Transformer2DModel(
+            patch_size=2,
+            in_channels=4,
+            hidden_size=48,
+            num_layers=1,
+            num_refiner_layers=1,
+            num_attention_heads=6,
+            num_kv_heads=2,
+            multiple_of=8,
+            ffn_dim_multiplier=1.0,
+            axes_dim_rope=(2, 2, 4),
+            axes_lens=(8, 8, 8),
+            text_feat_dim=32,
+        )
+        .cuda()
+        .to(torch.bfloat16)
+        .eval()
+    )
+    hidden = torch.randn(2, 4, 4, 4, device="cuda", dtype=torch.bfloat16)
+    text = torch.randn(2, 4, 32, device="cuda", dtype=torch.bfloat16)
+    mask = torch.ones(2, 4, device="cuda", dtype=torch.bool)
+    freqs = RotaryPosEmbedReal.get_freqs_real((2, 2, 4), (8, 8, 8), 10000)
+    calls = []
+
+    def threshold(default):
+        calls.append(default)
+        return 9
+
+    if force_native:
+        monkeypatch.setattr(mammoth_dit, "_fused_cuda_supported", lambda *args, **kwargs: False)
+    supported = mammoth_dit._fused_cuda_supported(hidden, hidden, 8, 8, interleaved=True)
+    tables: list[mammoth_dit.RotaryEmbedding] = []
+    refiners = model._apply_refiners
+    layers = model._apply_transformer_layers
+
+    def record_refiners(text, text_mask, context, img, img_mask, noise, temb):
+        tables.extend((context, noise))
+        return refiners(text, text_mask, context, img, img_mask, noise, temb)
+
+    def record_layers(hidden, mask, joint, temb):
+        tables.append(joint)
+        return layers(hidden, mask, joint, temb)
+
+    monkeypatch.setattr(model, "_apply_refiners", record_refiners)
+    monkeypatch.setattr(model, "_apply_transformer_layers", record_layers)
+    monkeypatch.setattr(mammoth_dit, "fused_qk_norm_rope_min_tokens", threshold)
+    args = (hidden, torch.ones(2, device="cuda"), text, freqs, mask)
+    with torch.no_grad():
+        fused_output = model(*args)
+    context, noise, joint = tables
+    assert calls == ([0] if supported else [])
+    assert all(len(rotary) == 3 for rotary in (context, noise, joint))
+    assert context[2] is None and noise[2] is None
+    if supported:
+        assert joint[2].shape == (2 * 8, 8)
+    else:
+        assert joint[2] is None
+    for rotary in (context, noise, joint):
+        if rotary[2] is not None:
+            expected = torch.cat((rotary[0][..., 0::2], rotary[1][..., 0::2]), dim=-1).reshape(-1, 8).float()
+            torch.testing.assert_close(rotary[2], expected, atol=0, rtol=0)
+
+    # Exercise the prepared tuple through every real refiner and joint block,
+    # then compare the same model forward with the unsupported-device fallback.
+    with torch.no_grad():
+        monkeypatch.setattr(mammoth_dit, "_fused_cuda_supported", lambda *args, **kwargs: False)
+        native_output = model(*args)
+    assert fused_output.shape == hidden.shape
+    torch.testing.assert_close(fused_output, native_output, atol=0.0625, rtol=0.02)

tests/diffusion/models/mammoth_moda2/test_pipeline_mammothmoda2_dit.py

@@ -33,7 +33,10 @@ from vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit import (
     get_mammoth_moda2_pre_process_func,
 )
 from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.sched.step_scheduler import StepScheduler
+from vllm_omni.diffusion.worker.input_batch import InputBatch
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+from vllm_omni.diffusion.worker.utils import StepRequestState
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform
 
@@ -90,12 +93,12 @@ def test_root_weight_source_forwards_revision() -> None:
     assert _root_weight_source(config).revision == "rev-7"
 
 
-def test_pipeline_declares_native_components_and_batch_request_mode() -> None:
+def test_pipeline_declares_native_components_and_batch_modes() -> None:
     assert MammothModa2DiTPipeline._dit_modules == ["gen_transformer"]
     assert MammothModa2DiTPipeline._encoder_modules == ["gen_image_condition_refiner"]
     assert MammothModa2DiTPipeline._vae_modules == ["gen_vae"]
     assert MammothModa2DiTPipeline.supports_request_batch is True
-    assert MammothModa2DiTPipeline.supports_step_execution is False
+    assert MammothModa2DiTPipeline.supports_step_execution is True
 
 
 def test_mammoth_postprocess_denormalizes_nonnegative_raw_vae_output() -> None:
@@ -342,6 +345,21 @@ class _FakeTransformer(nn.Module):
         return torch.zeros_like(hidden_states)
 
 
+class _FakeImageRefiner(nn.Module):
+    def __init__(self, num_queries: int) -> None:
+        super().__init__()
+        self.anchor = nn.Parameter(torch.zeros(1))
+        self.num_queries = num_queries
+
+    def forward(self, image_embeds, attention_mask):
+        assert attention_mask.shape == image_embeds.shape[:2]
+        return image_embeds.new_zeros(
+            image_embeds.shape[0],
+            self.num_queries,
+            image_embeds.shape[-1],
+        )
+
+
 @dataclass
 class _FakeVaeConfig:
     scaling_factor: float | None = None
@@ -590,6 +608,25 @@ def test_pre_process_registers_batch_compatibility_key() -> None:
     assert req.batch_compatibility_key == ("mammoth_moda2_dit", 32, 48, 7)
 
 
+def test_step_pre_process_allows_different_step_counts() -> None:
+    config = _od_config()
+    config.step_execution = True
+    pre_process = get_mammoth_moda2_pre_process_func(config)
+    first = _batch(
+        request_id="short",
+        sampling=OmniDiffusionSamplingParams(height=32, width=48, num_inference_steps=2),
+    ).requests[0]
+    second = _batch(
+        request_id="long",
+        sampling=OmniDiffusionSamplingParams(height=32, width=48, num_inference_steps=7),
+    ).requests[0]
+
+    pre_process(first)
+    pre_process(second)
+
+    assert first.batch_compatibility_key == second.batch_compatibility_key == ("mammoth_moda2_dit", 32, 48)
+
+
 def test_pre_process_rejects_missing_ar_conditions() -> None:
     pre_process = get_mammoth_moda2_pre_process_func(_od_config())
     req = _batch(request_id="req-bad-ar", prompt={"prompt": "test"}).requests[0]
@@ -1020,3 +1057,367 @@ def test_admission_and_inference_threshold_resolution_matches(llm_cfg_patch: dic
     text_cond, image_cond = pipeline._split_request_conditions(parsed_at)
     assert text_cond.shape == (1, 8)
     assert image_cond.shape == (1, 8)
+
+
+def _step_state(batch: DiffusionRequestBatch) -> StepRequestState:
+    request = batch.requests[0]
+    return StepRequestState(
+        request_id=request.request_id,
+        sampling=request.sampling_params,
+        prompt=request.prompt,
+    )
+
+
+@pytest.mark.parametrize("step_execution", [False, True])
+@pytest.mark.parametrize(
+    ("standard_steps", "extra_args", "extra_info", "expected_steps"),
+    [
+        (None, {}, {}, 50),
+        (None, {"num_inference_steps": 30}, {}, 30),
+        (50, {"num_inference_steps": 30}, {}, 30),
+        (None, {}, {"num_inference_steps": [20]}, 20),
+        (10, {}, {"num_inference_steps": [20]}, 10),
+    ],
+)
+def test_pre_process_normalizes_steps_before_scheduler_admission(
+    step_execution: bool,
+    standard_steps: int | None,
+    extra_args: dict,
+    extra_info: dict,
+    expected_steps: int,
+) -> None:
+    config = _od_config()
+    config.step_execution = step_execution
+    request = _batch(
+        sampling=OmniDiffusionSamplingParams(num_inference_steps=standard_steps, extra_args=extra_args),
+    ).requests[0]
+    request.prompt["additional_information"].update(extra_info)
+    request = get_mammoth_moda2_pre_process_func(config)(request)
+
+    assert request.sampling_params.num_inference_steps == expected_steps
+    parsed = _pipeline_shell()._parse_request(DiffusionRequestBatch([request]))
+    assert parsed.num_inference_steps == expected_steps
+    if step_execution:
+        scheduler = StepScheduler()
+        assert scheduler.add_request(request) == request.request_id
+        assert scheduler._request_progress[request.request_id].total_steps == expected_steps
+    else:
+        assert request.batch_compatibility_key[-1] == expected_steps
+
+
+def test_step_protocol_matches_request_mode() -> None:
+    pipeline = _pipeline_shell()
+    pipeline.gen_transformer = _FakeTransformer()
+    pipeline.gen_image_condition_refiner = None
+    pipeline.gen_vae = _FakeVae()
+    pipeline.gen_freqs_cis = torch.zeros(1)
+    batch = _batch(
+        sampling=OmniDiffusionSamplingParams(
+            height=32,
+            width=48,
+            seed=42,
+            guidance_scale=1.0,
+            num_inference_steps=2,
+        )
+    )
+
+    module = "vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit"
+    with (
+        patch(f"{module}.FlowMatchEulerDiscreteScheduler", side_effect=lambda: _FakeScheduler()),
+        patch(f"{module}.randn_tensor", side_effect=lambda shape, **kwargs: torch.zeros(shape, dtype=kwargs["dtype"])),
+    ):
+        request_output = pipeline.forward(batch)[0]
+        state = pipeline.prepare_encode(_step_state(batch))
+        while not state.denoise_completed:
+            input_batch = InputBatch.make_batch([state])
+            noise_pred = pipeline.denoise_step(input_batch, states=[state])
+            pipeline.step_scheduler(state, noise_pred)
+        step_output = pipeline.post_decode(state)
+
+    torch.testing.assert_close(step_output.output, request_output.output)
+    assert state.step_index == 2
+
+
+@pytest.mark.parametrize(
+    ("guidance_scales", "expected_predictions", "expected_calls"),
+    [
+        ((1.0, 4.0), (2.0, 5.0), 2),
+        ((4.0, 1.0), (5.0, 2.0), 2),
+        ((1.0, 1.0), (2.0, 2.0), 1),
+    ],
+)
+def test_step_batch_handles_cfg_and_non_cfg_requests(
+    guidance_scales: tuple[float, float],
+    expected_predictions: tuple[float, float],
+    expected_calls: int,
+) -> None:
+    pipeline = _pipeline_shell()
+    pipeline.gen_transformer = _FakeTransformer()
+    pipeline.gen_image_condition_refiner = None
+    pipeline.gen_freqs_cis = torch.zeros(1)
+    config = _od_config()
+    config.step_execution = True
+    pre_process = get_mammoth_moda2_pre_process_func(config)
+    scheduler = StepScheduler()
+    requests = [
+        pre_process(
+            _batch(
+                request_id=f"cfg-{index}",
+                sampling=OmniDiffusionSamplingParams(
+                    height=32,
+                    width=48,
+                    seed=index,
+                    guidance_scale=4.0,
+                    num_inference_steps=2,
+                    extra_args={"text_guidance_scale": scale},
+                ),
+            ).requests[0]
+        )
+        for index, scale in enumerate(guidance_scales)
+    ]
+    assert scheduler._build_sampling_params_key(requests[0]) == scheduler._build_sampling_params_key(requests[1])
+
+    def predict(*, hidden_states, text_hidden_states, **kwargs):
+        return torch.full_like(hidden_states, 2.0 if text_hidden_states.shape[1] else 1.0)
+
+    module = "vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit"
+    with (
+        patch(f"{module}.FlowMatchEulerDiscreteScheduler", side_effect=lambda: _FakeScheduler()),
+        patch.object(pipeline.gen_transformer, "forward", side_effect=predict) as transformer_forward,
+    ):
+        states = [pipeline.prepare_encode(_step_state(DiffusionRequestBatch([request]))) for request in requests]
+        predictions = pipeline.denoise_step(InputBatch.make_batch(states), states=states)
+
+    for prediction, expected in zip(predictions, expected_predictions):
+        torch.testing.assert_close(prediction, torch.full_like(prediction, expected))
+    assert transformer_forward.call_count == expected_calls
+
+
+@pytest.mark.parametrize(
+    ("sampling", "extra_info", "expected_guidance"),
+    [
+        (OmniDiffusionSamplingParams(num_inference_steps=2), {}, 9.0),
+        (OmniDiffusionSamplingParams(num_inference_steps=2), {"text_guidance_scale": [3.5]}, 3.5),
+    ],
+)
+def test_step_protocol_resolves_guidance_like_request_mode(
+    sampling: OmniDiffusionSamplingParams,
+    extra_info: dict,
+    expected_guidance: float,
+) -> None:
+    """Guidance from defaults or additional_information must survive step mode.
+
+    prepare_encode() must parse the already-initialized step state directly
+    instead of re-wrapping its sampling params in a fresh OmniDiffusionRequest,
+    whose __post_init__ promotes the auto-filled 1.0 guidance_scale to an
+    explicitly-provided value and silently disables model-default CFG.
+    """
+    pipeline = _pipeline_shell()
+    pipeline.gen_transformer = _FakeTransformer()
+    pipeline.gen_image_condition_refiner = None
+    pipeline.gen_vae = _FakeVae()
+    pipeline.gen_freqs_cis = torch.zeros(1)
+    batch = _batch(sampling=sampling)
+    batch.prompts[0]["additional_information"].update(extra_info)
+    assert sampling.guidance_scale_provided is False
+
+    module = "vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit"
+    with (
+        patch(f"{module}.FlowMatchEulerDiscreteScheduler", side_effect=lambda: _FakeScheduler()),
+        patch(f"{module}.randn_tensor", side_effect=lambda shape, **kwargs: torch.zeros(shape, dtype=kwargs["dtype"])),
+    ):
+        request_spec = pipeline._parse_request(batch)
+        request_output = pipeline.forward(batch)[0]
+        state = pipeline.prepare_encode(_step_state(batch))
+        while not state.denoise_completed:
+            noise_pred = pipeline.denoise_step(InputBatch.make_batch([state]), states=[state])
+            pipeline.step_scheduler(state, noise_pred)
+        step_output = pipeline.post_decode(state)
+
+    assert request_spec.text_guidance_scale == expected_guidance
+    assert state.extra["mammoth_text_guidance_scale"] == expected_guidance
+    assert state.negative_prompt_embeds is not None
+    assert sampling.guidance_scale_provided is False
+    torch.testing.assert_close(step_output.output, request_output.output)
+
+
+def test_refiner_output_uses_query_length_mask_in_request_and_step_modes() -> None:
+    pipeline = _pipeline_shell()
+    pipeline.gen_transformer = _FakeTransformer()
+    pipeline.gen_image_condition_refiner = _FakeImageRefiner(num_queries=3)
+    pipeline.gen_vae = _FakeVae()
+    pipeline.gen_freqs_cis = torch.zeros(1)
+    batch = _batch(
+        sampling=OmniDiffusionSamplingParams(
+            height=32,
+            width=48,
+            seed=42,
+            guidance_scale=1.0,
+            num_inference_steps=1,
+        )
+    )
+
+    module = "vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit"
+    with (
+        patch(f"{module}.FlowMatchEulerDiscreteScheduler", side_effect=lambda: _FakeScheduler()),
+        patch(f"{module}.randn_tensor", side_effect=lambda shape, **kwargs: torch.zeros(shape, dtype=kwargs["dtype"])),
+    ):
+        assert len(pipeline.forward(batch)) == 1
+        state = pipeline.prepare_encode(_step_state(batch))
+
+    assert state.prompt_embeds is not None
+    assert state.prompt_embeds_mask is not None
+    assert state.prompt_embeds.shape[1] == 5
+    assert state.prompt_embeds_mask.shape == state.prompt_embeds.shape[:2]
+    assert state.prompt_embeds_mask.all()
+
+
+def test_step_protocol_keeps_request_schedulers_and_progress_independent() -> None:
+    pipeline = _pipeline_shell()
+    pipeline.gen_transformer = _FakeTransformer()
+    pipeline.gen_image_condition_refiner = None
+    pipeline.gen_vae = _FakeVae()
+    pipeline.gen_freqs_cis = torch.zeros(1)
+    short = _step_state(
+        _batch(
+            request_id="short",
+            sampling=OmniDiffusionSamplingParams(
+                height=32,
+                width=48,
+                seed=1,
+                guidance_scale=1.0,
+                num_inference_steps=1,
+            ),
+        )
+    )
+    long = _step_state(
+        _batch(
+            request_id="long",
+            sampling=OmniDiffusionSamplingParams(
+                height=32,
+                width=48,
+                seed=2,
+                guidance_scale=1.0,
+                num_inference_steps=2,
+            ),
+        )
+    )
+
+    module = "vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit"
+    with (
+        patch(f"{module}.FlowMatchEulerDiscreteScheduler", side_effect=lambda: _FakeScheduler()),
+        patch(f"{module}.randn_tensor", side_effect=lambda shape, **kwargs: torch.zeros(shape, dtype=kwargs["dtype"])),
+    ):
+        pipeline.prepare_encode(short)
+        pipeline.prepare_encode(long)
+        predictions = pipeline.denoise_step(InputBatch.make_batch([short, long]), states=[short, long])
+        pipeline.step_scheduler(short, predictions[:1])
+        pipeline.step_scheduler(long, predictions[1:])
+
+    assert short.scheduler is not long.scheduler
+    assert short.denoise_completed
+    assert not long.denoise_completed
+    assert short.step_index == long.step_index == 1
+
+
+class _InputDependentTransformer(_FakeTransformer):
+    """Expose dropped masks, timesteps, conditioning and latent rows in outputs."""
+
+    def forward(self, *, hidden_states, timestep, text_hidden_states, text_attention_mask, **kwargs):
+        self.calls += 1
+        conditioning = (text_hidden_states * text_attention_mask.unsqueeze(-1)).sum(dim=(1, 2)) / 100
+        return hidden_states * 0.125 + timestep[:, None, None, None] + conditioning[:, None, None, None]
+
+
+def _step_test_pipeline() -> MammothModa2DiTPipeline:
+    pipeline = _pipeline_shell()
+    pipeline.gen_transformer = _InputDependentTransformer()
+    pipeline.gen_image_condition_refiner = None
+    pipeline.gen_vae = _FakeVae()
+    pipeline.gen_freqs_cis = torch.zeros(1)
+    return pipeline
+
+
+def _step_test_request(request_id: str, *, seed: int = 11, steps: int = 4, text_tokens: int = 2):
+    return _batch(
+        request_id=request_id,
+        prompt={
+            "prompt": "",
+            "additional_information": {
+                "full_hidden_states": torch.arange((text_tokens + 2) * 8, dtype=torch.float32).reshape(-1, 8) / 10,
+                "full_token_ids": [10] * text_tokens + [100, 101],
+                "answer_start_index": text_tokens,
+            },
+        },
+        sampling=OmniDiffusionSamplingParams(
+            height=32,
+            width=48,
+            seed=seed,
+            guidance_scale=4.0,
+            num_inference_steps=steps,
+            extra_args={"cfg_range": [0.25, 0.75]},
+        ),
+    )
+
+
+def test_step_protocol_seed_conditioning_and_request_parity() -> None:
+    pipeline = _step_test_pipeline()
+    outputs = []
+    for seed, text_tokens in ((11, 2), (11, 2), (12, 2), (11, 5)):
+        batch = _step_test_request("solo", seed=seed, text_tokens=text_tokens)
+        state = pipeline.prepare_encode(_step_state(batch))
+        expected_noise = torch.randn(state.latents.shape, generator=torch.Generator().manual_seed(seed))
+        torch.testing.assert_close(state.latents, expected_noise, rtol=0, atol=0)
+        while not state.denoise_completed:
+            prediction = pipeline.denoise_step(InputBatch.make_batch([state]), states=[state])
+            pipeline.step_scheduler(state, prediction)
+        output = pipeline.post_decode(state).output
+        reference = pipeline.forward(batch)[0].output
+        torch.testing.assert_close(output, reference)
+        outputs.append(output)
+    torch.testing.assert_close(outputs[0], outputs[1], rtol=0, atol=0)
+    assert not torch.allclose(outputs[0], outputs[2])
+    assert not torch.allclose(outputs[0], outputs[3])
+
+
+def test_step_protocol_handles_arrival_reordering_and_retirement() -> None:
+    pipeline = _step_test_pipeline()
+    long = pipeline.prepare_encode(_step_state(_step_test_request("long")))
+    solo_long = pipeline.prepare_encode(_step_state(_step_test_request("solo-long")))
+    solo_short = pipeline.prepare_encode(_step_state(_step_test_request("solo-short", seed=22, steps=2, text_tokens=5)))
+    short = None
+    cached_batch = None
+    completed = {}
+    retired_latents = None
+
+    # Match runner ordering: newly admitted requests precede existing requests.
+    for tick in range(4):
+        if tick == 1:
+            short = pipeline.prepare_encode(_step_state(_step_test_request("short", seed=22, steps=2, text_tokens=5)))
+            assert short.scheduler is not long.scheduler
+        active = [short, long] if tick in (1, 2) else [long]
+        references = [solo_short, solo_long] if tick in (1, 2) else [solo_long]
+        cached_batch = InputBatch.make_batch(active, cached_batch=cached_batch)
+        if tick == 1:
+            assert long.current_timestep != short.current_timestep
+            assert cached_batch.prompt_embeds_mask.sum(dim=1).tolist() == [7, 4]
+        predictions = pipeline.denoise_step(cached_batch, states=active)
+        for index, (state, reference) in enumerate(zip(active, references)):
+            solo_prediction = pipeline.denoise_step(InputBatch.make_batch([reference]), states=[reference])
+            torch.testing.assert_close(predictions[index : index + 1], solo_prediction)
+            pipeline.step_scheduler(state, predictions[index : index + 1])
+            pipeline.step_scheduler(reference, solo_prediction)
+            torch.testing.assert_close(state.latents, reference.latents)
+            assert state.step_index == reference.step_index == state.scheduler.step_index
+            if state.denoise_completed:
+                completed[state.request_id] = pipeline.post_decode(state).output
+                torch.testing.assert_close(completed[state.request_id], pipeline.post_decode(reference).output)
+        if tick == 2:
+            assert short.denoise_completed and not long.denoise_completed
+            assert list(completed) == ["short"]
+            retired_latents = short.latents.clone()
+
+    assert list(completed) == ["short", "long"]
+    assert (short.step_index, long.step_index) == (2, 4)
+    torch.testing.assert_close(short.latents, retired_latents, rtol=0, atol=0)
+    assert not torch.allclose(completed["short"], completed["long"])

tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py

@@ -71,12 +71,12 @@ def _small_od_config(*, class_name=None):
 
 def test_fp8_scope_and_prefix_propagation(monkeypatch):
     from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
-    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
     from vllm.model_executor.layers.quantization.utils.quant_utils import (
         is_layer_skipped,
     )
 
     from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
+    from vllm_omni.quantization.fp8_config import DiffusionFp8Config
 
     monkeypatch.setattr(h3, "ColumnParallelLinear", _FakeLinear)
     monkeypatch.setattr(h3, "MergedColumnParallelLinear", _FakeLinear)
@@ -92,7 +92,7 @@ def test_fp8_scope_and_prefix_propagation(monkeypatch):
         "blocks.0.adaln_proj.linear",
         "final_layer.adaln_proj.linear",
     }
-    fp8_config = Fp8Config(ignored_layers=sorted(ignored_layers))
+    fp8_config = DiffusionFp8Config(ignored_layers=sorted(ignored_layers))
     model = h3.MiniMaxH3DiTModel(
         _small_od_config(),
         quant_config=fp8_config,

tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py

@@ -102,7 +102,11 @@ def _make_state(request_id: str, model, branch, video_rows, audio_rows, sigmas_v
     state.latents = video_rows.clone()
     state.timesteps = torch.tensor([1.0 - sigma for sigma in sigmas_video[:-1]], dtype=torch.float32)
     state.step_index = 0
+    from vllm_omni.diffusion.models.minimax_h3.sampling import create_h3_sample_solver
+
     state.extra = {
+        mod._STEP_VIDEO_SOLVER: create_h3_sample_solver("euler", sigmas_video),
+        mod._STEP_AUDIO_SOLVER: create_h3_sample_solver("euler", sigmas_audio),
         mod._STEP_BRANCH: branch,
         # Co-batched requests must share one DiT instance, or denoise_step()
         # treats the batch as mixed-task and falls back to one forward each.

tests/diffusion/offloader/test_layerwise_backend.py

@@ -542,6 +542,24 @@ class TestLayerwiseComponentSelection:
         for name, tensor in pipeline.state_dict().items():
             torch.testing.assert_close(tensor, expected[name])
 
+    def test_shutdown_skips_block_restore(self, patched_offload_runtime, monkeypatch):
+        # Restoring every block copies the whole DiT back onto the device, which
+        # outlived the executor's shutdown grace period (issue 7625).
+        pipeline = nn.Module()
+        pipeline.transformer = _SingleBlockModel(num_blocks=3)
+        backend = _layer_backend()
+        backend.enable(pipeline)
+        restores = Mock()
+        for hook in backend._dit_hooks:
+            monkeypatch.setattr(hook, "restore_next_block", restores)
+
+        backend.shutdown()
+
+        restores.assert_not_called()
+        assert not backend.enabled
+        assert not backend._dit_hooks
+        assert not backend._hooked_dit_blocks
+
     def test_encoder_disable_failure_keeps_host_masters_for_retry(
         self,
         patched_offload_runtime,

tests/diffusion/quantization/test_fp8_config.py

@@ -15,6 +15,74 @@ def test_build_quant_config_fp8():
     assert config is not None
     assert config.get_name() == "fp8"
     assert config.activation_scheme == "dynamic"
+    assert not config.is_checkpoint_fp8_serialized
+
+
+def test_build_quant_config_fp8_uses_native_online_quantization():
+    from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig
+    from vllm.model_executor.layers.quantization.utils.quant_utils import kFp8StaticTensorSym
+
+    from vllm_omni.quantization import build_quant_config
+
+    config = build_quant_config("fp8", ignored_layers=["proj_out"])
+    assert isinstance(config, OnlineQuantizationConfig)
+    assert config.args.linear.weight == kFp8StaticTensorSym
+    assert config.args.moe.weight == kFp8StaticTensorSym
+    assert config.ignored_layers == ["proj_out"]
+
+
+def test_build_quant_config_upstream_online_fp8_shorthand():
+    from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig
+    from vllm.model_executor.layers.quantization.utils.quant_utils import kFp8StaticTensorSym
+
+    from vllm_omni.quantization import build_quant_config
+
+    config = build_quant_config("fp8_per_tensor", ignore=["proj_out"])
+    assert isinstance(config, OnlineQuantizationConfig)
+    assert config.args.linear.weight == kFp8StaticTensorSym
+    assert config.ignored_layers == ["proj_out"]
+
+
+def test_online_fp8_rejects_static_activation():
+    from vllm_omni.quantization import build_quant_config
+
+    with pytest.raises(ValueError, match="activation_scheme='dynamic'"):
+        build_quant_config("fp8", activation_scheme="static")
+
+
+def test_build_online_quant_config_preserves_explicit_native_args():
+    from vllm.config.quantization import QuantizationConfigArgs, QuantSpec
+    from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig
+
+    from vllm_omni.quantization import build_quant_config
+
+    args = QuantizationConfigArgs(linear=QuantSpec(weight="fp8_per_tensor_static"), ignore=["proj_out"])
+    config = build_quant_config("online", args=args)
+    assert isinstance(config, OnlineQuantizationConfig)
+    assert config.args is args
+    assert config.ignored_layers == ["proj_out"]
+
+
+def test_build_online_quant_config_requires_args():
+    from vllm_omni.quantization import build_quant_config
+
+    with pytest.raises(ValueError, match="requires quantization config arguments"):
+        build_quant_config("online")
+
+
+def test_serialized_checkpoint_replaces_online_fp8_config():
+    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
+
+    from vllm_omni.quantization import build_quant_config
+    from vllm_omni.quantization.factory import resolve_quant_config_from_disk
+
+    config = resolve_quant_config_from_disk(
+        build_quant_config("fp8"),
+        {"quant_method": "fp8", "is_checkpoint_fp8_serialized": True, "activation_scheme": "static"},
+    )
+    assert isinstance(config, Fp8Config)
+    assert config.is_checkpoint_fp8_serialized
+    assert config.activation_scheme == "static"
 
 
 def test_build_quant_config_none():
@@ -39,7 +107,7 @@ def test_build_quant_config_invalid():
 def test_build_quant_config_dict():
     from vllm_omni.quantization import build_quant_config
 
-    config = build_quant_config({"method": "fp8", "activation_scheme": "static"})
+    config = build_quant_config({"method": "fp8", "is_checkpoint_fp8_serialized": True, "activation_scheme": "static"})
     assert config is not None
     assert config.get_name() == "fp8"
     assert config.activation_scheme == "static"
@@ -48,7 +116,7 @@ def test_build_quant_config_dict():
 def test_build_quant_config_dict_not_mutated():
     from vllm_omni.quantization import build_quant_config
 
-    original = {"method": "fp8", "activation_scheme": "static"}
+    original = {"method": "fp8", "is_checkpoint_fp8_serialized": True, "activation_scheme": "static"}
     copy = original.copy()
     build_quant_config(original)
     assert original == copy
@@ -99,7 +167,7 @@ def test_build_quant_config_per_component_inner_dict_not_mutated():
     """Inner component dicts should not be mutated by build_quant_config."""
     from vllm_omni.quantization import build_quant_config
 
-    inner = {"method": "fp8", "activation_scheme": "static"}
+    inner = {"method": "fp8", "is_checkpoint_fp8_serialized": True, "activation_scheme": "static"}
     original = inner.copy()
     build_quant_config({"transformer": inner, "vae": None})
     assert inner == original
@@ -115,20 +183,18 @@ def test_flat_dict_not_misdetected_as_per_component():
 
 
 def test_build_quant_config_passthrough():
-    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
-
     from vllm_omni.quantization import build_quant_config
+    from vllm_omni.quantization.fp8_config import DiffusionFp8Config
 
-    fp8 = Fp8Config(is_checkpoint_fp8_serialized=False, activation_scheme="dynamic")
+    fp8 = DiffusionFp8Config()
     assert build_quant_config(fp8) is fp8
 
 
 def test_component_config_routing():
-    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
-
     from vllm_omni.quantization import ComponentQuantizationConfig
+    from vllm_omni.quantization.fp8_config import DiffusionFp8Config
 
-    fp8 = Fp8Config(is_checkpoint_fp8_serialized=False, activation_scheme="dynamic")
+    fp8 = DiffusionFp8Config()
     config = ComponentQuantizationConfig(
         component_configs={"transformer": fp8, "vae": None},
     )
@@ -140,11 +206,10 @@ def test_component_config_routing():
 
 
 def test_component_config_with_default():
-    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
-
     from vllm_omni.quantization import ComponentQuantizationConfig
+    from vllm_omni.quantization.fp8_config import DiffusionFp8Config
 
-    fp8 = Fp8Config(is_checkpoint_fp8_serialized=False, activation_scheme="dynamic")
+    fp8 = DiffusionFp8Config()
     config = ComponentQuantizationConfig(
         component_configs={"vae": None},
         default_config=fp8,
@@ -167,7 +232,7 @@ def test_integration_dict():
 
     config = OmniDiffusionConfig(
         model="test",
-        quantization_config={"method": "fp8", "activation_scheme": "static"},
+        quantization_config={"method": "fp8", "is_checkpoint_fp8_serialized": True, "activation_scheme": "static"},
     )
     assert config.quantization_config is not None
     assert config.quantization_config.get_name() == "fp8"

tests/diffusion/quantization/test_online_fp8_gpu.py

@@ -0,0 +1,72 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Exercise online FP8 loading and inference through the native CUDA path."""
+
+import pytest
+import torch
+from torch.nn import functional as F
+from vllm.config import set_current_vllm_config
+from vllm.distributed import (
+    destroy_distributed_environment,
+    destroy_model_parallel,
+    init_distributed_environment,
+    initialize_model_parallel,
+)
+from vllm.model_executor.layers.linear import ReplicatedLinear, UnquantizedLinearMethod
+from vllm.model_executor.layers.quantization.online.fp8 import Fp8PerTensorOnlineLinearMethod
+from vllm.utils.network_utils import get_open_port
+from vllm.utils.torch_utils import set_default_torch_dtype
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.vllm_config import create_diffusion_vllm_config
+from vllm_omni.quantization import build_quant_config
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
+
+
+@pytest.mark.parametrize("method", ["fp8", "fp8_per_tensor"])
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_online_fp8_loads_bf16_weights_and_runs_native_kernel(method):
+    ignored_key = "ignored_layers" if method == "fp8" else "ignore"
+    config = build_quant_config(method, **{ignored_key: ["skip"]})
+    od_config = OmniDiffusionConfig(model="test", dtype=torch.bfloat16, quantization_config=config)
+    vllm_config = create_diffusion_vllm_config(torch.device("cuda:0"), od_config)
+
+    try:
+        init_distributed_environment(
+            world_size=1,
+            rank=0,
+            local_rank=0,
+            distributed_init_method=f"tcp://127.0.0.1:{get_open_port()}",
+            backend="nccl",
+        )
+        with set_current_vllm_config(vllm_config), set_default_torch_dtype(torch.bfloat16), torch.device("cuda:0"):
+            initialize_model_parallel(tensor_model_parallel_size=1, pipeline_model_parallel_size=1)
+            layer = ReplicatedLinear(
+                128, 128, bias=False, params_dtype=torch.bfloat16, quant_config=config, prefix="linear", disable_tp=True
+            )
+            assert isinstance(layer.quant_method, Fp8PerTensorOnlineLinearMethod)
+            assert layer.weight.is_meta
+
+            generator = torch.Generator(device="cuda").manual_seed(42)
+            weights = torch.randn(128, 128, generator=generator, dtype=torch.bfloat16) / 128**0.5
+            inputs = torch.randn(32, 128, generator=generator, dtype=torch.bfloat16)
+            expected = F.linear(inputs, weights).float()
+
+            # This is the loader installed by the native online quantizer.
+            # It materializes and quantizes the BF16 checkpoint in one pass.
+            layer.weight.weight_loader(layer.weight, weights)
+            assert layer.weight.dtype == torch.float8_e4m3fn
+            actual = layer(inputs)[0].float()
+            relative_mse = (actual - expected).square().mean() / expected.square().mean()
+            assert relative_mse.item() < 0.002
+
+            ignored = ReplicatedLinear(
+                128, 128, bias=False, params_dtype=torch.bfloat16, quant_config=config, prefix="skip", disable_tp=True
+            )
+            assert isinstance(ignored.quant_method, UnquantizedLinearMethod)
+            assert ignored.weight.dtype == torch.bfloat16
+    finally:
+        destroy_model_parallel()
+        destroy_distributed_environment()

tests/diffusion/test_diffusion_engine_cleanup.py

@@ -31,6 +31,7 @@ def _make_request(request_id: str) -> OmniDiffusionRequest:
 
 def _make_engine() -> DiffusionEngine:
     engine = DiffusionEngine.__new__(DiffusionEngine)
+    engine.od_config = SimpleNamespace(distributed_executor_backend=None)
     engine.scheduler = RequestScheduler()
     engine.scheduler.initialize(SimpleNamespace())
     engine.executor = SimpleNamespace(shutdown=Mock())

tests/diffusion/test_multiproc_engine_concurrency.py

@@ -1290,6 +1290,44 @@ class TestStageDiffusionClientErrorPropagation:
         client = self._make_client()
         client.check_health()
 
+    def test_local_shutdown_uses_existing_signal_path(self):
+        client = self._make_client()
+        client._zmq_ctx = MagicMock()
+        client._proc_manager.distributed_executor_backend = "mp"
+        client._proc_manager.wait_for_shutdown = MagicMock()
+        client._proc_manager.shutdown = MagicMock()
+        client.shutdown()
+        client._proc_manager.wait_for_shutdown.assert_not_called()
+        client._proc_manager.shutdown.assert_called_once_with(timeout=10)
+
+    def test_shutdown_waits_for_clean_subprocess_exit(self):
+        client = self._make_client()
+        client._zmq_ctx = MagicMock()
+        client._proc_manager.distributed_executor_backend = "ray"
+        client._proc_manager.wait_for_shutdown = MagicMock(return_value=True)
+        client._proc_manager.shutdown = MagicMock()
+
+        client.shutdown()
+
+        client._request_socket.send.assert_called_once()
+        client._proc_manager.wait_for_shutdown.assert_called_once_with(timeout=10)
+        client._proc_manager.shutdown.assert_not_called()
+        client._request_socket.close.assert_called_once_with(linger=0)
+        client._response_socket.close.assert_called_once_with(linger=0)
+        client._zmq_ctx.term.assert_called_once_with()
+
+    def test_shutdown_uses_signals_only_after_clean_exit_timeout(self):
+        client = self._make_client()
+        client._zmq_ctx = MagicMock()
+        client._proc_manager.distributed_executor_backend = "ray"
+        client._proc_manager.wait_for_shutdown = MagicMock(return_value=False)
+        client._proc_manager.shutdown = MagicMock()
+
+        client.shutdown()
+
+        client._proc_manager.wait_for_shutdown.assert_called_once_with(timeout=10)
+        client._proc_manager.shutdown.assert_called_once_with(timeout=10)
+
     def test_get_output_raises_engine_dead_when_dead(self):
         """When ``_engine_dead`` is True and the output queue is empty,
         ``get_diffusion_output_nowait`` must raise ``EngineDeadError``."""
@@ -1783,3 +1821,29 @@ class TestDrainResponsesDeathSentinel:
         assert output.request_id == "req-fail"
         assert output.error == "gpu fault"
         assert output.finished is True
+
+
+@pytest.mark.parametrize("backend", ["ray", "mp"])
+def test_engine_close_interrupts_only_ray_before_join(backend):
+    engine = DiffusionEngine.__new__(DiffusionEngine)
+    engine.od_config = SimpleNamespace(distributed_executor_backend=backend)
+    engine._cv = threading.Condition()
+    engine._closed = False
+    engine._shutdown_complete = False
+    engine.stop_event = threading.Event()
+    engine._out_streams = {}
+    interrupted = threading.Event()
+    engine.executor = MagicMock()
+    engine.executor.shutdown.side_effect = interrupted.set
+    engine.scheduler = MagicMock(_native_prefetch_enabled=False)
+    engine.worker_thread = MagicMock()
+    engine.worker_thread.is_alive.side_effect = [True, False]
+
+    def join(timeout):
+        assert engine._closed
+        assert interrupted.is_set() == (backend == "ray")
+
+    engine.worker_thread.join.side_effect = join
+    engine.close()
+    assert engine._shutdown_complete
+    engine.scheduler.close.assert_called_once()

tests/diffusion/test_ray_executor.py

@@ -0,0 +1,586 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""CPU tests of Ray placement, rank assignment, RPC output, and cleanup."""
+
+import threading
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.data import DiffusionOutput
+from vllm_omni.diffusion.executor import ray_executor as ray_module
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def test_cross_node_workers_get_global_rank_and_explicit_rendezvous(monkeypatch):
+    monkeypatch.setenv("NCCL_SOCKET_IFNAME", "eth-test")
+    monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "3,4,5")
+    actors = []
+    for ip in ["10.0.0.2", "10.0.0.1", "10.0.0.2"]:
+        actor = Mock()
+        actor.get_node_ip.remote.return_value = ip
+        actor.get_open_port.remote.return_value = 23456
+        actors.append(actor)
+    factory = Mock()
+    factory.remote.side_effect = actors
+    fake_ray = Mock()
+    fake_ray.remote.return_value = lambda cls: factory
+    fake_ray.get.side_effect = lambda values, **kwargs: values
+    monkeypatch.setattr(ray_module, "ray", fake_ray)
+    monkeypatch.setattr(ray_module, "get_ip", lambda: "10.0.0.1")
+    monkeypatch.setattr(ray_module, "PlacementGroupSchedulingStrategy", Mock())
+    executor = object.__new__(ray_module.RayDiffusionExecutor)
+    executor.od_config = SimpleNamespace(num_gpus=3, ray_worker_env={"CUSTOM_PLUGIN_SETTING": "enabled"})
+    executor.workers = []
+    executor._init_workers(object())
+    assert fake_ray.remote.call_args.kwargs["max_concurrency"] == 1
+    assert fake_ray.remote.call_args.kwargs["concurrency_groups"] == {"health": 1, "monitor": 1}
+    worker_env = fake_ray.remote.call_args.kwargs["runtime_env"]["env_vars"]
+    assert worker_env["NCCL_SOCKET_IFNAME"] == "eth-test"
+    assert worker_env["CUSTOM_PLUGIN_SETTING"] == "enabled"
+    assert "CUDA_VISIBLE_DEVICES" not in worker_env
+    assert len(fake_ray.get.call_args_list) == 3
+    assert all(call.kwargs["timeout"] == ray_module._WORKER_INIT_TIMEOUT_S for call in fake_ray.get.call_args_list)
+    assert [worker.ip for worker in executor.workers] == ["10.0.0.1", "10.0.0.2", "10.0.0.2"]
+    for rank, metadata in enumerate(executor.workers):
+        metadata.worker.init_worker.remote.assert_called_once_with(executor.od_config, rank, "tcp://10.0.0.1:23456")
+
+
+def test_worker_env_forwards_stage_overrides_without_driver_identity(monkeypatch):
+    from unittest.mock import patch
+
+    driver_env = {
+        "HF_TOKEN": "test-token",
+        "VLLM_PLUGINS": "test-plugin",
+        "NCCL_DEBUG": "WARN",
+        "DIFFUSION_ATTENTION_BACKEND": "FLASH_ATTN",
+        "PYTHONPATH": "/test/modules",
+        "UNRELATED_SECRET": "do-not-copy",
+        "HOME": "/driver/home",
+    }
+    protected = {
+        "CUDA_VISIBLE_DEVICES": "7",
+        "VLLM_HOST_IP": "10.0.0.1",
+        "VLLM_HOST_PORT": "1234",
+        "VLLM_NIXL_SIDE_CHANNEL_HOST": "10.0.0.1",
+        "LOCAL_RANK": "7",
+        "RANK": "7",
+        "WORLD_SIZE": "8",
+        "MASTER_ADDR": "driver-host",
+        "MASTER_PORT": "1234",
+        "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1",
+    }
+    config = SimpleNamespace(ray_worker_env={**protected, "NCCL_DEBUG": "INFO", "CUSTOM_PLUGIN_SETTING": "stage"})
+    with patch.dict(ray_module.os.environ, {**driver_env, **protected}, clear=True):
+        env = ray_module._worker_env(config)
+        assert env == {key: value for key, value in driver_env.items() if key not in {"UNRELATED_SECRET", "HOME"}} | {
+            "NCCL_DEBUG": "INFO",
+            "CUSTOM_PLUGIN_SETTING": "stage",
+        }
+        assert ray_module.os.environ["NCCL_DEBUG"] == "WARN"
+
+
+@pytest.mark.parametrize("rank", [0, 1])
+def test_collective_executes_all_ranks_but_only_output_rank_returns_tensors(rank):
+    actor = ray_module.RayDiffusionWorkerWrapper(rank)
+    value = torch.ones(2, requires_grad=True)
+    actor.worker = Mock()
+    actor.worker.execute_method.return_value = DiffusionOutput(output=value)
+    result = actor.execute_rpc("execute_model", output_rank=0, exec_all_ranks=True)
+    actor.worker.execute_method.assert_called_once()
+    assert result["replied"] == (rank == 0)
+    if rank == 0:
+        assert not result["result"].output.requires_grad
+        torch.testing.assert_close(result["result"].output, value)
+    else:
+        assert result["result"] is None
+
+
+def test_finalizer_kills_workers_without_waiting_and_releases_owned_group(monkeypatch):
+    events = []
+    fake_ray = Mock()
+    fake_ray.get.side_effect = lambda futures, **kw: events.append("wait")
+    fake_ray.kill.side_effect = lambda worker, **kw: events.append("kill")
+    fake_ray.util.remove_placement_group.side_effect = lambda group: events.append("remove")
+    monkeypatch.setattr(ray_module, "ray", fake_ray)
+    actor = Mock()
+    actor.shutdown.remote.side_effect = lambda: events.append("shutdown")
+    resources = ray_module._RayExecutorResources([ray_module._RayWorkerMetadata(actor)], object(), True)
+    resources()
+    assert events == ["kill", "remove"]
+    fake_ray.kill.assert_called_once_with(actor, no_restart=True)
+    actor.shutdown.remote.assert_not_called()
+    fake_ray.get.assert_not_called()
+    resources()
+    assert events == ["kill", "remove"]
+    assert resources.workers == []
+
+
+def test_wrapper_shutdown_releases_worker_once():
+    actor = ray_module.RayDiffusionWorkerWrapper(0)
+    worker = Mock()
+    actor.worker = worker
+    actor.shutdown()
+    actor.shutdown()
+    worker.shutdown.assert_called_once_with()
+    assert not actor.check_alive()
+
+
+def _make_executor(fake_ray, actors, group=None, owns_group=True):
+    executor = object.__new__(ray_module.RayDiffusionExecutor)
+    executor._closed = False
+    executor._is_failed = False
+    executor._state_lock = threading.Lock()
+    executor._failure_callbacks = []
+    executor.workers = [ray_module._RayWorkerMetadata(actor, rank=rank) for rank, actor in enumerate(actors)]
+    executor._resources = ray_module._RayExecutorResources(executor.workers, group, owns_group)
+    executor._finalizer = executor._resources
+    return executor
+
+
+@pytest.mark.parametrize("stuck_rank", [None, 1])
+def test_shutdown_drains_workers_before_killing_them(monkeypatch, stuck_rank):
+    events = []
+    actors = [Mock(), Mock()]
+    refs = []
+    for rank, actor in enumerate(actors):
+        ref = f"shutdown-{rank}"
+        refs.append(ref)
+        actor.shutdown.remote.side_effect = lambda ref=ref: events.append(ref) or ref
+    fake_ray = Mock()
+
+    def wait(pending, num_returns, timeout):
+        events.append("wait")
+        assert num_returns == len(pending)
+        assert timeout == ray_module._WORKER_SHUTDOWN_TIMEOUT_S
+        return [ref for ref in pending if ref != f"shutdown-{stuck_rank}"], [
+            ref for ref in pending if ref == f"shutdown-{stuck_rank}"
+        ]
+
+    fake_ray.wait.side_effect = wait
+    fake_ray.kill.side_effect = lambda worker, **kw: events.append("kill")
+    fake_ray.util.remove_placement_group.side_effect = lambda group: events.append("remove")
+    monkeypatch.setattr(ray_module, "ray", fake_ray)
+    executor = _make_executor(fake_ray, actors, group=object())
+    callback = Mock()
+    executor._failure_callbacks = [callback]
+
+    executor.shutdown()
+
+    assert events == [*refs, "wait", "kill", "kill", "remove"]
+    assert [call.args[0] for call in fake_ray.get.call_args_list] == [
+        ref for ref in refs if ref != f"shutdown-{stuck_rank}"
+    ]
+    assert executor._closed
+    assert not executor._is_failed
+    callback.assert_not_called()
+    executor.shutdown()
+    assert fake_ray.kill.call_count == len(actors)
+
+
+def test_shutdown_kills_workers_when_drain_fails(monkeypatch):
+    fake_ray = Mock()
+    fake_ray.wait.side_effect = RuntimeError("Ray is shutting down")
+    monkeypatch.setattr(ray_module, "ray", fake_ray)
+    actors = [Mock(), Mock()]
+    executor = _make_executor(fake_ray, actors)
+    executor.shutdown()
+    assert [call.args[0] for call in fake_ray.kill.call_args_list] == actors
+
+
+def test_shutdown_preserves_borrowed_placement_group(monkeypatch):
+    fake_ray = Mock()
+    monkeypatch.setattr(ray_module, "ray", fake_ray)
+    ray_module._RayExecutorResources([], object(), False)()
+    fake_ray.util.remove_placement_group.assert_not_called()
+
+
+@pytest.mark.parametrize("failure_kind", ["GetTimeoutError", "RayActorError", "RayTaskError"])
+@pytest.mark.parametrize("output_rank", [None, 0])
+def test_rpc_failure_kills_peers_and_rejects_subsequent_work(monkeypatch, failure_kind, output_rank):
+    fake_ray = Mock()
+    fake_ray.exceptions = SimpleNamespace(
+        **{name: type(name, (Exception,), {}) for name in ("GetTimeoutError", "RayActorError", "RayTaskError")}
+    )
+    failure = getattr(fake_ray.exceptions, failure_kind)("rank failed while peers are still running")
+    fake_ray.get.side_effect = failure
+    monkeypatch.setattr(ray_module, "ray", fake_ray)
+    actors = [Mock(), Mock()]
+    group = object()
+    executor = _make_executor(fake_ray, actors, group=group)
+    callback = Mock()
+    executor._failure_callbacks = [callback]
+
+    expected_error = TimeoutError if failure_kind == "GetTimeoutError" else ray_module.EngineDeadError
+    with pytest.raises(expected_error) as raised:
+        executor.collective_rpc("execute_model", unique_reply_rank=output_rank, exec_all_ranks=True)
+
+    assert raised.value.__cause__ is failure
+    assert executor.is_dead
+    assert executor._is_failed
+    assert executor._closed
+    callback.assert_called_once_with()
+    assert [call.args[0] for call in fake_ray.kill.call_args_list] == actors
+    fake_ray.util.remove_placement_group.assert_called_once_with(group)
+    # Peers may be stuck in a collective, so failure skips the graceful drain.
+    fake_ray.wait.assert_not_called()
+    for actor in actors:
+        actor.shutdown.remote.assert_not_called()
+
+    with pytest.raises(RuntimeError, match="closed"):
+        executor.collective_rpc("execute_model", unique_reply_rank=output_rank, exec_all_ranks=True)
+    for actor in actors:
+        actor.execute_rpc.remote.assert_called_once()
+    fake_ray.get.assert_called_once()
+    executor.shutdown()
+    assert fake_ray.kill.call_count == len(actors)
+    callback.assert_called_once_with()
+
+
+def test_actor_death_during_shutdown_does_not_report_failure():
+    executor = object.__new__(ray_module.RayDiffusionExecutor)
+    executor._closed = True
+    executor._is_failed = False
+    executor._state_lock = threading.Lock()
+    executor._finalizer = Mock()
+    callback = Mock()
+    executor._failure_callbacks = [callback]
+    executor._fail(RuntimeError("actor terminated"))
+    assert not executor._is_failed
+    callback.assert_not_called()
+    executor._finalizer.assert_not_called()
+
+
+class _InlineThread:
+    """Run the monitor synchronously so tests observe its effects."""
+
+    def __init__(self, target, **kwargs):
+        self._target = target
+
+    def start(self):
+        self._target()
+
+
+@pytest.mark.parametrize("closed_first", [False, True])
+def test_worker_monitor_fails_executor_when_idle_actor_dies(monkeypatch, closed_first):
+    fake_ray = Mock()
+    death = RuntimeError("actor died")
+    actors = [Mock(), Mock()]
+    sentinels = [actor.wait_for_exit.options.return_value.remote.return_value for actor in actors]
+
+    def wait(refs, num_returns):
+        assert refs == sentinels
+        assert num_returns == 1
+        if closed_first:
+            executor._closed = True
+        return [sentinels[1]], [sentinels[0]]
+
+    fake_ray.wait.side_effect = wait
+    fake_ray.get.side_effect = death
+    monkeypatch.setattr(ray_module, "ray", fake_ray)
+    monkeypatch.setattr(ray_module.threading, "Thread", _InlineThread)
+    group = object()
+    executor = _make_executor(fake_ray, actors, group=group)
+    callback = Mock()
+    executor.register_failure_callback(callback)
+
+    executor._start_worker_monitor()
+
+    for actor in actors:
+        actor.wait_for_exit.options.assert_called_once_with(concurrency_group="monitor")
+    if closed_first:
+        assert not executor._is_failed
+        callback.assert_not_called()
+        fake_ray.kill.assert_not_called()
+        return
+    fake_ray.get.assert_called_once_with(sentinels[1])
+    assert executor._is_failed
+    assert executor._closed
+    callback.assert_called_once_with()
+    assert [call.args[0] for call in fake_ray.kill.call_args_list] == actors
+    fake_ray.util.remove_placement_group.assert_called_once_with(group)
+    fake_ray.wait.assert_called_once()
+    with pytest.raises(RuntimeError, match="closed"):
+        executor.check_health()
+
+
+def test_worker_monitor_does_not_keep_executor_alive(monkeypatch):
+    fake_ray = Mock()
+    monkeypatch.setattr(ray_module, "ray", fake_ray)
+    started = []
+    monkeypatch.setattr(
+        ray_module.threading,
+        "Thread",
+        lambda target, **kwargs: started.append(target) or SimpleNamespace(start=lambda: None),
+    )
+    executor = _make_executor(fake_ray, [Mock()])
+    executor._start_worker_monitor()
+    (monitor,) = started
+    assert all(executor is not cell.cell_contents for cell in monitor.__closure__ or ())
+
+
+@pytest.mark.parametrize(
+    "bundles, required, error",
+    [
+        ([{"GPU": 1}] * 4, 4, None),
+        ([{"GPU": 1}] * 8, 4, None),
+        ([{"CPU": 1}, {"GPU": 0.5}, {"GPU": 1}], 1, None),
+        ([{"GPU": 1}] * 2, 4, "requires 4 full GPUs.*only 2"),
+        ([{"GPU": 0.5}] * 4, 4, "requires 4 full GPUs.*only 0"),
+        ([{"CPU": 4}], 1, "requires 1 full GPUs.*only 0"),
+        ([{"GPU": 4}], 4, None),
+        ([{"GPU": 2}, {"GPU": 2}], 4, None),
+        ([{"GPU": 1.5}, {"GPU": 1.5}], 3, "requires 3 full GPUs.*only 2"),
+    ],
+)
+def test_reused_placement_group_requires_full_gpu_capacity(monkeypatch, bundles, required, error):
+    group = SimpleNamespace(bundle_specs=bundles)
+    fake_ray = Mock()
+    fake_ray.util.get_current_placement_group.return_value = group
+    monkeypatch.setattr(ray_module, "ray", fake_ray)
+    executor = object.__new__(ray_module.RayDiffusionExecutor)
+    executor.od_config = SimpleNamespace(num_gpus=required)
+
+    if error is not None:
+        with pytest.raises(ValueError, match=error):
+            executor._get_or_create_placement_group()
+    else:
+        actual_group, owned = executor._get_or_create_placement_group()
+        assert actual_group is group
+        assert not owned
+
+    fake_ray.util.placement_group.assert_not_called()
+    fake_ray.util.remove_placement_group.assert_not_called()
+
+
+@pytest.mark.parametrize("failure_phase", ["creation", "discovery", "port", "initialization"])
+def test_worker_startup_failure_cleans_up_registered_actors(monkeypatch, failure_phase):
+    actors = [Mock(), Mock()]
+    for actor in actors:
+        actor.get_node_ip.remote.return_value = "10.0.0.1"
+        actor.get_open_port.remote.return_value = 23456
+    factory = Mock()
+    error = RuntimeError("startup failed")
+    factory.remote.side_effect = [actors[0], error] if failure_phase == "creation" else actors
+    fake_ray = Mock()
+    fake_ray.is_initialized.return_value = True
+    fake_ray.remote.return_value = lambda cls: factory
+    call_count = 0
+
+    def get(values, **kwargs):
+        nonlocal call_count
+        call_count += 1
+        assert kwargs["timeout"] == ray_module._WORKER_INIT_TIMEOUT_S
+        if call_count == {"discovery": 1, "port": 2, "initialization": 3}.get(failure_phase):
+            raise error
+        return values
+
+    fake_ray.get.side_effect = get
+    monkeypatch.setattr(ray_module, "ray", fake_ray)
+    monkeypatch.setattr(ray_module, "get_ip", lambda: "10.0.0.1")
+    monkeypatch.setattr(ray_module, "PlacementGroupSchedulingStrategy", Mock())
+    executor = object.__new__(ray_module.RayDiffusionExecutor)
+    executor.od_config = SimpleNamespace(num_gpus=2)
+    group = object()
+    executor._get_or_create_placement_group = Mock(return_value=(group, True))
+
+    with pytest.raises(RuntimeError, match="startup failed"):
+        executor._init_executor()
+
+    expected = actors[:1] if failure_phase == "creation" else actors
+    assert [call.args[0] for call in fake_ray.kill.call_args_list] == expected
+    for actor in expected:
+        actor.shutdown.remote.assert_not_called()
+        actor.wait_for_exit.options.assert_not_called()
+    fake_ray.util.remove_placement_group.assert_called_once_with(group)
+    assert executor.workers == []
+    assert executor._closed
+
+
+@pytest.mark.parametrize("failure_type", [TimeoutError, RuntimeError, KeyboardInterrupt])
+@pytest.mark.parametrize("cleanup_fails", [False, True])
+def test_placement_group_wait_failure_cleans_up_and_preserves_error(monkeypatch, failure_type, cleanup_fails):
+    fake_ray = Mock()
+    fake_ray.exceptions.GetTimeoutError = TimeoutError
+    fake_ray.util.get_current_placement_group.return_value = None
+    group = fake_ray.util.placement_group.return_value
+    original_error = failure_type("readiness failed")
+    fake_ray.get.side_effect = original_error
+    if cleanup_fails:
+        fake_ray.util.remove_placement_group.side_effect = RuntimeError("cleanup failed")
+    monkeypatch.setattr(ray_module, "ray", fake_ray)
+    executor = object.__new__(ray_module.RayDiffusionExecutor)
+    executor.od_config = SimpleNamespace(num_gpus=2)
+
+    expected_type = ValueError if failure_type is TimeoutError else failure_type
+    with pytest.raises(expected_type) as raised:
+        executor._get_or_create_placement_group()
+
+    if failure_type is TimeoutError:
+        assert raised.value.__cause__ is original_error
+        assert "Cannot reserve 2 Ray GPU bundle(s)" in str(raised.value)
+    else:
+        assert raised.value is original_error
+    fake_ray.util.remove_placement_group.assert_called_once_with(group)
+
+
+def test_health_probes_use_separate_concurrency_group(monkeypatch):
+    fake_ray = Mock()
+    fake_ray.get.return_value = [True, True]
+    monkeypatch.setattr(ray_module, "ray", fake_ray)
+    actors = [Mock(), Mock()]
+    executor = _make_executor(fake_ray, actors)
+
+    executor.check_health()
+
+    for actor in actors:
+        actor.check_alive.options.assert_called_once_with(concurrency_group="health")
+        actor.check_alive.options.return_value.remote.assert_called_once_with()
+        actor.check_alive.remote.assert_not_called()
+    fake_ray.get.assert_called_once_with(
+        [actor.check_alive.options.return_value.remote.return_value for actor in actors],
+        timeout=ray_module._HEALTH_CHECK_TIMEOUT_S,
+    )
+    assert not executor._is_failed
+
+
+@pytest.mark.parametrize("probe_result", ["error", "not_alive"])
+def test_health_failure_releases_actors_and_notifies(monkeypatch, probe_result):
+    fake_ray = Mock()
+    if probe_result == "error":
+        fake_ray.get.side_effect = RuntimeError("actor died")
+    else:
+        fake_ray.get.return_value = [True, False]
+    monkeypatch.setattr(ray_module, "ray", fake_ray)
+    actors = [Mock(), Mock()]
+    group = object()
+    executor = _make_executor(fake_ray, actors, group=group)
+    callback = Mock()
+    executor.register_failure_callback(callback)
+
+    with pytest.raises(ray_module.EngineDeadError):
+        executor.check_health()
+
+    assert executor._is_failed
+    assert executor._closed
+    callback.assert_called_once_with()
+    assert [call.args[0] for call in fake_ray.kill.call_args_list] == actors
+    fake_ray.util.remove_placement_group.assert_called_once_with(group)
+
+
+def _make_request_executor(monkeypatch, prompts, component="text_encoder", allgather=True):
+    executor = object.__new__(ray_module.RayDiffusionExecutor)
+    executor._closed = False
+    executor._is_failed = False
+    executor._state_lock = threading.Lock()
+    executor._failure_callbacks = []
+    executor._finalizer = Mock()
+    executor.od_config = SimpleNamespace(
+        parallel_config=SimpleNamespace(data_parallel_size=len(prompts)),
+        diffusion_offload_config={
+            "mode": "layer",
+            "components": [component],
+            "layer_options": {component: {"weight_transfer": "allgather" if allgather else "rank-local"}},
+        },
+    )
+    monkeypatch.setattr(ray_module, "build_request_batch_sampling_params_key", lambda req: ())
+    requests = []
+    for index, prompt in enumerate(prompts):
+        request_id = str(index)
+        req = SimpleNamespace(request_id=request_id, prompt=prompt, sampling_params=SimpleNamespace(extra_args=None))
+        requests.append(SimpleNamespace(request_id=request_id, req=req, diffusion_kv_metadata=None))
+    return executor, SimpleNamespace(scheduled_new_reqs=requests, kv_prefetch_job=None)
+
+
+@pytest.mark.parametrize("field", ["prompt_embeds", "negative_prompt_embeds"])
+def test_text_encoder_allgather_rejects_mismatched_embedding_presence(monkeypatch, field):
+    executor, scheduled = _make_request_executor(monkeypatch, ["first", {"prompt": "second", field: object()}])
+    executor.collective_rpc = Mock()
+    with pytest.raises(ValueError, match="same positive/negative prompt embedding fields"):
+        executor.execute_request(scheduled)
+    executor.collective_rpc.assert_not_called()
+
+
+def test_text_encoder_allgather_accepts_matching_precomputed_embedding_paths(monkeypatch):
+    executor, scheduled = _make_request_executor(
+        monkeypatch, [{"prompt_embeds": object()}, {"prompt_embeds": object()}]
+    )
+    results = [DiffusionOutput(output="first"), DiffusionOutput(output="second")]
+    executor.collective_rpc = Mock(return_value=[{"dp_rank": rank, "output": out} for rank, out in enumerate(results)])
+    output = executor.execute_request(scheduled)
+    assert [item.result for item in output.runner_outputs] == results
+    executor.collective_rpc.assert_called_once()
+
+
+@pytest.mark.parametrize("allgather", [False, True])
+def test_single_request_allgather_timeout_option(monkeypatch, allgather):
+    executor, scheduled = _make_request_executor(monkeypatch, ["prompt"], allgather=allgather)
+    executor.collective_rpc = Mock(return_value=DiffusionOutput(output="done"))
+    executor.execute_request(scheduled)
+    kwargs = executor.collective_rpc.call_args.kwargs
+    if allgather:
+        assert kwargs["timeout"] == ray_module._DLO_DP_WAVE_TIMEOUT_S
+    else:
+        assert "timeout" not in kwargs
+
+
+def test_single_request_allgather_timeout_fails_executor_and_cleans_up(monkeypatch):
+    executor, scheduled = _make_request_executor(monkeypatch, ["prompt"])
+    callback = Mock()
+    executor.register_failure_callback(callback)
+    executor.collective_rpc = Mock(side_effect=TimeoutError("AllGather timed out"))
+    output = executor.execute_request(scheduled)
+    assert executor._is_failed
+    assert executor._closed
+    executor._finalizer.assert_called_once()
+    callback.assert_called_once()
+    assert "AllGather timed out" in output.runner_outputs[0].result.error
+
+
+@pytest.mark.parametrize("wrapper", ["media", "output", "batch"])
+def test_move_to_cpu_handles_typed_video_media(wrapper):
+    from vllm_omni.diffusion.media import (
+        DiffusionMediaOutput,
+        VideoMediaOutput,
+        VideoTensorEncoding,
+        VideoTensorLayout,
+        VideoTensorSpec,
+        VideoValueRange,
+    )
+    from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput
+
+    tensor = torch.ones((1, 3, 2, 4, 4), requires_grad=True)
+    media = DiffusionMediaOutput(
+        video=VideoMediaOutput(
+            tensor=tensor,
+            spec=VideoTensorSpec(
+                layout=VideoTensorLayout.BCTHW,
+                encoding=VideoTensorEncoding.NORMALIZED_FLOAT,
+                value_range=VideoValueRange.ZERO_TO_ONE,
+            ),
+        ),
+        prepared_for_transport=True,
+    )
+    value = media
+    if wrapper != "media":
+        value = DiffusionOutput(media=media, to_cpu=False)
+    if wrapper == "batch":
+        value = BatchRunnerOutput.from_list([RunnerOutput(request_id="video", finished=True, result=value)])
+
+    moved = ray_module._move_to_cpu(value)
+    if wrapper == "batch":
+        moved = moved.runner_outputs[0].result
+    if wrapper != "media":
+        moved = moved.media
+
+    moved.validate()
+    assert moved.video.tensor.device.type == "cpu"
+    assert not moved.video.tensor.requires_grad
+    assert moved.video.spec == media.video.spec
+    assert moved.prepared_for_transport
+    torch.testing.assert_close(moved.video.tensor, tensor.detach())
+    assert media.video.tensor is tensor
+    assert tensor.requires_grad

tests/diffusion/test_stage_diffusion_proc.py

@@ -191,3 +191,62 @@ async def test_proc_process_request_with_batching_async_output():
         time_gap = elapsed_time - base_time
         assert time_gap > time_gap_std - eps and time_gap < time_gap_std + eps
         base_time = elapsed_time
+
+
+class _FailingExecutor:
+    def __init__(self, is_dead: bool = False) -> None:
+        self.is_dead = is_dead
+        self.callbacks = []
+
+    def register_failure_callback(self, callback) -> None:
+        self.callbacks.append(callback)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("failed_before_loop", [False, True])
+async def test_run_loop_exits_when_executor_fails_between_requests(failed_before_loop):
+    """An idle executor failure must reach the client without waiting for a request."""
+    import threading
+
+    import zmq
+    import zmq.asyncio
+
+    ctx = zmq.asyncio.Context()
+    request_socket = ctx.socket(zmq.PUSH)
+    response_socket = ctx.socket(zmq.PULL)
+    request_port = request_socket.bind_to_random_port("tcp://127.0.0.1")
+    response_port = response_socket.bind_to_random_port("tcp://127.0.0.1")
+    executor = _FailingExecutor(is_dead=failed_before_loop)
+    stage_proc = StageDiffusionProc.__new__(StageDiffusionProc)
+    stage_proc._engine = type("_Engine", (), {"executor": executor})()
+    stage_proc._active_tasks = None
+    stage_proc._fatal_event = None
+    try:
+        loop_task = asyncio.create_task(
+            stage_proc.run_loop(f"tcp://127.0.0.1:{request_port}", f"tcp://127.0.0.1:{response_port}")
+        )
+        if not failed_before_loop:
+            for _ in range(500):
+                if executor.callbacks:
+                    break
+                await asyncio.sleep(0.01)
+            assert executor.callbacks, "run_loop did not register an executor failure callback"
+            # Executor monitors report failures from their own threads.
+            monitor = threading.Thread(target=executor.callbacks[0])
+            monitor.start()
+            monitor.join()
+
+        with pytest.raises(RuntimeError, match="permanent failure"):
+            await asyncio.wait_for(loop_task, timeout=5)
+        assert await asyncio.wait_for(response_socket.recv(), timeout=5) == StageDiffusionProc.DIFFUSION_PROC_DEAD
+        assert len(executor.callbacks) == 1
+
+        # A late failure after run_loop has exited must not raise.
+        closed_loop = asyncio.new_event_loop()
+        closed_loop.close()
+        stage_proc._watch_executor_failure(closed_loop)
+        executor.callbacks[-1]()
+    finally:
+        request_socket.close(linger=0)
+        response_socket.close(linger=0)
+        ctx.term()

tests/diffusion/test_worker_wrapper_base.py

@@ -770,3 +770,60 @@ class TestCustomPipelineWorkerExtension:
         # Verify load_model was called twice with different pipelines
         assert wrapper.worker.load_model.call_count == 2
         assert wrapper.worker.init_lora_manager.call_count == 2
+
+
+@pytest.mark.parametrize(
+    ("init_method", "rank", "local_rank", "expected_addr", "expected_port"),
+    [
+        (None, 1, 1, "localhost", "12345"),
+        ("tcp://10.0.0.2:23456", 7, 0, "unrelated-host", "9999"),
+    ],
+)
+def test_worker_rendezvous_preserves_local_device_rank(
+    monkeypatch, init_method, rank, local_rank, expected_addr, expected_port
+):
+    """TCP rendezvous must not confuse a global rank with the actor's GPU index."""
+    import os
+    from contextlib import nullcontext
+    from unittest.mock import Mock
+
+    from vllm_omni.diffusion.worker import diffusion_worker as worker_module
+
+    class ReachedRendezvousError(Exception):
+        pass
+
+    worker = object.__new__(DiffusionWorker)
+    worker.rank = rank
+    worker.local_rank = local_rank
+    worker.distributed_init_method = init_method
+    worker.od_config = SimpleNamespace(num_gpus=8, master_port=12345, moe_backend="auto")
+    config = SimpleNamespace(kernel_config=SimpleNamespace(ir_op_priority=None))
+    platform = Mock()
+    platform.get_torch_device.return_value = "cuda:0"
+    rendezvous = Mock(side_effect=ReachedRendezvousError)
+    monkeypatch.setenv("MASTER_ADDR", "unrelated-host")
+    monkeypatch.setenv("MASTER_PORT", "9999")
+    # Restore these variables after init_device publishes the worker identity.
+    for key in ("LOCAL_RANK", "RANK", "WORLD_SIZE"):
+        monkeypatch.setenv(key, "stale")
+    monkeypatch.setattr(worker_module, "current_omni_platform", platform)
+    monkeypatch.setattr(worker_module, "create_diffusion_vllm_config", Mock(return_value=config))
+    monkeypatch.setattr(worker_module, "_resolve_ir_op_priority", Mock(return_value=[]))
+    monkeypatch.setattr(worker_module, "set_forward_context", lambda **kwargs: nullcontext())
+    monkeypatch.setattr(worker_module, "set_current_vllm_config", lambda config: nullcontext())
+    monkeypatch.setattr(worker_module, "init_distributed_environment", rendezvous)
+
+    with pytest.raises(ReachedRendezvousError):
+        worker.init_device()
+
+    rendezvous.assert_called_once_with(
+        world_size=8,
+        rank=rank,
+        distributed_init_method=init_method or "env://",
+        local_rank=local_rank,
+    )
+    platform.get_torch_device.assert_called_once_with(local_rank)
+    assert os.environ["MASTER_ADDR"] == expected_addr
+    assert os.environ["MASTER_PORT"] == expected_port
+    assert os.environ["LOCAL_RANK"] == str(local_rank)
+    assert os.environ["RANK"] == str(rank)

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -11,6 +11,7 @@ import pytest
 import torch
 from pytest_mock import MockerFixture
 from vllm.sampling_params import SamplingParams
+from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue
 from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
 from vllm.v1.metrics.stats import PrefillStats, PromptTokenStats
 from vllm.v1.request import Request, RequestStatus
@@ -635,6 +636,45 @@ def build_adapter(monkeypatch, mocker: MockerFixture):
     return _build
 
 
+@pytest.mark.parametrize("active_stream_window", [0, 2])
+def test_scheduler_chunk_gate_preserves_kv_holding_queue(build_adapter, active_stream_window):
+    adapter, _ = build_adapter(active_stream_window=active_stream_window)
+    scheduler = OmniARScheduler.__new__(OmniARScheduler)
+    scheduler.waiting = create_request_queue(SchedulingPolicy.FCFS)
+    scheduler.kv_holding_waiting = create_request_queue(SchedulingPolicy.FCFS)
+    scheduler.deferred_waiting = set()
+    scheduler.running = []
+    scheduler.chunk_transfer_adapter = adapter
+    scheduler.input_coordinator = None
+    holder = Request("holder", [1, 2, 3], SamplingParams(max_tokens=1), pooling_params=None)
+    holder.num_computed_tokens = 2
+    fresh = Request("fresh", [1], SamplingParams(max_tokens=1), pooling_params=None)
+    scheduler.requests = {request.request_id: request for request in (holder, fresh)}
+    scheduler._enqueue_waiting_request(fresh)
+    scheduler._enqueue_waiting_request(holder)
+
+    scheduler._process_pending_omni_inputs(model_mode="ar")
+
+    assert not scheduler.waiting and not scheduler.kv_holding_waiting
+    assert list(adapter.waiting_for_chunk_waiting_requests) == [holder, fresh]
+    assert holder.status == fresh.status == RequestStatus.WAITING_FOR_CHUNK
+    scheduler._restore_omni_wait_queues()
+    assert list(scheduler.kv_holding_waiting) == [holder]
+    assert list(scheduler.waiting) == [fresh]
+
+    adapter._finished_load_reqs.add(holder.request_id)
+    scheduler._process_pending_omni_inputs(model_mode="ar")
+
+    assert holder.status == RequestStatus.WAITING
+    assert fresh.status == RequestStatus.WAITING_FOR_CHUNK
+    assert list(scheduler.kv_holding_waiting) == [holder]
+    assert not scheduler.waiting
+    assert adapter.requests_with_ready_chunks == {holder.request_id}
+    scheduler._restore_omni_wait_queues()
+    assert list(scheduler.kv_holding_waiting) == [holder]
+    assert list(scheduler.waiting) == [fresh]
+
+
 def _dequeue_load_entry(adapter, request):
     """Register and take the entry as the background recv loop would."""
     adapter.load_async(request)
@@ -2445,7 +2485,8 @@ def test_finish_requests_reclaims_resumable_segment_and_reuses_capacity(
     scheduler.requests = {request.request_id: request}
     scheduler.running = [request] if placement == "running" else []
     scheduler.waiting = DummyWaitingQueue([request] if placement == "waiting" else [])
-    scheduler.skipped_waiting = DummyWaitingQueue([request] if placement.startswith("skipped") else [])
+    scheduler.kv_holding_waiting = DummyWaitingQueue([request] if placement.startswith("skipped") else [])
+    scheduler.deferred_waiting = set()
     scheduler.num_waiting_for_streaming_input = int(streaming_counter_owned)
 
     freed = []
@@ -2469,7 +2510,7 @@ def test_finish_requests_reclaims_resumable_segment_and_reuses_capacity(
     assert not adapter.waiting_for_chunk_running_requests
     assert scheduler.running == []
     assert list(scheduler.waiting) == []
-    assert list(scheduler.skipped_waiting) == []
+    assert list(scheduler.kv_holding_waiting) == []
     assert scheduler.num_waiting_for_streaming_input == 0
 
     fresh_a = _req("req-fresh-a", RequestStatus.WAITING)
@@ -2492,7 +2533,8 @@ def test_finish_requests_does_not_reopen_off_queue_resumable_segment(build_adapt
     scheduler.requests = {request.request_id: request}
     scheduler.running = []
     scheduler.waiting = DummyWaitingQueue()
-    scheduler.skipped_waiting = DummyWaitingQueue()
+    scheduler.kv_holding_waiting = DummyWaitingQueue()
+    scheduler.deferred_waiting = set()
     scheduler._free_request = lambda *args, **kwargs: pytest.fail("off-queue request was freed")
 
     assert (
@@ -2566,6 +2608,7 @@ def test_generation_scheduler_calls_cleanup_on_finished(monkeypatch, mocker: Moc
 
     scheduler = mocker.MagicMock()
     scheduler._first_chunk_express = False
+    scheduler.aux_output_connector = None
     scheduler._express_min_slack_s = 0.0
     scheduler.chunk_transfer_adapter = adapter_mock
     scheduler.connector = None
@@ -2646,6 +2689,7 @@ def test_ar_scheduler_defers_cleanup_and_queues_save_on_finished(mocker: MockerF
     adapter_mock.save_async = lambda *a, **kw: save_calls.append((a, kw))
 
     scheduler = mocker.MagicMock()
+    scheduler.aux_output_connector = None
     scheduler.chunk_transfer_adapter = adapter_mock
     bind_omits_transfer_helpers(scheduler)
     scheduler.connector = None
@@ -2744,6 +2788,8 @@ def test_omni_ar_scheduler_finish_requests(mocker: MockerFixture):
     sched.requests = {}
     sched.running = []
     sched.waiting = []
+    sched.kv_holding_waiting = []
+    sched.deferred_waiting = set()
 
     with patch.object(VLLMScheduler, "finish_requests", _super_finish):
         OmniARScheduler.finish_requests(sched, ["r1"], RequestStatus.FINISHED_ABORTED)
@@ -2759,7 +2805,8 @@ def _parked_sender_scheduler(adapter, session):
     scheduler.requests = {session.request_id: session}
     scheduler.running = []
     scheduler.waiting = DummyWaitingQueue()
-    scheduler.skipped_waiting = DummyWaitingQueue([session])
+    scheduler.kv_holding_waiting = DummyWaitingQueue([session])
+    scheduler.deferred_waiting = set()
     scheduler.num_waiting_for_streaming_input = 1
     scheduler.finished_req_ids_dict = {}
 

tests/e2e/accuracy/ltx/test_ltx_official_similarity.py

@@ -103,7 +103,7 @@ class LTXAccuracyCase:
     spatial_upsampler: LTXArtifact | None = None
     distilled_lora: LTXArtifact | None = None
     # Official uses block streaming; Omni uses layerwise offload. Both preserve
-    # bf16 arithmetic while keeping full-resolution two-stage cases on one H100.
+    # bf16 arithmetic while keeping full-resolution cases on one H100.
     enable_layerwise_offload: bool = False
 
 
@@ -182,6 +182,7 @@ TWO_STAGE_CASES = (
         gemma_model_id="Lightricks/LTX-2",
         gemma_model_revision=LTX2_REVISION,
         gemma_model_env="VLLM_TEST_LTX2_MODEL",
+        enable_layerwise_offload=True,
     ),
     LTXAccuracyCase(
         name="ltx23_distilled_i2v",

tests/e2e/accuracy/sana_video/conftest.py

@@ -0,0 +1,26 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+
+
+@pytest.fixture
+def sana_model_parallel(tmp_path):
+    """Construct native parallel layers with real single-rank GPU groups."""
+    from vllm_omni.diffusion.distributed.parallel_state import (
+        destroy_distributed_env,
+        init_distributed_environment,
+        initialize_model_parallel,
+    )
+
+    init_distributed_environment(
+        world_size=1,
+        rank=0,
+        local_rank=0,
+        distributed_init_method=f"file://{tmp_path / 'sana-distributed'}",
+    )
+    try:
+        initialize_model_parallel()
+        yield
+    finally:
+        destroy_distributed_env()

tests/e2e/accuracy/sana_video/test_sana_video_golden.py

@@ -202,6 +202,7 @@ def test_sana_video_transformer_matches_frozen_golden(
     variant: str,
     _hardware,
     accuracy_artifact_root: Path,
+    sana_model_parallel,
 ) -> None:
     del _hardware
     from vllm.config import LoadConfig

tests/e2e/accuracy/sana_video/test_sana_video_pipeline_alignment.py

@@ -189,6 +189,7 @@ def _cfg_prediction(
 def test_sana_video_quick_pipeline_stages_match_diffusers(
     variant: str,
     _hardware,
+    sana_model_parallel,
 ) -> None:
     """Compare the first denoising step without replacing native attention.
 

tests/e2e/online_serving/test_flux_kontext_expansion.py

@@ -23,7 +23,8 @@ BASE_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"})
 def _get_diffusion_feature_cases(model: str):
     return [
         pytest.param(
-            OmniServerParams(model=model),
+            # The unsharded transformer alone exceeds the L4's 22 GiB.
+            OmniServerParams(model=model, server_args=["--enable-layerwise-offload"]),
             id="base",
             marks=BASE_FEATURE_MARKS,
         ),

tests/e2e/online_serving/test_moss_tts_expansion.py

@@ -28,7 +28,7 @@ import pytest
 
 from tests.helpers.mark import hardware_test
 from tests.helpers.runtime import OmniServerParams
-from tests.helpers.stage_config import get_deploy_config_path
+from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
 
 pytestmark = [
     pytest.mark.advanced_model,
@@ -84,7 +84,12 @@ tts_server_params = [
     pytest.param(
         OmniServerParams(
             model=MODEL,
-            stage_config_path=get_deploy_config_path("moss_tts_realtime.yaml"),
+            stage_config_path=modify_stage_config(
+                get_deploy_config_path("moss_tts_realtime.yaml"),
+                # Leave room on the single L4 for the codec weights, reference
+                # encoder, and speech-serving activations in addition to KV.
+                updates={"stages": {0: {"gpu_memory_utilization": 0.35}}},
+            ),
             server_args=["--disable-log-stats"],
         ),
         id="moss_tts_realtime",

tests/e2e/online_serving/test_qwen2_5_omni_expansion.py

@@ -21,7 +21,14 @@ models = ["Qwen/Qwen2.5-Omni-7B"]
 
 # Single CI deploy YAML; rocm/xpu deltas are picked automatically via the
 # platforms: section in vllm_omni/deploy/ci/qwen2_5_omni.yaml.
-stage_configs = [modify_stage_config(get_deploy_config_path("ci/qwen2_5_omni.yaml"))]
+stage_configs = [
+    modify_stage_config(
+        get_deploy_config_path("ci/qwen2_5_omni.yaml"),
+        # This suite requests two cards. Share the Talker's card with
+        # Code2Wav and leave room for its whole-sequence DiT workspace.
+        updates={"stages": {1: {"gpu_memory_utilization": 0.3}, 2: {"devices": "1"}}},
+    )
+]
 
 # Create parameter combinations for model and stage config
 test_params = [

tests/engine/duplex/test_aura_resumable_hooks.py

@@ -370,6 +370,7 @@ def test_draining_empty_eos_completes_owning_response_before_shortcut() -> None:
             stage_port=port,
             model_state=SimpleNamespace(clear_continuation=lambda: None),
             services=SimpleNamespace(spawn=lambda *_a, **_k: None),
+            run=SimpleNamespace(closing=False),
         ),
         out,
         close_from_runtime=_noop,
@@ -461,6 +462,7 @@ def test_draining_completion_drops_finished_response_books_only() -> None:
             stage_port=_Port(),
             model_state=SimpleNamespace(clear_continuation=lambda: None),
             services=SimpleNamespace(spawn=lambda *_a, **_k: None),
+            run=SimpleNamespace(closing=False),
         ),
         out,
         close_from_runtime=_noop,

tests/engine/duplex/test_delivery.py

@@ -0,0 +1,123 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import asyncio
+import json
+from dataclasses import replace
+
+import pytest
+
+from vllm_omni.engine.duplex.delivery import DuplexOutputBuffer, DuplexOutputOverflowError
+from vllm_omni.engine.duplex.events import AudioDelta, ResponseDone, SessionClosed, TranscriptDelta
+from vllm_omni.engine.duplex.realtime_events import RealtimeProjectionState, project_internal_event
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("event_type", "delta_field", "content_field"),
+    [
+        ("response.text.delta", "delta", "text"),
+        ("response.output_audio.delta", "text", "transcript"),
+    ],
+)
+async def test_queued_response_creation_keeps_payload_and_byte_count(event_type, delta_field, content_field):
+    state = RealtimeProjectionState(session_id="s")
+    events = project_internal_event(state, {"type": "response.created", "response_id": "r", "modalities": ["text"]})
+    payloads = [event.to_realtime() for event in events]
+    byte_count = sum(
+        len(json.dumps(payload, ensure_ascii=False, separators=(",", ":")).encode()) for payload in payloads
+    )
+    output = DuplexOutputBuffer(max_bytes=byte_count, max_events=len(events))
+    for event in events:
+        output.put(event)
+    assert output.pending_bytes == byte_count
+
+    text = "Growing response. " * 512
+    project_internal_event(state, {"type": event_type, "response_id": "r", delta_field: text})
+    item = next(iter(state.conversation_items.values()))
+    assert item["content"][0][content_field] == text
+    assert output.pending_bytes == byte_count
+    delivered = [await output.get() for _ in events]
+    assert [event.to_realtime() for event in delivered] == payloads
+    assert output.pending_events == output.pending_bytes == 0
+
+
+@pytest.mark.asyncio
+async def test_cancel_removes_only_matching_audio_and_invalidates_held_audio():
+    output = DuplexOutputBuffer(max_bytes=4096, max_events=8)
+    held = AudioDelta(response_id="old", epoch=0, delta="AAAA")
+    text = TranscriptDelta(response_id="old", epoch=0, delta="kept")
+    other = replace(held, response_id="other")
+    newer = replace(held, epoch=1)
+    done = ResponseDone(response_id="old", response={"status": "cancelled"})
+    for event in (held, replace(held), text, other, newer, done):
+        output.put(event)
+    assert await output.get() is held
+    assert output.is_valid(held)
+    output.invalidate("old", through_epoch=0)
+    assert output.pending_events == 4
+    with output.guard(held) as valid:
+        assert not valid
+    assert [await output.get() for _ in range(4)] == [text, other, newer, done]
+    assert output.pending_events == output.pending_bytes == 0
+
+
+@pytest.mark.asyncio
+async def test_full_reserve_cannot_block_closure_or_reorder_valid_audio():
+    output = DuplexOutputBuffer(max_bytes=4096, max_events=1, reserve_events=1)
+    audio = AudioDelta(response_id="r", epoch=0, delta="AAAA")
+    done = ResponseDone(response_id="r", response={"status": "completed"})
+    output.put(audio)
+    output.put(done)
+    with pytest.raises(DuplexOutputOverflowError):
+        output.put(done)
+    closed = SessionClosed(session_id="s", reason="client_close")
+    output.close(closed)
+    output.close(SessionClosed(reason="duplicate"))
+    assert not output.put(audio)
+    assert [await output.get() for _ in range(4)] == [audio, done, closed, None]
+    assert output.pending_events == output.pending_bytes == 0
+
+
+@pytest.mark.asyncio
+async def test_oversized_close_still_ends_stream_without_reopening_it():
+    output = DuplexOutputBuffer(max_bytes=128, max_events=1)
+    closed = SessionClosed(session_id="s", details={"large": "x" * 5000})
+    output.close(closed)
+    terminal = await output.get()
+    assert isinstance(terminal, SessionClosed)
+    assert terminal.event_id == closed.event_id
+    assert terminal.session_id == "s"
+    assert terminal.reason == "close_details_exceed_output_limit"
+    assert terminal.details == {}
+    assert await output.get() is None
+    output.close(closed)
+    assert await output.get() is None
+
+
+def test_byte_limit_rejects_audio_without_changing_pending_budget():
+    output = DuplexOutputBuffer(max_bytes=128, max_events=8)
+    with pytest.raises(DuplexOutputOverflowError):
+        output.put(AudioDelta(delta="AAAA" * 64))
+    assert output.pending_events == output.pending_bytes == 0
+
+
+@pytest.mark.asyncio
+async def test_cancelled_waiter_can_be_replaced_and_woken_from_another_thread():
+    output = DuplexOutputBuffer(max_bytes=4096, max_events=8)
+    abandoned = asyncio.create_task(output.get())
+    await asyncio.sleep(0)
+    abandoned.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await abandoned
+    waiting = asyncio.create_task(output.get())
+    await asyncio.sleep(0)
+    event = AudioDelta(response_id="r", epoch=0, delta="AAAA")
+    await asyncio.to_thread(output.put, event)
+    assert await asyncio.wait_for(waiting, timeout=2) is event
+    closing = asyncio.create_task(output.get())
+    await asyncio.sleep(0)
+    await asyncio.to_thread(output.close)
+    assert await asyncio.wait_for(closing, timeout=2) is None