S16Decoding Acceleration·152 lines

Mixture of Experts

An MoE layer activates 8 of 256 experts per token. Compute drops, memory does not, and routing becomes a load-balancing problem.

router
top-k gating
expert dispatch
token grouping
expert parallelism

Motivation

The problem

Bigger models are better, and bigger models are proportionally slower, because every parameter participates in every token. Mixture-of-experts breaks that coupling.

The idea: replace the FFN, which holds about two-thirds of all parameters, with N copies of it, and route each token to only k of them. Mixtral 8×7B stores 47 billion parameters and uses 13 billion per token. DeepSeek-V3 stores 671 billion and uses 37 billion.

For an inference engine this is a mixed blessing: compute scales with active parameters, but memory scales with total parameters. Any expert might be needed by the next token, so all of them must be resident.

Core idea

The solution

A router, one small linear layer, scores the experts for each token. Take the top k, run only those, and combine their outputs weighted by the router's softmax.

an MoE layerpython

def moe_layer(x, router_w, experts, top_k=2):
    # x: [n_tokens, d_model]
    logits = x @ router_w                       # [n_tokens, n_experts]
    idx = np.argpartition(-logits, top_k, axis=-1)[:, :top_k]
    weights = softmax(np.take_along_axis(logits, idx, -1), axis=-1)

    out = np.zeros_like(x)
    for e in range(len(experts)):
        # GROUP: all tokens routed to expert e, processed in one matmul
        rows, slot = np.where(idx == e)
        if len(rows) == 0:
            continue
        y = experts[e](x[rows])                 # one big matmul, not many small ones
        out[rows] += weights[rows, slot][:, None] * y

    return out

The grouping is the kernel

A naive implementation loops over tokens and runs one tiny matmul each: hundreds of launches doing almost nothing, which on a GPU is catastrophic. The real work is sorting tokens by expert so each expert gets one contiguous batch, running the grouped matmuls, then scattering the results back. quant.cpp describes this as a three-phase route → inverse-index → gather/compute/scatter pipeline.

Mechanics

How it works

Diagramrouting, dispatch, and the imbalance problem

Left: what MoE buys and what it costs. Right: the operational risk, where one popular expert sets the pace for every GPU.

Load imbalance is the operational problem

Routers are not naturally uniform. Left alone they collapse onto a few popular experts, and MoE training adds an auxiliary load-balancing loss to stop that. Even a well-trained router is only balanced on average; any particular batch can be skewed.

This matters most under expert parallelism, where different experts live on different GPUs. The step finishes when the busiest GPU finishes, so a 3× imbalance means most of your fleet is idle. The mitigations are a capacity factor (cap tokens per expert, drop or reroute the overflow), and at very large scale, replicating hot experts.

Why MoE is a low-concurrency win

At batch 1 you read k experts and skip the rest; the sparsity is real and decode really is faster. At batch 512, with tokens routing independently, nearly every expert is needed by somebody, so you read the entire model anyway.

So the advantage shrinks as concurrency rises, exactly like speculative decoding. Both trade abundant memory capacity for scarce bandwidth, and both stop paying off once the batch is large enough to saturate the machine.

Explore

Try it

Simulatorexpert routing and load balancedrag router skew, then toggle the balancing loss

experts64top-k2tokens in batch128capacity factor1.25router skew0.55

total params: 362.9G
active / token: 13.4G
sparsity: 27.0×
dropped tokens: 131 (51.2%)

expert load · capacity 5 tokens each

capacity

perfectly balanced

tokens processed
over capacity — dropped or rerouted

Imbalance factor 17.25×. Under expert parallelism each expert lives on a different GPU, so the step takes as long as the busiest expert. A 3× imbalance means the other GPUs idle two-thirds of the time. Turn on the load-balancing loss to see what training-time regularisation buys at inference time.

the MoE bargain

parameters in memory362.9G

parameters read per token13.4G

Compute drops by the sparsity factor. Memory does not drop at all, because every expert has to stay resident against the chance that a token routes to it. An MoE gives you a small model's speed on a large model's memory bill.

batch size changes everything

batch 1

batch 8

batch 64

batch 512

At batch 1 you read 2 experts and the sparsity is real. At batch 512 nearly every expert is touched by someone, so you read the whole model anyway and the MoE advantage evaporates. Sparsity is a low-concurrency benefit.

seed5

Push router skew to 1.0 with the balancing loss off: a handful of experts take most of the traffic and the capacity line starts dropping tokens. It does so silently, since a dropped token skips the FFN and passes through on its residual connection. Then enable the balancing loss and watch the histogram flatten and the drops vanish.

Build it

Implementation

code/s16_moe.py (excerpt)python

def grouped_dispatch(x, expert_ids, n_experts):
    """
    Sort tokens by expert so each expert sees one contiguous slice.
    This is the difference between a usable MoE kernel and a toy one.
    """
    flat = expert_ids.reshape(-1)                  # [n_tokens * top_k]
    order = np.argsort(flat, kind="stable")        # group by expert
    sorted_experts = flat[order]

    counts = np.bincount(sorted_experts, minlength=n_experts)
    offsets = np.concatenate([[0], np.cumsum(counts)])

    return order, offsets     # offsets[e]:offsets[e+1] is expert e's slice

A dropped token is not an error, and that is the danger

When an expert exceeds its capacity, the overflow tokens skip that expert entirely and pass through on the residual connection. No exception, no log line, slightly worse output. If your MoE quality is mysteriously below the reference implementation, instrument the drop rate first. A capacity factor set too low is the usual culprit.

Run it locally

Implements routing, grouped dispatch and combination for a small MoE, asserts the grouped path matches a naive per-token loop exactly, and measures load imbalance, drop rate and effective sparsity across batch sizes and capacity factors.

$ python code/s16_moe.py

Expect: An exact-match assertion between grouped and naive dispatch (256 tiny matmuls become 16 big ones), an imbalance histogram, and a batch-size sweep showing effective sparsity collapsing from 8× to 1× by batch 64.

Only NumPy is required — setup instructions.

Production notes

In production

vLLM / SGLang — fused MoE kernels with grouped GEMM and expert parallelism across GPUs.
DeepSeek-V3 — 256 routed experts plus a shared expert that every token uses, and an auxiliary-loss-free balancing scheme based on per-expert bias terms.
quant.cpp — token-grouped expert dispatch on CPU, with TQ_NO_MLOCK to let the OS page out cold expert weights. On a laptop the interesting question is which experts fit in RAM, not which GPU they live on.
Expert offloading — keep hot experts on GPU and cold ones in host memory, fetching on demand. Only viable when routing is predictable enough to prefetch.

Exercises

1
Measure effective sparsity as a function of batch size: at what batch does the fraction of experts touched per step exceed 90%?
2
Implement capacity-factor dropping and plot output quality against capacity factor. Find where drops start mattering.
3
Simulate expert parallelism across 8 GPUs and measure the step time as max-over-GPUs rather than mean. Quantify the throughput lost to a 2× imbalance.

Continue

What's next

MoE forced the question of what to do when a model does not fit on one GPU. S17 answers it properly: tensor parallelism, pipeline parallelism, and the communication costs that decide which one you want.

Check yourself

Questions

Answer before you look. A wrong answer tells you more than a right one, because the explanation points at the part of the chapter to reread.

Check yourselfQuestion 1 of 6

An MoE activates 2 of 64 experts per token. What scales with the active count and what scales with the total?

Score 0/6