性能飞跃：TensorRT-LLM 在线负载均衡策略与实测解析

摘要：之前文章已介绍引入大规模 EP 的初衷和 TensorRT-LLM 的大规模专家并行架构设计与创新实现。本篇将具体介绍 TensorRT-LLM 在线负载均衡策略与实测解析。

之前文章已介绍引入大规模 EP 的初衷和 TensorRT-LLM 的大规模专家并行架构设计与创新实现。本篇将具体介绍 TensorRT-LLM 在线负载均衡策略与实测解析。

第二篇：TensorRT-LLM 的大规模专家并行架构设计与创新实现「链接」

EP负载均衡器

TensorRT-LLM 采取了一系列实现 EP 负载均衡的功能，包括几个关键组件：

Python 接口

Python 接口层提供了一个对用户友好的 PyTorch / Python 原生接口，用于访问 MoE负载均衡的实现，例如用于 GPU / CPU 同步逻辑的 Python 包装器 (wrapper)、在线数据统计采集。

C++ 扩展

C++ 扩展作为 PyTorch / Python 接口与 C++ / CUDA 核心实现之间的桥梁。

Host 逻辑核心实现

Host 侧核心逻辑实现了以下关键部分：

负载平衡算法复制算法放置算法MoE 权重更新的编排逻辑MoE 权重更新逻辑

GPU 逻辑核心实现

GPU 核心逻辑包含以下组件：

在线流量统计采集为降低 CPU-GPU 往返同步开销，我们选择在 GPU 侧实现在线流量统计逻辑。专家路由逻辑需增强 MoE 路由逻辑以适应动态 EP 平衡影响。

已实现 GPU / CPU 同步组件。详情请参见 PR 4384[2] 和 PR 4495[3]。

基于这些核心实用程序，TensorRT-LLM 中提供了两种版本的 EP 负载均衡器：离线 EP 负载均衡器和在线 EP 负载均衡器。

在线 EP 负载均衡器

对于生产部署需求，推荐使用在线 EP 负载均衡器，因为它能够动态适应在线流量模式变化，从而更能保障性能。

但在线 EP 负载均衡器面临以下挑战。

首先，负载均衡引入了动态专家放置。单个专家的位置可能根据当前工作负载发生变化。例如，如果最初分配到 rank 0 的专家 0 和专家 1 均成为热门专家，负载均衡策略可能会将其重新分配到不同 rank，与冷门专家并列，这就需要及时更新权重数据。

我们旨在使在线负载均衡器能够快速响应请求模式的变化，并调整专家分配以避免负载失衡问题。重要的是，我们不希望平衡过程中干扰在线推理执行过程，也不希望采用“批量”策略更新权重。

大型 MoE 模型 (如 DeepSeek R1) 解码阶段的批处理大小通常较小，这使得 CUDA graph 成为一种有效的加速方法，尤其是在需要高单位用户 TPS 的情况下。这种优势在 Grace Blackwell 机架式系统等平台上表现得更为明显，因此我们希望整个负载均衡机制与 CUDA graph 兼容。

为避免无效化预捕获的 CUDA graph，我们通过将新专家权重写入相同显存位置的方式进行就地权重更新，而非替换张量指针。这确保了模型引擎中权重张量的地址保持不变。

在此设计中，每个专家槽位作为容器用于存储专家权重，与特定专家解耦。专家槽位的数量必须大于或等于专家总数，以确保每个专家始终至少有一个可用槽位。热门专家可能占用多个槽位。每个槽位通过 SlotId 进行标记。

由于 MoE 模型的路由逻辑输出 ExpertId（而非 SlotId），我们维护了一个从 ExpertId 到 SlotId 的路由表，该表由负载均衡策略定期更新。负载均衡路由模块使用当前路由表（专家复制信息和槽位）将每个 Token 映射到合适的专家槽位。

为实现权重更新的非阻塞处理并避免“停止世界 (Stop-The-World)”现象，我们采用分层更新的策略。在某一层的前向传播完成后，我们在其下一次前向传播开始前执行该层的权重平衡操作；若本次迭代中发生权重更新，则同一层的下一次前向传播需等待更新完成后再进行。

由于前向执行通常由单个 Python 线程调用一系列 PyTorch 操作驱动，我们将权重更新例程卸载到后台 C++ 线程。Python 侧仅初始化专家槽位，并在共享 Host 显存中注册专家权重。

在前向执行期间，我们在 MoE 计算的前后插入轻量化锁定 / 解锁内核，以及用于统计数据采集和将 SlotId 分配到 ExpertId 的内核。这些内核必须短且支持重叠，以尽量减少性能影响。只要 CPU 权重更新线程能及时完成工作，锁定 / 解锁操作将非常短暂。除路由内核外，所有内核均为轻量级且可轻松与不同 CUDA 流中的前向内核重叠，而路由内核是主要优化重点。

我们使用 MNNVL 进行专家调度和合并时的跨 GPU 通信。专家权重存储在 Host 显存中，并通过 C2C 传输至 GPU 显存以支持异步更新。多线程 Host 复制引擎管理此过程，自动检测 NUMA 拓扑并选择最优 CPU 核，实现与模型前向传递的完全异步。

对于不支持 C2C 但带有 PCIe 的服务器，若需跨节点通信，网络和权重更新可能会争夺 PCIe 带宽，需进行额外调优和设计考量。我们尚未为 PCIe 服务器实现复制引擎，该功能已列入未来任务清单。

离线 EP 负载均衡器

在线 EP 均衡器更适合部署在生产环境，可及时响应在线流量变化。但离线 EP 均衡器提供了一种适合性能研究 / 调试和验证的轻量级方式。您可以参考此 PR[1] 了解离线 EP 负载均衡器的详细实现方式。此外，我们还提供了一个工具用于采集专家激活分布统计数据，可作为推导 EP 负载均衡放置策略的输入。您可参考此文档[4] 了解更多细节，以及如何通过端到端方式运行离线 EP 负载均衡器。

端到端评估

EP 负载均衡器的作用

离线 EP 负载均衡器

如上篇图 1 所示，在机器翻译数据集上，MoE 第 36 层存在严重的专家负载失衡问题，因此我们使用该层展示 EP 负载均衡器的效果。我们仍在 32 个 GPU 上运行 DeepSeek-R1 并采用 32 路专家并行。

图 1. 第 36 层各接收 rank（x 轴）和迭代次数（y 轴）的传送 Token 数（未启用 EPLB）

图 2. 第 36 层各专家（x 轴）和迭代次数（y 轴）的传送 Token 数（未启用 EPLB）

图 1 显示了 50 次迭代中各接收 rank 收到的 Token 数，这可以代表每个 rank 的工作负载。rank 13 接收的 Token 数量明显多于其他所有 rank，且这种失衡的工作负载分布在整个迭代过程中几乎保持不变。图 2 将工作负载细分至专家。显然，rank 13 上的两个热门专家导致该 rank 上的压力过大。

基于上述统计数据，我们可以执行离线 EPLB。可行的策略之一是保持 32 路专家并行，同时将每个 rank 的专家槽位从 8 个增加到 9 个。这将产生 32 个冗余专家和总共 288 个专家槽位。图 3 和图 4 显示了 EPLB 后的传送 Token 数。显然，每个 rank 的 Token 分布更加平衡，且不再存在热门专家。

图 3. 第 36 层各接收 rank（x 轴）和迭代次数（y 轴）的传送 Token 数（启用 EPLB，每个 rank 9 个专家槽位，EP 为 32 路）

图 4. 第 36 层各专家（x 轴）和迭代次数（y 轴）的传送 Token 数（启用 EPLB，每个 rank 9 个专家槽位，EP 为 32 路）

另一项 EPLB 策略是保持每 rank 8 个专家槽位，同时将专家并行度提升至 36 路。该策略同样产生 32 个冗余专家和总计 288 个专家槽位。如图 5 和 6 所示，工作负载在各 rank 或专家槽位之间也实现了平衡。

图 5. 第 36 层各接收 rank（x 轴）和迭代次数（y 轴）的传送 Token 数（启用 EPLB，每个 rank 8 个专家槽位，EP 为 36 路）

图 6. 第 36 层各专家（x 轴）和迭代次数（y 轴）的传送 Token 数（启用 EPLB，每个 rank 8 个专家槽位，EP 为 36 路）

每个层和迭代可通过简单的指标（如标准差或失衡比率）测量负载失衡程度。基于所有 rank（或专家）传送 Token 数的情况下，失衡比率的计算方式为（最大值−均值） / 均值，表示最热门 rank（或专家）接收的过多工作负载。完美平衡负载的失衡比率为 0。

表 1 列出了上述情况下的标准差和失衡比率。每个数值均为各层各迭代指标的平均值。未启用 EPLB 时，负载失衡严重——最热门 rank 接收的传送 Token 数平均比整体均值多 1.56 倍。EPLB 可有效降低负载失衡——最热门 rank 接收的传送 Token 数平均仅比均值多约 0.11 倍。

表 1. 标准差与失衡比率（各层及各迭代指标的平均值）

在线 EP 负载均衡器

在上一部分中，我们展示了离线 EP 负载均衡器的效果。接下来，我们将基于我们实现的在线 EP 负载均衡器，进一步分析在线环境下 EP 均衡的动态模式。我们仍以机器翻译数据集、DeepSeek R1 模型、第 36 层（如上篇图 1 所示）为例来理解这一在线行为：

图 7. 从 rank 0 发送到所有 rank 的 Token 数（EP 为 32 路，本地 batch size=256，256 个专家槽位（无复制），每个 rank 有 8 个专家）

从图 7 可以看出，从迭代 1963 开始，由于 EPLB 已生效，原来的最热门 rank 13 不再是最热门 rank，原先发送到 rank 13 的工作负载已被重新分配到 rank 0 和 rank 1。

在图 7 中，仅通过在线 EPLB 进行了放置调整。如下图所示，若进一步引入专家复制，可进一步提升平衡效果：

图 8. 从 rank 0 发送到所有 rank 的 Token 数（EP 为 32 路，本地 batch size=256，288 个专家槽位（有复制），每个 rank 有 9 个专家）

显然，在使用 EPLB 时引入专家复制可进一步提升 EP 平衡效果。可设计更复杂的实验，观察在线 EPLB 在在线服务过程中的定期效果，以动态平衡 EP 工作负载。我们欢迎社区向我们反馈任何有趣的 EPLB 模式观察结果。

性能研究

请注意：本部分中展示的所有代表性工作负载均来自 DeepSeek R1 推理执行中提取的性能追踪。我们仍在对端到端的性能进行调优 / 优化，并将在未来的技术博客中进行讨论。

我们将通过一些具有代表性的工作负载说明大规模 EP 对性能的影响。

图 9. EP 对 MoE GroupGEMM 和 EP 通信的影响

如图 9 所示，当 EP 大小从 4 增加到 72 时，MoE GroupGEMM 的计算时间有所减少，而 EP 通信时间 (在 EP4 / EP8 时使用 Reduce / Scatter，在 EP>8 时使用 All2All) 几乎保持不变。当 EP 大小从 18 增加到 72 时，加速幅度有所下降。我们正在优化这个方面。

接下来，我们将通过一些具有代表性的工作负载了解 EPLB 对性能的影响。

图 10. EPLB 对性能的影响

如图 10 所示，当 EP 大小增加时，EPLB 能够在 MoE GroupGEMM 和 EP 通信上都实现大幅性能提升。

复现步骤

本部分中所述的复现步骤所需的代码和脚本已合并到主分支。

EP 负载均衡器的效果

请参见 EP 负载均衡器示例，了解如何复现离线 EP 负载均衡器的结果。

第 1 步：运行推理和采集统计数据

如要生成负载再平衡所需的统计数据，请在目标数据集上运行模型，并在推理过程中统计传送的专家 ID。计数过程完成后，统计数据将被保存供进一步处理。

设置一些环境变量：

export MODEL_NAME=deepseek-ai/DeepSeek-R1export MODEL_PATH=# Set the expert statistic data pathexport EXPERT_STATISTIC_PATH=./expert_statistic# Enable counting of routed expert IDs from iteration 100 to iteration 200export EXPERT_STATISTIC_ITER_RANGE=100-200

按照基准测试文档准备数据集并将其保存为 ./dataset.json。

在准备好的数据集上运行 32 路专家并行推理。有关在 Slurm 上运行 trtllm-bench 的详细信息，请参见 LLM API MGMN 示例。

cat > ./extra_llm_api_options.yaml

推理完成后，查看：

$EXPERT_STATISTIC_PATH 中的导出统计文件。

运行：

examples/ep_load_balancer/report_load_statistics.py 脚本，以显示标准差和失衡比率指标：

python examples/ep_load_balancer/report_load_statistics.py --expert_statistic_path $EXPERT_STATISTIC_PATH

输出如下：

Load statistics: mean std imbalance-ratio3 1024.0 187.955200 0.4980434 1024.0 202.728516 0.5376025 1024.0 209.339981 0.458676...58 1024.0 570.880676 2.46101459 1024.0 341.339447 0.71749860 1024.0 381.045471 1.119648average 1024.0 491.651199 1.564272

第 2 步：生成 EPLB 配置

使用提供的脚本：

examples/ep_load_balancer/generate_eplb_config.py

将采集到的统计数据转换为 EPLB 配置文件。指定目标专家并行度大小 (--ep_size) 和部署时使用的总槽位数 (--num_slots)。例如，如果我们选择保持每个 rank 8 个专家槽位，同时将专家并行度增加到 36 路，则应有 32 个冗余专家和总共 288 个专家槽位。

python examples/ep_load_balancer/generate_eplb_config.py \ --ep_size 36 \ --num_slots 288 \ --expert_statistic_path $EXPERT_STATISTIC_PATH \ --output_path ./moe_load_balancer.yaml

./moe_load_balancer.yaml 文件如下：

initial_global_assignments: 3: [138, 81, 60, ..., 69, 250, 77] 4: [24, 243, 72, ..., 90, 251, 52] 5: [120, 162, 246, ..., 14, 192, 171] ... 58: [67, 70, 160, ..., 212, 103, 125] 59: [45, 142, 152, ..., 99, 205, 49] 60: [34, 162, 119, ..., 234, 26, 129]num_slots: 288layer_updates_per_iter: 0

第 3 步：使用 EPLB 配置运行推理

设置一些环境变量：

# Set a new expert statistic data pathexport EXPERT_STATISTIC_PATH=./expert_statistic_eplb# Enable counting of routed expert IDs from iteration 100 to iteration 200export EXPERT_STATISTIC_ITER_RANGE=100-200

使用 EPLB 配置运行 36 路专家并行推理：

cat > ./extra_llm_api_options_eplb.yaml

再次运行脚本：

examples/ep_load_balancer/report_load_statistics.py

输出如下：

Load statistics: mean std imbalance-ratio3 1024.0 37.612328 0.0819474 1024.0 42.367714 0.0932565 1024.0 42.623219 0.092623...58 1024.0 49.167507 0.11342059 1024.0 44.529514 0.09231460 1024.0 48.408348 0.101029average 1024.0 53.976442 0.115378

注意：统计专家 ID 会大幅影响性能，因此在进行基准测试或生产环境运行时，请记得通过取消设置 EXPERT_STATISTIC_ITER_RANGE 来禁用该功能。

其他

共享显存清理：为实现在线负载平衡，所有专家权重都存储在共享 Host 内存中。同一节点上的 4 个 rank 共享相同的专家权重以节省显存。通常，这些共享 Host 显存在进程退出时会被清理。但如果发生异常退出，它们可能无法被清理。在这种情况下，可能需要手动检查 /dev/shm 目录并删除 /dev/shm/moe_shared_* 中的任何文件（如存在）。

拓展思考

我们高度认可 DeepSeek 团队的系统创新。他们将大规模 EP 支持引入内部推理系统，并以开放态度与社区分享工程洞察，这种做法和精神极其宝贵并且大大提升了推理系统设计性能。同时需要指出的是，系统设计与优化（如大规模 EP）并无万能解决方案。根据当前性能分析，当计划应用大规模 EP 时，需考虑以下因素：

MoE GroupGEMM 计算时间是否为端到端的性能瓶颈？大规模 EP 主要通过减轻专家权重负载压力缩短 MoE GroupGEMM 的执行时间，从而提升 MoE GroupGEMM 层的计算强度。如果在工作负载设置中 MoE GroupGEMM 计算不是瓶颈，那么大规模 EP 可能帮助不大。延迟限制大规模 EP 主要在存在严格延迟限制时发挥作用，尤其是在具有更大显存容量的 GPU 上。对于显存容量较小的 GPU，在延迟约束较弱的场景下，大规模 EP 仍可通过提升并发度和实现更高的 Token / 秒 / GPU 数发挥作用。可用的硬件规格大规模 EP 的优化配置取决于 GPU 规格，包括显存带宽、容量、GPU 间带宽及算力等。这些因素决定了是否采用大规模 EP 以及理想的并行度。系统复杂性和生产部署限制在没有故障容错保证的情况下，大规模 EP 可能会增加在线系统的故障率。当某些大规模 EP 服务实例故障时，即使可以通过集群级协调将流量传送到其他运行中的服务实例，大规模 EP 单实例部署所需的大量 GPU 也会增加系统级部署难度。

未来，我们计划总结并分享更多大规模 EP 技术部署的最佳实践。

请根据自身判断决定是否将大规模 EP 集成到系统中，以及在使用时确定何种 EP 规模和具体部署设置最适合您的需求。

当前 TensorRT-LLM 大规模 EP 的实现方式尚不完善，仍存在已知的限制（欢迎至 TensorRT-LLM 社区帮助进行改进）。例如，我们需要：

更广泛的平台支持扩展对其他 GPU 硬件的支持。目前大型 EP 支持仅覆盖 NVFP4 数据精度，需要开展进一步的工作将支持范围扩展到 FP8 和 INT8 / INT4 数据精度。性能进一步的性能调优和优化。进行更多接近生产流量的工作负载验证。我们期望 TensorRT-LLM 社区在这方面的反馈，以帮助我们基于更具体的工作负载校准 TensorRT-LLM 大规模 EP 实现。对与其他推理核心功能的组合进行全面验证，例如分离服务、推测性解码、验证更多 MoE 模型系列等。易用性易于定制

➢ 我们认为大规模 EP 可至少分为两层：

➔ 由推理引擎开发者开发的内核层。该层包含自定义 EP 通信内核、CPU 与 GPU 之间的同步逻辑以及 MoE 权重重新分配逻辑。

➔ 策略层，可由推理引擎开发者与机器学习研究人员共同开发。该层包含用于以不同方法收集在线流量统计数据的工具，以及专家最佳复制与放置算法。

➢ 基于此理解，我们计划使接近策略层的组件更易于社区用户扩展和自定义，以鼓励涌现出更好的想法。

我们希望能够基于用户输入的部署要求 (ISL / OSL、延迟约束、硬件规格等)，自动推荐最佳 EP 设置。容错性由于大规模 EP 部署方案可能导致在线部署系统故障率增加，这可能会增加与 NVIDIA GPU 上端到端大语言模型 (LLM) 推理系统的多个组件进行跨层交互的需求，包括低级通信内核、集群级编排器和调度器等。我们正在与 NVIDIA 的多个工程团队积极合作，推动相关进展。

我们认为当前实现方式可作为一种合理的端到端大规模 EP 实现方案，并鼓励社区尝试新思路和性能验证。我们期待着社区能够提供反馈，帮助我们在该领域快速进步。我们正在积极追踪此 GitHub 问题中的 TensorRT-LLM 大规模 EP 执行情况，以保证对社区的公开透明。

致谢

上述大规模 EP 工作是团队又一项巨大成果，涉及内核级优化、运行时增强以及系统性能分析和调优。我们无法逐一感谢每位贡献者，但我们为这支具有奉献精神的工程师团队感到骄傲，正是他们的集体专业知识推动了 TensorRT-LLM 性能的大幅提升。

此次团队合作为我们提供了提高大语言模型推理中 GPU 利用率的宝贵洞察。希望本文介绍的技术和经验能够帮助开发者社区在关键 LLM 推理应用中更加充分地利用 NVIDIA GPU 的性能。

参考

[1] PR 4695:

[2] PR 4384：

[3] PR 4495:

[4] 文档：

https://github.com/NVIDIA/TensorRT-LLM/tree/feat/large-ep/examples/ep_load_balancer#offline-ep-load-balancer

作者