Batching & GPU kernels

Running the model & the world

The idea

A GPU is a throughput machine. It doesn’t want a thousand small, scattered matrix multiplies; it wants one big contiguous one, where thousands of compute units stream through memory in order. Most of inference engineering is reshaping work until it looks like that: serving systems batch many users’ requests into one forward pass, and kernels are written so the data each multiply needs sits side by side in memory.

Mixture-of-Experts makes this concrete. After the router assigns each token to its top-k experts, the naive layout is a mess: tokens for expert 7 scattered among tokens for expert 0. The standard fix is to sort tokens by expert, record the offset where each expert’s slice begins, and then run one grouped multiply that walks expert by expert through contiguous blocks - gather, group, multiply, scatter back.

Reading a real PR with this page

mlx#4390 (“[CUDA] Grouped mm”) is exactly this story in production code. In our vocabulary: MLX’s MoE path used an index-taking gather-multiply, which meant recomputing token offsets for every projection of the expert MLP (gate, up, down - modern MLPs come in threes, see below). The PR adds a grouped_mm operation that takes activations already sorted by expert, the per-expert token offsets, and the stacked expert weights - one op, one contiguous pass per projection, ported from NVIDIA’s grouped-matmul kernel. Every concept in that sentence is a page on this site: MoE routing, the MLP projections, and why the GPU cares.

The story, animated

Why a GPU hates crumbs, what the router leaves mixed in memory, scattered reads vs one sorted pass, the offsets a grouped multiply walks, why gate / up / down reuse them, and the same stacking idea at the serving layer — slide by slide.

Watch

From the full film - the fused-kernel bottleneck · 169s, from the film series.

Where it sits

Worth having first: The machine underneath · Mixture of Experts