APEX: Adaptive Expert Prefetching for Memory-Efficient Edge MoE Inference

2026-08-12

AI Summary

This work addresses the memory bandwidth bottleneck in Mixture-of-Experts (MoE) model inference on edge devices, primarily caused by off-chip expert parameter loading. To mitigate this, the authors propose APEX, a novel framework featuring a confidence-based adaptive expert prefetching mechanism. A lightweight prefetch router predicts candidate experts ahead of the attention module, and a dynamic confidence model orchestrates parameter prefetching to maximize overlap between computation and data loading. APEX operates in two modes—correctness-preserving and stall-free—both achieving substantial efficiency gains with negligible accuracy degradation. Experiments show that, in correctness-preserving mode, per-token latency is reduced by up to 26% and energy-delay product (EDP) improves by 41%; the stall-free mode further enhances throughput while incurring only marginal accuracy loss.

Abstract

Mixture-of-Experts (MoE) models are attractive for edge deployment because they provide high model capacity while activating only a small subset of parameters per token, improving compute efficiency. However, MoE inference at the edge is fundamentally limited by memory. Expert parameters are large and often reside in off-chip memory due to capacity, cost, and power constraints, putting expert loading to the critical path. We present APEX: Adaptive Expert Prefetching, a predictive resource management framework that overlaps expert loading with useful computation. APEX introduces a lightweight prefetch router that predicts candidate experts before the attention block to dynamically fetch additional experts using a learned confidence model. This adaptive strategy achieves over 99% overlap accuracy, significantly outperforming fixed top-k prefetching techniques. APEX supports two execution modes: a correctness-preserving mode that guarantees exact routing semantics, and a stall-free mode that eliminates residual stalls by operating on available experts with negligible impact on application accuracy. Across multiple MoE models, the correctness-preserving mode reduces per-token latency by up to 26% and improves energy-delay product (EDP) by up to 41% over state-of-the-art baselines, while the stall-free mode provides additional efficiency gains with negligible impact on application accuracy. These results establish adaptive, confidence-driven expert prefetching as an effective approach for efficient MoE inference on edge systems.

Problem

• Mixture-of-Experts
• edge inference
• memory bottleneck
• expert loading
• resource management

Innovation

• Adaptive Prefetching
• Mixture-of-Experts
• Edge Inference
• Memory Efficiency
• Confidence-based Routing