vLLM TML Inkling サポート

vLLM TML Inkling サポート

vLLM は現在、Thinking Machines Lab が開発した 1T パラメータのマルチモーダルモデルである TML Inkling の Day-0 サポートを提供しています。この統合により、thinkingmachines/Inkling-NVFP4 と thinkingmachines/Inkling (BF16) バージョンの両方に対して高性能な推論が可能になり、ネイティブコンテキスト長が最大 100 万トークンまでのテキスト、画像、音声入力をサポートします。

パフォーマンスとハードウェアベンチマーク

4 つの NVIDIA GB200 GPU の構成では、vLLM は Multi-Token Prediction (MTP) を使用すると 380 tok/s/user、使用しない場合は 140 tok/s/user までのスループットを達成します。これらの結果は、SPEED-Bench からの 8K 入力トークンのプロンプトと、リクエストあたり 1K 出力トークンを使用して測定されました。

現在のサポートは NVIDIA Blackwell と Hopper GPU に最適化されていますが、vLLM は AMD GPU を含むその他のハードウェアへのサポート拡張に積極的に取り組んでいます。これらの GPU は現在、モデルの相対注意メカニズム用の専用カーネルを必要としています。

TML Inkling モデルアーキテクチャ

TML Inkling は 1T パラメータのネイティブマルチモーダルデコーダー専用 Transformer です。そのアーキテクチャにはいくつかの特徴的なコンポーネントがあります：

  • Multimodal Inputs: モデルはテキスト、軽量 hMLP エンコーダーを介した画像、および dMel 埋め込みを介した音声を受け入れます。
  • Hybrid Attention: バックボーンは 66 層からなり、1M コンテキスト長での効率を維持するために 11 層のフルアテンション層と 55 層のスライドウィンドウアテンション層で構成されます。すべての層はヘッドサイズ 128 の Grouped-Query Attention (GQA) を使用します。
  • Relative Attention: Inkling は Rotary Positional Embeddings (RoPE) の代わりに、プリソフトマックスアテンションロジットに加算される学習済みの相対位置項を使用します。
  • Short Convolution (Sconv): 各層は、注意キー、注意値、注意出力、および MoE 出力に適用されるウィンドウサイズ 4 の sconv モジュール 4 つを使用し、オーバーヘッドを最小限に抑えたローカルアテンションを提供します。
  • Mixture of Experts (MoE): 各層は 256 個のルーティングエキスパート (top-6) と 2 個の共有エキスパートを特徴とします。Inkling は「エキスプシンク」を導入し、ここで 2 つの共有エキスパートはルーティングスコア計算に参加して確率質量を吸収しますが、top-6 の選択からは除外されます。
  • Speculative Decoding (MTP): モデルは 8 つの連鎖した MTP ヘッド (単層 Transformer) を含み、これによりフォワードステップごとに最大 9 トークンの生成が可能になります。

In the Inkling-NVFP4 バリアントでは、ルーティングエキスパートのみが NVFP4 に量子化され、共有エキ、共有エキスパートと qkvr linears は BF16 のままです。

vLLM テクニカル最適化

vLLM は、Inkling の独自のアーキテクチャに対応するためにいくつかの専門的な最適化を実装しています：

Sconv キャッシュと TP シャーディング

短い畳み込み (sconv) キャッシュを管理するために、vLLM はそれを仮想スライドウィンドウアテンション層の KV キャッシュとして扱い、統一された KV キャッシュマネージャーに統合します。

GPU 全体で sconv の計算とキャッシュの重複を避けるために、vLLM は Sconv-aware TP sharding を使用します。出力投後の標準的な all-reduce の代わりに、vLLM はチャネル次元に対して reduce-scatter と all-gather を使用します。これにより、各 GPU はチャネル次元に適用されたシーケンスパラレルism と同様に、自身のチャネルスライスのみを保存および計算します。

低レイテンシコレクティブとカーネル

  • Fused Collectives: vLLM は Lamport プロトコル (データ値ポーリング) に基づく低レイテンシの reduce-scatter と all-gather カーネルを使用し、バッチサイズ 1 のカーネル時間を 40 µs から 8 µs に削減します。
  • FA4 Kernel: 相対注意のメモリアクセスパターンを最適化するために、vLLM は「sheared-bias」技術を用いた新しい FA4 カーネルを統合します。vLLM はバッチサイズ、TP サイズ、および KV 長に基づいて num_splits ファクターを動的に選択します。
  • MTP KV Cache Management: MTP ヘッドは以前のドラフトトークンに依存するため、vLLM はベースモデルの隠れ状態をキャッシュし、リジェクションサンプリングに従って受け入れられたトークンで MTP ヘッドを再実行します。

精度と検証

vLLM の実装は、複数のベンチマークにおいて参照実装と一致します：

  • Multimodal & Reasoning: MMAU (Audio)、MMMU-Pro (Vision)、BFCL (Tool calling)、および HLE (Reasoning) を介して検証済みです。
  • Long Context: NIAH を使用して検証済み。vLLM は 221K トークンまで参照と完全に一致し、513K トークンまで約 1 パーセンテージポイント (pp) 内に収まります。極端な長さ (800K+ では) では、実行間のばらつきが大きくなることが観測されました。

フューチャーロードマップ

vLLM は TML Inkling に対して以下の拡張を計画しています：

  • FP8 Global Attention: FA4 カーネルの変更を通じて計算と KV キャッシュのボトルネックを削減するため、グローバルアテンションでの FP8 の使用を検討しています。
  • CUDA Graphs: プリフィル中の CPU オーバーヘッドをなくすために、画像および音声エンコーダーに CUDA グラフを適用しています。
  • AMD Support: AMD GPU との互換性を可能にするために必要な相対注意カーネルを開発しています。

Sources