高效且可攜的專家混合通訊

我們推出了一款高效能、可攜式且開源的專家混合（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 環境中進行高效 Token 分派與組合的方法。

簡介

專家混合（MoE）模型（如 DeepSeek R1 和 Mixtral 8x7B）透過限制每個 Token 激活的權重數量，優於稠密模型。MoE 模型並非透過增加層數或擴大每個解碼器層中多層感知器（MLP）的線性層大小來實現高參數計數，而是將傳統的 MLP 取代為多個專家與一個路由器。例如，在 DeepSeek R1 的 6710 億參數中，推理期間僅有 370 億參數會與給定的 Token 相乘。這大幅降低了解碼 Token 所需的運算量，並相較於稠密模型改善了延遲。

與稠密模型相比，MoE 模型在推理方面帶來了一些額外的挑戰。雖然專家本身是小型 MLP 層，但每個解碼器層都包含一個路由器，用以決定將 Token 分派給哪些專家，每個 Token 會被分派給多個專家。路由器通常是一個產生機率分佈的小型線性層。通常會選擇 Top-K 分數的專家，最終的激活值則計算為加權平均值，即將專家輸出乘以源自路由器機率分佈的權重後求和。

為了將分散式 MoE 系統中的延遲降至最低，可以利用跨多個裝置的並行處理，但這會引入通訊挑戰。諸如 Mixtral 8x7B 或 Llama-70B 等模型適合於單一節點上的 8 個裝置內，受益於快速且低延遲的 NVLink 互連（高達 900Gbps）。然而，更大的模型需要使用 InfiniBand（峰值為 400Gbps）將專家分片（Sharding）至多個節點，這會帶來額外的延遲挑戰。

放大

在本文中，我們探討了專家並行處理的問題，並描述了我們實現 GPU 核心的過程，這些核心能高效地將 Token 分派至持有分配專家之裝置，並透過計算專家輸出的加權總和將其收回。雖然存在針對特定網路硬體進行激進優化的實作，但我們提出了一種高度可攜的實作，它依賴最少量的 NVSHMEM 原語，卻能提供比標準 All-to-All 通訊快 10 倍的效能。我們的實作在單節點配置上實現了頂尖的效能，同時在各種網路環境中保持了優異的可攜性。

分片與並行處理

為了實現高效推理，模型權重必須保留在裝置記憶體中，同時須留出足夠的空間給前向傳遞模型所需的激活值、KV 快取及其他緩衝區。功能最強大的模型超出了即使是最強大 GPU 的容量，因此推理必須分散至多個裝置，這些裝置可共同儲存模型權重。基於權重的分片方案，必須使用不同的通訊與運算方案來卸載運算並同步裝置。

放大

專家並行處理（EP）如第一張圖所示，僅對專家運算進行並行化。不同的專家被指派至持有其權重的不同裝置。路由後，Token 會被傳送至對應的裝置，隨後收集並累加結果。路由的複雜度取決於模型其他部分的並行程度：複製其他層可能會完全消除路由的需求，因為每個 Rank 都可以從本地複製的路由表中選擇 Token。然而，如果僅有一個 Rank 執行路由，則需要廣播（Broadcast）以將 Token、索引與權重分派至各自的專家。最後，全體收集（All-Gather）或全體廣播（All-to-All）會將輸出 Token 與繼續執行模型的任何 Rank 進行同步。此類實作相對簡單，因為 PyTorch 已經公開了所需的原語，儘管透過將累加與從專家處收集 Token 的收集操作進行融合，可獲得一些優勢。

放大

專家專用並行處理無法理想地擴展，因為當模型執行非專家層（如注意力、標準化與採樣）時，叢集中的節點可能會處於閒置狀態。然而，這些層（主要是注意力）的運算也可受益於張量並行處理（TP）。大多數模型依賴多頭注意力，這意味著注意力頭及其對應的 Q、K 與 V 投影也可以跨裝置進行分片，並在不同層之間複製或收集切片。如果注意力分散至所有裝置，全體收集可以同步激活值，從而允許路由被複製，這需要類似於專家並行處理情況的同步原語才能實現高效實作。然而，在此層級的並行處理有限制，因為將注意力頭的數量減少至特定閾值以下將產生邊際遞減效應。

放大

為了最佳化所有裝置的使用並支援高達 128 或 256 個 GPU 的高度專家並行處理，需要資料並行處理（DP）。在此方案下，模型分佈的裝置會分組以同時處理請求，計算注意力並維護跨其本地組分片的 KV 快取。這些並行組的多個執行個體協同進行專家評估，每個執行個體託管不同子集的專家權重。基於注意力頭的數量，一個組通常可擴展至整個節點的大小，因為節點內通訊較快。例如，在上圖中，一個 DP Rank 獨立服務兩個請求 A 與 B，處理請求的注意力、標準化與任何其他簿記工作。另一個 DP Rank 則處理不同的請求 C。然而，第一個節點託管了一半的專家，而另一個節點託管了另一半，因此路由後，來自 A 與 B 的 Token 可能會被傳送至第二個節點，反之亦然。這導致了一個稀疏通訊問題：每個裝置可能向任何其他目的地 Rank 傳送不同數量的 Token。現有的 PyTorch 原語（主要是 all_to_all）並非特別適用，因為它們可能需要某種形式的填充（Padding）、GPU 至 CPU 的同步與元資料廣播。為了有效地實作通訊，需要自訂核心來處理路由資訊，並在無 CPU 干預的情況下從裝置發起通訊。路由後，分派核心必須將 Token 傳送至它們所路由到的 Rank，而在組合端，必須收集屬於當前 DP 組中請求的激活值。此外，必須在可能擁有多個裝置的 DP 組內平衡工作量，以最小化廣播，並正確地將 Token 複製到所有 Rank 上以供後續層使用。

NVSHMEM

NVSHMEM 是一種 NVIDIA 特有的 OpenSHMEM 實作，提供可攜式的裝置間通訊設施，抽象化了底層硬體的複雜性。該 API 可以表達裝置對裝置的讀寫，這些讀寫被映射至個別傳輸層的原語。NVLink 與 RDMA 均受支援，賦予了一定程度的可攜性，並允許 NVSHMEM 核心在節點內或跨多個節點的叢集中運作。在我們的核心中，我們使用遠端記憶體寫入與訊號操作來傳輸資料並同步參與的裝置。

NVSHMEM 操作圍繞對稱記憶體的概念構建：它們在已配置於參與裝置間通訊的所有裝置上的緩衝區上運作。本地裝置可以使用在其自身本地緩衝區內衍生的位址，連同目標 Rank 的索引來指定操作的目的地。下圖說明了此概念：兩個 GPU 配置了相同大小的對稱緩衝區，並保留指向它們的 src 指標。第一個 Rank 想要傳送 3 個整數至第二個 Rank，並將它們放置於緩衝區開頭：nvshmem_int_put_nbi 從本地緩衝區衍生開始位址，並指定目標裝置。第二個 Rank 從其自己的緩衝區衍生一個偏移量，傳送一個元素至第一個裝置，偏移一個位置。雖然目的地位址必須始終是使用 nvshmem_alloc 配置的對稱緩衝區，但來源緩衝區可以是裝置記憶體的任意區域，前提是它們已向 NVSHMEM 預先註冊。

雖然 NVSHMEM 提供了廣泛的原語，但我們的核心僅依賴 3 個函式，並在其上構建所有同步與防禦機制。

放大

nvshmemx_putmem_signal_nbi_warp：將一個資料區塊從一個裝置傳輸至另一個裝置，同時在遠端裝置上設定一個旗標。該操作會設定 (NVSHMEM_SIGNAL_SET) 或遞增 (NVSHMEM_SIGNAL_ADD) 一個 64 位元的位置。在整個記憶體區塊傳輸完畢後，旗標會被更新。如果遠端裝置觀測到旗標的變更，它便能安全地存取其自身記憶體中的緩衝區。此函式對於合併資料傳輸與同步非常有用。

nvshmemx_signal_op 在單一記憶體位置（通常為 64 位元旗標）上運作，以原子方式設定或遞增它。它在傳送元資料與同步裝置方面非常有用。

nvshmem_uint64_wait_until 用於訊號的接收端，以輪詢一個旗標直到遠端更新它為止。

透過 nvshmem_quiet 或 nvshmem_fence 確保操作間的任何排序形式代價高昂，因此我們對核心的實作避免使用任何其他 NVSHMEM 原語，並將除個別傳送-接收對之外的任何依賴降至最低。傳送資料時，我們始終使用函式的非阻塞版本，而無需等待資料從本地 Rank 傳送出去。跨核心組的其他隱式同步確保了當 NIC 需要存取這些緩衝區時，它們不會被破壞性地覆寫。

NVSHMEM 使用 GPUDirect RDMA，以便 NIC 可以直接讀寫 GPU 記憶體進行資料交換，而無需 CPU 參與。此外，在 ConnectX NIC 上，NVSHMEM 支援 GPU 啟動的通訊（也稱為 Infiniband GPUDirect Async 或 IBGDA），這允許 GPU 直接向 NIC 發佈 RDMA 操作，並直接輪詢 NIC 的完成佇列。在不支援 GPU 啟動通訊的平台上，NVSHMEM 會在主機 CPU 上生成一個「代理」討論串來發起通訊，並代表 GPU 輪詢完成情況。無論是否支援 GPU 啟動的通訊，NVSHMEM 程式皆具可攜性。然而，GPU 啟動的通訊因完全繞過了對 CPU 的迂迴處理，而大幅縮短了延遲。

可攜式核心

我們透過一對分派與組合核心來實作 MoE 通訊。分派核心負責在每個 Rank 上讀取 Token 與路由資訊，將它們分派至適當的專家。組合核心收集專家產生的激活值並將其送回各自的來源 Rank，同時根據路由器計算的權重，計算來自所選專家的加權平均值。核心進一步拆分為傳送與接收組件，以允許資料傳輸與運算重疊。傳送核心皆為非阻塞與非同步：它們只需分派所有對遠端的寫入操作。在另一端，組合核心僅讀取記憶體，並等待所有必要的資料傳輸完畢。在將工作分派至 NIC 後，當資料在連線上非同步傳輸時，GPU 可以在本地執行其他有用的工作，例如應用共享專家或計算注意力。

每一對核心都有自己的對稱記憶體緩衝區，每個 Rank 配置私人儲存空間以接收來自所有其他 Rank 的資料。從 DP 組可以分派的最大 Token 數量在各 Rank 間是固定的，這也設定了每個 Rank 可從每個 DP Rank 為每個本地專家接收的上限。這允許傳送者 Rank 在目的地 Rank 上導出一個唯一位址進行寫入，而無需它們之間進行任何同步。資料接收後，核心會重新排列 Token，將其以接收核心所需的最優格式佈局於記憶體中。雖然緩衝區具有相當大的維度，但它們在模型的所有連續層中會被重複使用。

放大

唯一的全域同步形式（如上圖所示）是在組合-接收核心中實作。一旦資料從傳送端分派，Rank 就會返回並在收到從其他 Rank 所需的所有 Token 時繼續執行，並在處理後立即進入組合-傳送核心。組合-接收中的屏障確保了沒有 Rank 可以超前並開始分派-傳送，而同時還有任何其他 Rank 仍在等待接收資料，這可能會導致破壞性的重疊。為了簡單起見，同步在組合核心中完成；作為一種優化，相同的屏障可以在分派-接收與組合-傳送核心之間實作，藉由允許 NIC 在 GPU 執行其他核心時有更多時間穩定傳輸，進而進一步降低延遲。

兩個核心皆拆分於裝置的所有可用 SM 上：分派-傳送與組合-接收核心必須跨每個 Rank 的最大 Token 計數 (max_m) 進行並行處理，而分派-接收與組合-傳送核心必須處理來自每個 Rank 與本地專家數量的 Token，最大值為 max_m * (num_experts // EP) * (EP // TP)。

分派

分派核心負責將 Token 與區塊縮放因子分派給託管 Token 所路由至的專家之對等節點。傳送端依賴 Warp 專業化（Warp Specialization）來並行處理兩個任務：聚合路由資訊以確定從每個 Rank 傳送了多少 Token，以及將 Token 打包成訊息以分派至遠端。接收端首先等待所有 Token 計數被接收，然後將負載從共享記憶體讀取並解包至將傳遞給後續分組 GEMM 核心的張量中。它還將資訊儲存至與組合核心共享的緩衝區中，以指示每個 Token 應被送回何處。此資訊是必須的，因為接收核心以非確定性順序在連續緩衝區中重新排列 Token。

放大

在傳送端部分，每個區塊上的一個 32 執行緒 Warp 負責讀取路由資訊並計算每個目的地專家預期從此 Rank 接收的 Token 數量，使用 nvshmem_signal_op 傳送計數加 1。計數被遞增 1，因為在遠端從零轉換至非零，標誌著來自來源 Rank 的計數已收到。同時，剩餘的 Warps 合作將 Token 平行複製到所有區塊的對稱記憶體中，將激活值、縮放因子及其在本地 Rank 上的索引打包成一個連續區塊。組合傳送者需要此索引以確定寫入 Token 的位址。接下來，在透過 Warp 組間的屏障確保所有資料皆已複製後，Warps 再次獨立運作，每個皆並行地將相同的緩衝區傳送至不同的專家。Token 使用 nvshmemx_putmem_signal_nbi_warp 傳送，該函式也會在遠端裝置上以原子方式增加來自本地 Rank 的已傳送 Token 計數。在 DP 組內，由於每個 Rank 都擁有待傳送 Token 的複本，因此分派是均勻平衡的，每個裝置皆傳送 Token 的子集。

在接收端，所有核心首先透過等待總 Token 計數傳送過來與傳送者 Rank 同步。之後，它們都等待以原子方式遞增的已傳送 Token 計數穩定至總計數，這標誌著來自來源 Rank 的所有負載也已傳送過來，這歸功於 putmem 呼叫的語義。核心使用 nvshmem_uint64_wait_until 對計數進行輪詢，並跨所有區塊與執行緒進行操作並行化。隨後，跨區塊屏障確保除非所有資料皆已正確接收，否則沒有區塊會讀取緩衝區。Token 分佈於區塊間並透過以原子方式遞增的計數器同步，從對稱記憶體複製到常規連續緩衝區，這些緩衝區成為實現專家的後續核心之輸入。來源 Rank、專家索引與 Token 索引被分開儲存，精確指出組合核心必須傳送激活值的位置。儘管 DP 組內的 Token 是從不同的裝置傳送的，但它們全被分組在一起以傳遞給對應的專家。

組合

組合核心再次拆分為傳送與接收兩半：傳送者將未量化的 16 位元激活值複製到遠端裝置，而接收核心則等待資料傳送完畢，並在本地計算所有專家貢獻的加權平均值。此外，它們還充當屏障以同步分派-組合序列：每個 Rank 在進入傳送核心時會在每個對等節點上設定旗標，若接收核心未觀察到旗標已設定，則不被允許返回。同步的延遲極小，因為它與實際的通訊與運算重疊。

放大

在傳送端，核心並行遍歷分配給所有本地專家的 Token 列表，並將其寫入目的地 Rank 的緩衝區。目標 Rank、專家索引與 Token 索引係從散佈核心（Scatter Kernels）在收到 Token 後填入的每個 Token 緩衝區中讀取。每個傳送者對每個專家都有其自己的私人記憶體區域可供寫入（如上圖所示），避免了同步需求。與分派類似，組合操作會以原子方式遞增目的地 Rank 上的每個 Token 計數器，以指示資料的接收：當計數器與一個 Token 分派至的專家數量相符時，便可存取該 Token 的內容。

在接收核心中，Token 列表跨多個區塊進行並行遍歷，透過輪詢由訊號操作設定的旗標來等待其內容送達。送達後，從私人緩衝區讀取負載，並由路由表指示從哪個緩衝區讀取以及分配給每個專家何種權重。結果隨後寫入外部配置的張量中，核心在所有裝置通過屏障後結束執行。跨 DP 組，所有 Rank 皆接收每個專家激活值的複本，以計算其自己的複本。

效能基準測試

我們在一個包含 128 個 GPU 的叢集上評估我們的核心，這些 GPU 分佈在 16 個節點上，並透過 InfiniBand 與 CX-7 NIC 連接。我們同時測試了 GPUDirect Async (IBGDA) 與具 CPU 代理的可靠連接 (RC)。我們將其與高度優化的 DeepSeek 實作，以及 PyTorch（透過 NCCL）或 NVSHMEM 提供的稠密原語進行比較。

放大

雖然在稠密 NVSHMEM All-to-All 操作中，IBRC 與 IBGDA 的效能相似（6378 µs 對 6180 µs），但 IBGDA 在稀疏核心下速度顯著更快。加總分派與組合，IBGDA 使用 902 µs，而 IBRC 需要 3223 µs——延遲改善了 3.6 倍。雖然 All-to-All 受限於頻寬，但稀疏核心廣播的資料量少了好幾個數量級，因而受限於延遲。藉由直接從 GPU 觸發網路傳輸，無需 CPU 代理來協調 GPU 與 NIC，端到端延遲大幅降低。

雖然我們的可攜式核心比高度優化的 DeepSeek 核心慢約 2 倍，但與稠密核心相比，它們將延遲降低了 10 倍（902 µs 對 9944 µs）。

此外，與函式庫原語不同，拆分為傳送者-接收者組件也允許隱藏部分延遲。

放大

在單節點 (EP8) 上，NVSHMEM 利用 NVLINK 進行傳輸，提供比節點間網路更低的延遲與更高的吞吐量。我們的可攜式核心在單節點上比 DeepEP 快約 2.5 倍（186 µs 對 481 µs）。

進一步改進

此處描述的核心優於 ML 框架的內建原語，並在不針對特定裝置間傳輸（如 NVLink 或 InfiniBand）進行過度專業化的情況下提供不錯的效能。除了已提及的機會外，透過以更專業的版本替換通訊原語，還可以獲得進一步的效能提升。例如，在 NVLink 上，使用對稱記憶體可以替換為跨裝置共享的緩衝區，從而消除部分複製並允許更細粒度的同步方案（針對個別 Token 而非 Token 批次）。在 InfiniBand 上，一種實作可以直接存取底層佇列對，從而消除傳輸預期與已傳送 Token 計數的需求，並在 Token 到達時將其從佇列沖刷至輸出緩衝區。然而，此類實作是以可攜性為代價的，而我們的實作允許在投入開發精力優化特定硬體之前，對具有特定配置的叢集進行評估。

結論

我們推出了一款高效能、可攜式的 MoE 通訊函式庫，與標準 All-to-All 通訊相比，其效能快 10 倍，同時保持了跨不同硬體配置的相容性。在採用 NVLink 的單節點部署中，我們的解決方案證明了比先前實作低 2.5 倍的延遲。

我們的方法透過關鍵創新平衡了效能與可攜性，這些創新包括 GPU 啟動的通訊支援、實現運算-通訊重疊的分割核心架構，以及使用最少量 NVSHMEM 原語的高效 Token 分派。儘管在多節點設定上比高度專業化的實作慢約 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}
}