[Bugfix] Fix MiniCPM-o 4.5 seeded turns and window E2E inputs (#8160)

Signed-off-by: Gao Han <hgaoaf@connect.ust.hk>

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

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/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/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/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/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 \

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/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_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py

@@ -25,6 +25,7 @@ class _FakeDecoder(nn.Module):
         self.total_upsample = total_upsample
         self.decode_calls: list[dict[str, object]] = []
         self.batched_decode_calls: list[dict[str, object]] = []
+        self.batched_decode_codes: list[torch.Tensor] = []
         self.decode_codes: list[torch.Tensor] = []
         self.cudagraph_calls: list[dict[str, int | torch.device]] = []
 
@@ -64,6 +65,7 @@ class _FakeDecoder(nn.Module):
         left_context_size: int = 25,
         max_batch_size: int = 0,
     ) -> list[torch.Tensor]:
+        self.batched_decode_codes.append(codes.detach().cpu().clone())
         self.batched_decode_calls.append(
             {
                 "chunk_size": chunk_size,
@@ -203,6 +205,26 @@ def test_forward_uses_decoder_audio_contract_without_context():
     torch.testing.assert_close(audio, expected)
 
 
+def test_forward_reads_current_model_intermediate_buffer_for_full_payload():
+    """Full-payload sync mode must decode connector codec ids, not placeholders."""
+    model = _make_model()
+    placeholder_ids = torch.zeros(12, dtype=torch.long)
+    payload_codes = torch.arange(12, dtype=torch.long) + 17
+
+    model.forward(
+        input_ids=placeholder_ids,
+        runtime_additional_information=[],
+        model_intermediate_buffer=[
+            {
+                "codes": {"audio": payload_codes},
+                "meta": {"left_context_size": 0},
+            }
+        ],
+    )
+
+    torch.testing.assert_close(model.decoder.batched_decode_codes[-1], payload_codes.reshape(1, _NUM_QUANTIZERS, 6))
+
+
 @pytest.mark.parametrize("skipped_length", [0, 1], ids=["empty", "malformed"])
 def test_stateless_batch_uses_original_request_index_for_left_context(skipped_length):
     model = _make_model()

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

tests/model_tests/diffusion/model_settings.py

@@ -135,6 +135,11 @@ DIFFUSION_TEST_SETTINGS = {
         model="black-forest-labs/FLUX.2-dev",
         builder=diff_model_builders.tiny_flux2_builder,
         supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE],
+        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",
@@ -146,4 +151,13 @@ DIFFUSION_TEST_SETTINGS = {
         builder=diff_model_builders.tiny_qwen_image_edit_plus_builder,
         supported_tasks=[DiffusionTasks.IMAGE_TO_IMAGE],
     ),
+    "StableDiffusion3Pipeline": DiffusionModelTestOpts(
+        model="stabilityai/stable-diffusion-3.5-medium",
+        builder=diff_model_builders.tiny_sd3_builder,
+        supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE],
+        extra_test_groups=[
+            [DiffusionAccs.CACHE_DIT, DiffusionAccs.LAYERWISE_OFFLOAD],
+            [DiffusionAccs.CFG_PARALLEL, DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CPU_OFFLOAD],
+        ],
+    ),
 }

tests/model_tests/diffusion/test_alignment.py

@@ -56,7 +56,6 @@ EXCLUDED_MODELS = [
     "MingImageDiffusionPipeline",
     "MingImageLayeredDiffusionPipeline",
     "InternVLAA1Pipeline",
-    "StableDiffusion3Pipeline",
     "HunyuanImage3ForCausalMM",
     "ErnieImagePipeline",
     "NextStep11Pipeline",

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,

vllm_omni/diffusion/models/ltx2/ltx2_runtime.py

@@ -85,7 +85,7 @@ def _run_ltx_vocoder(vocoder: nn.Module, generated_mel: torch.Tensor) -> torch.T
 
         input_dtype = generated_mel.dtype
         module_dtype = next(vocoder.parameters()).dtype
-        if device_type == "mps":
+        if device_type in ("mps", "xpu", "cpu"):
             if module_dtype != torch.float32:
                 vocoder.float()
             try:

vllm_omni/diffusion/models/magi2/layers.py

@@ -19,8 +19,13 @@ import torch.distributed as dist
 import torch.nn as nn
 import torch.nn.functional as F
 
+from vllm_omni.diffusion.layers.mhc import MHCMix, MHCPostResidual, sinkhorn_knopp
+
 from .parallel import Magi2ParallelGroup, get_magi2_tp_group
 
+_mhc_post_residual = MHCPostResidual()
+_mhc_mix = MHCMix()
+
 
 def swiglu7(
     x: torch.Tensor,
@@ -329,14 +334,6 @@ class ElementWiseFourierEmbed(nn.Module):
 MHCTensorTuple = tuple[torch.Tensor, torch.Tensor, torch.Tensor]
 
 
-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
-
-
 class MHCHandler:
     """Exact four-stream manifold-constrained hyper-connection math."""
 
@@ -397,13 +394,19 @@ class MHCHandler:
     ) -> tuple[torch.Tensor, torch.Tensor]:
         alpha_post, bias_post, post_logits = post
         alpha_residual, bias_residual, residual_logits = residual
-        post_coefficients = 2.0 * torch.sigmoid(alpha_post * self.matmul_scale * post_logits + bias_post.unsqueeze(0))
-        residual_matrix = sinkhorn_knopp(
-            alpha_residual * self.matmul_scale * residual_logits.float() + bias_residual.unsqueeze(0).float(),
-            self.sinkhorn_iterations,
-            self.sinkhorn_epsilon,
+        prepare = _mhc_post_residual if not torch.compiler.is_compiling() else _mhc_post_residual.forward_native
+        return prepare(
+            post_logits,
+            residual_logits,
+            alpha_post,
+            bias_post,
+            alpha_residual,
+            bias_residual,
+            scale=self.matmul_scale,
+            iterations=self.sinkhorn_iterations,
+            epsilon=self.sinkhorn_epsilon,
+            out_dtype=out_dtype,
         )
-        return post_coefficients.to(out_dtype), residual_matrix.to(out_dtype)
 
     def hyper_connect(
         self,
@@ -415,9 +418,8 @@ class MHCHandler:
         self._check_multi(residual_streams)
         if branch_output.ndim != 2 or branch_output.shape[-1] != self.hidden_size:
             raise ValueError("invalid mHC branch-output shape")
-        branch = torch.einsum("tn,tc->tnc", post_coefficients, branch_output)
-        mixed = torch.einsum("tij,tjc->tic", residual_matrix, residual_streams)
-        return mixed + branch
+        mix = _mhc_mix if not torch.compiler.is_compiling() else _mhc_mix.forward_native
+        return mix(residual_streams, branch_output, post_coefficients, residual_matrix)
 
     def _check_multi(self, tensor: torch.Tensor) -> None:
         expected = (self.num_streams, self.hidden_size)

vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py

@@ -631,8 +631,33 @@ class Transformer2DModel(ModelMixin, ConfigMixin):
                 dtype=torch.bool,
                 device=image_hidden_states.device,
             )
-            text_hidden_states = torch.cat([text_hidden_states, image_hidden_states], dim=1)
-            text_attention_mask = torch.cat([text_attention_mask, image_attention_mask], dim=1)
+            # When batch requests have unequal text lengths, right-padded text_hidden_states
+            # followed by concatenated image_hidden_states places padding in the middle of
+            # the row: [valid text, pad, valid image]. Since downstream rope_embedder and
+            # joint_hidden_states prefix-pack using encoder_seq_lengths (sum of valid mask),
+            # middle padding causes image queries to be truncated.
+            # Compact each row so valid tokens are placed contiguously at the start.
+            if batch_size > 1 and not text_attention_mask.all():
+                text_lens = text_attention_mask.sum(dim=-1).tolist()
+                img_lens = image_attention_mask.sum(dim=-1).tolist()
+                max_valid_len = max(t + img for t, img in zip(text_lens, img_lens))
+                compact_hidden_states = text_hidden_states.new_zeros(
+                    batch_size, max_valid_len, text_hidden_states.shape[-1]
+                )
+                compact_attention_mask = torch.zeros((batch_size, max_valid_len), dtype=torch.bool, device=device)
+                for i in range(batch_size):
+                    t_len = text_lens[i]
+                    i_len = img_lens[i]
+                    if t_len > 0:
+                        compact_hidden_states[i, :t_len] = text_hidden_states[i, :t_len]
+                    if i_len > 0:
+                        compact_hidden_states[i, t_len : t_len + i_len] = image_hidden_states[i, :i_len]
+                    compact_attention_mask[i, : t_len + i_len] = True
+                text_hidden_states = compact_hidden_states
+                text_attention_mask = compact_attention_mask
+            else:
+                text_hidden_states = torch.cat([text_hidden_states, image_hidden_states], dim=1)
+                text_attention_mask = torch.cat([text_attention_mask, image_attention_mask], dim=1)
 
         img_tokens = rearrange(hidden_states, "b c (h p1) (w p2) -> b (h w) (p1 p2 c)", p1=p, p2=p)
         img_tokens = self.x_embedder(img_tokens)

vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py

@@ -24,7 +24,9 @@ from vllm_omni.diffusion.data import DiffusionCacheConfig, DiffusionOutput, Omni
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
+from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.transformers_utils.configs.mammoth_moda2 import Mammothmoda2Config
 
 from .mammothmoda2_dit_model import SimpleQFormerImageRefiner, Transformer2DModel
@@ -75,8 +77,235 @@ def _root_weight_source(
     )
 
 
+def _pad_cond_sequence(
+    embeds: list[torch.Tensor],
+    masks: list[torch.Tensor],
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Right-pad per-request ``[1, T_i, H]`` embeds into ``[B, T_max, H]``.
+
+    The transformer derives per-row ``encoder_seq_lengths`` from the boolean
+    mask, so padding must sit on the right and masked positions are excluded
+    from the packed joint sequence.
+    """
+    if not embeds:
+        raise ValueError("Cannot pad an empty conditioning sequence list")
+    if len(embeds) != len(masks):
+        raise ValueError(f"Conditioning embeds/mask count mismatch: {len(embeds)} vs {len(masks)}")
+
+    # Fast-path for single-item batch: bypass padding tensor allocation
+    if len(embeds) == 1:
+        emb0, m0 = embeds[0], masks[0]
+        if emb0.ndim != 3 or emb0.shape[0] != 1:
+            raise ValueError(f"Conditioning embeds[0] must be [1, T, H], got {tuple(emb0.shape)}")
+        if m0.ndim != 2 or m0.shape != emb0.shape[:2]:
+            raise ValueError(f"Conditioning embeds/mask length mismatch: {tuple(emb0.shape)} vs {tuple(m0.shape)}")
+        return emb0, m0
+
+    hidden = embeds[0].shape[-1]
+    lengths = [int(e.shape[1]) for e in embeds]
+    max_len = max(lengths)
+
+    # Fast-path for homogeneous lengths: concatenate directly in C++
+    if all(length == max_len for length in lengths):
+        for i, (emb, m) in enumerate(zip(embeds, masks)):
+            if emb.ndim != 3 or emb.shape[0] != 1 or emb.shape[-1] != hidden:
+                raise ValueError(f"Conditioning embeds[{i}] must be [1, T, {hidden}], got {tuple(emb.shape)}")
+            if m.ndim != 2 or m.shape[0] != 1 or m.shape[1] != max_len:
+                raise ValueError(f"Conditioning embeds/mask length mismatch: {tuple(emb.shape)} vs {tuple(m.shape)}")
+        return torch.cat(embeds, dim=0), torch.cat(masks, dim=0)
+
+    out = embeds[0].new_zeros((len(embeds), max_len, hidden))
+    mask = torch.zeros((len(embeds), max_len), dtype=torch.bool, device=out.device)
+    for i, (emb, m) in enumerate(zip(embeds, masks)):
+        if emb.ndim != 3 or emb.shape[0] != 1 or emb.shape[-1] != hidden:
+            raise ValueError(f"Conditioning embeds[{i}] must be [1, T, {hidden}], got {tuple(emb.shape)}")
+        if m.ndim != 2 or m.shape[0] != 1 or m.shape[1] != emb.shape[1]:
+            raise ValueError(f"Conditioning embeds/mask length mismatch: {tuple(emb.shape)} vs {tuple(m.shape)}")
+        out[i, : emb.shape[1]] = emb[0]
+        mask[i, : m.shape[1]] = m[0].bool()
+    return out, mask
+
+
+def _resolve_request_sampling_and_dims(
+    prompt: dict,
+    sampling: OmniDiffusionSamplingParams | None,
+    request_id: str,
+) -> tuple[int, int, float, tuple[float, float], int]:
+    """Shared extraction and validation of geometry and sampling knobs."""
+    dimensions = []
+    for name in ("height", "width"):
+        value = DiffusionRequestBatch.get_prompt_field(prompt, name)
+        if value is None and sampling is not None:
+            value = getattr(sampling, name, None)
+        if value is None:
+            value = 1024
+        try:
+            dimensions.append(int(value))
+        except (TypeError, ValueError, OverflowError) as exc:
+            raise ValueError(f"Invalid image size {name}={value!r} for request {request_id}") from exc
+    height, width = dimensions
+    if height <= 0 or width <= 0:
+        raise ValueError(f"Invalid image size: {height}x{width} for request {request_id}")
+    if height % 16 != 0 or width % 16 != 0:
+        raise ValueError(f"Image size must be multiples of 16, got {height}x{width} for request {request_id}")
+
+    info = prompt.get("additional_information")
+    request_info = info if isinstance(info, dict) else {}
+    extra_args = sampling.extra_args or {} if sampling else {}
+
+    guidance = extra_args.get("text_guidance_scale")
+    if guidance is None and sampling is not None:
+        guidance = sampling.guidance_scale if sampling.guidance_scale_provided else None
+    if guidance is None:
+        guidance = _first_request_value(request_info.get("text_guidance_scale"))
+    if guidance is None:
+        guidance = 9.0
+    try:
+        text_guidance_scale = float(guidance)
+    except (TypeError, ValueError, OverflowError) as exc:
+        raise ValueError(f"Invalid text_guidance_scale for request {request_id}") from exc
+
+    raw_num_inference_steps = extra_args.get("num_inference_steps")
+    if raw_num_inference_steps is None and sampling is not None:
+        raw_num_inference_steps = sampling.num_inference_steps
+    if raw_num_inference_steps is None:
+        raw_num_inference_steps = _first_request_value(request_info.get("num_inference_steps"))
+    if raw_num_inference_steps is None:
+        raw_num_inference_steps = 50
+    try:
+        num_inference_steps = int(raw_num_inference_steps)
+    except (TypeError, ValueError, OverflowError) as exc:
+        raise ValueError(f"Invalid num_inference_steps for request {request_id}") from exc
+    if num_inference_steps <= 0:
+        raise ValueError(f"num_inference_steps must be positive for request {request_id}")
+
+    cfg_range = extra_args.get("cfg_range")
+    if cfg_range is None:
+        cfg_range = request_info.get("cfg_range")
+    if cfg_range is None:
+        cfg_range = [0.0, 1.0]
+    if not isinstance(cfg_range, (list, tuple)) or len(cfg_range) != 2:
+        raise ValueError(f"cfg_range requires two values for request {request_id}")
+    try:
+        cfg_start, cfg_end = float(cfg_range[0]), float(cfg_range[1])
+    except (TypeError, ValueError, OverflowError) as exc:
+        raise ValueError(f"cfg_range requires two values convertible to floats for request {request_id}") from exc
+    if not 0 <= cfg_start <= cfg_end <= 1:
+        raise ValueError(f"cfg_range must satisfy 0 <= start <= end <= 1 for request {request_id}")
+
+    return height, width, text_guidance_scale, (cfg_start, cfg_end), num_inference_steps
+
+
+def _resolve_gen_vocab_start_index(od_config: OmniDiffusionConfig | None) -> int:
+    """Resolve the visual token threshold through the normalized model configuration."""
+    if od_config is not None and getattr(od_config, "tf_model_config", None) is not None:
+        try:
+            cfg = _build_mammoth_config(od_config)
+            llm_cfg = getattr(cfg, "llm_config", None)
+            if llm_cfg is not None:
+                val = getattr(llm_cfg, "gen_vocab_start_index", None)
+                if val is not None:
+                    return int(val)
+                text_cfg = getattr(llm_cfg, "text_config", None)
+                if text_cfg is not None and getattr(text_cfg, "gen_vocab_start_index", None) is not None:
+                    return int(text_cfg.gen_vocab_start_index)
+        except Exception:
+            pass
+
+        # Fallback inspection of raw dictionary if _build_mammoth_config fails (e.g. mock test objects)
+        raw_cfg = od_config.tf_model_config.to_dict()
+        if isinstance(raw_cfg, dict):
+            llm_cfg = raw_cfg.get("llm_config")
+            if isinstance(llm_cfg, dict):
+                text_cfg = llm_cfg.get("text_config")
+                if isinstance(text_cfg, dict) and text_cfg.get("gen_vocab_start_index") is not None:
+                    return int(text_cfg["gen_vocab_start_index"])
+                if llm_cfg.get("gen_vocab_start_index") is not None:
+                    return int(llm_cfg["gen_vocab_start_index"])
+    return 152064
+
+
+def _validate_request_for_admission(
+    request: OmniDiffusionRequest,
+    od_config: OmniDiffusionConfig | None = None,
+    gen_vocab_start_index: int | None = None,
+) -> tuple[int, int, int]:
+    """Validate request at admission time to fail-fast before entering scheduler queue.
+
+    Returns (height, width, num_inference_steps).
+    """
+    request_id = request.request_id
+    prompt = request.prompt if isinstance(request.prompt, dict) else {}
+    sampling = request.sampling_params
+    if sampling is not None and getattr(sampling, "num_outputs_per_prompt", 1) != 1:
+        raise ValueError(
+            f"MammothModa2 requires num_outputs_per_prompt == 1, got {sampling.num_outputs_per_prompt} "
+            f"for request {request_id}"
+        )
+    if not request.is_dummy_run():
+        info = prompt.get("additional_information")
+        if not isinstance(info, dict):
+            raise ValueError(f"Missing additional_information AR conditions for request {request_id}")
+        full_hidden_states = info.get("full_hidden_states")
+        full_token_ids = info.get("full_token_ids")
+        if not isinstance(full_hidden_states, torch.Tensor) or not isinstance(full_token_ids, list):
+            raise ValueError(f"Expected full_hidden_states tensor and full_token_ids list for request {request_id}")
+        try:
+            answer_start_index = int(info.get("answer_start_index"))
+        except (TypeError, ValueError, OverflowError) as exc:
+            raise ValueError(f"Invalid answer_start_index for request {request_id}") from exc
+        if full_hidden_states.ndim != 2:
+            raise ValueError(f"Expected 2D full_hidden_states for request {request_id}")
+        if full_hidden_states.shape[0] != len(full_token_ids):
+            raise ValueError(f"AR hidden-state/token-count mismatch for request {request_id}")
+        if not 0 <= answer_start_index <= len(full_token_ids):
+            raise ValueError(f"answer_start_index outside token range for request {request_id}")
+        try:
+            int_tokens = [int(token_id) for token_id in full_token_ids]
+        except (TypeError, ValueError, OverflowError) as exc:
+            raise ValueError(f"Invalid full_token_ids for request {request_id}") from exc
+
+        if gen_vocab_start_index is None:
+            gen_vocab_start_index = _resolve_gen_vocab_start_index(od_config)
+        answer_tokens = int_tokens[answer_start_index:]
+        if not any(token_id >= gen_vocab_start_index for token_id in answer_tokens):
+            raise ValueError(
+                f"MammothModa2 AR stage produced no visual-token hidden states for request {request_id}; "
+                f"the DiT stage requires at least one generated visual token. Generated token ids: {answer_tokens[:32]}"
+            )
+
+    height, width, _, _, num_inference_steps = _resolve_request_sampling_and_dims(prompt, sampling, request_id)
+    return height, width, num_inference_steps
+
+
+def get_mammoth_moda2_pre_process_func(od_config: OmniDiffusionConfig | None = None):
+    """Admission preprocessor: fail-fast per-request validation and grouping.
+
+    Validates AR conditions, dimensions, and sampling knobs at admission so
+    malformed requests are rejected individually with their request_id before
+    entering the scheduler queue. Admitted requests are assigned a
+    batch_compatibility_key based on output geometry and inference steps.
+    """
+    gen_vocab_start_index = _resolve_gen_vocab_start_index(od_config)
+
+    def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest:
+        height, width, num_inference_steps = _validate_request_for_admission(
+            request, od_config=od_config, gen_vocab_start_index=gen_vocab_start_index
+        )
+        request.batch_compatibility_key = (
+            "mammoth_moda2_dit",
+            height,
+            width,
+            num_inference_steps,
+        )
+        return request
+
+    return pre_process_func
+
+
 @dataclass(frozen=True)
 class _MammothRequest:
+    index: int
     request_id: str
     full_hidden_states: torch.Tensor
     full_token_ids: list[int]
@@ -88,6 +317,15 @@ class _MammothRequest:
     num_inference_steps: int
     seed: int | None
     generator: torch.Generator | list[torch.Generator] | None
+    generator_device: torch.device | str | None = None
+
+
+@dataclass(frozen=True)
+class _RequestConditioning:
+    text_embeds: torch.Tensor
+    text_mask: torch.Tensor
+    image_embeds: torch.Tensor
+    image_mask: torch.Tensor
 
 
 class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
@@ -103,7 +341,7 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
     _encoder_modules: ClassVar[list[str]] = ["gen_image_condition_refiner"]
     _vae_modules: ClassVar[list[str]] = ["gen_vae"]
 
-    supports_request_batch = False
+    supports_request_batch = True
     supports_step_execution = False
 
     # Load only gen_* weights; ignore llm_model.* to prevent loading the entire LLM backbone in the DiT stage.
@@ -234,20 +472,42 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
             nn.Linear(in_features, out_features, bias=True),
         )
 
-    def _parse_request(self, req: DiffusionRequestBatch) -> _MammothRequest:
-        if req.num_reqs != 1:
-            raise ValueError("MammothModa2 diffusion requires exactly one request")
-        request_id = req.request_id
-        sampling = req.sampling_params
-        if sampling.num_outputs_per_prompt != 1:
-            raise ValueError(f"MammothModa2 requires num_outputs_per_prompt=1 for request {request_id}")
+    @staticmethod
+    def _group_requests(specs: list[_MammothRequest]) -> list[list[int]]:
+        """Group co-scheduled requests by (height, width, num_inference_steps).
 
-        prompt = req.prompts[0]
-        prompt = prompt if isinstance(prompt, dict) else {}
+        Preserves arrival order: the first-seen request of each bucket
+        determines the group's execution order.
+        """
+        buckets: dict[tuple[int, int, int], list[int]] = {}
+        for idx, spec in enumerate(specs):
+            key = (spec.height, spec.width, spec.num_inference_steps)
+            buckets.setdefault(key, []).append(idx)
+        return list(buckets.values())
+
+    def _parse_request(self, req: DiffusionRequestBatch, index: int = 0) -> _MammothRequest:
+        request = req.requests[index]
+        request_id = request.request_id
+        prompt = request.prompt if isinstance(request.prompt, dict) else {}
+        sampling = request.sampling_params
+        if sampling is not None and getattr(sampling, "num_outputs_per_prompt", 1) != 1:
+            raise ValueError(
+                f"MammothModa2 requires num_outputs_per_prompt == 1, got {sampling.num_outputs_per_prompt} "
+                f"for request {request_id}"
+            )
         info = prompt.get("additional_information")
-        if req.is_dummy_run():
-            full_hidden_states = torch.zeros((2, self._llm_hidden_size), dtype=torch.float32, device="cpu")
-            full_token_ids = [0, int(self.config.llm_config.gen_vocab_start_index)]
+        if request.is_dummy_run():
+            gen_start = 152064
+            if hasattr(self, "config") and hasattr(self.config, "llm_config") and self.config.llm_config is not None:
+                val = getattr(self.config.llm_config, "gen_vocab_start_index", None)
+                if val is not None:
+                    gen_start = val
+                else:
+                    text_cfg = getattr(self.config.llm_config, "text_config", None)
+                    if text_cfg is not None and getattr(text_cfg, "gen_vocab_start_index", None) is not None:
+                        gen_start = text_cfg.gen_vocab_start_index
+            full_hidden_states = torch.zeros((2, self._llm_hidden_size), dtype=torch.float32)
+            full_token_ids = [0, int(gen_start)]
             answer_start_index = 1
         else:
             if not isinstance(info, dict):
@@ -271,70 +531,18 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
             except (TypeError, ValueError, OverflowError) as exc:
                 raise ValueError(f"Invalid full_token_ids for request {request_id}") from exc
 
-        dimensions = []
-        for name in ("height", "width"):
-            value = DiffusionRequestBatch.get_prompt_field(prompt, name)
-            if value is None:
-                value = getattr(sampling, name)
-            if value is None:
-                value = 1024
-            try:
-                dimensions.append(int(value))
-            except (TypeError, ValueError, OverflowError) as exc:
-                raise ValueError(f"Invalid image size {name}={value!r} for request {request_id}") from exc
-        height, width = dimensions
-        if height <= 0 or width <= 0:
-            raise ValueError(f"Invalid image size: {height}x{width} for request {request_id}")
-        if height % 16 != 0 or width % 16 != 0:
-            raise ValueError(f"Image size must be multiples of 16, got {height}x{width} for request {request_id}")
+        height, width, text_guidance_scale, cfg_range, num_inference_steps = _resolve_request_sampling_and_dims(
+            prompt, sampling, request_id
+        )
 
-        request_info = info if isinstance(info, dict) else {}
-        extra_args = sampling.extra_args or {}
-        guidance = extra_args.get("text_guidance_scale")
-        if guidance is None:
-            guidance = sampling.guidance_scale if sampling.guidance_scale_provided else None
-        if guidance is None:
-            guidance = _first_request_value(request_info.get("text_guidance_scale"))
-        if guidance is None:
-            guidance = 9.0
-        try:
-            text_guidance_scale = float(guidance)
-        except (TypeError, ValueError, OverflowError) as exc:
-            raise ValueError(f"Invalid text_guidance_scale for request {request_id}") from exc
-        raw_num_inference_steps = extra_args.get("num_inference_steps")
-        if raw_num_inference_steps is None:
-            raw_num_inference_steps = sampling.num_inference_steps
-        if raw_num_inference_steps is None:
-            raw_num_inference_steps = _first_request_value(request_info.get("num_inference_steps"))
-        if raw_num_inference_steps is None:
-            raw_num_inference_steps = 50
-        try:
-            num_inference_steps = int(raw_num_inference_steps)
-        except (TypeError, ValueError, OverflowError) as exc:
-            raise ValueError(f"Invalid num_inference_steps for request {request_id}") from exc
-        if num_inference_steps <= 0:
-            raise ValueError(f"num_inference_steps must be positive for request {request_id}")
-        cfg_range = extra_args.get("cfg_range")
-        if cfg_range is None:
-            cfg_range = request_info.get("cfg_range")
-        if cfg_range is None:
-            cfg_range = [0.0, 1.0]
-        if not isinstance(cfg_range, (list, tuple)) or len(cfg_range) != 2:
-            raise ValueError(f"cfg_range requires two values for request {request_id}")
-        try:
-            cfg_start, cfg_end = float(cfg_range[0]), float(cfg_range[1])
-        except (TypeError, ValueError, OverflowError) as exc:
-            raise ValueError(f"cfg_range requires two values convertible to floats for request {request_id}") from exc
-        if not 0 <= cfg_start <= cfg_end <= 1:
-            raise ValueError(f"cfg_range must satisfy 0 <= start <= end <= 1 for request {request_id}")
-
-        generator = sampling.generator
+        generator = sampling.generator if sampling else None
         if isinstance(generator, list) and len(generator) != 1:
             raise ValueError(
                 f"MammothModa2 single-output request mode requires exactly one generator for request {request_id}"
             )
 
         return _MammothRequest(
+            index=index,
             request_id=request_id,
             full_hidden_states=full_hidden_states,
             full_token_ids=full_token_ids,
@@ -342,10 +550,11 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
             height=height,
             width=width,
             text_guidance_scale=text_guidance_scale,
-            cfg_range=(cfg_start, cfg_end),
+            cfg_range=cfg_range,
             num_inference_steps=num_inference_steps,
-            seed=sampling.seed,
+            seed=sampling.seed if sampling else None,
             generator=generator,
+            generator_device=getattr(sampling, "generator_device", None) if sampling else None,
         )
 
     def _split_ar_conditions(
@@ -363,13 +572,23 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
         caller no longer needs to pass them. Mirrors the masking the bespoke
         MammothModa2 example performed via ar2dit.
         """
-        gen_vocab_start_index = int(self.config.llm_config.gen_vocab_start_index)
-        visual_ids = [
-            int(self.config.image_token_id),
-            int(self.config.video_token_id),
-            int(self.config.vision_start_token_id),
-            int(self.config.vision_end_token_id),
-        ]
+        gen_vocab_start_index = getattr(self.config.llm_config, "gen_vocab_start_index", None)
+        if gen_vocab_start_index is None:
+            text_cfg = getattr(self.config.llm_config, "text_config", None)
+            gen_vocab_start_index = getattr(text_cfg, "gen_vocab_start_index", 152064)
+        gen_vocab_start_index = int(gen_vocab_start_index)
+        cached_visual_ids = getattr(self, "_cached_visual_ids_tensor", None)
+        if cached_visual_ids is None:
+            cached_visual_ids = torch.tensor(
+                [
+                    int(self.config.image_token_id),
+                    int(self.config.video_token_id),
+                    int(self.config.vision_start_token_id),
+                    int(self.config.vision_end_token_id),
+                ],
+                dtype=torch.long,
+            )
+            self._cached_visual_ids_tensor = cached_visual_ids
 
         device = full_hidden_states.device
         token_ids = torch.tensor(full_token_ids, dtype=torch.long, device=device)
@@ -377,7 +596,8 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
         questions_mask = positions < answer_start_index
         answers_mask = ~questions_mask
         gen_token_mask = token_ids >= gen_vocab_start_index
-        visual_token_mask = torch.isin(token_ids, torch.tensor(visual_ids, dtype=torch.long, device=device))
+        visual_ids_device = cached_visual_ids.to(device=device, non_blocking=True)
+        visual_token_mask = torch.isin(token_ids, visual_ids_device)
         text_mask = questions_mask & ~(visual_token_mask | gen_token_mask)
         image_mask = answers_mask & gen_token_mask
 
@@ -389,9 +609,8 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
         image_cond = full_hidden_states[image_mask].contiguous()
         return text_cond, image_cond
 
-    @torch.inference_mode()
-    def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
-        request = self._parse_request(req)
+    def _split_request_conditions(self, request: _MammothRequest) -> tuple[torch.Tensor, torch.Tensor]:
+        """Split AR hidden states and fail fast before touching model params."""
         text_cond, image_cond = self._split_ar_conditions(
             full_hidden_states=request.full_hidden_states,
             full_token_ids=request.full_token_ids,
@@ -404,104 +623,181 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
                 f"request {request.request_id}; the DiT stage requires at least one generated visual token. "
                 f"Generated token ids: {answer_token_ids[:32]}"
             )
+        return text_cond, image_cond
 
-        # Move to model device/dtype.
-        model_device = next(self.parameters()).device
+    def _build_conditioning(
+        self,
+        text_cond: torch.Tensor,
+        image_cond: torch.Tensor,
+        model_device: torch.device,
+        target_dtype: torch.dtype,
+    ) -> _RequestConditioning:
+        text_embeds = text_cond.to(device=model_device, dtype=target_dtype, non_blocking=True).contiguous().unsqueeze(0)
+        text_mask = torch.ones((1, text_embeds.shape[1]), dtype=torch.bool, device=model_device)
+        image_embeds = (
+            image_cond.to(device=model_device, dtype=target_dtype, non_blocking=True).contiguous().unsqueeze(0)
+        )
+        image_mask = torch.ones((1, image_embeds.shape[1]), dtype=torch.bool, device=model_device)
+        return _RequestConditioning(
+            text_embeds=text_embeds,
+            text_mask=text_mask,
+            image_embeds=image_embeds,
+            image_mask=image_mask,
+        )
+
+    @staticmethod
+    def _make_latent_generators(
+        specs: list[_MammothRequest],
+        fallback_device: torch.device,
+    ) -> list[torch.Generator] | None:
+        generators: list[torch.Generator] = []
+        for s in specs:
+            if s.generator is not None:
+                gen = s.generator[0] if isinstance(s.generator, list) else s.generator
+            elif s.seed is not None:
+                device = s.generator_device or fallback_device
+                gen = torch.Generator(device=device).manual_seed(s.seed)
+            else:
+                return None
+            generators.append(gen)
+        return generators
+
+    def _denoise_group(
+        self,
+        specs: list[_MammothRequest],
+        conds: list[_RequestConditioning],
+        model_device: torch.device,
+    ) -> list[torch.Tensor]:
+        batch = len(specs)
+        height = specs[0].height
+        width = specs[0].width
+        num_inference_steps = specs[0].num_inference_steps
+        target_dtype = conds[0].text_embeds.dtype
+
+        # Conditioning collation. Token order is text-then-image per request,
+        # so concatenate first and right-pad the combined sequence.
+        nested_image_embedder = getattr(self.gen_transformer.time_caption_embed, "image_embedder", None)
         if self.gen_image_condition_refiner is not None:
-            target_dtype = next(self.gen_image_condition_refiner.parameters()).dtype
-        else:
-            target_dtype = next(self.gen_transformer.parameters()).dtype
-
-        text_cond = text_cond.to(device=model_device, dtype=target_dtype, non_blocking=True).contiguous()
-        image_cond = image_cond.to(device=model_device, dtype=target_dtype, non_blocking=True).contiguous()
-
-        text_embeds = text_cond.unsqueeze(0)  # [1, T_text, H]
-        text_attention_mask = torch.ones(
-            (1, text_embeds.shape[1]),
-            dtype=torch.bool,
-            device=text_embeds.device,
-        )
-
-        image_embeds = image_cond.unsqueeze(0)  # [1, T_img, H]
-        image_attention_mask = torch.ones(
-            (1, image_embeds.shape[1]),
-            dtype=torch.bool,
-            device=image_embeds.device,
-        )
-
-        # Apply optional refiner ONLY on image condition tokens.
-        if self.gen_image_condition_refiner is not None and image_embeds.shape[1] > 0:
-            image_embeds = self.gen_image_condition_refiner(image_embeds, ~image_attention_mask.bool())
-            image_attention_mask = torch.ones(
+            image_embeds, image_mask = _pad_cond_sequence(
+                [c.image_embeds for c in conds],
+                [c.image_mask for c in conds],
+            )
+            # Apply optional refiner ONLY on image condition tokens.
+            if image_embeds.shape[1] > 0:
+                image_embeds = self.gen_image_condition_refiner(image_embeds, ~image_mask.bool())
+            refined_image_mask = torch.ones(
                 image_embeds.shape[:2],
                 dtype=torch.bool,
                 device=image_embeds.device,
             )
-
-        nested_image_embedder = getattr(self.gen_transformer.time_caption_embed, "image_embedder", None)
-        if nested_image_embedder is None:
-            prompt_embeds = torch.cat([text_embeds, image_embeds], dim=1)
-            prompt_attention_mask = torch.cat([text_attention_mask, image_attention_mask], dim=1)
+            seq_embeds = []
+            seq_masks = []
+            for i, c in enumerate(conds):
+                cur_img_embed = image_embeds[i : i + 1]
+                cur_img_mask = refined_image_mask[i : i + 1]
+                seq_embeds.append(torch.cat([c.text_embeds, cur_img_embed], dim=1))
+                seq_masks.append(torch.cat([c.text_mask, cur_img_mask], dim=1))
+            prompt_embeds, prompt_attention_mask = _pad_cond_sequence(seq_embeds, seq_masks)
             ar_image_embeds = None
             ar_image_attention_mask = None
+        elif nested_image_embedder is not None:
+            prompt_embeds, prompt_attention_mask = _pad_cond_sequence(
+                [c.text_embeds for c in conds],
+                [c.text_mask for c in conds],
+            )
+            ar_image_embeds, ar_image_attention_mask = _pad_cond_sequence(
+                [c.image_embeds for c in conds],
+                [c.image_mask for c in conds],
+            )
         else:
-            prompt_embeds = text_embeds
-            prompt_attention_mask = text_attention_mask
-            ar_image_embeds = image_embeds
-            ar_image_attention_mask = image_attention_mask
+            prompt_embeds, prompt_attention_mask = _pad_cond_sequence(
+                [torch.cat([c.text_embeds, c.image_embeds], dim=1) for c in conds],
+                [torch.cat([c.text_mask, c.image_mask], dim=1) for c in conds],
+            )
+            ar_image_embeds = None
+            ar_image_attention_mask = None
 
-        # Empty unconditional prompt for classifier-free guidance.
+        # Empty unconditional prompt for classifier-free guidance, shared by
+        # all rows; rows that never take the CFG branch select their cond
+        # prediction directly via torch.where below.
+        needs_uncond = any(s.text_guidance_scale > 1.0 for s in specs)
         negative_prompt_embeds = None
         negative_prompt_attention_mask = None
-        if request.text_guidance_scale > 1.0:
+        if needs_uncond:
             hidden_size = int(prompt_embeds.shape[-1])
             negative_prompt_embeds = torch.zeros(
-                (1, 0, hidden_size),
+                (batch, 0, hidden_size),
                 dtype=target_dtype,
-                device=prompt_embeds.device,
+                device=model_device,
             )
             negative_prompt_attention_mask = torch.zeros(
-                (1, 0),
+                (batch, 0),
                 dtype=torch.bool,
-                device=prompt_embeds.device,
+                device=model_device,
             )
 
-        generator = request.generator
-        if generator is None and request.seed is not None:
-            generator = torch.Generator(device=model_device).manual_seed(request.seed)
-
-        # Output image size (px), passed from stage input processor.
-        height, width = request.height, request.width
         vae_scale_factor = 16
-
         latent_channels = int(self.gen_transformer.config.in_channels)
-        shape = (1, latent_channels, 2 * height // vae_scale_factor, 2 * width // vae_scale_factor)
-        latents = randn_tensor(shape, generator=generator, device=prompt_embeds.device, dtype=prompt_embeds.dtype)
+        shape = (batch, latent_channels, 2 * height // vae_scale_factor, 2 * width // vae_scale_factor)
+        single_shape = (1, latent_channels, 2 * height // vae_scale_factor, 2 * width // vae_scale_factor)
+        generators = self._make_latent_generators(specs, model_device)
+        if generators is not None:
+            noise_list = []
+            for gen in generators:
+                gen_device = gen.device if hasattr(gen, "device") else model_device
+                n = randn_tensor(single_shape, generator=gen, device=gen_device, dtype=target_dtype)
+                noise_list.append(n.to(device=model_device))
+            latents = torch.cat(noise_list, dim=0)
+        else:
+            latents = randn_tensor(shape, device=model_device, dtype=target_dtype)
 
         scheduler = FlowMatchEulerDiscreteScheduler()
-
         scheduler.set_timesteps(
-            num_inference_steps=request.num_inference_steps,
-            device=prompt_embeds.device,
+            num_inference_steps=num_inference_steps,
+            device=model_device,
             num_tokens=latents.shape[-2] * latents.shape[-1],
         )
 
-        # Run diffusion loop (CFG supported when text_guidance_scale > 1.0)
+        scale_vec = torch.tensor(
+            [s.text_guidance_scale for s in specs],
+            device=model_device,
+            dtype=target_dtype,
+        ).view(batch, 1, 1, 1)
+        cfg_specs = [(s.text_guidance_scale > 1.0, float(s.cfg_range[0]), float(s.cfg_range[1])) for s in specs]
+
         total_steps = max(1, len(scheduler.timesteps))
+        any_cfg_active = any(s.text_guidance_scale > 1.0 for s in specs)
         self._prepare_cache_dit_for_request(
             num_inference_steps=total_steps,
-            cfg_active=request.text_guidance_scale > 1.0,
+            cfg_active=any_cfg_active,
         )
-        # cache-dit distinguishes the conditional vs unconditional passes
-        # purely by transformer-forward parity, so when it is active the
-        # uncond pass must run on every step. Outside cfg_range CFG is
-        # neutralized via scale=1.0 (uncond + 1.0 * (cond - uncond) == cond),
-        # matching the skip optimization's cond-only update.
         requires_paired_cfg = self._cache_dit_runtime.is_enabled and getattr(
             self, "_cache_dit_requires_paired_cfg", False
         )
+
+        # Precompute mixed-CFG active masks across timesteps to avoid CPU list
+        # comprehensions and H2D copies in the step loop.
+        precomputed_active_masks = None
+        all_active_per_step = None
+        any_active_per_step = None
+        if needs_uncond:
+            precomputed_active_masks = torch.empty((total_steps, batch, 1, 1, 1), device=model_device, dtype=torch.bool)
+            all_active_per_step = []
+            any_active_per_step = []
+            for step_idx in range(total_steps):
+                frac = step_idx / total_steps
+                step_mask = [is_active and (lo <= frac <= hi) for is_active, lo, hi in cfg_specs]
+                all_active_per_step.append(all(step_mask))
+                any_active_per_step.append(any(step_mask))
+                mask_t = torch.tensor(step_mask, device=model_device, dtype=torch.bool)
+                precomputed_active_masks[step_idx] = mask_t.view(batch, 1, 1, 1)
+
+        # Precompute timesteps in latent dtype to eliminate per-step cast overhead
+        timesteps_dtype = scheduler.timesteps.to(latents.dtype)
+
+        # Run diffusion loop (CFG supported when text_guidance_scale > 1.0)
         for i, t in enumerate(scheduler.timesteps):
-            timestep = t.expand(latents.shape[0]).to(latents.dtype)
+            timestep = timesteps_dtype[i].expand(batch)
             model_pred = self.gen_transformer(
                 hidden_states=latents,
                 timestep=timestep,
@@ -512,10 +808,7 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
                 ar_image_attention_mask=ar_image_attention_mask,
                 freqs_cis=self.gen_freqs_cis,
             )
-            guidance_scale = (
-                request.text_guidance_scale if request.cfg_range[0] <= i / total_steps <= request.cfg_range[1] else 1.0
-            )
-            run_uncond = guidance_scale > 1.0 and negative_prompt_embeds is not None
+            run_uncond = (any_active_per_step is not None and any_active_per_step[i]) if needs_uncond else False
             if requires_paired_cfg and negative_prompt_embeds is not None:
                 run_uncond = True
             if run_uncond:
@@ -527,18 +820,65 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
                     ref_image_hidden_states=None,
                     freqs_cis=self.gen_freqs_cis,
                 )
-                model_pred = model_pred_uncond + guidance_scale * (model_pred - model_pred_uncond)
+                # Fused CFG blend: torch.lerp fuses (uncond + scale * (cond - uncond)) into 1 kernel
+                blended = torch.lerp(model_pred_uncond, model_pred, scale_vec)
+                if all_active_per_step is not None and all_active_per_step[i]:
+                    # Fast path: all requests in the batch are active.
+                    model_pred = blended
+                else:
+                    # Inactive rows keep their conditional prediction exactly:
+                    # gating the blend with torch.where avoids propagating
+                    # uncond-branch NaN/rounding into CFG-free rows.
+                    active_tensor = precomputed_active_masks[i]
+                    model_pred = torch.where(active_tensor, blended, model_pred)
             latents = scheduler.step(model_pred, t, latents, return_dict=False)[0]
-            latents = latents.to(dtype=prompt_embeds.dtype)
+            latents = latents.to(dtype=target_dtype)
 
-        # VAE decode
+        # VAE decode (in-place scaling to reduce peak VRAM before decode)
         if self.gen_vae.config.scaling_factor is not None:
-            latents = latents / self.gen_vae.config.scaling_factor
+            latents.div_(self.gen_vae.config.scaling_factor)
         if self.gen_vae.config.shift_factor is not None:
-            latents = latents + self.gen_vae.config.shift_factor
-        image = self.gen_vae.decode(latents, return_dict=False)[0]
+            latents.add_(self.gen_vae.config.shift_factor)
+        vae_dtype = next(self.gen_vae.parameters()).dtype
+        image = self.gen_vae.decode(latents.to(dtype=vae_dtype), return_dict=False)[0]  # [B, C, H, W]
+        return [image[i : i + 1] for i in range(batch)]
 
-        return DiffusionOutput(output=image)
+    @torch.inference_mode()
+    def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]:
+        specs = [self._parse_request(req, i) for i in range(req.num_reqs)]
+
+        # Validate AR conditions before touching model parameters so malformed
+        # requests fail fast instead of surfacing device/dtype errors.
+        raw_conds = [self._split_request_conditions(spec) for spec in specs]
+
+        model_device = next(self.parameters()).device
+        if self.gen_image_condition_refiner is not None:
+            target_dtype = next(self.gen_image_condition_refiner.parameters()).dtype
+        else:
+            target_dtype = next(self.gen_transformer.parameters()).dtype
+
+        conds = [
+            self._build_conditioning(text_cond, image_cond, model_device, target_dtype)
+            for text_cond, image_cond in raw_conds
+        ]
+
+        groups = self._group_requests(specs)
+        if len(groups) == 1 and len(groups[0]) == len(specs):
+            images = self._denoise_group(specs, conds, model_device)
+            return [DiffusionOutput(output=img) for img in images]
+
+        outputs: list[DiffusionOutput | None] = [None] * len(specs)
+        for group in groups:
+            images = self._denoise_group(
+                [specs[i] for i in group],
+                [conds[i] for i in group],
+                model_device,
+            )
+            for index, image in zip(group, images):
+                outputs[index] = DiffusionOutput(output=image)
+        if any(output is None for output in outputs):
+            raise RuntimeError("DiT batching produced no image for at least one scheduled request")
+        return [output for output in outputs if output is not None]
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         loader = AutoWeightsLoader(self)

vllm_omni/diffusion/models/ming_image/pipeline.py

@@ -107,6 +107,11 @@ class MingImageDiffusionPipeline(ZImagePipeline):
     _encoder_modules: ClassVar[list[str]] = ["conditioning"]
     _vae_modules: ClassVar[list[str]] = ["vae"]
 
+    @property
+    def do_classifier_free_guidance(self) -> bool:
+        # The vendor Ming-Image pipeline only enables CFG above 1.0
+        return self._guidance_scale > 1
+
     def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = "") -> None:
         del prefix
         nn.Module.__init__(self)
@@ -190,6 +195,7 @@ class MingImageDiffusionPipeline(ZImagePipeline):
         self.image_processor = VaeImageProcessor(
             vae_scale_factor=self.vae_scale_factor * 2,
             do_convert_rgb=False,
+            resample="bilinear",
         )
         self.setup_diffusion_pipeline_profiler(
             enable_diffusion_pipeline_profiler=od_config.enable_diffusion_pipeline_profiler
@@ -314,18 +320,10 @@ class MingImageDiffusionPipeline(ZImagePipeline):
 
         height = int(extra_args.get("height") or sampling.height or 1024)
         width = int(extra_args.get("width") or sampling.width or 1024)
-        steps = int(extra_args.get("steps") or sampling.num_inference_steps or self.default_num_inference_steps)
-        cfg = float(
-            extra_args["cfg"]
-            if extra_args.get("cfg") is not None
-            else sampling.guidance_scale
-            if sampling.guidance_scale is not None
-            else self.default_guidance_scale
-        )
+        steps = int(sampling.num_inference_steps or self.default_num_inference_steps)
+        cfg = float(sampling.guidance_scale if sampling.guidance_scale is not None else self.default_guidance_scale)
         seed = extra_args.get("seed", sampling.seed)
-        generator = (
-            torch.Generator(device=self.device).manual_seed(int(seed)) if seed is not None else sampling.generator
-        )
+        generator = torch.Generator(device="cpu").manual_seed(int(seed)) if seed is not None else sampling.generator
 
         ref_latent = self._encode_reference(reference, height, width)
         positive = [item for item in cap_feats]
@@ -333,7 +331,7 @@ class MingImageDiffusionPipeline(ZImagePipeline):
         self._pending_prompt_embeds = positive
         self._pending_negative_prompt_embeds = negative
 
-        apply_cfg = cfg > 0
+        apply_cfg = cfg > 1
         context_direct = (
             torch.cat([direct_condition, torch.zeros_like(direct_condition)], dim=0) if apply_cfg else direct_condition
         )

vllm_omni/diffusion/models/minimax_h3/denoise_loop.py

@@ -16,6 +16,7 @@ from typing import Any
 import torch
 
 from vllm_omni.diffusion.attention.backends.abstract import VideoTokenLayout, VideoTokenSpan
+from vllm_omni.diffusion.cancellation import check_request_cancellation
 from vllm_omni.diffusion.forward_context import (
     set_forward_context_denoise_step_idx,
     set_forward_context_denoise_timestep,
@@ -366,6 +367,7 @@ def minimax_h3_denoise_loop(
 
     num_steps = len(sigmas_video) - 1
     for step in range(num_steps):
+        check_request_cancellation()
         step_cm = step_profiler(step) if step_profiler is not None else nullcontext()
         with step_cm:
             s_v, s_v_next = sigmas_video[step], sigmas_video[step + 1]
@@ -450,6 +452,7 @@ def minimax_h3_denoise_loop(
                 audio_rows[~audio_update] = audio_anchor  # per-step audio ref reset
             if on_step is not None:
                 on_step(step, video_rows, audio_rows)
+            check_request_cancellation(synchronize=True)
 
     minimax_h3_publish_denoise_progress(None, None, None)
     return video_rows, audio_rows

vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py

@@ -27,6 +27,7 @@ from vllm_omni.diffusion.cache.cachedit import (
     CacheDiTBackend,
     RequestScopedCacheDiTRuntime,
 )
+from vllm_omni.diffusion.cancellation import check_request_cancellation
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.parallel_state import get_world_group, init_world_group
 from vllm_omni.diffusion.distributed.utils import get_local_device
@@ -604,6 +605,7 @@ class MiniMaxH3Pipeline(
     """CFG-distilled joint video/audio generation for MiniMax H3."""
 
     supports_step_execution: ClassVar[bool] = True
+    supports_request_cancellation: ClassVar[bool] = True
 
     _dit_modules: ClassVar[list[str]] = ["transformer", "transformers_ref"]
     _encoder_modules: ClassVar[list[str]] = ["text_encoder"]
@@ -2619,15 +2621,18 @@ class MiniMaxH3Pipeline(
     def forward(self, request: DiffusionRequestBatch) -> DiffusionOutput:
         if len(request.prompts) != 1:
             raise OmniClientError("MiniMax H3 supports one request at a time")
+        check_request_cancellation()
         context = self._prepare_request_inputs(
             request.prompts[0],
             request.sampling_params,
         )
+        check_request_cancellation()
         denoise_kwargs = self._denoise_kwargs(context)
         num_outputs = context["num_outputs"]
         videos = []
         audios = []
         for output_seed in _minimax_h3_output_seeds(context["seed"], num_outputs):
+            check_request_cancellation()
             output_kwargs = {**denoise_kwargs, "seed": output_seed}
             if context.get("continuation") is None:
                 video_latent, audio_latent = self.diffuse(**output_kwargs)
@@ -2640,6 +2645,7 @@ class MiniMaxH3Pipeline(
                     overlap_frames=overlap_frames,
                     text_conditioning=context.get("continuation_text_conditioning"),
                 )
+            check_request_cancellation()
             if context["preencode_mp4"]:
                 videos.append(
                     self.decode_to_mp4(

vllm_omni/diffusion/offloader/base.py

@@ -244,5 +244,12 @@ class OffloadBackend(ABC):
         """
         raise NotImplementedError
 
+    def shutdown(self) -> None:
+        """Release offload resources at process exit.
+
+        Backends may skip work that only matters for a later ``enable``.
+        """
+        self.disable()
+
     def is_enabled(self) -> bool:
         return self.enabled

vllm_omni/diffusion/offloader/distributed_layerwise_backend.py

@@ -1889,10 +1889,11 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         skipped_allgather = bool(allgather_hooks) and not restore_allgather_weights
         if skipped_allgather:
             # Startup rollback cannot safely enter a collective that a failed
-            # peer may never reach. Those blocks cannot be reconstructed, so
-            # make accidental reuse explicit instead of accepting zero weights.
+            # peer may never reach, and shutdown has no later user. Those blocks
+            # cannot be reconstructed, so make accidental reuse explicit
+            # instead of accepting zero weights.
             self._poisoned_reason = (
-                "Distributed layerwise offload startup skipped AllGather weight restoration; "
+                "Distributed layerwise offload skipped AllGather weight restoration; "
                 "recreate the backend and reload the pipeline before retrying"
             )
 
@@ -1981,6 +1982,11 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
     def disable(self) -> None:
         self._disable(restore_allgather_weights=True)
 
+    def shutdown(self) -> None:
+        # Restoring AllGather blocks rebuilds the full DiT in every rank's host
+        # memory, which can outlast the executor's shutdown grace period.
+        self._disable(restore_allgather_weights=False)
+
     @staticmethod
     def _allocate_shared_buffers(
         hooks: list[DistributedLayerwiseOffloadHook],

vllm_omni/diffusion/registry.py

@@ -746,6 +746,7 @@ _DIFFUSION_PRE_PROCESS_FUNCS = {
     "SanaWmPipeline": "get_sana_wm_pre_process_func",
     "Cosmos3OmniDiffusersPipeline": "get_cosmos3_pre_process_func",
     "Cosmos3OmniPipeline": "get_cosmos3_pre_process_func",
+    "MammothModa2DiTPipeline": "get_mammoth_moda2_pre_process_func",
 }
 
 

vllm_omni/diffusion/request.py

@@ -5,7 +5,7 @@
 from __future__ import annotations
 
 import random
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any
 
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
@@ -81,6 +81,9 @@ class OmniDiffusionRequest:
     kv_recv_ms: float = 0.0
     # Time spent waiting for initial admission by the diffusion scheduler.
     scheduler_queue_wait_ms: float | None = None
+    # Engine-owned shared-memory signal for cooperative full-forward workers.
+    # Only the engine creates/unlinks it; workers attach readers until return.
+    cancellation_signal: str | None = field(default=None, repr=False)
 
     def __post_init__(self):
         """Initialize dependent fields after dataclass initialization."""

vllm_omni/diffusion/utils/video_decode.py

@@ -0,0 +1,91 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Decode video file paths to RGB frames.
+
+Path suffix checks and bounded first/last-N imageio decoding. Model-specific
+conditioning and decode-budget stay in each pipeline's preprocess layer.
+
+Currently used by Cosmos3 V2V preprocess and transfer ``control_path``.
+Unifying with ``OmniVideoBackend`` / ``_decode_video_bytes`` is a follow-up.
+"""
+
+from __future__ import annotations
+
+import os
+from collections import deque
+from pathlib import Path
+from typing import Any, Literal
+
+import numpy as np
+import PIL.Image
+
+VIDEO_EXTENSIONS = {".avi", ".m4v", ".mkv", ".mov", ".mp4", ".webm"}
+
+
+def is_video_file_path(value: Any) -> bool:
+    """True when ``value`` is a filesystem path with a known video suffix.
+
+    Serving may leave uploaded references as temp paths such as
+    ``/tmp/vllm_omni_video_reference_*.mp4`` instead of decoded frames.
+    """
+    if not isinstance(value, str | os.PathLike):
+        return False
+    suffix = os.path.splitext(os.fspath(value))[1].lower()
+    return suffix in VIDEO_EXTENSIONS
+
+
+def _frame_to_uint8_rgb(value: Any) -> np.ndarray:
+    if isinstance(value, PIL.Image.Image):
+        return np.array(value.convert("RGB"), dtype=np.uint8, copy=True)
+    if isinstance(value, np.ndarray):
+        array = value
+        if array.ndim == 3 and array.shape[0] in (3, 4) and array.shape[-1] not in (3, 4):
+            array = np.transpose(array[:3], (1, 2, 0))
+        if np.issubdtype(array.dtype, np.floating):
+            if array.size and (array.min() < 0.0 or array.max() > 1.0):
+                array = np.clip(array, -1.0, 1.0) * 0.5 + 0.5
+            array = (np.clip(array, 0.0, 1.0) * 255.0).round().astype(np.uint8)
+        return array[..., :3].astype(np.uint8)
+    raise TypeError(f"Video decode expected an RGB frame, got {type(value)!r}.")
+
+
+def decode_path_video_frames(
+    path: str | Path,
+    *,
+    max_frames: int | None = None,
+    keep: Literal["first", "last"] = "first",
+) -> list[np.ndarray]:
+    """Decode a video file to RGB uint8 frames (H, W, C).
+
+    ``keep='first'`` stops after ``max_frames``. ``keep='last'`` still scans the
+    file and retains only the tail window.
+    """
+    media_path = Path(path)
+    if not media_path.exists():
+        raise FileNotFoundError(f"Video path does not exist: {media_path}")
+    if keep not in {"first", "last"}:
+        raise ValueError("Video keep must be either 'first' or 'last'.")
+    if max_frames is not None and int(max_frames) <= 0:
+        raise ValueError("Video max_frames must be positive.")
+    try:
+        import imageio.v3 as iio
+    except ImportError as exc:
+        raise ImportError(
+            "Video path decoding requires imageio. Install imageio[ffmpeg] or provide decoded frames."
+        ) from exc
+
+    limit = None if max_frames is None else int(max_frames)
+    if keep == "last" and limit is not None:
+        window: deque[np.ndarray] = deque(maxlen=limit)
+        for frame in iio.imiter(media_path):
+            window.append(_frame_to_uint8_rgb(frame))
+        frames = list(window)
+    else:
+        frames = []
+        for frame in iio.imiter(media_path):
+            frames.append(_frame_to_uint8_rgb(frame))
+            if limit is not None and len(frames) >= limit:
+                break
+    if not frames:
+        raise ValueError(f"Video path produced no frames: {media_path}")
+    return frames

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -31,8 +31,9 @@ from vllm_omni.diffusion.cache.prompt_embed_cache import (
     resolve_prompt_embed_cache_config,
 )
 from vllm_omni.diffusion.cache.selector import get_cache_backend
+from vllm_omni.diffusion.cancellation import check_request_cancellation, request_cancellation_scope
 from vllm_omni.diffusion.compile import regionally_compile
-from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.data import DiffusionOutput, DiffusionRequestAbortedError, OmniDiffusionConfig
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
 from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata
 from vllm_omni.diffusion.diffusion_kv.model_runner_backend import DiffusionKVModelRunnerBackend
@@ -853,15 +854,29 @@ class DiffusionModelRunner(DiffusionStagePayloadMixin):
                     in_diffusion_kv_memory_profile=in_diffusion_kv_memory_profile,
                 ),
                 paged_kv_context,
+                request_cancellation_scope(
+                    [getattr(req, "cancellation_signal", None) for req in reqs],
+                    enabled=getattr(self.pipeline, "supports_request_cancellation", False) is True,
+                ),
             ):
                 with record_function(record_name):
-                    raw_outputs = self.pipeline.forward(batch)
-                    outputs = _normalize_pipeline_outputs(
-                        raw_outputs,
-                        expected_count=len(reqs),
-                        allow_single_output=allow_single_output,
-                        pipeline_name=type(self.pipeline).__name__,
-                    )
+                    try:
+                        check_request_cancellation()
+                        raw_outputs = self.pipeline.forward(batch)
+                        outputs = _normalize_pipeline_outputs(
+                            raw_outputs,
+                            expected_count=len(reqs),
+                            allow_single_output=allow_single_output,
+                            pipeline_name=type(self.pipeline).__name__,
+                        )
+                    except DiffusionRequestAbortedError as exc:
+                        # The checkpoint aborts only a fully cancelled wave;
+                        # a mixed batch must keep running for its live peers.
+                        logger.info(
+                            "Stopped cancelled diffusion request(s) %s at a model execution boundary",
+                            [req.request_id for req in reqs],
+                        )
+                        outputs = [DiffusionOutput(aborted=True, abort_message=str(exc)) for _ in reqs]
                 with record_function("prepare_output_for_transport"):
                     outputs = [
                         self._prepare_output_for_transport(output, req.sampling_params)

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -1048,9 +1048,9 @@ class DiffusionWorker:
             try:
                 offload_backend = getattr(model_runner, "offload_backend", None)
                 if offload_backend is not None:
-                    offload_backend.disable()
+                    offload_backend.shutdown()
             except Exception:
-                logger.exception("Failed to disable diffusion offload backend during shutdown")
+                logger.exception("Failed to shut down diffusion offload backend during shutdown")
             try:
                 if mgr is not None:
                     mgr.close()

vllm_omni/engine/duplex/intermediate.py

@@ -3,7 +3,11 @@
 
 from __future__ import annotations
 
-from typing import TypedDict
+from collections.abc import Mapping
+from typing import TYPE_CHECKING, TypedDict
+
+if TYPE_CHECKING:
+    from vllm_omni.engine.duplex.contracts import DuplexFence
 
 
 class DuplexIntermediateBuffer(TypedDict, total=False):
@@ -32,6 +36,49 @@ class DuplexIntermediateBuffer(TypedDict, total=False):
     mel_spec: object
 
 
+def build_duplex_append_prompt(
+    *,
+    request_id: str,
+    fence: DuplexFence,
+    session_config: Mapping[str, object],
+    runtime_config: Mapping[str, object],
+    seq: int,
+    turn_seq: int,
+    payload: object,
+    final: bool,
+    prompt_token_ids: list[int],
+    model_fields: Mapping[str, object],
+) -> dict[str, object]:
+    """Wrap model-selected tokens and payload in the shared append envelope.
+
+    Models own token selection and extra worker fields; the framework owns
+    request identity, sequencing and shallow config snapshots.
+    """
+    return {
+        "prompt_token_ids": prompt_token_ids,
+        "model_intermediate_buffer": {
+            "request_id": request_id,
+            "global_request_id": [fence.session_id],
+            "duplex": {
+                **model_fields,
+                "data_plane": True,
+                "fence": fence,
+                "session_id": fence.session_id,
+                "epoch": fence.epoch,
+                "seq": seq,
+                "turn_id": fence.turn_id,
+                "turn_seq": turn_seq,
+                "mode": "append_audio_chunk",
+                "payload": payload,
+                "final": final,
+                "session_config": dict(session_config),
+                "runtime_config": dict(runtime_config),
+                "scheduler_token_budget": len(prompt_token_ids),
+            },
+        },
+    }
+
+
 def build_duplex_intermediate_buffer(
     *,
     request_id: str,
@@ -62,12 +109,12 @@ def build_duplex_intermediate_buffer(
     return buffer
 
 
-def set_ref_audio(buffer: dict[str, object], waveform: object, sample_rate_hz: int) -> None:
+def set_ref_audio(buffer: DuplexIntermediateBuffer, waveform: object, sample_rate_hz: int) -> None:
     buffer.setdefault("codes", {})["ref"] = waveform
     buffer.setdefault("meta", {})["ref_audio_sr"] = int(sample_rate_hz)
 
 
-def set_tts_handoff(buffer: dict[str, object], token_ids: object | None, hidden_states: object | None) -> None:
+def set_tts_handoff(buffer: DuplexIntermediateBuffer, token_ids: object | None, hidden_states: object | None) -> None:
     """Store the AR-to-TTS handoff used by the full-duplex stage bridge."""
     if token_ids is not None:
         buffer.setdefault("ids", {})["tts"] = token_ids

vllm_omni/engine/duplex/messages.py

@@ -65,6 +65,11 @@ class TouchDuplexSessionMessage(EngineQueueMessage, kw_only=True):
     control_id: str
     session_id: str
     activity: str
+    #: For a detach: the lease generation the caller opened or resumed
+    #: against. The manager refuses to detach a newer lease, so a connection
+    #: giving up its own resume can never put the lease a later resume owns
+    #: into disconnect grace. ``None`` detaches whatever generation is current.
+    expected_lease_generation: int | None = None
 
 
 class DuplexSessionCommandMessage(EngineQueueMessage, kw_only=True):

vllm_omni/engine/duplex/plugin.py

@@ -38,6 +38,21 @@ class DuplexRuntimeConfigError(ValueError):
         self.code = code
 
 
+def reject_private_runtime_keys(
+    extra_body: object,
+    private_runtime_config_keys: frozenset[str],
+    *,
+    message: str,
+    error_cls: type[DuplexRuntimeConfigError] = DuplexRuntimeConfigError,
+) -> None:
+    """Reject client overrides while preserving each plugin's error contract."""
+    if not isinstance(extra_body, dict):
+        return
+    private_keys = sorted(private_runtime_config_keys.intersection(extra_body))
+    if private_keys:
+        raise error_cls(message + ", ".join(private_keys))
+
+
 def reject_changed_runtime_value(
     new_value: object,
     current_value: object,
@@ -146,6 +161,64 @@ class DuplexModelSessionState(ABC):
     def clear_continuation(self) -> None: ...
 
 
+@dataclass(slots=True)
+class DefaultDuplexModelSessionState(DuplexModelSessionState):
+    """Framework-owned flag anatomy, implemented once.
+
+    The flag set above is the session runner's contract with the model (commit
+    retention, deferred response/creates, silence-continuation bookkeeping); it
+    is identical for every lockstep or frame-locked model, so a plugin only
+    supplies its input packetizer as ``audio_buffer``.
+    """
+
+    audio_buffer: PcmAppendBuffer
+    input_since_commit: bool = False
+    speech_since_commit: bool = False
+    context_locked: bool = False
+    committed_audio_payload: dict[str, object] | None = None
+    committed_audio_operation_id: str | None = None
+    committed_audio_reserved_bytes: int = 0
+    deferred_response_create: bool = False
+    deferred_precreate_response: bool = False
+    continuation_owner_id: str | None = None
+    continuation_units: int = 0
+    pending_silence_task: asyncio.Task[bool] | None = None
+    pending_silence_owner_id: str | None = None
+    # Deadline-aligned silence continuation state: the monotonic submission
+    # time of the most recent native input unit and the next silence
+    # continuation deadline. A real (non-silence) input resets the chain.
+    last_native_submit_monotonic: float | None = None
+    silence_deadline_monotonic: float | None = None
+
+    def retain_committed_audio(
+        self,
+        payload: dict[str, object],
+        *,
+        operation_id: str | None,
+        reserved_bytes: int = 0,
+    ) -> None:
+        self.committed_audio_payload = payload
+        self.committed_audio_operation_id = operation_id
+        self.committed_audio_reserved_bytes += max(0, int(reserved_bytes))
+
+    def clear_committed_audio(self) -> int:
+        reserved_bytes = self.committed_audio_reserved_bytes
+        self.committed_audio_payload = None
+        self.committed_audio_operation_id = None
+        self.committed_audio_reserved_bytes = 0
+        self.deferred_response_create = False
+        self.deferred_precreate_response = False
+        return reserved_bytes
+
+    def clear_continuation(self) -> None:
+        self.continuation_owner_id = None
+        self.continuation_units = 0
+        self.pending_silence_task = None
+        self.pending_silence_owner_id = None
+        self.last_native_submit_monotonic = None
+        self.silence_deadline_monotonic = None
+
+
 class DuplexDataPlane(ABC):
     """Projects raw stage outputs of one model into internal duplex events."""
 
@@ -397,6 +470,20 @@ class DuplexModelPlugin(ABC):
         del config, current, item
         return None
 
+    def runtime_config_after_model_output(
+        self,
+        current: Mapping[str, object],
+        output_metadata: Mapping[str, object],
+    ) -> dict[str, object] | None:
+        """Return a runtime-config patch after a model output is observed.
+
+        Plugins may use model-owned output metadata to retire server-side
+        runtime state that has been consumed by the worker. The default keeps
+        the framework unaware of model-specific metadata.
+        """
+        del current, output_metadata
+        return None
+
 
 def load_duplex_plugin(path: str, encode_audio: EncodeAudio) -> DuplexModelPlugin:
     module_name, separator, attribute_name = path.rpartition(".")
@@ -448,6 +535,7 @@ def coerce_int(value: object) -> int | None:
 
 
 __all__ = [
+    "DefaultDuplexModelSessionState",
     "DuplexDataPlane",
     "DuplexModelPlugin",
     "DuplexModelSessionState",

vllm_omni/engine/duplex/session/append_task.py

@@ -13,6 +13,13 @@ That compensation is the reason this is an object. The five steps of an append
 all need the same eleven values, so as closures over ``_start_append`` they read
 as one function with five entry points and no way to test any of them. Named
 fields make the captured state explicit and the compensation a method.
+
+The precreated response has two other owners the compensation defers to. When
+the runner cancels the append on purpose (``DuplexSessionTasks.cancel_append_tasks``,
+for a ``response.cancel``, a barge-in or a close), the runner ends that response
+itself right after, with the cancel status the client asked for. When the
+session is already closing, the close ends it. Everywhere else, the append is
+the last thing that knows the response will never be filled, and it fails it.
 """
 
 from __future__ import annotations
@@ -53,6 +60,10 @@ class AppendAttempt:
     retained_committed_payload: dict[str, object] | None
     #: Response reserved before submission, to be failed if submission does not happen.
     precreated_response_id: str | None
+    #: Whether ``_start_append`` bound the session's active request to this
+    #: append (a final or response-bound one), so an append that never happens
+    #: has to unbind it again.
+    owns_request: bool = False
     #: Commits timing state once the runtime accepts the append (before any
     #: returned output event can clear the continuation chain).
     on_append_accepted: Callable[[float], None] | None = None
@@ -75,11 +86,31 @@ class AppendAttempt:
         ):
             self.ctx.session.release_input_bytes(model_state.clear_committed_audio())
 
-    def abandon(self) -> None:
-        """Give back everything this append reserved but never used."""
+    def abandon(self, *, reason: str = "runtime_append_failed", response_taken_over: bool = False) -> None:
+        """Give back everything this append reserved but never used.
+
+        That includes the response it precreated: once the append is off,
+        nothing will ever fill it, and a client that saw ``response.created``
+        is owed a ``response.done``. It is left alone only when somebody else
+        ends it: the runner that cancelled this append on purpose
+        (``response_taken_over``) or the close already under way.
+        """
         if self.pcm_reservation is not None:
             self.pcm_reservation.rollback()
         self.discard_retained_audio()
+        if response_taken_over:
+            return
+        session = self.ctx.session
+        if self.ctx.run.closing or session.state != DuplexSessionState.OPEN:
+            # The close ends the active response and request itself, with the
+            # session's one terminal event after them.
+            return
+        if self.owns_request and session.epoch == self.epoch:
+            # Compare-before-clear: only if the session still points at this
+            # append's request, exactly as the submitted path does.
+            session.clear_request(self.request_id)
+        if self._precreated_response_still_active():
+            self._fail_precreated_response(reason)
 
     def release_on_failure(self, done: asyncio.Task[bool]) -> None:
         """Done-callback: a cancelled or failed append must not hold the audio."""
@@ -105,7 +136,22 @@ class AppendAttempt:
     # ------------------------------------------------------------------ #
 
     async def run_in_wire_order(self, predecessor: asyncio.Task[bool] | None) -> bool:
-        """Wait for the previous append, then submit this one if it still applies."""
+        """Wait for the previous append, then submit this one if it still applies.
+
+        A cancellation is compensated here, wherever it lands (waiting for the
+        predecessor or mid-submit): the reservations go back, and the
+        precreated response is failed unless the runner cancelled the append
+        itself, in which case it ends the response with the status the
+        cancellation was for.
+        """
+        try:
+            return await self._run(predecessor)
+        except asyncio.CancelledError:
+            taken_over = self.ctx.tasks.cancelled_by_runner(asyncio.current_task())
+            self.abandon(reason="append_cancelled", response_taken_over=taken_over)
+            raise
+
+    async def _run(self, predecessor: asyncio.Task[bool] | None) -> bool:
         if predecessor is not None:
             try:
                 predecessor_ok = await predecessor
@@ -127,7 +173,11 @@ class AppendAttempt:
             # Called off, not failed: the chain behind it still runs.
             return True
         if self.pcm_reservation is not None and not self.pcm_reservation.active:
-            self.discard_retained_audio()
+            # The client cleared the input buffer while this append waited its
+            # turn: the audio is gone, so is the response it was going to
+            # answer, and that is the client's doing, not a runtime failure.
+            # The reservation rollback itself is already a no-op.
+            self.abandon(reason="input_cleared")
             return False
         return await self._submit()
 
@@ -152,8 +202,6 @@ class AppendAttempt:
                 self.discard_retained_audio()
             else:
                 self.abandon()
-            if not append_ok and self._precreated_response_still_active():
-                self._fail_precreated_response()
             if not append_ok and session.state == DuplexSessionState.CLOSED:
                 self.ctx.run.runtime_closed = True
                 return False
@@ -171,11 +219,11 @@ class AppendAttempt:
                     self.out.emit_events([session.signal_turn(DuplexTurnEventType.USER_STARTED.value)])
             return append_ok
         except asyncio.CancelledError:
-            if self.pcm_reservation is not None:
-                self.pcm_reservation.rollback()
             raise
         except Exception as exc:
-            self.abandon()
+            # Compensate while the session is still open, so the client gets
+            # the failed ``response.done`` before the error and the close.
+            self.abandon(reason="runtime_append_task_failed")
             logger.exception("Native duplex append task failed: %s", exc)
             self.model.send_runtime_error("runtime_append_task_failed", exc)
             if session.state != DuplexSessionState.CLOSED:
@@ -188,7 +236,7 @@ class AppendAttempt:
             and self.ctx.session.active_response_id == self.precreated_response_id
         )
 
-    def _fail_precreated_response(self) -> None:
+    def _fail_precreated_response(self, reason: str) -> None:
         session = self.ctx.session
         session.end_response(commit_text=False)
         self.out.emit(
@@ -199,7 +247,7 @@ class AppendAttempt:
                 "epoch": session.epoch,
                 "committed": False,
                 "status": "failed",
-                "status_details": {"type": "failed", "reason": "runtime_append_failed"},
+                "status_details": {"type": "failed", "reason": reason},
                 "playback": session.playback.as_dict(),
             }
         )

vllm_omni/engine/duplex/session/context.py

@@ -20,7 +20,7 @@ from __future__ import annotations
 import asyncio
 from collections.abc import Awaitable, Callable
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Protocol, TypeVar
+from typing import TYPE_CHECKING, Any, Protocol, TypeVar
 
 if TYPE_CHECKING:
     from vllm.outputs import RequestOutput
@@ -117,6 +117,10 @@ class DuplexSessionTasks:
     append_tasks: dict[asyncio.Task[bool], DuplexAppendTaskMeta] = field(default_factory=dict)
     append_tail: asyncio.Task[bool] | None = None
     active_response_task: asyncio.Task[None] | None = None
+    #: Appends the runner cancelled on purpose (a cancel, barge-in or close).
+    #: The runner ends the response such an append precreated, so the append
+    #: leaves it alone on its way out instead of failing it.
+    runner_cancelled: set[asyncio.Task[bool]] = field(default_factory=set)
 
     def track_append_task(
         self,
@@ -127,16 +131,31 @@ class DuplexSessionTasks:
         response_bound: bool,
     ) -> None:
         self.append_tasks[task] = DuplexAppendTaskMeta(epoch, final, response_bound)
-        task.add_done_callback(self.append_tasks.pop)
+        task.add_done_callback(self._forget_append_task)
+
+    def _forget_append_task(self, task: asyncio.Task[bool]) -> None:
+        self.append_tasks.pop(task, None)
+        self.runner_cancelled.discard(task)
+
+    def cancelled_by_runner(self, task: asyncio.Task[Any] | None) -> bool:
+        """Whether ``task`` is an append the runner cancelled through ``cancel_append_tasks``."""
+        return task is not None and task in self.runner_cancelled
 
     def has_response_bound_append_tasks(self) -> bool:
         return any(meta.response_bound for meta in self.append_tasks.values())
 
     async def cancel_append_tasks(self, timeout_s: float = 0.25, *, response_bound_only: bool = False) -> bool:
+        """Cancel the tracked appends; the caller owns their responses from here on.
+
+        Every caller ends the active response itself right after (with the
+        cancel status, or silently as part of a close), so the cancelled
+        appends are told not to fail the response they precreated.
+        """
         tasks = [task for task, meta in self.append_tasks.items() if not response_bound_only or meta.response_bound]
         if not tasks:
             return False
         cancelled_tail = self.append_tail if self.append_tail in tasks else None
+        self.runner_cancelled.update(tasks)
         for task in tasks:
             task.cancel()
         try:

vllm_omni/engine/duplex/session/engine_session.py

@@ -280,9 +280,9 @@ class DuplexEngineSession:
         self._validate_fence(self.fence)
         self.lease.detach(self._clock())
 
-    def resume_lease(self, *, expected_lease_generation: int) -> int:
+    def resume_lease(self, *, expected_lease_generation: int, control_id: str | None = None) -> int:
         self._validate_fence(self.fence)
-        return self.lease.resume(self._clock(), expected_generation=expected_lease_generation)
+        return self.lease.resume(self._clock(), expected_generation=expected_lease_generation, control_id=control_id)
 
     def begin_lease_operation(self, fence: DuplexFence, operation_id: str) -> None:
         self._validate_fence(fence)

vllm_omni/engine/duplex/session/lease.py

@@ -40,6 +40,12 @@ class DuplexLeaseState:
     detached_at: float | None = None
     active_operations: set[str] = field(default_factory=set)
     terminal_reason: str | None = None
+    #: Control id of the resume that produced ``generation`` (``None`` for the
+    #: open). A resume replayed with the same id is answered with the current
+    #: generation instead of bumping it again: the caller lost the first
+    #: answer (its RPC waiter was cancelled or timed out) and is finding out
+    #: whether the resume landed.
+    resumed_by: str | None = None
 
     @property
     def expires_at(self) -> float | None:
@@ -61,11 +67,16 @@ class DuplexLeaseState:
         self.touch(now, DuplexLeaseActivity.DETACH)
         self.detached_at = now
 
-    def resume(self, now: float, *, expected_generation: int) -> int:
+    def resume(self, now: float, *, expected_generation: int, control_id: str | None = None) -> int:
         self._require_open()
+        if control_id is not None and control_id == self.resumed_by:
+            # The same resume asking again: it already produced this
+            # generation, so answer without changing anything.
+            return self.generation
         if expected_generation != self.generation:
             raise ValueError(f"duplex lease generation mismatch: expected {self.generation}, got {expected_generation}")
         self.generation += 1
+        self.resumed_by = control_id
         self.detached_at = None
         self.last_activity = now
         return self.generation
@@ -104,6 +115,7 @@ class DuplexLeaseState:
             return False
         self.terminal_reason = reason
         self.generation += 1
+        self.resumed_by = None
         self.detached_at = None
         self.active_operations.clear()
         return True

vllm_omni/engine/duplex/session/manager.py

@@ -523,32 +523,74 @@ class DuplexSessionManager:
                 # and the stage resource straight back.
                 self._abandoned_opens.discard(session_id)
                 await self._close_runner(runner, kind="close", reason="open_abandoned")
+        except asyncio.CancelledError:
+            # A cancel (engine teardown, task cancellation) while awaiting the
+            # plugin or the result sink is not an ``Exception``: without this
+            # branch the runner stayed in ``runners`` with its Stage0
+            # reservation, and since admission counts runners, every cancelled
+            # open burned one slot for good. Nobody is answered: the caller's
+            # RPC waiter is being torn down with us or times out.
+            logger.info("open_duplex_session cancelled for %s; rolling the admission back", session_id)
+            await self._rollback_open(session_id, session=session, runner=runner)
+            raise
         except Exception as exc:
             error_code, _, _ = self._control_error(exc)
             if error_code in {"resource_exhausted", "session_exists"}:
                 logger.info("open_duplex_session rejected: %s", exc)
             else:
                 logger.exception("open_duplex_session failed: %s", exc)
-            if runner is not None and self.runners.get(session_id) is runner:
-                self.runners.pop(session_id, None)
-                try:
-                    await runner.shutdown()
-                except Exception:
-                    logger.exception("duplex open rollback: runner shutdown failed for %s", session_id)
-            if session is not None:
-                reserved = session.release_all_requests()
-                if reserved:
-                    try:
-                        await self.stage_port.cleanup(list(reserved))
-                    except Exception:
-                        logger.warning("duplex open rollback: request cleanup pending for %s", session_id)
-                self._unregister_session_requests(session_id)
+            await self._rollback_open(session_id, session=session, runner=runner)
             await self._put_result(message, operation="open", ok=False, error=exc)
         finally:
             if holds_admission_slot:
                 self._admitting.discard(session_id)
             self._abandoned_opens.discard(session_id)
 
+    async def _rollback_open(
+        self,
+        session_id: str,
+        *,
+        session: DuplexEngineSession | None,
+        runner: DuplexSessionRunner | None,
+    ) -> None:
+        """Undo whatever an open that will not complete already did.
+
+        The bookkeeping that decides admission (``runners`` and the session's
+        request reservations) is undone synchronously first, so a second
+        cancellation landing in the awaits below cannot leave the slot burned.
+        The Stage0 request ids are recorded as a pending request cleanup
+        before anything is awaited: if the runner shutdown or the cleanup
+        itself is cancelled or fails, the reaper retries the cleanup instead
+        of the orchestrator keeping the request state forever.
+        """
+        owned_runner = runner is not None and self.runners.get(session_id) is runner
+        if owned_runner:
+            self.runners.pop(session_id, None)
+        pending: _PendingRequestCleanup | None = None
+        key: tuple[str, int] | None = None
+        if session is not None:
+            reserved = tuple(session.release_all_requests())
+            self._unregister_session_requests(session_id)
+            if reserved:
+                key = (session_id, session.lease_generation)
+                pending = _PendingRequestCleanup(
+                    session_id=session_id,
+                    lease_generation=session.lease_generation,
+                    request_ids=reserved,
+                    abort=False,
+                )
+                self._pending_request_cleanups[key] = pending
+        if owned_runner and runner is not None:
+            try:
+                await runner.shutdown()
+            except Exception:
+                logger.exception("duplex open rollback: runner shutdown failed for %s", session_id)
+        if pending is not None and key is not None:
+            try:
+                await self._complete_request_cleanup(key, pending)
+            except Exception:
+                logger.warning("duplex open rollback: request cleanup pending for %s; the reaper retries", session_id)
+
     def _require_runner(self, session_id: str) -> DuplexSessionRunner:
         runner = self.runners.get(session_id)
         if runner is None:
@@ -668,7 +710,13 @@ class DuplexSessionManager:
             runner = self._require_runner(message.session_id)
             session = runner.session
             try:
-                session.resume_lease(expected_lease_generation=message.expected_lease_generation)
+                # Keyed by the control id, so a caller that lost the answer
+                # (waiter cancelled or timed out) can replay the same resume
+                # to learn whether it landed instead of resuming twice.
+                session.resume_lease(
+                    expected_lease_generation=message.expected_lease_generation,
+                    control_id=message.control_id,
+                )
             except ValueError as exc:
                 raise DuplexSessionError(str(exc), code="session_resume_conflict") from exc
             await self._put_result(message, operation="resume", ok=True, session=session)
@@ -683,6 +731,14 @@ class DuplexSessionManager:
             session = runner.session
             activity = DuplexLeaseActivity(message.activity)
             if activity is DuplexLeaseActivity.DETACH:
+                expected = message.expected_lease_generation
+                if expected is not None and expected != session.lease_generation:
+                    # The caller is giving up a lease it no longer holds: a
+                    # later resume owns the current one and must keep it.
+                    raise DuplexSessionError(
+                        f"duplex lease generation mismatch: expected {session.lease_generation}, got {expected}",
+                        code="session_resume_conflict",
+                    )
                 session.detach_lease()
             else:
                 session.touch_lease(activity)

vllm_omni/engine/duplex/session/model_channel.py

@@ -569,6 +569,12 @@ class ModelChannel:
                 ):
                     return
         engine_output = self._build_stage_output(item)
+        runtime_config = self._ctx.plugin.runtime_config_after_model_output(
+            dict(session.runtime_config),
+            item.context.segment_output_metadata,
+        )
+        if runtime_config is not None:
+            session.replace_runtime_config(runtime_config)
         drain_result = {"data_plane_outputs": [engine_output]}
         close_reason, emitted_response = await self._send_model_output_events(
             drain_result, expected_epoch=expected_epoch

vllm_omni/engine/duplex/session/runner.py

@@ -1107,6 +1107,7 @@ class DuplexSessionRunner:
             operation_id=operation_id,
             retained_committed_payload=retained_committed_payload,
             precreated_response_id=session.active_response_id if precreate_response else None,
+            owns_request=final or precreate_response,
             on_append_accepted=on_append_accepted,
             before_append=before_append,
         )

vllm_omni/engine/duplex_omni_engine.py

@@ -214,9 +214,12 @@ class DuplexOmniEngine(AsyncOmniEngine):
         session_id: str,
         *,
         expected_lease_generation: int,
+        control_id: str | None = None,
         timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S,
     ) -> DuplexControlResultMessage:
-        control_id = uuid.uuid4().hex
+        # A caller-supplied id makes the resume replayable: the engine answers
+        # a repeated id with the generation that resume produced.
+        control_id = control_id or uuid.uuid4().hex
         return self._execute_control(
             ResumeDuplexSessionMessage(
                 control_id=control_id,
@@ -234,6 +237,7 @@ class DuplexOmniEngine(AsyncOmniEngine):
         session_id: str,
         *,
         expected_lease_generation: int,
+        control_id: str | None = None,
         timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S,
     ) -> DuplexControlResultMessage:
         loop = asyncio.get_running_loop()
@@ -242,6 +246,7 @@ class DuplexOmniEngine(AsyncOmniEngine):
             lambda: self._resume_session(
                 session_id,
                 expected_lease_generation=expected_lease_generation,
+                control_id=control_id,
                 timeout=timeout,
             ),
         )
@@ -251,6 +256,7 @@ class DuplexOmniEngine(AsyncOmniEngine):
         session_id: str,
         *,
         activity: str,
+        expected_lease_generation: int | None = None,
         timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S,
     ) -> DuplexControlResultMessage:
         control_id = uuid.uuid4().hex
@@ -259,6 +265,7 @@ class DuplexOmniEngine(AsyncOmniEngine):
                 control_id=control_id,
                 session_id=session_id,
                 activity=activity,
+                expected_lease_generation=expected_lease_generation,
             ),
             control_id=control_id,
             operation="touch",
@@ -271,12 +278,18 @@ class DuplexOmniEngine(AsyncOmniEngine):
         session_id: str,
         *,
         activity: str,
+        expected_lease_generation: int | None = None,
         timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S,
     ) -> DuplexControlResultMessage:
         loop = asyncio.get_running_loop()
         return await loop.run_in_executor(
             None,
-            lambda: self._touch_session(session_id, activity=activity, timeout=timeout),
+            lambda: self._touch_session(
+                session_id,
+                activity=activity,
+                expected_lease_generation=expected_lease_generation,
+                timeout=timeout,
+            ),
         )
 
     def _submit_command(self, session_id: str, command: DuplexCommand) -> None: