LAVa: Layer-wise KV Cache Eviction with Dynamic Budget Allocation

发表时间: 2025-11 · EMNLP 2025 Findings

原文: https://aclanthology.org/2025.findings-emnlp.737

作者/机构：Yiqun Shen, Song Yuan, Zhengze Zhang, Xiaoliang Wang, Daxin Jiang, Cam-Tu Nguyen (南京大学，阶跃星辰)

速读

一句话结论 本文提出了一种名为 LAVa 的免训练 KV Cache 压缩方法，通过最小化 Transformer 残差流中的层级注意力输出损失，首次实现了同时在注意力头和网络层维度的动态预算分配，在保持长文本理解能力的同时大幅降低了显存占用并提升了推理速度。

要解决什么问题 大语言模型在处理 128K 甚至更长的上下文时，KV Cache 会占用海量显存，成为推理加速的绝对卡点。现有的缓存压缩方法（如 H2O、SnapKV）大多基于启发式观察，例如仅依赖累积的注意力分数来决定保留哪些 Token，缺乏严格的理论基础。更核心的机制卡点在于“预算分配的僵化”：不同任务对缓存的需求分布完全不同，但现有方法要么只能在单个网络层内动态分配各个注意力头的预算（如 AdaKV），要么只能跨层动态分配预算但要求每个头预算固定（如 CAKE、PyramidKV）。由于缺乏一个统一的理论指标来跨头、跨层公平比较缓存条目的重要性，现有方法往往需要拼凑多个启发式策略和超参数，无法在单一框架下实现完全动态的头与层预算分配，导致模型在面对不同类型任务（如文本生成与信息抽取）时，压缩后的效果会出现明显波动。

怎么做的 作者将 KV Cache 压缩问题转化为“最小化 Transformer 残差流信息损失”的优化问题。由于直接优化全局的 Logits 损失计算成本过高，LAVa 选择优化局部的“层级注意力输出损失”。通过推导该损失的 $L_1$ 范数上界，作者发现评估一个 Token 的重要性不仅要看注意力权重，还必须考虑其 Value 向量的幅度。基于此，LAVa 包含两个协同工作的核心机制：
第一步是计算 LAVa 分数并实现动态头预算。对于第 $l$ 层、第 $h$ 头的第 $i$ 个 Token，其重要性分数由最近 $w$ 个解码步的注意力分数之和，乘以该头在整个序列上 Value 向量 $L_1$ 范数的最大值构成：

$$s_{l,h}[i] = \frac{\max_{k \in [N]} \|V_{l,h}[k]\|_1}{w} \sum_{j=N-w}^N A_{l,h}^j[i]$$


由于引入了 Value 向量的缩放因子，不同注意力头产生的分数被拉平到了同一量纲。因此，LAVa 直接将该层所有头的分数展平为一个一维数组，进行全局排序并剔除分数最低的条目。这一操作在执行缓存驱逐的同时，自然实现了层内各头的动态预算分配。
第二步是实现动态层预算。LAVa 认为，如果某一层各个 Token 的分数分布越均匀（即不确定性越高），说明该层越难决定该丢弃谁，因此应该分配更多的预算。具体而言，将该层的 LAVa 分数归一化后计算信息熵 $e_l$：

$$e_l = \frac{-\sum_{h,i} (\hat{s}_{l,h}[i] \log \hat{s}_{l,h}[i])}{H \times N}$$
随后，根据各层熵值 $e_l$ 在所有层中的占比，按比例动态划分总预算。整个过程无需训练，也不依赖任何需要人工调优的超参数。

效果如何 实验在 Mistral-7B-Instruct-v0.2、Qwen2.5 系列（7B/14B/32B）和 Llama3-8B-Instruct 模型上展开，评测基准包括 LongBench、Needle In A Haystack、Ruler 和 InfiniteBench。对比基线涵盖了代表固定层/固定头预算的 SnapKV、代表固定层/动态头预算的 Ada-SnapKV、代表递减层/固定头预算的 PyramidKV、代表递减层/动态头预算的 Ada-PyramidKV，以及代表动态层/固定头预算的 CAKE。在 LongBench 评测中，当总预算限制为 128HL 时，LAVa 在 Mistral-7B 上取得 36.74 的平均分，显著优于 Ada-SnapKV（35.82）和 CAKE（35.06）。在 A800 80GB 硬件上使用 FlashAttention-2 测试，LAVa 在 128K 上下文长度下相比全量缓存实现了 9 倍的解码提速，且有效避免了 OOM。实验还揭示了一个重要结论：动态层预算对生成类任务（如代码补全）至关重要，而动态头预算对抽取类任务（如问答）更为关键，LAVa 凭借完全动态的特性在两类任务上均保持了最佳性能。该方法的代价是需要额外存储 Value 向量的范数和层级分数，相比 SnapKV 会增加约 0.6% 的显存和 0.01% 的计算开销；此外，作者承认在生成类任务上，压缩缓存与全量缓存之间仍存在一定性能差距，且目前尚未集成到 vLLM 等主流推理框架中。

主要贡献

在处理长上下文的大型语言模型（LLM）推理中，KV Cache（键值缓存）被广泛用于加速计算，但其高昂的内存需求使得缓存压缩技术变得不可或缺。现有的缓存压缩方法大多是启发式的，依赖于经验观察（如累积注意力分数）而非坚实的理论基础，并且缺乏动态的预算分配机制。

为了解决这一局限性，本文提出了一个用于缓存压缩的统一框架，其核心目标是最小化Transformer残差流中的信息损失。基于该框架，本文分析了层注意力输出损失（Layer Attention Output Loss），并推导出一个新的指标，用于在跨注意力头（heads）之间比较缓存条目，从而实现具有动态头预算的逐层压缩。此外，通过对比跨层的信息，本文还实现了动态的层预算分配。

LAVa是首个同时实现缓存驱逐和动态预算分配的统一策略，与以往方法不同，它既不需要训练，也不依赖于多种策略的复杂组合。在多个基准测试（LongBench、Needle-In-A-Haystack、Ruler和InfiniteBench）上的实验证明了其卓越性。此外，实验揭示了一个新的洞察：动态层预算对于生成任务（如代码补全）至关重要，而动态头预算在提取任务（如抽取式问答）中发挥着关键作用。作为一个完全动态的压缩方法，LAVa在各类任务中始终保持着顶级的性能表现。

解码器架构中的信息流
解码器架构中的信息流

背景知识与设计原则

KV Cache初始化与符号定义：KV Cache在预填充（prefilling）阶段初始化，按标准方式计算初始提示中token的Key和Value。假设存在包含$(N - 1)$个先前token的KV Cache，LLM有$L$层，每层$H$个头。模型和头的维度分别为$d$和$d_h = d / H$；$K_l, V_l$是第$l$层到当前时间步（第$N$个token）的KV Cache，大小为$[H, (N - 1), d_h]$。

解码过程的残差流视角：根据文献【6，Information flow routes: Automatically interpreting language models at scale+2024+EMNLP】，LLM解码可被视为在当前（第$N$个）残差流上操作。假设$x_l^N$是第$l$层的当前输入，首先计算对应的$Q_l^N, K_l^N, V_l^N$：$Q_l^N = x_l^N W_l^Q; K_l^N = x_l^N W_l^K; V_l^N = x_l^N W_l^V$。接着更新层级KV Cache：$K_l = Cat[K_l, K_l^N], V_l = Cat[V_l, V_l^N]$。然后计算第$l$层在第$N$步的注意力分数：$A_l^N = Cat_{h \in [H]}(A_{l,h}^N)$，其中$A_{l,h}^N = Softmax(\frac{Q_{l,h}^N (K_{l,h})^T}{\sqrt{d_h}})$。第$l$层的注意力输出计算如下：$y_l^N = Cat_{h \in [H]}(A_{l,h}^N V_{l,h}) W_l^O \in \mathbb{R}^{1 \times d}$。层输出$x_{l+1}^N$计算为$x_{l+1}^N = y_l^N + FFN(y_l^N)$，作为下一层的输入。在最后一层，利用去嵌入层（$W^M \in \mathbb{R}^{d \times |\mathcal{V}|}$）获得用于下一个token采样的概率向量$p^N$。

基于信息损失的统一缓存驱逐框架：给定KV Cache，压缩可以看作是在$K V$张量中掩蔽条目。形式上，为第$l$层和第$h$头定义注意力掩码$\mathcal{T}_{l,h}$：当保留$K_{l,h}[i]$和$V_{l,h}[i]$时，$\mathcal{T}_{l,h}[i]=1$；驱逐时为$0$。目标是找到一种KV Cache驱逐策略，以最小化所有后续残差流在最后一层逻辑值（$p^N$）的信息损失。设$\mathcal{P}$表示该逻辑值损失，$\mathbb{B}$为内存约束。预算分配和缓存驱逐的统一问题定义为：$\min_{B, s \in \mathcal{F}} \mathcal{P}(\mathbf{x}_1^{1 \ldots N}, s, \mathbb{B})$，其中$\mathcal{F}$表示所有评分函数的空间。此处$B_{l,h}$代表第$l$层第$h$头的预算，$B_l$代表第$l$层的总预算。最终约束确保保留窗口大小为$w$内的最新token。

局部信息损失替代方案：由于计算未来未见token的损失是不切实际的，且最小化全局逻辑损失的方法在评分函数计算昂贵时对于在线推理不可行，一个更可行的替代方案是关注局部信息并应用局部KV Cache驱逐。例如，头部注意力损失（Head Attention Loss）可用于逐头驱逐，此时评分函数是轻量级的，依赖于简单的统计特征（如逐头的注意力权重）。

方法细节

优化问题公式化：统一的预算分配和缓存驱逐问题可以形式化为：
$\min_{\mathcal{T}, \mathcal{B}} \mathcal{P}(\mathbf{x}_1^{1 \ldots N}, \mathcal{T}, \mathcal{B})$
约束条件为：$\sum_{i \in [N]} \mathcal{T}_{l,h}[i] = \mathcal{B}_{l,h}$；$\sum_{h \in [H]} \mathcal{B}_{l,h} = \mathcal{B}_l$；$\sum_{l \in [L]} \mathcal{B}_l = \mathbb{B}$。

层注意力输出损失的定义：基于层注意力输出损失（Layer Attention Output Loss）设计新算法。引理1指出，基于$L_p$范数，由于注意力掩码$\mathcal{T}$导致的层注意力输出损失，在当前（第$N$个）残差流的第$l$层测量如下：
$\mathcal{P}(\pmb{x}_1^{1 \dots N}, \mathcal{T}, \mathcal{B}) = \|\pmb{y}_l^N - \hat{\pmb{y}}_l^N\|_p = \|Cat_h [ (A_{l,h}^N - \frac{A_{l,h}^N \odot \pmb{\mathcal{T}}_{l,h}}{\|A_{l,h}^N \odot \pmb{\mathcal{T}}_{l,h}\|_1}) V_{l,h} ] W_l^O\|_p$
其中$\odot$表示逐元素乘法，$\hat{y}_l^N$表示用$\mathcal{T}$掩蔽KV Cache后获得的层注意力输出。

层注意力输出损失的上界分析：定理1给出了$L_1$范数的新上界：
$\|y_l^N - \hat{y}_l^N\|_1 \leq 2 \hat{C} \sum_{h \in [H]} \sum_{i \in [N]} A_{l,h}^N[i] \bar{V}_{l,h} (1 - \mathcal{T}_{l,h}[i])$
其中$\hat{C} = \|W_l^{O^T}\|_1$是一个独立于第$l$层内任何头或token的常数；$\bar{V}_{l,h} = \max_{k \in [N]} \|V_{l,h}[k]\|_1$是一个依赖于头的值。

贪婪驱逐策略与评分函数合理性：给定固定预算$B_l$，考虑一种贪婪算法，每次迭代驱逐一个缓存条目，直到满足缓存预算。为了最小化上界，驱逐由评分函数$s_{l,h}[i] = A_{l,h}^N[i] \bar{V}_{l,h}$给出的分数最小的条目。值得注意的是，该函数结合了一个依赖于头的值$\bar{V}_{l,h}$，在跨不同头比较KV Cache条目时不应忽略该值。这不同于AdaKV（【7，Ada-kv: Optimizing kv cache eviction by adaptive budget allocation for efficient llm inference+2024+arXiv】），后者考虑了层注意力输出损失但未考虑Value。这也为将Value引入评分提供了理论依据，而VATP（【11，Attention score is not all you need for token importance indicator in KV cache reduction: Value also matters+2024+EMNLP】）是在启发式中利用了这一点。因为该指标本质上基于逐层视角，使得能够设计动态预算分配策略。

LAVa分数定义：为了提高KV Cache驱逐的性能，受SnapKV（【16，SnapKV: LLM knows what you are looking for before generation+2024+NeurIPS】）启发，结合最近$w$个残差流的信息，产生一个新的评分函数。定义1：第$l$层、第$h$头的token $i$的层级注意力和值（LAVa）分数定义如下：
$s_{l,h}[i] = \frac{\max_{k \in [N]} \|V_{l,h}[k]\|_1}{w} \sum_{j=N-w}^N A_{l,h}^j[i]$

动态头预算的逐层缓存驱逐：基于LAVa评分函数，开发了层级KV Cache驱逐算法。该算法仅驱逐最近窗口$[N - w, N]$之外的条目，有效地保留了最新token。驱逐方法在第$l$层内的所有注意力头上运行。具体而言，将层中所有头的LAVa分数展平为一维数组$s_l$。然后比较并对跨所有头的$\mathcal{B}_l$个缓存条目进行排名以进行层级驱逐，从而在执行驱逐时有效地获得动态头预算。

基于不确定性的层预算分配：在确定驱逐哪个缓存条目时具有更大不确定性的层应分配更大的预算。基于LAVa分数，在第$l$层和第$h$头驱逐token $k$的概率通过归一化LAVa评分值获得：$\hat{s}_{l,h}[i] = \frac{s_{l,h}[i]}{\sum_{k,h} s_{l,h}[k]}$。第$l$层的不确定性然后通过归一化熵测量如下：$e_l = \frac{- \sum_{h,i} (\hat{s}_{l,h}[i] \log \hat{s}_{l,h}[i])}{H \times N}$。

在线预填充与层预算调整：为了解决预填充后（压缩前）内存峰值很高的问题，采用逐层执行预填充和缓存驱逐。从CAKE（【22，CAKE: Cascading and adaptive KV cache eviction with layer preferences+2025+ICLR】）中汲取灵感：在预填充第$l$层后，重新压缩较低的层（$< l$）。因此，使用相同的LAVa分数对较低层进行多次压缩，但预算被调整，随着内存与更多被预填充的层共享，预算随着时间的推移变小。

适配GQA架构：对于分组查询注意力（GQA）（【1，GQA: Training generalized multi-query transformer models from multi-head checkpoints+2023+EMNLP】），在应用LAVa分数时采取保守方法：token的组级分数确定为其在相应组内的头级分数的最大值。换句话说，只要条目对于组内的至少一个头是重要的，就倾向于保留该条目。

方法细节中引用的参考文献汇总

  • 【1】Ainslie et al. 2023. GQA: Training generalized multi-query transformer models from multi-head checkpoints. 发表会议: EMNLP. 引用内容: 描述了现代LLM广泛采用的分组查询注意力（GQA）机制。
  • 【6】Ferrando and Voita. 2024. Information flow routes: Automatically interpreting language models at scale. 发表会议: EMNLP. 引用内容: 描述了解码过程可以被视为在当前残差流上进行操作。
  • 【7】Feng et al. 2024. Ada-kv: Optimizing kv cache eviction by adaptive budget allocation for efficient llm inference. 引用内容: 指出AdaKV虽然考虑了层注意力输出损失，但忽略了Value的影响。
  • 【11】Guo et al. 2024. Attention score is not all you need for token importance indicator in KV cache reduction: Value also matters. 发表会议: EMNLP. 引用内容: 提及VATP方法在启发式设计中利用了Value的重要性。
  • 【16】Li et al. 2024. SnapKV: LLM knows what you are looking for before generation. 发表会议: NeurIPS. 引用内容: 启发了LAVa结合最近$w$个残差流信息的设计。
  • 【22】Qin et al. 2025. CAKE: Cascading and adaptive KV cache eviction with layer preferences. 发表会议: ICLR. 引用内容: 启发了LAVa的动态层预算调整机制，即在预填充时重新压缩较低层以控制内存峰值。

实验环境

  • 数据集：LongBench（包含21个数据集，涵盖单文档QA、多文档QA、摘要、少样本学习、合成任务和代码补全等6大类，支持中英双语）、Needle In A Haystack、Ruler和InfiniteBench。
  • 模型架构：Mistral-7B-Instruct-v0.2（32k上下文）、Qwen2.5-7/14/32B-Instruct（32k上下文）、Llama3-8B-Instruct（8k上下文）。所有模型均采用GQA架构。
  • 硬件配置：A800 80GB GPU。
  • 软件配置：采用FlashAttention-2实现。所有的token分数向量均采用最大池化（max pooling）操作进行平滑，核大小（kernel size）设置为7。
  • 基线方法：PyramidKV、SnapKV、Ada-SnapKV、Ada-PyramidKV和CAKE。

实验结果

LongBench主实验评估：
- 实验内容：在Mistral-7B-Instruct-v0.2上，评估不同缓存预算（128HL, 256HL, 512HL, 1024HL）下各方法的表现。
- 实验结果：LAVa在所有预算下均超越了基线方法，在较小预算下优势更为明显。在无需超参数调整的方法中（SnapKV, Ada-SnapKV, LAVa），LAVa表现最佳（例如在$\mathbb{B} = 128HL$时，LAVa平均得分为36.74，而Ada-SnapKV为35.82）。
- 分析结论：LAVa能有效保持模型在长文本理解上的性能。特别地，LAVa和CAKE在代码相关任务（如RepoBench-P）上表现出色。相关数据参考表2。

任务类型分析：
- 实验内容：将LongBench的20个数据集分为提取任务（如QA）和生成任务（如摘要和代码补全），分析Qwen和Mistral模型在不同策略下的平均得分。
- 实验结果：提取任务受压缩影响较小；生成任务中不同策略的性能差距更大。CAKE和LAVa在生成任务上优于固定层预算的方法，但CAKE在提取任务上表现较差。
- 分析结论：动态层预算对于生成任务至关重要，而动态头预算对于文本提取任务至关重要。LAVa在两类任务中均保持了顶级性能。
生成和提取任务的结果

延迟与峰值内存评估：
- 实验内容：在Mistral-7B-Instruct-v0.2上，使用1024HL分配预算，评估不同上下文长度下的峰值内存和解码延迟。
- 实验结果：LAVa引入的解码成本可忽略不计，在128K上下文长度下相比全缓存（Full Cache）实现了超过$9 \times$的加速。峰值内存保持在合理水平，避免了全缓存的OOM问题。
- 分析结论：LAVa在显著降低内存消耗的同时减少了延迟，且由于无需参数调优，比PyramidKV和CAKE更易于部署。
峰值内存使用和解码延迟

动态预算分配的消融研究：
- 实验内容：引入LAVa（-layer dynamic，强制统一层预算）和LAVa（-head dynamic，固定头预算，不进行跨头比较）进行对比。
- 实验结果：移除头级别或层级别的动态分配都会导致性能下降。
- 分析结论：头和层级别的动态预算分配对于性能都是必不可少的，进一步证实了层预算对生成任务重要，头预算对提取任务重要的结论。
LongBench上的消融研究

LAVa分数有效性分析：
- 实验内容：将LAVa的层分配替换为PyramidKV或Uniform，与Ada-PyramidKV和AdaKV进行胜率对比。
- 实验结果：LAVa分数在大多数情况下获得了显著更高的胜场数。
- 分析结论：证明了LAVa评分函数的有效性。
LAVa分数与AdaKV分数的比较

结论

本文将当前的KV Cache压缩方法统一到一个框架中，该框架基于最小化Transformer残差流中信息损失的原则。通过分析层注意力输出损失，提出了LAVa，这是一种新颖的逐层压缩方法，支持完全动态的头和层预算分配。实验表明，动态层预算对于生成任务至关重要，而动态头预算对于提取任务很重要。作为一种完全动态的压缩方法，LAVa在不同的任务类型和LLM架构中始终保持最高性能，同时在128K上下文长度下实现了与全缓存相比$9 \times$的加速。未来的方向包括探索基于该框架的新压缩算法，以及将框架扩展到模型压缩。

附录与补充细节

相关工作对比：
大多数现有的KV Cache压缩方法（如依赖最近性、累积注意力分数等）通常假设统一预算，或将KV Cache驱逐和预算分配视为独立问题，需要结合独立策略。DuoAttention需要训练并限制掩码模式；Dodo通过训练LoRA适配器来决定掩码。与LAVa最密切相关的是AdaKV，它旨在最小化层输出扰动，但仅将推导出的指标局部应用于头预算分配。相比之下，LAVa提出了一个用于逐层缓存驱逐并具有动态层预算的指标，且基于完整的理论推导而非纯启发式（如VATP仅启发式地将Value范数与注意力分数相乘）。

时间复杂度与内存使用分析：
- 时间复杂度：以SnapKV为基准，原缓存计算需$O(HN^2d_h)$，近期注意力分数重计算需$O(HNwd_h)$，驱逐需$O(HN \log B_{l,h})$。LAVa额外需要$O(HNd_h)$来计算Value范数，层级驱逐需$O(HN \log B_l)$。当$N=10,000$时，LAVa相比SnapKV的额外计算量仅约为$0.01\%$。
- 内存使用：动态层预算方法（LAVa和CAKE）需要在预填充期间分配内存。LAVa需要额外存储层分数，其大小为$O(LHB_l)$。在典型设置下，LAVa的额外内存使用量仅为SnapKV峰值内存的$0.6\%$。

局限性：
尽管提出了一个具有多个优化机会的统一框架，但理论分析和实验仅集中在一个方向上。为了进一步缩小与全缓存设置的性能差距（特别是在生成任务中），应探索其他方法。此外，需要进一步研究为什么动态层预算对生成任务至关重要。最后，除了FlashAttention-2，该方法尚未集成到其他广泛使用的推理框架（如vLLM）中。