跨模型KV Cache迁移：闭式线性映射实现Prefill复用

如果你最近在做 LLM 推理优化，应该绕不开 KV Cache 这个词。我仔细读了一个论文标题：Cross-Model KV Cache Transfer in LLM Families: A Closed-Form Linear Mapping for Prefill Reuse。这个题目表达得很直接：把同一个模型家族里某个模型 prefill 阶段已经算好的 KV Cache，通过一个闭式线性映射传给同族另一个模型，省掉一次完整的 prefill 计算。它的价值一眼就能看出来：长上下文场景下 prefill 的时间成本很高，如果同族模型之间能复用 prefill 结果，多模型调度、A/B 测试、Agent 编排这些场景就能省下大量算力。

不过先说明一点：我目前看到的只是论文标题，不是完整实验报告。所以下面这篇内容，是把这条技术路线拆成“它解决什么问题、为什么可能成立、怎么验证、有哪些坑、什么时候值得用”来解读。文章给出的更多是通用验证流程和工程判断方法，具体效果如何，要拿你手上的模型实际跑一遍才能下结论。

1. 先搞清楚 KV Cache 和 Prefill Reuse 在解决什么问题

1.1 为什么 prefill 阶段会成为瓶颈

生成式 LLM 的推理过程一般分成两个阶段。第一个阶段是 prefill，也称预填充。用户输入一整段 prompt 后，模型并行处理所有 token，为每个 token 计算出 Key 和 Value，并存入 KV Cache。第二个阶段是 decode，也就是逐 token 生成。每一步生成新 token 时，模型都会从 KV Cache 里读取历史 token 的 K、V，避免重新计算已经处理过的内容。

prefill 是计算密集型阶段，decode 是访存密集型阶段。长 prompt 下 prefill 的计算量会变得非常大。比如一份几千 token 的文档，模型第一次完整读取时，要跑一遍整个自注意力计算，并把中间结果存下来。如果每次用户提问、每次 Agent 调用工具、每次模型版本切换都要重新跑一次 prefill，GPU 算力会被吃掉一大块。

KV Cache 就是为了缓解这个问题。常规做法里，同一个模型、同一个输入序列的 KV Cache 可以复用。但一旦换成另一个模型，缓存就失效了。原因很简单：KV Cache 和模型权重是绑定的，模型 A 的 K、V 数值不能直接喂给模型 B。Cross-Model KV Cache Transfer 想做的，就是打破“缓存只能同模型复用”这个限制。

1.2 KV Cache 可复用的边界在哪里

先看常规复用边界，这样容易对比出跨模型复用的难度。

• 同一个模型、同一个配置、同一段输入：KV Cache 可以直接复用。
• 同一个模型、不同 batch，但存在公共前缀：公共前缀部分的 KV 可以复用。
• 同一个模型、不同采样参数：采样参数不影响 KV Cache，输入一致就能复用。
• 同一个模型、不同 batch 顺序：需要按 token id 对齐，理论上也能复用。

跨模型场景就没有这么多默认支持了。第一个障碍是维度。source 模型和 target 模型的 hidden size、num_key_value_heads、head_dim 可能不同。第二个障碍是语义空间。就算两个模型的 hidden size 一样，它们内部表示所在的坐标系也可能完全错位。第三个障碍是注意力结构。MHA、GQA、MLA 之间存在实质差异，K 和 V 的组织方式不一样，位置编码也可能不同，直接搬运自然不可行。

所以跨模型 KV Cache 迁移，本质上需要完成一次“表示空间变换”。把 source 模型内部空间里的 KV 张量，变换到 target 模型内部空间里一个可用的状态。这篇论文标题给出的工具，是一个 closed-form linear mapping，也就是闭式线性映射。

1.3 跨模型迁移的两种路线

要解决跨模型 KV 迁移，业界大致会想到两条路线。

第一条是训练一个神经网络适配器。拿大量文本让 source 模型和 target 模型分别生成 KV Cache，然后训练一个小网络把 source 侧映射到 target 侧。这种方案理论上能抓住非线性关系，但缺点也明显：需要训练数据、训练时间、调参成本，而且额外网络会在推理链路上增加延迟。对“省 prefill”这件事来说，引入一个复杂适配器有点得不偿失。

第二条就是闭式线性映射。它的核心假设非常直接：target 模型某一层的 KV Cache，约等于 source 模型同一层 KV Cache 乘以一个线性矩阵 W。求解 W 不需要反向传播，不需要多次迭代，只需要把收集好的 X、Y 矩阵丢进最小二乘求解器，一次就能得到解析解，所以叫 closed-form。

这条路线之所以值得关注，是因为它很便宜。拟合一个 W 的代价远低于训练一个适配器，部署时也只要对 KV Cache 做一次矩阵乘法。代价是它依赖线性假设。如果两个模型的空间变换关系已经很接近线性，那这个方案能拿到的收益非常明显；如果假设不成立，效果就会打折扣。所以接下来的核心问题变成：同族模型之间，这个线性假设到底成不成立。

2. 同族模型之间的关系决定线性映射是否可行

2.1 什么是 LLM 家族，同族模型共享了什么

“LLM Families”通常指同一个基础模型衍生出来的一系列变体。比如一个开源社区放出某个 base 模型，接着有人做了 instruction tuning 得到 instruct 版，有人做了 RLHF 或 DPO 对齐得到 chat 版，还有人在此基础上做量化和蒸馏，形成 8-bit、4-bit 版本或更小的蒸馏模型。这些模型可能权重规模不同、输出风格不同、指令遵循能力不同，但它们共享同一个 tokenizer、同一套词表，甚至大部分网络结构都来自同一个初始权重。

这个“同源”属性非常关键。因为同族模型在训练起点上一致，后续微调只是对权重做有限幅度的改动。这会让它们的内部语义空间不会完全漂移。两个模型在看到同一段文本时，注意力层的表示虽然数值不同，但很可能存在一种稳定的空间对应关系。

反过来看跨家族模型，比如 Llama 家族和 Qwen 家族之间，tokenizer 不同、词表不同、架构细节不同，连输入文本切出来的 token id 序列都对不齐。这时候想直接做 KV 迁移，首先要解决 token 对齐问题，线性映射的适用性会大幅下降。所以论文标题里强调 in LLM Families，是一个非常强的边界条件。

2.2 闭式线性映射的核心假设

可以用公式把这个假设写得更清楚。

假设 source 模型某一层的 KV Cache 张量是 X，target 模型同一层的 KV Cache 张量是 Y。闭式线性映射假设存在一个矩阵 W，使得：

Y ≈ X W^T

如果 X 是 n 行 d_s 维，Y 是 n 行 d_t 维，那么 W 就是 d_t 行 d_s 列的矩阵。要求解 W，最自然的方式是最小二乘：

W = (X^T X)^{-1} X^T Y

更稳定的做法是直接用 SVD 或 PyTorch 里的 lstsq，避免因 X^T X 条件数过大导致数值不稳定。

这个假设为什么在同族模型里可能成立？因为同族模型面对相同的输入分布，优化的是相同的语言建模目标，只是微调策略不同。大量模型表征研究工作发现，不同初始化或不同训练方式得到的模型，其内部表征之间往往存在线性或近似线性的变换关系。KV Cache 虽然不是最终输出，但它也是模型计算图里的一环，同样可能满足这个性质。

但“可能成立”不等于“一定成立”。层跟层之间的线性度很可能不同。有的层本身语义稳定，映射误差很小；有的层对微调非常敏感，线性拟合误差很大。所以实际验证时，不能只测一层，更不能直接拟合一个全局大矩阵。

2.3 从单层映射到整体迁移

工程落地时，不要试图一次性拟合一个覆盖所有层的联合映射。正确路线是逐层处理。

首选做法是先测单层映射效果。比如 source 模型和 target 模型都是 12 层结构，那么每一层单独拟合一个 W_l。这样能直观看到哪些层容易映射、哪些层难以映射。如果某些层的余弦相似度很高，说明 prefill 复用可以从这些层开始；如果某些层完全不行，就考虑跳过或只用真实 KV。

这里要提一个容易忽略的点：prefill 复用只省 prefill，decode 阶段仍然由 target 模型自己的权重逐渐生成。即使 KV Cache 全部迁移成功，每生成一个新 token，target 模型还是会做一次前向计算。所以这个方案针对的是“首 token 延迟”和“重复 prefill 计算量”，不是把 decode 阶段也省掉。理解了这一点，才不会对这个方案产生不切实际的期待。

3. 从论文标题到可复现流程：先跑通单层映射

3.1 准备一个最小验证环境

如果你想验证“闭式线性映射到底能不能用”，我建议先搭一个最小环境，不要一开始就上生产级别的大模型。

环境依赖大致是：

• Python 3.10 以上
• PyTorch 2.x
• transformers、accelerate、datasets
• 一张 16GB 以上显存的 GPU；如果只测 1B 模型，普通消费级显卡也够用
• 如果测 7B 模型并限制上下文在 1024 以内，16GB 到 24GB 都可以

模型对的选择更重要。我建议优先选同一家族的小模型，比如同一个基础模型的 base 版和 instruct 版，或者一个 1B 原版和它的量化版。优先测 1B 到 3B 模型，原因是跑得快、样本量大、定位问题快。不要一上来就开 70B。

可以按这个配置来安排测试：

3.2 让两个模型处理完全相同的 token 序列

这一步是很多复现失败的最常见原因。线性映射成立的前提是，X 和 Y 在维度上对齐、在语义位置上对齐。

对齐 token 序列的做法：

• 尽量使用同一个 tokenizer。
• 两个模型处理完全相同的文本。
• 固定 max_length，固定 padding 策略，最好用 right padding 或 per-sample 单条处理。
• 不要混用 truncation 策略，否则序列长度不一致，张量形状直接对不上。

如果 source 和 target 的 tokenizer 不一致，那 token id 序列大概率不同。即使文本相同，切分结果也会不同，KV Cache 的位置语义完全错位。这时候线性映射没有意义，先处理 tokenizer 对齐再继续。

3.3 收集 source 和 target 的 KV Cache

收集 KV Cache 的代码逻辑很简单，主要留意 transformers 版本差异。老版本里 past_key_values 是 tuple of tuple，新版本可能是集中管理的 tensor。统一处理方式是把每一层的 K 和 V 转成独立 tensor，记录 shape。

def extract_kv_list(model, inputs, layer_idx): with torch.no_grad(): outputs = model(**inputs, use_cache=True) pkv = outputs.past_key_values # 根据实际版本解析，这里假设 K 和 V 是第一个 tuple 的第 0、1 项 K, V = pkv[layer_idx][0].detach().float(), pkv[layer_idx][1].detach().float() return K, V

返回的 K 和 V，shape 通常是 [batch, heads, seq_len, head_dim]。如果 batch 内每条样本长度一致，可以直接处理；如果长度不一致，模型可能已经做了 padding，需要额外处理 mask，避免把 padding 位置也当成有效数据参与拟合。

3.4 用最小二乘拟合线性映射

把 source 侧和 target 侧的数据都收集完之后，将它们转换成二维矩阵。一个常见做法是把每条样本的 K、V 拼接或分开展开。比如把 K 展开成 [batch * heads * seq_len, head_dim]，再把 X 拼成一个大矩阵，Y 拼成对应的大矩阵，然后求解。

import torch def fit_linear_mapping(X, Y): # X: [sample_count, source_dim] # Y: [sample_count, target_dim] W, _ = torch.linalg.lstsq(X, Y) return W # 示例：将 source_kv 和 target_kv 展开 X = torch.cat([k.reshape(-1, src_dim) for k in source_kv_list], dim=0) Y = torch.cat([k.reshape(-1, tgt_dim) for k in target_kv_list], dim=0) W = fit_linear_mapping(X, Y)

为什么要用 lstsq 而不是直接求逆？因为直接求逆要吃 (X^T X) 矩阵的条件数，而矩阵分解会更稳定。如果样本量不够，可以加一个很小的岭正则项，变成岭回归：

W = (X^T X + λI)^{-1} X^T Y

这样能避免过拟合。需要注意的是，W 的求解用的是浮点精度。如果 X、Y 本身是 fp16，建议先转成 fp32 再拟合，否则数值误差会被求逆过程放大。

3.5 评估单层映射效果

拟合完 W，不能只看训练集上的损失，否则很容易被高拟合分数误导。我一般分两个层次评估。

第一层看映射拟合质量。将 W 作用到 source KV 上，跟 target 的真实 KV 算余弦相似度和 MSE。如果某个 layer 的相似度能到 0.95 以上，说明映射方向比较乐观。如果只有 0.5，说明该层线性假设偏弱。

第二层看生成质量。这里更接近真实使用方式：把映射后的 KV 作为 past_key_values 喂给 target 模型，然后只对最后一个 token 做前向，得到 next token logits。

def eval_pseudo_kv(target_model, input_ids, pseudo_kv): last_ids = input_ids[:, -1:] with torch.no_grad(): out = target_model(last_ids, past_key_values=pseudo_kv, use_cache=True) return out.logits[:, -1, :]

然后对比目标模型正常 prefill 后得到的 logits，可以用 KL 散度、余弦相似度或 top-5 重合率。这一步能直接看出迁移后的 KV 到底能不能支撑生成。不过要特别强调：测试时建议先只替换单层 KV，其他层保留真实 KV。这样如果效果崩了，能快速定位是这一层的映射问题，还是整体替换逻辑有问题。

4. 不要把线性映射当成万能转换器：边界、风险与排查

4.1 KV 头数和 head_dim 不一致怎么办

同族模型并不等于结构完全一致。比如 source 是 GQA，target 是 MHA，KV 头数就不同。GQA 的 KV 头数少，MHA 的 KV 头数多，直接拼接会导致维度不匹配。

常规处理思路是先用广播或 repeat 把 GQA 的 KV 头展开到 MHA 对应分组，然后再拟合。如果是 MHA 到 GQA 的映射，可能需要把多个头合并或先做平均。这些都会影响映射精度，所以最容易出效果的模型对，往往是结构设计一致、只是微调方式不同的同族模型。

如果 head_dim 不同，比如 source 是 128，target 是 64，可以按 token 维度单独做线性映射，但需要更多样本才能拟合出稳定矩阵。这种场景下建议先做小规模实验，确认收益足够大再继续。

4.2 量化与 dtype 会影响映射矩阵稳定性

现在很多模型使用 fp16、bf16 或 int4/int8 量化部署。KV Cache 在推理时通常以半精度保存。这里存在一个很常见的坑：如果你直接用 fp16 的 K、V 去求最小二乘解，数值误差会被放大。求逆和 SVD 都对低精度数据相对敏感，最后得到的 W 可能偏差很大。

建议这样处理：

• 收集 KV Cache 时，统一转为 fp32。
• 在 fp32 精度下拟合 W。
• 应用映射时，再把映射后的 KV 转回目标模型实际使用的 dtype。
• 如果目标模型是量化模型，要确认它的 KV Cache 是否有额外的 scale 和 zero point，有的话先还原再映射，映射完再重新量化。

还要提醒一个容易忽略的问题：有些推理框架对外部传入的 past_key_values 支持不完整，比如服务化接口只允许模型自己维护 KV Cache，不接受外部注入。这种情况下，即使线性映射本身正确，工程上也很难接入。开始做方案评估前，先确认目标模型的 serving 框架是否开放了这个口子。

4.3 常见失败现象与排查链路

遇到结果不对，不要第一时间怀疑“线性映射不可行”。很多问题是数据对齐、参数配置或接口兼容性导致的。下面是我比较常用的排查顺序。

除此之外，还有一个必须先查的点：source 和 target 的 attention mask 是否一致。padding 位置不同，会导致 KV 的有效 token 数量不同，也会让线性映射的对应关系错乱。我见过不少项目，模型选得没问题，线性假设也成立，就是因为 padding 策略不统一，最后结果一塌糊涂。

排查顺序可以按这个链路来：先核对输入 token，再核对层索引和 shape，再核对 dtype，再核对 attention mask，再核对评估指标，最后才怀疑算法假设。

4.4 位置编码差异是隐藏变量

KV Cache 里的 K、V 已经包含了位置信息，特别是 RoPE 这类相对位置编码，会在 Q 和 K 上做旋转。如果 source 和 target 模型的 rope_theta 不同，或者位置编码实现细节不同，K 的数值含义会产生偏移。这样即使 token id 完全一致，同一位置的 K、V 也可能不在同一个坐标系里。

同族模型多数情况下会沿用相同的位置编码参数，但并不是绝对的。有些微调版本会为了长文本能力调整旋转基频。所以在收集 KV Cache 前，先确认两个模型的 config 里位置编码相关参数完全一致。如果不同，线性映射的难度会明显上升。

5. 工程上哪些场景值得做，哪些不值得做

5.1 值得做的高价值场景

第一个场景是同族模型的 A/B 测试。假设你在对比 base 模型、instruct 模型、量化模型在同一份长文档上的输出。常规做法是每个模型都跑一遍 prefill，计算量成倍增加。如果 KV Cache 能跨模型迁移，只需要 source 模型做一次 prefill，其它模型直接接收映射后的 KV，首 token 延迟和 GPU 占用都会大幅下降。

第二个场景是长文档多 Agent 分析。Agent 并行处理同一份材料时，经常会有多个模型分别理解全文。如果这几个模型属于同一家族，完全可以先把 source KV 提取出来，再分别映射给不同 target 模型。这比每个模型各自读取文档再 prefill 要高效得多。

第三个场景是推理引擎灰度切换。比如线上白天用原版高精度模型，夜间切到量化版降低资源成本。如果两个模型属于同族，共享同一套 KV Cache 转换链路，切换时就不需要把用户的整段历史上下文重新 prefill 一遍。

5.2 不值得做的低价值场景

第一类，跨家族模型。Llama 到 Qwen、Qwen 到 DeepSeek 这类场景，tokenizer、结构、训练数据差异都太大，线性映射基本不成立。即使强行做，也需要复杂的 token 对齐和层对齐，收益会被处理成本吃掉。

第二类，短 prompt 场景。如果用户输入只有几十个 token，prefill 消耗本来就小，省下的延迟几乎感知不到，还要承担映射质量风险和维护成本，性价比很低。

第三类，质量高度敏感的任务。代码生成、数学证明、合同审查这类任务，对输出质量波动容忍度很低。KV Cache 迁移本质上是近似，哪怕首 token 延迟降低了，如果生成结果出现细微偏差，代价可能远高于省下的一点算力。

第四类，权重频繁更新的场景。映射矩阵 W 和模型权重强相关。模型版本一更新，旧的 W 基本需要重新校准。如果模型每周都发布新版本，维护成本会快速上升。

5.3 工程落地前要处理的四件事

如果决定要做，我建议先把下面这四件事处理完。

第一，离线校准映射矩阵。不要用线上随机 prompt 在线拟合 W，这样不可控。准备一份覆盖目标场景的校验集，离线把所有层或目标层的 W 求好，再发布到线上。

第二，缓存键设计。跨模型缓存不能只按模型名加 prompt hash。还要把模型版本、量化方式、dtype、rope_theta、mapping_version 都纳入缓存键，否则容易出现旧缓存被新模型错误命中的情况。

{ "source_model": "base-v1.0-fp16", "target_model": "chat-v1.0-fp16", "source_rope_theta": 10000, "runtime_dtype": "float16", "mapping_version": 3 }

第三，质量回退策略。映射质量下降时，要能自动回到 target 模型正常 prefill。不要为了省算力让用户长期承受质量劣化。可以设置一个阈值，比如映射后 logits 与真实 logits 的 KL 散度过大时，直接走原始路径。

第四，观测指标。记录 prefill 节省时间、缓存命中率、首 token 延迟、完整生成质量指标。最好能按模型版本和业务场景分开统计，这样能快速定位是哪个环节出了问题。

6. 我的一些经验和判断

6.1 先跑通单层，再谈整体迁移

这个方向最忌讳一上来就做全层映射。全层替换之后，如果结果崩了，你很难判断是 W 拟合得不好，还是层与层之间的误差在累积，或者 KV 注入接口本身有问题。

我的建议是先选一层。用 200 到 500 条同分布样本拟合 W，评估该层替换后的 logits 质量。如果单层效果稳定，再逐步增加替换层数。每增加一层，记录一次 logits 相似度和生成质量。这样你能看到一个非常清晰的误差累积曲线，也知道该在哪个位置停下来。

6.2 评估指标要分层，不要只看一个数

第一层是拟合指标。看 MSE、余弦相似度，确认矩阵 W 有没有把 source KV 拉到 target KV 附近。

第二层是 next token logits。把映射后的 KV 注入 target 模型，看它与正常 prefill 的 logits 有多接近。这个指标比单纯看 KV 相似度更接近实际生成效果。

第三层是完整任务指标。跑一段短文本生成，看困惑度变化，或者跑几个代表性下游任务，看分数波动。注意短 prompt 下指标可能看不出差别，要专门测长上下文场景。

我踩过的坑是：KV 余弦相似度很高，但 next token 的 top-5 已经变了。因为 KV 的小误差会在 attention 加权后放大，最终影响排名的变化。所以不要被单独的余弦相似度迷惑。

6.3 这个方向后续会怎么走

Cross-Model KV Cache Transfer 这类思路不会止步于简单线性映射。后续可能有几个演进方向。一是逐层自适应映射，不同层采用不同复杂度的变换。二是低秩映射，减少 W 的参数体积，同时保持拟合效果。三是结合量化感知映射，让 W 直接适配低比特 KV Cache。四是推理框架层面的集成，把跨模型缓存路由变成像 prefix cache 一样透明的能力。

对普通开发者来说，最值得先做的是理解 prefill 复用这个思想。KV Cache 不是模型权重本身，而是一次前向计算的中间产物。如果它能在保证质量的前提下被安全迁移，降低的就不仅是一个模型的服务成本，而是整个多模型生态的调度成本。

如果你也想试，我建议先拿同一家族的两个小模型跑一遍单层映射。把 token 对齐、KV 收集、矩阵拟合、logits 评估这条链路走通，再考虑接到自己的推理框架里。很多问题不是线性映射不成立，而是数据对齐和目标模型结构没有处理干净。踩过几次之后你会认同这个判断。