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

優雅的工具,用以處理新興的 LLM 通訊模式

作者Perplexity Research

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

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

分散式 Prefill 和 Decode:使用分頁寫入和 UVM 監控器的逐層 KvCache 傳輸,允許 prefill 和 decode 實例獨立擴展。

在 2 秒內完成 RL 後訓練的權重傳輸:訓練和推論節點之間極速的權重同步。

在 AWS EFA 上實現兆級參數模型:使用分散和障礙基元進行低延遲的 MoE 調度與合併操作,在 ConnectX-7 上實現最先進的延遲,並成為第一個與 AWS EFA 相容的可行實作。

本文專注於 TransferEngine 本身,這個支援這些系統的通訊庫。在 arXiv 上閱讀完整的研究論文,並在 GitHub 上探索程式碼。

動機

集體通訊在結構化平行模式(例如張量平行和資料平行)方面表現出色,但它們施加了限制其在應對新興工作負載時適用性的條件:

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

統一的張量形狀:集體通訊假設所有參與者交換形狀和資料類型相同的張量。在 MoE 路由中,發送到每個專家的權杖數量取決於資料並在執行階段發現,這迫使密集的全對全集體通訊填充到最壞情況下的緩衝區大小。

操作排序:即使應用程式只需要知道一組操作已經完成,集體通訊 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 啟動一個 worker 執行緒。每個 worker 會管理一個 DomainGroup,該 DomainGroup 會根據硬件平台協調 1 至 4 個 RDMA NIC(單一 ConnectX-7 可提供 400 Gbps,而 EFA 則需要四個 100 Gbps 介面卡或兩個 200 Gbps 介面卡)。在 DomainGroup 內,每個 Domain 都是特定於硬件的,並管理單一 NIC,負責 queue pair 管理、工作提交及完成輪詢。

每個 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:發布接收緩衝區池。當訊息到達時,會使用接收到的資料叫用回呼。回呼完成後,緩衝區會自動歸還給池。

單向寫入

對於大量資料傳輸,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 監控器解決了這個問題:統一虛擬記憶體位置,可由 GPU 核心更新(甚至在 CUDA 圖形內),並可由 CPU 使用 GDRCopy 有效輪詢。分配 UVM 監控器後,會註冊一個回呼,每當 GPU 更新該值時就會被叫用。

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

完成通知

TransferEngine 為發起端(submit_write 的呼叫者)和接收端(寫入操作的目標)提供完成通知機制。

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

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

這種接收端機制提供了可靠的同步,而不需要訊息排序。它在 ConnectX NIC(保證透過 RC 傳輸進行有序交付)和 EFA NIC(透過 SRD 傳輸亂序交付訊息)上同樣運作良好。

實施

TransferEngine 是用 Rust 實作的,仔細最佳化了分配、執行緒和同步,以最小化延遲並實現高吞吐量。引擎為每個 DomainGroup 衍生一個工作執行緒,每個執行緒都釘選到裝置所連接的 NUMA 節點上的 CPU 核心,從而最小化排程和記憶體存取延遲。特定於域的資料結構在釘選後分配,以確保在正確的 NUMA 節點中保留記憶體。一個工作執行緒最多處理 4 個 Domain(每個管理一個單獨的 NIC),而另一個專用執行緒則輪詢 GPU 以更新 UVM 監控器。跨執行緒通訊是透過無鎖佇列完成的。

API 將請求轉發到服務適當裝置的 DomainGroup,其中第一個也服務主機。在緊湊迴圈中,域工作執行緒輪詢新工作,優先處理新請求的提交。根據硬體和配置,請求會在可用的 Domain 之間進行分片和負載平衡。複合請求的第一個寫入會立即發布到 NIC 的傳送佇列。一旦新的請求用盡,工作執行緒就會繼續推進掛起的操作,發布寫入以填滿硬體管線。隨後,輪詢完成佇列以查詢已完成的傳輸和立即計數器增量。事件被聚合以提供每個傳輸的通知,將傳輸移交給所有群組共用的專用回呼執行緒。

DomainGroup 內部分片是很靈活的。傳輸可以按索引針對特定的 NIC。單個寫入可以分割。分頁傳輸、分散和障礙操作(這些操作全部轉化為多個寫入)可以在所有 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 完成都會按發布順序消耗工作請求,從而允許高級別的 Recv 語義,同時支援 WriteImm 而不受干擾。

工作請求範本化:為大量傳輸和對等群組預先填入工作請求結構的通用欄位。

工作請求鏈結:透過 ibv_send_wr 的 next 指標連結最多 4 個工作請求,減少對 NIC 的門鈴響鈴次數。

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