Alloc-MoE: Budget-Aware Expert Activation Allocation for Efficient Mixture-of-Experts Inference

TL;DR

Alloc-MoE introduces a unified framework for optimizing expert activation distribution in Mixture-of-Experts models, significantly reducing inference latency while preserving performance, especially in resource-limited settings.

Contribution

It proposes a novel activation budget concept and a coordinated optimization method at layer and token levels to improve efficiency without performance loss.

Findings

Achieves 1.15x prefill speedup on DeepSeek-V2-Lite.

Achieves 1.34x decode speedup with maintained performance.

Maintains model accuracy under constrained expert activation budgets.

Abstract

Mixture-of-Experts (MoE) has become a dominant architecture for scaling large language models due to their sparse activation mechanism. However, the substantial number of expert activations creates a critical latency bottleneck during inference, especially in resource-constrained deployment scenarios. Existing approaches that reduce expert activations potentially lead to severe model performance degradation. In this work, we introduce the concept of \emph{activation budget} as a constraint on the number of expert activations and propose Alloc-MoE, a unified framework that optimizes budget allocation coordinately at both the layer and token levels to minimize performance degradation. At the layer level, we introduce Alloc-L, which leverages sensitivity profiling and dynamic programming to determine the optimal allocation of expert activations across layers. At the token level, we propose Alloc-T, which dynamically redistributes activations based on routing scores, optimizing budget allocation without increasing latency. Extensive experiments across multiple MoE models demonstrate that Alloc-MoE maintains model performance under a constrained activation budget. Especially, Alloc-MoE achieves $1.15\times$ prefill and $1.34\times$ decode speedups on DeepSeek-V2-Lite at half of the original budget.