Faster Video Diffusion with Trainable Sparse Attention
Authors: Peiyuan Zhang
Yongqi Chen
Haofeng Huang
Will Lin
Zhengzhong Liu
Ion Stoica
Eric P. Xing
Hao Zhang
Venue: NeurIPS 2025 poster ()
Keywords: Diffusion Model;Video Generation;Efficient Attention;Sparse Attention
Primary area: deep_learning
TL;DR: Trainable sparse attention for video diffusion model.

Abstract: Scaling video diffusion transformers (DiTs) is limited by their quadratic 3D attention, even though most of the attention mass concentrates on a small subset of positions.  We turn this observation into VSA, a trainable, hardware-efficient sparse attention that replaces full attention at both training and inference. In VSA, a lightweight coarse stage pools tokens into tiles and identifies high-weight critical tokens; a fine stage computes token-level attention only inside those tiles subjecting to block computing layout to ensure hard efficiency. This leads to a single differentiable kernel that trains end-to-end, requires no post-hoc profiling, and sustains 85\% of FlashAttention3 MFU. We perform a large sweep of ablation studies and scaling-law experiments by pretraining DiTs from 60M to 1.4B parameters. VSA reaches a Pareto point that cuts training FLOPS by 2.53$\times$ with no drop in diffusion loss. Retrofitting the open-source Wan2.1-1.3B model speeds up attention time by 6$\times$ and lowers end-to-end generation time from 31s to 18s with comparable quality, while for the 14B model, end-to-end generation time is reduced from 1274s to 576s.
Furthermore, we introduce a preliminary study of Sparse-Distill, the first method to enable  sparse attention and distillation concurrently, achieving 50.9x speed up for Wan-1.3B while maintaining quality.
These results establish trainable sparse attention as a practical alternative to full attention and a key enabler for further scaling of video diffusion models.  Code is available at https://github.com/hao-ai-lab/FastVideo.

Faster Video Diffusion with Trainable Sparse Attention
Peiyuan Zhang"
Yongqi Chen
Haofeng Huang 1- Will Lin' Eric P. Xing"
Zhengzhong Liu
Ion Stoica
Hao Zhang'
UC San Diego MBBZAAI
'UC Berkeley
Abstract
Scaling video diffusion transformers (DiT') is limited by their quadratic 3D tention, even though most of the attention mass concentrates on small subset of positions. We turn this observation into VSA, a trainable, hardware-efficient sparse attention that replaces full attention at both training and inference. In VSA, lightweight coarse stage pools tokens into tiles and identifies high-weight crit- ical tokens; fine stage computes token-level attention only inside those tiles subjecting to block computing layout to ensure hard efficiency. This leads to a single differentiable kernel that trains end-to-end, requires post-hoc profiling, and sustains 85% FlashAttentinn3 MFU. We perform large sweep of ablation studies and scaling-law experiments by pretraining DiTs from 60M to 1.4B pa- rameters. VSA reaches a Pareto point that cuts training FLOPS by 2.53 with no drop in diffusion loss. Retrofitting the open- up attention time by 6x and lowers end-to-end generation time from 31s to 18s with comparable quality, while for the 14B model, end-to-end generation time is reduced from 1274s to 576s. Furthermore, we introduce a preliminary study of Sparse-Distill, the first method enable sparse attention and distillation concur- rently, achieving 50.9x speed up for Wan- 1.3B while maintaining quality. These results establish trainable sparse attention as practical alternative full attention and a key enabler for further scaling of video diffusion models. Code is available https://github. wttptuuw..oo..oaaaaaaaaaaa aato-aallab/aaaaiiiiae
Wan2. 1-1.3B model speeds
Introduction
Attention computation is the primary bottleneck when scaling video Diffusion Transformers (DiT) [34, 28]. Even a seemingly short 5-second 720p clip unfolds into more than 100K tokens [29, 20] once flattened as sequence. Consequently, -tate-of-the-art video DiTs [20, 35, 43, 26] expend the majority of compute on attention when training on full-resolution, sengseeuencc data; the trained DiTs remain painfully slow at inference. Fortunately, recent studies [37, 48, 6, 47] reveal inherent sparsity in DiTs trained using full attention: only a tiny subset of entries in the attention matrix Softmax (QKT/Q /Vd), which we refer to as critical tokens, significantly influence outputs, while the vast majority approach zero. This inherent sparsity calls for developing native, trainable sparse attention mechanism purpose-built for video DiTs.
Most prior work approaches sparsity class training primitive, Sliding Tile Attention (STA) [48] and Sparge Attention [47], for example, begin with model trained under full attention and then substitute each head with fixed profile- derived sparse mask only at inference time [37, 38). Because the sparsity pattern is decided after training, these methods leave the bulk of training cost untouched and introduce a train-test mismatch: the DIT' learns parameters dense context yet is evaluated in a sparse one. That mismatch caps
post-hoc speedup for pretrained DiTs rather than first-
Equal contribution. Work performed during an internship at UC San Diego
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

best-case quality at the dense model's 's ceiling and, in practice, often erodes quality once sparsity pushed beyond gentle budget. Consequently, state-of-the-art DiTs still default to quadratic 3D attention despite its prohibitive cost [43, 35, 20, 2, 11].
Designing a trainable sparse attention for video DiT's faces fundamental chicken-and-egg dilemma: identifying critical token positions traditionally requires computing the full attention matrix, which then erases any computational gains and defeats the purpose of sparse attention. Conversely, resorting to cheap heuristics and not precisely identifying critical tokens may miss high-weight regions and yield suboptimal results. More importantly, any practical attention implementation must honor the as Flash Attention (FA) [5] otherwise theoretical savings do not translate to wall-clock speedup. The central research question therefore: how can we predict critical tokens accurately, subject to hardware-aligner block structure, without
layouts expected by modern GPU kernels such
paying the quadratic cost we aim to avoid?
This paper presents VSA A (Video Sparse Attention), a trainable, hardware-aligned sparse attention mechanism for video DiT', drawing inspiration from recent developments in large language mod- els [45, 25, 30J. At the high level, VSA is a hierarchical granular attention, illustrated in Figure I. The coarse stage first aggregates a cube containing (4, 4, 4) tokens into a single representation to compute cube-to-cube dense attention. Since attention operates on a pooled (short) sequence, it is lightweight, yet simultaneously predicts which cube contains critical tokens while modeling global context. A fine stage then applies token-level attention only inside the top-X selected cube. The final output of attention combines both stages through a differentiable gate. Because VSA end-to-end trainable, it identifies critical tokens not by heuristics, but by learning from data. To ensure hardware efficiency, VSA meticulously designed to map a spatial-temporal cube to a kernel-level tile [48]. This ensures that tokens within cube are loaded on the same GPU SM and adhere to the block sparse compute layout ($2.2).
One critical parameter in VSA is the tile size. Small tiles let the coarse stage localize critical tokens close to token resolution, while the fine stage attends only to those pinpointed cubes. This sharpens sparsity and improves model quality at the cost of fragmenting work into tiny blocks, which hurts the GPU throughput. In contrast, large tiles boost arithmetic intensity. but the fine stage must then process whole cubes even if only a few tokens inside the cubes matter, thus blurring sparsity (Figure (b)). Another key parameter whether to inject dedicated local or spatial-temporaa patterns into the fine stage. Our systematic ablation studies reveal that an effective VSA configuration combines global coarse stage with a freely-selectable fine stage. We found that tile size of 64 and 87.5% attention sparsity achieves performance comparable to full attention while maintaining efficient kernel execution. Furthermore, purposely injecting locality heuristics proved unnecessary. A large sweep of scaling experiments, in which we pretrain video DiTs from scratch (60M to 1.4B parameters with up to x 10 FLOPS) with 16K sequence length, uncovers a Pareto frontier where VSA achieves near 8x reduction in attention FLOPS and 2.53x reduction in total training FLOPS.
To support VSA's 's hierarchical sparse attention, we prototype GPU kernel where the coarse stage kernel fuses softmax, Top-K. selection, and block indexing into single pass, and the fine stage leverages a block-sparse attention kernel based on FA [5]. This leads to our VSA implementation that retains 85% of FA3's MFU [31]. We further retrofit VSA into SoTA open source DIT, Wan2. 1.3B [35], originally trained with full attention. This integration speeds up attention time by 6x and reduces end-to-end inference latency from 31s to 18s (1.7x) on H100. As a result, VSA enables the first video DiT where attention accounts for only 20% of runtime during both training and inference without quality degradation.
To the best of our knowledge, VSA is the first trainable sparse attention approach that, based on extensive experiments totaling around 90k H200 hours, shows better scaling than full attention on DiTs. Finally, we hope that our explicit ablation of the key parameters (e.g., tile size, critical token prediction, locality prior, sparsity, etc.) will enable more targeted exploration of sparse attention in scaling video DiTs.
2

---

Coted Output
Coarse CorrseOuputt
Resaooenng
Attertioicc. K.
Consowetion
(b) Impact of Tile Size
Gibbi
Fibel
cobe
Guoed
Cuse Fartbon
0
6
8
91011
(8) VSA Architecture
Cube Partition
Figure 1: Overview of VSA. (a) VSA introduce hierarchical attention with sparsity and different granularity
(coarse & fine). (b) Larger sizes (left) blur attention pattern while small tiles coarse stage localize critical tokens close to token resolution. The red dots indicate critical tokens. (c) An illustration of (2, 2, 2) cube partition (in practice we use (4, 4, 4) in VSA).
2
Methods
This section introduces VSA, our sparse attention designed to reduce both training and inference costs of video DiTs. We begin by detailing the design space, key components, and core motivations behind sparse attention in $ 2.1. Although the final method is presented in $ 2.2, emphasize that design exploration and ablation studies (deferred to $ 3) are central to understanding the effectiveness of VSA. Finally, 2.3 describes how to adapt VSA to pretrained Video DiTs originally trained with full attention.
2.1 Sparse Attention Design Space
Modern video DiTs use 3D full attention capture dependencies across the entire video volume. Given a video latent shape (T, H, W). it is flattened into a ID sequence of length - THW by mapping each token location (t, w, h) in the 3D latent to its position n in the ID sequence following n = tHW + hW + W. Then full attention is applied across the entire 1D sequence, allowing each token interact with all the others. Let Q K, V E R Lxd denote the query, key, and value matrices for a single attention head; let M E (0, 0) 1. denote the attention mask specifying the allowed connections between each pair of tokens. The attention output O is then calculated as:
S QK VTTe
= Sofumaxx + M), -=AV
Block Size vs. Hardware Efficiency. In full attention, all entries in M are zero. Sparse atlention and theoretically reduces the total FLOPS by allowing the computation of the corresponding elements in both QK" and AV to be skipped. However, modern accelerators are optimized for dense computation, making unstructured sparsity ineffective for real speedup. Block- sparse attention [5] addresses this by structuring the sparsity to align with the hardware capabilities. In this approach, the attention mask M is divided into tiles of size (By. B&) 2, with all entries of
introduces -o entries
'ee the word block and tile interchangeably in this paper. They refer submatrix that GPU threadblock loads into SRAM when performing matrix multiplication.
'Technically (B,, B&) can be non-square with different values for B, and Br. To keep the notation simple,
assume that B Bg Be
3

---

each tile sharing the same value. This enables each tile in GPU SM to be processed as dense block or skipped entirely, maximizing hardware efficiency. The tile size B is key design parameter: smaller tiles allow flexible, fine-grained sparsity but are less efficient on hardware, while larger tiles improve throughput but restrict the model to coarser attention patterns, potentially reducing modeling expressiveness (see Figure (b)). Thus, selecting B involves a tradeoff between expressiveness and efficiency. In practice, we find that modest reductions in speed can be acceptable if they yield significant improvements in generation quality.
Prediction Cost vs. Coverage Quality in Critical Token Selection. Recent studies have shown that the attention score matrix A is naturally sparse [37, 48, 6, 47], with most values close to zero. This suggests that, by constructing a mask M that preserves only the high-value regions of A the so-called critical tokens- we can closely approximate full attention while significantly reducing computation. A central design choice is how much computation to spend identifying these critical tokens. Computing full attention scores provides the most accurate selection, but largely negates computational savings, as only the AV operation benefits from sparsity. contrast, fixed patterns (e.g., windowed or spatiotemporal attention) incur no prediction cost but often miss informative tokens. Inspired by NSA [45] and MoBA [25], we propose a lightweight, trainable, coarse-granule attention module to estimate the locations of critical tokens without fully computing A. The main challenge is to balance prediction accuracy with computational efficiency for practical use in DiT architectures.
Maintaining Global and Local Context in Sparse Attention. A key challenge with sparse attention is its restricted receptive field, which can limit the model's ability to capture global context. One approach to address this to augment sparse attention with a lightweight global module that captures coarse global signals. Conversely, incorporating local context- motivated by the locality priors commonly used in vision models such as CN also poteantially enhance feature learning. We empirically ablate both strategies to assess their impact on video generation quality in $ 3.1.
2.2
VSA: Video Sparse Attention
VSA employs cube-based partitioning strategy followed by a two-stage attention mechanism to efficiently process video latents. The method first divides the input video latent into spatially and temporally contiguous cubes, then processes these cubes through coarse stage that identifies important regions and fine stage that performs detailed token-level attention within these regions. This design enables efficient computation while maintaining the ability to capture both global and local dependencies in the video data.
Given a video latent with shape (T, H. W), VSA divides it into multiple cubes, each with shape (C,, Ch. Cwe (Figure (c)). VSA co-designs the sparse attention algorithm and its kernel imple- mentation by mapping each cube in the video latent into single tile on GPU SM, where the tile size B = C, x Cp x C.. We assume the video latent shape (T, H, W) is an integer multiple of the tile size and define (Ne, NN,, =
When flattening the 3D video latent into a ID sequence, each token at position (t,, h, w) assigned a ID index n using the following mapping:
T C.
W
Ch C..)
n- (1a. No
Now
+
w
B (t mod C.)OC + (h mod Ca)Cu + (w mod Ow).
Ch
Building on this based partitioning, VSA implements its two-stage attention mechanism to efficiently predict critical token locations without computing the full attention matrix A, as shown in Figure (a). In the coarse stage, we apply mean pooling over each (C,, Cn, C..)) cube to obtain cube-level representations, producing Qe, Ke, V. E R xd This stage then computes attention scores A. E IR 4 & and outputs O.. The attention mask M is derived by selecting the Top-K entries per row in Ae and setting others to -oo, followed by broadcasting to full-resolution mask of size L x L. This mask naturally conforms a block sparse structure because the coarse stage operates on cube-level representations. When broadcasting the mask from B x L
B to I, x L,, each selected entry in A, expands into B X B block in M This block-sparse pattern is crucial for hardware efficiency as allows the next stage to process attention in contiguous blocks that align with GPU memory access patterns and enable efficient parallel computation.
In practice do not broadcast to full-resolution mask but only input the selected block index to fine stage.

---

Next, in the fine stage, this mask M guides fine- grained attention computation for Q K. V E Rexd yielding output Oç. Finally, outputs from both stages are combined to obtain the final output O:
02 C Ge + O
where Ge and Gp are gating vectors obtained from linear projections of the input hidden states. Since the coarse stage introduces negligible computational cost (less than 1% of total FLOPS), the overall sparsity can be approximated by T. cannot be directly applied to the coarse stage, resulting in increased latency. $ 2.4 discusses how we mitigate this overhead. Appendix B gives pseudo code implementation of VSA.
KBB. However, due to the row-wise Top-X selection, FlashAttention
In 3.1 and 3.2, we show that using smaller B leads more expressive sparse attention pattern and improved performance, but comes at the cost of slower attention kernel execution. Setting B 64 with (C,, Ch. Cwe) = (4, 4, 4) provides favorable trade-off between expressiveness and efficiency. We further demonstrate that combining a coarse stage for global context with a fine stage for token-level sparse attention is both necessary and
modules for local context modeling offer minimal additional benefit. Setting KC = 32 consistently yields strong performance across a wide range of sequence lengths. VSA adopts these hyperparameters as default.
2.3 Sparse Adaptation & Distillation
VSA is designed to train video DiTs from scratch, reducing both training and inference FLOPS. It can also be adapted to pretrained video DiTs originally trained with full attention. However, directly replacing full attention with VSA leads unstable training. We hypothesize two main reasons: (1) the gating projection weights G are not present in the full attention checkpoint and are randomly initialized, and (2) VSA differs significantly from full attention, introducing a coarse stage and sparsity in the fine stage. To address this, we develop an annealing strategy that smoothly transitions the model from full attention to VSA. We initialize the weights of the coarse gate Ge to zero. and remove the fine gate GF (equivalent to G3 = 1). We also initialize the sparsity level by setting K == r. effectively making VSA equivalent to full attention the start of training. As training progresses, we gradually reduce IC to the target sparsity level. Meanwhile, G. is updated through training, enabling the model to learn how to balance the contributions of the coarse- and fine stages.
In & 3.3, we demonstrate that Wan-2.1 can be efficiently converted into a sparse-attention model by finetuning with standard flow-matching loss for only a small number of steps. We further conduct a preliminary study on integrating sparse attention with distillation [44]. In this setup, the student trained to act simultancously as a few-step generator and a sparse generator. while the teacher remains unchanged with full attention. Importantly, we preserve the original distillation loss and all hyperparameters, requiring no modifications beyond adapting the student. To our knowledge, VSA is the first sparse attention method shown to be compatible with distillation, whereas prior approaches that exploit diffusion time-step redundancy may fail in the extremely low-step distillation regime.
2.4 Kernel Implementation
VSA requires implementing both forward and backward kernels. We write block-sparse attention kernel with ThunderKittenss [32] for fine stage. Despite using a relatively small tile size of 64, our kernel achieves over 85% of FA's MFU ($3.4). The coarse stage, illustrated in Figure 1, requires row-wise Top-K selection over the cube-level attention matrix. This step necessitates materializing the attention matrix, which precludes direct use of FA-style fused kernels.
One possible workaround is to modify the FA kernel to incorporate in-kernel bitonic sorting for Top-K, thereby avoiding materialization. However, such fusion demands intrusive kernel rewriting and careful tuning. We instead ask: is this complexity necessary? For VSA, the coarse stage operates on (4, 4, 4) cubes, reducing sequence length by 64x, e.g., a 100K-token sequence reduces 1.5K. At this scale, the memory overhead from materialization is negligible. FLOPS wise, the coarse stage contributes less than 0.2% of total attention compute, and we show in $3.4 that its runtime accounts for only 14%, even when the fine stage is 87.5% sparse. This makes further kermel fusion unnecessary. Nonetheless, we still make some efforts to speed up the coarse stage. Our block-sparse kernel consumes block indices, not binary masks. Therefore, converting the Top-K. mask into index form incurs additional overhead. To mitigate this, we fuse softmax, Top-K. selection, and -o-deex conversion into single kernel. This fused kernel reduces coarse stage runtime modestly($3..4)
5

---

Exp ID
Case
Loss Opt
Loss Over
Exp ID
Case
Loss
1
Compress KV
0.15281
0.14282
7
L
0.13330
2
Spatial Temporal
0.13574
0.13034
8
F (no Oc)
0.13296
3
Spatial Full
0.13555
0.12811
9
C&L
0.13220
C&F
0.13162
4
Strided Window
0.13271
0.12716
C&F&L
0.13194
5
Full
0.13877
0.12703
C&F&L&E
0.13124
6
VSA
0.13162
0.12687
13
C&(++L)
0.13192
(a) VSA v.s. other altention.
(b) Different attention design.
B
(Cr. Ca. Cw)
Exp ID
C Pooling
File
TFLOPS
Loss
Exp ID
Pooling
Loss
256
128
(4, 8, 8)
(4, 8, 4)
15
14
256x256
128x128
256x256
128x128
478
444
0.13244
0.13375
19
Conv
0.27787
16
64x64
256x256
478
0.13328
Max
0.13929
64
16
(2, 4,
(4, 4, 4)
17
18
64x64
16x64
64x64
16x64
408
181
0.13155
0.13162
Avg.
0.13162
(c) B to (Ct, Ch. Cw).
(d) Tile size ablation.
(e) Critical token prediction
study.
Table I: Ablation results for key design parameters of VSA.
Experiments
3.1 Ablation Studies
We conduct extensive pretraining experiments to ablate various design choices. Our experiments are based on the Wan2.1 model architecture, a sfate-of-the-art open-source video DIT. Unless otherwise specified, we train models with 120M parameters from scratch for 4.5 x 1020 FLOPS using video latents of shape (16, 32, 32) from the Vchitect- T2V-attveeee dataset [10]. We determined 4.5 x 1020 to be compute-optimal for our setup by following established scaling laws [15, 19]: when comparing models of different sizes under fixed compute budgets, we observe that a 120M parameter model with full attention trained at 4.5 x 1020 FLOPS achieves better performance than smaller models (60M) with the same compute budget, while increasing compute beyond this point yields diminishing returns. For VSA and its variants, we set the number of selected KV tiles K. to 32, achieving an attention sparsity of approximately 87 87.5%. Detailed experimental setup and the rationale behind our experiment design can be found in Appendix C. Key findings from these ablations are presented below.
Data- Peta--ependeett Trainable Sparsity Wins Over Fixed-Patterns. We first investigate why full attention predominates despite sparse alternatives. Table shows existing sparse methods (Exp 1-4) outperform full attention (Exp 5) with compute-optimal training budget (4.5 x 1020 FLOPS), but this advantage reverses with extended training (4 x 10 FLOPS). VSA (Exp 6) outperforms both previous fixed-pattern methods and full attention. To understand why, in Table lb(b) examine two key factors: pattern type and stage contributions. We compare data-dependent patterns against fixed local patterns ("L") that use (, 3) window. Simultancously, the coarse stage by including or excluding its output O (denoted as "C") from the final attention output, versus using only the fine stage output (denoted "F"). Data-dependent patterns consistently outperform fixed patterns, both without the gated residual (Exp vs. 8) and with (Exp 9 vs. 10), demonstrating the inherent advantage of adaptive sparsity and the gated residual.
investigate the impact of
Global Information is Necessary; Locality Priors Offer Limited Gains. We tested three ap- proaches for incorporating local context: (1) adding separate local stage for (3, 3,3) window attention (Exp 11), (2) explicitly excluding ("E") cubes selected by the local stage from the fine stage (Exp 12), and (3) forcing the fine stage to include local cubes, without separate local stage (Exp 13), All three variations performed similarly to the simpler & F architecture, indicating that explicit local modeling provides minimal benefit. VSA therefore adopts the simpler C & F architecture (Exp 10), which effectively captures both global and local information.
Finegrained Attention Map Leads to Better Performance But a Slower Kernel. As analyzed in Section 2.1, , the tile size in V is a critical hyperparameter that balances computational efficiency against model performance. It directly affects two key aspects: (1) how accurately critical tokens can be identified through the coarse stage's 's cube size, and (2) the granularity of attention in the fine stage. Hardware constraints partially dictate this Coaumttrr-NIDDA Hopper GPUs optimize for

---

(a) VSA Full Attention
(b) Loss-FLOPS Pareto Frontier
0.16
0.16
VSA
Fuil
0.15
Loss
Validation
0.14
0.14
0.13
0.12
0.12
5000
10000
15000 Step
20000
25000
30000
10 5 1020
1 x 10
3 10 5 1011
FLOPS
(c) Top-K v.s. Sequence Length
(d) Validation Loss v.s. FLOPS (Top-K)
0.16
0.22
Seq=B114
lop-1b
Validation Loss
0.15
Seq=24576
Seq-16384
Velidation Loss
0.20
Tep-52
Top-64
Scq-61440
0.18
Top128
0.14
10 FLOPS
0.16
0.13
0.14
0.12
Top3
Top16
Top32 Top-K
Top61
Top128
0.12
1020
10
FLOPS
Figure 2: VSA scaling experiments. (a): Video DiT trained with VSA achieves similar curve compared to one trained with full attention. (b): VSA consistently produces a better Pareto frontier when scaling model size up to 1.4B. (c) & (d): The optimal Top-X value (dictating sparsity) depends on both sequence length and training compute. A larger K. is needed for larger training budget.
matrix multiplications with dimensions divisible by 16, and smaller tiles generally reduce arithmetic intensity. Our benchmarks in Table 1d show that decreasing tile size from 256 256 to 64 x 16 significantly reduces Model FLOPS Utilization (MFU).
Training with various tile sizes (Table Ic) reveals smaller tiles consistently reduce model loss through finer attention granularity. This improvement stems from the finer granularity allowing the coarse stage to more accurately predict critical-token cubes and the fine stage to focus attention on smaller, more relevant regions. Experiment 16 specifically tests mismatched granularity between stages, confirming that both coarse and fine stage granularity significantly impact performance. Here, the coarse stage used smaller pooling cubes (C,, Ck. Cw)) = (4,4, 4,44 (effectively B = 64) while the fine stage operated larger tiles corresponding to (Ce, Cn, Cw) (4, 8, 8) (effectively B = 256). To reconcile the finer granularity of the coarse stage's ' predicted attention map with the coarser block-sparse attention in the fine stage, applied additional 21, average pooling before selecting top-IC entries.
Balancing these findings, we selected 64 x 64 tiles ((Cr, Cn, Cw) = (4, 4, 4)) our default configu- ration. While 64 16 tiles offer slightly better performance, they run 2.26x slower (Exp 18 vs. 17), making this tradcoff unfavorable.
Mean Pooling Is Sufficient. We also examined different pooling methods for the coarse stage. Table le shows that average pooling outperforms both max pooling and convolutional approaches, with the latter causing training instability.
3.2 Scaling Studies
To validate VSA, we pretrained a 4 410M video DiT with latent shape (16, 32, 32) (16,384 tokens), larger than our 120M ablation models. Figure 2(a) shows VSA achieves nearly identical loss to full attention despite 87.5% sparsity (KC = 32 out of 256 cubes), while reducing attention FLOPS by 8x and end-to-end training FLOPS by 2.53x. Further scaling experiments from 60M to 1.4B parameters (Figure 2(b)) confirm that VSA consistently produces a better Pareto frontier than full attention. The parallel fitted curves indicate that VSA maintains its 2.53x FLOPS reduction across scales. Each model was trained with compute budgets up to 4 x 1021 FLOPS on 128 H200 GPUs with sequence length of 16K. To our knowledge, VSA the first trainable sparse attention for video DiTs demonstrating superior performance compared to full attention under rigorous scaling
7

---

VSA-1.36
Model
Qual.
Sem.
Total
SVG-1.38
55.5%
25.0%
19.5%
VSA-14B
Dri-Wan
83.71%
77.98%
82.56%
FullAtt-144 148
VS
27.0%
47.0%
26.0%
Sparse-Cistill
Full f.t. 84.07%
81.85%
83.63%
V5
27.0%
50.0%
23.0%
Distil-only
VSA f.t.
83.60%
79.47%
82.77%
20%
60%
80%
100%
Win Fil Loss
(a) VSA Wan- 1.3B results on VBench. With sparse adaptation, VSA is able to achieve similar score to full attention counterpart and even slightly outperform the original model.
(b) Top: VSA vs. SVG human evaluation. SVG has 82.5% attention sparsity with (fp0.03, A000225, s0.1). Middle: VSA v.s. full attention with finetuning. Bottom: VSA v.s. full attention with distillation.
evaluation. While we leave comprehensive scaling studies longer sequence lengths to future work, our fine-tuned Wan 2.1 already shows 1.7 x inference speedup at 23K sequence length ($3.3).
An important design question for VSA is determining the optimal sparsity level via the Top-IC parameter. In Figure 2(c), we pretrained 120M models with varying sequence lengths under fixed 4.5 1020 FLOPS budget. Surprisingly, K. 32 performs consistently well across sequence lengths of 8192, 16384, and 24675, but underperforms compared to IC 16 at 61440 sequence length. This contradicts the conventional intuition that longer sequences require higher K. values. Further investigation with increased compute budget at 61440 sequence length (Figure 2(c)) reveals that
K = 32 eventually outperforms K 16 x 1021 FLOPS, with similar patterns at other lengths. These findings suggest that optimal IC' depends on both sequence length and training budget. We hypothesize that the ideal Top-X. increases with available compute, converging to full attention with infinite resources. However, precisely predicting optimal IC given budget, model size, and sequence length remains an open question. One promising direction is to explicitly model sparsity as an additional axis in the scaling law framework [19, 15], alongside model size and total FLOPS. Incorporating inference costs further complicates this analysis, as higher values may improve training loss but increase inference overhead. We leave a comprehensive treatment of these tradeoffs to future work.
3.3 Sparse Adaptation & Distillation
To evaluate VSA's effectiveness in a post-training setup, we finetune Wan2.1-1.3B with VSA on synthetic data generated by Wan-14B with video latent 16 28 52 (480P). We set K. to 32, corresponding to a 91.2% attention sparsity. As shown in Table 3a, VSA achieves even higher VBench [16] score compared to the original Wan-1.3B. We hypothesize training with synthetic data from larger model may contribute to this boost. To ensure a fair comparison, we also finetune Wan-1.3B using the same synthetic data. The results show that all models perform closely on VBench, indicating that VSA can retain generation quality despite significant attention sparsity. We additionally compared VSA to SVG [37], training-free attention sparsification method, under extreme sparsity. Figure 3b Top shows that VSA is preferred even though it has a higher sparsity, demonstrating the effectiveness of training with sparse attention. With VSA, the DiT inference time of Wan-1.30 33 drops from 31s (full attention with torch compile) to 18s.
To further validate the effectiveness of VSA across different model size and resolution, we scale our study to the Wan-14B model, finetuning on 720P synthetic data (latent 20 x 48 80) 90% sparsity. For this setup, we conduct human preference study on 200 randomly sampled MovieGen prompts [29] to compliment our 1.3B VBench results. As shown Figure 3b Middle, human evalua- tion demonstrates that VSA preserves generation quality compared to the official full attention model, indicating that VSA is capable of maintaining high-quality performance larger model scales. As pilot study, we further explore whether sparse attention can complement other acceleration techniques, particularly distillation. In this setup, the student model in DMD2 [44] employs VSA while the teacher remains unchanged with full attention. All sparse distillation losses and hyperparameters are kept identical to full-attention distillation, as detailed C.6. This simple substitution yields better efficiency (50.9x acceleration) with quality drop, shown in Figure 3b Bottom.
8

---

3,4 Kernel Performance
As Figure 4b shows, VSA's fine block sparse kernel approaches the theoretical limit with nearly 7x speedup over FlashAttention-3 at long sequence lengths (85% MFU over FA3). Even after accounting for the coarse stage computations, VSA still maintains over 6x speedup. In contrast, FlexAttentio [8] with an identical block-sparse mask (64x64 block size) achieves only 2x speedup. Applying VSA 's speedup Wan-1.33 and Hunyuan brings 2-3x inference speedup, as shown Figure 4a.
(Zleseine) VSA
Heskeessful VSA (ite Only)
Fleskith-Spane
0000
-... Theorelial
Wan
Acceleration Ratio
complie
Wan
+VSA
Hunyuan
4-
Original
9-
Hunyuan
Hunyuan
+VSA
0.0
0.25 Normalized Execution Time
0.5
0.75
1.0
8192
16384
32768
65536
131072
Sequence Length
(a)
(b)
Figure 4: Kernel benchmarks. (a): Runtime breakdown of single transformer black for Wan and Hunyuan. VSA reduces attention latency by 6x. (b): Speed VSA fixed 87.5% sparsity under various sequence length with head dim 64. VSA approach the theorectical 8x speedup over FA3.
3.5 Inspecting VSA
To gain deeper insight into VSA's mechanism, we inspect the block-sparse attention maps generated by the coarse stage of our finetuned 1.3B model. As illustrated in Figure 5(a-f), the predicted attention patterns are highly dynamic, confirming our hypothesis that effective sparse attention must be data-dependent rather than relying predefined structures. Even within a single layer, different attention heads often exhibit markedly distinct behaviors. Many observed patterns echo established heuristics, such as local attention focused tokens near the query (akin sliding tile attention), or couporal-widthe plane (e). or the samporal-height plane. Conversely, some patterns deviate from simple heuristics, displaying highly global characteristics (b) or a combination of local and global focus (c).
attention concentrating on tokens within the same frame (d), the same
We quantify the accuracy of critical token prediction calculated as the sum of attention scores within the top-32 cubes selected by the coarse stage. As a baseline, random selection of 32 cubes from the 386 total (for a (16, 28, 52) latent) captures only 8% of the attention score, as shown by the red plane in Figure 5(e). In stark contrast, VSA maintains high accuracy rate, consistently achieving least 60% in most layers and timesteps, and reaching as high as 90% some instances. This underscores VSA's strong capability in identifying critical tokens. Critically, even if the fine stage misses a small portion of the attention weight, the direct output from the coarse stage can potentially compensate for this. Further examination of Figure 5 (c) reveals systematic variations in prediction accuracy. Accuracy tends to increase monotonically with the timestep. Across transformer layers, however, the accuracy rate exhibits zig-zag pattern. These accuracy dynamics across layers and timesteps suggest avenues for future optimizations with adaptive Top-X. value.
Related Work
Sparse Attention in LLMs. There has been a proliferation of fixed-pattern sparse attention mech- anisms in large language models [3, 46, 1, 7, 13]. In practice, however, most LLM training (over 90% of total FLOPs) happens on short sequences (<32K tokens) under the "train-short, adapt-long" paradigm [40, 23, 12, 42], so sparse attention saw little uptake beyond sliding window attention variants like in Mistral [17]. With LLMs targeting contexts beyond 1M tokens, sparse attention has seen renewed interest. Recent work primarily targets inference-time speedups [39, 49, 18, 41]. while the latest methods explore trainable, dynamic sparsity patterns (MoBA [25], NSA [45]) to enable efficient end-to-end training extreme-length sequences. We draw inspiration from them; However,

---

Step 1 Layer Head 8
Step Layer 1 Head
Step Layer Head
Acc.
1.0
(x) Local Head
(b) Global Head
(c) Local Global Head
Step Layer Head
Step I Layer 20 Head 5
Step Layer
0.4
0.2
40
25
20
15
20
Timestep
Layer Idx
10
0
(d) HW Head
Head
(1) Head
(e) Attn Recell
Figure 5: Visualization of the attention pattern of VSA. (a)-(f): VSA dynamically select different cubes to attend, where the blue cube indicates query and red cubes indicated selected key and values.(e): VSA critical-token prediction accuracy.
VSA differs from MoBA by directly contributing the coarse- poare-graiiee attention output to the final representation and using smaller blocks compatible with efficient block-sparse kernels. Compared to NSA, the nature of video and bidirectional attention avoids grouped query constraints of the attention pattern. We discuss similarities and differences in depth in SE.
Sparse Attention in Video DiTs. Recent work has explored applying sparse attention post-hoc to DiTs pretrained with full attention [48, 6, 47, 37, 14] at inference. However, we argue that the for trainable sparse attention in video DiT's is both distinct from and more urgent than LLMs. First, video DiT demands far longer sequences, e-g., 100K-token context yields only 5s video, making scaling inherently more costly than language models. Second, unlike LLMs, where long-context adaptation is small fraction total training, state-of-the-art video DiTs [20, 35, 29, 33] dedicate most of their compute budget to full-resolution, long-sequence training. As a result, these models remain bottlenecked by quadratic attention at both training and inference. This calls for trainable sparse attention mechanisms, like VSA, as core design of video DiT's, not post-hoc fix. DSV [33] also explores adding sparsity to attention during DiT training, but their multi-stage and profiler-based design may complicate the training pipeline, which we further discuss in $E.
5 Limitation and Conclusion
We present VSA, a trainable and hardware- hardware-efffcient efficient sparse attention tailored for scaling DiTs. Unlike prior work that applies sparsity post-hoc, VSA jointly learns to predict and apply attention sparsity training and remains compatible with block-sparse compute layout. VSA currently operates with fixed cube size of (4, 4, 4), which requires video latent dimensions to be divisible by 4. While this may restrict the set compatible resolutions, it can be addressed practice by generating a slightly larger latent and cropping the target shape. Another open question is how to determine the optimal sparsity level. While our scaling experiments ($3.2) provide preliminary insights, a complete understanding may require extending scaling laws explicitly account for sparsity, alongside model size and training compute. Across diverse model sizes (60M to 1.4B) and budgets (up to 4 x 1021 FLOPS), show that VSA matches the performance of full attention at 2.53x lower training cost, and achieves 85% MFU of FA3. When integrated with Wan2. 1-1.BB, it reduces end-to-end latency by 1.7x. We hope this work establishes trainable sparse attention as a practical and scalable alternative to full attention in further scaling video DiTs.
Acknowledgments
We would like to thank Wei Zhou, Kevin Lin, Matthew Noto, Wenxuan Tan, and Jinzhe Pan for helpful discussion. The work is supported by UCSD HDSI, Nvidia, and a faculty research award from Google. The computing resources were provided by MBZUAI IFM and Nvidia's s donation.
10

---

References
[1] Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05150, 2020.
[2] Shoufa Chen, Chongjian Ge, Yuqi Zhang, Yida Zhang. Fengda Zhu, Hao Yang. Hongxiang Hao, Hui Wu, Zhichao Lai, Yifei Hu, et al. Goku: Flow based video generative foundation models. arXiv preprint arXiv:2502.044940 2025.
[3] Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10009 2019.
[4] Hyung Won Chung, Noah Constant, Xavier Garcia, Adam Roberts, Yi Tay, Sharan Narang, and Orhan Firat. Unimax: Fairer and more effective language sampling for large-scale multilingual pretraining. arXiv preprini arXiv:2304.0915), 2023.
[5] Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness, 2022.
[6] Hangliang Ding, Dacheng Li, Runlong Su, Peiyuan Zhang, Zhijie Deng, Ion Stoica, and Hao Zhang. Efficient-vdit: Efficient video diffusion transformers with attention tile. arXiv preprint arXiv:2502.06155, 2025.
[7] Jiayu Ding, Shuming Ma, Li Dong, Xingxing Zhang, Shaohan Huang, Wenhui Wang, Nanning Zheng, and Furu Wei. Longnet: Scaling transformers to 1,000,000,000 tokens. arXiv preprint arXivv2307.02488. 2023.
[8] Juechu Dong, Boyuan Feng, Driss Guessous, Yanbo Liang, and Horace He. Flex attention: A pro- gramming model for generating optimized attention kernels. arXiv preprint arXiv:2412.05490 2024.
[9] Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, et al. Scaling rectified flow trans- formers for high-resolution image synthesis. In Forty-first international conference machine learning, 2024.
[10] Weichen Fan, Chenyang Si, Junhao Song, Zhenyu Yang, Yinan He, Long Zhuo, Ziqi Huang, Ziyue Dong, Jingwen He, Dongwei Pan, et al. Vchitect-2.0: Parallel transformer for scaling up video diffusion models. arXiv preprint mrXiv:2501.084533 2025.
[11] Genmo. Mochi 1. stttps://gitee http://ggtttuuu.oooggommoo// wwt.////mmmm//mmmmee...oomm 2024.
[12] Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al. The llama
3 herd of models. arXiv preprint arXiv:2407.21783 2024.
[13] Mandy Guo, Joshua Ainslie, David Uthus, Santiago Ontanon, Jianmo Ni, Yun-Hsuan Sung, and Yinfei Yang. Longt5: Efficient text-to-text transformer for long sequences. arXiv preprint arXiv:2112.07916 2021.
[14] Ali Hassani, Fengzhe Zhou, Aditya Kane, Jiannan Huang, Chich-Yun Yun Chen, Min Shi, Steven Walton, Markus Hoehnerbach, Vijay Thakkar, Michael Isaev, et al. Generalized neighbor- hood attention: Multi-dimensional sparse attention at the speed of light. arXiv preprint 2504 16922, 2025.
[15] Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas. Lisa Anne Hendricks. Johannes Welbl, Aidan Clark, al. Training compute-oper -pttimal large language models. In Proceedings of the 36th International Conference on Neural Information Processing Systems, pages 30016-30030, 2022.
[16] Ziqi Huang, Yinan He, Jiashuo Yu, Fan Zhang, Chenyang Si, Yuming Jiang, Yuanhan Zhang, Tianxing Wu, Qingyang Jin, Nattapol Chanpaisit, et al. Vbench: Comprehensive benchmark suite for video generative models. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition, pages 21807-21818 2024.

---

[17] Albert Q. Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, al. Mistral 7b, 2023.
[18] Huiqiang Jiang, Yucheng Li, Chengruidong Zhang, Qianhui Wu, Xufang Luo, Surin Ahn, Zhenhua Han, Amir Abdi, Dongsheng Li, Chin-Yew Lin, et al. Minference 1.0: Accelerating pre-filling for long-con congecontext Ilms via dynamic sparse attention. Advances in Neural Information Processing Systems, 37:52481-52511 2024.
[19] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint wxXi::2001008300 2020.
[20] Weijie Kong, Qi Tian, Zijian Zhang, Rox Min. Zuozhuo Dai, Jin Zhou, Jiangfeng Xiong, Xin Li, Bo Wu, Jianwei Zhang, et al. Hunyuanvideo: A systematic framework for large video generative models. arXiv preprint (rXiv:2412.03660) 2024.
[21] Bin Lin, Yunyang Ge, Xinhua Cheng, Zongjian Li, Bin Zhu, Shaodong Wang, Xianyi He, Yang Ye, Shenghai Yuan, Liuhan Chen, al. Open-sora plan: Open-source large video generation model. arXiv preprini 24Xi::24/2.01133 2024.
[22] Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In The Eleventh International Conference on Learning Representations.
[23] Aixin Liu, Bei Feng, Bin Wang, Bingxuan Wang, Bo Liu, Chenggang Zhao, Chengqi Dengr, Chong Ruan, Damai Dai, Daya Guo, al. Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model. arXiv preprint arXiv:2405.044 2024.
[24] Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209.03003 2022.
[25] Enzhe Lu, Zhejun Jiang, Jingyuan Liu, Yulun Du, Tao Jiang, Chao Hong, Shaowei Liu, Weiran He, Enming Yuan, Yuzhi Wang, et al. Moba: Mixture of block attention for long-context 1lms. arXiv preprint arXiv: 13189, 2025.
[26] Guoqing Ma. Haoyang Huang, Kun Yan, Liangyu Chen, Nan Duan, Shengming Yin, Changyi Wan, Ranchen Ming. Xiaoniu Song, Xing Chen, al. Step-video-12v technical report: The practice, challenges, and future of video foundation model. arXiv preprint arXiv:2502.10248, 2025.
[27] Xin Ma, Yaohui Wang, Gengyun Jia, Xinyuan Chen, Ziwei Liu, Yuan-Fang Li, Cunjian Chen, and Yu Qiao. Latte: Latent diffusion transformer for video generation. CoRR, 2024.
[28] William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF international conference on computer vision, pages
2023.
[29] Adam Polyak, Amit Zohar, Andrew Brown, Andros Tjandra, Animesh Sinha, Ann Lee, Apoory Vyas, Bowen Shi, Chih Yao Ma, Ching- Yao Chuang, David Yan, et al. Movie gen: A cast of media foundation models, 2025.
[30] Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content- sparse attention with routing transformers. Transactions of the Association for Computational Linguistics, 9:53-68, 2021.
[31] Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, and Tri Dao. Flashattention-33: Fast and accurate attention with asynchrony and low-precision. Advances in Neural Information Processing Systems, 37:68658-6866 2024.
[32] Benjamin F Spector, Simran Arora, Aaryan Singhal, Daniel Y Fu, and Christopher Ré. Thun- derkittens: Simple, fast, and adorable ai kemels. arXiv preprint arXi:24410.23399 2024.
[33] Xin Tan, Yuetao Chen, Yimin Jiang, Xing Chen, Kun Yan, Nan Duan, Yibo Zhu, Daxin Jiang, and Hong Xu. Dsv: Exploiting dynamic sparsity to accelerate large-scale video dit training, 2025.
12

---

[34] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
[35] Ang Wang, Baole Ai, Bin Wen, Chaojie Mao, Chen-Wei Xie, Di Chen, Feiwu Yu, Haiming Zhao, Jianxiao Yang, Jianyuan Zeng, et al. Wan: Open and advanced large-scale video generative models, arXiv preprint arXiv:2503. 20324.34444 14, 2025.
[36] Yaohui Wang, Xinyuan Chen, Xin Ma, Shangehen Zhou, Ziqi Huang, Yi Wang, Ceyuan Yang. Yinan He, Jiashuo Yu, Peiqing Yang. et al. Lavie: High-quality video generation with cascaded latent diffusion models. International Journal of Computer Vision, (33(5):3059-3775 2025.
[37] Haocheng Xi, Shuo Yang. Yilong Zhao, Chenfeng Xu, Muyang Li, Xiuyu Li, Yujun Lin, Han Cai, Jintao Zhang, Dacheng Li, Jianfei Chen, Ion Stoica, Kurt Keutzer, and Song Han. Sparse videogen: Accelerating video diffusion transformers with spatial-temporal sparsity, 2025.
[38] Yifei Xia, Suhan Ling, Fangcheng Fu, Yujie Wang, Huixia Li, Xuefeng Xiao, and Bin Cui. Training-free and adaptive sparse attention for efficient long video generation, 2025.
[39] Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. arXiv preprint arXiv:2309. / /7453, 2023.
[40] Wenhan Xiong, Jingyu Liu, Igor Molybog, Hejia Zhang, Prajjwal Bhargava, Rui Hou, Louis Martin, Rashi Rungta, Karthik Abinav Sankararaman, Barlas Oguz, et al. Effective long-context scaling of foundation models. In Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers), pages 4643-4663, 2024.
[41] Ruyi Xu, Guangxuan Xiao, Haofeng Huang, Junxian Guo, and Song Han. Xattention: Block sparse attention with antidiagonal scoring. arXiv preprint 2.X.:22503..6222 2025.
[42] An Yang. Baosong Yang. Beichen Zhang. Binyuan Hui, Bo Zheng, Bowen Yu, Chengyuan Li, Dayiheng Liu, Fei Huang, Haoran Wei, et al. Qwen2. 5 technical report. arXiv preprint arXiv:2412.15113 15115, 2024.
[43] Zhuoyi Yang. Jiayan Teng, Wendi Zheng, Ming Ding, Shiyu Huang, Jiazheng Xu, Yuanming Yang. Wenyi Hong, Xiaohan Zhang, Guanyu Feng, et al. Cogvideox: Text-to-video diffusion models with expert transformer. CoRR, 2024.
[44] Tianwei Yin, Michaël Gharbi, Taesung Park, Richard Zhang, Eli Shechtman, Fredo Durand, and Bill Freeman. Improved distribution matching distillation for fast image synthesis. Advances in neural information processing systems, 37:47455-47487, 2024.
[45] Jingyang Yuan, Huazuo Gao, Damai Dai, Junyu Luo, Liang Zhao, Zhengyan Zhang, Zhenda Xie, YX Wei, Lean Wang, Zhiping Xiao, et al. Native sparse attention: Hardware-aligned and natively trainable sparse attention. arXiv preprini arXiv:2502. 1089, 2025.
[46] Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santi- ago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. Big bird: Transformers for longer sequences. Advances in neural information processing systems, 13:17283-17297, 2020.
[47] Jintao Zhang, Chendong Xiang, Haofeng Huang, Jia Wei, Haocheng Xi, Jun Zhu, and Jianfei Chen. Spargeattn: Accurate sparse attention accelerating any model inference, 2025.
[48] Peiyuan Zhang, Yongqi Chen, Runlong Su, Hangliang Ding, Ion Stoica, Zhenghong Liu, and Hao Zhang. Fast video generation with sliding tile attention, 2025.
[49] Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, et al. H2o: Heavy-hitter oracle for efficient generative inference of large language models. Advances in Neural Information Processing Systems, 36:34661-36666 1347611,11 10, 2023.
13

---

[50] Zangwei Zheng, Xiangyu Peng, Tianji Yang, Chenhui Shen, Shenggui Li. Hongxin Liu, Yukun Zhou, Tianyi Li, and Yang You. Open-sora: Democratizing efficient video production for all. arXiv preprini arXiv:2412.20404, 2024.
[51] Lei Zhu, Xinjiang Wang. Zhanghan Ke, Wayne Zhang, and Rynson WH Lau, Biformer: Vision transformer with bi-level routing attention. In Proceedings of the IEEE/CVF conference on 2023.
computer vision and pattern recognition, pages
14

---

A
Qualitative Examples
Prompt: An astronaut walking between stone buildings.
Training Step
Training 750 Step
Training 1500 Step
Training 2250 Step
(a) Qualitative example across training steps
(b) Complete video example 3000 training steps
Prompt: A tilt-up from city street, ascending to show the skyline with its mix of modern and historic architecture.
Training 0 Step
Training Step
Training Step
Training 2250 Step
(a) Qualitative example training steps
(b) Complete video example 3000 training steps
Figure 6: Qualitative examples. In (a), we sample the same middle frame at each step of the video. In (b), we uniformly sample four frames across the video.
We qualitatively illustrate the finetuning process ($2.3$3) Wan-1.3B in Figure All frames are sampled from validation videos at selected training steps with I. = 32. At the start the training, the model exhibits noticeable artifacts when switching from full attention VSA, reflecting the change in attention structure. As training progresses, the model gradually adapts to the sparse attention mechanism and recovers the ability to generate coherent videos.
B Pseudocode of VSA
We provide pseudocode in pytorch- like API for easier understanding of VSA.
def tx)
return rearrange "b h (n_t ts_t n_h ts_h n ts_w) d -> b ( att n.h n_ ts_t te_h ts_w) d", n_t=4, n_h=8, n_v-8, ts_t=4, ta_h=4, ts_w=4)
def untile(x):
return rearrange (x, "b (n_t n ne ta_t ta_h ts_w) d -> b n.t ta_t n_h ta_h n. ts_w) d", n_t4,, n_h=8, n_w=8, ts_t=4, ts_h=4, t8_w=4)
q. k, v, gate tile(q), tile(k), tile(v), tile(g) mtarse_attn_atte gate gate. chunk (2, dim=1)
15

---

Coarse stage
B, H, L, q shape
block 64 topk 32
q-c q. (B. H, L//block, block, D) mean (dim-3)
k_c k. view H, L//block, block, D) mean (dim=3)
v_c vivi H, L//block, block, mean (dim=3)
score torch matmul torch. (q_c, &_c. (-2 -1)) / (D ** 0.5)
torch nnch..n..uute funcciionall softmar (score, dim=-1)
march.mammll (score, v_c) output, coarse output_coares view Н, L//block, 1, D) repeat (1, 1, 1, block, 1). (Be H, L, D)
patput_coarse
Keep only top-k blocks topk_vale vals, topk_idx = score. took (topk, dim=-1) score torch. zeros_ (score) score. scorter (-1, topk_idx, topk_vals)
score score view (B, н, L//block, L//block, 1, 1) 1). repeat 1, 1, 1, block, block) score score. permute 4,, reshape (B, H, L, L)
Fine stage
QK = torch
norch.matmul. matmul (q, k. trannspose (-2, -1)) / (D ** 0.5)
QK = QK baakkedfiil(aattaa _ask == 0, float (" ("-inf"))
QK = torch. mn. functiiiiina..oo softmax dim=-1)
output_fine torch matmul (QK,
Combine stages with residual connection Sadden_stated = output_coarse * coarse_attn_gar + output_fine
fine_attn_gate
hidden_states untile (hidddeddttntsees)) transpose (1, 2) flatten (2,
Listing Pseudocode ViLAS in pytorch-lin API with kernel optimization, assuming cube size (4,4,4)
vidco size (16,32,32). Note that the tilc and untile operation can moved the beginning and end of transformer to avoid calling them for each attention.
C Experimental Details
document the detailed experiments setups for results presented in Section
C.1 Model Architecture
We follow the architecture of Wan2.1 [35] for all experiments. Ablation studies are conducted on 20M-parameter model initialized using the GPT-NeoX scheme, which leads to faster convergence than the default Py Py'oorhh initialization. The model adopts the pretrained UMT5-XXL. [4] as the text encoder and the pretrained VAE from Wan2.1 for video tokenization. The architecture includes two types of attention: (1) self-attention among video tokens and (2) cross-attention for injecting textual information. Sparse attention applied only to the self-attention layers. Detailed model configurations are provided in Table 2a.
C.2 Ablation Experiments Setup
train on long sequences shape 61 512 x 512, motivated by two factors. First, attention dominates the computational cost at this scale, making it the primary bottleneck. Second, sparse attention must evaluated under long-context settings to demonstrate effectiveness; short sequences do not present sufficient challenge.
16

---

Table 2: Model configuration and training hyperparamelers used for the ablation studies.
Model Config
Value
Hyperparameter
Value
Head Dim
64
Learning Rate
6 x 10
FFN Dim
3072
I.R Scheduler
Constant
Cross Attn Norm
True
Warmup Steps
100
Freq Dim
256
Batch Size
1024
In Channels
16
Video Latent Shape
16 x 32 x 32
Out Channels
16
Sequence Length
16,384
Num Heads
12
Attention FLOPs Ratio
Num Layers
12
Weight Decay
1 x 10
Patch Size
1 2 2
Adam W Betas
(0.9, 0.95)
QK Norm
RMS Norm (across heads)
Objective
Flow Matching [24, 22] LogitNormallO.00 1.0) [9]
Epsilon
1 10
Timestep Sampler
Text Dim
4096
Total Traning FLOPS
4.5 x 1000
(a) 120M model used for ablation studies.
(b) Training hyperparameters
To establish strong baseline, we perform a grid search over batch sizes (512, 1024, 2048) and
learning rates {5 10-5, 1 x 10-4, 2 x 10-2, 6 x 10-4). The best hyperparameters is used for all ablation variants. Training conducted under fixed compute budget of 4.5 x 1020 FLOPs, which we find to be sufficient for training a 120M-paaameter model a 120M model with full attention outperforms 60M model trained with 4 10 FLOPs, indicating FLOPS budget around 4.5 x 10 is compute-optimal for 120M model. Each ablation job takes around 10 hours on 64 Nvidia H200 GPU. Full training hyperparameters are provided in Table 2b.
C.3 Baseline Attention Variants
Spatial-Temporal A widely adopted approach in early video generation works, including Open- Sora [50], OpenSora-PPan [21], LaVie [36], and Latte [27]. We alternate between spatial and temporal attention across layers.
Spatial-Full In spatial-tempora attention, the temporal stage can become overly sparse. For example, with latent shape of (16, 32, 32), the temporal attention accounts for less than 1% of the FLOPs of full 3D attention. To mitigate this, we design a variant with four spatial layers and one full-attention layer every five layers.
Compress KV
This variant pools only the key and value tokens using 2 x 2 2 average pooling, reducing attention FLOPs by 8x. The query tokens remain at full resolution. This setup mimics the coarse-grained stage of VSA with smaller pooling size and pooling on query tokens.
Strided Window Inspired by Swin Transformer, propose strided window attention that increases token interaction on top of spatial-tempont attention. Let W and W2 denote the spatial and temporal window sizes. For spatial attention, query attends to all tokens in the same frame and to those in the same temporal window (W, 2). For temporal attention, a token attends to the same spatial location and others in the same spatial window (W, = 8).
Conv Pooling Instead of mean pooling for block-level token aggregation, we use 3D convolution with kernel size and stride of (4, 4, 4) (same as the block size). The output channel dimension matches the head dimension.
C.4 FLOPS Calculation
Elementwise operations such as LayerNorm and RoPE contribute negligibly to the total computational cost in transformers. Following the approximation in [15], we omit these operations and estimate the model FLOPs as 6ND, where N the number model parameters and is the number of input tokens.
17

---

However, for video DiTs trained on long sequences, attention computation becomes a dominant cost. We therefore incorporate the attention FLOPs following the formulation from FlashAttention [5]:
FL.Pss = 6ND + 4. D S A 3.. L.
where D is the number of tokens, S is the sequence length, A is the number of attention heads, H is the head dimension, and L is the number of transformer layers. For sparse attention, we adjust the attention portion according to their sparse patter..
C.5 Sparse Adaptation Setup
To bridge the gap between full and sparse attention, we adopt a sparsity decay schedule that gradually reduces the number of cubes used in attention computation. The model is first trained with full attention for the initial 50 steps, to accommodate the changed resolution and aspect ratio. Thereafter, we decrease the number of attended cubes by 10 (i.e., reduce Top-K by 4) every 50 steps, until reaching the target sparsity level (In our setting Top-K = 32). Unlike directly training the model with extremely sparse attention, our progressive decay schedule enables a smooth transition and mitigates training instability.
In the finetuning experiments for Wan-1. 1.3B, we trained on 80,000 synthetically generated videos from Wan-14B, each with resolution of 448 832 and 61 frames. To accelerate training and reduce memory usage, we preprocessed both the VAE latents and text encodings. Training was conducted on 32 H200 GPUs using DDP as the parallelism strategy. We set the per-GPU batch size to 1, applied gradient accumulation of 2, and used a learning rate of 1e-5. The training ran for 4,000 steps.
In the finetuning experiments for Wan-14B, set the final sparsity to 0.9 and trained on 200,000 synthetic videos from Wan-14B, each with a resolution of 768 x 1280 and 77 frames. Training was conducted on 64 H200 GPUs using DDP the parallelism strategy. We set the global batch size to 64, and learning rate to 1e-5. The training ran for 4,000 steps.
C.6 Sparse Distillation Setup
Our Sparse Distillation Wan-1.3B uses 64 H200 with per-GPU batch size of 1, and runs for 12 hours. We initialize the generator model with the pretrained weights of the base model and then replace the attention module with VSA (sparsity=0.8), while keeping real score and fake score models the same as the base model. All DMD-related hyperparameters are held fixed, including the number of denoising steps, the generator update ratio, and the guidance scale for real-score model. After 4,000 steps training, the 3-step generator achieves a 50.9x reduction in denoising time relative to the baseline model and generate 5-second video in approximately 5 seconds single H200.
D Coarse Stage Runtime
For shorter sequence lengths, the coarse stage over- head is more pronounced. Our profiling experiments using nsys (Table 3) reveal that Top-K. selection dom- inates this runtime. Although we fused the kernels for attention scaling, softmax, and Top-K operations reduce memory traffic and improve data locality, it only provided modest improvements. Since the coarse stage overhead becomes negligible at longer sequence lengths- our primary target we did not pursue further optimizations in this work. However, coarse stage acceleration remains an important direc- tion for future research.
Table 3: Coarse stage runtime (us).
Breakdown
w/o fusion
w/ fusion
QK" scale
0.046
0.046
0.060
softmax
0.095
0.912
topk
0.869
PV
0.045
0.045
E Discussion and Extended Related Work
VSA builds upon insights from prior work trainable sparse attention mechanisms in both language modeling [45, 25] and computer vision [51, 33). This section situates VSA within this landscape, highlighting key similarities and differentiating design choices.
18

---

MoBA [25]; VSA shares conceptual similarities with MoBA, particularly in: (I) employing coarse- grained stage that utilizes mean pooling, akin to MoBA's gating mechanism, and (2) using attention scores from this pooled representation to guide block selection for sparse attention. However, key divergence lies in the utilization of the coarse-grained stage output. While MoBA employs pooled attention solely for block selection, VSA incorporates the output of its coarse-graines stage directly into the final attention output, potentially enriching global context representation. More critically, MoBA's implementation, which relies token gathering and variable-length FlashAttention, con- strains it to larger tile sizes (e.g., 512). This can limit the granularity of its sparsity patterns and reduce speedup efficacy, especially for sequences like those at 128K. In contrast, VSA is implemented with -packe attention leveraging smaller, fixed-size blocks (c.g., 64x64 as per our findings), aiming for better balance between performance (Table 1d) and practical world-clock speedup (Section 3.1).
NSA [45]: The two-stage (coarse/compress and fine/select) architecture of VSA bears resemblance to NSA's design. However, fundamental differences arise from their target domains. NSA is tailored for causal language model decoding, where typically only a single query token is processed at time. This necessitates specific strategies like group query attention to enable efficient kernel implementation(NSA A can only pool Key- Cey-Value pairs). Video DiTs, operating bidirectionally on entire sequences, do not face the same single-query constraint, allowing VSA to apply pooling more broadly and employ distinct sparse patterns for different attention heads without resorting to grouped queries. Furthermore, NSA includes additional sliding window stage for local information, component VSA found unnecessary for video generation in our setup (Table 1b).
BiFormer [51 Similar to BiFormer, VSA utilizes coarse -raniee, tile-to-tile attention mechanism. However, in BiFormer, this coarse attention serves only to derive the dynamic sparse pattern for the subsequent token-to-token attention and does not directly contribute to the final output. Our ablations (Table 1b) indicate that for VSA, the output of the coarse-graines stage is paramount for achieving optimal performance. Additionally, BiFormer's 's original implementation lacked direct FlashAttention compatibility, impacting its throughput compared to VSA 's design, which optimized for hardware- aligned block-sparse operations.
DSV [33]: DSV represents pioneering work in exploring trainable sparse attention specifically for video DiT training. Both VSA and DSV aim to reduce the cost of identifying load-bearing regions. DSV achieves this by introducing dedicated low-rank attention predictors with a reduced head dimension, which are trained in separate, multi-stage process and are not fully end-to-end integrated with the main DiT training. VSA, on the other hand, reduces this cost by performing attention on spatially pooled representations (e.g.. from 4x4x4 cube of tokens) within its coarse-grained stage. Crucially, VSA is designed to be end-to-end trainable, requiring minimal modifications to existing DiT training frameworks, unlike the more complex training system designed for DSV's predictors.
F Broader Impact
VSA aims to make high-quality video generation more accessible by significantly reducing the training and inference cost of video diffusion models. Our work may help democratize video creation tools for a broader range of users and use cases, including education, animation, and independent media production. However, as with many generative technologies, VSA also presents potential risks. In particular, the ability to generate realistic videos at scale may increase the risk of malicious applications, such as generating deepfakes or misleading media content. We emphasize the importance of developing robust detection tools, usage guidelines, and ethical standards in parallel with technical advances to ensure the responsible deployment of such models.
19

---

NeurIPS Paper Checklist
The checklist is designed to encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No], or [NA].
[NA] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1-2 sentence) justificcton right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.
The reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation.
While "[Yes] " is generally preferable to "[No]", ", it is perfectly acceptable to answer "[No] provided proper justification is given (e.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering "[No] or "[NA] is not grounds for rejection. While the questions are phrased in binary way, we acknowledge that the true answer often more nuanced, so please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS Paper Checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only use the provided macros for your answers.
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The main claims made in the abstract and introduction accurately reflect the paper's contributions and scope.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes]
20

---

Justification:
The paper discuss the limitations of the work performed by the authors.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approxinations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proof's
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer:
[NA]
Justification:
No thcoretical result.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: The paper fully disclose all the information needed to reproduce the main experimental results of the paper. The code is opensourced.
Guidelines:
21

---

The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification:
The paper provide open access to the data and code, with sufficient instructions faithfully reproduce the main experimental results.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.cm cc/
menides/CodeSubml.
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https:
/
http://ggiiiee//eeeuuuummi (oodeSumisssiii siionPolcyy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
22

---

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
The paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification:
We follow previous work's standard evaluation process to evaluate our model.
Guidelines:
The answer NA means that the paper does include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall
run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: The paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments.
23

---

Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https: /ttp///euripp...p cc/ppbblccEttiiccuuueeeeeeeeeee
Answer: [Yes]
Justification: The research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
Broader impacts
Question: Does the paper discuss both potential positive socictal impacts and negative socictal impacts of the work performed?
Answer: [Yes]
Justification: The paper discuss both potential positive societal impacts and negative societal impacts of the work performed.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g.. disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
24

---

Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The model we release do not provide additional risks than the open-source base model we use.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (c.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
The creators or original owners of assets (e.g., code, data, models), used in the paper, are properly credited and the license and terms of use are explicitly mentioned and properly respected.
Justification:
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URI
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, haperswithcode com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if has changed) should provided.
If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: New assets introduced in the paper are well documented and is the documenta- tion provided alongside the assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
25

---

The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/teview based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMSs if it is important, original, or non-standard component of the core methods in this research? Note that the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: The core method development this research does not involve LLMS as any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer LLM policy (https: httpp::///euripsst www.rrp..../oooeeeeeeeee cc/Confereecees2285 for what should or should not be described.
26


Reviews:
[{"review_id": "Xbha4iLY77", "reviewer": "Reviewer_oL6k", "summary": "The paper introduces ViSA, a trainable sparse attention mechanism for video diffusion transformers (DiTs) that reduces the computational bottleneck of quadratic 3D attention. ViSA employs a two-stage approach to identify critical tokens and compute attention efficiently, achieving a 2.53× reduction in training FLOPS and a 6× speedup in attention time with comparable quality.", "questions": "NA", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\n1. ViSA proposes a novel trainable sparse attention mechanism for video DiTs, learning critical token positions during training rather than sparsifying only at inference, thus better preserving model performance.\n\n2. The design of ViSA fully takes into account the hardware characteristics of modern GPUs, ensuring compatibility with efficient attention implementations such as FlashAttention through block-sparse computational layouts, thereby translating theoretical computational savings into actual runtime acceleration.\n\n3. The paper includes extensive ablation studies and scaling experiments, demonstrating ViSA's robustness across various configurations (e.g., tile sizes, sparsity levels) and model sizes.\n\n4. This paper is well written and easy to follow.\n\nWeakness: Although ViSA has limitations like fixed cube size and optimal sparsity determination, the authors provide practical solutions like cropping that still significantly enhance video DiT training and inference efficiency.", "was_revised": true, "final_justification": "This paper extends the study of sparse attention to the training stage, achieving competitive performance while enabling VSA to adapt to various resolutions. Its effectiveness in video generation highlights broad practical value, and I therefore recommend acceptance.", "rebuttal": "We thank reviewer oL6k for the positive feedback and for highlighting the effectiveness of ViSA’s trainable sparse attention design. \n\nRegarding the noted limitation on fixed cube size and resolution compatibility, we have updated our implementation of ViSA to support **arbitrary video resolutions via a variable-size block kernel**. Specifically, when the input video shape is not divisible by the default cube size of (4, 4, 4), we now partition it into regular cubes wherever possible and pad the remaining non-full cubes to match the expected size. These padded cubes are processed normally in the coarse stage. For the fine stage, we improve our sparse attention kernel to explicitly mask out padded tokens, ensuring they never contribute to attention computation. This update removes the resolution constraint, enabling ViSA to support any video shape. The overall framework, including the differentiable coarse-to-fine attention and compatibility with block-sparse FlashAttention-style layouts, remains unchanged."}, {"review_id": "BmwqKqpXhL", "reviewer": "Reviewer_WVmh", "summary": "The paper \"Faster Video Diffusion with Trainable Sparse Attention\" introduces VISA (Video Diffusion with Trainable Sparse Attention), a novel method aimed at accelerating video diffusion models by reducing the computational cost of attention mechanisms. The key contribution is a trainable, hardware-efficient sparse attention mechanism that replaces full attention during both training and inference. VISA employs a hierarchical attention approach, consisting of a coarse stage that identifies critical tokens and a fine stage that computes token-level attention within these critical regions. The method is designed to be end-to-end trainable and compatible with modern GPU architectures, achieving significant speedups without compromising generation quality. The authors demonstrate that VISA reduces training FLOPS by 2.53× and inference time by 1.7× compared to full attention, while maintaining comparable performance.", "questions": "- Generalization to Other Datasets: How does VISA perform on other types of video data, such as higher-resolution videos or videos with different content characteristics? Can the authors provide additional experiments or analysis to demonstrate the generalizability of their method?\n\n- Training Overhead: Can the authors provide more details on the computational overhead introduced by the coarse stage during training? How does this overhead compare to the overall computational savings achieved by VISA?", "limitations": "The authors have adequately addressed the limitations of their work by discussing the current constraints, such as the fixed cube size and the dependency on specific hardware architectures. They also acknowledge the potential negative societal impact of their work, such as the risk of generating misleading or harmful content, and emphasize the need for robust detection tools and ethical guidelines. Constructive suggestions for improvement could include exploring methods to make VISA more adaptable to different hardware platforms and conducting additional experiments to validate its generalizability to diverse datasets.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\n- The paper is technically sound, with a well-supported methodology and comprehensive experiments. The authors provide extensive ablation studies and scaling-law experiments, demonstrating the effectiveness of VISA across different model sizes and training budgets. The results show significant improvements in computational efficiency without a drop in diffusion loss.\n\nWeaknesses: \n- Generalizability: While the experiments demonstrate impressive results on specific datasets and model sizes, it is unclear how well VISA generalizes to other types of video data or different diffusion model architectures. Additional experiments on diverse datasets would strengthen the paper's claims.\n- Computational Overhead During Training: Although VISA reduces overall computational cost, the introduction of the coarse stage adds some overhead. The authors could provide more detailed analysis on the trade-offs between the benefits of sparse attention and the additional complexity introduced during training.", "was_revised": true, "final_justification": "I keep my original score", "rebuttal": "We thank reviewer WVmh for highlighting the importance of evaluating generalizability and training overhead. Below, we address those concerns.\n\n> Generalization\n\n\nWe conducted additional experiments to evaluate ViSA’s robustness across different models, resolutions, and training recipes. We then performed blind human evaluation over 100 prompts from MovieGen-Bench.\n\n- Low-step Setup via Distillation: We distilled Wan-1.3B-T2V from 50 diffusion steps to 3 steps with both full attention and ViSA (sparsity 0.8) using the DMD distillation recipe. For ViSA, we simply swapped the attention module in the student transformer, while keeping the teacher and critic unchanged. \n\n  |               | Full Attention Win | Tie | ViSA Win |\n  |---------------|--------------------|-----|----------|\n  | Distilled (1.3B) | 23                 | 50  | 27       |\n\n- High-resolution Finetuning: We fine-tuned Wan2.1-14B-T2V at resolution 77×768×1280 with ViSA (sparsity 0.9) and compared against full attention using the same evaluation protocol:\n\n  |               | Full Attention Win | Tie | ViSA Win |\n  |---------------|--------------------|-----|----------|\n  | Hi-res (14B)   | 27                 | 47  | 26       |\n\nThese results show that ViSA achieves comparable generation quality to full attention across diverse model sizes, resolutions, and training recipes, despite computing only a fraction of the full attention cost. **We believe ViSA is the first sparse attention mechanism to demonstrate effectiveness under extremely low diffusion steps in DMD distillation**. In contrast, prior sparse attention methods were only evaluated with 50-step schedules, often keeping the first few steps as full attention.\n\n\n> Additional Complexity\n\nWe clarify that ViSA **reduces**, rather than increases, training complexity:\n\n- In Figure 4b, we compare the runtime of ViSA (fine-only) and ViSA (coarse + fine). When the sequence length exceeds 30K (typical for SoTA video DiTs), the coarse stage overhead becomes negligible compared to the computational savings achieved by VISA. In Section 2.4, we mentioned that course stage's \"runtime accounts for only 14%, even when the fine stage is 87.5% sparse.\"\n- Figure 2b shows that ViSA reduces training FLOPs by 2.53× while maintaining loss parity, consistently across scales.\n- From an engineering complexity perspective, ViSA is implemented as a standalone kernel package (like FlashAttention) and can be directly plugged into DiT codebases."}, {"review_id": "fgfqbVEI2G", "reviewer": "Reviewer_dQZn", "summary": "This paper introduces ViSA, a trainable, hardware-efficient sparse attention mechanism for video diffusion transformers. By using a two-stage coarse-to-fine block-sparse design aligned with GPU kernels, VISA achieves faster attention inference without sacrificing generation quality. Extensive experiments demonstrate the effectiveness of ViSA.", "questions": "Please see the weaknesses for the main questions. In addition, the anonymous github repo is empty.", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "### **Strengths**\n1. The trainable sparse attention is novel, which can eliminate train-test mismatch.\n2. The block-sparse design is hardware-efficient and achieves real-world speedups compatible with Flash Attention.\n3. Extensive experiments, including scaling laws and ablations, validate ViSA’s effectiveness across model sizes and training budgets.\n\n### **Weaknesses**\n1. The tile size and sparsity are predefined hyperparameters that require empirical tuning and may not generalize well to other resolutions or multi-scale training and inference.\n2. Despite good quantitative results, the lack of qualitative visualizations makes it hard to judge the real impact of the sparse attention mechanism, especially given VBench’s limited coverage.\n3. The VBench results reported seem inconsistent, with Wan’s baseline scores much worse than previously published (should be 84.92/80.10/83.96 for Quality/Semantic/Total).", "was_revised": true, "final_justification": "My concerns are mostly addressed and I raise my rating to boarderline accept.", "rebuttal": "We thank reviewer dQZn for their  thoughtful feedback. We are glad that the reviewer appreciated our core contributions, including ViSA's novel trainable sparse attention mechanism, its hardware-efficient block-sparse design, and the comprehensive experimental validation. Below, we address the raised concerns regarding generalization, qualitative evaluation, and VBench reproducibility.\n\n> Tile Size Selection\n\nThe tile size of 64 is indeed predefined. However, we believe this is one of the most important contributions of the paper. Through extensive ablation studies (Table 1(d)), we identify that a tile size of 64 achieves an great trade-off between attention granularity and hardware efficiency. **Rather than a limitation, we consider this fixed tile size a feature** that makes ViSA hardware-aligned and quality-optimized.\n\n> Sparsity Selection \n\nWe acknowledge this as a limitation and explicitly discuss it in §3.2 and our Limitations section. Rigorously predicting the optimal sparsity level during pretraining likely requires an extension of the scaling law framework to include sparsity as an additional axis—a direction that is beyond the scope of our current training budget. That said, we empirically find that **sparsity levels between 0.8 and 0.9 consistently yield strong performance across a range of settings, including pretraining/finetuning/distillation, different model sizes, and various input resolutions**. This robustness is demonstrated both in our experiments in the paper and in the additional evaluations included below.\n\n> Generalization to Other Resolution\n\n\nFollowing your advice, we evaluate ViSA on two different configurations and report blind human preference scores (avoiding reliance on VBench) using 100 randomly sampled prompts from MovieGen-Bench:\n\n1. ViSA finetuned on Wan2.1-14B-T2V, 77×768×1280 resolution, sparsity = 0.9: \n\n  |               | Full Attention Win | Tie | ViSA Win |\n  |---------------|--------------------|-----|----------|\n  | Hi-res (14B)   | 27                 | 47  | 26       |\n\n2. ViSA on Wan2.1-1.3B-T2V with 3-step DMD distillation, 61×448×832 resolution, sparsity = 0.8:  We applied the DMD distillation recipe to Wan-1.3B-T2V using both full attention and ViSA (sparsity 0.8). For the ViSA variant, we replaced the student transformer's attention module with ViSA, leaving the teacher and critic components unchanged. After distillation, our model reduce inference diffusion steps from 50 to 3.\n\n\n  |               | Full Attention Win | Tie | ViSA Win |\n  |---------------|--------------------|-----|----------|\n  | Distilled (1.3B) | 23                 | 50  | 27       |\n\nThe results suggest that ViSA generalizes well across architectures, input resolutions, and distillation settings. To the best of our knowledge, **ViSA is the first sparse attention mechanism to show strong performance under extremely low diffusion steps in DMD distillation**. In comparison, previous sparse attention methods were only tested with 50-step schedules, often retaining full attention in the early steps.\n\n\n> Visualization\n\nOur submission includes some visualization in Figure 7, Appendix F. While we are unable to include more videos/links/attachments in rebuttal due to this year’s NeurIPS policy, we directly addressed this concern through fully blind human preference studies presentend above. **Our fully-blind human evaluation confirms that ViSA-generated videos are visually stunning and indistinguishable from those generated with full attention**.  We will publicly release all code, checkpoints, and visualizations to support open evaluation beyond VBench.\n\n> VBench Scores and Reproducibility\n\n\nWe would like to clarify that the VBench results in our submission were produced strictly following the VBench's official evaluation protocol, including the prompt provided by VBench.\n\nHowever, after going through VBench’s github issue, we discovered that the original Wan team applies an prompt rewriter not included in the official VBench repo’s guidelines, which substantially alters evaluation outcomes—particularly the Semantic score. Unfortunately, the Wan team has not released their rewritten prompts or their exact evaluation pipeline(random seed,guidance scale,scheduler shift...), which makes precise replication of their results impossible. We made our best effort to approximate their setup by rewriting VBench's prompts using Wan's official rewriter and re-evaluating both Full Attention and ViSA. The updated results are as follows:\n\n\n| Model      | Quality | Semantic | Total |\n|------------|---------|----------|-------|\n| Full Attn  | 83.71   | 77.98    | 82.56 |\n| ViSA       | 83.60   | 79.47    | 82.77 |\n\nThese updated scores are now much closer to those reported for Wan. We emphasize that our evaluation has always been consistent and fair: ViSA and full attention were always evaluated by us under identical setups. While we cannot exactly match Wan’s numbers due to their unreleased setup and inherent randomness in generation, we will clearly document this caveat and provide full evaluation details in the final version. We also hope our additional results on human evaluation can complement the VBench scores."}]

Author rebuttal:

Meta-review: This paper introduces a trainable sparse attention mechanism to accelerate video diffusion models. It adopts a two-stage approach to first identify critical tokens and then compute attention efficiently. The reviewers generally agreed that the proposed method is novel and hardware-efficient, and the evaluations are extensive. They raised some questions regarding generalizability, computational overhead, and lack of qualitative visualizations, which were addressed by the authors in the rebuttal and discussions. The final ratings are 1 "accept" and 2 "borderline accept". It is recommended to accept this paper for its novelty and efficiency.
Decision: Accept (poster)