Ada-KV: Optimizing KV Cache Eviction by Adaptive Budget Allocation for Efficient LLM Inference

Conference: NeurIPS'25

Github: https://github.com/FFY0/AdaKV

My Thoughts

这篇工作的idea在各个head之间adaptive 分配 budget, 直觉上就能知道这种设计是有效的，实验也验证了这一点。

比较惊喜的是作者实现了With efficient CUDA kernel implementations来解决variable-sized cache elements across attention heads，从而来真正实现了计算加速，以及其他KV Cache工作在也实现了对AdaKV集成，再次证明了其作为通用增强模块的价值。

1. Motivation

大型语言模型（LLM）在各个领域表现出色，但由于长序列推理所需的不断增长的键值（KV）cache，面临着效率挑战。LLM 的广泛应用推动了其处理扩展序列能力的发展。例如，GPT 支持长达 128K 的序列，Claude3 支持 200K，Gemini-Pro-1.5 甚至支持高达 2M 个 token。然而，这种 token 长度的增长带来了显著的挑战，尤其是在推理过程中cache大小的急剧膨胀。对于一个 8B 的 LLM，处理一个 2M token 的序列可能需要高达 256GB 的cache，这严重影响了 GPU 内存效率和计算运行时效率。

现有的 KV cache驱逐方法通常在所有注意力head上均匀分配压缩预算，忽略了每个head独特的注意力模式。如 Figure 1a 所示，不同head的注意力集中度（concentration）差异巨大：一些head（sparse heads）的注意力高度集中在少数几个 token 上，而另一些head（dispersed heads）的注意力则分布得更广。这种均匀分配导致了效率低下——要么在稀疏集中的head上浪费cache预算，要么在分散分布的head上造成显著的驱逐损失。

2. Relative Work

2.1 Cache Eviction Methods

cache驱逐方法主要分为两类：滑动窗口驱逐和 Top-k 驱逐。

滑动窗口方法（如 StreamingLLM）简单地保留初始cache元素和滑动窗口内的元素，但这种无差别的驱逐会显著降低生成质量。
Top-k 驱逐方法（如 H2O, SnapKV, Pyramid）基于注意力权重识别并保留 k 个关键cache元素。然而，现有 Top-k 方法通常在不同head上均匀分配总预算。Ada-KV 通过自适应预算分配来增强这些方法。

2.2 Sparse Attention Methods

稀疏注意力方法与 KV cache驱逐在根本上不同：前者保留所有cache，但在计算时只选择性地使用关键子集，因此不减少内存占用。而 KV cache驱逐直接移除非关键条目，从而减小内存占用。这两种技术是正orthogonal的，未来可以结合使用。

3. Contribution

自适应预算分配：这篇工作指出了当前 KV cache驱逐方法的一个关键限制：均匀预算分配忽略了各个head的独特注意力模式。为此，这篇工作提出了 Ada-KV，这是首个head级别的自适应预算分配策略，通过提高各个head的预算利用率，带来更高效的cache驱逐。
理论分析：这篇工作为cache驱逐建立了一个理论框架，定义了驱逐损失并推导了其上界。该框架解释了先前方法的优化目标，指导了 Ada-KV 的设计，实现了有原则的自适应预算分配。
具体实现：高效的 CUDA 内核实现，Ada-KV 具有即插即用的兼容性，可以集成到现有方法中。这篇工作在 Ruler 和 LongBench 的 29 个数据集上，在问答感知（question-aware）和问答不可知（question-agnostic）两种场景下都证明了其显著的性能提升。

4. Method

Ada-KV 的核心思想是：将所有head的注意力权重“拉平”成一个大的列表，然后在这个全局列表中选出权重最大的 B 个元素。每个head i 分配到的预算 $B_i^*$ 就等于其权重在全局 Top-B 中出现的次数。这确保了总预算被分配给了全局最重要的cache元素，从而最小化了理论上的驱逐损失上界。

4.1 理论基础：驱逐损失的上界

论文首先为 KV Cache 驱逐建立了一个理论框架。

符号定义:
- 考虑一个包含 h 个注意力head的多head自注意力层。
- 对于第 i 个head，其注意力权重向量为 $A_i \in \mathbb{R}^{1 \times n}$，其中 n 是当前上下文长度。
- 驱逐决策由一个指示变量 $I_i \in \mathbb{R}^{1 \times n}$ 表示，其中 $I_{ij} = 1$ 表示保留第 j 个 KV cache元素，否则为 0。
- 第 i 个head的预算为 $B_i$，满足 $\sum_{j=1}^{n} I_{ij} = B_i$。整个层的总预算为 $B = \sum_{i=1}^{h} B_i$。
- 驱逐前后的自注意力输出分别为 y 和 $\hat{y}$。
驱逐损失 (Eviction Loss): 定义为驱逐前后输出的 L1 距离： $$ L1 \text{ Eviction Loss} = |y - \hat{y}|_1 $$
上界推导: 论文证明了该损失存在一个上界 $\epsilon$： $$ L1 \text{ Eviction Loss} \leq \epsilon = 2hC - 2C \sum_{i=1}^{h} \sum_{j=1}^{n} I_{ij} A_{ji} $$ 其中 $C = \max{|V_i W_i^O|_{\infty}}$ 是一个常数。
对现有方法的解释: 这个上界揭示了现有 Top-k 驱逐方法的优化目标。为了最小化 $\epsilon$，需要最大化 $\sum_{i} \sum_{j} I_{ij} A_{ji}$，即在每个head i 内部，选择注意力权重 $A_{ji}$ 最大的 Bi 个元素保留。这正是 Top-k 驱逐策略。

4.2 自适应预算分配策略 (Ada-KV)

Algorithm 1 描述了 Ada-KV 的核心分配逻辑。

Algorithm 1: Ada-KV: 自适应预算分配

输入: 总预算 B; 每个head i 的注意力权重 {A_i}
输出: 分配的预算 {B_i^*}

1: 将所有head的注意力权重拼接成一个大向量 A = Cat({A_i})
2: 从 A 中选出最大的 B 个权重: Top-k(A, k = B)
3: 统计每个head i 在 Top-B 中被选中的次数: {f_i}
4: 设置分配的预算为 {B_i^* = f_i}
返回 分配的预算 {B_i^*}

理论优势: Theorem 3.3 证明了这种分配策略能够获得所有可能分配方案中最小的上界 $\epsilon^{**}$。

4.3 与现有方法的集成 (Ada-SnapKV / Ada-Pyramid)

Ada-KV 被设计为“即插即用”，可以无缝集成到现有的 Top-k 驱逐方法中。Algorithm 2 展示了如何将其集成到 SnapKV/Pyramid 中。

Algorithm 2: Ada-SnapKV/Ada-Pyramid 在单层中的流程

输入: 总预算 B, 观察窗口内的 token X_win ∈ R^{win*d}, 
      观察窗口内的cache {K_win, V_win}, 观察窗口外的cache {K_i, V_i}
输出: 保留的cache {K̂_i, V̂_i}

1: for i ← 1 to h do
2:   Q_win_i = X_win W_i^Q
3:   Ā_i = softmax(Q_win_i K_i^T)
4:   Ā_i = Ā_i.maxpooling(dim=1).mean(dim=0)  // 聚合观察窗口内的注意力
5: end for
6: B = B - winsize × h  // 减去观察窗口内强制保留的cache开销
7: 使用 Algorithm 1(B, {Ā_i}) 推导预算分配 {B_i^*}
8: 安全防护: {B_i^*} = α × {B_i^*} + (1 − α) × (B/h) // α 默认为 0.2
9: 根据 {B_i^*} 确定 Top-k 驱逐决策 {I_i^*}
10: 根据 {I_i^*} 从 {K_i, V_i} 中选择 {K̂_i, V̂_i}
11: {K̂_i, V̂_i} = Cat({K̂_i, V̂_i}, {K_win, V_win}) // 与观察窗口cache拼接
返回 保留的cache {K̂_i, V̂_i}

详细解释：

输入
- B：总预算（总的 KV cache 可保留的容量）。
- X_win ∈ R^{win × d}：观察窗口内的 token embedding，窗口大小为 win，embedding 维度为 d。
- K_win, V_win：观察窗口内的 key/value cache。
- K_i, V_i：观察窗口外的 key/value cache（通常是历史 token 对应的 cache）。
输出
- K̂_i, V̂_i：最终保留的 cache。

步骤解析

1–5 行：计算观察窗口内的注意力聚合

1: for i ← 1 to h do
2:   Q_win_i = X_win W_i^Q
3:   Ā_i = softmax(Q_win_i K_i^T)
4:   Ā_i = Ā_i.maxpooling(dim=1).mean(dim=0)
5: end for

目的：评估每个头 i 对观察窗口外的 KV cache 的重要性。
解释：
1. Q_win_i = X_win W_i^Q 计算第 i 个头的查询矩阵。
2. Ā_i = softmax(Q_win_i K_i^T) 计算观察窗口内 token 对每个历史 KV 的注意力分数。
3. Ā_i = Ā_i.maxpooling(dim=1).mean(dim=0) 聚合注意力信息：
  - maxpooling(dim=1)：取每个观察窗口 token 对历史 KV 中的最大注意力。
  - mean(dim=0)：对观察窗口 token 求平均，得到每个历史 KV 的重要性分数。
结果：每个 head 得到一个向量 Ā_i，表示该 head 对各个历史 KV 的注意力重要性。

Ā_i = Ā_i.maxpooling(dim=1).mean(dim=0)

Ā_i 是第 i 个注意力头对 历史 KV cache 的注意力分数矩阵。
形状大概是 (win, num_cache)：
- win：观察窗口内的 token 数量
- num_cache：历史 KV cache 的数量

也就是说，每一行表示一个观察窗口 token 对所有历史 KV 的注意力分布。

1️⃣ maxpooling(dim=1)

对 每一行（每个观察窗口 token）进行 最大池化：
- 取这个 token 对历史 KV 的 最大注意力值。
结果：
- 原本 (win, num_cache) → (win,) 向量
- 这个向量表示 每个 token 最关注的历史 KV 的重要性。

直观理解：我们只关心每个 token 在历史 KV 中的 “最重要的记忆”，而不是平均分散的注意力。

2️⃣ mean(dim=0)

对 所有观察窗口 token 求平均：
- 把每个 token 的最大注意力值加起来再除以 token 数量。
结果：
- 一个标量或向量（视具体实现），表示 该注意力头整体对某些历史 KV 的重要性评分。
作用：
- 聚合观察窗口内所有 token 的信息，得到一个 全局的历史 KV 重要性向量。

6 行：扣除观察窗口 cache 的开销

6: B = B - winsize × h

解释：观察窗口内的 cache K_win, V_win 是 必须保留的，所以总预算要减去它们占用的空间。

7–8 行：预算分配

7: 使用 Algorithm 1(B, {Ā_i}) 推导预算分配 {B_i^*}
8: {B_i^*} = α × {B_i^*} + (1 − α) × (B/h)

步骤 7：
- 调用另一个算法（Algorithm 1）根据各头的重要性向量 Ā_i 来分配剩余预算 B。
- 得到每个头分配到的预算 B_i^*（可保留的 KV 数量）。
步骤 8：安全防护：
- 用一个平滑系数 α 做线性插值，避免某些头的预算过低。
- 公式：
$$ B_i^* = \alpha \cdot B_i^* + (1 - \alpha) \cdot \frac{B}{h} $$
- 默认 α=0.2，即 20% 依赖算法分配，80% 保证每个头有基础份额。

9–10 行：选择保留的 KV

9: 根据 {B_i^*} 确定 Top-k 驱逐决策 {I_i^*}
10: 根据 {I_i^*} 从 {K_i, V_i} 中选择 {K̂_i, V̂_i}

步骤 9：
- 对每个头，选择 Top-k KV，其中 k = B_i^*。
- I_i^* 是被保留的 KV 的索引。
步骤 10：
- 根据索引 I_i^* 从历史 cache 中提取对应的 K̂_i, V̂_i。

11 行：与观察窗口 cache 拼接

11: {K̂_i, V̂_i} = Cat({K̂_i, V̂_i}, {K_win, V_win})

4.4 高效实现

自适应预算分配会导致不同head的cache长度可变，给高效计算带来挑战。Ada-KV 通过以下方式解决：

扁平化cache布局 (Flattened Cache Storage Layout): 将一个层内所有head的 KV cache拼接成一个大的张量。
可变长度 FlashAttention: 利用 flash_attn_varlen_func 来高效处理这种可变长度的注意力计算。
自定义 CUDA 内核: 实现了高效的cache更新操作，以支持在扁平化布局下动态插入新的 KV 对。如 Figure 2 所示，更新过程分为 malloc、copy old value 和 insert new value 三个阶段。
GQA 兼容性: 针对现代 LLM（如 Llama, Mistral）广泛使用的 Grouped Query Attention (GQA)，Ada-KV 通过在每个 GQA 组内使用平均注意力权重作为选择标准，避免了冗余计算。

5. Evaluation

5.1 实验设置

基座模型: Llama-3.1-8B-Instruct 和 Mistral-7B-Instruct-v0.2（均使用 GQA）。
数据集: Ruler (13个任务) 和 LongBench (16个数据集)。
评估场景:
- Question-aware: 压缩时已知问题（标准场景）。
- Question-agnostic: 更具挑战性的场景，压缩时不知道问题，更能反映真实应用（如提示cache）。
基线方法: SOTA 的 SnapKV 和 Pyramid，以及滑动窗口方法 StreamingLLM。
参数: 观察窗口大小为32，安全系数 $\alpha=0.2$。

5.2 Ruler 基准测试结果

总体表现: 如 Figure 3 所示，Ada-SnapKV 和 Ada-Pyramid 始终优于 原始方法，尤其在 question-agnostic 场景下优势更显著。例如，在 Llama-3.1-8B 上，20% cache预算时，Ada-SnapKV 将 SnapKV 的得分从 44.02 大幅提升至 53.29。
子任务分析: Figure 4 展示了 question-agnostic 场景下的详细结果。Ada-KV 在几乎所有任务上都带来了显著提升。特别在困难的 NIAH 任务（如 S-NIAH-3, MK-NIAH-2）上，Ada-SnapKV 在 80% 预算下几乎实现了无损性能（得分 97.6 和 99.6），而原始 SnapKV 则有明显下降。

5.3 LongBench 基准测试结果

固定预算评估: Figure 5 显示，在 question-aware 场景下，Ada-变体持续优于基线。在 question-agnostic 场景下，所有方法性能都有所下降，但 Ada-变体依然保持领先。
按比例预算评估: Table 1 和 Table 2（Llama-70B）展示了按任务领域的详细结果。Ada-SnapKV 在绝大多数领域（18个领域中的15个）都有效缓解了压缩带来的性能损失，证明了其通用性和鲁棒性。

5.4 计算效率分析

Ada-KV 的核心目标是在不牺牲效率的前提下提升质量。Figure 6 的结果表明，得益于高效的 CUDA 实现和可变长度 FlashAttention，Ada-SnapKV 在峰值内存占用和解码延迟方面与原始的 SnapKV 基本持平，并且都远优于使用完整cache的情况。

5.5 广泛适用性

由于其即插即用的特性，Ada-KV 已被后续多项工作（如 CriticalKV, DefensiveKV）采用。Table 3 显示，将 Ada-KV 应用于这些新方法后，性能得到了进一步提升，再次证明了其作为通用增强模块的价值。

6. Conclusion

本研究重新审视了用于高效 LLM 推理的cache驱逐策略，并发现了一个被忽视的关键因素：跨注意力head的自适应预算分配。在驱逐损失上界的理论分析指导下，这篇工作提出了 Ada-KV，这是首个用于优化 KV cache驱逐方法的自适应预算分配策略。这篇工作通过将其集成到两个现有的 SOTA 方法中，引入了 Ada-SnapKV 和 Ada-Pyramid。除了常见的问答感知压缩场景，这篇工作还在更具挑战性且较少被探索的问答不可知压缩场景中评估了这些方法。使用 Ruler 和 LongBench 基准，这篇工作证明了自适应预算分配在这两种设置下的有效性。这篇工作的结果不仅暴露了当前cache驱逐策略的局限性，也突显了自适应分配在增强cache驱逐方面的潜力。

My Thoughts#

1. Motivation#

2. Relative Work#

2.1 Cache Eviction Methods#

2.2 Sparse Attention Methods#

3. Contribution#

4. Method#

4.1 理论基础：驱逐损失的上界#

4.2 自适应预算分配策略 (Ada-KV)#

4.3 与现有方法的集成 (Ada-SnapKV / Ada-Pyramid)#

1–5 行：计算观察窗口内的注意力聚合#

1️⃣ maxpooling(dim=1)#

2️⃣ mean(dim=0)#

6 行：扣除观察窗口 cache 的开销#

7–8 行：预算分配#

9–10 行：选择保留的 KV#

11 行：与观察窗口 cache 拼接#

4.4 高效实现#

5. Evaluation#

5.1 实验设置#

5.2 Ruler 基准测试结果#

5.3 LongBench 基准测试结果#

5.4 计算效率分析#

5.5 广泛适用性#

6. Conclusion#