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

AB-Attention. Accurate Block Attention for Accelerating Video Diffusion Transformers. By Bin Xu and eight others, from InfRec, the Chinese University of Hong Kong, Shenzhen, and six other institutions. Published on the twenty-ninth of October.

Video diffusion transformers produce high-quality video, but their inference is dominated by one cost: full 3D attention, which grows quadratically. Dynamic block-sparse attention is the usual remedy. It computes exact attention only on the key and value blocks selected from the current activations, and skips the rest. The trouble arrives when sparsity gets aggressive, and the abstract names three ways fidelity degrades. Blocks are scored from coarse summaries, or formed by global clustering that is itself costly. Blocks that get pruned simply vanish from the softmax. And whole tiles are kept or dropped together, even when only a few of their keys actually matter.

Accurate Block Attention, or AB-Attention, is a training-free method aimed at each of those three failures, with one goal: approximate dense attention accurately while skipping most of it. The central move is to stop treating sparse attention as truncation. Instead of computing a subset and discarding the remainder, it treats the selected blocks as a sample, so that a single pass yields an estimate of the dense result.

There are three pieces. The Intra-Block Clustering Selector, IBCS, estimates how important each block is by clustering queries only within each query block, without permuting tokens, which keeps the estimate faithful yet cheap. Sampling with Inclusion-Probability Weighting, SIP, then draws blocks in proportion to those estimates and reweights each by its inverse inclusion probability. That reweighting is what turns a pruned computation into an estimate of dense attention. Finally, the kernel gathers key and value blocks much smaller than a tile and packs them into full tensor-core tiles, so selection becomes finer grained at no extra tensor-core cost.

Tests run on Wan2.1 and HunyuanVideo. AB-Attention skips eighty-nine to ninety-two percent of attention, while the baselines compute between one point eight and three point six times as much.

Despite doing less work, it improves PSNR by up to four point four decibels over the strongest baseline, and reaches up to two point two three times end-to-end speedup.

The usual trade in sparse attention is speed against fidelity, with fidelity losing as sparsity climbs. Here the claim is that the loss comes from how blocks are chosen and how the dropped ones are handled, and that fixing selection, weighting and tile granularity lets a sparser computation land closer to dense attention than denser competitors do.