通过预填充构建查询中心索引表,显著加速长文本大模型推理。
CSAttention: Centroid-Scoring Attention for Accelerating LLM Inference

- 预填充阶段构建固定大小的查询索引表,解码时用查表替代全扫描。
- 在32K-128K上下文、95%稀疏度下,推理速度比最优基线快4.6倍。
- 无需训练,适用于需重复使用上下文的智能体与问答场景。
长上下文大模型越来越多依赖可复用的预填充提示用于智能体和领域问答,导致注意力机制与键值缓存成为解码阶段的主要瓶颈。尽管稀疏注意力能降低计算与传输开销,但在高稀疏度下常因查询与键的分布偏移而损失精度。本文提出无训练的中心评分注意力(CSAttention),专为可复用上下文的高吞吐服务优化。其采用存储换计算策略,将计算集中在一次性离线预填充阶段,实现跨多轮查询的开销分摊,同时极大压缩每步解码延迟。具体而言,CSAttention在预填充阶段构建以查询为中心的查找表,其大小在解码阶段保持不变,使在线解码可替换全上下文扫描为高效表查找与GPU友好的分数累加。大量实验表明,CSAttention在精度上接近全注意力表现。在高稀疏度(95%)和长上下文(32K–128K)设置下,其在模型准确率与推理速度上均持续优于现有最先进稀疏注意力方法,在128K上下文长度下,较最准确基线最高实现4.6倍推理加速。
原文摘要 · Abstract (English)
Long-context LLMs increasingly rely on extended, reusable prefill prompts for agents and domain Q&A, pushing attention and KV-cache to become the dominant decode-time bottlenecks. While sparse attention reduces computation and transfer costs, it often struggles to maintain accuracy at high sparsity levels due to the inherent distribution shift between Queries and Keys. We propose Centroid-Scoring Attention (CSAttention), a training-free sparse attention method optimized for high-throughput serving of reusable contexts. CSAttention adopts a storage-for-computation strategy tailored to the offline-prefill/online-decode setting: it front-loads computation into a one-time offline prefill phase that can be amortized across multiple queries, while aggressively optimizing per-step decoding latency. Specifically, CSAttention constructs query-centric lookup tables during offline prefill, whose size remains fixed during decoding, and enables online decoding to replace full-context scans with efficient table lookups and GPU-friendly score accumulation. Extensive experiments demonstrate that CSAttention achieves near-identical accuracy to full attention. Under high sparsity (95%) and long-context settings (32K-128K), CSAttention consistently outperforms state-of-the-art sparse attention methods in both model accuracy and inference speed, achieving up to 4.6x inference speedup over the most accurate baseline at a context length of 128K.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。