Why architecture matters here
MoE serving architecture matters because the model shape is different. Dense-model tricks (tensor parallelism, sequence parallelism) still apply, but you now have a new form of parallelism — expert parallelism — with its own communication and load balance concerns. Get it right and you serve a 400B model at the compute cost of a 40B one. Get it wrong and you serve it at the wallclock time of a 400B one.
Cost matters more than most workloads. A poorly balanced MoE wastes 30-40% of GPUs on idle experts while a hot expert is the bottleneck. Right load balancing plus capacity tuning captures those idle FLOPs.
Reliability under uneven traffic is where MoE bites. Some prompts have vocabulary distributions that hit specific experts hard; others hit them lightly. The serving stack must handle both without dropping tokens or hitting a hot expert wall.
The architecture: every stage explained
Walk the diagram top to bottom.
Input tokens. The batch of tokens for the current step. In continuous batching, tokens from many concurrent requests are batched together.
Router. A linear layer per MoE block outputs logits over experts. Top-k (typically k=1 or 2) selects the experts each token visits. Modern models use k=2 for quality and k=1 for max efficiency.
Expert Assignment. Each token now has a list of experts to visit. Total assignments = batch × sequence × k. Assignments are the input to dispatch.
Dispatch (all-to-all). Under expert parallelism, experts live on different GPUs. Dispatch is an all-to-all collective: each GPU sends tokens targeted at experts hosted elsewhere, and receives tokens targeted at its own experts. Bandwidth-heavy step.
Expert Parallelism. Experts split across GPUs by expert index. E.g., 64 experts on 8 GPUs = 8 experts per GPU. Different from tensor parallelism which splits within a single tensor.
Expert MLP compute. Each GPU runs its experts' forward passes on the tokens dispatched to them. Standard dense MLP math per expert.
Combine (all-to-all). Reverse of dispatch. Expert outputs return to the GPU that originated each token. Bandwidth cost matches dispatch.
Capacity + Drop. Each expert has a per-batch capacity (e.g., 1.2x expected token count). Tokens beyond capacity are dropped (their contribution is zero for this block). Dropping vs capacity is a tuning axis.
Load Balancing. Auxiliary loss during training encourages the router to spread tokens evenly. At serving time, some models use random shuffling or dynamic rebalancing to prevent one hot expert.
Runtimes. vLLM's MoE support, TRT-LLM's expert parallelism, DeepSpeed-MoE, and Megablocks all implement the pipeline with different optimizations. Choose based on hardware, model, and integration needs.
End-to-end MoE forward pass
Trace a forward pass. Batch of 1024 tokens hits an MoE layer. Router outputs top-2 expert IDs per token: 2048 assignments total.
Dispatch: on 8 GPUs each hosting 8 of 64 experts, each GPU calculates which assignments target its experts and packs the corresponding tokens into send buffers. All-to-all: every GPU sends packed tokens to every other GPU based on target expert. Each GPU now has a batch of tokens to run through its 8 local experts.
Capacity: expert #17 got 200 tokens but capacity is 180. The extra 20 are dropped (contribution zero for this layer). Metric captured.
Expert compute: each GPU runs its 8 experts' MLPs on their respective token slices. Fully parallel; no cross-GPU comm during this phase.
Combine: reverse all-to-all. Expert outputs move back to token origin GPUs. Router weights combine top-2 outputs per token: y = w1·e1(x) + w2·e2(x).
Result: same token count out, different math for each. Compute cost = 2 of 64 experts activated ≈ 3% of a dense equivalent. Communication cost: 2 all-to-alls per MoE layer.