The paper introduces AB-Attention, a training-free method that accurately approximates dense 3D full attention in video diffusion transformers at high sparsity levels, addressing fidelity degradation in dynamic block-sparse attention caused by coarse block scoring, costly global clustering, and binary tile pruning. The approach comprises three components: an Intra-Block Clustering Selector (IBCS) that faithfully estimates block importance by clustering queries only within each query block without token permutation; Sampling with Inclusion-Probability Weighting (SIP), which samples blocks proportionally to IBCS estimates and reweights them by inverse inclusion probability to estimate dense attention in a single pass rather than truncating it; and a kernel that gathers key/value blocks smaller than a tile into full tensor-core tiles, refining selection without extra tensor-core cost. Evaluated on Wan2.1 and HunyuanVideo, AB-Attention skips 89–92% of attention computation while baselines compute 1.8–3.6× more, yet it improves PSNR by up to 4.4 dB over the strongest baseline and achieves up to 2.23× end-to-end speedup.
Read original
huggingface/daily-papers