前言

開源 DeepSeek R1 模型的創(chuàng)新架構(gòu)包含多頭潛在注意力機(jī)制 (MLA) 和大型稀疏混合專家模型 (MoE)，其顯著提升了大語言模型 (LLM) 的推理效率。但要充分發(fā)揮這種創(chuàng)新架構(gòu)的潛力，軟硬件的協(xié)同優(yōu)化也至關(guān)重要。本文將深入解析 NVIDIA 在基于 Blackwell GPU 的TensorRT-LLM框架內(nèi)為 DeepSeek R1 吞吐量優(yōu)化場景 (TPS / GPU) 開發(fā)的優(yōu)化策略。文內(nèi)將詳細(xì)闡述各項(xiàng)優(yōu)化措施的設(shè)計(jì)思路。另一篇關(guān)于降低延遲的博客（見下）已詳細(xì)解釋了 TensorRT-LLM 如何通過提升 R1 性能實(shí)現(xiàn)最佳 TPS / USER。

這些優(yōu)化措施顯著提升了 DeepSeek R1 在 Blackwell 上的吞吐量。其在 ISL / OSL 1K / 2K 數(shù)據(jù)集上的性能從 2 月的約 2000 TPS / GPU 提升至 4600 TPS / GPU。這些優(yōu)化措施具有通用性，可適用于其他 ISL / OSL 配置。其大致分為三個(gè)方面：MLA 層、MoE 層和運(yùn)行時(shí)。

精度策略

DeepSeek R1 吞吐量場景的混合精度策略與延遲優(yōu)化場景的策略基本一致，具體差異如下：

使用 FP8 KV 緩存和 FP8 注意力機(jī)制，而非 BF16 精度。

使用 FP4 Allgather 提高通信帶寬利用率。

本文中使用的 Checkpoint 位于 nvidia / DeepSeek-R1-FP4，由 NVIDIA 模型優(yōu)化器生成。常用數(shù)據(jù)集在該 FP4 Checkpoint 和 TensorRT-LLM 執(zhí)行中準(zhǔn)確率分?jǐn)?shù)為：

** 注意這些評(píng)估結(jié)果存在運(yùn)行間波動(dòng)，因此 FP4 數(shù)據(jù)的分?jǐn)?shù)略高。我們認(rèn)為 FP4 在這些數(shù)據(jù)集上的精度與 FP8 相當(dāng)。

該 Checkpoint 中的 MoE 層已量化為 FP4。將 MoE 層權(quán)重量化為 FP4 具有以下優(yōu)點(diǎn)：

充分利用 NVIDIA Blackwell GPU 第五代 Tensor Core 的 FLOPS（每秒浮點(diǎn)運(yùn)算）

將 MoE 權(quán)重所需的顯存負(fù)載減少近一半。在該場景中，MoE 部分在 Decoding 階段仍受顯存限制，且 DeepSeek R1 模型中 97% 的權(quán)重來自 MoE 層。

減少模型權(quán)重的顯存占用，從而釋放更多 GPU 顯存用于 KV 緩存，進(jìn)而增加最大并發(fā)數(shù)。DeepSeek R1 模型的原始 FP8 量化 Checkpoint 大小約為 640GB，而 NVIDIA 提供的 DeepSeek R1 FP4 量化模型僅為約 400GB。

我們?cè)?GSM8K 數(shù)據(jù)集上對(duì) FP8 KV 緩存和 FP8 注意力內(nèi)核的精度進(jìn)行了評(píng)估，未觀察到明顯的準(zhǔn)確率下降。有關(guān)準(zhǔn)確率數(shù)據(jù)，請(qǐng)參見 FP8 KV 緩存部分的表格。如果用戶在自己的數(shù)據(jù)集上觀察到準(zhǔn)確率差異，仍可選擇使用 BF16 KV 緩存和注意力機(jī)制。

并行策略

DeepSeek R1 吞吐量場景的并行策略與延遲場景的策略有所不同。

在接下來的部分，我們將解釋為何選擇數(shù)據(jù)并行 (DP) 和專家并行 (EP)，而非張量并行 (TP)。

權(quán)重吸收與 MQA

MLA 的核心理念是通過對(duì)注意力鍵 (K) 和值 (V) 進(jìn)行低秩聯(lián)合壓縮，減少推理過程中的 KV 緩存大小?；?MLA 公式，向下投影的 KV Latent 被向上投影到多個(gè)頭并與向上投影的 Q 結(jié)合，形成常規(guī)的多頭注意力機(jī)制 (MHA)。由于矩陣乘法的性質(zhì)，K 的向上投影權(quán)重矩陣 (W^UK) 可先與 Q 的向上投影權(quán)重矩陣 (W^Q) 相乘，再將結(jié)果與 Q 相乘。V 的向上投影權(quán)重矩陣 (W^UV) 與注意力輸出投影矩陣 W^O 也可在注意力輸出后相乘。DeepSeek-V2 技術(shù)報(bào)告將該技術(shù)稱為“吸收”。在權(quán)重被吸收后，MLA 等價(jià)于多查詢注意力 (MQA)。參見 DeepSeek-V2 技術(shù)論文獲取詳細(xì)公式和解釋，下圖所示的是 TensorRT-LLM 中權(quán)重吸收 MLA 的計(jì)算流程。

該圖片來源于 Github:Optimizing DeepSeek R1 Throughput on NVIDIA Blackwell GPUs: A Deep Dive for Developers一文，若您有任何疑問或需要使用該圖片，請(qǐng)聯(lián)系該文作者

在 Decoding 階段，權(quán)重吸收顯著減少向上投影 K 和 V 所需的數(shù)學(xué) FLOPS，這是因?yàn)檫@些 KV 向上投影所需的 FLOPS 與 KV 緩存長度成線性關(guān)系，而 Decoding 階段的 Q 向量長度始終為 1。KV 緩存歷史越長，所需的 FLOPS 越多，且由于僅保存了投影的 KV Latent，每個(gè) Decoding Token 都需要重復(fù)進(jìn)行向上投影，這進(jìn)一步增加了所需的 FLOPS。在 Prefill 階段，吸收權(quán)重的版本改變了 Q 和 KV 的維度，增加了注意力機(jī)制所需的 FLOPS。根據(jù) Roofline 分析，在輸入長度為 256 或以上的情況下，非吸收版本更有利于 Prefill 階段。TensorRT-LLM MLA 實(shí)現(xiàn)為 Prefill 和 Decoding 階段分別選擇了不同的高度優(yōu)化內(nèi)核，詳見 MLA。

注意力模塊數(shù)據(jù)并行 (ADP)

選擇注意力數(shù)據(jù)并行 (ADP) 的核心原因是對(duì) MQA（由不同 GPU 計(jì)算不同的注意力 Q 頭）采用張量并行 (TP) 會(huì)重復(fù)復(fù)制 KV 緩存顯存，從而限制系統(tǒng)能夠?qū)崿F(xiàn)的并發(fā)性。重復(fù)因子等于 TP 組大小，因此 TP8 時(shí)為 8 倍。并發(fā)性低會(huì)影響高性能系統(tǒng) (如 NVIDIA DGX B200) 的吞吐量。

對(duì)于使用 8 個(gè) B200 GPU 的 DeepSeek R1 FP4 Checkpoint，每個(gè) GPU 的權(quán)重和激活占用約 80GB 顯存，每個(gè) GPU 的可用 KV 緩存將為 100GB。假設(shè) ISL 為 1K，OSL 為 2K，每次請(qǐng)求將消耗約 200MB KV 緩存，使每個(gè) GPU 的最大并發(fā)數(shù)為 500。單節(jié)點(diǎn) 8 GPU 系統(tǒng)的全局并發(fā)性為 4000。使用 ATP 時(shí)，全局并發(fā)性將降至 500。

硅片實(shí)驗(yàn)表明，在保持所有其他因素不變的情況下，ADP 技術(shù)在最大吞吐量場景下可提速 400%。

MoE 專家并行 (EP)

DeepSeek R1 MoE 設(shè)計(jì)包含 256 個(gè)小型稀疏專家模型和 1 個(gè)共享專家模型，這些專家模型的 GEMM 問題規(guī)模如下。

這些專家模型可采用張量并行 (TP) 或?qū)＜也⑿?(EP) 的方式實(shí)現(xiàn)。目前的消融實(shí)驗(yàn)表明，由于 GEMM 問題規(guī)模更小，專家并行在 GEMM FLOPS 方面的表現(xiàn)更好。并且與 AllReduce 相比，專家并行可節(jié)省 GPU 通信帶寬，這是因?yàn)橹恍鑼?Token 發(fā)送至該 Token 對(duì)應(yīng)的活躍專家所在的 GPU，而張量并行 (TP) 則需對(duì)所有 GPU 間的 Token 執(zhí)行 AllReduce 操作。此外，為了將 DeepSeek R1 推理擴(kuò)展到 GB200 NVL72 等系統(tǒng)并充分利用聚合顯存帶寬和 Tensor Core FLOPS，需要使用大型專家并行。我們正在積極推進(jìn)該技術(shù)的實(shí)現(xiàn)。

硅片性能測(cè)試表明，在保持其他因素不變的情況下，專家并行 (EP) 可在 1K / 2K 最大吞吐量場景下將速度提升 142%。

MLA 層優(yōu)化

除了上述提到的并行策略和精度策略外，我們?cè)?MLA 模塊內(nèi)的層 / 內(nèi)核中進(jìn)行了以下優(yōu)化。

注意力內(nèi)核優(yōu)化

該優(yōu)化措施使端到端速度較 2 月基線實(shí)現(xiàn)版本提升了 20%。其使用高吞吐量生成 MLA 內(nèi)核，所涉及的技術(shù)包括使用 Blackwell GPU 的 Tensor Core 5th MMA 指令的 2CTA 組變體，通過 interleaved tile 將 MLA 與 softmax 重疊，并針對(duì) DeepSeek R1 問題規(guī)模微調(diào)內(nèi)核選擇啟發(fā)式算法。

FP8 KV 緩存

這項(xiàng)重要的優(yōu)化措施能夠在并發(fā)性相同的情況下使端到端吞吐量提升 6%。FP8 KV 緩存的另一個(gè)優(yōu)點(diǎn)是將 KV 緩存大小減半，從而支持更大的并發(fā)性，并使用比運(yùn)行 BF16 數(shù)據(jù)類型更快的 FP8 注意力內(nèi)核。我們建議用戶始終啟用 FP8 KV 緩存以提高性能。在上下文階段，KV 被量化為 FP8 并保存到 KV 緩存池中。在生成階段，Q 和 KV 均被量化為 FP8 并使用 FP8 多查詢注意力機(jī)制 (MQA)。在 GSM8k 數(shù)據(jù)集上的評(píng)估顯示，準(zhǔn)確率無明顯下降。量化通常使用靜態(tài)的每張量 FP8，縮放因子默認(rèn)設(shè)置為 1.0，但 KV 緩存縮放因子也可通過在目標(biāo)數(shù)據(jù)集上校準(zhǔn)生成。以下是不同組合在 GSM8K 數(shù)據(jù)集上的準(zhǔn)確率指標(biāo)。

手動(dòng) GEMM 策略調(diào)優(yōu)

此優(yōu)化措施針對(duì) cuBLAS 中默認(rèn)啟發(fā)式算法在模型中存在特定 GEMM 形狀時(shí)表現(xiàn)不佳的情況。我們開發(fā)了一個(gè)內(nèi)部工具用于離線查找適合這些特定形狀的最佳算法，然后在運(yùn)行時(shí)使用cublasLtMatmulAPI 應(yīng)用此特定優(yōu)化算法。當(dāng)通用啟發(fā)式算法無法為所有特定場景找到最高效的內(nèi)核時(shí)，這一系統(tǒng)優(yōu)化十分必要。我們也正與 cuBLAS 團(tuán)隊(duì)積極合作，進(jìn)一步改進(jìn)啟發(fā)式算法，以便始終能夠?qū)崿F(xiàn)開箱即用 (OOTB) 的最佳性能。調(diào)優(yōu)詳情參見 cublasScaledMM.cpp。

水平融合

這涉及融合 Q / KV 向下投影 GEMM 操作與 K 張量的 RoPE 維度。詳情參見 modeling_deepseekv3.py。水平融合可減少內(nèi)核啟動(dòng)開銷并增大 GEMM 問題規(guī)模，從而提高硬件利用率。這是常用的低延遲優(yōu)化和吞吐量優(yōu)化技術(shù)。

雙流優(yōu)化

在 MLA 中有一些小操作可并行運(yùn)行，例如 Q 范數(shù)和 KV 范數(shù)。這些操作無法充分使用 GPU 的計(jì)算能力和顯存帶寬，因此可通過在并行 CUDA 流中運(yùn)行來提升速度。

MoE 層優(yōu)化

已對(duì) MoE 層采取以下優(yōu)化措施。

在RouterGEMM 中使用混合 I/O 數(shù)據(jù)類型

通過避免類型轉(zhuǎn)換 (Casting) 操作并直接使用混合輸入和輸出數(shù)據(jù)類型（例如 BF16 輸入和 FP32 輸出）執(zhí)行 GEMM，將端到端速度提升了 4%。這樣就無需將輸入顯式轉(zhuǎn)換為輸出類型，并節(jié)省了顯存帶寬。

Top-K 內(nèi)核融合

該優(yōu)化措施將端到端速度提升了 7.4%。在 DeepSeek R1 中，從 256 個(gè)專家模型中選出前 8 名，其選擇過程分為兩個(gè)階段：首先選擇若干個(gè)專家組，然后在這些組中選出前 8 個(gè)專家。DeepSeek R1 采用了一些額外技術(shù)來實(shí)現(xiàn)更好的專家負(fù)載平衡，包括在 topK 復(fù)雜度中添加偏置和縮放因子。所有這些操作在未融合時(shí)都會(huì)產(chǎn)生 18 個(gè) PyTorch 操作，詳見 Deepseekv3RoutingImpl。融合這些 Top-K 計(jì)算所涉及的多個(gè)內(nèi)核可顯著縮短整體計(jì)算時(shí)間。與使用 18 個(gè)原生 PyTorch 操作相比，融合可將操作數(shù)量減少到僅有 2 個(gè)內(nèi)核。根據(jù)在 B200 平臺(tái)上的測(cè)量結(jié)果，融合這些內(nèi)核可將目標(biāo)設(shè)置下的內(nèi)核時(shí)間從 252us 縮短至 15us。

FP4 AllGather 優(yōu)化

該優(yōu)化措施使端到端速度提升了 4%，將 BF16 AllGather 操作替換為了 FP4 版本。由于這種通信原語使用較低精度，因此網(wǎng)絡(luò)傳輸?shù)臄?shù)據(jù)量有所減少，而通信效率大幅提升。此外，由于原始的 BF16 張量在 AllGather 通信后會(huì)被轉(zhuǎn)換為 FP4 格式，此優(yōu)化措施不會(huì)對(duì)準(zhǔn)確性產(chǎn)生任何影響。在內(nèi)核層面，從 BF16 切換到 FP4 AllGather 使性能提高了約 3 倍。

CUTLASS 組 GEMM 優(yōu)化

該優(yōu)化措施將端到端速度提升了 1.3%。若干 CUTLASS 方向的優(yōu)化同時(shí)適用于低延遲和高吞吐場景。更新 CUTLASS 至最新版本即可將 MoE 組 GEMM 的內(nèi)核性能提升 13% ，端到端 TPS / GPU 提升 1.3%。

多流優(yōu)化

在雙流中運(yùn)行共享和路由專家，結(jié)合 MLA 模塊中的其他多流優(yōu)化措施將端到端速度提升了 5.3%。

運(yùn)行時(shí)優(yōu)化

這些優(yōu)化措施旨在推理系統(tǒng)中的整體執(zhí)行流程、調(diào)度和資源管理。它們?cè)?DeepSeek R1 模型和其他 TensorRT-LLM 支持的模型之間共享。下面將介紹一些提升 B200 上 DeepSeek R1 性能的消融實(shí)驗(yàn)。

CUDA Graph

CUDA Graph 將吞吐量場景中的端到端性能提升了 22%。

CUDA Graph 允許捕獲一連串 CUDA 操作并將其作為單個(gè)單元啟動(dòng)，從而大幅減少內(nèi)核啟動(dòng)開銷。這對(duì)于具有大量小型內(nèi)核的模型特別有益，尤其是在 PyTorch 流程上，因?yàn)?Python 主機(jī)代碼的執(zhí)行速度通常慢于 C++。由于 CUDA Graph 凍結(jié)了內(nèi)核啟動(dòng)參數(shù)（這些參數(shù)通常與張量形狀相關(guān)），因此只能安全地用于靜態(tài)形狀，這意味著不同批次大小需要捕獲不同的 CUDA Graph。每個(gè) Graph 都會(huì)產(chǎn)生一些顯存使用和捕獲時(shí)間的開銷，因此無法為所有可能的批次捕獲所有可能的 CUDA Graph。對(duì)于未捕獲的批次大小，將執(zhí)行 PyTorch 的即時(shí)模式代碼。

TensorRT-LLM 中有一個(gè)名為“CUDA Graph 填充”(CUDA Graph padding) 的功能，它在 CUDA Graph 的數(shù)量和命中率之間取得了很好的平衡；該功能嘗試將批次填充到最近捕獲的 CUDA Graph。通常情況下建議啟用 CUDA Graph 填充功能以提高 CUDA Graph 命中率，但填充本身會(huì)因 Token 計(jì)算的浪費(fèi)而帶來一些開銷。

用戶可通過設(shè)置cuda_graph_config: enable_padding: False來禁用 CUDA Graph 填充功能以查看性能優(yōu)勢(shì)。參見 API：Pytorch 后端配置

https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/pyexecutor/config.py#L41

重疊調(diào)度器

該優(yōu)化措施將端到端性能提升了 4%，故應(yīng)默認(rèn)啟用。此調(diào)度器管理不同操作（如計(jì)算和通信）的執(zhí)行，從而在 GPU 和網(wǎng)絡(luò)上有效重疊操作。其原理是通過在等待數(shù)據(jù)傳輸時(shí)執(zhí)行計(jì)算（或反之）來隱藏延遲，從而提高整體硬件利用率。在 TensorRT-LLM 中，重疊調(diào)度已默認(rèn)在提交時(shí)啟用。如果存在特殊情況導(dǎo)致其無法正常工作，用戶仍可通過將 disable_overlap_scheduler 設(shè)置為 true 來禁用此功能。

顯存優(yōu)化

該優(yōu)化措施減少了 4GB 顯存占用。所使用的技術(shù)包括針對(duì) Hopper 架構(gòu)的分塊 MoE 以及修復(fù) CUDA 上下文初始化的錯(cuò)誤。這些方法降低了模型權(quán)重或中間張量的顯存占用，從而支持更大的批次大小或序列長度，避免了顯存不足 (OOM) 的錯(cuò)誤。

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

參見性能優(yōu)化實(shí)踐：

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

未來工作

大型專家并行 (EP)

上下文分塊

增加通信重疊

致謝

本文詳細(xì)介紹了如何在 Blackwell GPU 上顯著提升 DeepSeek R1 的吞吐量，這是我們工程團(tuán)隊(duì)共同努力的成果。他們通過深入研究 MLA 層、MoE 層和運(yùn)行時(shí)優(yōu)化措施，將 TPS / GPU 提高了近 2.3 倍。我們衷心感謝參與這次集中優(yōu)化工作的所有工程師，他們的集體智慧在進(jìn)一步突破 TensorRT-LLM 吞吐量性能方面起到了至關(guān)重要的作用。相信分享這些大幅提高吞吐量的針對(duì)性策略，將有助于開發(fā)者群體在 NVIDIA 硬件上部署高負(fù)載的 LLM 推理任務(wù)。

作者

張國銘

NVIDIA 性能架構(gòu)師，目前主要從事大模型推理架構(gòu)和優(yōu)化。