CoSA：通过代理内核共同设计的稀疏注意力加速长上下文推理

自注意力的二次成本使得长上下文推理成本过高，而基于代理的块稀疏注意力已成为一种实用的补救措施。现有方法通常依赖代理来预测二进制稀疏掩码，并依赖内核来使用该掩码并执行稀疏注意力计算。这种方法在适度的预算下是有效的。然而，随着预算的收紧，估计代理不可避免地会丢弃一些显着块，而内核只能机械地应用稀疏掩模，导致模型精度明显下降。我们提出 CoSA，这是代理内核协同设计下的两阶段 无需训练 稀疏注意力，它将内核感知代理 (KAP) 与有序跳过内核 (OSK) 结合起来。在第一阶段，KAP 在适度预算下选择块，并生成一个有序掩码，该掩码规定了内核内循环中访问 KV 页的顺序。在第二阶段，OSK 应用此掩码，并在给定在线 softmax 统计数据的紧缩预算下跳过更多块。在主流 LLM 主干网和长上下文基准测试中，CoSA 以更低的预算获得了更高的准确性。令人印象深刻的是，CoSA 在 128K 上下文长度下实现了 4.93$\times$ 注意力加速，并将端到端首次词元时间减少了 2.53$\times$，性能下降可以忽略不计。代码可在 https://github.com/Tencent/AngelSlim 获取。