NCCL EP 论文解读

汹萃热 · 昨天 18:08

NCCL-EP: NCCL统一的专家并行API

相关的代码在：Code
相关的论文：NCCL EP: Towards a Unified Expert Parallel Communication API for NCCL

本来很期待这篇论文能带来什么新的insights，但是最后感觉非常遗憾，很多地方都说的很不明白，在最后性能比DeepEP差了一点的时候解释的原因竟然是说“DeepEP采用pytorch的性能观测来分析性能，我们的性能差在Kernel唤起时带来的launch延迟”，这个说法显然是无法立足的。笔者猜测性能的差距很有可能在 GIN 上，但是笔者没有任何设备进行可靠的验证...总而言之，还是再继续做知识的“盗窃工”吧。
AI声明：本文采用glm-5.1辅助完成（虽然大部分都是笔者自己写和核对的但是AI确实指明了很多大致的方向。今后采用了AI阅读源码的blog笔者都会给出显著的AI声明）。

NCCL-EP: NCCL统一的专家并行API
- 摘要
- 背景介绍
  - 贡献
- 设备API，GIN
- NCCL-EP架构和API设计
  - 与其他现有EP并行库的对比
  - 主要设计想法
  - 核心功能
    - 函数功能分类
    - send_only与双buffer下的执行流水线简介
  - 低延迟Kernel
    - Moe模型
    - DeepEP 通信模型
    - NCCL-EP在 DeepEP 上做出的调整和优化
  - 高吞吐Kernel
    - Kernel的架构
    - Dispatch kernel
    - Combine Kernel
    - GIN
    - Buffer管理
      - Dispatch Shared Buffer
      - Combine Shared Buffer
  - LL模式和HT模式的区别
- 与Megatron-LM（训练架构）和vllm（推理架构）的集成
- 性能评估
- 总结

摘要

NCCL-EP提供了一个NCCL原生的EP并行通信库API。

统一的ncclEpDispatch,ncclEpCombine源语
支持LL（低延迟）和HT（高吞吐）模式。
- 前者应用于小batch模式（decode过程），采用直接的alltoall RDMA+NVLink mesh连接方式，通过双buffer实现dispatch和combine的交叠。
- 后者通过大batch模式（训练+prefill过程），采用层次化通信模式，将token通过NVLink域聚集在一起，再进行RDMA通信。
- 两种模式都使用了Device API，用于节点内和节点间通信，利用了拓扑感知性和优化后的GPU发起部署。

背景介绍

传统的Moe通信模式：NCCL的alltoall，CPU协调的流水线。
- token经过routing决定的topK进行重排，通过AlltoAll通信，然后在计算完后反重排。
- Megatron-Core和Tutel都采用这样的方案作为dispatcher。
现有的工作：设备发起通信
- 将这些分散的步骤融合成统一的dispatch和combine核心，GPU直接发起RDMA源语同时不将控制权转交给主机。
- 减少延迟，实现核内量化、数据移动交叠、输出布局优化。
NCCL-EP主要提供了：
- 基于NCCL设备API部署。
- 统一了ncclEpDispatch和ncclEpCombine源语，基于工作负载特征选择合适的算法。
- 提供了灵活性，具有C/Python接口，抽象了硬件复杂度。

HT模式：高吞吐模式，主要用于训练+prefill阶段。token数目在4096以上。使用了层次化的通信手段：NVLink在节点内聚合token，然后再通过RDMA传输。
LL模式：低延迟模式，主要用于decode阶段。使用了完全的alltoall mesh连接性，以及双缓冲区通信模式。

贡献

统一API接口和算法模式。
- ncclEpDispatch/ncclEpCombine 通过 ncclEpCreateGroup 在创建时选择算法。
两层的资源管理。将资源从操作中分离出来。
- ncclEpCreateGroup_t 管理跨操作共享的通信buffer和网络连接。
- ncclEpHandle_t 捕获前向传播前的路由决定。
设备发起通信。依托GPU发起通信，通过nccl设备API发起。
- 允许了在排除CPU的情况下开启异步RDMA通信的能力。
异构的通信道路。根据需要选择最佳的路径：NVLink用于节点内，RDMA用于节点间。Nvlink不可用时回退到PCIe。
LL设计优化。提供了一个减少LL内核中内存占用的机制。

接下来的结构大致如下：

我们跳过有关Moe架构的介绍，仅看看设备API和GIN相关。
介绍NCCL EP内的架构，API设计。
介绍LL是怎么优化内存占用的。（感觉这部分更有趣）
介绍HT是怎么部署的。（HybridEP相关）
介绍兼容性和如何继承进入ML框架，以及部署考量。
我们跳过性能介绍，而是分析它是如何评估性能的。
相关工作，以及GIN。

设备API，GIN

Device API主要有这三个部分，本图片来自于GPU Initiated Network

LSA部分：提供通过PCIe或者Nvlink直接访问远程GPU的内存方式。这种一般都是机内的片间通信。通过 ncclCommGetLocalSymm 来获取伙伴GPU的内存地址，再通过经典的load/store来进行远程内存访问。
GIN部分：允许GPU直接对远程节点执行RDMA操作，具有put/get/SignalAdd/WaitOnSignal操作，主要负责单边写，单边读，原子增加远程信号量用于同步，等待信号量叨叨预值的作用。这些操作完全在GPU上执行，减轻了CPU代理县城的开销，允许在核内细粒度交叠通信和计算。这对于Moe的逐token的工作负载非常有效。（CPU上下文不需要频繁切换）
Multimem部分：内存的多播操作，对于广播/规约都很有效。

这样，NCCL提供了完整的设备API生态，有别于单独的NVSHMEM或DOCA GPUNetIO，NCCL EP利用了NCCL独有的拓扑探测，网络插件架构，同时也获得了设备侧发起通信的能力，解决了生态割裂，设备绑定，API复杂的三大难题。
NCCL-EP架构和API设计

与其他现有EP并行库的对比

主要设计想法

原生C-API与python绑定。这样允许NCCL直接集成到C++运行时（TensorRT-LLM，vLLM原生后端）的同时，允许python通过c类型的wrapper绑定。
用户控制的资源生命周期。用户可以显式控制资源分配和销毁。同时提供了可预测的内存使用量，允许集成自主内存分配器，对于推理服务系统跨模型管理内存池非常重要。
基于Stream的异步执行。所有的NCCL-EP操作在CUDA stream上异步执行，与NCCL执行模型保持一致（所有的操作都通过GroupStart和GroupEnd，通过enqueue和dequeue完成）。避免了隐式的同步开销，影响性能。
硬件抽象。抽象了在dispatch和combine内的操作，允许了未来对网络后端的拓展。
GPU加速通信。通过计算/通信融合到cuda核中利用GPU并行性加速。Kernels 映射到 GPU 架构上，将token处理过程分布到可用的SM上。

核心功能

函数功能分类

资源管理
- ncclEpCreateGroup
  - 从NCCL comm中创建Moe通信组，配置算法模式(HT/LL)，专家数量，缓冲区大小，集合调用。
  - 可以创建 ncclEpGroup_t 进行资源管理。
    - 代表了上层资源的容器，通过现存的NCCL comm创建。这个通信组的配置取决于算法模式(HT/LL)。专家的数量，每个rank最大的token数目，buffer大小相关的参数。它包含了通信算法，通信缓冲区，网络资源(比如QP)三个部分。
  - 通信组的创建是一个集合操作，每个rank都需要有合理的配置，持续整个模型周期。
  - 你可以通过 alloc_fn, free_fn 来集成自定义的内存池。NCCL EP将会使用这些来进行针对不同模型的内存管理。
- ncclEpGroupDestroy
  - 释放通信组的资源，包括缓冲区和网络连接。
- ncclEpCreateHandle
  - 用routing信息创建每个前向传播的handle。
  - 在每个pass的过程中维护当前轮次的状态信息。通过ncclEpCreateHandle()和路由信息topk_idx创建。
  - 在HT模式下，handle的创建将触发元数据在各个rank之间的交互，计算token的分布。在LL模式，元数据触发则会在dispatch中隐式进行。
  - 在匹配的dispatch和combine模式中将会持续共享。
- ncclEpHandleDestroy
  - 在dispatch/combine完成后释放handle。
通信
- ncclEpDispatch
  - 将token发送给制定的专家。send_only=1启动阶段执行模式，也就是dispatch将会在发起dispatch传输后，不等待操作结束就释放GPU资源，允许计算/通信交叠。
  - dispatch 将会把token发送给指派的专家。接收三个tensor：
    - Inputs: token和路由需要的metadata。
    - Outputs：预先分配的接收buffer。
    - 本地tensors：可选的逐rank metadata，例如token数目。
- ncclCombine
  - 将token收集起来，规约成原有的顺序。这里采用了存储在handle中的routing状态，来实现反排列。同样也具有send_only来启用阶段执行模式。
- ncclEpComplete
  - 结束阶段操作(LL模式)，将会阻塞等待那些异步执行的操作直到完成。
- ncclEpHandleGetNumRecvTokens
  - 查询收到的token数量，用于buffer分配。这样可以实现根据输出的buffer大小管理。
算法模式
- LL(低延迟模式)：主要为了优化延迟。
  - 数据维度呈现为3维，专家主序的数据排布。在LL模式下，可以通过send_only来启动阶段化执行，允许专家交叠计算和数据传输。并且，LL通过双buffer来交叠执行dispatch和combine阶段。
- HT(高吞吐模式)：主要为了优化吞吐，需要利用聚合和网络架构优势。
  - 精准的buffer大小需要通过在初始化Handle的时候交换元数据。
  - 数据维度呈现2维模式，并且带有每个专家拥有的token数目信息，保证了可重现的训练下的确定性排序。

张量抽象
- 为了实现前面所述的不同的数据维度和排布方式，以及显式的计数，NCCL设计了多维数据布局来描述数据并且区分他们的功能。
- ncclNDTensor_t 就是用于描述这些张量的。在这其中，主要的信息
  - 布局信息（维度数目和每个维度的大小），
  - 语义角色（当前张量的用途）
  - 数据类型（FP32/BF16/FP16/FP8）
  - 描述非连续向量特征的stride。
- ncclEpTensorTag_t：在创建Tensor的时候用于描述张量在dispatch和combine中用途的标识。例如：区分输入输出，区分token数据和元数据，辅助NCCLEP验证张量大小。
  - _TOKENS: 输入/输出数据。
  - _TOPK_IDX: top-k 专家的索引。
  - _TOPK_WEIGHTS: top-k 路由权重。
  - _SCALES: 缩放数据。这里用于量化过程中的scale。
  - _RECV_EXPERT_COUNTER_DEVICE: 设备侧每个专家token数量。
  - _RECV_EXPERT_COUNTER_DEVICE: 主机侧每个专家的token数量。
  - _TOKENS_PER_EXPERTS: 每个专家的token计数。
  - _NONE: 没有使用/默认值
- Tensor Layouts: 张量布局通过ndim来确定维度数目，sizes来确定各个维度的大小，strides描述张量非连续分布或者连续分布的方式。
  - 对于token: 2D [num_tokens x hidden]
  - 对于路由元数据: 2D [num_tokens x top_k]
  - 对于输入数据，数据布局均相同。
  - 对于输出数据：
    - 在低延迟模式下为3D的数据，[num_local_experts x tokens_per_expert x hidden] 为其布局，通过接收token的专家来进行聚合。专家主序布局允许直接输入到GEMM计算中。
    - 在高吞吐模式下为2D的数据，[total_recv_tokens x hidden] 为其布局。来自所有的专家的token都将被拼接起来。另一个输出将会带有每个expert token的计数信息，表明了tensor边界。
    - Dispatch模式和Combine模式下的输出张量都根据算法的对应具有上述的张量布局。

使用方式：HT下的训练和LL下的推理

HT是论文中的图。
LL模式是笔者根据vllm和deepep猜测而来的，如有纰漏恳请批评指正。主要参考了：

vllm/model_executor/layers/fused_moe/modular_kernel.py

vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ll.py

vllm/v1/worker/ubatching.py

DeepEP/deep_ep/buffer.py

DeepEP/csrc/deep_ep.cpp

DeepEP/csrc/kernels/internode_ll.cu

双批次重叠

融合Moe模块化内核

在论文中，给出了HT模式使用时的示例代码和LL模式使用的示例代码。

// 训练（HT模式），在推理prefill时可以去除backward部分代码。
// Once per model
ncclEpCreateGroup(&group, comm, &config, stream);
// Forward pass: creating one handle/layer in ep[]
for (int l = 0; l < L; ++l) {
topk = route(tokens);
ncclEpCreateHandle(&ep[l], group, topk);
ncclEpDispatch(ep[l], tokens, exp_in, ...);
expert_ffn_forward(exp_in, exp_out);
ncclEpCombine(ep[l], exp_out, tok_upd, ...);
tokens = normalize(tokens, tok_upd);
}
// Backward pass: reusing handles from ep[]
loss = calc_loss(ideal, tokens);
for (int l = L - 1; l >= 0; --l) {
ncclEpDispatch(ep[l], loss, exp_in, ...);
expert_ffn_backward(exp_in, exp_out);
ncclEpCombine(ep[l], exp_out, loss_upd, ...);
loss = update_loss(loss, loss_upd);
ncclEpHandleDestroy(ep[l]);
}
// 推理decode
// Micro-batch 0: dispatch (send only)
ncclEpDispatch(ep[0], tokens[0], exp_in[0], ..., /*send_only=*/1, stream);
// Micro-batch 1: dispatch while 0 transfers
ncclEpDispatch(ep[1], tokens[1], exp_in[1], ..., /*send_only=*/1, stream);
// Complete micro-batch 0 receives and run experts
ncclEpComplete(ep[0], NULL, stream);
expert_ffn_forward(exp_in[0], exp_out[0]);
ncclEpCombine(ep[0], exp_out[0], tok_upd[0], ..., /*send_only=*/1, stream);
// Complete micro-batch 1 receives, and run experts
ncclEpComplete(ep[1], NULL, stream);
expert_ffn_forward(exp_in[1], exp_out[1]);
ncclEpCombine(ep[1], exp_out[1], tok_upd[1], ..., /*send_only=*/1, stream);
ncclEpComplete(ep[0], NULL, stream);
// Continue pipeline...

复制代码

send_only与双buffer下的执行流水线简介

这一章节是作者自己的理解。如有纰漏恳请批评指正。

文章中提到了“启用send_only可以实现计算-通信交叠。我们结合代码来稍微看一下send_only的实际表现。
Key 1: send_only 只作用在LL模式下的ncclDispatch和ncclCombine模式下。
Key 2: send_only实现了专家计算和数据传输过程的重叠。
在开启send_only的情况下，send和recv阶段将分开成两个GPU Kernel进行执行。这样结合代码，我们的运行流水线如下图所示，可以实现“专家计算和数据传输”的交叠。

Key 3: 结合双buffer实现真正的计算/通信融合
这和我们看到过的DeepEP的那张图是一样的。只不过需要注意的是实际上Dispatch和combine的send/recv的执行时间很短（只进行快速的处理和解包与内存搬运），因此实际上的宏观时间占比非常少。这样我们在实际上的GPU执行流程是这样的，也就是现有DeepEP/ncclEP的在LL下的可能执行模式：

低延迟Kernel

NCCL EP LL 模式面向推理过程中的decode阶段，一般都是小batch（1-128个tokens），降低端到端的延迟是它的主要优化目标。NCCL-EP借鉴了DeepEP，并依托原有的NCCL设备侧API集成到NCCL通信库中。
Moe模型

NCCL建立的模型如下：

假设了E个专家逐块分布在N个rank上，每个rank有 $L=\lceil E/N \rceil$ 的专家。
在每个rank $r$ 上，注意力子层生成 $B$ 大小的token。
对于每个token $t\in[0, B)$, Moe的路由定义了一个向量 $R(r,t)$ 表示 $t$ 将会被路由到的K个专家。

通信发生：DP -> EP -> EP

每个数据并行的rank $r$ 将与所有的专家进行通信，形成E个有效的专家-rank通信对 $\Gamma_r^{DP}=\{(e,r)\,|\,e\in[0,E)\}$，并且 $rem^{DP}(e,r)=\lceil e/N \rceil 是远程的rank。(DP -> EP)
对于每个专家 $e$ 将会与所有的rank进行通信。这样，$r$ 拥有L个专家也会形成 $E$ 个通信对 $\Gamma_r^E=\{(e,r')\,|\,r'\in[0,N)\}$, 远程通信对定义为 $rem^E(e,r)=r$.
在每个通信对之间的通信token数目设为 $m^s(e,r), s\in\{R,E\}$.

DeepEP 通信模型

拓扑
- LL采用全mesh模式，每个GPU都会通过最高效的方式与其他GPU进行通信(机内nvlink/PCIe, 机间RDMA)，允许路由的高灵活性来满足Moe中的动态专家分布。
通信buffer
- LL模式分配了两个内部的buffer，用于双缓冲区模式来交叠 dispatch 和 combine 操作。具体的交叠方式可以看前面的双buffer部分。
- 在DeepEP中，我们的buffer分布（不要看论文了，论文写的太抽象，以下是作者的理解）
- 注意：此处在recv中的msg则对应了发送端的slot.
- 对于DeepSeek v3来说，我们的 hidden = 7168，这样我们就有：
  - 对于Dispatch: 14352Bytes(BF16) 或者 7408Bytes(FP8)
  - 对于Combine: 14336Bytes(BF16) 或者 8960~14336(LogFMT 压缩)

对于LogFMT，读者可以详询：Insights into DeepSeek-V3: Scaling Challenges and Reflections on Hardware for AI Architectures
对于BF16，FP8数据的大小请详询：Transformer Engine。笔者不太了解量化，这部分还是知识短板(其实读者通信存储计算和量化都不懂，哈哈)。
对于DeepSeek采用的量化策略，主要分为Block量化和tile量化。这两种量化分别作用与模型权重和输入输出激励。具体可以参考Deepseek-v3 Technical Report图六和图七。

通信模式：dispatch和combine具有类似的结构。
- 他们都被分成 send 和 recv 阶段
- 并且通过细粒度的计数器 update-and-flush 操作。
- dispatch/send + combine/recv -> DP侧唤起
- dispatch/recv + combine/send -> EP侧唤起。
- 我们将其泛化成 $s\in\{DP, E\}$
发送阶段。
- 在发送区域，我们主要经过：准备数据/发送数据/信号通知处理三个阶段。
- 对于Dispatch：
  - 首先，我们将数据进行读取并写入缓冲区。并且执行核内量化（如果开启了FP8），通过 bar.sync 等待所有线程准备完成数据。
  - 其次，利用RDMA对数据进行发送。在LL核心中，所有的GPU都会执行通信操作，通过RDMA或者Nvlik/PCIe进行通信操作。
  - 最后，将会利用 update+flush 操作来实现信号通知。
接收阶段。
- 在接收的时候，rank r将会观测所有有效的专家-rank通信对的计数器。当计数器更新时，说明对端写入操作已经完成，这时候我们就可以从中提取数据了。

DeepEP怎么做的？我们结合代码简单分析了他们的通信方式。

我们需要注意：图上是针对了dispatch过程中的send和recv过程分别进行讲解而做的示意图。实际上所有需要执行dispatch的设备sm都会参与两个过程。并且，在使用了 send_only=1 (DeepEP中是recv_hook来决定是否需要分阶段进行)的时候，不需要block sync（因为block会退出）。

并行模式：我们分成三个部分：Dispatch/send, Dispatch/recv+Combine/send, Combine/recv 来进行分析。
- 并行模式通用：自下而上的并行的层级为 lane 0 -> warp -> warp group(expert) -> SM(block). 每个rank上将会分配多个 warp group，也就是多个专家。
- Dispatch/send：我们主要分为三个部分：发送token，计数token以及消息通知，
  - 发送token：一个SM，假设其有 $w$ 个warp，那么其 $[0, w-1]$ 的warp中的所有的线程都会参与token的发送。每个SM中的前 $w-1$ 的warp都会参与发送。
  - 计数token：每个SM的最后一个warp负责针对每个warp组的专家计数。因为我们每个SM上都有可能有 $w_g < w$ 个warps组成一个组，这个组专门负责计数应该发给对应专家的token数量。特别地，如果这是SM0，还要负责获取QP数量和清楚缓冲区的工作。
  - 发送通知（更新counter）：每一个warp组的warp 0(第一个warp)的thread 0(lane 0)负责发送信号并且通知对端发送完成。
- Dispatch/recv+combine/send: 等待token+复制token两个部分，其余部分和上面的dispatch/send相似。
  - 等待token：这是通过每个warp组内的warp 1(第二个warp)的thread 0(lane 0)来进行等待。可以看到这里实现了类似“自旋锁”的等待方式。当超时的时候将标记成为故障rank。
  - 复制token：所有的warp都将负责数据的拷贝和量化处理。特别地，每个warp的lane 0将负责加载相关的rank信息。
  - 发送token/计数token/发送通知的并行方式类似。
- Combine/Recv: 此处将会重新组织warps和线程。最后一个warp用于TMA加载数据，其余的warp用于规约计算。
- 所有以上的组织方式都形成了完美的以warp为粒度的多层级流水线：发送token -> 计数token -> 发送通知 -> 等待token -> 复制token/处理 -> 发送token -> 加载token -> 规约计算 -> ...

NCCL-EP在 DeepEP 上做出的调整和优化

设备唤起的通信：
- DeepEP采用IBGDA+nvshmem的方式，GPU唤起IB协议通信发送token和信号，并通过nvshmem这种跨rank共享内存的方式来实现rank之间的同步。主要的原理如下：GPU的SM直接将数据生产写入显存(1)，通过nvshmem的方式发起通信。此时，nvshmem直接在GPU显存中构造WQ(2)，SM更新DBR(3)并直接通知网卡DB(4)，网卡从WQ中poll得到WQE(5)，根据WQE中的元数据去获取显存中的Data(6)并发送出去(7)，发送完成后写入CQ。
- NCCL-EP采用的是自己的设备侧API。通过GIN来实现将不同的设备侧API整合形成网络，并通过(1)数据移动操作 (2)完成追踪操作(本地完成counter，远程完成signal, rank同步屏障barrier) (3) 状态管理操作(重置计数器/信号)来建立与DeepEP对等的语义。后端主要采用了GDAKI+DOCA GPUNetIO的方式执行。（在RDMA下GDAKI就是IBGDA，但是这里GDAKI可以控制以太网，也就意味着普通的网络协议也可以采用GDAKI）
通信buffer的布局优化：
- 我们的DeepEP中采用的是双buffer机制，将分配所有专家，所有batch，所有rank-expert对的双缓冲区大小$O(EBP)$。但是，不是所有的专家都会占用所有的slot，这样将会导致通信buffer的浪费。
- 因此，NCCL-EP打算拓展消息的消息头，避免基于布局的路由。在路由头中，我们的路由地址仅依赖于token和rank，不依赖于主要的专家。因此dispatch buffer将会缩小到 $O(N\cdot B\cdot P)$，我们只根据rank进行布局。
- 对于combine，这样的效果将更好。我们原先的总体buffer为 2EBP, 现在缩小到了 NBP+BKP。这样在N=64，E=512，K=8（K就是top-K数目）下，我们可以得到14倍的优化。

为什么我要强调“打算”呢？是因为我还没有看到对应的NCCL_EP的相关代码是怎么做的。具体参见issue，后面我看到了再补吧。（又挖坑）

高吞吐Kernel

这一部分NCCL主要是采用了HybridEP来实现。（不过HybridEP本来就是他们家的hh）
Kernel的架构

主要采用了warp specialization 持续性kernel（和deepep一样的），让每个block都专有地占有一个单独的SM，并且实现完整的数据流水线。主要分为三类warp组：RDMA组，G2S组以及S2G组。

G2S: global memory -> shared memory; S2G 反之。

接下来逐步分析每个kernel是如何工作的。

Dispatch kernel

此处的图片来源于 Scalable Training of Mixture-of-Experts Models with Megatron Core.

初始化异步屏障。mbarrier 用于实现在CTA中的thread异步机制。这主要是为了TMA的异步拷贝进行细粒度的同步。
RDMA warp group: 在具有多个节点的时候，节点间的传输将通过 N2N_warp_group_device_function 来进行传输。在这其中，元数据将首先在具有相同的本地id的node之间进行交换，随后才会在节点内转发数据。这样实现了降低跨节点流量同时允许跨节点和节点内传输的交叠。后端是GPUNetIO。
G2S warp group: 在等待来自远程的节点同步后（通过mbarrier），以及本地的数据同步后，我们将会把数据从全局内存搬运到共享内存中。
smem cyclic fifo: 经典的双缓冲区共享内存，通过mbarrier进行屏障同步，多阶段实现G2S -> S2G的交叠。num_of_stages形成了环形队列，stage指针用于表明写入的当前的共享缓存fifo地址。parity代表不同的信号，分别用于防止WAR问题(加载完成)和RAW问题(可重复使用)，正好在生产者(G2S)和消费者(S2G)中形成arrive-wait语义。
- 在此处，本地的token和远程传输来的token将根据路由表写入到正确的全局内存中。
S2G warp group: 将数据从共享内存中搬运回全局内存，准备进行下一步的combine。

Combine Kernel

在combine中，因为涉及到规约操作，因此此处设计的warp group将进一步细分。我们主要还是把握住三类warp: RDMA warp组，G2S组，规约warp组。但是此处，对于G2S和规约warp组将有节点内和节点间的区别。因此总共有五类warp groups。

HybridEP实现了将规约融合到通信kernel中。我们可以看到，在节点间存在warp组进行规约，在节点内页存在规约，最后将规约得到的数据直接通过RDMA warp组写入到远程内存中。
在跨节点操作中，HybridEP回首先执行跨节点规约操作，通过RDMA进行传输后再进行节点内的规约操作。这是因为RDMA传输时间较长，通过内存屏障细粒度控制可以更好掩盖通信时间。

GIN

同样，nccl-ep只是把原先的直接采用Doca-GPUNetIO的后端换成了自己的Gin API而已。这里我们按下不表。
Buffer管理

在HybridEP中，分别有 dispatch 的动态内存管理 buffer 和 combine 的动态内存管理buffer。dispatch针对数据类型，多节点和是否为前向传播阶段进行了单独的不同实现。combine则针对多节点，是否为反响传播阶段进行了不同的实现。
Dispatch Shared Buffer

主要包括了：

节点内token的buffer(Stages * 隐藏层维度)
节点内的稀疏 -> 稠密的映射buffer的映射表，表示在当前rank下的输出位置。（双缓冲实现G2S->S2G交叠）
prob缓冲区（仅用于前向传播阶段，是专家的top-k概率）
缩放值(FP8需要的)
路由表（决定哪些节点需要传输给远程节点，只有多节点使用）
内存屏障信号m_barrier
- 数据的内存屏障: G2S -> S2G 的生产完成，以及 S2G -> G2S 的消费完成。多阶段。
- 映射表的屏障，用于TMA。
- S2G组内的屏障，要求S2G完全消费映射表，将数据写入到对应位置后才能继续进行。
RDMA的MR信息以及跨节点的数量（确定我每次需要构造多少个WQE）。

Combine Shared Buffer

主要包括了

四组Token缓冲区：节点间/节点内的G2S/S2G缓冲区
四组Prob缓冲区：仅在反向传播中启动。
三组内存屏障，分别保证节点内的G2S内存搬运完成，节点间的G2S内存搬运完成，以及在节点内操作完成后通知RDMA warp进行RDMA发送。
RDMA的MR消息以及跨节点数量（确定我每次需要构造多少个WQE）。
标记当前阶段的token是否是reduce的最后一个元素。用于规约warp。

LL模式和HT模式的区别

与Megatron-LM（训练架构）和vllm（推理架构）的集成

与megatron的集成如下：

与vllm的集成猜测应该大致如下，参考fused Moe Kernel
对应的结构简单来说大致如下（忽略了很多细节）

性能评估

在Dispatch吞吐上，仅有8节点表现略低于DeepEP。在combine上则表现不如DeepEP。
在延时上，vllm集成后的DeepEP表现均优于NCCLep。

他们声称这是因为kernel唤起有开销，而DeepEP采用的是pytorch kineto的性能检测。NCCLEP反应的是主机侧C API的时间。（emm...这个难道不能通过预热+多次iteration取平均来得到结果吗？对这个evaluation抱有怀疑）

总结

NCCL EP集成并利用NCCL本身具有的优势，统一API，通过C进行高性能的实现和python绑定提供接口，继承了DeepEP LL和Hybrid EP HT的方式来实现了属于自己的一套Moe通信库。但是其中的优化点目前并没有在NCCL中体现出来，其性能表现看起来也有些差强人意。希望后续能有更好的发展，笔者也会持续跟进。

来源：程序园用户自行投稿发布，如果侵权，请联系站长删除
免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！

账号		自动登录	找回密码
密码			立即注册

NCCL EP 论文解读

相关帖子

签约作者

NCCL EP 论文解读

相关帖子

相关推荐

签约作者