在 AWS EFA 上啟用萬億參數模型

在 Perplexity，我們為我們的產品、API 和研究團隊使用最好的模型。像 Kimi-K2 這樣的大型開源 Mixture-of-Experts 模型帶來了特殊的挑戰，因為擁有 8 個 NVIDIA H200 GPU 的最大推理節點無法有效容納它們，因而需要進行多節點部署。我們提出了一組針對專家平行的核心 (kernels)，在 ConnectX-7 上實現了最先進的延遲，超越了 DeepEP 的效能。相同的核心也是第一批在 AWS Elastic Fabric Adapter (EFA) 上實現可行延遲的核心，從而實現了萬億參數模型的部署。

在 GitHub 上試用我們的核心，並在 arXiv 上閱讀完整的研究論文。

Mixture-of-Experts (MoE) 已成為擴展到數千億或萬億參數的標準架構，同時在推理時也能提供合理的延遲。MoE 模型用一組專家和一個小型線性路由層取代了 transformer 的密集層，該線性層決定了將 token 與哪個專家相乘。透過將與同一專家相關的 token 分組，問題變得稀疏且易於平行化，因為不同的專家現在可以託管在分佈於多個節點的不同 GPU 上。

與其他形式的平行化（例如 Tensor Parallelism (TP) 和 Data Parallelism (DP)）不同，後者可以使用諸如 torch.distributed 和 NCCL 等集合通訊函式庫輕鬆實現，而 MoE 路由涉及稀疏的對等點對對等點 (peer-to-peer) 通訊。分發核心 (dispatch kernel) 必須分割傳入的 token 集合並將其分發到託管專家的 rank，而其對偶問題則是合併核心 (combine kernel)，它必須接收其處理後的對應項並計算其加權平均值。這些傳輸不會對應到集合通訊，且需要自定義的核心來實現低延遲。

在 Perplexity，我們一直在開發可攜式的節點間核心以及專門的節點內核心，並在之前關於 MoE 通訊的文章中更深入地介紹了我們的工作。我們還在另一篇網誌文章中展示了多節點部署可以為大型 MoE 模型帶來更低的延遲和更高的吞吐量。

在本文中，我們深入探討了我們的新節點間核心的實現，這大大改善了基於 NVSHMEM 的先前工作的效能。除了展示在 ConnectX-7 上的最先進結果之外，我們還證明了在我們的推理執行階段 (runtime) 內，這些核心可以用於有效提供最大的開源模型。

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

儘管 InfiniBand 仍然非常快速，但它僅提供 400Gbps (50GB/s) 的吞吐量，從而為 MoE 路由增加了數十到數百微秒的延遲。幸運的是，透過使用共享專家、微批次 (micro-batching) 和計算與通訊重疊，可以隱藏增加的延遲，從所以合理的成本提供可行的吞吐量。DeepSeek 透過部署在 H800 GPU 叢集和 NVIDIA ConnectX-7 InfiniBand 轉接卡上的 DeepSeek-R1 開創了這些技術。

在 AWS 叢集上，由於使用了其自訂的 Elastic Fabric Adapter (EFA)，情況有所不同。儘管這些轉接卡在訓練工作負載中通常使用的集合通訊上提供了 400Gbps 的峰值吞吐量，但在 MoE 分發和合併過程中交換的訊息大小方面，它們略遜於 ConnectX-7。此外，EFA 不支援 GPUDirect Async，需要一個 CPU 代理執行緒 (thread) 將 GPU 連接到 NIC 以便發起傳輸。除了增加複雜性之外，額外的 PCIe 事務還為所有傳輸增加了價值微秒的額外負荷 (overhead)。像 NVSHMEM 這樣的框架允許核心編寫者提供與裝置無關的實現，然而通用 API 會產生額外負荷，從而在成本和延遲方面使部署變得無效。

我們先前的核心透過 NVSHMEM 將底層傳輸實現抽象化。它們依靠 nvshmemx_putmem_signal_nbi_warp 使用原子計數器 (atomic counters) 來同步對等點，以逐個傳輸 token。雖然這在帶有 IBGDA 的 ConnectX-7 NIC 上提供了合理的效能，但基於代理的 ConnectX-7 IBRC 實現要慢得多，而 EFA 實現則未能提供低於一毫秒的延遲。

我們將這些核心用作底層硬體的測試平台，同時也確定了要公開給我們的模型的最適介面。我們假設核心與 NIC 或代理執行緒之間的頻繁傳輸和互動阻礙了效能。此外，我們的實現除了可靠性之外，不依賴來自硬體的任何保證，這表明我們可以透過降低同步原語的額外負荷來加快速度。

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

採用主機代理的 RDMA 與 NVLink 混合通訊，顯示發起與接收的 rank、代理工作執行緒及 NIC。

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

MoE 路由由分發核心執行，該核心從模型接收 token 和要路由到的每 token 專家，並透過網路將其傳輸到託管專家的 rank。在分組的 GEMM 核心將 token 與專家權重相乘之後，合併核心將 token 傳輸回其來源 rank，並在將其返回模型之前計算其加權平均值。

分發和合併核心均分為發送端和接收端兩半，以促進重疊和微批次。發送端將 token 打亂到發送緩衝區中，以便透過單個 RDMA 寫入操作將其寫入對等點。執行發送端後，GPU 可以執行其他工作（例如共享專家相乘），同時透過網路的傳輸在背景中等待。接收核心會阻塞，直到接收到所有傳輸。在背景中，主機端的背景工作執行緒 (worker thread) 處理傳輸：當發送緩衝區已被填入時，它會收到來自 GPU 的信號，此後它會分發 RDMA 寫入。一旦 NIC 發出完成信號，它就會解除阻塞接收端，接收端會將 token 返回給模型。

這些核心與主機緊密耦合，透過統一記憶體 (unified memory) 和 GDRCopy 相互通訊。CUDA 統一記憶體允許將裝置記憶體或主機記憶體對應到兩者的記憶體空間中，從而透明地處理透過 PCIe 進行的通訊。

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

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

分發核心的發送緩衝區設計對於實現涉及主機代理的可行延遲至關重要。如果有 N 個 rank，每個 rank 分發 T 個 token，每個 token 路由到 E 中的 R 個專家，則在最壞的情況下（如果全部路由到它），一個 rank 可以接收最多 N * T * max(R, E / N) 個 token。此外，只有當發送者連續佈置其寫入時，才能達到此最小界限，因為私有接收緩衝區會耗費太多記憶體。為了讓發送者能夠在沒有衝突的情況下連續佈置其寫入，必須將路由資訊的完整知識分發到每個 rank。因此，我們的主機分發核心首先交換每個專家的 token 計數，此後它們各自可以確定其寫入的位置以及它們從哪個 rank 接收 token 的偏移量。

單純交換路由資訊會導致不必要的的高延遲，因為每個 rank 都必須等待寫入到達，然後才能分發自己的任何寫入。相反，我們在每個接收端 rank 上保留少量的私有每發送者空間，並在路由資訊的寫入旁邊分發一些 token。這確保了在整個分發過程的大部分時間內都能充分利用頻寬。交換路由遠端後，其餘的 token 透過每個遠端對等點的單次寫入進行傳輸。路由資訊的交換和固定大小的緩衝區，以及填充它們的傳輸，如下圖所示。

MoE 分派的緩衝區佈局，顯示 token 資料與路由中繼資料如何排列，以及如何在兩個 rank 之間交換。

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

對於 NVLink，節點內的發送和接收緩衝區會被對應到對等點中，以便它們能夠直接從相鄰的 rank 推送或拉取 token。

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

工作從分發發送開始，核心聚合路由資訊並透過單個區塊內共用記憶體中的原子累積來計算每專家的 token 計數。一旦 token 計數可用，它們就會被複製到主機記憶體，並且立即通知代理開始將路由資訊分發到所有其他對等 rank。同時，分發繼續計算 token 計數的累積總和，以識別將 token 打包到連續的每對等發送緩衝區所需的偏移量。一旦偏移量可用，所有要透過 EFA 發送的 token 都會被打包到發送緩衝區中。代理收到通知，將寫入發佈到節點間對等點的發送緩衝區中。在節點內部，核心還立即透過 NVLink 將 token 推送到同一節點內的私有緩衝區中。

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

一旦一個 rank 從所有對等點接收到路由資訊，它就會集中該資訊並計算識別用於分發和合併的 token 所需的各種偏移量。這項工作是在傳輸已經掛起時在 CPU 上完成的，因此它不是延遲關鍵型。隨後，一旦所有偏移量已知，就會發佈超出私有緩衝區大小的其餘 token 的寫入，並且分發的發送階段就完成了。

在接收端，核心首先透過 NVLink 從超出私有緩衝區容量的同一節點拉取 token。接下來，它會阻塞，直到主機代理發出所有傳輸完成的信號。接收到信號後，token 會從接收緩衝區打亂到傳遞給分組 GEMM 的輸出張量中。由於緩衝區在分發和合併之間重複使用，因此 NVLink 和 RDMA 屏障 (barriers) 確保在允許下一個核心覆寫它們之前已讀取接收緩衝區。

合併核心遵循大致相似的模式，打亂 token、向代理發送信號並等待完成。然而，用於發佈其寫入的中繼資料隨手可得，並且所有遠端偏移量都是已知的，因此所有在節點內傳輸的 token 都透過 NVLink 推送到接收器，而節點間的傳輸則透過單次寫入進行。在返回結果之前，核心累積來自接收緩衝區的 token，將其寫入輸出張量中。

為了對我們的實現進行原型設計，我們最初重新使用了為分散式預填充 (prefill) 和解碼而建構的 Transfer Engine，以便發出和追蹤 RDMA 寫入。然而，與每個節點向同一個對等點發布大量寫入的分散式解碼不同，MoE 分發是在一次操作中向所有對等點發布小型傳輸。開箱即用，效能並非最佳，我們的實現落後 ConnectX-7 上的 DeepEP 約 20 微秒。因此，我們專門化了 TransferEngine 以支援兩個操作（scatter 和 barrier），從而優化與大量對等點的互動。

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

儘管最初我們僅在使用 libfabric 的 EFA 上開發 TransferEngine，但在使用 libibverbs 在 TransferEngine 中新增 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 轉接卡以及標準 DeepSeek-V3 和 Kimi-K2 配置（具有 7168 的隱藏維度、區塊縮放的 fp8 分發和 bf16 合併）的 H200 系統上評估了分發和合併延遲。我們考慮分發的批次大小為 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。這主要是因為我們僅在合併的接收端累積 token，而 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 配置。

我們也考慮了微批次，儘管它僅在非常高的批次大小或使用一個以上的草稿 token 時才改善延遲。為了從將分組的 GEMM 問題一分為二中獲益，token 批次大小必須足夠密集。大部分時間仍花在 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 吞吐量

未來工作

參考資料

引用