AB-Attention: Accurate Block Attention for Accelerating Video Diffusion Transformers
Authors: Bin Xu, Mo Zhang, Zhengfeiyang Zhang, Xinfeng Yang, Yixing Ma, Wei Wang, HeXin Wang, Pengfei Zhou, Wangbo Zhao
Organizations: InfRec · The Chinese University of Hong Kong, Shenzhen · Shandong University · The University of Tokyo · Fudan University · University of California, Berkeley · The Hong Kong University of Science and Technology · National University of Singapore
Video diffusion transformers generate high-quality videos, but the quadratic cost of 3D full attention dominates their inference. Dynamic block-sparse attention reduces this cost by computing exact attention only on key/value blocks selected from current activations. However, its fidelity degrades under aggressive sparsity: blocks are scored from coarse summaries or formed by costly global clustering, pruned blocks vanish from the softmax, and whole tiles are kept or dropped even when only a few of their keys matter. To address these limitations, we propose Accurate Block Attention (AB-Attention), a training-free method that accurately approximates dense attention at high sparsity. Its Intra-Block Clustering Selector (IBCS) estimates block importance faithfully yet cheaply by clustering queries only within each query block, without permuting tokens. Sampling with Inclusion-Probability Weighting (SIP) then samples blocks in proportion to these estimates and reweights them by inverse inclusion probability, so that one pass estimates dense attention rather than truncating it. Finally, its kernel gathers key/value blocks much smaller than a tile into full tensor-core tiles, refining selection at no extra tensor-core cost. On Wan2.1 and HunyuanVideo, AB-Attention skips 89--92% of attention, whereas the baselines compute 1.8--3.6× as much, yet it still improves PSNR by up to 4.4~dB over the strongest baseline and achieves up to 2.23× end-to-end speedup. Our code is available at https://github.com/InfRec-Cardinal/AB-Attention.
Figures & tables
Figure 2: Attention mass, recall, and error on Wan2.1-T2V-14B at 89.16% sparsity. (a) Sorted cumulative attention mass of 4-token key blocks; the band spans the 10th–90th percentiles of query blocks. (b) Top- k recall: the fraction of each query block’s dense attention mass covered by the selected key blocks. Finer blocks keep more mass, and IBCS halves the gap of mean-pool to the oracle at b=4 . (c) Relative ℓ2 error of the attention output against dense attention. Sampling with SIP on IBCS scores lowers the error 41% below oracle top- k at b=4 .
Figure 3: Overview of AB-Attention. (a) IBCS clusters each query block Qi into centroids Qˉi with weights wi , and scores every key/value block by the weighted attention of the centroids to its pooled keys. (b) SIP samples each block with probability πij proportional to its score Sij , so no block is excluded, and weights it by 1/πij to estimate full attention. (c) The kernel gathers the sampled blocks into full tensor-core tiles and attends to them with logits shifted by −logπij .
Method
Perceptual Metrics
VBench
Efficiency
PSNR ↑
SSIM ↑
LPIPS ↓
AQ ↑
BC ↑
IQ ↑
MS ↑
SC ↑
Sparsity ↑
Latency (s) ↓
Speedup ↑
Wan2.1 1.3B
∞
1.0000
0.0000
59.11
95.96
65.48
98.99
95.96
0%
962.39
1.00×
Radial
27.856
0.8672
0.1704
59.65
95.97
66.56
99.02
95.63
73.63%
549.91
1.75×
DFSAttn
27.734
0.8508
0.2193
57.83
95.78
67.30
98.93
95.45
81.03%
548.75
1.75×
SVG2
30.021
0.8858
0.1576
54.59
94.82
58.25
98.90
93.84
80.12%
536.47
1.79×
AB-Attention
34.436
0.9328
0.0734
58.02
95.75
63.92
98.96
95.84
89.16%
479.11
2.01×
Table 1: Quality and efficiency comparison of AB-Attention and baselines on Wan2.1-T2V-1.3B , Wan2.1-T2V-14B , and HunyuanVideo-T2V-13B . Shaded rows report the full-attention reference.
Figure 4: Qualitative comparison on Wan2.1-T2V-14B . Panels (a)–(d) are four Penguin Benchmark prompts, abbreviated with ellipses. Within a panel, rows are methods and columns are frames 12, 40, and 68 of the same video. Headers show sparsity and full-video PSNR relative to the dense reference.
Table 5
Appendix figures & tables1 asset
Supplementary material from the paper’s appendix.
Appendix
Block size
PSNR ↑
SSIM ↑
LPIPS ↓
4
29.705
0.9154
0.0956
8
29.619
0.9150
0.0956
16
29.455
0.9108
0.0985
32
29.201
0.9092
0.1015
64
29.069
0.9072
0.1032
128
28.686
0.9006
0.1113
Appendix
Table 3: Fidelity to dense attention for each selection block size on Wan2.1-T2V-14B (100 prompts). Every row executes 8192 key tokens per query block in 128-token tiles (89.16% sparsity). Block size 4 is the full AB-Attention, and block size 128 is tile-level selection.
Diffusion transformers have achieved remarkable success in high-quality video generation, yet their reliance on spatiotemporal 3D full attention incurs prohibitive computational cost due to the quadratic complexity of attention. Block sparse attention is a common approach to mitigate this by focusing computation on important regions. However, attention maps in DiTs exhibit inherently dynamic and fine-grained sparsity, which causes existing block sparse attention methods to degrade significantly in quality, especially at high sparsity ratios. In this paper, we revisit block sparse attention and derive a theoretical lower bound on attention recall to characterize the key factors governing its effectiveness. Guided by these insights, we propose DFSAttn, a training-free sparse attention framework that enables dynamic, fine-grained sparsification efficiently. DFSAttn incorporates three core designs: Hilbert curve-based token reordering to achieve fine-grained sparsity while preserving efficient GPU execution, hierarchical block scoring for accurate block importance estimation, and sparse mask caching with adaptive ratios to balance accuracy and efficiency. Experimental results demonstrate that DFSAttn consistently outperforms prior methods under high sparsity, achieving up to 2.1× end-to-end speedup while maintaining high generation quality. Our code is open-sourced and available at https://github.com/jessica-hujie/DFSAttn.
Diffusion Transformers (DiTs) have become a dominant architecture for video generation, but their efficiency is limited by the quadratic complexity of full attention. Sparse attention reduces this cost by retrieving important blocks and computing attention only within them, but inaccurate retrieval can either degrade generation quality or yield unnecessary computation. We identify two retrieval mismatches in methods that retrieve blocks using the averaged representations of query and key blocks: (i) query-side aggregation mismatch, where averaging queries before Softmax fails to preserve their individual attention preferences, and (ii) key-side clustering metric mismatch, where standard Euclidean clustering in the original key space can group keys with dissimilar QK scores under the current query, so their average representation may not accurately represent how the current query scores individual keys. These mismatches can lead to inaccurate block retrieval. To address these mismatches, we propose PARK, a training-free sparse attention method for accurate block retrieval. PARK retains every original query, independently normalizes its attention over key blocks, and then averages these distributions within each query block. It also uses information from the current queries to transform keys before clustering, so that keys receiving similar QK scores are grouped together. A fused GPU kernel further reduces the overhead of block retrieval. Experiments on HunyuanVideo and Wan demonstrate that PARK improves block retrieval accuracy and preserves generation quality while accelerating inference, achieving the best quality-efficiency trade-off among the compared sparse attention methods.
Full 3D attention is costly in video Diffusion Transformers. Block-sparse attention reduces this cost by computing only selected attention blocks. However, existing methods either overlook the potential to optimize sparsity allocation or rely on heuristic rules or local metrics to guide it. In this work, we observe that sparsifying attention in different layers and at different denoising steps affects final-video fidelity to substantially different degrees. Yet the identities of the most sensitive layer-step pairs remain nearly unchanged across inputs. Building on these findings, we propose Allocating Layer-step Latency by Output Tests (ALLOT), a training-free sparsity allocation method. Treating each layer-step pair as a cell, ALLOT estimates each cell's contribution to final-video error through (i) cell probing and (ii) trajectory calibration, then jointly optimizes cell densities under a measured-latency budget. It preserves existing selection strategies and kernels, replacing only the allocation component. At equal or shorter wall-clock time, ALLOT outperforms every compared allocation strategy with XAttention, with an 11.8 dB PSNR gain on HunyuanVideo. Extensive experiments across four additional published pipelines on Wan2.1, Wan2.2, and HunyuanVideo show that ALLOT outperforms or ties the default in all nine settings.