TensorRT-LLM 的大規模專家并行架構設計與創新實現

之前文章已介紹引入大規模 EP 的初衷，本篇將繼續深入介紹 TensorRT-LLM 的大規模專家并行架構設計與創新實現。

上篇文章參考，點擊閱讀：

大規模專家并行 (EP) 在 TensorRT-LLM 的設計動機與系統分析

高層次設計介紹

根據引入大規模 EP 的初衷部分的詳細分析與研究，可以明確觀察到 EP 中的專家失衡是大規模 EP 的常見模式。這種 EP 失衡會通過以下方式顯著降低整體系統性能：

• 熱門 EP rank 將消耗更多顯存（用于激活值），這會限制推理過程中調度的有效最大批處理大小。

• 更多數據將從熱門 EP rank 被發送和接收。

這些問題將導致系統級擁塞效應，即熱門 EP rank 將延遲整體端到端執行。

為確保大規模 EP 能穩定運行，需通過精心設計盡可能減少 EP 失衡問題。整體設計如下：

圖 1. TensorRT-LLM 大規模 EP 的高層次設計

此設計同時包含 CPU 和 GPU 兩側邏輯：

• CPU 側

• 使用復制與放置算法（復制與放置計算組件）實現更均衡的 EP 策略。這些算法是經典算法，更適合 CPU 計算。此外，將此計算卸載至 CPU 可減少對 GPU 的干擾。未來可探索基于機器學習的算法，并可能需要額外設計考量。復制與放置計算組件將生成“放置信息”，該信息將被 GPU路由邏輯和 CPU更新權重與放置組件共同使用。由 GPU 上運行的統計組件生成的統計數據將被用作復制與放置計算組件的輸入。

• 編排流程（更新權重與放置組件）將 MoE 權重從 CPU 內存更新并重新加載到 GPU 設備顯存。該組件還將使用由復制與放置計算組件生成的放置信息。我們的可擴展設計允許通過 MNNVL 或 NIC 從遠程 GPU 顯存重新加載 MoE 權重。

• GPU 側

• 這是推理的主要執行工作流。我們在設計中引入了以下新的 GPU 組件：

• EP 通信內核，在上篇圖 11 中為分發合并組件。

• 在線流量數據統計采集器（統計組件）。該組件采集統計數據復制與放置計算組件使用。

• MoE 路由邏輯（路由組件）。該組件將 Token 發送至激活的專家，并且需要進行調整以支持 MoE 權重的動態放置。它使用復制與放置計算組件生成的放置信息

• MoE 計算邏輯 (MoE組件) 也需進行相應調整。

• CPU 和 GPU 組件之間需要仔細同步，以確保整個執行過程的有效性，尤其是為了避免卡頓以及無效或次優執行。

我們為更新權重與放置組件提供了兩種設計方案：

• 批量方案

• 在此方案中，當 MoE 權重重新分配邏輯啟動時，當前服務實例上的推理過程將不得不暫停，直至 MoE 權重重新分配過程完成。我們估計這可能導致約0.5 至 1 秒的在線服務暫停，最壞情況下會引發請求超時。此類超時或暫停可通過系統級措施來緩解，例如將請求傳送至其他服務實例或通過請求重試來應對。

• 分層方案

圖 2. 分層 MoE 權重重新分配示例

在當前系統中，我們選擇采用分層方案以盡量減少對在線用戶體驗的影響。批量方案應更易于實現，但本文將不再討論。為了正確實現分層方案，需仔細評估不同底層硬件的性能以確定具體實現方案。圖 3 展示了系統節點中不同硬件組件的通信帶寬。

圖 3. 系統高層次拓撲結構

以 DeepSeek R1 模型為例，采用 FP4 精度時，每個 MoE 專家占用 24MiB 顯存空間。每層包含 256 個專家，總共包含 58 個 MoE 層加 1 個 MTP 層。因此，為實現 EP 平衡所需重新分配的 MoE 權重最大總量為 348GiB。每個節點為每個 Grace CPU 提供 480GB LPDDR5X 顯存。在 NUMA 域內，總計可提供 960GB Host 顯存。一個節點可在其 CPU Host 顯存中完整承載如 DeepSeek R1 LLM 等模型的全部 MoE 權重。基于此，MoE 權重重新分配可通過將對應的 MoE 權重從 CPU Host 顯存移動至 GPU 設備顯存來實現。

假設我們將50ms的跨 Token 延遲 (ITL) 作為主要延遲約束。通過粗略估算，可以計算出在每次解碼迭代中，可從 MoE 權重池（可保存在 Grace CPU 顯存或另一節點上的 GPU 顯存中）移動到 Blackwell GPU（用于實際 MoE 推理）的專家權重數量為：

圖 4. 在以下 50ms ITL 限制下，每次迭代理論上需要更新的專家數量（使用不同硬件作為存儲完整 MoE 權重的池）

基于此分析，若依賴每個節點上的 Grace CPU 內存來存儲 MoE 權重池，則每次解碼迭代中，最多可將 300 個專家的權重重新分配至同一節點上的每個 GPU。假設目標是在 5 次解碼迭代內完成整個模型 MoE 權重再平衡，以下為具體用例研究：

• 用例 1（專家分配均衡，不進行專家復制）

• 64 個 GPU，每個 GPU 分配 4 個專家

• 58 層，每個 GPU 分配 232 個專家

• 每次迭代需要 47 次專家更新，所有方法均可滿足延遲目標。

• 用例 2（專家分配均衡并進行復制）

• 64 或 72 個 GPU，每個 GPU 分配 5 個專家

• 58 層，每個 GPU 分配 290 個專家

• 每次迭代需要 58 次專家更新，所有方法均可滿足延遲目標。

• 用例 3（專家分配均衡并進行復制）

• 36 個 GPU，每個 GPU 分配 8 個專家

• 58 層，每個 GPU 分配 464 個專家

• 每次迭代需要 93 次專家更新，所有方法均可滿足延遲目標。

綜上所述，根據理論分析，采用 Grace CPU 內存作為存儲完整大小 MoE 權重的池，應能使我們在 5 次解碼迭代內實現 EP（專家并行）的再平衡。如果將要求放寬至 10 次或以上迭代，系統實現將變得更加靈活。

接下來我們將介紹大規模 EP 系統的詳細實現方式。

EP 通信內核

我們評估了多種實現大規模 EP 所需 EP 通信內核的途徑，包括 DeepEP、其他解決方案以及重新開發一種方法。

當前的技術決策是：

• 我們實現了一組新的自定義 EP 通信內核。

• 對于其他系統（如 Hopper），我們選擇直接集成 DeepEP 并進行一些可能的增強。

考慮因素：

• DeepEP 是由 DeepSeek 團隊完成的一項出色成果。我們在啟動 TensorRT-LLM 大規模 EP 工作時，最初把重點放在 Grace Blackwell 機架式系統上。我們選擇實現自己的定制 EP 通信內核，因為這更便于引入需要 Grace Blackwell 機架式系統功能的優化措施。

• 當我們開始在 Hopper 上啟用大規模 EP 工作時，我們得出的結論是 DeepEP 可以適應并滿足我們在該平臺上的需求。

我們也在積極評估將通信內核整合為單一解決方案以簡化系統架構的可能性，并將持續向社區更新進展。接下來，我們將進一步探討自定義 EP 通信內核實現中引入的優化措施。

在系統中引入 EP 通信內核的初衷

在解碼階段與預填充解碼 (PD) 分離的場景中，我們觀察到批處理大小可能不會很大，因此延遲成為一個重要考慮因素。在此背景下，我們非常需要實現與 CUDA graph 的兼容。NCCL 是一個優秀的 GPU 通信庫，為我們提供了高效的通信內核和基本操作。目前，其 Send 和 Recv 操作在調用 ncclSend / ncclRecv 時，需要顯式指定數據大小。但在大規模專家并行 (large-EP) 場景中，待傳輸的數據大小根據模型在每次迭代中的輸出動態確定。當前 NCCL 通信接口需要同步將通信大小發回 CPU，并以對應數據大小從 CPU 發起 NCCL 調用。這將破壞 CUDA graph 兼容性。這一限制迫使我們開發與 CUDA graph 兼容，且能直接從 GPU 顯存接受通信大小的高性能通信內核。我們還希望這些內核能夠充分利用 MNNVL 的顯存帶寬。

EP 通信內核的實現

我們的內核采用與 NCCL 的 LL128 原語類似的通信方法。由于這種方法在延遲和帶寬之間取得了良好的平衡，因此非常適合 LLM 推理。我們的自定義內核可直接從 GPU 顯存讀取通信大小并兼容 CUDA graph，即使數據大小在不同運行中變化也不例外。

我們的實現方式是使用 CUDA 的驅動程序 API 通過 MNNVL 建立點對點 (P2P) 緩沖區作為工作區。每個 GPU 都可以訪問其他 GPU 的工作區。工作區被劃分為多個通道，每個通道分配給遠程 GPU 作為寫入緩沖區。這些寫入緩沖區以 FIFO 方式使用，通過標志同步 FIFO 狀態以避免數據損壞。詳細信息請參見 PR 3504:

https://github.com/NVIDIA/TensorRT-LLM/pull/3504

下一篇我們將繼續介紹 TensorRT-LLM 在線負載均衡策略與實測的解析。

作者

楊東旭

現任職于 NVIDIA Compute Arch 部門。主要負責 LLM 推理系統的開發和性能優化。加入 NVIDIA 之前，曾從事搜索系統的 GPU 加速和開發工作。

喬顯杰

NVIDIA Compute Arch 部門高級架構師，主要負責 LLM 推理的性能評估和優化。加入 NVIDIA 之前，他曾從事推薦系統的 GPU 加速研發工作。

謝開宇

NVIDIA Compute Arch 部門高級架構師，主要負責 TensorRT-LLM 項目的開發，專注在系統性能和優化工作。

朱恩偉

NVIDIA DevTech 部門高級工程師，主要負責 TensorRT-LLM 項目的開發和性能優化。

陳曉明

NVIDIA Compute Arch 部門的首席架構師和高級經理，對深度學習模型的算法軟硬件協同設計感興趣，最近從事大語言模型推理的性能建模、分析和優化。