KeyDiff：基于键相似性的长上下文 LLM 推理中的 KV 缓存驱逐
人工智能
2026-01-21 v4
摘要
我们展示了在LLM推理过程中，几何上具有鲜艳特征的键倾向于获得高注意力得分。基于这一现象，我们提出KeyDiff，一种仅基于键相似性的KV缓存驱逐方法。与其他KV缓存驱逐方法不同，KeyDiff能够在严格的资源约束下处理任意长度的提示，并高效生成响应。我们通过将键的多样性与注意力得分建立关系，为KeyDiff提供了理论依据。这些结果表明KeyDiff能够高效识别最需要保留的关键token。值得注意的是，KeyDiff不依赖注意力得分，这使得使用像FlashAttention等优化注意力机制成为可能。在严格的内存限制下，我们在Llama和Qwen模型系列上验证了KeyDiff的有效性，通过观察在LongBench上Llama 3.1-8B和Llama 3.2-3B的性能差距不足0.04%（相当于8K缓存预算下的约23% KV缓存减量），以及在Deepseek-R1-Distill-Llama-8B上在Math500推理基准测试中实现接近基线的性能，我们还观察到相较于其他token驱逐方法，端到端推理延迟降低最高可达30%。
引用
@article{arxiv.2504.15364,
title = {KeyDiff: Key Similarity-Based KV Cache Eviction for Long-Context LLM Inference in Resource-Constrained Environments},
author = {Junyoung Park and Dalton Jones and Matthew J Morse and Raghavv Goel and Mingu Lee and Chris Lott},
journal= {arXiv preprint arXiv:2504.15364},
year = {2026}
}
备注
37 pages, 19 figures, NeurIPS 2025