Xinghao Wang Pengyu Wang Dong Zhang Chenkun Tan Shaojun Zhou Zhaoxiang Liu Shiguo Lian Fangxu Liu Kai Song Xipeng Qiu

摘要

扩展大型语言模型（LLMs）的上下文长度虽能带来显著优势，但其计算成本高昂。这一高昂开销主要源于自注意力机制，其复杂度随序列长度呈 O(N²) 增长，成为内存占用和延迟方面的主要瓶颈。值得庆幸的是，注意力矩阵在实际中通常具有稀疏性，尤其是在处理长序列时，这为优化提供了契机。块稀疏注意力（Block-sparse attention）应运而生，成为一种有前景的解决方案：它将序列划分为多个块，并跳过其中部分块的计算。然而，该方法的效果高度依赖于底层注意力模式，可能导致块级稀疏性不足。例如，某一区块内查询所需的关键令牌可能分散在大量其他区块中，从而引发计算冗余。针对这一问题，本文提出了一种即插即用的新型方法——置换块稀疏注意力（Permuted Block-Sparse Attention, PBS-Attn），该方法利用注意力机制的置换特性，提升块级稀疏性，显著增强 LLM 预填充（prefilling）阶段的计算效率。我们在多个具有挑战性的真实世界长上下文数据集上进行了全面实验，结果表明，PBS-Attn 在模型准确率方面持续优于现有块稀疏注意力方法，且性能接近全注意力（full attention）基线。借助我们自研的置换型 FlashAttention 内核，PBS-Attn 在长上下文预填充任务中实现了高达 2.75 倍的端到端加速，充分验证了其实际可行性。代码已开源，地址为：https://github.com/xinghaow99/pbs-attn