中文

CSAttention:基于质心评分的注意力机制加速大语言模型推理

机器学习 2026-04-13 v1 人工智能

摘要

长上下文大语言模型(LLM)越来越依赖于扩展且可重复使用的预填充提示,用于代理和领域问答,将注意力和 KV 缓存推至解码时瓶颈的主导地位。虽然稀疏注意力可减少计算和传输成本,但常常在稀疏度较高时难以维持准确性,因为查询(Queries)与键(Keys)之间固有的分布迁移。我们提出了质心评分注意力(CSAttention),这是一种针对可重复上下文高吞吐服务优化的无训练稀疏注意力方法。CSAttention 采用以存储为计算的策略,针对离线预填充/在线解码的场景进行优化:它将计算前置到一次离线预填充阶段,可跨多个查询 amortized(消耗),同时积极优化每一步解码的延迟。具体而言,CSAttention 在离线预填充期间构建查询中心的查找表,其大小在解码期间保持固定,使在线解码能够用高效的表查找和 GPU 友好得分累积来取代完整的上下文扫描。大量实验表明,CSAttention 在准确性上几乎与完整注意力持平。在高稀疏度(95%)和长上下文设置(32K-128K)下,CSAttention 在模型准确性和推理速度方面 consistently 优于当前最先进的稀疏注意力方法,在上下文长度为 128K 时,相对于最准确的基线实现最高可达 4.6 倍的推理加速。

关键词

引用

@article{arxiv.2604.08584,
  title  = {CSAttention: Centroid-Scoring Attention for Accelerating LLM Inference},
  author = {Chuxu Song and Zhencan Peng and Jiuqi Wei and Chuanhui Yang},
  journal= {arXiv preprint arXiv:2604.08584},
  year   = {2026}
}