arXiv:2510.21270cs.CLcs.AI2025-10被引 2

通过词元置换提升块稀疏注意力效率,加速长文本处理。

Sparser Block-Sparse Attention via Token Permutation

  • 用词元置换优化注意力块的稀疏结构,提升计算效率。
  • 在长序列预填充中实现最高2.75倍速度提升,精度接近全注意力。
  • 无需修改模型架构,可直接替换现有注意力模块,适合长文本应用。

扩大大语言模型的上下文长度虽有益处,但计算成本高昂,主要源于自注意力机制随序列长度呈$O(N^2)$增长,成为内存与延迟瓶颈。尽管注意力矩阵通常稀疏,尤其是长序列时,但现有块稀疏注意力方法的效果依赖于注意力模式,常因关键键值词元分散在多个块中导致计算冗余。本文提出可插拔的置换块稀疏注意力(PBS-Attn),利用注意力的置换特性增强块级稀疏性,提升模型预填充阶段的计算效率。在真实世界长上下文数据集上的实验表明,PBS-Attn在模型精度上持续优于现有块稀疏注意力方法,并接近全注意力基线。借助定制的置换式FlashAttention内核,其端到端预填充速度提升最高达2.75倍,验证了实际可行性。代码已公开于https://github.com/xinghaow99/pbs-attn。

原文摘要 · Abstract (English)

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

注意力机制长文本处理高效推理

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。