AB-Attention speeds up video generation by up to 2.23× and skips 92% of attention

The InfRec team, together with researchers from CUHK Shenzhen, Shandong University, the University of Tokyo, Fudan University, Berkeley, HKUST and NUS, has released AB-Attention. It is a way to speed up video generation in diffusion transformers (Wan2.1 and HunyuanVideo) without fine-tuning the weights.

The most expensive operation in these models is attention: every piece of the video is compared against every other piece across space and time, so cost grows quadratically with sequence length. AB-Attention skips 89–92% of that compute. The resulting video stays closer to "honest" full attention than the output of competing methods.

PSNR (a metric of similarity to the reference clip, in decibels; higher means closer) is 1.6–4.4 dB better, and end-to-end generation on a single NVIDIA H20 GPU runs 1.

84–2.23 times faster.

Where earlier sparse schemes break down

Sparse attention computes only some of the query–key pairs, not all of them. Early methods fixed the pattern in advance, for example local 3D windows. They cannot know what matters for a particular prompt, attention head or denoising step, so they have to keep a lot of unnecessary work. Dynamic methods choose key and value blocks based on current activations, but the authors count three weak spots in them:

• Importance estimates are crude or expensive. Averaging a block into a single vector mixes tokens that behave differently, and global clustering (SVG2) requires k-means over the whole sequence plus token permutation. • Top-k selection cuts off the tail. A discarded block disappears from both the numerator of attention and the softmax normalization. Better scoring won't fix this: even ideal top-k selection at 10.8% of keys retains only about 69% of the attention mass. • The unit of selection is a compute tile. A block is kept or dropped as a whole, even when only a few of its keys are needed.

AB-Attention's three techniques

The first technique, IBCS (Intra-Block Clustering Selector), estimates the importance of key blocks. Queries are split into blocks of 128 tokens. Within each block they are grouped into 8 clusters, each with a centroid and a weight equal to the share of queries in that cluster. A key block's score is the weighted attention of the centroids to its averaged keys. The clusters are used only for estimation; the actual computation runs in the original token order. As a result, a block that only one cluster attends to is not lost, as it would be if the whole block were averaged. The recall gap between block averaging and ideal top-k is cut in half.

The second technique, SIP (Sampling with Inclusion-Probability Weighting), replaces hard selection with sampling. A block enters the sample with a probability proportional to its score. The most important blocks are always taken, and the rest have a nonzero chance. Each selected block is multiplied by the inverse of its inclusion probability: a block with a 25% probability counts as four. This is the classic Horvitz–Thompson estimator, and it yields unbiased estimates of both the numerator and the softmax normalization. Attention error comes out 41% lower than with ideal top-k at the same number of blocks. No extra correction pass is needed: the computation remains ordinary softmax attention with a logit shift.

The third technique is fine-grained selection on "chunky" hardware. Key and value blocks are only 4 tokens each, but the kernel gathers the selected blocks by index into full-size tensor-core tiles. Tokens are not permuted and no extra copies are made. Shrinking the block from 128 to 4 tokens raises the share of attention mass retained by ideal selection from 59% to 69%, while the number of tiles at a fixed budget does not grow. The sampling plan is recomputed once every four denoising steps; the attention itself runs at every step.

How it was measured and what came out

The authors compared three models at 720p: two versions of Wan2.1 (81 frames) and HunyuanVideo (129 frames), on 600 prompts from the Penguin Video benchmark. The reference was a video generated with full attention using the same prompt and generation seed. Besides PSNR, they computed SSIM and LPIPS, plus five scores of the clip itself: aesthetics, background consistency, image quality, motion smoothness and subject consistency. The competitors were RadialAttention (a static mask with energy decay), SVG2 and DFSAttn, all running on the official code and the FlashInfer backend with FlashAttention-3. For every method, the first transformer block and the first 10 of 50 steps run with full attention.

• Sparsity: 89.16% on Wan2.1 and 92.35% on HunyuanVideo. • Competitors compute 1.8–3.6 times more attention. • AB-Attention is best on PSNR, SSIM and LPIPS across all three models. • End-to-end speedup: 1.84–2.23× on a single H20 in BF16.

What this means in practice

The method needs no fine-tuning, so it can be plugged into existing weights. For teams running Wan2.1 or HunyuanVideo, that means almost half the time per clip at quality close to the reference. The gain will grow on longer and higher-resolution videos, where attention accounts for a larger share of latency. There are limitations: the measurements were made on an H20 and on three models, and the kernel and sampling plan are designed for NVIDIA tensor cores. The code is open on GitHub in the InfRec-Cardinal/AB-Attention repository, so the results can be rechecked on your own hardware.