grouped-gemm-moe

grouped-gemm-moe is a variable-size batched expert GEMM with the Mixture-of-Experts token permutation fused into the matmul, loadable through kernels. A MoE layer routes each token to a few of E experts, so the work is a batch of GEMMs with variable row counts, one per expert, wrapped by a gather that brings each expert's tokens together and a scatter that returns the results and combines them by routing weight. Done the direct way the gather writes a [num_assignments, hidden] activation tensor and the scatter writes another [num_assignments, hidden] output tensor, both pure memory traffic. This kernel fuses the permutation into the GEMM: rows are gathered as they are read into shared memory and results are scattered and routing-weight-combined as they are written, so neither intermediate is ever materialized. One tiled GEMM spans every expert; each threadblock owns a tile of a single expert's rows, so the variable group sizes are respected without padding and experts with no tokens cost nothing. The backward is fused the same way, so the layer is trainable end to end: gradients flow to the tokens, the expert weights, and the routing weights.

Usage

import torch
from kernels import get_kernel

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

# fused Mixture-of-Experts forward, differentiable end to end
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)  # from your gate
y = moe.moe(x, W, expert_ids, topk_weights)                     # [T, N] float32
y.sum().backward()                                              # grads to x, W, topk_weights

# bare variable-size grouped GEMM (rows already grouped by expert)
a = torch.randn(1000, K, device="cuda")
group_sizes = torch.tensor([200, 0, 500, 300], device="cuda")  # sums to 1000
out = moe.grouped_gemm(a, W[:4].detach(), group_sizes)         # [1000, N]

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

API

How it works

The routing is turned into a sort by expert (a cheap index op in torch): the S = T*k token-to-expert assignments are ordered so each expert's rows are contiguous, giving a row -> token map, a per-row routing weight, and per-expert group sizes. From the group sizes the launcher lays out one row-tile per BM rows of each expert and records which expert each tile belongs to. The GEMM kernel then, per tile:

• gathers on load the activation rows through the row -> token map straight into shared memory, so the permuted [S, K] activation tensor is never stored;
• multiplies against that tile's expert weight and accumulates in fp32 — for bf16 on the tensor cores (a 128x128 mma.sync grouped GEMM with a cp.async multi-stage pipeline streaming the next K chunk while the current one computes), or a register-blocked tile loop for the exact fp32 path;
• scatters on store each result into the output row of its original token, scaled by the routing weight, with an atomic add, so the expanded [S, N] output tensor is never stored and the top-k combine happens in place.

The same kernel, with the permutation maps omitted, is a bare variable-size grouped GEMM (rows already grouped, written in order), which matches torch._grouped_mm.

The backward reuses the same fused machinery. The gradient to the tokens is the fused forward run on the incoming gradient with each expert weight transposed, so the gather-on-load and scatter-combine carry the permutation exactly as in the forward. The gradient to the expert weights is a per-expert A^T B reduction over each expert's assignments (a second grouped kernel). The gradient to the routing weights is the inner product of the incoming gradient with the unweighted per-assignment product, obtained from the bare grouped GEMM.

Correctness

Measured on L4, H200, and RTX PRO 6000 through get_kernel:

• Grouped GEMM matches torch._grouped_mm. The variable-size grouped GEMM agrees with PyTorch's own grouped matmul to a relative Frobenius error of 2e-5 (bf16), and with an independent per-group reference to max relative error 4e-3 (fp32) and 3.9e-3 (bf16), including with uneven and empty groups and for contraction/output dimensions that are not multiples of the tile.
• Fused MoE matches a reference. The fused dispatch matches a pure-torch MoE (gather, per-expert matmul, scatter, routing-weight combine) to max relative error 2.8e-4 for top-1 and 3.8e-4 for top-2 routing in fp32, and relative Frobenius error 2.3e-3 in bf16.
• The fusion is exact. The fused forward equals the explicit gather-then-grouped-GEMM-then-scatter path bit-for-bit in fp32 (max abs difference 0), confirming the gather-on-load and scatter-combine-on-store carry the permutation without approximation; because it accumulates in fp32 with no round-trip through the operand dtype, it is strictly at least as accurate as the unfused path.
• The backward matches autograd. The fused gradients to the tokens, the expert weights, and the routing weights match PyTorch autograd of an equivalent MoE to max relative error 3e-4, 2e-4, and 8e-5 respectively, and agree with finite differences. The per-expert weight-gradient kernel matches an explicit reference to 4e-4.

Scope

The contribution is a trainable variable-size grouped GEMM with the MoE gather and scatter-combine fused into both the forward and the backward, so no [num_assignments, hidden] intermediates are materialized in the forward and no padding is used, and the layer differentiates end to end (tokens, expert weights, routing weights), spanning every architecture from compute capability 8.0 to 12.0 from one source. The bf16 forward computes each tile on the tensor cores with a cp.async-pipelined mma.sync grouped GEMM (the gather and scatter fused around it), about 19x faster than the portable tiled path on an H200; the fp32 forward keeps the exact tiled path, so the bit-exact fp32 and torch._grouped_mm-matching results are unchanged. On top of the tensor-core throughput the fused dispatch is the structural win: no materialized intermediates (the memory saved is 2 * S * (K + N) bytes, growing with token count and hidden size), no padding, and zero cost for empty experts. Gating, meaning the softmax and top-k selection that produce the expert ids and routing weights, is the caller's; the routing sort runs in torch.

Requirements and limits

• NVIDIA GPU with compute capability 8.0+.
• float32 or bfloat16 activations and weights; contraction and output dimensions are arbitrary (tile-boundary guarded).

References

Gale et al., "MegaBlocks: Efficient Sparse Training with Mixture-of-Experts" (block-sparse grouped GEMM without padding); Tan et al., "Scattered Mixture-of- Experts" / ScatterMoE (fused gather-scatter GEMM); Lepikhin et al., "GShard" and Fedus et al., "Switch Transformer" (token routing); torch._grouped_mm.

License

Apache-2.0.