PolyKV: 异构保留与分配的KV缓存压缩

# PolyKV：用于KV缓存压缩的异构保留与分配
来源：https://arxiv.org/html/2606.15157
Chao Fei Panos Kalnis 阿卜杜拉国王科技大学 \{chao\.fei, panos\.kalnis\}@kaust\.edu\.sa

###### 摘要

KV缓存压缩对于降低长上下文大语言模型推理的内存成本至关重要。然而，现有方法通常对所有Transformer层应用单一的压缩策略和统一的缓存预算。这种统一设计忽略了不同层在预填充和解码阶段可能扮演不同角色，因此可能需要不同的驱逐策略和缓存容量。我们提出PolyKV，一个逐层KV缓存优化框架，它考虑了方法选择和预算分配的设计空间。PolyKV基于层级别信号为每一层分配合适的KV压缩策略，同时在固定总预算下分配非均匀预算。这种表述能够实现现有KV缓存方法的异构组合。在LLaMA-3.1-8B和Qwen3-8B上的实验表明，在相同的512 token平均KV预算下，PolyKV分别恢复了最强单一策略基线与FullKV之间LongBench性能差距的54.5%和25.7%。在128–1024预算的扫描中，PolyKV相较于最强基线持续提升1.7%–6.4%，相当于恢复FullKV差距的40.0%–54.5%。

## 1 引言

随着智能体应用、检索增强生成以及其他需要模型处理长上下文的负载的兴起，长上下文LLM推理正变得越来越重要。然而，更长的上下文显著增加了KV缓存的大小，造成了推理过程中的主要内存瓶颈。不断增长的KV缓存可能耗尽HBM容量，降低可实现的批处理大小，并拖慢token生成速度。

KV缓存驱逐通过减少缓存内存使用量来缓解这一瓶颈，同时保留长上下文质量。现有方法已引入一系列有效的驱逐模式。重击者和持久性方法保留获得大量或稳定注意力权重的token(Zhang等人，2023 (https://arxiv.org/html/2606.15157#bib.bib19)；Liu等人，2023 (https://arxiv.org/html/2606.15157#bib.bib8))。注意力汇聚方法保留汇聚token以及最近的token(Xiao等人，2024b (https://arxiv.org/html/2606.15157#bib.bib17))。提示观察方法在生成前识别出重要的提示token(Li等人，2024 (https://arxiv.org/html/2606.15157#bib.bib7))。这些方法共同表明，缓存驱逐模式是长上下文推理中质量和内存使用之间权衡的关键因素。然而，典型的KV驱逐方法主要定义有用的驱逐模式，并且通常被实例化为模型范围的策略。这种统一处理忽视了Transformer层的异质性，隐含地假设不同层应遵循相同的驱逐规则并接受相当的缓存容量。

Transformer层并非同质的计算模块。先前的工作表明，不同层和头在注意力模式、表示抽象、上下文角色以及对剪枝的敏感性方面存在差异(Tenney等人，2019 (https://arxiv.org/html/2606.15157#bib.bib13)；Jawahar等人，2019 (https://arxiv.org/html/2606.15157#bib.bib6)；Clark等人，2019 (https://arxiv.org/html/2606.15157#bib.bib2)；Voita等人，2019 (https://arxiv.org/html/2606.15157#bib.bib14)；Michel等人，2019 (https://arxiv.org/html/2606.15157#bib.bib9))。这种异质性表明全局KV驱逐规则是一种粗糙的抽象。全局策略对所有层施加相同的驱逐偏好和相当的缓存容量。因此，我们将KV驱逐视为一个逐层的设计空间。这个设计空间暴露了两个耦合的选择。第一个是驱逐什么，因为不同层可能偏好不同的驱逐模式。第二个是保留多少，因为不同层对KV缓存缩减的敏感性不同。近来自适应KV方法已开始揭示非均匀驱逐的价值(Ge等人，2024 (https://arxiv.org/html/2606.15157#bib.bib5)；Cai等人，2024 (https://arxiv.org/html/2606.15157#bib.bib1)；Feng等人，2025 (https://arxiv.org/html/2606.15157#bib.bib3)；Zhou等人，2025 (https://arxiv.org/html/2606.15157#bib.bib20)；Shen等人，2025 (https://arxiv.org/html/2606.15157#bib.bib11)；Tang等人，2025 (https://arxiv.org/html/2606.15157#bib.bib12)；Xiao等人，2024a (https://arxiv.org/html/2606.15157#bib.bib16)；Fu等人，2025 (https://arxiv.org/html/2606.15157#bib.bib4))。然而，它们仅探索了设计空间的一部分，而非将驱逐模式和缓存容量视为耦合的逐层设计选择。

参见标题图1：KV缓存压缩中的逐层异质性（LongBench上）。(A) 压缩误差在不同层和方法之间变化。(B) 我们在相同总预算下比较了均匀KV分配与使用PyramidKV的逐层分配。图1 (https://arxiv.org/html/2606.15157#S1.F1) 显示两个维度在不同层间都有显著变化。面板A显示，最低误差的驱逐模式随层深度变化，表明没有单一的驱逐规则在整个模型中统一最优。面板B显示，在大多数LongBench任务上，将相同的总KV预算重新分配到各层优于均匀分配。这些观察共同促使我们将驱逐模式和缓存容量视为耦合的逐层设计选择。

我们提出PolyKV，一个用于设计和评估逐层KV缓存驱逐策略的框架。PolyKV通过允许为每一层分配一个驱逐模式和一个缓存预算，使这两个设计维度明确化。通过模块化接口，PolyKV集成了广泛使用的驱逐模式，并支持对驱逐模式和缓存容量的细粒度控制。这个统一框架使我们能够系统地构建和评估将两个维度与每层特性对齐的驱逐策略。

为了在不穷尽搜索整个设计空间的情况下实例化该框架，PolyKV使用一个小型校准集进行离线校准。它收集两类特定层的信号。驱逐信号表征一层依赖哪些上下文模式，指导驱逐模式的选择。缓存敏感信号估计一层受KV缩减影响的程度，指导缓存容量的分配。PolyKV将这些信号转换为跨层的固定异构驱逐策略。在部署时，推理遵循这个校准后的策略，而不是进行在线策略搜索，从而使得层感知的KV驱逐实用化，同时保持运行时过程静态。

我们进行了全面的实验，以证明PolyKV暴露的逐层设计空间可以转化为更好的质量。在固定的总KV预算下，我们在LLaMA-3.1-8B和Qwen3-8B上，在LongBench和RULER上比较了PolyKV与代表性的KV缓存驱逐基线。在LongBench上，在相同的512 token平均KV预算下，PolyKV在两个模型家族上都优于最强的单一策略基线，分别恢复了LLaMA-3.1-8B和Qwen3-8B上与FullKV之间压缩导致的差距的54.5%和25.7%。在LLaMA-3.1-8B上的预算扫描进一步显示，从128到1024平均KV token，PolyKV持续改进，相较于最强基线提升1.7%–6.4%。

我们的贡献有三方面：

- •我们将KV缓存驱逐重新表述为一个逐层设计空间。该空间耦合了两个选择：保留什么以及每层在共享内存约束下获得多少KV容量。
- •我们提出PolyKV，一个通过以层粒度组合驱逐模式和缓存预算来操作化该设计空间的框架。PolyKV通过统一接口集成多种驱逐原语，并使用离线校准信号指导逐层策略构建。
- •我们通过实验表明，在固定内存预算下，逐层异构策略能提高质量。在LongBench上，在相同的512 token平均KV预算下，PolyKV分别恢复了LLaMA-3.1-8B和Qwen3-8B上最强单一策略基线与FullKV之间差距的54.5%和25.7%。在LLaMA-3.1-8B上128–1024预算扫描中，PolyKV相较于最强基线持续提升1.7%–6.4%。

## 2 相关工作

### 2.1 KV缓存驱逐与Transformer异质性

KV缓存驱逐通过仅保留历史键值状态的一个子集来减少推理内存。现有方法可以理解为不同的token驱逐模式，每种模式由不同的驱逐模式驱动。累积注意力方法使用累积的注意力分数对历史token进行评分，并保护一个最近的窗口(Zhang等人，2023 (https://arxiv.org/html/2606.15157#bib.bib19)；2024 (https://arxiv.org/html/2606.15157#bib.bib18)；Wan等人，2025 (https://arxiv.org/html/2606.15157#bib.bib15))。近期查询窗口方法则使用最近查询窗口的注意力对历史token进行评分，可选地应用池化或平滑，并保留前k个历史token以及最近的token(Li等人，2024 (https://arxiv.org/html/2606.15157#bib.bib7)；Cai等人，2024 (https://arxiv.org/html/2606.15157#bib.bib1))。持久性方法识别那些在最近查询步骤中反复超过注意力阈值的token，从而偏好重要性随时间保持稳定的token(Liu等人，2023 (https://arxiv.org/html/2606.15157#bib.bib8))。StreamingLLM代表另一种驱逐模式。它不是依赖基于注意力的显著性分数，而是根据它们的位置保留固定的初始汇聚token和最近的token(Xiao等人，2024b (https://arxiv.org/html/2606.15157#bib.bib17))。TOVA使用当前最后一个查询token的注意力来选择全局前k个缓存的token，使驱逐依赖于即时的解码状态(Oren等人，2024 (https://arxiv.org/html/2606.15157#bib.bib10))。

这些方法依赖于不同的驱逐模式，但它们都向模型中的每一层应用相同的策略。这种统一处理是粗糙的，因为Transformer块是异质的。先前的分析表明，LLM各层在表示抽象、注意力模式、上下文角色以及对剪枝的敏感性方面存在差异(Tenney等人，2019 (https://arxiv.org/html/2606.15157#bib.bib13)；Jawahar等人，2019 (https://arxiv.org/html/2606.15157#bib.bib6)；Clark等人，2019 (https://arxiv.org/html/2606.15157#bib.bib2)；Voita等人，2019 (https://arxiv.org/html/2606.15157#bib.bib14)；Michel等人，2019 (https://arxiv.org/html/2606.15157#bib.bib9))。对于KV缓存驱逐，这表明不同层可能偏好不同的驱逐偏好和不同的驱逐缓存预算。层异质性因此将驱逐转化为两个相关的设计问题。第一个是每层应该使用哪种驱逐模式，第二个是每层应该获得多少缓存预算。

### 2.2 自适应与非均匀KV缓存驱逐

自适应和非均匀的KV缓存驱逐方法，挑战了所有层应具有相同缓存预算的假设。有些方法在模型深度上分配不同的缓存预算，为不同层分配递减或异构的预算(Ge等人，2024 (https://arxiv.org/html/2606.15157#bib.bib5)；Cai等人，2024 (https://arxiv.org/html/2606.15157#bib.bib1)；Feng等人，2025 (https://arxiv.org/html/2606.15157#bib.bib3)；Fu等人，2025 (https://arxiv.org/html/2606.15157#bib.bib4))。其他方法使用源自提示、任务或生成时间统计的启发式指标自适应缓存预算(Zhou等人，2025 (https://arxiv.org/html/2606.15157#bib.bib20)；Shen等人，2025 (https://arxiv.org/html/2606.15157#bib.bib11)；Tang等人，2025 (https://arxiv.org/html/2606.15157#bib.bib12)；Xiao等人，2024a (https://arxiv.org/html/2606.15157#bib.bib16))。这些工作共同表明，均匀的缓存预算通常次优，定制化的预算可以提升性能。

这些方法探索了异构预算的重要性，但它们的设计空间仍然有限。许多自适应方法主要调整分配多少缓存预算，而忽略了应给每层分配哪种驱逐模式。PolyKV采取互补的方法，沿两个维度扩展设计空间：驱逐模式和缓存预算。它旨在在相同的总体缓存预算下为每一层识别更好的策略。这允许不同层遵循不同的驱逐规则并获得不同的缓存预算，从而更好地匹配每层的异构特征。

## 3 方法

参见标题图2：PolyKV架构。(A) 配置文件按层范围和推理阶段指定压缩计划。(B) 运行时钩子暴露不同KV缓存方法所需的执行点。(C) 组合包装器在阶段和层之间分派到不同方法。本节首先定义逐层驱逐设计空间，然后描述PolyKV如何利用层级别信号来缩减该设计空间。我们进一步介绍PolyKV的实现细节，整体架构如图2 (https://arxiv.org/html/2606.15157#S3.F2) 所示。

### 3.1 逐层KV驱逐表述

考虑一个具有LL个Transformer层和两个阶段Φ={pre,dec}\\Phi=\{\\mathrm\{pre\},\\mathrm\{dec\}\}的模型，分别对应预填充和解码。在生成步骤tt，层ll维护一个KV缓存。PolyKV表示一个阶段感知的逐层驱逐策略为

Π={({rl,ϕ}ϕ∈Φ,{bl,ϕ}ϕ∈Φ,)}l=1L,\\Pi=\\left\\\{\\left\(\\\{r\_\{l,\\phi\}\\\}\_\{\\phi\\in\\Phi\},\\\{b\_\{l,\\phi\}\\\}\_\{\\phi\\in\\Phi\},\\right\)\\right\\\}\_\{l=1\}^\{L\},其中rl,ϕ∈Rr\_\{l,\\phi\}\\in\\mathcal\{R\}是在阶段ϕ\\phi期间分配给层ll的驱逐模式，bl,ϕ∈Z+b\_\{l,\\phi\}\\in\\mathbb\{Z\}\_\{\+\}是分配给层ll的预算。

总KV预算受约束于

∑l=1Lbl,ϕ=Bϕ,bl,ϕ≥bmin。\\sum\_\{l=1\}^\{L\}b\_\{l,\\phi\}=B\_\{\\phi\},\\qquad b\_\{l,\\phi\}\\geq b\_\{\\min\}.\(1\)其中BB匹配相应基线使用的总预算，bminb\_\{\\min\}是每层的最小预算大小。等价地，如果所有层在平均预算b¯\\bar\{b\}下比较，则B=Lb¯B=L\\bar\{b\}。这里，blb\_\{l\}表示层ll的缓存预算。

在生成步骤tt，层ll的实际缓存大小最多为min⁡(t,bl)\\min\(t,b\_\{l\}\)。FullKV被视为一种特殊的驱逐模式，其预算等于当前完整上下文长度。因此，任何分配了FullKV的层都不会豁免预算核算；其完整的KV成本包含在总预算中。

在传统的模型范围驱逐方法中，所有层具有相同的驱逐模式和相同的缓存预算。自适应驱逐方法具有不同的缓存预算，但它们仍然遵循共同的驱逐模式。PolyKV通过允许驱逐模式和缓存预算跨层变化来扩大设计空间。

为了实例化这个设计空间，我们在PolyKV框架内纳入一组代表性的KV驱逐基线作为驱逐模式。PolyKV的目标是探索，当以逐层方式分配并且具有不同的缓存预算时，现有驱逐模式能否被更有效地使用。

算法 1 PolyKV 策略构建1:候选保留模式R\\mathcal\{R\}; 阶段