论文

KV-Lingo:通过蒸馏学习 KV-Cache 转换器

KV-Lingo: Learning KV-Cache Translators with Distillation

摘要

大语言模型用键值 (KV) 缓存表示上下文。缓存是特定于模型的:对于相同的文本,具有不同架构或权重的模型会产生不兼容的表示。这使得在共享上下文上切换模型的成本很高:尽管上下文已经被一个模型处理过,但传入的模型必须再次处理它以构建自己的缓存。我们引入了 KV-Lingo,一种将源模型的 KV 缓存转换为目标模型可以读取的缓存的方法。 KV-Lingo 由一组线性映射组成,通常目标模型的每一层都有一个,它们独立应用于所有词元的键和值表示。我们使用蒸馏来训练这些映射,最大限度地减少目标模型的本机缓存预测与翻译后缓存预测之间的差异。我们考虑了多个跨越多种尺寸和架构的模型对,在通用文本语料库上为每对模型训练一名翻译人员。由此产生的转换器在小到大和大到小传输中都保留了强大的下游性能。由于切换会花费线性映射和单个解码步骤而不是预填充,因此用缓存翻译替换重新预填充可以将模型切换后第一个词元的时间缩短 9.6 倍(在 Apple M3 Ultra 上的 Qwen 模型的 64 个词元提示下),在 H100 上的 32k 上下文长度下缩短高达 29 倍。这些优点使得 KV-Lingo 对于动态模型路由特别有用:上下文可以由一个模型处理,并仅在需要时才传递给另一个模型,而无需重新预填充共享前缀。我们最终证明 KV-Lingo 可用于无缝模型切换,在多轮评估中的重复切换中保持接近重新预填充。