arXiv:2503.12491cs.CL2025-03ICLR被引 66

CAKE动态分配各层缓存,用时序感知机制显著降低内存占用。

CAKE: Cascading and Adaptive KV Cache Eviction with Layer Preferences

  • 按层分析注意力动态,分层分配缓存资源
  • 仅需3.2%缓存即保持模型性能,低内存下优势明显
  • 适合长文本生成、资源受限部署场景

大型语言模型(LLMs)在处理长序列时表现出色,推动了键值(KV)缓存的需求。尽管近期的缓存淘汰方法减轻了推理负担,但往往未能根据具有不同注意力模式的各层合理分配资源。本文提出级联自适应KV缓存淘汰(CAKE),将缓存淘汰问题建模为“蛋糕切分”问题。CAKE通过考虑空间与时间维度上的注意力动态,评估各层偏好,相应分配合理缓存大小,并以级联方式管理内存约束。该方法实现全局缓存分配视角,在维持内存预算的同时,自适应地分配多样化注意力机制的资源。此外,CAKE引入新的淘汰指标,考虑了随时间变化的词元重要性,克服了现有方法忽视时序动态的局限。在LongBench和NeedleBench上的全面实验表明,CAKE仅需3.2%的KV缓存即可保持模型性能,且在多种模型和内存约束下持续优于当前基线,尤其在低内存环境下表现突出。当使用FlashAttention-2处理128K token上下文时,解码延迟比完整缓存快10倍以上。代码已开源:https://github.com/antgroup/cakekv。

原文摘要 · Abstract (English)

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.

KV缓存长文本生成推理优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。