なぜ今、TPU上のPyTorchネイティブ推論なのか

TPUは長年にわたり大規模AIワークロードを高性能・高効率で処理してきました。しかし開発者からは常に同じ声が上がっていました。「JAXバックエンドは良いが、自分たちのサービングスタックをPyTorchでカスタマイズし、デバッグしたい」。エージェント型ワークロードやマルチモーダルモデルが本格化するにつれ、「動く」だけでなく「効率的にサービングする」ことが競争力の核心になりました。

そこでGoogleはTorchTPUを開発しました。PyTorchをTPU上でネイティブに実行するレイヤーです。その上にvLLM(Inferactとの協業)とSGLang(RadixArkとの協業)のエンジンを新たに構築しました。目標は明確です。PyTorchコミュニティが使い慣れたツールのまま、モデルをTPUに載せて高速にサービングできるようにすること。

本記事はGoogle開発者ブログの根拠資料を基に、実務観点で再構成したものです。

TorchTPUの仕組み:ラッパーなしのネイティブデバイス

TorchTPUはTPUを「PyTorchネイティブデバイス」にします。要点は以下です。

  • テンソルはtpuデバイス上の標準PyTorchテンソルとして管理され、ラッパー型はありません。
  • コードはeager実行、またはtorch.compileでXLAコンパイルされます。
  • 集団通信(collective)はtorch.distributedをそのまま通ります。
  • 各デバイスが独自のPyTorchプロセスを動かすため、両エンジンとも各フレームワークのワーカーモデル・スケジューラ・KVキャッシュマネージャ・APIサーバをそのまま維持し、下層のみTPU実装に差し替えます。

ほとんどのモデルは両エンジンで標準のアップストリームPyTorchモデル定義をそのまま使用します。TPU固有の最適化はカーネルとレイヤ実装にのみ入ります。Pallasで書かれたカスタムカーネルは、PyTorchのカスタムオペレータとしてコンパイルグラフ内から直接呼び出されます。この「プラットフォーム・プラグイン」設計により、フレームワークの再利用を最大化しつつ、TPU最適化を独立して素早く進化させられます。

本番機能とエージェント型ワークロードへの対応

両エンジンとも本番ワークロードの全範囲を狙います。特にエージェント型ワークロードは複数機能を同時に要求します。エージェントは拡大するマルチターン会話を繰り返し再送信するため、再計算ではなくキャッシュ済みプレフィックスを探す必要があります。両エンジンは再利用可能なプレフィックスをHBMに保持し、ホストオフロードを有効にすると古いプレフィックスを2次階層のホストメモリへ自動移動します。

ここで2つの機能がKVキャッシュを「計算したデバイスの外」へ移します。

  • 分離サービング(disaggregated serving): プリフィルとデコードを別々のTPUホストで実行。リクエストのKVキャッシュがプリフィルホストからデコードホストへ移動する必要があります。
  • 階層型KVオフロード(tiered KV offload): HBMに収まらないKVキャッシュをホストDRAMとGoogle Cloud Managed Lustreのセカンダリストレージへ移動し、後で同じプレフィックスを再利用する際に戻します。

いずれも転送が計算と重ならないとレイテンシに直結します。両エンジンはこれをTPU Syncで解決します。TPU上のKVキャッシュ転送と管理のためのオープンソースC++ライブラリです。PyTorchテンソル背後のデバイスバッファを直接読み取るため、ホスト側のステージングコピーが不要で、Python外の専用スレッドで転送を実行しプリフィル・デコードを継続させます。vLLM-TorchTPUはvLLMのKVコネクタインターフェース経由で公開し、SGLang-TorchTPUは分離転送バックエンド兼HiCache I/Oバックエンドとして利用します。

両エンジンはすでにプレビュー顧客とともに本番規模で稼働しています。性能の透明性のため、コミュニティ標準であるSemiAnalysis InferenceXプロジェクトで公開ベンチマークを実施しています。

TPU server rack powering large-scale PyTorch inference engines vLLM and SGLang Dev Environment Setup

ベンチマークとパートナーシップ:数字で見る現状

Ironwood TPU(TPU 7x)上でQwen 3.5を動かした初期提出を基に、エンジニアリングチームがスタックをさらに最適化し、100万トークンあたり最低コスト2.01セントという業界をリードする数値を記録しました。エージェント型のロングコンテキスト・マルチターンワークロードを評価するAgentXトラックにも新たなSOTA結果を公開し、Ironwood TPUがパレートフロンティアの大部分で先行しています。

このバックエンドは真空から生まれたものではありません。PyTorchを受け入れ、オープンソースメンテナと直接協業することで、アップストリーム整合とコミュニティファーストを保っています。

  • SGLangチームとの協業: TPU向けエンジンを構築し、Metaの本番環境で既に稼働。MetaのPyTorch-TPUリードZongwei Zhou氏は「PyTorchネイティブのサービングスタックが既存AIインフラに自然に統合される」と述べています。
  • RadixArk: TorchTPUが多様なシリコンへSGLangを展開する開発速度を大幅に高めたと評価。Banghua Zhu CTOは「顧客へ効率の利得と柔軟なコンピュート選択をそのまま届けられる」とコメント。
  • Kimi K3の事例: 2.8兆パラメータのKimi K3向け高性能Pallasメガカーネルを統合し、デコードステップ全体を単一カーネルで実行。Ironwood TPU上で単一ユーザーに毎秒700トークン以上をサービングします。
  • Inferact: vLLM-TorchTPUスタックでKimi K3の初回AgentX予備提出を主導しました。
# TorchTPU上でvLLMエンジンを使う概念例(疑似コード)
from vllm import LLM, SamplingParams

# TPUデバイス上でネイティブPyTorchテンソルとして動作
llm = LLM(
    model="Qwen/Qwen3.5",
    device="tpu",              # TPUネイティブデバイス
    enable_prefix_caching=True, # エージェント型マルチターン対応
    kv_cache_dtype="auto",
)

sampling = SamplingParams(temperature=0.7, max_tokens=512)
out = llm.generate(["TPU上のPyTorchネイティブ推論を説明して"], sampling)
print(out[0].outputs[0].text)

上記コードは概念理解用の例です。実際のAPIは各エンジンの公式ドキュメントを確認してください。

注意点と限界:導入前に押さえるべきこと

  1. JAXバックエンドとの共存: vLLM-JAXとSGLang-JAXは引き続き維持・サポートされます。既存のJAXベースパイプラインを無理に置き換える必要はありません。
  2. モデルカバレッジ: Qwen、DeepSeek、GLM、Kimi、GPT-OSS、Gemmaなど人気オープンモデル群はサポートされますが、新規モデルのDay 0対応は拡張中です。社内ファインチューニングモデルはカーネル互換性の検証が必要です。
  3. 分離・オフロードの複雑さ: TPU Syncが転送を隠蔽しますが、プリフィル/デコードホストの分離はネットワークトポロジとスケジューリング設計を要求します。小規模ワークロードではむしろオーバーヘッドになり得ます。
  4. ベンチマークの解釈: 2.01セント/100万トークンは特定モデル・TPU構成・バッチ条件での数値です。自社ワークロードにそのまま当てはめず、InferenceXリーダーボードの条件を確認してください。

次のステップ学習の方向性

  • TorchTPUアーキテクチャ: torch.compile + XLAバックエンドがTPUグラフをどうコンパイルするかを学ぶと、カスタムカーネル最適化の勘所が掴めます。
  • Pallasカーネル: TPU向けカスタムカーネルの書き方を習得すれば、メガカーネル事例が理解できます。
  • KVキャッシュ管理: プレフィックスキャッシング、階層型オフロード、分離サービングはLLMサービングエンジニアリングの中心テーマです。
  • vLLM/SGLangの内部構造: スケジューラとKVコネクタインターフェースを読み解けば、TPUバックエンドがどこに差し込まれるか明確になります。

日本国内で大規模推論基盤を運用するチームであれば、TPU導入時にコストモデル(予約 vs オンデマンド)やGKE連携、社内PyTorchパイプラインとの整合性を先に検討することを推奨します。特に既存のJAXベース資産とのマイグレーションコストは必ず見積もってください。

Cloud infrastructure diagram showing disaggregated prefill and decode on TPU hosts Programming Illustration

vLLM-TorchTPU vs SGLang-TorchTPU 比較

項目vLLM-TorchTPUSGLang-TorchTPU
開発パートナーInferactRadixArk
TPUネイティブ基盤TorchTPUTorchTPU
KVキャッシュ転送vLLM KVコネクタインターフェースでTPU Syncを公開分離転送バックエンド + HiCache I/Oバックエンド
プレフィックスキャッシングHBM + ホストオフロードの2階層HBM + ホストオフロードの2階層
分離サービング対応対応
代表事例Kimi K3 AgentX予備提出Meta本番、Kimi K3メガカーネル700+ tok/s
カスタムカーネルPallasカーネルをPyTorchカスタムオペレータとして呼び出し同左
エコシステム整合アップストリームvLLMに整合アップストリームSGLangに整合
適したチームvLLMベースのサービングスタックを持つSGLangベースのエージェント型サービング中心

両エンジンともJAXベースバックエンド(vLLM-JAX、SGLang-JAX)と並行運用されます。新規導入なら、ワークロード特性(エージェント型マルチターン vs 単発大量バッチ)とチームの既存フレームワーク習熟度で選択してください。

AI agent workload running on TPU with PyTorch native backend for high throughput Algorithm Concept Visual

結論:TPUエコシステムの「PyTorchネイティブ」転換点

今回の発表の本質は新しいエンジン2つではなく、TPUがPyTorchネイティブデバイスになったという点です。ラッパーなしのテンソル、torch.compile/XLA、torch.distributedそのまま。そこにPallasカスタムカーネルとTPU Syncが加わり、エージェント型ワークロードのKVキャッシュボトルネックに正面から対処します。

実務観点の助言は3つです。

  1. 小さく始める。 既存PyTorchモデルを1つTorchTPU上のvLLMに載せ、InferenceX条件と自社指標を比較してください。
  2. KVキャッシュ戦略を先に設計する。 プレフィックス再利用パターンが明確なら、階層型オフロードと分離が大きな利得になります。
  3. JAX資産は維持する。 両陣営とも引き続きサポートされるため、検証済みワークロードはそのままに、新規エージェント型サービングからPyTorchネイティブへ移すハイブリッドが現実的です。

あわせて読みたい記事

本コンテンツは、信頼性の高い情報源をもとにAIツールを活用して作成され、編集者によるレビューを経て公開されています。専門家によるアドバイスの代替となるものではありません。