FlashMoE: Fast Distributed MoE in a Single Kernel
FlashMoE: Fast Distributed MoE in a Single Kernel
作者/机构: Osayamen Jonathan Aimuyo (Cornell University), Byungsoo Oh (Cornell University), Rachee Singh (Cornell University)
A1 主要贡献
核心问题: 现有的专家混合（MoE）模型实现存在GPU利用率低、延迟开销大以及无法利用任务局部性等问题。这些问题主要源于CPU管理的调度、主机启动的通信以及频繁的内核启动。
研究目标: 克服上述限制，开发一个完全驻留在GPU上的MoE算子，该算子将专家计算和GPU间通信融合到一个单一的持久化GPU内核中。
创新点:
1. 单内核融合（Megakernel）: 开发了FlashMoE，这是一个将所有MoE计算（门函数、专家FFN、合并）和通信任务集成到一个单一持久化GPU内核中的巨型内核。该内核在MoE算子的整个生命周期内保持活跃，从而消除了由CPU协调的多次内核启动，显著减少了CPU的参与。
2. 细粒度流水线: 在融合内核内部，FlashMoE实现了调度（dispatch）、计算（compute）和合并（combine）阶段的细粒度流水线，消除了内核启动开销和空闲间隙。
3. 异步且高效的通信: FlashMoE摒弃了传统的批量同步集合通信（如AlltoAll），转而采用单边的、由设备发起的GPU间(R)DMA传输。这种方法避免了因单个慢GPU拖累所有其他GPU的“掉队者效应”。
4. 有效载荷效率（Payload Efficiency）: 通过只向激活了专家的GPU发送非填充的令牌，FlashMoE消除了稀疏激活层中膨胀或冗余的网络载荷，从而节省了通信带宽和计算资源。
5. 内核内调度与并行:
* 基于Actor的模型: FlashMoE实现了一种响应式编程模型，将GPU线程块专门化为处理器（Processor）、调度器（Scheduler）和订阅者（Subscriber）三种角色。
* 分块并行（Tile Parallelism）: 将输入令牌矩阵划分为称为“分块”（tile）的独立单元，由线程块处理，但由warp进行管理（调度和构建），以最大化GPU的算术强度和并行度。
* 专用OS块: 指定一个专用操作系统（OS）块（4个warp）来执行管理任务，包括向处理器调度计算工作和从其他GPU接收的消息中解码计算任务。
图 1: FlashMoE 性能。
图 2: (a) 不含MoE的Transformer块，(b) 含MoE的Transformer块，以及(c) 含分布式MoE和专家并行的Transformer块。T、E和O分别代表输入令牌、专家和输出激活。
图 3: FlashMoE与现有先进技术的比较。这些技术要么不重叠通信和计算（左上），要么进行部分重叠（左中）。FlashMoE是一个持久化内核，融合了MoE算子的所有计算和通信（左下）。FlashMoE实现了设备发起的计算（门、专家FFN、缩放）和通信任务（右）。
表 1: 我们通过Nsight Systems [17] 分析报告了MoE实现所启动的GPU操作数量。我们统计了在2个A100 GPU上单个MoE层（门 → 调度 → 专家 → 合并）的操作，其中每个GPU有32个专家。FlashMoE是第一个将分布式MoE层完全融合到单个GPU内核中的工作。
A3 背景知识与关键观察
同步通信的弊端。目前在MoE框架中使用的AlltoAll或AllGather通信是一种同步集合操作，其完成需要所有相关GPU的参与。在这种模式下，GPU之间处理速度或内核调度的差异会引发“掉队者效应”（straggler effect），这对(1)集合操作的性能和(2)端到端性能都是有害的（见图13），因为停滞的GPU在集合操作结束前无法继续执行下游的依赖或独立任务。我们在附录A中详细阐述了这个问题。
图 4: 4a显示了在2个A100（单向带宽300 GB/s）上运行100次MoE前向传播的平均GPU利用率，由于内核启动间隙和非重叠通信，我们观察到高达90%的空闲时间。
内核启动开销。我们比较了FlashMoE与现有基线方法的内核启动开销。表1显示了单次前向传播过程中的内核启动次数：FlashMoE仅启动一个持久化内核，而基线方法为完成相同的计算最多启动了550个短生命周期的内核。图4使用NSight Systems捕获的CUDA API调用轨迹进行了可视化比较，展示了FlashMoE和DeepEP之间的差异。DeepEP表现出大量小型CUDA API调用，在各个算子之间频繁停顿，导致GPU空闲时间增加（图4a）。相比之下，FlashMoE通过避免启动开销和同步间隙来维持高GPU利用率——实现了93.17%的GPU利用率，而DeepEP仅为14%。实验结果详见第4节，相关工作讨论见附录B。
A2 方法细节
FlashMoE的核心设计。现代分布式MoE系统存在两个局限性：(1) 关键路径上频繁的“多对多”（AlltoAll或AllGather）集合通信，以及(2) 重复内核启动带来的显著开销。我们在FlashMoE中解决了这些问题，它是一个完全融合的MoE算子，作为单个持久化GPU内核实现。与以往的方法【[12] Shulai Zhang, 等. Comet: Finegrained computation-communication overlapping for mixture-of-experts. MLSys ’25.】, 【[1] DeepSeek-AI. Deepseek-v3 technical report, 2025.】, 【[11] Samyam Rajbhandari, 等. DeepSpeed-MoE: Advancing mixtureof-experts inference and training to power next-generation AI scale. ICML 2022.】, 【[13] NVIDIA. Megatron-lm, 2025. v0.11.0.】, 【[20] Changho Hwang, 等. Tutel: Adaptive mixture-of-experts at scale. MLSys 2023.】, 【[10] Chenyu Jiang, 等. Lancet: Accelerating mixture-of-experts training via whole graph computation-communication overlapping. MLSys 2024.】, 【[21-25]】不同，FlashMoE是第一个实现完全融合的分布式MoE内核的解决方案，通过仅需一次内核启动，完全消除了内核启动开销（见表1）。
图 5: FlashMoE 融合内核
基于Actor的模型。FlashMoE的设计基于并发计算的Actor模型【[26] Gul A. Agha. Actors: A model of concurrent computation in distributed systems. 1985.】, 【[27] Carl Hewitt, Peter Bishop, and Richard Steiger. A universal modular actor formalism for artificial intelligence. IJCAI’73.】, 【[28] Irene Greif. SEMANTICS OF COMMUNICATING PARALLEL PROCESSES. PhD thesis, Massachusetts Institute of Technology, 1975.】。我们通过将GPU线程块和warp专门化为三种不同的actor角色来实现此模型：(1) 处理器（Processor）（§F.1），(2) 订阅者（Subscriber）（§F.3），和 (3) 调度器（Scheduler）（§F.2）。处理器执行计算（GEMM和逐元素操作）和分块（tile）通信。我们使用CUTLASS【[19] Vijay Thakkar, 等. CUTLASS, 2025.】作为高性能BLAS例程的基础设施，并使用NVSHMEM进行内核启动的通信【[18] NVIDIA. Nvidia openshmem library (nvshmem), 2025. v3.2.5.】。订阅者和调度器执行管理功能。具体来说，调度器将计算任务分配给可用的线程块。我们的关键创新在于使调度器既是多线程的，从而实现高调度吞吐量，又是工作保持（work-conserving）的，确保始终保持较高的GPU SM利用率。另一方面，订阅者将从对等GPU接收到的分块数据包解码为任务描述符（§3.1）。
Actor资源分配与架构。在GPU上的N个线程块中，我们将N-1个专门化为处理器角色。我们将最后一个块专门化为操作系统（OS）。在此块内，我们将三个warp专门化为订阅者角色，一个warp为调度器角色。这种线程块在actor间的分配是经过深思熟虑的：我们的目标是用少量资源执行管理任务，同时保留大部分资源用于执行MoE计算任务。图5总结了FlashMoE的架构及其组成的actor，而算法1则给出了该系统非常接近代码的翻译。请注意，$A \in R^{S \times H}$ 是输入令牌矩阵；$O \in R^{S \times H}$ 是输出矩阵；$X \in R^{E \times H \times D}$ 是专家权重的三维张量，其中E表示执行GPU的本地专家数量，H是嵌入维度，D是FFN中间维度，S是序列长度。$T_{\phi} \in (N \times R)^{E \times C}$ 是一个路由表数据结构，其中 $T_{\phi}(e, c) = (i, w)$ 表示槽c中的令牌i被分派给专家e。w是合并权重（公式2），C是专家容量。$T_{\phi}$ 的元组结构是一个实现细节。$G_{\phi} \in R^{S \times E}$ 捕获了门函数产生的亲和度分数（公式3）。
Input: A, O ∈ RS×H , X ∈ RE×H×D, N
1 begin
2 Tϕ, Gϕ ← FusedGate(A)
3 if blockId + 1 < N then
4 Dispatch(Tϕ, A)
5 processor::start()
6 else
7 if warpID == 0 then
8 scheduler::start()
9 else
10 subscriber::start(Tϕ, Gϕ, O, X)
11 end if
12 end if
13 end
Actor间的交互。FlashMoE以分块（tile）的粒度分解MoE的计算和通信，分块是张量的一个静态大小的分区，以实现并行执行和任务的高效重叠。每个分块映射到一个由任务描述符封装的离散工作单元。订阅者从它接收到的远程分块数据包中解码这些任务描述符。同时，调度器接收到关于可用任务的通知，并将它们分派给处理器actor执行，这些处理器执行由这些任务定义的计算，即前馈网络（FFN）和专家合并操作。图6展示了actor交互链，演示了FlashMoE如何强制执行分布式MoE（DMoE）的功能依赖关系。
图 6: DMoE功能依赖性表示为Actor交互链。我们用Sb、Sh和P分别表示订阅者、调度器和处理器Actor。对于任何actor a ∈ {Sb, Sh, P}，ai表示GPU i上的一个actor。我们定义Dji为操作符，其中GPU j向GPU i分发分块数据包。此图以分块的粒度表示任务依赖，即GEMM0、GEMM1、合并和通信产生一个输出分块。通知通过共享内存（订阅者↔调度器）或全局内存（调度器↔处理器或GPU间通信）传播的信号发生。注意，单向的GPU间传输（数据包或单个分块）与一个信号耦合，以通知接收GPU j上的Sjb消息已送达。
确定Tile维度。在FlashMoE中选择合适的tile（分块）维度对于确保高效的GPU利用率至关重要。过小的tile无法充分利用GPU，而过大的tile会产生寄存器压力，导致性能下降的寄存器溢出到本地内存。经过仔细的参数扫描，我们选择了(128, 64)的tile维度。我们的关键见解是：增加tile宽度会显著提高每个线程的寄存器使用量，可能引发昂贵的溢出；在不调整线程数的情况下增加tile高度会增加每个线程的工作量，损害性能。将每个块的线程数增加到超过我们固定的128个线程会减少并发块的数量，对SM占用率产生负面影响。更大的线程块大小还会增加块内同步（__syncthreads()屏障）的开销，进一步降低性能。因此，我们选择的tile维度在寄存器使用、共享内存约束和GPU占用率之间取得了平衡，以提供最佳性能。
3.1 任务抽象化计算
计算算子。FFN算子是Transformer架构中广泛使用的标准逐位置前馈网络【[6] Ashish Vaswani, 等. Attention is all you need. NeurIPS 2017.】，由两个线性变换组成，中间由一个非线性激活函数$\phi$（例如GELU或ReLU）分隔：
这里，$W_1$和$W_2$代表可学习的权重矩阵，而$b_1$和$b_2$是偏置。专家合并操作，用于像GShard【[29] Dmitry Lepikhin, 等. Gshard: Scaling giant models with conditional computation and automatic sharding. ICLR 2021.】和DeepSeek【[1] DeepSeek-AI. Deepseek-v3 technical report, 2025.】这样的架构中，通过计算基于亲和度分数的加权组合来合并多个专家的输出：
统一任务抽象。我们将FFN和合并操作统一到一个称为“任务”的通用抽象之下。任务为订阅者、调度器和处理器之间通信分块级别的工作提供了统一的接口。形式上，一个任务描述符 $t \in T$ 定义为一个元组：
$t = (\mathcal{M}, \star, \phi)$
其中$M$是一组元数据（例如，设备ID，分块索引），$\star$是一个二元张量操作（具体来说是矩阵乘法$\cdot$或哈达玛积$\odot$），而$\phi$是一个逐元素的激活函数（例如，ReLU或恒等函数）。我们定义一个任务t作用于输入张量A, B, D，产生输出张量C，如下所示：
$F_t(A, B, C, D) := C \leftarrow \phi(A \star_t B + D)$
算子$\star_t$（从$\star$实例化）可能会根据任务元数据$M$表现不同，并且$A \star_t B$的结果会累加到D中。我们在附录E中提供了一个任务元数据的例子。
任务实现。在实践中，我们将公式4定义的每个任务实现为单个融合的__device__修饰函数，处理器（算法2）在运行时调用它。对t进行融合意味着将$\phi$和后续的加法操作应用于存储二元算子$\star_t$结果的寄存器中。为了说明其灵活性，我们展示了如何使用这个任务框架来表示FFN和专家合并操作。注意，为简化起见，我们省略了矩阵乘法符号（·）。此外，$\phi_1$可以是任何激活函数，而$\phi_2$是恒等函数。FFN表示为：
3.2 用于GPU间通信的对称张量布局
Actor间通信。在单个GPU设备内部，FlashMoE中的actor通过GPU的内存子系统进行通信。具体来说，调度器和订阅者actor通过快速的共享内存交换数据，而其他actor对则通过全局内存进行通信。对于跨多个设备的通信，FlashMoE使用设备发起的通信，利用单边PGAS（分区全局地址空间）编程模型【[30] Katherine Yelick, 等. Productivity and performance using partitioned global address space languages. PASCO 2007.】。然而，在PGAS中实现可扩展且正确的单边内存访问而不进行昂贵的同步是一个众所周知的挑战【[1] DeepSeek-AI. Deepseek-v3 technical report, 2025.】, 【[31] Size Zheng, 等. 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是专家并行的世界大小，R标识通信轮次（即两轮，一轮用于令牌分派，一轮用于合并），B是暂存缓冲区的数量，E是本地专家的数量，C是上调后的专家容量（§3.2.1），H是令牌嵌入维度。我们实现非阻塞通信的核心思想是时间缓冲。具体来说，我们为底层令牌矩阵超额配置了至少$2 \cdot r$倍的内存，其中r是依赖图中的通信轮次数，因子2解释为每轮通信中用于传入和传出数据的单独缓冲区。对于MoE模型，我们有$2 \cdot r = 4$。这种内存使用的适度增加消除了单边数据传输期间的同步需求。图7b说明了这种对称张量布局内的单元如何被索引和用于直接内存访问（DMA）和远程DMA（RDMA）操作。正如定理3.1所强调的，这种在$L$上的索引方案是允许完全非阻塞访问从而避免同步的底层机制，因为所有访问都是无写冲突的。证明见§C。
图 7: 对称张量布局
定理 3.1。对称张量布局 L 是写-写无冲突的。
构建L。为了构建$L$，我们从原始的令牌缓冲区$T \in R^{S \times H}$开始，其中S是序列长度，H是令牌嵌入维度。我们首先将序列维度S重组为三个子维度，分别代表专家容量（C）、本地专家槽（E）和专家并行世界大小（W），使得：
$C \cdot E \cdot W = C \cdot E' = S'$, where $S' \geq S$ and $E' \geq E_W$
在专家均匀分布的典型情况下（如图7a所示），我们有 $S' = S$ 和 $E' = E_W$，其中$E_W$是模型中的专家总数。因此，令牌缓冲区的大小为 $Size(T) = S' \cdot H$。在图7a中，每个标记为$E_i$（其中$i \in {0, . . . , 3}$）的单元是一个大小为(C, H)的矩阵。我们扩展了先前的工作【[29] Dmitry Lepikhin, 等. Gshard: Scaling giant models with conditional computation and automatic sharding. ICLR 2021.】, 【[12] Shulai Zhang, 等. 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)$，对于流行模型的推理，这贡献的内存开销≤内存容量的2%。我们在§D中提供了详细的分解。
3.2.1 用于有效载荷效率的原地填充
原地填充。由于MoE分派中令牌的动态和不均匀分布【[32] Quentin Anthony, 等. Blackmamba: Mixture of experts for state-space models, 2024.】，GPU通常接收到的令牌数量少于其预定义的专家容量。当前的MoE框架【[11] Samyam Rajbhandari, 等. DeepSpeed-MoE: Advancing mixtureof-experts inference and training to power next-generation AI scale. ICML 2022.】通常在计算前用空令牌填充这些缓冲区，这不必要地增加了通信载荷并降低了性能。相比之下，我们提出了原地填充（in-place padding），直接在本地对称张量缓冲区内执行填充，从而消除了多余的网络通信。
容量对齐。正如我们在图7a中作为参考所示，每个单元$E_i$的大小根据专家容量C确定。我们进一步对齐此容量，以确保其可被分块大小$b_M = 128$整除，从而保证处理器线程在消费远程令牌时能够进行安全和对齐的内存读取。这种原地填充策略略微增加了$L$的内存占用，如下所述：
A4 实验环境
硬件配置:
- GPU: 8块 NVIDIA H100 80G GPU，通过NVLink互连。
- CPU/内存: 20个vCPU，125 GB RAM。
软件配置:
* 框架/库: PyTorch 2.6.0, CUDA 12.8, NVSHMEM。
* 操作系统: Ubuntu 22.04。
模型配置:
* 架构: MoE Transformer模型。
* 参数: 16个注意力头，嵌入维度2048，FFN中间层大小2048。
* 并行策略: 分布式数据并行（DDP）和专家并行（Expert Parallelism）。
* 路由策略: Top-2路由，容量因子为1.0。
实验设置:
* 评估范围: 仅评估单个MoE层的前向传播。
* 测量方法: 运行32次预热后，取32次运行的平均时间。
* 精度: FlashMoE使用FP32精度，而所有基线系统使用FP16。这一差异对FlashMoE不利，因为其通信量和计算工作量是基线的两倍。
对比基线:
1. Comet: 使用cudaMemcpyPeerAsync进行通信。
2. FasterMoE: 使用NCCL进行通信。
3. Megatron-CUTLASS: Megatron-LM的一个版本。
4. Megatron-TE: 使用NVIDIA Transformer Engine的Megatron-LM。
A4 实验结果
前向传播延迟（§ 4.1）
* 实验内容: 在4 GPU和8 GPU设置下，测量不同序列长度下的前向传播延迟。
* 实验结果: FlashMoE在所有配置下都显著优于基线，尤其是在长序列上。在4个H100上，处理16K令牌时，速度比Megatron-TE快4.6倍。在8个H100上，速度比性能急剧下降的基线快6.4倍。
* 分析结论: FlashMoE通过其融合内核和高效通信设计，能够有效控制随着序列长度和GPU数量增加而带来的通信开销。
图 8: 随着每个GPU的令牌数量增加，前向传播延迟的变化。
GPU利用率（§ 4.2）
* 实验内容: 测量前向传播期间的流式多处理器（SM）利用率。
* 实验结果: FlashMoE实现了93.17%的平均SM利用率，比FasterMoE（9.67%）高9倍以上，比其他基线也高出2.2倍至6.8倍。
* 分析结论: 全融合的内核架构和计算与通信任务的细粒度流水线消除了内核启动和同步造成的空闲间隙，使得SM在整个执行过程中保持繁忙。
图 9: SM利用率，定义为SM中至少有一个warp在执行的周期数与总周期数的比率[17]。数值表示100次迭代的平均SM利用率。
吞吐量（§ 4.3）
* 实验内容: 测量在扩展GPU数量时的系统吞吐量（MTokens/s）。
* 实验结果: FlashMoE的吞吐量随GPU数量线性扩展，在8个GPU上达到17.7 MTokens/s，比FasterMoE高5.7倍以上。
* 分析结论: 尽管使用FP32（而基线使用FP16），FlashMoE依然取得了高吞吐量，表明其性能提升源于最大化硬件利用率和消除主机驱动的低效，而非依赖于较低的精度。
图 10: 扩展GPU数量时的吞吐量，计算公式为 T × NG。
重叠效率（§ 4.4）
* 实验内容: 通过测量弱扩展效率来评估通信和计算的重叠程度。
* 实验结果: 随着GPU数量增加，基线方法的重叠效率显著下降（在≥4个GPU时低于50%）。FlashMoE在4和8个GPU上分别展现出高达3.88倍和4倍的效率提升，且延迟增长稳定。
* 分析结论: FlashMoE基于actor的设计和异步数据移动实现了近乎理想的计算-通信重叠。
图 11: 弱扩展效率。我们将重叠效率Oe定义为Oe = T(2)/T(NG)，其中T(NG)是在NG个GPU上的延迟，T(2)是在2个GPU上的延迟。
专家数量扩展性（§ 4.5）
* 实验内容: 在固定序列长度（16K）下，评估随着专家数量（从8到128）增加时的前向传播延迟。
* 实验结果: FlashMoE保持了低且一致的延迟。相比之下，基线方法的延迟随专家数量增加呈超线性增长。在128个专家时，FlashMoE在4个H100上比基线快4倍，在8个H100上快6.6倍。
* 分析结论: FlashMoE的有效载荷高效通信和调度器驱动的内核内分派使其能够在不产生其他系统所见的通信和协调惩罚的情况下，持续支持专家并行，这证明了其对超稀疏MoE配置的可扩展性。
图 12: 随着专家数量增加，前向传播延迟的变化。
A7 补充细节
工程复杂性。完全融合的持久化内核需要深厚的GPU和分布式系统专业知识；未来的工作可能会研究编译器/DSL抽象来降低这一门槛。
FP16效率低下。我们的FP16路径由于调整不足而并非最优（§H）。我们预计通过自动调优的GEMM算子如cuBLASDx【[35] NVIDIA. cublasdx, 2025.】或CUTLASS构建器来解决这一问题。
训练支持。这项工作主要针对推理；要支持训练，需要将反向计算和梯度通信与新的记账和任务描述符融合在一起。
A5 结论
本文介绍了FlashMoE，这是首个将整个分布式MoE算子融合到单个持久化GPU内核中的工作。该内核通过actor风格的并发、warp专门化和异步(R)DMA技术，统一了计算、通信和调度。我们解决了先前系统中两个主要的瓶颈——由CPU管理的同步通信和碎片化的多内核执行。实验证明，FlashMoE在分布式MoE任务中实现了高达6倍的加速、9倍的GPU利用率提升和5.7倍的吞吐量提升。展望未来，我们预见到一个从CPU编排转向完全自主的、GPU原生流水线的转变——将这种融合方法扩展到训练及更广泛的领域。
A6 附录
A 补充动机
在图14中，我们展示了在32个A100和8个V100 GPU上分布式训练一个1.3B GPT-3 MoE模型时，AlltoAll内核运行时间的经验累积分布和原始分布。我们用这个结果来论证掉队者效应的严重性和普遍性。在图14b中，我们观察到与平均实际内核时间相比，P95通信性能下降了1.32倍。这种性能下降相对温和，因为底层硬件是一台针对“软件抖动”【[36] NERSC. Network - NERSC Documentation. 2025.】进行过良好调优的超级计算机。然而，在单节点虚拟机（VM）中，我们观察到更严重的p95性能损失，达到11倍。与先前的高性能计算工作【[37] C. Bell, 等. Optimizing bandwidth limited problems using one-sided communication and overlap. IPDPS 2006.】, 【[38] Yuxin Chen, 等. Atos: A task-parallel gpu scheduler for graph analytics. ICPP 2023.】一致，我们认为，消除这种同步集合通信中固有的障碍将允许GPU将观察到的空闲时间重新用于下游计算，如图13所示。
图 13: 重叠调度（底部）显示了顺序调度（顶部）中的空闲时间如何被重新用于计算。FlashMoE实现了重叠调度。
表 2: 同步All-to-All通信中的掉队者延迟。我们捕获了在许多步骤中由掉队者引起的延迟分布。设实际时间ta为所有GPU中最快的内核执行时间，总时间t为记录的最大步骤时间。我们将延迟定义为t和ta之间的最大差值。注意延迟即是空闲时间。对于1x8 V100，我们分析了1750个步骤，对于8x4 A100，我们分析了600个步骤。原始分布见图14。
图 14: 同步AlltoAll的掉队者效应。M × N A100或V100表示M个节点内每个节点有N个GPU。每个GPU在每个AlltoAll步骤中与所有其他GPU通信。我们捕获了在许多步骤中由掉队者引起的延迟分布。实际时间ta表示所有GPU中最快的内核执行时间，相反，总时间t是记录的最大步骤时间，而延迟是t和ta之间的最大差值。注意延迟即是空闲时间。
B 相关工作
计算-通信重叠与内核融合。为了减少分布式DNN训练中同步带来的通信开销，许多研究工作都致力于增加计算和通信的重叠。对于没有MoE层的通用Transformer模型，许多工作【[39-47]】为划分和调度计算与通信操作提供了见解和技术，旨在实现更细粒度的重叠。为了应对MoE训练中AlltoAll通信和专家并行带来的挑战，Tutel【[48] Changho Hwang, 等. Tutel: Adaptive mixture-of-experts at scale. MLSys 2023.】和FasterMoE【[14] Jiaao He, 等. Fastermoe: modeling and optimizing training of large-scale dynamic pre-trained models. PPoPP 2022.】将AlltoAll与专家计算重叠。Lancet【[49] Chenyu Jiang, 等. Lancet: Accelerating mixture-of-experts training via whole graph computation-communication overlapping. MLSys 2024.】额外地使前向传播中的非MoE计算和后向传播中的权重梯度计算能够与AlltoAll重叠。尽管实现了重叠，但这些方法的性能在实践中受限于带屏障的阻塞式同步集合通信。相比之下，FlashMoE通过异步、设备发起的数据传输与分块计算在单个内核内完全重叠，从根本上消除了这些低效。FlashMoE进一步与COMET【[12] Shulai Zhang, 等. Comet: Finegrained computation-communication overlapping for mixture-of-experts. MLSys ’25.】和DeepEP【[1] DeepSeek-AI. Deepseek-v3 technical report, 2025.】等SOTA工作区分开来，后者也使用这种形式的内核发起通信，但粒度较粗且没有完全的内核融合。
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$。此规则描述了设备内的暂存写操作。
我们重申定理3.1并概述其证明如下。
定理 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_2 = (p_2, r_2, b_2, e_2, c_2)$
注意，进程内的写操作总是有不同的$c_j$坐标，其中$j \in \{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 内存开销
我们测量了FlashMoE的对称张量L和运行时记账状态所需的GPU内存。内存开销主要取决于分块大小、专家容量（EC）和专家数量（E）。表3总结了在推理过程中，近期MoE模型【[50-55]】的内存开销，显示FlashMoE保持了适度且可预测的内存占用。特别是，对称张量（ST）相对于总推理内存需求，最多只增加了2.15%的额外内存。
表 3: FlashMoE的内存开销（分块大小 bM = 128, Size(T) = Tokens × 4KB）。
E 任务实现
图 15: Task结构体。TaskType ∈ {GEMM0, GEMM1, Combine}
F Actors
算法2：处理器Actor：由一个块执行
1 begin
2 tQ ← GetTQ()
3 signal ← 0
4 // shared memory variables
5 task ← {}
6 interrupt ← False
7 complete ← False
8 while interrupt == False do
9 if warpId == 0 then
10 if threadId == 0 then
11 awaitTaskFromScheduler(interrupt, signal)
12 FencedNotifyRQ(ready)
13 end if
14 syncwarp()
15 warpReadTQ(tQ, signal, task)
16 end if
17 syncthreads()
18 if interrupt == False then
19 switch task.Type do
20 case GEMM0 do
21 // fused GEMM, epilogue and async tile staging
22 fGET(GEMM0, task)
23 if threadId == 0 then
24 complete ← NotifyTileCompletion()
25 end if
26 syncthreads()
27 if complete == True then
28 NotifySchedulerNextGEMM(tQ)
29 end if
30 end case
31 case GEMM1 do
32 // fused GEMM, epilogue and async tile transfer
33 fGET(GEMM1, task)
34 end case
35 case Combine do
36 combine(task)
37 end case
38 end switch
39 end if
40 end while
41 end算法3：调度器Actor：由一个warp执行
1 begin
2 scheduled ← 0
3 tT B ← 0
4 tqState ← {}
5 pT DB ← GetProcessorDoorbell()
6 sT DB ← GetSubscriberDoorbell()
7 taskBound ← GetTaskBound()
8 tT B ← AtomicLoad(taskBound)
9 // circular buffer ready queue
10 rQ ← {}
11 // Populate ready queue with Processor ids
12 PopulateRQ(rQ)
13 while scheduled < tT B do
14 lt ← 0
15 do in parallel
16 Sweep doorbells and populate observed task counts into tqState
17 Aggregate locally observed task counts into lt
18 end
19 qS, taskT ally ← 0
20 // qS is the inclusive output
21 WarpInclusiveSum(lt, qS, tasktally)
22 while tasktally > 0 do
23 Repopulate rQ with ready processor ids
24 do in parallel
25 Starting at rQ[qS], signal processors about task indices from tqState
26 end
27 end while
28 if threadId == 0 then
29 tT B ← AtomicLoad(taskBound)
30 end if
31 tT B ← WarpBroadcast(tT B)
32 end while
33 InterruptSubscribers()
34 InterruptProcessors()
35 end算法4：订阅者Actor：由一个或多个warp执行
G 实现
表 4: FlashMoE的实现指标。
H FP16内存吞吐量
图 16: 这里我们报告了FlashMoE的FP16（上）和FP32（下）变体的总A100内存吞吐量。值得注意的是，在相同工作负载下，FP16实现发出的共享内存指令大约是其FP32对应版本的2倍。我们将这种低效率归因于FlashMoE在处理半精度数据时共享内存布局不佳。虽然这个瓶颈可以通过改进布局策略来解决，但我们将其留给未来的工作。
💬 评论讨论
欢迎在这里分享您的想法和见解！