2025_NIPS_Spotlight Attention: Towards Efficient LLM Generation via Non-linear Hashing-based KV Cach

文章总结与翻译

一、主要内容

该研究聚焦于大型语言模型（LLMs）推理阶段的KV缓存优化问题，针对解码阶段因KV缓存频繁内存交互导致的效率瓶颈，提出了一种名为Spotlight Attention的高效推理方法。

1. 核心背景：LLMs推理分为预填充和解码两个阶段，解码阶段因KV缓存的内存交换成为主要瓶颈（如LLaMA2-7B在A100 GPU上解码阶段GPU利用率不足10%）。现有KV缓存优化方法存在缺陷：静态剪枝不适应长序列任务，动态永久淘汰可能过早移除关键令牌，而基于线性哈希的令牌级检索（如MagicPIG）需冗长哈希码，存储和计算成本高。

2. 方法设计：

   • 采用非线性MLP哈希函数替代传统线性哈希，解决LLMs中查询和键在嵌入空间呈正交锥分布导致的线性哈希编码效率低问题，哈希码长度缩短至少5倍。
   • 提出基于Bradley-Terry排序损失的轻量训练框架，仅需少量校准数据，在16GB内存GPU上8小时即可完成训练，且不冻结LLM主干网络。
   • 实现专用CUDA内核，通过位运算优化哈希检索，单A100 GPU上512K令牌的哈希检索耗时低于100