FlashDMoE: Fast Distributed MoE in a Single Kernel

FlashDMoE：在单个内核中实现快速分布式MoE
作者/机构：Osayamen Jonathan Aimuyo, Byungsoo Oh, Rachee Singh (Cornell University)

A1 主要贡献

核心问题与研究目标：目前的混合专家（Mixture-of-Experts, MoE）模型虽然通过计算稀疏性实现了模型规模的亚线性计算成本增长，但现有实现存在GPU利用率低、显著的延迟开销和无法利用任务局部性等问题。这些问题主要源于CPU管理的调度、主机启动的通信以及频繁的内核启动。例如，现有工作观察到通信操作占用了高达68%的总运行时间，期间GPU完全空闲。此外，每次MoE前向传播需要启动大量（最多550个）内核，这引入了不必要的同步点和内存往返开销（见表1），并因动态路由模式而无法使用CUDA图进行优化。

图2：(a) 无MoE的Transformer块，(b) 带MoE的Transformer块，以及(c) 带分布式MoE和专家并行的Transformer块。T、E和O分别代表输入令牌、专家和输出激活。

创新点：为解决上述问题，本文提出了FlashDMoE，这是一个完全驻留在GPU上的MoE算子，它将专家计算和GPU间通信融合到一个单一的持久化GPU内核中（图3）。这种设计消除了内核启动开销，并通过细粒度的流水线化调度、计算和合并阶段来减少空闲间隙。
主要创新点如下：
1. 单一持久化内核设计：FlashDMoE将所有DMoE（分布式MoE）的计算和通信任务集成到一个内核中，该内核在整个MoE算子生命周期内保持活动状态，从而消除了CPU协调的多次内核启动开销。
2. 内核内调度与瓦片并行：FlashDMoE实现了瓦片级并行，并将线程块专门化为“处理器”（Processor）和“操作系统”（OS）两种角色。OS块负责调度计算工作和解码来自其他GPU的消息，确保GPU SM（流式多处理器）在DMoE算子生命周期内保持高利用率。
3. 异步且载荷高效的通信：FlashDMoE摒弃了传统的批量同步集合通信（如AllToAll），转而采用单边的、设备启动的GPU间(R)DMA传输。这种方法引入了载荷高效通信，只向真正被激活专家的GPU发送非填充的令牌，从而节省了通信带宽和计算资源。
4. 技术挑战的解决：为实现单一内核设计，FlashDMoE解决了多个技术挑战，包括轻量级计算依赖管理、优化SM占用率、实现设备内BLAS操作（使用【索引编号17，CUTLASS】）、最小化同步开销，以及利用NVSHMEM【索引编号16，Nvidia openshmem library (nvshmem), 2025】实现设备启动的DMA/RDMA通信。

图3：比较FlashDMoE与现有先进技术。左上图为不重叠通信与计算，左中图为部分重叠。FlashDMoE是一个持久化内核，融合了MoE算子的所有计算与通信（左下图）。FlashDMoE实现了设备启动的计算（门控、专家FFN、缩放）和通信任务（右图）。

表1：内核融合比较。我们的方法是第一个将DMoE层完全融合到单个GPU内核中的。我们报告了使用Nsight Systems分析得出的GPU操作。我们在2个A100上对单个层（门控 → 分发 → 专家 → 合并）进行计数，每个GPU有32个专家。

A3 背景知识与关键观察

同步通信与掉队者问题。当前MoE框架中使用的AlltoAll或AllGather通信是一种同步集合操作，需要所有参与的GPU共同完成。GPU之间处理速度或内核调度的差异会导致“掉队者效应”（straggler effect），这不仅损害了集合操作的性能，也影响了端到端性能，因为停滞的GPU必须等待集合操作完成后才能继续执行后续任务（如图4所示）。具体数据显示（图15），在一个包含32个A100 GPU的1.3B GPT-3 MoE模型分布式训练中，与平均内核执行时间相比，P95通信性能下降了1.32倍。在单节点虚拟机（VM）中，P95性能损失更是高达11倍（表2）。与先前的高性能计算工作【索引编号19，Optimizing bandwidth limited problems using one-sided communication and overlap, 2006, IPDPS】、【索引编号20，Atos: A task-parallel gpu scheduler for graph analytics, 2023, ICPP】一致，本文主张消除同步集合通信中的固有屏障，从而允许GPU将观察到的空闲时间重新用于下游计算，如图4所示。

图4：重叠调度（下图）展示了如何将顺序调度（上图）中的空闲时间用于计算。FlashDMoE实现了重叠调度。

表2：同步All-to-All通信中的掉队者延迟。我们捕捉了多步中由掉队者引起的延迟分布。设实际时间 $t_a$ 为所有GPU中最快的内核执行时间，总时间 $t$ 为记录的最大步骤时间。我们将延迟定义为 $t$ 和 $t_a$ 之间的最大差值。注意延迟即为空闲时间。对于1x8 V100，我们分析了1750步；对于8x4 A100，我们分析了600步。原始分布见图15。

内核启动开销。我们比较了FlashDMoE和现有基线的内核启动开销。表1显示了单次前向传播过程中的内核启动次数：FlashDMoE仅启动一个持久化内核，而基线方法则启动多达550个短生命周期的内核来完成相同计算。图5通过Nsight Systems捕捉的CUDA API跟踪记录直观地展示了FlashDMoE与DeepEP之间的差异。DeepEP显示出大量微小的CUDA API调用，并且在各个算子之间频繁停顿，导致GPU空闲时间增加（图5a）。相比之下，FlashDMoE通过避免启动开销和同步间隙来保持高GPU利用率——实现了93.17%的GPU利用率，而DeepEP仅为14%。

图5：5a展示了在2个通过NVLink连接的A100（单向带宽300 GB/s）上，对100次MoE前向传播的平均GPU利用率。尽管互连带宽很高，我们观察到高达90%的空闲时间，这归因于内核启动间隙和非重叠通信。

A2 方法细节

FlashDMoE设计概述。现代分布式MoE系统存在两大局限：(1) 关键路径上频繁的AlltoAll或AllGather集合通信，以及(2) 反复内核启动带来的显著开销。我们通过FlashDMoE来解决这些问题，它是一个完全融合的MoE算子，实现为单个持久化GPU内核。与以往方法【索引编号11，Comet: Finegrained computation-communication overlapping for mixture-of-experts, MLSys ’25】、【索引编号1，Deepseek-v3 technical report, 2025】、【索引编号10，DeepSpeed-MoE: Advancing mixtureof-experts inference and training to power next-generation AI scale, 2022, ICML】、【索引编号12，Megatron-lm, 2025】、【索引编号21，Tutel: Adaptive mixture-of-experts at scale, 2023, MLSys】、【索引编号8，Lancet: Accelerating mixture-of-experts training via whole graph computation-communication overlapping, 2024, MLSys】、【索引编号22，Fastermoe: modeling and optimizing training of large-scale dynamic pre-trained models, 2022, PPoPP】、【索引编号23，Flexmoe: Scaling large-scale sparse pre-trained model training via dynamic device placement, 2023, Proc. ACM Manag. Data】、【索引编号24，Schemoe: An extensible mixture-of-experts distributed training system with tasks scheduling, 2024, EuroSys】、【索引编号25，Harnessing inter-gpu shared memory for seamless moe communication-computation fusion, 2025, PPoPP】、【索引编号26，Flashattention: Fast and memory-efficient exact attention with io-awareness, 2022, NeurIPS】不同，FlashDMoE是首个实现完全融合的分布式MoE内核的解决方案，仅需一次内核启动，从而完全消除了内核启动开销（见表1）。

图6：FlashDMoE融合内核

输入: A, O ∈ R^(S×H), X ∈ R^(E×H×D), N
1 开始
2   Tϕ, Gϕ ← FusedGate(A)
3   如果 blockId + 1 < N 那么
4     Dispatch(Tϕ, A)
5     processor::start()
6   否则
7     如果 warpID == 0 那么
8       scheduler::start()
9     否则
10      subscriber::start(Tϕ, Gϕ, O, X)
11    结束如果
12  结束如果
13 结束

图7：表示为Actor交互链的DMoE功能依赖关系。我们用Sb、Sh和P分别表示Subscriber、Scheduler和Processor actor。对于任何actor a ∈ {Sb, Sh, P}，$a_i$ 标识GPU i上的一个actor。我们将 $D^j_i$ 定义为GPU j向GPU i分发瓦片数据包的操作。此图表达了瓦片粒度的任务依赖关系，即GEMM0、GEMM1、合并和通信产生一个输出瓦片。通知通过共享内存（subscriber ↔ scheduler）或全局内存（scheduler ↔ processor或GPU间通信）中的信号传播。注意，单边的GPU间传输（数据包或单个瓦片）与一个信号耦合，以通知接收GPU j上的 $S^j_b$ 消息已送达。

基于Actor的模型。FlashDMoE的设计基于并发计算的Actor模型【索引编号27，Actors: A model of concurrent computation in distributed systems, 1985】、【索引编号28，A universal modular actor formalism for artificial intelligence, 1973, IJCAI】、【索引编号29，SEMANTICS OF COMMUNICATING PARALLEL PROCESSES, 1975】。我们通过将GPU线程块和warp专门化为三种不同的actor角色来实现此模型：(1) 处理器(Processor)（§E.1），(2) 订阅者(Subscriber)（§E.3），和(3) 调度器(Scheduler)（§E.2）。处理器执行计算（GEMM和逐元素操作）和瓦片通信。我们使用CUTLASS【索引编号17，CUTLASS】作为高性能BLAS例程的基础设施，并使用NVSHMEM【索引编号16，Nvidia openshmem library (nvshmem), 2025】进行内核启动的通信。订阅者和调度器执行管理功能。具体来说，调度器将计算任务分配给可用的线程块。我们的关键创新是使调度器多线程化，从而实现高调度吞吐量，并采用工作保持（work-conserving）机制，确保始终保持高GPU SM利用率。另一方面，订阅者将来自对等GPU的瓦片数据包解码为任务描述符（§3.1）。在GPU上的N个线程块中，我们将N-1个专门化为处理器角色。我们将最后一个块专门化为操作系统（OS）。在此块内，我们将三个warp专门化为订阅者角色，一个warp专门化为调度器角色。这种线程块在actor间的划分是有意为之的：我们的目标是使用少量资源进行管理任务，同时将大部分资源保留用于执行MoE计算任务。图6总结了FlashDMoE架构及其组成的actor，而算法1则给出了该系统的非常接近的代码实现。注意，A ∈ $R^{S×H}$ 是输入令牌矩阵；O ∈ $R^{S×H}$ 是输出矩阵；X ∈ $R^{E×H×D}$ 是专家权重的3D张量，其中E表示执行GPU的本地专家数量，H是嵌入维度，D是FFN中间维度，S是序列长度。$T_ϕ$ ∈ $R^{2E×C}$ 是一个路由表数据结构，其中 $T_ϕ(e, c) = (i, w)$ 表示槽c中的令牌i被分派到专家e。w是合并权重（公式2），C是专家容量。$T_ϕ$ 的元组结构是一个实现细节。$G_ϕ$ ∈ $R^{S×E}$ 捕捉了门控函数产生的亲和度分数（公式3）。

FlashDMoE中的Actor间交互。FlashDMoE在瓦片（tensor的静态大小分区）粒度上分解MoE的计算和通信，以实现并行执行和任务的高效重叠。每个瓦片映射到一个由任务描述符封装的离散工作单元。订阅者从其接收的远程瓦片数据包中解码这些任务描述符。同时，调度器接收关于可用任务的通知，并将它们分派给执行这些任务所定义计算的处理器actor，即前馈网络（FFN）和专家合并操作。图7展示了actor交互链，演示了FlashDMoE如何强制执行DMoE的功能依赖关系。

图8：GPU内存层次结构。倒金字塔（左）显示了加载/存储访问延迟【索引编号30，Benchmarking and Dissecting the Nvidia Hopper GPU Architecture, 2024, IPDPS】、【索引编号31，Demystifying the nvidia ampere architecture through microbenchmarking and instruction-level analysis, 2022】、【索引编号32，Ptx isa: Version 8.7, 2025】。上表概述了不同内存层级的容量（针对A100 GPU）。共享内存和寄存器容量是FlashDMoE的静态配置。右图显示了可访问范围：片上寄存器作用域为单个线程；片上共享内存对块内所有线程可见；片外全局内存可被设备上所有线程访问。

在FlashDMoE中确定瓦片维度。在FlashDMoE中选择合适的瓦片维度对于确保高效的GPU利用率至关重要。尺寸过小的瓦片会使GPU利用不足，而尺寸过大的瓦片则会产生寄存器压力，导致性能下降的寄存器溢出到本地内存。经过仔细的参数扫描，我们选择的瓦片维度为(128, 64)。我们的关键见解是：增加瓦片宽度会显著提高每个线程的寄存器使用量，可能引发昂贵的溢出；在不调整线程数的情况下增加瓦片高度会增加每个线程的工作负载，从而损害性能。将每个块的线程数增加到超过我们固定的128个线程会减少并发块的数量，从而对SM占用率产生负面影响。更大的线程块大小还会增加块内同步（__syncthreads()屏障）的开销，进一步降低性能。因此，我们选择的瓦片维度在寄存器使用、共享内存约束和GPU占用率之间取得了平衡，以提供最佳性能。

3.1 用于计算的任务抽象

计算算子。FFN算子是Transformer架构中广泛使用的标准逐位置前馈网络【索引编号5，Attention is all you need, 2017, NeurIPS】，由两个线性变换和一个非线性激活函数ϕ（例如GELU或ReLU）组成：

$$\text{FFN}(x) = W_2 \cdot \phi(xW_1 + b_1) + b_2$$

这里，$W_1$ 和 $W_2$ 代表可学习的权重矩阵，而 $b_1$ 和 $b_2$ 是偏置。专家合并操作，用于GShard【索引编号33，Gshard: Scaling giant models with conditional computation and automatic sharding, 2021, ICLR】和DeepSeek【索引编号1，Deepseek-v3 technical report, 2025】等架构中，通过计算基于亲和度分数的加权组合来合并多个专家的输出：

$$C_{i}=\sum_{j=1}^{k} g_{i, e}$$ $$\mathbf{h}_i=\sum_{j=1}^k \frac{g_{i, e}}{C_i} \cdot \mathbf{h}_i^k$$

在这些方程中，$i ∈ 0, S-1$ 代表一个输入令牌的索引，$e = E_{i,k}$ 标识为令牌i选择的第k个专家，$g_{i,e}$ 是亲和度分数，表示专家e对令牌i的相关性。

统一的任务抽象。我们将FFN和合并操作统一在一个称为任务的通用抽象之下。任务为订阅者、调度器和处理器之间通信瓦片级工作提供了一个统一的接口。形式上，一个任务描述符 $t ∈ T$ 定义为一个元组：

$$t=(\mathcal{M}, \star, \phi)$$

其中M是一组元数据（例如，设备ID，瓦片索引），⋆是一个二元张量操作（具体来说是矩阵乘法·或哈达玛积⊙），而ϕ是一个逐元素的激活函数（例如，ReLU或恒等函数）。

我们定义一个任务t，它对输入张量A, B, D进行操作，产生输出张量C，如下所示：

$$\mathcal{F}_{t}(A, B, C, D):=C \leftarrow \phi\left(A \star_{t} B+D\right)$$

算子 $⋆_t$（从⋆实例化）可能会根据任务元数据M表现不同，并且 $A ⋆_t B$ 的结果会累加到D中。我们在§D中提供了一个任务元数据的示例。

在实践中，我们将由公式4定义的每个任务实现为一个单一融合的__device__修饰的函数，处理器（算法2）在运行时调用它。对t进行融合意味着将ϕ和后续的加法操作应用于存储二元算子 $⋆_t$ 结果的寄存器。为了说明其灵活性，我们展示了如何使用这个任务框架来表示FFN和专家合并操作。注意，为简化起见，我们省略了矩阵乘法符号（·）。此外，$ϕ_1$ 可以是任何激活函数，而 $ϕ_2$ 是恒等函数。FFN表示为：

$$\begin{aligned} \begin{aligned} t_{1} & =\left(\mathcal{M}, \cdot, \phi_{1}\right), \quad t_{2}=\left(\mathcal{M}, \cdot, \phi_{2}\right), \\ \mathcal{F}_{t_{1}}\left(A, B_{1}, C_{1}, D_{1}\right) & :=C_{1} \leftarrow \phi_{1}\left(A B_{1}+D_{1}\right), \\ \mathcal{F}_{t_{2}}\left(C_{1}, B_{2}, C_{2}, D_{2}\right) & :=C_{2} \leftarrow \phi_{2}\left(C_{1} B_{2}+D_{2}\right) . \end{aligned} \end{aligned}$$

而专家合并操作形式化为：

$$\begin{aligned} \begin{gathered} t_{3}=\left(\mathcal{M}, \odot, \phi_{2}\right), \\ \mathcal{F}_{t_{3}}(A, S, C, C):=C \leftarrow \phi_{2}(A \odot S+C) . \end{gathered} \end{aligned}$$

3.2 用于GPU间通信的对称张量布局

(b) DMA（上）和(a) 2个专家并行进程间的对称张量布局。RDMA（下）通信的状态机。

Actor通信机制。在单个GPU设备内，FlashDMoE中的actor通过GPU的内存子系统进行通信（见图8）。具体来说，调度器和订阅者actor通过快速的共享内存交换数据，而其他actor对则通过全局内存通信。对于跨多个设备的通信，FlashDMoE使用设备启动的通信，利用单边PGAS（分区全局地址空间）编程模型【索引编号34，Productivity and performance using partitioned global address space languages, 2007, PASCO】。然而，在PGAS中实现可扩展且正确的单边内存访问而不引入昂贵的同步是一个已知的挑战【索引编号1，Deepseek-v3 technical report, 2025】、【索引编号35，Triton-distributed: Programming overlapping kernels on distributed ai systems with the triton compiler, 2025】。我们通过一种可证明正确且可扩展的解决方案来应对这一挑战：一个对称张量布局L，支持完全非阻塞的内存访问。我们将L定义为：

$$L \in \mathbb{R}^{P \times R \times B \times E \times C \times H}$$

其中：P是专家并行的world size，R标识通信轮次（即两轮，一轮用于令牌分发，一轮用于合并），B是暂存缓冲区数量，E是本地专家数量，C是向上扩展的专家容量（§3.2.1），H是令牌嵌入维度。我们实现非阻塞通信的核心思想是时间缓冲。具体来说，我们为底层的令牌矩阵超配内存至少2·r倍，其中r是依赖图中的通信轮次数，因子2解释了每个通信轮次中用于传入和传出数据的独立缓冲区。对于MoE模型，我们有2·r = 4。这种内存使用的适度增加消除了单边数据传输期间的同步需求。图9b说明了如何在这个对称张量布局中索引和使用单元进行直接内存访问（DMA）和远程DMA（RDMA）操作。正如定理3.1所强调的，这种在L上的索引方案是允许完全非阻塞访问并避免同步的根本机制，因为所有访问都是无写冲突的。证明见§C。

定理3.1。对称张量布局L是无写-写冲突的。

对称布局的构建。为了构建L，我们从原始的令牌缓冲区 T ∈ $R^{S×H}$ 开始，其中S是序列长度，H是令牌嵌入维度。我们首先将序列维度S重组为三个子维度，分别代表专家容量（C）、本地专家槽位（E）和专家并行world size（W），使得：

$$C \cdot E \cdot W = C \cdot E' = S', \quad \text{where} \quad S' \ge S \text{ and } E' \ge E_W$$

在典型的均匀专家分布情况下（如图9a所示），我们有 S' = S 和 E' = $E_W$，其中 $E_W$ 是模型中的专家总数。因此，令牌缓冲区的大小为 Size(T) = S' · H。在图9a中，每个标记为 $E_i$（i ∈ {0, . . . , 3}）的单元是一个大小为 (C, H) 的矩阵。我们扩展了先前的工作【索引编号33，Gshard: Scaling giant models with conditional computation and automatic sharding, 2021, ICLR】、【索引编号11，Comet: Finegrained computation-communication overlapping for mixture-of-experts, MLSys ’25】，引入了额外的时间维度R（通信轮次）和B（暂存缓冲区）。每个通信轮次有两个固定的暂存槽：一个用于传出令牌，另一个用于传入令牌。每个由维度P索引的槽形成一个形状为 (S', H) 的张量。因此，张量大小Size(L)通常至少是原始令牌缓冲区大小的四倍，在均匀专家分布的情况下恰好是四倍。根据经验，我们发现：

$$Size(L) \approx 4 \cdot Size(T)$$

3.2.1 用于载荷效率的原地填充

原地填充策略。由于MoE分发中令牌的动态和不均匀分布【索引编号36，Blackmamba: Mixture of experts for state-space models, 2024】，GPU通常接收到的令牌数量少于其预定义的专家容量。当前的MoE框架【索引编号10，DeepSpeed-MoE: Advancing mixtureof-experts inference and training to power next-generation AI scale, 2022, ICML】通常在计算前用空令牌填充这些缓冲区，这不必要地增加了通信载荷并降低了性能。相反，我们提出了原地填充，直接在本地对称张量缓冲区内执行填充，从而消除了多余的网络通信。

容量对齐。如图9a所示，每个单元Ei的大小根据专家容量C来确定。我们进一步对齐此容量，以确保其可被瓦片块大小 $b_M$ = 128 整除，从而保证处理器线程在消费远程令牌时能够进行安全和对齐的内存读取。这种原地填充策略略微增加了L的内存占用，如下所述：

$$\begin{aligned} Size(L) \approx \begin{cases} 4 \cdot Size(T), & \frac{S}{E} \geq bM \\ 4 \cdot \frac{bM \cdot E}{S} \cdot Size(T), & \text{otherwise} \end{cases} \end{aligned}$$

A4 实验环境

• 硬件配置：实验在一台配备8个NVIDIA H100 80G GPU的服务器上进行，GPU之间通过NVLink互连。服务器拥有125GB内存和20个vCPU。
• 软件配置：使用PyTorch 2.6.0、CUDA 12.8和Ubuntu 22.04。
• 模型配置：所有实验均使用MoE Transformer模型，配置为16个注意力头，嵌入维度为2048，FFN中间层大小为2048。
• 实验设置：所有实验均采用分布式数据并行（DDP）和专家并行（EP）。我们只执行单层MoE的前向传播，并在32次预热后测量32次传播的平均运行时间。路由策略为top-2，容量因子为1.0。
• 对比基线：FlashDMoE与多个先进的MoE系统进行比较：(1) Comet【索引编号11，Comet: Finegrained computation-communication overlapping for mixture-of-experts, MLSys ’25】，(2) FasterMoE【索引编号13，Fastermoe: modeling and optimizing training of large-scale dynamic pre-trained models, 2022, PPoPP】，(3) Megatron-CUTLASS【索引编号37，Efficient large-scale language model training on gpu clusters using megatron-lm, 2021】，以及(4) Megatron-TE：带有Transformer Engine的Megatron-LM【索引编号38，Transformer engine】。Comet依赖cudaMemcpyPeerAsync【索引编号39，Flux’s overlap performance is worse than non-overlap [4xrtx 4090], 2025】，而FasterMoE和Megatron-LM则完全使用NCCL进行通信。

A4 实验结果

评估前提。在实验中，我们观察到Comet在8个GPU上的性能异常差，因此在8 GPU的评估中排除了其结果。一个重要前提是，FlashDMoE使用FP32精度进行评估，而所有基线都使用FP16。这样做是因为基线不支持FP32，且时间限制使我们无法将系统调整到FP16的最佳性能。这种精度差异实际上使FlashDMoE处于不利地位，因为其通信和计算精度要求加倍，因此我们的结果可以看作是性能的保守下限。尽管如此，FlashDMoE仍然超越了所有基线。

前向延迟。我们测量了FlashDMoE在不同序列长度下（4和8个GPU）的前向延迟（图10）。FlashDMoE始终优于所有基线，尤其是在长序列下优势更为明显。在4个GPU上，当令牌数为16K时，其速度比Megatron-TE快4.6倍，比FasterMoE快2.6倍。在8个GPU上，增益更加显著，速度比基线快6.4倍，因为基线性能因通信成本随令牌缓冲区增加而急剧下降。

图10：随着每个GPU的令牌数量增加，前向延迟的变化。

GPU利用率。为了量化GPU效率，我们测量了前向传播期间的流式多处理器（SM）利用率（图11）。FlashDMoE实现了93.17%的平均SM利用率，比FasterMoE（9.67%）高9倍以上，比DeepEP+Megatron-LM（13.55%）高6.8倍，比Megatron-TE（59.11%）高4倍，比Comet（42.31%）高2.2倍。这得益于其完全融合的内核架构和计算与通信任务的细粒度流水线。

图11：SM利用率对比，定义为SM有至少一个warp在执行的周期数与总周期数的比率【索引编号40，NVIDIA Nsight Systems Metrics】。数值为100次迭代的平均SM利用率。所有实验使用T=8K和E=64，在两个A100上进行。

重叠效率。我们通过测量弱扩展效率来评估FlashDMoE重叠通信和计算的程度（图12b）。以2个GPU为基准，Megatron-CUTLASS和Megatron-TE的重叠效率在≥4个GPU时降至0.5以下。而FlashDMoE在4和8个GPU上分别实现了高达3.88倍和4倍的效率提升。图12a进一步说明了这一点，FlashDMoE的延迟增长平稳，而基线则表现出近似线性的延迟放大。这证实了FlashDMoE基于actor的设计和异步数据移动实现了近乎理想的重叠。

图12：随着每个GPU的令牌数量增加，前向延迟的变化。我们将重叠效率Oe定义为Oe = T(2)/T(NG)，其中T(NG)是NG个GPU的延迟，T(2)是2个GPU的延迟。

吞吐量。如图13所示，FlashDMoE的吞吐量随GPU数量线性扩展，在8个GPU上达到17.7 MTokens/s。这比FasterMoE高5.7倍以上，比Megatron-TE和Megatron-CUTLASS高4.9倍。值得注意的是，这些结果是在FlashDMoE完全使用FP32的情况下取得的，而基线使用FP16。

图13：随着GPU数量增加的吞吐量。我们计算吞吐量为 T*NG/latency，其中NG是GPU的数量。

专家可扩展性。我们分析了在固定序列长度（T=16K）下，FlashDMoE随专家数量增加的扩展性。如图14所示，即使专家数量从8个增加到128个，FlashDMoE仍保持低且均匀的延迟。相反，基线由于内核启动开销增加而表现出超线性的延迟增长。在4个H100上，FlashDMoE的性能比基线高出4倍，在8个H100上高出6.6倍（均为128个专家时）。

图14：随着专家数量增加的前向延迟。

内存开销。我们测量了FlashDMoE所需的GPU内存，包括对称张量L和运行时簿记状态。内存开销主要取决于瓦片大小、专家容量（EC）和专家数量（E）。表3总结了各种配置下的内存开销，证实了FlashDMoE保持了适度和可预测的内存占用。

表3：FlashDMoE的内存开销（瓦片大小 bM = 128），Size(T) = Tokens * 4KB）。

A7 补充细节 (局限性与未来工作)

编程复杂性。开发完全融合的持久化内核是一项非常重要的工程任务，需要深入了解GPU架构、同步、分布式协议和内存层次结构。这种高门槛限制了其采用。未来的工作可以考虑编译器级别的抽象或DSL来普及这种技术。

FP16支持与共享内存访问模式。尽管现代GPU原生支持半精度计算，但将FlashDMoE适配到FP16对于处理器的计算算子来说并非易事。具体来说，我们手动调整的交错共享内存布局对于我们用来实现设备内GEMM的CUTLASS的Collective Mainloop算子来说不是最高效的模板参数。这种次优配置降低了内存吞吐量，如§H所示。在Ampere及更早的GPU上克服这一点需要仔细研究最佳布局，但对于Hopper及以上的GPU，我们预计在未来的改进中使用CUTLASS提供的构建器接口。

缺乏反向传播和训练支持。虽然这项工作专注于推理，但要支持训练，需要将反向计算和梯度通信融合到内核中。这需要对内存簿记和任务描述符定义进行非平凡的修改。尽管如此，这仍然是扩展该系统以完全支持端到端训练的一个令人兴奋的方向。

A5 结论

本文介绍了FlashDMoE，这是第一个将整个混合专家（MoE）算子融合到单个持久化GPU内核中的系统。我们指出，现有的MoE实现存在两个关键效率问题：(1) 由CPU管理的同步通信导致互连利用率不足，以及(2) 通过多个GPU内核进行碎片化执行，引入了开销和同步延迟。

相比之下，FlashDMoE通过将计算、通信和调度嵌入一个统一的内核中，实现了GPU的自主性。它利用了actor风格的并发、warp专门化和异步（R）DMA，以实现细粒度的通信-计算重叠。

我们的评估表明，与最先进的系统相比，FlashDMoE在分布式MoE上实现了高达6倍的速度提升、9倍的GPU利用率改善和5.7倍的吞吐量增加。FlashDMoE挑战了分布式深度学习中主流的执行范式，并为构建未来的GPU原生系统提供了一个引人注目的模板。

尽管仍存在一些局限性，如编程复杂性和缺乏FP16支持，但这项工作为内核内分布式计算的新时代奠定了基础。未来的系统可以在此基础上构建，实现整个训练流水线的内核融合，从而迎来从CPU编排到完全自主GPU执行的设计转变。

A6 附录

A 相关工作

计算-通信重叠与内核融合。为减少分布式DNN训练中同步的通信开销，许多研究致力于增加计算与通信的重叠。对于不含MoE层的通用Transformer模型，许多工作【索引编号41，Breaking the computation and communication abstraction barrier in distributed machine learning workloads, 2022, ASPLOS】、【索引编号42，Overlap communication with dependent computation via decomposition in large deep learning models, 2022, ASPLOS】、【索引编号43，Centauri: Enabling efficient scheduling for communication-computation overlap in large model training via communication partitioning, 2024, ASPLOS】、【索引编号44，T3: Transparent tracking & triggering for fine-grained overlap of compute & collectives, 2024, ASPLOS】、【索引编号45，{MegaScale}: Scaling large language model training to more than 10,000 {GPUs}, 2024, NSDI】、【索引编号46，CO2: Efficient distributed training with full communication-computation overlap, 2024, ICLR】、【索引编号47，Better together: Jointly optimizing {ML} collective scheduling and execution planning using {SYNDICATE}, 2023, NSDI】、【索引编号48，Harnessing inter-gpu shared memory for seamless moe communication-computation fusion, 2025, PPoPP】、【索引编号49，Optimizing distributed ml communication with fused computation-collective operations, 2024, SC24】提供了划分和调度计算与通信操作的技术，旨在实现更细粒度的重叠。为了应对MoE训练中AllToAll通信和专家并行带来的挑战，Tutel【索引编号50，Tutel: Adaptive mixture-of-experts at scale, 2023, MLSys】和FasterMoE【索引编号13，Fastermoe: modeling and optimizing training of large-scale dynamic pre-trained models, 2022, PPoPP】将AllToAll与专家计算重叠。Lancet【索引编号51，Lancet: Accelerating mixture-of-experts training via whole graph computation-communication overlapping, 2024, MLSys】进一步使前向传播中的非MoE计算和后向传播中的权重梯度计算都能与AllToAll重叠。尽管实现了重叠，但这些方法的性能在实践中受限于带有屏障的阻塞式同步集体通信。相比之下，FlashDMoE通过在单个内核内实现异步、设备启动的数据传输与瓦片化计算的重叠，从根本上消除了这些低效问题。FlashDMoE与COMET【索引编号11，Comet: Finegrained computation-communication overlapping for mixture-of-experts, MLSys ’25】和DeepEP【索引编号1，Deepseek-v3 technical report, 2025】等前沿工作进一步区分开来，后者也使用这种形式的内核启动通信，但粒度较粗且没有实现完全的内核融合。

B 动机图示

图15：同步AllToAll的掉队者效应。M × N A100或V100表示M个节点中每个节点内有N个GPU。每个GPU在每个AllToAll步骤中与其他所有GPU通信。我们捕捉了多步中由掉队者引起的延迟分布。实际时间 $t_a$ 表示所有GPU中最快的内核执行时间，而总时间 $t$ 是记录的最大步骤时间，延迟是 $t$ 和 $t_a$ 之间的最大差值。注意延迟即为空闲时间。

C 定理3.1的证明

初始定义。我们首先给出证明所需的两个必要定义。

定义 C.1。定义一个写操作为 $w(p_s, p_t, i)$，其中 $p_s$ 是源进程，i是一个有序元组，表示位于目标进程 $p_t$ 上的L的索引坐标。当存在至少两个不同的、未同步的、并发的写操作 $w_1(p_{s1}, p_{t1}, i_1)$ 和 $w_2(p_{s2}, p_{t2}, i_2)$，满足 $p_{t1} = p_{t2}$ 且索引坐标 $i_1 = i_2$ 但 $p_{s1} \neq p_{s2}$ 时，发生写-写冲突。

定义 C.2。对于任何源进程 $p_s$，一个有效的索引坐标 $i = (p^*, r, b, e, c)$ 满足以下条件：
1. 对于设备间写操作，必须满足 $p^* = p_s$ 且 $b=1$。这也适用于自循环写操作 $w(p_t, p_t, i)$。
2. 对于任何写操作 $w(p_s, p_t, i)$，如果 $b=0$，则 $p_s = p_t$。此规则描述了设备内的暂存写操作。

定理C.1（重述）。对称张量布局L是无写-写冲突的。

证明。假设每个索引坐标i映射到L中一个不同的内存段。接下来，我们通过反证法证明，使用有效的i访问L时不存在写-写冲突。为简化，我们仅用索引坐标描述写操作。假设存在至少两个写操作 $w_1(p_{s1}, p_{t1}, i_1)$, $w_2(p_{s2}, p_{t2}, i_2)$，其中 $p_{t1} = p_{t2}$ 且目标坐标 $i_1, i_2$ 有效，$i_1 = i_2$（按字典序），其展开如下：

$$i_1 = (p_1, r_1, b_1, e_1, c_1), i_1 = (p_2, r_2, b_2, e_2, c_2)$$

对于消息暂存状态，尽管 $i_1 = i_2$，但最终的内存段分别位于 $p_{s1}$ 和 $p_{s2}$ 的不同物理缓冲区中。因此，在此状态下没有冲突，因为进程内写操作总是有不同的 $c_j$ 坐标（j ∈ {0, C-1}）。对于进程间传输，我们有两种情况：

情况1：$p_{s1} = p_{s2}$。此时，$w_1$ 和 $w_2$ 是相同的操作。这与冲突的定义（要求 $p_{s1} \neq p_{s2}$）相矛盾。实际上，这种重复的写操作永远不会发生。

情况2：$p_{s1} \neq p_{s2}$。为确保 $i_1$ 和 $i_2$ 的有效性，必须有 $p_1 = p_{s1}$ 和 $p_2 = p_{s2}$。然而，这意味着 $i_1 \neq i_2$，从而产生矛盾。证毕。□

D 任务实现

图16：任务结构体。TaskType ∈ {GEMM0, GEMM1, Combine}

E Actors

E.1 处理器

算法2：处理器Actor：由一个块执行

1 开始
2   tQ ← GetTQ()
3   signal ← 0
4   // 共享内存变量
5   task ← {}
6   interrupt ← False
7   complete ← False
8   当 interrupt == False 时循环
9     如果 warpId == 0 那么
10      如果 threadId == 0 那么
11        awaitTaskFromScheduler(interrupt, signal)
12        FencedNotifyRQ(ready)
13      结束如果
14      syncwarp()
15      warpReadTQ(tQ, signal, task)
16    结束如果
17    syncthreads()
18    如果 interrupt == False 那么
19      根据 task.Type 进行选择
20        情况 GEMM0:
21          // 融合的GEMM、收尾和异步瓦片暂存
22          fGET(GEMM0, task)
23          如果 threadId == 0 那么
24            complete ← NotifyTileCompletion()
25          结束如果
26          syncthreads()
27          如果 complete == True 那么
28            NotifySchedulerNextGEMM(tQ)
29          结束如果
30        结束情况
31        情况 GEMM1:
32          // 融合的GEMM、收尾和异步瓦片传输
33          fGET(GEMM1, task)
34        结束情况
35        情况 Combine:
36          combine(task)
37        结束情况
38      结束选择
39    结束如果
40  结束循环
41 结束

E.2 调度器

算法3：调度器Actor：由一个warp执行

1 开始
2   scheduled ← 0
3   tTB ← 0
4   tqState ← {}
5   pTDB ← GetProcessorDoorbell()
6   sTDB ← GetSubscriberDoorbell()
7   taskBound ← GetTaskBound()
8   tTB ← AtomicLoad(taskBound)
9   // 循环缓冲区就绪队列
10  rQ ← {}
11  // 用处理器ID填充就绪队列
12  PopulateRQ(rQ)
13  当 scheduled < tTB 时循环
14    lt ← 0
15    并行执行
16      扫描门铃并用观察到的任务计数填充tqState
17      将本地观察到的任务计数聚合到lt中
18    结束并行
19    qS, taskTally ← 0
20    // qS是包含性输出
21    WarpInclusiveSum(lt, qS, tasktally)
22    当 tasktally > 0 时循环
23      用就绪的处理器ID重新填充rQ
24      并行执行
25        从rQ[qS]开始，向处理器发送来自tqState的任务索引信号
26      结束并行
27    结束循环
28    如果 threadId == 0 那么
29      tTB ← AtomicLoad(taskBound)
30    结束如果
31    tTB ← WarpBroadcast(tTB)
32  结束循环
33  InterruptSubscribers()
34  InterruptProcessors()
35 结束

E.3 订阅者

F 多节点评估

F.1 设置

实验目的。本实验旨在多节点环境下评估FlashDMoE。我们使用4个节点，每个节点包含4个通过NVLink完全互连的A100 GPU。跨节点通信时，每个GPU使用一个提供25 GB/s带宽的NIC。我们将专家数量设为16，并为每个GPU分配一个专家，因此本地专家数量为1。我们将MIV（最大内向流量）正式定义如下：

$$MIV = \frac{Tokens}{Experts} * local\_experts * precision * hidden\_size * 2 * n_{rg}$$

其中 $n_{rg}$ 是远程对等体的数量，乘数2表示通信轮次（分发和合并）。本实验中 $n_{rg}$ = 12。

F.2 结果

图17：多节点延迟评估。嵌入维度为1024，FFN中间层大小为4096。我们将最大内向流量（MIV）定义为单个内向流量事件中NIC接收数据的最坏情况上限。

延迟与网络故障。我们观察到，随着令牌数量的增加，延迟呈亚线性增长。然而，当令牌数 > 2048时，应用程序因未能接收到预期消息而无法终止。我们推测这是由于网络硬件层的缓冲区溢出所致，这在产生大量和大消息的应用程序中很常见【索引编号18，Network - NERSC Documentation, 2025】。我们注意到，这个问题可以通过调整硬件配置来解决【索引编号52，fi_cxi, 2025】，但我们将此探索视为与本工作正交的练习。

G 实现

表4：FlashDMoE的实现指标。

H FP16内存吞吐量

图18：此处我们报告了FlashDMoE的FP16（上）和FP32（下）变体的总A100内存吞吐量。值得注意的是，在相同工作负载下，FP16实现发出的共享内存指令大约是其FP32对应版本的2倍。我们将这种低效归因于FlashDMoE在处理半精度数据时共享内存布局不佳。虽然这个瓶颈可以通过改进布局策略来解决，但由于时间限制，我们将其留给未来的工作。