Fast MoE Inference via Predictive Prefetching and Expert Replication
For practitioners deploying large MoE models, this work offers a practical method to significantly improve inference efficiency without major performance loss.
The paper tackles the problem of suboptimal GPU utilization and high latency in MoE inference due to expert load imbalance. Their dynamic replication strategy achieves near-complete GPU utilization and up to 3x speedup while preserving 90-95% of baseline performance.
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