Expert Parallelismを活用したAMD Instinct MI300X GPUでの毎秒21K出力トークンのDeepSeek推論

概要

大規模言語モデル(LLM)における主要なブレークスルーの一つは、Mixture-of-Experts(MoE)アーキテクチャの採用です。単一の大規模なfeed-forward network(FFN)レイヤーの代わりに、モデルはそれぞれexpertと呼ばれる複数の小規模FFNレイヤーを含みます。各トークンに対して、軽量なgatingネットワークが各ステップで処理するexpertの小さなサブセットを動的に選択します。この条件付き計算により、MoEモデルは効率的な計算利用を維持しながら、非常に大きなパラメータ数にスケールすることが可能です。例えば、2025年1月のリリース時に大きな話題を呼んだDeepSeek-R1 671Bモデルと、2025年8月にリリースされたOpenAIの人気オープンソースGPT-OSS 120Bモデルは、いずれもMoEアーキテクチャを使用しています。

しかし、大規模かつスパースな設計のため、DeepSeek-R1のようなモデルを大規模に効率的にサービングするには、高度な最適化技術が必要です。特に、Expert Parallelism(EP)を適用することで、高いスループット(すなわち、1秒あたりの総出力トークン数)を達成できます。これは、(1) GPU メモリにexpertの一部のみが格納されるため、より大きなバッチサイズが可能となり、(2) 個々のexpertのパラメータがロード後により多く再利用されるため、FLOPS per byteが向上するためです。しかし、活性化されたexpertを担当するGPUに活性化値を転送するdispatch/combineと呼ばれる細粒度の通信パターンを効率的に実装すること、およびexpert間の不均衡を解消することが課題となります。

AMDのソフトウェアパートナーであるMorehは、最近ROCmソフトウェアスタック上でEPを実装し、DeepSeek-R1推論を高スループットで実行できることを実証しました。LLM推論でこのような高性能を実際に達成するには、スタックの複数のレイヤーにわたって綿密に最適化されたソフトウェアが必要です。Morehチームは、ライブラリおよびモデル実装の両レベルでさまざまな最適化を適用しています。その結果、8x AMD Instinct MI300X GPUを搭載したサーバーで>21,000 tokens/secのデコーディングスループットを達成しました。

並列化手法

DeepSeek R1モデルでは、各decoderブロックに256個のexpertが含まれ、各トークンに対してそのうち8個が活性化されます。さらに、ルーティングに関係なく常に実行されるshared expertと、DeepSeekが導入した新しいアテンションメカニズムであるMulti-Head Latent Attention(MLA)レイヤーがあります。EPは通常、data parallelism(DP)と組み合わせて実装されます。各GPUは異なるexpertのサブセットを保持し実行します。一方、shared expertやMLAのように常に実行されるコンポーネントについては、各GPUが別々のリクエストバッチを並列に処理します。これに基づき、ノードあたりDP=8およびEP=8の並列化構成を適用しました。つまり、各GPUは32個のexpertを担当します。

さらに、prefill-decode disaggregation――すなわちprefillとdecodeステージを別々のノードで実行すること――を適用すると、クラスターシステムの全体的なスループットをさらに向上させることができます。prefillステージは入力シーケンス全体を一度に処理し、より計算バウンド(compute-bound)であるのに対し、decodeステージは出力トークンを自己回帰的に1つずつ生成し、よりメモリバウンド(memory-bound)です。実行を分離するだけでなく、prefillとdecodeの異なる計算特性に合わせて並列化と最適化の構成を調整することで、GPU効率をさらに向上させることができます。本記事では、一般的なアプリケーションではGPUの大部分がデコーディングに割り当てられるため、主にdecodeステージのパフォーマンス評価に焦点を当てます。

パフォーマンス最適化

新しい並列化手法を実装するだけでは、最良のパフォーマンスが自動的に得られるわけではありません。ハードウェアの潜在能力を完全に引き出すには、ライブラリからモデル実装に至るソフトウェアスタック全体を、新しい計算および通信パターンに対応するよう最適化する必要があります。Morehチームは、vLLMをベースに以下のような複数の最適化を実装しました。

GPU間ロードバランシング

EP実装の主要な課題の一つは、ワークロードがGPU間で均等に分配されず、ルーティング結果によって変動することです。MoEモデルはこの不均衡を可能な限り最小化するよう学習されますが、依然として本質的な限界があります。DeepSeekもこの問題を認識し、対処するためにEPLBと呼ばれる独自のロードバランシング手法を使用したと報告しています。明示的なロードバランシング戦略を適用せず、expertを単純に8つの連続したチャンクに分割した場合――すなわち、expert 0-31を最初のGPUに、32-63を2番目に、……、224-255を最後のGPUに割り当てた場合――GPU間で最大2倍のワークロード不均衡が観察されました。我々は各expertの活性化頻度を測定し、256個のexpertを各セットの総活性化頻度が可能な限り均衡するように32個ずつ8セットにグループ化するアルゴリズムを設計しました。(これは均衡数分割問題の変形と見なすことができます。)このアルゴリズムは、各セットが連続的に配置されるようexpertを並べ替え、その上でEPを適用します。expertの順列が変更されたため、gatingネットワークも新しいexpert順列を反映したルーティング結果を生成する必要があります。その結果、GPU間のワークロード不均衡を5%以内に削減することに成功しました。我々の知る限り、これはAMD GPU上での初のEPロードバランシング実装です。実際のサービングシナリオでは、活性化頻度を定期的に再測定し、それに応じて新しいexpert順列を生成することで、継続的なロードバランシングを実現できます。

計算-通信オーバーラップ

EPを適用する際のもう一つの課題は、dispatch/combine all-to-all通信によるオーバーヘッドを最小化することです。バッチを2つのマイクロバッチに分割して同時に実行すると、一方のマイクロバッチの通信を他方のマイクロバッチの計算とオーバーラップさせることができます。バッチサイズが十分に大きい場合――すなわち、高スループットのシナリオでは――このアプローチは計算パフォーマンスを犠牲にすることなく通信レイテンシを効果的に隠すことができます。このアプローチにはいくつかの実装バリエーションがあり、その中から我々はvLLMのDBO(Dual Batch Overlap)システムを採用しました。DBOシステムをMORI-EPを使用するよう修正し、特にMORI-EPライブラリの通信操作がDBOの2つのワーカースレッドから呼び出され、他の計算と適切に同期されるよう強化しました。

ライブラリ最適化

Morehチームは、GPU効率を最大化し、スループット(tokens/sec)とレイテンシ(最初のトークンまでの時間およびトークン間レイテンシ)の両方を改善するために、いくつかのライブラリ最適化も適用しました。以下は最適化の代表的な例です:

• 最適なGEMMおよびAttention Kernel選択:オンラインプロファイリングや手動チューニングを必要とせず、最適なGEMMおよびAttentionカーネルを動的に選択します。
• Fused MoE Kernel最適化:特に小さなバッチサイズにおいて、AMDのAITERライブラリよりも優れたパフォーマンスを提供する、高度に最適化されたfused MoEカーネルを実装します。
• FP8 KV Cacheサポート:KV cacheをFP8形式で格納およびロードできるMulti-head Latent Attention(MLA)カーネルを実装し、特に長いコンテキストのシナリオでパフォーマンスを向上させます。
• 垂直および水平Kernel Fusion:垂直fusion(例:fused RoPEカーネル)と水平fusion(例:shared expertにおける複数GEMMの統合)の両方を活用し、カーネルオーバーヘッドを削減して計算効率を向上させます。

これらのライブラリレベルの最適化およびEPとは独立したパフォーマンス改善の詳細については、Morehの技術レポートをご参照ください。

HIP Graphsの活用

HIP Graphsは、複数のGPU操作を静的グラフにキャプチャして一括で発行することで、CPUランタイムオーバーヘッドを削減する技術です。これは、個々の操作が比較的短く、GPUを完全に活用することが重要なdecodeステージにおいて特に不可欠です。しかし、静的グラフを構築するには、操作間で受け渡されるテンソルサイズが固定されている必要があります。EPでは、各expertの入力サイズは本質的に動的でルーティング結果によって決定されるため、これを静的グラフとして表現することは非常に困難です。我々はvLLMのDeepSeekモデル実装を修正し、EPをサポートしつつテンソルサイズを可能な限り静的にしました。これにより、HIP Graphを最大限に活用してCPUオーバーヘッドを最小化することができました。MoRIライブラリの操作をグラフに組み込む際にはいくつかの技術的課題がありましたが、最終的にMORI-EPのdispatch/combine操作をグラフ実行の一部として含めることに成功しました。

実験方法

実験は以下の仕様のMI300Xサーバーで実施しました。

• CPU: 2x AMD EPYC 9474F 48-core 3.6 GHz
• Main memory: 2,304 GB
• GPU: 8x AMD Instinct MI300X OAM GPU 192 GB
• Server: Gigabyte G593-ZX1-AAX1
• Operating system: Ubuntu 22.04.4 LTS
• ROCm version: 6.4.1

評価方法論はLMSysのブログ記事を参考に設計しました。同記事において、EPパフォーマンス評価の主要指標(prefill/decode disaggregationを前提)は、decodeノードあたりの総スループット(tokens/sec)とトークン間レイテンシです。LMSysは12ノードにわたってprefill/decode disaggregationを実装しましたが、我々は単一ノード環境でもdecodeパフォーマンスを測定することを目指しました。これを実現するため、vLLMに以下の修正を加えました:

• サーバー側では、prefill時間を除外し、decode時間とトークン数のみを蓄積してdecodeスループットを個別に測定しました。トークン間レイテンシは、標準的なベンチマーキング手法を使用してクライアント側で引き続き測定できます。
• vLLM V1エンジンのスケジューラは、元々prefillとdecodeリクエストを単一バッチで実行しており(mixed prefillと呼ばれる)、prefill時間を除外することが不可能でした。prefillとdecodeリクエストが別々のキューに発行され、独立して実行されるよう修正しました。

以下のコマンドを使用して、(修正された)vLLMサーバーにリクエストを送信し、パフォーマンスを測定できます。

vllm bench serve \
  --backend vllm \
  --model "deepseek-ai/DeepSeek-R1" \
  --metric-percentiles "10,25,50,75,90" \
  --percentile-metrics "itl,tps,ttft,e2el" \
  --port 8001 \
  --num-prompts 2048 \
  --max-concurrency 2048 \
  --ignore-eos \
  --dataset-name sharegpt \
  --dataset-path /app/dataset/ShareGPT_V3_unfiltered_cleaned_split.json \
  --sharegpt-input-len 2000 \
  --sharegpt-output-len 100

実験結果

レイテンシとスループットのトレードオフを観察するため、4つの異なる同時実行レベル(1024、2048、3072、4096)でスループット(ノードあたり)とトークン間レイテンシを測定しました。結果は以下の通りです。

最高の同時実行レベルで21,224.6 tokens/secのスループットを達成しました。これはSGLangチームが8x NVIDIA H100 GPUサーバーで報告した22,282 tokens/secとほぼ同等であり、我々の実装が最先端のEPパフォーマンスを提供することを確認しています。一方、最小レイテンシに最適化された構成(中央値トークン間レイテンシ91.48 ms)においても、スループットの低下は15%未満であり、システムが目標SLO(Service Level Objectives)に応じて最大同時実行数を柔軟に調整できることを示しています。

結論

Morehは、MoEモデル向けのさまざまな技術を適用・最適化することで、AMD Instinct MI300シリーズGPU上でEPによる高スループット推論を成功裏に実装しました。また、本記事で説明した効率的なEP実装を含む、さまざまな分散推論技術をAMD Instinct GPUクラスター上で自動的に適用するMoAI Inference Frameworkを提供しています。詳細については、Morehのウェブサイトをご覧ください。