vLLM TorchTPU のご紹介

本日リリースするのは vLLM TorchTPU (vllm-torchtpu)、Google Cloud TPU上でvLLMを実行するための新しいバックエンドです。これはGoogleの TorchTPU stack上で実行され、すでにvLLMに存在するモデル定義を実行します。

本日リリースするのは vLLM TorchTPU (vllm-torchtpu)、Google Cloud TPU上でvLLMを実行するための新しいバックエンドです。これはGoogleの TorchTPU stack上で実行され、すでにvLLMに存在するモデル定義を実行します。そのため、現在vLLMで動作しているほとんどのモデルは、モデルコードを変更することなくTPUで提供できます。

この研究は次のものから生まれたものです: InferactとGoogleのエンジニアリングパートナーシップ先月発表しました。このプロジェクトの目標は次のとおりです。

TPUへの新しいモデルの取り込みを、vLLMの既存のモデルコードを再利用することで容易にします。
Pallas カーネルによる attention、Mixture-of-Experts、量子化 matmul を活用し、TPU の推論性能をオープンソースで引き上げる。
vLLMの体験をそのままに:他のハードウェアで使っているのと同じCLI、OpenAI互換サーバー、スケジューラ、KVキャッシュマネージャーをそのままご利用いただけます。

仕組み

vllm-torchtpu は、ツリー外の vLLM プラットフォームプラグインです。インストールすると、vLLM が自動的にそれを検出します。 vllm.platform_plugins エントリーポイント。

vLLM の標準モデルランナーを通じて実行され、以下でコンパイルされたモデルです。 torch.compile(backend="tpu")、XLAにダウンコンパイルされます。各TPUデバイスは、GPU上でvLLMが実行されるのと同じように、それぞれ専用のワーカープロセスを受け持ち、デバイス間の通信は torch.distributedつまり、スケジューラ、ワーカー、KVキャッシュマネージャー、APIサーバーはすべて変更されていない上流のvLLMです。TPU固有の処理はその下の層にあり、主にカスタムオペレーションとして登録され、コンパイル済みグラフの内部から呼び出される一連のJAX/Pallasカーネルにあります。

はじめに

インストール
```
pip install vllm-torchtpu
```
これにはCloud TPU VM上でPython 3.12が必要で、対応するバージョンのvLLMがインストールされます。 torch, torch_tpu、そして libtpu プラグインとともに。インストールガイドもカバーしています uv、ソースビルド、そして Docker.
モデルを提供する
```
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 グラフのコンパイルが行われます。これには数分かかることがあります。Set VLLM_XLA_CACHE_PATH 再起動をまたいでコンパイル済みアーティファクトを再利用するには(参照: コンパイル & ウォームアップキャッシュ).

投機的デコーディング、プレフィックスキャッシング、マルチLoRA、構造化出力、マルチモーダル入力は、いずれも通常の vLLM のフラグで有効化できます。詳細な例（オフラインバッチ推論、マルチモーダルモデル、マルチLoRA など）については、クイックスタートをご覧ください。

カーネル

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 量子化行列乗算

スケールアウト

私たちは、今日のサービングシステムに最も大きな負荷をかけるワークロード、特にマルチターンのエージェンティックなトラフィックやマルチモーダルトラフィックに重点を置いてきました。これらではコンテキストが成長し続け、ツールの出力が長くなることがあります。

並列化

vLLM の 6 つの並列化ディメンションはすべて、シングルホストおよびマルチホストの TPU スライスの両方でサポートされており、組み合わせることもできます：

ディメンション	フラグ
テンソル並列（TP）	`--tensor-parallel-size`
データ並列( DP )	`--data-parallel-size`
パイプライン並列(PP)	`--pipeline-parallel-size`
Expert Parallelism(EP)	`--enable-expert-parallel`
Prefill Context Parallelism（PCP）	`--prefill-context-parallel-size`
Decode Context Parallelism(DCP)	`--decode-context-parallel-size`

大規模なMoEモデルでは、通常、エキスパート並列と組み合わせたデータ並列のアテンション(attention)から最高のスループットが得られます(例えば、 DP=8, TP=1 「単一の v7x ホスト上で EP を使用する場合)。これにより、各デバイスへの注意の局所性が保たれ、エキスパートがスライス全体に分散されます。長いコンテキストでは、PCP と DCP がプレフィルと KV キャッシュをデバイス間で分割します。この並列処理ガイドさらに詳しい詳細と設定例が掲載されています。

KVキャッシュのオフロードと分離型サービング

vllm-torchtpu 用途 TPU SyncPythonランタイムの外部でTPUデバイスバッファを直接読み取ることができるオープンソースのC++ライブラリです。vLLMのKVコネクタインターフェースを通じて、これにより以下が可能になります:

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

以下をご覧ください: 分離型サービングガイドセットアップする方法。

スペキュレーティブデコーディングとプレフィックスキャッシング

投機的デコーディングは、MTP、EAGLE-3、DFlash、DSpark、n-gram、およびsuffix decodingをサポートしています(ガイド). 自動プレフィックスキャッシュは、Qwen3.5やKimi K3のようなハイブリッドモデルでも機能し、linear-attention状態がKVキャッシュとともにキャッシュされます。

モデル、機能、そして今後の展望

vllm-torchtpu vLLMがサポートするほとんどのモデルに対応でき、Qwen、DeepSeek、Kimi、GLM、Llama、GPT-OSS、Gemmaファミリー向けにTPU固有のカーネル、量子化、並列化の作業が行われています。TPU v7x（Ironwood）が推奨ハードウェアであり、すべてのプルリクエストごとにCIでテストされています。v6e、v5e、v5pは実験的です。

その対応モデルと機能 validatedなチェックポイントと機能の全リスト、およびそれぞれの有効化方法はこのページに記載されています。The nightly サポートマトリクス CIからTPU v7xの正確性およびベンチマーク結果を追跡します。

次は何が起こるのか?

私たちは、主要モデルファミリー(Qwen、DeepSeek、GLM、Kimi、Gemma、GPT-OSS)の新しいオープンモデルが、リリースされたその日にTPU上で動作することを目指しています。近日公開される予定の一部については、公開ロードマップ:

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

コントリビューションを歓迎します!

コントリビューター向けTPU容量

また、コミュニティコントリビューションTPU容量プールも開放します。あなたが以下で構築している場合 vllm-torchtpu、無料のTPU利用時間を申請でき、InferactのvLLMコアメンテナーが設計やPRレビューをサポートします。

コントリビューター向けTPUアクセスはこちらから申請してください。

試してみる

ご質問がある場合や参加をご希望の場合は、 #sig-tpu の vLLM Developer Slack で私たちを見つけてください.

原文の出典

vLLM Blog

原文を読む

内容について

原文の公開と権利は出典元に帰属します。

機械翻訳 · 原文をご参照ください