UBEP replaces BSP All-to-All for MoE on multi-tier superpods with dependency-driven kernel decomposition, topology-aware token scheduling, and Data-as-Flag atomics, cutting All-to-All latency up to 52.4% and TPOT up to 11.1%.
UniEP: Unified Expert-Parallel MoE MegaKernel for LLM Training
2 Pith papers cite this work. Polarity classification is still indexing.
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.
fields
cs.DC 2years
2026 2representative citing papers
HyperParallel-MoE reduces Dispatch-to-Combine MoE-FFN latency by up to 1.58x on Ascend A3 clusters via tile-level heterogeneous scheduling that overlaps communication, matrix, and vector computation inside a single kernel launch.
citing papers explorer
-
UBEP: Re-architecting Expert Parallelism Communication Library for Production Superpods
UBEP replaces BSP All-to-All for MoE on multi-tier superpods with dependency-driven kernel decomposition, topology-aware token scheduling, and Data-as-Flag atomics, cutting All-to-All latency up to 52.4% and TPOT up to 11.1%.
-
HyperParallel-MoE: Multi-Core Interleaved Scheduling for Fast MoE Training on Ascend NPUs
HyperParallel-MoE reduces Dispatch-to-Combine MoE-FFN latency by up to 1.58x on Ascend A3 clusters via tile-level heterogeneous scheduling that overlaps communication, matrix, and vector computation inside a single kernel launch.