Weave boosts MoE GPU efficiency with fine-grained dynamic scheduling
Weave: Fine-Grained Dynamic SM Scheduling in an MoE Megakernel for Compute-Communication Overlap
Distributed, Parallel, and Cluster Computing
Summary
Running complex AI models called Mixture-of-Experts (MoE) on GPUs is slow because the tasks need lots of communication and computing, which can leave parts of the GPU idle. The authors developed Weave, a system that dynamically decides how to split GPU resources for communication and computation based on the workload in real time. This improves how well the GPU is used and speeds up processing by nearly three times on average for MoE layers. Weave’s smart scheduling reduces waiting times and wasted resources during AI model inference.
What this means in practice
- •For machine learning engineers: Accelerate inference of large MoE models on multi-GPU systems by dynamically optimizing GPU resource allocation.
- •For data center operators: Improve GPU utilization and throughput in data centers running AI workloads with MoE models through fine-grained scheduling.
Authors
Ziyu Huang, Yangjie Zhou, Chenhao Zhu, Zihan Liu, Jinyu Liu, Shulai Zhang, Xingxun Tang, Hongzhe Yan, Xinhao Luo, Minyi Guo, Xiu Lin, Yinghao Yu, Guodong Yang, Liping Zhang, Shixuan Sun, Jingwen Leng
Abstract
Mixture-of-Experts (MoE) inference under expert parallelism (EP) turns each MoE layer into a distributed computation with costly dispatch and combine communication. State-of-the-art systems reduce this cost through communication-computation overlap, splitting the GPU's SMs for communication and computation respectively. However, this approach still leaves GPU resources wasted along two dimensions. Spatially, the best SM split is determined by each layer's routing result and varies across layers and GPUs, so fixed policies mismatch the workload and waste either NVLink bandwidth or compute throughput. Temporally, complex MoE data dependencies introduce bubbles that leave SMs idle. We present Weave, to our knowledge the first MoE overlap system that performs fine-grained dynamic SM scheduling - deciding per layer and per GPU by routing results at runtime. Once routing completes, each layer's communication and computation volumes become known; Weave exploits this predictability through a lightweight cost model running inside the persistent megakernel: a spatial scheduler partitions SMs into communication workers and computation workers to match the communication/computation throughput ratio, and a temporal scheduler coordinates the two worker groups to minimize SM idleness. On 4x H100 SXM GPUs across six mainstream MoE models, Weave achieves a 2.89x geometric-mean MoE-layer speedup and a 1.33x geometric-mean end-to-end speedup over five state-of-the-art baselines.