NVIDIA、線形マッピングでモデル間のKV cache転送を実現——モデル切り替え時のprefillを最速25分の1に短縮

NVIDIAの研究チームは、コストに基づくルーティング、モデルのアップグレード、長い対話における動的なモデル切り替えを想定し、受け側のモデルがprefillを再実行するコストを省くモデル間KV cache転送手法を提案した。既存のprefix cacheは通常、同一モデルでしか再利用できない。この手法では、転送元と転送先のモデルでKV head数と各headの次元が一致する場合、転送先の各layerおよびheadに対して個別のマッピングを構築する。

研究者らはまず、転送先の各layerを最もよく予測できる複数のlayerを転送元モデルから選び、それらのkeyとvalueを連結したうえで、閉形式のridge regressionにより変換行列を求める。マッピングがキャリブレーション時に観測した位置だけに適合するのを防ぐため、keyから先にRoPEによる回転を除去し、位置情報を含まない内容空間で変換を行った後、転送先モデルのRoPEを再適用する。キャリブレーションに使用するのは、FineWeb-Eduから抽出した各1,024 tokenのデータ500件のみで、勾配ベースの学習は不要だ。

Qwen3、Llama 3.1、Ministral 3における小規模モデルから大規模モデルへの6組のペアでは、4組が5つのbenchmarkにおいて、転送先モデル自身がprefillした場合の平均精度の73～98%を維持した。マッピング速度はprefillの再実行より2.7～25倍高速だった。一方、性能が低かった2組では維持率が最低約42%にとどまった。非線形MLPは、一部の失敗したペアでHellaSwagの精度維持率を最大37 percentage points改善できたが、単純かつ学習不要という利点は失われる。

エンジニアリング上の制約も無視できない。線形mapper自体が依然として10.1億～33.6億parametersを持ち、4～12 GBのメモリを占有する。また、モデルペアごとのキャリブレーションには8基のH100で約47～87分を要する。実験対象も、同一ファミリーに属し、tokenizerを共有し、KVのshapeが一致するモデルに限られている。今後の焦点は、マッピングをvLLMやSGLangなどのpaged cache runtimeへ統合し、長いcontext、multi-tenant環境、量子化KVにおける実際のend-to-end効果を検証することだ。