PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Bin Xu, Mo Zhang, Zhengfeiyang Zhang, Xinfeng Yang, Yixing Ma, Wei Wang, HeXin Wang, Pengfei Zhou, Wangbo Zhao

arXiv:2610.39663v1cs.CV

TL;DR

Video diffusion transformers rely on full attention, whose quadratic cost grows with video length, while dynamic sparse methods have limits in scoring, truncation, and selection granularity. AB-Attention uses local query clustering, probability-weighted sampling, and fine-grained block gathering to approximate dense attention. On Wan2.1 and HunyuanVideo, it skips most attention computation while reporting higher fidelity and end-to-end speedups than baselines.

  • Problem

    Full attention's quadratic cost makes video diffusion inference increasingly expensive with higher resolution and more frames, while dynamic sparse methods face coarse or costly scoring, truncation, and tile-sized selection.

    Source: §1 · p. 3

  • Method

    AB-Attention clusters queries within each query block to score key/value blocks, samples them with inverse inclusion-probability reweighting, and gathers smaller blocks into full tensor-core tiles.

    Source: §1 · p. 3

  • Results

    AB-Attention skips 89–92% of attention computation while baselines compute 1.8–3.6× as much, improves PSNR by 1.6–4.4 dB over the strongest baselines, and reaches 1.84–2.23× end-to-end speedups.

    Source: §1 · p. 3, §6 · p. 12

  • Takeaways & Limitations

    On Wan2.1 and HunyuanVideo, AB-Attention stays closest to dense attention while achieving the highest end-to-end speedups.

    Source: §6 · p. 12

    The current kernel supports H20 GPUs and forward inference only, and evaluation covers three text-to-video diffusion transformers.

    Source: §6 · p. 12

Abstract

from arXiv · show

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.

1. Introduction

Video diffusion attention is costly because its quadratic computation repeats across layers and denoising steps, while dynamic sparse methods face coarse or expensive scoring, truncation, and tile-sized selection. AB-Attention combines local query clustering, probability-weighted block sampling, and finer selection; it reports higher fidelity and speed at high sparsity.

1. Introduction

  • Full attention grows quadratically with sequence length and becomes increasingly expensive as video resolution and frame count rise, making it a major inference bottleneck.The computation repeats across transformer layers and denoising steps.Source: §1 · p. 3
  • Dynamic block-sparse methods adapt key/value block selection to current activations, unlike predefined layouts, but incur online scoring costs.Prior methods use mean-pooled tokens or antidiagonal attention samples, while SVG2 forms blocks through clustering and token permutation.Source: §1 · p. 3
  • Dynamic methods face coarse or costly importance estimates, top-k truncation that removes blocks from the numerator and normalizer, and tile-sized selection that can retain irrelevant keys.Global clustering can require sequence-wide k-means and token permutation; correction methods add a separate approximation branch.Source: §1 · p. 3
  • IBCS forms lightweight query clusters within each query block, while SIP samples blocks in proportion to their scores and reweights by inverse inclusion probability for unbiased numerator and normalizer estimates.The kernel gathers smaller-than-tile key/value blocks into full tensor-core tiles without reordering tokens.Source: §1 · p. 3
  • AB-Attention skips 89–92% of attention computation, while baselines compute 1.8–3.6× as much, and improves PSNR by 1.6–4.4 dB over the strongest baselines.It reports 1.84–2.23× end-to-end speedups on a single H20 GPU.Source: §1 · p. 3

2. Preliminaries

Dense attention computes interactions across all tokens, while sparse attention masks selected interactions to reduce arithmetic. Block-sparse attention groups tokens into blocks, trading finer selection against hardware utilization.

2. Preliminaries

  • For n tokens, dense attention uses query, key, and value matrices for one head and has O(n^2d) complexity, which makes it costly for high-resolution videos.FlashAttention reduces data movement and intermediate storage but still computes all query–key interactions.Source: §2 · p. 4
  • Sparse attention reduces arithmetic by computing only selected query–key pairs, conceptually masking skipped interactions in the logits.A mask value of 0 keeps an interaction, while −∞ skips it.Source: §2 · p. 4
  • Block-sparse attention keeps or skips each query–key block pair as a whole, with larger blocks improving hardware utilization and smaller blocks enabling finer selection.Retained block pairs can be evaluated using dense matrix operations; smaller blocks may reduce execution efficiency.Source: §2 · p. 4

3. Methodology

AB-Attention combines within-block query clustering, inclusion-weighted sampling, and fine-grained key/value selection to estimate dense attention in a sparse pass while retaining tiled execution.

3.1. Intra-Block Clustering Selector (IBCS)

  • Smaller key/value blocks improve selection granularity: shrinking them from 128 to 4 tokens raises oracle top-k retained attention mass from 59% to 69%.Query blocks remain aligned with execution tiles to avoid unioned selections or tensor-core underutilization.Source: §3.1 · p. 5
  • IBCS clusters queries within each query block, scores key/value blocks from weighted centroid attention, and halves mean pooling’s recall gap to oracle top-k.Attention still runs in the original token order; clustering is used only for scoring.Source: §3.1 · pp. 5–6

3.2. Sampling with Inclusion-Probability Weighting (SIP)

  • SIP samples blocks in proportion to their scores and reweights sampled contributions by inverse inclusion probability, accounting for omitted blocks in one sparse attention pass.Its attention error is 41% below oracle top-k at the same block size.Source: §3.2 · p. 6

3.3. Fine-Grained Selection with Tiled Execution

  • The fused kernel gathers sampled small blocks into full tensor-core tiles, applies inclusion-probability logit shifts, and performs attention in one pass without token reordering.Packing sampled blocks preserves the number of tensor-core tile computations at a fixed sampling budget while allowing finer selection.Source: §3.3 · p. 7
  • The selector batches scoring across query and key/value blocks; sampling plans are rebuilt every C denoising steps and reused while attention uses current activations each step.Algorithm 1 summarizes the method for one attention head at one denoising step.Source: §3.3 · p. 7, §3.2 · p. 7

4. Experiments

The experiments compare AB-Attention with sparse-attention baselines on three video diffusion models and evaluate quality and efficiency. Across the models, AB-Attention delivers the strongest fidelity while skipping most attention computation; ablations show that its clustering, sampling, and fine-grained selection each contribute to fidelity.

4.1. Experimental Setup

  • Evaluation covers Wan2.1-T2V-1.3B, Wan2.1-T2V-14B, and HunyuanVideo-T2V-13B using 600 Penguin Benchmark prompts and compares against RadialAttention, SVG2, and DFSAttn.Videos contain 81 frames on Wan2.1 and 129 on HunyuanVideo.Source: §4.1 · pp. 7–8
  • Quality is measured against same-prompt, same-seed dense videos with PSNR, SSIM, and LPIPS; VBench assesses five video dimensions, and speedup is dense-to-sparse end-to-end latency.AB-Attention uses 4-token key/value blocks, 128-query blocks, eight query clusters, and fixed sparsity of 89.16% on Wan2.1 or 92.35% on HunyuanVideo.Source: §4.1 · pp. 7–8

4.2. Main Results

  • AB-Attention achieves the best PSNR, SSIM, and LPIPS on all three models, improving PSNR by 1.6–4.4 dB over the strongest baseline while skipping 89–92% of attention computation.Every competing baseline computes 1.8–3.6 times as much attention.Source: Table 1 · p. 8, §4.2 · p. 8

4.3. Ablation Study

  • 1.02 dB lower PSNR from tile-level selection supports finer selection: PSNR rises as key/value blocks shrink toward AB-Attention’s 4-token blocks.Packing the smaller sampled blocks into full execution tiles adds no tensor-core work.Source: Figure 5 · p. 10, §4.3 · p. 10

4.2. Main Results

  • 1.84–2.23× end-to-end speedups put AB-Attention ahead of every baseline on every model; on HunyuanVideo, generation time falls from 4399 s to 1975 s.DFSAttn reaches 2.07× on HunyuanVideo while computing 2.4 times as much attention as AB-Attention.Source: §4.2 · p. 10

4.3. Ablation Study

  • 0.91 dB lower PSNR from positional query pooling shows that grouping queries by feature similarity improves block scoring over averaging consecutive query intervals.The comparison keeps eight representatives and the overall scoring cost unchanged; SSIM and LPIPS also degrade with pooling.Source: §4.3 · p. 10
  • Removing SIP’s inverse-inclusion weights lowers PSNR by 4.00 dB, and removing scores as well costs another 3.25 dB and more than doubles LPIPS.The weights account for blocks outside the sample, while scores direct the fixed budget toward blocks receiving more attention.Source: §4.3 · p. 10
  • At the same 89.16% sparsity, replacing SIP costs 3.99–7.25 dB of PSNR, versus 1.02 dB for tile-level selection and 0.91 dB for query pooling.Every tested variant is worse than AB-Attention on PSNR, SSIM, and LPIPS.Source: §4.3 · pp. 10–11, Table 2 · p. 10

5. Related Work

Related work spans static and activation-driven sparse attention, methods that correct for pruned interactions, and complementary approaches that accelerate other parts of diffusion inference.

Sparse Attention

  • Static methods impose geometric or energy-based masks, while dynamic methods use activations for block scoring, clustering, or adaptive policies.Examples include Sliding Tile Attention [4], Radial Attention, XAttention, SVG2 [8], and SPADE [33].Source: §5 · p. 11

Correction-Based Sparse Attention

  • Correction methods approximate pruned attention contributions, whereas AB-Attention estimates the numerator and normalizer from one weighted sample of all blocks.Rectified SpaAttn, PISA [10], and SVG-EAR compute selected blocks exactly and approximate the rest.Source: §5 · p. 11

Other Acceleration Methods

  • Reducing denoising steps, reusing computation, shortening sequences, and quantizing attention are complementary to sparse attention.The passage identifies DPM-Solver [45], TeaCache [47], and SageAttention [51] among these approaches.Source: §5 · p. 11

6. Conclusion

AB-Attention estimates dense attention through block sampling rather than truncation, combining local query clustering with inverse-probability weighting. Its appendix details the clustering, systematic sampling, and weighted estimation procedures.

6. Conclusion

  • AB-Attention uses query clusters formed within each query block and samples blocks proportionally to their scores, reweighting them to estimate the attention numerator and normalizer without truncation.The kernel packs selected key/value blocks smaller than a tile into full tensor-core tiles.Source: §6 · pp. 11–12

A.1. Query Centroids and Pooled Keys

  • For IBCS, one k-means iteration starts from farthest-first query centers, while each four-token key block is pooled into two keys to halve scoring cost.The pooled keys each represent two keys, so key-side size weights cancel in the softmax.Source: §A.1 · p. 19

A.2. Systematic Sampling Without Replacement

  • Systematic sampling returns exactly the allocated number of distinct blocks with their specified inclusion probabilities, unlike independent draws with variable sample size or replacement draws that can duplicate blocks.Scores are floored so each inclusion probability remains positive, and the budget fills whole 128-token tiles.Source: §A.2 · p. 19

A.3. Weighted Attention Estimation

  • For fixed activations, inverse-inclusion-weighted estimates of the attention numerator and denominator are unbiased, although their ratio is generally biased.The sampled-block attention logits are shifted by −log inclusion probability, and tiling preserves the selected indices and weights.Source: §A.3 · p. 19

B.1. AB-Attention

The implementation applies sparse selection to video-token attention while retaining exact attention for text-related paths. Its block and sampling configurations differ between Wan2.1 and HunyuanVideo.

B.1. AB-Attention

  • Wan2.1 replaces video-token self-attention with AB-Attention while leaving text cross-attention dense; HunyuanVideo samples video keys for video queries but keeps prompt-key and prompt-query attention exact.Source: §B.1 · p. 19
  • Both Wan2.1 models use 128-token query blocks, four-token key/value blocks, eight query clusters per block, and 2,048 sampled blocks per query block.The sample is 8,192 keys per query block.Source: §B.1 · pp. 19–20
  • HunyuanVideo samples 2,272 blocks per query block, equal to 9,088 keys and 7.65% of its 29,700 video blocks; Wan2.1 samples 10.84% of its 18,900 blocks.Attention remains dense in the first transformer block and denoising steps 1–10; thereafter the plan is rebuilt every four steps.Source: §B.1 · p. 20

B.2. Baselines

The baseline settings use method-specific sparsity controls and block or cluster layouts.

B.2. Baselines

  • RadialAttention uses decay factors of 0.2 on Wan2.1 and 0.95 on HunyuanVideo; SVG2 retains clusters up to cumulative attention 0.9, while DFSAttn refreshes its retained block fraction from 0.3 down to 0.1.SVG2 uses a minimum key ratio of 0.1; DFSAttn scores 16-token sub-blocks and refreshes every 12 steps.Source: §B.2 · p. 20

C. Evaluation Protocol

The evaluation specifies fidelity metrics, sparsity accounting, and end-to-end latency measurement. Its block-size sweep reports that fidelity metrics do not degrade as selection blocks shrink.

C. Evaluation Protocol

  • PSNR, SSIM, and VGG-backbone LPIPS are computed on RGB frames in [0, 1], averaged over video frames and then prompts; five VBench dimensions are averaged over prompts.Source: C. Evaluation Protocol · p. 20
  • Sparsity counts only sparse denoising steps and is calculated from each method’s corresponding mask, budget, or recorded block selections.Source: C. Evaluation Protocol · p. 20
  • Table 3 compares fidelity to dense attention across selection block sizes on Wan2.1-T2V-14B (100 prompts), with every row executing 8,192 key tokens per query block in 128-token tiles at 89.16% sparsity.Block size 4 is full AB-Attention, while block size 128 is tile-level selection.Source: Table 3 · p. 20
  • End-to-end latency is measured for one video on NVIDIA H20 GPUs across text encoding, 50 denoising steps, and VAE decoding, including AB-Attention plan construction and sampling.Model loading, initial-noise loading, and video export are excluded.Source: C. Evaluation Protocol · p. 20

D. Additional Ablation Results

Across the block-size sweep, shrinking the selection block from 128 to 4 tokens does not degrade PSNR, SSIM, or LPIPS.

D. Additional Ablation Results

  • PSNR, SSIM, and LPIPS do not degrade as the selection block shrinks from 128 to 4 tokens.Table 3 reports these metrics for the block-size sweep in Figure 5.Source: D. Additional Ablation Results · p. 20

Related papers

Loading 2610.39663v1…