vLLM TorchTPU: Google Cloud TPU 推論のための新しい vLLM バックエンド

vLLM TorchTPU は、vLLM の既存のモデルコード、CLI、スケジューラ、KV キャッシュマネージャー、OpenAI 互換サーバーを変更せずに、vLLM が Google Cloud TPU 上でモデルを提供できるようにする、新しいツリー外 vLLM プラットフォームプラグインです。これは、vLLM で既に動作しているほとんどのモデルが、モデル定義を編集することなく TPU で実行できることを意味し、また、アテンション、Mixture-of-Experts、量子化された matmul 用の手書きの Pallas カーネルをオープンソースにするため、重要です。

vLLM TorchTPU とは

vLLM TorchTPU は、Google Cloud TPU 上で vLLM を実行するためのバックエンドである vllm-torchtpu パッケージです。これは、Google の TorchTPU スタック（TPU 上で PyTorch をネイティブに実行する）上に構築されており、vLLM に既に存在するモデル定義を実行します。その結果、現在 vLLM で動作するほとんどのモデルは、モデルコードを変更することなく TPU で提供できます。

このプロジェクトは、リリースの前月に発表された Inferact と Google のエンジニアリングパートナーシップから生まれました。リリース投稿では、プロジェクトの 3 つの目標が示されています。

  1. vLLM の既存のモデルコードを再利用して、新しいモデルを TPU に簡単に導入できるようにする。
  2. アテンション、Mixture-of-Experts、量子化された matmul 用の Pallas カーネルを使用して、オープンソースで TPU 推論パフォーマンスを向上させる。
  3. vLLM のエクスペリエンスを同じに保つ：他のハードウェアで使用されるのと同じ CLI、OpenAI 互換サーバー、スケジューラ、KV キャッシュマネージャー。

TPU バックエンドが vLLM にどのようにプラグインされるか

vllm-torchtpu は、ツリー外の vLLM プラットフォームプラグインです。インストールされると、vLLM は vllm.platform_plugins エントリポイントを通じて自動的にそれを認識するため、vLLM のフォークやパッチ適用されたコピーは必要ありません。

モデルは vLLM の標準モデルランナーを通じて実行され、torch.compile(backend="tpu") でコンパイルされます。これにより、グラフが XLA に低減されます。各 TPU デバイスは、vLLM が GPU で実行されるのと同じように、独自のワーカープロセスを取得し、デバイス間の通信は torch.distributed を通じて行われます。

実際的な結果：スケジューラ、ワーカー、KV キャッシュマネージャー、API サーバーはすべて、変更されていないアップストリームの vLLM です。TPU 固有の作業はこれらのコンポーネントの下にあり、主にカスタム ops として登録され、コンパイルされたグラフ内から呼び出される一連の JAX/Pallas カーネルにあります。

はじめに：インストールとモデルの提供

インストールは、Python 3.12 を実行している Cloud TPU VM での単一の pip コマンドです。

pip install vllm-torchtpu

パッケージは、プラグインと一緒に、一致する vLLM、torch、torch_tpu、libtpu をインストールします。インストールガイドでは、uv、ソースビルド、Docker についても説明しています。

発表からの代表的な提供コマンド（FP8 Qwen3-Coder MoE モデルを 8 つの TPU チップでエキスパート並列処理と FP8 KV キャッシュを使用）:

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

最初の起動時に、構成されたシェイプの TPU グラフがコンパイルされます。これには数分かかる場合があります。VLLM_XLA_CACHE_PATH を設定すると、コンパイルされたアーティファクトが再起動間で再利用され、起動のたびに再コンパイルされるのを防ぎます。

投機的デコード、プレフィックスキャッシュ、マルチ LoRA、構造化出力、マルチモーダル入力はすべて、通常の vLLM フラグで有効にできます。クイックスタートでは、オフラインバッチ推論、マルチモーダルモデル、マルチ LoRA などの追加の例について説明しています。

Pallas カーネル：TPU 用に手動調整されるもの

vLLM TorchTPU は、コンパイラに完全に依存するのではなく、TPU での手動調整の恩恵を受ける操作にカスタム Pallas カーネルを使用します。発表されたカーネルセットには以下が含まれます：

  • MHA および GQA モデル用の Ragged Paged Attention v3。
  • DeepSeek、Kimi、GLM 用の高密度およびスパースな Multi-head Latent Attention (MLA)。
  • ハイブリッド線形アテンションモデル用の Gated DeltaNet および Kimi Delta Attention。
  • エキスパート並列処理を備えた融合 Mixture-of-Experts。
  • FP8、INT8、MXFP4、NVFP4 量子化 matmul。

スケールアウト：6 つの並列処理次元

vLLM の 6 つの並列処理次元すべてが、シングルホストおよびマルチホストの TPU スライスの両方でサポートされており、1 つのデプロイメントで組み合わせることができます。

次元 フラグ
テンソル並列処理 (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

大規模な Mixture-of-Experts モデルの場合、プロジェクトは、データ並列アテンションとエキスパート並列処理を組み合わせることで、通常、最高のスループットが得られると報告しています。例えば、単一の v7x ホストで EP を有効にした DP=8, TP=1 です。この構成により、アテンションは各デバイスにローカルに保たれ、エキスパートはスライス全体に分散されます。長いコンテキストの場合、PCP と DCP はプリフィルと KV キャッシュをデバイス間で分割します。

重点は、現在のサービスシステムに最も負荷をかけるワークロード、特にコンテキストが成長し続け、ツール出力が長くなる可能性があるマルチターンのエージェントおよびマルチモーダルトラフィックに置かれています。

KV キャッシュオフロードと分離されたサービス

vllm-torchtpu は、Python ランタイムの外部で TPU デバイスバッファを直接読み取ることができるオープンソースの C++ ライブラリである TPU Sync 上に KV 移動機能を構築します。vLLM の KV コネクタインターフェースを通じて配線され、これにより 2 つの機能が得られます：

  • ホスト DRAM へのゼロコピー KV キャッシュオフロードにより、長いマルチターンセッションのプレフィックスキャッシュが HBM 容量によって制限されません。
  • KV 転送オーバーヘッドが低い、TPU ホスト間の分離されたプリフィル/デコードサービス。プリフィルとデコードは異なるレイアウトを使用できます（例えば、プリフィルは PCP=4, TP=2、デコードは DP=8, TP=1）。KV キャッシュは転送中に再シャーディングされます。

投機的デコードとプレフィックスキャッシュ

投機的デコードは、MTP、EAGLE-3、DFlash、DSpark、n-gram、サフィックスデコードをサポートしています。自動プレフィックスキャッシュは、Qwen3.5 や Kimi K3 などのハイブリッドモデルでも機能し、線形アテンション状態が KV キャッシュと一緒にキャッシュされます。

サポートされているモデルとハードウェア

vLLM TorchTPU は、vLLM でサポートされているほとんどのモデルを提供でき、Qwen、DeepSeek、Kimi、GLM、Llama、GPT-OSS、Gemma ファミリー向けの TPU 固有のカーネル、量子化、並列処理作業が含まれています。

TPU v7x (Ironwood) が推奨ハードウェアであり、すべてのプルリクエストで CI でテストされています。TPU v6e、v5e、v5p は実験的と説明されています。ナイトリーメンテナンスマトリックスは、TPU v7x の正確性と CI からのベンチマーク結果を追跡し、Supported Models & Features ページには、検証済みチェックポイントと各機能を有効にする方法がリストされています。

ロードマップ：次のステップ

表明された野心は、主要なファミリー（Qwen、DeepSeek、GLM、Kimi、Gemma、GPT-OSS）の新しいオープンモデルが、リリース当日に TPU で実行されるべきであるということです。公開ロードマップには以下がリストされています：

  • 新しいオープンウェイトモデルをリリースから 24 時間以内にオンボードし、ベンチマークし、サポートマトリックスを公開する自動 CI パイプライン。
  • プレフィックス、キャッシュ、セッションを意識したルーティング、および llm-d TPU メモリ階層ルーティング。
  • TorchTPU エコシステム内の他のプロジェクトと共有される Tokamax カーネル用の PyTorch インターフェース。
  • vllm-omni による拡散モデルのサービス。

プロジェクトのコントリビューティングガイドを通じて、コントリビューションが明示的に歓迎されています。

コントリビューター TPU キャパシティプール

ローンチと並行して、プロジェクトはコミュニティコントリビューション TPU キャパシティプールを開設しています。vllm-torchtpu 上で開発する開発者は、無料の TPU 時間を申請でき、Inferact のコア vLLM メンテナーが設計と PR レビューを支援します。申請は、発表からリンクされている公開 Google Form を通じて行われます。

コードとドキュメントの場所

  • クイックスタート、Cloud TPU セットアップ、インストール（pip、uv、ソース、Docker）のドキュメントは、vllm-torchtpu リポジトリにあります。
  • デプロイメントリソースには、TPU v7x (Ironwood) サービングレシピ、tpu-recipes リポジトリの Ironwood 上の vLLM TorchTPU レシピ、Kubernetes デプロイメントガイドが含まれます。
  • ディスカッションは、vLLM Developer Slack の #sig-tpu チャンネルで行われます。

Sources