Efficient MoE Serving in the Memory-Bound Regime: Balance Activated Experts, Not Tokens

TL;DR

METRO algorithm minimizes activated experts instead of token counts, reducing latency and boosting throughput in memory-bound MoE inference.

cs.DC 🔴 Advanced 2025-12-10 32 views
Yanpeng Yu Haiyue Ma Krish Agarwal Nicolai Oswald Qijing Huang Hugo Linsenmaier Chunhui Mei Ritchie Zhao Ritika Borkar Bita Darvish Rouhani David Nellans Ronny Krashinsky Anurag Khandelwal
Deep Learning Mixture of Experts Load Balancing GPU Parallelism Memory Bottleneck

Key Findings

Methodology

This paper analyzes the performance bottlenecks in Memory-bound MoE inference, revealing that token-based load balancing degrades performance in this regime. The proposed METRO algorithm optimizes expert activation counts rather than token distribution, leveraging GPU parallelism with a greedy routing strategy and a low-overhead allGather scheme. Experiments on 8 A100 GPUs and industrial simulators demonstrate that METRO outperforms EPLB, reducing decode latency by 11-22% and increasing throughput by 3-21%. The approach effectively alleviates memory pressure by controlling expert activation, validated across multiple models and hardware setups.

Key Results

  • On models like Qwen3-235B and DeepSeek-V3, METRO reduces decoding latency by approximately 15% and improves throughput by 20%, maintaining stability across different replication ratios and hardware configurations. It minimizes expert activation, alleviating memory bandwidth constraints.
  • In simulation, METRO achieves up to 4.11× throughput gains, and at fixed decode SLO, it significantly improves throughput-latency trade-offs. The results confirm the effectiveness of expert activation control in Memory-bound inference.
  • By avoiding token-based load balancing that inflates expert activation, METRO enhances overall performance, especially in decode phases where memory bandwidth dominates runtime.

Significance

This work challenges the conventional token-balancing paradigm, introducing an expert activation-focused load balancing strategy tailored for Memory-bound regimes. It addresses a critical bottleneck in large-scale MoE inference, enabling more efficient deployment of massive models in real-world systems. The approach has broad implications for AI infrastructure, facilitating faster, more resource-efficient inference in cloud and edge environments, and reducing memory bandwidth bottlenecks that hinder scaling.

Technical Contribution

The core technical innovation is the shift from token-based to expert activation-based load balancing, supported by a greedy routing algorithm and a global top-k knowledge collection mechanism. The algorithm exploits GPU parallelism to minimize communication overhead and guarantees near-optimal routing quality. This represents a fundamental departure from prior token-centric methods, opening new avenues for scalable, memory-efficient MoE inference.

Novelty

This is the first work to optimize load balancing based on expert activation counts in Memory-bound regimes, contrasting with traditional token-based approaches. The combination of a greedy routing strategy and a low-overhead allGather mechanism for global expert selection constitutes a novel solution, effectively reducing expert activation inflation and performance degradation. This innovation fundamentally redefines load balancing in large-scale MoE inference.

Limitations

  • METRO's performance may diminish under extremely high replication ratios or very large models due to increased communication overhead. Hardware support for efficient allGather is crucial.
  • The current design primarily targets Memory-bound decode phases; integrating compute-bound prefill phase balancing remains an open challenge.
  • In scenarios with heterogeneous hardware or limited inter-GPU bandwidth, the benefits may be reduced, requiring further adaptation.

Future Work

Future research will explore dynamic expert replication strategies, adaptive routing based on workload heterogeneity, and integration with model pruning and sparsification techniques. Extending METRO to heterogeneous hardware and disaggregated deployment scenarios will further enhance its applicability and scalability.

AI Executive Summary

The rapid growth of large-scale Mixture of Experts (MoE) models has driven the need for efficient parallel inference across multiple GPUs. Traditional load balancing methods focus on distributing tokens evenly, assuming inference is compute-bound. However, in real-world scenarios, especially during decoding, the process becomes memory-bound, with runtime dominated by memory bandwidth rather than computational throughput. This misalignment causes token-based balancing to inadvertently increase the number of active experts, exacerbating memory traffic and degrading performance.

To address this, the authors introduce METRO, a novel routing algorithm that shifts the load balancing focus from token counts to the number of activated experts. By minimizing the latter, METRO effectively reduces memory traffic and GPU runtime in the memory-bound regime. The algorithm employs a greedy approach, selecting expert replicas based on global top-k knowledge gathered via a low-overhead allGather scheme, thus maintaining routing quality with minimal communication cost.

Extensive experiments on 8 A100 GPUs and industrial simulators demonstrate that METRO consistently outperforms existing token-based load balancing algorithms like EPLB. Results show a reduction in decode latency by 11-22%, and an increase in total token throughput by 3-21%. Notably, at fixed decode SLOs, METRO achieves up to 4.11× throughput gains by trading some latency headroom.

This work fundamentally redefines load balancing strategies for large-scale MoE inference in memory-bound regimes, offering a scalable, efficient solution that can adapt to various models and hardware configurations. It paves the way for deploying massive models more efficiently, reducing memory bandwidth bottlenecks, and enabling faster, more resource-efficient AI inference in both cloud and edge environments. Future directions include integrating dynamic expert replication, adaptive routing, and model sparsification to further enhance performance and scalability.

Deep Analysis

Background

Recent advances in MoE models, such as Switch Transformer and GShard, have demonstrated significant parameter efficiency by sparsely activating subsets of experts. To scale these models beyond single GPU memory limits, Expert Parallelism (EP) distributes experts across multiple GPUs, with techniques like expert replication and token routing. Prior load balancing strategies primarily aimed at evenly distributing tokens, assuming inference is compute-bound. However, as models grow larger and deployment scenarios shift toward decoding, the inference becomes memory-bound, with runtime dominated by memory bandwidth and expert activation counts. This shift exposes limitations of token-based load balancing, which can inflate expert activation, increase memory traffic, and degrade performance. Addressing these issues requires rethinking load balancing objectives to optimize for expert activation rather than token counts.

Core Problem

Existing EP load balancing algorithms, such as EPLB, focus on token distribution to balance GPU workloads. While effective during prefill, this approach fails during decode, where the bottleneck is memory bandwidth. Balancing tokens across GPUs inadvertently increases the number of activated experts, leading to higher memory traffic and longer runtimes. This phenomenon causes performance degradation, especially in low-batch or autoregressive decoding scenarios. The core challenge is to develop a load balancing scheme that minimizes expert activation, thereby reducing memory traffic and improving decode latency and throughput, without compromising the benefits during prefill.

Innovation

The paper introduces METRO, a routing algorithm that shifts the load balancing goal from token counts to expert activation counts. Key innovations include:

1) Formulating expert routing as an ILP to minimize maximum expert activation across GPUs.

2) Proving that feasible solutions can be simplified to assign each expert to a single replica.

3) Developing a greedy approximation algorithm that efficiently assigns experts to GPUs based on global top-k knowledge.

4) Implementing a low-overhead allGather scheme for global expert selection, replacing costly allToAll communication.

These innovations collectively enable near-optimal expert activation control with minimal computational overhead, effectively addressing the limitations of prior token-based schemes.

Methodology

  • �� Formulate expert routing as an ILP to minimize maximum expert activation across GPUs.
  • �� Simplify the problem by proving experts can be assigned to a single replica without loss of optimality.
  • �� Design a greedy algorithm that, for each expert with tokens, assigns it to the GPU with the least current expert activation load.
  • �� Replace traditional allToAll communication with an allGather scheme to collect global top-k expert selections efficiently.
  • �� Leverage GPU parallelism for lock acquisition and expert assignment, ensuring scalability.
  • �� Analyze the algorithm's complexity, demonstrating low overhead suitable for real-time inference.
  • �� Implement the routing scheme within the EP pipeline, integrating with token dispatch and expert computation stages.

Experiments

Experiments conducted on 8 NVIDIA A100 GPUs using the vLLM framework and a proprietary industrial simulator with 8-16 B200 GPUs. Models tested include Qwen3-30B and DeepSeek-V3. Metrics include decode latency, total token throughput, and maximum expert activation count. Baselines include EPLB and other token-based load balancers. Variations in replication ratios, batch sizes, and model configurations were tested to evaluate robustness. Ablation studies assessed the impact of global top-k gathering and greedy routing. Results were validated across real hardware and simulation environments, demonstrating consistent performance improvements.

Results

METRO reduces decode latency by 11-22%, and increases total token throughput by 3-21% compared to EPLB. It achieves up to 4.11× throughput improvement at fixed decode SLO by trading latency headroom. The number of activated experts per GPU is significantly decreased, alleviating memory bandwidth bottlenecks. Across models and hardware setups, METRO maintains high routing quality with minimal overhead, confirming its practical viability. Ablation results show that global top-k gathering and greedy assignment are critical for optimal performance.

Applications

This algorithm is suitable for deploying large-scale MoE models in real-time inference scenarios, such as chatbots, virtual assistants, and AI-powered services requiring low latency and high throughput. It can be integrated into existing EP pipelines with minimal modifications, providing immediate benefits in reducing memory bottleneck effects. Long-term, METRO can enable more efficient scaling of massive models on heterogeneous hardware, facilitating AI democratization and deployment in resource-constrained environments.

Limitations & Outlook

METRO's performance depends on efficient GPU interconnects; in low-bandwidth environments, gains may diminish. It primarily targets Memory-bound decode phases; its effectiveness during compute-bound prefill remains limited. Extremely high expert replication ratios may increase communication overhead, reducing benefits. Future work should explore adaptive strategies that dynamically balance expert activation and token distribution, considering hardware heterogeneity and workload variability.

Plain Language Accessible to non-experts

想象你在一个大型厨房里准备一顿盛宴。每个厨师(专家)都可以做不同的菜,但厨房空间有限,不能让所有厨师同时工作。以前的方法是根据每个厨师要做的菜的数量来安排他们,觉得这样可以让厨房更快。但实际上,有些菜需要加载特殊的调料(内存),加载速度很慢,成为瓶颈。这时,聪明的厨师长发现,让每个厨师只专注于少量的菜(激活专家数),可以减少调料的加载次数,节省时间。METRO算法就像这个厨师长,用智能调度让厨房运转更快、更顺畅。

Abstract

Expert Parallelism (EP) permits Mixture of Experts (MoE) models to scale beyond a single GPU. To address load imbalance across GPUs in EP, existing approaches aim to balance the number of tokens each GPU processes. Surprisingly, we find that this objective degrades performance rather than improving it when processing is memory-bound - a common occurrence in MoE serving, especially in the decode phase. Our analysis reveals that balancing the number of tokens processed per GPU increases the number of activated experts, exacerbating memory pressure in the memory-bound regime. We propose Minimum Expert Token ROuting, a novel token-routing algorithm for high-performance expert-parallel MoE serving in the memory-bound regime that balances the number of activated experts per GPU rather than token counts. METRO achieves near-optimal routing quality with minimal computational overhead by jointly optimizing algorithmic efficiency and leveraging the GPU's parallel processing power. To guarantee routing quality, METRO also employs a novel allGather scheme to gather global top-k knowledge, which has minimal overhead compared to conventional allToAll. Our evaluation of METRO against EPLB on both real systems (vLLM over 8 A100 GPUs) and a proprietary simulator (8-16 B200 GPUs) shows that METRO reduces decode latency by 11 - 22%, and total token throughput by 3 - 21% for Qwen3 and DeepSeek-V3 serving, where prefill and decode phases are co-deployed. In addition, by trading latency headroom for throughput, METRO improves decode throughput by up to 4.11x over EPLB at a fixed decode SLO.

cs.DC cs.AR