MoE Capacity Factor: Why Your Tokens are Being Dropped
The Mechanics of Token Dropping
In an MoE layer, each expert has a static buffer. The size is determined by this formula:
capacity = capacity_factor * tokens * top_k / num_expertsWith a capacity_factor of 1.0, an expert can only handle its mathematically "fair share" of tokens. If the router sends more than that, the excess tokens are dropped. They bypass the expert MLP entirely and continue through the residual stream. Essentially, a 100-layer network suddenly has one of its blocks acting as a no-op for that specific token.
The reason we use fixed buffers instead of dynamic ones is purely for hardware efficiency. Static shapes allow for optimized batched matmuls and avoid the nightmare of recompiling kernels on TPU/XLA or dealing with ragged collectives.
Why This Happens in Production
Load imbalance is the default state for routers; they tend to favor certain experts (the "rich get richer" effect). While auxiliary load-balancing loss helps, it doesn't eliminate the problem.
The critical part is that drop rates depend on the token distribution of the current batch. This is why you won't see this in single-sequence evals but will see it under heavy production traffic. Even worse, dropping is order-dependent. A token's survival depends on how many tokens preceding it in the batch picked the same expert. This explains why the same sequence can produce different outputs depending on its position in a batch.
How to Fix It: A Practical Guide
If you're seeing quality dips under load, try these fixes in order of complexity:
1. Increase the capacity_factor: The quickest fix, though it increases memory overhead.
2. Implement Dropless Kernels: Switch to block-sparse kernels (like MegaBlocks-style grouped GEMM) that handle variable-length expert batches.
3. Advanced Balancing: Move toward auxiliary-loss-free balancing (similar to the per-expert routing bias used in DeepSeek-V3) to reduce the quality tax associated with forced balancing.
Pro Tip: You must instrument your drop rate per layer and per expert. If you aren't logging these metrics, you are flying blind and will never know why your model is underperforming in the real world.