vLLM向けのPyTorchネイティブTPUバックエンド

［画像: vLLM TorchTPU］(https://github.com/vllm-project/vllm-torchtpu/blob/main/docs/assets/torch_tpu_dark_mode_short.png#gh-dark-mode-only) ［画像: vLLM TorchTPU］(https://github.com/vllm-project/vllm-torchtpu/blob/main/docs/assets/torch_tpu_light_mode_short.png#gh-light-mode-only)

ドキュメント | ブログ | ユーザーフォーラム | 開発者Slack (#sig-tpu)

---

vLLM-TorchTPUについて

vLLM TorchTPUは、vLLM向けのPyTorchネイティブTPUバックエンドです。これは、vllm-torchtpuとしてパッケージ化された、ツリー外のvLLMプラットフォームプラグインであり、TorchTPU（torch_tpu、torch.compile(backend="tpu")経由）と、アテンション、Mixture-of-Experts、量子化された行列乗算のためのカスタムJAX/Pallasカーネルライブラリを通じて、Google Cloud TPU上でvLLMの標準的なPyTorchモデル定義を実行します。

目標はシンプルです：モデルコードの変更なしに、TPU速度で動作する、あなたが既に使っているvLLMと同じもの — 同じCLI、同じOpenAI互換サーバー、同じモデル定義 —。

• PyTorchファースト。 モデルはvLLMのネイティブPyTorchランナーを通じて実行されます。TPUでのサービングのために他のフレームワークで再実装する必要はありません。
• TPUチューニング済みカーネル。 Pallasで記述された、Ragged Paged Attention、密および疎なMLA、Gated DeltaNet、融合MoE（SparseCoreオフロードを含む）、およびFP8/INT8/MXFP4量子化行列乗算。
• スケールアウトが組み込まれています。 TP、DP、PP、EP、およびプリフィル/デコードコンテキスト並列処理は、シングルホストおよびマルチホストTPUスライス全体で構成され、プリフィル/デコードサービングは分離されています。
• 標準的なvLLMエクスペリエンス。 スペキュラティブデコーディング、プレフィックスキャッシング、マルチLoRA、構造化出力、マルチモーダル入力は、通常のvLLMフラグを通じて機能します。

注意

JAXバックエンドをお探しですか？ vLLM向けのオリジナルのTPUバックエンドであるtpu-inferenceは、vLLM JAXにリブランドされました。これは、TPU上で純粋なJAXモデル実装をサービングするために引き続き利用可能です。vllm-torchtpuは、TPU上でPyTorchモデル定義をサービングするための推奨パスです。

サポートされているハードウェア

プロジェクトのセットアップ、クォータ、プロビジョニングコマンドについては、Cloud TPUセットアップを参照してください。

始め方

インストール

vllm-torchtpuは、Cloud TPU VM上でPython 3.12を必要とします。1つのコマンドで、vllm-torchtpuをPyPIからインストールし、対応するvLLMをインストールします。

pip install \
  --extra-index-url https://download.pytorch.org/whl/cpu \
  vllm-torchtpu

または、Docker Hub (vllm/vllm-torchtpu)から事前ビルドされたコンテナをプルします。

docker pull vllm/vllm-torchtpu:latest

uv、編集可能なソースインストール、ローカルvllmチェックアウトに対するクロスリポジトリ開発、およびDockerイメージについては、完全なインストールガイドとDockerガイドを参照してください。

モデルのサービング

他のvLLMプラットフォームで実行する場合と全く同じように、OpenAI互換サーバーを起動します。

vllm serve meta-llama/Llama-3.2-1B-Instruct \
  --tensor-parallel-size 1 \
  --max-model-len 4096

次にリクエストを送信します。

curl http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "meta-llama/Llama-3.2-1B-Instruct",
    "messages": [{"role": "user", "content": "Explain how TPUs accelerate matrix multiplication." }],
    "max_tokens": 128
  }'

ヒント

最初の起動時には、設定された形状のTPUグラフがコンパイルされるため、数分かかることがあります。再起動間でコンパイル済みアーティファクトを永続化するには、VLLM_XLA_CACHE_PATHを設定してください — コンパイルとウォームアップキャッシュを参照してください。

クイックスタートでは、vllm.LLMを使用したオフラインバッチ推論、TP/EPによるマルチデバイスサービング、マルチモーダル例をカバーしています。

サポートされているモデルと機能

vLLM-TorchTPUは、vLLMでサポートされているほとんどのPyTorchモデルをサービングでき、Qwen3 / Qwen3.5、DeepSeek-V3 / R1 / V4、Kimi-K2.5 / K3、GLM-5、Llama 3 / 4、GPT-OSS、Gemma 3 / 4、Mixtral、Qwen-VLおよびGemma VLファミリー、およびvllm-omni経由の拡散モデルに対して、TPUカーネル、量子化、並列処理のファーストクラスサポートを提供します。

• サポートされているモデルと機能 — 検証済みのモデルファミリー、機能ステータス、および量子化フォーマット。
• ナイトリーサポートマトリックス — TPU v7xのユニット、正しさ、およびベンチマークステータス。ナイトリーCIから自動更新されます。

レシピとデプロイメント

• TPU v7x (Ironwood) サービングレシピ
• TPU v6e (Trillium) サービングレシピ
• Kubernetesデプロイメント
• マルチホストサービング · 並列処理とスケーリング · 分離サービング (P/D)
• 量子化 · スペキュラティブデコーディング

🤝 貢献

バグレポート、モデルや機能のリクエスト、ドキュメントの修正、コードなど、あらゆる種類の貢献を歓迎します。

• 🐞 バグ報告または機能リクエスト: 新規イシューを開いてください。
• 👀 プルリクエストのレビュー: オープンなPRをレビューすることで、品質維持にご協力ください。
• 💻 コードの記述: 貢献ガイドから始めてください。環境セットアップ、DCOサインオフ、ready-ラベルCIゲート、およびPRライフサイクルについて説明しています。最初の貢献向けイシューは、始めるのに最適な場所です。

開発者向け参考情報: テストとCI · プロファイリング · ベンチマーク · コーディングガイドライン · 最適化のベストプラクティス

💬 お問い合わせ

• 技術的な質問や機能リクエストについては、GitHubイシューを開いてください。
• 他のユーザーとのディスカッションについては、vLLMフォーラムのTPUサポートトピックを使用してください。
• 貢献や開発の調整については、vLLM開発者Slack (#sig-tpu)に参加してください。