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

This reverts commit a94b2580a1.

docs/design/feature/attention_backend_selection.md

@@ -61,7 +61,8 @@ unstructured keyword dictionaries:
 - `quant` is consumed by FlashInfer and TRTLLM with backend-specific value
   validation;
 - `skip_softmax` is consumed by TRTLLM; and
-- `block_sparse` is consumed by block-sparse backends such as RainFusion.
+- `block_sparse` is consumed by block-sparse backends such as RainFusion; and
+- `fastvideo_vsa_topk` is consumed by FastVideo VSA.
 
 A backend reads only the fields it owns and rejects incompatible values. New
 options should be added to a shared typed spec only when more than one backend
@@ -87,3 +88,19 @@ An implementation change is complete only when it:
    corresponding user guide; and
 6. updates the overview matrix without moving algorithm details back into the
    landing page.
+
+## FastVideo VSA model metadata
+
+FastVideo VSA support is scoped to `FastVideo/FastWan2.2-TI2V-5B-Diffusers`, whose text-to-video and image-to-video modes both use `Wan22Pipeline`. It operates on a flattened DiT sequence but partitions tokens in
+the original latent video grid. Wan integrations therefore attach the
+post-patch `(T, H, W)` shape as `vsa_dit_seq_shape` attention metadata. The
+backend validates that `T * H * W` equals the sequence length, derives the
+runtime block count, and selects the configured top-k key/value blocks for
+each query block.
+
+Checkpoint-specific compensation remains model-owned. Wan layers pass an
+optional learned `gate_compress` projection through attention metadata when
+the checkpoint contains those weights. Checkpoints without the projection do
+not allocate or execute it. When top-k selects every block, native checkpoints
+route to SDPA, while FastVideo DMD checkpoints preserve the VSA all-block path
+to retain checkpoint semantics.

docs/user_guide/diffusion/attention_backends.md

@@ -20,6 +20,7 @@ For the internal selector, registry, and platform contract, see
 | Install and use SageAttention 2.2 or SageAttention3 | [SageAttention](attention_backends/sage.md) |
 | Match training or rollout kernels loaded from Hugging Face | [Hugging Face Hub Backends](attention_backends/huggingface_hub.md) |
 | Use block-sparse video attention on Ascend NPU | [RainFusion](attention_backends/rainfusion.md) |
+| Use FastVideo VSA with FastWan2.2-TI2V-5B on CUDA | [FastVideo VSA](attention_backends/fastvideo_vsa.md) |
 
 ## Backend options
 
@@ -35,6 +36,7 @@ For the internal selector, registry, and platform contract, see
 | `FLASH_ATTN_HUB` | Hub kernel | FlashAttention 2 from Hugging Face `kernels` | [Hugging Face Hub Backends](attention_backends/huggingface_hub.md) |
 | `FLASH_ATTN_3_HUB` | Hub kernel | FlashAttention 3 from Hugging Face `kernels` on Hopper or newer | [Hugging Face Hub Backends](attention_backends/huggingface_hub.md) |
 | `RAINFUSION_ATTN` | Block sparse | MindIE-SD RainFusion video attention on Ascend NPU | [RainFusion](attention_backends/rainfusion.md) |
+| `FASTVIDEO_VSA` | Block sparse | FastVideo variable sparse self-attention for FastWan2.2-TI2V-5B on CUDA | [FastVideo VSA](attention_backends/fastvideo_vsa.md) |
 
 ## Configuration
 
@@ -105,6 +107,7 @@ Backend-specific typed blocks are documented with their consumers:
   and [TRTLLM SAGE](attention_backends/trtllm.md#sage-quantization).
 - `skip_softmax`: [TRTLLM Skip-Softmax](attention_backends/trtllm.md#skip-softmax).
 - `block_sparse`: [RainFusion](attention_backends/rainfusion.md#configuration).
+- `fastvideo_vsa_topk`: [FastVideo VSA](attention_backends/fastvideo_vsa.md#choose-top-k).
 
 ## Platform defaults
 

docs/user_guide/diffusion/attention_backends/fastvideo_vsa.md

@@ -0,0 +1,125 @@
+# FastVideo VSA
+
+FastVideo Variable Sparse Attention (VSA) accelerates the `FastVideo/FastWan2.2-TI2V-5B-Diffusers` model's
+self-attention by partitioning the post-patch latent grid into spatiotemporal
+blocks. For every query block, VSA scores the key/value blocks and computes
+attention only against the selected top-k blocks.
+
+The supported checkpoint provides both text-to-video and image-to-video modes through `Wan22Pipeline`; the separate Wan I2V-14B, S2V, and VACE pipelines are outside this backend's supported scope.
+
+VSA is a CUDA-only, explicitly selected backend. It requires the
+`fastvideo-kernel` package and currently supports non-causal self-attention
+with equal query and key/value sequence lengths. Unsupported shapes, masks,
+dtypes, sequence-parallel execution, or kernel failures fall back to
+`TORCH_SDPA` and emit a warning with the reason.
+
+## Enable the backend
+
+For online serving, select the backend with the existing attention backend
+flag. Use `--fastvideo-vsa-topk` to set the number of key/value blocks retained
+for every query block:
+
+```bash
+vllm-omni serve <model> \
+  --diffusion-attention-backend FASTVIDEO_VSA \
+  --fastvideo-vsa-topk 64
+```
+
+The backwards-compatible environment variable selects the backend with the
+default `topk=64`:
+
+```bash
+export DIFFUSION_ATTENTION_BACKEND=FASTVIDEO_VSA
+vllm-omni serve <model>
+```
+
+To tune top-k, pass the CLI backend and top-k flags together as shown above,
+or use the equivalent structured configuration:
+
+```bash
+vllm-omni serve <model> \
+  --diffusion-attention-config \
+  '{"default":{"backend":"FASTVIDEO_VSA","fastvideo_vsa_topk":64}}'
+```
+
+Do not combine `--diffusion-attention-backend` with an explicit
+`diffusion_attention_config.default.backend`. The top-k value must be positive
+and is valid only when the default backend is `FASTVIDEO_VSA`.
+
+For a deploy YAML stage, use either the shorthand fields:
+
+```yaml
+stages:
+  - stage_id: 0
+    diffusion_attention_backend: FASTVIDEO_VSA
+    fastvideo_vsa_topk: 64
+```
+
+or the structured configuration:
+
+```yaml
+stages:
+  - stage_id: 0
+    diffusion_attention_config:
+      default:
+        backend: FASTVIDEO_VSA
+        fastvideo_vsa_topk: 64
+```
+
+## Choose top-k
+
+At runtime the backend logs the sequence shape and derived block count:
+
+```text
+FASTVIDEO_VSA routing: seq_len=27280, dit_seq_shape=(31, 22, 40),
+block_size=(4, 8, 8), num_blocks=120, topk=64, keep_ratio=53.3%,
+checkpoint_mode=native, route=VSA
+```
+
+Use `num_blocks` as the upper bound when tuning top-k. Top-k is the number of
+key/value blocks retained for each query block, not the number of tokens or
+the total number of blocks processed by the layer.
+
+- A smaller top-k increases sparsity and may reduce attention computation, but
+  it can remove relevant blocks and reduce visual quality or temporal
+  consistency. Routing and padding overhead can also make a smaller value
+  slower for some shapes.
+- A larger top-k retains more context and generally approaches dense-attention
+  quality, but reduces the potential speedup and increases memory traffic.
+- `topk > num_blocks` is invalid for the runtime shape and falls back to SDPA.
+- For a native checkpoint, `topk == num_blocks` routes to SDPA because scoring
+  every block provides no sparsity benefit.
+- For a FastVideo DMD checkpoint, `topk == num_blocks` stays on the VSA
+  all-block path so the checkpoint keeps its trained compensation semantics.
+
+There is no universally optimal value. Resolution, frame count, GPU, kernel
+version, and checkpoint all affect both quality and latency. Start from the
+logged `num_blocks`, test several keep ratios on the target workload, and
+compare output quality against the same checkpoint running dense attention.
+
+## Checkpoint behavior
+
+The backend does not expose a user-selectable gate mode. Wan checkpoints that
+contain learned `to_gate_compress` weights use that projection automatically.
+When those weights are absent, the unused projection is removed and VSA uses
+the sparse branch without learned compensation.
+
+FastVideo DMD checkpoints use their fixed distilled timestep schedule. Native
+Wan checkpoints keep their normal scheduler and inference-step configuration;
+selecting VSA does not turn a native checkpoint into a distilled model.
+
+## Verify routing and fallback
+
+Check the startup and first-forward logs instead of assuming that selecting
+the backend guarantees sparse execution:
+
+- `route=VSA` means top-k block selection is active.
+- `route=VSA_ALL_BLOCKS` means the FastVideo DMD checkpoint retained all
+  blocks through the VSA kernel.
+- `route=SDPA` or `FASTVIDEO_VSA falling back to SDPA: ...` means dense SDPA
+  executed; the warning includes the reason.
+
+The current kernel requires CUDA tensors in FP16 or BF16, 256-token blocks,
+standard `head_size**-0.5` scaling, equal Q/K/V head counts, no attention mask,
+and no active sequence-parallel context. NPU and XPU paths do not execute the
+FastVideo VSA CUDA kernel.

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

@@ -80,7 +80,7 @@ between the legacy and direct planar paths
 ([full result](https://github.com/vllm-project/vllm-omni/pull/6288#issuecomment-5337546499)):
 
 | Metric | Legacy | Direct planar | Change |
-|---|---:|---:|---:|
+| --- | ---: | ---: | ---: |
 | Median wall time | 1.805 s | 1.394 s | -22.8% |
 | Median process CPU time | 3.613 s | 3.207 s | -11.2% |
 | Peak RSS | 3182 MiB | 2794 MiB | -387 MiB (-12.2%) |
@@ -183,7 +183,7 @@ vllm serve "${MODEL}" \
 Use the profile that matches the per-GPU memory capacity:
 
 | Profile | GPUs | Starting shape | Resident DiT blocks | Attention | Execution | Status |
-|---|---:|---:|---:|---|---|---|
+| --- | ---: | ---: | ---: | --- | --- | --- |
 | `rtx5090` | 2 x 32 GB | 1344x768 | 20 | cuDNN attention | eager | Target-hardware validated |
 | `rtx4090` | 2 x 24 GB | 1024x576 | 12 | cuDNN attention | eager | Capacity-proxy starting point |
 
@@ -202,8 +202,8 @@ increasing it on a different request shape.
 At vLLM-Omni commit `ae6577ea`, one full 50-step T2VA request completed on
 2 x RTX 5090 without OOM:
 
-| Shape | Frames | Client E2E | Sampled peak/GPU | Output validation |
-|---:|---:|---:|---:|---|
+| Shape    | Frames        | Client E2E | Sampled peak/GPU       | Output validation                                            |
+| -------: | ------------: | ---------: | ---------------------: | -----------------------------------------------------------: |
 | 1344x768 | 124 at 24 FPS | 8 min 38 s | approximately 22.6 GiB | H.264 video + 32 kHz stereo AAC; full `ffmpeg` decode passed |
 
 This is a single end-to-end validation run, not a warmed multi-run latency
@@ -272,7 +272,7 @@ vllm serve "${MODEL}" \
 Do not add `--enforce-eager` to this performance configuration. The first
 request includes regional compilation; warm the server once before measuring
 steady-state latency. H3 is CFG-distilled, so `--cfg-parallel-size` must remain
-1. The H3 VAE supports its native `tile` mode, not
+`1`. The H3 VAE supports its native `tile` mode, not
 `spatial_shard_height` or `spatial_shard_width`.
 
 ### Attention Backends
@@ -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
@@ -530,7 +587,7 @@ vLLM-Omni with MiniMax H3 support, BF16. gfx942 rows measured with the
 `0.26.0+rocm723` wheel (HIP 7.2).
 
 | Workload | Configuration | Observed result |
-|----------|---------------|-----------------|
+| ---------- | --------------- | ----------------- |
 | T2VA, 1344x768, 209 frames, 50 steps | 4x gfx942 (MI300X), FLASH_ATTN, USP4, text-enc TP4, VAE PP4 tile | encode 0.09 s, denoise 244.04 s, decode 4.15 s, 267.42 s client E2E; H.264 24 FPS + 32 kHz stereo AAC |
 | FL2VA, 1344x768, 209 frames, 50 steps | 4x gfx942 (MI300X), FLASH_ATTN, USP4, text-enc TP4, VAE PP4 tile | encode 13.98 s, denoise 257.58 s, decode 4.11 s, 287.07 s client E2E; H.264 24 FPS + 32 kHz stereo AAC |
 | T2VA, 832x480, ~4 s, 40 steps | 1x gfx950 (MI350), FLASH_ATTN, CPU offload | valid MP4 (H.264 + synced audio); ~0.73 s/denoise-step (~1.37 it/s), ~55 s client E2E incl. warmup |
@@ -712,7 +769,7 @@ at most 15 seconds combined.
 ## Official input matrix and limits
 
 | Task | Supported references | Limits |
-|------|----------------------|--------|
+| ------ | ---------------------- | -------- |
 | T2VA | text only | prompt must be non-empty |
 | FL2VA | first image, last image, or ordered first+last images | at most 2 images; `frame_indices` is `[0]`, `[-1]`, or `[0,-1]` |
 | Ref2VA | image-only, image+image, image+video, video+audio, and mixed image/video/audio | images ≤9, videos ≤3, audios ≤3, total references ≤12; audio requires a visual reference |
@@ -743,13 +800,17 @@ Omitting `quality` preserves the startup default: it uses the reference path
 normally, or the server-configured profile when the server was started with
 `--cache-backend cache_dit`.
 
+Turbo is independent of this quality switch: `quality` selects a Cache-DiT
+policy, while Turbo changes the active LoRA weights and sampling schedule. See
+[LoRA](#lora) below.
+
 The following result was measured on 4× NVIDIA H200 with SP4, text-encoder
 TP4, 1344×768, 124 frames, 24 FPS, and 50 inference steps. One full
 `lossless` warmup was excluded, followed by three fixed prompt/seed pairs in
 balanced switch order.
 
 | `quality` | Median inference latency | Speedup | SSIM vs `lossless` | PSNR vs `lossless` | Expected trade-off |
-|---|---:|---:|---:|---:|---|
+| --- | ---: | ---: | ---: | ---: | --- |
 | `lossless` | 85.49 s | 1.00× | 1.0000 | exact | Native reference path |
 | `high` | 63.36 s | 1.35× | 0.9709 | 34.98 dB | Faster with measured same-seed deviation |
 
@@ -759,10 +820,49 @@ balanced switch order.
 > workload. The values above apply to this deployment and are not universal
 > guarantees. `lossless` remains the exact reference path.
 
+## LoRA
+
+### Turbo LoRA
+
+Only the native Diffusers 4-step FL2VA/T2VA v1.0 artifact is supported:
+
+```text
+lightx2v/Minimax-h3-Turbo/minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors
+```
+
+Download only that file:
+
+```bash
+export TURBO_DIR=/path/to/minimax-h3-turbo
+export TURBO_FILE=minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors
+hf download lightx2v/Minimax-h3-Turbo "${TURBO_FILE}" --local-dir "${TURBO_DIR}"
+export TURBO_LORA="${TURBO_DIR}/${TURBO_FILE}"
+```
+
+Start from a non-offloaded FL2VA server command and add
+`--task-type fl2va --lora-backend peft --lora-path "${TURBO_LORA}"`.
+`--lora-path` preloads the adapter; each request still activates it and uses
+the published sampling settings:
+
+```bash
+-F 'num_inference_steps=5' \
+-F 'flow_shift=6' \
+-F 'extra_params={"task":"t2va","duration":4.4,"audio_flow_shift":3.0}' \
+-F "lora={\"name\":\"h3-turbo-v1.0\",\"path\":\"${TURBO_LORA}\",\"scale\":1.0}"
+```
+
+For FL2VA, change `task` and add `input_reference` as shown above. The 8-step,
+ComfyUI, Ref2VA, and v1.1 artifacts are not supported. This integration is
+dynamic-only and does not support prefusion, DLO, or LoRA composition.
+It also rejects model-level, layerwise, and distributed layerwise offload;
+the legacy dynamic LoRA tensors do not participate in those weight lifecycles.
+The five requested sigma points produce the four denoiser evaluations expected
+by the Turbo artifact.
+
 ## Key parameters
 
 | Parameter | Recommended value | Notes |
-|-----------|-------------------|-------|
+| ----------- | ------------------- | ------- |
 | `quality` | omitted or `lossless` | Request-level quality intent; `high` dynamically installs H3's conservative Cache-DiT profile |
 | `extra_params.force_refresh_step_hint` | omitted | Optional positive 1-based denoising-step hint for an active Cache-DiT request; pair with `extra_params.force_refresh_step_policy`=`once` or `repeat` |
 | `task` | `t2va`, `fl2va`, or `ref2va` | Passed in `extra_params`; selects the task-specific DiT |
@@ -787,9 +887,9 @@ Users can also use a ComfyUI frontend to interact with a hosted MiniMax-H3 servi
 The four-GPU recommendation was measured on four NVIDIA B300 GPUs with one
 excluded warmup followed by three requests.
 
-| Workload | Configuration | Observed result |
-|----------|---------------|-----------------|
-| FL2VA, 209 frames, 1248x768 | no offload, U4, VPP4 tile, regional compile | 86.964 s mean HTTP client latency |
+| Workload                               | Configuration                               | Observed result                      |
+| -------------------------------------- | ------------------------------------------- | ------------------------------------ |
+| FL2VA, 209 frames, 1248x768            | no offload, U4, VPP4 tile, regional compile | 86.964 s mean HTTP client latency    |
 | Two-video Ref2VA, 362 frames, 1344x768 | no offload, U4, VPP4 tile, regional compile | 784.394 s accounted model-stage mean |
 
 These measurements describe the validated shapes rather than a general
@@ -881,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/config/test_config_factory.py

@@ -1304,6 +1304,7 @@ stages:
             ({"model_type": "step_audio_2"}, None, "step_audio_2"),
             (None, {"_class_name": "HunyuanVideo15Pipeline"}, "hunyuan_video_15"),
             (None, {"_class_name": "WanPipeline"}, "wan2_2_ti2v"),
+            (None, {"_class_name": "WanDMDPipeline"}, "wan2_2_ti2v"),
         ],
     )
     def test_migrated_models_are_discovered_without_explicit_deploy(

tests/config/test_omni_config.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Tests for the additive structured Omni config."""
 
 from __future__ import annotations
@@ -1152,6 +1152,35 @@ def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_ar
     assert stage.diffusion_config.diffusion_kv_max_rows_per_request == 2
 
 
+def test_from_pipeline_config_forwards_fastvideo_vsa_topk(tmp_path, monkeypatch):
+    from vllm_omni.platforms import current_omni_platform
+
+    monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True)
+    deploy_path = tmp_path / "dreamzero_fastvideo_vsa.yaml"
+    deploy_path.write_text(
+        "\n".join(
+            [
+                "pipeline: dreamzero",
+                "async_chunk: false",
+                "stages:",
+                "  - stage_id: 0",
+                "    diffusion_attention_backend: FASTVIDEO_VSA",
+                "    fastvideo_vsa_topk: 96",
+            ]
+        )
+    )
+
+    stage = _from_pipeline_key(
+        "dreamzero",
+        deploy_config_path=str(deploy_path),
+    ).stage_by_id(0)
+
+    attention_config = stage.diffusion_config.diffusion_attention_config
+    assert attention_config.default is not None
+    assert attention_config.default.backend == "FASTVIDEO_VSA"
+    assert attention_config.default.fastvideo_vsa_topk == 96
+
+
 def test_from_pipeline_config_rejects_reserved_diffusion_kv_mode(tmp_path):
     deploy_path = tmp_path / "dreamzero_reserved_diffusion_kv.yaml"
     deploy_path.write_text(

tests/core/sched/test_omni_ar_scheduler_aborted_queue_sweep.py

@@ -1,89 +0,0 @@
-"""Tests for ``OmniARScheduler._drop_aborted_queued_requests``.
-
-``schedule()`` sweeps ``FINISHED_ABORTED`` requests out of the queues before
-handing control to upstream, because upstream ``Scheduler.schedule()`` raises
-``RuntimeError: Invalid request status: FINISHED_ABORTED`` for any request it
-admits in a finished state -- which kills the stage's engine core, not just
-the request.
-
-Regression for the sweep missing ``skipped_waiting``, the third queue
-upstream admits from. An aborted duplex session parked there was re-selected
-by ``_select_waiting_queue_for_scheduling`` on a later tick and crashed the
-stage.
-"""
-
-from __future__ import annotations
-
-import pytest
-from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue
-from vllm.v1.request import RequestStatus
-
-import vllm_omni.core.sched.omni_ar_scheduler as ar_sched_mod
-from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-_POLICIES = [
-    pytest.param(SchedulingPolicy.FCFS, id="fcfs"),
-    pytest.param(SchedulingPolicy.PRIORITY, id="priority"),
-]
-
-
-class _StubRequest:
-    """Minimal ``Request`` stub with the surface the sweep exercises.
-
-    Carries the ``(priority, arrival_time)`` ordering ``PriorityRequestQueue``
-    relies on so the same stub works under both scheduling policies.
-    """
-
-    def __init__(self, request_id: str, status: RequestStatus, arrival_time: float = 0.0) -> None:
-        self.request_id = request_id
-        self.status = status
-        self.priority = 0
-        self.arrival_time = arrival_time
-
-    def __lt__(self, other: _StubRequest) -> bool:
-        return (self.priority, self.arrival_time) < (other.priority, other.arrival_time)
-
-
-def _make_scheduler(policy: SchedulingPolicy, *, waiting=(), skipped_waiting=(), running=()):
-    """Bare scheduler with just the surface the sweep reads."""
-    scheduler = OmniARScheduler.__new__(OmniARScheduler)
-    scheduler.waiting = create_request_queue(policy)
-    for req in waiting:
-        scheduler.waiting.add_request(req)
-    scheduler.skipped_waiting = create_request_queue(policy)
-    for req in skipped_waiting:
-        scheduler.skipped_waiting.add_request(req)
-    scheduler.running = list(running)
-    return scheduler
-
-
-@pytest.mark.parametrize("policy", _POLICIES)
-def test_schedule_sweeps_skipped_waiting_before_upstream_selection(
-    monkeypatch: pytest.MonkeyPatch,
-    policy: SchedulingPolicy,
-) -> None:
-    """The public schedule path must remove aborted skipped requests first."""
-    aborted = _StubRequest("req-aborted", RequestStatus.FINISHED_ABORTED)
-    scheduler = _make_scheduler(policy, skipped_waiting=[aborted])
-    scheduler.chunk_transfer_adapter = None
-    scheduler.input_coordinator = None
-    scheduler.num_waiting_for_streaming_input = 0
-    scheduler._process_pending_omni_inputs = lambda model_mode: None
-    scheduler._resync_streaming_input_counter = lambda: None
-    scheduler._should_defer_waiting_admission = lambda: False
-    scheduler._restore_omni_wait_queues = lambda: None
-    scheduler._postprocess_omni_schedule_output = lambda *args, **kwargs: None
-    scheduler.get_finished_requests_needing_kv_transfer = lambda: {}
-    scheduler._wrap_omni_scheduler_output = lambda output, **kwargs: output
-
-    sentinel = object()
-
-    def upstream_schedule(_self, _throttle_prefills: bool = False):
-        assert list(scheduler.skipped_waiting) == []
-        return sentinel
-
-    monkeypatch.setattr(ar_sched_mod.VLLMScheduler, "schedule", upstream_schedule)
-
-    assert OmniARScheduler.schedule(scheduler) is sentinel

tests/core/sched/test_omni_ar_scheduler_streaming.py

@@ -28,9 +28,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 def _make_scheduler(*, stage_id: int = 0) -> OmniARScheduler:
     sched = OmniARScheduler.__new__(OmniARScheduler)
     sched._new_prompt_len_snapshot = {}
-    sched.vllm_config = SimpleNamespace(
-        model_config=SimpleNamespace(stage_id=stage_id, max_model_len=None),
-    )
+    sched.vllm_config = SimpleNamespace(model_config=SimpleNamespace(stage_id=stage_id))
     sched.num_waiting_for_streaming_input = 0
     sched.log_stats = False
     sched.chunk_transfer_adapter = None
@@ -212,81 +210,6 @@ def test_running_decode_step_without_inter_stage_payload_does_not_raise() -> Non
     sched.chunk_transfer_adapter.save_async.assert_not_called()
 
 
-def test_stale_async_frame_is_dropped_before_output_processing() -> None:
-    session = _make_request()
-    session.status = RequestStatus.RUNNING
-    session.num_in_flight_tokens = 2
-    session.num_computed_tokens = session.num_tokens + 1
-    session.num_output_placeholders = 1
-    session.async_tokens_to_discard = 1
-    session.sampling_params = SimpleNamespace(num_logprobs=1)
-    num_computed_tokens = session.num_computed_tokens
-    num_output_placeholders = session.num_output_placeholders
-
-    sched = MagicMock()
-    sched.requests = {session.request_id: session}
-    sched.perf_metrics = None
-    sched.structured_output_manager.should_advance.return_value = False
-
-    def discard_stale_output(request: Request, token_ids: list[int]) -> tuple[list[int], bool]:
-        request.async_tokens_to_discard = 0
-        return token_ids, False
-
-    sched._update_request_with_output.side_effect = discard_stale_output
-    sched._process_kv_transfer_trigger.return_value = False
-    sched.chunk_transfer_adapter = MagicMock()
-    sched.running = [session]
-    sched.waiting_for_transfer_free = set()
-    sched.transfer_triggered_requests = set()
-    sched.active_kv_transfers = set()
-    sched.pending_stop_after_extraction = set()
-    sched.connector = None
-    sched.kv_cache_manager.take_events.return_value = None
-    sched.finished_req_ids_dict = {}
-    sched.make_stats.return_value = None
-
-    scheduler_output = MagicMock(spec=SchedulerOutput)
-    scheduler_output.num_scheduled_tokens = {session.request_id: 1}
-    scheduler_output.scheduled_spec_decode_tokens = {}
-    scheduler_output.num_invalid_spec_tokens = 0
-
-    model_runner_output = MagicMock(spec=ModelRunnerOutput)
-    model_runner_output.sampled_token_ids = [[42]]
-    model_runner_output.logprobs = None
-    model_runner_output.prompt_logprobs_dict = {}
-    model_runner_output.pooler_output = None
-    model_runner_output.multimodal_outputs = None
-    model_runner_output.inter_stage_outputs = [{"hidden": object()}]
-    model_runner_output.num_nans_in_logits = None
-    model_runner_output.kv_connector_output = None
-    model_runner_output.cudagraph_stats = None
-    model_runner_output.req_id_to_index = {session.request_id: 0}
-    model_runner_output.routed_experts = None
-
-    OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output)
-
-    assert session.async_tokens_to_discard == 0
-    assert session.status == RequestStatus.RUNNING
-    assert session.num_computed_tokens == num_computed_tokens
-    assert session.num_output_placeholders == num_output_placeholders
-    sched.chunk_transfer_adapter.save_async.assert_not_called()
-
-    session.sampling_params = SamplingParams(max_tokens=8)
-    next_payload = {"hidden": object()}
-    sched._update_request_with_output.side_effect = None
-    sched._update_request_with_output.return_value = ([43], False)
-    model_runner_output.sampled_token_ids = [[43]]
-    model_runner_output.inter_stage_outputs = [next_payload]
-
-    OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output)
-
-    sched.chunk_transfer_adapter.save_async.assert_called_once_with(
-        next_payload,
-        session,
-        False,
-    )
-
-
 def test_stage0_streaming_update_discards_outstanding_async_placeholder_token() -> None:
     sched = _make_scheduler(stage_id=0)
     session = _make_request()
@@ -431,34 +354,3 @@ def test_ready_async_chunk_prompt_replacement_releases_stale_kv_once() -> None:
     assert sched.chunk_transfer_adapter.replaced_streaming_prompt_ids == set()
     assert sched.chunk_transfer_adapter.requests_with_ready_chunks == {session.request_id}
     assert sched.chunk_transfer_adapter.requests_num_chunks_sent == {}
-
-
-def test_native_duplex_prompt_reset_prevents_stage1_context_overflow() -> None:
-    sched = _make_scheduler(stage_id=1)
-    sched.vllm_config.model_config.max_model_len = 4096
-    sched.chunk_transfer_adapter = SimpleNamespace(
-        receives_chunks=False,
-        segment_finished_requests=set(),
-    )
-    session = _make_request()
-    session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
-    session.resumable = True
-    session.prompt_token_ids = [0] * 3
-    session._all_token_ids.clear()
-    session._all_token_ids.extend([0] * 4095)
-    session._output_token_ids.extend([1] * 4080)
-    session.num_prompt_tokens = 3
-    session.num_computed_tokens = 4095
-
-    update = _make_update([0] * 16)
-    update.model_intermediate_buffer = {"native_duplex": True}
-
-    sched._update_request_as_session(session, update)
-
-    assert session.prompt_token_ids == [0] * 16
-    assert list(session._all_token_ids) == [0] * 16
-    assert session._output_token_ids == []
-    assert session.num_prompt_tokens == 16
-    assert session.num_computed_tokens == 0
-    assert session.model_intermediate_buffer == update.model_intermediate_buffer
-    assert session.status == RequestStatus.WAITING

tests/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py

@@ -76,9 +76,6 @@ def _make_scheduler(
         else:
             hf_config.voxcpm2_runtime_config = runtime_config
     sched.vllm_config = SimpleNamespace(model_config=model_config)
-    # Upstream ``Scheduler.__init__`` always creates ``skipped_waiting``; the
-    # abort sweep at the top of ``schedule()`` reads it.
-    sched.skipped_waiting = _MockQueue()
     return sched
 
 

tests/core/sched/test_omni_scheduler_streaming_input_counter.py

@@ -1,151 +0,0 @@
-"""Tests for the ``num_waiting_for_streaming_input`` resync placed in
-``OmniARScheduler.finish_requests`` / ``OmniGenerationScheduler.finish_requests``
-and in ``OmniARScheduler.schedule``.
-
-Upstream keeps that counter incrementally (``+1`` when a resumable request
-parks in a waiting queue as ``WAITING_FOR_STREAMING_REQ``, ``-1`` when the
-next update arrives or the request is finished) and
-``Scheduler.get_num_unfinished_requests`` *subtracts* it from the waiting
-queues. Omni rewrites ``request.status`` outside both hooks -- the
-chunk-transfer adapter's park/restore and
-``_realign_request_status_to_queues`` -- so a counted request can leave
-``WAITING_FOR_STREAMING_REQ`` without the matching decrement and inflate the
-counter for the lifetime of the process.
-
-An inflated counter does not degrade throughput, it hangs the stage:
-``EngineCore.has_work()`` reads false while a live request sits in
-``waiting``, so the engine blocks in ``input_queue.get()`` and never calls
-``schedule()`` again. Regression for the Qwen3-Omni duplex symptom where the
-server answered the first conversation after a boot and was silent from the
-second on while still reporting healthy.
-"""
-
-from __future__ import annotations
-
-import pytest
-from vllm.v1.request import RequestStatus
-
-import vllm_omni.core.sched.omni_ar_scheduler as ar_sched_mod
-import vllm_omni.core.sched.omni_generation_scheduler as gen_sched_mod
-from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
-from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-class _StubRequest:
-    def __init__(self, request_id: str, status: RequestStatus) -> None:
-        self.request_id = request_id
-        self.status = status
-
-    def is_finished(self) -> bool:
-        return RequestStatus.is_finished(self.status)
-
-
-_SCHEDULER_PARAMS = [
-    pytest.param(OmniARScheduler, ar_sched_mod, id="ar"),
-    pytest.param(OmniGenerationScheduler, gen_sched_mod, id="generation"),
-]
-
-
-def _make_scheduler(scheduler_cls, *, requests, running, waiting, counter, skipped=None):
-    scheduler = scheduler_cls.__new__(scheduler_cls)
-    scheduler.chunk_transfer_adapter = None
-    scheduler.input_coordinator = None
-    scheduler.requests = requests
-    scheduler.running = running
-    scheduler.waiting = waiting
-    scheduler.skipped_waiting = skipped if skipped is not None else []
-    scheduler.num_waiting_for_streaming_input = counter
-    return scheduler
-
-
-@pytest.mark.parametrize(("scheduler_cls", "scheduler_mod"), _SCHEDULER_PARAMS)
-def test_finish_requests_clears_counter_left_by_a_stomped_status(
-    monkeypatch: pytest.MonkeyPatch,
-    scheduler_cls,
-    scheduler_mod,
-) -> None:
-    """The abort case that wedged the talker.
-
-    The request was counted while parked for streaming input, then had its
-    status rewritten to ``RUNNING`` by omni's queue surgery, so upstream's
-    ``finish_requests`` took the running branch and never decremented. With
-    no parked requests left, the counter must land at 0.
-    """
-    aborted = _StubRequest("req-aborted", RequestStatus.RUNNING)
-    scheduler = _make_scheduler(
-        scheduler_cls,
-        requests={},  # upstream already dropped it
-        running=[],
-        waiting=[],
-        counter=1,
-    )
-
-    monkeypatch.setattr(
-        scheduler_mod.VLLMScheduler,
-        "finish_requests",
-        lambda self, request_ids, finished_status: [aborted],
-    )
-
-    scheduler_cls.finish_requests(scheduler, [aborted.request_id], RequestStatus.FINISHED_ABORTED)
-
-    assert scheduler.num_waiting_for_streaming_input == 0
-
-
-def test_resync_counts_parked_requests_in_skipped_waiting() -> None:
-    """``get_num_unfinished_requests`` sums ``waiting`` and ``skipped_waiting``,
-    so the resync has to look at both or it under-counts and reintroduces the
-    spin the counter exists to prevent."""
-    parked = _StubRequest("req-skipped", RequestStatus.WAITING_FOR_STREAMING_REQ)
-    scheduler = _make_scheduler(
-        OmniARScheduler,
-        requests={parked.request_id: parked},
-        running=[],
-        waiting=[],
-        counter=0,
-        skipped=[parked],
-    )
-
-    scheduler._resync_streaming_input_counter()
-
-    assert scheduler.num_waiting_for_streaming_input == 1
-
-
-def test_ar_schedule_resyncs_before_delegating_upstream(monkeypatch: pytest.MonkeyPatch) -> None:
-    """Pins placement in ``schedule``: the chunk-transfer adapter rewrites
-    status during ``process_pending_chunks``, so the resync has to run after
-    that and before upstream schedules, while ``self.waiting`` is still the
-    real queue."""
-    parked = _StubRequest("req-resumed", RequestStatus.RUNNING)
-    scheduler = OmniARScheduler.__new__(OmniARScheduler)
-    scheduler.chunk_transfer_adapter = None
-    scheduler.input_coordinator = None
-    scheduler.requests = {parked.request_id: parked}
-    scheduler.running = [parked]
-    scheduler.waiting = []
-    scheduler.skipped_waiting = []
-    scheduler.num_waiting_for_streaming_input = 1  # leaked by the status rewrite
-    scheduler.policy = "fcfs"
-
-    seen: dict[str, int] = {}
-
-    monkeypatch.setattr(
-        ar_sched_mod.OmniSchedulerMixin, "_consume_pending_connector_output", lambda self, model_mode: None
-    )
-    monkeypatch.setattr(ar_sched_mod.OmniSchedulerMixin, "_process_pending_input_timeouts", lambda self: None)
-    monkeypatch.setattr(OmniARScheduler, "_should_defer_waiting_admission", lambda self: False)
-
-    def fake_schedule(self, throttle_prefills=False):
-        seen["counter"] = self.num_waiting_for_streaming_input
-        raise _StopScheduleError
-
-    class _StopScheduleError(Exception):
-        pass
-
-    monkeypatch.setattr(ar_sched_mod.VLLMScheduler, "schedule", fake_schedule)
-
-    with pytest.raises(_StopScheduleError):
-        OmniARScheduler.schedule(scheduler)
-
-    assert seen["counter"] == 0

tests/diffusion/attention/test_attention_config.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """Tests for per-role attention backend configuration (RFC: per-role-attention-backend).
 
@@ -90,6 +90,14 @@ class TestAttentionSpec:
         with pytest.raises(ValueError, match=match):
             AttentionSpec(**spec)
 
+    def test_fastvideo_vsa_topk_serialized(self):
+        spec = AttentionSpec(backend="FASTVIDEO_VSA", fastvideo_vsa_topk=96)
+        assert spec.backend_kwargs() == {"topk": 96}
+
+    def test_fastvideo_vsa_topk_rejected_for_other_backend(self):
+        with pytest.raises(ValueError, match="only supported by the FASTVIDEO_VSA"):
+            AttentionSpec(backend="TORCH_SDPA", fastvideo_vsa_topk=96)
+
     def test_block_sparse_defaults_applied_when_backend_selected(self):
         spec = AttentionSpec(backend="RAINFUSION_ATTN")
         assert spec.block_sparse.sparsity == 0.8

tests/diffusion/attention/test_fastvideo_vsa.py

@@ -0,0 +1,203 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import sys
+import types
+from typing import Any
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.backends.fastvideo_vsa import (
+    FastVideoVSAImpl,
+)
+from vllm_omni.diffusion.attention.backends.registry import (
+    DiffusionAttentionBackendEnum,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def test_fastvideo_vsa_backend_is_registered():
+    assert DiffusionAttentionBackendEnum.FASTVIDEO_VSA.get_path().endswith("fastvideo_vsa.FastVideoVSABackend")
+
+
+def test_fastvideo_vsa_tiles_3d_sequence_and_untiles(monkeypatch):
+    calls: dict[str, Any] = {}
+    fake_module = types.ModuleType("fastvideo_kernel")
+
+    def fake_video_sparse_attn_bshd(
+        q,
+        k,
+        v,
+        variable_block_sizes,
+        q_variable_block_sizes,
+        compress_attn_weight,
+        topk,
+        block_size,
+    ):
+        calls["q_shape"] = tuple(q.shape)
+        calls["vbs"] = variable_block_sizes.detach().cpu().tolist()
+        calls["q_vbs"] = q_variable_block_sizes.detach().cpu().tolist()
+        calls["compress_sum"] = float(compress_attn_weight.detach().cpu().sum())
+        calls["compress_shape"] = tuple(compress_attn_weight.shape)
+        calls["topk"] = topk
+        calls["block_size"] = block_size
+        return q + k + v
+
+    setattr(fake_module, "video_sparse_attn_bshd", fake_video_sparse_attn_bshd)
+    monkeypatch.setitem(sys.modules, "fastvideo_kernel", fake_module)
+
+    impl = FastVideoVSAImpl(
+        num_heads=2,
+        head_size=8,
+        softmax_scale=8**-0.5,
+        causal=False,
+        backend_kwargs={
+            "topk": 1,
+            "block_size": (4, 8, 8),
+            "min_seq_len": 1,
+            "disable_when_sp_active": False,
+        },
+    )
+    query = torch.randn(1, 300, 2, 8)
+    key = torch.randn_like(query)
+    value = torch.randn_like(query)
+    attn_metadata = AttentionMetadata(extra={"vsa_dit_seq_shape": (3, 10, 10)})
+    monkeypatch.setattr(impl, "_fallback_reason", lambda *args, **kwargs: None)
+
+    output = impl.forward_cuda(query, key, value, attn_metadata)
+
+    assert output.shape == query.shape
+    assert calls["q_shape"] == (1, 1024, 2, 8)
+    assert calls["vbs"] == [192, 48, 48, 12]
+    assert calls["q_vbs"] == [192, 48, 48, 12]
+    assert calls["compress_shape"] == (1, 1024, 2, 8)
+    assert calls["compress_sum"] == 0.0
+    assert calls["topk"] == 1
+    assert calls["block_size"] == (4, 8, 8)
+
+
+def test_fastvideo_vsa_uses_learned_gate_when_provided(monkeypatch):
+    calls: dict[str, Any] = {}
+    fake_module = types.ModuleType("fastvideo_kernel")
+
+    def fake_video_sparse_attn_bshd(
+        q,
+        k,
+        v,
+        variable_block_sizes,
+        q_variable_block_sizes,
+        compress_attn_weight,
+        topk,
+        block_size,
+    ):
+        calls["compress_sum"] = float(compress_attn_weight.detach().cpu().sum())
+        calls["compress_shape"] = tuple(compress_attn_weight.shape)
+        return q + k + v
+
+    setattr(fake_module, "video_sparse_attn_bshd", fake_video_sparse_attn_bshd)
+    monkeypatch.setitem(sys.modules, "fastvideo_kernel", fake_module)
+
+    impl = FastVideoVSAImpl(
+        num_heads=2,
+        head_size=8,
+        softmax_scale=8**-0.5,
+        causal=False,
+        backend_kwargs={
+            "topk": 1,
+            "block_size": (4, 8, 8),
+            "min_seq_len": 1,
+            "disable_when_sp_active": False,
+        },
+    )
+    query = torch.randn(1, 300, 2, 8)
+    gate = torch.ones_like(query)
+    metadata = AttentionMetadata(extra={"vsa_dit_seq_shape": (3, 10, 10), "gate_compress": gate})
+    monkeypatch.setattr(impl, "_fallback_reason", lambda *args, **kwargs: None)
+
+    output = impl.forward_cuda(query, query, query, metadata)
+
+    assert output.shape == query.shape
+    assert calls["compress_shape"] == (1, 1024, 2, 8)
+    # Only the 300 real tokens carry ones; padded gate slots stay zero.
+    assert calls["compress_sum"] == gate.numel()
+
+
+def test_fastvideo_vsa_falls_back_without_dit_shape(monkeypatch):
+    calls: dict[str, Any] = {}
+
+    impl = FastVideoVSAImpl(
+        num_heads=2,
+        head_size=8,
+        softmax_scale=8**-0.5,
+        causal=False,
+        backend_kwargs={
+            "topk": 1,
+            "block_size": (4, 8, 8),
+            "min_seq_len": 1,
+            "disable_when_sp_active": False,
+        },
+    )
+
+    def fake_fallback(query, key, value, attn_metadata, reason):
+        calls["reason"] = reason
+        return torch.zeros_like(query)
+
+    monkeypatch.setattr(impl, "_fallback", fake_fallback)
+
+    query = torch.randn(1, 512, 2, 8)
+    output = impl.forward_cuda(query, query, query, AttentionMetadata())
+
+    assert output.shape == query.shape
+    assert calls["reason"] == "vsa_dit_seq_shape metadata is required"
+
+
+def test_fastvideo_vsa_falls_back_for_mask(monkeypatch):
+    calls: dict[str, Any] = {}
+
+    impl = FastVideoVSAImpl(
+        num_heads=2,
+        head_size=8,
+        softmax_scale=8**-0.5,
+        causal=False,
+        backend_kwargs={"min_seq_len": 1},
+    )
+
+    def fake_fallback(query, key, value, attn_metadata, reason):
+        calls["reason"] = reason
+        return torch.zeros_like(query)
+
+    monkeypatch.setattr(impl, "_fallback", fake_fallback)
+
+    query = torch.randn(1, 512, 2, 8)
+    mask = torch.ones(1, 512, dtype=torch.bool)
+    output = impl.forward_cuda(query, query, query, AttentionMetadata(attn_mask=mask))
+
+    assert output.shape == query.shape
+    assert calls["reason"]
+
+
+def test_fastvideo_vsa_allows_topk_equal_to_num_blocks():
+    query = torch.randn(1, 512, 2, 8)
+    metadata = AttentionMetadata(extra={"vsa_dit_seq_shape": (4, 8, 16)})
+
+    all_blocks = FastVideoVSAImpl(
+        num_heads=2,
+        head_size=8,
+        softmax_scale=8**-0.5,
+        causal=False,
+        backend_kwargs={"topk": 2, "block_size": (4, 8, 8), "min_seq_len": 1},
+    )
+    too_many = FastVideoVSAImpl(
+        num_heads=2,
+        head_size=8,
+        softmax_scale=8**-0.5,
+        causal=False,
+        backend_kwargs={"topk": 3, "block_size": (4, 8, 8), "min_seq_len": 1},
+    )
+
+    # CPU/float32 is rejected later, but k=N itself must not trigger fallback.
+    assert all_blocks._fallback_reason(query, query, query, metadata) == "dtype torch.float32 is not supported"
+    assert too_many._fallback_reason(query, query, query, metadata) == "topk 3 > num_blocks 2"

tests/diffusion/lora/test_lora_manager.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from __future__ import annotations
 
@@ -20,9 +20,10 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 class _DummyLoRALayer:
-    def __init__(self, n_slices: int, output_slices: tuple[int, ...]):
+    def __init__(self, n_slices: int, output_slices: tuple[int, ...], tp_size: int = 1):
         self.n_slices = n_slices
         self.output_slices = output_slices
+        self.tp_size = tp_size
         self.set_calls: list[
             tuple[list[torch.Tensor | None] | torch.Tensor, list[torch.Tensor | None] | torch.Tensor]
         ] = []
@@ -236,6 +237,14 @@ def test_lora_manager_activates_fused_lora_on_packed_layer():
 
 def test_lora_manager_activates_packed_lora_from_sublayers():
     pipeline = torch.nn.Module()
+    bound_names = None
+
+    def validate_binding(*, lora_model, bound_lora_names):
+        nonlocal bound_names
+        assert lora_model.id == 1
+        bound_names = bound_lora_names
+
+    pipeline._validate_diffusion_lora_binding = validate_binding
     pipeline.stacked_params_mapping = [
         (".to_qkv", ".to_q", "q"),
         (".to_qkv", ".to_k", "k"),
@@ -279,6 +288,70 @@ def test_lora_manager_activates_packed_lora_from_sublayers():
     assert torch.allclose(lora_b_list[0], torch.ones((2, rank)) * 3 * 2.0)
     assert torch.allclose(lora_b_list[1], torch.ones((1, rank)) * 4 * 2.0)
     assert torch.allclose(lora_b_list[2], torch.ones((1, rank)) * 4 * 2.0)
+    assert bound_names == frozenset(loras)
+
+
+def test_lora_manager_rolls_back_all_layers_when_activation_fails():
+    class _FailingLoRALayer(_DummyLoRALayer):
+        fail = False
+
+        def set_lora(self, index: int, lora_a, lora_b):
+            if self.fail:
+                raise RuntimeError("bind failed")
+            super().set_lora(index, lora_a, lora_b)
+
+    manager = DiffusionLoRAManager(
+        pipeline=torch.nn.Module(),
+        device=torch.device("cpu"),
+        dtype=torch.bfloat16,
+        max_cached_adapters=2,
+    )
+    first = _DummyLoRALayer(n_slices=1, output_slices=(2,))
+    second = _FailingLoRALayer(n_slices=1, output_slices=(2,))
+    manager._lora_modules = {
+        "transformer.first": first,
+        "transformer.second": second,
+    }
+
+    def adapter(adapter_id: int):
+        loras = {
+            name: LoRALayerWeights(
+                module_name=name,
+                rank=2,
+                lora_alpha=2,
+                lora_a=torch.full((2, 2), float(adapter_id)),
+                lora_b=torch.full((2, 2), float(adapter_id)),
+            )
+            for name in manager._lora_modules
+        }
+        return type(
+            "LM",
+            (),
+            {
+                "id": adapter_id,
+                "loras": loras,
+                "get_lora": lambda self, key: self.loras.get(key),
+            },
+        )()
+
+    manager._registered_adapters = {1: adapter(1), 2: adapter(2)}
+    manager._activate_adapter(1, scale=1.0)
+    first_calls_after_success = len(first.set_calls)
+
+    second.fail = True
+    with pytest.raises(RuntimeError, match="bind failed"):
+        manager._activate_adapter(2, scale=1.0)
+
+    assert manager._active_adapter_id is None
+    assert first.reset_calls == 1
+    assert second.reset_calls == 1
+
+    # The failed activation must not leave the old adapter eligible for the
+    # fast path: reactivating it binds every layer again.
+    second.fail = False
+    manager._activate_adapter(1, scale=1.0)
+    assert manager._active_adapter_id == 1
+    assert len(first.set_calls) == first_calls_after_success + 2
 
 
 def _dummy_lora_request(adapter_id: int) -> LoRARequest:

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

@@ -1,10 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 import json
 import sys
 from multiprocessing.reduction import ForkingPickler
 from types import SimpleNamespace
+from typing import Any, TypeVar
 from unittest.mock import Mock, patch
 
 import numpy as np
@@ -14,6 +15,14 @@ import torch.nn as nn
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
+_ItemT = TypeVar("_ItemT")
+_ResultT = TypeVar("_ResultT")
+
+
+def _append_and_return(items: list[_ItemT], item: _ItemT, result: _ResultT) -> _ResultT:
+    items.append(item)
+    return result
+
 
 def test_h3_prepares_resolved_cache_state_immediately_before_denoise():
     from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
@@ -185,6 +194,7 @@ def test_startup_task_rejects_unsupported_value():
 
 def test_combined_task_inference_and_transformer_routing():
     from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.errors import OmniClientError
 
     pipeline = object.__new__(MiniMaxH3Pipeline)
     torch.nn.Module.__init__(pipeline)
@@ -196,6 +206,9 @@ def test_combined_task_inference_and_transformer_routing():
     assert pipeline._resolve_task(None, {"image": object()}) == "fl2va"
     assert pipeline._resolve_task(None, {"audio": object()}) == "ref2va"
     assert pipeline._resolve_task(None, {"video": object()}) == "ref2va"
+    assert pipeline._resolve_task("fl2va", {"image": object()}, has_turbo_lora=True) == "fl2va"
+    with pytest.raises(OmniClientError, match="supports T2VA/FL2VA requests only"):
+        pipeline._resolve_task("ref2va", {}, has_turbo_lora=True)
     pipeline.partition = "ref2va"
     pipeline.supported_tasks = frozenset({"ref2va"})
     assert pipeline._resolve_task(None, {"image": object()}) == "ref2va"
@@ -260,7 +273,7 @@ def test_pipeline_loads_task_selected_dits_and_shared_components_once(
         ):
             (tmp_path / partition_name / component).mkdir()
 
-    created = {"dit": [], "text_encoder": [], "video_vae": [], "audio_vae": []}
+    created: dict[str, list[Any]] = {"dit": [], "text_encoder": [], "video_vae": [], "audio_vae": []}
 
     class FakeModule(torch.nn.Module):
         def __init__(self, *args, **kwargs):
@@ -278,15 +291,15 @@ def test_pipeline_loads_task_selected_dits_and_shared_components_once(
 
         return create
 
-    tokenizer_calls = []
-    processor_calls = []
+    tokenizer_calls: list[Any] = []
+    processor_calls: list[Any] = []
     tokenizer_cls = Mock(spec=pipeline_module.Qwen2TokenizerFast)
-    tokenizer_cls.from_pretrained.side_effect = lambda *args, **kwargs: (
-        tokenizer_calls.append((args, kwargs)) or object()
+    tokenizer_cls.from_pretrained.side_effect = lambda *args, **kwargs: _append_and_return(
+        tokenizer_calls, (args, kwargs), object()
     )
     processor_cls = Mock(spec=pipeline_module.Qwen3VLProcessor)
-    processor_cls.from_pretrained.side_effect = lambda *args, **kwargs: (
-        processor_calls.append((args, kwargs)) or object()
+    processor_cls.from_pretrained.side_effect = lambda *args, **kwargs: _append_and_return(
+        processor_calls, (args, kwargs), object()
     )
     monkeypatch.setattr(pipeline_module, "MiniMaxH3DiTModel", FakeDiT)
     monkeypatch.setattr(
@@ -626,7 +639,7 @@ def test_partially_constructed_pipeline_falls_back_to_the_uniform_schedule():
 
 
 def test_distilled_forward_reports_denoising_steps_not_sigma_boundaries():
-    diffuse_calls = []
+    diffuse_calls: list[dict[str, Any]] = []
     base_schedule = [1.0, 0.7, 0.4, 0.15, 0.0]
     pipeline = _distilled_pipeline(diffuse_calls, {"fl2va": base_schedule, "ref2va": None})
 
@@ -643,7 +656,7 @@ def test_distilled_forward_reports_denoising_steps_not_sigma_boundaries():
 
 
 def test_distilled_forward_accepts_the_matching_explicit_step_count():
-    diffuse_calls = []
+    diffuse_calls: list[dict[str, Any]] = []
     pipeline = _distilled_pipeline(diffuse_calls, {"fl2va": [1.0, 0.7, 0.4, 0.15, 0.0], "ref2va": None})
 
     pipeline.forward(_t2va_batch(num_inference_steps=4))
@@ -713,8 +726,10 @@ def test_cudnn_packed_attention_uses_python_length_without_padding_mask():
     )
 
     assert output.shape == q.shape
-    assert attention.attention.metadata.attn_mask is None
-    assert attention.attention.metadata.extra["valid_kv_length"] == 5
+    metadata = attention.attention.metadata
+    assert metadata is not None
+    assert metadata.attn_mask is None
+    assert metadata.extra["valid_kv_length"] == 5
 
 
 def test_packed_attention_skips_mask_for_packed_mask_free_backend():
@@ -754,9 +769,11 @@ def test_packed_attention_skips_mask_for_packed_mask_free_backend():
         packed_total=8,
     )
 
-    assert attention.attention.metadata.attn_mask is None
-    assert attention.attention.metadata.extra["valid_kv_length"] == 5
-    assert attention.attention.metadata.extra["npu_attn_varlen"] is True
+    metadata = attention.attention.metadata
+    assert metadata is not None
+    assert metadata.attn_mask is None
+    assert metadata.extra["valid_kv_length"] == 5
+    assert metadata.extra["npu_attn_varlen"] is True
 
 
 def test_packed_attention_keeps_padding_mask_for_other_backends():
@@ -796,12 +813,90 @@ def test_packed_attention_keeps_padding_mask_for_other_backends():
         packed_total=8,
     )
 
+    metadata = attention.attention.metadata
+    assert metadata is not None
     assert torch.equal(
-        attention.attention.metadata.attn_mask,
+        metadata.attn_mask,
         torch.tensor([[True, True, True, True, True, False, False, False]]),
     )
 
 
+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
 
@@ -1484,12 +1579,12 @@ def test_r7_r8_ref2va_video_segment_matrix(monkeypatch, tmp_path, case, start_ti
         "audio_codecs": (),
         "file_size": 1024,
     }
-    transcode_calls = []
+    transcode_calls: list[Any] = []
     monkeypatch.setattr(reference_video_module, "_probe_video", lambda _path: metadata)
     monkeypatch.setattr(
         reference_video_module,
         "_transcode_reference_video",
-        lambda source, **kwargs: transcode_calls.append((source, kwargs)) or "prepared.mp4",
+        lambda source, **kwargs: _append_and_return(transcode_calls, (source, kwargs), "prepared.mp4"),
     )
 
     prepared = reference_video_module.prepare_reference_videos(
@@ -1510,7 +1605,7 @@ def test_ref2va_qwen_sampling_uses_one_selective_decode(monkeypatch):
     from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
 
     decoded = np.arange(3 * 2 * 2 * 3, dtype=np.uint8).reshape(3, 2, 2, 3)
-    decode_calls = []
+    decode_calls: list[Any] = []
     monkeypatch.setattr(
         reference_video_module,
         "_probe_video",
@@ -1523,7 +1618,7 @@ def test_ref2va_qwen_sampling_uses_one_selective_decode(monkeypatch):
     monkeypatch.setattr(
         reference_video_module,
         "_decode_video_frames_ffmpeg",
-        lambda path, **kwargs: decode_calls.append((path, kwargs)) or decoded,
+        lambda path, **kwargs: _append_and_return(decode_calls, (path, kwargs), decoded),
     )
 
     sampled = reference_video_module.sample_reference_video_frames("prepared.mp4")
@@ -1584,7 +1679,7 @@ def test_ref2va_load_video_frames_uses_ffmpeg_without_decord(monkeypatch):
     from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
 
     expected = np.zeros((2, 2, 3, 3), dtype=np.uint8)
-    decode_calls = []
+    decode_calls: list[Any] = []
     monkeypatch.setitem(sys.modules, "decord", None)
     monkeypatch.setattr(
         reference_video_module,
@@ -1598,7 +1693,7 @@ def test_ref2va_load_video_frames_uses_ffmpeg_without_decord(monkeypatch):
     monkeypatch.setattr(
         reference_video_module,
         "_decode_video_frames_ffmpeg",
-        lambda path, **kwargs: decode_calls.append((path, kwargs)) or expected,
+        lambda path, **kwargs: _append_and_return(decode_calls, (path, kwargs), expected),
     )
 
     actual = reference_video_module.load_video_frames("prepared.mp4")
@@ -1638,7 +1733,7 @@ def test_ref2va_ffmpeg_decode_rejects_invalid_metadata(kwargs, message):
 def test_ref2va_probe_uses_container_frame_count_without_scan(monkeypatch):
     from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
 
-    calls = []
+    calls: list[Any] = []
     probe = {
         "streams": [
             {
@@ -1656,7 +1751,11 @@ def test_ref2va_probe_uses_container_frame_count_without_scan(monkeypatch):
     monkeypatch.setattr(
         reference_video_module.subprocess,
         "run",
-        lambda command, **kwargs: calls.append((command, kwargs)) or SimpleNamespace(stdout=json.dumps(probe)),
+        lambda command, **kwargs: _append_and_return(
+            calls,
+            (command, kwargs),
+            SimpleNamespace(stdout=json.dumps(probe)),
+        ),
     )
 
     metadata = reference_video_module._probe_video("prepared.mp4")
@@ -1669,7 +1768,7 @@ def test_ref2va_probe_uses_container_frame_count_without_scan(monkeypatch):
 def test_ref2va_probe_counts_frames_only_when_metadata_is_missing(monkeypatch):
     from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
 
-    calls = []
+    calls: list[Any] = []
     stream = {
         "codec_type": "video",
         "width": 4,
@@ -1689,8 +1788,11 @@ def test_ref2va_probe_counts_frames_only_when_metadata_is_missing(monkeypatch):
     monkeypatch.setattr(
         reference_video_module.subprocess,
         "run",
-        lambda command, **kwargs: calls.append((command, kwargs))
-        or SimpleNamespace(stdout=json.dumps(first_probe if len(calls) == 1 else second_probe)),
+        lambda command, **kwargs: _append_and_return(
+            calls,
+            (command, kwargs),
+            SimpleNamespace(stdout=json.dumps(first_probe if not calls else second_probe)),
+        ),
     )
 
     metadata = reference_video_module._probe_video("prepared.mp4")
@@ -1750,12 +1852,16 @@ def test_ref2va_two_video_recipe_tolerates_container_rounding(monkeypatch, tmp_p
             "file_size": 1024,
         },
     }
-    transcode_calls = []
+    transcode_calls: list[Any] = []
     monkeypatch.setattr(reference_video_module, "_probe_video", lambda path: metadata[str(path)])
     monkeypatch.setattr(
         reference_video_module,
         "_transcode_reference_video",
-        lambda source, **kwargs: transcode_calls.append((source, kwargs)) or f"{source}.prepared.mp4",
+        lambda source, **kwargs: _append_and_return(
+            transcode_calls,
+            (source, kwargs),
+            f"{source}.prepared.mp4",
+        ),
     )
 
     prepared = reference_video_module.prepare_reference_videos(

tests/diffusion/models/minimax_h3/test_minimax_h3_lora.py

@@ -0,0 +1,381 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from safetensors.torch import save_file
+from vllm.lora.lora_weights import PackedLoRALayerWeights
+
+from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
+from vllm_omni.diffusion.models.minimax_h3 import lora as lora_module
+from vllm_omni.diffusion.models.minimax_h3.lora import (
+    load_minimax_h3_turbo_lora,
+)
+from vllm_omni.errors import OmniClientError
+from vllm_omni.lora.request import LoRARequest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+_TINY_TARGET_DIMS = {
+    "attn.to_q": (2, 3),
+    "attn.to_k": (2, 3),
+    "attn.to_v": (2, 3),
+    "attn.to_out.0": (3, 2),
+    "ff.net.0.proj": (2, 4),
+    "ff.net.2": (2, 2),
+}
+
+
+@pytest.fixture(autouse=True)
+def _use_tiny_h3_dimensions(monkeypatch):
+    monkeypatch.setattr(lora_module, "_TURBO_TARGET_DIMS", _TINY_TARGET_DIMS)
+
+
+def _request(path) -> LoRARequest:
+    return LoRARequest(
+        lora_name="turbo",
+        lora_int_id=1,
+        lora_path=str(path),
+    )
+
+
+def _write_tiny_turbo(
+    path,
+    *,
+    alpha: str = "128",
+    key_format: str = "minimax-h3-diffusers",
+    omit_target: str | None = None,
+    shape_overrides: dict[str, tuple[int, int]] | None = None,
+) -> None:
+    rank = 128
+    tensors = {}
+    overrides = shape_overrides or {}
+    target_suffixes = (
+        "attn.to_q",
+        "attn.to_k",
+        "attn.to_v",
+        "attn.to_out.0",
+        "ff.net.0.proj",
+        "ff.net.2",
+    )
+    for prefix, block_count in (
+        ("transformer_blocks", 50),
+        ("token_refiner.refiner_blocks", 2),
+    ):
+        for block_index in range(block_count):
+            for suffix in target_suffixes:
+                target = f"{prefix}.{block_index}.{suffix}"
+                if target == omit_target:
+                    continue
+                input_dim, output_dim = _TINY_TARGET_DIMS[suffix]
+                a_name = f"{target}.lora_A.default.weight"
+                b_name = f"{target}.lora_B.default.weight"
+                tensors[a_name] = torch.ones(overrides.get(a_name, (rank, input_dim)))
+                if suffix == "ff.net.0.proj":
+                    if b_name in overrides:
+                        tensors[b_name] = torch.ones(overrides[b_name])
+                    else:
+                        tensors[b_name] = torch.cat(
+                            (
+                                torch.ones(output_dim // 2, rank),
+                                torch.full((output_dim // 2, rank), 2.0),
+                            ),
+                            dim=0,
+                        )
+                else:
+                    tensors[b_name] = torch.ones(overrides.get(b_name, (output_dim, rank)))
+    save_file(
+        tensors,
+        str(path),
+        metadata={"alpha": alpha, "key_format": key_format},
+    )
+
+
+def test_h3_turbo_loads_through_legacy_lora_model_and_swaps_ffn(tmp_path):
+    path = tmp_path / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"
+    _write_tiny_turbo(path)
+
+    loaded = load_minimax_h3_turbo_lora(
+        partition="fl2va",
+        lora_request=_request(path),
+        lora_path=path,
+        dtype=torch.float32,
+    )
+
+    assert loaded is not None
+    lora_model, peft_helper = loaded
+    assert peft_helper.r == 128
+    assert peft_helper.lora_alpha == 128
+    assert len(lora_model.loras) == 312
+    assert "blocks.0.attn.to_q" in lora_model.loras
+    assert "blocks.0.mlp.fc1" in lora_model.loras
+    fc1 = lora_model.get_lora("blocks.0.mlp.fc1")
+    assert isinstance(fc1, PackedLoRALayerWeights)
+    assert len(fc1.lora_a) == 2
+    assert len(fc1.lora_b) == 2
+    torch.testing.assert_close(fc1.lora_a[0], torch.ones(128, 2))
+    torch.testing.assert_close(fc1.lora_a[1], torch.ones(128, 2))
+    torch.testing.assert_close(fc1.lora_b[0], torch.full((2, 128), 2.0))
+    torch.testing.assert_close(fc1.lora_b[1], torch.ones(2, 128))
+
+
+@pytest.mark.parametrize(
+    ("tensor_suffix", "bad_shape"),
+    [
+        ("attn.to_q.lora_A.default.weight", (128, 1)),
+        ("attn.to_q.lora_A.default.weight", (128, 3)),
+        ("attn.to_q.lora_B.default.weight", (2, 128)),
+        ("attn.to_q.lora_B.default.weight", (4, 128)),
+    ],
+)
+def test_h3_turbo_rejects_under_and_oversized_global_shapes(tmp_path, tensor_suffix, bad_shape):
+    path = tmp_path / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"
+    tensor_name = f"transformer_blocks.0.{tensor_suffix}"
+    _write_tiny_turbo(path, shape_overrides={tensor_name: bad_shape})
+
+    with pytest.raises(ValueError, match="invalid global shape"):
+        load_minimax_h3_turbo_lora(
+            partition="fl2va",
+            lora_request=_request(path),
+            lora_path=path,
+            dtype=torch.float32,
+        )
+
+
+def test_legacy_manager_uses_the_h3_model_loader_without_changing_its_interface(tmp_path):
+    path = tmp_path / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"
+    _write_tiny_turbo(path)
+
+    class _Pipeline:
+        def _load_diffusion_lora_adapter(self, **kwargs):
+            return load_minimax_h3_turbo_lora(partition="fl2va", **kwargs)
+
+    manager = object.__new__(DiffusionLoRAManager)
+    manager.pipeline = _Pipeline()
+    manager.dtype = torch.float32
+    manager._expected_lora_modules = {"to_q", "fc1"}
+
+    lora_model, peft_helper = manager._load_adapter(_request(path))
+
+    assert lora_model.id == 1
+    assert peft_helper.lora_alpha == 128
+    assert len(lora_model.loras) == 312
+    assert "blocks.0.attn.to_q" in lora_model.loras
+    assert "blocks.0.mlp.fc1" in lora_model.loras
+
+
+def test_h3_turbo_rejects_wrong_alpha_and_ref2va(tmp_path):
+    wrong_alpha = tmp_path / "wrong_alpha" / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"
+    wrong_alpha.parent.mkdir()
+    _write_tiny_turbo(wrong_alpha, alpha="8")
+    with pytest.raises(ValueError, match="requires alpha=128"):
+        load_minimax_h3_turbo_lora(
+            partition="fl2va",
+            lora_request=_request(wrong_alpha),
+            lora_path=wrong_alpha,
+            dtype=torch.float32,
+        )
+
+    valid = tmp_path / "valid" / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"
+    valid.parent.mkdir()
+    _write_tiny_turbo(valid)
+    with pytest.raises(ValueError, match="supports FL2VA/T2VA only"):
+        load_minimax_h3_turbo_lora(
+            partition="ref2va",
+            lora_request=_request(valid),
+            lora_path=valid,
+            dtype=torch.float32,
+        )
+
+
+@pytest.mark.parametrize(
+    "offload_mode",
+    [
+        "model-level CPU offload (--enable-cpu-offload)",
+        "layerwise offload (--enable-layerwise-offload)",
+        "distributed layerwise offload (--enable-distributed-layerwise-offload)",
+    ],
+)
+def test_h3_turbo_rejects_offload_modes(tmp_path, offload_mode):
+    path = tmp_path / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"
+    _write_tiny_turbo(path)
+
+    with pytest.raises(ValueError, match="does not support"):
+        load_minimax_h3_turbo_lora(
+            partition="fl2va",
+            lora_request=_request(path),
+            lora_path=path,
+            dtype=torch.float32,
+            unsupported_offload_mode=offload_mode,
+        )
+
+
+def test_h3_turbo_accepts_only_the_declared_v1_artifact(tmp_path):
+    unsupported = tmp_path / "minimax_h3_fl2v_turbo_8step_v1.0_bf16.safetensors"
+    _write_tiny_turbo(unsupported)
+
+    with pytest.raises(ValueError, match="supports only"):
+        load_minimax_h3_turbo_lora(
+            partition="fl2va",
+            lora_request=_request(unsupported),
+            lora_path=unsupported,
+            dtype=torch.float32,
+        )
+
+    supported = tmp_path / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"
+    _write_tiny_turbo(supported)
+    # Directory resolution selects the declared artifact even when another
+    # v1.0 checkpoint is present beside it.
+    assert (
+        load_minimax_h3_turbo_lora(
+            partition="fl2va",
+            lora_request=_request(tmp_path),
+            lora_path=tmp_path,
+            dtype=torch.float32,
+        )
+        is not None
+    )
+
+
+def test_h3_turbo_rejects_a_truncated_declared_artifact(tmp_path):
+    path = tmp_path / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"
+    _write_tiny_turbo(path, omit_target="transformer_blocks.49.ff.net.2")
+
+    with pytest.raises(ValueError, match="target set does not match"):
+        load_minimax_h3_turbo_lora(
+            partition="fl2va",
+            lora_request=_request(path),
+            lora_path=path,
+            dtype=torch.float32,
+        )
+
+
+def test_non_h3_checkpoint_falls_back_to_the_generic_peft_loader(tmp_path):
+    path = tmp_path / "other.safetensors"
+    _write_tiny_turbo(path, key_format="other")
+
+    assert (
+        load_minimax_h3_turbo_lora(
+            partition="fl2va",
+            lora_request=_request(path),
+            lora_path=path,
+            dtype=torch.float32,
+        )
+        is None
+    )
+
+    declared = tmp_path / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"
+    _write_tiny_turbo(declared, key_format="other")
+    with pytest.raises(ValueError, match="requires safetensors metadata"):
+        load_minimax_h3_turbo_lora(
+            partition="fl2va",
+            lora_request=_request(declared),
+            lora_path=declared,
+            dtype=torch.float32,
+        )
+
+
+def test_only_an_active_recognized_turbo_adapter_restricts_ref2va(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.diffusion.models.minimax_h3 import (
+        pipeline_minimax_h3 as pipeline_module,
+    )
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.partition = "combined"
+    pipeline.supported_tasks = frozenset({"t2va", "fl2va", "ref2va"})
+    pipeline._turbo_lora_adapter_ids = set()
+    request = _request("generic-peft")
+
+    def load():
+        return pipeline._load_diffusion_lora_adapter(
+            lora_request=request,
+            lora_path=request.lora_path,
+            dtype=torch.float32,
+        )
+
+    def resolve(scale):
+        sampling = SimpleNamespace(lora_request=request, lora_scale=scale)
+        return pipeline._resolve_task(
+            "ref2va",
+            {},
+            has_turbo_lora=pipeline._has_active_turbo_lora(sampling),
+        )
+
+    recognized = (object(), object())
+    monkeypatch.setattr(pipeline_module, "load_minimax_h3_turbo_lora", lambda **_: recognized)
+    assert load() is recognized
+    assert resolve(0.0) == "ref2va"
+    with pytest.raises(OmniClientError, match="supports T2VA/FL2VA requests only"):
+        resolve(1.0)
+
+    # Simulate manager eviction followed by a generic PEFT adapter reusing the
+    # same client-supplied ID. A real reload must replace the stale kind.
+    monkeypatch.setattr(pipeline_module, "load_minimax_h3_turbo_lora", lambda **_: None)
+    assert load() is None
+    assert resolve(1.0) == "ref2va"
+
+
+def test_h3_turbo_requires_all_loaded_targets_to_bind():
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline._turbo_lora_adapter_ids = {1}
+    lora_model = SimpleNamespace(id=1, loras={"to_q": object(), "to_k": object()})
+
+    pipeline._validate_diffusion_lora_binding(
+        lora_model=lora_model,
+        bound_lora_names=frozenset(lora_model.loras),
+    )
+    with pytest.raises(ValueError, match="bound=1/2"):
+        pipeline._validate_diffusion_lora_binding(
+            lora_model=lora_model,
+            bound_lora_names=frozenset({"to_q"}),
+        )
+
+
+@pytest.mark.parametrize(
+    ("num_inference_steps", "extra_args", "error"),
+    [
+        (None, {"flow_shift": 6.0, "audio_flow_shift": 3.0}, "num_inference_steps=5"),
+        (4, {"flow_shift": 6.0, "audio_flow_shift": 3.0}, "num_inference_steps=5"),
+        ("5", {"flow_shift": 6.0, "audio_flow_shift": 3.0}, "num_inference_steps=5"),
+        (5, {"flow_shift": 7.0, "audio_flow_shift": 3.0}, "flow_shift=6"),
+        (5, {"flow_shift": "bad", "audio_flow_shift": 3.0}, "flow_shift=6"),
+        (5, {"flow_shift": 6.0, "audio_flow_shift": 4.0}, "audio_flow_shift=3"),
+        (5, {"flow_shift": 6.0, "audio_flow_shift": []}, "audio_flow_shift=3"),
+    ],
+)
+def test_h3_turbo_rejects_unsupported_sampling(num_inference_steps, extra_args, error):
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    pipeline.default_video_shift = 12.0
+    pipeline.default_audio_shift = 3.0
+    sampling = SimpleNamespace(
+        num_inference_steps=num_inference_steps,
+        extra_args=extra_args,
+    )
+
+    with pytest.raises(OmniClientError, match=error):
+        pipeline._validate_turbo_sampling(sampling)
+
+
+def test_h3_turbo_accepts_five_sigma_points_for_four_nfe():
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    pipeline.default_video_shift = 12.0
+    pipeline.default_audio_shift = 3.0
+    pipeline._validate_turbo_sampling(
+        SimpleNamespace(
+            num_inference_steps=5,
+            extra_args={"flow_shift": 6.0, "audio_flow_shift": 3.0},
+        )
+    )

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/models/wan2_2/test_wan22_pipeline_diffuse.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
+import importlib
 from contextlib import contextmanager
 from types import SimpleNamespace
 
@@ -9,6 +10,7 @@ import torch
 from torch import nn
 
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import Wan22Pipeline
+from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanSelfAttention
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -82,6 +84,7 @@ def _make_pipeline() -> Wan22Pipeline:
     pipeline.vae_scale_factor_spatial = 8
     pipeline.boundary_ratio = 0.875
     pipeline.expand_timesteps = False
+    pipeline.is_dmd = False
     pipeline._guidance_scale = None
     pipeline._guidance_scale_2 = None
     pipeline._num_timesteps = None
@@ -94,7 +97,7 @@ def _make_pipeline() -> Wan22Pipeline:
 
 
 def _make_sampling(**overrides):
-    values = {
+    values: dict[str, object] = {
         "height": None,
         "width": None,
         "num_frames": 1,
@@ -338,3 +341,125 @@ def test_diffuse_runs_prediction_and_scheduler_for_each_timestep() -> None:
         (3.0, 3, 28.0, False),
     ]
     assert torch.equal(result, torch.full_like(latents, 10.0))
+
+
+class _StubDMDScheduler:
+    def __init__(self) -> None:
+        self.predict_clean_calls: list[tuple[float, float, float]] = []
+        self.add_noise_calls: list[tuple[float, float, float]] = []
+
+    def predict_clean(self, model_output, sample, timestep):
+        self.predict_clean_calls.append((float(model_output.mean()), float(sample.mean()), float(timestep)))
+        return sample - model_output
+
+    def add_noise(self, clean_sample, noise, timestep):
+        self.add_noise_calls.append((float(clean_sample.mean()), float(noise.mean()), float(timestep)))
+        return clean_sample + 10.0
+
+
+def test_diffuse_dmd_predicts_clean_and_renoises_between_steps(monkeypatch) -> None:
+    pipeline = _make_pipeline()
+    pipeline.is_dmd = True
+    pipeline.scheduler = _StubDMDScheduler()
+    latents = torch.zeros((1, 1, 1, 1, 1), dtype=torch.float32)
+    timesteps = torch.tensor([1000.0, 757.0, 522.0])
+
+    pipeline.predict_noise_maybe_with_cfg = lambda **kwargs: torch.ones_like(latents)  # type: ignore[method-assign]
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2.randn_tensor",
+        lambda *args, **kwargs: torch.full(args[0], 2.0, dtype=kwargs["dtype"]),
+    )
+
+    result = pipeline.diffuse(
+        latents=latents,
+        timesteps=timesteps,
+        prompt_embeds=torch.zeros(1, 8),
+        negative_prompt_embeds=None,
+        guidance_low=1.0,
+        guidance_high=1.0,
+        boundary_timestep=None,
+        dtype=torch.float32,
+        attention_kwargs={},
+        generator=torch.Generator(device="cpu").manual_seed(1),
+    )
+
+    assert pipeline.scheduler.predict_clean_calls == [
+        (1.0, 0.0, 1000.0),
+        (1.0, 9.0, 757.0),
+        (1.0, 18.0, 522.0),
+    ]
+    assert pipeline.scheduler.add_noise_calls == [
+        (-1.0, 2.0, 757.0),
+        (8.0, 2.0, 522.0),
+    ]
+    torch.testing.assert_close(result, torch.tensor([[[[[17.0]]]]]))
+
+
+def _make_gate_loading_pipeline():
+    pipeline = Wan22Pipeline.__new__(Wan22Pipeline)
+    nn.Module.__init__(pipeline)
+    gate = WanSelfAttention.__new__(WanSelfAttention)
+    nn.Module.__init__(gate)
+    gate.to_gate_compress = nn.Linear(1, 1)
+    pipeline.gate_holder = gate
+    return pipeline, gate
+
+
+@pytest.mark.parametrize(
+    ("module_name", "class_name"),
+    [
+        ("pipeline_wan2_2", "Wan22Pipeline"),
+        ("pipeline_wan2_2_i2v", "Wan22I2VPipeline"),
+        ("pipeline_wan2_2_s2v", "Wan22S2VPipeline"),
+        ("pipeline_wan2_2_vace", "Wan22VACEPipeline"),
+    ],
+)
+def test_wan_pipeline_loaders_share_optional_gate_cleanup(monkeypatch, module_name, class_name) -> None:
+    module = importlib.import_module(f"vllm_omni.diffusion.models.wan2_2.{module_name}")
+    pipeline_cls = getattr(module, class_name)
+    pipeline = pipeline_cls.__new__(pipeline_cls)
+    expected = {"loaded"}
+
+    def fake_loader(model, weights):
+        assert model is pipeline
+        assert list(weights) == [("weight", torch.ones(1))]
+        return expected
+
+    monkeypatch.setattr(module, "load_wan_weights_with_optional_gate", fake_loader)
+
+    assert pipeline_cls.load_weights(pipeline, iter((("weight", torch.ones(1)),))) is expected
+
+
+def test_load_weights_removes_unloaded_vsa_gate(monkeypatch) -> None:
+    pipeline, gate = _make_gate_loading_pipeline()
+
+    class _Loader:
+        def __init__(self, model):
+            del model
+
+        def load_weights(self, weights):
+            return {name for name, _ in weights}
+
+    monkeypatch.setattr("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2.AutoWeightsLoader", _Loader)
+    pipeline.load_weights(iter((("other.weight", torch.ones(1)),)))
+
+    assert pipeline.has_gate_compress_weights is False
+    assert gate.to_gate_compress is None
+
+
+def test_load_weights_keeps_trained_vsa_gate(monkeypatch) -> None:
+    pipeline, gate = _make_gate_loading_pipeline()
+    original_gate = gate.to_gate_compress
+
+    class _Loader:
+        def __init__(self, model):
+            del model
+
+        def load_weights(self, weights):
+            return {name for name, _ in weights}
+
+    monkeypatch.setattr("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2.AutoWeightsLoader", _Loader)
+    pipeline.load_weights(iter((("gate_holder.to_gate_compress.weight", torch.ones(1)),)))
+
+    assert pipeline.has_gate_compress_weights is True
+    assert gate.to_gate_compress is original_gate

tests/diffusion/models/wan2_2/test_wan_euler_scheduler_dmd.py

@@ -0,0 +1,32 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.wan2_2.scheduling_wan_euler import WanEulerScheduler
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def test_dmd_predict_clean_and_add_noise_match_flow_formulas() -> None:
+    scheduler = WanEulerScheduler(num_train_timesteps=1000, shift=8.0)
+    timestep = torch.tensor(757.0)
+    sigma = scheduler.sigma_for_timestep(timestep)
+    sample = torch.tensor([2.0, -1.0])
+    model_output = torch.tensor([0.5, 0.25])
+    noise = torch.tensor([-0.5, 1.5])
+
+    clean = scheduler.predict_clean(model_output, sample, timestep)
+    expected_clean = sample - sigma * model_output
+    torch.testing.assert_close(clean, expected_clean)
+
+    renoised = scheduler.add_noise(clean, noise, timestep)
+    expected_renoised = (1.0 - sigma) * expected_clean + sigma * noise
+    torch.testing.assert_close(renoised, expected_renoised)
+
+
+def test_dmd_sigma_requires_scalar_timestep() -> None:
+    scheduler = WanEulerScheduler(num_train_timesteps=1000, shift=8.0)
+    with pytest.raises(ValueError, match="scalar timestep"):
+        scheduler.sigma_for_timestep(torch.tensor([757.0, 522.0]))

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/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -546,7 +546,7 @@ def test_personaplex_sender_cleanup_drops_delayed_frame_state(build_adapter):
     assert second.codes is None
 
 
-def test_save_async_skips_stale_resumable_chunk_within_segment(build_adapter):
+def test_save_async_skips_stale_resumable_chunk_until_dedup_is_reset(build_adapter):
     adapter, _ = build_adapter(stage_id=1)
     request = _req("req-stream", RequestStatus.WAITING, external_req_id="ext-stream")
     request.resumable = True
@@ -565,51 +565,6 @@ def test_save_async_skips_stale_resumable_chunk_within_segment(build_adapter):
     assert adapter.requests_num_chunks_sent["ext-stream"] == 0
 
 
-def test_save_async_drops_late_previous_segment_after_boundary_reset(build_adapter):
-    adapter, _ = build_adapter(stage_id=1)
-    request = _req("req-stream", RequestStatus.WAITING, external_req_id="ext-stream")
-    request.resumable = True
-    request._omni_segment_generation = 0
-    adapter.requests_num_chunks_sent["ext-stream"] = 111
-
-    adapter.save_async(multimodal_output=None, request=request, is_segment_finished=True)
-
-    late = _req("req-stream-late", RequestStatus.WAITING, external_req_id="ext-stream")
-    late.resumable = True
-    late.num_computed_tokens = 111
-    late._omni_segment_generation = 0
-    adapter.save_async(multimodal_output=None, request=late, is_segment_finished=False)
-
-    next_segment = _req("req-stream-next", RequestStatus.WAITING, external_req_id="ext-stream")
-    next_segment.resumable = True
-    next_segment.num_computed_tokens = 3
-    next_segment._omni_segment_generation = 1
-    adapter.save_async(multimodal_output=None, request=next_segment, is_segment_finished=False)
-
-    assert len(adapter._pending_save_reqs) == 2
-    assert adapter.requests_num_chunks_sent["ext-stream"] == 3
-    assert adapter._pending_save_reqs[-1]["request"] is next_segment
-
-
-def test_send_single_request_preserves_queued_preboundary_chunk(build_adapter):
-    adapter, connector = build_adapter(stage_id=1)
-    adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct(meta=MetaStruct())
-    request = _req("req-queued-old", RequestStatus.WAITING, external_req_id="ext-queued-old")
-    request.resumable = True
-    request._omni_segment_generation = 0
-
-    adapter.save_async(multimodal_output=None, request=request, is_segment_finished=False)
-    adapter.save_async(multimodal_output=None, request=request, is_segment_finished=True)
-
-    old_task = adapter._pending_save_reqs.popleft()
-    adapter._send_single_request(old_task)
-    connector.put.assert_called_once()
-
-    boundary_task = adapter._pending_save_reqs.popleft()
-    adapter._send_single_request(boundary_task)
-    assert connector.put.call_count == 2
-
-
 def test_send_single_request_cleans_up_after_finished_payload(build_adapter, monkeypatch):
     adapter, _ = build_adapter(stage_id=1)
     request = _req("req-finished", RequestStatus.FINISHED_STOPPED, external_req_id="ext-finished")

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,
@@ -30,7 +31,6 @@ from vllm_omni.engine.messages import (
 from vllm_omni.engine.orchestrator import (
     Orchestrator,
     OrchestratorRequestState,
-    StreamingSegmentState,
     _build_terminal_empty_output,
 )
 from vllm_omni.engine.stage_pool import StagePool
@@ -806,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,
     *,
@@ -1952,7 +2016,7 @@ async def test_resumable_segment_boundary_builds_stage_metrics() -> None:
         final_stage_id=0,
     )
     req_state.streaming.enabled = True
-    req_state.streaming.segments[0] = StreamingSegmentState(finished=True)
+    req_state.streaming.segment_finished = True
     req_state.stage_submit_ts[0] = time.time()
     orchestrator.request_states = {"req-stream": req_state}
     orchestrator.stage_pools = [pool]

tests/engine/test_orchestrator_segment_stage_keying.py

@@ -1,130 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Streaming segment-boundary state must be keyed by stage.
-
-``Orchestrator._orchestration_loop`` polls every stage into the same
-``OrchestratorRequestState``. While the segment-boundary fields lived in a single
-flat slot per request, whichever stage polled most recently overwrote the others,
-and the per-stage consumers then read back a different stage's boundary.
-"""
-
-from __future__ import annotations
-
-from dataclasses import dataclass, field
-
-import pytest
-from vllm.sampling_params import SamplingParams
-
-from vllm_omni.engine.orchestrator import (
-    Orchestrator,
-    OrchestratorRequestState,
-    StreamingSegmentState,
-)
-from vllm_omni.experimental.fullduplex.engine.contracts import DuplexRequestIdentity
-from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-THINKER_STAGE = 0
-TALKER_STAGE = 1
-
-
-@dataclass
-class _StageMetrics:
-    pipeline_timings: dict[str, float] = field(default_factory=dict)
-
-
-@dataclass
-class _StageOutput:
-    """Stand-in for a processed stage output, with the fields the router reads."""
-
-    request_id: str
-    finished: bool = False
-    error: object | None = None
-
-
-class _RecordingPool:
-    """Stage pool that records the requests it was asked to build metrics for."""
-
-    def __init__(self) -> None:
-        self.built_for: list[str] = []
-
-    def build_stage_metrics(self, outputs: list[_StageOutput], **_kwargs: object) -> _StageMetrics:
-        self.built_for.extend(output.request_id for output in outputs)
-        return _StageMetrics()
-
-
-def _duplex_req_state(request_id: str = "req-duplex") -> OrchestratorRequestState:
-    req_state = OrchestratorRequestState(
-        request_id=request_id,
-        sampling_params_list=[SamplingParams(max_tokens=1), SamplingParams(max_tokens=1)],
-        final_stage_id=TALKER_STAGE,
-    )
-    req_state.streaming.enabled = True
-    req_state.duplex_identity = DuplexRequestIdentity(
-        session_id="session-1",
-        fence=DuplexFence(session_id="session-1"),
-    )
-    return req_state
-
-
-def test_duplex_output_context_reads_only_its_own_stage_segment() -> None:
-    req_state = _duplex_req_state()
-    req_state.streaming.segments[THINKER_STAGE] = StreamingSegmentState(
-        finished=True,
-        token_ids=[11, 22],
-        output_metadata={"stage": "thinker"},
-    )
-    # The talker polls after the thinker and is still mid-segment.
-    req_state.streaming.segments[TALKER_STAGE] = StreamingSegmentState(finished=False)
-
-    thinker = Orchestrator._duplex_output_context(req_state, stage_id=THINKER_STAGE)
-    talker = Orchestrator._duplex_output_context(req_state, stage_id=TALKER_STAGE)
-
-    assert thinker is not None
-    assert talker is not None
-
-    # These are the assertions that fail on the flat field: the talker is written
-    # second, so the thinker loses its boundary to whichever stage polled last.
-    assert thinker.segment_finished is True
-    assert thinker.segment_token_ids == (11, 22)
-    assert thinker.segment_output_metadata == {"stage": "thinker"}
-
-    assert talker.segment_finished is False
-    assert talker.segment_token_ids == ()
-    assert talker.segment_output_metadata == {}
-
-
-@pytest.mark.asyncio
-async def test_segment_boundary_on_one_stage_does_not_build_metrics_on_another() -> None:
-    req_state = _duplex_req_state()
-    req_state.streaming.segments[THINKER_STAGE] = StreamingSegmentState(finished=True)
-
-    talker_pool = _RecordingPool()
-    orchestrator = object.__new__(Orchestrator)
-    orchestrator.request_states = {req_state.request_id: req_state}
-    orchestrator.stage_pools = [_RecordingPool(), talker_pool]
-
-    routed: list[_StageMetrics | None] = []
-
-    async def record_route(
-        _stage: int,
-        _replica: int,
-        _output: _StageOutput,
-        _state: OrchestratorRequestState,
-        metrics: _StageMetrics | None,
-    ) -> None:
-        routed.append(metrics)
-
-    orchestrator._route_output = record_route
-
-    # The talker's own output is neither finished nor at a segment boundary, so
-    # it must not be treated as one just because the thinker reached one.
-    await orchestrator._handle_processed_outputs(
-        TALKER_STAGE,
-        0,
-        [_StageOutput(request_id=req_state.request_id)],
-    )
-
-    assert talker_pool.built_for == []
-    assert routed == [None]

tests/engine/test_orchestrator_stage_input_bridge.py

@@ -17,7 +17,6 @@ from vllm.sampling_params import SamplingParams
 from vllm_omni.engine.orchestrator import (
     Orchestrator,
     OrchestratorRequestState,
-    StreamingSegmentState,
     _OrchestratorDuplexStagePort,
 )
 from vllm_omni.engine.stage_pool import StagePool
@@ -359,7 +358,7 @@ async def test_streaming_segment_does_not_complete_final_output_stage() -> None:
         final_output_stage_ids={0},
     )
     req_state.streaming.enabled = True
-    req_state.streaming.segments[0] = StreamingSegmentState(finished=True)
+    req_state.streaming.segment_finished = True
     output = SimpleNamespace(
         request_id=req_state.request_id,
         finished=True,

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/openai/test_duplex_protocol.py

@@ -462,16 +462,6 @@ def test_minicpmo_native_capabilities_separate_model_state_from_core_kv_lease():
     assert caps["stage_handoff_transport"] == "scheduler_data_plane"
 
 
-def test_duplex_overlap_policy_auto_remains_a_compatibility_value():
-    from vllm_omni.experimental.fullduplex.openai.protocol import (
-        DuplexOverlapPolicy,
-        DuplexSessionConfig,
-    )
-
-    assert DuplexOverlapPolicy.AUTO.value == "auto"
-    assert DuplexSessionConfig._normalize_overlap_policy("auto") == "auto"
-
-
 def test_minicpmo_native_capabilities_do_not_overclaim_single_session_deployment():
     caps = DuplexCapabilities.minicpmo45_native(max_sessions=1).as_dict()
 

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

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from types import SimpleNamespace
 
@@ -400,7 +400,9 @@ def test_serve_cli_accepts_diffusion_attention_backend():
             "Qwen/Qwen-Image",
             "--omni",
             "--diffusion-attention-backend",
-            "FLASH_ATTN",
+            "FASTVIDEO_VSA",
+            "--fastvideo-vsa-topk",
+            "96",
         ]
     )
 
@@ -408,10 +410,12 @@ def test_serve_cli_accepts_diffusion_attention_backend():
     stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0]
     diffusion_attention_config = stage_cfg["engine_args"]["diffusion_attention_config"]
 
-    assert args.diffusion_attention_backend == "FLASH_ATTN"
+    assert args.diffusion_attention_backend == "FASTVIDEO_VSA"
+    assert args.fastvideo_vsa_topk == 96
     assert isinstance(diffusion_attention_config, AttentionConfig)
     assert diffusion_attention_config.default is not None
-    assert diffusion_attention_config.default.backend == "FLASH_ATTN"
+    assert diffusion_attention_config.default.backend == "FASTVIDEO_VSA"
+    assert diffusion_attention_config.default.backend_kwargs() == {"topk": 96}
 
 
 def test_serve_cli_accepts_request_batch_max_wait_ms():

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

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Config factories for vllm-omni, e.g., StageConfigFactory."""
 
 from __future__ import annotations
@@ -226,7 +226,10 @@ class StageConfigFactory:
                     # If we have a resolver, call it with the optional hf_config
                     # to get the default pipeline config for this key
                     pipeline_cfg = obj(hf_config) if callable(obj) else obj
-                    if pipeline_cfg is not None and pipeline_cfg.diffusers_class_name == class_name:
+                    if pipeline_cfg is not None and class_name in (
+                        pipeline_cfg.diffusers_class_name,
+                        *pipeline_cfg.diffusers_class_aliases,
+                    ):
                         logger.info(
                             "Detected pipeline %r from model_index.json (_class_name=%r)",
                             pipeline_cfg.model_type,
@@ -433,6 +436,7 @@ class StageConfigFactory:
         load-balancer policy (``None`` when no strategy set one) travels with the
         stages instead of through a mutable out-param.
         """
+        deploy_cfg: DeployConfig | None
         if user_deploy_config is not None:
             deploy_cfg = user_deploy_config
         elif deploy_config_path is not None:
@@ -440,9 +444,12 @@ class StageConfigFactory:
             assert deploy_cfg is not None
         elif pipeline_cfg.default_deploy_config_name is not None:
             deploy_cfg = load_deploy_config(_DEPLOY_DIR / pipeline_cfg.default_deploy_config_name)
+            assert deploy_cfg is not None
         else:
             deploy_cfg = DeployConfig()
 
+        assert deploy_cfg is not None
+
         cli_async_chunk = cli_overrides.get("async_chunk")
         if cli_async_chunk is not None:
             deploy_cfg.async_chunk = bool(cli_async_chunk)

vllm_omni/config/omni_config.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Structured vLLM-Omni configuration classes.
 
 This module is additive for Phase 2 of RFC #4021.
@@ -126,7 +126,7 @@ def _enforce_keyword_only_init(cls: type[Any]) -> type[Any]:
 
     cls.__init__ = keyword_only_init
     cls.__signature__ = keyword_only_signature
-    cls.__match_args__ = ()
+    setattr(cls, "__match_args__", ())
     return cls
 
 
@@ -447,7 +447,7 @@ class OmniStageSchedulerConfig(_TrackExplicitConfigFields, VllmSchedulerConfig):
     max_num_seqs: int | None = Field(default=None, ge=1)
     max_num_batched_tokens: int | None = Field(default=None, ge=1)
     max_model_len: int | None = Field(default=None, ge=-1)
-    is_encoder_decoder: InitVar[bool] = False
+    is_encoder_decoder: InitVar[bool] = False  # type: ignore[assignment]
     enable_chunked_prefill: bool | None = None
     async_scheduling: bool | None = None
 
@@ -519,8 +519,8 @@ class OmniStageParallelConfig(_TrackExplicitConfigFields, VllmParallelConfig):
             self.data_parallel_rank,
             self.data_parallel_rpc_port,
             self.worker_cls,
-            self.all2all_backend,
-            self.disable_custom_all_reduce,
+            getattr(self, "all2all_backend", None),
+            getattr(self, "disable_custom_all_reduce", None),
         )
         if self.data_parallel_size_local is not None and self.data_parallel_size_local > self.data_parallel_size:
             raise ValueError(
@@ -741,7 +741,7 @@ class _DiffusionConfigProjection:
         from vllm_omni.diffusion.data import normalize_omni_diffusion_kwargs
 
         normalized_kwargs = normalize_omni_diffusion_kwargs(kwargs)
-        valid_fields = {f.name for f in fields(cls)}
+        valid_fields = {f.name for f in fields(cast(Any, cls))}
         return cls(**{k: v for k, v in normalized_kwargs.items() if k in valid_fields})
 
     def __post_init__(self) -> None:
@@ -893,7 +893,7 @@ class _DiffusionConfigProjection:
                 setattr(self, name, _copy_value(getattr(omni_diffusion_config, name)))
 
 
-_DIFFUSION_CONFIG_FIELDS = frozenset(f.name for f in fields(_DiffusionConfigProjection))
+_DIFFUSION_CONFIG_FIELDS = frozenset(f.name for f in fields(cast(Any, _DiffusionConfigProjection)))
 
 # Current OmniDiffusionConfig still contains a flat mix of shared engine,
 # runtime, parallel, and diffusion-specific knobs. Keep this classification
@@ -956,6 +956,7 @@ _STAGE_DEPLOY_ENGINE_FIELDS: tuple[str, ...] = tuple(_STAGE_DEPLOY_FIELDS)
 _DIFFUSION_BACKCOMPAT_ENGINE_FIELDS = frozenset(
     {
         "diffusion_attention_backend",
+        "fastvideo_vsa_topk",
         "kv_cache_dtype",
         "kv_cache_skip_layers",
         "kv_cache_skip_steps",
@@ -1140,7 +1141,7 @@ def _global_stage_cli_fields() -> frozenset[str]:
     )
     externally_consumed = (
         _NON_STAGE_ENGINE_CLI_FIELDS
-        | frozenset(f.name for f in fields(VllmOmniOrchestratorConfig))
+        | frozenset(f.name for f in fields(cast(Any, VllmOmniOrchestratorConfig)))
         | (orchestrator_field_names() - _STAGE_ENGINE_FIELDS)
     )
     return candidates - externally_consumed
@@ -1238,7 +1239,7 @@ def _stage_sampling_params(
 
 def _orchestrator_cli_overrides(cli_overrides: Mapping[str, Any]) -> dict[str, Any]:
     overrides: dict[str, Any] = {}
-    for config_field in fields(VllmOmniOrchestratorConfig):
+    for config_field in fields(cast(Any, VllmOmniOrchestratorConfig)):
         name = config_field.name
         if name == "deploy_config_path":
             continue
@@ -1443,8 +1444,8 @@ def _with_resolved_processors(
     input_proc: str | None,
     next_stage_proc: str | None,
 ) -> StageConfigType:
-    stage_config._resolved_custom_process_input_func = input_proc
-    stage_config._resolved_custom_process_next_stage_input_func = next_stage_proc
+    setattr(stage_config, "_resolved_custom_process_input_func", input_proc)
+    setattr(stage_config, "_resolved_custom_process_next_stage_input_func", next_stage_proc)
     return stage_config
 
 
@@ -1558,13 +1559,16 @@ def _build_stage_config(
         builder = _STAGE_CONFIG_BUILDERS[topology.execution_type]
     except KeyError as exc:
         raise ValueError(f"Unsupported stage execution type: {topology.execution_type!r}") from exc
-    return builder(
-        pipeline,
-        deploy,
-        topology,
-        stage_deploy,
-        engine,
-        model=model,
+    return cast(
+        StageConfigType,
+        builder(
+            pipeline,
+            deploy,
+            topology,
+            stage_deploy,
+            engine,
+            model=model,
+        ),
     )
 
 
@@ -1608,7 +1612,7 @@ def _build_model_config(
         kwargs["model_subdir"] = topology.model_subdir
     if "tokenizer_subdir" not in kwargs and topology.tokenizer_subdir is not None:
         kwargs["tokenizer_subdir"] = topology.tokenizer_subdir
-    return OmniStageModelConfig(
+    return cast(Any, OmniStageModelConfig)(
         default_sampling_params=default_sampling_params,
         duplex_max_sessions=duplex_max_sessions,
         **kwargs,
@@ -1699,7 +1703,7 @@ def _build_connector_config(
 ) -> OmniStageConnectorConfig:
     output_connectors = stage_deploy.output_connectors if stage_deploy is not None else None
     input_connectors = stage_deploy.input_connectors if stage_deploy is not None else None
-    return OmniStageConnectorConfig(
+    return cast(Any, OmniStageConnectorConfig)(
         async_chunk=bool(deploy.async_chunk),
         omni_kv_config=_copy_value(engine.get("omni_kv_config")),
         output_connectors=_copy_value(output_connectors) if output_connectors else None,
@@ -1848,11 +1852,11 @@ class VllmOmniConfig:
             for topology in pipeline_cfg.stages
         )
 
-        orchestrator_config = VllmOmniOrchestratorConfig(
+        orchestrator_config = cast(Any, VllmOmniOrchestratorConfig)(
             deploy_config_path=loaded_deploy_config_path,
             **_orchestrator_cli_overrides(cli_overrides),
         )
-        return cls(
+        return cast(Any, cls)(
             pipeline_config=pipeline_cfg,
             stage_configs=stage_configs,
             orchestrator_config=orchestrator_config,

vllm_omni/config/stage_config.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Stage configuration system for vLLM-Omni."""
 
 from __future__ import annotations
@@ -234,6 +234,7 @@ class PipelineConfig:
     # this value to auto-detect the pipeline.  Only needed for diffusers-style
     # multi-component repos (e.g. GLM-Image).  ``None`` = not a diffusers model.
     diffusers_class_name: str | None = None
+    diffusers_class_aliases: tuple[str, ...] = ()
     endpoint_restrictions: tuple[EndpointRestriction, ...] = ()
     # Optional model-owned duplex planner loaded by the stable engine runtime.
     duplex_runtime_extension: str | None = None
@@ -356,6 +357,7 @@ class StageDeployConfig:
     diffusers_call_kwargs: dict[str, Any] | None = None
     diffusion_quantization_config: str | None = None
     diffusion_attention_backend: str | None = None
+    fastvideo_vsa_topk: int | None = None
     diffusion_attention_config: dict[str, Any] | None = None
 
     # Diffusion execution, cache, and VAE behavior.

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -218,21 +218,16 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         if stop_after_transfer and req_id in self.requests_needing_kv_transfer:
             self.pending_stop_after_extraction.add(req_id)
 
-    def _drop_aborted_queued_requests(self) -> None:
-        for queue in (self.waiting, self.skipped_waiting):
-            aborted = [req for req in queue if req.status == RequestStatus.FINISHED_ABORTED]
-            if aborted:
-                queue.remove_requests(aborted)
-        self.running[:] = [req for req in self.running if req.status != RequestStatus.FINISHED_ABORTED]
-
     def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
         # Remove FINISHED_ABORTED requests before the upstream scheduler sees
         # them. Upstream vllm raises RuntimeError on this status; omni allows
         # async abort (e.g. client disconnect during TTS streaming) to leave
         # requests in the waiting/running queues temporarily.
-        self._drop_aborted_queued_requests()
+        for queue in (self.waiting, self.running):
+            for req in list(queue):
+                if getattr(req, "status", None) == RequestStatus.FINISHED_ABORTED:
+                    queue.remove(req)
         self._process_pending_omni_inputs(model_mode="ar")
-        self._resync_streaming_input_counter()
 
         original_waiting = None
         if self._should_defer_waiting_admission():
@@ -359,11 +354,6 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
 
             req_index = model_runner_output.req_id_to_index[req_id]
             generated_token_ids = sampled_token_ids[req_index] if sampled_token_ids else []
-
-            if generated_token_ids and getattr(request, "async_tokens_to_discard", 0) > 0:
-                self._update_request_with_output(request, generated_token_ids)
-                continue
-
             status_before_stop = request.status
             new_logprobs = None
             logprob_validation_failed = False
@@ -661,7 +651,6 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         if self.chunk_transfer_adapter and self.chunk_transfer_adapter.receives_chunks:
             self.chunk_transfer_adapter.requests_num_chunks_sent.pop(session.external_req_id, None)
             if stage_id != 0:
-                session._omni_segment_generation = int(getattr(session, "_omni_segment_generation", 0) or 0) + 1
                 # Downstream async-chunk stages receive real payloads from the
                 # connector. This update only resumes polling for the next segment.
                 self.chunk_transfer_adapter.segment_finished_requests.discard(session.request_id)
@@ -693,30 +682,6 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
             self._release_replaced_streaming_prompt_cache(session)
             self._replace_streaming_session(session, update)
             return
-        native_duplex_update = any(
-            isinstance(info, dict) and info.get("native_duplex") is True for info in update_infos
-        )
-        if native_duplex_update:
-            max_model_len = int(getattr(self.vllm_config.model_config, "max_model_len", 0) or 0)
-            current_prompt_len = len(getattr(session, "prompt_token_ids", None) or ())
-            computed_tokens = int(getattr(session, "num_computed_tokens", 0) or 0)
-            kept_output_len = max(0, computed_tokens - current_prompt_len)
-            next_prompt_len = len(update.prompt_token_ids or ())
-            candidate_len = current_prompt_len + kept_output_len + next_prompt_len
-            # Keep one context position available for the next sampled codec token.
-            if max_model_len > 0 and candidate_len >= max_model_len:
-                logger.info(
-                    "Resetting native duplex Stage-1 prompt before context overflow: "
-                    "req=%s current=%s kept_output=%s next=%s max_model_len=%s",
-                    session.request_id,
-                    current_prompt_len,
-                    kept_output_len,
-                    next_prompt_len,
-                    max_model_len,
-                )
-                self._replace_streaming_session(session, update)
-                return
-        session._omni_segment_generation = int(getattr(session, "_omni_segment_generation", 0) or 0) + 1
         super()._update_request_as_session(session, update)
         if hasattr(update, "model_intermediate_buffer"):
             session.model_intermediate_buffer = update.model_intermediate_buffer

vllm_omni/core/sched/omni_scheduler_mixin.py

@@ -146,7 +146,6 @@ class OmniSchedulerMixin:
 
     def _replace_streaming_session(self, session: Request, update: StreamingUpdate) -> None:
         """Replace a downstream stage's placeholder with its next payload."""
-        session._omni_segment_generation = int(getattr(session, "_omni_segment_generation", 0) or 0) + 1
         adapter = getattr(self, "chunk_transfer_adapter", None)
         if adapter is not None:
             adapter.segment_finished_requests.discard(session.request_id)
@@ -646,7 +645,6 @@ class OmniSchedulerMixin:
         )
         finished = super().finish_requests(request_ids, finished_status)
         self._purge_finished_from_running(target_request_ids)
-        self._resync_streaming_input_counter()
 
         for request in finished:
             self._free_input_coordinator_request(request.request_id)
@@ -812,18 +810,3 @@ class OmniSchedulerMixin:
             return resumable_segment_stop and req.request_id not in target_request_ids
 
         self.running[:] = [req for req in self.running if keep_running(req)]
-
-    def _resync_streaming_input_counter(self) -> None:
-        counter = getattr(self, "num_waiting_for_streaming_input", None)
-        if counter is None:
-            return
-        parked = sum(
-            1
-            for queue in (getattr(self, "waiting", None), getattr(self, "skipped_waiting", None))
-            if queue
-            for request in queue
-            if request.status == RequestStatus.WAITING_FOR_STREAMING_REQ
-        )
-        if parked != counter:
-            logger.debug("[Omni] resynced streaming-input counter %s -> %s", counter, parked)
-            self.num_waiting_for_streaming_input = parked

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/fastvideo_vsa.py

@@ -0,0 +1,431 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import functools
+import math
+from collections.abc import Mapping
+from typing import Any
+
+import torch
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.attention.backends.abstract import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+)
+from vllm_omni.diffusion.attention.backends.sdpa import SDPAImpl
+from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
+
+logger = init_logger(__name__)
+
+
+# Keep the external FastVideo pybind/CUDA kernel opaque to torch.compile.
+# This mirrors the SageAttention3 backend pattern: tracing the raw extension
+# through Dynamo can reach Inductor scheduling with unstable internal op names
+# (e.g. KeyError: "op12").  The custom op gives Dynamo a single Tensor->Tensor
+# boundary and lets Inductor schedule the surrounding Wan block normally.
+if not hasattr(torch.ops.vllm_omni, "fastvideo_vsa_bshd"):
+
+    @torch.library.custom_op("vllm_omni::fastvideo_vsa_bshd", mutates_args=())
+    def _fastvideo_vsa_bshd_op(
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        variable_block_sizes: torch.Tensor,
+        q_variable_block_sizes: torch.Tensor,
+        compress_attn_weight: torch.Tensor,
+        topk: int,
+        block_t: int,
+        block_h: int,
+        block_w: int,
+    ) -> torch.Tensor:
+        from fastvideo_kernel import video_sparse_attn_bshd
+
+        return video_sparse_attn_bshd(
+            query,
+            key,
+            value,
+            variable_block_sizes=variable_block_sizes,
+            q_variable_block_sizes=q_variable_block_sizes,
+            topk=topk,
+            block_size=(block_t, block_h, block_w),
+            compress_attn_weight=compress_attn_weight if compress_attn_weight.numel() else None,
+        )
+
+    @_fastvideo_vsa_bshd_op.register_fake
+    def _(
+        query,
+        key,
+        value,
+        variable_block_sizes,
+        q_variable_block_sizes,
+        compress_attn_weight,
+        topk,
+        block_t,
+        block_h,
+        block_w,
+    ):
+        del (
+            key,
+            value,
+            variable_block_sizes,
+            q_variable_block_sizes,
+            compress_attn_weight,
+            topk,
+            block_t,
+            block_h,
+            block_w,
+        )
+        return torch.empty_like(query)
+
+
+_fastvideo_vsa_bshd_op = torch.ops.vllm_omni.fastvideo_vsa_bshd
+
+
+@functools.lru_cache(maxsize=32)
+def _get_tile_partition_indices(
+    dit_seq_shape: tuple[int, int, int],
+    tile_size: tuple[int, int, int],
+    device: torch.device,
+) -> torch.Tensor:
+    t_size, h_size, w_size = dit_seq_shape
+    tile_t, tile_h, tile_w = tile_size
+    indices = torch.arange(t_size * h_size * w_size, device=device, dtype=torch.long).reshape(t_size, h_size, w_size)
+    tiles = []
+    for tile_t_idx in range(math.ceil(t_size / tile_t)):
+        for tile_h_idx in range(math.ceil(h_size / tile_h)):
+            for tile_w_idx in range(math.ceil(w_size / tile_w)):
+                tiles.append(
+                    indices[
+                        tile_t_idx * tile_t : min((tile_t_idx + 1) * tile_t, t_size),
+                        tile_h_idx * tile_h : min((tile_h_idx + 1) * tile_h, h_size),
+                        tile_w_idx * tile_w : min((tile_w_idx + 1) * tile_w, w_size),
+                    ].flatten()
+                )
+    return torch.cat(tiles, dim=0)
+
+
+@functools.lru_cache(maxsize=32)
+def _construct_variable_block_sizes(
+    dit_seq_shape: tuple[int, int, int],
+    tile_size: tuple[int, int, int],
+    device: torch.device,
+) -> torch.Tensor:
+    num_tiles = tuple(math.ceil(seq_dim / tile_dim) for seq_dim, tile_dim in zip(dit_seq_shape, tile_size))
+
+    def _sizes(dim_len: int, tile: int, n_tiles: int) -> torch.Tensor:
+        sizes = torch.full((n_tiles,), tile, dtype=torch.int32, device=device)
+        remainder = dim_len - (n_tiles - 1) * tile
+        sizes[-1] = remainder if remainder > 0 else tile
+        return sizes
+
+    t_sizes = _sizes(dit_seq_shape[0], tile_size[0], num_tiles[0])
+    h_sizes = _sizes(dit_seq_shape[1], tile_size[1], num_tiles[1])
+    w_sizes = _sizes(dit_seq_shape[2], tile_size[2], num_tiles[2])
+    return (t_sizes[:, None, None] * h_sizes[None, :, None] * w_sizes[None, None, :]).reshape(-1)
+
+
+@functools.lru_cache(maxsize=32)
+def _get_non_pad_index(variable_block_sizes: torch.Tensor, max_block_size: int) -> torch.Tensor:
+    num_blocks = variable_block_sizes.shape[0]
+    device = variable_block_sizes.device
+    starts = torch.arange(num_blocks, device=device) * max_block_size
+    padded_index = starts[:, None] + torch.arange(max_block_size, device=device)[None, :]
+    valid = torch.arange(max_block_size, device=device)[None, :] < variable_block_sizes[:, None]
+    return padded_index[valid]
+
+
+@torch.compiler.disable
+def _get_tile_metadata(
+    dit_seq_shape: tuple[int, int, int],
+    tile_size: tuple[int, int, int],
+    block_elements: int,
+    device: torch.device,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    tile_partition_indices = _get_tile_partition_indices(dit_seq_shape, tile_size, device)
+    variable_block_sizes = _construct_variable_block_sizes(dit_seq_shape, tile_size, device)
+    non_pad_index = _get_non_pad_index(variable_block_sizes, block_elements)
+    untile_combined_index = non_pad_index[torch.argsort(tile_partition_indices)]
+    return tile_partition_indices, variable_block_sizes, non_pad_index, untile_combined_index
+
+
+def _get_vsa_dit_seq_shape(attn_metadata: AttentionMetadata | None) -> tuple[int, int, int] | None:
+    if attn_metadata is None:
+        return None
+    value = attn_metadata.extra.get("vsa_dit_seq_shape")
+    if value is None:
+        return None
+    if not isinstance(value, (list, tuple)) or len(value) != 3:
+        return None
+    return (int(value[0]), int(value[1]), int(value[2]))
+
+
+def _get_gate_compress(attn_metadata: AttentionMetadata | None) -> torch.Tensor | None:
+    if attn_metadata is None:
+        return None
+    value = attn_metadata.extra.get("gate_compress")
+    return value if isinstance(value, torch.Tensor) else None
+
+
+def _preserve_vsa_all_blocks(attn_metadata: AttentionMetadata | None) -> bool:
+    if attn_metadata is None:
+        return False
+    return attn_metadata.extra.get("preserve_vsa_all_blocks") is True
+
+
+class FastVideoVSABackend(AttentionBackend):
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        # FastVideo VSA is intended for video DiT head sizes such as 64/128.
+        # Keep this permissive and let the runtime fallback handle unsupported
+        # cases from the installed fastvideo-kernel build.
+        return []
+
+    @staticmethod
+    def get_name() -> str:
+        return "FASTVIDEO_VSA"
+
+    @staticmethod
+    def get_impl_cls() -> type[FastVideoVSAImpl]:
+        return FastVideoVSAImpl
+
+
+class FastVideoVSAImpl(AttentionImpl):
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        softmax_scale: float,
+        causal: bool = False,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        qkv_layout: str | None = None,
+        backend_kwargs: Mapping[str, Any] | None = None,
+        **extra_impl_args,
+    ) -> None:
+        backend_kwargs = backend_kwargs or {}
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
+        self.head_size = head_size
+        self.softmax_scale = softmax_scale
+        self.causal = causal
+        self.qkv_layout = qkv_layout
+
+        self.topk = int(backend_kwargs.get("topk", 64))
+        self.block_size = self._parse_block_size(backend_kwargs.get("block_size", (4, 8, 8)))
+        self.block_elements = self.block_size[0] * self.block_size[1] * self.block_size[2]
+        self.min_seq_len = int(backend_kwargs.get("min_seq_len", self.block_elements * 2))
+        self.fallback_on_error = bool(backend_kwargs.get("fallback_on_error", True))
+        self.disable_when_sp_active = bool(backend_kwargs.get("disable_when_sp_active", True))
+
+        self.sdpa_fallback = SDPAImpl(
+            num_heads=num_heads,
+            head_size=head_size,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            num_kv_heads=num_kv_heads,
+            qkv_layout=qkv_layout,
+        )
+
+        if self.block_elements != 256:
+            logger.warning(
+                "FASTVIDEO_VSA currently uses fastvideo_kernel.video_sparse_attn_bshd, "
+                "which supports only 256-token blocks. Configured block_size=%s "
+                "(product=%d) will fall back to SDPA.",
+                self.block_size,
+                self.block_elements,
+            )
+
+    @staticmethod
+    def _parse_block_size(value: Any) -> tuple[int, int, int]:
+        if isinstance(value, int):
+            return (value, value, value)
+        if isinstance(value, (list, tuple)) and len(value) == 3:
+            return (int(value[0]), int(value[1]), int(value[2]))
+        raise ValueError(f"FASTVIDEO_VSA block_size must be an int or length-3 tuple/list, got {value!r}")
+
+    def _fallback(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None,
+        reason: str,
+    ) -> torch.Tensor:
+        logger.warning_once("FASTVIDEO_VSA falling back to SDPA: %s", reason)
+        return self.sdpa_fallback.forward(query, key, value, attn_metadata)
+
+    def _fallback_reason(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None,
+    ) -> str | None:
+        if self.causal:
+            return "causal attention is not supported"
+        if self.block_elements != 256:
+            return f"block_elements must be 256, got {self.block_elements}"
+        if self.topk <= 0:
+            return f"topk must be positive, got {self.topk}"
+        if query.ndim != 4 or key.ndim != 4 or value.ndim != 4:
+            return f"expected [B, S, H, D] tensors, got {query.shape}, {key.shape}, {value.shape}"
+        if query.shape[0] != key.shape[0] or query.shape[0] != value.shape[0]:
+            return "batch dimensions must match"
+        if query.shape[2:] != key.shape[2:] or query.shape[2:] != value.shape[2:]:
+            return "head/head_dim dimensions must match"
+        if query.shape[1] != key.shape[1] or query.shape[1] != value.shape[1]:
+            return "initial VSA backend supports self-attention with Sq == Skv only"
+        if query.shape[1] < self.min_seq_len:
+            return f"sequence length {query.shape[1]} is below min_seq_len {self.min_seq_len}"
+        dit_seq_shape = _get_vsa_dit_seq_shape(attn_metadata)
+        if dit_seq_shape is None:
+            return "vsa_dit_seq_shape metadata is required"
+        if math.prod(dit_seq_shape) != query.shape[1]:
+            return f"vsa_dit_seq_shape product {math.prod(dit_seq_shape)} != sequence length {query.shape[1]}"
+        num_blocks = math.prod(
+            math.ceil(seq_dim / tile_dim) for seq_dim, tile_dim in zip(dit_seq_shape, self.block_size)
+        )
+        if self.topk > num_blocks:
+            return f"topk {self.topk} > num_blocks {num_blocks}"
+        if query.dtype not in (torch.float16, torch.bfloat16):
+            return f"dtype {query.dtype} is not supported"
+        if key.dtype != query.dtype or value.dtype != query.dtype:
+            return "q/k/v dtypes must match"
+        if query.device.type != "cuda" or key.device.type != "cuda" or value.device.type != "cuda":
+            return "q/k/v must be CUDA tensors"
+        expected_scale = self.head_size**-0.5
+        if abs(float(self.softmax_scale) - float(expected_scale)) > 1e-6:
+            return f"softmax_scale {self.softmax_scale} differs from FastVideo VSA scale {expected_scale}"
+        if attn_metadata is not None and attn_metadata.attn_mask is not None:
+            return "attention masks are not supported"
+        if attn_metadata is not None and attn_metadata.full_attn_spans is not None:
+            return "piecewise/full attention spans are not supported"
+        if self.num_heads != self.num_kv_heads:
+            return "GQA/MQA is not supported"
+        if self.disable_when_sp_active and is_forward_context_available():
+            ctx = get_forward_context()
+            if getattr(ctx, "sp_active", False):
+                return "sequence parallel context is active"
+        return None
+
+    def forward_cuda(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None = None,
+    ) -> torch.Tensor:
+        reason = self._fallback_reason(query, key, value, attn_metadata)
+        if reason is not None:
+            return self._fallback(query, key, value, attn_metadata, reason)
+
+        seq_len = query.shape[1]
+        dit_seq_shape = _get_vsa_dit_seq_shape(attn_metadata)
+        assert dit_seq_shape is not None
+        num_blocks = math.prod(
+            math.ceil(seq_dim / tile_dim) for seq_dim, tile_dim in zip(dit_seq_shape, self.block_size)
+        )
+        preserve_all_blocks = _preserve_vsa_all_blocks(attn_metadata)
+        use_native_sdpa = self.topk == num_blocks and not preserve_all_blocks
+        route = "SDPA" if use_native_sdpa else "VSA_ALL_BLOCKS" if self.topk == num_blocks else "VSA"
+        checkpoint_mode = "fastvideo_dmd" if preserve_all_blocks else "native"
+        logger.info_once(
+            "FASTVIDEO_VSA routing: seq_len=%d, dit_seq_shape=%s, block_size=%s, num_blocks=%d, "
+            "topk=%d, keep_ratio=%.1f%%, checkpoint_mode=%s, route=%s",
+            seq_len,
+            dit_seq_shape,
+            self.block_size,
+            num_blocks,
+            self.topk,
+            100.0 * self.topk / num_blocks,
+            checkpoint_mode,
+            route,
+        )
+        if use_native_sdpa:
+            return self._fallback(
+                query,
+                key,
+                value,
+                attn_metadata,
+                f"topk {self.topk} selects all blocks for a native checkpoint",
+            )
+
+        try:
+            tile_partition_indices, variable_block_sizes, non_pad_index, untile_combined_index = _get_tile_metadata(
+                dit_seq_shape,
+                self.block_size,
+                self.block_elements,
+                query.device,
+            )
+
+            padded_len = variable_block_sizes.numel() * self.block_elements
+            target_shape = (query.shape[0], padded_len, query.shape[2], query.shape[3])
+            query_tiled = torch.zeros(target_shape, device=query.device, dtype=query.dtype)
+            key_tiled = torch.zeros_like(query_tiled)
+            value_tiled = torch.zeros_like(query_tiled)
+            query_tiled[:, non_pad_index] = query[:, tile_partition_indices]
+            key_tiled[:, non_pad_index] = key[:, tile_partition_indices]
+            value_tiled[:, non_pad_index] = value[:, tile_partition_indices]
+            # Gate behavior is checkpoint-driven, not user-configured.
+            # Wan VSA layers always provide a gate projection. Its zero
+            # initialization makes checkpoints without gate weights sparse-only;
+            # checkpoints containing to_gate_compress weights use the learned gate.
+            gate_compress = _get_gate_compress(attn_metadata)
+            if gate_compress is None:
+                gate_compress = torch.zeros_like(query)
+            elif gate_compress.shape != query.shape:
+                raise ValueError(f"gate_compress shape {gate_compress.shape} must match query shape {query.shape}")
+            gate_tiled = torch.zeros_like(query_tiled)
+            gate_tiled[:, non_pad_index] = gate_compress[:, tile_partition_indices]
+            compress_attn_weight = gate_tiled
+
+            output = _fastvideo_vsa_bshd_op(
+                query_tiled.contiguous(),
+                key_tiled.contiguous(),
+                value_tiled.contiguous(),
+                variable_block_sizes,
+                variable_block_sizes,
+                compress_attn_weight,
+                self.topk,
+                self.block_size[0],
+                self.block_size[1],
+                self.block_size[2],
+            )
+            return output[:, untile_combined_index].contiguous()
+        except Exception as exc:
+            if not self.fallback_on_error:
+                raise
+            return self._fallback(query, key, value, attn_metadata, f"VSA kernel failed: {exc}")
+
+    def forward_npu(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None = None,
+    ) -> torch.Tensor:
+        return self.sdpa_fallback.forward_npu(query, key, value, attn_metadata)
+
+    def forward_xpu(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None = None,
+    ) -> torch.Tensor:
+        return self.sdpa_fallback.forward_xpu(query, key, value, attn_metadata)
+
+    def forward_musa(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None = None,
+    ) -> torch.Tensor:
+        return self.sdpa_fallback.forward_musa(query, key, value, attn_metadata)

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/attention/backends/registry.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Diffusion attention backend registry.
 
 This module provides an enum-based registry for diffusion attention backends,
@@ -23,7 +23,7 @@ logger = init_logger(__name__)
 class _DiffusionBackendEnumMeta(EnumMeta):
     """Metaclass for DiffusionAttentionBackendEnum to provide better error messages."""
 
-    def __getitem__(cls, name: str) -> "DiffusionAttentionBackendEnum":
+    def __getitem__(cls, name: str) -> "DiffusionAttentionBackendEnum":  # type: ignore[override]
         """Get backend by name with helpful error messages."""
         try:
             return super().__getitem__(name)  # type: ignore[return-value]
@@ -62,6 +62,7 @@ class DiffusionAttentionBackendEnum(Enum, metaclass=_DiffusionBackendEnumMeta):
     SAGE_ATTN_3 = "vllm_omni.diffusion.attention.backends.sage_attn3.SageAttention3Backend"
     CUDNN_ATTN = "vllm_omni.diffusion.attention.backends.cudnn_attn.CuDNNAttentionBackend"
     FLASHINFER_ATTN = "vllm_omni.diffusion.attention.backends.flashinfer_attn.FlashInferAttentionBackend"
+    FASTVIDEO_VSA = "vllm_omni.diffusion.attention.backends.fastvideo_vsa.FastVideoVSABackend"
     FLASH_ATTN_HUB = "vllm_omni.diffusion.attention.backends.flash_attn_hub.FlashAttentionHubBackend"
     FLASH_ATTN_3_HUB = "vllm_omni.diffusion.attention.backends.flash_attn_hub.FlashAttention3HubBackend"
     TRTLLM_ATTN = "vllm_omni.diffusion.attention.backends.trtllm_attn.TrtllmAttentionBackend"

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/data.py

@@ -1,6 +1,6 @@
 # adapted from sglang and fastvideo
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 import copy
 import math
 import os
@@ -77,11 +77,13 @@ def normalize_omni_diffusion_kwargs(kwargs: Mapping[str, Any]) -> dict[str, Any]
     # Handle "diffusion_attention_backend" shorthand: merge into
     # diffusion_attention_config before field filtering.
     diffusion_attn_backend = normalized.pop("diffusion_attention_backend", None)
-    if diffusion_attn_backend is not None:
+    fastvideo_vsa_topk = normalized.pop("fastvideo_vsa_topk", None)
+    if diffusion_attn_backend is not None or fastvideo_vsa_topk is not None:
         existing = normalized.get("diffusion_attention_config")
         normalized["diffusion_attention_config"] = parse_attention_config(
             existing,
             attention_backend=diffusion_attn_backend,
+            fastvideo_vsa_topk=fastvideo_vsa_topk,
         )
 
     # Check environment variable as fallback for cache_backend.
@@ -104,6 +106,7 @@ def parse_kv_cache_skip_selector(
     """Parse a non-negative index selector such as "0-9,20,25-30"."""
     if selector is None:
         return None
+    values: set[int]
     if isinstance(selector, set):
         values = selector
     elif isinstance(selector, (list, tuple)):
@@ -112,7 +115,7 @@ def parse_kv_cache_skip_selector(
         text = selector.strip()
         if not text:
             return None
-        values: set[int] = set()
+        values = set()
         for chunk in text.split(","):
             token = chunk.strip()
             if not token:
@@ -245,6 +248,7 @@ class DiffusionParallelConfig:
         assert self.pipeline_parallel_size > 0, "Pipeline parallel size must be > 0"
         assert self.data_parallel_size is None or self.data_parallel_size > 0, "Data parallel size must be > 0"
         assert self.tensor_parallel_size > 0, "Tensor parallel size must be > 0"
+        assert self.sequence_parallel_size is not None
         assert self.sequence_parallel_size > 0, "Sequence parallel size must be > 0"
         assert self.ulysses_degree > 0, "Ulysses degree must be > 0"
         assert self.ring_degree > 0, "Ring degree must be > 0"
@@ -362,6 +366,7 @@ class DiffusionParallelConfig:
             self.world_size = world_size
             return 1
 
+        assert self.sequence_parallel_size is not None
         non_dp_size = (
             self.pipeline_parallel_size
             * self.tensor_parallel_size
@@ -1063,6 +1068,7 @@ class OmniDiffusionConfig:
 
         if self.diffusion_compile_granularity == "full":
             incompatible_features = []
+            assert self.parallel_config.sequence_parallel_size is not None
             if self.parallel_config.use_hsdp:
                 incompatible_features.append("HSDP")
             if self.parallel_config.sequence_parallel_size > 1:
@@ -1262,6 +1268,7 @@ class OmniDiffusionConfig:
         if self.model_class_name is None and self.diffusion_load_format == "diffusers":
             self.model_class_name = "DiffusersAdapterPipeline"
 
+        assert self.model is not None
         try:
             config_dict = get_diffusion_model_index(
                 self.model,
@@ -1338,6 +1345,7 @@ class OmniDiffusionConfig:
 
                 from vllm_omni.diffusion.utils.hf_utils import _looks_like_hidream_o1
 
+                assert self.model is not None
                 is_hidream_o1 = _looks_like_hidream_o1(self.model, cfg)
                 if self.model_class_name == "HiDreamO1ImagePipeline" and not is_hidream_o1:
                     raise ValueError(f"Checkpoint {self.model} does not have the HiDream-O1 signature")
@@ -1395,6 +1403,7 @@ class OmniDiffusionConfig:
                 elif model_type == "vla":
                     from vllm_omni.diffusion.utils.hf_utils import _looks_like_dreamzero
 
+                    assert self.model is not None
                     if _looks_like_dreamzero(self.model):
                         self.model_class_name = "DreamZeroPipeline"
                         self.set_tf_model_config(TransformerConfig())
@@ -1593,10 +1602,11 @@ class AttnQuantSpec:
         for name, v in (("dtype_qk", self.dtype_qk), ("dtype_vo", self.dtype_vo)):
             if v is not None and v not in self._VALID_DTYPES:
                 raise ValueError(f"quant.{name}={v!r} unsupported; use one of {sorted(self._VALID_DTYPES)}.")
-        for name, v in (("q_block_size", self.q_block_size), ("k_block_size", self.k_block_size)):
-            if v not in self._VALID_BLOCK_SIZES:
+        for block_name, block_value in (("q_block_size", self.q_block_size), ("k_block_size", self.k_block_size)):
+            if block_value not in self._VALID_BLOCK_SIZES:
                 raise ValueError(
-                    f"quant.{name}={v!r} unsupported; kernels exist only for {sorted(self._VALID_BLOCK_SIZES)}."
+                    f"quant.{block_name}={block_value!r} unsupported; "
+                    f"kernels exist only for {sorted(self._VALID_BLOCK_SIZES)}."
                 )
 
     @property
@@ -1638,6 +1648,7 @@ class AttentionSpec:
     backend: str
     skip_softmax: SkipSoftmaxSpec | None = None
     quant: AttnQuantSpec | None = None
+    fastvideo_vsa_topk: int | None = None
     block_sparse: BlockSparseSpec | None = None
     skip_calibration: dict | None = field(default=None, repr=False)
 
@@ -1657,6 +1668,11 @@ class AttentionSpec:
                 f"quant is only supported by the TRTLLM_ATTN and FLASHINFER_ATTN backends, but "
                 f"backend={self.backend!r}. Remove quant or set a supported backend."
             )
+        if self.fastvideo_vsa_topk is not None:
+            if self.backend.upper() != "FASTVIDEO_VSA":
+                raise ValueError("fastvideo_vsa_topk is only supported by the FASTVIDEO_VSA backend.")
+            if self.fastvideo_vsa_topk <= 0:
+                raise ValueError("fastvideo_vsa_topk must be positive.")
         if self.backend.upper() in BLOCK_SPARSE_BACKENDS:
             # Selecting the backend is the opt-in; without an explicit block the
             # defaults apply rather than silently running dense.
@@ -1698,12 +1714,15 @@ class AttentionSpec:
             if q.flashinfer_backend is not None:
                 quant_kw["flashinfer_backend"] = q.flashinfer_backend
             kw["quant"] = quant_kw
+        if self.fastvideo_vsa_topk is not None:
+            kw["topk"] = self.fastvideo_vsa_topk
         if self.block_sparse is not None:
             bs = self.block_sparse
             kw["sparsity"] = bs.sparsity
             kw["start_step"] = bs.start_step
             if bs.skip_layer_indices:
                 kw["skip_layers"] = sorted(bs.skip_layer_indices)
+
         return kw or None
 
 
@@ -1773,7 +1792,7 @@ class AttentionConfig:
             normalized[role] = node
             return
 
-        spec_keys = {"backend", "skip_softmax", "quant", "block_sparse"}
+        spec_keys = {"backend", "skip_softmax", "quant", "fastvideo_vsa_topk", "block_sparse"}
         node_dict = dict(node)
         node_keys = set(node_dict)
         if node_keys & spec_keys:
@@ -1812,6 +1831,7 @@ def parse_attention_config(
     attention_config: AttentionConfig | Mapping[str, Any] | None = None,
     *,
     attention_backend: str | None = None,
+    fastvideo_vsa_topk: int | None = None,
 ) -> AttentionConfig:
     """Pure type-conversion: coerce *attention_config* to an AttentionConfig.
 
@@ -1839,6 +1859,14 @@ def parse_attention_config(
         if attention_backend.lower() != "auto":
             normalized.default = AttentionSpec(backend=attention_backend)
 
+    if fastvideo_vsa_topk is not None:
+        if normalized.default is None:
+            raise ValueError("--fastvideo-vsa-topk requires --diffusion-attention-backend FASTVIDEO_VSA.")
+        if normalized.default.backend.upper() != "FASTVIDEO_VSA":
+            raise ValueError("--fastvideo-vsa-topk is only valid with the FASTVIDEO_VSA backend.")
+        normalized.default.fastvideo_vsa_topk = fastvideo_vsa_topk
+        normalized.default.__post_init__()
+
     return normalized
 
 

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/lora/manager.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 import math
 import time
@@ -289,11 +289,35 @@ class DiffusionLoRAManager:
         lora_path = get_adapter_absolute_path(lora_request.lora_path)
         logger.debug("Resolved LoRA path: %s", lora_path)
 
-        peft_helper = PEFTHelper.from_local_dir(
-            lora_path,
-            max_position_embeddings=None,  # no need in diffusion
-            tensorizer_config_dict=lora_request.tensorizer_config_dict,
-        )
+        model_loader = getattr(self.pipeline, "_load_diffusion_lora_adapter", None)
+        loaded = None
+        if callable(model_loader):
+            loaded = model_loader(
+                lora_request=lora_request,
+                lora_path=lora_path,
+                dtype=self.dtype,
+            )
+
+        if loaded is None:
+            peft_helper = PEFTHelper.from_local_dir(
+                lora_path,
+                max_position_embeddings=None,  # no need in diffusion
+                tensorizer_config_dict=lora_request.tensorizer_config_dict,
+            )
+
+            lora_model = LoRAModel.from_local_checkpoint(
+                lora_path,
+                expected_lora_modules=self._expected_lora_modules,
+                peft_helper=peft_helper,
+                lora_model_id=lora_request.lora_int_id,
+                device="cpu",  # consistent w/ vllm's behavior
+                dtype=self.dtype,
+                model_vocab_size=None,
+                tensorizer_config_dict=lora_request.tensorizer_config_dict,
+                weights_mapper=None,
+            )
+        else:
+            lora_model, peft_helper = loaded
 
         logger.info(
             "Loaded PEFT config: r=%d, lora_alpha=%d, target_modules=%s",
@@ -302,18 +326,6 @@ class DiffusionLoRAManager:
             peft_helper.target_modules,
         )
 
-        lora_model = LoRAModel.from_local_checkpoint(
-            lora_path,
-            expected_lora_modules=self._expected_lora_modules,
-            peft_helper=peft_helper,
-            lora_model_id=lora_request.lora_int_id,
-            device="cpu",  # consistent w/ vllm's behavior
-            dtype=self.dtype,
-            model_vocab_size=None,
-            tensorizer_config_dict=lora_request.tensorizer_config_dict,
-            weights_mapper=None,
-        )
-
         logger.info(
             "Loaded LoRA model: id=%d, num_modules=%d, modules=%s",
             lora_model.id,
@@ -522,13 +534,17 @@ class DiffusionLoRAManager:
         matches_scale = self._adapter_scales.get(adapter_id) == rounded_scale
         return is_active and matches_scale
 
-    def _activate_adapter(self, adapter_id: int, scale: float) -> None:
-        if self._is_active_at_scale(adapter_id, scale):
-            logger.debug("Adapter %d already active at scale %.3f skipping", adapter_id, scale)
-            return
+    def _bind_adapter_weights(self, lora_model: LoRAModel, scale: float) -> None:
+        binding_validator = getattr(self.pipeline, "_validate_diffusion_lora_binding", None)
+        lora_names_by_id = (
+            {id(weights): name for name, weights in lora_model.loras.items()} if callable(binding_validator) else {}
+        )
+        bound_lora_names: set[str] = set()
 
-        logger.info("Activating adapter: id=%d", adapter_id)
-        lora_model = self._registered_adapters[adapter_id]
+        def _record_bound(weights: LoRALayerWeights | PackedLoRALayerWeights) -> None:
+            name = lora_names_by_id.get(id(weights))
+            if name is not None:
+                bound_lora_names.add(name)
 
         # activate weights in each LoRA layer
         for full_module_name, lora_layer in self._lora_modules.items():
@@ -570,6 +586,9 @@ class DiffusionLoRAManager:
                         lora_b_list.append(sub_lora.lora_b * scale)
 
                     lora_layer.set_lora(index=0, lora_a=lora_a_list, lora_b=lora_b_list)
+                    for sub_lora in sub_loras:
+                        if sub_lora is not None:
+                            _record_bound(sub_lora)
                     logger.debug(
                         "Activated packed LoRA for %s via submodules=%s (scale=%.2f)",
                         full_module_name,
@@ -588,6 +607,7 @@ class DiffusionLoRAManager:
                     for b in lora_weights.lora_b
                 ]
                 lora_layer.set_lora(index=0, lora_a=lora_a_list, lora_b=lora_b_list)
+                _record_bound(lora_weights)
                 logger.debug(
                     "Activated packed LoRA for %s (scale=%.2f)",
                     full_module_name,
@@ -618,6 +638,7 @@ class DiffusionLoRAManager:
                 lora_a_list = [lora_weights.lora_a] * n_slices
                 lora_b_list = [b * scale for b in b_splits]
                 lora_layer.set_lora(index=0, lora_a=lora_a_list, lora_b=lora_b_list)
+                _record_bound(lora_weights)
                 logger.debug(
                     "Activated fused LoRA for packed layer %s (scale=%.2f)",
                     full_module_name,
@@ -627,6 +648,7 @@ class DiffusionLoRAManager:
 
             scaled_lora_b = lora_weights.lora_b * scale
             lora_layer.set_lora(index=0, lora_a=lora_weights.lora_a, lora_b=scaled_lora_b)
+            _record_bound(lora_weights)
             logger.debug(
                 "Activated LoRA for %s: lora_a shape=%s, lora_b shape=%s, scale=%.2f",
                 full_module_name,
@@ -635,6 +657,33 @@ class DiffusionLoRAManager:
                 scale,
             )
 
+        if callable(binding_validator):
+            binding_validator(
+                lora_model=lora_model,
+                bound_lora_names=frozenset(bound_lora_names),
+            )
+
+    def _reset_lora_layers(self) -> None:
+        for lora_layer in self._lora_modules.values():
+            lora_layer.reset_lora(0)
+
+    def _activate_adapter(self, adapter_id: int, scale: float) -> None:
+        if self._is_active_at_scale(adapter_id, scale):
+            logger.debug("Adapter %d already active at scale %.3f skipping", adapter_id, scale)
+            return
+
+        logger.info("Activating adapter: id=%d", adapter_id)
+        lora_model = self._registered_adapters[adapter_id]
+        # Binding overwrites slot 0 incrementally. Invalidate the fast-path
+        # state before the first mutation and leave every wrapper inactive if
+        # any set_lora() call or model validator fails.
+        self._active_adapter_id = None
+        try:
+            self._bind_adapter_weights(lora_model, scale)
+        except Exception:
+            self._reset_lora_layers()
+            raise
+
         self._active_adapter_id = adapter_id
         self._update_adapter_scale(adapter_id, scale)
 
@@ -643,8 +692,7 @@ class DiffusionLoRAManager:
             logger.debug("All adapters already inactive")
             return
         logger.info("Deactivating all adapters: %d layers", len(self._lora_modules))
-        for lora_layer in self._lora_modules.values():
-            lora_layer.reset_lora(0)
+        self._reset_lora_layers()
         self._active_adapter_id = None
         logger.debug("All adapters deactivated")
 

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/lora.py

@@ -0,0 +1,205 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import math
+from pathlib import Path
+
+import torch
+from safetensors import safe_open
+from vllm.lora.lora_model import LoRAModel
+from vllm.lora.lora_weights import PackedLoRALayerWeights
+from vllm.lora.peft_helper import PEFTHelper
+from vllm.model_executor.models.utils import WeightsMapper
+
+from vllm_omni.lora.request import LoRARequest
+
+_TURBO_RANK = 128
+_TURBO_ALPHA = 128
+_TURBO_HIDDEN_SIZE = 5376
+_TURBO_ATTENTION_INNER_SIZE = 7168
+_TURBO_FFN_HIDDEN_SIZE = 14336
+_TURBO_FILENAME = "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"
+_LORA_A_SUFFIX = ".lora_A.default.weight"
+_LORA_B_SUFFIX = ".lora_B.default.weight"
+_TURBO_TARGETS = frozenset({"to_q", "to_k", "to_v", "out_proj", "fc1", "fc2"})
+_TURBO_RAW_TARGET_SUFFIXES = (
+    "attn.to_q",
+    "attn.to_k",
+    "attn.to_v",
+    "attn.to_out.0",
+    "ff.net.0.proj",
+    "ff.net.2",
+)
+_TURBO_TARGET_DIMS = {
+    "attn.to_q": (_TURBO_HIDDEN_SIZE, _TURBO_ATTENTION_INNER_SIZE),
+    "attn.to_k": (_TURBO_HIDDEN_SIZE, _TURBO_ATTENTION_INNER_SIZE),
+    "attn.to_v": (_TURBO_HIDDEN_SIZE, _TURBO_ATTENTION_INNER_SIZE),
+    "attn.to_out.0": (_TURBO_ATTENTION_INNER_SIZE, _TURBO_HIDDEN_SIZE),
+    "ff.net.0.proj": (_TURBO_HIDDEN_SIZE, 2 * _TURBO_FFN_HIDDEN_SIZE),
+    "ff.net.2": (_TURBO_FFN_HIDDEN_SIZE, _TURBO_HIDDEN_SIZE),
+}
+_TURBO_EXPECTED_RAW_TARGETS = frozenset(
+    f"{prefix}.{block_index}.{suffix}"
+    for prefix, block_count in (
+        ("transformer_blocks", 50),
+        ("token_refiner.refiner_blocks", 2),
+    )
+    for block_index in range(block_count)
+    for suffix in _TURBO_RAW_TARGET_SUFFIXES
+)
+_TURBO_TARGET_PATTERN = (
+    r"^transformer\.(?:token_refiner\.blocks|blocks)\.\d+\."
+    r"(?:attn\.(?:to_q|to_k|to_v|out_proj)|mlp\.(?:fc1|fc2))$"
+)
+
+_TURBO_WEIGHTS_MAPPER = WeightsMapper(
+    orig_to_new_substr={
+        "token_refiner.refiner_blocks.": "token_refiner.blocks.",
+        "transformer_blocks.": "blocks.",
+        ".attn.to_out.0.": ".attn.out_proj.",
+        ".ff.net.0.proj.": ".mlp.fc1.",
+        ".ff.net.2.": ".mlp.fc2.",
+        ".lora_A.default.": ".lora_A.",
+        ".lora_B.default.": ".lora_B.",
+    }
+)
+
+
+def _select_turbo_file(artifact_path: str | Path) -> Path | None:
+    path = Path(artifact_path)
+    if path.is_file():
+        return path if path.suffix == ".safetensors" else None
+    if not path.is_dir():
+        return None
+
+    candidate = path / _TURBO_FILENAME
+    return candidate if candidate.is_file() else None
+
+
+def _validate_and_convert_tensors(checkpoint) -> dict[str, torch.Tensor]:
+    tensors: dict[str, torch.Tensor] = {}
+    pairs: dict[str, set[str]] = {}
+    raw_targets: set[str] = set()
+    for name in checkpoint.keys():
+        if name.endswith(_LORA_A_SUFFIX):
+            raw_target = name[: -len(_LORA_A_SUFFIX)]
+            side = "a"
+        elif name.endswith(_LORA_B_SUFFIX):
+            raw_target = name[: -len(_LORA_B_SUFFIX)]
+            side = "b"
+        else:
+            raise ValueError(f"Unconsumed MiniMax-H3 Turbo tensor: {name!r}")
+        raw_targets.add(raw_target)
+
+        mapped_name = _TURBO_WEIGHTS_MAPPER.apply_list([name])[0]
+        mapped_target = mapped_name.rsplit(".lora_", 1)[0]
+        if mapped_target.rsplit(".", 1)[-1] not in _TURBO_TARGETS:
+            raise ValueError(f"Unsupported MiniMax-H3 Turbo target: {raw_target!r}")
+        target_sides = pairs.setdefault(mapped_target, set())
+        if side in target_sides:
+            raise ValueError(f"Duplicate MiniMax-H3 Turbo tensor for {mapped_target}.{side}")
+        target_sides.add(side)
+
+        tensor = checkpoint.get_tensor(name)
+        if tensor.ndim != 2:
+            raise ValueError(f"MiniMax-H3 Turbo LoRA tensors must be matrices, got {name}={tuple(tensor.shape)}")
+        suffix = next((suffix for suffix in _TURBO_RAW_TARGET_SUFFIXES if raw_target.endswith(suffix)), None)
+        if suffix is None:
+            raise ValueError(f"MiniMax-H3 Turbo LoRA contains unsupported target: {raw_target}")
+        input_dim, output_dim = _TURBO_TARGET_DIMS[suffix]
+        expected_shape = (_TURBO_RANK, input_dim) if side == "a" else (output_dim, _TURBO_RANK)
+        if tuple(tensor.shape) != expected_shape:
+            raise ValueError(
+                f"MiniMax-H3 Turbo tensor has invalid global shape: {name}={tuple(tensor.shape)}, "
+                f"expected={expected_shape}"
+            )
+        if side == "b" and ".ff.net.0.proj." in name:
+            value, gate = tensor.chunk(2, dim=0)
+            tensor = torch.cat((gate, value), dim=0).contiguous()
+        tensors[name] = tensor
+
+    incomplete = sorted(target for target, sides in pairs.items() if sides != {"a", "b"})
+    if incomplete:
+        raise ValueError(f"Incomplete MiniMax-H3 Turbo LoRA pairs: {incomplete}")
+    missing = sorted(_TURBO_EXPECTED_RAW_TARGETS - raw_targets)
+    unexpected = sorted(raw_targets - _TURBO_EXPECTED_RAW_TARGETS)
+    if missing or unexpected:
+        raise ValueError(
+            "MiniMax-H3 Turbo target set does not match the supported v1.0 artifact: "
+            f"missing={len(missing)} {missing[:5]}, unexpected={len(unexpected)} {unexpected[:5]}"
+        )
+    return tensors
+
+
+def _pack_h3_turbo_fc1(lora_model: LoRAModel) -> None:
+    """Represent H3's fused gate/up projection without generic layout guesses."""
+
+    for module_name, weights in tuple(lora_model.loras.items()):
+        if not module_name.endswith(".mlp.fc1"):
+            continue
+        gate_b, up_b = weights.lora_b.chunk(2, dim=0)
+        lora_model.loras[module_name] = PackedLoRALayerWeights(
+            module_name=module_name,
+            rank=weights.rank,
+            lora_alphas=[weights.lora_alpha, weights.lora_alpha],
+            lora_a=[weights.lora_a, weights.lora_a],
+            lora_b=[gate_b.contiguous(), up_b.contiguous()],
+            scaling=[weights.scaling, weights.scaling],
+        )
+
+
+def load_minimax_h3_turbo_lora(
+    *,
+    partition: str,
+    lora_request: LoRARequest,
+    lora_path: str | Path,
+    dtype: torch.dtype,
+    unsupported_offload_mode: str | None = None,
+) -> tuple[LoRAModel, PEFTHelper] | None:
+    """Load the published LightX2V Turbo v1.0 through the legacy manager."""
+
+    lora_file = _select_turbo_file(lora_path)
+    if lora_file is None:
+        return None
+    with safe_open(lora_file, framework="pt", device="cpu") as checkpoint:
+        metadata = checkpoint.metadata() or {}
+        if metadata.get("key_format") != "minimax-h3-diffusers":
+            if lora_file.name == _TURBO_FILENAME:
+                raise ValueError(
+                    "MiniMax-H3 Turbo v1.0 requires safetensors metadata key_format='minimax-h3-diffusers'"
+                )
+            return None
+        if lora_file.name != _TURBO_FILENAME:
+            raise ValueError(f"MiniMax-H3 Turbo supports only {_TURBO_FILENAME!r}, got {lora_file.name!r}")
+        raw_alpha = metadata.get("alpha")
+        try:
+            alpha = float(raw_alpha) if raw_alpha is not None else math.nan
+        except ValueError as exc:
+            raise ValueError(f"MiniMax-H3 Turbo alpha must be numeric, got {raw_alpha!r}") from exc
+        if alpha != _TURBO_ALPHA:
+            raise ValueError(f"MiniMax-H3 Turbo v1.0 requires alpha={_TURBO_ALPHA}, got {raw_alpha!r}")
+        if partition == "ref2va":
+            raise ValueError("MiniMax-H3 Turbo LoRA supports FL2VA/T2VA only")
+        if unsupported_offload_mode is not None:
+            raise ValueError(f"MiniMax-H3 Turbo dynamic LoRA does not support {unsupported_offload_mode}")
+        tensors = _validate_and_convert_tensors(checkpoint)
+
+    peft_helper = PEFTHelper.from_dict(
+        {
+            "r": _TURBO_RANK,
+            "lora_alpha": _TURBO_ALPHA,
+            "target_modules": _TURBO_TARGET_PATTERN,
+        }
+    )
+    lora_model = LoRAModel.from_lora_tensors(
+        lora_model_id=lora_request.lora_int_id,
+        tensors=tensors,
+        peft_helper=peft_helper,
+        device="cpu",
+        dtype=dtype,
+        weights_mapper=_TURBO_WEIGHTS_MAPPER,
+    )
+    _pack_h3_turbo_fc1(lora_model)
+    return lora_model, peft_helper

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,
         )
@@ -954,6 +1002,14 @@ class MiniMaxH3DiTModel(nn.Module):
     # quantization or LoRA to map onto. Address the fused layers directly, e.g.
     # ignored_layers=["blocks.0.attn.qkv_proj"].
     packed_modules_mapping = {}
+    # Turbo LoRA checkpoints publish separate Q/K/V adapters. This declaration
+    # lets the legacy diffusion LoRA manager bind them to the packed QKV layer;
+    # it does not change the fused base-checkpoint loading path above.
+    stacked_params_mapping = (
+        (".attn.qkv_proj", ".attn.to_q", "q"),
+        (".attn.qkv_proj", ".attn.to_k", "k"),
+        (".attn.qkv_proj", ".attn.to_v", "v"),
+    )
 
     def _validate_tp_config(self, *, arch: MiniMaxH3DiTArchConfig, tp_size: int) -> None:
         if tp_size < 1:
@@ -1137,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,
         *,
@@ -1152,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
@@ -1192,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)
@@ -1261,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"))
@@ -1297,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,
@@ -1331,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

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """vLLM-Omni pipeline for MiniMax H3 FL2VA and Ref2VA partitions."""
 
 from __future__ import annotations
@@ -11,7 +12,7 @@ from collections.abc import Iterable, Mapping, Sequence
 from contextlib import contextmanager, nullcontext
 from itertools import groupby
 from pathlib import Path
-from typing import Any, ClassVar
+from typing import TYPE_CHECKING, Any, ClassVar
 
 import numpy as np
 import torch
@@ -56,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,
 )
@@ -68,13 +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 .minimax_h3_transformer import MiniMaxH3DiTModel
+from .lora import load_minimax_h3_turbo_lora
+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,
@@ -101,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,
@@ -108,8 +124,18 @@ 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:
+    from vllm.lora.lora_model import LoRAModel
+    from vllm.lora.peft_helper import PEFTHelper
+
+    from vllm_omni.lora.request import LoRARequest
+
 MINIMAX_H3_FPS = 24
 MINIMAX_H3_AUDIO_SAMPLE_RATE = 32000
 MINIMAX_H3_IMGVID_COND_TIMESTEP = 0.999
@@ -130,6 +156,9 @@ MINIMAX_H3_MAX_REFERENCE_IMAGE_BYTES = 30 * 1024 * 1024
 MINIMAX_H3_REFERENCE_IMAGE_FORMATS = frozenset({"jpeg", "png", "webp", "heic", "heif"})
 MINIMAX_H3_MIN_OUTPUT_SECONDS = 4.0
 MINIMAX_H3_MAX_OUTPUT_SECONDS = 15.0
+MINIMAX_H3_TURBO_SIGMA_POINTS = 5
+MINIMAX_H3_TURBO_VIDEO_SHIFT = 6.0
+MINIMAX_H3_TURBO_AUDIO_SHIFT = 3.0
 MINIMAX_H3_DOWNLOAD_PATTERNS = [
     "FL2VA/**",
     "Ref2VA/model_index.json",
@@ -190,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)
@@ -433,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,
     *,
@@ -537,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"]
@@ -553,12 +698,89 @@ 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
     # uniform path available to partially constructed pipelines.
     _base_schedule_by_partition: ClassVar[Mapping[str, DMD2SigmaSchedule | None]] = {}
 
+    def _load_diffusion_lora_adapter(
+        self,
+        *,
+        lora_request: LoRARequest,
+        lora_path: str | Path,
+        dtype: torch.dtype,
+    ) -> tuple[LoRAModel, PEFTHelper] | None:
+        # A cache eviction may be followed by a different adapter reusing the
+        # same client-supplied ID. Every real load replaces the classification.
+        self._turbo_lora_adapter_ids.discard(lora_request.lora_int_id)
+        od_config = getattr(self, "od_config", None)
+        offload_modes = []
+        if getattr(od_config, "enable_cpu_offload", False):
+            offload_modes.append("model-level CPU offload (--enable-cpu-offload)")
+        if getattr(od_config, "enable_layerwise_offload", False):
+            offload_modes.append("layerwise offload (--enable-layerwise-offload)")
+        if getattr(od_config, "enable_distributed_layerwise_offload", False):
+            offload_modes.append("distributed layerwise offload (--enable-distributed-layerwise-offload)")
+        loaded = load_minimax_h3_turbo_lora(
+            partition=self.partition,
+            lora_request=lora_request,
+            lora_path=lora_path,
+            dtype=dtype,
+            unsupported_offload_mode=" or ".join(offload_modes) or None,
+        )
+        if loaded is not None:
+            self._turbo_lora_adapter_ids.add(lora_request.lora_int_id)
+        return loaded
+
+    def _validate_diffusion_lora_binding(
+        self,
+        *,
+        lora_model: LoRAModel,
+        bound_lora_names: frozenset[str],
+    ) -> None:
+        if lora_model.id not in self._turbo_lora_adapter_ids:
+            return
+        missing = sorted(set(lora_model.loras) - bound_lora_names)
+        if missing:
+            raise ValueError(
+                "MiniMax-H3 Turbo LoRA binding is incomplete: "
+                f"bound={len(bound_lora_names)}/{len(lora_model.loras)}, missing={missing[:5]}"
+            )
+
+    def _has_active_turbo_lora(self, sampling: Any) -> bool:
+        lora_request = sampling.lora_request
+        return (
+            lora_request is not None
+            and not math.isclose(0.0, float(sampling.lora_scale))
+            and lora_request.lora_int_id in self._turbo_lora_adapter_ids
+        )
+
+    def _validate_turbo_sampling(self, sampling: Any) -> None:
+        extra = sampling.extra_args or {}
+        sigma_points = sampling.num_inference_steps
+        if sigma_points != MINIMAX_H3_TURBO_SIGMA_POINTS:
+            raise OmniClientError(
+                "MiniMax-H3 Turbo requires num_inference_steps=5 (five sigma points produce four denoiser evaluations)"
+            )
+        try:
+            video_shift = float(extra.get("flow_shift", self.default_video_shift))
+        except (TypeError, ValueError) as exc:
+            raise OmniClientError(f"MiniMax-H3 Turbo requires flow_shift={MINIMAX_H3_TURBO_VIDEO_SHIFT:g}") from exc
+        if not math.isclose(video_shift, MINIMAX_H3_TURBO_VIDEO_SHIFT):
+            raise OmniClientError(f"MiniMax-H3 Turbo requires flow_shift={MINIMAX_H3_TURBO_VIDEO_SHIFT:g}")
+        try:
+            audio_shift = float(extra.get("audio_flow_shift", self.default_audio_shift))
+        except (TypeError, ValueError) as exc:
+            raise OmniClientError(
+                f"MiniMax-H3 Turbo requires audio_flow_shift={MINIMAX_H3_TURBO_AUDIO_SHIFT:g}"
+            ) from exc
+        if not math.isclose(audio_shift, MINIMAX_H3_TURBO_AUDIO_SHIFT):
+            raise OmniClientError(f"MiniMax-H3 Turbo requires audio_flow_shift={MINIMAX_H3_TURBO_AUDIO_SHIFT:g}")
+
     def adopt_cache_dit_backend(self, backend: CacheDiTBackend) -> None:
         """Adopt runner-installed generic Cache-DiT for request transitions."""
 
@@ -589,6 +811,7 @@ class MiniMaxH3Pipeline(
             getattr(od_config, "task_type", None),
             str(od_config.model),
         )
+        self._turbo_lora_adapter_ids: set[int] = set()
         model_root = _resolve_minimax_h3_model_root(
             str(od_config.model),
             od_config.revision,
@@ -781,6 +1004,8 @@ class MiniMaxH3Pipeline(
         self,
         requested: str | None,
         multi_modal_data: dict[str, Any],
+        *,
+        has_turbo_lora: bool = False,
     ) -> str:
         if requested is None:
             # A Ref2VA-only startup has no FL2VA transformer; preserve its
@@ -798,6 +1023,8 @@ class MiniMaxH3Pipeline(
             raise OmniClientError(
                 f"checkpoint partition {self.partition!r} supports {sorted(self.supported_tasks)}, got task={task!r}"
             )
+        if task == "ref2va" and has_turbo_lora:
+            raise OmniClientError("MiniMax-H3 Turbo LoRA supports T2VA/FL2VA requests only")
         return task
 
     def _resolve_shape(
@@ -1448,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,
@@ -1472,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,
@@ -1572,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,
@@ -1617,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,
@@ -1637,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] = {}
@@ -1650,12 +1959,32 @@ 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 {}
-        task = self._resolve_task(extra.get("task"), multi_modal_data)
+        has_turbo_lora = self._has_active_turbo_lora(sampling)
+        task = self._resolve_task(
+            extra.get("task"),
+            multi_modal_data,
+            has_turbo_lora=has_turbo_lora,
+        )
+        if has_turbo_lora:
+            self._validate_turbo_sampling(sampling)
 
         raw_image = multi_modal_data.get("image")
         raw_videos = multi_modal_data.get("video")
@@ -1721,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,
@@ -1863,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)
@@ -1900,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/models/wan2_2/pipeline_wan2_2.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from __future__ import annotations
 
@@ -33,7 +33,7 @@ from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin, _is_rank_zero
 from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler
 from vllm_omni.diffusion.models.wan2_2.scheduling_wan_euler import WanEulerScheduler
-from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanTransformer3DModel
+from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanSelfAttention, WanTransformer3DModel
 from vllm_omni.diffusion.postprocess import interpolate_video_tensor
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
@@ -44,6 +44,8 @@ from vllm_omni.platforms import current_omni_platform
 logger = logging.getLogger(__name__)
 DEBUG_PERF = False
 WAN_SAMPLE_SOLVER_CHOICES = {"unipc", "euler"}
+FASTWAN_DMD_TIMESTEPS = (1000.0, 757.0, 522.0)
+FASTWAN_DMD_SCHEDULER_SHIFT = 8.0
 
 
 def build_wan_scheduler(sample_solver: str, flow_shift: float) -> Any:
@@ -64,6 +66,30 @@ def build_wan_scheduler(sample_solver: str, flow_shift: float) -> Any:
     )
 
 
+def load_wan_weights_with_optional_gate(model: nn.Module, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+    """Load Wan weights and discard optional VSA gates absent from the checkpoint."""
+    gate_param_names = {name for name, _ in model.named_parameters() if ".to_gate_compress." in name}
+    has_gate_compress_weights = False
+
+    def tracked_weights():
+        nonlocal has_gate_compress_weights
+        for name, weight in weights:
+            if ".to_gate_compress." in name:
+                has_gate_compress_weights = True
+            yield name, weight
+
+    loaded_weights = AutoWeightsLoader(model).load_weights(tracked_weights())
+    setattr(model, "has_gate_compress_weights", has_gate_compress_weights)
+    if not has_gate_compress_weights:
+        for module in model.modules():
+            if isinstance(module, WanSelfAttention):
+                module.to_gate_compress = None
+
+    # Optional gate parameters are absent from public FullAttn/DMD checkpoints.
+    loaded_weights.update(gate_param_names)
+    return loaded_weights
+
+
 def resolve_wan_sample_solver(req: OmniDiffusionRequest, default: str = "unipc") -> str:
     extra_args = getattr(req.sampling_params, "extra_args", {}) or {}
     raw = extra_args.get("sample_solver", default)
@@ -314,6 +340,7 @@ class Wan22Pipeline(
 
         # Read model_index.json to detect expand_timesteps mode (for TI2V-5B)
         self.expand_timesteps = False
+        self.is_dmd = False
         self.has_transformer_2 = False
         if local_files_only:
             model_index_path = os.path.join(model, "model_index.json")
@@ -321,6 +348,7 @@ class Wan22Pipeline(
                 with open(model_index_path) as f:
                     model_index = json.load(f)
                     self.expand_timesteps = model_index.get("expand_timesteps", False)
+                    self.is_dmd = model_index.get("_class_name") == "WanDMDPipeline"
             # Check if this is a two-stage model (MoE with transformer_2)
             transformer_2_path = os.path.join(model, "transformer_2")
             self.has_transformer_2 = os.path.exists(transformer_2_path)
@@ -333,6 +361,7 @@ class Wan22Pipeline(
                 with open(model_index_path) as f:
                     model_index = json.load(f)
                     self.expand_timesteps = model_index.get("expand_timesteps", False)
+                    self.is_dmd = model_index.get("_class_name") == "WanDMDPipeline"
                     # Check transformer_2 from model_index
                     transformer_2_info = model_index.get("transformer_2", [None, None])
                     self.has_transformer_2 = transformer_2_info[0] is not None
@@ -422,6 +451,10 @@ class Wan22Pipeline(
         else:
             self.transformer_2 = None
 
+        for transformer in (self.transformer, self.transformer_2):
+            if transformer is not None:
+                transformer.preserve_vsa_all_blocks = self.is_dmd
+
         # Store the active transformer config
         if load_transformer:
             self.transformer_config = self.transformer.config
@@ -430,8 +463,14 @@ class Wan22Pipeline(
         else:
             raise RuntimeError("No transformer loaded")
 
-        self._sample_solver = "unipc"
-        self._flow_shift = od_config.flow_shift if od_config.flow_shift is not None else 5.0
+        self._sample_solver = "euler" if self.is_dmd else "unipc"
+        self._flow_shift = (
+            FASTWAN_DMD_SCHEDULER_SHIFT
+            if self.is_dmd
+            else od_config.flow_shift
+            if od_config.flow_shift is not None
+            else 5.0
+        )
         self.scheduler = build_wan_scheduler(self._sample_solver, self._flow_shift)
 
         self.vae_scale_factor_temporal = self.vae.config.scale_factor_temporal if getattr(self, "vae", None) else 4
@@ -480,6 +519,7 @@ class Wan22Pipeline(
         attention_kwargs: dict[str, Any],
         latent_condition: torch.Tensor | None = None,
         first_frame_mask: torch.Tensor | None = None,
+        generator: torch.Generator | None = None,
     ) -> torch.Tensor | AsyncLatents:
         if attention_kwargs is None:
             attention_kwargs = {}
@@ -561,7 +601,20 @@ class Wan22Pipeline(
                     cfg_normalize=False,
                 )
 
-                latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg)
+                if self.is_dmd:
+                    pred_clean = self.scheduler.predict_clean(noise_pred, latents, t).to(noise_pred.dtype)
+                    if step_idx + 1 < len(timesteps):
+                        noise = randn_tensor(
+                            latents.shape,
+                            generator=generator,
+                            device=latents.device,
+                            dtype=pred_clean.dtype,
+                        )
+                        latents = self.scheduler.add_noise(pred_clean, noise, timesteps[step_idx + 1])
+                    else:
+                        latents = pred_clean
+                else:
+                    latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg)
                 pbar.update()
 
         return latents
@@ -600,13 +653,21 @@ class Wan22Pipeline(
         mod_value = self.vae_scale_factor_spatial * patch_size[1]  # 16*2=32 for TI2V, 8*2=16 for I2V
         height = (height // mod_value) * mod_value
         width = (width // mod_value) * mod_value
-        num_steps = 40 if common.num_inference_steps is None else common.num_inference_steps
+        if self.is_dmd:
+            # The checkpoint was distilled for these three transitions. Ignore
+            # request-level step counts, including the engine's 1-step warmup.
+            num_steps = len(FASTWAN_DMD_TIMESTEPS)
+        else:
+            num_steps = 40 if common.num_inference_steps is None else common.num_inference_steps
 
         output_type = common.output_type or "np"
         num_outputs_per_prompt = common.num_outputs_per_prompt or 1
         attention_kwargs: dict | None = None
 
-        guidance_low, guidance_high = resolve_wan_guidance_scales(common, default_guidance_scale=4.0)
+        guidance_low, guidance_high = resolve_wan_guidance_scales(
+            common,
+            default_guidance_scale=1.0 if self.is_dmd else 4.0,
+        )
 
         # record guidance for properties
         self._guidance_scale = guidance_low
@@ -685,17 +746,19 @@ class Wan22Pipeline(
             current_omni_platform.synchronize()
             _t_text_enc_ms = (time.perf_counter() - _t_text_enc_start) * 1000
 
-        first_request = req.requests[0]
-        sample_solver = resolve_wan_sample_solver(first_request, default=self._sample_solver)
-        flow_shift = resolve_wan_flow_shift(first_request, self.od_config)
-        if sample_solver != self._sample_solver or abs(flow_shift - self._flow_shift) > 1e-6:
-            self.scheduler = build_wan_scheduler(sample_solver, flow_shift)
-            self._sample_solver = sample_solver
-            self._flow_shift = flow_shift
+        if self.is_dmd:
+            timesteps = torch.tensor(FASTWAN_DMD_TIMESTEPS, device=device, dtype=torch.float32)
+        else:
+            first_request = req.requests[0]
+            sample_solver = resolve_wan_sample_solver(first_request, default=self._sample_solver)
+            flow_shift = resolve_wan_flow_shift(first_request, self.od_config)
+            if sample_solver != self._sample_solver or abs(flow_shift - self._flow_shift) > 1e-6:
+                self.scheduler = build_wan_scheduler(sample_solver, flow_shift)
+                self._sample_solver = sample_solver
+                self._flow_shift = flow_shift
 
-        # Timesteps
-        self.scheduler.set_timesteps(num_steps, device=device)
-        timesteps = self.scheduler.timesteps
+            self.scheduler.set_timesteps(num_steps, device=device)
+            timesteps = self.scheduler.timesteps
         self._num_timesteps = len(timesteps)
         boundary_timestep = None
         if boundary_ratio is not None:
@@ -817,6 +880,7 @@ class Wan22Pipeline(
             attention_kwargs=attention_kwargs,
             latent_condition=latent_condition,
             first_frame_mask=first_frame_mask,
+            generator=generator,
         )
 
         # Wan2.2 is prone to out of memory errors when predicting large videos
@@ -1003,9 +1067,7 @@ class Wan22Pipeline(
         return latents
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        """Load weights using AutoWeightsLoader for vLLM integration."""
-        loader = AutoWeightsLoader(self)
-        return loader.load_weights(weights)
+        return load_wan_weights_with_optional_gate(self, weights)
 
     def check_inputs(
         self,

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from __future__ import annotations
 
@@ -16,7 +16,6 @@ import torchvision.transforms.functional as TF
 from diffusers.utils.torch_utils import randn_tensor
 from torch import nn
 from transformers import AutoTokenizer, CLIPImageProcessor, CLIPVisionModel, UMT5EncoderModel
-from vllm.model_executor.models.utils import AutoWeightsLoader
 from vllm.sequence import IntermediateTensors
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
@@ -36,6 +35,7 @@ from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
     build_wan_scheduler,
     create_transformer_from_config,
     load_transformer_config,
+    load_wan_weights_with_optional_gate,
     resolve_wan_flow_shift,
     resolve_wan_guidance_scales,
     resolve_wan_sample_solver,
@@ -1014,9 +1014,7 @@ class Wan22I2VPipeline(
             raise ValueError("`guidance_scale_2` is only supported when `boundary_ratio` is set.")
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        """Load weights using AutoWeightsLoader for vLLM integration."""
-        loader = AutoWeightsLoader(self)
-        return loader.load_weights(weights)
+        return load_wan_weights_with_optional_gate(self, weights)
 
 
 # ---------------------------------------------------------------------------

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """
 Wan2.2 Speech-to-Video (S2V) Pipeline for vLLM-Omni.
@@ -22,7 +22,6 @@ from diffusers.utils.torch_utils import randn_tensor
 from torch import nn
 from torchvision import transforms
 from transformers import AutoTokenizer, UMT5Config, UMT5EncoderModel, Wav2Vec2ForCTC, Wav2Vec2Processor
-from vllm.model_executor.models.utils import AutoWeightsLoader
 from vllm.multimodal.media.audio import load_audio
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
@@ -36,6 +35,7 @@ from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
 from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
     load_transformer_config,
+    load_wan_weights_with_optional_gate,
     retrieve_latents,
 )
 from vllm_omni.diffusion.models.wan2_2.wan2_2_s2v_transformer import (
@@ -1475,6 +1475,4 @@ class Wan22S2VPipeline(
         return outputs
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        """Load weights using AutoWeightsLoader for vLLM integration."""
-        loader = AutoWeightsLoader(self)
-        return loader.load_weights(weights)
+        return load_wan_weights_with_optional_gate(self, weights)

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """
 VACE (Video Creation and Editing) Pipeline for WAN models.
@@ -22,12 +22,12 @@ import PIL.Image
 import torch
 from vllm.logger import init_logger
 from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
-from vllm.model_executor.models.utils import AutoWeightsLoader
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.models.interface import SupportImageInput
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
     Wan22Pipeline,
+    load_wan_weights_with_optional_gate,
     retrieve_latents,
 )
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
@@ -747,6 +747,4 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput):
         )
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        """Load weights using AutoWeightsLoader for vLLM integration."""
-        loader = AutoWeightsLoader(self)
-        return loader.load_weights(weights)
+        return load_wan_weights_with_optional_gate(self, weights)

vllm_omni/diffusion/models/wan2_2/scheduling_wan_euler.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from __future__ import annotations
 
@@ -110,6 +110,42 @@ class WanEulerScheduler:
         self._step_index = None
         self._begin_index = None
 
+    def sigma_for_timestep(
+        self,
+        timestep: float | torch.Tensor,
+        *,
+        device: torch.device | str | None = None,
+        dtype: torch.dtype = torch.float32,
+    ) -> torch.Tensor:
+        """Return the flow sigma nearest to a (possibly distilled) timestep."""
+        timestep_t = torch.as_tensor(timestep, device=self.timesteps.device, dtype=self.timesteps.dtype)
+        if timestep_t.numel() != 1:
+            raise ValueError(f"Expected a scalar timestep, got shape {tuple(timestep_t.shape)}")
+        index = torch.argmin(torch.abs(self.timesteps - timestep_t))
+        return self.sigmas[index].to(device=device or self.timesteps.device, dtype=dtype)
+
+    def predict_clean(
+        self,
+        model_output: torch.Tensor,
+        sample: torch.Tensor,
+        timestep: float | torch.Tensor,
+    ) -> torch.Tensor:
+        """Convert a flow/noise prediction into the corresponding clean latent."""
+        sample_fp32 = sample.to(torch.float32)
+        sigma = self.sigma_for_timestep(timestep, device=sample.device, dtype=torch.float32)
+        return sample_fp32 - sigma * model_output.to(torch.float32)
+
+    def add_noise(
+        self,
+        clean_sample: torch.Tensor,
+        noise: torch.Tensor,
+        timestep: float | torch.Tensor,
+    ) -> torch.Tensor:
+        """Re-noise a clean latent at a distilled flow timestep."""
+        sigma = self.sigma_for_timestep(timestep, device=noise.device, dtype=torch.float32)
+        noisy = (1.0 - sigma) * clean_sample.to(torch.float32) + sigma * noise.to(torch.float32)
+        return noisy.to(noise.dtype)
+
     def scale_model_input(self, sample: torch.Tensor, timestep: int | None = None) -> torch.Tensor:  # noqa: ARG002
         return sample
 

vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 import math
 from collections.abc import Iterable
@@ -419,6 +419,24 @@ class WanSelfAttention(nn.Module):
             prefix=prefix,
         )
 
+        # FastVideo VSA checkpoints may add a learned projection that gates
+        # the compressed global branch. Zero initialization makes checkpoints
+        # without these weights sparse-only, with no user-facing mode switch.
+        self.to_gate_compress: ColumnParallelLinear | None = None
+        if self.attn.attn_backend.get_name() == "FASTVIDEO_VSA":
+            self.to_gate_compress = ColumnParallelLinear(
+                dim,
+                self.inner_dim,
+                bias=True,
+                gather_output=False,
+                return_bias=False,
+                quant_config=quant_config,
+                prefix=f"{prefix}.to_gate_compress" if prefix else "to_gate_compress",
+            )
+            nn.init.zeros_(self.to_gate_compress.weight)
+            if self.to_gate_compress.bias is not None:
+                nn.init.zeros_(self.to_gate_compress.bias)
+
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -441,6 +459,15 @@ class WanSelfAttention(nn.Module):
         key = key.unflatten(2, (self.num_kv_heads, self.head_dim))
         value = value.unflatten(2, (self.num_kv_heads, self.head_dim))
 
+        if self.to_gate_compress is not None:
+            gate_result = self.to_gate_compress(hidden_states)
+            gate_compress = gate_result[0] if isinstance(gate_result, tuple) else gate_result
+            gate_compress = gate_compress.unflatten(2, (self.num_heads, self.head_dim))
+            if attn_metadata is None:
+                attn_metadata = AttentionMetadata(extra={"gate_compress": gate_compress})
+            else:
+                attn_metadata.extra["gate_compress"] = gate_compress
+
         # Apply rotary embeddings
         if rotary_emb is not None:
             freqs_cos, freqs_sin = rotary_emb
@@ -709,6 +736,8 @@ class WanTransformerBlock(nn.Module):
         temb: torch.Tensor,
         rotary_emb: tuple[torch.Tensor, torch.Tensor],
         hidden_states_mask: torch.Tensor | None = None,
+        vsa_dit_seq_shape: tuple[int, int, int] | None = None,
+        preserve_vsa_all_blocks: bool = False,
     ) -> torch.Tensor:
         if temb.ndim == 4:
             # temb: batch_size, seq_len, 6, inner_dim (wan2.2 ti2v)
@@ -729,7 +758,12 @@ class WanTransformerBlock(nn.Module):
 
         # 1. Self-attention
         norm_hidden_states = self.norm1(hidden_states, scale_msa, shift_msa).type_as(hidden_states)
-        self_attn_metadata = AttentionMetadata(attn_mask=hidden_states_mask)
+        self_attn_extra = {}
+        if vsa_dit_seq_shape is not None:
+            self_attn_extra["vsa_dit_seq_shape"] = vsa_dit_seq_shape
+        if preserve_vsa_all_blocks:
+            self_attn_extra["preserve_vsa_all_blocks"] = True
+        self_attn_metadata = AttentionMetadata(attn_mask=hidden_states_mask, extra=self_attn_extra)
         attn_output = self.attn1(norm_hidden_states, rotary_emb, self_attn_metadata)
         hidden_states = (hidden_states + attn_output * gate_msa).type_as(hidden_states)
 
@@ -910,6 +944,9 @@ class WanTransformer3DModel(nn.Module):
             pos_embed_seq_len=pos_embed_seq_len,
         )
 
+        # DMD/FastVideo checkpoints retain VSA semantics when top-k selects every block.
+        self.preserve_vsa_all_blocks = False
+
         # 3. Transformer blocks — partitioned across PP stages via vLLM's `make_layers`.
         # It computes the [start_layer, end_layer) slice for this rank and fills the remaining slots
         # with PPMissingLayer so that weight names stay globally consistent.
@@ -1044,8 +1081,19 @@ class WanTransformer3DModel(nn.Module):
             )
 
         # Transformer blocks
+        # Preserve the post-patch (T, H, W) grid so VSA can partition
+        # the flattened DiT sequence into spatiotemporal blocks.
+        vsa_dit_seq_shape = (post_patch_num_frames, post_patch_height, post_patch_width)
         for block in self.blocks[self.start_layer : self.end_layer]:
-            hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, rotary_emb, hidden_states_mask)
+            hidden_states = block(
+                hidden_states,
+                encoder_hidden_states,
+                timestep_proj,
+                rotary_emb,
+                hidden_states_mask,
+                vsa_dit_seq_shape,
+                self.preserve_vsa_all_blocks,
+            )
 
         if not is_pipeline_last_stage():
             # Non-last PP stage: hand the token sequence to the caller via IntermediateTensors.

vllm_omni/diffusion/models/wan2_2/wan2_2_vace_transformer.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """VACE variant of WanTransformer3DModel for conditional video generation."""
 
@@ -61,6 +61,7 @@ class VaceWanTransformerBlock(WanTransformerBlock):
         temb: torch.Tensor,
         rotary_emb: tuple[torch.Tensor, torch.Tensor],
         hidden_states_mask: torch.Tensor | None = None,
+        vsa_dit_seq_shape: tuple[int, int, int] | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         if self.proj_in is not None:
             control_hidden_states = self.proj_in(control_hidden_states)
@@ -72,6 +73,7 @@ class VaceWanTransformerBlock(WanTransformerBlock):
             temb,
             rotary_emb,
             hidden_states_mask,
+            vsa_dit_seq_shape,
         )
 
         conditioning_states = self.proj_out(control_hidden_states)
@@ -250,6 +252,7 @@ class WanVACETransformer3DModel(WanTransformer3DModel):
             )
 
         # VACE: embed context and run conditioning blocks
+        vsa_dit_seq_shape = (post_patch_num_frames, post_patch_height, post_patch_width)
         vace_hints = None
         if vace_context is not None and self.vace_blocks is not None:
             full_seq_len = hidden_states.shape[1] * sp_size
@@ -263,6 +266,7 @@ class WanVACETransformer3DModel(WanTransformer3DModel):
                     timestep_proj,
                     rotary_emb,
                     hidden_states_mask,
+                    vsa_dit_seq_shape,
                 )
                 vace_hints.append(conditioning_states)
 
@@ -278,6 +282,7 @@ class WanVACETransformer3DModel(WanTransformer3DModel):
                 timestep_proj,
                 rotary_emb,
                 hidden_states_mask,
+                vsa_dit_seq_shape,
             )
             if vace_hints is not None and self.vace_layers_mapping is not None and i in self.vace_layers_mapping:
                 vace_idx = self.vace_layers_mapping[i]

vllm_omni/diffusion/registry.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 import importlib
 
@@ -61,6 +61,11 @@ _DIFFUSION_MODELS = {
         "pipeline_wan2_2",
         "Wan22Pipeline",
     ),
+    "WanDMDPipeline": (
+        "wan2_2",
+        "pipeline_wan2_2",
+        "Wan22Pipeline",
+    ),
     "WanVACEPipeline": (
         "wan2_2",
         "pipeline_wan2_2_vace",
@@ -459,6 +464,7 @@ def _apply_sequence_parallel_if_enabled(model, od_config: OmniDiffusionConfig) -
 
     try:
         sp_size = od_config.parallel_config.sequence_parallel_size
+        assert sp_size is not None
         if sp_size <= 1:
             return
 
@@ -539,6 +545,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
     "OvisImagePipeline": "get_ovis_image_post_process_func",
     "BooguImagePipeline": "get_boogu_image_post_process_func",
     "WanPipeline": "get_wan22_post_process_func",
+    "WanDMDPipeline": "get_wan22_post_process_func",
     "WanVACEPipeline": "get_wan22_vace_post_process_func",
     "LTX2Pipeline": "get_ltx2_post_process_func",
     "LTX2TwoStagePipeline": "get_ltx2_post_process_func",
@@ -612,6 +619,7 @@ _DIFFUSION_PRE_PROCESS_FUNCS = {
     "LongCatVideoAvatarPipeline": "get_longcat_video_avatar_pre_process_func",
     "QwenImageLayeredPipeline": "get_qwen_image_layered_pre_process_func",
     "WanPipeline": "get_wan22_pre_process_func",
+    "WanDMDPipeline": "get_wan22_pre_process_func",
     "WanVACEPipeline": "get_wan22_vace_pre_process_func",
     "WanImageToVideoPipeline": "get_wan22_i2v_pre_process_func",
     "WanS2VPipeline": "get_wan22_s2v_pre_process_func",
@@ -703,6 +711,7 @@ def register_diffusion_model(
 
 def _load_process_func(od_config: OmniDiffusionConfig, func_name: str):
     """Load and return a process function from the appropriate module."""
+    assert od_config.model_class_name is not None
     mod_folder, mod_relname, _ = _DIFFUSION_MODELS[od_config.model_class_name]
     if mod_relname == "":
         # Full module path (registered via register_diffusion_model)

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/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py

@@ -95,9 +95,6 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         #   been read").
         self._held_non_active: deque[Any] = deque()
         self.requests_num_chunks_sent: dict[str, int] = defaultdict(int)
-        # Boundary tasks advance this before the background sender runs, so a
-        # late old-segment frame cannot restore the previous watermark.
-        self._segment_generation: dict[str, int] = defaultdict(int)
         self._pending_streaming_prefills: dict[str, dict] = {}
         # Monotonic timestamp of when each request last began waiting for a
         # chunk, refreshed every time one arrives.  Read by
@@ -191,31 +188,11 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
             is_segment_finished: whether the segment of request is finished
         """
         is_finished = request.is_finished() and not request.resumable
-        if not hasattr(self, "_segment_generation"):
-            self._segment_generation = defaultdict(int)
-        external_req_id = request.external_req_id
-        raw_generation = getattr(request, "_omni_segment_generation", 0)
-        try:
-            generation = int(raw_generation)
-        except (TypeError, ValueError):
-            generation = 0
-        expected_generation = self._segment_generation.get(external_req_id, generation)
-        if generation < expected_generation:
-            logger.warning(
-                "Skip late save_async for request %s, segment_generation=%s, expected=%s",
-                external_req_id,
-                generation,
-                expected_generation,
-            )
-            return
-        self._segment_generation[external_req_id] = generation
 
         confirmed_num_computed_tokens = self._confirmed_num_computed_tokens(request)
 
         # If the request is preempted, skip the already saved chunks.
-        if not is_segment_finished and confirmed_num_computed_tokens < self.requests_num_chunks_sent.get(
-            request.external_req_id, 0
-        ):
+        if confirmed_num_computed_tokens < self.requests_num_chunks_sent.get(request.external_req_id, 0):
             logger.warning(
                 f"Enqueue save_async for request {request.external_req_id}, "
                 f"request.num_computed_tokens={request.num_computed_tokens}, "
@@ -224,18 +201,13 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
             )
             return
 
+        self.requests_num_chunks_sent[request.external_req_id] = confirmed_num_computed_tokens
         task = {
             "multimodal_output": multimodal_output,
             "request": request,
             "is_finished": is_finished,
             "is_segment_finished": is_segment_finished,
-            "segment_generation": generation,
         }
-        if is_segment_finished:
-            self._segment_generation[external_req_id] = generation + 1
-            self.requests_num_chunks_sent.pop(external_req_id, None)
-        else:
-            self.requests_num_chunks_sent[external_req_id] = confirmed_num_computed_tokens
         self._pending_save_reqs.append(task)
         with self._save_cond:
             self._save_cond.notify()
@@ -453,6 +425,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
 
         if is_segment_finished:
             self.code_prompt_token_ids.pop(external_req_id, None)
+            self.requests_num_chunks_sent.pop(external_req_id, None)
             self.ramp_chunk_count.pop(external_req_id, None)
             cached_ic = getattr(self, "_cached_ic", None)
             if cached_ic is not None:
@@ -528,7 +501,6 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         self.request_payload.pop(external_req_id, None)
         self.code_prompt_token_ids.pop(external_req_id, None)
         self.requests_num_chunks_sent.pop(external_req_id, None)
-        self._segment_generation.pop(external_req_id, None)
         self.ramp_chunk_count.pop(external_req_id, None)
         self._pending_streaming_prefills.pop(external_req_id, None)
 

vllm_omni/engine/arg_utils.py

@@ -1,9 +1,13 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 import argparse
 import json
 import os
 import tempfile
+from collections.abc import Callable
 from dataclasses import dataclass, field, fields
-from typing import Any
+from typing import Any, cast
 
 from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs
 from vllm.logger import init_logger
@@ -298,24 +302,26 @@ class OmniEngineArgs(EngineArgs):
         }
         stage_connector_config["extra"]["stage_id"] = self.stage_id
 
+        hf_overrides = cast(dict[str, Any] | Callable[[Any], Any] | None, getattr(self, "hf_overrides", None))
         # If model_arch is specified, inject it into hf_overrides so vLLM can
         # resolve the architecture even when config.json lacks 'architectures'.
         # Also inject model_type so AutoConfig can resolve the correct config
         # class for models with empty or missing config.json (e.g. CosyVoice3).
         if self.model_arch:
-            if self.hf_overrides is None:
-                self.hf_overrides = {}
-            if isinstance(self.hf_overrides, dict):
-                self.hf_overrides.setdefault("architectures", [self.model_arch])
-                if "model_type" not in self.hf_overrides:
+            if hf_overrides is None:
+                hf_overrides = {}
+                self.hf_overrides = hf_overrides
+            if isinstance(hf_overrides, dict):
+                hf_overrides.setdefault("architectures", [self.model_arch])
+                if "model_type" not in hf_overrides:
                     model_type = _ARCH_TO_MODEL_TYPE.get(self.model_arch)
                     if model_type is not None:
-                        self.hf_overrides.setdefault("model_type", model_type)
+                        hf_overrides.setdefault("model_type", model_type)
 
                 # Stage wrappers (e.g. Code2Wav) may need max_model_len larger
                 # than the base checkpoint's text max_position_embeddings.
                 if self.model_arch == "Qwen3TTSCode2Wav" and self.max_model_len is not None:
-                    self.hf_overrides.setdefault("talker_config", {}).setdefault(
+                    hf_overrides.setdefault("talker_config", {}).setdefault(
                         "max_position_embeddings", int(self.max_model_len)
                     )
 
@@ -331,33 +337,37 @@ class OmniEngineArgs(EngineArgs):
                 if model_type is not None:
                     self._patch_empty_hf_config(model_type)
 
+        tokenizer = cast(str | None, getattr(self, "tokenizer", None))
+
         # NemotronVoiceChat: the checkpoint's only tokenizer subfolder is the
         # auxiliary rnnt_tokenizer/ (ASR loss vocab), which the generic
         # subfolder auto-detect below would wrongly pick up. The text tokenizer
         # is the Nemotron backbone's, resolved from a local override or its
         # HF id (cache-friendly).
-        if not self.tokenizer and self.model_arch in (
+        if not tokenizer and self.model_arch in (
             "NemotronVoiceChatThinkerForConditionalGeneration",
             "NemotronVoiceChatTalkerForConditionalGeneration",
             "NemotronVoiceChatCode2Wav",
         ):
             self.tokenizer = os.environ.get("NEMOTRON_VOICECHAT_LLM_PATH") or "nvidia/NVIDIA-Nemotron-Nano-9B-v2"
-            logger.info("NemotronVoiceChat: using text tokenizer from %s", self.tokenizer)
+            tokenizer = self.tokenizer
+            logger.info("NemotronVoiceChat: using text tokenizer from %s", tokenizer)
 
         # Auto-detect tokenizer for models that store it in a subdirectory
         # rather than the root (e.g. CosyVoice3 uses CosyVoice-BlankEN/).
-        if not self.tokenizer and self.model:
+        if not tokenizer and self.model:
             model_path = self.model
             if os.path.isdir(model_path) and not os.path.isfile(os.path.join(model_path, "tokenizer_config.json")):
                 for subfolder in sorted(os.listdir(model_path)):
                     candidate = os.path.join(model_path, subfolder)
                     if os.path.isdir(candidate) and os.path.isfile(os.path.join(candidate, "tokenizer_config.json")):
                         self.tokenizer = candidate
+                        tokenizer = candidate
                         logger.info("Auto-detected tokenizer at %s", candidate)
                         break
             elif not os.path.isdir(model_path):
-                subfolder = _TOKENIZER_SUBFOLDER_MAP.get(self.model_arch)
-                if subfolder:
+                tokenizer_subfolder = _TOKENIZER_SUBFOLDER_MAP.get(self.model_arch or "")
+                if tokenizer_subfolder:
                     # Download just the tokenizer files from the subfolder
                     try:
                         from huggingface_hub import snapshot_download
@@ -365,15 +375,15 @@ class OmniEngineArgs(EngineArgs):
                         local_dir = snapshot_download(
                             model_path,
                             allow_patterns=[
-                                f"{subfolder}/tokenizer*",
-                                f"{subfolder}/tokenization_*",
-                                f"{subfolder}/special_tokens*",
-                                f"{subfolder}/vocab*",
-                                f"{subfolder}/merges*",
-                                f"{subfolder}/added_tokens*",
+                                f"{tokenizer_subfolder}/tokenizer*",
+                                f"{tokenizer_subfolder}/tokenization_*",
+                                f"{tokenizer_subfolder}/special_tokens*",
+                                f"{tokenizer_subfolder}/vocab*",
+                                f"{tokenizer_subfolder}/merges*",
+                                f"{tokenizer_subfolder}/added_tokens*",
                             ],
                         )
-                        candidate = os.path.join(local_dir, subfolder)
+                        candidate = os.path.join(local_dir, tokenizer_subfolder)
                         if os.path.isdir(candidate):
                             self.tokenizer = candidate
                             logger.info("Downloaded tokenizer from %s/%s", model_path, subfolder)
@@ -532,6 +542,7 @@ class OrchestratorArgs:
     hsdp_shard_size: int = -1
     hsdp_replicate_size: int = 1
     diffusion_attention_backend: str | None = None
+    fastvideo_vsa_topk: int | None = None
     diffusion_attention_config: str | None = None
     diffusion_compile_granularity: str | None = None
     diffusion_compile_dynamic: bool | None = None

vllm_omni/engine/async_omni_engine.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """
 Async Omni Engine for vLLM-Omni multi-stage runtime.
 
@@ -47,6 +50,7 @@ from vllm_omni.engine.async_engine_utils import (
 )
 from vllm_omni.engine.messages import (
     AbortRequestMessage,
+    AbortResultMessage,
     AddCompanionRequestMessage,
     CollectiveRPCRequestMessage,
     CollectiveRPCResultMessage,
@@ -375,9 +379,13 @@ class AsyncOmniEngine:
                     getattr(self, "_duplex_runtime_extension_path", None)
                 )
                 if duplex_runtime_extension is not None:
+                    stage_clients = []
+                    for pool in self.stage_pools:
+                        assert pool.stage_client is not None
+                        stage_clients.append(pool.stage_client)
                     validate_duplex_runtime_extension(
                         duplex_runtime_extension,
-                        sampling_defaults=tuple(pool.stage_client.default_sampling_params for pool in self.stage_pools),
+                        sampling_defaults=tuple(client.default_sampling_params for client in stage_clients),
                     )
 
             orchestrator = Orchestrator(
@@ -685,6 +693,7 @@ class AsyncOmniEngine:
             )
 
             # Full input processing (tokenization, multimodal, etc.)
+            assert self.input_processor is not None
             _t_preprocess = time.perf_counter()
             try:
                 request = self.input_processor.process_inputs(
@@ -763,11 +772,13 @@ class AsyncOmniEngine:
             Exception: Whatever prompt expansion or input processing raised.
                 The caller is expected to let it reach the client.
         """
+        assert self.prompt_expand_func is not None
         expanded = self.prompt_expand_func(original_prompt, stage0_params)
         if not expanded:
             return []
 
         companions: list[AddCompanionRequestMessage] = []
+        assert self.input_processor is not None
         for ep in expanded:
             cid = f"{parent_id}{ep.request_id_suffix}"
             companion_prompt = ep.prompt
@@ -864,7 +875,7 @@ class AsyncOmniEngine:
                 "mag_max_skip_steps": 5,
                 "mag_retention_ratio": 0.1,
             }
-        if cache_backend in ("step_cache"):
+        if cache_backend == "step_cache":
             return {
                 "step_cache_dit_enabled": True,
                 "velocity_sim_thresholds": [0.95, 0.93],
@@ -1018,10 +1029,12 @@ class AsyncOmniEngine:
         if (
             kwargs.get("diffusion_attention_config") is not None
             or kwargs.get("diffusion_attention_backend") is not None
+            or kwargs.get("fastvideo_vsa_topk") is not None
         ):
             attention_config = parse_attention_config(
                 kwargs.get("diffusion_attention_config"),
                 attention_backend=kwargs.get("diffusion_attention_backend"),
+                fastvideo_vsa_topk=kwargs.get("fastvideo_vsa_topk"),
             )
 
         stage_engine_args = {
@@ -1223,6 +1236,7 @@ class AsyncOmniEngine:
                 if (
                     kwargs.get("diffusion_attention_config") is not None
                     or kwargs.get("diffusion_attention_backend") is not None
+                    or kwargs.get("fastvideo_vsa_topk") is not None
                 ):
                     has_stage_attention = (
                         hasattr(cfg.engine_args, "diffusion_attention_config")
@@ -1232,6 +1246,7 @@ class AsyncOmniEngine:
                         cfg.engine_args.diffusion_attention_config = parse_attention_config(
                             kwargs.get("diffusion_attention_config"),
                             attention_backend=kwargs.get("diffusion_attention_backend"),
+                            fastvideo_vsa_topk=kwargs.get("fastvideo_vsa_topk"),
                         )
                 quantization_config = kwargs.get("diffusion_quantization_config") or kwargs.get("quantization_config")
                 if quantization_config is not None:
@@ -1722,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,
@@ -198,38 +199,18 @@ class OrchestratorRequestState:
     running_counter_registered: bool = False
 
 
-@dataclass
-class StreamingSegmentState:
-    """Streaming segment-boundary state for one stage of a request.
-
-    Every stage of a multi-stage pipeline reaches its own segment boundaries
-    independently, so this is tracked per stage rather than per request.
-    """
-
-    # Flag of segment of streaming input finished
-    finished: bool = False
-    # Tokens from the current raw segment boundary. The vLLM output processor
-    # does not guarantee that EngineCoreOutput.new_token_ids survives on the
-    # processed RequestOutput used by the routing layer.
-    token_ids: list[int] = field(default_factory=list)
-    output_metadata: dict[str, Any] = field(default_factory=dict)
-
-
 @dataclass
 class StreamingInputState:
     # Flag of streaming input request
     enabled: bool = False
-    # Keyed by stage. ``_orchestration_loop`` records a boundary per stage poll
-    # into this one ``req_state``, while every consumer reads it for a specific
-    # ``stage_id``, so a flat slot is only correct as long as each read happens in
-    # the same iteration as its own stage's write. It is not for a stage that
-    # never records one (the ``poll_diffusion_output`` branch reaches
-    # ``_handle_processed_outputs`` without the raw-output loop) or for any
-    # consumer reading outside the writing iteration.
-    segments: dict[int, StreamingSegmentState] = field(default_factory=dict)
-    # Streaming update prompt length. NOT stage-keyed: its only consumer reads it
-    # through the streaming context on the stage-input-bridge path, not by stage
-    # id. Do not assume the rest of this struct is stage-safe.
+    # Flag of segment of streaming input finished
+    segment_finished: bool = False
+    # Tokens from the current raw segment boundary. The vLLM output processor
+    # does not guarantee that EngineCoreOutput.new_token_ids survives on the
+    # processed RequestOutput used by the routing layer.
+    segment_token_ids: list[int] = field(default_factory=list)
+    segment_output_metadata: dict[str, Any] = field(default_factory=dict)
+    # Streaming update prompt length
     new_prompt_len_snapshot: int | None = None
     # Model/bridge-specific runtime states (e.g., thinker->talker)
     bridge_states: dict[str, Any] = field(default_factory=dict)
@@ -237,14 +218,6 @@ class StreamingInputState:
     # while the downstream input processor consumes upstream token output.
     source_token_decoder: Callable[..., str] | None = None
 
-    def segment(self, stage_id: int | None) -> StreamingSegmentState:
-        """Return ``stage_id``'s segment state, or a fresh empty one if unreported.
-
-        Fresh rather than shared: ``output_metadata`` is handed out by reference.
-        """
-        segment = self.segments.get(stage_id)
-        return segment if segment is not None else StreamingSegmentState()
-
 
 class _OrchestratorDuplexStagePort:
     """Adapts generic stage pools to the model-neutral duplex control plane."""
@@ -845,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."""
@@ -994,18 +989,17 @@ class Orchestrator:
                                 req_state = self.request_states.get(getattr(eco, "request_id", None))
                                 if req_state is None or not req_state.streaming.enabled:
                                     continue
-                                segment_finished = bool(getattr(eco, "is_segment_finished", False))
+                                req_state.streaming.segment_finished = bool(getattr(eco, "is_segment_finished", False))
+                                req_state.streaming.segment_token_ids = (
+                                    self._coerce_int_list(getattr(eco, "new_token_ids", None))
+                                    if req_state.streaming.segment_finished
+                                    else []
+                                )
                                 raw_mm = self._completion_multimodal_output(eco, None)
-                                req_state.streaming.segments[stage_id] = StreamingSegmentState(
-                                    finished=segment_finished,
-                                    token_ids=(
-                                        self._coerce_int_list(getattr(eco, "new_token_ids", None))
-                                        if segment_finished
-                                        else []
-                                    ),
-                                    output_metadata=(
-                                        dict(raw_mm) if segment_finished and isinstance(raw_mm, dict) else {}
-                                    ),
+                                req_state.streaming.segment_output_metadata = (
+                                    dict(raw_mm)
+                                    if req_state.streaming.segment_finished and isinstance(raw_mm, dict)
+                                    else {}
                                 )
                                 req_state.streaming.new_prompt_len_snapshot = getattr(
                                     eco,
@@ -1073,7 +1067,7 @@ class Orchestrator:
                 continue
 
             stage_metrics = None
-            segment_finished = req_state.streaming.enabled and req_state.streaming.segment(stage_id).finished
+            segment_finished = req_state.streaming.enabled and req_state.streaming.segment_finished
             if output.finished or segment_finished:
                 stage_metrics = pool.build_stage_metrics(
                     [output],
@@ -1407,7 +1401,6 @@ class Orchestrator:
         req_id = output.request_id
         finished = output.finished
         submit_ts = req_state.stage_submit_ts.get(stage_id)
-        segment_finished = req_state.streaming.enabled and req_state.streaming.segment(stage_id).finished
         # CFG companion: stash output so parent can bundle [parent, *companions]
         # into source_outputs for the bridge (e.g. thinker2imagegen).
         if finished and self._cfg_tracker.is_companion(req_id):
@@ -1417,7 +1410,11 @@ class Orchestrator:
             return
 
         request_finished = False
-        if finished and self.stage_pools[stage_id].final_output and not segment_finished:
+        if (
+            finished
+            and self.stage_pools[stage_id].final_output
+            and not (req_state.streaming.enabled and req_state.streaming.segment_finished)
+        ):
             req_state.finished_final_output_stage_ids.add(stage_id)
             final_output_stage_ids = req_state.final_output_stage_ids or {req_state.final_stage_id}
             request_finished = final_output_stage_ids.issubset(req_state.finished_final_output_stage_ids)
@@ -1429,7 +1426,7 @@ class Orchestrator:
         # filter out again (the official implementation returns exactly one
         # result per audio chunk).
         is_duplex_stage0_segment = (
-            stage_id == 0 and self._is_duplex_session_request(req_state) and req_state.streaming.segment(0).finished
+            stage_id == 0 and self._is_duplex_session_request(req_state) and req_state.streaming.segment_finished
         )
         if self.stage_pools[stage_id].final_output and not is_duplex_stage0_segment:
             await self.output_async_queue.put(
@@ -1441,10 +1438,7 @@ class Orchestrator:
                     metrics=stage_metrics,
                     finished=(
                         request_finished
-                        or (
-                            self._is_duplex_session_request(req_state)
-                            and req_state.streaming.segment(stage_id).finished
-                        )
+                        or (self._is_duplex_session_request(req_state) and req_state.streaming.segment_finished)
                     ),
                     stage_submit_ts=submit_ts,
                 )
@@ -1479,7 +1473,7 @@ class Orchestrator:
             return
 
         if (
-            (finished or segment_finished)
+            (finished or (req_state.streaming.enabled and req_state.streaming.segment_finished))
             and stage_id < req_state.final_stage_id
             and (not self.async_chunk or not self._stage_receives_async_chunks(stage_id + 1))
             and (not self._next_stage_already_submitted(stage_id, req_state) or req_state.streaming.enabled)
@@ -1618,9 +1612,7 @@ class Orchestrator:
     def _duplex_output_context(
         req_state: OrchestratorRequestState,
         *,
-        # Required: segment state is stage-keyed, so a caller that omitted this
-        # would silently read "no boundary" rather than the intended stage's.
-        stage_id: int | None,
+        stage_id: int | None = None,
     ) -> DuplexOutputContext | None:
         identity = req_state.duplex_identity
         if identity is None:
@@ -1631,16 +1623,15 @@ class Orchestrator:
         )
 
         fence = req_state.duplex_stage_fences.get(stage_id, identity.fence) if stage_id is not None else identity.fence
-        segment = req_state.streaming.segment(stage_id)
         return DuplexOutputContext(
             identity=DuplexRequestIdentity(
                 session_id=identity.session_id,
                 fence=fence,
             ),
             final_stage_id=req_state.final_stage_id,
-            segment_finished=req_state.streaming.enabled and segment.finished,
-            segment_token_ids=tuple(segment.token_ids),
-            segment_output_metadata=segment.output_metadata,
+            segment_finished=req_state.streaming.enabled and req_state.streaming.segment_finished,
+            segment_token_ids=tuple(req_state.streaming.segment_token_ids),
+            segment_output_metadata=req_state.streaming.segment_output_metadata,
         )
 
     @staticmethod

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:

vllm_omni/entrypoints/cli/serve.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """
 Omni serve command for vLLM-Omni.
 
@@ -606,6 +609,12 @@ class OmniServeCommand(CLISubcommand):
             "May be combined with --diffusion-attention-config.per_role.* overrides, "
             "but mutually exclusive with --diffusion-attention-config.default.backend.",
         )
+        omni_config_group.add_argument(
+            "--fastvideo-vsa-topk",
+            type=int,
+            default=None,
+            help="Number of key/value blocks selected per query block by FASTVIDEO_VSA.",
+        )
         omni_config_group.add_argument(
             "--diffusion-attention-config",
             "-dac",

vllm_omni/model_executor/models/wan2_2/pipeline.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Wan2.2 TI2V pipeline topology."""
 
 from vllm_omni.config.stage_config import (
@@ -12,6 +12,7 @@ WAN2_2_TI2V_PIPELINE = PipelineConfig(
     model_type="wan2_2_ti2v",
     model_arch="WanPipeline",
     diffusers_class_name="WanPipeline",
+    diffusers_class_aliases=("WanDMDPipeline",),
     stages=(
         StagePipelineConfig(
             stage_id=0,