Bin Xu, Mo Zhang, Zhengfeiyang Zhang, Xinfeng Yang, Yixing Ma, Wei Wang, HeXin Wang, Pengfei Zhou, Wangbo Zhao
8 min
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.
Sam: That’s the central distinction. Earlier correction methods approximate the discarded remainder through an additional computation path. AB-Attention instead uses a weighted sample to estimate the full calculation in a single sparse attention pass.
Alex: Before the weighting, how does it decide which blocks deserve a place in that sample?
Sam: It groups similar queries within each query block. Each group gets a representative and a weight reflecting how many queries it contains. Those representatives score small key blocks, without rearranging the original token sequence.
Alex: Why group by similarity rather than just average neighbouring queries? Both seem like cheap summaries.
Sam: Neighbours can want different information. Averaging them can wash out a preference held by only part of the block. Local clustering preserves separate preferences, while avoiding the sequence-wide clustering and token permutation used by SVG2, a prior dynamic method.
Alex: Now suppose those scores are imperfect. Does the method still permanently exclude blocks with low scores?
Sam: No. The scores set inclusion probabilities: each block’s chance of entering the sample. Highly scored blocks may always enter, while the others retain a chance. The sampler returns a fixed number of distinct blocks, keeping the computation budget predictable.
Alex: Random sampling alone could still miss useful information. What makes it an estimate of full attention rather than another way to prune?
Sam: The inverse-probability weights. The paper’s example is a block sampled with probability one quarter: when selected, its contribution counts four times. Rarely selected blocks receive more weight when they appear, so the sample accounts statistically for contributions it leaves out.
Alex: Does that make the final attention output unbiased? That would be a stronger guarantee than simply improving an approximation.
Sam: Not the final output. For fixed activations, the estimated weighted sum and estimated normalization are individually unbiased. But attention divides one by the other, and that ratio is generally biased. The authors explicitly distinguish those claims.
Alex: And the practical attraction is that this correction doesn’t require computing a separate approximation for every discarded block?
Sam: The weights become shifts to the sampled attention scores before softmax, the normalization that makes attention weights sum to one. The kernel then performs the usual attention accumulation. Only sampled blocks receive the exact attention computation.
Alex: There’s still a hardware problem. Small blocks give finer choices, but aren’t larger blocks more efficient to execute?
Sam: The method separates selection from execution. It chooses small key-and-value blocks, then gathers them into full hardware tiles on the fly. At a fixed sampling budget, that improves selection without adding tensor-core tile computations.
Alex: So it avoids both global token reordering and tiny matrix operations. Does it also repeat the selection work at every denoising step?
Sam: It reuses sampling plans across nearby steps to reduce that overhead. Attention itself is recomputed from the current activations every step. Reusing the plan does not mean reusing the attention output.
Alex: Let’s move from the mechanism to the test. How broad was the evaluation, and what counted as success?
Sam: The authors evaluated the Wan models and HunyuanVideo using the Penguin Benchmark’s six hundred prompts. Each sparse video was compared with a full-attention video generated from the same prompt and random seed. They also assessed video quality independently of that reference.
Alex: Which comparisons make the result persuasive? Beating a weak shortcut wouldn’t tell us much.
Sam: They compared against RadialAttention’s fixed masks and the dynamic methods DFSAttn and SVG2. Baselines used their released default settings, with a common attention execution backend. Across every tested model, AB-Attention was fastest and closest to the dense reference on all reported reference-based metrics.
Alex: Give us one quality number, but say what it actually measures.
Sam: The largest improvement over the strongest baseline was about four-point-four decibels in peak signal-to-noise ratio, a measure of reference closeness. That is not a direct measure of viewer preference. On the separate video-quality dimensions, the method stayed close to full attention, rather than consistently improving on it.
Alex: The method changes several things together. Can the experiments tell us whether weighting really drives the improvement?
Sam: In a smaller prompt subset on the larger Wan model, removing inverse-probability weights cost about four decibels at the same sparsity. Deterministic top-k selection performed similarly poorly. Replacing clustering or using tile-sized selection blocks also hurt, but less.
Alex: That supports the proposed mechanism within this setting. Where would you hesitate to carry the result further?
Sam: The kernel was built for NVIDIA H20 GPUs and currently supports forward inference only, not training. The evaluation covers a limited family of text-to-video transformers. I’d treat the speedups as demonstrated for those conditions, not as a hardware-independent promise.
Alex: Who should read the full document, and where should they begin?
Sam: Researchers accelerating video generation should start with Main Results, checking both latency and reference fidelity. Then read Sampling with Inclusion-Probability Weighting and the Ablation Study. Kernel developers should also examine Fine-Grained Selection with Tiled Execution. For everyone else, remember this: skipping a connection need not mean pretending its contribution is zero.
Alex: That’s a useful distinction to carry into the next approximation you evaluate.