Fast MoE Inference via Predictive Prefetching and Expert Replication

2026-05-12

AI Summary

This work addresses the low GPU utilization, load imbalance, and high latency in Mixture-of-Experts (MoE) inference caused by sparse expert activation. To overcome these limitations, the authors propose a prediction-based dynamic expert replication mechanism that anticipates overloaded experts and replicates them on-the-fly, enabling concurrent token processing across multiple layers. By integrating predictive prefetching, dynamic load balancing, and co-optimization with the MoE architecture, the method significantly enhances parallelism and hardware efficiency. Evaluated on Switch-base-128/256 models, the approach achieves near 100% GPU utilization, accelerates inference by up to 3×, and retains 90–95% of the original model’s performance.

Abstract

The Mixture of Experts (MoE) architecture has become a fundamental building block in state-of-the-art large language models (LLMs), improving domain-specific expertise in LLMs and scaling model capacity without proportionally increasing their computational overhead. However, MoE inference often suffers from suboptimal GPU utilization, load imbalance, and elevated latency arising from multiple tokens waiting on the same experts for their computation which arises from sparsity of expert activation. To address these challenges, we propose a dynamic expert replication strategy that predicts which experts are likely to be overloaded and replicates them for upcoming batches of tokens. The replicated experts process batch tokens concurrently across layers, which leads to improved parallelism, shorter GPU idle time, and significantly faster inference. Experimental evaluations conducted on large-scale MoE models, including Switch-base-128 and Switch-base-256, demonstrate that our method achieves near-complete GPU utilization (approx 100%), leading to upto 3x improvement in inference speed while preserving approximately 90-95% of the performance of baseline architectures

Problem

• Mixture of Experts
• GPU utilization
• load imbalance
• inference latency
• expert sparsity

Innovation

• Mixture of Experts
• expert replication
• predictive prefetching
• GPU utilization
• MoE inference