PyTorch CPU grouped matmul may stop looping over experts
An AMD proposal would run mixture-of-experts grouped GEMMs as one oneDNN call on Zen CPUs, with about a 2x gain on published expert shapes.
On CPU, PyTorch's grouped matrix multiply still has no kernel of its own. Unlike the CUDA and MPS paths, a call falls back to one ordinary matmul per expert. An AMD engineer has proposed a backend that instead sends the whole call through oneDNN's grouped matmul, which the zen64 integration executes on AMD Zen CPUs with ZenDNN.
The target is mixture-of-experts inference on CPU servers. A layer routes each token to a few of many experts, then runs two grouped multiplies, the gate and up projection and the down projection. Expert weights would be packed once at load, through a new reorder operation modeled on the existing linear-weight reorder, and reused across calls. Call sites would not change, and dense weights that are not packed would keep today's fallback. The zen64 pieces depend on a separate PyTorch build integration already in progress. oneDNN 3.14 already has the grouped primitive this path would call, including a layout that can be pre-packed and reused.
On expert shapes from Qwen3-30B-A3B, Mixtral-8x7B, and gpt-oss-20b (42 cases, bfloat16, 32 cores of an AMD EPYC 9755), the author reports the new backend beat the fallback in 41 of 42 cases: a 1.99x geometric mean on the same build, and 1.85x against stock oneDNN.