DepthKV: Layer-Dependent KV Cache Pruning for Long-Context LLM Inference

📄 arXiv: 2604.24647v1 📥 PDF

作者: Zahra Dehghanighobadi, Asja Fischer

分类: cs.CL, cs.AI

发布日期: 2026-04-27

💡 一句话要点

DepthKV：针对长文本LLM推理的分层KV缓存剪枝方法

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

关键词: 长文本推理 KV缓存 剪枝 大语言模型 层相关 内存优化 自回归推理

📋 核心要点

  1. 现有长文本LLM推理中，KV缓存的内存占用随序列长度线性增长，成为性能瓶颈。
  2. DepthKV提出一种层相关的KV缓存剪枝框架，根据各层对剪枝的敏感度动态分配KV缓存预算。
  3. 实验表明，DepthKV在相同剪枝率下优于统一剪枝，有效提升了KV缓存的利用率。

📝 摘要（中文）

长文本推理是大语言模型（LLMs）的关键能力，支持长文档理解、摘要和代码生成等应用。然而，高效的自回归推理依赖于键值（KV）缓存，其内存占用随序列长度线性增长，导致主要的内存瓶颈。为了缓解这种开销，KV缓存剪枝方法在推理期间丢弃具有低注意力分数的缓存token。大多数现有方法在所有层上应用统一的剪枝率，隐含地假设所有层对整体模型性能的贡献相等。我们表明这种假设并非最优，因为各层对剪枝的敏感度差异显著。我们提出了DepthKV，一个层相关的剪枝框架，它基于各层的敏感度在各层之间分配固定的全局KV预算，而不是使用统一的分配。在多个模型和任务中，DepthKV在相同的全局剪枝率下始终优于统一剪枝，证明了通过层相关的分配更有效地利用了KV缓存预算。

🔬 方法详解

问题定义：论文旨在解决长文本LLM推理过程中，KV缓存占用大量内存，导致推理速度降低的问题。现有方法通常采用统一的剪枝策略，即对所有层采用相同的剪枝比例，忽略了不同层对模型性能的贡献差异，导致KV缓存的利用效率不高。

核心思路：DepthKV的核心思想是根据不同层对模型性能的敏感度，动态地分配KV缓存的剪枝比例。对模型性能影响较小的层，可以进行更激进的剪枝，而对模型性能影响较大的层，则保留更多的KV缓存。通过这种方式，可以在相同的全局剪枝率下，更有效地利用KV缓存，从而提高推理速度。

技术框架：DepthKV框架主要包含以下几个步骤：1. 敏感度分析：通过实验或理论分析，确定每一层对模型性能的敏感度。敏感度可以根据层对最终预测结果的影响程度来衡量。2. 预算分配：根据各层的敏感度，将全局KV缓存预算分配给每一层。敏感度越高的层，分配到的预算越多。3. KV缓存剪枝：在每一层，根据分配到的预算，对KV缓存进行剪枝。可以使用各种剪枝策略，例如基于注意力分数的剪枝。

关键创新：DepthKV的关键创新在于提出了层相关的KV缓存剪枝策略。与传统的统一剪枝策略相比，DepthKV能够更有效地利用KV缓存，从而提高推理速度。这种方法能够自适应地调整每一层的剪枝比例，从而更好地适应不同模型的特点。

关键设计：DepthKV的关键设计包括：1. 敏感度指标：如何准确地衡量每一层对模型性能的敏感度是一个关键问题。论文可能采用了某种特定的指标，例如基于梯度的敏感度分析方法。2. 预算分配策略：如何根据各层的敏感度，合理地分配全局KV缓存预算也是一个关键问题。论文可能采用了某种优化算法，例如基于动态规划的预算分配方法。3. 剪枝策略：论文可能采用了某种特定的剪枝策略，例如基于注意力分数的剪枝方法，并对其进行了改进。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

论文实验结果表明，DepthKV在多个模型和任务上均优于统一剪枝方法。在相同的全局剪枝率下，DepthKV能够取得更高的模型性能。具体的性能提升幅度取决于模型、任务和剪枝率，但总体趋势是DepthKV能够更有效地利用KV缓存，从而提高推理速度和模型精度。

🎯 应用场景

DepthKV可应用于各种需要处理长文本的LLM应用场景，例如长文档摘要、代码生成、机器翻译、对话系统等。通过降低KV缓存的内存占用，DepthKV可以提高推理速度，降低部署成本，并支持更大规模的模型和更长的上下文长度。该技术对于资源受限的设备（如移动设备）上的LLM部署尤为重要。

📄 摘要（原文）

Long-context reasoning is a critical capability of large language models (LLMs), enabling applications such as long-document understanding, summarization, and code generation. However, efficient autoregressive inference relies on the key-value (KV) cache, whose memory footprint grows linearly with sequence length, leading to a major memory bottleneck. To mitigate this overhead, KV cache pruning methods discard cached tokens with low attention scores during inference. Most existing methods apply a uniform pruning ratio across layers, implicitly assuming that all layers contribute equally to overall model performance. We show that this assumption is suboptimal, as layers differ significantly in their sensitivity to pruning. We propose DepthKV, a layer-dependent pruning framework that allocates a fixed global KV budget across layers based on their sensitivity, rather than using a uniform allocation. Across multiple models and tasks, DepthKV consistently outperforms uniform pruning at the same global pruning ratio, demonstrating more effective utilization of the KV cache budget through layer-dependent allocation.