Hybrid-EP: An Efficient MoE Communication Implementation

郁凡, 刘童, NVIDIA GPU加速计算专家团队, 高级工程师 | NVIDIA AI Open Day, Nov. 7th, 2025

议程 (Agenda)

• Hybrid-EP 简介
• 详细实现
• 基准测试结果
• 在 Megatron-Core 中的应用
• 未来路线图

Hybrid-EP 简介

Hybrid-EP 设计目标

• 在混合专家（Mix-of-Expert, MoE）模型架构中高效支持专家并行（Expert Parallel, EP）通信。
• 达到硬件网络带宽极限，同时最小化流式多处理器（SM）的使用，以便为其他计算核心的并发执行留出更多SM资源。
• 充分利用最佳的 NVIDIA 硬件和软件技术。

Hybrid-EP 功能性

Hybrid-EP 库被设计为一个专家并行（EP）通信库，其中包含两个通信算子：dispatch 和 combine。这两个通信算子通常用于混合专家（MoE）模型架构中。

上图展示了 Hybrid-EP 在 Transformer 层中的前向传播和后向传播过程。
- 前向传播:
1. Router 生成本地路由图（Local routing map）。
2. 通过 Allgather 操作汇聚生成全局路由图（Global routing map）。
3. 元数据预处理（Metadata preprocessing）后，dispatch 算子根据路由信息将 Tokens 分发给相应的 Expert MLP。
4. combine 算子将处理后的 Tokens 聚合起来，送入下一个 Attention 层。

• 后向传播:
  1. dispatch 算子将梯度分发。
  2. combine 算子将梯度聚合。

Hybrid-EP 路由图 (Routing Map)

• 在 Hybrid-EP 中，每个 rank（GPU）的 router 会生成一个本地路由图，该图指示了每个 expert 是否需要各个 attention 输出的 token。

• 在 Hybrid-EP 中，每个 RANK 意味着一个 GPU，每个 NODE 意味着一个 peer2peer 域（例如一个 NVLink 域）。

  • 1 x Grace Blackwell NVL72:

    • NUM_OF_RANKS_PER_NODE = 72
    • NUM_OF_NODES = 1

  • 2 x DGX Hopper:

    • NUM_OF_RANKS_PER_NODE = 8
    • NUM_OF_NODES = 2

• 右图展示了一个路由图的示例。

Hybrid-EP Dispatch 演示

下图演示了 Dispatch 操作，它通过 Multicast 将数据分发，然后进行 permute（置换）操作，将 token 发送到目标专家所在的设备。

Hybrid-EP Combine 演示

下图演示了 Combine 操作，它是 Dispatch 的逆过程。它首先进行 unpermute（逆置换），然后通过 Reduce add 操作将来自不同设备的数据聚合起来。

Hybrid-EP 支持的数据格式

• Token 数据格式:

  • 形状 (Shape): HIDDEN_DIM.
  • 数据类型 (Data type):
    • Dispatch: BF16 (不带缩放因子), FP8 (带缩放因子).
    • Combine: BF16.

• 概率向量 (Prob vector) 数据格式:

  • 形状 (Shape):

    • Attention 侧: NUM_OF_TOTAL_EXPERT.
    • Expert MLP 侧: NUM_OF_EXPERTS_PER_NODE (带填充).

  • 数据类型 (Data type): FP32.

• 缩放因子 (Scaling factor) 数据格式:

  • 形状 (Shape): HIDDEN_DIM / 128 (DeepSeek 配方).
  • 数据类型 (Data type): FP32.

详细实现

Hybrid-EP 通信核函数设计

• Hybrid-EP 中的 dispatch 和 combine 核函数是 warp-specialized（Warp 专用）和 persistent（持久化）的核函数。
• Dispatch 核函数: 每个 CUDA block 有 2(无 RDMA) / 3(有 RDMA) 个 warp。
• Combine 核函数: 每个 CUDA block 有 5(无 RDMA) / 11(有 RDMA) 个 warp。
• 持久化核函数: 每个 CUDA block 必须独占一个 SM。
• 一个 CUDA block 内的不同 warp 组构成一个数据流水线。
• CUDA blocks 是完全独立的；数据被均匀地分发到不同的 CUDA blocks 中形成数据块（chunks）。

Hybrid-EP Dispatch 数据流水线设计

• 下图展示了 Hybrid-EP 的 dispatch 核函数数据流水线设计。
• 数据从 Attn input（token, prob, scaling factor）流入，通过 G2S warp group 进入 SMEM Cyclic FIFO，然后由 S2G warp group 处理后输出为 Expert output。
• RDMA warp group 负责与其他节点上相同 rank id 的 ranks 进行数据收发。

Hybrid-EP Combine 数据流水线设计

• 下图展示了 Hybrid-EP 的 combine 核函数数据流水线设计。
• 数据从 Expert input（token, prob）流入，经过 Intra-node G2S warp group，进入 SMEM Intra-node G2S Cyclic FIFO。Intra-node red warp group 和 Inter-node red warp group 进行归约操作。最终由 Inter-node G2S warp group 处理，输出到 Attn output。
• RDMA warp group 负责与其他节点上相同 rank id 的 ranks 进行数据收发。

基准测试结果

Hybrid-EP 性能测试配置和条件

• 测试配置:

  • Token: HIDDEN_DIM = 8192, data type = BF16.
  • 无概率向量或缩放因子。
  • NUM_OF_ATTN_TOKENS_PER_RANK = 4096.
  • NUM_OF_EXPERTS_PER_RANK = 2.
  • 具有均匀分布的随机路由图。
  • TOPK = 4 或 8.

• 硬件配置:

  • 单节点 DGX Hopper。
  • 4x DGX Hopper，每对 GPU 之间有 400Gbps RDMA NIC 连接。
  • Grace Blackwell NVL36x1。

Hybrid-EP 性能

在单节点 DGX Hopper 上的性能

该图表展示了在单节点 DGX Hopper 上的性能测试结果。

• 左图 (TOPK = 4): 随着 CUDA 块数量的增加，调度（dispatch）算法和合并（combine）算法的带宽都随之提升。在16个CUDA块时，调度算法带宽达到350.66 GB/s，合并算法带宽达到357.78 GB/s。
• 右图 (TOPK = 8): 趋势与 TOPK = 4 相似。在16个CUDA块时，调度算法带宽达到368.9 GB/s，合并算法带宽达到368.87 GB/s。
  在两种 TOPK 设置下，合并算法的性能都略优于或持平于调度算法。

在由 400Gbps 网络连接的 4x DGX Hopper 上的性能

该图表展示了在4台通过400Gbps网络连接的 DGX Hopper 上的性能测试结果。图中比较了网卡总线带宽（nic bus bandwidth）和算法带宽（algorithm bandwidth）。
- 左图 (TOPK = 4): 算法带宽显著高于网卡总线带宽。在8个CUDA块时，合并算法的带宽达到了83.2 GB/s。
- 右图 (TOPK = 8): 同样地，算法带宽优于网卡总线带宽。在8个CUDA块时，合并算法的带宽达到了123.4 GB/s。
在跨节点场景下，合并算法的性能仍然略微优于调度算法。

在 Grace Blackwell NVL36x1 上的性能

该图表展示了在 Grace Blackwell NVL36x1 上的性能测试结果。
- 左图 (TOPK = 4): 性能随 CUDA 块数量的增加而显著提升。在32个CUDA块时，调度算法带宽达到693.02 GB/s，合并算法带宽达到625.54 GB/s。
- 右图 (TOPK = 8): 性能同样表现出强劲的增长。在32个CUDA块时，调度算法带宽达到714.66 GB/s，合并算法带宽达到657.72 GB/s。
注：仅供技术讨论和参考，性能可能因产品组合不同而异。

在 Megatron-Core 中的应用

方法论

原始的 Hybrid-EP 内核是 C++ 风格的、基于模板的实现。所有的输入和输出都应该是指针。为了使其成为一个可用的基于 Python 的算子，还需要几个步骤：

1. 通过 Hybrid-EP 计算所有类型缓冲区所需的缓冲区大小并分配足够的内存。

   • (1) 当前，Hybrid-EP 假设最坏情况来计算大小。

2. 预处理相关的寄存器/固定（pin）操作以准备相应的缓冲区。

3. 交换内存句柄。
4. 使用适当的值初始化元数据。
5. 使用 JIT (Just-In-Time compilation) 来避免模板问题。

我们已经发布了 Hybrid-EP 作为 DeepEP 的一个分支，并且 Hybrid-EP 已经集成到 Megatron-Core 的 dev 分支中。

缓冲区管理

上图展示了缓冲区管理的流程：
1. 初始化: HybridEPBuffer 接收 hidden_dim, num_of_tokens_per_rank, num_local_experts, use_fp8 等参数来创建 HybridEPBuffer-CPP 实例。
2. 管理: HybridEPBuffer-CPP 负责管理缓冲区，包括：尺寸计算、内存分配（Malloc）、缓冲区注册、句柄交换。主要管理的有输入/输出缓冲区等。
3. 运行: 当接收到 Torch::Tensor 类型的 token, probs, routing_map 后，HybridEPBuffer 会生成一个 Instance Config。
4. JIT编译与执行: Instance Config 用于 JIT 编译内核，然后进行预处理（Preprocess）和调度（Dispatch）。
5. 输出: 调度结果的句柄（Handle）用于反向传播或合并（backward/combine）过程。

缓冲区管理（续）

为了使缓冲区能被其他 rank 访问，需要进行特殊的注册，并且缓冲区句柄必须在 EP (Expert Parallelism) 组内进行交换。
流程如下：当前 rank 的普通 GPU 内存 -> 注册/获取句柄 -> 通信 -> 在其他 rank 打开句柄 -> 可访问。

因此，输入缓冲区和输出缓冲区应在不同场景下进行注册：
- 节点内（intra-node）使用普通缓冲区。
- 通过 RDMA (Remote Direct Memory Access) 通信时使用注册缓冲区。

Memcpy 优化

朴素的 D2D (Device-to-Device) 拷贝实现存在以下问题：
- 注册的缓冲区不能由 torch 的分配器管理。
- 注册过程耗时，不能在每次迭代中执行。
- 结果或输入有时需要存储在注册缓冲区中。在某些情况下，需要进行 D2D 拷贝以将数据传输到 PyTorch 张量中。

上图展示了朴素的实现流程，其中包含了两次 D2D 拷贝。D2D 拷贝引入的开销：
- 仅在 RDMA 上，与通信核相比，会增加约 5% 的额外延迟。
- 从输入到专家以及从专家到输出的整个过程，与通信核相比，会增加约 20% 的额外延迟（在 Hopper GPU 上基于 DeepSeek V3 模型测量）。

Memcpy 优化（续）

为了避免 D2D 拷贝，我们将 permute 和 unpermute 操作添加到了 Hybrid-EP 中以隐藏延迟。这也为未来将 permutation 融合到 Hybrid-EP 内核中提供了便利。
优化后的流程中不再需要 D2D 拷贝。

另一种方案是在 Python 级别暴露注册缓冲区的指针。这种方法不被接受，因为它不符合 PyTorch 的风格，并且需要在框架层面进行更多修改。

Memcpy 优化（续）

对于 FP8 训练，cuBLAS FP8 GEMM 要求在 M 维度上对齐。这种填充（padding）行为实际上是一次 D2D 拷贝，可以被融合到 permute 内核中以提高效率。

如上图所示，在优化前，Dispatch, Permute, Padding 是独立的步骤，耗时分别为 600us, 150us, 180us。优化后，Padding 被融合，Dispatch 耗时 600us，新的 Permute 耗时 170us。此测量基于 Grace Blackwell GPU 上的 DeepSeek V3 模型 (TP1PP8EP32)。

无同步（Sync Free）

上图展示了 GPU 和 CPU 之间的交互。torch.empty() 需要在主机端（host）知道大小值，这导致了两次同步。这里存在一个依赖关系：permute 预处理步骤需要 buffer-size 作为输入，并产生 token-per-expert 数据。

无同步（续）

Hybrid-EP 模型可以将所有输入和输出数据都放在设备端，这可以消除所有的同步。如果用户能为缓冲区大小提供一个合适的上界，就可以实现这一点。

设置输出张量大小的可能方法：
- 使用最坏情况来计算上界。
- Token drop。

我们目前正在开发一种新策略，以在 Megatron 中启用全层 CUDA Graph，其中 hybrid-ep 有望成为一个关键组件。

在 Megatron-Core 中的应用总结

当前 Hybrid-EP 的特性:
- 与 NVL72 兼容。
- 融合了预处理、permute 和 pad 操作。
- 与 CUDA-Graph 兼容。
- 支持 JIT。
- 支持节点内和节点间场景。
- 注意: 当前的 RMDA 通信要求用户为每个 GPU 手动分配最近的 NIC；尚不支持自动检测。

Hybrid-EP 在 LLM 训练上的收益:
测试平台：Grace Blackwell GPU，DeepSeek V3 模型（无 MTP）。
- 强制均衡 (Force balance)
- Hybrid-EP
- 融合 FP8 填充 (Fuse FP8 padding)

性能数据:

Model	Precision	Dispatcher	Feature List	TFLOPS/GPU
DeepSeek-V3	MXFP8	DeepEP	Baseline	829.53
DeepSeek-V3	MXFP8	Hybrid-EP	Hybrid-EP	898.47
DeepSeek-V3	MXFP8	Hybrid-EP	Fuse FP8 Padding	943.56

• Hybrid-EP 带来了约 70 TFLOPS 的提升。
• 融合 FP8 填充带来了约 44 TFLOPS 的提升。

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

未来路线图

路线图

• 优化 NVLink 网络在特殊情况下的内核性能：例如小隐藏层尺寸（<= 7KB token size 等）和非常大的 NVLink 域（EP64 等）。
• 支持自动 NIC 拓扑检测和选择。
• 提升预处理内核和路由图 all gather 的性能。
• 提升可用性：支持 RoCE，并在 IBGDA 不支持时可以回退到 IBRC。

结束

演示文稿以NVIDIA品牌标志页结束。