# UniEP: Unified Expert-Parallel MoE MegaKernel for LLM Training arXiv: https://arxiv.org/abs/2604.19241 · HTML: https://mlsanity.com/paper/2604.19241/ ## Key facts - Cited by 5 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 25 of its 42 references that are in the map. - First posted on arXiv 21 Apr 2026. ## 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, 25 of 25) - [Qwen3 Technical Report](https://mlsanity.com/paper/2505.09388/) (2025) - [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023) - [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022) - [DeepSeek-VL: Towards Real-World Vision-Language Understanding](https://mlsanity.com/paper/2403.05525/) (2024) - [Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM](https://mlsanity.com/paper/2104.04473/) (2021) - [Kimi K2: Open Agentic Intelligence](https://mlsanity.com/paper/2507.20534/) (2025) - [DeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training to Power Next-Generation AI Scale](https://mlsanity.com/paper/2201.05596/) (2022) - [FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving](https://mlsanity.com/paper/2501.01005/) (2025) - [MegaBlocks: Efficient Sparse Training with Mixture-of-Experts](https://mlsanity.com/paper/2211.15841/) (2022) - [Kimi Linear: An Expressive, Efficient Attention Architecture](https://mlsanity.com/paper/2510.26692/) (2025) - [TorchTitan: One-stop PyTorch native solution for production ready LLM pre-training](https://mlsanity.com/paper/2410.06511/) (2024) - [FLUX: Fast Software-based Communication Overlap On GPUs Through Kernel Fusion](https://mlsanity.com/paper/2406.06858/) (2024) - [Comet: Fine-grained Computation-communication Overlapping for Mixture-of-Experts](https://mlsanity.com/paper/2502.19811/) (2025) - [Mirage: A Multi-Level Superoptimizer for Tensor Programs](https://mlsanity.com/paper/2405.05751/) (2024) - [Breaking the Computation and Communication Abstraction Barrier in Distributed Machine Learning Workloads](https://mlsanity.com/paper/2105.05720/) (2021) - [TACCL: Guiding Collective Algorithm Synthesis using Communication Sketches](https://mlsanity.com/paper/2111.04867/) (2021) - [MegaScale-Infer: Serving Mixture-of-Experts at Scale with Disaggregated Expert Parallelism](https://mlsanity.com/paper/2504.02263/) (2025) - [MegaScale-MoE: Large-Scale Communication-Efficient Training of Mixture-of-Experts Models in Production](https://mlsanity.com/paper/2505.11432/) (2025) - [Synthesizing Optimal Collective Algorithms](https://mlsanity.com/paper/2008.08708/) (2020) - [Triton-distributed: Programming Overlapping Kernels on Distributed AI Systems with the Triton Compiler](https://mlsanity.com/paper/2504.19442/) (2025) - [TileLink: Generating Efficient Compute-Communication Overlapping Kernels using Tile-Centric Primitives](https://mlsanity.com/paper/2503.20313/) (2025) - [GC3: An Optimizing Compiler for GPU Collective Communication](https://mlsanity.com/paper/2201.11840/) (2022) - [HeterMoE: Efficient Training of Mixture-of-Experts Models on Heterogeneous GPUs](https://mlsanity.com/paper/2504.03871/) (2025) - [HybridEP: Scaling Expert Parallelism to Cross-Datacenter Scenario via Hybrid Expert/Data Transmission](https://mlsanity.com/paper/2510.19470/) (2025) - [MoEBlaze: Breaking the Memory Wall for Efficient MoE Training on Modern GPUs](https://mlsanity.com/paper/2601.05296/) (2026) ## Built on by (most cited first, 5 of 5) - [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) - [HyperParallel-MoE: Multi-Core Interleaved Scheduling for Fast MoE Training on Ascend NPUs](https://mlsanity.com/paper/2605.23764/) (2026) - [Route-Block Membership Selects Packed-AWQ Arithmetic: A Controlled Single-Fixture Mechanism Study](https://mlsanity.com/paper/2607.26316/) (2026) - [Compute-Optimal Is Not Cluster-Optimal: Systems-Aware Scaling for Sparse Mixture-of-Experts](https://mlsanity.com/paper/2608.10605/) (2026) - [Weave: Fine-Grained Dynamic SM Scheduling in an MoE Megakernel for Compute-Communication Overlap](https://mlsanity.com/paper/2609.21483/) (2026) ## Abstract (arXiv) 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×-1.38× speedups over state-of-the-art work, effectively mitigating communication bottlenecks while maintaining the rigorous accuracy standards required for production LLM training.