Fast MoE Inference via Predictive Prefetching and Expert Replication

Technical Summary: Fast MoE Inference via Predictive Prefetching and Expert Replication

Problem Statement
The Mixture of Experts (MoE) architecture has become a cornerstone for scaling Large Language Models (LLMs) by enabling massive parameter counts without proportional increases in computational overhead. However, MoE inference faces critical inefficiencies that hinder practical deployment:

1. Suboptimal GPU Utilization: Due to the inherent sparsity of expert activation, only a small subset of experts is active at any given time, leaving the majority of GPU resources idle.
2. Load Imbalance and Latency: Dynamic routing often results in uneven load distribution, where multiple tokens queue for the same physical expert, causing token stalling and increased latency.
3. Routing Overhead: Traditional softmax-based routing mechanisms introduce computational complexity and communication overhead (all-to-all routing) that do not scale favorably with increasing expert counts.

Existing optimization techniques, such as distillation or kernel-level enhancements, often offer limited improvements or compromise model accuracy.

Methodology: MoE-MPMC
The authors propose MoE-MPMC (Mixture of Experts with Multiple Producers Multiple Consumers), a framework that addresses these inefficiencies through predictive prefetching and dynamic expert replication. The system operates using two concurrent threads:

1. Inference Thread:

   • Retrieves a pre-computed hash table containing expert assignments and replication counts for the current batch.
   • Dynamically loads only the predicted active experts into GPU memory, offloading unused experts to CPU memory.
   • Executes the inference using a modified MoE layer where active experts are physically replicated on the GPU. This allows multiple tokens requiring the same expert to be processed concurrently by distinct physical replicas, eliminating queueing delays.

2. Hash-Building Thread:

   • Runs in parallel with the inference thread, processing the next batch of tokens.
   • Utilizes a Simple Recurrent Unit (SRU) with sparse attention to predict which experts will be activated for the upcoming batch.
   • Replaces the LSTM-based predictor used in prior work (SiDA-MoE) with SRU to achieve faster training and inference speeds due to reduced inter-dependencies and simpler gating mechanisms.
   • Computes replication counts for each predicted expert based on token demand and updates the hash table, which is then enqueued for the inference thread.

Key Contributions
The paper outlines three primary contributions:

1. SRU-Based Prediction: The introduction of the Simple Recurrent Unit (SRU) to rapidly predict expert selections for each token. This replaces slower LSTM predictors, significantly speeding up the prediction phase while maintaining accuracy.
2. Deep Physical Replication: A strategy to replicate predicted experts physically on the GPU based on token demand. This ensures that almost no token experiences waiting delays during batch processing, thereby maximizing parallelism and GPU utilization.
3. Capacity Capping: A practical mechanism to cap expert replication counts according to predefined GPU memory constraints. This ensures that replication is optimized within hardware limits to achieve near-optimal resource utilization without causing overflow.

Experimental Results
The authors evaluated MoE-MPMC on large-scale SwitchTransformers models (Switch-base-8, 64, 128, and 256) using datasets like SST2 and MRPC. Key findings include:

• GPU Utilization: The method achieves near-complete GPU utilization (up to ~100% in some configurations), a significant improvement over baseline SwitchTransformers (which showed utilization as low as 1.3% for 128 experts) and the SiDA-MoE approach.
• Inference Speed: The approach yields up to a 3x improvement in inference speed compared to baseline architectures.
• Accuracy: The method preserves approximately 90-95% of the performance of baseline architectures, with accuracy maintained better than SiDA-MoE due to the superior prediction capabilities of SRU over LSTM.
• Throughput: Throughput increases significantly as the number of experts increases, driven by the ability to process multiple tokens in parallel via expert replication.

Significance and Claims
The paper claims that MoE-MPMC represents a significant step toward the practical and economically viable deployment of MoE-based large-scale language models. By shifting from reactive load balancing to a proactive strategy of predictive prefetching and replication, the method:

• Mitigates memory access bottlenecks and token stalling.
• Optimizes GPU resource utilization without altering the underlying MoE computation graph.
• Demonstrates that predicting expert selections is feasible and effective when fine-tuning and pretraining occur concurrently with token prediction tasks.

The authors position this work as an advancement over reactive methods (like SiDA-MoE) by anticipating future demands and dynamically adjusting resource allocation, thereby addressing the typical pitfalls of idle GPU memory and load imbalances in sparse architectures.