[Doc] Keep ROCm run examples on the published v0.28.0 image tag

The release pipeline builds and publishes only the CUDA image
(.buildkite/release/release-pipeline.yml builds docker/Dockerfile.cuda;
publish-release-images.sh targets vllm/vllm-omni). The prebuilt ROCm
image is uploaded out-of-band, and v0.28.0 is the latest tag confirmed
live on Docker Hub (#7405). Revert the two docker run examples to
v0.28.0 and state why the tag lags the 0.30.0 wheels, so users do not
hit a missing-tag pull failure. Source-build examples stay on the
upstream vllm/vllm-openai-rocm:v0.30.0 base, which is published.

Signed-off-by: hsliu_ustc <hsliu_ustc@noreply.gitcode.com>

.buildkite/cuda/test-ready.yml

@@ -69,6 +69,7 @@ steps:
         source_file_dependencies: diffusion_distributed_attention
         commands:
           - timeout 20m pytest -sv tests/diffusion/distributed -m 'core_model and cuda and L4 and cards_2' --run-level "core_model"
+          - timeout 10m pytest -sv tests/diffusion/attention/test_ulysses_uaa.py -m 'core_model and cuda and L4 and cards_2' --run-level "core_model"
         mirror_hardwares: l4_2
 
       - label: "Diffusion · Distributed & Attention Test · 4-GPU"

README.md

@@ -16,6 +16,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
 
 *Latest News* 🔥
 
+- [2026/09] We released [0.30.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.30.0), rebased onto vLLM 0.30.0, featuring a unified full-duplex serving framework around engine-owned sessions for [MiniCPM-o 4.5](recipes/OpenBMB/MiniCPM-o-4_5.md) and AURA, native cross-stage KV and multimodal payload transfer (Mooncake AR-to-DiT handoff and NIXL connectors), interactive world-model serving with [LingBot World](recipes/Robbyant/LingBot-World-2.0.md), and realtime [MiniMax H3](recipes/MiniMaxAI/MiniMax-H3.md) Turbo inference on NVIDIA Blackwell and Ascend 950.
 - [2026/08] We released [0.28.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.28.0), featuring production-ready [MiniMax H3](recipes/MiniMaxAI/MiniMax-H3.md) serving on GPU and NPU, a unified AR/DiT paged KV cache runtime, and enhanced realtime full-duplex serving for the [MiniCPM-o series](recipes/OpenBMB/MiniCPM-o-4_5.md).
 - [2026/08] [VeRL-Omni](https://github.com/verl-project/verl-omni) `v0.2.0` is released: faster diffusion RL powered by vLLM-Omni (request-level/step-wise batching with FA3), rebuilt Qwen3-Omni multimodal training (DPO & GSPO), plus LTX-2.3, Qwen-Image-Edit support and more. See the [release notes](https://github.com/verl-project/verl-omni/releases/tag/v0.2.0).
 - [2026/08] We released [0.26.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.26.0) - aligned with the vLLM 0.26 release line, featuring [MiniMax H3](recipes/MiniMaxAI/MiniMax-H3.md) joint video/audio generation, an experimental full-duplex realtime runtime for [MiniCPM-o 4.5](recipes/OpenBMB/MiniCPM-o-4_5.md), distributed layerwise diffusion offload, and broader model, hardware, streaming, TTS, and quantization support.
@@ -58,9 +59,9 @@ vLLM-Omni is flexible and easy to use with:
 vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, including:
 
 - **Omni-modality models** (e.g. Qwen3-Omni, MiniCPM-o 4.5, Cosmos3, HunyuanImage, BAGEL)
-- **TTS models** (e.g. Qwen3-TTS, IndexTTS 2.5, CosyVoice3)
-- **Diffusion models** — image, video, and audio generation (e.g. MiniMax H3, LTX-2.5, SANA-Video, Wan2.2)
-- **Robot-policy and action models** (e.g. π0, GR00T-N1.7, DreamZero-DROID, InternVLA-A1)
+- **TTS models** (e.g. Qwen3-TTS, Tencent AuK, Breeze-TTS-2, CosyVoice3)
+- **Diffusion models** — image, video, and audio generation (e.g. MiniMax H3, LingBot World, MAGI-2, LTX-2.5, Wan2.2)
+- **Robot-policy and action models** (e.g. π0.5, GR00T-N1.7, DreamZero-DROID, InternVLA-A1)
 
 ## Getting Started
 

benchmarks/kernels/_magi2_gpu_bench.py

@@ -0,0 +1,194 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Cold-cache torch.cuda timing and correctness gates for MAGI-2 microbenchmarks."""
+
+import argparse
+import importlib
+import importlib.metadata
+import json
+import logging
+import math
+import os
+import subprocess
+import sys
+from collections.abc import Callable
+from pathlib import Path
+from types import ModuleType
+from typing import Any
+
+LOGGER = logging.getLogger(__name__)
+
+
+def emit(**fields: Any) -> None:
+    LOGGER.info(json.dumps(fields, sort_keys=True))
+
+
+def add_arguments(parser: argparse.ArgumentParser) -> None:
+    parser.add_argument("--tokens", type=int, nargs="+", default=[1, 4, 16, 64, 129, 4096])
+    parser.add_argument("--dtype", choices=["bfloat16", "float16", "float32"], default="bfloat16")
+    parser.add_argument("--num-tests", type=int, default=30, help="Measured samples per implementation.")
+    parser.add_argument("--warmup", type=int, default=10, help="Interleaved warmup calls before sampling.")
+    parser.add_argument("--seed", type=int, default=42)
+
+
+def load_musa(args: argparse.Namespace, parser: argparse.ArgumentParser) -> Any:
+    if min(*args.tokens, args.num_tests) <= 0 or args.warmup < 0:
+        parser.error("tokens and num-tests must be positive; warmup must be nonnegative")
+    logging.basicConfig(level=logging.INFO, format="%(message)s", stream=sys.stdout)
+    torch = importlib.import_module("torch")
+    cuda = getattr(torch, "cuda")
+    # torchada is needed only on MUSA; after it is loaded, use torch.cuda for
+    # both CUDA and MUSA as documented by torchada.
+    if hasattr(torch.version, "musa") and torch.version.musa is not None:
+        try:
+            importlib.import_module("torchada")
+        except ModuleNotFoundError as error:
+            if error.name != "torchada":
+                raise
+    has_gpu = cuda.is_available() or (
+        hasattr(torch.version, "musa") and torch.version.musa is not None and torch.musa.is_available()
+    )
+    if not has_gpu:
+        emit(status="skip", reason="CUDA/MUSA device is unavailable", torch=torch.__version__)
+        return None
+    torch.manual_seed(args.seed)
+    cuda.manual_seed_all(args.seed)
+    return torch
+
+
+def report_environment(torch: Any, args: argparse.Namespace, module: ModuleType) -> None:
+    root = Path(__file__).resolve().parents[2]
+    packages = {}
+    for name in ("torch", "torch_musa", "torchada", "triton", "vllm", "vllm-omni"):
+        try:
+            packages[name] = importlib.metadata.version(name)
+        except importlib.metadata.PackageNotFoundError:
+            packages[name] = None
+    cuda = getattr(torch, "cuda")
+    properties = cuda.get_device_properties(cuda.current_device())
+    try:
+        source_sha = subprocess.check_output(["git", "-C", str(root), "rev-parse", "HEAD"], text=True).strip()
+        source_dirty = bool(subprocess.check_output(["git", "-C", str(root), "status", "--porcelain"], text=True))
+    except (subprocess.CalledProcessError, FileNotFoundError):
+        source_sha = os.environ.get("MAGI2_BENCH_SOURCE_SHA", "unknown")
+        source_dirty = None
+    emit(
+        status="environment",
+        command=sys.argv,
+        source_root=str(root),
+        source_sha=source_sha,
+        source_dirty=source_dirty,
+        kernel_source=module.__file__,
+        device=cuda.get_device_name(cuda.current_device()),
+        sm_count=getattr(properties, "multi_processor_count", None),
+        packages=packages,
+        env={name: os.environ.get(name) for name in ("MUSA_VISIBLE_DEVICES", "TORCHDYNAMO_DISABLE")},
+        config=vars(args),
+        mode="eager operator diagnostic; no compile or graph capture",
+        comparison="forward_native versus the default CustomOp dispatch",
+        timing="torch.cuda.Event; alternating native/fused calls; vanilla 512 MiB flush outside timed regions",
+        timing_scope="whole callable, including device launch gaps; not isolated kernel duration",
+        l2_flush=True,
+        l2_flush_bytes=512 * 1024 * 1024,
+    )
+
+
+class _LaunchProbe:
+    def __init__(self, kernel: Any) -> None:
+        self.kernel = kernel
+        self.count = 0
+
+    def __getitem__(self, grid: Any) -> Any:
+        self.count += 1
+        return self.kernel[grid]
+
+
+def check_pair(
+    torch: Any,
+    reference: Callable,
+    fused: Callable,
+    module: ModuleType,
+    kernels: tuple[str, ...],
+    *,
+    rtol: float,
+    atol: float,
+) -> dict:
+    cuda = getattr(torch, "cuda")
+    expected = reference()
+    probes = {name: _LaunchProbe(getattr(module, name)) for name in kernels}
+    try:
+        for name, probe in probes.items():
+            setattr(module, name, probe)
+        actual = fused()
+        cuda.synchronize()
+    finally:
+        for name, probe in probes.items():
+            setattr(module, name, probe.kernel)
+    if any(probe.count != 1 for probe in probes.values()):
+        raise RuntimeError(f"Expected one launch per fused kernel; got {[(k, p.count) for k, p in probes.items()]}")
+    expected = expected if isinstance(expected, tuple) else (expected,)
+    actual = actual if isinstance(actual, tuple) else (actual,)
+    errors = []
+    for result, baseline in zip(actual, expected, strict=True):
+        for value in (baseline, result):
+            if not torch.isfinite(value).all().item():
+                raise AssertionError("Nonfinite benchmark output")
+            if value.numel() > 1 and (value == value.flatten()[0]).all().item():
+                raise AssertionError("Constant benchmark output")
+        torch.testing.assert_close(result, baseline, rtol=rtol, atol=atol)
+        difference = (result.float() - baseline.float()).abs()
+        errors.append(
+            {
+                "max_abs": difference.max().item(),
+                "max_rel": (difference / baseline.float().abs().clamp_min(1e-6)).max().item(),
+            }
+        )
+    return {"parity": "pass", "errors": errors, "rtol": rtol, "atol": atol, "kernels": list(kernels)}
+
+
+def _percentile(samples: list[float], quantile: float) -> float:
+    ordered = sorted(samples)
+    position = (len(ordered) - 1) * quantile
+    lower, upper = math.floor(position), math.ceil(position)
+    return ordered[lower] + (ordered[upper] - ordered[lower]) * (position - lower)
+
+
+def measure_pair(torch: Any, reference: Callable, fused: Callable, args: argparse.Namespace) -> dict:
+    cuda = getattr(torch, "cuda")
+    functions = (reference, fused)
+    order: list[int] = []
+
+    def alternate() -> Any:
+        mode = len(order) % 2
+        order.append(mode)
+        return functions[mode]()
+
+    times_ms = []
+    for _ in range(args.warmup):
+        alternate()
+    cuda.synchronize()
+    order.clear()
+    flush = torch.empty(512 * 1024 * 1024, dtype=torch.uint8, device="cuda")
+    for _ in range(2 * args.num_tests):
+        flush.zero_()
+        cuda.synchronize()
+        start = cuda.Event(enable_timing=True)
+        end = cuda.Event(enable_timing=True)
+        start.record()
+        alternate()
+        end.record()
+        end.synchronize()
+        times_ms.append(start.elapsed_time(end))
+    samples: tuple[list[float], list[float]] = ([], [])
+    for mode, milliseconds in zip(order[-len(times_ms) :], times_ms, strict=True):
+        if not math.isfinite(milliseconds) or milliseconds <= 0:
+            raise RuntimeError(f"Invalid GPU event sample: {milliseconds}")
+        samples[mode].append(milliseconds * 1000)
+    stats = {}
+    for label, values in zip(("native", "fused"), samples, strict=True):
+        if len(values) != args.num_tests:
+            raise RuntimeError("GPU timing harness returned an unexpected sample count")
+        p50 = _percentile(values, 0.5)
+        stats[label] = {"median_us": p50, "p50_us": p50, "p90_us": _percentile(values, 0.9), "samples": len(values)}
+    stats["speedup_p50"] = stats["native"]["p50_us"] / stats["fused"]["p50_us"]
+    return stats

benchmarks/kernels/benchmark_magi2_mhc.py

@@ -0,0 +1,141 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Compare native and fused MAGI-2 mHC post-processing on CUDA or MUSA.
+
+Run from the checkout root with PyTorch and the selected backend dependencies.
+torchada is loaded automatically only on MUSA::
+
+    python -m benchmarks.kernels.benchmark_magi2_mhc \
+        --tokens 1 4 16 64 129 4096 --hidden-size 3072 --num-tests 30
+
+The released config supplies hidden=3072, four streams and 20 Sinkhorn steps.
+Token counts are synthetic; supply per-rank sizes from the target trace when
+comparing production shapes. FP32 post/residual logits are strided views of the
+same [tokens, 24] projection, as in MHCHandler.compute_logits. Mix inputs use
+the selected dtype. Both mix implementations consume identical coefficients.
+
+The three comparisons cover coefficient preparation, stream mixing, and their
+combined pipeline (including intermediate allocations, excluding projection,
+normalization and apply_pre). This is an eager operator diagnostic, without
+compile or graph capture. It does not measure full model performance.
+No CUDA/MUSA device (or no PyTorch) produces a JSON skip record and exit status zero.
+"""
+
+import argparse
+import math
+from functools import partial
+from typing import Any
+
+from benchmarks.kernels._magi2_gpu_bench import (
+    add_arguments,
+    check_pair,
+    emit,
+    load_musa,
+    measure_pair,
+    report_environment,
+)
+
+
+def benchmark_shape(torch: Any, mhc: Any, args: argparse.Namespace, tokens: int) -> None:
+    dtype = getattr(torch, args.dtype)
+    device = torch.device("cuda")
+    packed = torch.randn(tokens, 24, device=device, dtype=torch.float32)
+    # Preserve the token stride of the model's shared projection allocation.
+    post_inputs = (
+        packed[:, 4:8],
+        packed[:, 8:].view(tokens, 4, 4),
+        torch.tensor([0.8], device=device),
+        torch.randn(4, device=device),
+        torch.tensor([1.1], device=device),
+        torch.randn(4, 4, device=device),
+    )
+    streams = torch.randn(tokens, 4, args.hidden_size, device=device, dtype=dtype)
+    branch = torch.randn(tokens, args.hidden_size, device=device, dtype=dtype)
+    post_op, mix_op = mhc.MHCPostResidual(), mhc.MHCMix()
+    scale = 1 / math.sqrt(4 * args.hidden_size)
+    kwargs = dict(scale=scale, iterations=args.sinkhorn_iterations, epsilon=1e-12, out_dtype=dtype)
+    prepare_reference = partial(post_op.forward_native, *post_inputs, **kwargs)
+    prepare_fused = partial(post_op, *post_inputs, **kwargs)
+    coefficients = prepare_reference()
+    mix_reference = partial(mix_op.forward_native, streams, branch, *coefficients)
+    mix_fused = partial(mix_op, streams, branch, *coefficients)
+    post_kernel, mix_kernel = "_mhc_post_residual_kernel", "_mhc_mix_kernel"
+    cases = {
+        "post_residual": (prepare_reference, prepare_fused, (post_kernel,)),
+        "mix": (mix_reference, mix_fused, (mix_kernel,)),
+        "pipeline": (
+            lambda: mix_op.forward_native(streams, branch, *prepare_reference()),
+            lambda: mix_op(streams, branch, *prepare_fused()),
+            (post_kernel, mix_kernel),
+        ),
+    }
+    rtol = {torch.bfloat16: 1.6e-2, torch.float16: 2e-3, torch.float32: 1e-5}[dtype]
+    for name, (reference, fused, kernels) in cases.items():
+        if args.op != "all" and args.op != name:
+            continue
+        # BF16 MUSA custom ops can quantize values near zero to the adjacent
+        # representable value; retain the strict relative gate while allowing
+        # one BF16 ulp in the absolute comparison.
+        atol = (
+            (1e-2 if dtype == torch.bfloat16 else 1e-6)
+            if name == "post_residual"
+            else (1e-2 if dtype == torch.bfloat16 else 1e-5)
+        )
+        parity = check_pair(torch, reference, fused, mhc, kernels, rtol=rtol, atol=atol)
+        timing = measure_pair(torch, reference, fused, args)
+        # Logical minimum traffic: input reads plus output writes. Native
+        # intermediates and repeat parameter loads can add actual DRAM traffic.
+        element_bytes = streams.element_size()
+        prepare_bytes = tokens * 20 * (4 + element_bytes) + 22 * 4
+        mix_bytes = tokens * (9 * args.hidden_size + 20) * element_bytes
+        minimum_bytes = {"post_residual": prepare_bytes, "mix": mix_bytes, "pipeline": prepare_bytes + mix_bytes}[name]
+        # Sigmoid/exp costs and max reductions are not represented as FLOPs.
+        prepare_flops = tokens * (72 + 64 * args.sinkhorn_iterations)
+        mix_flops = tokens * args.hidden_size * 36
+        scalar_flops = {"post_residual": prepare_flops, "mix": mix_flops, "pipeline": prepare_flops + mix_flops}[name]
+        emit(
+            status="pass",
+            op=name,
+            tokens=tokens,
+            hidden_size=args.hidden_size,
+            num_streams=4,
+            dtype=args.dtype,
+            logits_dtype="float32",
+            packed_shape=list(packed.shape),
+            post_stride=list(post_inputs[0].stride()),
+            residual_stride=list(post_inputs[1].stride()),
+            scale=scale,
+            sinkhorn_iterations=args.sinkhorn_iterations,
+            epsilon=1e-12,
+            minimum_io_bytes=minimum_bytes,
+            scalar_flops_excluding_exp_and_max=scalar_flops,
+            effective_fused_gbps=minimum_bytes / timing["fused"]["p50_us"] / 1000,
+            **parity,
+            **timing,
+        )
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    add_arguments(parser)
+    parser.add_argument("--hidden-size", type=int, default=3072)
+    parser.add_argument("--sinkhorn-iterations", type=int, default=20)
+    parser.add_argument("--op", choices=["all", "post_residual", "mix", "pipeline"], default="all")
+    args = parser.parse_args()
+    if args.hidden_size <= 0 or not 0 <= args.sinkhorn_iterations <= 64:
+        parser.error("hidden-size must be positive; sinkhorn-iterations must be in [0, 64]")
+    torch = load_musa(args, parser)
+    if torch is None:
+        return
+    cuda = getattr(torch, "cuda")
+    from vllm_omni.diffusion.layers import mhc
+
+    report_environment(torch, args, mhc)
+    with torch.inference_mode():
+        for tokens in args.tokens:
+            benchmark_shape(torch, mhc, args, tokens)
+            cuda.empty_cache()
+
+
+if __name__ == "__main__":
+    main()

docs/README.md

@@ -56,9 +56,9 @@ vLLM-Omni is flexible and easy to use with:
 vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, including:
 
 - Omni-modality models (e.g. Qwen3-Omni, MiniCPM-o 4.5, Cosmos3, HunyuanImage, BAGEL)
-- TTS models (e.g. Qwen3-TTS, IndexTTS 2.5, CosyVoice3)
-- Diffusion models — image, video, and audio generation (e.g. MiniMax H3, LTX-2.5, SANA-Video, Wan2.2)
-- Robot-policy and action models (e.g. π0, GR00T-N1.7, DreamZero-DROID, InternVLA-A1)
+- TTS models (e.g. Qwen3-TTS, Tencent AuK, Breeze-TTS-2, CosyVoice3)
+- Diffusion models — image, video, and audio generation (e.g. MiniMax H3, LingBot World, MAGI-2, LTX-2.5, Wan2.2)
+- Robot-policy and action models (e.g. π0.5, GR00T-N1.7, DreamZero-DROID, InternVLA-A1)
 
 For more information, checkout the following:
 

docs/configuration/README.md

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

docs/configuration/environment_variables.md

@@ -162,7 +162,7 @@ their keys only.
 ## Inherited vLLM variables
 
 vLLM-Omni also reads variables through its aligned vLLM dependency. Refer to
-the [vLLM 0.29 environment-variable reference](https://docs.vllm.ai/en/v0.29.0/configuration/env_vars.html)
+the [vLLM 0.30 environment-variable reference](https://docs.vllm.ai/en/v0.30.0/configuration/env_vars.html)
 for their definitions. This includes vLLM launch, cache, logging, plugin, ROCm,
 XPU, ModelScope, and FlashInfer workspace settings.
 

docs/design/fullduplex.md

@@ -242,10 +242,12 @@ vllm_omni/
 │           └── overlap_policy.py / commit_policy.py / playback_ledger.py
 ├── config/stage_config.py           PipelineConfig.duplex_plugin; DuplexSessionRuntimeConfig
 ├── model_executor/models/minicpmo_4_5/duplex/plugin.py   MiniCPMO45DuplexPlugin (+ data_plane, input, policy, ...)
+├── model_executor/models/nemotron_voicechat/duplex/plugin.py   NemotronVoiceChatDuplexPlugin (+ data_plane, input, capabilities)
 └── clients/
     ├── duplex.py                    DuplexClientBase (ABC), DuplexClient (websocket), client-side events
     ├── inline_duplex.py             InlineDuplexClient (in-process, over DuplexOmni)
     ├── minicpmo_4_5.py              MiniCPM-o 4.5 session preset
+    ├── nemotron_voicechat.py        Nemotron VoiceChat session preset
     └── personaplex.py               PersonaPlex session preset
 ```
 
@@ -283,8 +285,13 @@ command  handle.submit(DuplexCommand)
          -> 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()
-detach   DuplexOmni.detach_session -> touch(DETACH): engine-owned disconnect grace; expiry -> SessionExpired
-resume   DuplexOmni.resume_session(expected_lease_generation) -> lease CAS; the existing handle is re-entered
+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
+         answers with the generation it produced); the existing handle is re-entered. A caller cancelled
+         mid-RPC observes the outcome afterwards (replaying until the engine answers) and settles a landed
+         resume: the generation goes to the resume waiting to activate, else to the connection still
+         attached, else the lease is detached again
 close    DuplexOmni.close_session -> close RPC; the manager tears the runner down, then the stage cleanup
          (abort submitted requests, release reserved ids), then SessionClosed, then the RPC result.
          SessionClosed is emitted after the cleanup attempt (in a finally), so it is also sent when the
@@ -383,7 +390,7 @@ policy that used to be two separately configured objects:
 | Half | Members |
 | --- | --- |
 | engine policy | `configure_sampling_params(runtime_config, defaults)`, `plan_append(...) -> DuplexAppendPlan` (the resumable Stage0 prompt for one unit), `decide_output(...) -> DuplexOutputDecision \| None` (e.g. the listen decision on a finished Stage0 segment) |
-| session policy | `capabilities(max_sessions)`, `validate_client_extra_body`, `prepare_runtime_config(config, model_config)` (server-owned runtime keys, reference audio resolution), `runtime_config_for_update`, `runtime_config_for_function_output`, `create_session_state() -> DuplexModelSessionState`, `data_plane: DuplexDataPlane` (projects raw stage outputs into internal events), `data_plane_context(...)` |
+| session policy | `capabilities(max_sessions)`, `validate_client_extra_body`, `prepare_runtime_config(config, model_config)` (server-owned runtime keys, reference audio resolution), `runtime_config_for_update`, `runtime_config_for_function_output`, `runtime_config_after_model_output` (consumption acknowledgement), `create_session_state() -> DuplexModelSessionState`, `data_plane: DuplexDataPlane` (projects raw stage outputs into internal events), `data_plane_context(...)` |
 
 `DuplexOmniEngine._validate_deployment` loads the plugin before any stage
 starts; `DuplexSessionManager.__init__` validates it against the stage
@@ -391,12 +398,12 @@ sampling defaults once the stage pools exist. Plugin hooks that may block
 (`prepare_runtime_config` fetching `ref_audio`) are awaited in `open()` and
 offloaded from the loop.
 
-The MiniCPM-o 4.5 plugin (`model_executor/models/minicpmo_4_5/duplex/plugin.py`)
-is the only integration in this framework version. PersonaPlex and Nemotron
-VoiceChat still carry their pre-framework duplex code (runtime extension plus
-serving adapter) and are therefore not served over this framework yet: their
-pipelines declare no `duplex_plugin`, so they run turn-based until the
-follow-up PRs port them (RFC vllm-omni#7181, PR 2/3).
+Frame-based plugins use `engine/duplex/intermediate.py::build_duplex_append_prompt`
+for the shared request identity, sequencing and config snapshots. Token budgets,
+PCM framing and model-specific worker fields remain in the plugins.
+
+See [supported models and deployments](../serving/full_duplex_api.md#enable-full-duplex)
+for the current plugin integrations and deployment configurations.
 
 ## Serving
 

docs/getting_started/installation/README.md

@@ -1,7 +1,7 @@
 # Installation
 
 !!! important
-    vLLM-Omni is released against the matching upstream vLLM major/minor version. The 0.30 development line requires vLLM 0.30.x; the source-install instructions pin vLLM 0.30.0. Published vLLM-Omni 0.28.0 wheels and images still require vLLM 0.28.x.
+    vLLM-Omni is released against the matching upstream vLLM major/minor version. The vLLM-Omni 0.30.x release line uses vLLM 0.30.x; the source-install instructions pin vLLM 0.30.0.
 
 vLLM-Omni supports the following hardware platforms:
 

docs/getting_started/installation/gpu.md

@@ -34,9 +34,7 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
 
 ### Pre-built wheels
 
-The source-install instructions target vLLM 0.30.0. Until vLLM-Omni 0.30.0 artifacts are published, the CUDA/ROCm pre-built examples below remain on the matching 0.28.0 release pair.
-
-Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, 0.22.0, 0.23.0, 0.24.0, 0.25.0, 0.26.0, and 0.28.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
+Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, 0.22.0, 0.23.0, 0.24.0, 0.25.0, 0.26.0, 0.28.0, and 0.30.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
 
 === "NVIDIA CUDA"
 

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

@@ -5,7 +5,7 @@
 # --8<-- [end:requirements]
 # --8<-- [start:set-up-using-python]
 
-vLLM-Omni depends on the matching major/minor release of vLLM. The 0.30 development line uses vLLM 0.30.x. Published 0.28.0 wheels and images use vLLM 0.28.x.
+vLLM-Omni depends on the matching major/minor release of vLLM. The vLLM-Omni 0.30.x release line uses vLLM 0.30.x.
 
 !!! note
     PyTorch installed via `conda` will statically link `NCCL` library, which can cause issues when vLLM tries to use `NCCL`. See <gh-issue:8420> for more details.
@@ -18,22 +18,22 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new**
 
 #### Installation of vLLM
 
-These pre-built wheel instructions install the published vLLM-Omni 0.28.0 release. For the 0.30 development line, use the source-install instructions below.
+These pre-built wheel instructions install the published vLLM-Omni 0.30.0 release.
 
 vLLM-Omni is built based on vLLM. Please install it with command below.
 ```bash
-uv pip install vllm==0.28.0 --torch-backend=auto
+uv pip install vllm==0.30.0 --torch-backend=auto
 ```
 
 #### Installation of vLLM-Omni
 
 ```bash
-uv pip install vllm-omni==0.28.0
+uv pip install vllm-omni==0.30.0
 ```
 
 To run Gradio demos, also install the optional extras:
 ```bash
-uv pip install 'vllm-omni[demo]==0.28.0'
+uv pip install 'vllm-omni[demo]==0.30.0'
 ```
 
 # --8<-- [end:pre-built-wheels]
@@ -111,7 +111,7 @@ docker run --runtime nvidia --gpus 2 \
     --env "HF_TOKEN=$HF_TOKEN" \
     -p 8091:8091 \
     --ipc=host \
-    vllm/vllm-omni:v0.28.0 \
+    vllm/vllm-omni:v0.30.0 \
     vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```
 

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

@@ -10,9 +10,8 @@
 
 For ROCm, vLLM-Omni currently recommends the setup steps through Docker Images.
 
-vLLM-Omni depends on the matching major/minor release of vLLM. The 0.30
-development line uses vLLM 0.30.x. Published 0.28.0 wheels and images use vLLM
-0.28.x.
+vLLM-Omni depends on the matching major/minor release of vLLM. The vLLM-Omni
+0.30.x release line uses vLLM 0.30.x.
 
 The Dockerfile's `BASE_IMAGE` pin applies only to Docker builds. The
 `vllm-omni` package does not install vLLM as a dependency, so non-Docker source
@@ -23,12 +22,12 @@ installing vLLM-Omni, as shown below.
 
 #### Installation of vLLM
 
-These pre-built wheel instructions install the published vLLM-Omni 0.28.0 release. For the 0.30 development line, use the source-install instructions below.
+These pre-built wheel instructions install the published vLLM-Omni 0.30.0 release.
 
 vLLM-Omni is built based on vLLM. Please install it with command below.
 
 ```bash
-uv pip install vllm==0.28.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.28.0/rocm723
+uv pip install vllm==0.30.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.30.0/rocm723
 ```
 
 #### Installation of vLLM-Omni
@@ -37,7 +36,7 @@ uv pip install vllm==0.28.0+rocm723 --extra-index-url https://wheels.vllm.ai/roc
 # we need to add --no-build-isolation as the torch
 # is not obtained from pypi, we have to install using the
 # torch installed in our environment
-uv pip install vllm-omni==0.28.0
+uv pip install vllm-omni==0.30.0
 
 # Optional if want to run Qwen3 TTS
 uv pip uninstall onnxruntime # should be removed before we can install onnxruntime-rocm
@@ -150,6 +149,8 @@ vllm-omni-rocm
 
 vLLM-Omni offers an official docker image for deployment. These images are built on top of vLLM docker images and available on Docker Hub as [vllm/vllm-omni-rocm](https://hub.docker.com/r/vllm/vllm-omni-rocm/tags). The version of vLLM-Omni indicates which release of vLLM it is based on.
 
+The prebuilt ROCm image is published separately from the release pipeline: `v0.28.0` is the latest tag currently available on Docker Hub, and newer-tag availability is tracked in [#7405](https://github.com/vllm-project/vllm-omni/issues/7405).
+
 #### Launch vLLM-Omni Server
 
 Here's an example deployment command that has been verified on 2 x MI300's:

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

@@ -53,7 +53,7 @@ docker run --rm \
 
 # Inside the container, install vLLM-Omni from source
 cd /vllm-workspace
-git clone -b v0.30.0rc1 https://github.com/vllm-project/vllm-omni.git
+git clone -b v0.30.0 https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni
 pip install -v -e . --no-build-isolation
 # or VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e .

docs/models/supported_models.md

@@ -19,9 +19,9 @@ th {
 
 | Architecture | Models | Example HF Models | NVIDIA GPU | AMD GPU | Ascend NPU | Intel GPU | Recipe |
 | -------------- | -------- | ------------------- | ------------ | --------- | ----- | ----------- | -------- |
-| `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen3-Omni.md) |
+| `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | ✅︎ | | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen3-Omni.md) |
 | `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
-| `MingFlashOmniForConditionalGeneration` + `MingImagePipeline` | Ming-flash-omni-2.0 (omni-speech + imagegen<sup>1</sup>) | `Jonathan1909/Ming-flash-omni-2.0` | ✅︎ | | | | — |
+| `MingFlashOmniForConditionalGeneration` + `MingImagePipeline` | Ming-flash-omni-2.0 (omni-speech + imagegen<sup>1</sup>) | `Jonathan1909/Ming-flash-omni-2.0` | ✅︎ | | | ✅︎ | — |
 | `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` | ✅︎ | ✅︎ | | ✅︎ | — |
 | `InternVLAA1Pipeline` | InternVLA-A1 | `InternRobotics/InternVLA-A1-3B` | ✅︎ | ✅︎ | | | — |
 | `Gr00tN1d7Pipeline` | GR00T N1.7 | `nvidia/GR00T-N1.7-3B` | ✅︎ | | | | — |
@@ -33,9 +33,9 @@ th {
 | `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2511 | `Qwen/Qwen-Image-Edit-2511` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `GlmImagePipeline` | GLM-Image | `zai-org/GLM-Image` | ✅︎ | ✅︎ | | | — |
-| `ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
-| `Krea2Pipeline` | Krea 2 (Raw + Turbo) | `krea/Krea-2-Raw`, `krea/Krea-2-Turbo` | ✅︎ | | | ✅︎ | — |
-| `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Tongyi-MAI/Z-Image-Turbo.md) |
+| `Krea2Pipeline` | Krea 2 (Raw + Turbo) | `krea/Krea-2-Raw`, `krea/Krea-2-Turbo` | ✅︎ | | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/krea/Krea-2.md) |
+| `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Wan-AI/Wan2.2-T2V.md) |
 | `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `LingBotWorldCausalDMDPipeline` | LingBot-World 2.0 v1 (experimental) | `robbyant/lingbot-world-v2-14b-causal-fast-diffusers` | Experimental | | | | — |
 | `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
@@ -43,16 +43,16 @@ th {
 | `Wan22S2VPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `Wan22VACEPipeline` | Wan2.2-VACE | `Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers` | ✅︎ | | | ✅︎ | — |
-| `LTX2Pipeline` | LTX-2 / LTX-2.3 one-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | | — |
-| `LTX2TwoStagePipeline` | LTX-2 / LTX-2.3 ordinary two-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` + matching Lightricks LoRA and upsampler | ✅︎ | ✅︎ | | | — |
-| `LTX2DistilledOneStagePipeline` | LTX-2 / LTX-2.3 merged-distilled one-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled`, `diffusers/LTX-2.3-Distilled-Diffusers` | ✅︎ | ✅︎ | | | — |
-| `LTX2DistilledTwoStagePipeline` | LTX-2 / LTX-2.3 merged-distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled`, `diffusers/LTX-2.3-Distilled-Diffusers` + matching Lightricks upsampler | ✅︎ | ✅︎ | | | — |
-| `LTX2Pipeline` | LTX-2.5 Full/SFT one-stage T2V and I2V | `Lightricks/LTX-2.5-Diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.5.md) |
+| `LTX2Pipeline` | LTX-2 / LTX-2.3 one-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.md) |
+| `LTX2TwoStagePipeline` | LTX-2 / LTX-2.3 ordinary two-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` + matching Lightricks LoRA and upsampler | ✅︎ | ✅︎ | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.md) |
+| `LTX2DistilledOneStagePipeline` | LTX-2 / LTX-2.3 merged-distilled one-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled`, `diffusers/LTX-2.3-Distilled-Diffusers` | ✅︎ | ✅︎ | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.md) |
+| `LTX2DistilledTwoStagePipeline` | LTX-2 / LTX-2.3 merged-distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled`, `diffusers/LTX-2.3-Distilled-Diffusers` + matching Lightricks upsampler | ✅︎ | ✅︎ | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.md) |
+| `LTX2Pipeline` | LTX-2.5 Full/SFT one-stage T2V and I2V | `Lightricks/LTX-2.5-Diffusers` | ✅︎ | | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.5.md) |
 | `LTX2TwoStagePipeline` | LTX-2.5 Full/SFT two-stage T2V and I2V | `Lightricks/LTX-2.5-Diffusers` + `Lightricks/LTX-2.5` LoRA | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.5.md) |
 | `LTX2DistilledOneStagePipeline` | LTX-2.5 merged-distilled one-stage T2V and I2V | `Lightricks/LTX-2.5-Diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.5.md) |
 | `LTX2DistilledTwoStagePipeline` | LTX-2.5 merged-distilled two-stage T2V and I2V | `Lightricks/LTX-2.5-Diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.5.md) |
 | `LingBotVideoPipeline` | LingBot-Video dense and MoE T2I, T2V, TI2V | `robbyant/lingbot-video-dense-1.3b`, `robbyant/lingbot-video-moe-30b-a3b` | ✅︎ | | | | — |
-| `SanaVideoPipeline` | SANA-Video-2B T2V (native 480p/720p; Diffusers adapter validated at both resolutions) | `Efficient-Large-Model/SANA-Video_2B_480p_diffusers`, `Efficient-Large-Model/SANA-Video_2B_720p_diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/NVIDIA/SANA-Video-2B.md) |
+| `SanaVideoPipeline` | SANA-Video-2B T2V (native 480p/720p; Diffusers adapter validated at both resolutions) | `Efficient-Large-Model/SANA-Video_2B_480p_diffusers`, `Efficient-Large-Model/SANA-Video_2B_720p_diffusers` | ✅︎ | | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/NVIDIA/SANA-Video-2B.md) |
 | `SanaImageToVideoPipeline` | SANA-Video-2B I2V (native and Diffusers adapter validated at 480p/720p) | `Efficient-Large-Model/SANA-Video_2B_480p_diffusers`, `Efficient-Large-Model/SANA-Video_2B_720p_diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/NVIDIA/SANA-Video-2B.md) |
 | `MiniMaxH3Pipeline` | MiniMax H3 T2VA, first/last-keyframe FL2VA, and mixed-reference Ref2VA | `MiniMaxAI/MiniMax-H3` | ✅︎ | ✅︎ | | | [Published](https://recipes.vllm.ai/MiniMaxAI/MiniMax-H3) |
 | `Magi2Pipeline` | MAGI-2 native Preview-stage T2VA and I2VA (272p/540p) | `sand-ai/MAGI-2-preview` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/SandAI/MAGI-2-preview-L20X.md) |
@@ -65,18 +65,18 @@ th {
 | `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `LongCatVideoAvatarPipeline` | LongCat-Video-Avatar-1.5 A2V/AI2V (native single-speaker and multi-speaker AI2V/AVC) | `meituan-longcat/LongCat-Video-Avatar-1.5` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/meituan-longcat/LongCat-Video-Avatar-1.5.md) |
 | `BooguImagePipeline` | Boogu-Image | `Boogu/Boogu-Image-0.1-Base`, `Boogu/Boogu-Image-0.1-Edit`, `Boogu/Boogu-Image-0.1-Edit-Turbo` | ✅︎ | | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Boogu/Boogu-Image.md) |
-| `BooguImageTurboPipeline` | Boogu-Image Turbo distilled T2I | `Boogu/Boogu-Image-0.1-Turbo` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Boogu/Boogu-Image.md) |
+| `BooguImageTurboPipeline` | Boogu-Image Turbo distilled T2I | `Boogu/Boogu-Image-0.1-Turbo` | ✅︎ | | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Boogu/Boogu-Image.md) |
 | `StableDiffusionXLPipeline` | Stable-Diffusion-XL | `stabilityai/stable-diffusion-xl-base-1.0` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
-| `StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | ✅︎ | ✅︎ | | ✅︎ | — |
+| `StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | ✅︎ | ✅︎ | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/StabilityAI/Stable-Diffusion-3.5.md) |
 | `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ | — |
 | `AuKForConditionalGeneration` | AuK / AuK-Flash (instruction-driven speech generation and editing) | `tencent/AuK`, `tencent/AuK-Flash` (with `Qwen/Qwen2.5-Omni-3B`; assemble with `tools/prepare_auk_checkpoint.py`) | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Tencent/AuK-H100.md) |
 | `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | ✅︎ | — |
 | `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | ✅︎ | — |
 | `BreezeForConditionalGeneration` | Breeze-TTS-2 (voice design, voice clone, voice direction, CFG) | `BreezeBlue/Breeze-TTS-2` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/BreezeBlue/Breeze-TTS-2.md) |
-| `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | | — |
+| `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | ✅︎ | — |
 | `MammothModa2ForConditionalGeneration` | MammothModa2-Dev (AR-only image understanding) | `bytedance-research/MammothModa2-Dev` | ✅︎ | | | | — |
-| `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
-| `FluxKontextPipeline` | FLUX.1-Kontext-dev | `black-forest-labs/FLUX.1-Kontext-dev` | ✅︎ | ✅︎ | | | — |
+| `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/black-forest-labs/FLUX.2-klein.md) |
+| `FluxKontextPipeline` | FLUX.1-Kontext-dev | `black-forest-labs/FLUX.1-Kontext-dev` | ✅︎ | ✅︎ | | ✅︎ | — |
 | `FluxPipeline` | FLUX.1-dev | `black-forest-labs/FLUX.1-dev` | ✅︎ | ✅︎ | | ✅︎ | — |
 | `FluxPipeline` | FLUX.1-schnell | `black-forest-labs/FLUX.1-schnell` | ✅︎ | ✅︎ | | ✅︎ | — |
 | `OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` | ✅︎ | ✅︎ | | ✅︎ | — |
@@ -97,13 +97,13 @@ th {
 | `NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` | ✅︎ | ✅︎ | | ✅︎ | — |
 | `MiMoAudioModel` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` | ✅︎ | ✅︎ | | | — |
 | `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | | — |
-| `Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | ✅︎ | ✅︎ | | | — |
+| `Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | ✅︎ | ✅︎ | | ✅︎ | — |
 | `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | ✅︎ | — |
-| `Audio8TTSSlowARForConditionalGeneration` | Audio8 TTS Preview | `Audio8/Audio8-TTS-Preview-0.6b` | ✅︎ | | | | — |
-| `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only), SenseNova-U1.5 | `SenseNova/SenseNova-U1-8B-MoT`, `sensenova/SenseNova-U1.5-8B-MoT` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/SenseNova/SenseNova-U1.5.md) |
+| `Audio8TTSSlowARForConditionalGeneration` | Audio8 TTS Preview | `Audio8/Audio8-TTS-Preview-0.6b` | ✅︎ | | | ✅︎ | — |
+| `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only), SenseNova-U1.5 | `SenseNova/SenseNova-U1-8B-MoT`, `sensenova/SenseNova-U1.5-8B-MoT` | ✅︎ | | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/SenseNova/SenseNova-U1.5.md) |
 | `LancePipeline` | Lance | `bytedance-research/Lance` | ✅︎ | | | | — |
-| `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | | — |
-| `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | | — |
+| `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Tencent/HunyuanVideo-1.5.md) |
+| `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Tencent/HunyuanVideo-1.5.md) |
 | `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | | — |
 | `CovoAudioForConditionalGeneration` | Covo-Audio-Chat | `tencent/Covo-Audio-Chat` | ✅︎ | | | | — |
 | `MiniCPMO45OmniForConditionalGeneration` | MiniCPM-o 4.5 | `openbmb/MiniCPM-o-4_5` | ✅︎ | | ✅︎ | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/OpenBMB/MiniCPM-o-4_5.md) |

docs/serving/README.md

@@ -150,17 +150,22 @@ interchangeably.
 | Robot policy inference | `WS /v1/realtime/robot/openpi` | Send MessagePack observations and receive action arrays | [OpenPI Robot Policy](openpi_api.md) |
 
 All six routes are model- or configuration-dependent. In particular,
-`/v1/realtime` is full duplex when the model is a duplex model (its pipeline
-declares a `duplex_plugin` and the deploy configuration sets
-`session_mode: duplex`); a stock Realtime client needs no vendor query
-parameter, and `duplex=0` is the explicit opt-out. Such a server serves the
-websocket route plus `POST /v1/chat/completions`, and no other turn-based
-HTTP route. On a duplex server that chat route is not the turn-based path:
-each request runs on a short-lived duplex session, so it holds one of
-`duplex_session.max_sessions` for its lifetime and answers at the model's
-real-time pace -- see [Full Duplex](full_duplex_api.md). Clients should
-also verify the duplex capability payload because the query-parameter form
-falls back to the ordinary realtime handler when duplex is unavailable.
+`/v1/realtime` is full duplex when the server is a duplex server: the model's
+pipeline declares a `duplex_plugin` and its deploy configuration sets
+`session_mode: duplex` (with `session_mode: turn` the same model boots the
+ordinary turn-based serving stack instead). On a duplex server a stock
+Realtime client needs no vendor query parameter; `duplex=0` selects the
+turn-based Realtime handler, which a duplex server does not mount, so that
+connection is refused with `Realtime API is not available`. Such a server
+serves the websocket route plus `POST /v1/chat/completions`, `/v1/models`
+and `/health`, and no other turn-based HTTP route. The chat route is the
+ordinary chat service running on the duplex engine: a request is a
+turn-based generation on the same stages, served alongside the live
+websocket sessions; it opens no duplex session and holds no
+`duplex_session.max_sessions` slot -- see [Full Duplex](full_duplex_api.md).
+On a server that is not duplex, an explicit `?duplex=1` is refused rather
+than answered by the turn-based handler, so a client that requested duplex
+never silently gets the other protocol.
 
 ## Related Endpoints
 

docs/serving/full_duplex_api.md

@@ -72,12 +72,21 @@ does not switch engines. The Python `Omni` / `AsyncOmni` APIs are unchanged.
     `session.created.session.capabilities` is present before treating the
     connection as full duplex.
 
-**MiniCPM-o 4.5** (`vllm_omni/deploy/minicpmo_4_5.yaml`) is the only model
-served over this endpoint today. PersonaPlex and Nemotron VoiceChat still carry
-their pre-framework duplex code: their pipelines declare no `duplex_plugin`, so
-they run turn-based until the follow-up PRs port them to the plugin contract
+**MiniCPM-o 4.5** (`vllm_omni/deploy/minicpmo_4_5.yaml`), **Qwen3-Omni**
+(`vllm_omni/deploy/qwen3_omni_duplex.yaml`), and **Nemotron VoiceChat**
+(`vllm_omni/deploy/nemotron_labs_voicechat_duplex.yaml`) are served
+over this endpoint today. PersonaPlex still carries its pre-framework duplex
+code: its pipeline declares no `duplex_plugin`, so it runs turn-based until the
+follow-up PR ports it to the plugin contract
 (RFC [vllm-omni#7181](https://github.com/vllm-project/vllm-omni/issues/7181)).
 
+For Nemotron VoiceChat Python clients, use
+`vllm_omni.clients.nemotron_voicechat.create_duplex_session_config()` with
+`DuplexClient`. The preset selects 16 kHz float32 input, 22.05 kHz PCM16 output,
+and automatic responses. Append 1,280-sample (80 ms) input frames; when using
+`stream_pcm`, set `chunk_ms=80`. Pass `instructions` and optional `tools` to
+the preset before opening the session.
+
 JoyVL is a separate HTTP interaction orchestrator and does not use these
 WebSocket endpoints. See [Standalone Experimental Servers](standalone_servers.md).
 

docs/serving/realtime_duplex_api.md

@@ -8,11 +8,9 @@ history stays honest, and a dropped connection can resume the same session.
 This page covers how to run a duplex deployment, how to drive it from Python
 with `vllm_omni.clients.duplex.DuplexClient`, and the complete wire contract.
 
-The endpoint is served for models that ship a duplex plugin — including
-MiniCPM-o 4.5 and Qwen3-Omni — and just for deploy configurations that declare
-`session_mode: duplex`. PersonaPlex and Nemotron VoiceChat are ported to the
-plugin contract in follow-up PRs and are not served over this endpoint yet. The runtime architecture
-is described in [Full-Duplex Runtime (MiniCPM-o 4.5)](../design/fullduplex.md).
+The endpoint requires a duplex plugin and a deployment with `session_mode: duplex`.
+See [supported models and deployments](full_duplex_api.md#enable-full-duplex)
+and the [runtime architecture](../design/fullduplex.md).
 
 ## Qwen3-Omni conversation history
 
@@ -1340,9 +1338,6 @@ them out into typed events before they reach a client.
 
 ## Known Limitations
 
-- Only MiniCPM-o 4.5 is served over this endpoint today; PersonaPlex and
-  Nemotron VoiceChat arrive with the follow-up PRs that port them to the
-  plugin contract.
 - Several surfaces are capability-gated per model (see *Capability
   negotiation by model* above): PersonaPlex does not support session resume,
   barge-in, or audio truncation; Nemotron VoiceChat does not support barge-in

docs/serving/videos_api.md

@@ -120,7 +120,18 @@ request for execution.
 `DELETE /v1/videos/{video_id}` issues a bounded engine abort
 (`VLLM_OMNI_ABORT_TIMEOUT`, default 2s), then cancels the frontend
 task. Cancellation cleanup is also bounded and best-effort: it confirms
-the abort was queued, and the current request batch may still drain.
+the abort was submitted. In request execution mode, MiniMax-H3 checks
+cancellation at model boundaries,
+including before and after input preparation, after each denoising step, and
+before decode, and skips the remaining generation when cancelled. Input
+preparation or decoding already in progress may finish before the next boundary.
+Successful steps do not force device synchronization for cancellation.
+Parallel ranks agree before stopping; independent requests sharing a distributed
+AllGather offload wave can stop that wave early only when all its requests are
+cancelled, so a cancelled request cannot strand its live peers in a collective.
+In step execution mode, cancellation is handled by the scheduler between steps;
+the request-mode component-boundary checks do not apply within an active step.
+Other pipelines may still drain their current request batch.
 The job is then re-read so a completed save is not orphaned.
 
 ### Synchronous Response

recipes/LTX/LTX-2.5.md

@@ -89,6 +89,7 @@ Install matching vLLM and vLLM-Omni versions, and ensure `ffmpeg` and
 | NVIDIA B200 or H200 | Capacity-based recommendation; not yet verified | All four canonical pipelines |
 | NVIDIA GB200 or GB300 | Capacity-based recommendation; not yet verified | All four canonical pipelines |
 | NVIDIA H100 80 GB | FP8 recipe | Distilled one-stage at 960x544 |
+| Intel Arc Pro B70 32 GB | Verified | Offline Full/SFT one-stage T2V at 768x512 |
 
 The 1920x1088 two-stage examples require about 114 GB of peak GPU memory.
 80 GB GPUs do not have enough safety margin for the canonical two-stage
@@ -133,6 +134,33 @@ SP attention head counts must be divisible by `ulysses_degree`. Use each
 phase's `H, W` (half/full resolution for two-stage); non-divisible shapes fail.
 Each GPU must still fit the resident weights. Warm up with the production shape.
 
+### 1x Intel Arc Pro B70 32 GB (XPU)
+
+Offline Full/SFT one-stage T2V at 768x512, 17 frames. Select the legacy
+convolutional VAE with `ltx2_use_conv_vae`; the default DiffVAE decoder loads a
+NATTEN kernel, which has no XPU build. FP8 with layerwise offload streams both
+the transformer and the Gemma-4 text encoder.
+
+Environment: torch 2.14.0+xpu, vLLM 0.29.0, vLLM-Omni `main` at `d8d162d8`.
+
+```bash
+python examples/offline_inference/text_to_video/text_to_video.py \
+  --model Lightricks/LTX-2.5-Diffusers \
+  --model-class-name LTX2Pipeline \
+  --stage-overrides '{"0":{"extras":{"ltx2_use_conv_vae":true}}}' \
+  --prompt "A serene lakeside sunrise with mist over the water." \
+  --num-frames 17 \
+  --quantization fp8 \
+  --enable-layerwise-offload \
+  --vae-use-tiling \
+  --vae-use-slicing \
+  --enforce-eager \
+  --output ltx25_output.mp4
+```
+
+Peaks at 16.3 GiB in 163 s. Two-stage, distilled, and I2V paths were not
+exercised on this hardware.
+
 ## Offline inference
 
 Choose values from the pipeline table. For example, the distilled two-stage

recipes/LTX/LTX-2.md

@@ -165,6 +165,63 @@ curl -X POST http://localhost:8000/v1/videos/sync \
 Use `image_reference` for a URL or JSON-safe image reference. Do not provide it
 together with `input_reference`.
 
+## Hardware Support
+
+## XPU
+
+### 1x Intel Arc Pro B70 (32 GB)
+
+Offline one-stage T2V at 768x512, 17 frames. FP8 with layerwise offload streams
+both the transformer and the Gemma text encoder, which otherwise stays resident
+beside it.
+
+#### Environment
+
+- OS: Linux
+- Python: 3.10+
+- torch: 2.14.0+xpu
+- vLLM: 0.29.0
+- vLLM-Omni: `main` at `d8d162d8`
+
+#### Command
+
+```bash
+python examples/offline_inference/text_to_video/text_to_video.py \
+  --model diffusers/LTX-2.3-Diffusers \
+  --prompt "A serene lakeside sunrise with mist over the water." \
+  --num-frames 17 \
+  --quantization fp8 \
+  --enable-layerwise-offload \
+  --vae-use-tiling \
+  --vae-use-slicing \
+  --enforce-eager \
+  --output ltx23_output.mp4
+```
+
+The same flags cover every one-stage checkpoint; only `--model` changes. The
+distilled checkpoints additionally need their class and an explicit size,
+because their preset defaults to 1536x1024:
+
+```bash
+  --model diffusers/LTX-2.3-Distilled-Diffusers \
+  --model-class-name LTX2DistilledOneStagePipeline \
+  --height 512 --width 768
+```
+
+| Checkpoint | Peak | Time |
+| --- | ---: | ---: |
+| `Lightricks/LTX-2` | 14.8 GiB | 160 s |
+| `diffusers/LTX-2.3-Diffusers` | 16.4 GiB | 167 s |
+| `dg845/LTX-2.3-Diffusers` | 16.4 GiB | 155 s |
+| `diffusers/LTX-2.3-Distilled-Diffusers` | 14.9 GiB | 128 s |
+
+#### Verification
+
+Each command writes its `--output` MP4. Confirm it decodes and that the sampled
+frames match the prompt.
+
+
+
 ## Guidance
 
 One-stage and ordinary Stage 1 support independent video/audio CFG,

recipes/MammothModa2/MammothModa2.md

@@ -59,7 +59,9 @@ the measured results are summarized below.
 ### Optional FP8 AR KV cache
 
 For CUDA deployments, `mammoth_moda2_fp8_kv.yaml` is an opt-in preset that
-stores the Stage 0 AR KV cache as FP8 E4M3. Stage 1 remains on
+keeps the Stage 0 AR KV cache of decoder layer 0 in BF16 and stores the other
+27 layers as FP8 E4M3 (`kv_cache_dtype_skip_layers: ["0"]`; write the layer
+indices as quoted strings). Stage 1 remains on
 `kv_cache_dtype=auto`; its DiT execution is unaffected. This setting quantizes
 only the autoregressive KV cache. It is neither FP8 weight/activation
 quantization nor vLLM-Omni diffusion KV-cache quantization.
@@ -78,9 +80,9 @@ python examples/offline_inference/text_to_image/text_to_image.py \
   --output mammoth_t2i.png
 ```
 
-The preset was validated on one NVIDIA H800 80GB with CUDA and
-FlashAttention 3. The native and FP8 runs used the same model, code revision,
-and downstream configuration.
+The preset was first validated on one NVIDIA H800 80GB with CUDA and
+FlashAttention 3, with every layer in FP8. The native and FP8 runs used the
+same model, code revision, and downstream configuration.
 
 | Metric | Native BF16 (`kv_cache_dtype=auto`) | FP8 E4M3 (`fp8_e4m3`) | Change |
 | --- | ---: | ---: | ---: |
@@ -97,9 +99,30 @@ twice as many tokens because FP8 reduces the bytes per cached token. This is a
 capacity/concurrency tradeoff: the measured AR and end-to-end latencies were
 higher than the native-BF16 baseline.
 
-1024x1024 fixed-seed smoke test completed successfully with no obvious visual
-failure. FP8 is lossy, so numerical or image-quality equivalence with BF16 is
-not implied.
+On H800, a 1024x1024 fixed-seed smoke test with every layer in FP8 completed
+successfully with no obvious visual failure. FP8 is lossy, so numerical or
+image-quality equivalence with BF16 is not implied.
+
+On one A800 80GB (vLLM 0.30.0), FP8 KV layers run on FlashInfer and BF16 layers
+on FlashAttention 2. Text-to-image at 1024x1024, 50 steps, guidance 4.0, seeds 42
+and 1-5, with a studio tabby cat prompt and a peephole-view Samoyed prompt;
+an image counts when it shows the prompted subject and scene.
+
+| Stage 0 KV cache | GPU KV cache size | Cat images that follow the prompt | Samoyed images that follow the prompt |
+| --- | ---: | ---: | ---: |
+| BF16 on all 28 layers (`auto`) | 147,408 tokens | 6/6 | 6/6 |
+| FP8 E4M3 on all 28 layers | 294,816 tokens | 1/6 | 0/6 |
+| Layer 0 BF16, other 27 layers FP8 E4M3 (this preset) | 284,640 tokens | 6/6 | 6/6 |
+
+The KV cache sizes above still count the 28 attention layers of the replaced
+Qwen-VL language model, which #8095 removes: with it, BF16 goes from 147,408 to
+294,816 tokens and all-FP8 from 294,816 to 589,632. The prompt-following
+columns do not depend on it.
+
+With all 28 layers in FP8, most cat images become a framed print on a wall.
+Keeping layer 0 in BF16 restores them; keeping only layer 27, which has the
+largest key magnitude, does not. The H800 numbers above were measured with all 28
+layers in FP8; this preset has been run on A800 only.
 
 ### 1x L40S 48GB
 

recipes/Qwen/Qwen-Image.md

@@ -7,7 +7,8 @@
 - Vendor: Qwen
 - Model: `Qwen/Qwen-Image`
 - Task: Text-to-image generation
-- Mode: Online serving with optional step-wise continuous batching
+- Mode: Online serving with optional step-wise continuous batching; offline
+  inference
 - Maintainer: Community
 
 ## When to use this recipe
@@ -25,11 +26,14 @@ the benchmark assets already bundled in this repository.
   [`examples/online_serving/text_to_image/README.md`](../../examples/online_serving/text_to_image/README.md)
 - Related benchmark:
   [`benchmarks/diffusion/diffusion_benchmark_serving.py`](../../benchmarks/diffusion/diffusion_benchmark_serving.py)
+- Offline inference example used by the XPU section:
+  [`examples/offline_inference/text_to_image/text_to_image.py`](../../examples/offline_inference/text_to_image/text_to_image.py)
 
 ## Hardware Support
 
-This recipe currently documents one CUDA GPU serving configuration. Extend it
-with more hardware sections as community validation lands.
+This recipe documents CUDA GPU serving configurations and one Intel XPU
+offline-inference configuration. Extend it with more hardware sections as
+community validation lands.
 
 ## GPU
 
@@ -235,3 +239,59 @@ CUTLASS was the fastest validated backend for this checkpoint.
 - The benchmark table intentionally reports Qwen-Image-2512 as
   concurrency-1 request-level data; do not compare it directly with an online
   HTTP concurrency benchmark.
+
+## XPU
+
+### 1x Intel Arc Pro B70 (32 GB)
+
+Offline text-to-image at 1024x1024, with layerwise offload streaming the
+transformer blocks so the 20B model fits one card.
+
+#### Environment
+
+- OS: Linux
+- Python: 3.10+
+- torch: 2.13.0+xpu
+- vLLM: 0.29.0 (`98dff2a8`)
+- vLLM-Omni: `main` at `4c7a98c2`
+
+#### Command
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model Qwen/Qwen-Image \
+  --prompt "a cup of coffee on the table" \
+  --num-inference-steps 50 \
+  --enable-layerwise-offload \
+  --vae-use-tiling \
+  --vae-use-slicing \
+  --enforce-eager \
+  --output qwen_image_output.png
+```
+
+Given spare cards, `--enable-cpu-offload` shards the transformer instead of
+streaming it. It shortens the denoise loop but not the run, because the
+per-request model swap costs more than the loop saves.
+
+```bash
+  --tensor-parallel-size 4 \
+  --enable-cpu-offload
+```
+
+| Offload | Cards | Peak VRAM | Time |
+| --- | ---: | ---: | ---: |
+| `--enable-layerwise-offload` | 1 | 16.6 GiB | 50 s |
+| `--enable-cpu-offload` | 2 | 26.4 GiB | 85 s |
+| `--enable-cpu-offload` | 4 | 19.8 GiB | 70 s |
+| `--enable-cpu-offload` | 8 | 16.9 GiB | 87 s |
+
+#### Verification
+
+Confirm `qwen_image_output.png` is written and looks coherent for the prompt.
+
+#### Notes
+
+- Memory usage: peak 16.6 GiB, about 55 s per image.
+- Known limitations: only offline generation was qualified. Step-wise continuous
+  batching, ModelOpt quantization, and online serving are out of scope for this
+  profile.

recipes/Qwen/Qwen3-Omni.md

@@ -53,7 +53,7 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
   --deploy-config /path/to/your_qwen3_omni_overrides.yaml
 ```
 
-#### Runtime tuning
+### Runtime tuning
 
 Prefer CLI overrides for day-to-day tuning:
 
@@ -75,7 +75,7 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
   }'
 ```
 
-#### Stage-based launch (one stage per process)
+### Stage-based launch (one stage per process)
 
 Use three terminals (one per stage). Start with the default commands below, then
 add `--max-num-seqs` only if you need explicit per-stage concurrency control.
@@ -135,7 +135,7 @@ CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
 
 If you use custom deploy YAML, add `--deploy-config` to each stage command.
 
-#### Verification
+### Verification
 
 After server startup, run a multimodal example client:
 
@@ -185,7 +185,7 @@ python examples/online_serving/qwen3_omni/openai_realtime_client.py \
   --output-wav realtime_output.wav
 ```
 
-#### Benchmark with `vllm bench`
+### Benchmark with `vllm bench`
 
 After the server is up, you can run online serving benchmarks with
 `vllm bench serve --omni`.
@@ -240,8 +240,54 @@ vllm bench serve \
   --percentile-metrics ttft,tpot,itl,e2el
 ```
 
-#### Notes
+### Notes
 
 - `/v1/realtime` is unsupported while `async_chunk` is enabled.
 - The default deploy uses `SharedMemoryConnector`; this is for single-host
   stage wiring.
+
+## XPU
+
+### 5x Intel Arc Pro B70 (32 GB)
+
+Offline end-to-end text query returning both text and speech. The 62 GiB of
+BF16 thinker weights need four cards, so the `xpu:` block of
+`vllm_omni/deploy/qwen3_omni_moe.yaml` runs the thinker at TP=4 on cards 0-3
+and the talker and code2wav on card 4. It is the registry default here, so
+`--deploy-config` is not needed.
+
+#### Environment
+
+- OS: Linux
+- Python: 3.10+
+- torch: 2.13.0+xpu
+- vLLM: 0.29.0 (`98dff2a8`)
+- vLLM-Omni: `main` at `4c7a98c2`
+
+#### Command
+
+The raised timeouts absorb the first decode step, which JIT-compiles the
+Triton MoE kernels.
+
+```bash
+VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS=1800 \
+python examples/offline_inference/qwen3_omni/end2end.py \
+  --model Qwen/Qwen3-Omni-30B-A3B-Instruct \
+  --query-type text \
+  --stage-init-timeout 1800 \
+  --init-timeout 1800 \
+  --output-wav qwen3omni_output
+```
+
+#### Verification
+
+`qwen3omni_output/` holds one `.txt` answer and the matching `.wav`. Confirm
+the transcript of the audio tracks the text response.
+
+#### Notes
+
+- Runtime: about 260 s for a single text query.
+- The thinker MoE layers have no tuned Triton config for this shape and log
+  `Using default MoE config. Performance might be sub-optimal!`.
+- Known limitations: only `--query-type text` was qualified. The audio, image,
+  and video query types and online serving are out of scope for this profile.

recipes/README.md

@@ -45,6 +45,7 @@ recipes/
 | [`Baidu/ERNIE-Image.md`](./Baidu/ERNIE-Image.md) | Text-to-image online serving (ERNIE-Image 8B) | 1x or 2x RTX 4090 24GB |
 | [`Bagel/BAGEL-7B-MoT.md`](./Bagel/BAGEL-7B-MoT.md) | Text-to-image with shared online/offline examples | 1x A100 80GB / 2x CUDA GPUs |
 | [`black-forest-labs/FLUX.1-schnell.md`](./black-forest-labs/FLUX.1-schnell.md) | Offline text-to-image generation with CPU offload | 1x NVIDIA GeForce RTX 5090 32GB |
+| [`black-forest-labs/FLUX.2-klein.md`](./black-forest-labs/FLUX.2-klein.md) | Offline image editing (FLUX.2-klein-4B) | 1x Intel Arc Pro B70 32GB |
 | [`Boogu/Boogu-Image.md`](./Boogu/Boogu-Image.md) | Text-to-image online serving (Boogu-Image-0.1-Base) | 1x A100/H100 40GB+ |
 | [`BosonAI/Higgs-Audio-V3-TTS.md`](./BosonAI/Higgs-Audio-V3-TTS.md) | Online + offline multilingual TTS with voice cloning | 1x H100 80GB |
 | [`BreezeBlue/Breeze-TTS-2.md`](./BreezeBlue/Breeze-TTS-2.md) | Online + offline TTS (plain / voice design / voice clone / voice direction, streaming PCM) | 1x NVIDIA L20 48GB |
@@ -62,9 +63,9 @@ recipes/
 | [`MiniMaxAI/MiniMax-H3.md`](./MiniMaxAI/MiniMax-H3.md) | T2VA, FL2VA, image+audio Ref2VA, and multi-video Ref2VA serving | 1x GPU with CPU offload / 4x B300 |
 | [`MiniMaxAI/MiniMax-H3-Spark-GB10.md`](./MiniMaxAI/MiniMax-H3-Spark-GB10.md) | MiniMax H3 T2VA and Ref2VA on DGX Spark | 1x or 2x DGX Spark (GB10) |
 | [`MiniMaxAI/MiniMax-H3-MUSA.md`](./MiniMaxAI/MiniMax-H3-MUSA.md) | MiniMax H3 T2VA, FL2VA, and Ref2VA on MUSA | 4x MTT S5000 (TP4 Ref2VA profile) |
-| [`krea/Krea-2.md`](./krea/Krea-2.md) | Text-to-image (Turbo + Raw), offline + online, with LoRA | 1x H100 80GB |
-| [`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 |
-| [`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 |
+| [`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 |
 | [`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 |
@@ -82,10 +83,10 @@ recipes/
 | [`NVIDIA/Nemotron-Labs-Audex.md`](./NVIDIA/Nemotron-Labs-Audex.md) | TTS / text-to-audio / audio understanding / cascaded S2S (Audex 2B + 30B-A3B) | 1x H100 80GB |
 | [`NVIDIA/NemotronLabs-VoiceChat.md`](./NVIDIA/NemotronLabs-VoiceChat.md) | Offline speech-to-speech voice chat (11B, frame-locked 12.5 Hz, 3-stage thinker/talker/code2wav) | 1x H100 80GB |
 | [`NVIDIA/SANA-WM.md`](./NVIDIA/SANA-WM.md) | First-frame image-to-video serving with camera control | 1x 24GB+ CUDA GPU; ~40GB disk |
-| [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt mixed FP8/NVFP4 | 1x A100 80GB / 2x B200 |
+| [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt mixed FP8/NVFP4; offline text-to-image on XPU | 1x A100 80GB / 2x B200 / 1x Intel Arc Pro B70 32GB |
 | [`Qwen/Qwen-Image-2512.md`](./Qwen/Qwen-Image-2512.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt FP8 / mixed FP8/NVFP4 | 1x A800 80GB / 2x B200 |
 | [`Qwen/Qwen-Image-Edit.md`](./Qwen/Qwen-Image-Edit.md) | Text-guided single-image editing | 1x or 2x H200 141GB |
-| [`Qwen/Qwen3-Omni.md`](./Qwen/Qwen3-Omni.md) | Online serving for multimodal chat | 1x A100 80GB |
+| [`Qwen/Qwen3-Omni.md`](./Qwen/Qwen3-Omni.md) | Online serving for multimodal chat; offline end-to-end text-to-speech on XPU | 1x A100 80GB / 5x Intel Arc Pro B70 32GB |
 | [`Qwen/Qwen3-TTS.md`](./Qwen/Qwen3-TTS.md) | Text-to-speech serving (CustomVoice / VoiceDesign / Base) | 1x H100/A100 80GB / 1x AMD MI300X |
 | [`SenseNova/SenseNova-U1.md`](./SenseNova/SenseNova-U1.md) | Unified image generation and understanding | 1x or 2x H200 / 1x AMD MI300X |
 | [`SenseNova/SenseNova-U1.5.md`](./SenseNova/SenseNova-U1.5.md) | Unified image generation and understanding, with an 8-step distilled LoRA | 1x A800 80GB |
@@ -93,13 +94,15 @@ recipes/
 | [`Tencent/AuK-H100.md`](./Tencent/AuK-H100.md) | Offline instruction-driven TTS, speech editing, enhancement and separation (AuK base and Flash) | 1x H100 80GB |
 | [`Tencent/Covo-Audio-Chat.md`](./Tencent/Covo-Audio-Chat.md) | Online serving for audio chat | 1x A100 80GB |
 | [`Tencent/HunyuanImage-3.0-Instruct.md`](./Tencent/HunyuanImage-3.0-Instruct.md) | DiT-only text-to-image serving and benchmark, including ModelOpt mixed FP8/NVFP4 | 4x H100/H800 80GB / 2x B200 |
-| [`Wan-AI/Wan2.2-T2V.md`](./Wan-AI/Wan2.2-T2V.md) | Text-to-video serving (Wan2.2 14B) with Skip-Softmax sparse attention | 1x datacenter Blackwell (B200/B300) |
+| [`Tencent/HunyuanVideo-1.5.md`](./Tencent/HunyuanVideo-1.5.md) | Offline text-to-video (HunyuanVideo-1.5 Diffusers) | 1x Intel Arc Pro B70 32GB |
+| [`Tongyi-MAI/Z-Image-Turbo.md`](./Tongyi-MAI/Z-Image-Turbo.md) | Offline few-step text-to-image (Z-Image-Turbo) | 1x Intel Arc Pro B70 32GB |
+| [`Wan-AI/Wan2.2-T2V.md`](./Wan-AI/Wan2.2-T2V.md) | Text-to-video serving (Wan2.2 14B) with Skip-Softmax sparse attention | 8x H20/H100/A100 / 1x datacenter Blackwell (B200/B300) / 1x Intel Arc Pro B70 32GB |
 | [`Wan-AI/Wan2.2-TI2V-5B-A100.md`](./Wan-AI/Wan2.2-TI2V-5B-A100.md) | Unified 720p text-to-video and image-to-video serving (Wan2.2 5B) | 1x A100-SXM4-80GB |
 | [`Wan-AI/Wan2.2-I2V.md`](./Wan-AI/Wan2.2-I2V.md) | Image-to-video serving (Wan2.2 14B) | 8x Ascend NPU (A2/A3) |
 | [`Wan-AI/Wan2.2-S2V.md`](./Wan-AI/Wan2.2-S2V.md) | Speech-to-video serving (Wan2.2 14B) | 2x A100/H100 80GB |
 | [`Wan-AI/Wan2.1-VACE.md`](./Wan-AI/Wan2.1-VACE.md) | Unified T2V, I2V, V2LF, FLF2V, inpaint, and R2V | 1x RTX 5090 (1.3B) / 1x L40S 48GB with layerwise offload (14B) |
 | [`XiaomiMiMo/MiMo-Audio.md`](./XiaomiMiMo/MiMo-Audio.md) | Offline + online omni audio (TTS, ASR, dialogue) | 1x RTX 5090 / 5090D 32GB |
-| [`StabilityAI/Stable-Diffusion-3.5.md`](./StabilityAI/Stable-Diffusion-3.5.md) | Text-to-image serving (SD 3.5-medium and SD 3.5-large) | 1x RTX A6000 48GB |
+| [`StabilityAI/Stable-Diffusion-3.5.md`](./StabilityAI/Stable-Diffusion-3.5.md) | Text-to-image serving (SD 3.5-medium and SD 3.5-large) | 1x RTX A6000 48GB / 1x Intel Arc Pro B70 32GB |
 | [`zai-org/GLM-TTS.md`](./zai-org/GLM-TTS.md) | Online serving for Chinese/English zero-shot voice-cloned TTS | 1x A40 48GB / 2x non-standard RTX 4090 48GB (49,140 MiB each) |
 | [`GLM/GLM-Image.md`](./GLM/GLM-Image.md) | Online serving for image generation | 1x A800 80GB / 2x A800 80GB |
 | [`JD/JoyAI-VL-Interaction.md`](./JD/JoyAI-VL-Interaction.md) | Real-time streaming video-language interaction (proactive speak/silence/delegate) | 1x GPU 24GB+ |

recipes/StabilityAI/Stable-Diffusion-3.5.md

@@ -137,3 +137,39 @@ print('saved output.png')
 - Image quality is noticeably better than medium, with sharper details.
 - `negative_prompt` requires `guidance_scale > 1` to take effect; SD3 defaults to 1.0 (no CFG).
 - 2048×2048 fits within 48 GB but leaves little headroom (~45.5 GiB used out of 49 GiB).
+
+## XPU
+
+### 1x Intel Arc Pro B70 (32 GB) — stable-diffusion-3.5-medium
+
+Offline text-to-image at 1024x1024 with BF16 weights fully resident.
+
+#### Environment
+
+- OS: Linux
+- Python: 3.10+
+- torch: 2.13.0+xpu
+- vLLM: 0.29.0 (`98dff2a8`)
+- vLLM-Omni: `main` at `4c7a98c2`
+
+#### Command
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model stabilityai/stable-diffusion-3.5-medium \
+  --prompt "a sunset over mountains, photorealistic" \
+  --num-inference-steps 28 \
+  --vae-use-tiling \
+  --output sd35_medium_output.png
+```
+
+#### Verification
+
+Confirm `sd35_medium_output.png` is written and matches the prompt.
+
+#### Notes
+
+- Memory usage: peak 19.5 GiB against 29.8 GiB free.
+- Known limitations: only `stable-diffusion-3.5-medium` offline generation was
+  qualified. `stable-diffusion-3.5-large` and online serving are out of scope
+  for this profile.

recipes/Tencent/HunyuanVideo-1.5.md

@@ -0,0 +1,72 @@
+# HunyuanVideo 1.5
+
+> Text-to-video with the HunyuanVideo-1.5 Diffusers checkpoints
+
+## Summary
+
+- Vendor: Tencent (community Diffusers ports)
+- Model: `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v`
+- Task: Text-to-video
+- Mode: Offline inference
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe to generate a short clip with HunyuanVideo-1.5 on a memory-
+constrained card, trading output size for capacity.
+
+## References
+
+- Model card:
+  <https://huggingface.co/hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v>
+- Shared runnable example:
+  [`examples/offline_inference/text_to_video/text_to_video.py`](../../examples/offline_inference/text_to_video/text_to_video.py)
+- CPU offload guide:
+  [`docs/user_guide/diffusion/cpu_offload.md`](../../docs/user_guide/diffusion/cpu_offload.md)
+
+## Hardware Support
+
+This recipe documents one validated Intel XPU configuration. Extend it with
+more hardware sections as community validation lands.
+
+## XPU
+
+### 1x Intel Arc Pro B70 (32 GB)
+
+BF16 weights with model-level CPU offload at a reduced 320x240, 17-frame
+shape. The 480p and 720p T2V and I2V checkpoints all use the same command; only
+`--model` changes.
+
+#### Environment
+
+- OS: Linux
+- Python: 3.10+
+- torch: 2.13.0+xpu
+- vLLM: 0.29.0 (`98dff2a8`)
+- vLLM-Omni: `main` at `4c7a98c2`
+
+#### Command
+
+```bash
+python examples/offline_inference/text_to_video/text_to_video.py \
+  --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v \
+  --prompt "A serene lakeside sunrise with mist over the water." \
+  --height 240 --width 320 --num-frames 17 \
+  --num-inference-steps 50 \
+  --guidance-scale 6.0 \
+  --enable-cpu-offload \
+  --vae-use-tiling \
+  --vae-use-slicing \
+  --output hunyuan_video_15_output.mp4
+```
+
+#### Verification
+
+Confirm `hunyuan_video_15_output.mp4` decodes and that the sampled frames match
+the prompt.
+
+#### Notes
+
+- Memory usage: 18.3 GiB loaded, 25.7 GiB peak, about 39 s per clip.
+- Known limitations: the native 720p geometry does not fit a 32 GB card with
+  CPU offload. Online serving is out of scope for this profile.

recipes/Tongyi-MAI/Z-Image-Turbo.md

@@ -0,0 +1,68 @@
+# Z-Image Turbo
+
+> Few-step text-to-image with `Tongyi-MAI/Z-Image-Turbo`
+
+## Summary
+
+- Vendor: Tongyi-MAI
+- Model: `Tongyi-MAI/Z-Image-Turbo`
+- Task: Text-to-image
+- Mode: Offline inference
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe to generate 1024x1024 images with the distilled Z-Image Turbo
+checkpoint through the shared offline text-to-image example.
+
+## References
+
+- Model card: <https://huggingface.co/Tongyi-MAI/Z-Image-Turbo>
+- Shared runnable example:
+  [`examples/offline_inference/text_to_image/text_to_image.py`](../../examples/offline_inference/text_to_image/text_to_image.py)
+- CPU offload guide:
+  [`docs/user_guide/diffusion/cpu_offload.md`](../../docs/user_guide/diffusion/cpu_offload.md)
+
+## Hardware Support
+
+This recipe documents one validated Intel XPU configuration. Extend it with
+more hardware sections as community validation lands.
+
+## XPU
+
+### 1x Intel Arc Pro B70 (32 GB)
+
+BF16 weights with model-level CPU offload and VAE tiling/slicing, 25 steps.
+
+#### Environment
+
+- OS: Linux
+- Python: 3.10+
+- torch: 2.13.0+xpu
+- vLLM: 0.29.0 (`98dff2a8`)
+- vLLM-Omni: `main` at `4c7a98c2`
+
+#### Command
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model Tongyi-MAI/Z-Image-Turbo \
+  --prompt "a cup of coffee on the table" \
+  --num-inference-steps 25 \
+  --enable-cpu-offload \
+  --vae-use-tiling \
+  --vae-use-slicing \
+  --enforce-eager \
+  --output z_image_turbo_output.png
+```
+
+#### Verification
+
+Confirm `z_image_turbo_output.png` is written as a 1024x1024 PNG matching the
+prompt.
+
+#### Notes
+
+- Memory usage: peak 15.6 GiB, about 58 s per image.
+- Known limitations: only offline generation was qualified. Online serving is
+  out of scope for this profile.

recipes/Wan-AI/Wan2.2-T2V.md

@@ -85,14 +85,78 @@ For online serving client examples and request formats, see
 #### Notes
 
 - **Key flags:**
-  - `--omni` — enables vLLM-Omni diffusion serving.
-  - `--use-hsdp` — Hybrid Sharded Data Parallelism for the 14B DiT weights.
-  - `--usp <N>` — Unified (Ulysses) Sequence Parallelism degree.
-  - `--cfg-parallel-size <N>` — CFG parallelism; set to 2 for the official model,
+    - `--omni` — enables vLLM-Omni diffusion serving.
+    - `--use-hsdp` — Hybrid Sharded Data Parallelism for the 14B DiT weights.
+    - `--usp <N>` — Unified (Ulysses) Sequence Parallelism degree.
+    - `--cfg-parallel-size <N>` — CFG parallelism; set to 2 for the official model,
     omit for distilled checkpoints.
-  - `--vae-patch-parallel-size 8` / `--vae-use-tiling` — parallel + tiled VAE
+    - `--vae-patch-parallel-size 8` / `--vae-use-tiling` — parallel + tiled VAE
     decoding; disabling patch parallelism can significantly increase VAE latency.
 
+## XPU
+
+### 1x Intel Arc Pro B70 (32 GB)
+
+Offline text-to-video at a reduced 576x320, 17-frame shape. Layerwise offload
+streams the two 14B expert transformers; VAE slicing bounds the decode
+activations.
+
+#### Environment
+
+- OS: Linux
+- Python: 3.10+
+- torch: 2.13.0+xpu
+- vLLM: 0.29.0 (`98dff2a8`)
+- vLLM-Omni: `main` at `4c7a98c2`
+
+#### Command
+
+```bash
+python examples/offline_inference/text_to_video/text_to_video.py \
+  --model Wan-AI/Wan2.2-T2V-A14B-Diffusers \
+  --prompt "A serene lakeside sunrise with mist over the water." \
+  --negative-prompt "blurry, noisy, artifacts, distorted, low quality, streaks, sparkles, blocky, pixelated" \
+  --height 320 --width 576 --num-frames 17 \
+  --num-inference-steps 60 \
+  --guidance-scale 4.0 \
+  --flow-shift 5.0 \
+  --boundary-ratio 0.875 \
+  --enable-layerwise-offload \
+  --vae-use-slicing \
+  --enforce-eager \
+  --seed 12345 \
+  --output wan22_t2v_output.mp4
+```
+
+Given spare cards, `--enable-cpu-offload` shards the two expert transformers
+instead of streaming them, which nearly halves the denoise loop. Four is the
+fastest count; eight is slower.
+
+```bash
+  --tensor-parallel-size 4 \
+  --enable-cpu-offload
+```
+
+| Offload | Cards | Peak VRAM | Denoise |
+| --- | ---: | ---: | ---: |
+| `--enable-layerwise-offload` | 1 | 9.9 GiB | 160 s |
+| `--enable-cpu-offload` | 2 | 18.9 GiB | 113 s |
+| `--enable-cpu-offload` | 4 | 12.4 GiB | 86 s |
+
+#### Verification
+
+Confirm `wan22_t2v_output.mp4` decodes and that frame 0 is clean.
+
+#### Notes
+
+- Key flags: `--boundary-ratio 0.875` sets the high-noise to low-noise expert
+  switch. `--flow-shift 5.0` with `--guidance-scale 4.0` over 60 steps clears
+  the frame-0 VAE corruption seen at the stock 50-step, `flow_shift=12.0`
+  settings.
+- Known limitations: only offline T2V at 576x320 was qualified. Larger shapes,
+  TI2V, HSDP, CFG parallelism, and online serving are out of scope for this
+  profile.
+
 ## Accelerating with TRTLLM_ATTN + Skip-Softmax (datacenter Blackwell)
 
 On **datacenter Blackwell** (sm_100 / sm_103), the mask-free Wan pipeline
@@ -153,7 +217,7 @@ instead of `target_sparsity`):
 ### Skip-Softmax controls
 
 | Key | Valid values | Meaning |
-|---|---|---|
+| --- | --- | --- |
 | `target_sparsity` | finite, `[0, 1]` | Operating point on the checkpoint's calibrated curve (`a·exp(b·s)`). |
 | `threshold` | finite, `≥ 0` | Direct skip threshold, no calibration needed. Mutually exclusive with `target_sparsity`. |
 | `disabled_until_timestep` | finite, `[0, 1]` | Keeps early, high-noise steps dense; skip turns on once the normalized timestep `t ≤ D`. |
@@ -169,7 +233,7 @@ Measured on B300 / SM103, 1280×720 / 81f / 50 steps, torch.compile, official
 ModelOpt calibration (speedup and divergence relative to dense):
 
 | config | speedup | LPIPS (divergence vs dense) |
-|---|---|---|
+| --- | --- | --- |
 | dense | 1.000x | - |
 | `s=0.65, D=0.86` | 1.062x | 0.041 |
 | `s=0.75, D=1.00` | 1.181x | 0.377 |
@@ -181,6 +245,6 @@ ModelOpt calibration (speedup and divergence relative to dense):
   `flashinfer`. Elsewhere, or without FlashInfer, selecting `TRTLLM_ATTN` raises
   rather than silently degrading.
 - **Known limitations:**
-  - Skip-Softmax is not supported under ring sequence parallelism; use Ulysses SP
+    - Skip-Softmax is not supported under ring sequence parallelism; use Ulysses SP
     (`--usp`), or run Skip-Softmax without ring.
-  - `TRTLLM_ATTN` is BF16-only (no attention-level quantization).
+    - `TRTLLM_ATTN` is BF16-only (no attention-level quantization).

recipes/black-forest-labs/FLUX.2-klein.md

@@ -0,0 +1,69 @@
+# FLUX.2-klein
+
+> Text-guided image editing with `black-forest-labs/FLUX.2-klein-4B`
+
+## Summary
+
+- Vendor: Black Forest Labs
+- Model: `black-forest-labs/FLUX.2-klein-4B`
+- Task: Image editing (image-to-image)
+- Mode: Offline inference
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe to run single-image edits with FLUX.2-klein-4B through the
+shared offline image-edit example.
+
+## References
+
+- Model card: <https://huggingface.co/black-forest-labs/FLUX.2-klein-4B>
+- Shared runnable example:
+  [`examples/offline_inference/image_to_image/image_edit.py`](../../examples/offline_inference/image_to_image/image_edit.py)
+- CPU offload guide:
+  [`docs/user_guide/diffusion/cpu_offload.md`](../../docs/user_guide/diffusion/cpu_offload.md)
+
+## Hardware Support
+
+This recipe documents one validated Intel XPU configuration. Extend it with
+more hardware sections as community validation lands.
+
+## XPU
+
+### 1x Intel Arc Pro B70 (32 GB)
+
+BF16 weights with model-level CPU offload. `--image` takes any RGB image.
+
+#### Environment
+
+- OS: Linux
+- Python: 3.10+
+- torch: 2.13.0+xpu
+- vLLM: 0.29.0 (`98dff2a8`)
+- vLLM-Omni: `main` at `4c7a98c2`
+
+#### Command
+
+```bash
+python examples/offline_inference/image_to_image/image_edit.py \
+  --model black-forest-labs/FLUX.2-klein-4B \
+  --image test_input.png \
+  --prompt "Add a sunset sky with orange and purple clouds" \
+  --num-inference-steps 28 \
+  --guidance-scale 3.5 \
+  --enable-cpu-offload \
+  --enforce-eager \
+  --seed 42 \
+  --output image_edit_output.png
+```
+
+#### Verification
+
+Confirm `image_edit_output.png` follows the prompt while preserving the rest
+of the input image.
+
+#### Notes
+
+- Memory usage: 7.7 GiB loaded in 8.4 s, about 18 s per edit.
+- Known limitations: only offline editing of the 4B checkpoint was qualified.
+  `FLUX.2-klein-9B` and online serving are out of scope for this profile.

recipes/inclusionAI/Ming-Image.md

@@ -27,7 +27,7 @@ For Design-Layer, the first returned image is the reconstructed composite and th
 - Python: 3.10+
 - CUDA: 13.0
 - vLLM version: 0.29.0
-- vLLM-Omni version or commit: 59db6a428
+- vLLM-Omni version or commit: 6daf5b30f
 - 2x H100 80GB (1xH100 to be validated)
 
 ### Commands
@@ -52,7 +52,7 @@ curl -s http://127.0.0.1:8091/v1/chat/completions \
     "modalities": ["image"],
     "extra_body": {
       "height": 1024, "width": 1024,
-      "steps": 12, "cfg": 1.0, "seed": 42
+      "num_inference_steps": 12, "guidance_scale": 1.0, "seed": 42
     }
   }' \
   | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' \
@@ -79,7 +79,7 @@ jq -n \
     modalities: ["image"],
     extra_body: {
       height: 1024, width: 1024,
-      steps: 12, cfg: 1.0, seed: 42
+      num_inference_steps: 12, guidance_scale: 1.0, seed: 42
     }
   }' |
 curl -sS http://127.0.0.1:8091/v1/chat/completions \
@@ -97,7 +97,7 @@ Note that the prompt should better depict each layer to be decomposed, we will r
 ```bash
 MODEL=inclusionAI/Ming-Image-0.1-Design-Layer
 INPUT_IMAGE=/path/to/input.png
-PROMPT="Decompose this design into editable visual layers"
+PROMPT="Decompose this image into 6 layers with the following specifications: Number of layers: 6 \nLayer 1: Central text ... Layer 2: ... Layer 6: ..."
 
 jq -n \
   --arg model "$MODEL" \
@@ -113,7 +113,7 @@ jq -n \
     extra_body: {
       num_layers: 6,
       height: 1024, width: 1024,
-      steps: 12, cfg: 2.0, seed: 42
+      num_inference_steps: 12, guidance_scale: 2.0, seed: 42
     }
   }' |
 curl -sS http://127.0.0.1:8091/v1/chat/completions \

recipes/krea/Krea-2.md

@@ -141,3 +141,58 @@ confirm a valid 1024×1024 PNG is written.
   no-CFG, True for Raw CFG).
 - Known limitations: TP / SP / CFG-Parallel are not yet wired for Krea 2. HSDP,
   layerwise CPU offload, and VAE-patch-parallel (decode) are supported.
+
+## XPU
+
+### 1x Intel Arc Pro B70 (32 GB)
+
+Offline text-to-image for the distilled Turbo checkpoint at 1024x1024, with
+layerwise offload streaming the transformer blocks.
+
+#### Environment
+
+- OS: Linux
+- Python: 3.10+
+- torch: 2.13.0+xpu
+- vLLM: 0.29.0 (`98dff2a8`)
+- vLLM-Omni: `main` at `4c7a98c2`
+
+#### Command
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model krea/Krea-2-Turbo \
+  --prompt "a photo of a sunset over mountains" \
+  --height 1024 --width 1024 \
+  --num-inference-steps 8 \
+  --enable-layerwise-offload \
+  --vae-use-tiling \
+  --enforce-eager \
+  --output krea2_turbo_output.png
+```
+
+Given spare cards, `--enable-cpu-offload` replaces layerwise streaming. At 8
+steps the swap dominates the run, and cards past the second shard nothing
+further.
+
+```bash
+  --tensor-parallel-size 2 \
+  --enable-cpu-offload
+```
+
+| Offload | Cards | Peak VRAM | Time |
+| --- | ---: | ---: | ---: |
+| `--enable-layerwise-offload` | 1 | 12.3 GiB | 27 s |
+| `--enable-cpu-offload` | 2 | 25.6 GiB | 59 s |
+| `--enable-cpu-offload` | 4 | 25.6 GiB | 59 s |
+
+#### Verification
+
+Confirm `krea2_turbo_output.png` is written as a 1024x1024 PNG matching the
+prompt.
+
+#### Notes
+
+- Memory usage: 7.8 GiB loaded, 12.3 GiB peak, about 26 s per image.
+- Known limitations: only offline Turbo generation was qualified. Krea-2-Raw,
+  LoRA, Cache-DiT, and online serving are out of scope for this profile.

tests/clients/test_model_session_configs.py

@@ -11,6 +11,9 @@ from vllm_omni.clients.duplex import AudioFormat
 from vllm_omni.clients.minicpmo_4_5 import (
     create_duplex_session_config as create_minicpmo45_session_config,
 )
+from vllm_omni.clients.nemotron_voicechat import (
+    create_duplex_session_config as create_nemotron_voicechat_session_config,
+)
 from vllm_omni.clients.personaplex import (
     create_duplex_session_config as create_personaplex_session_config,
 )
@@ -38,12 +41,22 @@ def test_personaplex_session_config_matches_deployment():
     assert payload["instructions"] == "You are calm."
 
 
+def test_nemotron_voicechat_session_config_matches_deployment():
+    tools = [{"type": "function", "name": "lookup", "parameters": {"type": "object"}}]
+    config = create_nemotron_voicechat_session_config(instructions="hi", tools=tools)
+    assert config.input_audio == AudioFormat("pcm_f32le", 16_000)
+    assert config.output_audio == AudioFormat("pcm16", 22_050)
+    payload = config.to_session_payload(model="nemotron")
+    assert payload["model"] == "nemotron" and payload["instructions"] == "hi"
+    assert payload["extra_body"] == {"auto_response": True, "realtime_tools": tools}
+
+
 @pytest.mark.parametrize(
     "duplex_package",
     [
-        # PersonaPlex and Nemotron VoiceChat rejoin this list with the follow-up
-        # PRs that port them to the duplex plugin framework.
+        # PersonaPlex rejoins with its port to the duplex plugin framework.
         "vllm_omni.model_executor.models.minicpmo_4_5.duplex",
+        "vllm_omni.model_executor.models.nemotron_voicechat.duplex",
     ],
 )
 def test_model_duplex_packages_stay_clear_of_client_library(duplex_package: str):

tests/config/test_config_factory.py

@@ -2857,6 +2857,28 @@ class TestPlatformOverrides:
             config = deploy.stages[0].compilation_config or {}
             assert "+rotary_embedding" not in config.get("custom_ops", [])
 
+    def test_qwen3_omni_talker_sampling_is_seeded(self):
+        """The only stochastic Qwen3-Omni stage must stay reproducible.
+
+        Dropping this seed (as #4986 did) leaves the talker sampling codec
+        tokens unseeded at temperature 0.9, which reopened the audio-vs-text
+        nightly failures in #6090. Pin it so a cleanup cannot remove it again
+        without failing here.
+        """
+        deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
+        pipeline = resolve_pipeline_config(
+            "qwen3_omni_moe",
+            Q3_OMNI_ALL_STAGES_HF_CONFIG,
+        )
+        assert isinstance(pipeline, PipelineConfig)
+
+        for platform in ("cpu", "cuda", "musa", "npu", "rocm", "xpu"):
+            deploy = _apply_platform_overrides(load_deploy_config(deploy_path), platform=platform)
+            stages = merge_pipeline_deploy(pipeline, deploy)
+            talker_sampling = stages[1].yaml_extras["default_sampling_params"]
+            assert talker_sampling["temperature"] > 0.0, "talker is expected to sample, not decode greedily"
+            assert talker_sampling.get("seed") == 42, f"talker sampling lost its seed on {platform}"
+
     def test_minicpmo_4_5_cuda_caps_talker_kv_cache(self):
         pipeline = resolve_pipeline_config("minicpmo_4_5")
         assert isinstance(pipeline, PipelineConfig)

tests/config/test_omni_config.py

@@ -127,7 +127,7 @@ def test_mammothmoda2_diffusion_stage_projects_native_backend_config() -> None:
     assert stage.diffusion_config.model_class_name == "MammothModa2DiTPipeline"
     assert stage.diffusion_config.model == "/models/MammothModa2-Preview"
     assert stage.diffusion_config.step_execution is False
-    assert stage.scheduler_config.max_num_seqs == 1
+    assert stage.scheduler_config.max_num_seqs == 8
     assert stage.connector_config.omni_kv_config == {"need_recv_cache": False}
 
 

tests/dfx/reliability/invalid_param_test/test_invalid_omni_chat.py

@@ -15,8 +15,6 @@ from tests.helpers.stage_config import get_deploy_config_path
 
 pytestmark = [pytest.mark.slow, pytest.mark.omni]
 
-_SKIP_ISSUE_3649 = pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/3649")
-
 
 def _minimal_chat_json(omni_server: OmniServer) -> dict[str, object]:
     """Minimal valid chat body; individual tests override one offending field."""
@@ -121,9 +119,8 @@ def _chat_completions_request_without_expectations(omni_server: OmniServer, case
         pytest.param(
             "modalities_list_bad",
             400,
-            ("modalities", "value_error", ""),
+            ("modalities", "list of strings"),
             id="modalities_list_bad_element",
-            marks=_SKIP_ISSUE_3649,
         ),
         pytest.param(
             "response_format_json_schema_incomplete",
@@ -131,7 +128,7 @@ def _chat_completions_request_without_expectations(omni_server: OmniServer, case
             ("response_format", "BadRequestError", "json_schema"),
             id="invalid_response_format_json_schema",
         ),
-        pytest.param("logprobs_wrong_type", 400, "logprobs", id="logprobs_wrong_type", marks=_SKIP_ISSUE_3649),
+        pytest.param("logprobs_wrong_type", 400, ("logprobs", "boolean"), id="logprobs_wrong_type"),
         pytest.param(
             "logprobs_top_without_enabled",
             400,

tests/diffusion/attention/test_ulysses_uaa.py

@@ -5,13 +5,16 @@ from __future__ import annotations
 
 import os
 import tempfile
+from contextlib import AbstractContextManager
 from types import SimpleNamespace
 
 import numpy as np
 import pytest
 import torch
 
+from tests.helpers.mark import hardware_test
 from tests.helpers.runtime import get_distributed_init_method
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.attention.parallel.ulysses import UlyssesParallelAttention
 from vllm_omni.diffusion.config import set_current_diffusion_config
@@ -578,6 +581,7 @@ def _run_attention_case(
     num_kv_heads: int | None = None,
     has_joint: bool = False,
     joint_strategy: str = "front",
+    mask_layout: str | None = None,
 ) -> None:
     device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
     current_omni_platform.set_device(device)
@@ -620,6 +624,7 @@ def _run_attention_case(
             q_full = torch.from_numpy(payload["q"]).to(device=device)
             k_full = torch.from_numpy(payload["k"]).to(device=device)
             v_full = torch.from_numpy(payload["v"]).to(device=device)
+            mask_full = torch.from_numpy(payload["mask"]).to(device=device) if "mask" in payload else None
             if has_joint:
                 joint_q = torch.from_numpy(payload["joint_q"]).to(device=device)
                 joint_k = torch.from_numpy(payload["joint_k"]).to(device=device)
@@ -640,6 +645,7 @@ def _run_attention_case(
                     v = torch.cat([v_full, joint_v], dim=1)
             else:
                 q, k, v = q_full, k_full, v_full
+            mask = mask_full
         else:
             if split_sizes is None:
                 # NOTE: torch.chunk may return fewer than `world_size` chunks for some
@@ -659,14 +665,18 @@ def _run_attention_case(
                 q = torch.split(q_full, split_sizes, dim=1)[local_rank].contiguous()
                 k = torch.split(k_full, split_sizes, dim=1)[local_rank].contiguous()
                 v = torch.split(v_full, split_sizes, dim=1)[local_rank].contiguous()
+            if mask_full is None:
+                mask = None
+            elif mask_layout == "global":
+                mask = mask_full
+            else:
+                raise ValueError("mask_layout must be 'global' when a mask is provided.")
             # The Attention layer only enables SP communication when ForwardContext.sp_active is True.
             # In production this is managed by SequenceParallelSplitHook/GatherHook, but here we
             # shard manually for testing.
             get_forward_context()._sp_shard_depth = 1
 
             if has_joint:
-                from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
-
                 attn_metadata = AttentionMetadata(
                     joint_query=joint_q,
                     joint_key=joint_k,
@@ -682,6 +692,8 @@ def _run_attention_case(
             raise ValueError(f"Invalid sdp_kernel_mode: {sdp_kernel_mode!r}")
 
         with sdp_ctx:
+            if mask is not None:
+                attn_metadata = AttentionMetadata(attn_mask=mask)
             out_local = attn(q, k, v, attn_metadata=attn_metadata).contiguous()
 
         if world_size == 1:
@@ -908,7 +920,7 @@ def _run_contract_skip_case(
         ctx = get_forward_context()
         ctx._sp_shard_depth = 1
 
-        def make_sdp_ctx() -> object:
+        def make_sdp_ctx() -> AbstractContextManager[None]:
             # Fresh instance per phase: the @contextmanager instance releases
             # its state on first exit and cannot be re-entered.
             return torch.backends.cuda.sdp_kernel(enable_flash=False, enable_mem_efficient=False, enable_math=True)
@@ -1011,6 +1023,105 @@ def test_ulysses_uaa_contract_skips_length_gather_matches_slow_path(num_heads: i
             pass
 
 
+@pytest.mark.core_model
+@pytest.mark.parametrize("mask_layout", ["global"])
+@hardware_test(res={"cuda": "L4"}, num_cards=2)
+def test_ulysses_uaa_2d_mask_layout_matches_baseline(
+    mask_layout: str,
+) -> None:
+    world_size = 2
+    batch_size = 2
+    seq_len = 5
+    split_sizes = [2, 3]
+    num_heads = 3
+    head_size = 8
+
+    base_init_method = get_distributed_init_method()
+    sp_init_method = get_distributed_init_method()
+
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".npz") as f_in:
+        input_file = f_in.name
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".npy") as f_base:
+        baseline_file = f_base.name
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".npy") as f_sp:
+        sp_file = f_sp.name
+
+    try:
+        torch.manual_seed(0)
+        q = torch.randn(batch_size, seq_len, num_heads, head_size, dtype=torch.float32)
+        k = torch.randn(batch_size, seq_len, num_heads, head_size, dtype=torch.float32)
+        v = torch.randn(batch_size, seq_len, num_heads, head_size, dtype=torch.float32)
+        mask = torch.tensor(
+            [
+                [True, True, True, True, False],
+                [True, True, True, False, False],
+            ],
+            dtype=torch.bool,
+        )
+        np.savez(
+            input_file,
+            q=q.numpy(),
+            k=k.numpy(),
+            v=v.numpy(),
+            mask=mask.numpy(),
+        )
+
+        # Baseline (no SP), mask is used verbatim.
+        torch.multiprocessing.spawn(
+            _run_attention_case,
+            args=(
+                1,
+                base_init_method,
+                input_file,
+                baseline_file,
+                num_heads,
+                head_size,
+                1,
+                "strict",
+                1,
+                None,
+                "math",
+                None,
+                False,
+                "front",
+                "global",
+            ),
+            nprocs=1,
+        )
+        # SP (Ulysses-P with UAA), mask is shared globally across ranks.
+        torch.multiprocessing.spawn(
+            _run_attention_case,
+            args=(
+                world_size,
+                sp_init_method,
+                input_file,
+                sp_file,
+                num_heads,
+                head_size,
+                world_size,
+                "advanced_uaa",
+                1,
+                split_sizes,
+                "math",
+                None,
+                False,
+                "front",
+                mask_layout,
+            ),
+            nprocs=world_size,
+        )
+
+        baseline = torch.from_numpy(np.load(baseline_file, allow_pickle=False))
+        sp = torch.from_numpy(np.load(sp_file, allow_pickle=False))
+        torch.testing.assert_close(sp, baseline, atol=1e-5, rtol=1e-5)
+    finally:
+        for path in (input_file, baseline_file, sp_file):
+            try:
+                os.remove(path)
+            except OSError:
+                pass
+
+
 @pytest.mark.parametrize(
     "num_heads,num_kv_heads,joint_len",
     [

tests/diffusion/conftest.py

@@ -1,15 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import pytest
-import torch
-
-
-@pytest.fixture(autouse=True)
-def restore_torch_default_dtype():
-    """Keep process-global dtype changes from leaking between diffusion tests."""
-    default_dtype = torch.get_default_dtype()
-    try:
-        yield
-    finally:
-        torch.set_default_dtype(default_dtype)

tests/diffusion/distributed/test_boogu_image_sp.py

@@ -0,0 +1,317 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import os
+import socket
+import tempfile
+
+import numpy as np
+import pytest
+import torch
+from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.diffusion.config import set_current_diffusion_config
+from vllm_omni.diffusion.data import (
+    DiffusionParallelConfig,
+    OmniDiffusionConfig,
+    TransformerConfig,
+)
+from vllm_omni.diffusion.distributed.parallel_state import (
+    destroy_distributed_env,
+    init_distributed_environment,
+    initialize_model_parallel,
+)
+from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig
+from vllm_omni.diffusion.forward_context import (
+    get_forward_context,
+    set_forward_context,
+)
+from vllm_omni.diffusion.hooks.sequence_parallel import apply_sequence_parallel
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.diffusion,
+    pytest.mark.parallel,
+]
+
+
+def _find_free_port() -> int:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+        sock.bind(("127.0.0.1", 0))
+        return int(sock.getsockname()[1])
+
+
+def _set_dist_env(rank: int, world_size: int, port: int) -> None:
+    os.environ.update(
+        {
+            "RANK": str(rank),
+            "LOCAL_RANK": str(rank),
+            "WORLD_SIZE": str(world_size),
+            "MASTER_ADDR": "127.0.0.1",
+            "MASTER_PORT": str(port),
+            "DIFFUSION_ATTENTION_BACKEND": "TORCH_SDPA",
+        }
+    )
+
+
+def _transformer_config() -> TransformerConfig:
+    return TransformerConfig.from_dict(
+        {
+            "patch_size": 2,
+            "in_channels": 4,
+            "hidden_size": 64,
+            "num_layers": 2,
+            "num_double_stream_layers": 1,
+            "num_refiner_layers": 1,
+            "num_attention_heads": 4,
+            "num_kv_heads": 1,
+            "multiple_of": 32,
+            "norm_eps": 1e-5,
+            "axes_dim_rope": [8, 4, 4],
+            "axes_lens": [64, 32, 32],
+            "instruction_feature_configs": {
+                "instruction_feat_dim": 32,
+                "reduce_type": "mean",
+                "num_instruction_feature_layers": 1,
+            },
+            "prompt_tuning_configs": {"use_prompt_tuning": False},
+            "timestep_scale": 1.0,
+        }
+    )
+
+
+def _od_config(world_size: int) -> OmniDiffusionConfig:
+    return OmniDiffusionConfig(
+        model="boogu-tiny",
+        dtype=torch.float32,
+        tf_model_config=_transformer_config(),
+        parallel_config=DiffusionParallelConfig(
+            sequence_parallel_size=world_size,
+            ulysses_degree=world_size,
+            ulysses_mode="advanced_uaa",
+        ),
+    )
+
+
+def _checkpoint_name(name: str) -> str:
+    if ".to_out." in name:
+        name = name.replace(".to_out.", ".to_out.0.")
+    for projection in (
+        "img_to_q",
+        "img_to_k",
+        "img_to_v",
+        "instruct_to_q",
+        "instruct_to_k",
+        "instruct_to_v",
+        "instruct_out",
+        "img_out",
+    ):
+        token = f".img_instruct_attn.{projection}."
+        if token in name:
+            return name.replace(
+                token,
+                f".img_instruct_attn.processor.{projection}.",
+            )
+    return name
+
+
+def _random_checkpoint(model: torch.nn.Module) -> list[tuple[str, torch.Tensor]]:
+    generator = torch.Generator(device=next(model.parameters()).device).manual_seed(2026)
+    return [
+        (
+            _checkpoint_name(name),
+            torch.empty_like(param).uniform_(
+                -0.02,
+                0.02,
+                generator=generator,
+            ),
+        )
+        for name, param in model.named_parameters()
+    ]
+
+
+def _worker(
+    rank: int,
+    world_size: int,
+    port: int,
+    output_file: str,
+) -> None:
+    device = torch.device(f"{current_omni_platform.device_type}:{rank}")
+    current_omni_platform.set_device(device)
+    _set_dist_env(rank, world_size, port)
+    init_distributed_environment(world_size=world_size, rank=rank)
+    initialize_model_parallel(
+        data_parallel_size=1,
+        cfg_parallel_size=1,
+        sequence_parallel_size=world_size,
+        ulysses_degree=world_size,
+        ring_degree=1,
+        tensor_parallel_size=1,
+        pipeline_parallel_size=1,
+    )
+
+    try:
+        from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
+            BooguImageDoubleStreamRotaryPosEmbed,
+            BooguImageTransformer2DModel,
+        )
+
+        od_config = _od_config(world_size)
+        with (
+            set_current_vllm_config(VllmConfig(device_config=DeviceConfig(device=str(device)))),
+            set_forward_context(omni_diffusion_config=od_config),
+            set_current_diffusion_config(od_config),
+        ):
+            source = BooguImageTransformer2DModel(od_config).to(device)
+            checkpoint = _random_checkpoint(source)
+            del source
+
+            model = BooguImageTransformer2DModel(od_config).to(device)
+            assert model.load_weights(checkpoint) == set(dict(model.named_parameters()))
+            model.eval()
+
+            if world_size > 1:
+                apply_sequence_parallel(
+                    model,
+                    SequenceParallelConfig(ulysses_degree=world_size),
+                    model._sp_plan,
+                )
+                get_forward_context().sp_plan_hooks_applied = True
+
+            generator = torch.Generator(device=device).manual_seed(17)
+            latents = torch.randn(
+                1,
+                4,
+                8,
+                8,
+                device=device,
+                generator=generator,
+            )
+            timestep = torch.tensor([0.5], device=device)
+            instruction = torch.randn(
+                1,
+                7,
+                32,
+                device=device,
+                generator=generator,
+            )
+            instruction_mask = torch.ones(
+                1,
+                7,
+                dtype=torch.bool,
+                device=device,
+            )
+            ref_image = [
+                [
+                    torch.randn(
+                        4,
+                        6,
+                        10,
+                        device=device,
+                        generator=generator,
+                    )
+                ]
+            ]
+            # Two differently sized references: their packed segments straddle
+            # the SP rank boundary, which is what the per-rank reference-refiner
+            # mask has to reproduce.
+            ref_images_multi = [
+                [
+                    torch.randn(4, 6, 10, device=device, generator=generator),
+                    torch.randn(4, 4, 8, device=device, generator=generator),
+                ]
+            ]
+            freqs_real = BooguImageDoubleStreamRotaryPosEmbed.get_freqs_real(
+                model.axes_dim_rope,
+                model.axes_lens,
+                theta=10000,
+            )
+
+            with torch.inference_mode():
+                t2i = model(
+                    latents,
+                    timestep,
+                    instruction,
+                    freqs_real,
+                    instruction_mask,
+                )
+                edit = model(
+                    latents,
+                    timestep,
+                    instruction,
+                    freqs_real,
+                    instruction_mask,
+                    ref_image_hidden_states=ref_image,
+                )
+                edit_multi = model(
+                    latents,
+                    timestep,
+                    instruction,
+                    freqs_real,
+                    instruction_mask,
+                    ref_image_hidden_states=ref_images_multi,
+                )
+
+            if rank == 0:
+                np.savez(
+                    output_file,
+                    t2i=t2i.cpu().numpy(),
+                    edit=edit.cpu().numpy(),
+                    edit_multi=edit_multi.cpu().numpy(),
+                )
+    finally:
+        destroy_distributed_env()
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=2)
+def test_boogu_sp_matches_single_gpu() -> None:
+    if not current_omni_platform.is_cuda() or current_omni_platform.get_device_count() < 2:
+        pytest.skip("BOOGU SP test requires two CUDA GPUs.")
+
+    output_files = []
+    try:
+        for _ in range(2):
+            with tempfile.NamedTemporaryFile(
+                delete=False,
+                suffix=".npz",
+            ) as handle:
+                output_files.append(handle.name)
+
+        torch.multiprocessing.spawn(
+            _worker,
+            args=(1, _find_free_port(), output_files[0]),
+            nprocs=1,
+        )
+        torch.multiprocessing.spawn(
+            _worker,
+            args=(2, _find_free_port(), output_files[1]),
+            nprocs=2,
+        )
+
+        with (
+            np.load(output_files[0], allow_pickle=False) as baseline,
+            np.load(output_files[1], allow_pickle=False) as sp_output,
+        ):
+            for key in ("t2i", "edit", "edit_multi"):
+                # fp32 SDPA: the observed SP1-vs-SP2 gap is ~3e-8, so 2e-4
+                # would also pass with the SP masking logic entirely removed
+                # (that mutation measures ~5e-7). Mask correctness is pinned
+                # structurally in tests/diffusion/models/boogu_image/
+                # test_sp_layout.py; this bound just stays close to the real
+                # reduction-order noise instead of hiding logic errors.
+                np.testing.assert_allclose(
+                    sp_output[key],
+                    baseline[key],
+                    rtol=1e-5,
+                    atol=1e-5,
+                )
+    finally:
+        for path in output_files:
+            try:
+                os.remove(path)
+            except OSError:
+                pass

tests/diffusion/distributed/test_sp_plan_hooks.py

@@ -115,6 +115,47 @@ class TestSequenceParallelPlanValidation:
         with pytest.raises(ValueError, match="split_output=True"):
             validate_sp_plan(plan)
 
+    def test_keyed_metadata_requires_auto_pad(self):
+        plan = {
+            "prepare": {
+                0: SequenceParallelInput(
+                    split_dim=1,
+                    split_output=True,
+                    shard_group="image",
+                ),
+            },
+            "output": SequenceParallelOutput(
+                gather_dim=1,
+                shard_group="image",
+            ),
+        }
+
+        with pytest.raises(ValueError, match="requires auto_pad=True"):
+            validate_sp_plan(plan)
+
+    def test_shard_group_is_scoped_to_one_split_boundary(self):
+        plan = {
+            "prepare_image": {
+                0: SequenceParallelInput(
+                    split_dim=1,
+                    split_output=True,
+                    auto_pad=True,
+                    shard_group="image",
+                ),
+            },
+            "prepare_rope": {
+                0: SequenceParallelInput(
+                    split_dim=1,
+                    split_output=True,
+                    auto_pad=True,
+                    shard_group="image",
+                ),
+            },
+        }
+
+        with pytest.raises(ValueError, match="multiple split boundaries"):
+            validate_sp_plan(plan)
+
 
 @pytest.mark.cpu
 class TestGetSpPlanFromModel:
@@ -1006,6 +1047,8 @@ class TestDimensionValidation:
                 return "NO_MASK_BACKEND"
 
         monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: 2)
+        monkeypatch.setattr(parallel_state, "get_sequence_parallel_rank", lambda: 0)
+        monkeypatch.setattr(parallel_state, "get_ring_parallel_world_size", lambda: 1)
         monkeypatch.setattr(selector, "get_attn_backend_for_capability", lambda **kwargs: NoMaskBackend)
 
         hook = SequenceParallelSplitHook(
@@ -1015,7 +1058,7 @@ class TestDimensionValidation:
         tensor = torch.randn(1, 3, 8)
 
         with pytest.raises(ValueError, match="does not support attention_mask"):
-            hook._shard_with_auto_pad(tensor, dim=1)
+            hook._shard_with_auto_pad(tensor, dim=1, shard_group=None)
 
     def test_auto_pad_allows_explicit_cudnn(self, monkeypatch):
         from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, OmniDiffusionConfig
@@ -1040,7 +1083,7 @@ class TestDimensionValidation:
         )
 
         with set_forward_context(omni_diffusion_config=od_config):
-            result = hook._shard_with_auto_pad(tensor, dim=1)
+            result = hook._shard_with_auto_pad(tensor, dim=1, shard_group=None)
 
         # seq_len=3 padded to 4, then sharded by world_size=2 → rank 0 gets 2.
         assert result.shape == (1, 2, 8)
@@ -1070,7 +1113,7 @@ class TestDimensionValidation:
 
         with set_forward_context(omni_diffusion_config=od_config):
             with pytest.raises(ValueError, match="does not support attention_mask"):
-                hook._shard_with_auto_pad(tensor, dim=1)
+                hook._shard_with_auto_pad(tensor, dim=1, shard_group=None)
 
 
 @pytest.mark.cpu
@@ -1227,6 +1270,158 @@ class TestEqualPadContract:
             assert not ctx.sp_rank_local_seq_lens_equal
 
 
+@pytest.mark.cpu
+class TestKeyedShardMetadata:
+    """Per-shard_group padding metadata, for models with several sequences.
+
+    A model that splits image, reference and context tokens at one boundary
+    needs each one's pre-padding global length to rebuild masks; the singleton
+    sp_original_seq_len can only describe one of them.
+    """
+
+    def test_auto_pad_records_the_global_length_under_its_key(self, monkeypatch):
+        from vllm_omni.diffusion.attention import selector
+        from vllm_omni.diffusion.distributed import parallel_state
+        from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig
+        from vllm_omni.diffusion.forward_context import (
+            get_forward_context,
+            get_sp_shard_original_seq_len,
+            set_forward_context,
+        )
+        from vllm_omni.diffusion.hooks.sequence_parallel import (
+            SequenceParallelSplitHook,
+        )
+
+        class MaskBackend:
+            supports_attention_mask = True
+
+            @staticmethod
+            def get_name():
+                return "test"
+
+        monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: 2)
+        monkeypatch.setattr(parallel_state, "get_sequence_parallel_rank", lambda: 1)
+        monkeypatch.setattr(parallel_state, "get_ring_parallel_world_size", lambda: 1)
+        monkeypatch.setattr(selector, "get_attn_backend_for_role", lambda **_: (MaskBackend(), None))
+
+        hook = SequenceParallelSplitHook(
+            {
+                0: SequenceParallelInput(
+                    split_dim=1,
+                    expected_dims=3,
+                    split_output=True,
+                    auto_pad=True,
+                    shard_group="image",
+                )
+            },
+            SequenceParallelConfig(ulysses_degree=2),
+        )
+        hook.initialize_hook(nn.Identity())
+
+        with set_forward_context():
+            hook.pre_forward(nn.Identity())
+            shard = hook.post_forward(nn.Identity(), torch.arange(5).reshape(1, 5, 1))
+            ctx = get_forward_context()
+
+            assert shard.shape[1] == 3
+            # The *pre*-padding length, which is what a mask has to clip to.
+            assert ctx.sp_shard_metadata["image"] == 5
+            assert get_sp_shard_original_seq_len("image") == 5
+            assert get_sp_shard_original_seq_len("context") is None
+
+    def test_gather_trims_to_its_key_and_leaves_other_groups(self, monkeypatch):
+        from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig
+        from vllm_omni.diffusion.forward_context import (
+            get_forward_context,
+            set_forward_context,
+        )
+        from vllm_omni.diffusion.hooks import sequence_parallel
+        from vllm_omni.diffusion.hooks.sequence_parallel import (
+            SequenceParallelGatherHook,
+        )
+
+        monkeypatch.setattr(
+            sequence_parallel,
+            "sp_gather",
+            lambda tensor, dim, validate: torch.cat([tensor, tensor], dim=dim),
+        )
+        hook = SequenceParallelGatherHook(
+            SequenceParallelOutput(gather_dim=1, expected_dims=3, shard_group="image"),
+            SequenceParallelConfig(ulysses_degree=2),
+        )
+
+        with set_forward_context():
+            ctx = get_forward_context()
+            ctx.sp_shard_metadata["image"] = 5
+            ctx.sp_shard_metadata["context"] = 7
+            ctx._sp_shard_depth = 1
+            output = hook.post_forward(nn.Identity(), torch.zeros(1, 3, 4))
+
+            # 2x3 gathered, trimmed back to "image"'s recorded 5.
+            assert output.shape == (1, 5, 4)
+            # A concurrent group is untouched.
+            assert ctx.sp_shard_metadata == {"context": 7}
+
+    def test_gather_without_registered_metadata_is_an_error(self, monkeypatch):
+        from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig
+        from vllm_omni.diffusion.forward_context import set_forward_context
+        from vllm_omni.diffusion.hooks import sequence_parallel
+        from vllm_omni.diffusion.hooks.sequence_parallel import (
+            SequenceParallelGatherHook,
+        )
+
+        monkeypatch.setattr(
+            sequence_parallel,
+            "sp_gather",
+            lambda tensor, dim, validate: torch.cat([tensor, tensor], dim=dim),
+        )
+        hook = SequenceParallelGatherHook(
+            SequenceParallelOutput(gather_dim=1, expected_dims=3, shard_group="image"),
+            SequenceParallelConfig(ulysses_degree=2),
+        )
+
+        with set_forward_context(), pytest.raises(RuntimeError, match="shard_group='image'"):
+            hook.post_forward(nn.Identity(), torch.zeros(1, 3, 4))
+
+    def test_conflicting_lengths_in_one_group_are_rejected(self, monkeypatch):
+        from vllm_omni.diffusion.attention import selector
+        from vllm_omni.diffusion.distributed import parallel_state
+        from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig
+        from vllm_omni.diffusion.forward_context import set_forward_context
+        from vllm_omni.diffusion.hooks.sequence_parallel import (
+            SequenceParallelSplitHook,
+        )
+
+        class MaskBackend:
+            supports_attention_mask = True
+
+            @staticmethod
+            def get_name():
+                return "test"
+
+        monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: 2)
+        monkeypatch.setattr(parallel_state, "get_sequence_parallel_rank", lambda: 1)
+        monkeypatch.setattr(parallel_state, "get_ring_parallel_world_size", lambda: 1)
+        monkeypatch.setattr(selector, "get_attn_backend_for_role", lambda **_: (MaskBackend(), None))
+
+        spm = SequenceParallelInput(
+            split_dim=1,
+            expected_dims=3,
+            split_output=True,
+            auto_pad=True,
+            shard_group="image",
+        )
+        hook = SequenceParallelSplitHook({0: spm, 1: spm}, SequenceParallelConfig(ulysses_degree=2))
+        hook.initialize_hook(nn.Identity())
+
+        with set_forward_context(), pytest.raises(ValueError, match="same global sequence length"):
+            hook.pre_forward(nn.Identity())
+            hook.post_forward(
+                nn.Identity(),
+                (torch.arange(5).reshape(1, 5, 1), torch.arange(6).reshape(1, 6, 1)),
+            )
+
+
 @pytest.mark.cpu
 class TestSequenceParallelConfig:
     """Test SequenceParallelConfig dataclass."""

tests/diffusion/layers/test_mhc.py

@@ -0,0 +1,374 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.layers import mhc
+from vllm_omni.diffusion.models.magi2 import layers as model_layers
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.core_model]
+
+
+@pytest.fixture(autouse=True)
+def clear_mhc_failure_cache():
+    """Keep the process-lifetime production cache isolated between tests."""
+
+    mhc._FAILED_MHC_KERNELS.clear()
+    mhc._WARNED_MHC_KERNELS.clear()
+    yield
+    mhc._FAILED_MHC_KERNELS.clear()
+    mhc._WARNED_MHC_KERNELS.clear()
+
+
+def coefficients(tokens=7, streams=4):
+    torch.manual_seed(71)
+    # Actual model split views: post and residual share a wider FP32 row.
+    packed = torch.randn(tokens, streams * (streams + 2))
+    post = packed[:, streams : 2 * streams]
+    residual = packed[:, 2 * streams :].view(tokens, streams, streams)
+    return post, residual, torch.tensor([0.8]), torch.randn(streams), torch.tensor([1.1]), torch.randn(streams, streams)
+
+
+def reference(args, *, scale=0.125, iterations=20, epsilon=1e-12, out_dtype=torch.float32):
+    post, residual, ap, bp, ar, br = args
+    post = 2.0 * torch.sigmoid(ap * scale * post + bp.unsqueeze(0))
+    residual = ar * scale * residual.float() + br.unsqueeze(0).float()
+    residual = torch.exp(residual - residual.amax(dim=(-2, -1), keepdim=True))
+    for _ in range(iterations):
+        residual = residual / (residual.sum(-2, keepdim=True) + epsilon)
+        residual = residual / (residual.sum(-1, keepdim=True) + epsilon)
+    return post.to(out_dtype), residual.to(out_dtype)
+
+
+def mix_reference(args):
+    streams, output, post, matrix = args
+    branch = torch.einsum("tn,tc->tnc", post, output)
+    mixed = torch.einsum("tij,tjc->tic", matrix, streams)
+    return mixed + branch
+
+
+def mix_inputs(tokens=7, hidden=70, dtype=torch.float32, streams=4):
+    torch.manual_seed(41)
+    return (
+        torch.randn(tokens, streams, hidden).to(dtype),
+        torch.randn(tokens, hidden).to(dtype),
+        torch.rand(tokens, streams).to(dtype),
+        torch.randn(tokens, streams, streams).to(dtype),
+    )
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("streams", [3, 4])
+@pytest.mark.parametrize("iterations", [-1, 0, 1, 20, 65])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+def test_native_coefficient_contract(streams, iterations, dtype):
+    args = coefficients(streams=streams)
+    kwargs = dict(scale=0.125, iterations=iterations, out_dtype=dtype)
+    actual = mhc.MHCPostResidual().forward_cuda(*args, **kwargs)
+    for a, b in zip(actual, reference(args, **kwargs)):
+        torch.testing.assert_close(a, b, rtol=0, atol=0)
+
+
+@pytest.mark.cpu
+def test_native_gradients_and_reused_parameters():
+    op = mhc.MHCPostResidual()
+    args = tuple(x.detach().clone().requires_grad_() for x in coefficients())
+    expected_args = tuple(x.detach().clone().requires_grad_() for x in args)
+    result = op.forward_cuda(*args, scale=0.125, out_dtype=torch.float32)
+    expected = reference(expected_args)
+    sum(x.square().sum() for x in result).backward()
+    sum(x.square().sum() for x in expected).backward()
+    for a, b in zip(args, expected_args):
+        torch.testing.assert_close(a.grad, b.grad, rtol=0, atol=0)
+    with torch.no_grad():
+        args[3].add_(2)
+        result = op.forward_cuda(*args, scale=0.125, out_dtype=torch.float32)
+        for a, b in zip(result, reference(args)):
+            torch.testing.assert_close(a, b, rtol=0, atol=0)
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+def test_native_mix_dtype_and_noncontiguous(dtype):
+    args = mix_inputs(dtype=dtype)
+    args = (args[0][..., ::2], args[1][..., ::2], *args[2:])
+    torch.testing.assert_close(mhc.MHCMix().forward_cuda(*args), mix_reference(args), rtol=0, atol=0)
+
+
+@pytest.mark.cpu
+def test_meta_shapes_and_empty():
+    args = tuple(x.to("meta") for x in coefficients())
+    post, matrix = mhc.MHCPostResidual().forward_cuda(*args, scale=0.125, out_dtype=torch.bfloat16)
+    assert post.shape == (7, 4) and matrix.shape == (7, 4, 4)
+    assert post.dtype == matrix.dtype == torch.bfloat16
+    mix = mhc.MHCMix().forward_cuda(*[x.to("meta") for x in mix_inputs()])
+    assert mix.shape == (7, 4, 70) and mix.device.type == "meta"
+    args = coefficients(tokens=0)
+    post, matrix = mhc.MHCPostResidual().forward_cuda(*args, scale=0.125, out_dtype=torch.float32)
+    assert post.shape == (0, 4) and matrix.shape == (0, 4, 4)
+    assert mhc.MHCMix().forward_cuda(*mix_inputs(tokens=0)).shape == (0, 4, 70)
+
+
+@pytest.mark.cpu
+def test_sinkhorn_alias_and_rectangular_native_contract():
+    assert model_layers.sinkhorn_knopp is mhc.sinkhorn_knopp
+    logits = torch.randn(2, 3, 4)
+    expected = torch.exp(logits - logits.amax((-2, -1), keepdim=True))
+    torch.testing.assert_close(mhc.sinkhorn_knopp(logits, 0, 1e-12), expected, rtol=0, atol=0)
+
+
+@pytest.mark.cpu
+def test_model_wiring_uses_fused_eager_path(monkeypatch):
+    calls = []
+
+    class Prepare:
+        def __call__(self, *args, **kwargs):
+            calls.append("dispatch")
+            return reference(args, **kwargs)
+
+        def forward_native(self, *args, **kwargs):
+            calls.append("native")
+            return reference(args, **kwargs)
+
+    monkeypatch.setattr(model_layers, "_mhc_post_residual", Prepare())
+    handler = model_layers.MHCHandler(4, 16)
+    post, residual, ap, bp, ar, br = coefficients()
+    handler.compute_post_residual((ap, bp, post), (ar, br, residual), out_dtype=torch.float32)
+    assert calls == ["dispatch"]
+
+
+@pytest.mark.cpu
+def test_compile_keeps_native_graph():
+    torch._dynamo.reset()
+    targets: list[str] = []
+
+    def backend(gm, _inputs):
+        targets.extend(str(node.target) for node in gm.graph.nodes)
+        return gm.forward
+
+    op = torch.compile(mhc.MHCPostResidual(), fullgraph=True, backend=backend)
+    args = coefficients()
+    actual = op(*args, scale=0.125, iterations=2, out_dtype=torch.float32)
+    for a, b in zip(actual, reference(args, iterations=2)):
+        torch.testing.assert_close(a, b, rtol=0, atol=0)
+    assert any("sigmoid" in x for x in targets) and any("sum" in x for x in targets)
+    assert not any("triton" in x for x in targets)
+
+
+DEVICES = [pytest.param("cuda", marks=pytest.mark.cuda), pytest.param("musa", marks=pytest.mark.musa)]
+
+
+def device(name):
+    api = getattr(torch, name, None)
+    if api is None or not api.is_available():
+        pytest.skip(f"requires {name}")
+    return torch.device(name)
+
+
+def spy(monkeypatch, name):
+    original = getattr(mhc, name)
+    launches = []
+
+    class Spy:
+        def __getitem__(self, grid):
+            launches.append(grid)
+            return original[grid]
+
+    monkeypatch.setattr(mhc, name, Spy())
+    return launches
+
+
+@pytest.mark.parametrize("device_name", DEVICES)
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("tokens,iterations", [(0, 20), (1, 0), (17, 20)])
+def test_coefficients_gpu(monkeypatch, device_name, dtype, tokens, iterations):
+    dev = device(device_name)
+    # Preserve the model's strided split views on the GPU too.
+    args = coefficients(tokens)
+    packed = torch.zeros(tokens, 24, device=dev)
+    packed[:, 4:8] = args[0].to(dev)
+    packed[:, 8:] = args[1].reshape(tokens, 16).to(dev)
+    args = (packed[:, 4:8], packed[:, 8:].view(tokens, 4, 4), *(x.to(dev) for x in args[2:]))
+    launches = spy(monkeypatch, "_mhc_post_residual_kernel")
+    actual = mhc.MHCPostResidual()(*args, scale=0.125, iterations=iterations, out_dtype=dtype)
+    getattr(torch, device_name).synchronize()
+    assert len(launches) == bool(tokens)
+    tolerance = {torch.float32: 1e-5, torch.float16: 2e-3, torch.bfloat16: 1.6e-2}[dtype]
+    for a, b in zip(actual, reference(args, iterations=iterations, out_dtype=dtype)):
+        torch.testing.assert_close(a, b, rtol=tolerance, atol=1e-6)
+
+
+@pytest.mark.parametrize("device_name", DEVICES)
+def test_sinkhorn_gpu_shift_and_extremes(device_name):
+    dev = device(device_name)
+    args = tuple(x.to(dev) for x in coefficients())
+    op = mhc.MHCPostResidual()
+    post, matrix = op(*args, scale=1.0, out_dtype=torch.float32)
+    shifted = (args[0], args[1] + 16.0, *args[2:])
+    shifted_post, shifted_matrix = op(*shifted, scale=1.0, out_dtype=torch.float32)
+    torch.testing.assert_close(shifted_matrix, matrix, rtol=1e-5, atol=3e-6)
+    torch.testing.assert_close(shifted_post, post, rtol=0, atol=0)
+    torch.testing.assert_close(matrix.sum(-1), torch.ones_like(post), rtol=1e-5, atol=1e-5)
+    # A well-conditioned case converges in both dimensions at the requested
+    # iteration count. Extreme matrices below need not converge in 20 steps.
+    balanced_args = list(args)
+    balanced_args[1] = args[1] * 0.125
+    balanced_args[5] = torch.zeros_like(args[5])
+    _, balanced = op(*balanced_args, scale=1.0, out_dtype=torch.float32)
+    for dim in (-2, -1):
+        torch.testing.assert_close(balanced.sum(dim), torch.ones_like(post), rtol=1e-5, atol=1e-5)
+    extreme = list(args)
+    extreme[1] = torch.eye(4, device=dev).expand(7, -1, -1).contiguous() * 1000
+    extreme[1][0, 0, 0] = float("nan")
+    extreme[1][1, 0, 0] = float("inf")
+    for iters in (0, 20):
+        actual = op(*extreme, scale=1.0, iterations=iters, out_dtype=torch.float32)
+        expected = reference(extreme, scale=1.0, iterations=iters)
+        for a, b in zip(actual, expected):
+            torch.testing.assert_close(a, b, rtol=1e-5, atol=1e-5, equal_nan=True)
+
+
+@pytest.mark.parametrize("device_name", DEVICES)
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("tokens,hidden", [(0, 70), (2, 70), (17, 3072)])
+def test_mix_gpu(monkeypatch, device_name, dtype, tokens, hidden):
+    dev = device(device_name)
+    args = tuple(x.to(dev) for x in mix_inputs(tokens, hidden, dtype))
+    launches = spy(monkeypatch, "_mhc_mix_kernel")
+    actual = mhc.MHCMix()(*args)
+    getattr(torch, device_name).synchronize()
+    assert len(launches) == bool(tokens)
+    tolerance = {torch.float32: 1e-5, torch.float16: 2e-3, torch.bfloat16: 1.6e-2}[dtype]
+    torch.testing.assert_close(actual, mix_reference(args), rtol=tolerance, atol=1e-5)
+
+
+@pytest.mark.parametrize("device_name", DEVICES)
+def test_mix_preserves_separate_bf16_rounding(device_name):
+    dev = device(device_name)
+    streams = torch.ones(1, 4, 1, device=dev, dtype=torch.bfloat16)
+    branch = torch.full((1, 1), 1.0078125, device=dev, dtype=torch.bfloat16)
+    post = torch.full((1, 4), 1.03125, device=dev, dtype=torch.bfloat16)
+    matrix = torch.eye(4, device=dev, dtype=torch.bfloat16).unsqueeze(0)
+    args = (streams, branch, post, matrix)
+    expected = mix_reference(args)
+    once = (
+        torch.einsum("tij,tjc->tic", matrix.float(), streams.float())
+        + torch.einsum("tn,tc->tnc", post.float(), branch.float())
+    ).bfloat16()
+    assert not torch.equal(expected, once), "fixture must distinguish the cast boundaries"
+    torch.testing.assert_close(mhc.MHCMix()(*args), expected, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("device_name", DEVICES)
+def test_gpu_unsupported_layout_falls_back(monkeypatch, device_name):
+    dev = device(device_name)
+
+    class Forbidden:
+        def __getitem__(self, _grid):
+            pytest.fail("unsupported layouts must not launch the kernel")
+
+    monkeypatch.setattr(mhc, "_mhc_post_residual_kernel", Forbidden())
+    monkeypatch.setattr(mhc, "_mhc_mix_kernel", Forbidden())
+    args = tuple(x.to(dev) for x in coefficients(streams=3))
+    actual = mhc.MHCPostResidual()(*args, scale=0.125, out_dtype=torch.float32)
+    for a, b in zip(actual, reference(args)):
+        torch.testing.assert_close(a, b, rtol=0, atol=0)
+    args = tuple(x.to(dev) for x in mix_inputs())
+    args = (args[0][..., ::2], args[1][..., ::2], *args[2:])
+    torch.testing.assert_close(mhc.MHCMix()(*args), mix_reference(args), rtol=0, atol=0)
+    # A singleton alpha with extra rank changes broadcasting, so it cannot be
+    # treated like the scalar alpha loaded by the eager kernel.
+    broadcast_args = list(x.to(dev) for x in coefficients())
+    broadcast_args[2] = broadcast_args[2].view(1, 1, 1)
+    actual = mhc.MHCPostResidual()(*broadcast_args, scale=0.125, out_dtype=torch.float32)
+    for a, b in zip(actual, reference(broadcast_args)):
+        torch.testing.assert_close(a, b, rtol=0, atol=0)
+
+
+@pytest.mark.cpu
+def test_no_triton_uses_native_fallback(monkeypatch):
+    monkeypatch.setattr(mhc, "HAS_TRITON", False)
+    args = coefficients(tokens=2)
+    actual = mhc.MHCPostResidual().forward_cuda(*args, scale=0.125, out_dtype=torch.float32)
+    for got, expected in zip(actual, reference(args)):
+        torch.testing.assert_close(got, expected, rtol=0, atol=0)
+
+    mix_args = mix_inputs(tokens=2)
+    torch.testing.assert_close(mhc.MHCMix().forward_cuda(*mix_args), mix_reference(mix_args), rtol=0, atol=0)
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize(
+    ("device_type", "native_required"),
+    [("cpu", True), ("meta", True), ("hip", True), ("xpu", True), ("cuda", False), ("musa", False)],
+)
+def test_unvalidated_backend_uses_native(device_type, native_required):
+    tensor = SimpleNamespace(device=SimpleNamespace(type=device_type), requires_grad=False)
+    assert mhc._native_required((tensor,)) is native_required
+
+
+@pytest.mark.cpu
+def test_kernel_failure_is_cached_and_falls_back(monkeypatch):
+    monkeypatch.setattr(mhc, "HAS_TRITON", True)
+    monkeypatch.setattr(mhc, "_native_required", lambda tensors: False)
+    warnings = []
+    monkeypatch.setattr(mhc.logger, "warning", lambda *args, **kwargs: warnings.append((args, kwargs)))
+
+    class Broken:
+        def __getitem__(self, _grid):
+            def launch(*args, **kwargs):
+                raise RuntimeError("synthetic kernel failure")
+
+            return launch
+
+    monkeypatch.setattr(mhc, "_mhc_post_residual_kernel", Broken())
+    args = tuple(x.contiguous() for x in coefficients(tokens=2))
+    op = mhc.MHCPostResidual()
+    expected = reference(args)
+    for _ in range(2):
+        actual = op.forward_cuda(*args, scale=0.125, out_dtype=torch.float32)
+        for got, want in zip(actual, expected):
+            torch.testing.assert_close(got, want, rtol=0, atol=0)
+    assert ("cpu", 0, "post_residual") in mhc._FAILED_MHC_KERNELS
+    assert len(warnings) == 1
+    assert "post_residual" in str(warnings[0])
+
+
+@pytest.mark.cpu
+def test_kernel_oom_is_propagated_without_caching(monkeypatch):
+    monkeypatch.setattr(mhc, "HAS_TRITON", True)
+    monkeypatch.setattr(mhc, "_native_required", lambda tensors: False)
+
+    class OutOfMemory:
+        def __getitem__(self, _grid):
+            def launch(*args, **kwargs):
+                raise RuntimeError("synthetic CUDA out of memory")
+
+            return launch
+
+    monkeypatch.setattr(mhc, "_mhc_post_residual_kernel", OutOfMemory())
+    args = tuple(x.contiguous() for x in coefficients(tokens=2))
+    with pytest.raises(RuntimeError, match="out of memory"):
+        mhc.MHCPostResidual().forward_cuda(*args, scale=0.125, out_dtype=torch.float32)
+    assert ("cpu", 0, "post_residual") not in mhc._FAILED_MHC_KERNELS
+
+
+@pytest.mark.parametrize("device_name", DEVICES)
+def test_gpu_compiled_parameters_do_not_get_cached(monkeypatch, device_name):
+    dev = device(device_name)
+
+    class Forbidden:
+        def __getitem__(self, _grid):
+            pytest.fail("compilation must retain the native expression")
+
+    monkeypatch.setattr(mhc, "_mhc_post_residual_kernel", Forbidden())
+    op = torch.compile(mhc.MHCPostResidual(), backend="eager", fullgraph=True)
+    args = list(x.to(dev) for x in coefficients())
+    for bias in (0.0, 2.0, -1.0):
+        args[3].fill_(bias)
+        actual = op(*args, scale=0.125, out_dtype=torch.float32)
+        for a, b in zip(actual, reference(args)):
+            torch.testing.assert_close(a, b, rtol=0, atol=0)

tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py

@@ -365,8 +365,12 @@ def test_load_weights_maps_mllm_and_preserves_transformer():
     ("parallel_config", "cache_backend", "message"),
     [
         (DiffusionParallelConfig(tensor_parallel_size=2), "none", "Tensor parallelism"),
-        (DiffusionParallelConfig(ulysses_degree=2), "none", "Sequence parallelism"),
-        (DiffusionParallelConfig(ring_degree=2), "none", "Sequence parallelism"),
+        (DiffusionParallelConfig(ring_degree=2), "none", "Ulysses only"),
+        (
+            DiffusionParallelConfig(ulysses_degree=2, cfg_parallel_size=2),
+            "none",
+            "CFG parallelism is not validated",
+        ),
         (
             DiffusionParallelConfig(use_hsdp=True, hsdp_shard_size=2),
             "none",
@@ -421,6 +425,96 @@ def test_constructor_accepts_cfg_parallel(mock_dependencies, cfg_parallel_size):
     assert hasattr(pipeline, "predict_noise_with_multi_branch_cfg")
 
 
+def test_constructor_accepts_ulysses_uaa(
+    mock_dependencies,
+    monkeypatch,
+):
+    from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import (
+        BooguImagePipeline,
+    )
+
+    monkeypatch.setattr(
+        f"{_MODULE}.current_omni_platform.is_cuda",
+        lambda: True,
+    )
+    od_config = OmniDiffusionConfig(
+        model="dummy-boogu",
+        tf_model_config=TransformerConfig(
+            params={
+                "num_attention_heads": 28,
+                "num_kv_heads": 7,
+            }
+        ),
+        dtype=torch.float32,
+        parallel_config=DiffusionParallelConfig(
+            ulysses_degree=2,
+            ulysses_mode="advanced_uaa",
+        ),
+    )
+
+    pipeline = BooguImagePipeline(od_config=od_config)
+    assert pipeline.transformer is mock_dependencies["transformer"]
+
+
+def test_constructor_requires_uaa_for_boogu_gqa(
+    mock_dependencies,
+    monkeypatch,
+):
+    from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import (
+        BooguImagePipeline,
+    )
+
+    monkeypatch.setattr(
+        f"{_MODULE}.current_omni_platform.is_cuda",
+        lambda: True,
+    )
+    od_config = OmniDiffusionConfig(
+        model="dummy-boogu",
+        tf_model_config=TransformerConfig(
+            params={
+                "num_attention_heads": 28,
+                "num_kv_heads": 7,
+            }
+        ),
+        dtype=torch.float32,
+        parallel_config=DiffusionParallelConfig(
+            ulysses_degree=2,
+            ulysses_mode="strict",
+        ),
+    )
+
+    with pytest.raises(ValueError, match="advanced_uaa"):
+        BooguImagePipeline(od_config=od_config)
+
+
+def test_constructor_rejects_non_cuda_sequence_parallelism(
+    mock_dependencies,
+    monkeypatch,
+):
+    from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import (
+        BooguImagePipeline,
+    )
+
+    monkeypatch.setattr(
+        f"{_MODULE}.current_omni_platform.is_cuda",
+        lambda: False,
+    )
+    od_config = OmniDiffusionConfig(
+        model="dummy-boogu",
+        tf_model_config=TransformerConfig(params={}),
+        dtype=torch.float32,
+        parallel_config=DiffusionParallelConfig(
+            ulysses_degree=2,
+            ulysses_mode="advanced_uaa",
+        ),
+    )
+
+    with pytest.raises(NotImplementedError, match="requires CUDA"):
+        BooguImagePipeline(od_config=od_config)
+
+    mock_dependencies["mllm_wrapper"].model.to.assert_not_called()
+
+
 # ---------------------------------------------------------------------------
 # Prompt-encoding tests (deterministic fakes, no constructor)
 # ---------------------------------------------------------------------------

tests/diffusion/models/boogu_image/test_sp_layout.py

@@ -0,0 +1,132 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""L1 unit tests for BOOGU's SP shard-layout arithmetic.
+
+These cover the property the model relies on to build attention masks without
+communicating: every rank can derive *every* rank's segment bounds from the
+global sequence length alone.
+"""
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.boogu_image.sp_layout import (
+    ShardLayout,
+    rank_concat_mask_or_none,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+class TestShardLayout:
+    def test_padding_and_local_extent(self):
+        layout = ShardLayout(original_seq_len=5, world_size=2, rank=0)
+        assert layout.padded_seq_len == 6
+        assert layout.local_seq_len == 3
+        assert layout.padding_size == 1
+        assert layout.bounds(0) == (0, 3)
+        assert layout.bounds(1) == (3, 6)
+
+    def test_valid_lengths_clip_to_each_rank(self):
+        layout = ShardLayout(original_seq_len=5, world_size=2, rank=0)
+        # Sample 0 has 5 valid tokens (3 on rank 0, 2 on rank 1); sample 1 has 2.
+        assert layout.valid_lengths([5, 2], rank=0) == [3, 2]
+        assert layout.valid_lengths([5, 2], rank=1) == [2, 0]
+
+    def test_segment_lengths_split_across_the_boundary(self):
+        layout = ShardLayout(original_seq_len=5, world_size=2, rank=0)
+        # Two reference images of 2 and 3 tokens; the second straddles rank 0/1.
+        assert layout.segment_lengths([[2, 3]], rank=0) == [[2, 1]]
+        assert layout.segment_lengths([[2, 3]], rank=1) == [[0, 2]]
+
+    @pytest.mark.parametrize("world_size", [1, 2, 4])
+    def test_per_rank_lengths_sum_to_the_global_length(self, world_size):
+        layout = ShardLayout(original_seq_len=13, world_size=world_size, rank=0)
+        for global_length in (0, 1, 7, 13):
+            total = sum(layout.valid_lengths([global_length], rank=r)[0] for r in range(world_size))
+            assert total == global_length
+
+
+class TestRankConcatMask:
+    """The mask must describe the sequence Ulysses forms after its all-to-all.
+
+    That sequence is rank0's local shard followed by rank1's, etc. -- not the
+    natural global ordering -- because each rank packs [context|ref|noise]
+    locally before attention.
+    """
+
+    def test_mask_is_concatenation_of_per_rank_segments(self):
+        mask = rank_concat_mask_or_none(
+            [[2], [1]],
+            capacity=3,
+            like=torch.zeros(1, 1),
+        )
+        # rank 0 contributes [T, T, F], rank 1 contributes [T, F, F]
+        assert mask.tolist() == [[True, True, False, True, False, False]]
+
+    def test_all_valid_mask_is_dropped(self):
+        assert (
+            rank_concat_mask_or_none(
+                [[3], [3]],
+                capacity=3,
+                like=torch.zeros(1, 1),
+            )
+            is None
+        )
+
+    def test_required_keeps_an_all_valid_mask(self):
+        mask = rank_concat_mask_or_none(
+            [[3], [3]],
+            capacity=3,
+            like=torch.zeros(1, 1),
+            required=True,
+        )
+        assert mask is not None
+        assert bool(mask.all())
+
+
+class TestRankConcatEquivalence:
+    """The built mask must equal an all-gather of per-rank local masks.
+
+    BOOGU used to hand rank-local masks to the attention layer, which
+    all-gathered them into the global sequence. The model now builds that
+    global mask directly -- saving one collective per attention call -- so the
+    two constructions must agree exactly. This is asserted structurally
+    because the end-to-end SP1-vs-SP2 comparison is numerically insensitive to
+    masking (a run with masks entirely disabled still matches to ~5e-7).
+    """
+
+    @staticmethod
+    def _all_gather_of_local_masks(per_rank_lengths, capacity):
+        """What the deleted `dist.all_gather` path used to produce."""
+        batch_size = len(per_rank_lengths[0])
+        local_masks = []
+        for lengths in per_rank_lengths:
+            local = torch.zeros(batch_size, capacity, dtype=torch.bool)
+            for i, length in enumerate(lengths):
+                local[i, :length] = True
+            local_masks.append(local)
+        return torch.cat(local_masks, dim=1)
+
+    @pytest.mark.parametrize(
+        ("global_len", "world_size", "sample_lengths"),
+        [
+            (5, 2, [5, 2]),
+            (8, 2, [8, 1]),
+            (13, 4, [13, 7, 0]),
+            (16, 4, [16]),
+            (7, 1, [7, 3]),
+        ],
+    )
+    def test_matches_all_gather_of_local_masks(self, global_len, world_size, sample_lengths):
+        layout = ShardLayout(original_seq_len=global_len, world_size=world_size, rank=0)
+        per_rank = [layout.valid_lengths(sample_lengths, rank=r) for r in range(world_size)]
+
+        built = rank_concat_mask_or_none(
+            per_rank,
+            layout.local_seq_len,
+            like=torch.zeros(1, 1),
+            required=True,
+        )
+        assert torch.equal(built, self._all_gather_of_local_masks(per_rank, layout.local_seq_len))

tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py

@@ -1121,7 +1121,7 @@ def test_pipeline_init_passes_tokenizer_attrs_into_transformer(
 def test_preprocess_i2v_image_and_action_video_inputs() -> None:
     from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import get_cosmos3_pre_process_func
 
-    preprocess = get_cosmos3_pre_process_func(SimpleNamespace())
+    preprocess = get_cosmos3_pre_process_func(SimpleNamespace(model_config={"guardrails": False}, tf_model_config=None))
     i2v = SimpleNamespace(
         prompt={"prompt": "A slow camera push.", "multi_modal_data": {"image": Image.new("RGB", (320, 160))}},
         sampling_params=make_sampling_params(height=None, width=None, extra_args={}),
@@ -1155,6 +1155,130 @@ def test_preprocess_i2v_image_and_action_video_inputs() -> None:
     assert additional["condition_frame_indexes_vision"] == [0, 1]
 
 
+def test_preprocess_v2v_decodes_uploaded_video_path(tmp_path) -> None:
+    """Serving may pass multipart uploads as /tmp/...mp4 path lists (#8073)."""
+    imageio = pytest.importorskip("imageio.v3")
+
+    from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import get_cosmos3_pre_process_func
+
+    frames = [np.full((16, 32, 3), i * 40, dtype=np.uint8) for i in range(6)]
+    video_path = tmp_path / "vllm_omni_video_reference_test.mp4"
+    imageio.imwrite(video_path, frames, fps=4, codec="libx264")
+
+    preprocess = get_cosmos3_pre_process_func(SimpleNamespace(model_config={"guardrails": False}, tf_model_config=None))
+    request = SimpleNamespace(
+        prompt={"prompt": "Continue.", "multi_modal_data": {"video": [str(video_path)]}},
+        sampling_params=make_sampling_params(
+            height=16,
+            width=32,
+            extra_args={"condition_frame_indexes_vision": [0, 1], "condition_video_keep": "first"},
+        ),
+    )
+
+    additional = preprocess(request).prompt["additional_information"]
+    # condition_frame_indexes_vision=[0,1] => 5 pixel frames after VAE indexing.
+    assert tuple(additional["preprocessed_video"].shape) == (1, 3, 5, 16, 32)
+    assert additional["condition_frame_indexes_vision"] == [0, 1]
+
+
+def test_decode_path_video_frames_honors_max_frames_and_keep(tmp_path, monkeypatch: pytest.MonkeyPatch) -> None:
+    pytest.importorskip("imageio.v3")
+    import imageio.v3 as iio
+
+    from vllm_omni.diffusion.utils.video_decode import decode_path_video_frames
+
+    source = tmp_path / "clip.mp4"
+    source.write_bytes(b"placeholder")
+    seen = {"n": 0}
+
+    def fake_imiter(_path):
+        for idx in range(20):
+            seen["n"] += 1
+            yield np.full((4, 4, 3), idx, dtype=np.uint8)
+
+    monkeypatch.setattr(iio, "imiter", fake_imiter)
+
+    first = decode_path_video_frames(source, max_frames=5, keep="first")
+    assert len(first) == 5
+    assert seen["n"] == 5
+    assert first[0][0, 0, 0] == 0
+    assert first[-1][0, 0, 0] == 4
+
+    seen["n"] = 0
+    last = decode_path_video_frames(source, max_frames=5, keep="last")
+    assert len(last) == 5
+    assert seen["n"] == 20
+    assert last[0][0, 0, 0] == 15
+    assert last[-1][0, 0, 0] == 19
+
+
+def test_preprocess_v2v_video_path_honors_prompt_level_keep_and_indexes(
+    tmp_path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    pytest.importorskip("imageio.v3")
+    import imageio.v3 as iio
+
+    from vllm_omni.diffusion.models.cosmos3 import pipeline_cosmos3 as cosmos_pipeline
+    from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import get_cosmos3_pre_process_func
+    from vllm_omni.diffusion.utils.video_decode import decode_path_video_frames
+
+    source = tmp_path / "vllm_omni_video_reference_prompt.mp4"
+    source.write_bytes(b"placeholder")
+    captured: dict[str, Any] = {}
+
+    def fake_imiter(_path):
+        for idx in range(20):
+            yield np.full((16, 32, 3), idx, dtype=np.uint8)
+
+    def spy_decode(*args, **kwargs):
+        frames = decode_path_video_frames(*args, **kwargs)
+        captured["keep"] = kwargs.get("keep")
+        captured["max_frames"] = kwargs.get("max_frames")
+        captured["n"] = len(frames)
+        captured["first"] = int(frames[0][0, 0, 0])
+        captured["last"] = int(frames[-1][0, 0, 0])
+        return frames
+
+    monkeypatch.setattr(iio, "imiter", fake_imiter)
+    monkeypatch.setattr(cosmos_pipeline, "decode_path_video_frames", spy_decode)
+
+    preprocess = get_cosmos3_pre_process_func(SimpleNamespace(model_config={"guardrails": False}, tf_model_config=None))
+    last_keep = SimpleNamespace(
+        prompt={
+            "prompt": "Continue.",
+            "condition_video_keep": "last",
+            "condition_frame_indexes_vision": [0, 1],
+            "multi_modal_data": {"video": [str(source)]},
+        },
+        sampling_params=make_sampling_params(height=16, width=32, extra_args={}),
+    )
+    additional = preprocess(last_keep).prompt["additional_information"]
+    assert captured["keep"] == "last"
+    assert captured["max_frames"] == 5
+    assert captured["n"] == 5
+    assert captured["first"] == 15
+    assert captured["last"] == 19
+    assert tuple(additional["preprocessed_video"].shape) == (1, 3, 5, 16, 32)
+    assert additional["condition_frame_indexes_vision"] == [0, 1]
+
+    wider_indexes = SimpleNamespace(
+        prompt={
+            "prompt": "Continue.",
+            "condition_frame_indexes_vision": [0, 1, 2],
+            "multi_modal_data": {"video": [str(source)]},
+        },
+        sampling_params=make_sampling_params(height=16, width=32, extra_args={}),
+    )
+    additional = preprocess(wider_indexes).prompt["additional_information"]
+    assert captured["keep"] == "first"
+    assert captured["max_frames"] == 9
+    assert captured["n"] == 9
+    assert captured["first"] == 0
+    assert captured["last"] == 8
+    assert tuple(additional["preprocessed_video"].shape) == (1, 3, 9, 16, 32)
+    assert additional["condition_frame_indexes_vision"] == [0, 1, 2]
+
+
 def test_transfer_config_media_helpers_and_preprocess_budget(monkeypatch: pytest.MonkeyPatch) -> None:
     from vllm_omni.diffusion.models.cosmos3 import transfer
     from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import (
@@ -1224,7 +1348,7 @@ def test_transfer_config_media_helpers_and_preprocess_budget(monkeypatch: pytest
     )
     assert tuple(loaded.shape) == (3, 2, 8, 8)
 
-    preprocess = get_cosmos3_pre_process_func(SimpleNamespace())
+    preprocess = get_cosmos3_pre_process_func(SimpleNamespace(model_config={"guardrails": False}, tf_model_config=None))
 
     class FramesWithFps(list):
         fps = 12.5

tests/diffusion/models/mammoth_moda2/test_pipeline_mammothmoda2_dit.py

@@ -4,6 +4,7 @@
 from contextlib import contextmanager
 from dataclasses import dataclass
 from types import SimpleNamespace
+from unittest.mock import patch
 
 import cache_dit
 import pytest
@@ -26,7 +27,10 @@ from vllm_omni.diffusion.models.mammoth_moda2 import pipeline_mammothmoda2_dit
 from vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit import (
     MammothModa2DiTPipeline,
     _build_mammoth_config,
+    _MammothRequest,
+    _pad_cond_sequence,
     _root_weight_source,
+    get_mammoth_moda2_pre_process_func,
 )
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
@@ -62,9 +66,7 @@ def _od_config() -> OmniDiffusionConfig:
 
 def test_build_mammoth_config_uses_shared_transformer_projection() -> None:
     config = _build_mammoth_config(_od_config())
-    assert config.model_type == "mammothmoda2"
-    assert config.get_text_config().hidden_size == 8
-    assert config.gen_dit_config == {"hidden_size": 8, "in_channels": 4}
+    assert config.gen_dit_config["hidden_size"] == 8
 
 
 def test_build_mammoth_config_rejects_empty_shared_projection() -> None:
@@ -88,11 +90,11 @@ def test_root_weight_source_forwards_revision() -> None:
     assert _root_weight_source(config).revision == "rev-7"
 
 
-def test_pipeline_declares_native_components_and_single_request_mode_only() -> None:
+def test_pipeline_declares_native_components_and_batch_request_mode() -> 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 False
+    assert MammothModa2DiTPipeline.supports_request_batch is True
     assert MammothModa2DiTPipeline.supports_step_execution is False
 
 
@@ -216,13 +218,17 @@ def test_parse_request_standard_fields_precede_request_level_fallbacks() -> None
     assert parsed.num_inference_steps == 7
 
 
-def test_parse_request_rejects_multiple_requests_and_outputs() -> None:
+def test_parse_request_parses_each_request_and_rejects_multi_output() -> None:
     pipeline = _pipeline_shell()
     batch = _batch()
     batch.requests.append(_batch(request_id="req-8").requests[0])
-    with pytest.raises(ValueError, match="exactly one request"):
-        pipeline._parse_request(batch)
-    with pytest.raises(ValueError, match="num_outputs_per_prompt=1"):
+    req0 = pipeline._parse_request(batch, index=0)
+    req1 = pipeline._parse_request(batch, index=1)
+    assert req0.request_id == "req-7"
+    assert req1.request_id == "req-8"
+    assert req0.index == 0
+    assert req1.index == 1
+    with pytest.raises(ValueError, match="num_outputs_per_prompt == 1"):
         pipeline._parse_request(_batch(sampling=OmniDiffusionSamplingParams(num_outputs_per_prompt=2)))
 
 
@@ -347,10 +353,14 @@ class _FakeVae(nn.Module):
         super().__init__()
         self.anchor = nn.Parameter(torch.zeros(1))
         self.config = _FakeVaeConfig()
+        self.last_decode_dtype: torch.dtype | None = None
 
     def decode(self, latents, return_dict=False):
         assert return_dict is False
-        return (torch.zeros(1, 3, 32, 48, dtype=latents.dtype),)
+        self.last_decode_dtype = latents.dtype
+        b = latents.shape[0]
+        base = latents[:, :1, :1, :1]
+        return (torch.zeros(b, 3, 32, 48, dtype=latents.dtype) + base,)
 
 
 class _FakeScheduler:
@@ -392,8 +402,9 @@ def test_forward_returns_diffusion_output_with_request_sampling(mocker) -> None:
         )
     )
 
-    assert isinstance(result, DiffusionOutput)
-    assert result.output.shape == (1, 3, 32, 48)
+    assert len(result) == 1
+    assert isinstance(result[0], DiffusionOutput)
+    assert result[0].output.shape == (1, 3, 32, 48)
     assert captured["shape"] == (1, 4, 4, 6)
     assert captured["seed"] == 42
     assert scheduler.requested_steps == 2
@@ -527,7 +538,7 @@ def test_forward_transitions_request_scoped_cache_dit_across_requests(request: p
     # Reference run for phase 2: identical request while hooks were never
     # installed on this pipeline.
     assert not pipeline.is_cache_dit_enabled()
-    reference = pipeline.forward(_batch(sampling=unguided_sampling)).output
+    reference = pipeline.forward(_batch(sampling=unguided_sampling))[0].output
     assert seq_lens == [cond_tokens] * 8
     assert torch.isfinite(reference.float()).all()
 
@@ -551,7 +562,7 @@ def test_forward_transitions_request_scoped_cache_dit_across_requests(request: p
     # Phase 2: no-CFG request disables the hooks and must reproduce the
     # never-enabled reference bit-for-bit.
     seq_lens.clear()
-    unguided = pipeline.forward(_batch(sampling=unguided_sampling)).output
+    unguided = pipeline.forward(_batch(sampling=unguided_sampling))[0].output
     assert seq_lens == [cond_tokens] * 8
     assert not pipeline.is_cache_dit_enabled()
     assert not BlockAdapter.is_cached(pipeline.gen_transformer)
@@ -560,7 +571,7 @@ def test_forward_transitions_request_scoped_cache_dit_across_requests(request: p
     # Phase 3: CFG with a new step count and image size re-installs the hooks
     # on a fresh context; output shape follows the new request geometry.
     seq_lens.clear()
-    regenerated = pipeline.forward(_batch(sampling=changed_cfg_sampling, height=48, width=64)).output
+    regenerated = pipeline.forward(_batch(sampling=changed_cfg_sampling, height=48, width=64))[0].output
     assert seq_lens == [cond_tokens, uncond_tokens] * 10
     assert pipeline.is_cache_dit_enabled()
     assert BlockAdapter.is_cached(pipeline.gen_transformer)
@@ -570,3 +581,442 @@ def test_forward_transitions_request_scoped_cache_dit_across_requests(request: p
     assert torch.isfinite(regenerated.float()).all()
 
     hook.remove()
+
+
+def test_pre_process_registers_batch_compatibility_key() -> None:
+    pre_process = get_mammoth_moda2_pre_process_func(_od_config())
+    req = _batch().requests[0]
+    pre_process(req)
+    assert req.batch_compatibility_key == ("mammoth_moda2_dit", 32, 48, 7)
+
+
+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]
+    with pytest.raises(ValueError, match="Missing additional_information AR conditions.*req-bad-ar"):
+        pre_process(req)
+
+
+@pytest.mark.parametrize(
+    ("height", "width", "message"),
+    [
+        (0, 32, "Invalid image size.*req-dim"),
+        (32, -1, "Invalid image size.*req-dim"),
+        (30, 32, "multiples of 16.*req-dim"),
+        (32, 31, "multiples of 16.*req-dim"),
+    ],
+)
+def test_pre_process_rejects_invalid_dimensions(height, width, message) -> None:
+    pre_process = get_mammoth_moda2_pre_process_func(_od_config())
+    batch = _batch(request_id="req-dim")
+    batch.prompts[0].update(height=height, width=width)
+    with pytest.raises(ValueError, match=message):
+        pre_process(batch.requests[0])
+
+
+def test_pre_process_rejects_non_positive_steps() -> None:
+    pre_process = get_mammoth_moda2_pre_process_func(_od_config())
+    req = _batch(
+        request_id="req-zero-steps",
+        sampling=OmniDiffusionSamplingParams(num_inference_steps=0),
+    ).requests[0]
+    with pytest.raises(ValueError, match="num_inference_steps must be positive.*req-zero-steps"):
+        pre_process(req)
+
+
+def test_pre_process_allows_dummy_run() -> None:
+    pre_process = get_mammoth_moda2_pre_process_func(_od_config())
+    req = _batch(
+        request_id="dummy_req_id",
+        prompt={"prompt": "dummy run"},
+        sampling=OmniDiffusionSamplingParams(height=512, width=512, num_inference_steps=20),
+    ).requests[0]
+    pre_process(req)
+    assert req.batch_compatibility_key == ("mammoth_moda2_dit", 512, 512, 20)
+
+
+def test_forward_casts_latents_to_vae_dtype() -> None:
+    pipeline = _pipeline_shell()
+    pipeline.gen_transformer = _FakeTransformer()
+    pipeline.gen_image_condition_refiner = None
+    vae = _FakeVae()
+    vae.anchor = nn.Parameter(torch.zeros(1, dtype=torch.float16))
+    pipeline.gen_vae = vae
+    pipeline.gen_freqs_cis = torch.zeros(1)
+    scheduler = _FakeScheduler()
+
+    module = "vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit"
+    with (
+        patch(f"{module}.FlowMatchEulerDiscreteScheduler", return_value=scheduler),
+        patch(
+            f"{module}.randn_tensor",
+            side_effect=lambda s, **kw: torch.zeros(s, device=kw.get("device"), dtype=kw.get("dtype")),
+        ),
+    ):
+        result = pipeline.forward(_batch())
+
+    assert len(result) == 1
+    assert vae.last_decode_dtype == torch.float16
+
+
+def test_group_requests_by_geometry_and_steps() -> None:
+    def make_req(idx: int, h: int, w: int, s: int) -> _MammothRequest:
+        return _MammothRequest(
+            index=idx,
+            request_id=f"r{idx}",
+            full_hidden_states=torch.zeros(1, 8),
+            full_token_ids=[1],
+            answer_start_index=0,
+            height=h,
+            width=w,
+            text_guidance_scale=1.0,
+            cfg_range=(0.0, 1.0),
+            num_inference_steps=s,
+            seed=None,
+            generator=None,
+            generator_device=None,
+        )
+
+    reqs = [make_req(0, 32, 32, 5), make_req(1, 32, 32, 5), make_req(2, 64, 64, 5)]
+    groups = MammothModa2DiTPipeline._group_requests(reqs)
+    assert sorted(groups) == [[0, 1], [2]]
+
+
+def test_pad_cond_sequence_right_pads() -> None:
+    embeds = [torch.ones(1, 2, 8), torch.ones(1, 4, 8) * 2]
+    masks = [torch.ones(1, 2, dtype=torch.bool), torch.ones(1, 4, dtype=torch.bool)]
+    padded, mask = _pad_cond_sequence(embeds, masks)
+    assert padded.shape == (2, 4, 8)
+    assert mask.tolist() == [[True, True, False, False], [True, True, True, True]]
+
+
+def test_forward_batched_matches_solo_output() -> 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)
+    scheduler = _FakeScheduler()
+
+    def fake_randn_tensor(shape, *, generator, device, dtype):
+        if isinstance(generator, list):
+            tensors = [
+                torch.full((1, *shape[1:]), float(g.initial_seed() if g else 0), device=device, dtype=dtype)
+                for g in generator
+            ]
+            return torch.cat(tensors, dim=0)
+        val = float(generator.initial_seed() if generator else 0)
+        return torch.full(shape, val, device=device, dtype=dtype)
+
+    module = "vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit"
+    with (
+        patch(f"{module}.FlowMatchEulerDiscreteScheduler", return_value=scheduler),
+        patch(f"{module}.randn_tensor", side_effect=fake_randn_tensor),
+    ):
+        req1 = _batch(
+            request_id="r1",
+            sampling=OmniDiffusionSamplingParams(
+                height=32, width=48, seed=10, guidance_scale=1.0, num_inference_steps=2
+            ),
+        ).requests[0]
+        req2 = _batch(
+            request_id="r2",
+            sampling=OmniDiffusionSamplingParams(
+                height=32, width=48, seed=20, guidance_scale=1.0, num_inference_steps=2
+            ),
+        ).requests[0]
+
+        batched_out = pipeline.forward(DiffusionRequestBatch([req1, req2]))
+        solo1 = pipeline.forward(DiffusionRequestBatch([req1]))
+        solo2 = pipeline.forward(DiffusionRequestBatch([req2]))
+
+    assert len(batched_out) == 2
+    torch.testing.assert_close(batched_out[0].output, solo1[0].output)
+    torch.testing.assert_close(batched_out[1].output, solo2[0].output)
+    assert not torch.equal(batched_out[0].output, batched_out[1].output)
+
+
+def test_refiner_output_length_preserved_when_input_len_differs_from_queries() -> None:
+    pipeline = _pipeline_shell()
+    pipeline.gen_transformer = _FakeTransformer()
+
+    class _MockRefiner(nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.anchor = nn.Parameter(torch.zeros(1))
+
+        def forward(self, features, mask=None):
+            b, _, d = features.shape
+            return torch.zeros(b, 8, d, device=features.device, dtype=features.dtype)
+
+    pipeline.gen_image_condition_refiner = _MockRefiner()
+    pipeline.gen_vae = _FakeVae()
+    pipeline.gen_freqs_cis = torch.zeros(1)
+    scheduler = _FakeScheduler()
+
+    captured_cond_lens = []
+
+    def fake_forward(*, hidden_states, text_hidden_states, **kwargs):
+        captured_cond_lens.append(text_hidden_states.shape[1])
+        return torch.zeros_like(hidden_states)
+
+    module = "vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit"
+    with (
+        patch(f"{module}.FlowMatchEulerDiscreteScheduler", return_value=scheduler),
+        patch(
+            f"{module}.randn_tensor",
+            side_effect=lambda s, **kw: torch.zeros(s, device=kw.get("device"), dtype=kw.get("dtype")),
+        ),
+        patch.object(pipeline.gen_transformer, "forward", fake_forward),
+    ):
+        prompt = {
+            "prompt": "",
+            "additional_information": {
+                "full_hidden_states": torch.zeros(6, 8),
+                "full_token_ids": [10, 11, 100, 101, 102, 103],
+                "answer_start_index": 2,
+            },
+        }
+        batch = _batch(
+            prompt=prompt,
+            sampling=OmniDiffusionSamplingParams(height=32, width=48, guidance_scale=1.0, num_inference_steps=1),
+        )
+        pipeline.forward(batch)
+
+    # 2 text tokens + 8 refined queries = 10 total conditioning tokens seen by transformer
+    assert captured_cond_lens == [10]
+
+
+def test_nested_image_embedder_compacts_unequal_text_lengths() -> None:
+    from vllm_omni.diffusion.models.mammoth_moda2.mammothmoda2_dit_model import (
+        Transformer2DModel,
+    )
+
+    model = object.__new__(Transformer2DModel)
+    nn.Module.__init__(model)
+    object.__setattr__(model, "_internal_dict", SimpleNamespace(patch_size=2))
+
+    dim = 8
+    batch_size = 2
+    max_text_len = 4
+
+    class _MockTimeCaptionEmbed(nn.Module):
+        def forward(self, timestep, text_hidden_states, dtype):
+            return torch.zeros(batch_size, dim), text_hidden_states
+
+        def image_embedder(self, feats, mask=None):
+            return torch.full((batch_size, 3, dim), 99.0)
+
+    model.time_caption_embed = _MockTimeCaptionEmbed()
+    model.x_embedder = nn.Identity()
+
+    def mock_rope(freqs, mask, *args):
+        return None, None, None, None, mask.sum(dim=-1).tolist(), []
+
+    model.rope_embedder = mock_rope
+
+    hidden_states = torch.zeros(batch_size, 4, 4, 4)
+    timestep = torch.tensor([1.0, 1.0])
+    text_hidden_states = torch.zeros(batch_size, max_text_len, dim)
+    text_hidden_states[0, 0, :] = 1.0
+    text_hidden_states[0, 1, :] = 2.0
+    text_hidden_states[1, 0, :] = 10.0
+    text_hidden_states[1, 1, :] = 20.0
+    text_hidden_states[1, 2, :] = 30.0
+    text_hidden_states[1, 3, :] = 40.0
+
+    text_attention_mask = torch.tensor(
+        [
+            [True, True, False, False],
+            [True, True, True, True],
+        ]
+    )
+
+    ar_image_hidden_states = torch.zeros(batch_size, 5, dim)
+
+    res = model._prepare_embeddings(
+        hidden_states=hidden_states,
+        timestep=timestep,
+        text_hidden_states=text_hidden_states,
+        text_attention_mask=text_attention_mask,
+        freqs_cis=torch.zeros(1),
+        batch_size=batch_size,
+        height=4,
+        width=4,
+        ar_image_hidden_states=ar_image_hidden_states,
+    )
+
+    compacted_text = res[1]
+    compacted_mask = res[2]
+    encoder_seq_lengths = res[9]
+
+    assert encoder_seq_lengths == [5, 7]
+    assert compacted_text.shape == (2, 7, dim)
+    assert compacted_text[0, 0, 0].item() == 1.0
+    assert compacted_text[0, 1, 0].item() == 2.0
+    assert compacted_text[0, 2, 0].item() == 99.0
+    assert compacted_text[0, 3, 0].item() == 99.0
+    assert compacted_text[0, 4, 0].item() == 99.0
+    assert compacted_mask[0].tolist() == [True, True, True, True, True, False, False]
+    assert compacted_mask[1].tolist() == [True, True, True, True, True, True, True]
+
+
+def test_pre_process_rejects_missing_visual_tokens_and_multi_output() -> None:
+    pre_process = get_mammoth_moda2_pre_process_func(_od_config())
+
+    multi_out = _batch(
+        request_id="req-multi",
+        sampling=OmniDiffusionSamplingParams(num_outputs_per_prompt=2),
+    ).requests[0]
+    with pytest.raises(ValueError, match="num_outputs_per_prompt == 1.*req-multi"):
+        pre_process(multi_out)
+
+    no_vis_prompt = {
+        "prompt": "",
+        "additional_information": {
+            "full_hidden_states": torch.zeros(3, 8),
+            "full_token_ids": [10, 11, 12],
+            "answer_start_index": 1,
+        },
+    }
+    no_vis_req = _batch(request_id="req-no-vis", prompt=no_vis_prompt).requests[0]
+    with pytest.raises(ValueError, match="no visual-token hidden states.*req-no-vis"):
+        pre_process(no_vis_req)
+
+    valid_prompt = {
+        "prompt": "",
+        "additional_information": {
+            "full_hidden_states": torch.zeros(3, 8),
+            "full_token_ids": [10, 100, 101],
+            "answer_start_index": 1,
+        },
+    }
+    valid_req = _batch(request_id="req-valid", prompt=valid_prompt).requests[0]
+    pre_process(valid_req)
+    assert valid_req.batch_compatibility_key is not None
+
+
+def test_cfg_precomputation_uses_cpu_booleans_without_device_sync() -> 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)
+    scheduler = _FakeScheduler()
+
+    raw_any = torch.Tensor.any
+    tensor_any_called = []
+
+    def wrapped_any(self, *args, **kwargs):
+        tensor_any_called.append(self)
+        return raw_any(self, *args, **kwargs)
+
+    module = "vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit"
+    with (
+        patch(f"{module}.FlowMatchEulerDiscreteScheduler", return_value=scheduler),
+        patch(
+            f"{module}.randn_tensor",
+            side_effect=lambda s, **kw: torch.zeros(s, device=kw.get("device"), dtype=kw.get("dtype")),
+        ),
+        patch.object(torch.Tensor, "any", wrapped_any),
+    ):
+        req1 = _batch(
+            request_id="r1",
+            sampling=OmniDiffusionSamplingParams(
+                height=32,
+                width=48,
+                seed=1,
+                guidance_scale=4.0,
+                num_inference_steps=5,
+                extra_args={"cfg_range": [0.2, 0.8]},
+            ),
+        ).requests[0]
+        req2 = _batch(
+            request_id="r2",
+            sampling=OmniDiffusionSamplingParams(
+                height=32, width=48, seed=2, guidance_scale=1.0, num_inference_steps=5
+            ),
+        ).requests[0]
+
+        pipeline.forward(DiffusionRequestBatch([req1, req2]))
+
+    assert len(tensor_any_called) == 0
+
+
+def test_dummy_request_splits_valid_ar_conditions_without_error() -> None:
+    """Startup warmup dummy run must split into valid non-empty AR conditions."""
+    pipeline = _pipeline_shell()
+    batch = _batch(
+        request_id="dummy_req_id",
+        prompt={"prompt": "dummy run"},
+        sampling=OmniDiffusionSamplingParams(height=512, width=512, seed=1, guidance_scale=0.0, num_inference_steps=2),
+    )
+    parsed = pipeline._parse_request(batch)
+    text_cond, image_cond = pipeline._split_request_conditions(parsed)
+    assert text_cond.shape == (1, pipeline._llm_hidden_size)
+    assert image_cond.shape == (1, pipeline._llm_hidden_size)
+
+
+@pytest.mark.parametrize(
+    ("llm_cfg_patch", "expected_threshold"),
+    [
+        # Parent-level llm_config threshold
+        ({"gen_vocab_start_index": 200}, 200),
+        # Nested text_config threshold
+        ({"text_config": {"model_type": "mammothmoda2_qwen2_5_vl_text", "gen_vocab_start_index": 350}}, 350),
+        # Default-derived threshold (fallback to vocab_size / 152064)
+        ({}, 152064),
+    ],
+)
+def test_admission_and_inference_threshold_resolution_matches(llm_cfg_patch: dict, expected_threshold: int) -> None:
+    """Verify identical threshold resolution for parent-level, nested, and default configs."""
+    raw = _raw_config()
+    raw["llm_config"] = {"model_type": "mammothmoda2_qwen2_5_vl", **llm_cfg_patch}
+    od_cfg = OmniDiffusionConfig(
+        model="/models/MammothModa2-Preview",
+        model_class_name="MammothModa2DiTPipeline",
+        tf_model_config=TransformerConfig.from_dict(raw),
+    )
+    # 1. Check admission preprocessor resolution
+    pre_process = get_mammoth_moda2_pre_process_func(od_cfg)
+
+    # Token just below threshold -> rejected
+    below_prompt = {
+        "prompt": "",
+        "additional_information": {
+            "full_hidden_states": torch.zeros(2, 8),
+            "full_token_ids": [10, expected_threshold - 1],
+            "answer_start_index": 1,
+        },
+    }
+    below_req = _batch(request_id="req-below", prompt=below_prompt).requests[0]
+    with pytest.raises(ValueError, match="no visual-token hidden states.*req-below"):
+        pre_process(below_req)
+
+    # Token at or above threshold -> accepted
+    at_prompt = {
+        "prompt": "",
+        "additional_information": {
+            "full_hidden_states": torch.zeros(2, 8),
+            "full_token_ids": [10, expected_threshold],
+            "answer_start_index": 1,
+        },
+    }
+    at_req = _batch(request_id="req-at", prompt=at_prompt).requests[0]
+    pre_process(at_req)
+    assert at_req.batch_compatibility_key is not None
+
+    # 2. Check inference pipeline model config resolution
+    pipeline = _pipeline_shell()
+    pipeline.config = _build_mammoth_config(od_cfg)
+
+    # Token below threshold rejected by pipeline
+    parsed_below = pipeline._parse_request(_batch(request_id="req-inf-below", prompt=below_prompt))
+    with pytest.raises(ValueError, match="no visual-token hidden states.*req-inf-below"):
+        pipeline._split_request_conditions(parsed_below)
+
+    # Token at threshold accepted by pipeline
+    parsed_at = pipeline._parse_request(_batch(request_id="req-inf-at", prompt=at_prompt))
+    text_cond, image_cond = pipeline._split_request_conditions(parsed_at)
+    assert text_cond.shape == (1, 8)
+    assert image_cond.shape == (1, 8)

tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py

@@ -3494,3 +3494,48 @@ def test_long_video_shape_requires_explicit_opt_in(duration):
     assert frames >= duration * 24 and frames % 17 == 5
     assert video_t == (frames - 5) // 17 * 5 + 2
     assert audio_t == round(frames / 24 * 40)
+
+
+@pytest.mark.parametrize("preencode", [False, True])
+@pytest.mark.parametrize("cancel_phase", ["before_prepare", "prepare", "diffuse"])
+def test_request_cancellation_at_prepare_and_decode_boundaries(preencode, cancel_phase, monkeypatch):
+    from vllm_omni.diffusion.cancellation import RequestCancellationRegistry, request_cancellation_scope
+    from vllm_omni.diffusion.data import DiffusionRequestAbortedError
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.platforms import current_omni_platform
+
+    synchronize = Mock()
+    monkeypatch.setattr(current_omni_platform, "synchronize", synchronize)
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    registry = RequestCancellationRegistry()
+    signal = registry.create("request")
+
+    def prepare(*args):
+        if cancel_phase == "prepare":
+            registry.cancel(["request"])
+        return {"num_outputs": 1, "seed": 1101, "preencode_mp4": preencode}
+
+    def diffuse(**kwargs):
+        registry.cancel(["request"])
+        return torch.zeros(1), torch.zeros(1)
+
+    pipeline._prepare_request_inputs = Mock(side_effect=prepare)
+    pipeline._denoise_kwargs = Mock(return_value={})
+    pipeline.diffuse = Mock(side_effect=diffuse)
+    pipeline.decode = Mock()
+    pipeline.decode_to_mp4 = Mock()
+    try:
+        if cancel_phase == "before_prepare":
+            registry.cancel(["request"])
+        with request_cancellation_scope([signal]), pytest.raises(DiffusionRequestAbortedError):
+            pipeline.forward(_t2va_batch())
+        synchronize.assert_called_once_with()
+        if cancel_phase == "before_prepare":
+            pipeline._prepare_request_inputs.assert_not_called()
+        if cancel_phase != "diffuse":
+            pipeline.diffuse.assert_not_called()
+        pipeline.decode.assert_not_called()
+        pipeline.decode_to_mp4.assert_not_called()
+    finally:
+        registry.close()

tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py

@@ -171,6 +171,41 @@ def test_step_execution_matches_request_mode_denoise_loop(num_steps, mocker):
     torch.testing.assert_close(state.extra[mod._STEP_AUDIO_ROWS], reference_audio)
 
 
+def test_request_mode_cancellation_stops_before_next_denoise_step(monkeypatch):
+    from vllm_omni.diffusion.cancellation import RequestCancellationRegistry, request_cancellation_scope
+    from vllm_omni.diffusion.data import DiffusionRequestAbortedError
+    from vllm_omni.diffusion.models.minimax_h3.denoise_loop import minimax_h3_denoise_loop
+    from vllm_omni.platforms import current_omni_platform
+
+    # This test runs real packing/Euler updates on CPU with the small DiT above.
+    monkeypatch.setattr(current_omni_platform, "synchronize", lambda: None)
+    branch, video_rows, audio_rows = _make_branch(text_len=9, latent_t=2, latent_h=4, latent_w=6, audio_t=3, seed=5)
+    registry = RequestCancellationRegistry()
+    signal = registry.create("request")
+    steps = []
+
+    def cancel_after_first_step(step, video, audio):
+        steps.append(step)
+        registry.cancel(["request"])
+
+    try:
+        with request_cancellation_scope([signal]), pytest.raises(DiffusionRequestAbortedError):
+            minimax_h3_denoise_loop(
+                model=_SegmentMeanModel(),
+                positive=branch,
+                initial_video_rows=video_rows,
+                initial_audio_rows=audio_rows,
+                keyframe_cond_rows=None,
+                sigmas_video=_sigmas(6, 12.0),
+                sigmas_audio=_sigmas(6, 3.0),
+                device=torch.device("cpu"),
+                on_step=cancel_after_first_step,
+            )
+    finally:
+        registry.close()
+    assert steps == [0]
+
+
 def test_step_execution_matches_request_mode_with_latent_edits():
     """Masked model rows, row timesteps, and scheduler math match both paths."""
     from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod

tests/diffusion/offloader/test_distributed_layerwise_backend.py

@@ -55,6 +55,7 @@ from vllm_omni.diffusion.offloader.offload_plan import (
 )
 from vllm_omni.diffusion.offloader.plan_resolver import ResolvedComponent, resolve_offload_plan
 from vllm_omni.diffusion.offloader.startup import OffloadStartupState, attach_offload_startup_state
+from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
 from vllm_omni.host_weight_runtime import MappedHostRegion
 from vllm_omni.platforms import current_omni_platform
 
@@ -2866,6 +2867,34 @@ class TestDistributedComponentSelection:
         with pytest.raises(RuntimeError, match="recreate the backend and reload the pipeline"):
             backend.enable(nn.Module())
 
+    def test_worker_shutdown_skips_allgather_restore(self, patched_offload_runtime, mocker):
+        # Restoring AllGather blocks rebuilds the full DiT on every rank's host,
+        # which outlived the executor's shutdown grace period (issue 8081).
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dp_size=2,
+                components=frozenset({"dit", "text_encoder"}),
+                dlo_transfers={"dit": "allgather", "text_encoder": "rank-local"},
+            ),
+            torch.device("cpu"),
+        )
+        allgather_hook = Mock(dp_size=2, next_block=nn.Linear(1, 1))
+        rank_local_hook = Mock(dp_size=1, next_block=nn.Linear(1, 1))
+        backend._all_hook_groups = [[allgather_hook], [rank_local_hook]]
+        backend.enabled = True
+        mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.destroy_distributed_env")
+        worker = DiffusionWorker.__new__(DiffusionWorker)
+        worker.model_runner = SimpleNamespace(offload_backend=backend, kv_transfer_manager=None)
+
+        worker.shutdown()
+
+        allgather_hook.restore_next_block_to_cpu.assert_not_called()
+        rank_local_hook.restore_next_block_to_cpu.assert_called_once_with()
+        assert not backend.enabled
+        assert not backend._all_hook_groups
+
     def test_disable_runs_collectives_before_best_effort_local_cleanup(
         self,
         patched_offload_runtime,

tests/diffusion/test_diffusion_engine_cleanup.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
 
 import asyncio
 import queue
@@ -432,3 +432,68 @@ def test_finalize_aborted_request_without_async_output_skips_drop() -> None:
 
     assert output.aborted is True
     engine.executor.drop_output.assert_not_called()
+
+
+@pytest.mark.parametrize("scheduler_fails", [False, True])
+def test_fail_engine_releases_cancellation_signals_after_workers_stop(scheduler_fails):
+    from multiprocessing.shared_memory import SharedMemory
+
+    from vllm_omni.diffusion.cancellation import RequestCancellationRegistry
+
+    engine = _make_engine()
+    engine._fail_pending_rpcs = Mock()
+    registry = RequestCancellationRegistry()
+    engine._request_cancellations = registry
+    name = registry.create("running")
+
+    def shutdown():
+        # Cancellation is visible before shutdown, and late readers can attach.
+        reader = SharedMemory(name=name)
+        try:
+            assert reader.buf[0] == 1
+        finally:
+            reader.close()
+
+    engine.executor.shutdown.side_effect = shutdown
+    engine.scheduler.close = Mock(side_effect=RuntimeError("scheduler cleanup failed") if scheduler_fails else None)
+    try:
+        if scheduler_fails:
+            with pytest.raises(RuntimeError, match="scheduler cleanup failed"):
+                engine._fail_engine(RuntimeError("engine failed"))
+        else:
+            engine._fail_engine(RuntimeError("engine failed"))
+            engine.close()
+            engine.executor.shutdown.assert_called_once_with()
+        with pytest.raises(FileNotFoundError):
+            SharedMemory(name=name)
+        assert not registry._signals
+    finally:
+        registry.close()
+
+
+def test_close_retains_cancellation_signal_until_stuck_worker_stops():
+    from multiprocessing.shared_memory import SharedMemory
+
+    from vllm_omni.diffusion.cancellation import RequestCancellationRegistry
+
+    engine = _make_engine()
+    registry = RequestCancellationRegistry()
+    engine._request_cancellations = registry
+    name = registry.create("running")
+    engine.worker_thread = Mock()
+    engine.worker_thread.is_alive.side_effect = [True, True, False, False]
+    try:
+        engine.close()
+        engine.executor.shutdown.assert_not_called()
+        reader = SharedMemory(name=name)
+        try:
+            assert reader.buf[0] == 1
+        finally:
+            reader.close()
+        engine.close()
+        engine.executor.shutdown.assert_called_once_with()
+        with pytest.raises(FileNotFoundError):
+            SharedMemory(name=name)
+        assert not registry._signals
+    finally:
+        registry.close()

tests/diffusion/test_diffusion_model_runner.py

@@ -1411,3 +1411,44 @@ def test_vllm_set_forward_context_implementation(monkeypatch):
             ),
         ),
     ], ERROR_MESSAGE
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+@pytest.mark.parametrize("cancel_all", [False, True])
+def test_execute_model_batch_cancellation_preserves_live_peer(monkeypatch, cancel_all):
+    from vllm_omni.diffusion.cancellation import (
+        RequestCancellationRegistry,
+        check_request_cancellation,
+    )
+    from vllm_omni.platforms import current_omni_platform
+
+    monkeypatch.setattr(current_omni_platform, "synchronize", lambda: None)
+
+    class CancellableBatchPipeline(_BatchPipeline):
+        supports_request_cancellation = True
+
+        def forward(self, batch):
+            check_request_cancellation()
+            return super().forward(batch)
+
+    monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(model_runner_module, "current_omni_platform", _fake_platform_for_peak_memory())
+    pipeline = CancellableBatchPipeline(outputs=[DiffusionOutput(output="a"), DiffusionOutput(output="b")])
+    runner = _make_batch_runner(pipeline)
+    sched = _make_scheduler_output(num_reqs=2)
+    registry = RequestCancellationRegistry()
+    try:
+        for entry in sched.scheduled_new_reqs:
+            entry.req.cancellation_signal = registry.create(entry.request_id)
+        registry.cancel(["req-0", "req-1"] if cancel_all else ["req-0"])
+        result = runner.execute_model_batch(sched, runner.od_config)
+        assert len(result.runner_outputs) == 2
+        assert [output.result.aborted for output in result.runner_outputs] == [cancel_all, cancel_all]
+        if not cancel_all:
+            assert result.runner_outputs[1].request_id == "req-1"
+            assert result.runner_outputs[1].result.output == "b"
+        else:
+            assert pipeline.last_batch is None
+    finally:
+        registry.close()

tests/diffusion/test_request_cancellation.py

@@ -0,0 +1,276 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Cancellation transport, ownership and collective agreement regressions."""
+
+import multiprocessing as mp
+import queue
+import threading
+from concurrent.futures import ThreadPoolExecutor
+from multiprocessing.shared_memory import SharedMemory
+from unittest.mock import patch
+
+import pytest
+
+from vllm_omni.diffusion.cancellation import (
+    RequestCancellationRegistry,
+    check_request_cancellation,
+    request_cancellation_scope,
+)
+from vllm_omni.diffusion.data import DiffusionRequestAbortedError, OmniDiffusionConfig
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+@pytest.fixture
+def registry(monkeypatch):
+    from vllm_omni.platforms import current_omni_platform
+
+    # These tests exercise the cancellation protocol without device work.
+    monkeypatch.setattr(current_omni_platform, "synchronize", lambda: None)
+    signals = RequestCancellationRegistry()
+    yield signals
+    signals.close()
+
+
+def test_cancel_is_request_scoped_and_id_reuse_is_fresh(registry):
+    first = registry.create("first")
+    other = registry.create("other")
+    registry.cancel(["missing", "first"])
+    with pytest.raises(ValueError, match="Duplicate cancellation request"):
+        registry.create("first")
+    with request_cancellation_scope([first]):
+        with pytest.raises(DiffusionRequestAbortedError):
+            check_request_cancellation()
+    with request_cancellation_scope([other]):
+        check_request_cancellation()
+    registry.finish("first")
+    registry.finish("first")
+    with pytest.raises(FileNotFoundError):
+        SharedMemory(name=first)
+    replacement = registry.create("first")
+    with request_cancellation_scope([replacement]):
+        check_request_cancellation()
+
+
+def test_scope_restores_outer_state_and_does_not_leak_to_threads(registry):
+    signal = registry.create("request")
+    registry.cancel_all()
+    with request_cancellation_scope([signal]):
+        with request_cancellation_scope([None]):
+            check_request_cancellation()
+        with ThreadPoolExecutor(max_workers=1) as executor:
+            executor.submit(check_request_cancellation).result(timeout=10)
+        with pytest.raises(DiffusionRequestAbortedError):
+            check_request_cancellation()
+    check_request_cancellation()
+
+
+def test_mixed_batch_keeps_uncancelled_request_running(registry):
+    first = registry.create("first")
+    other = registry.create("other")
+    with request_cancellation_scope([first, other]):
+        registry.cancel(["first"])
+        check_request_cancellation()
+        registry.cancel(["other"])
+        with pytest.raises(DiffusionRequestAbortedError):
+            check_request_cancellation()
+
+
+def test_checkpoint_observes_cancel_during_device_wait(registry, monkeypatch):
+    from vllm_omni.platforms import current_omni_platform
+
+    first = registry.create("first")
+    signal = registry.create("request")
+    registry.cancel(["first"])
+    # A local cancellation triggers the device wait. A peer can be cancelled
+    # during that wait; re-read before deciding whether the whole batch stops.
+    monkeypatch.setattr(current_omni_platform, "synchronize", lambda: registry.cancel(["request"]))
+    with request_cancellation_scope([first, signal]):
+        with pytest.raises(DiffusionRequestAbortedError):
+            check_request_cancellation(synchronize=True)
+
+
+@pytest.mark.parametrize("synchronize_early", [False, True])
+def test_checkpoint_drains_once_before_abort_only(registry, monkeypatch, synchronize_early):
+    from unittest.mock import Mock
+
+    from vllm_omni.platforms import current_omni_platform
+
+    synchronize = Mock()
+    monkeypatch.setattr(current_omni_platform, "synchronize", synchronize)
+    signal = registry.create("request")
+    with request_cancellation_scope([signal]):
+        check_request_cancellation(synchronize=synchronize_early)
+        synchronize.assert_not_called()
+        # Simulate work queued after a clean checkpoint. Cancellation at the
+        # next checkpoint must drain it before exception unwinding drops refs.
+        pending_work = [object()]
+        synchronize.side_effect = pending_work.clear
+        registry.cancel(["request"])
+        with pytest.raises(DiffusionRequestAbortedError):
+            try:
+                check_request_cancellation(synchronize=synchronize_early)
+            finally:
+                assert not pending_work
+        synchronize.assert_called_once_with()
+
+
+def test_close_removes_all_owned_signals(registry):
+    names = [registry.create(str(i)) for i in range(3)]
+    registry.cancel_all()
+    registry.close()
+    registry.close()
+    for name in names:
+        with pytest.raises(FileNotFoundError):
+            SharedMemory(name=name)
+
+
+def _admission_engine(registry):
+    from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
+    from vllm_omni.diffusion.sched.request_scheduler import RequestScheduler
+
+    engine = object.__new__(DiffusionEngine)
+    engine.scheduler = RequestScheduler()
+    engine.scheduler.initialize(OmniDiffusionConfig())
+    engine._request_cancellations = registry
+    engine._cv = threading.Condition(threading.RLock())
+    engine._closed = False
+    engine._out_streams = {}
+    engine.abort_queue = queue.Queue()
+    return engine
+
+
+def _request(request_id):
+    from vllm_omni.diffusion.request import OmniDiffusionRequest
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+    return OmniDiffusionRequest(
+        prompt="A fox in the snow.", sampling_params=OmniDiffusionSamplingParams(), request_id=request_id
+    )
+
+
+def test_engine_signals_worker_before_busy_loop_drains_abort(registry):
+    from vllm_omni.diffusion.sched.interface import DiffusionRequestStatus
+
+    engine = _admission_engine(registry)
+    request = _request("running")
+    engine._add_prepared_request(request)
+    engine.scheduler.schedule()
+    signal = request.cancellation_signal
+    assert signal is not None
+    with request_cancellation_scope([signal]):
+        engine.abort("running")
+        # The busy loop is still occupied by execution. Only the signal has
+        # changed; scheduler mutation must stay on that loop's thread.
+        assert engine.scheduler.get_request_state("running").status == DiffusionRequestStatus.RUNNING
+        with pytest.raises(DiffusionRequestAbortedError):
+            check_request_cancellation()
+    engine._process_aborts_queue()
+    assert engine._finalize_finished_request("running").aborted
+    assert engine.scheduler.get_request_state("running") is None
+    with pytest.raises(FileNotFoundError):
+        SharedMemory(name=signal)
+
+
+def test_rejected_admission_releases_signal(registry, monkeypatch):
+    engine = _admission_engine(registry)
+    allocated_names = []
+
+    def reject(request):
+        allocated_names.append(request.cancellation_signal)
+        raise ValueError("admission rejected")
+
+    monkeypatch.setattr(engine.scheduler, "add_request", reject)
+    request = _request("rejected")
+    with pytest.raises(ValueError, match="admission rejected"):
+        engine._add_prepared_request(request)
+    assert request.cancellation_signal is None
+    with pytest.raises(FileNotFoundError):
+        SharedMemory(name=allocated_names[0])
+    registry.create("rejected")
+
+
+def _collective_worker(rank, name, rendezvous, ready, first_check, second_check, results):
+    import torch
+    import torch.distributed as dist
+
+    from vllm_omni.diffusion.distributed import parallel_state
+    from vllm_omni.platforms import current_omni_platform
+
+    dist.init_process_group("gloo", init_method=f"file://{rendezvous}", rank=rank, world_size=2)
+    # Exercise a real CPU collective without constructing GPU model groups.
+    group = object.__new__(parallel_state.GroupCoordinator)
+    group.world_size = 2
+    group.cpu_group = dist.group.WORLD
+    parallel_state._WORLD = group
+    try:
+        with (
+            torch.inference_mode(),
+            patch.object(current_omni_platform, "synchronize") as synchronize,
+            request_cancellation_scope([name]),
+        ):
+            dist.barrier()
+            if rank == 0:
+                ready.set()
+            assert first_check.wait(120)
+            check_request_cancellation()
+            results.put((rank, "partial_wave_kept_running"))
+            # A replica missing its token (or all local requests) still votes
+            # false. Both ranks must reach each collective without deadlocking.
+            for absent in ([None], []):
+                with request_cancellation_scope([name] if rank == 0 else absent):
+                    check_request_cancellation()
+            results.put((rank, "missing_signal_kept_running"))
+            synchronize.assert_not_called()
+            assert second_check.wait(120)
+            try:
+                check_request_cancellation()
+            except DiffusionRequestAbortedError:
+                synchronize.assert_called_once_with()
+                results.put((rank, "fully_cancelled_wave_stopped"))
+            else:
+                raise AssertionError("All cancelled ranks must stop")
+    finally:
+        dist.destroy_process_group()
+
+
+def test_spawned_ranks_agree_before_leaving_collective_wave(registry, tmp_path):
+    ctx = mp.get_context("spawn")
+    names = [registry.create(str(rank)) for rank in range(2)]
+    ready, first_check, second_check = ctx.Event(), ctx.Event(), ctx.Event()
+    results = ctx.Queue()
+    workers = [
+        ctx.Process(
+            target=_collective_worker,
+            args=(rank, names[rank], str(tmp_path / "rendezvous"), ready, first_check, second_check, results),
+        )
+        for rank in range(2)
+    ]
+    for worker in workers:
+        worker.start()
+    try:
+        assert ready.wait(120)
+        registry.cancel(["0"])
+        first_check.set()
+        assert {results.get(timeout=120) for _ in range(2 * len(workers))} == {
+            (0, "partial_wave_kept_running"),
+            (1, "partial_wave_kept_running"),
+            (0, "missing_signal_kept_running"),
+            (1, "missing_signal_kept_running"),
+        }
+        registry.cancel(["1"])
+        second_check.set()
+        assert {results.get(timeout=120) for _ in workers} == {
+            (0, "fully_cancelled_wave_stopped"),
+            (1, "fully_cancelled_wave_stopped"),
+        }
+        for worker in workers:
+            worker.join(timeout=30)
+            assert worker.exitcode == 0
+    finally:
+        for worker in workers:
+            if worker.is_alive():
+                worker.terminate()
+            worker.join(timeout=30)
+        results.close()
+        results.join_thread()

tests/diffusion/test_worker_wrapper_base.py

@@ -290,7 +290,7 @@ class TestWorkerWrapperBaseDelegation:
     def test_worker_shutdown_releases_model_before_cumem_pools(self, mocker: MockerFixture, native_kv):
         events: list[str] = []
         offload_backend = mocker.Mock()
-        offload_backend.disable.side_effect = lambda: events.append("offload")
+        offload_backend.shutdown.side_effect = lambda: events.append("offload")
         kv_manager = mocker.Mock()
         kv_manager.close.side_effect = lambda: events.append("kv")
         worker = DiffusionWorker.__new__(DiffusionWorker)
@@ -389,7 +389,7 @@ class TestWorkerWrapperBaseDelegation:
         mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.shutdown_kv_connector")
         offload_backend = mocker.Mock()
         if cleanup_fails:
-            offload_backend.disable.side_effect = RuntimeError("offload cleanup failed")
+            offload_backend.shutdown.side_effect = RuntimeError("offload cleanup failed")
         worker = DiffusionWorker.__new__(DiffusionWorker)
         weights = torch.ones(1)
         weights_ref = weakref.ref(weights)

tests/e2e/online_serving/helpers/minicpmo_window_e2e.py

@@ -38,10 +38,13 @@ async def run_window_turn(
         temperature=0.0,
         extra_body={
             "sliding_window_mode": mode,
-            "context_max_units": 1,
+            # Keep a complete ~5.5 s fixture utterance available to the
+            # native listen/speak policy. These limits still rebuild KV
+            # during the two-repeat input and its response.
+            "context_max_units": 8,
             "context_previous_max_tokens": 16,
-            "basic_window_high_tokens": 64,
-            "basic_window_low_tokens": 32,
+            "basic_window_high_tokens": 192,
+            "basic_window_low_tokens": 128,
         },
     )
     client = DuplexClient(
@@ -58,6 +61,10 @@ async def run_window_turn(
                 is_speech=True,
                 video_frames=video_frames,
             )
+            # The fixture ends in speech. A native duplex commit seals input
+            # but does not force an answer; keep the microphone clock running
+            # through a pause so the model can choose to speak.
+            await client.stream_pcm(bytes(5 * 32000), chunk_ms=200, realtime=True, is_speech=False)
             await client.commit(final=True)
             await wait_for_condition(
                 lambda: collector.count("input_audio_buffer.committed") > 0 or bool(collector.errors()),

tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py

@@ -6,17 +6,16 @@ import asyncio
 import base64
 import json
 import math
-import uuid
 import wave
 from collections.abc import Sequence
 from contextlib import asynccontextmanager, suppress
 from pathlib import Path
-from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
 
 import numpy as np
 from scipy.signal import resample_poly
 
-from vllm_omni.experimental.fullduplex.client import RealtimeDuplexClient, wait_for, write_pcm16_wav
+from vllm_omni.clients.duplex import DuplexClient, EventCollector, wait_for_condition, write_pcm16_wav
+from vllm_omni.clients.nemotron_voicechat import create_duplex_session_config
 
 INPUT_SAMPLE_RATE_HZ = 16_000
 OUTPUT_SAMPLE_RATE_HZ = 22_050
@@ -45,13 +44,6 @@ DEFAULT_FUNCTION_INSTRUCTIONS = (
 )
 
 
-def _url(base_url: str, model: str, session_id: str) -> str:
-    parts = urlsplit(base_url)
-    query = dict(parse_qsl(parts.query, keep_blank_values=True))
-    query.update(duplex="1", model=model, autostart="0", session_id=session_id)
-    return urlunsplit((parts.scheme, parts.netloc, parts.path, urlencode(query), parts.fragment))
-
-
 def _read_wav(path: Path, *, input_channel: int = 0) -> np.ndarray:
     with wave.open(str(path), "rb") as wav_file:
         channels = wav_file.getnchannels()
@@ -73,50 +65,42 @@ def _read_wav(path: Path, *, input_channel: int = 0) -> np.ndarray:
     return np.ascontiguousarray(pcm, dtype="<f4")
 
 
-async def _stream(client: RealtimeDuplexClient, pcm: np.ndarray, *, max_frames: int | None, realtime: bool) -> int:
+async def _stream(client: DuplexClient, pcm: np.ndarray, *, max_frames: int | None, realtime: bool) -> int:
     count = math.ceil(pcm.size / FRAME_SAMPLES)
     if max_frames is not None:
         count = min(count, max_frames)
     for seq in range(count):
         frame = pcm[seq * FRAME_SAMPLES : (seq + 1) * FRAME_SAMPLES]
         frame = np.pad(frame, (0, FRAME_SAMPLES - frame.size)).astype("<f4")
-        await client.send(
-            {
-                "type": "input_audio_buffer.append",
-                "audio": base64.b64encode(frame).decode("ascii"),
-                "format": "pcm_f32le",
-                "sample_rate_hz": INPUT_SAMPLE_RATE_HZ,
-                "duration_ms": 80,
-                "audio_end_ms": (seq + 1) * 80,
-            }
-        )
+        await client.append_audio(frame.tobytes())
         if realtime:
             await asyncio.sleep(FRAME_PERIOD_S)
     return count
 
 
-def _events(client: RealtimeDuplexClient, event_type: str) -> list[dict[str, object]]:
-    return [event for event in client.events.events if event.get("type") == event_type]
+def _events(collector: EventCollector, event_type: str) -> list[dict[str, object]]:
+    return [event for event in collector.events if event.get("type") == event_type]
 
 
 async def _return_function_output_when_ready(
-    client: RealtimeDuplexClient,
+    client: DuplexClient,
+    collector: EventCollector,
     *,
     output: str,
     timeout_s: float,
 ) -> tuple[str, int]:
-    await wait_for(
-        lambda: bool(client.events.errors()) or client.events.count("response.function_call_arguments.done") > 0,
+    await wait_for_condition(
+        lambda: bool(collector.errors()) or collector.count("response.function_call_arguments.done") > 0,
         timeout_s=timeout_s,
         label="function call to execute",
     )
-    if client.events.errors():
-        raise AssertionError(f"function call failed before tool execution: {client.events.errors()}")
-    function_done = _events(client, "response.function_call_arguments.done")[-1]
+    if collector.errors():
+        raise AssertionError(f"function call failed before tool execution: {collector.errors()}")
+    function_done = _events(collector, "response.function_call_arguments.done")[-1]
     call_id = function_done.get("call_id")
     if not isinstance(call_id, str) or not call_id:
         raise AssertionError(f"completed function call has no call_id: {function_done}")
-    event_count_before_output = len(client.events.events)
+    event_count_before_output = len(collector.events)
     await client.send(
         {
             "type": "conversation.item.create",
@@ -130,22 +114,26 @@ async def _return_function_output_when_ready(
     return call_id, event_count_before_output
 
 
-def _write_events(path: Path, client: RealtimeDuplexClient) -> None:
+def _write_events(path: Path, collector: EventCollector) -> None:
     path.write_text(
-        "".join(json.dumps(event, ensure_ascii=False) + "\n" for event in client.events.events),
+        "".join(json.dumps(event, ensure_ascii=False) + "\n" for event in collector.events),
         encoding="utf-8",
     )
 
 
 @asynccontextmanager
-async def _managed_client(client: RealtimeDuplexClient, *, timeout_s: float):
-    async with client:
-        try:
+async def _managed_client(client: DuplexClient, collector: EventCollector):
+    consume_task = asyncio.create_task(collector.consume(client))
+    # Subscribe before the handshake so the collector sees session.created.
+    await asyncio.sleep(0)
+    try:
+        async with client:
             yield
-        finally:
-            if client.events.count("session.created") and not client.events.count("session.closed"):
-                with suppress(Exception):
-                    await client.close_session(timeout_s=min(timeout_s, 30.0))
+        await consume_task
+    finally:
+        consume_task.cancel()
+        with suppress(asyncio.CancelledError):
+            await consume_task
 
 
 async def run(args: argparse.Namespace) -> dict[str, object]:
@@ -156,33 +144,22 @@ async def run(args: argparse.Namespace) -> dict[str, object]:
     else:
         instructions = args.instructions
     tools = DEFAULT_FUNCTION_TOOLS if args.expect_function_call else None
-    session_id = f"nemotron-voicechat-{uuid.uuid4().hex}"
-    client = RealtimeDuplexClient(_url(args.url, args.model, session_id))
-    async with _managed_client(client, timeout_s=args.timeout_s):
-        session_payload: dict[str, object] = {
-            "session_id": session_id,
-            "model": args.model,
-            "modalities": ["audio", "text"],
-            "input_audio_format": "pcm_f32le",
-            "output_audio_format": "pcm16",
-            "instructions": instructions,
-            "idle_timeout_s": args.timeout_s,
-            "turn_detection": None,
-            "extra_body": {"auto_response": True},
-        }
-        if tools is not None:
-            session_payload["tools"] = tools
-        await client.send({"type": "session.update", "session": session_payload})
-        await wait_for(
-            lambda: client.events.count("session.created") > 0 or bool(client.events.errors()),
-            timeout_s=args.timeout_s,
-            label="session.created",
-        )
-        if client.events.errors():
-            raise AssertionError(f"session setup failed: {client.events.errors()}")
-        created = _events(client, "session.created")[-1]
-        session = created.get("session")
-        capabilities = session.get("capabilities") if isinstance(session, dict) else None
+    collector = EventCollector()
+    client = DuplexClient(
+        args.url,
+        model=args.model,
+        config=create_duplex_session_config(
+            instructions=instructions,
+            tools=tools,
+            idle_timeout_s=args.timeout_s,
+        ),
+        handshake_timeout_s=args.timeout_s,
+        reconnect=None,
+        heartbeat_interval_s=None,
+    )
+    async with _managed_client(client, collector):
+        session_id = client.session_id
+        capabilities = client.session_info.get("capabilities")
         expected = dict(
             implementation_level="model_native_duplex",
             chunk_period_ms=80,
@@ -197,6 +174,7 @@ async def run(args: argparse.Namespace) -> dict[str, object]:
             asyncio.create_task(
                 _return_function_output_when_ready(
                     client,
+                    collector,
                     output=args.function_output,
                     timeout_s=args.timeout_s,
                 )
@@ -206,19 +184,19 @@ async def run(args: argparse.Namespace) -> dict[str, object]:
         )
         pcm = _read_wav(Path(args.input_wav), input_channel=args.input_channel)
         frame_count = await _stream(client, pcm, max_frames=args.max_frames, realtime=not args.no_realtime)
-        completed_responses_at_commit = client.events.count("response.done")
-        await client.send({"type": "input_audio_buffer.commit", "final": True})
-        await wait_for(
+        completed_responses_at_commit = collector.count("response.done")
+        await client.commit()
+        await wait_for_condition(
             lambda: (
-                bool(client.events.errors())
+                bool(collector.errors())
                 or (
-                    client.events.count("response.function_call_arguments.done") > 0
+                    collector.count("response.function_call_arguments.done") > 0
                     if args.expect_function_call
                     else (
-                        client.events.count("response.output_audio.delta") >= args.minimum_audio_chunks
+                        collector.count("response.output_audio.delta") >= args.minimum_audio_chunks
                         and (
                             args.allow_incomplete_response
-                            or client.events.count("response.done") > completed_responses_at_commit
+                            or collector.count("response.done") > completed_responses_at_commit
                         )
                     )
                 )
@@ -227,9 +205,9 @@ async def run(args: argparse.Namespace) -> dict[str, object]:
             label="model output",
         )
         await asyncio.sleep(args.drain_s)
-        if client.events.errors():
-            raise AssertionError(f"Realtime session emitted errors: {client.events.errors()}")
-        done_events = _events(client, "response.done")
+        if collector.errors():
+            raise AssertionError(f"Realtime session emitted errors: {collector.errors()}")
+        done_events = _events(collector, "response.done")
         if not args.expect_function_call and not args.allow_incomplete_response:
             response = done_events[-1].get("response") if done_events else None
             status = response.get("status") if isinstance(response, dict) else None
@@ -237,24 +215,22 @@ async def run(args: argparse.Namespace) -> dict[str, object]:
                 raise AssertionError(f"response did not complete successfully: {done_events[-1:]}")
 
         function_events = [
-            event for event in client.events.events if str(event.get("type", "")).startswith("response.function_call")
+            event for event in collector.events if str(event.get("type", "")).startswith("response.function_call")
         ]
         function_items = [
-            event
-            for event in _events(client, "response.output_item.done")
-            if isinstance(event.get("item"), dict) and event["item"].get("type") == "function_call"
+            item
+            for event in _events(collector, "response.output_item.done")
+            if isinstance((item := event.get("item")), dict) and item.get("type") == "function_call"
         ]
         if args.expect_function_call and not any(
             event.get("type") == "response.function_call_arguments.done" for event in function_events
         ):
             raise AssertionError(f"no completed function call: {function_events}")
         if args.expect_function_call:
-            matching_items = [
-                event for event in function_items if event["item"].get("name") == args.expected_function_name
-            ]
+            matching_items = [item for item in function_items if item.get("name") == args.expected_function_name]
             if not matching_items:
                 raise AssertionError(f"expected {args.expected_function_name!r}, got {function_items}")
-            function_item = matching_items[-1]["item"]
+            function_item: dict[str, object] = matching_items[-1]
             try:
                 function_arguments = json.loads(str(function_item.get("arguments", "")))
             except json.JSONDecodeError as exc:
@@ -275,7 +251,7 @@ async def run(args: argparse.Namespace) -> dict[str, object]:
                 assert returned_call_id == call_id
 
                 def tool_result_completed() -> bool:
-                    later = client.events.events[event_count_before_output:]
+                    later = collector.events[event_count_before_output:]
                     transcript = "".join(
                         str(event.get("delta", ""))
                         for event in later
@@ -287,23 +263,23 @@ async def run(args: argparse.Namespace) -> dict[str, object]:
                         and (args.expected_post_tool_text is None or args.expected_post_tool_text.lower() in transcript)
                     )
 
-                await wait_for(
-                    lambda: bool(client.events.errors()) or tool_result_completed(),
+                await wait_for_condition(
+                    lambda: bool(collector.errors()) or tool_result_completed(),
                     timeout_s=args.timeout_s,
                     label="completed response after function output",
                 )
-                if client.events.errors():
-                    raise AssertionError(f"function output failed: {client.events.errors()}")
+                if collector.errors():
+                    raise AssertionError(f"function output failed: {collector.errors()}")
                 await asyncio.sleep(args.drain_s)
 
-        audio = client.events.audio_bytes()
-        audio_events = _events(client, "response.output_audio.delta")
+        audio = collector.audio_bytes()
+        audio_events = _events(collector, "response.output_audio.delta")
         rates = {event.get("sample_rate_hz") for event in audio_events}
         if not args.expect_function_call and audio and rates != {OUTPUT_SAMPLE_RATE_HZ}:
             raise AssertionError(f"unexpected output sample rates: {rates}")
         if not args.expect_function_call and args.minimum_audio_chunks and not audio:
             raise AssertionError("model produced no audio")
-        expected_bytes = 2 * OUTPUT_SAMPLE_RATE_HZ * expected["chunk_period_ms"] // 1000
+        expected_bytes = 2 * OUTPUT_SAMPLE_RATE_HZ * int(str(expected["chunk_period_ms"])) // 1000
         packet_sizes = [len(base64.b64decode(str(event.get("delta", "")), validate=True)) for event in audio_events]
         if not args.expect_function_call and any(size != expected_bytes for size in packet_sizes):
             raise AssertionError(f"audio deltas are not fixed 80 ms PCM16 packets: {packet_sizes}")
@@ -314,17 +290,17 @@ async def run(args: argparse.Namespace) -> dict[str, object]:
                 f"model output RMS {audio_rms:.6f} is below {args.minimum_audio_rms:.6f}; "
                 "received packets contain only silence"
             )
-    _write_events(output_dir / "events.jsonl", client)
+    _write_events(output_dir / "events.jsonl", collector)
     if audio:
         write_pcm16_wav(output_dir / "output.wav", audio, sample_rate_hz=OUTPUT_SAMPLE_RATE_HZ)
-    result = {
+    result: dict[str, object] = {
         "ok": True,
-        "session_id": session_id,
+        "session_id": session_id if isinstance(session_id, str) else None,
         "input_frames": frame_count,
         "capabilities": capabilities,
         "event_counts": {
-            event_type: client.events.count(event_type)
-            for event_type in sorted({str(event.get("type")) for event in client.events.events})
+            event_type: collector.count(event_type)
+            for event_type in sorted({str(event.get("type")) for event in collector.events})
         },
         "audio_bytes": len(audio),
         "audio_rms": audio_rms,

tests/e2e/online_serving/test_nemotron_voicechat_duplex.py

@@ -29,16 +29,7 @@ _SERVER_ENV: dict[str, str] = {"VLLM_OMNI_INPUT_WAIT_TIMEOUT_S": str(SERVER_INPU
 if TOKENIZER:
     _SERVER_ENV["NEMOTRON_VOICECHAT_LLM_PATH"] = TOKENIZER
 
-pytestmark = [
-    pytest.mark.full_model,
-    pytest.mark.omni,
-    # Since #7413 a duplex-capable model is served by the duplex engine only when
-    # its pipeline declares a duplex_plugin. The Nemotron VoiceChat pipeline does
-    # not, so the server starts the turn engine and the /v1/realtime websocket
-    # falls through to vllm's speech-to-text realtime handler, which rejects the
-    # native duplex client's session.update ("Missing required field: model").
-    pytest.mark.skip(reason="Nemotron VoiceChat is not served by the duplex engine yet (no duplex_plugin)"),
-]
+pytestmark = [pytest.mark.full_model, pytest.mark.omni]
 
 
 @hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=1)
@@ -71,6 +62,10 @@ def test_native_duplex_turn_taking_streams_model_audio(omni_server, tmp_path: Pa
     )
     result = asyncio.run(run(args))
     assert result["ok"] is True
-    assert result["input_frames"] == 190
-    assert result["event_counts"]["response.speak"] > 0
-    assert result["audio_bytes"] >= result["input_frames"] * 1764 * 2 // 4
+    input_frames = result["input_frames"]
+    event_counts = result["event_counts"]
+    audio_bytes = result["audio_bytes"]
+    assert isinstance(input_frames, int) and input_frames == 190
+    assert isinstance(event_counts, dict) and int(str(event_counts.get("response.speak", 0))) > 0
+    assert isinstance(audio_bytes, int)
+    assert audio_bytes >= input_frames * 1764 * 2 // 4

tests/engine/duplex/test_duplex_lease.py

@@ -85,6 +85,32 @@ def test_detach_grace_and_resume_advance_the_lease_generation() -> None:
         session.resume_lease(expected_lease_generation=0)
 
 
+def test_a_resume_replayed_under_its_control_id_answers_without_resuming_again() -> None:
+    """A caller that lost the answer (cancelled or timed-out waiter) asks again with the same id.
+
+    The engine must tell it the generation that resume produced, not bump the
+    lease a second time, and must still refuse any other resume against the
+    old generation.
+    """
+    clock = FakeMonotonicClock(10.0)
+    session = _session("sid-replay", clock)
+    session.detach_lease()
+
+    assert session.resume_lease(expected_lease_generation=0, control_id="rpc-a") == 1
+    clock.advance(5.0)
+    session.detach_lease()
+
+    assert session.resume_lease(expected_lease_generation=0, control_id="rpc-a") == 1
+    assert session.lease_generation == 1
+    assert session.lease.detached_at == 15.0, "a replay is an answer, not a new resume"
+    with pytest.raises(ValueError, match="lease generation mismatch"):
+        session.resume_lease(expected_lease_generation=0, control_id="rpc-b")
+
+    assert session.resume_lease(expected_lease_generation=1, control_id="rpc-b") == 2
+    with pytest.raises(ValueError, match="lease generation mismatch"):
+        session.resume_lease(expected_lease_generation=0, control_id="rpc-a")
+
+
 def test_active_operation_prevents_mid_transaction_expiry() -> None:
     clock = FakeMonotonicClock(0.0)
     session = _session("sid-operation", clock, idle_ttl_s=5.0)

tests/engine/duplex/test_session_manager.py

@@ -296,6 +296,7 @@ class Harness:
     clock: FakeClock
     output_sink: asyncio.Queue = field(default_factory=asyncio.Queue)
     result_sink: asyncio.Queue = field(default_factory=asyncio.Queue)
+    _control_seq: int = 0
 
     @classmethod
     def create(
@@ -344,20 +345,30 @@ class Harness:
         )
         return await self.result()
 
-    async def resume(self, session_id: str, *, expected_lease_generation: int) -> DuplexControlResultMessage:
+    async def resume(
+        self, session_id: str, *, expected_lease_generation: int, control_id: str | None = None
+    ) -> DuplexControlResultMessage:
+        # Distinct by default: the engine answers a repeated control id as a
+        # replay of the same resume, which only the replay test wants.
+        self._control_seq += 1
         await self.manager.handle(
             ResumeDuplexSessionMessage(
-                control_id=f"resume-{session_id}-{expected_lease_generation}",
+                control_id=control_id or f"resume-{session_id}-{expected_lease_generation}-{self._control_seq}",
                 session_id=session_id,
                 expected_lease_generation=expected_lease_generation,
             )
         )
         return await self.result()
 
-    async def touch(self, session_id: str, activity: str) -> DuplexControlResultMessage:
+    async def touch(
+        self, session_id: str, activity: str, *, expected_lease_generation: int | None = None
+    ) -> DuplexControlResultMessage:
         await self.manager.handle(
             TouchDuplexSessionMessage(
-                control_id=f"touch-{session_id}-{activity}", session_id=session_id, activity=activity
+                control_id=f"touch-{session_id}-{activity}-{expected_lease_generation}",
+                session_id=session_id,
+                activity=activity,
+                expected_lease_generation=expected_lease_generation,
             )
         )
         return await self.result()
@@ -1656,6 +1667,168 @@ async def test_reaper_loop_survives_one_cleanup_failure(first_cleanup_delay: flo
         await asyncio.wait_for(task, timeout=5.0)
 
 
+# --------------------------------------------------------------------------- #
+# Cancelled open (#7636 Issue 1)                                              #
+# --------------------------------------------------------------------------- #
+
+
+class _BlockingResultSink(asyncio.Queue):
+    """A result sink whose ``put`` parks, so an open can be cancelled after it admitted the session."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.entered = asyncio.Event()
+        self.gate = asyncio.Event()
+
+    async def put(self, item: object) -> None:
+        self.entered.set()
+        await self.gate.wait()
+        await super().put(item)
+
+
+async def test_an_open_cancelled_while_awaiting_the_plugin_frees_the_admission_slot() -> None:
+    """``CancelledError`` is not an ``Exception``: the rollback used to be skipped.
+
+    Admission counts ``runners | closing | admitting``, so every open cancelled
+    mid-flight (engine teardown, task cancellation) burned one slot for good;
+    after ``max_sessions`` of them every new session was refused.
+    """
+    async with Harness.create(max_sessions=1) as harness:
+        harness.plugin.blocked_session_ids.add("blocked")
+        config = DuplexSessionConfig(model="fake-model", instructions="blocked")
+        open_task = asyncio.create_task(harness.open("sid-cancelled", config))
+        await asyncio.wait_for(harness.plugin.runtime_config_started.wait(), timeout=1.0)
+
+        open_task.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await open_task
+
+        assert harness.manager.get("sid-cancelled") is None
+        assert harness.manager.active_count() == 0
+        assert harness.result_sink.empty(), "a cancelled open answers nobody"
+        assert (await harness.open("sid-replacement")).ok is True
+
+
+async def test_an_open_cancelled_after_admission_releases_the_runner_and_its_stage_reservation() -> None:
+    """The worst landing spot for the cancel: the runner is registered and Stage0 is reserved.
+
+    Both have to be undone, or the runner stays in ``runners`` (one slot gone)
+    and the Stage0 request stays reserved in the orchestrator.
+    """
+    async with Harness.create(max_sessions=1) as harness:
+        sink = _BlockingResultSink()
+        original_sink = harness.manager._result_sink
+        harness.manager._result_sink = sink
+        open_task = asyncio.create_task(
+            harness.manager.handle(
+                OpenDuplexSessionMessage(
+                    control_id="open-cancelled",
+                    session_id="sid-cancelled",
+                    session_config=DuplexSessionConfig(model="fake-model"),
+                )
+            )
+        )
+        await asyncio.wait_for(sink.entered.wait(), timeout=1.0)
+        assert "sid-cancelled" in harness.manager.runners
+        assert [context.request_id for context in harness.stage_port.ensure_calls] == [
+            stage0_request_id("sid-cancelled")
+        ]
+
+        open_task.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await open_task
+        harness.manager._result_sink = original_sink
+
+        assert "sid-cancelled" not in harness.manager.runners
+        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
+        assert (await harness.open("sid-replacement")).ok is True
+
+
+async def test_an_open_cancelled_again_during_its_rollback_leaves_the_stage_cleanup_to_the_reaper() -> None:
+    """A second cancellation landing in the rollback's awaits must not orphan the Stage0 request.
+
+    The runner, the reservations and the request index are already gone by
+    then, so nothing would ever clean the orchestrator's request state up.
+    The rollback records the ids as a pending request cleanup before it
+    awaits anything, and the reaper's retry finishes the job.
+    """
+    async with Harness.create(max_sessions=1) as harness:
+        sink = _BlockingResultSink()
+        original_sink = harness.manager._result_sink
+        harness.manager._result_sink = sink
+        open_task = asyncio.create_task(
+            harness.manager.handle(
+                OpenDuplexSessionMessage(
+                    control_id="open-cancelled-twice",
+                    session_id="sid-cancelled",
+                    session_config=DuplexSessionConfig(model="fake-model"),
+                )
+            )
+        )
+        await asyncio.wait_for(sink.entered.wait(), timeout=1.0)
+        runner = harness.manager.runners["sid-cancelled"]
+        shutdown_started = asyncio.Event()
+
+        async def parked_shutdown() -> None:
+            shutdown_started.set()
+            await asyncio.Event().wait()
+
+        runner.shutdown = parked_shutdown  # type: ignore[method-assign]
+
+        open_task.cancel()
+        await asyncio.wait_for(shutdown_started.wait(), timeout=1.0)
+        open_task.cancel()  # lands inside the rollback, before the stage cleanup
+        with pytest.raises(asyncio.CancelledError):
+            await open_task
+        harness.manager._result_sink = original_sink
+
+        assert "sid-cancelled" not in harness.manager.runners
+        assert harness.manager.active_count() == 0
+        assert harness.stage_port.cleanup_calls == [], "the cleanup await was never reached"
+        key = ("sid-cancelled", 0)
+        assert key in harness.manager._pending_request_cleanups
+
+        await harness.manager.reap_expired()
+        assert harness.stage_port.cleanup_calls == [([stage0_request_id("sid-cancelled")], False)]
+        assert key not in harness.manager._pending_request_cleanups
+        assert (await harness.open("sid-replacement")).ok is True
+
+
+# --------------------------------------------------------------------------- #
+# Detach fenced on the lease generation                                       #
+# --------------------------------------------------------------------------- #
+
+
+async def test_detach_is_refused_for_a_lease_the_caller_no_longer_holds() -> None:
+    """A connection giving up its lease must not start the grace for the lease a later resume owns.
+
+    Resume is a CAS on the lease generation; detach is now fenced the same
+    way, so a stale caller (its resume was superseded by another connection's)
+    is refused instead of detaching the winner.
+    """
+    async with Harness.create() as harness:
+        await harness.open("sid-fence")
+        session = harness.session("sid-fence")
+        assert (await harness.resume("sid-fence", expected_lease_generation=0)).ok is True
+        assert session.lease_generation == 1
+
+        stale = await harness.touch("sid-fence", "detach", expected_lease_generation=0)
+        assert stale.ok is False
+        assert stale.error_code == "session_resume_conflict"
+        assert session.lease.detached_at is None, "the current lease is untouched"
+
+        current = await harness.touch("sid-fence", "detach", expected_lease_generation=1)
+        assert current.ok is True
+        assert session.lease.detached_at is not None
+
+        # An unfenced detach keeps its meaning: whatever generation is current.
+        assert (await harness.resume("sid-fence", expected_lease_generation=1)).ok is True
+        assert (await harness.touch("sid-fence", "detach")).ok is True
+        assert session.lease.detached_at is not None
+
+
 # --------------------------------------------------------------------------- #
 # Cancellation waits                                                          #
 # --------------------------------------------------------------------------- #
@@ -1685,3 +1858,20 @@ async def test_cancel_append_tasks_absorbs_a_task_that_outlives_the_wait() -> No
     assert await tasks.cancel_append_tasks(timeout_s=0.05) is True
     assert tasks.append_tail is None
     await asyncio.wait_for(asyncio.gather(task, return_exceptions=True), timeout=5.0)
+
+
+async def test_a_resume_replayed_under_its_control_id_reports_the_generation_it_produced() -> None:
+    """The client side of an abandoned resume: it lost the answer and asks again with the same id."""
+    async with Harness.create() as harness:
+        await harness.open("sid-replay")
+        session = harness.session("sid-replay")
+        first = await harness.resume("sid-replay", expected_lease_generation=0, control_id="rpc-a")
+        assert first.ok is True and first.lease_generation == 1
+
+        replay = await harness.resume("sid-replay", expected_lease_generation=0, control_id="rpc-a")
+        assert replay.ok is True
+        assert replay.lease_generation == 1
+        assert session.lease_generation == 1, "a replay does not resume again"
+
+        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"

tests/engine/duplex/test_session_runner.py

@@ -73,6 +73,9 @@ class RecordingStagePort(DuplexStagePort):
         self.cleanups: list[tuple[list[str], bool]] = []
         self.aborts: list[list[str]] = []
         self.fail_submit: Exception | None = None
+        #: When set, ``submit`` parks on it after signalling ``submit_started``.
+        self.submit_gate: asyncio.Event | None = None
+        self.submit_started = asyncio.Event()
 
     @property
     def stage_count(self) -> int:
@@ -87,6 +90,9 @@ class RecordingStagePort(DuplexStagePort):
     async def submit(self, submission: DuplexStageSubmission) -> DuplexStageSubmissionResult:
         if self.fail_submit is not None:
             raise self.fail_submit
+        if self.submit_gate is not None:
+            self.submit_started.set()
+            await self.submit_gate.wait()
         self.submissions.append(submission)
         return DuplexStageSubmissionResult(
             request_id=submission.context.request_id,
@@ -163,7 +169,7 @@ class Harness:
 
     def deliver(
         self,
-        output: object,
+        output: SimpleNamespace,
         *,
         stage_id: int = 1,
         segment_finished: bool = False,
@@ -186,7 +192,7 @@ class Harness:
             raise AssertionError("stage output is missing request_id")
         return self.runner.on_stage_output(stage_id, output, metrics, request_id=request_id, context=context)
 
-    async def deliver_and_settle(self, output: object, **kwargs: Any) -> list[DuplexEvent]:
+    async def deliver_and_settle(self, output: SimpleNamespace, **kwargs: Any) -> list[DuplexEvent]:
         self.deliver(output, **kwargs)
         return await self.settle()
 
@@ -1624,3 +1630,133 @@ async def test_server_vad_speech_stopped_still_commits_a_turn_mode_session() ->
         assert len(_final_submissions(h)) == 1, "the detector's stop commits the turn and starts the response"
     finally:
         await close_harness(h)
+
+
+# --------------------------------------------------------------------------- #
+# Cancelled append and its precreated response (#7636 Issue 2)                #
+# --------------------------------------------------------------------------- #
+
+
+async def _commit_with_pending_response(h: Harness) -> str:
+    """Commit a turn whose final append parks in the stage port; return the precreated response id."""
+    await h.run(append_audio())
+    h.port.submit_gate = asyncio.Event()
+    h.submit(commands.Commit(create_response=True))
+    await asyncio.wait_for(h.port.submit_started.wait(), timeout=2.0)
+    # The parked append keeps the runner "busy": settle only until the mailbox is quiet.
+    events = await h.settle(timeout_s=0.3)
+    assert "response.created" in types(events)
+    response_id = h.session.active_response_id
+    assert response_id is not None
+    assert h.runner.tasks.append_tail is not None and not h.runner.tasks.append_tail.done()
+    return response_id
+
+
+@pytest.mark.asyncio
+async def test_an_append_cancelled_from_outside_fails_the_response_it_precreated() -> None:
+    """A cancelled append gives back the response it precreated, as its docstring promises.
+
+    The ``CancelledError`` branch used to roll back only the PCM reservation,
+    so ``active_response_id`` stayed pinned to a response nobody would ever
+    finish: the client had seen ``response.created`` and waited forever.
+    """
+    h = await open_harness(auto_response=False)
+    try:
+        response_id = await _commit_with_pending_response(h)
+
+        pending = h.runner.tasks.append_tail
+        assert pending is not None
+        pending.cancel()
+        events = await h.settle()
+
+        done = find(events, "response.done")
+        assert done.response_id == response_id
+        assert done.status == "failed"
+        assert done.response["status_details"]["reason"] == "append_cancelled"
+        assert h.session.active_response_id is None
+        assert h.session.state == DuplexSessionState.OPEN
+
+        # The cancelled append counts as a failed one: the chain behind it is
+        # refused with the documented error, not left hanging.
+        h.port.submit_gate = None
+        await h.run(append_audio())
+        events = await h.run(commands.Commit(create_response=True))
+        assert "response.created" not in types(events)
+        assert find(events, "error").code == "commit_aborted"
+        assert h.session.active_response_id is None
+    finally:
+        await close_harness(h)
+
+
+@pytest.mark.asyncio
+async def test_response_cancel_of_a_pending_append_ends_the_response_once_as_cancelled() -> None:
+    """When the runner cancels the append itself, it owns the response: one terminal, status cancelled.
+
+    The append must not also fail the response on its way out, or the client
+    would get a failed ``response.done`` for a response it cancelled.
+    """
+    h = await open_harness(auto_response=False)
+    try:
+        response_id = await _commit_with_pending_response(h)
+
+        events = await h.run(commands.CancelResponse())
+
+        terminals = [event for event in events if event.type == "response.done"]
+        assert [event.response_id for event in terminals] == [response_id]
+        assert terminals[0].status == "cancelled"
+        assert h.session.active_response_id is None
+        assert h.session.state == DuplexSessionState.OPEN
+    finally:
+        await close_harness(h)
+
+
+@pytest.mark.asyncio
+async def test_closing_the_session_with_a_pending_append_emits_no_failed_response() -> None:
+    """A close ends the active response itself; the append it cancels stays quiet."""
+    h = await open_harness(auto_response=False)
+    try:
+        await _commit_with_pending_response(h)
+
+        events = await h.run(commands.CloseSession(reason="client_close"))
+        events += await h.settle()
+
+        assert [event.type for event in events if event.type == "response.done"] == []
+        assert "session.closed" in types(events)
+        assert h.session.active_response_id is None
+    finally:
+        await close_harness(h)
+
+
+@pytest.mark.asyncio
+async def test_input_clear_racing_a_queued_append_drops_it_as_the_clients_doing() -> None:
+    """An ``input_audio_buffer.clear`` that lands while an append waits its turn is not a runtime failure.
+
+    The clear deactivates the queued append's reservation; when its turn
+    comes it gives everything back and stops, without an error and without
+    failing the session, and its reason is the client's clear rather than
+    ``runtime_append_failed``.
+    """
+    h = await open_harness()
+    try:
+        h.port.submit_gate = asyncio.Event()
+        h.submit(append_audio())
+        await asyncio.wait_for(h.port.submit_started.wait(), timeout=2.0)
+        h.submit(append_audio())
+        await h.settle(timeout_s=0.3)
+        assert len(h.runner.tasks.append_tasks) == 2, "the second append is queued behind the parked one"
+        queued = h.runner.tasks.append_tail
+        assert queued is not None and not queued.done()
+
+        h.submit(commands.ClearInput())
+        events = await h.settle(timeout_s=0.3)
+        assert types(events) == ["input_audio_buffer.cleared"]
+
+        h.port.submit_gate.set()
+        events = await h.settle()
+        assert queued.done() and queued.result() is False, "the cleared append never ran"
+        assert len(h.port.submissions) == 1, "only the append that was already in flight reached the stage"
+        assert [event.type for event in events if event.type in {"error", "response.done", "session.closed"}] == []
+        assert h.session.pending_input_bytes == 0
+        assert h.session.state == DuplexSessionState.OPEN
+    finally:
+        await close_harness(h)

tests/engine/duplex/test_session_runner_nemotron_voicechat.py

@@ -0,0 +1,167 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Essential Nemotron regressions through the real session runner and plugin."""
+
+import asyncio
+from types import SimpleNamespace
+
+import pytest
+import pytest_asyncio
+import torch
+
+from tests.engine.duplex.test_session_runner import (
+    SESSION_ID,
+    Harness,
+    RecordingStagePort,
+    _fake_encode_audio,
+    append_audio,
+    types,
+)
+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.messages import OpenDuplexSessionMessage
+from vllm_omni.engine.duplex.plugin import load_duplex_plugin
+from vllm_omni.engine.duplex.session.manager import DuplexSessionManager
+from vllm_omni.model_executor.models.nemotron_voicechat.duplex import plugin as plugin_module
+from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_thinker import (
+    NemotronVoiceChatThinkerForConditionalGeneration as Thinker,
+)
+from vllm_omni.model_executor.models.nemotron_voicechat.pipeline import NEMOTRON_VOICECHAT_PIPELINE
+from vllm_omni.protocol.realtime.items import normalize_conversation_item
+from vllm_omni.utils.mm_outputs import build_mm_cpu
+from vllm_omni.worker.output.payload_build import build_omni_mm_payload
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest_asyncio.fixture
+async def harness(monkeypatch):
+    monkeypatch.setenv("NEMOTRON_VOICECHAT_LLM_PATH", "test-tokenizer")
+    special_ids = dict(
+        nvc_text_bos_id=0,
+        nvc_text_eos_id=1,
+        nvc_text_pad_id=12,
+        nvc_function_sotc_id=20,
+        nvc_function_eotc_id=21,
+        nvc_function_eotr_id=22,
+        nvc_tokenizer_ref="test-tokenizer",
+    )
+    tokenizer = SimpleNamespace(
+        encode=lambda text, **kwargs: [42],
+        decode=lambda ids, **kwargs: '[{"name":"lookup","arguments":{"x":1}}]',
+    )
+    monkeypatch.setattr(plugin_module, "_load_tokenizer_runtime", lambda _: (special_ids, tokenizer))
+    plugin = load_duplex_plugin(NEMOTRON_VOICECHAT_PIPELINE.duplex_plugin, _fake_encode_audio)
+    port = RecordingStagePort(stage_count=3)
+    output: asyncio.Queue = asyncio.Queue()
+    results: asyncio.Queue = asyncio.Queue()
+    manager = DuplexSessionManager(
+        plugin=plugin,
+        stage_port=port,
+        output_sink=output,
+        result_sink=results,
+        runtime_config=DuplexSessionRuntimeConfig(),
+        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))
+        result = await asyncio.wait_for(results.get(), timeout=2)
+        assert result.ok, result
+        h = Harness(manager, port, output, results, manager.runners[SESSION_ID])
+        await h.settle()
+        yield h
+    finally:
+        await manager.shutdown()
+
+
+async def deliver(h, token_id=12, **metadata):
+    output = SimpleNamespace(
+        request_id=h.stage0_request_id(),
+        finished=False,
+        multimodal_output={},
+        outputs=[SimpleNamespace(text="", token_ids=[token_id], multimodal_output={})],
+    )
+    return await h.deliver_and_settle(
+        output,
+        stage_id=0,
+        segment_finished=True,
+        segment_token_ids=[token_id],
+        segment_output_metadata=metadata,
+    )
+
+
+@pytest.mark.asyncio
+async def test_frame_buffering_prefill_and_resumable_append(harness):
+    h = harness
+    await h.run(append_audio(samples=1000))
+    assert h.port.submissions == []
+    await h.run(append_audio(samples=280))
+    await h.run(append_audio(samples=1280))
+    first, second = h.port.submissions
+    assert first.prompt["prompt_token_ids"] == [0, 42, 1, 12]
+    assert second.prompt["prompt_token_ids"] == [12]
+    assert first.prompt["model_intermediate_buffer"]["duplex"]["source_input_seq"] == 1
+    assert second.prompt["model_intermediate_buffer"]["duplex"]["source_input_seq"] == 2
+    assert not first.already_submitted and second.already_submitted
+    assert first.context.request_id == second.context.request_id == h.stage0_request_id()
+    assert "error" not in types(h.events)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("request_ids", [("active",), ("other", "active"), ("active", "other")])
+async def test_tool_batches_reach_worker_and_drain_before_ninth_result(harness, request_ids):
+    h = harness
+    await h.run(append_audio(samples=1280))
+    worker_state = {"function_response_generation": 0, "forced_function_tokens": []}
+    thinker = Thinker.__new__(Thinker)
+    torch.nn.Module.__init__(thinker)
+    thinker._sessions = {"active": worker_state, "other": {"function_response_generation": 2}}
+    thinker._use_function_head, thinker._dtype = True, torch.float32
+    thinker.function_head = torch.nn.Identity()
+    key = "nvc_function_response_consumed_generation"
+    assert key in thinker.omni_client_multimodal_output_keys
+    infos = [{"request_id": rid} for rid in request_ids]
+    hidden = torch.zeros((len(infos) + 1, 2))
+    for generation in range(1, 10):
+        for token_id in (20, 7, 21):
+            events = await deliver(h, nvc_function_token=[token_id])
+        done = next(event for event in events if event.type == "response.function_call_arguments.done")
+        item = {"type": "function_call_output", "call_id": done.call_id, "output": "20"}
+        events = await h.run(commands.CreateItem(item=normalize_conversation_item(item)))
+        assert "error" not in types(events)
+        runtime = h.session.runtime_config
+        assert runtime["nvc_function_response_generation"] == generation
+        assert runtime["nvc_function_response_batches"][-1] == {"generation": generation, "token_ids": [42]}
+        if generation < 8:
+            continue
+        if generation == 8:
+            with pytest.raises(plugin_module.NemotronVoiceChatClientRuntimeConfigError, match="faster than"):
+                h.runner.plugin.runtime_config_for_function_output(h.session.config, runtime, item)
+        # Copy into model state exactly once, even if the same runtime snapshot is replayed.
+        Thinker._sync_forced_function_response(worker_state, runtime)
+        Thinker._sync_forced_function_response(worker_state, runtime)
+        assert worker_state["forced_function_tokens"] == [42] * generation
+        assert worker_state["function_response_generation"] == generation
+        output = thinker.make_omni_output(hidden, model_intermediate_buffer=infos)
+        for idx, rid in enumerate(request_ids):
+            payload = build_omni_mm_payload(
+                combined_multimodal_outputs=None,
+                mm_cpu=build_mm_cpu(output.multimodal_outputs),
+                rid=rid,
+                idx=idx,
+                start=0 if idx == 0 else 2,
+                end=2 if idx == 0 else 3,
+                audio_sparse_output=False,
+                sparse_mm_index={},
+                hidden_seq_len=len(hidden),
+                scheduled_seq_len=len(hidden),
+            )
+            assert int(payload[key]) == (generation if rid == "active" else 2)
+            if rid == "active":
+                await deliver(h, **payload)
+        assert h.session.runtime_config["nvc_function_response_batches"] == []
+        assert h.session.runtime_config["nvc_function_response_generation"] == generation
+    assert "error" not in types(h.events)

tests/engine/test_stage_engine_args.py

@@ -20,6 +20,7 @@ from vllm.config import ParallelConfig as VllmParallelConfig
 from vllm.config import ProfilerConfig as VllmProfilerConfig
 from vllm.config import SchedulerConfig as VllmSchedulerConfig
 from vllm.engine.arg_utils import EngineArgs
+from vllm.model_executor.models.utils import extract_layer_index
 
 from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
 from vllm_omni.config.omni_config import (
@@ -404,7 +405,9 @@ def test_mammoth_fp8_kv_deploy_projects_only_ar_stage(monkeypatch):
     )
 
     assert deploy.stages[0].engine_extras["kv_cache_dtype"] == "fp8_e4m3"
+    assert deploy.stages[0].engine_extras["kv_cache_dtype_skip_layers"] == ["0"]
     assert "kv_cache_dtype" not in deploy.stages[1].engine_extras
+    assert "kv_cache_dtype_skip_layers" not in deploy.stages[1].engine_extras
 
     legacy_args = [build_legacy_engine_args_dict(stage, "test-model") for stage in legacy_stages]
     typed_args = [
@@ -420,14 +423,22 @@ def test_mammoth_fp8_kv_deploy_projects_only_ar_stage(monkeypatch):
 
     assert ar_stage.stage_pipeline_config.execution_type == StageExecutionType.LLM_AR
     assert legacy_args[0]["kv_cache_dtype"] == "fp8_e4m3"
+    assert legacy_args[0]["kv_cache_dtype_skip_layers"] == ["0"]
     assert ar_stage.cache_config.cache_dtype == "fp8_e4m3"
+    assert ar_stage.cache_config.kv_cache_dtype_skip_layers == ["0"]
     assert "cache_dtype" in ar_stage.cache_config._omni_explicit_fields
     assert typed_args[0]["kv_cache_dtype"] == "fp8_e4m3"
+    assert typed_args[0]["kv_cache_dtype_skip_layers"] == ["0"]
+    # vLLM matches the list against the index it parses from each attention prefix.
+    assert str(extract_layer_index("ar.language_model.layers.0.self_attn.attn")) == "0"
 
     assert "kv_cache_dtype" not in legacy_args[1]
+    assert "kv_cache_dtype_skip_layers" not in legacy_args[1]
     assert dit_stage.cache_config.cache_dtype == "auto"
+    assert dit_stage.cache_config.kv_cache_dtype_skip_layers == []
     assert "cache_dtype" not in dit_stage.cache_config._omni_explicit_fields
     assert "kv_cache_dtype" not in typed_args[1]
+    assert "kv_cache_dtype_skip_layers" not in typed_args[1]
 
 
 def test_typed_llm_projection_rejects_explicit_fields_owned_by_another_boundary():

tests/entrypoints/duplex/test_duplex_serving.py

@@ -8,6 +8,7 @@ from __future__ import annotations
 import asyncio
 import base64
 import json
+from collections.abc import Awaitable, Callable
 from contextlib import suppress
 from typing import Any
 
@@ -131,7 +132,15 @@ class FakeOmni:
         self.handles: dict[str, FakeHandle] = {}
         self.resumed: list[tuple[str, int]] = []
         self.detached: list[str] = []
+        #: Detaches the engine refused because the caller's lease generation was stale.
+        self.detach_refused: list[tuple[str, int]] = []
         self.open_error: DuplexSessionError | None = None
+        #: When set, ``resume_session`` parks on it once the engine applied the
+        #: resume: the lease generation is bumped, the result is still in flight.
+        self.resume_gate: asyncio.Event | None = None
+        self.resume_started = asyncio.Event()
+        #: Settlements of resumes whose caller was cancelled mid-RPC (DuplexOmni's contract).
+        self.compensations: list[asyncio.Task[None]] = []
 
     async def open_session(self, config: Any) -> FakeHandle:
         self.opened.append(dict(config))
@@ -146,13 +155,46 @@ class FakeOmni:
     def get_session(self, session_id: str) -> FakeHandle | None:
         return self.handles.get(session_id)
 
-    async def resume_session(self, session_id: str, *, expected_lease_generation: int) -> FakeHandle:
+    async def resume_session(
+        self,
+        session_id: str,
+        *,
+        expected_lease_generation: int,
+        on_abandoned: Callable[[int], Awaitable[None]] | None = None,
+    ) -> FakeHandle:
         self.resumed.append((session_id, expected_lease_generation))
         handle = self.handles[session_id]
-        handle.lease_generation = expected_lease_generation + 1
+        lease_generation = expected_lease_generation + 1
+        handle.lease_generation = lease_generation
+        if self.resume_gate is not None:
+            self.resume_started.set()
+            gate = self.resume_gate
+            try:
+                await gate.wait()
+            except asyncio.CancelledError:
+                # DuplexOmni's contract: the engine applied the resume; once
+                # its answer arrives the landed generation is settled off the
+                # cancelled task, through the caller's callback if it gave one.
+                async def settle() -> None:
+                    await gate.wait()
+                    if on_abandoned is not None:
+                        await on_abandoned(lease_generation)
+                    else:
+                        with suppress(DuplexSessionError):
+                            await self.detach_session(session_id, expected_lease_generation=lease_generation)
+
+                self.compensations.append(asyncio.create_task(settle()))
+                raise
         return handle
 
-    async def detach_session(self, session_id: str) -> None:
+    async def detach_session(self, session_id: str, *, expected_lease_generation: int | None = None) -> None:
+        handle = self.handles[session_id]
+        if expected_lease_generation is not None and expected_lease_generation != handle.lease_generation:
+            # The engine's fence: a detach of a lease the caller no longer holds is refused.
+            self.detach_refused.append((session_id, expected_lease_generation))
+            raise DuplexSessionError(
+                "duplex lease generation mismatch", code="session_resume_conflict", session_id=session_id
+            )
         self.detached.append(session_id)
 
 
@@ -607,3 +649,380 @@ async def test_a_resume_that_fails_to_activate_does_not_detach_the_live_attachme
                 pending.cancel()
                 with suppress(asyncio.CancelledError, Exception):
                     await pending
+
+
+# --------------------------------------------------------------------------- #
+# Cancelled resume and the engine disconnect grace (#7636 Issue 3)            #
+# --------------------------------------------------------------------------- #
+
+
+@pytest.mark.asyncio
+async def test_a_resume_cancelled_during_activation_puts_the_engine_lease_back_into_grace() -> None:
+    """``_resume`` resumes the engine lease first; a cancel mid-activation must detach it again.
+
+    The rollback used to catch ``Exception`` only. A handler task cancelled
+    while sending ``session.resumed`` or a replay entry left the engine with
+    ``detached_at=None`` and the outer handler with no attachment to clean up:
+    the reaper reclaimed nothing after the grace, and with one slot the next
+    open was refused with ``resource_exhausted``.
+    """
+    omni = FakeOmni()
+    handler = _handler(omni)
+    ws, handle, task = await _open(handler, omni)
+    token = ws.sent[0]["resume_token"]
+    ws.disconnect()
+    await asyncio.wait_for(task, timeout=2.0)
+    assert omni.detached == [handle.session_id]
+    registry = handler._attachment_registry
+
+    activating = asyncio.Event()
+
+    async def parked_resume(*args, **kwargs):
+        activating.set()
+        await asyncio.Event().wait()
+
+    original_resume = registry.resume
+    registry.resume = parked_resume  # type: ignore[method-assign]
+    try:
+        ws2, task2 = await _resume(handler, handle.session_id, token)
+        await asyncio.wait_for(activating.wait(), timeout=2.0)
+        assert omni.resumed == [(handle.session_id, 0)], "the engine lease was resumed before activation"
+
+        task2.cancel()
+        with suppress(asyncio.CancelledError):
+            await task2
+    finally:
+        registry.resume = original_resume  # type: ignore[method-assign]
+
+    # The lease is detached again, so the disconnect grace runs for it.
+    assert omni.detached == [handle.session_id, handle.session_id]
+    assert not await registry.has_attachment(handle.session_id)
+    assert handle.close_reasons == []
+    # And the session is still resumable with the token the cancelled attempt presented.
+    ws3, task3 = await _resume(handler, handle.session_id, token)
+    resumed = await ws3.wait_for("session.resumed")
+    assert resumed["session_id"] == handle.session_id
+    ws3.disconnect()
+    await asyncio.wait_for(task3, timeout=2.0)
+
+
+@pytest.mark.asyncio
+async def test_a_resume_cancelled_after_activation_detaches_the_attachment_it_made() -> None:
+    """A cancel after the registry activated the new attachment must undo that attachment too.
+
+    Otherwise the session stays attached to a socket nobody serves, with the
+    engine lease resumed: no grace, no expiry, until the idle TTL.
+    """
+    omni = FakeOmni()
+    handler = _handler(omni)
+    ws, handle, task = await _open(handler, omni)
+    token = ws.sent[0]["resume_token"]
+    registry = handler._attachment_registry
+
+    # The takeover notifies the replaced socket after activation; park there.
+    notifying = asyncio.Event()
+
+    async def parked_send(payload: dict[str, Any]) -> None:
+        notifying.set()
+        await asyncio.Event().wait()
+
+    ws.send_json = parked_send  # type: ignore[method-assign]
+    ws2, task2 = await _resume(handler, handle.session_id, token)
+    await asyncio.wait_for(notifying.wait(), timeout=2.0)
+    assert await registry.is_current_attachment(handle.session_id, 2), "the new socket is attached"
+
+    task2.cancel()
+    with suppress(asyncio.CancelledError):
+        await task2
+
+    assert not await registry.has_attachment(handle.session_id)
+    assert omni.detached == [handle.session_id]
+    assert handle.close_reasons == []
+
+    task.cancel()
+    with suppress(asyncio.CancelledError, Exception):
+        await task
+
+
+@pytest.mark.asyncio
+async def test_an_abandoned_resume_never_detaches_the_lease_a_later_resume_owns() -> None:
+    """Two connections resuming the same session: the loser's rollback must not touch the winner's lease.
+
+    Interleaving: A has activated and is notifying the socket it replaced; B,
+    holding A's rotated token, has completed its engine resume but its RPC
+    result has not reached the handler yet, so nothing is attached for B.
+    Cancelling A detaches A's registry generation, which used to make the
+    engine detach unconditional: it landed on the lease B had just resumed,
+    and the reaper expired B after the grace even though it heartbeats. The
+    engine now refuses a detach fenced on a generation it has moved past.
+    """
+    omni = FakeOmni()
+    handler = _handler(omni)
+    ws, handle, task = await _open(handler, omni)
+    token = ws.sent[0]["resume_token"]
+    registry = handler._attachment_registry
+
+    # A takes over and parks while notifying the replaced socket (after activation).
+    notifying = asyncio.Event()
+
+    async def parked_send(payload: dict[str, Any]) -> None:
+        notifying.set()
+        await asyncio.Event().wait()
+
+    ws.send_json = parked_send  # type: ignore[method-assign]
+    ws_a, task_a = await _resume(handler, handle.session_id, token)
+    await asyncio.wait_for(notifying.wait(), timeout=2.0)
+    resumed_a = await ws_a.wait_for("session.resumed")
+    assert handle.lease_generation == 1
+
+    # B resumes with A's rotated token: the engine applies it, the result is still in flight.
+    omni.resume_gate = asyncio.Event()
+    ws_b, task_b = await _resume(handler, handle.session_id, resumed_a["resume_token"])
+    await asyncio.wait_for(omni.resume_started.wait(), timeout=2.0)
+    assert handle.lease_generation == 2
+
+    task_a.cancel()
+    with suppress(asyncio.CancelledError):
+        await task_a
+    assert omni.detached == [], "A's rollback must not detach the lease B resumed"
+    assert omni.detach_refused == [(handle.session_id, 1)]
+    assert not await registry.has_attachment(handle.session_id), "A's attachment is gone"
+
+    omni.resume_gate.set()
+    resumed_b = await ws_b.wait_for("session.resumed")
+    assert resumed_b["attachment_generation"] == 3
+    assert await registry.is_current_attachment(handle.session_id, 3)
+    assert omni.detached == []
+
+    # B's own disconnect detaches the lease B holds.
+    ws_b.disconnect()
+    await asyncio.wait_for(task_b, timeout=2.0)
+    assert omni.detached == [handle.session_id]
+    assert omni.detach_refused == [(handle.session_id, 1)]
+    task.cancel()
+    with suppress(asyncio.CancelledError, Exception):
+        await task
+
+
+@pytest.mark.asyncio
+async def test_an_abandoned_takeover_leaves_the_live_connection_serving_and_owning_the_lease() -> None:
+    """A takeover cancelled mid-RPC must not put the connection it never replaced into disconnect grace.
+
+    A is attached at generation 0. B's engine resume lands generation 1 and B
+    is cancelled before it reaches the registry, so A stays the current
+    attachment. Settling B's generation must hand it to A rather than detach
+    it: A keeps serving, and A's own disconnect later detaches generation 1,
+    which is the lease the engine actually holds.
+    """
+    omni = FakeOmni()
+    handler = _handler(omni)
+    ws, handle, task = await _open(handler, omni)
+    token = ws.sent[0]["resume_token"]
+    registry = handler._attachment_registry
+
+    omni.resume_gate = asyncio.Event()
+    ws_b, task_b = await _resume(handler, handle.session_id, token)
+    await asyncio.wait_for(omni.resume_started.wait(), timeout=2.0)
+    assert handle.lease_generation == 1
+
+    task_b.cancel()
+    with suppress(asyncio.CancelledError):
+        await task_b
+    omni.resume_gate.set()
+    await asyncio.gather(*omni.compensations)
+
+    assert omni.detached == [], "A is still serving: its lease must not enter disconnect grace"
+    assert omni.detach_refused == []
+    assert await registry.is_current_attachment(handle.session_id, 1), "A is still the attachment"
+    handle.deliver(AudioDelta(session_id=handle.session_id, response_id="r1", delta="aGk="))
+    await ws.wait_for("response.output_audio.delta")
+
+    # A now owns generation 1: its disconnect detaches the lease the engine holds.
+    ws.disconnect()
+    await asyncio.wait_for(task, timeout=2.0)
+    assert omni.detached == [handle.session_id]
+    assert omni.detach_refused == []
+
+
+@pytest.mark.asyncio
+async def test_a_send_failure_detaches_the_failed_sockets_own_lease_not_a_pending_resumes() -> None:
+    """The pump's detach is fenced on the lease of the socket whose send failed, captured with it.
+
+    The pump holds the registry's outbound lock while sending to A. B's
+    engine resume lands generation 1 meanwhile and B waits for that lock to
+    activate. When A's send fails, the pump drops A and must detach A's
+    generation 0, which the engine refuses, not B's generation 1, which would
+    put B into disconnect grace the moment it activates.
+    """
+    omni = FakeOmni()
+    handler = _handler(omni)
+    ws, handle, task = await _open(handler, omni)
+    token = ws.sent[0]["resume_token"]
+    registry = handler._attachment_registry
+
+    sending = asyncio.Event()
+    fail_send = asyncio.Event()
+
+    async def parked_failing_send(payload: dict[str, Any]) -> None:
+        sending.set()
+        await fail_send.wait()
+        raise RuntimeError("socket closed")
+
+    ws.send_json = parked_failing_send  # type: ignore[method-assign]
+    handle.deliver(AudioDelta(session_id=handle.session_id, response_id="r1", delta="aGk="))
+    await asyncio.wait_for(sending.wait(), timeout=2.0)
+
+    # B: the engine resume lands; activation waits for the outbound lock the pump holds.
+    ws_b, task_b = await _resume(handler, handle.session_id, token)
+    await asyncio.sleep(0.05)
+    assert omni.resumed == [(handle.session_id, 0)]
+    assert handle.lease_generation == 1
+    assert await registry.is_current_attachment(handle.session_id, 1), "B has not activated yet"
+
+    fail_send.set()
+    resumed_b = await ws_b.wait_for("session.resumed")
+    assert resumed_b["attachment_generation"] == 2
+    assert omni.detached == [], "A's lease (generation 0) is refused, B's lease is untouched"
+    assert omni.detach_refused == [(handle.session_id, 0)]
+    assert await registry.is_current_attachment(handle.session_id, 2)
+
+    # B serves; its own disconnect detaches the lease it resumed.
+    ws_b.disconnect()
+    await asyncio.wait_for(task_b, timeout=2.0)
+    assert omni.detached == [handle.session_id]
+    ws.disconnect()
+    with suppress(asyncio.CancelledError, Exception):
+        await asyncio.wait_for(task, timeout=2.0)
+
+
+@pytest.mark.asyncio
+async def test_an_abandoned_resume_goes_to_the_resume_waiting_to_activate_not_the_socket_it_replaces() -> None:
+    """Three connections: A sending, B resumed and waiting for the outbound lock, C abandoned mid-RPC.
+
+    C's generation 2 must be parked for B's activation. If it were handed to
+    A instead, A's send failure would detach generation 2 (accepted by the
+    engine), and B would then activate on a lease already in disconnect grace
+    and expire despite heartbeating.
+    """
+    omni = FakeOmni()
+    handler = _handler(omni)
+    ws, handle, task = await _open(handler, omni)
+    token = ws.sent[0]["resume_token"]
+    registry = handler._attachment_registry
+
+    sending = asyncio.Event()
+    fail_send = asyncio.Event()
+
+    async def parked_failing_send(payload: dict[str, Any]) -> None:
+        sending.set()
+        await fail_send.wait()
+        raise RuntimeError("socket closed")
+
+    ws.send_json = parked_failing_send  # type: ignore[method-assign]
+    handle.deliver(AudioDelta(session_id=handle.session_id, response_id="r1", delta="aGk="))
+    await asyncio.wait_for(sending.wait(), timeout=2.0)
+
+    # B: engine resume lands generation 1; activation waits for the lock the pump holds.
+    ws_b, task_b = await _resume(handler, handle.session_id, token)
+    await asyncio.sleep(0.05)
+    assert handle.lease_generation == 1
+
+    # C: same unrotated token, engine resume lands generation 2, cancelled mid-RPC.
+    omni.resume_gate = asyncio.Event()
+    ws_c, task_c = await _resume(handler, handle.session_id, token)
+    await asyncio.wait_for(omni.resume_started.wait(), timeout=2.0)
+    assert handle.lease_generation == 2
+    task_c.cancel()
+    with suppress(asyncio.CancelledError):
+        await task_c
+    omni.resume_gate.set()
+    await asyncio.gather(*omni.compensations)
+    assert omni.detached == [] and omni.detach_refused == [], "generation 2 is parked for B"
+
+    # A's send fails: the pump drops A and detaches A's own generation 0, refused.
+    fail_send.set()
+    resumed_b = await ws_b.wait_for("session.resumed")
+    assert resumed_b["attachment_generation"] == 2
+    assert omni.detach_refused == [(handle.session_id, 0)]
+    assert omni.detached == [], "the lease B serves is untouched"
+    assert await registry.is_current_attachment(handle.session_id, 2)
+
+    # B serves generation 2 (inherited from C): its disconnect detaches that lease.
+    ws_b.disconnect()
+    await asyncio.wait_for(task_b, timeout=2.0)
+    assert omni.detached == [handle.session_id]
+    assert omni.detach_refused == [(handle.session_id, 0)]
+    ws.disconnect()
+    with suppress(asyncio.CancelledError, Exception):
+        await asyncio.wait_for(task, timeout=2.0)
+
+
+@pytest.mark.asyncio
+async def test_a_resume_cancelled_during_replay_settles_the_newer_lease_it_was_handed() -> None:
+    """Activation rollback must give back the lease the provisional attachment actually owned.
+
+    A resumed generation 1, delivered ``session.resumed`` and is blocked
+    replaying events. B, with A's rotated token, resumed generation 2 and was
+    cancelled mid-RPC; its settlement handed 2 to A. When A is then cancelled
+    during replay, the registry rolls A back; detaching A's own generation 1
+    would be refused and generation 2 would never enter disconnect grace.
+    """
+    omni = FakeOmni()
+    handler = _handler(omni)
+    ws, handle, task = await _open(handler, omni)
+    token = ws.sent[0]["resume_token"]
+    registry = handler._attachment_registry
+    ws.disconnect()
+    await asyncio.wait_for(task, timeout=2.0)
+    assert omni.detached == [handle.session_id]
+    # Journaled while detached: A's activation has something to replay.
+    handle.deliver(AudioDelta(session_id=handle.session_id, response_id="r1", delta="one"))
+    await asyncio.sleep(0.05)
+
+    # A: session.resumed goes out, the replay entry parks.
+    ws_a = FakeWebSocket({"duplex": "1", "resume": "1"})
+    replaying = asyncio.Event()
+    sends = 0
+    original_send = ws_a.send_json
+
+    async def send_parking_on_replay(payload: dict[str, Any]) -> None:
+        nonlocal sends
+        sends += 1
+        if sends == 2:
+            replaying.set()
+            await asyncio.Event().wait()
+        await original_send(payload)
+
+    ws_a.send_json = send_parking_on_replay  # type: ignore[method-assign]
+    task_a = asyncio.create_task(handler.handle_realtime_session(ws_a))
+    ws_a.feed(
+        {
+            "type": "session.resume",
+            "session_id": handle.session_id,
+            "resume_token": token,
+            "last_received_server_event_seq": 0,
+        }
+    )
+    await asyncio.wait_for(replaying.wait(), timeout=2.0)
+    resumed_a = await ws_a.wait_for("session.resumed")
+    assert handle.lease_generation == 1
+
+    # B: A's rotated token, engine resume lands generation 2, cancelled mid-RPC.
+    omni.resume_gate = asyncio.Event()
+    ws_b, task_b = await _resume(handler, handle.session_id, resumed_a["resume_token"])
+    await asyncio.wait_for(omni.resume_started.wait(), timeout=2.0)
+    assert handle.lease_generation == 2
+    task_b.cancel()
+    with suppress(asyncio.CancelledError):
+        await task_b
+    omni.resume_gate.set()
+    await asyncio.gather(*omni.compensations)
+    assert omni.detached == [handle.session_id], "generation 2 went to A's pending activation"
+
+    # A is cancelled during replay: its rollback settles generation 2, the lease A owned.
+    task_a.cancel()
+    with suppress(asyncio.CancelledError):
+        await task_a
+    assert not await registry.has_attachment(handle.session_id)
+    assert omni.detached == [handle.session_id, handle.session_id]
+    assert omni.detach_refused == [], "generation 1 was never detached: the engine would refuse it"

tests/entrypoints/openai/test_duplex_session_attachment.py

@@ -552,3 +552,90 @@ async def test_registry_resume_cancelled_mid_delivery_rolls_back_like_a_failure(
         close=close,
     )
     assert recovered.attachment_generation == 3
+
+
+@pytest.mark.asyncio
+async def test_release_hands_back_the_lease_generation_the_dropped_connection_was_serving(mocker) -> None:
+    """The lease generation travels with the attachment, so a detach is fenced on the right lease."""
+    registry = DuplexSessionAttachmentRegistry(replay_ttl_s=60.0, replay_max_bytes_per_session=4096)
+    send = mocker.AsyncMock()
+    created = await registry.create("sid-lease", send=send, close=mocker.AsyncMock(), lease_generation=0)
+
+    # An abandoned takeover hands its generation to the connection still attached.
+    assert await registry.settle_lease_generation("sid-lease", 1) is None
+    assert await registry.settle_lease_generation("sid-lease", 0) is None, "generations only move forward"
+
+    released = await registry.release_attachment("sid-lease", attachment_generation=created.attachment_generation)
+    assert released is not None
+    assert (released.attachment_generation, released.lease_generation) == (1, 1)
+    assert await registry.release_attachment("sid-lease", attachment_generation=None) is None
+    assert await registry.settle_lease_generation("sid-lease", 2) == 2, "nobody is attached: the caller detaches it"
+
+    resumed = await registry.resume(
+        "sid-lease",
+        resume_token=created.resume_token.plaintext,
+        last_received_server_event_seq=0,
+        send=send,
+        close=mocker.AsyncMock(),
+        lease_generation=3,
+    )
+    released = await registry.release_attachment("sid-lease", attachment_generation=resumed.attachment_generation)
+    assert released is not None and released.lease_generation == 3
+
+
+@pytest.mark.asyncio
+async def test_an_orphaned_generation_waits_for_the_resume_that_is_about_to_activate(mocker) -> None:
+    """While a resume is pending, an orphan goes to its activation, not to the socket it will replace."""
+    registry = DuplexSessionAttachmentRegistry(replay_ttl_s=60.0, replay_max_bytes_per_session=4096)
+    send = mocker.AsyncMock()
+    created = await registry.create("sid-pending", send=send, close=mocker.AsyncMock(), lease_generation=0)
+
+    await registry.begin_resume("sid-pending")
+    assert await registry.settle_lease_generation("sid-pending", 2) is None, "parked for the pending activation"
+    released = await registry.release_attachment("sid-pending", attachment_generation=created.attachment_generation)
+    assert released is not None and released.lease_generation == 0, "the replaced socket keeps its own lease"
+
+    resumed = await registry.resume(
+        "sid-pending",
+        resume_token=created.resume_token.plaintext,
+        last_received_server_event_seq=0,
+        send=send,
+        close=mocker.AsyncMock(),
+        lease_generation=1,
+    )
+    assert await registry.end_resume("sid-pending") is None, "the activation absorbed the orphan"
+    released = await registry.release_attachment("sid-pending", attachment_generation=resumed.attachment_generation)
+    assert released is not None and released.lease_generation == 2
+
+    # A pending resume that never activates leaves the orphan to whoever ends the claim.
+    await registry.begin_resume("sid-pending")
+    assert await registry.settle_lease_generation("sid-pending", 3) is None
+    assert await registry.end_resume("sid-pending") == 3
+
+
+@pytest.mark.asyncio
+async def test_a_rolled_back_activation_orphans_the_lease_it_was_handed(mocker) -> None:
+    """A provisional attachment may own a newer generation than its own resume produced."""
+    registry = DuplexSessionAttachmentRegistry(replay_ttl_s=60.0, replay_max_bytes_per_session=4096)
+    created = await registry.create("sid-rollback", send=mocker.AsyncMock(), close=mocker.AsyncMock())
+    await registry.detach("sid-rollback", attachment_generation=created.attachment_generation)
+
+    async def failing_activation(payload):
+        raise RuntimeError("socket closed during replay")
+
+    await registry.begin_resume("sid-rollback")
+    with pytest.raises(RuntimeError):
+        await registry.resume(
+            "sid-rollback",
+            resume_token=created.resume_token.plaintext,
+            last_received_server_event_seq=0,
+            send=failing_activation,
+            close=mocker.AsyncMock(),
+            activation_payload_factory=lambda token, generation: {"type": "session.resumed"},
+            lease_generation=1,
+        )
+    # Meanwhile a resume abandoned mid-RPC handed generation 2 to the provisional attachment.
+    # (Order does not matter: both land in the orphan slot while the claim is open.)
+    assert await registry.settle_lease_generation("sid-rollback", 2) is None
+    assert await registry.settle_lease_generation("sid-rollback", 1) is None
+    assert await registry.end_resume("sid-rollback") == 2, "the newest lease is the one to detach"

tests/entrypoints/openai_api/test_api_server_guards.py

@@ -821,6 +821,46 @@ def test_speech_without_handler_preserves_not_found_http_error() -> None:
     assert exc_info.value.detail == "The model does not support Speech API"
 
 
+@pytest.mark.parametrize(
+    ("field", "value", "detail"),
+    [
+        ("modalities", [123], "modalities must be a list of strings"),
+        ("logprobs", "yes", "logprobs must be a boolean"),
+    ],
+)
+def test_chat_completion_raw_body_guards_reject_lax_types(field, value, detail) -> None:
+    """The HTTP boundary rejects values that upstream Pydantic would coerce."""
+    app = FastAPI()
+    app.state.openai_serving_chat = None
+    app.state.serving_tokenization = None
+    app.add_api_route("/v1/chat/completions", api_server.create_chat_completion, methods=["POST"])
+    client = TestClient(app)
+
+    payload = {
+        "model": "demo-model",
+        "messages": [{"role": "user", "content": "hello"}],
+        "stream": False,
+        field: value,
+    }
+    response = client.post("/v1/chat/completions", json=payload)
+
+    assert response.status_code == 400
+    assert detail in response.json()["detail"]
+
+
+@pytest.mark.parametrize(
+    "raw_body",
+    [
+        {"modalities": None},
+        {"logprobs": None},
+        {"modalities": None, "logprobs": None},
+    ],
+)
+def test_chat_completion_raw_body_guards_allow_null_defaults(raw_body) -> None:
+    """Explicit JSON null keeps the upstream request model's default behavior."""
+    api_server._validate_chat_completion_raw_body(raw_body)
+
+
 @pytest.mark.asyncio
 async def test_multi_api_rejects_runtime_voice_upload() -> None:
     app = FastAPI()

tests/entrypoints/test_async_omni.py

@@ -7,6 +7,7 @@ from types import SimpleNamespace
 
 import anyio
 import pytest
+from vllm.engine.protocol import StreamingInput
 from vllm.lora.request import LoRARequest
 from vllm.sampling_params import RequestOutputKind, SamplingParams
 
@@ -67,6 +68,7 @@ def get_async_omni_instance(fake_add_request=_noop, fake_abort_request=_noop) ->
     omni._paused = False
     omni.engine = SimpleNamespace(
         num_stages=1,
+        stage_configs=[],
         add_request_async=fake_add_request,
         abort_async=fake_abort_request,
     )
@@ -473,6 +475,61 @@ def test_generate_accepts_request_after_repeated_cancellations():
     asyncio.run(run_test())
 
 
+@pytest.mark.cpu
+@pytest.mark.parametrize("has_chunk", [False, True])
+@pytest.mark.parametrize("input_fails", [False, True])
+def test_streaming_input_terminal_submit_failure_reaches_generate(mocker, has_chunk, input_fails):
+    async def run_test():
+        submissions: list[tuple[str, bool]] = []
+
+        async def add_request_async(*, resumable, **kwargs):
+            del kwargs
+            submissions.append(("add", resumable))
+            if not resumable:
+                raise RuntimeError("terminal streaming submit failed")
+
+        async def add_streaming_update_async(*, resumable, **kwargs):
+            del kwargs
+            submissions.append(("update", resumable))
+            assert not resumable
+            raise RuntimeError("terminal streaming submit failed")
+
+        omni = get_async_omni_instance(fake_add_request=add_request_async)
+        omni.engine.add_streaming_update_async = add_streaming_update_async
+        omni.engine.stage_clients = []
+        omni.engine.model_config = object()
+        del omni._process_orchestrator_results
+        mocker.patch(
+            "vllm_omni.entrypoints.async_omni.extract_prompt_components",
+            return_value=(None, None, None),
+        )
+        error_metadata = mocker.spy(async_omni_mod, "client_error_metadata")
+
+        async def input_stream():
+            if has_chunk:
+                yield StreamingInput(prompt={"prompt": "chunk"})
+            if input_fails:
+                raise RuntimeError("streaming input failed")
+
+        async def consume():
+            async for _ in omni.generate(
+                prompt=input_stream(),
+                request_id="stream-submit-error",
+                sampling_params_list=[SamplingParams(output_kind=RequestOutputKind.DELTA)],
+                output_modalities=["text"],
+            ):
+                pass
+
+        expected_error = "streaming input failed" if input_fails else "terminal streaming submit failed"
+        with pytest.raises(RuntimeError, match=expected_error):
+            await asyncio.wait_for(consume(), timeout=1.0)
+        assert error_metadata.call_count == 1
+        assert submissions == ([("add", True), ("update", False)] if has_chunk else [("add", False)])
+        assert omni.request_states == {}
+
+    asyncio.run(run_test())
+
+
 @pytest.mark.cpu
 def test_generate_cancel_bounds_cleanup_abort(monkeypatch):
     """Cancelled generate() must not wait forever on a wedged orchestrator abort."""

tests/entrypoints/test_duplex_omni.py

@@ -43,7 +43,20 @@ class FakeDuplexEngine(FakeAsyncOmniEngine):
         self.opened: list[tuple[str, DuplexSessionConfig]] = []
         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).
+        self.resume_control_ids: list[str | None] = []
+        #: Per session: the control id that produced the current lease generation, and that generation.
+        self.lease: dict[str, tuple[str | None, int]] = {}
+        #: How many resume RPCs, after applying the resume, answer with a
+        #: timeout: the engine landed it but the waiter never saw the result.
+        self.resume_timeouts_remaining = 0
         self.touched: list[tuple[str, str]] = []
+        #: The lease-generation fence of every detach, in order (``None`` = unfenced).
+        self.detach_generations: list[int | None] = []
+        #: When set, ``resume_session_async`` parks on it after signalling ``resume_started``:
+        #: the engine has applied the resume, its result has not reached the caller.
+        self.resume_gate: asyncio.Event | None = None
+        self.resume_started = asyncio.Event()
         self.commands: list[tuple[str, commands.DuplexCommand]] = []
         self.open_error: DuplexSessionError | None = None
         self.close_error: DuplexSessionError | None = None
@@ -71,12 +84,33 @@ class FakeDuplexEngine(FakeAsyncOmniEngine):
         self.emit(session_id, SessionClosed(reason=reason))
         return self._result("close", session_id)
 
-    async def resume_session_async(self, session_id, *, expected_lease_generation, timeout=None):
+    async def resume_session_async(self, session_id, *, expected_lease_generation, control_id=None, timeout=None):
         self.resumed.append((session_id, expected_lease_generation))
-        return self._result("resume", session_id, lease_generation=expected_lease_generation + 1)
+        self.resume_control_ids.append(control_id)
+        resumed_by, generation = self.lease.get(session_id, (None, 0))
+        if control_id is not None and control_id == resumed_by:
+            # The engine's replay answer: the generation this resume produced.
+            # Its answer can be lost to a timeout exactly like the first one.
+            if self.resume_timeouts_remaining > 0:
+                self.resume_timeouts_remaining -= 1
+                raise DuplexSessionError("duplex resume timed out", code="timeout", retryable=True)
+            return self._result("resume", session_id, lease_generation=generation)
+        if expected_lease_generation != generation:
+            raise DuplexSessionError("duplex lease generation mismatch", code="session_resume_conflict")
+        generation += 1
+        self.lease[session_id] = (control_id, generation)
+        if self.resume_gate is not None:
+            self.resume_started.set()
+            await self.resume_gate.wait()
+        if self.resume_timeouts_remaining > 0:
+            self.resume_timeouts_remaining -= 1
+            raise DuplexSessionError("duplex resume timed out", code="timeout", retryable=True)
+        return self._result("resume", session_id, lease_generation=generation)
 
-    async def touch_session_async(self, session_id, *, activity, timeout=None):
+    async def touch_session_async(self, session_id, *, activity, expected_lease_generation=None, timeout=None):
         self.touched.append((session_id, activity))
+        if activity == "detach":
+            self.detach_generations.append(expected_lease_generation)
         return self._result("touch", session_id)
 
     async def submit_command_async(self, session_id, command):
@@ -324,3 +358,141 @@ async def test_shutdown_closes_every_handle(monkeypatch) -> None:
     assert first.close_reason == "shutdown"
     assert omni.sessions == {}
     assert engine.shutdown_called
+
+
+@pytest.mark.asyncio
+async def test_detach_session_forwards_the_lease_generation_fence(monkeypatch) -> None:
+    omni, engine = _make_omni(monkeypatch)
+    try:
+        handle = await omni.open_session()
+        await omni.detach_session(handle.session_id, expected_lease_generation=3)
+        await omni.detach_session(handle.session_id)
+        assert engine.detach_generations == [3, None]
+    finally:
+        omni.shutdown()
+
+
+@pytest.mark.asyncio
+async def test_a_resume_cancelled_mid_rpc_is_adopted_and_detached_once_it_lands(monkeypatch) -> None:
+    """The resume RPC runs in an executor: cancelling its waiter does not stop the engine.
+
+    Without compensation the engine ends up with the lease resumed (no
+    disconnect grace) for a connection that is gone, and the handle keeps the
+    old generation, so a later reconnect resumes against a stale lease. The
+    outcome has to be observed off the cancelled task: adopt the generation
+    the resume produced and detach exactly that lease.
+    """
+    omni, engine = _make_omni(monkeypatch)
+    try:
+        handle = await omni.open_session()
+        engine.resume_gate = asyncio.Event()
+        resume = asyncio.create_task(omni.resume_session(handle.session_id, expected_lease_generation=0))
+        await asyncio.wait_for(engine.resume_started.wait(), timeout=2.0)
+
+        resume.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await resume
+        # Nothing to give back yet: the engine has not answered.
+        assert handle.lease_generation == 0
+        assert engine.detach_generations == []
+
+        engine.resume_gate.set()
+        await asyncio.gather(*list(omni._resume_compensations))
+        assert handle.lease_generation == 1, "the generation the resume produced is adopted"
+        assert engine.touched == [(handle.session_id, "detach")]
+        assert engine.detach_generations == [1], "the detach is fenced on that generation"
+    finally:
+        omni.shutdown()
+
+
+@pytest.mark.asyncio
+async def test_a_resume_cancelled_then_timed_out_is_settled_by_replaying_it(monkeypatch) -> None:
+    """A timeout after the cancel leaves the outcome unknown: the resume is replayed under its id.
+
+    The engine may have applied the resume after the waiter gave up, and the
+    RPC router drops the late answer. Replaying the same control id makes the
+    engine report the generation that resume produced (or land it now), so
+    the handle adopts the real generation and the lease is detached again.
+    """
+    omni, engine = _make_omni(monkeypatch)
+    try:
+        handle = await omni.open_session()
+        engine.resume_gate = asyncio.Event()
+        engine.resume_timeouts_remaining = 1
+        resume = asyncio.create_task(omni.resume_session(handle.session_id, expected_lease_generation=0))
+        await asyncio.wait_for(engine.resume_started.wait(), timeout=2.0)
+
+        resume.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await resume
+        engine.resume_gate.set()
+        await asyncio.gather(*list(omni._resume_compensations))
+
+        assert len(engine.resume_control_ids) == 2, "the resume was replayed once"
+        assert engine.resume_control_ids[0] == engine.resume_control_ids[1] is not None
+        assert engine.lease[handle.session_id][1] == 1, "the replay did not resume a second time"
+        assert handle.lease_generation == 1
+        assert engine.detach_generations == [1]
+    finally:
+        omni.shutdown()
+
+
+@pytest.mark.asyncio
+async def test_an_abandoned_resume_is_settled_through_the_callers_callback(monkeypatch) -> None:
+    """A caller that owns attachments decides what the landed generation means; DuplexOmni does not detach."""
+    omni, engine = _make_omni(monkeypatch)
+    try:
+        handle = await omni.open_session()
+        settled: list[int] = []
+
+        async def on_abandoned(lease_generation: int) -> None:
+            settled.append(lease_generation)
+
+        engine.resume_gate = asyncio.Event()
+        resume = asyncio.create_task(
+            omni.resume_session(handle.session_id, expected_lease_generation=0, on_abandoned=on_abandoned)
+        )
+        await asyncio.wait_for(engine.resume_started.wait(), timeout=2.0)
+        resume.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await resume
+        engine.resume_gate.set()
+        await asyncio.gather(*list(omni._resume_compensations))
+
+        assert settled == [1]
+        assert handle.lease_generation == 1
+        assert engine.detach_generations == [], "the callback owns the decision"
+    finally:
+        omni.shutdown()
+
+
+@pytest.mark.asyncio
+async def test_an_abandoned_resume_keeps_replaying_across_repeated_timeouts(monkeypatch) -> None:
+    """A replay can time out too; giving up then would reopen the lost-answer window.
+
+    The engine applied the resume, the abandoned RPC and the first replay both
+    lose their answers to timeouts, the second replay is answered: the handle
+    adopts the generation and the lease is settled, with the lease bumped once.
+    """
+    omni, engine = _make_omni(monkeypatch)
+    try:
+        omni._resume_replay_backoff_s = (0.01, 0.02)
+        handle = await omni.open_session()
+        engine.resume_gate = asyncio.Event()
+        engine.resume_timeouts_remaining = 2
+        resume = asyncio.create_task(omni.resume_session(handle.session_id, expected_lease_generation=0))
+        await asyncio.wait_for(engine.resume_started.wait(), timeout=2.0)
+
+        resume.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await resume
+        engine.resume_gate.set()
+        await asyncio.wait_for(asyncio.gather(*list(omni._resume_compensations)), timeout=5.0)
+
+        assert len(engine.resume_control_ids) == 3, "the abandoned RPC, a replay that timed out, the replay answered"
+        assert len(set(engine.resume_control_ids)) == 1
+        assert engine.lease[handle.session_id][1] == 1, "the lease was bumped once"
+        assert handle.lease_generation == 1
+        assert engine.detach_generations == [1]
+    finally:
+        omni.shutdown()

tests/helpers/fixtures/clean.py

@@ -1,7 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
-"""Opt-in cleanup fixtures for GPU memory, pipeline registry, and speaker cache."""
+"""Cleanup fixtures for process-global test state.
+
+Torch dtype/device restoration is autouse. GPU memory, the pipeline registry,
+and the speaker cache stay opt-in.
+"""
 
 from __future__ import annotations
 
@@ -10,6 +14,26 @@ import pytest
 from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
 
 
+@pytest.fixture(autouse=True)
+def restore_torch_default_dtype_and_device():
+    """Restore process-global dtype and device after every test.
+
+    ``get_default_device()`` reports ``cpu`` both when no default device was
+    set and when ``cpu`` was set explicitly. Writing that value back enables
+    default-device dispatch, so snapshot the active context instead.
+    """
+    import torch
+
+    default_dtype = torch.get_default_dtype()
+    device_context = getattr(torch._GLOBAL_DEVICE_CONTEXT, "device_context", None)
+    default_device = None if device_context is None else torch.get_default_device()
+    try:
+        yield
+    finally:
+        torch.set_default_dtype(default_dtype)
+        torch.set_default_device(default_device)
+
+
 @pytest.fixture
 def clean_gpu_memory_between_tests():
     """Opt-in GPU pre/post hooks for a test (no environment-variable gate).

tests/helpers/tiny_model.py

@@ -44,14 +44,16 @@ def build_tiny_from_configs(
             config or transform) that don't have vendored configs.
         configs_dir: Path to a directory containing vendored config files
             (model_index.json and per-component config.json).
-        transform: Maps a component (subfolder) name to a function that
-            takes that component's raw config dict (fetched from model_id)
-            and returns a shrunk one, applied only when that component has
-            no vendored config in configs_dir. Lets each pipeline encode
-            exactly the shrink it needs (layer counts, and any component-
-            internal width that isn't read by another component), including
-            keeping sibling config fields in sync (e.g. layer_types matching
-            num_hidden_layers).
+        transform: Maps a component (subfolder) name or None to a function that
+            transforms a config dict. For subfolder keys, the function takes
+            that component's raw config dict (fetched from model_id) and returns
+            a shrunk one, applied only when that component has no vendored config
+            in configs_dir. For None key, the function transforms the pipeline's
+            model_index.json config dict. Lets each pipeline encode exactly the
+            transforms it needs (layer counts, config-level fields, and any
+            component-internal width that isn't read by another component),
+            including keeping sibling config fields in sync (e.g. layer_types
+            matching num_hidden_layers).
 
     Returns:
         Path to the saved tiny model directory with safetensors weights.
@@ -72,7 +74,7 @@ def build_tiny_from_configs(
         subfolder = kwargs.get("subfolder")
         load_kwargs = {"subfolder": subfolder} if subfolder else {}
         config = cls.load_config(pretrained_model_name_or_path, **load_kwargs)
-        if subfolder in transform:
+        if subfolder is not None and subfolder in transform:
             config = transform[subfolder](dict(config))
         return cls.from_config(config)
 
@@ -86,7 +88,7 @@ def build_tiny_from_configs(
         subfolder = kwargs.get("subfolder")
         load_kwargs = {"subfolder": subfolder} if subfolder else {}
         config_dict, _ = PretrainedConfig.get_config_dict(pretrained_model_name_or_path, **load_kwargs)
-        if subfolder in transform:
+        if subfolder is not None and subfolder in transform:
             config_dict = transform[subfolder](config_dict)
         config = cls.config_class.from_dict(config_dict)
         return cls(config)
@@ -96,6 +98,8 @@ def build_tiny_from_configs(
     # can vendor just one component without also needing its own model_index.json.
     vendored_index = config_dir is not None and (config_dir / "model_index.json").exists()
     config_dict = DiffusionPipeline.load_config(config_dir if vendored_index else model_id)
+    if None in transform:
+        config_dict = transform[None](dict(config_dict))
     pipeline_cls = _get_pipeline_class(DiffusionPipeline, config=config_dict)
 
     init_dict, _, _ = pipeline_cls.extract_init_dict(config_dict)

tests/model_executor/models/mammoth_moda2/test_mammoth_moda2_ar_kv_layers.py

@@ -0,0 +1,145 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""The AR stage must register attention layers for its own language model only.
+
+The Qwen-VL parent ``__init__`` builds a stock language model before MammothModa2
+replaces it. Its attention layers stay in ``static_forward_context`` unless they are
+removed, and the KV cache is then sized and allocated for both models. The parent
+``__init__`` and the replacement are stubbed here so only the registry is exercised.
+"""
+
+import gc
+import weakref
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+from vllm.model_executor.models.utils import StageMissingLayer
+
+import vllm_omni.model_executor.models.mammoth_moda2.mammoth_moda2 as mm2
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+NUM_LAYERS = 3
+VISION_LAYER = "ar.visual.blocks.0.attn"
+AR_CLASSES = [
+    (mm2.MammothModa2ARForConditionalGeneration, mm2.Qwen2_5_VLForConditionalGeneration),
+    (mm2.MammothModa2Qwen3ARForConditionalGeneration, mm2.Qwen3VLForConditionalGeneration),
+]
+AR_IDS = ["preview_qwen2_5_vl", "dev_qwen3_vl"]
+
+
+class _FakeAttention(nn.Module):
+    def __init__(self, registry: dict, name: str):
+        super().__init__()
+        registry[name] = self
+
+
+class _FakeLM(nn.Module):
+    def __init__(self, registry: dict, prefix: str):
+        super().__init__()
+        self.layers = nn.ModuleList(
+            _FakeAttention(registry, f"{prefix}.layers.{i}.self_attn.attn") for i in range(NUM_LAYERS)
+        )
+        self.make_empty_intermediate_tensors = lambda *args, **kwargs: None
+
+
+class _FakeCompiledLM(_FakeLM):
+    cleanups = 0
+
+    def cleanup(self):
+        type(self).cleanups += 1
+
+
+def _build(
+    monkeypatch,
+    ar_class,
+    parent,
+    *,
+    encoder_only=False,
+    parent_lm=_FakeLM,
+    parent_registers=True,
+    pp_rank=(1, True),
+    stray_parent_layer=False,
+):
+    registry: dict = {}
+    built = {}
+
+    def parent_init(self, *, vllm_config, prefix=""):
+        nn.Module.__init__(self)
+        self.visual = _FakeAttention(registry, VISION_LAYER)
+        stock = nn.Module()
+        stock.model = parent_lm(registry if parent_registers else {}, f"{prefix}.language_model.model")
+        built["parent_lm"] = weakref.ref(stock)
+        if stray_parent_layer:
+            # A parent-prefixed layer that is not part of the module being dropped.
+            registry[f"{prefix}.language_model.model.layers.9.self_attn.attn"] = nn.Module()
+        self.language_model = StageMissingLayer("language_model", stock) if encoder_only else stock
+
+    # The model module only imports get_current_vllm_config with this fix; keep the
+    # reverted source testable so the registry assertion is what fails there.
+    monkeypatch.setattr(
+        mm2,
+        "get_current_vllm_config",
+        lambda: SimpleNamespace(compilation_config=SimpleNamespace(static_forward_context=registry)),
+        raising=False,
+    )
+    monkeypatch.setattr(parent, "__init__", parent_init)
+    monkeypatch.setattr(mm2, "init_vllm_registered_model", lambda *, prefix, **kwargs: _FakeLM(registry, prefix))
+    monkeypatch.setattr(mm2, "MammothModa2Qwen3ForCausalLM", lambda *, vllm_config, prefix: _FakeLM(registry, prefix))
+    monkeypatch.setattr(mm2, "TorchCompileWithNoGuardsWrapper", _FakeCompiledLM, raising=False)
+    monkeypatch.setattr(torch.accelerator, "empty_cache", lambda: None)
+    world_size, is_first_rank = pp_rank
+    monkeypatch.setattr(
+        mm2, "get_pp_group", lambda: SimpleNamespace(world_size=world_size, is_first_rank=is_first_rank)
+    )
+
+    hf_config = SimpleNamespace(llm_config=SimpleNamespace(text_config=SimpleNamespace()))
+    vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(hf_config=hf_config, architectures=["MammothModa2Model"])
+    )
+    vllm_config.with_hf_config = lambda hf, architectures=None: vllm_config
+    model = ar_class(vllm_config=vllm_config, prefix="ar")
+    return model, registry, built
+
+
+@pytest.mark.parametrize("encoder_only", [False, True], ids=["default", "mm_encoder_only"])
+@pytest.mark.parametrize(("ar_class", "parent"), AR_CLASSES, ids=AR_IDS)
+def test_ar_stage_registers_only_the_replacement_language_model(monkeypatch, ar_class, parent, encoder_only):
+    model, registry, built = _build(monkeypatch, ar_class, parent, encoder_only=encoder_only)
+
+    assert set(registry) == {f"ar.language_model.layers.{i}.self_attn.attn" for i in range(NUM_LAYERS)} | {VISION_LAYER}
+    gc.collect()
+    assert built["parent_lm"]() is None
+
+
+@pytest.mark.parametrize(("ar_class", "parent"), AR_CLASSES, ids=AR_IDS)
+def test_compiled_parent_language_model_releases_its_hook(monkeypatch, ar_class, parent):
+    _FakeCompiledLM.cleanups = 0
+    _build(monkeypatch, ar_class, parent, parent_lm=_FakeCompiledLM)
+    assert _FakeCompiledLM.cleanups == 1
+
+
+@pytest.mark.parametrize(("ar_class", "parent"), AR_CLASSES, ids=AR_IDS)
+def test_unexpected_parent_registry_fails_loudly(monkeypatch, ar_class, parent):
+    with pytest.raises(RuntimeError, match="Unexpected attention registry"):
+        _build(monkeypatch, ar_class, parent, parent_registers=False)
+
+
+@pytest.mark.parametrize(("ar_class", "parent"), AR_CLASSES, ids=AR_IDS)
+def test_later_pipeline_rank_without_parent_layers_starts(monkeypatch, ar_class, parent):
+    _, registry, _ = _build(monkeypatch, ar_class, parent, parent_registers=False, pp_rank=(2, False))
+    assert set(registry) == {f"ar.language_model.layers.{i}.self_attn.attn" for i in range(NUM_LAYERS)} | {VISION_LAYER}
+
+
+@pytest.mark.parametrize(("ar_class", "parent"), AR_CLASSES, ids=AR_IDS)
+def test_first_pipeline_rank_without_parent_layers_fails(monkeypatch, ar_class, parent):
+    with pytest.raises(RuntimeError, match="Unexpected attention registry"):
+        _build(monkeypatch, ar_class, parent, parent_registers=False, pp_rank=(2, True))
+
+
+@pytest.mark.parametrize(("ar_class", "parent"), AR_CLASSES, ids=AR_IDS)
+def test_later_pipeline_rank_still_rejects_a_leftover_parent_layer(monkeypatch, ar_class, parent):
+    with pytest.raises(RuntimeError, match="Unexpected attention registry"):
+        _build(monkeypatch, ar_class, parent, parent_registers=False, pp_rank=(2, False), stray_parent_layer=True)

tests/model_executor/models/ming_image/test_config_and_pipeline.py

@@ -5,20 +5,27 @@ from types import SimpleNamespace
 
 import pytest
 import torch
+from vllm.multimodal.processing import ProcessorInputs
 
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.models.ming_image.pipeline import (
     MingImageDiffusionPipeline,
     _validate_variant_config,
 )
+from vllm_omni.diffusion.models.z_image.pipeline_z_image import ZImagePipeline
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.model_executor.models.ming_flash_omni.ming_flash_omni_thinker import (
     MingFlashOmniThinkerForConditionalGeneration,
+    MingFlashOmniThinkerMultiModalProcessor,
 )
 from vllm_omni.model_executor.models.ming_image import checkpoint
-from vllm_omni.model_executor.models.ming_image.model import MingImageForConditionalGeneration
+from vllm_omni.model_executor.models.ming_image.model import (
+    MingImageForConditionalGeneration,
+    MingImageMultiModalProcessor,
+)
 from vllm_omni.model_executor.models.ming_image.pipeline import MING_IMAGE_PIPELINE
 from vllm_omni.model_executor.stage_input_processors.ming_image import thinker2image
 from vllm_omni.model_extras import (
@@ -26,6 +33,7 @@ from vllm_omni.model_extras import (
     should_init_extra_args_for_non_diffusion_stages,
 )
 from vllm_omni.transformers_utils.configs.ming_flash_omni import BailingMM2Config
+from vllm_omni.transformers_utils.processors.ming import MingImageProcessor
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -119,6 +127,21 @@ def test_pipeline_rejects_multiple_outputs_per_prompt():
         pipeline.forward(request)
 
 
+@pytest.mark.parametrize(
+    ("guidance_scale", "ming_expected", "z_image_expected"),
+    ((0.0, False, False), (1.0, False, True), (1.0001, True, True), (2.0, True, True)),
+)
+def test_ming_cfg_is_only_enabled_above_one(guidance_scale, ming_expected, z_image_expected):
+    # Ming-Image follows the vendor threshold; the shared Z-Image base keeps diffusers semantics.
+    ming_pipeline = MingImageDiffusionPipeline.__new__(MingImageDiffusionPipeline)
+    ming_pipeline._guidance_scale = guidance_scale
+    z_image_pipeline = ZImagePipeline.__new__(ZImagePipeline)
+    z_image_pipeline._guidance_scale = guidance_scale
+
+    assert ming_pipeline.do_classifier_free_guidance is ming_expected
+    assert z_image_pipeline.do_classifier_free_guidance is z_image_expected
+
+
 def _source_output(prefix_len: int = 4):
     prompt_ids = [11] * prefix_len + [157158] + [157157] * 256 + [157159]
     length = len(prompt_ids)
@@ -157,6 +180,78 @@ def test_bridge_rejects_text_negative_condition():
         )
 
 
+def test_image_generation_template_matches_upstream_layout():
+    processor = MingImageProcessor.__new__(MingImageProcessor)
+    processor.tokenizer = SimpleNamespace(eos_token="<eos>")
+
+    text = processor._apply_image_generation_template("draw a cabin")
+    assert text == (
+        "<role>SYSTEM</role>你是一个友好的AI助手。\n\ndetailed thinking off"
+        "<eos><role>HUMAN</role>draw a cabin"
+        "<eos><role>ASSISTANT</role>"
+    )
+    assert processor._apply_image_generation_template(text) == text
+    assert "<role>HUMAN</role><IMAGE>edit this image<eos>" in processor._apply_image_generation_template(
+        "edit this image",
+        has_reference_image=True,
+    )
+
+
+def test_multimodal_processor_applies_generation_template(monkeypatch: pytest.MonkeyPatch):
+    template_calls = []
+
+    def _format(prompt, *, has_reference_image):
+        template_calls.append((prompt, has_reference_image))
+        return "formatted prompt"
+
+    tokenizer = SimpleNamespace(
+        decode=lambda token_ids, **kwargs: "edit this image",
+        encode=lambda text, **kwargs: [101, 102],
+    )
+    processor = MingImageMultiModalProcessor.__new__(MingImageMultiModalProcessor)
+    processor.info = SimpleNamespace(
+        get_tokenizer=lambda: tokenizer,
+        get_hf_processor=lambda: SimpleNamespace(_apply_image_generation_template=_format),
+    )
+    captured = {}
+
+    def _parent_apply(_self, inputs, timing_ctx):
+        captured["inputs"] = inputs
+        return "processed"
+
+    monkeypatch.setattr(MingFlashOmniThinkerMultiModalProcessor, "apply", _parent_apply)
+    inputs = ProcessorInputs(
+        prompt=[11, 12],
+        mm_data_items=None,
+        hf_processor_mm_kwargs={"modalities": ["img2img"]},
+    )
+
+    assert processor.apply(inputs, None) == "processed"
+    assert template_calls == [("edit this image", True)]
+    assert captured["inputs"].prompt == [101, 102]
+    assert captured["inputs"].hf_processor_mm_kwargs["modalities"] == ["image"]
+
+
+@pytest.mark.parametrize(
+    ("pipeline_config", "expected_class_name"),
+    ((MING_IMAGE_PIPELINE, "MingImageDiffusionPipeline"), (None, None)),
+)
+def test_engine_od_config_falls_back_to_registered_pipeline_class(monkeypatch, pipeline_config, expected_class_name):
+    # Ming-Image ships neither a root config.json nor model_index.json for now,
+    # so class-name resolution returns None; serving must still find the extras (e.g. num_layers).
+    monkeypatch.setattr("vllm_omni.diffusion.data.resolve_model_class_name", lambda model, **kwargs: None)
+    engine = object.__new__(AsyncOmniEngine)
+    engine.model = "inclusionAI/Ming-Image-0.1-Design-Layer"
+    engine.pipeline_config = pipeline_config
+    engine._diffusion_od_config_view = None
+
+    od_config = engine.get_diffusion_od_config()
+
+    assert od_config.model_class_name == expected_class_name
+    assert ("num_layers" in get_extra_body_params(od_config.model_class_name)) is (expected_class_name is not None)
+    assert engine.get_diffusion_od_config() is od_config
+
+
 def test_two_stage_topology_and_request_metadata():
     assert MING_IMAGE_PIPELINE.model_type == "ming_image"
     assert [stage.model_stage for stage in MING_IMAGE_PIPELINE.stages] == ["mllm", "dit"]
@@ -168,6 +263,8 @@ def test_two_stage_topology_and_request_metadata():
     assert MING_IMAGE_PIPELINE.hf_architectures == ("MingImageForConditionalGeneration",)
     for class_name in MING_IMAGE_PIPELINE.diffusers_class_aliases + (MING_IMAGE_PIPELINE.diffusers_class_name,):
         assert "num_layers" in get_extra_body_params(class_name)
+        assert "steps" not in get_extra_body_params(class_name)
+        assert "cfg" not in get_extra_body_params(class_name)
         assert should_init_extra_args_for_non_diffusion_stages(class_name)
 
 

tests/model_executor/models/minicpmo_4_5/duplex/test_stage0_runtime_load.py

@@ -0,0 +1,61 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""The MiniCPM-o 4.5 Stage-0 duplex runtime is built with the model, not in the first session."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.minicpmo_4_5 import minicpmo_4_5_omni
+from vllm_omni.model_executor.models.minicpmo_4_5.duplex import compat
+from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import (
+    MiniCPMO45OmniForConditionalGeneration,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _StageModel(torch.nn.Module):
+    def make_empty_intermediate_tensors(self):
+        return None
+
+
+@pytest.mark.parametrize(
+    ("model_stage", "session_mode", "expected_builds"),
+    [
+        ("llm", "duplex", 1),
+        ("llm", "turn", 0),
+        ("tts", "duplex", 0),
+    ],
+)
+def test_init_builds_duplex_runtime_only_for_duplex_thinker(
+    monkeypatch: pytest.MonkeyPatch,
+    model_stage: str,
+    session_mode: str,
+    expected_builds: int,
+) -> None:
+    monkeypatch.setattr(minicpmo_4_5_omni, "init_vllm_registered_model", lambda **kwargs: _StageModel())
+    monkeypatch.setattr(compat, "patch_minicpmo_remote_config", lambda config: None)
+    build_devices: list[torch.device] = []
+    monkeypatch.setattr(
+        MiniCPMO45OmniForConditionalGeneration,
+        "_duplex_data_plane_helper",
+        lambda self: build_devices.append(torch.get_default_device()),
+    )
+    vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(
+            hf_config=SimpleNamespace(),
+            multimodal_config=None,
+            model_stage=model_stage,
+            session_mode=session_mode,
+        )
+    )
+
+    # The loader constructs the model under the target-device context.
+    with torch.device("meta"):
+        MiniCPMO45OmniForConditionalGeneration(vllm_config=vllm_config)
+
+    assert build_devices == [torch.device("cpu")] * expected_builds

tests/model_executor/models/minicpmo_4_5/test_talker_batching.py

@@ -4,12 +4,14 @@
 
 from __future__ import annotations
 
+from contextlib import nullcontext
 from dataclasses import dataclass
 from typing import Any
 
 import pytest
 import torch
 import torch.nn as nn
+from torch.overrides import TorchFunctionMode
 
 from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import (
     MiniCPMO45OmniForConditionalGeneration,
@@ -684,7 +686,8 @@ def test_make_omni_output_packs_the_previous_codec_id() -> None:
 
 
 @pytest.mark.parametrize("device", ["cpu", "cuda"])
-def test_make_omni_output_reads_history_from_runner_cpu_buffer(monkeypatch, device) -> None:
+@pytest.mark.parametrize("default_device", [None, "cpu"])
+def test_make_omni_output_reads_history_from_runner_cpu_buffer(device, default_device) -> None:
     if device == "cuda" and not torch.cuda.is_available():
         pytest.skip("requires CUDA codec output")
     from types import SimpleNamespace
@@ -719,16 +722,20 @@ def test_make_omni_output_reads_history_from_runner_cpu_buffer(monkeypatch, devi
     assert source.device.type == "cpu"
     torch.testing.assert_close(source, torch.tensor([[3]], dtype=torch.long))
 
-    original_tolist = torch.Tensor.tolist
     readback_devices = []
 
-    def cpu_only_tolist(tensor):
-        readback_devices.append(tensor.device.type)
-        assert tensor.device.type == "cpu", "codec history must use the CPU transport delta"
-        return original_tolist(tensor)
+    class CpuOnlyToList(TorchFunctionMode):
+        def __torch_function__(self, func, types, args=(), kwargs=None):
+            if func is torch.Tensor.tolist:
+                tensor = args[0]
+                readback_devices.append(tensor.device.type)
+                assert tensor.device.type == "cpu", "codec history must use the CPU transport delta"
+            return func(*args, **(kwargs or {}))
 
-    with monkeypatch.context() as patch:
-        patch.setattr(torch.Tensor, "tolist", cpu_only_tolist)
+    # A Tensor.tolist monkeypatch is re-entered by default-device dispatch.
+    # Observe dispatch directly so each CPU readback is counted once.
+    device_context = torch.device(default_device) if default_device is not None else nullcontext()
+    with device_context, CpuOnlyToList():
         output = model.make_omni_output(
             torch.ones(1, 2, device=device),
             model_intermediate_buffer=infos,

tests/model_executor/models/nemotron_voicechat/duplex/conftest.py

@@ -1,20 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-
-"""Nemotron VoiceChat duplex is not on the duplex plugin framework yet.
-
-RFC vllm-omni#7181 splits the port: this PR lands the unified framework with
-MiniCPM-o 4.5, and the follow-up PR moves Nemotron VoiceChat onto
-``DuplexModelPlugin``. Until then the modules under test still import the
-pre-framework duplex runtime / serving adapter that this PR removed, so
-collecting them would fail at import time and abort the whole pytest session.
-The tests themselves are unchanged and come back with the port.
-
-Listed per file rather than globbed, so the coverage loss is visible and shrinks
-file by file as the port lands. ``test_input.py`` imports nothing that was
-removed and keeps running.
-"""
-
-#: Both import ``DuplexFence`` from ``engine.duplex.messages``, which no longer
-#: exports it; ``test_data_plane`` also imports ``engine.duplex.runtime``.
-collect_ignore = ["test_data_plane.py", "test_runtime_contract.py"]

tests/model_executor/models/nemotron_voicechat/duplex/test_data_plane.py

@@ -7,9 +7,8 @@ from types import SimpleNamespace
 
 import numpy as np
 import pytest
-from vllm_omni.engine.duplex.runtime import duplex_resource_request_id
 
-from vllm_omni.engine.duplex.messages import DuplexFence
+from vllm_omni.engine.duplex.contracts import DuplexFence, duplex_resource_request_id
 from vllm_omni.model_executor.models.nemotron_voicechat.duplex.data_plane import (
     NemotronVoiceChatDataPlaneContext,
     NemotronVoiceChatDataPlaneSession,

tests/model_executor/models/nemotron_voicechat/duplex/test_input.py

@@ -45,6 +45,15 @@ def test_irregular_browser_packets_emit_exact_ordered_80ms_frame() -> None:
     np.testing.assert_array_equal(decoded, np.arange(1280, dtype=np.float32))
     assert buffer.pending_byte_count == 0
 
+    later = _append(buffer, np.ones(1280), "later")
+    later.commit()
+    later.rollback()
+    reservation.rollback()
+    assert not buffer.has_reserved()
+    retry = _append(buffer, [], "retry")
+    assert retry.payload == reservation.payload
+    assert not buffer.has_pending()
+
 
 def test_append_rejects_invalid_frame_contract() -> None:
     cases = (

tests/model_executor/models/nemotron_voicechat/duplex/test_runtime_contract.py

@@ -1,115 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-
-from __future__ import annotations
-
-import base64
-from types import SimpleNamespace
-
-import numpy as np
-import pytest
-from vllm.sampling_params import SamplingParams
-
-from vllm_omni.engine.duplex.contracts import DuplexInputMode
-from vllm_omni.engine.duplex.messages import DuplexFence
-from vllm_omni.model_executor.models.nemotron_voicechat.duplex.runtime import (
-    NemotronVoiceChatDuplexRuntimeExtension,
-)
-from vllm_omni.model_executor.models.nemotron_voicechat.duplex.serving_adapter import (
-    NemotronVoiceChatServingRuntimeAdapter,
-    _render_tool_response,
-)
-from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_thinker import (
-    NemotronVoiceChatThinkerForConditionalGeneration,
-)
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-def _frame() -> dict[str, object]:
-    raw = np.zeros(1280, dtype=np.float32).tobytes()
-    return {
-        "type": "audio",
-        "audio": base64.b64encode(raw).decode("ascii"),
-        "format": "pcm_f32le",
-        "sample_rate_hz": 16000,
-    }
-
-
-def _plan(extension, *, input_seq: int):
-    return extension.plan_append(
-        request_id="req",
-        fence=DuplexFence("sid", incarnation=2, epoch=3),
-        session_config={},
-        runtime_config={
-            "nvc_prompt_token_ids": [0, 42, 1],
-            "nvc_text_pad_id": 12,
-            "nvc_max_model_len": 8192,
-        },
-        seq=input_seq,
-        turn_seq=input_seq,
-        mode=DuplexInputMode.APPEND_AUDIO_CHUNK,
-        payload=_frame(),
-        final=False,
-        sampling_params=SamplingParams(),
-    )
-
-
-def test_first_append_prefills_prompt_then_each_append_consumes_one_frame() -> None:
-    extension = NemotronVoiceChatDuplexRuntimeExtension()
-
-    params = extension.configure_sampling_params(
-        runtime_config={}, defaults=(SamplingParams(temperature=0.7, max_tokens=4),)
-    )[0]
-    assert params.temperature == 0.0 and params.max_tokens == 1 and params.top_p == 1.0 and params.top_k == 0
-
-    first = _plan(extension, input_seq=1)
-    later = _plan(extension, input_seq=2)
-
-    assert first.prompt["prompt_token_ids"] == [0, 42, 1, 12]
-    assert later.prompt["prompt_token_ids"] == [12]
-    assert first.prompt["model_intermediate_buffer"]["duplex"]["source_input_seq"] == 1
-    assert later.prompt["model_intermediate_buffer"]["duplex"]["source_input_seq"] == 2
-
-
-def test_append_rejects_stage0_context_overflow() -> None:
-    with pytest.raises(ValueError, match="max_model_len"):
-        _plan(NemotronVoiceChatDuplexRuntimeExtension(), input_seq=8190)
-
-
-def test_function_output_becomes_versioned_nvidia_channel_tokens() -> None:
-    encoded: list[str] = []
-    adapter = NemotronVoiceChatServingRuntimeAdapter(lambda *_: None)
-    adapter._tokenizer = SimpleNamespace(
-        encode=lambda text, **_kwargs: encoded.append(text) or [31, 32, 33],
-    )
-
-    first = adapter.runtime_config_for_function_output(
-        {"type": "function_call_output", "call_id": "call-1", "output": '{"result":20}'},
-        {},
-    )
-    second = adapter.runtime_config_for_function_output(
-        {"type": "function_call_output", "call_id": "call-2", "output": "plain text"},
-        first,
-    )
-
-    assert _render_tool_response('{"result":20}') == '<TOOL_RESPONSE>[{"result":20}]</TOOL_RESPONSE>'
-    assert encoded == [
-        '<TOOL_RESPONSE>[{"result":20}]</TOOL_RESPONSE>',
-        '<TOOL_RESPONSE>["plain text"]</TOOL_RESPONSE>',
-    ]
-    assert first["nvc_function_response_generation"] == 1
-    assert second["nvc_function_response_generation"] == 2
-    assert second["nvc_function_response_token_ids"] == [31, 32, 33]
-    assert second["nvc_function_response_call_id"] == "call-2"
-    assert second["nvc_function_response_batches"] == [
-        {"generation": 1, "call_id": "call-1", "token_ids": [31, 32, 33]},
-        {"generation": 2, "call_id": "call-2", "token_ids": [31, 32, 33]},
-    ]
-
-    session = {"function_response_generation": 0, "forced_function_tokens": []}
-    NemotronVoiceChatThinkerForConditionalGeneration._sync_forced_function_response(session, second)
-
-    assert session["function_response_generation"] == 2
-    assert session["forced_function_token"] == 31
-    assert session["forced_function_tokens"] == [31, 32, 33, 31, 32, 33]

tests/model_tests/diffusion/case_filtering.py

@@ -20,7 +20,12 @@ import itertools
 import pytest
 
 from tests.helpers.mark import hardware_marks
-from tests.model_tests.diffusion.config_types import DiffusionAccs, DiffusionModelTestOpts, get_required_device_count
+from tests.model_tests.diffusion.config_types import (
+    DiffusionAccs,
+    DiffusionModelTestOpts,
+    ModelTypeMarker,
+    get_required_device_count,
+)
 from vllm_omni.platforms import current_omni_platform
 
 # These tests are intended to run on L4 GPUs since they test parallelism (i.e., need GPUs),
@@ -32,7 +37,12 @@ from vllm_omni.platforms import current_omni_platform
 MAX_CI_DEVICES = 4
 
 
-def get_test_group_marks(model_name: str, test_group: list[DiffusionAccs] | None, model_marks: list | None) -> list:
+def get_test_group_marks(
+    model_name: str,
+    test_group: list[DiffusionAccs] | None,
+    model_marks: list | None,
+    model_type_marker: ModelTypeMarker,
+) -> list:
     """Build the full set of pytest marks for a test group.
 
     For now, single device groups default to core model, and multi device groups
@@ -40,6 +50,7 @@ def get_test_group_marks(model_name: str, test_group: list[DiffusionAccs] | None
     the max number of GPUs for the CI to prevent accidentally adding configurations that
     would be skipped in the CI."""
     marks = list(model_marks) if model_marks is not None else []
+    marks.append(getattr(pytest.mark, model_type_marker.value))
 
     required_devices = get_required_device_count(test_group)
     if required_devices > MAX_CI_DEVICES:
@@ -83,7 +94,12 @@ def get_model_parametrization(model_name: str, test_info: DiffusionModelTestOpts
             # The divergence check compares generated outputs, which only the offline test does.
             *([] if online else [test_info.check_i2v_t2v_divergence]),
             id=f"{model_name}[{'+'.join(test_group)}]" if test_group else model_name,
-            marks=get_test_group_marks(model_name, test_group, test_info.marks),
+            marks=get_test_group_marks(
+                model_name,
+                test_group,
+                test_info.marks,
+                model_type_marker=test_info.model_type_marker,
+            ),
         )
         for test_group in test_groups
     ]

tests/model_tests/diffusion/config_types.py

@@ -42,6 +42,14 @@ class DiffusionTasks(str, Enum):
     # Text to audio, etc should be added here as needed
 
 
+class ModelTypeMarker(str, Enum):
+    """Model type markers for pytest parametrization."""
+
+    DIFFUSION = "diffusion"
+    OMNI = "omni"
+    TTS = "tts"
+
+
 class DiffusionModelTestOpts(NamedTuple):
     """Configuration for one Diffusion model's tests."""
 
@@ -59,6 +67,10 @@ class DiffusionModelTestOpts(NamedTuple):
     # Actual tasks which controls the tests actually run
     supported_tasks: list[DiffusionTasks]
 
+    # Function module marker applied per-test-case via parametrization.
+    # Set to ModelTypeMarker.OMNI or ModelTypeMarker.TTS for non-diffusion models.
+    model_type_marker: ModelTypeMarker
+
     # Additional acceleration groups to run beyond the base case (no acceleration).
     # The base case is always run for every model in the test settings. None means
     # we only run the base case.

tests/model_tests/diffusion/diff_model_builders.py

@@ -256,6 +256,28 @@ def tiny_flux_kontext_builder() -> str:
     )
 
 
+def _shrink_sd3_transformer(config: dict) -> dict:
+    config["num_layers"] = 2
+    config["attention_head_dim"] = 32
+    config["num_attention_heads"] = 4
+    config["caption_projection_dim"] = 128
+    config["dual_attention_layers"] = [0]
+    return config
+
+
+def tiny_sd3_builder() -> str:
+    return build_tiny_from_configs(
+        "StableDiffusion3Pipeline",
+        "stabilityai/stable-diffusion-3.5-medium",
+        transform={
+            "text_encoder": _shrink_flux_clip_text_encoder,
+            "text_encoder_2": _shrink_flux_clip_text_encoder,
+            "text_encoder_3": _shrink_flux_t5_text_encoder,
+            "transformer": _shrink_sd3_transformer,
+        },
+    )
+
+
 def tiny_flux2_builder() -> str:
     def shrink_text_encoder(config: dict) -> dict:
         config["tie_word_embeddings"] = False
@@ -280,3 +302,69 @@ def tiny_flux2_builder() -> str:
             "transformer": partial(_shrink_dit_rope_config, num_single_layers=2, joint_attention_dim=96),
         },
     )
+
+
+def _shrink_krea2_text_encoder(config: dict) -> dict:
+    tc = config["text_config"]
+    old_head_dim = tc["head_dim"]
+    tc["num_hidden_layers"] = 3
+    tc["hidden_size"] = 32
+    tc["intermediate_size"] = 64
+    tc["num_attention_heads"] = 2
+    tc["num_key_value_heads"] = 2
+    tc["head_dim"] = 16
+    if "rope_parameters" in tc:
+        rp = tc["rope_parameters"]
+        if "mrope_section" in rp:
+            factor = old_head_dim / 16
+            rp["mrope_section"] = [max(1, round(d / factor)) for d in rp["mrope_section"]]
+    vc = config["vision_config"]
+    vc["depth"] = 2
+    vc["intermediate_size"] = 64
+    vc["num_heads"] = 2
+    vc["out_hidden_size"] = 32
+    if "deepstack_visual_indexes" in vc:
+        vc["deepstack_visual_indexes"] = [0, 1]
+    return config
+
+
+def _shrink_krea2_transformer(config: dict) -> dict:
+    config["num_layers"] = 2
+    config["attention_head_dim"] = 32
+    config["num_attention_heads"] = 4
+    config["num_key_value_heads"] = 2
+    config["intermediate_size"] = 128
+    config["timestep_embed_dim"] = 32
+    config["text_hidden_dim"] = 32
+    config["num_text_layers"] = 3
+    config["text_num_attention_heads"] = 2
+    config["text_num_key_value_heads"] = 2
+    config["text_intermediate_size"] = 64
+    config["num_layerwise_text_blocks"] = 1
+    config["num_refiner_text_blocks"] = 1
+    config["axes_dims_rope"] = [8, 12, 12]
+    return config
+
+
+def _shrink_krea2_vae(config: dict) -> dict:
+    config["base_dim"] = 32
+    return config
+
+
+def _patch_krea2_model_index(config: dict) -> dict:
+    config["text_encoder_select_layers"] = [0, 1, 2]
+    return config
+
+
+def tiny_krea2_builder() -> str:
+    """Build a tiny Krea2 model with random weights using diffusers building blocks."""
+    return build_tiny_from_configs(
+        "Krea2Pipeline",
+        "krea/Krea-2-Turbo",
+        transform={
+            None: _patch_krea2_model_index,
+            "text_encoder": _shrink_krea2_text_encoder,
+            "transformer": _shrink_krea2_transformer,
+            "vae": _shrink_krea2_vae,
+        },
+    )

tests/model_tests/diffusion/model_settings.py

@@ -7,6 +7,7 @@ from tests.model_tests.diffusion.config_types import (
     DiffusionAccs,
     DiffusionModelTestOpts,
     DiffusionTasks,
+    ModelTypeMarker,
 )
 
 # This object defines the (tiny) model configurations for common tests.
@@ -37,12 +38,14 @@ DIFFUSION_TEST_SETTINGS = {
         model=real_anima_model,
         builder=tiny_anima_builder,
         supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE],
+        model_type_marker=ModelTypeMarker.DIFFUSION,
         checkpoint_filename=CHECKPOINT_FILENAME,
     ),
     "Flux2KleinPipeline": DiffusionModelTestOpts(
         model="black-forest-labs/FLUX.2-klein-4B",
         builder=diff_model_builders.tiny_flux2_klein_builder,
         supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE, DiffusionTasks.IMAGE_TO_IMAGE],
+        model_type_marker=ModelTypeMarker.DIFFUSION,
         extra_test_groups=[
             [DiffusionAccs.HSDP, DiffusionAccs.TEA_CACHE],
             [DiffusionAccs.SEQUENCE_PARALLEL, DiffusionAccs.CACHE_DIT, DiffusionAccs.LAYERWISE_OFFLOAD],
@@ -53,6 +56,7 @@ DIFFUSION_TEST_SETTINGS = {
         model="Lightricks/LTX-2",
         builder=diff_model_builders.tiny_ltx2_builder,
         supported_tasks=[DiffusionTasks.TEXT_TO_VIDEO, DiffusionTasks.IMAGE_TO_VIDEO],
+        model_type_marker=ModelTypeMarker.DIFFUSION,
         extra_test_groups=[
             [DiffusionAccs.HSDP, DiffusionAccs.CACHE_DIT],
             [DiffusionAccs.SEQUENCE_PARALLEL, DiffusionAccs.CACHE_DIT, DiffusionAccs.LAYERWISE_OFFLOAD],
@@ -63,6 +67,7 @@ DIFFUSION_TEST_SETTINGS = {
         model="Efficient-Large-Model/SANA-Video_2B_480p_diffusers",
         builder=diff_model_builders.tiny_sana_video_builder,
         supported_tasks=[DiffusionTasks.TEXT_TO_VIDEO],
+        model_type_marker=ModelTypeMarker.DIFFUSION,
         extra_test_groups=[
             [DiffusionAccs.TENSOR_PARALLEL],
             [DiffusionAccs.CFG_PARALLEL],
@@ -82,6 +87,7 @@ DIFFUSION_TEST_SETTINGS = {
         model="Efficient-Large-Model/SANA-Video_2B_480p_diffusers",
         builder=diff_model_builders.tiny_sana_video_i2v_builder,
         supported_tasks=[DiffusionTasks.IMAGE_TO_VIDEO],
+        model_type_marker=ModelTypeMarker.DIFFUSION,
         extra_test_groups=[
             [DiffusionAccs.CFG_PARALLEL],
             [DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CFG_PARALLEL],
@@ -95,11 +101,13 @@ DIFFUSION_TEST_SETTINGS = {
         model="Qwen/Qwen-Image",
         builder=diff_model_builders.tiny_qwen_image_builder,
         supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE],
+        model_type_marker=ModelTypeMarker.DIFFUSION,
     ),
     "LongCatImagePipeline": DiffusionModelTestOpts(
         model="meituan-longcat/LongCat-Image",
         builder=diff_model_builders.tiny_longcat_image_builder,
         supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE],
+        model_type_marker=ModelTypeMarker.DIFFUSION,
         extra_test_groups=[
             [DiffusionAccs.TEA_CACHE],
             [DiffusionAccs.SEQUENCE_PARALLEL, DiffusionAccs.CACHE_DIT, DiffusionAccs.LAYERWISE_OFFLOAD],
@@ -110,6 +118,7 @@ DIFFUSION_TEST_SETTINGS = {
         model="meituan-longcat/LongCat-Image-Edit",
         builder=diff_model_builders.tiny_longcat_image_edit_builder,
         supported_tasks=[DiffusionTasks.IMAGE_TO_IMAGE],
+        model_type_marker=ModelTypeMarker.DIFFUSION,
         extra_test_groups=[
             [DiffusionAccs.TEA_CACHE],
             [DiffusionAccs.SEQUENCE_PARALLEL, DiffusionAccs.CACHE_DIT, DiffusionAccs.LAYERWISE_OFFLOAD],
@@ -120,11 +129,13 @@ DIFFUSION_TEST_SETTINGS = {
         model="black-forest-labs/FLUX.1-schnell",
         builder=diff_model_builders.tiny_flux_builder,
         supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE],
+        model_type_marker=ModelTypeMarker.DIFFUSION,
     ),
     "FluxKontextPipeline": DiffusionModelTestOpts(
         model="black-forest-labs/FLUX.1-Kontext-dev",
         builder=diff_model_builders.tiny_flux_kontext_builder,
         supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE, DiffusionTasks.IMAGE_TO_IMAGE],
+        model_type_marker=ModelTypeMarker.DIFFUSION,
         extra_test_groups=[
             [DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CPU_OFFLOAD],
             [DiffusionAccs.CFG_PARALLEL, DiffusionAccs.CPU_OFFLOAD],
@@ -135,15 +146,39 @@ DIFFUSION_TEST_SETTINGS = {
         model="black-forest-labs/FLUX.2-dev",
         builder=diff_model_builders.tiny_flux2_builder,
         supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE],
+        model_type_marker=ModelTypeMarker.DIFFUSION,
+        extra_test_groups=[
+            [DiffusionAccs.CACHE_DIT, DiffusionAccs.LAYERWISE_OFFLOAD],
+            [DiffusionAccs.CFG_PARALLEL, DiffusionAccs.SEQUENCE_PARALLEL, DiffusionAccs.CPU_OFFLOAD],
+            [DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.VAE_PATCH_PARALLEL],
+        ],
     ),
     "QwenImageEditPipeline": DiffusionModelTestOpts(
         model="Qwen/Qwen-Image-Edit",
         builder=diff_model_builders.tiny_qwen_image_edit_builder,
         supported_tasks=[DiffusionTasks.IMAGE_TO_IMAGE],
+        model_type_marker=ModelTypeMarker.DIFFUSION,
     ),
     "QwenImageEditPlusPipeline": DiffusionModelTestOpts(
         model="Qwen/Qwen-Image-Edit-2511",
         builder=diff_model_builders.tiny_qwen_image_edit_plus_builder,
         supported_tasks=[DiffusionTasks.IMAGE_TO_IMAGE],
+        model_type_marker=ModelTypeMarker.DIFFUSION,
+    ),
+    "StableDiffusion3Pipeline": DiffusionModelTestOpts(
+        model="stabilityai/stable-diffusion-3.5-medium",
+        builder=diff_model_builders.tiny_sd3_builder,
+        supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE],
+        model_type_marker=ModelTypeMarker.DIFFUSION,
+        extra_test_groups=[
+            [DiffusionAccs.CACHE_DIT, DiffusionAccs.LAYERWISE_OFFLOAD],
+            [DiffusionAccs.CFG_PARALLEL, DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CPU_OFFLOAD],
+        ],
+    ),
+    "Krea2Pipeline": DiffusionModelTestOpts(
+        model="krea/Krea-2-Turbo",
+        builder=diff_model_builders.tiny_krea2_builder,
+        supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE],
+        model_type_marker=ModelTypeMarker.DIFFUSION,
     ),
 }

tests/model_tests/diffusion/test_alignment.py

@@ -56,12 +56,10 @@ EXCLUDED_MODELS = [
     "MingImageDiffusionPipeline",
     "MingImageLayeredDiffusionPipeline",
     "InternVLAA1Pipeline",
-    "StableDiffusion3Pipeline",
     "HunyuanImage3ForCausalMM",
     "ErnieImagePipeline",
     "NextStep11Pipeline",
     "FluxDMD2Pipeline",
-    "Krea2Pipeline",
     "QwenImageDMD2Pipeline",
     "OmniGen2Pipeline",
     "HeliosPipeline",

tests/model_tests/diffusion/test_common_offline.py

@@ -21,8 +21,8 @@ from tests.model_tests.diffusion.task_runners import (
     run_and_validate_text_to_video_request,
 )
 
-# NOTE : Hardware marks are added dynamically based on test requirements
-pytestmark = [pytest.mark.diffusion, pytest.mark.xdist]
+# NOTE: Hardware and model type marks are added dynamically based on test requirements and model type
+pytestmark = [pytest.mark.xdist]
 
 
 @pytest.mark.parametrize(

tests/model_tests/diffusion/test_common_online.py

@@ -31,8 +31,8 @@ from tests.model_tests.diffusion.task_runners import (
 )
 from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata
 
-# NOTE : Hardware marks are added dynamically based on test requirements
-pytestmark = [pytest.mark.diffusion, pytest.mark.xdist]
+# NOTE: Hardware and model type marks are added dynamically based on test requirements and model type
+pytestmark = [pytest.mark.xdist]
 
 
 @pytest.fixture(autouse=True)

tests/worker/test_native_duplex_hooks.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 
 import ast
 import inspect
+from dataclasses import replace
 from types import SimpleNamespace
 
 import numpy as np
@@ -1071,6 +1072,158 @@ def _stage0_vision_runtime():
     return runtime
 
 
+def _minicpmo_seeded_silence_sampling_case(initial_user_text: str | None, *, force_listen_count: int = 0):
+    from vllm_omni.engine.duplex.contracts import DuplexFence
+    from vllm_omni.model_executor.models.minicpmo_4_5.duplex.plugin import build_duplex_data_plane_prompt
+    from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import _MiniCPMO45Stage0SessionState
+
+    runtime = _stage0_vision_runtime()
+    state = _MiniCPMO45Stage0SessionState(session_id="sid-seeded-text")
+    runtime_config = {"initial_user_text": initial_user_text}
+    runtime._prepare_session_context(state, {"instructions": "Speak exactly."}, runtime_config=runtime_config)
+    prefill = runtime._stage_prefill_embeddings_only(
+        state, np.zeros(4, dtype=np.float32), epoch=0, seq=1, is_speech=False
+    )
+    assert prefill["success"] is True
+    prompt = build_duplex_data_plane_prompt(
+        request_id="req-seeded-text",
+        fence=DuplexFence(state.session_id),
+        session_config={"extra_body": {"force_listen_count": force_listen_count}},
+        runtime_config=runtime_config,
+        seq=1,
+        turn_seq=1,
+        payload={"type": "audio", "is_speech": False},
+        final=False,
+    )
+    payload = prompt["model_intermediate_buffer"]["duplex"]["payload"]
+    model, row = _minicpmo_duplex_policy_case(state, payload)
+    model._minicpmo45_native_duplex_token_ids_cache = runtime._special_token_ids()
+    return runtime, state, model, row
+
+
+@pytest.mark.parametrize("initial_user_text", [None, "", "Get the trust fund to the bank early."])
+@pytest.mark.parametrize("force_listen_count", [0, 1])
+def test_minicpmo_seeded_text_reaches_sampling_on_silence(initial_user_text, force_listen_count):
+    """A seeded user turn must reach the model despite its silent audio clock."""
+    runtime, state, model, row = _minicpmo_seeded_silence_sampling_case(
+        initial_user_text, force_listen_count=force_listen_count
+    )
+    logits = torch.zeros((1, 32), dtype=torch.float32)
+    logits[0, 20] = 10.0
+    original_logits = logits.clone()
+
+    model.prepare_duplex_sampling(logits, SimpleNamespace(), (row,))
+
+    if initial_user_text and not force_listen_count:
+        assert torch.equal(logits, original_logits)
+        assert state.pending_speech_context is True
+    else:
+        assert logits[0, runtime.listen_token_id] == 0
+        assert torch.isfinite(logits).sum() == 1
+
+
+def test_minicpmo_seeded_text_is_consumed_before_post_turn_silence():
+    """The opening text allows one answer; silence after its end stays quiet."""
+    from dataclasses import replace
+
+    runtime, state, model, row = _minicpmo_seeded_silence_sampling_case("Read this sentence.")
+    token_ids = runtime._special_token_ids()
+    logits = torch.zeros((1, 32), dtype=torch.float32)
+    model.prepare_duplex_sampling(logits, SimpleNamespace(), (row,))
+    assert torch.isfinite(logits).all()
+
+    model._record_minicpmo45_duplex_terminator(0, runtime.tts_bos_token_id, token_ids)
+    assert state.pending_speech_context is True
+    model._record_minicpmo45_duplex_terminator(0, 20, token_ids)
+    assert state.pending_speech_context is False
+    model._record_minicpmo45_duplex_terminator(0, runtime.turn_eos_token_id, token_ids)
+    assert state.current_turn_ended is True
+
+    prefill = runtime._stage_prefill_embeddings_only(
+        state, np.zeros(4, dtype=np.float32), epoch=0, seq=2, is_speech=False
+    )
+    assert prefill["success"] is True
+    next_logits = torch.zeros((1, 32), dtype=torch.float32)
+    model.prepare_duplex_sampling(next_logits, SimpleNamespace(), (replace(row, seq=row.seq + 1),))
+    assert next_logits[0, runtime.listen_token_id] == 0
+    assert torch.isfinite(next_logits).sum() == 1
+
+
+def test_minicpmo_seeded_text_survives_explicit_initial_listen():
+    """The force-listen prefix delays the seed without losing its answer."""
+    from dataclasses import replace
+
+    runtime, state, model, row = _minicpmo_seeded_silence_sampling_case("Read this sentence.", force_listen_count=1)
+    token_ids = runtime._special_token_ids()
+    logits = torch.zeros((1, 32), dtype=torch.float32)
+    model.prepare_duplex_sampling(logits, SimpleNamespace(), (row,))
+    assert torch.isfinite(logits).sum() == 1
+    model._record_minicpmo45_duplex_terminator(0, runtime.listen_token_id, token_ids)
+
+    prefill = runtime._stage_prefill_embeddings_only(
+        state, np.zeros(4, dtype=np.float32), epoch=0, seq=2, is_speech=False
+    )
+    assert prefill["success"] is True
+    next_logits = torch.zeros((1, 32), dtype=torch.float32)
+    next_row = replace(row, seq=row.seq + 1, payload={"is_speech": False, "force_listen": False})
+    model.prepare_duplex_sampling(next_logits, SimpleNamespace(), (next_row,))
+    assert torch.isfinite(next_logits).all()
+    assert state.pending_speech_context is True
+
+
+@pytest.mark.parametrize("initial_text", [None, "", "Please say hello.", "请说你好。"])
+@pytest.mark.parametrize("speech_envelope", [False, True])
+def test_minicpmo_seeded_context_keeps_assistant_turn_open_until_eos(initial_text, speech_envelope, mocker):
+    from vllm.v1.sample.metadata import SamplingMetadata
+
+    from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import _MiniCPMO45Stage0SessionState
+
+    runtime = _stage0_vision_runtime()
+    state = _MiniCPMO45Stage0SessionState(session_id="seeded-context")
+    runtime._prepare_session_context(state, {}, runtime_config={"initial_user_text": initial_text})
+    result = runtime._stage_prefill_embeddings_only(
+        state, np.zeros(4, dtype=np.float32), epoch=0, seq=1, is_speech=False
+    )
+    assert result["success"] is True
+    model, row = _minicpmo_duplex_policy_case(state, {"is_speech": False})
+    logits = torch.zeros((1, 16), dtype=torch.float32)
+    logits[0, 10] = 20.0
+    original_logits = logits.clone()
+    sampling_metadata = mocker.Mock(spec=SamplingMetadata)
+    model.prepare_duplex_sampling(logits, sampling_metadata, (row,))
+    token_ids = model._minicpmo45_native_duplex_token_ids_cache
+    sampled = model._finalize_minicpmo45_native_duplex_sample(0, 7, token_ids)
+
+    if not initial_text:
+        assert sampled == 7
+        assert logits[0, 7].item() == 0.0
+        assert torch.isneginf(logits[0, :7]).all()
+        assert torch.isneginf(logits[0, 8:]).all()
+        return
+
+    assert torch.equal(logits, original_logits)
+    assert sampled == 8
+    if speech_envelope:
+        model._record_minicpmo45_duplex_terminator(0, 8, token_ids)
+        assert state.current_turn_ended is False
+        assert state.pending_speech_context is True
+    model._record_minicpmo45_duplex_terminator(0, 10, token_ids)
+    assert state.pending_speech_context is False
+    model._record_minicpmo45_duplex_terminator(0, 9, token_ids)
+
+    # A later silent unit must not re-admit the seed after the turn ended.
+    result = runtime._stage_prefill_embeddings_only(
+        state, np.zeros(4, dtype=np.float32), epoch=0, seq=2, is_speech=False
+    )
+    assert result["success"] is True
+    logits = original_logits.clone()
+    model.prepare_duplex_sampling(logits, sampling_metadata, (replace(row, seq=4),))
+    assert logits[0, 7].item() == 0.0
+    assert torch.isneginf(logits[0, :7]).all()
+    assert torch.isneginf(logits[0, 8:]).all()
+    assert model._finalize_minicpmo45_native_duplex_sample(0, 7, token_ids) == 7
+
+
 def test_minicpmo_stage0_puts_every_frame_of_an_append_in_one_unit():
     """Official streaming_prefill feeds the whole frame_list into one unit.
 

tests/worker_v2/test_omni_ar_model_runner.py

@@ -217,6 +217,51 @@ def test_build_async_chunk_outputs_slices_padded_axis_and_splits_channels() -> N
     assert client is None and torch.equal(inter_stage[0]["codes.audio"], req_codes[0])
 
 
+@pytest.mark.parametrize("async_chunk", [False, True])
+@pytest.mark.parametrize(
+    "hidden_padded,codes_padded",
+    [(False, False), (True, False), (False, True), (True, True)],
+    ids=["unpadded", "hidden_padded", "codes_padded", "both_padded"],
+)
+@pytest.mark.parametrize("lengths", [(1, 1, 1), (3, 1, 1)], ids=["decode", "mixed_prefill_decode"])
+def test_async_output_slices_request_payloads_with_graph_padding(
+    monkeypatch, mocker, async_chunk, hidden_padded, codes_padded, lengths
+):
+    """Sync and async transfers must exclude other requests and graph padding."""
+    from vllm.v1.worker.gpu.input_batch import InputBatch
+
+    monkeypatch.setattr(torch.cuda, "set_stream", lambda _stream: None)
+    total = sum(lengths)
+    padded_total = 8 if hidden_padded or codes_padded else total
+    offsets = np.cumsum([0, *lengths])
+    hidden_rows = padded_total if hidden_padded else total
+    code_rows = padded_total if codes_padded else total
+    hidden = torch.arange(hidden_rows * 4, dtype=torch.float32).reshape(hidden_rows, 4)
+    codes = torch.arange(code_rows * 16).reshape(code_rows, 16)
+    # Reference frames are request-local, even if their length matches the
+    # padded token count. They must not be sliced along the batch token axis.
+    refs = [torch.arange(padded_total * 16).reshape(padded_total, 16), torch.empty(0), torch.empty(0)]
+    batch = mocker.Mock(spec=InputBatch)
+    batch.query_start_loc_np = offsets
+    batch.num_scheduled_tokens = np.array(lengths)
+    batch.num_reqs = len(lengths)
+    batch.num_tokens_after_padding = padded_total
+    output = _async_output(
+        req_ids=[f"req-{i}" for i in range(len(lengths))],
+        sampler_output=SamplerOutput(torch.ones(len(lengths), 1, dtype=torch.long), None, None, None),
+        text_hidden=hidden,
+        multimodal_outputs={"codes": {"audio": codes, "ref": refs}},
+        input_batch=batch,
+        async_chunk=async_chunk,
+    ).get_output()
+
+    for i, payload in enumerate(output.inter_stage_outputs):
+        torch.testing.assert_close(payload["codes.audio"], codes[offsets[i] : offsets[i + 1]])
+        torch.testing.assert_close(payload["codes.ref"], refs[i])
+        if not async_chunk:
+            torch.testing.assert_close(payload["hidden"], hidden[offsets[i] : offsets[i + 1]])
+
+
 def test_async_chunk_output_stages_mm_on_copy_stream_before_get_output(monkeypatch) -> None:
     monkeypatch.setattr(torch.cuda, "set_stream", lambda _stream: None)
     calls = []

vllm_omni/clients/nemotron_voicechat.py

@@ -0,0 +1,29 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Client-side duplex session preset for Nemotron VoiceChat."""
+
+from __future__ import annotations
+
+from vllm_omni.clients.duplex import AudioFormat, SessionConfig
+
+__all__ = ["create_duplex_session_config"]
+
+
+def create_duplex_session_config(
+    *,
+    instructions: str | None = None,
+    tools: list[dict[str, object]] | None = None,
+    idle_timeout_s: float | None = None,
+    extra_body: dict[str, object] | None = None,
+) -> SessionConfig:
+    """Use 16 kHz float32 input and 22.05 kHz PCM16 output; append 80 ms frames."""
+    session_extra_body = dict(extra_body or {})
+    if tools is not None:
+        session_extra_body["realtime_tools"] = tools
+    return SessionConfig(
+        input_audio=AudioFormat("pcm_f32le", 16_000),
+        output_audio=AudioFormat("pcm16", 22_050),
+        instructions=instructions,
+        idle_timeout_s=idle_timeout_s,
+        extra_body=session_extra_body,
+    )

vllm_omni/config/stage_config.py

@@ -312,10 +312,10 @@ class PipelineConfig:
     # ``hf_config_predicate=lambda c: getattr(c, "version", "") == "4.5"``
     # to avoid misrouting 2.6 checkpoints.
     hf_config_predicate: Callable[[Any], bool] | None = None
-    # Diffusers pipeline class name: for models that ship a ``model_index.json``
-    # (no root ``config.json``), the ``_class_name`` field is matched against
-    # this value to auto-detect the pipeline.  Only needed for diffusers-style
-    # multi-component repos (e.g. GLM-Image).  ``None`` = not a diffusers model.
+    # Canonical diffusion runtime class for this registered pipeline.
+    # Serving uses it as a metadata fallback; model_index.json discovery
+    # also matches its _class_name against this value and its aliases.
+    # None delegates class discovery to the checkpoint metadata.
     diffusers_class_name: str | None = None
     diffusers_class_aliases: tuple[str, ...] = ()
     endpoint_restrictions: tuple[EndpointRestriction, ...] = ()

vllm_omni/deploy/mammoth_moda2.yaml

@@ -15,7 +15,8 @@ stages:
 
   - stage_id: 1
     devices: "0"
-    max_num_seqs: 1
+    max_num_seqs: 8
+    request_batch_max_wait_ms: 20
     gpu_memory_utilization: 0.3
     enforce_eager: true
     trust_remote_code: true

vllm_omni/deploy/mammoth_moda2_fp8_kv.yaml

@@ -5,3 +5,6 @@ base_config: mammoth_moda2.yaml
 stages:
   - stage_id: 0
     kv_cache_dtype: fp8_e4m3
+    # Layer 0 stays BF16: with every layer in FP8, most text-to-image samples
+    # lose the prompt (A800 results in recipes/MammothModa2/MammothModa2.md).
+    kv_cache_dtype_skip_layers: ["0"]

vllm_omni/deploy/qwen3_omni_moe.yaml

@@ -46,6 +46,16 @@ stages:
     input_connectors:
       from_stage_0: connector_of_shared_memory
     default_sampling_params:
+      # The talker is the only stochastic stage (the thinker and code2wav both
+      # decode greedily at temperature 0). Without a seed, every request draws
+      # fresh codec tokens, and a small fraction of draws derail into fluent
+      # nonsense part-way through an utterance while the thinker text stays
+      # correct -- the audio-vs-text nightly failures in #6090. This seed was
+      # present from the original deploy schema and was dropped as collateral
+      # in #4986; keep it so Qwen3-Omni speech is reproducible for a given
+      # request. It bounds sampling only: it does not make the talker immune
+      # to the underlying degeneration, which is tracked separately.
+      seed: 42
       temperature: 0.9
       top_k: 50
       max_tokens: 4096

vllm_omni/diffusion/cancellation.py

@@ -0,0 +1,141 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Cooperative cancellation for local, full-forward diffusion workers.
+
+The engine owns each signal until the worker has returned. Unlike an executor
+RPC, writing the signal does not wait behind the forward being cancelled.
+Opted-in pipelines call ``check_request_cancellation`` at safe boundaries.
+"""
+
+from __future__ import annotations
+
+import threading
+from collections.abc import Iterable, Iterator, Sequence
+from contextlib import contextmanager
+from contextvars import ContextVar
+from multiprocessing.shared_memory import SharedMemory
+
+
+class RequestCancellationRegistry:
+    """Engine-owned signals; scheduler mutation remains on the engine thread."""
+
+    def __init__(self) -> None:
+        self._signals: dict[str, SharedMemory] = {}
+        self._lock = threading.Lock()
+
+    def create(self, request_id: str) -> str:
+        with self._lock:
+            if request_id in self._signals:
+                raise ValueError(f"Duplicate cancellation request: {request_id}")
+            signal = SharedMemory(create=True, size=1)
+            signal.buf[0] = 0
+            self._signals[request_id] = signal
+            return signal.name
+
+    def cancel(self, request_ids: Iterable[str]) -> None:
+        with self._lock:
+            for request_id in request_ids:
+                signal = self._signals.get(request_id)
+                if signal is not None:
+                    signal.buf[0] = 1
+
+    def cancel_all(self) -> None:
+        with self._lock:
+            for signal in self._signals.values():
+                signal.buf[0] = 1
+
+    @staticmethod
+    def _dispose(signal: SharedMemory) -> None:
+        try:
+            signal.unlink()
+        except FileNotFoundError:
+            # A process-level resource tracker may already have removed it.
+            pass
+        finally:
+            signal.close()
+
+    def finish(self, request_id: str) -> None:
+        """Release only after execution returns, including aborted execution."""
+        with self._lock:
+            signal = self._signals.pop(request_id, None)
+            if signal is not None:
+                self._dispose(signal)
+
+    def close(self) -> None:
+        """Release remaining signals after the executor has shut down."""
+        with self._lock:
+            for signal in self._signals.values():
+                self._dispose(signal)
+            self._signals.clear()
+
+
+_current_signals: ContextVar[tuple[SharedMemory | None, ...] | None] = ContextVar(
+    "diffusion_request_cancellation", default=None
+)
+
+
+@contextmanager
+def request_cancellation_scope(signal_names: Sequence[str | None], *, enabled: bool = True) -> Iterator[None]:
+    """Attach worker readers without taking ownership of the engine's names."""
+    signals: list[SharedMemory | None] = []
+    try:
+        for name in signal_names if enabled else ():
+            signals.append(SharedMemory(name=name) if name is not None else None)
+        # An opted-in rank without a local signal must still join every
+        # checkpoint collective and vote "not cancelled". None means only
+        # that this execution does not use cooperative cancellation at all.
+        token = _current_signals.set(tuple(signals) if enabled else None)
+        try:
+            yield
+        finally:
+            _current_signals.reset(token)
+    finally:
+        for signal in signals:
+            if signal is not None:
+                signal.close()
+
+
+def check_request_cancellation(*, synchronize: bool = False) -> None:
+    """Stop a cancelled execution wave without stranding a peer's collectives.
+
+    ``synchronize=True`` drains queued device work only when cancellation has
+    been requested locally, before re-reading the flags. Every abort drains queued
+    device work before unwinding. Successful steps retain asynchronous execution
+    when no local cancellation is observed.
+    Independent requests coupled by an AllGather offload wave must all be cancelled
+    before the wave can exit; cancelling one must not abort its live peers.
+    """
+    signals = _current_signals.get()
+    if signals is None:
+        return
+
+    import torch
+
+    locally_requested = any(signal is not None and signal.buf[0] for signal in signals)
+    device_drained = False
+    if synchronize and locally_requested:
+        from vllm_omni.platforms import current_omni_platform
+
+        current_omni_platform.synchronize()
+        device_drained = True
+
+    cancelled = bool(signals) and all(signal is not None and signal.buf[0] for signal in signals)
+    if torch.distributed.is_initialized():
+        from vllm_omni.diffusion.distributed.parallel_state import get_world_group
+
+        group = get_world_group()
+        if group.world_size > 1:
+            # Reading a shared byte alone is insufficient: ranks can observe a
+            # concurrent write on different steps. Make one collective decision.
+            flag = torch.tensor([int(cancelled)], dtype=torch.int32, device="cpu")
+            torch.distributed.all_reduce(flag, op=torch.distributed.ReduceOp.MIN, group=group.cpu_group)
+            cancelled = bool(flag.item())
+
+    if cancelled:
+        from vllm_omni.diffusion.data import DiffusionRequestAbortedError
+
+        if not device_drained:
+            from vllm_omni.platforms import current_omni_platform
+
+            current_omni_platform.synchronize()
+        raise DiffusionRequestAbortedError("Request cancelled at a model execution boundary")

vllm_omni/diffusion/diffusion_engine.py

@@ -25,6 +25,7 @@ from vllm.utils.import_utils import resolve_obj_by_qualname
 from vllm.v1.engine.exceptions import EngineDeadError
 from vllm.v1.kv_cache_interface import KVCacheConfig
 
+from vllm_omni.diffusion.cancellation import RequestCancellationRegistry
 from vllm_omni.diffusion.data import (
     DIFFUSION_REQUEST_LIFECYCLE_KEY,
     DIFFUSION_REQUEST_STARTED,
@@ -180,6 +181,19 @@ def _max_num_seqs(od_config: OmniDiffusionConfig) -> int:
         return 1
 
 
+def supports_request_cancellation(od_config: OmniDiffusionConfig) -> bool:
+    """Whether the local pipeline checks cooperative cancellation boundaries."""
+    model_cls = _resolve_custom_pipeline_cls(getattr(od_config, "custom_pipeline_args", None))
+    if model_cls is None:
+        name = (
+            "DiffusersAdapterPipeline"
+            if uses_diffusers_adapter(od_config)
+            else getattr(od_config, "model_class_name", None)
+        )
+        model_cls = DiffusionModelRegistry._try_load_model_cls(name)
+    return getattr(model_cls, "supports_request_cancellation", False) is True
+
+
 def _uses_dlo_dp_concurrency(od_config: OmniDiffusionConfig) -> bool:
     parallel_config = getattr(od_config, "parallel_config", None)
     dp_size = getattr(parallel_config, "data_parallel_size", 1)
@@ -230,6 +244,8 @@ class DiffusionEngine:
     # don't hit AttributeError when _busy_loop accesses them.
     dp_concurrent: bool = False
     _scheduling_paused: bool = False
+    # Disabled until runtime initialization resolves the pipeline capability.
+    _request_cancellations: RequestCancellationRegistry | None = None
 
     def __init__(
         self,
@@ -393,6 +409,9 @@ class DiffusionEngine:
         self._shutting_down = False
         self._shutdown_complete = False
         self.abort_queue: queue.Queue[str] = queue.Queue()
+        self._request_cancellations = (
+            RequestCancellationRegistry() if supports_request_cancellation(self.od_config) else None
+        )
         self._rpc_queue: queue.Queue[_RpcTask] = queue.Queue()
         # pause_scheduler(mode="keep"): no new batch is scheduled while set.
         self._scheduling_paused = False
@@ -913,6 +932,8 @@ class DiffusionEngine:
             # path, and repeating executor.shutdown() is unsafe.
             self._shutting_down = True
             self._closed = True
+            if self._request_cancellations is not None:
+                self._request_cancellations.cancel_all()
             if self.stop_event is not None:
                 self.stop_event.set()
             streams = list(self._out_streams.values())
@@ -925,7 +946,13 @@ class DiffusionEngine:
         # If Worker shutdown fails, retain the Scheduler reservations. A
         # remote producer may still be writing into those allocations.
         self.executor.shutdown()
-        self.scheduler.close()
+        try:
+            self.scheduler.close()
+        finally:
+            # Workers are down, so readers cannot attach or use the signals.
+            # Release them even if scheduler cleanup itself fails.
+            if self._request_cancellations is not None:
+                self._request_cancellations.close()
         self._shutdown_complete = True
 
     def _emit_finished_outputs(
@@ -1106,7 +1133,15 @@ class DiffusionEngine:
             if self._closed:
                 raise RuntimeError("DiffusionEngine is closed.")
             queue: asyncio.Queue[DiffusionOutput] = asyncio.Queue()
-            request_id = self.scheduler.add_request(request)
+            if self._request_cancellations is not None:
+                request.cancellation_signal = self._request_cancellations.create(request.request_id)
+            try:
+                request_id = self.scheduler.add_request(request)
+            except BaseException:
+                if self._request_cancellations is not None:
+                    self._request_cancellations.finish(request.request_id)
+                    request.cancellation_signal = None
+                raise
             self._out_streams[request_id] = queue
             self._cv.notify_all()
 
@@ -1517,6 +1552,8 @@ class DiffusionEngine:
                 return
             if not self._closed:
                 self._closed = True
+                if self._request_cancellations is not None:
+                    self._request_cancellations.cancel_all()
                 if self.stop_event is not None:
                     self.stop_event.set()
                 pending_streams = list(self._out_streams.values())
@@ -1532,6 +1569,9 @@ class DiffusionEngine:
             if worker_thread.is_alive():
                 worker_thread.join(timeout=10)
             if worker_thread.is_alive():
+                # Keep cancellation names available: an in-flight worker may
+                # not have attached its readers yet. A later close releases
+                # them after execution and executor shutdown have completed.
                 logger.warning(
                     "Worker thread did not terminate within 10s; scheduler and executor shutdown will be deferred."
                 )
@@ -1550,6 +1590,8 @@ class DiffusionEngine:
         else:
             self.scheduler.close()
             self.executor.shutdown()
+        if self._request_cancellations is not None:
+            self._request_cancellations.close()
         self._shutdown_complete = True
 
     def abort(self, request_id: str | Iterable[str]) -> None:
@@ -1558,6 +1600,10 @@ class DiffusionEngine:
         with self._cv:
             if self._closed:
                 return
+            if self._request_cancellations is not None:
+                # Do not queue this behind the full-forward executor call.
+                # Scheduler state is still mutated only by the busy loop.
+                self._request_cancellations.cancel(request_ids)
             for req_id in request_ids:
                 self.abort_queue.put(req_id)
             self._cv.notify_all()
@@ -1601,6 +1647,8 @@ class DiffusionEngine:
         state = self.scheduler.get_request_state(request_id)
         popped_state = self.scheduler.pop_request_state(request_id)
         state = state or popped_state
+        if self._request_cancellations is not None:
+            self._request_cancellations.finish(request_id)
 
         if state is None:
             raise RuntimeError(f"Diffusion scheduler lost state for request {request_id}.")

vllm_omni/diffusion/distributed/sp_plan.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM and The HuggingFace Team
 # Type definitions in this module are adapted from HuggingFace diffusers library:
 #   diffusers/src/diffusers/models/_modeling_parallel.py
@@ -36,7 +37,7 @@ Example:
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, TypeGuard
 
 if TYPE_CHECKING:
     import torch
@@ -220,10 +221,13 @@ class SequenceParallelInput:
             This is useful for layers whose outputs should be split after preprocessing
             (e.g., RoPE embeddings).
         auto_pad: If True, automatically pad the tensor if its size along split_dim
-            is not divisible by world_size. Creates an attention mask to indicate
-            valid vs padding positions. The mask is stored in ForwardContext.
+            is not divisible by world_size. Padding metadata is stored in
+            ForwardContext so models can construct attention masks when needed.
             Note: Ring attention does not support attention mask, so auto_pad
             should only be used with Ulysses SP.
+        shard_group: Optional key shared by tensors representing the same global
+            sequence. Keyed groups track independent padding metadata; omitting
+            it preserves the legacy single-sequence padding behavior.
 
     Example:
         # Split hidden_states along sequence dimension (dim 1)
@@ -240,12 +244,13 @@ class SequenceParallelInput:
     expected_dims: int | None = None
     split_output: bool = False
     auto_pad: bool = False
+    shard_group: str | None = None
 
     def __repr__(self) -> str:
         return (
             f"SequenceParallelInput(split_dim={self.split_dim}, "
             f"expected_dims={self.expected_dims}, split_output={self.split_output}, "
-            f"auto_pad={self.auto_pad})"
+            f"auto_pad={self.auto_pad}, shard_group={self.shard_group!r})"
         )
 
 
@@ -262,6 +267,8 @@ class SequenceParallelOutput:
         gather_dim: The dimension along which to gather the tensor.
         expected_dims: Expected number of dimensions. If provided, validates that
             the tensor has this many dimensions before gathering.
+        shard_group: Optional key whose padding metadata determines the gathered
+            sequence length. Omitting it preserves legacy singleton trimming.
 
     Example:
         # Gather output along sequence dimension (dim 1)
@@ -270,9 +277,13 @@ class SequenceParallelOutput:
 
     gather_dim: int
     expected_dims: int | None = None
+    shard_group: str | None = None
 
     def __repr__(self) -> str:
-        return f"SequenceParallelOutput(gather_dim={self.gather_dim}, expected_dims={self.expected_dims})"
+        return (
+            f"SequenceParallelOutput(gather_dim={self.gather_dim}, "
+            f"expected_dims={self.expected_dims}, shard_group={self.shard_group!r})"
+        )
 
 
 @dataclass(frozen=True)
@@ -372,12 +383,14 @@ SequenceParallelModelPlan = dict[str, SequenceParallelInputType | SequenceParall
 # =============================================================================
 
 
-def _is_valid_input_config(value: object) -> bool:
+def _is_valid_input_config(value: object) -> TypeGuard[AnySequenceParallelInput]:
     """Check if a value is a valid input configuration type."""
     return isinstance(value, (SequenceParallelInput, SequenceParallelPartialInput))
 
 
-def _is_valid_input_config_list(value: object) -> bool:
+def _is_valid_input_config_list(
+    value: object,
+) -> TypeGuard[list[AnySequenceParallelInput] | tuple[AnySequenceParallelInput, ...]]:
     """Check if a value is a list/tuple of valid input configurations."""
     if not isinstance(value, (list, tuple)):
         return False
@@ -396,6 +409,27 @@ def validate_sp_plan(plan: SequenceParallelModelPlan) -> None:
     if not isinstance(plan, dict):
         raise ValueError(f"_sp_plan must be a dict, got {type(plan).__name__}")
 
+    # A shard_group names one split boundary's global sequence. Reusing it
+    # across boundaries would silently overwrite the recorded length.
+    shard_group_owners: dict[str, str] = {}
+
+    def _validate_input(sp_input: AnySequenceParallelInput, *, module_id: str) -> None:
+        if not isinstance(sp_input, SequenceParallelInput) or sp_input.shard_group is None:
+            return
+        if not sp_input.auto_pad:
+            raise ValueError(
+                f"SequenceParallelInput for module {module_id!r} declares "
+                f"shard_group={sp_input.shard_group!r}, but keyed shard "
+                "metadata currently requires auto_pad=True."
+            )
+        owner = shard_group_owners.setdefault(sp_input.shard_group, module_id)
+        if owner != module_id:
+            raise ValueError(
+                f"shard_group={sp_input.shard_group!r} is used by multiple "
+                f"split boundaries ({owner!r} and {module_id!r}). Use a unique "
+                "name for each boundary."
+            )
+
     for module_id, module_plan in plan.items():
         if not isinstance(module_id, str):
             raise ValueError(f"_sp_plan keys must be strings, got {type(module_id).__name__}")
@@ -407,6 +441,8 @@ def validate_sp_plan(plan: SequenceParallelModelPlan) -> None:
             if all(isinstance(x, SequenceParallelOutput) for x in module_plan):
                 continue
             if _is_valid_input_config_list(module_plan):
+                for sp_input in module_plan:
+                    _validate_input(sp_input, module_id=module_id)
                 # List of inputs for a specific parameter (when output is tuple)
                 continue
 
@@ -428,8 +464,10 @@ def validate_sp_plan(plan: SequenceParallelModelPlan) -> None:
                             f"Integer keys (output indices) require split_output=True, "
                             f"got split_output=False for module '{module_id}'[{key}]"
                         )
+                    _validate_input(value, module_id=module_id)
                 elif _is_valid_input_config_list(value):
-                    pass  # Valid list of input configs
+                    for sp_input in value:
+                        _validate_input(sp_input, module_id=module_id)
                 else:
                     raise ValueError(
                         f"Input spec values must be SequenceParallelInput/PartialInput or list thereof, "

vllm_omni/diffusion/forward_context.py

@@ -64,6 +64,10 @@ class ForwardContext:
     sp_padding_size: int = 0
     # Original sequence length before padding (for removing padding in gather)
     sp_original_seq_len: int | None = None
+    # Pre-padding global sequence length per SequenceParallelInput shard_group,
+    # for models that auto-pad several independent sequences in one boundary.
+    # The singleton fields above remain for single-sequence models.
+    sp_shard_metadata: dict[str, int] = field(default_factory=dict)
 
     # Set by registry when _sp_plan hooks are applied.
     # When True, sp_active is determined by _sp_shard_depth (for _sp_plan hooks)
@@ -127,6 +131,13 @@ def is_forward_context_available() -> bool:
     return _forward_context is not None
 
 
+def get_sp_shard_original_seq_len(shard_group: str) -> int | None:
+    """Pre-padding global length of `shard_group`, or None if it was not split."""
+    if not is_forward_context_available():
+        return None
+    return get_forward_context().sp_shard_metadata.get(shard_group)
+
+
 def build_local_sp_padding_mask(
     batch_size: int,
     local_seq_len: int,

vllm_omni/diffusion/hooks/sequence_parallel.py

@@ -182,6 +182,12 @@ class SequenceParallelSplitHook(ModelHook):
         self.module_forward_metadata: ModuleForwardMetadata | None = None
         # Cache for text lengths resolved from kwargs
         self._text_len_cache: dict[str, int] = {}
+        self._shard_groups = {
+            sp_input.shard_group
+            for value in metadata.values()
+            for sp_input in (value if isinstance(value, (list, tuple)) else (value,))
+            if isinstance(sp_input, SequenceParallelInput) and sp_input.shard_group is not None
+        }
 
     def initialize_hook(self, module: nn.Module) -> nn.Module:
         cls = _unwrap_module(module).__class__
@@ -193,6 +199,16 @@ class SequenceParallelSplitHook(ModelHook):
         args_list = list(args)
         # Clear text length cache for this forward pass
         self._text_len_cache.clear()
+        if self._shard_groups:
+            from vllm_omni.diffusion.forward_context import (
+                get_forward_context,
+                is_forward_context_available,
+            )
+
+            if is_forward_context_available():
+                ctx = get_forward_context()
+                for shard_group in self._shard_groups:
+                    ctx.sp_shard_metadata.pop(shard_group, None)
 
         for name, spm in self.metadata.items():
             # Skip if this is a split_output entry (handled in post_forward)
@@ -376,7 +392,11 @@ class SequenceParallelSplitHook(ModelHook):
         if isinstance(sp_input, SequenceParallelInput):
             # Full split with optional auto-padding
             if sp_input.auto_pad:
-                return self._shard_with_auto_pad(x, sp_input.split_dim)
+                return self._shard_with_auto_pad(
+                    x,
+                    sp_input.split_dim,
+                    sp_input.shard_group,
+                )
             _maybe_validate_strict_divisibility(dim=sp_input.split_dim, seq_len=x.size(sp_input.split_dim))
             return sp_shard(x, sp_input.split_dim, validate=False)
         elif isinstance(sp_input, SequenceParallelPartialInput):
@@ -397,13 +417,17 @@ class SequenceParallelSplitHook(ModelHook):
         else:
             raise ValueError(f"Unsupported input config type: {type(sp_input).__name__}")
 
-    def _shard_with_auto_pad(self, x: torch.Tensor, dim: int) -> torch.Tensor:
-        """Shard tensor with automatic padding and attention mask creation.
+    def _shard_with_auto_pad(
+        self,
+        x: torch.Tensor,
+        dim: int,
+        shard_group: str | None,
+    ) -> torch.Tensor:
+        """Pad and shard a tensor while recording its global sequence length.
 
-        When sequence length is not divisible by SP world size, this method:
-        1. Pads the tensor to make it divisible
-        2. Creates an attention mask indicating valid vs padding positions
-        3. Stores the mask and padding info in ForwardContext
+        Models with several independent sequences read the recorded length back
+        via `shard_group` to build masks. Legacy singleton padding fields are
+        also populated for single-sequence models.
         """
         from vllm_omni.diffusion.attention.selector import get_attn_backend_for_capability
         from vllm_omni.diffusion.distributed.parallel_state import (
@@ -418,10 +442,20 @@ class SequenceParallelSplitHook(ModelHook):
             return x
 
         seq_len = x.size(dim)
+        rank = get_sequence_parallel_rank()
         remainder = seq_len % world_size
 
+        if shard_group is not None and is_forward_context_available():
+            ctx = get_forward_context()
+            existing = ctx.sp_shard_metadata.setdefault(shard_group, seq_len)
+            if existing != seq_len:
+                raise ValueError(
+                    f"All tensors in shard_group={shard_group!r} must have the "
+                    f"same global sequence length, but got {existing} and {seq_len}."
+                )
+
         if remainder == 0:
-            # No padding needed
+            # Keyed groups record their length even when no padding is needed.
             return sp_shard(x, dim, validate=False)
 
         # Check backend compatibility
@@ -464,8 +498,11 @@ class SequenceParallelSplitHook(ModelHook):
         x_padded = torch.cat([x, padding], dim=dim)
 
         # Store padding info in forward context (only once, for primary tensor)
-        # Attention layers will create masks dynamically using this info
-        if is_forward_context_available():
+        # Attention layers will create masks dynamically using this info.
+        # Keyed groups record their length in sp_shard_metadata above and must
+        # not clobber the legacy singleton fields, whose semantics assume a
+        # single sequence per boundary.
+        if shard_group is None and is_forward_context_available():
             ctx = get_forward_context()
             # Only set if not already set (first auto_pad tensor wins)
             if ctx.sp_original_seq_len is None:
@@ -477,7 +514,6 @@ class SequenceParallelSplitHook(ModelHook):
                 )
 
         # Shard the padded tensor
-        rank = get_sequence_parallel_rank()
         return x_padded.chunk(world_size, dim=dim)[rank]
 
 
@@ -523,12 +559,6 @@ class SequenceParallelGatherHook(ModelHook):
         if len(output) != len(self.metadata):
             raise ValueError(f"Expected {len(self.metadata)} outputs, got {len(output)}.")
 
-        # Check if padding was applied during split
-        original_seq_len = None
-        if is_forward_context_available():
-            ctx = get_forward_context()
-            original_seq_len = ctx.sp_original_seq_len
-
         actually_gathered = False
 
         for i, spm in enumerate(self.metadata):
@@ -546,6 +576,15 @@ class SequenceParallelGatherHook(ModelHook):
             gathered = sp_gather(x, spm.gather_dim, validate=False)
 
             # Remove padding if it was applied
+            original_seq_len = None
+            if is_forward_context_available():
+                ctx = get_forward_context()
+                if spm.shard_group is None:
+                    original_seq_len = ctx.sp_original_seq_len
+                else:
+                    original_seq_len = ctx.sp_shard_metadata.get(spm.shard_group)
+                    if original_seq_len is None:
+                        raise RuntimeError(f"No SP shard metadata was registered for shard_group={spm.shard_group!r}.")
             if original_seq_len is not None and gathered.size(spm.gather_dim) > original_seq_len:
                 gathered = gathered.narrow(spm.gather_dim, 0, original_seq_len)
                 logger.debug(f"Removed padding: gathered shape {gathered.shape} (original_seq_len={original_seq_len})")
@@ -559,6 +598,9 @@ class SequenceParallelGatherHook(ModelHook):
             ctx._sp_shard_depth = max(0, ctx._sp_shard_depth - 1)
             if ctx._sp_equal_pad_stack:
                 ctx._sp_equal_pad_stack.pop()
+            for spm in self.metadata:
+                if spm is not None and spm.shard_group is not None:
+                    ctx.sp_shard_metadata.pop(spm.shard_group, None)
 
         return output[0] if is_tensor else type(output)(output)
 

vllm_omni/diffusion/layers/mhc.py

@@ -0,0 +1,279 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# Copyright (c) 2026 SandAI. All Rights Reserved.
+
+"""Native and eager Triton mHC post-processing, without projection or tensor caches."""
+
+import threading
+
+import torch
+from vllm.logger import init_logger
+from vllm.triton_utils import HAS_TRITON, tl, triton
+
+from vllm_omni.diffusion.layers.custom_op import CustomOp
+
+_FAILED_MHC_KERNELS: set[tuple[str, int, str]] = set()
+_WARNED_MHC_KERNELS: set[tuple[str, int, str]] = set()
+_MHC_WARNING_LOCK = threading.Lock()
+logger = init_logger(__name__)
+
+
+def _mhc_kernel_key(tensor: torch.Tensor, name: str) -> tuple[str, int, str]:
+    return (tensor.device.type, tensor.device.index or 0, name)
+
+
+def _is_mhc_oom_error(error: Exception) -> bool:
+    message = str(error).lower()
+    error_name = type(error).__name__.lower()
+    return isinstance(error, MemoryError) or "out of memory" in message or "outofmemoryerror" in error_name
+
+
+def _cache_mhc_kernel_failure(failure_key: tuple[str, int, str], error: Exception) -> None:
+    _FAILED_MHC_KERNELS.add(failure_key)
+    with _MHC_WARNING_LOCK:
+        should_warn = failure_key not in _WARNED_MHC_KERNELS
+        _WARNED_MHC_KERNELS.add(failure_key)
+    if should_warn:
+        logger.warning(
+            "mHC fused kernel failed for key=%s; falling back to native: %s",
+            failure_key,
+            repr(error),
+        )
+
+
+def sinkhorn_knopp(matrix_logits: torch.Tensor, iterations: int, epsilon: float) -> torch.Tensor:
+    matrix = torch.exp(matrix_logits - matrix_logits.amax(dim=(-2, -1), keepdim=True))
+    for _ in range(iterations):
+        matrix = matrix / (matrix.sum(dim=-2, keepdim=True) + epsilon)
+        matrix = matrix / (matrix.sum(dim=-1, keepdim=True) + epsilon)
+    return matrix
+
+
+@triton.jit
+def _mhc_post_residual_kernel(
+    post_ptr,
+    residual_ptr,
+    alpha_post_ptr,
+    bias_post_ptr,
+    alpha_residual_ptr,
+    bias_residual_ptr,
+    post_out_ptr,
+    residual_out_ptr,
+    post_stride,
+    residual_stride,
+    scale,
+    epsilon,
+    iterations: tl.constexpr,
+):
+    token = tl.program_id(0).to(tl.int64)
+    index = tl.arange(0, 4)
+    post = tl.load(post_ptr + token * post_stride + index)
+    ap = tl.load(alpha_post_ptr)
+    bp = tl.load(bias_post_ptr + index)
+    post = 2.0 * tl.sigmoid(ap * scale * post + bp)
+    offsets = index[:, None] * 4 + index[None, :]
+    matrix = tl.load(residual_ptr + token * residual_stride + offsets)
+    ar = tl.load(alpha_residual_ptr)
+    br = tl.load(bias_residual_ptr + offsets)
+    matrix = ar * scale * matrix + br
+    # Native amax propagates any NaN to the whole token's normalization.
+    has_nan = tl.sum(tl.sum((matrix != matrix).to(tl.int32), axis=0), axis=0) != 0
+    maximum = tl.max(tl.max(matrix, axis=0), axis=0)
+    maximum = tl.where(has_nan, float("nan"), maximum)
+    matrix = tl.exp(matrix - maximum)
+    for _ in tl.static_range(iterations):
+        matrix = matrix / (tl.sum(matrix, axis=0)[None, :] + epsilon)
+        matrix = matrix / (tl.sum(matrix, axis=1)[:, None] + epsilon)
+    tl.store(post_out_ptr + token * 4 + index, post)
+    tl.store(residual_out_ptr + token * 16 + offsets, matrix)
+
+
+@triton.jit
+def _mhc_mix_kernel(streams_ptr, branch_ptr, post_ptr, matrix_ptr, out_ptr, hidden, block_size: tl.constexpr):
+    token = tl.program_id(0).to(tl.int64)
+    channels = tl.program_id(1).to(tl.int64) * block_size + tl.arange(0, block_size)
+    valid = channels < hidden
+    stream_ids = tl.arange(0, 4)
+    mixed = tl.zeros((4, block_size), tl.float32)
+    for j in tl.static_range(4):
+        x = tl.load(streams_ptr + (token * 4 + j) * hidden + channels, mask=valid, other=0).to(tl.float32)
+        weight = tl.load(matrix_ptr + token * 16 + stream_ids * 4 + j).to(tl.float32)
+        mixed += weight[:, None] * x[None, :]
+    branch = tl.load(branch_ptr + token * hidden + channels, mask=valid, other=0).to(tl.float32)
+    post = tl.load(post_ptr + token * 4 + stream_ids).to(tl.float32)
+    branch = post[:, None] * branch[None, :]
+    # Preserve native einsum materialization before the final addition.
+    mixed = mixed.to(out_ptr.dtype.element_ty).to(tl.float32)
+    branch = branch.to(out_ptr.dtype.element_ty).to(tl.float32)
+    tl.store(
+        out_ptr + (token * 4 + stream_ids[:, None]) * hidden + channels[None, :], mixed + branch, mask=valid[None, :]
+    )
+
+
+def _native_required(tensors: tuple[torch.Tensor, ...]) -> bool:
+    return (
+        torch.compiler.is_compiling()
+        or tensors[0].device.type not in ("cuda", "musa")
+        or (torch.is_grad_enabled() and any(t.requires_grad for t in tensors))
+    )
+
+
+class MHCPostResidual(CustomOp):
+    """Prepare post coefficients and the residual matrix with FP32 arithmetic."""
+
+    @staticmethod
+    def forward_native(
+        post_logits: torch.Tensor,
+        residual_logits: torch.Tensor,
+        alpha_post: torch.Tensor,
+        bias_post: torch.Tensor,
+        alpha_residual: torch.Tensor,
+        bias_residual: torch.Tensor,
+        *,
+        scale: float,
+        iterations: int = 20,
+        epsilon: float = 1e-12,
+        out_dtype: torch.dtype,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        post = 2.0 * torch.sigmoid(alpha_post * scale * post_logits + bias_post.unsqueeze(0))
+        residual = sinkhorn_knopp(
+            alpha_residual * scale * residual_logits.float() + bias_residual.unsqueeze(0).float(),
+            iterations,
+            epsilon,
+        )
+        return post.to(out_dtype), residual.to(out_dtype)
+
+    def forward_cuda(
+        self,
+        post_logits: torch.Tensor,
+        residual_logits: torch.Tensor,
+        alpha_post: torch.Tensor,
+        bias_post: torch.Tensor,
+        alpha_residual: torch.Tensor,
+        bias_residual: torch.Tensor,
+        *,
+        scale: float,
+        iterations: int = 20,
+        epsilon: float = 1e-12,
+        out_dtype: torch.dtype,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        tensors = (post_logits, residual_logits, alpha_post, bias_post, alpha_residual, bias_residual)
+        if _native_required(tensors):
+            return self.forward_native(
+                *tensors, scale=scale, iterations=iterations, epsilon=epsilon, out_dtype=out_dtype
+            )
+        failure_key = _mhc_kernel_key(post_logits, "post_residual")
+        if (
+            not HAS_TRITON
+            or failure_key in _FAILED_MHC_KERNELS
+            or not (
+                post_logits.ndim == 2
+                and post_logits.shape[1] == 4
+                and residual_logits.shape == (post_logits.shape[0], 4, 4)
+                and post_logits.stride(1) == 1
+                and residual_logits.stride()[1:] == (4, 1)
+                and alpha_post.numel() == 1
+                and alpha_residual.numel() == 1
+                and alpha_post.ndim <= 1
+                and alpha_residual.ndim <= 1
+                and bias_post.shape == (4,)
+                and bias_residual.shape == (4, 4)
+                and all(t.is_contiguous() for t in tensors[2:])
+                and all(t.dtype == torch.float32 and t.device == post_logits.device for t in tensors)
+                and out_dtype in (torch.float32, torch.bfloat16, torch.float16)
+                and isinstance(scale, (int, float))
+                and isinstance(epsilon, (int, float))
+                and isinstance(iterations, int)
+                and 0 <= iterations <= 64
+            )
+        ):
+            return self.forward_native(
+                *tensors, scale=scale, iterations=iterations, epsilon=epsilon, out_dtype=out_dtype
+            )
+        tokens = post_logits.shape[0]
+        post_out = torch.empty((tokens, 4), device=post_logits.device, dtype=out_dtype)
+        residual_out = torch.empty((tokens, 4, 4), device=post_logits.device, dtype=out_dtype)
+        if tokens:
+            try:
+                _mhc_post_residual_kernel[(tokens,)](
+                    *tensors,
+                    post_out,
+                    residual_out,
+                    post_logits.stride(0),
+                    residual_logits.stride(0),
+                    scale,
+                    epsilon,
+                    iterations=iterations,
+                    num_warps=1,
+                    enable_fp_fusion=False,
+                )
+            except Exception as error:
+                if _is_mhc_oom_error(error):
+                    raise
+                _cache_mhc_kernel_failure(failure_key, error)
+                return self.forward_native(
+                    *tensors, scale=scale, iterations=iterations, epsilon=epsilon, out_dtype=out_dtype
+                )
+        return post_out, residual_out
+
+    forward_npu = forward_native
+
+
+class MHCMix(CustomOp):
+    """Mix four streams and add a separately rounded branch term."""
+
+    @staticmethod
+    def forward_native(
+        streams: torch.Tensor,
+        branch_output: torch.Tensor,
+        post_coefficients: torch.Tensor,
+        residual_matrix: torch.Tensor,
+    ) -> torch.Tensor:
+        branch = torch.einsum("tn,tc->tnc", post_coefficients, branch_output)
+        mixed = torch.einsum("tij,tjc->tic", residual_matrix, streams)
+        return mixed + branch
+
+    def forward_cuda(
+        self,
+        streams: torch.Tensor,
+        branch_output: torch.Tensor,
+        post_coefficients: torch.Tensor,
+        residual_matrix: torch.Tensor,
+    ) -> torch.Tensor:
+        tensors = (streams, branch_output, post_coefficients, residual_matrix)
+        if _native_required(tensors):
+            return self.forward_native(*tensors)
+        failure_key = _mhc_kernel_key(streams, "mix")
+        if (
+            not HAS_TRITON
+            or failure_key in _FAILED_MHC_KERNELS
+            or not (
+                streams.ndim == 3
+                and streams.shape[1] == 4
+                and branch_output.shape == (streams.shape[0], streams.shape[2])
+                and post_coefficients.shape == streams.shape[:2]
+                and residual_matrix.shape == (streams.shape[0], 4, 4)
+                and all(t.is_contiguous() and t.device == streams.device and t.dtype == streams.dtype for t in tensors)
+                and streams.dtype in (torch.float32, torch.bfloat16, torch.float16)
+            )
+        ):
+            return self.forward_native(*tensors)
+        output = torch.empty_like(streams)
+        if output.numel():
+            try:
+                _mhc_mix_kernel[(streams.shape[0], triton.cdiv(streams.shape[2], 256))](
+                    *tensors,
+                    output,
+                    streams.shape[2],
+                    block_size=256,
+                    num_warps=4,
+                    enable_fp_fusion=False,
+                )
+            except Exception as error:
+                if _is_mhc_oom_error(error):
+                    raise
+                _cache_mhc_kernel_failure(failure_key, error)
+                return self.forward_native(*tensors)
+        return output
+
+    forward_npu = forward_native

vllm_omni/diffusion/models/boogu_image/boogu_image_transformer.py

@@ -38,13 +38,22 @@ from vllm.triton_utils import HAS_TRITON
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.sp_plan import (
+    SequenceParallelInput,
+    SequenceParallelOutput,
+)
 from vllm_omni.diffusion.layers.fused_qk_norm_rope import (
     _fused_cuda_supported,
     fused_qk_norm_rope,
     fused_qk_norm_rope_min_tokens,
 )
 from vllm_omni.diffusion.layers.norm import RMSNorm
-from vllm_omni.diffusion.models.utils import make_attention_mask
+from vllm_omni.diffusion.models.boogu_image.sp_layout import (
+    RotaryEmbedding,
+    ShardLayout,
+    pack_local_rotary,
+    rank_concat_mask_or_none,
+)
 from vllm_omni.platforms import current_omni_platform
 
 if TYPE_CHECKING:
@@ -52,10 +61,6 @@ if TYPE_CHECKING:
 
 logger = init_logger(__name__)
 
-# (cos, sin) — optionally (cos, sin, packed_table) where packed_table is the
-# fp32 [tokens, head_dim] = [cos(theta) | sin(theta)] layout the fused
-# qk-norm+RoPE op consumes; the eager path ignores the third element.
-RotaryEmbedding = tuple[torch.Tensor, ...]
 RotaryFrequencyTables = list[RotaryEmbedding]
 
 
@@ -540,12 +545,16 @@ def _concat_instruction_image_features(
     assert len(img_tensors) == len(instruct_tensors)
 
     batch_size = img_tensors[0].shape[0]
-    max_seq_len = max(seq_lengths)
+    joint_capacity = instruct_tensors[0].shape[1] + img_tensors[0].shape[1]
 
     concatenated_list = []
     for img_tensor, instruct_tensor in zip(img_tensors, instruct_tensors):
         feature_dim = img_tensor.shape[-1]
-        concatenated = img_tensor.new_zeros(batch_size, max_seq_len, feature_dim)
+        concatenated = img_tensor.new_zeros(
+            batch_size,
+            joint_capacity,
+            feature_dim,
+        )
         for i, (encoder_seq_len, seq_len) in enumerate(zip(encoder_seq_lengths, seq_lengths)):
             concatenated[i, :encoder_seq_len] = instruct_tensor[i, :encoder_seq_len]
             concatenated[i, encoder_seq_len:seq_len] = img_tensor[i, : seq_len - encoder_seq_len]
@@ -558,16 +567,24 @@ def _split_instruction_image_features(
     hidden_states: torch.Tensor,
     encoder_seq_lengths: list[int],
     seq_lengths: list[int],
+    *,
+    instruct_capacity: int,
+    img_capacity: int,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Unpack a joint sequence back into (instruction, image) streams."""
     batch_size = hidden_states.shape[0]
     feature_dim = hidden_states.shape[-1]
 
-    max_instruct_len = max(encoder_seq_lengths)
-    max_img_len = max(seq_len - encoder_seq_len for seq_len, encoder_seq_len in zip(seq_lengths, encoder_seq_lengths))
-
-    instruct_hidden_states = hidden_states.new_zeros(batch_size, max_instruct_len, feature_dim)
-    img_hidden_states = hidden_states.new_zeros(batch_size, max_img_len, feature_dim)
+    instruct_hidden_states = hidden_states.new_zeros(
+        batch_size,
+        instruct_capacity,
+        feature_dim,
+    )
+    img_hidden_states = hidden_states.new_zeros(
+        batch_size,
+        img_capacity,
+        feature_dim,
+    )
 
     for i, (encoder_seq_len, seq_len) in enumerate(zip(encoder_seq_lengths, seq_lengths)):
         img_len = seq_len - encoder_seq_len
@@ -772,7 +789,11 @@ class BooguImageJointAttention(nn.Module):
         attn_output = attn_output.flatten(2, 3).to(dtype)
 
         instruct_attn_out, img_attn_out = _split_instruction_image_features(
-            attn_output, encoder_seq_lengths, seq_lengths
+            attn_output,
+            encoder_seq_lengths,
+            seq_lengths,
+            instruct_capacity=instruct_hidden_states.shape[1],
+            img_capacity=img_hidden_states.shape[1],
         )
         instruct_projected, _ = self.instruct_out(instruct_attn_out)
         img_projected, _ = self.img_out(img_attn_out)
@@ -1139,6 +1160,20 @@ _BOOGU_STACKED_PARAMS_MAPPING = (
 )
 
 
+class _BooguImageSPBoundary(nn.Module):
+    """Identity module used by framework SP split hooks."""
+
+    def forward(self, *tensors: torch.Tensor) -> tuple[torch.Tensor, ...]:
+        return tensors
+
+
+class _BooguImageSPGather(nn.Module):
+    """Identity module used by the framework SP gather hook."""
+
+    def forward(self, tensor: torch.Tensor) -> torch.Tensor:
+        return tensor
+
+
 class BooguImageTransformer2DModel(nn.Module):
     """Boogu-Image transformer with mixed stream topology.
 
@@ -1156,6 +1191,46 @@ class BooguImageTransformer2DModel(nn.Module):
         "BooguImageSingleStreamTransformerBlock",
     ]
     _layerwise_offload_blocks_attrs = ["single_stream_layers", "double_stream_layers"]
+    # Boundary outputs are hidden states, masks, and RoPE; each triplet is in
+    # image/reference/context order so tensors sharing a sequence reuse metadata.
+    # Twelve entries: 3 hidden states (dim=3), 3 masks (dim=2), then the
+    # rotary cos/sin pairs for image/reference/context (dim=3 each). Real-valued
+    # RoPE returns tuples, so each rotary embedding contributes two tensors.
+    _sp_plan = {
+        "sp_shard_boundary": {
+            index: SequenceParallelInput(
+                split_dim=1,
+                expected_dims=dims,
+                split_output=True,
+                auto_pad=True,
+                shard_group=shard_group,
+            )
+            for index, (dims, shard_group) in enumerate(
+                zip(
+                    (3, 3, 3, 2, 2, 2, 3, 3, 3, 3, 3, 3),
+                    (
+                        "image",
+                        "reference",
+                        "context",
+                        "image",
+                        "reference",
+                        "context",
+                        "image",
+                        "image",
+                        "reference",
+                        "reference",
+                        "context",
+                        "context",
+                    ),
+                )
+            )
+        },
+        "sp_gather_boundary": SequenceParallelOutput(
+            gather_dim=1,
+            expected_dims=3,
+            shard_group="image",
+        ),
+    }
 
     def __init__(
         self,
@@ -1226,6 +1301,8 @@ class BooguImageTransformer2DModel(nn.Module):
             axes_lens=axes_lens,
             patch_size=patch_size,
         )
+        self.sp_shard_boundary = _BooguImageSPBoundary()
+        self.sp_gather_boundary = _BooguImageSPGather()
 
         self.x_embedder = nn.Linear(
             in_features=patch_size * patch_size * in_channels,
@@ -1472,6 +1549,7 @@ class BooguImageTransformer2DModel(nn.Module):
         l_effective_ref_img_len: list[list[int]],
         l_effective_img_len: list[int],
         temb: torch.Tensor,
+        per_rank_ref_lengths: list[list[list[int]]] | None = None,
     ):
         """Embed image patches and run the refiner blocks.
 
@@ -1481,9 +1559,7 @@ class BooguImageTransformer2DModel(nn.Module):
         avoiding a degenerate zero-length attention.
         """
         batch_size = len(hidden_states)
-        max_combined_img_len = max(
-            img_len + sum(ref_img_len) for img_len, ref_img_len in zip(l_effective_img_len, l_effective_ref_img_len)
-        )
+        max_combined_img_len = hidden_states.shape[1] + ref_image_hidden_states.shape[1]
 
         hidden_states = self.x_embedder(hidden_states)
         ref_image_hidden_states = self.ref_image_patch_embedder(ref_image_hidden_states)
@@ -1502,6 +1578,14 @@ class BooguImageTransformer2DModel(nn.Module):
         flat_l_effective_ref_img_len = list(itertools.chain(*l_effective_ref_img_len))
         num_ref_images = len(flat_l_effective_ref_img_len)
         max_ref_img_len = max(flat_l_effective_ref_img_len)
+        sp_sharded = per_rank_ref_lengths is not None and len(per_rank_ref_lengths) > 1
+        if sp_sharded:
+            # Keep the packed reference batch at the padded local capacity so
+            # every SP rank retains the equal-length contract.
+            max_ref_img_len = max(
+                max_ref_img_len,
+                padded_ref_img_mask.shape[1],
+            )
 
         if max_ref_img_len > 0:
             batch_ref_img_mask = ref_image_hidden_states.new_zeros(num_ref_images, max_ref_img_len, dtype=torch.bool)
@@ -1544,9 +1628,27 @@ class BooguImageTransformer2DModel(nn.Module):
                 # table; the forward-level tuples do not flow into this batch.
                 batch_ref_img_rotary_emb = _with_packed_rope_table(batch_ref_img_rotary_emb)
 
+            if sp_sharded:
+                # Ulysses will concatenate each rank's packed reference batch,
+                # so the mask must cover all ranks' segments.
+                ref_attention_mask = rank_concat_mask_or_none(
+                    [list(itertools.chain(*rank_lengths)) for rank_lengths in per_rank_ref_lengths],
+                    max_ref_img_len,
+                    like=batch_ref_img_mask,
+                )
+            else:
+                ref_attention_mask = (
+                    batch_ref_img_mask
+                    if any(length != max_ref_img_len for length in flat_l_effective_ref_img_len)
+                    else None
+                )
+
             for layer in self.ref_image_refiner:
                 batch_ref_image_hidden_states = layer(
-                    batch_ref_image_hidden_states, batch_ref_img_mask, batch_ref_img_rotary_emb, batch_temb
+                    batch_ref_image_hidden_states,
+                    ref_attention_mask,
+                    batch_ref_img_rotary_emb,
+                    batch_temb,
                 )
 
             # Restore reference-image sequence layout.
@@ -1614,11 +1716,11 @@ class BooguImageTransformer2DModel(nn.Module):
             context_rotary_emb,
             ref_img_rotary_emb,
             noise_rotary_emb,
-            rotary_emb,
-            encoder_seq_lengths,
-            seq_lengths,
-            combined_img_rotary_emb,
-            combined_img_seq_lengths,
+            _,
+            global_encoder_seq_lengths,
+            _,
+            _,
+            _,
         ) = self.rope_embedder(
             freqs_real,
             instruction_attention_mask,
@@ -1629,10 +1731,91 @@ class BooguImageTransformer2DModel(nn.Module):
             device,
         )
 
+        # RoPE is real-valued (cos, sin) tuples; unpack them to shard cos and
+        # sin independently through the split hook, then re-pair below.
+        noise_cos, noise_sin = noise_rotary_emb
+        ref_cos, ref_sin = ref_img_rotary_emb
+        context_cos, context_sin = context_rotary_emb
+        (
+            hidden_states,
+            ref_image_hidden_states,
+            instruction_hidden_states,
+            img_mask,
+            ref_img_mask,
+            instruction_attention_mask,
+            noise_cos,
+            noise_sin,
+            ref_cos,
+            ref_sin,
+            context_cos,
+            context_sin,
+        ) = self.sp_shard_boundary(
+            hidden_states,
+            ref_image_hidden_states,
+            instruction_hidden_states,
+            img_mask,
+            ref_img_mask,
+            instruction_attention_mask,
+            noise_cos,
+            noise_sin,
+            ref_cos,
+            ref_sin,
+            context_cos,
+            context_sin,
+        )
+        noise_rotary_emb = (noise_cos, noise_sin)
+        ref_img_rotary_emb = (ref_cos, ref_sin)
+        context_rotary_emb = (context_cos, context_sin)
+
+        image_layout = ShardLayout.resolve("image", local_seq_len=img_mask.shape[1])
+        reference_layout = ShardLayout.resolve("reference", local_seq_len=ref_img_mask.shape[1])
+        context_layout = ShardLayout.resolve("context", local_seq_len=instruction_attention_mask.shape[1])
+        sp_world_size = image_layout.world_size
+        sp_rank = image_layout.rank
+
+        context_mask_required = context_layout.padding_size > 0 or any(
+            length != context_layout.original_seq_len for length in global_encoder_seq_lengths
+        )
+        noise_mask_required = image_layout.padding_size > 0 or any(
+            length != image_layout.original_seq_len for length in l_effective_img_len
+        )
+        ref_mask_required = reference_layout.padding_size > 0 or any(
+            sum(lengths) != reference_layout.original_seq_len for lengths in l_effective_ref_img_len
+        )
+
+        global_img_lengths = list(l_effective_img_len)
+        # Ulysses concatenates the rank-local sequences after its all-to-all, so
+        # attention masks must describe that rank-concatenated layout. Every
+        # rank's shard bounds are pure arithmetic over the global lengths, so
+        # each rank derives all ranks' segments without any communication.
+        ranks = range(sp_world_size)
+        per_rank_img_lengths = [image_layout.valid_lengths(global_img_lengths, rank=r) for r in ranks]
+        per_rank_ref_lengths = [reference_layout.segment_lengths(l_effective_ref_img_len, rank=r) for r in ranks]
+        per_rank_encoder_lengths = [context_layout.valid_lengths(global_encoder_seq_lengths, rank=r) for r in ranks]
+
+        local_img_lengths = per_rank_img_lengths[sp_rank]
+        local_ref_lengths = per_rank_ref_lengths[sp_rank]
+        local_encoder_lengths = per_rank_encoder_lengths[sp_rank]
+        (
+            rotary_emb,
+            combined_img_rotary_emb,
+            seq_lengths,
+            combined_img_seq_lengths,
+        ) = pack_local_rotary(
+            context_rotary_emb,
+            ref_img_rotary_emb,
+            noise_rotary_emb,
+            local_encoder_lengths,
+            local_ref_lengths,
+            local_img_lengths,
+        )
+
         if _FUSED_QK_NORM_ROPE:
             # One packed [cos|sin] table per rotary embedding that reaches an
             # attention, built once per forward; the tuples grow a third
             # element that the fused path consumes and the eager path ignores.
+            # Under SP, cos/sin are already the rank-local sharded tensors, so
+            # the packed table matches the fused kernel's local Q/K geometry.
             # (ref_img_rotary_emb is packed inside img_patch_embed_and_refine,
             # on the rebuilt per-reference-image batch tuple.)
             context_rotary_emb = _with_packed_rope_table(context_rotary_emb)
@@ -1640,32 +1823,68 @@ class BooguImageTransformer2DModel(nn.Module):
             rotary_emb = _with_packed_rope_table(rotary_emb)
             combined_img_rotary_emb = _with_packed_rope_table(combined_img_rotary_emb)
 
+        context_attention_mask = rank_concat_mask_or_none(
+            per_rank_encoder_lengths,
+            context_layout.local_seq_len,
+            like=instruction_attention_mask,
+            required=context_mask_required,
+        )
+        noise_attention_mask = rank_concat_mask_or_none(
+            per_rank_img_lengths,
+            image_layout.local_seq_len,
+            like=img_mask,
+            required=noise_mask_required,
+        )
+
         # Context refinement.
         for layer in self.context_refiner:
-            instruction_hidden_states = layer(instruction_hidden_states, instruction_attention_mask, context_rotary_emb)
+            instruction_hidden_states = layer(
+                instruction_hidden_states,
+                context_attention_mask,
+                context_rotary_emb,
+            )
 
         # Image patch embedding and refinement.
         combined_img_hidden_states = self.img_patch_embed_and_refine(
             hidden_states,
             ref_image_hidden_states,
-            img_mask,
+            noise_attention_mask,
             ref_img_mask,
             noise_rotary_emb,
             ref_img_rotary_emb,
-            l_effective_ref_img_len,
-            l_effective_img_len,
+            local_ref_lengths,
+            local_img_lengths,
             temb,
+            per_rank_ref_lengths=per_rank_ref_lengths,
         )
 
         instruct_hidden_states = instruction_hidden_states
         img_hidden_states = combined_img_hidden_states
 
-        # Joint mask for [instruct + image].
-        joint_attention_mask = make_attention_mask(hidden_states, seq_lengths)
+        # Joint mask for [instruct + image], over the rank-concatenated sequence.
+        per_rank_combined_img_lengths = [
+            [sum(ref_lengths) + img_length for ref_lengths, img_length in zip(ref_per_sample, img_per_sample)]
+            for ref_per_sample, img_per_sample in zip(per_rank_ref_lengths, per_rank_img_lengths)
+        ]
+        per_rank_seq_lengths = [
+            [encoder_length + combined_length for encoder_length, combined_length in zip(enc_per_sample, comb)]
+            for enc_per_sample, comb in zip(per_rank_encoder_lengths, per_rank_combined_img_lengths)
+        ]
+        joint_attention_mask = rank_concat_mask_or_none(
+            per_rank_seq_lengths,
+            rotary_emb[0].shape[1],
+            like=hidden_states,
+            required=(context_mask_required or noise_mask_required or ref_mask_required),
+        )
 
         # Dual-stream (double-stream) stage.
         if self.num_double_stream_layers > 0:
-            img_attention_mask = make_attention_mask(hidden_states, combined_img_seq_lengths)
+            img_attention_mask = rank_concat_mask_or_none(
+                per_rank_combined_img_lengths,
+                combined_img_rotary_emb[0].shape[1],
+                like=hidden_states,
+                required=noise_mask_required or ref_mask_required,
+            )
 
             for layer in self.double_stream_layers:
                 img_hidden_states, instruct_hidden_states = layer(
@@ -1676,13 +1895,17 @@ class BooguImageTransformer2DModel(nn.Module):
                     combined_img_rotary_emb,
                     rotary_emb,
                     temb,
-                    encoder_seq_lengths,
+                    local_encoder_lengths,
                     seq_lengths,
                 )
 
         # Fuse streams to joint sequence.
-        joint_hidden_states = hidden_states.new_zeros(batch_size, max(seq_lengths), self.hidden_size)
-        for i, (encoder_seq_len, seq_len) in enumerate(zip(encoder_seq_lengths, seq_lengths)):
+        joint_hidden_states = hidden_states.new_zeros(
+            batch_size,
+            rotary_emb[0].shape[1],
+            self.hidden_size,
+        )
+        for i, (encoder_seq_len, seq_len) in enumerate(zip(local_encoder_lengths, seq_lengths)):
             joint_hidden_states[i, :encoder_seq_len] = instruct_hidden_states[i, :encoder_seq_len]
             joint_hidden_states[i, encoder_seq_len:seq_len] = img_hidden_states[i, : seq_len - encoder_seq_len]
 
@@ -1694,12 +1917,23 @@ class BooguImageTransformer2DModel(nn.Module):
         # Output projection.
         hidden_states = self.norm_out(hidden_states, temb)
 
+        # Extract this rank's noise-image shard and gather it once.
+        local_img_output = hidden_states.new_zeros(
+            batch_size,
+            img_mask.shape[1],
+            hidden_states.shape[-1],
+        )
+        for i, (img_len, seq_len) in enumerate(zip(local_img_lengths, seq_lengths)):
+            local_img_output[i, :img_len] = hidden_states[i, seq_len - img_len : seq_len]
+
+        image_tokens = self.sp_gather_boundary(local_img_output)
+
         # Reshape back to image format.
         p = self.patch_size
         output = []
-        for i, (img_size, img_len, seq_len) in enumerate(zip(img_sizes, l_effective_img_len, seq_lengths)):
+        for i, (img_size, img_len) in enumerate(zip(img_sizes, global_img_lengths)):
             height, width = img_size
-            img_tokens = hidden_states[i][seq_len - img_len : seq_len]
+            img_tokens = image_tokens[i, :img_len]
             img_output = rearrange(
                 img_tokens,
                 "(h w) (p1 p2 c) -> c (h p1) (w p2)",

vllm_omni/diffusion/models/boogu_image/pipeline_boogu_image.py

@@ -62,6 +62,7 @@ from vllm_omni.diffusion.models.utils import create_transformers_model_with_vllm
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, split_diffusion_output_by_request
 from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
+from vllm_omni.platforms import current_omni_platform
 from vllm_omni.quantization.component_config import resolve_component_quant_config
 
 logger = init_logger(__name__)
@@ -382,7 +383,36 @@ class BooguImagePipeline(CFGParallelMixin, nn.Module, ProgressBarMixin, Supports
         if parallel_config.tensor_parallel_size > 1:
             raise NotImplementedError("Tensor parallelism is not supported by BooguImagePipeline.")
         if (parallel_config.sequence_parallel_size or 1) > 1:
-            raise NotImplementedError("Sequence parallelism is not supported by BooguImagePipeline.")
+            if parallel_config.cfg_parallel_size > 1:
+                # SP × CFG-Parallel is unvalidated: the combined path runs
+                # without crashing but drifts far from the SP=1 baseline
+                # (~19 dB PSNR at 4 steps), so guidance semantics differ.
+                # Remove this guard only alongside a parity test.
+                raise NotImplementedError(
+                    "BooguImagePipeline: sequence parallelism combined with CFG parallelism is not validated."
+                )
+            if parallel_config.ring_degree > 1 or getattr(parallel_config, "allgather_degree", 1) > 1:
+                raise NotImplementedError(
+                    "BooguImagePipeline currently supports sequence parallelism through Ulysses only."
+                )
+            if not current_omni_platform.is_cuda():
+                raise NotImplementedError("BooguImagePipeline sequence parallelism currently requires CUDA.")
+            if parallel_config.ulysses_mode == "strict":
+                # Subscript rather than get(): a missing key means the config
+                # schema has drifted and we would silently validate against a
+                # stale default, not the actual head count.
+                tf_params = self.od_config.tf_model_config.params
+                num_heads = tf_params["num_attention_heads"]
+                num_kv_heads = tf_params["num_kv_heads"]
+                if (
+                    num_heads % parallel_config.ulysses_degree != 0
+                    or num_kv_heads % parallel_config.ulysses_degree != 0
+                ):
+                    raise ValueError(
+                        "BooguImagePipeline GQA heads are not divisible by the "
+                        "configured Ulysses degree; set "
+                        "parallel_config.ulysses_mode='advanced_uaa'."
+                    )
         if parallel_config.use_hsdp:
             raise NotImplementedError("HSDP is not supported by BooguImagePipeline.")
         if self.od_config.cache_backend not in (None, "", "none"):

vllm_omni/diffusion/models/boogu_image/sp_layout.py

@@ -0,0 +1,188 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Shard-layout arithmetic for BOOGU Image sequence parallelism.
+
+BOOGU splits three independent sequences (noise image, reference images, and
+instruction context) at one SP boundary. The framework records each sequence's
+pre-padding global length under a ``shard_group`` key; everything else -- how
+that global length maps onto a given rank's shard -- is arithmetic, so every
+rank can derive the layout of *all* ranks without communicating.
+
+That property is what lets the model build attention masks for the global,
+rank-concatenated sequence that Ulysses produces after its all-to-all, instead
+of all-gathering rank-local masks inside the attention layer.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+
+from vllm_omni.diffusion.forward_context import get_sp_shard_original_seq_len
+
+# (cos, sin) -- optionally (cos, sin, packed_table) where packed_table is the
+# fp32 [tokens, head_dim] = [cos(theta) | sin(theta)] layout the fused
+# qk-norm+RoPE op consumes; the eager path ignores the third element.
+RotaryEmbedding = tuple[torch.Tensor, ...]
+
+
+@dataclass(frozen=True, slots=True)
+class ShardLayout:
+    """How one globally padded sequence maps onto the SP ranks."""
+
+    original_seq_len: int
+    world_size: int
+    rank: int
+
+    @classmethod
+    def resolve(cls, shard_group: str, *, local_seq_len: int) -> ShardLayout:
+        """Read a boundary's layout from the ForwardContext.
+
+        Falls back to an unsharded layout (world_size=1) when SP is off or the
+        tensor was never split, in which case `local_seq_len` is the whole
+        sequence.
+        """
+        from vllm_omni.diffusion.distributed.parallel_state import (
+            get_sequence_parallel_rank,
+            get_sequence_parallel_world_size,
+        )
+
+        original_seq_len = get_sp_shard_original_seq_len(shard_group)
+        if original_seq_len is None:
+            return cls(original_seq_len=local_seq_len, world_size=1, rank=0)
+        return cls(
+            original_seq_len=original_seq_len,
+            world_size=get_sequence_parallel_world_size(),
+            rank=get_sequence_parallel_rank(),
+        )
+
+    @property
+    def padded_seq_len(self) -> int:
+        ws = self.world_size
+        return ((self.original_seq_len + ws - 1) // ws) * ws
+
+    @property
+    def local_seq_len(self) -> int:
+        return self.padded_seq_len // self.world_size
+
+    @property
+    def padding_size(self) -> int:
+        return self.padded_seq_len - self.original_seq_len
+
+    def bounds(self, rank: int) -> tuple[int, int]:
+        """Half-open [start, end) span of the global sequence owned by `rank`."""
+        start = rank * self.local_seq_len
+        return start, start + self.local_seq_len
+
+    def valid_lengths(self, global_lengths: list[int], *, rank: int) -> list[int]:
+        """Per-sample valid prefix lengths clipped to `rank`'s shard."""
+        start, end = self.bounds(rank)
+        return [max(0, min(int(length), end) - start) for length in global_lengths]
+
+    def segment_lengths(self, global_segments: list[list[int]], *, rank: int) -> list[list[int]]:
+        """Per-sample contiguous segments intersected with `rank`'s shard."""
+        start, end = self.bounds(rank)
+        local: list[list[int]] = []
+        for sample in global_segments:
+            sample_local: list[int] = []
+            offset = 0
+            for length in sample:
+                segment_end = offset + int(length)
+                sample_local.append(max(0, min(segment_end, end) - max(offset, start)))
+                offset = segment_end
+            local.append(sample_local)
+        return local
+
+
+def rank_concat_mask_or_none(
+    per_rank_lengths: list[list[int]],
+    capacity: int,
+    *,
+    like: torch.Tensor,
+    required: bool = False,
+) -> torch.Tensor | None:
+    """Mask over the rank-concatenated sequence Ulysses builds post all-to-all.
+
+    ``per_rank_lengths[r][i]`` is sample ``i``'s valid length inside rank ``r``'s
+    shard, and every shard has the same ``capacity`` (guaranteed by auto_pad),
+    so the global sequence is ``world_size * capacity`` long. Returns ``None``
+    for an all-valid mask unless a padding contract requires one.
+    """
+    if not required and all(int(length) == capacity for lengths in per_rank_lengths for length in lengths):
+        return None
+
+    batch_size = len(per_rank_lengths[0])
+    mask = like.new_zeros(batch_size, len(per_rank_lengths) * capacity, dtype=torch.bool)
+    for rank, lengths in enumerate(per_rank_lengths):
+        base = rank * capacity
+        for i, length in enumerate(lengths):
+            mask[i, base : base + int(length)] = True
+    return mask
+
+
+def pack_local_rotary(
+    context_rotary_emb: RotaryEmbedding,
+    ref_img_rotary_emb: RotaryEmbedding,
+    noise_rotary_emb: RotaryEmbedding,
+    encoder_seq_lengths: list[int],
+    ref_img_seq_lengths: list[list[int]],
+    img_seq_lengths: list[int],
+) -> tuple[RotaryEmbedding, RotaryEmbedding, list[int], list[int]]:
+    """Pack local RoPE as ``[context, references, noise]``.
+
+    Capacities stay equal across ranks; returned lengths exclude padding.
+    Rotary embeddings arrive as ``(cos, sin)`` pairs -- the same packing is
+    applied component-wise so both halves stay index-aligned.
+    """
+    context_cos, context_sin = context_rotary_emb
+    ref_cos, ref_sin = ref_img_rotary_emb
+    noise_cos, noise_sin = noise_rotary_emb
+
+    combined_img_seq_lengths = [
+        sum(ref_lengths) + img_length
+        for ref_lengths, img_length in zip(
+            ref_img_seq_lengths,
+            img_seq_lengths,
+        )
+    ]
+    seq_lengths = [
+        context_length + combined_length
+        for context_length, combined_length in zip(
+            encoder_seq_lengths,
+            combined_img_seq_lengths,
+        )
+    ]
+
+    def _pack_component(
+        context_t: torch.Tensor, ref_t: torch.Tensor, noise_t: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        batch_size = context_t.shape[0]
+        rotary_dim = context_t.shape[-1]
+        combined_img_capacity = ref_t.shape[1] + noise_t.shape[1]
+        combined_img = context_t.new_zeros(batch_size, combined_img_capacity, rotary_dim)
+        rotary = context_t.new_zeros(
+            batch_size,
+            context_t.shape[1] + combined_img_capacity,
+            rotary_dim,
+        )
+        for i, (context_length, ref_lengths, img_length) in enumerate(
+            zip(encoder_seq_lengths, ref_img_seq_lengths, img_seq_lengths)
+        ):
+            ref_length = sum(ref_lengths)
+            combined_length = ref_length + img_length
+            combined_img[i, :ref_length] = ref_t[i, :ref_length]
+            combined_img[i, ref_length:combined_length] = noise_t[i, :img_length]
+            rotary[i, :context_length] = context_t[i, :context_length]
+            rotary[i, context_length : context_length + combined_length] = combined_img[i, :combined_length]
+        return rotary, combined_img
+
+    rotary_cos, combined_img_cos = _pack_component(context_cos, ref_cos, noise_cos)
+    rotary_sin, combined_img_sin = _pack_component(context_sin, ref_sin, noise_sin)
+
+    return (
+        (rotary_cos, rotary_sin),
+        (combined_img_cos, combined_img_sin),
+        seq_lengths,
+        combined_img_seq_lengths,
+    )

vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py

@@ -72,6 +72,7 @@ from vllm_omni.diffusion.models.schedulers.scheduling_flow_unipc_multistep impor
 from vllm_omni.diffusion.offloader.config import OffloadStrategy, resolve_offload
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.utils.video_decode import decode_path_video_frames, is_video_file_path
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 from vllm_omni.entrypoints.openai.video_api_utils import positive_float
 from vllm_omni.experimental.world_models.adapters.state_cosmos3_adapter import (
@@ -465,10 +466,39 @@ def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig):
                     return nested
         return video
 
-    def _video_payload_to_frames(video: Any) -> list[Any]:
+    def _decode_video_file_to_frames(
+        path: str | os.PathLike,
+        *,
+        max_frames: int | None,
+        keep: str,
+    ) -> list[PIL.Image.Image]:
+        rgb_frames = decode_path_video_frames(path, max_frames=max_frames, keep=keep)
+        return [PIL.Image.fromarray(frame).convert("RGB") for frame in rgb_frames]
+
+    def _expand_video_payload_item(
+        item: Any,
+        *,
+        max_frames: int | None,
+        keep: str,
+    ) -> list[Any]:
+        if is_video_file_path(item):
+            return _decode_video_file_to_frames(item, max_frames=max_frames, keep=keep)
+        return [item]
+
+    def _video_payload_to_frames(
+        video: Any,
+        *,
+        max_frames: int | None,
+        keep: str,
+    ) -> list[Any]:
         video = _unwrap_video_payload(video)
+        if is_video_file_path(video):
+            return _decode_video_file_to_frames(video, max_frames=max_frames, keep=keep)
         if isinstance(video, list):
-            return video
+            frames: list[Any] = []
+            for item in video:
+                frames.extend(_expand_video_payload_item(item, max_frames=max_frames, keep=keep))
+            return frames
         if isinstance(video, torch.Tensor):
             tensor = video.detach().cpu()
             if tensor.ndim == 5:
@@ -531,11 +561,31 @@ def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig):
         if "additional_information" not in prompt:
             prompt["additional_information"] = {}
 
+        extra = _extra_args(request)
+        transfer_requested = action_mode is None and has_transfer_hints(extra)
+        condition_frame_indexes_vision = normalize_condition_frame_indexes_vision(
+            extra.get("condition_frame_indexes_vision", prompt.get("condition_frame_indexes_vision"))
+        )
+        condition_video_keep = normalize_condition_video_keep(
+            extra.get("condition_video_keep", prompt.get("condition_video_keep"))
+        )
+
         raw_video_frames: list[Any] | None = None
         transfer_input_fps: float | None = None
         if raw_video is not None:
             transfer_input_fps = _video_payload_fps(raw_video)
-            raw_video_frames = _video_payload_to_frames(raw_video)
+            decode_extra = dict(extra)
+            decode_extra["condition_frame_indexes_vision"] = list(condition_frame_indexes_vision)
+            decode_extra["condition_video_keep"] = condition_video_keep
+            decode_spec = Cosmos3OmniDiffusersPipeline.reference_video_decode_spec(
+                num_frames=getattr(request.sampling_params, "num_frames", None),
+                extra_args=decode_extra,
+            )
+            raw_video_frames = _video_payload_to_frames(
+                raw_video,
+                max_frames=decode_spec.max_frames,
+                keep=decode_spec.keep,
+            )
             if not raw_video_frames:
                 raise TypeError("Cosmos3 video input must be a non-empty list of PIL images or image paths.")
 
@@ -544,8 +594,6 @@ def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig):
             image = _pil_to_rgb(raw_video_frames[0])
         else:
             image = _pil_to_rgb(raw_image)
-        extra = _extra_args(request)
-        transfer_requested = action_mode is None and has_transfer_hints(extra)
 
         # Resolve missing H/W.
         if transfer_requested:
@@ -610,22 +658,13 @@ def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig):
                     dtype=torch.float32,
                 )
             else:
-                condition_frame_indexes_vision = normalize_condition_frame_indexes_vision(
-                    extra.get(
-                        "condition_frame_indexes_vision",
-                        prompt.get("condition_frame_indexes_vision"),
-                    )
-                )
-                keep = normalize_condition_video_keep(
-                    extra.get("condition_video_keep", prompt.get("condition_video_keep"))
-                )
                 max_frames = condition_pixel_frame_count(condition_frame_indexes_vision)
                 prompt["additional_information"]["preprocessed_video"] = _preprocess_condition_video(
                     raw_video_frames,
                     int(target_h),
                     int(target_w),
                     max_frames,
-                    keep,
+                    condition_video_keep,
                 )
                 prompt["additional_information"]["condition_frame_indexes_vision"] = list(
                     condition_frame_indexes_vision

vllm_omni/diffusion/models/cosmos3/transfer.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
 """Cosmos3 transfer inference helpers.
 
 The reference Cosmos Framework transfer path accepts one or more control hints
@@ -25,6 +25,8 @@ import PIL.Image
 import torch
 import torch.nn.functional as F
 
+from vllm_omni.diffusion.utils.video_decode import decode_path_video_frames
+
 TRANSFER_HINT_KEYS: tuple[str, ...] = ("edge", "blur", "depth", "seg", "wsm")
 _TRANSFER_HINT_COMMON_FIELDS = frozenset({"control_path", "control", "control_weight"})
 _TRANSFER_HINT_FIELDS: dict[str, frozenset[str]] = {
@@ -538,23 +540,9 @@ def _path_media_to_uint8_cthw(path: str | Path, max_frames: int | None) -> torch
         array = _pil_to_uint8_rgb(media_path)
         return torch.from_numpy(array).permute(2, 0, 1).unsqueeze(1).contiguous()
 
-    try:
-        import imageio.v3 as iio
-    except ImportError as exc:
-        raise ImportError(
-            "Cosmos3 transfer video control_path loading requires imageio. "
-            "Install imageio[ffmpeg] or provide decoded control frames."
-        ) from exc
-
-    frames: list[torch.Tensor] = []
-    limit = max_frames if max_frames is not None else None
-    for frame in iio.imiter(media_path):
-        frames.append(torch.from_numpy(_pil_to_uint8_rgb(frame)).permute(2, 0, 1))
-        if limit is not None and len(frames) >= int(limit):
-            break
-    if not frames:
-        raise ValueError(f"Cosmos3 transfer control_path produced no frames: {media_path}")
-    return torch.stack(frames, dim=1).contiguous()
+    rgb_frames = decode_path_video_frames(media_path, max_frames=max_frames, keep="first")
+    stacked = [torch.from_numpy(frame).permute(2, 0, 1) for frame in rgb_frames]
+    return torch.stack(stacked, dim=1).contiguous()
 
 
 def media_to_uint8_cthw(value: Any, *, height: int, width: int, max_frames: int | None = None) -> torch.Tensor:

vllm_omni/diffusion/models/ltx2/ltx2_components.py

@@ -30,6 +30,7 @@ from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
 from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
+from vllm_omni.diffusion.offloader.offload_plan import OffloadPlan
 from vllm_omni.transformers_utils.repo_utils import hf_api
 
 if TYPE_CHECKING:
@@ -633,6 +634,13 @@ def _place_aux_components(pipeline: Any) -> None:
         module.to(pipeline.device)
 
 
+def _declare_text_encoder_offload_plan(pipeline: Any) -> None:
+    """Make the text encoder streamable; offloading the DiT alone leaves it resident."""
+    language_model = getattr(getattr(pipeline.text_encoder, "model", None), "language_model", None)
+    if language_model is not None and hasattr(language_model, "layers"):
+        pipeline._offload_plan = OffloadPlan(encoder_block_attrs={"text_encoder": ("model.language_model.layers",)})
+
+
 def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None:
     """Build the common LTX component graph selected by ``component_profile``."""
     profile: LTXComponentProfile = pipeline.component_profile
@@ -680,6 +688,7 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None:
         dtype=dtype,
         revision=revision,
     )
+    _declare_text_encoder_offload_plan(pipeline)
     _install_connector_attention(
         pipeline.connectors,
         preserve_learned_register_mask=profile.preserve_connector_attention_mask,