MxMoE: Mixed-precision Quantization for MoE with Accuracy and Performance Co-Design

TL;DR

This work addresses the deployment challenges of large Mixtr-Experts (MoE) models by introducing MxMoE, a framework for hardware-aware, mixed-precision quantization that accounts for both algorithmic sensitivity and dynamic expert activation. It discovers that quantization sensitivity varies across linear blocks and experts, motivating a linear-block level bitwidth allocation and a hardware-aware optimization that jointly considers accuracy and execution time. The approach combines an ILP-based bitwidth allocation with automatic generation of mixed-precision Group-GEMM kernels and a tile scheduler to achieve substantial speedups while preserving accuracy, outperforming GPTQ and uniform-precision baselines across multiple MoE models and settings. Overall, MxMoE demonstrates practical, scalable gains for deploying MoE-based LLMs on GPUs, enabling lower memory footprints and faster inference through principled co-design and automated kernel generation.

Abstract

Mixture-of-Experts (MoE) models face deployment challenges due to their large parameter counts and computational demands. We explore quantization for MoE models and highlight two key insights: 1) linear blocks exhibit varying quantization sensitivity, and 2) divergent expert activation frequencies create heterogeneous computational characteristics. Based on these observations, we introduce MxMoE, a mixed-precision optimization framework for MoE models that considers both algorithmic and system perspectives. MxMoE navigates the design space defined by parameter sensitivity, expert activation dynamics, and hardware resources to derive efficient mixed-precision configurations. Additionally, MxMoE automatically generates optimized mixed-precision GroupGEMM kernels, enabling parallel execution of GEMMs with different precisions. Evaluations show that MxMoE outperforms existing methods, achieving 2.4 lower Wikitext-2 perplexity than GPTQ at 2.25-bit and delivering up to 3.4x speedup over full precision, as well as up to 29.4% speedup over uniform quantization at equivalent accuracy with 5-bit weight-activation quantization. Our code is available at https://github.com/cat538/MxMoE.

Figures (6)

• Figure 1: (a) Quantization loss across experts in DeepSeekV2-Lite's 11th layer under various quantization schemes (top), and across linear components (Gate_proj/Down_proj) under the w4a4_g-1_sym configuration (bottom). The quantization notation wxay_gz_b denotes x-bit weights, y-bit activations, group size z (-1 indicates per-channel/token) with symmetric (sym) or asymmetric (asym) quantization. The quantization loss metric is formally defined in Section \ref{['ssec:quant-loss']}. (b) Roofline performance analysis for RTX 4090 GPU (left) and expert activation frequency distribution in DeepSeekV2-Lite's 11th layer (right).
• Figure 2: Comparison of the computation throughput of low-precision MoE block. W4 denotes 4-bit per-channel symmetric weight-only quantization, while W8A8 refers to 8-bit per-channel symmetric weight-activation quantization. The problem consists of 60 experts, each with a shape of $[N,K] = [2816,2048]$ (from Qwen2_MoE1.5), with each token activating 4 experts. The total number of input tokens is set to 512.
• Figure 3: Overview of MxMoE.
• Figure 4: MxMoE ensures that tile configurations for different precisions have the same number of warps.
• Figure 5: Computational throughput of MoE blocks across models and precision settings. W4A16 denotes 4-bit weight-only per-channel asymmetric quantization; W8A8 and W4A4 denote 8/4-bit weight-activation per-channel symmetric quantization. Number followed MxMoE represents the average bitwidth for weight and activation. The $\star$ symbols indicate corresponding perplexity values on WikiText2.