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

AI summary

Overview

Research area: Computer Vision — efficient inference for video diffusion transformers (sparse attention acceleration).

Technical level: Advanced. The paper combines a statistical estimator (Horvitz–Thompson sampling with inclusion-probability weights) with a fused GPU attention kernel, and assumes familiarity with block-sparse attention, softmax attention, and diffusion sampling loops.

Scope: A training-free block-sparse attention method that estimates dense 3D attention from a weighted sample of small key/value blocks, evaluated on Wan2.1-T2V-1.3B, Wan2.1-T2V-14B, and HunyuanVideo-T2V-13B.

What This Paper Is About

Video diffusion transformers generate high-quality video, but their 3D full attention cost grows quadratically with sequence length and dominates inference time across layers and denoising steps. Existing dynamic block-sparse methods save compute by keeping only a top-k set of key/value blocks, but this both prunes real attention mass out of the softmax and relies on coarse block scores, so fidelity collapses when sparsity is pushed high. The paper's goal is a training-free method that, at very high sparsity, stays close to dense attention rather than approximating a truncated version of it.

Key Contributions

1. Intra-Block Clustering Selector (IBCS): scores key/value blocks using a small set of query centroids clustered within each query block (one k-means iteration from farthest-first centers), avoiding sequence-wide k-means and token permutation while keeping block rankings faithful. Reported to halve the top-k recall gap of mean-pooling to the oracle at a key block size of 4.

2. Sampling with Inclusion-Probability Weighting (SIP): replaces top-k truncation with score-proportional sampling of a fixed number of distinct blocks (systematic sampling), then reweights each sampled block by its inverse inclusion probability. This makes one sparse attention pass an unbiased Horvitz–Thompson estimate of both the attention numerator and the softmax normalizer, and it is implemented as ordinary softmax with logits shifted by −log π.

3. Fine-grained selection with tiled execution: selection is performed on key/value blocks of only 4 tokens, which the fused kernel gathers by index and packs into full 128-token tensor-core tiles, so finer selection costs no extra tensor-core work and requires no token reordering.

4. Sampling-plan reuse: plans are rebuilt every C = 4 denoising steps and reused in between, while attention itself is recomputed from current activations at every step.

Main Findings

• Highest fidelity at the highest sparsity: AB-Attention achieves the best PSNR, SSIM, and LPIPS of any method on all three models, improving PSNR by 1.6–4.4 dB over the strongest baseline, while skipping 89.16% of attention on both Wan2.1 models and 92.35% on HunyuanVideo.
• Baselines compute far more attention: competing methods compute 1.8–3.6× as much attention; for example, DFSAttn on HunyuanVideo reaches 2.07× speedup only while computing 18.71% of dense attention versus AB-Attention's 7.65%.
• Best speedups: 2.01× on Wan2.1-T2V-1.3B (479.11 s vs 962.39 s), 1.84× on Wan2.1-T2V-14B (2571.52 s vs 4743.52 s), and 2.23× on HunyuanVideo-T2V-13B (1974.61 s vs 4399.03 s), all on a single H20 GPU.
• Concrete PSNR values on Wan2.1-T2V-14B: 30.788 for AB-Attention versus 29.236 (SVG2), 27.708 (DFSAttn), and 26.696 (RadialAttention).
• Concrete PSNR values on HunyuanVideo-T2V-13B: 34.347 for AB-Attention versus 30.251 (DFSAttn), 30.230 (SVG2), and 27.686 (RadialAttention).
• SIP is the largest single contributor: in the ablation on Wan2.1-T2V-14B with 100 prompts at 89.16% sparsity, AB-Attention scores 29.705 PSNR, removing the inverse-probability weights drops to 25.706 (−4.00 dB), uniform sampling without weights drops to 22.452 (−7.25 dB), and deterministic top-k gives 25.711 (−3.99 dB).
• Finer selection helps monotonically: tile-level selection gives 28.686 PSNR (−1.02 dB); the block-size sweep rises from 28.686 at 128 tokens through 29.069 (64), 29.201 (32), 29.455 (16), 29.619 (8) to 29.705 (4), with SSIM and LPIPS improving in the same direction.
• Query clustering beats positional pooling: averaging eight consecutive 16-query intervals instead of clustering costs 0.91 dB (28.793 PSNR).
• Attention mass has a long tail: the oracle top-k keeps only about 69% of attention mass with 10.8% of keys at a key block size of 4, and shrinks key/value blocks from 128 to 4 tokens raises that recall from 59% to 69%; sampling with SIP on IBCS scores lowers relative ℓ2 attention error 41% below oracle top-k at b = 4.

Methodology in Plain English

The method starts from a simple observation: at aggressive sparsity, the standard trick of keeping the "top-k" most important blocks throws away a long tail of real attention weight, and no better ranking can bring it back. Instead of keeping a fixed set of blocks, AB-Attention treats the problem as estimation: it draws a random sample of key/value blocks, computes attention exactly on those, and scales each one up to stand in for the blocks it represents.

Two things make this practical. First, cheap but accurate scoring. Rather than summarizing an entire query block with one average vector (which blurs queries that look at different things), the method groups the queries inside each block into a handful of clusters using a single k-means step, and scores each key/value block by how much attention mass those cluster centroids would put on it, weighted by cluster size. This needs no global clustering and no reordering of tokens.

Second, the sampling itself. Each block gets an inclusion probability proportional to its score, with high-scoring blocks assigned probability 1 (always kept). A systematic draw then returns exactly the budgeted number of distinct blocks, so the sample fills whole execution tiles. Each sampled block's logits are shifted by minus the log of its inclusion probability, which turns the whole thing into a single ordinary softmax attention pass whose numerator and denominator are unbiased estimates of the dense ones.

Finally, the block sizes are chosen asymmetrically. Selection happens on very small 4-token key/value blocks so that only the few keys that matter are kept, but the kernel gathers those small blocks and packs them into full tensor-core tiles on the fly, so the finer granularity costs nothing in tensor-core computation and requires no materialized reordering. Plans are rebuilt every four denoising steps to save the scoring overhead, but attention is recomputed each step.

Evaluation compares each sparse video to the dense video generated from the same prompt and seed, using PSNR, SSIM, LPIPS, plus five VBench dimensions, across 600 Penguin Benchmark prompts at 720p with 81 frames (Wan2.1) or 129 frames (HunyuanVideo).

Why This Matters

Impact on research. The paper reframes sparse attention from selection to estimation, which decouples three things prior work tied together: how finely you select, how accurately you score, and how efficiently you execute. The finding that inverse-probability weighting matters far more than any ranking improvement (4.00 dB versus 0.91–1.02 dB) is a direct challenge to the top-k paradigm that most dynamic sparse attention methods share, and it connects diffusion-acceleration work to classical survey-sampling estimators.

Real-world applications:

• Cost reduction for hosted video generation: cutting HunyuanVideo generation from 4399 s to 1975 s and Wan2.1-T2V-14B from 4743.52 s to 2571.52 s on a single H20 GPU lowers the GPU-seconds billed per generated clip.
• Higher resolution and longer videos: because the savings scale with sequence length, the same quality budget can be spent on more frames or larger frames rather than on attention.
• Interactive or batch creative tooling: text-to-video for advertising, previsualization, and storyboarding, where per-clip latency of roughly an hour is a practical blocker.
• Drop-in deployment: the method is training-free, so existing Wan2.1 and HunyuanVideo checkpoints can be accelerated without retraining or fine-tuning.

Industry relevance. The sparsity levels reported (89.16% and 92.35%) with reported PSNR gains over SVG2 and DFSAttn make this directly comparable to production sparse-attention stacks, and the authors' released code at github.com/InfRec-Cardinal/AB-Attention targets the same FlashAttention-3 / FlashInfer backend that the baselines use.

Future Directions

• Hardware generality: the authors state the kernel is built for the H20 GPUs used in their experiments; porting to other accelerators is listed as future work.
• Backward support and training: the current implementation supports only forward inference; a trainable version with backward support would enable fine-tuning under AB-Attention.
• Broader model coverage: evaluation covers three text-to-video diffusion transformers, and the authors plan to extend to a broader range of models.
• Open question — how far the estimator generalizes: since the method's guarantee rests on inclusion probabilities being positive and proportional to attention mass, it is untested on attention distributions where scores are flat, on other modalities (image, audio, 3D), or in combination with orthogonal techniques such as quantization (SageAttention) and step or cache reduction (TeaCache, DyDiT), which the paper positions as complementary but does not evaluate jointly.

Target Audience

Researchers and engineers working on efficient generative video models — particularly those implementing or benchmarking sparse attention kernels for diffusion transformers — will get the most from this paper, since the ablation isolates the value of the weighting scheme versus the ranking scheme and the block-size sweep directly informs kernel design choices. It is also relevant to practitioners deploying Wan2.1 or HunyuanVideo who need concrete latency and fidelity trade-offs, and to students of sampling-based approximation who want a worked example of Horvitz–Thompson estimation inside a GPU attention kernel. It is not beginner-friendly: readers should already understand softmax attention, block-sparse execution, and the diffusion denoising loop.

Authors’ 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.