AMD Instinct MI355X上のKimi K3: AgentXにおけるvLLMサービングがNVIDIA HGX B300のGPUあたりスループットに匹敵

AgentXベンチマークの50トークン/秒/ユーザーP90インタラクティビティにおいて、vLLMでKimi K3をサービスする8GPU構成のAMD Instinct MI355Xシステムは、NVIDIA HGX B300よりもGPUあたりの総トークン数が24.4%多く、ベンチマークのモデル化された所有コスト前提では1ドルあたりのトークン数が87.4%多いことを実現しています。この結果が重要なのは、ソフトウェアの成果だからです。MI355Xの公開されている高密度低精度マトリックスピークはB300よりも低いため、この優位性は、キャッシュ分割、MoEコレクティブ再構築、カーネル融合、AITERアセンブリアテンションパス、KDA並列化、ドラフト深度チューニングといったサービングスタックのエンジニアリングによるものであり、生のハードウェアスループットによるものではありません。

AgentXでの主要数値

50トークン/秒/ユーザーP90インタラクティビティにおいて、AMD Instinct MI355Xは約7,567トークン/秒/GPUをサービスし、NVIDIA HGX B300の6,083と比較して24.4%のスループット優位性を示します。SemiAnalysisのモデル化TCO(GPU時間あたり$1.50対$2.26)を使用すると、TCO 1ドルあたりの総トークン数は1,816万対969万となり、1ドルあたりのトークン数が87.4%多くなります。50トークンのインタラクティビティ目標では、この比較によりモデル化されたトークンあたりコストが46.6%低くなります。

このベンチマークのインタラクティビティは、P90応答トークン間レイテンシの逆数であり、最初のトークンまでの時間は含まれません。総トークン数には、再利用された入力と生成された出力の両方が含まれます。

公開仕様(GPUあたり) AMD Instinct MI355X NVIDIA B300 GPU
高密度低精度マトリックスピーク 10.1 PFLOP/s MXFP4 13.5 PFLOP/s FP4
GPUメモリ 288 GB HBM3E 288 GB HBM3E
ピークメモリ帯域幅 8 TB/s 8 TB/s

2つのアクセラレータは、公開されているメモリ容量と帯域幅が同一ですが、B300の方が公開されている高密度低精度マトリックスピークが高くなっています。このサービング性能の達成は、AMDとvLLMコミュニティ全体のコラボレーターによる持続的なソフトウェアエンジニアリングを反映しています。

AgentXが測定するもの

公開のAgentXベンチマークは、実際のコーディングエージェントセッショントレースを合成トークンでリプレイし、リクエスト長、タイミング、プレフィックス共有パターンを保持します。このワークロードは、無制限キャッシュで理論上のプレフィックス再利用率96%以上を提供します。投機的実行は、固定された合成受け入れ基準を使用します。

エージェントターンはプレフィックスのほとんどを再利用するため、このベンチマークは、会話状態を常駐させ、再計算を避けるシステムに報酬を与えます。公開曲線は完全なサービングスタックを比較します。各最適化の横に報告された測定値は、それぞれの記載された構成の下での個別テストからのものです。

Kimi K3の内部: MLA、KDA、MoE

Kimi K3は、2種類のアテンションとスパースMoEを組み合わせています。戻ってくるエージェントターンはプレフィックスのほとんどを再利用できますが、短い継続部分の処理でも3種類のデータに触れます。

  • 24のMLAレイヤーは、会話の進行に伴って成長する圧縮KV履歴を保持します。新しいクエリは依然としてその履歴を参照します。
  • 69のKDAレイヤーは、固定次元の状態行列を前方に運び、新しいトークンが到着するたびに更新します。
  • MoEレイヤーは、各トークンをルーターを通じてエキスパートネットワークを選択します。Kimi K3はトークンごとに896のエキスパートから16を選択し、エキスパートの重みはリクエスト間で共有されます。

MLAは会話が成長するにつれてより多くの履歴を保持します。KDAは固定次元の状態を更新するため、以前のチェックポイントを保持することは、その状態の追加コピーを保持することを意味します。最適化作業は3つのスコープに分けられます:ターン間(キャッシュ再利用とオフロード)、モデルステップ内(エキスパート射影、カーネルハンドオフ、MLA/KDAスケジューリング)、そしてドラフト深度の選択です。

1. キャッシュ: 会話状態を再利用可能に保つ

デコードコンテキスト並列処理(DCP)はGPUメモリ内の重複履歴を削減し、CPUオフロードは再利用可能な状態をGPU外に保持します。

TP8では、各GPUはKimi K3の96のMLAクエリヘッドのうち12を処理しますが、圧縮されたKV履歴全体を保存します。テンソル並列処理はヘッドを分割しますが、キャッシュは分割しません。**8方向デコードコンテキスト並列処理(DCP8)**では、各GPUはキャッシュされた位置の8分の1を保存します。デコード中、GPUはクエリを交換し、ローカル履歴に対して計算し、ソフトマックス重みを使用して部分的な答えをマージします。各部分的な答えは正規化情報を運び、マージがシャードを正しく重み付けできるようにします。PR #51705は、ROCm AITER DCPパスをマルチトークン検証に拡張します。

142,000トークンの会話の場合、推定MLAペイロードは24レイヤー×142,000位置×576 FP8バイト≈GPUあたり1.96 GBです。DCP8はこれを割り当てオーバーヘッドの前に約245 MBに削減します。各GPUは履歴の8分の1に対して8倍のヘッドを処理し、主要なアテンション演算はほぼ変わらずに残ります。この追加の分割はMLA履歴に適用されます。KDA状態は、ヘッド間の既存のパーティションを維持します。

CPUオフロードはHBMに収まらなくなった状態を保持します

状態がGPUメモリに収まらなくなった場合、AMDレシピのネイティブvLLM CPUオフロードは、MLA履歴とKDAチェックポイントをホストDRAMに保持します。

キャッシュヒットには、一致する再開ポイントが必要です。MLA履歴がトークン3,072に達しても、最新の互換性のあるKDAチェックポイントが2,048にある場合、リクエストは2,048で再開し、残りのプレフィックスを再計算する必要があります。アテンションブロックテールをKDAチェックポイントと一緒に保存すると、再開位置が揃います(PR #54736)。その状態を復元することは、プレフィックスの再計算を避けるためにGPUへの転送をトレードオフします。

2. MoE: GPU間の重複作業を排除

1つのGPUでは、ルーティングされたエキスパート出力射影は、コンシューマーが期待する完全なベクトルを返します。元のROCmパスでは、以前の削減と正規化により、8つのGPUすべてが同じ潜在入力を持ち、それぞれが完全な重みを保持していました。つまり、同じ射影が8回実行されました。

出力座標はGPU間で分割できますが、すべてのランクは依然として完全な結果を必要とします。Kimi K3はすでに共有エキスパート部分のためのオールリデュースを持っています。各GPUは、所有するルーティングされた出力のみを共有エキスパート部分に追加し、既存のオールリデュースが新しいコレクティブなしで完全な結果を組み立てます(PR #51253)。ROCm実装は、対応するCUDA戦略に基づいています。各出力座標は、ルーティングされた寄与を正確に1回受け取り、完全な射影重みは複製されたままです。

PRのテストでは、8つのAMD Instinct MI355X GPU(TP8)で、同時実行数8〜128、ランダムな8K入力/1K出力リクエスト、プレフィックスキャッシュ無効で、出力スループットは3.05〜4.91%向上しました。

GPUメモリを解放して、より多くのキャッシュ状態を保持

各ランクが使用する行のみをロードすることで、TP8の92のMoEレイヤー全体でGPUあたり約3.85 GiBが解放されます(PR #59591)。融合パスの計算は同じままです。追加のスペースはMLA履歴とKDAチェックポイントを保持できます。

この変更は、主要な公開結果に使用されたイメージの後に導入されました。その別のAgentX実験では、同時実行数32と48では1%未満のスループット向上でしたが、70では28.9%向上し、報告されたプレフィックスキャッシュヒット率は74.8%から80.1%に上昇しました。高同時実行での大きな利得は、キャッシュプレッシャーと一致しています。測定は8つのAMD Instinct MI355X GPU、TP8/DCP8、FP8 KV、30分間の実行、CPUオフロードや投機的デコードなしで行われました。繰り返し実行の不確実性は報告されていません。

3. メモリトラフィック: 操作を融合し、不要なコピーを削除

各GPU内で、次に注目すべきはカーネル間で渡されるデータです。

Kimi K3のシングルトークンKDAパスは、3つの計算カーネルと2つの出力コピーを実行しました。畳み込みは、Q/K正規化、減衰、ベータゲート、リカレント更新をすでに融合したカーネルの前に実行され、出力正規化はその後実行されました。融合HIPパスはこれらの操作を1回の起動で実行し、宛先バッファに直接書き込みます(PR #50654)。中間ハンドオフはオンチップに留まりますが、畳み込み履歴とリカレント状態は次のトークンのためにGPUメモリに送られます。

AttnResは、以前のレイヤーブロック出力を実行中の残差と混合します。この混合をRMSNormと組み合わせることで、一時的な混合物をレジスタに保持します(PR #50593)。後で保存される実行中の残差はBF16丸めを保持します。混合物は正規化を通じてFP32のままです。AttnRes融合により、8つのAMD Instinct MI355X GPU(TP8)で、100K入力/1K出力リクエスト、プレフィックスキャッシュ無効で、出力スループットは14.2%向上しました。ロングプリフィルレイテンシは実質的に変わりませんでした。

使用するカーネルにのみ入力をコピー

ROCm行列乗算ディスパッチャーは、カーネルを選択する前に入力をパックしていました。フォールバックを選択した場合、そのカーネルは元の入力を読み取り、コピーは未使用のままでした。パッキングを特殊化されたブランチに移動することで、これらのコピーが削除されました(PR #58566)。

プロファイルにより、無駄が明らかになりました:デコードステップごとに69の不要なコピー、KDAレイヤーごとに1つ。それぞれが1,792バイトしか移動しませんでしたが、約4.8マイクロ秒かかりました。起動オーバーヘッドが支配的だったため、これらのバイトを高速に移動してもほとんど効果がありませんでした。PRのTP8、同時実行1、6ドラフト、384Kコンテキストプロファイルでは、削除されたコピーはステップごとに0.339 msのコピーカーネル時間を占めていました。

4. MLA: AMD AITERアセンブリによる高速ロングコンテキストデコード

Kimi K3はTP8でGPUあたり12のクエリヘッドを持ちます。16にパディングすることで、AMD AITERのアセンブリデコードパスが利用可能になりました。報告されたロングコンテキストテストでは、出力トークンあたりの時間は43.8から23.6ミリ秒に短縮されました。

比較はTP8、DCPなし、同時実行1で行われました。PR #50578のアダプターは4つのクエリヘッドをコピーし、16ヘッドレイアウトにパックして、その出力を破棄します。DCP8は異なる96ヘッド形状を提示します。これは、各クエリヘッドがキャッシュされた位置に対して独立してアテンションするため、モデルの操作を保持します。別のヘッドを追加しても、元のヘッドのソフトマックスは変わりません。このデコードパスではクエリのみがパディングされ、保存された圧縮KV履歴は同じサイズのままです。

入力/出力トークン 元のGluonパス パディングされたアセンブリパス
1K / 1K 22.9 ms/トークン 22.8 ms/トークン
32K / 1K 43.8 ms/トークン 23.6 ms/トークン

これらは、8つのAMD Instinct MI355X GPU、TP8、同時実行1、MXFP4モデル重み、BF16アテンション、通常のシングルトークンデコードでのPR報告の平均出力トークンあたりの時間です。短い入力のケースはほとんど変わりません。長い入力のケースでは、出力トークンあたりの時間がほぼ半分になります。

メカニズムはスケジューリングです。調査された歴史的なコードでは、保守的なスケジューリングヒントによりGluonパスがリクエストごとに1つのワークグループに制限され、そのワークグループが会話全体を処理していました。パディングによりAMD AITERのアセンブリパスが利用可能になり、履歴を並列部分に分割して結果を組み合わせるための異なるポリシーが提供されました。調査されたアセンブリポリシーは16のパーティションを選択しました。履歴パーティションは1つのGPU内のワークグループに割り当てられた作業です。KVストレージは変わりません。

5. KDA: プリフィルと投機的検証の並列化

MLAは保持された履歴をワークグループ間で分割できます。KDAはより難しいスケジューリング問題を提示します:各トークンは前の状態に依存します。PR #54038の融合プリフィルは、各トークングループの結合変換を、その入力状態が既知になる前に準備します。

KDAは行列変換でこれを行います。グループサマリーの順序付きスキャンが各グループの開始状態を決定し、グループはその後トークンを並列にリプレイします。トークン順序はそのままですが、各グループは前のグループのリプレイを待つ必要がなくなります。グループ化は2回目のウォークを追加し、融合チャンクパスを2回の起動から4回に拡張します。これは、少数のシーケンスやヘッドが並列作業をほとんど公開しない場合に役立ちます。それ以外の場合は、追加のパスが節約以上にコストがかかる可能性があります。各グループ内では、リカレントウォークはチャンク間で状態を行列アキュムレータレジスタに保持し、グローバルメモリへの書き込みと再ロードを回避します。

融合KDAは、別の3バックエンド比較で、4Kおよび16K入力長でTritonよりも5.3〜5.6%多い入力スループットを提供しました:8つのAMD Instinct MI355X GPU、TP8、同時実行32、プレフィックスキャッシュ無効、1トークン出力。これは完全なバックエンドを測定します。

投機的検証のための小さなタイル

短い検証バッチには異なる分割が必要です。TP8では、各GPUは128の値列を持つ12のKDAヘッドを持ちます。各ヘッドを4つの32列タイルに分割すると、256の計算ユニットで48のワークグループしか起動されません。PR #58045は4列タイルに切り替え、カウントを384に引き上げます。各タイルは異なる値列を所有し、同じトークンシーケンスを順にウォークします。

1つのAMD Instinct MI355X GPU(TP8次元、1シーケンス、4検証位置)での別のグラフリプレイカーネルテストでは、時間はレイヤーあたり**16.20から5.61マイクロ秒(2.89倍)**に短縮されました。これには、より狭いタイル、シングルウェーブワークグループ、均一長の特殊化、ホイストされたトークン不変作業が含まれます。ラグドバッチは動的パスを保持します。テストは出力とリカレント状態を比較します。

6. 投機的デコード: ドラフトするトークン数を選択

検証が速くなると、チェックする価値のあるドラフトトークン数が変わる可能性があります。有用なレートは、コミットされた出力を完全なドラフト・検証・コミットサイクル時間で割ったものです。

AgentXの固定受け入れ基準では、3つのドラフトはサイクルあたり3.00の期待コミットトークンを提供し、4つのドラフトは回復またはボーナストークンを含めて3.36を提供します。つまり、サイクルあたり12%多いコミット出力です。したがって、4つのドラフトは、完全なサイクルが3つのドラフトよりも12%未満長い場合にのみスループットを向上させます。その12%のしきい値はコスト予算です。完全なサイクルの成長は投稿で測定されていません。

公開のAMDレシピは、DCP1で3〜7のドラフトを使用し、より大きなDCP8同時実行設定では投機を無効にします。

技術ノート: KDAプロローグの精度とオンチップメモリ

KDAプロローグは64×64チャンク内依存行列の逆行列を計算し、後の行列積はこの変換を使用してトークン依存性を考慮します。逆行列の中間値は成長してキャンセルする可能性があるため、早期に丸めると最終結果が必要とする情報が失われます。実装はAMDのネイティブFP32行列命令を使用し、敏感な逆行列のためにFP32入力と累積を保持し、完了した逆行列のみを後続の積のためにBF16に丸めます。

FP32累積は、開始前に失われた情報を回復できません。256から256.5を引くと、元のFP32入力が保持されている場合は0.5になりますが、入力が最初にBF16に丸められている場合は0になります(256.5はその大きさで256に丸められます)。これらのおもちゃの値は、逆行列がネイティブの16×16×4 FP32行列命令を使用し、入力と累積の両方の精度を保持する理由を示しています。

プロローグにはローカルメモリ予算もあります。AMDはオンチップ共有スクラッチをLDSと呼びます。QとKタイルは異なる時間に必要とされるためストレージを再利用し、三角入力は占有する10の16×16ブロックのみを保存します。ワークグループは79.5 KiBを要求し、80 KiBに丸められます。これはAMD CDNA 4計算ユニットの160 KiB LDS予算の半分です。割り当ては1,280バイト単位に丸められるため、2つの80 KiB割り当ては収まりますが、2つの81.25 KiB割り当ては収まりません。これは容量計算です。レジスタや他のリソースが追加の制限を課す可能性があります。

MoE側では、別の技術ノートが精度をカバーしています:PR #53940のA4W4統合は、アクティベーションを8ビットから4ビットに変更し、重みは4ビットのままです。総トラフィックにはBF16中間出力とエキスパートタイルパディングも含まれます。これは、大規模なサーバーバッチでもエキスパートごとに数行しか残らない可能性があるためです。

ベンチマーク構成と注意事項

比較ターゲットでのサービング構成:

P90インタラクティビティ(トークン/秒/ユーザー) AMD Instinct MI355X NVIDIA HGX B300
30および50 SimpleCPUオフロード; DCP8 / ドラフトなしからDCP1 / 3ドラフト Mooncakeオフロード; DCP8
150 GPUのみ; DCP1; 5〜7ドラフト Mooncakeオフロード; DCP8

両システムとも8 GPU、TP8、FP4重み、FP8 KVを使用します。値は、一致するインタラクティビティで公開された各オフロードまたはGPUのみのシリーズ内で補間されます。同時実行、DCP、ドラフト深度は、ブラケット実行間で異なる場合があります。

提出物はvLLM 36768d1bf(AMD)および5894fdf98(NVIDIA)を使用します。構成されたホストキャッシュ予算は、それぞれGPUあたり224.875および281 GBです。リプレイはcc-traces-weka-062126データセットと固定AIPerfハーネスを使用し、ウォームアップは除外され、固定合成受け入れ曲線の下で1時間のプロファイリングが行われます。ダウンロードされたデータセットリビジョンは記録されておらず、クローズドループリプレイは異なる完了リクエストミックスを生成する可能性があります。

コストモデルはSemiAnalysisの大規模ハイパースケーラー所有モデルを使用し、公開されたGPUコスト定数はAMD Instinct MI355XでGPU時間あたり$1.50、NVIDIA HGX B300で$2.26です。総トークン/$ = 総トークン/秒/GPU × 3,600 ÷ モデル化された$/GPU時間。一部の資本と利用前提は未公開のままです。

スナップショットは2026年10月9日に取得されたSemiAnalysis InferenceXデータで、提出日は2026年9月30日(AMD)および2026年8月16日(NVIDIA)です。PRベンチマークは記載されたワークロードを使用し、PR #59591は主要なイメージの後に実行された別の実験です。

次に最適化すべきこと

進行中のプリフィル/デコード分離作業は、別の注目ポイントを追加します:ワーカーバランスと、転送がデータが必要になる前に完了するかどうかです。

別のモデルやワークロードの場合、投稿は展開されたリクエスト形状から始めて、3つのチェックを行うことを推奨します:

  1. ランクを比較します。 複数のGPUが、それを組み立てるコレクティブの前に同じ結果を計算していますか?
  2. バッファを追跡します。 各一時テンソルには、選択したパスにコンシューマーがありますか?データはオンチップに留まりますか?
  3. 独立した作業を数えます。 実際のヘッド、トークン、タイルは、負荷を分散するのに十分なワークグループを生成しますか?

必要な応答速度で完全なリクエストを測定し、出力と保存された状態の両方を確認します。より高速なカーネルは、次のトークンを妨げる作業を短縮する場合に重要です。

Kimi K3の取り組みは、AMD、Embedded LLM、および広範なコミュニティからの50以上のマージされたAMD/ROCm関連vLLM PRに及び、AITER開発、統合、レビュー、ベンチマークに加えて、アテンション、リカレンス、MoE、量子化、グラフ実行、キャッシュサポート、プラットフォーム統合にわたる貢献が含まれます。

Sources