Improving speed and efficiency of Mixture-of-Experts models on GPUs
Analytical Resource Management for Fine-grained MoE Computation-Communication Overlap
Distributed, Parallel, and Cluster Computing
Summary
Mixture-of-Experts (MoE) models use multiple specialized parts to make decisions, but running them on GPUs can be slow due to limits on how the computer balances work and communication. The paper presents a way to carefully plan how tasks are divided and timed on GPUs to make sure communication happens alongside computation without delays. The authors’ method adjusts resource use dynamically without needing extra trial runs, speeding up inference significantly compared to previous approaches. This approach works well on multiple models and GPU setups, showing better performance across various tasks.
Mixture-of-Experts (MoE)GPUstreaming multiprocessor (SM)cooperative thread array (CTA)computation-communication overlapkernel occupancyCOMETGEMM operatorresource managementparallelism
Authors
Hongyu Liu, Minyu Cui, Miquel Pericas
Abstract
Fine-grained computation--communication overlap in distributed Mixture-of-Experts (MoE) inference allows communication to begin as partial compute results become ready. However, cooperative thread arrays (CTAs) performing computation and communication contend for finite residency capacity on streaming multiprocessors (SMs). Because a resident CTA generally retains its allocated SM resources until completion, CTAs that cannot be co-resident must wait for resources, resulting in wave-like execution. A fixed resource partition cannot adapt to changes in input size, routed expert load, and kernel configuration, potentially causing a communication backlog or reducing expert compute parallelism. We present a wave-quantized analytical model and launch-time resource manager for dependency-coupled overlap pipelines. Using routed-tile counts, kernel occupancy, GPU residency constraints, and split-level readiness dependencies, it selects the communication-CTA count and resource partition before each launch without candidate execution, per-workload profiling, or kernel recompilation. We integrate the method into the public COMET A100 implementation in FLUX. We evaluate three MoE models on four NVIDIA A100 GPUs under several parallelism strategies at the GEMM2+GatherRS operator, complete post-router MoE layer, and complete-model prefill levels. Across 15 real-p90 workloads, the analytical selector achieves 3.22 percent mean regret relative to the measured oracle with a mean solver overhead of 0.157 microseconds. Over COMET, our method achieves geometric-mean speedups of 2.528x at the GEMM2+GatherRS operator, 1.771x at the complete post-router MoE layer, and 1.185x for complete-model prefill, with maxima of 4.218x, 2.584x, and 1.439x, respectively. At every feasible TP=2/EP=2 sequence length of at least 4,096, our implementation outperforms COMET, Megatron core-TE, and FastMoE TP+NCCL.