Dense scaling hits a hard wall: every parameter you add is a parameter you pay for on every token, in both VRAM and bandwidth. Mixture-of-Experts breaks that coupling by making total parameters and active parameters different numbers. The trade is that a token’s path is no longer fixed — a router decides — and routing is where MoE both wins and dies. Mixtral-8x7B (December 2023) made the pattern mainstream: 47B total parameters, 12.9B active per token, and roughly 6x faster decoding than a dense Llama-2-70B while staying competitive with it on quality. The math behind the router is deceptively small and disproportionately important.

How a sparse MoE layer works

You replace the dense feed-forward block with N parallel experts (each a standard MLP) plus a router: a single linear layer that maps the token embedding to N logits. Softmax those logits, pick the top-K experts (Mixtral uses K=2), renormalize the two probabilities, and compute the weighted sum of the experts’ outputs. That weighted sum is added through the residual stream like a normal FFN. The router is a rounding error on FLOPs — one d_model × N matmul — but its decisions control everything downstream, so the auxiliary load-balancing loss (Switch-Transformer style, coefficients in the 0.001–0.01 range) exists to stop the router from routing everything to one favorite expert and underusing the rest of the hardware you paid for.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
import torch
import torch.nn as nn

class MoE(nn.Module):
    def __init__(self, d_model, d_ff, n_experts, top_k=2, capacity_factor=1.25):
        super().__init__()
        self.router = nn.Linear(d_model, n_experts)
        self.experts = nn.ModuleList([
            nn.Sequential(nn.Linear(d_model, d_ff), nn.GELU(), nn.Linear(d_ff, d_model))
            for _ in range(n_experts)
        ])
        self.top_k, self.capacity = top_k, int(capacity_factor * n_experts)

    def forward(self, x, aux_loss_weight=0.01):
        B, L, D = x.shape
        flat = x.view(-1, D)
        N = flat.size(0)
        probs = torch.softmax(self.router(flat), dim=-1)
        top_probs, top_idx = probs.topk(self.top_k, dim=-1)
        top_probs = top_probs / top_probs.sum(dim=-1, keepdim=True)

        out = torch.zeros_like(flat)
        for e in range(self.n_experts):
            mask = (top_idx == e)                     # [N, top_k] boolean
            if not mask.any():
                continue
            rows, slots = mask.nonzero(as_tuple=True)
            rows, slots = rows[:self.capacity], slots[:self.capacity]  # capacity check
            w = top_probs[rows, slots]
            out[rows] += w[:, None] * self.experts[e](flat[rows])

        # aux load-balancing loss: penalize concentration of router mass
        frac = mask.float().mean(0) if mask.any() else torch.zeros(self.n_experts, device=x.device)
        aux_loss = aux_loss_weight * self.n_experts * (frac * probs.mean(0)).sum()
        return out.view(B, L, D) + x, aux_loss

The two lines that bite in production are the capacity check (rows[:capacity]) and the aux loss. The capacity factor (1.25 above) is a buffer on how many tokens each expert processes per batch: tokens past capacity are dropped to the residual stream. Raise it and you eliminate drops but waste the batch; lower it and you trade throughput for routing fidelity. We tune it per workload — long-document generation tolerates 1.0–1.1, bursty chat needs 1.25.

What quality actually costs

The reason MoE feels free is that the compute per token is dominated by the active experts, not the stored ones. On a single A100-80G, Mixtral decodes at roughly 12–13 tokens/s for a 47B-parameter model — 3x what the raw parameter count suggests, because only the two selected experts’ weights cross memory bandwidth per token. But “free” has a floor: even inactive experts occupy VRAM. At fp16, Mixtral’s weights alone are ~94 GB, so you cannot host it on one 80 GB card. Serving engines handle this with expert parallelism: shard experts across GPUs so each GPU holds a subset, and route tokens across the interconnect. vLLM’s Mixtral path, for instance, wants tensor-parallel width 8 — one expert per GPU — with the caveat that you now need 8 GPUs (or 4 with double the per-GPU memory) and a fast interconnect, because every token’s experts may live on different devices.

Production lessons

• Audit routing distributions, don’t trust the aux loss. We have caught load imbalance that the auxiliary term tolerated but that showed up as one GPU pegging thermals while its neighbors idled. Export per-expert token counts per batch; a healthy router stays within ~20% of uniform.
• The router is a hidden eval dimension. It forms semantic preferences (which experts absorb code, which absorb prose), and those preferences shift with fine-tuning. Retrain or re-route after adaptation, or you serve stale partitions.
• Capacity drops are silent quality loss. Drops flow through the residual stream and read as model confusion, not as a routing bug. Log drop rates per batch and alert above ~2%.
• Don’t hand-roll expert parallelism for MoE. NCCL all-to-all under ragged per-expert token counts is where latency hides; use engines with a proven grouped-GEMM path (vLLM, TGI) and profile before adding your own sharding.

MoE is the pragmatic way to keep scaling past the dense wall, but the router converts your model’s quality into a systems problem you can actually measure — which, frankly, is the most pleasant kind of problem an ML engineer can have.