Step3 推論最適化：AMD Instinct MI308X で NVIDIA H20 比 1.30 倍の Decode スループットを実現

背景

StepFun の Step3 は、321B パラメータの Mixture-of-Experts (MoE) マルチモーダルモデルで、トークンあたり 38B のパラメータがアクティベートされます。61 層で構成され、そのうち 56 層が 3-in-48 のエキスパート選択を使用する MoE 層であり、Multi-Matrix Factorization Attention (MFA) を導入することで、KV-cache の要求量を DeepSeek V3 のトークンあたりの attention コストの約 22% に削減しています。

StepFun は NVIDIA H20 GPU で Step3 を提供しており、代替として AMD Instinct MI308X の評価を希望していました。Step3 がオープンソースとして公開される前に、Moreh は Step3 と同じアーキテクチャを持つプライベートモデルの MI308X 上での推論最適化を依頼されました。これは Moreh のカスタムモデル最適化サービスの一例であり、独自のモデルアーキテクチャに対して Moreh vLLM を適応させるものです。

なぜ Decode に MI308X を選ぶのか

AMD Instinct MI308X は、中国市場向けに提供される MI300X の派生モデルです。計算コア数は MI300X の 1/4 ですが、同じ HBM3e メモリ容量と帯域幅を維持しています。これにより MI308X は LLM 推論の decode フェーズに特に適しています。decode フェーズは計算バウンドではなくメモリ帯域幅バウンドだからです。トークンは自己回帰的に一つずつ生成され、ボトルネックはメモリからのモデル重みと KV-cache の読み込みにあり — 行列乗算の実行ではありません。

最適化技術

• カスタム HIP attention kernel：デフォルトの vLLM Triton attention kernel が最大のボトルネックであり、GPU 時間の約 50% を占めていました。Step3 の MFA 構成（64 query heads、1 KV head、head dimension 256）に最適化されたカスタム HIP attention kernel を data parallelism で開発しました。このカーネルにより、decode batch の attention レイテンシが 72%、混合 prefill/decode batch では 37% 削減されました。
• CUDA graph：GPU kernel レイテンシが大幅に削減された後、CPU 側のオーバーヘッドが decode ステップの次のボトルネックとなりました。DP8-EP8 parallelism を使用する Step3 モデルに対して完全な CUDA graph キャプチャを有効にし、decode スループットを約 2,900 から 4,100 tok/s に改善しました。
• 混合 BF16–FP8 blockscale 量子化：BF16 と FP8 blockscale 計算の両方に対する網羅的な GEMM チューニングにより、最適な精度–効率のトレードオフを実現しました。
• 最適化された MoE one-stage kernel：Step3 の MoE 層に対するカスタム kernel 最適化で、inter_dim パラメータに焦点を当てています。
• Shared-expert MLP 融合：shared-expert MLP を MoE 層内に統合し、冗長な計算を削減して推論レイテンシを改善しました。
• MoRI EP 統合：AMD GPU 上での効率的な expert-parallel all-to-all 通信のために MoRI ライブラリを統合しました。

性能結果

8× MI308X 上の Moreh vLLM を、StepFun が報告した 8× NVIDIA H20 上の数値と比較してベンチマークを行いました。同じテスト構成を使用：ISL=4096、OSL=256、Concurrency=256、DP8-EP8 parallelism（attention に 8 way data parallelism、MoE に 8 way expert parallelism）。

結果は、2 つのフェーズ間で明確な差異を示しています：

• Decode：Moreh vLLM は MI308X 上で 4,082 tok/s を達成 — StepFun の H20 ベースラインと比較してスループットが 1.30× 向上、レイテンシが 1.30× 低減（63 ms vs. 82 ms）。
• Prefill：H20 は計算バウンドの prefill フェーズで優位性を維持（13,780 vs. 9,601 tok/s）。これはより強力なオンチップキャッシュサブシステムを考慮すると想定通りです。

prefill–decode 分離アーキテクチャを採用した本番環境では、decode フェーズに最も多くの GPU が割り当てられます。MI308X の優れた decode 性能は、大規模サービングにおけるコスト効率に直結します。

まとめ

今回の取り組みは、AMD Instinct MI308X が Moreh vLLM のモデル固有の最適化と組み合わせることで、大規模 MoE モデルにおいて NVIDIA H20 を上回る decode スループットを実現できることを示しています。MI308X は計算能力に対して高いメモリ帯域幅を持ち、本番 LLM サービングデプロイメントで GPU 割り当ての大部分を占める decode フェーズにおいて、コスト効率の高い選択肢となります。

Moreh は独自モデルおよびファインチューニング済みモデル向けのカスタム vLLM 最適化を提供しています。AMD GPU でのモデル評価をご検討中の場合は、お問い合わせください。最適化のご支援について議論いたします。