Sparser Block-Sparse Attention via Token Permutation

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 (**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 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 $\mathbf{2.75\times}$ in long-context prefilling, confirming its practical viability. Code available at \url{https://github.com/xinghaow99/pbs-attn}.

Long-context language models can read long documents, code, and conversations, but this is slow because attention becomes expensive as the input grows. Block-sparse attention speeds this up by skipping unnecessary blocks, but important tokens can scatter across many blocks, limiting the savings. We propose Permuted Block-Sparse Attention, which rearranges tokens within causal segments so important tokens are grouped into denser blocks. This makes block-sparse attention more efficient while preserving correctness and requiring no model retraining. On long-context benchmarks, our method keeps accuracy close to full attention and achieves up to 2.75× faster prefilling with custom GPU kernels, making long-context LLMs more practical and efficient.