Improving Video Sparse Attention with Fine-grained Router and Sparse Rebasing
Authors: Peiyuan Zhang, Guoqiang Wei, Yilong Zhao, Zixiang Zhang, Wei Zhou, Will Lin, Heng Zhang, Xiaonan Nie, Yan Zeng, Hao Zhang
Organizations: University of California, San Diego · Work done during an internship at ByteDance Seed · ByteDance Seed · University of California, Berkeley · Georgia Institute of Technology
We present VSA2, a frontier trainable sparse attention for video DiTs. VSA2 includes a variety of new architectural features and training procedures that we apply across all stages of the DiT development cycle, including pretraining, RL, and inference, to produce a DiT with comparable or better quality than a full attention counterpart. Architecturally, VSA2 introduces a fine-grained router that improves the precision of identifying critical tokens and supports dynamic computation by allowing each query to attend to a variable number of key-value pairs. In training, we identify a Hard-to-Easy Curriculum, where models trained under high sparsity and later evaluated with lower sparsity during inference not only generalize effectively, but also outperform models trained with full attention in motion quality. VSA2 is also flexible: it can replace full attention during the middle of progressive low-to-high resolution pretraining, rebasing early-stage full-attention checkpoints. Experiments show that VSA2 reduces attention computation by half over VSA with lower loss. On 720p videos, it accelerates attention by 8.9x and end-to-end generation by 4.62x compared to the FlashAttention-3 baseline, while achieving comparable or better video quality.
Figures & tables
Figure 1Figure 2
Exp
Stage
Configuration
T2V
I2V
Train TopK
Infer TopK
Attn Sparsity
E2E Speedup
Motion
Following
Aesthetics
Motion
Following
1
480p
64
64
0.90
2.09 ×
-1.34%
1.34%
-0.67%
8.72%
1.34%
2
480p
64
128
0.82
1.92 ×
22.1%
-3.36%
-4.7%
6.71%
-8.72%
3
480p
32
128
0.82
1.92 ×
7.38%
-15.4%
-4.7%
11.4%
-9.4%
4
480p
32
64
0.90
2.09 ×
12.1%
-10.1%
-7.38%
4.03%
-4.03%
5
480p RL
64
64
0.90
2.09 ×
-4.03%
1.34%
-2.01%
8.05%
-4.7%
Table 2: Human evaluation results of VSA2 against full attention across various stages of video DiT training. A positive score means the model trained with VSA2 has better quality than its full attention counterpart at that stage.
Figure 4: Training loss and reward scores of VSA2 vs. full attention across various stages of video DiT training. With a sparsity ratio of 90% to 95%, VSA2 shows lower loss and similar reward scores compared to full attention.
Figure 5
Figure 7: Ablation studies of various design choices in video sparse attention. All models are trained from scratch except those in plot (f), which are initialized from a 256p video checkpoint.
Figure 8: Speed comparison on a single H800 GPU with batch size 1, number of heads 20, head dimension 128, and top64.
Appendix figures & tables3 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 9: More text-to-video samples of VSA2 vs. full attention.
Figure 10: More image-to-video samples of VSA2 vs. full attention.
Figure 11: Trained on 5–12 s videos, VSA2 can directly generate 30 s videos.
Diffusion Transformers (DiTs) enable high-quality video generation but suffer from substantial inference latency, primarily attributable to the computationally expensive full spatio-temporal attention. While sparse attention methods offer potential solutions, existing approaches face an inherent flexibility--efficiency dilemma: predefined masks lack the flexibility to capture diverse attention patterns, while runtime-determined masks introduce overheads and sacrifice hardware efficiency. We identify the lack of a unified structural characterization of DiT attention as a key limitation of existing methods, and establish that video DiT attention exhibits periodic diagonal stripe structures along both temporal and spatial dimensions. To formally encode these structured patterns within a single efficient kernel, we present {\bf PSA}, a parameterized stripe attention that formalizes the observed stripe regularity, unifying diverse attention patterns for efficient mask generation. This unified representation enables a single hardware-efficient CUDA kernel to process all sparse patterns, achieving FlashAttention-3-level Model FLOPs Utilization. To determine optimal sparsity configurations, we propose a training-free offline search algorithm that automatically maximizes sparsity under a specified error tolerance for each attention head. Experiments on HunyuanVideo and Wan~2.1 demonstrate that PSA achieves 1.57× and 1.37× end-to-end speedups over FlashAttention-3 baselines, with acceptable visual quality degradation.
Video diffusion transformers (vDiTs) generate high quality but pay quadratic self-attention cost, making inference prohibitive at video-token scales. The challenge is input-adaptive sparsity: selecting critical Q/K/V tokens with negligible overhead and executing them for end-to-end gains. We present SPADE, a training-free sparse-attention engine of three parts: (i) vDiT-SSR, a specification defining 3D blocking candidates and formalizing dynamic masks via Summarizer/Estimator expressions; (ii) runtime scheme generation using SICS and a head-wise policy; and (iii) an executor with low-overhead index search, flash block-sparse attention, and kernel grouping. Across Hunyuan-Video and Wan 2.1/2.2 for text-to-video and image-to-video generation, SPADE raises sparsity and speed while preserving quality, accelerating attention by 2.26x-3.40x and end-to-end inference by 1.49x-1.80x. Our code is open-sourced at https://github.com/6somehow/DAC-SPADE.
Diffusion Transformers (DiTs) dominate video generation but suffer from high inference cost due to the quadratic complexity of self-attention. Sparse attention mitigates this by skipping unnecessary computation in the attention matrix, but existing methods apply a fixed sparsity ratio throughout denoising and struggle to accurately determine which computations to skip. To address these limitations, we analyze attention dynamics through attention entropy and mean attention distance, and find that attention becomes increasingly concentrated across denoising steps. We further observe that blocks selected early continue to capture a large fraction of the total attention weight at later steps. These findings suggest two acceleration opportunities: the sparsity ratio should grow as denoising progresses, and earlier search results can be reused to reduce the search space. Based on these findings, we propose Growing Sparsity and Reduced Search (GSRS), a training-free framework that retains more computation early and progressively reduces it later, while restricting subsequent searches to an early candidate set to reduce search overhead. GSRS adopts Top-p selection for adaptive sparsity control, supported by a precise evaluation scheme that computes exact attention weights within the FlashAttention computation at negligible overhead. Experiments on HunyuanVideo and Wan 2.1 validate that GSRS achieves 1.55×--2.21× speedup while maintaining generation quality comparable to full attention.