Google、TPU向けPyTorchネイティブのvLLMとSGLang推論エンジンを発表

両エンジンはTorchTPUの上に構築されている。TorchTPUはTPUをPyTorchのネイティブデバイスとして扱い、テンソルはラッパー型なしの標準PyTorchテンソルとして管理される。コードはeager実行またはXLAを使ったtorch.compileで動き、集団通信はtorch.distributedを通る。ワーカーモデル、スケジューラ、KV cacheマネージャ、APIサーバーは各フレームワークのものを維持し、その下にTPU実装を加える設計である。Pallasで書いたカスタムカーネルは、コンパイル済みグラフ内のカスタムオペレータとしてPyTorchから呼ばれる。

対応機能は幅広い。自動プレフィックスキャッシュ、ホストメモリやManaged LustreへのKV cacheオフロード、ツール呼び出し、MTPやEAGLE-3などによる投機的デコーディング、expert parallelismを伴うdata-parallel attention、コンテキスト並列、prefillとdecodeを別ホストに分ける分離型サービングなどが挙がっている。KV cacheの移動にはオープンソースのC++ライブラリTPU Syncを使い、ホスト側の追加コピーなしにPyTorchテンソルの背後のデバイスバッファを直接読む。

性能面では、SemiAnalysis InferenceXで公開ベンチマークを行っており、エージェント型の長コンテキストやマルチターンを評価するAgentXトラックにも新たな結果を出したとしている。Ironwood TPUはPareto frontierの大半で先行するという。両エンジンはプレビュー顧客の本番環境で稼働しており、SGLang-TorchTPUはMetaで本番稼働していると、Metaの担当者がコメントしている。

vLLM側では、Inferactと2.8Tパラメータの Kimi K3向けにPallasメガカーネルを統合し、decodeステップ全体を1つのカーネルで実行して、Ironwood TPU上で単一ユーザーに700トークン/秒超を提供したと述べている。vLLMはTPU計算資源の無償プールも用意し、両エンジンへ積極的に貢献する開発者に提供する。Qwen、DeepSeek、GLM、Kimi、GPT-OSS、Gemma4などの主要なオープンモデルに対応し、新モデルのDay 0対応にも取り組む。既存のvLLM-JaxとSGLang-JAXも引き続き保守される。