TensorRT-LLM Large-scale Expert Parallelism Optimizations

Enwei Zhu (朱恩伟), NVIDIA 加速计算专家团队 高级工程师
Jinyang Yuan (袁劲飏), NVIDIA 加速计算专家团队 高级工程师

议程 (Agenda)

• 动机: 为何需要大规模专家并行 (EP)
• 大规模专家并行的挑战
• 专家并行负载均衡器
• 通信优化及其他
• 端到端性能

动机：为何需要大规模专家并行 (Large-scale EP)？

大规模专家并行（EP）可以在以下方面提供帮助：

1. 提升计算强度：在给定的单GPU工作负载下，扩展EP规模可以减少分组GEMM（通用矩阵乘法）的内存访问，从而提高计算强度。
2. 改善内存密集型操作性能：对于分组GEMM中受内存带宽限制的范围（即高每用户吞吐率/tps/user），大规模EP可以提升其性能。
3. 提升并发能力：对于受内存容量限制的GPU，大规模EP可以通过实现更高的并发度，在帕累托曲线的左侧区域（高并发）提供帮助。

下图展示了在GB200上的预期帕累托曲线，比较了EP规模为4和8时的性能差异。

注：此图仅供技术讨论和参考，实际性能可能因产品组合不同而异。

大规模专家并行的挑战

通信 (Communication)

在不同并行策略组合下，通信开销是关键挑战。

• 小规模EP:

  • Attention TP + MoE EP: 在MoE部分使用AllReduce通信。
  • Attention DP + MoE EP: 在MoE部分使用AllGather后接ReduceScatter的通信模式。

• 大规模EP:

  • Attention TP: 可扩展性不佳。
  • Attention DP + MoE EP: AllGather和ReduceScatter的组合是否仍然高效存疑。

AllGather + ReduceScatter 模式

这种模式下，每个令牌（token）的隐藏状态（hidden states）会被发送到所有的计算排名（rank），但实际上只有拥有被选中专家（topK experts）的rank才需要这些数据，造成了通信冗余。

AlltoAllv 模式

AlltoAllv是一种更高效的通信原语。在这种模式下，每个令牌的隐藏状态只会被发送到那些拥有其所需专家的rank，从而减少了不必要的数据传输。

专家负载不均衡 (Expert Load Imbalance)

专家负载不均衡是大规模EP面临的另一个核心挑战。

• 在某些层中，负载不均衡可能成为性能热点。
• 负载较低的rank需要等待其他rank完成通信，导致性能瓶颈。
• 负载不均衡的行为在不同层和不同数据集中表现各异。
• 可能的原因:
  • 某些专家比其他专家“热门”得多。
  • 多个中等热度的专家被分配到了同一个rank上。

下图展示了在翻译任务（模型：DeepSeek-R1，EP规模：32）中，不同rank和层的令牌分配数量热力图，以及特定rank上专家负载的柱状图，清晰地揭示了负载不均衡现象。

专家并行负载均衡器 (Expert Parallelism Load Balancer)

为了解决专家负载不均衡问题，我们设计了专家并行负载均衡器（EPLB）。

离线负载均衡器 (Offline Load Balancer)

设计思路

• 离线统计: 在服务启动前，通过运行推理来离线收集专家使用情况的统计数据。

• 生成配置文件: 基于收集到的统计数据，生成EPLB配置文件。

  • 该文件定义了从逻辑槽位（slot ID）到物理专家（expert ID）的映射关系。
  • 可以使用冗余专家来更好地平衡负载。
  • 槽位总数等于正常专家数加上冗余专家数。
  • 选择映射关系的目标是使分配到每个rank的令牌数尽可能接近平均值。

• GPU数量解耦: 在线服务和离线统计收集所用的GPU数量可以不相同。

• 配置约束:
  • 离线收集统计时，专家数量需要能被EP规模整除。
  • 在线服务时，槽位数量（而非专家数量）需要能被EP规模整除。

实验结果

• 有效性: 离线负载均衡器在在线数据分布与离线数据分布相似时非常有效。
• 效果: 负载在槽位/rank之间变得均衡，不再出现热点槽位/rank。

下图对比了禁用和启用EPLB后的专家负载情况。启用后（右图），负载分布明显更加均匀。

在线负载均衡器 (Online Load Balancer)

观察

离线均衡器依赖于数据分布的稳定性。然而，实际场景中数据分布可能动态变化，这促使我们研究在线负载均衡。

• 迭代间相似性: 在相邻的几次迭代中，专家负载不均衡的模式非常相似。下图展示了在翻译数据集和GSM8K数据集上，专家负载热力图随迭代（y轴）的变化，呈现出明显的垂直条纹，表明负载模式具有短期稳定性。

• 请求间差异性: 在单个请求内部，负载模式在相邻迭代中相似，但不同请求之间的模式差异很大，即使是来自同一数据集的请求。

设计思路

基于以上观察，我们设计了在线负载均衡器，以适应动态变化的负载模式。

• 在线统计: 在服务期间在线收集专家统计数据，采用带衰减率的在线求和方法，更加关注最近的令牌。
• 冗余专家: 同样可以利用冗余专家来更好地平衡负载，槽位总数等于正常专家数加冗余专家数，且槽位总数需能被EP规模整除。

• 动态更新:

  • 每层中从槽位ID到专家ID的映射信息会进行周期性的在线更新。
  • 选择映射的目标是使分配到每个rank的令牌数接近平均值。
  • 不同rank上的专家权重也相应地在线更新。

• 后台执行: 映射信息和专家权重的更新在后台执行，与其它层的推理计算重叠，以减少对性能的影响。

• CUDA Graphs支持: 相关内核经过专门设计以支持CUDA Graphs。

后台更新机制

• 可以在每次迭代中选择更新特定数量的层（例如，下图中的2层）。
• 每个MoE层都会被周期性地更新（例如，下图中需要29次迭代来更新第3至第60层）。
• 更新操作由一个独立的“更新线程”执行，该线程利用“推理线程”在处理某些层时的空闲时间片来更新其他层的权重，从而实现计算与更新的重叠。

在线负载均衡器：设计

在线负载均衡器通过在运行时动态调整专家（Expert）的放置和权重，以解决负载不均衡问题。其工作流程涉及CPU和GPU的协同，具体如下：

• CPU端 在时间步 n 的工作包括：

  1. 等待并接收来自GPU的结束信号 (Wait End Signal)。
  2. 根据GPU提供的统计数据 (Statistics Data)，计算专家副本和放置策略 (Replication & Placement Compute)。
  3. 更新权重和放置信息 (Update Weights & Placement)。
  4. 将更新后的放置信息 (Placement Info) 和 MoE 权重 (MoE Weights) 发送给GPU，并为下一个时间步 n+1 发送更新开始信号 (Update Start Signals)。

• GPU端 在时间步 n 的工作流程如下：

  1. 接收CPU的启动信号 (GPU Start Signal)。
  2. 执行 Wait -> Statistics (统计负载信息) -> Expert-Slot Mapping (根据Placement Info进行映射) -> Prepare -> Dispatch -> MoE (使用MoE Weights进行计算) -> Combine -> End。
  3. 将收集到的 Statistics Data 发送回CPU，用于下一步的决策。

• 一个关键约束是，在GPU执行从接收启动信号到发送结束信号之间的核函数时，主机（CPU）不应更改放置信息和MoE权重，以确保数据一致性。

在线负载均衡器：结果

在线负载均衡器能有效解决负载不均衡问题。实验结果表明：
- 在权重更新后，负载在不同的槽（slots）/等级（ranks）之间变得均衡，不再出现热点槽/等级。
- 下图展示了在翻译数据集上的效果，其中红色横线表示权重更新发生的层。更新前，存在明显的负载不均衡（黄色条带）；更新后，负载分布变得非常均匀（紫色/蓝色区域）。
- 实验设置：翻译数据集，ep_size=32, local_bs=256, num_slots=288 (32个冗余专家)。

专家并行负载均衡器：结论

下表对比了离线和在线两种负载均衡器的优缺点及适用场景：

通信优化及其他

在解决了负载均衡问题后，下一个优化重点是通信开销。

通信优化：基线分析 (AllGather + ReduceScatter)

在优化前，基线采用 AllGather + ReduceScatter 的通信模式。性能分析显示：
- 计算部分：MoE的分组GEMM（2 moe_gemm）和Attention（attn）核函数的执行时间表现良好。随着专家并行（EP）规模的增大，GEMM的耗时甚至会减少，而Attention的耗时保持不变。
- 通信部分：AllGather和ReduceScatter操作的耗时随着EP规模的增大而显著增加。然而，在每个GPU批处理大小固定的情况下，所需的通信消息大小应为常数，这表明当前的通信方式存在优化空间。

性能剖析图显示，AllGather（152us + 149us）和 ReduceScatter（390us）占用了大量的执行时间，成为性能瓶颈。

通信优化：AlltoAll 初步实现

为了优化通信，引入了 AlltoAll 操作，其特点包括：
- 通过NVLink进行点对点（P2P）访问。
- 支持设备内存中的发送/接收计数，实现 AlltoAll 风格的通信。
- 支持CUDA Graph。

新的流程变为 Prepare -> A2A -> MoE -> A2A -> Local Reduce。性能剖析显示，All2All 的通信时间（fp4: 18us+12us, bf16: 174us）远低于之前的 AllGather 和 ReduceScatter。但该实现可能会因 recv_m ~selectedExpert * m 而引入新的负载不均衡问题。

通信优化：AlltoAll 优化

在初步实现的基础上，进行了三项进一步的优化：
1. 使用 AlltoAll 重构 prepare AllGather：将用于准备专家ID和令牌尺度的AllGather操作也替换为更高效的AlltoAll。
2. 低精度合并：在发送端进行量化（Quant），在接收端进行反量化（DeQuant），以减少数据传输量。
3. 融合 AlltoAll 核函数：将多个小的AlltoAll核函数融合成一个，减少核函数启动开销。

优化1的效果显著，如下图左侧所示，“current”（当前）版本的耗时远低于“previous”（之前）版本。优化2和3的实现细节如右侧图所示，通过共享内存（SM）和P2P工作空间进行高效的数据交换。

其他优化

MoE 辅助核函数优化

• 问题：在优化前，诸如Permute/Activation/Unpermute等辅助核函数的执行时间会随着EP规模的增大而增加。理论上，在每个GPU批处理大小固定的情况下，这些核函数的工作负载应该是恒定的。

• 根本原因：在MoE GEMM之前，M个令牌被扩展为M*topK个令牌。平均而言，只有1/EP的扩展令牌是有效的（即被路由到当前rank上的专家）。当EP规模很大时，大多数线程块被启动用于处理无效令牌，造成了严重的资源浪费。例如，对于DS-R1在EP32的设置下，96个扩展令牌中只有3个在rank 0上是有效的。

• 解决方案：

  1. 只为有效令牌启动线程块。
  2. 使用持久化线程块（persistent thread blocks）来处理数据依赖的有效令牌数量，并保持与CUDA Graph的兼容性。

• 优化效果：优化后，expandInputRows、dyActivation和finalizeMoeRouting等辅助核函数的执行时间不再随EP规模增加而增加，而是保持为一个常数，显著提升了性能。

MTP LM Head 张量并行 (Tensor Parallelism)

• 场景：在DeepSeek R1 MTP-3模型中，采用32路注意力数据并行（DP）+ MoE专家并行（EP），每个rank的批处理大小为32。
• 问题：在这种设置下，MTP（Mixture of Tensors and Pipelines）模块的LM Head GEMM操作批处理大小较小（32），容易成为内存带宽瓶颈（memory-bound），而非计算瓶颈（math-bound）。经验表明，256个令牌是内存与计算瓶颈的边界。
• 优化思路：在内存带宽受限的情况下，减少内存负载可以提高性能。

• 解决方案（TP=4）：
  1. 引入4路张量并行（TP）。在每4个GPU上执行AllGather，将每个rank的本地批处理大小从32增加到32*4=128，使GEMM操作更接近计算瓶颈。
  2. 将LM Head的权重沿输出维度切分以支持4路TP。
  3. 执行本地LM Head GEMM，然后通过AllGather和切片操作合并logits。
  4. 为进一步减少通信，先在本地计算argmax，然后仅对argmax的结果进行AllGather和全局argmax。

• 更多低精度计算

  • 在 NVFP4 中实现 Attention WO GEMM
  • FP8 上下文 FMHA

• 更多融合与重叠

  • 融合 Q/K/V 拼接和 MLA
  • 融合 Q/K RoPE 和 NoPE 拷贝/拼接
  • 融合来自路由和共享专家的输出归约

• 主机侧优化

  • 减少 pybind 和进程间通信开销
  • 流式处理区间

端到端性能

• 在 ISL/OSL = 1k/1k，MTP（Model Tensor Parallelism）禁用的情况下的帕累托曲线。

• 在 ISL/OSL = 8k/1k，MTP 禁用的情况下的帕累托曲线。

• 在 ISL/OSL = 8k/1k，MTP 启用的情况下的帕累托曲线。