Cake: Cascading and Adaptive KV Cache Eviction with Layer Preferences

发表时间: 2025-03 · arXiv:2503.12491 (ICLR 2025)

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

作者/机构：Ziran Qin, Weiyao Lin (上海交通大学)；Yuchen Cao, Wen Hu, Shixuan Fan, Ke Cheng, Jianguo Li (蚂蚁集团)；Mingbao Lin (独立研究员)

速读

一句话结论
本文提出了一种名为 CAKE 的级联自适应 KV Cache 驱逐方法，通过全局视角的层级显存动态分配和容忍注意力偏移的词元保留策略，在仅使用 3.2% KV Cache 的极端显存限制下，依然维持了长上下文大语言模型的无损表现，并大幅提升了推理吞吐量。

要解决什么问题
大语言模型在处理超长上下文时，KV Cache 的线性增长会导致极大的推理期显存压力。为了在固定显存预算下进行推理，现有的 KV Cache 驱逐方法通常会丢弃部分不重要的键值对。然而，原有做法卡在两个机制缺陷上：首先，现有方法大多对所有模型层分配相同大小的缓存（均匀分配），或者采用固定的金字塔形状分配，完全忽略了不同层在注意力机制上的巨大差异。实际上，有些层的注意力分布很广，有些层则高度集中。其次，在决定丢弃哪些词元时，传统方法仅依赖静态的累加或平均注意力分数，忽略了注意力随时间推移的动态变化（时间偏移）。这种做法很容易过早地把当前看似不重要、但在后续解码步骤中会突然变得关键的词元踢出显存，从而导致模型在长文本检索和复杂推理时出现严重的信息丢失和外推失效。

怎么做的
核心思路是将固定预算下的显存分配视作“切蛋糕”问题，根据每一层对 KV Cache 的实际偏好进行动态自适应分配，并配合一种新的词元保留指标来绕开传统方法的信息丢失卡点。该方法由三个关键部件构成：
第一，层级偏好评估。作者发现，注意力分散（需要保留更广的上下文）和注意力频繁转移（需要支持复杂的时间模式）的层更需要显存。因此，方法提取了近期窗口内的注意力权重，计算其空间分散度（信息熵 $\mathcal{H}$）和时间偏移度（方差 $\mathcal{V}$），定义出第 $l$ 层的偏好得分 $\mathcal{P}_l$：

$$ \mathcal{P}_l = \mathcal{H}^{\frac{1}{\tau_1}} \cdot \mathcal{V}^{\frac{1}{\tau_2}} $$


基于此得分，系统会打破均匀分配的限制，将总缓存预算 $B_{\mathrm{total}}$ 按比例自适应地分配给各层，计算出该层的目标缓存大小 $B_l$。
第二，级联缓存管理。如果等所有层计算完偏好再分配，预填充阶段的峰值显存会暴涨。为此，设计了一种级联机制：在预填充阶段，每计算完一个新层的注意力，就利用当前已获得的各层偏好得分重新瓜分一次显存预算，并同步更新已处理层的缓存。这种做法在数学上被证明与全局一次性分配的驱逐结果完全等价，但成功将峰值显存严格控制在了预算红线之内。
第三，容忍注意力偏移的驱逐指标。为了防止误删重要词元，方法放弃了单一的注意力均值，引入了多维度的重要性评估指标 $\mathbf{I}_l$。对于词元 $n$，其保留优先级由持续重要性（均值）和注意力波动性（方差）共同决定：

$$ \mathbf{I}_l[n] = \mathrm{Mean}(\mathbf{A}_l[:, n]) + \gamma \cdot \mathrm{Var}(\mathbf{A}_l[:, n]) $$
最终，系统只保留该指标得分最高的 $B_l$ 个词元的 KV Cache。均值保底长期记忆，方差捕捉突发关键信息，两者结合有效避免了动态重要词元被提前驱逐。

效果如何
实验在 NVIDIA A100 80GB GPU 上进行，覆盖了 Llama2、Llama3.1、Mistral、Qwen2.5 和 Gemma 等多种架构（参数量从 7B 到 70B 不等）。对比基线包括代表均匀分配路线的 StreamingLLM（保留首尾）、H2O（累加注意力）、TOVA（最后词元注意力）、SnapKV（近期注意力聚类），以及代表固定非均匀分配路线的 PyramidKV（金字塔形分配）。在 LongBench（长文本理解）和 NeedleBench（大海捞针检索与推理）测试中，CAKE 展现出极强的压制力。在仅保留 3.2% KV Cache 的极端低显存设置下，CAKE 依然能维持模型处理 32K 词元长上下文的能力，甚至在单针检索任务中超越了全量缓存的表现。在 128K 上下文长度配合 FlashAttention-2 的测试中，CAKE 相比全量缓存方案降低了约 48.63% 的峰值显存，同时解码延迟实现了 10 倍以上的加速。代价与局限性方面，作者承认该方法目前仅在“层”级别优化了缓存分配，尚未深入到更细粒度的“注意力头”级别。此外，在极度压缩的显存预算下，由于硬性容量限制无法承载多重线索所需的信息量，模型在复杂的多针检索任务中依然会出现性能的断崖式下跌。

主要贡献

大型语言模型（LLMs）在处理长序列方面表现出色，但这显著增加了对键值（KV）缓存的需求。随着上下文长度扩展至128K甚至更长，KV缓存大小呈线性增长，导致推理时的内存负担急剧加重。尽管近期在无需额外训练的KV缓存驱逐（Eviction）方法上取得了一定进展，通过移除不重要的KV对来缓解推理负担，但这些方法通常在不同层之间分配统一的缓存大小，未能根据不同注意力模式合理分配资源。这种忽略层级特定需求的统一分配方式，在严格的内存限制下会严重损害模型性能。

为了解决这一问题，本文将固定内存预算下为不同层分配最佳缓存大小的挑战框架化为“切蛋糕问题（cake-slicing problem）”，并提出了级联自适应KV缓存驱逐方法（Cascading and Adaptive KV cache Eviction，简称CAKE）。该研究的核心目标是通过全局视角评估各层对KV缓存的偏好，自适应地分配资源，并在保持内存预算的同时优化整体性能。

本文的主要创新点包括：
1. 分析了注意力动态变化，揭示了空间分散度（spatial dispersion）和时间偏移（temporal shifts）的特征，从而提出了一种针对特定层缓存大小需求的度量标准。
2. 提出了一种自适应缓存分配策略，能够基于层级偏好从全局视角优化整体缓存分配。
3. 设计了一种级联缓存管理方法，在预填充（prefilling）阶段动态调整KV缓存，在不牺牲驱逐性能的前提下，实现了与统一分配策略相当的内存使用效率。
4. 引入了一种新的驱逐指标，该指标综合考虑了Token的持续重要性和时间变异性，有效提升了Token驱逐的性能。

广泛的实验表明，CAKE在仅使用3.2% KV缓存的情况下，依然能维持模型在LongBench和NeedleBench基准测试中的性能，并在低内存设置下持续优于现有的基线方法。此外，在处理128K Token上下文并结合FlashAttention-2时，CAKE的解码延迟相比全缓存方法实现了超过10倍的加速。

图1：注意力模式在空间（a, b）和时间（c, d）特征上的变化。左侧提供示例，右侧为Mistral不同层的真实示例。
图1：注意力模式在空间（a, b）和时间（c, d）特征上的变化。左侧提供示例，右侧为Mistral不同层的真实示例。
图2：CAKE与现有缓存分配策略的比较说明。(a)统一缓存分配；(b)固定形状缓存分配；(c)CAKE中使用的偏好优先自适应缓存分配。
图2：CAKE与现有缓存分配策略的比较说明。(a)统一缓存分配；(b)固定形状缓存分配；(c)CAKE中使用的偏好优先自适应缓存分配。

背景知识与关键观察

KV缓存操作基础。对于单个注意力头，其权重矩阵定义为$\mathbf{W}_Q, \mathbf{W}_K, \mathbf{W}_V \in \mathbb{R}^{D \times D}$，其中$D$表示模型的隐藏维度。给定提示词嵌入$\mathbf{X} \in \mathbb{R}^{S \times D}$（$S$为序列长度），注意力模块包含提示预填充和Token解码两个阶段。在提示预填充阶段，首先计算查询（Query）、键（Key）和值（Value）状态：$\mathbf{Q} = \mathbf{X}\mathbf{W}_Q, \mathbf{K} = \mathbf{X}\mathbf{W}_K, \mathbf{V} = \mathbf{X}\mathbf{W}_V$。随后，注意力模块的输出确定为$\operatorname{Attn}(\mathbf{Q}, \mathbf{K}, \mathbf{V}) = \mathbf{A}\mathbf{V}$，其中注意力权重$\mathbf{A} = \mathrm{Softmax}(\frac{\mathbf{Q}\mathbf{K}^T}{\sqrt{D}}) \in \mathbb{R}^{S \times S}$。键值状态$\mathbf{K}$和$\mathbf{V}$随后被存储在缓存中，建立起KV缓存。在Token解码阶段，针对每个解码步骤$i$，新的Token嵌入$\mathbf{x}_i \in \mathbb{R}^{1 \times D}$生成新的KV对$\mathbf{x}_i\mathbf{W}_K, \mathbf{x}_i\mathbf{W}_V \in \mathbb{R}^{1 \times D}$，并与现有的KV缓存拼接以更新缓存：$\mathbf{K} = \mathrm{Concat}(\mathbf{K}, \mathbf{x}_i\mathbf{W}_K), \mathbf{V} = \mathrm{Concat}(\mathbf{V}, \mathbf{x}_i\mathbf{W}_V)$。尽管KV缓存缓解了计算需求，但其随序列长度线性增长的特性给有限的缓存预算带来了巨大挑战。

KV缓存驱逐的局限性。推理期间的KV缓存驱逐通过策略性减少缓存来提升效率而不改变注意力机制。早期的StreamingLLM和LM-Infinite优先保留首尾Token，但这容易忽略序列中间的重要Token。随后的研究引入了更复杂的指标，如累积注意力分数（H2O）、最后Token注意力分数（TOVA）、平均注意力分数或对近期注意力分数进行聚类（SnapKV）。然而，这些方法通常采用统一分配策略，无法最大化缓存利用率。PyramidInfer和PyramidKV采用金字塔形分配，而D2O根据当前层的注意力密度调整缓存大小。这些方法要么依赖预定义的静态分配，要么仅基于局部层的注意力，未能全面捕捉注意力机制的复杂动态，也缺乏对层级缓存偏好的全局视角。

注意力动态的空间与时间分析。为探究解码器架构模型中的注意力动态，本文对注意力权重矩阵$\mathbf{A}$进行了空间和时间维度的量化分析。空间分析关注行$\mathbf{A}[i, :]$，揭示了第$i$个Token在单步生成时对其他Token的注意力分布情况。本文采用空间注意力分散度（Spatial attention dispersion）来量化，定义为$\mathbf{A}[i, :]$的熵：$H(\mathbf{A}) = -\sum_{i=0}^{S-1} \mathbf{A}[i, :] \log(\mathbf{A}[i, :])^T$。该值越高表明注意力分布越均匀，越低则表明注意力高度集中。时间分析关注列$\mathbf{A}[:, j]$，追踪对第$j$个Token的注意力随时间步的演变。本文采用时间注意力偏移（Temporal attention shift）来量化，定义为$\mathbf{A}[:, j]$的方差之和：$V(\mathbf{A}) = \sum_{j=0}^{S-1} \mathrm{Var}(\mathbf{A}[:, j])$。该值越高表示注意力在不同位置间发生显著转移，越低则表示注意力保持稳定。

注意力动态的实验观察。基于LongBench数据集的长上下文样本，对多个LLM层的注意力动态进行了分析。实验结果表明，在不同的层、模型和上下文中，注意力分散度和偏移量存在显著差异。这些发现突显了LLM注意力机制的复杂性，证明了采用统一或固定模式的缓存分配策略是低效的，迫切需要一种能够有效处理注意力动态变化的定制化KV缓存管理策略。

图3：注意力动态分析。(a)空间注意力分散度（上）和时间注意力偏移（下）的热力图。(b)空间注意力分散度（上）和时间注意力偏移（下）在不同层和模型间的变化。
图3：注意力动态分析。(a)空间注意力分散度（上）和时间注意力偏移（下）的热力图。(b)空间注意力分散度（上）和时间注意力偏移（下）在不同层和模型间的变化。

方法细节

偏好优先的自适应分配策略。由于注意力机制在不同层、模型和上下文中存在变异性，统一或固定模式的缓存分配策略在有限内存预算下效率低下。为了应对这一挑战，本文引入了偏好优先的自适应分配策略，该策略从全局注意力模式的视角出发，综合考虑每一层的独特特征并自适应地分配缓存大小。具体而言，本文为每一层的KV缓存需求定义了一个偏好度量指标，该指标同时结合了空间分散度和时间偏移：$\mathcal{P} = \mathcal{H}^{\frac{1}{\tau_1}} \cdot \mathcal{V}^{\frac{1}{\tau_2}}$，其中$\mathcal{H} = \mathrm{H}(\mathbf{A}[-S_w:, -S_w:])$，$\mathcal{V} = \mathrm{V}(\mathbf{A}[-S_w:, -S_w:])$。在该公式中，$\mathcal{P}$代表注意力层缓存大小的偏好分数。因为分散的注意力（高$\mathcal{H}$）需要保留更广泛的上下文，而频繁的偏移（高$\mathcal{V}$）需要支持复杂的时间模式，所以具有高偏好分数$\mathcal{P}$的层能从更大的KV缓存中获益更多以维持性能。考虑到注意力分散和偏移在不同模型和内存约束下的重要性不同，本文引入了温度参数$\tau_1$和$\tau_2$来调整它们对$\mathcal{P}$的影响（【引用文献1，Snapkv: Llm knows what you are looking for before generation+2024+arXiv】，【引用文献2，Pyramidinfer: Pyramid kv cache compression for high-throughput llm inference+2024+arXiv】）。本文将计算焦点集中在预填充注意力权重矩阵$\mathbf{A}$的最近窗口大小为$S_w$的子矩阵$\mathbf{A}[-S_w:, -S_w:]$上，这一设计受到近期研究的启发，表明通过分析最近查询的注意力模式可以有效捕捉解码模式。

层级缓存预算的动态归一化。本文使用特定层的偏好分数来实现偏好优先的自适应分配策略。这些分数会根据不同的模型和上下文动态调整，从而确保了适应性。每一层的缓存大小通过归一化偏好分数并据此分配总内存预算来设定。这将生成一个自适应缓存大小集合$\mathbf{B} = \{B_l\}_{l=0}^{L-1}$，其计算公式为：$B_l = \frac{\mathcal{P}_l}{\sum_{k=0}^{L-1} \mathcal{P}_k} \cdot B_{\mathrm{total}}$，其中$\mathcal{P}_l$是第$l$层的偏好分数，$B_{\mathrm{total}}$是总缓存预算。这种方法针对每一层的独特性质和输入上下文，优化了缓存大小的分配。

偏好引导的级联缓存管理动机。尽管偏好优先的自适应分配策略是最优的，但它最初需要全面获取所有层的预填充注意力权重。这导致峰值内存使用量达到$\mathcal{O}(S \cdot L)$，可能会超过设定的缓存预算$B_{\mathrm{total}}$。为了克服这个问题，本文进一步开发了偏好引导的级联缓存管理方法。该方法在预填充期间动态维护缓存预算，有效地将峰值内存使用量降低至目标值。

级联缓存管理算法执行流程。本文的缓存管理过程将预填充过程划分为$L$个阶段，每个模型层对应一个阶段，并在偏好分数的引导下级联地管理缓存内存。在每个阶段$m$，当一个新层完成其预填充计算时，系统会根据当前已获得的偏好分数重新分配预算，并在所有已处理的层之间根据重新分配的预算更新KV缓存。具体过程如算法1所示：首先初始化缓存集合、指标集合、预算集合和偏好分数集合。接着从阶段0遍历到$L-1$，在每一层$m$，将当前层的$\{\mathbf{K}_m, \mathbf{V}_m\}$加入缓存集合，计算偏好分数$\mathcal{P}_m$，并使用所选驱逐方法的局部注意力权重计算指标$\mathbf{I}_m^{(m)}$。然后，更新偏好分数和指标集合，并计算当前阶段各层的预算$B_l^{(m)}$。随后，遍历从层0到层$m$，针对每一层$l$，提取当前缓存、预算和指标，执行驱逐操作$\mathrm{EVICT}$以获取更新后的缓存$\{\hat{\mathbf{K}}_l, \hat{\mathbf{V}}_l\}^{(m)}$和指标$\hat{\mathbf{I}}_l^{(m)}$，并将其作为下一阶段的基础。

分配预算的单调递减性质。该算法的方法确保了在预填充阶段保持恒定的内存使用，同时维持与标准策略等效的缓存分布和驱逐结果。理论命题1保证了分配给每一层的缓存预算随着阶段的推进呈单调递减趋势。对于任意层$l \in [L]$，从阶段$l$到$L-1$，分配的预算大小严格满足：$B_l^{(m+1)} < B_l^{(m)}, m \in [l, L-1]$。其中$B_l^{(m)}$是基于当前获取的偏好分数在阶段$m$为层$l$重新分配的缓存预算：$B_l^{(m)} = \frac{\mathcal{P}_l}{\sum_{k=0}^m \mathcal{P}_k} \cdot B_{\mathrm{total}}$。这一性质确保了无论使用何种驱逐方法，更新后的KV缓存$\{\mathbf{K}_l, \mathbf{V}_l\}^{(m+1)}$始终是前一阶段缓存$\{\mathbf{K}_l, \mathbf{V}_l\}^{(m)}$的真子集。

驱逐操作的具体定义。在介绍等价性定理之前，本文明确了算法相关的操作细节。给定指标向量$\mathbf{I}_l \in \mathbb{R}^S$，它衡量了第$l$层Token的重要性。对于每个阶段$m \in [L]$，驱逐操作记为$\mathrm{EVICT}(\cdot)$，它从当前缓存$\mathbf{K}_l^{(m)}, \mathbf{V}_l^{(m)} \in \mathbb{R}^{B_l^{(m-1)} \times D}$中保留与$\mathbf{I}_l^{(m)} \in \mathbb{R}^{B_l^{(m-1)}}$中得分最高的$B_l^{(m)}$个位置相对应的KV对$\hat{\mathbf{K}}_l^{(m)}, \hat{\mathbf{V}}_l^{(m)} \in \mathbb{R}^{B_l^{(m)} \times D}$。指标向量$\mathbf{I}_l^{(m)}$也更新为$\hat{\mathbf{I}}_l^{(m)} \in \mathbb{R}^{B_l^{(m)}}$，仅保留与被保留位置相对应的元素。更新后的缓存和指标将作为下一阶段$m+1$进一步更新的基础。

级联驱逐等价性定理。定理1证明了，对于任意层$l \in [L]$，通过从阶段$l$到阶段$L-1$的级联驱逐获得的KV缓存$\{\mathbf{K}_l, \mathbf{V}_l\}^{(L-1)}$，等价于在完整的KV缓存$\{\mathbf{K}_l, \mathbf{V}_l\} \in \mathbb{R}^{S \times D}$上使用目标缓存预算$B_l$执行一次驱逐操作的结果：$\{\mathbf{K}_l, \mathbf{V}_l\}^{(L-1)} = \mathrm{EVICT}\big(\{\mathbf{K}_l, \mathbf{V}_l\}^{(m)}, B_l^{(m)}, \mathbf{I}_l^{(m)}\big)_{m=l}^{L-1} = \mathrm{EVICT}\big(\{\mathbf{K}_l, \mathbf{V}_l\}, B_l, \mathbf{I}_l\big)$。这表明偏好引导的级联缓存管理在执行级联操作的同时，实现了与基础偏好优先自适应分配策略完全等效的KV缓存驱逐结果，并且在理论上兼容现有的KV驱逐技术。

抗注意力偏移的驱逐指标设计。当前顶级的驱逐指标，例如累积或平均注意力分数（【引用文献3，H2o: Heavy-hitter oracle for efficient generative inference of large language models+2024+NeurIPS】，【引用文献4，Scissorhands: Exploiting the persistence of importance hypothesis for llm kv cache compression at test time+2024+NeurIPS】，【引用文献5，On the efficacy of eviction policy for key-value constrained generative language model inference+2024+arXiv】），将Token的注意力曲线简化为单一数值以识别重要Token。然而，这种方法可能会忽略重要性波动的Token，因为它未能捕捉到注意力转移的动态变化。这可能导致过早地将Token从缓存中驱逐，从而影响模型未来检索相关信息的能力。为此，本文提出了一种能够容忍注意力偏移的鲁棒驱逐策略。该策略使用一个多维度的指标，同时考虑了持续重要性和注意力变异性。

驱逐指标的计算公式。对于第$l$层，计算驱逐指标$\mathbf{I}_l \in \mathbb{R}^S$，其中每个元素$\mathbf{I}_l[n]$表示Token $n$的重要性，计算方式如下：如果$n < S - S_w$，则$\mathbf{I}_l[n] = \mathrm{Mean}(\mathbf{A}_l[-S_w:, n]) + \gamma \cdot \mathrm{Var}(\mathbf{A}_l[-S_w:, n])$；否则，$\mathbf{I}_l[n] = \Omega$。公式中，$\mathrm{Mean}(\cdot)$和$\mathrm{Var}(\cdot)$分别衡量持续重要性和注意力变异性，$\gamma$用于调整它们的影响权重。受近期研究启发，本文赋予一个任意大的值$\Omega$以确保保留最近的$S_w$个Token。此外，系统使用一个池化层对$\mathbf{I}_l[:S - S_w]$进行聚类，以保持上下文连贯性并防止信息碎片化。最后执行驱逐操作$\mathrm{EVICT}\big(\{\mathbf{K}_l, \mathbf{V}_l\}, B_l, \mathbf{I}_l\big)$，保留与$\mathbf{I}_l$中前$B_l$个最高分相对应的KV对$\{\hat{\mathbf{K}}_l, \hat{\mathbf{V}}_l\}$：$\hat{\mathbf{K}}_l = \mathbf{K}_l[\mathbf{D}_l, :], \hat{\mathbf{V}}_l = \mathbf{V}_l[\mathbf{D}_l, :]$，其中$\mathbf{D}_l = \mathrm{TopK}(\mathbf{I}_l, B_l)$用于选择$\mathbf{I}_l$中前$B_l$个分数的索引$\mathbf{D}_l \in \mathbb{R}^{B_l}$。

图4：三个阶段中偏好引导的级联缓存管理算法图解。
图4：三个阶段中偏好引导的级联缓存管理算法图解。

实验环境

数据集名称、规模及用途：
1. LongBench：专注于长上下文理解，包含6个类别的16个数据集（单/多文档问答、摘要、少样本学习、合成任务和代码补全），输入长度范围从1,235到18,409个Token。用于评估不同内存预算下的长文本任务性能。
2. NeedleBench：测试复杂上下文中的检索和推理能力，包含三个子任务：单针检索（Single-Needle Retrieval）、多针检索（Multi-Needle Retrieval）和多针推理（Multi-Needle Reasoning），上下文长度为8K和32K。

模型架构关键参数：
实验涵盖5个主要的开源LLM（参数量7B-70B），支持4k-128k的Token上下文。
1. 多头注意力（Multi-head attention）架构：Llama2-Chat（7B/13B）、Gemma-Instruct（7B）。
2. 分组查询注意力（Grouped-query attention, GQA）架构：Llama3-Instruct（8B/70B）、Mistral-v0.3（7B）、Qwen2.5-Instruct（7B/32B）。

硬件与软件配置：
- 硬件：所有实验均在NVIDIA A100 80GB GPU上运行。
- 软件：代码基于PyTorch实现，并集成了FlashAttention-2以优化内存和吞吐量。基线方法包括统一分配策略（StreamingLLM, H2O, TOVA, SnapKV）和非统一分配策略（PyramidKV）。总内存预算设置范围从$64L$到$2048L$（$L$为层数）。

实验结果

LongBench数据集评估。实验对比了CAKE与基线方法在16个数据集上的性能（Fig 5, Table 1）。在Llama2-7B-Chat、Llama3.1-8B-Instruct和Mistral-7B-Instruct-v0.3模型上，CAKE在各种内存约束下均优于其他方法。特别是在低内存场景（如$B_{\mathrm{total}} = 128L$）下，CAKE通过动态分配内存，表现出显著的优势；在分配$512L$或更高预算时，CAKE的性能几乎与全缓存模型持平。结论表明，CAKE的自适应分配策略和鲁棒驱逐指标能有效适应不同模型架构和任务。

图5：不同缓存预算下LongBench 16个数据集的平均得分。
图5：不同缓存预算下LongBench 16个数据集的平均得分。

NeedleBench数据集评估。在检索和推理任务中（Fig 6），CAKE仅使用3.2%的缓存大小就保留了模型处理32K Token长上下文的能力，并在单针检索任务中甚至超越了全缓存基线。在更复杂的多针检索任务中，CAKE显著优于现有方法。结论表明，CAKE的设计平衡了长期重要性和短期相关性，避免了过早丢弃对后续复杂检索至关重要的信息。

图6：NeedleBench 32K上的性能比较。
图6：NeedleBench 32K上的性能比较。

内存消耗与吞吐量分析。基于Mistral-7B结合FlashAttention-2进行测试（Fig 7）。在峰值内存使用方面，CAKE在128K上下文长度下，相比全缓存实现减少了约48.63%的峰值内存。在解码延迟方面，随着输入长度增加，全缓存方法的延迟急剧上升，而CAKE保持稳定的解码速度。处理128K上下文时，CAKE的解码延迟比全缓存方法实现了超过10倍的加速。结论证实CAKE在大幅节省内存的同时极大地提升了推理效率。

图7：A100 80GB上的峰值内存使用和解码延迟。
图7：A100 80GB上的峰值内存使用和解码延迟。

与现有KV驱逐方法的兼容性。将CAKE的偏好优先自适应分配策略（P2A）应用于H2O和SnapKV（Fig 8, Table 2）。在Llama2-7B-Chat上，配备P2A策略的方法在几乎所有任务上均持续提升了性能。结论表明，CAKE的分配策略具有高度的通用性，可作为增强现有驱逐方法的通用框架。

图8：使用或不使用偏好优先缓存分配策略（缩写为“P2A”）的方法的性能。
图8：使用或不使用偏好优先缓存分配策略（缩写为“P2A”）的方法的性能。

消融实验。在分配策略方面（Table 2），对比了统一、金字塔、随机和CAKE的分配策略，CAKE稳定优于所有基线（29.29分）。在驱逐指标方面（Table 3），对比了仅均值、仅方差、均值乘方差和均值加方差（CAKE）。均值加方差实现了最佳性能（29.29分）。结论表明，结合长期重要性（均值）和注意力分布变化（方差）的加法组合能做出最明智的驱逐决策。

结论

本文提出了级联自适应KV缓存驱逐方法（CAKE），这是一种用于优化LLM中KV缓存驱逐的新颖方法。CAKE通过全局视角分析特定层的注意力模式，利用层偏好引导的级联缓存管理来动态分配缓存大小。此外，CAKE引入了一种新的驱逐指标，同时考虑了Token重要性的长期影响和时间变异性，从而实现更明智的Token选择。在LongBench和NeedleBench上的实验突显了CAKE在不同模型和内存约束下（尤其是低内存场景）的卓越性能。该方法在提升长上下文任务性能的同时，显著提高了推理效率。未来的工作可以探索将头级别（head-level）的注意力模式与层级变化相结合，并研究将CAKE与其他KV缓存优化技术（如缓存合并和剪枝）相集成，以进一步突破资源受限推理的边界。

附录补充细节

与量化方法的对比与兼容性。KV缓存驱逐和量化（降低存储值的位精度）是正交且互补的策略。本文将CAKE与最先进的非对称量化方法KIVI以及通道级量化方法KCVC进行了对比。在相同的压缩率下，CAKE在25%和12.5%的压缩率下均优于全缓存和两种量化方法。当将KIVI-INT4与CAKE结合使用时，在12.5%的压缩率下取得了32.51的得分，优于相同存储开销下的KIVI-INT2（32.17）。甚至在激进的6.25%压缩率下，组合方法仍保持了32.48的强劲性能（Fig 9）。这证实了驱逐与量化策略集成的有效性。

图9：在LongBench数据集上使用Llama2-7B-Chat比较CAKE、KV缓存量化方法（KIVI和KCVC）及其组合（KIVI与CAKE）。
图9：在LongBench数据集上使用Llama2-7B-Chat比较CAKE、KV缓存量化方法（KIVI和KCVC）及其组合（KIVI与CAKE）。

在其他模型架构上的扩展评估。为验证CAKE的泛化能力，在Qwen2.5-7B-Instruct和Gemma-7B-Instruct上进行了测试。在低内存（$128L$）和高内存（$1024L$）场景下，CAKE均持续优于基线方法。在高预算设置下，CAKE在Gemma上的表现甚至超过了全缓存基线（34.18 vs. 34.09）。这进一步验证了CAKE跨模型架构的适应性。

在更大规模模型上的扩展评估。实验扩展到了13B至70B参数的模型，包括Llama2-13B、Qwen2.5-32B和Llama3-70B。CAKE在所有模型规模上均优于基线方法。特别是在高内存设置（$1024L$）下，CAKE在Llama2-13B（29.98 vs. 29.95）和Llama3-70B（45.83 vs. 45.79）上实现了比全缓存更好的性能，表明该方法能良好地扩展到大型模型。

多针检索任务的深入分析。在NeedleBench 32K的多针检索任务中，当缓存预算极小（如$256L$）时，所有方法的性能都会急剧下降，因为严格的缓存限制无法支持多个“针”所需的信息量。然而，通过图表（Fig 10, Fig 11, Fig 12）可以看出，相比于其他方法，CAKE在缓解这一问题上展现出了最优越的能力，其结合持续重要性和变异性的策略在长上下文中具有更强的抗注意力偏移能力。

图10：在$B_{\mathrm{total}} = 1024L$下，使用不同KV驱逐方法在NeedleBench 32K多针检索任务（EN）上Mistral-7B-Instruct-v0.3的性能比较。
图10：在$B_{\mathrm{total}} = 1024L$下，使用不同KV驱逐方法在NeedleBench 32K多针检索任务（EN）上Mistral-7B-Instruct-v0.3的性能比较。
图11：在$B_{\mathrm{total}} = 512L$下，使用不同KV驱逐方法在NeedleBench 32K多针检索任务（EN）上Mistral-7B-Instruct-v0.3的性能比较。
图11：在$B_{\mathrm{total}} = 512L$下，使用不同KV驱逐方法在NeedleBench 32K多针检索任务（EN）上Mistral-7B-Instruct-v0.3的性能比较。
图12：在$B_{\mathrm{total}} = 512L$下，使用不同KV驱逐方法在NeedleBench 32K多针检索任务（EN）上Llama3.1-8B-Instruct的性能比较。
图12：在$B_{\mathrm{total}} = 512L$下，使用不同KV驱逐方法在NeedleBench 32K多针检索任务（EN）上Llama3.1-8B-Instruct的性能比较。

效率时间分解分析。在提示预填充和Token解码阶段的时间分解评估中，CAKE在显著降低解码阶段延迟的同时，实现了与SnapKV和PyramidKV相当的提示预填充时间。尽管CAKE使用了级联缓存管理，但由于层间KV缓存驱逐的并行执行特性，其动态更新过程可以合并为单个驱逐操作，因此总执行时间最终等同于执行$L$次驱逐操作的时间。

温度参数的影响。本文评估了调节空间分散度和时间偏移影响的温度参数$\tau_1$和$\tau_2$。在$128L$预算下，CAKE在不同配置下均持续优于SnapKV基线。通过低成本调整$\tau_1$和$\tau_2$，可以更好地捕捉特定模型的注意力模式（Fig 13），无需昂贵的重新训练，这在工业部署中具有极高的实用价值。

图13：不同$\tau_1$和$\tau_2$配置对Mistral-7B-Instruct-v0.3和Llama3.1-8B-Instruct模型性能比较的影响。图表显示了改变$\tau_1$和$\tau_2$对模型性能的影响，并标出了基线（SnapKV）和最佳得分。
图13：不同$\tau_1$和$\tau_2$配置对Mistral-7B-Instruct-v0.3和Llama3.1-8B-Instruct模型性能比较的影响。图表显示了改变$\tau_1$和$\tau_2$对模型性能的影响，并标出了基线（SnapKV）和最佳得分。

注意力动态的详细可视化。通过聚合各注意力头的权重，本文可视化了不同层、模型和上下文的注意力模式（Fig 14, Fig 15）。
- 层间差异：无论模型或输入上下文如何，不同层在注意力分散度和偏移上呈现显著差异，证明了按需分配缓存资源的必要性。
- 模型间差异：Mistral通常在其层叠的开头和结尾表现出较高的注意力分散度；而Llama3则主要在早期层显示出高分散度。这强调了缓存策略需要考虑模型特定架构。
- 上下文间差异：不同任务类型甚至同一任务的不同数据集之间，注意力偏移模式存在显著差异，导致对KV缓存的层偏好不同。CAKE通过全局视角综合考虑这些动态变化，能够有效适应不同场景。

图14：Llama3.1-8B-Instruct在不同任务数据集上的注意力权重可视化结果。
图14：Llama3.1-8B-Instruct在不同任务数据集上的注意力权重可视化结果。
图15：Mistral-7B-Instruct-v0.3在单文档问答任务数据集上的注意力权重可视化结果。
图15：Mistral-7B-Instruct-v0.3在单文档问答任务数据集上的注意力权重可视化结果。

方法细节参考文献汇总

在方法细节章节中，CAKE的设计借鉴并引用了以下关键文献：
- 【引用文献1，Snapkv: Llm knows what you are looking for before generation+2024+arXiv】：在设计偏好优先的自适应分配策略时，引用该文献以支持将计算焦点集中在最近查询窗口$S_w$的注意力模式上，证明其能有效捕捉解码模式；并在设计驱逐指标时，借鉴其保留最近$W$个Token的思想，以及在实验设定中采用其$S_w=32$和池化层聚类的设定。
- 【引用文献2，Pyramidinfer: Pyramid kv cache compression for high-throughput llm inference+2024+arXiv】：与引用文献1相同，用于支撑分析最近查询窗口注意力权重的合理性，并在驱逐指标中赋予最近Token极大值$\Omega$以确保其不被驱逐。
- 【引用文献3，H2o: Heavy-hitter oracle for efficient generative inference of large language models+2024+NeurIPS】：在分析现有驱逐指标局限性时，指出此类基于累积注意力分数的方法将注意力曲线简化为单一数值，从而忽略了注意力转移的动态变化。
- 【引用文献4，Scissorhands: Exploiting the persistence of importance hypothesis for llm kv cache compression at test time+2024+NeurIPS】：同样作为现有顶级驱逐指标的代表被引用，用以说明现有方法在捕捉重要性波动Token方面的不足。
- 【引用文献5，On the efficacy of eviction policy for key-value constrained generative language model inference+2024+arXiv】：作为使用平均注意力分数进行驱逐的代表性研究被引用，指出了单一平均值指标在处理时间变异性时的局限性，从而引出CAKE结合均值与方差的鲁棒指标设计。