CAKE: Cascading and Adaptive KV Cache Eviction with Layer Preferences

📄 arXiv: 2503.12491v2 📥 PDF

作者: Ziran Qin, Yuchen Cao, Mingbao Lin, Wen Hu, Shixuan Fan, Ke Cheng, Weiyao Lin, Jianguo Li

分类: cs.CL

发布日期: 2025-03-16 (更新: 2025-12-24)

备注: Accepted by ICLR 2025

🔗 代码/项目: GITHUB

💡 一句话要点

CAKE：基于层偏好的级联自适应KV缓存淘汰方法，提升长序列LLM推理效率

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: KV缓存 缓存淘汰 长序列模型 推理优化 注意力机制

📋 核心要点

  1. 现有KV缓存淘汰方法无法根据不同层的注意力模式合理分配资源，导致性能瓶颈。
  2. CAKE将KV缓存淘汰视为“切蛋糕问题”，根据层偏好自适应分配缓存，并级联管理内存。
  3. 实验表明，CAKE仅用3.2%的KV缓存即可维持模型性能，并在低内存下显著优于现有方法。

📝 摘要（中文）

大型语言模型（LLM）在处理长序列方面表现出色，从而推动了对键值（KV）缓存的需求。虽然最近在淘汰KV缓存方面所做的努力已经减轻了推理负担，但它们通常无法在具有不同注意力模式的层之间合理地分配资源。本文介绍了一种新的方法，即级联自适应KV缓存淘汰（CAKE），它将KV缓存淘汰视为一个“切蛋糕问题”。CAKE通过考虑空间和时间维度上的注意力动态来评估特定层的偏好，相应地为各层分配合理的缓存大小，并以级联方式管理内存约束。这种方法能够全局地看待缓存分配，自适应地在不同的注意力机制之间分配资源，同时保持内存预算。CAKE还采用了一种新的淘汰指标，该指标考虑了token随时间推移的重要性变化，解决了现有方法忽略时间动态的局限性。在LongBench和NeedleBench上的综合实验表明，CAKE仅使用3.2%的KV缓存即可维持模型性能，并且在各种模型和内存约束下，特别是在低内存设置中，始终优于当前的基线。此外，当使用FlashAttention-2处理128K token的上下文时，CAKE在解码延迟方面实现了超过10倍的加速。我们的代码可在https://github.com/antgroup/cakekv获取。

🔬 方法详解

问题定义：论文旨在解决大型语言模型（LLM）推理过程中，KV缓存的有效管理问题。现有方法，如简单的LRU或FIFO，无法充分利用不同层之间的注意力模式差异，导致缓存资源分配不合理，影响长序列处理的效率和性能。尤其是在低内存环境下，如何最大化利用有限的缓存资源成为一个关键挑战。

核心思路：CAKE的核心思路是将KV缓存淘汰问题转化为一个“切蛋糕”问题，即如何根据不同层的需求（层偏好）合理分配有限的缓存资源。通过分析每一层在空间和时间维度上的注意力动态，评估其对KV缓存的需求程度，并据此分配缓存大小。这种自适应分配策略旨在最大化整体性能，同时满足内存约束。

技术框架：CAKE的技术框架主要包含以下几个阶段：1) 层偏好评估：通过分析每一层在空间和时间维度上的注意力动态，计算层偏好指标。2) 缓存大小分配：根据层偏好指标，为每一层分配相应的缓存大小，确保资源分配的合理性。3) 级联淘汰：当内存不足时，采用级联淘汰策略，优先淘汰对整体性能影响较小的层的缓存。4) 时间敏感的淘汰指标：设计新的淘汰指标，考虑token随时间推移的重要性变化，避免淘汰重要token。

关键创新：CAKE的关键创新在于：1) 层偏好感知的缓存分配：不同于以往的全局统一淘汰策略，CAKE能够根据每一层的注意力模式自适应地分配缓存资源。2) 级联淘汰机制：通过级联淘汰，能够在内存受限的情况下，最大程度地保留对整体性能至关重要的层的缓存。3) 时间敏感的淘汰指标：考虑了token的重要性随时间变化的特性，避免了对重要token的误淘汰。

关键设计：CAKE的关键设计包括：1) 层偏好指标的计算方法：具体如何量化空间和时间维度上的注意力动态，并将其转化为层偏好指标。2) 缓存大小分配策略：如何根据层偏好指标，确定每一层的缓存大小，并保证整体内存约束。3) 级联淘汰的优先级确定：如何确定每一层的淘汰优先级，确保优先淘汰对整体性能影响较小的层。4) 时间敏感的淘汰指标的具体公式：如何结合token的生成时间和其他相关信息，设计一个能够反映token重要性随时间变化的指标。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

CAKE在LongBench和NeedleBench上进行了全面评估，实验结果表明，CAKE仅使用3.2%的KV缓存即可维持模型性能，并在各种模型和内存约束下，始终优于当前的基线方法。特别是在低内存设置中，性能提升尤为显著。此外，当使用FlashAttention-2处理128K token的上下文时，CAKE在解码延迟方面实现了超过10倍的加速。

🎯 应用场景

CAKE具有广泛的应用前景，尤其适用于需要处理长序列的LLM推理场景，例如：长文本生成、对话系统、代码生成等。通过更有效地管理KV缓存，CAKE可以显著降低推理延迟，提高资源利用率，使得LLM能够在资源受限的环境下高效运行。此外，CAKE的自适应缓存分配策略也为未来的模型优化和硬件设计提供了新的思路。

📄 摘要（原文）

Large language models (LLMs) excel at processing long sequences, boosting demand for key-value (KV) caching. While recent efforts to evict KV cache have alleviated the inference burden, they often fail to allocate resources rationally across layers with different attention patterns. In this paper, we introduce Cascading and Adaptive KV cache Eviction (CAKE), a novel approach that frames KV cache eviction as a "cake-slicing problem." CAKE assesses layer-specific preferences by considering attention dynamics in both spatial and temporal dimensions, allocates rational cache size for layers accordingly, and manages memory constraints in a cascading manner. This approach enables a global view of cache allocation, adaptively distributing resources across diverse attention mechanisms while maintaining memory budgets. CAKE also employs a new eviction indicator that considers the shifting importance of tokens over time, addressing limitations in existing methods that overlook temporal dynamics. Comprehensive experiments on LongBench and NeedleBench show that CAKE maintains model performance with only 3.2% of the KV cache and consistently outperforms current baselines across various models and memory constraints, particularly in low-memory settings. Additionally, CAKE achieves over 10x speedup in decoding latency compared to full cache when processing contexts of 128K tokens with FlashAttention-2. Our code is available at https://github.com/antgroup/cakekv.