Capacity-Aware Inference: Mitigating the Straggler Effect in Mixture of
  Experts
Authors: name: Shwai He, name: Weilin Cai, name: Jiayi Huang, name: Ang Li

The Mixture of Experts (MoE) is an effective architecture for scaling large
language models by leveraging sparse expert activation, optimizing the
trade-off between performance and efficiency. However, under expert
parallelism, MoE suffers from inference inefficiencies due to imbalanced
token-to-expert assignment, where some experts are overloaded while others
remain underutilized. This imbalance leads to poor resource utilization and
increased latency, as the most burdened expert dictates the overall delay, a
phenomenon we define as the \textit{Straggler Effect}. To mitigate
this, we propose Capacity-Aware Inference, including two key techniques: (1)
\textit{Capacity-Aware Token Drop}, which discards overloaded tokens
to regulate the maximum latency of MoE, and (2) \textit{Capacity-Aware
Token Reroute}, which reallocates overflowed tokens to underutilized experts,
balancing the token distribution. These techniques collectively optimize both
high-load and low-load expert utilization, leading to a more efficient MoE
inference pipeline. Extensive experiments demonstrate the effectiveness of our
methods, showing significant improvements in inference efficiency, e.g., 0.2\%
average performance increase and a 1.94times inference speedup on
Mixtral-8times7B-Instruct.