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

高層次設(shè)計(jì)介紹

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

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

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

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

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

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

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

CPU 側(cè)

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

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

GPU 側(cè)

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

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

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

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

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

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

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

批量方案

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

分層方案

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

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

圖 3. 系統(tǒng)高層次拓?fù)浣Y(jié)構(gòu)

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

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

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

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

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

64 個(gè) GPU，每個(gè) GPU 分配 4 個(gè)專家

58 層，每個(gè) GPU 分配 232 個(gè)專家

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

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

64 或 72 個(gè) GPU，每個(gè) GPU 分配 5 個(gè)專家

58 層，每個(gè) GPU 分配 290 個(gè)專家

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

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

36 個(gè) GPU，每個(gè) GPU 分配 8 個(gè)專家

58 層，每個(gè) GPU 分配 464 個(gè)專家

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

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

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

EP 通信內(nèi)核

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

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

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

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

考慮因素：

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

當(dāng)我們開始在 Hopper 上啟用大規(guī)模 EP 工作時(shí)，我們得出的結(jié)論是 DeepEP 可以適應(yīng)并滿足我們?cè)谠撈脚_(tái)上的需求。

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

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

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

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

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

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

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

作者

楊東旭

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

喬顯杰

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

謝開宇

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

朱恩偉

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

陳曉明

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