Principle:Unslothai Unsloth MoE Grouped GEMM
| Knowledge Sources | |
|---|---|
| Domains | MoE, Linear_Algebra, GPU_Computing |
| Last Updated | 2026-02-07 08:40 GMT |
Overview
Algorithm for performing batched matrix multiplications across multiple MoE experts in a single fused GPU kernel with automatic differentiation support.
Description
Grouped GEMM (General Matrix Multiplication) replaces the traditional approach of sequential per-expert matrix multiplications with a single fused kernel that processes all experts simultaneously. In a Mixture-of-Experts model, each token is routed to a subset of experts, requiring different subsets of inputs to be multiplied by different weight matrices. Grouped GEMM tiles this computation across all experts using persistent tile scheduling, where each GPU streaming multiprocessor (SM) processes tiles from multiple expert groups.
The key innovation is fusing token permutation (gathering tokens by expert assignment) directly into the GEMM kernel, avoiding separate permute/unpermute memory operations and reducing memory bandwidth pressure.
Usage
Apply this principle when implementing MoE layers where the per-expert GEMM is the computational bottleneck. It is the core computation strategy for efficient MoE training and inference on GPUs.
Theoretical Basis
For E experts, each with weight matrix , and token subsets :
The grouped GEMM computes all E multiplications in a single kernel launch:
Forward: Failed to parse (syntax error): {\displaystyle Y = \text{GroupedGEMM}(X, W, \text{m\_sizes})} Backward dX: Backward dW:
Pseudo-code Logic:
# Abstract grouped GEMM with persistent tiling
for tile in all_tiles_across_all_experts:
expert_id = get_expert_for_tile(tile)
X_tile = load_input_tile(X, expert_id, tile)
W_tile = load_weight_tile(W, expert_id, tile)
acc += X_tile @ W_tile.T
if tile_complete:
store_output(Y, expert_id, tile, acc)