271,104 Routes: What I Missed About Mixture-of-Experts

271,104.

That is the number of token-to-expert assignments I pulled out of one forward pass through IBM's Granite 3.1 1B-A400M model.

The model has 24 MoE layers. Each layer has 32 routed experts. Every token selects eight of them. I passed 1,412 real tokens through the model, which gave me:

1,412 tokens * 24 layers * 8 selected experts = 271,104 routes

Before I ran that experiment, I understood the basic description of Mixture-of-Experts:

Keep a large pool of FFNs, choose a small subset for each token, and only run the selected ones.

What I didn't understand was what happened after you choose the experts.

The router code is tiny. It produces expert ID, passes it through a softmax, and weights with a projection, a topk() call, and usually some kind of normalization.

But topk() does not move a token anywhere. The runtime still has to duplicate token rows, remember where they came from, group them by expert, move them to the right device, run the expert FFNs, return the outputs, apply the gate weights, and restore the original token order.

So this is a worklog/rant of my process learning MoEs from the inside out.

The question I kept coming back to was:

If only a few experts run for each token, where does the rest of the cost go?

My final answer that we will explore in this blog is:

MoE layer = route + pack + move + compute + combine

Though most of this post is me unpacking the parts after route.

1. The MoE Block

Figure 1. A dense block reuses one FFN for every token. A sparse MoE block keeps attention dense, adds a router, and activates only a small subset of a larger expert pool.

The name Mixture-of-Experts gave me the wrong picture at first. An expert isn't a smaller language model that specializes in a task like coding. It is one feed-forward network inside a Transformer block.

Attention still lets tokens exchange information across the sequence. Normalization and residual connections are still there. The main change is in the FFN path:

shared attention
    -> router
    -> selected expert FFNs
    -> combine
    -> shared residual path

The important part is that an FFN processes token rows independently. That means the runtime can reorder those rows, send different rows through different FFNs, and put the outputs back afterward. ^{1, 2, 3, 10}

For one token state x, the layer is usually written as:

y(x) = sum over selected experts i of p_i(x) * E_i(x)

In plain English:

• pick a few experts;
• run the token through them;
• multiply each output by its gate weight;
• add the results.

This runnable version makes that sequence literal. The Python loops are intentionally slow so it is obvious that only the selected experts execute:

import torch

torch.manual_seed(0)
x = torch.randn(3, 4)
experts = [torch.nn.Linear(4, 4, bias=False) for _ in range(4)]
expert_ids = torch.tensor([[0, 2], [1, 2], [3, 0]])
gates = torch.tensor([[0.7, 0.3], [0.4, 0.6], [0.8, 0.2]])

y = torch.zeros_like(x)
for token in range(x.size(0)):
    for slot, expert_id in enumerate(expert_ids[token]):
        y[token] += gates[token, slot] * experts[expert_id](x[token])

assert y.shape == x.shape
assert torch.allclose(gates.sum(dim=-1), torch.ones(3))

If a layer has 32 experts and uses top-2 routing, the token runs through two expert FFNs. The other 30 do not run for that token.

That is why it is known as sparse compute.

The experts themselves still do the matrix multiplications. The sparsity is in which experts receive which token rows.

This distinction ended up mattering more than I expected, because the runtime still has to turn sparse assignments into large expert batches.

The model wants conditional execution. The GPU still wants large regular matmuls. The runtime has to make both of those statements true at the same time.

2. The Router

Figure 2. Router logits become scores, top-k experts are selected, and their outputs are combined using gate weights.

Once the block was clear, I wrote the smallest router I could:

import torch

torch.manual_seed(7)
tokens = torch.randn(4, 6)
hidden_size, num_experts, top_k = 6, 32, 2
router = torch.nn.Linear(hidden_size, num_experts, bias=False)
logits = router(tokens)
top_logits, expert_ids = logits.topk(k=top_k, dim=-1)
weights = top_logits.softmax(dim=-1)

assert expert_ids.shape == (len(tokens), top_k)
assert torch.allclose(weights.sum(dim=-1), torch.ones(4))

For every token, this gives me two useful things:

• expert_ids: which experts were selected;
• weights: how much each selected output contributes.

Suppose one token gets:

expert_ids = [8, 26]
weights    = [0.72, 0.28]

The layer runs that token through experts 8 and 26, then combines the two outputs using those weights. ^{1, 2, 3, 10}

But that doesn't mean that expert 8 is “the expert for this token.” It means expert 8 received the largest selected weight for this token state, in this layer, under this router.

I am being annoyingly specific here because routing heatmaps make it very easy to promote a pattern into a “personality” test for expert IDs.

Typically the projection is small, but its output changes a lot:

• which transformations the token receives;
• which experts get training examples;
• how many rows each expert has to process;
• which devices may receive those rows.

This is where the router stopped looking like a classifier to me.

It is part of the model, but it is also a learned scheduler.

With T tokens and top-k routing, the expert side receives roughly:

T * k token-expert assignments

Granite selects eight experts per token. So every token creates eight assignments at every MoE layer. ^{15, 16}

The router has now made a decision, and it hasn't done the work required to execute it.

3. Dispatch: Turning Routes Into GPU Work

Figure 3. A real MoE layer permutes token vectors into expert batches, runs grouped expert matrix multiplications, then restores token order and combines top-k results.

Right after topk(), the assignments are still arranged by token.

The experts need them arranged by expert.

Suppose three tokens choose these experts:

token 0 -> experts 2 and 7
token 1 -> experts 7 and 4
token 2 -> experts 2 and 4

The runtime then needs to turn that into:

expert 2 -> token 0, token 2
expert 4 -> token 1, token 2
expert 7 -> token 0, token 1

With top-k routing, the same original token appears in several expert groups. The runtime therefore needs a route table that remembers the:

• original token;
• selected expert;
• position inside the expert batch;
• gate weight.

This runnable permutation groups the assignments by expert and then restores their original assignment order:

import torch

x = torch.tensor([[1.0, 1.1], [2.0, 2.1], [3.0, 3.1]])
expert_ids = torch.tensor([[2, 7], [7, 4], [2, 4]])
token_ids = torch.arange(len(x)).repeat_interleave(2)
flat_experts = expert_ids.flatten()
order = flat_experts.argsort(stable=True)
packed = x.index_select(0, token_ids[order])
expert_ends = torch.bincount(flat_experts, minlength=8).cumsum(0)

restored = torch.empty_like(packed)
restored[order] = packed
assert torch.equal(restored, x.index_select(0, token_ids))
assert expert_ends.tolist() == [0, 0, 2, 2, 4, 4, 4, 6]

After that, the data path looks like this:

token-order activations
    -> duplicate rows for top-k routes
    -> sort or group by expert
    -> move rows to the expert owner
    -> run expert FFNs
    -> move outputs back
    -> undo the grouping
    -> apply gate weights and add
    -> original token order

The token changes sets. In Granite's case, one token row becomes eight routed assignments per MoE layer.

A literal implementation could launch one tiny FFN operation for every assignment. It would be correct and almost useless on a GPU.

The point of packing is to gather all rows for one expert into a useful batch so the expert can run an ordinary dense matmul.

This is the part I had been skipping over when I said “the router sends a token to an expert.”

The router does not send anything. It produces metadata. The runtime turns that metadata into memory movement and GPU work.

When one expert gets too many tokens

Figure 4. Fixed expert buffers make shapes regular, but skew can overflow one expert while leaving padded slots elsewhere.

Routers do not naturally produce equal-size expert batches.

One expert may receive 40 assignments while another receives 3. That creates two separate problems.

The first is shape. Many implementations reserve a fixed number of slots for each expert because predictable tensors are easier to run efficiently. That limit is called expert capacity.

If an expert gets more assignments than it has slots, the runtime has to drop some routes, reroute them, reserve more space, or use a dynamic kernel that can handle variable loads.

None of those choices is free.

Dropping changes the computation the router asked for. Extra capacity creates padding. Dynamic packing keeps the routes but needs more complicated metadata and kernels.

The second problem is feedback.

If one expert is slightly better for the current traffic, the router may send it more tokens. More tokens give it more useful updates. It improves faster, so the router sends it even more traffic.

Now the model has a learning problem and the cluster has a straggler problem at the same time.

The hot expert gets more training data. The other experts get less. And the whole layer waits for the busiest expert batch to finish.

This is why load balancing is not just a utilization metric. It changes which experts learn and how long the layer takes to run. ^{3, 4}

A balanced router can still learn useless assignments. An unbalanced router can still have good language-model loss. I need both model metrics and traffic metrics to tell those cases apart.

When the experts live on different GPUs

If every expert lives on one device, dispatch is a local permutation.

If experts are split across devices, the same route becomes network traffic.

The token activations start on the GPUs that own the sequences. The selected experts may live somewhere else. So an MoE layer usually needs two communication phases:

source token owners
    -> dispatch all-to-all
    -> local expert computation
    -> combine all-to-all
    -> source token owners

The first trip sends token rows to the devices that own the selected experts. It is called an all-to-all because every rank may send a different number of rows to every other rank.

The second trip returns the expert outputs so the original token owners can apply the gate weights and continue the Transformer.

The amount of data matters, but so do the message sizes, number of peers, network topology, packing overhead, and the slowest rank. The router creates the traffic pattern. The cluster decides how expensive that pattern is. ^{2, 8, 9}

This is where “activate fewer parameters” turns into a distributed-systems problem.

4. The Granite Routing Experiment

At this point I wanted to stop being stuck in tutorial hell with no practical application.

I used IBM's Granite 3.1 1B-A400M base model. The checkpoint has 24 MoE layers, 32 routed experts per layer, and selects eight experts for every token.

I used 48 fixed prompts: 12 each for code, math, prose, and multilingual text. Each prompt stayed in its own batch element. After truncating at 96 tokens and ignoring padding, I had 1,412 real tokens.

That produced the 271,104 routes from the title. This was inference only. I did not generate text or train the model. ^{15, 16}

The extraction itself was small:

with torch.inference_mode():
    output = model(
        **batch,
        output_router_logits=True,
        use_cache=False,
    )

for router_logits in output.router_logits:
    top_logits, expert_ids = router_logits.float().topk(8, dim=-1)
    weights = top_logits.softmax(dim=-1)

Granite applies softmax after selecting the top eight logits, so the eight selected weights for each token add up to one.

I checked that invariant directly:

gate_sums = weights.sum(dim=-1)
torch.testing.assert_close(
    gate_sums,
    torch.ones_like(gate_sums),
    atol=1e-6,
    rtol=0,
)

The largest error I saw was 2.38e-7. I also reran fixed prompts from every domain. The expert IDs matched exactly and the largest weight difference was zero.

That is a property of this router. It is not really a rule that every MoE has to follow.

Experiment 1. Expert selection share across 24 layers for code, math, prose, and multilingual prompts.

What showed up

The first result was that the domains left different routing fingerprints.

Expert 8 had the largest overall selection share for code, math, and multilingual tokens. Expert 26 led for prose. But the leading shares were only about 4% to 5%, compared with 3.125% for perfectly uniform traffic.

So what I found was:

Expert 8 was selected more often for code in this sample.

The heatmap did not give me permission to rename it code_expert_final_final.

The second result was that multilingual routing was much more concentrated than prose routing.

Averaged across layers, the largest selected weight was 53.1% for multilingual tokens and 33.9% for prose. Normalized routing entropy moved in the opposite direction: 65.7% for multilingual and 84.2% for prose.

In other words, the multilingual routes tended to put more of the mixture weight on one selected expert. The prose routes spread that weight more evenly across the eight selected experts.

The third result was that code and math produced the most similar expert-load patterns. I measured that with mean layer-wise Jensen–Shannon distance: 0.169. Code and multilingual were the farthest apart at 0.287. Smaller means more similar here.

Those numbers describe traffic patterns, so they do not show what an expert “understands” per se.

The fourth result was that load skew depended heavily on the layer.

At layers 2 and 3, one expert appeared in every sampled token's top eight in each domain. That gave it four times the average expert load, which is the maximum possible when eight of 32 experts are selected.

Later layers behaved differently across domains.

Experiment 2. Router confidence, gate entropy, maximum-to-mean load, and the average weight assigned to each selected expert rank.

What did not show up

This was one checkpoint, one small hand-built prompt set, and unequal token counts after tokenization.

More importantly, I measured routing frequency and gate weights.

A routing heatmap tells me which experts received traffic. It does not tell me what computation those experts performed or whether one expert caused a result.

To make a better claim, I would need ablations, expert swaps, output comparisons, or feature-level analysis. That is a lot of work, though. Setting this up was already fun but storage-intensive :(

5. What I Learned

I started this expecting the router to be the hard part.

The router is important, yes, but the code that selects experts is not where most of the mechanism lives.

The route changes three things at once.

First, it changes the model's computation. Different tokens run through different FFNs, and only the selected experts receive token-specific gradients. ^{1, 3, 4}

Second, it changes the shape of the GPU work. The runtime has to turn irregular assignments into expert batches large enough to run efficiently.

Third, it changes the traffic pattern. If experts live on different devices, the route decides where activations move and which rank may become the straggler.

That is why I like to remember MoEs as this line I saw somewhere on X:

MoE layer = route + pack + move + compute + combine

“Sparse” describes the expert arithmetic selected for one token. It does not mean the other weights disappear. The full expert pool still has to live somewhere in memory, be sharded across devices, or be fetched when needed. ^{10, 13}

It also does not guarantee a speedup.

With very small expert batches, launch overhead and weight movement can dominate. With large expert-parallel groups, communication can dominate. Once expert batches are large enough, the matmuls usually become the main cost. ^{7, 8, 9, 13}

The bottleneck moves.

So the first questions I would ask when looking at an MoE are now:

• How many token-expert assignments does each layer create?
• How are those assignments packed into expert batches?
• What happens when one expert receives too many or too few rows?
• Which routes cross devices, and how large are the messages?
• Are the experts learning different functions, or only receiving different traffic?

I originally had separate sections on alternate routing schemes, upcycling, parallelism layouts, specialized kernels, serving tricks, and every paper I read on the way down.

They are useful topics. They also turned one question into a small textbook.

They belong in separate posts.

The main thing I wanted from this one was a clean picture of what happens to a token after the router makes a choice.

The answer is that topk() is the beginning, not the mechanism.

One small routing decision expands into a model decision, a batching problem, a memory shuffle, and sometimes two network collectives.

That is where the rest of the cost goes.

The next experiment

The next thing I want to run is an actual serving profile across several batch sizes and expert-parallel layouts.

I want one trace that lines up:

• router counts;
• permutation and packing;
• first all-to-all;
• expert matmuls;
• second all-to-all;
• combine;
• per-expert load.

That would let me separate costs created by the model's routes from costs created by the runtime and the cluster topology.

Thank you all for reading if you made it this far :)