Skip to content
Markdown

MoE expert backends and grouped-GEMM kernels

Scope: the expert-compute kernel layer of a Mixture-of-Experts model. How the expert FFNs are scheduled on the GPU (loop, batched GEMM, grouped GEMM, block-sparse, fully fused), the packed 3D weight layout those kernels require and the checkpoint conversion that produces it, the six backends HuggingFace Transformers registers behind experts_implementation, and the silent-corruption traps in group offsets, sort inversion, and gate_up packing. This page sits below MoE sparse scaling, MoE routing and load balancing and expert parallelism for inference: those pages decide which expert a token goes to and which GPU holds it, this one covers what happens once the token arrives.

Code blocks come in two kinds. The torch and CLI snippets are unexecuted reference templates against upstream APIs (no GPU here); pin versions and verify against the repo. The numpy blocks are executed on stock python3 + numpy 2.4.6, assert their own results, and include the adversarial cases: an off-by-one in the group offsets, a skipped permutation inversion, and a mis-read gate_up layout. The roofline block is explicitly a model, not a measurement.

What it is

Every MoE backend implements the same MoE function, subject to its numerical precision: project each token through the selected experts, then combine their outputs with the routing weights. Backends differ in layout, schedule, precision, communication, and fusion boundaries; those choices decide tensor-core utilization and memory traffic. Five designs, in increasing order of sophistication:

Design Schedule Weight traffic Wins when
Loop over experts (eager) one GEMM per expert that received tokens one logical expert operand per active expert debugging, correctness reference
Batched GEMM (batched_mm) replicate each token k times, gather the selected expert's weights per token-expert pair, one torch.bmm expert operand materialized per pair small token counts (decode)
Grouped GEMM (grouped_mm, SonicMoE) sort tokens by expert id, one GEMM per contiguous group via offsets, then invert the permutation one logical expert operand per active expert large token counts (prefill, training), memory-constrained
Block-sparse (MegaBlocks) reformulate the whole layer as block-sparse matmuls sparse expert operands avoids token dropping and dense capacity padding
Fully fused (DeepGEMM Mega MoE) dispatch, up-proj, activation, down-proj, combine in one kernel fused, communication overlapped Blackwell FP4 expert parallelism

MegaBlocks is the design worth understanding for its own sake: it exists because the older formulations "force a tradeoff between model quality and hardware efficiency, as users must choose between dropping tokens from the computation or wasting computation and memory on padding", and block-sparse kernels "never drop tokens" while still mapping efficiently to the hardware.8

In HuggingFace Transformers, this layer is src/transformers/integrations/moe.py (588 lines as of 2026-07-13). It is not a class hierarchy. There is no MoeExperts base class: use_experts_implementation(...) is a decorator applied to each model's own experts class, which swaps that class's forward for a lookup in ALL_EXPERTS_FUNCTIONS, an ExpertsInterface(GeneralInterface) registry.1 The decorator carries four flags that reconcile per-model weight layouts with the shared kernels: is_concatenated, is_transposed, has_bias, has_gate. Roughly 46 model directories under src/transformers/models carry the decorator (GitHub code search, 2026-07-13), Mixtral and gpt-oss and the DeepSeek and Qwen MoE families among them.

The six registered backends

The Transformers blog announcing this feature says three backends exist.10 Upstream main registers six. Read the source, not the blog:

experts_implementation Kernel Registered in Hard requirement
eager per-expert loop special-cased in ExpertsInterface.get_interface none; not fullgraph-compilable
batched_mm torch.bmm over duplicated weights ALL_EXPERTS_FUNCTIONS none
grouped_mm torch.nn.functional.grouped_mm / torch._grouped_mm ALL_EXPERTS_FUNCTIONS see the torch gating table below
deepgemm TMA-aligned grouped GEMM, kernels-community/deep-gemm ALL_EXPERTS_FUNCTIONS CUDA, SM90+
sonicmoe CuteDSL grouped GEMM from Dao-AILab/quack, via kernels-community/sonic-moe ALL_EXPERTS_FUNCTIONS CUDA SM90+, gated activation
deepgemm_megamoe one fused kernel: EP dispatch + gated MLP + EP combine FP8ExpertsInterface only (integrations/finegrained_fp8.py) Blackwell SM100+, FP4 weights, a process group

The last row is the one that repays attention. DeepGEMM's Mega MoE "fuses and overlaps EP dispatch, linear 1 (FP8xFP4), SwiGLU, linear 2 (FP8xFP4), and EP combine into a single mega-kernel, overlapping NVLink communication and tensor core computation".9 It is not in the base ExpertsInterface._global_mapping; it is registered only in the FP8 interface, and get_correct_experts_implementation validates a requested name against the union of both mappings plus eager.1 The default when nothing is requested is grouped_mm, silently downgraded to eager if _grouped_mm_can_dispatch() fails.1

The packed layout, and where it comes from

MoE checkpoints serialize each expert as its own tensors (model.layers.3.mlp.experts.0.gate_proj.weight, ...experts.1.gate_proj.weight, and so on up to the expert count, which is 256 routed experts on DeepSeek-V3), which is exactly the layout the kernels cannot use: they want one contiguous 3D tensor per projection. Transformers PR #41580 ("Refactor weight loading", merged 2025-11-13) introduced a generic WeightConverter to bridge the two. For Mixtral, MergeModulelist(dim=0) stacks the per-expert tensors and Concatenate(dim=1) fuses w1 and w3 into one mlp.experts.gate_up_proj; every op declares a reverse_op, so SplitModulelist puts the checkpoint back the way it found it and save_pretrained() stays HF-compatible.5 The Experts Backend interface itself landed later, in PR #42697 ("batched and grouped experts implementations", merged 2026-01-05).

The layout is per-model, and the flags are what reconcile it. GptOssExperts declares gate_up_proj as nn.Parameter[num_experts, hidden_size, 2 * intermediate_size] and down_proj as [num_experts, intermediate_size, hidden_size], and is decorated @use_experts_implementation(is_concatenated=False, is_transposed=True, has_bias=True). is_concatenated=False means the gate and up columns are interleaved (gate_up[..., ::2] and gate_up[..., 1::2]), not the [gate | up] concatenation the default flag assumes and the Mixtral WeightConverter produces.4 Reading one layout with the other assumption produces an output of the right shape, with no NaN and no exception. The block in How to develop with it quantifies exactly how wrong it is.

Why use it

  • It is the difference between busy and idle tensor cores, and it is orthogonal to routing. A perfectly balanced router still runs slowly through the wrong kernel. Capacity factor, EPLB and token dropping (covered in MoE routing and load balancing) do not fix a kernel that pays a 16-row tile for a 1-row group.
  • The operand-materialization tradeoff is directional. batched_mm gathers a weight slice per token-expert pair, whereas grouped_mm presents one expert matrix with a variable-size token group. At 96 tokens with top-2 routing, the executed block counts 192 pair-wise weight operands against 7 active-expert operands. This is a logical count, not a claim that hardware reads each weight tile from HBM exactly once.
  • It is a one-line, reload-free switch. model.set_experts_implementation("eager") changes the kernel on a live model, which makes bisecting a numerical bug against the reference implementation cheap.
  • The fused backends buy overlap, not just speed. Mega MoE hides the NVLink dispatch and combine of expert parallelism inside the GEMM rather than around it, which is the same idea as the overlap techniques in comms-compute overlap but pushed into a single kernel.

When to use it (and when not)

  • Leave grouped_mm as the default for training and prefill: no weight duplication, one tile read per active expert, and the memory-efficient choice on a constrained card.
  • Do not hand-pick batched_mm for decode. Transformers already does it (see below), and doing it manually on top of the automatic switch buys nothing.
  • Use eager as a correctness oracle, not in production. It is the only backend that cannot be compiled with fullgraph=True, because finding which experts were hit is a data-dependent operation.2
  • Use deepgemm for DeepSeek-style FP8/FP4 checkpoints on Hopper or Blackwell; sonicmoe for bf16 gated experts on Hopper, which upstream describes as state-of-the-art throughput there, "especially for training".2
  • deepgemm_megamoe has a narrow domain: Blackwell SM100+, FP4-packed weights with UE8M0 scales, running under expert parallelism. Outside that domain it is not an option, and set_experts_implementation refuses to flip in or out of it at runtime.
  • Do not reach for this layer to fix an imbalance problem. A hot expert stalls the layer no matter which kernel runs it. That is routing and EP placement work.
  • Do not expect experts_implementation to give you kernels-hub kernels, or vice versa. They are two separate mechanisms; see Two mechanisms, one blog paragraph.

Architecture

The checkpoint layout, the packing step, and the kernel dispatch are three distinct stages, and each is a place where a model can be silently wrong.

flowchart TB
  CKPT["Checkpoint on disk<br/>experts.0.gate_proj ... experts.255.gate_proj"]
  CONV["WeightConverter (PR #41580)<br/>MergeModulelist(dim=0) + Concatenate(dim=1)"]
  PACK["Packed 3D parameter<br/>gate_up_proj[E, H, 2I] &middot; down_proj[E, I, H]"]
  FLAGS["@use_experts_implementation flags<br/>is_concatenated &middot; is_transposed<br/>has_bias &middot; has_gate"]
  DISP{"ExpertsInterface<br/>config._experts_implementation"}
  EAG["eager<br/>loop over hit experts"]
  BAT["batched_mm<br/>torch.bmm, weights duplicated per pair"]
  GRP["grouped_mm<br/>sort by expert, offsets, invert permutation"]
  DG["deepgemm / sonicmoe<br/>external grouped-GEMM kernels"]
  MEGA["deepgemm_megamoe<br/>EP dispatch + MLP + combine fused"]
  CKPT --> CONV --> PACK --> DISP
  FLAGS --> DISP
  DISP --> EAG
  DISP --> BAT
  DISP --> GRP
  DISP --> DG
  DISP --> MEGA
  GRP -.->|"decode stage, GPU only"| BAT

The dotted edge is the operational surprise: on GPU, a model loaded with grouped_mm switches itself to batched_mm for the decode stage of generation and restores the original backend when generation finishes. On CPU it stays on grouped_mm, which upstream considers the more efficient backend at every input size there. The switch recurses into sub-configs, so a vision-language model's text_config experts switch too, and only grouped_mm entries are touched.23 Anyone reading a decode profile and looking for the grouped GEMM they configured will not find it.

How to use it

Loading, switching, and reading back the backend (reference template, unexecuted; verify against transformers main, which moves fast here):

# Reference template. Requires transformers (v5 line) + torch >= 2.9 and a CUDA GPU.
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen1.5-MoE-A2.7B",
    dtype="bfloat16",
    experts_implementation="batched_mm",   # default is grouped_mm
)

model.set_experts_implementation("eager")  # runtime switch, no reload
model.get_experts_implementation()         # {"": "eager"}; one entry per sub-config

Backbone-specific selection takes a dict, keyed by sub-config name, which is how a VLM runs grouped_mm on its text tower and eager on its vision tower.2

How to develop with it

The contract that matters is not throughput, it is equivalence. All backends must compute the same function, and the two ways to break grouped GEMM do not raise. The block below implements the same MoE layer three ways (eager loop, batched GEMM with duplicated weights, grouped GEMM with sort/offsets/inverse-permutation), asserts they agree, then injects the two real bugs: group offsets shifted by one row, and a forgotten permutation inversion. The routing is adversarial: one expert receives zero tokens (its group is empty and must issue no GEMM) and one is hot.

# Runnable on system python3 (numpy 2.4.6). Core contract of the expert-backend layer:
# eager, batched_mm and grouped_mm are three schedules of ONE function and must agree to
# floating-point rounding. The two classic grouped-GEMM bugs corrupt tokens without raising.
import numpy as np

rng = np.random.default_rng(0)
E, H, I, T, K = 8, 16, 32, 96, 2  # experts, hidden, intermediate, tokens, top-k


def swiglu(gate: np.ndarray, up: np.ndarray) -> np.ndarray:
    return (gate / (1.0 + np.exp(-gate))) * up


def eager(x, gu, dn, idx, w):
    """`eager`: loop over experts, gather that expert's tokens, one GEMM each."""
    out = np.zeros_like(x)
    pair_e, pair_t = idx.reshape(-1), np.repeat(np.arange(x.shape[0]), K)
    for e in range(E):
        rows = np.flatnonzero(pair_e == e)
        if rows.size == 0:
            continue  # expert got no tokens: no GEMM issued
        h = x[pair_t[rows]] @ gu[e]                      # (n, 2I)
        h = swiglu(h[:, :I], h[:, I:]) @ dn[e]           # (n, H)
        np.add.at(out, pair_t[rows], h * w.reshape(-1)[rows, None])
    return out


def batched_mm(x, gu, dn, idx, w):
    """`batched_mm`: replicate each token K times, GATHER the selected expert weights
    (duplicating them per token-expert pair), one batched GEMM over S = T*K."""
    pair_e = idx.reshape(-1)
    xs = np.repeat(x, K, axis=0)[:, None, :]             # (S, 1, H)
    h = (xs @ gu[pair_e])[:, 0, :]                       # (S, 2I), weights duplicated S times
    h = (swiglu(h[:, :I], h[:, I:])[:, None, :] @ dn[pair_e])[:, 0, :]
    return (h * w.reshape(-1)[:, None]).reshape(x.shape[0], K, -1).sum(axis=1)


def grouped_mm(x, gu, dn, idx, w, off_bug=0, skip_unsort=False):
    """`grouped_mm`: SORT the S pairs by expert id, one GEMM per contiguous group
    delimited by offsets, then INVERT the permutation. off_bug/skip_unsort inject the
    two silent-corruption bugs."""
    pair_e, S = idx.reshape(-1), x.shape[0] * K
    perm = np.argsort(pair_e, kind="stable")             # sort tokens by expert
    xs = np.repeat(x, K, axis=0)[perm]
    counts = np.bincount(pair_e, minlength=E)
    offs = np.concatenate(([0], np.cumsum(counts)))      # group boundaries, len E+1
    ys = np.zeros((S, x.shape[1]))
    for e in range(E):
        lo, hi = offs[e] + off_bug, offs[e + 1] + off_bug   # off-by-one -> shifted groups
        lo, hi = max(lo, 0), min(hi, S)
        if lo >= hi:
            continue
        h = xs[lo:hi] @ gu[e]                            # ONE GEMM for the whole group
        ys[lo:hi] = swiglu(h[:, :I], h[:, I:]) @ dn[e]
    ys = ys * w.reshape(-1)[perm][:, None]
    if not skip_unsort:                                  # invert the sort permutation
        inv = np.empty_like(perm)
        inv[perm] = np.arange(S)
        ys = ys[inv]
    return ys.reshape(x.shape[0], K, -1).sum(axis=1)


x = rng.standard_normal((T, H))
gu = rng.standard_normal((E, H, 2 * I)) * 0.1           # concatenated [gate | up]
dn = rng.standard_normal((E, I, H)) * 0.1
logits = rng.standard_normal((T, E))
logits[:, 5] = -1e9                                      # adversarial: expert 5 gets 0 tokens
logits[:, 1] += 4.0                                      # adversarial: expert 1 is hot
idx = np.argsort(-logits, axis=1)[:, :K]
w = rng.random((T, K))
w /= w.sum(axis=1, keepdims=True)
counts = np.bincount(idx.reshape(-1), minlength=E)
assert counts[5] == 0 and counts[1] == counts.max(), "adversarial routing not hit"

ref = eager(x, gu, dn, idx, w)
bat = batched_mm(x, gu, dn, idx, w)
grp = grouped_mm(x, gu, dn, idx, w)
assert np.allclose(ref, bat, atol=1e-12) and np.allclose(ref, grp, atol=1e-12)

# --- Adversarial: both bugs are SILENT. No exception, no NaN, wrong tokens. ---
bad_off = grouped_mm(x, gu, dn, idx, w, off_bug=1)       # offsets shifted by one row
bad_srt = grouped_mm(x, gu, dn, idx, w, skip_unsort=True)  # forgot to invert the permutation
wrong_off = int((~np.isclose(bad_off, ref, atol=1e-9)).any(axis=1).sum())
wrong_srt = int((~np.isclose(bad_srt, ref, atol=1e-9)).any(axis=1).sum())
assert np.isfinite(bad_off).all() and np.isfinite(bad_srt).all(), "bugs raise no error"
assert wrong_off > 0 and wrong_srt > 0

# Logical weight operands expose the batched/grouped tradeoff: batched_mm gathers one slice
# per token-expert PAIR; grouped_mm presents one operand per active expert. This does not count
# physical HBM transactions, tiling, or cache reuse inside the GEMM kernel.
tiles_batched, tiles_grouped = T * K, int((counts > 0).sum())

print("expert token counts  :", counts.tolist(), "(expert 5 empty, expert 1 hot)")
print("grouped_mm vs eager  : max|d| = %.2e  (identical GEMM grouping: bit-exact)" % np.abs(ref - grp).max())
print("batched_mm vs eager  : max|d| = %.2e  (same function, different reduction order)" % np.abs(ref - bat).max())
print("BUG offsets+1        : %2d/%d token rows wrong, max|d| = %.3f, exception: none" % (wrong_off, T, np.abs(bad_off - ref).max()))
print("BUG unsort skipped   : %2d/%d token rows wrong, max|d| = %.3f, exception: none" % (wrong_srt, T, np.abs(bad_srt - ref).max()))
print("logical expert-weight operands: batched_mm %d (one per pair) vs grouped_mm %d (one per active expert)" % (tiles_batched, tiles_grouped))

Executed output:

expert token counts  : [16, 96, 25, 13, 19, 0, 14, 9] (expert 5 empty, expert 1 hot)
grouped_mm vs eager  : max|d| = 0.00e+00  (identical GEMM grouping: bit-exact)
batched_mm vs eager  : max|d| = 8.33e-17  (same function, different reduction order)
BUG offsets+1        :  6/96 token rows wrong, max|d| = 0.223, exception: none
BUG unsort skipped   : 96/96 token rows wrong, max|d| = 0.205, exception: none
logical expert-weight operands: batched_mm 192 (one per pair) vs grouped_mm 7 (one per active expert)

The off-by-one is the dangerous one. It corrupts 6 of 96 token rows, leaves the other 90 exactly right, raises nothing, and produces finite plausible numbers. A smoke test passes. A perplexity eval on a small sample may pass. A token-level numerical comparison against eager catches it; use tolerances appropriate to the production dtype and kernel reduction order.

The layout trap and the hard contracts

The second silent-corruption class is the packed layout itself, and the fix is a boolean flag. The block below packs a tensor the gpt-oss way (gate and up columns interleaved), reads it back with the default concatenated assumption, and quantifies the damage. It then asserts the two hard bounds the distributed path enforces: GroupedGemmParallel raises when the expert count is not divisible by the world size, and top-k cannot exceed the expert count.

# Runnable on system python3 (numpy 2.4.6). Two contracts the packed layout imposes:
# (1) the gate_up split is per-model (interleaved vs concatenated) and reading it wrong is
# silent; (2) expert-parallel sharding and top-k have hard divisibility/range bounds.
import numpy as np

rng = np.random.default_rng(7)
H, I = 8, 6  # hidden, intermediate


def split_concatenated(gate_up: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """transformers default (`is_concatenated=True`): gate_up is [gate | up]."""
    return gate_up[..., : gate_up.shape[-1] // 2], gate_up[..., gate_up.shape[-1] // 2 :]


def split_interleaved(gate_up: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """GptOssExperts._apply_gate (`is_concatenated=False`): [g0, u0, g1, u1, ...]."""
    return gate_up[..., ::2], gate_up[..., 1::2]


def swiglu(gate: np.ndarray, up: np.ndarray) -> np.ndarray:
    return (gate / (1.0 + np.exp(-gate))) * up


# A checkpoint packed the gpt-oss way: columns alternate gate, up, gate, up, ...
gate = rng.standard_normal((H, I))
up = rng.standard_normal((H, I))
gate_up_interleaved = np.empty((H, 2 * I))
gate_up_interleaved[:, ::2] = gate
gate_up_interleaved[:, 1::2] = up

g_ok, u_ok = split_interleaved(gate_up_interleaved)      # correct flag
g_bad, u_bad = split_concatenated(gate_up_interleaved)   # is_concatenated=True on a gpt-oss layout
assert np.array_equal(g_ok, gate) and np.array_equal(u_ok, up)

x = rng.standard_normal((4, H))
y_ok = swiglu(x @ gate, x @ up)
y_bad = swiglu(x @ g_bad, x @ u_bad)
assert np.isfinite(y_bad).all(), "the wrong split does not raise and does not produce NaN"
wrong = int((~np.isclose(y_ok, y_bad, atol=1e-9)).sum())
rel = float(np.abs(y_bad - y_ok).max() / np.abs(y_ok).max())

# Shapes match, so nothing downstream complains: the tensors are the same size.
assert y_ok.shape == y_bad.shape == (4, I)


def shard_experts(num_experts: int, world_size: int) -> int:
    """GroupedGemmParallel.shard_tensor: local experts per rank, hard divisibility check."""
    if num_experts % world_size != 0:
        raise ValueError(
            f"Global number of experts must be divisible by number of devices: "
            f"{num_experts} % {world_size} != 0"
        )
    return num_experts // world_size


def route(num_experts: int, top_k: int) -> None:
    if top_k > num_experts:
        raise ValueError(f"top_k ({top_k}) cannot exceed num_experts ({num_experts})")


assert shard_experts(256, 8) == 32
for bad_ws in (3, 7, 24):
    try:
        shard_experts(256, bad_ws)
        raise AssertionError(f"world_size={bad_ws} should have raised")
    except ValueError:
        pass
try:
    route(num_experts=8, top_k=9)
    raise AssertionError("top_k > num_experts should have raised")
except ValueError:
    pass

print("interleaved layout read as concatenated:")
print("  shapes agree:", y_ok.shape == y_bad.shape, "| finite:", bool(np.isfinite(y_bad).all()), "| exception: none")
print("  wrong elements: %d/%d, max relative error %.2f" % (wrong, y_ok.size, rel))
print("EP contract: 256 experts / 8 ranks -> %d local experts" % shard_experts(256, 8))
print("EP contract: 256 experts / 3, 7 or 24 ranks -> ValueError (all three)")
print("routing contract: top_k=9 with num_experts=8 -> ValueError")

Executed output:

interleaved layout read as concatenated:
  shapes agree: True | finite: True | exception: none
  wrong elements: 24/24, max relative error 1.38
EP contract: 256 experts / 8 ranks -> 32 local experts
EP contract: 256 experts / 3, 7 or 24 ranks -> ValueError (all three)
routing contract: top_k=9 with num_experts=8 -> ValueError

Every element is wrong, the shapes agree, and nothing raises. This is the entire job of the is_concatenated flag, and it is why a new model's experts class cannot be decorated by copy-paste.

Why batched_mm wins at decode: an arithmetic model

The decode-stage switch is worth understanding rather than memorising. The block below is a roofline model, not a measurement (there is no GPU here). It does not try to prove the switch; it asks what the switch implies, by pricing the two backends on a DeepSeek-V3-shaped MoE layer and reporting how much fixed overhead grouped_mm's sort, histogram, gather and inverse permutation would have to cost for batched_mm to win.

# Runnable on system python3 (numpy 2.4.6). A ROOFLINE MODEL, not a measurement. No GPU was
# used. It does not "prove" the decode-stage switch; it asks what the switch implies. The
# question: how much fixed overhead must grouped_mm's sort + histogram + gather + inverse
# permutation cost for batched_mm to be the better decode kernel?
import numpy as np

# Stated assumptions, so they can be argued with:
PEAK = 989e12    # H100 SXM dense BF16 tensor-core peak FLOP/s (vendor number, never reached)
BW = 3.35e12     # H100 SXM HBM3 bandwidth, bytes/s
BM = 16          # grouped-GEMM tile rows. vLLM's tuned fused_moe config at M=1 selects
                 # BLOCK_SIZE_M=16, so a group holding 1 token still pays a 16-row tile.
E, K, H, I, B = 256, 8, 7168, 2048, 2  # DeepSeek-V3-shaped MoE layer, bf16


def expert_bytes() -> int:
    """gate_up + down weights of ONE expert."""
    return (H * 2 * I + I * H) * B


def active_experts(pairs: int) -> float:
    """Expected distinct experts touched by `pairs` uniformly-routed token-expert pairs."""
    return E * (1.0 - (1.0 - 1.0 / E) ** pairs)


def flop_rows(rows: float) -> float:
    return rows * 2.0 * (H * 2 * I + I * H)


def t_batched(tokens: int) -> float:
    """One weight slice gathered per token-expert PAIR: weight traffic scales with S."""
    s = tokens * K
    return max(flop_rows(s) / PEAK, s * expert_bytes() / BW)


def t_grouped(tokens: int) -> float:
    """Each active expert tile read ONCE, but every group pays a full BM-row tile."""
    s = tokens * K
    ea = active_experts(s)
    return max(flop_rows(max(s, ea * BM)) / PEAK, ea * expert_bytes() / BW)


rows = []
for tokens in (1, 4, 16, 64, 256, 1024, 4096):
    tb, tg = t_batched(tokens), t_grouped(tokens)
    rows.append((tokens, tb * 1e6, tg * 1e6, (tb - tg) * 1e6))  # microseconds

# The GEMM alone does not separate the two backends at decode: at 1 token nearly every pair
# lands on a distinct expert, so both read ~the same weights and both are bandwidth-bound.
budget_decode = rows[0][3]   # us of grouped-side overhead that batched_mm can absorb
budget_prefill = rows[-1][3]
assert abs(budget_decode) < 50, "at 1 token the two rooflines are within tens of us"
assert budget_prefill > 1000, "at 4096 tokens grouped_mm's advantage is milliseconds"

print("MODEL ONLY (H100 roofline, no GPU). One MoE layer, E=%d K=%d H=%d I=%d bf16." % (E, K, H, I))
print("expert weights: %.1f MB each; all %d experts: %.1f GB" % (expert_bytes() / 1e6, E, E * expert_bytes() / 1e9))
print()
print("tokens  distinct experts  batched(us)  grouped(us)  batched - grouped (us)")
for tokens, tb, tg, d in rows:
    print("%6d  %16.1f  %11.1f  %11.1f  %+22.1f" % (tokens, active_experts(tokens * K), tb, tg, d))
print()
print("Read the last column as an overhead budget:")
print("  at 1 token, grouped_mm's roofline lead is only %.1f us, so any sort/histogram/" % budget_decode)
print("  gather/scatter overhead above that makes batched_mm the faster decode kernel.")
print("  at 4096 tokens the model leaves a %.0f us overhead budget before batched_mm breaks even." % budget_prefill)

Executed output:

MODEL ONLY (H100 roofline, no GPU). One MoE layer, E=256 K=8 H=7168 I=2048 bf16.
expert weights: 88.1 MB each; all 256 experts: 22.5 GB

tokens  distinct experts  batched(us)  grouped(us)  batched - grouped (us)
     1               7.9        210.3        207.5                    +2.9
     4              30.1        841.4        792.4                   +49.0
    16             100.9       3365.5       2652.4                  +713.1
    64             221.5      13461.8       5823.5                 +7638.3
   256             255.9      53847.4       6728.7                +47118.7
  1024             256.0     215389.4       6730.9               +208658.5
  4096             256.0     861557.6       6730.9               +854826.7

Read the last column as an overhead budget:
  at 1 token, grouped_mm's roofline lead is only 2.9 us, so any sort/histogram/
  gather/scatter overhead above that makes batched_mm the faster decode kernel.
  at 4096 tokens the model leaves a 854827 us overhead budget before batched_mm breaks even.

At one token per step, eight pairs land on nearly eight distinct experts, so the model assigns both backends almost the same logical expert-weight traffic. The computed difference is only 2.9 microseconds; grouped execution loses if its sort, histogram, gather, and inverse scatter exceed that budget. At 4096 tokens, the model gives grouped execution a much larger break-even budget because pair-wise materialization duplicates expert operands. These are conditional roofline results, not measured launch or HBM-transaction counts.

How to maintain it

The torch version gate, and the silent fallback

_can_use_grouped_mm() is the function to read before filing a performance bug.1 Its gates, as of transformers main on 2026-07-13:

Condition Consequence
torch.nn.functional.grouped_mm present (torch >= 2.10) CUDA SM80+ dispatches
torch._grouped_mm present, torch >= 2.9 CUDA SM80+ dispatches
torch._grouped_mm present, torch < 2.9 SM90+ only; Ampere silently falls back
CPU, torch <= 2.8 no CPU kernel at all; always falls back
CPU, torch <= 2.10, misaligned pointers 16-byte alignment required; memmapped safetensors are not guaranteed to satisfy it
compiling with dynamo and weights not bf16 falls back

Every one of those failures lands in _grouped_mm_fallback, a python for loop of torch.mm calls registered as a torch.library custom op with its own autograd. It is correct and it is slow, and it does not warn. A grouped-GEMM MoE that is mysteriously at eager speed is almost always here. The custom-op wrapper exists because torch.compiler.disable does not work with fullgraph=True, so inductor cannot optimise the matmuls inside that loop.1

Two mechanisms, one blog paragraph

experts_implementation= and use_kernels=True are different systems, and the Transformers blog runs them together.10

  • experts_implementation=... selects a function from ExpertsInterface / FP8ExpertsInterface: the six backends above.
  • use_kernels=True activates the kernels hub mapping in integrations/hub_kernels.py, which swaps whole layer classes for compiled kernels downloaded from the Hub. This is where MegaBlocks lives: MegaBlocksMoeMLP maps to kernels-community/megablocks for cuda (training and inference), rocm, xpu, and cpu (as CPUMegaBlocksMoeMLP).7

Two things in that mapping deserve a note in a runbook. Llama4TextMoe mapped to kernels-community/moe is commented out, with the comment "# NOTE: No longer maintained". And sonic-moe is pinned to a non-default branch, revision="ep-support".7 Which concrete model classes bind to the MegaBlocks kernel at runtime is a mapping entry, not something traced end-to-end here; verify on your own model before claiming block-sparse execution.

The vLLM naming table

vLLM's MoE kernels live in vllm/model_executor/layers/fused_moe/ and use a different vocabulary for the same tensors. Keep this table next to any weight-conversion or profiling script:

Concept HuggingFace Transformers vLLM
packed gate+up expert weights mlp.experts.gate_up_proj w13_weight
packed down-projection weights mlp.experts.down_proj w2_weight
kernel selection experts_implementation="grouped_mm" Triton fused_moe + tuned configs

vLLM ships 326 hand-tuned Triton config JSONs under fused_moe/configs/ (count on main, 2026-07-13), named E=<num_experts>,N=<intermediate>,device_name=<GPU>[,dtype=<...>].json. Each file is keyed by M, the token count, and each entry pins BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, GROUP_SIZE_M, num_warps and num_stages. At M=1 a typical entry selects BLOCK_SIZE_M: 16, which is the tile-quantization fact the roofline model above uses: a one-token group still pays a 16-row tile. A model shape with no matching config file falls back to a default and loses the tuning, which is a quiet throughput regression on any new expert count or intermediate size.

Maintenance habits

  • Re-diff integrations/moe.py on every transformers bump. Six backends existed on 2026-07-13; three existed when the announcement blog was written. This file changes faster than its documentation.
  • Bisect numerics against eager at token level, using dtype-aware tolerances rather than a perplexity score. The executed block above shows why: an offset bug corrupted 6 of 96 rows.
  • Assert the packed layout after any conversion. The reverse_op on every WeightConverter operation makes a round-trip test cheap: convert, invert, compare to the original checkpoint tensors.
  • Pin the kernels-hub revisions. A hub kernel pinned to a branch (ep-support) rather than a commit can move under you.

How to run it in production

  • Expect the decode profile to show batched_mm. If you configured grouped_mm and Nsight shows bmm, the backend switched itself. That is the intended behaviour, not a bug, and it reverses when generation ends.
  • enable_expert_parallel=True in transformers is not DeepEP. This distinction changes the capacity model, and it is covered in detail on expert parallelism for inference: the transformers ep_plan uses RouterParallel plus GroupedGemmParallel, which replicate the router on every rank, mask non-local experts to a sentinel index, and all-reduce the partial outputs. There is no all-to-all dispatch or combine. Every rank still pushes every token through its local experts.6 It is simpler and NVLink-friendly, but it does not reduce per-rank token count, and a capacity plan that assumed DeepEP semantics will be wrong.
  • Divisibility is a hard failure, in two different places. GroupedGemmParallel.shard_tensor raises Global number of experts must be divisible by number of devices at shard time, so a bad world size fails while the weights are being loaded. RouterParallel enforces the same arithmetic with a different message (The number of experts must be divisible by number of ep_size) from _prepare_output_fn, so that one fires at forward time.6 Both are loud, but only the first is a load-time failure. A 6-GPU box cannot serve a 256-expert model under this plan.
  • Blackwell scale format is a first-forward failure, and that is good. On SM100+, the deepgemm and deepgemm_megamoe kernels require power-of-two UE8M0 expert scales; a checkpoint quantized with plain float32 scales raises a ValueError on the first forward "instead of silently corrupting the output".2 Prefer loud failures at this layer; most of the others are silent.

Reading the published benchmarks

Two numbers circulate with this feature and neither means what it appears to mean. Both were checked against the primary source on 2026-07-13.

  • The loading benchmark does not use an MoE model. The Transformers blog's loading section says "The focus is on loading speed of large MoE models, which is often a bottleneck in training and inference", then benchmarks Qwen/Qwen1.5-110B-Chat (66.24s on v4.57.6 to 20.71s on v5, and 10.1s with tensor parallelism).10 That checkpoint's config.json declares "architectures": ["Qwen2ForCausalLM"], "model_type": "qwen2", and has no num_experts field at all: it is a dense 110B model.11 Those figures measure the generic v5 async loading pipeline and exercise the expert-packing path not at all. Do not use them to size an MoE load.
  • The "~12x faster MoE training" figure is not reproducible from Unsloth's own gpt-oss BF16 table. The blog lists it as one of four standalone bullets (~12x faster MoE training, >35% VRAM reduction, ~6x longer context, 12-30x overall speedup compared to v4), and links to Unsloth's guide.10 That guide's best v5-relative row is 7.3x: 712.33 ms against Transformers v5's 5226.86 ms at 8192 context, with VRAM at 47.43 against 73.80 GB (35.73% saved), on a B200, at LoRA rank 64 (not a full fine-tune). The 12-30x is a v4-relative composite, and 12 is its bottom, not its top.12 Report the row with its hardware, its baseline version and its LoRA rank, or do not report it.

The blog's own back-of-the-envelope figure is worth keeping, because it is honest about being one: gpt-oss-20b has roughly 3.6B active parameters, so on an M3 Ultra with about 800 GB/s of memory bandwidth, 800 / (3.6 * 2) predicts about 111 tok/s in bf16, against roughly 115 tok/s observed.10 That is a memory-bandwidth sanity check on active parameters, not a Transformers throughput benchmark, and the blog concedes it would be faster with native mxfp4 kernels. It is the same arithmetic that makes expert offload to CPU work.

Failure modes

  • Silent grouped-GEMM fallback. Wrong torch version, wrong compute capability, a CPU tensor on torch <= 2.8, or a misaligned memmapped safetensor pointer, and _grouped_mm_fallback runs a python loop of torch.mm. Correct output, eager-class throughput, no warning. Check _can_use_grouped_mm() conditions against your torch build before profiling anything else.
  • Layout flag mismatch. An experts class decorated with the wrong is_concatenated or is_transposed produces right-shaped, finite, entirely wrong activations. Quantified above: 24 of 24 elements wrong, no exception.
  • Off-by-one group offsets. Corrupts a subset of token rows (6 of 96 in the executed example) and passes a smoke test. Token-level comparison against eager with dtype-aware tolerances finds it.
  • Missing permutation inversion. Corrupts every token row while leaving the tensor finite and correctly shaped, because the output is a valid permutation of the right values.
  • Profiler shows the wrong kernel at decode. The automatic grouped_mm to batched_mm switch is scoped to generation and restored afterwards. Profiling a generate() call and a raw forward() call will show different kernels for the same model.
  • Assuming enable_expert_parallel=True means DeepEP. Masked-router plus all-reduce, not all-to-all dispatch and combine. Per-rank token count does not fall.
  • Expert count not divisible by world size. Hard ValueError at shard time. This one is loud, which puts it in the minority on this page.
  • Hub-kernel drift. Llama4TextMoe is commented out of the kernels mapping as unmaintained, and sonic-moe is pinned to a branch. A silently-not-applied hub kernel looks exactly like a slow model.

References

  • HuggingFace Transformers, Experts backends (the six backends, decode-stage switching, experts_implementation / set_experts_implementation / get_experts_implementation, torch.compile compatibility, Blackwell UE8M0 requirement) — https://huggingface.co/docs/transformers/main/en/experts_interface
  • HuggingFace Transformers, src/transformers/integrations/moe.py (ExpertsInterface, ALL_EXPERTS_FUNCTIONS, use_experts_implementation, batched_mm_experts_forward, grouped_mm_experts_forward, _can_use_grouped_mm, _grouped_mm_fallback) — https://github.com/huggingface/transformers/blob/main/src/transformers/integrations/moe.py
  • HuggingFace Transformers, Expert parallelism (DistributedConfig(enable_expert_parallel=True), ep_plan, GroupedGemmParallel, all-reduce combine) — https://huggingface.co/docs/transformers/main/en/expert_parallelism
  • HuggingFace Transformers PR #41580, "Refactor weight loading" (WeightConverter, MergeModulelist, Concatenate, SplitModulelist; merged 2025-11-13) — https://github.com/huggingface/transformers/pull/41580
  • HuggingFace Transformers PR #42697, "batched and grouped experts implementations" (the Experts Backend interface; merged 2026-01-05) — https://github.com/huggingface/transformers/pull/42697
  • HuggingFace, Mixture of Experts in Transformers (the announcement blog; three backends, the loading benchmark, the M3 Ultra estimate) — https://huggingface.co/blog/moe-transformers
  • MegaBlocks: Efficient Sparse Training with Mixture-of-Experts, arXiv:2211.15841 (block-sparse MoE kernels; "never drops tokens"; up to 40% over Tutel, 2.4x over Megatron-LM) — https://arxiv.org/abs/2211.15841
  • databricks/megablocks — https://github.com/databricks/megablocks · kernels-community/megablocks (the Hub kernel) — https://huggingface.co/kernels-community/megablocks
  • DeepSeek DeepGEMM (Mega MoE: fuses EP dispatch, FP8xFP4 linear 1, SwiGLU, linear 2, and EP combine into one kernel, overlapping NVLink with tensor-core compute) — https://github.com/deepseek-ai/DeepGEMM
  • kernels-community/sonic-moe (CuteDSL grouped-GEMM MoE) — https://huggingface.co/kernels-community/sonic-moe · Dao-AILab/quack — https://github.com/Dao-AILab/quack
  • vLLM fused MoE kernels and tuned Triton configs (w13_weight / w2_weight, E=...,N=...,device_name=...json) — https://github.com/vllm-project/vllm/tree/main/vllm/model_executor/layers/fused_moe
  • Unsloth, Faster MoE training (the real gpt-oss BF16 table: 712.33 ms vs 5226.86 ms at 8192 ctx on B200, LoRA rank 64) — https://unsloth.ai/docs/basics/faster-moe
  • GShard, arXiv:2006.16668 — https://arxiv.org/abs/2006.16668 · Switch Transformer, arXiv:2101.03961 — https://arxiv.org/abs/2101.03961 · DeepSeekMoE, arXiv:2401.06066 — https://arxiv.org/abs/2401.06066

Related: MoE sparse scaling · MoE routing and load balancing · Expert parallelism for inference · MoE semantic parallelism and co-scheduling · NeMo AutoModel · Model weight loading in inference engines · Tensor Cores and mixed precision · Comms/compute overlap

  1. src/transformers/integrations/moe.py on main, read 2026-07-13 (588 lines): ExpertsInterface._global_mapping = {"deepgemm", "batched_mm", "grouped_mm", "sonicmoe"}; eager is special-cased in get_interface; use_experts_implementation is a decorator (there is no MoeExperts class) carrying is_concatenated, is_transposed, has_bias, has_gate; _grouped_mm dispatches to torch.nn.functional.grouped_mm or torch._grouped_mm and otherwise to torch.ops.transformers.grouped_mm_fallback, a torch.library.custom_op with register_autograd. modeling_utils.get_correct_experts_implementation defaults to grouped_mm, validates against ["eager"] + ALL_EXPERTS_FUNCTIONS | ALL_FP8_EXPERTS_FUNCTIONS, and downgrades to eager if _grouped_mm_can_dispatch() raises only when grouped_mm was not explicitly requested: the handler is except (ValueError, ImportError) as e: if requested_experts == "grouped_mm": raise e, so an explicit experts_implementation="grouped_mm" re-raises rather than downgrading. deepgemm_megamoe is registered only in FP8ExpertsInterface._global_mapping in integrations/finegrained_fp8.py. ↩↩↩↩↩

  2. HuggingFace Transformers docs, "Experts backends" (docs/source/en/experts_interface.md on main, 2026-07-13): the six-backend table; "On GPU, a model loaded with experts_implementation="grouped_mm" automatically switches to "batched_mm" for the decode stage of generation... On CPU, grouped_mm stays active throughout generation"; SonicMoE "State-of-the-art throughput on Hopper (SM90+) for bfloat16 experts... especially for training"; the eager backend "uses a data-dependent operation... not compatible with full graph compilation"; Blackwell UE8M0 scales "raises a ValueError on the first forward instead of silently corrupting the output". ↩↩↩↩↩

  3. src/transformers/generation/utils.py on main, 2026-07-13: _optimize_model_for_decode is a context manager that snapshots get_experts_implementation(), switches every grouped_mm entry to batched_mm when self.device.type != "cpu", and restores the original mapping in its finally block. ↩

  4. src/transformers/models/gpt_oss/modeling_gpt_oss.py on main, 2026-07-13: @use_experts_implementation(is_concatenated=False, is_transposed=True, has_bias=True) on GptOssExperts; gate_up_proj is nn.Parameter(torch.empty(num_experts, hidden_size, 2 * intermediate_size)), down_proj is (num_experts, intermediate_size, hidden_size); _apply_gate splits with gate_up[..., ::2], gate_up[..., 1::2] (interleaved). ↩

  5. src/transformers/conversion_mapping.py and core_model_loading.py on main, 2026-07-13: the Mixtral entry maps mlp.experts.*.w1.weight and mlp.experts.*.w3.weight to mlp.experts.gate_up_proj with operations=[MergeModulelist(dim=0), Concatenate(dim=1)], and mlp.experts.*.w2.weight to mlp.experts.down_proj with [MergeModulelist(dim=0)]. Every ConversionOps subclass declares a reverse_op; MergeModulelist's is SplitModulelist. ↩

  6. src/transformers/integrations/tensor_parallel.py on main, 2026-07-13: GroupedGemmParallel.shard_tensor raises "Global number of experts must be divisible by number of devices" and slices param[start:end] on dim 0. RouterParallel.shard_tensor returns the full parameter (the router is replicated); its divisibility check lives instead in _prepare_output_fn, which raises a different message ("The number of experts must be divisible by number of ep_size") at forward time. _prepare_output_fn also masks non-local scores to 0.0 and non-local indices to a sentinel, and its docstring states "After the expert forward, an all_reduce sums partial outputs across EP ranks to produce the full result." No all-to-all appears in this path. ↩↩

  7. src/transformers/integrations/hub_kernels.py on main, 2026-07-13: MegaBlocksMoeMLP maps to kernels-community/megablocks for cuda (Mode.TRAINING and Mode.INFERENCE), rocm, xpu, and cpu (layer CPUMegaBlocksMoeMLP); the Llama4TextMoe to kernels-community/moe entry is commented out under "# NOTE: No longer maintained"; "sonic-moe": {"repo_id": "kernels-community/sonic-moe", "revision": "ep-support"}. ↩↩

  8. Gale et al., MegaBlocks: Efficient Sparse Training with Mixture-of-Experts, arXiv:2211.15841 (abstract): existing formulations "force a tradeoff between model quality and hardware efficiency, as users must choose between dropping tokens from the computation or wasting computation and memory on padding"; the block-sparse reformulation "never drops tokens and maps efficiently to modern hardware", with "end-to-end training speedups of up to 40% over MoEs trained with the state-of-the-art Tutel library and 2.4x over DNNs trained with the highly-optimized Megatron-LM framework". ↩

  9. deepseek-ai/DeepGEMM README (2026-07-13): "Mega MoE fuses and overlaps EP dispatch, linear 1 (FP8xFP4), SwiGLU, linear 2 (FP8xFP4), and EP combine into a single mega-kernel, overlapping NVLink communication and tensor core computation. It requires multi-process launch with symmetric memory." Released 2026.04.16. ↩

  10. HuggingFace blog, "Mixture of Experts in Transformers" (read 2026-07-13): "Three backends are currently provided"; "The focus is on loading speed of large MoE models, which is often a bottleneck in training and inference", benchmarking Qwen/Qwen1.5-110B-Chat at 66.24s (v4.57.6), 20.71s (v5) and 10.1s (v5 + TP); the "~12x faster MoE training" claim linking to Unsloth; and the M3 Ultra estimate "800 / (3.6 * 2)... about 111 tokens per second. The actual performance number we get is ~115 tok/s", with the caveat that "speed would be even faster if we used kernels for the native mxfp4 quantization the model uses". As of 2026-07-13 the post is roughly five months behind its own subject: three backends rather than six, no experts_implementation kwarg, no base_model_ep_plan, no decode-stage switching. ↩↩↩↩↩

  11. Qwen/Qwen1.5-110B-Chat config.json (fetched 2026-07-13): "architectures": ["Qwen2ForCausalLM"], "model_type": "qwen2", hidden_size 8192, num_hidden_layers 80, and no num_experts or num_experts_per_tok key. It is a dense model, so the blog's loading benchmark does not exercise the MoE expert-packing path. ↩

  12. Unsloth, "Faster MoE" docs (read 2026-07-13). gpt-oss (20B) BF16 on NVIDIA B200, LoRA rank 64 with all LoRA modules on the MoE layers: at 8192 context, Unsloth 712.33 ms vs Transformers v5 5226.86 ms (7.3x) and 47.43 GB vs 73.80 GB VRAM (35.73% saved); at 1024 context the same table shows 1.4x. The "12-30x" figure is quoted against Transformers v4. ↩