ProxyKV: 跨模型代理剪枝实现高效长上下文LLM推理

# 跨模型代理剪枝用于高效长上下文LLM推理
来源: https://arxiv.org/html/2605.16360

###### 摘要

大型语言模型（LLM）中的高效长上下文推理受到键值（KV）缓存内存瓶颈的严重制约，然而现有的剪枝方法迫使人们在低延迟启发式方法（牺牲精度）和高精度重建方法（带来高昂的预填充开销）之间做出选择。为了弥合这种评分成本与精度之间的差距，我们提出了ProxyKV，一种跨模型代理剪枝框架，它将重要性评分卸载给一个轻量级的家族内小模型代理，该代理异步于大模型目标运行。为了弥合异构模型之间的架构差距，我们设计了HybridAxialMapper，它将时间特征提取与跨头对齐解耦，同时引入了多粒度混合损失，将学习目标从严格的回归转变为相对排序一致性。在跨Llama-3.1、Qwen-2.5和Qwen-3家族（目标参数从7B到32B）的LongBench、SCBench和RULER基准测试中，ProxyKV在整体上匹配了KVZip（恢复了其平均精度的约98.7%），同时在Llama-3.1-8B上实现了高达3.21倍的预填充加速（双GPU；共享单GPU约1.5倍），并在Qwen-2.5-7B上保持了高达170k token上下文长度的加速。

## 1 引言

处理长上下文序列并进行推理的能力已成为现代大型语言模型（LLM）的关键需求，其应用范围从全仓库代码分析（Jiang等人，2024）到复杂的多跳文档推理（Minaee等人，2024）。为了实现高效的自回归生成，Transformer架构利用键值（KV）缓存来存储过去的token激活，有效地消除了冗余计算（Kwon等人，2023）。然而，随着序列长度N的增加，这种不可或缺的机制引入了严重的内存瓶颈。KV缓存占用的线性增长常常超过单个GPU的物理内存容量，导致频繁的I/O开销和整体推理延迟的急剧增加（Li等人，2024a）。

KV缓存剪枝应运而生以缓解这一瓶颈，但现有的查询无关方案在评分效率与剪枝精度之间面临权衡：启发式方法如H2O（Zhang等人，2023）和SnapKV（Li等人，2024c）速度很快，但其局部观测窗口会遗漏全局语义依赖；而基于重建的方法如KVZip（Kim等人，2025）通过全局上下文重建实现了高精度，但代价是在目标模型上进行了一次代价高昂的二次预填充过程。

如图1 (a, b)所示，这两种范式锚定在效率-精度谱的两端。我们假设，受IAM（Zhao等人，2025a）和SmallKV（Zhao等人，2025b）观察到的家族内注意力相关性的启发，可以在不依赖目标模型上大量计算的情况下实现重建方法的高精度评分；然而，现有的静态头对头对齐过于刚性，无法弥合头数和层深度方面的结构差异。

我们提出ProxyKV（图1 (c)），一种跨模型代理剪枝框架，它将评分卸载给一个轻量级的家族内小模型代理，该代理异步于目标的临界路径执行。为了处理架构差异，我们的HybridAxialMapper将时间特征提取与跨头对齐解耦，并且多粒度混合损失将学习目标从严格的值回归转变为相对排序一致性——更直接地针对Top-K剪枝质量进行对齐。我们的贡献：(i) 一个基于异步代理的剪枝框架，消除了目标临界路径上的评分瓶颈，在Llama-3.1-8B上实现了高达3.21倍的预填充加速；(ii) HybridAxialMapper和一个五项多粒度混合损失，共同弥合了异构模型之间的结构差距；(iii) 跨Llama-3.1、Qwen-2.5和Qwen-3家族（7B–32B目标）的实证研究表明，ProxyKV恢复了约98.7%的KVZip oracle性能，并在长达170k token的上下文长度中保持了加速。

## 2 相关工作

请参考图注：图1：三种KV缓存剪枝范式：SnapKV (a) 启发式，KVZip (b) 重建，ProxyKV (c) 异步代理。

#### 启发式和架构剪枝。

基于规则的方法通过局部模式识别非必要token：StreamingLLM（Xiao等人，2023）保留注意力汇聚点；H2O（Zhang等人，2023）、Scissorhands（Liu等人，2023）和AhaKV（Gu等人，2025）使用累积分数或最近注意力模式。架构变体进一步利用结构：FastGen（Ge等人，2023）应用头级差异化策略，PyramidKV（Cai等人，2024）分配层级预算，SnapKV（Li等人，2024c）通过提示结束观测窗口聚类键。ProxyKV则蒸馏全局重建信号，生成更细粒度的剪枝，在查询无关场景中表现稳健。

#### 学习和代理剪枝。

数据驱动的策略追求自适应计算：CoT-Influx（Huang等人，2024）使用强化学习进行粗细粒度选择，而GateSkip（Laitenberger等人，2025）和LTP（Kim等人，2022）引入了可微分门控。最近，KVzap（Jegou and Jeblick, 2026）通过每层隐藏状态的代理近似重建oracle；我们不将其作为基线，因为 (i) 其训练成本远高于ProxyKV，(ii) 其发布的实现采用了与KVZip oracle不同的token选择标准，因此无法直接比较，(iii) 它不支持在固定保留/压缩比下进行评估，而这是LongBench、SCBench和RULER上的标准协议。与这些在目标执行内部运行代理的方法不同，ProxyKV将评分完全卸载到推理临界路径之外，且训练开销更低。

#### 跨模型KV缓存对齐。

利用家族内相关性，IAM（Zhao等人，2025a）跨规模复用注意力相似性，SmallKV（Zhao等人，2025b）用较小模型补偿剪枝损失，SpeContext（Xu等人，2025）对齐检索以进行推测性预取。这些方法依赖于固定对之间的刚性头级或层级映射；ProxyKV训练一个*可学习的跨轴*映射，联合建模时间上下文和头轴对齐，能够在具有不同头数和深度的架构之间迁移。ProxyKV在概念上也与推测性解码（Kwon等人，2023）相邻，但其推测的是预填充阶段的*重要性分数分布*，而非生成阶段的token。

## 3 系统描述与问题形式化

#### KV缓存剪枝背景。

Transformer LLM解码会缓存过去token的键/值张量，以避免冗余计算。随着上下文长度N的增长，缓存成为内存瓶颈，这促使进行剪枝，只保留信息最丰富的KV对，并在*效率*（低评分延迟）与*精度*（长上下文推理质量）之间取得平衡。

#### 系统概述。

ProxyKV是一个部署时的流水线，包含三个组件：大模型目标M_l（拥有预填充临界路径和整个解码阶段）、一个家族内小模型代理M_s（在单独的执行流上异步运行以提取跨头注意力特征X）、以及HybridAxialMapper Φ_θ（将X映射到目标对齐的重要性分数Ŷ）。具体来说，X是在每个代理层上，沿查询轴累积的softmax后注意力概率，即 X_{b,l,h,n} = ∑_{q=1}^N Softmax(QK^⊤/√D)_{b,l,h,q,n}，它将一个二次的N×N注意力图压缩成每个（批次、层、头）长度为N的单一每键重要性向量。在推理时，输入上下文同时分发到两个模型；Φ_θ接收代理特征并生成目标形状的分数，驱动Top-K剪枝掩码，之后目标在*没有*代理参与其临界路径的情况下进行解码。该流水线支持两种部署模式（图1(c)）：双GPU模式，即目标/代理在独立设备上并行运行；单GPU模式，它们通过独立的CUDA流共享同一设备；代理的KV缓存在预填充结束时释放，因此代理侧的内存开销是短暂的（图12）。

#### 问题形式化。

设M_l有L_l层，每层H_l个头，输入长度N，每层KV张量 K, V ∈ R^{B×H_l×N×D}；设M_s有L_s层，每层H_s ≪ H_l个头。基于重建的oracle（例如KVZip）通过M_l上的二次教师强制预填充推导出真实注意力分数 Y ∈ R^{B×L_l×H_l×N}（图1(b)），这会增加TTFT。ProxyKV则通过可学习的跨轴映射 Ŷ = Φ_θ(X) ≈ Y 来近似Y，其中 X ∈ R^{B×L_s×H_s×N}，并通过在保留比ρ∈(0,1]下对Ŷ进行Top-K阈值化生成二元掩码 M ∈ {0,1}^{B×L_l×H_l×N}。映射器Φ_θ根据第4.1节的层配对计划，独立应用于每对（代理层，目标层）；因此以下所有每层公式描述了对单个这样的对的操作；我们删除了X, Y, M中的显式L索引以简化符号。我们用*注意力分数*指代Y，用*重要性分数*指代Ŷ。训练目标为：

θ* = arg min_θ E_(X,Y)[L(Φ_θ(X), Y)], (1)
满足 t_score(M_s, Φ_θ) ≤ t_prefill(M_l), family(M_s) = family(M_l),

其中 t_score(M_s, Φ_θ) = t_proxy(M_s) + t_Φ_θ 是端到端运行代理和映射器的总时间，异步预算约束确保两者都不在目标的临界路径上（否则加速会消失）；家族内约束保留了Φ_θ利用的逐层注意力相关性。损失L是第4.2节形式化的多粒度混合损失；部署度量是在相同ρ下对KVZip oracle的任务精度恢复。

## 4 方法

请参考图注：图2：ProxyKV概述：一个异步的小模型代理M_s将特征输入HybridAxialMapper，后者为大模型目标M_l生成对齐的重要性分数Ŷ，无需二次预填充过程。

基于第3节系统流水线的构建，HybridAxialMapper Φ_θ通过一个三阶段流水线处理代理特征X——(1) *时间特征提取*：一维卷积主干，(2) *时间轴上下文编码器*：Transformer编码器，(3) *头轴交叉注意力*：可学习的目标查询（图2）——预测的分数在指定预算B下生成二元掩码M。以下两个小节详细介绍映射器架构和多粒度混合损失。

### 4.1 HybridAxialMapper架构

#### 设计原理。

跨模型对齐必须协调两个耦合的轴：一个*时间*轴，沿此轴token显著性随序列演变；以及一个*头*轴，沿此轴代理和目标注意力模式以不同粒度划分（H_s ≪ H_l）。先前的静态头对头方案（Zhao等人，2025a, b）通过固定代理和目标头之间的排列来纠缠这些轴，这阻碍了向具有不同头数或层深度的架构迁移。我们的设计则将映射分解为一个产生与头无关的潜在表示的时间阶段和一个学习代理到目标聚合的头轴阶段，因此相同的配方适用于任何家族内对。在层轴上，我们采用深度比例配对计划，将目标层 l_l ∈ [1, L_l] 与代理层 l_s = ⌈l_l · L_s / L_l⌉ 配对，并为每一对运行Φ_θ的一次独立前向传播；映射器参数在所有L_l对上共享，因此Φ_θ仍然是一个单一网络，其前向传播每次请求被调用L_l次，其输出沿层轴堆叠，恢复第3节中的B×L_l×H_l×N分数张量。

三阶段流水线将时间依赖提取与特定头的对齐解耦。

#### 阶段1：时间特征提取。

一维卷积主干（两个堆叠的Conv1D层，k=3, p=1，每层后接BatchNorm和GELU）将原始代理特征