ACE Journal

Mixture-of-Experts Routing at Inference Time

Who this is for: ML engineers and inference-platform owners deploying or operating an open-weight MoE model (the Mixtral, DeepSeek, and Qwen MoE families, or your own). The assumption is that you understand transformer feed-forward layers and have served dense LLMs on vLLM, TensorRT-LLM, or DeepSpeed, but you have not yet had to reason about expert routing under production load. If you have ever seen an MoE deployment quietly lose quality on long prompts and had no metric pointing at why, this is the layer that explains it.

The regression that hides in the residual stream

A team swaps a dense 13B model for an open-weight MoE with eight experts and top-2 routing. On the benchmark suite it looks great: same latency class, higher quality, a fraction of the active parameters. They ship it. Two weeks later a support thread says the model has gotten worse at long documents. Not broken, just worse, vaguer, more likely to lose a thread of reasoning past a few thousand tokens. The eval suite, all short prompts, still passes. Nothing errors. The GPUs are not saturated. There is no obvious place to look.

What is happening is that under production batching, certain experts are overloaded, their capacity is exceeded, and the overflow tokens are dropped. A dropped token skips its expert and passes through on its residual connection alone, a different and worse computation than the model was trained to do. On short prompts the effect disappears into the noise of the eval. On long prompts it compounds layer after layer until the output visibly drifts. The router, the smallest and cheapest module in the model, is what broke, and nobody was watching it because at training time it had been kept honest by losses that no longer exist.

That is the through-line of this article. The MoE router is deceptively simple to describe and carries outsized consequences at inference time. We will walk the mechanics of sparse routing, build a worked top-k router, dig into the hard part (token dropping as a silent correctness issue), look at how it fits a real serving stack, position the routing variants against each other, and end with the failure modes worth knowing before you deploy.

How sparse routing works

A dense transformer runs every token through the same feed-forward network (FFN). An MoE layer replaces that single FFN with N expert FFNs and a router. For each token, the router produces a score for every expert, keeps the top k (almost always 2, sometimes 1), runs only those k experts, and combines their outputs weighted by the router’s scores. The other N - k experts do no work for that token. That is the entire trick: parameters scale with N, but compute per token scales with k.

Concretely, for a token hidden state h, the router is a single learned matrix W_g of shape [hidden, N]. It projects h to N logits, takes a softmax, selects the top k experts, renormalizes their gate values, and forms the output as the gate-weighted sum of those experts applied to h. Figure 1 shows the path one token takes through a top-2 layer.

flowchart LR
    H[token hidden state h] --> G[router: h x W_g -> N logits]
    G --> S[softmax + top-k select]
    S -->|gate g1| E1[expert 1]
    S -->|gate g3| E3[expert 3]
    E1 --> C[weighted sum: g1*E1 + g3*E3]
    E3 --> C
    C --> O[layer output for h]
    G -.->|experts 2,4..N not run| X[skipped]

Figure 1. The path of a single token through a top-2 MoE layer. Note that only the two selected experts execute; the rest are skipped entirely, which is where the compute savings come from and where the routing decision becomes load-bearing.

The decisive detail for serving is what trains this router and what does not. During training, auxiliary losses push the router toward spreading tokens across experts. The original sparsely-gated MoE work [1] introduced a load-balancing penalty, and the Switch Transformer [2] simplified routing to top-1 with a load-balancing loss that has become standard. These losses are part of the optimization objective. At inference time they are gone. The router runs purely from its learned W_g, and on inputs that differ from the training distribution there is nothing actively pulling it back toward balance. The balance you observed during training was partly an artifact of an objective you are no longer optimizing.

Worked example: a top-k router and the imbalance it produces

Here is a minimal top-k router in PyTorch. It is the gating logic at the heart of every MoE layer, stripped of the expert FFNs so you can see only the routing decision.

import torch
import torch.nn.functional as F

class TopKRouter(torch.nn.Module):
    def __init__(self, hidden: int, num_experts: int, k: int = 2):
        super().__init__()
        self.gate = torch.nn.Linear(hidden, num_experts, bias=False)
        self.k = k

    def forward(self, h: torch.Tensor):
        # h: [tokens, hidden]
        logits = self.gate(h)                       # [tokens, num_experts]
        gates = F.softmax(logits, dim=-1)
        topk_gates, topk_idx = gates.topk(self.k, dim=-1)
        topk_gates = topk_gates / topk_gates.sum(-1, keepdim=True)  # renormalize
        return topk_idx, topk_gates                 # who runs, and with what weight

Given a batch of token hidden states, topk_idx tells the engine which experts to run for each token and topk_gates is how to weight their outputs. To see imbalance, count how many tokens each expert is asked to process across a batch.

def expert_load(topk_idx: torch.Tensor, num_experts: int) -> torch.Tensor:
    # topk_idx: [tokens, k] -> counts per expert
    return torch.bincount(topk_idx.flatten(), minlength=num_experts)

router = TopKRouter(hidden=512, num_experts=8, k=2)
batch = torch.randn(4096, 512)            # 4096 tokens
idx, _ = router(batch)
load = expert_load(idx, num_experts=8)
print(load.tolist())                      # rarely uniform; some experts get far more

With a freshly initialized router and random inputs the load is roughly uniform, but a trained router fed real, correlated traffic is a different story. Tokens are not independent draws. A batch dominated by code, or one language, or a single document’s vocabulary tends to route to the same handful of experts, because those experts specialized in exactly that kind of token during training. The router is doing its job, sending similar tokens to the expert that handles them best, and the side effect is that on skewed traffic a few experts get most of the tokens. Per-expert capacity in a forward pass is fixed, so when load concentrates the overloaded experts hit their ceiling.

We can make the capacity ceiling explicit. Each expert can take at most capacity tokens in a pass; tokens beyond that are overflow.

def dropped_tokens(load: torch.Tensor, capacity: int) -> int:
    overflow = (load - capacity).clamp(min=0)   # tokens past capacity, per expert
    return int(overflow.sum())

# capacity is usually set as a factor of the average load
avg = load.float().mean()
capacity = int(avg * 1.25)                       # capacity factor 1.25
print("dropped:", dropped_tokens(load, capacity))

The capacity here is the same knob production engines expose. vLLM’s fused MoE path and DeepSpeed-MoE [3] both let you set a capacity factor, the multiple of average per-expert load each expert may accept. Set it to 1.0 and you assume perfect balance, which real traffic rarely delivers, so you drop tokens off the overloaded experts. Set it high and every expert reserves memory for a worst case that mostly does not happen. The right value is a property of your traffic, not your training config, which is the first sign that MoE serving is an operational problem and not just a model-loading one.

The hard part: token dropping is a silent correctness bug

Everything above is the setup for the one insight a practitioner only learns by getting burned. When an expert is over capacity, the engine has to handle the overflow tokens, and the common default is to drop them. A dropped token is not skipped from the layer entirely; the residual connection still carries it forward. What it loses is the expert contribution. Where a non-dropped token gets h + (g1*E1(h) + g3*E3(h)), a dropped token gets h plus nothing, or in a top-2 case where only one expert overflowed, h plus the single surviving expert. The token still has a value, it is just the wrong value, computed by a model that is now smaller than the one you trained. Figure 2 traces both branches.

flowchart TD
    T[token routed to expert E] --> Q{E at capacity?}
    Q -->|no| P[E processes token] --> A[output: h + gate * E h]
    Q -->|yes| D{overflow policy}
    D -->|drop| R[output: h only, expert skipped]
    D -->|reroute| N[send to 2nd-choice expert] --> A
    R --> W[silent quality loss, no error raised]

Figure 2. What happens to a token routed to an expert that is already full. The drop branch is the dangerous one. It produces a valid-looking output with no exception and no log entry, which is exactly why the resulting quality loss is so hard to trace.

Three properties make this nasty. First, it is silent: no exception, no NaN, tensor shapes unchanged. Second, it is distribution-dependent. The same model under the same capacity factor drops zero tokens on one traffic mix and many on another, so it passes your evals and fails in production. Third, it compounds. The error from one dropped token at one layer is small, but a long sequence has many tokens and the model has many MoE layers, so the per-layer perturbations accumulate into a measurable shift in the output distribution. That is why the symptom shows up first on long inputs and complex reasoning, not on short prompts.

The practical response is the discipline you would apply to any silent corruption: measure first, then mitigate. The single most actionable step is to count dropped tokens per forward pass and export it as a metric. You cannot fix what you cannot see, and most teams running MoE in production have never looked at this number.

# instrumentation probe: log per-expert load and drops each forward pass
def routing_stats(topk_idx, num_experts, capacity_factor=1.25):
    load = expert_load(topk_idx, num_experts)
    cap = int(load.float().mean() * capacity_factor)
    drops = dropped_tokens(load, cap)
    return {
        "max_expert_load": int(load.max()),
        "min_expert_load": int(load.min()),
        "imbalance_ratio": float(load.max() / load.float().mean()),
        "dropped_tokens": drops,
        "drop_fraction": drops / int(load.sum()),
    }

With those numbers in hand the mitigations order themselves. If drop_fraction is consistently nonzero, raise the capacity factor until it returns to near zero on representative traffic, accepting the memory cost as the price of correctness. If that is too expensive, change the overflow policy from drop to reroute, sending overflow tokens to their second-choice expert instead of skipping the expert tier. Rerouting trades a perfect assignment for a worse-but-real one, which beats no expert at all. The imbalance_ratio tells you whether the problem is fundamental skew (a few experts always hot, pointing at the model and your traffic) or transient bursts (which a slightly higher capacity factor absorbs). The metric is not a dashboard you admire; it is the input to the decision.

Integration: where routing lives in a serving stack

In a real deployment the router is not something you call directly. It is buried inside the MoE layers of a model loaded by an inference engine, and the levers you actually touch are engine configuration and monitoring. The flow from a request to an answer, and the points where routing behavior matters, looks like Figure 3.

flowchart LR
    Req[incoming requests] --> B[batch scheduler]
    B --> F[forward pass]
    subgraph layer [MoE layer x many]
        Rt[router top-k] --> Cap[capacity check + overflow policy]
        Cap --> Ex[expert compute]
        Ex --> AA[all-to-all gather/scatter<br/>if expert-parallel]
    end
    F --> layer
    layer --> Out[tokens out]
    Cap -.->|drop count| M[(metrics: load,<br/>drops, imbalance)]
    Rt -.->|per-expert load| M

Figure 3. Where routing sits in an MoE serving path. Watch the two dotted lines into the metrics store: per-expert load and dropped-token counts are the signals that tell you whether the configuration above them is healthy on live traffic.

Two configuration surfaces dominate day to day. The first is the capacity factor, tuned against observed drop fractions rather than copied from the training recipe. The second appears once the model is too big for one device. MoE models with many experts do not fit on a single GPU, so experts are sharded across devices in a scheme called expert parallelism: each device holds a subset of experts, and after the router decides where every token goes, an all-to-all collective ships each token to the device holding its expert and ships the results back. This layers with tensor parallelism inside each expert group, a combination GShard [4] established for large MoE training and inference frameworks inherited.

The all-to-all is the part that surprises people. Its cost depends on the routing distribution, because the router decides how much data crosses the fabric. On NVLink-connected GPUs within a node it is cheap enough to hide. Across PCIe or a slower inter-node fabric it can dominate the forward pass, and skewed routing makes it worse by concentrating traffic on the links into the hot experts’ devices. So before committing to a topology, measure all-to-all throughput under a realistic routing distribution, not a uniform synthetic one, because uniform traffic flatters the interconnect in a way your users will not. The day-to-day workflow is to load the model, run representative traffic, watch the routing metrics and collective timings, and tune capacity factor and placement against what you see rather than against the model card.

The routing variants people actually choose between fall along a single axis: how rigidly the router commits a token to an expert, and how much that commitment can be undone when an expert is full.

Token-choice top-k is the default, the scheme in this article and in the Switch Transformer [2] and Mixtral [5] lineages. Each token independently picks its top experts. It is simple and trains well, and its weakness is the one we have spent the article on: a token’s expert is chosen with no regard for whether that expert has room, so popular experts overflow and tokens drop. Capacity factor and reroute policies are bolt-on fixes for a scheme that ignores load at routing time.

Expert-choice routing [6] inverts the assignment: each expert picks the tokens it will take, up to its capacity. By construction no expert exceeds capacity and load is balanced, removing the dropped-token problem on the expert side. The trade is that a token can now be picked by zero experts, the dual failure, and the scheme is awkward in autoregressive decoding because the choice depends on seeing a batch of tokens at once. It is a genuinely different answer to the same imbalance, not a strict improvement.

Dropless routing, popularized by the MegaBlocks work [7], reframes the expert computation as block-sparse matrix multiplication so variable per-expert loads need no fixed capacity at all. No capacity means no drops, at the cost of a more complex kernel and irregular compute that is harder to schedule. For engines that adopt the kernels it removes the capacity-factor tuning loop entirely; for an off-the-shelf stack that assumes fixed-capacity batched experts it is not a knob you can simply flip on.

The honest summary is that token-choice top-k won on simplicity and trainability, and the entire operational burden of this article, capacity, drops, reroute, monitoring, is the bill for that simplicity. Expert-choice and dropless routing pay it differently, in scheduling complexity or decode-time awkwardness. If you run an open-weight model you almost certainly inherit token-choice and its tuning, so knowing its failure modes is not optional.

Where it breaks

MoE routing fails at inference time in recognizable ways. Knowing them before you deploy is most of avoiding them.

Silent token dropping on skewed traffic. The headline failure. A capacity factor tuned on one traffic mix drops tokens on another, and because the drop is silent the first symptom is a vague quality regression on long inputs. Without a dropped-token metric there is nothing to point at. The fix is to instrument drops, then raise capacity or switch overflow to reroute.

Out-of-distribution routing collapse. With the balancing loss gone at inference, the router can collapse toward a few experts far harder on unfamiliar inputs than it ever did in training. A new language, code style, or adversarial prompt all concentrate load. The imbalance ratio is the early warning; there is no clean fix beyond capacity headroom and, longer term, fine-tuning the router on representative traffic.

All-to-all dominating the forward pass. On expert-parallel deployments across a slow fabric, the routing collective can cost more than the expert compute it serves, and skewed routing makes it worse. Teams that benchmark with uniform synthetic routing miss this and are surprised in production. Benchmark with realistic distributions and prefer high-bandwidth interconnect for the expert-parallel dimension.

Capacity factor set from the training config. The most common misconfiguration is copying the training-time capacity factor into serving as if it were a model property. It is a traffic property. The right value comes from observed drop fractions on your own load and drifts as traffic changes, so it is a setting to monitor, not set once.

Batch composition coupling unrelated requests. Because routing and capacity are computed over a whole batch, one request’s tokens influence whether another’s get dropped. A burst of one kind of traffic can degrade quality for unrelated requests sharing the batch. The mitigation is capacity headroom plus the awareness that MoE quality is a per-batch property, not a per-request one.

Quantization interacting with routing. Quantizing an MoE model can perturb the router’s logits enough to flip top-k selections at the margin, sending some tokens to different experts than the full-precision model would. The effect is usually small, but it means quality under quantization should be checked with routing-aware metrics, not just dense-style perplexity, since the failure is in which experts run, not only how each computes.

Further reading

Start with the original sparsely-gated MoE paper [1] for the routing-and-balancing formulation, then the Switch Transformer [2] for the top-1 simplification and the load-balancing loss that became standard. GShard [4] is the canonical source for expert parallelism and the all-to-all communication pattern that governs large-MoE serving. For the inference-engine view of capacity and expert parallelism, the DeepSpeed-MoE work [3] is the clearest write-up of the serving-time trade-offs. Mixtral [5] is the open-weight model that made these questions practical for most teams. For the routing alternatives, read the expert-choice routing paper [6] and the MegaBlocks dropless approach [7] back to back; they are the two cleanest answers to the imbalance problem this article centers on. The vLLM documentation [8] is the best running reference for how a production engine exposes these knobs.

Written with AI assistance, editorially reviewed.

References

[1] N. Shazeer et al., “Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer.” https://arxiv.org/abs/1701.06538

[2] W. Fedus, B. Zoph, and N. Shazeer, “Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity.” https://arxiv.org/abs/2101.03961

[3] S. Rajbhandari et al., “DeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training to Power Next-Generation AI Scale.” https://arxiv.org/abs/2201.05596

[4] D. Lepikhin et al., “GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding.” https://arxiv.org/abs/2006.16668

[5] A. Q. Jiang et al., “Mixtral of Experts.” https://arxiv.org/abs/2401.04088

[6] Y. Zhou et al., “Mixture-of-Experts with Expert Choice Routing.” https://arxiv.org/abs/2202.09368

[7] T. Gale et al., “MegaBlocks: Efficient Sparse Training with Mixture-of-Experts.” https://arxiv.org/abs/2211.15841

[8] vLLM, “vLLM Documentation.” https://docs.vllm.ai/en/latest/

[9] J. Pan et al., “A Survey on Mixture of Experts in Large Language Models.” https://arxiv.org/abs/2407.06204

[10] DeepSeek-AI, “DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models.” https://arxiv.org/abs/2401.06066