优化大语言模型KV缓存驱逐:自适应分配提升预算利用率

大型语言模型（LLMs）在多个领域均实现了突出的成效，被普遍运用于对话系统、文档摘要以及代码生成等各类自然语言处理任务。然而，随着处理序列长度的提升，LLMs在计算资源消耗和内存需求方面面临显著增长的难题。为应对这些挑战，研究者引入了键值（Key-Value, KV）缓存机制，通过保留先前令牌的中间运算结果，以降低自回归生成过程中的计算压力。

虽然KV缓存技术提升了LLM的运行效率，但在处理长序列推理任务时，缓存容量可能超出模型参数规模，从而引发较大的内存占用和较高的输入/输出（I/O）耗时。为缓解这一问题，研究者提出了KV缓存淘汰策略，通过移除部分非关键缓存条目以缩减缓存体积，并尽可能降低对生成质量的影响。

本文对当前缓存淘汰机制的核心理念进行了深入分析，并提出了一种简洁且高效的自适应分配算法，旨在优化不同注意力头之间的资源分配比例，从而增强生成效果。基于该算法框架，我们设计出两种新型自适应淘汰方案：Ada-SnapKV与Ada-Pyramid。通过在16个数据集上的广泛测试验证，我们证实了这些方法能够在压缩过程中有效降低质量下降幅度，并确立了新的最先进性能标准。

KV缓存驱逐的原理重新审视

KV缓存与多头注意力机制优化

为深入掌握KV缓存驱逐机制的本质，有必要先回溯LLMs中KV缓存与多头注意力机制的运算流程。在自回归生成模式中，LLMs在每一步骤