CompressKV：语义检索引导的 KV 缓存压缩，实现资源高效长上下文 LLM 推理

背景

随着大型语言模型（LLM）在长上下文任务中的应用日益广泛，推理过程中的内存占用和解码成本成为了主要的瓶颈。特别是在基于分组查询注意力（Grouped-Query Attention, GQA）的架构中，键值（Key-Value, KV）缓存的大小随着上下文长度的增加而线性增长。这种增长不仅限制了模型在资源受限硬件上的可持续部署，也极大地增加了推理延迟和硬件成本。

现有的 KV 缓存淘汰（eviction）方法通常采用启发式策略，对所有注意力头（attention heads）的令牌进行统一评分，从而决定保留哪些 KV 对。然而，这种方法存在一个根本性的缺陷：它忽略了不同注意力头在功能上的差异性。在 GQA 模型中，不同的注意力头负责捕捉不同层面的语义信息。简单的全局评分往往会导致那些承载关键语义信息的令牌被错误地淘汰，从而显著降低 LLM 的性能表现。因此，如何在不牺牲模型精度的前提下，高效地压缩 KV 缓存，成为了当前研究的一个关键挑战。

核心内容

为了解决上述问题，研究团队提出了 CompressKV，这是一个专为基于 GQA 的 LLM 设计的资源高效 KV 缓存压缩框架。CompressKV 的核心创新在于其“语义检索引导”的机制，它不再盲目地聚合所有头的注意力分数，而是通过识别特定的“语义检索头”来指导令牌的保留。

1. 识别语义检索头（Semantic Retrieval Heads, SRHs）

CompressKV 的关键步骤是识别出那些能够捕捉提示词（prompt）中关键信息的注意力头，即语义检索头（SRHs）。这些头具有特定的功能特性：

• 首尾捕捉能力：SRHs 能够有效捕捉提示词的初始令牌和最终令牌。
• 中间上下文证据捕捉：SRHs 还能捕捉中间上下文中具有语义重要性的证据令牌。

通过定位这些特定的头，CompressKV 能够更精准地判断哪些令牌对于维持模型性能至关重要，从而避免淘汰那些对回答长上下文问题至关重要的信息。

2. 基于 SRHs 的令牌选择

一旦确定了 SRHs，CompressKV 便利用这些头的注意力分数来筛选需要保留的 KV 对。与以往方法对所有头进行平均或最大池化不同，CompressKV 仅依赖 SRHs 的反馈来选择令牌。这种方法确保了被保留的令牌在语义上是高价值的，从而在大幅减少缓存大小的同时，最大限度地保持模型的推理能力。

3. 分层缓存预算分配

除了令牌选择，CompressKV 还解决了缓存预算在不同网络层之间的分配问题。该方法利用离线估计的每层淘汰误差（layer-wise eviction error），动态地分配缓存预算。这意味着，对于更容易因淘汰关键信息而导致性能下降的网络层，CompressKV 会分配更多的缓存空间，从而在整体资源受限的情况下实现最优的性能平衡。

4. 实验验证

研究团队在 LongBench 和 Needle-in-a-Haystack 等基准测试上对 CompressKV 进行了广泛评估。实验结果表明，CompressKV 在各种内存预算下均一致地优于现有的 KV 缓存淘汰方法。具体数据亮点包括：

• 在 LongBench 问答任务中，仅使用 3% 的全量 KV 缓存，即可保留超过 97% 的全缓存性能。
• 在 Needle-in-a-Haystack 任务中，仅使用 0.7% 的 KV 存储，即可达到 90% 的准确率。

这些结果有力地证明了 CompressKV 在长上下文 LLM 推理中实现了更好的资源-性能权衡。

关键要点

• 问题痛点：现有 GQA 模型的 KV 缓存淘汰方法忽略注意力头的功能差异，导致关键语义令牌被误删，损害模型性能。
• 核心机制：引入“语义检索头”（SRHs）概念，专门识别捕捉提示词首尾及中间关键语义证据的注意力头。
• 精准压缩：仅基于 SRHs 的注意力分数来选择保留的 KV 对，而非聚合所有头的分数，提高了令牌选择的语义准确性。
• 分层优化：根据离线估计的每层淘汰误差，动态分配各层的缓存预算，优化整体资源利用效率。
• 极致效率：
  • LongBench 任务：3% 缓存占用 -> 97% 性能保留。
  • Needle-in-a-Haystack 任务：0.7% 缓存占用 -> 90% 准确率。
• 开源贡献：代码已公开，便于社区复现和进一步研究。

意义与影响

CompressKV 的提出对于推动大语言模型在边缘设备和资源受限环境中的部署具有重要意义。

首先，它解决了长上下文推理中内存瓶颈的关键技术难题。通过大幅降低 KV 缓存的内存占用（低至 0.7%-3%），使得在显存有限的 GPU 或边缘设备上运行超长上下文模型成为可能，降低了硬件门槛和运营成本。

其次，CompressKV 提供了一种更精细化的 KV 缓存管理思路。它强调了注意力头的异质性（heterogeneity），指出“所有头同等重要”的假设是错误的。这一洞察为后续研究如何更细粒度地优化注意力机制和缓存策略提供了新的方向。

最后，CompressKV 在性能与资源之间取得了显著的平衡。在保持极高准确率的同时，将存储需求降低了两个数量级，这对于构建可持续、可扩展的 AI 基础设施至关重要。随着长上下文应用（如长文档分析、长期记忆代理等）的需求增长，此类高效压缩技术将成为 LLM 工程化落地不可或缺的一环。