AMD Instinct MI355X 上での MiniMax M3 の最適化

vLLM は、AMD Instinct MI355X 上での MiniMax M3 のサービング性能を大幅に最適化し、低並行性での MXFP8 標準サービングにおいて出力スループットを最大 3.14 倍に向上させました。これは、GEMM のタイリングやカーネル結合から分散システムのオーケストレーションに至るまで、スタック全体にわたるボトルネックを段階的に特定・解決することで達成されました。

AMD Instinct MI355X 上のパフォーマンスベンチマーク

最適化の取り組みにより、MiniMax M3 のさまざまな設定においてスループットとレイテンシが大幅に改善されました。

  • MXFP8 標準サービング（並行性 32）： 出力トークン数/s/GPU は 109.1 から 342.4 に上昇（3.14 倍）。中央値の最初のトークン到達時間（TTFT）は 1.46 秒から 0.67 秒に低下し、平均出力トークンあたり時間（TPOT）は 69.1 ミリ秒から 22.1 ミリ秒に低下しました。
  • MXFP8 標準サービング（並行性 128）： 出力トークン数/s/GPU は 297.8 から 623.7 に上昇（2.09 倍）。中央値の TTFT は 3.53 秒から 1.54 秒に低下し、平均 TPOT は 100.7 ミリ秒から 48.8 ミリ秒に低下しました。
  • MXFP4 サービス（並行性 128）： TP4/EP1 で出力トークン数/s/GPU は 212.1 から 716.8 に上昇しました。さらに TP2/EP1 での最適化により 943.5 出力トークン数/s/GPU に達し、初期結果と比べて 4.45 倍の向上を達成しました。
  • 予測的デコード（EAGLE3）： TP4/EP1 で並行性 128 の状態で 682.4 出力トークン数/s/GPU を達成しました。
  • P/D 離散化： 並行性 512 で 6,370.5 合計トークン数/s/GPU を達成し、中央値の TTFT は 1.32 秒でした。

カーネルおよびオペレータの最適化

パフォーマンス向上は、局所的な形状解析を体系的に実施し、余分な計算を排除することで実現されました。

局所的な形状とタイリング

vLLM は、プレフィル（大 M）とデコード（小 M）の異なる状態に対応する GEMM ランチャを最適化しました。デコードではより狭い N タイルを選択し、ループ回数を減らすために大きな K ステップを採用することで、TP8 8K/1K の出力スループットを 7.8–9.4% 向上させました。また、グループ化された MoE プログラムの順序を再調整することで、隣接するプログラムが GPU キャッシュからアクティベーション行やエキスパート重みのタイルを再利用できるようにし、TP4 テストで 1.08×–1.46× の向上を達成しました。

シェアードエキスパートの統合

MiniMax M3 のシェアードエキスパートは当初、別個の密行列 MLP パスとして実行されていました。vLLM はこのシェアードエキスパートをルーティングエキスパートテーブルに統合し、グループ化 GEMM が両方を同時に処理できるようにしました。これにより、別々の起動や中間トラフィックが削減され、並行性 1 では出力スループットが 30.2% 向上し、並行性 128 では 5.6% 向上しました。

スパースアテンションとインデックス再利用

スパースアテンションのオーバーヘッドを削減するために、vLLM はインデックス共有を実装しました。後続のレイヤーが前のレイヤーの top-k ブロック決定を再利用することで、並行性 1 での平均 TPOT を約 10% 低下させました。さらに、vLLM は MiniMax M3 の 128 トークンの論理ブロックを AITER の 16 トークンページにコピーせずにマッピングするスパースページアダプタを開発し、TP4、並行性 256 で MXFP8 では 5.56%、MXFP4 では 6.93% の出力スループット向上を達成しました。

分散システムと予測的実行の最適化

最適化は個々のカーネルにとどまらず、複数 GPU の連携および予測的デコードの統合にも及びました。

EAGLE3 予測的デコード

vLLM は AMD ハードウェア上で EAGLE3 ドラフトモデルを統合し、リクエストレベルのインデックスバッチ処理を実装して、すべてのドラフト位置を一度に処理しました。これによりインデックスカーネルの性能が最大 48.9% 向上しました。さらに、AITER のスパースページアテンションを複数トークン検証に拡張することで、TP4 テストで MXFP8 では 7.90%、MXFP4 では 8.32% の出力スループット向上を達成しました。

プレフィル/デコード（P/D）離散化

vLLM は、各レイヤーごとに正確な転送ジオメトリとバイトオフセットを導出し、正しい KV の引き渡しを保証することで P/D 離散化を実装しました。システムチューニングの結果、すべてのワーカーを TP4 に移動し、8K/1K タスクでは 2 つの TP4 プレフィルワーカーと 1 つの TP4 デコードワーカーの比率を採用することで、並行性 512 で 6,370.5 合計トークン数/s/GPU、中央値 TTFT 1.32 秒を達成しました。

エージェントワークロードの分析

不規則な出力と再利用可能なプレフィックスを含むエージェントコードを対象とした AgentX ベンチマークを使用して、vLLM は MXFP4、EAGLE3-GQA、プレフィックスキャッシュを TP4 で並行性 28 でテストしました。システムは 127.4 出力トークン/s/GPU、合計 509.5 出力トークン/s を達成し、p50 TTFT は 645 ミリ秒、p50 TPOT は 41.3 ミリ秒でした。分析の結果、GPU キャッシュヒット率は 92.1%、KV キャッシュ利用率は 88.5% であり、今後の最適化は GEMM のさらなるチューニングではなく、プレフィックスの整合性とエヴィクションポリシーに焦点を当てるべきであることが示唆されました。

最適化手法のチェックリスト

将来のモデルサービングパスの最適化のために、vLLM は以下の技術的チェックリストを提案します。

  1. 局所的な形状解析： シャーディング後に局所的な形状のヒストグラムを記録し、複製されたヘッドやルーティングされたトークン数を含める。
  2. 繰り返しの推定： 各レイヤーの作業量にレイヤー数、出力トークン数、アクティブなリクエスト数を乗算して計算する。
  3. バイトプレーンの分離： 計算用テンソル、永続状態、通信の区別を行う。
  4. ファストパスの検証： すべてのファストパスの実行条件を記録し、ディスパッチトレースで実際の実行を検証する。
  5. 正しさのゲート： パフォーマンスゲートと併せて正しさのチェックを実装する。
  6. 反復的なプロファイリング： 葉のカーネルがフラット化された後は、上位レベルのキュー、所有権、キャッシュ容量、OS の制限を検査する。

Sources