Analytical resource manager boosts MoE inference on GPUs
Analytical Resource Management for Fine-grained MoE Computation-Communication Overlap
Distributed, Parallel, and Cluster Computing
Summary
Large AI models called Mixture-of-Experts split tasks among many smaller parts, which need both computing and talking to each other on GPUs. The authors found that managing the balance between computing and communication better can speed up these models. They created a model that decides how many tasks run simultaneously without trial runs or extra tuning, making the process faster and more efficient. Testing on NVIDIA GPUs showed significant speed improvements compared to previous methods.
What this means in practice
- •For machine learning engineers: Optimize GPU resource allocation to speed up large Mixture-of-Experts model inference without extra profiling or code changes.
- •For cloud infrastructure teams: Improve cloud GPU usage efficiency for serving AI models with fine-grained control over computation and communication tasks.
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.