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

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.

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 |
 

docs/configuration/stage_configs.md

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

docs/design/feature/attention_execution_contract_poc.md

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

docs/design/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/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/models/supported_models.md

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

docs/serving/realtime_duplex_api.md

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

docs/user_guide/diffusion_features.md

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

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]

recipes/MammothModa2/MammothModa2.md

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

recipes/MiniMaxAI/MiniMax-H3.md

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

recipes/README.md

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

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

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

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

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

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/config/test_omni_config.py

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

tests/conftest.py

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

tests/core/sched/test_omni_ar_scheduler_streaming.py

@@ -23,6 +23,7 @@ from vllm.v1.engine import FinishReason
 from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue
 from vllm.v1.request import Request, RequestStatus, StreamingUpdate
 from vllm_omni.core.sched.omni_ar_scheduler import OmniARAsyncScheduler, OmniARScheduler
+from vllm_omni.model_executor.stage_input_processors.cosyvoice3 import talker2code2wav_async_chunk
 from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import (
     OmniChunkTransferAdapter,
 )
@@ -288,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
@@ -297,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()
@@ -317,13 +325,13 @@ def test_running_decode_step_without_inter_stage_payload_does_not_raise() -> Non
     sched.make_stats.return_value = None
     bind_omits_transfer_helpers(sched)
 
-    scheduler_output = MagicMock(spec=SchedulerOutput)
+    scheduler_output = mocker.MagicMock(spec=SchedulerOutput)
     scheduler_output.num_scheduled_tokens = {session.request_id: 1}
     scheduler_output.scheduled_spec_decode_tokens = {}
     scheduler_output.num_invalid_spec_tokens = 0
 
-    model_runner_output = MagicMock(spec=ModelRunnerOutput)
-    model_runner_output.sampled_token_ids = [[42]]
+    model_runner_output = mocker.MagicMock(spec=ModelRunnerOutput)
+    model_runner_output.sampled_token_ids = [new_token_ids]
     model_runner_output.logprobs = None
     model_runner_output.prompt_logprobs_dict = {}
     model_runner_output.pooler_output = None
@@ -336,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:

tests/diffusion/ar_diffusion/test_history_staging_gpu.py

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

tests/diffusion/ar_diffusion/test_paged_attention.py

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

tests/diffusion/attention/test_trtllm_attn.py

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

tests/diffusion/attention/test_trtllm_contract.py

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

tests/diffusion/cache/test_cache_backends.py

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

tests/diffusion/cache/test_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/models/mammoth_moda2/test_dit_attention.py

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

tests/diffusion/models/mammoth_moda2/test_dit_attention_cuda.py

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

tests/diffusion/models/mammoth_moda2/test_pipeline_mammothmoda2_dit.py

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

tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py

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

tests/diffusion/offloader/test_layerwise_backend.py

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

tests/diffusion/test_diffusion_engine_cleanup.py

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

tests/diffusion/test_multiproc_engine_concurrency.py

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

tests/diffusion/test_ray_executor.py

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

tests/diffusion/test_stage_diffusion_proc.py

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

tests/diffusion/test_worker_wrapper_base.py

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

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -31,10 +31,10 @@ from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
-@pytest.fixture
-def shm_sender(build_adapter):
+@pytest.fixture(params=[0, 4 << 20], ids=["per-key-shm", "host-ring"])
+def shm_sender(build_adapter, request):
     adapter, _ = build_adapter(stage_id=0)
-    connector = SharedMemoryConnector({"stage_id": 0})
+    connector = SharedMemoryConnector({"stage_id": 0, "host_ring_bytes": request.param})
     adapter.connector = connector
     adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct(
         codes=CodesStruct(audio=torch.tensor([1], dtype=torch.long))
@@ -111,13 +111,14 @@ def test_shm_release_reclaims_undrained_chunks_after_finish(shm_sender):
 
     adapter.release_shm_resources(ext_id)
     # The release is queued behind sends; nothing is unlinked on the caller.
-    assert f"{ext_id}_0_1" in connector._pending_keys
+    assert adapter._reclaimed_shm_total == 0
     while adapter._pending_save_reqs:
         adapter._send_single_request(adapter._pending_save_reqs.popleft())
 
     assert adapter._reclaimed_shm_total == 2
     assert not any(k.startswith(f"{ext_id}_0_") for k in connector._pending_keys)
     assert connector.get("0", "1", f"{ext_id}_0_1") is None
+    assert connector.get("0", "1", f"{ext_id}_0_2") is None
     assert connector.get("0", "1", sibling_key)[0] == "sibling"
 
 

tests/distributed/omni_connectors/test_shm_connector.py

@@ -19,7 +19,9 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 @pytest.fixture()
 def connector():
-    c = SharedMemoryConnector({})
+    # These cases inspect the per-key segment/lock protocol. The enabled ring
+    # protocol and its pressure fallback are covered in test_shm_ring.py.
+    c = SharedMemoryConnector({"host_ring_bytes": 0})
     yield c
     c.close()
 

tests/distributed/omni_connectors/test_shm_ring.py

@@ -0,0 +1,527 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Cross-process ownership, pressure and compatibility of the default host path."""
+
+import fcntl
+import multiprocessing
+import os
+import time
+import uuid
+from dataclasses import asdict
+from multiprocessing import shared_memory
+from pathlib import Path
+
+import msgspec
+import numpy as np
+import pytest
+import torch
+from PIL import Image
+from vllm.outputs import CompletionOutput, RequestOutput
+
+from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayloadStruct
+from vllm_omni.distributed.omni_connectors.connectors.shm_connector import SharedMemoryConnector
+from vllm_omni.distributed.omni_connectors.utils.serialization import OmniSerializer
+from vllm_omni.distributed.omni_connectors.utils.tensor_frame import prepare_tensor_frame
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.fixture
+def edge():
+    scope = uuid.uuid4().hex
+    sender = SharedMemoryConnector({"stage_id": 0, "wakeup_scope": scope, "host_ring_bytes": 4096})
+    receiver = SharedMemoryConnector({"stage_id": 1, "wakeup_scope": scope, "host_ring_bytes": 4096})
+    yield sender, receiver
+    receiver.close()
+    sender.close()
+
+
+@pytest.mark.parametrize(
+    "dtype",
+    [
+        torch.bool,
+        torch.uint8,
+        torch.int16,
+        torch.int32,
+        torch.int64,
+        torch.float16,
+        torch.bfloat16,
+        torch.float32,
+        torch.float64,
+        torch.complex64,
+    ],
+)
+def test_owned_tensor_tree_preserves_values_shapes_and_types(edge, dtype):
+    sender, receiver = edge
+    tensor = torch.arange(24).to(dtype).reshape(4, 6).T
+    data = {
+        "codes": {"audio": tensor},
+        "other": (tensor[0], [torch.empty(0, dtype=dtype), torch.tensor(1)]),
+        "__tensor__": False,
+        "meta": {"finished": torch.tensor(False)},
+    }
+    ok, size, metadata = sender.put("0", "1", "tree", data)
+    assert ok and "host_ring" in metadata
+    output, read_size = receiver.get("0", "1", "tree", metadata)
+    assert read_size == size and output["__tensor__"] is False
+    assert type(output["other"]) is list and type(output["other"][1]) is list
+    assert output["other"][1][0].shape == (0,)
+    assert output["other"][1][1].shape == ()
+    assert output["codes"]["audio"].dtype == dtype
+    assert torch.equal(output["codes"]["audio"], tensor)
+    tensor.fill_(0)
+    assert torch.count_nonzero(output["codes"]["audio"]) > 0
+    for index in range(100):
+        sender.put("0", "1", f"wrap{index}", {"codes": torch.full((56,), index)})
+        assert receiver.get("0", "1", f"wrap{index}")[0]["codes"][0] == index
+    assert torch.count_nonzero(output["codes"]["audio"]) > 0
+
+
+def test_conjugate_tensor_and_scalar_finish(edge):
+    sender, receiver = edge
+    tensor = torch.tensor([1 + 2j, 3 - 4j]).conj()
+    sender.put("0", "1", "conjugate", {"condition": tensor, "finished": torch.tensor(True)})
+    result = receiver.get("0", "1", "conjugate")[0]
+    assert torch.equal(result["condition"], tensor)
+    assert result["finished"].shape == () and result["finished"].item() is True
+
+
+def test_shared_typed_payload_keeps_omitted_defaults_and_cpu_tensors(edge):
+    sender, receiver = edge
+    payload = OmniPayloadStruct(
+        codes=CodesStruct(audio=torch.arange(28).reshape(1, 28)),
+        meta=MetaStruct(finished=torch.tensor(False), left_context_size=0),
+        kv_metadata={"binary_handle": b"\x00\xff\x10"},
+    )
+    ok, _, metadata = sender.put("0", "1", "typed", payload)
+    assert ok and "host_ring" in metadata
+    result = receiver.get("0", "1", "typed")[0]
+    assert set(result) == {"codes", "meta", "kv_metadata"}
+    assert set(result["codes"]) == {"audio"}
+    assert torch.equal(result["codes"]["audio"], payload.codes.audio)
+    assert result["meta"]["left_context_size"] == 0
+    assert result["kv_metadata"]["binary_handle"] == b"\x00\xff\x10"
+
+
+def test_source_can_be_overwritten_as_soon_as_put_returns(edge):
+    sender, receiver = edge
+    source = torch.arange(56)
+    sender.put("0", "1", "source-reuse", source)
+    source.fill_(-1)
+    assert receiver.get("0", "1", "source-reuse")[0].tolist() == list(range(56))
+
+
+@pytest.mark.parametrize("dtype", [np.bool_, np.int32, np.float16, np.float64, np.complex64])
+def test_mixed_numpy_image_and_tensor_keep_shared_wire_semantics(edge, dtype):
+    sender, receiver = edge
+    array = np.arange(24).astype(dtype).reshape(4, 6)[::-1, ::2]
+    image = Image.fromarray(np.arange(18, dtype=np.uint8).reshape(2, 3, 3))
+    payload = {"codes": torch.arange(28), "array": array, "image": image, "slice": slice(1, 7, 2)}
+    _, _, metadata = sender.put("0", "1", "mixed-native", payload)
+    assert "host_ring" in metadata
+    output = receiver.get("0", "1", "mixed-native", metadata)[0]
+    assert np.array_equal(output["array"], array)
+    assert output["array"].dtype == array.dtype
+    assert not output["array"].flags.writeable
+    assert output["image"].mode == image.mode and output["image"].tobytes() == image.tobytes()
+    assert output["slice"] == [1, 7, 2]
+    assert torch.equal(output["codes"], payload["codes"])
+    array.fill(0)
+    image.paste(0, (0, 0, 3, 2))
+    assert np.count_nonzero(output["array"]) and any(output["image"].tobytes())
+
+
+def test_large_array_rejects_unused_framing_snapshot(monkeypatch):
+    def unexpected_snapshot(value):
+        pytest.fail("a too-large array must not be snapshotted while testing ring admission")
+
+    monkeypatch.setattr(OmniSerializer.encoder, "_enc_hook", unexpected_snapshot)
+    assert prepare_tensor_frame({"array": np.zeros(1024)}, max_bytes=1024) is None
+
+
+def test_native_scalar_restore_preserves_nested_markers_and_input_containers():
+    array = np.arange(3)
+    image = Image.fromarray(np.array([[1, 2]], dtype=np.uint8))
+    tensor = torch.arange(2)
+    payload = {"ids": list(range(8192)), "mixed": [1.5, None, (array, {"image": image, "tensor": tensor})]}
+    decoded = msgspec.msgpack.decode(OmniSerializer.serialize(payload))
+    output = OmniSerializer.restore(decoded)
+    assert output is not decoded and output["ids"] is not decoded["ids"]
+    assert output["ids"] == payload["ids"]
+    assert np.array_equal(output["mixed"][2][0], array)
+    assert output["mixed"][2][1]["image"].tobytes() == image.tobytes()
+    assert torch.equal(output["mixed"][2][1]["tensor"], tensor)
+    assert isinstance(decoded["mixed"][2][0], dict)
+
+
+def test_fallback_miss_does_not_register_an_unclaimed_allocation(edge, monkeypatch, caplog):
+    sender, receiver = edge
+    key = "fallback-claim-" + uuid.uuid4().hex
+    _, _, metadata = sender.put("0", "1", key, torch.arange(4096))
+    assert "shm" in metadata
+    original = shared_memory.SharedMemory
+    calls = []
+
+    def tracked_open(*args, **kwargs):
+        calls.append((args, kwargs))
+        return original(*args, **kwargs)
+
+    monkeypatch.setattr(shared_memory, "SharedMemory", tracked_open)
+    with open(f"/dev/shm/shm_{key}_lockfile.lock", "rb+") as lock:
+        fcntl.flock(lock, fcntl.LOCK_EX)
+        assert receiver.get("0", "1", key) is None
+        assert not calls
+    assert torch.equal(receiver.get("0", "1", key)[0], torch.arange(4096))
+    assert receiver.get("0", "1", key, metadata) is None
+    assert not any(record.levelname == "ERROR" for record in caplog.records)
+
+
+def test_tensor_frames_preserve_shared_output_reconstruction(edge):
+    sender, receiver = edge
+    completion = asdict(CompletionOutput(index=0, text="", token_ids=[1], cumulative_logprob=0.0, logprobs=None))
+    completion["multimodal_output"] = {"codes": torch.arange(28)}
+    payload = {
+        "request_id": "output",
+        "prompt": None,
+        "prompt_token_ids": [1],
+        "prompt_logprobs": None,
+        "outputs": [completion],
+        "finished": False,
+    }
+    _, _, metadata = sender.put("0", "1", "output", payload)
+    assert "host_ring" in metadata
+    result = receiver.get("0", "1", "output", metadata)[0]
+    assert isinstance(result, RequestOutput)
+    assert isinstance(result.outputs[0], CompletionOutput)
+    assert torch.equal(result.outputs[0].multimodal_output["codes"], torch.arange(28))
+
+
+def test_large_tensor_does_not_share_owner_with_small_metadata():
+    with SharedMemoryConnector({}) as connector:
+        payload = {"hidden": torch.ones(16384), "finished": torch.tensor(False)}
+        assert connector.put("0", "1", "large-owner", payload)[0]
+        output = connector.get("0", "1", "large-owner")[0]
+        assert output["finished"].untyped_storage().nbytes() == 1
+        assert output["hidden"].untyped_storage().nbytes() == 65536
+
+
+def test_full_ring_and_large_frame_use_existing_shm(edge):
+    sender, receiver = edge
+    payload = {"codes": torch.arange(56)}
+    metadata = {}
+    for index in range(12):
+        ok, _size, metadata[index] = sender.put("0", "1", f"full{index}", payload)
+        assert ok
+    assert any("host_ring" in item for item in metadata.values())
+    assert any("shm" in item for item in metadata.values())
+    # Unclaimed earlier frames pin the ring; later requests still arrive.
+    for index in reversed(range(12)):
+        assert torch.equal(receiver.get("0", "1", f"full{index}")[0]["codes"], payload["codes"])
+    ok, _size, large_metadata = sender.put("0", "1", "large", torch.arange(4096))
+    assert ok and "shm" in large_metadata
+    assert torch.equal(receiver.get("0", "1", "large", large_metadata)[0], torch.arange(4096))
+    assert "host_ring" in sender.put("0", "1", "recovered", payload)[2]
+    assert receiver.get("0", "1", "recovered") is not None
+
+
+def test_reput_across_ring_and_fallback_and_stale_descriptor(edge):
+    sender, receiver = edge
+    _, _, old = sender.put("0", "1", "reuse", torch.tensor([11]))
+    _, _, current = sender.put("0", "1", "reuse", torch.tensor([22]))
+    assert receiver.get("0", "1", "reuse", old) is None
+    assert receiver.get("0", "1", "reuse", current)[0].item() == 22
+    sender.put("0", "1", "reuse", torch.tensor([33]))
+    sender.put("0", "1", "reuse", torch.full((4096,), 44))
+    assert receiver.get("0", "1", "reuse")[0][0].item() == 44
+    sender.put("0", "1", "reuse", torch.full((4096,), 55))
+    sender.put("0", "1", "reuse", torch.tensor([66]))
+    assert receiver.get("0", "1", "reuse")[0].item() == 66
+    assert not os.path.exists("/dev/shm/reuse")
+
+
+def test_cancel_without_receiver_scan_returns_credit_and_prefix_is_exact(edge):
+    sender, receiver = edge
+    for key in ("req_0_0", "req_0_1", "req_0_10other", "request_0_0"):
+        assert sender.put("0", "1", key, torch.arange(56))[0]
+    assert sender.cleanup_prefix("req_0_") == 2
+    assert receiver.get("0", "1", "req_0_0") is None
+    assert receiver.get("0", "1", "req_0_10other") is not None
+    assert receiver.get("0", "1", "request_0_0") is not None
+    assert "host_ring" in sender.put("0", "1", "next", torch.arange(56))[2]
+
+
+def test_flat_scope_and_metadata_work_across_different_parents(monkeypatch):
+    monkeypatch.setattr(os, "getppid", lambda: 100)
+    sender = SharedMemoryConnector({"stage_id": 0, "wakeup_scope": "flat-scope"})
+    monkeypatch.setattr(os, "getppid", lambda: 200)
+    receiver = SharedMemoryConnector({"stage_id": 1, "wakeup_scope": "flat-scope"})
+    try:
+        assert sender._wake_directory == receiver._wake_directory
+        assert sender.put("0", "1", "scope", torch.tensor([7]))[0]
+        assert receiver.get("0", "1", "scope")[0].item() == 7
+    finally:
+        sender.close()
+        receiver.close()
+
+
+def test_empty_poll_reuses_discovery_and_never_waits_for_a_ring_lock(edge, monkeypatch):
+    sender, receiver = edge
+    sender.put("0", "1", "prime-poll", torch.arange(3))
+    assert receiver.get("0", "1", "prime-poll") is not None
+
+    def unexpected_work(*args, **kwargs):
+        pytest.fail("an empty, unchanged channel must not rescan discovery or acquire its frame lock")
+
+    monkeypatch.setattr("vllm_omni.distributed.omni_connectors.connectors.shm_ring.glob.iglob", unexpected_work)
+    monkeypatch.setattr(fcntl, "flock", unexpected_work)
+    for index in range(128):
+        assert receiver.get("0", "1", f"missing-poll-{index}") is None
+    monkeypatch.undo()
+    sender.put("0", "1", "new-poll", torch.arange(5))
+    assert receiver.get("0", "1", "new-poll")[0].tolist() == list(range(5))
+
+
+def test_new_producer_is_discovered_after_a_miss_even_when_the_clock_does_not_advance(edge, monkeypatch):
+    sender, receiver = edge
+    sender.put("0", "1", "first-producer", torch.arange(3))
+    assert receiver.get("0", "1", "first-producer") is not None
+    assert receiver.get("0", "1", "second-producer") is None
+    old_version = receiver._host_ring._discovery_versions["0", "1"]
+    monkeypatch.setattr("vllm_omni.distributed.omni_connectors.connectors.shm_ring.time.time_ns", lambda: 0)
+    with SharedMemoryConnector(sender.config) as second:
+        second.put("0", "1", "second-producer", torch.arange(5))
+        assert receiver.get("0", "1", "second-producer")[0].tolist() == list(range(5))
+        new_version = receiver._host_ring._discovery_versions["0", "1"]
+        assert new_version[2] == old_version[2] + 1
+
+
+def test_one_receiver_discovers_each_input_edge_at_the_same_registry_version(edge):
+    sender, receiver = edge
+    with SharedMemoryConnector({**sender.config, "stage_id": 2}) as second:
+        sender.put("0", "1", "fanin-first", torch.arange(3))
+        second.put("2", "1", "fanin-second", torch.arange(5))
+        assert receiver.get("0", "1", "fanin-first")[0].tolist() == list(range(3))
+        assert receiver.get("2", "1", "fanin-second")[0].tolist() == list(range(5))
+
+
+def test_discovery_attaches_all_producers_even_when_the_first_one_has_the_requested_key(edge, monkeypatch):
+    sender, receiver = edge
+    with SharedMemoryConnector(sender.config) as second:
+        first_metadata = sender.put("0", "1", "multi-first", torch.arange(3))[2]
+        second_metadata = second.put("0", "1", "multi-second", torch.arange(5))[2]
+        paths = [
+            f"/dev/shm/{receiver._host_ring.directory}/{metadata['host_ring']['name']}"
+            for metadata in (first_metadata, second_metadata)
+        ]
+        monkeypatch.setattr(
+            "vllm_omni.distributed.omni_connectors.connectors.shm_ring.glob.iglob", lambda _: iter(paths)
+        )
+        assert receiver.get("0", "1", "multi-first")[0].tolist() == list(range(3))
+        assert receiver.get("0", "1", "multi-second")[0].tolist() == list(range(5))
+
+
+def test_closed_channel_releases_reader_cache_while_the_producer_process_is_still_alive(edge):
+    sender, receiver = edge
+    sender.put("0", "1", "close-cache", torch.arange(3))
+    assert receiver.get("0", "1", "close-cache") is not None
+    channel = next(iter(receiver._host_ring.readers.values()))
+    fd = channel.fd
+    sender.close()
+    assert receiver.get("0", "1", "not-published-after-close") is None
+    assert not receiver._host_ring.readers
+    with pytest.raises(OSError):
+        os.fstat(fd)
+
+
+def test_disabled_ring_and_unsupported_tree_keep_legacy_serializer(edge):
+    sender, receiver = edge
+    with SharedMemoryConnector({"host_ring_bytes": 0}) as legacy:
+        _, _, metadata = legacy.put("0", "1", "legacy", {"codes": torch.arange(3)})
+        assert "shm" in metadata
+        assert receiver.get("0", "1", "legacy", metadata)[0]["codes"].tolist() == [0, 1, 2]
+    assert prepare_tensor_frame({"codes": torch.arange(3), "opaque": object()}) is None
+    assert sender.put("0", "1", "falsey", False)[0]
+    assert receiver.get("0", "1", "falsey")[0] is False
+
+
+def test_tensor_free_lists_reuse_native_serialization():
+    payload = {"prompt_token_ids": list(range(8192)), "condition": [0.25] * 8192}
+    with SharedMemoryConnector({}) as connector:
+        frame = prepare_tensor_frame(payload)
+        assert isinstance(frame, bytes)
+        assert frame == connector.serialize_obj(payload)
+        assert connector.put("0", "1", "native-lists", payload)[0]
+        assert connector.get("0", "1", "native-lists")[0] == payload
+
+
+def test_scalar_lists_and_literal_extensions_preserve_wire_semantics():
+    literal = msgspec.msgpack.Ext(42, bytes(20))
+    payload = {"condition": [0.25] * 8192, "codes": torch.arange(28), "extension": literal}
+    with SharedMemoryConnector({}) as connector:
+        _, _, metadata = connector.put("0", "1", "mixed-lists", payload)
+        assert "host_ring" in metadata
+        result = connector.get("0", "1", "mixed-lists")[0]
+        assert result["condition"] == payload["condition"]
+        assert torch.equal(result["codes"], payload["codes"])
+        assert result["extension"].code == literal.code
+        assert result["extension"].data == literal.data
+
+
+def _locked_writer(name, ready, release):
+    fd = os.open(f"/dev/shm/{name}", os.O_RDWR)
+    try:
+        fcntl.flock(fd, fcntl.LOCK_EX)
+        ready.set()
+        assert release.wait(10)
+    finally:
+        os.close(fd)
+
+
+@pytest.mark.parametrize("use_metadata", [False, True])
+def test_deadline_receive_never_waits_for_another_process_ring_lock(edge, use_metadata):
+    sender, receiver = edge
+    _, _, metadata = sender.put("0", "1", "locked", torch.arange(3))
+    ctx = multiprocessing.get_context("spawn")
+    ready, release = ctx.Event(), ctx.Event()
+    writer = ctx.Process(target=_locked_writer, args=(metadata["host_ring"]["name"], ready, release))
+    writer.start()
+    try:
+        assert ready.wait(15)
+        start = time.monotonic()
+        assert (
+            receiver.get_with_deadline("0", "1", "locked", metadata if use_metadata else None, deadline=start + 1)
+            is None
+        )
+        assert time.monotonic() - start < 1
+    finally:
+        release.set()
+        writer.join(15)
+        if writer.is_alive():
+            writer.kill()
+            writer.join()
+    assert writer.exitcode == 0
+    assert receiver.get("0", "1", "locked", metadata if use_metadata else None)[0].tolist() == [0, 1, 2]
+
+
+def _competing_reader(scope, barrier, results):
+    with SharedMemoryConnector({"stage_id": 1, "wakeup_scope": scope}) as receiver:
+        barrier.wait(timeout=20)
+        result = receiver.get("0", "1", "shared")
+        results.put(None if result is None else result[0]["codes"].tolist())
+
+
+def test_two_reader_processes_claim_exactly_once_and_close_does_not_unlink_owner():
+    scope = uuid.uuid4().hex
+    ctx = multiprocessing.get_context("spawn")
+    barrier, results = ctx.Barrier(3), ctx.Queue()
+    with SharedMemoryConnector({"stage_id": 0, "wakeup_scope": scope}) as sender:
+        _, _, metadata = sender.put("0", "1", "shared", {"codes": torch.arange(56)})
+        readers = [ctx.Process(target=_competing_reader, args=(scope, barrier, results)) for _ in range(2)]
+        for reader in readers:
+            reader.start()
+        try:
+            barrier.wait(timeout=30)
+            outputs = [results.get(timeout=20) for _ in readers]
+            assert outputs.count(None) == 1
+            assert next(output for output in outputs if output is not None) == list(range(56))
+            assert os.path.exists(f"/dev/shm/{metadata['host_ring']['name']}")
+        finally:
+            for reader in readers:
+                reader.join(15)
+                if reader.is_alive():
+                    reader.kill()
+                    reader.join()
+            results.close()
+        assert all(reader.exitcode == 0 for reader in readers)
+
+
+def test_ring_setup_failure_keeps_connector_usable(edge, monkeypatch):
+    sender, receiver = edge
+
+    def no_space(*args):
+        raise OSError(28, "No space left on device")
+
+    monkeypatch.setattr(os, "posix_fallocate", no_space)
+    ok, _, metadata = sender.put("0", "1", "no-ring-space", torch.arange(3))
+    assert ok and "shm" in metadata
+    assert receiver.get("0", "1", "no-ring-space")[0].tolist() == [0, 1, 2]
+
+
+def test_discovery_stamp_failure_keeps_fallback_and_cleanup_usable(edge, monkeypatch):
+    sender, receiver = edge
+
+    def stamp_denied(*args, **kwargs):
+        raise PermissionError("directory timestamp update denied")
+
+    monkeypatch.setattr(os, "utime", stamp_denied)
+    ok, _, metadata = sender.put("0", "1", "no-ring-stamp", torch.arange(3))
+    assert ok and "shm" in metadata
+    assert not sender._host_ring.producers
+    assert not list(Path(f"/dev/shm/{sender._host_ring.directory}").iterdir())
+    assert receiver.get("0", "1", "no-ring-stamp")[0].tolist() == [0, 1, 2]
+    sender.close()
+
+
+def test_close_is_idempotent_and_removes_owned_ring(edge):
+    sender, receiver = edge
+    _, _, metadata = sender.put("0", "1", "close", torch.arange(3))
+    receiver.get("0", "1", "close")
+    sender.close()
+    sender.close()
+    receiver.close()
+    receiver.close()
+    assert not os.path.exists(f"/dev/shm/{metadata['host_ring']['name']}")
+
+
+def _orphan_producer(scope, pipe):
+    producer = SharedMemoryConnector({"stage_id": 0, "wakeup_scope": scope})
+    metadata = producer.put("0", "1", "restart", torch.tensor([11]))[2]
+    pipe.send(metadata)
+    pipe.recv()
+    pipe.close()
+    os._exit(0)
+
+
+def test_dead_producer_payload_cannot_reach_a_reused_key():
+    scope = uuid.uuid4().hex
+    ctx = multiprocessing.get_context("spawn")
+    parent, child = ctx.Pipe()
+    process = ctx.Process(target=_orphan_producer, args=(scope, child))
+    process.start()
+    child.close()
+    old = parent.recv()
+    try:
+        with SharedMemoryConnector({"stage_id": 1, "wakeup_scope": scope}) as receiver:
+            # Cache the live channel without claiming the pending frame.
+            assert receiver.get("0", "1", "not-published") is None
+            channel = next(iter(receiver._host_ring.readers.values()))
+            fd = channel.fd
+            parent.send("exit")
+            process.join(15)
+            assert process.exitcode == 0
+            assert receiver.get("0", "1", "restart", old) is None
+            assert not receiver._host_ring.readers
+            with pytest.raises(OSError):
+                os.fstat(fd)
+            with SharedMemoryConnector({"stage_id": 0, "wakeup_scope": scope}) as replacement:
+                replacement.put("0", "1", "restart", torch.tensor([22]))
+                assert receiver.get("0", "1", "restart")[0].item() == 22
+    finally:
+        if process.is_alive():
+            process.kill()
+            process.join()
+        # The multiprocessing parent still owns this test's resource tracker.
+        # Retire only the intentionally orphaned allocation and its exact marker.
+        segment = shared_memory.SharedMemory(name=old["host_ring"]["name"])
+        segment.close()
+        segment.unlink()
+        directory = old["host_ring"]["name"].rsplit("_", 4)[0]
+        try:
+            os.unlink(f"/dev/shm/{directory}/{old['host_ring']['name']}")
+        except FileNotFoundError:
+            pass
+        try:
+            os.rmdir(f"/dev/shm/{directory}")
+        except FileNotFoundError:
+            pass
+        parent.close()

tests/e2e/accuracy/ltx/test_ltx_official_similarity.py

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

tests/engine/duplex/test_aura_resumable_hooks.py

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

tests/engine/duplex/test_delivery.py

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

tests/engine/duplex/test_session_manager.py

@@ -25,13 +25,17 @@ from vllm_omni.engine.duplex.contracts import (
     DuplexStageSubmissionResult,
     duplex_resource_request_belongs_to_session,
 )
+from vllm_omni.engine.duplex.delivery import DuplexOutputBuffer
 from vllm_omni.engine.duplex.events import (
+    AudioDelta,
     DuplexEvent,
     ErrorEvent,
+    ResponseDone,
     SessionClosed,
     SessionCreated,
     SessionExpired,
     SessionHeartbeatAck,
+    TranscriptDelta,
 )
 from vllm_omni.engine.duplex.messages import (
     CloseDuplexSessionMessage,
@@ -297,6 +301,7 @@ class Harness:
     output_sink: asyncio.Queue = field(default_factory=asyncio.Queue)
     result_sink: asyncio.Queue = field(default_factory=asyncio.Queue)
     _control_seq: int = 0
+    outputs: dict[str, DuplexOutputBuffer] = field(default_factory=dict)
 
     @classmethod
     def create(
@@ -330,11 +335,20 @@ class Harness:
         await self.manager.shutdown()
 
     async def open(self, session_id: str, config: DuplexSessionConfig | None = None) -> DuplexControlResultMessage:
+        limits = self.manager.runtime_config
+        output_buffer = self.outputs.setdefault(
+            session_id,
+            DuplexOutputBuffer(
+                max_bytes=limits.max_pending_output_bytes_per_session,
+                max_events=limits.max_pending_output_events_per_session,
+            ),
+        )
         await self.manager.handle(
             OpenDuplexSessionMessage(
                 control_id=f"open-{session_id}",
                 session_id=session_id,
                 session_config=config or DuplexSessionConfig(model="fake-model"),
+                output_buffer=output_buffer,
             )
         )
         return await self.result()
@@ -381,8 +395,14 @@ class Harness:
         assert isinstance(result, DuplexControlResultMessage)
         return result
 
-    def events(self, session_id: str | None = None) -> list[DuplexEvent]:
+    async def events(self, session_id: str | None = None) -> list[DuplexEvent]:
         events: list[DuplexEvent] = []
+        for identity, output in self.outputs.items():
+            while output.pending_events:
+                event = await output.get()
+                assert event is not None and event.session_id == identity
+                if output.is_valid(event) and (session_id is None or identity == session_id):
+                    events.append(event)
         while not self.output_sink.empty():
             message = self.output_sink.get_nowait()
             assert isinstance(message, DuplexSessionEventMessage)
@@ -396,6 +416,22 @@ class Harness:
         assert session is not None
         return session
 
+    async def begin_output_response(self, session_id: str) -> str:
+        """Initialize real response, text and audio projection before filling output."""
+        runner = self.manager.runners[session_id]
+        response_id = runner.session.begin_response()
+        runner.session.append_assistant_text("abcdefghij")
+        payloads: list[dict[str, object]] = [
+            {"type": "response.created", "modalities": ["audio", "text"]},
+            {"type": "response.text.delta", "delta": "abcdefghij"},
+            {"type": "response.output_audio.delta", "audio": "AAAAAA==", "format": "pcm16"},
+        ]
+        for payload in payloads:
+            runner.emit({**payload, "response_id": response_id, "epoch": runner.session.epoch})
+            await self.events(session_id)
+        assert not runner.run.closing
+        return response_id
+
     def capture_submissions(self, session_id: str) -> list[DuplexCommand]:
         """Replace the runner mailbox with a list so admitted commands can be inspected."""
         runner = self.manager.runners[session_id]
@@ -433,7 +469,7 @@ async def test_open_answers_with_capabilities_and_emits_session_created() -> Non
         assert result.public_session["voice"] == "test"
         assert result.error_code is None
 
-        events = harness.events("sid-open")
+        events = await harness.events("sid-open")
         assert isinstance(events[0], SessionCreated)
         assert events[0].epoch == 0
         assert events[0].session["id"] == "sid-open"
@@ -587,7 +623,7 @@ async def test_plugin_runtime_config_rejection_rolls_back_the_open(error: Except
         assert harness.manager.get("sid-runtime") is None
         assert harness.manager.active_count() == 0
         assert harness.stage_port.ensure_calls == []
-        assert harness.events() == []
+        assert await harness.events() == []
 
         # Nothing was registered under the rejected id and the slot is free again.
         plugin.runtime_config_error = None
@@ -603,7 +639,7 @@ async def test_command_for_unknown_session_emits_unknown_session_error() -> None
     async with Harness.create() as harness:
         harness.command("sid-missing", Heartbeat(event_id="evt-hb"))
 
-        events = harness.events("sid-missing")
+        events = await harness.events("sid-missing")
         assert len(events) == 1
         error = events[0]
         assert isinstance(error, ErrorEvent)
@@ -618,11 +654,11 @@ async def test_command_for_closed_session_emits_unknown_session_error() -> None:
     async with Harness.create() as harness:
         await harness.open("sid-closed")
         assert (await harness.close("sid-closed")).ok is True
-        harness.events()
+        await harness.events()
 
         harness.command("sid-closed", Heartbeat(event_id="evt-late"))
 
-        events = harness.events("sid-closed")
+        events = await harness.events("sid-closed")
         assert [type(event) for event in events] == [ErrorEvent]
         assert events[0].code == "unknown_session"
         assert events[0].related_event_id == "evt-late"
@@ -631,14 +667,14 @@ async def test_command_for_closed_session_emits_unknown_session_error() -> None:
 async def test_admitted_command_reaches_the_session_runner_in_order() -> None:
     async with Harness.create() as harness:
         await harness.open("sid-runner")
-        harness.events()
+        await harness.events()
         harness.clock.advance(5.0)
 
         harness.command("sid-runner", Heartbeat(event_id="evt-1"))
         harness.command("sid-runner", Heartbeat(event_id="evt-2"))
         await _settle()
 
-        events = harness.events("sid-runner")
+        events = await harness.events("sid-runner")
         assert [type(event) for event in events] == [SessionHeartbeatAck, SessionHeartbeatAck]
         assert all(event.epoch == 0 for event in events)
         assert harness.session("sid-runner").lease.last_activity == 5.0
@@ -647,7 +683,7 @@ async def test_admitted_command_reaches_the_session_runner_in_order() -> None:
 async def test_append_audio_is_backpressured_by_pending_input_bytes() -> None:
     async with Harness.create(max_pending_input_bytes_per_session=8) as harness:
         await harness.open("sid-bytes")
-        harness.events()
+        await harness.events()
         session = harness.session("sid-bytes")
         submitted = harness.capture_submissions("sid-bytes")
         assert session.reserve_input_bytes(6, limit=8)
@@ -655,7 +691,7 @@ async def test_append_audio_is_backpressured_by_pending_input_bytes() -> None:
         harness.command("sid-bytes", AppendAudio(audio=b"pcm", event_id="evt-over"))
         harness.command("sid-bytes", AppendAudio(audio=b"pc", event_id="evt-fits"))
 
-        events = harness.events("sid-bytes")
+        events = await harness.events("sid-bytes")
         assert len(events) == 1
         assert isinstance(events[0], ErrorEvent)
         assert events[0].code == "input_backpressure"
@@ -670,7 +706,7 @@ async def test_append_audio_is_backpressured_by_pending_input_bytes() -> None:
 async def test_commit_holds_a_pending_turn_reservation_until_the_runner_dequeues_it() -> None:
     async with Harness.create(max_pending_turns_per_session=1) as harness:
         await harness.open("sid-turns")
-        harness.events()
+        await harness.events()
         session = harness.session("sid-turns")
         submitted = harness.capture_submissions("sid-turns")
 
@@ -679,7 +715,7 @@ async def test_commit_holds_a_pending_turn_reservation_until_the_runner_dequeues
 
         assert [command.event_id for command in submitted] == ["evt-commit-1"]
         assert session.pending_input_turns == 1
-        events = harness.events("sid-turns")
+        events = await harness.events("sid-turns")
         assert len(events) == 1
         assert isinstance(events[0], ErrorEvent)
         assert events[0].code == "input_backpressure"
@@ -689,7 +725,7 @@ async def test_commit_holds_a_pending_turn_reservation_until_the_runner_dequeues
         session.release_pending_turn()
         harness.command("sid-turns", Commit(event_id="evt-commit-3"))
         assert [command.event_id for command in submitted] == ["evt-commit-1", "evt-commit-3"]
-        assert harness.events("sid-turns") == []
+        assert await harness.events("sid-turns") == []
 
 
 def test_accepts_only_duplex_message_types() -> None:
@@ -704,7 +740,14 @@ def test_accepts_only_duplex_message_types() -> None:
     )
     try:
         config = DuplexSessionConfig()
-        assert manager.accepts(OpenDuplexSessionMessage(control_id="c", session_id="s", session_config=config))
+        assert manager.accepts(
+            OpenDuplexSessionMessage(
+                control_id="c",
+                session_id="s",
+                session_config=config,
+                output_buffer=DuplexOutputBuffer(max_bytes=2 * 1024 * 1024, max_events=512),
+            )
+        )
         assert manager.accepts(CloseDuplexSessionMessage(control_id="c", session_id="s"))
         assert manager.accepts(ResumeDuplexSessionMessage(control_id="c", session_id="s", expected_lease_generation=0))
         assert manager.accepts(TouchDuplexSessionMessage(control_id="c", session_id="s", activity="heartbeat"))
@@ -729,7 +772,7 @@ def test_accepts_only_duplex_message_types() -> None:
 async def test_close_releases_stage_resources_and_is_idempotent() -> None:
     async with Harness.create() as harness:
         await harness.open("sid-close")
-        harness.events()
+        await harness.events()
         session = harness.session("sid-close")
         request_id = stage0_request_id("sid-close")
         session.bind_stage_request(0, request_id, fence=session.fence)
@@ -750,7 +793,7 @@ async def test_close_releases_stage_resources_and_is_idempotent() -> None:
         assert harness.manager.runner_for_request_id(request_id) is None
         assert harness.stage_port.cleanup_calls == [([request_id], True), ([reserved_context.request_id], False)]
         assert harness.plugin.data_plane.closed_sessions == ["sid-close"]
-        closed = [event for event in harness.events("sid-close") if isinstance(event, SessionClosed)]
+        closed = [event for event in await harness.events("sid-close") if isinstance(event, SessionClosed)]
         assert len(closed) == 1
         assert closed[0].reason == "client_close"
         assert closed[0].is_terminal is True
@@ -760,7 +803,7 @@ async def test_close_releases_stage_resources_and_is_idempotent() -> None:
         assert again.lease_generation is None
         assert again.public_session is None
         assert harness.stage_port.cleanup_calls[2:] == []
-        assert harness.events() == []
+        assert await harness.events() == []
 
 
 async def test_close_retains_the_admission_slot_until_stage_cleanup_succeeds() -> None:
@@ -858,7 +901,7 @@ async def test_reaper_expires_idle_sessions_and_emits_session_expired() -> None:
     async with Harness.create(idle_ttl_s=2.0, disconnect_grace_s=1.0) as harness:
         await harness.open("sid-idle")
         await harness.open("sid-active")
-        harness.events()
+        await harness.events()
         harness.clock.advance(1.0)
         assert (await harness.touch("sid-active", DuplexLeaseActivity.HEARTBEAT.value)).ok is True
         harness.clock.advance(1.1)
@@ -868,7 +911,7 @@ async def test_reaper_expires_idle_sessions_and_emits_session_expired() -> None:
         assert harness.manager.get("sid-idle") is None
         assert harness.manager.get("sid-active") is not None
         assert harness.manager.active_count() == 1
-        expired = harness.events("sid-idle")
+        expired = await harness.events("sid-idle")
         assert len(expired) == 1
         assert isinstance(expired[0], SessionExpired)
         assert expired[0].reason == "idle_ttl_expired"
@@ -878,7 +921,7 @@ async def test_reaper_expires_idle_sessions_and_emits_session_expired() -> None:
             "session_id": "sid-idle",
             "reason": "idle_ttl_expired",
         }
-        assert harness.events("sid-active") == []
+        assert await harness.events("sid-active") == []
         assert harness.stage_port.cleanup_calls == [([stage0_request_id("sid-idle")], False)]
         assert harness.result_sink.empty()
 
@@ -888,7 +931,7 @@ async def test_reaper_expires_idle_sessions_and_emits_session_expired() -> None:
 async def test_reaper_expires_detached_sessions_after_the_disconnect_grace() -> None:
     async with Harness.create(idle_ttl_s=300.0, disconnect_grace_s=1.0) as harness:
         await harness.open("sid-detached")
-        harness.events()
+        await harness.events()
         assert (await harness.touch("sid-detached", DuplexLeaseActivity.DETACH.value)).ok is True
         harness.clock.advance(0.5)
         assert await harness.manager.reap_expired() == 0
@@ -897,7 +940,7 @@ async def test_reaper_expires_detached_sessions_after_the_disconnect_grace() ->
         assert await harness.manager.reap_expired() == 1
 
         assert harness.manager.get("sid-detached") is None
-        expired = harness.events("sid-detached")
+        expired = await harness.events("sid-detached")
         assert [type(event) for event in expired] == [SessionExpired]
         assert expired[0].reason == "disconnect_grace_expired"
         assert (await harness.resume("sid-detached", expected_lease_generation=0)).error_code == "unknown_session"
@@ -906,7 +949,7 @@ async def test_reaper_expires_detached_sessions_after_the_disconnect_grace() ->
 async def test_session_closed_is_emitted_only_after_stage_cleanup_so_a_reopen_is_admitted() -> None:
     async with Harness.create(max_sessions=1) as harness:
         await harness.open("sid-closing")
-        harness.events()
+        await harness.events()
         session = harness.session("sid-closing")
         session.bind_stage_request(0, stage0_request_id("sid-closing"), fence=session.fence)
         harness.stage_port.cleanup_gate = asyncio.Event()
@@ -915,11 +958,11 @@ async def test_session_closed_is_emitted_only_after_stage_cleanup_so_a_reopen_is
         await asyncio.wait_for(harness.stage_port.cleanup_started.wait(), timeout=1.0)
         # The runner is gone but the stage request is still being aborted:
         # no session.closed yet, so a client cannot race the admission slot.
-        assert harness.events("sid-closing") == []
+        assert await harness.events("sid-closing") == []
 
         harness.stage_port.cleanup_gate.set()
         assert (await close_task).ok is True
-        events = harness.events("sid-closing")
+        events = await harness.events("sid-closing")
         assert [type(event) for event in events] == [SessionClosed]
         assert events[0].reason == "client_close"
         assert (await harness.open("sid-replacement")).ok is True
@@ -928,7 +971,7 @@ async def test_session_closed_is_emitted_only_after_stage_cleanup_so_a_reopen_is
 async def test_wire_close_command_frees_the_slot_before_session_closed() -> None:
     async with Harness.create(max_sessions=1) as harness:
         await harness.open("sid-wire-close")
-        harness.events()
+        await harness.events()
         session = harness.session("sid-wire-close")
         request_id = stage0_request_id("sid-wire-close")
         session.bind_stage_request(0, request_id, fence=session.fence)
@@ -937,7 +980,7 @@ async def test_wire_close_command_frees_the_slot_before_session_closed() -> None
         harness.command("sid-wire-close", CloseSession(reason="client_close"))
         await asyncio.wait_for(harness.stage_port.cleanup_started.wait(), timeout=1.0)
         assert harness.manager.get("sid-wire-close") is None
-        assert harness.events("sid-wire-close") == []
+        assert await harness.events("sid-wire-close") == []
 
         harness.stage_port.cleanup_gate.set()
         for _ in range(50):
@@ -945,7 +988,7 @@ async def test_wire_close_command_frees_the_slot_before_session_closed() -> None
                 break
             await asyncio.sleep(0.01)
         assert harness.stage_port.cleanup_calls == [([request_id], True)]
-        events = harness.events("sid-wire-close")
+        events = await harness.events("sid-wire-close")
         assert [type(event) for event in events] == [SessionClosed]
         assert events[0].reason == "client_close"
         assert (await harness.open("sid-replacement")).ok is True
@@ -954,7 +997,7 @@ async def test_wire_close_command_frees_the_slot_before_session_closed() -> None
 async def test_wire_close_stops_the_worker_and_refuses_later_control_ops() -> None:
     async with Harness.create(max_sessions=1) as harness:
         await harness.open("sid-wire")
-        harness.events()
+        await harness.events()
         runner = harness.manager.runners["sid-wire"]
         worker = runner._worker
         assert worker is not None
@@ -975,7 +1018,7 @@ async def test_wire_close_stops_the_worker_and_refuses_later_control_ops() -> No
                 break
             await asyncio.sleep(0.01)
         assert worker.done(), "the mailbox worker must exit after a wire close"
-        events = harness.events("sid-wire")
+        events = await harness.events("sid-wire")
         assert [type(event) for event in events] == [ErrorEvent, SessionClosed] or [
             type(event) for event in events
         ] == [SessionClosed, ErrorEvent]
@@ -1003,7 +1046,7 @@ async def test_concurrent_opens_cannot_exceed_max_sessions() -> None:
 async def test_runtime_close_without_stage_requests_frees_the_slot() -> None:
     async with Harness.create(max_sessions=1) as harness:
         await harness.open("sid-runtime")
-        harness.events()
+        await harness.events()
         runner = harness.manager.runners["sid-runtime"]
         runner.session.release_all_requests()
 
@@ -1014,7 +1057,7 @@ async def test_runtime_close_without_stage_requests_frees_the_slot() -> None:
             await asyncio.sleep(0.01)
 
         assert harness.manager.get("sid-runtime") is None
-        assert [type(event) for event in harness.events("sid-runtime")] == [SessionClosed]
+        assert [type(event) for event in await harness.events("sid-runtime")] == [SessionClosed]
         assert (await harness.open("sid-next")).ok is True
 
 
@@ -1028,7 +1071,7 @@ async def test_close_of_a_forgotten_session_is_idempotent() -> None:
 async def test_append_bytes_are_reserved_at_admission_until_the_runner_dequeues() -> None:
     async with Harness.create(max_sessions=1, max_pending_input_bytes_per_session=8) as harness:
         await harness.open("sid-bytes")
-        harness.events()
+        await harness.events()
         session = harness.session("sid-bytes")
         runner = harness.manager.runners["sid-bytes"]
         # Park the worker so admitted commands stay in the mailbox.
@@ -1040,7 +1083,7 @@ async def test_append_bytes_are_reserved_at_admission_until_the_runner_dequeues(
         harness.command("sid-bytes", AppendAudio(audio=b"12345"))
         assert session.pending_input_bytes == 5
         harness.command("sid-bytes", AppendAudio(audio=b"12345"))
-        errors = [event for event in harness.events("sid-bytes") if isinstance(event, ErrorEvent)]
+        errors = [event for event in await harness.events("sid-bytes") if isinstance(event, ErrorEvent)]
         assert [error.code for error in errors] == ["input_backpressure"]
         gate.set()
 
@@ -1049,7 +1092,7 @@ async def test_append_admission_counts_audio_and_video_frame_bytes() -> None:
     """Manager reserves len(audio)+Σlen(frame); video bytes count toward the same limit."""
     async with Harness.create(max_sessions=1, max_pending_input_bytes_per_session=20) as harness:
         await harness.open("sid-av")
-        harness.events()
+        await harness.events()
         session = harness.session("sid-av")
         runner = harness.manager.runners["sid-av"]
         gate = asyncio.Event()
@@ -1072,7 +1115,7 @@ async def test_append_admission_counts_audio_and_video_frame_bytes() -> None:
             "sid-av",
             AppendAudio(audio=b"x" * 10, video_frames=("yyyyyyyyyy",), event_id="evt-over"),
         )
-        errors = [event for event in harness.events("sid-av") if isinstance(event, ErrorEvent)]
+        errors = [event for event in await harness.events("sid-av") if isinstance(event, ErrorEvent)]
         assert [error.code for error in errors] == ["input_backpressure"]
         assert errors[0].related_event_id == "evt-over"
         assert session.pending_input_bytes == expected
@@ -1082,7 +1125,7 @@ async def test_append_admission_counts_audio_and_video_frame_bytes() -> None:
 async def test_expired_session_retains_the_admission_slot_until_cleanup_succeeds() -> None:
     async with Harness.create(max_sessions=1, idle_ttl_s=1.0) as harness:
         await harness.open("sid-expired")
-        harness.events()
+        await harness.events()
         session = harness.session("sid-expired")
         request_id = stage0_request_id("sid-expired")
         session.bind_stage_request(0, request_id, fence=session.fence)
@@ -1093,14 +1136,14 @@ async def test_expired_session_retains_the_admission_slot_until_cleanup_succeeds
 
         assert harness.manager.get("sid-expired") is None
         assert session.state == DuplexSessionState.CLOSED
-        assert [type(event) for event in harness.events("sid-expired")] == [SessionExpired]
+        assert [type(event) for event in await harness.events("sid-expired")] == [SessionExpired]
         assert (await harness.open("sid-replacement")).error_code == "resource_exhausted"
 
         assert await harness.manager.reap_expired() == 1
 
         assert harness.stage_port.cleanup_calls == [([request_id], True)] * 2
         assert session.resource_request_ids() == []
-        assert harness.events("sid-expired") == []
+        assert await harness.events("sid-expired") == []
         assert (await harness.open("sid-replacement")).ok is True
 
 
@@ -1131,13 +1174,13 @@ async def test_expired_cleanup_failure_does_not_block_other_sessions() -> None:
     async with Harness.create(stage_port=OneStuckStagePort(), idle_ttl_s=1.0) as harness:
         await harness.open("sid-stuck")
         await harness.open("sid-fine")
-        harness.events()
+        await harness.events()
         harness.clock.advance(2.0)
 
         assert await harness.manager.reap_expired() == 1
 
         assert harness.manager.active_count() == 0
-        assert {type(event) for event in harness.events()} == {SessionExpired}
+        assert {type(event) for event in await harness.events()} == {SessionExpired}
         assert (await harness.open("sid-fine")).ok is True
         assert (await harness.open("sid-stuck")).error_code == "session_exists"
 
@@ -1150,7 +1193,7 @@ async def test_expired_cleanup_failure_does_not_block_other_sessions() -> None:
 async def test_close_sessions_for_request_ids_expires_the_owner_and_retries_cleanup() -> None:
     async with Harness.create(max_sessions=1) as harness:
         await harness.open("sid-request")
-        harness.events()
+        await harness.events()
         session = harness.session("sid-request")
         request_id = stage0_request_id("sid-request")
         session.bind_stage_request(0, request_id, fence=session.fence)
@@ -1162,7 +1205,7 @@ async def test_close_sessions_for_request_ids_expires_the_owner_and_retries_clea
         assert harness.manager.get("sid-request") is None
         assert session.lease.terminal_reason == "request_cleanup"
         await _settle()
-        expired = harness.events("sid-request")
+        expired = await harness.events("sid-request")
         assert [type(event) for event in expired] == [SessionExpired]
         assert expired[0].reason == "request_cleanup"
         assert session.state == DuplexSessionState.CLOSED
@@ -1256,7 +1299,12 @@ async def test_control_dispatch_is_ordered_per_session_without_blocking_other_se
         blocked_config = DuplexSessionConfig(model="fake-model", instructions="sid-blocked")
 
         manager.dispatch(
-            OpenDuplexSessionMessage(control_id="blocked-open", session_id="sid-blocked", session_config=blocked_config)
+            OpenDuplexSessionMessage(
+                control_id="blocked-open",
+                session_id="sid-blocked",
+                session_config=blocked_config,
+                output_buffer=DuplexOutputBuffer(max_bytes=2 * 1024 * 1024, max_events=512),
+            )
         )
         await asyncio.wait_for(plugin.runtime_config_started.wait(), timeout=1.0)
         manager.dispatch(
@@ -1271,6 +1319,7 @@ async def test_control_dispatch_is_ordered_per_session_without_blocking_other_se
                 control_id="independent-open",
                 session_id="sid-independent",
                 session_config=DuplexSessionConfig(model="fake-model"),
+                output_buffer=DuplexOutputBuffer(max_bytes=2 * 1024 * 1024, max_events=512),
             )
         )
         manager.dispatch(
@@ -1338,16 +1387,18 @@ async def test_shutdown_closes_every_runner_and_stops_dispatch_tasks() -> None:
     await harness.open("sid-a")
     await harness.open("sid-b")
     sessions = [harness.session("sid-a"), harness.session("sid-b")]
-    harness.events()
+    await harness.events()
 
     await harness.manager.shutdown()
 
     assert harness.manager.active_count() == 0
     assert harness.manager.runners == {}
+    assert harness.manager._outputs == {}
+    assert harness.manager._output_failed == set()
     assert all(session.state == DuplexSessionState.CLOSED for session in sessions)
     assert harness.manager.runner_for_request_id(stage0_request_id("sid-b")) is None
     # Shutdown is silent: no session.closed / session.expired is emitted.
-    assert harness.events() == []
+    assert await harness.events() == []
     assert harness.stage_port.cleanup_calls == []
 
 
@@ -1414,12 +1465,12 @@ async def test_expiry_emits_a_terminal_event_even_when_a_close_deferred_it() ->
         runner = harness.manager.runners["sid-deferred"]
         # What the runner does for a wire close: tear down, defer the terminal.
         await runner.close("client_close", emit_closed=False)
-        harness.events("sid-deferred")
+        await harness.events("sid-deferred")
 
         await runner.expire("request_cleanup", emit_expired=True)
 
         terminal = [
-            event for event in harness.events("sid-deferred") if isinstance(event, SessionClosed | SessionExpired)
+            event for event in await harness.events("sid-deferred") if isinstance(event, SessionClosed | SessionExpired)
         ]
         assert len(terminal) == 1, terminal
         assert isinstance(terminal[0], SessionExpired)
@@ -1466,7 +1517,7 @@ async def test_a_runtime_close_racing_a_manager_close_still_emits_one_terminal()
     """
     async with Harness.create(max_sessions=1) as harness:
         await harness.open("sid-race")
-        harness.events()
+        await harness.events()
         session = harness.session("sid-race")
         session.bind_stage_request(0, stage0_request_id("sid-race"), fence=session.fence)
         runner = harness.manager.runners["sid-race"]
@@ -1486,7 +1537,9 @@ async def test_a_runtime_close_racing_a_manager_close_still_emits_one_terminal()
 
         await harness.close("sid-race")
 
-        terminal = [event for event in harness.events("sid-race") if isinstance(event, SessionClosed | SessionExpired)]
+        terminal = [
+            event for event in await harness.events("sid-race") if isinstance(event, SessionClosed | SessionExpired)
+        ]
         assert len(terminal) == 1, terminal
 
 
@@ -1498,7 +1551,9 @@ async def test_a_terminal_event_is_emitted_once_even_if_expiry_runs_twice() -> N
         await runner.expire("idle_ttl_expired", emit_expired=True)
         await runner.expire("idle_ttl_expired", emit_expired=True)
 
-        terminal = [event for event in harness.events("sid-once") if isinstance(event, SessionClosed | SessionExpired)]
+        terminal = [
+            event for event in await harness.events("sid-once") if isinstance(event, SessionClosed | SessionExpired)
+        ]
         assert len(terminal) == 1, terminal
 
 
@@ -1590,12 +1645,12 @@ async def test_wire_close_of_an_idle_session_still_emits_session_closed() -> Non
     """
     async with Harness.create(max_sessions=1) as harness:
         await harness.open("sid-idle-close")
-        harness.events()
+        await harness.events()
 
         harness.command("sid-idle-close", CloseSession(reason="client_close"))
         events: list[DuplexEvent] = []
         for _ in range(200):
-            events.extend(harness.events("sid-idle-close"))
+            events.extend(await harness.events("sid-idle-close"))
             if events:
                 break
             await asyncio.sleep(0.01)
@@ -1725,11 +1780,13 @@ async def test_an_open_cancelled_after_admission_releases_the_runner_and_its_sta
                     control_id="open-cancelled",
                     session_id="sid-cancelled",
                     session_config=DuplexSessionConfig(model="fake-model"),
+                    output_buffer=DuplexOutputBuffer(max_bytes=2 * 1024 * 1024, max_events=512),
                 )
             )
         )
         await asyncio.wait_for(sink.entered.wait(), timeout=1.0)
         assert "sid-cancelled" in harness.manager.runners
+        assert "sid-cancelled" in harness.manager._outputs
         assert [context.request_id for context in harness.stage_port.ensure_calls] == [
             stage0_request_id("sid-cancelled")
         ]
@@ -1740,6 +1797,7 @@ async def test_an_open_cancelled_after_admission_releases_the_runner_and_its_sta
         harness.manager._result_sink = original_sink
 
         assert "sid-cancelled" not in harness.manager.runners
+        assert "sid-cancelled" not in harness.manager._outputs
         assert harness.manager.active_count() == 0
         assert harness.stage_port.cleanup_calls == [([stage0_request_id("sid-cancelled")], False)]
         assert stage0_request_id("sid-cancelled") not in harness.manager._request_index
@@ -1764,6 +1822,7 @@ async def test_an_open_cancelled_again_during_its_rollback_leaves_the_stage_clea
                     control_id="open-cancelled-twice",
                     session_id="sid-cancelled",
                     session_config=DuplexSessionConfig(model="fake-model"),
+                    output_buffer=DuplexOutputBuffer(max_bytes=2 * 1024 * 1024, max_events=512),
                 )
             )
         )
@@ -1779,6 +1838,7 @@ async def test_an_open_cancelled_again_during_its_rollback_leaves_the_stage_clea
 
         open_task.cancel()
         await asyncio.wait_for(shutdown_started.wait(), timeout=1.0)
+        assert "sid-cancelled" not in harness.manager._outputs
         open_task.cancel()  # lands inside the rollback, before the stage cleanup
         with pytest.raises(asyncio.CancelledError):
             await open_task
@@ -1875,3 +1935,127 @@ async def test_a_resume_replayed_under_its_control_id_reports_the_generation_it_
 
         other = await harness.resume("sid-replay", expected_lease_generation=0, control_id="rpc-b")
         assert other.ok is False and other.error_code == "session_resume_conflict"
+
+
+async def test_output_overflow_closes_only_its_session_and_releases_capacity(mocker) -> None:
+    async with Harness.create(max_sessions=2, max_pending_output_events_per_session=8) as harness:
+        assert (await harness.open("slow")).ok
+        assert (await harness.open("normal")).ok
+        await harness.events()
+        assert not harness.manager.runners["slow"].run.closing
+        assert not harness.manager.runners["normal"].run.closing
+        slow = harness.session("slow")
+        normal = harness.session("normal")
+        slow_response = await harness.begin_output_response("slow")
+        normal_response = normal.begin_response()
+        request_id = stage0_request_id("slow")
+        slow.bind_request(request_id)
+        slow.mark_audio_sent(1000)
+        slow.acknowledge_playback(500, 500)
+        close_stream = mocker.spy(harness.plugin.data_plane, "close_stream")
+
+        for _ in range(5):
+            harness.manager.runners["slow"].emit(
+                {
+                    "type": "response.output_audio.delta",
+                    "response_id": slow_response,
+                    "audio": "AAAAAA==",
+                    "format": "pcm16",
+                }
+            )
+        harness.manager.emit(normal, [AudioDelta(response_id=normal_response, delta="CCCC")])
+        await asyncio.wait_for(asyncio.gather(*tuple(harness.manager._dispatched_control_tasks)), timeout=2.0)
+
+        events = await harness.events()
+        failed = [event for event in events if event.session_id == "slow"]
+        # Queued transcript deltas remain ordered; stale audio and the failed
+        # response's undelivered completion markers are discarded.
+        assert [type(event) for event in failed] == [TranscriptDelta] * 4 + [ErrorEvent, ResponseDone, SessionClosed]
+        assert failed[-3].code == "output_backpressure"
+        done = failed[-2]
+        assert isinstance(done, ResponseDone)
+        assert done.response_id == slow_response and done.status == "failed"
+        metadata = done.response["metadata"]
+        assert isinstance(metadata, dict) and metadata["committed"] is False
+        assert failed[-1].reason == "output_backpressure"
+        assert slow.history == ()
+        assert slow.active_response_id is None
+        assert harness.stage_port.abort_calls == [[request_id]]
+        close_stream.assert_called_once_with(request_id)
+        delivered = [event for event in events if event.session_id == "normal"]
+        assert len(delivered) == 1 and isinstance(delivered[0], AudioDelta)
+        assert delivered[0].response_id == normal_response and delivered[0].delta == "CCCC"
+        assert harness.manager.get("slow") is None
+        assert harness.manager.get("normal") is normal
+        assert harness.manager.active_count() == 1
+        assert (await harness.open("replacement")).ok
+
+
+@pytest.mark.parametrize(
+    "capacity,end_before_emit,expected_status,closed",
+    [
+        (8, True, "failed", True),
+        (8, False, "failed", True),
+        (14, True, "completed", True),
+        (64, True, "completed", False),
+    ],
+    ids=["already-ended-overflow", "active-overflow", "overflow-after-ending", "sufficient-capacity"],
+)
+async def test_completion_projection_retains_exactly_one_response_done(
+    capacity: int, end_before_emit: bool, expected_status: str, closed: bool
+) -> None:
+    async with Harness.create(max_pending_output_events_per_session=capacity) as harness:
+        assert (await harness.open("completion")).ok
+        await harness.events()
+        response_id = await harness.begin_output_response("completion")
+        runner = harness.manager.runners["completion"]
+        for _ in range(4):
+            runner.emit(
+                {
+                    "type": "response.output_audio.delta",
+                    "response_id": response_id,
+                    "audio": "AAAAAA==",
+                    "format": "pcm16",
+                }
+            )
+        assert harness.outputs["completion"].pending_events == 8
+        runner.session.mark_audio_sent(1000)
+        runner.session.acknowledge_playback(500, 500)
+        if end_before_emit:
+            message = runner.session.end_response(commit_text=True)
+            assert message == {"role": "assistant", "content": "abcde"}
+            runner.session.register_history_item(f"item_{response_id}", message)
+            assert runner.session.history == (message,)
+
+        runner.emit(
+            {"type": "response.done", "response_id": response_id, "status": "completed", "committed": end_before_emit}
+        )
+        await asyncio.wait_for(asyncio.gather(*tuple(harness.manager._dispatched_control_tasks)), timeout=2.0)
+
+        events = await harness.events()
+        endings = [event for event in events if isinstance(event, ResponseDone)]
+        assert len(endings) == 1
+        assert endings[0].response_id == response_id and endings[0].status == expected_status
+        metadata = endings[0].response["metadata"]
+        assert isinstance(metadata, dict) and metadata["committed"] is (expected_status == "completed")
+        assert any(isinstance(event, SessionClosed) for event in events) is closed
+        if expected_status == "failed":
+            assert [type(event) for event in events] == [TranscriptDelta] * 4 + [
+                ErrorEvent,
+                ResponseDone,
+                SessionClosed,
+            ]
+            assert endings[0].response["status_details"] == {"type": "failed", "reason": "output_backpressure"}
+            assert runner.session.history == ()
+        else:
+            assert [type(event) for event in events[:8]] == [AudioDelta, TranscriptDelta] * 4
+            assert runner.session.history == ({"role": "assistant", "content": "abcde"},)
+            if closed:
+                # The ending was accepted; overflowing on rate_limits.updated
+                # must not replace it, emit a second ending or undo its history.
+                assert [type(event) for event in events[-3:]] == [ResponseDone, ErrorEvent, SessionClosed]
+                assert events[-3] is endings[0]
+            else:
+                assert not any(isinstance(event, ErrorEvent) for event in events)
+                assert events[-2] is endings[0]
+                assert events[-1].type == "rate_limits.updated"

tests/engine/duplex/test_session_runner.py

@@ -6,7 +6,7 @@
 The runner is driven exactly the way ``DuplexOrchestrator`` drives it: typed
 commands go through ``DuplexSessionManager.dispatch``, stage outputs are pushed
 with ``runner.on_stage_output`` and everything the session says is read back
-from the manager's output sink as typed events. The stage port is a recording
+from the session's output buffer as typed events. The stage port is a recording
 fake; the model plugin is the real MiniCPM-o 4.5 one so append planning and
 output projection are exercised end to end.
 """
@@ -38,7 +38,8 @@ from vllm_omni.engine.duplex.contracts import (
     DuplexStageSubmissionResult,
     duplex_resource_request_id,
 )
-from vllm_omni.engine.duplex.events import DuplexEvent
+from vllm_omni.engine.duplex.delivery import DuplexOutputBuffer
+from vllm_omni.engine.duplex.events import AudioDelta, DuplexEvent
 from vllm_omni.engine.duplex.messages import (
     CloseDuplexSessionMessage,
     DuplexControlResultMessage,
@@ -72,6 +73,8 @@ class RecordingStagePort(DuplexStagePort):
         self.submissions: list[DuplexStageSubmission] = []
         self.cleanups: list[tuple[list[str], bool]] = []
         self.aborts: list[list[str]] = []
+        self.abort_started = asyncio.Event()
+        self.abort_gate: asyncio.Event | None = None
         self.fail_submit: Exception | None = None
         #: When set, ``submit`` parks on it after signalling ``submit_started``.
         self.submit_gate: asyncio.Event | None = None
@@ -105,6 +108,9 @@ class RecordingStagePort(DuplexStagePort):
 
     async def abort_requests(self, request_ids: list[str]) -> None:
         self.aborts.append(list(request_ids))
+        self.abort_started.set()
+        if self.abort_gate is not None:
+            await self.abort_gate.wait()
 
 
 def _fake_encode_audio(audio: object, sample_rate_hz: int, response_format: str, speed: float | None) -> str | None:
@@ -120,6 +126,7 @@ class Harness:
     manager: DuplexSessionManager
     port: RecordingStagePort
     output: asyncio.Queue[Any]
+    output_buffer: DuplexOutputBuffer
     results: asyncio.Queue[Any]
     runner: DuplexSessionRunner
     events: list[DuplexEvent] = field(default_factory=list)
@@ -139,6 +146,12 @@ class Harness:
         collected: list[DuplexEvent] = []
         while True:
             drained = False
+            while self.output_buffer.pending_events:
+                event = await self.output_buffer.get()
+                assert event is not None
+                if self.output_buffer.is_valid(event):
+                    collected.append(event)
+                drained = True
             while not self.output.empty():
                 message = self.output.get_nowait()
                 if isinstance(message, DuplexSessionEventMessage):
@@ -215,12 +228,17 @@ async def open_harness(
     port = RecordingStagePort(stage_count=stage_count)
     output: asyncio.Queue[Any] = asyncio.Queue()
     results: asyncio.Queue[Any] = asyncio.Queue()
+    limits = runtime_config or DuplexSessionRuntimeConfig()
+    output_buffer = DuplexOutputBuffer(
+        max_bytes=limits.max_pending_output_bytes_per_session,
+        max_events=limits.max_pending_output_events_per_session,
+    )
     manager = DuplexSessionManager(
         plugin=plugin,
         stage_port=port,
         output_sink=output,
         result_sink=results,
-        runtime_config=runtime_config or DuplexSessionRuntimeConfig(),
+        runtime_config=limits,
         model_config=None,
         log_stats=log_stats,
         clock=clock,
@@ -232,10 +250,21 @@ async def open_harness(
         instructions="You are a concise assistant.",
         extra_body=body,
     )
-    await manager.handle(OpenDuplexSessionMessage(control_id="c-open", session_id=SESSION_ID, session_config=config))
+    await manager.handle(
+        OpenDuplexSessionMessage(
+            control_id="c-open", session_id=SESSION_ID, session_config=config, output_buffer=output_buffer
+        )
+    )
     result = await asyncio.wait_for(results.get(), timeout=2.0)
     assert isinstance(result, DuplexControlResultMessage) and result.ok, result
-    harness = Harness(manager=manager, port=port, output=output, results=results, runner=manager.runners[SESSION_ID])
+    harness = Harness(
+        manager=manager,
+        port=port,
+        output=output,
+        output_buffer=output_buffer,
+        results=results,
+        runner=manager.runners[SESSION_ID],
+    )
     await harness.settle()
     return harness
 
@@ -710,6 +739,8 @@ async def test_stale_epoch_output_is_dropped_after_barge_in() -> None:
 
 async def test_barge_in_aborts_draining_tts_as_well_as_the_active_request() -> None:
     h = await open_harness()
+    release_abort = asyncio.Event()
+    h.port.abort_gate = release_abort
     try:
         await h.run(append_audio())
         request_id = h.stage0_request_id()
@@ -730,7 +761,24 @@ async def test_barge_in_aborts_draining_tts_as_well_as_the_active_request() -> N
         h.session.request_resources[(1, "still-live")] = DuplexRequestResource(
             stage_id=1, request_id="still-live", fence=older, submitted=True
         )
-        events = await h.run(commands.BargeIn())
+        h.manager.emit(
+            h.session,
+            [
+                AudioDelta(response_id="resp-old", delta="AAAA"),
+                AudioDelta(response_id="resp-old", delta="BBBB"),
+                AudioDelta(response_id=h.session.active_response_id, delta="CCCC"),
+            ],
+        )
+        held = await h.output_buffer.get()
+        assert held is not None and h.output_buffer.is_valid(held)
+        h.submit(commands.BargeIn())
+        await asyncio.wait_for(h.port.abort_started.wait(), timeout=2.0)
+        # Both queued audio and an event held by the consumer become stale
+        # before the stage abort can yield, including the older draining turn.
+        assert not h.output_buffer.is_valid(held)
+        assert h.output_buffer.pending_events == 0
+        release_abort.set()
+        events = await h.settle()
         assert h.port.aborts == [[request_id, "duplex-drain-tts"]]
         assert not h.session.is_draining_request("duplex-drain-tts")
         done_ids = [
@@ -744,6 +792,36 @@ async def test_barge_in_aborts_draining_tts_as_well_as_the_active_request() -> N
         assert (0, request_id) not in h.session.request_resources
         assert (1, "still-live") in h.session.request_resources
     finally:
+        release_abort.set()
+        await close_harness(h)
+
+
+@pytest.mark.asyncio
+async def test_close_invalidates_draining_audio_before_stage_abort_finishes() -> None:
+    h = await open_harness()
+    release_abort = asyncio.Event()
+    h.port.abort_gate = release_abort
+    try:
+        h.session.bind_draining_request("duplex-drain-tts", "resp-old")
+        h.manager.emit(
+            h.session,
+            [AudioDelta(response_id="resp-old", delta="AAAA"), AudioDelta(response_id="resp-old", delta="BBBB")],
+        )
+        held = await h.output_buffer.get()
+        assert held is not None and h.output_buffer.is_valid(held)
+        h.submit(commands.CloseSession())
+        await asyncio.wait_for(h.port.abort_started.wait(), timeout=2.0)
+        # Session close suppresses audio.cancelled notifications; it must still
+        # invalidate the pending output without relying on those later events.
+        assert not h.output_buffer.is_valid(held)
+        assert h.output_buffer.pending_events == 0
+        assert h.port.aborts == [["duplex-drain-tts"]]
+        release_abort.set()
+        events = await h.settle()
+        assert "session.closed" in types(events)
+        assert h.manager.active_count() == 0
+    finally:
+        release_abort.set()
         await close_harness(h)
 
 

tests/engine/duplex/test_session_runner_nemotron_voicechat.py

@@ -21,6 +21,7 @@ from tests.engine.duplex.test_session_runner import (
 from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig
 from vllm_omni.engine.duplex import commands
 from vllm_omni.engine.duplex.config import DuplexSessionConfig
+from vllm_omni.engine.duplex.delivery import DuplexOutputBuffer
 from vllm_omni.engine.duplex.messages import OpenDuplexSessionMessage
 from vllm_omni.engine.duplex.plugin import load_duplex_plugin
 from vllm_omni.engine.duplex.session.manager import DuplexSessionManager
@@ -57,20 +58,36 @@ async def harness(monkeypatch):
     port = RecordingStagePort(stage_count=3)
     output: asyncio.Queue = asyncio.Queue()
     results: asyncio.Queue = asyncio.Queue()
+    limits = DuplexSessionRuntimeConfig()
+    output_buffer = DuplexOutputBuffer(
+        max_bytes=limits.max_pending_output_bytes_per_session,
+        max_events=limits.max_pending_output_events_per_session,
+    )
     manager = DuplexSessionManager(
         plugin=plugin,
         stage_port=port,
         output_sink=output,
         result_sink=results,
-        runtime_config=DuplexSessionRuntimeConfig(),
+        runtime_config=limits,
         model_config=SimpleNamespace(hf_config=SimpleNamespace(stt_cfg={}), max_model_len=8192),
     )
     try:
         config = DuplexSessionConfig(model="nemotron", instructions="hi", extra_body={"auto_response": True})
-        await manager.handle(OpenDuplexSessionMessage(control_id="open", session_id=SESSION_ID, session_config=config))
+        await manager.handle(
+            OpenDuplexSessionMessage(
+                control_id="open", session_id=SESSION_ID, session_config=config, output_buffer=output_buffer
+            )
+        )
         result = await asyncio.wait_for(results.get(), timeout=2)
         assert result.ok, result
-        h = Harness(manager, port, output, results, manager.runners[SESSION_ID])
+        h = Harness(
+            manager=manager,
+            port=port,
+            output=output,
+            output_buffer=output_buffer,
+            results=results,
+            runner=manager.runners[SESSION_ID],
+        )
         await h.settle()
         yield h
     finally:

tests/engine/duplex/test_session_runner_personaplex.py

@@ -34,6 +34,7 @@ from tests.engine.duplex.test_session_runner import (
 from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig
 from vllm_omni.engine.duplex import commands
 from vllm_omni.engine.duplex.config import DuplexSessionConfig
+from vllm_omni.engine.duplex.delivery import DuplexOutputBuffer
 from vllm_omni.engine.duplex.messages import DuplexControlResultMessage, OpenDuplexSessionMessage
 from vllm_omni.engine.duplex.session.manager import DuplexSessionManager
 from vllm_omni.model_executor.models.personaplex.duplex import stage0
@@ -55,12 +56,17 @@ async def open_personaplex_harness(*, extra_body: dict[str, object] | None = Non
     port = RecordingStagePort(stage_count=2)
     output: asyncio.Queue[Any] = asyncio.Queue()
     results: asyncio.Queue[Any] = asyncio.Queue()
+    limits = DuplexSessionRuntimeConfig()
+    output_buffer = DuplexOutputBuffer(
+        max_bytes=limits.max_pending_output_bytes_per_session,
+        max_events=limits.max_pending_output_events_per_session,
+    )
     manager = DuplexSessionManager(
         plugin=plugin,
         stage_port=port,
         output_sink=output,
         result_sink=results,
-        runtime_config=DuplexSessionRuntimeConfig(),
+        runtime_config=limits,
         model_config=SimpleNamespace(model="/models/personaplex-7b-v1"),
     )
     config = DuplexSessionConfig(
@@ -70,10 +76,21 @@ async def open_personaplex_harness(*, extra_body: dict[str, object] | None = Non
         voice="NATF2.pt",
         extra_body=dict(extra_body or {}),
     )
-    await manager.handle(OpenDuplexSessionMessage(control_id="c-open", session_id=SESSION_ID, session_config=config))
+    await manager.handle(
+        OpenDuplexSessionMessage(
+            control_id="c-open", session_id=SESSION_ID, session_config=config, output_buffer=output_buffer
+        )
+    )
     result = await asyncio.wait_for(results.get(), timeout=2.0)
     assert isinstance(result, DuplexControlResultMessage) and result.ok, result
-    harness = Harness(manager=manager, port=port, output=output, results=results, runner=manager.runners[SESSION_ID])
+    harness = Harness(
+        manager=manager,
+        port=port,
+        output=output,
+        output_buffer=output_buffer,
+        results=results,
+        runner=manager.runners[SESSION_ID],
+    )
     await harness.settle()
     return harness
 

tests/engine/test_async_omni_engine_stage_init.py

@@ -71,6 +71,34 @@ def test_stage_runtime_env_accepts_typed_runtime_config(monkeypatch):
     assert env_key not in os.environ
 
 
+@pytest.mark.parametrize("typed", [False, True])
+def test_build_ray_diffusion_config_preserves_explicit_stage_env(monkeypatch, typed):
+    from vllm_omni.engine import stage_init_utils as init_mod
+
+    env = {"CUSTOM_PLUGIN_SETTING": "stage", "OMP_NUM_THREADS": 2}
+    runtime_cfg = OmniStageRuntimeConfig(env=env) if typed else {"env": env}
+    metadata = types.SimpleNamespace(
+        runtime_cfg=runtime_cfg, stage_id=1, cfg_kv_collect_func=None, default_sampling_params=None
+    )
+    config = types.SimpleNamespace(
+        distributed_executor_backend="ray", parallel_config=types.SimpleNamespace(world_size=2)
+    )
+    monkeypatch.setattr(init_mod, "build_engine_args_dict", lambda *args: {})
+    monkeypatch.setattr(init_mod, "OmniDiffusionConfig", lambda **kwargs: config)
+    monkeypatch.setattr(
+        init_mod,
+        "current_omni_platform",
+        types.SimpleNamespace(device_control_env_var=None, get_device_count=lambda: 0),
+    )
+
+    result = init_mod.build_diffusion_config("model", {}, metadata)
+
+    assert result is config
+    assert result.ray_worker_env == {"CUSTOM_PLUGIN_SETTING": "stage", "OMP_NUM_THREADS": "2"}
+    assert result.num_gpus == 2
+    assert env["OMP_NUM_THREADS"] == 2
+
+
 @pytest.mark.parametrize(
     "invalid_key,invalid_value",
     [("INVALID=ENV", "value"), ("INVALID\0ENV", "value"), ("INVALID_ENV", "value\0")],
@@ -156,7 +184,7 @@ def _make_llm_metadata(
         default_sampling_params=types.SimpleNamespace(name=f"sp-{stage_id}-{replica_id}"),
         custom_process_input_func=None,
         engine_input_source=[] if stage_id == 0 else [stage_id - 1],
-        engine_output_type="token_ids",
+        engine_output_type="text",
         replica_id=replica_id,
         is_comprehension=is_comprehension,
     )
@@ -712,7 +740,11 @@ def test_launch_diffusion_stage_replica_preserves_configured_max_num_seqs(monkey
         assert kwargs["visible_devices"] == "1"
         return [42]
 
-    od_config = types.SimpleNamespace(max_num_seqs=4, parallel_config=types.SimpleNamespace(world_size=1))
+    od_config = types.SimpleNamespace(
+        max_num_seqs=4,
+        distributed_executor_backend="mp",
+        parallel_config=types.SimpleNamespace(world_size=1),
+    )
     monkeypatch.setattr(startup_mod, "build_diffusion_config", lambda *args: od_config)
     monkeypatch.setattr(startup_mod, "acquire_device_locks", acquire)
     monkeypatch.setattr(
@@ -806,6 +838,7 @@ def test_launch_diffusion_stage_replica_preserves_step_execution_max_num_seqs(mo
     od_config = types.SimpleNamespace(
         max_num_seqs=8,
         step_execution=True,
+        distributed_executor_backend="mp",
         parallel_config=types.SimpleNamespace(world_size=1),
     )
     monkeypatch.setattr(startup_mod, "build_diffusion_config", lambda *args: od_config)
@@ -1297,6 +1330,32 @@ def test_stage_runtime_multi_api_failure_shuts_down_before_exceptional_context_e
     ]
 
 
+@pytest.mark.parametrize("engine_output_type", ["", "tokens", "text+speech", "AUDIO"])
+def test_stage_runtime_rejects_invalid_output_type_during_initialization(mocker, engine_output_type):
+    runtime = StageRuntime(
+        stage_configs=[types.SimpleNamespace()],
+        model="dummy-model",
+        config_path="dummy-config",
+        stage_init_timeout=1,
+        async_chunk=False,
+    )
+    cfg = types.SimpleNamespace(model_config=types.SimpleNamespace(skip_tokenizer_init=True))
+    stage_plan = _make_llm_plan(0, stage_id=0, vllm_config=cfg)
+    stage_plan.replicas[0].metadata.engine_output_type = engine_output_type
+    stage_client = mocker.Mock()
+    mocker.patch.object(runtime, "_prepare_stage_plans", return_value=[stage_plan])
+    initialize_replicas = mocker.patch.object(runtime, "_initialize_stage_replicas", return_value={0: [stage_client]})
+
+    # Keep output-processor construction real: no request is needed to parse
+    # the type, even though the engine replica is already initialized here.
+    with pytest.raises(ValueError, match="Unknown modality"):
+        runtime.initialize()
+
+    initialize_replicas.assert_called_once()
+    stage_client.shutdown.assert_called_once_with()
+    assert runtime.stage_pools == []
+
+
 def test_stage_runtime_passes_log_stats_to_output_processor(monkeypatch):
     import vllm_omni.engine.stage_runtime as runtime_mod
 

tests/engine/test_duplex_omni_engine.py

@@ -14,6 +14,7 @@ import pytest
 
 from vllm_omni.engine.duplex.commands import Heartbeat
 from vllm_omni.engine.duplex.config import DuplexCapabilities, DuplexSessionConfig
+from vllm_omni.engine.duplex.delivery import DuplexOutputBuffer
 from vllm_omni.engine.duplex.messages import (
     CloseDuplexSessionMessage,
     DuplexControlResultMessage,
@@ -72,14 +73,16 @@ async def test_open_session_runs_the_correlated_rpc_with_the_typed_config() -> N
     result = _ok("open", capabilities=DuplexCapabilities(), public_session={"id": "sid"}, lease_generation=0)
     engine = _engine(result)
     config = DuplexSessionConfig(model="m")
+    output_buffer = DuplexOutputBuffer(max_bytes=2 * 1024 * 1024, max_events=512)
 
-    returned = await engine.open_session_async("sid", config, timeout=3.0)
+    returned = await engine.open_session_async("sid", config, output_buffer=output_buffer, timeout=3.0)
 
     assert returned is result
     ((key, message, timeout),) = engine.rpc_client.calls
     assert key == ("duplex", message.control_id)
     assert isinstance(message, OpenDuplexSessionMessage)
     assert message.session_id == "sid" and message.session_config is config
+    assert message.output_buffer is output_buffer
     assert timeout == 3.0
     # A momentarily full request queue is backpressure, not a failed open: the
     # control RPC blocks rather than raising an untyped queue.Full at the caller.
@@ -116,7 +119,9 @@ async def test_failed_control_result_maps_to_a_typed_session_error() -> None:
         )
     )
     with pytest.raises(DuplexSessionError) as excinfo:
-        await engine.open_session_async("sid", DuplexSessionConfig())
+        await engine.open_session_async(
+            "sid", DuplexSessionConfig(), output_buffer=DuplexOutputBuffer(max_bytes=2 * 1024 * 1024, max_events=512)
+        )
     assert excinfo.value.code == "resource_exhausted"
     assert excinfo.value.retryable is True
     assert excinfo.value.session_id == "sid"
@@ -133,7 +138,9 @@ async def test_engine_error_and_timeout_map_to_session_errors() -> None:
 
     engine = _engine(TimeoutError("duplex open timed out for session sid"))
     with pytest.raises(DuplexSessionError) as excinfo:
-        await engine.open_session_async("sid", DuplexSessionConfig())
+        await engine.open_session_async(
+            "sid", DuplexSessionConfig(), output_buffer=DuplexOutputBuffer(max_bytes=2 * 1024 * 1024, max_events=512)
+        )
     assert excinfo.value.code == "timeout" and excinfo.value.retryable is True
 
 

tests/engine/test_duplex_orchestrator.py

@@ -24,6 +24,7 @@ from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig
 from vllm_omni.engine.duplex import commands
 from vllm_omni.engine.duplex.config import DuplexSessionConfig, DuplexSessionState
 from vllm_omni.engine.duplex.contracts import DuplexFence, duplex_resource_request_id
+from vllm_omni.engine.duplex.delivery import DuplexOutputBuffer
 from vllm_omni.engine.duplex.messages import (
     CloseDuplexSessionMessage,
     DuplexControlResultMessage,
@@ -78,10 +79,15 @@ def _build(
     return orchestrator, clients, rpc_q, output_q
 
 
-def _open_message(extra_body: dict[str, object] | None = None) -> OpenDuplexSessionMessage:
+def _open_message(
+    extra_body: dict[str, object] | None = None,
+    *,
+    output_buffer: DuplexOutputBuffer | None = None,
+) -> OpenDuplexSessionMessage:
     return OpenDuplexSessionMessage(
         control_id=f"open-{SESSION_ID}",
         session_id=SESSION_ID,
+        output_buffer=output_buffer or DuplexOutputBuffer(max_bytes=2 * 1024 * 1024, max_events=512),
         session_config=DuplexSessionConfig(
             model="openbmb/MiniCPM-o-4_5",
             modalities=["text"],
@@ -369,12 +375,13 @@ async def test_forwarded_stage_requests_are_bound_and_barge_in_aborts_them() ->
 @pytest.mark.asyncio
 async def test_session_owned_outputs_reach_the_runner_and_never_the_client_queue() -> None:
     orchestrator, _, rpc_q, output_q = _build(stages=2)
-    await _open(orchestrator, rpc_q)
+    output_buffer = DuplexOutputBuffer(max_bytes=2 * 1024 * 1024, max_events=512)
+    await _open(orchestrator, rpc_q, output_buffer=output_buffer)
     request_id = _stage0_request_id()
     await _submit(orchestrator, _append_audio())
     request_state = orchestrator.request_states[request_id]
-    while not output_q.empty():
-        output_q.get_nowait()
+    while output_buffer.pending_events:
+        await output_buffer.get()
 
     consumed = await orchestrator._intercept_stage_output(1, 0, _tts_output(request_id), request_state, None, None)
     await _settle(orchestrator)
@@ -382,8 +389,9 @@ async def test_session_owned_outputs_reach_the_runner_and_never_the_client_queue
     assert consumed is True
     session = orchestrator.session_manager.get(SESSION_ID)
     assert session is not None and session.active_response_id is not None
-    types = [message.event.type for message in [output_q.get_nowait() for _ in range(output_q.qsize())]]
+    types = [(await output_buffer.get()).type for _ in range(output_buffer.pending_events)]
     assert "response.created" in types and "response.output_audio.delta" in types
+    assert output_q.empty()
 
     orphan = DuplexOrchestratorRequestState(
         request_id="duplex-s.b3RoZXI.e.0.r.stage0",
@@ -401,7 +409,8 @@ async def test_session_owned_outputs_reach_the_runner_and_never_the_client_queue
 @pytest.mark.asyncio
 async def test_forward_failure_closes_the_owning_session() -> None:
     orchestrator, clients, rpc_q, output_q = _build(stages=2)
-    await _open(orchestrator, rpc_q)
+    output_buffer = DuplexOutputBuffer(max_bytes=2 * 1024 * 1024, max_events=512)
+    await _open(orchestrator, rpc_q, output_buffer=output_buffer)
     request_id = _stage0_request_id()
     await _submit(orchestrator, _append_audio())
     request_state = orchestrator.request_states[request_id]
@@ -417,7 +426,8 @@ async def test_forward_failure_closes_the_owning_session() -> None:
     assert SESSION_ID not in orchestrator.session_manager.runners
     assert session.state == DuplexSessionState.CLOSED
     assert orchestrator.session_manager.active_count() == 0
-    types = [message.event.type for message in [output_q.get_nowait() for _ in range(output_q.qsize())]]
+    types = [(await output_buffer.get()).type for _ in range(output_buffer.pending_events)]
+    types.extend(output_q.get_nowait().event.type for _ in range(output_q.qsize()))
     # The stage failure is reported before the session expires, never after.
     assert types.index("error") < types.index("session.expired")
     assert types[-1] == "session.expired"
@@ -428,7 +438,8 @@ async def test_aura_forward_failure_keeps_the_session() -> None:
     from vllm_omni.model_executor.models.aura_omni.duplex.plugin import AuraDuplexPlugin
 
     orchestrator, clients, rpc_q, output_q = _build(stages=2, plugin=AuraDuplexPlugin(_encode_audio))
-    await _open(orchestrator, rpc_q)
+    output_buffer = DuplexOutputBuffer(max_bytes=2 * 1024 * 1024, max_events=512)
+    await _open(orchestrator, rpc_q, output_buffer=output_buffer)
     request_id = next(iter(orchestrator.request_states))
     await _submit(orchestrator, _append_audio())
     request_state = orchestrator.request_states[request_id]
@@ -442,7 +453,8 @@ async def test_aura_forward_failure_keeps_the_session() -> None:
     assert request_id not in orchestrator.request_states
     assert SESSION_ID in orchestrator.session_manager.runners
     assert session.state != DuplexSessionState.CLOSED
-    types = [message.event.type for message in [output_q.get_nowait() for _ in range(output_q.qsize())]]
+    types = [(await output_buffer.get()).type for _ in range(output_buffer.pending_events)]
+    types.extend(output_q.get_nowait().event.type for _ in range(output_q.qsize()))
     assert "error" in types
     assert "session.expired" not in types
     await orchestrator.session_manager.shutdown()
@@ -604,6 +616,7 @@ async def test_turn_plugin_processes_multimodal_prompt_before_stage_submission(m
             control_id="open-qwen",
             session_id=SESSION_ID,
             session_config=DuplexSessionConfig(model="qwen", modalities=["text", "audio"]),
+            output_buffer=DuplexOutputBuffer(max_bytes=2 * 1024 * 1024, max_events=512),
         )
     )
     assert (await rpc_q.get()).ok

tests/engine/test_output_modality.py

@@ -1,3 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Unit tests for Phase 1 foundation types (RFC #1601).
 
 Note: Uses importlib to load modules directly, bypassing the vllm_omni
@@ -7,10 +9,14 @@ package __init__ which requires the vllm base package.
 import importlib.util
 import sys
 from pathlib import Path
+from typing import TYPE_CHECKING
 
 import pytest
 import torch
 
+if TYPE_CHECKING:
+    from vllm_omni.outputs.mm_outputs import MultimodalPayload
+
 # ── Load modules without triggering vllm_omni.__init__ ─────────────
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -19,6 +25,7 @@ _OUTPUTS_DIR = Path(__file__).resolve().parents[2] / "vllm_omni" / "outputs"
 
 def _load_module(name: str, filepath: Path):
     spec = importlib.util.spec_from_file_location(name, filepath)
+    assert spec is not None and spec.loader is not None
     mod = importlib.util.module_from_spec(spec)
     sys.modules[name] = mod
     spec.loader.exec_module(mod)
@@ -37,45 +44,155 @@ _mm_mod = _load_module(
     "vllm_omni.outputs.mm_outputs",
     _OUTPUTS_DIR / "mm_outputs.py",
 )
+_accumulation_mod = _load_module(
+    "vllm_omni.outputs.multimodal_accumulation",
+    _OUTPUTS_DIR / "multimodal_accumulation.py",
+)
 
 OutputModality = _om_mod.OutputModality
 TensorAccumulationStrategy = _om_mod.TensorAccumulationStrategy
 get_accumulation_strategy = _om_mod.get_accumulation_strategy
 register_key_accumulation_strategy = _om_mod.register_key_accumulation_strategy
-MultimodalPayload = _mm_mod.MultimodalPayload
+if not TYPE_CHECKING:
+    MultimodalPayload = _mm_mod.MultimodalPayload
 MultimodalCompletionOutput = _mm_mod.MultimodalCompletionOutput
 
 
-def test_output_modality_parsing_and_flags():
-    """Test OutputModality enum: from_string, aliases, compounds, properties, and accumulation strategy."""
-    # Defaults
-    assert OutputModality.from_string(None) == OutputModality.TEXT
-    assert OutputModality.from_string("") == OutputModality.TEXT
+@pytest.mark.parametrize(
+    ("value", "expected"),
+    [
+        (None, OutputModality.TEXT),
+        ("text", OutputModality.TEXT),
+        ("image", OutputModality.IMAGE),
+        ("audio", OutputModality.AUDIO),
+        ("latent", OutputModality.LATENT),
+        ("token_ids", OutputModality.TOKEN_IDS),
+        ("text+token_ids", OutputModality.TEXT | OutputModality.TOKEN_IDS),
+        ("token_ids,latent", OutputModality.TOKEN_IDS | OutputModality.LATENT),
+        ("text+image", OutputModality.TEXT | OutputModality.IMAGE),
+        ("text,audio", OutputModality.TEXT | OutputModality.AUDIO),
+        ("latent+image,audio", OutputModality.LATENT | OutputModality.IMAGE | OutputModality.AUDIO),
+        ("image+text", OutputModality.TEXT | OutputModality.IMAGE),
+        ("audio+audio", OutputModality.AUDIO),
+    ],
+)
+def test_output_modality_parses_canonical_names(value, expected):
+    assert OutputModality.from_string(value) == expected
 
-    # Direct names and case insensitivity
-    assert OutputModality.from_string("image") == OutputModality.IMAGE
-    assert OutputModality.from_string("Audio") == OutputModality.AUDIO
 
-    # Aliases
-    assert OutputModality.from_string("speech") == OutputModality.AUDIO
-    assert OutputModality.from_string("latents") == OutputModality.LATENT
-    assert OutputModality.from_string("pixel_values") == OutputModality.IMAGE
+@pytest.mark.parametrize(
+    "value",
+    [
+        "speech",
+        "images",
+        "latents",
+        "wav",
+        "waveform",
+        "pixel_values",
+        "pixels",
+        "tokens",
+        "text+speech",
+        "token_ids+tokens",
+        "TOKEN_IDS",
+        "token_ids ",
+        "Audio",
+        "TEXT",
+        "text+Image",
+        " audio",
+        "audio ",
+        "text +audio",
+        "text, audio",
+        "text\taudio",
+        "audio\n",
+        "",
+        " ",
+        "+audio",
+        "audio,",
+        "text++audio",
+        "text,,audio",
+        "text+,audio",
+        "video",
+        "text+unknown",
+    ],
+)
+def test_output_modality_rejects_noncanonical_names(value):
+    with pytest.raises(ValueError, match="Unknown modality"):
+        OutputModality.from_string(value)
 
-    # Compound
+
+def test_output_modality_flags_and_accumulation_strategy():
     compound = OutputModality.from_string("text+image")
     assert compound.has_text and compound.has_multimodal
 
     # Flag properties
     assert OutputModality.TEXT.has_text and not OutputModality.TEXT.has_multimodal
     assert OutputModality.IMAGE.has_multimodal and not OutputModality.IMAGE.has_text
+    assert OutputModality.TOKEN_IDS.has_multimodal and not OutputModality.TOKEN_IDS.has_text
+    assert OutputModality.TOKEN_IDS != OutputModality.LATENT
 
     # Accumulation strategy
     assert get_accumulation_strategy(OutputModality.AUDIO) == TensorAccumulationStrategy.CONCAT_LAST
     assert get_accumulation_strategy(OutputModality.IMAGE) == TensorAccumulationStrategy.CONCAT_DIM0
+    assert get_accumulation_strategy(OutputModality.TOKEN_IDS) == TensorAccumulationStrategy.CONCAT_DIM0
 
-    # Unknown raises
-    with pytest.raises(ValueError, match="Unknown modality"):
-        OutputModality.from_string("video")
+
+@pytest.mark.parametrize("producer_key", [None, "model_outputs", "hidden", "token_ids"])
+def test_token_id_payload_preserves_discrete_primary_output(producer_key):
+    token_ids = torch.tensor([[12, 85], [206, 7]], dtype=torch.long)
+    raw = token_ids if producer_key is None else {producer_key: token_ids}
+
+    payload = MultimodalPayload.from_raw(raw, "token_ids")
+
+    assert payload is not None
+    assert set(payload) == {"token_ids"}
+    torch.testing.assert_close(payload["token_ids"], token_ids)
+
+
+def test_token_id_payload_keeps_explicit_latents_and_image_ids_separate():
+    token_ids = torch.tensor([12, 85, 206], dtype=torch.long)
+    latent = torch.tensor([[0.1, 0.2], [0.3, 0.4]])
+    prior_image = torch.tensor([4, 5], dtype=torch.long)
+
+    payload = MultimodalPayload.from_raw(
+        {"model_outputs": token_ids, "latent": latent, "ids.prior_image": prior_image},
+        "token_ids",
+    )
+
+    assert payload is not None
+    assert set(payload) == {"token_ids", "latent", "ids.prior_image"}
+    torch.testing.assert_close(payload["token_ids"], token_ids)
+    torch.testing.assert_close(payload["latent"], latent)
+    torch.testing.assert_close(payload["ids.prior_image"], prior_image)
+
+
+def test_token_id_payload_consolidates_dim0_and_respects_key_overrides():
+    reference_key = "__test_token_ids__.reference"
+    register_key_accumulation_strategy(reference_key, TensorAccumulationStrategy.REPLACE)
+    first_ids = torch.tensor([[1, 2], [3, 4]], dtype=torch.long)
+    last_ids = torch.tensor([[5, 6]], dtype=torch.long)
+    first_ref = torch.tensor([[7, 8]], dtype=torch.long)
+    last_ref = torch.tensor([[9, 10]], dtype=torch.long)
+    payload = MultimodalPayload()
+    payload.tensors["token_ids"] = [first_ids, last_ids]
+    payload.tensors[reference_key] = [first_ref, last_ref]
+
+    payload.consolidate_tensors(OutputModality.TOKEN_IDS)
+
+    torch.testing.assert_close(payload["token_ids"], torch.cat([first_ids, last_ids], dim=0))
+    torch.testing.assert_close(payload[reference_key], last_ref)
+
+
+def test_delta_drain_retains_token_ids_and_latents():
+    token_ids = torch.tensor([12, 85], dtype=torch.long)
+    latent = torch.tensor([[0.1, 0.2]])
+    payload = MultimodalPayload.from_dict({"token_ids": token_ids, "latent": latent, "audio": torch.ones(4)})
+    assert payload is not None
+
+    _accumulation_mod.drain_delta_payload(payload)
+
+    assert set(payload) == {"token_ids", "latent"}
+    torch.testing.assert_close(payload["token_ids"], token_ids)
+    torch.testing.assert_close(payload["latent"], latent)
 
 
 def test_multimodal_payload_and_completion_output():
@@ -107,7 +224,7 @@ def test_multimodal_payload_and_completion_output():
 def test_output_modality_printed_examples(capsys):
     """Printed examples for output modality types."""
     print("\n=== OutputModality Parsing ===")
-    for s in [None, "", "image", "Audio", "speech", "latents", "pixel_values", "text+image"]:
+    for s in [None, "text", "image", "audio", "latent", "text+image", "text,audio"]:
         print(f"  from_string({s!r:20s}) -> {OutputModality.from_string(s)}")
 
     print("\n=== Flag Properties ===")

tests/engine/test_output_processor.py

@@ -3,6 +3,7 @@
 """Regression tests for OmniRequestState multimodal DELTA drain and consolidation guard."""
 
 from dataclasses import dataclass
+from enum import Flag
 from types import SimpleNamespace
 from unittest.mock import MagicMock
 
@@ -81,6 +82,26 @@ def test_output_modality_name_uses_string_value():
     assert str(OutputModalityNames.AUDIO) == "audio"
 
 
+@pytest.mark.parametrize(
+    ("name", "expected"),
+    [
+        ("TOKEN_IDS", "token_ids"),
+        ("TEXT|TOKEN_IDS", "token_ids"),
+        ("LATENT|TOKEN_IDS", "latent"),
+        ("IMAGE|TOKEN_IDS", "image"),
+        ("AUDIO|TOKEN_IDS", "audio"),
+    ],
+)
+def test_modality_type_string_handles_reloaded_token_id_enum(name, expected):
+    # A separate Flag class reproduces enum identity changes after module reload.
+    ReloadedOutputModality = Flag("ReloadedOutputModality", ["TEXT", "IMAGE", "AUDIO", "LATENT", "TOKEN_IDS"])
+    value = ReloadedOutputModality(0)
+    for member in name.split("|"):
+        value |= ReloadedOutputModality[member]
+
+    assert output_processor._modality_to_type_string(value) == expected
+
+
 def test_init_empty_dict():
     """Ensure mm_accumulated is initially empty."""
     assert _make_state(RequestOutputKind.CUMULATIVE).mm_accumulated == {}
@@ -574,6 +595,92 @@ def _make_no_detok_state(output_kind: RequestOutputKind):
     return OmniRequestState(**_NO_DETOK_STATE_KWARGS, output_kind=output_kind)
 
 
+def test_token_id_delta_retains_ids_and_latents_until_final_consolidation(mocker):
+    state = _make_state(RequestOutputKind.DELTA)
+    state.detokenizer = mocker.MagicMock(
+        output_token_ids=[10],
+        get_next_output_text=mocker.MagicMock(return_value=""),
+        num_output_tokens=mocker.MagicMock(return_value=1),
+    )
+    first_ids = torch.tensor([[1, 2]], dtype=torch.long)
+    last_ids = torch.tensor([[3, 4]], dtype=torch.long)
+    first_latent = torch.tensor([[0.1, 0.2, 0.3]])
+    last_latent = torch.tensor([[0.4, 0.5, 0.6]])
+    state.add_multimodal_tensor({"model_outputs": first_ids, "latent": first_latent}, mm_type="token_ids")
+
+    first = state.make_request_output([10], None, None, None)
+
+    assert first is not None and not isinstance(first, PoolingRequestOutput)
+    first_completion = first.outputs[0]
+    assert set(state.mm_accumulated) == {"token_ids", "latent"}
+    torch.testing.assert_close(first_completion.multimodal_output["token_ids"], first_ids)
+    torch.testing.assert_close(first_completion.multimodal_output["latent"], first_latent)
+    assert first_completion.cumulative_token_ids == [10]
+
+    state.detokenizer.output_token_ids.append(11)
+    state.add_multimodal_tensor({"model_outputs": last_ids, "latent": last_latent}, mm_type="token_ids")
+    final = state.make_request_output([11], None, FinishReason.STOP, None)
+
+    assert final is not None and not isinstance(final, PoolingRequestOutput)
+    completion = final.outputs[0]
+    assert set(state.mm_accumulated) == {"token_ids", "latent"}
+    torch.testing.assert_close(completion.multimodal_output["token_ids"], torch.cat([first_ids, last_ids], dim=0))
+    torch.testing.assert_close(completion.multimodal_output["latent"], torch.cat([first_latent, last_latent], dim=0))
+    # Completion tokens remain on their own cumulative path, independent of
+    # tensor payload IDs; the earlier completion keeps its original snapshot.
+    assert completion.cumulative_token_ids == [10, 11]
+    assert first_completion.cumulative_token_ids == [10]
+
+
+@pytest.mark.parametrize(
+    ("engine_output_type", "expected_key", "expected_values"),
+    [
+        (None, "text", [[0, 1, 2], [3, 4, 5]]),
+        ("text", "text", [[0, 1, 2], [3, 4, 5]]),
+        ("image", "image", [[0, 1, 2], [3, 4, 5]]),
+        ("latent", "latent", [[0, 1, 2], [3, 4, 5]]),
+        ("token_ids", "token_ids", [[0, 1, 2], [3, 4, 5]]),
+        ("audio", "audio", [[0, 1, 2, 3, 4, 5]]),
+        ("text+image", "image", [[0, 1, 2], [3, 4, 5]]),
+        ("text,latent", "latent", [[0, 1, 2], [3, 4, 5]]),
+        ("text+audio", "audio", [[0, 1, 2, 3, 4, 5]]),
+        ("text+token_ids", "token_ids", [[0, 1, 2], [3, 4, 5]]),
+        ("latent,token_ids", "latent", [[0, 1, 2], [3, 4, 5]]),
+        ("image+token_ids", "image", [[0, 1, 2], [3, 4, 5]]),
+        ("audio+token_ids", "audio", [[0, 1, 2, 3, 4, 5]]),
+    ],
+)
+def test_engine_output_type_controls_payload_key_and_concatenation(engine_output_type, expected_key, expected_values):
+    """Canonical config values retain output naming and tensor merge behavior."""
+    processor = MultimodalOutputProcessor(
+        tokenizer=None,
+        log_stats=False,
+        engine_core_output_type=engine_output_type,
+    )
+    state = _make_no_detok_state(RequestOutputKind.FINAL_ONLY)
+    processor.request_states[state.request_id] = state
+    processor.external_req_ids[state.external_req_id].append(state.request_id)
+
+    first = OmniEngineCoreOutput(
+        request_id=state.request_id,
+        new_token_ids=[10],
+        multimodal_output={"model_outputs": torch.tensor([[0, 1, 2]])},
+    )
+    assert not processor.process_outputs([first]).request_outputs
+
+    last = OmniEngineCoreOutput(
+        request_id=state.request_id,
+        new_token_ids=[11],
+        finish_reason=FinishReason.STOP,
+        multimodal_output={"model_outputs": torch.tensor([[3, 4, 5]])},
+    )
+    result = processor.process_outputs([last])
+    assert len(result.request_outputs) == 1
+    payload = result.request_outputs[0].outputs[0].multimodal_output
+    assert set(payload) == {expected_key}
+    torch.testing.assert_close(payload[expected_key], torch.tensor(expected_values))
+
+
 def test_no_detokenizer_completion_output():
     """_new_completion_output works when detokenizer is None (generation stages)."""
     s = _make_no_detok_state(RequestOutputKind.CUMULATIVE)

tests/engine/test_single_stage_mode.py

@@ -1219,7 +1219,11 @@ class TestSingleStageReplicaInitialization:
         runtime._init_visible_devices_baseline = "0"
 
         mocker.patch.object(runtime_mod, "inject_kv_stage_info")
-        od_config = SimpleNamespace(max_num_seqs=4, parallel_config=SimpleNamespace(world_size=1))
+        od_config = SimpleNamespace(
+            max_num_seqs=4,
+            parallel_config=SimpleNamespace(world_size=1),
+            distributed_executor_backend="mp",
+        )
         mocker.patch("vllm_omni.engine.stage_engine_startup.build_diffusion_config", return_value=od_config)
         mock_register = mocker.patch(
             "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master",
@@ -1352,7 +1356,11 @@ class TestSingleStageReplicaInitialization:
         runtime._init_visible_devices_baseline = "0"
 
         mocker.patch.object(runtime_mod, "inject_kv_stage_info")
-        od_config = SimpleNamespace(max_num_seqs=None, parallel_config=SimpleNamespace(world_size=1))
+        od_config = SimpleNamespace(
+            max_num_seqs=None,
+            parallel_config=SimpleNamespace(world_size=1),
+            distributed_executor_backend="mp",
+        )
         mocker.patch("vllm_omni.engine.stage_engine_startup.build_diffusion_config", return_value=od_config)
         mocker.patch(
             "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master",

tests/engine/test_stage_engine_args.py

@@ -226,7 +226,7 @@ def _engine_arg_inputs(tmp_path: Path) -> tuple[PipelineConfig, DeployConfig, st
                 execution_type=StageExecutionType.LLM_AR,
                 requires_multimodal_data=True,
                 hf_config_name="thinker_config",
-                engine_output_type="hidden_states",
+                engine_output_type="latent",
                 model_arch="Qwen3OmniMoeForConditionalGeneration",
                 model_subdir="ar-model",
                 tokenizer_subdir="ar-tokenizer",
@@ -238,7 +238,7 @@ def _engine_arg_inputs(tmp_path: Path) -> tuple[PipelineConfig, DeployConfig, st
                 model_stage="talker",
                 execution_type=StageExecutionType.LLM_GENERATION,
                 input_sources=(0,),
-                engine_output_type="audio_tokens",
+                engine_output_type="latent",
             ),
             StagePipelineConfig(
                 stage_id=2,

tests/engine/test_stage_metadata.py

@@ -1,5 +1,9 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 import copy
 import operator
+from dataclasses import replace
 
 import pytest
 from vllm.sampling_params import SamplingParams
@@ -33,7 +37,7 @@ def _metadata_inputs() -> tuple[PipelineConfig, DeployConfig]:
                 execution_type=StageExecutionType.LLM_AR,
                 owns_tokenizer=True,
                 requires_multimodal_data=True,
-                engine_output_type="token_ids",
+                engine_output_type="text",
                 custom_process_input_func="operator.add",
                 prompt_expand_func="operator.mul",
             ),
@@ -42,7 +46,7 @@ def _metadata_inputs() -> tuple[PipelineConfig, DeployConfig]:
                 model_stage="talker",
                 execution_type=StageExecutionType.LLM_GENERATION,
                 input_sources=(0,),
-                engine_output_type="audio_tokens",
+                engine_output_type="latent",
                 custom_process_input_func="operator.sub",
                 sync_process_input_func="operator.floordiv",
             ),
@@ -160,6 +164,26 @@ def test_extract_stage_metadata_preserves_legacy_one_argument_api():
     assert metadata.custom_process_input_func is operator.add
 
 
+@pytest.mark.parametrize("engine_output_type", [None, "", "latent", "token_ids", "text+token_ids", "text+audio"])
+def test_extract_stage_metadata_preserves_declared_engine_output_type(engine_output_type):
+    pipeline, deploy = _metadata_inputs()
+    pipeline = replace(
+        pipeline,
+        stages=(replace(pipeline.stages[0], engine_output_type=engine_output_type), *pipeline.stages[1:]),
+    )
+    omni_config = VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=copy.deepcopy(deploy))
+    legacy_config = merge_pipeline_deploy(pipeline, deploy)[0].to_omegaconf()
+
+    # An explicit empty value must reach the output processor, where it is
+    # rejected during startup, instead of silently becoming the None default.
+    assert legacy_config.engine_args.get("engine_output_type") == engine_output_type
+    assert extract_legacy_stage_metadata(legacy_config).engine_output_type == engine_output_type
+    assert (
+        extract_stage_metadata_from_omni_stage_config(omni_config.stage_by_id(0)).engine_output_type
+        == engine_output_type
+    )
+
+
 def test_extract_stage_metadata_defaults_missing_engine_input_source():
     pipeline, deploy = _metadata_inputs()
     legacy_config = merge_pipeline_deploy(

tests/entrypoints/duplex/test_duplex_serving.py

@@ -18,6 +18,7 @@ from fastapi import WebSocketDisconnect
 from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig
 from vllm_omni.engine.duplex import commands
 from vllm_omni.engine.duplex.config import DuplexCapabilities
+from vllm_omni.engine.duplex.delivery import DuplexOutputBuffer
 from vllm_omni.engine.duplex.events import AudioDelta, DuplexEvent, SessionClosed, SessionCreated
 from vllm_omni.engine.duplex.messages import DuplexSessionError
 from vllm_omni.entrypoints.duplex.realtime_input import RealtimeEnvelope, parse_resume_request
@@ -94,10 +95,13 @@ class FakeHandle:
         self.closed = False
         self.close_reasons: list[str] = []
         self.commands: list[commands.DuplexCommand] = []
-        self._outbox: asyncio.Queue[DuplexEvent] = asyncio.Queue()
+        self._outbox = DuplexOutputBuffer(max_bytes=2 * 1024 * 1024, max_events=512)
 
     def deliver(self, event: DuplexEvent) -> None:
-        self._outbox.put_nowait(event)
+        self._outbox.put(event)
+
+    def output_guard(self, event: DuplexEvent):
+        return self._outbox.guard(event)
 
     async def submit(self, command: commands.DuplexCommand) -> None:
         if self.closed:
@@ -107,6 +111,10 @@ class FakeHandle:
     async def events(self):
         while True:
             event = await self._outbox.get()
+            if event is None:
+                return
+            if not self._outbox.is_valid(event):
+                continue
             yield event
             if event.is_terminal:
                 return

tests/entrypoints/openai/test_duplex_session_attachment.py

@@ -9,6 +9,8 @@ import json
 
 import pytest
 
+from vllm_omni.engine.duplex.delivery import DuplexOutputBuffer
+from vllm_omni.engine.duplex.events import AudioDelta, ResponseDone
 from vllm_omni.entrypoints.duplex.session_attachment import (
     DuplexEventJournal,
     DuplexJournalGapError,
@@ -22,6 +24,36 @@ from vllm_omni.entrypoints.duplex.session_attachment import (
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("journal", [False, True])
+async def test_audio_invalidated_while_waiting_for_send_is_not_sequenced_or_journaled(mocker, journal: bool):
+    registry = DuplexSessionAttachmentRegistry(replay_ttl_s=60.0, replay_max_bytes_per_session=4096)
+    output = DuplexOutputBuffer(max_bytes=4096, max_events=8)
+    sent: list[dict[str, object]] = []
+    accepted = mocker.Mock()
+
+    await registry.create("s", send=mocker.AsyncMock(side_effect=sent.append), close=mocker.AsyncMock())
+    audio = AudioDelta(session_id="s", response_id="r", epoch=0, delta="AAAA")
+    output.put(audio)
+    assert await output.get() is audio
+    state = registry._sessions["s"]
+    async with state.outbound_lock:
+        delivery = asyncio.create_task(
+            registry.send_event(
+                "s", audio.to_realtime(), journal=journal, on_accepted=accepted, event_guard=lambda: output.guard(audio)
+            )
+        )
+        await asyncio.sleep(0)
+        assert not delivery.done()
+        output.invalidate("r", through_epoch=0)
+    assert await delivery is None
+    accepted.assert_not_called()
+    done = ResponseDone(response_id="r", response={"status": "cancelled"})
+    await registry.send_event("s", done.to_realtime())
+    assert [(entry["type"], entry["server_event_seq"]) for entry in sent] == [("response.done", 1)]
+    assert [entry.payload["event_id"] for entry in state.journal.replay_after(0)] == [done.event_id]
+
+
 class _Clock:
     def __init__(self) -> None:
         self.value = 0.0

tests/entrypoints/openai_api/test_image_server.py

@@ -438,6 +438,90 @@ def async_omni_stage_configs_only_client():
     return TestClient(app)
 
 
+@pytest.fixture
+def mammoth_moda2_test_client():
+    """MammothModa2-style pipeline: AR (llm) + generation-LLM DiT stage.
+
+    The DiT stage is ``stage_type == "llm"`` with ``final_output_type ==
+    "image"`` (no classical diffusion stage) — the legacy topology
+    MammothModa2 served before #7134 migrated it to the shared diffusion
+    runtime. Regression guard for the 503 (stage discovery) and the T2I
+    envelope / max_tokens wiring on that compat path.
+    """
+    from fastapi import FastAPI
+
+    from vllm_omni.entrypoints.async_omni import AsyncOmni
+    from vllm_omni.entrypoints.openai.api_server import router
+    from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
+
+    class FakeAsyncOmniClass(AsyncOmni):
+        def __init__(self):
+            stage_configs = [
+                SimpleNamespace(
+                    stage_type="llm",
+                    is_comprehension=False,
+                    engine_args=SimpleNamespace(model_arch="MammothModa2ForConditionalGeneration"),
+                ),
+                SimpleNamespace(
+                    stage_type="llm",
+                    is_comprehension=False,
+                    final_output=True,
+                    final_output_type="image",
+                    engine_args=SimpleNamespace(model_arch="MammothModa2ForConditionalGeneration"),
+                ),
+            ]
+            default_sampling_params_list = [
+                SamplingParams(temperature=1.0, max_tokens=16),
+                OmniDiffusionSamplingParams(),
+            ]
+            self.engine = SimpleNamespace(
+                stage_configs=stage_configs,
+                default_sampling_params_list=default_sampling_params_list,
+                # Real AsyncOmni engines always carry stage_clients (read by
+                # get_diffusion_od_config); edits-endpoint tests walk through
+                # helpers that touch it before any stage-type gate.
+                stage_clients=[],
+            )
+            self.default_sampling_params_list = default_sampling_params_list
+            self.captured_sampling_params_list = None
+            self.captured_prompt = None
+            self._images = [Image.new("RGB", (64, 64), color="green")]
+
+        async def generate(self, prompt, request_id, sampling_params=None, sampling_params_list=None, **kwargs):
+            if sampling_params_list is not None:
+                self.captured_sampling_params_list = sampling_params_list
+            else:
+                self.captured_sampling_params_list = [sampling_params]
+            self.captured_prompt = prompt
+            images = [img.copy() for img in self._images]
+            yield MockGenerationResult(images)
+
+        def __class_getitem__(cls, item):
+            return cls
+
+    app = FastAPI()
+    app.include_router(router)
+
+    engine = FakeAsyncOmniClass()
+    chat_handler = object.__new__(OmniOpenAIServingChat)
+    chat_handler.engine_client = engine
+    chat_handler._diffusion_engine = None
+    app.state.openai_serving_chat = chat_handler
+    app.state.engine_client = engine
+    app.state.stage_configs = [
+        SimpleNamespace(stage_type="llm"),
+        SimpleNamespace(stage_type="llm", final_output=True, final_output_type="image"),
+    ]
+    app.state.openai_serving_models = _DiffusionServingModels(
+        [BaseModelPath(name="Mammoth/MammothModa2-Preview", model_path="Mammoth/MammothModa2-Preview")]
+    )
+    app.state.args = Namespace(
+        default_sampling_params='{"1": {"num_inference_steps":50, "generator_device":"cpu"}}',
+        max_generated_image_size=1024 * 1024,
+    )
+    return TestClient(app)
+
+
 @pytest.fixture
 def streaming_image_edit_client():
     """Create a multi-stage client whose engine yields AR text before image output."""
@@ -2418,3 +2502,190 @@ def test_image_edits_omitted_bot_task_stop_tokens_match_prompt_default(
         f"{ar_params.stop_token_ids} -- this is the full ratio range, meaning "
         "resolve_stop_token_ids disagreed with build_prompt_tokens's default."
     )
+
+
+def test_generate_images_mammoth_moda2_llm_dit_pipeline_accepted(mammoth_moda2_test_client):
+    """Regression: /v1/images/generations must not 503 on generation-LLM DiT stages.
+
+    A generation-LLM DiT stage (``stage_type == "llm"`` with
+    ``final_output_type == "image"``, MammothModa2's topology before #7134):
+    stage discovery previously required a classical diffusion stage and
+    returned 503 (issue #7199).
+    """
+    response = mammoth_moda2_test_client.post(
+        "/v1/images/generations",
+        json={
+            "prompt": "a mammoth in the snow",
+            "n": 1,
+            "size": "1024x1024",
+            "seed": 7,
+        },
+    )
+    assert response.status_code == 200, response.text
+    data = response.json()
+    assert len(data["data"]) == 1
+
+
+def test_generate_images_mammoth_moda2_llm_dit_rejects_multiple_outputs(mammoth_moda2_test_client):
+    """Regression: n > 1 must be rejected on the LLM-typed DiT path.
+
+    The LLM engine path for generation-LLM DiT stages produces exactly one
+    image per request; accepting ``n > 1`` used to silently drop the extra
+    outputs (review on PR #7293). The request must fail fast with 400 instead.
+    """
+    response = mammoth_moda2_test_client.post(
+        "/v1/images/generations",
+        json={
+            "prompt": "a mammoth in the snow",
+            "n": 4,
+            "size": "1024x1024",
+            "seed": 7,
+        },
+    )
+    assert response.status_code == 400, response.text
+    assert "exactly one image" in response.text
+
+
+def test_generate_images_mammoth_moda2_t2i_envelope_and_max_tokens(mammoth_moda2_test_client):
+    """Regression: the registry T2I builder drives the AR prompt + max_tokens.
+
+    The AR stage must receive the <|image start|>W*H<|image token|> envelope,
+    the structural additional_information, and max_tokens sized from the AR
+    grid (ar_height * (ar_width + 1) + 1) instead of SamplingParams' default.
+    """
+    response = mammoth_moda2_test_client.post(
+        "/v1/images/generations",
+        json={
+            "prompt": "a mammoth in the snow",
+            "n": 1,
+            "size": "1024x1024",
+            "seed": 7,
+        },
+    )
+    assert response.status_code == 200, response.text
+
+    engine = mammoth_moda2_test_client.app.state.engine_client
+    captured_prompt = engine.captured_prompt
+    assert captured_prompt["prompt"].endswith("<|image start|>64*64<|image token|>")
+    addi = captured_prompt["additional_information"]
+    assert addi["omni_task"] == ["t2i"]
+    assert addi["ar_width"] == [64]
+    assert addi["ar_height"] == [64]
+
+    captured = engine.captured_sampling_params_list
+    assert captured is not None
+    assert len(captured) == 2
+    assert captured[0].max_tokens == 64 * 65 + 1
+    assert captured[0].seed == 7
+
+
+def test_generate_images_mammoth_moda2_extra_args_routed_to_both_stages(mammoth_moda2_test_client):
+    """Regression: declared extra_body params reach AR and DiT extra_args.
+
+    text_guidance_scale / cfg_range / num_inference_steps are declared by the
+    MammothModa2 registry spec; the DiT stage (llm-typed) reads them from
+    sampling_params.extra_args, so the serving layer must route them there.
+    """
+    response = mammoth_moda2_test_client.post(
+        "/v1/images/generations",
+        json={
+            "prompt": "a mammoth in the snow",
+            "n": 1,
+            "size": "1024x1024",
+            "seed": 7,
+            "text_guidance_scale": 4.5,
+            "cfg_range": [0.0, 0.8],
+            "num_inference_steps": 30,
+        },
+    )
+    assert response.status_code == 200, response.text
+
+    engine = mammoth_moda2_test_client.app.state.engine_client
+    captured = engine.captured_sampling_params_list
+    assert captured is not None
+    assert len(captured) == 2
+    for stage_params in captured:
+        extra_args = getattr(stage_params, "extra_args", None) or {}
+        assert extra_args.get("text_guidance_scale") == 4.5
+        assert extra_args.get("cfg_range") == [0.0, 0.8]
+        assert extra_args.get("num_inference_steps") == 30
+
+
+def test_generate_images_mammoth_moda2_migrated_topology_keeps_t2i_builder(mammoth_moda2_test_client):
+    """Regression: the T2I builder must stay active after #7134's migration.
+
+    Once MammothModa2's DiT stage is migrated to the shared diffusion runtime
+    (``stage_type == "diffusion"``, see #7134), the pipeline contains a
+    classical diffusion stage while stage 0 still needs the AR envelope/grid
+    metadata. An earlier guard skipped the registry T2I builder whenever any
+    diffusion stage was present, which would silently drop the
+    <|image start|>W*H<|image token|> envelope and additional_information
+    after the migration. Dispatch must be by resolved model class, not by
+    pipeline topology.
+    """
+    engine = mammoth_moda2_test_client.app.state.engine_client
+    # Simulate the post-#7134 topology in place: stage 1 becomes a classical
+    # diffusion stage (final_output metadata moves onto it unchanged).
+    engine.engine.stage_configs[1] = SimpleNamespace(
+        stage_type="diffusion",
+        is_comprehension=False,
+        final_output=True,
+        final_output_type="image",
+        engine_args=SimpleNamespace(model_arch="MammothModa2ForConditionalGeneration"),
+    )
+    mammoth_moda2_test_client.app.state.stage_configs = list(engine.engine.stage_configs)
+
+    response = mammoth_moda2_test_client.post(
+        "/v1/images/generations",
+        json={
+            "prompt": "a mammoth in the snow",
+            "n": 1,
+            "size": "1024x1024",
+            "seed": 7,
+        },
+    )
+    assert response.status_code == 200, response.text
+
+    captured_prompt = engine.captured_prompt
+    assert captured_prompt["prompt"].endswith("<|image start|>64*64<|image token|>"), (
+        "the registry T2I builder must still wrap the AR prompt when the "
+        "pipeline contains a classical diffusion stage (post-#7134 topology)"
+    )
+    addi = captured_prompt["additional_information"]
+    assert addi["omni_task"] == ["t2i"]
+    assert addi["ar_width"] == [64]
+    assert addi["ar_height"] == [64]
+
+    captured = engine.captured_sampling_params_list
+    assert captured is not None
+    assert captured[0].max_tokens == 64 * 65 + 1
+
+
+def test_image_edits_mammoth_moda2_llm_dit_pipeline_rejected(mammoth_moda2_test_client):
+    """Regression: /v1/images/edits stays diffusion-only for now.
+
+    MammothModa2 is admitted to /v1/images/generations via the LLM-typed DiT
+    stage, but image edits (img2img) were never validated for it. Keep the
+    pre-existing diffusion-stage gate on the edits endpoint so the widened
+    generations acceptance does not silently admit it here.
+    """
+    img = make_test_image_bytes((64, 64))
+    response = mammoth_moda2_test_client.post(
+        "/v1/images/edits",
+        files=[("image", img)],
+        data={"prompt": "make it neon", "size": "512x512"},
+    )
+    assert response.status_code == 503, response.text
+    assert "No diffusion stage found" in response.json()["detail"]
+
+
+def test_is_image_generation_stage_matrix():
+    """Unit-test the stage classifier used by the image serving endpoints."""
+    from vllm_omni.entrypoints.openai.utils import is_image_generation_stage
+
+    assert is_image_generation_stage(SimpleNamespace(stage_type="diffusion"))
+    assert is_image_generation_stage(SimpleNamespace(stage_type="llm", final_output=True, final_output_type="image"))
+    assert is_image_generation_stage({"stage_type": "llm", "final_output": True, "final_output_type": "images"})
+    assert not is_image_generation_stage(SimpleNamespace(stage_type="llm", final_output=True, final_output_type="text"))
+    assert not is_image_generation_stage(SimpleNamespace(stage_type="llm", final_output=False))
+    assert not is_image_generation_stage(SimpleNamespace(stage_type="llm", final_output=True, final_output_type=None))

tests/entrypoints/test_duplex_omni.py

@@ -15,6 +15,7 @@ from tests.entrypoints.test_omni_entrypoints import FakeAsyncOmniEngine
 from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig
 from vllm_omni.engine.duplex import commands
 from vllm_omni.engine.duplex.config import DuplexCapabilities, DuplexSessionConfig
+from vllm_omni.engine.duplex.delivery import DuplexOutputBuffer
 from vllm_omni.engine.duplex.events import (
     AudioDelta,
     DuplexEvent,
@@ -41,6 +42,7 @@ class FakeDuplexEngine(FakeAsyncOmniEngine):
         self.duplex_session_config = DuplexSessionRuntimeConfig(max_sessions=2)
         self.duplex_capabilities = DuplexCapabilities(supports_session_resume=True)
         self.opened: list[tuple[str, DuplexSessionConfig]] = []
+        self.outputs: dict[str, DuplexOutputBuffer] = {}
         self.closed: list[tuple[str, str]] = []
         self.resumed: list[tuple[str, int]] = []
         #: The control id of every resume RPC, in order (a replay repeats one).
@@ -64,10 +66,11 @@ class FakeDuplexEngine(FakeAsyncOmniEngine):
     def _result(self, operation: str, session_id: str, **fields: Any) -> DuplexControlResultMessage:
         return DuplexControlResultMessage(control_id="c", operation=operation, session_id=session_id, ok=True, **fields)
 
-    async def open_session_async(self, session_id, session_config, *, timeout=None):
+    async def open_session_async(self, session_id, session_config, *, output_buffer, timeout=None):
         self.opened.append((session_id, session_config))
         if self.open_error is not None:
             raise self.open_error
+        self.outputs[session_id] = output_buffer
         return self._result(
             "open",
             session_id,
@@ -119,7 +122,11 @@ class FakeDuplexEngine(FakeAsyncOmniEngine):
     def emit(self, session_id: str, event: DuplexEvent) -> None:
         # The session manager binds the session identity to every event it emits.
         event = replace(event, session_id=session_id)
-        self.output_q.put(DuplexSessionEventMessage(session_id=session_id, event=event))
+        if isinstance(event, SessionClosed):
+            self.outputs.pop(session_id, None)
+            self.output_q.put(DuplexSessionEventMessage(session_id=session_id, event=event))
+        elif session_id in self.outputs:
+            self.outputs[session_id].put(event)
 
 
 def _make_omni(monkeypatch: pytest.MonkeyPatch) -> tuple[DuplexOmni, FakeDuplexEngine]:
@@ -152,6 +159,7 @@ async def test_open_session_allocates_the_id_and_ignores_client_ids(monkeypatch)
         assert handle.session_id != "mine"
         ((session_id, config),) = engine.opened
         assert session_id == handle.session_id
+        assert engine.outputs[session_id] is handle._outbox
         assert isinstance(config, DuplexSessionConfig)
         assert config.model == "dummy-model" and config.instructions == "hi"
         assert handle.capabilities is engine.duplex_capabilities
@@ -169,6 +177,32 @@ async def test_open_session_allocates_the_id_and_ignores_client_ids(monkeypatch)
         omni.shutdown()
 
 
+@pytest.mark.asyncio
+async def test_events_skips_audio_invalidated_after_dequeue(monkeypatch) -> None:
+    omni, engine = _make_omni(monkeypatch)
+    try:
+        handle = await omni.open_session()
+        output = engine.outputs[handle.session_id]
+        original_get = output.get
+
+        async def get_then_invalidate() -> DuplexEvent | None:
+            event = await original_get()
+            if isinstance(event, AudioDelta) and event.response_id == "cancelled":
+                output.invalidate("cancelled", through_epoch=0)
+            return event
+
+        monkeypatch.setattr(output, "get", get_then_invalidate)
+        engine.emit(handle.session_id, AudioDelta(response_id="cancelled", epoch=0, delta="old"))
+        engine.emit(handle.session_id, AudioDelta(response_id="current", epoch=1, delta="new"))
+
+        events = await _collect(handle, 1)
+
+        assert isinstance(events[0], AudioDelta)
+        assert events[0].response_id == "current" and events[0].delta == "new"
+    finally:
+        omni.shutdown()
+
+
 @pytest.mark.asyncio
 async def test_open_failure_drops_the_pending_handle_and_raises_the_engine_code(monkeypatch) -> None:
     omni, engine = _make_omni(monkeypatch)

tests/examples/offline_inference/test_text_to_image_cli_forwarding.py

@@ -0,0 +1,140 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""
+Offline inference regression tests: explicit CLI parallel-knob forwarding.
+
+Covers the argparse defaults -> ``build_parallel_knob_kwargs`` -> ``main()``'s
+``omni_kwargs`` assembly in examples/offline_inference/text_to_image/text_to_image.py.
+Three invariants:
+
+1. The shipped argparse must default the parallel knobs to ``None`` so unset
+   flags do not enter ``omni_kwargs`` and a deploy YAML value stays in effect.
+2. Explicitly supplied values equal to the old argparse defaults (e.g.
+   ``--ulysses-degree 1``, ``--ulysses-mode strict``) must still be forwarded
+   so they can override the deploy YAML, instead of being silently dropped
+   (review feedback on PR #7293).
+3. The AR grid ``max_tokens`` budget must stay owned by the shared
+   ``ar_grid_max_tokens`` helper used by both this example and the online
+   serving path (review feedback on PR #7293).
+
+Both invariants are asserted against the real script (its ``parse_args()`` and
+its ``build_parallel_knob_kwargs``), not a re-implemented copy. Invariant 3 is
+asserted directly against ``ar_grid_max_tokens`` itself.
+"""
+
+import sys
+from pathlib import Path
+
+import pytest
+
+_script_dir = str(Path(__file__).parent.parent.parent.parent / "examples" / "offline_inference" / "text_to_image")
+sys.path.insert(0, _script_dir)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_KNOBS = (
+    "ulysses_degree",
+    "ulysses_mode",
+    "ring_degree",
+    "cfg_parallel_size",
+    "vae_patch_parallel_size",
+)
+
+
+def _parse_real_args(cli_flags: list[str]):
+    """Run the example's real argparse against a synthetic argv."""
+    import text_to_image as t2i
+
+    saved = sys.argv
+    sys.argv = ["text_to_image.py", *cli_flags]
+    try:
+        return t2i.parse_args()
+    finally:
+        sys.argv = saved
+
+
+class TestParallelKnobForwarding:
+    """Drive the example's real argparse + real forwarding helper."""
+
+    def test_unset_knobs_forward_nothing(self):
+        """No CLI flags -> no keys, so the deploy YAML values stay in effect."""
+        import text_to_image as t2i
+
+        args = _parse_real_args([])
+        assert t2i.build_parallel_knob_kwargs(args) == {}
+
+    def test_explicit_default_equivalent_values_are_forwarded(self):
+        """--ulysses-degree 1 / --ulysses-mode strict must override the YAML."""
+        import text_to_image as t2i
+
+        args = _parse_real_args(["--ulysses-degree", "1", "--ulysses-mode", "strict"])
+        assert t2i.build_parallel_knob_kwargs(args) == {
+            "ulysses_degree": 1,
+            "ulysses_mode": "strict",
+        }
+
+        args = _parse_real_args(["--ring-degree", "1", "--cfg-parallel-size", "1", "--vae-patch-parallel-size", "1"])
+        assert t2i.build_parallel_knob_kwargs(args) == {
+            "ring_degree": 1,
+            "cfg_parallel_size": 1,
+            "vae_patch_parallel_size": 1,
+        }
+
+    def test_non_default_values_are_forwarded(self):
+        import text_to_image as t2i
+
+        args = _parse_real_args(
+            [
+                "--ulysses-degree",
+                "2",
+                "--ulysses-mode",
+                "advanced_uaa",
+                "--ring-degree",
+                "2",
+                "--cfg-parallel-size",
+                "2",
+                "--vae-patch-parallel-size",
+                "2",
+            ]
+        )
+        assert t2i.build_parallel_knob_kwargs(args) == {
+            "ulysses_degree": 2,
+            "ulysses_mode": "advanced_uaa",
+            "ring_degree": 2,
+            "cfg_parallel_size": 2,
+            "vae_patch_parallel_size": 2,
+        }
+
+    def test_script_defaults_are_none(self):
+        """The shipped argparse must default to None so unset == defer to YAML."""
+        import text_to_image as t2i
+
+        parser = t2i.build_parser()
+        defaults = {key: parser.get_default(key) for key in _KNOBS}
+        assert defaults == {key: None for key in _KNOBS}
+
+
+class TestARGridMaxTokens:
+    """The AR grid budget formula must stay shared between serving and offline.
+
+    ``ar_grid_max_tokens`` (vllm_omni.diffusion.utils.param_utils) is the
+    single owner of the ``ar_height * (ar_width + 1) + 1`` contract used by
+    both the online serving path (serving_chat.py) and this example script,
+    so a grid-contract change cannot desynchronize one copy (review feedback
+    on PR #7293).
+    """
+
+    def test_grid_budget_matches_formula(self):
+        from vllm_omni.diffusion.utils.param_utils import ar_grid_max_tokens
+
+        assert ar_grid_max_tokens(64, 64) == 64 * 65 + 1
+        assert ar_grid_max_tokens(32, 48) == 48 * 33 + 1
+
+    def test_grid_budget_none_for_absent_or_invalid_grid(self):
+        from vllm_omni.diffusion.utils.param_utils import ar_grid_max_tokens
+
+        # Callers treat None as "keep the existing max_tokens default".
+        assert ar_grid_max_tokens(0, 0) is None
+        assert ar_grid_max_tokens(0, 64) is None
+        assert ar_grid_max_tokens(64, 0) is None

tests/helpers/fixtures/cpu_threads.py

@@ -0,0 +1,52 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import time
+from collections.abc import Callable, Generator
+from typing import Protocol
+
+import pytest
+
+
+class _TorchThreadAPI(Protocol):
+    def get_num_threads(self) -> int: ...
+
+    def set_num_threads(self, num_threads: int) -> None: ...
+
+
+def _single_threaded_cpu(
+    request: pytest.FixtureRequest,
+    load_torch: Callable[[], _TorchThreadAPI],
+) -> Generator[None, None, None]:
+    """Implement the marker guard separately so its state handling is testable."""
+    if request.node.get_closest_marker("cpu") is None:
+        yield
+        return
+
+    torch = load_torch()
+    previous_threads = torch.get_num_threads()
+    started = time.perf_counter()
+    try:
+        torch.set_num_threads(1)
+        print(f"{request.node.nodeid}: CPU reference threads {previous_threads} -> {torch.get_num_threads()}")
+        yield
+    finally:
+        torch.set_num_threads(previous_threads)
+        print(
+            f"{request.node.nodeid}: CPU reference elapsed {time.perf_counter() - started:.3f}s; "
+            f"restored threads {torch.get_num_threads()}"
+        )
+
+
+@pytest.fixture
+def single_threaded_cpu(request: pytest.FixtureRequest) -> Generator[None, None, None]:
+    """Bound opted-in CPU tests while leaving GPU-marked cases unchanged."""
+
+    def load_torch() -> _TorchThreadAPI:
+        import torch
+
+        return torch
+
+    yield from _single_threaded_cpu(request, load_torch)

tests/helpers/tests/test_cpu_threads.py

@@ -0,0 +1,89 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+
+from tests.helpers.fixtures.cpu_threads import _single_threaded_cpu
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _ThreadController:
+    def __init__(self, num_threads: int, *, fail_when_bounding: bool = False):
+        self.num_threads = num_threads
+        self.fail_when_bounding = fail_when_bounding
+        self.transitions: list[int] = []
+
+    def get_num_threads(self) -> int:
+        return self.num_threads
+
+    def set_num_threads(self, num_threads: int) -> None:
+        self.num_threads = num_threads
+        self.transitions.append(num_threads)
+        if self.fail_when_bounding and num_threads == 1:
+            raise RuntimeError("could not set thread count")
+
+
+def _request(*, cpu_marked: bool):
+    marker = object() if cpu_marked else None
+    node = SimpleNamespace(
+        nodeid="tests/helpers/tests/test_cpu_threads.py::probe",
+        get_closest_marker=lambda name: marker if name == "cpu" else None,
+    )
+    return SimpleNamespace(node=node)
+
+
+def _finish_fixture(fixture) -> None:
+    with pytest.raises(StopIteration):
+        next(fixture)
+
+
+def test_registered_fixture_caps_cpu_marked_test(single_threaded_cpu):
+    import torch
+
+    assert torch.get_num_threads() == 1
+
+
+def test_single_threaded_cpu_caps_and_restores(capsys):
+    torch = _ThreadController(8)
+    fixture = _single_threaded_cpu(_request(cpu_marked=True), lambda: torch)
+
+    next(fixture)
+    assert torch.num_threads == 1
+    _finish_fixture(fixture)
+
+    assert torch.num_threads == 8
+    assert torch.transitions == [1, 8]
+    output = capsys.readouterr().out
+    assert "CPU reference threads 8 -> 1" in output
+    assert "restored threads 8" in output
+
+
+def test_single_threaded_cpu_is_noop_without_cpu_marker():
+    loaded = False
+
+    def load_torch():
+        nonlocal loaded
+        loaded = True
+        return _ThreadController(8)
+
+    fixture = _single_threaded_cpu(_request(cpu_marked=False), load_torch)
+    next(fixture)
+    _finish_fixture(fixture)
+
+    assert not loaded
+
+
+def test_single_threaded_cpu_restores_if_bounding_raises():
+    torch = _ThreadController(8, fail_when_bounding=True)
+    fixture = _single_threaded_cpu(_request(cpu_marked=True), lambda: torch)
+
+    with pytest.raises(RuntimeError, match="could not set thread count"):
+        next(fixture)
+
+    assert torch.num_threads == 8
+    assert torch.transitions == [1, 8]

tests/model_executor/models/cosyvoice3/test_cosyvoice3_f0_weight_norm.py

@@ -10,6 +10,7 @@ from torch.nn.utils import parametrize
 from vllm_omni.model_executor.models.cosyvoice3.code2wav_core import hifigan
 
 pytestmark = [
+    pytest.mark.usefixtures("single_threaded_cpu"),
     pytest.mark.core_model,
     pytest.mark.filterwarnings("ignore:.*weight_norm.*:FutureWarning"),
 ]

tests/model_executor/models/cosyvoice3/test_cosyvoice3_hift_weight_norm.py

@@ -10,6 +10,7 @@ from torch.nn.utils import parametrize
 from vllm_omni.model_executor.models.cosyvoice3.code2wav_core import hifigan
 
 pytestmark = [
+    pytest.mark.usefixtures("single_threaded_cpu"),
     pytest.mark.core_model,
     pytest.mark.cpu,
     # The legacy weight_norm API warns on construction; the "legacy" params

tests/model_executor/models/cosyvoice3/test_cosyvoice3_incremental_hift.py

@@ -31,9 +31,8 @@ from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import (
     _wrapped_slice,
 )
 from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3_code2wav import CosyVoice3Code2Wav
-from vllm_omni.platforms import current_omni_platform
 
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.usefixtures("single_threaded_cpu")]
 
 CHUNK_LEN = 24
 TOTAL_MEL = 96
@@ -341,10 +340,6 @@ def test_incremental_hift_matches_reference_for_voiced_f0(config):
 
 @pytest.mark.parametrize("config", CONFIGS)
 @pytest.mark.parametrize("first_len,step_len", [(10, 6), (10, 3), (20, 3), (10, 1)])
-@pytest.mark.skipif(
-    current_omni_platform.is_rocm(),
-    reason="exhaustive CPU reference matrix exceeds the AMD CI time budget",
-)
 def test_incremental_hift_matches_reference_for_voiced_f0_small_chunks(config, first_len, step_len):
     """Same as test_incremental_hift_matches_reference_for_voiced_f0, but with small,
     sub-receptive-field chunk sizes (down to 1 mel frame) instead of the uniform

tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py

@@ -3,32 +3,32 @@
 
 from __future__ import annotations
 
-import functools
+from collections import OrderedDict
 from threading import Lock
 from types import SimpleNamespace
-from typing import TYPE_CHECKING
 
 import pytest
 import torch
 import torch.nn as nn
+from vllm.config import VllmConfig
+from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.outputs import SamplerOutput
 from vllm.v1.sample.logits_processor.state import LogitsProcessors
 from vllm.v1.sample.metadata import SamplingMetadata
+from vllm.v1.worker.gpu_input_batch import InputBatch
+from vllm.v1.worker.gpu_model_runner import GPUModelRunner
+
+from vllm_omni.config.model import OmniModelConfig
+from vllm_omni.data_entry_keys import to_struct
+from vllm_omni.model_executor.models.cosyvoice3 import cosyvoice3, cosyvoice3_talker
+from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3 import CosyVoice3Model, _normalize_request_conditioning
+from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config
+from vllm_omni.worker import gpu_ar_model_runner
+from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
+from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
-if TYPE_CHECKING:
-    from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3 import CosyVoice3Model
-
-
-@functools.lru_cache(maxsize=1)
-def _cosyvoice3_model_and_runner():
-    """Defer heavy Omni/vLLM imports until a test runs (avoids duplicate CustomOp init)."""
-    from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3 import CosyVoice3Model
-    from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
-
-    return CosyVoice3Model, GPUARModelRunner
-
 
 class _DummyCode2Wav:
     def __init__(
@@ -89,7 +89,6 @@ def _make_code2wav_model(
     num_samples: int = 32,
     outputs: list[tuple[torch.Tensor, dict[str, object] | None]] | None = None,
 ) -> CosyVoice3Model:
-    CosyVoice3Model, _ = _cosyvoice3_model_and_runner()
     model = object.__new__(CosyVoice3Model)
     nn.Module.__init__(model)
     model.model_stage = "cosyvoice3_code2wav"
@@ -115,7 +114,6 @@ def _make_code2wav_model(
 
 
 def _make_talker_model() -> CosyVoice3Model:
-    CosyVoice3Model, _ = _cosyvoice3_model_and_runner()
     model = object.__new__(CosyVoice3Model)
     nn.Module.__init__(model)
     model.model_stage = "cosyvoice3_talker"
@@ -160,6 +158,162 @@ def _make_sampling_metadata(
     )
 
 
+@pytest.fixture
+def output_payload_vllm_config(tmp_path, request):
+    hf_config = CosyVoice3Config()
+    hf_config.llm.update(llm_input_size=16, llm_output_size=16, speech_token_size=32)
+    # Exercise the real config fields without model downloads or engine setup.
+    model_config = object.__new__(OmniModelConfig)
+    model_config.hf_config = hf_config
+    model_config.model_stage = "cosyvoice3_talker"
+    model_config.model = str(tmp_path)
+    model_config.async_chunk = getattr(request, "param", True)
+    model_config.enable_return_routed_experts = False
+    model_config.engine_output_type = "latent"
+    model_config.stage_connector_config = {"name": "SharedMemoryConnector", "extra": {"role": "sender"}}
+    vllm_config = object.__new__(VllmConfig)
+    vllm_config.model_config = model_config
+    return vllm_config
+
+
+@pytest.fixture
+def output_payload_talker(monkeypatch, output_payload_vllm_config):
+    class DummyEncoder(nn.Module):
+        def __init__(self, **kwargs):
+            super().__init__()
+
+        def forward(self, inputs_embeds, positions):
+            return inputs_embeds
+
+    class DummyTalker(nn.Module):
+        def __init__(self, *, llm: nn.Module, **kwargs):
+            super().__init__()
+            self.llm = llm
+
+    monkeypatch.setattr(cosyvoice3_talker, "VLLMQwen2Encoder", DummyEncoder)
+    monkeypatch.setattr(cosyvoice3_talker, "CosyVoice3LM", DummyTalker)
+    monkeypatch.setattr(CosyVoice3Model, "_create_llm_vllm_config", lambda self, config: config)
+
+    return CosyVoice3Model(vllm_config=output_payload_vllm_config)
+
+
+@pytest.mark.parametrize(
+    "output_payload_vllm_config", [True, False], indirect=True, ids=["async_chunks", "sync_chunks"]
+)
+@pytest.mark.parametrize("async_scheduling", [True, False], ids=["async_scheduling", "sync_scheduling"])
+def test_talker_skips_unused_hidden_payload_without_async_materialization(
+    output_payload_talker, output_payload_vllm_config, async_scheduling
+):
+    runner = object.__new__(GPUARModelRunner)
+    runner.use_async_scheduling = async_scheduling
+    runner.omni_prefix_cache = None
+    runner.speculative_config = None
+    runner.model_config = output_payload_vllm_config.model_config
+    runner.model = output_payload_talker
+    # load_model caches this policy before the output path runs.
+    runner._pooler_payload_include_hidden_flag = bool(getattr(runner.model, "omni_pooler_payload_include_hidden", True))
+
+    assert runner._model_omni_pooler_payload_include_hidden() is (not runner.model_config.async_chunk)
+    assert runner._should_use_async_omni_output() is False
+
+
+@pytest.mark.parametrize("is_prefill", [True, False], ids=["prefill", "decode"])
+@pytest.mark.parametrize(
+    "output_payload_vllm_config", [True, False], indirect=True, ids=["async_chunks", "sync_chunks"]
+)
+def test_talker_inline_output_preserves_conditioning_and_tokens(
+    output_payload_talker, output_payload_vllm_config, monkeypatch, mocker, is_prefill
+):
+    """Token-only output retains prefill conditioning and sampled codec IDs."""
+
+    runner = object.__new__(GPUARModelRunner)
+    runner.model = output_payload_talker
+    runner._pooler_payload_include_hidden_flag = bool(getattr(runner.model, "omni_pooler_payload_include_hidden", True))
+    runner.vllm_config = output_payload_vllm_config
+    runner.model_config = output_payload_vllm_config.model_config
+    runner._async_chunk = runner.model_config.async_chunk
+    runner.omni_prefix_cache = None
+    runner.supports_mm_inputs = False
+    runner.routed_experts_initialized = False
+    runner.model_intermediate_buffer = {}
+    runner.input_batch = object.__new__(InputBatch)
+    runner.input_batch._req_ids = ["r1", "r2"]
+    runner.input_batch.req_id_to_index = {"r1": 0, "r2": 1}
+    monkeypatch.setattr(GPUARModelRunner, "_resolve_pooler_payload_req_ids", lambda self, req_ids: ("latent", req_ids))
+    monkeypatch.setattr(GPUARModelRunner, "get_omni_connector_output", lambda self: None)
+
+    conditioning = {
+        "speech_token": torch.tensor([[11, 12], [21, 0]], dtype=torch.long),
+        "speech_token_len": torch.tensor([2, 1], dtype=torch.long),
+        "speech_feat": torch.arange(16, dtype=torch.float32).reshape(2, 4, 2),
+        "embedding": torch.tensor([[0.1, 0.2], [0.3, 0.4]]),
+    }
+    expected = {key: value.clone() for key, value in conditioning.items()}
+    hidden_states = output_payload_talker.forward(
+        input_ids=torch.tensor([1, 2, 3]),
+        positions=torch.arange(3),
+        inputs_embeds=torch.ones(3, 16),
+        **(conditioning if is_prefill else {}),
+    )
+    model_output = output_payload_talker.make_omni_output(hidden_states, **(conditioning if is_prefill else {}))
+
+    cpu_materialization = mocker.spy(gpu_ar_model_runner, "_to_cpu_contiguous")
+    scheduler_output = object.__new__(SchedulerOutput)
+    scheduler_output.total_num_scheduled_tokens = 3
+    scheduler_output.num_scheduled_tokens = {"r1": 2, "r2": 1}
+    output = runner._build_omni_model_runner_output_from_snapshot(
+        scheduler_output=scheduler_output,
+        hidden_states=model_output.text_hidden_states,
+        staged_hidden_states_cpu=None,
+        multimodal_outputs=model_output.multimodal_outputs,
+        req_ids_output_copy=["r1", "r2"],
+        req_id_to_index_output_copy={"r1": 0, "r2": 1},
+        valid_sampled_token_ids=[[101], [102]],
+        logprobs_lists=None,
+        prompt_logprobs_dict={},
+        num_nans_in_logits=None,
+        kv_connector_output=None,
+        ec_connector_output=None,
+        cudagraph_stats=None,
+        kv_extracted_req_ids=None,
+        num_scheduled_tokens_np=torch.tensor([2, 1], dtype=torch.int32).numpy(),
+        query_start_loc_cpu=torch.tensor([0, 2], dtype=torch.long),
+    )
+    assert output.req_ids == ["r1", "r2"]
+    assert output.sampled_token_ids == [[101], [102]]
+    hidden_storage = model_output.text_hidden_states.untyped_storage().data_ptr()
+    hidden_copied = any(
+        call.args[0].untyped_storage().data_ptr() == hidden_storage for call in cpu_materialization.call_args_list
+    )
+    assert hidden_copied is (not runner._async_chunk)
+    if runner._async_chunk:
+        assert output.multimodal_outputs is None
+    if not is_prefill and runner._async_chunk:
+        assert not output.inter_stage_outputs or all(not item for item in output.inter_stage_outputs)
+        return
+
+    assert len(output.inter_stage_outputs) == 2
+    for idx, prompt_len in enumerate([2, 1]):
+        item = output.inter_stage_outputs[idx]
+        assert ("hidden" not in item) is runner._async_chunk
+        if not is_prefill:
+            continue
+        # The downstream processor removes padding using speech_token_len.
+        torch.testing.assert_close(item["embed.speech_token"], expected["speech_token"][idx : idx + 1])
+        torch.testing.assert_close(item["embed.speech_feat"], expected["speech_feat"][idx : idx + 1])
+        assert item["embed.speech_token_len"].item() == prompt_len
+        torch.testing.assert_close(item["embed.embedding"], expected["embedding"][idx : idx + 1])
+
+
+def test_talker_preserves_packed_payload_policy(output_payload_talker, output_payload_vllm_config, monkeypatch):
+    output_payload_vllm_config.model_config.async_chunk = False
+    monkeypatch.setattr(cosyvoice3, "cosyvoice3_packed_inference_enabled", lambda: True)
+    model = CosyVoice3Model(vllm_config=output_payload_vllm_config)
+
+    assert model.omni_pooler_payload_include_hidden is False
+    assert not getattr(model, "use_async_omni_output", False)
+
+
 def test_forward_prefers_token_offset_when_present():
     model = _make_code2wav_model()
 
@@ -561,7 +715,6 @@ def test_gpu_ar_model_runner_prefers_model_sampler_when_opted_in():
             # runs for all sampling paths
             self.updated = True
 
-    _, GPUARModelRunner = _cosyvoice3_model_and_runner()
     runner = object.__new__(GPUARModelRunner)
     runner.input_batch = _DummyInputBatch()
 
@@ -602,7 +755,6 @@ def test_gpu_ar_model_runner_supplies_req_output_history_to_model_sampler():
             # runs for all sampling paths
             self.update_async_called = True
 
-    _, GPUARModelRunner = _cosyvoice3_model_and_runner()
     runner = object.__new__(GPUARModelRunner)
     runner.input_batch = _DummyInputBatch()
 
@@ -649,7 +801,6 @@ def test_gpu_ar_model_runner_repairs_async_placeholders_for_model_sampler():
             # runs for all sampling paths (model sampler + fallback sampler)
             self.update_async_called = True
 
-    _, GPUARModelRunner = _cosyvoice3_model_and_runner()
     runner = object.__new__(GPUARModelRunner)
     runner.input_batch = _DummyInputBatch()
 
@@ -722,8 +873,6 @@ def test_embed_input_ids_preserves_interleaved_request_boundaries(order):
 
 
 def test_mrv2_conditioning_follows_encoded_prompts_to_batch_rows():
-    from collections import OrderedDict
-
     model = _make_talker_model()
     speech = nn.Embedding.from_pretrained(torch.arange(100 * 4, dtype=torch.float32).reshape(100, 4))
     text = nn.Embedding.from_pretrained(torch.arange(100 * 4, dtype=torch.float32).reshape(100, 4) + 1000)
@@ -791,10 +940,8 @@ def test_embed_input_ids_rejects_cross_request_placeholder_block():
 
 
 def test_full_response_singleton_uses_the_optimized_batch_path(monkeypatch):
-    import vllm_omni.model_executor.models.cosyvoice3.cosyvoice3 as module
-
     monkeypatch.setenv("COSYVOICE3_BATCH_FLOW", "1")
-    monkeypatch.setattr(module, "cosyvoice3_packed_inference_enabled", lambda: True)
+    monkeypatch.setattr(cosyvoice3, "cosyvoice3_packed_inference_enabled", lambda: True)
     model = _make_code2wav_model()
     model.forward(
         input_ids=torch.tensor([0, 1, 2]),
@@ -816,10 +963,6 @@ def test_full_response_singleton_uses_the_optimized_batch_path(monkeypatch):
 
 
 def test_full_response_adapter_uses_live_prefill_then_empty_decode(monkeypatch):
-    from vllm.v1.worker.gpu_model_runner import GPUModelRunner
-
-    from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
-
     model = _make_talker_model()
     hidden = torch.ones(2, 4)
     runner = object.__new__(OmniGPUModelRunner)
@@ -840,13 +983,11 @@ def test_full_response_adapter_uses_live_prefill_then_empty_decode(monkeypatch):
 
 
 def test_packed_flow_preserves_torch_estimator_when_speaker_trt_is_enabled(monkeypatch):
-    import vllm_omni.model_executor.models.cosyvoice3.cosyvoice3 as module
-
     model = _make_code2wav_model()
     model._code2wav_trt_done = False
-    monkeypatch.setattr(module.torch.cuda, "is_available", lambda: True)
-    monkeypatch.setattr(module, "cosyvoice3_packed_inference_enabled", lambda: True)
-    monkeypatch.setattr(module, "_cosyvoice3_trt_enabled", lambda: True)
+    monkeypatch.setattr(cosyvoice3.torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(cosyvoice3, "cosyvoice3_packed_inference_enabled", lambda: True)
+    monkeypatch.setattr(cosyvoice3, "_cosyvoice3_trt_enabled", lambda: True)
 
     def unexpected_estimator_resolution():
         raise AssertionError("Packed Flow must not resolve or replace its torch estimator")
@@ -904,21 +1045,18 @@ def test_sampling_mode_preserves_or_merges_control_logits(mode):
 
 
 def test_cancelled_stream_releases_only_its_vocoder_state():
-    model_cls, _ = _cosyvoice3_model_and_runner()
     model = SimpleNamespace(
         _stream_audio_cache_lock=Lock(),
         _stream_vocoder_cache_by_req={"cancelled": {"mel": torch.ones(2)}, "live": {"mel": torch.zeros(2)}},
     )
-    model_cls.on_requests_finished(model, {"cancelled", "unknown"})
-    model_cls.on_requests_finished(model, {"cancelled"})
+    CosyVoice3Model.on_requests_finished(model, {"cancelled", "unknown"})
+    CosyVoice3Model.on_requests_finished(model, {"cancelled"})
     assert set(model._stream_vocoder_cache_by_req) == {"live"}
 
 
 @pytest.mark.parametrize("mrv2,packed", [(False, False), (False, True), (True, False)])
 def test_talker_output_contract(monkeypatch, mrv2, packed):
-    import vllm_omni.model_executor.models.cosyvoice3.cosyvoice3 as mod
-
-    monkeypatch.setattr(mod, "cosyvoice3_packed_inference_enabled", lambda: packed)
+    monkeypatch.setattr(cosyvoice3, "cosyvoice3_packed_inference_enabled", lambda: packed)
     model = _make_talker_model()
     hidden = torch.zeros(2, 4)
     model.model = SimpleNamespace(llm=lambda embeddings, positions: embeddings)
@@ -932,9 +1070,6 @@ def test_talker_output_contract(monkeypatch, mrv2, packed):
 
 
 def test_request_conditioning_normalizes_singletons_without_changing_schema():
-    from vllm_omni.data_entry_keys import to_struct
-    from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3 import _normalize_request_conditioning
-
     token = torch.tensor([[7, 8]], dtype=torch.int32)
     feat = torch.randn(1, 4, 80)
     speaker = torch.randn(1, 192)
@@ -949,8 +1084,6 @@ def test_request_conditioning_normalizes_singletons_without_changing_schema():
 
 
 def test_request_conditioning_rejects_unsplit_voices():
-    from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3 import _normalize_request_conditioning
-
     with pytest.raises(ValueError, match="unsplit batch"):
         _normalize_request_conditioning({"embed": {"embedding": [torch.ones(1, 192), torch.zeros(1, 192)]}})
 

tests/model_executor/models/indextts2/test_talker_v2_5.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from types import SimpleNamespace
 
@@ -215,8 +215,8 @@ def test_mel_head_preserves_official_checkpoint_bias_and_hidden_payload_policy(
 
     assert seen["bias"] is True
     assert talker.mel_head.bias is not None
-    assert talker.omni_pooler_payload_include_hidden is use_gpt_latent
-    assert talker.omni_payload_at_request_end is (model_type == "indextts2_5")
+    assert talker.omni_pooler_payload_include_hidden is False
+    assert talker.omni_payload_at_request_end is True
 
 
 def test_v2_load_weights_loads_exact_mel_head_bias_parameter(monkeypatch):

tests/model_executor/models/mammoth_moda2/test_mammoth_moda2_teacache.py

@@ -0,0 +1,274 @@
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.diffusion.cache.cachedit import RequestScopedCacheDiTRuntime
+from vllm_omni.diffusion.cache.teacache.config import TeaCacheConfig
+from vllm_omni.diffusion.cache.teacache.extractors import extract_mammoth_moda2_context
+from vllm_omni.diffusion.data import DiffusionOutput
+from vllm_omni.diffusion.models.mammoth_moda2 import pipeline_mammothmoda2_dit as mammoth_pipeline_module
+from vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit import MammothModa2DiTPipeline
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.cpu]
+
+
+def test_mammoth_moda2_calibrated_teacache_defaults():
+    config = TeaCacheConfig(transformer_type="MammothModa2Transformer2DModel")
+
+    assert config.rel_l1_thresh == 0.075
+    assert config.coefficients == [
+        -1761.242764481119,
+        859.398730831851,
+        -126.78044436937441,
+        9.16173963457422,
+        -0.1102826051200547,
+    ]
+
+
+class _FakeScheduler:
+    def set_timesteps(self, num_inference_steps, device, num_tokens):  # noqa: ARG002
+        self.timesteps = torch.arange(num_inference_steps, 0, -1, device=device, dtype=torch.float32)
+
+    def step(self, model_pred, t, latents, return_dict=False):  # noqa: ARG002
+        return (latents,)
+
+
+class _FakeTransformer(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.config = SimpleNamespace(in_channels=4)
+        self.time_caption_embed = SimpleNamespace(image_embedder=None)
+        self.param = nn.Parameter(torch.zeros(()))
+        self.branches: list[str | None] = []
+        self.calls: list[dict] = []
+
+    def forward(self, hidden_states, **kwargs):
+        self.branches.append(kwargs.get("teacache_branch"))
+        self.calls.append(kwargs)
+        return torch.zeros_like(hidden_states)
+
+
+class _FakeVAE(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.param = nn.Parameter(torch.zeros(()))
+        self.config = SimpleNamespace(scaling_factor=None, shift_factor=None)
+
+    def decode(self, latents, return_dict=False):  # noqa: ARG002
+        return (latents,)
+
+
+def _build_pipeline(monkeypatch):
+    monkeypatch.setattr(mammoth_pipeline_module, "FlowMatchEulerDiscreteScheduler", _FakeScheduler)
+    monkeypatch.setattr(
+        mammoth_pipeline_module,
+        "randn_tensor",
+        lambda shape, *, device, dtype, generator=None: torch.zeros(shape, device=device, dtype=dtype),
+    )
+
+    pipe = MammothModa2DiTPipeline.__new__(MammothModa2DiTPipeline)
+    nn.Module.__init__(pipe)
+    pipe.gen_transformer = _FakeTransformer()
+    pipe.gen_vae = _FakeVAE()
+    pipe.gen_image_condition_refiner = None
+    pipe.gen_freqs_cis = []
+    pipe._llm_hidden_size = 8
+    pipe._cache_dit_config = None
+    pipe._cache_dit_runtime = RequestScopedCacheDiTRuntime(pipe)
+    pipe.config = SimpleNamespace(
+        llm_config=SimpleNamespace(gen_vocab_start_index=100),
+        image_token_id=900,
+        video_token_id=901,
+        vision_start_token_id=902,
+        vision_end_token_id=903,
+    )
+    return pipe
+
+
+def _run_pipeline(pipe, *, text_guidance_scale, cfg_range, num_inference_steps=4, num_requests=1):
+    request = OmniDiffusionRequest(
+        request_id="req-teacache",
+        prompt={
+            "prompt": "",
+            "height": 32,
+            "width": 32,
+            "additional_information": {
+                "full_hidden_states": torch.randn(3, 8),
+                "full_token_ids": [10, 11, 100],
+                "answer_start_index": 2,
+            },
+        },
+        sampling_params=OmniDiffusionSamplingParams(
+            height=32,
+            width=32,
+            guidance_scale=text_guidance_scale,
+            num_inference_steps=num_inference_steps,
+            extra_args={"cfg_range": cfg_range},
+        ),
+    )
+    requests = [
+        OmniDiffusionRequest(
+            request_id=f"req-teacache-{i}",
+            prompt=request.prompt,
+            sampling_params=request.sampling_params,
+        )
+        for i in range(num_requests)
+    ]
+    outputs = pipe(DiffusionRequestBatch(requests))
+    assert isinstance(outputs, list)
+    assert len(outputs) == num_requests
+    assert all(isinstance(output, DiffusionOutput) and output.output.shape == (1, 4, 4, 4) for output in outputs)
+    assert all(call["timestep"].shape == (num_requests,) for call in pipe.gen_transformer.calls)
+    return pipe.gen_transformer.branches
+
+
+@pytest.mark.parametrize("num_requests", [1, 2])
+def test_mammoth_moda2_non_cfg_passes_positive_teacache_branch(monkeypatch, num_requests):
+    pipe = _build_pipeline(monkeypatch)
+
+    branches = _run_pipeline(pipe, text_guidance_scale=1.0, cfg_range=[0.0, 1.0], num_requests=num_requests)
+
+    assert branches == ["positive", "positive", "positive", "positive"]
+
+
+def test_mammoth_moda2_dev_forwards_ar_image_conditioning_to_positive_branch(monkeypatch):
+    """Dev keeps AR image tokens separate and forwards them to the positive DiT call."""
+    pipe = _build_pipeline(monkeypatch)
+    pipe.gen_transformer.time_caption_embed.image_embedder = object()
+
+    _run_pipeline(
+        pipe,
+        text_guidance_scale=1.0,
+        cfg_range=[0.0, 1.0],
+        num_inference_steps=1,
+    )
+
+    assert len(pipe.gen_transformer.calls) == 1
+    call = pipe.gen_transformer.calls[0]
+    assert call["teacache_branch"] == "positive"
+    assert call["text_hidden_states"].shape == (1, 2, 8)
+    assert call["ar_image_hidden_states"].shape == (1, 1, 8)
+    assert torch.equal(call["ar_image_attention_mask"], torch.ones(1, 1, dtype=torch.bool))
+
+
+@pytest.mark.parametrize("num_requests", [1, 2])
+def test_mammoth_moda2_cfg_passes_positive_then_negative_teacache_branch(monkeypatch, num_requests):
+    pipe = _build_pipeline(monkeypatch)
+
+    branches = _run_pipeline(pipe, text_guidance_scale=4.0, cfg_range=[0.0, 1.0], num_requests=num_requests)
+
+    assert branches == [
+        "positive",
+        "negative",
+        "positive",
+        "negative",
+        "positive",
+        "negative",
+        "positive",
+        "negative",
+    ]
+
+
+@pytest.mark.parametrize("num_requests", [1, 2])
+def test_mammoth_moda2_cfg_range_only_uses_negative_inside_range(monkeypatch, num_requests):
+    pipe = _build_pipeline(monkeypatch)
+
+    branches = _run_pipeline(pipe, text_guidance_scale=4.0, cfg_range=[0.5, 1.0], num_requests=num_requests)
+
+    assert branches == [
+        "positive",
+        "positive",
+        "positive",
+        "negative",
+        "positive",
+        "negative",
+    ]
+
+
+def test_mammoth_moda2_teacache_forwards_ar_image_conditioning():
+    """The TeaCache extractor follows the full embedding contract, including the expanded text mask."""
+    hidden_size = 4
+    hidden_states = torch.zeros(1, hidden_size, 1, 1)
+    timestep = torch.zeros(1)
+    text_hidden_states = torch.zeros(1, 1, hidden_size)
+    text_attention_mask = torch.ones(1, 1, dtype=torch.bool)
+    freqs_cis = object()
+    ar_image_hidden_states = torch.randn(1, 2, hidden_size)
+    ar_image_attention_mask = torch.ones(1, 2, dtype=torch.bool)
+
+    temb = torch.zeros(1, hidden_size)
+    prepared_text_hidden_states = torch.randn(1, 3, hidden_size)
+    prepared_text_attention_mask = torch.ones(1, 3, dtype=torch.bool)
+    img_tokens = torch.randn(1, 1, hidden_size)
+    img_mask = torch.ones(1, 1, dtype=torch.bool)
+    context_rotary_emb = object()
+    noise_rotary_emb = object()
+    rotary_emb = object()
+
+    first_layer = SimpleNamespace(
+        norm1=Mock(side_effect=lambda joint_hidden_states, _temb: (joint_hidden_states,)),
+    )
+    module = SimpleNamespace(
+        layers=[first_layer],
+        config=SimpleNamespace(hidden_size=hidden_size, patch_size=1),
+        _validate_inputs=Mock(return_value=(1, 1, 1)),
+        _prepare_embeddings=Mock(
+            return_value=(
+                temb,
+                prepared_text_hidden_states,
+                prepared_text_attention_mask,
+                img_tokens,
+                img_mask,
+                1,
+                context_rotary_emb,
+                noise_rotary_emb,
+                rotary_emb,
+                [3],
+                [4],
+            )
+        ),
+        _apply_refiners=Mock(return_value=(prepared_text_hidden_states, img_tokens)),
+    )
+
+    ctx = extract_mammoth_moda2_context(
+        module,
+        hidden_states=hidden_states,
+        timestep=timestep,
+        text_hidden_states=text_hidden_states,
+        text_attention_mask=text_attention_mask,
+        ref_image_hidden_states=None,
+        ar_image_hidden_states=ar_image_hidden_states,
+        ar_image_attention_mask=ar_image_attention_mask,
+        freqs_cis=freqs_cis,
+        teacache_branch="positive",
+    )
+
+    module._prepare_embeddings.assert_called_once_with(
+        hidden_states,
+        timestep,
+        text_hidden_states,
+        text_attention_mask,
+        freqs_cis,
+        1,
+        1,
+        1,
+        ar_image_hidden_states,
+        ar_image_attention_mask,
+    )
+    module._apply_refiners.assert_called_once_with(
+        prepared_text_hidden_states,
+        prepared_text_attention_mask,
+        context_rotary_emb,
+        img_tokens,
+        img_mask,
+        noise_rotary_emb,
+        temb,
+    )
+    assert ctx.hidden_states.shape == (1, 4, hidden_size)
+    assert ctx.extra_states["teacache_branch"] == "positive"

tests/model_executor/models/qwen3_omni/test_duplex_plugin.py

@@ -27,6 +27,7 @@ from vllm_omni.engine.duplex.commands import (
     UpdateSession,
 )
 from vllm_omni.engine.duplex.config import DuplexSessionConfig
+from vllm_omni.engine.duplex.delivery import DuplexOutputBuffer
 from vllm_omni.engine.duplex.messages import OpenDuplexSessionMessage
 from vllm_omni.engine.duplex.plugin import DuplexRuntimeConfigError
 from vllm_omni.engine.duplex.session.manager import DuplexSessionManager
@@ -77,6 +78,7 @@ async def open_qwen():
     port = RecordingStagePort(stage_count=3)
     output: asyncio.Queue[Any] = asyncio.Queue()
     results: asyncio.Queue[Any] = asyncio.Queue()
+    output_buffer = DuplexOutputBuffer(max_bytes=2 * 1024 * 1024, max_events=512)
     manager = DuplexSessionManager(
         plugin=plugin,
         stage_port=port,
@@ -86,10 +88,21 @@ async def open_qwen():
         model_config=None,
     )
     config = DuplexSessionConfig(model="qwen", modalities=["text", "audio"], overlap_policy="barge_in_on_speech")
-    await manager.handle(OpenDuplexSessionMessage(control_id="open", session_id=SESSION_ID, session_config=config))
+    await manager.handle(
+        OpenDuplexSessionMessage(
+            control_id="open", session_id=SESSION_ID, session_config=config, output_buffer=output_buffer
+        )
+    )
     result = await results.get()
     assert result.ok, result
-    harness = Harness(manager, port, output, results, manager.runners[SESSION_ID])
+    harness = Harness(
+        manager=manager,
+        port=port,
+        output=output,
+        output_buffer=output_buffer,
+        results=results,
+        runner=manager.runners[SESSION_ID],
+    )
     await harness.settle()
     return harness
 

tests/model_executor/models/qwen3_omni/test_qwen3_omni_deepstack_device.py

@@ -0,0 +1,89 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Deepstack buffers must remain usable when the vision tower is skipped."""
+
+import pytest
+import torch
+from torch import nn
+from transformers import Qwen3OmniMoeThinkerConfig
+from vllm.config import DeviceConfig, ModelConfig, MultiModalConfig, SchedulerConfig, VllmConfig
+from vllm.model_executor.models.utils import StageMissingLayer
+
+import vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker as thinker_module
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _LanguageModel(nn.Module):
+    def make_empty_intermediate_tensors(self):
+        raise AssertionError("No pipeline-parallel allocation is needed by this test")
+
+
+class _DecoderLayer(nn.Module):
+    def forward(self, positions, hidden_states, residual):
+        return hidden_states, residual
+
+
+class _Norm(nn.Module):
+    def forward(self, hidden_states, residual):
+        return hidden_states, residual
+
+
+@pytest.mark.parametrize("vision_limit", [0, 1])
+@pytest.mark.parametrize("num_tokens", [4, 8])
+def test_deepstack_buffers_use_execution_device(mocker, vision_limit, num_tokens):
+    """Exercise the real constructor, tower-skip context and LM addition."""
+    config = Qwen3OmniMoeThinkerConfig(
+        text_config={"hidden_size": 8},
+        vision_config={"out_hidden_size": 8, "deepstack_visual_indexes": [0, 1]},
+    )
+    multimodal_config = MultiModalConfig(limit_per_prompt={"image": vision_limit, "video": vision_limit, "audio": 1})
+    model_config = mocker.Mock(spec=ModelConfig, hf_config=config, multimodal_config=multimodal_config)
+    model_config.get_multimodal_config.return_value = multimodal_config
+    vllm_config = mocker.Mock(
+        spec=VllmConfig,
+        model_config=model_config,
+        device_config=DeviceConfig(device="cpu"),
+        scheduler_config=mocker.Mock(spec=SchedulerConfig, max_num_batched_tokens=4),
+        quant_config=None,
+    )
+    vllm_config.with_hf_config.return_value = vllm_config
+
+    # Replace only the large networks. vLLM's real _mark_tower_model context
+    # still builds the disabled visual module on meta and removes its weights.
+    mocker.patch.object(thinker_module, "Qwen3OmniMoeAudioEncoder", side_effect=lambda *args, **kwargs: nn.Linear(8, 8))
+    mocker.patch.object(thinker_module, "Qwen3Omni_VisionTransformer", side_effect=lambda **kwargs: nn.Linear(8, 8))
+    mocker.patch.object(thinker_module, "Qwen3MoeLLMForCausalLM", return_value=_LanguageModel())
+    thinker = thinker_module.Qwen3OmniMoeThinkerForConditionalGeneration(vllm_config=vllm_config)
+    assert isinstance(thinker.visual, StageMissingLayer) is (vision_limit == 0)
+
+    # Resizing must also retain the buffer device rather than the ambient one.
+    with torch.device("meta"):
+        deepstack = thinker._get_deepstack_input_embeds(num_tokens)
+
+    model = object.__new__(thinker_module.Qwen3MoeLLMModel)
+    nn.Module.__init__(model)
+    model.layers = nn.ModuleList([_DecoderLayer(), _DecoderLayer()])
+    model.start_layer = 0
+    model.end_layer = 2
+    model.norm = _Norm()
+    mocker.patch.object(thinker_module, "get_pp_group", return_value=mocker.Mock(is_first_rank=True, is_last_rank=True))
+    inputs = torch.ones(num_tokens, 8)
+    # Before the fix, audio-only construction reaches this actual forward
+    # with meta buffers and raises the reported cross-device addition error.
+    output, _ = thinker_module.Qwen3MoeLLMModel.forward(
+        model,
+        input_ids=None,
+        positions=torch.arange(num_tokens),
+        inputs_embeds=inputs,
+        deepstack_input_embeds=deepstack,
+    )
+    torch.testing.assert_close(output, inputs)
+    assert len(thinker.deepstack_input_embeds) == 2
+    assert thinker.deepstack_input_embeds_num_tokens == 0
+    for buffer in thinker.deepstack_input_embeds:
+        assert buffer.device == torch.device("cpu")
+        assert not buffer.is_meta
+        assert buffer.shape == (num_tokens, 8)
+        assert buffer.dtype == torch.get_default_dtype()
+        assert torch.count_nonzero(buffer).item() == 0

tests/model_executor/models/qwen3_tts/test_qwen3_tts_incremental_decode.py

@@ -1,9 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 import copy
 from collections import Counter
 from types import SimpleNamespace
+from typing import Any
 
 import pytest
 import torch
@@ -22,7 +23,7 @@ from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.modeling_qwen3_tts
     Qwen3TTSTokenizerV2DecoderTransformerModel,
 )
 
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.usefixtures("single_threaded_cpu")]
 
 # Shared test helpers
 
@@ -341,7 +342,7 @@ def test_incremental_request_audio_matches_full_decode(mode):
         (1, decoder.config.num_quantizers, prefix_frames + suffix_frames),
     )
     suffix_codes = codes[..., prefix_frames:]
-    caches = {"prefix_frames": prefix_frames}
+    caches: dict[str, Any] = {"prefix_frames": prefix_frames}
 
     with torch.no_grad():
         first = decoder(codes[..., : prefix_frames + 1], caches=caches)
@@ -383,7 +384,7 @@ def test_single_frame_xvec_prefix_keeps_all_next_chunk_conv_frames():
     decoder = Qwen3TTSTokenizerV2Decoder(config).eval()
     decoder._incremental_chunk_frames = 25
     codes = torch.randint(0, config.codebook_size, (1, config.num_quantizers, 26))
-    caches = {"prefix_frames": 0}
+    caches: dict[str, Any] = {"prefix_frames": 0}
 
     with torch.no_grad():
         decoder(codes[..., :1], caches=caches)
@@ -421,7 +422,7 @@ def test_xvec_rolling_matches_truncated_suffix_window():
     decoder = Qwen3TTSTokenizerV2Decoder(config).eval()
     decoder._incremental_chunk_frames = 25
     codes = torch.randint(0, config.codebook_size, (1, config.num_quantizers, 101))
-    caches = {"prefix_frames": 0}
+    caches: dict[str, Any] = {"prefix_frames": 0}
 
     with torch.no_grad():
         decoder(codes[..., :1], caches=caches)
@@ -470,7 +471,7 @@ def test_icl_rolling_matches_reference_plus_truncated_suffix_window():
     )
     prefix_codes = codes[..., :prefix_frames]
     suffix_codes = codes[..., prefix_frames:]
-    caches = {"prefix_frames": prefix_frames}
+    caches: dict[str, Any] = {"prefix_frames": prefix_frames}
 
     with torch.no_grad():
         decoder(codes[..., : prefix_frames + 1], caches=caches)
@@ -975,7 +976,7 @@ def test_batched_chunked_decode_groups_exact_phases(monkeypatch):
 
     prefix_frames = 48
     codes_list = [torch.zeros(1, 2, prefix_frames + 1)]
-    caches = [{"prefix_frames": prefix_frames}]
+    caches: list[dict[str, Any]] = [{"prefix_frames": prefix_frames}]
     for previous, _suffix_frames, cached_frames in (
         (1, 26, 1),
         (26, 51, 26),
@@ -1064,7 +1065,7 @@ def test_missing_graph_phase_falls_back_instead_of_returning_empty(monkeypatch):
     wrapper._record_graph_fallback = lambda *_args: None
 
     codes = [torch.zeros(1, 2, 25)]
-    caches = [{}]
+    caches: list[dict[str, Any]] = [{}]
 
     assert wrapper._decode_icl_prefix_batch(codes, caches) is None
     assert wrapper._decode_suffix_batch("xvec", 50, codes, caches, [25]) is None
@@ -1078,9 +1079,12 @@ def test_xvec_first_chunk_uses_prefixless_initializer(monkeypatch):
     wrapper._icl_previous_frames_by_target = {26: 1}
     wrapper._xvec_previous_frames_by_target = {26: 1}
     calls = []
-    wrapper.decoder = _decoder_stub(
-        _decode_xvec_first_chunk=lambda codes, cache: calls.append("xvec") or codes[:, :1, :],
-    )
+
+    def decode_first_chunk(codes, cache):
+        calls.append("xvec")
+        return codes[:, :1, :]
+
+    wrapper.decoder = _decoder_stub(_decode_xvec_first_chunk=decode_first_chunk)
 
     output = wrapper._batched_request_decode(
         [torch.zeros(1, 2, 25)],

tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py

@@ -6,6 +6,7 @@ from types import SimpleNamespace
 
 import pytest
 import torch
+from vllm.sampling_params import SamplingParams
 
 from vllm_omni.data_entry_keys import serialize_payload
 from vllm_omni.model_executor.stage_input_processors.cosyvoice3 import (
@@ -14,6 +15,7 @@ from vllm_omni.model_executor.stage_input_processors.cosyvoice3 import (
     text2flow_token_only,
 )
 from vllm_omni.outputs.mm_outputs import MultimodalPayload
+from vllm_omni.request import OmniRequest
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -376,6 +378,61 @@ def test_talker2code2wav_async_chunk_emits_terminal_eof_without_duplicate_audio(
     assert payload_final.codes.audio.tolist() == []
 
 
+def test_interleaved_token_only_requests_keep_conditioning_and_eof_separate():
+    transfer_manager = _transfer_manager()
+    requests = [
+        OmniRequest(
+            request_id=f"token-only-{index}",
+            external_req_id=f"token-only-{index}",
+            prompt_token_ids=[1],
+            sampling_params=SamplingParams(max_tokens=16),
+            pooling_params=None,
+            arrival_time=0.0,
+            block_hasher=None,
+        )
+        for index in range(2)
+    ]
+    conditioning = [
+        {
+            "embed": {
+                "speech_token": torch.tensor([[101 + index, 201 + index]]),
+                "speech_feat": torch.arange(8, dtype=torch.float32).reshape(1, 4, 2) + index,
+                "embedding": torch.tensor([[0.1 + index, 0.2 + index]]),
+            }
+        }
+        for index in range(2)
+    ]
+
+    requests[0].append_output_token_ids([8, 9])
+    requests[1].append_output_token_ids([18, 19])
+    first_chunks = [
+        talker2code2wav_async_chunk(transfer_manager, payload, request)
+        for request, payload in zip(requests, conditioning)
+    ]
+    for index, chunk in enumerate(first_chunks):
+        assert chunk is not None
+        assert chunk.codes.audio.tolist() == ([8, 9] if index == 0 else [18, 19])
+        torch.testing.assert_close(chunk.embed.embedding, conditioning[index]["embed"]["embedding"])
+
+    requests[0].append_output_token_ids([10, 11])
+    final_a = talker2code2wav_async_chunk(transfer_manager, None, requests[0], is_finished=True)
+    requests[1].append_output_token_ids([20])
+    pending_b = talker2code2wav_async_chunk(transfer_manager, None, requests[1])
+    final_b = talker2code2wav_async_chunk(transfer_manager, None, requests[1], is_finished=True)
+
+    assert final_a.codes.audio.tolist() == [8, 9, 10, 11]
+    assert final_a.meta.left_context_size == 2
+    assert final_a.meta.finished.item() is True
+    assert final_a.embed is None
+    assert pending_b is None
+    assert final_b.codes.audio.tolist() == [18, 19, 20]
+    assert final_b.meta.left_context_size == 2
+    assert final_b.meta.finished.item() is True
+    assert final_b.embed is None
+    for request in requests:
+        assert talker2code2wav_async_chunk(transfer_manager, None, request, is_finished=True) is None
+
+
 def test_full_payload_retains_true_length_for_padded_conditioning():
     from vllm_omni.data_entry_keys import to_struct
 

tests/model_executor/stage_input_processors/test_glm_image.py

@@ -1,12 +1,15 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Unit tests for GLM-Image stage input processor."""
 
+from collections.abc import Mapping
 from types import SimpleNamespace
 
 import pytest
 import torch
 
+from vllm_omni.data_entry_keys import unflatten_payload
+from vllm_omni.model_executor.models.glm_image.pipeline import GLM_IMAGE_PIPELINE
 from vllm_omni.model_executor.stage_input_processors.glm_image import (
     _first_source_image,
     _has_source_image,
@@ -15,6 +18,7 @@ from vllm_omni.model_executor.stage_input_processors.glm_image import (
     ar2diffusion,
     compute_max_tokens,
 )
+from vllm_omni.outputs.mm_outputs import MultimodalPayload
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -24,7 +28,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 # =============================================================================
 
 
-def _source_output(token_ids: list[int], mm_output: dict | None = None):
+def _source_output(token_ids: list[int], mm_output: Mapping | None = None):
     """Create a minimal AR output mock."""
     return SimpleNamespace(
         outputs=[SimpleNamespace(token_ids=token_ids, cumulative_token_ids=token_ids)],
@@ -276,6 +280,84 @@ class TestParseGeneratedTokens:
 
 
 class TestAr2Diffusion:
+    @pytest.mark.parametrize("include_hidden", [False, True], ids=["tokens_only", "with_generic_hidden"])
+    def test_t2i_token_id_output_preserves_cumulative_tokens(self, include_hidden):
+        """DiT consumes the cumulative token sequence, not the generic tensor payload."""
+        output_type = GLM_IMAGE_PIPELINE.stages[0].engine_output_type
+        assert output_type == "token_ids"
+
+        # Default GPU execution omits hidden; some runner paths still attach it.
+        # The shared normalizer renames it without changing the tensor's values.
+        mm_output = None
+        if include_hidden:
+            hidden = torch.arange(6, dtype=torch.float32).reshape(2, 3)
+            mm_output = MultimodalPayload.from_raw({"hidden": hidden}, output_type)
+            assert mm_output is not None
+            assert set(mm_output) == {"token_ids"}
+            torch.testing.assert_close(mm_output["token_ids"], hidden)
+
+        # At 64x64 the AR sequence contains a 256-token preview, four target
+        # image tokens, and EOS. The final emitted delta contains only EOS.
+        token_ids = [9] * 256 + [1, 2, 3, 4] + [16385]
+        source_output = _source_output(token_ids, mm_output)
+        source_output.outputs[0].token_ids = [16385]
+        prompt = {"prompt": "a cat", "mm_processor_kwargs": {"target_h": 64, "target_w": 64}}
+
+        result = ar2diffusion([source_output], prompt=prompt)
+
+        assert result is not None
+        assert (result["prompt"], result["height"], result["width"]) == ("a cat", 64, 64)
+        torch.testing.assert_close(
+            result["extra"]["prior_token_ids"],
+            torch.tensor([1, 1, 2, 2, 1, 1, 2, 2, 3, 3, 4, 4, 3, 3, 4, 4]),
+        )
+        assert result["extra"]["prior_token_image_ids"] is None
+
+    @pytest.mark.parametrize("include_hidden", [False, True], ids=["explicit_ids_only", "with_generic_hidden"])
+    def test_i2i_token_id_output_preserves_source_image_tokens(self, include_hidden):
+        """Explicit image IDs survive normalization with or without hidden output."""
+        from PIL import Image
+
+        output_type = GLM_IMAGE_PIPELINE.stages[0].engine_output_type
+        assert output_type == "token_ids"
+
+        prior_image_ids = torch.tensor([5, 6, 7, 8])
+        raw_output = {"ids.prior_image": prior_image_ids}
+        hidden = torch.arange(6, dtype=torch.float32).reshape(2, 3)
+        if include_hidden:
+            raw_output["hidden"] = hidden
+        mm_output = MultimodalPayload.from_raw(raw_output, output_type)
+        assert mm_output is not None
+        # The output processor restores dotted runner keys before emitting
+        # the payload consumed by the stage input processor.
+        mm_output = MultimodalPayload.from_dict(unflatten_payload(mm_output.to_dict()))
+        assert mm_output is not None
+        assert set(mm_output) == ({"ids", "token_ids"} if include_hidden else {"ids"})
+        torch.testing.assert_close(mm_output["ids"]["prior_image"], prior_image_ids)
+        if include_hidden:
+            torch.testing.assert_close(mm_output["token_ids"], hidden)
+
+        source_output = _source_output([1, 2, 3, 4, 16385], mm_output)
+        source_output.outputs[0].token_ids = [16385]
+        image = Image.new("RGB", (64, 64))
+        prompt = {
+            "prompt": "edit this",
+            "mm_processor_kwargs": {"target_h": 64, "target_w": 64},
+            "multi_modal_data": {"image": image},
+        }
+
+        result = ar2diffusion([source_output], prompt=prompt, requires_multimodal_data=True)
+
+        assert result is not None
+        assert (result["prompt"], result["height"], result["width"]) == ("edit this", 64, 64)
+        assert result["pil_image"] is image
+        torch.testing.assert_close(
+            result["extra"]["prior_token_ids"],
+            torch.tensor([1, 1, 2, 2, 1, 1, 2, 2, 3, 3, 4, 4, 3, 3, 4, 4]),
+        )
+        assert len(result["extra"]["prior_token_image_ids"]) == 1
+        torch.testing.assert_close(result["extra"]["prior_token_image_ids"][0], prior_image_ids)
+
     def test_basic_t2i(self):
         """Test basic text-to-image pipeline: AR -> Diffusion."""
         # 1024x1024 t2i: small(256) + large(1024) + EOS

tests/model_extras/test_mammothmoda2.py

@@ -0,0 +1,72 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import pytest
+
+from vllm_omni.model_extras import build_text_to_image_prompt
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.cpu,
+    pytest.mark.parametrize(
+        "model_class_name",
+        ["MammothModa2DiTPipeline", "MammothModa2ForConditionalGeneration", "Mammothmoda2Model"],
+    ),
+]
+
+
+@pytest.mark.parametrize(
+    ("height", "width", "expected_height", "expected_width"),
+    [
+        (None, None, 1024, 1024),
+        (None, 512, 1024, 512),
+        (512, None, 512, 1024),
+        (16, 16, 16, 16),
+        (16, 32, 16, 32),
+        (32, 16, 32, 16),
+        (256, 256, 256, 256),
+        (512, 768, 512, 768),
+        (768, 512, 768, 512),
+    ],
+)
+def test_text_to_image_dimensions(
+    model_class_name: str, height: int | None, width: int | None, expected_height: int, expected_width: int
+) -> None:
+    result = build_text_to_image_prompt(
+        model_class_name,
+        {"prompt": "a cat", "modalities": ["image"]},
+        height=height,
+        width=width,
+    )
+
+    info = result["additional_information"]
+    assert info["image_height"] == [expected_height]
+    assert info["image_width"] == [expected_width]
+    assert info["ar_height"] == [expected_height // 16]
+    assert info["ar_width"] == [expected_width // 16]
+    assert result["prompt"].endswith(f"<|image start|>{expected_width // 16}*{expected_height // 16}<|image token|>")
+
+
+@pytest.mark.parametrize("dimension", ["height", "width"])
+@pytest.mark.parametrize("value", [-16, 0, 1, 15, 17, 513])
+def test_text_to_image_rejects_invalid_dimensions(model_class_name: str, dimension: str, value: int) -> None:
+    dimensions = {"height": 512, "width": 512}
+    dimensions[dimension] = value
+
+    with pytest.raises(ValueError, match=rf"MammothModa2 {dimension} must be a positive multiple of 16, got {value}"):
+        build_text_to_image_prompt(model_class_name, {"prompt": "a cat", "modalities": ["image"]}, **dimensions)
+
+
+@pytest.mark.parametrize("dimension", ["height", "width"])
+@pytest.mark.parametrize("value", [0, 513])
+def test_text_to_image_rejects_invalid_dimension_when_other_is_omitted(
+    model_class_name: str, dimension: str, value: int
+) -> None:
+    with pytest.raises(ValueError, match=rf"MammothModa2 {dimension} must be a positive multiple of 16, got {value}"):
+        build_text_to_image_prompt(
+            model_class_name,
+            {"prompt": "a cat", "modalities": ["image"]},
+            **{dimension: value},
+        )

tests/protocol/realtime/test_realtime_codec_single_source.py

@@ -218,14 +218,17 @@ def test_the_error_code_vocabulary_is_tier3() -> None:
     assert not hasattr(realtime_errors, "REALTIME_ERROR_TYPES_BY_CODE")
 
 
-def test_the_tier1_error_event_reports_only_openai_classes() -> None:
+@pytest.mark.parametrize("code", ["resource_exhausted", "output_backpressure"])
+def test_the_tier1_error_event_reports_only_openai_classes(code: str) -> None:
     from vllm_omni.protocol.duplex import events as duplex_wire_events
     from vllm_omni.protocol.realtime import events as realtime_events
 
     # Tier 1 knows the envelope shape but not our code vocabulary.
-    assert realtime_events.ErrorEvent(code="resource_exhausted").error_type == "invalid_request_error"
+    assert realtime_events.ErrorEvent(code=code).error_type == "invalid_request_error"
     # Tier 2 resolves it through our table.
-    assert duplex_wire_events.ErrorEvent(code="resource_exhausted").error_type == "rate_limit_error"
+    error = duplex_wire_events.ErrorEvent(code=code)
+    assert error.error_type == "rate_limit_error"
+    assert error.to_realtime()["error"]["type"] == "rate_limit_error"
 
 
 @pytest.mark.parametrize("name", _TIER1_COMMANDS + _TIER2_COMMANDS + _TIER3_COMMANDS)

tests/tools/test_sheetsage2_transcribe.py

@@ -0,0 +1,193 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Contracts for the standalone audio-to-score/YuE2 handoff tool."""
+
+import json
+import sys
+from pathlib import Path
+
+import pytest
+
+from tools import sheetsage2_transcribe as cli
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.fixture
+def arguments(tmp_path):
+    audio = tmp_path / "song.wav"
+    audio.write_bytes(b"audio is decoded by the upstream transcriber")
+    return cli.create_parser().parse_args([str(audio), "--output-dir", str(tmp_path / "score"), "--trust-remote-code"])
+
+
+@pytest.fixture
+def backend(mocker):
+    model = mocker.Mock()
+    model.config = mocker.Mock(
+        spec_set=["model_type", "base_model_revision"], model_type="sheetsage2", base_model_revision="parent-sha"
+    )
+    model.eval.return_value = model
+    model.to.return_value = model
+    model.transcribe.return_value = {"abc": "X:1\nM:4/4\nK:C\nC4|\n", "midi": b"MThd\x00\x00\x00\x06"}
+    loader = mocker.Mock(return_value=model)
+    torch = mocker.Mock(
+        float32=object(),
+        __version__="test",
+        get_num_threads=lambda: 8,
+        set_num_threads=mocker.Mock(),
+        accelerator=mocker.Mock(is_available=lambda: False),
+    )
+    mocker.patch.dict(
+        sys.modules,
+        {"torch": torch, "transformers": mocker.Mock(AutoModel=mocker.Mock(from_pretrained=loader))},
+    )
+    return mocker.Mock(model=model, loader=loader, torch=torch)
+
+
+@pytest.mark.parametrize("melody_only,expected_cot", [(True, "melody"), (False, "full")])
+def test_score_handoff_preserves_unicode_and_selects_matching_cot(arguments, backend, melody_only, expected_cot):
+    arguments.lyrics_file = arguments.audio.with_suffix(".txt")
+    arguments.lyrics_file.write_text("[Verse]\n江边的清晨", encoding="utf-8")
+    arguments.style = "  Chinese folk, bamboo flute  "
+    arguments.yue2_model = "music-server"
+    arguments.melody_only = melody_only
+    arguments.seed = 123
+
+    output = cli.transcribe(arguments)
+
+    request = json.loads((output / "yue2_request.json").read_text())
+    assert request == {
+        "model": "music-server",
+        "input": "[Verse]\n江边的清晨",
+        "instructions": "Chinese folk, bamboo flute",
+        "response_format": "wav",
+        "stream": False,
+        "seed": 123,
+        "extra_params": {"cot": expected_cot, "abc": backend.model.transcribe.return_value["abc"]},
+    }
+    assert backend.model.transcribe.call_args.kwargs["melody_only"] == melody_only
+    assert "max_new_tokens" not in request  # Leave natural song length to the server.
+
+
+def test_default_model_pins_code_and_loads_adapters_in_fp32(arguments, backend):
+    output = cli.transcribe(arguments)
+    options = backend.loader.call_args.kwargs
+    assert options["revision"] == options["code_revision"] == cli.DEFAULT_REVISION
+    assert options["torch_dtype"] is backend.torch.float32
+    backend.model.to.assert_called_once_with("cpu")
+    assert not (output / "yue2_request.json").exists()
+
+
+def test_local_offline_checkpoint_forwards_parent_without_default_hub_revision(arguments, backend, tmp_path, mocker):
+    arguments.model = str(tmp_path / "SheetSage2")
+    snapshot = Path(arguments.model)
+    snapshot.mkdir()
+    source = snapshot / "modeling_sheetsage2.py"
+    direct = snapshot / "exports_sheetsage2.py"
+    transitive = snapshot / "chord_spelling_sheetsage2.py"
+    source.write_text("from .exports_sheetsage2 import export\n")
+    direct.write_text("from .chord_spelling_sheetsage2 import chord\n")
+    transitive.write_text("chord = 'C'\n")
+    dynamic_modules = mocker.Mock()
+    dynamic_modules.get_relative_import_files.return_value = [str(direct), str(transitive)]
+    mocker.patch.dict(sys.modules, {"transformers.dynamic_module_utils": dynamic_modules})
+    calls = mocker.Mock()
+    calls.attach_mock(dynamic_modules.get_cached_module_file, "prime")
+    calls.attach_mock(backend.loader, "load")
+    arguments.base_model_path = tmp_path / "MERT"
+    arguments.base_model_path.mkdir()
+    arguments.local_files_only = True
+    cli.transcribe(arguments)
+    dynamic_modules.get_relative_import_files.assert_called_once_with(str(source))
+    dynamic_modules.get_cached_module_file.assert_has_calls(
+        [
+            mocker.call(arguments.model, direct.name, local_files_only=True),
+            mocker.call(arguments.model, transitive.name, local_files_only=True),
+        ]
+    )
+    assert [call[0] for call in calls.mock_calls] == ["prime", "prime", "load"]
+    backend.model.transcribe.assert_called_once()
+    options = backend.loader.call_args.kwargs
+    assert options["revision"] is None
+    assert options["local_files_only"] is True
+    assert options["base_model_path"] == str(arguments.base_model_path)
+
+
+@pytest.mark.parametrize("limit", [0, -1, float("inf"), float("nan")])
+def test_invalid_duration_fails_before_model_loading(arguments, backend, limit):
+    arguments.max_seconds = limit
+    with pytest.raises(ValueError, match="finite and positive"):
+        cli.transcribe(arguments)
+    backend.loader.assert_not_called()
+    assert not arguments.output_dir.exists()
+
+
+def test_nonempty_output_is_not_reused(arguments, backend):
+    arguments.output_dir.mkdir()
+    old = arguments.output_dir / "yue2_request.json"
+    old.write_text("old request")
+    with pytest.raises(ValueError, match="new or empty"):
+        cli.transcribe(arguments)
+    assert old.read_text() == "old request"
+    backend.loader.assert_not_called()
+
+
+def test_explicit_code_trust_is_required(arguments, backend):
+    arguments.trust_remote_code = False
+    with pytest.raises(ValueError, match="--trust-remote-code"):
+        cli.transcribe(arguments)
+    backend.loader.assert_not_called()
+
+
+@pytest.mark.parametrize(
+    "result,message",
+    [
+        ({"abc": "", "midi": b"MThd"}, "ABC export failed"),
+        ({"abc": "X:1", "abc_error": "invalid beat grid", "midi": b"MThd"}, "invalid beat grid"),
+        ({"abc": "X:1", "midi": b"invalid"}, "MIDI export failed"),
+    ],
+)
+def test_failed_export_never_writes_a_yue2_request(arguments, backend, result, message):
+    arguments.lyrics_file = arguments.audio.with_suffix(".txt")
+    arguments.lyrics_file.write_text("lyrics")
+    arguments.style = "folk"
+    backend.model.transcribe.return_value = result
+    with pytest.raises(RuntimeError, match=message):
+        cli.transcribe(arguments)
+    assert not (arguments.output_dir / "yue2_request.json").exists()
+    assert not (arguments.output_dir / "manifest.json").exists()
+
+
+def test_style_requires_lyrics_before_loading(arguments, backend):
+    arguments.style = "folk"
+    with pytest.raises(ValueError, match="supplied together"):
+        cli.transcribe(arguments)
+    backend.loader.assert_not_called()
+
+
+def test_seed_requires_request_context_before_loading(arguments, backend):
+    arguments.seed = 123
+    with pytest.raises(ValueError, match="supply --lyrics-file and --style"):
+        cli.transcribe(arguments)
+    backend.loader.assert_not_called()
+    assert not arguments.output_dir.exists()
+
+
+@pytest.mark.parametrize("seed", [-1, 2**63])
+def test_out_of_range_seed_fails_before_loading(arguments, backend, seed):
+    arguments.lyrics_file = arguments.audio.with_suffix(".txt")
+    arguments.lyrics_file.write_text("lyrics", encoding="utf-8")
+    arguments.style = "folk"
+    arguments.seed = seed
+    with pytest.raises(ValueError, match="between 0 and"):
+        cli.transcribe(arguments)
+    backend.loader.assert_not_called()
+    assert not arguments.output_dir.exists()
+
+
+def test_cli_reports_export_failure_as_nonzero(arguments, backend, capsys):
+    backend.model.transcribe.return_value = {"abc": "", "midi": b"MThd"}
+    with pytest.raises(SystemExit) as exc:
+        cli.main([str(arguments.audio), "--output-dir", str(arguments.output_dir), "--trust-remote-code"])
+    assert exc.value.code == 1
+    assert "ABC export failed" in capsys.readouterr().err

tests/worker/test_gpu_ar_model_runner.py

@@ -15,6 +15,8 @@ import numpy as np
 import pytest
 import torch
 from vllm.sampling_params import SamplingParams
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.outputs import SamplerOutput
 from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.utils import CpuGpuBuffer
 from vllm.v1.worker import gpu_input_batch
@@ -25,7 +27,7 @@ from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechReques
 from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
 from vllm_omni.entrypoints.openai.tts_adapters.base import PreparedRequest
 from vllm_omni.outputs import OmniModelRunnerOutput
-from vllm_omni.worker import sparse_audio
+from vllm_omni.worker import gpu_ar_model_runner, sparse_audio
 from vllm_omni.worker.gpu_ar_model_runner import (
     ExecuteModelState,
     GPUARModelRunner,
@@ -614,8 +616,10 @@ def test_build_omni_output_copies_hidden_for_partial_downstream_batch(monkeypatc
     assert output.multimodal_outputs is None
 
 
-def test_process_additional_information_uses_snapshot_request_order(monkeypatch):
+@pytest.mark.parametrize("include_hidden", [True, False])
+def test_process_additional_information_uses_snapshot_request_order(monkeypatch, include_hidden):
     runner = _make_async_output_runner()
+    runner._pooler_payload_include_hidden_flag = include_hidden
     seen = []
 
     class PostprocessModel:
@@ -958,6 +962,11 @@ def test_sample_tokens_tail_only_prefix_cache_uses_staged_cpu_hidden_states(monk
     monkeypatch.setattr(GPUARModelRunner, "_should_accumulate_full_payload_output", lambda self: False)
     monkeypatch.setattr(GPUARModelRunner, "get_omni_connector_output", lambda self: None)
 
+    monkeypatch.setattr(
+        GPUARModelRunner,
+        "_snapshot_scheduler_output_for_async_omni_output",
+        lambda *args: pytest.fail("inline output must not snapshot scheduler metadata"),
+    )
     output = GPUARModelRunner.sample_tokens(runner, grammar_output=None)
 
     # Non-async-chunk now ships the full payload to the next stage, so
@@ -1946,3 +1955,164 @@ def test_prepare_hook_keeps_legacy_signature_and_orders_opt_in_metadata(accepts_
         assert seen["discard"] == [True, False]
     else:
         assert set(seen) == {"req_ids"}
+
+
+@pytest.fixture
+def token_only_builder(monkeypatch):
+    runner = _make_async_output_runner()
+    runner._pooler_payload_include_hidden_flag = False
+    runner.supports_mm_inputs = True
+    monkeypatch.setattr(GPUARModelRunner, "_resolve_pooler_payload_req_ids", lambda self, req_ids: ("latent", req_ids))
+    scheduler = SchedulerOutput.make_empty()
+    scheduler.total_num_scheduled_tokens = 2
+    scheduler.num_scheduled_tokens = {"r1": 1, "r2": 1}
+    kwargs = dict(
+        scheduler_output=scheduler,
+        hidden_states=torch.tensor([[1.0], [2.0]]),
+        staged_hidden_states_cpu=None,
+        multimodal_outputs=None,
+        req_ids_output_copy=["r1", "r2"],
+        req_id_to_index_output_copy={"r1": 0, "r2": 1},
+        valid_sampled_token_ids=[[11], [22]],
+        logprobs_lists=object(),
+        prompt_logprobs_dict={"r1": object()},
+        num_nans_in_logits={"r1": 0},
+        kv_connector_output=object(),
+        ec_connector_output=object(),
+        cudagraph_stats=object(),
+        kv_extracted_req_ids=["r2"],
+        num_scheduled_tokens_np=np.array([1, 1], dtype=np.int32),
+        query_start_loc_cpu=torch.tensor([0, 1], dtype=torch.int32),
+    )
+    return runner, kwargs
+
+
+@pytest.mark.parametrize("empty_mm", [None, {}])
+def test_token_only_output_preserves_sampling_and_connector_metadata(token_only_builder, mocker, empty_mm):
+    runner, kwargs = token_only_builder
+    kwargs["multimodal_outputs"] = empty_mm
+    payload_builder = mocker.spy(runner, "_build_omni_pooler_payload")
+    output = runner._build_omni_model_runner_output_from_snapshot(**kwargs)
+
+    assert output.req_ids == ["r1", "r2"]
+    assert output.req_id_to_index == {"r1": 0, "r2": 1}
+    assert output.sampled_token_ids == [[11], [22]]
+    assert output.inter_stage_outputs is None
+    assert output.multimodal_outputs is None
+    assert output.pooler_output is None
+    assert output.logprobs is kwargs["logprobs_lists"]
+    for name in (
+        "prompt_logprobs_dict",
+        "num_nans_in_logits",
+        "kv_connector_output",
+        "ec_connector_output",
+        "cudagraph_stats",
+        "kv_extracted_req_ids",
+    ):
+        assert getattr(output, name) is kwargs[name]
+    payload_builder.assert_not_called()
+
+
+def test_token_only_output_keeps_prefix_cache_merge(token_only_builder, mocker):
+    runner, kwargs = token_only_builder
+    runner.omni_prefix_cache = object()
+    cached = {"codes.audio": {"r1": torch.tensor([31]), "r2": torch.tensor([32])}}
+    merge = mocker.patch.object(
+        runner, "_prepare_prefix_cache_pooler_payload_sources", return_value=(None, None, cached)
+    )
+    output = runner._build_omni_model_runner_output_from_snapshot(**kwargs)
+
+    merge.assert_called_once()
+    assert output.inter_stage_outputs[0]["codes.audio"].tolist() == [31]
+    assert output.inter_stage_outputs[1]["codes.audio"].tolist() == [32]
+
+
+@pytest.mark.parametrize("snapshot", [True, False], ids=["background", "inline"])
+def test_query_metadata_lifetime_matches_materialization_mode(snapshot):
+    runner = object.__new__(GPUARModelRunner)
+    runner.query_start_loc = object.__new__(CpuGpuBuffer)
+    runner.query_start_loc.cpu = torch.tensor([0, 2, 3], dtype=torch.int32)
+    offsets = runner._get_query_start_loc_cpu(snapshot=snapshot)
+    runner.query_start_loc.cpu[1] = 9
+
+    if snapshot:
+        assert offsets.tolist() == [0, 2, 3]
+    else:
+        assert offsets is runner.query_start_loc.cpu
+        assert offsets.tolist() == [0, 9, 3]
+
+
+def test_background_output_survives_step_metadata_reuse(monkeypatch, mocker):
+    """Delay the real builder until offsets, batch IDs and counts are reused."""
+
+    runner = _make_async_output_runner()
+    runner.use_async_scheduling = True
+    runner.speculative_config = None
+    runner.model.use_async_omni_output = True
+    runner.broadcast_pp_output = False
+    runner.kv_connector_output = None
+    runner.device = torch.device("cpu")
+    runner.async_output_copy_stream = None
+    runner.input_batch = mocker.Mock(
+        spec=InputBatch, req_ids=["r1", "r2"], req_id_to_index={"r1": 0, "r2": 1}, vocab_size=10
+    )
+    runner.query_start_loc = object.__new__(CpuGpuBuffer)
+    runner.query_start_loc.cpu = torch.tensor([0, 1, 3], dtype=torch.int32)
+    runner._omni_num_scheduled_tokens_np = np.array([1, 2], dtype=np.int32)
+    scheduler = SchedulerOutput.make_empty()
+    scheduler.num_scheduled_tokens = {"r1": 1, "r2": 2}
+    scheduler.total_num_scheduled_tokens = 3
+    hidden = torch.tensor([[1.0], [2.0], [3.0]])
+    runner.execute_model_state = ExecuteModelState(
+        scheduler, None, None, None, hidden, None, None, None, None, None, {}, None
+    )
+    sampler = SamplerOutput(sampled_token_ids=torch.tensor([[7], [8]]), logprobs_tensors=None)
+    mocker.patch.object(runner, "_sample", return_value=sampler)
+    mocker.patch.object(runner, "_update_states_after_model_execute")
+    # Match upstream bookkeeping's ownership contract: these are already copies.
+    mocker.patch.object(
+        runner,
+        "_bookkeeping_sync",
+        return_value=(
+            {},
+            None,
+            None,
+            [],
+            {},
+            runner.input_batch.req_ids.copy(),
+            runner.input_batch.req_id_to_index.copy(),
+            [],
+        ),
+    )
+    mocker.patch.object(runner, "_run_post_sample_talker_mtp", return_value={})
+    mocker.patch.object(runner, "eplb_step")
+    mocker.patch.object(runner, "get_omni_connector_output", return_value=None)
+    monkeypatch.setattr(gpu_ar_model_runner, "get_pp_group", lambda: mocker.Mock(world_size=1))
+    mocker.patch.object(runner, "_resolve_pooler_payload_req_ids", side_effect=lambda ids: ("audio", ids))
+    # Device copying is covered separately; this test isolates metadata lifetime.
+    mocker.patch.object(
+        runner,
+        "_snapshot_omni_output_tensors_for_async_output",
+        return_value=gpu_ar_model_runner._OmniOutputTensorSnapshot(
+            hidden_states=hidden, staged_hidden_states_cpu=None, multimodal_outputs={}
+        ),
+    )
+    async_output = mocker.Mock(
+        spec=OmniAsyncGPUModelRunnerOutput,
+        sampled_token_ids_cpu=sampler.sampled_token_ids,
+        async_copy_ready_event=object(),
+    )
+    constructor = mocker.patch.object(gpu_ar_model_runner, "OmniAsyncGPUModelRunnerOutput", return_value=async_output)
+
+    assert runner.sample_tokens(None) is async_output
+    runner.query_start_loc.cpu.fill_(99)
+    runner._omni_num_scheduled_tokens_np.fill(99)
+    scheduler.num_scheduled_tokens.clear()
+    runner.input_batch.req_ids[:] = ["next-step"]
+    runner.input_batch.req_id_to_index.clear()
+    output = constructor.call_args.kwargs["model_runner_output_builder"]()
+
+    assert output.req_ids == ["r1", "r2"]
+    assert output.req_id_to_index == {"r1": 0, "r2": 1}
+    torch.testing.assert_close(output.inter_stage_outputs[0]["hidden"], hidden[:1])
+    torch.testing.assert_close(output.inter_stage_outputs[1]["hidden"], hidden[1:])

tests/worker_v2/test_omni_ar_model_runner.py

@@ -52,7 +52,7 @@ def _async_output(req_ids=("req-0",), **overrides) -> OmniAsyncOutput:
 
 
 @pytest.mark.parametrize("compact_width", [1, 2])
-def test_async_output_blocking_event_and_routing_masks(monkeypatch, compact_width) -> None:
+def test_async_output_blocking_event_preserves_masks_and_aux_output(monkeypatch, compact_width) -> None:
     event_kwargs = []
     monkeypatch.setattr(torch.cuda, "set_stream", lambda _stream: None)
 
@@ -68,7 +68,7 @@ def test_async_output_blocking_event_and_routing_masks(monkeypatch, compact_widt
         vocab_size=4,
     )
     sampler_output = SamplerOutput(
-        torch.tensor([[2], [0]]), None, None, torch.tensor([1, 0]), torch.tensor([0, 0]), masks
+        torch.tensor([[2], [0]]), None, None, torch.tensor([1, 0]), torch.tensor([2, 0]), masks
     )
     copied: list[tuple[np.ndarray, np.ndarray]] = []
 
@@ -98,7 +98,9 @@ def test_async_output_blocking_event_and_routing_masks(monkeypatch, compact_widt
     np.testing.assert_array_equal(output.aux_output_connector_output["prefill"].rows, [[4, 5]])
     assert output.aux_output_connector_output["decode"].token_start == 0
     np.testing.assert_array_equal(output.sampling_masks.token_ids, [0, 2])
-    assert len(copied) == 1 and copied[0][0].tolist() == [1, 0]
+    assert len(copied) == 1
+    np.testing.assert_array_equal(copied[0][0], [1, 0])
+    np.testing.assert_array_equal(copied[0][1], [2, 0])
 
 
 @pytest.mark.parametrize("needs_history", [False, True])

tests/worker_v2/test_omni_model_state.py

@@ -603,8 +603,13 @@ def test_eager_decode_without_a_frame_fails_loudly():
         state._eager_state._apply_eager_frames(batches, torch.zeros((1, _EAGER_DIM)), _EagerBatch([1]), None)
 
 
-def test_run_preprocess_records_rows_that_keep_a_sample():
+@pytest.mark.parametrize("has_stream_decoder", [True, False])
+def test_run_preprocess_records_rows_that_keep_a_sample(has_stream_decoder):
     state = _make_eager_state()
+    if not has_stream_decoder:
+        # Eager-MTP talkers without an in-Talker codec (Qwen3-Omni) do not
+        # define the attribute at all.
+        del state.model.stream_decoder
     _fill_buffers(state, "chunk", "final", "decode")
     state._eager_ready = {2: "decode"}
     state._eager_embeds[2] = 4.0

tools/requirements/sheetsage2.txt

@@ -0,0 +1,13 @@
+# Isolated preprocessing environment; do not install into the Omni server venv.
+# Matches m-a-p/SheetSage2 at cafc0df1021e14f49e928c4b345f5959d414ef64.
+torch==2.8.0
+torchaudio==2.8.0
+transformers==4.45.2
+huggingface-hub==0.36.0
+safetensors==0.5.3
+numpy==1.24.3
+scipy==1.13.1
+mir_eval==0.8.2
+pretty_midi==0.2.10
+mido==1.3.3
+setuptools==78.1.1

tools/sheetsage2_transcribe.py

@@ -0,0 +1,178 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Convert music to a score and optionally prepare a YuE2 speech request.
+
+Run this preprocessing tool in the separate environment documented in
+recipes/m-a-p/SheetSage2-H200.md. It uses the upstream transcriber, not an Omni
+engine, and does not import vLLM or change the serving environment.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import math
+from pathlib import Path
+
+DEFAULT_MODEL = "m-a-p/SheetSage2"
+DEFAULT_REVISION = "cafc0df1021e14f49e928c4b345f5959d414ef64"
+
+
+def create_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("audio", type=Path, help="Local music recording (decoded by FFmpeg)")
+    parser.add_argument("--output-dir", required=True, type=Path, help="New or empty output directory")
+    parser.add_argument("--model", default=DEFAULT_MODEL, help="SheetSage2 Hub ID or local snapshot")
+    parser.add_argument("--revision", help="Model/code revision; the default Hub model uses a pinned commit")
+    parser.add_argument("--base-model-path", type=Path, help="Local MERT-v2 parent for offline loading")
+    parser.add_argument("--local-files-only", action="store_true")
+    parser.add_argument(
+        "--trust-remote-code", action="store_true", help="Allow the upstream model's custom Python code"
+    )
+    parser.add_argument("--device", default="auto", help="auto, cpu, cuda:0, etc.")
+    parser.add_argument(
+        "--dtype", choices=("bf16", "fp32"), default="bf16", help="Inference autocast; weights load in FP32"
+    )
+    parser.add_argument("--melody-only", action="store_true", help="Retain vocal/instrumental melodies and omit chords")
+    parser.add_argument(
+        "--max-seconds", type=float, help="Transcribe only this many seconds; otherwise process the song"
+    )
+    parser.add_argument("--lyrics-file", type=Path, help="With --style, write yue2_request.json for /v1/audio/speech")
+    parser.add_argument("--style", help="Target musical arrangement for the optional YuE2 request")
+    parser.add_argument("--yue2-model", default="m-a-p/YuE2-3B", help="Served model name in the optional request")
+    parser.add_argument("--seed", type=int, help="Optional YuE2 generation seed")
+    return parser
+
+
+def validate_args(args: argparse.Namespace) -> str | None:
+    if not args.audio.is_file():
+        raise ValueError(f"Audio file does not exist: {args.audio}")
+    if not args.trust_remote_code:
+        raise ValueError(
+            "SheetSage2 uses custom model code; pass --trust-remote-code after reviewing the model repository"
+        )
+    if args.max_seconds is not None and (not math.isfinite(args.max_seconds) or args.max_seconds <= 0):
+        raise ValueError("--max-seconds must be finite and positive")
+    if args.output_dir.exists() and (not args.output_dir.is_dir() or any(args.output_dir.iterdir())):
+        raise ValueError("--output-dir must be a new or empty directory")
+    if args.base_model_path is not None and not args.base_model_path.is_dir():
+        raise ValueError(f"MERT-v2 directory does not exist: {args.base_model_path}")
+    if (args.lyrics_file is None) != (args.style is None):
+        raise ValueError("--lyrics-file and --style must be supplied together")
+    if args.lyrics_file is None:
+        if args.seed is not None:
+            raise ValueError("--seed applies to the YuE2 request; supply --lyrics-file and --style")
+        return None
+    lyrics = args.lyrics_file.read_text(encoding="utf-8").strip()
+    if not lyrics or not args.style.strip():
+        raise ValueError("Lyrics and --style must not be blank")
+    if not args.yue2_model.strip():
+        raise ValueError("--yue2-model must not be blank")
+    if args.seed is not None and not 0 <= args.seed < 2**63:
+        raise ValueError("--seed must be between 0 and 2**63 - 1")
+    return lyrics
+
+
+def build_yue2_request(abc: str, lyrics: str, args: argparse.Namespace) -> dict[str, object]:
+    request: dict[str, object] = {
+        "model": args.yue2_model,
+        "input": lyrics,
+        "instructions": args.style.strip(),
+        "response_format": "wav",
+        "stream": False,
+        "extra_params": {"cot": "melody" if args.melody_only else "full", "abc": abc},
+    }
+    if args.seed is not None:
+        request["seed"] = args.seed
+    return request
+
+
+def transcribe(args: argparse.Namespace) -> Path:
+    # Validate before loading dependencies/weights or touching an existing output.
+    lyrics = validate_args(args)
+    import torch
+    from transformers import AutoModel
+
+    if Path(args.model).is_dir():
+        # Transformers 4.45 only copies direct imports for local snapshots.
+        # Prime transitive dependencies too (e.g. chord_spelling_sheetsage2),
+        # using its cache API without modifying the upstream source files.
+        # The pinned snapshot keeps the entry point and all relative imports
+        # as sibling .py files; local snapshots must preserve that flat layout.
+        from transformers.dynamic_module_utils import get_cached_module_file, get_relative_import_files
+
+        source = Path(args.model) / "modeling_sheetsage2.py"
+        for dependency in get_relative_import_files(str(source)):
+            get_cached_module_file(args.model, Path(dependency).name, local_files_only=True)
+
+    torch.set_num_threads(min(4, torch.get_num_threads()))
+    device = args.device
+    if device == "auto":
+        device = torch.accelerator.current_accelerator().type if torch.accelerator.is_available() else "cpu"
+    revision = args.revision
+    if revision is None and args.model == DEFAULT_MODEL:
+        revision = DEFAULT_REVISION
+    load_options = {
+        "trust_remote_code": args.trust_remote_code,
+        "revision": revision,
+        "code_revision": revision,
+        "local_files_only": args.local_files_only,
+        # Upstream merges the MERT adapters before any reduced-precision cast.
+        "torch_dtype": torch.float32,
+    }
+    if args.base_model_path is not None:
+        load_options["base_model_path"] = str(args.base_model_path)
+    model = AutoModel.from_pretrained(args.model, **load_options).eval().to(device)
+    if model.config.model_type != "sheetsage2":
+        raise ValueError("--model must identify a SheetSage2 checkpoint")
+    args.output_dir.mkdir(parents=True, exist_ok=True)
+    result = model.transcribe(
+        str(args.audio),
+        output_dir=str(args.output_dir),
+        melody_only=args.melody_only,
+        dtype=args.dtype,
+        max_seconds=args.max_seconds,
+    )
+    abc = result.get("abc")
+    if result.get("abc_error") or not isinstance(abc, str) or not abc.strip():
+        raise RuntimeError(
+            f"ABC export failed: {result.get('abc_error') or 'empty score'}; no YuE2 request was written"
+        )
+    midi = result.get("midi")
+    if not isinstance(midi, bytes) or not midi.startswith(b"MThd"):
+        raise RuntimeError("MIDI export failed; no YuE2 request was written")
+    if lyrics is not None:
+        request = build_yue2_request(abc, lyrics, args)
+        (args.output_dir / "yue2_request.json").write_text(
+            json.dumps(request, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
+        )
+    manifest = {
+        "model": args.model,
+        "revision": revision,
+        "base_model_revision": getattr(model.config, "base_model_revision", None),
+        "audio": str(args.audio.resolve()),
+        "melody_only": args.melody_only,
+        "max_seconds": args.max_seconds,
+        "device": str(device),
+        "inference_dtype": args.dtype,
+        "torch_version": torch.__version__,
+        "backend": "upstream SheetSage2 transcribe",
+    }
+    (args.output_dir / "manifest.json").write_text(json.dumps(manifest, indent=2) + "\n", encoding="utf-8")
+    return args.output_dir
+
+
+def main(argv: list[str] | None = None) -> None:
+    parser = create_parser()
+    args = parser.parse_args(argv)
+    try:
+        output_dir = transcribe(args)
+    except (ValueError, RuntimeError, OSError, ImportError) as exc:
+        parser.exit(1, f"SheetSage2: {exc}\nSee recipes/m-a-p/SheetSage2-H200.md for the isolated environment setup.\n")
+    print(f"Saved score.abc, transcription.mid and annotations to {output_dir.resolve()}")
+    if args.lyrics_file is not None:
+        print(f"YuE2 request: {output_dir.resolve() / 'yue2_request.json'}")
+
+
+if __name__ == "__main__":
+    main()

vllm_omni/config/model.py

@@ -105,7 +105,7 @@ class OmniModelConfig(ModelConfig):
          worker_type: Model Type, e.g., "ar" or "generation"
          engine_output_type: Optional output type specification for the engine.
              Used to route outputs to appropriate processors (e.g., "image",
-             "audio", "latents"). If None, output type is inferred.
+             "audio", "latent", "token_ids"). If None, output type is inferred.
          stage_connector_config: Stage connector configuration dictionary.
              Contains "name" (connector name), "extra" (extra connector config).
          task_type: Model-defined startup task type. Each model validates its

vllm_omni/config/omni_config.py

@@ -1145,6 +1145,9 @@ _DIFFUSION_MOVED_SHARED_FIELDS = frozenset(
         "disable_autocast",
     }
 )
+# Runtime-populated OmniDiffusionConfig state (init=False) that is never
+# user-configurable and therefore not part of the projection.
+_DIFFUSION_INTERNAL_FIELDS = frozenset({"ray_worker_env"})
 
 
 _STAGE_DEPLOY_ENGINE_FIELDS: tuple[str, ...] = tuple(_STAGE_DEPLOY_FIELDS)

vllm_omni/config/stage_config.py

@@ -536,6 +536,8 @@ class DuplexSessionRuntimeConfig:
     resume_replay_max_bytes_per_session: int = 8 * 1024 * 1024
     max_pending_input_bytes_per_session: int = 16 * 1024 * 1024
     max_pending_turns_per_session: int = 4
+    max_pending_output_bytes_per_session: int = 2 * 1024 * 1024
+    max_pending_output_events_per_session: int = 512
     max_sessions: int = 1
     # Unread by the plugin framework. It used to bound the per-session
     # completed-append table that made a retried append RPC submit once; the
@@ -561,6 +563,8 @@ class DuplexSessionRuntimeConfig:
             "resume_replay_max_bytes_per_session": self.resume_replay_max_bytes_per_session,
             "max_pending_input_bytes_per_session": self.max_pending_input_bytes_per_session,
             "max_pending_turns_per_session": self.max_pending_turns_per_session,
+            "max_pending_output_bytes_per_session": self.max_pending_output_bytes_per_session,
+            "max_pending_output_events_per_session": self.max_pending_output_events_per_session,
             "max_sessions": self.max_sessions,
             "completed_append_cache_size": self.completed_append_cache_size,
         }
@@ -1046,7 +1050,7 @@ def _build_engine_args(
     engine_args["retains_state_across_chunks"] = ps.retains_state_across_chunks
     if ps.execution_type == StageExecutionType.DIFFUSION and ps.model_arch:
         engine_args.setdefault("model_class_name", ps.model_arch)
-    if ps.engine_output_type:
+    if ps.engine_output_type is not None:
         engine_args["engine_output_type"] = ps.engine_output_type
     if next_stage_proc:
         engine_args["custom_process_next_stage_input_func"] = next_stage_proc

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -432,6 +432,9 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         pooler_outputs = model_runner_output.pooler_output
         mm_outputs = getattr(model_runner_output, "multimodal_outputs", None)
         inter_stage_outputs = getattr(model_runner_output, "inter_stage_outputs", None)
+        # Token-only processors need sampled IDs even without a tensor payload.
+        processor = getattr(getattr(self, "chunk_transfer_adapter", None), "custom_process_next_stage_input_func", None)
+        requires_token_updates = getattr(processor, "requires_token_updates", False) is True
         num_nans_in_logits = model_runner_output.num_nans_in_logits
         kv_connector_output = model_runner_output.kv_connector_output
         ec_connector_output = getattr(model_runner_output, "ec_connector_output", None)
@@ -689,7 +692,9 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
             # Only when this step might save. Read additional_information
             # before _free_request rewrites it.
             omits_chunk_transfer = False
-            if self.chunk_transfer_adapter is not None and (inter_stage_output is not None or stopped):
+            if self.chunk_transfer_adapter is not None and (
+                inter_stage_output is not None or (new_token_ids and requires_token_updates) or stopped
+            ):
                 omits_chunk_transfer = self._request_omits_chunk_transfer_to_next_stage(request)
             # Capture before resumable stop handling can clear token history.
             output_token_ids: Any = getattr(request, "output_token_ids", None)
@@ -824,7 +829,12 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
             if (
                 self.chunk_transfer_adapter is not None
                 and not omits_chunk_transfer
-                and (inter_stage_output is not None or is_segment_finished or finished)
+                and (
+                    inter_stage_output is not None
+                    or (new_token_ids and requires_token_updates)
+                    or is_segment_finished
+                    or finished
+                )
             ):
                 save_kwargs = {
                     "new_token_ids": new_token_ids,

vllm_omni/deploy/mammoth_moda2.yaml

@@ -19,10 +19,12 @@ stages:
     request_batch_max_wait_ms: 20
     gpu_memory_utilization: 0.3
     enforce_eager: true
+    step_execution: false
     trust_remote_code: true
     enable_prefix_caching: false
     # Optional Cache-DiT acceleration for the DiT stage (~2x faster denoise,
-    # training-free). Uncomment to enable; guidance=1.0 requests bypass it.
+    # training-free). Request mode only; uncomment to enable.
+    # guidance=1.0 requests bypass it. Use mammoth_moda2_step.yaml for step mode.
     # cache_backend: cache_dit
     # cache_config:
     #   Fn_compute_blocks: 1

vllm_omni/deploy/mammoth_moda2_step.yaml

@@ -0,0 +1,10 @@
+# MammothModa2 AR→DiT with step execution and continuous batching.
+
+base_config: mammoth_moda2.yaml
+
+stages:
+  - stage_id: 1
+    step_execution: true
+    request_batch_max_wait_ms: 0
+    # Step mode does not support diffusion cache backends.
+    cache_backend: none

vllm_omni/diffusion/attention/backends/trtllm_attn.py

@@ -15,6 +15,16 @@ from vllm_omni.diffusion.attention.backends.abstract import (
     AttentionMetadata,
     PackedPaddingMetadata,
 )
+from vllm_omni.diffusion.attention.capabilities import (
+    CapabilityResult,
+    CompilationMode,
+    ExecutionContext,
+    ExecutionPathResult,
+    MaskMode,
+    PackingMode,
+    ParallelStrategy,
+)
+from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
 
@@ -88,6 +98,96 @@ except Exception as e:  # pragma: no cover - import guard
     )
 
 
+if not hasattr(torch.ops.vllm_omni, "trtllm_ragged_attention"):
+
+    @torch.library.custom_op(
+        "vllm_omni::trtllm_ragged_attention",
+        mutates_args=("workspace_buffer",),
+    )
+    def _trtllm_ragged_attention_op(
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        workspace_buffer: torch.Tensor,
+        seq_lens: torch.Tensor,
+        cum_seq_lens_q: torch.Tensor,
+        cum_seq_lens_kv: torch.Tensor,
+        sage_q_sf: torch.Tensor | None,
+        sage_k_sf: torch.Tensor | None,
+        sage_v_sf: torch.Tensor | None,
+        max_q_len: int,
+        max_kv_len: int,
+        batch_size: int,
+        bmm1_scale: float,
+        bmm2_scale: float,
+        skip_softmax_threshold_scale_factor: float,
+        sage_q_block_size: int,
+        sage_k_block_size: int,
+        is_causal: bool,
+    ) -> torch.Tensor:
+        sage_kwargs = {}
+        if sage_q_sf is not None:
+            sage_kwargs = {
+                "sage_attn_sfs": (sage_q_sf, sage_k_sf, None, sage_v_sf),
+                "num_elts_per_sage_attn_blk": (sage_q_block_size, sage_k_block_size, 0, 1),
+            }
+        return trtllm_ragged_attention_deepseek(
+            query=query,
+            key=key,
+            value=value,
+            workspace_buffer=workspace_buffer,
+            seq_lens=seq_lens,
+            max_q_len=max_q_len,
+            max_kv_len=max_kv_len,
+            bmm1_scale=bmm1_scale,
+            bmm2_scale=bmm2_scale,
+            o_sf_scale=-1.0,
+            batch_size=batch_size,
+            window_left=-1,
+            cum_seq_lens_q=cum_seq_lens_q,
+            cum_seq_lens_kv=cum_seq_lens_kv,
+            enable_pdl=False,
+            is_causal=is_causal,
+            return_lse=False,
+            skip_softmax_threshold_scale_factor=(
+                None if skip_softmax_threshold_scale_factor < 0.0 else skip_softmax_threshold_scale_factor
+            ),
+            skip_all_rows_active_check=True,
+            **sage_kwargs,
+        )
+
+    @_trtllm_ragged_attention_op.register_fake
+    def _trtllm_ragged_attention_fake(
+        query,
+        key,
+        value,
+        workspace_buffer,
+        seq_lens,
+        cum_seq_lens_q,
+        cum_seq_lens_kv,
+        sage_q_sf,
+        sage_k_sf,
+        sage_v_sf,
+        max_q_len,
+        max_kv_len,
+        batch_size,
+        bmm1_scale,
+        bmm2_scale,
+        skip_softmax_threshold_scale_factor,
+        sage_q_block_size,
+        sage_k_block_size,
+        is_causal,
+    ):
+        # Match FlashInfer's default output dtype for quantized queries.
+        out_dtype = (
+            torch.bfloat16 if query.dtype in (torch.float8_e4m3fn, torch.float8_e5m2, torch.int8) else query.dtype
+        )
+        return torch.empty((*query.shape[:-1], value.shape[-1]), dtype=out_dtype, device=query.device)
+
+
+_trtllm_ragged_attention_op = torch.ops.vllm_omni.trtllm_ragged_attention
+
+
 _QK_QUANT_DTYPES = {
     "int8": torch.int8,
     "fp8_e4m3": torch.float8_e4m3fn,
@@ -156,9 +256,53 @@ def _workspace_bytes() -> int:
     return getattr(envs, "VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE", 394 * 1024 * 1024)
 
 
+def _is_trtllm_blackwell_device(query: torch.Tensor) -> bool:
+    # Capability inspection is eager; leave device queries outside forward.
+    if query.device.type != "cuda" or not current_omni_platform.is_cuda():
+        return False
+    capability = current_omni_platform.get_device_capability(query.device.index)
+    return capability is not None and tuple(capability) in ((10, 0), (10, 3))
+
+
+def _normalize_trtllm_metadata(
+    attn_metadata: AttentionMetadata | None,
+) -> tuple[dict, bool, PackedPaddingMetadata | None]:
+    """Inspect host metadata only; never read device tensor values."""
+    extra = getattr(attn_metadata, "extra", {}) if attn_metadata is not None else {}
+    packed_keys = ("cu_seqlens_q", "cu_seqlens_k", "max_seqlen_q", "max_seqlen_k")
+    present_packed_keys = [key for key in packed_keys if key in extra]
+    if present_packed_keys and len(present_packed_keys) != len(packed_keys):
+        missing = sorted(set(packed_keys) - set(present_packed_keys))
+        raise ValueError(f"Incomplete packed TRTLLM attention metadata; missing {missing}")
+    has_packed_metadata = len(present_packed_keys) == len(packed_keys)
+
+    attn_mask = getattr(attn_metadata, "attn_mask", None) if attn_metadata is not None else None
+    packed_padding = getattr(attn_metadata, "packed_padding", None) if attn_metadata is not None else None
+    if packed_padding is not None and not isinstance(packed_padding, PackedPaddingMetadata):
+        raise ValueError("packed_padding must be PackedPaddingMetadata")
+    if packed_padding is not None and not has_packed_metadata:
+        raise ValueError("PackedPaddingMetadata requires complete packed TRTLLM attention metadata")
+    if attn_mask is not None:
+        raise ValueError(
+            "TRTLLM_ATTN does not support attn_mask. Represent structural suffix padding "
+            "with packed-padding metadata, or select a mask-capable backend such as "
+            "CUDNN_ATTN or TORCH_SDPA."
+        )
+
+    return extra, has_packed_metadata, packed_padding
+
+
 class TrtllmAttentionBackend(AttentionBackend):
     accept_output_buffer: bool = True
 
+    @classmethod
+    def resolve_capabilities(cls, context: ExecutionContext) -> ExecutionPathResult:
+        # Construction decides SAGE/skip settings; caller-provided kernel identity
+        # cannot establish the initialized path or a fullgraph guarantee.
+        return ExecutionPathResult.unmigrated(
+            cls.get_name(), replace(context, kernel_variant=None), path="uninitialized"
+        )
+
     @classmethod
     def supports_packed_mask_free(cls) -> bool:
         return True
@@ -231,6 +375,96 @@ class TrtllmAttentionImpl(AttentionImpl):
                 )
             self._sage_quantize_fn = flashinfer.trtllm_sage_attention_quantize
 
+    def resolve_execution_path(
+        self,
+        context: ExecutionContext,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None,
+    ) -> ExecutionPathResult:
+        extra, packed, packed_padding = _normalize_trtllm_metadata(attn_metadata)
+        packing_mode = PackingMode.NONE
+        if packed_padding is not None:
+            packing_mode = PackingMode.PACKED_PADDING
+        elif packed:
+            packing_mode = PackingMode.MULTI_DOCUMENT
+        context = replace(
+            context,
+            kernel_variant="trtllm_sage" if self.quant.enabled else "trtllm",
+            dtype=str(query.dtype).removeprefix("torch."),
+            causal=self.causal,
+            mask_mode=MaskMode.NONE,
+            packing_mode=packing_mode,
+            piecewise=attn_metadata is not None and attn_metadata.full_attn_spans is not None,
+            kv_cache_dtype=extra.get("kv_cache_dtype"),
+        )
+        if self.quant.enabled:
+            path = "sage"
+        elif self.skip.configured:
+            path = "skip_softmax"
+        elif packed_padding is not None:
+            path = "packed_padding"
+        elif packed:
+            path = "packed_varlen"
+        else:
+            path = "dense"
+        result = ExecutionPathResult.unmigrated("TRTLLM_ATTN", context, path="trtllm_" + path)
+        if not HAS_FLASHINFER:
+            return replace(
+                result,
+                support=CapabilityResult.unsupported(
+                    "TRTLLM_ATTN requires FlashInfer; install it or select another backend."
+                ),
+            )
+        if (
+            context.platform != "cuda"
+            or not _is_trtllm_blackwell_device(query)
+            or path != "dense"
+            or self.causal
+            or context.dtype != "bfloat16"
+            or context.piecewise
+            or context.paged_kv
+            or context.kv_cache_dtype not in (None, "auto", "float")
+            or context.parallel_strategy is not ParallelStrategy.NONE
+            or context.outer_boundaries
+        ):
+            return replace(
+                result,
+                support=CapabilityResult.unmigrated(
+                    "Only dense noncausal BF16 TRTLLM on Blackwell SM100/SM103 "
+                    "without parallel or outer boundaries is migrated."
+                ),
+            )
+        if any(t.ndim != 4 for t in (query, key, value)):
+            reason = "Q, K, and V must have rank 4 (batch, sequence, heads, head dimension)."
+        elif query.dtype != key.dtype or query.dtype != value.dtype:
+            reason = "Q, K, and V dtypes must match."
+        elif query.device != key.device or query.device != value.device:
+            reason = "Q, K, and V devices must match."
+        elif key.shape != value.shape or query.shape[0] != key.shape[0] or query.shape[-1] != key.shape[-1]:
+            reason = "K/V shapes, Q/K batch sizes, and Q/K head dimensions must match."
+        elif any(size == 0 for t in (query, key, value) for size in t.shape):
+            reason = "Q, K, and V dimensions must be nonzero."
+        else:
+            reason = None
+        if reason:
+            return replace(
+                result,
+                support=CapabilityResult.unsupported(
+                    "TRTLLM_ATTN: " + reason + " Select compatible inputs or another backend."
+                ),
+            )
+        # This is a verified geometry, not an exhaustive list of kernel limits.
+        if query.shape[-1] != 128 or query.shape[2] != key.shape[2]:
+            return replace(
+                result,
+                support=CapabilityResult.unmigrated(
+                    "TRTLLM fullgraph validation covers head dimension 128 with equal Q/K/V head counts."
+                ),
+            )
+        return replace(result, support=CapabilityResult.supported(), compilation_mode=CompilationMode.CUSTOM_OP)
+
     def set_layer_calibration(self, a: float, b: float) -> None:
         self.skip = replace(self.skip, a=a, b=b)
 
@@ -361,26 +595,7 @@ class TrtllmAttentionImpl(AttentionImpl):
         value: torch.Tensor,
         attn_metadata: AttentionMetadata | None = None,
     ) -> torch.Tensor:
-        extra = getattr(attn_metadata, "extra", {}) if attn_metadata is not None else {}
-        packed_keys = ("cu_seqlens_q", "cu_seqlens_k", "max_seqlen_q", "max_seqlen_k")
-        present_packed_keys = [key for key in packed_keys if key in extra]
-        if present_packed_keys and len(present_packed_keys) != len(packed_keys):
-            missing = sorted(set(packed_keys) - set(present_packed_keys))
-            raise ValueError(f"Incomplete packed TRTLLM attention metadata; missing {missing}")
-        has_packed_metadata = len(present_packed_keys) == len(packed_keys)
-
-        attn_mask = getattr(attn_metadata, "attn_mask", None) if attn_metadata is not None else None
-        packed_padding = getattr(attn_metadata, "packed_padding", None) if attn_metadata is not None else None
-        if packed_padding is not None and not isinstance(packed_padding, PackedPaddingMetadata):
-            raise ValueError("packed_padding must be PackedPaddingMetadata")
-        if packed_padding is not None and not has_packed_metadata:
-            raise ValueError("PackedPaddingMetadata requires complete packed TRTLLM attention metadata")
-        if attn_mask is not None:
-            raise ValueError(
-                "TRTLLM_ATTN does not support attn_mask. Represent structural suffix padding "
-                "with packed-padding metadata, or select a mask-capable backend such as "
-                "CUDNN_ATTN or TORCH_SDPA."
-            )
+        extra, has_packed_metadata, packed_padding = _normalize_trtllm_metadata(attn_metadata)
 
         if not HAS_FLASHINFER:
             raise ImportError(
@@ -449,10 +664,8 @@ class TrtllmAttentionImpl(AttentionImpl):
 
         _skip_factor = self._resolve_skip_factor(max_kv_len)
 
-        # SAGE kwargs are only understood by newer FlashInfer builds; pass them exclusively when
-        # SAGE quant is active (which already requires the kernel, checked at init) so the dense
-        # path stays compatible with older builds that lack these parameters.
-        sage_kwargs: dict = {}
+        sage_q_sf = sage_k_sf = sage_v_sf = None
+        sage_q_block_size = sage_k_block_size = 0
         if self.quant.enabled:
             q, k, v, sage_attn_sfs, sage_block_sizes = self.quant.quantize(
                 q,
@@ -462,30 +675,29 @@ class TrtllmAttentionImpl(AttentionImpl):
                 cu_seq_lens_q,
                 cu_seq_lens_kv,
             )
-            sage_kwargs["sage_attn_sfs"] = sage_attn_sfs
-            sage_kwargs["num_elts_per_sage_attn_blk"] = sage_block_sizes
+            sage_q_sf, sage_k_sf, _, sage_v_sf = sage_attn_sfs
+            sage_q_block_size, sage_k_block_size, _, _ = sage_block_sizes
 
-        out = trtllm_ragged_attention_deepseek(
-            query=q,
-            key=k,
-            value=v,
-            workspace_buffer=workspace,
-            seq_lens=seq_lens,
-            max_q_len=max_q_len,
-            max_kv_len=max_kv_len,
-            bmm1_scale=bmm1_scale,
-            bmm2_scale=bmm2_scale,
-            o_sf_scale=-1.0,
-            batch_size=batch,
-            window_left=-1,
-            cum_seq_lens_q=cu_seq_lens_q,
-            cum_seq_lens_kv=cu_seq_lens_kv,
-            enable_pdl=False,
-            is_causal=self.causal,
-            return_lse=False,
-            skip_softmax_threshold_scale_factor=_skip_factor,
-            skip_all_rows_active_check=True,
-            **sage_kwargs,
+        out = _trtllm_ragged_attention_op(
+            q,
+            k,
+            v,
+            workspace,
+            seq_lens,
+            cu_seq_lens_q,
+            cu_seq_lens_kv,
+            sage_q_sf,
+            sage_k_sf,
+            sage_v_sf,
+            max_q_len,
+            max_kv_len,
+            batch,
+            bmm1_scale,
+            bmm2_scale,
+            -1.0 if _skip_factor is None else _skip_factor,
+            sage_q_block_size,
+            sage_k_block_size,
+            self.causal,
         )
         if out.shape[0] != output_tokens:
             padded_out = torch.zeros(

vllm_omni/diffusion/cache/teacache/backend.py

@@ -8,8 +8,9 @@ This module provides the TeaCache backend that implements the CacheBackend
 interface using the hooks-based TeaCache system.
 """
 
-from typing import Any
+from typing import Any, Protocol
 
+import torch.nn as nn
 from vllm.logger import init_logger
 
 from vllm_omni.diffusion.cache.base import CacheBackend
@@ -20,6 +21,11 @@ from vllm_omni.diffusion.data import DiffusionCacheConfig
 logger = init_logger(__name__)
 
 
+class _MammothModa2TeaCachePipeline(Protocol):
+    gen_transformer: nn.Module
+    transformer: nn.Module
+
+
 def enable_hunyuan_image3_teacache(pipeline: Any, config: DiffusionCacheConfig) -> None:
     """
     Enable TeaCache for HunyuanImage3 model.
@@ -57,6 +63,23 @@ def enable_bagel_teacache(pipeline: Any, config: DiffusionCacheConfig) -> None:
     )
 
 
+def enable_mammoth_moda2_teacache(pipeline: _MammothModa2TeaCachePipeline, config: DiffusionCacheConfig) -> None:
+    """Enable TeaCache for MammothModa2's DiT stage."""
+    teacache_config = TeaCacheConfig(
+        transformer_type="MammothModa2Transformer2DModel",
+        rel_l1_thresh=config.rel_l1_thresh,
+        coefficients=config.coefficients,
+    )
+    transformer = pipeline.gen_transformer
+    apply_teacache_hook(transformer, teacache_config)
+    pipeline.transformer = transformer
+
+    logger.info(
+        f"TeaCache applied with rel_l1_thresh={teacache_config.rel_l1_thresh}, "
+        f"transformer_class={teacache_config.transformer_type}"
+    )
+
+
 def enable_sensenova_u1_teacache(pipeline: Any, config: DiffusionCacheConfig) -> None:
     """Enable TeaCache for SenseNova-U1 denoising forwards."""
     teacache_config = TeaCacheConfig(
@@ -124,6 +147,7 @@ CUSTOM_TEACACHE_ENABLERS = {
     "BagelPipeline": enable_bagel_teacache,
     "Flux2KleinPipeline": enable_flux2_klein_teacache,
     "HunyuanImage3Pipeline": enable_hunyuan_image3_teacache,
+    "MammothModa2DiTPipeline": enable_mammoth_moda2_teacache,
     "MiniMaxH3Pipeline": enable_minimax_h3_teacache,
     "SenseNovaU1Pipeline": enable_sensenova_u1_teacache,
 }

vllm_omni/diffusion/cache/teacache/config.py

@@ -83,6 +83,15 @@ _MODEL_COEFFICIENTS = {
     ],
     # LongCat Image transformer coefficients
     "LongCatImageTransformer2DModel": [652.5980, -424.1615, 84.5526, -4.5923, 0.1694],
+    # MammothModa2 DiT coefficients fitted from full-compute traces
+    # at 1024x1024 with 50 denoising steps.
+    "MammothModa2Transformer2DModel": [
+        -1761.242764481119,
+        859.398730831851,
+        -126.78044436937441,
+        9.16173963457422,
+        -0.1102826051200547,
+    ],
     # MiniMax-H3 FL2VA coefficients.
     "MiniMaxH3DiTModel": [
         2.283704065852778e03,
@@ -95,6 +104,7 @@ _MODEL_COEFFICIENTS = {
 
 _DEFAULT_REL_L1_THRESH = 0.2
 _MODEL_DEFAULT_REL_L1_THRESH = {
+    "MammothModa2Transformer2DModel": 0.075,
     "MiniMaxH3DiTModel": 0.17,
 }
 

vllm_omni/diffusion/cache/teacache/extractors.py

@@ -20,6 +20,7 @@ from typing import Any
 
 import torch
 import torch.nn as nn
+from einops import rearrange
 from vllm.logger import init_logger
 
 from vllm_omni.diffusion.forward_context import get_forward_context
@@ -823,6 +824,114 @@ def extract_longcat_context(
     )
 
 
+def extract_mammoth_moda2_context(
+    module: nn.Module,
+    hidden_states: torch.Tensor,
+    timestep: torch.Tensor,
+    text_hidden_states: torch.Tensor,
+    freqs_cis: torch.Tensor,
+    text_attention_mask: torch.Tensor,
+    ref_image_hidden_states: list[list[torch.Tensor]] | None = None,
+    ar_image_hidden_states: torch.Tensor | None = None,
+    ar_image_attention_mask: torch.Tensor | None = None,
+    return_dict: bool = False,
+    teacache_branch: str | None = None,
+) -> CacheContext:
+    """Extract cache context for MammothModa2's Lumina-style DiT."""
+    if not hasattr(module, "layers") or len(module.layers) == 0:
+        raise ValueError("Module must have layers")
+
+    batch_size, height, width = module._validate_inputs(
+        hidden_states,
+        text_hidden_states,
+        text_attention_mask,
+        ref_image_hidden_states,
+        return_dict,
+    )
+
+    (
+        temb,
+        text_hidden_states,
+        text_attention_mask,
+        img_tokens,
+        img_mask,
+        img_len,
+        context_rotary_emb,
+        noise_rotary_emb,
+        rotary_emb,
+        encoder_seq_lengths,
+        seq_lengths,
+    ) = module._prepare_embeddings(
+        hidden_states,
+        timestep,
+        text_hidden_states,
+        text_attention_mask,
+        freqs_cis,
+        batch_size,
+        height,
+        width,
+        ar_image_hidden_states,
+        ar_image_attention_mask,
+    )
+
+    text_hidden_states, img_tokens = module._apply_refiners(
+        text_hidden_states,
+        text_attention_mask,
+        context_rotary_emb,
+        img_tokens,
+        img_mask,
+        noise_rotary_emb,
+        temb,
+    )
+
+    max_seq_len = max(seq_lengths)
+    attention_mask = img_tokens.new_zeros(batch_size, max_seq_len, dtype=torch.bool)
+    joint_hidden_states = img_tokens.new_zeros(batch_size, max_seq_len, module.config.hidden_size)
+    for i, (encoder_seq_len, seq_len) in enumerate(zip(encoder_seq_lengths, seq_lengths)):
+        attention_mask[i, :seq_len] = True
+        joint_hidden_states[i, :encoder_seq_len] = text_hidden_states[i, :encoder_seq_len]
+        joint_hidden_states[i, encoder_seq_len : encoder_seq_len + img_len] = img_tokens[i, :img_len]
+
+    modulated_input = module.layers[0].norm1(joint_hidden_states, temb)[0]
+
+    def run_transformer_blocks() -> tuple[torch.Tensor]:
+        h = module._apply_transformer_layers(joint_hidden_states, attention_mask, rotary_emb, temb)
+        return (h,)
+
+    def postprocess(h: torch.Tensor) -> torch.Tensor:
+        h = module.norm_out(h, temb)
+        p = module.config.patch_size
+        img_hidden_states = torch.stack(
+            [
+                h[i, encoder_seq_len : encoder_seq_len + img_len]
+                for i, encoder_seq_len in enumerate(encoder_seq_lengths)
+            ],
+            dim=0,
+        )
+        return rearrange(
+            img_hidden_states,
+            "b (h w) (p1 p2 c) -> b c (h p1) (w p2)",
+            h=height // p,
+            w=width // p,
+            p1=p,
+            p2=p,
+        )
+
+    extra_states = {}
+    if teacache_branch is not None:
+        extra_states["teacache_branch"] = teacache_branch
+
+    return CacheContext(
+        modulated_input=modulated_input,
+        hidden_states=joint_hidden_states,
+        encoder_hidden_states=None,
+        temb=temb,
+        run_transformer_blocks=run_transformer_blocks,
+        postprocess=postprocess,
+        extra_states=extra_states,
+    )
+
+
 def extract_stable_audio_context(
     module: nn.Module,
     hidden_states: torch.Tensor,
@@ -1556,6 +1665,7 @@ EXTRACTOR_REGISTRY: dict[str, Callable] = {
     "StableAudioDiTModel": extract_stable_audio_context,
     "Flux2Transformer2DModel": extract_flux2_context,
     "LongCatImageTransformer2DModel": extract_longcat_context,
+    "MammothModa2Transformer2DModel": extract_mammoth_moda2_context,
     "FluxTransformer2DModel": extract_flux_context,
     "SenseNovaU1ForCausalLM": extract_sensenova_u1_context,
     "MiniMaxH3DiTModel": extract_minimax_h3_context,

vllm_omni/diffusion/cache/teacache/hook.py

@@ -145,7 +145,15 @@ class TeaCacheHook(ModelHook):
         #   - cfg_rank 0: positive branch
         #   - cfg_rank > 0: negative branch
         # Without CFG-parallel, branches alternate within a single rank
-        if getattr(module, "do_true_cfg", False):
+        branch_hint = None
+        if getattr(ctx, "extra_states", None):
+            branch_hint = ctx.extra_states.get("teacache_branch")
+
+        if branch_hint is not None:
+            if branch_hint not in ("positive", "negative"):
+                raise ValueError(f"Invalid teacache_branch={branch_hint!r}; expected 'positive' or 'negative'.")
+            cache_branch = branch_hint
+        elif getattr(module, "do_true_cfg", False):
             cfg_parallel_size = get_classifier_free_guidance_world_size()
             if cfg_parallel_size > 1:
                 cfg_rank = get_classifier_free_guidance_rank()

vllm_omni/diffusion/data.py

@@ -1005,6 +1005,9 @@ class OmniDiffusionConfig:
     # Worker extension class for custom functionality
     worker_extension_cls: str | None = None
 
+    # Internal transport of explicit stage runtime.env to remote Ray actors.
+    ray_worker_env: dict[str, str] = field(default_factory=dict, init=False, repr=False)
+
     # Custom pipeline arguments for custom pipelines
     custom_pipeline_args: dict[str, Any] | None = None
 

vllm_omni/diffusion/diffusion_engine.py

@@ -669,8 +669,10 @@ class DiffusionEngine:
                 self._poll_native_kv()
                 worker_execution_completed = True
             except Exception as exc:
+                # Shutdown can interrupt an in-flight executor call. Exit through
+                # the loop tail so pending RPCs are failed instead of left waiting.
                 if self._closed:
-                    return
+                    break
                 worker_execution_completed = False
                 logger.error(
                     "Execution failed for diffusion requests %s", sched_output.scheduled_request_ids, exc_info=True
@@ -1569,6 +1571,10 @@ class DiffusionEngine:
         for stream in pending_streams:
             self._put_queue_output(stream, closed_output)
 
+        # Interrupt remote inference before waiting for the thread blocked on Ray.
+        if self.od_config.distributed_executor_backend == "ray":
+            self.executor.shutdown()
+
         worker_thread = self.worker_thread
         if worker_thread is not None:
             if worker_thread.is_alive():

vllm_omni/diffusion/executor/abstract.py

@@ -49,7 +49,9 @@ class DiffusionExecutor(ABC):
                 )
             executor_class = distributed_executor_backend
         elif distributed_executor_backend == "ray":
-            raise NotImplementedError("ray backend is not yet supported.")
+            from vllm_omni.diffusion.executor.ray_executor import RayDiffusionExecutor
+
+            executor_class = RayDiffusionExecutor
         elif distributed_executor_backend == "mp":
             from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor