MoE token dispatch — when sparse computation becomes a network problem

Yesterday we explored training the Mixture-of-Experts router: the router must choose useful experts without overloading a few and starving the rest. Day 75 introduced expert parallelism, where different GPUs hold different experts.

Today we'll follow a token through that distributed system and ask a practical question:

If an MoE model computes only two of eight experts per token, why isn't it automatically four times faster?

Because selecting an expert is not the same as executing it. The token may need to travel across GPUs, wait for other tokens to form an efficient batch, and return before the Transformer can continue.

The journey of one token¶

Suppose an MoE layer has eight experts, one per GPU, and uses top-2 routing. A token's hidden state \(h\) starts on GPU 0, but the router selects experts 3 and 6.

The system must execute something like:

GPU 0: token h
       |
       +--> router chooses experts 3 and 6
       |
       +--> send h to GPU 3 --> expert 3 --> return y3
       |
       +--> send h to GPU 6 --> expert 6 --> return y6
       |
       +--> combine: y = g3*y3 + g6*y6
       |
       v
    next Transformer operation

The router's weights \(g_3\) and \(g_6\) determine how the two expert outputs contribute. Both outputs must be available before the token's MoE result is complete.

At scale, implementations do not usually send thousands of tiny individual messages. They pack tokens by destination, exchange larger buffers with an all-to-all collective, execute experts, then perform a reverse exchange.

That produces the familiar pattern:

\[ \boxed{\text{route}\rightarrow\text{dispatch}\rightarrow \text{expert compute}\rightarrow\text{return}\rightarrow\text{combine}} \]

Why top-2 routing duplicates work and traffic¶

Let \(T\) be the number of tokens entering an MoE layer, \(k\) the number of selected experts per token, and \(d\) the hidden width.

There are

\[ A=Tk \]

token-to-expert assignments. Each assignment needs a copy of the token's hidden state at its destination, even if several experts were selected for the same token.

If activations use \(b\) bytes per element, the total logical dispatch payload is

\[ M_{\mathrm{dispatch}}=Tkdb. \]

The return payload is approximately the same size if every expert returns a \(d\)-dimensional output.

This does not mean all those bytes cross the network. Some assignments may target experts on the token's current GPU. Others may share links or be handled by more elaborate routing schemes.

If experts are evenly placed across \(E\) GPUs, one expert per GPU, and expert choices are independent of where tokens start, the expected remote fraction is

\[ f_{\mathrm{remote}}=1-\frac{1}{E}. \]

Under those assumptions, total remote traffic for dispatch and return is approximately

\[ M_{\mathrm{remote}}\approx 2Tkdb\left(1-\frac{1}{E}\right). \]

The factor of two is for the outbound hidden states and inbound expert outputs—not for the two selected experts, which are already counted by \(k\).

Worked example: a 4,096-token prefill¶

Consider a layer with:

\(T=4096\) tokens, \(d=4096\), BF16 activations (\(b=2\)).
Eight experts, each on a separate GPU.
Top-2 routing, approximately balanced across experts.

The number of assignments is

\[ A=4096(2)=8192. \]

Each hidden state is \(4096(2)=8192\) bytes, or 8 KiB. The total logical dispatch is therefore

\[ 8192(8192)=64\text{ MiB}. \]

If approximately \(7/8\) of assignments are remote, dispatch moves 56 MiB across GPUs. Returning the outputs moves another 56 MiB:

\[ M_{\mathrm{remote}}\approx112\text{ MiB} \]

across the eight-GPU group for one MoE layer.

With perfectly balanced sources and destinations, that's about 14 MiB of total send-plus-receive traffic per GPU. At a hypothetical 100 GB/s of effective communication bandwidth per GPU, the pure bandwidth floor would be roughly

\[ t_{\mathrm{network}}\gtrsim \frac{14\text{ MiB}}{100\text{ GB/s}} \approx0.15\text{ ms}. \]

This is an illustrative accounting, not a prediction of collective latency. Full-duplex links, network topology, contention, buffer packing, and synchronization change the result. Two separate all-to-all operations also incur startup costs.

The key observation is that a single MoE layer can create substantial network traffic even though each token executes only two experts.

The second problem: making expert computation efficient¶

Balanced routing would give each expert

\[ 8192/8=1024 \]

token assignments. That is a reasonably large batch for expert matrix multiplications.

Now imagine the same model during decode with only 32 new tokens across the serving batch. Top-2 routing produces just 64 assignments, or eight per expert on average.

The experts have not become smaller. Their weight matrices may still contain hundreds of megabytes of parameters, but each is now multiplying a matrix with only a few token rows.

This connects to Day 69's roofline model. With a large expert batch, weights can be reused across many tokens, and the computation may approach efficient matrix-matrix multiplication. With a tiny batch, the GPU may spend much of its time loading weights and launching kernels.

A useful approximation for an expert with \(N_e\) parameters and \(n_e\) assigned tokens is

\[ I_{\mathrm{weights}}\approx \frac{2n_eN_e}{b_wN_e} =\frac{2n_e}{b_w}, \]

where \(b_w\) is bytes per weight. With BF16 weights, \(b_w=2\), giving roughly \(n_e\) FLOPs per weight byte, ignoring activation traffic and other work.

So an expert batch of eight has about 8 FLOPs/byte from weight reuse, while a batch of 1,024 has about 1,024 FLOPs/byte. The difference is enormous.

Why grouped GEMM helps—but cannot create tokens¶

An MoE layer often has many experts receiving different numbers of tokens. Running a separate GPU kernel for each tiny expert batch can waste launch overhead and leave compute units underutilized.

A grouped GEMM implementation packages several experts' matrix multiplications into one coordinated operation. It can handle different expert batch sizes while reducing launch overhead and improving scheduling.

But grouped GEMM cannot eliminate fundamental problems:

If few tokens arrive, there may simply be too little work.
If one expert receives far more tokens than others, that expert can determine completion time.
If expert weights are spread across GPUs, token exchange still costs time.

Grouped GEMM improves the execution of the expert work after routing; it does not make routing free.

Three ways serving systems improve the economics¶

First, batch tokens across requests. Continuous batching can accumulate decode tokens from many independent requests, giving each expert more work per step. The tradeoff is familiar: waiting to form a larger batch can increase per-request latency.

Second, place or replicate popular experts strategically. If two GPUs hold the same expert, a dispatcher may choose the closer or less-loaded replica. That can reduce communication and hotspots, but duplicates expert weights and complicates placement. Replication must preserve the intended expert computation; it is not the same as changing the router's selected expert.

Third, overlap communication and computation. Systems can partition tokens into chunks and begin expert work as chunks arrive, rather than waiting for the entire exchange. This requires careful coordination and may not help much when batches are already tiny.

All three approaches shift the balance among memory capacity, network traffic, throughput, and latency.

Prefill and decode want different expert layouts¶

During prefill, thousands of tokens may be routed together. Expert batches are large, so the savings from sparse computation can outweigh dispatch overhead.

During decode, only one new token per active request enters each step. Expert batches can be small and irregular, and two all-to-all exchanges per MoE layer may add significant latency.

That suggests a connection to Day 68's disaggregated prefill and decode: a system might benefit from different expert-parallel configurations or replication strategies for the two phases. However, moving between layouts introduces its own model-state and KV-cache coordination costs.

There is no universal winner. A topology that maximizes prefill throughput may be a poor choice for low-concurrency decode.

Mental model: a specialist hospital with a pneumatic-tube system¶

Imagine a hospital with eight specialist departments. Every case is reviewed by two departments, chosen by a triage desk.

The hospital has much more expertise than any single department, and most cases avoid visiting the other six. That's MoE's sparse-computation advantage.

But records must travel to the selected departments and results must return. When hundreds of cases arrive together, each department can process a substantial queue efficiently. When only a few cases arrive, the paperwork and coordination can take longer than the specialist review.

Expert parallelism is the hospital's physical layout; all-to-all dispatch is the pneumatic-tube system; grouped GEMM is efficient batch processing within each department.

The deeper lesson¶

MoE does not simply exchange eight expert computations for two. It exchanges dense, predictable local computation for sparse, dynamically routed distributed computation.

\[ \boxed{\text{Sparse FLOPs are valuable only when routing and execution remain efficient.}} \]

The practical question is: At the actual number of tokens per inference step, does skipping expert computation save more time than token movement, small expert batches, and load imbalance cost?

That is why MoE model architecture and inference-engine design cannot be optimized independently.