arXivDaily arXiv每日学术速递 周一至周五更新
arXiv周末暂无论文更新,休息一下吧,周末愉快~~
arXiv 2407.11550cs.CLcs.AI
收藏

Ada-KV:通过自适应预算分配优化KV缓存驱逐以实现高效LLM推理

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

AI总结:

本文提出Ada-KV,一种按头自适应预算分配的KV缓存驱逐策略,通过理论损失上界指导优化,即插即用地提升现有方法在长序列推理中的质量,并在多个数据集上验证有效性。

摘要:

大型语言模型在各个领域表现出色,但由于长序列推理所需的键值(KV)缓存不断增长,面临效率挑战。近期研究尝试通过在运行时驱逐大量非关键缓存元素来减少KV缓存大小,同时保持生成质量。然而,这些方法通常在所有注意力头上均匀分配压缩预算,忽略了每个头独特的注意力模式。本文建立了驱逐前后注意力输出之间的理论损失上界,解释了先前缓存驱逐方法的优化目标,同时指导自适应预算分配的优化。基于此,我们提出Ada-KV,首个按头自适应的预算分配策略。它提供即插即用的优势,能与先前缓存驱逐方法无缝集成。在Ruler的13个数据集和LongBench的16个数据集上进行的广泛评估,均在问题感知和问题无关场景下进行,展示了相较于现有方法的显著质量提升。我们的代码可在https://github.com/FFY0/AdaKV获取。

英文摘要:

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.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑