LookaheadKV: Fast and Accurate KV Cache Eviction by Glimpsing into the Future without Generation
AI 摘要
LookaheadKV通过预测重要性得分,无需生成草稿即可实现快速准确的KV缓存淘汰。
主要贡献
- 提出LookaheadKV框架,无需草稿生成即可预测KV缓存的重要性。
- 设计参数高效的模块,用于提升重要性得分预测的准确性。
- 实验表明,LookaheadKV在长文本理解任务中表现优异,并显著降低了淘汰成本。
方法论
通过训练参数高效的模块来预测KV缓存的重要性得分,无需生成耗时的草稿,实现快速准确的淘汰。
原文摘要
Transformer-based large language models (LLMs) rely on key-value (KV) caching to avoid redundant computation during autoregressive inference. While this mechanism greatly improves efficiency, the cache size grows linearly with the input sequence length, quickly becoming a bottleneck for long-context tasks. Existing solutions mitigate this problem by evicting prompt KV that are deemed unimportant, guided by estimated importance scores. Notably, a recent line of work proposes to improve eviction quality by "glimpsing into the future", in which a draft generator produces a surrogate future response approximating the target model's true response, and this surrogate is subsequently used to estimate the importance of cached KV more accurately. However, these approaches rely on computationally expensive draft generation, which introduces substantial prefilling overhead and limits their practicality in real-world deployment. To address this challenge, we propose LookaheadKV, a lightweight eviction framework that leverages the strength of surrogate future response without requiring explicit draft generation. LookaheadKV augments transformer layers with parameter-efficient modules trained to predict true importance scores with high accuracy. Our design ensures negligible runtime overhead comparable to existing inexpensive heuristics, while achieving accuracy superior to more costly approximation methods. Extensive experiments on long-context understanding benchmarks, across a wide range of models, demonstrate that our method not only outperforms recent competitive baselines in various long-context understanding tasks, but also reduces the eviction cost by up to 14.5x, leading to significantly faster time-to-first-token. Our code is available at https://github.com/SamsungLabs/LookaheadKV.