换模型不必重算 prefill:英伟达用线性映射搬 KV cache,提速最高 25 倍

Cross-Model KV Cache Transfer in LLM Families: A Closed-Form Linear Mapping for Prefill Reuse

Taekyung Heo, Rasoul Shafipour, Ritchie Zhao, Maximilian Golub, Mohammad Mahdi Kamani, Ritika Borkar, Makesh Tarun Chandran, Pantea Zardoshti, Bita Darvish Rouhani

cs.LG

2026-08-05

同家族不同规模模型的 KV cache 之间存在强线性关系,英伟达用无需梯度训练的 ridge 映射跨模型搬运、跳过 prefill,六组对里四组保住 73–98% 精度,prefill 提速 2.7–25 倍。

这篇在解决什么

线上服务常在同家族不同规模的模型之间切换:便宜的用小模型、难的升级到大模型(cost-quality cascading),对话中途换模型,或按 query 路由到不同档位。每次切换,接收方都得把整段上下文重新过一遍所有层、重算一遍 KV cache,这段计算叫 prefill。上下文一长,prefill 就是最贵的那段。

此前想复用 KV cache 的工作要么得给每对模型单独训练融合网络(C2C),要么学一组 adapter 把两边投影进共享隐空间(LatentAlign),要么干脆要求两个模型架构完全一致(DroidSpeak)。这篇想绕开梯度训练,只靠一个闭式解的线性映射把源模型的 KV cache 搬到目标模型上,直接跳过 prefill。

方法

核心是一个逐 head 的 ridge 回归映射器,三步:

前提是 matched-KV:源和目标模型的 KV head 数、每个 head 的维度必须一致(GQA 结构要对齐),层数和参数量可以不同。所以它只管同家族内部,跨家族不行。

线性为什么够用:作者发现同家族不同规模模型的 KV cache 之间存在很强的线性关系。在 Qwen3 14B→32B 上,单层源 KV 能解释目标 key 56% 的方差、value 32%;多用几层源,解释度升到 79%(key)和 65%(value),最高的单个格子 R² 达到 0.81。

结果

六个模型对、三个家族(Qwen3、Llama 3.1、Ministral 3),五个 benchmark 上的精度保留率(目标模型独立 prefill 的精度记为 100%):

模型对平均保留HellaSwagMMLUGSM8K
Qwen3 14B→32B97.5%97.6%95.0%95.6%
Qwen3 8B→32B87.5%95.2%88.5%68.8%
Llama 3.1 8B→70B72.8%94.4%73.3%18.2%
Ministral 3B→8B76.2%93.3%69.4%36.6%
Ministral 3B→14B44.2%68.0%32.0%3.2%
Ministral 8B→14B41.6%58.7%32.7%1.6%

四组保住了 73–98% 的平均精度,两组(Ministral 的 3B→14B、8B→14B)掉到 42–44%。最扎眼的是 GSM8K:即便算成功的几组,Llama 3.1 8B→70B 的保留率也只剩 18.2%,只发挥出目标模型原本数学能力的不到五分之一。选择题、常识这类对激活精度不敏感的任务对 KV 搬运很宽容,一到需要精确推理的数学题就露馅。

速度:在 Qwen3 14B↔32B 上,小→大方向 prefill 延迟降 4–25 倍,32K token 时从 6.98 秒压到 0.28 秒;大→小方向降 3–7 倍。70 个测试格子里映射器每一格都快。代价是映射器本身不小,1.0–3.4B 参数、4–12GB 存储。

非线性兜底:那两组失败的对,把 ridge 换成一个两层 MLP(每层 1024 单元),HellaSwag 保留率从 68.0%/58.7% 拉回到 92.3%/95.5%,涨 24–37 个百分点。有意思的是,ridge 本来就行的对换 MLP 反而略差。作者的解释是 MLP 把残差误差从 attention 敏感的子空间挪走了。

为什么重要

跑模型级联的人(小模型接客、难题升级大模型,或按 query 路由分发)能用这套东西省掉换模型时的整段 prefill 重算。对长上下文场景(几万 token),省的是几秒到十几秒延迟和对应算力。

「重要」要打折看。三个限制卡得很死:只跨同家族(必须 matched-KV),只在 dense full-attention 上验证过,数学推理类任务即使成功也明显掉点。真正能放心用,是同家族换规模、且任务对激活精度不敏感的场景(分类、摘要、常识问答、多轮对话)。多轮这块作者在 CoQA 上验证过,Qwen3 14B↔32B 跑 10 轮,小→大方向累积偏差 1.7 个百分点,大→小方向每轮线性漂移 0.33 个百分点,十轮内不崩,但超长会话里大→小的漂移会持续累积。

局限与存疑

作者自己承认:标定只用 FineWeb-Edu 一个领域,换 CodeAlpaca 在 HellaSwag 上掉 5.24 个百分点(Wikipedia 还在噪声内);只测了 matched-KV 的对,mismatched 的没碰;只覆盖 dense full-attention,带 sliding window / sparse attention 的混合架构在范围外。

读下来还有两点存疑。一是可预测性最强的拟合指标(R²)反而预测不了下游效果:Llama 3.1 8B→70B 拟合得很好,下游却退化。作者发现 attention-output 余弦相似度才预测得准(Pearson r=+0.57,R² 是 −0.20),说明「误差落在哪个子空间」比「误差多大」更要命。这个观察关键,但也意味着没法在标定阶段就判断一对模型值不值得做,得跑完下游才知道。二是 headline「保住 73–98% 精度」在数学推理上是过度乐观的概括,表里 GSM8K 全面塌方,落地得按任务类型分别看。

术语

原文与代码

社区讨论

相关论文

全部论文解读