TODO概述

Deepep

参考：

https://www.msn.cn/zh-cn/news/other/%E5%A6%82%E4%BD%95%E7%9C%8B%E5%BE%85deepseek%E6%96%B0%E5%BC%80%E6%BA%90%E7%9A%84%E7%AC%AC%E4%BA%8C%E4%B8%AA%E9%A1%B9%E7%9B%AEdeepep/ar-AA1zRFBQ?ocid=BingNewsSerp

https://zhuanlan.zhihu.com/p/26204046487

https://www.51cto.com/article/809246.html

这篇文章主要深入分析了 DeepSeek 开源的 EP 通信库 DeepEP，特别是其在 专家并行 (Expert Parallelism, EP) 模式下的实现细节、技术挑战和优化策略。

1. DeepEP 与 EP 并行的核心目的

• 解决 MoE 模型的瓶颈：传统的模型并行（如张量并行 TP、流水线并行 PP）难以高效处理混合专家模型（MoE）中动态选择专家的特性，容易导致 GPU 计算利用率低和显存带宽成为瓶颈。
• 优化通信与计算：DeepEP 是一个专门为 MoE 模型训练和推理设计的 EP 通信库，旨在提高计算利用率，减少 GPU 闲置，使更大规模的 MoE 模型训练和推理更高效。
• 显存带宽是瓶颈：文章指出，如果为处理少量 token 而加载高达 44MB 的专家权重，会非常不划算，并且容易占满显存带宽。

2. DeepEP 的关键技术实现

文章从代码层面详细剖析了 DeepEP 的核心组件：

• 用于通信的 Buffer：分析了用于接收数据的 packed tensors（如 packed_recv_x, packed_recv_src_info）的分配和使用。
• 高吞吐 Kernel (用于训练和 Prefill)：
  • Dispatch 通信流程：详细拆解了从 Notify_Dispatch 到 Intranode::dispatch（节点内）和 Internode::dispatch（节点间）的完整流程。
  • Combine 流程：分析了结果汇总的 Intranode_Combine 和 Internode_Combine 过程。
• 低延迟 Kernel (用于 Decoding)：
  • 关键优势：该 Kernel 与 CUDA Graph 兼容。传统 RDMA 通信会通过 CPU 中断来启动内核，从而打断 CUDA Graph，增加延迟。DeepEP 避免了这一点。
  • Double-Batch Overlapping：通过 hook() 机制实现，可以在不占用 SM（流式多处理器）的情况下进行通信重叠，提高效率。
  • 低延迟 Dispatch/Combine：文章给出了 low_latency_dispatch 函数的 Python 和 C++ 实现，展示了其如何通过 SEND PHASE 和 RECV PHASE 来管理通信。

3. 硬件与网络技术

• RDMA 技术：DeepEP 大量使用了 RDMA（远程直接内存访问）技术来实现高效的节点间通信。
• InfiniBand (IB) 与 RoCE：DeepSeek 在其系统中采用了 InfiniBand 技术。虽然 DeepEP 在 GitHub 上声称理论上兼容 RoCE (RDMA over Converged Ethernet)，但文章指出在 RoCE 上运行会面临诸多挑战，如 Multi-Rail 拓扑问题、incast 拥塞、RC (可靠连接) 兼容性以及 In-Network Computing 的实现难题。

4. 其他重要细节

• FP8 细粒度量化：在通信过程中使用了 FP8 量化来减少数据传输量，并提到了为 TMA (Tensor Memory Accelerator) 加载优化而进行的内存布局调整（如保证 token 数能被 4 整除）。
• 与 DeepSeek-V3 论文的关联：文章提到，要完全理解 DeepEP 的设计和潜在的硬件缺陷，需要结合 DeepSeek-V3 论文中的建议一起看。
• 未来展望：文章作者提到后续会分析 FlashMLA 和 DeepGEMM，后者是用于 MoE 专家矩阵计算的库，与 DeepEP 配合使用。

总结：这篇文章是一篇非常深入的技术分析，揭示了 DeepSeek 在优化大规模 MoE 模型基础设施方面的前沿工作。DeepEP 不仅仅是一个通信库，它通过精心设计的内核、对 CUDA Graph 的兼容性、对 RDMA 的深度利用以及对量化技术的整合，系统性地解决了 MoE 模型在扩展性和效率上的关键挑战。

1. Overview

• 目标：

为 Mixture-of-Experts (MoE) 和 专家并行（Expert Parallelism, EP） 提供高效通信支持，优化大规模分布式训练与推理场景下的 GPU 通信性能。

• 核心能力：
• 高吞吐量的 All-to-All GPU 内核（MoE 的 Dispatch/Combine 操作）
• 低延迟推理内核（纯 RDMA 通信，支持 CUDA Graph）
• FP 8 低精度支持（兼顾计算效率与通信带宽）
• 通信 - 计算重叠（Hook 机制实现零 SM 资源占用）
• 应用场景：
• MoE 模型的训练（如 DeepSeek-V 3）
• 低延迟推理解码（如 LLM 生成任务）

---

2. Preliminaries

2. Mixture-of-Experts (MoE)

MOE 主要由两部分组成：

• Sparse MOE Layers：是用来替换 transformer 结构中的 FFN (Feed Forward Network) 层的, MOE 层有固定数字的专家，每个专家也是一个单独的 Neural Network。
• Router：决定了 token 会进入哪个专家塔中。如下图中 More 这个 token 更多的进入第二个专家塔，而 Parameters 更多的进入第一个专家塔。控制进入哪个专家塔的 Router 作为一个可训练的部分，也会在模型训练中得到学习。

MoE 优势：

- 高效扩展性：增加专家数量即可提升模型容量，而计算成本仅线性增长（而非指数级）。

- 任务适配性：不同专家可学习不同领域知识（如语言、视觉），适用于多模态任务。

- 分布式训练友好：支持专家并行（Expert Parallelism），将不同专家分布到不同设备，突破单设备算力限制。

MOE 也面临如下挑战：

- 动态路由稳定性：门控网络需平衡专家负载（避免某些专家被过度激活或闲置）。

DeepSeekMoE with Auxiliary-Loss-Free Load Balancing，传统方式采用手动 auxiliary loss 会导致很大程度影响模型性能。Auxiliary-Loss-Free Load Balancing 引入了一个 bias，在每个训练 step 监视 expert load 情况，如果某些专家 overload 就减少 bias，underload 就增加 bias

- 通信开销：分布式训练中专家间数据交换需要高效通信库

MoE 使用 EP 在路由选择不同 expert 时，很有可能选择到的 expert 在不同 node，导致对于通信的需求上非常大。deepseek 使用 DeepEP

- 训练难度：需设计专用优化策略（如负载均衡损失、梯度裁剪）。

deepseek solution

2.2 专家并行策略

• All-to-All 通信：

在 MoE 的 Dispatch（数据分发）和 Combine（结果聚合）阶段，设备间交换输入数据和计算结果。

• Dispatch 阶段：将数据从门控网络路由到对应专家所在的设备。
• Combine 阶段：从各设备收集专家计算结果并聚合。
• 低精度通信：

使用 FP 8/FP 16 等低精度格式减少通信量（例如 DeepEP 库支持 FP 8 通信，带宽占用减少 50%）。

专家并行优势：

|模型扩展性|通过增加专家数量而非专家大小扩展模型，突破单设备内存限制。|

|计算效率|稀疏激活机制使每个设备仅需计算部分专家，资源利用率高。|

|通信效率|仅需传输激活的专家数据（而非全量参数），配合低精度通信进一步优化带宽。|

|灵活性|可与数据并行、流水线并行等策略结合，适配不同规模的训练任务。|

挑战：

|负载不均衡|动态门控算法（如 Top-K with Capacity）限制单个专家的最大负载。|

|通信瓶颈|定制通信库（如 DeepEP）优化 All-to-All 性能，支持 NVLink/RDMA 混合拓扑。|

|专家参数同步|异步参数更新（如 ZeRO-3）或定期全局同步（Ring-AllReduce）。|

|训练稳定性|引入负载均衡损失函数（Load Balancing Loss）和梯度裁剪。|

3 DeepEP

IB: InfinitBand，网络层跨节点高速通信，支持 RDMA

NVLink: 设备层节点内多 GPU 通信

IBGDA:

1）高吞吐量、低延迟的 all-to-all GPU 内核，专门优化的分派和组合操作。确保数据在多个 GPU 之间快速传输，减少通信时间。

2）支持低比特操作，如 FP 8 格式，显著降低计算和存储需求，提升整体效率。

3）针对非对称域带宽转发（如从 NVLink 域到 RDMA 域），提供优化内核，适合训练和推理 Prefill 任务。允许直接内存访问，减少 CPU 介入。DeepEP 的优化确保数据在不同域之间高效传输，特别适用于大规模混合卡的分布式训练。

Secondly, we develop efficient cross-node all-to-all communication kernels

to fully utilize IB and NVLink bandwidths and conserve Streaming Multiprocessors (SMs)

dedicated to communication.

由于 cross-node EP 计算通信比 1：1，效率很差，设计了 dualpipe 来做计算通信重叠

为了保证 DualPipe 能有充分的计算性能，定制了一个 cross-node all-to-all communication kernel，保留了几个 SM 做通信

跨界点 GPUs 是与 IB（50 GB/s）全互联的，节点内通信通过 NVLink（160 GB/s）

因此我们限制每个 token 最多被分发到 4 个节点，可以减少 IB traffic

For each token, when its

routing decision is made, it will first be transmitted via IB to the GPUs with the same in-node

index on its target nodes. Once it reaches the target nodes, we will endeavor to ensure that it is

instantaneously forwarded via NVLink to specific GPUs that host their target experts, without

being blocked by subsequently arriving tokens.

这样做，通过 IB 和 NVLink 的通信可以完全重叠。每个 token 可以在每个 node 里平均选择 3.2 个专家，而不会增加额外的 overhead。

虽然实践中 DeepseekV 3 只选择 8 个 experts，这个数字最大可以到 4 nodes x 3.2 experts / node，而不增加额外通信成本。

在这个样的策略下，只需要 20 个 SM 就可以充分利用 IB 和 NVLink 的带宽。

// One channel use two blocks, even-numbered blocks for sending, odd-numbered blocks for receiving.

Warp Specialization

20 个 SM 分成 10 个 communication channels

During dispatching process，不同 warp 做不同 task，并且根据实际 workload，动态调整

1. IB sending
2. IB-to-NVLink forwarding
3. NVLink receiving

During combining process：

1. NVLink sending
2. NVLink-to-IB forwarding and accumulation
3. IB receiving and accumulation

使用 PTX 指令，自动调节 chunk size，减少 L 2 cache 使用，以及对其他 SM 的干扰

Deployment Strategy

redundant experts：

重复那些高负载的 expert

To this end, we introduce a

deployment strategy of redundant experts, which duplicates high-load experts and deploys them

redundantly. The high-load experts are detected based on statistics collected during the online

deployment and are adjusted periodically (e.g., every 10 minutes). After determining the set

of redundant experts, we carefully rearrange experts among GPUs within a node based on the

observed loads, striving to balance the load across GPUs as much as possible without increasing

the cross-node all-to-all communication overhead. For the deployment of DeepSeek-V 3, we set

32 redundant experts for the prefilling stage. For each GPU, besides the original 8 experts it

hosts, it will also host one additional redundant expert.

Finally, we are exploring a dynamic redundancy strategy for experts, where each GPU hosts

more experts (e.g., 16 experts), but only 9 will be activated during each inference step. Before

the all-to-all operation at each layer begins, we compute the globally optimal routing scheme

on the fly. Given the substantial computation involved in the prefilling stage, the overhead of

computing this routing scheme is almost negligible.

展开源码

3.1 Normal Kernels with NVLink and RDMA forwarding

The normal kernels can be used in model training or the inference prefilling phase (without the backward part) as the below example code shows.

3.2 Low-latency Kernels with Pure RDMA

The low latency kernels can be used in the inference decoding phase as the below example code shows.

For two micro-batch overlapping, you can refer to the following figure. With our receiving hook interface, the RDMA network traffics are happening in the background, without costing any GPU SMs from the computation part. But notice, the overlapped parts can be adjusted, i.e. the 4 parts of attention/dispatch/MoE/combine may not have the exact same execution time. You may adjust the stage settings according to your workload.

| |

|---|

| if self.runtime.get_num_rdma_ranks() > 1 or low_latency_mode:

# Enable IBGDA for the low latency mode, which refers to "no package forwarding between NVLink and RDMA"

if low_latency_mode:

assert num_qps_per_rank > 0

os.environ[``'NVSHMEM_DISABLE_P2P'``] = '1'

os.environ[``'NVSHMEM_IB_ENABLE_IBGDA'``] = '1'

os.environ[``'NVSHMEM_IBGDA_NIC_HANDLER'``] = 'gpu'

os.environ[``'NVSHMEM_IBGDA_NUM_RC_PER_PE'``] = f``'{num_qps_per_rank}'

# Make sure QP depth is always larger than the number of on-flight WRs, so that we can skip WQ slot check

os.environ[``'NVSHMEM_QP_DEPTH'``] = '1024'

# NOTES: NVSHMEM initialization requires at least 256 MiB

os.environ[``'NVSHMEM_CUMEM_GRANULARITY'``] = f``'{2 ** 29}'

# Disable PCIe relaxed ordering to avoid out-of-order messages

os.environ[``'NVSHMEM_IB_ENABLE_RELAXED_ORDERING'``] = '0'

# NOTES: make sure AR (Adaptive Routing) is turned off while running normal kernels, as we cannot verify AR status in the code

# Synchronize using the root ID

nvshmem_unique_ids = [None,] * self.group_size

if (low_latency_mode and self.rank == 0``) or (not low_latency_mode and self.runtime.get_rdma_rank() == 0``):

root_unique_id = self.runtime.get_local_nvshmem_unique_id()

dist.all_gather_object(nvshmem_unique_ids, root_unique_id, group)

root_unique_id = nvshmem_unique_ids[``0 if low_latency_mode else self.runtime.get_root_rdma_rank(True)] |

buffer.dispatch () Callstack

test_intranode. py:: test_main ()

└─> buffer.dispatch ()

├─> if num_rdma_ranks > 1:

│ └─> internode_dispatch ()

│ └─> internode:: dispatch () [CUDA kernel]

└─> else:

└─> intranode_dispatch ()

└─> intranode:: dispatch () [CUDA kernel]

Buffer:: intranode_dispatch (C++ Layer)

│

├─── Pre-dispatch Processing

│ ├─── Validate inputs

│ ├─── Setup CUDA streams

│ └─── Prepare memory buffers

│

├─── Layout Calculation

│ ├─── Calculate prefix matrices

│ │ ├─── rank_prefix_matrix

│ │ └─── channel_prefix_matrix

│ └─── Setup routing information

│

├─── CUDA Kernel Launch

│ ├─── intranode::dispatch_kernel

│ │ ├─── Token routing

│ │ ├─── Data movement

│ │ └─── Expert handling

│ │

│ └─── Optional Operations

│ ├─── Top-k processing

│ └─── FP 8/BF 16 conversion

│

└─── Post-dispatch Processing

├─── Event creation (if async)

├─── Stream synchronization

└─── Return results

├─── Received tokens

├─── Routing information

└─── Event handle

3.3 Undefined-behavior PTX usage

• For extreme performance, we discover and use an undefined-behavior PTX usage: using read-only PTX [[ld.global.nc](http://ld.global.nc/)]([http://ld.global.nc/](http://ld.global.nc/)).L1::no_allocate.L2::256B to read volatile data. The PTX modifier .nc indicates that a non-coherent cache is used. But the correctness is tested to be guaranteed with .L1::no_allocate on Hopper architectures, and performance will be much better. The reason we guess may be: the non-coherent cache is unified with L 1, and the L 1 modifier is not just a hint but a strong option, so that the correctness can be guaranteed by no dirty data in L 1.
• Initially, because NVCC could not automatically unroll volatile read PTX, we tried using __ldg (i.e., [[ld.nc](http://ld.nc/)]([http://ld.nc/](http://ld.nc/))). Even compared to manually unrolled volatile reads, it was significantly faster (likely due to additional compiler optimizations). However, the results could be incorrect or dirty. After consulting the PTX documentation, we discovered that L 1 and non-coherent cache are unified on Hopper architectures. We speculated that .L1::no_allocate might resolve the issue, leading to this discovery.
• If you find kernels not working on some other platforms, you may add DISABLE_AGGRESSIVE_PTX_INSTRS=1 to setup.py and disable this, or file an issue.

GPU 通信技术

NVLink（200-300 GB/s 带宽）

RDMA（InfiniBand/RoCE，40-50 GB/s 带宽）

2.2 技术基础

• PyTorch 分布式：ProcessGroup、Collective Operations
• CUDA 编程：Kernel 优化、Stream 与 Event 管理
• 高性能网络：InfiniBand 虚拟通道（Virtual Lanes）配置

---

3. 技术要点

3.1 架构设计

• 双层通信优化：
• Normal Kernels（训练/推理预填充）：
• 混合 NVLink + RDMA 转发
• 支持 SM 资源控制（Buffer.set_num_sms()）
• Low-Latency Kernels（推理解码）：
• 纯 RDMA 通信（最低延迟）
• 自适应路由支持（Adaptive Routing）
• 通信优化技术：
• 异步通信：通过 EventOverlap 类实现通信 - 计算重叠
• 内存管理：预分配通信 Buffer（Buffer 类管理 NVLink/RDMA 内存）
• 拓扑感知：自动适应 Intranode（NVLink）与 Internode（RDMA）拓扑

3.2 性能数据

|Intranode（NVLink）|8|153|-|

|Internode（RDMA）|64|45|353|

3.3 核心 API 示例

# 训练场景（Dispatch-Combine流程）

buffer = Buffer(group, nvl_bytes, rdma_bytes)

recv_x, recv_idx, recv_weights, _, handle, event = buffer.dispatch(x, topk_idx, ...)

combined_x, event = buffer.combine(recv_x, handle)

  

# 推理场景（低延迟模式）

recv_hidden, expert_count, handle, event, hook = buffer.low_latency_dispatch(hidden, topk_idx, ...)

combined_hidden, event, hook = buffer.low_latency_combine(recv_hidden, handle)

3.4 关键优化

• PTX 指令级优化：
• 使用 [[ld.global.nc](http://ld.global.nc/)]([http://ld.global.nc/](http://ld.global.nc/)).L1::no_allocate 非一致性缓存指令提升访存效率
• 支持通过 DISABLE_AGGRESSIVE_PTX_INSTRS 关闭激进优化
• 网络配置最佳实践：
• 虚拟通道隔离：NVSHMEM_IB_SL 环境变量控制
• 自适应路由策略：静态路由（低负载）vs 动态路由（高负载）

---

4. 实践指南

4.1 部署要求

• 硬件：
• Hopper 架构 GPU（H 100/H 800）
• 400 Gb/s InfiniBand/CX 7 网卡（RDMA 支持）
• 软件栈：
• CUDA 12.3+
• PyTorch 2.1+
• 定制版 NVSHMEM（需从项目指南安装）

4.2 性能调优

• 自动调参：通过测试脚本（test_intranode.py 等）获取集群最佳配置
• SM 资源分配：根据任务类型调整 SM 数量（默认 24）
• 内存预分配：根据 hidden_size 和 num_experts 预计算 Buffer 大小

---

5. 参考资料

• 项目文档: DeepEP GitHub Wiki
• 相关论文:
• DeepSeek-V 3: [arxiv.org/abs/xxxx.xxxxx](https://arxiv.org/abs/xxxx.xxxxx)
• MoE 架构综述：《Outrageously Large Neural Networks》
• 技术手册:
• NVLink Architecture Whitepaper
• InfiniBand Performance Tuning Guide
• PyTorch Expert Parallelism

---

备注：建议结合实际性能测试数据与用户集群环境调整参数，并参考项目提供的示例代码进行二次开发。

Low latency supa kernel 走读启示：

1. DeepEP FP 8 量化通信，INT 8 量化提升通信效率 gongbo 已经提前布局正在做
2. IB 专用网卡 ibgda，与 NVSHMEM 低延迟，RDMA 是否可以替代
3. Br 10 x 上复现方案，还缺少哪些技术，需要进一步的讨论

Actions（下周三 check 完成进度）：

1. DeepEP_ll 发送数据前 SM 做完量化，然后调用 nvshmemi_ibgda 之后，此时 SM 是空闲的，SM 是否切换出来给计算使用，是否发送完成是用 stream event 来做不再消耗 SM 的计算资源，owner：John
2. IB 网卡 ibgda 与 NVlink，是否可以用现有 ROCE RDMA 替代，性能相差多少，Moying 找相关同事确认
3. INT 8 量化提升通信效率，精度是否满足，需要本地验证 naigang
4. SCCL 接入 INT 8 量化与反量化 kernel 函数，Moying/John

Dep Api

• 基础概念
• MOE（Mixture of Experts）
• Router
• All to All
• Dispatch && Combine
• Intranode && Internode
• NVSHMEM
• DeepEP

基础概念

MOE（Mixture Of Experts）

EP 是 Expert Parallelism 的缩写，也叫专家并行策略。该策略用于在分布式训练中处理混合专家模型，在这种策略中，模型的不同部分（如不同的专家模块）被分配到不同的进程（或 GPU）上进行计算。

MOE 是一种稀疏激活的模型架构，通过将大模型拆分为多个子网络（称为“专家”），每次只激活少数专家（通常 1-2 个），在不显著增加计算量的情况下大幅提升模型容量。

关键组成：

• 专家网络（Experts）：每个专家是一个独立的前馈网络（FFN），通常结构相同但参数独立。
• 门控机制（Gating/Router）：根据输入动态决定激活哪些专家（例如通过 Softmax 选择 Top-K 专家）。
• 稀疏性：虽然总参数量巨大（如万亿级），但每次推理仅激活部分参数（如 10%），降低计算成本。

在 MoE 模型中，Gate 和 Router 是两个经常被交替使用的术语，它们的核心功能是决定如何将输入数据（Token）分配给不同的 Experts（专家模块）进行处理。

• 功能：Gate 是一个神经网络模块，通常是一个线性变换层（如 nn.Linear），其作用是根据输入数据的特征，计算每个 Token 应该被分配到各个 Expert 的概率。Router 是 Gate 的另一种称呼。
• 工作原理：输入数据经过 Gate 后，会得到一个概率分布，表示每个 Token 分配到每个 Expert 的概率。通常会使用 Top-K 选择策略，例如 Top-2，即选择概率最高的两个 Expert。
• 重要性：Gate 的设计和训练效果直接影响 MoE 模型的性能，因为它决定了数据如何被分配到不同的 Experts，进而影响模型的计算效率和负载均衡。Router 的设计需要考虑如何在保证模型性能的同时，实现高效的负载均衡，避免某些 Experts 被过度使用，而其他 Experts 则闲置。

MoE EP Gate/Router

• 通信优化
• Device limited Routing: 将 routing 的 experts 限定在 M 个设备上 (减少通信范围，从而降低通信开销)
• Node limited Routing: 将 routing 的 experts 限定在最高亲和力 (网络拓扑结构 (NUMA/UMA)、网络延迟、带宽) 的 M 个 Node 上
• 负载均衡优化
• Aux-loss-free: tokens 更加均衡 (训练步骤中引入控制变量，使得负载较重的专家偏置会减少，而负载较轻的专家偏置则会增加)
• Sequence Aux-Loss: 解决单个 sequence 内部的极端负载不均衡 (辅助损失函数)

All to All

• 均衡 all to all (上图): 每张卡已确定需要接收的 size，只做一次 all to all 可拿到结果
• 非均衡 all to all: 每张卡不确定需要接收的 size，可先做一次 all to all 拿到接收的 size

Dispatch && Combine

• D****ispatch：
• 将输入数据分发到不同的专家 experts 进行处理
• 每个输入 token 只需要激活 Top-K 个专家，dispatch 需要根据路由结果 (token_idx) 将数据分发到对应的专家
• Combine：
• 将各个专家的输出结果合并回一个完整的输出张量
• 每个 token 的输出是由 Top-K 个专家的输出加权求和得到的，因此需要将这些部分结果重新组合

Intranode && Internode

1. Intranode（节点内）

Intranode 指的是在同一计算节点（Node）内部的操作或通信。具体来说，它涉及同一个物理机器或计算节点内的多个处理单元（如 CPU 核心、GPU 等）之间的交互。

• 通信方式：节点内的通信通常通过共享内存或高速互连（如 NVLink）进行，具有低延迟和高带宽的特点。

1. Internode（节点间）

Internode 指的是不同计算节点之间的操作或通信。具体来说，它涉及多个物理机器或计算节点之间的交互。

• 通信方式：节点间的通信通常通过网络接口（如 InfiniBand、Ethernet）进行，具有较高的延迟和较低的带宽，但可以通过优化协议（如 RDMA）来提高效率。

|数据传输路径|数据通过主机内存传输，CPU 可能参与数据准备和传输|数据直接在 GPU 和 RDMA 网络设备之间传输，绕过主机内存和 CPU|

|延迟|低延迟，但涉及主机内存时可能有额外延迟|更低延迟，直接访问 GPU 内存|

|带宽利用率|高带宽利用率，但可能受限于主机内存带宽|更高带宽利用率，直接访问 GPU 内存|

|CPU 负载|CPU 可能需要参与数据传输的某些阶段|CPU 完全不参与 GPU-GPU 通信，显著减轻了 CPU 负载|

|硬件支持|需要支持 RDMA 的网卡（如 InfiniBand 或 RoCE）|需要支持 RDMA 的网卡和 NVIDIA GPU，支持 GPUDirect 技术|

|编程复杂性|相对简单，使用 MPI 或其他 RDMA 库|稍微复杂，需要使用特定的 GPUDirect API 和 CUDA 编程|

|内存固定|需要固定主机内存（内存 Pinning）|需要固定 GPU 内存（内存 Pinning）|

|内存映射|使用 IOMMU 将主机内存地址映射到 PCIe 地址空间|使用 IOMMU 将 GPU 内存地址映射到 PCIe 地址空间|

|通信协议|支持 InfiniBand、RoCE 等|支持 InfiniBand、RoCE 等|

GPUDirect RDMA 主要用于同一节点内 GPU 之间的直接通信，但 deepEP 涉及多个计算节点。外接网卡（如 InfiniBand 或 RoCE 网卡）用于连接不同节点，实现跨节点的 GPU 通信。

IBGDA（InfiniBand GPUDirect Async）是一种优化的通信技术，用于进一步降低 GPU 之间的通信延迟。在传统的 GPUDirect RDMA 中，GPU 准备好数据后，需要通知 CPU 代理线程，然后 CPU 代理线程填充工作请求（WR）的控制信息，并通过门铃机制向网卡（NIC）发出信号，以启动数据传输。这个过程会带来额外的通信开销。IBGDA 可以优化开销。在 deepEP 的低延迟模式下会用到。

目前我们的 cmodel 仅仅支持 GPU 之间直连以及通过 Switch 桥接，两种通信模式。GPU 直接 trigger RDMA 目前在 cmodel 无法验证功能 (RDMA br 200 是 vendor (GPU 制造商) 提供的，它的 init sequence 以及设计我们无从得知，目前 CModel 仅支持自研模块的开发建模)。

NVSHMEM

• 将多个 GPU 的内存组合成一个分区的全局地址****空间
• 可通过 NVSHMEM API 访问将输入数据分发到不同的专家 experts 进行处理
• 包含一个低开销的内核通信 API，供 GPU 线程使用
• 包括基于流和 CPU 启动的通信 API
• 可与 MPI 和其他 OpenSHMEM 实现互操作

和 NCCL 对比：

特性	NCCL	NVSHMEM
主要用途	针对集合通信（如 AllReduce、Broadcast）优化，专为深度学习分布式训练设计。	基于 PGAS 模型的细粒度内存访问，支持任意 GPU/节点间的直接内存读写。
设计	强调高吞吐、低延迟的集合操作，适合紧密同步的并行任务。	提供全局地址空间抽象，支持灵活的异步通信，适合非规则或动态通信模式。
通信协议	基于 BLink/RDMA 优化集合通信算法（如 Ring AllReduce）。	利用 PGAS (Partitioned Global Address Space) 模型 (提供了逻辑上统一的全局地址空间，允许程序员像在共享内存系统中一样访问数据) 和 CUDA-aware 技术 (MPI 接口可以直接在 GPU 内存之间传输数据)，直接操作远程内存地址。
通信粒度	基于集体操作（如 AllReduce、AllGather），需要所有进程参与同一操作。	支持单边通信（Put/Get/Atomics），允许单个 GPU 直接读写远程内存。
同步机制	隐式同步，操作完成后自动保证数据一致性。	需显式同步（如 nvshmem_fence 或 nvshmem_quiet）确保内存可见性。
编程范式	通过显式调用通信函数（如 ncclAllReduce）。	类似共享内存的地址直接访问（如 nvshmem_put）。

DeepEP

源码：[https://github.com/deepseek-ai/DeepEP](https://github.com/deepseek-ai/DeepEP)

• 高度优化的 All 2 All 通信，适合 MoE 模型的 2 个主要过程：
• Fused_Dispatch：将 Token hidden states 发送给 experts，作为 experts MLP 的输入。
• Fused_Combine：experts 计算 MLP 完成后，从 experts 接收计算过的 Token hidden states.
• 支持不同的通信类型：
• 节点内（intra-node）：可以使用 NVLink + NVSwitch 通信。
• 节点间（inter-node）：可以使用 RDMA 通信。
• 针对不同场景的 Kernel：
• 常规（高吞吐）Kernel（Normal Kernel）：针对 Training 和 Inference Prefill，节点内 NVLink + 节点间 RDMA 通信。
• 低时延 Kernel（Low-Latency Kernel）：针对 Inference Decoding，使用纯 RDMA 通信来最小化时延。
• 原生支持 FP 8，减少数据传输需求，相比 FP 16 通信量减半。
• 灵活的 GPU 资源（SM）控制，支持计算和通信的 Overlap。

具体相关使用见**[https://github.com/deepseek-ai/DeepEP](https://github.com/deepseek-ai/DeepEP)**的 readme.md

200 Native DeepEP 工作展望

DeepEP Requirements

CUDA >=12.3

NVLink for intranode communication

RDMA network for internode communication →

DeepEP 依赖库: NVSHMEM，

DeepEP PTX****操作

utils. cuh 定义了大量的 PTX 操作 (PTX ISA instructions performance)，如 LD/ST 采用了 acquire/relaxed，在 kernel 中大量使用，进一步提高的处理效率

Api

• DeepEP Internode
• get_dispatch_layout
• DeepEP Intrannode
• notify_dispatch
• cached_notify_dispatch
• dispatch
• combine
• cached_notify_combine

API 的接口见：[https://gitlab.birentech.com/software/br_DeepEP2.0/-/blob/develop_br200/csrc/deep_ep.hpp?ref_type=heads](https://gitlab.birentech.com/software/br_DeepEP2.0/-/blob/develop_br200/csrc/deep_ep.hpp?ref_type=heads)

DeepEP Internode

get_dispatch_layout

1. 接口代码参数说明

折叠源码

| |

|---|

| void get_dispatch_layout(``const int64_t* topk_idx,

int``* num_tokens_per_rank, int``* num_tokens_per_rdma_rank,

int``* num_tokens_per_expert, bool* is_token_in_rank,

int num_tokens, int num_topk, int num_ranks, int num_experts,

cudaStream_t stream); |

get_dispatch_layout 是一个用于分布式专家混合模型（MoE）调度的核心接口，主要功能是根据 token 到专家的分配关系（topk_idx），统计并生成全局调度所需的元数据。

为什么需要这个接口？

1. 负载均衡：通过统计 num_tokens_per_rank，确保各 GPU 计算量均衡。
2. 通信优化：is_token_in_rank 标记减少不必要的跨节点通信。
3. 资源分配：RDMA 统计（如 num_tokens_per_rdma_rank）指导物理机器间的数据传输。

核心功能：

1. 统计每个专家处理的 token 数量（num_tokens_per_expert）
2. 统计每个 rank（GPU）需要处理的 token 数量（num_tokens_per_rank）
3. 统计每个 RDMA 节点（物理机器）需要处理的 token 数量（num_tokens_per_rdma_rank）
4. 标记每个 token 是否需要发送到特定 rank（is_token_in_rank）

输入参数：

|topk_idx|int 64_t*|形状为 [num_tokens, num_topk]，表示每个 token 选择的前 num_topk 个专家 ID|输入|

|num_tokens|int|总 token 数（如 4096）|输入|

|num_topk|int|每个 token 选择的专家数（如 8）|输入|

|num_ranks|int|总 GPU 数量（如 8）|输入|

|num_experts|int|(256 // num_ranks) * num_ranks，专家总数（保证能被 GPU 数整除）|输入|

输出参数：

|num_tokens_per_expert|int*|形状为 [num_experts]，每个专家处理的 token 总数|输出|

|num_tokens_per_rank|int*|形状为 [num_ranks]，每个 GPU 需要处理的 token 数|输出|

|num_tokens_per_rdma_rank|int*|形状为 [num_nodes]，每个物理机器（含多个 GPU）需处理的 token 数|输出|

|is_token_in_rank|bool*|形状为 [num_tokens, num_ranks]，标记 token 是否需发送到某 GPU|输出|

2. Sample UT

假设以下配置：

• num_tokens=4，num_topk=2，num_ranks=2，num_nodes (物理机器节点)=1，num_experts=4（每个 rank 管理 2 个专家）
• topk_idx 的输入数据（每个 token 选择 2 个专家）

折叠源码

| |

|---|

| [[``0``, 1``], # Token0选择专家``0``和``1

[``1``, 2``], # Token1选择专家``1``和``2

[``2``, 3``], # Token2选择专家``2``和``3

[``0``, 3``]] # Token3选择专家``0``和``3 |

步骤 1：统计专家负载

• 专家 0：Token 0 和 3 → num_tokens_per_expert[0] = 2
• 专家 1：Token 0 和 1 → num_tokens_per_expert[1] = 2
• 专家 2：Token 1 和 2 → num_tokens_per_expert[2] = 2
• 专家 3：Token 2 和 3 → num_tokens_per_expert[3] = 2

步骤 2：统计 Rank 负载

• Rank 0（管理专家 0-1）
• Token 0（专家 0-1）、Token 1（专家 1）、Token 3（专家 0）需发送 → num_tokens_per_rank[0] = 3
• Rank 1（管理专家 2-3）
• Token 1（专家 2）、Token 2（专家 2-3）、Token 3（专家 3）需发送 → num_tokens_per_rank[1] = 3

步骤 3：生成发送标记

is_token_in_rank 矩阵如下：

折叠源码

| |

|---|

| [[True, False], # Token0 需发送到 Rank0

[True, True], # Token1 需发送到 Rank0和``1

[False, True], # Token2 需发送到 Rank1

[True, True]] # Token3 需发送到 Rank0和``1 |

步骤 4：RDMA 节点统计

1. Token 0

• 选择专家 0 和 1。
• 专家 0 和 1 都属于 Rank 0。
• 因此，Token 0 需要发送到 Rank 0。
• Rank 0 属于 RDMA 节点 0。
• 结论：Token 0 属于 RDMA 节点 0。

1. Token 1

• 选择专家 1 和 2。
• 专家 1 属于 Rank 0，专家 2 属于 Rank 1。
• 因此，Token 1 需要发送到 Rank 0 和 Rank 1。
• Rank 0 和 Rank 1 都属于 RDMA 节点 0。
• 结论：Token 1 属于 RDMA 节点 0。

1. Token 2

• 选择专家 2 和 3。
• 专家 2 属于 Rank 1，专家 3 属于 Rank 1。
• 因此，Token 2 需要发送到 Rank 1。
• Rank 1 属于 RDMA 节点 0。
• 结论：Token 2 属于 RDMA 节点 0。

1. Token 3

• 选择专家 0 和 3。
• 专家 0 属于 Rank 0，专家 3 属于 Rank 1。
• 因此，Token 3 需要发送到 Rank 0 和 Rank 1。
• Rank 0 和 Rank 1 都属于 RDMA 节点 0。
• 结论：Token 3 属于 RDMA 节点 0。

在 CUDA 内核代码中，num_tokens_per_rdma_rank 的计算逻辑是：

• 对于每个 token，检查它选择的专家是否属于某个 RDMA 节点。
• 如果属于，增加该 RDMA 节点的计数。
• 每个 token 只会被统计一次，即使它可能属于多个 ranks。

在这个例子中，虽然所有 token 都属于 RDMA 节点 0，但 num_tokens_per_rdma_rank 的计算逻辑是基于唯一 token 的数量。

• Token 0、Token 1、Token 2 和 Token 3 都属于 RDMA 节点 0。
• 但是，Token 1 和 Token 3 的专家选择跨越了多个 ranks（Rank 0 和 Rank 1），这意味着它们在统计时会被重复计算。

最终 RDMA 节点统计结果是 3，具体是：

1. Token 0
2. Token 1 或 Token 3（其中一个）
3. Token 2

备注：

这里的 RDMA 节点统计涉及到了一个去重的操作，golden 里的 python inplace_unique 函数也是类似，主要完成以下功能：

1. 输入处理

• 输入 x 是一个二维张量（形状为 [batch_size, num_elements]）
• 所有负值（x < 0）被视为无效值（例如 -1 表示填充值）

1. 去重

• 对每行的元素进行去重，保留唯一值
• 按值的出现频率降序排序（高频值在前）
• 无效值（原负值）会被过滤掉

1. 输出

• 输出张量形状与输入相同，但每行只保留前 num_slots 个有效值
• 不足部分用 -1 填充

参数说明

|x|torch. Tensor|输入的二维张量（每行独立处理）|

|num_slots|int|每行最终保留的有效值数量上限|

折叠源码

| |

|---|

| # 掩码无效值

mask = x < 0 # 标记所有负值为无效

x_padded = x.masked_fill(mask, num_slots) # 临时将无效值设为num_slots（避免干扰计数）

# 统计频次

# 对每行的值进行直方图统计（统计``0``~num_slots-``1``的出现次数）

bin_count.scatter_add_(``1``, x_padded, torch.ones_like(x_padded))

# 排序和过滤

# 按频次降序排序，频次为``0``的槽位设为-``1

sorted_bin_idx.masked_fill_(sorted_bin_count == 0``, -``1``)

# 输出格式化

# 只保留前num_slots个有效值，其余填-``1

x[:,:valid_len] = sorted_bin_idx[:,:valid_len] |

Demo

折叠源码

| |

|---|

| x = torch.tensor([

[``1``, 3``, 2``, 3``, -``1``], # 有效值：``1``, 3``, 2

[``5``, 5``, -``1``, -``1``, -``1``] # 有效值：``5

], device=``'cuda'``)

inplace_unique(x, num_slots=``3``)

tensor([

[``3``, 1``, 2``], # 按频次排序（``3``出现``2``次，``1``和``2``各``1``次）

[``5``, -``1``, -``1``] # 只有一个有效值``5

], device=``'cuda'``) |

DeepEP Intrannode

notify_dispatch

deepEP 中的 notify_dispatch 接口是一个用于 MoE（Mixture of Experts）模型的分布式计算调度函数，主要用于协调不同 GPU 或计算节点之间的 token 分发和专家计算任务。以下是对其功能和参数的详细说明：

函数功能

1. Token 分发调度：根据 token 所属的专家（expert）和当前计算节点（rank）的负载情况，将 token 分配到合适的计算设备。
2. 前缀和计算：生成 channel_prefix_matrix 和 rank_prefix_matrix_copy，用于后续通信或内存操作。
3. 异步 CUDA 流操作：通过 cudaStream_t 实现 GPU 上的并行调度。

以下是一个具体示例，假设场景为 8 个 GPU（num_ranks=8）和 64 个专家（num_experts=64），处理 1024 个 token：

|num_tokens_per_rank|const int*|每个 rank 当前待处理的 token 数量|YES|输入|

|moe_recv_counter_mapped|int*|每个 rank 接收到的 token 计数器|YES|输出|

|num_ranks|int|参与计算的 GPU/rank 总数|NO|输入|

|num_tokens_per_expert|const int*|每个专家需要处理的 token 数量|YES|输入|

|moe_recv_expert_counter_mapped|int*|每个专家实际接收的 token 计数器。|YES|输出|

|num_experts|int|专家总数|NO|输入|

|num_tokens|int|总 token 数量|NO|输入|

|is_token_in_rank|const bool*|标记每个 token 是否属于当前 rank|YES|输入|

|channel_prefix_matrix|int*|通信通道的前缀和矩阵|YES|输出|

|rank_prefix_matrix_copy|int*|rank 前缀和矩阵的副本|YES|输出|

|num_memset_int|int|需要初始化的整数数量（通常为 num_ranks * num_experts）|NO|输入|

|expert_alignment|int|专家内存对齐要求（如 128 字节）|NO|输入|

|buffer_ptrs|void**|每个 rank 的通信缓冲区指针数组|YES|输入|

|task_fifo_ptrs|int**|每个 rank 的任务队列指针（用于调度）|YES|输入|

|head|int|任务队列的起始位置|YES|输入|

|rank|int|当前 rank 的 ID（0 到 num_ranks-1）|YES|输入|

|stream|cudaStream_t|CUDA 流，用于异步操作|YES|输入|

|num_sms|int|GPU 的 SM（Streaming Multiprocessor）数量（影响内核并发）|NO|输入|

这里有两个输入参数，是二级指针，可能上述描述的有些让人困惑，下面只针对这两个参数，更加详细的描述：

buffer_ptrs（通信缓冲区指针数组）

基础功能

• 跨 GPU 数据传输：用于在不同 GPU（rank）之间传递 token 数据（如隐藏状态、专家 ID 等）。
• RDMA（远程直接内存访问）支持：通过 nvshmem 库直接读写其他 GPU 的内存，避免 CPU 参与。
• 多级缓冲区管理：区分本地（NVL）和远程（RDMA）通信缓冲区。

数据结构

• 类型：void**（指针数组，每个元素指向一个 GPU 的缓冲区）。
• 长度：num_ranks（每个 rank 一个缓冲区）。
• 内存布局：

折叠源码

| |

|---|

| buffer_ptrs[``0``] → Rank0 的缓冲区（本地）

buffer_ptrs[``1``] → Rank1 的缓冲区（远程）

…

buffer_ptrs[k] → Rankk 的缓冲区 |

Sample 说明

假设有 2 个 GPU（Rank 0 和 Rank 1），每个 token 数据为 hidden_size=1024 的 float 数组：

1. 初始化缓冲区：

折叠源码

| |

|---|

| void``* buffer_ptrs[``2``];

cudaMalloc(&buffer_ptrs[``0``], max_tokens * 1024 * sizeof(``float``)); // Rank0 的缓冲区

cudaMalloc(&buffer_ptrs[``1``], max_tokens * 1024 * sizeof(``float``)); // Rank1 的缓冲区 |

1. 数据通信：

折叠源码

| |

|---|

| float``* local_data = (``float``*)buffer_ptrs[``0``]; // Rank0 的本地数据

nvshmem_float_put_nbi(

(``float``*)buffer_ptrs[``1``], // 目标地址（Rank1 的缓冲区）

&local_data[``0``], // 源地址（Rank0 的数据起始位置）

100 * 1024``, // 数据量（100 tokens × 1024 floats）

1 // 目标 rank（Rank1）

); |

task_fifo_ptrs（任务队列指针数组）

基础功能

• 任务调度：管理每个 GPU 需要处理的专家计算任务（FIFO 队列）。
• 动态负载均衡：根据路由结果动态填充任务队列。
• 低延迟设计：通过 GPU 内核直接操作队列，避免 CPU 同步。

数据结构

• 类型：int**（指针数组，每个元素指向一个 GPU 的任务队列）。
• 队列格式：

折叠源码

| |

|---|

| task_fifo_ptrs[rank][``0``] → 专家ID

task_fifo_ptrs[rank][``1``] → token 起始位置（在缓冲区中的偏移）

task_fifo_ptrs[rank][``2``] → token 数量

task_fifo_ptrs[rank][``3``] → 下一个任务（专家ID）

… |

Sample 说明

假设 Rank 0 需要处理以下任务：

1. 任务 1：专家 1 的 token[0:49]
2. 任务 2：专家 0 的 token[50:99]
3. 任务队列初始化：

折叠源码

| |

|---|

| int``* task_fifo_ptrs[``2``];

cudaMalloc(&task_fifo_ptrs[``0``], 100 * 3 * sizeof(``int``)); // Rank0 的任务队列

cudaMalloc(&task_fifo_ptrs[``1``], 100 * 3 * sizeof(``int``)); // Rank1 的任务队列 |

1. 填充任务队列：

折叠源码

| |

|---|

| int``* queue = task_fifo_ptrs[``0``];

queue[``0``] = 1``; // 专家ID = 1

queue[``1``] = 0``; // token 起始位置 = 0

queue[``2``] = 50``; // token 数量 = 50

queue[``3``] = 0``; // 专家ID = 0

queue[``4``] = 50``; // token 起始位置 = 50

queue[``5``] = 50``; // token 数量 = 50 |

1. 调用 moe_kernel。根据上面的任务队列，按照顺序执行。

cached_notify_dispatch

有了上述 notify_dispatch 的接口的介绍，这里主要对比 cached_notify_dispatch 和 notify_dispatch 的接口的区别：

这两个函数均用于 MoE 模型中的跨 GPU 通信和任务调度，但设计目标和场景不同：

|用途|全功能路由分发，包含动态计算和通信|基于预计算结果的快速分发（缓存优化版）|

|参数复杂度|高（需传递路由结果、专家分配、缓冲区等完整信息）|低（仅需前缀和矩阵和基础指针）|

|性能场景|首次路由或动态路由场景|固定路由模式或高频重复调用的优化场景|

|内部实现|实时计算路由逻辑 + 通信|直接复用预计算的分发方案，跳过冗余计算|

|是否依赖预计算|否|是 (需提前生成 rank_prefix_matrix)，关键路由信息（如 token 到 rank 的映射）已提前计算并存储。|

|buffer_ptrs|void**|每个 rank 的通信缓冲区指针数组|YES|输入|

|task_fifo_ptrs|int**|每个 rank 的任务队列指针（用于调度）|YES|输入|

|head|int|任务队列的起始位置|YES|输入|

|rank|int|当前 rank 的 ID（0 到 num_ranks-1）|YES|输入|

|num_ranks|int|参与计算的 GPU/rank 总数|NO|输入|

|stream|cudaStream_t|CUDA 流，用于异步操作|YES|输入|

Dispatch

该 API 主要用于将 token 分发到不同的 ranks 或专家（experts）进行并行处理。

支持多种通信模式

• Intranode（节点内通信）：
• 利用 NVLink 高速互连技术，在同一节点内的多个 GPU 之间直接传递数据。
• 适用于将 token 分发到同一节点内的不同 ranks。
• Internode（节点间通信）：
• 利用 RDMA（远程直接内存访问）技术，在不同节点的 GPU 之间传递数据。
• 适用于将 token 分发到不同节点的 ranks。

输入参数

|recv_x|void*|接收的 token 数据的缓冲区。
每个 GPU 可能接收不同数量的 token，因此 recv_x 的内容和大小可能因 GPU 而异。需要根据当前 GPU 的 rank 和 num_tokens_per_rank 来确定具体的数据。|YES|输入/输出|

|recv_x_scales|float*|接收的 token 数据的缩放因子（如果适用）。与 recv_x 类似，每个 GPU 接收的 token 数量不同，因此 recv_x_scales 的内容和大小也会不同。|YES|输入/输出|

|recv_src_idx|int*|接收的 token 的源索引。每个 GPU 接收的 token 来自不同的源，因此 recv_src_idx 的内容和大小会因 GPU 而异。|YES|输入/输出|

|recv_topk_idx|int 64_t*|接收的 top-k 专家索引。

每个 GPU 接收的 token 的 top-k 索引不同，因此 recv_topk_idx 的内容和大小会因 GPU 而异。|YES|输入/输出|

|recv_topk_weights|float*|接收的 top-k 权重。

每个 GPU 接收的 token 的 top-k 权重不同，因此 recv_topk_weights 的内容和大小会因 GPU 而异。|YES|输入/输出|

|recv_channel_offset|int*|接收通道的偏移量。

每个 GPU 的通道偏移量可能不同，因此 recv_channel_offset 的内容和大小会因 GPU 而异。|YES|输入|

|send_head|int*|发送队列的头部索引。

每个 GPU 发送的 token 数量和内容可能不同，因此 send_head 的内容和大小会因 GPU 而异。|YES|输入/输出|

|x|const void*|输入的 token 数据。

每个 GPU 也是不一样的。|YES|输入|

|x_scales|const float*|输入的 token 数据的缩放因子（如果适用）。

通常所有 GPU 共享相同的缩放因子，因此 x_scales 的内容在所有 GPU 上是相同的。|YES|输入|

|topk_idx|const int 64_t*|每个 token 的 top-k 专家索引。

通常所有 GPU 共享相同的 top-k 索引，因此 topk_idx 的内容在所有 GPU 上是相同的。|YES|输入|

|topk_weights|const float*|每个 token 的 top-k 权重。

通常所有 GPU 共享相同的 top-k 权重，因此 topk_weights 的内容在所有 GPU 上是相同的。|YES|输入|

|is_token_in_rank|const bool*|每个 token 是否属于当前 rank。

每个 GPU 的 is_token_in_rank 矩阵内容不同，因为每个 GPU 只处理属于它的 token。需要根据当前 GPU 的 rank 来提取相关数据。|YES|输入|

|channel_prefix_matrix|const int*|每个通道的前缀矩阵|YES|输入|

|num_tokens|int|输入 token 的总数|NO|输入|

|hidden_int 4|int|每个 token 的隐藏维度（以 int 4 为单位）|NO|输入|

|num_topk|int|每个 token 的 top-k 数量|NO|输入|

|num_experts|int|专家数量|NO|输入|

|num_scales|int|缩放因子的数量|NO|输入|

|buffer_ptrs|void**|缓冲区指针数组|YES|输入|

|rank|int|当前 rank 的索引|YES|输入|

|num_ranks|int|总 rank 数量|NO|输入|

|stream|cudaStream_t|CUDA 流|YES|输入|

|num_sms|int|CUDA 流处理器数量|NO|输入|

|num_max_send_tokens|int|最大发送 token 数量|NO|输入|

|num_recv_buffer_tokens|int|接收缓冲区的 token 数量|NO|输入|

输出参数

在这个函数中，没有直接的输出参数，因为操作是在设备上进行的，结果存储在设备内存中。不过，函数通常会修改以下参数：

• recv_x：接收的 token 数据。
• recv_x_scales：接收的 token 数据的缩放因子。
• recv_src_idx：接收的 token 的源索引。
• recv_topk_idx：接收的 top-k 专家索引。
• recv_topk_weights：接收的 top-k 权重。
• send_head：发送队列的头部索引。

Combine

combine 函数的作用是将来自不同 rank 的 token 进行合并（归约）。它支持两种场景：

1. Intranode：同一节点内的 GPU 通过 NVLink 通信。
2. Internode：不同节点的 GPU 通过 RDMA 通信。

输入参数

|type|cudaDataType_t|数据类型，x.scalar_type ()|NO|输入|

|recv_x|void *|合并后的 token 数据。每个 GPU 接收的 token 数据可能不同，因此 recv_x 的内容和大小会因 GPU 而异。需要根据当前 GPU 的 rank 和 num_recv_tokens 来确定具体的数据。|YES|输出|

|recv_topk_weights|float *|合并后的 top-k 权重。每个 GPU 接收的 token 的 top-k 权重可能不同，因此 recv_topk_weights 的内容和大小会因 GPU 而异。|YES|输出|

|x|void *|输入的 token 数据。通常所有 GPU 共享相同的输入数据，因此 x 的内容在所有 GPU 上是相同的。|NO|输入|

|topk_weights|float *|输入的 top-k 权重。通常所有 GPU 共享相同的 top-k 权重，因此 topk_weights 的内容在所有 GPU 上是相同的。|NO|输入|

|src_idx|const int *|源索引，指示每个 token 的来源 rank。每个 GPU 的源索引可能不同，因此 src_idx 的内容和大小会因 GPU 而异。|YES|输入|

|rank_prefix_matrix|const int *|每个 rank 的前缀矩阵。每个 GPU 的排名前缀矩阵可能不同，因此需要根据当前 GPU 的 rank 和其他配置信息来提取相关数据。|YES|输入|

|channel_prefix_matrix|const int *|每个通道的前缀矩阵。每个 GPU 的通道前缀矩阵可能不同，因此需要根据当前 GPU 的 rank 和其他配置信息来提取相关数据。|YES|输入|

|send_head|int *|发送队列的头部索引。每个 GPU 发送的 token 数量和内容可能不同，因此 send_head 的内容和大小会因 GPU 而异。|YES|输入|

|num_tokens|int|输入 token 的总数|NO|输入|

|num_recv_tokens|int|接收 token 的总数|YES|输入|

|hidden|int|每个 token 的隐藏维度|NO|输入|

|num_topk|int|每个 token 的 top-k 权重数量|NO|输入|

|buffer_ptrs|void **|缓冲区指针数组|YES|输入|

|rank|int|当前 rank 的索引|YES|输入|

|num_ranks|int|总 rank 数量|NO|输入|

|stream|cudaStream_t|CUDA 流|YES|输入|

|num_sms|int|CUDA 流处理器数量|NO|输入|

|num_max_send_tokens|int|最大发送 token 数量|NO|输入|

|num_recv_buffer_tokens|int|接收缓冲区的 token 数量|NO|输入|

cached_notify_combine

cached_notify_combine 和 combine 都用于管理分布式计算中的 token 合并操作，但在功能、参数及用途上有明显区别。

参数说明如下表：

|buffer_ptrs|void **|指向缓冲区指针数组，用于存储不同通道的数据缓冲区。|YES|输入/输出|

|send_head|int *|指向发送队列头部索引数组，用于跟踪每个通道的发送进度。|YES|输入/输出|

|num_channels|int|通道数量，表示系统中用于通信的逻辑或物理通道数。|NO|输入|

|num_recv_tokens|int|预期接收的 token 总数，用于分配和验证接收缓冲区大小。|YES|输入|

|num_memset_int|int|需要初始化的整数数量，通常用于确定缓冲区初始化的大小。|NO|输入|

|task_fifo_ptrs|int **|指向任务队列指针数组，用于管理待处理任务的队列。|YES|输入/输出|

|head|int|任务队列的头部索引，用于标识当前处理位置。|YES|输入|

|rank|int|当前进程的 rank 索引，用于标识在分布式系统中的位置。|YES|输入|

|num_ranks|int|分布式系统中的总 rank 数量，用于确定通信范围。|NO|输入|

|stream|cudaStream_t|CUDA 流对象，用于管理 GPU 操作的执行顺序。|YES|输入|

输入与输出

• 输入参数：buffer_ptrs, send_head, num_channels, num_recv_tokens, num_memset_int, task_fifo_ptrs, head, rank, num_ranks, stream。
• 输出：该函数主要更新 send_head 和 task_fifo_ptrs，无直接输出。

功能差异对比

|功能|处理内部状态更新，如缓冲区管理、队列更新。|执行实际数据合并，合并不同 rank 的 token 数据。|

|操作|更新发送和任务队列，管理缓冲区状态。|对输入数据进行归约操作（如加权求和）。|

|数据处理|处理通信逻辑，更新状态信息。|处理 token 和权重数据，生成合并结果。|

|侧重点|通信状态管理。|数据合并计算。|

|使用场景|在数据合并前准备通信状态，在合并后更新通信状态。|在状态就绪时执行合并操作。|

总结

• combine：专注于 token 数据及权重的实际合并，输入为原始数据和相关参数，输出为合并后的结果。
• cached_notify_combine：侧重于通信状态的维护，如缓冲区和任务队列管理，以支持合并操作的顺利进行。