Faster Video Diffusion with Trainable Sparse Attention
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.53times with no drop in
diffusion loss. Retrofitting the open-source Wan-2.1 model speeds up attention
time by 6times and lowers end-to-end generation time from 31s to 18s with
comparable 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.
682c27e2fffb36958f8cd84e
Peiyuan Zhang
admin_assigned
2025-05-20T08:36:24.007Z
63565cc56d7fcf1bedb7d347
https://cdn-avatars.huggingface.co/v1/production/uploads/63565cc56d7fcf1bedb7d347/XGcHP4VkO_oieA1gZ4IAX.jpeg
Zhang Peiyuan
user
PY007
682c27e2fffb36958f8cd84f
Haofeng Huang
admin_assigned
2025-05-20T08:36:29.763Z
67ea1f6693f71dd8167a2d22
https://cdn-avatars.huggingface.co/v1/production/uploads/no-auth/H_upra_XVG1AoBKUe9ArV.png
haofeng huang
user
haofeng666
682c27e2fffb36958f8cd850
Yongqi Chen
admin_assigned
2025-05-20T08:36:36.258Z
65416817271d3bc4d70f6745
https://cdn-avatars.huggingface.co/v1/production/uploads/65416817271d3bc4d70f6745/1YkW0MpuufejvxqksVMIx.jpeg
Yongqi Chen
user
BrianChen1129
682c27e2fffb36958f8cd851
Will Lin
682c27e2fffb36958f8cd852
Zhengzhong Liu
admin_assigned
2025-05-20T08:36:46.746Z
62fbdc67c776fd8821ae3f2d
https://cdn-avatars.huggingface.co/v1/production/uploads/62fbdc67c776fd8821ae3f2d/cI7iAZOL40RUYluo5ZVTU.png
Zhengzhong Liu
user
hunterhector
682c27e2fffb36958f8cd853
Ion Stoica
682c27e2fffb36958f8cd854
Eric P. Xing
admin_assigned
2025-05-20T08:37:02.570Z
64ff67722ad36636be6c4542
https://cdn-avatars.huggingface.co/v1/production/uploads/noauth/sLIrNelAWPVOy4e3oo5LB.jpeg
Eric Xing
user
EricX003
682c27e2fffb36958f8cd855
Hao Zhang
Faster Video Diffusion with Trainable Sparse Attention
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.53times with no drop in
diffusion loss. Retrofitting the open-source Wan-2.1 model speeds up attention
time by 6times and lowers end-to-end generation time from 31s to 18s with
comparable 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.
Trainable sparse attention (VSA) reduces computational cost in video diffusion transformers with minimal impact on performance, enabling efficient scaling of the models.
video diffusion transformers DiTs quadratic 3D attention VSA token-level attention block computing FlashAttention3 MFU FLOPS diffusion loss Wan-2.1 model