TensorRT-LLM 的大规模专家并行架构设计与创新实现

之前文章已介绍引入大规模 EP 的初衷，本篇将继续深入介绍 TensorRT-LLM 的大规模专家并行架构设计与创新实现。

上篇文章参考，点击阅读：

大规模专家并行 (EP) 在 TensorRT-LLM 的设计动机与系统分析

高层次设计介绍

根据引入大规模 EP 的初衷部分的详细分析与研究，可以明确观察到 EP 中的专家失衡是大规模 EP 的常见模式。这种 EP 失衡会通过以下方式显著降低整体系统性能：

• 热门 EP rank 将消耗更多显存（用于激活值），这会限制推理过程中调度的有效最大批处理大小。

• 更多数据将从热门 EP rank 被发送和接收。

这些问题将导致系统级拥塞效应，即热门 EP rank 将延迟整体端到端执行。

为确保大规模 EP 能稳定运行，需通过精心设计尽可能减少 EP 失衡问题。整体设计如下：

图 1. TensorRT-LLM 大规模 EP 的高层次设计

此设计同时包含 CPU 和 GPU 两侧逻辑：

• CPU 侧

• 使用复制与放置算法（复制与放置计算组件）实现更均衡的 EP 策略。这些算法是经典算法，更适合 CPU 计算。此外，将此计算卸载至 CPU 可减少对 GPU 的干扰。未来可探索基于机器学习的算法，并可能需要额外设计考量。复制与放置计算组件将生成“放置信息”，该信息将被 GPU路由逻辑和 CPU更新权重与放置组件共同使用。由 GPU 上运行的统计组件生成的统计数据将被用作复制与放置计算组件的输入。

• 编排流程（更新权重与放置组件）将 MoE 权重从 CPU 内存更新并重新加载到 GPU 设备显存。该组件还将使用由复制与放置计算组件生成的放置信息。我们的可扩展设计允许通过 MNNVL 或 NIC 从远程 GPU 显存重新加载 MoE 权重。

• GPU 侧

• 这是推理的主要执行工作流。我们在设计中引入了以下新的 GPU 组件：

• EP 通信内核，在上篇图 11 中为分发合并组件。

• 在线流量数据统计采集器（统计组件）。该组件采集统计数据复制与放置计算组件使用。

• MoE 路由逻辑（路由组件）。该组件将 Token 发送至激活的专家，并且需要进行调整以支持 MoE 权重的动态放置。它使用复制与放置计算组件生成的放置信息

• MoE 计算逻辑 (MoE组件) 也需进行相应调整。

• CPU 和 GPU 组件之间需要仔细同步，以确保整个执行过程的有效性，尤其是为了避免卡顿以及无效或次优执行。

我们为更新权重与放置组件提供了两种设计方案：

• 批量方案

• 在此方案中，当 MoE 权重重新分配逻辑启动时，当前服务实例上的推理过程将不得不暂停，直至 MoE 权重重新分配过程完成。我们估计这可能导致约0.5 至 1 秒的在线服务暂停，最坏情况下会引发请求超时。此类超时或暂停可通过系统级措施来缓解，例如将请求传送至其他服务实例或通过请求重试来应对。

• 分层方案

图 2. 分层 MoE 权重重新分配示例

在当前系统中，我们选择采用分层方案以尽量减少对在线用户体验的影响。批量方案应更易于实现，但本文将不再讨论。为了正确实现分层方案，需仔细评估不同底层硬件的性能以确定具体实现方案。图 3 展示了系统节点中不同硬件组件的通信带宽。

图 3. 系统高层次拓扑结构

以 DeepSeek R1 模型为例，采用 FP4 精度时，每个 MoE 专家占用 24MiB 显存空间。每层包含 256 个专家，总共包含 58 个 MoE 层加 1 个 MTP 层。因此，为实现 EP 平衡所需重新分配的 MoE 权重最大总量为 348GiB。每个节点为每个 Grace CPU 提供 480GB LPDDR5X 显存。在 NUMA 域内，总计可提供 960GB Host 显存。一个节点可在其 CPU Host 显存中完整承载如 DeepSeek R1 LLM 等模型的全部 MoE 权重。基于此，MoE 权重重新分配可通过将对应的 MoE 权重从 CPU Host 显存移动至 GPU 设备显存来实现。

假设我们将50ms的跨 Token 延迟 (ITL) 作为主要延迟约束。通过粗略估算，可以计算出在每次解码迭代中，可从 MoE 权重池（可保存在 Grace CPU 显存或另一节点上的 GPU 显存中）移动到 Blackwell GPU（用于实际 MoE 推理）的专家权重数量为：

图 4. 在以下 50ms ITL 限制下，每次迭代理论上需要更新的专家数量（使用不同硬件作为存储完整 MoE 权重的池）

基于此分析，若依赖每个节点上的 Grace CPU 内存来存储 MoE 权重池，则每次解码迭代中，最多可将 300 个专家的权重重新分配至同一节点上的每个 GPU。假设目标是在 5 次解码迭代内完成整个模型 MoE 权重再平衡，以下为具体用例研究：

• 用例 1（专家分配均衡，不进行专家复制）

• 64 个 GPU，每个 GPU 分配 4 个专家

• 58 层，每个 GPU 分配 232 个专家

• 每次迭代需要 47 次专家更新，所有方法均可满足延迟目标。

• 用例 2（专家分配均衡并进行复制）

• 64 或 72 个 GPU，每个 GPU 分配 5 个专家

• 58 层，每个 GPU 分配 290 个专家

• 每次迭代需要 58 次专家更新，所有方法均可满足延迟目标。

• 用例 3（专家分配均衡并进行复制）

• 36 个 GPU，每个 GPU 分配 8 个专家

• 58 层，每个 GPU 分配 464 个专家

• 每次迭代需要 93 次专家更新，所有方法均可满足延迟目标。

综上所述，根据理论分析，采用 Grace CPU 内存作为存储完整大小 MoE 权重的池，应能使我们在 5 次解码迭代内实现 EP（专家并行）的再平衡。如果将要求放宽至 10 次或以上迭代，系统实现将变得更加灵活。

接下来我们将介绍大规模 EP 系统的详细实现方式。

EP 通信内核

我们评估了多种实现大规模 EP 所需 EP 通信内核的途径，包括 DeepEP、其他解决方案以及重新开发一种方法。

当前的技术决策是：

• 我们实现了一组新的自定义 EP 通信内核。

• 对于其他系统（如 Hopper），我们选择直接集成 DeepEP 并进行一些可能的增强。

考虑因素：

• DeepEP 是由 DeepSeek 团队完成的一项出色成果。我们在启动 TensorRT-LLM 大规模 EP 工作时，最初把重点放在 Grace Blackwell 机架式系统上。我们选择实现自己的定制 EP 通信内核，因为这更便于引入需要 Grace Blackwell 机架式系统功能的优化措施。

• 当我们开始在 Hopper 上启用大规模 EP 工作时，我们得出的结论是 DeepEP 可以适应并满足我们在该平台上的需求。

我们也在积极评估将通信内核整合为单一解决方案以简化系统架构的可能性，并将持续向社区更新进展。接下来，我们将进一步探讨自定义 EP 通信内核实现中引入的优化措施。

在系统中引入 EP 通信内核的初衷

在解码阶段与预填充解码 (PD) 分离的场景中，我们观察到批处理大小可能不会很大，因此延迟成为一个重要考虑因素。在此背景下，我们非常需要实现与 CUDA graph 的兼容。NCCL 是一个优秀的 GPU 通信库，为我们提供了高效的通信内核和基本操作。目前，其 Send 和 Recv 操作在调用 ncclSend / ncclRecv 时，需要显式指定数据大小。但在大规模专家并行 (large-EP) 场景中，待传输的数据大小根据模型在每次迭代中的输出动态确定。当前 NCCL 通信接口需要同步将通信大小发回 CPU，并以对应数据大小从 CPU 发起 NCCL 调用。这将破坏 CUDA graph 兼容性。这一限制迫使我们开发与 CUDA graph 兼容，且能直接从 GPU 显存接受通信大小的高性能通信内核。我们还希望这些内核能够充分利用 MNNVL 的显存带宽。

EP 通信内核的实现

我们的内核采用与 NCCL 的 LL128 原语类似的通信方法。由于这种方法在延迟和带宽之间取得了良好的平衡，因此非常适合 LLM 推理。我们的自定义内核可直接从 GPU 显存读取通信大小并兼容 CUDA graph，即使数据大小在不同运行中变化也不例外。

我们的实现方式是使用 CUDA 的驱动程序 API 通过 MNNVL 建立点对点 (P2P) 缓冲区作为工作区。每个 GPU 都可以访问其他 GPU 的工作区。工作区被划分为多个通道，每个通道分配给远程 GPU 作为写入缓冲区。这些写入缓冲区以 FIFO 方式使用，通过标志同步 FIFO 状态以避免数据损坏。详细信息请参见 PR 3504:

https://github.com/NVIDIA/TensorRT-LLM/pull/3504

下一篇我们将继续介绍 TensorRT-LLM 在线负载均衡策略与实测的解析。

作者

杨东旭

现任职于 NVIDIA Compute Arch 部门。主要负责 LLM 推理系统的开发和性能优化。加入 NVIDIA 之前，曾从事搜索系统的 GPU 加速和开发工作。

乔显杰

NVIDIA Compute Arch 部门高级架构师，主要负责 LLM 推理的性能评估和优化。加入 NVIDIA 之前，他曾从事推荐系统的 GPU 加速研发工作。

谢开宇

NVIDIA Compute Arch 部门高级架构师，主要负责 TensorRT-LLM 项目的开发，专注在系统性能和优化工作。

朱恩伟

NVIDIA DevTech 部门高级工程师，主要负责 TensorRT-LLM 项目的开发和性能优化。

陈晓明

NVIDIA Compute Arch 部门的首席架构师和高级经理，对深度学习模型的算法软硬件协同设计感兴趣，最近从事大语言模型推理的性能建模、分析和优化。