Expert Capacity

Expert capacity is the maximum number of tokens that can be assigned to a single expert in a Mixture of Experts forward pass. Its primary purpose is to enable efficient batched computation on hardware (like GPUs) by creating fixed-size, dense tensors from sparse token assignments. Without a defined capacity, a variable and unpredictable number of tokens could be routed to each expert, making parallelization inefficient. Setting this parameter involves a direct trade-off: higher capacity reduces the risk of dropped tokens but increases memory usage and computational waste for underutilized experts.

Expert capacity is typically set using a derived hyperparameter called the capacity factor. The calculation is:

• expert_capacity = (batch_size * sequence_length * k) / num_experts * capacity_factor

Where k is the number of experts selected per token (e.g., top-2). A capacity factor greater than 1.0 (e.g., 1.25, 2.0) provides a buffer above the theoretical average load. This buffer accommodates the natural imbalance in token routing, preventing tokens from being dropped when an expert receives slightly more than its 'fair share' of assignments. The capacity factor directly controls the trade-off between computational efficiency and model accuracy.

When the number of tokens routed to an expert exceeds its predefined expert capacity, the system must handle the overflow. The standard solution is token dropping, where excess tokens are not processed by the expert. Common handling strategies include:

• Pass-through: The token's representation is passed unchanged to the next layer.
• Sent to a fallback expert (e.g., a shared expert or the first expert).
• Scaling router scores to fit within capacity. Dropped tokens represent a direct loss of information and degrade model performance. Therefore, setting expert capacity is an optimization to minimize dropping while controlling memory overhead.

Expert capacity has a direct, linear impact on GPU memory consumption during inference. Each expert's activation buffer must be allocated to hold the maximum possible number of tokens (expert_capacity * hidden_dim). Overallocating capacity leads to:

• Increased memory footprint: Storing large, partially empty buffers.
• Computational waste: Performing matrix multiplications on zero-padded sections of the buffer. Underallocating capacity leads to dropped tokens and accuracy loss. Systems like vLLM with MoE support and Fused MoE Kernels optimize this by managing memory pools and combining operations to reduce the overhead of this fixed-size allocation.

Expert capacity works in tandem with load balancing techniques during training. The router is trained with an auxiliary load balancing loss to encourage a uniform distribution of tokens across experts. Effective load balancing reduces the variance in tokens per expert, which allows for a lower capacity factor to be used safely. Without good load balancing, a few 'popular' experts would consistently hit their capacity limit, forcing a higher capacity setting to avoid excessive dropping, thereby increasing memory waste. Thus, the required expert capacity is a reflection of the router's imbalance.

In production serving systems, expert capacity is a key lever for inference cost optimization. Engineers tune it based on:

• Observed routing distributions from production traffic.
• Target latency and throughput SLAs.
• GPU memory constraints. Advanced systems may employ dynamic capacity strategies, where the capacity is adjusted per request batch based on real-time router statistics, or use expert parallelism to scale capacity across multiple devices. The goal is to find the minimal capacity that keeps the rate of dropped tokens below an acceptable threshold for the application.

A hyperparameter that defines a soft upper limit on the number of tokens assigned to an expert, calculated as (batch_size * sequence_length * k) / num_experts multiplied by the factor (e.g., 1.25). It provides a computational buffer to prevent dropped tokens while aiming to keep expert batches full for efficient GPU utilization. A factor too low increases drop rates; too high wastes memory and can increase communication overhead in expert-parallel systems.

Tokens that cannot be processed by their assigned expert because it has reached its expert capacity limit. This is a critical failure mode in MoE inference. Common handling strategies include:

• Passing the token unchanged to the next layer.
• Routing it to a fallback or overflow expert.
• Using an auxiliary buffer. Dropped tokens degrade model quality and represent wasted computation, making capacity tuning vital for inference cost optimization.

A set of techniques to ensure uniform computational load across all experts. Without it, the router may collapse, favoring a few experts and leaving others underutilized, which negates the efficiency benefits of MoE. Key methods include:

• Auxiliary Load Loss: An added training loss term that penalizes imbalance.
• Noise Top-k Gating: Adds tunable noise to router logits to encourage exploration.
• Expert Clipping: Enforcing capacity limits. Effective load balancing is prerequisite for setting optimal expert capacity.

The dominant collective communication pattern in expert parallelism. After the router assigns tokens to experts on different devices, an all-to-all operation scatters tokens from all devices to the devices hosting their assigned experts. After computation, another all-to-all gathers the outputs. The volume of this communication is directly constrained by expert capacity. High capacity can lead to larger, more imbalanced transfers, increasing router latency and creating network bottlenecks.

The core efficiency principle of MoE architectures, where only a small, dynamic subset of the model's total parameters (the chosen experts) is activated per token. Expert capacity governs how this sparse computation is batched. While activation is sparse at the token level, within an expert, computation is dense and batched across all tokens assigned to it. The goal is to set capacity high enough to enable large, efficient batches but low enough to preserve the memory and FLOPs advantages of sparsity.

The adaptation of dynamic batching techniques to MoE models. In systems like vLLM, requests are added to and removed from a running batch as they finish generation. For MoE, this is complex because the routing pattern changes each step. Expert capacity must be managed dynamically across the continuous batch, requiring sophisticated scheduling to minimize padding and dropped tokens while maximizing GPU utilization. This is a key focus for inference serving architectures.