MoE: dynamism the compiler cannot remove

Because both the launch configuration and the pointer arguments depend on values computed inside the forward pass.

Take the natural implementation:

logits = router(x)                     # [T, E]
idx = logits.topk(k, dim=-1).indices   # [T, k] — data dependent
for e in range(E):
    sel = (idx == e).any(-1)           # runtime cardinality
    out[sel] += experts[e](x[sel])     # shape unknown until executed

Three separate blockers, and fixing one does not fix the others:

1. x[sel] has a data-dependent shape. The number of rows is a runtime value, so the GEMM’s $M$ dimension — and therefore the grid dimensions of the launch — cannot be recorded. A graph records a fixed launch configuration.
2. The index tensors are data. Even with fixed shapes, sel is computed on device, so any host logic depending on it needs a sync — and no host code runs during replay.
3. Allocation. x[sel] allocates a new tensor whose size is unknown, and allocation is forbidden during replay (Module 9).

Capacity-based MoE fixes all three at once, which is why it exists. Fix each expert’s buffer at $C \cdot Tk/E$ rows, scatter tokens into it with a fixed-shape scatter, run $E$ GEMMs of identical known shape, and gather back. Every shape is now a compile-time constant, every allocation is static, and there is no host logic in the loop. The layer becomes capturable and the whole of Part L applies again.

The price is stated plainly by the arithmetic in the table: with realistic routing skew you need $C \approx 1.5$ to avoid dropping tokens, so roughly a third of the expert compute is on padding. You bought a static graph with wasted FLOPs — which, at batch 1 where launch overhead dominates (Module 9), is very often the right trade even though it looks wasteful.

Capacity $C$ sets each expert’s buffer to $C \cdot Tk/E$ tokens, and it trades wasted compute against dropped tokens, mediated entirely by how imbalanced your router is.

• $C$ too low: the busiest expert overflows and excess tokens are dropped — they skip the FFN entirely and pass through on the residual. Training degrades subtly; inference produces measurably worse output for exactly the tokens the router thought were most specialised.
• $C$ too high: every expert’s buffer is sized for the worst case while most are half empty, so you pay FLOPs on padding in every expert, every layer, every step.

The distribution is what decides it. With perfect balance $C = 1.0$ suffices. Real routers skew because experts specialise, and a common shape is a busiest expert at 1.3–1.8× the mean — hence $C = 1.25$–$2.0$ in practice.

The levers that let you run a lower $C$:

• Load-balancing loss. An auxiliary term penalising the dot product of routing probability mass and actual token fractions, weighted around 0.01. It flattens the distribution and directly buys you a lower capacity factor.
• Noisy / stochastic routing during training, which prevents early winner-take-all dynamics.
• Expert-choice routing, which inverts the assignment: instead of each token picking $k$ experts, each expert picks its top $T k / E$ tokens. Load is then perfectly balanced by construction and $C = 1.0$ is exact — at the cost that some tokens get more experts than others and some get none, which is fine for training and awkward for autoregressive inference.
• Shared experts (DeepSeek’s design), where a always-on expert handles common patterns so the routed experts specialise more cleanly.

At inference the calculus shifts, because dropping a token is a visible quality regression rather than a training nuisance. Many serving systems run uncapped with grouped GEMM instead, accepting ragged shapes and losing CUDA Graphs for that layer — which is exactly the trade the next probe is about.

By moving the dynamism from the launch configuration into kernel arguments.

A grouped (or batched-variable) GEMM launches one kernel over a fixed grid sized for the total work, and each block reads a small descriptor array — pointers, and per-group $M$, $N$, $K$ — to discover which group it belongs to and where its data lives. The grid does not depend on the split of work between groups, only on the total, which is $Tk$ and therefore known. So no expert needs padding and no tokens are dropped: every block does useful work on some expert’s rows.

MegaBlocks framed this as block-sparse matrix multiplication, which is the clean way to see it — the MoE layer is one big block-sparse GEMM whose sparsity pattern is the routing assignment.

What you give up:

• Load balance within the kernel. Blocks assigned to a large group finish later, so the kernel’s duration is set by the largest group. You removed the padding waste, not the tail — and at extreme skew the tail can cost as much as the padding would have.
• CUDA Graphs, unless you are careful. The descriptor arrays are device tensors computed at runtime, so the graph can only be captured if those buffers are at fixed addresses with fixed maximum sizes. It is doable — allocate worst-case descriptor buffers and let the grid over-provision — but it is engineering rather than a decorator.
• Kernel complexity. You are now in hand-written CUTLASS or Triton territory (Module 7), not something Inductor generates.
• Autotuning difficulty. The best tile shape depends on group sizes that vary per batch, so a single configuration must serve a distribution.

The honest summary: grouped GEMM is the right answer for throughput-oriented inference where the padding waste is the dominant cost, and capacity-based padding remains competitive at small batch where launch overhead and graph capture matter more than FLOPs.

Because expert parallelism puts a latency-bound collective on the critical path twice per layer, and unlike the compute it does not shrink when you add GPUs.

With experts sharded across devices, each MoE layer must dispatch every token to whichever device holds its chosen experts, then combine the results back. That is two all-to-alls. Intra-node over NVLink they cost tens of microseconds; inter-node they are hundreds, and a 60-layer model with MoE every other layer performs 60 of them per forward pass.

Three properties make this worse than the raw number suggests:

1. It is a synchronisation point. Every rank waits for the slowest, so the most loaded expert’s device sets the pace — routing imbalance converts directly into collective latency, which is why the max expert load matters far more than the mean.
2. The message sizes are small and irregular. $T k d_{\text{model}}$ bytes split across $E$ destinations means many small transfers, which is the regime where interconnect latency rather than bandwidth binds.
3. Adding devices adds hops. Scaling expert parallelism increases the number of participants and often crosses more node boundaries, so the collective gets slower as you scale the thing you scaled it to make faster.

The mitigations are all about overlap and locality: overlap dispatch with the attention compute of the same layer, so the collective hides behind work that does not depend on it; place frequently co-selected experts on the same node to keep traffic intra-node; use shared experts that need no routing at all; and reduce $k$, since traffic scales linearly in it.

The framing worth extracting, and it is the hardware series’ argument arriving in a distributed setting: this is a latency problem wearing a bandwidth problem’s clothing. The links are not full; you are waiting on round trips and on the slowest participant.

Batch 1 is the case where you should accept substantial waste, because launch overhead is the dominant cost (Module 9) and the FLOPs are nearly free.

Concretely:

1. Fix capacity and pad. At batch 1 with a short sequence, $T$ is small and $Tk/E$ may be under one token per expert. Set every expert’s buffer to a fixed small size — often just a handful of rows — and accept that most are empty. The absolute waste is tiny because the tensors are tiny.
2. Never drop. At inference, dropping is a quality regression. Size capacity to the worst case you will accept and assert on overflow.
3. Make routing a fixed-shape scatter. Compute the top-$k$ on device, build a fixed-shape assignment matrix, and use index_add_/scatter into preallocated buffers — no boolean indexing, no data-dependent allocation.
4. Bucket the sequence length and capture one graph per bucket, sharing pools between them (Module 9).
5. Verify. Count cudaGraphLaunch calls per decode step: one per bucket is the target. TORCH_LOGS=graph_breaks should be empty for the MoE layer.

And the sanity check that decides whether any of this was worth it: measure the dense-equivalent time. At batch 1 an MoE layer that pads to full capacity is doing close to dense-model compute anyway, so if the graphed MoE is not clearly beating a dense model of the same active size, the architecture is buying you nothing at this batch size. MoE’s advantage is a throughput advantage; at batch 1 you are paying its complexity for a much smaller share of its benefit.

That is the honest conclusion of this module: MoE is a throughput architecture, and every technique here is about recovering enough staticness to stop it being a latency disaster.