LookaheadKV: 通过窥视未来而无需生成实现快速且准确的 KV 缓存驱逐

基于 Transformer 的大型语言模型 (LLM) 依赖于键值 (KV) 缓存以避免自回归推理期间的冗余计算。尽管此机制大大提高了效率，但缓存大小会随输入序列长度线性增长，快速成为长上下文任务的瓶颈。现有的解决方案通过据估计的重要性得分驱动驱逐被视为不重要的提示 KV。值得注意的是，最近的一系列工作提出通过“窥视未来”来提高驱逐质量，其中草稿生成器产生近似目标模型真实响应的 surrogate future response，然后 subsequently 用于更准确地估计已缓存 KV 的重要性。然而，这些方法依赖于计算昂贵的草稿生成，引入显著的 prefilling 开销，限制了其在实际部署中的实用性。为此，我们提出 LookaheadKV，一个轻量级驱逐框架，它利用 surrogate future response 的优势而无需显式草稿生成。LookaheadKV 通过在 transformer 层中增添参数高效模块来预测高准确性的真实重要性得分。我们的设计确保与现有低成本启发式方法相当的运行时开销，同时实现优于更昂贵近似方法的准确性。广泛的实验在长上下文理解基准测试中进行，涵盖多种模型，证明我们的方法不仅在各种长上下文理解任务中超过最新有竞争力的基线，还将驱逐成本降低最高可达 14.5 倍，显著加快 first-token 生成速度。我们的代码可在 https://github.com/SamsungLabs/LookaheadKV 查阅。