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

翻译：扩展大型语言模型（LLM）的上下文长度能带来显著优势，但计算成本高昂。这种成本主要源于自注意力机制，其相对于序列长度的 $O(N^2)$ 复杂度构成了内存和延迟的主要瓶颈。幸运的是，注意力矩阵通常是稀疏的，尤其是在长序列中，这暗示了优化的可能性。块稀疏注意力作为一种有前景的解决方案应运而生，它将序列划分为块，并跳过其中一部分块的计算。然而，该方法的有效性高度依赖于底层的注意力模式，这可能导致次优的块级稀疏性。例如，单个块内查询所需的重要键令牌可能分散在许多其他块中，从而导致计算冗余。在本工作中，我们提出了置换块稀疏注意力（\textbf{PBS-Attn}），这是一种即插即用的方法，它利用注意力的置换特性来增加块级稀疏度，从而提升LLM预填充的计算效率。我们在具有挑战性的真实世界长上下文数据集上进行了全面的实验，结果表明PBS-Attn在模型准确性上持续优于现有的块稀疏注意力方法，并与全注意力基线结果高度接近。借助我们定制的置换FlashAttention内核，PBS-Attn在长上下文预填充中实现了端到端高达 $2.75\times$ 的加速，证实了其实用可行性。代码发布于 https://github.com/xinghaow99/pbs-attn。