高效且便攜的混合專家模型通訊

我們推出了一個高效能、便攜且開源的混合專家模型（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 的 671B 個參數中，在推論過程中只有 37B 個會與給定的 Token 相乘。與密集模型相比，這大大減少了解碼 Token 所需的運算量並改善了延遲。

與密集模型相比，MoE 模型對推論提出了一些額外的挑戰。雖然專家本身是小型 MLP 層，但每個解碼器層都包含一個路由器，用於決定將 Token 分派給哪些專家，每個 Token 都會被分派給多個專家。路由器通常是一個小型線性層，可產生機率分佈。通常會挑選具有最高 K 分數的專家，並將最終激活計算為加權平均值，將專家輸出乘以從路由器機率分佈派生的權重進行加總。

為了最大程度地減少分佈式 MoE 系統中的延遲，可以在多個裝置之間利用平行化，但這引入了通訊挑戰。諸如 Mixtral 8x7B 或 Llama-70B 之類的模型適合放入單一節點上的 8 個裝置中，受益於快速且低延遲的 NVLink 互連（高達 900Gbps）。然而，較大的模型需要使用 InfiniBand（峰值為 400Gbps）將專家分片到多個節點上，這會帶來額外的延遲挑戰。

放大

在本文中，我們探討了專家平行化的問題，描述了我們對 GPU 核心的實現，這些核心能有效地將 Token 分派給持有指定專家的裝置，並透過計算專家輸出的加權和將它們收回。雖然存在針對特定網路硬體的積極優化實現，但我們提出了一個高度便攜的實現，它依賴一組最少的 NVSHMEM 原語，同時仍能提供比標準 all-to-all 通訊快 10 倍的效能。我們的實現單節點配置上實現了最先進的效能，同時在各種網路環境中保持了出色的便攜性。

分片與平行化

為了進行高效的推論，模型的權重必須保存在裝置記憶體中，同時還要留出足夠的空間給模型的向前傳遞所需的激活、KV 快取和其他緩衝區。最強大的模型的容量甚至超過了最強大 GPU 的容量，因此推論必須分佈在多個裝置上，這些裝置可以共同儲存模型的權重。基於權重的分片方案，必須使用不同的通訊和運算方案來卸載運算並同步裝置。

放大

專家平行化（EP）如第一張圖所示，僅使專家運算平行化。不同的專家被指派給持有其權重的不同裝置。路由後，Token 被傳送到相應的裝置，隨後收集並累積結果。路由的複雜性取決於模型其他部分的平行化程度：複製其他層可以完全消除對路由的需求，因為每個 Rank 可以從本地複製的路由表中選擇 Token。但是，如果只有一個 Rank 執行路由，則需要廣播來將 Token、索引和權重分派給各自的專家。最後，全聚合（all-gather）或全對全廣播將輸出 Token 與繼續模型執行的任何 Rank 進行同步。這樣的實現相對簡單，因為 torch 已經公開了所需的原語，儘管透過將累積與從專家收集 Token 的聚合操作進行融合，可以獲得一些好處。

放大

專家專屬平行化無法以理想方式擴展，因為當模型正在執行非專家層（例如注意力、規範和取樣）時，叢集中的節點可能會處於閒置狀態。然而，這些層（主要是注意力）的運算也可以受益於張量平行化（TP）。大多數模型依賴多頭注意力（multi-head attention），這意味著注意力頭及其相應的 Q、K 和 V 投影也可以在裝置之間進行分片，從而在各種層之間複製或聚合切片。如果注意力分佈在所有裝置上，全聚合可以同步激活，從而允許複製路由，這需要類似於專家平行化情況的同步原語來實現高效的實現。然而，此層級的平行化存在限制，因為將注意力頭的數量減少到低於特定閾值會產生遞減的報酬。

放大

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

NVSHMEM

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

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

雖然 NVSHMEM 提供了廣泛的原語，但我們的核心僅依賴 3 個函式，並在其之上建構所有同步和隔離（fencing）。

放大

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 上，NVSHEMEM 支援 GPU 發起的通訊（也稱為 Infiniband GPUDirect Async 或 IBGDA），這允許 GPU 直接向 NIC 發佈 RDMA 操作並直接輪詢 NIC 的完成佇列。在不支援 GPU 發起的通訊的平台上，NVSHMEM 會在主機 CPU 上衍生（spawn）一個「代理」執行緒，以代表 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 時，它就會返回並繼續執行，並且一旦處理完畢就會進入組合傳送核心。組合-接收中的屏障（barrier）確保了沒有任何 Rank 可以領先並開始分派-傳送，而同時任何其他 Rank 仍在等待接收資料，這可能會導致破壞性的重疊。為了簡單起見，同步是在組合核心中完成的，作為一種優化，可以在分派接收和組合傳送核心之間實現相同的屏障，透過允許 NIC 有更多時間在 GPU 執行其他核心時穩定傳輸，從而進一步減少延遲。

兩個核心分佈在裝置的所有可用 SM 上：分派傳送和組合接收核心必須跨每個 Rank 的最大 Token 數（max_m）進行平行化，而分派接收和組合傳送核心必須處理來自每個 Rank 和本地專家的這麼多 Token，最多為 max_m * (num_experts // EP) * (EP // TP)。

分派

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

放大

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

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

組合

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

放大

在傳送端，核心平行走訪分配給所有本地專家的 Token 列表，將它們寫入目標 Rank 上的緩衝區。目標 Rank、專家索引和 Token 索引是從接收 Token 時由散射（scatter）核心填入的每個 Token 緩衝區中讀取的。正如上圖所示，每個傳送者都有自己的私有記憶體區域供每個專家寫入，從而避免了同步的需求。與分派類似，組合以原子方式增加目標 Rank 上的每個 Token 計數器，以指示資料的接收：當計數器與分派 Token 的專家數量相匹配時，就可以存取 Token 內容。

在接收核心中，Token 列表在多個區塊中平行走訪，透過輪詢由發訊操作設定的旗標來等待其內容到達。到達後，從私有緩衝區讀取有效負載，路由表指示從哪個緩衝區讀取以及為每個專家指派什麼權重。然後將結果寫入外部配置的張量中，一旦所有裝置都透過屏障，核心就會完成執行。在 DP 群組中，所有 Rank 都接收每個專家激活的複本，以計算其自己的複本。

基準測試

我們在透過 InfiniBand 和 CX-7 NIC 連接的 16 個節點上分佈的 128 個 GPU 叢集上評估我們的核心。我們嘗試了 GPUDirect Async (IBGDA) 和具有 CPU 代理的可靠連接 (RC)。我們將它們與高度優化的 DeepSeek 實現以及 PyTorch（透過 NCCL）或 NVSHMEM 提供的密集原語進行比較。

放大

儘管在密集 NVSHMEM all-to-all 操作上，IRBC 和 IBGDA 的效能相似（6378 微秒對 6180 微秒），但 IBGDA 在稀疏核心上的速度明顯更快。將分派和組合加起來，IBGDA 使用 902 微秒，而 IBRC 需要 3223 微秒 - 延遲改善了 3.6 倍。雖然 all-to-all 是頻寬受限的，但稀疏核心廣播的資料量要少幾個數量級，因此受延遲限制。透過直接從 GPU 觸發網路傳輸，而不需要 CPU 代理來協調 GPU 和 NIC，端到端延遲顯著降低。

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

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

放大

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

進一步改進

這裡描述的核心優於 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://www.perplexity.ai/hub/blog/efficient-and-portable-mixture-of-experts-communication}
}