Accelerating Mixture-of-Experts language model inference via plug-and-play lookahead gate on a single GPU

Abstract

The widespread adoption of large language models (LLMs) has encouraged researchers to explore strategies for running these models more efficiently, such as the mixture of experts (MoE) method, which aims to increase the knowledge capacity of the model without substantially increasing its computational costs, as only a fraction of the model components are active for each token. However, this approach also increases the size of the model, which makes it challenging to run these models even on high-end GPUs. Quantization and offloading strategies have been used to enable the execution of MoE in resource-constrained environments, however, the time overhead introduced by offloading remains a bottleneck. In this paper, we propose a plug-and-play lookahead gate that predicts in advance the experts to be used in the next few layers. Furthermore, to mitigate the misalignment problem arising from cross-layer prediction, we introduce an alignment training method, layer-wise gate alignment, enhancing the prediction hit rate while maintaining low resource requirements. Moreover, we present a speculative expert scheduling strategy to accelerate the end-to-end inference process of MoE models. To validate our approach, we established an inference framework for quantized MoE and conducted extensive experiments. The results demonstrate the effectiveness of our proposed methods, with throughput improvements of 57.72%, 60.00%, and 62.26% under 4, 3, and 2-bit quantization conditions for experts, respectively, compared with the Mixtral-offloading method.