8月10日、Dataconomyが「Google open-sources TPU Raiden for faster AI inference」と題した記事を公開した。AnthropicやOpenAIといった主要AIラボがGoogle TPUの採用を本格化させる中、Nvidia GPU向けには整備されてきた推論最適化のソフトウェア基盤がTPUには欠けているという課題が顕在化していた。GoogleがオープンソースとしてリリースしたKVキャッシュ転送ライブラリ「TPU Raiden」は、その空白を埋めようとする取り組みだ。

KVキャッシュ転送の専用ライブラリ

GoogleはLLM推論向けの最適化ライブラリ「TPU Raiden」をオープンソースとして公開した。リポジトリはApache-2.0ライセンスのもとGitHub上で公開されており、現時点では本番環境での一般利用を想定したものではなく、開発継続中のプロジェクトと位置付けられている。

Raidenが解決しようとする問題を理解するには、大規模LLM推論の構造を押さえる必要がある。推論処理は通常、入力プロンプトを処理する「プリフィル」フェーズと、トークンを逐次生成する「デコード」フェーズの2段階に分かれる。本番環境ではこの2フェーズを別々のハードウェアプールに分離して処理する「disaggregated serving（分散サービング）」が採用されることが多い。この構成では、プリフィル完了後にKVキャッシュ（アテンション計算の中間状態を保持するデータ）を高速に転送する仕組みが不可欠になる。Raidenはまさにこの転送レイヤーに特化したライブラリとして設計されている。

Anthropic・OpenAIのTPU採用が背景に

GoogleのTPUユーザー基盤はここ最近で急拡大している。AnthropicはGoogleとの契約を通じて大規模なTPUチップへのアクセスを確保しており（※原文では "millions of TPU chips" と表現されているが、具体的な台数・金額については元記事の原文を参照されたい）、OpenAIもGoogle TPUを一部の処理に利用し始めたことが報じられている。こうした需要の急増を受け、Nvidia GPU周辺にvLLMやSGLangなどのフレームワークを通じて整備されてきたようなソフトウェアエコシステムを、TPU向けにも構築する必要性が高まっていた。

今回のRaidenのリリースは、そうした需要に直接応える動きと見ることができる。

NvidiaのNIXLに相当する位置付け

業界アナリスト機関SemiAnalysisがこのリリースをいち早く取り上げ、Raidenが推論スタックの中でNvidiaのNIXLライブラリと同じ層に位置すると指摘した。NIXL（NVIDIA Inference Xfer Library）はvLLMやSGLangといった主要な推論フレームワークと深く統合されており、disaggregated servingにおけるKVキャッシュ転送の標準的な実装として機能している。

RaidenがそれらのフレームワークにNIXL同等の形で統合されれば、すでにNvidia GPU上でdisaggregated servingを運用している企業がTPUへ移行する際のエンジニアリングコストを下げられる可能性がある。互換性のある抽象レイヤーが存在することで、フレームワーク側がハードウェアごとに転送ロジックを個別実装しなくて済むためだ。

Raidenが提供する機能モジュール

リポジトリのドキュメントによると、Raidenは以下の機能モジュールを含む。

• チップ間直接転送：同一マシン内のTPUチップ間でKVキャッシュを直接移動
• クロスVM転送：ネットワーク越しに異なるVM間でキャッシュを転送
• オフロード：TPUメモリからホストRAMへキャッシュブロックを退避

この3モードは、プリフィルとデコードを分離した構成において発生するさまざまなトポロジに対応するものだ。プリフィル側とデコード側が同一ホスト内に収まる場合はチップ間転送、異なるノードに分散している場合はクロスVM転送が使われる。オフロードはTPU HBM（High Bandwidth Memory）の容量を補うためのフォールバック手段として機能する。

さらに共有メモリモードも備えており、モデルサーバーが再起動してもKVキャッシュをDRAM上に保持し続けることができる。これはローリングアップデートのような通常の本番更新時にキャッシュを失わずに済む機能であり、サービス連続性の観点で実用上の意味が大きい。

推論インフラのオープン化という文脈

このRaidenのリリースは、業界全体の動きとも重なる。リリースの数日前の8月4日には、NvidiaがcuFile APIとストレージスタックをクロスベンダーのGitHub organizationを通じてオープン化すると発表しており、その初期メンテナーにはGoogle、Intel、Metaが名を連ねている。推論インフラのオープン化・標準化をめぐる動きが、ほぼ同時期に複数のベンダーから相次いでいる形だ。

Raidenはまだ開発途上にあるが、主要なLLM推論フレームワークとの統合が進めば、TPUを推論基盤として選択する上での実装コストを下げる重要なピースになり得る。

詳細はGoogle open-sources TPU Raiden for faster AI inferenceを参照していただきたい。