vLLM における Kimi K3: NVIDIA GB300 向け専用 CuTe DSL デコードが 842 トークン/秒を達成
vLLM は Kimi K3 向けの新しいデコードパスを公開しました。このパスでは、モデルのデコード操作のほぼすべてを、NVIDIA GB300 向けにチューニングされたカスタム CuTe DSL カーネルに置き換え、DSpark 投機的デコードを使用して 16 台の GB300 GPU でバッチ 1 あたり 842 トークン/秒 を達成しています。これは、同じ構成の標準 vLLM と比較して 1.73 倍高速 です。この投稿では、この作業の背後にあるプロセス変更も紹介しています。提案されている「高速トラックモデル最適化 RFC」(vllm-project/vllm issue #59665) は、行単位の人間によるレビューを、分離、範囲の制限、体系的な正確性検証、再現可能なパフォーマンス証拠と交換するものです。
この貢献は、8 または 16 GPU の NVIDIA Blackwell GPU、1 から 64 のデコードバッチ、および長いコンテキストのためのオプションのデコードコンテキスト並列処理を対象としています。GPU はマルチキャストをサポートする単一の NVLink ドメインを共有します。Kimi K3 のプルリクエストは vllm-project/vllm #61867 です。
Kimi K3 デコードパス: レイヤーあたり 23 カーネルではなく 13 カーネル
専用デコードパスは、Kimi K3 デコードステップのほぼすべての操作を融合し、カスタマイズされた CuTe DSL カーネルに置き換えます。レイヤーごとに測定すると、標準 vLLM は 130 マイクロ秒と 23 カーネル を使用するのに対し、専用デコードパスは 69 マイクロ秒と 13 カーネル を使用します。
最適化作業はエージェント支援で行われました。コーディングエージェントを使用して、複雑な GPU カーネルを作成し、演算子間の融合を実行し、計算と通信をオーバーラップさせ、バッチサイズと GPU 構成にわたって実行を調整しました。結果として得られるパスは、8 または 16 GPU にわたるアテンション、MoE、および通信をカバーします。
1. ワークロードに合わせてカーネル形状を調整
デコードバッチサイズが異なると、異なるカーネル選択が必要になります。アテンション入力プロジェクションでは、Kimi K3 パスはバッチサイズ 1〜2 で専用の小行列カーネルを使用し、より大きなバッチでは分割 K カーネルを使用します。これにより、行列乗算がより多くの GPU ワーカーに分割されます。これらのプロジェクションは、GPU リソースを完全に活用し、メモリ読み取り、通信、計算を可能な限りオーバーラップするように、適切なサイズのカーネルを起動するように調整されています。
KDA は、各 GPU のバッチサイズとアテンションヘッド数に合わせてスケジュールを調整した最適化された CuTe DSL デコードカーネルも使用します。
長いコンテキストの場合、デコードコンテキスト並列処理 (DCP) は KV キャッシュを GPU 間で分割します。各 GPU はローカル部分でアテンションを計算し、調整された融合カーネルが部分結果を v-proj とゲーティングと結合します。
2. 計算と通信のオーバーラップ
Kimi K3 パスは、アテンション出力プロジェクションを all-reduce と後続のアテンション残差計算に直接融合します。GPU カーネルは部分結果を NVLink アクセス可能なバッファに公開し、受信カーネルはそれらを結合して計算を続行します。一部の小バッチ構成では、プロデューサーとコンシューマーを単一のカーネルに融合します。データ交換には、オーバーヘッドが最小限のフェンスレスアプローチを使用します。
より大きなバッチでは、実装は 2 段階のリダクションに切り替わります。GPU はまず出力の割り当てられた部分を結合し、次に結果を配布します。これにより、GPU 間のデータ移動が削減されます。
3. MoE スケジューリングとオーバーラップの改善
エキスパートはすでにテンソル並列処理を使用していますが、ベースラインはすべての GPU で完全なルータースコアと潜在入力プロジェクションを計算します。Kimi K3 パスはこれらの計算を TP ランク間で分割し、各ランクのシェアを 1 つの行列乗算に結合します。GPU は NVLink 経由で結果を交換します。これにより、エキスパート実行が開始される前の繰り返し計算、重み読み取り、カーネル起動が削減されます。
共有エキスパートは別のストリームで実行されます。その作業は、構成に応じてルーティングまたはルーティングされたエキスパート計算とオーバーラップします。スケジュールは、各操作が使用する GPU リソースの量を制御し、競合なしに並列処理のバランスを取るように調整されています。実装の詳細は experts_runtime.py にあります。
4. 次の重みをキャッシュにプリフェッチ
小バッチ推論はメモリバウンドであることがよくあります。つまり、メモリから重みをロードする速度によって制限されます。通信などの一部のフェーズでは、利用可能なメモリ帯域幅を完全に使用できない場合があります。通信が実行されている間、サイドストリーム上の専用カーネルが次の操作の重みをプリフェッチするため、通信が完了すると、これらの重みのキャッシュヒット率がはるかに高くなります。
エージェント最適化された推論コードが手動レビューを上回る理由
エージェント最適化された推論コードは、vLLM のレビューおよびメンテナンスワークフローを上回るペースで進んでいます。すべてのカーネルのすべての行を手動で理解し、潜在的な問題を特定し、サポートされているすべての構成をテストすることは非常に時間がかかります。その作業をスキップすると、本番環境で失敗する可能性のあるバグのあるソフトウェアをリリースするリスクがあります。
提案された高速トラック RFC は、パフォーマンス作業の編成方法を変更することでこの問題に対処します。エージェントによる最適化に加えて、vLLM は信頼性とパフォーマンスを自動的に検証するシステムと、設計によって分離とリスク最小化を強制するように構造化されたソフトウェアを望んでいます。述べられている指針は、リスクを管理しながらレビュー負担を軽減することです。
私たちの指針は、リスクを管理しながらレビュー負担を軽減することです。
このアプローチは 4 つの原則に基づいています。
原則 1: 変更を分離する
モデル固有の最適化はモデルのコード内に留まります。これにより、影響を受けるパスを特定しやすくなり、レビュー担当者が考慮すべき爆発半径が減少します。分離は実行時にも維持する必要があります。レビュー担当者は、最適化されたパスがいつ選択されるか、どの状態を変更するか、実行が標準パスにどのように戻るかを追跡できる必要があります。
分離は意図的に一部の最適化を除外します。入力準備、サンプリング、KV キャッシュ管理は、分離が除外するオーバーヘッドを引き起こす可能性があります。これは大部分が意図的です。これらのコンポーネントはほとんどの vLLM デプロイメントで共有されており、特定のワークロードに最適化するために変更すると、他の数十の安定した vLLM 構成を壊すリスクがあります。これらのコンポーネントの作業には、高速トラックが許可するよりもはるかに密接にコアメンテナーが関与する必要があります。
原則 2: 範囲を制限するが、ユーザーの機能を維持する
最適化は、すべてのデプロイメントを改善する必要はありません。特定の GPU、並列処理設定、またはバッチサイズの範囲を対象にすることができます。明確な境界により、作業の調整が容易になり、検証する構成のセットが定義されます。
Kimi K3 の場合、初期ハードウェア検証は、TP8 および TP16 の GB300、最大 64 のデコードバッチ、さまざまなドラフト長の投機的デコード、および単一 NVLink ドメイン内のデコードコンテキスト並列構成をカバーします。
これらの境界は最適化が適用される場所を説明しますが、機能カバレッジを削減するものではありません。本番デプロイメントには、エンドポイントを有用にする基本的な機能が必要です。フルコンテキストサービス、温度サンプリング、構造化出力などの機能は、「最適化して除外」したり「オプトアウト」したりすることはできません。これらは第一級のサービス機能であり、カーネル形状や数値チェックとともに検証計画に属します。
原則 3: 正確性を検証する
カーネルテストとモデル評価は異なる質問に答えます。カーネルテストは、形状と入力値にわたる数値結果をチェックします。モデル評価は、完全な実行パスがモデルの回答を保持するかどうかをチェックします。両方が必要です。
テストは実際に変更されたコードを実行する必要があります。小バッチ最適化には小バッチでの評価が必要です。アテンションカーネルの変更には長いコンテキストのカバレッジが必要です。標準パスを暗黙的に使用するテストは、最適化に関する証拠をほとんど提供せず、一般的な罠として説明されています。
Kimi K3 の開発は、これが重要である理由を示しています。監査により、コンパイラがアクティベーション読み取りを同期待機の前に移動したことが発見されました。修正により、そのロードタイプは静的ウェイトに制限されました。このようなエラーは単純な機能チェックを生き残る可能性があり、対象を絞った調査が必要です。
MRCR や AIME25 などの長いコンテキストおよび長い生成評価も、モック入力や一貫した起動タイミングを使用したスタンドアロンカーネルテストでは再現できない、アテンションおよび通信カーネルの敏感な数値問題を明らかにしました。
原則 4: パフォーマンスを広くチェックする
見出しのスピードアップは出発点です。検証では、変更が影響するワークロード全体でどのように動作するかを示す必要があります。つまり、バッチサイズ、コンテキスト長、サポートされている並列処理構成をスイープする必要があります。MoE 測定には実際のプロンプトを含める必要があります。エキスパートの選択が読み取られるウェイトと再利用できる作業量に影響するためです。純粋なデコード結果は、プリフィルとデコードを混合するワークロードと区別する必要があります。
信頼性の高い検証には GPU 時間が必要です。関連する Bring-Your-Own-GPU 評価 RFC (vllm-project/vllm issue #61698) は、貢献者が自分のハードウェアで実行できる標準評価イメージを提案しています。再現可能な結果バンドルと GitHub PR にリンクされた提出物により、証拠がよりアクセスしやすくなります。異なる評価層と共有コミュニティコンピューティングにより、GPU 需要を管理し、プロジェクト資金による CI への圧力を軽減しながら、カバレッジを拡大できます。
パフォーマンス結果
DSpark によるデコードパフォーマンス
DSpark を使用すると、Kimi K3 は制御された AL=6 ベンチマークで 16 台の NVIDIA GB300 GPU で 842 トークン/秒 に達します。これは、同じ投機的デコード設定で 488 トークン/秒を測定した標準 vLLM よりも 1.73 倍高速 です。
セットアップ: TP16、7 つのドラフトトークンを持つ DSpark、合成受入長 6、DL=7、ISL=1、OSL=1k、バッチサイズ 1。TPS は総デコードスループットを測定し、プリフィルを除外します。
投機なしのデコードパフォーマンス
改善はテストされたバッチ範囲に及びます。TP16 はバッチ 16〜64 で 1.62〜1.71 倍のゲイン を維持し、TP8 は同じ範囲で 1.31〜1.43 倍 を提供します。
| バッチサイズ | TP8 トークン/秒 | TP8 スピードアップ | TP16 トークン/秒 | TP16 スピードアップ |
|---|---|---|---|---|
| 1 | 193.3 | 1.58 倍 | 256.4 | 1.88 倍 |
| 2 | 317.8 | 1.50 倍 | 436.7 | 1.92 倍 |
| 4 | 540.6 | 1.55 倍 | 702.5 | 1.81 倍 |
| 8 | 820.1 | 1.45 倍 | 1,157.9 | 1.79 倍 |
| 16 | 1,274.8 | 1.43 倍 | 1,758.3 | 1.71 倍 |
| 32 | 1,702.6 | 1.33 倍 | 2,550.9 | 1.64 倍 |
| 64 | 2,375.9 | 1.31 倍 | 3,606.1 | 1.62 倍 |
これらの数値は投機的デコードなしです。スループットは TP グループ全体に対するものです。
AgentX での改善
最適化された Kimi K3 実装は、AgentX の同時実行数 1 で 1.8 倍高い P90 インタラクティビティ を提供します: 403.2 対 223.7 トークン/秒/ユーザー。これは、16 台の GB300 GPU での最適化された実行と、8 台の GB300 GPU での以前の vLLM 結果を比較したものです。
AgentX は、長いコンテキスト、複数の会話ターン、ツールの一時停止、並列サブエージェント、KV キャッシュ再利用など、実際のコーディングエージェントセッションからのトラフィックパターンを再生します。プライベートコンテンツを合成テキストに置き換えながらこれらのワークロードパターンを維持し、現実的なエージェントトラフィック下でサービスシステムをテストします。
次のステップ
DeepSeek-V4.1-Flash
DeepSeek-V4.1-Flash (vllm-project/vllm PR #61868) は Kimi K3 の設計パターンに従います。同じオプトインの decode_backend: latency_specialized スイッチの背後で、この実装は最大 64 の同時リクエスト、プリフィルと混合バッチのシームレスなフォールバック、および投機的デコードをサポートします。単一のホットスポットを追いかける代わりに、そのエグゼキュータはカーネルとホストワークを削減します: すべてのサブレイヤー境界で all-reduce + mHC を融合し、低レイテンシ用に構築された CuTe DSL カーネルで大バッチ MoE と GEMM を置き換え、インデクサーの作業を重複排除してシャーディングし、標準 vLLM が CPU で熱心に実行するステップごとの作業を CUDA グラフとして再生します。DeepSeek AgentX の結果は 4 台の GB300 GPU で報告されています。
Muse Spark Open
Meta の Muse Spark Open リリースから 2 日以内に、vLLM の初期結果はすでに ユーザーあたり 950 トークン/秒 を超え、7 つのドラフトトークン (K=7) と受入長 6 (AL=6) を達成しました。さらなる最適化が進行中です。
GLM5.3
vLLM は、低同時実行シナリオでの GLM5.3 のデコードレイテンシを最適化する同様のメカニズムもリリースしています。最適化の詳細とパフォーマンス結果は別の投稿で提供されます。
各貢献について、vLLM は具体的なパフォーマンス向上、本番機能サポート、再現可能な正確性とパフォーマンスチェックに引き続き焦点を当てると述べています。
謝辞
vLLM は、この作業への貢献に対して Rohan Varma (NVIDIA) と Alex Steiner (NVIDIA)、モデルリファクタリングとフィードバックおよびサポートに対して Woosuk Kwon (Inferact)、Michael Goin、Lucas Wilkinson (Red Hat) に感謝しました。