効率的かつポータブルなMixture-of-Experts通信

Mixture-of-Experts (MoE) 通信向けの高性能かつポータブルなオープンソースライブラリを紹介します。本ライブラリは、標準的なAll-to-All通信プリミティブと比較して10倍の高速化を実現します。

当社の実装には、優れたMoE通信効率を実現するいくつかの重要な技術革新が含まれています。

GPU開始型通信 (IBGDA): GPUとNIC間の直接通信をサポートし、CPUの関与をバイパスすることでレイテンシを大幅に削減します。

通信と計算のオーバーラップ: 送信ステージと受信ステージを分離したスプリットカーネルアーキテクチャにより、ネットワーク転送と並行して計算を実行可能です。

最速のシングルノード性能: シングルノード構成において、従来の最速実装と比較して2.5倍低いレイテンシを実現しています。

効率的かつポータブルなマルチノード性能: 当社の実装は、標準的なAll-to-All通信より最大10倍高速な速度を達成します。高度に特化された実装よりは約2倍低速ですが、当社の手法はNVSHMEMバージョンやネットワーク環境（NVLink、CX-7、EFA）間での優れた移植性を提供します。

このライブラリは完全にオープンソースであり、https://github.com/ppl-ai/pplx-kernelsで公開されています。

本記事では、大規模MoEモデルにおけるエキスパート並列処理の課題を探求し、分散GPU環境全体で効率的にトークンをディスパッチおよび結合する当社の手法について説明します。

はじめに

DeepSeek R1やMixtral 8x7BなどのMixture-of-Experts (MoE) モデルは、トークンごとにアクティブ化される重みの数を制限することで、高密度モデルを改良しています。層の数や各デコーダ層の多層パーセプトロン (MLP) 内の線形層のサイズを増やすことでパラメータ数を増やすのではなく、MoEモデルは従来のMLPを複数のエキスパートとルーターに置き換えます。例えば、DeepSeek R1の671Bパラメータのうち、推論時に特定のトークンに乗算されるのはわずか37Bです。これにより、トークンのデコードに必要な計算量が大幅に削減され、高密度モデルと比較してレイテンシが改善されます。

MoEモデルは、高密度モデルと比較して推論にさらなる課題をもたらします。エキスパート自体は小さなMLP層ですが、各デコーダ層にはトークンのディスパッチ先エキスパートを決定するルーターが含まれており、各トークンは複数のエキスパートにディスパッチされます。ルーターは通常、確率分布を生成する小さな線形層です。一般的には、スコアの上位K個のエキスパートが選択され、最終的なアクティベーションは、ルーターの確率分布から導出された重みをエキスパートの出力に乗算して合計した加重平均として計算されます。

分散MoEシステムのレイテンシを最小化するには、複数のデバイス間で並列処理を活用できますが、これには通信上の課題が伴います。Mixtral 8x7BやLlama-70Bのようなモデルは、単一ノード内の8台のデバイスに収まり、高速かつ低レイテンシのNVLink相互接続（最大900Gbps）の恩恵を受けます。しかし、より大規模なモデルでは、InfiniBand（最大400Gbps）を使用して複数のノード間でエキスパートをシャーディングする必要があり、さらなるレイテンシの課題が生じます。

拡大

本記事ではエキスパート並列処理の問題を探求し、割り当てられたエキスパートを持つデバイスへ効率的にトークンをディスパッチし、エキスパート出力の加重和を計算して回収するGPUカーネルの実装について説明します。特定のネットワークハードウェアを対象とした過度に最適化された実装は存在しますが、当社は最小限のNVSHMEMプリミティブに依存しながらも、標準的なAll-to-All通信より10倍高速な性能を発揮する非常にポータブルな実装を提示します。当社の実装は、様々なネットワーク環境間で優れた移植性を維持しつつ、シングルノード構成において最先端の性能を達成しています。

シャーディングと並列処理

効率的な推論のためには、モデルの重みをデバイスメモリに保持しつつ、モデルのフォワードパスに必要なアクティベーション、KVキャッシュ、その他のバッファ用のスペースを十分に確保する必要があります。最も有能なモデルは最高性能のGPUの容量さえも超えるため、推論を複数のデバイスに分散させ、それら全体でモデルの重みを格納しなければなりません。重みのシャーディングスキームに基づき、計算のオフロードとデバイスの同期を行うために、異なる通信および計算スキームを使用する必要があります。

拡大

最初の図に示すエキスパート並列処理 (EP) は、エキスパート計算のみを並列化します。異なるエキスパートがそれぞれの重みを保持する異なるデバイスに割り当てられます。ルーティング後、トークンは対応するデバイスに送信され、その後に結果が回収および蓄積されます。ルーティングの複雑さは、モデルの他の部分における並列処理の度合いに依存します。各ランクがローカルに複製されたルーティングテーブルからトークンを選択できるため、他の層を複製することでルーティングの必要性を完全に排除できる可能性があります。しかし、1つのランクのみがルーティングを実行する場合、トークン、インデックス、重みをそれぞれのエキスパートにディスパッチするためにブロードキャストが必要です。最後に、オールギャザーまたはオールトゥオールブロードキャストにより、モデルの実行を継続するランクと出力トークンを同期させます。torchは必要なプリミティブを公開しているため、このような実装は比較的単純ですが、エキスパートからトークンを収集するギャザー操作と蓄積を融合させることで、多少の利点が得られる可能性があります。

拡大

エキスパートのみの並列処理は理想的にはスケールしません。モデルがアテンション、ノルム、サンプリングなどのエキスパート以外の層を実行している間、クラスタ内のノードがアイドル状態になる可能性があるためです。しかし、主にアテンションであるこれらの層の計算も、テンソル並列処理 (TP) から恩恵を受けることができます。ほとんどのモデルはマルチヘッドアテンションに依存しており、これはアテンションヘッドとその対応するQ、K、Vプロジェクションもデバイス間でシャーディングし、様々な層間でスライスを複製または収集できることを意味します。アテンションがすべてのデバイスに分散されている場合、オールギャザーによってアクティベーションを同期させ、ルーティングを複製することが可能になります。これにより、効率的な実装のためにエキスパート並列処理の場合と同様の同期プリミティブが必要になります。ただし、アテンションヘッドの数を一定のしきい値未満に減らすと収益が減少するため、このレベルでの並列処理には限界があります。

拡大

すべてのデバイスを最大限に活用し、最大128または256基のGPUという非常に高いレベルのエキスパート並列処理をサポートするには、データ並列処理 (DP) が必要です。このスキームの下では、モデルが分割されるデバイスがグループ化されてリクエストを同時に処理し、アテンションを計算し、ローカルグループ全体でシャーディングされたKVキャッシュを維持します。これらの並列グループの複数のインスタンスがエキスパート評価で協力し、それぞれが異なるエキスパート重みのサブセットをホストします。アテンションヘッドの数に基づき、グループは通常、ノード間通信が高速であるため、1つのノード全体のサイズまで拡張できます。例えば、上記の図では、1つのDPランクがリクエストAとBを独立して処理し、アテンション、ノルム、およびリクエストに関するその他の計算処理を処理します。もう一方のDPランクは個別のリクエストCを処理します。ただし、最初のノードがエキスパートの半分をホストし、もう一方のノードが残りの半分をホストしているため、ルーティング後のAとBのトークンは2番目のノードに送信される可能性があり、その逆も同様です。これはスパース通信の問題を引き起こします。各デバイスが宛先ランクごとに異なる数のトークンを送信する可能性があるためです。torchの既存のプリミティブ、主にall_to_allは、パディングやGPUとCPUの同期、メタデータのブロードキャストが必要になる可能性があるため、特に適しているとは言えません。通信を効果的に実装するには、カスタムカーネルを使用してルーティング情報を処理し、CPUの介入なしにデバイスから通信を開始する必要があります。ルーティング後、ディスパッチカーネルはルーティング先のランクにトークンを送信する必要があり、結合側では現在のDPグループ内のリクエストに属するアクティベーションを収集する必要があります。さらに、ブロードキャストを最小限に抑え、後続の層のためにすべてのランクでトークンを正しく複製するために、複数のデバイスを持つ可能性のあるDPグループ内で作業のバランスを取る必要があります。

NVSHMEM

NVSHMEMはNVIDIA固有のOpenSHMEM実装であり、基礎となるハードウェアの複雑さを抽象化するポータブルなデバイス間通信機能を提供します。APIは、個々のトランスポート層のプリミティブにマッピングされるデバイス間読み取りおよび書き込みを表現できます。NVLinkとRDMAの両方がサポートされており、移植性を確保し、NVSHMEMカーネルがノード内または複数のノードのクラスタ全体で動作できるようにします。当社のカーネルでは、データを転送して参加デバイスを同期させるために、リモートメモリ書き込みおよびシグナリング操作を使用しています。

NVSHMEM操作は対称メモリの概念を中心に構築されており、デバイス間通信に参加するすべてのデバイスに割り当てられたバッファ上で動作します。ローカルデバイスは、自身のローカルバッファ内で導出されたアドレスとターゲットランクのインデックスを併用して、操作の宛先を指定できます。下の図はこの概念を示しています。両方のGPUが同じサイズの対称バッファを割り当て、srcポインタを保持します。最初のランクは3つの整数を2番目のランクに送信し、それらをバッファの先頭に配置したいと考えています。nvshmem_int_put_nbiはローカルバッファから開始アドレスを導出し、ターゲットデバイスを指定します。2番目のランクは自身のバッファからオフセットを導出し、1つの要素を最初のデバイスに送信し、1つオフセットさせます。宛先アドレスは常にnvshmem_allocを使用して割り当てられた対称バッファである必要がありますが、ソースバッファはNVSHMEMに事前登録されていれば、デバイスメモリの任意の領域を使用できます。

NVSHMEMは幅広いプリミティブを提供していますが、当社のカーネルはわずか3つの関数に依存しており、すべての同期とフェンシングをそれらの上に構築しています。

拡大

nvshmemx_putmem_signal_nbi_warp: 1つのデバイスから別のデバイスへデータのブロックを転送し、同時にリモートデバイス上のフラグを設定します。この操作は64ビットの場所を設定(NVSHMEM_SIGNAL_SET)またはインクリメント(NVSHMEM_SIGNAL_ADD)します。フラグはメモリブロック全体が転送された後に更新されます。リモートデバイスがフラグの変更を観測した場合、自身のメモリ内のバッファに安全にアクセスできます。この関数は、データ転送と同期を合体させるのに役立ちます。

nvshmemx_signal_opは単一のメモリ場所（通常は64ビットフラグ）に対して動作し、アトミックに設定またはインクリメントします。メタデータを送信したりデバイスを同期させたりするのに役立ちます。

nvshmem_uint64_wait_untilは受信側で使用され、リモートが更新するまでフラグをポーリングします。

nvshmem_quietまたはnvshmem_fenceを通じて操作間の順序を保証することはコストがかかるため、当社のカーネル実装は他のNVSHMEMプリミティブの使用を避け、個々の送受信ペア以外の依存関係を最小限に抑えています。データを送信する際は、常に非ブロッキングバージョンの関数を使用し、ローカルランクからデータが送信されるのを待機することさえしません。カーネルグループ全体にわたる他の暗黙的な同期により、NICがアクセスする必要がある間、これらのバッファが破壊的に上書きされないことが保証されます。

NVSHMEMはGPUDirect RDMAを使用するため、NICはCPUを介さずにデータ交換のために直接GPUメモリを読み書きできます。さらに、ConnectX NICにおいて、NVSHMEMはGPU開始型通信（Infiniband GPUDirect AsyncまたはIBGDAとしても知られる）をサポートしており、これによりGPUはRDMA操作をNICに直接投稿し、NICの完了キューを直接ポーリングできます。GPU開始型通信をサポートしていないプラットフォームでは、NVSHMEMはホストCPU上で「プロキシ」スレッドを生成し、GPUの代わりに通信を開始して完了をポーリングします。NVSHMEMプログラムは、GPU開始型通信がサポートされているかどうかにかかわらずポータブルです。しかし、GPU開始型通信はCPUへの迂回を完全にバイパスするため、レイテンシを大幅に削減します。

ポータブルカーネル

当社はディスパッチカーネルと結合カーネルのペアを通じてMoE通信を実装しています。ディスパッチカーネルは各ランクでトークンとルーティング情報を読み取り、適切なエキスパートにディスパッチする役割を担います。結合カーネルはエキスパートによって生成されたアクティベーションを回収してソースランクに送り返し、ルーターによって計算された重みに基づいて選択されたエキスパートからの加重平均を計算します。カーネルは、データ転送と計算をオーバーラップさせるために、さらに送信コンポーネントと受信コンポーネントに分割されています。送信カーネルはすべて非ブロッキングかつ非同期であり、リモートへの書き込みをディスパッチするだけです。反対側では、結合カーネルはメモリから読み取るのみであり、必要なデータがすべて転送されるまで待機します。NICに作業をディスパッチした後、データがネットワーク経由で非同期に転送されている間、GPUは共有エキスパートの適用やアテンションの計算など、他の有用な作業をローカルで実行できます。

各カーネルペアは独自の対称メモリバッファを持ち、各ランクは他のすべてのランクからデータを受信するためのプライベートストレージを割り当てます。DPグループからディスパッチできるトークンの最大数はランク間で固定されており、これにより各ランクが各DPランクから各ローカルエキスパートに対して受信できる上限も設定されます。これにより、送信側ランクは、ランク間の同期を必要とせずに、書き込み先の宛先ランク上の固有アドレスを導出できます。データ受信後、カーネルは受信側カーネルが必要とする最適な形式でメモリに配置するためにトークンをシャッフルします。バッファはかなりのサイズですが、モデルのすべての連続する層で再利用されます。

拡大

上の図に示す唯一のグローバル同期形式は、結合受信カーネルに実装されています。送信側からデータがディスパッチされると、ランクは他のランクから必要なすべてのトークンを受信した時点で実行を再開し、処理され次第すぐに結合送信カーネルに入ります。結合受信のバリアにより、他のランクがまだデータ受信を待機している間にランクが先走ってディスパッチ送信を開始し、破壊的なオーバーラップを引き起こす可能性を排除します。同期は単純化のために結合カーネル内で行われますが、最適化として、ディスパッチ受信と結合送信カーネルの間で同じバリアを実装することもでき、GPUが他のカーネルを実行している間にNICが転送を完了させる時間を増やすことでレイテンシをさらに削減できます。

両方のカーネルはデバイスのすべての利用可能なSMに分割されています。ディスパッチ送信と結合受信カーネルはランクごとの最大トークン数(max_m)で並列化する必要がありますが、ディスパッチ受信と結合送信カーネルは各ランクおよびローカルエキスパートからその数を処理する必要があり、最大でmax_m * (num_experts // EP) * (EP // TP)となります。

ディスパッチ

ディスパッチカーネルは、トークンとブロックのスケーリング係数を、トークンのルーティング先であるエキスパートをホストするピアにディスパッチする役割を担います。送信側は、ルーティング情報を集約して各ランクから送信されるトークン数を決定し、リモートにディスパッチするメッセージにトークンをパッキングするという2つのタスクを並列化するために、ワープスペシャライゼーションに依存しています。受信側はまずトークン数の合計が受信されるのを待ち、次に共有メモリからペイロードを読み取り、ダウンストリームのグループ化されたGEMMカーネルに渡されるテンソルにアンパックします。また、結合カーネルと共有するバッファに、各トークンをどこに送り返すべきかを示す情報を格納します。受信カーネルが決定論的ではない順序で連続バッファ内のトークンをシャッフルするため、この情報が必要となります。

拡大

送信側では、各ブロック上の32スレッドの単一ワープがルーティング情報を読み取り、各宛先エキスパートがこのランクから受信することが予想されるトークン数をカウントし、nvshmem_signal_opを使用してカウント数プラス1を送信します。リモート側でゼロから非ゼロへの遷移はソースランクからカウントを受信したことを意味するため、カウントは1インクリメントされます。並行して、残りのワープは協力して、すべてのアクティベーション、スケーリング係数、およびローカルランク上のインデックスを連続したチャンクにパッキングし、すべてのブロック内の対称メモリへ並行してトークンをコピーします。インデックスは、結合送信側がトークンの書き込み先アドレスを決定するために必要です。次に、ワープグループ全体でバリアを介してすべてのデータがコピーされたことを確認した後、ワープは再び独立して動作し、それぞれが同じバッファを並行して異なるエキスパートに送信します。トークンはnvshmemx_putmem_signal_nbi_warpを使用して送信され、これによりリモートデバイス上のローカルランクから送信されたトークン数もアトミックに増加します。DPグループ内では、各ランクが送信されるトークンのレプリカを所有しているため、ディスパッチは均等にバランスされ、各デバイスがトークンのサブセットを送信します。

受信側では、すべてのカーネルがまず送信側ランクからトークン数の合計が送られてくるのを待機して同期します。その後、全員がアトミックにインクリメントされた送信トークン数が合計数に達するまで待機します。これは、putmem呼び出しのセマンティクスのおかげで、ソースランクからのすべてのペイロードも送信されたことを示します。カーネルは、すべてのブロックとスレッドにわたって操作を並列化し、nvshmem_uint64_wait_untilを使用してカウントをポーリングします。その後、クロスブロックバリアにより、すべてのデータが正しく受信されない限りどのブロックもバッファを読み取らないことが保証されます。ブロック全体に分散され、アトミックにインクリメントされるカウンターを介して同期されたトークンは、対称メモリから、エキスパートを実装するダウンストリームカーネルの入力となる通常の連続バッファへコピーされます。ソースランク、エキスパートインデックス、トークンインデックスは別々に格納され、結合カーネルがアクティベーションを送信すべき場所を正確に示します。DPグループ内のトークンは異なるデバイスから送信されますが、すべてがまとめられ、対応するエキスパートに渡されます。

結合

結合カーネルは再び送信側と受信側の半分に分割されています。送信側は量子化されていない16ビットのアクティベーションをリモートデバイスへコピーし、受信カーネルはデータが送られてくるのを待機して、すべてのエキスパートの貢献の加重平均をローカルで計算します。さらに、それらはディスパッチ結合シーケンスを同期させるバリアとしても機能します。各ランクは結合送信カーネルへのエントリで各ピア上のフラグを設定し、受信カーネルはフラグが設定されていることを確認しない限り戻ることができません。同期のレイテンシは最小限であり、実際の通信と計算とオーバーラップします。

拡大

送信側では、カーネルがすべてのローカルエキスパートに割り当てられたトークンのリストを並行して走査し、宛先ランク上のバッファにそれらを書き込みます。ターゲットランク、エキスパートインデックス、トークンインデックスは、トークン受信時にスキャッターカーネルによって投入されたトークンごとのバッファから読み取られます。上の図に示すように、各送信者はエキスパートごとに独自のプライベートメモリ領域を持って書き込みを行うため、同期の必要がありません。ディスパッチと同様に、結合は宛先ランク上のトークンごとのカウンターをアトミックにインクリメントしてデータの受信を示します。カウンターがトークンがディスパッチされたエキスパートの数と一致すると、トークンの内容にアクセスできるようになります。

受信カーネルでは、複数のブロックにわたってトークンのリストを並行して走査し、シグナリング操作によって設定されたフラグをポーリングすることでコンテンツの到着を待ちます。到着すると、ルーティングテーブルがどのバッファから読み取るか、どのエキスパートにどのような重みを割り当てるかを示し、プライベートバッファからペイロードが読み取られます。結果は外部で割り当てられたテンソルに書き込まれ、すべてのデバイスがバリアを通過するとカーネルは実行を終了します。DPグループ全体で、すべてのランクが各エキスパートのアクティベーションのコピーを受け取り、独自のレプリカを計算します。

ベンチマーク

当社は、InfiniBandとCX-7 NICを介して接続された16ノードにわたる128基のGPUクラスタ上でカーネルを評価します。GPUDirect Async (IBGDA) とCPUプロキシを使用した信頼性のある接続 (RC) の両方を試します。それらを、高度に最適化されたDeepSeek実装、およびPyTorch（NCCL経由）やNVSHMEMによって提供される高密度プリミティブと比較します。

拡大

高密度NVSHMEMのオールトゥオール操作では、IBRCとIBGDAの性能は同等ですが（6378µs対6180µs）、スパースカーネルではIBGDAの方が大幅に高速です。ディスパッチと結合を合計すると、IBGDAは902µsを使用するのに対し、IBRCは3223µsを要し、レイテンシが3.6倍向上します。オールトゥオールは帯域幅の制限を受けますが、スパースカーネルは桁違いに少ないデータをブロードキャストするため、レイテンシに制限されます。GPUとNICを調整するためのCPUプロキシを必要とせず、GPUから直接ネットワーク転送をトリガーすることで、エンドツーエンドのレイテンシが大幅に削減されます。

当社のポータブルカーネルは、高度に最適化されたDeepSeekカーネルより約2倍低速ですが、高密度カーネルと比較してレイテンシを10倍改善します（902µs対9944µs）。

さらに、送信側と受信側のコンポーネントへの分割により、ライブラリのプリミティブとは異なり、レイテンシの一部を隠蔽することも可能です。

拡大

シングルノード (EP8) では、NVSHMEMが輸送にNVLINKを使用し、ノード間ネットワーキングよりも低いレイテンシと高いスループットを実現します。当社のポータブルカーネルは、シングルノードでDeepEPより約2.5倍高速です（186µs対481µs）。

さらなる改善

ここで説明したカーネルは、MLフレームワークの組み込みプリミティブを上回る性能を発揮し、NVLinkやInfiniBandなどの特定のデバイス間トランスポートに過度に特化することなく、十分なパフォーマンスを提供します。すでに言及した機会に加えて、通信プリミティブをより特殊化されたバージョンに置き換えることで、さらなるパフォーマンス向上が可能です。例えば、NVLink全体では、対称メモリの使用をデバイス間で共有されるバッファに置き換えることで、一部のコピーを排除し、トークンバッチではなく個々のトークン間で、よりきめ細かい同期スキームを可能にすることができます。InfiniBand全体では、実装は基礎となるキューペアに直接アクセスできる可能性があり、予想されるトークン数と送信トークン数の両方を転送したり、キューから出力バッファにトークンが到着した時点でフラッシュしたりする必要を排除できます。しかし、そのような実装は移植性を犠牲にするというコストが伴います。一方、当社の実装では、特定のハードウェアに最適化するための開発努力を投じる前に、特定の構成を持つクラスタを評価できます。

結論

私たちは、多様なハードウェア構成間で互換性を維持しながら、標準的なAll-to-All通信と比較して10倍高速な性能を実現する、MoE通信向けの高性能かつポータブルなライブラリを提示しました。NVLinkを使用したシングルノード展開において、当社のソリューションは従来の実装より2.5倍低いレイテンシを示します。

当社のアプローチは、GPU開始型通信のサポート、計算と通信のオーバーラップを可能にするスプリットカーネルアーキテクチャ、最小限のNVSHMEMプリミティブを使用した効率的なトークンディスパッチなど、主要な技術革新を通じて、性能と移植性のバランスを取っています。マルチノードセットアップでは高度に特化された実装より約2倍低速ですが、当社のライブラリは様々なネットワーク環境（NVLink、CX-7、EFA）間で優れた柔軟性を提供します。

MoEモデルのスケールが続くにつれ、当社のような効率的な通信戦略は実用的な展開においてますます重要になります。完全にオープンソースである当社の実装は、https://github.com/ppl-ai/pplx-kernelsで入手可能です。

参考文献

DeepEP

DeepSeek-V3テクニカルレポート

NVIDIA NVSHMEM

引用

@online{perplexity2025EfficientAndPortableMixtureOfExpertsCommunication,
  author = {{Perplexity} and Nandor Licker and Lequn Chen and Phi Hoang},
  title = {{Efficient and portable mixture-of-experts communication}},
  date = {2025-04-02},
  url = {https://perplexity.coolkeypoint.com/hub/blog/efficient-and-portable-mixture-of-experts-communication}
}