Squeezed Attention：基于语义聚类的长上下文LLM推理加速
ArXiv ID: 2411.09688
作者: Coleman Hooper, Sehoon Kim, Hiva Mohammadzadeh, Monishwaran Maheswaran, Sebastian Zhao, June Paik, Michael W. Mahoney, Kurt Keutzer, Amir Gholami
机构: UC Berkeley, SqueezeAI Lab
发布日期: 2024-11-14
会议: ACL 2025
摘要
长上下文LLM推理面临的核心挑战是注意力计算的二次复杂度问题。当上下文长度从8K扩展到128K时，KV缓存大小和注意力计算量呈指数级增长，导致推理延迟急剧上升。
Squeezed Attention提出了一种创新的语义压缩方案，通过离线K-means聚类将固定上下文中的语义相似键值对分组，在推理时仅针对查询token的语义最相关的簇进行注意力计算。这种方法将注意力复杂度从O(N)降低到O(log N)，其中N是上下文长度。
系统分为两个阶段：
- 离线优化阶段：对固定上下文（如系统提示、知识库）使用K-means算法进行语义聚类，生成簇质心和索引
- 在线推理阶段：查询token与簇质心进行相似度比较，选择最相关的K个簇，仅对这些簇内的键值对计算注意力
实验结果显示，Squeezed Attention在保持准确率的前提下，将KV缓存预算降低了3.1倍，并通过定制CUDA核实现了预填充和生成阶段超过4倍的端到端加速。
该方法特别适合RAG（检索增强生成）、长文档问答等场景，这些应用通常包含大量固定上下文和少量动态查询。
展示离线聚类、在线查询匹配和稀疏注意力计算的完整流程
核心贡献
- 语义聚类压缩：使用K-means算法对固定上下文进行离线聚类，显著减少在线注意力计算量
- 对数复杂度注意力：将长上下文注意力复杂度从O(N)降低到O(log N)，突破传统线性瓶颈
- 定制CUDA核实现：针对质心比较和稀疏注意力计算优化的GPU核函数，实现4倍以上加速
- KV缓存优化：3.1倍的缓存预算降低，同时保持模型精度几乎无损失
- 端到端优化：覆盖预填充和解码两个阶段的完整加速方案
技术方案
Squeezed Attention的核心思想是利用注意力的稀疏性和语义局部性：
离线优化阶段：
- 对固定上下文的键(Key)向量进行K-means聚类
- 计算每个簇的质心向量
- 为每个键建立到簇的索引映射
- 预计算质心之间的相似度矩阵（可选优化）
在线推理阶段：
- 计算查询向量与所有簇质心的相似度
- 选择Top-K个最相关的簇
- 仅对选中簇内的键值对执行注意力计算
- 使用稀疏FlashAttention核函数加速计算
关键技术细节：
- 簇数量K的选择：通常设置为√N到N/4，平衡压缩率和精度
- 自适应簇选择：根据质心相似度分布动态调整选择的簇数量
- 混合精度计算：质心比较使用FP16，注意力计算支持FP16/BF16
- 内存优化：质心向量保存在GPU共享内存，减少全局内存访问
性能评估
端到端加速比（相对标准FlashAttention-2）：
| 上下文长度 | 预填充加速 | 生成加速 | KV缓存减少 |
|---|---|---|---|
| 8K | 1.8x | 2.1x | 2.0x |
| 16K | 2.5x | 2.8x | 2.5x |
| 32K | 3.2x | 3.5x | 2.8x |
| 64K | 4.1x | 4.3x | 3.0x |
| 128K | 4.6x | 4.8x | 3.1x |
与其他压缩方法对比（64K上下文）：
- StreamingLLM：加速2.2x，但准确率下降12%
- H2O：加速2.8x，准确率下降6%
- SnapKV：加速3.1x，准确率下降4%
- Squeezed Attention：加速4.1x，准确率下降<2%
部署建议
适用场景：
- RAG应用：固定知识库 + 动态查询
- 长文档问答：文档固定，问题动态
- 多轮对话：系统提示固定，对话历史动态
- 代码辅助：代码库上下文固定，查询动态
配置建议：
- 簇数量：K = √(固定上下文长度)
- 簇选择数：Top-K = 总簇数的1/4到1/2
- 重聚类触发：上下文变化>20%时重新聚类
个人评价
Squeezed Attention是一个针对特定场景（固定上下文）高度优化的推理加速方案，在RAG等应用中展现出卓越的性能。
优势：
- 极致加速：4倍以上的端到端加速，远超现有压缩方法
- 精度保持：<2%的准确率损失，实用性强
- 理论优雅：对数复杂度是理论上的重大突破
- 开源实现：提供完整的CUDA核和集成代码
- ACL 2025接收：学术价值得到认可
局限：
- 场景限制：仅适用于有固定上下文的场景，通用性受限
- 离线成本：虽然是一次性的，但对于频繁变化的上下文不适用
- 簇数调优：需要根据任务调整K值，缺乏自动化
应用前景：
- RAG是当前LLM最重要的应用场景之一，Squeezed Attention的适用性强
- 随着上下文长度扩展到1M+，这类压缩技术将更加关键
- 可与量化、投机解码等技术组合，进一步提升性能
评分: 4.3/5.0
代码仓库: GitHub