在 AWS EFA 上實現萬億參數模型

在 Perplexity，我們為產品、API 和研究團隊採用了最佳模型。大型開源專家混合 (MoE) 模型（例如 Kimi-K2）帶來了特殊的挑戰，因為配備 8x NVIDIA H200 GPU 的最大推論節點無法有效容納這些模型，因此必須進行多節點部署。我們提出了一套用於專家平行處理的核心程式，這些核心程式在 ConnectX-7 上實現了業界領先的延遲表現，超越了 DeepEP 的效能。這些核心程式也是首批在 AWS Elastic Fabric Adapter (EFA) 上實現可行延遲表現的程式，從而實現了萬億參數模型的部署。

在 GitHub 上嘗試我們的核心程式，並在 arXiv 上閱讀完整研究報告。

專家混合 (MoE) 已成為擴展至數千億或萬億參數的標準架構，同時在推論時也能提供合理的延遲。MoE 模型將 Transformer 的稠密層替換為一組專家模型和一個小型線性路由層，後者負責決定要將權杖 (token) 乘以哪些專家模型。透過將與同一專家模型相關的權杖分組，問題變得稀疏且易於平行處理，因為不同的專家模型現在可以託管在分佈於多個節點的不同 GPU 上。

與其他形式的平行處理（如張量平行處理 (TP) 和資料平行處理 (DP)，這些可以透過 torch.distributed 和 NCCL 等集體通訊程式庫輕鬆實現）不同，MoE 路由涉及稀疏的點對點通訊。調度核心 (dispatch kernel) 必須分割傳入的權杖集合並將其分發到託管專家模型的排位 (ranks)，而其對偶問題即組合核心 (combine kernel) 必須接收其已處理的對應部分並計算其加權平均值。這些傳輸無法對應到集體通訊，需要客製化核心程式才能實現低延遲。

在 Perplexity，我們一直致力於開發可攜式節點間核心程式和專業節點內核心程式，並在先前關於 MoE 通訊的文章中更深入地介紹了我們的工作。我們還在另一篇部落格文章中展示了多節點部署可以為大型 MoE 模型帶來更低的延遲和更高的輸送量。

在本文中，我們將深入探討我們新的節點間核心程式的實作，該實作大幅改進了我們先前基於 NVSHMEM 的工作效能。除了展示在 ConnectX-7 上實現的業界領先成果外，我們還證明了在我們的推論執行時期中，這些核心程式可用於有效地服務最大的開源模型。

張量平行處理和專家平行處理可以在單一節點內有效擴展，因為 NVLink 提供了高輸送量 (900GB/s) 和微秒級的延遲。雖然 NVIDIA GB200 和 GB300 系統具備足夠的記憶體來容納萬億參數模型的權重和 KV 快取，但在 AWS p5en 執行個體（配備多達 8 個 H200 GPU）上的部署受限於 1120GB 的 HBM，這必須在模型權重和 KV 快取之間進行分配。因此，某些模型需要透過 InfiniBand 進行節點間部署。

儘管 InfiniBand 速度極快，但它僅提供 400Gbps (50GB/s) 的輸送量，這為 MoE 路由增加了數十到數百微秒的延遲。幸運的是，透過使用共享專家模型、微批次處理 (micro-batching) 和計算-通訊重疊，可以隱藏這些增加的延遲，從而以具有競爭力的成本提供可行的輸送量。DeepSeek 透過部署在 H800 GPU 叢集和 NVIDIA ConnectX-7 InfiniBand 配接器上的 DeepSeek-R1 開創了這些技術。

在 AWS 叢集上，由於使用其客製化的 Elastic Fabric Adapters (EFA)，情況有所不同。即使這些配接器在訓練工作負載中通常使用的集體通訊中能提供 400Gbps 的峰值輸送量，但在 MoE 分發和組合期間交換的訊息大小方面，它們略遜於 ConnectX-7。此外，EFA 不支援 GPUDirect Async，因此需要一個 CPU 代理討論串 (proxy thread) 作為 GPU 與 NIC 之間的橋樑以發起傳輸。除了增加複雜性外，額外的 PCIe 交易還為所有傳輸增加了數微秒的額外開銷。諸如 NVSHMEM 之類的架構允許核心程式編寫者提供獨立於裝置的實作，然而通用 API 產生的開銷使得部署在成本和延遲方面都缺乏效率。

我們先前的核心程式透過 NVSHMEM 將底層傳輸實作抽象化。它們依賴 nvshmemx_putmem_signal_nbi_warp 以原子計數器同步對等點，從而逐一傳輸權杖。雖然這在配備 IBGDA 的 ConnectX-7 NIC 上提供了合理的效能，但基於代理的 ConnectX-7 IBRC 實作速度顯著變慢，而 EFA 實作則未能達到毫秒以下的延遲表現。

我們將這些核心程式作為底層硬體的測試平台，同時鎖定最適合向我們的模型公開的介面。我們推測，核心程式與 NIC 或代理討論串之間的頻繁傳輸與互動阻礙了效能。此外，我們的實作除了可靠性之外不依賴硬體的任何保證，這表明我們可以透過減少同步原語的開銷來加快處理速度。

根據這些觀察，我們設計了基於 EFA 的核心程式並公開了相同的介面，在底層最佳化了裝置核心程式與主機代理的互動，以達到峰值效能。

具有主機代理的 RDMA 與 NVLink 混合通訊，顯示發起與接收的 rank、代理工作程序和 NIC。

新核心程式遵循混合 CPU-GPU 架構，其中 GPU 核心程式與裝置上的模型進行互動，而執行在主機 CPU 上的代理討論串則處理與 NIC 的互動。新核心程式重複使用了我們為 KV 快取傳輸建構的相同 TransferEngine。

MoE 路由由調度核心程式執行，該核心程式接收權杖以及模型路由所需的每個權杖對應的專家模型，並透過網路將其傳送到託管專家模型的排位。在分組 GEMM 核心程式將權杖與專家權重相乘後，組合核心程式將權杖傳輸回其來源排位，並在將其返回給模型之前計算其加權平均值。

調度核心程式和組合核心程式都被拆分為發送端和接收端兩部分，以促進重疊和微批次處理。發送端將權杖洗牌到發送緩衝區，以便透過單一 RDMA 寫入操作將其寫入對等點。在執行發送端後，GPU 可以執行其他工作（例如共享專家模型乘法），同時網路傳輸在背景中掛起。接收端核心程式會阻塞，直到所有傳輸都已接收。在背景中，主機端工作討論串處理傳輸：當發送緩衝區已填滿時，它會收到來自 GPU 的訊號，之後它會分發 RDMA 寫入。一旦 NIC 發出完成訊號，它就會解除阻塞接收端，接收端將權杖傳回給模型。

這些核心程式與主機緊密耦合，透過統一記憶體和 GDRCopy 相互通訊。CUDA 統一記憶體允許將裝置記憶體或主機記憶體對映到兩者的位址空間中，透過 PCIe 透明地處理通訊。

雖然統一記憶體適用於大量傳輸，但低延遲輪詢是透過 GDRCopy 完成的。GDRCopy 使用 GPUDirect RDMA 以微秒級的延遲讀取和寫入 GPU 記憶體。

除了 RDMA 之外，節點內的核心程式還利用 NVLink 來交換權杖。這可以卸載相當大比例的流量，EP32 最多可達四分之一，EP64 最多可達八分之一，從而實現整體更低的延遲。

發送緩衝區對於調度核心程式的設計對於實現涉及主機代理的可行延遲至關重要。如果有 N 個排位，每個排位分發 T 個權杖，並且每個權杖路由到 R 個專家模型（從 E 個專家模型中選出），在最壞的情況下，如果所有權杖都被路由到同一個排位，則一個排位最多可以接收 N * T * max(R, E / N) 個權杖。此外，只有當發送端連續佈局其寫入時，才能達到此最小邊界，因為私有接收緩衝區會耗盡太多記憶體。為了讓發送端能夠在沒有衝突的情況下連續佈局其寫入，路由資訊的完整知識必須分發到每個排位。因此，我們的調度核心程式首先交換每個專家模型的權杖計數，之後它們各自可以確定寫入的位置以及從每個排位接收權杖的偏移量。

僅僅交換路由資訊會導致不必要的長延遲，因為每個排位在發送自己的寫入之前都必須等待寫入到達。相反，我們在每個接收排位上預留少量的私有每發送端空間，並在路由資訊寫入的同時分發部分權杖。這確保了在整個調度過程中的大部分時間裡都能完全利用頻寬。在路由交換後，其餘權杖透過每個遠端對等點的單一寫入進行傳輸。路由資訊的交換、固定大小的緩衝區以及填充它們的傳輸如下圖所示。

MoE 派送的緩衝區配置，顯示 token 資料與路由中繼資料如何在兩個 rank 之間排列與交換。

組合核心程式要簡單得多，因為路由資訊在調度階段已對它們可用。它們可以透過每個對等點的單一寫入複製權杖，之後它們從接收緩衝區讀取，以將權杖的加權平均值計算到輸出張量中。

對於 NVLink，節點內的發送和接收緩衝區會被對映到對等點，以便它們能夠直接從相鄰排位推入或拉取權杖。

下圖詳細說明了 GPU、CPU 和 NIC 之間的時序和互動，展示了一次分發和組合的迭代與模型核心程式交錯進行。

工作從分發發送開始，核心程式聚合路由資訊並透過單一區塊內的共享記憶體中的原子累加來計算每個專家模型的權杖計數。一旦獲得權杖計數，它們就會複製到主機記憶體，並立即通知代理開始將路由資訊分發到所有其他對等排位。同時，分發過程繼續計算權杖計數的累加和，識別將權杖封裝到連續的每對等點發送緩衝區所需的偏移量。一旦獲得偏移量，所有要透過 EFA 發送的權杖都會被封裝到發送緩衝區中。通知代理將寫入內容發佈到節點間對等點的發送緩衝區中。在節點內，核心程式還透過 NVLink 立即將權杖推送到同一節點內的私有緩衝區。

MoE 派送的時序與協調，顯示分階段的 GPU、CPU 與 NIC 作業及其同步。

一旦排位從所有對等點接收到路由資訊，它就會將其集中並計算識別分發和組合權杖所需的各種偏移量。這項工作在傳輸掛起時在 CPU 上完成，因此對延遲不敏感。隨後，一旦所有偏移量已知，剩餘權杖（超過私有緩衝區大小）的寫入內容就會發佈，分發的發送階段完成。

在接收端，核心程式首先透過 NVLink 從同一節點拉取超過私有緩衝區容量的權杖。接下來，它會阻塞，直到主機代理發出所有傳輸完成的訊號。一旦收到訊號，權杖就會從接收緩衝區洗牌到傳遞給組 GEMM 的輸出張量中。由於緩衝區在分發和組合之間重複使用，因此 NVLink 和 RDMA 屏障都能確保接收緩衝區在允許下一個核心程式覆蓋它們之前已被讀取。

組合核心程式遵循大致相似的模式，即洗牌權杖、向代理發出訊號並等待完成。然而，發佈其寫入內容的中繼資料很容易獲得，並且所有遠端偏移量都是已知的，因此所有節點內傳輸的權杖都透過 NVLink 推送到接收器，而節點間權杖則透過單一寫入進行傳輸。在返回結果之前，核心程式會累加來自接收緩衝區的權杖，並將其寫入輸出張量。

為了對我們的實作進行原型設計，我們最初重複使用了為分拆預填充 (disaggregate prefill) 和解碼建構的 Transfer Engine，以發出和追蹤 RDMA 寫入。然而，與每個節點向同一個對等點發佈大量寫入的分拆解碼不同，MoE 分發在一次操作中向所有對等點發佈小額傳輸。開箱即用的效能並不理想，我們的實作在 ConnectX-7 上比 DeepEP 滯後約 20us。因此，我們專門化了 TransferEngine 以支援散射 (scatter) 和屏障 (barrier) 兩種操作，以最佳化與大量對等點的互動。

散射和屏障操作都公開了對等群組，這些群組可以預先註冊以預先填充底層傳輸實作所使用的資料結構（工作請求 (WR) 範本化）。散射將來自同一個來源緩衝區的不同切片複製到對等群組所有成員的不同偏移量中。屏障透過立即值向所有對等點發出訊號。大量批次處理和處理這些操作可在關鍵路徑上節省寶貴的微秒時間。在 EFA 上（其中兩個 200Gbps NIC 聚合提供 400Gbps 頻寬），我們將對等群組分片到兩個 NIC 上，而不是分片每次傳輸的位元組。

儘管最初我們僅使用 libfabric 在 EFA 上開發了 TransferEngine，但在 TransferEngine 中使用 libibverbs 加入 ConnectX-7 支援後，MoE 核心程式無需更改任何程式碼即可在 ConnectX-7 上運作。與基於可擴展可靠資料報 (SRD) 的 EFA 相比，ConnectX-7 需要額外的連線設定和對等管理工作。我們使用不可靠資料報 (UD) 佇列對來引導可靠連線 (RC)。我們為每個遠端對等點建立兩個佇列對，一個用於雙向 SEND/RECV，另一個用於單向 WRITE_IMM，因為 RECV 和 WRITE_IMM 都按順序消耗 RECV 操作，而這兩種操作需要不同的緩衝區大小。此外，我們放寬了 ConnectX-7 和 GPU 記憶體之間的 PCIe 交易順序。與 EFA 類似，我們實作了 WR 範本化最佳化。此外，當向同一個對等點發送多個小訊息時，我們使用 WR 鏈結來提高訊息速率。

由於 GPU 端的工作在兩個配接器上是相同的，我們將底層傳輸實作推向極限，以超越 DeepEP 的效能，從而強調 EFA 與 ConnectX-7 MoE 路由之間的差距是由底層硬體的能力所決定的。

我們在配備 ConnectX-7 和 EFA 配接器的 H200 系統上，針對標準 DeepSeek-V3 和 Kimi-K2 配置評估了分發和組合延遲，配置包括 7168 的隱藏維度、區塊縮放的 fp8 分發和 bf16 組合。我們考慮分發的批次大小為 128，組合的批次大小為 4096。在 ConnectX-7 上，我們與 DeepEP 和我們先前基於 NVSHMEM 的 pplx-kernels 進行比較，同時使用 GPUDirect Async (IBGDA) 和主機代理 (IBRC)。我們在 EP16、EP32 和 EP64 上分別實現了 459us、582us 和 692us 的組合延遲。我們勝過依賴通用代理的其他基於 EFA 的核心程式（例如 UCCL-EP，其分別達到 519us、966us 和 1159us）。

在解碼時，儘管我們的核心程式比 DeepEP 慢幾微秒，但我們實現了整體更低的延遲，因為我們的組合核心程式速度快得多。雖然我們使用了主機代理，但由於我們的批次傳輸，我們的效能超過了所有基於 IBGDA 的無代理實作。我們還證明了代理實作的重要性，因為我們高度專業化的代理比通用 NVSHMEM 實作快一個數量級。儘管在傳輸 256KB 封包（接近 MoE 路由交換的典型封包大小）時，EFA 的速度是 ConnectX-7 的兩倍，但兩個配接器之間的差距整體並不顯著。

EP64 MoE 預填充的派送與合併延遲，對比 pplx 與 DeepEP 通訊路徑。

儘管我們針對解碼最佳化的核心程式可以在不進行任何額外調整的情況下擴展到預填充批次大小，但效能落後於 DeepEP。這主要是因為我們僅在組合的接收端累加權杖，而 DeepEP 在發送端累加以減少傳輸的資料量。然而，這種最佳化可能會對準確性產生影響。此外，由於核心程式在一次寫入中傳輸資料，我們為發送和接收緩衝區使用了更多記憶體。

為了證明我們核心程式的可行性，我們對我們內部推論引擎的端到端效能進行了基準測試，該引擎在多個配備 Elastic Fabric Adapters (EFA) 和 H200 GPU 的 p5en 執行個體上服務 DeepSeek-V3 (671B 參數) 和 Kimi-K2 (1T 參數)。我們在草稿驗證模式下執行模型，並假設使用單一 MTP 層的接受率為 80% 來報告延遲。

在 DeepSeek-V3 上，我們與執行高效僅限 NVLink 分發和組合核心程式的 EP=8 節點內基準進行比較。節點間部署整體上要麼匹配要麼超過了單節點部署的效能，特別是在中型批次上，證明了跨節點擴展確實能提高整體輸送量。

在 Kimi-K2 上，這些核心程式實現了服務。由於其規模，該模型無法使用 H200 節點上的資料平行處理 (DP) 放入同一個節點內，因此需要進行節點間部署才能服務該模型。由於模型具有較少的注意力頭 (attention heads)，我們實現了比 DeepSeek-V3 更好的延遲，並且可以在 EFA 上以可行的延遲服務模型。

DeepSeek-V3 每 GPU 吞吐量與輸出速度對比，適用於 EP32、EP16 和 EP8 配置。

Kimi-K2 每 GPU 吞吐量與輸出速度對比，適用於 EP32 和 EP16 配置。

我們也考慮了微批次處理，儘管它僅在非常大的批次大小或使用多於一個草稿權杖時才改善了延遲。為了從將分組 GEMM 問題拆分為兩個中受益，權杖批次大小必須足夠稠密。大部分時間仍花在 MoE 分發和組合上，這表明對核心程式的任何進一步改進都將相應地改善延遲並促進微批次處理。

我們正與 AWS 工程師密切合作，以提高 EFA 的效能。就我們而言，我們將遵循 libfabric 的更新，這將減少資料平面的開銷。此外，我們將嘗試 efa-direct 以進一步減少使用者空間的開銷並直接與硬體互動。

使用 NVIDIA Magnum IO NVSHMEM 和 GPUDirect Async 改善 HPC 系統的網路效能

BibTeX

@online{perplexity2025EnablingTrillionParameterModelsOnAwsEfa,
  author = {{Perplexity} and Nandor Licker and Kevin Hu and Vladimir Zaytsev and Lequn Chen},
  title = {{Enabling trillion-parameter models on AWS EFA}},
  date = {2025-11-04},
  url = {https://www.perplexity.ai/hub/blog/enabling-trillion-parameter-models-on-aws-efa}
}

簡介

基於 InfiniBand 的推論

基於 NVSHMEM 的 pplx-kernels

基於 EFA 的 MoE 全對全通訊 (all-to-all)

發送與接收緩衝區

分發與組合

用於 MoE 路由的 TransferEngine

評估

分發/組合延遲

Kimi-K2 / DeepSeek-V3 輸送量

未來工作

參考資料

引用