RapidGEMM: An Efficient Tile-agnostic GEMM for Mixture-of-Experts Training

Abstract

Grouped GEMM (General Matrix Multiplication) is the core primitive in performing each expert’s matrix multiplications on tokens routed to them in MoE (Mixture-of-Experts) training. However, existing implementations either drop tokens or zero-pad each expert’s token slice to meet GEMM tiling and alignment requirements. While dropping tokens avoids some padding overhead, it degrades model accuracy, whereas zero-padding wastes compute and memory on zero-filled slots and adds unpadding overhead. To address this problem, we design RapidGEMM, a new grouped GEMM kernel, for efficient MoE model training. Specifically, RapidGEMM adopts two no-drop no-padding strategies: (i) overflow-safe loads and predicated stores in the forward pass and (ii) overflow-bound loads and bulk stores in the backward. This design eliminates per-expert padding while producing the same result as the approach with padding and effectively overlaps asynchronous tensor-level transfers with MMA (Matrix Multiply-Accumulate) to avoid warp stalls. On B200 GPUs, the BF16 RapidGEMM kernel achieves up to 1.20 × speedup in the forward pass and 1.27 × in the backward pass, and speeds up both passes of a MoE layer by 1.12 × on average over the state-of-the-art grouped GEMM kernels. Code available at https://doi.org/https://github.com/RapidsAtHKUST/RapidGEMM.