~/FlashDMoE：分布式 MoE 执行范式的变革

2025-07-22 · mlsys · 3 min read

Originally published on Zhihu 知乎 · 2025-07-22

引言

继续尝试用 沈向洋、华刚：读科研论文的三个层次、四个阶段与十个问题 - 知乎 的十个问题读论文。

近年来，混合专家模型（Mixture-of-Experts, MoE）已成为扩展大型语言模型（LLM）能力的主流架构，但其分布式实现（DMoE）长期受系统效率低下困扰。

1. 这篇文章究竟讲了什么问题？比方说你设计一个算法，它的 input 和 output 是什么？

首先需要定义 DMoE 算子

1.1 DMoE 算子定义

DMoE 层是整个分析的基本单元，其工作流程可以被清晰地定义：

输入:

分布在多 GPU 上的 token 嵌入向量、分片存储的专家网络集合

处理流程:

• 门控：每个 GPU 为本地 token 选择少数专家；
• 分发：通过 All-to-All 通信将 token 发送到专家所在 GPU
• 计算：专家对接收的 token 执行计算
• 合并：将结果通过 All-to-All 传回原 GPU 并加权合并。

输出：

最终的 token 表征张量

1.2 问题：系统性效率低下

论文明确指出，当前所有主流的 DMoE 实现，尽管享受了 MoE 架构带来的理论计算稀疏性，却普遍遭受着严重的性能问题。具体表现为：

• GPU 利用率低下（SM 单元大量空闲）；
• 延迟开销显著（端到端时间被非计算开销主导）；
• 无法高效重叠计算与通信，难以利用任务局部性。

1.3 原因分析

作者们精准地找出上述缺陷，并将其归因于现有 DMoE 系统在架构设计上的几个根本性缺陷：

• CPU 管理调度：由 CPU 依次启动门控、通信、计算等阶段的独立 GPU 核函数；
• 频繁核函数启动：每次启动存在延迟，各阶段间产生空闲间隙；
• 主机发起通信：All-to-All 通信由 CPU 发起，引入同步和延迟。

通信与调度开销抵消了 MoE 计算稀疏性的优势，瓶颈从 “计算” 转向 “执行编排与数据移动”

2. 这个问题的性质是什么？...为什么它“仍然重要”

(对应问题二)

优化分布式计算本身是一个经典的高性能计算（HPC）问题，并非全新。但 DMoE 的优化问题在当前人工智能时代背景下，其重要性被前所未有地放大了。

MoE 已成为构建顶尖大模型（如 DBRX、DeepSeek-v3）的核心技术，其效率直接影响前沿大模型领域的发展。提升训练/推理效率可以节约大量成本和电力，并且支持更大参数量的模型架构探索。

3. 论文的核心科学假设

(对应问题三)

FlashDMoE 的核心科学假设是：当前 DMoE 效率低下是由 “多核函数、CPU 编排” 执行模型导致的，通过将整个 DMoE 流程（门控、计算、通信）融合到单个持久化 GPU 核函数中，可系统性消除这些问题。

GPU 线程块自主协同，无需外部指令即可动态响应数据与资源变化

4. 相关研究与技术背景

(对应问题四：有哪些与这篇文章相关的研究？)

论文将其方法与一系列代表了当前最先进水平的系统进行了比较，这些系统构成了它试图超越的现状：

• DeepSpeed-MoE / Tutel
• Megatron-LM
• FasterMoE / Comet
• MegaBlocks

这几个工作都还在用多个核函数的方式，最后一个重构稀疏计算为块稠密计算，抽象层次高于 FlashDMoE

要深入理解 FlashDMoE，必须了解其所利用的底层技术：

• 核函数融合 (Kernel Fusion) vs. 持久化核函数 (Persistent Kernels):
• 核函数融合 是一种标准的 GPU 优化技术，它将多个连续的、独立的核函数合并成一个更大的核函数，目的是减少因数据往返于全局内存而产生的开销（数据可以保留在寄存器或共享内存中），并消除多次核函数启动的延迟
• 持久化核函数 是一种设计模式，即启动一个核函数后，它会进入一个循环，持续运行以处理不断到来的工作项。这种模式非常适合对延迟敏感的流式工作负载，因为它将一次性的启动开销分摊到了整个生命周期中。

FlashDMoE 将 DMoE 层的所有逻辑操作（门控、计算、通信）融合进一个单一的持久化核函数中，该核函数自主管理整个操作的生命周期。这个操作和 Mirage 的后续工作 Mirage Persistent Kernel 很像。

• 通信原语：集合通信 (Collectives) vs. 单边通信 (One-Sided / PGAS):
• 集合通信 (例如，通过 NCCL 实现的 AlltoAll): 这是所有基准系统采用的方式。这类操作是批量同步的，即所有参与的进程必须都到达通信调用点后，数据交换才能开始。这种方式编程模型简单，但容易受到“掉队者效应”（straggler effect）的影响，并且通常要求收发的数据大小对称，从而导致了效率低下的填充（padding）问题。
• 单边通信 / PGAS (例如，通过 NVSHMEM 实现的 (R)DMA): 这是 FlashDMoE 采用的方式。在分区全局地址空间（Partitioned Global Address Space, PGAS）模型中，一个进程可以直接对另一个远程进程的内存进行读或写操作，而无需目标进程的显式协调。

5. FlashDMoE 的架构创新

(对应问题五：这篇文章提出的问题解决方案中，核心贡献是什么？)

多个创新组件的集成，形成统一执行模型

单一持久化核函数

将整个 DMoE 算子封装在一个 CUDA 核函数中，仅启动一次并持续驻留，彻底消除多次核函数启动与 CPU-GPU 同步开销。

核函数内基于演员模型的并发模型

GPU 线程块 / 束被赋予三种角色：

• 处理器：执行 GEMM 等计算任务，发起数据传输；
• 调度器：管理任务队列，动态分配任务给空闲处理器；
• 订阅者：监听入站数据包，解码元数据并生成新任务。

异步、设备发起的通信

基于 NVSHMEM 构建 PGAS 空间，支持 GPU 直接通过 RDMA 写入远程内存，通信异步且由设备发起，源 GPU 无需等待目标响应即可继续计算，实现细粒度通信与计算重叠。

载荷高效的数据处理

在基于 AlltoAll 的系统中，一个主要的低效来源是令牌填充（token padding）。由于集合通信的对称性约束，所有 GPU 必须发送和接收大小相同的缓冲区。如果某个专家接收到的令牌数量少于其预先分配的容量，发送方 GPU 就必须用无意义的零值令牌来填充数据包，以凑满缓冲区大小。

FlashDMoE 完全消除了这种浪费。由于它使用直接的、单边的写入操作，它只发送那些真正需要被送到某个专家的令牌。这极大地节省了网络带宽，并避免了接收方 GPU 在无用的空数据上浪费计算周期。

尽管 FlashDMoE 仍会在本地进行小规模的“就地填充”（in-place padding），以便将数据规整为适合 GEMM 高效计算的“瓦片”（tile）形状，但这与填充到整个专家容量并进行网络传输相比，其开销可以忽略不计。

瓦片级并行与依赖管理 (Tile-Level Parallelism and Dependency Management)

为最大化计算与通信的重叠，FlashDMoE 将输入数据和计算任务拆分为“瓦片”（tile）级别的细粒度单元（通常为 128 或 256 个 token 组成的子批次）。这种拆分使得：

• 计算流水线化：一个瓦片的专家计算可与另一个瓦片的通信并行执行。例如，当 GPU 正在计算瓦片 A 的专家输出时，可同时将已完成的瓦片 B 结果通过 RDMA 发送到目标 GPU，无需等待所有瓦片计算完成。
• 通信重叠最大化：由于瓦片粒度远小于整个批次，通信操作能更早启动，且不会被单个大任务阻塞。这种流处理模式显著提升了硬件资源的时间利用率。

动态负载均衡

订阅者线程束实时跟踪各 GPU 的负载状态，调度器可根据当前网络和计算资源 availability，动态调整瓦片的传输目标（如将部分 token 路由到负载较轻的专家副本），避免单一 GPU 过载。

6. 实验设计

(对应问题六：实验是如何设计的？)

8 卡 H100，选择 Comet、FasterMoE、Megatron-CUTLASS 和 Megatron-TE 做基准。分别测试了 2、4、8 卡情况和最多 128 个专家的情况。序列长度最长 16k，top-k 路由选择 k - 2.

选择前向延迟、吞吐量、平均 SM 利用率和重叠效率（当问题规模随 GPU 数量增加时，系统隐藏通信成本的能力）作为评价指标。

7. 数据与可复现性

(对应问题七：实验是在什么样的数据集基础上运行的？)

使用合成工作负载（synthetic workloads），即通过程序生成具有特定维度（如序列长度、嵌入维度）的输入张量来进行测试。

代码开源了但是是纯 cuda 从头开始写的，并且用的卡是 H100

8. 实验结果的批判性分析

(对应问题八：实验结果能否有力地支持假设？)

结果整理如下表。

卡数增加延迟几乎不变，很妙啊

不过实验部分一个需要审慎评估的点是，作者明确指出 FlashDMoE 使用 FP32 精度进行评估，而所有基准系统均使用 FP16 精度。作者认为，使用 FP32 对 FlashDMoE 来说是一个不利因素（数据传输量和计算量翻倍），因此他们取得的优异结果是一个“保守的下界”。

• 作者本人在社区的讨论中以及一些外部评述都提到，由于未能为 CUTLASS 调整共享内存布局，其 FP16 性能并非最优。

9. 论文核心贡献的综合阐述

(对应问题九：这篇文章的贡献是什么？)

设计、实现并成功验证了首个完全融合的、GPU 常驻的 DMoE 算子，一种用于 GPU 动态任务调度的核函数内演员模型，一种高效的、感知载荷的通信协议。

10. 未来工作、影响与展望

(对应问题十：下一步可以做什么？)

当前实现仅支持推理，下一步可以尝试支持训练。

可以支持 fp16

可以与编译器结合，尝试像 Mirage 那样用优化器自动生成 CUDA 代码。

// EOF