# FlashMoE: Fast Distributed MoE in a Single Kernel

arXiv: https://arxiv.org/abs/2506.04667 · HTML: https://mlsanity.com/paper/2506.04667/

## Key facts

- Cited by 15 papers in the map (data through 14 Oct 2026).
- Placed in Efficient MoE training and inference systems, one of 12 directions in Mixture-of-experts.
- Builds on 19 of its 44 references that are in the map.
- First posted on arXiv 5 Jun 2025.

## Placed in

- Mixture-of-experts / [Efficient MoE training and inference systems](https://mlsanity.com/moe/efficient-moe-training-and-inference-systems/)

## Builds on (most cited first, 19 of 19)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Qwen3 Technical Report](https://mlsanity.com/paper/2505.09388/) (2025)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Mixtral of Experts](https://mlsanity.com/paper/2401.04088/) (2024)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM](https://mlsanity.com/paper/2104.04473/) (2021)
- [Muon is Scalable for LLM Training](https://mlsanity.com/paper/2502.16982/) (2025)
- [DeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training to Power Next-Generation AI Scale](https://mlsanity.com/paper/2201.05596/) (2022)
- [Tutel: Adaptive Mixture-of-Experts at Scale](https://mlsanity.com/paper/2206.03382/) (2022)
- [MegaScale: Scaling Large Language Model Training to More Than 10,000 GPUs](https://mlsanity.com/paper/2402.15627/) (2024)
- [FlexMoE: Scaling Large-scale Sparse Pre-trained Model Training via Dynamic Device Placement](https://mlsanity.com/paper/2304.03946/) (2023)
- [Comet: Fine-grained Computation-communication Overlapping for Mixture-of-Experts](https://mlsanity.com/paper/2502.19811/) (2025)
- [A Hybrid Tensor-Expert-Data Parallelism Approach to Optimize Mixture-of-Experts Training](https://mlsanity.com/paper/2303.06318/) (2023)
- [Breaking the Computation and Communication Abstraction Barrier in Distributed Machine Learning Workloads](https://mlsanity.com/paper/2105.05720/) (2021)
- [BlackMamba: Mixture of Experts for State-Space Models](https://mlsanity.com/paper/2402.01771/) (2024)
- [Lancet: Accelerating Mixture-of-Experts Training via Whole Graph Computation-Communication Overlapping](https://mlsanity.com/paper/2404.19429/) (2024)
- [CO2: Efficient Distributed Training with Full Communication-Computation Overlap](https://mlsanity.com/paper/2401.16265/) (2024)
- [T3: Transparent Tracking & Triggering for Fine-grained Overlap of Compute & Collectives](https://mlsanity.com/paper/2401.16677/) (2024)
- [Optimizing Distributed ML Communication with Fused Computation-Collective Operations](https://mlsanity.com/paper/2305.06942/) (2023)

## Built on by (most cited first, 15 of 15)

- [ParallelKittens: Systematic and Practical Simplification of Multi-GPU AI Kernels](https://mlsanity.com/paper/2511.13940/) (2025)
- [Event Tensor: A Unified Abstraction for Compiling Dynamic Megakernel](https://mlsanity.com/paper/2604.13327/) (2026)
- [L3: Large Lookup Layers](https://mlsanity.com/paper/2601.21461/) (2026)
- [TaxBreak: Unmasking the Hidden Costs of LLM Inference Through Overhead Decomposition](https://mlsanity.com/paper/2603.12465/) (2026)
- [Efficient Mixture-of-Experts LLM Inference with Apple Silicon NPUs](https://mlsanity.com/paper/2604.18788/) (2026)
- [Searching for Plans You Can Actually Build: A Realizability-Aware Full-Space Optimizer for MoE Training and Serving](https://mlsanity.com/paper/2607.18631/) (2026)
- [Fine-grained Computation-Communication Overlap via Tile-level Signaling and Scheduling for Mixture-of-Experts](https://mlsanity.com/paper/2607.19539/) (2026)
- [MoE-Hub: Taming Software Complexity for Seamless MoE Overlap with Hardware-Accelerated Communication on Multi-GPU Systems](https://mlsanity.com/paper/2605.05888/) (2026)
- [Federation of Experts: Communication Efficient Distributed Inference for Large Language Models](https://mlsanity.com/paper/2605.06206/) (2026)
- [MonoMoE: An Efficient Fused Mega-kernel for Quantized MoE Decoding](https://mlsanity.com/paper/2609.04244/) (2026)
- [Analytical Resource Management for Fine-grained MoE Computation-Communication Overlap](https://mlsanity.com/paper/2609.07536/) (2026)
- [Mixture-of-Kittens: MoE Megakernel for NVL72s](https://mlsanity.com/paper/2609.36070/) (2026)
- [Purlin: Separating Orchestration from the Datapath of Collectives](https://mlsanity.com/paper/2609.36954/) (2026)
- [Trident: Unifying Guarded Dispatch and Host Execution for PyTorch Triton Workloads](https://mlsanity.com/paper/2609.37241/) (2026)
- [RESOLVE: Language-Agnostic Validation of GPU Kernels Through Testing, Reduction, and Proof](https://mlsanity.com/paper/2610.05683/) (2026)

## Abstract (arXiv)

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 eliminates bulk-synchronous collectives for one-sided, device-initiated, inter-GPU (R)DMA transfers, thereby unlocking payload efficiency by eliminating bloated or redundant network payloads in sparsely activated layers. When evaluated on an 8-H100 GPU node with MoE models comprising up to 128 experts and 16K token sequences, FlashMoE achieves up to 9x higher GPU utilization, 6x lower latency, 5.7x higher throughput, and 4x better overlap efficiency compared to state-of-the-art baselines, despite using FP32, whereas the 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.