[FEAT]Add FastVideo VSA backend for Wan2.2 (#4820)

Signed-off-by: bjf-frz <frz123db@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

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.

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/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/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/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():

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/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/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/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/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/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.
 
@@ -375,9 +378,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 +692,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 +771,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 +874,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 +1028,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 +1235,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 +1245,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:

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,