vLLMのKimi K3.1 Proデイ0サポート:スパースアテンション、DCP、ゼロミスマッチRL

TL;DR

vLLMは、Moonshot AIの2.8Tパラメータ(アクティブ104B)のオープンウェイトのマルチモーダルMoEであるKimi K3.1 Proに対して、NVIDIA GPU、AMD GPU、Google TPU、Alibaba T-HEAD Zhenwu PPU上でデイ0サポートを提供しました。このリリースが重要なのは、K3.1 Proが、ハイブリッドKDA + MLAモデルの24のフルアテンション層すべてにスパースアテンション(DSA)を追加したため、vLLMはこのハイブリッドスパースアーキテクチャを効率的にサービス提供するために、KVキャッシュ、プリフィルシャーディング、デコード並列処理、低レイテンシデコードパスを再設計する必要があったからです。

Kimi K3.1 Proの変更点

Kimi K3.1 ProはKimi K3のアーキテクチャを維持しつつ、すべてのフルアテンション層にスパースアテンションインデクサーを追加しています。モデルは93層のハイブリッドで、69のKDA層と24のフルアテンション層、896のMoEエキスパート(16アクティブ)、MXFP4 MoEウェイト、1Mトークンのコンテキストウィンドウを維持しています。

Kimi K3 Kimi K3.1 Pro
層 93: 69 KDA + 24 フルアテンション 93: 69 KDA + 24 フルアテンション
フルアテンション層 高密度MLA MLA + DSAインデクサー(top-2048)、NoPE
インデクサー なし 32ヘッド × 128次元、FP8キー
MoE 896エキスパート、16アクティブ、MXFP4 変更なし
コンテキストウィンドウ 1M 1M
総パラメータ数 2.8T 2.8T

中心となる技術的変更はアテンションです。各スパース層は、MLAアテンションの隣に小さなインデクサーを保持します。インデクサーは、トークンごとに1つのFP8キー(スケールを含めて132バイト)を格納し、32の軽量ヘッドを使用して、各クエリに対して以前のすべてのトークンをスコアリングします。次に、上位2,048トークンを保持し、MLAアテンションはその2,048トークンに対してのみ正確に実行されます。vLLMは、DSAがDeepSeek-V3.2で導入されたことに言及し、詳細なインデクサーの説明については以前のDeepSeek-V3.2の投稿を参照しています。

スパースアテンションが長いコンテキストのデコード帯域幅を削減

DSAの利点は、メモリ帯域幅に制約される長いコンテキストのデコードで最も大きくなります。高密度MLAは、コンテキストトークンごとに層ごとに576バイトのFP8潜在変数を読み取ります。DSAは依然としてコンテキスト全体をスキャンしますが、そのスキャン中に132バイトのインデックスキーのみを読み取り、その後、2,048トークンに対してのみ完全な潜在変数を読み取ります。長いコンテキストでは、これは高密度MLAが層ごとに読み取るバイト数の約23%になります。

これがK3.1 Proのサービス提供の核となる主張です。モデルは1Mトークンのコンテキストを維持しながら、長いコンテキストでのデコードステップの層ごとのメモリトラフィックを、高密度MLAのトラフィックの約4分の1に削減します。

3種類の状態のためのパック型KVキャッシュ

vLLMは、各キャッシュブロックのMLA潜在変数、インデクサーキー、KDA状態をまとめてパックし、すべての層に同じサイズのページを与えるのではなく、各DSA層は、幅の異なる2つのトークンごとのキャッシュを格納します。MLA潜在変数はFP8で576バイト、NVFP4で352バイト、インデックスキーはFP8で132バイト、MXFP4で68バイトです。どちらの幅も他方の倍数ではなく、69のKDA層はそれぞれ、1層のトークンページよりもはるかに大きいリクエストごとのリカレント状態を追加します。従来の層ごとに同じサイズのキャッシュでは、小さなページを最大のものにパディングする必要があります。

vLLMは代わりに、キャッシュの1ブロックを、そのブロックのMLAページとインデックスキーページを24のDSA層すべてについて保持し、ブロックを共有するKDA状態も保持する、単一の連続した行にします。行内のページは異なるサイズにできます。唯一の制約は、行がKDA状態を保持することです。ブロックサイズは2のべき乗であるため、異なるTPサイズで実行されているプリフィルとデコードインスタンスのブロックは均等に分割され、KV転送のためにペアにできます。

アテンションカーネルは、期待する小さな固定サイズのページを引き続き認識します。vLLMは、各ステップの前にブロックテーブルとキャッシュスロットをカーネルページアドレスに変換するため、FlashInferとFlashMLAのスパースMLAカーネルとDeepGEMMのインデクサーカーネルは、パックされたキャッシュを変更なしで読み取ります。

このレイアウトにより、次の3つのことが可能になります。

  • 任意のスパースMLAカーネルとKV形式。 Blackwellでは、デフォルトはFP8潜在変数を使用するFlashInferのスパースMLAです。FlashInferのPrimsTSカーネルはNVFP4潜在変数も提供します。
  • 分離型サービス提供。 各ブロックはメモリの1つの連続した範囲であるため、プリフィル/デコード転送は、数十の小さな層ごとの読み取りではなく、1つの大きなRDMA読み取りで移動します。
  • 共有インフラストラクチャ。 同じパック型プランナーが、同じ混合ページサイズの問題を持つキャッシュを持つGLM-5.3-FlashとDeepSeek-V4にサービスを提供します。

トークンあたりのKVキャッシュフットプリント

テンソル並列処理では、MLA潜在変数とインデックスキーの両方がすべてのGPUにレプリケートされます。KDA状態はK3と同様にリクエストごとです。

GPUあたりのトークンあたりのKVキャッシュ(24フルアテンション層) バイト
Kimi K3、FP8 MLA潜在変数 13,824
Kimi K3.1 Pro、FP8潜在変数 + FP8インデックスキー 16,992
Kimi K3.1 Pro、NVFP4潜在変数 + MXFP4インデックスキー 10,080

潜在変数とインデックスキーは独立して量子化できます。NVFP4潜在変数だけでも、K3.1 ProはFP8のK3を下回ります。オプトインのMXFP4インデックスキーは、さらに11%多くのトークンを収容でき、長いコンテキストでインデクサーを10〜20%高速化します。両方を使用すると、8つのGB300 GPUはFP8の1.7倍のトークンを保持でき、DCP8はそれを7倍に拡大します。つまり、同時に約36の完全な1Mトークンコンテキストを処理でき、FP8のTP8では3つです。

トークンシャーディングされたプリフィルがファーストトークンまでの時間を改善

大きなプリフィルステップでは、vLLMはアテンションをヘッドではなくトークンでシャーディングします。TP8では、各GPUはK3.1 Proの96アテンションヘッドのうち12を所有します。これはスパースMLAプリフィルには不向きです。MLAの潜在KVはすべてのヘッドで共有されるため、各GPUは12ヘッドにサービスを提供するために完全なキャッシュを読み取り、スパースカーネルはそれらの12ヘッドを最小タイルにパディングします。インデクサーはヘッドで分割されていないため、すべてのGPUがすべてのトークンをスコアリングし、その2次コストが8回繰り返されます。

vLLMは代わりに、all-to-allを使用して、各GPUのクエリを「すべてのトークン、12ヘッド」から「トークンの1/8、すべての96ヘッド」に移動します。各GPUは、自身のトークンに対してインデクサーとスパースMLAを実行し、2番目のall-to-allが出力をヘッドレイアウトに戻します。KVキャッシュはTPですでにレプリケートされているため、他に変更はありません。MoE、シーケンス並列処理、キャッシュレイアウトはそのままです。因果アテンションのバランスを取るために、各リクエストは16の部分に分割され、GPU rは部分rと15 - rを取得します。これは、プリフィルコンテキスト並列処理で使用されるのと同じバランスの取れた分割です。両方のall-to-allは、ピアの対称メモリに直接書き込みます。

8× GB300、TP8、実ウェイト、一度に1リクエストの場合:

プロンプトトークン TTFT、ヘッドシャーディング TTFT、トークンシャーディング スピードアップ
27K 0.82秒 0.71秒 1.16×
108K 3.72秒 2.89秒 1.29×
340K 15.89秒 9.42秒 1.69×

エージェント的なターンで一般的な、長いキャッシュされたプレフィックスへの短い追加も恩恵を受けます。150Kトークンのキャッシュされたプレフィックスへの2Kトークンの追加では、インデクサーとスパースMLAカーネルは、FlashInferスパースMLAを使用した8× GB300で、56ミリ秒ではなくGPUあたり約15ミリ秒かかります。コストは交換バッファで、GPUあたり約0.8 GiBのKVキャッシュです。

スパースインデクサーを使用したデコードコンテキスト並列処理

テンソル並列処理はMLA潜在変数をすべてのGPUにレプリケートするため、GPUを追加してもKV容量は増加しません。デコードコンテキスト並列処理(DCP)は、代わりにキャッシュをトークンでシャーディングします。WのGPU rは、W番目ごとのトークンを格納します。高密度MLAの場合、各GPUは自身のシャードにアテンションし、部分出力はログサム指数(LSE)によって結合されます。DSAは1つのステップを追加します。各GPUは自身のトークンの1/Wのみをランク付けできるため、グローバルなtop-2048は、アテンションが開始される前にすべてのGPUでマージする必要があります。

約10倍少ないトラフィックで正確なtop-kマージ

一般的なマージは、すべてのGPUの完全なローカルtop-2048を、FP32スコアとトークンIDのペアとして、他のすべてのGPUに送信します。DCP8の8× GB300で、それぞれ5クエリトークンの64のAgentXのようなリクエストがある場合、671マイクロ秒のアテンションモジュールのうち258マイクロ秒かかります。

vLLMははるかに少ないデータを送信します。各GPUは、BF16スコアから直接、最良の1.5 × 2048 / W候補のみを選択し、各候補をスコアとともに32ビットにパックします。これにより、約10倍少ないバイト数になります。次に、すべてのGPUが、収集された候補から同じグローバルtop-2048を基数選択します。マージは正確なままです。GPUが送信したよりも多くのtop-2048を保持している場合、選択がそれを検出し、そのGPUはその行の完全なローカルtop-kを再送信します。

正確な1ラウンドのマージは、一般的なマージよりも1.27〜1.51倍高速です。バッチ32以上、128Kコンテキストでは、DCP8でのK3.1 Proのスパースアテンションは、K3の高密度MLAよりも高速です。

部分的なアテンション出力は、対称メモリ上の2つのCuTe DSLカーネルでLSEによって交換および結合されます。8× GB300の320のBF16形状で、既存のCUDA実装よりも最大3.7倍高速で、ビット単位で同一の出力を生成します。vLLMは、これがDCPでのすべてのMLAモデルに役立つと述べています。

マージなしのDCPプリフィル

通常のDCPプリフィルは、プロンプトトークンごとにマージとLSE結合を支払い、すべてのGPUでアテンション作業を繰り返すことになります。DCPグループがTPグループである場合、vLLMは代わりに、ステップのコンテキスト(MLA潜在変数とインデックスキーの両方)をすべてのGPUから対称メモリワークスペース(FP8でGPUあたり1Mトークンあたり約0.7 GB)に収集します。次に、上記のトークンシャーディングされたプリフィルをその上で実行します。マージや結合は残らず、出力はDCPなしのトークンシャーディングされたプリフィルとビット単位で一致します。

1つのDSA層、8× GB300でのDCP8、AgentX混合16Kトークンステップ DCPプリフィル 収集コンテキストプリフィル
FP8、FlashInfer trtllm-gen 24.8ミリ秒 5.8ミリ秒(4.3×)
FP8、FlashInfer PrimsTS 24.8ミリ秒 5.7ミリ秒(4.3×)
NVFP4、FlashInfer PrimsTS 25.7ミリ秒 6.2ミリ秒(4.1×)

エンドツーエンドでは、単一リクエストのプリフィルはDCP8でTP8と同じくらい高速です。27K、108K、340Kトークンでそれぞれ1.10、4.40、14.61秒であり、8× GB300のFP8 KVでは1.10、4.39、14.50秒です。

DCPを使用するタイミング

DCPは、ステップごとのレイテンシをKV容量とトレードオフします。DCP8は、GPUあたりのMLAおよびインデックスキーキャッシュを8分の1に削減し、KDA状態はすでにヘッドで分割されています。コンテキストが収まる場合、通常のTPは低同時実行でより高速にデコードします。DCPは、非常に長いコンテキストや多数の長期エージェントセッションなど、KVが制限となる場合に効果を発揮します。8× GB300では、DCP8はFP8 KVの21.8Mトークンを保持し、TP8は3.04Mトークンです。NVFP4潜在変数では、32.1M対4.49Mです。同時実行1では、TPOTはTP8で9.35ミリ秒、DCP8で11.48ミリ秒(32Kコンテキスト)、128Kでは9.45対11.58ミリ秒です。融合デコードパスにより、DCP8は5.96ミリ秒と6.04ミリ秒になります。

TP8の8つのGPUは、完全な1Mトークンリクエストも処理できます。8× GB300のFP8 KVでは、キャッシュは3.04Mトークンを保持し、つまり2.9の同時1Mトークンリクエストを処理できます。

DSA層のための融合低レイテンシデコード

Kimi K3では、vLLMは、小さなバッチ向けに調整されたCuTe DSLカーネルでほぼすべての演算子を置き換えるオプトインデコードパスを提供しました。通信を計算に融合し、MoEスケジュールを再設計し、ウェイトをL2にプリフェッチします。K3.1 Proでは、このパスはDSA層もカバーしています。DSA層は次のように実行されます。

  • クエリ、KV潜在変数、出力ゲート、インデクサーのキーとヘッドウェイトのための1つのパックされた入力射影。
  • 融合インデクサーフロント:クエリノルム、インデクサーのクエリ射影、ヘッドウェイトを折り込んだFP8またはMXFP4量子化、新しいトークンのインデックスキーを1つのカーネルで処理。
  • MLAクエリ(アップ射影と吸収されたキー射影を含み、新しいトークンの潜在変数をKVキャッシュに書き込む)を、インデクサーの前に実行するのではなく、インデクサーとオーバーラップさせて実行。
  • 融合インデクサー:キーバランスの取れたスコアラーと正確なtop-k選択を2つのカーネルで実行し、アテンションが読み取る物理KV行を生成。
  • スパースアテンション:FlashInferスパースMLA(trtllm-genまたはPrimsTS、バッチサイズに応じて選択)またはFlashMLAスパース。DCPでバッチ1〜8の場合、形状に合わせて構築されたCuTe DSLカーネルを使用し、バッチ1で層あたり11.9マイクロ秒(21.4マイクロ秒に対して)。
  • 値射影とシグモイドゲートを1つのカーネルで実行し、K3の出力射影と融合オールリデュースを続けて実行。

DCPでは、クエリはすべてのGPUに公開され、各GPUのtop-k候補が交換されて正確にマージされ、部分出力はLSEで結合されます。これらはすべて、カーネル内のNVLinkメールボックスを介して行われ、パス上にNCCLコレクティブはありません。

融合デコードパスは、8× GB300のTP8と1Kトークン入力で、バッチ1で1.63倍、バッチ64でも1.18倍高速です。DSpark投機的デコードと7つのドラフトトークンを使用すると、融合パスは同時実行1でのトークンあたりの時間を2.86ミリ秒から2.14ミリ秒に短縮し、GSM8Kの精度は変わりません。DCP8では、融合パスは同時実行1でのTPOTを32Kコンテキストで11.48ミリ秒から5.96ミリ秒に、128Kで11.58ミリ秒から6.04ミリ秒に短縮します。

DSparkによる投機的デコード

vLLMはK3.1 Pro用のDSparkスペキュレーターをトレーニングし、プレビューチェックポイントであるInferact/Kimi-K3.1-Pro-DSpark-previewをリリースしました。温度1.0、top-p 0.95の現実的なサンプリングで、プレビューは7つのドラフトトークンでGSM8Kで平均受理長4.597を達成します。8× GB300のTP8でバッチサイズ1の場合、これは3.0倍のデコード高速化です。スペキュレーターは、長いコンテキストのBEAMベンチマークで測定された、K3.1 Proの最大1Mトークンの完全なコンテキストウィンドウ全体で精度を維持します。

vLLMはこのチェックポイントをプレビューのみと説明し、より正確で効率的なスペキュレーターを近日中にリリースするためにトレーニングしていると述べています。

NVIDIA GB300およびRubinでのエージェントサービス提供

vLLMは、デイ0イメージをNVIDIA GB300およびNVIDIA RubinでSemiAnalysis AgentXベンチマークを使用してベンチマークしました。

NVIDIA GB300

フロンティアのレイテンシ重視の端では、Kimi K3.1 Proはユーザーあたり毎秒375.9トークンを維持します。フロンティアに沿ってさらに進むと、各ユーザーを毎秒100トークン以上に保ちながら、GPUあたり毎秒19.2Kトークンを処理します。スループット端では、GPUあたり毎秒37.8Kトークンに達し、最もインタラクティブなポイントの約26倍ですが、各ユーザーは依然として毎秒約54トークンを確認できます。GB300スイープでは、4つの投機的トークンと合成受理長3.36トークンを使用しました。

NVIDIA Rubin

NVIDIA Rubinでは、Kimi K3.1 Proはフロンティアのレイテンシ重視の端でユーザーあたり毎秒815.6トークンを維持します。フロンティアに沿ってさらに進むと、各ユーザーを毎秒100トークン以上に保ちながら、GPUあたり毎秒46.0Kトークンを処理します。スループット端では、GPUあたり毎秒68.7Kトークンに達し、最もインタラクティブなポイントの約51倍ですが、各ユーザーは依然として毎秒約75トークンを確認できます。同時実行64までのポイントは、合成受理長3.84トークンの7つの投機的トークンを使用します。より高い同時実行では、受理長3.0の3つの投機的トークンを使用します。

デイ0でのゼロミスマッチRL

vLLMのデイ0サポートは、初めてゼロミスマッチのRLを提供します。verlを使用したLoRA RL、ロールアウトエンジンとしてのvLLM、トレーナーとしてのMegatronで、トレーナーのログ確率がロールアウトとビット単位で一致します。重要度サンプリングやルーティングリプレイでギャップを制限する代わりに、vLLMはトレーナーのフォワードパスがロールアウトをビット単位で再現するようにすることで、カーネルレベルでギャップを排除します。

  • モデル全体でのLoRA RL。 ランク32のLoRAは、すべてのアテンション射影とすべての896エキスパートをカバーし、ベースモデルとアダプターの両方がvLLMとMegatron間でビット単位で同一です。
  • バッチ不変カーネル、ネイティブに。 VLLM_BATCH_INVARIANT=1を使用すると、トークンのログ確率はバッチに依存せず、DeepSelect top-kなど、vLLMの高性能カーネルの多くは設計上バッチ不変です。

モデル全体で100ステップ以上にわたって、ロールアウト/トレーナーKLはすべてのステップで正確に0.0であり、AIME mean@8は0.451から0.664に上昇します。vLLMは、バッチ不変パスのパフォーマンス最適化は進行中であり、ソースコードはInferact/zero-mismatch-rlにあると述べています。

AMD Instinct MI355Xサポート

AMD Instinct MI355Xでは、vLLMはMXFP4インデックス処理とスパースアテンションおよびデコードコンテキスト並列処理を組み合わせています。これらの最適化により、長いコンテキストのメモリ要件が削減され、プリフィルとデコードの効率が向上します。

ネイティブMXFP4インデクサー

インデクサーは、アテンションのために2,048トークンを選択する前にコンテキストをスキャンします。MXFP4は、各インデックスキーをFP8の132バイトからスケールを含めて68バイトに削減し、インデックスキャッシュストレージを48.5%削減します。スコアリングは、FP32累積とヘッド間の加重削減を備えたネイティブFP4行列命令を使用します。キーは、投機的検証中にクエリ行間で再利用され、繰り返し読み取りを減らします。AgentXスイープでは、MXFP4インデクサー構成は、同時実行48で総ワークロードスループットを5.3%向上させました。

コンテキスト並列スパースアテンションプリフィル

プリフィル中、GPUはインデックスキーを収集し、クエリ行を分割します。各GPUは、割り当てられた行を完全なコンテキストに対してスコアリングし、グローバルtop-kを直接選択するため、ランクごとの候補のマージを回避します。メインアテンションは、通信を削減しワークスペース予算に収まる場合にKVを個別に収集し、ローカルヘッドのすべてのクエリ行を処理できます。

入力/出力トークン 総スループット(トークン/秒) 出力スループット(トークン/秒) 中央値TTFT(秒) 中央値TPOT(ミリ秒)
8K / 1K 5,357.8 → 6,015.4(+12%) 595.3 → 668.4 2.462 → 1.875 24.571 → 22.063
128K / 100 6,950 → 17,678(+154%) 5.30 → 13.48 268.7 → 105.5 331.5 → 129.6

セットアップ:MI355X、TP8/DCP8、同時実行16、64プロンプト、プレフィックスキャッシュ無効。

FlyDSLスパースMLAは、選択されたトークンをFP8 KVキャッシュから直接読み取ります。各ワークグループは、ローカルアテンションヘッド全体で1つのクエリ行を処理し、プリフィルステップの多数のクエリ行にわたって並列処理を提供します。小さなバッチはTritonパスを維持し、各行内で追加の並列処理を公開します。

実装 ベンチマーク時間(秒) 中央値TTFT(秒) 出力スループット(トークン/秒)
Triton 432.8 96.2 14.79
FlyDSL 404.4 89.8 15.83
改善 6.6%低減 6.6%低減 7.0%向上

セットアップ:MI355X、TP8/DCP8、同時実行16、64プロンプト、ISL/OSL 128K/100。

融合低レイテンシデコード

小さな投機的検証バッチでは、DSAインデックス処理とスパースアテンションが別々に実行され、その後、アテンション出力射影、AttnRes、テンソル並列オールリデュースを含む完全な潜在MoE計算を単一の起動に融合する永続モノカーネルが続きます。これにより、カーネル起動オーバーヘッドが削減され、MoEウェイトのロードが先行する削減および残差処理とオーバーラップします。特殊なパスは、7つのドラフトトークンを持つ単一リクエストを含む、ステップごとに最大8つの検証トークンをサポートします。平均TPOTは、FP8インデクサー、8K/1K、同時実行1、7つのドラフトトークン、合成受理長3.84のTP8/DCP1で、6.02ミリ秒から4.40ミリ秒に改善され、27.0%削減されます。

GPU間通信

デコード中、各GPUはトークンIDをローカルに保ちながら候補スコアを交換します。共有しきい値と一貫したタイ処理により、選択されるローカル候補が決定されます。2,048候補の場合、スコアのみを交換すると、候補ペイロードがGPUあたりクエリあたり16 KiBから8 KiBに半分になります。AITERは、HIP IPCを介してピアバッファを読み取り、必要な出力レイアウトに直接書き込む専用GPUカーネルでDCPオールギャザーを高速化し、データ収集とレイアウト変換を組み合わせます。8つのMI355X GPUでのスタンドアロンマイクロベンチマークでは、AITERパスは、テストされたクエリおよび候補テンソル形状で、RCCLと比較してオールギャザーレイテンシを53〜61%削減しました。

Google TPUサポート

vLLMは、vllm-torchtpuを通じてTPUでKimi K3.1 Proをデイ0でサポートし、TPU v7xの16チップスライスで検証されています。レイアウトはMoE用のTP32/EP32で、長いコンテキストのためにMLA潜在変数とインデックスキーをトークンでシャーディングするDCPを備えています。DSAの作業はTPUの2種類の計算ユニットに分割されます。TensorCoreがFP8インデックスキーをスコアリングし、SparseCoreが上位2,048トークンを選択して、その潜在変数をTensorCoreのMXUがアテンションを実行するための高密度バッファに収集します。GPUと同様に、DSAとKDA状態は1つのブロックインデックス付きKVプールを共有し、ハイブリッドキャッシュ全体でプレフィックスキャッシュを機能させ続けます。

Alibaba T-HEAD Zhenwu PPUサポート

Alibaba T-HEADチームは、vllm-sailプラグインを通じて、Alibaba T-HEAD Zhenwu PPU(M890および810Eの両方)でKimi K3.1 Proを実行するためのサポートを追加しました。検証済みの構成は次のとおりです。

  • 810E:2ノード、TP32 / EP32、W4A16 MoEバックエンドを使用。
  • M890:4ノード、TP8 / DP4 / EP32、DeepEPを使用。

vLLMは、Zhenwu PPUが主流のオープンソース大規模モデルに対してデイ0サポートを提供するのはこれが初めてであり、継続的なサポートが近日中に提供されると述べています。

クイックスタート

NVIDIA GPUでは、vLLMはテンソル並列処理8、fastsafetensorsロード、Kimi K3ツール呼び出しおよび推論パーサー、5つの投機的トークンを持つDSpark投機的構成を使用するvllm serveコマンドを提供します。

vllm serve moonshotai/Kimi-K3.1-Pro \
  --tensor-parallel-size 8 \
  --trust-remote-code \
  --load-format fastsafetensors \
  --enable-auto-tool-choice \
  --tool-call-parser kimi_k3 \
  --reasoning-parser kimi_k3 \
  --speculative-config '{"model":"Inferact/Kimi-K3.1-Pro-DSpark-preview","num_speculative_tokens":5,"method":"dspark","draft_sample_method":"probabilistic"}'

AMD GPUでは、コマンドはAITERと高速GPU safetensorsロードを有効にし、テンソル並列処理8を設定し、シーケンスを16に制限し、GPUメモリ使用率を0.85に設定し、同じDSpark投機的構成を使用します。

VLLM_ROCM_USE_AITER=1 SAFETENSORS_FAST_GPU=1 \
vllm serve moonshotai/Kimi-K3.1-Pro \
      --served-model-name kimi-k31-pro \
      --tensor-parallel-size 8 \
      --trust-remote-code \
      --load-format fastsafetensors \
      --max-num-seqs 16 \
      --gpu-memory-utilization 0.85 \
      --speculative-config '{"model":"Inferact/Kimi-K3.1-Pro-DSpark-preview","num_speculative_tokens":5,"method":"dspark","draft_sample_method":"probabilistic"}'

このリリースがハイブリッドスパースモデルのサービス提供について示すこと

Kimi K3.1 Proは、トークンごとまたはリクエストごとの3種類の異なる状態(MLA潜在変数、DSAインデックスキー、KDAリカレント状態)を持つハイブリッドモデルです。vLLMのデイ0サポートは、このようなモデルを効率的にサービス提供するには、単一のアテンションカーネルではなく、KVキャッシュレイアウト、プリフィルシャーディング、デコード並列処理、低レイテンシカーネル融合にわたる協調的な変更が必要であることを示しています。

パック型ブロックキャッシュは、それを可能にする抽象化です。アテンションカーネルが固定サイズのページを維持できるようにしながら、サービス提供スタックは、RDMA転送とプレフィックスキャッシュのために混合幅の状態を連続して格納します。トークンシャーディングされたプリフィルは、テンソル並列処理でのMLAのヘッドレプリケーション問題に対処し、正確なtop-kマージを持つDCPは、長いコンテキストでのTPのKV容量制限に対処します。融合デコードパスとDSpark投機的デコードは、エージェントワークロードが依存する小さなバッチレイテンシをターゲットにしています。

ゼロミスマッチRLの結果は推論パフォーマンスとは別ですが、トレーニングインフラストラクチャにとって重要です。vLLMは、ベースモデルと、すべてのアテンション射影とすべての896エキスパートをカバーするランク32のLoRAアダプターの両方について、vLLMロールアウトとMegatronトレーニングの間でビット単位で同一のログ確率を報告し、100ステップにわたってKLが正確に0.0です。これにより、重要度サンプリングやルーティングリプレイで制限するのではなく、RL不安定性の原因を排除します。

プラットフォームカバレッジ

デイ0サポートは、NVIDIA GPU、AMD GPU、Google TPU、Alibaba T-HEAD Zhenwu PPUの4つのハードウェアファミリに及びます。NVIDIAパスには、SemiAnalysis AgentXでのGB300およびRubinベンチマークが含まれます。AMDパスには、MXFP4インデックス処理、コンテキスト並列スパースプリフィル、融合低レイテンシデコード、AITERオールギャザーのMI355X最適化が含まれます。TPUパスは、TP32/EP32とDCPを備えた16チップTPU v7xスライスでvllm-torchtpuを通じて実行されます。T-HEADパスは、810EおよびM890構成でvllm-sailプラグインを通じて実行されます。

vLLMは、Kimi K3.1 Proをオープンソース化したMoonshotに感謝し、デイ0サポートを提供したInferact、AMD、Alibaba T-HEAD、および支援してくれたNVIDIAに感謝します。

Sources