之前文章已介紹引入大規(guī)模 EP 的初衷，本篇將繼續(xù)深入介紹 TensorRT-LLM 的大規(guī)模專家并行架構(gòu)設(shè)計與創(chuàng)新實現(xiàn)。

高層次設(shè)計介紹

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

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

更多數(shù)據(jù)將從熱門 EP rank 被發(fā)送和接收。

這些問題將導(dǎo)致系統(tǒng)級擁塞效應(yīng)，即熱門 EP rank 將延遲整體端到端執(zhí)行。

為確保大規(guī)模 EP 能穩(wěn)定運行，需通過精心設(shè)計盡可能減少 EP 失衡問題。整體設(shè)計如下：

圖 1. TensorRT-LLM 大規(guī)模 EP 的高層次設(shè)計

此設(shè)計同時包含 CPU 和 GPU 兩側(cè)邏輯：

CPU 側(cè)

使用復(fù)制與放置算法(復(fù)制與放置計算組件)實現(xiàn)更均衡的 EP 策略。這些算法是經(jīng)典算法，更適合 CPU 計算。此外，將此計算卸載至 CPU 可減少對 GPU 的干擾。未來可探索基于機器學(xué)習(xí)的算法，并可能需要額外設(shè)計考量。復(fù)制與放置計算組件將生成“放置信息”，該信息將被 GPU 路由邏輯和 CPU 更新權(quán)重與放置組件共同使用。由 GPU 上運行的統(tǒng)計組件生成的統(tǒng)計數(shù)據(jù)將被用作復(fù)制與放置計算組件的輸入。

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

GPU 側(cè)

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

EP 通信內(nèi)核，在上篇圖 11 中為分發(fā)和合并組件。

在線流量數(shù)據(jù)統(tǒng)計采集器(統(tǒng)計組件)。該組件采集統(tǒng)計數(shù)據(jù)供復(fù)制與放置計算組件使用。

MoE 路由邏輯(路由組件)。該組件將 Token 發(fā)送至激活的專家，并且需要進行調(diào)整以支持 MoE 權(quán)重的動態(tài)放置。它使用復(fù)制與放置計算組件生成的放置信息。

MoE 計算邏輯 (MoE 組件) 也需進行相應(yīng)調(diào)整。

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

我們?yōu)楦聶?quán)重與放置組件提供了兩種設(shè)計方案：

批量方案

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

分層方案

圖 2. 分層 MoE 權(quán)重重新分配示例

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

圖 3. 系統(tǒng)高層次拓撲結(jié)構(gòu)

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

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

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

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

用例 1(專家分配均衡，不進行專家復(fù)制)

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

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

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

用例 2(專家分配均衡并進行復(fù)制)

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

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

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

用例 3(專家分配均衡并進行復(fù)制)

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

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

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

綜上所述，根據(jù)理論分析，采用 Grace CPU 內(nèi)存作為存儲完整大小 MoE 權(quán)重的池，應(yīng)能使我們在 5 次解碼迭代內(nèi)實現(xiàn) EP(專家并行)的再平衡。如果將要求放寬至 10 次或以上迭代，系統(tǒng)實現(xiàn)將變得更加靈活。

接下來我們將介紹大規(guī)模 EP 系統(tǒng)的詳細實現(xiàn)方式。

EP 通信內(nèi)核

我們評估了多種實現(xiàn)大規(guī)模 EP 所需 EP 通信內(nèi)核的途徑，包括 DeepEP、其他解決方案以及重新開發(fā)一種方法。

當(dāng)前的技術(shù)決策是：

我們實現(xiàn)了一組新的自定義 EP 通信內(nèi)核。

對于其他系統(tǒng)(如 Hopper)，我們選擇直接集成 DeepEP 并進行一些可能的增強。

考慮因素：

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

當(dāng)我們開始在 Hopper 上啟用大規(guī)模 EP 工作時，我們得出的結(jié)論是 DeepEP 可以適應(yīng)并滿足我們在該平臺上的需求。

我們也在積極評估將通信內(nèi)核整合為單一解決方案以簡化系統(tǒng)架構(gòu)的可能性，并將持續(xù)向社區(qū)更新進展。接下來，我們將進一步探討自定義 EP 通信內(nèi)核實現(xiàn)中引入的優(yōu)化措施。

在系統(tǒng)中引入 EP 通信內(nèi)核的初衷

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

EP 通信內(nèi)核的實現(xiàn)

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

我們的實現(xiàn)方式是使用 CUDA 的驅(qū)動程序 API 通過 MNNVL 建立點對點 (P2P) 緩沖區(qū)作為工作區(qū)。每個 GPU 都可以訪問其他 GPU 的工作區(qū)。工作區(qū)被劃分為多個通道，每個通道分配給遠程 GPU 作為寫入緩沖區(qū)。這些寫入緩沖區(qū)以 FIFO 方式使用，通過標志同步 FIFO 狀態(tài)以避免數(shù)據(jù)損壞。

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

作者

楊東旭

現(xiàn)任職于 NVIDIA Compute Arch 部門。主要負責(zé) LLM 推理系統(tǒng)的開發(fā)和性能優(yōu)化。加入 NVIDIA 之前，曾從事搜索系統(tǒng)的 GPU 加速和開發(fā)工作。

喬顯杰

NVIDIA Compute Arch 部門高級架構(gòu)師，主要負責(zé) LLM 推理的性能評估和優(yōu)化。加入 NVIDIA 之前，他曾從事推薦系統(tǒng)的 GPU 加速研發(fā)工作。

謝開宇

NVIDIA Compute Arch 部門高級架構(gòu)師，主要負責(zé) TensorRT-LLM 項目的開發(fā)，專注在系統(tǒng)性能和優(yōu)化工作。

朱恩偉

NVIDIA DevTech 部門高級工程師，主要負責(zé) TensorRT-LLM 項目的開發(fā)和性能優(yōu)化。

陳曉明

NVIDIA Compute Arch 部門的首席架構(gòu)師和高級經(jīng)理，對深度學(xué)習(xí)模型的算法軟硬件協(xié)同設(shè)計感興趣，最近從事大語言模型推理的性能建模、分析和優(yōu)化。