vLLM、TML InklingのDay-0サポート開始

このたび、vLLM が TML Inkling モデルを Day 0 で正式にサポートすることを発表できることを嬉しく思います。`thinkingmachines/Inkling-NVFP4` および `thinkingmachines/Inkling` (BF16) の両モデルが、最適化されたパフォーマンスと完全な機能互換性を備えてサポートされています。

TML Inkling は Thinking Machines Lab が訓練した 1T パラメータのマルチモーダルモデルです。このモデルはテキスト、画像、音声入力をネイティブに受け付け、最大1M のコンテキスト長でテキストを生成します。相対アテンション、短距離畳み込み、共有エキスパートシンクといった複数の新規アーキテクチャコンポーネントを導入しており、これらはすべて vLLM に効率的に統合されました。

vLLM を利用すると、4 GB200 GPU 上で MTP を使用した場合はユーザーあたり最大380 tok/s、MTP なしの場合は140 tok/sでモデルを実行できます。また、LoRA、TP/DP/EP/PP パラライゼーション、プレフィックスキャッシング、分離型サービングなど、機能互換性も完全に提供します。包括的なベンチマークを通じて、両モデルの精度とツール解析能力を検証済みです。

統合に関する PR は こちら で確認できます。以下の手順でモデルを実行してください:

export VLLM_USE_V2_MODEL_RUNNER=1
export FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED=1

vllm serve thinkingmachines/Inkling-NVFP4 \
      --tokenizer-mode inkling \
      --reasoning-parser inkling \
      --tool-call-parser inkling \
      --enable-auto-tool-choice \
      --tensor-parallel-size 8 \
      --speculative-config '{"method": "mtp", "num_speculative_tokens": 8}' \
      --kernel-config.enable_flashinfer_autotune=False \
      --trust-remote-code

TL;DR

vLLM は TML Inkling に対して強力な Day-0 サポートを提供します:

対応モデル：`thinkingmachines/Inkling-NVFP4` と `thinkingmachines/Inkling` (BF16) の両方をサポートしています。

ハードウェア：NVIDIA Blackwell および Hopper GPU に対応。より広範なハードウェアへの対応も進行中です。

モダリティ：テキスト・画像・音声の入力に対し、テキストを出力します。

コンテキスト長：ネイティブで最大 1M トークンを処理可能（Tinker では 64K と 256K のコンテキストウィンドウを提供）。

機能：LoRA、推測的デコーディング (MTP)、TP/DP/EP/PP、プレフィックスキャッシング、非集約型サービングなどに対応しています。

最適化：Sconv を意識した TP シャーディング、低遅延の融合された集合演算、カーネル融合、マルチストリーミング、PDL などを実装。

パフォーマンス：4 GB200 GPU 上で、MTP 使用時は 380 tok/s/user、非使用時は 140 tok/s/user を達成。

精度：MMAU、MMMU-Pro、BFCL、NIAH-1M、HLE によるモデル品質とツール解析の検証済み。

モデルアーキテクチャ

図 1. TML Inkling モデルアーキテクチャ（RMSNorm や残差接続などの一部の演算は省略）。

モダリティ。 TML Inkling は、1 兆パラメータを備えるネイティブなマルチモーダルモデルです。テキストと画像に加え、オーディオ入力も受け付け、テキストを生成します。このモデルは、TML のインタラクションモデルのプレビューで説明されている通り、極めて軽量な画像エンコーダー（hMLP）とオーディオ埋め込み（dMel）を使用しています。得られた埋め込み表現は、デコーダ型の Transformer バックボーンによって処理されます。

アテンション。 バックボーンには 66 レイヤーが用意されており、そのうち 11 レイヤーがフルアテンション、残りの 55 レイヤーがスライディングウィンドウ・アテンションです。このようにスライディングウィンドウ・アテンションを多用している点が、モデルの100 万トークンというコンテキスト長を実現する効率性の源となっています。すべてのアテンション層は、ヘッドサイズ 128 のグループ化クエリ型アテンション（GQA）を採用しています。

Inkling の特徴的な設計選択の一つが、位置情報の表現に「相対アテンション」を採用している点です。RoPE を使用するのではなく、ソフトマックス前のアテンションロジットに学習可能な相対位置項を追加します。詳細は TML のブログ記事をご覧ください。

Sconv。 Inkling は、ウィンドウサイズ 4 の「短縮畳み込み（sconv）」を積極的に活用しています。各レイヤーには 4 つの sconv モジュールが含まれており、アテンションキー、アテンション値、アテンション出力、そして MoE の出力に対してそれぞれ適用されます。Sconv は計算量とメモリオーバーヘッドを最小限に抑えつつ、小さな局所アテンションのような役割を果たします。

MoE. 各層には 256 のルーティング専門家が（トップ 6 が選択される）配置され、さらに 2 つの共有専門家も加わります。その結果、1 トークンあたり合計 8 個の専門家が処理に関与します。既存モデルとの決定的な違いは、「エキスパート・シンク（expert sink）」という概念を導入している点です。この仕組みでは、2 つの共有専門家はルーティングスコアの計算に組み込まれ、確率質量の一部を吸収する役割を果たしますが、トップ 6 の候補からは除外されます。

thinkingmachines/Inkling-NVFP4 では、ルーティング専門家のみが NVFP4 に量子化されており、共有専門家や QKVR リニア層など他のすべてのパラメータは BF16 で保持されています。一方、thinkingmachines/Inkling の場合も MoE 重みは BF16 です。

MTP。 Inkling は推測デコーディング用に 8 つの MTP ヘッド を搭載しており、これにより 1 フォワードステップで最大 9 トークンを生成できます。これらの MTP ヘッドは連鎖（chained）構造になっており、各ヘッドが前のヘッドから出力された隠れ状態とサンプリングされたドラフトトークンを入力として受け取ります。各 MTP ヘッドは、フルまたはスライディングウィンドウアテンションと密な MLP を備えた 1 レイヤーの Transformer です。すべての MTP 重みも BF16 で実装されています。

vLLM Integration & Optimization

vLLM は一連の最適化を通じてこのモデルを効率的に実装しています。主な特徴は以下の通りです：

sconv キャッシュの管理

短畳み込み（short convolution）では、直前の W-1 トークンの隠れ状態を保持する必要があります。vLLM はこれを、仮想のスライディングウィンドウアテンション層の KV キャッシュとして扱うことで管理しています。これにより、vLLM の統一された KV キャッシュマネージャーを通じて sconv キャッシュを洗練された形で処理できます。ウィンドウ外にある状態は削除可能としてマークされ、プレフィックスキャッシュも sconv キャッシュとシームレスに連携します。

sconv を意識した TP シャーディング

このモデルに対する単純な TP（Tensor Parallelism）の実装としては、すべての GPU で「全結合（all-reduce：例では o_proj の後）」→「sconv」→「残差接続」→「RMSNorm」という順序が考えられます。しかし、これだと sconv が各 GPU 上の完全な隠れ状態に対して適用され、sconv の計算量とキャッシュの両方がランク間で重複してしまいます。

この重複を解消するため、vLLM はモデルを異なる方法でシャーディングします。sconv はチャネル次元に沿って独立して動作するため、チャネル方向に sconv を分割します。具体的には、全結合（all-reduce）の代わりに、チャネル次元に対して「reduce-scatter」と「all-gather」を行います。各 GPU は sconv キャッシュの一部のみを保持し、自らのチャネルスライスの計算のみを実行します。この考え方はシーケンス並列化に似ていますが、シャーディングの対象がトークン次元ではなくチャネル次元である点が異なります。

低遅延の融合型集合通信。 vLLM はこの新しいシャードリング方式を最適化するため、複数の融合カーネルを実装しています。具体的には、FlashInfer の低遅延オールリダクションカーネルにおける Lamport プロトコル設計を拡張することで、低遅延のリデューススキャッターとオールギャザーカーネル（周辺演算との融合）を開発しました。Lamport プロトコルにより、カーネルは明示的なバリアではなくデータ値のポーリングで同期を行うため、バッチサイズ 1 の場合のカーネル実行時間を 40 µs から 8 µs（5 倍短縮） に削減しています。

シアーバイアス付き FA4。 リラティブアテンションはメモリアクセスパターンを複雑にし、アテンションカーネルの計算パイプラインを大幅に遅延させます。これを克服するため、TML は Colfax Research と共同で、シアーバイアス技術を採用した新しい FA4 カーネル をリリースし、vLLM が直接統合しました。さらに vLLM は、バッチサイズ、TP サイズ、KV 長を考慮して構成ごとに FA4 の num_splits パラメータを選択し、パフォーマンスの最大化を図っています。

MTP KV キャッシュの再計算。 MTP ヘッドはそれぞれ、直前のヘッドが生成したドラフトトークンを入力として受け取るため、ドラフトトークンが拒否されるとその KV キャッシュは陳腐化してしまいます。vLLM はこれを慎重に処理します。最後に数トークンのベースモデルの隠れ状態をキャッシュし、拒否サンプリング後に承認されたトークンを用いて MTP ヘッドを再実行する仕組みです。

これらに加え、vLLM のモデル実装にはカーネル融合、PDL、マルチストリーミングが追加されており、驚異的な速度を実現しています。詳細は our PR をご確認ください。

パフォーマンス

上記の広範な最適化により、vLLM は 4× GB200 GPU 環境で、MTP8（平均受容長 4.5）を使用する場合 380 tok/s/user、MTP なしの場合は 140 tok/s/user を達成しました。測定は SPEED-Bench からサンプリングした 8K トークンの入力プロンプトに対して、各リクエストで 1K トークンを生成する条件で行われました。

精度評価

vLLM の実装が正しいことを確認するため、あらゆるモダリティと機能を網羅した包括的なベンチマークを実施しました：

• オーディオ: MMAU

• ビジョン: MMMU-Pro

• ツール呼び出し: BFCL

• 推論: HLE

• 長文コンテキスト: NIAH

vLLM はあらゆる項目でリファレンス実装と同等の性能を発揮しました。特に長文コンテキストでは、221K トークンまでリファレンスと完全に一致し、513K まで誤差は約 1 パーセントポイント以内です。最も極端なコンテキスト長（800K 以上）においては、このベンチマークの NIAH スコアでランごとのばらつきが大きくなる傾向がありますが、その領域での再現性の向上に取り組んでいます。

ロードマップ

前述の通り、vLLM は TML Inkling に対して Day-0 のサポートをすでに提供しています。今後はさらにいくつかの改善に取り組んでいきます。

• グローバルアテンションでの FP8 対応: 現在、Inkling はグローバルアテンションに BF16 を使用していますが、これが計算能力と KV キャパシティの両面でボトルネックになる可能性があります。そこで、新しい FA4 カーネルを改修することで FP8 の導入を検討しています。

• 画像エンコーダーとオーディオエンコーダーへの CUDA Graphs 適用: 現在のところ、画像およびオーディオエンコーダーはイーガーモードで動作しています。これらはプレフィル段階で実行されるため通常は大きな問題にはなりませんが、CPU オーバーヘッドを完全に排除するために、これらのエンコーダーにも CUDA Graphs を適用する計画です。

• AMD GPU への対応: 新しい相対アテンション機構に専用のカーネルが必要となるため、現時点ではこのモデルに対する AMD GPU のサポートは提供されていませんが、近日中にサポートを開始する予定です。