Title: Capacity-Aware Inference: Mitigating the Straggler Effect in Mixture of Experts
Primary area: foundation or frontier models, including LLMs
Keywords: Mixture of Experts, Load Balancing, Computation Efficiency

Abstract: The Mixture of Experts (MoE) is an effective architecture for scaling large language models by leveraging sparse expert activation to balance performance and efficiency. However, under expert parallelism, MoE suffers from inference inefficiencies due to imbalanced token-to-expert assignment, where underloaded experts complete computations early but must wait for overloaded experts, leading to global delays. We define this phenomenon as the \textbf{\textit{Straggler Effect}}, as the most burdened experts dictate the overall inference latency. To address this, we first propose \textit{\textbf{Capacity-Aware Token Drop}}, which enforces expert capacity limits by discarding excess tokens from overloaded experts, effectively reducing load imbalance with minimal performance impact (e.g., 
 speedup with only 
 degradation on OLMoE).
Next, given the presence of low-load experts remaining well below the capacity threshold, we introduce \textit{\textbf{Capacity-Aware Expanded Drop}}, which allows tokens to include additional local experts in their candidate set before enforcing strict local capacity constraints, thereby improving load balance and enhancing the utilization of underused experts. Extensive experiments on both language and multimodal MoE models demonstrate the effectiveness of our approach, yielding substantial gains in expert utilization, model performance, and inference efficiency, e.g., applying Expanded Drop to Mixtral-8
7B-Instruct yields a {0.2%} average performance improvement and a {1.85
} inference speedup. The code will be released upon acceptance.

TL;DR: We propose capacity-aware inference to mitigate imbalanced token assignments during inference, significantly improving efficiency without compromising performance.