针对长上下文大模型推理重新定义 KV 缓存淘汰问题

# 重构长上下文 LLM 推理中的 KV 缓存淘汰问题
来源：https://arxiv.org/html/2605.07234
Tho Mai 韩国大田 KAIST [email protected] &Joo-Young Kim 韩国大田 KAIST [email protected]

###### 摘要

大型语言模型（LLMs）支持长上下文推理，但由于键-值（KV）缓存的增长，面临着巨大的内存和运行时开销。现有的 KV 缓存淘汰方法主要依赖局部注意力权重，忽略了价值表示、输出投影以及头间交互的影响。在本工作中，我们将 KV 缓存淘汰从传统的按头平均权重的方法，重构为一个感知输出的、按层的矩阵乘法近似问题。我们引入了 LaProx，这是一种新颖的淘汰策略，它显式地对注意力图与投影价值状态之间的乘法相互作用进行建模，从而在考虑头间依赖关系的同时，准确地量化 token 的贡献。基于这一指标，我们提出了第一种统一的淘汰策略，为 token 分配全局可比较的重要性分数，从而实现模型范围内的选择，而不是局部的按头决策。在 LongBench 和 Needle-In-A-Haystack 等长上下文基准的 19 个数据集上的实验结果表明，我们的方法在仅保留 5% 的 KV 缓存的情况下维持了模型性能，并在所有配置下始终优于先前的工作。值得注意的是，与现有最先进基线相比，我们的方法在极端压缩场景下最多可将准确率损失降低 2 倍，且开销极小。

## 1 引言

大型语言模型（LLMs）的最近进展显著扩展了它们处理长上下文的能力，在多轮对话<sup>40</sup>、问答<sup>25</sup>、代码生成<sup>18</sup>和文档理解<sup>50</sup>等应用中实现了强大的性能。为了加速自回归推理，Transformer 会缓存来自前序 token 的键和值状态，避免重复计算注意力。虽然这种键-值（KV）缓存对于高效解码至关重要，但其大小随上下文长度线性增长，迅速成为长上下文设置中内存使用和延迟的主要瓶颈。虽然头合并或架构修改<sup>2</sup>等技术可以在训练期间部分缓解这些成本，但它们通常与部署中常用的固定、预训练模型不兼容。因此，在推理时高效管理 KV 缓存——无需重新训练或改变模型参数——成为在现实内存和硬件约束下可扩展且成本效益高的长上下文 LLM 部署的关键挑战<sup>53</sup>。

为了在受限的内存预算下运行大型语言模型，一种常见的策略是在推理过程中动态减少键-值（KV）缓存的大小，通过淘汰被认为影响较小的条目。 prior work has shown that, in practice, only a small subset of cached tokens meaningfully contributes to the attention output<sup>52</sup>; <sup>43</sup>，这促使了一类基于淘汰的方法，选择性保留关键条目而丢弃其余部分。早期方法利用了注意力权重高度集中的经验观察，即少数 token 始终接收大部分注意力质量。基于这一现象，几种方法通过随时间平均注意力分数来识别重要的缓存条目，随后的改进引入了观察窗口、池化机制<sup>30</sup>或自适应预算分配<sup>34</sup>以更好地保留显著信息。然而，这些策略往往是启发式的，缺乏对缓存条目关键性的原则性公式化。因此，注意力行为、价值表示及其对最终模型输出的联合影响之间的精确关系仍未得到充分表征。

在本文中，我们将 KV 缓存淘汰重新表述为一个在固定预算下保持层注意力输出的优化问题。通过显式地将输出建模为注意力、值和输出矩阵的乘积，我们超越了传统的仅基于注意力的启发式方法，使我们能够根据缓存条目对形成最终层输出的乘法相互作用的实际贡献对其进行排序。关键的是，这一公式表明，token 的重要性从根本上与每层内形成的聚合表示以及模型最终输出耦合在一起。这一观察结果表明，淘汰在模型级别管理最为有效，而不是通过孤立的按头决策。基于这一见解，我们提出了一种新颖的淘汰策略，能够实现更有效的全局缓存选择。我们的贡献总结如下：

1. 我们证明仅凭注意力权重提供的 token 重要性度量是不完整的，准确的选择必须考虑输出信息以及注意力层本身的结构。
2. 我们揭示现有的独立按头淘汰次优，因为它忽略了头间和层间交互，并表明淘汰应在模型级别进行。
3. 我们引入了 LayerApproximated Cache (LaProx)，这是一种新的淘汰策略，通过同时评估跨头和跨层的 token 来近似层的输出，无需任何校准。
4. 在长上下文基准上的广泛评估表明，所提出的方法始终优于基于注意力的淘汰策略，证实了我们提议的有效性。

## 2 背景和相关工作

### 2.1 注意力和 KV 缓存操作基础

为清晰起见，我们使用多头注意力（MHA）描述机制，并省略层索引，注意该公式同样适用于所有 Transformer 注意力层。设 $\mathbf{X} \in \mathbb{R}^{S \times D}$ 为长度 $S$ 序列的 token 嵌入，其中 $D$ 是模型隐藏维度。每个注意力头在维度为 $d_h$ 的子空间上运行，其中 $D = H \cdot d_h$ 对应 $H$ 个头。投影矩阵 $\mathbf{W}_Q^{(h)}, \mathbf{W}_K^{(h)}, \mathbf{W}_V^{(h)} \in \mathbb{R}^{D \times d_h}$ 将共享隐藏表示映射到特定于头的查询、键和值状态。在处理提示期间，每个头计算

$\mathbf{Q}^{(h)} = \mathbf{X}\mathbf{W}_Q^{(h)}, \mathbf{K}^{(h)} = \mathbf{X}\mathbf{W}_K^{(h)}, \mathbf{V}^{(h)} = \mathbf{X}\mathbf{W}_V^{(h)}$ (1)

以及注意力权重

$\mathbf{A}^{(h)} = \operatorname{Softmax}\left(\frac{\mathbf{Q}^{(h)}\{\mathbf{K}^{(h)}\}^\top}{\sqrt{d_h}}\right)$ (2)

然后拼接每头的注意力输出，

$\mathbf{AV} = \operatorname{Concat}\left(\mathbf{A}^{(1)}\mathbf{V}^{(1)}, \dots, \mathbf{A}^{(H)}\mathbf{V}^{(H)}\right)$ (3)

并投影以产生最终注意力输出，

$\mathbf{O} = \mathbf{AV}\mathbf{W}_O$ (4)

经过投影 $\mathbf{W}_O$ 后，最终层输出通过残差连接集成：

$\mathbf{Y} = \operatorname{Norm}(\mathbf{O} + \mathbf{X})$ (5)

其中 $\mathbf{X}$ 是输入恒等映射，$\operatorname{Norm}$ 表示归一化函数。

在自回归解码期间，在每一步解码 $i$，仅将新生成的 token 嵌入 $\mathbf{x}_i \in \mathbb{R}^{1 \times D}$ 投影以获得其按头的查询、键和值状态。为了避免对过去 token 的重新计算，新的键-值对附加到缓存中

$\mathbf{K}^{(h)} \leftarrow \operatorname{Concat}\left(\mathbf{K}^{(h)}, \mathbf{x}_i\mathbf{W}_K^{(h)}\right), \quad \mathbf{V}^{(h)} \leftarrow \operatorname{Concat}\left(\mathbf{V}^{(h)}, \mathbf{x}_i\mathbf{W}_V^{(h)}\right)$ (6)

并且查询 $\mathbf{q}_i^{(h)} = \mathbf{x}_i\mathbf{W}_Q^{(h)}$ 使用公式 2 对缓存的键进行注意力计算。

虽然 KV 缓存显著减少了解码期间的计算，但缓存随序列长度线性增长，导致长上下文推理中巨大的内存和注意力开销。

### 2.2 KV 缓存淘汰

推理期间的 KV 缓存淘汰在不修改注意力机制的情况下减少内存和计算开销。其目标是在移除低影响 token 的同时保留重要 token。早期方法如 StreamingLLM<sup>45</sup> 采用基于窗口的策略，保留注意力 sink 和近期 token，而 LongFormer<sup>4</sup> 使用两种类型的滑动窗口与一些预选择的位置配合。虽然高效，但这些方法可能会丢弃长序列中间的信息性 token，降低长上下文性能。其他工作，包括 H2O<sup>52</sup> 和 Scissorhands<sup>32</sup>，使用累积注意力分数对 KV 条目进行排名，以更好地捕捉 token 重要性。基于这一系列工作，SnapKV<sup>30</sup> 和 CAKE<sup>34</sup> 通过在观察窗口内平均注意力并应用池化操作进一步提高了性能，取得了最先进（SOTA）的结果。

除了 token 选择外，几项研究探索了非均匀缓存预算分配。按层的方法如 PyramidInfer<sup>47</sup> 和 PyramidKV<sup>5</sup> 根据网络深度分配预算，而 D2O<sup>42</sup> 和 CAKE<sup>34</sup> 使用特定层的注意力方差调整缓存大小。在头级别，AdaKV<sup>11</sup> 在头分数上应用 top-k 选择并带有经验保障，而 HeadKV<sup>13</sup> 使用校准程序在推理前确定固定的每头预算。

少数工作超越了注意力分数。例如，LAVa<sup>36</sup> 和 CAOTE<sup>15</sup> 在其淘汰指标中利用价值表示但忽略了输出投影；与此同时，CriticalKV<sup>12</sup> 依赖两个经验保障用输出信息重新缩放平均注意力分数，忽略了注意力层的实际公式。

尽管现有方法具有竞争力的结果，但它们主要依赖注意力权重进行淘汰和预算分配，或者启发式地利用输出信息，而不考虑层的实际公式。此外，这些方法仅限于按头执行淘汰，忽略了跨头和跨层交互。相比之下，本工作提出了一种原则性的淘汰准则，结合了注意力概率和 $\mathbf{V}\mathbf{W}_O$ 贡献以及跨头交互，提供了更准确的 token 重要性度量。

## 3 动机

参见图 (a) $\mathbf{A}$ 和 $\mathbf{V}\mathbf{W}_O$ 模式。
参见图 (b) 平均强度。

图 1: $\mathbf{A}$ 和 $\mathbf{V}\mathbf{W}_O$ 的模式和幅度。在本节中，我们调查注意力权重 ($\mathbf{A}$) 与价值-输出投影 ($\mathbf{V}\mathbf{W}_O$) 之间的关系。具体而言，我们检查仅 $\mathbf{A}$ 的平均值是否可以作为注意力层输出的忠实代理，即注意力权重 $\mathbf{A}$ 是否足以表征整个乘积 $\mathbf{A}\mathbf{V}\mathbf{W}_O$。这种方法假设满足两个关键条件：(1) $\mathbf{A}$ 和 $\mathbf{V}\mathbf{W}_O$ 的模式良好对齐，以及 (2) $\mathbf{A}$ 的幅度不被 $\mathbf{V}\mathbf{W}_O$ 主导。

实验设置。我们的分析使用 Mistral-7B-Instruct-v0.3 模型进行。为了可视化清晰，我们仅显示连续的 token 子集。

观察。图 1(a)<sup>1</sup> 报告了 $\|\mathbf{A}\|$ 和 $\|\mathbf{V}\mathbf{W}_O\|$ 的每 token 归一化幅度。虽然这两个量共享一些高分 token（如 #25 或 #49-50），但它们的整体模式差异显著。许多步骤甚至显示相反的峰值；例如，token #37 和 #39 具有高 $\mathbf{V}\mathbf{W}_O$ 值但低 $\mathbf{A}$ 值。这表明 $\mathbf{A}$ 和 $\mathbf{V}\mathbf{W}_O$ 以不同方式评估 token 重要性，且不能相互替代。

此外，图 1(b)<sup>2</sup> 揭示 $\mathbf{A}$ 的值范围远小于 $\mathbf{V}\mathbf{W}_O$。由于注意力权重是归一化概率，其值局限于狭窄范围，而 $\mathbf{V}\mathbf{W}_O$ 具有更宽的值范围，且在更深层中扩展。

这些观察结果表明，仅凭注意力权重不足以表示注意力层输出，从而促使在缓存淘汰决策中纳入价值和输出投影。

## 4 方法论

### 4.1 淘汰指标

公式 3 和 4 表明，标准 MHA 定义为所有头输出的拼接后跟随线性投影。虽然输出投影混合了来自所有头的注意力信息，但计算可以精确分解为独立头贡献的总和。

算法 1 淘汰分数计算
输入：查询 $\mathbf{Q}$，KV 缓存 $(\mathbf{K}, \mathbf{V})$，投影 $\mathbf{W}_O$，预算 $B_{total}$，观察窗口 $w$
输出：压缩 KV 缓存 $(\tilde{\mathbf{K}}, \tilde{\mathbf{V}})$
// 计算注意力权重和投影值

$\mathbf{A} \leftarrow \operatorname{Softmax}\left(\frac{\mathbf{Q}[-w:,] \mathbf{K}^\top}{\sqrt{d_k}}\right)$

$\mathbf{H} \leftarrow \mathbf{V}\mathbf{W}_O$

// 评分 token

$T \leftarrow$ 缓存 token 数量

for $i=0$ to $T$ do

if $i < T-w$ then

$\mathbf{p}[i] \leftarrow \|\mathbf{A}[:,i]\|_2 \cdot \|\mathbf{H}[i,:]\|_2$

else

$\mathbf{p}[i] \leftarrow \infty$

endif

endfor

// 淘汰 token

$\mathcal{S} \leftarrow \operatorname{TopK}(\mathbf{p}, B_{total})$

$(\tilde{\mathbf{K}}, \tilde{\mathbf{V}}) \leftarrow (\mathbf{K}[\mathcal{S}], \mathbf{V}[\mathcal{S}])$

return $(\tilde{\mathbf{K}}, \tilde{\mathbf{V}})$

注记 4.1<sup>3</sup> 显示通过集成 $\mathbf{V}\mathbf{W}_O$...