GoogleのTPU Sync、KVキャッシュデータの移動により推論レイテンシを削減

Googleは、第8世代Tensor Processing Units（TPU）でのAI推論を高速化するために設計された新しいオープンソースソフトウェアフレームワーク「TPU Sync」を発表しました。2026年11月9日に発表されたこのツールは、異なるメモリ階層間でのキー・バリュー（KV）キャッシュデータの移動を最適化することで、大規模言語モデルのサービングにおけるメモリボトルネックに対処します。同社によれば、特定の高パラメータ環境では、このアプローチにより最初のトークン生成までの時間を90%以上削減できるということです。

何が起きたのか

今回のリリースは、Googleのより広範なハードウェアアップデートに伴うもので、特に数百万トークンのコンテキストウィンドウや強化学習ワークロードで大規模モデルを実行する開発者を対象としています。コンテキスト長が増加するにつれ、キー・バリュー（KV）キャッシュの管理コストとレイテンシが重要なボトルネックとなります。TPU Syncは、不要なコピーを行うことなく、アクセラレータメモリ、ホストDRAM、ストレージ、ネットワーク間でデータを効率的に移動させることで、この問題を解決することを目指しています。

社内ベンチマークでは、Googleは大幅な性能向上を実証しました。2350億パラメータのQwen3モデルでは、KVキャッシュの移動時間が数秒から約230ミリ秒に短縮されました。これにより、軽負荷時の平均最初のトークン生成時間は60%低下しました。さらに大きな4800億パラメータのモデルでテストを行った場合、レイテンシ削減率は90%に達しました。64,000トークンのコンテキストを用いたプレフィックスキャッシュテストでも同様の改善が見られ、スループットは50%以上増加し、P99インタートークンレイテンシは30%低下しました。

ただし、これらの結果には留意点があります。性能は特定のワークロード、キャッシュデータの再利用量、利用可能なネットワーク帯域幅によって変動します。また、すべてのユースケースに対して本番稼働可能とはまだ見なされていません。GoogleはAPIが依然として進化中であることを警告しており、メインブランチを追跡するのではなく、特定の検証済みコミットまたはタグ付きリリースにデプロイメントを固定することを推奨しています。

仕組み

TPU SyncはC++で記述されており、TPUホストのメモリ階層全体にわたるデータ転送を管理することで動作します。ゼロコピー転送とネイティブのPjRtBufferハードウェアディスクリプターを活用し、冗長なデータ複製を排除します。この手法により、システムはハードウェアインターフェースの理論的限界に近づきます。Ironwood TPUホストでのテストでは、ホストからデバイスへ、およびデバイスからホストへの総合スループットが223 GB/sに達し、PCIe Gen5の限界値である256 GB/sに迫りました。

このフレームワークは、TorchTPU、vLLM、SGLangなど既存の推論エコシステムと統合されています。この互換性により、開発者はサービングスタック全体を再構築することなく、TPU上で分離型（disaggregated）ワークロードを実行できます。分離型アーキテクチャでは、推論のプリフィルフェーズとデコードフェーズを分割し、それぞれを独立して最適化できるようにします。従来、外部TPUサービングにはこの分離を効果的に処理するためのインフラストラクチャーがありませんでしたが、TPU Syncはこれらのフェーズ間の状態とデータ移動を管理するために必要なメカニズムを提供します。

主要な詳細

• TPU Syncは、KVキャッシュデータの移動を管理するためにC++で記述されたオープンソースフレームワークです。
• ベンチマークでは、4800億パラメータモデルにおいて最初のトークン生成までの時間が最大90%削減されることが示されています。
• Ironwood TPUホストでのスループットは223 GB/sに達し、PCIe Gen5の理論的限界に近づいています。
• ソフトウェアは、特定のJAXおよびPyTorchバージョン要件を満たすCloud TPU v7xおよびv6eをサポートしています。
• TorchTPU経由で、vLLMやSGLangなどの人気推論エンジンとの統合が可能です。
• APIが変更中のため、Googleは検証済みコミットまたはタグ付きリリースの使用を推奨しています。

なぜ重要なのか

大規模言語モデルを用いて製品を開発するエンジニアリングチームにとって、推論レイテンシはユーザー体験と運用コストに直接影響します。インタラクティブなアプリケーションにとって「最初のトークン生成までの時間」は重要な指標であり、これを大幅に短縮することは、これまで実用的ではなかったコンテキストウィンドウを実現可能にする可能性があります。メモリ移動を管理するソフトウェア層に対処することで、Googleは成熟したツールとパフォーマンス最適化により長年AI推論分野を支配してきたNvidiaのCUDAエコシステムとのギャップを埋めようとしています。

この進展は、分離型サービングアーキテクチャの重要性が高まっていることも浮き彫りにしています。モデルのスケーリングに伴い、プリフィルとデコードを別個の問題として扱うことで、リソース利用率を向上させることができます。TPU SyncはGoogleハードウェア上でこれを可能にするインフラストラクチャーを提供し、特定のワークロードにおいてNvidiaのGB200およびGB300システムに対するコスト効率の良い代替案となる可能性があります。しかし、完全にチューニングされたNvidiaセットアップとの直接的な比較はまだ限られており、開発者は現在のソフトウェアの未成熟さとその利点を天秤にかける必要があります。

実施できること

• KVキャッシュ最適化の恩恵を受ける可能性のある、大きなコンテキストウィンドウや高パラメータモデルを含む現在のワークロードを評価してください。
• Cloud TPU v7xまたはv6eインスタンスを使用して非本番環境でTPU Syncをテストし、実際のレイテンシ改善度を測定してください。
• 進化するAPIによる破壊的な変更を避けるため、依存関係を特定のタグ付きリリースまたは検証済みコミットに固定してください。
• 既存の推論スタックを確認し、TorchTPU、vLLM、またはSGLangとの統合が可能かどうかを検討して、導入を容易にしてください。
• ネットワーク帯域幅とキャッシュデータの再利用率を監視してください。これらの要因は得られる性能向上に大きく影響します。
• 本番移行を計画する前に、APIの安定性とより広範なハードウェアサポートに関するGoogleからの今後のアップデートに注意してください。