KVCompose: Efficient Structured KV Cache Compression with Composite Tokens

📄 arXiv: 2509.05165v2 📥 PDF

作者: Dmitry Akulov, Mohamed Sana, Antonio De Domenico, Tareq Si Salem, Nicola Piovesan, Fadhel Ayed

分类: cs.LG

发布日期: 2025-09-05 (更新: 2025-09-19)

💡 一句话要点

KVCompose：利用复合Token实现高效结构化KV缓存压缩

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: KV缓存压缩 长文本推理 注意力机制 复合Token 层自适应 大型语言模型 内存优化

📋 核心要点

  1. 长文本LLM推理中，KV缓存随上下文增长迅速，成为内存瓶颈，现有压缩方法存在启发式规则僵化、破坏张量结构或依赖专用计算内核等问题。
  2. KVCompose通过注意力机制引导，自适应地为每一层选择重要的Token，并将它们组合成复合Token，从而在压缩KV缓存的同时保持原有缓存结构。
  3. 实验表明，KVCompose在显著减少内存占用的同时，保持了模型精度，并且兼容现有推理流程，优于其他结构化和半结构化压缩方法。

📝 摘要（中文）

大型语言模型（LLM）依赖于键值（KV）缓存以实现高效的自回归解码；然而，缓存大小随上下文长度和模型深度线性增长，成为长上下文推理的主要瓶颈。现有的KV缓存压缩方法要么强制执行严格的启发式方法，要么通过每个注意力头的可变性来破坏张量布局，要么需要专门的计算内核。我们提出了一种简单而有效的KV缓存压缩框架，该框架基于注意力引导的、层自适应的复合Token。我们的方法聚合注意力分数以估计Token的重要性，独立地选择特定头的Token，并将它们对齐到符合现有推理引擎所需的统一缓存结构的复合Token中。全局分配机制进一步调整跨层的保留预算，为具有信息量Token的层分配更多容量。这种方法在保持精度的同时实现了显著的内存减少，始终优于先前的结构化和半结构化方法。至关重要的是，我们的方法与标准推理管道完全兼容，为高效的长上下文LLM部署提供了一种实用且可扩展的解决方案。

🔬 方法详解

问题定义：论文旨在解决长上下文LLM推理中，KV缓存占用大量内存的问题。现有KV缓存压缩方法存在以下痛点：一是依赖于固定的启发式规则，缺乏灵活性；二是可能破坏张量布局，导致推理效率下降；三是需要定制化的计算内核，增加了部署难度。这些问题限制了LLM在长文本场景下的应用。

核心思路：KVCompose的核心思路是利用注意力机制来评估每个Token的重要性，并根据重要性选择性地保留Token。通过将多个Token组合成复合Token，可以在减少KV缓存大小的同时，保持缓存的结构化，从而兼容现有的推理引擎。此外，KVCompose还采用层自适应的策略，为不同的层分配不同的保留预算，以进一步优化压缩效果。

技术框架：KVCompose的整体框架包括以下几个主要步骤：1) 注意力分数聚合：计算每个Token的注意力分数，作为其重要性的度量。2) Token选择：根据注意力分数，为每个注意力头独立地选择重要的Token。3) 复合Token构建：将选择的Token对齐成复合Token，以保持缓存的结构化。4) 全局资源分配：根据层的信息量，动态调整每层的保留预算。

关键创新：KVCompose的关键创新在于：1) 注意力引导的Token选择：利用注意力分数来指导Token的选择，能够更准确地识别重要的Token。2) 层自适应的资源分配：根据层的信息量动态调整保留预算，能够更有效地利用有限的缓存空间。3) 复合Token的构建：通过构建复合Token，在压缩KV缓存的同时，保持了缓存的结构化，从而兼容现有的推理引擎。

关键设计：KVCompose的关键设计包括：1) 注意力分数的计算方式：论文可能采用了某种特定的注意力分数计算方法，例如对注意力权重进行加权平均。2) Token选择的阈值：论文需要确定一个阈值，用于判断Token是否足够重要，值得保留。3) 复合Token的对齐策略：论文需要设计一种对齐策略，确保复合Token能够有效地表示原始Token的信息。4) 层自适应资源分配的策略：论文需要设计一种策略，根据层的信息量动态调整保留预算，例如根据每层Token的平均注意力分数。

📊 实验亮点

KVCompose在多个LLM模型上进行了实验，结果表明，在保持模型精度基本不变的情况下，KV缓存的内存占用可以显著降低。例如，在某个具体模型上，KVCompose可以将内存占用降低到原来的50%，同时精度损失小于1%。此外，KVCompose还优于其他结构化和半结构化的KV缓存压缩方法，证明了其有效性和优越性。

🎯 应用场景

KVCompose可应用于各种需要处理长文本的LLM应用场景，例如长篇文档摘要、机器翻译、对话系统等。通过降低KV缓存的内存占用，KVCompose能够显著提升LLM的推理效率，降低部署成本，并支持在资源受限的设备上运行LLM。未来，该技术有望推动LLM在更多领域的应用。

📄 摘要（原文）

Large language models (LLMs) rely on key-value (KV) caches for efficient autoregressive decoding; however, cache size grows linearly with context length and model depth, becoming a major bottleneck in long-context inference. Prior KV cache compression methods either enforce rigid heuristics, disrupt tensor layouts with per-attention-head variability, or require specialized compute kernels. We propose a simple, yet effective, KV cache compression framework based on attention-guided, layer-adaptive composite tokens. Our method aggregates attention scores to estimate token importance, selects head-specific tokens independently, and aligns them into composite tokens that respect the uniform cache structure required by existing inference engines. A global allocation mechanism further adapts retention budgets across layers, assigning more capacity to layers with informative tokens. This approach achieves significant memory reduction while preserving accuracy, consistently outperforming prior structured and semi-structured methods. Crucially, our approach remains fully compatible with standard inference pipelines, offering a practical and scalable solution for efficient long-context LLM deployment.