AB-Attention speeds up video diffusion transformers up to 2.23x, no training

Video diffusion transformers produce high-quality video, but the quadratic cost of 3D full attention dominates their inference. One common remedy is dynamic block-sparse attention, which computes exact attention only on key/value blocks chosen from the current activations. The paper's authors say this approach loses fidelity when sparsity becomes aggressive, for three reasons: 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 this, the paper proposes Accurate Block Attention (AB-Attention), a training-free method that approximates dense attention accurately at high sparsity. It has three parts. First, the Intra-Block Clustering Selector (IBCS) estimates block importance by clustering queries only within each query block, without permuting tokens; the authors describe this as faithful yet cheap. Second, Sampling with Inclusion-Probability Weighting (SIP) samples blocks in proportion to those importance estimates and reweights them by inverse inclusion probability, so that a single pass estimates dense attention rather than truncating it. Third, the kernel gathers key/value blocks much smaller than a tile into full tensor-core tiles, which refines selection at no extra tensor-core cost.
On Wan2.1 and HunyuanVideo, AB-Attention skips 89 to 92% of attention, while the baselines compute 1.8 to 3.6 times as much. Even so, it improves PSNR by up to 4.4 dB over the strongest baseline and reaches up to 2.23x end-to-end speedup. The code is published at github.com/InfRec-Cardinal/AB-Attention.
Key facts
- AB-Attention is a training-free block-sparse attention method for video diffusion transformers, aimed at keeping dense-attention quality at high sparsity.
- It combines an Intra-Block Clustering Selector (IBCS), Sampling with Inclusion-Probability Weighting (SIP), and a kernel that packs small key/value blocks into full tensor-core tiles.
- On Wan2.1 and HunyuanVideo it skips 89 to 92% of attention; the baselines compute 1.8 to 3.6 times as much.
- Reported results: up to 4.4 dB higher PSNR than the strongest baseline and up to 2.23x end-to-end speedup.
- Code is available on GitHub at InfRec-Cardinal/AB-Attention.
Why it matters
Full 3D attention is the main inference cost in video diffusion transformers, so any method that skips most of it without hurting quality directly targets the slowest part of video generation. The paper's pitch is that earlier block-sparse methods degrade as sparsity rises, and AB-Attention is built to hold up at 89 to 92% skipped attention. The sampling step is the key idea: instead of dropping low-ranked blocks, it samples blocks by importance and reweights them so the result estimates dense attention.
Who it affects
Teams that run or serve video diffusion transformers such as Wan2.1 and HunyuanVideo, the two models evaluated in the paper, and researchers working on sparse attention and inference efficiency for video generation. Because the method is training-free, it is framed as something applied to existing models rather than a new model to adopt.
How to use it
The authors publish code at https://github.com/InfRec-Cardinal/AB-Attention. The method is described as training-free, so no retraining is part of the approach. The abstract describes a custom kernel that packs small key/value blocks into full tensor-core tiles, so the speedup relies on that kernel and not only on the selection logic.
How solid is it
These are the authors' own reported results from a paper abstract, and the headline numbers are "up to" figures. The abstract names no authors or institutions. No hardware (GPU model), resolution, video length or batch size is given for the speedup and PSNR measurements. It does not say which baselines are compared, nor which baseline is the strongest. It does not say which of the two models gives the 4.4 dB or the 2.23x figure. No quality metrics other than PSNR are mentioned, and no absolute PSNR values or latency numbers are given.
Risks and caveats
The 4.4 dB and 2.23x figures are maxima, not typical results, and cannot be tied to a specific model or setup from the abstract. PSNR is the only quality measure mentioned, so perceived video quality is not covered by the stated evidence. The speedup is end-to-end but its conditions are unstated, so results on other hardware or settings may differ. Independent replication is not mentioned in the source.
“a training-free method that accurately approximates dense attention at high sparsity”
— From the paper's abstract describing AB-Attention