如何在NVIDIA Blackwell GPU上优化DeepSeek R1吞吐量

前言

开源 DeepSeek R1 模型的创新架构包含多头潜在注意力机制 (MLA) 和大型稀疏混合专家模型 (MoE)，其显著提升了大语言模型 (LLM) 的推理效率。但要充分发挥这种创新架构的潜力，软硬件的协同优化也至关重要。本文将深入解析 NVIDIA 在基于 Blackwell GPU 的TensorRT-LLM框架内为 DeepSeek R1 吞吐量优化场景 (TPS / GPU) 开发的优化策略。文内将详细阐述各项优化措施的设计思路。另一篇关于降低延迟的博客（见下）已详细解释了 TensorRT-LLM 如何通过提升 R1 性能实现最佳 TPS / USER。

这些优化措施显著提升了 DeepSeek R1 在 Blackwell 上的吞吐量。其在 ISL / OSL 1K / 2K 数据集上的性能从 2 月的约 2000 TPS / GPU 提升至 4600 TPS / GPU。这些优化措施具有通用性，可适用于其他 ISL / OSL 配置。其大致分为三个方面：MLA 层、MoE 层和运行时。

精度策略

DeepSeek R1 吞吐量场景的混合精度策略与延迟优化场景的策略基本一致，具体差异如下：

使用 FP8 KV 缓存和 FP8 注意力机制，而非 BF16 精度。

使用 FP4 Allgather 提高通信带宽利用率。

本文中使用的 Checkpoint 位于 nvidia / DeepSeek-R1-FP4，由 NVIDIA 模型优化器生成。常用数据集在该 FP4 Checkpoint 和 TensorRT-LLM 执行中准确率分数为：

** 注意这些评估结果存在运行间波动，因此 FP4 数据的分数略高。我们认为 FP4 在这些数据集上的精度与 FP8 相当。

该 Checkpoint 中的 MoE 层已量化为 FP4。将 MoE 层权重量化为 FP4 具有以下优点：

充分利用 NVIDIA Blackwell GPU 第五代 Tensor Core 的 FLOPS（每秒浮点运算）

将 MoE 权重所需的显存负载减少近一半。在该场景中，MoE 部分在 Decoding 阶段仍受显存限制，且 DeepSeek R1 模型中 97% 的权重来自 MoE 层。

减少模型权重的显存占用，从而释放更多 GPU 显存用于 KV 缓存，进而增加最大并发数。DeepSeek R1 模型的原始 FP8 量化 Checkpoint 大小约为 640GB，而 NVIDIA 提供的 DeepSeek R1 FP4 量化模型仅为约 400GB。

我们在 GSM8K 数据集上对 FP8 KV 缓存和 FP8 注意力内核的精度进行了评估，未观察到明显的准确率下降。有关准确率数据，请参见 FP8 KV 缓存部分的表格。如果用户在自己的数据集上观察到准确率差异，仍可选择使用 BF16 KV 缓存和注意力机制。

并行策略

DeepSeek R1 吞吐量场景的并行策略与延迟场景的策略有所不同。

在接下来的部分，我们将解释为何选择数据并行 (DP) 和专家并行 (EP)，而非张量并行 (TP)。

权重吸收与 MQA

MLA 的核心理念是通过对注意力键 (K) 和值 (V) 进行低秩联合压缩，减少推理过程中的 KV 缓存大小。基于 MLA 公式，向下投影的 KV Latent 被向上投影到多个头并与向上投影的 Q 结合，形成常规的多头注意力机制 (MHA)。由于矩阵乘法的性质，K 的向上投影权重矩阵 (W^UK) 可先与 Q 的向上投影权重矩阵 (W^Q) 相乘，再将结果与 Q 相乘。V 的向上投影权重矩阵 (W^UV) 与注意力输出投影矩阵 W^O 也可在注意力输出后相乘。DeepSeek-V2 技术报告将该技术称为“吸收”。在权重被吸收后，MLA 等价于多查询注意力 (MQA)。参见 DeepSeek-V2 技术论文获取详细公式和解释，下图所示的是 TensorRT-LLM 中权重吸收 MLA 的计算流程。

该图片来源于 Github:Optimizing DeepSeek R1 Throughput on NVIDIA Blackwell GPUs: A Deep Dive for Developers一文，若您有任何疑问或需要使用该图片，请联系该文作者

在 Decoding 阶段，权重吸收显著减少向上投影 K 和 V 所需的数学 FLOPS，这是因为这些 KV 向上投影所需的 FLOPS 与 KV 缓存长度成线性关系，而 Decoding 阶段的 Q 向量长度始终为 1。KV 缓存历史越长，所需的 FLOPS 越多，且由于仅保存了投影的 KV Latent，每个 Decoding Token 都需要重复进行向上投影，这进一步增加了所需的 FLOPS。在 Prefill 阶段，吸收权重的版本改变了 Q 和 KV 的维度，增加了注意力机制所需的 FLOPS。根据 Roofline 分析，在输入长度为 256 或以上的情况下，非吸收版本更有利于 Prefill 阶段。TensorRT-LLM MLA 实现为 Prefill 和 Decoding 阶段分别选择了不同的高度优化内核，详见 MLA。

注意力模块数据并行 (ADP)

选择注意力数据并行 (ADP) 的核心原因是对 MQA（由不同 GPU 计算不同的注意力 Q 头）采用张量并行 (TP) 会重复复制 KV 缓存显存，从而限制系统能够实现的并发性。重复因子等于 TP 组大小，因此 TP8 时为 8 倍。并发性低会影响高性能系统 (如 NVIDIA DGX B200) 的吞吐量。

对于使用 8 个 B200 GPU 的 DeepSeek R1 FP4 Checkpoint，每个 GPU 的权重和激活占用约 80GB 显存，每个 GPU 的可用 KV 缓存将为 100GB。假设 ISL 为 1K，OSL 为 2K，每次请求将消耗约 200MB KV 缓存，使每个 GPU 的最大并发数为 500。单节点 8 GPU 系统的全局并发性为 4000。使用 ATP 时，全局并发性将降至 500。

硅片实验表明，在保持所有其他因素不变的情况下，ADP 技术在最大吞吐量场景下可提速 400%。

MoE 专家并行 (EP)

DeepSeek R1 MoE 设计包含 256 个小型稀疏专家模型和 1 个共享专家模型，这些专家模型的 GEMM 问题规模如下。

这些专家模型可采用张量并行 (TP) 或专家并行 (EP) 的方式实现。目前的消融实验表明，由于 GEMM 问题规模更小，专家并行在 GEMM FLOPS 方面的表现更好。并且与 AllReduce 相比，专家并行可节省 GPU 通信带宽，这是因为只需将 Token 发送至该 Token 对应的活跃专家所在的 GPU，而张量并行 (TP) 则需对所有 GPU 间的 Token 执行 AllReduce 操作。此外，为了将 DeepSeek R1 推理扩展到 GB200 NVL72 等系统并充分利用聚合显存带宽和 Tensor Core FLOPS，需要使用大型专家并行。我们正在积极推进该技术的实现。

硅片性能测试表明，在保持其他因素不变的情况下，专家并行 (EP) 可在 1K / 2K 最大吞吐量场景下将速度提升 142%。

MLA 层优化

除了上述提到的并行策略和精度策略外，我们在 MLA 模块内的层 / 内核中进行了以下优化。

注意力内核优化

该优化措施使端到端速度较 2 月基线实现版本提升了 20%。其使用高吞吐量生成 MLA 内核，所涉及的技术包括使用 Blackwell GPU 的 Tensor Core 5th MMA 指令的 2CTA 组变体，通过 interleaved tile 将 MLA 与 softmax 重叠，并针对 DeepSeek R1 问题规模微调内核选择启发式算法。

FP8 KV 缓存

这项重要的优化措施能够在并发性相同的情况下使端到端吞吐量提升 6%。FP8 KV 缓存的另一个优点是将 KV 缓存大小减半，从而支持更大的并发性，并使用比运行 BF16 数据类型更快的 FP8 注意力内核。我们建议用户始终启用 FP8 KV 缓存以提高性能。在上下文阶段，KV 被量化为 FP8 并保存到 KV 缓存池中。在生成阶段，Q 和 KV 均被量化为 FP8 并使用 FP8 多查询注意力机制 (MQA)。在 GSM8k 数据集上的评估显示，准确率无明显下降。量化通常使用静态的每张量 FP8，缩放因子默认设置为 1.0，但 KV 缓存缩放因子也可通过在目标数据集上校准生成。以下是不同组合在 GSM8K 数据集上的准确率指标。

手动 GEMM 策略调优

此优化措施针对 cuBLAS 中默认启发式算法在模型中存在特定 GEMM 形状时表现不佳的情况。我们开发了一个内部工具用于离线查找适合这些特定形状的最佳算法，然后在运行时使用cublasLtMatmulAPI 应用此特定优化算法。当通用启发式算法无法为所有特定场景找到最高效的内核时，这一系统优化十分必要。我们也正与 cuBLAS 团队积极合作，进一步改进启发式算法，以便始终能够实现开箱即用 (OOTB) 的最佳性能。调优详情参见 cublasScaledMM.cpp。

水平融合

这涉及融合 Q / KV 向下投影 GEMM 操作与 K 张量的 RoPE 维度。详情参见 modeling_deepseekv3.py。水平融合可减少内核启动开销并增大 GEMM 问题规模，从而提高硬件利用率。这是常用的低延迟优化和吞吐量优化技术。

双流优化

在 MLA 中有一些小操作可并行运行，例如 Q 范数和 KV 范数。这些操作无法充分使用 GPU 的计算能力和显存带宽，因此可通过在并行 CUDA 流中运行来提升速度。

MoE 层优化

已对 MoE 层采取以下优化措施。

在RouterGEMM 中使用混合 I/O 数据类型

通过避免类型转换 (Casting) 操作并直接使用混合输入和输出数据类型（例如 BF16 输入和 FP32 输出）执行 GEMM，将端到端速度提升了 4%。这样就无需将输入显式转换为输出类型，并节省了显存带宽。

Top-K 内核融合

该优化措施将端到端速度提升了 7.4%。在 DeepSeek R1 中，从 256 个专家模型中选出前 8 名，其选择过程分为两个阶段：首先选择若干个专家组，然后在这些组中选出前 8 个专家。DeepSeek R1 采用了一些额外技术来实现更好的专家负载平衡，包括在 topK 复杂度中添加偏置和缩放因子。所有这些操作在未融合时都会产生 18 个 PyTorch 操作，详见 Deepseekv3RoutingImpl。融合这些 Top-K 计算所涉及的多个内核可显著缩短整体计算时间。与使用 18 个原生 PyTorch 操作相比，融合可将操作数量减少到仅有 2 个内核。根据在 B200 平台上的测量结果，融合这些内核可将目标设置下的内核时间从 252us 缩短至 15us。

FP4 AllGather 优化

该优化措施使端到端速度提升了 4%，将 BF16 AllGather 操作替换为了 FP4 版本。由于这种通信原语使用较低精度，因此网络传输的数据量有所减少，而通信效率大幅提升。此外，由于原始的 BF16 张量在 AllGather 通信后会被转换为 FP4 格式，此优化措施不会对准确性产生任何影响。在内核层面，从 BF16 切换到 FP4 AllGather 使性能提高了约 3 倍。

CUTLASS 组 GEMM 优化

该优化措施将端到端速度提升了 1.3%。若干 CUTLASS 方向的优化同时适用于低延迟和高吞吐场景。更新 CUTLASS 至最新版本即可将 MoE 组 GEMM 的内核性能提升 13% ，端到端 TPS / GPU 提升 1.3%。

多流优化

在双流中运行共享和路由专家，结合 MLA 模块中的其他多流优化措施将端到端速度提升了 5.3%。

运行时优化

这些优化措施旨在推理系统中的整体执行流程、调度和资源管理。它们在 DeepSeek R1 模型和其他 TensorRT-LLM 支持的模型之间共享。下面将介绍一些提升 B200 上 DeepSeek R1 性能的消融实验。

CUDA Graph

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

CUDA Graph 允许捕获一连串 CUDA 操作并将其作为单个单元启动，从而大幅减少内核启动开销。这对于具有大量小型内核的模型特别有益，尤其是在 PyTorch 流程上，因为 Python 主机代码的执行速度通常慢于 C 。由于 CUDA Graph 冻结了内核启动参数（这些参数通常与张量形状相关），因此只能安全地用于静态形状，这意味着不同批次大小需要捕获不同的 CUDA Graph。每个 Graph 都会产生一些显存使用和捕获时间的开销，因此无法为所有可能的批次捕获所有可能的 CUDA Graph。对于未捕获的批次大小，将执行 PyTorch 的即时模式代码。

TensorRT-LLM 中有一个名为“CUDA Graph 填充”(CUDA Graph padding) 的功能，它在 CUDA Graph 的数量和命中率之间取得了很好的平衡；该功能尝试将批次填充到最近捕获的 CUDA Graph。通常情况下建议启用 CUDA Graph 填充功能以提高 CUDA Graph 命中率，但填充本身会因 Token 计算的浪费而带来一些开销。

用户可通过设置cuda_graph_config: enable_padding: False来禁用 CUDA Graph 填充功能以查看性能优势。参见 API：Pytorch 后端配置

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

重叠调度器

该优化措施将端到端性能提升了 4%，故应默认启用。此调度器管理不同操作（如计算和通信）的执行，从而在 GPU 和网络上有效重叠操作。其原理是通过在等待数据传输时执行计算（或反之）来隐藏延迟，从而提高整体硬件利用率。在 TensorRT-LLM 中，重叠调度已默认在提交时启用。如果存在特殊情况导致其无法正常工作，用户仍可通过将 disable_overlap_scheduler 设置为 true 来禁用此功能。

显存优化

该优化措施减少了 4GB 显存占用。所使用的技术包括针对 Hopper 架构的分块 MoE 以及修复 CUDA 上下文初始化的错误。这些方法降低了模型权重或中间张量的显存占用，从而支持更大的批次大小或序列长度，避免了显存不足 (OOM) 的错误。

如何复现

参见性能优化实践：

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)

上下文分块

增加通信重叠

致谢

本文详细介绍了如何在 Blackwell GPU 上显著提升 DeepSeek R1 的吞吐量，这是我们工程团队共同努力的成果。他们通过深入研究 MLA 层、MoE 层和运行时优化措施，将 TPS / GPU 提高了近 2.3 倍。我们衷心感谢参与这次集中优化工作的所有工程师，他们的集体智慧在进一步突破 TensorRT-LLM 吞吐量性能方面起到了至关重要的作用。相信分享这些大幅提高吞吐量的针对性策略，将有助于开发者群体在 NVIDIA 硬件上部署高负载的 LLM 推理任务。

作者

张国铭

NVIDIA 性能架构师，目前主要从事大模型推理架构和优化。

• NVIDIA NVIDIA 关注

  关注

  14

  文章

  5343

  浏览量

  106782