[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>

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