AB-Attention: Accurate Block Attention for Accelerating Video Diffusion Transformers

Bin Xu; Mo Zhang; Zhengfeiyang Zhang; Xinfeng Yang; Yixing Ma; Wei Wang; HeXin Wang; Pengfei Zhou

cs.CVOct 29, 2026

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,

Organizations: InfRec · The Chinese University of Hong Kong, Shenzhen · Shandong University · The University of Tokyo ·

Abstract

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 $\times$ as much, yet it still improves PSNR by up to 4.4~dB over the strongest baseline and achieves up to 2.23 $\times$ 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

\ell_{2}

error of the attention output against dense attention. Sampling with SIP on IBCS scores lowers the error 41% below oracle top-

k

b{=}4

Figure 3: Overview of AB-Attention. (a) IBCS clusters each query block

Q_{i}

into centroids

\bar{Q}_{i}

with weights

w_{i}

, and scores every key/value block by the weighted attention of the centroids to its pooled keys. (b) SIP samples each block with probability

\pi_{ij}

proportional to its score

S_{ij}

, so no block is excluded, and weights it by

1/\pi_{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\pi_{ij}

Method	Perceptual Metrics			VBench					Efficiency
Method	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$	AQ $\uparrow$	BC $\uparrow$	IQ $\uparrow$	MS $\uparrow$	SC $\uparrow$	Sparsity $\uparrow$	Latency (s) $\downarrow$	Speedup $\uparrow$
Wan2.1 1.3B	$\infty$	1.0000	0.0000	59.11	95.96	65.48	98.99	95.96	0%	962.39	$1.00\times$
Radial	27.856	0.8672	0.1704	59.65	95.97	66.56	99.02	95.63	73.63%	549.91	$1.75\times$
DFSAttn	27.734	0.8508	0.2193	57.83	95.78	67.30	98.93	95.45	81.03%	548.75	$1.75\times$
SVG2	30.021	0.8858	0.1576	54.59	94.82	58.25	98.90	93.84	80.12%	536.47	$1.79\times$
AB-Attention	34.436	0.9328	0.0734	58.02	95.75	63.92	98.96	95.84	89.16%	479.11	$2.01\times$

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.

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Block size	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$
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 Model Inference Acceleration

Diffusion model inference acceleration reduces the time and compute needed to generate or reconstruct outputs while preserving quality.

See all 509 papers

DFSAttn: Dynamic Fine-grained Sparse Attention for Efficient Video Generation
May 22, 2026Jie Hu, Zixiang Gao, Yutong He +1Video Diffusion Models Diffusion Transformer
1Peking University.
PARK: Accurate Block Retrieval for Sparse Attention in Video Diffusion Transformers
Sep 30, 2026Yun Dai, Jiarui Wen, Huiping Zhuang +2Video Diffusion Models Diffusion Transformer
South China University of Technology
Sparsity Should Be ALLOTted by Output: Training-Free Allocation of Block-Sparse Attention for Video Diffusion Transformers
Oct 21, 2026Tianze Xia, Xiaoling Zhou, Ruya Lan +3Sparse Attention Diffusion Model Inference Acceleration
UCAS · CASIA · KlingAI Research +1

May 22, 2026J/K move · Enter open · S save

DFSAttn: Dynamic Fine-grained Sparse Attention for Efficient Video Generation

Jie Hu, Zixiang Gao, Yutong He, Kun Yuan

1Peking University.

Video Diffusion Models Diffusion Transformer Transformer Attention Sparse Attention Diffusion Model Inference Acceleration Block-Sparse Attention

Open paper AskPDF arXiv

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

\times

end-to-end speedup while maintaining high generation quality. Our code is open-sourced and available at https://github.com/jessica-hujie/DFSAttn.

Abstract

Diffusion Model Inference Acceleration

Explore similar work

DFSAttn: Dynamic Fine-grained Sparse Attention for Efficient Video Generation

PARK: Accurate Block Retrieval for Sparse Attention in Video Diffusion Transformers

Sparsity Should Be ALLOTted by Output: Training-Free Allocation of Block-Sparse Attention for Video Diffusion Transformers