Sparser Block-Sparse Attention via Token Permutation
Authors: name: Xinghao Wang; name: Pengyu Wang; name: Dong Zhang; name: Chenkun Tan; name: Shaojun Zhou; name: Zhaoxiang Liu; name: Shiguo Lian; name: Fangxu Liu; name: Kai Song; name: Xipeng Qiu
Published: 2025-10-24 09:11:50

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 (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.75times in long-context
prefilling, confirming its practical viability. Code available at
https://github.com/xinghaow99/pbs-attn