Speculating Experts Accelerates Inference for Mixture-of-Experts

TL;DR

This paper introduces a speculative expert prefetching method for Mixture-of-Experts models that overlaps expert loading with computation, significantly reducing inference time without sacrificing accuracy.

Contribution

It proposes a novel expert prefetching scheme that predicts future experts using internal representations, enabling efficient memory transfer overlap during inference.

Findings

Achieves up to 14% reduction in time per output token.

Reliable expert prediction maintains downstream task accuracy.

Lightweight estimators improve expert prediction hit rates.

Abstract

Mixture-of-Experts (MoE) models have gained popularity as a means of scaling the capacity of large language models (LLMs) while maintaining sparse activations and reduced per-token compute. However, in memory-constrained inference settings, expert weights must be offloaded to CPU, creating a performance bottleneck from CPU-GPU transfers during decoding. We propose an expert prefetching scheme that leverages currently computed internal model representations to speculate future experts, enabling memory transfers to overlap with computation. Across multiple MoE architectures, we demonstrate that future experts can be reliably predicted by these internal representations. We also demonstrate that executing speculated experts generally maintains downstream task accuracy, thus preserving more effective compute-memory overlap by eliminating the need to re-fetch true router-selected experts. Integrated into an optimized inference engine, our approach achieves up to 14\% reduction in time per output token (TPOT) over on-demand loading of experts from CPU memory. For MoEs where speculative execution alone yields suboptimal accuracy, we further examine lightweight estimators that improve expert prediction hit rates, thereby reducing performance degradation. Our code is released in open-source at https://github.com/axonn-ai/yalis/tree/offload_prefetch.