Google、TPU間でKVキャッシュを高速転送するOSS「TPU Sync」を公開

KVキャッシュ（キーバリューキャッシュ）とは、言語モデルがプロンプトを読み込む際に作る作業記憶のことです。処理済みのトークンごとに中間値を保存しておくことで、新しい単語を生成するたびに再計算せずに済みます。会話や文書が長くなるほどこのキャッシュは大きくなり、通常はアクセラレータチップ上の希少で高価な広帯域メモリ（HBM）に置かれます。

Googleは、このキャッシュを自社のAIプロセッサであるTPU（テンソル処理ユニット）の間で移動・管理するための、オープンソースのネイティブC++ライブラリ「TPU Sync」を公開しました。TechBytesによると、これはエンジニアのFenghui Zhang氏とWeiren Yu氏が書いたGoogle Cloudブログの記事を要約したもので、このライブラリは強化学習のワークフローにおけるモデルパラメータの同期も担います。SemiAnalysisは、TPU Syncが以前はTPU-raidenと呼ばれていたと述べており、Open Source For YouはこのプロジェクトをTPU RaidenとしてApache 2.0ライセンスで公開されたと紹介しています。

なぜKVキャッシュを移動させる必要があるのか

Open Source For Youによると、言語モデルの推論には二つの段階があります。モデルがプロンプト全体を読み込むプリフィル（prefill）と、回答をトークンごとに生成するデコード（decode）です。分離型サービングはこの二つの段階を別々のハードウェアプールで実行するため、プリフィルの後にKVキャッシュをデコード担当のチップへ送る必要があります。SemiAnalysisは、転送ライブラリがなければTPUの外部向けサービングは集約モードに限られると指摘しています。Google自身は、Geminiのために社内で長年分離型サービングを運用してきました。

もう一つの必要性はオフロードです。SemiAnalysisによると、チップのHBMが全ユーザーのキャッシュを収容できない場合、データをHBMからホストのDRAMへ移し、また戻す必要があります。The Next Platformが2026年11月9日に伝えたところでは、GoogleのAI・コンピューティングインフラ担当副社長兼ゼネラルマネージャーであるMark Lohmeyer氏が電話会見で、エージェント型のワークロードではTPUが時間の大部分で遊んでいることが多いと述べました。同氏によれば、TPU Syncはエージェントが一時停止したときに作業記憶をチップの外へ移し、再開時に復元します。これによりアクセラレータが他の作業に使えるようになり、エージェント型ワークロードあたりのコストが下がります。The Next Platformは、TPU Syncはソフトウェアなので、Googleが以前の世代のTPUシステムにも移植できると付け加えています。

Googleはどのような性能を示しているのか

以下の数値は、2026年11月9日付のTechBytesによるもので、同媒体はこれをGoogle Cloudブログに帰しています。Googleの元記事そのものは確認できなかったため、これらはベンダーが示した数値を二次的な媒体が伝えたものとして扱う必要があります。

• TechBytesによると、パラメータ数が2350億から4800億のオープンウェイトのMixture of Expertsモデルで、キーバリュー転送の遅延が90%以上短縮されます。
• TechBytesによると、最初のトークンが出るまでの時間（TTFT、回答の最初の単語が現れるまでの待ち時間）は、プレフィックスキャッシュと分離型サービングのベンチマークで、標準構成と比べて平均90%短縮されます。
• TechBytesによると、raw_transferコンポーネントはIronwood TPU v7xホスト上で、デバイスとホストの間を双方向合計で最大223 GB/s転送でき、PCIe Gen5の物理的な上限256 GB/sに迫ります。
• TechBytesによると、データセンターのネットワークでは、200 Gbit/sのネットワークカード1枚で約192 Gbit/s、NUMAを考慮したストライピングを使った2枚で350 Gbit/sを維持します。

TPU Syncは内部でどう動くのか

TechBytesによると、TPU Syncはデータのレイアウトを保ったままのダイレクトメモリアクセス（DMA、プロセッサを介さずにメモリをコピーするハードウェアの仕組み）と、インプレースのバッファエイリアシングを組み合わせています。このライブラリはJAXのPJRTバッファやPyTorchのテンソルと直接やり取りするため、Pythonのシリアライズや中間コピーを避けられ、転送はXLAのコンパイルグラフの外で行われます。SemiAnalysisは同じ考え方を、PJRTバッファのネイティブなハードウェア記述子を取り出すことで実現するゼロコピー転送だと説明しています。

TechBytesによると、ランタイムは三つのモジュールで構成されます。グローバルなキーバリューストアはプロンプトのハッシュを物理的な保存ブロックに対応づける台帳として働き、最も長く使われていないものの追い出しを制御します。クラスタコントローラは実行計画とリース管理を担います。ローカルのキャッシュマネージャはデバイス上のブロックの状態を追跡し、ピン留めされたホストメモリを管理し、/dev/shmの共有メモリにウォームブート用のコピーを保持します。メモリ階層はHBM、ピン留めされたホストDRAM、ローカルSSD、Google Cloud Managed Lustreにまたがります。

統合について、TechBytesはTorchTPU経由のPyTorch向け公式バインディング、JAX向け公式バインディング、vLLM、SGLang、llm-dのサービングスタック向けコネクタ、そしてPathwaysやマルチコントローラのJAXパイプライン向けの低レベルのプリミティブを挙げています。

限界と成熟度はどうなっているのか

TechBytesは、HBMとホストDRAMの間の転送は本番利用が可能だとする一方、Managed Lustreや外部のPOSIXバックエンドのような共有ストレージへの階層化は、まだ実験的な検証段階にあるとしています。強化学習向けの重み同期モジュールも実験的で、最大のネットワークスループットを得るには、NUMAを考慮したストライピングを正しく設定した200 Gbit/sのネットワークカード2枚が必要です。日付のないOpen Source For Youの記事はさらに踏み込み、このプロジェクトはまだ活発に開発中で、一般的な本番利用を想定していないとしており、成熟度の見方は情報源によって異なります。

SemiAnalysisは競合の観点も加えています。同社のInferenceXプレビューは、集約型のTPU v7を集約型のNvidia B200およびB300と比較しており、分離型サービングではNvidiaのGB200とGB300 NVL72システムが曲線の一部で現在リードしていると述べています。TPU-Syncに基づく分離型サービングが外部スタックで最適化された段階で、分離型のTPU v7との比較が予定されており、SemiAnalysisはTPU上のMooncake Storeのサポートが、TPU-Syncのプリミティブの上に構築されると見ています。

企業にとって何が変わるのか

Google Cloud上で長いコンテキストのアシスタントやエージェントを運用または利用する日本企業にとって、KVキャッシュの移動は自分で作り込むものから、設定して使うものに変わります。実務上の手順は、自社のプロンプトでvLLMまたはSGLangとのコネクタを試し、導入前後で最初のトークンまでの時間を測り、頼る前にどのメモリ階層がまだ実験的とされているかを確認することです。強化学習で学習を行うチームは、自社のトレーナーとサンプラーの構成で重み同期モジュールを検証すべきです。公表された数値のどれも、自社のワークロードでのベンチマークの代わりにはなりません。