🤖 AI Summary
This work addresses the performance bottleneck in Mixture-of-Experts (MoE) models during multi-GPU deployment, where serial execution of expert computation and all-to-all communication exposes communication latency on the critical path, limiting GPU utilization. The authors propose a producer-consumer co-design that employs block-level scheduling: persistent compute kernels prioritize processing remote critical blocks, while dedicated streaming multiprocessor (SM) partitions run persistent communication kernels that initiate fine-grained communication based on block readiness. This enables efficient overlap between computation and return-phase communication without modifying underlying operators or communication primitives. Evaluated on a 4×A100 platform, the approach achieves up to 2.74× speedup for MoE layers and 2.64× end-to-end acceleration, demonstrating consistent performance gains and correctness across diverse GEMM shapes and routing strategies.
📝 Abstract
Mixture-of-Experts (MoE) architectures increase model capacity without proportionally increasing computation cost and have become a key building block for scaling large language models (LLMs) to trillion-parameter regimes. Efficient deployment of these MoE models relies on distributed execution across multiple GPUs, where each MoE layer involves two all-to-all communications: dispatching tokens to expert ranks and returning the expert outputs to their source ranks. Conventional MoE implementations launch this return all-to-all after expert compute completes, exposing communication latency on the critical path and reducing GPU utilization. We present a fine-grained approach that overlaps expert compute with the second all-to-all via tile-level signaling and scheduling. Our producer-consumer co-design combines: (1) a persistent per-rank computation kernel (producer) that covers all local experts on the rank to eliminate repeated kernel launch overhead and prioritizes remote-critical tiles, and (2) a persistent communication kernel (consumer) on a small dedicated partition of streaming multiprocessors (SMs) that issues segment-granular transfers as tiles become ready. Our co-design avoids intrusive changes to the underlying computation operators or communication primitives, making it practical for improving distributed MoE execution efficiency on multi-GPU systems. On a 4-A100 GPU platform, evaluated on three MoE models against four state-of-the-art MoE systems, our approach achieves up to 2.64x end-to-end speedup and 2.74x MoE-layer speedup. Compared with a conventional non-overlap baseline, our approach consistently improves both operator- and MoE-layer-level performance across varying GEMM shapes, router modes, and a broad range of producer/consumer SM partitions, while preserving correctness.