WindowKV：基于任务自适应的群体级KV缓存窗口选择机制，用于高效实现大型语言模型（LLM）推理

《Expert Systems with Applications》：WindowKV: Task-Adaptive Group-Wise KV Cache Window Selection for Efficient LLM Inference

　　```

摘要

随着大型语言模型（LLMs）在长上下文推理能力方面的进步，键值（KV）缓存已成为其基础组件之一。尽管近期研究集中在优化KV缓存占用的内存上，但它们忽视了两个关键因素：保持语义连贯性和在压缩过程中考虑任务特定特征。为了解决这些局限性，我们提出了一种新颖的任务自适应KV缓存窗口选择方法，称为WindowKV。WindowKV根据任务特定特征动态选择由连续令牌组成的局部语义窗口，确保保留的KV缓存能够捕捉到连续且重要的上下文。此外，我们还引入了一种组内层KV缓存索引共享策略，以减少计算开销，从而在性能和效率之间取得平衡。我们在LongBench基准测试中严格评估了WindowKV，结果表明，它在使用原始KV缓存仅12%的存储空间的情况下，仍能保持与完整KV缓存相当的性能，并显著降低了内存需求。此外，我们的方法在Needle-in-a-Haystack评估中也取得了最佳效果，彰显了其有效性和稳健性。（我们的代码可在以下链接获取：GitHub）。

引言

在LLMs时代，需要长上下文理解的任务，如长文本阅读理解（Trivedi, Balasubramanian, Khot, & Sabharwal (2022)）、上下文学习（Dong et al. (2024b)）、文档总结（Huang, Cao, Parulian, Ji, & Wang (2021)）和代码完成（Zheng et al. (2023)）变得尤为重要。因此，能够处理较长上下文长度的LLMs变得越来越普遍（Huang et al. (2023); Xu, Bi, & Liu (2026); Yao, He, & Xu (2026)）。例如，像GPT-4和DeepSeek-V3这样的模型支持多达128K令牌的上下文长度，而Claude-3.5和Yi将这一能力扩展到200K令牌（DeepSeek-AI (2024b); Fu et al. (2024); OpenAI (2023)）。然而， transformer架构中的自注意力机制随着上下文长度的增加而呈现出二次复杂性（Vaswani et al. (2017)），从而导致长上下文场景下的推理延迟显著增加。一种减轻这种延迟的潜在方法是将前一个令牌的键值（KV）状态缓存起来，从而避免重新计算历史上下文（Waddington, Colmenares, Kuang, & Song (2013)）。然而，随着输入上下文长度和层数的增加，存储KV状态所需的内存也会大幅增加（Shi, Zhang, Yao, Li, & Zhao (2024)）。例如，在LLaMA2-7B模型中，存储100K令牌的KV缓存需要超过50GB的内存（Touvron et al. (2023），而2K令牌的上下文则只需要不到1GB的内存（Wu, Wang, Xiao, Peng, & Fu (2025)）。

近期研究通过修改注意力架构（Ainslie et al. (2023); DeepSeek-AI (2024a); Shazeer (2019）或实现层间KV缓存共享（Brandon, Mishra, Nrusimha, Panda, & Ragan-Kelley (2024); Sun et al. (2024); Wu & Tu (2024）来缓解上述内存限制。然而，这些方法需要额外的训练。相比之下，另一条研究路线是在推理阶段压缩KV缓存。例如，一些方法在固定层预算下丢弃非必要令牌的KV状态（Adnan et al. (2024); Ge et al. (2024); Xiao, Tian, Chen, Han, & Lewis (2024); Zhang et al. (2023)）。然而，这些方法忽略了层间注意力稀疏性的差异。PyramidInfer（Yang et al. (2024b）和PyramidKV（Cai et al. (2024）观察到，较低层中密集注意力较为普遍，而在较高层中稀疏注意力占主导（Wang et al. (2023)）。他们提出在不同层之间分配不同的KV缓存大小，以维持金字塔状结构。

然而，上述方法各自独立选择KV缓存令牌，所选离散令牌会破坏上下文的语义一致性。这也与人类的阅读习惯相矛盾。在处理长文本时，人们通常不是逐个令牌阅读，而是以窗口为单位处理信息（Rayner (1998）。此外，这些方法在所有任务中采用统一的压缩策略，限制了它们的任务特定适应性和整体性能。实际上，根据人类经验，不同任务的信息处理方式有很大差异。例如，在问答任务中，这可以被视为信息定位任务¹，需要处理整个窗口以捕捉全面的语义细节，从而便于生成准确的答案。相比之下，在摘要任务中，目标是信息聚合²，需要从每个窗口中提取最显著的信息以生成简洁的摘要。上述挑战激发了我们提出任务自适应KV缓存窗口选择方法的想法。此外，为了提高计算效率，我们整合了组内层KV缓存索引共享策略，以便在受限的内存预算下更好地维持模型的性能。

总之，我们的贡献如下：

•

与之前选择离散令牌的KV缓存压缩方法不同，我们提出了一种任务自适应的窗口选择方法WindowKV。WindowKV根据任务特定特征动态压缩KV缓存，同时保持局部窗口内的语义连贯性。

•

此外，我们提出了一种组内层KV缓存索引共享策略，以减少计算开销，在性能和效率之间取得平衡。

•

我们在LongBench和Needle-in-a-Haystack基准测试上进行了广泛的实验。结果表明，WindowKV在各种骨干LLMs和KV缓存配置上取得了最佳效果，并在Needle-in-a-Haystack测试中超越了其他基线方法。

章节片段

训练和推理中的KV缓存效率

大量的先前研究探索了修改transformer架构以缩小KV缓存大小的方法，包括按头部修改的方法（Ainslie et al. (2023); DeepSeek-AI (2024a); Shazeer (2019）和按层修改的方法（Brandon et al. (2024); Sun et al. (2024); Wu & Tu (2024）。这些方法通过在不同层之间共享或压缩KV状态（例如，跨层注意力Brandon et al. (2024)和层压缩KV缓存Wu & Tu (2024)）或重新设计Transformer来减少KV缓存内存。

KV缓存压缩概述

在基于自回归transformer的LLMs中，生成第个令牌需要注意力模块访问所有前置$i?1$个令牌的KV状态。为了提高计算效率并避免重复计算，这些KV状态在初次计算后会被缓存。这种缓存机制通过消除重复计算的需要显著加快了推理过程。然而，KV缓存可能会对内存产生很大需求，尤其是在处理长文本时。

实验

在本节中，我们进行了全面的实验来评估WindowKV的有效性。

结论

在这项工作中，我们提出了WindowKV，旨在解决现有方法中的两个关键限制：保持语义连贯性和在压缩过程中考虑任务特定特征。通过根据任务特定特征动态选择局部语义窗口，WindowKV确保保留的KV缓存能够保持连续且重要的上下文。在LongBench上的评估表明，WindowKV的性能与完整KV缓存相当。

左宇辉：概念化、方法论、软件、验证、撰写——原始草稿。魏思博：方法论、软件、形式分析、撰写——审稿与编辑。张晨：概念化、撰写——审稿与编辑。刘卓瑞：调查。卢文鹏：资源、撰写——审稿与编辑。宋大伟：撰写——审稿与编辑、资金获取。

利益冲突声明

作者声明他们没有已知的可能会影响本文报告工作的竞争性财务利益或个人关系。

左宇辉|魏思博|张晨|刘卓瑞|卢文鹏|宋大伟

``````