动态选择关键语义窗口,用12%内存保持模型性能
WindowKV: Task-Adaptive Group-Wise KV Cache Window Selection for Efficient LLM Inference
- 按任务特征动态选连续词块,保留核心上下文
- 仅用原内存12%即达全缓存性能,长文本任务表现优
- 适合需要低内存推理的工业级大模型部署
随着大语言模型长上下文推理能力的发展,KV缓存已成为基础组件。然而其巨大的显存占用使得缓存压缩成为工业场景高效推理的关键。现有研究虽聚焦于降低缓存内存占用,却忽视了语义连贯性与任务特异性。为此,我们提出任务自适应的分组式KV缓存窗口选择方法WindowKV。该方法根据任务特征动态选取包含连续词元的局部语义窗口,确保保留的KV缓存具备连续且关键的上下文信息。此外,引入组内层间缓存索引共享策略以降低计算开销,实现性能与效率的平衡。我们在LongBench基准上进行严格评估,结果表明,WindowKV在仅使用原缓存12%内存的情况下,仍能保持与完整缓存相当的性能,显著降低内存需求。同时,在Needle-in-a-Haystack测试中也取得当前最优结果,验证了方法的有效性与鲁棒性。
原文摘要 · Abstract (English)
With the advancements in long-context inference capabilities of large language models (LLMs), the KV cache has become one of the foundational components. However, its substantial GPU memory consumption makes KV cache compression a key technique for enabling efficient LLM inference in industrial scenarios. While recent studies have focused on optimizing the memory occupied by the KV cache, they overlook two critical factors: preserving semantic coherence and considering task-specific characteristic during compression. To address these limitations, we propose a novel task-adaptive KV cache window selection method, WindowKV. WindowKV dynamically selects local semantic windows consisting of consecutive tokens, according to task-specific characteristics, ensuring the retained KV cache captures continuous, essential context. Additionally, we introduce an intra-group layer KV cache indices sharing strategy to reduce computational overhead, achieving a balance between performance and efficiency. We rigorously evaluate WindowKV on the LongBench benchmark, and the results demonstrate that it maintains a performance comparable to full KV cache retention while using only 12% of the original KV cache, significantly reducing memory requirements. Furthermore, our method also achieves state-of-the-art results in the Needle-in-a-Haystack evaluation, highlighting its effectiveness and robustness.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。