Expert routing in practice — why MoE needs a traffic controller

Yesterday we introduced Mixture-of-Experts (MoE): instead of sending every token through one dense MLP, a router chooses a small number of expert MLPs for each token.

The appealing equation was:

But that description hides the hardest operational question:

What happens when thousands of tokens choose experts at the same time?

An MoE layer is not just a sparse neural network. It is also a real-time traffic-routing system.

Today we will look more closely at the router, capacity, load balancing, and why a theoretically sparse model can still run inefficiently.

The router makes a decision for every token¶

Suppose an MoE layer has \(E\) experts.

For a token representation:

the router can compute:

where:

The vector \(z\) contains one score per expert.

For eight experts:

token x

      router
        │
        ▼

E0   E1   E2   E3   E4   E5   E6   E7
.03  .07  .61  .04  .02  .18  .03  .02

With top-2 routing, the token is sent to:

E2 and E5

The router itself is tiny compared with the expert MLPs.

The difficult part is what happens when we repeat this decision for every token in a large batch.

Routing creates a token-to-expert matrix¶

Suppose we have six tokens and four experts.

The router might produce:

          E0   E1   E2   E3

token A    X
token B              X
token C         X
token D              X
token E              X
token F    X

Expert loads are then:

E0: 2
E1: 1
E2: 3
E3: 0

Already we have uneven work.

Now scale this to:

tokens and perhaps:

The router is dynamically creating a distributed workload.

That workload changes from batch to batch.

Why imbalance hurts even if total FLOPs are unchanged¶

Imagine four GPUs, each holding one expert.

A batch routes:

GPU 0 / Expert 0: 1000 tokens
GPU 1 / Expert 1: 1000 tokens
GPU 2 / Expert 2: 1000 tokens
GPU 3 / Expert 3: 1000 tokens

If processing 1000 tokens takes 10 ms, the expert phase finishes in roughly:

Now consider:

GPU 0 / Expert 0: 2500 tokens
GPU 1 / Expert 1:  700 tokens
GPU 2 / Expert 2:  500 tokens
GPU 3 / Expert 3:  300 tokens

The total number of expert-token operations is still:

But the layer cannot proceed until Expert 0 finishes its 2500 tokens.

If time scales roughly with token count, the phase now takes around:

So:

Parallel hardware cares about how evenly the work is distributed.

Why not simply force perfectly uniform routing?¶

We could tell the router:

exactly 25% of tokens must go to each of four experts

That would make the systems engineer happy.

But it could hurt the model.

Suppose a batch contains mostly source code and one expert has learned features especially useful for code.

Forcing the router to distribute those tokens uniformly may send many tokens to less suitable experts.

So there are two objectives:

and:

They are related, but they are not identical.

MoE training has to negotiate between them.

A simple load-balancing intuition¶

Suppose \(f_i\) is the fraction of tokens actually routed to expert \(i\), and \(p_i\) is the average router probability assigned to expert \(i\).

A balancing objective can encourage both quantities to be spread across experts.

One common family of auxiliary losses has the rough form:

The exact formula varies by architecture, but the purpose is more important than the details.

If one expert receives both high router probability and a large fraction of tokens, the auxiliary loss pushes against that concentration.

Training then optimizes something like:

The language-model loss says:

predict tokens well.

The balancing term says something like:

while doing that, do not route almost everything to the same experts.

The coefficient \(\lambda\) controls the tradeoff.

This is interesting because the loss function is now partly optimizing systems efficiency, not merely predictive quality.

Capacity factor turns routing into a queueing problem¶

Suppose \(T\) tokens are routed across \(E\) experts, with top-1 routing.

Perfectly balanced load is:

Instead of provisioning exactly that many slots, systems may allocate:

where \(\alpha\) is the capacity factor.

If:

then average load is:

tokens per expert.

Capacity becomes:

Each expert therefore has room for 25% more tokens than perfectly balanced routing would require.

That extra space absorbs some routing imbalance.

But it is not free.

Too little capacity versus too much capacity¶

Suppose Expert 3 receives \(1400\) tokens but capacity is \(1280\).

There are \(120\) assignments beyond capacity.

Different MoE designs handle this differently. They may drop assignments, reroute tokens, use backup experts, or use routing schemes designed to avoid overflow.

But every option introduces a tradeoff.

If capacity is too small:

more overflow
more dropped/rerouted work
possible quality problems

If capacity is too large:

larger buffers
more padding
more memory
potentially worse communication efficiency

So the capacity factor is analogous to reserving extra checkout lanes in a supermarket.

Too few lanes create queues.

Too many lanes waste space and staff.

Top-1 versus top-2 routing¶

Another major design choice is how many experts each token activates.

With top-1:

Each token uses one expert.

With top-2:

Each token uses two experts and combines their outputs.

Top-2 potentially gives the model more expressive computation per token, but expert work approximately doubles.

If there are \(T\) tokens, the number of expert assignments is roughly:

For \(T=8192\), top-1 creates \(8192\) expert assignments.

Top-2 creates:

That means more:

expert FLOPs
token movement
communication
buffer space

So \(k\) is not merely a modeling hyperparameter.

It is a systems hyperparameter.

The all-to-all step is really a permutation¶

Imagine four GPUs initially hold tokens by batch position:

GPU 0: A B C D
GPU 1: E F G H
GPU 2: I J K L
GPU 3: M N O P

After routing, suppose experts are distributed one per GPU.

The desired arrangement might be:

Expert/GPU 0: A F J N
Expert/GPU 1: B E K O
Expert/GPU 2: C G I P
Expert/GPU 3: D H L M

The system must transform the first layout into the second.

That is essentially a distributed permutation.

After expert computation, the outputs must be sent back so that tokens return to their original logical positions:

expert layout
     ↓
return all-to-all
     ↓
original token order

So every MoE layer may contain two large data rearrangements:

This is why fast all-to-all communication is so important for large MoE systems.

Token locality can matter¶

Suppose a token is already on GPU 2.

If its chosen expert also lives on GPU 2, no network transfer may be necessary.

If its expert lives on GPU 7, the token representation must cross the interconnect.

This suggests an intriguing possibility:

Could routing consider communication locality as well as model quality?

Some systems and research directions do exactly this in various forms: expert placement, hierarchical routing, expert replication, and topology-aware dispatch can all reduce expensive traffic.

But there is a danger.

If routing strongly prefers nearby experts merely to save communication, it may stop choosing the experts that produce the best model output.

Again:

MoE repeatedly forces us to balance the two.

Expert replication can relieve hot spots¶

Suppose Expert 7 becomes extremely popular.

Instead of storing only one copy:

GPU 7 → Expert 7

we might replicate it:

GPU 7  → Expert 7 copy A
GPU 15 → Expert 7 copy B

Now traffic can be split between the replicas.

This costs additional parameter memory but can improve throughput.

The tradeoff resembles caching in ordinary distributed systems:

But training makes replication harder than ordinary read-only caching because copies must remain synchronized as parameters update.

Routing behaves differently during inference¶

Training usually processes many tokens simultaneously.

That gives the router a large pool of work:

Large token counts make it easier to form efficient expert batches.

Autoregressive inference is different.

During decode, each active sequence contributes roughly one new token per decoding step.

Suppose only eight user requests are currently active.

A top-1 MoE layer may have only \(8\) tokens to distribute across dozens of experts.

The result might be:

Expert 2: 1 token
Expert 9: 2 tokens
Expert 17: 1 token
Expert 31: 1 token
...

Those are tiny matrix multiplications.

GPUs dislike tiny matrix multiplications because there is not enough parallel work to saturate the hardware.

This means MoE serving benefits enormously from batching tokens across many concurrent requests.

Continuous batching helps MoE twice¶

We previously saw that continuous batching improves LLM serving by keeping the GPU supplied with active sequences.

For MoE models, it has another benefit.

More simultaneous tokens create larger per-expert batches.

Instead of:

Expert 12 receives 2 tokens

a larger serving batch might produce:

Expert 12 receives 64 tokens

The expert's matrix multiplication becomes much more efficient.

So batching improves GPU utilization and also expert utilization.

This is a good example of how inference optimizations interact with model architecture.

A concrete mental model: airport security¶

Imagine an airport with eight specialized security lanes.

A dispatcher looks at each passenger and selects the appropriate lane.

If passengers distribute evenly, all lanes stay busy.

But suppose the dispatcher sends 70% of passengers to lane 3:

lane 3: enormous queue
other lanes: mostly idle

Adding more total screening capacity does not help much if traffic concentrates on one lane.

You can respond in several ways:

encourage better distribution
open extra capacity for popular lanes
replicate a popular lane
redirect overflow
change the dispatch policy

For an MoE model:

passengers       = tokens
dispatcher       = router
security lanes   = experts
lane capacity    = expert capacity
airport hallways = GPU network

The router is therefore doing much more than choosing "which expert knows the answer."

It is continuously scheduling work across a distributed machine.

The deeper lesson¶

Yesterday's simple MoE picture was:

The more complete picture is:

token
  ↓
router scores
  ↓
top-k selection
  ↓
capacity constraints
  ↓
distributed permutation
  ↓
expert queues
  ↓
expert computation
  ↓
return permutation
  ↓
token continues

The router therefore sits at the boundary between machine learning and distributed scheduling.

Its decisions affect:

model quality
specialization
network traffic
GPU utilization
memory usage
latency
throughput

That gives us today's main principle:

This is a broader lesson about modern AI systems. As models become more sparse and distributed, the distinction between "the neural network" and "the runtime executing the neural network" becomes increasingly blurry. Architectural choices determine communication patterns, and systems constraints increasingly influence what architectures are practical.