Pre-Attention Expert Prediction and Prefetching for Mixture-of-Experts Large Language Models

TL;DR

This paper introduces a lightweight pre-attention expert prediction method for MoE large language models, significantly improving expert prediction accuracy and enabling efficient prefetching, including in the first layer.

Contribution

It proposes a novel ranking-preserving approach using simple linear functions and pre-attention activations for accurate expert prediction in MoE LLMs.

Findings

Achieves over 93% prediction accuracy on multiple models

Improves prediction accuracy by about 15% over state-of-the-art methods

Supports expert prefetching in the first layer

Abstract

Mixture-of-Experts (MoE) Large Language Models (LLMs) efficiently scale-up the model while keeping relatively low inference cost. As MoE models only activate part of the experts, related work has proposed expert prediction and caching methods to prefetch the experts for faster inference. However, existing approaches utilize the activations from the previous layer for prediction, incurring low accuracy and leave the first layer unoptimized. Applying complex layers or even training standalone networks for better prediction introduces high computation overhead. In this paper, we propose pre-attention expert prediction to achieve accurate and lightweight expert prefetching. The key insight is that some functions in LLMs are ranking-preserving, indicating that matching the ranking of selected experts using simple linear functions is possible. Therefore, we utilize the activations before the attention block in the same layer with 2 linear functions and ranking-aware loss to achieve accurate prediction, which also supports prefetching in the first layer. Our lightweight, pre-attention expert routers achieve 93.03% accuracy on DeepSeek V2 Lite, 94.69% on Qwen3-30B, and 97.62% on Phi-mini-MoE, showing about 15% improvement on absolute accuracy over the state-of-the-art methods.