FloE: On-the-Fly MoE Inference on Memory-constrained GPU

TL;DR

FloE addresses the memory bottleneck of MoE inference on consumer GPUs by enabling on-the-fly computation with a hybrid compression pipeline that combines contextual sparsity for gate/down projections and ultra-low-bit quantization for up projections. Two sparsity predictors—an inter-expert learned predictor and an intra-expert reuse-based predictor—pre-fetch compressed weights to enable pipelined transfer and computation, while a Triton-based sparse kernel and compact asynchronous transfer maximize data throughput. Empirically, FloE achieves a 9.3× compression per expert, enables deployment on GPUs with as little as 11 GB VRAM, and delivers up to 48.7× end-to-end acceleration over DeepSpeed-MII on a single GeForce RTX 3090 with only 4.4–7.6% average degradation. These results demonstrate FloE’s potential to make on-the-fly MoE inference practical on memory-constrained devices, narrowing the gap to fully GPU-resident baselines while significantly reducing memory footprint by up to 8.5×.

Abstract

With the widespread adoption of Mixture-of-Experts (MoE) models, there is a growing demand for efficient inference on memory-constrained devices. While offloading expert parameters to CPU memory and loading activated experts on demand has emerged as a potential solution, the large size of activated experts overburdens the limited PCIe bandwidth, hindering the effectiveness in latency-sensitive scenarios. To mitigate this, we propose FloE, an on-the-fly MoE inference system on memory-constrained GPUs. FloE is built on the insight that there exists substantial untapped redundancy within sparsely activated experts. It employs various compression techniques on the expert's internal parameter matrices to reduce the data movement load, combined with low-cost sparse prediction, achieving perceptible inference acceleration in wall-clock time on resource-constrained devices. Empirically, FloE achieves a 9.3x compression of parameters per expert in Mixtral-8x7B; enables deployment on a GPU with only 11GB VRAM, reducing the memory footprint by up to 8.5x; and delivers a 48.7x inference speedup compared to DeepSpeed-MII on a single GeForce RTX 3090 - all with only a 4.4$\%$ - 7.6$\%$ average performance degradation.

Figures (12)

• Figure 1: Comparison of MoE inference offloading systems for memory-constrained GPUs: (a) Naive MoE Inference with Offloading; (b) Advanced MoE Inference with Offloading; (c) On-the-Fly MoE Inference (FloE).
• Figure 2: Activation distributions of Mixtral-8$\times$7B’s three hidden states at experts $\mathcal{E}_{0,0}$ (shallow layer), $\mathcal{E}_{15,7}$ (middle layer), and $\mathcal{E}_{35,4}$ (deep layer).
• Figure 3: Compression sensitivity of expert parameters: (a) Sparsification sensitivity; (b) Quantization sensitivity.
• Figure 4: Next layer cosine similarity (blue), intra-expert predictor precision (yellow), inter-expert predictor recall (red), and the outlier corresponding to the cosine similarity at the 0-th layer (gray).
• Figure 5: Process of FloE's compact asynchronous transfer: compacting weights layout in DRAM for reduced access latency and multi-threaded packaging of activated experts to enable asynchronous data transfer.

Theorems & Definitions (20)

• Theorem 3.1: informal
• Theorem A.1: informal
• Theorem A.2
• Remark A.3
• Lemma A.4
• proof
• Lemma A.5
• proof
• Remark A.6
• Lemma A.7