Fast Distributed MoE in a Single Kernel
Osayamen Aimuyo, Byungsoo Oh, Rachee Singh
The computational sparsity of Mixture-of-Experts (MoE) models allows them to scale efficiently, achieving sub-linear growth in computational costs as model sizes increase. However, distributing the experts across multiple GPUs introduces significant communication overheads, undermining some of their computational benefits. These overheads arise from frequent GPU kernel launches and synchronous communication across distributed experts. To overcome these limitations, we develop FlashDMoE, a novel MoE operator that fuses all computation and communication tasks into a single persistent GPU kernel. FlashDMoE employs device-initiated asynchronous communication and payload-efficient data transfers to dynamically route computation only to GPUs hosting actively selected experts. By leveraging fine-grained, tile-level parallelism within the kernel, FlashDMoE mitigates performance degradation caused by communication stragglers. Our evaluation on multiple ML servers, each with 8 NVIDIA H100 GPUs, demonstrates that FlashDMoE achieves 1.25-15X speedup over state-of-the-art MoE implementations.