针对长文本推理的缓存瓶颈,提出按层敏感度动态分配缓存空间。
DepthKV: Layer-Dependent KV Cache Pruning for Long-Context LLM Inference
- 按各层对缓存丢失的敏感度差异分配缓存预算,而非统一裁剪。
- 在相同缓存比例下,多个模型和任务上均优于传统均匀裁剪方法。
- 适合需要处理长文本的场景,如文档理解、代码生成等。
长文本推理是大语言模型的关键能力,支撑文档理解、摘要生成和代码生成等应用。然而,自回归推理依赖键值(KV)缓存,其内存开销随序列长度线性增长,形成主要瓶颈。现有方法通过丢弃注意力分数低的缓存项来缓解,但普遍采用各层统一裁剪比例,隐含假设所有层对性能影响相同。我们发现该假设不成立,各层对裁剪的敏感度差异显著。为此提出DepthKV,基于层敏感度将固定全局缓存预算分配至各层,而非均匀分配。在多个模型和任务上,DepthKV在相同全局裁剪率下持续优于均匀裁剪,证明了分层动态分配能更高效利用缓存资源。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。