Revert "Revert "[CI]Fix Wan DMD pipeline test alignment" (#6574)"

This reverts commit a94b2580a1.

docs/user_guide/diffusion/execution_modes.md

@@ -101,8 +101,16 @@ set `DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA` or configure
 `--max-num-seqs >1`. See the
 [HunyuanImage-3.0 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Tencent/HunyuanImage-3.0-Instruct.md)
 for its validated configuration. Helios supports single-request step execution only: use
-`--step-execution --max-num-seqs 1` for Helios. Consult the selected pipeline's
-documentation and source for the latest support status.
+`--step-execution --max-num-seqs 1` for Helios. MiniMax H3 supports step-wise
+continuous batching by packing co-batched requests into one sequence that keeps
+a separate attention document per request; that layout needs a backend which
+honors the packed `cu_seqlens` metadata, so run it with
+`--diffusion-attention-backend FLASH_ATTN` (other backends stay correct but
+fall back to one transformer forward per request). Batching H3 does not improve
+its throughput — see the measured numbers in the
+[MiniMax-H3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/MiniMaxAI/MiniMax-H3.md)
+— so keep `--max-num-seqs 1` unless you need step-level scheduling. Consult the
+selected pipeline's documentation and source for the latest support status.
 
 ## Streaming Output
 

recipes/MiniMaxAI/MiniMax-H3.md

@@ -380,6 +380,63 @@ No restart is needed: `task=fl2va` routes to `FL2VA/transformer`, while
 `task=ref2va` routes to
 `Ref2VA/transformer`. T2VA uses the FL2VA DiT.
 
+### Step execution and continuous batching
+
+H3 implements the step-wise execution contract, so the scheduler can admit and
+retire requests between denoise steps instead of running one request end to
+end. Add the feature gate, then raise `--max-num-seqs` to co-batch:
+
+```bash
+--step-execution --max-num-seqs 4
+```
+
+Co-batched requests are packed into a single sequence that keeps one attention
+document per request, so a batch costs one DiT forward. That packing requires
+`--diffusion-attention-backend FLASH_ATTN`; other backends fall back to one
+forward per request. `--max-num-seqs 1` keeps the conservative single-request
+step path. Cache acceleration (`--cache-backend`) is not available in step mode.
+
+!!! warning "Co-batching does not improve H3 throughput for large simultaneous requests"
+    Keep `--max-num-seqs 1` unless you specifically need scheduler-level control
+    (admitting and retiring requests between denoise steps). Measured on two
+    H100s (TP2, BF16, 672x384, 209 frames, 30 steps, 4 requests at concurrency
+    4 with `--request-rate inf`, i.e. all four submitted at time 0; one packed
+    request is 16384 rows):
+
+    | Configuration | Wall time | Mean latency | Peak memory |
+    |---------------|-----------|--------------|-------------|
+    | request mode | 174.8 s | 111.5 s | 72.4 GB |
+    | `--step-execution --max-num-seqs 1` | 179.0 s | 113.8 s | 72.4 GB |
+    | `--step-execution --max-num-seqs 4` | 182.1 s | 175.7 s | 78.3 GB |
+
+    A single H3 denoise step is a compute-bound dense GEMM over an already long
+    packed sequence, so fusing N requests costs N times the FLOPs and buys almost
+    no amortization — unlike LLM decoding, which is memory-bandwidth bound.
+    Going from one request per step to four cuts the per-request denoise cost
+    only from 1.323 s to 1.291 s (2.4%), which the step-mode bookkeeping then
+    spends. Mean latency degrades further because co-batched requests finish
+    together instead of staggered. Quantization moves the absolute numbers
+    without changing this: with online `int8` the same workload runs in 153.3 s
+    at 56.9 GB (request mode), and `--max-num-seqs 4` is still 5.0% slower than
+    request mode.
+
+    Two workloads outside this table are unmeasured and may behave differently:
+
+    - **Staggered arrivals (admission latency).** A single 16384-row H3 request
+      already saturates the GPUs, so submitting all requests at time 0 is the
+      one arrival pattern where co-batching cannot win: it can only bunch
+      completions. In request mode a new request queues behind the whole
+      in-flight generation (~45 s at this size); step mode admits at the next
+      denoise-step boundary (~1.3 s). Whether that shows up as a wall-clock
+      benefit under Poisson arrivals is not yet measured for H3. To reproduce,
+      run request mode vs `--step-execution --max-num-seqs 4` with
+      `diffusion_benchmark_serving.py --request-rate 0.05` (roughly one request
+      every 20 s) and report mean / p95 latency, which then includes queueing.
+    - **Small requests.** The numbers above are drawn from 672x384 / 209-frame
+      requests. A short clip at lower resolution packs a few thousand rows and
+      may not saturate the hardware; co-batching may amortize better there.
+      This is also unmeasured.
+
 ### Online FP8 quantization
 
 MiniMax H3 supports online FP8 quantization of both the DiT and the Qwen3-VL
@@ -924,7 +981,12 @@ vllm serve "${MODEL_ROOT}/FL2VA" \
 - TeaCache is calibrated for FL2VA only; Ref2VA requests run uncached.
 - Combined serving requires sibling `FL2VA` and `Ref2VA` directories, loads
   both task-specific DiTs, and loads shared components once from `FL2VA`.
-- H3 currently executes one generation request per diffusion batch.
+- Request mode executes one generation request per diffusion batch. Use
+  `--step-execution` with `--max-num-seqs N` to admit several requests at once
+  (see
+  [Execution modes](https://github.com/vllm-project/vllm-omni/blob/main/docs/user_guide/diffusion/execution_modes.md));
+  step
+  mode does not support `cache_backend`.
 - The first regional-compile request is a warmup and should not be included in
   steady-state performance measurements.
 - The serving path accepts fewer references than the model supports. H3 documents up

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -5,6 +5,7 @@
 Tests for the DiffusersPipelineLoader.
 """
 
+import json
 from types import SimpleNamespace
 
 import pytest
@@ -102,6 +103,138 @@ def _make_dlo_online_quant_config() -> OmniDiffusionConfig:
     )
 
 
+@pytest.mark.parametrize("offline", [False, True])
+def test_prepare_weights_honors_component_index_and_explicit_override(tmp_path, mocker, offline):
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+
+    snapshot = tmp_path / "snapshot"
+    transformer = snapshot / "transformer"
+    transformer.mkdir(parents=True)
+    indexed_files = [f"diffusion_pytorch_model-{index:05d}-of-00002.safetensors" for index in (1, 2)]
+    stale_file = "diffusion_pytorch_model-00001-of-00008.safetensors"
+    index_path = transformer / "diffusion_pytorch_model.safetensors.index.json"
+    index_path.write_text(
+        json.dumps({"weight_map": {f"weight.{index}": filename for index, filename in enumerate(indexed_files)}})
+    )
+    for filename in indexed_files + [stale_file]:
+        (transformer / filename).touch()
+
+    mocker.patch.object(loader_mod.huggingface_hub.constants, "HF_HUB_OFFLINE", offline)
+
+    def download_index(*, filename, **_kwargs):
+        if filename == "transformer/diffusion_pytorch_model.safetensors.index.json":
+            return str(index_path)
+        raise loader_mod.huggingface_hub.errors.EntryNotFoundError(filename)
+
+    hub_api = mocker.Mock()
+    hub_api.hf_hub_download.side_effect = download_index
+    mocker.patch.object(loader_mod, "hf_api", return_value=hub_api)
+    indexed_download = mocker.patch.object(loader_mod, "download_weights_from_hf_specific", return_value=str(snapshot))
+    generic_download = mocker.patch.object(loader_mod, "download_weights_from_hf", return_value=str(snapshot))
+
+    loader = _make_loader_with_weights([])
+    cache_dir = str(tmp_path / "cache")
+    loader.load_config.download_dir = cache_dir
+    folder, files, use_safetensors = loader._prepare_weights(
+        "org/model",
+        subfolder="transformer",
+        revision="revision",
+        fall_back_to_pt=True,
+        allow_patterns_overrides=None,
+    )
+
+    assert folder == str(transformer)
+    assert [str(transformer / filename) for filename in indexed_files] == files
+    assert use_safetensors
+    assert hub_api.hf_hub_download.call_count == 2
+    hub_api.hf_hub_download.assert_any_call(
+        repo_id="org/model",
+        filename="transformer/diffusion_pytorch_model.safetensors.index.json",
+        cache_dir=cache_dir,
+        revision="revision",
+        local_files_only=offline,
+    )
+    indexed_download.assert_called_once_with(
+        model_name_or_path="org/model",
+        cache_dir=cache_dir,
+        allow_patterns=[f"transformer/{filename}" for filename in indexed_files],
+        revision="revision",
+        ignore_patterns=loader.load_config.ignore_patterns,
+        require_all=True,
+    )
+
+    _, override_files, _ = loader._prepare_weights(
+        "org/model",
+        subfolder="transformer",
+        revision="revision",
+        fall_back_to_pt=True,
+        allow_patterns_overrides=[stale_file],
+    )
+    assert override_files == [str(transformer / stale_file)]
+    generic_download.assert_called_once_with(
+        "org/model",
+        cache_dir,
+        [stale_file],
+        "revision",
+        subfolder="transformer",
+        ignore_patterns=loader.load_config.ignore_patterns,
+    )
+
+
+def test_prepare_local_weights_honors_component_index(tmp_path):
+    transformer = tmp_path / "transformer"
+    transformer.mkdir()
+    indexed_file = "diffusion_pytorch_model-00001-of-00001.safetensors"
+    stale_file = "diffusion_pytorch_model-00001-of-00008.safetensors"
+    (transformer / "diffusion_pytorch_model.safetensors.index.json").write_text(
+        json.dumps({"weight_map": {"weight": indexed_file}})
+    )
+    for filename in (indexed_file, stale_file):
+        (transformer / filename).touch()
+
+    folder, files, use_safetensors = _make_loader_with_weights([])._prepare_weights(
+        tmp_path,
+        subfolder="transformer",
+        revision=None,
+        fall_back_to_pt=True,
+        allow_patterns_overrides=None,
+    )
+
+    assert folder == str(transformer)
+    assert files == [str(transformer / indexed_file)]
+    assert use_safetensors
+
+
+def test_prepare_weights_rejects_polluted_offline_cache_without_index(tmp_path, mocker):
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+
+    snapshot = tmp_path / "snapshot"
+    transformer = snapshot / "transformer"
+    transformer.mkdir(parents=True)
+    for shard_count in (4, 8):
+        for shard in range(1, shard_count + 1):
+            (transformer / f"diffusion_pytorch_model-{shard:05d}-of-{shard_count:05d}.safetensors").touch()
+
+    mocker.patch.object(loader_mod.huggingface_hub.constants, "HF_HUB_OFFLINE", True)
+    hub_api = mocker.Mock()
+    hub_api.hf_hub_download.side_effect = loader_mod.huggingface_hub.errors.EntryNotFoundError("index is not cached")
+    mocker.patch.object(loader_mod, "hf_api", return_value=hub_api)
+    mocker.patch.object(loader_mod, "download_weights_from_hf", return_value=str(snapshot))
+
+    loader = _make_loader_with_weights([])
+    with pytest.raises(ValueError, match="conflicting shard totals"):
+        loader._prepare_weights(
+            "org/model",
+            subfolder="transformer",
+            revision="revision",
+            fall_back_to_pt=True,
+            allow_patterns_overrides=None,
+        )
+
+    assert hub_api.hf_hub_download.call_count == 2
+    assert all(call.kwargs["local_files_only"] for call in hub_api.hf_hub_download.call_args_list)
+
+
 def test_strict_check_only_validates_source_prefix_parameters():
     model = _DummyPipelineModel(source_prefix="transformer.")
     loader = _make_loader_with_weights(["transformer.weight"])

tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py

@@ -821,6 +821,82 @@ def test_packed_attention_keeps_padding_mask_for_other_backends():
     )
 
 
+def _fake_packed_attention(
+    backend_name: str,
+    *,
+    prefix_kv_slicing: bool = False,
+    supports_multi_doc: bool = False,
+):
+    from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
+        MiniMaxH3Attention,
+    )
+
+    class FakeBackend:
+        supports_prefix_kv_slicing = prefix_kv_slicing
+
+        @staticmethod
+        def get_name() -> str:
+            return backend_name
+
+        @classmethod
+        def supports_multi_doc_packed_varlen(cls) -> bool:
+            return supports_multi_doc
+
+    class FakeAttention(torch.nn.Module):
+        attn_backend = FakeBackend
+
+        def __init__(self):
+            super().__init__()
+            self.metadata = None
+
+        def forward(self, query, key, value, metadata):
+            self.metadata = metadata
+            return query
+
+    attention = object.__new__(MiniMaxH3Attention)
+    torch.nn.Module.__init__(attention)
+    attention.attention = FakeAttention()
+    return attention
+
+
+def test_packed_attention_drops_prefix_semantics_for_a_co_batched_layout():
+    """A batch's valid rows are block-diagonal, so no prefix can describe them."""
+    attention = _fake_packed_attention("FLASH_ATTN", supports_multi_doc=True)
+    q = torch.randn(12, 2, 4)
+
+    attention._run_packed_attention(
+        q,
+        q,
+        q,
+        # Two requests: rows + padding tail each.
+        cu_seqlens=torch.tensor([0, 5, 6, 11, 12], dtype=torch.int32),
+        max_seqlen=5,
+        packed_total=12,
+        num_requests=2,
+    )
+
+    assert attention.attention.metadata.attn_mask is None
+    assert attention.attention.metadata.extra["valid_kv_length"] == 12
+
+
+@pytest.mark.parametrize("backend_name", ["TORCH_SDPA", "FLASH_ATTN"])
+def test_packed_attention_rejects_backends_without_multi_doc_capability(backend_name):
+    """Backend names alone must not allow cross-request attention."""
+    attention = _fake_packed_attention(backend_name, supports_multi_doc=False)
+    q = torch.randn(12, 2, 4)
+
+    with pytest.raises(ValueError, match=backend_name):
+        attention._run_packed_attention(
+            q,
+            q,
+            q,
+            cu_seqlens=torch.tensor([0, 5, 6, 11, 12], dtype=torch.int32),
+            max_seqlen=5,
+            packed_total=12,
+            num_requests=2,
+        )
+
+
 def test_reference_image_resize_contract():
     from PIL import Image
 

tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py

@@ -0,0 +1,471 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""MiniMax H3 step-wise execution (continuous batching) contract tests.
+
+These run on CPU against a stand-in DiT, so they cover the packing, the
+per-request scheduler math, and the runner-facing state wiring without needing
+checkpoint weights.
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+_HIDDEN = 8
+
+
+class _SegmentMeanModel:
+    """Stand-in DiT whose output depends on packed-document boundaries.
+
+    Every row is mixed with the mean of its own ``cu_seqlens`` document and with
+    its own timestep, so a wrong document boundary, row offset, or timestep
+    assignment shows up as a numeric difference instead of passing silently.
+    """
+
+    def __call__(self, **kwargs):
+        x = kwargs["x"][0]
+        audio_x = kwargs["audio_x"][0]
+        bounds = kwargs["packed_seq_params"]["cu_seqlens_q"].tolist()
+        img_pos = kwargs["img_pos_info"]["position_ids"]
+        audio_pos = kwargs["audio_pos_info"]["position_ids"]
+
+        pooled_video = torch.zeros_like(x)
+        pooled_audio = torch.zeros_like(audio_x)
+        for start, stop in zip(bounds[:-1], bounds[1:]):
+            if stop <= start:
+                continue
+            pooled_video[start:stop] = x[start:stop].mean(dim=0, keepdim=True)
+            pooled_audio[start:stop] = audio_x[start:stop].mean(dim=0, keepdim=True)
+
+        row_timesteps = kwargs["unique_timesteps"][kwargs["inverse_indices"]].unsqueeze(-1)
+        video = (pooled_video + 0.5 * x + row_timesteps)[img_pos]
+        audio = (pooled_audio + 0.5 * audio_x + row_timesteps)[audio_pos]
+        if not kwargs.get("skip_mask_out_condition", False):
+            video = video * kwargs["update_mask"].view(-1).unsqueeze(-1)
+        return video, audio
+
+
+def _make_branch(*, text_len: int, latent_t: int, latent_h: int, latent_w: int, audio_t: int, seed: int):
+    from vllm_omni.diffusion.models.minimax_h3.denoise_loop import MiniMaxH3DenoiseBranch
+    from vllm_omni.diffusion.models.minimax_h3.packed_sequence import minimax_h3_packed_sequence
+
+    packed = minimax_h3_packed_sequence(
+        text_len=text_len,
+        latent_t=latent_t,
+        latent_h=latent_h,
+        latent_w=latent_w,
+        audio_t=audio_t,
+        include_keyframe_cond=False,
+    )
+    generator = torch.Generator().manual_seed(seed)
+    text_embeddings = torch.randn(text_len, _HIDDEN, generator=generator, dtype=torch.float32)
+    branch = MiniMaxH3DenoiseBranch(
+        packed=packed,
+        text_embeddings=text_embeddings,
+        token_tags=packed["token_tags"],
+        device=torch.device("cpu"),
+    )
+    video_rows = torch.randn(int(branch.img_pos.shape[0]), 96, generator=generator, dtype=torch.float32)
+    audio_rows = torch.randn(int(branch.audio_pos.shape[0]), 32, generator=generator, dtype=torch.float32)
+    return branch, video_rows, audio_rows
+
+
+def _sigmas(num_points: int, shift: float) -> list[float]:
+    from vllm_omni.diffusion.models.minimax_h3.time_request import minimax_h3_time_shift_sigmas
+
+    return minimax_h3_time_shift_sigmas(num_steps=num_points, shift_scale=shift)
+
+
+def _make_state(request_id: str, model, branch, video_rows, audio_rows, sigmas_video, sigmas_audio):
+    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod
+    from vllm_omni.diffusion.worker.utils import StepRequestState
+
+    state = StepRequestState(request_id=request_id, sampling=SimpleNamespace())
+    state.latents = video_rows.clone()
+    state.timesteps = torch.tensor([1.0 - sigma for sigma in sigmas_video[:-1]], dtype=torch.float32)
+    state.step_index = 0
+    state.extra = {
+        mod._STEP_BRANCH: branch,
+        # Co-batched requests must share one DiT instance, or denoise_step()
+        # treats the batch as mixed-task and falls back to one forward each.
+        mod._STEP_TRANSFORMER: model,
+        mod._STEP_AUDIO_ROWS: audio_rows.clone(),
+        mod._STEP_COND_ANCHOR: None,
+        mod._STEP_AUDIO_ANCHOR: None,
+        mod._STEP_SIGMAS_VIDEO: sigmas_video,
+        mod._STEP_SIGMAS_AUDIO: sigmas_audio,
+    }
+    return state
+
+
+def _step_pipeline(model, *, packed_batch_supported: bool = True):
+    """A pipeline instance carrying only what the step methods touch."""
+    from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    pipeline.transformer = model
+    pipeline.device = torch.device("cpu")
+    pipeline._transformer_for_task = lambda task: model
+    pipeline._packed_batch_supported = lambda transformer: packed_batch_supported
+    return pipeline
+
+
+def test_step_execution_matches_request_mode_denoise_loop():
+    """Stepping through the contract must reproduce the request-mode loop."""
+    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod
+    from vllm_omni.diffusion.models.minimax_h3.denoise_loop import minimax_h3_denoise_loop
+
+    model = _SegmentMeanModel()
+    branch, video_rows, audio_rows = _make_branch(text_len=9, latent_t=2, latent_h=4, latent_w=6, audio_t=3, seed=5)
+    sigmas_video = _sigmas(6, 12.0)
+    sigmas_audio = _sigmas(6, 3.0)
+
+    reference_video, reference_audio = minimax_h3_denoise_loop(
+        model=model,
+        positive=branch,
+        initial_video_rows=video_rows,
+        initial_audio_rows=audio_rows,
+        keyframe_cond_rows=None,
+        sigmas_video=sigmas_video,
+        sigmas_audio=sigmas_audio,
+        device=torch.device("cpu"),
+    )
+
+    pipeline = _step_pipeline(model)
+    state = _make_state("req-0", model, branch, video_rows, audio_rows, sigmas_video, sigmas_audio)
+    input_batch = SimpleNamespace(states=(state,))
+
+    steps = 0
+    while not state.denoise_completed:
+        noise_pred = pipeline.denoise_step(input_batch, states=[state])
+        pipeline.step_scheduler(state, noise_pred)
+        steps += 1
+
+    assert steps == len(sigmas_video) - 1
+    assert state.total_steps == steps
+    torch.testing.assert_close(state.latents, reference_video)
+    torch.testing.assert_close(state.extra[mod._STEP_AUDIO_ROWS], reference_audio)
+
+
+def test_batched_step_execution_matches_independent_requests():
+    """Two co-batched requests must land where they would have landed alone."""
+    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod
+
+    model = _SegmentMeanModel()
+    specs = [
+        dict(text_len=9, latent_t=2, latent_h=4, latent_w=6, audio_t=3, seed=6),
+        # Exactly 64 packed rows exercises the no-padding-tail boundary case.
+        dict(text_len=46, latent_t=7, latent_h=2, latent_w=4, audio_t=2, seed=7),
+    ]
+    # Different step counts, so the batch composition changes mid-flight.
+    schedules = [(_sigmas(6, 12.0), _sigmas(6, 3.0)), (_sigmas(4, 12.0), _sigmas(4, 3.0))]
+
+    pipeline = _step_pipeline(model)
+
+    alone: list[tuple[torch.Tensor, torch.Tensor]] = []
+    for spec, (sigmas_video, sigmas_audio) in zip(specs, schedules):
+        branch, video_rows, audio_rows = _make_branch(**spec)
+        state = _make_state("solo", model, branch, video_rows, audio_rows, sigmas_video, sigmas_audio)
+        while not state.denoise_completed:
+            pipeline.step_scheduler(state, pipeline.denoise_step(SimpleNamespace(states=(state,)), states=[state]))
+        alone.append((state.latents, state.extra[mod._STEP_AUDIO_ROWS]))
+
+    states = []
+    for index, (spec, (sigmas_video, sigmas_audio)) in enumerate(zip(specs, schedules)):
+        branch, video_rows, audio_rows = _make_branch(**spec)
+        if index == 1:
+            assert branch.used_len == branch.seq_len
+        states.append(_make_state(f"req-{index}", model, branch, video_rows, audio_rows, sigmas_video, sigmas_audio))
+
+    active = list(states)
+    while active:
+        noise_pred = pipeline.denoise_step(SimpleNamespace(states=tuple(active)), states=active)
+        offset = 0
+        for state in active:
+            rows = state.latents.shape[0]
+            pipeline.step_scheduler(state, noise_pred[offset : offset + rows])
+            offset += rows
+        assert offset == noise_pred.shape[0]
+        # Finished requests leave the batch, exactly like the runner drops them.
+        active = [state for state in active if not state.denoise_completed]
+
+    for state, (expected_video, expected_audio) in zip(states, alone):
+        torch.testing.assert_close(state.latents, expected_video)
+        torch.testing.assert_close(state.extra[mod._STEP_AUDIO_ROWS], expected_audio)
+
+
+def test_both_modes_publish_denoise_progress_for_gated_attention():
+    """TRTLLM's skip gate and RAINFUSION's warmup stay dense without these."""
+    from vllm_omni.diffusion import forward_context as fc
+    from vllm_omni.diffusion.models.minimax_h3.denoise_loop import minimax_h3_denoise_loop
+
+    model = _SegmentMeanModel()
+    branch, video_rows, audio_rows = _make_branch(text_len=9, latent_t=2, latent_h=4, latent_w=6, audio_t=3, seed=14)
+    sigmas_video = _sigmas(4, 12.0)
+    sigmas_audio = _sigmas(4, 3.0)
+
+    published: list[tuple[int | None, float | None]] = []
+
+    class _Recorder:
+        def __init__(self):
+            self.denoise_step_idx = None
+            self.denoise_timestep = None
+
+        def __setattr__(self, name, value):
+            object.__setattr__(self, name, value)
+            if name == "denoise_timestep":
+                published.append((self.denoise_step_idx, value))
+
+    recorder = _Recorder()
+    published.clear()  # drop the pair __init__ recorded
+    original = fc._forward_context
+    fc._forward_context = recorder
+    try:
+        minimax_h3_denoise_loop(
+            model=model,
+            positive=branch,
+            initial_video_rows=video_rows,
+            initial_audio_rows=audio_rows,
+            keyframe_cond_rows=None,
+            sigmas_video=sigmas_video,
+            sigmas_audio=sigmas_audio,
+            device=torch.device("cpu"),
+        )
+        request_mode = list(published)
+
+        published.clear()
+        pipeline = _step_pipeline(model)
+        state = _make_state("req-0", model, branch, video_rows, audio_rows, sigmas_video, sigmas_audio)
+        while not state.denoise_completed:
+            pipeline.step_scheduler(state, pipeline.denoise_step(SimpleNamespace(states=(state,)), states=[state]))
+        step_mode = list(published)
+    finally:
+        fc._forward_context = original
+
+    num_steps = len(sigmas_video) - 1
+    expected = [(step, sigmas_video[step]) for step in range(num_steps)]
+    assert request_mode == expected + [(None, None)]
+    # Step mode has no loop to close, so it publishes only the per-step pairs.
+    assert step_mode == expected
+
+
+def test_mixed_step_batch_leaves_gated_attention_dense():
+    """A batch spanning different steps has no single progress point to publish."""
+    from vllm_omni.diffusion import forward_context as fc
+
+    model = _SegmentMeanModel()
+    first, first_video, first_audio = _make_branch(text_len=9, latent_t=2, latent_h=4, latent_w=6, audio_t=3, seed=15)
+    second, second_video, second_audio = _make_branch(
+        text_len=5, latent_t=3, latent_h=6, latent_w=4, audio_t=2, seed=16
+    )
+    sigmas_video, sigmas_audio = _sigmas(6, 12.0), _sigmas(6, 3.0)
+    states = [
+        _make_state("req-0", model, first, first_video, first_audio, sigmas_video, sigmas_audio),
+        _make_state("req-1", model, second, second_video, second_audio, sigmas_video, sigmas_audio),
+    ]
+    states[1].step_index = 2  # admitted later, so it trails the batch
+
+    recorder = SimpleNamespace(denoise_step_idx="unset", denoise_timestep="unset")
+    original = fc._forward_context
+    fc._forward_context = recorder
+    try:
+        _step_pipeline(model).denoise_step(SimpleNamespace(states=tuple(states)), states=states)
+    finally:
+        fc._forward_context = original
+
+    assert recorder.denoise_step_idx is None
+    assert recorder.denoise_timestep is None
+
+
+def test_prepare_encode_seeds_runner_visible_state(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod
+
+    branch, video_rows, audio_rows = _make_branch(text_len=9, latent_t=2, latent_h=4, latent_w=6, audio_t=3, seed=8)
+    sigmas_video = _sigmas(6, 12.0)
+    sigmas_audio = _sigmas(6, 3.0)
+    context = {
+        "height": 96,
+        "width": 64,
+        "latent_t": 2,
+        "latent_h": 4,
+        "latent_w": 6,
+        "audio_t": 3,
+        **{key: None for key in mod._MINIMAX_H3_DENOISE_INPUT_KEYS},
+    }
+
+    pipeline = _step_pipeline(_SegmentMeanModel())
+    monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_extract_prompt", staticmethod(lambda _: ("a prompt", {})))
+    monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_prepare_request_inputs", lambda self, **_: context)
+    monkeypatch.setattr(
+        mod.MiniMaxH3Pipeline,
+        "_build_denoise_inputs",
+        lambda self, **_: {
+            "branch": branch,
+            "video_rows": video_rows,
+            "audio_rows": audio_rows,
+            "cond_anchor": None,
+            "audio_anchor": None,
+            "sigmas_video": sigmas_video,
+            "sigmas_audio": sigmas_audio,
+        },
+    )
+
+    from vllm_omni.diffusion.worker.utils import StepRequestState
+
+    state = StepRequestState(
+        request_id="req-0",
+        sampling=SimpleNamespace(num_outputs_per_prompt=1),
+        prompt="a prompt",
+    )
+    pipeline.prepare_encode(state)
+
+    # The runner slices the batched velocity by this row count.
+    assert state.latents.shape == (int(branch.img_pos.shape[0]), 96)
+    assert state.total_steps == len(sigmas_video) - 1
+    assert state.step_index == 0
+    assert state.do_true_cfg is False
+    torch.testing.assert_close(state.current_timestep, torch.tensor(1.0 - sigmas_video[0]))
+    assert state.extra[mod._STEP_BRANCH] is branch
+    assert state.extra[mod._STEP_SHAPE]["height"] == 96
+
+
+def test_prepare_encode_rejects_request_mode_only_features():
+    """Multi-output and DLO have no representation in the step contract."""
+    from vllm_omni.diffusion.worker.utils import StepRequestState
+    from vllm_omni.errors import OmniClientError
+
+    # A request state carries exactly one latent tensor.
+    multi_output = StepRequestState(
+        request_id="req-0",
+        sampling=SimpleNamespace(num_outputs_per_prompt=2),
+        prompt="a prompt",
+    )
+    with pytest.raises(OmniClientError, match="one output per request"):
+        _step_pipeline(_SegmentMeanModel()).prepare_encode(multi_output)
+
+    # Distributed layerwise offload streams the DiT around a whole denoise loop.
+    single_output = StepRequestState(
+        request_id="req-1",
+        sampling=SimpleNamespace(num_outputs_per_prompt=1),
+        prompt="a prompt",
+    )
+    dlo_pipeline = _step_pipeline(_SegmentMeanModel())
+    dlo_pipeline._dlo_residency_controller = object()
+    with pytest.raises(ValueError, match="distributed layerwise offload"):
+        dlo_pipeline.prepare_encode(single_output)
+
+
+def test_prepare_encode_rejects_high_quality_cache_dit():
+    """quality=high installs a transformer-scoped Cache-DiT profile that
+    would leak across interleaved or co-batched step-mode requests."""
+    from vllm_omni.diffusion.worker.utils import StepRequestState
+    from vllm_omni.errors import OmniClientError
+
+    state = StepRequestState(
+        request_id="req-hi",
+        sampling=SimpleNamespace(num_outputs_per_prompt=1, quality="high"),
+        prompt="a prompt",
+    )
+    with pytest.raises(OmniClientError, match="quality=high"):
+        _step_pipeline(_SegmentMeanModel()).prepare_encode(state)
+
+
+def _fake_attention_module(
+    *,
+    use_ring: bool,
+    backend: str = "FLASH_ATTN",
+    supports_multi_doc: bool = True,
+):
+    """Build a bare MiniMaxH3Attention with the attributes ``_packed_batch_supported`` reads.
+
+    ``supports_multi_doc`` models the platform-dependent capability probe on
+    ``AttentionBackend`` (e.g. FLASH_ATTN returns True on CUDA/ROCm/MUSA and
+    False on NPU/XPU); ``backend`` remains only for logging/back-compat.
+    """
+    from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import MiniMaxH3Attention
+
+    attn = object.__new__(MiniMaxH3Attention)
+    attn.attention = SimpleNamespace(
+        attn_backend=SimpleNamespace(
+            get_name=lambda: backend,
+            supports_multi_doc_packed_varlen=lambda: supports_multi_doc,
+        ),
+        use_ring=use_ring,
+    )
+    return attn
+
+
+class _FakeTransformer:
+    """The minimal ``modules()`` protocol ``_packed_batch_supported`` walks."""
+
+    def __init__(self, attentions):
+        self._attentions = list(attentions)
+
+    def modules(self):
+        return iter([self, *self._attentions])
+
+
+@pytest.mark.parametrize(
+    "attention",
+    [
+        _fake_attention_module(use_ring=True),
+        _fake_attention_module(use_ring=False, backend="XFORMERS", supports_multi_doc=False),
+        _fake_attention_module(use_ring=False, backend="FLASH_ATTN", supports_multi_doc=False),
+    ],
+)
+def test_packed_batch_rejects_backends_that_cannot_isolate_requests(attention):
+    from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline
+
+    assert MiniMaxH3Pipeline._packed_batch_supported(_FakeTransformer([attention])) is False
+
+
+def test_broadcast_rank0_exception_single_rank_reraises():
+    """Single-rank execution has no group; the helper just reraises."""
+    from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
+        _broadcast_rank0_exception,
+    )
+    from vllm_omni.errors import OmniClientError
+
+    _broadcast_rank0_exception(None)
+    with pytest.raises(OmniClientError, match="bad ref"):
+        _broadcast_rank0_exception(OmniClientError("bad ref"))
+
+
+def test_broadcast_rank0_exception_propagates_to_non_zero_ranks(monkeypatch):
+    """A rank-0 error becomes a matching client error on every other DiT rank."""
+    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod
+    from vllm_omni.errors import OmniClientError
+
+    def fake_rank_world(rank):
+        return lambda: (object(), rank, 4)
+
+    def make_broadcast(rank0_payload):
+        def fake_broadcast(payload_list, *, src, group):
+            payload_list[0] = rank0_payload
+
+        return fake_broadcast
+
+    monkeypatch.setattr(mod, "_dit_rank_world", fake_rank_world(0))
+    err = OmniClientError("invalid reference-video file", status_code=422, error_type="UnprocessableEntityError")
+    rank0_payload = {
+        "type": type(err).__name__,
+        "message": str(err),
+        "status_code": err.status_code,
+        "error_type": err.error_type,
+    }
+    monkeypatch.setattr(mod.dist, "broadcast_object_list", make_broadcast(rank0_payload))
+    with pytest.raises(OmniClientError) as rank0_info:
+        mod._broadcast_rank0_exception(err)
+    assert rank0_info.value is err
+
+    monkeypatch.setattr(mod, "_dit_rank_world", fake_rank_world(2))
+    with pytest.raises(OmniClientError) as rank2_info:
+        mod._broadcast_rank0_exception(None)
+    assert rank2_info.value.status_code == 422
+    assert rank2_info.value.error_type == "UnprocessableEntityError"
+    assert "invalid reference-video file" in str(rank2_info.value)

tests/diffusion/test_diffusion_scheduler.py

@@ -25,7 +25,7 @@ from vllm_omni.diffusion.sched import (
     StepScheduler,
 )
 from vllm_omni.diffusion.sched.interface import CachedRequestData, NewRequestData
-from vllm_omni.diffusion.worker.utils import RunnerOutput
+from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
@@ -1559,6 +1559,39 @@ class TestStepScheduler:
         assert sched_output.num_running_reqs == 2
         assert sched_output.num_waiting_reqs == 0
 
+    def test_mixed_batch_error_finishes_only_failed_request(self) -> None:
+        scheduler = StepScheduler()
+        scheduler.initialize(SimpleNamespace(max_num_seqs=2))
+
+        req_bad = scheduler.add_request(_make_step_request("bad", num_inference_steps=2))
+        req_good = scheduler.add_request(_make_step_request("good", num_inference_steps=2))
+        first = scheduler.schedule()
+
+        finished = scheduler.update_from_output(
+            first,
+            BatchRunnerOutput.from_list(
+                [
+                    _make_step_output(
+                        req_bad,
+                        step_index=0,
+                        finished=True,
+                        error="invalid request input",
+                    ),
+                    _make_step_output(req_good, step_index=1),
+                ]
+            ),
+        )
+
+        assert finished == {req_bad}
+        assert scheduler.get_request_state(req_bad).status == DiffusionRequestStatus.FINISHED_ERROR
+        assert scheduler.get_request_state(req_good).status == DiffusionRequestStatus.RUNNING
+        assert scheduler.get_request_state(req_good).req.sampling_params.step_index == 1
+
+        second = scheduler.schedule()
+        assert second.finished_req_ids == {req_bad}
+        assert _new_ids(second) == []
+        assert _cached_ids(second) == [req_good]
+
     def test_step_batch_rejects_different_sampling_key(self) -> None:
         scheduler = StepScheduler()
         scheduler.initialize(SimpleNamespace(max_num_seqs=3))

tests/diffusion/test_diffusion_step_pipeline.py

@@ -47,6 +47,7 @@ from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
 from vllm_omni.diffusion.worker.input_batch import InputBatch
 from vllm_omni.diffusion.worker.utils import RunnerOutput, StepRequestState
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
+from vllm_omni.errors import OmniClientError
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform
 
@@ -102,6 +103,22 @@ class _StepPipeline:
         return DiffusionOutput(output=torch.tensor([state.step_index], dtype=torch.float32))
 
 
+class _PerRequestErrorStepPipeline(_StepPipeline):
+    def prepare_encode(self, state, **kwargs):
+        if state.prompt == "fail-prepare":
+            raise OmniClientError(
+                "invalid request input",
+                status_code=422,
+                error_type="UnprocessableEntityError",
+            )
+        return super().prepare_encode(state, **kwargs)
+
+    def denoise_step(self, input_batch, **kwargs):
+        del kwargs
+        self.denoise_calls += 1
+        return torch.ones_like(input_batch.latents)
+
+
 class _ProfilingStepPipeline(_StepPipeline):
     enable_diffusion_pipeline_profiler = True
 
@@ -697,6 +714,72 @@ class TestRunner:
         result = DiffusionModelRunner.execute_stepwise(runner, scheduler_output)
         assert len(result) == 2
 
+    def test_prepare_error_isolated_from_healthy_request(self, monkeypatch):
+        runner = _make_runner()
+        runner.pipeline = _PerRequestErrorStepPipeline()
+        bad_req = _make_step_request()
+        bad_req.request_id = "req-bad"
+        bad_req.prompt = "fail-prepare"
+        good_req = _make_step_request()
+        good_req.request_id = "req-good"
+        good_req.prompt = "valid"
+        monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+
+        result = DiffusionModelRunner.execute_stepwise(
+            runner,
+            _make_batch_scheduler_output([bad_req, good_req]),
+        )
+
+        bad_output = result.get_request_output("req-bad")
+        assert bad_output.finished is True
+        assert bad_output.result is not None
+        assert bad_output.result.error == "invalid request input"
+        assert bad_output.result.error_status_code == 422
+        assert bad_output.result.error_type == "UnprocessableEntityError"
+        good_output = result.get_request_output("req-good")
+        assert good_output.finished is False
+        assert good_output.step_index == 1
+        assert good_output.result is None
+        assert set(runner.state_cache) == {"req-good"}
+        assert runner.pipeline.denoise_calls == 1
+        assert runner.pipeline.scheduler_calls == 1
+
+    def test_prepare_error_on_other_rank_is_isolated_via_all_reduce(self, monkeypatch):
+        """P1-B: rank-0-only pipeline work (e.g. H3 reference-video prep)
+        can raise on rank 0 while non-zero ranks return cleanly and then hang
+        on a downstream ``dist.broadcast``. Beyond the pipeline-side guard,
+        the runner all-reduces a per-request failure flag across the DiT
+        group so a rank that ``prepare_encode`` returned cleanly on still
+        skips the request and does not enter the denoise step batch.
+        """
+        runner = _make_runner()
+        runner.pipeline = _StepPipeline()  # locally always succeeds
+        monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+
+        # Simulate a peer DiT rank having reported failure. We patch the
+        # helper directly (rather than ``torch.distributed`` broadly) because
+        # ``execute_stepwise`` also queries ``torch.distributed.get_rank``.
+        monkeypatch.setattr(
+            model_runner_module,
+            "_dit_any_rank_failed",
+            lambda local_failed: True,
+        )
+
+        req = _make_step_request(num_inference_steps=1)
+        result = DiffusionModelRunner.execute_stepwise(
+            runner,
+            _make_scheduler_output(req),
+        )
+
+        output = result.get_request_output("req-1")
+        assert output.finished is True
+        assert output.result is not None
+        assert output.result.error is not None
+        assert "another DiT rank" in output.result.error
+        # A rank that skipped the request must not have driven a denoise step.
+        assert runner.pipeline.denoise_calls == 0
+        assert runner.state_cache == {}
+
     def test_receives_kv_payload_before_prepare_encode(self, monkeypatch):
         runner = _make_runner()
         captured: dict[str, object] = {}

tests/engine/test_async_omni_engine_abort_ack.py

@@ -0,0 +1,156 @@
+"""Unit tests for acknowledged abort on AsyncOmniEngine."""
+
+from __future__ import annotations
+
+import asyncio
+import queue
+import threading
+from types import SimpleNamespace
+
+import pytest
+from pytest_mock import MockerFixture
+
+from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
+from vllm_omni.engine.messages import AbortRequestMessage, AbortResultMessage
+from vllm_omni.engine.rpc_result_router import CorrelatedRpcClient
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_engine(request_q: queue.Queue, rpc_q: queue.Queue) -> AsyncOmniEngine:
+    engine = object.__new__(AsyncOmniEngine)
+    engine.request_queue = SimpleNamespace(sync_q=request_q)
+    engine.rpc_output_queue = SimpleNamespace(sync_q=rpc_q)
+    engine._correlated_rpc_client = CorrelatedRpcClient(request_q, rpc_q)
+    return engine
+
+
+def test_abort_result_message_correlation_key():
+    msg = AbortResultMessage(rpc_id="abc", success=True)
+    assert msg.rpc_correlation_key == ("abort", "abc")
+
+
+def test_sync_abort_is_fire_and_forget():
+    request_q: queue.Queue = queue.Queue()
+    rpc_q: queue.Queue = queue.Queue()
+    engine = _make_engine(request_q, rpc_q)
+
+    engine.abort(["req-1"])
+
+    msg = request_q.get_nowait()
+    assert isinstance(msg, AbortRequestMessage)
+    assert msg.request_ids == ["req-1"]
+    assert msg.rpc_id is None
+    assert rpc_q.empty()
+    engine._correlated_rpc_client.close()
+
+
+def test_abort_async_waits_for_ack(mocker: MockerFixture):
+    request_q: queue.Queue = queue.Queue()
+    rpc_q: queue.Queue = queue.Queue()
+    engine = _make_engine(request_q, rpc_q)
+    mocker.patch(
+        "vllm_omni.engine.async_omni_engine.uuid.uuid4",
+        return_value=SimpleNamespace(hex="abort-rpc-1"),
+    )
+
+    async def _run() -> None:
+        task = asyncio.create_task(engine.abort_async(["req-a", "req-b"], timeout=1))
+        msg = await asyncio.to_thread(request_q.get, True, 1)
+        assert isinstance(msg, AbortRequestMessage)
+        assert msg.request_ids == ["req-a", "req-b"]
+        assert msg.rpc_id == "abort-rpc-1"
+        assert not task.done()
+
+        rpc_q.put(AbortResultMessage(rpc_id="abort-rpc-1", success=True))
+        await task
+
+    try:
+        asyncio.run(_run())
+    finally:
+        engine._correlated_rpc_client.close()
+
+
+def test_abort_async_raises_on_orchestrator_error(mocker: MockerFixture):
+    request_q: queue.Queue = queue.Queue()
+    rpc_q: queue.Queue = queue.Queue()
+    engine = _make_engine(request_q, rpc_q)
+    mocker.patch(
+        "vllm_omni.engine.async_omni_engine.uuid.uuid4",
+        return_value=SimpleNamespace(hex="abort-rpc-err"),
+    )
+
+    async def _run() -> None:
+        task = asyncio.create_task(engine.abort_async(["req-x"], timeout=1))
+        await asyncio.to_thread(request_q.get, True, 1)
+        rpc_q.put(
+            AbortResultMessage(
+                rpc_id="abort-rpc-err",
+                success=False,
+                error="stage abort failed",
+            )
+        )
+        with pytest.raises(RuntimeError, match="stage abort failed"):
+            await task
+
+    try:
+        asyncio.run(_run())
+    finally:
+        engine._correlated_rpc_client.close()
+
+
+def test_abort_async_times_out_without_result(mocker: MockerFixture):
+    request_q: queue.Queue = queue.Queue()
+    rpc_q: queue.Queue = queue.Queue()
+    engine = _make_engine(request_q, rpc_q)
+    mocker.patch(
+        "vllm_omni.engine.async_omni_engine.uuid.uuid4",
+        return_value=SimpleNamespace(hex="abort-rpc-timeout"),
+    )
+
+    async def _run() -> None:
+        with pytest.raises(TimeoutError, match="abort timed out"):
+            await engine.abort_async(["req-timeout"], timeout=0.05)
+
+    try:
+        asyncio.run(_run())
+    finally:
+        engine._correlated_rpc_client.close()
+
+
+def test_abort_async_preserves_request_queue_backpressure(mocker: MockerFixture):
+    class SignallingQueue(queue.Queue):
+        def __init__(self) -> None:
+            super().__init__(maxsize=1)
+            self.put_attempted = threading.Event()
+
+        def put(self, item, block=True, timeout=None):
+            self.put_attempted.set()
+            return super().put(item, block=block, timeout=timeout)
+
+    request_q = SignallingQueue()
+    request_q.put("queue-is-full")
+    request_q.put_attempted.clear()
+    rpc_q: queue.Queue = queue.Queue()
+    engine = _make_engine(request_q, rpc_q)
+    mocker.patch(
+        "vllm_omni.engine.async_omni_engine.uuid.uuid4",
+        return_value=SimpleNamespace(hex="blocked-abort"),
+    )
+
+    async def _run() -> None:
+        task = asyncio.create_task(engine.abort_async(["req-blocked"], timeout=1))
+        assert await asyncio.to_thread(request_q.put_attempted.wait, 1)
+        assert not task.done()
+
+        assert request_q.get(timeout=1) == "queue-is-full"
+        msg = await asyncio.to_thread(request_q.get, True, 1)
+        assert isinstance(msg, AbortRequestMessage)
+        assert msg.rpc_id == "blocked-abort"
+        rpc_q.put(AbortResultMessage(rpc_id="blocked-abort", success=True))
+        await task
+
+    try:
+        asyncio.run(_run())
+    finally:
+        engine._correlated_rpc_client.close()

tests/engine/test_orchestrator.py

@@ -19,6 +19,7 @@ from vllm.v1.metrics.stats import IterationStats
 
 from vllm_omni.engine.messages import (
     AbortRequestMessage,
+    AbortResultMessage,
     AddCompanionRequestMessage,
     CollectiveRPCRequestMessage,
     CollectiveRPCResultMessage,
@@ -805,10 +806,74 @@ async def test_run_abort(orchestrator_factory) -> None:
         assert stages[0].abort_calls == [["req-abort"]]
         assert stages[1].abort_calls == []
         assert "req-abort" not in orchestrator_fixture.orchestrator.request_states
+        # Fire-and-forget abort must not emit an RPC result.
+        assert orchestrator_fixture.queues[2].sync_q.empty()
     finally:
         await _shutdown_orchestrator(orchestrator_fixture)
 
 
+@pytest.mark.asyncio
+async def test_run_abort_emits_result_when_rpc_id_set(orchestrator_factory) -> None:
+    stages = [
+        FakeStageClient(stage_type="llm", final_output=False),
+        FakeStageClient(stage_type="llm", final_output=True),
+    ]
+    processors = [
+        FakeOutputProcessor(request_outputs=[_build_request_output("req-ack", token_ids=[1], finished=True)]),
+        FakeOutputProcessor(request_outputs=[_build_request_output("req-ack", token_ids=[2], finished=True)]),
+    ]
+    orchestrator_fixture = orchestrator_factory(stages, output_processors=processors)
+    request = SimpleNamespace(request_id="req-ack", prompt_token_ids=[1, 2, 3])
+
+    try:
+        await _enqueue_add_request(
+            orchestrator_fixture,
+            request_id="req-ack",
+            prompt=request,
+            original_prompt={"prompt": "cancel with ack"},
+            sampling_params_list=[_sampling_params(), _sampling_params()],
+            final_stage_id=1,
+        )
+        await _wait_for(lambda: len(stages[0].add_request_calls) == 1)
+
+        orchestrator_fixture.request_sync_q.put_nowait(AbortRequestMessage(request_ids=["req-ack"], rpc_id="abort-1"))
+        result = await _get_rpc_message(orchestrator_fixture)
+        assert isinstance(result, AbortResultMessage)
+        assert result.rpc_id == "abort-1"
+        assert result.success is True
+        assert result.error is None
+        assert result.rpc_correlation_key == ("abort", "abort-1")
+        assert stages[0].abort_calls == [["req-ack"]]
+        assert "req-ack" not in orchestrator_fixture.orchestrator.request_states
+    finally:
+        await _shutdown_orchestrator(orchestrator_fixture)
+
+
+@pytest.mark.asyncio
+async def test_handle_abort_emits_error_result_on_failure() -> None:
+    rpc_queue: asyncio.Queue = asyncio.Queue()
+    orchestrator = Orchestrator(
+        request_async_queue=asyncio.Queue(),
+        output_async_queue=asyncio.Queue(),
+        rpc_async_queue=rpc_queue,
+        stage_pools=[],
+    )
+
+    async def boom(_request_ids, *, abort=False):
+        del _request_ids, abort
+        raise RuntimeError("cleanup exploded")
+
+    orchestrator._cleanup_request_ids = boom  # type: ignore[method-assign]
+
+    await orchestrator._handle_abort(AbortRequestMessage(request_ids=["req-fail"], rpc_id="abort-err"))
+
+    result = rpc_queue.get_nowait()
+    assert isinstance(result, AbortResultMessage)
+    assert result.rpc_id == "abort-err"
+    assert result.success is False
+    assert "cleanup exploded" in (result.error or "")
+
+
 def _duplex_open_message(
     session_id: str,
     *,

tests/engine/test_rpc_result_router.py

@@ -4,6 +4,7 @@ from typing import Literal
 import pytest
 
 from vllm_omni.engine.messages import (
+    AbortResultMessage,
     CollectiveRPCResultMessage,
     EngineQueueMessage,
     ErrorMessage,
@@ -93,6 +94,19 @@ def test_rpc_result_router_does_not_broadcast_uncorrelated_nonfatal_errors():
     router.close()
 
 
+def test_rpc_result_router_routes_abort_results_by_correlation_id():
+    source: queue.Queue = queue.Queue()
+    router = RpcResultRouter(source)
+    waiter = router.register(("abort", "abort-42"))
+
+    source.put(AbortResultMessage(rpc_id="abort-42", success=True))
+
+    result = waiter.get(timeout=1)
+    assert isinstance(result, AbortResultMessage)
+    assert result.rpc_id == "abort-42"
+    router.close()
+
+
 def test_rpc_result_router_close_unblocks_waiters_and_stops_consumer():
     source: queue.Queue = queue.Queue()
     router = RpcResultRouter(source)

tests/engine/test_stage_pool_collective_rpc.py

@@ -0,0 +1,71 @@
+"""Unit tests for StagePool.collective_rpc EngineCore control dispatch."""
+
+from __future__ import annotations
+
+import asyncio
+from types import SimpleNamespace
+from unittest.mock import AsyncMock
+
+import pytest
+
+from vllm_omni.engine.stage_pool import StagePool
+
+pytestmark = [pytest.mark.core_model]
+
+
+def _make_pool(*, stage_type: str = "llm", **client_methods: AsyncMock) -> tuple[StagePool, SimpleNamespace]:
+    client = SimpleNamespace(stage_type=stage_type, **client_methods)
+    if "collective_rpc_async" not in client_methods:
+        client.collective_rpc_async = AsyncMock(return_value={"via": "collective"})
+    pool = StagePool(0, [client])  # type: ignore[arg-type]
+    return pool, client
+
+
+@pytest.mark.cpu
+def test_collective_rpc_normalizes_none_args_on_control_helper():
+    async def run() -> None:
+        pause = AsyncMock(return_value="paused")
+        pool, client = _make_pool(pause_scheduler_async=pause)
+
+        result = await pool.collective_rpc(0, "pause_scheduler", args=None, kwargs={"mode": "abort"})
+
+        assert result == "paused"
+        pause.assert_awaited_once_with(mode="abort")
+        client.collective_rpc_async.assert_not_awaited()
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_collective_rpc_unrelated_async_helper_uses_collective_path():
+    async def run() -> None:
+        other = AsyncMock(return_value="should-not-run")
+        pool, client = _make_pool(reset_prefix_cache_async=other)
+
+        result = await pool.collective_rpc(0, "reset_prefix_cache", timeout=1.5, args=("x",), kwargs={"k": 1})
+
+        assert result == {"via": "collective"}
+        other.assert_not_awaited()
+        client.collective_rpc_async.assert_awaited_once_with(
+            method="reset_prefix_cache",
+            timeout=1.5,
+            args=("x",),
+            kwargs={"k": 1},
+        )
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_collective_rpc_control_helper_honors_timeout():
+    async def run() -> None:
+        async def slow_sleep(*_args, **_kwargs):
+            await asyncio.sleep(1.0)
+            return "slept"
+
+        pool, _client = _make_pool(sleep_async=slow_sleep)
+        result = await pool.collective_rpc(0, "sleep", timeout=0.01, args=(1,))
+        assert result != "slept"
+        assert result["supported"] is False
+
+    asyncio.run(run())

tests/entrypoints/test_async_omni.py

@@ -185,6 +185,58 @@ def test_abort_handles_internal_request_mapping(req_ids: list[str], cancel_prefi
     asyncio.run(run())
 
 
+@pytest.mark.cpu
+def test_abort_pops_request_states_only_after_ack():
+    async def run():
+        release = asyncio.Event()
+        seen_during_wait: list[int] = []
+
+        async def slow_abort_async(request_ids):
+            del request_ids
+            seen_during_wait.append(len(omni.request_states))
+            await release.wait()
+
+        omni = get_async_omni_instance(fake_abort_request=slow_abort_async)
+        omni.request_states["req-1-aaaa"] = SimpleNamespace(
+            request_id="req-1-aaaa",
+            external_request_id="req-1",
+            input_stream_task=None,
+        )
+
+        task = asyncio.create_task(omni.abort("req-1"))
+        await asyncio.sleep(0)
+        assert not task.done()
+        assert "req-1-aaaa" in omni.request_states
+        assert seen_during_wait == [1]
+
+        release.set()
+        await task
+        assert "req-1-aaaa" not in omni.request_states
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_abort_propagates_engine_errors_without_popping_state():
+    async def run():
+        async def failing_abort_async(request_ids):
+            del request_ids
+            raise RuntimeError("orchestrator abort failed")
+
+        omni = get_async_omni_instance(fake_abort_request=failing_abort_async)
+        omni.request_states["req-1-bbbb"] = SimpleNamespace(
+            request_id="req-1-bbbb",
+            external_request_id="req-1",
+            input_stream_task=None,
+        )
+
+        with pytest.raises(RuntimeError, match="orchestrator abort failed"):
+            await omni.abort("req-1")
+        assert "req-1-bbbb" in omni.request_states
+
+    asyncio.run(run())
+
+
 @pytest.mark.cpu
 def test_generate_accepts_request_after_repeated_cancellations():
     async def run_test():

tests/entrypoints/test_async_omni_pause_sleep_routing.py

@@ -0,0 +1,310 @@
+"""Unit tests for AR EngineCore vs diffusion worker pause/sleep routing."""
+
+from __future__ import annotations
+
+import asyncio
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, call
+
+import pytest
+
+from vllm_omni.diffusion.data import CuMemTag, OmniACK
+from vllm_omni.entrypoints.async_omni import AsyncOmni
+
+pytestmark = [pytest.mark.core_model]
+
+
+def _make_omni(*, stage_types: list[str]) -> AsyncOmni:
+    omni = object.__new__(AsyncOmni)
+    omni._pause_cond = asyncio.Condition()
+    omni._paused = False
+    omni._hold_admission_until_resume = False
+    omni._sleeping_tags = set()
+    omni._level2_sleeping = False
+    omni.event_resolver = SimpleNamespace(watch_task=lambda *a, **k: None, resolve=AsyncMock())
+    omni._final_output_handler = lambda: None
+
+    stage_clients = [SimpleNamespace(stage_type=stage_type) for stage_type in stage_types]
+    omni.engine = SimpleNamespace(
+        stage_clients=stage_clients,
+        stage_vllm_configs=[None] * len(stage_types),
+        collective_rpc_async=AsyncMock(return_value=[True]),
+    )
+    omni.collective_rpc = AsyncMock(return_value=[True])
+    return omni
+
+
+@pytest.mark.cpu
+def test_split_stage_ids_by_type():
+    omni = _make_omni(stage_types=["llm", "diffusion", "llm"])
+    ar_ids, diff_ids = omni._split_stage_ids_by_type()
+    assert ar_ids == [0, 2]
+    assert diff_ids == [1]
+
+
+@pytest.mark.cpu
+def test_split_stage_ids_by_type_rejects_out_of_range():
+    omni = _make_omni(stage_types=["llm", "diffusion"])
+    with pytest.raises(ValueError, match=r"Invalid stage_ids \[2\].*0\.\.1"):
+        omni._split_stage_ids_by_type([0, 2])
+    with pytest.raises(ValueError, match=r"Invalid stage_ids \[-1\]"):
+        omni._split_stage_ids_by_type([-1])
+
+
+@pytest.mark.cpu
+def test_pause_generation_routes_ar_via_collective_rpc():
+    async def run() -> None:
+        omni = _make_omni(stage_types=["llm", "diffusion"])
+        omni.reset_prefix_cache = AsyncMock(return_value=True)
+        omni.reset_mm_cache = AsyncMock()
+        omni.reset_encoder_cache = AsyncMock()
+
+        await omni.pause_generation(mode="abort", clear_cache=True)
+
+        assert omni._paused is True
+        assert omni._hold_admission_until_resume is True
+        omni.collective_rpc.assert_awaited_once_with(
+            method="pause_scheduler",
+            args=(),
+            kwargs={"mode": "abort", "clear_cache": True},
+            stage_ids=[0],
+        )
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_pause_generation_still_rpcs_when_already_paused():
+    async def run() -> None:
+        omni = _make_omni(stage_types=["llm", "diffusion"])
+        omni.reset_prefix_cache = AsyncMock(return_value=True)
+        omni.reset_mm_cache = AsyncMock()
+        omni.reset_encoder_cache = AsyncMock()
+        omni._paused = True
+
+        await omni.pause_generation(mode="abort", clear_cache=True, stage_ids=[0])
+
+        omni.collective_rpc.assert_awaited_once_with(
+            method="pause_scheduler",
+            args=(),
+            kwargs={"mode": "abort", "clear_cache": True},
+            stage_ids=[0],
+        )
+        omni.reset_prefix_cache.assert_awaited_once()
+        omni.reset_mm_cache.assert_awaited_once()
+        omni.reset_encoder_cache.assert_awaited_once()
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_resume_generation_resumes_ar_then_clears_frontend_pause():
+    async def run() -> None:
+        omni = _make_omni(stage_types=["llm", "diffusion"])
+        omni._paused = True
+
+        await omni.resume_generation()
+
+        omni.collective_rpc.assert_awaited_once_with(
+            method="resume_scheduler",
+            args=(),
+            kwargs=None,
+            stage_ids=[0],
+        )
+        assert omni._paused is False
+        assert omni._hold_admission_until_resume is False
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_sleep_routes_ar_via_collective_rpc_and_diffusion_to_worker_rpc():
+    async def run() -> None:
+        omni = _make_omni(stage_types=["llm", "diffusion"])
+        diffusion_ack = OmniACK(task_id="diff", status="SUCCESS", stage_id=1, rank=0)
+        omni._sleep_diffusion = AsyncMock(return_value=[diffusion_ack])
+
+        acks = await omni.sleep(stage_ids=[0, 1], level=1, mode="abort")
+
+        omni.collective_rpc.assert_awaited_once_with(
+            method="sleep",
+            args=(1, "abort"),
+            kwargs=None,
+            stage_ids=[0],
+        )
+        omni._sleep_diffusion.assert_awaited_once_with([1], 1)
+        assert {ack.stage_id for ack in acks} == {0, 1}
+        assert any(ack.metadata.get("path") == "engine_core" for ack in acks if ack.stage_id == 0)
+        assert CuMemTag.WEIGHTS.value in omni._sleeping_tags
+        assert CuMemTag.KV_CACHE.value in omni._sleeping_tags
+        # Sleep gates frontend admission for the trainer resume contract.
+        assert omni._paused is True
+        assert omni._hold_admission_until_resume is True
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_sleep_blocks_admission_before_engine_core_rpc():
+    """Problem 3: _paused must be set before awaiting sleep RPC."""
+
+    async def run() -> None:
+        omni = _make_omni(stage_types=["llm"])
+        paused_at_rpc: list[bool] = []
+
+        async def rpc_side_effect(**kwargs):
+            if kwargs.get("method") == "sleep":
+                paused_at_rpc.append(omni._paused)
+            return [True]
+
+        omni.collective_rpc = AsyncMock(side_effect=rpc_side_effect)
+
+        assert omni._paused is False
+        await omni.sleep(level=1, mode="abort")
+
+        assert paused_at_rpc == [True]
+        assert omni._paused is True
+        assert omni._hold_admission_until_resume is True
+        # Sleep must not also call pause_scheduler (EngineCore.sleep pauses).
+        assert all(c.kwargs.get("method") != "pause_scheduler" for c in omni.collective_rpc.await_args_list)
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_wake_up_routes_ar_via_collective_rpc_and_diffusion_to_worker_rpc():
+    async def run() -> None:
+        omni = _make_omni(stage_types=["llm", "diffusion"])
+        omni._paused = True
+        omni._hold_admission_until_resume = True
+        omni._sleeping_tags = {CuMemTag.WEIGHTS.value, CuMemTag.KV_CACHE.value}
+        diffusion_ack = OmniACK(task_id="diff", status="SUCCESS", stage_id=1, rank=0)
+        omni._wake_diffusion = AsyncMock(return_value=[diffusion_ack])
+
+        acks = await omni.wake_up(stage_ids=[0, 1])
+
+        omni.collective_rpc.assert_awaited_once()
+        wake_kwargs = omni.collective_rpc.await_args.kwargs
+        assert wake_kwargs["method"] == "wake_up"
+        assert wake_kwargs["stage_ids"] == [0]
+        assert set(wake_kwargs["kwargs"]["tags"]) == {
+            CuMemTag.WEIGHTS.value,
+            CuMemTag.KV_CACHE.value,
+        }
+        omni._wake_diffusion.assert_awaited_once()
+        assert {ack.stage_id for ack in acks} == {0, 1}
+        assert not omni._sleeping_tags
+        # Mixed/AR wake restores memory but does not resume frontend admission.
+        assert omni._paused is True
+        assert omni._hold_admission_until_resume is True
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_wake_up_does_not_resume_frontend_admission():
+    async def run() -> None:
+        omni = _make_omni(stage_types=["llm"])
+        await omni.sleep(level=1, mode="abort")
+        assert omni._paused is True
+        assert omni._hold_admission_until_resume is True
+
+        await omni.wake_up()
+
+        assert omni._paused is True
+        assert omni._hold_admission_until_resume is True
+        assert not omni._sleeping_tags
+        # Explicit resume is required after AR sleep/wake.
+        await omni.resume_generation()
+        assert omni._paused is False
+        assert omni._hold_admission_until_resume is False
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_sleep_level1_wake_without_tags_clears_all_sleeping_tags():
+    """sleep(level=1) tracks WEIGHTS+KV; untagged wake_up must clear both."""
+
+    async def run() -> None:
+        omni = _make_omni(stage_types=["llm"])
+
+        await omni.sleep(level=1, mode="abort")
+        assert CuMemTag.WEIGHTS.value in omni._sleeping_tags
+        assert CuMemTag.KV_CACHE.value in omni._sleeping_tags
+
+        await omni.wake_up()
+        assert not omni._sleeping_tags
+        assert omni._paused is True
+        assert omni._hold_admission_until_resume is True
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_sleep_diffusion_only_skips_engine_core_collective_rpc():
+    async def run() -> None:
+        omni = _make_omni(stage_types=["diffusion"])
+        omni._sleep_diffusion = AsyncMock(return_value=[OmniACK(task_id="d", status="SUCCESS", stage_id=0, rank=0)])
+
+        await omni.sleep(level=1)
+
+        # AR EngineCore sleep path must not run for diffusion-only.
+        assert call(method="sleep", args=(1, "abort"), kwargs=None, stage_ids=[0]) not in (
+            omni.collective_rpc.await_args_list
+        )
+        omni._sleep_diffusion.assert_awaited_once_with([0], 1)
+        assert omni._paused is True
+        assert omni._hold_admission_until_resume is False
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_wake_up_restores_admission_for_diffusion_only():
+    """Pure diffusion must keep sleep → wake → generate (no resume)."""
+
+    async def run() -> None:
+        omni = _make_omni(stage_types=["diffusion"])
+        omni._sleep_diffusion = AsyncMock(return_value=[OmniACK(task_id="d", status="SUCCESS", stage_id=0, rank=0)])
+        omni._wake_diffusion = AsyncMock(return_value=[OmniACK(task_id="w", status="SUCCESS", stage_id=0, rank=0)])
+
+        await omni.sleep(level=1)
+        assert omni._paused is True
+
+        async def wait_for_admission() -> None:
+            async with omni._pause_cond:
+                await omni._pause_cond.wait_for(lambda: not omni._paused)
+
+        waiter = asyncio.create_task(wait_for_admission())
+        await omni.wake_up()
+        await asyncio.wait_for(waiter, timeout=1.0)
+        assert omni._paused is False
+        assert omni._hold_admission_until_resume is False
+        assert not omni._sleeping_tags
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_pause_then_sleep_wake_keeps_admission_paused_for_diffusion():
+    """Explicit pause_generation still requires resume after diffusion wake."""
+
+    async def run() -> None:
+        omni = _make_omni(stage_types=["diffusion"])
+        omni.reset_prefix_cache = AsyncMock(return_value=True)
+        omni.reset_mm_cache = AsyncMock()
+        omni.reset_encoder_cache = AsyncMock()
+        omni._sleep_diffusion = AsyncMock(return_value=[OmniACK(task_id="d", status="SUCCESS", stage_id=0, rank=0)])
+        omni._wake_diffusion = AsyncMock(return_value=[OmniACK(task_id="w", status="SUCCESS", stage_id=0, rank=0)])
+
+        await omni.pause_generation()
+        await omni.sleep(level=1)
+        await omni.wake_up()
+
+        assert omni._paused is True
+        assert omni._hold_admission_until_resume is True
+        await omni.resume_generation()
+        assert omni._paused is False
+
+    asyncio.run(run())

tests/entrypoints/test_omni_sleep_mode.py

@@ -114,6 +114,10 @@ async def _ensure_awake(engine: AsyncOmni, stage_ids: list[int]) -> None:
         await engine.wake_up(stage_ids=stage_ids)
     except Exception as e:
         logger.warning("ensure_awake failed (stage_ids=%s): %s", stage_ids, e)
+    try:
+        await engine.resume_generation(stage_ids=stage_ids)
+    except Exception as e:
+        logger.warning("ensure_resume failed (stage_ids=%s): %s", stage_ids, e)
 
 
 def _build_llm_stages() -> tuple[list[dict], list[dict]]:
@@ -240,6 +244,9 @@ class TestOmniLlmSleepMode:
         try:
             await llm_engine.sleep(stage_ids=[0], level=1)
             await llm_engine.wake_up(stage_ids=[0], tags=["weights"])
+            # wake_up restores weights but keeps the AR admission hold.
+            # Resume so generate() can reach the sleeping-tags rejection.
+            await llm_engine.resume_generation(stage_ids=[0])
             with pytest.raises(RuntimeError, match="partially or fully asleep"):
                 async for _ in llm_engine.generate("test", sampling_params=SamplingParams(max_tokens=4)):
                     pass
@@ -490,6 +497,7 @@ async def test_multistage_llm_diffusion_sleep_wake(tp_size: int):
         assert len(acks) == 2
 
         await engine.wake_up(stage_ids=[0, 1])
+        await engine.resume_generation(stage_ids=[0, 1])
         async for _ in engine.generate("verify", sampling_params_list=[SamplingParams(), sp]):
             pass
     finally:

tests/model_tests/diffusion/test_alignment.py

@@ -26,6 +26,7 @@ EXCLUDED_MODELS = [
     "ZImagePipeline",
     "OvisImagePipeline",
     "WanPipeline",
+    "WanDMDPipeline",
     "WanVACEPipeline",
     "LTX2TwoStagePipeline",
     "LTX2DistilledOneStagePipeline",

vllm_omni/diffusion/attention/backends/abstract.py

@@ -36,6 +36,23 @@ class AttentionBackend(ABC):
         """
         return False
 
+    @classmethod
+    def supports_multi_doc_packed_varlen(cls) -> bool:
+        """Whether this backend keeps N-document packed boundaries isolated.
+
+        When True, the backend consumes ``AttentionMetadata.extra`` cu_seqlens
+        as a genuine block-diagonal attention plan (a dedicated varlen kernel,
+        not a padding-mask rebuild), so a caller may pack multiple real
+        requests into one forward without attention crossing document
+        boundaries. When False, callers packing more than one real document
+        must run one forward per document; otherwise a backend that only
+        supports a ``[real, pad]`` two-document contract, or that ignores
+        cu_seqlens outright, will silently attend across request boundaries.
+        Backends whose kernel selection is platform-dependent must consult
+        ``current_omni_platform``.
+        """
+        return False
+
     # ``OmniPlatformEnum`` values this backend runs on; None means unrestricted.
     # Platform resolution rejects an explicit selection outside this set, so a
     # hardware-specific backend fails before the model is built.

vllm_omni/diffusion/attention/backends/flash_attn.py

@@ -37,6 +37,20 @@ class FlashAttentionBackend(AttentionBackend):
         # constructing it there.
         return current_omni_platform.is_cuda() or current_omni_platform.is_npu()
 
+    @classmethod
+    def supports_multi_doc_packed_varlen(cls) -> bool:
+        # CUDA / ROCm / MUSA all route through ``forward_cuda``, which
+        # dispatches ``_forward_varlen_packed`` -> ``flash_attn_varlen_func``
+        # over the caller's cu_seqlens without a mask, so an arbitrary
+        # N-document packing keeps its boundaries. NPU's forward path
+        # (``_resolve_packed_seq_npu``) only accepts a ``[real, pad]``
+        # two-document layout (cu_seqlens shape <= 3) and otherwise falls
+        # back to a padding-mask rebuild that spans the whole packed row,
+        # so N>=2 real documents would silently attend across request
+        # boundaries. XPU never consumes cu_seqlens from ``extra``. Both
+        # therefore stay False here.
+        return current_omni_platform.is_cuda() or current_omni_platform.is_rocm() or current_omni_platform.is_musa()
+
     @classmethod
     def supports_attention_mask(cls) -> bool:
         return True

vllm_omni/diffusion/cache/teacache/extractors.py

@@ -1304,6 +1304,9 @@ def extract_minimax_h3_context(
     psp = _required_kwarg(kwargs, "packed_seq_params")
     cu_seqlens = module._psp_field(psp, "packed_seq_params", "cu_seqlens_q").to(torch.int32)
     max_seqlen = int(module._psp_field(psp, "packed_seq_params", "max_seqlen_q"))
+    # Mirror forward()'s host-side scalar so the refiner sees the number of
+    # packed requests without reading cu_seqlens off-device.
+    num_requests = int(module._psp_optional(psp, "num_requests", 1))
     refiner_psp = _required_kwarg(kwargs, "refiner_packed_seq_params")
     refiner_cu = module._psp_field(refiner_psp, "refiner_packed_seq_params", "cu_seqlens_q").to(torch.int32)
     refiner_max = int(module._psp_field(refiner_psp, "refiner_packed_seq_params", "max_seqlen_q"))
@@ -1330,6 +1333,7 @@ def extract_minimax_h3_context(
         text_pos=text_pos.to(device),
         refiner_cu_seqlens=refiner_cu.to(device),
         refiner_max_seqlen=refiner_max,
+        num_requests=num_requests,
         seq_len=seq_len,
         device=device,
         local_span=(0, seq_len),

vllm_omni/diffusion/forward_context.py

@@ -244,6 +244,17 @@ def set_forward_context_denoise_step_idx(step_idx: int | None) -> None:
         _forward_context.denoise_step_idx = step_idx
 
 
+def set_forward_context_denoise_timestep(timestep: float | None) -> None:
+    """Set the normalized (descending, 1 -> 0) denoise timestep.
+
+    Timestep-gated attention features read this; pipelines that drive their own
+    denoise loop can publish it directly instead of going through
+    :meth:`DenoiseProgressMixin.record_denoise_step`.
+    """
+    if _forward_context is not None:
+        _forward_context.denoise_timestep = None if timestep is None else float(timestep)
+
+
 class DenoiseProgressMixin:
     def record_denoise_step(
         self,

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -3,6 +3,7 @@
 import contextlib
 import dataclasses
 import glob
+import json
 import os
 import re
 import time
@@ -10,21 +11,20 @@ from collections.abc import Generator, Iterable, Sequence
 from pathlib import Path
 from typing import cast
 
+import huggingface_hub
 import torch
 from torch import nn
 from vllm.config.load import LoadConfig
 from vllm.logger import init_logger
 from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
 from vllm.model_executor.model_loader.weight_utils import (
-    download_safetensors_index_file_from_hf,
     download_weights_from_hf,
-    filter_duplicate_safetensors_files,
     filter_files_not_needed_for_inference,
     maybe_download_from_modelscope,
     multi_thread_safetensors_weights_iterator,
     safetensors_weights_iterator,
 )
-from vllm.transformers_utils.repo_utils import file_exists
+from vllm.transformers_utils.repo_utils import hf_api
 from vllm.utils.import_utils import resolve_obj_by_qualname
 from vllm.utils.torch_utils import set_default_torch_dtype
 
@@ -42,6 +42,7 @@ from vllm_omni.diffusion.model_loader.host_weight_plan import (
 from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import DiffusersAdapterPipeline
 from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
 from vllm_omni.diffusion.registry import initialize_model
+from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
 
 
 # download_gguf was removed from upstream vLLM (commit 6635279d8).
@@ -88,6 +89,23 @@ def _natural_sort_key(filepath: str) -> list:
 DIFFUSION_MODEL_WEIGHTS_INDEX = "diffusion_pytorch_model.safetensors.index.json"
 TRANSFORMER_WEIGHTS_INDEX = "model.safetensors.index.json"
 INDEX_FILES = [DIFFUSION_MODEL_WEIGHTS_INDEX, TRANSFORMER_WEIGHTS_INDEX]
+SHARDED_SAFETENSORS_PATTERN = re.compile(r"^(?P<family>.+)-\d+-of-(?P<count>\d+)\.safetensors$")
+
+
+def _validate_unindexed_safetensors_layout(weight_files: Sequence[str]) -> None:
+    """Reject ambiguous shard families when no index is available."""
+    shard_counts: dict[str, set[int]] = {}
+    for weight_file in weight_files:
+        match = SHARDED_SAFETENSORS_PATTERN.fullmatch(os.path.basename(weight_file))
+        if match is not None:
+            shard_counts.setdefault(match.group("family"), set()).add(int(match.group("count")))
+
+    conflicts = {family: sorted(counts) for family, counts in shard_counts.items() if len(counts) > 1}
+    if conflicts:
+        raise ValueError(
+            "Ambiguous unindexed safetensors checkpoint with conflicting shard totals: "
+            f"{conflicts}. Refusing to load potentially stale checkpoint shards."
+        )
 
 
 def _resolve_custom_pipeline_cls(custom_pipeline_name: str | type | None) -> type:
@@ -148,6 +166,60 @@ class DiffusersPipelineLoader:
         self.host_weight_plan = None
         return plan
 
+    @staticmethod
+    def _repo_relative_path(subfolder: str | None, filename: str) -> str:
+        if subfolder is None:
+            return filename
+        prefix = f"{subfolder.rstrip('/')}/"
+        return filename if filename.startswith(prefix) else f"{prefix}{filename.lstrip('/')}"
+
+    def _resolve_weight_index(
+        self,
+        model_name_or_path: Path | str,
+        subfolder: str | None,
+        revision: str | None,
+    ) -> list[str] | None:
+        """Resolve an index and return its authoritative shard manifest."""
+        is_local = os.path.isdir(model_name_or_path)
+        index_paths: list[tuple[str, Path]] = []
+        for index_file in INDEX_FILES:
+            repo_index_path = self._repo_relative_path(subfolder, index_file)
+            if is_local:
+                index_path = Path(model_name_or_path) / repo_index_path
+                if index_path.is_file():
+                    index_paths.append((index_file, index_path))
+                continue
+
+            try:
+                index_path = hf_api().hf_hub_download(
+                    repo_id=str(model_name_or_path),
+                    filename=repo_index_path,
+                    cache_dir=self.load_config.download_dir,
+                    revision=revision,
+                    local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,
+                )
+            except huggingface_hub.errors.EntryNotFoundError:
+                continue
+            index_paths.append((index_file, Path(index_path)))
+
+        if len(index_paths) > 1:
+            raise ValueError(
+                f"Multiple index files found in {model_name_or_path} with subfolder {subfolder}: "
+                f"{[index_file for index_file, _ in index_paths]}"
+            )
+        if not index_paths:
+            return None
+
+        index_file, index_path = index_paths[0]
+        with open(index_path) as f:
+            index = json.load(f)
+        weight_map = index.get("weight_map") if isinstance(index, dict) else None
+        if not isinstance(weight_map, dict) or not weight_map:
+            raise ValueError(f"Weight index {index_file} must contain a non-empty `weight_map`")
+        if not all(isinstance(filename, str) and filename for filename in weight_map.values()):
+            raise ValueError(f"Weight index {index_file} contains an invalid shard filename")
+        return sorted(set(weight_map.values()))
+
     def _prepare_weights(
         self,
         model_name_or_path: Path | str,
@@ -164,17 +236,11 @@ class DiffusersPipelineLoader:
         is_local = os.path.isdir(model_name_or_path)
         load_format = self.load_config.load_format
         use_safetensors = False
-        possible_index_files = [
-            f"{subfolder}/{index_file}" if subfolder is not None else index_file for index_file in INDEX_FILES
-        ]
-        available_index_file = [
-            f for f in possible_index_files if file_exists(model_name_or_path, f, revision=revision)
-        ]
-        if len(available_index_file) > 1:
-            raise ValueError(
-                f"Multiple index files found in {model_name_or_path} with subfolder {subfolder}: {available_index_file}"
-            )
-        index_file = available_index_file[0] if available_index_file else ""
+        indexed_weight_files = (
+            self._resolve_weight_index(model_name_or_path, subfolder, revision)
+            if allow_patterns_overrides is None
+            else None
+        )
 
         # only hf is supported currently
         if load_format == "auto":
@@ -192,7 +258,16 @@ class DiffusersPipelineLoader:
         if allow_patterns_overrides is not None:
             allow_patterns = allow_patterns_overrides
 
-        if not is_local:
+        if not is_local and indexed_weight_files is not None:
+            hf_folder = download_weights_from_hf_specific(
+                model_name_or_path=str(model_name_or_path),
+                cache_dir=self.load_config.download_dir,
+                allow_patterns=[self._repo_relative_path(subfolder, filename) for filename in indexed_weight_files],
+                revision=revision,
+                ignore_patterns=self.load_config.ignore_patterns,
+                require_all=True,
+            )
+        elif not is_local:
             hf_folder = download_weights_from_hf(
                 model_name_or_path,
                 self.load_config.download_dir,
@@ -207,31 +282,24 @@ class DiffusersPipelineLoader:
         if subfolder is not None:
             hf_folder = os.path.join(hf_folder, subfolder)
 
-        hf_weights_files: list[str] = []
-        for pattern in allow_patterns:
-            hf_weights_files += glob.glob(os.path.join(hf_folder, pattern))
-            if len(hf_weights_files) > 0:
-                # Decide by actual files rather than pattern name (patterns may include subfolders).
-                use_safetensors = any(f.endswith(".safetensors") for f in hf_weights_files)
-                break
-
-        if use_safetensors:
-            # For models like Mistral-7B-Instruct-v0.3
-            # there are both sharded safetensors files and a consolidated
-            # safetensors file. Using both breaks.
-            # Here, we download the `model.safetensors.index.json` and filter
-            # any files not found in the index.
-            if not is_local:
-                download_safetensors_index_file_from_hf(
-                    model_name_or_path,
-                    index_file,
-                    cache_dir=self.load_config.download_dir,
-                    subfolder=subfolder,
-                    revision=revision,
-                )
-            hf_weights_files = filter_duplicate_safetensors_files(hf_weights_files, hf_folder, index_file)
+        if indexed_weight_files is not None:
+            hf_weights_files = [os.path.join(hf_folder, filename) for filename in indexed_weight_files]
+            missing_files = [filename for filename in hf_weights_files if not os.path.isfile(filename)]
+            if missing_files:
+                raise FileNotFoundError(f"Weight files referenced in index but missing: {missing_files}")
+            use_safetensors = any(filename.endswith(".safetensors") for filename in hf_weights_files)
         else:
-            hf_weights_files = filter_files_not_needed_for_inference(hf_weights_files)
+            hf_weights_files = []
+            for pattern in allow_patterns:
+                hf_weights_files += glob.glob(os.path.join(hf_folder, pattern))
+                if hf_weights_files:
+                    # Decide by actual files rather than pattern name (patterns may include subfolders).
+                    use_safetensors = any(f.endswith(".safetensors") for f in hf_weights_files)
+                    break
+            if use_safetensors:
+                _validate_unindexed_safetensors_layout(hf_weights_files)
+            else:
+                hf_weights_files = filter_files_not_needed_for_inference(hf_weights_files)
 
         if len(hf_weights_files) == 0:
             raise RuntimeError(f"Cannot find any model weights with `{model_name_or_path}`")

vllm_omni/diffusion/models/minimax_h3/batched_packing.py

@@ -0,0 +1,155 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Batched multi-request packing for MiniMax H3 step-wise execution.
+
+Request mode forwards exactly one packed sequence per denoise step. Step mode
+(continuous batching) may hold several requests at once, so this module
+concatenates their packed layouts into a single sequence and rebuilds
+``cu_seqlens`` with one document per request plus that request's
+alignment-padding tail. Attention therefore never crosses a request boundary,
+while every request keeps its own RoPE coordinates, token tags, and timesteps.
+A batch of one is not rebuilt at all -- it delegates to the request-mode layout.
+
+Row order is the request order, so the concatenated video velocity returned by
+the DiT can be split back per request by row count alone -- which is exactly
+what the runner does with ``StepRequestState.latents``.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Any
+
+import torch
+
+from .denoise_loop import (
+    MINIMAX_H3_AUDIO_ROW_WIDTH,
+    MINIMAX_H3_VIDEO_ROW_WIDTH,
+    MiniMaxH3DenoiseBranch,
+)
+
+
+def minimax_h3_batched_forward_kwargs(
+    *,
+    branches: Sequence[MiniMaxH3DenoiseBranch],
+    video_rows: Sequence[torch.Tensor],
+    audio_rows: Sequence[torch.Tensor],
+    t_video: Sequence[float],
+    t_audio: Sequence[float],
+    imgvid_cond_timesteps: Sequence[float],
+    audio_ref_cond_timesteps: Sequence[float],
+) -> dict[str, Any]:
+    """Build one DiT forward kwargs dict covering every request in the batch.
+
+    A single request delegates to :meth:`MiniMaxH3DenoiseBranch.forward_kwargs`,
+    so step mode and request mode are the same call rather than two layouts that
+    happen to agree -- which also keeps ``video_token_layout`` (the sparse-
+    attention hint) on the single-request step path.
+    """
+    if len(branches) == 1:
+        return branches[0].forward_kwargs(
+            video_rows=video_rows[0],
+            audio_rows=audio_rows[0],
+            t_video=t_video[0],
+            t_audio=t_audio[0],
+            imgvid_cond_timestep=imgvid_cond_timesteps[0],
+            audio_ref_cond_timestep=audio_ref_cond_timesteps[0],
+        )
+
+    device = branches[0].device
+    total_seq = sum(branch.seq_len for branch in branches)
+
+    img_pos_parts: list[torch.Tensor] = []
+    audio_pos_parts: list[torch.Tensor] = []
+    text_pos_parts: list[torch.Tensor] = []
+    update_parts: list[torch.Tensor] = []
+    tag_parts: list[torch.Tensor] = []
+    position_id_parts: list[torch.Tensor] = []
+    text_embed_parts: list[torch.Tensor] = []
+    cu_bounds: list[int] = [0]
+    refiner_bounds: list[int] = [0]
+    max_seqlen = 0
+    refiner_max_seqlen = 0
+
+    # Non-media rows (text and alignment padding) inherit their request's video
+    # timestep, matching the single-request packed-sequence semantics.
+    timesteps = torch.empty(total_seq, dtype=torch.float32, device=device)
+
+    seq_offset = 0
+    text_offset = 0
+    for index, branch in enumerate(branches):
+        img_pos = branch.img_pos_dev + seq_offset
+        audio_pos = branch.audio_pos_dev + seq_offset
+        img_pos_parts.append(img_pos)
+        audio_pos_parts.append(audio_pos)
+        text_pos_parts.append(branch.text_pos_dev + seq_offset)
+        update_parts.append(branch.update_mask_dev)
+        tag_parts.append(branch.token_tags_dev)
+        position_id_parts.append(branch.img_position_ids_dev)
+        text_embed_parts.append(branch.text_embeddings_dev)
+
+        timesteps[seq_offset : seq_offset + branch.seq_len] = float(t_video[index])
+        timesteps[img_pos[branch.update_mask_dev]] = float(t_video[index])
+        timesteps[img_pos[~branch.update_mask_dev]] = float(imgvid_cond_timesteps[index])
+        timesteps[audio_pos[branch.audio_update_mask_dev]] = float(t_audio[index])
+        timesteps[audio_pos[~branch.audio_update_mask_dev]] = float(audio_ref_cond_timesteps[index])
+
+        # One document for the request's real rows and one for its padding tail,
+        # matching the single-request layout. The tail bound is emitted even when
+        # the 64-row alignment left it empty, so that a request always
+        # contributes exactly two bounds regardless of its length.
+        cu_bounds.append(seq_offset + branch.used_len)
+        cu_bounds.append(seq_offset + branch.seq_len)
+        max_seqlen = max(max_seqlen, branch.used_len, branch.seq_len - branch.used_len)
+
+        text_offset += branch.text_len
+        refiner_bounds.append(text_offset)
+        refiner_max_seqlen = max(refiner_max_seqlen, branch.text_len)
+        seq_offset += branch.seq_len
+
+    img_pos_all = torch.cat(img_pos_parts)
+    audio_pos_all = torch.cat(audio_pos_parts)
+
+    x = torch.zeros(1, total_seq, MINIMAX_H3_VIDEO_ROW_WIDTH, dtype=torch.float32, device=device)
+    x[0].index_copy_(0, img_pos_all, torch.cat([rows.to(device=device, dtype=torch.float32) for rows in video_rows]))
+    audio_x = torch.zeros(1, total_seq, MINIMAX_H3_AUDIO_ROW_WIDTH, dtype=torch.float32, device=device)
+    audio_x[0].index_copy_(
+        0,
+        audio_pos_all,
+        torch.cat([rows.to(device=device, dtype=torch.float32) for rows in audio_rows]),
+    )
+
+    unique_timesteps, inverse_indices = torch.unique(timesteps, sorted=True, return_inverse=True)
+    return {
+        "x": x,
+        "audio_x": audio_x,
+        "img_position_ids": torch.cat(position_id_parts)[None],
+        "update_mask": torch.cat(update_parts),
+        "token_tags": torch.cat(tag_parts),
+        "skip_mask_out_condition": False,
+        "prompt_embeds": torch.cat(text_embed_parts),
+        "img_pos_info": {"position_ids": img_pos_all},
+        "audio_pos_info": {"position_ids": audio_pos_all},
+        "text_pos_info": {"position_ids": torch.cat(text_pos_parts)},
+        "img_pos_for_infer_output_info": {"position_ids": img_pos_all},
+        "packed_seq_params": {
+            "cu_seqlens_q": torch.tensor(cu_bounds, dtype=torch.int32, device=device),
+            "max_seqlen_q": max_seqlen,
+            # Attention needs the request count on the host to know its valid
+            # rows are block-diagonal rather than a prefix. It describes the
+            # whole forward, so the DiT reads it once here and applies it to the
+            # token refiner too -- the refiner's text-only bounds cannot be told
+            # apart from a single padded request on their own.
+            "num_requests": len(branches),
+        },
+        "refiner_packed_seq_params": {
+            "cu_seqlens_q": torch.tensor(refiner_bounds, dtype=torch.int32, device=device),
+            "max_seqlen_q": refiner_max_seqlen,
+        },
+        "unique_timesteps": unique_timesteps,
+        "inverse_indices": inverse_indices,
+    }
+
+
+__all__ = [
+    "minimax_h3_batched_forward_kwargs",
+]

vllm_omni/diffusion/models/minimax_h3/denoise_loop.py

@@ -15,7 +15,10 @@ from typing import Any
 import torch
 
 from vllm_omni.diffusion.attention.backends.abstract import VideoTokenLayout
-from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx
+from vllm_omni.diffusion.forward_context import (
+    set_forward_context_denoise_step_idx,
+    set_forward_context_denoise_timestep,
+)
 
 from .scheduling_minimax_h3_euler_ancestral import (
     minimax_h3_euler_eta0_step,
@@ -31,6 +34,17 @@ MINIMAX_H3_VIDEO_ROW_WIDTH = 96
 MINIMAX_H3_AUDIO_ROW_WIDTH = 32
 
 
+def minimax_h3_publish_denoise_progress(step: int | None, sigma_video: float | None) -> None:
+    """Publish denoise progress for step-gated attention features.
+
+    Both execution modes must publish the same pair: the step index drives the
+    dense warmup of RAINFUSION_ATTN, and the normalized descending timestep
+    drives the TRTLLM_ATTN skip gate, which stays dense while it is unset.
+    """
+    set_forward_context_denoise_step_idx(step)
+    set_forward_context_denoise_timestep(sigma_video)
+
+
 class MiniMaxH3DenoiseBranch:
     """Static per-branch state: packed layout + fixed forward kwargs.
 
@@ -64,26 +78,39 @@ class MiniMaxH3DenoiseBranch:
         if int(token_tags.view(-1).shape[0]) != seq_len:
             raise ValueError(f"token_tags length {int(token_tags.view(-1).shape[0])} != seq_len {seq_len}")
         cu = packed["cu_seqlens"].to(torch.int32)
+        self.device = device
+        # ``used_len`` is the real (non-padding) document length; step-mode
+        # batching concatenates layouts and needs both bounds as Python ints so
+        # it can rebuild ``cu_seqlens`` without a device sync per step.
+        self.used_len = int(cu[1])
+        self.text_len = text_len
         self.img_pos_dev = self.img_pos.to(device)
         self.audio_pos_dev = self.audio_pos.to(device)
         self.update_mask_dev = self.update_mask.to(device)
         self.audio_update_mask_dev = self.audio_update_mask.to(device)
         self.x_base = torch.zeros(1, seq_len, MINIMAX_H3_VIDEO_ROW_WIDTH, dtype=torch.float32, device=device)
         self.audio_x_base = torch.zeros(1, seq_len, MINIMAX_H3_AUDIO_ROW_WIDTH, dtype=torch.float32, device=device)
-        text_pos_dev = packed["text_pos"].view(-1).to(torch.long).to(device)
+        self.text_pos_dev = packed["text_pos"].view(-1).to(torch.long).to(device)
+        self.token_tags_dev = token_tags.view(-1).to(torch.long).to(device)
+        self.img_position_ids_dev = packed["img_position_ids"].to(device)
+        self.text_embeddings_dev = text_embeddings.to(device)
         self.static_kwargs: dict[str, Any] = {
-            "img_position_ids": packed["img_position_ids"][None].to(device),
+            "img_position_ids": self.img_position_ids_dev[None],
             "update_mask": self.update_mask_dev,
-            "token_tags": token_tags.view(-1).to(torch.long).to(device),
+            "token_tags": self.token_tags_dev,
             "skip_mask_out_condition": False,
-            "prompt_embeds": text_embeddings.to(device),
+            "prompt_embeds": self.text_embeddings_dev,
             "img_pos_info": {"position_ids": self.img_pos_dev},
             "audio_pos_info": {"position_ids": self.audio_pos_dev},
-            "text_pos_info": {"position_ids": text_pos_dev},
+            "text_pos_info": {"position_ids": self.text_pos_dev},
             "img_pos_for_infer_output_info": {"position_ids": self.img_pos_dev},
             "packed_seq_params": {
                 "cu_seqlens_q": cu.to(device),
-                "max_seqlen_q": int(cu[1]),
+                "max_seqlen_q": self.used_len,
+                # One request: valid rows are a prefix, so attention may use a
+                # KV prefix length or a 1-D pad mask. See batched_packing for the
+                # co-batched layout, where neither can describe the valid rows.
+                "num_requests": 1,
             },
             "refiner_packed_seq_params": {
                 "cu_seqlens_q": torch.tensor([0, text_len, text_len], dtype=torch.int32, device=device),
@@ -136,34 +163,21 @@ class MiniMaxH3DenoiseBranch:
         }
 
 
-def minimax_h3_denoise_loop(
+def minimax_h3_prepare_denoise_rows(
     *,
-    model: Any,
     positive: MiniMaxH3DenoiseBranch,
     initial_video_rows: torch.Tensor,
     initial_audio_rows: torch.Tensor,
     keyframe_cond_rows: torch.Tensor | None,
-    audio_ref_rows: torch.Tensor | None = None,
-    sigmas_video: list[float],
-    sigmas_audio: list[float],
+    audio_ref_rows: torch.Tensor | None,
     device: torch.device,
-    imgvid_cond_noise_aug_for_inference: float = MINIMAX_H3_IMGVID_COND_TIMESTEP,
-    audio_cond_noise_aug_for_inference: float = MINIMAX_H3_AUDIO_REF_COND_TIMESTEP,
-    on_step_end: Callable[[int, torch.Tensor, torch.Tensor], None] | None = None,
-    on_step_start: Callable[[int, float, float], None] | None = None,
-    step_profiler: Callable[[int], AbstractContextManager] | None = None,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """Run the full denoise loop; returns final (video_rows, audio_rows).
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
+    """Validate the initial rows against the layout and pin the condition rows.
 
-    ``initial_video_rows`` covers all image rows of the positive layout. For a
-    conditional task, pass ``keyframe_cond_rows`` and/or ``audio_ref_rows`` to
-    pin those rows across every step. The model's raw positive velocity is the
-    update signal; MiniMax H3 only supports cfg-distilled checkpoints.
+    Returns ``(video_rows, audio_rows, cond_anchor, audio_anchor)`` with the
+    anchors already written into their rows, which is the state both the
+    request-mode loop and step-mode ``prepare_encode()`` start from.
     """
-    if len(sigmas_video) != len(sigmas_audio):
-        raise ValueError("video/audio sigma schedules must have equal length")
-    if len(sigmas_video) < 2:
-        raise ValueError("sigma schedules need at least 2 entries")
     n_cond = int((~positive.update_mask).sum())
     if keyframe_cond_rows is None:
         if n_cond != 0:
@@ -197,19 +211,58 @@ def minimax_h3_denoise_loop(
         video_rows[~update] = cond_anchor
     if audio_anchor is not None:
         audio_rows[~audio_update] = audio_anchor
+    return video_rows, audio_rows, cond_anchor, audio_anchor
+
+
+def minimax_h3_denoise_loop(
+    *,
+    model: Any,
+    positive: MiniMaxH3DenoiseBranch,
+    initial_video_rows: torch.Tensor,
+    initial_audio_rows: torch.Tensor,
+    keyframe_cond_rows: torch.Tensor | None,
+    audio_ref_rows: torch.Tensor | None = None,
+    sigmas_video: list[float],
+    sigmas_audio: list[float],
+    device: torch.device,
+    imgvid_cond_noise_aug_for_inference: float = MINIMAX_H3_IMGVID_COND_TIMESTEP,
+    audio_cond_noise_aug_for_inference: float = MINIMAX_H3_AUDIO_REF_COND_TIMESTEP,
+    on_step: Callable[[int, torch.Tensor, torch.Tensor], None] | None = None,
+    step_profiler: Callable[[int], AbstractContextManager] | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Run the full denoise loop; returns final (video_rows, audio_rows).
+
+    ``initial_video_rows`` covers all image rows of the positive layout. For a
+    conditional task, pass ``keyframe_cond_rows`` and/or ``audio_ref_rows`` to
+    pin those rows across every step. The model's raw positive velocity is the
+    update signal; MiniMax H3 only supports cfg-distilled checkpoints.
+    """
+    if len(sigmas_video) != len(sigmas_audio):
+        raise ValueError("video/audio sigma schedules must have equal length")
+    if len(sigmas_video) < 2:
+        raise ValueError("sigma schedules need at least 2 entries")
+    video_rows, audio_rows, cond_anchor, audio_anchor = minimax_h3_prepare_denoise_rows(
+        positive=positive,
+        initial_video_rows=initial_video_rows,
+        initial_audio_rows=initial_audio_rows,
+        keyframe_cond_rows=keyframe_cond_rows,
+        audio_ref_rows=audio_ref_rows,
+        device=device,
+    )
+    update = positive.update_mask_dev
+    audio_update = positive.audio_update_mask_dev
 
     num_steps = len(sigmas_video) - 1
     for step in range(num_steps):
         step_cm = step_profiler(step) if step_profiler is not None else nullcontext()
         with step_cm:
-            # Publish the step index so step-gated attention features (e.g. the
-            # dense warmup of RAINFUSION_ATTN) can see where we are.
-            set_forward_context_denoise_step_idx(step)
             s_v, s_v_next = sigmas_video[step], sigmas_video[step + 1]
             s_a, s_a_next = sigmas_audio[step], sigmas_audio[step + 1]
-            if on_step_start is not None:
-                # Attention gates use the scheduler-style descending timestep.
-                on_step_start(step, s_v, s_a)
+            # Publish where we are so step-gated attention features (the dense
+            # warmup of RAINFUSION_ATTN, the timestep gate of TRTLLM_ATTN) can
+            # see it. Gates use the scheduler-style descending timestep, which
+            # for a rectified-flow schedule is the video sigma.
+            minimax_h3_publish_denoise_progress(step, s_v)
             t_v, t_a = 1.0 - s_v, 1.0 - s_a
             imgvid_cond_t = max(t_v, float(imgvid_cond_noise_aug_for_inference))
             audio_ref_cond_t = max(t_a, float(audio_cond_noise_aug_for_inference))
@@ -250,10 +303,10 @@ def minimax_h3_denoise_loop(
             audio_rows[audio_update] = new_audio
             if audio_anchor is not None:
                 audio_rows[~audio_update] = audio_anchor  # per-step audio ref reset
-            if on_step_end is not None:
-                on_step_end(step, video_rows, audio_rows)
+            if on_step is not None:
+                on_step(step, video_rows, audio_rows)
 
-    set_forward_context_denoise_step_idx(None)
+    minimax_h3_publish_denoise_progress(None, None)
     return video_rows, audio_rows
 
 
@@ -264,4 +317,6 @@ __all__ = [
     "MINIMAX_H3_VIDEO_ROW_WIDTH",
     "MiniMaxH3DenoiseBranch",
     "minimax_h3_denoise_loop",
+    "minimax_h3_prepare_denoise_rows",
+    "minimax_h3_publish_denoise_progress",
 ]

vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py

@@ -56,6 +56,26 @@ if TYPE_CHECKING:
 logger = init_logger(__name__)
 
 
+# Packed multi-request forwards require the attention backend to actually
+# consume cu_seqlens as a block-diagonal plan (not a padding-mask rebuild that
+# spans the full packed row). The pipeline gates on this capability before
+# packing, and ``_run_packed_attention`` re-checks it per forward; a name-only
+# gate would let FLASH_ATTN's NPU/XPU code paths through even though those
+# variants would silently attend across request boundaries.
+def _attention_isolates_packed_requests(attention_layer: Any) -> bool:
+    """True if this attention layer keeps N-document packed boundaries.
+
+    Requires a backend advertising ``supports_multi_doc_packed_varlen`` *and*
+    that the layer is not running under ring sequence parallelism (the ring
+    kernel dispatches through its own attention that ignores the packed
+    cu_seqlens regardless of the configured backend).
+    """
+    backend = getattr(attention_layer, "attn_backend", None)
+    if backend is None or not backend.supports_multi_doc_packed_varlen():
+        return False
+    return not getattr(attention_layer, "use_ring", False)
+
+
 @dataclass
 class MiniMaxH3DiTArchConfig:
     num_layers: int = 50
@@ -433,6 +453,7 @@ class MiniMaxH3Attention(nn.Module):
         cu_seqlens: torch.Tensor,
         max_seqlen: int,
         packed_total: int,
+        num_requests: int = 1,
         video_layout: VideoTokenLayout | None = None,
     ) -> torch.Tensor:
         """Run packed attention as a small eager island.
@@ -442,27 +463,47 @@ class MiniMaxH3Attention(nn.Module):
         narrow lets regional compile fuse projections, norms, RoPE, and the
         surrounding DiT block without repeated graph breaks.
         """
-        # max_seqlen is already the first (real) packed document length. Do
-        # not read the CUDA cu_seqlens scalars here: this function runs once
-        # per layer and .item() would serialize every attention launch.
+        # max_seqlen is already the longest packed document length. Do not read
+        # the CUDA cu_seqlens scalars here: this function runs once per layer
+        # and .item() would serialize every attention launch. ``num_requests``
+        # is carried as a Python int for the same reason.
         if not 0 < max_seqlen <= packed_total:
             raise ValueError(
                 f"max_seqlen must be within the packed sequence, got {max_seqlen} for length {packed_total}"
             )
-        used = min(max_seqlen, packed_total)
         attn_mask = None
-        # Ring attention can dispatch to a different implementation from the
-        # configured backend, so the no-mask fast paths are local-only.
-        # supports_prefix_kv_slicing: backend slices K/V itself (cuDNN).
-        # supports_packed_mask_free: backend consumes the packed metadata
-        # without ever reading attn_mask (CUDA packed varlen, NPU
-        # npu_attn_varlen opt-in with its own fallback rebuild).
-        no_mask = not getattr(self.attention, "use_ring", False) and (
-            self.attention.attn_backend.supports_prefix_kv_slicing
-            or self.attention.attn_backend.supports_packed_mask_free()
-        )
-        if used < packed_total and not no_mask:
-            attn_mask = torch.arange(packed_total, device=q.device)[None] < used
+        if num_requests > 1:
+            # A step-mode batch packs one document per request, so its valid
+            # rows are block-diagonal rather than a prefix: neither a KV prefix
+            # length nor a 1-D key mask can describe them. Such a layout is
+            # only correct on a backend that actually attends by cu_seqlens as
+            # a block-diagonal plan. Check the capability (not the backend
+            # name): FLASH_ATTN's NPU/XPU variants would otherwise silently
+            # fall back to a padding-mask rebuild that spans the whole packed
+            # row and attend across request boundaries.
+            if not _attention_isolates_packed_requests(self.attention):
+                backend_name = self.attention.attn_backend.get_name()
+                raise ValueError(
+                    f"MiniMax H3 packed a {num_requests}-request batch, but the resolved "
+                    f"attention ({backend_name}, use_ring={getattr(self.attention, 'use_ring', False)}) "
+                    "does not isolate multi-document packed cu_seqlens. Run one request "
+                    "per forward on this backend."
+                )
+            used = packed_total
+        else:
+            used = min(max_seqlen, packed_total)
+            # Ring attention can dispatch to a different implementation from the
+            # configured backend, so the no-mask fast paths are local-only.
+            # supports_prefix_kv_slicing: backend slices K/V itself (cuDNN).
+            # supports_packed_mask_free: backend consumes the packed metadata
+            # without ever reading attn_mask (CUDA packed varlen, NPU
+            # npu_attn_varlen opt-in with its own fallback rebuild).
+            no_mask = not getattr(self.attention, "use_ring", False) and (
+                self.attention.attn_backend.supports_prefix_kv_slicing
+                or self.attention.attn_backend.supports_packed_mask_free()
+            )
+            if used < packed_total and not no_mask:
+                attn_mask = torch.arange(packed_total, device=q.device)[None] < used
         metadata = AttentionMetadata(
             attn_mask=attn_mask,
             extra={
@@ -498,6 +539,7 @@ class MiniMaxH3Attention(nn.Module):
         cu_seqlens: torch.Tensor,
         max_seqlen: int,
         packed_total: int | None = None,
+        num_requests: int = 1,
         sp_seq_lens: list[int] | None = None,
         video_layout: VideoTokenLayout | None = None,
     ) -> torch.Tensor:
@@ -533,9 +575,9 @@ class MiniMaxH3Attention(nn.Module):
                 self.q_norm.variance_epsilon,
             )
 
-        # The packed layout uses a second document for alignment padding.
-        # Local/Ulysses backends unpad it, while Ring keeps aligned rows for
-        # fixed-size P2P buffers.
+        # Each request contributes a document for its rows plus one for its
+        # alignment padding. Local/Ulysses backends unpad it, while Ring keeps
+        # aligned rows for fixed-size P2P buffers.
         out = self._run_packed_attention(
             q,
             k,
@@ -546,6 +588,7 @@ class MiniMaxH3Attention(nn.Module):
             # backend receives the global sequence after all-to-all, so carry
             # its Python length explicitly instead of inferring it from q.
             packed_total=packed_total if packed_total is not None else q.shape[0],
+            num_requests=num_requests,
             video_layout=video_layout,
         )
         out = out.reshape(total, self.num_heads * self.head_dim)
@@ -673,12 +716,14 @@ class MiniMaxH3TokenRefinerBlock(nn.Module):
         *,
         cu_seqlens: torch.Tensor,
         max_seqlen: int,
+        num_requests: int = 1,
     ) -> torch.Tensor:
         x = x + self.attn(
             self.norm1(x),
             rope_table=None,
             cu_seqlens=cu_seqlens,
             max_seqlen=max_seqlen,
+            num_requests=num_requests,
         )
         x = x + self.mlp(self.norm2(x))
         return x
@@ -711,9 +756,10 @@ class MiniMaxH3TokenRefiner(nn.Module):
         *,
         cu_seqlens: torch.Tensor,
         max_seqlen: int,
+        num_requests: int = 1,
     ) -> torch.Tensor:
         for block in self.blocks:
-            x = block(x, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen)
+            x = block(x, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen, num_requests=num_requests)
         return self.final_norm(x)
 
 
@@ -759,6 +805,7 @@ class MiniMaxH3DiTBlock(nn.Module):
         cu_seqlens: torch.Tensor,
         max_seqlen: int,
         packed_total: int,
+        num_requests: int = 1,
         sp_seq_lens: list[int] | None = None,
         video_layout: VideoTokenLayout | None = None,
     ) -> torch.Tensor:
@@ -793,6 +840,7 @@ class MiniMaxH3DiTBlock(nn.Module):
             cu_seqlens=cu_seqlens,
             max_seqlen=max_seqlen,
             packed_total=packed_total,
+            num_requests=num_requests,
             sp_seq_lens=sp_seq_lens,
             video_layout=video_layout,
         )
@@ -1145,6 +1193,11 @@ class MiniMaxH3DiTModel(nn.Module):
             raise ValueError(f"{key}.{field} is required")
         return value
 
+    @staticmethod
+    def _psp_optional(psp: Any, field: str, default: Any) -> Any:
+        value = psp.get(field) if isinstance(psp, dict) else getattr(psp, field, None)
+        return default if value is None else value
+
     def _embed(
         self,
         *,
@@ -1160,10 +1213,16 @@ class MiniMaxH3DiTModel(nn.Module):
         seq_len: int,
         device: torch.device,
         local_span: tuple[int, int],
+        num_requests: int = 1,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """Build this rank's packed multimodal embedding rows.
 
         Returns (decoder_input [S_local, H] bf16, t_emb [M, t_dim] fp32).
+
+        ``num_requests`` defaults to a single packed request so callers that
+        pre-date the continuous-batching change (e.g. TeaCache's extractor
+        contract) do not silently miss the kwarg. ``forward()`` reads the real
+        value from ``packed_seq_params["num_requests"]``.
         """
         local_start, local_len = local_span
         local_end = local_start + local_len
@@ -1200,6 +1259,7 @@ class MiniMaxH3DiTModel(nn.Module):
             text_embed,
             cu_seqlens=refiner_cu_seqlens,
             max_seqlen=refiner_max_seqlen,
+            num_requests=num_requests,
         )
         if text_local_indices is not None:
             text_embed = text_embed.index_select(0, text_local_indices)
@@ -1269,6 +1329,10 @@ class MiniMaxH3DiTModel(nn.Module):
         psp = _required_kwarg(kwargs, "packed_seq_params")
         cu_seqlens = self._psp_field(psp, "packed_seq_params", "cu_seqlens_q").to(torch.int32)
         max_seqlen = int(self._psp_field(psp, "packed_seq_params", "max_seqlen_q"))
+        # How many requests share this packed sequence. Carried as a host int so
+        # attention never reads cu_seqlens scalars off the device; a producer
+        # that omits it is packing a single request.
+        num_requests = int(self._psp_optional(psp, "num_requests", 1))
         refiner_psp = _required_kwarg(kwargs, "refiner_packed_seq_params")
         refiner_cu = self._psp_field(refiner_psp, "refiner_packed_seq_params", "cu_seqlens_q").to(torch.int32)
         refiner_max = int(self._psp_field(refiner_psp, "refiner_packed_seq_params", "max_seqlen_q"))
@@ -1305,6 +1369,7 @@ class MiniMaxH3DiTModel(nn.Module):
             text_pos=text_pos.to(device),
             refiner_cu_seqlens=refiner_cu.to(device),
             refiner_max_seqlen=refiner_max,
+            num_requests=num_requests,
             seq_len=seq_len,
             device=device,
             local_span=local_span,
@@ -1339,6 +1404,7 @@ class MiniMaxH3DiTModel(nn.Module):
                 cu_seqlens=cu_seqlens,
                 max_seqlen=max_seqlen,
                 packed_total=seq_len,
+                num_requests=num_requests,
                 video_layout=video_layout,
             )
         if local_len == seq_len:

vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py

@@ -57,7 +57,7 @@ from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import (
 )
 from vllm_omni.diffusion.sched.sigma_schedule import DMD2SigmaSchedule
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
-from vllm_omni.errors import OmniClientError
+from vllm_omni.errors import OmniClientError, client_error_from_metadata
 from vllm_omni.model_executor.model_loader.weight_utils import (
     download_weights_from_hf_specific,
 )
@@ -69,14 +69,24 @@ from vllm_omni.quantization.component_config import (
     resolve_encoder_quant_config as _resolve_encoder_quant_config,
 )
 
+from .batched_packing import minimax_h3_batched_forward_kwargs
 from .condition_noise import (
     minimax_h3_audio_cond_noise_aug_rows,
     minimax_h3_imgvid_cond_noise_aug_rows,
 )
-from .denoise_loop import MiniMaxH3DenoiseBranch, minimax_h3_denoise_loop
+from .denoise_loop import (
+    MiniMaxH3DenoiseBranch,
+    minimax_h3_denoise_loop,
+    minimax_h3_prepare_denoise_rows,
+    minimax_h3_publish_denoise_progress,
+)
 from .encoder import MiniMaxH3Qwen3VLEncoder
 from .lora import load_minimax_h3_turbo_lora
-from .minimax_h3_transformer import MiniMaxH3DiTModel
+from .minimax_h3_transformer import (
+    MiniMaxH3Attention,
+    MiniMaxH3DiTModel,
+    _attention_isolates_packed_requests,
+)
 from .packed_sequence import (
     minimax_h3_packed_sequence,
     minimax_h3_packed_sequence_ref2va_blocks,
@@ -103,6 +113,10 @@ from .reference_video import (
     validate_reference_audio_files,
     validate_reference_audio_waveforms,
 )
+from .scheduling_minimax_h3_euler_ancestral import (
+    minimax_h3_euler_eta0_step,
+    minimax_h3_rf_v_to_x0,
+)
 from .time_request import (
     MINIMAX_H3_SHAPE_PLANNER,
     minimax_h3_align_frame_count,
@@ -110,6 +124,10 @@ from .time_request import (
 )
 from .vae import MiniMaxH3AudioVAE, MiniMaxH3VideoVAE
 
+if TYPE_CHECKING:
+    from vllm_omni.diffusion.worker.input_batch import InputBatch
+    from vllm_omni.diffusion.worker.utils import StepRequestState
+
 logger = init_logger(__name__)
 
 if TYPE_CHECKING:
@@ -201,6 +219,68 @@ def _resolve_minimax_h3_model_root(
     )
 
 
+# Keys of ``_prepare_request_inputs`` that feed ``diffuse`` / ``_build_denoise_inputs``.
+_MINIMAX_H3_DENOISE_INPUT_KEYS = (
+    "task",
+    "text_embeddings",
+    "text_tags",
+    "seed",
+    "latent_t",
+    "latent_h",
+    "latent_w",
+    "audio_t",
+    "num_frames",
+    "num_steps",
+    "video_shift",
+    "audio_shift",
+    "base_schedule",
+    "visual_condition",
+    "visual_condition_shape",
+    "audio_condition",
+    "ref_audio_t",
+    "ref_blocks",
+    "visual_condition_shapes",
+    "audio_condition_lengths",
+    "keyframe_frame_indices",
+)
+
+# ``StepRequestState.extra`` keys owned by the step-execution path.
+_STEP_BRANCH = "minimax_h3_branch"
+_STEP_AUDIO_ROWS = "minimax_h3_audio_rows"
+_STEP_AUDIO_NOISE_PRED = "minimax_h3_audio_noise_pred"
+_STEP_SIGMAS_VIDEO = "minimax_h3_sigmas_video"
+_STEP_SIGMAS_AUDIO = "minimax_h3_sigmas_audio"
+_STEP_COND_ANCHOR = "minimax_h3_cond_anchor"
+_STEP_AUDIO_ANCHOR = "minimax_h3_audio_anchor"
+_STEP_SHAPE = "minimax_h3_shape"
+_STEP_TRANSFORMER = "minimax_h3_transformer"
+
+
+def _minimax_h3_step_schedule(state: StepRequestState) -> dict[str, float]:
+    """Return the sigma/timestep values this request needs for its current step.
+
+    Mirrors the per-iteration arithmetic of ``minimax_h3_denoise_loop`` so step
+    mode and request mode advance identically.
+    """
+    step = int(state.step_index)
+    sigmas_video = state.extra[_STEP_SIGMAS_VIDEO]
+    sigmas_audio = state.extra[_STEP_SIGMAS_AUDIO]
+    sigma_video = float(sigmas_video[step])
+    sigma_audio = float(sigmas_audio[step])
+    t_video = 1.0 - sigma_video
+    t_audio = 1.0 - sigma_audio
+    return {
+        "sigma_video": sigma_video,
+        "sigma_video_next": float(sigmas_video[step + 1]),
+        "sigma_audio": sigma_audio,
+        "sigma_audio_next": float(sigmas_audio[step + 1]),
+        "t_video": t_video,
+        "t_audio": t_audio,
+        "imgvid_cond_timestep": max(t_video, MINIMAX_H3_IMGVID_COND_TIMESTEP),
+        "audio_ref_cond_timestep": max(t_audio, MINIMAX_H3_AUDIO_REF_COND_TIMESTEP),
+    }
+
+
 def _read_base_schedule(release: Mapping[str, Any]) -> DMD2SigmaSchedule | None:
     """Read a partition's distilled schedule. An absent key means legacy uniform."""
     return DMD2SigmaSchedule.from_metadata(release)
@@ -444,6 +524,58 @@ def _dit_rank_world() -> tuple[Any, int, int]:
     return group, dist.get_rank(group), dist.get_world_size(group)
 
 
+def _broadcast_rank0_exception(exc: Exception | None) -> None:
+    """Synchronize a rank-0-only exception across every DiT rank.
+
+    H3 reference-video preparation runs only on rank 0; the other DiT ranks
+    return ``None`` without touching disk. When rank 0 raises inside that
+    path it exits :meth:`prepare_encode` before reaching the downstream
+    ``dist.broadcast`` calls, and non-zero ranks then hang on those
+    collectives forever. Every rank calls this helper right after the
+    rank-0-only work, before any subsequent collective, so all ranks either
+    raise the same error together or all continue.
+    """
+    group, rank, world_size = _dit_rank_world()
+    if world_size == 1:
+        if exc is not None:
+            raise exc
+        return
+    if rank == 0:
+        if exc is None:
+            payload: list[Any] = [None]
+        else:
+            payload = [
+                {
+                    "type": type(exc).__name__,
+                    "message": str(exc),
+                    "status_code": getattr(exc, "status_code", None),
+                    "error_type": getattr(exc, "error_type", None),
+                }
+            ]
+    else:
+        payload = [None]
+    dist.broadcast_object_list(payload, src=0, group=group)
+    info = payload[0]
+    if info is None:
+        return
+    if rank == 0:
+        assert exc is not None
+        raise exc
+    # Rebuild a matching client-facing error on non-zero ranks so the runner's
+    # per-request try/except records the same 4xx status as rank 0. The exact
+    # subclass need not survive the wire; the message and status suffice.
+    status_code = info.get("status_code")
+    error_type = info.get("error_type")
+    message = f"[rank 0] {info['type']}: {info['message']}"
+    if status_code is not None:
+        raise client_error_from_metadata(
+            message,
+            status_code=int(status_code),
+            error_type=error_type,
+        )
+    raise RuntimeError(message)
+
+
 def _broadcast_tensor(
     tensor: torch.Tensor | None,
     *,
@@ -548,6 +680,8 @@ class MiniMaxH3Pipeline(
 ):
     """CFG-distilled joint video/audio generation for MiniMax H3."""
 
+    supports_step_execution: ClassVar[bool] = True
+
     _dit_modules: ClassVar[list[str]] = ["transformer", "transformers_ref"]
     _encoder_modules: ClassVar[list[str]] = ["text_encoder"]
     _vae_modules: ClassVar[list[str]] = ["video_vae", "audio_vae"]
@@ -564,6 +698,9 @@ class MiniMaxH3Pipeline(
         "_encode_reference_audio_conditions",
         "diffuse",
         "decode",
+        "prepare_encode",
+        "denoise_step",
+        "post_decode",
     ]
     dummy_run_num_frames: ClassVar[int] = 0
     # Only distilled releases pin a schedule, so the default keeps the legacy
@@ -1538,7 +1675,7 @@ class MiniMaxH3Pipeline(
             if controller is not None and enabled:
                 controller.offload_resident_layers()
 
-    def diffuse(
+    def _build_denoise_inputs(
         self,
         *,
         task: str,
@@ -1562,7 +1699,12 @@ class MiniMaxH3Pipeline(
         visual_condition_shapes: list[tuple[int, int, int]] | None = None,
         audio_condition_lengths: list[int] | None = None,
         keyframe_frame_indices: list[int] | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
+    ) -> dict[str, Any]:
+        """Build the packed layout, initial rows, anchors, and sigma schedules.
+
+        Shared by request-mode :meth:`diffuse` and step-mode
+        :meth:`prepare_encode` so both paths start from identical state.
+        """
         initial_video, initial_audio = self._initial_noise(
             seed=seed,
             latent_t=latent_t,
@@ -1662,32 +1804,34 @@ class MiniMaxH3Pipeline(
             shift_scale=audio_shift,
             base_schedule=base_schedule,
         )
-        transformer = self._transformer_for_task(task)
-        # The static DLO plan keeps leading blocks resident only for the
-        # primary ``transformer``. In combined mode ``transformers_ref`` is
-        # fully streamed, so a Ref2VA request must not stage the inactive
-        # FL2VA transformer resident blocks.
-        with self._resident_dit_layers_on_device(enabled=transformer is self.transformer):
-            with self.progress_bar(total=len(video_sigmas) - 1) as progress:
-                video_rows, audio_rows = minimax_h3_denoise_loop(
-                    model=transformer,
-                    positive=branch,
-                    initial_video_rows=initial_video,
-                    initial_audio_rows=initial_audio,
-                    keyframe_cond_rows=visual_anchor,
-                    audio_ref_rows=audio_anchor,
-                    sigmas_video=video_sigmas,
-                    sigmas_audio=audio_sigmas,
-                    device=self.device,
-                    imgvid_cond_noise_aug_for_inference=(MINIMAX_H3_IMGVID_COND_TIMESTEP),
-                    audio_cond_noise_aug_for_inference=(MINIMAX_H3_AUDIO_REF_COND_TIMESTEP),
-                    on_step_start=lambda step, video_sigma, audio_sigma: self.record_denoise_step(
-                        step,
-                        normalized_timestep=video_sigma,
-                    ),
-                    on_step_end=lambda step, video, audio: progress.update(),
-                )
+        return {
+            "branch": branch,
+            # The request-mode loop moves these onto the device itself; step mode
+            # keeps them resident across steps, so normalize once for both.
+            "video_rows": initial_video.to(device=self.device, dtype=torch.float32),
+            "audio_rows": initial_audio.to(device=self.device, dtype=torch.float32),
+            "cond_anchor": (
+                None if visual_anchor is None else visual_anchor.to(device=self.device, dtype=torch.float32)
+            ),
+            "audio_anchor": (
+                None if audio_anchor is None else audio_anchor.to(device=self.device, dtype=torch.float32)
+            ),
+            "sigmas_video": video_sigmas,
+            "sigmas_audio": audio_sigmas,
+        }
 
+    def _unpack_denoised_rows(
+        self,
+        branch: MiniMaxH3DenoiseBranch,
+        video_rows: torch.Tensor,
+        audio_rows: torch.Tensor,
+        *,
+        latent_t: int,
+        latent_h: int,
+        latent_w: int,
+        audio_t: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Select the target rows and unpack them back into VAE latents."""
         target_video = video_rows[branch.update_mask_dev]
         video_latent = minimax_h3_unpatchify_video_tokens(
             target_video,
@@ -1707,6 +1851,83 @@ class MiniMaxH3Pipeline(
         )
         return video_latent, audio_latent
 
+    def diffuse(
+        self,
+        *,
+        task: str,
+        text_embeddings: torch.Tensor,
+        text_tags: torch.Tensor,
+        seed: int,
+        latent_t: int,
+        latent_h: int,
+        latent_w: int,
+        audio_t: int,
+        num_frames: int,
+        num_steps: int,
+        video_shift: float,
+        audio_shift: float,
+        base_schedule: Sequence[float] | None,
+        visual_condition: torch.Tensor | None,
+        visual_condition_shape: tuple[int, int, int] | None,
+        audio_condition: torch.Tensor | None,
+        ref_audio_t: int | None,
+        ref_blocks: list[dict[str, Any]] | None = None,
+        visual_condition_shapes: list[tuple[int, int, int]] | None = None,
+        audio_condition_lengths: list[int] | None = None,
+        keyframe_frame_indices: list[int] | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        inputs = self._build_denoise_inputs(
+            task=task,
+            text_embeddings=text_embeddings,
+            text_tags=text_tags,
+            seed=seed,
+            latent_t=latent_t,
+            latent_h=latent_h,
+            latent_w=latent_w,
+            audio_t=audio_t,
+            num_frames=num_frames,
+            num_steps=num_steps,
+            video_shift=video_shift,
+            audio_shift=audio_shift,
+            base_schedule=base_schedule,
+            visual_condition=visual_condition,
+            visual_condition_shape=visual_condition_shape,
+            audio_condition=audio_condition,
+            ref_audio_t=ref_audio_t,
+            ref_blocks=ref_blocks,
+            visual_condition_shapes=visual_condition_shapes,
+            audio_condition_lengths=audio_condition_lengths,
+            keyframe_frame_indices=keyframe_frame_indices,
+        )
+        branch = inputs["branch"]
+        transformer = self._transformer_for_task(task)
+        with self._resident_dit_layers_on_device(enabled=transformer is self.transformer):
+            with self.progress_bar(total=len(inputs["sigmas_video"]) - 1) as progress:
+                video_rows, audio_rows = minimax_h3_denoise_loop(
+                    model=transformer,
+                    positive=branch,
+                    initial_video_rows=inputs["video_rows"],
+                    initial_audio_rows=inputs["audio_rows"],
+                    keyframe_cond_rows=inputs["cond_anchor"],
+                    audio_ref_rows=inputs["audio_anchor"],
+                    sigmas_video=inputs["sigmas_video"],
+                    sigmas_audio=inputs["sigmas_audio"],
+                    device=self.device,
+                    imgvid_cond_noise_aug_for_inference=(MINIMAX_H3_IMGVID_COND_TIMESTEP),
+                    audio_cond_noise_aug_for_inference=(MINIMAX_H3_AUDIO_REF_COND_TIMESTEP),
+                    on_step=lambda step, video, audio: progress.update(),
+                )
+
+        return self._unpack_denoised_rows(
+            branch,
+            video_rows,
+            audio_rows,
+            latent_t=latent_t,
+            latent_h=latent_h,
+            latent_w=latent_w,
+            audio_t=audio_t,
+        )
+
     def decode(
         self,
         video_latent: torch.Tensor,
@@ -1727,11 +1948,9 @@ class MiniMaxH3Pipeline(
             audio = self.audio_vae.decode_latent(audio_latent)
         return video, audio
 
-    @torch.no_grad()
-    def forward(self, request: DiffusionRequestBatch) -> DiffusionOutput:
-        if len(request.prompts) != 1:
-            raise OmniClientError("MiniMax H3 supports one request at a time")
-        raw_prompt = request.prompts[0]
+    @staticmethod
+    def _extract_prompt(raw_prompt: Any) -> tuple[str, dict[str, Any]]:
+        """Split a request prompt into its text and multimodal parts."""
         if isinstance(raw_prompt, str):
             prompt = raw_prompt
             multi_modal_data: dict[str, Any] = {}
@@ -1740,8 +1959,21 @@ class MiniMaxH3Pipeline(
             multi_modal_data = raw_prompt.get("multi_modal_data") or {}
         if not prompt:
             raise OmniClientError("MiniMax H3 requires a non-empty prompt")
+        return prompt, multi_modal_data
 
-        sampling = request.sampling_params
+    def _prepare_request_inputs(
+        self,
+        *,
+        prompt: str,
+        multi_modal_data: dict[str, Any],
+        sampling: Any,
+    ) -> dict[str, Any]:
+        """Resolve the task and output shape, then run every request-level encode.
+
+        Shared by request-mode :meth:`forward` and step-mode
+        :meth:`prepare_encode`; the returned mapping feeds :meth:`diffuse` and
+        :meth:`_build_denoise_inputs` unchanged.
+        """
         quality = sampling.quality
         logger.debug("MiniMax H3 request quality=%s", quality)
         extra = sampling.extra_args or {}
@@ -1818,12 +2050,24 @@ class MiniMaxH3Pipeline(
             video_count = 0
             if raw_videos is not None:
                 video_count = len(raw_videos) if isinstance(raw_videos, (list, tuple)) else 1
-                prepared_videos = self._prepare_reference_videos(
-                    raw_videos,
-                    target_frame_count=num_frames,
-                    workdir=workdir,
-                    start_time_seconds=extra.get("start_time_seconds"),
-                )
+                # File-based reference-video prep runs only on rank 0; other
+                # ranks return None without touching disk. If rank 0 raises
+                # (e.g. invalid file, unsupported codec) it must not exit
+                # ``prepare_encode`` before the downstream broadcasts below --
+                # non-zero ranks would then deadlock on them forever. Capture
+                # the exception here and let every rank agree on the outcome
+                # before starting any subsequent collective.
+                prep_error: Exception | None = None
+                try:
+                    prepared_videos = self._prepare_reference_videos(
+                        raw_videos,
+                        target_frame_count=num_frames,
+                        workdir=workdir,
+                        start_time_seconds=extra.get("start_time_seconds"),
+                    )
+                except Exception as exc:
+                    prep_error = exc
+                _broadcast_rank0_exception(prep_error)
                 has_audio_tensor = torch.zeros(
                     video_count,
                     dtype=torch.long,
@@ -1960,33 +2204,60 @@ class MiniMaxH3Pipeline(
         )
         self._cache_dit_runtime.prepare(quality_plan.cache_dit)
         num_outputs = _resolve_minimax_h3_num_outputs(sampling.num_outputs_per_prompt)
+        return {
+            "task": task,
+            "height": height,
+            "width": width,
+            "num_frames": num_frames,
+            "latent_t": latent_t,
+            "latent_h": height // 16,
+            "latent_w": width // 16,
+            "audio_t": audio_t,
+            "text_embeddings": text_embeddings,
+            "text_tags": text_tags,
+            "visual_condition": visual_condition,
+            "visual_condition_shape": visual_shape,
+            "audio_condition": audio_condition,
+            "ref_audio_t": ref_audio_t,
+            "ref_blocks": ref_blocks,
+            "visual_condition_shapes": visual_shapes,
+            "audio_condition_lengths": audio_lengths,
+            "keyframe_frame_indices": keyframe_frame_indices,
+            "seed": seed,
+            "num_steps": num_steps,
+            "video_shift": video_shift,
+            "audio_shift": audio_shift,
+            "base_schedule": base_schedule,
+            "num_outputs": num_outputs,
+        }
+
+    @staticmethod
+    def _denoise_kwargs(context: dict[str, Any]) -> dict[str, Any]:
+        """Select the denoise-input arguments from a prepared request context."""
+        return {key: context[key] for key in _MINIMAX_H3_DENOISE_INPUT_KEYS}
+
+    @torch.no_grad()
+    def forward(self, request: DiffusionRequestBatch) -> DiffusionOutput:
+        if len(request.prompts) != 1:
+            raise OmniClientError("MiniMax H3 supports one request at a time")
+        prompt, multi_modal_data = self._extract_prompt(request.prompts[0])
+        context = self._prepare_request_inputs(
+            prompt=prompt,
+            multi_modal_data=multi_modal_data,
+            sampling=request.sampling_params,
+        )
+        denoise_kwargs = self._denoise_kwargs(context)
+        num_outputs = context["num_outputs"]
         videos = []
         audios = []
-        for output_seed in _minimax_h3_output_seeds(seed, num_outputs):
-            video_latent, audio_latent = self.diffuse(
-                task=task,
-                text_embeddings=text_embeddings,
-                text_tags=text_tags,
-                seed=output_seed,
-                latent_t=latent_t,
-                latent_h=height // 16,
-                latent_w=width // 16,
-                audio_t=audio_t,
-                num_frames=num_frames,
-                num_steps=num_steps,
-                video_shift=video_shift,
-                audio_shift=audio_shift,
-                base_schedule=base_schedule,
-                visual_condition=visual_condition,
-                visual_condition_shape=visual_shape,
-                audio_condition=audio_condition,
-                ref_audio_t=ref_audio_t,
-                ref_blocks=ref_blocks,
-                visual_condition_shapes=visual_shapes,
-                audio_condition_lengths=audio_lengths,
-                keyframe_frame_indices=keyframe_frame_indices,
+        for output_seed in _minimax_h3_output_seeds(context["seed"], num_outputs):
+            video_latent, audio_latent = self.diffuse(**{**denoise_kwargs, "seed": output_seed})
+            video, audio = self.decode(
+                video_latent,
+                audio_latent,
+                height=context["height"],
+                width=context["width"],
             )
-            video, audio = self.decode(video_latent, audio_latent, height=height, width=width)
             videos.append(video)
             audios.append(audio)
         video = torch.cat(videos, dim=0)
@@ -1997,6 +2268,302 @@ class MiniMaxH3Pipeline(
             stage_durations=(self.stage_durations if hasattr(self, "_stage_durations") else {}),
         )
 
+    # ------------------------------------------------------------------
+    # Step-wise execution (continuous batching)
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def _packed_batch_supported(transformer: MiniMaxH3DiTModel) -> bool:
+        """Whether every attention in this DiT honors multi-document cu_seqlens.
+
+        A packed batch is only isolated if *all* of them do: the token refiner
+        runs under its own attention role and can resolve to a different backend
+        from the DiT blocks. Ring sequence parallelism dispatches through
+        ``RingParallelAttention``, whose kernels ignore the packed
+        ``cu_seqlens`` metadata regardless of the configured backend; packing
+        multiple requests under ring would let attention cross document
+        boundaries, so any layer running ring disqualifies the batch.
+
+        The gate probes a per-backend capability rather than a fixed backend
+        name: FLASH_ATTN, for example, only isolates arbitrary N-document
+        packed cu_seqlens on CUDA/ROCm/MUSA. Its NPU path only accepts a
+        ``[real, pad]`` two-document layout and its XPU path ignores
+        cu_seqlens outright — either would silently attend across request
+        boundaries.
+        """
+        attentions = [module for module in transformer.modules() if isinstance(module, MiniMaxH3Attention)]
+        if not attentions:
+            return False
+        return all(_attention_isolates_packed_requests(module.attention) for module in attentions)
+
+    def prepare_encode(self, state: StepRequestState, **kwargs: Any) -> StepRequestState:
+        """Run every request-level stage once and seed the per-request step state."""
+        del kwargs
+        # Two request-mode features have no place in the shared step contract:
+        # a request state carries exactly one latent tensor, and distributed
+        # layerwise offload streams the DiT around one whole denoise loop rather
+        # than around a single scheduler-driven step.
+        num_outputs = _resolve_minimax_h3_num_outputs(state.sampling.num_outputs_per_prompt)
+        if num_outputs != 1:
+            raise OmniClientError(
+                f"MiniMax H3 step execution produces one output per request, got num_outputs_per_prompt={num_outputs}"
+            )
+        if getattr(self, "_dlo_residency_controller", None) is not None:
+            raise ValueError(
+                "MiniMax H3 step execution is not compatible with distributed layerwise offload; "
+                "the resident-layer window spans a whole denoise loop, so per-step streaming would "
+                "reload the DiT every step. Drop --step-execution or --enable-distributed-layerwise-offload."
+            )
+        # Request-scoped Cache-DiT (quality=high) mutates hook state on the
+        # shared transformer rather than on ``StepRequestState``. In step mode
+        # two requests can interleave denoise steps, or be co-batched into a
+        # single forward, and the second one would then re-enter the DiT with
+        # cache buffers shaped for the first. Reject the profile here rather
+        # than let it corrupt outputs at runtime; startup-configured Cache-DiT
+        # is already blocked in ``DiffusionModelRunner.execute_stepwise``.
+        if getattr(state.sampling, "quality", None) == "high":
+            raise OmniClientError(
+                "MiniMax H3 step execution does not support the high-quality Cache-DiT profile "
+                "(quality=high); its hooks live on the shared transformer, so interleaved or "
+                "co-batched requests would reuse incompatible cache state. Drop --step-execution "
+                "or omit quality=high."
+            )
+        prompt, multi_modal_data = self._extract_prompt(state.prompt)
+        context = self._prepare_request_inputs(
+            prompt=prompt,
+            multi_modal_data=multi_modal_data,
+            sampling=state.sampling,
+        )
+        inputs = self._build_denoise_inputs(**self._denoise_kwargs(context))
+
+        sigmas_video = inputs["sigmas_video"]
+        sigmas_audio = inputs["sigmas_audio"]
+        if len(sigmas_video) < 2:
+            raise OmniClientError(
+                f"MiniMax H3 step execution needs at least one denoise step, got num_inference_steps="
+                f"{len(sigmas_video) - 1}"
+            )
+
+        branch = inputs["branch"]
+        video_rows, audio_rows, cond_anchor, audio_anchor = minimax_h3_prepare_denoise_rows(
+            positive=branch,
+            initial_video_rows=inputs["video_rows"],
+            initial_audio_rows=inputs["audio_rows"],
+            keyframe_cond_rows=inputs["cond_anchor"],
+            audio_ref_rows=inputs["audio_anchor"],
+            device=self.device,
+        )
+
+        # The denoise loop consumes sigma pairs, so the schedule carries one more
+        # point than there are steps. ``timesteps`` holds the video branch because
+        # the shared contract gives a request exactly one timestep sequence; the
+        # audio schedule rides along in ``extra``.
+        state.timesteps = torch.tensor(
+            [1.0 - sigma for sigma in sigmas_video[:-1]],
+            dtype=torch.float32,
+            device=self.device,
+        )
+        state.step_index = 0
+        # Video rows are the batched tensor the runner slices per request; audio
+        # rows have a different width, so they stay request-private.
+        state.latents = video_rows
+        state.do_true_cfg = False  # H3 checkpoints are CFG-distilled.
+        state.extra.update(
+            {
+                _STEP_BRANCH: branch,
+                _STEP_TRANSFORMER: self._transformer_for_task(context["task"]),
+                _STEP_AUDIO_ROWS: audio_rows,
+                _STEP_COND_ANCHOR: cond_anchor,
+                _STEP_AUDIO_ANCHOR: audio_anchor,
+                _STEP_SIGMAS_VIDEO: sigmas_video,
+                _STEP_SIGMAS_AUDIO: sigmas_audio,
+                _STEP_SHAPE: {
+                    "height": context["height"],
+                    "width": context["width"],
+                    "latent_t": context["latent_t"],
+                    "latent_h": context["latent_h"],
+                    "latent_w": context["latent_w"],
+                    "audio_t": context["audio_t"],
+                },
+            }
+        )
+        return state
+
+    def denoise_step(
+        self,
+        input_batch: InputBatch,
+        *,
+        states: Sequence[StepRequestState] | None = None,
+        **kwargs: Any,
+    ) -> torch.Tensor | None:
+        """Run one denoise forward covering every request in the batch.
+
+        Requests are concatenated into a single packed sequence that keeps one
+        attention document each, so the whole batch costs one DiT forward.
+        Backends that ignore ``cu_seqlens`` cannot express that isolation, so
+        they fall back to one forward per request.
+        """
+        del kwargs
+        batch_states = list(states if states is not None else input_batch.states)
+
+        branches = [state.extra[_STEP_BRANCH] for state in batch_states]
+        video_rows = [state.latents for state in batch_states]
+        audio_rows = [state.extra[_STEP_AUDIO_ROWS] for state in batch_states]
+        schedules = [_minimax_h3_step_schedule(state) for state in batch_states]
+        transformers = [state.extra[_STEP_TRANSFORMER] for state in batch_states]
+        mixed_transformers = len({id(transformer) for transformer in transformers}) > 1
+
+        # Both execution modes must publish denoise progress for step-gated
+        # attention features. Requests can differ in both step index and sigma
+        # schedule, so a batch that is not at one single point has nothing to
+        # publish and those gates stay dense -- which is their safe default.
+        progress = {(state.step_index, schedule["sigma_video"]) for state, schedule in zip(batch_states, schedules)}
+        minimax_h3_publish_denoise_progress(*(progress.pop() if len(progress) == 1 else (None, None)))
+
+        if len(batch_states) > 1 and (mixed_transformers or not self._packed_batch_supported(transformers[0])):
+            if mixed_transformers:
+                logger.warning_once(
+                    "MiniMax H3 step batch contains requests for different task-specific DiTs; "
+                    "running %d requests one forward at a time.",
+                    len(batch_states),
+                )
+            elif any(
+                getattr(getattr(module, "attention", None), "use_ring", False)
+                for module in transformers[0].modules()
+                if isinstance(module, MiniMaxH3Attention)
+            ):
+                logger.warning_once(
+                    "MiniMax H3 step batching is disabled when ring attention is active: "
+                    "the ring kernels ignore packed cu_seqlens and would attend across request "
+                    "boundaries. Running %d requests one forward at a time.",
+                    len(batch_states),
+                )
+            else:
+                logger.warning_once(
+                    "MiniMax H3 step batching needs every attention on a backend that isolates "
+                    "packed multi-document cu_seqlens (see AttentionBackend."
+                    "supports_multi_doc_packed_varlen); running %d requests one forward at a time.",
+                    len(batch_states),
+                )
+            video_parts: list[torch.Tensor] = []
+            audio_parts: list[torch.Tensor] = []
+            for index, branch in enumerate(branches):
+                forward_kwargs = branch.forward_kwargs(
+                    video_rows=video_rows[index],
+                    audio_rows=audio_rows[index],
+                    t_video=schedules[index]["t_video"],
+                    t_audio=schedules[index]["t_audio"],
+                    imgvid_cond_timestep=schedules[index]["imgvid_cond_timestep"],
+                    audio_ref_cond_timestep=schedules[index]["audio_ref_cond_timestep"],
+                )
+                request_video, request_audio = transformers[index](**forward_kwargs)
+                video_parts.append(request_video)
+                audio_parts.append(request_audio)
+            video_velocity = torch.cat(video_parts)
+            audio_velocity = torch.cat(audio_parts)
+        else:
+            forward_kwargs = minimax_h3_batched_forward_kwargs(
+                branches=branches,
+                video_rows=video_rows,
+                audio_rows=audio_rows,
+                t_video=[schedule["t_video"] for schedule in schedules],
+                t_audio=[schedule["t_audio"] for schedule in schedules],
+                imgvid_cond_timesteps=[schedule["imgvid_cond_timestep"] for schedule in schedules],
+                audio_ref_cond_timesteps=[schedule["audio_ref_cond_timestep"] for schedule in schedules],
+            )
+            logger.debug(
+                "MiniMax H3 denoise step: %d request(s) packed into %d rows",
+                len(batch_states),
+                int(forward_kwargs["x"].shape[1]),
+            )
+            video_velocity, audio_velocity = transformers[0](**forward_kwargs)
+
+        # The shared contract carries one velocity tensor per step, and audio rows
+        # are a different width than video rows, so hand the audio branch to
+        # step_scheduler() through request-private state.
+        audio_parts_by_request = torch.split(audio_velocity, [int(branch.audio_pos.shape[0]) for branch in branches])
+        for state, request_audio in zip(batch_states, audio_parts_by_request, strict=True):
+            state.extra[_STEP_AUDIO_NOISE_PRED] = request_audio
+        return video_velocity
+
+    def step_scheduler(self, state: StepRequestState, noise_pred: torch.Tensor, **kwargs: Any) -> None:
+        """Apply one Euler-eta0 update to this request's video and audio rows."""
+        del kwargs
+        # denoise_step() stages the audio half of this step's velocity; popping
+        # it keeps a second step_scheduler() call from reusing a stale one.
+        audio_noise_pred = state.extra.pop(_STEP_AUDIO_NOISE_PRED)
+
+        branch = state.extra[_STEP_BRANCH]
+        schedule = _minimax_h3_step_schedule(state)
+        update = branch.update_mask_dev
+        audio_update = branch.audio_update_mask_dev
+        video_rows = state.latents
+        audio_rows = state.extra[_STEP_AUDIO_ROWS]
+        cond_anchor = state.extra[_STEP_COND_ANCHOR]
+        audio_anchor = state.extra[_STEP_AUDIO_ANCHOR]
+        device = video_rows.device
+
+        x0_video = minimax_h3_rf_v_to_x0(
+            video_rows[update],
+            noise_pred.float()[update],
+            torch.tensor(schedule["t_video"], dtype=torch.float32, device=device),
+        )
+        new_video = minimax_h3_euler_eta0_step(
+            video_rows[update],
+            x0_video,
+            sigma_curr=schedule["sigma_video"],
+            sigma_next=schedule["sigma_video_next"],
+        )
+        video_rows = video_rows.clone()
+        video_rows[update] = new_video
+        if cond_anchor is not None:
+            video_rows[~update] = cond_anchor  # per-step imgvid cond reset
+
+        x0_audio = minimax_h3_rf_v_to_x0(
+            audio_rows[audio_update],
+            audio_noise_pred.float()[audio_update],
+            torch.tensor(schedule["t_audio"], dtype=torch.float32, device=device),
+        )
+        new_audio = minimax_h3_euler_eta0_step(
+            audio_rows[audio_update],
+            x0_audio,
+            sigma_curr=schedule["sigma_audio"],
+            sigma_next=schedule["sigma_audio_next"],
+        )
+        audio_rows = audio_rows.clone()
+        audio_rows[audio_update] = new_audio
+        if audio_anchor is not None:
+            audio_rows[~audio_update] = audio_anchor  # per-step audio ref reset
+
+        state.latents = video_rows
+        state.extra[_STEP_AUDIO_ROWS] = audio_rows
+        state.step_index += 1
+
+    def post_decode(self, state: StepRequestState, **kwargs: Any) -> DiffusionOutput:
+        """Unpack the denoised rows and run the joint video/audio VAE decode."""
+        del kwargs
+        shape = state.extra[_STEP_SHAPE]
+        video_latent, audio_latent = self._unpack_denoised_rows(
+            state.extra[_STEP_BRANCH],
+            state.latents,
+            state.extra[_STEP_AUDIO_ROWS],
+            latent_t=shape["latent_t"],
+            latent_h=shape["latent_h"],
+            latent_w=shape["latent_w"],
+            audio_t=shape["audio_t"],
+        )
+        video, audio = self.decode(
+            video_latent,
+            audio_latent,
+            height=shape["height"],
+            width=shape["width"],
+        )
+        return DiffusionOutput(
+            output=(video, audio),
+            post_process_func=get_minimax_h3_post_process_func(self.od_config),
+            stage_durations=(self.stage_durations if hasattr(self, "_stage_durations") else {}),
+        )
+
 
 __all__ = [
     "MiniMaxH3Pipeline",

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -81,6 +81,31 @@ if TYPE_CHECKING:
 logger = init_logger(__name__)
 
 
+def _dit_any_rank_failed(local_failed: bool) -> bool:
+    """All-reduce a per-request failure flag across the DiT process group.
+
+    Every DiT rank must reach this point in lockstep; the caller wraps
+    ``pipeline.prepare_encode`` so both the success and the exception paths
+    call the helper. In single-rank execution this collapses to the local
+    flag with no collectives issued.
+    """
+    if not torch.distributed.is_initialized():
+        return local_failed
+    try:
+        from vllm_omni.diffusion.distributed.parallel_state import get_dit_group
+
+        group = get_dit_group()
+    except (AssertionError, ImportError):
+        group = None
+    if group is None:
+        return local_failed
+    signal = torch.tensor(1 if local_failed else 0, dtype=torch.int32)
+    if current_omni_platform.is_available():
+        signal = signal.to(device=current_omni_platform.device_type)
+    torch.distributed.all_reduce(signal, op=torch.distributed.ReduceOp.MAX, group=group)
+    return bool(signal.item())
+
+
 def _normalize_pipeline_outputs(
     outputs: object,
     *,
@@ -853,25 +878,70 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
 
         return resolved, new_request_ids
 
-    def _prepare_batch_inputs(self, states: list[StepRequestState], new_request_ids: list[str]) -> InputBatch:
+    def _prepare_batch_inputs(
+        self,
+        states: list[StepRequestState],
+        new_request_ids: list[str],
+    ) -> tuple[list[StepRequestState], InputBatch | None, list[RunnerOutput]]:
         # process new reqs
+        prepared_states: list[StepRequestState] = []
+        error_outputs: list[RunnerOutput] = []
         for state in states:
             if state.request_id in new_request_ids:
-                self._initialize_generator(state.sampling)
-                clear_pipeline_stage_durations(self.pipeline)
-                # encode
-                self.pipeline.prepare_encode(state)
-                merge_stage_durations(
-                    state,
-                    consume_pipeline_stage_durations(self.pipeline),
-                )
+                # Everything that runs before ``_dit_any_rank_failed`` must be
+                # inside the try: an exception in ``_initialize_generator`` or
+                # ``clear_pipeline_stage_durations`` on one rank would skip the
+                # all-reduce here while every peer proceeds into it, and the
+                # peers then hang on the NCCL collective until timeout.
+                per_req_exc: BaseException | None = None
+                try:
+                    self._initialize_generator(state.sampling)
+                    clear_pipeline_stage_durations(self.pipeline)
+                    # encode
+                    self.pipeline.prepare_encode(state)
+                    merge_stage_durations(
+                        state,
+                        consume_pipeline_stage_durations(self.pipeline),
+                    )
+                except Exception as exc:
+                    per_req_exc = exc
+                # Pipelines that do rank-0-only work (e.g. MiniMax H3
+                # reference-video prep) must broadcast per-request failures
+                # internally so downstream collectives stay in step; even so,
+                # cross-check that every DiT rank agrees so a rank-local error
+                # (or a future pipeline that omits the guard) does not leave
+                # the process group half-way through a new request.
+                if _dit_any_rank_failed(per_req_exc is not None):
+                    self.state_cache.pop(state.request_id, None)
+                    if per_req_exc is None:
+                        per_req_exc = RuntimeError(
+                            f"Stepwise preparation failed on another DiT rank for {state.request_id}"
+                        )
+                    logger.error(
+                        "Stepwise request preparation failed for %s: %s",
+                        state.request_id,
+                        per_req_exc,
+                        exc_info=isinstance(per_req_exc, Exception),
+                    )
+                    error_outputs.append(
+                        RunnerOutput(
+                            request_id=state.request_id,
+                            step_index=state.step_index,
+                            finished=True,
+                            result=DiffusionOutput.from_exception(per_req_exc),
+                        )
+                    )
+                    continue
+            prepared_states.append(state)
 
+        if not prepared_states:
+            return prepared_states, None, error_outputs
         input_batch = InputBatch.make_batch(
-            states,
+            prepared_states,
             cached_batch=getattr(self, "input_batch", None),
         )
         self.input_batch = input_batch
-        return input_batch
+        return prepared_states, input_batch, error_outputs
 
     def _update_states_after(
         self,
@@ -978,7 +1048,9 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                 and current_omni_platform.is_available()
             ):
                 current_omni_platform.reset_peak_memory_stats()
-            input_batch = self._prepare_batch_inputs(states, new_request_ids)
+            states, input_batch, runner_output_list = self._prepare_batch_inputs(states, new_request_ids)
+            if input_batch is None:
+                return BatchRunnerOutput.from_list(runner_output_list)
             attn_metadata = {}
 
             with set_forward_context(
@@ -995,7 +1067,6 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                         denoise_stage_durations,
                     )
 
-                runner_output_list = []
                 pipeline_interrupted = getattr(self.pipeline, "interrupt", False)
                 if noise_pred is None and pipeline_interrupted:
                     for state in states:
@@ -1048,6 +1119,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                             offset = offset + row_num
                         except Exception as per_req_exc:
                             offset = offset + row_num
+                            self.state_cache.pop(req.request_id, None)
                             logger.error(
                                 "Stepwise per-request error for %s: %s",
                                 req.request_id,
@@ -1059,7 +1131,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                                     request_id=req.request_id,
                                     step_index=req.step_index,
                                     finished=True,
-                                    result=DiffusionOutput(error=str(per_req_exc)),
+                                    result=DiffusionOutput.from_exception(per_req_exc),
                                 )
                             )
 

vllm_omni/engine/async_omni_engine.py

@@ -50,6 +50,7 @@ from vllm_omni.engine.async_engine_utils import (
 )
 from vllm_omni.engine.messages import (
     AbortRequestMessage,
+    AbortResultMessage,
     AddCompanionRequestMessage,
     CollectiveRPCRequestMessage,
     CollectiveRPCResultMessage,
@@ -1736,14 +1737,51 @@ class AsyncOmniEngine:
         return self.stage_metadata[stage_id]
 
     def abort(self, request_ids: list[str]) -> None:
-        """Send abort message to the Orchestrator."""
+        """Fire-and-forget abort: enqueue and return without waiting.
+
+        Prefer :meth:`abort_async` when the caller needs acknowledgment that
+        stage aborts, binding release, and orchestrator request cleanup finished.
+        """
         if self.request_queue is None:
             raise RuntimeError("request_queue is not initialized")
         self.request_queue.sync_q.put(AbortRequestMessage(request_ids=request_ids))
 
-    async def abort_async(self, request_ids: list[str]) -> None:
-        """Async abort API."""
-        self.abort(request_ids)
+    async def abort_async(
+        self,
+        request_ids: list[str],
+        timeout: float | None = None,
+    ) -> None:
+        """Abort requests and wait for orchestrator acknowledgment.
+
+        Unlike :meth:`abort`, this generates an ``rpc_id``, correlates the
+        :class:`AbortResultMessage` via :class:`CorrelatedRpcClient`, and
+        raises if the orchestrator reports failure or times out.
+        """
+        if self.request_queue is None:
+            raise RuntimeError("request_queue is not initialized")
+        transport = self._correlated_rpc_client
+        if transport is None:
+            raise RuntimeError("correlated RPC client is not initialized")
+
+        rpc_id = uuid.uuid4().hex
+        msg = AbortRequestMessage(request_ids=request_ids, rpc_id=rpc_id)
+
+        def _wait() -> AbortResultMessage:
+            result_msg = transport.execute(
+                ("abort", rpc_id),
+                msg,
+                timeout=timeout,
+                timeout_message=f"abort timed out after {timeout} seconds",
+                block_on_submit=True,
+            )
+            if not isinstance(result_msg, AbortResultMessage):
+                raise RuntimeError(f"unexpected abort result type: {type(result_msg).__name__}")
+            return result_msg
+
+        loop = asyncio.get_running_loop()
+        result_msg = await loop.run_in_executor(None, _wait)
+        if not result_msg.success:
+            raise RuntimeError(result_msg.error or "abort failed")
 
     def submit_interaction(
         self,

vllm_omni/engine/messages.py

@@ -42,6 +42,20 @@ class AddCompanionRequestMessage(EngineQueueMessage, kw_only=True):
 class AbortRequestMessage(EngineQueueMessage, kw_only=True):
     type: Literal["abort"] = "abort"
     request_ids: list[str]
+    # When set, Orchestrator emits AbortResultMessage on rpc_async_queue so
+    # callers can await acknowledgment. Omit for fire-and-forget abort().
+    rpc_id: str | None = None
+
+
+class AbortResultMessage(EngineQueueMessage, kw_only=True):
+    type: Literal["abort_result"] = "abort_result"
+    rpc_id: str
+    success: bool
+    error: str | None = None
+
+    @property
+    def rpc_correlation_key(self) -> tuple[str, str]:
+        return ("abort", self.rpc_id)
 
 
 class InteractionMessage(EngineQueueMessage, kw_only=True):

vllm_omni/engine/orchestrator.py

@@ -36,6 +36,7 @@ from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker
 from vllm_omni.engine.membership_controller import MembershipController
 from vllm_omni.engine.messages import (
     AbortRequestMessage,
+    AbortResultMessage,
     AddCompanionRequestMessage,
     CollectiveRPCRequestMessage,
     CollectiveRPCResultMessage,
@@ -817,13 +818,35 @@ class Orchestrator:
         )
 
     async def _handle_abort(self, msg: AbortRequestMessage) -> None:
-        """Handle an abort message from the main thread."""
+        """Handle an abort message from the main thread.
+
+        When ``msg.rpc_id`` is set, emit :class:`AbortResultMessage` after
+        stage aborts, binding release, and request cleanup complete so the
+        caller can await acknowledgment. Fire-and-forget aborts omit rpc_id.
+        """
         request_ids = msg.request_ids
-        # _cleanup_request_ids is CFG-aware: it expands aborted parents to
-        # their companions and fails a deferred parent whose companion is
-        # aborted before its output arrived.
-        await self._cleanup_request_ids(list(request_ids), abort=True)
-        logger.info("[Orchestrator] Aborted request(s) %s", request_ids)
+        error: str | None = None
+        try:
+            # _cleanup_request_ids is CFG-aware: it expands aborted parents to
+            # their companions and fails a deferred parent whose companion is
+            # aborted before its output arrived.
+            await self._cleanup_request_ids(list(request_ids), abort=True)
+            logger.info("[Orchestrator] Aborted request(s) %s", request_ids)
+        except Exception as exc:
+            error = str(exc)
+            logger.exception("[Orchestrator] Abort failed for request(s) %s", request_ids)
+            if msg.rpc_id is None:
+                raise
+        if msg.rpc_id is not None:
+            # Always emit a result when rpc_id is set so CorrelatedRpcClient
+            # waiters are unblocked even on failure.
+            await self.rpc_async_queue.put(
+                AbortResultMessage(
+                    rpc_id=msg.rpc_id,
+                    success=error is None,
+                    error=error,
+                )
+            )
 
     async def _handle_interaction(self, msg: InteractionMessage) -> None:
         """Handle a midway interaction for an active streaming diffusion request."""

vllm_omni/engine/stage_pool.py

@@ -94,6 +94,16 @@ class StagePool:
 
     DISPATCH_WAIT_TIMEOUT_S: float = 10.0
     DISPATCH_RETRY_INTERVAL_S: float = 0.1
+    # Only these EngineCore helpers may skip collective_rpc_async. A generic
+    # ``{method}_async`` on AsyncMPClient must not silently drop timeout.
+    _ENGINE_CORE_CONTROL_ASYNC_METHODS = frozenset(
+        {
+            "pause_scheduler",
+            "resume_scheduler",
+            "sleep",
+            "wake_up",
+        }
+    )
 
     def __init__(
         self,
@@ -1214,6 +1224,7 @@ class StagePool:
         kwargs: dict[str, Any] | None = None,
     ) -> dict[str, Any] | Any:
         """Dispatch a stage-scoped control-plane RPC to one physical route."""
+        args = tuple(args or ())
         kwargs = dict(kwargs or {})
         client = self.clients[replica_id]
         if client is None:
@@ -1222,6 +1233,14 @@ class StagePool:
                 "error": f"stage {self.stage_id} replica {replica_id} is not attached",
             }
         try:
+            if self.stage_type != "diffusion" and method in self._ENGINE_CORE_CONTROL_ASYNC_METHODS:
+                client_method = getattr(client, f"{method}_async", None)
+                if callable(client_method):
+                    result = client_method(*args, **kwargs)
+                    if timeout is not None:
+                        return await asyncio.wait_for(result, timeout=timeout)
+                    return await result
+
             return await client.collective_rpc_async(
                 method=method,
                 timeout=timeout,

vllm_omni/entrypoints/async_omni.py

@@ -146,6 +146,11 @@ class AsyncOmni(EngineClient, OmniBase):
         OmniBase.__init__(self, model=model, **kwargs)
         self._pause_cond: asyncio.Condition = asyncio.Condition()
         self._paused: bool = False
+        # True after pause_generation() or AR EngineCore sleep; wake_up must
+        # not reopen generate() until resume_generation(). Diffusion-only
+        # sleep uses _paused as a temporary admission gate and clears it
+        # on wake so sleep → wake → generate keeps working.
+        self._hold_admission_until_resume: bool = False
         self._sleeping_tags: set[str] = set()
         self._level2_sleeping: bool = False
         self._duplex_request_client: DuplexRequestClient | None = None
@@ -1009,6 +1014,30 @@ class AsyncOmni(EngineClient, OmniBase):
 
         return results
 
+    async def _engine_core_rpc(
+        self,
+        method: str,
+        *,
+        stage_ids: list[int],
+        args: tuple[Any, ...] = (),
+        kwargs: dict[str, Any] | None = None,
+    ) -> list[Any]:
+        """Call an AR EngineCore helper via collective_rpc (orchestrator loop).
+
+        StagePool resolves ``{method}_async`` on the AR client when present
+        (vLLM AsyncMPClient convention). Raises if any replica reports failure.
+        """
+        results = await self.collective_rpc(
+            method=method,
+            args=args,
+            kwargs=kwargs,
+            stage_ids=stage_ids,
+        )
+        for result in results:
+            if isinstance(result, dict) and result.get("error"):
+                raise RuntimeError(f"{method} failed: {result['error']}")
+        return results
+
     @staticmethod
     def _coerce_stage_bool(result: Any) -> bool:
         """Reduce a stage RPC result to a boolean.
@@ -1079,7 +1108,12 @@ class AsyncOmni(EngineClient, OmniBase):
         await self._abort(internal_req_ids)
 
     async def _abort(self, request_ids: list[str]) -> None:
-        """Submit request IDs to be aborted to the engine."""
+        """Abort request IDs via the engine and clean frontend state after ack.
+
+        Waits for orchestrator abort acknowledgment before popping
+        ``request_states`` so generate() cleanup stays consistent with
+        backend binding/request teardown. Orchestrator abort errors propagate.
+        """
         await self.engine.abort_async(request_ids)
         for rid in request_ids:
             state = self.request_states.pop(rid, None)
@@ -1089,23 +1123,77 @@ class AsyncOmni(EngineClient, OmniBase):
         if self.log_stats:
             logger.info("[AsyncOmni] Aborted request(s) %s", ",".join(request_ids))
 
+    def _split_stage_ids_by_type(self, stage_ids: list[int] | None = None) -> tuple[list[int], list[int]]:
+        """Split stage ids into AR/LLM (EngineCore) vs diffusion (worker RPC)."""
+        n_stages = len(self.engine.stage_clients)
+        if stage_ids is None:
+            stage_ids = list(range(n_stages))
+        else:
+            invalid = [sid for sid in stage_ids if not isinstance(sid, int) or sid < 0 or sid >= n_stages]
+            if invalid:
+                raise ValueError(
+                    f"Invalid stage_ids {invalid}; valid range is 0..{n_stages - 1}"
+                    if n_stages
+                    else f"Invalid stage_ids {invalid}; this engine has no stages"
+                )
+        ar_stage_ids: list[int] = []
+        diffusion_stage_ids: list[int] = []
+        for sid in stage_ids:
+            client = self.engine.stage_clients[sid]
+            if getattr(client, "stage_type", "llm") == "diffusion":
+                diffusion_stage_ids.append(sid)
+            else:
+                ar_stage_ids.append(sid)
+        return ar_stage_ids, diffusion_stage_ids
+
     async def pause_generation(
         self,
         *,
         mode: PauseMode = "abort",
         wait_for_inflight_requests: bool = False,
         clear_cache: bool = True,
+        stage_ids: list[int] | None = None,
     ) -> None:
-        """Pause generation."""
+        """Pause generation, mirroring vLLM AsyncLLM.pause_generation.
+
+        1. Stop frontend admission (``_paused``).
+        2. For AR/LLM stages, call EngineCore.pause_scheduler via the
+           Orchestrator loop (abort/wait/keep + optional cache clear).
+        3. Diffusion stages have no EngineCore scheduler — only frontend
+           admission is paused for them.
+
+        Note: ``sleep()`` already pauses the AR scheduler internally (same as
+        vLLM EngineCore.sleep). Call this API when you need pause *without*
+        freeing GPU memory (e.g. weight sync).
+        """
+        if wait_for_inflight_requests:
+            mode = "wait"
+
         async with self._pause_cond:
-            if self._paused:
-                return
+            # Keep running EngineCore pause + cache clear even when frontend
+            # admission is already paused (sleep or a prior pause_generation).
             self._paused = True
+            self._hold_admission_until_resume = True
 
-        # TODO: Implement request draining if wait_for_inflight_requests
+        ar_stage_ids, _diffusion_stage_ids = self._split_stage_ids_by_type(stage_ids)
+        if ar_stage_ids:
+            logger.info(
+                "[%s] Pausing AR stage(s) %s via EngineCore.pause_scheduler(mode=%s)",
+                self._name,
+                ar_stage_ids,
+                mode,
+            )
+            # Same API name as vLLM AsyncMPClient.pause_scheduler_async; routed
+            # through collective_rpc so it runs on the orchestrator event loop.
+            await self._engine_core_rpc(
+                "pause_scheduler",
+                stage_ids=ar_stage_ids,
+                kwargs={"mode": mode, "clear_cache": clear_cache},
+            )
 
+        # Frontend / sender-side cache clear (P0). EngineCore.pause_scheduler
+        # already clears AR-side caches when clear_cache=True.
         if clear_cache:
-            # Clear caches for all stages.
             await self.reset_prefix_cache(
                 reset_running_requests=not wait_for_inflight_requests,
                 reset_connector=True,
@@ -1113,14 +1201,20 @@ class AsyncOmni(EngineClient, OmniBase):
             await self.reset_mm_cache()
             await self.reset_encoder_cache()
 
-    async def resume_generation(self) -> None:
-        """Resume generation."""
+    async def resume_generation(self, stage_ids: list[int] | None = None) -> None:
+        """Resume generation after :meth:`pause_generation`."""
+        ar_stage_ids, _diffusion_stage_ids = self._split_stage_ids_by_type(stage_ids)
+        if ar_stage_ids:
+            logger.info("[%s] Resuming AR stage(s) %s via EngineCore", self._name, ar_stage_ids)
+            await self._engine_core_rpc("resume_scheduler", stage_ids=ar_stage_ids)
+
         async with self._pause_cond:
             self._paused = False
+            self._hold_admission_until_resume = False
             self._pause_cond.notify_all()
 
     async def is_paused(self) -> bool:
-        """Check if paused."""
+        """Check if frontend admission is paused."""
         async with self._pause_cond:
             return self._paused
 
@@ -1178,33 +1272,65 @@ class AsyncOmni(EngineClient, OmniBase):
     async def sleep(
         self, stage_ids: list[int] | None = None, level: int = 2, mode: PauseMode = "abort"
     ) -> list[OmniACK]:
+        """Put stages to sleep.
+
+        AR/LLM stages use EngineCore.sleep (pause scheduler, wait idle, then
+        offload/discard memory) — matching vLLM AsyncLLM.sleep.
+
+        Diffusion stages keep the existing worker-level handle_sleep_task RPC
+        because StageDiffusionProc does not expose EngineCore.pause_scheduler.
+
+        Frontend admission is blocked at the start of this call (``_paused``)
+        so pipelined :meth:`generate` cannot race into stages while sleep is
+        in flight. This does **not** invoke EngineCore.pause_scheduler again
+        (sleep already pauses the AR scheduler).
+
+        For AR / mixed engines, ``wake_up`` does **not** clear ``_paused``;
+        callers must :meth:`resume_generation` when ready (typical trainer
+        order: pause → abort → sleep → train → wake → resume). Diffusion-only
+        engines have no EngineCore pause to hold, so ``wake_up`` restores
+        admission and ``sleep → wake → generate`` keeps working.
+        """
+        # Block admission before any sleep RPC so generate() waits on
+        # _pause_cond during the drain/offload window. EngineCore.sleep will
+        # pause the AR scheduler itself; avoid a second pause_scheduler here.
+        async with self._pause_cond:
+            self._paused = True
+
         self._final_output_handler()
-        if stage_ids is None:
-            stage_ids = list(range(len(self.engine.stage_clients)))
-        total_workers = 0
-        for sid in stage_ids:
-            client = self.engine.stage_clients[sid]
-            # During the Diffusion phase, regardless of the TP amount,
-            # currently only a summary ACK is reported at Rank 0.
-            if getattr(client, "stage_type", "") == "diffusion":
-                total_workers += 1
-            else:
-                config = self.engine.stage_vllm_configs[sid]
-                actual_tp = config.parallel_config.tensor_parallel_size if config else 1
-                total_workers += actual_tp
+        ar_stage_ids, diffusion_stage_ids = self._split_stage_ids_by_type(stage_ids)
+        final_acks: list[OmniACK] = []
+        if ar_stage_ids:
+            self._hold_admission_until_resume = True
+            logger.info(
+                "[%s] Sleeping AR stage(s) %s via EngineCore.sleep(level=%s, mode=%s)",
+                self._name,
+                ar_stage_ids,
+                level,
+                mode,
+            )
+            await self._engine_core_rpc(
+                "sleep",
+                stage_ids=ar_stage_ids,
+                args=(level, mode),
+            )
+            # EngineCore.sleep has no OmniACK handshake; emit stage-level SUCCESS
+            # markers so callers/tests that count ACKs keep a stable API.
+            task_id = f"engine_core-sleep-{uuid.uuid4().hex[:8]}"
+            final_acks.extend(
+                OmniACK(
+                    task_id=task_id,
+                    status="SUCCESS",
+                    stage_id=sid,
+                    rank=0,
+                    metadata={"path": "engine_core", "level": level, "mode": mode},
+                )
+                for sid in ar_stage_ids
+            )
+
+        if diffusion_stage_ids:
+            final_acks.extend(await self._sleep_diffusion(diffusion_stage_ids, level))
 
-        task_id = str(uuid.uuid4())
-        self.event_resolver.watch_task(task_id, expected_count=total_workers)
-        logger.info(f"[{self._name}] Sleep initiated (Task: {task_id}). Awaiting {total_workers} ACKs...")
-        task = OmniSleepTask(level=level, task_id=task_id)
-        rpc_results = await self.collective_rpc(method="handle_sleep_task", args=(task,), stage_ids=stage_ids)
-        final_acks = []
-        for stage_res in rpc_results:
-            worker_acks = stage_res if isinstance(stage_res, list) else [stage_res]
-            for ack in worker_acks:
-                if ack is not None:
-                    await self.event_resolver.resolve(ack)
-                    final_acks.append(ack)
         if not hasattr(self, "_sleeping_tags"):
             self._sleeping_tags = set()
         self._sleeping_tags.update([CuMemTag.WEIGHTS.value, CuMemTag.KV_CACHE.value])
@@ -1212,7 +1338,36 @@ class AsyncOmni(EngineClient, OmniBase):
             self._level2_sleeping = True
         return final_acks
 
+    async def _sleep_diffusion(self, stage_ids: list[int], level: int) -> list[OmniACK]:
+        """Worker-level sleep RPC for diffusion stages only."""
+        # Diffusion reports one summary ACK at rank 0 regardless of TP.
+        total_workers = len(stage_ids)
+        task_id = str(uuid.uuid4())
+        self.event_resolver.watch_task(task_id, expected_count=total_workers)
+        logger.info("[%s] Sleep (diffusion) initiated (Task: %s).", self._name, task_id)
+        task = OmniSleepTask(level=level, task_id=task_id)
+        rpc_results = await self.collective_rpc(method="handle_sleep_task", args=(task,), stage_ids=stage_ids)
+        final_acks: list[OmniACK] = []
+        for stage_res in rpc_results:
+            worker_acks = stage_res if isinstance(stage_res, list) else [stage_res]
+            for ack in worker_acks:
+                if ack is not None:
+                    await self.event_resolver.resolve(ack)
+                    final_acks.append(ack)
+        return final_acks
+
     async def wake_up(self, stage_ids: list[int] | None = None, tags: list[str] | None = None) -> list[OmniACK]:
+        """Wake stages after sleep.
+
+        AR/LLM stages use EngineCore.wake_up (restore memory, auto-resume
+        scheduler). Diffusion stages keep the worker-level wake RPC.
+
+        Does **not** clear the frontend ``_paused`` admission gate when
+        :meth:`pause_generation` ran or AR stages were slept — call
+        :meth:`resume_generation` when the trainer is ready to admit new
+        requests. Diffusion-only ``sleep`` uses ``_paused`` only as a race
+        guard; this method restores admission after a successful wake.
+        """
         self._final_output_handler()
 
         if getattr(self, "_level2_sleeping", False):
@@ -1231,30 +1386,30 @@ class AsyncOmni(EngineClient, OmniBase):
             logger.info(f"[{self._name}] Requested tags {tags} are already warm. Skipping wake_up.")
             return []
 
-        if stage_ids is None:
-            stage_ids = list(range(len(self.engine.stage_clients)))
-        total_workers = 0
-        for sid in stage_ids:
-            client = self.engine.stage_clients[sid]
-            if getattr(client, "stage_type", "") == "diffusion":
-                total_workers += 1
-            else:
-                config = self.engine.stage_vllm_configs[sid]
-                total_workers += config.parallel_config.tensor_parallel_size if config else 1
-        task_id = str(uuid.uuid4())
-        self.event_resolver.watch_task(task_id, expected_count=total_workers)
-        logger.info(f"[{self._name}] Wake-up initiated (Task: {task_id}). Awaiting {total_workers} ACKs...")
-        task = OmniWakeTask(tags=requested_tags, task_id=task_id)
-        rpc_results = await self.collective_rpc(method="handle_wake_task", args=(task,), stage_ids=stage_ids)
-        final_acks = []
-        for stage_res in rpc_results:
-            worker_acks = stage_res if isinstance(stage_res, list) else [stage_res]
-            for ack in worker_acks:
-                if ack is not None:
-                    await self.event_resolver.resolve(ack)
-                    final_acks.append(ack)
-        current_omni_platform.synchronize()
-        await asyncio.sleep(0.1)
+        ar_stage_ids, diffusion_stage_ids = self._split_stage_ids_by_type(stage_ids)
+
+        final_acks: list[OmniACK] = []
+        if ar_stage_ids:
+            logger.info("[%s] Waking AR stage(s) %s via EngineCore", self._name, ar_stage_ids)
+            await self._engine_core_rpc(
+                "wake_up",
+                stage_ids=ar_stage_ids,
+                kwargs={"tags": requested_tags},
+            )
+            task_id = f"engine_core-wake-{uuid.uuid4().hex[:8]}"
+            final_acks.extend(
+                OmniACK(
+                    task_id=task_id,
+                    status="SUCCESS",
+                    stage_id=sid,
+                    rank=0,
+                    metadata={"path": "engine_core", "tags": list(requested_tags)},
+                )
+                for sid in ar_stage_ids
+            )
+
+        if diffusion_stage_ids:
+            final_acks.extend(await self._wake_diffusion(diffusion_stage_ids, requested_tags))
 
         for t in requested_tags:
             if hasattr(self, "_sleeping_tags"):
@@ -1263,7 +1418,37 @@ class AsyncOmni(EngineClient, OmniBase):
         # wake support (e.g. tags=["kv_cache"] only) is added in the future.
         if not getattr(self, "_sleeping_tags", None):
             self._level2_sleeping = False
-        logger.info(f"[{self._name}] All {len(final_acks)}/{total_workers} workers reported WARM for task {task_id}.")
+        logger.info(
+            "[%s] Wake-up complete for stage(s) %s.",
+            self._name,
+            ar_stage_ids + diffusion_stage_ids,
+        )
+        # Diffusion-only sleep uses `_paused` as a race guard. Restore
+        # generate() admission after memory is back. AR/mixed sleep and
+        # pause_generation keep the trainer hold until resume_generation.
+        if not getattr(self, "_hold_admission_until_resume", False):
+            async with self._pause_cond:
+                self._paused = False
+                self._pause_cond.notify_all()
+        return final_acks
+
+    async def _wake_diffusion(self, stage_ids: list[int], requested_tags: list[str]) -> list[OmniACK]:
+        """Worker-level wake RPC for diffusion stages only."""
+        total_workers = len(stage_ids)
+        task_id = str(uuid.uuid4())
+        self.event_resolver.watch_task(task_id, expected_count=total_workers)
+        logger.info("[%s] Wake-up (diffusion) initiated (Task: %s).", self._name, task_id)
+        task = OmniWakeTask(tags=requested_tags, task_id=task_id)
+        rpc_results = await self.collective_rpc(method="handle_wake_task", args=(task,), stage_ids=stage_ids)
+        final_acks: list[OmniACK] = []
+        for stage_res in rpc_results:
+            worker_acks = stage_res if isinstance(stage_res, list) else [stage_res]
+            for ack in worker_acks:
+                if ack is not None:
+                    await self.event_resolver.resolve(ack)
+                    final_acks.append(ack)
+        current_omni_platform.synchronize()
+        await asyncio.sleep(0.1)
         return final_acks
 
     async def is_sleeping(self) -> bool: