ST-MoE: Accelerating Large MoE Model Inference via Spatiotemporal Expert Prefetching

Mixture of Experts (MoE) is a mainstream approach for scaling large language models, but dynamic expert activation leads to severe expert loading latency issues. The ST-MoE framework, by mining the spatiotemporal correlation of expert activation and combining a lightweight prediction mechanism with reconfigurable hardware design, overlaps expert loading with computation, significantly improving the inference performance and energy efficiency of MoE models while maintaining model accuracy.

The MoE architecture achieves a balance between capacity and efficiency through sparse activation (each token activates only a few experts), parameter expansion (total parameters are several times that of dense models), and dynamic routing (gating networks determine which experts to activate). However, expert parameters are often stored in CPU memory or slow storage; on-demand loading requires steps like routing computation, memory check, and weight loading, leading to latency as a performance bottleneck and increased energy consumption.

The ST-MoE framework consists of two core components:

1. Lightweight runtime prediction mechanism: Uses spatial (correlation of expert selection across adjacent layers) and temporal (correlation of expert selection across consecutive tokens) predictors, combined with confidence evaluation to trigger prefetching, without changing the model's computation graph or output distribution;
2. Reconfigurable hardware support: Implements asynchronous loading via a dedicated prefetch engine, overlaps computation with loading to hide latency, and uses dynamic scheduling to adaptively adjust prefetching strategies.

ST-MoE brings significant benefits:

• Reduced inference latency: Overlapping expert loading with computation reduces waiting time;
• Increased throughput: Lower latency translates to more requests processed in batch scenarios;
• Optimized energy efficiency: Improves GPU utilization and reduces memory access energy consumption;
• Preserved accuracy: Only optimizes memory management without changing model weights or computation logic, so accuracy is consistent with the original implementation.

ST-MoE effectively solves the loading latency problem in MoE inference and improves performance and energy efficiency by deeply understanding the spatiotemporal correlation of MoE expert activation and converting it into prefetching opportunities. This work emphasizes that optimization comes not only from algorithmic innovation but also from insight into the essence of the problem, which is of great significance for inference optimization after MoE architecture becomes a standard for large models.

Application scenarios: Online services (reducing P99 latency), edge deployment (energy efficiency advantages), multi-tenant environments (improving cluster utilization); Future outlook: Learned prediction (using lightweight neural networks to improve accuracy), cross-request prefetching (sharing experts in multi-user scenarios), adaptive thresholds (dynamically adjusting prefetching strategies), heterogeneous storage hierarchical prefetching.