KVzip: Query-Agnostic KV Cache Compression with Context Reconstruction

发表时间: 2025-12 · arXiv:2505.23416 (NeurIPS 2025)

原文: https://arxiv.org/abs/2505.23416

Jang-Hyun Kim, Jinuk Kim, Sangwoo Kwon, Jae W. Lee, Sangdoo Yun, Hyun Oh Song
Seoul National University, Neural Processing Research Center, NAVER AI Lab

主要贡献

基于Transformer的大型语言模型（LLMs）在推理过程中会将上下文缓存为键值（KV）对。随着上下文长度的增加，KV缓存的大小急剧膨胀，导致了巨大的内存开销并增加了注意力机制的计算延迟。例如，在FP16精度下，Qwen2.5-14B缓存120K个Token需要约33GB的内存，这甚至超过了该模型在同等精度下的28GB参数存储需求。

现有的KV缓存驱逐方法（如SnapKV、PyramidKV）主要采用在推理过程中在线计算的“查询感知（query-aware）”KV对重要性评分，选择性地保留与当前查询最相关的KV对。虽然这种方法在单查询场景下有效，但在多查询场景下，由于保留的KV对过度拟合于初始查询，导致后续查询的性能显著下降。

为了解决这一问题，本文提出了KVzip，一种新颖的“查询无关（query-agnostic）”的KV缓存驱逐算法。该算法的核心研究目标是为给定的上下文优化出一个可重用的压缩KV缓存，从而在面对未来多样化的查询时实现高效推理。KVzip的创新点在于，它利用LLM本身从缓存的KV对中重建原始上下文，以此来量化KV对的重要性，随后驱逐重要性较低的KV对。这种方法特别适用于离线准备KV缓存的场景，例如保留用户指令和聊天历史的个性化对话代理，或利用预先计算的文档KV缓存进行检索的企业系统。

广泛的实证评估表明，KVzip将KV缓存大小减少了$3 - 4 \times$，并使FlashAttention的解码延迟降低了约$2 \times$，同时在问答、检索、推理和代码理解任务中的性能损失微乎其微。评估涵盖了LLaMA3.1、Qwen2.5和Gemma3等多种模型，上下文长度最高可达170K个Token。在多查询场景下，即使在$90\%$的缓存预算比例下，现有的查询感知KV驱逐方法也会出现性能退化，而KVzip显著优于这些现有方法。

背景知识与关键观察

符号与问题定义 考虑文本域$\mathcal{T}$和一个自回归的基于Transformer的LLM $f_{\mathrm{LM}} : \mathcal{T} \to \mathcal{T}$，该模型通过贪婪解码生成序列。模型包含$L$层，采用分组查询注意力（GQA），具有$H$个KV头，每个KV头由一组$G$个查询头关注。在推理期间，$f_{\mathrm{LM}}$将隐藏表示缓存为KV对以提高计算效率。给定一个被分词为$n_c$个Token的输入上下文$c \in \mathcal{T}$，预填充阶段生成一个包含$L \times H \times n_c$个KV对的缓存，记为$\mathrm{KV}_c$。使用该缓存的条件生成记为$f_{\mathrm{LM}}(\cdot \mid \mathrm{KV}_c)$。我们的目标是推导出一个紧凑的修剪缓存$\mathrm{KV}_{c,\mathrm{evicted}} \subseteq \mathrm{KV}_c$，使其满足：
$f_{\mathrm{LM}}(q|\mathrm{KV}_{c,\mathrm{evicted}}) \approx f_{\mathrm{LM}}(q|\mathrm{KV}_c), \forall q \in \mathcal{T}$。

现有方法的局限性 现有的KV驱逐方法，如SnapKV 【引文编号33，Snapkv: Llm knows what you are looking for before generation, NeurIPS 2024】和PyramidKV 【引文编号6，Pyramidkv: Dynamic kv cache compression based on pyramidal information funneling, 2024】，基于预填充期间给出的信息来压缩KV缓存。这些方法利用尾部上下文窗口内的查询来计算KV对的基于注意力的重要性评分，选择性地保留与这些查询相关的KV对。虽然对于单查询基准测试有效，但这些方法需要为每个新查询重复进行缓存预填充。另一方面，如果将先前压缩的KV缓存重用于后续查询，通常只会保留与初始查询相关的上下文KV对，无法泛化到不同的查询，导致性能显著下降（如图2所示）。

模型作为上下文编码-解码器的直觉 为了有效回答任意查询，压缩后的缓存$\mathrm{KV}_{c,\mathrm{evicted}}$和$f_{\mathrm{LM}}$应保留完整的上下文信息。我们的直觉是，可以通过明确提示$f_{\mathrm{LM}}$从$\mathrm{KV}_{c,\mathrm{evicted}}$中重建先前的上下文来验证这种完整性。如果$\mathrm{KV}_{c,\mathrm{evicted}}$使得$f_{\mathrm{LM}}$能够使用重复提示（repeat prompt）准确重建原始上下文$c$，我们就可以重新预填充原始缓存$\mathrm{KV}_c$并进行准确的推理。尽管在每次推理时重新生成原始缓存实际上是不可行的，但实证研究表明，即使不重建原始缓存，压缩后的缓存也表现出强大的泛化能力。

重建过程中的注意力稀疏性观察 与在初始预填充$\mathrm{KV}_c$期间计算的自注意力模式相比，在上下文重建期间获得的交叉注意力模式表现出更大的稀疏性。在预填充期间，模型在Token之间密集交互以编码全面的上下文信息。然而，在重建中，模型有效地利用了（1）存储在$\mathrm{KV}_c$中的高级表示和（2）编码在模型权重中的内部知识，从而减少了不必要的注意力查找。这种交叉注意力的稀疏性有效地识别并移除了冗余的KV对，优于依赖预填充期间获得的注意力分数的先前方法（如$\mathrm{H_2O}$ 【引文编号60，H2o: Heavy-hitter oracle for efficient generative inference of large language models, NeurIPS 2023】）。

跨任务的注意力重叠观察 通过比较不同任务（重复、问答、摘要和推理）的最大交叉注意力分数，我们发现，在重建中获得高注意力的KV特征在其他任务中也获得高注意力。相反，比较两个不同QA任务的热力图显示出特定的查询相关注意力变异性。这一观察证明，对重建至关重要的KV对在各种任务中都持续发挥作用，支持了KVzip的有效性。

方法细节

KVzip整体流程 我们算法的主要目标是为每个KV对分配一个重要性评分，以确定驱逐优先级。给定上下文长度$n_c$，KVzip为$\mathrm{KV}_c$中的KV对分配重要性评分$S \in \mathbb{R}^{L \times H \times n_c}$，随后驱逐得分最低的对。我们的方法支持非均匀和均匀的头部预算分配 【引文编号17，Ada-kv: Optimizing kv cache eviction by adaptive budget allocation for efficient llm inference, 2024；引文编号33，Snapkv: Llm knows what you are looking for before generation, NeurIPS 2024】。KVzip进一步适应了头部级别的驱逐策略，通过在序列维度$n_c$上取最大的对级别分数来计算头部级别分数 【引文编号53，Duoattention: Efficient long-context llm inference with retrieval and streaming heads, ICLR 2025】。

KV重要性评分机制 KVzip基于KV对在上下文重建中的贡献来量化其重要性。具体而言，我们通过教师强制解码（teacher-forced decoding）模拟重建，通过单次前向传递并行化，输入序列包含一个重复提示（repeat prompt）后跟原始上下文。我们将重要性评分定义为每个KV对在此前向传递期间接收到的最大注意力分数，这利用了接收到最小注意力的KV对对Transformer计算贡献甚微的见解 【引文编号60，H2o: Heavy-hitter oracle for efficient generative inference of large language models, NeurIPS 2023】。在形式上，给定长度为$n_c$的上下文，我们通过将长度为$n_{\mathrm{prompt}}$的重复提示与上下文拼接，构建长度为$n_{\mathrm{in}} = n_{\mathrm{prompt}} + n_c$的输入序列。通过带有$\mathrm{KV}_c$的$f_{\mathrm{LM}}$前向传递此输入，为第$l$层的第$h$个KV头生成$d$维分组查询特征$Q_{l,h} \in \mathbb{R}^{G \times n_{\mathrm{in}} \times d}$和键特征$K_{l,h} \in \mathbb{R}^{(n_c + n_{\mathrm{in}}) \times d}$。这些特征之间的分组注意力产生注意力矩阵$A_{l,h} = \mathrm{Softmax}(Q_{l,h}K_{l,h}^\intercal) \in \mathbb{R}_+^{G \times n_{\mathrm{in}} \times (n_c + n_{\mathrm{in}})}$。提取对应于$\mathrm{KV}_c$中键的条目，得到切片注意力矩阵$\bar{A}_{l,h} \in \mathbb{R}_+^{G \times n_{\mathrm{in}} \times n_c}$。最后，我们通过在分组查询上取最大值，计算第$l$层第$h$个KV头的重要性评分$S_{l,h} \in \mathbb{R}^{n_c}$：
$S_{l,h} = \max_{g=1, \ldots, G; \ i=1, \ldots, n_{\mathrm{in}}} \bar{A}_{l,h}[g,i]$。
我们将所有KV头上聚合的分数$S$称为最大交叉注意力分数。

技术挑战与分块评分解决方案 我们的方法将重复提示与上下文Token拼接，通过$f_{\mathrm{LM}}$处理此输入以获得注意力矩阵。然而，由于注意力矩阵随上下文长度$n_c$呈二次方扩展，直接计算长上下文的注意力矩阵是不可行的。虽然像FlashAttention这样的融合注意力内核通过分块计算注意力分数而不存储完整矩阵来减少内存开销，但我们的方法独特地要求在沿键维度进行Softmax归一化之后，沿查询维度进行最大化。这种跨维度的依赖性阻碍了将上述公式直接集成到现有的分块注意力算法中。为了应对这一挑战，我们引入了基于分块的评分（chunk-based scoring），独立地重建上下文片段。通过在固定大小的块中计算重要性评分，而不是在整个上下文上同时计算，计算复杂度从二次方$O(n_c^2)$降低到线性$O(m n_c)$，其中$m$表示块的大小。具体来说，我们将上下文Token划分为长度为$m$的固定块，将每个块与重复提示拼接，并通过$f_{\mathrm{LM}}$处理长度为$n_{\mathrm{in}} = n_{\mathrm{prompt}} + m$的结果输入。对于每个Transformer层，我们对对应于每个块的$\mathrm{KV}_c$中的键进行子采样，获得大小为$n_{\mathrm{in}} \times (m + n_{\mathrm{in}})$的较小注意力矩阵。如公式所示，对注意力矩阵进行切片并在分组查询上最大化，产生分块重要性评分。我们对每个块重复该过程，并汇总分数以获得$\mathrm{KV}_c$的完整重要性评分。我们将块大小设置为$m = 2\mathrm{K}$，这在不同上下文长度、模型和任务中保持不变。

复杂度分析与系统开销 假设重复提示长度可以忽略不计（即$n_{\mathrm{prompt}} \ll m$，因此$n_{\mathrm{in}} \approx m$），每个块的计算复杂度为$O(m^2)$。对所有$n_c / m$个块重复此计算，总复杂度为$O(m n_c)$，与上下文长度呈线性关系。峰值内存开销为$O(m^2)$，它不随$n_c$变化，并且与模型参数和KV缓存大小相比可以忽略不计。重要性评分引入了通过带有$\mathrm{KV}_c$的$f_{\mathrm{LM}}$计算分块输入的注意力查询和键的额外开销。FlashAttention每个块产生$O(n_c m + m^2 / 2)$的因果注意力FLOPs，导致所有$n_c / m$个块的总复杂度为$O(n_c^2 + n_c m / 2)$。这一成本大约是标准预填充因果注意力复杂度$O(n_c^2 / 2)$的两倍。为了提高效率，KVzip还支持上下文无关的驱逐，通过为每个模型分配静态的头部级别重要性评分，在部署后不产生压缩开销。经验评估证实，在压缩期间计算开销约为标准预填充的两倍，且额外内存极小（低于$2\%$）。

参考文献引用汇总

• 【引文编号6，Pyramidkv: Dynamic kv cache compression based on pyramidal information funneling, 2024】：在背景知识中引用，说明现有方法基于预填充信息计算分数。
• 【引文编号17，Ada-kv: Optimizing kv cache eviction by adaptive budget allocation for efficient llm inference, 2024】：在方法细节中引用，说明KVzip支持非均匀的头部预算分配策略。
• 【引文编号33，Snapkv: Llm knows what you are looking for before generation, NeurIPS 2024】：在背景知识中引用，说明其基于尾部查询计算分数；在方法细节中引用，说明支持均匀头部预算分配。
• 【引文编号53，Duoattention: Efficient long-context llm inference with retrieval and streaming heads, ICLR 2025】：在方法细节中引用，说明KVzip适应了其提出的头部级别驱逐策略。
• 【引文编号60，H2o: Heavy-hitter oracle for efficient generative inference of large language models, NeurIPS 2023】：在背景知识中引用，说明基于预填充的注意力分数不如重建有效；在方法细节中引用，支撑了接收最小注意力的KV对对计算贡献甚微的观点。

实验环境

• 驱逐结构：采用非均匀头部预算分配策略，保留跨所有注意力头重要性评分前$r\%$的KV对（系统提示保持完整）。
• 数据集：评估涵盖12个基准数据集，包括SQuAD、GSM8K、needle-in-a-haystack (NIAH) 以及SCBench中的9个任务（涵盖RULER和$\infty$Bench的任务）。上下文长度从100到170K个Token不等，覆盖长文档问答、检索、数学推理、上下文学习和代码理解。
• 模型配置：采用指令微调的LLMs，包括Qwen2.5-7B-1M、LLaMA3.1-8B和Gemma3-12B。GQA组大小从4到7不等。Gemma3采用全局和滑动窗口混合注意力机制。
• 硬件与软件：所有评估在单张NVIDIA A100 80GB GPU上进行，使用Bfloat16精度。此外，还集成了QServe量化框架（8位权重、8位激活和4位KV缓存）。基准方法包括$\mathrm{H_2O}$、SnapKV、PyramidKV和DuoAttention。

实验结果

任务泛化能力 使用Qwen2.5-7B-1M在12个基准数据集上的多查询评估结果表明，KVzip在各种任务中表现出强大的泛化能力。在检索密集型任务中，我们的方法在$30\%$缓存比例下仍能保持性能，而基线方法在$90\%$保留率时就出现显著退化。在上下文理解任务（如GSM8K）中，KVzip在低至$20\%$的比例下实现近乎无损的压缩。在具有高上下文冗余的任务中，KVzip容忍低至$10\%$的激进压缩且无性能下降，甚至因减少注意力干扰而显示出性能提升（见图9）。

模型规模与架构的适应性 在更大的模型（Qwen2.5-14B-1M）、不同的模型系列（LLaMA3.1-8B）和混合注意力架构（Gemma3-12B）上，KVzip同样证明了其泛化能力和卓越的压缩性能。对于Gemma，我们将KV驱逐专门应用于全局注意力层，结果证实其优于基线方法（见图10）。

KV量化集成效果 KVzip与KV缓存量化有效集成。在4位KV量化模型（LLaMA3-8B-W8A8KV4）上，KVzip在量化下保持稳健。在124K输入长度下，将4位量化与我们的$70\%$驱逐率相结合，有效地将缓存大小从16.3GB减少到1.2GB，且性能下降微乎其微。

上下文无关驱逐性能 KVzip支持上下文无关驱逐策略，只需每个模型进行一次重要性评分，部署后无压缩开销。与DuoAttention需要数小时的优化相比，KVzip使用自然语言教科书在不到一分钟内完成几次前向传递，获得了更优越的性能（见图11）。

补充细节

上下文重建的必要性分析 在重要性评分中，我们比较了使用重复提示结合上下文前$10\%$、后$10\%$或仅使用重复提示的输入。结果明确表明，完整的上下文重建对于防止KV驱逐导致的性能下降是必不可少的（见图12）。

超出任务解决的行为分析（隐私泄露风险） 进一步的分析揭示了KV驱逐引发的一个与隐私相关的有趣行为。对于涉及私人上下文信息的查询，LLaMA3.1-8B模型在使用完整KV缓存时拒绝响应，但在应用我们的压缩方法后却显著做出了响应。这是因为KVzip优先考虑重建所需的KV对并丢弃其他，这暗示了KV驱逐技术与浅层对齐问题（shallow-alignment）之间存在交叉，值得未来研究。

附录

重建块大小的影响 分析评分块大小$m$对性能的影响表明，在$0.3$的KV缓存比例下，1K和2K块大小之间的平均性能差异保持在$2\%$以下，证实了其影响微乎其微。因此，我们在所有实验中采用2K的块大小，以实现最佳计算效率。

重复提示的鲁棒性 实验比较了原始重复提示、释义版本和无重复提示。结果显示KVzip对提示变化具有鲁棒性。分析表明，保留的KV特征中$99.4\%$的最大注意力来源于重复的上下文，证实了提示词对评分的影响极小。

无Softmax的重要性评分变体 我们开发了一个不带Softmax归一化的KVzip变体，通过自定义Triton内核直接集成到FlashAttention中。虽然这减少了约$10\%$的前向计算开销，但导致压缩率下降了约$10\%$。

均匀KV头部预算 尽管KVzip在均匀头部预算分配下也优于基线，但非均匀分配通过更有效地捕获跨头部重要性的变化，实现了更卓越的压缩性能。

详细基准测试图表 附录提供了在RULER基准测试（图17）、Qwen2.5-14B-1M（图18）、LLaMA3.1-8B（图19）、Gemma3-12B（图20）、LLaMA3-8B-W8A8KV4（图21）以及多任务数据集（图22）上的详细评估结果，一致证明了KVzip在各种比例下的优越性。

结论

本文介绍了KVzip，一种查询无关的KV缓存驱逐算法，通过从KV对中重建原始上下文来有效优化可重用的压缩KV缓存。在涵盖不同任务、模型和长上下文基准的多查询设置的广泛评估中，KVzip展示了强大的压缩性能，在性能损失微乎其微的情况下将KV缓存大小减少了高达$70\%$，同时通过FlashAttention将解码注意力延迟显著提高了约$2 \times$。KVzip始终优于在$10\%$驱逐比例下就出现性能下降的现有KV驱逐方法。KVzip的实际适用性进一步扩展到量化模型和各种KV缓存结构，突显了其适应性和效率。未来的工作可以进一步探索无Softmax评分等硬件高效实现，以及KV驱逐与模型对齐安全性之间的潜在关联。