Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Load Balancing › Capacity Factor Constraints

Michael BrenndoerferUpdated January 4, 20266 min read

Language AI HandbookAll chapters

Part 33 of 67Mixture of Experts

Chapter 243Load Balancing › Capacity Factor Constraints

Reading section 9 of 13 · Capacity Factor Constraints · 6 minSection 9 of 13 · 6 min

Section 9 of 13 · 6 min read

Load Balancing › Capacity Factor ConstraintsOne section · Sections in this chapter9 / 13 ↓

Sections in this chapter13

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

One approach to limiting imbalance is to impose a hard capacity constraint on each expert. Rather than letting the router to send arbitrarily many tokens to a popular expert, we define a maximum limit that caps how many tokens any single expert can process. This constraint prevents the most extreme forms of imbalance by rejecting tokens that would overflow an expert's capacity.

Think of capacity constraints as a fire code for experts. Just as a concert venue has a maximum occupancy limit that prevents dangerous overcrowding (even if more people want to attend), each expert has a maximum token capacity that prevents it from being overwhelmed. Tokens that arrive at a full expert are turned away, just as latecomers to a sold-out venue are turned away at the door.

We define a maximum capacity limit for each expert using the following formula. Given the total number of tokens $N$ and the number of experts $E$ , the ideal load per expert is $N/E$ . The capacity factor $C$ controls how much we allow individual experts to exceed this ideal:

K_{\text{max}} = \left\lfloor C \cdot \frac{N}{E} \right\rfloor

where:

$K_{\text{max}}$ : the maximum number of tokens an expert is allowed to process
$C$ : the capacity factor (typically $>1.0$ ), which determines the buffer size
$N$ : the total number of tokens in the batch
$E$ : the number of experts
$\lfloor \cdot \rfloor$ : the floor function, so an integer result

This formula sets the limit to a multiple $C$ of the ideal uniform load ( $N/E$ ), letting for natural variance while capping extreme imbalance. The capacity factor $C$ provides flexibility: a value of 1.0 would enforce strict uniformity, while higher values like 1.25 or 1.5 allow experts to handle somewhat more than their fair share before hitting the limit. The floor function ensures we get a whole number of tokens, since partial token assignments are not meaningful.

Why does this formula make sense? Notice that when $C = 1$ , $K_{\text{max}} = N/E$ , exactly the uniform allocation. When $C = 2$ , experts can handle double their fair share before tokens start being dropped. The parameter $C$ lets you tune the tradeoff between strict enforcement (low $C$ , more dropped tokens, better balance) and lenient enforcement (high $C$ , fewer dropped tokens, worse balance). Most implementations use $C$ between 1.1 and 1.5 to allow modest overloading without creating severe imbalance.

In[22]:

Code

def apply_capacity_constraint(routing_probs, capacity_factor=1.25):
    """
    Apply capacity constraints by dropping excess tokens.

    Args:
        routing_probs: (num_tokens, num_experts) tensor of routing probabilities
        capacity_factor: How much over-subscription to allow

    Returns:
        assignment_mask: (num_tokens, num_experts) binary mask
        dropped_count: Number of tokens dropped due to capacity
        expert_counts: Tokens assigned per expert
        capacity_limit: The maximum tokens allowed per expert
    """
    num_tokens, num_experts = routing_probs.shape
    capacity_limit = int(capacity_factor * num_tokens / num_experts)

    # Get top-1 assignments
    top_expert = routing_probs.argmax(dim=-1)

    # Track assignments per expert
    expert_counts = torch.zeros(num_experts, dtype=torch.long)
    assignment_mask = torch.zeros(num_tokens, num_experts)
    dropped = 0

    for token_idx in range(num_tokens):
        expert_idx = top_expert[token_idx].item()
        if expert_counts[expert_idx] < capacity_limit:
            assignment_mask[token_idx, expert_idx] = 1
            expert_counts[expert_idx] += 1
        else:
            dropped += 1

    return assignment_mask, dropped, expert_counts, capacity_limit


# Create imbalanced probabilities to demonstrate capacity dropping
# (Simulating a state where Expert 0 is highly preferred)
logits = torch.randn(num_tokens, num_experts)
logits[:, 0] += 2.0  # Bias toward Expert 0
probs = F.softmax(logits, dim=-1)

mask, dropped, counts, capacity_limit = apply_capacity_constraint(
    probs, capacity_factor=1.25
)

Out[23]:

Console

Capacity per expert: 320
Tokens dropped due to capacity: 1169 (57.1%)

Tokens assigned per expert (after capacity constraint):
  Expert 0: 320
  Expert 1: 69
  Expert 2: 86
  Expert 3: 65
  Expert 4: 99
  Expert 5: 75
  Expert 6: 83
  Expert 7: 82

Out[24]:

Visualization

Bar chart showing expert token counts before capacity constraint, with Expert 0 exceeding the fair share baseline. — Token distribution without capacity constraints. Expert 0 far exceeds the fair share baseline (green dashed line), showing the overload that occurs when routing is unrestricted. Tokens pile up at the popular expert without limit.

Bar chart showing expert token counts after capacity constraint, with Expert 0 truncated at the capacity limit. — Token distribution after enforcing a capacity limit (red dotted line). Excess tokens from Expert 0 are dropped rather than processed, truncating peak load and preventing the most extreme imbalance. The capacity-capped Expert 0 bar now aligns with the red dotted line.

Capacity constraints prevent catastrophic imbalance by rejecting tokens that would overflow an expert's capacity. The mechanism is effective at bounding the maximum load any single expert can experience, which directly bounds the Load Imbalance Factor: with a capacity factor of 1.25, the LIF can never exceed 1.25. However, this approach has a significant downside: dropped tokens receive no expert processing and must either be handled by a shared residual pathway or simply contribute less to the model's output.

The dropped-token problem is significant. Each dropped token represents information that the model fails to process correctly. In a language-modeling context, this might mean the model makes worse predictions for the next token because some context tokens were not processed by their appropriate expert. At high drop rates, this can degrade model quality.

The next chapter explores auxiliary losses that encourage balance during training rather than enforcing it through hard constraints. These losses add a penalty term to the training objective that discourages the router from developing strong preferences for specific experts. The key advantage over capacity constraints is that auxiliary losses are proactive rather than reactive: they prevent imbalance from developing in the first place, so tokens are never dropped.