Grouped GEMM
Also known as grouped matmul, GroupedGEMM, ragged matmul
In plain English
A grouped GEMM runs many small matrix multiplies, one per expert, in a single launch so a mixture-of-experts layer does not pay per-expert overhead.
Technical definition
A grouped GEMM is a kernel that executes a set of independent matrix multiplications with different row counts in one launch, used in MoE layers where each expert receives a ragged group of routed tokens.
Engineering details
A mixture-of-experts layer produces irregular groups of tokens for each expert, and those ragged groups must be reshaped into something the matrix units can consume efficiently. On TPU the second grouped matmul version removes redundant tile computation, sizes transfers to the number of valid rows rather than the padded maximum, triple-buffers expert weights so the next group is in flight while the current one computes, and fuses group metadata generation into the kernel. The irregular token permutation moved to SparseCore. For small batches a dedicated path builds one-hot matrices and uses ordinary matmuls to permute and unpermute tokens because the general ragged path costs more than the work it arranges. With DP attention and expert parallelism, the kernel all-gathers token activations and routing metadata first and reduce-scatters weighted outputs back to each attention rank.
Why it matters
MoE efficiency on any accelerator reduces to how well the grouped GEMM tolerates ragged group sizes and how much of the routing work can be hidden. On TPU the added constraint is MXU tile geometry, so expert widths and token counts also need to fill 256-wide tiles.
How to read it in InferenceX
The SparseCore permutation rewrite reported 12% higher 8k1k throughput on Ironwood, the small-batch one-hot path gained 7.3% at concurrency 64 and 5.1% at 128, and merging the routing all-gathers saved roughly 80 microseconds per layer, about 4.64 ms per forward pass across DeepSeek-V3’s 58 layers.
Source material
See the concept in real benchmarks
TPU Inference Externalization Full Steam Ahead
InferenceX, Up to 50% Better Performance per Dollar, Rapid Externalization of TPU stack, Growing Customer Base, Ironwood, TPUv8i, Reducing CUDA Moat
DeepSeekV4 1.6T Day 0 to Day 43 Performance Over Time — Huawei, GB300 NVL72, MI355X, B200
Day 0 Inference Performance, InferenceX, 100x performance improvement in 26 Days, Cost per Million Tokens, Huawei 950DT Inference Trace Analysis