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

AMD introduced VSA (Video Sparse Attention), a hardware-efficient sparse attention mechanism implemented with CK Tile, achieving a 3.31× attention kernel-time speedup at 70% sparsity over FlashAttention on AMD Instinct MI308X GPUs for video diffusion inference. Developed by researchers from UC San Diego, MBZUAI, and UC Berkeley, VSA uses a two-stage coarse-to-fine attention approach with 3D space-filling curve token reordering to reduce computation while preserving visual quality.

VSA: Accelerating Video Diffusion Inference with Sparse Attention on AMD GPUs 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. 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. The Attention Bottleneck in Video Diffusion the-attention-bottleneck-in-video-diffusion Modern video diffusion models like Wan2.1 https://github.com/Wan-Video/Wan2.1 , HunyuanVideo https://github.com/Tencent/HunyuanVideo , and CogVideoX https://github.com/THUDM/CogVideo rely on transformer architectures where attention dominates both compute and memory costs. For a typical video generation task: Parameter | Typical Value | |---|---| | 832 × 480 | | 81 | | 32,768 tokens | | O N² = ~1 billion operations per layer | The standard scaled dot-product attention SDPA formula is: \ \text{Attention} Q, K, V = \text{Softmax}\left \frac{QK^T}{\sqrt{d k}}\right V\ 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 vsa-video-sparse-attention VSA Video Sparse Attention https://arxiv.org/abs/2505.13389 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. Core Principles core-principles VSA is built on three key insights: 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. Two-Stage Architecture two-stage-architecture 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. Pre-stage: 3D Space-Filling Curve SFC Token Reordering Before any attention computation, video tokens are reordered using a 3D space-filling curve SFC —a technique adopted from Jenga https://arxiv.org/abs/2505.16864 . 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. Stage 1: Coarse Selection 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. Stage 2: Fine Computation 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. Block-Sparse Encoding block-sparse-encoding VSA implements block-level sparse encoding that aligns with GPU execution characteristics: VSA Sparse Structure: +-- lut ptr Block index lookup table +-- valid block num Number of valid blocks per query row +-- kv block idx K/V block indices for each query block 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. CK Tile Implementation ck-tile-implementation We provide high-performance implementations of VSA optimized for AMD Instinct GPUs in AMD’s Composable Kernel CK Tile https://github.com/ROCm/composable kernel library. Key Components key-components Component | File Path | |---|---| | | | | | | | | Kernel Architecture kernel-architecture The CK Tile VSA kernel implements a three-stage pipeline with double buffering, enabling asynchronous overlap of computation and memory access: Stage 1: QK GEMM + Softmax Statistics Q tiles × K tiles → attention scores Compute running max M and sum L for online softmax Stage 2: Softmax + Post-ops Apply softmax normalization using M and L Stage 3: KV GEMM Softmax output × V tiles → attention output Accumulate with previous tiles Sparse Traversal sparse-traversal Unlike dense attention that iterates over all K/V blocks, VSA uses the LUT to jump directly to relevant blocks: // Pseudo-code for VSA kernel traversal for int i = 0; i < valid block num query block ; i++ { int kv block = kv block idx query block i ; // Load K/V tiles from kv block // Compute attention for this block pair } This eliminates wasted computation on blocks that would contribute negligible attention weight. Comparison with Jenga comparison-with-jenga Jenga https://arxiv.org/abs/2505.16864 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. Algorithmic Design algorithmic-design 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. Jenga is a two-component inference pipeline : 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. Sparse Encoding in CK Tile sparse-encoding-in-ck-tile When both methods are implemented as CK Tile kernels, the difference in how they store the sparse pattern becomes concrete: 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. Side-by-Side Summary side-by-side-summary Dimension | VSA | Jenga | |---|---|---| | Single Top-K from coarse cube-level attention | Union of 3 masks: data-driven Top-K + forced text-condition + forced adjacency | | | 3D Space-Filling Curve SFC reordering → uniform blocks | | Within-step attention only | Within-step AttenCarve + cross-step resolution ProRes + timestep skip | | Yes learnable gate between coarse and fine | No training-free, plug-and-play | | Captured implicitly by cube structure | Explicit Adjacency Mask enforces local block attention | | Not separately handled | Explicit Condition Mask preserves text-token attention | | Compact index list + valid count LUT | Full M×M one-hot block relation matrix | | 3.31× attention kernel vs FlashAttention measured by AMD, see | 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. Qualitative Visual Check qualitative-visual-check 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. Visual Comparison visual-comparison Flash Attention Dense, ~60.8 ms | CK VSA Sparse Sparse, ~25 ms avg | |---|---| Quality Notes quality-notes Implementation | Observation | Notes | |---|---|---| | Dense baseline | Computes all token pairs | | Visually close to baseline in this sample | No obvious artifacts observed in the sampled frames | 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. Performance Benchmarks performance-benchmarks 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 appendix-detailed-benchmark-configuration . Headline Result: Kernel Time by Sparsity Level headline-result-kernel-time-by-sparsity-level 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. End-to-End Impact end-to-end-impact Aggregated over the full 50-step inference using the sparsity distribution actually observed during generation see Sparsity Distribution During Inference 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. Summary summary 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 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. Appendix - Detailed Benchmark Configuration appendix-detailed-benchmark-configuration 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. Test Configuration test-configuration Parameter | Value | |---|---| | AMD Instinct™ MI308X GPU | | Wan2.1-T2V-1.3B | | Text-to-Video Generation | | 832 × 480 | | 81 | | 50 | | BF16 | Tensor Specifications tensor-specifications Parameter | Value | |---|---| | | | 1 | | 12 | | 32,768 tokens | | 128 | | 128 × 128 | | | | ~55-60 | 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 Distribution During Inference sparsity-distribution-during-inference 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. When CK VSA Sparse Attention Helps Most when-ck-vsa-sparse-attention-helps-most Sparsity range | Observed speedup vs FlashAttention | Practical guidance | |---|---|---| < 40% | ~1.6× | Marginal; dense FlashAttention is a reasonable fallback | 40-60% | ~1.9× – 2.8× | CK VSA recommended | 60% | 2.8× | CK VSA strongly recommended | Integration Guide integration-guide Prerequisites prerequisites GPU : AMD Instinct™ MI308X or other ROCm-compatible GPU ROCm : 6.3+ PyTorch : 2.3+ CK Tile : Latest from composable kernel https://github.com/ROCm/composable kernel Using VSA with CK Tile using-vsa-with-ck-tile The VSA implementation is available through the AITER https://github.com/ROCm/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. python from aiter.ops.jenga sparse attention import vsa sparse attention Prepare inputs TQ = torch.randn batch, heads, seq len, head dim, dtype=torch.bfloat16, device="cuda" TK = torch.randn batch, heads, seq len, head dim, dtype=torch.bfloat16, device="cuda" TV = torch.randn batch, heads, seq len, head dim, dtype=torch.bfloat16, device="cuda" Prepare LUT from Top-K selection Tkv block idx = ... batch, heads, num q blocks, max kv blocks block indices Tkv blocks = ... batch, heads, num q blocks valid block count per query Allocate output out = torch.zeros like TQ Compute VSA sparse attention output = vsa sparse attention TQ, TK, TV, Tkv block idx, LUT: K/V block indices for each Q block Tkv blocks, Number of valid K/V blocks per query block out, batch=batch, nhead=heads, nhead k=heads, seqlen q=seq len, seqlen k=seq len, hdim q=head dim, hdim v=head dim Key Input Parameters : Parameter | Shape | Description | |---|---|---| | | Query, Key, Value tensors BF16 | | | LUT storing K/V block indices for each Q block | | | Number of valid K/V blocks to compute per query block | CK Tile Source Code : Component | File Path | |---|---| VSA Kernel Example | | Dispatch Logic | | VSA Kernel | | VSA Pipeline | | Generating Sparsity Patterns generating-sparsity-patterns VSA requires upstream sparsity selection to generate the LUT. These can be: Heuristic-based : Use spatial-temporal locality to determine block importance Profile-based : Analyze attention patterns from sample runs to derive sparsity masks Dynamic : Compute coarse attention scores at runtime for selection 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. python def generate sparsity lut query, key, block size, top k ratio : """Generate LUT using coarse attention scores.""" Pool tokens into blocks q blocks = pool to blocks query, block size k blocks = pool to blocks key, block size Compute coarse attention coarse scores = torch.einsum 'bhqd,bhkd- bhqk', q blocks, k blocks Top-K selection per query block top k = int coarse scores.shape -1 top k ratio , lut = torch.topk coarse scores, top k, dim=-1 return lut, top k Acknowledgements acknowledgements 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. Additional Resources additional-resources Disclaimers disclaimers 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. © 2026 Advanced Micro Devices, Inc. All rights reserved.