PISA: Piecewise Sparse Attention Is Wiser for Efficient Diffusion Transformers

TL;DR

PISA tackles the quadratic bottleneck in diffusion transformers by introducing a training-free piecewise sparse attention that exacts computation for critical blocks and approximately processes the rest via block-wise Taylor expansions within the softmax. The core innovations include a global first-order correction with a covariance-aware block selection strategy and a hardware-friendly fused kernel that preserves normalization and pre-trained weights. The approach yields strong speedups (e.g., up to $2.57\times$ on Hunyuan-Video and Wan2.1-14B) while maintaining or surpassing state-of-the-art generation quality across video and image tasks, demonstrating a practical path to efficient, high-fidelity diffusion transformers. This work expands the sparse attention paradigm by unifying exact and approximate computations under the softmax framework, achieving near-lossless approximation with meaningful real-world benefits for scalable visual generation.

Abstract

Diffusion Transformers are fundamental for video and image generation, but their efficiency is bottlenecked by the quadratic complexity of attention. While block sparse attention accelerates computation by attending only critical key-value blocks, it suffers from degradation at high sparsity by discarding context. In this work, we discover that attention scores of non-critical blocks exhibit distributional stability, allowing them to be approximated accurately and efficiently rather than discarded, which is essentially important for sparse attention design. Motivated by this key insight, we propose PISA, a training-free Piecewise Sparse Attention that covers the full attention span with sub-quadratic complexity. Unlike the conventional keep-or-drop paradigm that directly drop the non-critical block information, PISA introduces a novel exact-or-approximate strategy: it maintains exact computation for critical blocks while efficiently approximating the remainder through block-wise Taylor expansion. This design allows PISA to serve as a faithful proxy to full attention, effectively bridging the gap between speed and quality. Experimental results demonstrate that PISA achieves 1.91 times and 2.57 times speedups on Wan2.1-14B and Hunyuan-Video, respectively, while consistently maintaining the highest quality among sparse attention methods. Notably, even for image generation on FLUX, PISA achieves a 1.2 times acceleration without compromising visual quality. Code is available at: https://github.com/xie-lab-ml/piecewise-sparse-attention.

Figures (11)

• Figure 1: PISA accelerates diverse generation tasks. Top (a, b): Wan2.1-14B video generation. PISA achieves 2.14$\times$ speedup over Dense Attention with no appreciable quality loss. Bottom (c): FLUX.1-dev text-to-image generation. PISA at higher sparsity ratio r=85% preserves better quality and structure than SpargeAttn zhang2025spargeattn.
• Figure 2: Visualization of attention patterns on Wan2.1-1.3B. PISA achieves 100% effective block coverage similar to full attention. This near-lossless approximation with only negligible computational overhead relative to standard sparse attention.
• Figure 3: Visualization of pre-softmax attention scores ($\boldsymbol{Q}\boldsymbol{K}^\top$) in Wan2.1-1.3B. The block-wise scores exhibit a symmetric bell-shaped distribution. uncritical blocks (Left) cluster in negative regions where the 1st-order Taylor expansion is highly accurate, whereas important blocks (Right) diverge. This property remains robust under "Safe-Exp" shift for numerical stability (Bottom).
• Figure 4: The algorithm pipeline of PISA. Prepare Phase: We pre-compute block-wise mean of the queries and keys ($\overline{Q},\overline{K}$), block-wise sum of value ($\hat{V}$) and the global $H$ in a single pass. A block-wise Top-K selection identifies critical blocks. Fused Attention Kernel: The kernel dynamically switches execution paths: selected blocks (e.g., indices 2, 3) undergo exact computation (Phase 1), while unselected blocks (e.g., indices 1, 4) are approximated using block-wise zeroth-order expansion (Phase 2). In Phase 3, the $H$ is applied to inject global first-order approximation. This design allows loading the global correction term once, avoiding memory-bound streaming.
• Figure 5: Kernel efficiency profile. (a) Latency comparison across sequence lengths at 12.5% density (87.5% sparsity) under two mainstream configurations with notation B-H-D (batch_size, num_heads, head_dim). (b) Relative speedup against FlashAttn-2/3 across varying densities and sequence lengths under the B2-H16-D128 configuration. The dashed line indicates the baseline performance.

Theorems & Definitions (5)

• Theorem 3.1: Error Analysis of Global First-Order Approximation
• Lemma 3.1: Residual operator norm bound
• proof
• Theorem 3.2: Error Analysis of Global First-Order Approximation
• proof