grouped-gemm-moe

A variable-size batched expert GEMM with the Mixture-of-Experts token permutation fused into the matmul, trainable end to end, loadable through kernels. The reference baselines are torch._grouped_mm, a pure-torch MoE (gather, per-expert matmul, scatter, combine), and PyTorch autograd, matched to a few parts in ten thousand with the fp32 fusion exact bit-for-bit.

A MoE layer routes each token to a few of E experts, so the work is a batch of GEMMs with variable row counts wrapped by a gather and a scatter. Done directly, both wrappers materialize [num_assignments, hidden] tensors, pure memory traffic, and padding to equal group sizes wastes compute. This kernel gathers rows as they are read into shared memory and scatters results with the routing-weight combine as they are written, so neither intermediate exists, the variable group sizes are respected without padding, an expert with no tokens costs nothing, and the backward carries the same fusion, so gradients flow to tokens, expert weights, and routing weights.

4,096 tokens routing top-2 into eight experts, real assignment counts per tile and one expert deliberately empty at zero cost; the fused path matches the eager MoE to 1e-6 and runs the bf16 forward in 0.72 ms against 1.24 ms eager, with no [assignments, hidden] intermediate ever materialized.

Usage

import torch
from kernels import get_kernel

moe = get_kernel("phanerozoic/grouped-gemm-moe", version=1, trust_remote_code=True)

T, K, N, E, k = 4096, 2048, 2048, 8, 2
x = torch.randn(T, K, device="cuda", requires_grad=True)        # tokens
W = torch.randn(E, N, K, device="cuda", requires_grad=True)     # expert weights
expert_ids = torch.stack([torch.randperm(E, device="cuda")[:k] for _ in range(T)])
topk_weights = torch.rand(T, k, device="cuda", requires_grad=True)
y = moe.moe(x, W, expert_ids, topk_weights)                     # [T, N]
y.sum().backward()                          # grads to x, W, topk_weights

# bare variable-size grouped GEMM (rows already grouped by expert)
out = moe.grouped_gemm(a, W[:4].detach(), group_sizes)

version selects the release branch; trust_remote_code is required by kernels for publishers without the trusted-publisher mark.

API

Method

The routing becomes a sort by expert; from the group sizes the launcher lays out one row-tile per BM rows of each expert. Per tile the GEMM kernel gathers activation rows through the row -> token map straight into shared memory, multiplies against that tile's expert weight with fp32 accumulation (a 128x128 mma.sync grouped GEMM with a cp.async pipeline for bf16, a register-blocked tile loop for the exact fp32 path), and scatters each result into its original token's output row scaled by the routing weight with an atomic add. The backward reuses the same machinery: the token gradient is the fused forward on the incoming gradient with transposed expert weights; the weight gradient is a per-expert A^T B reduction; the routing-weight gradient is an inner product with the unweighted product.

Measured

• The variable-size grouped GEMM agrees with torch._grouped_mm to 2e-5 relative Frobenius (bf16) and an independent per-group reference to 4e-3 max relative, including uneven and empty groups and non-tile-multiple dimensions.
• The fused dispatch matches a pure-torch MoE to 2.8e-4 (top-1) and 3.8e-4 (top-2) in fp32, 2.3e-3 relative Frobenius in bf16.
• The fusion is exact: the fused forward equals the explicit gather-grouped-GEMM-scatter path bit-for-bit in fp32 (max abs diff 0).
• The backward matches autograd to 3e-4 (tokens), 2e-4 (expert weights), 8e-5 (routing weights), and agrees with finite differences.
• The bf16 tensor-core forward is about 19x the portable tiled path on an H200; measured against an eager bf16 MoE loop at T = 4096, E = 8: 0.72 ms fused vs 1.24 ms eager.

Requirements and limits

• NVIDIA GPU with compute capability 8.0+ (one source spans 8.0 to 12.0).
• float32 or bfloat16 activations and weights; contraction and output dimensions arbitrary (tile-boundary guarded).
• Gating (softmax and top-k selection) is the caller's; the routing sort runs in torch.

References

Gale et al., "MegaBlocks" (block-sparse grouped GEMM without padding); Tan et al., ScatterMoE (fused gather-scatter GEMM); Lepikhin et al., "GShard"; Fedus et al., "Switch Transformer"; torch._grouped_mm.

License

Apache-2.0.