Token Sparse Attention: 基于交错token选择的高效长上下文推理

注意力机制的二次复杂度仍然是大语言模型长上下文推理的核心瓶颈。现有加速方法要么通过结构化模式稀疏化注意力图，要么在特定层永久移除token，这可能导致保留无关token或依赖不可逆的早期决策，而无法适应token重要性的逐层变化。本文提出了一种轻量且动态的token级稀疏化机制——Token Sparse Attention，在注意力计算过程中将每个head的Q、K、V压缩到一个缩减的token集合中，并在输出后将其解压回原始序列，从而允许后续层重新考虑token信息。此外，该方法揭示了token选择与稀疏注意力交汇的新设计点。该方法完全兼容密集注意力实现（如Flash Attention），并能无缝集成现有稀疏注意力内核。实验结果表明，Token Sparse Attention在保持精度损失低于1%的情况下，实现了最高达3.23倍的注意力加速效果，证明了动态且交错的token级稀疏化是可扩展长上下文推理的有效补充策略。