arXiv:2506.17121cs.CL2025-06被引 13

提出新指标评估长文本模型内存占用,显著降低存储开销。

Cache Me If You Can: How Many KVs Do You Need for Effective Long-Context LMs?

  • 用统一的KV足迹指标衡量内存使用与生命周期
  • 新方法PruLong在128K上下文下节省12%内存且不降性能
  • 适用于需高效处理长文本的部署场景

语言模型处理越来越长的上下文以完成如书籍摘要等任务,但随之带来关键-值(KV)缓存的内存成本上升。以往方法常通过丢弃KV来减少内存,但其适用性受限于理想条件,存在峰值内存高、性能下降等问题,且缺乏公平比较。本文提出“KV足迹”作为统一度量标准,综合考虑存储的KV数量及其在内存中的存留时间。在长达128K token的上下文中,我们评估各方法在保持长上下文理解与生成性能前提下的最小足迹。结果揭示了现有方法的高峰值内存问题。一类称为“后填充淘汰”的方法因无法在预填充阶段淘汰而造成高足迹;我们将其改进为支持预填充阶段淘汰,大幅降低足迹。接着研究“时效淘汰”方法,提出端到端优化的PruLong,学习哪些注意力头需保留完整KV缓存,哪些可省略。PruLong在保持挑战性回忆任务性能的同时,相较已有方法实现12%的足迹压缩。本工作厘清了长上下文推理方法的复杂关系,为未来最小化KV足迹指明方向。

原文摘要 · Abstract (English)

Language models handle increasingly long contexts for tasks such as book summarization, but this leads to growing memory costs for the key-value (KV) cache. Many prior works have proposed ways of discarding KVs from memory, but their approaches are tailored to favorable settings, obscuring caveats like high peak memory and performance degradation, and a fair comparison between methods is difficult. In this paper, we propose the *KV footprint* as a unified metric, which accounts for both the amount of KV entries stored and their lifespan in memory. We evaluate methods based on the smallest footprint they attain while preserving performance in both long-context understanding and generation, with context lengths of up to 128K tokens. This metric reveals the high peak memory of prior KV eviction methods. One class of methods -- *post-fill eviction* -- has a high footprint due to being incompatible with eviction during pre-filling. We adapt these methods to be able to evict KVs during pre-filling, achieving substantially lower KV footprints. We then turn to *recency eviction* methods, wherein we propose PruLong, an end-to-end optimization method for learning which attention heads need to retain the full KV cache and which do not. PruLong saves memory while preserving long-context performance, achieving 12% smaller KV footprint than prior methods while retaining performance in challenging recall tasks. Our paper clarifies the complex tangle of long-context inference methods and paves the way for future development to minimize the KV footprint.

长上下文内存优化注意力机制

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。