在大规模语言模型的长上下文推理场景中,Key-Value 缓存的内存消耗已经成为制约系统吞吐量的核心瓶颈。一个支持百万级 token 上下文的中等规模模型,其 KV 缓存所需的 GPU 显存可能超过数百 GB,远超单卡容量。传统的缓存淘汰策略,如保留最近 token 或基于位置的正则化方法,在长程依赖建模上表现乏力。近年来,基于注意力匹配(Attention Matching)的 KV 缓存压缩方法开始受到关注,其核心思想是:在压缩缓存的同时,尽可能保留对模型输出有重要影响的注意力模式,从而在内存占用与模型性能之间取得更好的平衡。
问题本质:为什么需要注意力感知的压缩策略
在自回归生成过程中,KV 缓存存储了每一个已处理 token 的键向量和值向量,供后续 token 的注意力计算复用。随着序列长度线性增长,缓存体积也以相同速度膨胀。问题的关键在于:并非所有历史 token 对未来预测的贡献都相等。语言模型中存在显著的语义冗余 —— 许多 token 在完成其上下文角色后,对后续推理的影响微乎其微。这一观察构成了 KV 缓存压缩的理论基础:我们可以有选择性地丢弃那些对模型输出影响较小的 KV 对,而保留那些被未来查询高频关注的 KV 对。
然而,准确地评估每个 KV 对的「重要性」面临根本性挑战。理论上,最优的评估指标应该是该 KV 对对残差流的实际贡献大小,这取决于两个因素:未来查询对该 key 的注意力权重,以及对应 value 向量经输出投影后的幅值。前者涉及尚未生成的 token,后者可在推理时直接计算。真正的困境在于,我们无法预知未来查询会如何分配注意力权重 —— 这些查询在压缩时尚未产生。这一限制使得直接基于真实注意力分数进行缓存淘汰变得不可行。
Expected Attention 的核心思路与理论推导
Expected Attention 提出了一种无需训练的解决方案,其核心创新在于:利用大型语言模型激活值的统计分布特性,对未来查询的注意力分布进行闭式估计。具体而言,研究者观察到现代 LLM 的隐藏状态近似服从高斯分布,这一现象在多个模型架构中得到了验证。既然隐藏状态服从高斯分布,经过线性投影变换后的查询向量同样继承高斯特性。通过对下一段位置的 RoPE 旋转矩阵取平均,可以得到一个位置无关的查询分布近似。
基于这一分布假设,Expected Attention 推导出每个 key 相对于未来查询的「期望注意力分数」的闭式表达式。对于给定的 key 向量 $k_i$ 和查询分布 $\bar {q} \sim \mathcal {N}(\bar {\mu}_q, \bar {\Sigma}_q)$,未归一化的期望注意力分数可以通过高斯矩生成函数获得:
$$\hat{z}i = \mathbb{E}{\bar{q}}\left[\exp\left(\frac{\bar{q}^T k_i}{\sqrt{d}}\right)\right] = \exp\left(\frac{\bar{\mu}_q^T k_i}{\sqrt{d}} + \frac{k_i^T \bar{\Sigma}_q k_i}{2d}\right)$$
进一步地,每个 KV 对对残差流的期望贡献幅度可以表示为:
$$|\widehat{\Delta h}_i| = (\hat{a}_i + \epsilon) |W_o v_i|$$
其中 $\hat {a}_i$ 是归一化后的期望注意力分数,$|W_o v_i|$ 是 value 向量经输出投影后的 L2 范数,$\epsilon$ 是一个小常数用于防止完全丢弃任何 token。这一指标综合考虑了未来查询的关注程度和 value 信息的丰富程度,为缓存淘汰提供了可靠的优先级排序依据。
工程实现的关键决策点
在将 Expected Attention 应用于实际推理系统时,需要关注以下几个工程实现的关键决策点。
首先是高斯参数估计策略。由于无法预知完整的未来查询序列,实现时采用一个固定大小的滑动窗口来估计查询分布的均值和协方差。论文中的实验表明,使用 128 个 token 的隐藏状态缓冲区进行统计估计已经能够提供足够准确的分布估计。对于超参数 $\epsilon$,默认值为 0.02,仅在 Needle in a Haystack 等极端检索任务中设置为 0 以避免边界情况下的误删。
其次是压缩触发频率与批量大小。在预填充阶段,通常在处理完整个输入 prompt 后执行一次性压缩,压缩比可设为 50% 至 90%。在解码阶段,由于缓存持续增长,系统设定一个最大缓存容量阈值 —— 当缓存达到该阈值后每隔固定步长(如 512 步)触发一次压缩淘汰。这一设计避免了每步都进行评分计算带来的开销,同时确保缓存大小始终控制在预算范围内。
第三是头级别与层级别的自适应压缩比率。不同注意力头在模型中扮演的角色存在显著差异:某些头负责捕获关键依赖关系,对缓存淘汰更敏感;而另一些头则相对冗余,允许更激进的压缩。Expected Attention 支持为每个头设置独立的压缩预算,更重要的头保留更多 KV 对。在实践中,可以基于注意力熵或层深度的先验知识来分配预算 —— 浅层的注意力头通常比深层更重要。
实践参数配置清单
以下是针对不同推理场景的推荐参数配置,可作为工程实现的起点。
对于长上下文问答任务(输入长度 32k 至 128k token),建议在预填充完成后执行压缩,压缩比设置为 50%。这意味着缓存体积减半,同时在 LongBench 等基准测试中模型性能可保持与未压缩版本相当。期望注意力评分窗口参数 $T$ 设为 512,即对后续 512 个位置的查询分布进行平均。
对于长程推理任务(如代码生成、数学推理),缓存需求更为苛刻,因为模型需要维护较长的思维链。这类场景下可以将压缩比提高至 75%(即 4 倍压缩)或更高(90% 对应 10 倍压缩)。需要注意的是,高压缩比会导致 NIAH 检索任务出现性能下降,因此对于同时包含检索与推理的复合任务,建议将压缩比控制在 50% 至 75% 之间。
对于流式生成场景(如实时对话),缓存压缩发生在解码过程中。建议设置缓存容量上限为 2GB(针对 8B 参数量级模型),每生成 512 个 token 评估一次淘汰需求。在此配置下,系统可以在单卡 H100 上运行 128k 上下文的对话模型,峰值显存相比未压缩版本降低约 15GB。
在多模型混合部署场景下,不同模型可采用不同的压缩策略以优化整体吞吐量。重要性较高的模型采用较低的压缩比(如 30%),而辅助模型可以使用更激进的压缩(如 70%)。通过这种差异化配置,可以在有限的 GPU 资源上同时服务更多并发请求。
与现有系统的集成考量
Expected Attention 的实现可以无缝集成到基于 Hugging Face Transformers 的推理 pipeline 中。框架通过 forward hook 的方式在每个注意力层的 forward 完成后介入,计算当前 KV 对的重要性分数并执行淘汰,无需修改模型架构本身。NVIDIA 发布的 KVPress 库提供了这一方法的参考实现以及超过 20 种压缩方法的基准测试代码,适合作为算法选型的验证工具。
需要注意的是,当前实现基于 PyTorch,未针对生产环境进行深度优化。在高吞吐场景下,建议将评分计算与淘汰逻辑迁移至 CUDA kernel 中实现,以减少 CPU-GPU 数据传输开销。此外,Expected Attention 可以与 KV 缓存量化方法(如 KVQuant、KIVI)正交叠加使用 —— 先通过注意力匹配淘汰低价值条目,再对保留下来的缓存进行量化,可进一步降低内存占用。
监控指标与回滚策略
部署基于注意力匹配的缓存压缩系统时,建议监控以下关键指标以确保服务稳定性。
重建误差(Reconstruction Error)是衡量压缩质量的直接指标,定义为压缩前后隐藏状态的 L2 距离。若重建误差在某一层突然增大,说明该层的压缩策略可能过于激进,需要调低对应层的压缩比或调整头级别的预算分配。实践中建议设置误差阈值 —— 当某层的平均重建误差超过基线的 1.5 倍时触发告警。
任务级指标退化是最终的业务衡量标准。在部署初期建议保留一份未压缩的缓存作为对照版本,定期对比相同输入下的输出质量差异。若特定任务类型的准确率下降超过预设阈值(如 2%),应自动回滚至更保守的压缩配置或禁用该任务的缓存压缩。
缓存命中率与淘汰率反映了压缩策略的有效性。若大量被淘汰的 KV 对在后续生成中又被重新访问,说明评分函数的排序结果与实际需求存在偏差。此时可考虑结合位置信息(如保留初始的 sink token)或增加历史窗口长度来改进评分准确性。
总结
基于注意力匹配的 KV 缓存压缩为长上下文 LLM 推理提供了一条无需训练、低实现成本的优化路径。Expected Attention 通过建模未来查询的分布,在理论上为每个 KV 对的重要性提供了可计算的量化指标。在实践中,通过合理配置压缩比、评分窗口大小和头级别预算,可以在 2 倍至 10 倍的缓存压缩范围内保持模型性能的可接受水平。建议工程团队从 50% 压缩比、512 评分窗口的基础配置开始验证,根据具体业务场景的任务指标进行微调。
资料来源:本文核心方法论源自 arXiv 论文《Expected Attention: KV Cache Compression by Estimating Attention from Future Queries Distribution》(arXiv:2510.00636),该工作由 NVIDIA 研究团队发布,同时参考了 KVPress 开源库的实现细节。