vLLM TT プラグインが Tenstorrent アクセラレータを LLM サービングに導入

TL;DR

vLLM TT プラグインは、vLLM サービングスタックに Tenstorrent アクセラレータを追加し、同じ OpenAI 互換 API を維持しながら、メッシュネイティブスケジューラ、デバイス内サンプリング、および単一プロセスラン データ並列設計を導入します。

TT プラグインの概要

このプラグインは、TT-Metal の ttnn パッケージがインポート可能になると自動的に有効化される、ツリー外のプラットフォームモジュールとして配布されます。クライアントコードやリクエスト形式の変更は必要ありません。OpenAI 互換 API はそのまま動作します。

対応モデルファミリー

このプラグインは、TT プレフィックスを付加した Tenstorrent ベースのアーキテクチャを登録します。モデル選択はモデル名ではなくアーキテクチャに基づいて行われるため、1つの登録で複数のリリースをカバーできます。現在対応しているのは以下の通りです:

モデルファミリー TT アーキテクチャクラス
Llama 3.1 / 3.2 / 3.3 TTLlamaForCausalLM
Llama 3.2 Vision TTMllamaForConditionalGeneration
Qwen 2.5 / Qwen 3 TTQwen2ForCausalLM, TTQwen3ForCausalLM
Qwen 3.5 / 3.6 TTQwen3_5ForConditionalGeneration
Qwen 2.5‑VL / 3‑VL TTQwen2_5_VLForConditionalGeneration, TTQwen3VLForConditionalGeneration
Mistral / Mistral 3 TTMistralForCausalLM, TTMistral3ForConditionalGeneration
Gemma 3 TTGemma3ForConditionalGeneration
Gemma 4 TTGemma4ForCausalLM, TTGemma4ForConditionalGeneration, TTGemma4UnifiedForConditionalGeneration
DeepSeek V3 TTDeepseekV3ForCausalLM
GPT‑OSS 20B / 120B TTGptOssForCausalLM

Llama 3.2 Vision、Qwen‑VL、Qwen 3.6、Mistral 3、Gemma 3 などのマルチモーダルモデルはすでにこのプラグイン経由で提供されています。

アーキテクチャ中心の登録

このプラグインにはモデルコードは含まれません。単にアーキテクチャ名を登録するのみです。実際の実装は TT-Metal にあり、各クラスは手書きの TTNN モデルをラップしています。登録はアーキテクチャに基づくため、1つのクラスで複数のモデルリリースを扱えます(例:TTQwen3_5ForConditionalGenerationQwen/Qwen3.6-27B をサポート)。

カスタムモデルはプラグインのソースを変更せずに追加できます。EXTRA_MODELS_DIRvllm_metadata.json とアダプタクラスを含むディレクトリに設定することで実現します。TT_VLLM_BUILTIN_MODELS=0 を設定すると、レジストリはユーザーが提供するモデルのみに制限されます。

Tenstorrent メッシュ vs. GPU 形状の推論スタック

Tenstorrent ハードウェアは、オンファブリックネットワークで接続されたコアとチップのメッシュ(例:n150、n300、QuietBox、Galaxy)です。プログラムは固定されたメッシュ形状に対してコンパイルされ、チップ間のデータ移動はコンパイルされたトレースに埋め込まれており、ホスト側のコレクティブ操作として発行されません。

このコンパイルモデルの主な結果は以下の通りです:

  • テンソル並列やパイプライン並列のランクなし – メッシュプログラムが並列性を直接エンコードします。MESH_DEVICE=TG フラグが通常の --tensor‑parallel-size 引数を置き換え、プラグインは -tp/-pp を拒否します。
  • ステップの粒度は完全なトレースプログラム – 各ステップは固定バッチサイズのキャプチャされたトレースを再実行するため、均一なバッチは非均一なバッチよりもはるかに安価です。
  • デバイス内でのサンプリングが可能 – メッシュプログラムは選択されたトークンを直接返すことができ、ホスト側のロジット転送が不要になります。

これらの違いは、vLLM のプラグインインターフェースに大幅な調整を必要としました。

プラグイン統合ポイント

vLLM のハードウェアプラグインメカニズム(2025年5月導入)は以下の2つのエントリポイントを提供します:

エントリポイントグループ 名前 対象
vllm.platform_plugins tt vllm_tt_plugin.entrypoints:platform_plugin
vllm.general_plugins tt_model_registry vllm_tt_plugin.entrypoints:register

platform_plugin()ttnn がインポート可能な場合にのみ TTPlatform インスタンスを返すため、純粋な CUDA 環境で誤って有効化されるのを防ぎます。

プラグインは vLLM の拡張ポイントを介して、Tenstorrent 固有のランタイムクラスを差し替えます:

vLLM 設定フィールド TT 実装
parallel_config.worker_cls vllm_tt_plugin.worker.TTWorker
scheduler_config.scheduler_cls vllm_tt_plugin.scheduler.TTScheduler または vllm_tt_plugin.lane_scheduler.TTLaneCoordinator

デバイス固有のオプションは、vLLM の汎用 --additional-config 名前空間を介して渡されます。例:

--additional-config.tt.sample_on_device_mode all
--additional-config.tt.fabric_config FABRIC_1D_RING

Tenstorrent 固有のコードは vLLM コアに存在しないため、アップストリームリリースとの前向き互換性が保証されます。

フェーズベーススケジューリング

vLLM のトークン予算スケジューラとは異なり、Tenstorrent パスは各スケジューリングステップを以下の3つの均一な結果のいずれかに制限します:

  1. prefill-only
  2. decode-only
  3. empty

混合 prefill と decode のバッチは許可されません。長いプロンプトは複数の prefill-only ステップに分割され、decode-only ステップが交互に挿入されて他のリクエストの進行を維持します。この設計により、コンパイルされたメッシュプログラムにとって不可欠なトレースの安定性が保たれます。

フェーズ分割の利点

  • 各ステップ形状ごとに1つのコンパイル済みトレースを再利用可能。
  • 大規模な GPU フレットで使われる「非統合サービング」パターンを、単一エンジン内で適用。

デメリット

  • decode リクエストは各 prefill チャンクを待つ必要があり、ステップレベルのレイテンシペナルティが発生。
  • スケジューラはステップ間でモードを切り替える必要があり、小さなポリシーオーバーヘッドが発生。

この設計は拡張可能であり、将来のバージョンで混合形状のトレースをキャプチャできるようにする可能性があります。

Galaxy 上の単一プロセスラン データ並列

Galaxy(32チップメッシュ)では、一部のモデルがメッシュ全体にまたがる 単一実行 プログラムとして実行されます。メッシュへの送信が1回しかできないため、従来のマルチプロセスデータ並列は適用できません。

解決策は プロセス内ラン DP です:

  • TTLaneCoordinator は1つのラン(デフォルトで4ラン)ごとに1つの TTScheduler を作成。
  • 各ランは独自の待機/実行キュー、KVキャッシュ、ブロックID空間を維持。
  • リクエストは最も負荷の低いランに割り当てられ、そのランに束縛されたままです。
  • 各ステップで、コーディネータは すべてのラン に対して共通のモード(prefill または decode)を選択。作業のないランは空のスライスを提供。
  • マージされたバッチは1回の送信でデバイスに送られ、結果は内部でランに再分割されます。

これにより、以前のマルチプロセスアプローチで問題となった高コストのプロセス間スキャターガザーが排除されます。

エッジケース

prefill ステップで KV 圧力によりトークンが0件に届いた場合、他のランに decode 作業があると、デッドロックを避けるために decode モードで再試行されます。

ユーザー向けフラグ

通常の vLLM フラグが再利用されます:

MESH_DEVICE=TG \
TT_LLAMA_TEXT_VER=llama3_70b_galaxy \
VLLM_RPC_TIMEOUT=900000 \
python examples/server_example_tt.py \
  --model "meta-llama/Llama-3.3-70B-Instruct" \
  --data_parallel_size 4 \
  --max_num_seqs 8 \
  --async-scheduling \
  --additional-config.tt.dispatch_core_axis col \
  --additional-config.tt.sample_on_device_mode all \
  --additional-config.tt.fabric_config FABRIC_1D_RING \
  --additional-config.tt.worker_l1_size 1344544 \
  --additional-config.tt.trace_region_size 220000000

--data_parallel_size 4 は現在、4つのプロセス内ランを作成し、それぞれが --max_num_seqs リクエストを処理できます。

デバイス内サンプリングと自動フォールバック

sample_on_device_mode が設定されている場合、メッシュプログラムはトークン選択を行い、直接トークンを返します。バッチがデバイスが表現できない機能(logprobs、ペナルティ、カスタムロジットプロセッサなど)を必要とする場合、そのバッチのみ vLLM のホスト側サンプラーにフォールバックします。always_compat_sampling フラグはデバッグ用にホスト側サンプリングを強制します。

非同期デコードオーバーラップ

このプラグインはデコード/ホストオーバーラップを提供しますが、非同期ホストリードバック としてのみ、別途デバイス実行スレッドとしてではありません:

  1. ブロッキングなしでデコード作業を送信(read_from_device=False)。
  2. 非ブロッキングホストリードバックを開始(async_read=True)。
  3. 結果のイベントを保存。
  4. 最終化時に、ホストテンソルに変換する前に ttnn.event_synchronize() で同期。

深さ2のインフラインキューにより、ステップ N+1 のスケジューリングがステップ N のリードバック進行中に可能になります。オーバーラップは安定した形状、デバイス内サンプリング、構造化出力の追跡が不要なステディステート生成においてのみ維持されます。prefill は同期のままです。

現在の制限

このプラグインは設定を早期に検証し、サポートされていない組み合わせを拒否します:

  • テンソル並列とパイプライン並列は、vLLM ランクではなくメッシュ形状で表現されます。
  • 予測デコード、LoRA、プロンプト logprobs はまだサポートされていません。
  • プレフィックスキャッシュは、モデルが宣言した場合にのみ利用可能です。
  • 非同期デコードオーバーラップには、モデルが宣言した機能が必要です。
  • 標準的なマルチプロセス DP は MoE モデルをサポートしていません。代わりにラン-DP が使用されます。
  • マルチホストサービングは実装されていません。

これらは現在の TT-Metal ランタイムおよびモデル実装の制限であり、ハードウェア上の硬い制約ではありません。

セットアップ手順

  1. 公式ガイドに従って TT-Metal をインストール。

  2. プラグインをクローンしてインストール:

    git clone https://github.com/tenstorrent/vllm-tt-plugin.git
    cd vllm-tt-plugin
    source docs/install-vllm-tt.sh
    

    このスクリプトは、TT-Metal 環境内でのバージョン 0.26.0 に対して vLLM をビルドします。

  3. モデルを提供:

    MESH_DEVICE=T3K VLLM_RPC_TIMEOUT=100000 python examples/server_example_tt.py
    
  4. 任意の OpenAI 互換クライアントでクエリ:

    curl http://localhost:8000/v1/completions \
      -H "Content-Type: application/json" \
      -d '{"model": "meta-llama/Llama-3.1-70B-Instruct", "prompt": "San Francisco is a", "max_tokens": 32}'
    

ロードマップ

  • さらに多くのモデルファミリーに非同期デコードサポートを拡張。
  • 追加モデルでのプレフィックスキャッシュとラン-DP RoPE ハンドリングを有効化。
  • メッシュ側のドラフト/検証パイプラインが安定したら、予測デコードを実装。
  • 複数ホストサービングを追加して、単一マシンを超えるスケーリングを実現。

謝辞

このプラグインは、Ascend チームが貢献した vLLM プラットフォームプラグインメカニズムと、Spyre チームのプラグイン可能スケジューラ設計に基づいています。メッシュアーキテクチャに対応できるよう、拡張ポイントを汎用的に保つために vLLM メンテナーズに感謝します。

貢献者:Viktor Puš、Tomasz Cheda、Sanjar Adylov、Salar Hosseini。ラン-DP ユーザーフェースや優先モデルファミリーに関するフィードバックは、GitHub のイシューまたは vLLM Slack で歓迎します。

Sources