Research map / Architectures and efficiency / Mixture of experts
FlashMoE: Fast Distributed MoE in a Single Kernel
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
FlashMoE fuses the entire distributed MoE operator (gate, dispatch, expert FFN, combine) into a single persistent GPU kernel using device-initiated one-sided NVSHMEM communication and a symmetric tensor layout.
Problem
Distributed MoE execution suffers from low GPU utilization, latency overhead, and straggler effects due to CPU-managed scheduling, host-initiated synchronous all-to-all communication, and many kernel launches per layer.
Approach
Replace multi-kernel, host-orchestrated MoE pipelines with one persistent kernel in which warp-specialized actors (processors, a scheduler, subscribers) schedule tile-level FFN/combine tasks and perform asynchronous device-initiated (R)DMA transfers over a conflict-free symmetric tensor layout, with in-place padding to avoid sending null tokens.
Key results
- On 8 H100s with 16K tokens, FlashMoE achieves up to 6.4x speedup over baselines in forward latency and 17.7 MTokens/s throughput, over 5.7x higher than FasterMoE.
- FlashMoE reaches 93.17% average SM utilization versus 9.67% for FasterMoE, 13.55% for DeepEP+Megatron-LM, 59.11% for Megatron-TE and 42.31% for Comet.
- FlashMoE launches a single kernel per MoE layer versus 33 to 550 GPU operations for baselines, and achieves up to 4x better overlap efficiency at 8 GPUs.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Fuse MoE computation and communication into one persistent GPU kernel with in-kernel scheduling and asynchronous device-initiated transfers | FlashMoE · MoE forward latency, 8 H100s | speedup over baselines: 6.4x | Helped |
| Fuse MoE computation and communication into one persistent GPU kernel with in-kernel scheduling and asynchronous device-initiated transfers | FlashMoE · MoE forward pass, 8 H100s | SM utilization: 93.17% compared to 9.67% for FasterMoE | Helped |
| Fuse MoE computation and communication into one persistent GPU kernel with in-kernel scheduling and asynchronous device-initiated transfers | FlashMoE · MoE throughput, 8 H100s | throughput: 17.7 MTokens/s; over 5.7x higher than FasterMoE | Helped |
| Use FlashMoE's actor-based design and asynchronous data movement | FlashMoE · MoE weak scaling, 8 GPUs | overlap efficiency: 4x higher efficiency | Helped |
| Use payload-efficient communication and scheduler-driven in-kernel dispatching | FlashMoE · MoE forward latency, 8 H100s, 128 experts | speedup over baselines: 6.6X | Helped |
Surprising findings
- FlashMoE runs in FP32 while baselines use FP16, yet still achieves higher throughput; its FP16 path is suboptimal due to shared memory layout issues.
Limitations
- Engineering complexity of fully fused persistent kernels requires deep GPU and distributed-systems expertise.
- FP16 path is suboptimal due to insufficient tuning and shared memory layout inefficiency.
- Targets inference only; training (backward pass) is not supported.
Experimental setup
Metrics: forward latency SM utilization throughput (MTokens/s) overlap efficiency expert scalability kernel launch count
Baselines: Comet FasterMoE Megatron-CUTLASS Megatron-TE DeepEP DeepSpeed-MoE
Builds on: NVSHMEM CUTLASS actor model
Closest papers by meaning
- MoE++: Accelerating Mixture-of-Experts Methods with Zero-Computation Experts (ICLR 2025 · similarity 0.90)
- Scaling Beyond the GPU Memory Limit for Large Mixture-of-Experts Model Training (ICML 2024 · similarity 0.90)
- Shortcut-connected Expert Parallelism for Accelerating Mixture of Experts (ICML 2025 · similarity 0.89)
- FlowMoE: A Scalable Pipeline Scheduling Framework for Distributed Mixture-of-Experts Training (NeurIPS 2025 · similarity 0.89)
- SpeedLoader: An I/O efficient scheme for heterogeneous and distributed LLM operation (NeurIPS 2024 · similarity 0.87)
- LSH-MoE: Communication-efficient MoE Training via Locality-Sensitive Hashing (NeurIPS 2024 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.