Perplexity、MoE通信ライブラリをオープンソース化：最大10倍の高速化を実現

pplx-kernelsは、Perplexityがオープンソースとして公開したMoE（Mixture-of-Experts）向け高速通信ライブラリです。大規模なMoEモデルを複数GPUで並列推論する際のボトルネックを解消し、標準的なAll-to-All通信と比べて最大10倍の転送速度を実現します。

DeepSeek R1やMixtralのようなMoEアーキテクチャでは、計算リソースそのものよりもGPU間のデータ転送待ちがレイテンシの主な原因になります。例えばDeepSeek R1は総パラメータ数が671Bに達するものの、1トークンあたりで実際に使われるのは37Bのみです。しかし、選ばれたエキスパートが別々のGPUに分散配置されているため、デバイス間の頻繁な通信が避けられません。

何ができるのか

従来のPyTorchによる分散処理では、トークンのルーティングに標準のall_to_allプリミティブを使うのが一般的でした。しかしMoE特有のスパースな通信パターンでは、パディングの発生やGPUとCPUの頻繁な同期、メタデータのブロードキャストによって無駄なオーバーヘッドが生じていました。

pplx-kernelsは、エキスパートへのトークン割り振り（ディスパッチ）と結果の重み付け集計（コンバイン）を極限まで効率化します。

単一ノード構成では、従来最速とされていた実装と比べてレイテンシを2.5倍削減しました。複数ノードにまたがるクラスタ環境でも、標準のAll-to-All比で最大10倍の通信速度を記録しています。

仕組みと技術的工夫

高速化の肝は、CPUを完全にバイパスすることと、通信と計算のオーバーラップにあります。

ConnectXネットワークカード環境では、GPU主導の通信（IBGDA）を利用します。GPUがホストCPUを介さずにNICへ直接RDMA処理を投入し、完了キューを監視するため、CPU介在の遅延を丸ごと削ぎ落とせます。

また、送信と受信のカーネルを分離するスプリットカーネル設計を採用しています。ネットワーク転送が走っている間にGPU本体で別の計算を同時進行させ、待ち時間を隠蔽します。

さらにNVSHMEMの利用方法も工夫されています。重いメモリフェンスや同期関数を排除し、わずか3つの非ブロッキング基本関数のみで同期と転送を完結させることで、最小限の依存関係を実現しました。

互換性と利用方法

特定のハードウェア専用に過度な最適化を施した手法と比べると約2倍の速度差がありますが、pplx-kernelsは移植性の高さを重視して作られています。

1ノード内の高速なNVLinkだけでなく、ノード間のInfiniBand、ConnectX-7、さらにはAWS EFA環境でもそのまま動作します。

コードはGitHub上のppl-ai/pplx-kernelsリポジトリで完全に公開されており、誰でも入手してビルドできます。

なぜ重要か

自前のクラスタでDeepSeek R1などの巨大MoEモデルを動かそうとしているチームにとって、通信レイテンシの短縮は最優先の課題です。

pplx-kernelsは、特定の閉じた専用環境に縛られることなく、手持ちの標準的なGPUクラスタ上で推論遅延を劇的に削れる現実的な選択肢を提供してくれます。既存の非効率な通信処理を置き換える価値は十分にあります。