Where Matters More Than What: Decoding-aligned KV Cache Compression via Position-aware Pseudo Queries

📄 arXiv: 2603.11564v1 📥 PDF

作者: Zhenxu Tian, Yi Su, Juntao Li, Min Zhang

分类: cs.CL

发布日期: 2026-03-12


💡 一句话要点

提出DapQ:通过位置感知伪查询实现解码对齐的KV缓存压缩

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

关键词: KV缓存压缩 大型语言模型 解码对齐 位置感知 伪查询

📋 核心要点

  1. 现有KV缓存压缩方法在预填充阶段评估token重要性,但忽略了解码过程中的实际需求,导致关键token被错误地驱逐。
  2. DapQ通过位置感知的伪查询模拟解码阶段的查询,从而建立更有效的观察窗口,更准确地评估token的重要性。
  3. 实验表明,DapQ在严格的内存约束下表现出色,例如在3% KV缓存预算下,NIAH数据集上实现了接近无损的99.5%性能。

📝 摘要(中文)

Key-Value (KV) 缓存对于高效的大型语言模型 (LLM) 推理至关重要,但过长的上下文会显著增加 KV 缓存的内存占用。现有的 KV 缓存压缩方法通常依赖于提示观察窗口内的输入端注意力模式来估计预填充阶段的 token 重要性。由于这些评估并非源自解码过程,因此它们无法保留用于未来生成的关键 token。直观地说,有效的观察窗口应该反映解码阶段的查询,以准确反映生成过程将关注哪些 token。然而,真实的解码查询在推理过程中本质上是不可用的。为了构建伪查询来近似它们,我们发现位置信息比语义内容起着更关键的作用。受此启发,我们提出了解码对齐的 KV 缓存压缩方法,通过位置感知伪查询 (DapQ),这是一种新颖且轻量级的驱逐框架,它利用位置感知伪查询来模拟输出 token,从而为重要性评估建立有效的观察窗口。它与实际生成上下文紧密对齐,并能够实现精确的 token 驱逐。在多个基准和 LLM 上的广泛评估表明,DapQ 实现了卓越的性能,尤其是在严格的内存约束下(例如,在 3% KV 缓存预算下,在 NIAH 上实现了接近无损的 99.5% 性能)。

🔬 方法详解

问题定义:论文旨在解决大型语言模型(LLM)推理过程中,由于上下文长度增加导致的KV缓存内存占用过高的问题。现有方法主要依赖于输入端的注意力模式来评估token的重要性,但这种评估方式忽略了解码阶段的实际查询需求,导致关键token被错误地驱逐,影响生成质量。

核心思路:论文的核心思路是构建与解码过程对齐的KV缓存压缩方法。通过模拟解码阶段的查询,更准确地评估token的重要性,从而保留对未来生成至关重要的token。论文发现位置信息在模拟解码查询中起着比语义信息更重要的作用。

技术框架:DapQ框架主要包含以下几个步骤:1. 使用位置感知伪查询生成器,模拟解码阶段的查询;2. 利用生成的伪查询,建立有效的观察窗口;3. 在观察窗口内,评估每个token的重要性;4. 根据重要性评估结果,驱逐不重要的token,压缩KV缓存。

关键创新:DapQ的关键创新在于使用位置感知的伪查询来模拟解码阶段的查询。与现有方法相比,DapQ更关注解码过程中的实际需求,从而能够更准确地评估token的重要性。此外,DapQ框架轻量级,易于集成到现有的LLM推理系统中。

关键设计:位置感知伪查询生成器是DapQ的关键组成部分。该生成器利用token的位置信息来构建伪查询,例如,可以使用token的相对位置或绝对位置作为伪查询的特征。此外,DapQ可以使用不同的重要性评估指标,例如注意力权重、梯度等。具体参数设置需要根据不同的LLM和数据集进行调整。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

DapQ在多个基准测试和LLM上进行了广泛评估,结果表明其性能优于现有方法。特别是在严格的内存约束下,DapQ表现出色,例如在3% KV缓存预算下,在NIAH数据集上实现了接近无损的99.5%性能。此外,DapQ的轻量级设计使其易于集成到现有的LLM推理系统中,具有很强的实用性。

🎯 应用场景

DapQ可应用于各种需要长上下文推理的大型语言模型,例如机器翻译、文本摘要、对话生成等。通过有效压缩KV缓存,DapQ可以降低LLM的内存需求,使其能够在资源受限的设备上运行,并提高推理效率,降低部署成本。该技术还有助于支持更长的上下文窗口,从而提升LLM处理复杂任务的能力。

📄 摘要(原文)

The Key-Value (KV) cache is crucial for efficient Large Language Models (LLMs) inference, but excessively long contexts drastically increase KV cache memory footprint. Existing KV cache compression methods typically rely on input-side attention patterns within a prompt observation window to estimate token importance during the prefill stage. They fail to preserve critical tokens for future generation since these assessments are not derived from the decoding process. Intuitively, an effective observation window should mirror the decoding-stage queries to accurately reflect which tokens the generation process will attend to. However, ground-truth decoding queries are inherently unavailable during inference. For constructing pseudo queries to approximate them, we find that positional information plays a more critical role than semantic content. Motivated by this insight, we propose decoding-aligned KV cache compression via position-aware pseudo queries (DapQ), a novel and lightweight eviction framework that leverages position-aware pseudo queries to simulate the output tokens, thereby establishing an effective observation window for importance assessment. It aligns closely with the actual generation context and enables precise token eviction. Extensive evaluations across multiple benchmarks and LLMs demonstrate that DapQ achieves superior performance, particularly under strict memory constraints (e.g., up to nearly lossless performance 99.5% on NIAH with 3% KV cache budgets).