用轻量代理模型加速大模型长文本推理,省内存还快
ProxyKV: Cross-Model Proxy Pruning for Efficient Long-Context LLM Inference

- 用小模型异步评估大模型关键记忆,避免直接计算开销
- 在长文本任务中恢复98.7%精度,预填充速度提升3.21倍
- 适合追求高效长上下文推理的AI部署场景
大型语言模型(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倍的预填充加速(双卡环境;单卡共享环境下约1.5倍),并在Qwen-2.5-7B上支持长达17万词元的上下文仍保持加速效果。
原文摘要 · Abstract (English)
Efficient long-context inference in Large Language Models (LLMs) is severely constrained by the Key-Value (KV) cache memory wall, yet existing pruning methods force a choice between low-latency heuristics that sacrifice precision and high-precision reconstruction methods that incur prohibitive prefilling overhead. To bridge this scoring-cost--accuracy gap, we propose ProxyKV, a cross-model proxy pruning framework that offloads importance scoring to a lightweight intra-family Small-Model Proxy executed asynchronously to the Large-Model Target. To bridge the architectural gap between heterogeneous models, we design the HybridAxialMapper, which disentangles temporal feature extraction from cross-head alignment, together with a Multi-Granularity Hybrid Loss that shifts the learning objective from rigid regression to relative ranking consistency. Across the Llama-3.1, Qwen-2.5, and Qwen-3 families spanning targets from 7B up to 32B parameters on LongBench, SCBench, and RULER, ProxyKV matches KVZip on aggregate (recovering $\sim$$98.7\%$ of its mean accuracy) while delivering up to a $3.21\times$ prefilling speedup on Llama-3.1-8B (dual-GPU; $\sim$$1.5\times$ shared single-GPU) and sustaining the speedup at contexts up to 170k tokens on Qwen-2.5-7B.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。