AWS EFAでの1兆パラメータモデルの実現

Perplexityでは、製品、API、研究チームのために最適なモデルを採用しています。Kimi-K2のような大規模なオープンソースのMixture-of-Expertsモデルには特有の課題があります。8基のNVIDIA H200 GPUを搭載した最大の推論ノードであっても、これらのモデルを効率的に収容することはできず、マルチノード展開が不可欠だからです。我々は、ConnectX-7上で最先端のレイテンシを達成し、DeepEPの性能を凌駕するエキスパート並列処理用のカーネル群を発表します。また、これらのカーネルは、AWS Elastic Fabric Adapter (EFA) において実用的なレイテンシを達成した最初の事例でもあり、これにより1兆パラメータ規模のモデル展開が可能となります。

GitHubで我々のカーネルを試し、arXivで詳細な研究論文を読んでください。

Mixture-of-Experts (MoE) は、推論時に合理的なレイテンシを提供しつつ、数千億から数兆のパラメータへとスケールするための標準的なアーキテクチャとして定着しました。MoEモデルは、Transformerの密な層を、一連のエキスパートと、トークンをどのエキスパートと乗算するかを決定する小さな線形ルーティング層に置き換えます。同一のエキスパートに関連付けられたトークンをグループ化することで、この問題は疎なものとなり、単純に並列化が可能になります。これは、異なるエキスパートを、複数のノードにまたがる異なるGPU上にホストできるようになったためです。

torch.distributedやNCCLといった集合通信ライブラリを使用して容易に実装できるTensor並列処理 (TP) やデータ並列処理 (DP) といった他の並列化形式とは異なり、MoEルーティングには疎なピア・ツー・ピア通信が関わります。ディスパッチカーネルは、入力されたトークン群を分割してエキスパートをホストするランクへディスパッチしなければならず、一方でその双対問題である結合カーネルは、処理されたトークンを受け取り、加重平均を計算しなければなりません。これらの転送は集合通信にはマッピングされず、低いレイテンシを達成するにはカスタマイズされたカーネルが必要となります。

Perplexityでは、ポータブルなノード間カーネルと専門化されたノード内カーネルの両方を開発しており、以前の記事であるMoE通信において我々の研究をより詳細に紹介しました。また、別のブログ記事では、マルチノード展開によって大規模なMoEモデルのレイテンシを下げ、スループットを向上させることができることを示しました。

本記事では、NVSHMEMベースの以前の研究成果からパフォーマンスを大幅に向上させた、新しいノード間カーネルの実装について深く掘り下げます。ConnectX-7で最先端の結果を示すだけでなく、このカーネルが我々の推論ランタイム内で最大のオープンソースモデルを効率的に提供するために使用できることを実証します。

Tensor並列処理とエキスパート並列処理は、NVLinkがマイクロ秒単位のレイテンシで高いスループット（900GB/s）を提供するため、単一ノード内で効果的にスケールできます。NVIDIA GB200およびGB300システムには、1兆パラメータモデルの重みとKVキャッシュを保持するのに十分なメモリが搭載されていますが、最大8基のH200 GPUを搭載するAWS p5enインスタンスでの展開は1120GBのHBMに制限されており、これをモデルの重みとKVキャッシュの両方で分割する必要があります。結果として、一部のモデルではInfiniBandを介したノード間展開が不可欠となります。

InfiniBandは依然として非常に高速ですが、スループットは400Gbps (50GB/s) しか提供できず、MoEルーティングに数十から数百マイクロ秒のレイテンシを追加します。幸いなことに、共有エキスパート、マイクロバッチ処理、計算と通信のオーバーラップを使用することで、追加されたレイテンシを隠蔽し、競争力のあるコストで実用的なスループットを実現できます。DeepSeekは、H800 GPUクラスターとNVIDIA ConnectX-7 InfiniBandアダプターに展開されたDeepSeek-R1を通じて、これらの技術を開拓しました。

AWSクラスターでは、独自のElastic Fabric Adapter (EFA) を使用するため状況が異なります。これらはトレーニングワークロードで一般的に使用される集合通信において400Gbpsのピークスループットを提供しますが、MoEのディスパッチおよび結合中に交換されるメッセージサイズにおいてはConnectX-7にわずかに及びません。さらに、EFAはGPUDirect Asyncをサポートしていないため、転送を開始するにはGPUとNICを仲介するCPUプロキシスレッドが必要です。複雑さが増すだけでなく、追加のPCIeトランザクションにより、すべての転送にマイクロ秒単位のオーバーヘッドが加わります。NVSHMEMのようなフレームワークを使用すれば、カーネル作成者はデバイスに依存しない実装を提供できますが、汎用APIはオーバーヘッドを発生させ、コストとレイテンシの両面で展開を非効率なものにします。

以前のカーネルは、NVSHMEMを通じて基礎となるトランスポート実装を抽象化していました。これらは、nvshmemx_putmem_signal_nbi_warpに依存してトークンを一つずつ転送し、アトミックカウンタを使用してピアを同期していました。これはIBGDAを備えたConnectX-7 NICでは妥当なパフォーマンスを提供しましたが、プロキシベースのConnectX-7 IBRC実装は大幅に遅く、EFA実装では1ミリ秒未満のレイテンシを達成できませんでした。

我々は、これらのカーネルを基礎となるハードウェアのテストベッドとして使用するとともに、モデルに公開する最も適切なインターフェースを特定しました。我々は、カーネルとNIC、あるいはプロキシスレッドとの間の頻繁な転送と相互作用がパフォーマンスを阻害していると仮説を立てました。さらに、我々の実装は信頼性以外にハードウェアからの保証に依存していなかったため、同期プリミティブのオーバーヘッドを削減することで高速化できる可能性が示唆されました。

これらの観察に基づき、我々は同一のインターフェースを公開するEFAベースのカーネルを考案し、最高のパフォーマンスを実現するために、デバイスカーネルとホストプロキシの相互作用を内部で最適化しました。

ホストプロキシを介したRDMAとNVLinkのハイブリッド通信。送信側と受信側のrank、プロキシワーカー、NICを示しています。

新しいカーネルはハイブリッドCPU-GPUアーキテクチャに従っており、GPUカーネルがデバイス上のモデルと対話し、ホストCPU上で実行されるプロキシスレッドがNICとの対話を処理します。新しいカーネルは、KVキャッシュ転送のために構築したTransferEngineを再利用しています。

MoEルーティングはディスパッチカーネルによって実行されます。このカーネルはトークンと、モデルからルーティング先のエキスパートを受け取り、エキスパートをホストするランクへネットワーク経由で送信します。グループ化されたGEMMカーネルがトークンとエキスパートの重みを乗算した後、結合カーネルがトークンをソースランクへ転送し、加重平均を計算してからモデルへ返します。

ディスパッチカーネルと結合カーネルはどちらも、オーバーラップとマイクロバッチ処理を促進するために、送信側と受信側の半分に分割されています。送信側はトークンを送信バッファにシャッフルして、単一のRDMA書き込み操作を通じてピアへ書き込めるようにします。送信側を実行した後、GPUはネットワーク経由の転送がバックグラウンドで保留されている間に、他の作業（共有エキスパートの乗算など）を実行できます。受信側カーネルは、すべての転送が受信されるまでブロックします。バックグラウンドでは、ホスト側のワーカースレッドが転送を処理します。送信バッファが満たされるとGPUから信号が送られ、その後RDMA書き込みをディスパッチします。NICが完了を通知すると、受信側のブロックが解除され、トークンがモデルへ返されます。

これらのカーネルはホストと密結合されており、統合メモリとGDRCopyの両方を通じて互いに通信します。CUDA統合メモリにより、デバイスメモリまたはホストメモリを両者のアドレススペースにマッピングし、PCIeを介した通信を透過的に処理できます。

統合メモリは一括転送に適していますが、低レイテンシのポーリングはGDRCopyを介して行われます。GDRCopyはGPUDirect RDMAを使用して、マイクロ秒単位のレイテンシでGPUメモリへの読み書きを行います。

RDMAに加えて、ノード内ではカーネルはNVLinkを活用してトークンを交換します。これにより、EP32では最大4分の1、EP64では最大8分の1のトラフィックをオフロードでき、全体として低いレイテンシを達成できます。

ディスパッチカーネル用の送信バッファの設計は、ホストプロキシを伴う実用的なレイテンシを達成するために不可欠です。N個のランクがそれぞれT個のトークンをディスパッチし、それぞれがE個のエキスパートのうちR個のエキスパートへルーティングされる場合、すべてのトークンが特定のランクへルーティングされる最悪のケースでは、ランクは最大N * T * max(R, E / N)個のトークンを受信します。さらに、この最小境界は、プライベートな受信バッファが過剰なメモリを消費するため、送信側が書き込みを連続的に配置できる場合にのみ達成可能です。送信側が競合なく書き込みを連続的に配置できるようにするには、ルーティング情報の完全な知識が各ランクに配布されていなければなりません。結果として、我々のディスパッチカーネルはまずエキスパートごとのトークン数を交換し、その後、各ランクが書き込み先と、どのランクからトークンを受信するかのオフセットを決定します。

単にルーティング情報を交換するだけでは、各ランクが独自の書き込みをディスパッチする前に、書き込みの到着を待機しなければならないため、不必要に高いレイテンシが生じます。その代わりに、各受信側ランク上に少量のプライベートな送信者ごとのスペースを確保し、ルーティング情報の書き込みと並行してトークンの一部をディスパッチします。これにより、ディスパッチプロセスの大部分を通じて完全な帯域幅利用率が保証されます。ルートが交換された後、残りのトークンはリモートピアごとに単一の書き込みによって転送されます。ルーティング情報の交換と固定サイズのバッファ、およびそれらを満たす転送については、以下の図で説明します。

MoEディスパッチのバッファ配置。トークンデータとルーティングメタデータがどのように配置され、2つのrank間で交換されるかを示しています。

結合カーネルは大幅に単純です。これは、ルーティング情報がディスパッチ段階からすでに利用可能であるためです。ピアごとに単一の書き込みでトークンをコピーし、その後受信バッファから読み取って、トークンの加重平均を出力テンソルへ計算します。

NVLinkの場合、ノード内の送信バッファと受信バッファは、近隣のランクから直接トークンをプッシュまたはプルできるように、ピアにマッピングされます。

以下の図は、GPU、CPU、NIC間のタイミングと相互作用の詳細を示しており、モデルのカーネルと交互に実行されるディスパッチと結合の1つの反復を示しています。

作業はディスパッチ送信で開始され、カーネルがルーティング情報を集約し、単一ブロック内の共有メモリでのアトミック加算によってエキスパートごとのトークン数を計算します。トークン数が利用可能になると、ホストメモリへコピーされ、直ちにプロキシへ通知されてすべての他のピアランクへルーティング情報をディスパッチし始めます。並行して、ディスパッチはトークン数の累積和の計算を継続し、連続的なピアごとの送信バッファにトークンをパックするために必要なオフセットを特定します。オフセットが利用可能になると、EFA経由で送信されるすべてのトークンが送信バッファへパックされます。プロキシは、ノード間ピアの送信バッファへの書き込みを投稿するように通知されます。ノード内では、カーネルはNVLinkを介して同じノード内のプライベートバッファへトークンを直ちにプッシュします。

MoEディスパッチのタイミングと連携。段階的なGPU、CPU、NICの操作とその同期を示しています。

MoEコンバインのタイミングと連携。GPU、CPU、NICの操作とその同期を示しています。

ランクがすべてのピアからルーティング情報を受信すると、それを中央集約し、ディスパッチと結合のためのトークンを特定するために必要な様々なオフセットを計算します。この作業は転送が保留されている間にCPU上で行われるため、レイテンシクリティカルではありません。その後、すべてのオフセットが判明すると、プライベートバッファサイズを超えるトークンの残りの書き込みが投稿され、ディスパッチの送信フェーズが完了します。

受信側では、カーネルはまず、NVLinkを介してプライベートバッファの容量を超える同じノード内のトークンをプルします。次に、ホストプロキシがすべての転送の完了を通知するまでブロックします。信号を受信すると、トークンは受信バッファからグループGEMMへ渡される出力テンソルへとシャッフルされます。バッファはディスパッチと結合の間で再利用されるため、NVLinkとRDMAのバリアの両方が、次のカーネルが受信バッファを上書きする前に読み取りが完了していることを保証します。

結合カーネルも大まかに同様のパターンに従い、トークンのシャッフル、プロキシへの信号送信、完了の待機を行います。しかし、書き込みを投稿するためのメタデータは即座に利用可能で、すべてのリモートオフセットも判明しているため、ノード間で転送されるすべてのトークンはNVLinkを介して受信側へプッシュされ、ノード間トークンは単一の書き込みを介して転送されます。結果を返す前に、カーネルは受信バッファからトークンを累積し、出力テンソルへ書き込みます。

実装のプロトタイプを作成するために、我々は当初、RDMA書き込みを発行および追跡するために、非集約型プリフィルおよびデコード用に構築されたTransferEngineを再利用しました。しかし、各ノードが同じピアに多数の書き込みを投稿する非集約型デコードとは異なり、MoEディスパッチはすべてのピアに対して一度の操作で小さな転送を投稿します。そのままではパフォーマンスが最適ではなく、ConnectX-7上でDeepEPに約20us遅れをとっていました。結果として、我々はTransferEngineを専門化し、多数のピアとの相互作用を最適化するために、スキャッタとバリアという2つの操作をサポートするようにしました。

スキャッタ操作とバリア操作の両方は、基礎となるトランスポート実装で使用されるデータ構造（Work Request (WR) テンプレート）を事前に生成するために、事前に登録可能なピアグループを公開します。スキャッタは、同じソースバッファからの異なるスライスを、ピアグループのすべてのメンバーの異なるオフセットへコピーします。バリアは、イミディエイト操作を介してすべてのピアに信号を送ります。これらの操作を一括でキューに入れ処理することで、クリティカルパス上の貴重なマイクロ秒を節約できます。2基の200Gbps NICが集約されて400Gbps帯域幅を提供するEFAでは、各転送のバイトをシャーディングする代わりに、グループピアを2つのNICにシャーディングします。

当初、TransferEngineはlibfabricを使用してEFA上でのみ開発されましたが、MoEカーネルはlibibverbsを使用してTransferEngineにConnectX-7サポートを追加した後、コードを1行も変更せずにConnectX-7上で動作しました。Scalable Reliable Datagram (SRD) ベースのEFAと比較して、ConnectX-7は接続設定とピア管理のため追加作業を必要とします。我々は、Reliable Connection (RC) をブートストラップするために、Unreliable Datagram (UD) キューペアを使用します。リモートピアごとに2つのキューペアを作成します。1つは双方向のSEND/RECV用、もう1つは片方向のWRITE_IMM用です。これは、RECVとWRITE_IMMの両方が順序どおりにRECV操作を消費する一方で、両方の操作で異なるバッファサイズを必要とするためです。さらに、ConnectX-7とGPUメモリ間のPCIeトランザクション順序を緩和します。EFAと同様に、WRテンプレートの最適化を実装します。さらに、同じピアに複数の小さなメッセージを送信する場合、メッセージレートを向上させるためにWRチェイニングを使用します。

GPU側の作業は両方のアダプターで同一であるため、我々は基礎となるトランスポート実装を限界まで押し上げ、DeepEPの性能を凌駕することで、EFAとConnectX-7のMoEルーティング間のギャップが基礎となるハードウェアの能力によって決定されることを強調しました。

我々は、ConnectX-7およびEFAアダプターを搭載したH200システム上で、7168の隠れ次元、ブロックスケールfp8ディスパッチ、bf16結合を備えた標準的なDeepSeek-V3およびKimi-K2構成で、ディスパッチおよび結合のレイテンシを評価しました。ディスパッチにはバッチサイズ128、結合には4096を考慮しています。ConnectX-7上では、DeepEPおよび我々の以前のNVSHMEMベースのpplxカーネルと比較し、GPUDirect Async (IBGDA) とホストプロキシ (IBRC) の両方を使用しました。EP16、EP32、EP64でそれぞれ459us、582us、692usの結合レイテンシを達成しました。我々は、519us、966us、1159usを達成するUCCL-EPのような汎用プロキシに依存する他のEFAベースのカーネルを上回っています。

EP64 MoEデコードにおけるディスパッチとコンバインのレイテンシ。示されている実装と通信経路を比較しています。

デコードにおいては、我々のカーネルはDeepEPよりも数マイクロ秒遅いにもかかわらず、結合カーネルが大幅に高速であるため、全体としてより低いレイテンシを達成しています。ホストプロキシを使用しているものの、一括転送のおかげで、すべてのIBGDAベースのプロキシレス実装のパフォーマンスを上回っています。また、プロキシの実装が重要であることも示しました。我々の高度に専門化されたプロキシは、汎用的なNVSHMEM実装よりも1桁高速です。256KBのパケットを転送する場合、EFAはConnectX-7の2倍遅いですが、これはMoEルーティングで交換される一般的なパケットサイズに近いため、2つのアダプター間のギャップは全体としてそれほど大きくありません。

EP64 MoEプレフィルにおけるディスパッチとコンバインのレイテンシ。pplxとDeepEPの通信経路を比較しています。

我々のデコード最適化カーネルは追加の調整なしでプリフィルバッチサイズにスケールできますが、パフォーマンスはDeepEPに及びません。これは主に、DeepEPが転送データ量を削減するために送信側で累積するのに対し、我々は結合の受信側でのみトークンを累積しているためです。ただし、この最適化は精度に影響を与える可能性があります。さらに、カーネルは1回の書き込みでデータを転送するため、送信バッファと受信バッファにより多くのメモリを使用します。

我々のカーネルの実用性を証明するため、Elastic Fabric Adapter (EFA) とH200 GPUを搭載した複数のp5enインスタンスでDeepSeek-V3 (671Bパラメータ) およびKimi-K2 (1Tパラメータ) を提供する自社推論エンジンのエンドツーエンド性能をベンチマークしました。モデルをドラフト検証モードで実行し、単一のMTP層を使用して80%の受け入れ率を想定したレイテンシを報告します。

DeepSeek-V3では、効率的なNVLinkのみのディスパッチおよび結合カーネルを実行するEP=8のノード内ベースラインと比較します。ノード間展開は全体として、単一ノード展開の性能と一致するか上回っており、特に中規模のバッチで顕著であり、ノード間スケーリングが全体のスループット向上をもたらすことを示しています。

Kimi-K2では、これらのカーネルによりサービス提供が可能になります。そのサイズのため、モデルはH200ノード上のデータ並列処理 (DP) を使用して単一ノード内に収まらないため、サービス提供にはノード間展開が必要です。モデルのアテンションヘッドが少ないため、DeepSeek-V3よりも優れたレイテンシを達成しており、EFA上で実用的なレイテンシでモデルを提供できます。

DeepSeek-V3のGPUごとのスループットと、EP32、EP16、EP8構成での出力速度の比較。

Kimi-K2のGPUごとのスループットと、EP32およびEP16構成での出力速度の比較。

マイクロバッチ処理も検討しましたが、非常に高いバッチサイズの場合や、複数のドラフトトークンを使用する場合にのみレイテンシが向上しました。グループ化されたGEMM問題を2つに分割するメリットを得るには、トークンバッチサイズが十分に高密度である必要があります。依然としてMoEディスパッチと結合に大部分の時間が費やされているため、カーネルへのさらなる改善は、結果としてレイテンシを改善し、マイクロバッチ処理を促進することになります。

我々はAWSエンジニアと緊密に連携し、EFA上のパフォーマンスを改善しています。我々の側では、データプレーンのオーバーヘッドを削減するlibfabricのアップデートに従います。さらに、efa-directを実験し、ユーザー空間のオーバーヘッドをさらに削減してハードウェアと直接対話するようにします。

NVIDIA Magnum IO NVSHMEMおよびGPUDirect Asyncを使用したHPCシステムのネットワークパフォーマンス向上

BibTeX

@online{perplexity2025EnablingTrillionParameterModelsOnAwsEfa,
  author = {{Perplexity} and Nandor Licker and Kevin Hu and Vladimir Zaytsev and Lequn Chen},
  title = {{Enabling trillion-parameter models on AWS EFA}},
  date = {2025-11-04},
  url = {https://www.perplexity.ai/hub/blog/enabling-trillion-parameter-models-on-aws-efa}
}

はじめに

InfiniBand上での推論

NVSHMEMベースのpplxカーネル

EFA上でのMoEオール・ツー・オール通信

送信バッファと受信バッファ

ディスパッチと結合

MoEルーティング用TransferEngine

評価

ディスパッチ/結合レイテンシ

Kimi-K2 / DeepSeek-V3 スループット

今後の展望

参考文献

引用