DeepSeek 发布开源第二弹！让MoE架构效率提升的神助攻【DeepEP】

摘要：本文将针对 DeepEP 项目进行深入浅出的功能解析与设计分析，并在此基础上提出一些潜在的优化思路。本报告分为三个主要部分：功能解析、创新设计点、可能的优化方案。为了便于理解，文中会适度引用部分代码片段或函数接口说明。

---

一、功能解析

DeepEP 旨在为 MoE（Mixture of Experts）及其专家并行（Expert-Parallel）场景提供高效的通信库，核心功能包括：

1. 分发（Dispatch）：将一定数量的 token（通常是深度学习中自注意力机制或 MoE 中的输入数据）按专家、节点或其他策略进行划分，并分发到指定的目标 GPU/节点。
2. 聚合（Combine）：在计算完成后，将分发到多个专家或节点的中间结果按原顺序或策略合并回源节点，以继续后续的计算流程。
3. 低时延模式（Low-Latency Mode）：该模式主要面向推理阶段或需要极低延迟的场景，通过 RDMA、IBGDA 等技术减少通信开销，并提供了一种 Hook 机制，可以在发起发送请求后延迟真正的接收过程，以与其他计算进行流水线重叠。

1. 整体结构概览

从代码组织来看，DeepEP 的核心由 C++/CUDA 层和 Python 层组成：

• C++/CUDA 层（见 csrc/deep_ep.cpp, csrc/deep_ep.hpp 等）

  1. 通过 pybind11 将核心通信操作和内存管理封装为 Python 可调用接口。
  2. 借助 NVSHMEM 以及 cudaIpc 等机制，实现了在 GPU 间共享或跨节点的高带宽/低延迟数据交换。
  3. 引入了 Buffer 类（deep_ep::Buffer）来统一管理通信流、FIFO 任务队列、本地和 RDMA 缓冲区，以及 CPU 端用作监控的计数器等。

• Python 层（见 deep_ep/buffer.py、deep_ep/utils.py 等）

  1. 提供了易于集成的高级接口，包括常见的 dispatch、combine 函数，屏蔽了底层通信细节。
  2. 定义了 Buffer（Python 版本）的封装类，在初始化时，根据用户参数分配 NVLink 或 RDMA 等所需的通信内存，并与底层 C++/CUDA 组件完成同步。
  3. 提供了 set_num_sms 等配置接口，让使用者可根据集群的 SM 资源进行性能调优。

整体而言，DeepEP 代码的核心流程是：

1. 在 Python 层初始化 Buffer 对象，并根据需求传入 num_nvl_bytes, num_rdma_bytes 等参数。
2. 在 C++ 层的 Buffer 构造函数中，申请 GPU 上/跨节点的通信内存（比如通过 cudaMalloc 或 NVSHMEM）。
3. 使用 sync 等接口，完成进程间、设备间的句柄交换与通信初始化。
4. 在分发/聚合等函数（例如 intranode_dispatch, internode_dispatch）中，启动对应的 kernel 执行，将数据在不同 GPU 或节点之间搬运。
5. 在低时延场景下，通过 low_latency_dispatch 与 low_latency_combine 走一条更轻量的 RDMA 通道，牺牲一定带宽以换取显著的端到端延迟优势。

2. 核心类与重要函数

下面以 deep_ep::Buffer 类中的部分函数为例，简要概括其功能。

（1）构造与销毁：Buffer::Buffer(...) 与 Buffer::~Buffer()

Buffer::Buffer(int rank, int num_ranks, int64_t num_nvl_bytes, int64_t num_rdma_bytes, bool low_latency_mode):
    rank(rank), num_ranks(num_ranks), ...
{
    // 申请本地或跨节点的 buffer；根据使用模式分配 FIFO、workplace 等
    // 与 NVSHMEM、cudaIpcMemHandle 等关联
}

• 构造函数里通过 cudaMalloc、cudaIpcGetMemHandle、NVSHMEM 等步骤建立了一个节点内/节点间的通信共享内存。
• 还会初始化 CPU 端的 moe_recv_counter、moe_recv_expert_counter 等，便于计算接收的 token 数量。
• 若处于低时延模式，还会针对 IBGDA 作额外初始化。

Buffer::~Buffer() noexcept(false) {
    // 回收资源，等待可能还在运行的 kernel 完成
    // 关闭远端 IPC 句柄 / 释放 NVSHMEM / cudaFree
}

（2）数据同步：Buffer::sync(...)

此函数主要在初始化阶段调用，用于在进程组内交换 IPC 句柄或 NVSHMEM 唯一 ID：

void Buffer::sync(const std::vector<int> &device_ids, const ...){
    // 1. 同步 IPC handles，打开远端显存句柄
    // 2. 如果需要 RDMA，则初始化 NVSHMEM 并分配对应内存
    // 3. 标记 ready，后续即可执行 dispatch/combine
}

在完成 sync 之前，buffer 尚不可用；一旦完成，is_available() 会返回 true。

（3）分发与聚合（通用）：dispatch / combine

在 deep_ep.cpp 中可以看到多对函数，如 intranode_dispatch / intranode_combine（仅走 NVLink）与 internode_dispatch / internode_combine（同时走 NVLink + RDMA 或纯 RDMA），它们大体做的事情相似：

1. 根据输入的张量（例如 x, topk_idx, num_tokens_per_rank 等）计算需要发送数据的拆分方式。
2. 在 GPU 上启动一个或多个 kernel，将数据打包并发送到各目标 GPU 的接收 buffer 中。
3. 通过计数器、barrier 或其他回调，等待所有数据到达后，完成目标张量的拼接或加和操作。

例如 intranode_dispatch 的签名（略简化）：

std::tuple<torch::Tensor, ..., std::vector<int>, torch::Tensor, ...>
Buffer::intranode_dispatch(
   const torch::Tensor& x, 
   const std::optional<torch::Tensor>& x_scales,
   const std::optional<torch::Tensor>& topk_idx,
   ...
) {
    // 1. 计算要发送给每个 rank 的 token 数
    // 2. 发起 NVLink 通信
    // 3. 返回接收后的张量
}

此过程中，代码会依赖 cudaStream 来实现异步执行，还可支持在调度流上进行内存申请及事件记录，以满足高性能需求。

（4）低时延分发与聚合：low_latency_dispatch / low_latency_combine

这是 DeepEP 的一大特色，主要用于推理阶段或需要极端低延迟的场景。在这类函数中：

1. 会将输入的 BF16 张量在发送端打包为 FP8 类型；
2. 通过 RDMA（以及 IBGDA 技术）直接发送到远端，不再使用 NVLink 或复杂的分片；
3. 接收端可以选择先发起发送请求，延后接收，避免占用 SM 资源，从而与其他计算做并行流水。

std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, std::optional<EventHandle>, std::optional<std::function<void()>>>
Buffer::low_latency_dispatch(..., bool return_recv_hook) {
    // 若 return_recv_hook = true，则只发起发送；调用者可在稍后合适的时间再执行 hook 完成真正接收
}

在代码中，可以看到会构造两个 buffer（偶数与奇数），交替使用，以避免同一时刻需要多次复用同一段内存时产生冲突。

---

二、创新设计点

1. 多种通信方式融合

   • 同时支持 节点内（NVLink IPC）和 节点间（RDMA / NVSHMEM）两种路径，可在大规模集群中高效部署。
   • 针对节点内多 GPU 通讯，可用 cudaIpcMemHandle 方式直接进行 P2P 读写，省去多次跳转。
   • 在节点间通信时，默认启用 NVSHMEM，自动屏蔽底层 RDMA 细节。

2. 低时延和高吞吐双管齐下

   • 提供了常规 all-to-all（面向训练或大批量推理）与低时延 all-to-all（面向高并发小批量推理）两类 kernel，实现了对“带宽”与“延迟”这两个关键指标的兼顾。
   • 普通模式下追求峰值吞吐量，利用 NVLink 和 RDMA 的带宽；低时延模式下利用 IBGDA 强化 RDMA 点对点性能，尽量缩短端到端通信时间。

3. Hook 机制解耦通信与计算

   • 在低时延模式下，可只“发起”通信请求，随后在需要的时刻再执行“接收”操作（return_recv_hook=true），这在自定义 pipeline 并行或 CUDA Graph 场景下具有较高自由度。

4. 显式流管理与事件机制

   • 允许将所有工作放到单独的 comm_stream 里，减少与默认计算流的互斥；
   • 支持 async_finish 参数，以事件进行同步或异步隔离，满足一边传输、一边计算的并行需求。

5. 代码可读性与可扩展性

   • 通过 Config 类封装了 chunk size、SM 数量等硬件相关的参数。
   • 采用 pybind11 进行接口导出，Python 层更容易被上层深度学习框架集成。

---

三、可能的优化方案

1. 进一步的自适应调优

   • 当前代码中，对于 chunk size、并行度、流数量等只是给了默认推荐值或手动配置。可以考虑基于集群与模型规模实现自动化调优，对 num_max_nvl_chunked_recv_tokens，num_max_rdma_chunked_send_tokens 等参数进行搜索与测量，以适配实际带宽和延迟特性。

2. 减少 CPU 等待的瓶颈

   • 在普通模式下，CPU 端需要先等待 GPU 写入的计数器（moe_recv_counter 等），这会导致调度上 CPU-GPU 间的同步点。可以尝试用 CUDA Graph 或者类似 GPU 触发回调的机制来取代 CPU 轮询，以减少 CPU 主线程的阻塞开销。

3. 支持更多网络拓扑结构

   • 目前主要针对单层 NVLink + IB 卡。可扩展到更复杂的多节点多 GPU 拓扑，例如多机多卡环形或其它高维拓扑。如果 GPU 直连顺序与 NVLink 拓扑不同，也可考虑在调度层进行映射优化。

4. 智能调度 + 拆包合并

   • 对于分发/聚合的数据，如果能够根据网络负载或专家负载自适应调度，也许能提升整体效率。比如在相邻 token 都发给同一专家时，自动合并发送请求。

5. 异构设备或跨版本兼容

   • 当前主要针对 Hopper（H100 / H800）等新架构以及 CX7。在其它 GPU（如 A100）或 Mellanox 网卡上使用 RDMA，需要对 IBGDA 和 AR（Adaptive Routing）的支持情况进一步验证。

---

结语

DeepEP 项目通过 CUDA 内核、NVLink IPC 与 NVSHMEM RDMA 等多种手段，极大提升了专家并行（MoE）场景下的分发与聚合效率。它在高吞吐和低时延方面都做了专门优化，提供了健全而灵活的接口设计。用户既可在大规模预训练等需要极高带宽的场景下使用常规模式，也可在在线推理等对延迟敏感的场景下选用低时延模式。

通过进一步的自动化调优、异构拓扑支持以及智能化的拆包合并策略，DeepEP 有望更好地适配未来规模更大、结构更复杂的分布式深度学习系统，为高效率的 MoE 模型训练和推理提供更加完备的通信解决方案。