PSA: Pyramid Sparse Attention for Efficient Video Understanding and Generation

TL;DR

PSA introduces a hierarchical, multi-level sparse attention mechanism that replaces binary keep/drop masks with a pyramid of pooled KV blocks, enabling wider attention coverage without increasing FLOPs. By combining adaptive level assignment based on block importance and cosine-similarity constraints, PSA preserves richer context for long video sequences while remaining hardware-friendly through a decoupled kernel. The approach demonstrates strong efficiency-quality trade-offs across video understanding and generation tasks, outperforming traditional block-sparse methods in training-free and distillation scenarios. Its plug-and-play design and hardware-optimized kernel make it a practical solution for scalable long-context video modeling.

Abstract

Attention mechanisms are the core of foundation models, but their quadratic complexity remains a critical bottleneck for scaling. This challenge has driven the development of efficient attention mechanisms, with sparsity emerging as the dominant paradigm. Current methods typically retain or discard entire key-value blocks with binary masks, resulting in substantial information loss under high sparsity. To mitigate this gap, we present Pyramid Sparse Attention (PSA), a versatile module applicable to both video understanding and generation tasks. Instead of binary masking, PSA introduces multi-level pooled KV representations, enabling finer mask granularity. Specifically, each query block dynamically allocates lower pooling levels to critical KV blocks and higher levels to less important ones, creating an informative interpolation between full retention and complete pruning. This design, analogous to fixed-point quantization and classical feature pyramid networks in computer vision, effectively mitigates information loss while preserving computational efficiency under a low compute budget. It works with a native, hardware-friendly kernel that leverages decoupled block-tile design to ensure efficient execution. Across video understanding and generation benchmarks, PSA preserves contextual information and visual fidelity, consistently outperforming or achieving comparable performance over existing sparse attention baselines with superior efficiency-quality trade-offs. Our code and model weights are publicly available at: http://ziplab.co/PSA

Figures (10)

• Figure 1: Comparison of attention mechanisms under identical compute budget. All methods use the same input $Q$, $K$, and $V$ tensors extracted from Wan2.1–1.3B wan2025wanopenadvancedlargescale denoising process. Computation Pattern (top-left two panels): Normalized block-wise FLOPs distribution. The two panels plot query blocks on the horizontal axis and key blocks on the vertical axis. Despite identical FLOPs ($20\%$ full), the proposed Pyramid Sparse Attention (PSA) allows each query block to attend to a much larger portion of KV blocks (70% active regions), whereas Block Sparse Attention (BSA) dao2022flashattentionfastmemoryefficientexactzhang2025spargeattentionaccuratetrainingfreesparsexu2025xattentionblocksparseattention restricts each query to only a narrow subset of KV blocks (20% active regions), concentrating FLOPs in limited areas. Attention Output (bottom row): Resulting attention visualizations. PSA closely matches the Full Attention baseline with minimal relative error ($<3\%$), while BSA shows noticeable distortions due to aggressive pruning.
• Figure 2: Adjacent Key Token Cosine Similarity. High cosine similarity between key tokens (Qwen2.5-VL, Wan2.1-1.3B) motivates hierarchical pooling: nearby visual tokens are highly similar.
• Figure 3: Overview of the Pyramid Sparse Attention (PSA) framework. PSA adaptively allocates attention computation across hierarchical KV representations (green; lighter shades denote coarser levels). The multi-level mask (blue) determines which KV level each query block attends to. As illustrated, the current attention block assigned to level 4 uses the coarsest KV representation $K_j^4$ and $V_j^4$.
• Figure 4: Qualitative comparison on Wan2.1-1.3B (Text-to-Video, 720p).
• Figure 5: A plane takes off over a city skyline; cut to a graffiti-covered steam train pulling into a station with billowing smoke. Cinematic colors and motion.