Google、階層型 KV キャッシュと非同期転送で推論待ち時間を短縮する TPU Sync をオープンソース化

Google Cloud は 11 月 9 日、長文コンテキスト推論と強化学習向けにモデル状態を移動する基盤として、TPU Sync をオープンソース化すると発表しました。KV キャッシュとモデルの重みを、異なるメモリ階層やノード間で移動させます。PyTorch と JAX 向けのパッケージを提供し、vLLM、SGLang、llm-d のエコシステムに接続できます。公式発表

技術の中核は、TPU ネイティブのテンソルレイアウトを維持しながら、PJRT バッファーのエイリアスと非同期 DMA を使って、HBM、ホストメモリ、リモートノード間でデータを移動することです。ネイティブ C++ のバックグラウンドスレッドが転送とメモリ操作を担い、Python によるシリアライズやグローバルインタープリターロックの介在を抑えます。これにより、データ移動を prefill や decode の計算と並行して進められます。ホスト DRAM は第 2 層のプレフィックスキャッシュとして機能し、HBM の容量が足りない場合に追い出されたブロックを保持し、必要に応じて読み戻します。リポジトリのアーキテクチャ説明

同じシステムプロンプトを繰り返し使うサービスや、複数ターンの会話を継続的に処理するサービスにとって、これは直接的な意味があります。計算済みのプレフィックスを保持すれば、prefill の再実行を減らせます。prefill と decode を分離した構成では、プロンプト処理を終えたノードから生成ノードへ KV ブロックをストリーミングすることもできます。Google によると、低並行度のテストで Qwen3 235B の初回トークン遅延が平均 60% 低下しました。これは特定のワークロードにおける社内測定結果であり、あらゆるサービスにそのまま当てはまるものではありません。テスト結果と条件

導入にあたっては、2 点に注意が必要です。まず、共有ストレージ階層と強化学習向け転送モジュールは、まだ実験的な機能です。また、リポジトリには API が急速に変化していると明記されており、タグ付きバージョン、または lkg.version が指定する検証済みコミットに固定することが推奨されています。現時点で記載されているハードウェア検証範囲は TPU v7x と v6e です。エンジニアリングチームは、自社のキャッシュヒット率、ホストメモリの予算、同時実行トラフィックに基づいて、テールレイテンシーとスループットの両方が改善するかを確認する必要があります。バージョンとハードウェアの説明