PrimeKV：无需注意力分数的关键KV缓存保留用于长上下文推理

《Neurocomputing》：PrimeKV: Preserving Critical KV Cache Without Attention Scores for Long-Context Inference

　　

**摘要**
- 提出了PrimeKV，一种具有理论依据的KV缓存压缩算法。
- 分析了单个及批量KV对移除情况下的输出扰动。
- 实现了无需注意力分数的KV缓存压缩。
- 在仅上下文场景中实现了高达244.6%的相对性能提升。
- 在查询漂移条件下，将多轮对话性能提升了高达28.62%。

**引言**
大语言模型（LLMs）[1]是一类包含数百亿甚至更多参数的深度神经网络，通常通过对海量无标注文本进行自监督学习来训练[2]。自2018年以来，包括Google、OpenAI、Meta、DeepSeek-AI、Alibaba等知名研究机构和企业，已陆续推出众多模型，如Gemini[3]、GPT[4]、Llama[5]、DeepSeek[6]、Qwen[7]等，几乎所有自然语言处理（NLP）任务上的表现均达到SOTA水平。用户可通过自然语言与系统交互，完成从理解到生成等多种任务，例如问答[8]、摘要[9]、翻译[10]、对话[11]等。

在自回归解码过程中，由于每一步自注意力[12]都需计算当前Query与所有历史已处理Token之间的兼容性，因此LLMs面临二次方累积的注意力复杂度。设提示长度为m，生成Token数为n，第i步解码时的序列长度为Li=m+i?1。若不使用KV缓存，模型将在每步解码时重新计算整个序列的Key与Value投影以及自注意力输出，导致单步复杂度为O(Li2)，累积解码复杂度为O(nm2+mn2+n3)。KV缓存[13]存储了所有已处理Token的Key与Value投影。因此，每步解码时，模型仅需为新生成Token计算Key与Value投影，并复用已缓存的历史投影。新的Query随后与所有Li个已缓存的Key和Value进行匹配，使得单步注意力复杂度降至O(Li)，累积解码复杂度降至O(nm+n2)。KV缓存大幅减少了冗余计算并降低了解码延迟，但同时也引入了显著的内存开销。缓存大小随上下文长度、层数、KV头数及每头隐藏维度线性增长。例如，一个采用分组查询注意力（GQA）[14]和8位KV量化[15]的65B参数模型，在处理512K Token的上下文时仍需要约86GB的GPU显存，超过了单张80GB H100 GPU的容量[13]。

受限于长上下文推理场景下对严格内存限制的同时保持高模型保真度的关键需求，我们提出了PrimeKV——一种无需训练、无需参数的KV缓存压缩算法，通过基于扰动的估计方法识别重要的KV对。PrimeKV利用KV缓存本身固有的特性来估计每个KV对的重要性，而无需显式的注意力分数或尾部观察窗口。由于其评分过程无需在压缩时物化或导出完整的注意力矩阵，因此在评分需求层面，PrimeKV与FlashAttention[16]、PagedAttention[17]等无需注意力分数的推理内核具有算法兼容性，独立于任何特定的内核实现或推理框架。

基于此框架，我们的工作做出了三项主要贡献。第一，我们分析了KV对移除所导致的输出变化，并以此为启发式动机，用于估计已缓存KV对的边际重要性。第二，基于该分析，我们提出了PrimeKV——一种基于Token级直接贡献估计的KV缓存压缩算法。PrimeKV既不需要注意力分数，也不需要超参数，在算法设计层面与现有的KV缓存效率技术（如GQA、多查询注意力（MQA）[18]及KV缓存量化）正交。第三，除传统的单轮对话任务（如LongBench[19]）外，我们首次在多轮对话场景下评估了KV缓存压缩算法，为实际应用中的交互场景提供了更现实的基准测试方法。

我们在RULER[24]、LongBench、基于L-Eval[25]构建的多轮对话数据集以及AIME25[26]上，对PrimeKV与主流算法（如Full KV（无KV缓存压缩）、StreamingLLM[20]、SnapKV[21]、PyramidKV[22]和KeyDiff[23]）进行了全面评估。实验结果表明，PrimeKV在仅上下文场景中表现出更好的性能，显著优于基于注意力的方法。此外，在大多数场景中，PrimeKV相比其他非基于注意力的替代方案也具备强大的竞争力。在具有挑战性的多轮对话任务中，PrimeKV取得了显著的性能优势，因为它不依赖固定的观察窗口注意力锚点，而是从两个互补的内在信号估计KV对的重要性：Key侧的不可替代性与Value侧的投影贡献。在AIME25上，PrimeKV在KV缓存压缩条件下的性能与Full KV相当，进一步证明了其在数学推理上的有效性。这些发现表明，PrimeKV能够在多样化的查询无关设置下提供有效的无需注意力分数的重要性估计，从而为长上下文KV缓存压缩提供了一个实用的算法解决方案。

**相关工作**
高效的KV缓存优化已成为高吞吐长上下文LLMs部署的关键研究方向。现有的主流KV缓存优化方法可系统性地划分为8种核心范式，每种范式均有来自顶级会议和高影响力研究的代表性验证工作：(1) 剪枝/驱逐/稀疏方法，保留语义和计算上关键的KV对，丢弃冗余Token，以SnapKV[21]作为标志性工作……

**方法**
为实现高效的KV缓存压缩，我们的核心目标是识别关键的KV对，使其在自注意力过程中更好地保留完整KV缓存的输出。在本节中，我们将从输出扰动的角度分析KV对的选择，并据此提出我们的压缩策略。具体而言，3.1节推导了KV对移除所导致的精确输出变化，并引入了直接贡献分数作为Token级排序的启发式动机。

**模型**
我们使用了五个覆盖不同模型家族、参数规模和架构设计的大语言模型，包括Llama-3.1-8B-Instruct（Llama-3.1-8B）[55]、Qwen3-8B[56]、Qwen3-30B-A3B-Instruct-2507（Qwen3-30B）[56]、Mistral-7B-Instruct-v0.3（Mistral-7B）[57]和DeepSeek-R1-Distill-Qwen-32B（DeepSeek-R1）[58]。骨干模型的选取基于各基准测试的特性和既有评估设置。对于RULER[24]，我们保留了Llama-3.1-8B和Qwen3-8B，因为它们是……

**理论意义**
从理论角度来看，本研究推导了KV对移除的精确输出扰动表达式，并将每个KV对的直接贡献作为重要性建模的启发式动机。该分析为在有限内存预算下构建Token级排序分数提供了一个有原则的视角。此外，针对实际场景中无法获得完整注意力分数的局限性，我们进一步提出了一种无需注意力的保留……

**局限性**
当前实验代码基于PyTorch和KVPress项目，主要聚焦于验证PrimeKV的算法行为。由于其评分过程无需物化的注意力分数，因此PrimeKV在算法层面与FlashAttention和PagedAttention等无需注意力分数的推理内核兼容。然而，这种兼容性仅限于算法设计。我们尚未实现定制的CUDA/Triton内核，也尚未将PrimeKV集成到……

**结论与未来工作**
PrimeKV是一种受输出扰动分析启发的无参数KV缓存压缩算法。它不仅在KVPress官方排行榜上取得了强劲表现，还证明了它能够在不依赖注意力分数的情况下保持高模型保真度。由于PrimeKV不需要显式的注意力分数，因此在算法层面与FlashAttention和PagedAttention等无需注意力分数的推理内核兼容。

在RULER任务中，PrimeKV表现出了……

**CRediT作者贡献声明**
Ilvovsky Dmitry：写作—审阅与编辑、监督、资源。
Zhipeng Zhang：写作—原稿、验证、软件、方法、调查、形式化分析、数据整理、概念化。

**关于生成式AI及AI辅助技术在写作过程中的声明**
在准备本文的过程中，作者使用了Grammarly以纠正语法错误。使用此工具后，作者根据需要审阅并编辑了内容，并对已发表文章的内容承担全部责任。

**利益冲突声明**
作者声明，他们没有已知的竞争性经济利益或个人关系可能影响本文所报告工作的独立性。

**致谢**
本文是国立高等经济学院（HSE University）基础研究项目计划的一部分，由研究项目（HSE-BR-2025-025）的输出。本文还部分得益于HSE University的HPC设施提供的计算资源支持。

Zhipeng Zhang于2017年获得电子信息工程学士学位，2024年获得信息系统与技术硕士学位。他目前正在俄罗斯莫斯科国立高等经济学院计算机科学系攻读人工智能与机器学习博士学位。他曾担任程序员、数字化转型顾问和讲师，在软件开发、数据工程、机器学习和人工智能方面具有丰富经验。

Zhipeng Zhang | Dmitry Ilvovsk