Cache friendly scheduling speeds up attention in long video generation
WaveAlign: Cache-Aware Query-Row Scheduling for Sparse Attention in Long-Video Generation
Distributed, Parallel, and Cluster Computing
Summary
Long videos made by AI models have very long sequences to process, which slows down their attention step. The authors found that reordering the processing order of certain pieces can make the computer’s cache work more efficiently, reducing slow memory traffic. They developed WaveAlign, a method that smartly groups and sorts these pieces to reuse data better in memory. This approach speeds up video generation without losing quality and fits into existing AI systems without major changes.
What this means in practice
- •For ai model developers: Improve speed and efficiency of attention computation in long video generation models by using cache-aware query scheduling.
- •For gpu system engineers: Optimize GPU memory usage and traffic during sparse attention workloads to boost throughput in video generation tasks.
Authors
Zijian Dai, Sen Han, Youhui Bai, Shannon Wang, Kan Wu, Jingkai Huang, Yuhang Wang, Jing Li, Cheng Li
Abstract
Long-video generation with diffusion transformers (DiTs) produces extremely long token sequences, making attention a dominant inference bottleneck. Dynamic sparse attention reduces computation, but its realized speedup remains limited because irregular query-row execution degrades L2 cache locality and increases HBM traffic. We present WaveAlign, a lightweight, cache-aware query-row reordering framework for dynamic sparse attention. WaveAlign formulates row ordering as an optimization problem and approximates it with two stages. The first stage derives a low-rank SVD representation of sparse-mask rows and groups query rows with similar K/V access patterns, increasing K/V overlap among concurrently scheduled rows. The second stage exploits streaming GPU scheduling by sorting rows within each wave in descending order of their K/V-block counts, so that short rows from the current wave are followed by long rows from the next. This aligns K/V accesses across wave boundaries and enables shared blocks to be reused before eviction. An adaptive skip module avoids unprofitable reordering. By only permuting query and mask rows, WaveAlign preserves sparse-attention semantics and requires no changes to existing methods or backend kernels. Across two GPU architectures, two video DiTs, and four sparse-attention methods, WaveAlign raises the L2 cache hit ratio from 28.48%--36.35% to 79.38%--89.06%, reduces HBM read traffic by up to 92.11%, and achieves up to 1.25x kernel and 1.17x end-to-end generation speedup without quality loss.