Mixture-of-Experts Kernels: Routing, Permutation, and Grouped GEMM
Mixture-of-Experts Kernels: Routing, Permutation, and Grouped GEMM
Most frontier open-weight models are now Mixture-of-Experts (MoE) models. DeepSeek-V3 and Kimi K2 have 256 and 384 routed experts per layer, Qwen3-235B-A22B and gpt-oss-120b have 128, and every token uses only 4 to 8 of them. The model holds hundreds of billions of parameters, but each token pays for only a few tens of billions. For a kernel engineer, that changes the central operation. Instead of one large GEMM per feed-forward layer, each layer runs dozens to hundreds of small GEMMs of different, data-dependent sizes, and the sizes are only known after the router has run on the GPU.
This chapter builds the MoE layer kernel by kernel. It writes a ground-truth reference, gets the routing semantics right (including the DeepSeek-style bias that is easy to misuse), builds the block-aligned indirection table that production kernels consume, writes a fused grouped GEMM in Triton that gathers and scatters rows itself, verifies it on a CPU with the Triton interpreter, compiles it for a Hopper GPU you do not own, and then uses shape arithmetic to explain why prefill and decode need different MoE kernels.
1. Why MoE is a kernel problem
A dense SwiGLU feed-forward layer computes \(W_2\,(\mathrm{silu}(W_1x)\odot W_3x)\) with one weight set. An MoE layer has \(E\) such weight sets (the experts) plus a small router \(W_g\in\mathbb{R}^{E\times H}\). For each token, the router scores all experts, keeps the top \(k\), and the layer output is the weighted sum of those \(k\) experts’ outputs.
| Model (config.json) | Routed experts \(E\) | Active \(k\) | Hidden \(H\) | Expert intermediate \(I\) | Router |
|---|---|---|---|---|---|
| Mixtral-8x7B | 8 | 2 | 4096 | 14336 | softmax, top-2 |
| gpt-oss-120b | 128 | 4 | 2880 | 2880 | softmax over top-4 |
| Qwen3-235B-A22B | 128 | 8 | 4096 | 1536 | softmax, renormalized top-8 |
| DeepSeek-V3 | 256 (+1 shared) | 8 | 7168 | 2048 | sigmoid + bias, 8 groups, top-4 groups |
| Kimi K2 | 384 (+1 shared) | 8 | 7168 | 2048 | sigmoid + bias |
Three properties make this hard to run fast:
- Ragged, data-dependent sizes. Expert \(e\) receives \(n_e\) token rows, with \(\sum_e n_e = T\cdot k\) for a batch of \(T\) tokens. The \(n_e\) change every step, live in GPU memory, and are often zero.
- Weights dominate bytes. One DeepSeek-V3 expert is \(3HI\approx 44\) M parameters. Every active expert must stream its whole weight set from HBM, whether it serves one token or a thousand.
- Imbalance. Routers are trained to balance load on average. Any given batch can still send most tokens to a few hot experts, and the slowest expert sets the latency of the layer.
The rest of the chapter is about turning “\(E\) ragged GEMMs” into one launch that keeps the tensor cores busy when there is enough work, and keeps HBM busy when there is not.
2. Mental model: route, align, grouped GEMM, combine
x [T, H] ──router GEMM──► logits [T, E] ──top-k──► topk_ids, topk_w [T, k]
│
moe_align: histogram ─► pad each expert to BLOCK_M ─► exclusive scan ─► scatter
▼
sorted_ids [padded rows] (flat row r = token*k + slot, or a sentinel)
expert_ids [padded rows / BLOCK_M] (one expert per BLOCK_M-row tile)
│
GEMM1: rows of x gathered by sorted_ids // k, B = W13[expert] ──► h [T*k, 2I] at row r
act : silu(h[:, :I]) * h[:, I:] ──► a [T*k, I]
GEMM2: rows of a gathered by sorted_ids, B = W2[expert], × topk_w[r] ──► y [T*k, H] at row r
combine: out[t] = Σ_slot y[t*k + slot] ──► out [T, H]There are two ways to feed the expert GEMMs:
| Design | Data movement | Used by |
|---|---|---|
| Permute, then grouped GEMM. Physically reorder token rows so each expert’s rows are contiguous, run a grouped GEMM with per-expert offsets, then un-permute. | Writes and reads a permuted copy of the activations (and of the outputs) | torch.nn.functional.grouped_mm, CUTLASS grouped GEMM, DeepGEMM’s grouped layouts, most expert-parallel stacks (the all-to-all already permutes) |
| Gather and scatter inside the GEMM. Keep activations in token order; each tile reads its rows through an index table and writes results back to their original positions. | No permuted copy; indexed loads instead of contiguous ones | vLLM and SGLang fused_moe Triton kernels |
Both need the same metadata: how many rows each expert has, and where each expert’s rows start. Computing that is a histogram followed by a prefix sum (scan), the same primitive that underlies radix sort.
3. Ground truth first
Every faster version in this chapter is compared against two deliberately simple implementations: a per-expert loop, and a dense version that runs every expert on every token and masks the result. They share no indexing code, so a bug in one does not hide in the other.
# file: moe_reference.py
# Ground truth for a top-k Mixture-of-Experts FFN layer (SwiGLU experts), in PyTorch on CPU.
# route: logits = x @ Wg^T -> softmax -> top-k -> renormalise the k weights
# experts: y_e(x) = W2_e (silu(W1_e x) * (W3_e x))
# combine: out(x) = sum_j w_j * y_{e_j}(x)
# Two independent implementations must agree: a per-expert loop and a dense "run every expert" version.
import torch
torch.manual_seed(0)
T, H, I, E, TOP_K = 64, 64, 96, 8, 2 # tokens, hidden, expert intermediate, experts, k
def route(x, w_gate, top_k, renormalize=True):
probs = torch.softmax(x @ w_gate.T, dim=-1) # [T, E], fp32
w, ids = torch.topk(probs, top_k, dim=-1) # [T, k] weights and expert ids
if renormalize:
w = w / w.sum(-1, keepdim=True)
return w, ids
def moe_loop(x, w_gate, w1, w3, w2, top_k):
"""Per-expert loop: gather the tokens routed to e, run e once, scatter-add weighted outputs."""
w, ids = route(x, w_gate, top_k)
out = torch.zeros_like(x)
for e in range(w1.shape[0]):
tok, slot = torch.nonzero(ids == e, as_tuple=True) # which tokens chose e, and in which slot
if tok.numel() == 0:
continue # an expert may receive no tokens at all
xe = x[tok]
ye = (torch.nn.functional.silu(xe @ w1[e].T) * (xe @ w3[e].T)) @ w2[e].T
out.index_add_(0, tok, ye * w[tok, slot, None])
return out, ids
def moe_dense(x, w_gate, w1, w3, w2, top_k):
"""Run every expert on every token, then keep only the routed ones. E/k times the FLOPs, but obviously right."""
w, ids = route(x, w_gate, top_k)
gate = torch.zeros(x.shape[0], w1.shape[0]).scatter(1, ids, w) # [T, E], k nonzeros per row
y_all = torch.stack([(torch.nn.functional.silu(x @ w1[e].T) * (x @ w3[e].T)) @ w2[e].T
for e in range(w1.shape[0])], dim=1) # [T, E, H]
return (gate[:, :, None] * y_all).sum(1)
def make_layer(H, I, E, seed=0):
g = torch.Generator().manual_seed(seed)
s = 1.0 / H ** 0.5
return (torch.randn(E, H, generator=g) * s, # router weight [E, H]
torch.randn(E, I, H, generator=g) * s, # W1 (gate proj) [E, I, H] nn.Linear layout
torch.randn(E, I, H, generator=g) * s, # W3 (up proj) [E, I, H]
torch.randn(E, H, I, generator=g) / I ** 0.5) # W2 (down proj) [E, H, I]
if __name__ == "__main__":
x = torch.randn(T, H)
wg, w1, w3, w2 = make_layer(H, I, E)
out_loop, ids = moe_loop(x, wg, w1, w3, w2, TOP_K)
out_dense = moe_dense(x, wg, w1, w3, w2, TOP_K)
counts = torch.bincount(ids.flatten(), minlength=E)
print("tokens per expert :", counts.tolist(), " (sum =", int(counts.sum()), "= T*k)")
print("max |loop - dense|:", f"{(out_loop - out_dense).abs().max().item():.2e}")
flops_moe = 2 * T * TOP_K * 3 * H * I
print(f"expert FLOPs: routed {flops_moe/1e6:.1f} MFLOP vs dense {flops_moe*E/TOP_K/1e6:.1f} MFLOP")python3 -m venv .venv && . .venv/bin/activate
pip install torch --index-url https://download.pytorch.org/whl/cpu
pip install triton numpy
python moe_reference.pyOutput (PyTorch 2.14.1 CPU):
tokens per expert : [19, 11, 15, 15, 21, 15, 18, 14] (sum = 128 = T*k)
max |loop - dense|: 0.00e+00
expert FLOPs: routed 4.7 MFLOP vs dense 18.9 MFLOPThe two implementations agree to the last bit here because they run the same row-wise products on the same CPU kernels. The FLOP line is the reason MoE exists: with \(k/E = 2/8\) the routed layer does a quarter of the dense work. DeepSeek-V3’s \(8/256\) is \(1/32\).
Note the indexing contract that every later kernel keeps: the \(T\cdot k\) routed copies of the tokens are numbered by a flat row \(r = t\cdot k + j\) (token \(t\), slot \(j\)). Expert loads are just bincount over the flattened topk_ids.
4. Routing semantics: the bias is for choosing, not for weighting
Softmax-then-top-\(k\) (Mixtral, Qwen3) is simple. DeepSeek-V3 and Kimi K2 use a router that kernel ports regularly get wrong:
- scores are sigmoid affinities, not a softmax;
- a per-expert correction bias, adjusted during training to balance load without an auxiliary loss, is added to the scores only to select experts;
- group-limited routing: experts are split into groups (8 groups for DeepSeek-V3). Each group is scored by the sum of its top-2 biased scores, only the best
topk_groupgroups (4) stay eligible, and top-\(k\) is taken inside them. In deployment this bounds how many nodes a token’s experts span; - the gate weights come from the unbiased sigmoid scores, are renormalized over the \(k\) chosen experts, and are multiplied by
routed_scaling_factor(2.5 for DeepSeek-V3).
# file: router_grouped.py
# DeepSeek-V3-style routing: sigmoid scores, a per-expert correction bias that steers SELECTION only
# (auxiliary-loss-free load balancing), group-limited top-k, renormalisation and a fixed scale.
import torch
torch.manual_seed(0)
T, E, K, N_GROUP, TOPK_GROUP, SCALE = 4, 16, 4, 4, 2, 2.5 # small stand-ins for 256 / 8 / 8 / 4 / 2.5
def route_grouped(logits, bias, use_bias_for_weights=False):
scores = torch.sigmoid(logits) # [T, E] affinities in (0, 1)
choice = scores + bias # biased scores: selection only
grp = choice.view(T, N_GROUP, -1).topk(2, dim=-1).values.sum(-1) # group score = sum of its top-2
keep_grp = grp.topk(TOPK_GROUP, dim=-1).indices # keep the best TOPK_GROUP groups
mask = torch.zeros(T, N_GROUP).scatter(1, keep_grp, 1.0).repeat_interleave(E // N_GROUP, dim=1)
ids = choice.masked_fill(mask == 0, float("-inf")).topk(K, dim=-1).indices
w = (choice if use_bias_for_weights else scores).gather(1, ids) # weights from UNBIASED scores
return ids, SCALE * w / w.sum(-1, keepdim=True)
logits = torch.randn(T, E)
bias = torch.zeros(E)
bias[torch.arange(E) % 4 == 0] = 0.5 # the balancer has been boosting under-used experts 0,4,8,12
ids0, w0 = route_grouped(logits, torch.zeros(E))
ids1, w1 = route_grouped(logits, bias)
ids2, w2 = route_grouped(logits, bias, use_bias_for_weights=True)
print("no bias ids:", ids0.tolist())
print("with bias ids:", ids1.tolist())
print("groups used :", (ids1 // (E // N_GROUP)).unique(dim=None).tolist(), "of", N_GROUP,
"| per token:", [sorted(set(r)) for r in (ids1 // (E // N_GROUP)).tolist()])
print("weights (right) :", [[round(v, 3) for v in r] for r in w1.tolist()])
print("weights (bug) :", [[round(v, 3) for v in r] for r in w2.tolist()])
print("max weight error from leaking the bias:", f"{(w1 - w2).abs().max().item():.3f}")python router_grouped.pyOutput:
no bias ids: [[13, 14, 4, 5], [7, 13, 12, 14], [13, 6, 12, 7], [11, 10, 2, 1]]
with bias ids: [[4, 12, 13, 14], [12, 4, 7, 13], [12, 4, 13, 6], [11, 8, 0, 10]]
groups used : [0, 1, 2, 3] of 4 | per token: [[1, 3], [1, 3], [1, 3], [0, 2]]
weights (right) : [[0.635, 0.48, 0.702, 0.683], [0.637, 0.514, 0.688, 0.661], [0.561, 0.543, 0.812, 0.585], [0.962, 0.388, 0.336, 0.815]]
weights (bug) : [[0.798, 0.685, 0.516, 0.501], [0.785, 0.692, 0.522, 0.501], [0.742, 0.729, 0.598, 0.431], [0.688, 0.633, 0.596, 0.583]]
max weight error from leaking the bias: 0.273Two things to read off it. First, the bias changes which experts are chosen. The boosted experts 0, 4, 8, and 12 move into the top-\(k\), and each token uses exactly two of the four groups. Second, computing the gate weights from the biased scores changes the layer output by up to 0.27 in weight per expert. Nothing crashes and every shape is right; the model just gets measurably worse. Fused router kernels (vLLM’s grouped_topk, SGLang’s biased group top-k) exist to do this whole block in one launch, and their unit tests check exactly this distinction.
5. The indirection table: histogram, pad, scan, scatter
A tiled GEMM wants every BLOCK_M-row tile to use one expert’s weights. So each expert’s run of rows is padded up to a multiple of BLOCK_M with a sentinel row index, and a second array records the expert that owns each tile. This is the contract of vLLM’s moe_align_block_size:
sorted_ids[p]: the flat row \(r\) at padded position \(p\), or the sentinel \(T\cdot k\) for padding;expert_ids[p // BLOCK_M]: the expert for that tile;- the padded total, so that tiles past the end can exit at once.
# file: moe_align.py
# Build the block-aligned indirection table a fused MoE GEMM consumes (same contract as vLLM's
# moe_align_block_size): sort the T*k (token, slot) pairs by expert, pad each expert's run to a
# multiple of BLOCK_M with a sentinel, and record which expert owns each BLOCK_M-row tile.
import torch
def moe_align(topk_ids, block_m, num_experts):
flat = topk_ids.flatten().to(torch.int64) # flat row r = token * k + slot
n = flat.numel()
counts = torch.bincount(flat, minlength=num_experts) # 1. histogram
padded = (counts + block_m - 1) // block_m * block_m # 2. round each run up to BLOCK_M
starts = torch.cumsum(padded, 0) - padded # 3. exclusive scan -> run start
total = int(padded.sum())
sorted_ids = torch.full((total,), n, dtype=torch.int32) # sentinel n = "padding row"
order = torch.argsort(flat, stable=True) # rows grouped by expert, stable
rank = torch.arange(n) - (torch.cumsum(counts, 0) - counts)[flat[order]]
sorted_ids[starts[flat[order]] + rank] = order.to(torch.int32) # 4. scatter into the runs
expert_ids = torch.repeat_interleave(torch.arange(num_experts, dtype=torch.int32),
padded // block_m) # one entry per BLOCK_M tile
return sorted_ids, expert_ids, total
if __name__ == "__main__":
# The worked example from vLLM's moe_align_block_size docstring.
topk_ids = torch.tensor([[2, 3, 4], [1, 2, 4], [1, 3, 4], [1, 2, 3]])
sorted_ids, expert_ids, total = moe_align(topk_ids, block_m=4, num_experts=5)
print("flat expert of each row :", topk_ids.flatten().tolist())
print("sorted_ids :", sorted_ids.tolist())
print("expert_ids :", expert_ids.tolist())
print("rows after padding:", total, " useful:", topk_ids.numel(),
f" waste: {1 - topk_ids.numel() / total:.0%}")python moe_align.pyOutput:
flat expert of each row : [2, 3, 4, 1, 2, 4, 1, 3, 4, 1, 2, 3]
sorted_ids : [3, 6, 9, 12, 0, 4, 10, 12, 1, 7, 11, 12, 2, 5, 8, 12]
expert_ids : [1, 2, 3, 4]
rows after padding: 16 useful: 12 waste: 25%This reproduces the worked example in vLLM’s moe_align_block_size docstring exactly. Expert 0 receives no rows, so it gets no tile and its weights are never read. Experts 1 to 4 have 3 rows each, padded to 4, and the sentinel 12 marks the padding.
On a GPU, the four numbered steps are one small kernel. The histogram uses shared-memory atomics, the scan over \(E\) counts runs inside one CTA (with \(E\le 384\) it fits easily), and each row is scattered to start[e] + rank. vLLM and SGLang ship CUDA versions of this kernel because, at small batch sizes, its cost is not negligible next to the GEMMs. The stable=True sort is a choice made here, not a requirement. Keeping rows in token order within an expert makes the gathered loads more local and the output deterministic.
6. A fused grouped GEMM in Triton
The kernel is an ordinary tiled GEMM with three changes:
- Tile to expert. Program
pid_mreadsexpert_ids[pid_m]and offsets the weight pointer byexpert * stride_be. - Gathered A rows. Row indices come from
sorted_ids. For GEMM1, the input has one row per token, shared by its \(k\) slots, so the kernel loads rowr // TOP_K. For GEMM2, the input already has one row per flat row, soTOP_K=1. - Scattered C rows, with the router weight folded in. Results are stored at row
r, so the output comes out in token order. GEMM2 multiplies bytopk_w[r]in the epilogue, which leaves the combine as a plain sum over \(k\).
Padding rows (r >= num_rows) are masked on load and on store. They cost tensor-core work but no memory traffic.
# file: fused_moe_triton.py
# A fused-MoE grouped GEMM in Triton: one launch computes every expert's GEMM by reading a
# block-aligned indirection table (moe_align.py). Each program owns BLOCK_M rows that all belong to
# ONE expert, gathers their token rows from A, multiplies by that expert's weight tile, and scatters
# the result to the rows' original (token, slot) positions. No permuted copy of A is ever written.
# Runs on a GPU, or on CPU for correctness work with: TRITON_INTERPRET=1 python fused_moe_triton.py
import os
import torch
import torch.nn.functional as F
import triton
import triton.language as tl
from moe_align import moe_align
from moe_reference import make_layer, moe_loop, route
@triton.jit
def moe_gemm_kernel(A, B, C, SORTED_IDS, EXPERT_IDS, TOPK_W,
num_rows, padded_rows, N, K,
stride_am, stride_be, stride_bn, stride_cm,
TOP_K: tl.constexpr, MUL_ROUTED_WEIGHT: tl.constexpr,
BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
pid_m = tl.program_id(0)
pid_n = tl.program_id(1)
if pid_m * BLOCK_M >= padded_rows: # grid is sized for the worst case; trailing tiles idle
return
offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
rows = tl.load(SORTED_IDS + offs_m) # flat (token*TOP_K + slot), or the sentinel num_rows
row_ok = rows < num_rows
a_rows = rows // TOP_K # GEMM1: A holds one row per token, shared by k slots
expert = tl.load(EXPERT_IDS + pid_m) # every row of this tile uses the same expert
offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
offs_k = tl.arange(0, BLOCK_K)
a_ptrs = A + a_rows[:, None] * stride_am + offs_k[None, :]
b_ptrs = B + expert * stride_be + offs_n[None, :] * stride_bn + offs_k[:, None] # B[e] is [N, K]
acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
for k0 in range(0, K, BLOCK_K):
k_ok = (k0 + offs_k) < K
a = tl.load(a_ptrs, mask=row_ok[:, None] & k_ok[None, :], other=0.0)
b = tl.load(b_ptrs, mask=k_ok[:, None] & (offs_n[None, :] < N), other=0.0)
acc += tl.dot(a, b)
a_ptrs += BLOCK_K
b_ptrs += BLOCK_K
if MUL_ROUTED_WEIGHT: # fold the router weight into the down projection
w = tl.load(TOPK_W + rows, mask=row_ok, other=0.0)
acc = acc * w[:, None]
c_ptrs = C + rows[:, None] * stride_cm + offs_n[None, :]
tl.store(c_ptrs, acc.to(C.dtype.element_ty), mask=row_ok[:, None] & (offs_n[None, :] < N))
def moe_gemm(a, b, out, sorted_ids, expert_ids, padded, topk_w, top_k, mul_w,
block_m, block_n=32, block_k=32):
"""a: [rows_a, K]; b: [E, N, K] (nn.Linear layout); out: [T*k, N] written at the flat row index."""
E, N, K = b.shape
grid = (expert_ids.numel(), triton.cdiv(N, block_n))
moe_gemm_kernel[grid](a, b, out, sorted_ids, expert_ids, topk_w,
out.shape[0], padded, N, K,
a.stride(0), b.stride(0), b.stride(1), out.stride(0),
TOP_K=top_k, MUL_ROUTED_WEIGHT=mul_w,
BLOCK_M=block_m, BLOCK_N=block_n, BLOCK_K=block_k)
def fused_moe(x, w_gate, w1, w3, w2, top_k, block_m=16):
T, H = x.shape
E, I, _ = w1.shape
topk_w, topk_ids = route(x, w_gate, top_k)
sorted_ids, expert_ids, padded = moe_align(topk_ids, block_m, E)
sorted_ids, expert_ids = sorted_ids.to(x.device), expert_ids.to(x.device)
w13 = torch.cat([w1, w3], dim=1).contiguous() # [E, 2I, H]: gate and up in one GEMM
h = torch.empty(T * top_k, 2 * I, device=x.device, dtype=x.dtype)
moe_gemm(x, w13, h, sorted_ids, expert_ids, padded, topk_w, top_k, False, block_m)
act = F.silu(h[:, :I]) * h[:, I:] # [T*k, I] (fused in production)
y = torch.empty(T * top_k, H, device=x.device, dtype=x.dtype)
moe_gemm(act, w2, y, sorted_ids, expert_ids, padded, topk_w.flatten().contiguous(), 1, True, block_m)
return y.view(T, top_k, H).sum(dim=1), padded # combine the k weighted slots
def grouped_mm_moe(x, w_gate, w1, w3, w2, top_k):
"""Third implementation: explicit permute -> torch grouped GEMM with offsets -> un-permute."""
T, H = x.shape
E = w1.shape[0]
topk_w, topk_ids = route(x, w_gate, top_k)
flat = topk_ids.flatten()
order = torch.argsort(flat, stable=True) # rows grouped by expert
xs = x[order // top_k] # the permuted copy fused_moe avoids
offs = torch.cumsum(torch.bincount(flat, minlength=E), 0).to(torch.int32)
h = F.silu(F.grouped_mm(xs, w1.transpose(-2, -1), offs=offs)) * \
F.grouped_mm(xs, w3.transpose(-2, -1), offs=offs)
ys = F.grouped_mm(h, w2.transpose(-2, -1), offs=offs) * topk_w.flatten()[order, None]
out = torch.zeros_like(x)
out.index_add_(0, order // top_k, ys) # un-permute + combine
return out
if __name__ == "__main__":
device = "cpu" if os.environ.get("TRITON_INTERPRET") == "1" else "cuda"
H, I = 64, 96
cases = [ # name, tokens, experts, top_k, router bias toward expert 0
("uniform E=8 k=2", 64, 8, 2, 0.0),
("skewed E=16 k=4", 50, 16, 4, 3.0),
("decode T=1 k=2", 1, 8, 2, 0.0),
("top-1 E=8 k=1", 37, 8, 1, 0.0),
]
for name, T, E, k, bias in cases:
torch.manual_seed(1)
wg, w1, w3, w2 = make_layer(H, I, E, seed=2)
u = torch.nn.functional.normalize(torch.randn(H), dim=0)
wg[0] += bias * u # make expert 0 popular:
x = torch.randn(T, H) + 1.5 * u # its logit rises by ~1.5*bias
ref, ids = moe_loop(x, wg, w1, w3, w2, k)
got, padded = fused_moe(*(t.to(device) for t in (x, wg, w1, w3, w2)), k)
gmm = grouped_mm_moe(x, wg, w1, w3, w2, k)
counts = torch.bincount(ids.flatten(), minlength=E)
print(f"{name}: loads={counts.tolist()} rows {T*k} -> padded {padded} "
f"|fused-ref|={(got.cpu() - ref).abs().max().item():.1e} "
f"|grouped_mm-ref|={(gmm - ref).abs().max().item():.1e}")TRITON_INTERPRET=1 python fused_moe_triton.pyOutput (Triton 3.8.0, PyTorch 2.14.1 CPU, interpreter mode):
uniform E=8 k=2: loads=[18, 8, 24, 21, 15, 12, 21, 9] rows 128 -> padded 192 |fused-ref|=1.1e-06 |grouped_mm-ref|=0.0e+00
skewed E=16 k=4: loads=[43, 5, 20, 11, 7, 7, 12, 10, 8, 15, 7, 5, 14, 20, 7, 9] rows 200 -> padded 320 |fused-ref|=1.2e-06 |grouped_mm-ref|=0.0e+00
decode T=1 k=2: loads=[1, 0, 0, 0, 0, 1, 0, 0] rows 2 -> padded 32 |fused-ref|=3.0e-07 |grouped_mm-ref|=0.0e+00
top-1 E=8 k=1: loads=[5, 1, 11, 7, 4, 3, 6, 0] rows 37 -> padded 112 |fused-ref|=9.5e-07 |grouped_mm-ref|=0.0e+00What the test matrix covers:
| Case | Why |
|---|---|
| Uniform, \(E=8\), \(k=2\) | The common case; every expert has several tiles’ worth of rows or a partial tile |
| Skewed, \(E=16\), \(k=4\) | Expert 0 takes 43 of 50 tokens. Multi-tile experts next to experts with 5 rows |
| Decode, \(T=1\) | Six experts have zero rows, so they get no tiles. 2 useful rows become 32 padded rows |
| Top-1, \(T=37\) | TOP_K=1 in GEMM1 (no row sharing). Expert 7 is empty. \(T\) is not a multiple of anything |
| Three implementations | Loop, fused Triton kernel, and F.grouped_mm with explicit permute and un-permute |
The fused kernel matches the loop to about \(10^{-6}\) in FP32: the same sums, accumulated in a different order. F.grouped_mm is PyTorch’s public grouped-GEMM API. Its documentation targets BF16 on SM80+ CUDA GPUs. The CPU build used here also accepted FP32 inputs, which makes it a useful third opinion while developing.
The interpreter executes tile semantics with NumPy. It proves the indexing is right, which is where MoE kernels break. It says nothing about speed.
7. Compile for Hopper, and what the compiler tells you about decode
# file: compile_check.py
# Ahead-of-time compile the fused-MoE GEMM for an H100 (sm_90) on a machine with no GPU,
# once per specialisation the layer launches (GEMM1: TOP_K=k, no weight; GEMM2: TOP_K=1, weight).
import triton
from triton.backends.compiler import GPUTarget
from triton.compiler import ASTSource
from fused_moe_triton import moe_gemm_kernel
target = GPUTarget("cuda", 90, 32)
sig = {"A": "*bf16", "B": "*bf16", "C": "*bf16", "SORTED_IDS": "*i32", "EXPERT_IDS": "*i32",
"TOPK_W": "*fp32", "num_rows": "i32", "padded_rows": "i32", "N": "i32", "K": "i32",
"stride_am": "i32", "stride_be": "i32", "stride_bn": "i32", "stride_cm": "i32",
"TOP_K": "constexpr", "MUL_ROUTED_WEIGHT": "constexpr",
"BLOCK_M": "constexpr", "BLOCK_N": "constexpr", "BLOCK_K": "constexpr"}
for name, top_k, mul_w, bm in [("gemm1 (gate+up)", 8, False, 64), ("gemm2 (down) ", 1, True, 64),
("gemm1 decode ", 8, False, 16)]:
src = ASTSource(fn=moe_gemm_kernel, signature=sig,
constexprs={"TOP_K": top_k, "MUL_ROUTED_WEIGHT": mul_w,
"BLOCK_M": bm, "BLOCK_N": 128, "BLOCK_K": 64})
k = triton.compile(src, target=target)
ptx = k.asm["ptx"]
print(f"{name} BLOCK_M={bm:<3} num_warps={k.metadata.num_warps} shared={k.metadata.shared:>6} B "
f"mma={'wgmma' if 'wgmma' in ptx else ('mma.sync' if 'mma.sync' in ptx else 'none')} "
f"ptx_lines={ptx.count(chr(10))}")python compile_check.pyOutput (Triton 3.8.0 on a CPU-only Linux box):
gemm1 (gate+up) BLOCK_M=64 num_warps=4 shared= 24576 B mma=wgmma ptx_lines=3937
gemm2 (down) BLOCK_M=64 num_warps=4 shared= 24576 B mma=wgmma ptx_lines=3936
gemm1 decode BLOCK_M=16 num_warps=4 shared= 18432 B mma=mma.sync ptx_lines=2294All three specializations lower to sm_90 code. The last column is the useful part. With BLOCK_M=64, Triton emits Hopper wgmma warpgroup MMAs. With BLOCK_M=16, it falls back to Ampere-style mma.sync, because wgmma works on tiles that are 64 rows tall per warpgroup. That is the structural problem of MoE decode in one line:
- At decode, an active expert typically has 1 to 4 rows. A 64-row tile is more than 90% padding (section 8), and a 16-row tile cannot use the Hopper MMA path at all.
- At those sizes the layer is bandwidth-bound anyway, so wasted MMA rows cost little. What matters is that each active expert’s weights are streamed once, at full HBM bandwidth, by enough CTAs to cover the SMs.
That is why production stacks keep separate MoE kernels or configurations for small batches. Common approaches are tuned configs with small BLOCK_M (vLLM keys its tuned Triton configs by \(E\), \(N\), and GPU), a per-token path (vLLM skips alignment when \(T\cdot k\cdot 4\le E\) and gives each routed row its own tile), swap-AB schedules that put the expert’s large \(N\) dimension on the MMA’s M side, and fixed-capacity “masked” layouts that keep shapes static for CUDA Graphs. The ahead-of-time compile is cheap enough to run in the agent loop for every candidate tile shape. Grep the PTX for the instruction you expect before paying for GPU time.
8. Shapes decide everything: padding, stragglers, and the roofline
One active expert in FP8 holds \(3HI\) bytes of weights and does \(2\cdot n_e\cdot 3HI\) FLOPs for \(n_e\) rows. Its arithmetic intensity against weight bytes is simply
\[ \text{FLOP per weight byte} \approx 2\,n_e \quad(\text{FP8 weights}),\qquad \approx n_e \quad(\text{BF16 weights}). \]
An H100 SXM’s ridge point is roughly \(1979\ \text{TFLOP/s} / 3.35\ \text{TB/s}\approx 590\) FLOP/byte for dense FP8. An expert needs hundreds of rows before it becomes compute-bound.
# file: moe_shapes.py
# Back-of-envelope numbers for one MoE layer: tokens per expert, padding waste of a block-aligned
# grouped GEMM, and arithmetic intensity of the expert weights (the number that decides whether
# the layer is bandwidth-bound or compute-bound).
import numpy as np
rng = np.random.default_rng(0)
def routed_counts(tokens, k, E, skew):
"""Expert loads for `tokens` tokens choosing k distinct experts; skew=0 is uniform, >0 is Zipf-like."""
p = 1.0 / np.arange(1, E + 1) ** skew
p /= p.sum()
counts = np.zeros(E, dtype=np.int64)
for _ in range(tokens):
counts[rng.choice(E, size=k, replace=False, p=p)] += 1
return counts
def padded_rows(counts, block_m):
return int(((counts + block_m - 1) // block_m * block_m).sum())
# DeepSeek-V3-like layer: hidden 7168, expert intermediate 2048, 256 routed experts, top-8.
H, I, E, K = 7168, 2048, 256, 8
w_bytes = 3 * H * I * 1 # one expert's W1, W3, W2 in FP8 (1 byte/param)
print(f"one expert: {3*H*I/1e6:.1f} M params = {w_bytes/2**20:.0f} MiB in FP8\n")
print(f"{'tokens':>6} {'skew':>4} {'tok/exp':>7} {'active':>6} {'max':>5} {'BLOCK_M':>7} {'rows':>7} {'waste':>6} "
f"{'FLOP/B(w)':>9}")
for tokens in (1, 16, 128, 1024, 8192):
for skew in (0.0, 1.0):
counts = routed_counts(tokens, K, E, skew)
active = int((counts > 0).sum())
for block_m in (16, 64):
rows = padded_rows(counts, block_m)
useful = tokens * K
# FLOPs actually needed vs bytes of expert weights that must be streamed (active experts only)
ai = 2 * useful * 3 * H * I / (active * w_bytes)
print(f"{tokens:>6} {skew:>4} {useful/E:>7.2f} {active:>6} {int(counts.max()):>5} {block_m:>7} {rows:>7} "
f"{1-useful/rows:>6.0%} {ai:>9.1f}")python moe_shapes.pyOutput:
one expert: 44.0 M params = 42 MiB in FP8
tokens skew tok/exp active max BLOCK_M rows waste FLOP/B(w)
1 0.0 0.03 8 1 16 128 94% 2.0
1 0.0 0.03 8 1 64 512 98% 2.0
1 1.0 0.03 8 1 16 128 94% 2.0
1 1.0 0.03 8 1 64 512 98% 2.0
16 0.0 0.50 103 4 16 1648 92% 2.5
16 0.0 0.50 103 4 64 6592 98% 2.5
16 1.0 0.50 80 13 16 1280 90% 3.2
16 1.0 0.50 80 13 64 5120 98% 3.2
128 0.0 4.00 252 11 16 4032 75% 8.1
128 0.0 4.00 252 11 64 16128 94% 8.1
128 1.0 4.00 205 110 16 3632 72% 10.0
128 1.0 4.00 205 110 64 13248 92% 10.0
1024 0.0 32.00 256 47 16 10032 18% 64.0
1024 0.0 32.00 256 47 64 16384 50% 64.0
1024 1.0 32.00 256 774 16 10176 19% 64.0
1024 1.0 32.00 256 774 64 20352 60% 64.0
8192 0.0 256.00 256 308 16 67424 3% 512.0
8192 0.0 256.00 256 308 64 73472 11% 512.0
8192 1.0 256.00 256 6455 16 67504 3% 512.0
8192 1.0 256.00 256 6455 64 73024 10% 512.0Reading the table, row by row:
- 1 to 16 tokens (decode): 8 to 103 experts are active, each with about one row. Intensity is 2 to 3 FLOP/byte, deep in bandwidth-bound territory. The padding “waste” of 90–98% is irrelevant: those MMA rows are free while the kernel waits on HBM. The real cost is \(\text{active}\times 42\) MiB of weights per layer. Below about 128 tokens, each extra token can wake experts that were idle, so per-token cost falls more slowly than in a dense model. Once nearly every expert is active, extra decode tokens add no weight bytes at all, and batching becomes almost free until the ridge.
- 1024 tokens: every expert is active, with 32 rows on average. Padding to 64 now wastes half the tensor-core work, so
BLOCK_M=16or 32 is better, and still bandwidth-bound at 64 FLOP/byte. - 8192 tokens (prefill): 256 rows per expert reach the ridge, and padding waste drops to 3–11%. This is where
wgmma-sized tiles, TMA, and warp specialization pay off, as for any large GEMM. - Skew (
skew=1.0): the hottest expert takes 774 of 1024 and 6455 of 8192 routed rows, so one expert holds most of the batch. A grid with one tile perBLOCK_Mrows still runs, but the hot expert’s tiles form a long tail if tiles are scheduled expert by expert. Persistent kernels that hand out tiles from a global queue, and the load-balancing router from section 4, both exist because of this column.
For expert parallelism across GPUs, the same arithmetic applies per GPU, with \(E\) replaced by the local expert count and \(T\) by the tokens the all-to-all delivers to that GPU.
9. Where this sits in production
| Component | Role |
|---|---|
vLLM / SGLang fused_moe (Triton) |
The gather/scatter design of section 6: moe_align_block_size, a fused_moe_kernel with FP8/INT8/INT4 weight variants, tuned configs per \((E, N, \text{GPU})\), and fused silu_and_mul and moe_sum kernels |
torch.nn.functional.grouped_mm |
Public PyTorch grouped GEMM over 2-D/3-D operands with offs. Used for MoE training in torchtitan |
| CUTLASS grouped GEMM | Many problem sizes in one persistent launch. Problem shapes and pointer arrays can live on the device |
| DeepGEMM grouped GEMMs | FP8 “m-grouped” kernels with a contiguous layout (each expert’s rows padded to the M-block alignment, for prefill and training) and a masked layout (fixed capacity per expert plus a valid-row count, for CUDA-Graph decode) |
| Triton grouped GEMM tutorial | A persistent Triton kernel that walks a device-side list of GEMM problems |
| MegaBlocks | “Dropless” MoE as block-sparse GEMMs. It removed the capacity-factor token dropping of earlier MoE systems |
| DeepEP | Expert-parallel all-to-all dispatch and combine. High-throughput kernels for prefill and training, and low-latency RDMA kernels for decode |
The capacity-factor approach deserves a sentence because generated code still produces it: early MoE systems gave every expert a fixed number of slots and dropped tokens that overflowed. That gives static shapes, but it silently changes the model’s output under load. Modern inference stacks are dropless: every routed row is computed.
10. Verification and benchmark plan
Functional
| Check | Catches |
|---|---|
| Loop vs dense vs fused vs grouped-GEMM agree | Index bugs in any one path |
Experts with 0 rows, 1 row, BLOCK_M−1, BLOCK_M, BLOCK_M+1 rows |
Tile-boundary and empty-expert bugs (reading weights of an expert with no tile) |
| \(T=1\) and \(T\cdot k < E\) | The decode path, and any per-token fallback |
| One expert receives every token | Grid sizing, the worst-case padded length, stragglers |
| Router: biased selection, unbiased weights, renormalization, scale factor | Silent quality bugs (section 4) |
| Sentinel rows never written; output rows written exactly once per slot | Races and out-of-bounds stores (run compute-sanitizer on a GPU) |
| BF16/FP8 inputs vs FP32 reference with tolerances | Accumulating in low precision; FP8 scale handling |
Performance
| Metric | Target |
|---|---|
| Decode: achieved bandwidth = active-expert weight bytes / layer time | Approaches HBM peak |
| Prefill: achieved TFLOP/s on useful rows only (exclude padding) | Approaches a dense GEMM of the same total size |
| Alignment + router + combine share of layer time | Small at prefill. Watch it at decode, where it can rival the GEMMs |
| Sweep over \(T\) and routing skew | No cliff where the tile config changes |
Benchmark with real router outputs captured from the model on real prompts, not randint expert IDs. Uniform random routing hides the straggler problem that real traffic creates.
11. Prompting an AI assistant for an MoE kernel
Generated MoE code tends to loop over experts in Python or to assume uniform loads. State the contract:
- Shapes and layout:
x [T, H],w13 [E, 2I, H]andw2 [E, H, I]innn.Linearlayout,topk_ids/topk_w [T, k], flat row \(r = t\cdot k + j\). - Metadata: produce
sorted_ids,expert_ids, and the padded count with a sentinel. Empty experts get no tiles. - Kernel: one launch per GEMM, tile-to-expert mapping, gathered A rows (
r // kfor GEMM1), scattered C rows, router weight in the GEMM2 epilogue, FP32 accumulation. - Router semantics if DeepSeek-style: sigmoid, bias for selection only, group-limited top-k, renormalize, scale.
- Tests first: interpreter-mode tests against a loop and a dense reference across the cases in section 10, plus an AOT compile for the target with a check that the expected MMA instruction appears.
- Metric: GB/s of weights at decode, useful TFLOP/s at prefill.
Pitfalls
| Pitfall | Result |
|---|---|
| Python loop over experts | \(E\) launches per layer, idle GPU between them |
| Capacity factor with token dropping | Static shapes, silently different outputs under load |
| Gate weights from bias-adjusted scores | Plausible but wrong outputs, no error |
Forgetting renormalization or routed_scaling_factor |
Output scaled wrongly by a per-token factor |
| Padding rows not masked on store | Writes past the output buffer, or over real rows |
| Launching tiles for empty experts | Weight bytes read for nothing, which hurts most at decode |
| One tile config for all batch sizes | Prefill-tuned 64-row tiles at decode, or mma.sync tiles at prefill |
| Benchmarking with uniform random routing | Hides hot-expert stragglers |
Checkpoint
- Explain why an MoE layer is \(E\) ragged GEMMs whose sizes are only known on the device.
- Write the
sorted_ids/expert_idscontract and compute it with a histogram and a scan. - Describe how a fused kernel gathers A rows for GEMM1 versus GEMM2, and where the router weight is applied.
- State the DeepSeek-V3 routing rule, and which scores the gate weights come from.
- Use “FLOP per weight byte \(\approx 2n_e\)” to say whether a given MoE layer is bandwidth- or compute-bound.
- Explain why 16-row tiles cannot use Hopper
wgmma, and name two ways decode kernels deal with small row counts.
Try this
- Change
moe_align.pyto computesorted_idswithoutargsort: one pass for the histogram, an exclusive scan, and an atomic-style running counter per expert. Check it against the current output. - Add an
expert_mapargument for expert parallelism: experts not on this GPU getexpert_ids = -1, and the kernel skips those tiles. Test with half the experts remote. - Fuse
silu(gate) * upinto the GEMM1 epilogue by computing the gate and up halves as two accumulators in the same program. Confirm the interpreter test still passes. - Run
compile_check.pywithBLOCK_M ∈ {16, 32, 64, 128}and record which tiles getwgmma, and the shared memory each needs. - On a GPU, capture real
topk_idsfrom a small MoE model, then sweep \(T\in\{1, 8, 64, 512, 4096\}\) and plot achieved GB/s and TFLOP/s forBLOCK_M=16andBLOCK_M=64.
Summary
An MoE layer replaces one GEMM with many small, ragged, data-dependent GEMMs. Making it fast starts with metadata: a histogram, a scan, and a scatter turn the router’s topk_ids into a block-aligned table that maps every tile to one expert. A single grouped-GEMM launch then either reads a permuted copy of the activations or gathers rows through the table, and scatters results back to token order with the router weight folded in. Routing semantics are part of correctness, especially the selection-only bias in DeepSeek-style routers. Shape arithmetic tells you which regime you are in: decode streams weights at about 2 FLOP per byte and needs bandwidth-oriented kernels, while prefill reaches the ridge only when experts see hundreds of rows. Verify the indexing on a CPU with the Triton interpreter, compile ahead of time for the target, and benchmark with real routing traces.
Sources (for further reading)
- DeepSeek-AI, “DeepSeek-V3 Technical Report” (arXiv:2412.19437): auxiliary-loss-free balancing with a selection bias, sigmoid gating, node-limited routing. Model
config.jsonfiles on Hugging Face for DeepSeek-V3, Kimi K2, Qwen3-235B-A22B, gpt-oss-120b, and Mixtral-8x7B (expert counts and routing fields). - vLLM source:
vllm/model_executor/layers/fused_moe/moe_align_block_size.py(docstring example reproduced in section 5) andfused_moe.py(fused_moe_kernel, the per-token assignment heuristic, tuned configs). - PyTorch documentation:
torch.nn.functional.grouped_mm(2.14); torchtitan MoE (_run_experts_grouped_mm). - Gale, Narayanan, Young, Zaharia, “MegaBlocks: Efficient Sparse Training with Mixture-of-Experts” (MLSys 2023, arXiv:2211.15841).
- Fedus, Zoph, Shazeer, “Switch Transformers” (arXiv:2101.03961): capacity factor and token dropping.
- DeepSeek-AI, DeepGEMM repository (m-grouped contiguous and masked FP8 GEMMs) and DeepEP repository (expert-parallel dispatch and combine kernels).
- NVIDIA CUTLASS documentation: grouped GEMM. PTX ISA:
wgmma.mma_asyncshapes (m64nNk16). - Triton documentation: tutorial “Group GEMM”,
TRITON_INTERPRET=1, andtriton.compilewithGPUTarget.