InferactとGoogleは2026年10月20日、Google Cloud TPU向けの新しいvLLMバックエンド「vLLM TorchTPU(vllm-torchtpu)」をリリースしました。GoogleのTorchTPUスタック上に構築されており、vLLMの既存モデル定義をそのまま使うため、現在vLLMで動く大半のモデルをモデルコードの変更なしにTPUで提供できるとしています。推奨ハードウェアはTPU v7x(Ironwood)です。
vLLMは大規模言語モデルの推論とAPI配信のためのライブラリで、PagedAttentionによるメモリ管理や連続バッチ処理、OpenAI互換APIを備えています。
何が新しいか
vllm-torchtpuはアウトオブツリーのvLLMプラットフォームプラグインです。インストールすると、vLLMがvllm.platform_pluginsエントリーポイント経由で自動検出します。モデルはvLLM標準のモデルランナーで動き、torch.compile(backend="tpu")でXLAへlowerされます。TPUデバイスごとにワーカープロセスを立て、デバイス間通信はtorch.distributedを使います。スケジューラー、ワーカー、KVキャッシュマネージャー、APIサーバーは無改変のupstream vLLMで、TPU固有の処理はカスタムopとして登録したJAX/Pallasカーネルが担います。本取り組みは、両社が前月に発表したエンジニアリングパートナーシップから生まれたものです。
カーネルと並列化
主なPallasカーネルは次のとおりです。
- MHA/GQA向けRagged Paged Attention v3
- DeepSeek、Kimi、GLM向けの密/疎Multi-head Latent Attention(MLA)
- ハイブリッド線形アテンション向けGated DeltaNet、Kimi Delta Attention
- エキスパート並列対応のFused MoE
- FP8、INT8、MXFP4、NVFP4の量子化行列積
vLLMの6つの並列化次元すべてに、シングルホスト/マルチホストのTPUスライスで対応し、組み合わせも可能です。
| 次元 | フラグ |
|---|---|
| TP | --tensor-parallel-size |
| DP | --data-parallel-size |
| PP | --pipeline-parallel-size |
| EP | --enable-expert-parallel |
| PCP | --prefill-context-parallel-size |
| DCP | --decode-context-parallel-size |
大規模MoEでは、v7x単一ホストでDP=8、TP=1+EPのデータ並列アテンション構成が最も高いスループットになることが多いとしています。長コンテキストではPCPとDCPでprefillとKVキャッシュをデバイス間に分割します。
KVキャッシュについては、TPUデバイスバッファをPython外から直接読むC++ライブラリ「TPU Sync」とvLLMのKV connectorを使い、ホストDRAMへのゼロコピーオフロードと、TPUホスト間の分離型prefill/decodeサービングを実現します。prefillをPCP=4、TP=2、decodeをDP=8、TP=1のように別レイアウトにでき、転送時にKVキャッシュを再シャーディングします。
speculative decodingはMTP、EAGLE-3、DFlash、DSpark、n-gram、suffix decodingに対応します。自動prefix cachingはQwen3.5やKimi K3のようなハイブリッドモデルでも動き、線形アテンション状態もKVキャッシュと一緒にキャッシュします。multi-LoRA、structured outputs、multimodal inputsも通常のvLLMフラグで有効化できます。
試し方・要件
Cloud TPU VM上のPython 3.12が必要です。以下でプラグインと対応するvLLM、torch、torch_tpu、libtpuが入ります。インストールガイドはuv、ソースビルド、Dockerも扱っています。
pip install vllm-torchtpu
サンプルの起動コマンドは次のとおりです。
vllm serve Qwen/Qwen3-Coder-30B-A3B-Instruct-FP8 \
--tensor-parallel-size 8 \
--enable-expert-parallel \
--kv-cache-dtype fp8 \
--max-model-len 16384 \
--async-scheduling
初回起動時は設定したshape向けにTPUグラフをコンパイルするため数分かかる場合があります。再起動後もコンパイル済み成果物を使うにはVLLM_XLA_CACHE_PATHを設定します。
ハードウェアはTPU v7x(Ironwood)が推奨で、全プルリクエストでCIテストしています。v6e、v5e、v5pは実験的対応です。TPU固有の最適化はQwen、DeepSeek、Kimi、GLM、Llama、GPT-OSS、Gemmaの各ファミリーで進めています。
今後の予定
公開ロードマップには、新しいオープンウェイトモデルを公開から24時間以内にオンボード・ベンチマークしサポート行列を公開する自動CIパイプライン、prefix/キャッシュ/セッションを考慮したルーティングとllm-dのTPUメモリ階層ルーティング、Tokamaxカーネル向けPyTorchインターフェース、vllm-omniによる拡散モデルのサービングが含まれます。あわせて、vllm-torchtpu上で開発するコントリビューター向けに無料TPU時間を提供する「Community Contribution TPU Capacity Pool」を開設し、InferactのvLLMコアメンテナーが設計とPRレビューを支援します。
まだコメントはありません。