FlashMoE: Fast Distributed MoE in a Single Kernel
Authors: Osayamen Jonathan Aimuyo
Byungsoo Oh
Rachee Singh
Venue: NeurIPS 2025 poster ()
Keywords: Distributed Machine Learning
Kernel Fusion
Mixture-of-Experts
Primary area: infrastructure
TL;DR: FlashMoE is the first work to completely fuse the Distributed Mixture-of-Experts operator into a single persistent GPU kernel with elegant actor-style concurrency and high-throughput tile parallelism, achieving up to 9× higher GPU utilization.

Abstract: The computational sparsity of Mixture-of-Experts (MoE) models enables sub-linear growth in compute cost as model size increases, thus offering a scalable path to training massive neural networks. However, existing implementations suffer from low GPU utilization, significant latency overhead, and a fundamental inability to leverage task locality, primarily due to CPU-managed scheduling, host-initiated communication, and frequent kernel launches. To overcome these limitations, we develop FlashMoE, a fully GPU-resident MoE operator that fuses expert computation and inter-GPU communication into a single persistent GPU kernel. FlashMoE enables fine-grained pipelining of dispatch, compute, and combine phases, eliminating launch overheads and reducing idle gaps. Unlike existing work, FlashMoE obviates bulk-synchronous collectives for one-sided, device-initiated, inter-GPU (R)DMA transfers, thus unlocking payload efficiency, where we eliminate bloated or redundant network payloads in sparsely activated layers. When evaluated on an 8-H100 GPU node with MoE models having up to 128 experts and 16K token sequences, FlashMoE achieves up to 9× higher GPU utilization, 6× lower latency, 5.7× higher throughput, and 4× better overlap efficiency compared to state-of-the-art baselines—despite using FP32 while baselines use FP16. FlashMoE shows that principled GPU kernel-hardware co-design is key to unlocking the performance ceiling of large-scale distributed ML. We provide code at https://github.com/osayamenja/FlashMoE.

FlashMoE: Fast Distributed MoE in a Single Kernel
Osayamen Jonathan Aimuyo Comell University oja7@cormell.cm edu
Byungsoo Oh Comell University bo239@cornell.cm
Rachee Singh Cornell University rs2293@cormell1.omm edu
Abstract
The computational sparsity of Mixture-of-Experts (MoE) models enables sub- linear growth in compute cost model size increases, thus offering scalable path to training massive neural networks. However, existing implementations suffer from low GPU utilization, significant latency overhead, and fundamen- tal inability to leverage task locality, primarily due to CPU-managed scheduling, host-initiated communication, and frequent kernel launches. To overcome these limitations, we develop FlashMoE, a fully GPU-resident MoE operator that fuses expert computation and inter-GPU communication into a single persistent GPU ker- nel. FlashMoE enables fine- ane-ggrined pipelining of dispatch, compute, and combine phases, eliminating launch overheads and reducing idle gaps. Unlike existing work, FlashMoE obviates bulk-synchronous collectives for one-sided, device-initiated, inter-GPU (R)DMA transfers, thus unlocking payload efficiency, where we elim- inate bloated or redundant network payloads in sparsely activated layers. When evaluated on an 8-H100 GPU node with MoE models having up to 128 experts and 16K token sequences, FlashMoE achieves up to 9x higher GPU utilization, 6x lower latency, 5.7 x higher throughput, and better overlap efficiency com- pared to state-of-the-arr baselines- FlashMoE shows that principled GPU kernel-hardware co-design is key to unlock- ing the performance ceiling of large-scale distributed ML. We provide code at https: //githgggtttt/tt com http//aaaaaaaaaaaalllaalaaa
using FP32 while baselines use FP16.
Ferword Lobenyy [t 22 22 11100
bocce
106-
95.17%
Forpard Pass E- 54
2 A1.00
beteer
FestsMoc
Mogatresit
Auarage 601
40-
1AR
Faa-Mat
Come:
Taste-Net
Megabonn
Deup2?
(a) GPU SM Utilization
(b) Scaling Tokens
Terword Labenyy T- 2K
better
Torwars Uibenry T 366 21 inate
batter
Megattos
Number of
(c) Weak Scaling across GPUs
(d) Across Experts
Figure I: FlashMoB performance.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Introduction
State-of-the-ard large language models
(LLMs) [15 have adopted the Mixture-of- Experts (MoB) architecture for its computational efficiency and strong performance across range of tasks. The traditional Transformer block consists of self-attention module followed by a dense feed-forward network (FFN) [6] In contrast, MoE architectures replace this single FFN (Figure 2[)) with many identically sized FFNs, known as experts (Figure 2b)). A trainable neural network, known as gate function, sparsely activates these experts by dynamically routing input tokens to the experts selected at runtime. This increase in model parameters due to more FFNs improves model quality without the corresponding increase in computational cost.
T
Tg
Attention
Attention
Attention
Attention
(Combine)
(a) Transformer
(b) MoE
(c) Distributed MoE
Figure Transformer blocks (a) without MoE, (b) with MoE, and (c) with distributed MoE and expert parallelism. т, E, and represent input tokens, experts, and output activations, respectively.
Communication overheads in MoE. As MoE
model sizes grow, GPU memory constraints prevent hosting all experts on single device. The standard practice is distribute experts across multiple GPUs using expert parallelism (EP), which requires token routing via many-to-many communication primitives like AlltoAll [ ure ()). Another round of AlloAll communication restores the permuted tokens processed by experts to their original order in the sequence. AllioAll communication is challenging to optimize on GPU networks and is highly sensitive to straggler delays - a phenomenon where single straggler GPU delays all others from making progress [8] These communication operations account for 68% the total runtime [9 [0], causing GPUs to remain idle (Figure top left).
4
3 /
(Fig-
Kernel launch overheads in MoE. To mitigate these communication bottlenecks, recent work pipelines computation with communication kernels (Figure 3 left middle). However, the effectiveness of these solutions is limited by the overhead of launching many kernels from the CPU. Specifically, existing implementations i launch a large number of kernels per a single layer pass (see Table Frequent kernel launches negatively affect performance by: (1) creating non-deterministic kernel start times across GPUs, exacerbating straggler issues; (2) introducing unnecessary synchronization points, causing GPUs to wait peers the CPU before proceeding; and (3) incurring repeated global memory round trips at kernel boundaries. Although CUDA graphs [I5] can partially mitigate the first issue in static workloads, they are incompatible with MoE's s dynamic expert routing patterns. Addressing the remaining issues requires novel solutions, which we provide in this work through complete kernel fusion and asynchronous sevice-initiated communication.
Overap
52
(Dkgatch)
Stale
Overop
A2A
CPU
Espert
GPU
Attii
Gate
Expert
Expert
AJA [Dispatch)
A7A (Combine)
A2A
Single Kamnl
CPU
Ours: FlashMoE
Baseline
GPU
Art
Gate
GPU karnel execution
MoE
Remote communication
Figure 3: Comparing FlashMoE with state-of of-the-art techniques that either do not overlap commu-
nication and computation (left, top) or do some overlap (left, middle). FlashMoE a persistent kernel that fuses all computation and communication of the MoE operator (left, bottom). FlashMoF implements device-initiated computation (gate, expert FFN, scale) and communication tasks (right).
Our Contributions: distributed MoE in single kernel. To overcome these fundamental inefficiencies in state-of-the-ar MoE models, we develop FlashMoE megakernel that integrates

---

all MoE computation and communication tasks into a single persistent GPU kernel i.e., kernel that remains active for the entirety of the MoE operator [Figure[]bottom left). Instead of multiple kernel launches coordinated by the CPU, FlashMoF requires launching only one kernel, significantly reducing the involvement of the CPU. Within the fused kernel, FlashMoF implements reactive programming model to achieve fine-grained parallelism and loosely coupled, non-blocking execution among tens of thousands GPU threads.
In-kernel Block scheduling and Tile paral- lelism. FlashMoE implements tile-level paral- lelism, meaning it partitions input token matri- ces into smaller, independent units called tiles, which are processed by blocks but managed (scheduled and constructed) by warps. We spe- cialize every thread block, except one, as pro- cessors perform compute. In addition, we designate dedicated Operating System (OS) block (4 warps) to perform administrative tasks of (1) scheduling computational work to pro- cessors (scheduler), and (2) decoding computa- tional tasks from messages received from other GPUs (subscriber). This design allows Flash- MoE to dynamically assign tasks to GPU blocks based on readiness, ensuring that no GPU SM remains idle throughout the lifetime of the MoE operator. FlashMoE selects tile dimensions to maximize GPU arithmetic intensity while benefitting from high-degree of parallelism.
MoE Implementation
GPU Ops
FlashMoB (this work)
1
COMET [12]
33
Megatron-LM CUTLASS 13
16
85
Megatron-LM TE
L3
16
261
Megatron-LMM DeepEP
432
DeepSpeedMoF 11
550
Table 1:
We report number of GPU operations launched by MoE implementations by profiling with Nsight Systems [7]. We count operations in a single MoE layer (Gate - Dispatch - Expert + Combine) on 2 A100 GPUs, where each GPU has 32 experts. FlashMoB is the first to fully fuse the distributed MoE layer into a single GPU kernel.
Asynchronous and payload-efficiens communication. By redesigning the MoE operator from the ground up, FlashMoE resolves fundamental inefficiencies inherent in the conventional MoE execution pipeline. One notable inefficiency is token padding during communication. To simplify programming complexity and due to symmetry constraints of collective communication APIs, existing implementations have to zero-pad token payloads to match predefined buffer sizes. This occurs when tokens asymmetrically routed to experts, resulting in GPUs receiving much less than the expected capacity. However, these null payloads waste communication bandwidth, bloat data transfer latency and may lead unnecessary computations on null matrices in some implementations. FlashMoE introduces payload-efficien communication by sending non-padded tokens only to GPUs with actively selected experts, conserving both communication and computational resources.
Technical challenges. Realizing the single-kernel design of FlashMoE required solving several technical challenges to achieve high performance: (1) lightweight computational dependency man- agement; (2) navigating optimal SM occupancy configurations: (3) implementing in-device BLAS operations; (4) minimizing inter- and intra-device synchronization overheads; (5) implementing transfer-awareness by leveraging DMA over Unified Virtual Addressing (UVA) when available. In addressing these challenges, FlashMoE's design presents radical departure from traditional synchronous AlltoAl/ collectives, where GPUs exhibit significant idle time during layer execution. For innititted communication, FlashMoE uses NVSHMEM (18] establish global address space across all GPUs for Direct Memory Access (DMA) communication. For in-device BLAS, FlashMoE develops custom nigh-performance GEMM operations via CUTLASS 19
Results. Our evaluations show that FlashMoE achieves 6x latency speedup. 9x higher GPU utilization, 4x better weak scaling efficiency and 5.7 increased throughput compared to state-of- the-art implementations. We project these performance gains becoming even better in molti-node scenarios, where inter-node communication occurs using lower bandwidth inter-node links (e.g.. RDMA, Infiniband).
2 Motivation
Synchronous Communication. AlloAAl or AllGather communication currently used in MoE frameworks is a synchronous collective operation, whose completion requires the participation of all involved GPUs. Here, disparities in processing speeds or kernel scheduling among GPUs induce
3

---

Fass K- A1.000
MoE Iteration (6.557 ms)
forward
FlashMoE
Aowerae
42.31%
MoË Iteration (19.561 ms]
9.57%
13.55%
Commet Mel
Megatron
Deept?
DeepEP
(a) GPU SM Utilization across baselines
(b) Kemel Launch overhead (CUDA API row)
Figure 4: 4a shows GPU utilization averaged across 100 MoE forward passes on 2 A110ss with 300 GB/s unidirectional bandwidth, where we observe up to 90% idle time, due to kernel launch gaps and non-overlapping communication.
straggler effect detrimental (Figure 13 to (1) the collective operation's 's performance and (2) E2E performance, as stalled GPUs cannot proceed to downstream dependent or independent tasks until the collective terminates. We expound on this problem in $
Kernel Launch Overhead. We compare the kernel launch overheads between FlashMoE and existing baselines. Table | shows the number of kernel launches during a single forward pass: FlashMoE launches exactly one persistent kernel, while the baselines launch up to 550 short-lived kernels to perform the same computation. Figure E provides a visual comparison using CUDA API traces captured by NSight Systems, illustrating the difference between FlashMoB and DeepEP. DeepEP exhibits many small CUDA API calls, with frequent stalls between individual operators, leading to increased GPU idle time FFag In contrast, FlashMoF maintains high GPU utilization by avoiding launch overhead and synchronization gaps achieving 93.17% GPU utilization compared to 14% for DeepEP. See $
for experimental results and SB for a discussion of related work.
Fused MoE Kernel Design
Modern distributed MoE systems suf- fer from two limitations: (1) frequent many-to-many (AlltoAll or AllGather) collectives on the critical path, and (2) significant overhead from repeated kernel launches. We address these in FlashMoE, fully fused MoE oper- ator implemented as a single persis- tent GPU kernel. Unlike previous ap- proaches [12] T IT 13 20 I0, 21- 22). FlashMoE is the first solution to implement a completely fused Dis- tributed MoE kernel, eliminating ker- nel launch overhead entirely by requir- ing only single kernel launch (see Table D..
Block
Bleck
Block
Backn-1
FusedGate
Remote Packets
Encodle Dispatch
We' war WND
OS
...
Scheduter
Subscriber
Devodee Tavis
Taw
Samd Tie to only
Actor-based model. The design of FlashMoH is based on the actor model of concurrent computation (26.28)) We implement this model by specializ- ing GPU thread blocks and warps into three distinct actor roles: (1) Processor ($F.)). p. (2) Subscriber
Figure 5: FlashMoE Fused Kernel
(F..), and (3) Scheduler F.2). The Processor performs compute (GEMMs and element- operations) and tile communication. We use CUTLASS [19 the underlying infrastructure for high-peeffrrance BLAS routines and NVSHMEM for kermel-initiated communication [18]. The Subscriber and Scheduler perform administrative functions. Specifically, the Scheduler assigns computational tasks to available thread blocks. Our key innovation is making the Scheduler both multithreaded, enabling high scheduling throughput, and work-conservingg ensuring consistently high GPU SM utilization. On the other hand, the Subscriber decodes tile packets from peer GPUs to task

---

Algorithm 1: FlashMoE Distributed MoE Fused Kernel
Input:
A, RR R
begin
Go.Go + (ate) if blockld 1 then Dispatch(Ta,.4 A) processor::star((
else
warppD == then scheduler:startt) else
subseriber:sttar((C,, Go. O)
11
end
12
end if
13
end
Dispatch packets
Notify Tasks
Schediale Task
Notify Tasks
Schodule Task
Sand Tile
Notify Tasks
Schociule
Task
D
S
S
P
S
pi
S
Sh
p
GEMM.
GEMM,
Combine
Figure 6: DMoE Functional Dependencies Expressed Chain of Actor Interactions. We denote S.. Sn, and P as the Subseriber, Scheduler and Processor actors, respectively. For any actor a E {Sp: Sh, P), a² identifies an actor on GPU i. We define D as the operator, where GPU dispatches packets of tiles GPU i. This diagram expresses task dependencies the granularity of tile, namely GEMM, GEMM, combine and communication produce an output tile. Notifications occur as signals propagated through shared memory (subscriber scheduler) or global memory (scheduler + processor or inter-GPU communication). Note one-sided inter-GPU transfers (packet or single tile) are coupled with signal to notify S on the receiving GPU j of the message's 's delivery.
descriptors ($3.1) Of the N thread blocks on GPU, we specialize N - 1 to adopt the Processor role. We specialize the last block as the Operating System (OS). Within this block, we specialize three warps for the Subscriber role and one warp for the Scheduler role. This split of thread blocks across actors is intentional: our goal is to use few resources for administrative tasks while reserving bulk of the resources for performing MoE computation tasks. Figure [igureeSsummarizer the FlashMoE architecture and its constituent actors, while Algorithm I gives a very close translation of the system in code. Note that A Sx H is the input token matrix; O Sx the output matrix; and X R Ex H x is a 3-D tensor of expert weights, where E denotes the number of local experts for the executing GPU, H is the embedding dimension, D is the FFN intermediate dimension and S is the sequence length. To E (N IR.) FXC"" is routing table data structure, where T' (e, (i, w)) indicates that token at slot с dispatches to expert w: is the combine weight (Equation 2 and C' is expert capacity. The tuple structure of To an implementation detail. Go S captures the affinity scores produced by the gate (Equation 3}. Inter-actor interactions in FlashMoE. FlashMoE decomposes MoE computation and communication the granularity tile, a statically sized partition of a tensor, to achieve parallel execution and efficient overlap of tasks. Each tile maps to discrete unit of work encapsulated by a task descriptor. The Subscriber decodes these task descriptors from the remote tile packets it receives. Concurrently, the Scheduler receives notifications about available tasks and dispatches them for execution to Processor actors that perform computations defined by these tasks, namely the feed-forward network (FFN) and expert-combine operations. Figure 6 show the chain of actor interactions, demonstrating how FlashMoE enforces DMoE functional dependencies.
Determining tile dimensions in FlashMoE. Selecting appropriate tile dimensions in FlashMoE crucial to ensure efficient GPU utilization. An undersized tile underutilizes the GPU, while excessively large tiles create register pressure, causing memory. After careful parameter sweeps, we choose tile dimensions of (128, 64). Our key insights are: increasing tile width significantly raises the register usage per thread, potentially triggering costly spills; increasing tile height without adjusting thread count increases workload per thread, harming performance. Raising the thread count per block beyond our fixed value of 128 threads reduces the
register spills to local
5

---

number of concurrent blocks, negatively affecting SM occupancy. Larger thread-bloc' sizes also increase overhead from intra-block synchronization syncthreads() barriers), further degrading performance. Thus, our chosen tile dimensions balance register usage, shared-memory constraints, and GPU occupancy to deliver optimal performance.
3.1 Task Abstraction for Computation
Computational operators. The FFN operator is a standard position-wise feed-forward network widely used in Transformer architectures [6). composed of two linear transformations separated by nonlinear activation Φ (e.g., GELU or ReLU):
(FN(2) W22 <(>W] b1) b2
(I)
Here, W1 and W22 represent learnable weight matrices, and b and by biases. The expert-combine
operation, used in architectures like GShard [29] and DeepSeek [I], merges outputs from multiple
experts by computing a weighted combination based on their affinity scores:
C
9b.e
(2)
==1
k
DLe C,
(3)
In these equations, i 0, represents an input token index, e= E, identifies the k-th expert selected for token i, and Si,e is the affinity score indicating how relevant expert is for token i.
Unified task abstraction. We unify the FFN and combine operations under common abstraction called task. Tasks provide uniform interface for communicating tile-level work among Subscribers, Schedulers, and Processors. Formally, a task descriptor
is defined as a tuple:
t= (M, Φ)
where M is set of metadata (e.g., device ID, tile index), is binary tensor operation (specifically, matrix multiplication or Hadamard product (), and an element-wise activation function (c.g.. ReLU or identity).
We define task t operating input tensors A, B, D, producing output tensor C, as follows: B, = C' + & (A B D)
(4)
The operator * (instantiated from +) may behave differently depending on the task metadata M, and the result of A +, B is accumulated into D. We provide example of task metadata in $
In practice, we implement each task defined by Equation as single fused ____ddddvee_____ decorated function which the Processor (Algorithm 2 invokes at runtime. Fusion for t entails applying and the succeeding addition operation to registers storing the results of the binary operator *;. To illustrate its flexibility, we show how the FFN and expert-combine operations can be expressed using this task framework. Note that we omit the matrix multiplication symbol (-) for simplicity. Also, di can be any activation function, while Φ2 is the identity function. The FFN is expressed as:
= (M, ; Φ1), to = (M,+.02) Fi, (A, D.) = C1 + Φ1 (ABı + -D), Fr, (C,, B2, C2, D2) = C2 d2 (C, D2). Whereas, the expert-combine operation is formalized as:
La =
3.2 Symmetric Tensor Layout for Inter-GPU Communication
Within a single GPU device, the actors in FlashMoF communicate through the GPU's memory subsystem. Specifically, the Scheduler and Subscriber actors exchange data via fast shared memory,

---

4.) Po receives
PiR,BIEs
Parallel Group
12
Dispatch
E
E
E2
E5
to & on
2. receives
3
Po
View
Pj View
R
R1
Rg
R1
Bg
B,
Bg
B
B
P. mends
alages
Dispatth to on
Eg
Ep
Eg
Eo
Eo
E2
E2
Eo
PIRiB
Ro
Po
Poo
4. processes
E
E;
E
E
E3
E
Start
Po Pa
Ep
Eo
E2
E2
E2
2.
P.
P1
Eg
Es
Ea
remote
(b) State machine for DMA (top) and RDMA (bottom) communication.
(a) Layout across 2 Expert-porallel Processes.
Figure 7: Symmetric Tensor Layout
while other actor pairs communicate through global memory. For communication across multiple devices, FlashMoE uses device- communicationn leveraging the one-sided PGAS (Partitioned Global Address Space) programming model [30]. However, achieving scalable and correct one-sided memory accesses in PGAS without costly synchronization is known challenge [1] 31). We address this challenge with a provably correct and scalable solution: a symmetric tensor layout T., supporting fully non-blocking memory accesses. We define I. as:
L LER PxRXxxxXXx
where: P is the expert parallel world size, identifies communication rounds (i.e.. two rounds. one for token dispatch and one for combine), B is number of staging buffers, E is the number of local experts, C is the upscaled expert capacity ($3.2.1) and H is the token embedding dimension. Our core insight to enable non-blocking communication is temporal buffering. Specifically, we overprovision memory for the underlying token matrix by at least 2- times, where is the number of communication rounds in the dependency graph, and the factor of 2 accounts for separate buffers for incoming and outgoing data within each communication round. For MoE models, we have 2 This modest increase in memory usage eliminates the need for synchronization during one-sided data transfers. Figure 7bb illustrates how cells within this symmetric tensor layout are indexed and used for Direct Memory Access (DMA) and Remote DMA (RDMA) operations. As Theorem 3.1 reinforces, this indexing scheme over is the underlying mechanism that allows for fully non-blocking accesses eliding synchronization because all accesses write conflict-free. See$ C for the proof.
Theorem 3.1. The symmetric tensor layout write-write conflict-free.
To construct L, we start from the original token buffer T E Sx where S is the sequence length and H is the token embedding dimension. We first reorganize the sequence dimension 5 into three dub-dimensions representing the expert capacity (C), local expert slots (E), and the expert parallel world size (W), st:
C. C. C-E.W-C.E'-3
where S' S and E' > Ew
In the typical case of uniform expert distribution (illustrated in Figure 7ap. we have S' and E" Ew. where Ew is the total number of experts in the model. Thus, the size of the token buffer is Size(T) S" H. In Figure 7a each cell labeled E, (with i {O.... ...... 3}) is matrix of size (C, If). Extending prior work (29 I2T. we introduce additional temporal dimensions R (communication rounds) and (staging buffers). Each communication round has two fixed staging slots: one for outgoing tokens and another for incoming tokens. Each slot, indexed by dimension P. forms a tensor of shape H)). Therefore, the tensor size Size(L) is generally at least four times the original token buffer size, becoming exactly four times larger in the case of uniform expert distribution. Empirically, we find Size(T.) (T.) x Size(T), contributing memory overhead inference of popular models. We present thorough breakdown in
of memory capacity for

---

3.2.1 In-place Padding for Payload Efficiency
Due to the dynamic and uneven distribution of tokens in MoE dispatch 32, GPUs commonly receive fewer tokens than their predefined expert capacity. Current MoE frameworks [ typically pad these buffers with null tokens before computation, unnecessarily increasing communication payloads and degrading performance. In contrast, we propose in-place pudding, performing padding directly within the local symmetric tensor buffers and thus eliminating excess network communication.
As we show in Figure 7ajas as a reference, each cell E, is sized according to the expert capacity C. We further align this capacity to ensure divisibility by the tile block size 6M = 128, guaranteeing safe and aligned memory reads by Processor threads consuming remote tokens. This in-place padding strategy slightly increases the memory footprint of L, described below:
4 STze(T),
B bM
Size(L) ~
4 bM-E Size(T), g
otherwise
Evaluation
We implement ($G) and evaluate FlashMoE across five metrics: Forward Latency ($
4.1 GPU
Utilization ($) 4.2). Overlap Efficiency ($ 4.4), Throughput ($ 4.3) and Expert Scalability ($ 4.5
We run experiments
server with 8 NVIDIA H100 80G GPUs interconnected via NVLink, 125 GB of RAM, and 20 vCPUs. We used PyTorch 2.6.0, CUDA 12.8, and Ubuntu 22.04. All experiments use MoE transformer models configured with 16 attention heads, an embedding dimension of 2048, and an FFN intermediate size of 2048. We apply Distributed Data Parallelism (DDP) and Expert Parallelism for all experiments. We execute only the forward pass over a single MoE layer and measure the average runtime of 32 passes after 32 warmup passes. We use top-2 routing with capacity factor of 1.0. We compare FlashMoE against several state-of-the-arr MoE systems: (1) Comet [1), (2) FasterMooF 14, (3) Regatron-CUTLASS 13] and (4) Megatron-TKe Megatron-
LM with Transformer Engine 33 Comet relies on audaMemcpyPeerAsyne 34 while FasterMoE
and Megatron-LM use NCCL exclusively for communication.
Desiderata. We observe Comet exhibiting anomalously bad performance values at 8 GPUs, so we exclude their results from evaluations at 8 GPUs and only include for results at < 4 GPUs. We evaluate FlashMoE using FP32 precision whereas all baselines
FP16. We do because (1) of incomplete fp16 tuning (H and (2) no baseline supports FP32. Note, this precision discrepancy disadvantages FlashMoE by doubling its communication volume and computational workload.
4.1 Forward Latency
Forward Latency 18 = H1000]
better
Forward Latency [8 3mm H1000 better 25-3
Megatron CUTLASS
224
FlaslttsE
CUTLASS
MegatronTT
Megatron-TE
Runtime (mat
1.5
30
12-4
1.-
23.0
10
5. 250
Number lokers
Number of lokers
(a) 4 H100s
(b) 8 H100s
Figure 8: Forward Latency as the Number of Tokens per GPU increases.
We first measure the forward latency of FlashMoE across different sequence lengths on both and 8 GPU setups (Figure 8). FlashMoE consistently outperforms all baselines, with especially notable improvements at longer sequence lengths. On 4 GPUs, it achieves up to 4.6x speedup over Megatron-TEF TE at 16K tokens, and 2.6x over FasterMoF. The gains are even more pronounced at 8 GPUs where FlashMoF. maintains low latency, exhibiting up to 6.4x speedup over baselines that degrade steeply due to increasing communication costs token buffers increase proportionally.

---

4.2
GPU Utilization
To quantify GPU efficiency, we measure Stream- ing Multiprocessor (SM) utilization during the forward pass (Figure 9. FlashMoE achieves 93.17% average SM utilization, over 9x higher than FasterMoF: (9.67%), 6.8x higher than DeepEP--Meggtron--L.
Forward
E- 2 A100s
93.17%
Avenge UnI (%)
00-
60-
59. 11%
42. 3175
(13.55%) 4x higher than Megatron- TE (59.11 and 2.2x higher than Comet (42.31%). This improvement stems from our fully fused kernel architecture and fine- grained pipelining of compute and communica- tion tasks. By eliminating idle gaps due to kernel launches and enabling un-kernel task scheduling, FlashMoE ensures SMs remain busy with pro- ductive work throughout execution.
40
20
9.67%
13.55%
RastNoE
Comtee
RasterMoE
Megal
Figure 9: SM utilization, defined as the ratio of cycles in which SMs have at least one warp in flight the total number of cycles [17]. Values represent the average SM utilization over 100 iterations.
4.3 Throughput
As shown in Figure 10 FlashMoE scales linearly with GPU count, reach- ing 17.7 MTokens/s at GPUs. This over 5.7x higher than FasterMoE and 4.9x higher than Megatron-TE and AUCLLSO Notably, these results are achieved despite Flash- MoE operating entirely in FP32, while baselines use FP/6. This indicates that FlashMoE 's design eliminates throughput bottlenecks not by exploit- ing lower precision, but by maximiz- ing hardware utilization and eliminat- ing host-driven inefficiencies.
Forward Pass T
B- E=
better
20.0
17.5
FasterMoE
Megatron-CUTLASS
Throughen (MT
Megatron-TE
12.5
10.0
5.50
7.5
2.5
395
2.50
1.42
1.
of H100s
Figure 10: Throughput when scaling the number of GPUs, computed TxNe
latency"
4.4
Overlap Efficiency
Forward Labency T 8K| 2
better
Forward TT
BK| E [ T better
25
Plasietct
CUUTrron-CUTLASS
Flash FM
20-
FasterMoE
Megatron-TE
M-C
(mc)
Overfap Efiiconcy
80
12.6
40
of H100s
of H100s
(a) Latency as Number of GPUs increases.
(b) Weak scaling efficiency
Figure 11: Weak scaling efficiency. We define Overlap Efficiency O₂ be O. (N2)/T(No) where T(Na) is the latency at Na GPUs and T(2) is the latency at 2 GPUs.
We evaluate the extent which FlashMoE overlaps communication and computation by measuring weak scaling efficiency as the number of GPUs increases (Figure lbb). We note that most baselines fail to execute at single GPU, hence why we use GPUs as the reference point. We observe that Megatron-
CitiutronCUTLLASS CUTLASS and Megatron- TE degrade significantly, with overlap efficiency dropping below and 8 GPUs, respectively. Figure 1aa further illuminates this efficiency, FlashMoB shows stable forward latency growth. These results corroborate that FlashMoE's actor-based design and asynchronous data movement achieve near-ideal overlap.
50% at > 4 GPUs FlashMoB gives up to 3.88x and higher efficiency

---

4.5
Expert Scalability
Forward Lakency |6K] | 4H100s
better
Forward Lakency [7 16K K- 8
better
FasteMsE
Megatron CUTLASS
Regatron-CUTLAS
FastMMs
Megatron CUTLASS
30-
Coreet
Megatron-T
Fazter MoE
Megatron-TE
TasterMod
Surtmme (mel)
22
1.5
16
32
64
128
16
32
64
120
Number of Experts
Number Experts
(a) H100s
(b) 8 H100s
Figure 12: Forward Latency the Number of experts increases.
We analyze how FlashMoE scales with increasing number of experts at fixed sequence length (T = 16K). Note that for the discussed plots, the number of experts on the x-axis is the total number across all GPUs. Each GPU gets 1/8th of this value. As seen in Figure 12 FlashMoE maintains low, uniform latency, as desired, even as the number of experts grows from to 128. In contrast, baselines exhibit superlinear latency increases due to increased kernel launch overheads. FlashMoE outperforms these baselines by up to 4X at H100s and 6.6X at H100s, both at 128 experts. FlashMoE payload-efficient communication and scheduler-driven in-kernel dispatching allow it to sustain expert parallelism without incurring the communication and orchestration penalties seen in other systems. These results reinforce FlashMoE 's scalability for ultra-sparse MoE configurations.
Limitations and Future Work
Engineering complexity. Fully fused, persistent kernels demand deep GPU distribute-ssystems expertise; future work may investigate compiler/DSL. abstractions to lower this barrier.
FP16 inefficiency. Our FP16 path is suboptimal ($H) due to insufficient tuning. We anticipate
addressing this gap with autotuned GEMM operators like cuBLASDx 35 CUTLASS builders.
Training support. This work targets inference; enabling training will require fusing backward computation and gradient communication with new bookkeeping and task descriptors.
6 Conclusion
We introduce FlashMoF, the first work to fuse the entire Distributed MoE operator into a single persistent GPU kernel that unifies computation, communication, and scheduling via actor-style concurrency, warp specialization, and asyne (R)DMA. We address two dominant bottlenecks in prior -ystems-CPU-managee synchronous communication and fragmented multi-kernel execution. Empirically, FlashMoE achieves up to 6x speedup, 9x higher GPU utilization, and 5.7x throughput for distributed MoE. Looking ahead, we shift from CPU orchestration fully autonomous, GPU-native pipelines extending this fusion approach to training and beyond.
7 Acknowledgements
This research is supported by NSF Award #2444537 and ACE, one seven centers in JUMP 2.0, Semiconductor Research Corporation (SRC) program sponsored by DARPA. This work also used resources of the National Energy Research Scientific Computing Center, DOE Office of Science User Facility supported by the Office of Science of the U.S. Department of Energy under Contract No. DE-AC02-0SCHI 1231 using NERSC award ASCR-ERCA -CC-EERCAP00000766 thank Dr. Giulia Guidi for providing access these NERSC supercomputing resources.
acknowledge and
10

---

References
[1] DeepSeek-Al. Deepseek-v3 technical report, 2025. URL htt...ttpss//axxxi.oorr//a//////
19437
[2] Meta AL The Ilama 4 herd: The beginning of new
new era of natively multimodal ai innovation,
2025. URL http:///.....tt....oo/..//aaaaaa
[3] Mosaic Research. Introducing dbrx: A new state-of-the-art open 1lm. 2024. URL https
/www. www.dd/awr..cck...oooooooo//aaaa http://www...../tocooiooooooooor sttp:/-tttt-art--p-paee open- 11m
[4] Snowflake AI Research. Snowflake arctic: The best 1lm for enterprise ai effi-
ciently intelligent, truly open, 2024. URL ltttps://www..... www../ cttpee.e..///////////////.
arctic-open -tt.eeefffiiitett-oootttaaaaaaae
-nnuagg--odells-snee lake/
[5] OpenAI. GPT-OSS: Open-weight language models for reasoning and agentic tasks. https //github. wtw http//pena//gttosss.t. August 2025. Models: gpt-oss-1205 and gpt-oss-20b.
[6] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, E ukasz Kaiser, and Illia Polosukhin. Attention all you need. In Ad- vances in Neural Information Processing Systems, volume 30. Curran Associates, Inc.,
2017. URL. htaps:://proceedingge wttpp/steepapeepeepeeaaeee.
pdf
[7] Siddharth Singh, Olatunji Ruwase, Ammar Ahmad Awan, Samyam Rajbhandari, Yuxiong He, and Abhinav Bhatele. A hybrid tensor-expert-dar parallelism approach optimize mixture-of- experts training. In Proceedings of the 37th ACM International Conference on Supercomputing, ICS '23, page 203-214, 14, New York, NY, USA, 2023. Association for Computing Machinery. ISBN 9798400700569. doi: 10. 1114/357719333593700 URL htt.hhtts://doo org/10. 10. 1145/ 1577193.3593700
23523 [8] Arjun Devraj, Eric Ding, Abhishek Vijaya Kumar, Robert Kleinberg, and Rachee Singh. Effi- cient allreduce with stragglers, 2025. URL https: wwtp:/xwwwwirxorx///xexeeeeeee/e
[9] Juncai Liu, Jessie Hui Wang, and Yimin Jiang. Janus: A unified distributed training framework for sparse mixture-of-expprts models. In Proceedings of the ACM SIGCOMM 2023 Conference, ACM SIGCOMM '23, page 486-498, New York, NY, USA, 2023. Association for Computing Machinery. ISBN 9798400702365. doi: 10. 0115533663226.3364486 URLhttps:://do org/ 10. N0.1145/3603269. 3604869
[10] Chenyu Jiang, Ye Tian, Zhen Jia, Shuai Zheng. Chuan Wu, and Yida Wang. Lancet: Accelerating -ixterss training via whole graph computation-commuiiaaiii overlapping. In Gib- bons, G. Pekhimenko, and C. De Sa, editors, Proceedings of Machine Learning and Systems, volume 6. pages 74-86, 2024. URL /tsphttes///rrooeeednngss .ml.s.. org wrwgraper........ arg/paper_fiies! files/ oorp Conference, pdf
[11] Samyam Rajbhandari, Conglong Li, Zhewei Yao, Minjia Zhang, Reza Yazdani Aminabadi, Ammar Ahmad Awan, Jeff Rasley, and Yuxiong He. DeepSpeed-MoE: Advancing mixture- of-experts inference and training to power next-generation Al scale. In Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 18332-1 18346. PMLR, 17-23 Jul 2022. URL https
//proceedings..o mlr. www.s/ressseeeee vwtps/vv22///a/ (etees/vv62/raabhheee wtwp:///rraaaaaaaaa22.2.2 html
[12] Shulai Zhang, Ningxin Zheng, Haibin Lin, Ziheng Jiang, Wenlei Bao, Chengquan Jiang, Qi Hou, Weihao Cui, Size Zheng, Li-Wen Chang. Quan Chen, and Xin Liu. Comet: Fine- grained cwwpptaaioonoooommicccccc https://aaxiz: org/ 2502 19811
overlapping for mixture-of-experts, In MLSys '25. URL
[13] NVIDIA. Mcgatron-Im, 2025, URL readme-ov-file 1...
[14] Jiaao He, Jidong Zhai, Tiago Antunes, Haojie Wang, Fuwen Luo, Shangfeng Shi, and Qin Li. Fastermoe: modeling and optimizing training of large-scale dynamic pre-trained models. In Proceedings of the 27th ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming (PPoPP 22), pages 120-134, 2022.

---

[15] Michael Wendt and Joshua Wyatt. Getting started with CUDA graphs. https://develop:: nvidia. www.moooooaoo..o..o cuda- graphs// 2019. Accessed: 2024-05-15.
[16]
Deepak Narayanan, Mohammad Shoeybi, Jared Casper, Patrick LeGresley, Mostofa Patwary, Vijay Korthikanti, Dmitri Vainbrand, Prethvi Kashinkunti, Julie Bernauer, Bryan Catanzaro, Amar Phanishayee, and Matei Zaharia. Efficient large-scale language model training on gpu clusters using megatron-Im. In Proceedings of the International Conference for High Performance Computing, Networking. Storage and Analysis. "21,,, New York, NY, USA, 2021. Association for Computing Machinery. ISBN 9781450384421. doi: 10.11451345811:34776099
URL https://doi. org/ /
[17]
NVIDIA. NVIDIA Nsight Systems Metries,
URL https://dbcs.nviiii
com/ https//taaas/srruuuueeeeeeeeee htt.:///.....aaa.........giaaaaa wrailable-metrice
[18] NVIDIA. Nvidia openshmem library (nvshmem), 2025. URL htt.::ttt//dos..nidddiiie
nvshmem/ane awtpem/m//p/iddeeee html v3.2.5.
[19] Vijay Thakkar, Pradeep Ramani, Cris Cecka, Aniket Shivam, Honghao Lu, Ethan Yan, Jack Kosaian, Mark Hoemmen, Haicheng Wu, Andrew Kerr, Matt Nicely, Duane Merrill, Dustyn Blasig, Fengqi Qiao, Piotr Majcher, Paul Springer, Markus Hohnerbach, Jin Wang, and Manish Gupta. CUTLASS, 2025. BtL/http::/ http//ggtttttt.tt com www.cccom//.//... NVIIII A/ cutlass.
[20] Changho Hwang, Wei Cui, Yifan Xiong, Ziyue Yang, Ze Liu, Han Hu, Zilong Wang, Rafael Salas, Jithin Jose, Prabhat Ram, Ho Yuen Chau, Peng Cheng, Fan Yang. Mao Yang. and Yongqiang Xiong. Tutel: Adaptive mixture-of-eeperts at scale. In D. Song, M. Carbin, and Т. Chen, editors, Proceedings of Machine Learning and Systems, volume 5, pages 269- 287. Curan, 2023. ItL ias.iannnsmssll..... org wwtppprfilleeppppeeeppeeeeppeeee file/
Sapper-minn wwwt-e-aapapaemmmpapeaeeee wwt.-mmlsssssssppmmmm pdf
[21] Jiaao He, Jidong Zhai, Tiago Antunes, Haojie Wang. Fuwen Luo, Shangfeng Shi, and Qin Li. Fastermoe: modeling and optimizing training of large-scale dynamic pre-trained mod- els. In Proceedings of the 27th ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming, PPoPP '22, page 120-134, New York, NY. USA, 2022. Association for Computing Machinery. ISBN 978 9781420392044, bstps://doi.org// 10. 1 3508418
doi: 10.114533502221.35006
URL
[22] Xiaonan Nie, Xupeng Miao, Zilong Wang. Zichao Yang. Jilong Xue, Lingxiao Ma, Gang Cao, and Bin Cui. Flexmoe: Scaling large- sparse pre-trained model training via dynamic device placement. Proc. ACM Manag. Data, 1(1), May 2023. doi: 10.1 145/35888489 URL
https://doi.. ot0e10..rg11.... 1145/3588963
[23] Shaohuai Shi, Xinglin Pan, Qiang Wang, Chengjian Liu, Xiaozhe Ren, Zhongzhe Hu, Yu Yang, Bo Li, and Xiaowen Chu. Schemoe: An extensible mixture-of-experts distributed training system with tasks scheduling. In Proceedings of the Nineteenth European Conference on Computer Systems, EuroSys "24, page 236-249, New York, NY, USA, 2024. Association for Computing Machinery. ISBN 9798400704376.
doi: 10. 10.1145/36277033365000 URL
[24] Hulin Wang, Yaqi Xia, Donglin Yang. Xiaobo Zhou, and Dazhao Cheng. Harnessing inter-gpu shared memory for seamless moe caaamunication-oompuaaaai
fusion. In Proceedings of the 30th ACM SIGPLAN Annual Symposium on Principles and Practice of Parallel Programming, PPoPP '25, page 170-182, New York, NY, USA, 2025, Association for Computing Machinery. ISBN 97984007 14436. 7710848.3710865
doi: 10.1145/37108488 01133700888..33 0w0.0088888778888888 10868. URL https: https:///oi. //doi. org/ 12///10.1 1145/
[25] Tri Dao, Dan Fu. Stefano Ernon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. In Advances in Neu- rat Information Processing Systems, volume 35, pages 16344-16359, Curran Associates, Inc., 2022. URL URL http://ppeaaaaeaaaaaaree.ooop/ap
12

---

[26] Gul A. Agha. Actors: A model of concurrent computation in distributed systems. Technical report, 1985. MIT Artificial Intelligence Laboratory Technical Reports.
[27] Carl Hewitt, Peter Bishop, and Richard Steiger. A universal modular actor formalism for artificial intelligence. IJCAI'3, page 235 245, San Francisco, CA, USA, 1973. Morgan Kaufmann Publishers Inc.
[28] Irene Greif. SEMANTICS OF COMMUNICATING PARALLEL PROCESSES. PhD thesis, Massachusetts Institute of Technology, 1975.
[29]
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. Gshard: Scaling giant models with con- ditional computation and automatic sharding, In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL httpps:///pperreeiiettt wene/foorumm htt////oommmmmmrreeeee Xd=grwe7XHTTmY
[30] Katherine Yelick, Dan Bonachea, Wei- Yu Chen, Phillip Colella, Kaushik Datta, Jason Duell, Susan L. Graham, Paul Hargrove, Paul Hilfinger, Parry Husbands, Costin Iancu, Amir Kamil, Rajesh Nishtala, Jimmy Su. Michael Welcome, and Tong Wen. Productivity and performance using partitioned global address space languages. In Proceedings of the 2007 International Workshop on Parallel Symbolic Computation, PASCO '07, page 24-32, New York. NY, USA, 2007. Association for Computing Machinery. ISBN 9781 9781595937414, 1278183. https::///: org/ 10. 1145/1278177 12115//1778177.121183 1278183
doi: 10.1145/1278177
[31] Size Zheng, Wenlei Bao, Qi Hou, Xuegui Zheng, Jin Fang, Chenhui Huang. Tianqi Li, Hao- jie Duanmu, Renze Chen, Ruifan Xu, Yifan Guo, Ningxin Zheng, Ziheng Jiang, Xinyi Di, Dongyang Wang. Jianxi Ye, Haibin Lin, Wi-Wen Chang, Liqiang Lu, Yun Liang, Jidong Zhai, and Xin Liu. Triton-distributed Programming overlapping kernels on distributed ai systems with the triton compiler, 2025. URL httttstt////rtttii org. abs 2504. 2540211442
[32] Quentin Anthony, Yury Tokpanov, Paolo Glorioso, and Beren Millidge. Blackmamba: Mixture of experts for state-space models, 2024. UKL_attpe//arriv..////////2222.2
[33] NVIDIA. Transformer engine. http://ttargggghha.ooooooo/////r
[34]
Bytedance. Flux's overlap performance worse than (4xrtx 4090], 2025. URL
[35]
NVIDIA. cublasdx, 2025. URL http:://woo...iiddaaa.c/o//ddaaa
www.niiiiii...
[36] NERSC. Network NERSC Documentation. httpss//docs htt.:/ddd....eer..oooooooooo goorss..or/perrformance.
network/ 2025. [Accessed 23-05-2025).
[37] C. Bell, D. Bonachea, R. Nishtala, and K. Yelick. Optimizing bandwidth limited problems using one-sided communication and overlap. Proceedings 20th IEEE International Parallel
& Distributed Processing Symposium, pages 10 pp.-, 2006. doi: 10.1 0..1009/PDPS.2006. 200PP..9.00..6633300
[38] Yuxin Chen, Benjamin Brock, Serban Porumbescu, Aydin Buluc, Katherine Yelick, and John Owens. Atos: A task- mask-parallel gpu scheduler for graph analytics. In Proceedings of the 5 st International Conference on Purallel Processing, ICPP "22, New York, NY, USA, 2023. Association for Computing Machinery. ISBN 9781 9781450397339, doi: 10. 10.1145/5550008.344555 355555 org/ 10. 1145/3545000 3545056
[39] Abhinav Jangda, Jun Huang, Guodong Liu, Amir Hossein Nodehi Sabet, Saeed Maleki, Youshan Miao, Madanlal Musuvathi, Todd Mytkowicz, and Olli Saarikivi. Breaking the computation and communication abstraction barrier in distributed machine learning workloads. In Proceedings of the 27th ACM International Conference on Architectural Support for Programming Languages and Operating Systems (ASPLOS 22), pages 402-416, 2022.
[40] Shibo Wang. Jinliang Wei, Amit Sabne, Andy Davis, Berkin Ilbeyi, Blake Hechtman, Dehao Chen, Karthik Srinivasa Murthy, Marcello Maggioni, Qiao Zhang, et al. Overlap communication with dependent computation via decomposition in large deep learning models. In Proceedings of the 28th ACM International Conference on Architectural Support for Programming I Languages and Operating Systems (ASPLOS 22), pages 93-106, 2022.
13

---

[41] Chang Chen, Xiuhong Li, Qianchao Zhu, Jiangfei Duan, Peng Sun, Xingcheng Zhang, and Chao Yang. Centauri: Enabling efficient scheduling for communication cummunication-compuatttii overlap in large model training via communication partitioning. In Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems (ASPLOS 24), pages 178-191, 2024.
[42] Suchita Pati, Shaizeen Aga, Mahzabeen Islam, Nuwan Jayasena, and Matthew D Sinclair. T3: Transparent tracking & triggering for fine- mne-grained overlap of compute & collectives. In Pro- ceedings of the 29th ACM International Conference Architectural Support for Programming Languages and Operating Systems (ASPLOS 24), pages 1 1146-1166 1164, 2024.
[43] Ziheng Jiang, Haibin Lin, Yinmin Zhong, Qi Huang. Yangrui Chen, Zhi Zhang, Yanghua Peng. Xiang Li, Cong Xie, Shibiao Nong. al. (MegaScale): Scaling large language model training to more than 10,000 {GPUs}. In 21st USENIX Symposium on Networked Systems Design and Implementation (NSDI 24), pages 745-760, 2024.
[44] Weigao Sun, Zhen Qin, Weixuan Sun, Shidi Li, Dong Li, Xuyang Shen, Yu Qiao, and Yiran Zhong. CO2: Efficient distributed training with full commuunccatoonccmppaatee overlap. In The Twelfth International Conference on Learning Representations (ICLR 24), 2024.
[45] Kshiteej Mahajan, Ching-Hsiange Chu, Srinivas Sridharan, and Aditya Akella. Better together: Jointly optimizing (ML} collective scheduling and execution planning using (SYNDICATE). In 20th USENIX Symposium on Networked Systems Design and Implementation (NSDI 23),
pages 809 824, 2023.
[46] Hulin Wang, Yaqi Xia, Donglin Yang, Xiaobo Zhou, and Dazhao Cheng. Harnessing inter-gpu shared memory for seamless moe connunicatonn-oompaaaio fusion. In Proceedings of the 30th ACM SIGPLAN Anmual Symposium on Principles and Practice of Parallel Programming. pages 170-182, 2025.
[47] Kishore Punniyamurthy, Khaled Hamidouche, and Bradford M Beckmann. Optimizing dis- tributed mi communication with fused computation-collecio operations. In SC24: Interna- tional Conference for High Performance Computing, Networking, Storage and Analysis, pages 1-17, 2024.
[48] Changho Hwang, Wei Cui, Yifan Xiong, Ziyue Yang, Ze Liu, Han Hu, Zilong Wang, Rafael Salas, Jithin Jose, Prabhat Ram, et al. Tutel: Adaptive mixture-offexppers scale. Proceedings of Machine Learning and Systems (MLSys 23), 5:269-288 28799228777188 2023.
[49] Chenyu Jiang, Ye Tian, Zhen Jia, Shuai Zheng, Chuan Wu, and Yida Wang. Lancet: Accelerating mixture-of-experts training via whole graph computation-communicc Proceedings of Machine Learning and Systems (MLSys 24), pages 74-86, 2024.
overlapping. In
[50] Jingyuan Liu, Jianlin Su, Xingcheng Yao, Zhejun Jiang, Guokun Lai, Yulun Du, Yidao Qin, Weixin Xu, Enzhe Lu, Junjie Yan, al. Muon scalable for Ilm training. arXiv preprint arXiv:2502.16900 2025.
[51] XAI. Grok-1.
https: btttps//hugg:ae otgpg//aaa....ooo..oo///////oo 2024.
[52] Snowflake.
Snowflake arctic. 2024.
https: attps://huggingfacc.. co/oooeffooo/o/eooeeeeeeeoeoo
Snowf lake/
[53] An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, al. Qwen3 technical report. arXiv preprint arXiv:2505.09388 2025.
[54] Albert Q Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, al. Mixtral of experts. arXiv preprint arXiv:2401.00 044088, 2024.
[55] Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, et al. Deepseek-y3 technical report. arXiv preprint 194337 2024.
14

---

NeurIPS Paper Checklist
Claims
Do the main claims made in the abstract and introduction accurately reflect paper's contributions and scope?
Question:
Answer: [Yes]
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [NA]
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: All reported results the evaluation section were obtained as the average of 32 executions preceded by 32 warmup runs.
8. Experiments compute resources
Question: For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Answer: [Yes]
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
15

---

Justification: We do not foresee immediate social or ethical impacts, but we acknowledge that increased compute efficiency could amplify access to large-scale models, which raises general considerations around prevalent issues such as environmental cost of training, and responsible downstream use. We recommend that users of our system consider these factors when integrating it into broader MI applications.
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
15. Institutional review board (IRB) approvals equivalent for research with human subjects
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Question:
Answer: [NA]
16.
Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
16

---

A
Supplementary Motivation
Gott
Communization
Expert(A, c)
Conmunization
Combina B)
Idle
2t
GPU
D
6c
St
Skew
Communication
Expert[B,
Communication
Combina(S) }))
A
GPU1
diste
Expert
coweiiee
Expert
cowbing
Put
2t
GPU
Put
Put
Skeuc
date
Expart
Expert (b)
combing (D)
cambine
(c)
Figure 13: Overlapped Schedule (bottom) showing how idle time from the sequential schedule (top) repurposed for computation. FlashMoE implements the overlapped schedule.
In Figure 14 we present empirical cumulative and raw distributions of AlltoAll kernel runtime from distributed training of 1.3B GPT-3 MoE model across 32 A100 and V V00 GPUs. We use this result to motivate the severity and prevalence of straggler effects. In Figure 14b] we observe P95 communication performance degradation of 1.32X when compared to the mean actual kernel time. This performance reduction is rather tame as the underlying hardware is supercomputer well-tuned against "software jitter" [36). However, we observe a more severe p95 performance loss of 11X in a single-node Virtual Machine (VM). In line with prior HPC works [37 38 we argue that obviating the inherent barrier this synchronous collective communication would allow GPUs to repurpose this observed idle time for downstream computation depicted in Figure B
Table Straggler Delay within Synchronous All-to-All communication. We capture the distribution of delay induced by stragglers across many steps. Let Actual Time ta denote the fastest kernel execution time across all GPUs, and Total Time be the maximum recorded step time. We define Delay as the maximum difference between t and t.. Note Delay is idle time. For the 1x8 V100, we profile 1750 steps and 600 steps for the 8x4 A100. See Figure for the raw distribution.
System
# Nodes
GPUs
Median
p95
Commercial VM (V100)
8
3.1x
11.4x
Supercomputer (A100)
8
32
1.09x
1.32x
B
Related Work
cuaputaaion-Communittt Overlap and Kernel Fusion. To reduce the communication over- heads of synchronization in distributed DNN training, many research efforts have been focused on increasing the overlap of computation and communication. For generic Transformer-baa without MoE. layers, many works [39477 have provided insights and techniques to partition and schedule computation and communication operations, aimed at finer-grained overlapping. address the challenges posed by AlltoAll communication and expert parallelism in MoE training, Tutel [48]
models
17

---

GPT-3 Mct 32x1. All Straggler Effect 8x4
GPT33 McE 32x1. All Straggler Effect 8x4
Distribution max delay
5.0
Total Time
3.22 ms
4.5
2vg actual
1.90ms
0.8
JI - 0.28m4
mediaan 0. 0.19s
P95 0.54ms
0.4
2.5
2.0
0.0
0.5
1.0
1.5
2.0
2.5
3.0
200
200
300
400
500
600
Dalay (ms)
All-to-All Stops
(a) ECDF
(b) Raw Distribution
GPT-3 MoE 8x350M All to-too All Straggler Effect: 1x8 V100
GPT-3 MoE 8x350M to-to-All Stragpler Effect: 1x8 V100
ECDF
Raw Distribution
Total Time
0.8
modion - 0:57mm
JI - 1.17mm
Time (mas)
mox delay 27.00 ms
P95 2.82ms
ogg ctual 0.2
0.4
5
10
15
20
25
250
500
750
10000
1250
1500
1750
Delay (ms)
All-toAAll
(c) ECDF
(d) Raw Distribution
Figure 14: Straggler effect of synchronous AlloAll. M N A100 or V100 denotes N GPUs within node across M nodes. Every GPU communicates with every other GPU per AlltoAll step. We capture the distribution of delay induced by stragglers across many steps. Actual Time ta denotes the fastest kernel execution time across all GPUs, conversely Total Time is the maximum recorded step time, while Delay is the maximum difference between and ta- Note Delay is idle time.
and FasterMoE [4] overlap AlltAAl with expert computation. Lancet 49 additionally enables both non-MoE computation in forward pass and weight gradient computation in backward pass to be overlapped with Alltoll. Despite overlapping, the performance of these approaches is limited in practice due to blocking synchronous collective communication with barriers. In contrast, Flash- MoE fundamentally eliminates these inefficiencies with asynchronous, device-initiated data transfers overlapped with tiled computation all within single kernel. FlashMoF further differentiates itself from SOTA works like COMET (12) and DeepEP [I], which also use this form of kernel-initiated communication but a coarse-grained granularity and without complete kernel fusion.
C Proof of Theorem 3.1
We begin with two necessary definitions vital to the proof.
Definition C.1. Define a write as p(,s..pt.) i), where ps is the source process and is an ordered tuple indicating the index coordinates for L, residing on the target process Pt- A write-write conftict occurs when there exist at least two distinct, un-synchronized, concurrent writes wy (Poy: (1,,21) and wz w2(poz Pizs i2), such that Pt, = Pt, and index coordinates i i2 but Ps, # Ps2
Definition C.2. For any source process Pa, valid index coordinate = (p*, T, b, e, c) satisfies the following:
For inter-device writes, it must hold that pa = Ps and b 1. Note this also applies to self-looping writes w w(ps, Pt,i).
18

---

2. For any write w(ps: 2t, if = 0, then Ps This rule describes intra-devics staging writes.
We restate Theorem
3.1 and outline its proof below.
Theorem C.I. The symmetric tensor layout L is write-write conflict-free.
Proof. As is the case for typical physical implementations, assume that each index coordinate maps to distinct memory segment in L. Next, we show by contradiction that no write-write conflicts can exist when accessing L using valid For simplicity, we only include the index coordinates when describing write. Assume that there exist at least two writes w/ (Ps, Ps, d1), W2(Psz, Ptz i2) with Pt. - Pts and valid destination coordinates i, i2, where i = i2 lexicographically and both are unpacked below.
(p1, ri, et, 61), 12 (pz, r2, 12,b2,e2,1
Note that intra- process writes always have distinct Cj coordinates, where j E {0, C 1}. For inter-process transfers, we have two cases.
Case 1: Psi Psz
Here, wh and th are identical operations. This contradicts the definition of a conflict, which requires that Ps; Psg- In practice, such repeat writes
Case 2: Pei # Pag
To ensure validity for i and iz, is the case that p1 Pst and p2 Psg. However, this implies that o i2 yielding a contradiction as desired.
Memory Overhead
We measure the GPU memory required for the symmetric tensor L and runtime bookkeeping state of FlashMoE. The memory overhead primarily depends on the tile size, expert capacity (EC), and the number of experts (E). Table summariizs the memory overhead across recent MoE models (50 during inference, showing that FlashMoE maintains a modest and predictable memory footprint. In particular, the symmetric tensor (ST) accounts for most 2. 2.55% additional memory relative the total inference memory requirements.
Table 3: Memory overhead of FlashMoB (tile size bM = 128, Size(T) Tokens x 4KB).
Model
Params
1
(GB)
Model (GB)
Overhead (%)
Moonlight-16RB 1athht--6B-AAt 16B-A3B
16B
8K
64
2K
1.38K
0.25
59
0.49
Grok-1
314B
8K
8
6K
32K
0.75
1169
0.15
Snowflake-Artte
479B
4K
128
7K
4.75K
1.75
1784
0.12
Swen3-235B-A22
235B
40K
128
4K
1.5K
3.00
875
0.38
Mixtral 8x7B
478
32K
4K
14K
2.00
175
2.15
DeepSeek-V3
685B
160K
256
2K
1.50
2551
0.11
19

---

E Task Implementation
1 #define GEMMs 2
2
struct allign (16) Task {
3
const byte* aData;
4
arraycconst byte*, GEMMs> bData; array<byte*, GEMMs> cData;
5
6
arraycconst byte*, GEMMs> dData; byte* rcData; uint64_t* flags;
7
8
9
uint M;
10
uint syncIdx;
11
uint tileIdx;
12
uint batchIdx;
13
uint peerIdx;
14
uint expertIdx;
15
uint isPeerRemote;
16
TaskType taskType;
17
uint16_t tileSize;
18
// Pad till 128-byte cache line
19
uint [adding[6] 0
20 }
Figure 15: Task Struct. Tysk E (GEMM, GEMM,, Combine)
20

---

Actors
F1 Processor
Algorithm Processor Actor: executed by block
begin
1Q+ + GetTQ()
signal +-0 0
// shared memory variables
task t 0
interrupt + False
complete False
while interrupt = False do if warpld = then
10
if threadld = 0 then
11
await waaiiai/oooosscceeeeeeeeeeeeeeee signal) FencedNotiff
13
end if syncwarp()
14
15
ParpReadTQ(tQ) signal, task)
16
end if syncthreads()
18
if interrupl = False then
switch Type do
20
case GEM do
1/ fused GEMM, epilogue and async tile staging task)
22
if threaddd then
complete < NotifyTiieCCompltion))
end if syncthreads()
27
if complete = True then NoitfyScheculerNexCCCEMMMM
29
end if end case case GEMM, do
30
31
// fused GEMM, epilogue and async tile transfer fGET(GEMM,, task)
33
34
end case
case Combine do combine(task)
36
end case end switch 37
end if end while
40
41 end
21

---

F.2
Scheduler
Algorithm 3: Scheduler Actor: executed by one warp
begin
scheduled +- 0
BTB < 0
tqState + 0
pTDB - Pooeeerllssoolllrrellllll
sTDB + GedSubbcrilerr
laskBoond + Get
B - (ts
// circular buffer ready queue
10
rQ 0
11
// Populate ready queue with Processor ids
12
PopulateRQ(rQ)
while scheduled < tTB do
14
6-
15
do in parallel
16
Sweep doorbells and populate observed task counts into State Aggregate locally observed task counts into It
18
end
19
gSS taskTally 0
20
// qs is the inclusive output
21
WarplnclusiveSummtt, qS, tasktally)
22
while tasktally > 0 do
23
Repopulate rQ with ready processor ids
24
do in parallel
Starting at rQ[g5], signal processors about task indices from State end end while 27
25
28
threadld == then
29
(TB + Atomiicooa((askBBonn
30
end if
B t- WarpBrondcast((T1)
32
end while
AnterruptSSbbscrienn
InterruptProcessorry
35 end
22

---

F33
Subscriber
Algorithm 4: Subscriber Actor: executed by three warps
To E (RR3) ExC Go ERXE O O X E REXHxD
Input:
begin
interrupt GetSharednntrrupp)) flags t G
iQ GetTQ()
// Predefined upper bound on the number of tasks.
// We modulate this value to the actual task count computed
// dispatch sigmals received from peer GPUs
task GetTaskBound()
while otd (interruppt))) == False do
10
// dispatch flags
11
do in parallel
12
Visit dispatch flags
13
Atomically retrieve signal
14
Signal is set and flag is not visited then
15
Mark visited
16
Bound, Signal)
Enforce memory consistency before consuming packet
18
Decode packet into set of GEM My task descriptors using X
Write task descriptors tQ Notify Scheduler decoded tasks
20
21
end if
22
end
23
Advance flags by number of dispatch flags length Atomically retrieve signal
24
25
// combine signals
26
do in parallel
27
Visit combine flags: one per tile
28
if Signal is set and flag not visited then
29
Mark visited
30
Enforce memory consistency before consuming packet
31
Decode packet into set of combine task descriptors using Top G O
32
Write task descriptors to tQ
Notify Scheduler of decoded tasks
34
end if
35
end
36
end while
end
23

---

G
Implementation
Table 4: Implementation metrics of FlashMoE.
Metric
Value
Total lines of code (CUDA/C++)
6820
Kernel stack frame size
0 B
Spill stores (per thread)
0
Spill loads (per thread)
0
Shared memory usage (per block)
46 KB
Registers per thread
255
Max active blocks per SM
2
Compilation time
53 seconds
Binary size
29 MB
H FP16 Memory Throughput
(a) Memory subsystem throughput for FP16
WANGE
2G
(b) Memory subsystem throughput for FP32
Figure 16: Here, we report the total A100 memory throughput for both FP16 (top) and FP32 (bottom) variants of FlashMoE. Notably, the FP16 implementation issues approximately 2x more shared memory instructions compared to its FP32 counterpart under identical workloads. We attribute this inefficiency to suboptimal shared memory layouts in FlashMoE when operating half half-precision data. While this bottleneck addressable through improved layout strategies, we leave its resolution future work.
24


Reviews:
[{"review_id": "9MLQRl3Y6e", "reviewer": "Reviewer_ST1R", "summary": "This paper introduces FlashDMoE, a novel system for executing the Mixture-of-Experts (MoE) layer in distributed deep learning environments. The authors identify two critical performance bottlenecks in existing MoE implementations: (1) high communication overhead and straggler sensitivity from synchronous \"ALLTOALL\" collectives, and (2) significant latency from the host CPU launching hundreds of individual GPU kernels.\n\nTo address these issues, the authors propose a paradigm shift: fusing the entire MoE operator—including gating, expert computation, and inter-GPU communication—into a single, persistent GPU kernel. This design eliminates kernel launch overhead and enables fine-grained, asynchronous execution. The core technical contributions are:\n\n1. An actor-based concurrency model within the kernel, where GPU warps are specialized into roles (Processor, Scheduler, Subscriber) to manage tasks.\n2. A device-initiated asynchronous communication protocol built on a novel symmetric tensor layout. This layout overprovisions memory to create a partitioned global address space (PGAS), enabling non-blocking data transfers without synchronization penalties.\n\nPayload-efficient communication via \"in-place padding\", which avoids transmitting zero-padded tokens across the network.\n\nThe authors evaluate FlashDMoE on an 8x H100 GPU server, demonstrating significant forward-pass latency speedups of 1.2x to 15x over state-of-the-art systems like DeepEP, Comet, and Megatron-LM.", "questions": "Q1. How sensitive is the performance of the symmetric tensor layout to the number of communication rounds ($\\tau$)? The paper states $2\\cdot\\tau$=4 for MoE, but how would the design and its memory overhead adapt to more complex computational graphs with more communication steps?\n\nQ2. The actor model dedicates one warp to the Scheduler role. Have you observed this scheduler becoming a bottleneck in any configuration, particularly with a very high number of experts or small tile sizes, leading to a high rate of task dispatching?\n\nQ3. The evaluation focuses exclusively on the forward pass for inference. Have you considered the implications of this design for the backward pass in a training scenario? Or, you can modify the architecture in a training-free manner? How would the actor model and communication scheme handle the additional complexities of gradient computation and accumulation?", "limitations": "N/A", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "S1. The concept of a fully fused, single-kernel MoE operator is a radical and elegant departure from conventional approaches. It fundamentally rethinks the execution model to eliminate well-known bottlenecks at their source, rather than merely optimizing around them. If this approach proves generalizable, it could influence the design of other complex communication-intensive operators in deep learning.\n\nS2. The experimental results are compelling. Achieving speedups of up to 15x and maintaining over 90% SM utilization is a testament to the effectiveness of the design. The system's ability to maintain near-constant latency when scaling the number of tokens (Figure 10a) is particularly noteworthy and highlights the success of the computation-communication overlap.\n\nS3. The symmetric tensor layout is a intuitive and insightful solution to the challenge of asynchronous, one-sided communication. Trading a predictable amount of memory for the complete removal of synchronization barriers is a powerful and well-justified engineering decision. The actor model provides a clean abstraction for managing the complex interactions within the fused kernel.\n\nS4. The authors conduct a comprehensive evaluation against multiple strong baselines. The experiments are well-designed to isolate performance across different scaling dimensions (tokens, experts, GPUs). The inclusion of micro-benchmarks that analyze GPU utilization (Fig 12a) and payload efficiency (Fig 12b) provides clear evidence for why the system is faster, strengthening the paper's claims.\n\nW1.The primary concern is the 4x memory overhead required for the token buffer to enable the symmetric tensor layout. The authors describe this as \"modest,\" but for models with very large sequence lengths or on memory-constrained hardware, this could be a prohibitive cost. The paper would benefit from a more detailed analysis of this trade-off, perhaps identifying the point at which this memory cost outweighs the performance benefits.\n\nW2. The claims of the paper are based entirely on a single-node, multi-GPU setup with high-bandwidth NVLink interconnect. The most challenging ALLTOALL bottlenecks and straggler effects often manifest in multi-node clusters connected over lower-bandwidth networks like InfiniBand. Without multi-node results, the paper's claims about scalability and straggler mitigation remain partially unsubstantiated. This is the most significant gap in the current evaluation.\n\nW3. The implementation is restricted to FP32 precision. Modern LLM inference and training workflows heavily rely on lower-precision formats (FP16, BF16, FP8) for performance and memory efficiency. This limitation currently makes FlashDMoE impractical for many state-of-the-art models.\n\nW4. The system is a highly specialized, low-level implementation (6,800+ lines of CUDA/C++). This presents a potential barrier to adoption and maintenance compared to solutions built on standard, well-documented libraries like NCCL. While not a flaw in the research, it's a practical consideration worth acknowledging.", "was_revised": true, "final_justification": "I will keep my score.", "rebuttal": "We thank you for your constructive comments and glowing commendations of our work. Below, we address your concerns (weaknesses W1–W4) and answer your thoughtful questions (Q1–Q3):\n\n## W1: Memory overhead for the Symmetric Tensor\n\nWe acknowledge the concern that FlashDMoE requires extra memory for the symmetric tensor (Section 3.2). We maintain that this overhead is modest in comparison with the existing memory requirements of MoE models. To show this empirically, we compute and analyze memory costs for popular large MoE models:\n\n| **Model**          | **Params** | **S** | **E** | **H** | **Intermediate** | **Symmetric Tensor (GB)** | **Memory for Model (GB)** | **Overhead Ratio (%)** |\n|--------------------|------------|-------|-------|-------|------------------|---------------------------|---------------------------|------------------------|\n| Moonlight-16B-A3B  | 16B        | 8K    | 64    | 2K    | 1.375K           | 0.25                      | 59                        | **0.49**               |\n| Grok-1             | 314B       | 8K    | 8     | 6K    | 32K              | 0.75                      | 1169                      | **0.15**               |\n| Snowflake-Arctic   | 479B       | 4K    | 128   | 7K    | 4.75K            | 1.75                      | 1784                      | **0.12**               |\n| Qwen3-235B-A22B    | 235B       | 40K   | 128   | 4K    | 1.5K             | 3                         | 875                       | **0.38**               |\n| Mixtral 8x7B       | 47B        | 32K   | 8     | 4K    | 14K              | 2                         | 175                       | **2.15**               |\n| DeepSeek-V3        | 685B       | 160K  | 256   | 7K    | 2K               | 1.5                       | 2551                      | **0.11**               |\n\n\nWe calculate the model’s memory requirement by multiplying its parameter count with precision (fp32 here but this choice is inconsequential to the ratio which we embolden). Note that the symmetric tensor demands **at most 2% additional memory** in comparison to the inference memory requirements. This overhead would even be more amortized in training, where the memory demand for the model is much greater. We will be sure to add a more thorough empirical analysis elaborating on the memory and performance tradeoffs exhibited in FlashDMoE.\n\n## W2: Multi-node results with FlashDMoE\n\nWe acknowledge that the main body of the paper did not present results of FlashDMoE on multi-node hardware. To address this concern, we successfully set up FlashDMoE on 4 GPU nodes, where each node contains 4 A100 80G GPUs (16 A100 GPUs in total), on a Supercomputer in the United States. Every GPU on a node has a dedicated NIC with 25 GB/s of bandwidth. There are 16 experts in this experiment. We provide preliminary results of the FlashDMoE operator latency on this multi-node setup:\n\n| Tokens | FlashDMoE Latency (ms) |\n|--------|--------------|\n| 128    | 1.22         |\n| 256    | 1.26         |\n| 512    | 1.29         |\n| 1024   | 1.43         |\n| 2048   | 1.74         |\n\nWe note that these results do not test all possible configurations in terms of number of tokens and experts. Due to limited time, these were the only configurations we managed to test. But, we hope they address some reviewer concerns regarding multi-node evaluation of FlashDMoE. We intend to run FlashDMoE across multiple nodes with a full set of configurations and compare with popular baselines in the near future. We will include these results in the camera-ready version of the paper (if accepted).\n\n## W3: Precision\nWe acknowledge the concern regarding FlashDMoE’s current use of FP32 precision. We note that our evaluation in the paper using FP32 represents a conservative estimate of FlashDMoE’s potential benefits. Despite operating in full precision, FlashDMoE substantially outperforms baseline systems that already utilize FP16 precision. Thus, transitioning FlashDMoE to FP16 precision will further improve these already significant performance gains. \n\nOn another note, we acknowledge the use of fp32 as a limitation in our paper and we highlight that FlashDMoE already supports half precision in bf16 and fp16. However, our fp16 implementations do not yet achieve peak hardware utilization due to (1) suboptimal shared memory layouts and (2) insufficient vectorization in custom computational operators within the Processor. To overcome these limitations, we plan to integrate recent advancements from Device Extensions (DX) libraries [1]. These libraries build on top of the CUTLASS APIs and provide a specialized configuration layer with highly optimized parameters, significantly improving both hardware memory throughput and tensor-core utilization. We anticipate adopting [1] in FlashDMoE and, in fact, already have ongoing work to that effect. In summary, this transition would involve deprecating the CollectiveMainLoop defined in `mmaConfig.cuh` and line 80 of `gemm.cuh`, in favor of a bespoke GEMM main loop that accumulates across the K dimension, where each iteration dispatches to the highly-tuned, block-scoped operator of [1]. Importantly, [1]’s adherence to returning GEMM results as fragments stored in fast register memory seamlessly preserves FlashDMoE’s ability to execute post-GEMM fusion operations. We plan to make these changes available in our public codebase.\n\n## W4: Maintainability of FlashDMoE\nFlashDMoE's specialized CUDA/C++ implementation may pose challenges for maintenance, but not necessarily for _adoption_ as long as we employ ergonomic user abstractions. An apt case in point is NCCL, whose implementation details are quite intricate but exist behind well-defined and documented APIs that have effectively facilitated its adoption in deep learning frameworks like PyTorch. We envision FlashDMoE following the same trajectory as we already have designed ergonomic, user-friendly APIs that abstract the complexity of the fused kernel. Moreover, our low-level approach was critical for achieving significant performance gains over the state-of-the-art. We will make substantial efforts to highlight these points in the camera-ready version of the paper (if accepted). \n\n## Q1: Performance sensitivity to increasing communication rounds\nThank you for the question. In principle, our temporal buffering technique is not a functional requirement for our operator; less temporal slots $t$ necessitate locks or busy-waiting to ensure conflict-free access. Thus, we conceptualize $t$ as a tunable parameter that the user can modulate to navigate this tradeoff optimally in the general case, including for computational graphs where $t > 4$.\n\n## Q2: Scheduler bottleneck\nExcellent question! The cost of dispatching 32 (32-way parallelism in a warp) tasks is equivalent to the latency of a global memory atomic operation, which is in the order of at least triple-digit nanoseconds. Even at tens of thousands of tasks, such latency only creeps to the double-digit microsecond regime, which is about the latency of computing a single 128x64 fp32 output tile. So, the bottleneck should not be task dispatching itself, but (1) how many processors are available at dispatch-time and (2) how quickly processors become available. Note that a processor's queue-depth is at most 1; it receives and processes only one task at a time and does not buffer. Thus, the more likely bottlenecking scenario is that the scheduler may have available tasks but no ready processors to dispatch to, as they are busy with computation or tile communication.\n\n## Q3: Extensibility to training\nFlashDMoE is generalizable to training with modifications only needing to occur for (1) the Subscriber to account for encoding tasks required for the additional gradient GEMMs, and (2) additional task types to capture these functions. Communication would remain as-is, using the same symmetric layout and initiated eagerly by thread blocks when tile computation completes.\n\n## References\n\n[1] NVIDIA. cuBLASDx."}, {"review_id": "9ZdnMF0InG", "reviewer": "Reviewer_MD3K", "summary": "This paper proposed FlashDMoE, a fully fused MoE operator that runs as a single GPU kernel. It integrates all compute and communication tasks directly on the GPU, bypassing the need for multiple kernel launches and improving GPU utilization. The innovative of this paper include 1. replacing hundreds of short-lived GPU kernels with one long-lived kernel; 2. breaking work into tiles processed in parallel across GPU threads; 3. communicating only with GPUs hosting selected experts. The experiment result shows it achieves 1.2× to 15× faster performance compared to SOTA method like DeepEP, Megatron-LM, FasterMoE, and Comet.", "questions": "1. How difficult would it be to adapt FlashDMoE to support other MoE routing strategies, such as top-k routing with dynamic capacity?\n2. How easily can FlashDMoE be integrated into popular frameworks like Transformers or PyTorch?\n3. Do you plan to open-source the implementation? If so, will it include utilities for configuring different model shape and architecture?", "limitations": "Please refer to weakness and questions.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. Novel Single-Kernel Design: FlashDMoE introduces a unified, persistent GPU kernel that eliminates host-side kernel launch overhead.\n2. Deep Systems Innovation: The paper presents a well-engineered actor-based GPU execution model with tile-level parallelism and efficient memory management.\n3. Substantial Performance Gains: The approach achieves 1.2× to 15× speedup over existing MoE systems and significantly improves GPU utilization.\n\nWeaknesses:\n1. Lack of Open Source and Flexibility: The code is not open-sourced, and adapting the implementation to new MoE architectures or shapes may be complex.\n2. Limited Evaluation Scope: All experiments are conducted on a single-node, 8-GPU system, with no multi-node benchmarks.\n3. Missing Baseline Comparisons: The paper does not compare against other recent expert-parallel fusion methods, such as [Optimizing Distributed ML Communication with Fused Computation-Collective Operations](https://arxiv.org/html/2305.06942v2).", "was_revised": true, "final_justification": "I would give this paper a borderline accept because it is a complete work with easy-use design, sufficient experiments, and well-written manuscripts.", "rebuttal": "We thank reviewer MD3K for their constructive feedback. Below, we address the concerns raised about open-source nature and flexibility of FlashDMoE, the limitations of our evaluation framework, and the absence of baseline comparisons. We hope this response highlights the robustness of our approach and its potential for future enhancements.\n\n## W1: Open Source Code\n\nWe echo the raised concern about open-source software, hence why we included our open-source, workable code in the supplementary material of this submission. We currently have a well-maintained, public GitHub repository (under a different name to preserve anonymity) with over 30 stars and a rapidly growing user base. To preserve the integrity of the review process and maintain anonymity, we refrained from including links to this repository in the manuscript but reiterate that we provide the code in supplemental materials.\n\n## W1: Flexibility of FlashDMoE\n\nFlashDMoE can already accommodate arbitrary dimensions for sequence length, token embedding dimensions, and FNN intermediate size. The included JSON file (`aristos_config.json`) in the included FlashDMoE codebase allows users to specify these parameters. Additionally, we support unbounded k for expert top-k selection, offer four different precisions (fp16, bf16, tf32, and fp32), and can accommodate up to 65,535 experts—significantly exceeding the architectural specifications of any state-of-the-art MoE model and the memory requirements of modern GPUs. \n\nWe note that in its current implementation, FlashDMoE does not support alternative MoE architectures like the shared expert approach used in DeepSeekMoE [1] or the Pyramid-Residual MoE technique introduced in DeepSpeedMoE [2]. Yet, we optimistically highlight that the key ingredients for unlocking these architectural shifts, and future ones, already exist in the design of FlashDMoE: tile-parallel implementation and task abstraction for computation. Specifically, (1) underlying these MoE variants are computational operators that are describable in our task framework (Section 3.1), and (2) for expert parallelism, the symmetric tensor (Section 3.2) will still support scalable inter- or intra-GPU communication among all three actors, whose algorithms remain structurally unchanged across these architectural variations.\n\n## W2: Limited Evaluation on multi-node setups\n\nWe acknowledge that the main body of the paper did not present results of FlashDMoE on multi-node hardware. To address this concern, we successfully set up FlashDMoE on 4 GPU nodes, where each node contains 4 A100 80G GPUs (16 A100 GPUs in total), on a Supercomputer in the United States. Every GPU on a node has a dedicated NIC with 25 GB/s of bandwidth. There are 16 experts in this experiment. We provide preliminary results of the FlashdMoE operator latency on this multi-node setup:\n\n| Tokens | FlashdMoE Latency (ms) |\n|--------|--------------|\n| 128    | 1.22         |\n| 256    | 1.26         |\n| 512    | 1.29         |\n| 1024   | 1.43         |\n| 2048   | 1.74         |\n\nWe note that these results do not test all possible configurations in terms of number of tokens and experts. Due to limited time, these were the only configurations we managed to test. But, we hope they address some reviewer concerns regarding multi-node evaluation of FlashDMoE. We intend to run FlashDMoE across multiple nodes with a full set of configurations and compare with cutting-edge baselines in the near future. We will include these results in the camera-ready version of the paper (if accepted).\n\n## W3: Missing Baseline Comparison\n\nThank you for pointing out this related work; we highlight that we discussed and cited this SC'24 paper in the related works section (Section C) of our manuscript. To the best of our knowledge, we could not find an open-source implementation provided by the authors. We will make an effort to replicate their fused GEMM+All-to-All kernel for comparison with FlashDMoE.\n\n## Q1: MoE Architectural variants\n\nPlease see **W1: Flexibility of FlashDMoE**. Furthermore, we will be sure to include a more thorough discussion about how FlashDMoE can adapt to these architectural variants in the camera-ready version of the paper (if accepted).\n\n\n## Q2: Integration to PyTorch\n\nThanks for asking! FlashDMoE exposes a minimalistic forward API, line 73 of `moe.cuh`, which we anticipate upstreaming as a Python API that integrates with frameworks like PyTorch, HuggingFace Transformers or Megtron-LM. We hope to complete this Python release before the camera-ready deadline.\n\n## Q3: Open Sourcing FlashDMoE\n\nKindly see **W1: Open Source Code**\n\n## References\n\n[1] DeepSeek-AI. DeepSeek-V2, arXiv, 2024.  \n\n[2] Rajbhandari et al. DeepSpeed-MoE, ICML ‘22."}, {"review_id": "1zvZyIdFTW", "reviewer": "Reviewer_KK8z", "summary": "Although MoE models are efficient because they only activate a few experts at a time, spreading experts across many GPUs causes heavy communication and frequent GPU launches, especially due to ALLTOALL operations. To solve this, FlashDMoE combines all computing and communication steps into one long-running GPU kernel, removing the need for CPU control.", "questions": "Have you trained an MoE model to compare the stability, throughput, and loss of FlashDMoE and Megatron during actual training?", "limitations": "yes", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n* It achieves significant speedups, with 1.2X-15X speedup over state-of-the-art MoE implementations in evaluations on multiple ML servers each with 8 NVIDIA H100 GPUs, and 3−10× improvement in forward pass latency on a single node with 8 NVlink-ed H100 GPUs compared to existing MoE systems.\n* It eliminates repeated kernel launches, launching only one persistent kernel during a single forward pass, while baselines launch up to 550 short-lived kernels, thus reducing kernel launch overheads significantly.\n* It mitigates the impact of communication stragglers through fine-grained, tile-level parallelism within the kernel, addressing the long-tail latency issue of synchronous ALLTOALL communication.\n* It improves payload efficiency by introducing in-place padding and symmetric tensor layouts, transmitting only active tokens, and reducing unnecessary data transfers, which leads to better runtime performance.\n\nWeaknesses:\n* FlashDMoE is currently only evaluated on a single node with 8 GPUs connected via NVLink. Its performance on multi-node systems, especially over slower interconnects like InfiniBand, is not yet tested.\n* To enable non-blocking communication, FlashDMoE over-provisions memory using a 4× symmetric tensor layout, which increases memory consumption compared to some other MoE implementations.\n* The system is benchmarked on synthetic workloads and toy MoE layers, without integration into full-scale production models like LLaMA or GPT. Thus, the practical benefit of end-to-end training/inference remains to be validated.", "was_revised": true, "final_justification": "I appreciate the authors' response. It has addressed my concerns, so I decide to keep my original rating.", "rebuttal": "We thank you for your constructive feedback on our work. We provide details to address your concerns (weaknesses W1–W3).\n\n## W1: Multi-node experiments with FlashDMoE\nWe acknowledge that the main body of the paper did not present results of FlashDMoE on multi-node hardware. To address this concern, we successfully set up FlashDMoE on 4 GPU nodes, where each node contains 4 A100 80G GPUs (16 A100 GPUs in total), on a Supercomputer in the United States. Every GPU on a node has a dedicated NIC with 25 GB/s of bandwidth. There are 16 experts in this experiment. We provide preliminary latency results on this multi-node setup:\n\n| Tokens | FlashDMoE Latency (ms) |\n|--------|--------------|\n| 128    | 1.22         |\n| 256    | 1.26         |\n| 512    | 1.29         |\n| 1024   | 1.43         |\n| 2048   | 1.74         |\n\nWe note that these results do not test all possible configurations in terms of number of tokens and experts. Due to limited time, these were the only configurations we managed to test. But, we hope they address some reviewer concerns regarding multi-node evaluation of FlashDMoE. We intend to run FlashDMoE across multiple nodes with a full set of configurations and compare with popular baselines in the near future. We will include these results in the camera-ready version of the paper (if accepted).\n\n## W2: Memory footprint of the Symmetric Tensor\nWe acknowledge the concern that FlashDMoE requires extra memory for the symmetric tensor (Section 3.2). We maintain that this overhead is modest in comparison with the existing memory requirements of MoE models. To show this empirically, we compute and analyze memory costs for popular large MoE models:\n\n| **Model**          | **Params** | **S** | **E** | **H** | **Intermediate** | **Symmetric Tensor (GB)** | **Memory for Model (GB)** | **Overhead Ratio (%)** |\n|--------------------|------------|-------|-------|-------|------------------|---------------------------|---------------------------|------------------------|\n| Moonlight-16B-A3B  | 16B        | 8K    | 64    | 2K    | 1.375K           | 0.25                      | 59                        | **0.49**               |\n| Grok-1             | 314B       | 8K    | 8     | 6K    | 32K              | 0.75                      | 1169                      | **0.15**               |\n| Snowflake-Arctic   | 479B       | 4K    | 128   | 7K    | 4.75K            | 1.75                      | 1784                      | **0.12**               |\n| Qwen3-235B-A22B    | 235B       | 40K   | 128   | 4K    | 1.5K             | 3                         | 875                       | **0.38**               |\n| Mixtral 8x7B       | 47B        | 32K   | 8     | 4K    | 14K              | 2                         | 175                       | **2.15**               |\n| DeepSeek-V3        | 685B       | 160K  | 256   | 7K    | 2K               | 1.5                       | 2551                      | **0.11**               |\n\nWe calculate the model’s memory requirement by multiplying its parameter count with precision (fp32 here but this choice is inconsequential to the ratio which we embolden). Note that the symmetric tensor demands at most 2% extra memory in comparison to the inference memory requirements. This overhead would even be more amortized in training, where the memory demand for the model is much greater. We will be sure to add a more thorough empirical analysis elaborating on the memory and performance tradeoffs exhibited in FlashDMoE.\n\n## W3: Integration \nWe acknowledge that our evaluation benchmarks MoE layers rather than full-scale production models like LLaMA or GPT. While our experiments rigorously quantify fundamental performance improvements from having a fully-fused distributed MoE implementation, validating these benefits through integration into production models remains our future work. We are actively investigating this line of work and are eagerly look forward to showcasing the performance of FlashDMoE in realistic training and inference scenarios."}, {"review_id": "TMemXbAps5", "reviewer": "Reviewer_pAJq", "summary": "This paper introduces FlashDMoE, a system designed to accelerate the execution of distributed Mixture-of-Experts (MoE) layers in large language models. \n\nThe work identifies two bottlenecks in existing MoE implementations: communication and kernel launch overheads.   \n\nTo solve these problems, the authors propose FlashDMoE, an MoE operator architected as a single persistent GPU kernel. Instead of relying on the CPU to launch a sequence of kernels, FlashDMoE launches a monolithic kernel that remains active throughout the entire MoE operation, managing all internal tasks directly on the GPU to eliminate kernel launch overhead by design, reducing the number of required GPU operations from hundreds to one.   \n\nThe core mechanisms of FlashDMoE are threefold:\n\n1. Actor-Based Concurrency\n2. Asynchronous, Device-Initiated Communication\n3. Payload-Efficient Data Transfer\n\nThe authors present a comprehensive evaluation of multi-GPU servers equipped with NVIDIA H100s. The results demonstrate that FlashDMoE achieves speedups ranging from 1.2x to 15x over state-of-the-art baselines. Furthermore, the system achieves over 90% Streaming Multiprocessor (SM) utilization, showcasing its effectiveness at eliminating GPU idle time.", "questions": "Please see above.", "limitations": "Yes.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 4, "originality": 3, "strengths_and_weaknesses": "Strength:\n\n1. The idea is intuitive and addresses bottlenecks in traditional LLM kernels.\n2. The experiments show good latency hiding and resource utilization results.\n\nWeakness\n\n1. While the idea of a monolithic kernel is new, the core techniques are not. The idea of integrating communication and computation tasks within fused GPU kernels and thread block specialization is well-studied and are in one of the baselines, Comet, that the authors compare with. Other than pushing the idea to a single kernel, which comes with its drawbacks of less flexibility, it is not clear what the additional challenges are.\n\n2. It's not clear how the additional activation buffer would affect per-device batch size in a larger-scale scenario. If this extra memory requirement causes higher GPU requirements, then a more detailed analysis of the efficiency-memory tradeoff is needed.\n\n3. FP32 is a suboptimal setup for both inference and training. \n\n4. It's not clear how much additional effort is required to support more diverse expert routing algorithms.\n\nMinor typos:\n\n1. Extra comma in line 43.\n2. Incomplete/fragmented sentence in line 138.\n3. Duplicated sentence in line 205.", "was_revised": true, "final_justification": "I am keeping my score.", "rebuttal": "We thank you for your constructive comments on our work. We provide details to address your concerns (weaknesses W1–W4):\n \n## W1: Novelty of FlashDMoE\nWe acknowledge that simple kernel fusion is a straightforward task. However, achieving fast, distributed kernel fusion is a significant challenge, particularly for MoE models, which expert researchers characterize as \"hard to train\" [1]. FlashDMoE is the first to introduce a fully fused, high-performance, systematic distributed MoE operator; to the best of our knowledge, previous systems [2, 3] have only achieved partial fusion through ad-hoc methods.\n\nOn one hand, efficient kernel fusion requires an in-depth understanding of the GPU memory hierarchy, its weak memory model, and extensive experience with its programming abstractions, particularly multidimensional thread-value layouts for coalesced access and the Matrix-Multiply-Accumulate (MMA) instructions provided by specialized libraries such as CUTLASS. On the other hand, implementing these advanced techniques within a distributed framework constitutes the core innovation of our work. Specifically, we present (1) an in-kernel actor-based operating system that enables high-throughput distributed tile parallelism and (2) a symmetric tensor layout that facilitates completely non-blocking and provably conflict-free memory access across distributed GPUs and tens of thousands of GPU threads.\n\n## W2: Memory Footprint\nWe acknowledge the concern that FlashDMoE requires extra memory for the symmetric tensor (Section 3.2). We maintain that this overhead is modest in comparison with the existing memory requirements of MoE models. To show this empirically, we compute and analyze memory costs for popular large MoE models:\n\n| **Model**          | **Params** | **S** | **E** | **H** | **Intermediate** | **Symmetric Tensor (GB)** | **Memory for Model (GB)** | **Overhead Ratio (%)** |\n|--------------------|------------|-------|-------|-------|------------------|---------------------------|---------------------------|------------------------|\n| Moonlight-16B-A3B  | 16B        | 8K    | 64    | 2K    | 1.375K           | 0.25                      | 59                        | **0.49**               |\n| Grok-1             | 314B       | 8K    | 8     | 6K    | 32K              | 0.75                      | 1169                      | **0.15**               |\n| Snowflake-Arctic   | 479B       | 4K    | 128   | 7K    | 4.75K            | 1.75                      | 1784                      | **0.12**               |\n| Qwen3-235B-A22B    | 235B       | 40K   | 128   | 4K    | 1.5K             | 3                         | 875                       | **0.38**               |\n| Mixtral 8x7B       | 47B        | 32K   | 8     | 4K    | 14K              | 2                         | 175                       | **2.15**               |\n| DeepSeek-V3        | 685B       | 160K  | 256   | 7K    | 2K               | 1.5                       | 2551                      | **0.11**               |\n\n\nWe calculate the model’s memory requirement by multiplying its parameter count with precision (fp32 here but this choice is inconsequential to the ratio which we embolden). Note that the symmetric tensor demands at most 2% extra memory in comparison to the inference memory requirements. This overhead would even be more amortized in training, where the memory demand for the model is much greater. We will be sure to add a more thorough empirical analysis elaborating on the memory and performance tradeoffs exhibited in FlashDMoE.\n\n## W3: Precision\nWe acknowledge the concern regarding FlashDMoE’s current use of FP32 precision. We note that our evaluation in the paper using FP32 represents a conservative estimate of FlashDMoE’s potential benefits. Despite operating in full precision, FlashDMoE substantially outperforms baseline systems that already utilize FP16 precision. Thus, transitioning FlashDMoE to FP16 precision will further improve these already significant performance gains. \n\nOn another note, we acknowledge the use of fp32 as a limitation in our paper and we highlight that FlashDMoE already supports half precision in bf16 and fp16. However, our fp16 implementations do not yet achieve peak hardware utilization due to (1) suboptimal shared memory layouts and (2) insufficient vectorization in custom computational operators within the Processor. To overcome these limitations, we plan to integrate recent advancements from Device Extensions (DX) libraries [4]. These libraries build on top of the CUTLASS APIs and provide a specialized configuration layer with highly optimized parameters, significantly improving both hardware memory throughput and tensor-core utilization. We anticipate adopting [4] in FlashDMoE and, in fact, already have ongoing work to that effect. In summary, this transition would involve deprecating the CollectiveMainLoop defined in `mmaConfig.cuh` and line 80 of `gemm.cuh`, in favor of a bespoke GEMM main loop that accumulates across the K dimension, where each iteration dispatches to the highly-tuned, block-scoped operator of [4]. Importantly, [4]’s adherence to returning GEMM results as fragments stored in fast register memory seamlessly preserves FlashDMoE’s ability to execute post-GEMM fusion operations. We plan to make these changes available in our public codebase.\n\n\n## W4: Diverse Expert Routing\nFlashDMoE already supports unbounded k for expert top-k selection and can accommodate up to 65,535 experts—significantly exceeding the architectural specifications of any state-of-the-art MoE model and the memory requirements of modern GPUs. \n\n### Typos\nWe appreciate you identifying the typos! We will fix them.  \n\n## References\n\n[1] Mosaic Research. _Introducing DBRX: A New State-of-the-Art Open LLM_. 2024.\n\n[2] Zhang et al. _COMET_. MLSys ‘25, Outstanding Paper Honorable Mention.\n\n[3] DeepSeek-AI. _DeepEP_. 2025.\n\n[4] NVIDIA. cuBLASDx."}, {"review_id": "We1v1MXDOL", "reviewer": "Reviewer_Kbkq", "summary": "This paper introduces FlashDMoE, a novel distributed MoE operator that fuses computation and communication into a single persistent GPU kernel. The work is motivated by the growing inefficiencies in conventional MoE models where distributed experts incur high communication overheads due to frequent kernel launches and synchronous communication operations. The paper addresses the challenges by employing device-initiated asynchronous communication, fine-grained tile-level parallelism, and an actor-based scheduling model, resulting in significant speedups (up to 15×) over existing state-of-the-art MoE implementations.", "questions": "See weaknesses.", "limitations": "Yes.", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\n1. The approach introduces a single persistent GPU kernel that fuses all computation and communication tasks. It successfully eliminates repeated kernel launch overhead and reduces idle time on GPUs. It achieves higher GPU utilization and demonstrably faster execution compared to multiple short-lived kernel launches in existing systems.\n\n2. The proposed method leverages device-initiated asynchronous communication to mitigate the straggler problem. It efficiently routes only active tokens to the GPUs hosting the relevant experts. It minimizes excessive data transfer by implementing payload-efficient communication.\n\n3. The design employs fine-grained tile-level parallelism and an actor-based concurrency model. It effectively partitions the workload into small, independent tasks that can be dynamically scheduled across GPU warps. It optimizes the balance between computation and administrative tasks for better overall performance.\n\n4. The experimental evaluations demonstrate strong performance improvements across multiple dimensions. It shows substantial speedups when scaling tokens, experts, and GPUs compared to existing MoE implementations. It substantiates the claims with detailed profiling and comparative analyses.\n\n\nWeaknesses:\n\n\n1. The current implementation targets FP32 precision exclusively. It does not explore or support lower-precision arithmetic such as FP16 or mixed precision, which are common in modern deep learning deployments. This limitation may impact its applicability and performance in production-scale environments where lower precision is preferred.\n\n2. While optimizing the efficiency of MoE at the CUDA level is helpful and practical, from the system perspective, the contribution of kernel fusion is relatively weak. A simple kernel fusion is trivial and might be not enough for NeurIPS.", "was_revised": false, "final_justification": "", "rebuttal": "We thank reviewer Kbkq for evaluating our paper and providing constructive feedback. Here, we aim to address the points raised regarding FP32 precision of FlashMoE and provide a more in-depth understanding of the novelty of FlashDMoE. \n\n## W1: Precision\nWe acknowledge the concern regarding FlashDMoE’s current use of FP32 precision. We note that our evaluation in the paper using FP32 represents a conservative estimate of FlashDMoE’s potential benefits. Despite operating in full precision, FlashDMoE substantially outperforms baseline systems that already utilize FP16 precision. Thus, transitioning FlashDMoE to FP16 precision will further improve these already significant performance gains. \n\nOn another note, we acknowledge the use of fp32 as a limitation in our paper and we highlight that FlashDMoE already supports half precision in bf16 and fp16. However, our fp16 implementations do not yet achieve peak hardware utilization due to (1) suboptimal shared memory layouts and (2) insufficient vectorization in custom computational operators within the Processor. To overcome these limitations, we plan to integrate recent advancements from Device Extensions (DX) libraries [1]. These libraries build on top of the CUTLASS APIs and provide a specialized configuration layer with highly optimized parameters, significantly improving both hardware memory throughput and tensor-core utilization. We anticipate adopting [1] in FlashDMoE and, in fact, already have ongoing work to that effect. In summary, this transition would involve deprecating the CollectiveMainLoop defined in `mmaConfig.cuh` and line 80 of `gemm.cuh`, in favor of a bespoke GEMM main loop that accumulates across the K dimension, where each iteration dispatches to the highly-tuned, block-scoped operator of [1]. Importantly, [1]’s adherence to returning GEMM results as fragments stored in fast register memory seamlessly preserves FlashDMoE’s ability to execute post-GEMM fusion operations. We plan to make these changes available in our public codebase.\n\n## W2: Novelty\nWe acknowledge that _simple_ kernel fusion is a straightforward task. However, achieving _fast, distributed_ kernel fusion is a significant challenge, particularly for MoE models, which expert researchers characterize as \"hard to train\" [2]. FlashDMoE is the first to introduce a fully fused, high-performance, systematic distributed MoE operator; to the best of our knowledge, previous systems [3, 4] have only achieved partial fusion through ad-hoc methods.\n\nOn one hand, efficient kernel fusion requires an in-depth understanding of the GPU memory hierarchy, its weak memory model, and extensive experience with its programming abstractions, particularly multidimensional thread-value layouts for coalesced access and the Matrix-Multiply-Accumulate (MMA) instructions provided by specialized libraries such as CUTLASS. On the other hand, implementing these advanced techniques within a distributed framework constitutes the core innovation of our work. Specifically, we present (1) an in-kernel actor-based operating system that enables high-throughput distributed tile parallelism and (2) a symmetric tensor layout that facilitates completely non-blocking and provably conflict-free memory access across distributed GPUs and tens of thousands of GPU threads.\n\nThus, the novelty of our contributions, combined with their superior performance compared to baseline works [3, 5] from prestigious venues, positions FlashDMoE, we believe, as an excellent fit for NeurIPS.\n\n## References\n\n[1] NVIDIA. cuBLASDx.\n\n[2] Mosaic Research. _Introducing DBRX: A New State-of-the-Art Open LLM_. 2024.\n\n[3] Zhang et al. _COMET_. MLSys ‘25, Outstanding Paper Honorable Mention.\n\n[4] DeepSeek-AI. _DeepEP_. 2025.\n\n[5] He et al. _FasterMoE_, PPopp’22."}]

Author rebuttal:

Meta-review: **1. Scientific Claims and Findings**

This paper presents FlashDMoE, a novel system for accelerating distributed Mixture-of-Experts (MoE) layers by addressing two critical bottlenecks in existing implementations: high communication overhead from synchronous "ALLTOALL" operations and latency from frequent GPU kernel launches. The core claim is that fusing the entire MoE workflow—including gating, expert computation, and inter-GPU communication—into a single, persistent GPU kernel eliminates these bottlenecks. Key technical innovations include an actor-based concurrency model (specializing GPU warps into roles like Processor and Scheduler), a symmetric tensor layout enabling asynchronous, device-initiated communication, and "in-place padding" for payload-efficient data transfer. Evaluations on an 8x H100 GPU server demonstrate 1.2x to 15x forward-pass latency speedups over baselines (e.g., DeepEP, Megatron-LM) with >90% GPU utilization.

**2. Major Strengths**

The single-kernel design rethinks MoE execution, eliminating kernel launch overhead and enabling fine-grained computation-communication overlap—an elegant departure from conventional multi-kernel approaches. It reports significant speedups (up to 15x) and high GPU utilization (>90%) across diverse scenarios (scaling tokens, experts, GPUs). The symmetric tensor layout and asynchronous, device-initiated communication effectively mitigate straggler effects and reduce synchronization penalties, with "in-place padding" minimizing unnecessary data transfer. Comprehensive benchmarks against strong baselines, including micro-analyses of GPU utilization and payload efficiency, strengthen the credibility of performance claims.

**3. Major Weaknesses**

The 4x memory cost of the symmetric tensor layout may be prohibitive for large sequences or memory-constrained hardware, with insufficient analysis of this trade-off. Evaluations are restricted to single-node, 8x H100 setups with NVLink; multi-node performance over limited-bandwidth networks remains unproven, leaving scalability claims partially unsubstantiated. The restriction to data precision of FP32 (vs. low-precision formats like FP16/BF16) limits applicability to modern LLM training/inference workflows.

**4. Main reasons of accept**

The paper is recommended to be accepted due to its novel single-kernel design for MoE layers that effectively eliminates critical bottlenecks, delivers significant performance speedups, and introduces innovative communication and concurrency solutions supported by rigorous evaluations.

**5. Summary of discussions**

Most of the reviewers are satisfactory with the rebuttal. Reviewer Kbkq who gave rating of 3 had concerns about the precision and novelty of the work, and was not convinced by the rebuttal. By considering all the review comments and discussions, I am more inclined to recommend acceptance.
Decision: Accept (poster)