RapidGEMM: An Efficient Tile-agnostic GEMM for Mixture-of-Experts Training
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....