ボトルネックを追う：AMD MI355XでMiniMax M3を高速化した方法

導入

MoEモデルの推論を速くする作業は、最速のカーネルを一つ見つければ終わるものではない。vLLMによるMiniMax M3とAMD Instinct MI355Xの検証は、ボトルネックが移動するたびに観測点を変える実践例だ。各rankに実際に届く形状は何か、層やtokenをまたいで同じ仕事を繰り返していないか、どのバイトが移動しているか、想定した高速パスが本当に選ばれているか、そしてカーネルが支配的でなくなった後にどのキューが伸びているかを確認する。

段階的に積み上がった性能向上

固定トポロジーのMXFP8標準構成では、並列度32の出力スループットが109.1から342.4 token/s/GPUへ伸び、3.14倍になった。中央値TTFTは1.46秒から0.67秒へ、平均TPOTは69.1ミリ秒から22.1ミリ秒へ低下した。並列度128では、同じTP4/EP1の4 GPU構成が297.8から623.7 token/s/GPUになった。MXFP4はTP4/EP1で212.1から716.8へ向上し、その後TP2/EP1で943.5に達した。ただし後者は配置密度の変化も含むため、固定トポロジーの速度向上とは区別する必要がある。EAGLE3は並列度128のTP4/EP1で682.4に達し、P/D分離とトポロジー再調整後は、並列度512で6,370.5 total token/s/GPU、中央値TTFT 1.32秒となった。

最適化から得られる五つの示唆

• モデル図ではなくローカル形状を見る。 TP分割、KVヘッドの複製、padding、routingが、各rankのM、N、Kを決める。TP8ではQヘッドが分割される一方、4個のKVヘッドと4個のインデックスヘッドは複製されるため、融合QKV投影のローカルNは1536であり、全体の次元を単純に8で割った値ではない。大きなMと小さなMを分けたlauncherにより、TP8の8K/1Kスループットは7.8%から9.4%向上した。
• 必要な計算でも、分離された実装は統合する。 共有エキスパートは当初、独立したdense MLPとして実行されていた。これをrouted expertの表に追加し、grouped GEMMでまとめて処理することで、モデルの計算内容を変えずにlaunchと中間データの移動を削減した。改善幅は並列度1で30.2%、並列度128で5.6%だった。
• 不変な処理をループの外へ出す。 MXFP8の重みとscaleの再配置はモデルロード時に実施するよう変更された。投機的デコードでは、MSA indexerをtokenごとのworkgroupからrequestごとのworkgroupへ変え、複数のdraft位置でkeyのロードを再利用した。インデックスカーネルは最大48.9%、エンドツーエンドでは約3.3%向上した。
• 制御用データの移動も測る。 疎な注意機構は演算量を減らす一方、top-k選択、物理ページの対応付け、メタデータ処理を追加する。近接する疎な層が似たblockを選ぶ性質を利用してインデックスを共有した結果、平均TPOTは並列度1で約10%、より高い並列度でも約4%低下した。
• カーネルが頭打ちならキューを追う。 graph executionはlaunchコストを下げるが、最適なバックエンドは入力長と並列度で変わる。native MXFP8、emulated linear、sparse paged attentionにも得意な領域がある。P/D分離でも、まずKVの受け渡しを検証し、その後に実際にリクエストが待つ場所へ容量を追加する必要がある。

意義

この事例が示すのは、疎なMoE推論に適用できる性能分析の型だ。「prefill」「decode」「TP4」といった名称だけでは、実行されるローカル形状やバックエンドを十分に表せない。また、カーネルの改善率をそのままサービス全体の改善率と見なすこともできない。インデックスカーネルが最大48.9%向上しても、エンドツーエンドでは約3.3%だったのは、別の処理がクリティカルパスに残っていたためだ。最も再利用しやすい方法は、形状を測り、重複を消し、データ移動とディスパッチを確認し、ボトルネックの移動先であるキューを追い続ける反復である。

出典：vLLM Blog