用可学习的保留头实现精准缓存淘汰，让消费级显卡跑长文本大模型
Locret: Enhancing Eviction in Long-Context LLM Inference with Trained Retaining Heads on Consumer-Grade Devices
- 通过可训练的保留头评估注意力缓存重要性，实现精准淘汰
- 在<10%性能损失下实现20倍缓存压缩，支持128K以上上下文
- 仅需不到1小时额外训练，适合在4090等消费级设备部署
扩展大语言模型输入上下文长度会显著增加计算成本和内存开销，主要源于注意力键值（KV）缓存的膨胀。现有KV缓存压缩方法存在策略低效、内存缩减有限等问题，导致大模型难以在消费级设备上进行长上下文推理，尤其在处理长序列流式输入时更为困难。这限制了复杂应用的落地，阻碍了大模型的普及。为此，我们提出Locret，首个支持分块预填充的缓存淘汰框架。通过可学习的保留头评估KV缓存单元的因果重要性，实现精确淘汰，提升长上下文推理效率。大量实证研究表明，Locret在内存效率和生成质量上均优于近期主流方法：在性能损失低于10%的前提下，实现最高达20倍的KV缓存压缩比；且可在单张NVIDIA 4090 GPU上完成128K以上长上下文推理，无需牺牲生成质量，仅需不到1个GPU小时的额外训练。
原文摘要 · Abstract (English)
Scaling the input context length of a large language model (LLM) incurs a significant increase in computation cost and memory footprint to maintain the attention key-value (KV) cache. Existing KV cache compression methods suffer from inefficient compression strategies and limited memory reduction effects, making it difficult for LLMs to conduct long-context inference on consumer-grade devices, especially when inferring long-context stream input. Such obstacles prevent consumer-grade devices from supporting more complex applications, creating challenges for the democratization of LLMs. To overcome this, we propose Locret, the first framework to create an eviction policy compatible with chunked prefill. By evaluating the causal importance of KV cache units by learnable retaining heads, Locret enables precise eviction of cache units, facilitating efficient long-context inference. In our extensive empirical studies, Locret outperforms the recent popular and competitive approaches in terms of memory efficiency and generation quality -- Locret achieves up to 20x of KV cache compression ratio within less than 10% performance loss. Furthermore, Locret achieves 128K+ long-context inference on a single NVIDIA 4090 GPU without compromising generation quality and only costs <1 GPU hour of additional training.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。