[Bugfix] Preserve host transport fallback when discovery cleanup fails

Signed-off-by: Sy03 <1370724210@qq.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/ci_source_file_dependencies.yml

@@ -16,6 +16,9 @@ x-qwen3-tts: &qwen3_tts
   - vllm_omni/worker/gpu_model_runner.py
   - vllm_omni/model_executor/models/qwen3_tts/
   - vllm_omni/model_executor/models/common/qwen3_code_predictor.py
+  - vllm_omni/model_executor/models/common/fused_sampling.py
+  - vllm_omni/model_executor/models/common/short_kv_attention.py
+  - vllm_omni/model_executor/models/common/talker_first_audio.py
   - vllm_omni/model_executor/models/common/snake_activation.py
   - vllm_omni/model_executor/models/common/alias_free_activation.py
   - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
@@ -85,11 +88,16 @@ x-qwen3-omni: &qwen3_omni
   - vllm_omni/entrypoints/openai/serving_chat.py
   - vllm_omni/model_executor/models/qwen3_omni/
   - vllm_omni/model_executor/models/common/qwen3_code_predictor.py
+  - vllm_omni/model_executor/models/common/fused_sampling.py
+  - vllm_omni/model_executor/models/common/short_kv_attention.py
+  - vllm_omni/model_executor/models/common/talker_first_audio.py
   - vllm_omni/model_executor/models/common/snake_activation.py
   - vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py
   - vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py
   - vllm_omni/model_executor/stage_input_processors/qwen3_omni.py
   - vllm_omni/deploy/qwen3_omni_moe.yaml
+  - vllm_omni/deploy/qwen3_omni_moe_mrv2.yaml
+  - vllm_omni/deploy/qwen3_omni_moe_mrv2_h200.yaml
   - vllm_omni/deploy/qwen3_omni_moe_thinking.yaml
 
 x-wan22: &wan22

.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:
@@ -124,6 +135,30 @@ steps:
           - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py tests/e2e/offline_inference/test_qwen3_omni_colocate_async.py -m 'advanced_model and cuda' --run-level 'advanced_model'
         mirror_hardwares: h100_2
 
+      - label: "Omni · Qwen3-Omni Async-Chunk SHM Leak"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_qwen3_omni_shm_leak.py
+          - tests/helpers/runtime.py
+          - tests/helpers/stage_config.py
+          - vllm_omni/core/sched/omni_ar_scheduler.py
+          - vllm_omni/distributed/omni_connectors/connectors/shm_connector.py
+          - vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py
+          - vllm_omni/engine/orchestrator.py
+          - vllm_omni/engine/stage_engine_core_proc.py
+          - vllm_omni/engine/stage_pool.py
+          - vllm_omni/entrypoints/openai/serving_chat.py
+          - vllm_omni/model_executor/models/qwen3_omni/
+          - vllm_omni/model_executor/models/common/qwen3_code_predictor.py
+          - vllm_omni/model_executor/models/common/snake_activation.py
+          - vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py
+          - vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py
+          - vllm_omni/model_executor/stage_input_processors/qwen3_omni.py
+          - vllm_omni/deploy/qwen3_omni_moe.yaml
+        timeout_in_minutes: 40
+        commands:
+          - pytest -s -v tests/e2e/online_serving/test_qwen3_omni_shm_leak.py -m 'advanced_model and cuda' --run-level 'advanced_model'
+        mirror_hardwares: h100_2
+
       - label: "Omni · MiniCPM-o 4.5 Test"
         source_file_dependencies: omni_minicpmo_4_5_function
         timeout_in_minutes: 50

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

.buildkite/cuda/test-ready.yml

@@ -25,6 +25,16 @@ steps:
           - "pytest -sv tests/model_executor/models/cosyvoice3/test_cosyvoice3_full_batch.py -k packed -m 'core_model and cuda' --run-level=core_model"
         mirror_hardwares: h100_1
 
+      - label: "Simple · MRV2 FA3 Graph Bounds Test"
+        timeout_in_minutes: 15
+        source_file_dependencies:
+          - tests/worker_v2/test_capture_query_bound_gpu.py
+          - tests/worker_v2/test_decode_query_bound_gpu.py
+          - vllm_omni/worker_v2/model_states/omni_model_state.py
+        commands:
+          - "pytest -sv tests/worker_v2/test_capture_query_bound_gpu.py tests/worker_v2/test_decode_query_bound_gpu.py -m 'core_model and cuda and H100' --run-level=core_model"
+        mirror_hardwares: h100_1
+
       - label: "Simple · Diffusion Test"
         timeout_in_minutes: 30
         commands:

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.
 
@@ -214,7 +214,7 @@ collection, examples, bug reports, or logs.
 
 ## Model-specific variables
 
-The inventory includes 66 variables read by a single model or pipeline family. They are
+The inventory includes 93 variables read by a single model or pipeline family. They are
 not listed as public usage options here because doing so would turn implementation
 escape hatches into an accidental compatibility contract.
 
@@ -223,10 +223,10 @@ Every audited model-specific name has a migration disposition in the
 
 | Disposition | Count | Required outcome |
 | --- | ---: | --- |
-| Promote | 38 | Move a stable setting into typed stage or model configuration. |
+| Promote | 64 | Move a stable setting into typed stage or model configuration. |
 | Request scope | 5 | Move request-varying behavior into a declared request-option schema. |
 | External | 0 | Retain only when a supported third-party library owns the contract. |
-| Internalize | 16 | Keep a debug or diagnostic switch out of public documentation and configuration. |
+| Internalize | 17 | Keep a debug or diagnostic switch out of public documentation and configuration. |
 | Deprecate/remove | 7 | Remove a compatibility escape hatch that has no continuing contract. |
 
 The disposition is a migration target, not a statement that the existing

docs/configuration/hsdp_pre_sharded_loading.md

@@ -0,0 +1,22 @@
+# Pre-sharded HSDP loading
+
+Use `--use-hsdp --hsdp-shard-size 8 --hsdp-weight-load-strategy pre_sharded`
+to shard transformer parameters before loading checkpoint values. Set the shard
+size to the available GPU count (or configure replication explicitly).
+The default strategy, `full`, retains the existing load-then-shard behavior.
+The same `hsdp_weight_load_strategy` option is available through Python and stage
+engine configuration.
+
+The pre-sharded strategy reads rank-local slices from ordinary Hugging Face
+safetensors checkpoints into FSDP-owned storage. No converted checkpoint is
+required. Cosmos 3 and Cosmos 3 Edge transformer blocks release their initialized
+parameter storage to meta as they are constructed; nonpersistent buffers retain
+their values.
+Encoders and VAEs follow their existing loading paths.
+
+This strategy requires the default diffusion pipeline loader, complete dedicated
+transformer checkpoint sources, and supported checkpoint key mappings. Tensor
+layouts must match the runtime layout; tensor transforms are unsupported.
+Checkpoint tensor dtypes must match runtime tensor dtypes at preflight; dtype conversion is unsupported.
+Quantization and LoRA are unsupported. Every worker needs checkpoint access.
+`num_weight_load_threads` controls the safetensors reader's threads per rank.

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
 
@@ -319,7 +400,9 @@ ROCm and MUSA. Select one of these deployment profiles:
 | `qwen3_tts_mrv2.yaml` | V2 | B1 | Explicit MRV2 profile (same runner selection as the default) |
 | `qwen3_tts_high_concurrency_mrv2.yaml` | V2 | B1, B2 | Opt-in throughput tuning |
 | `qwen3_tts_high_concurrency_mrv2_b4.yaml` | V2 | B1, B2, B3, B4 | Experimental throughput / buffered playback |
+| `qwen3_tts_high_concurrency_mrv2_single_gpu.yaml` | V2 | B1–B8 | Experimental two-stage deployment on one GPU, with MPS |
 | `qwen3_tts_high_concurrency.yaml` | V1 | Existing defaults | V1 high-concurrency control |
+| `qwen3_tts_fused_single_gpu.yaml` | V2 | N/A (in-Talker decoder) | Opt-in single-stage CUDA pipeline on one GPU |
 
 ```bash
 # MRV2 is the default; pass a copy with `model_runner: v1` to force V1.
@@ -350,12 +433,73 @@ supply `mtp_sampling_params` and `get_mtp_seed(sampling_params)` for model-local
 (with `mtp_sample_steps` and `mtp_sample_vocab_size`) as needed. Qwen3-TTS retains
 its existing `talker_mtp` entry point for V1.
 
+### Selecting one or two Qwen3-TTS stages
+
+Choose the deployment profile explicitly with `--deploy-config`. The default
+`qwen3_tts` pipeline retains two stages: the Talker publishes codec chunks to
+Code2Wav, which produces audio. Its stage devices can be placed on the same
+GPU or on separate GPUs. `qwen3_tts_fused` is a separate, opt-in pipeline for
+one CUDA GPU: the Talker owns the stateful codec decoder and emits PCM directly,
+without a Code2Wav engine or stage connector.
+
+```bash
+# Two stages on one GPU; the two-stage pipeline remains independently selectable.
+vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni \
+  --deploy-config vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2_single_gpu.yaml
+
+# One stage on one GPU; no inter-stage transport or Code2Wav scheduling.
+vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni \
+  --deploy-config vllm_omni/deploy/qwen3_tts_fused_single_gpu.yaml
+```
+
+The single-stage profile requires CUDA, V2, asynchronous chunks, TP=PP=1,
+an in-process worker and disabled prefix caching. Select the full profile;
+adding `talker_stream_decode: true` to a two-stage Talker is rejected because
+its latent output belongs to the Code2Wav input contract. Single-stage options
+live in `stages[].additional_config`; two-stage transport and first-frame
+options live in `connectors.*.extra`.
+
+The single-stage profile defaults `talker_stream_first_audio` to `false` for
+throughput. To opt into a separate first-frame delivery queue, save this
+overlay beside the bundled deployment files and pass its path to
+`--deploy-config`:
+
+```yaml
+base_config: qwen3_tts_fused_single_gpu.yaml
+stages:
+  - stage_id: 0
+    additional_config:
+      talker_stream_decode: true
+      talker_stream_first_audio: true
+      ref_code_context_frames: 72
+      code_predictor_kv_cache: true
+      code_predictor_fused_sampling: true
+      code_predictor_fused: true
+```
+
+This queue uses PCM from the single-stage decoder. The two-stage
+`talker_first_audio` option below uses its own first-frame decoder and does
+not require single-stage execution. The two options cannot be enabled
+together. Single-stage prefix-cache reset with
+`reset_running_requests=True` returns `False` while requests are active;
+wait for completion or abort them before retrying.
+
+Reference-cloning context is selected independently by each profile. The
+single-stage profile explicitly uses 72 reference frames, matching the
+`qwen3_tts.yaml` context boundary. A shorter
+`additional_config.ref_code_context_frames` (for example 25) can change
+timbre continuity even when WER is similar;
+the throughput measurements for CustomVoice do not qualify Base voice
+cloning quality. Reference layout, grouping and priming are model hooks;
+the shared worker does not select Qwen's context length or codebook layout.
+
 ### Qwen3-TTS first-frame delivery and rollback
 
 The standard `qwen3_tts.yaml`, `qwen3_tts_high_concurrency.yaml`,
 `qwen3_tts_mrv2.yaml` and `qwen3_tts_high_concurrency_mrv2.yaml` profiles
 enable the `talker_first_audio` connector option, as does the experimental
-single-GPU profile. This changes the default CUDA MRV2 streaming path:
+two-stage `qwen3_tts_high_concurrency_mrv2_single_gpu.yaml` profile.
+This changes the default CUDA MRV2 streaming path:
 residual prediction runs eagerly after the Talker sample, the Talker loads
 an additional first-frame decoder with its weights and CUDA graphs, and the
 orchestrator orders its audio before subsequent Code2Wav chunks.
@@ -451,3 +595,50 @@ for that comparison; this switch only controls automatic MPS management.
 
 MPS does not reserve a GPU. Use only assigned GPUs, explicitly place stages on
 the intended GPU, and warm the complete pipeline before measuring performance.
+
+## MOSS-TTS Local 1.5 with Model Runner V2
+
+`OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5` can opt into CUDA MRV2 using
+the shared runtime introduced for Qwen3-TTS:
+
+```bash
+vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 --omni \
+  --deploy-config vllm_omni/deploy/moss_tts_local_mrv2.yaml
+```
+
+This profile inherits the batching, codec graph buckets and 1-frame/15-frame
+chunk geometry from `moss_tts_local.yaml`, and selects V2 for both the Local
+Talker and codec stages. The Local depth predictor exposes the MRV2 `mtp`
+capabilities while retaining its V1 `talker_mtp` implementation, sampling
+defaults and explicit request-seed handling. Codec chunks use the native data
+plane; the internal Talker retains its final-only orchestrator output policy.
+
+On CUDA, the profile bounds stage-0 prefill work to 512 tokens per iteration.
+The codec retains its own CUDA Graph capture/replay, including batch buckets
+through 64, but disables Inductor compilation (`compilation_config.mode: 0`)
+of the stateful decoder to avoid lengthy compilation at startup. This is
+independent of `enforce_eager`; the profile keeps `enforce_eager: false`.
+It explicitly selects `cudagraph_mode: FULL`, which works without Inductor;
+the default `FULL_AND_PIECEWISE` would otherwise normalize to `NONE` and clear
+the codec's capture buckets when compilation is disabled.
+
+For sustained C128 serving on a large-memory CUDA GPU, select the
+separate `moss_tts_local_mrv2_high_concurrency.yaml` profile. It uses 128
+stream slots per stage, a 32 GiB Talker KV budget, and the codec's
+`triton_slot` backend with Inductor compilation and codec-owned CUDA graphs.
+The bounded KV budget leaves room for codec state and graphs; it does not
+guarantee that 128 maximum-length prompts fit simultaneously. See the
+[MOSS recipe](gh-file:recipes/OpenMOSS/MOSS-TTS.md#local-15-mrv2-and-slot-attention)
+for activation, backend comparisons, memory requirements and benchmark commands.
+
+Omitting `--deploy-config`, or selecting `moss_tts_local.yaml`, retains V1.
+NPU, XPU, ROCm and MUSA overrides also retain V1. This profile does not enable
+MRV2 for MOSS Delay, Realtime or Nano. Local 1.5 outputs 48 kHz stereo audio;
+set `VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE=48000` and
+`VLLM_OMNI_BENCH_AUDIO_CHANNELS=2` when benchmarking raw PCM.
+
+Event-driven orchestration remains independently selectable with
+`VLLM_OMNI_EVENT_DRIVEN_ORCH=0` or `1`. Keep the runner and deployment identical
+when comparing these modes. Model-runner selection does not change the
+orchestration default or enable experimental reference encoding, chunk ramps,
+generation-output draining or MPS.

docs/design/feature/attention_backend_selection.md

@@ -4,6 +4,9 @@ This document defines the selection and extension contract for diffusion
 attention backends. User-facing backend choices, installation, and tuning are
 in the [attention backend guides](../../user_guide/diffusion/attention_backends.md).
 
+The [execution contract PoC](attention_execution_contract_poc.md) describes the
+experimental post-selection capability and lifecycle contract from RFC #7226.
+
 ## Scope
 
 The contract applies to DiT and other diffusion attention layers. It is

docs/design/feature/attention_execution_contract_poc.md

@@ -0,0 +1,154 @@
+# Diffusion attention execution contract PoC
+
+This prototype accompanies [RFC #7226](https://github.com/vllm-project/vllm-omni/issues/7226)
+and extends the [attention selection design](attention_backend_selection.md).
+It demonstrates path-specific capabilities with dense BF16 FA4, NPU/ROCm
+routing examples, and tensor-state lifetime with a test-only attention module.
+
+## Execution contract
+
+`ExecutionContext` describes the requested execution path. `ExecutionPathResult`
+reports its identity, support status and reason, and compilation mode.
+`AttentionBackend.resolve_capabilities()` provides conservative pre-construction
+results. After initialization, `Attention.resolve_execution_path()` supplies the
+active parallel, paged-KV, and HSDP context; the backend combines it with the
+selected kernel and normalized metadata.
+
+Resolution runs outside compiled execution. FA4 reports `SUPPORTED` and
+`CUSTOM_OP` for dense, noncausal BF16 without parallel or HSDP boundaries when
+its kernel accepts the head dimensions. Dimension validation delegates to FA4's
+architecture-specific rules through `backends/utils/fa.py`. Kernel rejections
+report `UNSUPPORTED` with an actionable reason. Missing private validators and
+other unmigrated paths report `UNMIGRATED` with advisory `EAGER_ONLY` defaults.
+`requested_support()` checks a fullgraph request; it does not enforce selection
+or change existing execution.
+
+FA4 execution uses an opaque custom op. Its fake output preserves Q's batch,
+sequence, and head count and uses V's head dimension. Output is contiguous even
+when Q is noncontiguous. Metadata normalization is
+shared with dispatch. Producers must update published mask semantics when masks
+change; unpublished masks remain runtime-dependent to avoid synchronization.
+Callers resolve again when execution metadata changes.
+
+## NPU and ROCm worked examples
+
+The same `FLASH_ATTN` backend resolves MindIE paths on NPU and AITER paths on
+ROCm. Both examples report `UNMIGRATED`: their route is inspected, while hardware
+correctness and compilation remain unvalidated. `EAGER_ONLY` is an advisory
+default, not a measured compiler limitation. Existing execution is unchanged.
+
+| Input | NPU / MindIE | ROCm / AITER |
+| --- | --- | --- |
+| Dense | `npu_dense` | `rocm_dense` |
+| Padding mask | `npu_masked`: expand the key mask to `[B, 1, Q, K]` | `rocm_masked_varlen`: unpad, call varlen attention, then restore rows |
+| Packed inputs | Opt-in `[real, pad]` contract; choose varlen or prefix-KV slicing from the environment | Forward complete packed-document boundaries to varlen attention |
+| Incomplete packing | Rebuild a mask when possible; otherwise reject before the kernel | Shared CUDA-like normalization rejects incomplete metadata |
+| Unpublished mask semantics | Route to the masked wrapper | Keep resolution runtime-dependent without reading mask values |
+
+NPU resolution reuses `_resolve_packed_seq_npu`; ROCm reuses the metadata
+normalization used by dispatch. AITER identity comes from initialization rather
+than caller-provided context. Quantized, paged, piecewise, and parallel paths
+remain outside these worked examples.
+
+The new L1 tests use CPU tensors and replace only vendor kernel calls. They
+check path identity, forwarded masks/boundaries, fallback rejection, and
+conservative fullgraph-request handling. They do not test vendor numerics or
+compiler support. Existing L1 CI already collects this test file.
+
+```bash
+python -m pytest tests/diffusion/attention/test_flash_attn.py \
+  -k 'npu_contract or rocm_contract' -m 'core_model and cpu' -q
+```
+
+### Hardware validation required before migration
+
+On an Ascend/MindIE or ROCm/AITER runner, record device, driver, PyTorch, vendor
+library, and compiler versions, then validate the corresponding rows above:
+
+1. Compare eager outputs with FP32 SDPA for dense and padding-mask inputs,
+   including unequal Q/K lengths. Compare valid query rows; do not assume both
+   platforms define padded-query outputs identically.
+2. Compare packed outputs with separate per-document references. On NPU, cover
+   both varlen and laser prefix-slicing modes, plus mask reconstruction and the
+   missing-mask error. A padding fallback does not establish multi-document
+   isolation. On ROCm, cover multiple real documents.
+3. Check output shape, dtype, device, finite values, and unchanged inputs.
+   Start BF16 comparisons at `atol=rtol=1e-2` and have the platform maintainer
+   confirm the tolerance against its reference tests.
+4. Run `torch.compile(fullgraph=True, dynamic=True)` with the platform's supported
+   compiler over multiple sequence lengths. Record graph breaks and recompiles;
+   choose `TRACEABLE`, `CUSTOM_OP`, or verified `EAGER_ONLY` from that evidence.
+   A custom-op path additionally needs schema/fake checks and repeated replay.
+
+Promote only the paths validated on that platform. These hardware checks have
+not been run locally; ROCm/NPU fullgraph support is not claimed.
+
+## State lifetime example
+
+`test_attention_state_lifetime.py` prepares an owned query-scale tensor eagerly
+and passes it explicitly to an opaque attention op. The compiled module retains
+its state after the caller drops its reference and releases ownership when the
+compiled callable is deleted. Equivalent instances reuse a graph with their own
+state values.
+
+This example tests ownership and release, dynamic replay across instances,
+custom-op schema/fake behavior, and Inductor execution. It adds no production
+backend or state registry. Planning keys, caching, fallback-policy declarations,
+opaque handles, CUDA graph capture, and TRTLLM integration are deferred.
+
+## Validation
+
+```bash
+python -m pytest tests/diffusion/attention/test_flash_attn_compile.py \
+  tests/diffusion/attention/test_flash_attn.py \
+  tests/diffusion/attention/test_attention_capabilities.py \
+  tests/diffusion/attention/test_attention_state_lifetime.py -q -rs
+```
+
+The focused suite covers capability decisions and metadata changes, existing
+padding/mask regressions, tensor-state lifetime, and real FA4 fullgraph execution
+against FP32 SDPA and eager FA4. Representative Q/K and V dimensions are
+(32, 32), (64, 64), (80, 48), (192, 128), and (256, 256), with batches 1 and 2
+and sequence lengths up to 1,024. Invalid dimensions are checked against actual
+kernel errors. These samples exercise the contract rather than define support.
+BF16 comparisons use `atol=rtol=1e-2`; `torch.library.opcheck` checks schema and
+fake-output correctness.
+
+Single-graph reuse is asserted in the tensor-state test. The FA4 numerical tests
+retain compiler caches across shapes within each case but do not require one
+graph across different batch sizes, Q/K length equality, or head dimensions.
+
+Tested environment: GB300, PyTorch `2.13.0+cu130`, FA4 `4.0.0b18`, CUTLASS DSL
+`4.6.2` with CUDA 13 libraries, Quack `0.6.4`, and TVM FFI `0.1.11`.
+Real-kernel tests require CUDA and CuTe FA4; their availability is checked at runtime.
+
+Before the rebase, the focused suite passed all 97 tests without skips, including
+11 CPU NPU/ROCm routing tests. After rebasing onto upstream `e3be42e05`, the
+expanded suite reports 99 passed and 11 failed in the local vLLM 0.28 environment.
+All 11 failures are import errors: upstream now requires `compute_layout_strides`
+from vLLM, which this environment lacks. Both noncontiguous FA4 regression cases
+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/feature/stage_host_transport.md

@@ -0,0 +1,84 @@
+# Local stage host transport
+
+`SharedMemoryConnector` uses a persistent host ring for ordinary local stage
+payloads by default. Models keep the existing `put`/`get` and typed payload
+contracts; no model-specific sender, receiver or deployment change is required.
+This applies to both V1 and MRV2 paths using this connector, including CPU codec
+codes and CPU hidden states. Remote connectors use their existing transports.
+
+Each producer process and destination stage lazily allocates a 4 MiB ring.
+Small CPU tensor trees use a msgpack header and raw tensor regions. Plain
+serialized payloads can use the same ring. Frames larger than a quarter of the
+ring, and frames arriving while it is full, use the existing per-key SHM path.
+The fallback preserves delivery without blocking a model step on ring capacity.
+Large or unsupported tensor trees do not allocate an unused contiguous snapshot.
+
+## Ownership and delivery
+
+The producer copies the payload before `put` returns. Readers hold the ring's
+file lock while copying into independently owned CPU storage, then mark the
+frame consumed. Multiple receiving processes may compete for a key; one reader
+claims it. Existing TP fanout remains responsible for distributing that owned
+payload to peer ranks. This transport does not provide multicast.
+
+Small tensor leaves share one owned allocation. Large multi-leaf messages use
+separate allocations, so retaining small metadata does not retain a large
+hidden state. No view of ring storage escapes the receive boundary.
+
+Cancellation retires exact keys or tracked numeric chunk suffixes. Replacing
+an unclaimed key retires its previous frame, including transitions between
+the ring and per-key SHM. Descriptors include monotonic frame positions; an old
+descriptor cannot read a newer publication after wraparound. Keys must have
+one active publisher, as assumed by the existing stage request routing.
+
+Only producers register ring allocations with Python's shared-memory resource
+tracker. Readers map files without taking ownership. Producer process identity
+includes its start time, and available Linux pidfds detect producer exit, so
+old frames cannot be delivered after a producer restart. Normal close removes
+owned allocations and discovery markers. Resource tracking covers allocation
+cleanup when the owning application exits; an abnormal exit can leave small
+discovery markers, which contain no payload data.
+Live receivers release cached mappings after a producer exits and prune stale
+markers during discovery within their deployment scope.
+
+## Configuration and observation
+
+Connector extras are shared by the endpoints of an edge. For example:
+
+```yaml
+connectors:
+  shm:
+    name: SharedMemoryConnector
+    extra:
+      host_ring_bytes: 4194304
+```
+
+Set `host_ring_bytes: 0` at both endpoints to use the per-key path, including
+controlled comparisons or deployments with older receivers. Nonzero capacities
+must be between 4 KiB and 64 MiB. All ring endpoints must support this protocol.
+The launcher supplies a common deployment scope; discovery scans only that
+scope's directory. Receive attempts do not wait for another process's file lock.
+An unchanged empty channel skips its frame lock. Channel registration and
+retirement advance a serialized discovery stamp, so empty polling does not
+rescan the directory and new producers remain discoverable without a timer.
+Ring endpoints must share a Linux PID namespace. Deployments that share SHM
+across different PID namespaces must use `host_ring_bytes: 0`.
+
+`health()` exposes `host_ring_puts`, `host_ring_gets` and `host_ring_fallbacks`
+alongside the existing counters. Ring setup failure logs once per edge and uses
+per-key SHM. CPU payload size and the cost of serialization determine the benefit;
+reduced transfer latency does not imply higher model throughput.
+
+## Scope
+
+This mechanism reduces host serialization copies and per-message segment work.
+It does not retain worker outputs on the GPU, allocate a CUDA arena, or replace
+GPU stream synchronization. It also does not impose a global byte budget on
+queued snapshots, downstream model caches or client audio delivery. The bounded
+ring with a compatibility fallback is not end-to-end streaming backpressure.
+
+Regression tests are in
+`tests/distributed/omni_connectors/test_shm_ring.py`, collected by the existing
+`core_model and cpu` CI lane. They cover ownership, tensor types, typed payloads,
+pressure fallback, cancellation, overwrite, deadlines, competing processes and
+producer restart.

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/index.md

@@ -68,6 +68,7 @@ guides list selectable backends, platform defaults, installation, and tuning.
 The design contracts separate selection mechanics from backend algorithms:
 
 - [Attention Backend Selection](feature/attention_backend_selection.md)
+- [Attention Execution Contract PoC](feature/attention_execution_contract_poc.md)
 - [Skip-Softmax](feature/skip_softmax.md)
 
 #### CPU offloading
@@ -116,4 +117,5 @@ The pre-#5137 pages are preserved in the
 [legacy module archive](module/archive/README.md) for historical reference and
 are not active design contracts.
 
+- [Qwen3-Omni MRv2 performance](qwen3_omni_mrv2_performance.md)
 - [MiniCPM-o 4.5 turn-mode MRv2 performance](minicpm_o45_mrv2_performance.md)

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/design/qwen3_omni_mrv2_performance.md

@@ -0,0 +1,22 @@
+# Qwen3-Omni MRv2 performance
+
+This model adaptation builds on #8184's existing stage runtime and output
+contracts. It adds Talker CPU metadata parsing, sampled-token embedding
+handoff, an optional first-frame decoder, length-grouped Code2Wav graph
+buckets, and opt-in incremental short-KV code prediction. The H200 profile
+also enables conservative fused sampling and bounded cuDNN algorithm search.
+The shared predictor keeps the upstream re-prefill helper and its default
+sampling behavior; short KV is explicitly selected by the deployment profile.
+
+Use `vllm_omni/deploy/qwen3_omni_moe_mrv2_h200.yaml` for the experimental H200
+profile. Preserve a control with `code_predictor_fused_sampling=false` and
+`codec_cudnn_benchmark=false` when C64 P99 is the primary requirement.
+cuDNN search restores process settings on exit and has startup cost. Kernel
+arithmetic is not bitwise-equivalent; the first-frame decoder uses additional
+model memory. Full-duplex performance is not established by turn-mode tests.
+
+First-frame delivery requires `talker_first_audio: true` in connector extra,
+a one-frame initial codec chunk, CUDA, a local single-rank Talker and prefix caching disabled.
+Other deployments retain the regular Code2Wav path.
+`codec_fused_snake` explicitly opts CUDA profiles into fused decoder blocks;
+the base V1 profile retains its original blocks and indexed text cache.

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/serving/video_stream_api.md

@@ -16,6 +16,10 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct \
     --trust-remote-code
 ```
 
+To enable [incremental prefill](#incremental-prefill), set
+`enable_prefix_caching: true` on stage 0 (the thinker) in the deploy config. Otherwise
+the server falls back to rebuilding the prompt per query.
+
 ### Run the Example Client
 
 ```bash
@@ -54,7 +58,7 @@ WebSocket /v1/video/chat/stream
 
 When buffered frames include client-supplied `frame_id` values, the server emits `video.frames.consumed` after the first engine output. Its `frame_ids`, `frames`, and `latest_pts_ms` describe the images included in that query's prompt.
 
-Queries stride-sample the buffered frames, keeping the last frame, then exclude frames already known to have failed decoding. Excluded frames are not replaced with other buffered frames. Frames whose background decoding has not finished remain eligible through the image URL path.
+Queries stride-sample the buffered frames, keeping the last frame, then exclude frames already known to have failed decoding. With incremental prefill active, every buffered frame is included instead of stride-sampling, with the same decode-failure exclusion. Excluded frames are not replaced with other buffered frames. The incremental path waits for the whole window to finish decoding. The legacy path waits only for its sampled positions, so an unselected frame still decoding does not delay the query.
 
 **Bugfix compatibility note:** `video.frames.consumed` now excludes known decode failures that were already excluded from the prompt. Older versions could report those frames and their timestamps as consumed. The event name and fields are unchanged; an empty selection reports empty lists and `latest_pts_ms: null`.
 
@@ -64,10 +68,10 @@ Queries stride-sample the buffered frames, keeping the last frame, then exclude
 | ------- | ------ | --------- | ------------- |
 | `model` | string or null | null | Optional model name. Usually omitted because the server hosts one model. |
 | `modalities` | list[string] | `["text", "audio"]` | Output modalities. Use `["text"]`, `["audio"]`, or both. |
-| `num_frames` | integer, 1-128 | `4` | Number of buffered frames sampled for each query. |
-| `max_frames` | integer, 1-256 | `50` | Maximum retained frame buffer size. Oldest frames are evicted first. |
+| `num_frames` | integer, 1-128 | `4` | Number of buffered frames sampled for each query. Legacy path only: with incremental prefill active, every buffered frame is submitted and frame density is the client's responsibility (clients push discrete frames at their own rate). |
+| `max_frames` | integer, 1-256 | `50` | Maximum retained frame buffer size. Below the cap both paths keep every frame. Legacy evicts the oldest frame once full. Incremental prefill drops the oldest half at once, then appends until full again. |
 | `system_prompt` | string or null | null | Optional custom system prompt. |
-| `use_audio_in_video` | bool | `true` | Include streamed audio chunks in multimodal video understanding when audio is present. |
+| `use_audio_in_video` | bool | `true` | Incremental prefill: forwarded on warmup and every query when this is true, so multimodal hashes stay aligned. Legacy: forwarded only when this query has input audio. |
 | `sampling_params_list` | list or null | null | Optional per-stage parameter dictionaries. Each provided entry replaces that stage's deployment sampling settings. |
 | `enable_frame_filter` | bool | `true` | Enable EVS near-duplicate frame filtering. |
 | `frame_filter_threshold` | float, 0.0-1.0 | `0.95` | EVS similarity threshold. Higher keeps more frames; lower drops more near-duplicates. |
@@ -107,9 +111,42 @@ The server accepts these legacy field names and rewrites them before validation.
 
 EVS compares downsampled frames and drops near-duplicate frames before they enter the session frame buffer. `frame_filter_threshold` controls retention: higher values are more permissive and keep more frames; lower values are more aggressive and drop more similar frames.
 
+## Incremental Prefill
+
+Server-side only (stage-0 `enable_prefix_caching: true` and `modalities: ["text"]`).
+There is no `session.config` flag. Audio-output sessions stay on the legacy path.
+
+Input is unchanged: `video.frame` and `audio.chunk` are separate buffers. Warmup
+prefills `history + frames` as they arrive (`max_tokens=1`). When
+`use_audio_in_video` is true, warmup and query both pass that processor kwarg.
+`video.query` is the same prompt plus optional input audio and the question text,
+so the warmed vision prefix cache is reused and the query pays for the audio/text
+suffix. All buffered frames are
+submitted in arrival order (`num_frames` subsample is legacy-only). Once the buffer
+hits `max_frames`, the oldest half is dropped in one step and a warning is logged;
+until then the buffer matches the legacy path. After a turn,
+history compresses to the last two text messages and the next warmup follows the
+new prefix.
+
+### Legacy path (incremental prefill inactive)
+
+Stage-0 prefix caching is off, or `modalities` includes `"audio"`: no warmup,
+frames are re-sampled up to `num_frames` each query, and the query waits only
+for those sampled frames to decode. Prompt shape is the same
+(`image_pil` frames, optional trailing `input_audio`, text). Both paths keep the
+frame buffer across queries (`max_frames`) and compress history the same way.
+
 ## Known Limitations
 
-- Session KV reuse and incremental prefill are not implemented in this PR. Each `video.query` rebuilds the model prompt from the retained frame and audio buffers.
+- Past `max_frames`, incremental prefill keeps a suffix of the newest frames, not a resample of the whole stream. Rebuilding that longer window is left to a follow-up.
+- Incremental prefill applies only to text-only sessions with stage-0 prefix caching enabled
+  on `WS /v1/video/chat/stream`. It does not accelerate `/v1/realtime` camera input, which
+  uses duplex conversation items rather than `video.frame` / `video.query`. Audio-output
+  sessions on this endpoint rebuild the prompt each query.
+- Stage-0 prefix caching is safe for audio-output requests: on a hit, the omni prefix cache
+  restores the thinker hidden states the talker needs. A preempted request that resumes on a
+  hit does not get them back, so size the deployment to avoid preemption; see
+  [Prefix Caching](../design/feature/prefix_caching.md).
 - Back-to-back short replies can still expose an engine-layer scheduler race. The PR notes an observed workaround of at least 200 ms idle between turns when clients repeatedly see idle timeouts.
 - If the audio buffer exceeds the server limit, the server emits `Audio buffer overflow` and clears the currently buffered audio for the session.
 - The API is intended for Qwen3-Omni streaming video understanding; other models may not support the same multimodal processor arguments.

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/offline_inference/yue2/end2end.py

@@ -28,7 +28,9 @@ import torch
 from vllm_omni import Omni
 from vllm_omni.model_executor.models.yue2.yue2 import (
     ABC_SAMPLING,
+    CONTEXT,
     KEY_MAX_AUDIO_FRAMES,
+    KEY_MAX_HOLD_STEPS,
     KEY_MIN_TOKENS,
     KEY_PENALTY_WINDOW,
     KEY_PHASE,
@@ -41,10 +43,12 @@ from vllm_omni.model_executor.models.yue2.yue2 import (
     KEY_TOP_P,
     SEMANTIC_SAMPLING,
     STOP_TOKEN_IDS,
+    SYNTHESIS_HOLD_STEPS,
 )
 from vllm_omni.tokenizers.yue2_prompt import (
     abc_ids_from_generated,
     abc_prefix_ids,
+    semantic_frames,
     semantic_prefix_ids,
 )
 from vllm_omni.tokenizers.yue2_tokenizer import YuE2TextTokenizer
@@ -71,7 +75,13 @@ def sampling_params(engine, *, phase, seed, max_frames, prompt_ids):
         # conditioning prefix from the scheduled tokens alone.
         KEY_PREFIX_IDS: list(prompt_ids),
     }
-    params.max_tokens = preset["max_tokens"] if phase == "abc" else max_frames + 1
+    if phase == "abc":
+        params.max_tokens = preset["max_tokens"]
+    else:
+        # The song is synthesized on a side stream while the request is held;
+        # leave max_tokens headroom for the hold tokens.
+        params.max_tokens = min(max_frames + 1 + SYNTHESIS_HOLD_STEPS, CONTEXT - len(prompt_ids))
+        params.extra_args[KEY_MAX_HOLD_STEPS] = max(0, params.max_tokens - max_frames - 1)
     params.stop_token_ids = list(STOP_TOKEN_IDS)
     params.detokenize = False
     return params
@@ -144,6 +154,9 @@ def main() -> None:
         payload = {
             "prompt_token_ids": prompt_ids,
             "generated_token_ids": generated_ids,
+            # Codec frames only: without the end token and the hold tokens
+            # the request emits while its song is synthesized.
+            "semantic_frames": semantic_frames(generated_ids),
             "seed": args.seed,
             "cot": args.cot,
             "max_frames": args.max_frames,

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/OpenMOSS/MOSS-TTS.md

@@ -35,6 +35,10 @@ the requested `input` speech. Without a nonblank `ref_text`, Local uses
 audio-reference generation. Reference transcripts must match the reference
 audio; they are not style instructions.
 
+For the optional CUDA MRV2 runner, see the
+[Local 1.5 deployment profile](../../docs/configuration/stage_configs.md#moss-tts-local-15-with-model-runner-v2).
+The default Local deployment continues to use V1.
+
 ## References
 
 - Offline inference example: [`examples/offline_inference/text_to_speech/moss_tts/`](../../examples/offline_inference/text_to_speech/moss_tts/)
@@ -152,3 +156,296 @@ curl -X POST http://localhost:8091/v1/audio/speech \
 - No `ref_audio` required or accepted for MOSS-SoundEffect.
 - Input field maps to the `ambient_sound` parameter in the upstream processor.
 - Rate: ~12.5 tokens per second; longer descriptions produce longer audio.
+
+## Local 1.5 MRV2 and slot attention
+
+`MOSS-TTS-Local-Transformer-v1.5` supports the native CUDA MRV2 pipeline with
+an explicit deploy profile. The default Local profile continues to use V1.
+Both profiles below preserve 1-frame initial and 15-frame steady codec chunks
+and the model's sampling defaults.
+
+```bash
+vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 --omni \
+  --deploy-config vllm_omni/deploy/moss_tts_local_mrv2.yaml
+```
+
+This C64 profile bounds Talker prefill to 512 tokens and retains codec-owned
+CUDA graphs without compiling the codec with Inductor.
+
+For sustained high concurrency on a large-memory CUDA GPU, use:
+
+```bash
+CUDA_VISIBLE_DEVICES=0 OMP_NUM_THREADS=4 \
+vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 --omni \
+  --deploy-config vllm_omni/deploy/moss_tts_local_mrv2_high_concurrency.yaml \
+  --stage-init-timeout 1200 --init-timeout 1500
+```
+
+The high-concurrency profile places both stages on one GPU, sets each stage's
+capacity to 128 and fixes the Talker KV budget at 32 GiB. The Talker uses
+mixed FULL CUDA graphs with both its token budget and capture limit set to
+512, so new prompt chunks remain within graph coverage. It also enables
+`hf_overrides.mrv2_gpu_slot_state`: Local hidden states, audio codes and
+continuation control stay in GPU request slots, gathered through MRV2's
+batch mapping. Prefill conditioning and explicit per-request sampling seeds
+retain their original behavior. Published audio codes own their storage, so
+slot reuse cannot overwrite an in-flight output. Set the override to `false`
+to compare with the generic Omni model state. The profile selects
+`codec_attention_backend: triton_slot`, Inductor mode 3 with combo kernels
+disabled, and codec CUDA graph buckets through 128. The configuration was
+validated on one H200. The earlier capacity-256 variant used about 110 GiB of
+sampled peak GPU memory, including loading and warmup. Reduce the stage capacities,
+graph buckets and coalescing target together when adapting it to smaller GPUs. Cold codec
+compilation can take several minutes; subsequent starts can reuse the AOT
+cache. Graph capture is still performed at startup.
+
+The high-concurrency profile also coalesces ready codec streams with
+`connectors.shm.extra.generation_min_batch_size: 32` and
+`generation_max_wait_ms: 12`. The target controls when to stop waiting; it
+neither caps execution at 32 nor guarantees every codec group has 32 rows.
+First and final chunks share the bounded window at high concurrency. When
+fewer than 32 requests remain, dispatch is immediate. Pending outputs are
+retired before waiting, with the original deadline preserved, so waiting does
+not prevent the previous batch from releasing its in-flight state. Cancellation
+and input notifications keep their ordinary scheduler bookkeeping. This option
+requires a stateful native MRV2 generation stage with TP1/PP1. Set the wait to
+`0` to disable coalescing; other profiles keep the immediate-dispatch default.
+
+The high-concurrency profile enables two additional Local model-state
+optimizations under the CUDA stage-0 `hf_overrides`, alongside
+`mrv2_gpu_slot_state: true`:
+
+- `mrv2_batch_prefill: true` reuses the batch's text embeddings and combines
+  CPU reference-code slices and their destination positions into one pinned
+  upload. Host staging is reused only after the upload event completes;
+  GPU-resident references keep the original path.
+- `mrv2_direct_tokens: true` returns the already determined text token after
+  the Local audio/stop decision. It avoids constructing and sampling the
+  full text vocabulary. Audio-code and binary-stop sampling are unchanged.
+  Requests requiring distribution metadata or token constraints, and
+  unsupported execution modes, retain the normal sampler.
+
+Other profiles leave both flags disabled. Set both flags explicitly when
+comparing their combined effect. Deployment
+inheritance replaces a stage's `hf_overrides` mapping, so an overlay must also
+retain `mrv2_gpu_slot_state: true`. Measure first-packet latency as well as
+throughput: faster prompt preparation can change competition between the
+Talker and codec sharing the GPU.
+
+For the single-H200 capacity-128 throughput baseline, keep both switches on.
+In complete Seed-TTS EN1088 runs at client concurrency 128, four combined-switch
+rounds pooled 344.11 audio-s/s versus 335.78 for batch prefill alone; direct
+tokens alone had no reliable gain. The combined path had 685 ms mean first
+audio versus 679 ms across six control rounds, with lower mean completion time.
+A fixed-seed 128-row audio sample showed mean Whisper WER of 3.68% versus
+3.78% for the control. This sample does not establish speaker similarity or
+exclude small quality changes. At concurrency 256, the combined path had
+raised mean first-audio latency by about 0.33 s, so other profiles leave the
+switches off until their latency and quality are checked.
+
+The optional connector setting `generation_coalescing_policy: idle_wait`
+waits for an inbox notification when no stream is runnable, all receivers
+are parked and registered, and no output needs retirement. The first ready
+arrival starts a fresh coalescing window, preserving the batch budget.
+Control messages, including cancellations, can wait up to two windows
+instead of one (24 ms at the 12 ms setting). The default policy remains
+`fixed`; neither policy changes the global orchestration default.
+
+The codec backends differ in state access:
+
+- `sdpa` uses PyTorch attention after gathering and updating the ring cache.
+- `triton` replaces the attention calculation, retaining the ring-cache
+  gather/copy and explicit mask construction.
+- `triton_slot` writes surviving K/V directly into request slots, attends
+  directly to the ring, and advances the active slot offsets in three
+  ordered kernels. Graph-padding rows do not advance persistent state.
+
+The slot path preserves the existing chunk-complete ring semantics, including
+retaining the final cache-capacity tokens when a chunk exceeds the ring.
+It does not introduce another cache owner or change request-slot lifetime.
+Zero-length padding rows skip the attention computation and emit zeros.
+
+The high-concurrency profile includes intermediate codec batch buckets 6, 12
+and 24. To compare with power-of-two buckets, change the stage-1
+`cudagraph_capture_sizes` to `[1, 2, 4, 8, 16, 32, 64, 128]`.
+Keep the maximum bucket equal to the state capacity to retain terminal-tail
+coalescing. Smaller buckets reduce padding work but require more graphs;
+measure complete serving runs before selecting them for a deployment.
+
+For a codec output-path comparison, set
+`connectors.shm.extra.codec_gpu_stream_output: 0` to select the synchronous
+codec output path; `1` selects GPU output snapshots on a private codec
+stream. Both configurations use MRV2 and keep streaming audio responses.
+The option changes input ownership, metadata staging and snapshot handling
+as well as transfer scheduling, so its timing difference is not just D2H
+copy time.
+
+Those chunk-complete semantics truncate the causal window: every ring holds
+exactly `context` entries and a whole chunk is written before attention runs,
+so token `i` of a `T`-token chunk sees `capacity - T + i + 1` keys. With the
+tokenizer-v2 decoder's per-layer contexts (10/10/8/4/2/1 s, i.e. 125/250/400/
+400/400/400 tokens) a 15-frame chunk is 480 tokens at the last transformer, so
+its first 80 tokens attend to nothing and no chunk sees the previous one there.
+Padding a terminal tail to 15 frames therefore decodes it differently from an
+exact-length decode; the difference is deterministic (identical in fp32 and
+fp64), not bf16 noise. The connector option `codec_ring_headroom: 1` sizes each
+ring as `context + max_chunk_frames * tokens_per_frame` (including ramp shapes); chunked streaming
+then reproduces whole-sequence decoding bit-for-bit in fp64 and padded tails
+equal exact tails. It increases state/graph memory and capture time; historical C256 experiments
+observed about a 2% throughput cost. It is off by default and its perceptual
+quality benefit has not been established. Terminal tails are padded
+into the regular 15-frame graph bucket in either mode.
+The slot kernel is specific to the CUDA tokenizer-v2 decoder; other paths
+retain their existing attention implementation. Event-driven orchestration
+remains independently selectable and its default is unchanged.
+
+### Optional progressive chunks
+
+Set `codec_chunk_ramp` under `connectors.shm.extra` to insert smaller chunks
+before steady decoding, for example:
+
+```yaml
+connectors:
+  shm:
+    extra:
+      codec_chunk_ramp: [1, 4, 15]
+```
+
+The ramp's first entry overrides `initial_codec_chunk_frames`. After the last
+entry, chunks use `codec_chunk_frames`. Each request advances independently;
+the final partial chunk is flushed and an empty terminal carries no codec
+tokens. The codec captures every ramp length, sizes the maximum execution
+step for the largest entry, and uses the ramp's first entry for its dedicated
+first-chunk graph. Additional shapes increase startup and graph memory.
+
+Ramps are opt-in. They can reduce early playback underrun while increasing
+codec calls, first-audio latency or total generation time. The measured C128
+throughput presets retain 1→15; their published figures do not establish an
+additional benefit from enabling ramps. The processor and graph-shape tests
+run without GPU or model weights:
+
+```bash
+CUDA_VISIBLE_DEVICES= PYTHONPATH=. python -m pytest -q \
+  tests/model_executor/stage_input_processors/test_moss_tts_async_chunk.py \
+  -m 'core_model and cpu' --run-level core_model
+```
+
+### Low-latency and reference-encoding options
+
+`moss_tts_local_mrv2_low_latency.yaml` keeps both stage capacities at 128 and
+uses prefill 2048, GPU slot state, batch prefill, direct tokens and eager MTP.
+Its codec first-chunk fast path uses dedicated graphs and a stream handoff
+before the regular decoder reuses the request slot. A gate limits contention
+with regular codec work. Slot waits are bounded at 30 seconds. A failed or
+stalled decode raises an error rather than replaying a partially advanced
+slot; a failed output enqueue retains owned PCM for regular delivery.
+Closing the decoder rejects new jobs. Its regular-batch limit of 16 is internal dispatch
+policy; it does not change the client concurrency or stage capacities.
+
+Reference encoding runs in the API layer, independently of MRV2. Enable
+reference graphs explicitly with `VLLM_OMNI_MOSS_REF_GRAPHS=1`; the default
+eager encoder is retained when the option is absent. The graph path uses the
+loaded tokenizer's encoder/quantizer, length buckets and optional compilation
+(`VLLM_OMNI_MOSS_REF_COMPILE=0` disables compilation). It also uses windowed
+attention; `VLLM_OMNI_MOSS_REF_ATTN=sdpa` retains the original attention.
+Compilation and attention changes need not produce bit-identical codes.
+
+For multiple API processes, `VLLM_OMNI_MOSS_REF_CODES_SHARED_DIR` enables shared
+reference-code storage and, by default, a single encoder host with four
+workers. Use a dedicated directory per service, checkpoint and encoding
+configuration. Workers have their own graph resources; all graph captures
+complete before serving begins. `VLLM_OMNI_MOSS_REF_SHARED_ENCODER=0` retains
+separate encoders while sharing codes. More workers/graphs consume memory;
+they do not imply more GPUs. Example after selecting an available GPU and
+configuring any operator-managed MPS daemon:
+
+```bash
+VLLM_WORKER_MULTIPROC_METHOD=spawn \
+VLLM_OMNI_EVENT_DRIVEN_ORCH=1 \
+VLLM_OMNI_CONNECTOR_RECV_POLL_MS=1 \
+VLLM_OMNI_MOSS_REF_GRAPHS=1 \
+VLLM_OMNI_MOSS_REF_ENCODER_WORKERS=4 \
+VLLM_OMNI_MOSS_REF_BATCH_WINDOW_MS=0 \
+VLLM_OMNI_MOSS_REF_INFLIGHT=4 \
+VLLM_OMNI_MOSS_REF_HOST_WINDOW_MS=2 \
+VLLM_OMNI_MOSS_REF_CODES_SHARED_DIR=/dev/shm/moss-local-service \
+vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 --omni \
+  --api-server-count 4 \
+  --deploy-config vllm_omni/deploy/moss_tts_local_mrv2_low_latency.yaml \
+  --stage-init-timeout 1200 --init-timeout 1500
+```
+
+This example is a configurable serving profile, not the complete launch
+configuration of a historical benchmark. In particular, the low-latency YAML
+uses synchronous stage-1 scheduling; the retained source46 experiment used
+asynchronous stage-1 scheduling, FP8 backbone, MPS and additional encoder
+environment settings. Do not infer a throughput result from the YAML alone.
+
+### Cold versus hot reference measurements
+
+A cold reference misses the encoded-reference cache: the request pays for
+reference parsing/preparation and GPU encoding before speech generation.
+A hot reference reuses reference codes but still synthesizes new speech;
+generated audio is not cached. Pre-registering a reference moves encoding
+into registration and can make subsequent synthesis hot, but registration
+latency is still part of the first-use cost.
+
+Model startup/compilation is separate from reference coldness. EN1088 contains
+repeated references, so its first pass is not uniformly cold. For an all-unique
+cold pass, use its 666 first-occurrence references, keep order, warm execution
+with disjoint references, then repeat the same 666 requests. A second pass
+through independent API-local caches is not guaranteed hot; verify cache
+coverage before comparing it with a shared-cache result. Report requests/s
+alongside audio-s/s and retain failed, empty, long and near-silent outputs.
+
+### Reproduce the serving benchmark
+
+Use the complete Seed-TTS English test set, including reference audio and text.
+Set `SEED_TTS_DATA` to the directory containing `en/meta.lst` and its 1088
+entries. Run the native benchmark after the server is ready:
+
+```bash
+export VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE=48000
+export VLLM_OMNI_BENCH_AUDIO_CHANNELS=2
+export SEED_TTS_WER_EVAL=0
+
+for concurrency in 128; do
+  for phase in warm r1 r2; do
+    vllm bench serve --omni \
+      --model OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 \
+      --backend openai-audio-speech --endpoint /v1/audio/speech \
+      --dataset-name seed-tts --dataset-path "$SEED_TTS_DATA" \
+      --seed-tts-locale en --disable-shuffle --num-prompts 1088 \
+      --num-warmups 0 --ready-check-timeout-sec 0 \
+      --output-len 256 --max-concurrency "$concurrency" \
+      --request-rate inf --seed 42 \
+      --extra-body '{"task_type":"Base","max_new_tokens":256}' \
+      --save-result --save-detailed --result-dir results/moss-local \
+      --result-filename "c${concurrency}-${phase}.json"
+  done
+done
+```
+
+Discard the complete `warm` pass and combine measured runs as total generated
+audio seconds divided by total benchmark duration. Require 1088 successes,
+zero failures and 1088 nonempty-audio metric samples in each run. The output
+cap matches the benchmark protocol; success alone does not establish speech
+quality or that every sentence ended before the cap. Dataset seed 42 does
+not fix an independent sampling seed for every request.
+
+For an attention-only comparison, copy the high-concurrency YAML beside the
+original to preserve relative `base_config` resolution. Change its
+codec `compilation_config.mode` to `0` while retaining `cudagraph_mode: FULL`,
+and compare `codec_attention_backend: triton` against `triton_slot`. Keep all
+other settings, warmup and client concurrency identical. Restart the server
+between configurations and use separate result directories. Comparing the
+compiled slot profile to uncompiled Triton includes both changes.
+
+Kernel and MHA regression tests require CUDA, compatible vLLM/Triton packages,
+and no model weights:
+
+```bash
+python -m pytest -q tests/model_executor/models/moss_tts/test_slot_attention.py \
+  tests/model_executor/models/moss_tts/test_streaming_attention.py \
+  -m 'core_model and cuda' --run-level=core_model
+```

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)
@@ -40,9 +45,11 @@ user-downloaded weights and bundles none.
 | Output | 48 kHz stereo, whole song delivered when the semantic request finishes |
 
 Generation is one or two requests on a single AR stage: the abc phase writes
-the score, the semantic phase generates codec frames and, on its last step,
-solves the ODE and decodes the whole song in-engine (weights load once; the
-VAE loads at startup from `$YUE2_VAE` or the hub id).
+the score, the semantic phase generates codec frames, then queues the ODE and VAE on a
+high-priority CUDA stream. Other requests keep decoding while the finishing
+request emits HOLD tokens (up to 4096 steps, within its token/context budget).
+The scheduler is unchanged. Weights load once; the VAE loads at startup from
+`$YUE2_VAE` or the hub id.
 
 ## Running
 
@@ -67,15 +74,50 @@ ending.
 
 ## Notes
 
-- **Memory (RTX 4090, 24 GB, `gpu_memory_utilization: 0.70`):** the engine
-  reserves ~15.8 GiB after startup; with 4 concurrent requests pinned to the
-  9000-frame cap (worst case), peak usage during the terminal NAR/VAE
-  finishing pass reaches ~20.9 GiB, leaving ~3 GiB of headroom. At 0.85 the
-  same workload OOMs inside the finishing pass, so the deploy yaml pins 0.70.
-- **Known limitations:** 4 concurrent full-length requests are the verified
-  shape (`max_num_seqs: 4`); CUDA graph capture is a
-  follow-up; the abc phase runs eagerly
-  after prefill (its tokens are the product, not audio).
+The lifecycle e2e regression forces KV-cache preemption with chunked and
+unchunked recompute, mixes ABC and semantic sampling, and cancels a request
+after NAR work has started. It checks history alignment, released synthesis
+buffers, and successful generation after cancellation. Run with one H100/H200
+and cached YuE2-3B and YuE2-Vae weights:
+
+```bash
+export YUE2_MODEL_DIR=/path/to/YuE2-3B
+export YUE2_VAE=/path/to/YuE2-Vae
+CUDA_VISIBLE_DEVICES=0 python -m pytest tests/e2e/online_serving/test_yue2.py \
+    -k preemption_and_synthesis_abort -m 'slow and tts' --run-level full_model -q -s
+```
+
+The test retains the module's weekly TTS routing; it is not a per-PR CI gate.
+
+- **Memory:** `gpu_memory_utilization` budgets the vLLM engine; NAR K/V,
+  acoustic graph buffers and VAE activations also need room during synthesis.
+  The 0.70 / 4-slot defaults require validation on the target card. The earlier
+  RTX 4090 measurement (15.8 GiB after startup, 20.9 GiB peak at 9000 frames)
+  predates the async/compiled implementation; it does not validate this
+  version or long ABC prefixes on a 4090. Reduce the fraction if the target
+  workload needs more synthesis memory, especially for long ABC prefixes and
+  the 9000-frame cap. Releasing completed chunks bounds live NAR buffers,
+  while the shared graph allocator retains reserved storage for reuse.
+- **CUDA graphs:** the default deploy keeps the AR backbone eager. Both ABC
+  and semantic sampling use prewarmed power-of-two graph buckets. Custom
+  sampler captures have a bounded cache and fall back to eager sampling when
+  it fills. The NAR velocity pass uses a chunk graph on the serialized synthesis
+  stream. Chunk graphs share one allocator pool; completed chunks release
+  their engines and buffers after their own events. FA3 runs on Hopper; other
+  CUDA cards use SDPA, which is also warmed at startup.
+- **Serving:** run `vllm serve m-a-p/YuE2-3B --omni` with the default
+  `yue2.yaml`, which uses eager AR with 4 slots. Validate memory capacity
+  for the target card and workload.
+  For the tested single-H200 graph settings, reuse this deploy with
+  `--stage-overrides '{"0":{"max_num_seqs":32,"gpu_memory_utilization":0.5,"enforce_eager":false,"compilation_config":{"cudagraph_mode":"FULL_AND_PIECEWISE"}}}'`.
+  Decode uses FULL and mixed prefill uses PIECEWISE; VAE and the Python
+  synthesis queue remain outside these graphs.
+  Audio is delivered as a whole song; time to first audio equals completion
+  latency. Aborting a running song stops future work units and waits only for
+  the at-most-two submitted units before releasing its buffers.
+- **Preemption:** before drawing for a resumed row, model-owned sampling
+  reconciles the codec history and terminal flags with the engine's retained
+  sequence. Discarded draws cannot advance its penalty window or frame budget.
 - Audio is not expected to match the upstream torch reference bit-for-bit
   (fused vs eager kernels). Measured against the upstream torch reference
   (bd90e4c, same style/lyrics/seed 831001, cot=off, 200 frames): prompt token

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(
@@ -2928,9 +2978,17 @@ class TestPlatformOverrides:
 
         assert deploy.stages[0].engine_extras["block_size"] == 128
 
-    @pytest.mark.parametrize("platform", ["cuda", "npu"])
-    def test_recommended_native_runner_platform_defaults(self, platform):
-        filename, pipeline_key = "qwen3_tts_mrv2.yaml", "qwen3_tts"
+    @pytest.mark.parametrize(
+        "filename,pipeline_key",
+        [
+            ("qwen3_tts_mrv2.yaml", "qwen3_tts"),
+            ("moss_tts_local_mrv2.yaml", "moss_tts_local"),
+            ("moss_tts_local_mrv2_high_concurrency.yaml", "moss_tts_local"),
+            ("moss_tts_local_mrv2_low_latency.yaml", "moss_tts_local"),
+        ],
+    )
+    @pytest.mark.parametrize("platform", ["cuda", "npu", "xpu", "rocm", "musa"])
+    def test_recommended_native_runner_platform_defaults(self, platform, filename, pipeline_key):
         deploy = load_deploy_config(Path(get_deploy_config_path(filename)))
         deploy = _apply_platform_overrides(deploy, platform=platform)
         expect_v2 = platform == "cuda"
@@ -2946,6 +3004,60 @@ class TestPlatformOverrides:
             # v2 only engages the native plane on stages declaring support.
             assert all(ps.supports_native_mrv2_data_plane for ps in pipeline.stages)
 
+    def test_moss_local_mrv2_preserves_base_profile_and_variant_scope(self):
+        pipeline = resolve_pipeline_config("moss_tts_local")
+        base_path = Path(get_deploy_config_path(pipeline.default_deploy_config_name))
+        candidate_path = Path(get_deploy_config_path("moss_tts_local_mrv2.yaml"))
+        base = load_deploy_config(base_path)
+        candidate = load_deploy_config(candidate_path)
+        assert base.model_runner == "v1"
+        assert candidate.stages == base.stages
+        assert candidate.connectors == base.connectors
+        assert candidate.async_chunk == base.async_chunk
+        cuda = _apply_platform_overrides(candidate, platform="cuda")
+        stages = merge_pipeline_deploy(pipeline, cuda)
+        assert stages[0].yaml_engine_args["max_num_batched_tokens"] == 512
+        codec_args = stages[1].yaml_engine_args
+        assert codec_args["max_num_seqs"] == 64
+        assert codec_args["enforce_eager"] is False
+        assert codec_args["compilation_config"] == {
+            "mode": 0,
+            "cudagraph_mode": "FULL",
+            "cudagraph_capture_sizes": [1, 2, 4, 8, 16, 32, 64],
+            "cudagraph_num_of_warmups": 1,
+        }
+        for platform in ("npu", "xpu", "rocm", "musa"):
+            # Platform resolution mutates the deployment; each server loads
+            # a fresh profile before applying its own platform overrides.
+            fallback = _apply_platform_overrides(load_deploy_config(candidate_path), platform=platform)
+            original = _apply_platform_overrides(load_deploy_config(base_path), platform=platform)
+            assert fallback.stages == original.stages
+        for variant in ("moss_tts_delay", "moss_tts_realtime"):
+            assert not any(stage.supports_native_mrv2_data_plane for stage in resolve_pipeline_config(variant).stages)
+
+    def test_moss_local_high_concurrency_resolves_slot_compile_and_memory_budget(self):
+        pipeline = resolve_pipeline_config("moss_tts_local")
+        path = Path(get_deploy_config_path("moss_tts_local_mrv2_high_concurrency.yaml"))
+        deploy = _apply_platform_overrides(load_deploy_config(path), platform="cuda")
+        talker, codec = merge_pipeline_deploy(pipeline, deploy)
+        assert talker.yaml_engine_args["max_num_seqs"] == codec.yaml_engine_args["max_num_seqs"] == 128
+        assert talker.yaml_engine_args["hf_overrides"]["mrv2_batch_prefill"] is True
+        assert talker.yaml_engine_args["hf_overrides"]["mrv2_direct_tokens"] is True
+        assert talker.yaml_engine_args["hf_overrides"]["mrv2_gpu_slot_state"] is True
+        assert talker.yaml_engine_args["max_num_batched_tokens"] == 512
+        assert talker.yaml_engine_args["kv_cache_memory_bytes"] == 32 * 1024**3
+        assert codec.yaml_engine_args["hf_overrides"]["codec_attention_backend"] == "triton_slot"
+        compilation = codec.yaml_engine_args["compilation_config"]
+        assert compilation["mode"] == 3 and compilation["backend"] == "inductor"
+        assert compilation["cudagraph_mode"] == "FULL"
+        assert compilation["cudagraph_capture_sizes"] == [1, 2, 4, 6, 8, 12, 16, 24, 32, 64, 128]
+        assert compilation["inductor_compile_config"] == {"combo_kernels": False, "benchmark_combo_kernel": False}
+        base_path = Path(get_deploy_config_path("moss_tts_local.yaml"))
+        for platform in ("npu", "xpu", "rocm", "musa"):
+            actual = _apply_platform_overrides(load_deploy_config(path), platform=platform)
+            expected = _apply_platform_overrides(load_deploy_config(base_path), platform=platform)
+            assert actual.stages == expected.stages
+
     @pytest.mark.parametrize("runner,native", [("v1", False), ("v2", False), ("v2", True)])
     def test_mrv2_undeclared_transport_warns(self, monkeypatch, runner, native):
         from unittest.mock import Mock

tests/config/test_config_merge.py

@@ -35,3 +35,161 @@ def test_recursive_deploy_inheritance(tmp_path):
     merged = resolve_deploy_yaml(leaf)
     assert merged["speech_cache"] == {"nested": {"a": 1, "b": 0}}
     assert merged["async_chunk"] is False
+
+
+def test_single_stage_streaming_capability_uses_central_async_resolver():
+    from vllm_omni.config.stage_config import DeployConfig, _resolve_pipeline_async_chunk_enabled
+    from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_FUSED_PIPELINE
+
+    assert _resolve_pipeline_async_chunk_enabled(QWEN3_TTS_FUSED_PIPELINE, DeployConfig())
+    with pytest.raises(ValueError, match="requires async_chunk=True"):
+        _resolve_pipeline_async_chunk_enabled(QWEN3_TTS_FUSED_PIPELINE, DeployConfig(async_chunk=False))
+
+
+@pytest.mark.parametrize(
+    "pipeline_key,profile,stage_count",
+    [
+        ("qwen3_tts", "qwen3_tts_high_concurrency_mrv2_single_gpu.yaml", 2),
+        ("qwen3_tts_fused", "qwen3_tts_fused_single_gpu.yaml", 1),
+    ],
+)
+def test_qwen3_tts_deployment_selection_preserves_output_and_transport_contract(pipeline_key, profile, stage_count):
+    from tests.helpers.stage_config import get_deploy_config_path
+    from vllm_omni.config.pipeline_registry import resolve_pipeline_config
+    from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
+
+    pipeline = resolve_pipeline_config(pipeline_key)
+    assert pipeline is not None
+    deploy = load_deploy_config(get_deploy_config_path(profile))
+    stages = merge_pipeline_deploy(pipeline, deploy)
+    assert len(stages) == stage_count
+    talker = stages[0]
+    if stage_count == 1:
+        assert talker.final_output and talker.final_output_type == "audio"
+        assert talker.yaml_engine_args["engine_output_type"] == "audio"
+        assert "custom_process_next_stage_input_func" not in talker.yaml_engine_args
+        assert not talker.yaml_extras.get("output_connectors")
+        assert talker.yaml_engine_args["supports_running_prefix_cache_reset"] is False
+        assert talker.yaml_engine_args["additional_config"]["ref_code_context_frames"] == 72
+    else:
+        assert not talker.final_output
+        assert talker.yaml_engine_args["engine_output_type"] == "latent"
+        assert talker.yaml_engine_args["custom_process_next_stage_input_func"].endswith("talker2code2wav_async_chunk")
+        assert talker.yaml_extras["output_connectors"]
+        assert stages[1].input_sources == [0] and stages[1].final_output_type == "audio"
+        assert talker.yaml_engine_args.get("supports_running_prefix_cache_reset", True)
+
+
+def test_single_stage_running_reset_capability_cannot_be_overridden():
+    from vllm_omni.config.omni_config import VllmOmniConfig
+    from vllm_omni.config.stage_config import DeployConfig, StageDeployConfig, merge_pipeline_deploy
+    from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_FUSED_PIPELINE
+
+    deploy = DeployConfig(
+        stages=[StageDeployConfig(stage_id=0, engine_extras={"supports_running_prefix_cache_reset": True})]
+    )
+    legacy = merge_pipeline_deploy(QWEN3_TTS_FUSED_PIPELINE, deploy)[0]
+    legacy.runtime_overrides["supports_running_prefix_cache_reset"] = True
+    assert legacy.to_omegaconf().engine_args.supports_running_prefix_cache_reset is False
+    with pytest.raises(ValueError, match="no structured config owner: supports_running_prefix_cache_reset"):
+        VllmOmniConfig.from_pipeline_config(QWEN3_TTS_FUSED_PIPELINE, user_deploy_config=deploy)
+    structured = VllmOmniConfig.from_pipeline_config(QWEN3_TTS_FUSED_PIPELINE)
+    assert structured.stage_by_id(0).model_config.supports_running_prefix_cache_reset is False
+
+
+@pytest.fixture
+def single_stage_tts_config(mocker):
+    from vllm.config import VllmConfig
+
+    config = mocker.Mock(spec=VllmConfig)
+    config.additional_config = {"talker_stream_decode": True}
+    config.model_config = mocker.Mock(
+        stage_connector_config={},
+        engine_output_type="audio",
+        final_output=True,
+        supports_running_prefix_cache_reset=False,
+        use_v2_model_runner=True,
+        async_chunk=True,
+    )
+    config.parallel_config = mocker.Mock(
+        tensor_parallel_size=1, pipeline_parallel_size=1, distributed_executor_backend="uni"
+    )
+    config.cache_config = mocker.Mock(enable_prefix_caching=False)
+    config.device_config = mocker.Mock(device="cuda")
+    mocker.patch(
+        "vllm_omni.model_executor.models.qwen3_tts.stream_decode.current_omni_platform.is_cuda", return_value=True
+    )
+    return config
+
+
+def test_single_stage_tts_rejects_unsupported_decoder_configuration(single_stage_tts_config):
+    from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled
+
+    config = single_stage_tts_config
+    config.cache_config.enable_prefix_caching = True
+    with pytest.raises(ValueError, match="enable_prefix_caching=False"):
+        talker_stream_decode_enabled(config)
+    config.cache_config.enable_prefix_caching = False
+    assert talker_stream_decode_enabled(config)
+
+
+@pytest.mark.parametrize("option_location", ["additional_config", "connector"])
+@pytest.mark.parametrize("output_type", ["latent", "audio"])
+def test_two_stage_talker_rejects_single_stage_pcm_option(single_stage_tts_config, option_location, output_type):
+    from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled
+
+    config = single_stage_tts_config
+    config.model_config.engine_output_type = output_type
+    config.model_config.final_output = False
+    if option_location == "connector":
+        config.model_config.stage_connector_config = {"extra": config.additional_config}
+        config.additional_config = {}
+    with pytest.raises(ValueError, match="final audio-output Talker"):
+        talker_stream_decode_enabled(config)
+
+
+def test_two_stage_first_audio_remains_independent(single_stage_tts_config):
+    from vllm_omni.model_executor.models.qwen3_tts.first_audio import talker_first_audio_enabled
+    from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled
+
+    config = single_stage_tts_config
+    config.model_config.engine_output_type = "latent"
+    config.additional_config = {}
+    config.model_config.stage_connector_config = {"extra": {"talker_first_audio": True}}
+    assert not talker_stream_decode_enabled(config)
+    assert talker_first_audio_enabled(config)
+    config.model_config.stage_connector_config["extra"]["talker_stream_first_audio"] = True
+    with pytest.raises(ValueError, match="requires the single-stage"):
+        talker_stream_decode_enabled(config)
+
+
+def test_single_stage_pcm_requires_cuda_device_and_own_first_audio_option(single_stage_tts_config):
+    from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled
+
+    config = single_stage_tts_config
+    config.device_config.device = "cpu"
+    with pytest.raises(ValueError, match="requires CUDA"):
+        talker_stream_decode_enabled(config)
+    config.device_config.device = "cuda"
+    config.additional_config["talker_first_audio"] = True
+    with pytest.raises(ValueError, match="mutually exclusive"):
+        talker_stream_decode_enabled(config)
+    config.additional_config.pop("talker_first_audio")
+    config.additional_config["talker_stream_first_audio"] = True
+    assert talker_stream_decode_enabled(config)
+
+
+def test_single_stage_audio_talker_requires_pcm_option(single_stage_tts_config):
+    from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled
+
+    single_stage_tts_config.additional_config = {}
+    with pytest.raises(ValueError, match="requires talker_stream_decode=True"):
+        talker_stream_decode_enabled(single_stage_tts_config)
+
+
+def test_single_stage_pcm_requires_reset_capability_declaration(single_stage_tts_config):
+    from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled
+
+    single_stage_tts_config.model_config.supports_running_prefix_cache_reset = True
+    with pytest.raises(ValueError, match="must disable running prefix-cache reset"):
+        talker_stream_decode_enabled(single_stage_tts_config)

tests/config/test_environment_variables.py

@@ -167,9 +167,9 @@ def test_inventory_matches_reviewed_snapshot_counts():
         EnvironmentVariableCategory.PUBLIC_OMNI: 32,
         EnvironmentVariableCategory.INHERITED_VLLM: 20,
         EnvironmentVariableCategory.PLATFORM_EXTERNAL: 29,
-        EnvironmentVariableCategory.MODEL_SPECIFIC: 75,
+        EnvironmentVariableCategory.MODEL_SPECIFIC: 95,
         EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 21,
-        EnvironmentVariableCategory.INTERNAL: 6,
+        EnvironmentVariableCategory.INTERNAL: 7,
     }
 
     disposition_counts = Counter(
@@ -178,7 +178,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
         if item.category is EnvironmentVariableCategory.MODEL_SPECIFIC
     )
     assert {disposition: disposition_counts[disposition] for disposition in ModelEnvironmentVariableDisposition} == {
-        ModelEnvironmentVariableDisposition.PROMOTE: 46,
+        ModelEnvironmentVariableDisposition.PROMOTE: 66,
         ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 5,
         ModelEnvironmentVariableDisposition.EXTERNAL: 0,
         ModelEnvironmentVariableDisposition.INTERNALIZE: 17,

tests/config/test_omni_config.py

@@ -1046,6 +1046,7 @@ def test_sub_config_fields_match_structured_scopes():
         "interleave_mm_strings",
         "media_io_kwargs",
         "final_output",
+        "supports_running_prefix_cache_reset",
         "active_stream_window",
         "session_mode",
         "duplex_max_sessions",
@@ -1947,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

@@ -0,0 +1,258 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Bounded readiness coalescing must not allocate or consume request state."""
+
+import queue
+from types import SimpleNamespace
+
+import pytest
+from vllm.v1.core.sched.interface import PauseState
+from vllm.v1.request import RequestStatus
+
+from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
+from vllm_omni.outputs import OmniConnectorOutput
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_idle_wait_wakes_on_input_without_extending_coalescing_deadline(monkeypatch):
+    events = [notification("0"), notification("1", "2", "3")]
+    s, clock = scheduler(monkeypatch, delayed=[(0.004, events[0]), (0.002, events[1])])
+    s._generation_coalescing_policy = "idle_wait"
+    for request in s.requests.values():
+        request.status = RequestStatus.WAITING_FOR_CHUNK
+    assert s._drain_omni_connector_outputs() == events
+    assert clock[0] == pytest.approx(10.006)
+    assert s._omni_connector_output_inbox.waits == pytest.approx([0.012, 0.012])
+
+
+@pytest.mark.parametrize("pending", ["in_flight", "admission", "registration", "terminal", "empty"])
+def test_idle_wait_preserves_completion_and_control_progress(monkeypatch, pending):
+    s, clock = scheduler(monkeypatch)
+    s._generation_coalescing_policy = "idle_wait"
+    for request in s.requests.values():
+        request.status = RequestStatus.WAITING_FOR_CHUNK
+    if pending == "in_flight":
+        s.requests["0"].num_in_flight_tokens = 1
+    elif pending == "admission":
+        s.requests["0"].status = RequestStatus.WAITING
+    elif pending == "registration":
+        s.input_coordinator.pending_chunk_registrations = [object()]
+    elif pending == "terminal":
+        s._pending_data_plane_terminal_req_ids = {"0"}
+    else:
+        s.requests.clear()
+    assert s._drain_omni_connector_outputs() == []
+    assert clock[0] == 10
+    assert not s._omni_connector_output_inbox.waits
+
+
+def test_idle_wait_without_arrivals_is_bounded(monkeypatch):
+    s, clock = scheduler(monkeypatch)
+    s._generation_coalescing_policy = "idle_wait"
+    for request in s.requests.values():
+        request.status = RequestStatus.WAITING_FOR_CHUNK
+    assert s._drain_omni_connector_outputs() == []
+    assert clock[0] == pytest.approx(10.012)
+
+
+def notification(*ids, terminal=()):
+    return OmniConnectorOutput(chunk_ready_req_ids=set(ids), chunk_finished_req_ids=set(terminal))
+
+
+class Inbox:
+    def __init__(self, clock, initial=(), delayed=()):
+        self.clock = clock
+        self.initial = list(initial)
+        self.delayed = list(delayed)
+        self.waits = []
+
+    def get_nowait(self):
+        if self.initial:
+            return self.initial.pop(0)
+        raise queue.Empty
+
+    def get(self, timeout):
+        self.waits.append(timeout)
+        if not self.delayed or self.delayed[0][0] > timeout:
+            self.clock[0] += timeout
+            raise queue.Empty
+        delay, value = self.delayed.pop(0)
+        self.clock[0] += delay
+        return value
+
+
+def scheduler(monkeypatch, initial=(), delayed=()):
+    clock = [10.0]
+    monkeypatch.setattr("vllm_omni.core.sched.omni_generation_scheduler.time.monotonic", lambda: clock[0])
+    s = OmniGenerationScheduler.__new__(OmniGenerationScheduler)
+    s._generation_max_wait_s = 0.012
+    s._generation_min_batch_size = 4
+    s._pause_state = PauseState.UNPAUSED
+    s.requests = {
+        str(i): SimpleNamespace(request_id=str(i), num_in_flight_tokens=0, is_finished=lambda: False) for i in range(8)
+    }
+    s.running = list(s.requests.values())
+    s.input_coordinator = SimpleNamespace(
+        requests_with_ready_chunks=set(), finished_requests=set(), input_terminal_req_ids=set()
+    )
+    s._latest_omni_connector_output = None
+    s._omni_connector_output_inbox = Inbox(clock, initial, delayed)
+    return s, clock
+
+
+def test_arrivals_form_target_batch_in_order(monkeypatch):
+    events = [notification("0"), notification("1", "2"), notification("3")]
+    s, clock = scheduler(monkeypatch, events[:1], [(0.002, events[1]), (0.003, events[2])])
+    assert s._drain_omni_connector_outputs() == events
+    assert clock[0] == pytest.approx(10.005)
+    assert not s.input_coordinator.requests_with_ready_chunks
+    assert all(r.num_in_flight_tokens == 0 for r in s.requests.values())
+
+
+def test_duplicate_stale_and_finished_notifications_do_not_fill_batch(monkeypatch):
+    events = [notification("0", "gone", "1", "2"), notification("0")]
+    s, clock = scheduler(monkeypatch, events[:1], [(0.003, events[1])])
+    s.requests["1"].is_finished = lambda: True
+    s.requests["2"].is_finished = lambda: True
+    assert s._drain_omni_connector_outputs() == events
+    assert clock[0] == pytest.approx(10.012)
+    assert s._omni_connector_output_inbox.waits == pytest.approx([0.012, 0.009])
+
+
+@pytest.mark.parametrize("reason", ["enough", "empty", "low", "paused", "off"])
+def test_immediate_dispatch(monkeypatch, reason):
+    event = notification("0")
+    s, clock = scheduler(monkeypatch, [event])
+    if reason == "enough":
+        event.chunk_ready_req_ids.update(["1", "2", "3"])
+    elif reason == "empty":
+        event.chunk_ready_req_ids.clear()
+    elif reason == "low":
+        s.requests = {"0": s.requests["0"]}
+    elif reason == "paused":
+        s._pause_state = PauseState.PAUSED_ALL
+    elif reason == "off":
+        s._generation_max_wait_s = 0
+    assert s._drain_omni_connector_outputs() == [event]
+    assert not s._omni_connector_output_inbox.waits
+    assert clock[0] == 10
+
+
+def test_terminal_arrival_keeps_original_deadline(monkeypatch):
+    events = [notification("0"), notification("1", terminal=["1"])]
+    s, clock = scheduler(monkeypatch, events[:1], [(0.002, events[1])])
+    assert s._drain_omni_connector_outputs() == events
+    assert clock[0] == pytest.approx(10.012)
+
+
+def test_pending_ready_and_legacy_notification_preserved(monkeypatch):
+    s, _ = scheduler(monkeypatch)
+    s.input_coordinator.requests_with_ready_chunks.update(["0", "1", "2"])
+    last = notification("3")
+    s._latest_omni_connector_output = last
+    assert s._drain_omni_connector_outputs() == [last]
+    assert s._latest_omni_connector_output is None
+    assert not s._omni_connector_output_inbox.waits
+
+
+def test_deadline_is_new_for_each_batch_not_extended_by_arrivals(monkeypatch):
+    s, clock = scheduler(monkeypatch, [notification("0")], [(0.008, notification("1")), (0.008, notification("2"))])
+    assert len(s._drain_omni_connector_outputs()) == 2
+    assert clock[0] == pytest.approx(10.012)
+    s._omni_connector_output_inbox.initial.append(notification("3"))
+    assert len(s._drain_omni_connector_outputs()) == 2
+    assert clock[0] == pytest.approx(10.024)
+
+
+def test_cancelled_notification_is_filtered_before_coordinator(monkeypatch, mocker):
+    event = notification("0", "cancelled")
+    event.request_metadata = {"0": {"code_predictor_codes": [7]}, "cancelled": {"code_predictor_codes": [9]}}
+    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()
+    s._consume_pending_connector_output("generation")
+    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()
+    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"))])
+def test_invalid_config_rejected(monkeypatch, target, wait):
+    from vllm_omni.core.sched.omni_generation_scheduler import VLLMScheduler
+
+    def initialize(s, *args, **kwargs):
+        s.max_num_running_reqs = 256
+        s.vllm_config = SimpleNamespace(
+            model_config=SimpleNamespace(
+                retains_state_across_chunks=True,
+                stage_connector_config={"extra": {"generation_min_batch_size": target, "generation_max_wait_ms": wait}},
+            )
+        )
+
+    monkeypatch.setattr(VLLMScheduler, "__init__", initialize)
+    monkeypatch.setattr(OmniGenerationScheduler, "_init_omni_io_scheduling_state", lambda s: None)
+    with pytest.raises(ValueError):
+        OmniGenerationScheduler()
+
+
+@pytest.mark.parametrize("terminal", [False, True])
+def test_first_and_terminal_chunks_have_bounded_wait(monkeypatch, terminal):
+    event = notification("0", terminal=["0"] if terminal else [])
+    s, clock = scheduler(monkeypatch, [event])
+    s.running = s.running[1:]
+    assert s._drain_omni_connector_outputs() == [event]
+    assert clock[0] == pytest.approx(10.012)
+
+
+def test_retire_in_flight_output_before_waiting_without_resetting_deadline(monkeypatch):
+    events = [notification("0"), notification("1")]
+    s, clock = scheduler(monkeypatch, events[:1])
+    s.requests["7"].num_in_flight_tokens = 1
+    assert s._drain_omni_connector_outputs() == events[:1]
+    assert s._generation_defer_batch
+    assert not s._omni_connector_output_inbox.waits
+    s.input_coordinator.requests_with_ready_chunks.add("0")
+    # Simulate EngineCore retiring the previous batch after five milliseconds.
+    clock[0] += 0.005
+    s.requests["7"].num_in_flight_tokens = 0
+    s._omni_connector_output_inbox.initial.append(events[1])
+    assert s._drain_omni_connector_outputs() == events[1:]
+    assert not s._generation_defer_batch
+    assert s._omni_connector_output_inbox.waits == pytest.approx([0.007])
+    assert clock[0] == pytest.approx(10.012)
+    assert s._generation_batch_deadline is None
+
+
+def test_expired_deadline_dispatches_even_with_other_requests_in_flight(monkeypatch):
+    s, clock = scheduler(monkeypatch, [notification("0")])
+    s.requests["7"].num_in_flight_tokens = 1
+    s._generation_batch_deadline = clock[0] - 0.001
+    assert len(s._drain_omni_connector_outputs()) == 1
+    assert not s._generation_defer_batch
+    assert not s._omni_connector_output_inbox.waits
+
+
+def test_deferred_schedule_keeps_pending_request_and_allocates_nothing(monkeypatch, mocker):
+    from tests.core.sched.test_generation_scheduler_restore import _chunk_request, _make_generation_scheduler
+
+    request = _chunk_request("pending")
+    s = _make_generation_scheduler(request)
+    s._process_pending_omni_inputs = lambda **kwargs: setattr(s, "_generation_defer_batch", True)
+    allocate = mocker.spy(s.kv_cache_manager, "allocate_slots")
+    output = s.schedule()
+    assert output.total_num_scheduled_tokens == 0
+    assert not output.num_scheduled_tokens
+    allocate.assert_not_called()
+    assert list(s.waiting) == [request]
+    assert request.num_computed_tokens == request.num_in_flight_tokens == 0
+    assert s.chunk_transfer_adapter.restore_called

tests/core/sched/test_generation_scheduler_init.py

@@ -0,0 +1,107 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Exercise generation constructor policy after actual deploy/platform resolution."""
+
+from types import SimpleNamespace
+
+import pytest
+
+from tests.helpers.stage_config import get_deploy_config_path
+from vllm_omni.config.pipeline_registry import resolve_pipeline_config
+from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
+from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler, VLLMScheduler
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.fixture
+def construct_scheduler(monkeypatch, mocker):
+    # Only upstream cache allocation and connector construction are omitted.
+    # Keep the actual Omni I/O initialization, native decision and constructor.
+    mocker.patch("vllm_omni.core.sched.omni_scheduler_mixin.OmniChunkTransferAdapter")
+
+    def upstream_init(self, config):
+        self.vllm_config = config
+        self.max_num_running_reqs = config.scheduler_config.max_num_seqs
+
+    monkeypatch.setattr(VLLMScheduler, "__init__", upstream_init)
+    return OmniGenerationScheduler
+
+
+def _config(*, native=True, stateful=True, tp=1, pp=1, extras=None, capacity=128):
+    return SimpleNamespace(
+        model_config=SimpleNamespace(
+            use_v2_model_runner=native,
+            supports_native_mrv2_data_plane=True,
+            retains_state_across_chunks=stateful,
+            async_chunk=True,
+            stage_id=1,
+            stage_connector_config={"extra": extras or {}},
+        ),
+        parallel_config=SimpleNamespace(tensor_parallel_size=tp, pipeline_parallel_size=pp),
+        scheduler_config=SimpleNamespace(max_num_seqs=capacity),
+    )
+
+
+@pytest.mark.parametrize("profile", ["high_concurrency", "low_latency"])
+@pytest.mark.parametrize("platform", ["cuda", "npu", "xpu", "rocm", "musa"])
+def test_moss_profile_generation_constructor_after_platform_resolution(
+    construct_scheduler, monkeypatch, mocker, profile, platform
+):
+    # merge_pipeline_deploy owns platform resolution. Mock the detected device
+    # instead of applying an override first and then resolving the worker's
+    # actual platform a second time.
+    monkeypatch.setattr(current_omni_platform, "device_name", platform)
+    deploy = load_deploy_config(get_deploy_config_path(f"moss_tts_local_mrv2_{profile}.yaml"))
+    pipeline = resolve_pipeline_config("moss_tts_local")
+    codec = merge_pipeline_deploy(pipeline, deploy)[1]
+    extras = deploy.connectors["shm"]["extra"]
+    config = _config(
+        native=codec.yaml_engine_args["use_v2_model_runner"],
+        stateful=pipeline.stages[1].retains_state_across_chunks,
+        capacity=codec.yaml_engine_args["max_num_seqs"],
+        extras=extras,
+    )
+    warning = mocker.patch("vllm_omni.core.sched.omni_generation_scheduler.logger.warning")
+    scheduler = construct_scheduler(config)
+    # The raw extras persist through fallback; the live constructor must cope.
+    assert extras["generation_min_batch_size"] == 32 and extras["generation_max_wait_ms"] == 12
+    if platform == "cuda":
+        assert scheduler._native_data_plane and scheduler.input_coordinator is not None
+        assert scheduler._generation_min_batch_size == 32
+        assert scheduler._generation_max_wait_s == 0.012
+        assert scheduler._generation_max_regular_batch == (16 if profile == "low_latency" else 0)
+        warning.assert_not_called()
+        if profile == "low_latency":
+            assert extras["codec_first_chunk_fast_path"] == extras["codec_first_chunk_gate"] == 1
+    else:
+        assert not scheduler._native_data_plane and scheduler.chunk_transfer_adapter is not None
+        assert scheduler._generation_min_batch_size == 1
+        assert scheduler._generation_max_wait_s == scheduler._generation_max_regular_batch == 0
+        warning.assert_called_once()
+
+
+@pytest.mark.parametrize("stateful,tp,pp", [(False, 1, 1), (True, 2, 1), (True, 1, 2)])
+def test_native_batch_waiting_still_rejects_unsupported_state_or_topology(construct_scheduler, stateful, tp, pp):
+    config = _config(
+        stateful=stateful, tp=tp, pp=pp, extras={"generation_min_batch_size": 32, "generation_max_wait_ms": 12}
+    )
+    with pytest.raises(ValueError, match="stateful native MRV2 TP1/PP1"):
+        construct_scheduler(config)
+
+
+@pytest.mark.parametrize(
+    "extras,message",
+    [
+        ({"generation_min_batch_size": 0}, "generation_min_batch_size"),
+        ({"generation_min_batch_size": 129}, "generation_min_batch_size"),
+        ({"generation_max_wait_ms": -1}, "generation_max_wait_ms"),
+        ({"generation_max_wait_ms": float("nan")}, "generation_max_wait_ms"),
+        ({"generation_max_regular_batch": -1}, "generation_max_regular_batch"),
+        ({"generation_coalescing_policy": "unknown"}, "generation_coalescing_policy"),
+    ],
+)
+def test_non_native_fallback_still_rejects_invalid_values(construct_scheduler, extras, message):
+    with pytest.raises(ValueError, match=message):
+        construct_scheduler(_config(native=False, extras=extras))

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

@@ -1,10 +1,16 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from types import SimpleNamespace
 
 import pytest
 from vllm import SamplingParams
 from vllm.v1.engine import EngineCoreRequest
+from vllm.v1.request import RequestStatus
 
+from tests.helpers.omni_scheduler import bind_omits_transfer_helpers
 from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
 from vllm_omni.engine.async_engine_utils import apply_omni_final_stage_metadata
 from vllm_omni.engine.serialization import deserialize_additional_information
 
@@ -25,7 +31,7 @@ def _make_engine_request() -> EngineCoreRequest:
     )
 
 
-def _request_omits_kv_transfer(*, force_kv_transfer: bool) -> tuple[bool, dict]:
+def _omit_flags(*, force_kv_transfer: bool) -> tuple[bool, bool, dict]:
     tagged = apply_omni_final_stage_metadata(
         _make_engine_request(),
         final_stage_id=0,
@@ -37,21 +43,187 @@ def _request_omits_kv_transfer(*, force_kv_transfer: bool) -> tuple[bool, dict]:
         request_id="req",
         additional_information=tagged.additional_information,
     )
-    result = scheduler._request_omits_kv_transfer_to_next_stage(request)
+    omits_kv = scheduler._request_omits_kv_transfer_to_next_stage(request)
+    omits_chunk = scheduler._request_omits_chunk_transfer_to_next_stage(request)
     metadata = deserialize_additional_information(tagged.additional_information)
-    return result, metadata
+    return omits_kv, omits_chunk, metadata
 
 
 def test_stage_zero_request_omits_kv_transfer():
-    omits_transfer, metadata = _request_omits_kv_transfer(force_kv_transfer=False)
+    omits_kv, omits_chunk, metadata = _omit_flags(force_kv_transfer=False)
 
-    assert omits_transfer
+    assert omits_kv
+    assert omits_chunk
     assert "omni_force_kv_transfer" not in metadata
 
 
 def test_cfg_companion_forces_kv_transfer_without_downstream_payload():
-    omits_transfer, metadata = _request_omits_kv_transfer(force_kv_transfer=True)
+    omits_kv, omits_chunk, metadata = _omit_flags(force_kv_transfer=True)
 
-    assert not omits_transfer
+    assert not omits_kv
+    assert omits_chunk
     assert metadata["omni_final_stage_id"] == 0
     assert metadata["omni_force_kv_transfer"] is True
+
+
+class _ChunkRequest(SimpleNamespace):
+    def __hash__(self):
+        return hash(self.request_id)
+
+    def __eq__(self, other):
+        return isinstance(other, _ChunkRequest) and other.request_id == self.request_id
+
+
+def _make_chunk_request(final_stage_id: int, *, force_kv_transfer: bool = False):
+    tagged = apply_omni_final_stage_metadata(
+        _make_engine_request(),
+        final_stage_id=final_stage_id,
+        force_kv_transfer=force_kv_transfer,
+    )
+    request = _ChunkRequest(
+        request_id="req-shm-leak",
+        additional_information=tagged.additional_information,
+        num_in_flight_tokens=1,
+        num_stale_output_tokens=0,
+        sampling_params=SimpleNamespace(num_logprobs=None),
+        has_encoder_inputs=False,
+        pooling_params=None,
+        status=RequestStatus.RUNNING,
+        resumable=False,
+        num_output_placeholders=0,
+        spec_token_ids=[],
+        _output_token_ids=[],
+        num_computed_tokens=1,
+        stop_reason=None,
+        num_nans_in_logits=None,
+        client_index=0,
+        trace_headers=None,
+    )
+    request.is_finished = lambda: False
+    request.get_finished_reason = lambda: None
+    request.take_prefill_stats = lambda: None
+    request.take_events = lambda: None
+    return request
+
+
+def _run_finished_save_step(mocker, request, *, inter_stage_outputs=None):
+    adapter = mocker.MagicMock()
+    adapter._confirmed_num_computed_tokens.return_value = 1
+
+    sched = mocker.MagicMock()
+    sched.requests = {request.request_id: request}
+    sched.perf_metrics = None
+    sched.defer_block_free = False
+    sched.structured_output_manager.should_advance.return_value = False
+    sched._update_request_with_output.return_value = ([42], True)
+    sched._process_kv_transfer_trigger.return_value = False
+    sched._handle_stopped_request.return_value = True
+
+    def _free_request(req, delay_free_blocks=False):
+        sched._omits_kv_transfer_cache.pop(req.request_id, None)
+        return None, None
+
+    sched._free_request.side_effect = _free_request
+    sched.chunk_transfer_adapter = adapter
+    sched.running = [request]
+    sched.waiting_for_transfer_free = set()
+    sched.transfer_triggered_requests = set()
+    sched.active_kv_transfers = set()
+    sched.pending_stop_after_extraction = set()
+    sched.connector = None
+    sched.kv_cache_manager.take_events.return_value = None
+    sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
+    sched.finished_req_ids_dict = {}
+    sched.make_stats.return_value = None
+    bind_omits_transfer_helpers(sched)
+
+    scheduler_output = mocker.MagicMock()
+    scheduler_output.num_scheduled_tokens = {request.request_id: 1}
+    scheduler_output.total_num_scheduled_tokens = 1
+    scheduler_output.scheduled_spec_decode_tokens = {}
+    scheduler_output.num_invalid_spec_tokens = 0
+
+    model_runner_output = mocker.MagicMock()
+    model_runner_output.sampled_token_ids = [[42]]
+    model_runner_output.logprobs = None
+    model_runner_output.prompt_logprobs_dict = {}
+    model_runner_output.pooler_output = None
+    model_runner_output.num_nans_in_logits = None
+    model_runner_output.kv_connector_output = None
+    model_runner_output.cudagraph_stats = None
+    model_runner_output.req_id_to_index = {request.request_id: 0}
+    model_runner_output.routed_experts = None
+    model_runner_output.inter_stage_outputs = inter_stage_outputs
+
+    OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output)
+    return adapter
+
+
+@pytest.mark.parametrize(
+    ("final_stage_id", "force_kv_transfer", "expect_save"),
+    [
+        pytest.param(0, False, False, id="stage-zero-final"),
+        pytest.param(1, False, True, id="downstream-final"),
+    ],
+)
+def test_finish_saves_async_only_with_downstream_consumer(mocker, final_stage_id, force_kv_transfer, expect_save):
+    request = _make_chunk_request(final_stage_id, force_kv_transfer=force_kv_transfer)
+    adapter = _run_finished_save_step(mocker, request)
+
+    assert adapter.save_async.called is expect_save
+
+
+def test_stage_zero_final_skips_inter_stage_output(mocker):
+    warn = mocker.patch("vllm_omni.core.sched.omni_ar_scheduler.logger.warning")
+    request = _make_chunk_request(0)
+    adapter = _run_finished_save_step(mocker, request, inter_stage_outputs=[{"codes": 1}])
+
+    adapter.save_async.assert_not_called()
+    assert any("inter_stage_output is present" in str(call) for call in warn.call_args_list)
+
+
+def test_streaming_session_update_invalidates_omits_kv_transfer_cache():
+    scheduler = OmniARScheduler.__new__(OmniARScheduler)
+    scheduler._omits_kv_transfer_cache = {}
+    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)
+
+    class Session(SimpleNamespace):
+        __hash__ = object.__hash__
+
+    session = Session(
+        request_id="req",
+        additional_information=tagged.additional_information,
+        status=RequestStatus.WAITING,
+    )
+    assert scheduler._request_omits_chunk_transfer_to_next_stage(session)
+    assert "req" in scheduler._omits_kv_transfer_cache
+    update = SimpleNamespace(
+        additional_information={"omni_final_stage_id": 1},
+        model_intermediate_buffer=None,
+        arrival_time=0.0,
+        sampling_params=None,
+    )
+
+    OmniSchedulerMixin._finish_streaming_session_update(scheduler, session, update)
+
+    assert scheduler._omits_kv_transfer_cache == {}
+    assert session.additional_information == {"omni_final_stage_id": 1}
+    assert not scheduler._request_omits_chunk_transfer_to_next_stage(session)
+    assert not scheduler._request_omits_kv_transfer_to_next_stage(session)
+
+
+def test_omits_cache_misses_when_payload_object_is_replaced():
+    scheduler = OmniARScheduler.__new__(OmniARScheduler)
+    scheduler._omits_kv_transfer_cache = {}
+    tagged = apply_omni_final_stage_metadata(_make_engine_request(), final_stage_id=0)
+    request = SimpleNamespace(
+        request_id="req",
+        additional_information=tagged.additional_information,
+    )
+    assert scheduler._request_omits_chunk_transfer_to_next_stage(request)
+    request.additional_information = {"omni_final_stage_id": 1}
+
+    assert not scheduler._request_omits_chunk_transfer_to_next_stage(request)

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

@@ -40,6 +40,7 @@ from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.outputs import ModelRunnerOutput
 from vllm.v1.request import Request, RequestStatus, StreamingUpdate
 from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+from tests.helpers.omni_scheduler import bind_omits_transfer_helpers
 
 # isort: on
 
@@ -72,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,
@@ -103,10 +105,11 @@ def _make_drain_sched(session: Request) -> MagicMock:
     sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
     sched.finished_req_ids_dict = {}
     sched.make_stats.return_value = None
+    bind_omits_transfer_helpers(sched)
     return sched
 
 
-def _run_step(sched: MagicMock, session: Request, *, num_scheduled: int, token: int) -> bool:
+def _run_step(sched: MagicMock, session: Request, *, num_scheduled: int, token: int, spec=None) -> bool:
     """Feed one model-runner frame through update_from_output.
 
     Returns True when the frame's tokens were delivered (the append ran),
@@ -114,7 +117,7 @@ def _run_step(sched: MagicMock, session: Request, *, num_scheduled: int, token:
     """
     scheduler_output = MagicMock(spec=SchedulerOutput)
     scheduler_output.num_scheduled_tokens = {session.request_id: num_scheduled}
-    scheduler_output.scheduled_spec_decode_tokens = {}
+    scheduler_output.scheduled_spec_decode_tokens = {session.request_id: spec} if spec else {}
     scheduler_output.num_invalid_spec_tokens = 0
 
     model_runner_output = MagicMock(spec=ModelRunnerOutput)
@@ -134,6 +137,27 @@ def _run_step(sched: MagicMock, session: Request, *, num_scheduled: int, token:
     return sched._update_request_with_output.called
 
 
+def test_connector_prompt_replacement_drops_old_frame_and_delivers_new_frame() -> None:
+    from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
+
+    session = _make_session()
+    session.num_computed_tokens = 6
+    session.num_in_flight_tokens = 1
+    session.num_output_placeholders = 1
+    sched = _make_drain_sched(session)
+    sched.chunk_transfer_adapter.replaced_streaming_prompt_ids = {session.request_id}
+    sched.chunk_transfer_adapter.requests_with_ready_chunks = {session.request_id}
+    sched.chunk_transfer_adapter.requests_num_chunks_sent = {session.external_req_id: 2}
+
+    OmniSchedulerMixin._reset_ready_async_chunk_replacements(sched)
+    assert session.num_stale_output_tokens == 1 and session.drop_stale_output
+    assert session.num_output_placeholders == 0
+    sched._release_replaced_streaming_prompt_cache.assert_called_once_with(session)
+    assert _run_step(sched, session, num_scheduled=1, token=42) is False
+    assert session.num_stale_output_tokens == 0
+    assert _run_step(sched, session, num_scheduled=1, token=43) is True
+
+
 def test_exact_drain_delivers_new_segment_frame() -> None:
     """One in-flight decode frame at replacement: its late output drains the
     counter exactly and the new segment's first frame is delivered."""
@@ -225,3 +249,80 @@ def test_in_flight_prefill_chunk_drains_exactly_without_underflow() -> None:
     assert _run_step(sched, session, num_scheduled=5, token=42) is False  # late prefill chunk dropped
     assert session.num_stale_output_tokens == 0
     assert _run_step(sched, session, num_scheduled=1, token=43) is True  # new segment survives
+
+
+@pytest.mark.parametrize("drop", [False, True])
+def test_preemption_stale_output_follows_upstream_delivery_policy(drop):
+    request = _make_session()
+    request.status = RequestStatus.PREEMPTED
+    request.num_stale_output_tokens = 1
+    request.num_in_flight_tokens = 1
+    request.drop_stale_output = drop
+    sched = _make_drain_sched(request)
+    assert _run_step(sched, request, num_scheduled=1, token=42) is (not drop)
+    assert request.num_stale_output_tokens == 0
+    if not drop:
+        sched._update_request_with_output.assert_called_once_with(request, [42], is_stale=True)
+
+
+def test_preemption_stale_spec_rejection_does_not_roll_back_resumed_counters():
+    request = _make_session()
+    request.num_stale_output_tokens = 2
+    request.num_in_flight_tokens = 2
+    request.num_computed_tokens = 3
+    request.num_output_placeholders = 2
+    request.drop_stale_output = False
+    sched = _make_drain_sched(request)
+    assert _run_step(sched, request, num_scheduled=2, token=42, spec=[41])
+    assert request.num_computed_tokens == 3
+    assert request.num_output_placeholders == 2
+
+
+@pytest.mark.parametrize(
+    "arch,supports_reset,running,reset_running,blocked",
+    [
+        ("Qwen3TTSTalkerForConditionalGeneration", False, True, True, True),
+        ("Qwen3TTSTalkerForConditionalGeneration", False, False, True, False),
+        ("Qwen3TTSTalkerForConditionalGeneration", False, True, False, False),
+        ("Qwen3TTSTalkerForConditionalGeneration", True, True, True, False),
+        ("AnotherStatefulAudioModel", False, True, True, True),
+        ("AnotherStatefulAudioModel", True, True, True, False),
+    ],
+)
+def test_reset_running_streaming_codec_is_rejected_before_preemption(
+    mocker, arch, supports_reset, running, reset_running, blocked
+):
+    from vllm.config import VllmConfig
+    from vllm.v1.core.sched.scheduler import Scheduler
+
+    from vllm_omni.config.model import OmniModelConfig
+
+    scheduler = OmniARScheduler.__new__(OmniARScheduler)
+    scheduler.vllm_config = mocker.Mock(
+        spec=VllmConfig,
+        model_config=mocker.Mock(
+            spec=OmniModelConfig,
+            model_arch=arch,
+            supports_running_prefix_cache_reset=supports_reset,
+        ),
+    )
+    request = _make_session()
+    request.num_in_flight_tokens = 1
+    request.num_output_placeholders = 1
+    scheduler.running = [request] if running else []
+
+    def reset_cache(_self, reset_running_requests=False, reset_connector=False):
+        assert reset_running_requests is reset_running and reset_connector is True
+        return True
+
+    reset = mocker.patch.object(Scheduler, "reset_prefix_cache", autospec=True, side_effect=reset_cache)
+
+    assert scheduler.reset_prefix_cache(reset_running, reset_connector=True) is (not blocked)
+    if blocked:
+        reset.assert_not_called()
+        assert scheduler.running == [request]
+        assert request.status == RequestStatus.RUNNING
+        assert request.num_in_flight_tokens == request.num_output_placeholders == 1
+        assert not request.drop_stale_output
+    else:
+        reset.assert_called_once()

tests/core/sched/test_omni_ar_scheduler_streaming.py

@@ -23,9 +23,11 @@ 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,
 )
+from tests.helpers.omni_scheduler import bind_omits_transfer_helpers
 
 # isort: on
 
@@ -41,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()
@@ -175,6 +178,7 @@ def _run_resumable_segment_stop(
     sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
     sched.finished_req_ids_dict = {}
     sched.make_stats.return_value = None
+    bind_omits_transfer_helpers(sched)
 
     scheduler_output = MagicMock(spec=SchedulerOutput)
     scheduler_output.num_scheduled_tokens = {session.request_id: 1}
@@ -285,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
@@ -294,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()
@@ -312,14 +323,15 @@ def test_running_decode_step_without_inter_stage_payload_does_not_raise() -> Non
     sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
     sched.finished_req_ids_dict = {}
     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
@@ -332,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:
@@ -400,6 +420,7 @@ def test_stale_async_frame_is_dropped_before_output_processing() -> None:
     sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
     sched.finished_req_ids_dict = {}
     sched.make_stats.return_value = None
+    bind_omits_transfer_helpers(sched)
 
     scheduler_output = MagicMock(spec=SchedulerOutput)
     scheduler_output.num_scheduled_tokens = {session.request_id: 1}
@@ -1165,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,
@@ -1177,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
 
@@ -1204,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)
@@ -1212,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()
 
 
@@ -1268,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
 
 
@@ -1291,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}
@@ -1399,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
@@ -1428,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]
@@ -1447,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()
@@ -1469,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()
@@ -1492,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()
@@ -1525,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)
@@ -1551,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)
@@ -1559,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,13 +123,21 @@ 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
+    # Unset MagicMock attributes are truthy. Keep optional scheduling modes
+    # disabled so this fixture actually schedules its one-shot request.
+    sched._generation_defer_batch = False
+    sched._generation_first_chunk_only = False
+    sched._generation_max_regular_batch = 0
     sched._pause_state = PauseState.UNPAUSED
     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

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from __future__ import annotations
 
 from types import SimpleNamespace
@@ -49,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"]
@@ -76,7 +81,7 @@ def test_ar_free_request_cleans_input_coordinator_on_normal_free() -> None:
     scheduler = OmniARScheduler.__new__(OmniARScheduler)
     scheduler.input_coordinator = coordinator
     scheduler.chunk_transfer_adapter = None
-    scheduler._omits_kv_transfer_cache = {"req-free": True}
+    scheduler._omits_kv_transfer_cache = {"req-free": (1, True, False)}
     scheduler.encoder_cache_manager = SimpleNamespace(free=lambda request: None)
     scheduler.finished_req_ids = set()
     scheduler.finished_req_ids_dict = None

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

@@ -376,10 +376,17 @@ def test_lru_eviction_releases_blocks_and_notifies_pipeline():
 
 def test_budget_reduced_capacity_drives_runner_lru():
     pipeline = CapablePipeline(tiny_spec(capacity=2))
-    # One tiny LingBot-like session requires 128 bytes; two require 192.
+    # One tiny LingBot-like session requires 1,808 bytes; two require 2,592.
+    #
+    # These read 128 and 192 while the paging unit was the frame. tiny_spec
+    # declares one token per frame, so it used to page a single token at a
+    # time -- a block size no attention kernel accepts. Paging at 16 makes
+    # each page sixteen times larger. The block *counts* are unchanged, and
+    # so is what this test is about: a budget that fits fewer sessions than
+    # the pipeline asked for must drive the runner's LRU.
     runner = make_runner(
         pipeline,
-        available_bytes=128,
+        available_bytes=2048,
         gpu_memory_fraction=1.0,
     )
     kv = runner.kv_cache
@@ -665,3 +672,106 @@ def test_pipeline_without_warmup_provider_is_safely_skipped(monkeypatch):
     monkeypatch.setattr(runner, "execute_model", fail_if_called)
     runner._warmup_ar_rollout()
     assert execute.called is False
+
+
+# ── paging unit vs eviction unit ────────────────────────────────────────────
+
+
+def test_a_frame_that_is_already_a_legal_block_keeps_its_paging():
+    """Every resolution that runs today must page exactly as it does today."""
+    from vllm_omni.experimental.ar_diffusion.runner import paging_block_size
+
+    for tokens_per_frame in (16, 320, 1440, 4096):
+        assert tokens_per_frame % 16 == 0
+        assert paging_block_size(tokens_per_frame) == tokens_per_frame
+
+
+def test_a_frame_the_kernel_would_reject_pages_at_the_finest_legal_unit():
+    """832x480 is the checkpoint's own default and the kernel rejects it.
+
+    (480/16) x (832/16) = 1560 tokens per frame, and 1560 % 16 == 8, so a
+    frame-sized block is not something FlashAttention's paged kernel accepts.
+    """
+    from vllm_omni.experimental.ar_diffusion.runner import paging_block_size
+
+    assert 1560 % 16 == 8
+    assert paging_block_size(1560) == 16
+    # And whatever it returns is always something the kernel takes.
+    for tokens_per_frame in range(1, 200):
+        assert paging_block_size(tokens_per_frame) % 16 == 0
+
+
+# ── the kernel's constraint is data, not a constant ─────────────────────────
+
+
+def _multiple_of(base):
+    from vllm.v1.attention.backend import MultipleOf
+
+    return MultipleOf(base)
+
+
+def test_the_dispatching_module_answers_in_vllms_own_vocabulary():
+    """AR-Diffusion calls flash_attn itself, so no backend can be asked.
+
+    The answer has to come from the module that picks the kernel, and it has to
+    have the shape vLLM uses, so the policy that consumes it does not care
+    where it came from.
+    """
+    from vllm.v1.attention.backend import MultipleOf
+
+    from vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention import supported_kernel_block_sizes
+
+    advertised = supported_kernel_block_sizes()
+    assert advertised, "a kernel that accepts nothing cannot be paged for"
+    for entry in advertised:
+        assert isinstance(entry, int | MultipleOf)
+
+
+def test_a_kernel_that_takes_one_fixed_size_gets_that_size():
+    """hpc_attn advertises [64] -- not a multiple, a single legal value.
+
+    A frame of 1560 tokens is not 64, and 1560 % 64 == 24, so the frame cannot
+    be the block and the only thing left to page at is 64. Assuming multiples
+    of 16 here would hand the kernel a block size it rejects.
+    """
+    from vllm_omni.experimental.ar_diffusion.runner import paging_block_size
+
+    assert 1560 % 64 == 24
+    assert paging_block_size(1560, [64]) == 64
+    assert paging_block_size(64, [64]) == 64
+    assert paging_block_size(128, [64]) == 64  # a legal multiple is still not the advertised size
+
+
+def test_a_kernel_with_large_pages_still_keeps_a_whole_frame_when_it_fits():
+    """FlashInfer offers 128-and-up pages on Blackwell and not on Hopper.
+
+    The same checkpoint on the two cards must therefore page differently, which
+    is exactly what a constant cannot express.
+    """
+    from vllm_omni.experimental.ar_diffusion.runner import paging_block_size
+
+    hopper = [_multiple_of(16)]
+    blackwell = [_multiple_of(16), 128, 256]
+    assert paging_block_size(1440, hopper) == 1440  # a legal multiple of 16
+    assert paging_block_size(1440, blackwell) == 1440
+    assert paging_block_size(1560, hopper) == 16
+    assert paging_block_size(1560, blackwell) == 16  # finest legal unit, not the largest
+
+
+def test_whatever_comes_back_is_always_something_the_kernel_accepts():
+    from vllm_omni.experimental.ar_diffusion.runner import paging_block_size
+
+    for advertised, check in (
+        ([_multiple_of(16)], lambda n: n % 16 == 0),
+        ([64], lambda n: n == 64),
+        ([_multiple_of(32), 48], lambda n: n % 32 == 0 or n == 48),
+    ):
+        for tokens_per_frame in range(1, 300):
+            assert check(paging_block_size(tokens_per_frame, advertised))
+
+
+def test_a_kernel_that_advertises_nothing_is_an_error_not_a_guess():
+    from vllm_omni.experimental.ar_diffusion.runner import paging_block_size
+
+    with pytest.raises(ValueError, match="no legal block size"):
+        paging_block_size(1560, [])

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

@@ -625,6 +625,74 @@ def test_non_contiguous_branch_indices_rejected():
         )
 
 
+# --- eviction against the block table ---------------------------------------
+
+
+class _RecordingManager(ChunkWindowManager):
+    """Capture the block range eviction asks for, without a real block pool.
+
+    ``remove_skipped_blocks`` reads three attributes and calls one method. A
+    pool would add nothing: the bug under test is the conversion from the
+    spec's units into the block table's, and that is entirely visible here.
+    """
+
+    def __init__(self, spec):
+        self.kv_cache_spec = spec
+        self.sliding_window = spec.sliding_window
+        self.block_size = spec.block_size
+        self.freed: tuple[int, int] | None = None
+
+    def _remove_blocks_in_range(self, request_id, first_block, last_block):
+        self.freed = (first_block, last_block)
+
+
+def test_sink_is_preserved_when_a_frame_spans_many_blocks():
+    # The shipped geometry: 832x480 is 1560 tokens per frame, 1560 % 16 == 8 is
+    # not a legal block size, so paging falls back to 16 and a frame spans 98
+    # blocks. Every other eviction test uses chunk_size == BLOCK, where a frame
+    # is a block and a sink counted in frames indexes the block table correctly
+    # by accident -- which is why this went unnoticed.
+    chunk_size, sink_chunks, window_chunks = 1560, 9, 9
+    spec = ChunkWindowSpec(
+        block_size=16,
+        num_kv_heads=4,
+        head_size=64,
+        dtype=torch.float16,
+        sliding_window=window_chunks * chunk_size,
+        chunk_size=chunk_size,
+        window_chunks=window_chunks,
+        sink_chunks=sink_chunks,
+    )
+    manager = _RecordingManager(spec)
+
+    # Far enough past sink + window that eviction has definitely started.
+    manager.remove_skipped_blocks("req", total_computed_tokens=30 * chunk_size)
+    assert manager.freed is not None
+    first_freed, last_freed = manager.freed
+
+    # The assertion is on the token boundary rather than a block count, so it
+    # stays meaningful if the paging unit changes again.
+    sink_tokens = sink_chunks * chunk_size
+    assert first_freed * spec.block_size >= sink_tokens, (
+        f"eviction starts at token {first_freed * spec.block_size}, inside the {sink_tokens}-token sink"
+    )
+    # And it must not over-protect: at most one straddling block beyond the sink.
+    assert (first_freed - 1) * spec.block_size < sink_tokens
+    # Only wholly-skipped blocks are released.
+    skipped = manager.get_num_skipped_tokens(30 * chunk_size)
+    assert last_freed * spec.block_size <= sink_tokens + skipped
+
+
+def test_sink_conversion_is_identity_when_a_frame_is_a_block():
+    # Every geometry that paged one frame per block must be untouched.
+    spec = make_spec(chunk_size=BLOCK, window_chunks=2, sink_chunks=1)
+    manager = _RecordingManager(spec)
+    manager.remove_skipped_blocks("req", total_computed_tokens=8 * BLOCK)
+    assert manager.freed is not None
+    first_freed, _ = manager.freed
+    assert first_freed == spec.sink_chunks
+
+
 def test_staged_window_reuse_matches_a_full_gather():
     """Reusing the staged history must be byte-identical to re-gathering the whole window.
 

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,
@@ -23,9 +25,10 @@ from vllm_omni.experimental.ar_diffusion.kv_cache import (
 )
 from vllm_omni.experimental.ar_diffusion.kv_cache import paged_attention as paged_attention_module
 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
@@ -36,10 +39,26 @@ NEG = "negative"
 
 
 def make_state(
-    *, num_layers=1, window_chunks=2, dtype=torch.float32, device=torch.device("cpu"), reuse_history_staging=False
+    *,
+    num_layers=1,
+    window_chunks=2,
+    chunk_size=BLOCK,
+    sink_chunks=0,
+    reset_at_boundary=False,
+    dtype=torch.float32,
+    device=torch.device("cpu"),
+    reuse_history_staging=False,
 ):
+    """Build a cache. ``chunk_size`` defaults to the block size, but the two are
+    independent -- the shipped 832x480 gives 1560 tokens per frame against
+    16-token blocks, so a frame is 97.5 blocks."""
     cfg = ARDiffusionKVConfig(
-        enable=True, chunk_size=BLOCK, window_chunks=window_chunks, reuse_history_staging=reuse_history_staging
+        enable=True,
+        chunk_size=chunk_size,
+        window_chunks=window_chunks,
+        sink_chunks=sink_chunks,
+        reset_at_boundary=reset_at_boundary,
+        reuse_history_staging=reuse_history_staging,
     )
     kv = ARDiffusionKVCache(
         cfg,
@@ -101,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,
@@ -109,17 +139,20 @@ def _commit_video_span(
     n_chunks: int,
     dtype: torch.dtype,
     device: torch.device,
+    chunk_size: int = BLOCK,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    ctx = st.get_kv_caches(kv_branch, seq_len=n_chunks * BLOCK, commit_current=True)[0].forward_ctx
+    span = n_chunks * chunk_size
+    ctx = st.get_kv_caches(kv_branch, seq_len=span, commit_current=True)[0].forward_ctx
     ctx.ensure_video_slots(device)
-    k = torch.randn(1, n_chunks * BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
-    v = torch.randn(1, n_chunks * BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    k = torch.randn(1, span, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    v = torch.randn(1, span, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
     kv._k_pools[0][ctx.current_video_slot_mapping] = k[0]
     kv._v_pools[0][ctx.current_video_slot_mapping] = v[0]
     st.commit_paged_context(kv_branch)
     return k, v
 
 
+@pytest.mark.cpu
 def test_paged_context_allocates_lazily_and_commits_after_forward():
     _, st = make_state()
 
@@ -138,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()
 
@@ -152,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")
@@ -169,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")
@@ -245,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.
 
@@ -311,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")
@@ -332,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")
@@ -349,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
@@ -420,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")
@@ -443,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)
@@ -468,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]
@@ -475,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
@@ -489,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.
@@ -499,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."""
@@ -547,7 +593,657 @@ def test_custom_op_compiles_fullgraph_without_recompile_on_value_change(monkeypa
         torch._dynamo.reset()
 
 
-@pytest.mark.skipif(not _gpu_flash_attn_usable(), reason="usable GPU FlashAttention is required")
+# ── a frame that is not a whole number of blocks ────────────────────────────
+
+# 24 tokens per chunk against 16-token blocks leaves 8 over, the same remainder
+# the shipped 832x480 produces (1560 % 16 == 8). Every fixture above holds
+# chunk_size == BLOCK, so nothing there can reach this path.
+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.
+
+    With one committed chunk of 24 tokens the history stops 8 slots into its
+    last block. The committing path writes straight after it and is fine. The
+    scratch path used to restart at slot zero of a fresh region, which left
+    those 8 slots unwritten -- and since the kernel consumes the block table as
+    one contiguous run, it read them as if they were tokens and shifted every
+    token after them by 8. Shapes stayed correct throughout, so only the values
+    showed it.
+    """
+    torch.manual_seed(0)
+    device = torch.device("cpu")
+    dtype = torch.float32
+    # window_chunks=2 with one chunk of history keeps the whole sequence
+    # resident, so the window does not also round up here; that case is
+    # test_a_ragged_window_keeps_the_block_table_a_fixed_shape below.
+    kv, st = make_state(dtype=dtype, device=device, window_chunks=2, chunk_size=RAGGED_CHUNK)
+
+    history_k, history_v = _commit_video_span(
+        kv, st, kv_branch=POS, n_chunks=1, dtype=dtype, device=device, chunk_size=RAGGED_CHUNK
+    )
+
+    ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=commit_current)[0].forward_ctx
+    ctx.ensure_video_slots(device)
+    assert ctx.start_offset == RAGGED_CHUNK % BLOCK == 8
+
+    current_k = torch.randn(1, RAGGED_CHUNK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    current_v = torch.randn(1, RAGGED_CHUNK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    kv._k_pools[0][ctx.current_video_slot_mapping] = current_k[0]
+    kv._v_pools[0][ctx.current_video_slot_mapping] = current_v[0]
+
+    action_k = action_v = None
+    if action_len:
+        ctx.ensure_action_slots(action_len, device)
+        action_k = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+        action_v = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+        kv._k_pools[0][ctx.action_slot_mapping] = action_k[0]
+        kv._v_pools[0][ctx.action_slot_mapping] = action_v[0]
+
+    query = torch.randn(1, RAGGED_CHUNK + action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    block_table, query_start_loc, seq_lens, max_query_len, max_seq_len = ctx.build_block_table(
+        action_len=action_len, query_len=query.shape[1], device=device
+    )
+    paged = ar_diffusion_paged_attention(
+        query,
+        kv.key_cache(0),
+        kv.value_cache(0),
+        block_table=block_table,
+        query_start_loc=query_start_loc,
+        seq_lens=seq_lens,
+        max_query_len=max_query_len,
+        max_seq_len=max_seq_len,
+        softmax_scale=HEAD_DIM**-0.5,
+        causal=False,
+    )
+
+    new_k = torch.cat([history_k, current_k], dim=1)
+    new_v = torch.cat([history_v, current_v], dim=1)
+    if action_len:
+        new_k = torch.cat([new_k, action_k], dim=1)
+        new_v = torch.cat([new_v, action_v], dim=1)
+    ref = _dense_attention(query, new_k, new_v)
+
+    torch.testing.assert_close(paged, ref, rtol=1e-5, atol=1e-5)
+
+
+@pytest.mark.cpu
+def test_a_ragged_scratch_chunk_is_physically_continuous_with_its_history():
+    """The write targets themselves must leave no hole.
+
+    This is the property the values above depend on, asserted directly so a
+    regression names the cause rather than showing a numeric mismatch.
+    """
+    device = torch.device("cpu")
+    kv, st = make_state(device=device, window_chunks=2, chunk_size=RAGGED_CHUNK)
+    _commit_video_span(kv, st, kv_branch=POS, n_chunks=1, dtype=torch.float32, device=device, chunk_size=RAGGED_CHUNK)
+
+    ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=False)[0].forward_ctx
+    ctx.ensure_video_slots(device)
+
+    history_tail_block = ctx.history_block_ids[-1]
+    # The chunk starts inside the history's own last block, not at a fresh one.
+    assert ctx.current_video_block_ids[0] == history_tail_block
+    first_slot = int(ctx.current_video_slot_mapping[0])
+    assert first_slot == history_tail_block * BLOCK + 8
+
+    # And the action region must not be charged for that managed block.
+    ctx.ensure_action_slots(3, device)
+    assert ctx._scratch_blocks_used == len(ctx.current_video_block_ids) - 1
+    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.
+
+    With a sink the visible window is 2*24 + 24 = 72 tokens, which is 4.5
+    blocks. Deriving the width by flooring that would understate the capacity,
+    the max() against the live block count would take over, and the shape would
+    change as the window filled -- recompiling the graph. The same floor would
+    also let max_seq_len come out under the kv_len actually passed.
+    """
+    device = torch.device("cpu")
+    kv, st = make_state(device=device, window_chunks=2, sink_chunks=1, chunk_size=RAGGED_CHUNK)
+    # The window is deliberately not a whole number of blocks.
+    max_video_tokens = kv.spec.sliding_window + kv.spec.sink_chunks * kv.spec.chunk_size
+    assert max_video_tokens % BLOCK != 0
+
+    widths: set[int] = set()
+    kv_lens: list[int] = []
+    for _ in range(6):
+        ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=True)[0].forward_ctx
+        ctx.ensure_video_slots(device)
+        block_table, _, seq_lens, _, max_seq_len = ctx.build_block_table(
+            action_len=0, query_len=RAGGED_CHUNK, device=device
+        )
+        widths.add(int(block_table.shape[1]))
+        # A bound handed to the kernel has to actually bound the sequence.
+        assert int(seq_lens[0]) <= max_seq_len
+        kv_lens.append(int(seq_lens[0]))
+        st.commit_paged_context(POS)
+
+    assert len(widths) == 1, f"block table width varied across ticks: {sorted(widths)}"
+    # The window has to have actually filled, or none of the above was tested.
+    assert max(kv_lens) > min(kv_lens)
+    assert max(kv_lens) <= ctx.max_video_blocks * BLOCK
+    # Whole blocks are read, so each of the two boundaries -- the sink's end and
+    # the recent window's start -- can contribute up to BLOCK - 1 tokens the
+    # window no longer needs.
+    assert max(kv_lens) <= max_video_tokens + 2 * (BLOCK - 1)
+
+
+def _poison_pools(kv: ARDiffusionKVCache) -> None:
+    for pool in (*kv._k_pools, *kv._v_pools):
+        pool.fill_(float("nan"))
+
+
+def _write_prepared_chunk(
+    kv: ARDiffusionKVCache, ctx, *, dtype: torch.dtype, device: torch.device
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Write fresh K/V for a prepared chunk, poisoning the rest of its blocks first.
+
+    Blocks are handed out whole and a reused block keeps whatever it held, so a
+    slot read without having been written would not look unusual. Every slot of
+    the chunk's blocks that this chunk does not write is set to NaN first --
+    except the leading slots of a block the history still occupies -- so reading
+    one turns the attention output NaN.
+    """
+    mapping = ctx.current_video_slot_mapping
+    first_slot = int(mapping[0])
+    for index, block in enumerate(ctx.current_video_block_ids):
+        start = first_slot if index == 0 else block * BLOCK
+        kv._k_pools[0][start : (block + 1) * BLOCK] = float("nan")
+        kv._v_pools[0][start : (block + 1) * BLOCK] = float("nan")
+    k = torch.randn(ctx.seq_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    v = torch.randn(ctx.seq_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    kv._k_pools[0][mapping] = k
+    kv._v_pools[0][mapping] = v
+    return k, v
+
+
+def _window_positions(end: int, *, sink: int, window: int, written, history: int, resident: set[int]) -> list[int]:
+    """Token positions attention must read at ``end``.
+
+    Every token of the current chunk, plus every written history token in a
+    block that is still resident and holds a sink token or one of the window's
+    most recent tokens.
+    """
+    start = end - window
+    blocks = {p // BLOCK for p in range(min(sink, end))} | {p // BLOCK for p in range(max(start, 0), end)}
+    return [p for p in sorted(written) if p >= history or (p // BLOCK in blocks and p // BLOCK in resident)]
+
+
+def _model_block_indices(kv, adapter) -> dict[int, int]:
+    """Each resident block's index in model positions, by block id.
+
+    Compaction removes the evicted gap after the sink's blocks and shifts the
+    rest of the table down by it, so a table index past the sink is
+    ``compacted_tokens`` behind the position it holds.
+    """
+    block_size = kv.block_size
+    sink_blocks = -(-(kv.spec.sink_chunks * kv.spec.chunk_size) // block_size)
+    shift = adapter.compacted_tokens // block_size
+    return {
+        int(block): index if index < sink_blocks else index + shift
+        for index, block in enumerate(kv.block_table(adapter))
+        if block != kv.null_block_id
+    }
+
+
+@pytest.mark.parametrize(
+    ("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
+):
+    """Once the window slides, the table must still cover the window exactly.
+
+    The sink and the recent window are separate token ranges, and each can
+    straddle a block edge on its own: 24-token chunks against 16-token blocks
+    put a boundary 8 tokens past an edge on alternating chunks. Counting blocks
+    off the ends of the resident list dropped the block holding the window's
+    first tokens whenever both boundaries straddled, and the returned length
+    ran past the last written slot. Neither changes a shape, so this compares
+    values against dense attention over the tokens the window keeps -- after
+    eviction, with every unwritten slot poisoned, on both the scratch and the
+    committing forward of every chunk.
+
+    The reset case keeps only the sink across a boundary, so the recent window
+    reaches back over blocks that are no longer resident. Those must be skipped
+    by position rather than read.
+
+    Eight chunks are enough for compaction to drop the evicted gap from the
+    table, after which storage positions run behind model positions. Removing
+    a gap that is not a whole number of chunks moves where eviction snaps, and
+    in the reset case the window then keeps tokens a boundary should have
+    dropped.
+    """
+    torch.manual_seed(0)
+    device, dtype = torch.device("cpu"), torch.float32
+    kv, st = make_state(
+        device=device,
+        window_chunks=window_chunks,
+        sink_chunks=sink_chunks,
+        reset_at_boundary=reset_at_boundary,
+        chunk_size=RAGGED_CHUNK,
+    )
+    _poison_pools(kv)
+    sink, window = sink_chunks * RAGGED_CHUNK, int(kv.spec.sliding_window)
+    committed: dict[int, tuple[torch.Tensor, torch.Tensor]] = {}
+    evicted = compacted = False
+
+    for chunk in range(8):
+        history = int(st.adapter(POS).absolute_num_computed_tokens)
+        end = history + RAGGED_CHUNK
+        for commit_current in (False, True):
+            ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=commit_current)[0].forward_ctx
+            ctx.ensure_video_slots(device)
+            k, v = _write_prepared_chunk(kv, ctx, dtype=dtype, device=device)
+            written = dict(committed)
+            written.update({history + i: (k[i], v[i]) for i in range(RAGGED_CHUNK)})
+
+            query = torch.randn(1, RAGGED_CHUNK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+            block_table, query_start_loc, seq_lens, max_query_len, max_seq_len = ctx.build_block_table(
+                action_len=0, query_len=RAGGED_CHUNK, device=device
+            )
+            resident = set(_model_block_indices(kv, st.adapter(POS)).values())
+            positions = _window_positions(
+                end, sink=sink, window=window, written=written, history=history, resident=resident
+            )
+            label = f"chunk {chunk}, commit_current={commit_current}"
+            assert int(seq_lens[0]) == len(positions), f"{label}: kv_len does not match the written slots it covers"
+            paged = ar_diffusion_paged_attention(
+                query,
+                kv.key_cache(0),
+                kv.value_cache(0),
+                block_table=block_table,
+                query_start_loc=query_start_loc,
+                seq_lens=seq_lens,
+                max_query_len=max_query_len,
+                max_seq_len=max_seq_len,
+                softmax_scale=HEAD_DIM**-0.5,
+                causal=False,
+            )
+            ref = _dense_attention(
+                query,
+                torch.stack([written[p][0] for p in positions]).unsqueeze(0),
+                torch.stack([written[p][1] for p in positions]).unsqueeze(0),
+            )
+            torch.testing.assert_close(paged, ref, rtol=1e-5, atol=1e-5, msg=lambda m, label=label: f"{label}: {m}")
+        committed = written
+        st.commit_paged_context(POS)
+        compacted = compacted or st.adapter(POS).compacted_tokens > 0
+        evicted = evicted or compacted or kv.null_block_id in kv.block_table(st.adapter(POS))
+
+    assert evicted, "the window never slid, so nothing after eviction was tested"
+    assert compacted, "the table was never compacted, so storage and model positions never diverged"
+
+
+@pytest.mark.parametrize(
+    ("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
+):
+    """Compaction renumbers the block table; it must not change what eviction keeps.
+
+    Storage positions run ``compacted_tokens`` behind model positions, and
+    eviction snaps to chunk boundaries counted in storage positions. Removing a
+    gap that is not a whole number of chunks moves every later snap, so eviction
+    keeps a different set of tokens -- in the reset case, tokens a boundary
+    should have dropped, which the window then reads. Attention checked against
+    what is resident cannot see that, so the oracle here is the same rollout
+    with compaction switched off.
+    """
+
+    def resident_blocks_per_tick(compact: bool) -> tuple[list[set[int]], int]:
+        with monkeypatch.context() as patch:
+            if not compact:
+                patch.setattr(ChunkWindowManager, "compact_block_table", lambda self, request_id: 0)
+            kv, st = make_state(
+                window_chunks=window_chunks,
+                sink_chunks=sink_chunks,
+                reset_at_boundary=reset_at_boundary,
+                chunk_size=RAGGED_CHUNK,
+            )
+            ticks = []
+            for _ in range(12):
+                ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=True)[0].forward_ctx
+                ctx.ensure_video_slots(torch.device("cpu"))
+                ticks.append(set(_model_block_indices(kv, st.adapter(POS)).values()))
+                st.commit_paged_context(POS)
+            return ticks, st.adapter(POS).compacted_tokens
+
+    expected, _ = resident_blocks_per_tick(compact=False)
+    actual, compacted_tokens = resident_blocks_per_tick(compact=True)
+    assert compacted_tokens > 0, "the table was never compacted, so nothing was tested"
+    for tick, (want, got) in enumerate(zip(expected, actual)):
+        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.
+
+    A sink or window that is not a whole number of blocks spans one more block
+    than its token count suggests, so sizing staging from tokens left it short
+    and the first prepared forward refused to run. And because such a chunk
+    shares its first block with the history's tail, the history a probe sees is
+    not left untouched by the next probe, so it is restaged rather than reused.
+    """
+    monkeypatch.setenv(KV_GATHER_ENV, "1")
+    device = torch.device("cpu")
+    kv, st = make_state(
+        device=device, window_chunks=2, sink_chunks=1, chunk_size=RAGGED_CHUNK, reuse_history_staging=True
+    )
+    capacity = int(kv.history_staging[0][0].shape[0])
+    for chunk in range(6):
+        for probe in range(2):
+            ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=False)[0].forward_ctx
+            ctx.prepare(device=device, action_len=0, query_len=RAGGED_CHUNK)
+            assert int(ctx.max_seq_len) <= capacity
+            assert not ctx.reuse_history, f"chunk {chunk}, probe {probe}: a ragged chunk reused the staged history"
+        ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=True)[0].forward_ctx
+        ctx.prepare(device=device, action_len=0, query_len=RAGGED_CHUNK)
+        st.commit_paged_context(POS)
+
+
+@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.
+
+    A 24-token chunk from an empty history ends 8 slots into its second 16-token
+    block, so the run would read those 8 unwritten slots as the first action
+    tokens and never reach the last ones. Every shape stays right, which is why
+    this is refused rather than computed.
+    """
+    _, st = make_state(window_chunks=2, chunk_size=RAGGED_CHUNK)
+    ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=commit_current)[0].forward_ctx
+    with pytest.raises(ValueError, match="partly written video block"):
+        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.
+
+    A 9-frame sink and a 9-frame window are each 8 tokens past a 16-token edge
+    but 28080 tokens together -- a whole number of blocks -- so a width taken
+    from their sum was one block short on every tick once the window had
+    filled. Twelve ticks of three frames, checking every slot the kernel reads
+    against the token that slot is supposed to hold. No attention is computed:
+    at this size the bookkeeping is the whole question.
+    """
+    from vllm_omni.experimental.ar_diffusion.runner import paging_block_size
+
+    frame, frames_per_tick = (480 // 16) * (832 // 16), 3
+    block = paging_block_size(frame)
+    seq_len = frames_per_tick * frame
+    cfg = ARDiffusionKVConfig(enable=True, chunk_size=frame, window_chunks=9, sink_chunks=9)
+    kv = ARDiffusionKVCache(
+        cfg,
+        num_layers=1,
+        num_kv_heads=1,
+        head_size=8,
+        dtype=torch.float16,
+        block_size=block,
+        max_model_len=1 << 20,
+        available_bytes=1 << 27,
+        kv_branches=(ARDiffusionKVBranchSpec(POS, 0),),
+        session_capacity=1,
+        frames_per_block=frames_per_tick,
+        max_scratch_tokens_per_branch=seq_len,
+        device=torch.device("cpu"),
+    )
+    st = ARDiffusionKVState(kv, "s1", {POS: kv.begin_request("r-pos")}, num_layers=1)
+    sink, window = 9 * frame, int(kv.spec.sliding_window)
+    holds: dict[int, int] = {}  # slot -> token position last written there
+    widths: set[int] = set()
+    filled = compacted = False
+
+    for tick in range(12):
+        history = int(st.adapter(POS).absolute_num_computed_tokens)
+        end = history + seq_len
+        ctx = st.get_kv_caches(POS, seq_len=seq_len, commit_current=True)[0].forward_ctx
+        ctx.ensure_video_slots(torch.device("cpu"))
+        for offset, slot in enumerate(ctx.current_video_slot_mapping.tolist()):
+            holds[slot] = history + offset
+        table, _, seq_lens, _, _ = ctx.build_block_table(action_len=0, query_len=seq_len, device=torch.device("cpu"))
+        widths.add(int(table.shape[1]))
+        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: 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))
+        read = read[:kv_len]
+        wrong = [(slot, want) for slot, want in read if holds.get(slot) != want]
+        assert not wrong, f"tick {tick}: {len(wrong)} slots read hold a token other than the one their position implies"
+
+        start = end - window
+        kept = set(range(min(sink, end))) | set(range(max(start, 0), end))
+        missed = kept - {want for _, want in read}
+        assert not missed, f"tick {tick}: {len(missed)} tokens the window keeps were not read"
+
+        filled = filled or end > sink + window
+        st.commit_paged_context(POS)
+        compacted = compacted or st.adapter(POS).compacted_tokens > 0
+
+    assert filled, "the window never filled, so the case that fails was never reached"
+    assert compacted, "the table was never compacted, so storage and model positions never diverged"
+    assert len(widths) == 1, f"block table width varied across ticks: {sorted(widths)}"
+
+
+class _CountedBlockIds:
+    """Wraps a vLLM block manager's ``get_blocks`` and counts the block ids read.
+
+    Every path to a request's block ids goes through ``get_blocks``, so a copy of
+    the whole table is counted in full even if only a few entries are indexed.
+    """
+
+    def __init__(self, manager):
+        self.reads = 0
+        self._get_blocks = manager.get_blocks
+
+    def __call__(self, request_id):
+        from vllm.v1.core.kv_cache_manager import KVCacheBlocks
+
+        blocks = self._get_blocks(request_id)
+        return KVCacheBlocks(tuple([_CountedBlock(block, self) for block in group] for group in blocks.blocks))
+
+
+class _CountedBlock:
+    def __init__(self, block, counter: _CountedBlockIds):
+        self._block = block
+        self._counter = counter
+
+    @property
+    def block_id(self) -> int:
+        self._counter.reads += 1
+        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.
+
+    Evicted positions stay in the table as null entries until compaction drops
+    them a whole number of chunks at a time -- up to two frames' worth of blocks
+    at 832x480 -- so the table can still be far longer than the window. Only the
+    sink's blocks and the recent window's blocks can hold a kept token, so those
+    are all the read may touch. Compaction is switched off here so the table
+    grows to many times the window, and reads are counted where the ids leave
+    vLLM's block manager.
+    """
+    monkeypatch.setattr(ChunkWindowManager, "compact_block_table", lambda self, request_id: 0)
+    device = torch.device("cpu")
+    kv, st = make_state(device=device, window_chunks=2, sink_chunks=1, chunk_size=RAGGED_CHUNK)
+    reads: list[int] = []
+    table_lengths: list[int] = []
+    for _ in range(60):
+        ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=True)[0].forward_ctx
+        ctx.ensure_video_slots(device)
+        counted = _CountedBlockIds(kv.manager)
+        kv.manager.get_blocks = counted
+        try:
+            ctx.build_block_table(action_len=0, query_len=RAGGED_CHUNK, device=device)
+        finally:
+            del kv.manager.get_blocks
+        reads.append(counted.reads)
+        table_lengths.append(len(kv.block_table(st.adapter(POS))))
+        st.commit_paged_context(POS)
+
+    assert table_lengths[-1] > 10 * ctx.max_video_blocks, "the table never outgrew the window, so nothing was tested"
+    assert max(reads) <= ctx.max_video_blocks, (
+        f"read {max(reads)} table entries for a {ctx.max_video_blocks}-block window"
+    )
+
+
+@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.
+
+    A frame-sized block made the resolution a kernel-compatibility question:
+    1560 tokens per frame is not a multiple of 16, so FlashAttention's paged
+    kernel rejected it and the default resolution had no realtime path at all.
+    """
+    from vllm_omni.experimental.ar_diffusion.runner import paging_block_size
+
+    tokens_per_frame = (480 // 16) * (832 // 16)
+    assert tokens_per_frame == 1560
+    assert tokens_per_frame % 16 == 8, "the whole point is that a frame is not a legal block"
+
+    block_size = paging_block_size(tokens_per_frame)
+    cfg = ARDiffusionKVConfig(enable=True, chunk_size=tokens_per_frame, window_chunks=2)
+    kv = ARDiffusionKVCache(
+        cfg,
+        num_layers=1,
+        num_kv_heads=N_HEADS,
+        head_size=HEAD_DIM,
+        dtype=torch.float32,
+        block_size=block_size,
+        max_model_len=1 << 16,
+        available_bytes=1 << 28,
+        kv_branches=(ARDiffusionKVBranchSpec(POS, 0),),
+        session_capacity=1,
+        frames_per_block=2,
+        max_scratch_tokens_per_branch=block_size,
+        device=torch.device("cpu"),
+    )
+
+    assert kv.block_size % 16 == 0
+    # The eviction unit is still the frame; only the paging unit changed.
+    assert kv.spec.chunk_size == tokens_per_frame
+    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.
+
+    This is the case the issue reported failing at ~4e-02 on the scratch path.
+    It runs on CPU because the reference is dense attention, not a kernel.
+    """
+    torch.manual_seed(0)
+    device = torch.device("cpu")
+    dtype = torch.float32
+    tokens_per_frame = (480 // 16) * (832 // 16)
+    assert tokens_per_frame == 1560
+
+    kv, st = make_state(dtype=dtype, device=device, window_chunks=2, chunk_size=tokens_per_frame)
+    history_k, history_v = _commit_video_span(
+        kv, st, kv_branch=POS, n_chunks=1, dtype=dtype, device=device, chunk_size=tokens_per_frame
+    )
+
+    # The non-committing path is the one that was wrong: four of the five
+    # forwards per generated block take it.
+    ctx = st.get_kv_caches(POS, seq_len=tokens_per_frame, commit_current=False)[0].forward_ctx
+    ctx.ensure_video_slots(device)
+    assert ctx.start_offset == 8, "1560 % 16 == 8 is the whole reason this case exists"
+
+    current_k = torch.randn(1, tokens_per_frame, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    current_v = torch.randn(1, tokens_per_frame, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    kv._k_pools[0][ctx.current_video_slot_mapping] = current_k[0]
+    kv._v_pools[0][ctx.current_video_slot_mapping] = current_v[0]
+
+    query = torch.randn(1, tokens_per_frame, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    block_table, query_start_loc, seq_lens, max_query_len, max_seq_len = ctx.build_block_table(
+        action_len=0, query_len=tokens_per_frame, device=device
+    )
+    paged = ar_diffusion_paged_attention(
+        query,
+        kv.key_cache(0),
+        kv.value_cache(0),
+        block_table=block_table,
+        query_start_loc=query_start_loc,
+        seq_lens=seq_lens,
+        max_query_len=max_query_len,
+        max_seq_len=max_seq_len,
+        softmax_scale=HEAD_DIM**-0.5,
+        causal=False,
+    )
+    ref = _dense_attention(
+        query,
+        torch.cat([history_k, current_k], dim=1),
+        torch.cat([history_v, current_v], dim=1),
+    )
+    torch.testing.assert_close(paged, ref, rtol=1e-5, atol=1e-5)
+
+
+@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.
+
+    The CPU test above proves the addressing is right against a dense
+    reference. This one proves the real kernel agrees, at the real 1560
+    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
+    tokens_per_frame = (480 // 16) * (832 // 16)
+    assert tokens_per_frame == 1560
+
+    kv, st = make_state(dtype=dtype, device=device, window_chunks=2, chunk_size=tokens_per_frame)
+    history_k, history_v = _commit_video_span(
+        kv, st, kv_branch=POS, n_chunks=1, dtype=dtype, device=device, chunk_size=tokens_per_frame
+    )
+
+    layer_ctx = st.get_kv_caches(POS, seq_len=tokens_per_frame, commit_current=commit_current)[0]
+    ctx = layer_ctx.forward_ctx
+    current_k = torch.randn(1, tokens_per_frame, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    current_v = torch.randn(1, tokens_per_frame, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    query = torch.randn(1, tokens_per_frame, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+
+    ctx.prepare(device=device, action_len=0, query_len=tokens_per_frame)
+    assert ctx.start_offset == 8, "1560 % 16 == 8 is the whole reason this case exists"
+    assert kv.block_size == 16
+
+    inputs = layer_ctx.to_layer_inputs()
+    paged = paged_write_attn(inputs, query[0], current_k[0], current_v[0], None, None, HEAD_DIM**-0.5).unsqueeze(0)
+
+    new_k = torch.cat([history_k, current_k], dim=1)[:, -kv.spec.sliding_window :]
+    new_v = torch.cat([history_v, current_v], dim=1)[:, -kv.spec.sliding_window :]
+    ref = _dense_attention(query, new_k, new_v)
+
+    torch.testing.assert_close(paged, ref, rtol=2e-2, atol=2e-2)
+
+
+@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):
@@ -556,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_attention_capabilities.py

@@ -0,0 +1,359 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import sys
+from dataclasses import replace
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.attention.backends.abstract import (
+    AttentionBackend,
+    AttentionMetadata,
+)
+from vllm_omni.diffusion.attention.backends.flash_attn import (
+    FlashAttentionBackend,
+    FlashAttentionImpl,
+)
+from vllm_omni.diffusion.attention.backends.utils import fa
+from vllm_omni.diffusion.attention.capabilities import (
+    CapabilityResult,
+    CompilationMode,
+    ExecutionContext,
+    OuterBoundary,
+    ParallelStrategy,
+    SupportStatus,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+_validate_fa4_head_dims = fa.validate_fa4_head_dims
+
+
+@pytest.fixture(autouse=True)
+def kernel_validator(monkeypatch):
+    # Capability unit tests do not require FA4 or a CUDA device. Real-kernel
+    # validation is exercised by test_flash_attn_compile.py.
+    monkeypatch.setattr(fa, "validate_fa4_head_dims", lambda *_args: True)
+
+
+def _impl(*, kernel_variant: str | None = "fa4", causal: bool = False):
+    impl = FlashAttentionImpl.__new__(FlashAttentionImpl)
+    impl._kernel_variant = kernel_variant
+    impl.causal = causal
+    return impl
+
+
+def _resolve(
+    impl,
+    *,
+    context: ExecutionContext | None = None,
+    attn_metadata: AttentionMetadata | None = None,
+    dtype=torch.bfloat16,
+):
+    tensor = torch.empty((1, 16, 8, 64), dtype=dtype)
+    return impl.resolve_execution_path(
+        context or ExecutionContext(platform="cuda"),
+        tensor,
+        tensor,
+        tensor,
+        attn_metadata,
+    )
+
+
+class _UnmigratedBackend(AttentionBackend):
+    @staticmethod
+    def get_name() -> str:
+        return "UNMIGRATED_TEST"
+
+    @staticmethod
+    def get_impl_cls():
+        return None
+
+    @staticmethod
+    def get_metadata_cls():
+        return None
+
+    @staticmethod
+    def get_builder_cls():
+        return None
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return []
+
+
+def test_capability_result_is_tri_state():
+    assert CapabilityResult.supported().status is SupportStatus.SUPPORTED
+    assert CapabilityResult.unsupported("not available").status is SupportStatus.UNSUPPORTED
+    assert CapabilityResult.unmigrated().status is SupportStatus.UNMIGRATED
+    with pytest.raises(ValueError, match="actionable reason"):
+        CapabilityResult(SupportStatus.UNSUPPORTED)
+
+
+def test_unmigrated_backend_default_is_conservative():
+    context = ExecutionContext(platform="cuda", require_fullgraph=True)
+    result = _UnmigratedBackend.resolve_capabilities(context)
+
+    assert result.support.status is SupportStatus.UNMIGRATED
+    assert result.compilation_mode is CompilationMode.EAGER_ONLY
+    requested = result.requested_support(context)
+    assert requested.status is SupportStatus.UNSUPPORTED
+    assert "no verified fullgraph declaration" in requested.reason
+    assert result.support.status is SupportStatus.UNMIGRATED
+    assert result.requested_support(replace(context, require_fullgraph=False)).status is SupportStatus.UNMIGRATED
+
+
+@pytest.mark.parametrize("mode", list(CompilationMode))
+def test_fullgraph_request_checks_compilation_mode(mode):
+    result = replace(_resolve(_impl()), compilation_mode=mode)
+    context = ExecutionContext(platform="cuda", require_fullgraph=True)
+    expected = SupportStatus.UNSUPPORTED if mode is CompilationMode.EAGER_ONLY else SupportStatus.SUPPORTED
+    assert result.requested_support(context).status is expected
+    assert result.requested_support(replace(context, require_fullgraph=False)).status is SupportStatus.SUPPORTED
+
+
+def test_request_preserves_verified_rejection_reason():
+    result = replace(_resolve(_impl()), support=CapabilityResult.unsupported("unsupported head size"))
+    assert result.requested_support(ExecutionContext(platform="cuda", require_fullgraph=True)) == result.support
+
+
+def test_exact_cuda_fa4_dense_path_is_supported():
+    context = ExecutionContext(
+        platform="cuda",
+        require_fullgraph=True,
+    )
+    result = _resolve(_impl(), context=context)
+
+    assert result.path == "fa4_dense"
+    assert result.support.status is SupportStatus.SUPPORTED
+    assert result.compilation_mode is CompilationMode.CUSTOM_OP
+    assert result.requested_support(context).status is SupportStatus.SUPPORTED
+
+
+def test_causal_fa4_path_remains_unmigrated():
+    result = _resolve(_impl(causal=True))
+
+    assert result.support.status is SupportStatus.UNMIGRATED
+
+
+def test_fa4_delegates_dimensions_to_kernel(monkeypatch):
+    head_dims = (80, 48)
+    validator = Mock(return_value=True)
+    monkeypatch.setattr(fa, "validate_fa4_head_dims", validator)
+    query = torch.empty((1, 16, 8, head_dims[0]), dtype=torch.bfloat16)
+    value = torch.empty((1, 16, 8, head_dims[1]), dtype=torch.bfloat16)
+    result = _impl().resolve_execution_path(ExecutionContext(platform="cuda"), query, query, value, None)
+    assert result.support.status is SupportStatus.SUPPORTED
+    validator.assert_called_once_with(*head_dims, 8)
+
+
+def test_fa4_reports_kernel_rejection(monkeypatch):
+    monkeypatch.setattr(fa, "validate_fa4_head_dims", Mock(side_effect=AssertionError("kernel dimension constraint")))
+    result = _resolve(_impl())
+    assert result.support.status is SupportStatus.UNSUPPORTED
+    assert "kernel dimension constraint" in result.support.reason
+    assert "another backend" in result.support.reason
+
+
+def test_fa4_missing_validator_remains_unmigrated(monkeypatch):
+    monkeypatch.setattr(fa, "validate_fa4_head_dims", lambda *_args: False)
+    result = _resolve(_impl())
+    assert result.support.status is SupportStatus.UNMIGRATED
+    assert result.compilation_mode is CompilationMode.EAGER_ONLY
+
+
+def test_fa4_rejects_mismatched_query_key_dimensions():
+    query = torch.empty((1, 16, 8, 64), dtype=torch.bfloat16)
+    key = torch.empty((1, 16, 8, 32), dtype=torch.bfloat16)
+    result = _impl().resolve_execution_path(ExecutionContext(platform="cuda"), query, key, key, None)
+    assert result.support.status is SupportStatus.UNSUPPORTED
+    assert "Q and K head dimensions must match" in result.support.reason
+
+
+@pytest.mark.parametrize("arch", [90, 100, 120])
+def test_fa4_validator_adapter_uses_selected_architecture(monkeypatch, arch):
+    validator = Mock()
+    monkeypatch.setitem(
+        sys.modules,
+        "flash_attn.cute.interface",
+        SimpleNamespace(_get_device_arch=lambda: arch, _validate_head_dims=validator),
+    )
+    # Call the real adapter, not the fixture's stand-in.
+    verified = _validate_fa4_head_dims(80, 48, 8)
+    if arch == 120:
+        assert not verified
+        validator.assert_not_called()
+    else:
+        assert verified
+        validator.assert_called_once_with(80, 48, arch // 10, 8)
+
+
+def test_fa4_validator_adapter_handles_missing_private_api(monkeypatch):
+    monkeypatch.setitem(sys.modules, "flash_attn.cute.interface", SimpleNamespace())
+    assert not _validate_fa4_head_dims(64, 64, 8)
+
+
+def test_initialized_kernel_identity_overrides_caller_claim():
+    result = _resolve(
+        _impl(kernel_variant=None),
+        context=ExecutionContext(
+            platform="cuda",
+            kernel_variant="fa4",
+        ),
+    )
+
+    assert result.support.status is SupportStatus.UNMIGRATED
+
+
+def test_preconstruction_without_resolved_cuda_kernel_is_unmigrated():
+    result = FlashAttentionBackend.resolve_capabilities(ExecutionContext(platform="cuda"))
+
+    assert result.support.status is SupportStatus.UNMIGRATED
+    assert result.path == "unverified"
+
+
+@pytest.mark.parametrize(
+    "changes",
+    [
+        {"platform": "rocm"},
+        {"paged_kv": True},
+        {"parallel_strategy": ParallelStrategy.ULYSSES},
+        {"parallel_strategy": ParallelStrategy.RING},
+        {"parallel_strategy": ParallelStrategy.HYBRID_ULYSSES_RING},
+        {"outer_boundaries": frozenset({OuterBoundary.HSDP})},
+    ],
+)
+def test_unverified_flash_attention_variants_remain_unmigrated(changes):
+    result = _resolve(
+        _impl(),
+        context=replace(ExecutionContext(platform="cuda"), **changes),
+    )
+
+    assert result.support.status is SupportStatus.UNMIGRATED
+    assert result.compilation_mode is CompilationMode.EAGER_ONLY
+
+
+def test_unverified_dtype_remains_unmigrated():
+    result = _resolve(_impl(), dtype=torch.float32)
+
+    assert result.support.status is SupportStatus.UNMIGRATED
+
+
+@pytest.mark.parametrize(
+    "attn_metadata",
+    [
+        AttentionMetadata(full_attn_spans=[[(0, 8)]]),
+        AttentionMetadata(
+            extra={
+                "cu_seqlens_q": torch.tensor([0, 16], dtype=torch.int32),
+                "cu_seqlens_k": torch.tensor([0, 16], dtype=torch.int32),
+                "max_seqlen_q": 16,
+                "max_seqlen_k": 16,
+            }
+        ),
+        AttentionMetadata(extra={"kv_cache_dtype": "fp8"}),
+    ],
+)
+def test_metadata_normalization_prevents_dense_support(attn_metadata):
+    result = _resolve(_impl(), attn_metadata=attn_metadata)
+
+    assert result.support.status is SupportStatus.UNMIGRATED
+
+
+def test_unpublished_all_true_mask_remains_runtime_dependent():
+    metadata = AttentionMetadata(
+        attn_mask=torch.ones((1, 16), dtype=torch.bool),
+    )
+    result = _resolve(_impl(), attn_metadata=metadata)
+
+    assert result.path == "runtime_mask_dependent"
+    assert result.support.status is SupportStatus.UNMIGRATED
+
+
+def test_producer_published_noop_mask_resolves_dense_fa4():
+    metadata = AttentionMetadata(
+        attn_mask=torch.ones((1, 16), dtype=torch.bool),
+        extra={"attention_mask_mode": "none"},
+    )
+    result = _resolve(_impl(), attn_metadata=metadata)
+
+    assert result.path == "fa4_dense"
+    assert result.support.status is SupportStatus.SUPPORTED
+
+
+def test_fixed_shape_mask_semantic_change_invalidates_dense_resolution():
+    metadata = AttentionMetadata(
+        attn_mask=torch.ones((1, 16), dtype=torch.bool),
+        extra={"attention_mask_mode": "none"},
+    )
+    impl = _impl()
+    dense = _resolve(impl, attn_metadata=metadata)
+    assert dense.support.status is SupportStatus.SUPPORTED
+
+    metadata.attn_mask[:, 8:] = False
+    metadata.extra["attention_mask_mode"] = "padding"
+    masked = _resolve(impl, attn_metadata=metadata)
+    assert masked.support.status is SupportStatus.UNMIGRATED
+    request = ExecutionContext(platform="cuda", require_fullgraph=True)
+    assert masked.requested_support(request).status is SupportStatus.UNSUPPORTED
+
+
+def test_piecewise_dispatch_ignores_unused_incomplete_packed_metadata(monkeypatch):
+    from vllm_omni.diffusion.attention.backends import flash_attn
+    from vllm_omni.diffusion.attention.backends.utils import fa
+
+    metadata = AttentionMetadata(full_attn_spans=[[(0, 16)]], extra={"max_seqlen_q": 16})
+    impl = _impl()
+    impl.fa_deterministic = False
+    impl.softmax_scale = 0.125
+    assert _resolve(impl, attn_metadata=metadata).support.status is SupportStatus.UNMIGRATED
+
+    query = torch.empty((1, 16, 8, 64), dtype=torch.bfloat16)
+    monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True)
+    monkeypatch.setattr(fa, "flash_attn_func", lambda *args, **kwargs: query)
+    monkeypatch.setattr(flash_attn, "piecewise_attn", lambda *args, **kwargs: query)
+    assert impl.forward_cuda(query, query, query, metadata) is query
+
+
+@pytest.mark.parametrize("input_index", [1, 2])
+@pytest.mark.parametrize("mismatch", ["dtype", "device"])
+def test_fa4_rejects_incompatible_key_value(input_index, mismatch):
+    tensors = [torch.empty((1, 16, 8, 64), dtype=torch.bfloat16) for _ in range(3)]
+    tensors[input_index] = (
+        tensors[input_index].to(dtype=torch.float16) if mismatch == "dtype" else tensors[input_index].to("meta")
+    )
+    result = _impl().resolve_execution_path(ExecutionContext(platform="cuda"), *tensors, None)
+    assert result.support.status is SupportStatus.UNSUPPORTED
+    assert f"{mismatch}s must match" in result.support.reason
+
+
+@pytest.mark.parametrize("skip_quant", [False, True])
+def test_layer_resolution_uses_effective_kv_quantization(skip_quant):
+    from vllm_omni.diffusion.attention.layer import Attention
+    from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention
+
+    layer = Attention.__new__(Attention)
+    torch.nn.Module.__init__(layer)
+    layer.attention = _impl(kernel_variant=None)
+    layer._hsdp_compile_boundary_enabled = False
+    layer.skip_sequence_parallel = False
+    layer._no_parallel_strategy = NoParallelAttention()
+    layer.parallel_strategy = layer._no_parallel_strategy
+    layer.use_ring = False
+    layer.paged_kv_cache_role = None
+    layer._kv_cache_dtype = "fp8"
+    layer._disable_kv_quant = False
+    layer._kv_cache_skip_steps = None
+    layer._kv_cache_skip_layers = {0} if skip_quant else None
+    layer.layer_idx = 0
+    query = torch.empty((1, 16, 8, 64), dtype=torch.bfloat16)
+    metadata = AttentionMetadata(extra={"kv_cache_dtype": "fp8"})
+    result = layer.resolve_execution_path(ExecutionContext(platform="npu"), query, query, query, metadata)
+    assert result.path == ("npu_dense" if skip_quant else "npu_unverified")
+    assert result.support.status is SupportStatus.UNMIGRATED
+    assert metadata.extra == {"kv_cache_dtype": "fp8"}

tests/diffusion/attention/test_attention_state_lifetime.py

@@ -0,0 +1,118 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Stateful contract example, independent of an optional attention kernel.
+
+The prepared module owns tensor state and passes it explicitly to an opaque op.
+This deliberately does not model opaque Python/C++ planning handles.
+"""
+
+import gc
+import weakref
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+@torch.library.custom_op("vllm_omni_test::planned_attention", mutates_args=())
+def _planned_attention(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    query_scale: torch.Tensor,
+) -> torch.Tensor:
+    return F.scaled_dot_product_attention(query * query_scale, key, value).contiguous()
+
+
+@_planned_attention.register_fake
+def _planned_attention_fake(query, key, value, query_scale):
+    return query.new_empty((*query.shape[:-1], value.shape[-1]))
+
+
+class _PreparedAttention(torch.nn.Module):
+    """Initialization/planning runs eagerly; forward only executes tensors."""
+
+    def __init__(self, query_scale: torch.Tensor):
+        super().__init__()
+        self.register_buffer("query_scale", query_scale.detach().clone())
+
+    def forward(self, query, key, value):
+        return _planned_attention(query, key, value, self.query_scale)
+
+
+def _inputs(length):
+    generator = torch.Generator().manual_seed(length)
+    return tuple(torch.randn(2, 3, length, 8, generator=generator) for _ in range(3))
+
+
+def _reference(plan, inputs):
+    query, key, value = inputs
+    return F.scaled_dot_product_attention(query * plan.query_scale, key, value)
+
+
+def test_dynamic_replay_uses_each_instances_state_without_recompilation():
+    compile_count = 0
+
+    def counting_backend(graph_module, _example_inputs):
+        nonlocal compile_count
+        compile_count += 1
+        return graph_module.forward
+
+    # Same tensor metadata, different values: state must be a graph input,
+    # never a Python instance identifier or a captured first-instance constant.
+    plans = [_PreparedAttention(torch.full((1, 3, 1, 1), scale)) for scale in (0.5, 2.0)]
+    compiled = torch.compile(
+        lambda plan, q, k, v: plan(q, k, v),
+        backend=counting_backend,
+        fullgraph=True,
+        dynamic=True,
+    )
+    for length in (5, 9, 13):
+        inputs = _inputs(length)
+        for plan in plans:
+            torch.testing.assert_close(compiled(plan, *inputs), _reference(plan, inputs))
+        assert not torch.allclose(_reference(plans[0], inputs), _reference(plans[1], inputs))
+    assert compile_count == 1
+
+
+def test_compiled_callable_retains_state_until_released():
+    plan = _PreparedAttention(torch.full((1, 3, 1, 1), 0.5))
+    plan_ref = weakref.ref(plan)
+    state_ref = weakref.ref(plan.query_scale)
+    inputs = _inputs(5)
+    expected = _reference(plan, inputs)
+    compiled = torch.compile(plan, backend="eager", fullgraph=True, dynamic=True)
+    torch.testing.assert_close(compiled(*inputs), expected)
+
+    # Deleting the caller's reference must not invalidate subsequent replay.
+    del plan
+    gc.collect()
+    assert plan_ref() is not None
+    assert state_ref() is not None
+    torch.testing.assert_close(compiled(*inputs), expected)
+
+    del compiled
+    gc.collect()
+    assert plan_ref() is None
+    assert state_ref() is None
+
+
+def test_stateful_custom_op_schema_and_fake_implementation():
+    inputs = _inputs(5)
+    # Unequal Q/K lengths and V head dimensions exercise the fake output shape.
+    query, _, _ = inputs
+    key = torch.randn(2, 3, 7, 8)
+    value = torch.randn(2, 3, 7, 4)
+    scale = torch.full((1, 3, 1, 1), 0.5)
+    torch.library.opcheck(_planned_attention, (query, key, value, scale))
+
+
+def test_stateful_attention_compiles_with_inductor():
+    plan = _PreparedAttention(torch.full((1, 3, 1, 1), 0.5))
+    compiled = torch.compile(plan, fullgraph=True, dynamic=True)
+    for length in (5, 9):
+        inputs = _inputs(length)
+        torch.testing.assert_close(compiled(*inputs), _reference(plan, inputs))

tests/diffusion/attention/test_flash_attn.py

@@ -21,6 +21,7 @@ from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionImpl
 from vllm_omni.diffusion.attention.backends.sdpa import SDPAImpl
 from vllm_omni.diffusion.attention.backends.utils import fa  # noqa: E402
+from vllm_omni.diffusion.attention.capabilities import CompilationMode, ExecutionContext, SupportStatus
 from vllm_omni.platforms import current_omni_platform
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
@@ -432,6 +433,35 @@ def test_varlen_masked_routing_by_role(monkeypatch):
     assert cu_seqlens_k.tolist() == [0, 2, 5]
 
 
+@pytest.mark.cpu
+def test_published_noop_mask_uses_dense_path(monkeypatch):
+    calls = []
+
+    def fake_dense_func(q, k, v, **_kwargs):
+        calls.append("dense")
+        return q
+
+    def fail_varlen(*_args, **_kwargs):
+        raise AssertionError("published no-op mask must not use varlen attention")
+
+    monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True)
+    monkeypatch.setattr(fa, "IS_FLASH_ATTN_4", False)
+    monkeypatch.setattr(fa, "flash_attn_func", fake_dense_func)
+    monkeypatch.setattr(fa, "flash_attn_varlen_func", fail_varlen)
+
+    impl = FlashAttentionImpl(num_heads=2, head_size=8, softmax_scale=0.5, causal=False)
+    query = torch.randn(1, 4, 2, 8)
+    metadata = AttentionMetadata(
+        attn_mask=torch.ones((1, 4), dtype=torch.bool),
+        extra={"attention_mask_mode": "none"},
+    )
+
+    output = impl.forward_cuda(query, query, query, metadata)
+
+    assert torch.equal(output, query)
+    assert calls == ["dense"]
+
+
 @pytest.mark.cpu
 def test_piecewise_flash_attn_uses_varlen_fallback(monkeypatch):
     calls = []
@@ -550,6 +580,199 @@ def _fake_torch_npu(monkeypatch, *, npu_fusion_attention=None):
     return fake
 
 
+@pytest.fixture
+def npu_contract(monkeypatch, mocker):
+    monkeypatch.delenv("MINDIE_SD_FA_TYPE", raising=False)
+    monkeypatch.setattr(fa, "validate_fa4_head_dims", lambda *_args: pytest.fail("NPU must not validate FA4"))
+    query = torch.zeros(1, 8, 2, 4)
+    dense = mocker.Mock(return_value=torch.zeros_like(query))
+    packed = mocker.Mock(return_value=torch.zeros_like(query.squeeze(0)))
+    _fake_mindiesd(monkeypatch, attention_forward=dense, attention_forward_varlen=packed)
+    return _npu_impl(), query, dense, packed
+
+
+def _assert_platform_path(impl, platform, query, metadata, path):
+    context = ExecutionContext(platform=platform, require_fullgraph=True)
+    result = impl.resolve_execution_path(context, query, query, query, metadata)
+    assert result.path == path
+    assert result.kernel_variant == {"npu": "mindiesd", "rocm": "aiter"}[platform]
+    assert result.support.status is SupportStatus.UNMIGRATED
+    assert result.compilation_mode is CompilationMode.EAGER_ONLY
+    assert result.requested_support(context).status is SupportStatus.UNSUPPORTED
+
+
+def _real_pad_metadata():
+    cu = torch.tensor([0, 5, 8], dtype=torch.int32)
+    return AttentionMetadata(
+        extra={
+            "npu_attn_varlen": True,
+            "cu_seqlens_q": cu,
+            "cu_seqlens_k": cu,
+            "max_seqlen_q": 5,
+            "max_seqlen_k": 5,
+        }
+    )
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("kv_cache_dtype", [None, "float", "auto"])
+def test_npu_contract_dense_dispatch(npu_contract, kv_cache_dtype):
+    impl, query, dense, packed = npu_contract
+    metadata = AttentionMetadata(extra={"kv_cache_dtype": kv_cache_dtype})
+    _assert_platform_path(impl, "npu", query, metadata, "npu_dense")
+    impl.forward_npu(query, query, query, metadata)
+    dense.assert_called_once()
+    packed.assert_not_called()
+    assert dense.call_args.kwargs["attn_mask"] is None
+
+
+@pytest.mark.cpu
+def test_npu_contract_mask_layout(npu_contract):
+    impl, query, dense, packed = npu_contract
+    mask = torch.arange(8)[None] < 5
+    metadata = AttentionMetadata(attn_mask=mask)
+    _assert_platform_path(impl, "npu", query, metadata, "npu_masked")
+    impl.forward_npu(query, query, query, metadata)
+    dense.assert_called_once()
+    packed.assert_not_called()
+    torch.testing.assert_close(dense.call_args.kwargs["attn_mask"], mask[:, None, None, :].expand(1, 1, 8, 8))
+
+
+@pytest.mark.cpu
+def test_npu_contract_packed_dispatch(npu_contract):
+    impl, query, dense, packed = npu_contract
+    metadata = _real_pad_metadata()
+    _assert_platform_path(impl, "npu", query, metadata, "npu_packed_varlen")
+    impl.forward_npu(query, query, query, metadata)
+    dense.assert_not_called()
+    packed.assert_called_once()
+    assert packed.call_args.args[3:] == ([0, 5, 8], [0, 5, 8])
+
+
+@pytest.mark.cpu
+def test_npu_contract_prefix_dispatch(npu_contract, monkeypatch):
+    impl, query, dense, packed = npu_contract
+    metadata = _real_pad_metadata()
+    _assert_platform_path(impl, "npu", query, metadata, "npu_packed_varlen")
+    # The same metadata resolves differently after the kernel choice changes.
+    monkeypatch.setenv("MINDIE_SD_FA_TYPE", "ascend_laser_attention")
+    _assert_platform_path(impl, "npu", query, metadata, "npu_prefix_kv_slice")
+    impl.forward_npu(query, query, query, metadata)
+    dense.assert_called_once()
+    packed.assert_not_called()
+    torch.testing.assert_close(dense.call_args.args[1], query[:, :5])
+    torch.testing.assert_close(dense.call_args.args[2], query[:, :5])
+    assert dense.call_args.kwargs["attn_mask"] is None
+
+
+@pytest.mark.cpu
+def test_npu_contract_incomplete_packing_rebuilds_mask(npu_contract):
+    impl, query, dense, packed = npu_contract
+    # CUDA-like normalization rejects this metadata, but NPU can rebuild a mask.
+    metadata = AttentionMetadata(extra={"npu_attn_varlen": True, "max_seqlen_q": 5, "valid_kv_length": 5})
+    _assert_platform_path(impl, "npu", query, metadata, "npu_masked_fallback")
+    impl.forward_npu(query, query, query, metadata)
+    dense.assert_called_once()
+    packed.assert_not_called()
+    mask = (torch.arange(8) < 5)[None, None, None, :].expand(1, 1, 8, 8)
+    torch.testing.assert_close(dense.call_args.kwargs["attn_mask"], mask)
+
+
+@pytest.mark.cpu
+def test_npu_contract_unusable_packing_does_not_run_unmasked(npu_contract):
+    impl, query, dense, packed = npu_contract
+    metadata = AttentionMetadata(extra={"npu_attn_varlen": True})
+    _assert_platform_path(impl, "npu", query, metadata, "npu_masked_fallback")
+    with pytest.raises(ValueError, match="refusing to run unmasked attention"):
+        impl.forward_npu(query, query, query, metadata)
+    dense.assert_not_called()
+    packed.assert_not_called()
+
+
+@pytest.fixture
+def rocm_contract(monkeypatch, mocker):
+    query = torch.zeros(1, 8, 2, 4)
+    dense = mocker.Mock(return_value=torch.zeros_like(query))
+    packed = mocker.Mock(side_effect=lambda *args, **kwargs: torch.zeros_like(args[0] if args else kwargs["q"]))
+    monkeypatch.setattr(fa, "IS_FLASH_ATTN_4", False)
+    monkeypatch.setattr(fa, "IS_AITER", True)
+    monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True)
+    monkeypatch.setattr(fa, "flash_attn_func", dense)
+    monkeypatch.setattr(fa, "flash_attn_varlen_func", packed)
+    monkeypatch.setattr(fa, "validate_fa4_head_dims", lambda *_args: pytest.fail("ROCm must not validate FA4"))
+    return _npu_impl(), query, dense, packed
+
+
+@pytest.mark.cpu
+def test_rocm_contract_dense_dispatch(rocm_contract):
+    impl, query, dense, packed = rocm_contract
+    _assert_platform_path(impl, "rocm", query, None, "rocm_dense")
+    impl.forward_hip(query, query, query)
+    dense.assert_called_once()
+    packed.assert_not_called()
+
+
+@pytest.mark.cpu
+def test_rocm_contract_masked_unpadding(rocm_contract):
+    impl, query, dense, packed = rocm_contract
+    metadata = AttentionMetadata(attn_mask=torch.arange(8)[None] < 5, extra={"attention_mask_mode": "padding"})
+    _assert_platform_path(impl, "rocm", query, metadata, "rocm_masked_varlen")
+    impl.forward_hip(query, query, query, metadata)
+    dense.assert_not_called()
+    packed.assert_called_once()
+    assert packed.call_args.args[0].shape == (5, 2, 4)
+    torch.testing.assert_close(packed.call_args.kwargs["cu_seqlens_k"], torch.tensor([0, 5], dtype=torch.int32))
+
+
+@pytest.mark.cpu
+def test_rocm_contract_multiple_packed_documents(rocm_contract):
+    impl, query, dense, packed = rocm_contract
+    cu = torch.tensor([0, 2, 4, 8], dtype=torch.int32)
+    metadata = AttentionMetadata(
+        extra={
+            "cu_seqlens_q": cu,
+            "cu_seqlens_k": cu,
+            "max_seqlen_q": 4,
+            "max_seqlen_k": 4,
+        }
+    )
+    _assert_platform_path(impl, "rocm", query, metadata, "rocm_packed_varlen")
+    # NPU's [real, pad] path cannot consume this packing contract.
+    assert impl._resolve_packed_seq_npu(query, query, metadata.extra) is None
+    impl.forward_hip(query, query, query, metadata)
+    dense.assert_not_called()
+    packed.assert_called_once()
+    assert packed.call_args.kwargs["cu_seqlens_q"] is cu
+    assert packed.call_args.kwargs["cu_seqlens_k"] is cu
+
+
+@pytest.mark.cpu
+def test_rocm_contract_unknown_mask_stays_runtime_dependent(rocm_contract, monkeypatch):
+    impl, query, dense, packed = rocm_contract
+    metadata = AttentionMetadata(attn_mask=torch.ones(1, 8, dtype=torch.bool))
+    # Resolution may inspect shape/metadata, but must not read mask values.
+    with monkeypatch.context() as patch:
+        patch.setattr(torch, "any", lambda *_args, **_kwargs: pytest.fail("resolution read mask values"))
+        _assert_platform_path(impl, "rocm", query, metadata, "rocm_runtime_mask_dependent")
+    impl.forward_hip(query, query, query, metadata)
+    dense.assert_called_once()
+    packed.assert_not_called()
+
+
+@pytest.mark.cpu
+def test_rocm_contract_requires_initialized_aiter(monkeypatch):
+    monkeypatch.setattr(fa, "IS_FLASH_ATTN_4", False)
+    monkeypatch.setattr(fa, "IS_AITER", False)
+    impl = _npu_impl()
+    query = torch.zeros(1, 8, 2, 4)
+    result = impl.resolve_execution_path(
+        ExecutionContext(platform="rocm", kernel_variant="aiter"), query, query, query, None
+    )
+    assert result.path == "rocm_unverified"
+    assert result.kernel_variant is None
+    assert result.support.status is SupportStatus.UNMIGRATED
+
+
 # --- Test group A: boundary resolution (_resolve_packed_seq_npu) -------------
 
 

tests/diffusion/attention/test_flash_attn_compile.py

@@ -0,0 +1,308 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import functools
+from dataclasses import dataclass
+
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.diffusion.attention.backends.flash_attn import (
+    FlashAttentionBackend,
+    FlashAttentionImpl,
+)
+from vllm_omni.diffusion.attention.backends.utils import fa
+from vllm_omni.diffusion.attention.capabilities import (
+    CompilationMode,
+    ExecutionContext,
+    OuterBoundary,
+    ParallelStrategy,
+    SupportStatus,
+)
+from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention
+from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
+from vllm_omni.diffusion.forward_context import ForwardContext
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cuda]
+
+
+@pytest.fixture(autouse=True)
+def isolated_compiler_cache():
+    # Each case changes kernel/configuration on the same Attention.forward code
+    # object. Keep unrelated cases from exhausting its Dynamo recompile limit;
+    # retain the cache across all shapes and replays within an individual test.
+    torch.compiler.reset()
+    yield
+    torch.compiler.reset()
+
+
+@dataclass
+class _ParallelStrategyStub:
+    name: str
+    enabled: bool = True
+
+
+def _make_attention_layer(monkeypatch, *, head_size=64):
+    from vllm_omni.diffusion.attention import layer as attention_layer
+
+    monkeypatch.setattr(
+        attention_layer,
+        "get_attn_backend_for_role",
+        lambda **_kwargs: (FlashAttentionBackend, None),
+    )
+    monkeypatch.setattr(
+        attention_layer,
+        "build_parallel_attention_strategy",
+        lambda **_kwargs: NoParallelAttention(),
+    )
+    return attention_layer.Attention(num_heads=8, head_size=head_size, softmax_scale=head_size**-0.5, causal=False)
+
+
+@pytest.fixture
+def fake_fa4(monkeypatch, tmp_path):
+    marker = tmp_path / "kernel"
+    marker.write_text("loaded", encoding="utf-8")
+
+    @functools.cache
+    def cached_kernel_loader():
+        with open(marker, encoding="utf-8") as handle:
+            handle.read()
+
+    def fake_attention(query, key, value, **kwargs):
+        cached_kernel_loader()
+        return (
+            torch.nn.functional.scaled_dot_product_attention(
+                query.transpose(1, 2),
+                key.transpose(1, 2),
+                value.transpose(1, 2),
+                scale=kwargs["softmax_scale"],
+                is_causal=kwargs["causal"],
+            )
+            .transpose(1, 2)
+            .contiguous()
+        )
+
+    monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True)
+    monkeypatch.setattr(fa, "IS_FLASH_ATTN_4", True)
+    monkeypatch.setattr(fa, "flash_attn_func", fake_attention)
+    monkeypatch.setattr(fa, "validate_fa4_head_dims", lambda *_args: True)
+
+    return FlashAttentionImpl(
+        num_heads=8,
+        head_size=64,
+        softmax_scale=0.125,
+        causal=False,
+    )
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("noncontiguous", [False, True])
+def test_fa4_dense_dispatch_is_opaque_to_dynamic_torch_compile(fake_fa4, noncontiguous):
+    """Fast Dynamo/custom-op regression; Inductor coverage is separate."""
+    impl = fake_fa4
+    q = torch.randn(1, 16, 8, 64, device="cuda", dtype=torch.bfloat16)
+    if noncontiguous:
+        q = q.transpose(1, 2).contiguous().transpose(1, 2)
+        assert not q.is_contiguous()
+    path = impl.resolve_execution_path(
+        ExecutionContext(
+            platform="cuda",
+            kernel_variant="fa4",
+            dtype="bfloat16",
+            require_fullgraph=True,
+        ),
+        q,
+        q,
+        q,
+        None,
+    )
+    assert path.support.status is SupportStatus.SUPPORTED
+    assert path.compilation_mode is CompilationMode.CUSTOM_OP
+
+    compile_count = 0
+
+    def counting_backend(graph_module, _example_inputs):
+        nonlocal compile_count
+        compile_count += 1
+        return graph_module.forward
+
+    compiled = torch.compile(
+        lambda query, key, value: impl.forward_cuda(query, key, value),
+        backend=counting_backend,
+        fullgraph=True,
+        dynamic=True,
+    )
+    out = compiled(q, q, q)
+    q2 = torch.randn(1, 24, 8, 64, device="cuda", dtype=torch.bfloat16)
+    if noncontiguous:
+        q2 = q2.transpose(1, 2).contiguous().transpose(1, 2)
+    out2 = compiled(q2, q2, q2)
+    assert out.is_contiguous()
+    assert out2.is_contiguous()
+
+    assert out.shape == q.shape
+    assert out2.shape == q2.shape
+    torch.testing.assert_close(out, impl.forward_cuda(q, q, q))
+    torch.testing.assert_close(out2, impl.forward_cuda(q2, q2, q2))
+    assert compile_count == 1
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_fa4_production_attention_entry_compiles_with_inductor(
+    monkeypatch,
+    fake_fa4,
+):
+    """Compile through Attention.forward with the default Inductor backend."""
+    from vllm_omni.diffusion.attention import layer as attention_layer
+
+    layer = _make_attention_layer(monkeypatch)
+    layer.attention = fake_fa4
+    q = torch.randn(1, 16, 8, 64, device="cuda", dtype=torch.bfloat16)
+
+    resolved_contexts = []
+    resolve_execution_path = fake_fa4.resolve_execution_path
+
+    def record_execution_path(context, query, key, value, attn_metadata):
+        resolved_contexts.append(context)
+        return resolve_execution_path(
+            context,
+            query,
+            key,
+            value,
+            attn_metadata,
+        )
+
+    monkeypatch.setattr(fake_fa4, "resolve_execution_path", record_execution_path)
+    context = ExecutionContext(
+        platform="cuda",
+        kernel_variant="fa4",
+        dtype="bfloat16",
+        paged_kv=True,
+        parallel_strategy=ParallelStrategy.ULYSSES,
+        require_fullgraph=True,
+    )
+    path = layer.resolve_execution_path(context, q, q, q, None)
+    assert path.support.status is SupportStatus.SUPPORTED
+    assert path.compilation_mode is CompilationMode.CUSTOM_OP
+    assert resolved_contexts[-1].paged_kv is False
+    assert resolved_contexts[-1].parallel_strategy is ParallelStrategy.NONE
+
+    compiled = torch.compile(layer, fullgraph=True, dynamic=True)
+    out = compiled(q, q, q)
+
+    assert out.shape == q.shape
+    torch.testing.assert_close(out, layer(q, q, q))
+    q2 = torch.randn(1, 24, 8, 64, device="cuda", dtype=torch.bfloat16)
+    torch.testing.assert_close(compiled(q2, q2, q2), layer(q2, q2, q2))
+
+    layer.paged_kv_cache_role = "self"
+    monkeypatch.setattr(layer, "_active_paged_kv_adapter", lambda: object())
+    paged_path = layer.resolve_execution_path(context, q, q, q, None)
+    assert resolved_contexts[-1].paged_kv is True
+    assert paged_path.support.status is SupportStatus.UNMIGRATED
+    layer.paged_kv_cache_role = None
+
+    layer.use_ring = True
+    layer.parallel_strategy = _ParallelStrategyStub(name="ulysses")
+    hybrid_path = layer.resolve_execution_path(context, q, q, q, None)
+    assert resolved_contexts[-1].parallel_strategy is ParallelStrategy.HYBRID_ULYSSES_RING
+    assert hybrid_path.support.status is SupportStatus.UNMIGRATED
+    layer.use_ring = False
+    layer.parallel_strategy = NoParallelAttention()
+
+    layer._hsdp_compile_boundary_enabled = True
+    hsdp_path = layer.resolve_execution_path(context, q, q, q, None)
+    assert OuterBoundary.HSDP in resolved_contexts[-1].outer_boundaries
+    assert hsdp_path.support.status is SupportStatus.UNMIGRATED
+    assert hsdp_path.compilation_mode is CompilationMode.EAGER_ONLY
+
+    layer._hsdp_compile_boundary_enabled = False
+    runtime_context = ForwardContext(
+        omni_diffusion_config=OmniDiffusionConfig(
+            parallel_config=DiffusionParallelConfig(use_hsdp=True, hsdp_shard_size=1),
+        ),
+    )
+    monkeypatch.setattr(attention_layer, "is_forward_context_available", lambda: True)
+    monkeypatch.setattr(attention_layer, "get_forward_context", lambda: runtime_context)
+    layer.use_ring = True
+    layer.parallel_strategy = _ParallelStrategyStub(name="ring")
+    runtime_hsdp_path = layer.resolve_execution_path(context, q, q, q, None)
+    assert OuterBoundary.HSDP in resolved_contexts[-1].outer_boundaries
+    assert resolved_contexts[-1].parallel_strategy is ParallelStrategy.NONE
+    assert runtime_hsdp_path.support.status is SupportStatus.UNMIGRATED
+
+
+@hardware_test(res={"cuda": "B200"}, num_cards=1)
+@pytest.mark.parametrize("head_dims", [(32, 32), (64, 64), (80, 48), (192, 128), (256, 256)])
+def test_real_fa4_fullgraph_matches_sdpa_across_shapes(monkeypatch, head_dims):
+    """Validate real FA4 numerics through the complete compiled attention layer."""
+    capability = current_omni_platform.get_device_capability() if current_omni_platform.is_cuda() else None
+    if capability is None or capability.major < 10:
+        pytest.skip("Requires Blackwell")
+    if not fa.IS_FLASH_ATTN_4 or fa.flash_attn_func is None:
+        pytest.skip("Requires CuTe FlashAttention-4")
+
+    from torch.nn.attention import SDPBackend, sdpa_kernel
+
+    head_size, value_head_size = head_dims
+    layer = _make_attention_layer(monkeypatch, head_size=head_size)
+    compiled = torch.compile(layer, fullgraph=True, dynamic=True)
+    generator = torch.Generator(device="cuda").manual_seed(42)
+    context = ExecutionContext(platform="cuda", require_fullgraph=True)
+    # Tile boundaries, cross-attention lengths, batching, and a longer sequence.
+    for batch, q_length, kv_length in ((1, 16, 16), (1, 129, 257), (2, 257, 129), (1, 1024, 1024)):
+        query, key, value = (
+            torch.randn(batch, length, 8, dim, device="cuda", dtype=torch.bfloat16, generator=generator)
+            for length, dim in ((q_length, head_size), (kv_length, head_size), (kv_length, value_head_size))
+        )
+        path = layer.resolve_execution_path(context, query, key, value, None)
+        assert path.requested_support(context).status is SupportStatus.SUPPORTED
+        with sdpa_kernel(SDPBackend.MATH):
+            reference = torch.nn.functional.scaled_dot_product_attention(
+                query.transpose(1, 2).float(),
+                key.transpose(1, 2).float(),
+                value.transpose(1, 2).float(),
+                scale=head_size**-0.5,
+            ).transpose(1, 2)
+        actual = compiled(query, key, value)
+        assert actual.dtype == query.dtype
+        assert actual.device == query.device
+        torch.testing.assert_close(actual.float(), reference, atol=1e-2, rtol=1e-2)
+        torch.testing.assert_close(actual, layer(query, key, value))
+
+
+@hardware_test(res={"cuda": "B200"}, num_cards=1)
+@pytest.mark.parametrize("head_dims", [(192, 192), (65, 64), (64, 65)])
+def test_real_fa4_dimension_rejection_matches_kernel(monkeypatch, head_dims):
+    capability = current_omni_platform.get_device_capability() if current_omni_platform.is_cuda() else None
+    if capability is None or capability.major not in (10, 11) or not fa.IS_FLASH_ATTN_4:
+        pytest.skip("Requires SM100/SM110 CuTe FlashAttention-4")
+    head_size, value_head_size = head_dims
+    layer = _make_attention_layer(monkeypatch, head_size=head_size)
+    query = torch.randn(1, 16, 8, head_size, device="cuda", dtype=torch.bfloat16)
+    value = torch.randn(1, 16, 8, value_head_size, device="cuda", dtype=torch.bfloat16)
+    context = ExecutionContext(platform="cuda", require_fullgraph=True)
+    path = layer.resolve_execution_path(context, query, query, value, None)
+    assert path.requested_support(context).status is SupportStatus.UNSUPPORTED
+    with pytest.raises(AssertionError) as error:
+        layer(query, query, value)
+    assert str(error.value) in path.support.reason
+
+
+@hardware_test(res={"cuda": "B200"}, num_cards=1)
+@pytest.mark.parametrize("noncontiguous", [False, True])
+def test_real_fa4_custom_op_unequal_value_dimension_schema(noncontiguous):
+    if not current_omni_platform.is_cuda() or not fa.IS_FLASH_ATTN_4:
+        pytest.skip("Requires CuTe FlashAttention-4")
+    query = torch.randn(1, 17, 8, 80, device="cuda", dtype=torch.bfloat16)
+    key = torch.randn(1, 25, 8, 80, device="cuda", dtype=torch.bfloat16)
+    value = torch.randn(1, 25, 8, 48, device="cuda", dtype=torch.bfloat16)
+    if noncontiguous:
+        query, key, value = (t.transpose(1, 2).contiguous().transpose(1, 2) for t in (query, key, value))
+        assert not query.is_contiguous()
+    torch.library.opcheck(
+        torch.ops.vllm_omni.fa4_dense_attention.default,
+        (query, key, value, 80**-0.5, False, False),
+    )

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

@@ -0,0 +1,200 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Compare full and pre-sharded Cosmos loading with real two-rank FSDP."""
+
+from contextlib import ExitStack
+from datetime import timedelta
+from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+from torch import nn
+
+from tests.helpers.mark import hardware_test
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.parallel, pytest.mark.core_model, pytest.mark.gpu]
+
+
+class _SDPAAttention(nn.Module):
+    """Use deterministic native attention while testing loading and FSDP."""
+
+    def __init__(self, *, causal, softmax_scale, **kwargs):
+        super().__init__()
+        self.causal = causal
+        self.scale = softmax_scale
+
+    def forward(self, query, key, value):
+        return F.scaled_dot_product_attention(
+            query.transpose(1, 2),
+            key.transpose(1, 2),
+            value.transpose(1, 2),
+            is_causal=self.causal,
+            scale=self.scale,
+        ).transpose(1, 2)
+
+
+def _load_cosmos_worker(rank, rendezvous, checkpoint_dir, edge):
+    from tests.model_executor.helpers import bootstrap_vllm_layer_custom_op_modules
+
+    bootstrap_vllm_layer_custom_op_modules()
+
+    from safetensors.torch import save_file
+    from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config
+    from vllm.config.load import LoadConfig
+    from vllm.model_executor import parameter
+    from vllm.model_executor.layers import linear
+    from vllm.utils.torch_utils import set_default_torch_dtype
+
+    from vllm_omni.diffusion.distributed import hsdp
+    from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+    from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3, transformer_cosmos3_edge
+    from vllm_omni.platforms import current_omni_platform
+
+    device = torch.device("cuda", rank)
+    current_omni_platform.set_device(device)
+    dist.init_process_group("nccl", init_method=rendezvous, rank=rank, world_size=2, timeout=timedelta(seconds=120))
+    try:
+        with ExitStack() as stack:
+            stack.enter_context(set_current_vllm_config(VllmConfig(device_config=DeviceConfig(device="cuda"))))
+            stack.enter_context(set_default_torch_dtype(torch.bfloat16))
+            # HSDP spans both workers, while every vLLM linear uses TP=1.
+            for module in (linear, parameter):
+                stack.enter_context(patch.object(module, "get_tensor_model_parallel_rank", return_value=0))
+                stack.enter_context(patch.object(module, "get_tensor_model_parallel_world_size", return_value=1))
+            for module in (transformer_cosmos3, transformer_cosmos3_edge):
+                stack.enter_context(patch.object(module, "get_tensor_model_parallel_world_size", return_value=1))
+                stack.enter_context(patch.object(module, "FrameworkAttention", _SDPAAttention))
+            stack.enter_context(patch.object(transformer_cosmos3, "_get_ulysses_state", return_value=(1, 0, None)))
+            stack.enter_context(
+                patch.object(hsdp, "get_world_group", return_value=SimpleNamespace(world_size=2, rank_in_group=rank))
+            )
+
+            model_cls = (
+                transformer_cosmos3_edge.Cosmos3EdgeVFMTransformer
+                if edge
+                else transformer_cosmos3.Cosmos3VFMTransformer
+            )
+            tf_config = dict(
+                hidden_size=16,
+                num_hidden_layers=1,
+                num_attention_heads=2,
+                num_key_value_heads=2,
+                head_dim=8,
+                intermediate_size=32,
+                vocab_size=32,
+                latent_patch_size=2,
+                latent_channel=48,
+                temporal_compression_factor=4,
+                rope_scaling={"mrope_section": [2, 1, 1]},
+            )
+            if edge:
+                tf_config.update(
+                    backbone_type=transformer_cosmos3_edge.COSMOS3_EDGE_BACKBONE_TYPE,
+                    qk_norm_for_text=False,
+                    use_und_k_norm_for_gen=True,
+                )
+            config = SimpleNamespace(
+                tf_model_config=tf_config,
+                dtype=torch.bfloat16,
+                quantization_config=None,
+                lora_path=None,
+                num_weight_load_threads=1,
+                hsdp_weight_load_strategy="full",
+                parallel_config=SimpleNamespace(
+                    use_hsdp=True, hsdp_replicate_size=1, hsdp_shard_size=2, tensor_parallel_size=1
+                ),
+            )
+
+            class Pipeline(nn.Module):
+                def __init__(self):
+                    super().__init__()
+                    self.transformer = model_cls(config)
+                    self.transformer.register_buffer("scratch", torch.tensor([7.0]), persistent=False)
+                    self.vae = nn.Linear(2, 2)
+                    self.weights_sources = [
+                        DiffusersPipelineLoader.ComponentSource(checkpoint_dir, None, None, "transformer.", False)
+                    ]
+
+                def load_weights(self, weights):
+                    state = dict(weights)
+                    self.transformer.load_state_dict(
+                        {name.removeprefix("transformer."): t for name, t in state.items()}
+                    )
+                    self.transformer.post_load_weights()
+                    self.transformer.eval()
+                    self.transformer.validate_loaded_weights(set(state))
+                    return set(state)
+
+            full = Pipeline()
+            if rank == 0:
+                torch.manual_seed(42)
+                with torch.no_grad():
+                    for name, tensor in full.transformer.named_parameters():
+                        tensor.copy_(torch.randn_like(tensor) * 0.02)
+                        if "norm" in name:
+                            tensor.fill_(1)
+                save_file(full.transformer.state_dict(), str(Path(checkpoint_dir) / "model.safetensors"))
+            dist.barrier()
+
+            def load(model):
+                loader = DiffusersPipelineLoader(LoadConfig(), config)
+                with patch.object(loader, "_init_from_load_format", return_value=model):
+                    return loader._load_model_with_hsdp(device)
+
+            full = load(full)
+            config.hsdp_weight_load_strategy = "pre_sharded"
+            pre_sharded = load(Pipeline())
+            full_parameters = dict(full.transformer.named_parameters())
+            from torch.distributed.tensor import DTensor
+
+            for name, actual in pre_sharded.transformer.named_parameters():
+                expected = full_parameters[name]
+                if name.startswith("time_embedder."):
+                    assert not isinstance(actual, DTensor)
+                    assert actual.dtype == torch.float32
+                else:
+                    assert isinstance(actual, DTensor)
+                    actual = actual.full_tensor()
+                    expected = expected.full_tensor()
+                torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+            assert pre_sharded.vae.weight.device == device
+            torch.testing.assert_close(pre_sharded.transformer.scratch, full.transformer.scratch)
+            assert all(t.device.type != "meta" for t in pre_sharded.transformer.buffers())
+
+            inputs = dict(
+                hidden_states=torch.linspace(-1, 1, 48 * 4, device=device).reshape(1, 48, 1, 2, 2),
+                timestep=torch.tensor([1.0], device=device),
+                text_ids=torch.tensor([[1, 2]], device=device),
+                text_mask=torch.ones(1, 2, dtype=torch.long, device=device),
+                video_shape=(1, 2, 2),
+                fps=24.0,
+            )
+            with torch.no_grad():
+                # Also exercise the cached UND path on the second denoising step.
+                for step in (1.0, 0.5):
+                    inputs["timestep"].fill_(step)
+                    expected = full.transformer(**inputs)
+                    actual = pre_sharded.transformer(**inputs)
+                    assert torch.isfinite(actual).all()
+                    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+    finally:
+        dist.destroy_process_group()
+
+
+@pytest.mark.parametrize("edge", [False, True], ids=["cosmos3", "cosmos3_edge"])
+@hardware_test(res={"cuda": "L4"}, num_cards=2)
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.accelerator.device_count() < 2 or not dist.is_nccl_available(),
+    reason="requires two CUDA GPUs and NCCL",
+)
+def test_cosmos3_pre_sharded_matches_full_hsdp(tmp_path, edge):
+    torch.multiprocessing.spawn(
+        _load_cosmos_worker,
+        args=(f"file://{tmp_path / 'rendezvous'}", str(tmp_path), edge),
+        nprocs=2,
+        join=True,
+    )

tests/diffusion/distributed/test_hsdp.py

@@ -3,11 +3,13 @@
 """Unit tests for HSDP (Hybrid Sharded Data Parallel) configuration and utilities."""
 
 import gc
+from types import SimpleNamespace
 
 import pytest
 import torch
 import torch.distributed as dist
 import torch.nn as nn
+from torch.distributed.fsdp import fully_shard
 from torch.distributed.tensor import DeviceMesh, DTensor
 from vllm.utils.network_utils import get_file_store_init_method
 
@@ -18,6 +20,7 @@ from vllm_omni.diffusion.distributed.hsdp import (
     _unshardable_parameters,
     shard_model,
 )
+from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 
 pytestmark = [pytest.mark.diffusion, pytest.mark.parallel, pytest.mark.cpu, pytest.mark.core_model]
 
@@ -338,3 +341,92 @@ class TestHSDPShardConditions:
                 matched.append(name)
         assert "blocks.0" in matched
         assert "blocks.1" in matched
+
+
+def test_pre_sharded_to_empty_materializes_fsdp_internal_storage(cpu_process_group):
+    with torch.device("meta"):
+        model = nn.Module()
+        model.transformer = nn.Sequential(nn.Linear(2, 2, bias=False))
+    model.transformer.register_buffer("scratch", torch.tensor([3.0]), persistent=False)
+    shard_model(
+        model.transformer,
+        mesh=DeviceMesh("cpu", [0]),
+        hsdp_shard_conditions=[lambda name, _module: name == "0"],
+    )
+
+    load_plan = SimpleNamespace(
+        roots=[SimpleNamespace(name="transformer", module=model.transformer)],
+        groups=[SimpleNamespace(module=model.transformer[0])],
+        binding_names=frozenset({"transformer.0.weight"}),
+    )
+    sharded_weight = model.transformer[0].weight
+
+    DiffusersPipelineLoader._materialize_pre_sharded_hsdp_state(model, load_plan, torch.device("cpu"))
+
+    fsdp_param = fully_shard.state(model.transformer[0])._fsdp_param_group.fsdp_params[0]
+    assert model.transformer[0].weight is sharded_weight
+    assert model.transformer[0].weight.to_local().device.type == "cpu"
+    assert fsdp_param._sharded_param_data.device.type == "cpu"
+    assert (
+        fsdp_param._sharded_param_data.untyped_storage().data_ptr()
+        == model.transformer[0].weight.to_local().untyped_storage().data_ptr()
+    )
+    assert torch.equal(model.transformer.scratch, torch.tensor([3.0]))
+
+
+def test_pre_sharded_load_restores_checkpoint_values(cpu_process_group, tmp_path, monkeypatch):
+    from safetensors.torch import save_file
+    from vllm.config.load import LoadConfig
+
+    from vllm_omni.diffusion.distributed import hsdp as hsdp_module
+    from vllm_omni.diffusion.models.utils import release_module_parameters_to_meta
+
+    class Transformer(nn.Module):
+        _hsdp_shard_conditions = [lambda name, module: name == "block"]
+        _hsdp_ignored_modules = ["time_embedder"]
+
+        def __init__(self):
+            super().__init__()
+            self.block = nn.Linear(3, 4, bias=False)
+            self.time_embedder = nn.Linear(2, 2, bias=False, dtype=torch.bfloat16)
+            self.root_weight = nn.Parameter(torch.zeros(2))
+            self.register_buffer("scratch", torch.tensor([7.0]), persistent=False)
+            release_module_parameters_to_meta(self.block)
+
+        def post_load_weights(self):
+            self.time_embedder.to(torch.float32)
+
+    model = nn.Module()
+    model.transformer = Transformer()
+    model.vae = nn.Linear(2, 2)
+    tensors = {
+        "block.weight": torch.arange(12, dtype=torch.float32).reshape(4, 3),
+        "root_weight": torch.tensor([2.0, 3.0]),
+        "time_embedder.weight": torch.arange(4, dtype=torch.bfloat16).reshape(2, 2),
+    }
+    save_file(tensors, str(tmp_path / "model.safetensors"))
+    model.weights_sources = [DiffusersPipelineLoader.ComponentSource(str(tmp_path), None, None, "transformer.", False)]
+    config = SimpleNamespace(
+        dtype=torch.float32,
+        quantization_config=None,
+        lora_path=None,
+        hsdp_weight_load_strategy="pre_sharded",
+        num_weight_load_threads=1,
+        parallel_config=SimpleNamespace(use_hsdp=True, hsdp_replicate_size=1, hsdp_shard_size=1),
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), config)
+    monkeypatch.setattr(loader, "_init_from_load_format", lambda *args, **kwargs: model)
+    monkeypatch.setattr(hsdp_module, "get_world_group", lambda: SimpleNamespace(world_size=1, rank_in_group=0))
+    monkeypatch.setattr(hsdp_module, "current_omni_platform", SimpleNamespace(device_type="cpu"))
+    loaded = loader._load_model_with_hsdp(torch.device("cpu"))
+    assert loaded is model
+    for name, expected in tensors.items():
+        actual = dict(model.transformer.named_parameters())[name]
+        if name.startswith("time_embedder."):
+            assert not isinstance(actual, DTensor)
+            assert actual.dtype == torch.float32
+            torch.testing.assert_close(actual, expected.float())
+        else:
+            assert isinstance(actual, DTensor)
+            torch.testing.assert_close(actual.to_local(), expected)
+    torch.testing.assert_close(model.transformer.scratch, torch.tensor([7.0]))

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/model_loader/test_diffusers_loader.py

@@ -27,6 +27,7 @@ from vllm_omni.diffusion.model_loader.host_weight_plan import (
     HostWeightPlan,
     HostWeightPlanResult,
     TensorBinding,
+    build_checkpoint_binding_plan,
 )
 from vllm_omni.diffusion.model_loader.host_weights import source_identity as source_identity_module
 from vllm_omni.diffusion.models.helios import HeliosPipeline
@@ -2089,3 +2090,206 @@ def test_hsdp_broadcast_weight_load_online_quant_multiprocess():
             )
         assert os.path.exists(os.path.join(temp_dir, "rank_0_success.flag"))
         assert os.path.exists(os.path.join(temp_dir, "rank_1_success.flag"))
+
+
+def test_pre_sharded_hsdp_rejects_quantization_before_loading():
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        hsdp_weight_load_strategy="pre_sharded",
+        lora_path=None,
+        quantization_config=object(),
+        parallel_config=SimpleNamespace(
+            use_hsdp=True,
+            hsdp_replicate_size=1,
+            hsdp_shard_size=2,
+        ),
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    with pytest.raises(ValueError, match="does not support quantization"):
+        loader._load_model_with_hsdp(torch.device("cpu"))
+
+
+def test_pre_sharded_hsdp_strategy_dispatches_without_using_full_loader(mocker):
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+    from vllm_omni.diffusion.offloader.module_collector import PipelineModules
+
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        hsdp_weight_load_strategy="pre_sharded",
+        lora_path=None,
+        quantization_config=None,
+        parallel_config=SimpleNamespace(
+            use_hsdp=True,
+            hsdp_replicate_size=1,
+            hsdp_shard_size=2,
+        ),
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    discovered = PipelineModules(
+        dits=[model.transformer],
+        dit_names=["transformer"],
+        vaes=[],
+        encoders=[],
+        encoder_names=[],
+        resident_modules=[],
+        resident_names=[],
+    )
+    loader._init_from_load_format = mocker.Mock(return_value=model)  # type: ignore[method-assign]
+    loader.load_weights = mocker.Mock(side_effect=AssertionError("full loader must not run"))  # type: ignore[method-assign]
+    mocker.patch.object(loader_mod.ModuleDiscovery, "discover", return_value=discovered)
+    pre_sharded_load = mocker.patch.object(
+        loader,
+        "_load_model_with_pre_sharded_hsdp",
+        return_value=model,
+    )
+
+    assert loader._load_model_with_hsdp(torch.device("cpu")) is model
+    pre_sharded_load.assert_called_once()
+
+
+def test_pre_sharded_hsdp_rejects_tensor_transform():
+    def transform(tensor):
+        return tensor.transpose(-2, -1)
+
+    load_plan = SimpleNamespace(
+        bindings={
+            "transformer.weight": TensorBinding(
+                checkpoint_key="weight",
+                file_path="model.safetensors",
+                transform=transform,
+            )
+        }
+    )
+
+    with pytest.raises(ValueError, match="tensor transforms are unsupported"):
+        DiffusersPipelineLoader._validate_pre_sharded_hsdp_bindings(load_plan)
+
+
+def test_pre_sharded_hsdp_accepts_direct_safetensors_binding():
+    load_plan = SimpleNamespace(
+        bindings={
+            "transformer.weight": TensorBinding(
+                checkpoint_key="weight",
+                file_path="model.safetensors",
+            )
+        }
+    )
+
+    DiffusersPipelineLoader._validate_pre_sharded_hsdp_bindings(load_plan)
+
+
+def test_pre_sharded_hsdp_uses_fsdp_destination_directly():
+    target = torch.empty(8, 4, 3)
+    load_plan = SimpleNamespace(
+        bindings={
+            "transformer.weight": TensorBinding(
+                checkpoint_key="source_weight",
+                file_path="transformer/model.safetensors",
+            )
+        }
+    )
+
+    checkpoint_states, local_bytes = DiffusersPipelineLoader._prepare_pre_sharded_checkpoint_states(
+        load_plan,
+        {"transformer.weight": target},
+    )
+    checkpoint_target = checkpoint_states[Path("transformer")]["source_weight"]
+
+    assert checkpoint_target is target
+    assert checkpoint_target.shape == (8, 4, 3)
+    assert checkpoint_target.device.type == "cpu"
+    assert checkpoint_target.untyped_storage().data_ptr() == target.untyped_storage().data_ptr()
+    assert local_bytes == target.numel() * target.element_size()
+
+
+def test_pre_sharded_hsdp_checkpoint_writes_runtime_storage_directly():
+    runtime_target = torch.empty(2, 3)
+    load_plan = SimpleNamespace(
+        bindings={
+            "transformer.weight": TensorBinding(
+                checkpoint_key="source_weight",
+                file_path="transformer/model.safetensors",
+            )
+        }
+    )
+    checkpoint_tensor = torch.arange(6, dtype=torch.float32).reshape(2, 3)
+    checkpoint_states, _ = DiffusersPipelineLoader._prepare_pre_sharded_checkpoint_states(
+        load_plan,
+        {"transformer.weight": runtime_target},
+    )
+    checkpoint_states[Path("transformer")]["source_weight"].copy_(checkpoint_tensor)
+
+    assert torch.equal(runtime_target, checkpoint_tensor)
+
+
+def test_pre_sharded_hsdp_rejects_duplicate_checkpoint_destinations():
+    binding = TensorBinding(checkpoint_key="weight", file_path="transformer/model.safetensors")
+    load_plan = SimpleNamespace(
+        bindings={
+            "transformer.first": binding,
+            "transformer.second": binding,
+        }
+    )
+
+    with pytest.raises(ValueError, match="multiple runtime tensors"):
+        DiffusersPipelineLoader._prepare_pre_sharded_checkpoint_states(
+            load_plan,
+            {"transformer.first": torch.empty(2, 3), "transformer.second": torch.empty(2, 3)},
+        )
+
+
+def test_pre_sharded_hsdp_managed_module_traversal_is_child_first_and_skips_nested_roots():
+    root = nn.Module()
+    root.left = nn.Sequential(nn.Linear(2, 2), nn.ReLU())
+    root.right = nn.Sequential(nn.Linear(2, 2), nn.ReLU())
+
+    names = {module: name or "root" for name, module in root.named_modules()}
+    modules = DiffusersPipelineLoader._hsdp_managed_modules_post_order(
+        root,
+        nested_hsdp_roots={root.left},
+    )
+
+    assert [names[module] for module in modules] == ["right.0", "right.1", "right", "root"]
+
+
+@pytest.mark.parametrize("checkpoint_key", ["renamed", "unexpected"])
+def test_hsdp_checkpoint_plan_honors_remap_and_rejects_missing(tmp_path, checkpoint_key, monkeypatch):
+    import vllm_omni.diffusion.model_loader.host_weight_plan as plan_mod
+
+    monkeypatch.setattr(plan_mod, "get_direct_mmap_adapter", lambda _model: None)
+
+    class Pipeline(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.transformer = nn.Linear(2, 2, bias=False)
+
+        @staticmethod
+        def remap_checkpoint_key(name):
+            return "transformer.weight" if name == "transformer.renamed" else name
+
+    checkpoint = tmp_path / "model.safetensors"
+    save_file({checkpoint_key: torch.ones(2, 2)}, str(checkpoint))
+    source = SimpleNamespace(
+        model_or_path=str(tmp_path),
+        subfolder=None,
+        revision=None,
+        prefix="transformer.",
+    )
+    model = Pipeline()
+    result = build_checkpoint_binding_plan(
+        model,
+        dit_modules=(("transformer", model.transformer),),
+        sources=(source,),
+        model_path=None,
+        tensor_parallel_size=1,
+        online_quantization=False,
+    )
+
+    if checkpoint_key == "renamed":
+        assert result.plan is not None
+        assert result.plan.bindings["transformer.weight"].checkpoint_key == "renamed"
+    else:
+        assert result.plan is None
+        assert "no checkpoint binding" in result.fallback_reason

tests/diffusion/models/auk/test_auk_cudagraph.py

@@ -61,7 +61,10 @@ def test_single_request_graph_wrapper_cpu_falls_back_to_eager(cfg_strength: floa
 
     # The conditioning is prepared once per request and reused by every step.
     assert prepare_spy.call_count == 1
-    assert prepare_spy.call_args.args[-1] is (cfg_strength >= 1e-5)
+    # _prepare(x, x_mask, text, c_mask, ref, ref_mask, uses_cfg, timesteps)
+    assert prepare_spy.call_args.args[6] is (cfg_strength >= 1e-5)
+    # The time grid reaches prepare, so the adaLN modulations are computed once.
+    assert prepare_spy.call_args.args[7].numel() == 2
     assert step_spy.call_count == 2
     capture_spy.assert_not_called()
     assert not wrapper._cache
@@ -84,7 +87,8 @@ def test_graph_inputs_use_bounded_length_buckets() -> None:
     assert padded[4].shape == (1, 100, 4)
     assert padded[5].shape == (1, 100)
     assert [mask.sum().item() for mask in (padded[1], padded[3], padded[5])] == [65, 65, 51]
-    assert wrapper._key(padded[0], padded[2], padded[4], False) == (96, 96, 100, False)
+    assert wrapper._key(padded[0], padded[2], padded[4], False) == (96, 96, 100, False, 0)
+    assert wrapper._key(padded[0], padded[2], padded[4], False, 32) == (96, 96, 100, False, 32)
     assert wrapper.max_graphs == 32
 
 
@@ -153,18 +157,21 @@ def test_single_request_graph_replay_matches_eager_and_updates_inputs(cfg_streng
             inputs["ref"],
             inputs["ref_mask"],
         )
-        key = wrapper._key(bucketed[0], bucketed[2], bucketed[4], cfg_strength >= 1e-5)
+        # sample_latents passes the grid, so the key carries its step count (two steps here).
+        key = wrapper._key(bucketed[0], bucketed[2], bucketed[4], cfg_strength >= 1e-5, 2)
         assert key in wrapper._cache
 
         # The static context was refreshed for this request's conditioning.
         entry = wrapper._cache[key]
-        fresh = wrapper._prepare(*bucketed, cfg_strength >= 1e-5)
+        grid = torch.tensor([0.0, 0.4], device="cuda")
+        fresh = wrapper._prepare(*bucketed, cfg_strength >= 1e-5, grid)
         for static, want in zip(entry.static_ctx.tensors(), fresh.tensors(), strict=True):
             if want is None:
                 assert static is None
             else:
                 torch.testing.assert_close(static, want)
         torch.testing.assert_close(entry.static_timestep, torch.tensor(0.4, device="cuda"))
+        assert int(entry.static_step) == 1
 
     assert len(wrapper._cache) == 1
 

tests/diffusion/models/auk/test_auk_dit_fastpaths.py

@@ -0,0 +1,139 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""The AuK DiT fast paths keep the reference numerics.
+
+Two paths are checked: the grouped conv position embedding as one batched
+GEMM, and the adaLN modulations precomputed for a whole time grid, eager and
+under the per-step CUDA graph.
+"""
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.auk import auk_transformer as T
+from vllm_omni.diffusion.models.auk.auk_transformer import AuKTransformer, build_time_grid
+from vllm_omni.diffusion.models.auk.cudagraph_wrapper import AuKCUDAGraphWrapper
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_dit(device: str = "cpu") -> AuKTransformer:
+    torch.manual_seed(5)
+    return (
+        AuKTransformer(
+            dim=32,
+            heads=2,
+            dim_head=16,
+            ff_mult=2,
+            latent_dim=4,
+            text_hidden_dim=8,
+            num_layers=2,
+            num_single_layers=2,
+        )
+        .eval()
+        .to(device)
+    )
+
+
+def _reference_conv_pos_embed(module: T.ConvPosEmbedding, x: torch.Tensor, mask: torch.Tensor | None) -> torch.Tensor:
+    """The channels-first conv stack as the reference implementation writes it."""
+    keep = None if mask is None else mask.unsqueeze(1)
+    x = x.transpose(1, 2)
+    if keep is not None:
+        x = x.masked_fill(~keep, 0.0)
+    for layer in module.conv1d:
+        x = layer(x)
+        if keep is not None and isinstance(layer, torch.nn.Conv1d):
+            x = x.masked_fill(~keep, 0.0)
+    return x.transpose(1, 2)
+
+
+@torch.inference_mode()
+def test_grouped_conv_as_batched_gemm_matches_conv1d() -> None:
+    torch.manual_seed(0)
+    conv = torch.nn.Conv1d(64, 64, 31, groups=16, padding=15)
+    x = torch.randn(2, 37, 64)
+    expected = conv(x.transpose(1, 2)).transpose(1, 2)
+    torch.testing.assert_close(T._grouped_conv1d_btc(x, conv), expected, atol=1e-5, rtol=1e-5)
+
+
+@torch.inference_mode()
+@pytest.mark.parametrize("masked", [False, True])
+def test_conv_pos_embedding_matches_the_channels_first_stack(masked: bool) -> None:
+    torch.manual_seed(1)
+    module = T.ConvPosEmbedding(64).eval()
+    x = torch.randn(2, 37, 64)
+    mask = None
+    if masked:
+        mask = torch.ones(2, 37, dtype=torch.bool)
+        mask[1, 30:] = False
+    torch.testing.assert_close(module(x, mask), _reference_conv_pos_embed(module, x, mask), atol=1e-5, rtol=1e-5)
+
+
+@torch.inference_mode()
+@pytest.mark.parametrize("cfg_infer", [False, True])
+def test_modulation_table_matches_per_step_modulation(cfg_infer: bool) -> None:
+    dit = _make_dit()
+    text = torch.randn(1, 5, 8)
+    ref = torch.randn(1, 3, 4)
+    x = torch.randn(1, 6, 4)
+    grid = build_time_grid(nfe=4, sway_sampling_coef=-1.0, t_grid=None, device="cpu")
+    plain = dit.prepare(text, target_len=6, ref=ref, cfg_infer=cfg_infer)
+    tabled = dit.prepare(text, target_len=6, ref=ref, cfg_infer=cfg_infer, timesteps=grid[:-1])
+    assert tabled.modulation is not None and tabled.modulation.shape[0] == 4
+    for i in range(4):
+        expected = dit.step(x, grid[i], plain)
+        torch.testing.assert_close(dit.step(x, grid[i], tabled, step_index=i), expected, atol=1e-5, rtol=1e-5)
+        index = torch.tensor(i)
+        torch.testing.assert_close(dit.step(x, grid[i], tabled, step_index=index), expected, atol=1e-5, rtol=1e-5)
+
+
+def _padded_inputs(device: str) -> dict[str, torch.Tensor]:
+    generator = torch.Generator(device="cpu").manual_seed(3)
+    text = torch.randn(1, 7, 8, generator=generator)
+    c_mask = torch.ones(1, 7, dtype=torch.bool)
+    c_mask[:, 5:] = False
+    text[:, 5:] = 0.0
+    ref = torch.randn(1, 6, 4, generator=generator)
+    ref_mask = torch.ones(1, 6, dtype=torch.bool)
+    ref_mask[:, 4:] = False
+    mask = torch.ones(1, 9, dtype=torch.bool)
+    mask[:, 7:] = False
+    inputs = {"text": text, "c_mask": c_mask, "ref": ref, "ref_mask": ref_mask, "mask": mask}
+    return {name: value.to(device) for name, value in inputs.items()}
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph replay requires CUDA")
+@torch.inference_mode()
+@pytest.mark.parametrize("cfg_strength", [0.0, 2.0])
+def test_graph_with_time_grid_matches_eager_with_time_grid(cfg_strength: float) -> None:
+    dit = _make_dit("cuda")
+    inputs = _padded_inputs("cuda")
+    grid = build_time_grid(nfe=4, sway_sampling_coef=-1.0, t_grid=None, device="cuda")
+    x = torch.randn(1, 9, 4, device="cuda")
+    wrapper = AuKCUDAGraphWrapper(dit)
+    common = dict(
+        text=inputs["text"],
+        c_mask=inputs["c_mask"],
+        ref=inputs["ref"],
+        ref_mask=inputs["ref_mask"],
+        cfg_strength=cfg_strength,
+        timesteps=grid[:-1],
+    )
+    for i in range(4):
+        graph = wrapper(x=x, timestep=grid[i], new_request=i == 0, step_index=i, **common)
+        ctx = dit.prepare(
+            inputs["text"],
+            target_len=9,
+            c_mask=inputs["c_mask"],
+            ref=inputs["ref"],
+            ref_mask=inputs["ref_mask"],
+            cfg_infer=cfg_strength >= 1e-5,
+            timesteps=grid[:-1],
+        )
+        eager = wrapper._step(x, grid[i], ctx, cfg_strength, i)
+        torch.testing.assert_close(graph, eager, atol=1e-4, rtol=1e-4)
+    # One graph serves every step of the grid.
+    assert len(wrapper._cache) == 1
+    assert next(iter(wrapper._cache))[-1] == 4

tests/diffusion/models/auk/test_auk_step_context.py

@@ -127,7 +127,9 @@ def test_hoisted_euler_loop_matches_per_step_forward(cfg_strength: float) -> Non
             v = dit(x, time=timesteps[i], **inputs)
         x = x + (timesteps[i + 1] - timesteps[i]) * v
 
-    torch.testing.assert_close(got, x, rtol=0, atol=0)
+    # sample_latents precomputes the adaLN modulations of the whole grid in one
+    # GEMM per layer, whose rows round differently from one-row GEMMs in fp32.
+    torch.testing.assert_close(got, x, rtol=0, atol=1e-6)
     assert dit.text_cond is None and dit.text_uncond is None