近年來，大語言邏輯推理模型取得了顯著進(jìn)步，但也帶來了新的部署挑戰(zhàn)。其中，因復(fù)雜的“思考與邏輯推理”過程而引起的輸出序列長度 (OSL) 的加長已成為一大難題。OSL 的加長提高了對 token 間延遲 (Token-to-Token Latency, TTL) 的要求，往往會引發(fā)并發(fā)限制。在最極端的情況下，實(shí)時應(yīng)用會面臨單并發(fā)（最小延遲場景）這一特別棘手的問題。

本文將探討NVIDIATensorRT-LLM如何基于 8 個 NVIDIA Blackwell GPU 的配置，打破 DeepSeek-R1 在最小延遲場景中的性能紀(jì)錄：在 GTC 2025 前將 67 token / 秒 (TPS) 的速度提升至 253 TPS（提速3.7 倍），而目前這一速度已達(dá) 368 TPS（提速5.5 倍）。

實(shí)現(xiàn)配置

一、工作負(fù)載配置文件

輸入序列長度 (ISL)：1000 token

輸出序列長度 (OSL)：2000 token

二、模型架構(gòu)

DeepSeek-R1 的基礎(chǔ)主模型包含：3 個密集層（初始）和 58 個 MoE 層，此外還有 1 個多 token 預(yù)測 (Multi-Tokens Prediction, MTP) 層（相當(dāng)于 MoE 架構(gòu)）用于推測性解碼。我們的優(yōu)化配置將 MTP 層擴(kuò)展成 3 個層，采用自回歸方法探索其最大性能。

圖1: DeepSeek-R1 的基礎(chǔ)主模型

該圖片來源于 Github: Pushing Latency Boundaries: Optimizing DeepSeek-R1 Performance on NVIDIA Blackwell GPUs 一文，若您有任何疑問或需要使用該圖片，請聯(lián)系該文作者

三、精度策略

我們探索出了一種能夠更好平衡準(zhǔn)確度與性能的混合精度方案。

* TensorRT-LLM 已支持 FP8 Attention。但在該延遲場景下，低精度注意力計算并不能提升性能，因此我們?yōu)樽⒁饬δK選擇了 BF16 精度。

** NVFP4 模型檢查點(diǎn)由 NVIDIA TensorRT 模型優(yōu)化器套件生成。

*** RouterGEMM 使用 BF16 輸入 / 權(quán)重與 FP32 輸出來確保數(shù)值的穩(wěn)定性

四、并行策略

我們還在 8 個 Blackwell GPU 上嘗試并引入了混合并行策略。具體而言，該延遲場景的最佳策略為 “TP8EP2”，其定義如下：

五、一圖整合

現(xiàn)在，我們將所有內(nèi)容整合成一張圖，該圖表示的是解碼迭代中的一個 MoE 層。

該圖片來源于 Github: Pushing Latency Boundaries: Optimizing DeepSeek-R1 Performance on NVIDIA Blackwell GPUs 一文，若您有任何疑問或需要使用該圖片，請聯(lián)系該文作者

圖中的模塊包括：

• 輸入模塊：一個形狀為 [m, 7168] 的 BF16 張量，其中 m 表示 token 數(shù)量（例如使用 3 個 MTP 層時 m = 4），7168 為模型的隱藏大小。

• 模塊 1：Fuse_A_GEMM 拼接 WDQ、WDKV 和 WKR 的權(quán)重，以減少內(nèi)核調(diào)用開銷。

• 模塊 2：2 個 RMSNorm 對 Q / K 張量進(jìn)行歸一化。這些張量可以重疊在多個流上，也可以合并成單個分組 RMSNorm。

• 模塊 3：UQ_QR_GEMM 拼接 WUQ 和 WQR 的權(quán)重，以減少內(nèi)核調(diào)用開銷。

• 模塊 4：UK_BGEMM 在批量 GEMM 中使用 WUK。為防止權(quán)重規(guī)模膨脹和產(chǎn)生新的加載成本，我們未加入模塊 3 和 4。

• 模塊 5：Concat KVCache & applyRope 合并 K / V 緩存并應(yīng)用 ROPE（旋轉(zhuǎn)位置編碼）。

• 模塊 6：genAttention 在生成階段執(zhí)行 MLA，作用類似于 num_q_heads = 128 / TP8 = 16 的 MQA

• 模塊 7：UV_GEMM 執(zhí)行帶 WUV 權(quán)重的批量 GEMM。

• 模塊 8：WO_GEMM 使用 WO 權(quán)重運(yùn)行密集 GEMM。為避免增加權(quán)重加載的開銷，我們未加入模塊 7 和 8。

• 模塊 9：融合內(nèi)核將 oneshotAllReduce、Add_RMSNorm 和 DynamicQuant (BF16->NVFP4) 整合到單個內(nèi)核中。

• 模塊 10：routerGEMM & topK 處理路由器 GEMM (Router GEMM) 和 topK 選擇。

• 模塊 11：共享專家模型與模塊 10 和模塊 12 部分重疊。

• 模塊 12：稀疏專家模型通過分組 GEMM (Grouped GEMM) 實(shí)現(xiàn)專家層。

• 模塊 13：最終融合內(nèi)核同時執(zhí)行 localReduction、oneshotAllReduce 和 Add_RMSNorm 操作。

主要優(yōu)化

一、系統(tǒng)級優(yōu)化

1、CUDA Graph 與可編程依賴啟動

CUDA Graph 對于克服小型工作負(fù)載中的 CPU 開銷必不可少，而可編程依賴啟動可進(jìn)一步降低內(nèi)核啟動延遲。

2、MTP

基于 MTP 的兩種優(yōu)化措施：

1) 自回歸 MTP 層

根據(jù)我們的研究結(jié)果，3x MTP 層的配置性能最佳。

2) 寬松接受驗(yàn)證

邏輯推理模型 (如 DeepSeek R1) 的生成過程可以分為兩個階段：思考階段和實(shí)際輸出階段。在思考階段，如果啟用寬松接受 (Relax Acceptance) 模式，候選 token 處于候選集時即可被接受。該候選集基于 logits topN 和概率閾值生成。

• topN：從 logits 中采樣前 N 個 token。

• 概率閾值：基于 topN 個候選 token，只有概率大于 Top1 的概率減去 delta 的 token 時可保留在候選集。

在非思考階段，我們?nèi)圆捎脟?yán)格接受模式。

這是一種寬松的驗(yàn)證和比較方法，可以在對精度影響很小的情況下，提升接受率并帶來加速。

如需了解更多信息，請?jiān)L問：

multi-token-prediction-mtp

3、多流

我們引入了基于多流的優(yōu)化措施以隱藏部分內(nèi)核的開銷，例如：

• 將共享專家模型與稀疏專家模型重疊

• 將 Concat_KVCache 內(nèi)核與 GEMM 重疊

稀疏專家模型作為 GEMM (僅當(dāng) moe_backend=CUTLASS 時有效)

該圖片來源于 Github: Pushing Latency Boundaries: Optimizing DeepSeek-R1 Performance on NVIDIA Blackwell GPUs 一文，若您有任何疑問或需要使用該圖片，請聯(lián)系該文作者

現(xiàn)有的基于 CUTLASS 的稀疏專家模型流（如圖所示）將輸入的 token 分發(fā)到指定的專家模型，然后在每個專家模型的輸出上進(jìn)行索引式的局部歸約，最后進(jìn)行全局AllReduce。分發(fā)和索引局部歸約在低延遲場景下會產(chǎn)生高開銷。為解決此問題，我們提出將“稀疏專家模型作為 GEMM”處理，即將所有 token 發(fā)送至每個激活的專家模型，并在局部歸約前屏蔽不需要的輸出。由于分組 GEMM 受顯存限制，冗余 token 產(chǎn)生的額外計算開銷幾乎沒有影響，有效避免了昂貴的分發(fā)，同時減少開銷。

4、重新平衡稀疏專家模型

稀疏專家模型常用的并行化策略有兩種：專家并行 (EP) 和張量并行 (TP)。專家并行 (EP) 將每個專家模型分配到獨(dú)立的 GPU，以此實(shí)現(xiàn)高顯存和計算效率。但 token 放置依賴于數(shù)據(jù)，導(dǎo)致 GPU 間工作負(fù)載分布不均，并在 MoE 模塊后的 AllReduce 步驟中顯示額外開銷。張量并行 (TP) 將每個專家模型均勻劃分到多個 GPU，雖平衡了工作負(fù)載，但卻犧牲了數(shù)學(xué) / 顯存效率。

• 混合 ETP

結(jié)合 EP / TP 的混合方法可緩解上述問題。實(shí)驗(yàn)結(jié)果表明，TP4EP2 配置在實(shí)際中表現(xiàn)最佳。

• 智能路由器

另一方案是將所有專家模型權(quán)重存儲在由 4 個 GPU 組成的集群中，隨后將其復(fù)制到另一個 4 GPU 集群，智能路由器可將 token 動態(tài)地分配到各集群。該設(shè)計在不顯著影響本地顯存和計算效率的前提下，保持了工作負(fù)載分布的平衡。

二、內(nèi)核級優(yōu)化

1、注意力內(nèi)核

我們開發(fā)了定制的 MLA 注意力內(nèi)核，以便更好地使用 GPU 資源應(yīng)對延遲場景。

2、分組 GEMM

• CUTLASS 后端（默認(rèn)后端）

我們的默認(rèn) MoE 后端基于 CUTLASS，該后端具有靈活性和穩(wěn)定性，但可能不是最佳的性能方案。

• TensorRT-LLM 后端

另一個 MoE 后端是 TensorRT-LLM，其性能更優(yōu)。我們正在努力提高其靈活性和穩(wěn)定性，未來將作為延遲場景中分組 GEMM 計算的默認(rèn)后端。

3、通信內(nèi)核

對于小規(guī)模消息，受常規(guī) NCCL 延遲影響的 AllReduce 內(nèi)核效率低下，為此我們開發(fā)了一款定制化的一次性 AllReduce 內(nèi)核。該內(nèi)核通過先模仿初始廣播，然后進(jìn)行局部歸約的方式，利用 NVSwitch 的強(qiáng)大硬件能力在最小延遲場景中實(shí)現(xiàn)了更優(yōu)的性能。

4、密集 GEMM 優(yōu)化

我們重點(diǎn)優(yōu)化兩種密集 GEMM：Fuse_A_GEMM 和 RouterGEMM。因?yàn)檫@兩種 GEMM 占據(jù)了大部分執(zhí)行時間、顯存效率低下且難以分片（兩者均基于 DP）。

• Fuse_A_GEMM

我們開發(fā)了一個定制的 Fuse_A_GEMM，通過將大部分權(quán)重預(yù)先載入到共享顯存（通過 PDL 實(shí)現(xiàn)并與 oneshot-AllReduce 重疊），大幅提升了性能。當(dāng) num_tokens < 16 時，該內(nèi)核性能較默認(rèn)的 GEMM 實(shí)現(xiàn)有明顯提升。

該圖片來源于 Github: Pushing Latency Boundaries: Optimizing DeepSeek-R1 Performance on NVIDIA Blackwell GPUs 一文，若您有任何疑問或需要使用該圖片，請聯(lián)系該文作者

• RouterGEMM

我們通過使用內(nèi)部的 AI 代碼生成器，自動生成經(jīng)過優(yōu)化的 RouterGEMM 內(nèi)核。在 num_tokens ≤ 30 時，該內(nèi)核性能較默認(rèn)的 GEMM 實(shí)現(xiàn)有顯著提升。

該圖片來源于 Github: Pushing Latency Boundaries: Optimizing DeepSeek-R1 Performance on NVIDIA Blackwell GPUs 一文，若您有任何疑問或需要使用該圖片，請聯(lián)系該文作者

5、內(nèi)核融合

為了減少最小延遲場景中額外的全局顯存寫讀開銷，內(nèi)核融合必不可少。我們目前支持以下融合模式：

• 將兩個重疊的 RMS_Norm 融合成一個 GroupedRMSNorm

• 將 (LocalReduction) + AR + RMS_Norm + (Dynamic_Quant_BF16toNVFP4) 融合成一個內(nèi)核

• 將 Grouped GEMM_FC1 + 點(diǎn)激活 (當(dāng) moe_backend=TRTLLM 時) 融合成一個內(nèi)核

如何復(fù)現(xiàn)

https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/blogs/Best_perf_practice_on_DeepSeek-R1_in_TensorRT-LLM.md#b200-min-latency

需要注意的是，寬松接受模式是 Deepseek-R1 模型的特有模式。若要啟用該模式，需在準(zhǔn)備基準(zhǔn)數(shù)據(jù)集時設(shè)置 add_generation_prompt = True，示例代碼如下：

input_ids= tokenizer.encode(tokenizer.apply_chat_template(msg, tokenize=False, add_generation_prompt=True), add_special_tokens=False)

還需在 speculative_config 中設(shè)置 use_relaxed_acceptance_for_thinking: true, relaxed_topk: 10 和 relaxed_delta: 0.6。

后續(xù)工作

• 增加融合

• 增加重疊

• 增加對注意力內(nèi)核的優(yōu)化

• 增加對 MTP 的研究

結(jié)語

在延遲敏感型應(yīng)用中突破 DeepSeek R1 的性能極限是一項(xiàng)非凡的工程。本文詳細(xì)介紹的優(yōu)化措施是整個 AI 技術(shù)棧各個領(lǐng)域的協(xié)作成果，涵蓋了內(nèi)核級優(yōu)化、運(yùn)行時增強(qiáng)、模型量化技術(shù)、算法改進(jìn)以及系統(tǒng)性能分析與調(diào)優(yōu)。希望本文介紹的技術(shù)和最佳實(shí)踐，能夠幫助開發(fā)者社區(qū)在任務(wù)關(guān)鍵型 LLM 推理應(yīng)用中更充分地發(fā)揮 NVIDIA GPU 的性能。