LLMシステムのためのRDMAポイント・ツー・ポイント通信

新たなLLM通信パターンに対応するエレガントなツール

著者Perplexity Research

LLMシステムはこれまで、トレーニングや推論における構造化された並列パターンにおいて、集合通信に依存してきました。しかし、分散推論、Mixture-of-Experts(MoE)ルーティング、非同期強化学習のファインチューニングといった新たなワークロードでは、ポイント・ツー・ポイント通信の柔軟性が求められています。

私たちの最新の調査論文では、NVIDIA ConnectX-7やAWS Elastic Fabric Adapter(EFA)を含む、異なるネットワーク・インターフェース・コントローラー(NIC)間で動作する、移植性の高いRDMAベースの通信ライブラリ「TransferEngine」を紹介します。私たちは、Perplexityでの3つの本番環境システムにおいて、その有効性を実証しています。

分散プレフィルおよびデコード:ページド・ライトとUVMウォッチャーを使用したレイヤーごとのKvCache転送により、プレフィルインスタンスとデコードインスタンスが個別にスケーリング可能になります。

2秒未満でのRLポストトレーニング用ウェイト転送:トレーニングノードと推論ノード間での超高速ウェイト同期を実現します。

AWS EFA上で1兆パラメータモデルを実現:スキャッターおよびバリア・プリミティブを使用した低遅延のMoEディスパッチおよび結合操作により、ConnectX-7において最先端の遅延性能を達成し、AWS EFAと互換性のある初の実行可能な実装を実現しました。

本記事では、これらのシステムを支える通信ライブラリである「TransferEngine」自体に焦点を当てます。調査論文の全編はarXivで、コードはGitHubでご確認ください。

動機

集合通信は、テンソル並列やデータ並列のような構造化された並列パターンには優れていますが、新たなワークロードへの適用を制限する制約を課します。

固定メンバーシップ:集合通信では、初期化時にグローバルな通信ワールドを確立する必要があり、すべての参加者間での同期的なセットアップが求められます。分散推論の場合、トラフィックパターンに応じてプレフィルインスタンスとデコードインスタンスの数を動的にスケーリングさせる必要があるため、これが問題となります。

一様なテンソル形状:集合通信では、すべての参加者が同一の形状とデータ型のテンソルを交換することを前提としています。MoEルーティングでは、各エキスパートに送信されるトークン数はデータに依存し、実行時に決定されるため、高密度のオール・ツー・オール集合通信では最悪の場合のバッファサイズにパディングを強制されます。

操作の順序付け:集合APIは同期的なセマンティクスを強制しますが、アプリケーションが一連の操作が完了したことさえ把握できればよい場合も存在します。例えば、KvCache転送において、プレフィルインスタンスからデコードインスタンスへのすべてのページは、どのような順序で到着しても問題ありません。同様に、RLウェイト更新においても、すべてのウェイトテンソルは順序制約なしに独立して転送可能です。

RDMAベースのポイント・ツー・ポイント通信は、これらの制限を自然に解決しますが、採用を妨げてきた3つの障壁があります。

ライブラリの制限:torch.distributedやNCCLはSend/Recvのようなポイント・ツー・ポイントAPIを提供していますが、これらの操作は依然として集合通信のセマンティクスと制限を引き継いでいます。

ハードウェアの多様性:Reliable Connection(RC)トランスポートを備えたNVIDIA ConnectX NICは、最も一般的なRDMA構成であり、順序通りの配信を保証します。AWSのようなクラウドプロバイダーは異なるハードウェアを使用しており、Scalable Reliable Datagram(SRD)トランスポートを備えたElastic Fabric Adapter(EFA)はメッセージを順序不同で配信します。また、EFAで400 Gbpsの帯域幅を実現するには、4つの100 Gbps NIC(または2つの200 Gbps NIC)を集約する必要があります。

ソフトウェアの移植性:既存のRDMAライブラリは、特定のベンダーへの依存に苦しんでいます。DeepEPは、ConnectXでのみ利用可能なGPU開始型RDMA(IBGDA)を必要とします。NVSHMEMはEFA上で深刻な性能低下を示します。MooncakeやNIXLといったライブラリはEFAのサポートを欠いているか、開発の初期段階にあります。その結果、LLM推論のためのプロバイダーを横断する実行可能なソリューションは存在しません。

TransferEngine

TransferEngineは、4つの主要な機能を通じてこれらの障壁に対処します。

共通のトランスポート前提:TransferEngineは特定のNICをターゲットにするのではなく、基盤となるトランスポートが「信頼性は高いが順序は不定」であること、つまりRCトランスポートとSRDトランスポートの両方が提供する保証の共通部分のみを前提としています。

ImmCounter同期:メッセージの順序付けに依存せずに動作する、完了通知のための新しいプリミティブであり、順序付き(RC)トランスポートと順序不定(SRD)トランスポートの両方で同じコードを実行可能にします。

マルチNICサポート:GPUごとの複数のNICを透過的に管理し、帯域幅を最大化するために転送を自動的にシャーディングします。

最小限のAPIサーフェス:複雑な通信パターンを構築するために十分な、少数の操作セットです。

TransferEngine システム設計。リクエストとコールバックの経路、CPUワーカー、UVMバッファー、GPU、NICが、2つのGPU/NUMAドメインにまたがって示されています。

TransferEngine は GPU ごとに 1 つのワーカースレッドを生成します。各ワーカーは、ハードウェアプラットフォームに応じて 1 〜 4 個の RDMA NIC を調整する DomainGroup を管理します(単一の ConnectX-7 は 400 Gbps を提供し、EFA では 100 Gbps アダプター 4 枚または 200 Gbps アダプター 2 枚が必要です)。DomainGroup 内では、各 Domain はハードウェア固有で、単一の NIC を管理し、キューペア管理、ワーク送信、完了ポーリングを処理します。

各 TransferEngine インスタンスは、ピアの識別と検出のために単一の main address を公開します。リモートピアは、ネットワークアドレスをエンコードし、帯域外通信のためにシリアライズ可能な NetAddr 構造体を交換します。システム内のすべてのピアは、GPU ごとに同じ数の NIC を使用しなければなりません。この統一構成により、TransferEngine は転送をインテリジェントにシャーディングできます。転送を開始するとき、送信元と送信先の両方が利用可能な NIC を正確に把握しているため、すべてのアダプターにまたがる自動ロードバランシングが可能になります。

API設計

TransferEngineは、基盤となるRDMAインターフェースの複雑さを抽象化する最小限のAPIを公開しています。

rust
#[serde] struct NetAddr(Bytes); #[serde] struct MrDesc{ ptr: u64, rkeys: Vec<(NetAddr, u64)> } struct MrHandle(NonNull<c_void>); type Offset = u64; struct Pages{ indices: Vec<u32>, stride: u64, offset: Offset } struct PeerGroupHandle(u64); struct ScatterDst{ len: u64, src: Offset, dst: (MrDesc,Offset)} enum OnDone { Callback(fn () -> ()), Flag(Atomic<bool>) } trait TransferEngine { fn main_address() -> NetAddr; // Memory Region Management fn reg_mr(ptr, len, device) -> (MrHandle, MrDesc); // Two-sided Send/Recv fn submit_send(addr: NetAddr, msg: &[u8], cb: fn () -> ()); fn submit_recvs(len: u64, cnt: u64, cb: fn (&[u8]) -> ()); // One-sided Write fn expect_imm_count(imm: u32, count: u32, cb: fn () -> ()); fn submit_single_write(len: u64, imm: Option<u32> src: (MrHandle, Offset), dst: (MrDesc, Offset), OnDone); fn submit_paged_writes(page_len: u64, imm: Option<u32>, src: (MrHandle, Pages), dst: (MrDesc, Pages), OnDone); // One-sided Write to a group of peers fn add_peer_group(addrs: Vec<NetAddr>) -> PeerGroupHandle; fn submit_scatter(h: Option<PeerGroupHandle>, OnDone, imm: Option<u32>, src: MrHandle, dst: Vec<ScatterDst>); fn submit_barrier(h: Option<PeerGroupHandle>, OnDone, imm: u32, dst: Vec<MrDesc>); // Watcher for CPU-GPU synchronization fn alloc_uvm_watcher(cb: fn(u64,u64) -> ()) -> NonNull<u64>; }

このAPIは、いくつかの主要なコンポーネントで構成されています。

メモリ登録

RDMA 操作では、データ転送前にメモリ領域を NIC に登録する必要があります。reg_mr 関数は、ポインタ、長さ、デバイスを受け取り、2 つのハンドルを返します:

MrHandle:転送のソースとして使用されるローカルハンドル。

MrDesc:リモートピアと交換することで、彼らが登録されたメモリに書き込めるようにするシリアル化可能な記述子。

MrDesc 構造体には、メモリ領域へのリモートアクセスを許可する、リモートネットワークアドレスとリモートキーのリスト(NIC ごとに 1 つ)が含まれます。この設計は、マルチ NIC 構成を自然に扱えます。

双方向Send/Recv

小さなメッセージや RPC スタイルの制御プレーンメッセージ向けに、TransferEngine は submit_send と submit_recvs を提供します:

submit_send: NetAddr によって識別されるリモートピアにメッセージを送信します。メッセージは内部でコピーされるため、呼び出し元はすぐにバッファを再利用または解放できます。

submit_recvs:受信バッファのプールを投稿します。メッセージが到着すると、受信したデータと共にコールバックが呼び出されます。コールバック完了後、バッファは自動的にプールに返されます。

片方向Write

バルクデータ転送のために、TransferEngineはゼロコピーの片方向書き込み操作を提供します。

submit_single_write:ソースから宛先へ連続した領域を転送します。

submit_paged_writes:インデックス、ストライド、オフセットで指定された非連続のページを転送します。

両方の操作は、完了時に受信側へアトミックに配信されるオプションの32ビット即値を受け取ります。

このエンジンは、帯域幅を最大化するために、大きな転送を複数のNIC間で自動的にシャーディングします。

ピアグループとバルク操作

(MoEルーティングで一般的な)複数のピアをターゲットとする操作のために、TransferEngineは最適化されたプリミティブを提供します。

add_peer_group:高速検索のためにピアアドレスのセットを事前登録します。

submit_scatter:グループ内の各ピアに異なるデータスライスを送信します。

submit_barrier:すべてのピアに即値のみの通知を送信します(ゼロバイト転送)。

これらの操作は、内部データ構造をキャッシュし、RDMAワークリクエストをテンプレート化することで、反復的なセットアップのオーバーヘッドを回避し、低遅延に最適化されています。

UVMウォッチャー:GPUからCPUへのシグナリング

LLMシステムにおける重要な課題は、RDMA転送とGPUの進行を調整することです。理想的には、GPUが結果を生成するとすぐにデータ転送が開始されるべきですが、CPUからGPUをポーリングするのは非効率的です。

TransferEngine は、これを UVM Watchers で解決します。UVM Watchers は、GPU カーネルによって更新可能で(CUDA Graph 内でも)、GDRCopy を使って CPU が効率的にポーリングできる統合仮想メモリ上の場所です。UVM watcher が割り当てられると、GPU が値を更新するたびに呼び出されるコールバックが登録されます。

これにより、レイヤーごとのKvCache転送のようなパターンが可能になります。各トランスフォーマーレイヤーが完了すると、GPUはウォッチャーをインクリメントし、後続のレイヤーが計算を続けている間にそのレイヤーのKvCacheのRDMA転送をトリガーします。

完了通知

TransferEngine は、開始側(submit_write の呼び出し元)と受信側(Write 操作のターゲット)の両方に対して、完了通知機構を提供します。

イニシエーター側の通知:書き込み操作が送信されると、呼び出し元はコールバック関数またはアトミックフラグのいずれかを通じて完了通知を受けることができます。この通知は、データがソースバッファから正常に転送されたことを確認します。

受信側通知: 重要な革新は ImmCounter であり、32 ビットの immediate 番号で索引付けされたカウンタを追跡するコンポーネントです。Write 操作が immediate 値を伴う場合、データ到着時に受信側の ImmCounter が原子的にインクリメントされます。expect_imm_count 関数により、受信側は特定の immediate 値が目標カウントに達したときに発火するコールバックを登録できます。カウンタは GDRCopy を介して GPU と透過的に同期でき、直接ポーリングしたり、コールバックを通じて監視したりできます。

このレシーバー側のメカニズムは、メッセージの順序付けを必要とせずに信頼性の高い同期を提供します。ConnectX NIC(RCトランスポートでの順序通りの配信を保証)とEFA NIC(SRDトランスポートでの順序不定のメッセージ配信)の両方で等しく機能します。

実装

TransferEngine は Rust で実装されており、レイテンシを最小化し高スループットを実現するために、アロケーション、スレッド処理、同期を慎重に最適化しています。このエンジンは DomainGroup ごとに 1 つのワーカースレッドを生成し、それぞれをデバイスが接続された NUMA ノード上の CPU コアにピン留めすることで、スケジューリング遅延とメモリアクセス遅延の両方を最小化します。ドメイン固有のデータ構造は、ピン留め後に割り当てられ、メモリが正しい NUMA ノードに確保されるようにします。1 つのワーカースレッドが最大 4 つの Domain(それぞれ単一の NIC を管理)を処理し、別の専用スレッドが GPU をポーリングして UVM watcher を更新します。スレッド間通信はロックフリーキューを介して行われます。

API は要求を適切なデバイスを担当する DomainGroup に転送し、先頭のものはホストも担当します。ループを短く回しながら、ドメインワーカーは新しい作業をポーリングし、新規要求の投入を優先します。ハードウェアと設定に応じて、要求は利用可能な Domain 間でシャーディングされ、ロードバランスされます。複合要求の最初の Write は直ちに NIC の送信キューに投入されます。新規要求が尽きると、ワーカーは保留中の操作を進める処理に移り、ハードウェアのパイプラインを埋めるように Write を投入します。その後、完了キューがポーリングされ、終了した転送と immediate カウンタ増分を照会します。イベントは集約され、転送ごとの通知を配信し、転送をすべてのグループで共有される専用のコールバックスレッドに引き渡します。

DomainGroup 内でのシャーディングは柔軟です。転送はインデックスによって特定の NIC を対象にできます。単一の Write は分割可能です。ページ化転送、scatter、および barrier 操作は、いずれも複数の Write に変換され、すべての NIC にまたがってシャーディングできます。

ハードウェア固有の最適化

TransferEngine 内の Domain は、それらが制御するハードウェア向けに特化し、最適化されています:

AWS EFA: EFA 実装は libfabric を使用し、DomainGroup 内で NIC ごとに 1 つの fabric domain を管理します。主な実装判断と最適化には次のものがあります:

ゼロサイズの書き込み:EFAは、ゼロサイズの書き込みであっても有効なターゲット記述子を必要とするため、RDMA仕様から逸脱しています。そのため、TransferEngineはすべての転送に対して有効な記述子を強制します。

ワークリクエストのテンプレート化: バルク転送やピアグループ操作のために、投稿前に libfabric 構造体の共通フィールドを事前設定して保持し、操作ごとのオーバーヘッドを削減します

NVIDIA ConnectX-7: ConnectX-7 実装は RC トランスポートとともに libibverbs を使用します。主な実装判断と最適化には次のものがあります:

ハンドシェイク用の UD キューペア: 各 Domain は、RC 接続ハンドシェイクを交換するために Unreliable Datagram (UD) キューペアを使用します。

個別のRCキューペア:ピアごとに2つのRCキューペアを作成し、1つは双方向のSend/Recv操作用、もう1つは片方向のWriteおよびWriteImm操作用に使用します。RecvとWriteImmの完了は投稿順にワークリクエストを消費するため、この分離が必要です。これにより、WriteImmを干渉なしにサポートしつつ、高レベルのRecvセマンティクスが可能になります。

ワークリクエストのテンプレート化:バルク転送やピアグループ用のワークリクエスト構造体の共通フィールドを事前入力します。

ワークリクエストのチェーニング: next ポインタを介して最大 4 つのワークリクエストを ibv_send_wr にリンクし、NIC へのドアベルリング回数を削減します。

Relaxed ordering: IBV_ACCESS_RELAXED_ORDERING を有効にして、NIC と GPU メモリ間の PCIe トランザクションの順序を前後させることを許可します。

結論

異なるRDMAトランスポート間での保証の共通部分をターゲットとし、ImmCounterを介した順序非依存の同期を導入することで、TransferEngineはLLMシステムにおけるポイント・ツー・ポイント通信のための移植性の高い基盤を提供します。このライブラリはオープンソースでありGitHubで公開されています。コミュニティが大規模機械学習システムにおける新たな通信パターンに向けて、これを活用・構築していくことを推奨します。

引用

BibTeX
@online{perplexity2025RdmaPointToPointCommunicationForLlmSystems,
  author = {{Perplexity} and Nandor Licker and Kevin Hu and Vladimir Zaytsev and Lequn Chen},
  title = {{RDMA point-to-point communication for LLM systems}},
  date = {2025-11-05},
  url = {https://www.perplexity.ai/hub/blog/rdma-point-to-point-communication-for-llm-systems}
}