[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>
This commit is contained in:
bjf-frzandHongsheng Liu authored and GitHub committed 2026-08-24 14:48:31 +08:00
1 parent 3d035bfa19
commit 072bfc02dd
28 files changed
+1316 -107

No files matched your search

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