Jul, 2024

优化LLM中的KV缓存逐出策略:自适应分配以提升预算利用率

TL;DR大型语言模型在各个领域取得了很好的成绩,但由于长序列推理所需的广泛KV缓存限制了效率。许多努力尝试在运行时清除非关键缓存元素,从而在给定的内存预算内减小缓存大小,同时保持生成质量。我们重新审查了现有策略的基本原理,并发现当前的做法是在驱逐过程中均匀分配预算到不同的注意力头上,这往往会降低驱逐后生成质量。基于这些发现,我们提出了一种简单而有效的自适应分配算法,不仅在理论上确保其损失上限不超过之前均匀分配方法的上限,而且还能与自注意机制的特性相一致,从而实际降低上限。此外,将此算法与两种最先进的方法结合起来,得到了Ada-SnapKV和Ada-Pyramid。对16个数据集和“大海捞针”测试的广泛实验验证证实,Ada-SnapKV和Ada-Pyramid实现了进一步的改进,在最高性能的技术水平上建立了新的基准。