HeteroFold 让 Llama、Qwen、Ministral 跨家族直接共享 KV 缓存,32K 上下文比原生预填充快约 10.7 倍
核心概要
该工作提出 HeteroFold,一种在发送方与接收方模型均冻结的前提下、无需接收方预填充的跨家族 KV 缓存迁移方法,通过基于共享字符边界的 Token 对齐、跨层跨头 K/V 映射与面向接收方行为的校准,在 Llama-3.1-8B、Qwen3-4B、Ministral-3-14B 之间的六个迁移方向上于四个长上下文基准取得最佳缓存迁移表现,并在 HiddenBench 多智能体通信上达到与文本通信相当的水平,32K 上下文下 Llama-3.1-8B 到 Ministral-3-14B 的迁移比原生预填充快约 10.7 倍。
深度剖析
论文把分词器不匹配识别为跨家族免预填充 KV 复用的关键障碍,并提出 Token Alignment(TA),通过共享字符结束边界建立发送方与接收方的 token 对应关系。 此前的免预填充缓存迁移方法(Dense Latent、KV Ridge)的评测主要停留在同家族、分词器兼容的设定,未处理跨家族 token 对应问题。 消融实验显示,在 Ministral-3-14B 到 Llama-3.1-8B 方向上把 TA 换成同索引配对会造成最大幅度的性能下降;图 5 显示同索引配对的位置错位随上下文长度增长,而 TA 通过共享字符边界保持接近对齐。
论文提出跨家族 K/V 映射,融合多层发送方特征、跨头混合与 Recolor 矩匹配,在保持两个语言模型冻结的情况下构造接收方兼容的 K/V 状态。 相比单层映射与逐头独立映射,该方法显式处理模型深度、KV 结构与特征分布差异,并用 Recolor 匹配接收方特征的均值与协方差,而非仅做逐 token 重建。 消融显示单层与头局部映射均降低性能;表 D 显示三层发送方邻域相比单层把 ARC-C、HotpotQA、QuALITY 从 66.89/22.52/51.58 提升到 77.05/49.44/62.90;表 F 显示六个方向映射器参数为 2.01 亿至 2.52 亿,比 Dense Latent 少约 25%、比 KV Ridge 少约 62%。
论文指出仅靠 KV 重建误差不能反映接收方行为,并提出面向接收方的校准,直接匹配注意力模式与输出,学到的修正被折叠进固定仿射映射,推理时无需额外模块。 KV Ridge 的重建误差低于 HeteroFold,却会扭曲接收方注意力,说明需要以接收方下游计算为目标的校准,而非仅优化缓存数值重建。 图 2 与图 4(d) 显示 Recolor 之后仍存在注意力输出误差;表 C 显示非零修正秩相比仅 Recolor 提升 HotpotQA;附录 D.3 显示在 GSM8K 自回归解码中预测分歧未随解码长度持续增大。
论文在六个迁移方向、四个长上下文基准与五个短上下文基准上评估,并在 HiddenBench 多轮异构智能体通信中达到与文本通信相当的表现,同时降低接收方迁移延迟。 此前免预填充迁移缺少跨家族、跨分词器以及多轮生成消息通信的系统评估。 长上下文基准为 Qasper、HotpotQA、LoCoMo、QuALITY,短上下文为 ARC-Challenge、MMLU、WinoGrande、HellaSwag、GSM8K;HiddenBench 使用全部 65 个任务、三或四个智能体、15 轮通信;延迟测量在两张 H100 80GB、NVLink、batch size 2 下进行,32K 上下文下比原生预填充快约 10.7 倍,比 Dense Latent 与 KV Ridge 快 1.18 至 1.47 倍。
启示与展望
该结果面向发送方已处理共享上下文的多智能体场景:发送方与接收方分别位于两张 H100 80GB、NVLink 互联,batch size 为 2,发送方预填充、模型加载与离线校准不计入迁移延迟;若发送方尚未处理上下文,需额外计入发送方预填充与载荷捕获成本(论文给出 Llama 在 4K/16K/32K 下为 280/1328/3381 毫秒,Qwen 为 235/1163/3083 毫秒)。方法在 Llama-3.1-8B、Qwen3-4B、Ministral-3-14B 六个方向上验证,并在同家族 Qwen3 迁移中作为补充评估;校准使用 1600 条训练提示与 400 条留出提示,排除金标答案与解答。
论文报告 HiddenBench 上达到与文本通信相当的表现,但未给出逐项数值;跨家族迁移在短上下文设置中并非全部领先,仅在多数设置中占优;校准敏感性显示修正秩与校准语料选择会影响结果,例如仅用 Open-R1 校准时 HotpotQA 为 12.69,而等量混合为 41.97;自然语言重建中 HeteroFold 的有序重叠为 83.04、精确重建 16/100,说明缓存迁移并非无损;此外,缓存迁移在计入传入载荷拷贝后并不总能降低接收方额外峰值显存。
