UniEP: Unified Expert-Parallel MoE MegaKernel for LLM Training

2026-04-21

AI Summary

This work addresses key challenges in large-scale Mixture-of-Experts (MoE) training with Expert Parallelism (EP), including communication bottlenecks, the absence of a unified optimization framework, and numerical instability. To overcome these limitations, the authors propose MegaKernel, a novel approach that integrates EP communication and computation into a single parameterized operator, enabling automatic adaptation of diverse optimization strategies. Additionally, a deterministic token routing mechanism is introduced to guarantee numerical consistency and execution determinism during training. This method establishes, for the first time, a searchable parameter space for EP optimizations. Evaluated on NVIDIA Hopper GPU clusters, MegaKernel achieves speedups of 1.03–1.38× over baseline implementations while maintaining production-level model accuracy and substantially alleviating communication bottlenecks.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Search and Optimization: Learning to Search
• Natural Language Processing: Learning & Optimization for NLP

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation

Abstract

The exponential growth in Large Language Model (LLM) parameters has transformed model training into an increasingly resource-intensive endeavor. With the stagnation of Moore's Law and the widening disparity between computation throughput and communication bandwidth, expert parallelism (EP) has emerged as a critical strategy for scaling mixture-of-experts (MoE) models. However, despite numerous proposals for optimizing EP, ranging from communication compression to computation-communication overlap, adoption within production-grade frameworks like Megatron-LM remains conservative. Existing solutions often rely on ad-hoc, complex kernels that lack adaptability across diverse optimization configurations and frequently neglect numerical stability, failing to meet the strict precision requirements of large-scale training. In this paper, we introduce UniEP, a novel system that unifies diverse EP optimization strategies into a cohesive abstraction. UniEP fuses the MoE communication and computation into MegaKernels, effectively transforming complex architectural tuning into a unified parameter search space for automated adaptability. Crucially, UniEP incorporates a deterministic token ordering mechanism that guarantees numerical consistency with sequential execution, even under aggressive overlap schedules. We evaluate UniEP on GPU clusters equipped with NVIDIA Hopper GPUs. Our results demonstrate that UniEP achieves 1.03$\times$-1.38$\times$ speedups over state-of-the-art work, effectively mitigating communication bottlenecks while maintaining the rigorous accuracy standards required for production LLM training.

Problem

• Expert Parallelism
• Mixture-of-Experts
• Numerical Stability
• Large Language Model Training
• Communication Bottleneck

Innovation

• Expert Parallelism
• Mixture-of-Experts
• MegaKernel
• Numerical Stability
• Communication-Computation Overlap