摘要NCCL-EP提供了一个NCCL原生的EP并行通信库API。统一的ncclEpDispatch,ncclEpCombine源语支持LL低延迟和HT高吞吐模式。前者应用于小batch模式decode过程采用直接的alltoall RDMANVLink mesh连接方式通过双buffer实现dispatch和combine的交叠。后者通过大batch模式训练prefill过程采用层次化通信模式将token通过NVLink域聚集在一起再进行RDMA通信。两种模式都使用了Device API用于节点内和节点间通信利用了拓扑感知性和优化后的GPU发起部署。背景介绍传统的Moe通信模式NCCL的alltoallCPU协调的流水线。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。设备APIGINDevice API主要有这三个部分本图片来自于GPU Initiated NetworkLSA部分提供通过PCIe或者Nvlink直接访问远程GPU的内存方式。这种一般都是机内的片间通信。通过ncclCommGetLocalSymm来获取伙伴GPU的内存地址再通过经典的load/store来进行远程内存访问。GIN部分允许GPU直接对远程节点执行RDMA操作具有put/get/SignalAdd/WaitOnSignal操作主要负责单边写单边读原子增加远程信号量用于同步等待信号量叨叨预值的作用。这些操作完全在GPU上执行减轻了CPU代理县城的开销允许在核内细粒度交叠通信和计算。这对于Moe的逐token的工作负载非常有效。CPU上下文不需要频繁切换Multimem部分内存的多播操作对于广播/规约都很有效。这样NCCL提供了完整的设备API生态有别于单独的NVSHMEM或DOCA GPUNetIONCCL EP利用了NCCL独有的拓扑探测网络插件架构同时也获得了设备侧发起通信的能力解决了生态割裂设备绑定API复杂的三大难题。NCCL-EP架构和API设计与其他现有EP并行库的对比主要设计想法原生C-API与python绑定。这样允许NCCL直接集成到C运行时TensorRT-LLMvLLM原生后端的同时允许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_only1启动阶段执行模式也就是dispatch将会在发起dispatch传输后不等待操作结束就释放GPU资源允许计算/通信交叠。dispatch 将会把token发送给指派的专家。接收三个tensorInputs: 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.pyvllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ll.pyvllm/v1/worker/ubatching.pyDeepEP/deep_ep/buffer.pyDeepEP/csrc/deep_ep.cppDeepEP/csrc/kernels/internode_ll.cu双批次重叠融合Moe模块化内核在论文中给出了HT模式使用时的示例代码和LL模式使用的示例代码。// 训练HT模式在推理prefill时可以去除backward部分代码。// Once per modelncclEpCreateGroup(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 transfersncclEpDispatch(ep[1], tokens[1], exp_in[1], ..., /*send_only*/1, stream);// Complete micro-batch 0 receives and run expertsncclEpComplete(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 expertsncclEpComplete(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下的可能执行模式低延迟KernelNCCL EP LL 模式面向推理过程中的decode阶段一般都是小batch1-128个tokens降低端到端的延迟是它的主要优化目标。NCCL-EP借鉴了DeepEP并依托原有的NCCL设备侧API集成到NCCL通信库中。Moe模型NCCL建立的模型如下假设了E个专家逐块分布在N个rank上每个rank有 ⌈/⌉ 的专家。在每个rank 上注意力子层生成 大小的token。对于每个token ∈[0,), Moe的路由定义了一个向量 (,) 表示 将会被路由到的K个专家。通信发生DP - EP - EP每个数据并行的rank 将与所有的专家进行通信形成E个有效的专家-rank通信对 Γ{(,)|∈[0,)}并且 (,)⌈/⌉ 是远程的rank。(DP - EP)对于每个专家 将会与所有的rank进行通信。这样 拥有L个专家也会形成 个通信对 Γ{(,′)|′∈[0,)}, 远程通信对定义为 (,).在每个通信对之间的通信token数目设为 (,),∈{,}.DeepEP 通信模型拓扑LL采用全mesh模式每个GPU都会通过最高效的方式与其他GPU进行通信(机内nvlink/PCIe, 机间RDMA)允许路由的高灵活性来满足Moe中的动态专家分布。通信bufferLL模式分配了两个内部的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对于BF16FP8数据的大小请详询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侧唤起。我们将其泛化成 ∈{,}发送阶段。在发送区域我们主要经过准备数据/发送数据/信号通知处理三个阶段。对于Dispatch首先我们将数据进行读取并写入缓冲区。并且执行核内量化如果开启了FP8通过bar.sync等待所有线程准备完成数据。其次利用RDMA对数据进行发送。在LL核心中所有的GPU都会执行通信操作通过RDMA或者Nvlik/PCIe进行通信操作。最后将会利用updateflush操作来实现信号通知。接收阶段。在接收的时候rank r将会观测所有有效的专家-rank通信对的计数器。当计数器更新时说明对端写入操作已经完成这时候我们就可以从中提取数据了。DeepEP怎么做的我们结合代码简单分析了他们的通信方式。我们需要注意图上是针对了dispatch过程中的send和recv过程分别进行讲解而做的示意图。实际上所有需要执行dispatch的设备sm都会参与两个过程。并且在使用了send_only1(DeepEP中是recv_hook来决定是否需要分阶段进行)的时候不需要block sync因为block会退出。并行模式我们分成三个部分Dispatch/send, Dispatch/recvCombine/send, Combine/recv 来进行分析。并行模式通用自下而上的并行的层级为 lane 0 - warp - warp group(expert) - SM(block). 每个rank上将会分配多个 warp group也就是多个专家。Dispatch/send我们主要分为三个部分发送token计数token以及消
NCCL EP 论文解读
摘要NCCL-EP提供了一个NCCL原生的EP并行通信库API。统一的ncclEpDispatch,ncclEpCombine源语支持LL低延迟和HT高吞吐模式。前者应用于小batch模式decode过程采用直接的alltoall RDMANVLink mesh连接方式通过双buffer实现dispatch和combine的交叠。后者通过大batch模式训练prefill过程采用层次化通信模式将token通过NVLink域聚集在一起再进行RDMA通信。两种模式都使用了Device API用于节点内和节点间通信利用了拓扑感知性和优化后的GPU发起部署。背景介绍传统的Moe通信模式NCCL的alltoallCPU协调的流水线。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。设备APIGINDevice API主要有这三个部分本图片来自于GPU Initiated NetworkLSA部分提供通过PCIe或者Nvlink直接访问远程GPU的内存方式。这种一般都是机内的片间通信。通过ncclCommGetLocalSymm来获取伙伴GPU的内存地址再通过经典的load/store来进行远程内存访问。GIN部分允许GPU直接对远程节点执行RDMA操作具有put/get/SignalAdd/WaitOnSignal操作主要负责单边写单边读原子增加远程信号量用于同步等待信号量叨叨预值的作用。这些操作完全在GPU上执行减轻了CPU代理县城的开销允许在核内细粒度交叠通信和计算。这对于Moe的逐token的工作负载非常有效。CPU上下文不需要频繁切换Multimem部分内存的多播操作对于广播/规约都很有效。这样NCCL提供了完整的设备API生态有别于单独的NVSHMEM或DOCA GPUNetIONCCL EP利用了NCCL独有的拓扑探测网络插件架构同时也获得了设备侧发起通信的能力解决了生态割裂设备绑定API复杂的三大难题。NCCL-EP架构和API设计与其他现有EP并行库的对比主要设计想法原生C-API与python绑定。这样允许NCCL直接集成到C运行时TensorRT-LLMvLLM原生后端的同时允许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_only1启动阶段执行模式也就是dispatch将会在发起dispatch传输后不等待操作结束就释放GPU资源允许计算/通信交叠。dispatch 将会把token发送给指派的专家。接收三个tensorInputs: 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.pyvllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ll.pyvllm/v1/worker/ubatching.pyDeepEP/deep_ep/buffer.pyDeepEP/csrc/deep_ep.cppDeepEP/csrc/kernels/internode_ll.cu双批次重叠融合Moe模块化内核在论文中给出了HT模式使用时的示例代码和LL模式使用的示例代码。// 训练HT模式在推理prefill时可以去除backward部分代码。// Once per modelncclEpCreateGroup(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 transfersncclEpDispatch(ep[1], tokens[1], exp_in[1], ..., /*send_only*/1, stream);// Complete micro-batch 0 receives and run expertsncclEpComplete(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 expertsncclEpComplete(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下的可能执行模式低延迟KernelNCCL EP LL 模式面向推理过程中的decode阶段一般都是小batch1-128个tokens降低端到端的延迟是它的主要优化目标。NCCL-EP借鉴了DeepEP并依托原有的NCCL设备侧API集成到NCCL通信库中。Moe模型NCCL建立的模型如下假设了E个专家逐块分布在N个rank上每个rank有 ⌈/⌉ 的专家。在每个rank 上注意力子层生成 大小的token。对于每个token ∈[0,), Moe的路由定义了一个向量 (,) 表示 将会被路由到的K个专家。通信发生DP - EP - EP每个数据并行的rank 将与所有的专家进行通信形成E个有效的专家-rank通信对 Γ{(,)|∈[0,)}并且 (,)⌈/⌉ 是远程的rank。(DP - EP)对于每个专家 将会与所有的rank进行通信。这样 拥有L个专家也会形成 个通信对 Γ{(,′)|′∈[0,)}, 远程通信对定义为 (,).在每个通信对之间的通信token数目设为 (,),∈{,}.DeepEP 通信模型拓扑LL采用全mesh模式每个GPU都会通过最高效的方式与其他GPU进行通信(机内nvlink/PCIe, 机间RDMA)允许路由的高灵活性来满足Moe中的动态专家分布。通信bufferLL模式分配了两个内部的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对于BF16FP8数据的大小请详询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侧唤起。我们将其泛化成 ∈{,}发送阶段。在发送区域我们主要经过准备数据/发送数据/信号通知处理三个阶段。对于Dispatch首先我们将数据进行读取并写入缓冲区。并且执行核内量化如果开启了FP8通过bar.sync等待所有线程准备完成数据。其次利用RDMA对数据进行发送。在LL核心中所有的GPU都会执行通信操作通过RDMA或者Nvlik/PCIe进行通信操作。最后将会利用updateflush操作来实现信号通知。接收阶段。在接收的时候rank r将会观测所有有效的专家-rank通信对的计数器。当计数器更新时说明对端写入操作已经完成这时候我们就可以从中提取数据了。DeepEP怎么做的我们结合代码简单分析了他们的通信方式。我们需要注意图上是针对了dispatch过程中的send和recv过程分别进行讲解而做的示意图。实际上所有需要执行dispatch的设备sm都会参与两个过程。并且在使用了send_only1(DeepEP中是recv_hook来决定是否需要分阶段进行)的时候不需要block sync因为block会退出。并行模式我们分成三个部分Dispatch/send, Dispatch/recvCombine/send, Combine/recv 来进行分析。并行模式通用自下而上的并行的层级为 lane 0 - warp - warp group(expert) - SM(block). 每个rank上将会分配多个 warp group也就是多个专家。Dispatch/send我们主要分为三个部分发送token计数token以及消
