適用於 LLM 系統的 RDMA 點對點通訊

解決新興 LLM 通訊模式的優雅工具

作者Perplexity Research

LLM 系統傳統上依賴集體通訊來實現訓練和推論中的結構化平行模式。然而,新興的工作負載,例如解構式推論、專家混合 (MoE) 路由以及非同步強化學習微調,則需要點對點通訊的靈活性。

我們最新的研究論文介紹了 TransferEngine,這是一個基於 RDMA 的可攜式通訊程式庫,適用於不同的網路介面控制器 (NIC),包括 NVIDIA ConnectX-7 和 AWS Elastic Fabric Adapter (EFA)。我們在 Perplexity 的三個生產系統中展示了其有效性:

解構式 Prefill 與 Decode:使用分頁寫入和 UVM 監控器的逐層 KvCache 傳輸,使 Prefill 與 Decode 執行個體能夠獨立擴展。

不到 2 秒的 RL 後訓練權重傳輸:訓練與推論節點之間極速的權重同步。

在 AWS EFA 上啟用兆參數模型:使用散佈 (scatter) 與屏障 (barrier) 原語實現低延遲 MoE 調度與合併操作,在 ConnectX-7 上達到領先業界的延遲效能,並成為第一個與 AWS EFA 相容的可行實作。

本文聚焦於支撐這些系統的通訊程式庫 TransferEngine 本身。請在 arXiv 閱讀完整研究論文,並在 GitHub 探索程式碼。

動機

集體通訊在結構化平行模式(如張量平行與資料平行)方面表現優異,但它們所施加的限制限制了其對新興工作負載的適用性:

固定成員資格:集體通訊要求在初始化時建立全域通訊世界,並在所有參與者之間進行同步設定。對於解構式推論而言,這是個問題,因為 Prefill 與 Decode 執行個體的數量必須根據流量模式動態擴展。

統一張量形狀:集體通訊假設所有參與者交換相同形狀與資料類型的張量。在 MoE 路由中,發送給每個專家的 Token 數量取決於資料,且於執行階段發現,這迫使密集的全對全 (all-to-all) 集體通訊必須填補至最壞情況的緩衝區大小。

操作順序:集體通訊 API 強制執行同步語意,即使應用程式僅需知悉一組操作已完成亦然。例如,在 KvCache 傳輸中,來自 Prefill 執行個體傳送到 Decode 執行個體的所有頁面可以任何順序抵達。同樣地,在 RL 權重更新中,所有權重張量均可獨立傳輸,無須順序限制。

基於 RDMA 的點對點通訊自然能解決這些限制,但有三個障礙阻礙了其採用:

程式庫限制:雖然 torch.distributed 和 NCCL 提供了如 Send/Recv 等點對點 API,但這些操作仍繼承了集體通訊的語意與限制。

硬體多樣性:採用可靠連線 (RC) 傳輸的 NVIDIA ConnectX NIC 是最常見的 RDMA 設定,可保證順序傳遞。雲端供應商如 AWS 使用不同的硬體,具備可擴展可靠資料報 (SRD) 傳輸的 Elastic Fabric Adapter (EFA) 則以亂序傳遞訊息。EFA 同時需要聚合四個 100 Gbps NIC(或兩個 200 Gbps NIC)才能達到 400 Gbps 的頻寬。

軟體可攜性:現有的 RDMA 程式庫深受供應商鎖定所苦。DeepEP 需要僅在 ConnectX 上提供的 GPU 初始化 RDMA (IBGDA)。NVSHMEM 在 EFA 上表現出嚴重的效能下降。諸如 Mooncake 和 NIXL 等程式庫則缺乏 EFA 支援或仍處於開發早期階段。因此,目前尚無適用於 LLM 推論的跨供應商可行解決方案。

TransferEngine

TransferEngine 透過四項核心功能解決了這些障礙:

通用傳輸假設:TransferEngine 不針對特定 NIC,僅假設底層傳輸是可靠但無序的,亦即 RC 與 SRD 傳輸所提供保證的交集。

ImmCounter 同步:一種用於完成通知的新穎原語,無須依賴訊息順序,使得相同的程式碼能同時在有序 (RC) 與無序 (SRD) 的傳輸上執行。

多 NIC 支援:透明管理每個 GPU 的多個 NIC,自動分片傳輸以最大化頻寬。

最小化 API 介面:一套足以建構複雜通訊模式的小型操作集。

TransferEngine 系統設計,顯示請求與回呼路徑、CPU 工作執行緒、UVM 緩衝區、GPU 和 NIC,分佈於兩個 GPU/NUMA 網域之間。

TransferEngine 為每個 GPU 產生一個討論串。每個工作者管理一個 DomainGroup,根據硬體平台協調 1 到 4 個 RDMA NIC(單個 ConnectX-7 提供 400 Gbps,而 EFA 需要四個 100 Gbps 轉接器或兩個 200 Gbps 轉接器)。在 DomainGroup 內,每個 Domain 均為硬體專屬,並管理單個 NIC,處理佇列對管理、工作提交與完成輪詢。

每個 TransferEngine 執行個體公開一個用於對等識別與發現的 main address。遠端對等點交換編碼網路位址的 NetAddr 結構,並可針對頻外通訊進行序列化。系統中的所有對等點必須針對每個 GPU 使用相同數量的 NIC。這種統一的配置允許 TransferEngine 智慧地分片傳輸:當啟動傳輸時,來源與目的地均確切知悉哪些 NIC 可用,從而在所有轉接器之間實現自動負載平衡。

API 設計

TransferEngine 公開了一個最小化的 API,抽象化了底層 RDMA 介面的複雜性。

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 函數接收指標、長度與裝置,並回傳兩個控制代碼:

MrHandle:用作傳輸來源的本地控制代碼

MrDesc:可與遠端對等點交換的序列化描述符,使其能夠寫入註冊的記憶體

MrDesc 結構包含遠端網路位址與遠端金鑰清單(每個 NIC 一個),授與對該記憶體區域的遠端存取權。此設計自然地處理了多 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 Watcher:GPU 到 CPU 的信號傳輸

LLM 系統中的一項關鍵挑戰是協調 RDMA 傳輸與 GPU 進度。理想情況下,資料傳輸應在 GPU 產生輸出時立即開始,但從 CPU 輪詢 GPU 的效率並不高。

TransferEngine 透過 UVM Watchers 解決此問題:統一虛擬記憶體位置可由 GPU 核心(即使在 CUDA 圖中)更新,並由 CPU 使用 GDRCopy 有效輪詢。分配 UVM 監控器時,會註冊一個回呼,每當 GPU 更新該值時就會觸發。

這實現了如逐層 KvCache 傳輸的模式:在每個 transformer 層完成後,GPU 會遞增監控器,觸發該層 KvCache 的 RDMA 傳輸,同時後續層繼續計算。

完成通知

TransferEngine 為發起端(submit_write 的呼叫者)與接收端(Write 操作的目標)均提供了完成通知機制。

發起端通知:提交寫入操作時,呼叫者可選擇透過回呼函數或原子旗標在完成時獲得通知。此通知確認資料已成功從來源緩衝區傳輸。

接收端通知:關鍵創新在於 ImmCounter,這是一個追蹤由 32 位元立即數索引之計數器的組件。當寫入操作攜帶立即值時,接收者的 ImmCounter 會在資料送達時原子地遞增。expect_imm_count 函數允許接收者註冊一個回呼,該回呼會在特定立即值達到目標計數時觸發。計數器可透過 GDRCopy 與 GPU 透明同步、直接輪詢或透過回呼監控。

此接收端機制提供了可靠的同步,而無須訊息順序。它同樣適用於 ConnectX NIC(保證 RC 傳輸上的順序傳遞)與 EFA NIC(在 SRD 傳輸上亂序傳遞訊息)。

實作

TransferEngine 以 Rust 實作,審慎地最佳化了記憶體配置、多執行緒與同步,以最小化延遲並實現高吞吐量。該引擎為每個 DomainGroup 產生一個工作者討論串,每個討論串被固定 (pinned) 在裝置所連接的 NUMA 節點的 CPU 核心上,以最小化排程與記憶體存取延遲。特定於領域的資料結構會在固定後進行配置,以確保記憶體是在正確的 NUMA 節點中預留。一個工作者討論串處理最多 4 個 Domain(每個管理單個 NIC),而另一個專用討論串則輪詢 GPU 以更新 UVM 監控器。跨討論串通訊透過無鎖佇列完成。

API 將請求轉發至服務對應裝置的 DomainGroup,第一個 DomainGroup 同時服務主機。在緊湊迴圈中,領域工作者會輪詢新工作,並優先提交新請求。根據硬體與配置,請求會被分片並在可用的 Domain 之間進行負載平衡。複合請求的第一個 Write 會立即發佈至 NIC 的傳送佇列。一旦新請求耗盡,工作者將進行待處理的操作,發佈寫入以填滿硬體管線。隨後,輪詢完成佇列以查詢完成的傳輸與立即數計數器遞增。事件會被聚合以傳遞逐傳輸的通知,並將傳輸移交給所有群組共用的專用回呼討論串。

DomainGroup 內的分片具備靈活性。傳輸可依索引針對特定 NIC。單個 Write 可被分割。分頁傳輸、散佈與屏障操作(全部轉換為多個 Writes)均可跨所有 NIC 進行分片。

硬體特定最佳化

TransferEngine 內的 Domain 均針對其控制下的硬體進行了專門化與最佳化:

AWS EFA:EFA 實作使用 libfabric,在 DomainGroup 內為每個 NIC 管理一個 fabric 領域。關鍵實作決策與最佳化包括:

零大小寫入:EFA 背離了 RDMA 規範,即使對於零大小寫入也要求有效的目標描述符。因此,TransferEngine 強制所有傳輸必須具備有效的描述符。

工作請求模板化:在發佈前預先填入並保留 libfabric 結構的通用欄位,減少批次傳輸與對等群組操作的單操作開銷。

NVIDIA ConnectX-7:ConnectX-7 實作使用具備 RC 傳輸的 libibverbs。關鍵實作決策與最佳化包括:

用於交握的 UD 佇列對:每個 Domain 使用一個不可靠資料報 (UD) 佇列對來交換 RC 連線交握。

分離的 RC 佇列對:每個對等點建立 2 個 RC 佇列對,一個用於雙向 Send/Recv 操作,另一個用於單向 Write 與 WriteImm 操作。此分離是必要的,因為 Recv 與 WriteImm 完成皆以發佈順序消耗工作請求,從而在不干擾的情況下支援 WriteImm 並維持高階 Recv 語意。

工作請求模板化:為批次傳輸與對等群組預先填入工作請求結構的通用欄位。

工作請求鏈接:透過 ibv_send_wr 的 next 指標鏈接最多 4 個工作請求,減少對 NIC 的門鈴 (doorbell) 觸發次數。

放寬排序:啟用 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}
}