Reading:MoE-APEX: An Efficient MoE Inference System with Adaptive Precision Expert Offloading

Overall Introduce

MoE architecture use gating functions to select k experts for every token, expanding computation while increasing model parameters.

However, the number of parameters in complete MoE models is so huge that it’s difficult to deploy such models on edge devices.

At the same time, they find out that loading time accounts for 85.8% and 88.1% on RTX 4090 and Jetson Orin while inferencing, but computation time only accounts for a minimal portion, indicating that MoE models spend most of the time on loading while minimal time on computation.

Insights

Their key insight is that not all experts contribute equally to model outputs, making it possible to selectively replace less critical experts with low-precision variants during cache missing. To realize such idea, their contributions are summarized as follows:

  • Propose a token-level dynamic expert loading mechanism that reduces latency through low-precision replacement of less critical cache-miss experts, maintaining accuracy and flexibility.
  • Introduce a layer-level adaptive expert prefetching technique with high prediction accuracy, leveraging mixed-precision prefetching to optimize computation-communication overlap.
  • Develop a sequence-level cost-aware expert caching policy that combines model-specific locality characteristics with mixed-precision features to efficeintly manage the expert cache and minimize miss penalties.
  • Implement MoE-APEX on top of Llama.cpp with 8.5k addition lines of C/C++ code, and evaluate it on four MoE-based LLMs across three memory-limited platforms, demonstrating up to 9.75× speedup in decoding over state-of-the-art systems.

Motivations

  1. Expert loading dominates inference cost: During inference, experts loading process accounts for most of the time, over 80%, and computation only needs minimal portion of the time.To cope with this problem, some researchers employed dynamic gating to limit the number of experts loaded which needed significant accuracy trade-offs, or used experts skip method which resulted in notable degradation of model performance.
  2. Mixed precision expert preserves model accuracy: Given that skipping experts can make degradation of perfermance, they notice that replace the less important experts with their low-precision versions can not only maintain accuracy(less than 1%) but also reduce expert-loading cost(8× speedup).

MoE-APEX System

Overview

  1. Select the top-k required experts.
  2. Predicts the experts needed for subsequent layers.
  3. Check if the required experts are present in the expert cache and update the data record with its Policy Performer.
  4. If all the experts are in the cache, perform the expert computation on GPU cores.
  5. If some of the required experts are not in the cache, Dynamic Expert Loader uses the Expert Scorer to handle the cache miss by danamically generating the corresponding loading task and adding them to the Task Queue.
  6. Load the corresponding experts into the Expert Cache.
  7. Wait for all on-demand expert loading tasks to complete before conputing and then compute the outputs.

Token-Level Dynamic Expert Loading

They find that gating output magnitude is highly positively correlated with expert output contribution , with Pearson correlation coefficient of 0.99. So it’s practical to use as an indicator of expert importance.

In detail, they firstly rank the selected K experts in descending order of . Next, caculate the cumulative score for each expert as follows:

After caculation, they deviding the experts into three groups by setting two threshold numbers and . If , the expert is considered important and load its high-precision version; If , the expert is considered less important and load its low-precision version; If , the expert isn’t considered important, and finally skip it.

From the formula, we can find that the first expert is always considered important.

When completing deciding how to load experts, Expert Scorer dynamically generates mixed-precision loading tasks for each expert, adding them into Task Queue, and finally Expert Scheduler loads those experts as following picture:

Layer-Level Adaptive Expert Prefetching

Due to the layer-by-layer structure of LLMs, gating input of continuous layers has high cosine similarities, providing a method by leveraging the gating input from the current layer to predict the required experts for subsequent layers. According to experiments, they proved that the prediction method is reliable(Figure 8). So it’s practical to design Adaptive Expert Predictor to realize expert prefetching.

The process can be described as follows:

  1. Predict required experts of subsequent layers from the current layer.
  2. If all the experts required by next layer are in the cache, continue predicting next one, until finding expert missing in cache or finishing the whole prediction.
  3. Predicted experts will be loaded in different precisions and be masked to prevent them from being evicted from the expert cache, as they are high likely to be used in the subsequent layers.

As we all know, more gates in layers make more computation overhead. So in order to reduce the overhead of gating computation, they design Stacking Computer to compute all gating modules at once using several tensor operations, including stacking, matrix multiplication, and top-k selection, and to adaptively select the required experts for preloading, finally maintaining the accuracy while minimizing the prediction overhead.

Sequence-Level Cost-aware Expert Caching

Loading costs are various because of the different mixed-precision expert loading. Traditional method, like LRU or LFU are less efficient when dealing with such problem, so the team come up with Least Costly Used(Also refered as LCU) strategy.

First, they define expert cost in the following way:

where is the frequency of high-precision usage in the current sequence, is the frequency of low-precison usage, and and represent the bit-widths of the low- and high-precision versions. The final priority of expert is defined as follows:

where is current token ID, is the set that contains the experts used in last forward pass, and is the layer distance of expert to expert , computed as in which is the layer ID of expert , is the layer ID of the currently used expert , is the total number of layers in the model, and is a small constant(0.1 such as) to avoid division by zero when .

At the same time, Caching Manager will maintain low- and high-precision cache and record frequency of their usage. When cache missing happens, Policy Performer will replace the lowest-priority expert with required one according to the formula above.The Policy Performer resets all records at the start of each new sequence.

System Implementation

There are two modes:

  1. GPU-centric computing mode: The main thread handles input on GPU. Missing experts are loaded from CPU/SSD by scheduler thread.
  2. CPU-GPU cooperative computing mode: The missing experts are computed on CPU by helper thread, and results are sent to GPU.

This papper focuses on GPU-centric mode, for limited CPU resources on edge devices.

Experimental Evaluation

They evaluated Mixtral-8x7B, Phi-MoE, DeepSeek-MoE, and DeepSeekV2-Lite on Jetson AGX Orin, RTX 4090, and RTX 2080 Ti. The comparison baselines include Transformers, DeepSpeed-Inference, Llama.cpp, MoE-Offloading, MoE-Infinity, AdapMoE, and Fiddler.

Main Results

  • End-to-end Performance: MoE-APEX achieves the best decoding speed and prefill latency on all platforms. Compared to Llama.cpp, on Jetson Orin, the average decoding speedup is 12.0× (Mixtral-8x7B) and 18.57× (Phi-MoE); compared to MoE-Infinity, the speedups are 3.36× and 9.75×. On RTX 4090, compared to AdapMoE, the decoding speedups are 1.34× (Mixtral-8x7B) and 1.59× (Phi-MoE). On RTX 2080 Ti, DeepSeek-MoE and DeepSeekV2-Lite achieve decoding speedups of 1.49× and 1.68× compared to AdapMoE.
  • Model Accuracy: After adopting Float16+Int2 or Int8+Int2 mixed-precision strategies, the accuracy drop across all models and datasets does not exceed 1%.
  • Ablation Studies:
    1. Dynamic expert loading brings a 1.22×–1.53× speedup.
    2. Adaptive prefetching reduces prefill latency by approximately 10%; combined with dynamic loading, the decoding speedups are 1.10× (Mixtral-8x7B) and 1.15× (Phi-MoE).
    3. LCU reduces cache miss penalties by 4.85%–7.88% compared to LRU, and by 2.36%–3.10% compared to LFU; sequential-level LFU can improve the hit rate by 4.5%.
  • Overhead: The maximum additional CPU memory overhead is approximately 10.5–11GB, accounting for 12%–16% of the total model size, which is within an acceptable range.

Conclusion

Experiments show that MoE-APEX is a flexible and efficient inference system for deploying MoE LLM models on memory-constrained edge devices. By addressing the high cost associated with expert loading in existing MoE inference systems, MoE-APEX enables significant speedups in inference performance.

Overall, MoE-APEX represents a significant step forward in making complex MoE models more accessible and practical for edge computing environments.