VSA: Accelerating Video Diffusion Inference with Sparse Attention on AMD GPUs

Video generation powered by diffusion transformers has achieved remarkable quality, but the computational cost of attention mechanisms remains a critical bottleneck. With sequence lengths reaching tens of thousands of tokens in video generation tasks, the quadratic complexity of standard attention becomes prohibitively expensive.

9 min read. | 2212 total words.

Video generation powered by diffusion transformers has achieved remarkable quality, but the computational cost of attention mechanisms remains a critical bottleneck. With sequence lengths reaching tens of thousands of tokens in video generation tasks, the quadratic complexity of standard attention becomes prohibitively expensive.

This blog introduces VSA (Video Sparse Attention) implemented with CK Tile, a hardware-efficient sparse attention mechanism that significantly accelerates video diffusion inference. We demonstrate how VSA, implemented through AMD’s CK Tile library, delivers significant speedups across various sparsity levels, achieving a 3.31× attention kernel-time speedup at 70% sparsity over FlashAttention on AMD Instinct™ MI308X GPUs, with qualitative visual checks included as a sanity check.

Results may vary based on model, prompt, resolution, frame count, sequence length, sparsity level, inference settings, software versions, system configuration, and other factors.

Modern video diffusion models like Wan2.1, HunyuanVideo, and CogVideoX rely on transformer architectures where attention dominates both compute and memory costs. For a typical video generation task:

While FlashAttention optimizes memory access patterns, it still computes attention over all token pairs. However, research has shown that most attention mass concentrates in a small subset of positions—a property that sparse attention methods exploit to reduce computation.

VSA (Video Sparse Attention) is a hardware-efficient sparse attention mechanism designed specifically for video diffusion transformers. Developed by researchers from UC San Diego, MBZUAI, and UC Berkeley, VSA introduces a two-stage coarse-to-fine attention approach that dramatically reduces computation during inference by focusing token-level attention on selected spatial-temporal regions.

Attention Sparsity: In video diffusion, attention patterns exhibit strong spatial-temporal locality. Tokens primarily attend to nearby frames and spatial regions, making full attention wasteful.

Hardware Alignment: Sparse patterns must align with GPU tile sizes to achieve actual wall-clock speedups, not just theoretical FLOP reductions.

Adaptive Selection: VSA dynamically selects which blocks to attend to based on the attention patterns, adapting to the specific characteristics of video data.

VSA implements a two-stage coarse-to-fine selection mechanism. The goal is to avoid computing full attention over all video tokens while still preserving the most important spatial-temporal regions. The end-to-end data flow is summarized in the figure below.

Before any attention computation, video tokens are reordered using a 3D space-filling curve (SFC)—a technique adopted from Jenga. In native linear layout (T, H, W), tokens that are spatially adjacent in 3D but far apart in the flattened 1D sequence can end up in different attention blocks, breaking spatial locality. SFC reordering remaps token positions so that tokens close together in 3D space are also close together in the 1D sequence. This ensures that when the sequence is partitioned into fixed-size blocks for the attention kernel, each block corresponds to a contiguous spatial-temporal region of the video—which is the prerequisite for block-sparse patterns to be meaningful.

After SFC reordering, VSA groups neighboring video tokens into spatial-temporal cubes. In the VSA paper, a typical setting is (Ct, Ch, Cw) = (4, 4, 4), so each cube contains 64 tokens. Each cube is mean-pooled into one cube-level representation, producing cube-level Qc, Kc, and Vc.

The coarse stage then computes cube-to-cube attention scores. For each query cube, VSA selects the Top-K key/value cubes with the highest scores. These selected cube IDs define the block-sparse attention pattern used by the fine stage.

Conceptually, each selected cube-level entry expands into a B x B block in the full attention mask. In practice, VSA does not materialize this full-resolution mask. Instead, it passes the selected block indices directly to the fine-grained attention kernel.

The fine stage performs normal token-level attention, but only over the K/V cubes selected by the coarse stage. Unselected cubes are skipped entirely, reducing both memory traffic and attention computation while keeping the work aligned with block-sparse GPU kernels.

The final VSA output combines the coarse-stage output and the fine-stage output through learnable gates. This keeps global context from the coarse stage while using sparse token-level attention for the most important regions.

This block encoding format allows the kernel to skip irrelevant blocks entirely. Unlike a full 0/1 block mask, VSA stores only selected K/V block indices plus the number of valid blocks for each query block.

We provide high-performance implementations of VSA optimized for AMD Instinct GPUs in AMD’s Composable Kernel (CK) Tile library.

example/ck_tile/50_sparse_attn/vsa_sparse_attention.cpp

example/ck_tile/50_sparse_attn/fmha_fwd_trek.hpp

include/ck_tile/ops/sparse_attn/kernel/fmha_fwd_vsa_kernel.hpp

include/ck_tile/ops/sparse_attn/pipeline/block_fmha_pipeline_qr_ks_vs_async_vsa.hpp

This eliminates wasted computation on blocks that would contribute negligible attention weight.

Jenga is another recent work targeting efficient video diffusion inference, and its CK Tile implementation shares the same block-sparse attention infrastructure as VSA. However, VSA and Jenga differ substantially in their algorithmic design, sparse pattern selection strategy, and system-level scope—not just in their kernel encoding format.

VSA is a single-component sparse attention method. Its two stages (coarse and fine) both operate within the attention module of each transformer layer: the coarse stage selects which K/V cubes matter for each query, and the fine stage computes token-level attention only over those cubes. The final output is a learnable gate of the two stages.

AttenCarve (within-step sparse attention): Jenga first reorders tokens using a 3D space-filling curve (SFC) so that spatially adjacent video tokens are also adjacent in the 1D flattened sequence, then partitions them into M uniform blocks. The sparse block selection is the union of three masks: Importance Mask (B_top): data-dependent; uses block-level mean Q/K scores (similar in spirit to VSA’s coarse stage) to select Top-K relevant K/V blocks per query block. Condition Mask (B_cond): pre-computed; attends to text condition tokens to preserve cross-modal alignment. Adjacency Mask (B_adja): pre-computed; attends to spatially adjacent blocks to maintain local spatial coherence.

AttenCarve (within-step sparse attention): Jenga first reorders tokens using a 3D space-filling curve (SFC) so that spatially adjacent video tokens are also adjacent in the 1D flattened sequence, then partitions them into M uniform blocks. The sparse block selection is the union of three masks:

Importance Mask (B_top): data-dependent; uses block-level mean Q/K scores (similar in spirit to VSA’s coarse stage) to select Top-K relevant K/V blocks per query block.

Condition Mask (B_cond): pre-computed; attends to text condition tokens to preserve cross-modal alignment.

Adjacency Mask (B_adja): pre-computed; attends to spatially adjacent blocks to maintain local spatial coherence.

ProRes (cross-step resolution scheduling): early denoising steps run on low-resolution latents (fewer tokens); resolution is gradually increased to the target as denoising progresses. This reduces quadratic attention cost at the pipeline level, independently of AttenCarve.

VSA stores only the selected K/V block indices plus a valid-count per query block (compact index list / LUT). The kernel jumps directly to active blocks.

Jenga stores the full M×M one-hot block relation matrix B and skips cells where B[i][j]=0 during traversal.

VSA’s encoding is more compact when sparsity is high; Jenga’s encoding naturally represents the union of its three heterogeneous masks (Importance ∪ Condition ∪ Adjacency) without converting to a list.

3.31× attention kernel vs FlashAttention (measured by AMD, see Performance Benchmarks)

8.83× end-to-end on VBench (0.01% quality drop), as reported by the Jenga authors

Note on speedup numbers: the two figures measure different things. VSA’s 3.31× is a kernel-level timing comparison against FlashAttention at a fixed 70% sparsity. Jenga’s 8.83× is an end-to-end pipeline speedup that includes both AttenCarve and ProRes (reduced token count from lower resolution). A direct apples-to-apples comparison would require running both on the same model under the same conditions. The Jenga figures in this section, including the 8.83× speedup and the accompanying quality result, are those reported in the Jenga paper (arXiv:2505.16864); AMD has not independently reproduced or verified them. The descriptions of Jenga’s design in the table above are likewise drawn from that paper.

In our CK Tile implementation, both methods actually share the 3D SFC token reordering step upstream. The key difference lies in what happens after reordering: VSA uses a Top-K coarse attention score to build a compact LUT, while Jenga builds a full M×M block relation matrix from its three-mask union. The choice of kernel then follows naturally from the encoding: the VSA CK Tile kernel consumes the compact LUT, while the Jenga CK Tile kernel consumes the full block matrix.

We generated videos using the same prompt (“Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage.”) and the same random seed with dense FlashAttention and CK VSA Sparse Attention. The purpose of this check is to confirm that the sparse attention path does not introduce obvious visual artifacts in this sample.

Flash Attention (Dense, ~60.8 ms)

This is a qualitative sanity check rather than a full video-generation quality benchmark. We have not yet included quantitative metrics such as VBench in this post. A VBench-style evaluation would be useful future work to measure detail preservation, temporal consistency, and semantic alignment more rigorously.

We benchmarked CK VSA Sparse Attention against dense FlashAttention on a text-to-video generation task using the Wan2.1-T2V-1.3B model (832x480, 81 frames, 50 inference steps, BF16) on a single AMD Instinct™ MI308X GPU. Each attention call operates on Q/K/V of shape [1, 12, 32768, 128] with a 128 x 128 block size. Detailed tensor specifications and the per-step sparsity distribution observed during inference are listed in the Appendix: Detailed Benchmark Configuration.

The figure below is the headline comparison: CK VSA Sparse Attention versus dense FlashAttention at varying sparsity levels. Higher sparsity means fewer K/V blocks are selected by VSA, which directly translates into shorter attention kernel time. FlashAttention is essentially constant (~60.8 ms) because it computes dense attention regardless of the sparse pattern, while CK VSA’s runtime decreases as sparsity increases, reaching a 3.31× kernel-time speedup at 70% sparsity.

Aggregated over the full 50-step inference using the sparsity distribution actually observed during generation (see Sparsity Distribution During Inference below), the kernel-time gains translate into roughly ~37% lower end-to-end generation time (about 3 min with VSA versus ~4 min 47 s with FlashAttention) on this configuration.

The kernel timings above assume the selected block indices are already available to the sparse kernel. In a full deployment, the coarse-stage selection and LUT generation must also be accounted for; that overhead is designed to be lightweight and is amortized by the savings in the fine stage, but it should still be measured in any production evaluation.

We presented a CK Tile implementation of Video Sparse Attention (VSA) for video diffusion inference on AMD Instinct GPUs. On Wan2.1-T2V-1.3B at 32,768 tokens per attention call, CK VSA delivers up to 3.31× kernel-time speedup over dense FlashAttention at 70% sparsity, and roughly 37% lower end-to-end generation time on MI308X. We also compared the CK VSA index-list encoding with the existing CK Jenga block-mask encoding to clarify when each representation is preferable. As video diffusion models grow in resolution and length, hardware-aligned block-sparse attention such as VSA becomes increasingly important for practical deployment; a quantitative quality study (e.g., VBench) on top of these results is left as future work. For the full benchmark setup and the per-step sparsity distribution, please see the Appendix: Detailed Benchmark Configuration.

Results may vary based on model, prompt, resolution, frame count, sequence length, sparsity level, inference settings, software versions, system configuration, and other factors.

The main text uses a single condensed configuration paragraph. The full set of parameters and the per-step sparsity behavior are listed here for reproducibility.

Note: the (Ct, Ch, Cw) = (4, 4, 4) cubes mentioned in the VSA algorithm are used for the coarse Top-K selection; the 128 × 128 block size above refers to the GPU tile granularity at which the fine-stage CK Tile kernel iterates over the selected K/V cubes.

Sparsity is not constant across diffusion steps. Early diffusion steps select more K/V blocks (lower sparsity), while late steps become more selective (higher sparsity). The distribution we observed on this workload is shown below.

Combining this distribution with the per-sparsity kernel times in the main text gives a weighted-average attention kernel time of roughly ~25 ms per call for CK VSA versus ~60.8 ms for FlashAttention, which is the basis for the end-to-end speedup quoted above.

The VSA implementation is available through the AITER Python bindings, which wrap the CK Tile C++ kernels. Note that CK VSA and CK Jenga share the same sparse-attention dispatcher module in AITER, which is why the import path is named after jenga_sparse_attention; the underlying VSA kernel is still the index-list / LUT-based implementation described above.

example/ck_tile/50_sparse_attn/vsa_sparse_attention.cpp

example/ck_tile/50_sparse_attn/fmha_fwd_trek.hpp

include/ck_tile/ops/sparse_attn/kernel/fmha_fwd_vsa_kernel.hpp

include/ck_tile/ops/sparse_attn/pipeline/block_fmha_pipeline_qr_ks_vs_async_vsa.hpp

VSA requires upstream sparsity selection to generate the LUT. These can be:

Note: The following is illustrative pseudocode. pool_to_blocks is a placeholder for the actual cube-level mean-pool used by VSA; see the VSA paper for the production implementation.

The authors would like to thank the AMD CK Tile and AITER teams for their support in developing and optimizing the sparse attention kernels on AMD Instinct GPUs. We also thank the original VSA authors from UC San Diego, MBZUAI, and UC Berkeley for open-sourcing their work and making this collaboration possible.

Wan2.1 Text-to-Video Model

The information presented in this document is for informational purposes only and may contain technical inaccuracies, omissions, and typographical errors. The information contained herein is subject to change and may be rendered inaccurate for many reasons, including but not limited to product and roadmap changes, component and motherboard version changes, new model and/or product releases, product differences between differing manufacturers, software changes, BIOS flashes, firmware upgrades, or the like. Any computer system has risks of security vulnerabilities that cannot be completely prevented or mitigated. AMD assumes no obligation to update or otherwise correct or revise this information. However, AMD reserves the right to revise this information and to make changes from time to time to the content hereof without obligation of AMD to notify any person of such revisions or changes.

THIS INFORMATION IS PROVIDED “AS IS.” AMD MAKES NO REPRESENTATIONS OR WARRANTIES WITH RESPECT TO THE CONTENTS HEREOF AND ASSUMES NO RESPONSIBILITY FOR ANY INACCURACIES, ERRORS, OR OMISSIONS THAT MAY APPEAR IN THIS INFORMATION. AMD SPECIFICALLY DISCLAIMS ANY IMPLIED WARRANTIES OF NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR ANY PARTICULAR PURPOSE. IN NO EVENT WILL AMD BE LIABLE TO ANY PERSON FOR ANY RELIANCE, DIRECT, INDIRECT, SPECIAL, OR OTHER CONSEQUENTIAL DAMAGES ARISING FROM THE USE OF ANY INFORMATION CONTAINED HEREIN, EVEN IF AMD IS EXPRESSLY ADVISED OF THE POSSIBILITY OF SUCH DAMAGES.

Third-party content is licensed to you directly by the third party that owns the content and is not licensed to you by AMD. ALL LINKED THIRD-PARTY CONTENT IS PROVIDED “AS IS” WITHOUT A WARRANTY OF ANY KIND. USE OF SUCH THIRD-PARTY CONTENT IS DONE AT YOUR SOLE DISCRETION AND UNDER NO CIRCUMSTANCES WILL AMD BE LIABLE TO YOU FOR ANY THIRD-PARTY CONTENT. YOU ASSUME ALL RISK AND ARE SOLELY RESPONSIBLE FOR ANY DAMAGES THAT MAY ARISE FROM YOUR USE OF THIRD-PARTY CONTENT.

AMD, the AMD Arrow logo, and combinations thereof are trademarks of Advanced Micro Devices, Inc. Other product names used in this publication are for identification purposes only and may be trademarks of their respective companies.

Text integral sub licența deschisă a publicației · 2892 cuvinte

analiza aiqbench

VSA: Accelerating Video Diffusion Inference with Sparse Attention on AMD GPUs

Subiecte conexe