Parameterized stripe attention speeds up video generation models
Parameterized Stripe Attention for Efficient Video Generation
Computer Vision and Pattern RecognitionArtificial Intelligence
Summary
Video generation models that create new videos often take a long time because they have to look at many parts of a video at once, which is very slow. The authors found that the way these models pay attention follows a repeating stripe pattern in both space and time. They designed a new way to represent this pattern more efficiently, speeding up video generation while keeping the quality close to the original. Their method works well on two popular video generation datasets, making video creation faster without much loss in quality.
What this means in practice
- •For video software engineers: Speed up AI-based video generation tools by integrating parameterized stripe attention for faster inference with minimal quality loss.
- •For graphics hardware developers: Design efficient GPU kernels that implement unified sparse attention patterns to reduce computational cost in video models.
Authors
Xingyu Jia, Baole Ai, Ang Wang, Kang Zhao, Yong Li
Abstract
Diffusion Transformers (DiTs) enable high-quality video generation but suffer from substantial inference latency, primarily attributable to the computationally expensive full spatio-temporal attention. While sparse attention methods offer potential solutions, existing approaches face an inherent flexibility--efficiency dilemma: predefined masks lack the flexibility to capture diverse attention patterns, while runtime-determined masks introduce overheads and sacrifice hardware efficiency. We identify the lack of a unified structural characterization of DiT attention as a key limitation of existing methods, and establish that video DiT attention exhibits \textbf{periodic diagonal stripe structures} along both temporal and spatial dimensions. To formally encode these structured patterns within a single efficient kernel, we present {\bf PSA}, a parameterized stripe attention that formalizes the observed stripe regularity, unifying diverse attention patterns for efficient mask generation. This unified representation enables a single hardware-efficient CUDA kernel to process all sparse patterns, achieving FlashAttention-3-level Model FLOPs Utilization. To determine optimal sparsity configurations, we propose a training-free offline search algorithm that automatically maximizes sparsity under a specified error tolerance for each attention head. Experiments on HunyuanVideo and Wan~2.1 demonstrate that PSA achieves 1.57$\times$ and 1.37$\times$ end-to-end speedups over FlashAttention-3 baselines, with acceptable visual quality degradation.