重新表述长上下文 LLM 推理的 KV 缓存驱逐问题

大语言模型 (LLM) 支持长上下文推理，但由于键值 (KV) 缓存增长而遭受大量内存和运行时开销。现有的 KV Cache 驱逐方法主要依赖于局部注意力权重，忽略了值表示、输出投影和头间交互的影响。在这项工作中，我们将 KV 缓存驱逐从传统的逐头、权重平均方法重新表述为输出感知、逐层矩阵乘法近似问题。我们引入了 LaProx，一种新颖的驱逐策略，它显式地模拟了注意力图和预测值状态之间的乘法交互，以准确量化词元贡献，同时考虑到头部之间的依赖关系。在此指标的基础上，我们提出了第一个统一的逐出策略，该策略为词元分配全局可比较的重要性分数，从而实现模型范围的选择，而不是局部的、明智的决策。在长上下文基准 LongBench 和 Needle-In-A-Haystack 上的 19 个数据集上的实验结果表明，我们的方法仅用 5% 的 KV 缓存即可保持模型性能，并且在所有配置中始终优于先前的工作。值得注意的是，与现有最先进的基线相比，我们的方法在极端压缩场景下以最小的开销实现了高达 2 倍的精度损失减少。