Sparser Block-Sparse Attention via Token Permutation

TL;DR

This work tackles the $O(N^2)$ cost of self-attention for long-context LLMs by introducing Permuted Block-Sparse Attention (PBS-Attn), a plug-and-play method that rearranges tokens to produce a more favorable block-sparse structure while preserving causality through segmented permutation. It leverages attention permutation invariances to implement a query-aware key permutation that clusters salient key tokens within segments, enabling sparser block interactions without sacrificing accuracy. The authors provide a formal treatment of permutation properties, a practical segmentation strategy, and a custom permuted-FlashAttention kernel, along with extensive experiments on LongBench and LongBenchv2 showing near-full-attention performance with up to $2.75×$ end-to-end speedups for long-context prefilling. This approach offers a scalable and hardware-friendly path to more efficient long-context LLMs with broad practical impact for real-world long-sequence tasks.

Abstract

Scaling the context length of large language models (LLMs) offers significant benefits but is computationally expensive. This expense stems primarily from the self-attention mechanism, whose $O(N^2)$ complexity with respect to sequence length presents a major bottleneck for both memory and latency. Fortunately, the attention matrix is often sparse, particularly for long sequences, suggesting an opportunity for optimization. Block-sparse attention has emerged as a promising solution that partitions sequences into blocks and skips computation for a subset of these blocks. However, the effectiveness of this method is highly dependent on the underlying attention patterns, which can lead to sub-optimal block-level sparsity. For instance, important key tokens for queries within a single block may be scattered across numerous other blocks, leading to computational redundancy. In this work, we propose Permuted Block-Sparse Attention (\textbf{PBS-Attn}), a plug-and-play method that leverages the permutation properties of attention to increase block-level sparsity and enhance the computational efficiency of LLM prefilling. We conduct comprehensive experiments on challenging real-world long-context datasets, demonstrating that PBS-Attn consistently outperforms existing block-sparse attention methods in model accuracy and closely matches the full attention baseline. Powered by our custom permuted-FlashAttention kernels, PBS-Attn achieves an end-to-end speedup of up to $2.75\times$ in long-context prefilling, confirming its practical viability. Code available at https://github.com/xinghaow99/pbs-attn

Figures (9)

• Figure 1: Illustration of causal attention without (Left) and with (Right) segmented permutation with $B=1, S=4$. Segmented permutation enhances block-level sparsity via intra-segment permutation while preserving inter-segment causality. By restricting computation of blocks within on-diagonal segments (green blocks), we can safely skip inter-segment blocks (yellow blocks) for block-sparse attention.
• Figure 2: Comparison of attention maps for Llama-3.1-8B (layer 0, head 30) on an 8K LongBench example, showing the pattern without (\ref{['fig:llama_no_perm']}) and with (\ref{['fig:llama_perm']}) segmented permutation. The red overlay indicates blocks selected for block-sparse attention, and the attention coverage is calculated as the total attention scores covered by the selected blocks. More visualizations are provided in Appendix \ref{['sec:permutation-visualization']}.
• Figure 3: Speedup of various methods relative to FlashAttention, measured by time to first token (TTFT) on LongBenchv2 across various sequence lengths. To accommodate longer sequences under memory constraints, we employ tensor parallelism with tp_size of 2 and 8 for the 256K and 512K contexts, respectively.
• Figure 4: Block-level density on various context lengths with and without permutation. A relative sparsity improvement $\Delta$ is calculated.
• Figure 5: LongBench score vs. average block-level density at a context length of 32K.

Theorems & Definitions (9)

• Lemma 3.1: Key-Value Pair Permutation Invariance
• Lemma 3.2: Query Permutation Equivariance
• Theorem 3.3: Attention Permutation Invariance under Inverse Transformation
• Lemma A.1: Key-Value Pair Permutation Invariance
• proof
• Lemma A.2: Query Permutation Equivariance
• proof
• Theorem A.3: Attention Permutation Invariance under Inverse Transformation
• proof