高效且便携的专家混合通信

我们推出了一种用于专家混合（MoE）通信的高性能、便携式开源库，与标准的 All-to-All 通信原语相比，其性能提升了 10 倍。

我们的实现具有多项关键技术创新，可提供卓越的 MoE 通信效率：

GPU 发起的通信（IBGDA）：支持直接的 GPU 到网卡（NIC）通信，通过绕过 CPU 的参与显著降低了延迟。

通信与计算重叠：具有独立发送和接收阶段的分离内核架构，使计算能够与网络传输并发进行。

最快的单节点性能：在单节点配置下，延迟比之前最快的实现降低了 2.5 倍。

高效且便携的多节点性能：我们的实现速度比标准的 all-to-all 通信快 10 倍。尽管比高度专业化的实现慢约 2 倍，但我们的方法在不同 NVSHMEM 版本和网络环境（NVLink、CX-7 和 EFA）之间提供了更好的可移植性。

该库完全开源，可在 https://github.com/ppl-ai/pplx-kernels 获取。

在本文中，我们将探讨大规模 MoE 模型中专家并行的挑战，并描述我们在分布式 GPU 环境中实现高效 token 分发和组合的方法。

诸如 DeepSeek R1 和 Mixtral 8x7B 等专家混合（MoE）模型通过限制每个 token 激活的权重数量来改进稠密模型。MoE 模型没有通过增加层数或增加每个解码器层中多层感知机（MLP）的线性层大小来提高参数量，而是用多个专家和一个路由（Router）替换了传统的 MLP。例如，DeepSeek R1 的 6710 亿参数中，在推理过程中只有 370 亿参数与给定的 token 相乘。与稠密模型相比，这极大地减少了解码一个 token 所需的计算量并改善了延迟。

与稠密模型相比，MoE 模型在推理方面带来了一些额外的挑战。虽然专家本身是小型 MLP 层，但每个解码器层都包含一个路由，用于决定 token 分发给哪些专家，并且每个 token 被分发给多个专家。路由通常是一个产生概率分布的小型线性层。通常会选取 top-K 分数的专家，并将最终的激活值计算为加权平均值，即将专家输出乘以从路由的概率分布中得出的权重相加。

为了最小化分布式 MoE 系统中的延迟，可以跨多个设备利用并行性，但这引入了通信挑战。诸如 Mixtral 8x7B 或 Llama-70B 之类的模型可以容纳在单个节点的 8 个设备内，从而受益于快速且低延迟的 NVLink 互连（高达 900Gbps）。然而，更大的模型需要使用 InfiniBand（峰值 400Gbps）跨多个节点对专家进行分片，这引入了额外的延迟挑战。

在本文中，我们探讨了专家并行的难题，描述了我们的 GPU 内核实现，该实现能够有效地将 token 分发到持有分配专家的设备，并通过计算专家输出的加权和将其收集回来。虽然存在针对特定网络硬件进行过激进优化的实现，但我们提出了一种高度可移植的实现，它仅依赖最少量的 NVSHMEM 原语，但提供的性能却比标准 all-to-all 通信快 10 倍。我们的实现实现了单节点配置下的最先进性能，同时在各种网络环境中保持了出色的可移植性。

为了实现高效推理，模型的权重必须保留在设备内存中，同时还需要为前向传播所需的激活值、KV 缓存和其他缓冲区留出足够的空间。最强大的模型超出了即使是最强大的 GPU 的容量，因此推理必须分散在多个设备上，这些设备可以共同存储模型的权重。根据权重的分片方案，必须使用不同的通信和计算方案来卸载计算并同步设备。

专家并行（EP），如图一所示，仅对专家计算进行并行化。不同的专家被分配给持有其权重的不同设备。路由后，token 被发送到相应的设备，结果随后被收集和累积。路由的复杂性取决于模型其他部分的并行度：复制其他层可能会完全消除路由的需求，因为每个 rank 都可以从本地复制的路由表中选择 token。然而，如果只有一个 rank 运行路由，则需要进行广播以将 token、索引和权重分发到各自的专家。最后，all-gather 或 all-to-all 广播将输出的 token 与执行模型后续部分的任何 rank 进行同步。这种实现相对简单，因为 torch 已经公开了所需的原语，尽管通过将累积与从专家收集 token 的 gather 操作相融合可以获得一些益处。

专家独有并行（Expert-only parallelism）并不能理想地扩展，因为当模型运行非专家层（如注意力、归一化和采样）时，集群中的节点可能会处于空闲状态。然而，这些层（主要是注意力）的计算也可以受益于张量并行（TP）。大多数模型依赖于多头注意力，这意味着注意力头及其相应的 Q、K 和 V 投影也可以跨设备进行分片，在各层之间复制或收集切片。如果注意力分散在所有设备上，all-gather 可以同步激活值，从而允许复制路由，这需要类似于专家并行情况的同步原语以实现高效的实现。然而，在此级别上的并行性存在限制，因为将注意力头的数量减少到特定阈值以下会产生边际效益递减。

为了最好地利用所有设备并支持高达 128 或 256 个 GPU 的极高专家并行度，需要数据并行（DP）。在这种方案下，模型所分割的设备被分组以并发处理请求，在本地组中跨设备计算注意力和维护 KV 缓存。这些并行组的多个实例在专家评估上协作，每个实例承载专家权重的不同子集。基于注意力头的数量，一个组通常可以跨越整个节点的大小，因为节点内通信更快。例如，在上图中，一个 DP rank 独立地服务两个请求 A 和 B，处理这两个请求的注意力、归一化和其他任何簿记工作。另一个 DP rank 处理不同的请求 C。然而，第一个节点托管了一半的专家，而另一个节点托管了另一半，因此路由后来自 A 和 B 的 token 可能会被发送到第二个节点，反之亦然。这导致了一个稀疏通信问题：每个设备发送到任何其他目标 rank 的 token 数量可能不同。来自 torch 的现有原语（主要是 all_to_all）并不特别适用，因为它们可能需要某种形式的填充或 GPU 到 CPU 的同步以及元数据广播。为了有效地实现通信，需要自定义内核来处理路由信息并从设备发起通信，而无需 CPU 干预。路由后，分发内核必须将 token 发送到它们被路由到的 rank，而在合并侧，必须收集属于当前 DP 组中请求的激活值。此外，必须在可能拥有多个设备的 DP 组内平衡工作，以最小化广播并正确地在所有 rank 上复制 token 以供后续层使用。

NVSHMEM 是一种特定于 NVIDIA 的 OpenSHMEM 实现，提供便携式的设备间通信设施，抽象化了底层硬件的复杂性。API 可以表达设备到设备的读写，这些读写映射到各个传输层的原语。NVLink 和 RDMA 都受到支持，从而授予了一定程度的可移植性，并允许 NVSHMEM 内核在节点内或跨多个节点的集群中运行。在我们的内核中，我们使用远程内存写入和信号操作来传输数据并同步参与的设备。

NVSHMEM 操作围绕对称内存的概念构建：它们在已在参与设备间通信的所有设备上分配的缓冲区上运行。本地设备可以使用在其自身本地缓冲区内得出的地址以及目标 rank 的索引来指定操作的目的地。下图说明了这一概念：两个 GPU 分配相同大小的对称缓冲区，保留指向它们的 src 指针。第一个 rank 想要发送 3 个整数到第二个 rank，将它们放置在缓冲区的开头：nvshmem_int_put_nbi 从本地缓冲区得出起始地址，指定目标设备。第二个 rank 从其自己的缓冲区得出偏移量，发送一个元素到第一个设备，偏移量为 1。虽然目标地址必须始终是使用 nvshmem_alloc 分配的对称缓冲区，但源缓冲区可以是设备内存的任意区域，只要它们已在 NVSHMEM 中预注册。

虽然 NVSHMEM 提供了广泛的原语，但我们的内核仅依赖 3 个函数，在此基础上构建了所有同步和隔离（fencing）。

nvshmemx_putmem_signal_nbi_warp：将数据块从一个设备传输到另一个设备，同时在远程设备上设置一个标志。该操作要么设置（NVSHMEM_SIGNAL_SET）要么递增（NVSHMEM_SIGNAL_ADD）一个 64 位位置。标志在整个内存块传输完成后更新。如果远程设备观察到标志的变化，它就可以安全地访问其自身内存中的缓冲区。此函数对于合并数据传输和同步非常有用。

nvshmemx_signal_op 在单个内存位置（通常是一个 64 位标志）上运行，原子地设置或递增它。它在发送元数据和同步设备方面非常有用。

nvshmem_uint64_wait_until 在信号的接收端使用，用于轮询标志直到远程更新它。

通过 nvshmem_quiet 或 nvshmem_fence 确保操作之间的任何排序形式都是昂贵的，因此我们的内核实现避免了使用任何其他 NVSHMEM 原语，并最大限度地减少了除单个发送-接收对之外的任何依赖项。发送数据时，我们始终使用非阻塞版本的函数，无需等待数据从本地 rank 发送出去。跨内核组的其他隐式同步确保了当网卡需要访问这些缓冲区时，它们不会被破坏性地覆盖。

NVSHMEM 使用 GPUDirect RDMA，因此网卡无需 CPU 参与即可直接读写 GPU 内存进行数据交换。此外，在 ConnectX 网卡上，NVSHMEM 支持 GPU 发起的通信（也称为 Infiniband GPUDirect Async 或 IBGDA），这允许 GPU 直接向网卡发布 RDMA 操作并直接轮询网卡的完成队列。在不支持 GPU 发起通信的平台上，NVSHMEM 会在宿主 CPU 上生成一个“代理（proxy）”线程，代表 GPU 发起通信和轮询完成情况。无论是否支持 GPU 发起通信，NVSHMEM 程序都是便携式的。然而，GPU 发起的通信显著减少了延迟，因为它完全绕过了到 CPU 的迂回路径。

我们通过一对分发和合并内核实现 MoE 通信。分发内核负责读取每个 rank 上的 token 和路由信息，并将它们分发给相应的专家。合并内核收集专家产生的激活值并将它们发送回其源 rank，同时根据路由计算的权重计算所选专家的加权平均值。内核进一步拆分为发送和接收组件，以便允许数据传输与计算重叠。发送内核都是非阻塞和非同步的：它们只是将所有写入分发到远程。在另一端，合并内核仅从内存读取，等待所有必需的数据传输完成。在向网卡分发工作后，当数据异步传输时，GPU 可以在本地进行其他有用的工作，例如应用共享专家或计算注意力。

每对内核都有自己的对称内存缓冲区，每个 rank 分配私有存储以接收来自所有其他 rank 的数据。可以从 DP 组分发的最大 token 数量在各 rank 之间是固定的，这也设定了每个 rank 可以从每个 DP rank 接收每个本地专家数据的上限。这允许发送方 rank 在目标 rank 上得出唯一的地址进行写入，而无需在它们之间进行任何同步。数据接收后，内核在内存中对 token 进行洗牌（shuffle），以按接收内核所需的最佳格式进行排列。虽然缓冲区具有相当大的维度，但它们会在模型的所有连续层中重用。

唯一的全局同步形式（如上图所示）是在合并接收（combine-receive）内核中实现的。一旦数据从发送端分发，一个 rank 在接收到它需要的所有其他 rank 的 token 时就会返回并继续执行，并在处理完后立即进入合并发送（combine-send）内核。合并接收中的屏障（barrier）确保没有 rank 可以超前运行并启动分发发送（dispatch-send），而在任何其他 rank 仍在等待接收数据时，可能会导致破坏性的重叠。为简单起见，同步是在合并内核中完成的。作为一种优化，相同的屏障可以在分发接收（dispatch-receive）和合并发送（combine-send）内核之间实现，通过允许网卡有更多时间在 GPU 运行其他内核时结算传输，从而进一步降低延迟。

两个内核都跨设备的可用 SM 进行拆分：虽然分发发送和合并接收内核必须跨每个 rank 的最大 token 计数（max_m）进行并行化，但分发接收和合并发送内核必须处理来自每个 rank 和本地专家的相应数量的 token，最大值为 max_m * (num_experts // EP) * (EP // TP)。

分发内核负责将 token 和块缩放因子分发给托管 token 路由到的专家的对等点。发送方依赖 warp 特化（warp specialization）来并行化两个任务：聚合路由信息以确定每个 rank 发送多少 token，以及将 token 打包成消息以分发到远程。接收方首先等待接收所有 token 计数，然后它将有效载荷从共享内存读取并解包到将传递给下游 Grouped GEMM 内核的张量中。它还将信息存储在与合并内核共享的缓冲区中，以指示每个 token 应发送回哪里。由于接收内核以非确定性顺序在连续缓冲区中洗牌 token，因此需要此信息。

在发送端，每个块上的一个 32 线程的 warp 负责读取路由信息并计算每个目标专家预期从该 rank 接收的 token 数量，使用 nvshmem_signal_op 发送计数加 1。计数加 1，因为远程端从零到非零的转换标志着从源 rank 接收到计数。同时，剩余的 warp 合作将 token 并行复制到所有块中的对称内存中，将激活值、缩放因子及其在本地 rank 上的索引打包成一个连续块。合并发送方需要该索引来确定将要写入 token 的地址。接下来，在确保所有数据都已通过 warp 组之间的屏障复制后，warp 再次独立运行，每个 warp 并行地将相同的缓冲区发送到不同的专家。Token 使用 nvshmemx_putmem_signal_nbi_warp 发送，这也原子地增加了远程设备上来自本地 rank 的已发送 token 的计数。在 DP 组内，由于每个 rank 都拥有待发送 token 的副本，分发被均匀平衡，每个设备发出一部分 token。

在接收端，所有内核首先通过等待发送过来的总 token 计数来与发送方 rank 同步。之后，它们都等待原子递增的已发送 token 计数结算到总计数，这表明来自源 rank 的所有有效载荷也已发送，这得益于 putmem 调用的语义。内核使用 nvshmem_uint64_wait_until 对计数进行轮询，跨所有块和线程并行化该操作。随后，跨块屏障确保除非所有数据都已正确接收，否则没有块会读取缓冲区。分散在块中并通过原子递增计数器同步，token 从对称内存复制到常规连续缓冲区，这些缓冲区成为实现专家的下游内核的输入。源 rank、专家索引和 token 索引被单独存储，精确地指出合并内核必须发送激活值的位置。即使来自 DP 组内的 token 是从不同的设备发送的，它们也会被分组在一起以传递给相应的专家。

合并内核再次分为发送和接收两半：发送方将未量化的 16 位激活值复制到远程设备，接收内核等待数据发送完成，并在本地计算所有专家贡献的加权平均值。此外，它们还充当屏障，以同步分发-合并序列：每个 rank 在进入发送内核时在每个对等点上设置一个标志，如果接收内核没有观察到标志被设置，则不允许返回。同步的延迟很小，因为它与实际的通信和计算重叠。

在发送端，内核并行遍历分配给所有本地专家的 token 列表，将它们写入目标 rank 上的缓冲区。目标 rank、专家索引和 token 索引是从分散（scater）内核在接收 token 后填充的每个 token 缓冲区中读取的。每个发送方每个专家都有自己的私有内存区域要写入，如上图所示，避免了同步的需要。与分发类似，合并原子地递增目标 rank 上的每个 token 计数器以指示数据的接收：当计数器匹配 token 被分发的专家数量时，就可以访问 token 内容。

在接收内核中，token 列表跨多个块并行遍历，通过轮询信号操作设置的标志来等待其内容到达。到达后，从私有缓冲区读取有效载荷，路由表指示从哪个缓冲区读取以及为每个专家分配什么权重。结果随后写入外部分配的张量，内核在所有设备通过屏障后完成执行。跨 DP 组，所有 rank 都会接收每个专家激活值的副本，以计算它们自己的副本。

我们在由 16 个节点组成的 128 个 GPU 集群上评估我们的内核，这些节点通过 InfiniBand 和 CX-7 网卡连接。我们尝试了 GPUDirect Async (IBGDA) 和带有 CPU 代理的可靠连接 (RC)。我们将它们与高度优化的 DeepSeek 实现，以及 PyTorch（通过 NCCL）或 NVSHMEM 提供的稠密原语进行了比较。

虽然在稠密 NVSHMEM all-to-all 操作上，IBRC 和 IBGDA 的性能相似（6378 µs vs 6180 µs），但 IBGDA 在稀疏内核上明显更快。将分发和合并相加，IBGDA 使用 902 µs，而 IBRC 需要 3223 µs - 延迟提高了 3.6 倍。虽然 all-to-all 是带宽受限的，但稀疏内核广播的数据量少了几个数量级，受限于延迟。通过直接从 GPU 触发网络传输，无需 CPU 代理来协调 GPU 和网卡，端到端延迟显著降低。

虽然我们的便携式内核比高度优化的 DeepSeek 内核慢约 2 倍，但与稠密内核相比，它们将延迟降低了 10 倍（902 µs vs 9944 µs）。

此外，与库原语不同，拆分为发送方-接收方组件也允许隐藏部分延迟。

在单节点 (EP8) 上，NVSHMEM 利用 NVLINK 进行传输，提供比节点间网络更低的延迟和更高的吞吐量。我们的便携式内核在单节点上比 DeepEP 快约 2.5 倍（186 µs vs 481 µs）。

这里描述的内核性能优于 ML 框架的内置原语，并提供不错的性能，而无需为特定的设备间传输（例如 NVLink 或 InfiniBand）过度专业化。除了已经提到的机会外，通过将通信原语替换为更专业的版本，还可以获得进一步的性能提升。例如，在 NVLink 上，对称内存的使用可以用跨设备共享的缓冲区替换，从而消除一些复制并允许更细粒度的同步方案（针对单个 token 而不是 token 批次）。在 InfiniBand 上，实现可以直接访问底层的队列对，无需传输预期和已发送的 token 计数，并在 token 到达时将其从队列刷新到输出缓冲区。然而，这种实现是以可移植性为代价的，而我们的实现允许在提交开发工作以针对特定硬件进行优化之前对具有特定配置的集群进行评估。

我们推出了一种用于 MoE 通信的高性能、便携式库，与标准的 all-to-all 通信相比，它实现了 10 倍的性能提升，同时保持了跨不同硬件配置的兼容性。在具有 NVLink 的单节点部署上，我们的解决方案展示了比之前实现降低 2.5 倍的延迟。

我们的方法通过包括 GPU 发起的通信支持、实现计算-通信重叠的分离内核架构以及使用最少 NVSHMEM 原语的高效 token 分发等关键创新，平衡了性能与可移植性。虽然在多节点设置上比高度专业化的实现慢约 2 倍，但我们的库在各种网络环境（NVLink、CX-7 和 EFA）中提供了卓越的灵活性。

随着 MoE 模型继续扩展，像我们这样的高效通信策略对于实际部署将变得越来越重要。我们完全开源的实现可在 https://github.com/ppl-ai/pplx-kernels 获取。

DeepEP

DeepSeek-V3 技术报告

NVIDIA NVSHMEM

BibTeX

@online{perplexity2025EfficientAndPortableMixtureOfExpertsCommunication,
  author = {{Perplexity} and Nandor Licker and Lequn Chen and Phi Hoang},
  title = {{Efficient and portable mixture-of-experts communication}},
  date = {2025-04-02},
  url = {https://perplexity.coolkeypoint.com/hub/blog/efficient-and-portable-mixture-of-experts-communication}
}

简介

分片与并行

NVSHMEM

便携式内核

分发

合并

基准测试

进一步改进

结论

参考

引用