Expert Networks › Expert Capacity

When processing batches of sequences, MoE models face a unique challenge that dense transformers never encounter: different tokens route to different experts, but hardware prefers uniform workloads. GPUs and TPUs are most efficient when performing the same operation on identically-sized blocks of data. Batch matrix multiplications, in particular, assume that every element of the batch has the same shape and that no element requires extra synchronization. MoE routing violates this assumption by design: in a perfectly-functioning model, popular experts attract many tokens while niche experts process only a few. Without some mechanism to manage this variability, the hardware efficiency that motivated MoE in the first place can be partially or entirely lost to irregular computation patterns.

Think of expert capacity as a physical queue in front of each expert's desk in our library analogy. Each queue can hold at most $K$ patrons at a time. When a popular expert's queue is full and a new patron arrives, that patron cannot wait and instead must either go to their second-best option or handle their query on their own using only the general information they already have. This queue limit prevents the popular expert from becoming a bottleneck that makes everyone else wait, at the cost of occasional suboptimal service for tokens that get redirected or dropped. The capacity factor $C$ determines how long each queue can be relative to the perfectly-balanced case.

Expert capacity mechanisms manage this tension by imposing limits on how many tokens each expert can process per batch. The limits are not arbitrary restrictions: they are engineering constraints that make it possible to implement MoE layers as regular tensor operations on hardware that demands predictable memory layouts and computation graphs. Without these limits, implementing an MoE layer would require dynamic tensor shapes that change based on the routing decisions, making the computation far harder to optimize and parallelize.

The key insight is that capacity constraints create a tradeoff: tighter constraints improve hardware utilization and reduce memory overhead, but they also mean some tokens receive inferior processing (or none at all) when routing is imbalanced. Looser constraints preserve routing quality but waste memory and compute on empty buffer slots. Finding the right capacity factor is therefore a hyperparameter tuning problem with direct implications for both model quality and training efficiency.

Consider a batch of 1024 tokens being processed by an 8-expert MoE layer. If routing were perfectly balanced, each expert would receive exactly 128 tokens. In practice, some tokens cluster around certain experts while others receive few tokens. Without constraints, popular experts become computational bottlenecks while unpopular experts sit idle.

This load imbalance creates practical problems for efficient computation. Modern hardware, particularly GPUs, achieves peak throughput when performing regular, predictable operations on uniformly-sized data. When one expert must process 300 tokens while another processes only 50, we cannot fully utilize parallel processing capabilities. The expert with 300 tokens becomes a serialization point, and the expert with 50 tokens wastes potential compute cycles.

Load imbalance also creates memory management difficulties. In a static-graph compilation framework (such as XLA on TPU), every tensor must have a known shape at compile time. If expert queue lengths vary dynamically based on routing decisions, the compiler cannot determine tensor shapes in advance. The capacity mechanism solves this by pre-allocating a fixed-size buffer for each expert: we know at compile time that each expert will process at most $K$ tokens, so we can pre-allocate tensors of exactly that size. This static shape guarantee is what makes MoE layers compilable to efficient hardware kernels. The capacity constraint therefore is a load-balancing heuristic and a basic requirement for static-graph compilation.

To prevent bottlenecks where popular experts become overloaded, we calculate the expert capacity $K$ (the maximum tokens per expert) based on the total batch size and a buffer multiplier. This approach provides a principled way to balance routing flexibility against computational efficiency. We want each expert's buffer to be large enough to absorb typical routing variance without being so large that most of the allocated space goes unused.

The capacity formula computes how many tokens each expert is allowed to hold:

$$K = C \cdot \frac{T}{E}$$

where:

• $K$: the expert capacity (maximum tokens per expert for this batch)
• $C$: the capacity factor ($\geq 1.0$), a buffer multiplier that determines headroom above the perfectly-balanced allocation
• $T$: the total number of tokens in the batch
• $E$: the total count of available experts
• $T/E$: the baseline load per expert under perfectly balanced routing, serving as the reference point

The intuition behind this formula is straightforward. The ratio $T/E$ represents the average number of tokens each expert would receive under perfectly balanced routing. In our example with 1024 tokens and 8 experts, this baseline is 128 tokens per expert. The capacity factor $C$ then scales this baseline to determine how much deviation from perfect balance we allow. Setting $C = 1.25$ means each expert can accept up to 160 tokens (128 times 1.25). This provides a 25% buffer above the perfectly balanced load.

A capacity factor of 1.0 means each expert can handle exactly its "fair share" of tokens. Values greater than 1.0 provide headroom for imbalanced routing. The choice of $C$ affects model behavior in several ways:

• $C = 1.0$: Minimal buffer. Routing must be nearly perfectly balanced to avoid token dropping. Useful when load balancing losses are strong and routing is well-behaved.
• $C = 1.25$: Common choice in practice. This provides a 25% buffer for routing variance. This value was used in the original Switch Transformer paper and has become something of a default.
• $C = 2.0$: Large buffer. Rarely causes overflow but doubles the per-expert memory allocation. Useful for smaller models where routing is more difficult to balance.

Choosing the right capacity factor involves balancing memory efficiency against routing flexibility. A higher capacity factor allocates more buffer space for each expert's token queue, consuming more memory. However, it also reduces the probability that tokens are dropped due to capacity overflow, improving model quality. Empirically, the sweet spot for most models is between 1.0 and 1.5, with 1.25 being a reasonable starting point.

When an expert reaches capacity, additional tokens routed to it are typically "dropped": they bypass the MoE layer entirely and pass through via the residual connection. This ensures forward passes complete in a predictable amount of time, but it means some tokens receive less processing than intended. The model still produces an output for every dropped token, but that output is simply the unchanged input representation, as if the MoE layer were not there at all for those positions.

The drop rate provides a diagnostic signal during training. If 10% of tokens are dropped during training, the routing mechanism is not distributing load effectively, motivating the auxiliary losses we will explore in upcoming chapters on load balancing. You monitor drop rates as a health metric for your MoE training runs, as sustained high drop rates often correlate with degraded downstream performance. A good rule of thumb is that drop rates above 1 to 2% are worth investigating, and drop rates above 5% indicate a routing problem that should be addressed rather than tolerated.

Small amounts of token dropping are not necessarily harmful. Because the routing mechanism is differentiable (the router learns from gradients flowing back through the routing weights), moderate overflow can guide the router toward more balanced distributions. The prospect of dropped tokens creates an implicit incentive for the router to spread load more evenly, complementing explicit load-balancing loss terms. Consistently high drop rates indicate that the model is failing to use its expert capacity, and performance suffers.

Training typically uses fixed capacity factors to enable efficient batched operations. The predictable memory allocation allows for optimized kernel implementations and stable training dynamics. Fixed capacity also makes gradient computation more predictable, since the set of tokens that receive expert processing is determined before the forward pass rather than dynamically during it.

During inference, especially for single sequences, capacity constraints may be relaxed. Some implementations process all routed tokens regardless of capacity during inference, accepting variable compute time for better quality. This flexibility makes sense because inference often involves smaller batches where strict capacity constraints provide less benefit, and you typically prefer accuracy over strict latency guarantees. Single-request inference, for example, processes only one sequence at a time, so there is no competing workload to balance against and no risk of one expert starving another. For these cases, you can simply route every token to its top-k experts and process them all, recovering the full quality of the model.

Some production deployment strategies use a hybrid approach: during inference, maintain capacity constraints but set the capacity factor higher than during training. This preserves the memory and latency predictability of capacity-constrained routing while letting more routing flexibility than the training configuration. The optimal inference capacity factor is something you tune based on your latency budget and quality requirements.