2025-11-16T09:58:12.370377

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

Feng, Lv, Cao et al.
Large Language Models have excelled in various domains but face efficiency challenges due to the growing Key-Value (KV) cache required for long-sequence inference. Recent efforts aim to reduce KV cache size by evicting vast non-critical cache elements during runtime while preserving generation quality. However, these methods typically allocate compression budgets uniformly across all attention heads, ignoring the unique attention patterns of each head. In this paper, we establish a theoretical loss upper bound between pre- and post-eviction attention output, explaining the optimization target of prior cache eviction methods, while guiding the optimization of adaptive budget allocation. Base on this, we propose {\it Ada-KV}, the first head-wise adaptive budget allocation strategy. It offers plug-and-play benefits, enabling seamless integration with prior cache eviction methods. Extensive evaluations on 13 datasets from Ruler and 16 datasets from LongBench, all conducted under both question-aware and question-agnostic scenarios, demonstrate substantial quality improvements over existing methods. Our code is available at https://github.com/FFY0/AdaKV.
academic

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

基本信息

  • 论文ID: 2407.11550
  • 标题: Ada-KV: Optimizing KV Cache Eviction by Adaptive Budget Allocation for Efficient LLM Inference
  • 作者: Yuan Feng, Junlin Lv, Yukun Cao, Xike Xie, S. Kevin Zhou
  • 分类: cs.CL cs.AI
  • 发表时间/会议: 39th Conference on Neural Information Processing Systems (NeurIPS 2025)
  • 论文链接: https://arxiv.org/abs/2407.11550

摘要

大型语言模型(LLMs)在各个领域表现出色，但由于长序列推理中不断增长的Key-Value(KV)缓存需求，面临着效率挑战。最近的研究通过在运行时驱逐大量非关键缓存元素来减少KV缓存大小，同时保持生成质量。然而，这些方法通常在所有注意力头之间均匀分配压缩预算，忽略了每个头的独特注意力模式。本文建立了驱逐前后注意力输出之间的理论损失上界，解释了先前缓存驱逐方法的优化目标，同时指导自适应预算分配的优化。基于此，作者提出了Ada-KV，第一个头级自适应预算分配策略。该方法具有即插即用的优势，能够与现有的缓存驱逐方法无缝集成。

研究背景与动机

问题描述

随着大型语言模型处理序列长度的不断增长(如GPT支持128K、Claude3支持200K、Gemini-Pro-1.5支持2M tokens)，KV缓存的内存需求呈指数级增长。对于8B参数的LLM，处理单个2M token序列可能需要高达256GB的缓存，严重影响GPU内存效率和计算运行时效率。

现有方法的局限性

现有的缓存驱逐方法主要分为两类：

  1. 滑动窗口驱逐方法：简单保留初始和最近的缓存元素，但会显著降低生成质量
  2. Top-k驱逐方法：基于注意力权重选择关键缓存元素，但在所有注意力头之间均匀分配预算

关键问题是现有方法忽略了不同注意力头的独特特征：一些头具有稀疏的注意力集中模式，而另一些头的注意力分布更加分散。

研究动机

通过分析Llama-3.1-8B-Instruct模型，作者发现大多数注意力头仅用很小的缓存比例(如top 5%)就能保留几乎所有注意力权重，而分散的头则需要更大的缓存比例。这种不均匀的注意力集中模式为自适应预算分配提供了理论基础。

核心贡献

  1. 自适应预算分配策略：提出了首个头级自适应预算分配策略Ada-KV，能够根据各个注意力头的独特注意力模式动态调整预算分配
  2. 理论框架建立：建立了缓存驱逐的理论框架，定义了驱逐损失并推导其上界，解释了现有方法的优化目标并指导Ada-KV的设计
  3. 即插即用兼容性：Ada-KV具有即插即用特性，能够无缝集成到现有的缓存驱逐方法中，并通过高效的CUDA内核实现保持计算效率
  4. 全面实验验证：在Ruler和LongBench的29个数据集上进行了全面评估，在问题感知和问题无关两种场景下都显示出显著改进

方法详解

任务定义

给定一个多头自注意力层，在预算约束下选择要保留的KV缓存元素，使得驱逐后的注意力输出与原始输出之间的损失最小化。

理论基础

L1驱逐损失定义

作者将驱逐损失量化为驱逐前后自注意力机制输出之间的L1距离：

L1 Eviction Loss=∣∣y−y^∣∣1\text{L1 Eviction Loss} = ||y - \hat{y}||_1

其中yy和y^\hat{y}分别是驱逐前后的注意力输出。

损失上界推导

定理3.1：L1驱逐损失可以被ϵ\epsilon上界限制：

L1 Eviction Loss≤ϵ=2hC−2C∑i∈[1,h]∑j∈[1,n]IijAij\text{L1 Eviction Loss} \leq \epsilon = 2hC - 2C\sum_{i \in [1,h]}\sum_{j \in [1,n]} I_i^j A_i^j

其中C=max⁡{∥ViWiO∥∞}C = \max\{\|V_iW_i^O\|_\infty\}是常数，IijI_i^j是驱逐决策指示变量，AijA_i^j是注意力权重。

定理3.2：Top-k缓存驱逐方法在给定预算分配下能够最小化损失上界：

ϵ∗=2hC−2C∑i∈[1,h]∑Aij∈Top-k(Ai,k=Bi)Aij\epsilon^* = 2hC - 2C\sum_{i \in [1,h]}\sum_{A_i^j \in \text{Top-k}(A_i, k=B_i)} A_i^j

Ada-KV算法

算法1：自适应预算分配

输入：总预算B，各头注意力权重{A_i}
输出：分配预算{B_i^*}
1. 连接所有头的注意力权重：A = Cat({A_i})
2. 从A中选择top B个权重：Top-k(A, k=B)
3. 统计每个头被选中的权重数量：{f_i}
4. 设置分配预算：{B_i^* = f_i}

理论优势

定理3.3：自适应预算分配能够实现最小的损失上界：

ϵ∗∗=min⁡{Bi}ϵ∗\epsilon^{**} = \min_{\{B_i\}} \epsilon^*

与现有方法集成

作者展示了Ada-KV与两个SOTA方法的集成：

Ada-SnapKV和Ada-Pyramid

通过算法2，Ada-KV可以无缝集成到SnapKV和Pyramid中：

  1. 计算观察窗口内的注意力权重
  2. 使用Ada-KV算法分配预算
  3. 应用安全保护参数α = 0.2防止过度稀疏分配
  4. 执行Top-k驱逐决策

技术创新点

  1. 全局优化视角：将头级预算分配视为全局优化问题，而非局部优化
  2. 理论指导设计：基于严格的理论分析指导算法设计
  3. 计算效率保证：通过变长FlashAttention和扁平化缓存布局保持计算效率
  4. GQA兼容性：支持Group Query Attention，实现额外的缓存压缩

实验设置

数据集

  • Ruler基准：13个长序列任务，主要是Needle-in-a-Haystack测试的变体，评估16K长度
  • LongBench基准：16个数据集，涵盖单文档QA、多文档QA、摘要、少样本学习、合成任务和代码生成

基础模型

  • Llama-3.1-8B-Instruct
  • Mistral-7B-instruct-v0.2

评价指标

根据任务类型使用相应指标：F1分数(QA任务)、Rouge-L(摘要任务)、准确率(分类任务)、编辑相似度(代码任务)

对比方法

  • 基线方法：SnapKV、Pyramid、StreamingLLM
  • 增强版本：Ada-SnapKV、Ada-Pyramid

实验场景

  • 问题感知压缩：问题已知的标准场景
  • 问题无关压缩：更具挑战性的实际应用场景

实验结果

主要结果

Ruler基准测试

在问题无关场景下，使用Llama-3.1-8B-Instruct：

  • 80%缓存预算：Ada-SnapKV将SnapKV的分数从87.59提升到92.67
  • 20%缓存预算：Ada-SnapKV将SnapKV的分数从44.02提升到53.29

LongBench基准测试

在问题无关场景下：

  • Ada-SnapKV和Ada-Pyramid在所有固定预算设置下都持续改善生成质量
  • 在2048预算下接近无损性能

子任务分析

在困难的Needle-in-a-Haystack任务中：

  • S-NIAH-3任务(80%预算)：Ada-SnapKV将SnapKV从62.4提升到97.6
  • MK-NIAH-2任务(80%预算)：Ada-SnapKV将SnapKV从85.2提升到99.6

计算效率

Ada-SnapKV在固定1024预算下：

  • 峰值内存使用量与原始SnapKV相当
  • 解码延迟与原始SnapKV相当
  • 两者都显著优于全缓存情况

广泛应用验证

Ada-KV策略已被多个后续工作采用：

  • CriticalKV + Ada-KV：在20%缓存下从42.99提升到43.77
  • DefensiveKV + Ada-KV：在20%缓存下从43.78提升到46.68

相关工作

缓存驱逐方法

  • 滑动窗口方法：StreamingLLM等，简单但质量损失大
  • Top-k方法：H2O、SnapKV、Pyramid等，基于注意力权重选择关键元素

稀疏注意力方法

与缓存驱逐在概念上相关但方法不同：

  • 缓存驱逐：只保留KV缓存子集
  • 稀疏注意力：保留所有条目但选择性使用

其他相关技术

  • KV缓存量化：降低单个元素精度
  • 推测解码：使用减少缓存的模型生成草稿
  • 分页注意力：高效内存管理

结论与讨论

主要结论

  1. Ada-KV首次提出了头级自适应预算分配策略，显著改善了现有缓存驱逐方法的性能
  2. 理论分析为缓存驱逐建立了严格的框架，指导了算法设计
  3. 问题无关压缩场景揭示了现有方法的局限性，应受到更多关注

局限性

  1. 当前的头级分配仍限于单层内，未扩展到跨层分配
  2. 安全保护参数α需要在不同预算下权衡性能
  3. 理论分析基于L1距离，可能不完全反映实际生成质量

未来方向

  1. 扩展头级分配机制到跨层场景
  2. 开发相应的跨层理论分析
  3. 结合训练时的头重要性分析
  4. 与其他优化技术(如量化、稀疏注意力)的联合优化

深度评价

优点

  1. 理论贡献扎实：建立了完整的理论框架，从损失上界推导到算法设计逻辑清晰
  2. 方法简单有效：算法简洁易懂，即插即用特性使其易于采用
  3. 实验全面充分：在29个数据集上的全面评估，包含了被忽视的问题无关场景
  4. 实用价值高：已被多个后续工作采用，证明了方法的价值和影响力

不足

  1. 理论与实践的间隙：虽然理论上最小化了损失上界，但不能保证实际损失最小
  2. 超参数敏感性：安全保护参数α的选择需要经验调优
  3. 扩展性限制：目前仅考虑了单层内的预算分配
  4. 评估局限性：主要在中等规模模型上评估，大规模模型的效果有待验证

影响力

  1. 学术贡献：为KV缓存优化领域提供了新的研究方向
  2. 实用价值：即插即用特性使其易于在实际系统中部署
  3. 可复现性：提供了开源代码和详细的实现细节
  4. 启发性：为后续研究提供了理论框架和方法论指导

适用场景

  1. 长序列推理：特别适合需要处理长上下文的应用
  2. 资源受限环境：在GPU内存有限的情况下优化推理效率
  3. 实时系统：需要平衡质量和效率的在线服务
  4. 多轮对话：问题无关压缩场景特别适合对话系统

参考文献

论文引用了64篇相关文献，主要包括：

  • 大型语言模型基础工作(GPT-4, Claude, Gemini等)
  • KV缓存优化方法(H2O, SnapKV, Pyramid等)
  • 注意力机制优化(FlashAttention, 稀疏注意力等)
  • 长序列处理基准(Ruler, LongBench等)

总体评价：这是一篇高质量的研究论文，在理论贡献和实用价值之间取得了良好的平衡。Ada-KV方法简单而有效，理论分析严谨，实验验证充分。论文不仅解决了现有方法的重要局限性，还为未来研究提供了有价值的框架和方向。