Compute-communication overlap in MoE reduces latency through tiled scheduling and signaling. This directly impacts throughput and GPU utilization.
Modern Mixture-of-Experts (MoE) systems are limited not by compute, but by communication. In distributed execution, each layer requires two all-to-all operations, and the second—returning results—falls into the critical path. The classical scheme only initiates it after the GEMM is complete, causing the GPU to idle while waiting for data transfer. This is exacerbated by imbalance: the computational power of GPUs is growing faster than interconnect bandwidth, leading to an increased share of time spent on communication.
Attempts to overlap computation and communication (compute-communication overlap) typically use coarse-grained decomposition. GEMM and all-to-all are broken into chunks and executed in a pipeline. This approach is straightforward but has side effects: it reduces the efficiency of tensor cores on small chunks, introduces synchronization on the host side, and increases kernel launch overhead. An alternative is fusion, where compute and communication are combined into a single kernel. This reduces overhead but requires complex engineering and custom primitives.
The approach discussed here strikes a compromise: keeping separate kernels but synchronizing them at the tile level (tile-level signaling). The key idea is to start data transfer as soon as part of the GEMM is ready, without waiting for the entire operation to complete.
The solution is built as a producer-consumer model within the GPU. The producer is a persistent GEMM kernel that processes all experts on the rank without repeated launches. The consumer is a separate persistent kernel that handles data transfer via NVSHMEM. They operate in parallel on different SMs (streaming multiprocessors), eliminating resource contention.
A critical element is signaling at the tile level. After computing each tile in the epilogue of the GEMM kernel, a readiness flag is published. The consumer monitors these flags and initiates data transfer as soon as a sufficient volume is accumulated. This removes the dependency on host-side orchestration and reduces latency between compute and communication.
To make such a scheme feasible, a special data organization is introduced—a remote-owner-aligned layout. Each tile is pre-associated with one destination rank. This eliminates the need for routing at the row level and allows the transfer to be performed as one contiguous write. Additionally, a remote-first scheduling strategy is employed: data that needs to be sent to other GPUs is computed first. This increases the overlap window.
The granularity of transfer also becomes a parameter. Transferring each tile separately is inefficient due to low bandwidth utilization. Therefore, segmentation is introduced: several row bands are combined into one transfer. Small segments reduce start latency, while larger ones increase throughput. The balance is determined empirically.
At the execution level, the system divides SMs between compute and communication. A small portion is allocated for the consumer kernel. This creates a trade-off: too few resources and communication cannot keep up with compute; too many and GEMM slows down. The optimal range was found in the middle zone, where stable overlap is achieved without degrading computations.
Results show that such compute-communication overlap provides up to 2.64x end-to-end acceleration and up to 2.74x at the MoE layer level compared to the baseline implementation without overlap. When compared to other MoE systems, the gain depends on the model. In scenarios with high communication volume, the advantage is most noticeable. If the hidden layer size is smaller, the effect diminishes as the volume of data that can be “hidden” behind computations decreases.
Importantly, the solution does not require modifications to the underlying primitives (GEMM, all-to-all) and does not introduce a strict dependency on custom kernels. This makes it practical for integration into existing stacks.
From an engineering perspective, this is an evolutionary improvement: instead of changing the algorithm or model, the critical execution path is optimized. The approach demonstrates that further scaling of MoE will depend not only on model architecture but also on how effectively the system can hide communication overhead.
Information source
arXiv is the largest open preprint repository (since 1991, under the auspices of Cornell), where researchers quickly post working versions of papers; the materials are publicly accessible but do not undergo full peer review, so results should be considered preliminary and, where possible, checked against updated versions or peer‑reviewed journals. arxiv.org