MiLo: Efficient Quantized MoE Inference with Mixture of Low-Rank Compensators

MiLo: Efficient Quantized MoE Inference with Mixture of Low-Rank Compensators

Speakers: Beichen Huang, Yueming Yuan, Zelei Shao, Minjia Zhang

Conference: MLSys 2025

YouTube: https://www.youtube.com/watch?v=None

Overview

This article delves into MiLo, a novel approach for efficient quantized Mixture of Experts (MoE) inference, presented by Beichen Huang and Yueming Yuan at MLSys 2025. MiLo tackles the pressing challenge of deploying increasingly large MoE models—which are becoming prevalent due to their ability to scale model capacity while maintaining computational efficiency—on resource-constrained hardware, specifically single GPUs. The core innovation lies in its ability to perform extreme quantization, reducing model weights from FP16 to int3, without incurring severe accuracy loss or requiring extensive calibration data.

The talk introduces both the MiLo algorithm and the MiLo kernel. The algorithm leverages a mixture of low-rank compensators to mitigate the accuracy degradation typically associated with aggressive int3 quantization, particularly by focusing on the 'insignificant' weights where quantization loss is most pronounced. Complementing this, the MiLo kernel provides a custom, highly optimized int3 GEMM (General Matrix Multiply) backend, designed to overcome the inherent hardware limitations and unlock significant inference speedups. This work is crucial for enabling the widespread adoption and deployment of state-of-the-art MoE models in real-world applications where memory and computational budgets are tight.

Background

▶ Watch: Introduction to MiLo and the MoE scaling problem (0:00)

The landscape of large language models (LLMs) is rapidly evolving, with models growing exponentially in size and complexity. Mixture of Experts (MoE) architectures have emerged as a powerful paradigm for scaling up these models efficiently. Unlike traditional dense models where all parameters are activated for every input, MoE models replace standard feed-forward networks with a group of specialized "experts" and dynamically activate only a subset of them per input token. This allows MoE models to achieve a much larger parameter count, enhancing model capacity and performance, while keeping the computational cost per token roughly comparable to smaller dense models. Examples like the Mixture 8 * 7B model demonstrate the power of this approach.

However, the sheer size of these models, even with sparse activation, presents a significant deployment challenge. As highlighted in the talk, fitting a state-of-the-art MoE model onto a single GPU is often "almost impossible" due to memory limitations. This necessitates extreme model compression techniques, among which quantization stands out as a promising method. Quantization maps high-precision data (e.g., FP16, 16 bits per element) to lower-resolution formats (e.g., int4, 4 bits; or int3, 3 bits), thereby shrinking model size and potentially reducing computational requirements. The talk focuses specifically on the challenging int3 quantization regime.

Existing quantization methods can broadly be categorized into two groups:

1. Calibration-based methods: These require a calibration dataset to perform forward propagation and analyze weight distributions or activations. Examples include GPTQ, which calculates the Hessian matrix, and AWQ, which identifies salient weights based on activation statistics. While often achieving good quality, these methods suffer from significant drawbacks for large MoE models:

• Long execution time: Quantizing a Mixture 8 * 7B model with GPTQ can take over 5,000 seconds due to the need for extensive forward passes on calibration data.
• Calibration data bias: There's a risk of overfitting to the calibration dataset, potentially leading to suboptimal generalization.

1. Calibration-free methods: These methods do not require external data. The simplest is Round to Nearest (RTN). A more advanced option is Half Quadratic Quantization (HQQ), which frames quantization as an optimization problem solved using a half-quadratic solver. These methods are faster but can sometimes yield lower accuracy.

The primary challenges when applying quantization to MoE models, especially in the extreme int3 setting, are two-fold:

1. Severe accuracy loss with extreme quantization: While int4 quantization often retains model performance, pushing to int3 can lead to substantial degradation. For instance, GPTQ applied to a Mixture 8 * 7B model at int3 results in a 34% drop in performance (perplexity). This level of accuracy loss makes int3 quantization impractical for many real-world applications.
2. Calibration data bias and long execution time: As mentioned, calibration-based methods are prohibitively slow and carry the risk of overfitting, making them unsuitable for rapid experimentation or deployment of large MoE models.

MiLo's objective is to overcome these hurdles, enabling efficient int3 quantization for MoE models without requiring calibration data, while significantly mitigating the associated accuracy drop. The foundational insight driving MiLo is that the most significant quantization loss in int3 quantization primarily occurs in the "insignificant weights," while methods tend to capture outliers effectively. This observation forms the basis for introducing low-rank compensators to selectively restore precision where it matters most.

Key Findings

▶ Watch: MiLo's insight: low-rank compensators for quantization loss (4:40)

The MiLo project introduces a comprehensive solution for efficient int3 quantization of Mixture of Experts (MoE) models, addressing both accuracy degradation and inference speed limitations. The key findings and contributions can be summarized as follows:

1. MiLo Algorithm for Accuracy Preservation: MiLo successfully mitigates the severe accuracy loss typically observed with extreme int3 quantization in MoE models. It achieves this by introducing low-rank compensators that are added to the quantized weights. This approach specifically targets the "insignificant weights" where the majority of quantization errors occur, effectively restoring precision without significantly increasing the model size. The algorithm is calibration-free, relying on HQQ for initial quantization and Singular Value Decomposition (SVD) for compensator generation, making it fast and robust against data bias.

1. Adaptive Rank Selection Strategy: A crucial innovation within the MiLo algorithm is its adaptive rank selection strategy for the low-rank compensators. Instead of using a uniform rank, MiLo dynamically assigns different ranks to various parts of the model weights based on three key observations:

• Model Structure: Dense layers (e.g., attention weights) receive higher ranks due to their heavy-tailed distributions and constant activation.
• Expert Activation Frequency: More frequently activated experts are assigned higher ranks.
• Kurtosis Value: Weights with higher kurtosis values (indicating heavier-tailed distributions and thus higher quantization error) are compensated with higher ranks. This tailored approach significantly improves the quality of compensation and, consequently, the overall quantization accuracy.

1. MiLo Kernel for Accelerated Inference: MiLo introduces an efficient custom int3 quantized GEMM kernel designed to unlock real-world inference speedups. This kernel overcomes the fundamental challenge of modern hardware lacking native int3 support. It features an innovative bit-packing method that stores 32 three-bit values in three int32 vectors without bit wasting. The kernel also incorporates advanced optimizations like asynchronous global memory loads and MoE-specific tail shift tuning to maximize GPU utilization and reduce synchronization overhead.

1. Significant Performance Gains:

• Accuracy: The MiLo algorithm, especially with its adaptive rank strategies, consistently outperforms other quantization methods across a range of benchmarks for MoE models, demonstrating superior accuracy retention at int3.
• Speed: The MiLo kernel achieves "multiple times faster" GEMM operations compared to unfused native implementations and reports "higher TFLOPS than state-of-the-art kernels" designed for W4A16 quantization.
• End-to-End Latency Reduction: On the challenging Mixture 8 * 7B model, MiLo achieves a notable 20% reduction in end-to-end inference latency, making large MoE models practical for single-GPU deployment. This gain is partly attributed to fusing asymmetric dequantization with the zero point.

In essence, MiLo provides a holistic solution that enables the deployment of massive MoE models (like Mixture 8 * 7B) on a single GPU by making extreme int3 quantization both accurate and fast, without the typical drawbacks of calibration data.

Technical Deep Dive

▶ Watch: MiLo algorithm: HQQ, SVD compensation, and iterative optimization (5:40)

The technical prowess of MiLo lies in its dual-pronged approach: an innovative algorithm for accurate extreme quantization and a highly optimized kernel for efficient inference.

MiLo Algorithm: Mixture of Low-Rank Compensators

The MiLo algorithm is built upon the insight that extreme int3 quantization often causes significant accuracy loss primarily in the "insignificant weights," while effectively capturing the "outliers." To counteract this, MiLo introduces low-rank compensators which are essentially small matrices added to the quantized weights to "make up" for the lost precision.

The output weight matrix W_out is defined as the sum of the quantized weight matrix W_quant and a low-rank compensator U * V. Here, U is an M x R matrix and V is an R x N matrix, where R is the rank of the compensator. A higher rank R generally leads to better performance but with a slight increase in storage.

The MiLo algorithm proceeds in three main steps, with an iterative refinement process:

1. Initial Quantization: The original FP16 model weights are initially quantized using HQQ (Half Quadratic Quantization). HQQ is chosen because it is a calibration-free method, framing the quantization problem as an optimization problem. This ensures fast and efficient initial quantization without the need for a calibration dataset, addressing one of the key challenges of MoE quantization.

1. Compensation Generation: After initial quantization, a residual matrix is calculated by subtracting the quantized weights from the original high-precision weights (W_residual = W_original - W_quant). This residual matrix captures the quantization error. To generate the low-rank compensator, Singular Value Decomposition (SVD) is applied to this residual matrix. SVD decomposes the residual into U S V^T, where S contains singular values. By taking the top R singular values and corresponding singular vectors, the low-rank matrices U and V are derived, effectively approximating the residual with minimal additional parameters.

1. Mixture of Low-Rank Compensators and Adaptive Rank Selection: A key innovation is that the rank R is not fixed but is adaptively selected for different parts of the model. This creates a "mixture" of low-rank compensators. MiLo proposes three strategies for adaptive rank selection, based on distinct observations about MoE models:

• Strategy 1: Model Structure (Dense vs. Sparse Layers):
• Observation: Attention layers (dense layers) have a more heavy-tailed weight distribution compared to expert layers (sparse layers) and are always activated. Heavy-tailed distributions are known to increase quantization difficulty.
• Implementation: Assign a higher rank R to dense layers (e.g., attention weights) and a lower rank to sparse expert layers. This ensures that the most influential and difficult-to-quantize parts of the model receive more robust compensation.

• Strategy 2: Expert Activation Frequency:
• Observation: Not all experts in an MoE model are equally activated. Some experts are invoked much more frequently than others, especially in "deep MoE" models with fine-grained experts.
• Implementation: Assign a higher rank R to experts that are activated at a higher frequency. This prioritizes compensation for the experts that contribute most to the model's overall output, maximizing the impact of the limited compensation budget.

• Strategy 3: Kurtosis Value:
• Observation: The kurtosis value, a statistical measure of the "tailedness" of a distribution, indicates the potential for quantization error. A higher kurtosis value signifies a heavier-tailed distribution, which correlates with higher quantization error.
• Implementation: Assign a higher rank R to weights (or weight matrices) that exhibit a higher kurtosis value. This directly targets the weights that are inherently more challenging to quantize accurately.

The quantization and compensation steps are then iteratively optimized, meaning W_quant is refined based on U*V, and U*V is re-calculated based on the new residual. This iterative process leads to better overall quantization results. Finally, the low-rank compensators U and V themselves are quantized using simple asymmetric linear quantization to int3. This step is found to have negligible computation overhead and an almost negligible performance drop, ensuring that the compensators do not significantly increase the model size or computational burden.

MiLo Kernel: Efficient Int3 Quantized GEMM Backend

The MiLo kernel is designed to provide an efficient int3 x FP16 mixed-precision GEMM backend, crucial for realizing actual speedups. Native weight-quantized GEMM often involves two separate steps: dequantization and FP16 GEMM, which can be slower than unquantized FP16 GEMM due to memory access overhead. The goal is to fuse these operations into a single kernel for ideal speedup. While W4A16 (int4 weights, FP16 activations) kernels have achieved near-ideal speedups, efficient int3 kernels have been lacking.

The MiLo kernel addresses two main difficulties for int3 quantization:

1. Lack of Direct Int3 Hardware Support: Modern GPUs do not directly support int3 data storage or operations because 3 is not a power of two. This makes packing int3 values into standard 32-bit registers without wasting bits a significant challenge.

• Solution: Bit Packing: MiLo employs a clever bit-packing scheme. It uses three int32 vectors to pack 32 three-bit values. The high four bits of each of these three vectors form a cyclic buffer, which jointly stores four int3 values. This method ensures no bit wasting.
• Dequantization: During dequantization, the kernel sequentially applies dequantization operations to the three int32 vectors. A fourth vector is then constructed by concatenating the cyclic buffers. This process utilizes efficient binary operators and register-level parallelism to extract two int3 values simultaneously, minimizing overhead.

1. Efficient GPU Resource Utilization: It's challenging to fully utilize hierarchical GPU memory and other resources in a custom three-bit kernel.

• Optimization 1: Asynchronous Global Load: To hide memory latency, the kernel uses asynchronous global memory loads. This allows the expensive global memory access operations to overlap with computation, keeping the compute units busy.
• Optimization 2: MoE-Specific Tail Shift Tuning: An MLP layer often consists of multiple linear layers with different "shifts" (referring to the zero point in asymmetric quantization). For MoE models, these shifts can vary across experts. MiLo tunes the tail shift for each linear layer across different MoE models. This optimization specifically aims to reduce thread block synchronization within the GEMM kernel, which is a common bottleneck in highly parallel GPU computations, thereby improving overall throughput.

By meticulously addressing these architectural and algorithmic challenges, the MiLo kernel provides a robust and highly performant backend for int3 quantized MoE inference.

Experimental Setup & Results

▶ Watch: Adaptive rank strategy 2: Higher rank for active experts (7:50)

The evaluation of MiLo focused on demonstrating both the algorithm's ability to preserve accuracy under extreme int3 quantization and the kernel's ability to deliver significant inference speedups.

Models and Quantization Target:

The primary model used for end-to-end evaluation was the **Mixture 8 * 7B model, a large-scale MoE architecture. The target quantization precision was int3 for weights (W3), with activations remaining at FP16 (A16)**, making it a W3A16 mixed-precision scheme. This choice of int3 was deliberate, aiming to push the limits of quantization where accuracy degradation is most severe and where compensators offer the most value.

Baselines for Comparison:

• Accuracy Baseline: For algorithmic comparisons, GPTQ was used as a representative calibration-based quantization method. The talk highlighted GPTQ's significant perplexity drop (34%) when quantizing the Mixture 8 * 7B model to int3.
• Kernel Baselines: The MiLo kernel was compared against:
• Unfused native implementation: This represents a naive approach where dequantization and GEMM are separate steps, often leading to performance degradation compared to FP16.
• State-of-the-art W4A16 kernels: These are highly optimized kernels for int4 weight quantization, serving as a strong benchmark for throughput.
• End-to-End Baseline: For end-to-end latency measurements, MiLo was compared against a "MoE" baseline, presumably an unquantized or less optimized quantized version of the Mixture 8 * 7B model.

Metrics:

• Accuracy: Measured by perplexity (lower is better) for language models.
• Performance: Measured by TFLOPS (higher is better) for GEMM operations and end-to-end inference latency reduction (lower is better).

Headline Results:

1. Algorithmic Accuracy:

• The talk presented a graph showing that MiLo, particularly with its proposed adaptive rank strategies (and combinations thereof), consistently outperforms other methods in a range of benchmarks in terms of accuracy retention at int3. This directly addresses the 34% perplexity drop observed with GPTQ at int3. While specific perplexity numbers for MiLo were not explicitly stated, the visual evidence and qualitative claim strongly indicate superior accuracy.
• The effectiveness of the adaptive rank selection strategy was highlighted, demonstrating that tailoring compensation ranks based on model structure, expert activation frequency, and kurtosis values leads to better quantization quality.

1. MiLo Kernel Performance:

• GEMM Speedup: For all tested batch sizes, the MiLo kernel was shown to be "multiple times faster than the unfused native implementation baseline." This confirms the benefit of fused, custom int3 kernel design.
• TFLOPS: The MiLo kernel also reported "higher TFLOPS than the state-of-the-art kernel" for W4A16 quantization. This is a significant achievement, as int3 quantization is inherently more complex due to hardware limitations, yet MiLo manages to surpass the performance of optimized int4 kernels.

1. End-to-End Latency Reduction:

• On the Mixture 8 * 7B model, MiLo achieved a substantial 20% reduction in end-to-end inference latency compared to the baseline. This demonstrates the combined impact of the efficient algorithm and the optimized kernel. The presenters noted extra gains from fusing asymmetric dequantization with the zero point in the end-to-end evaluation.

Experimental Setup Details:

• Hardware: The measurements were primarily conducted on single-node machines (single GPU). The speakers explicitly stated that their major goal was to deploy models with limited memory on a single device, and distributed settings were not yet tested.
• Compensator Quantization: The low-rank compensator matrices (U and V) were also quantized to int3. Experiments showed that this resulted in "negligible quantization loss" and "negligible time consuming" due to the use of simple asymmetric linear quantization.

In summary, MiLo's experimental results strongly validate its effectiveness. It provides a robust solution for deploying large MoE models on single GPUs by enabling accurate and fast int3 quantization, overcoming significant challenges faced by prior methods.

Practical Implications

▶ Watch: Adaptive rank strategy 3: Kurtosis value indicates error (8:20)

MiLo presents profound practical implications for the deployment and operationalization of large Mixture of Experts (MoE) models, particularly in environments with constrained computational and memory resources.

Enabling Single-GPU Deployment of Large MoE Models: The most direct and impactful implication is that MiLo makes it feasible to deploy massive MoE models, such as the Mixture 8 * 7B, on a single GPU. Prior to MiLo, such deployments were "almost impossible" due to the prohibitive memory footprint of these models. By aggressively quantizing weights to int3, MiLo drastically reduces the model's memory footprint, opening up new possibilities for on-device, edge, or cost-effective cloud inference.

Viability of Extreme Quantization (Int3): MiLo elevates int3 quantization from a technique often associated with severe accuracy degradation (e.g., 34% perplexity drop with GPTQ) to a practical and viable option. By introducing low-rank compensators and adaptive strategies, MiLo ensures that this extreme level of compression can be achieved without compromising model performance to an unacceptable degree. This pushes the boundaries of model compression, allowing for even smaller model sizes and potentially higher throughput.

Elimination of Calibration Data Dependencies: For practitioners and infrastructure teams, the calibration-free nature of MiLo is a significant advantage. Calibration-based methods, like GPTQ, are notoriously slow for large MoE models (over 5,000 seconds for Mixture 8 * 7B), and carry the risk of calibration data bias or overfitting. MiLo bypasses these issues entirely, streamlining the quantization pipeline, reducing development cycles, and enhancing the robustness of quantized models. This means faster experimentation and more reliable deployments.

Significant Performance and Efficiency Gains:

• Reduced Latency: The 20% end-to-end latency reduction on a Mixture 8 * 7B model directly translates to faster response times for applications, improving user experience and enabling real-time use cases.
• Higher Throughput: The MiLo kernel's ability to achieve higher TFLOPS than state-of-the-art int4 kernels indicates that not only is memory reduced, but the actual computation is performed more efficiently. This means more queries can be processed per unit of time, optimizing hardware utilization.
• Cost Savings: Deploying larger models on fewer or less powerful GPUs can lead to substantial cost savings in cloud infrastructure or even enable entirely new on-device applications that were previously out of reach.

Tradeoffs and Limitations:

1. Focus on Single-Device Scenario: The current work primarily focuses on single-device deployment scenarios. While the memory reduction benefits would theoretically extend to distributed settings, the end-to-end performance gains and specific optimizations for multi-node inference have not yet been tested. This suggests that further research might be needed for optimal performance in distributed MoE setups.

1. Activations Remain FP16: MiLo currently quantizes only weights to int3, keeping activations at FP16. While this simplifies the problem and allows for maximum focus on the challenging int3 weight quantization, it leaves potential for further memory and compute savings if activations could also be quantized (e.g., to int8). The speakers clarified that int8 quantization for activations is already "good enough" with less room for compensator benefit, and their focus was on the "extreme" int3 weight quantization where compensators provide a much larger impact.

1. Minor Memory Cost for Compensators: While the compensator matrices (U and V) are low-rank and themselves quantized to int3, they do introduce a minor additional memory cost. However, this cost is explicitly justified by the significant accuracy improvements for int3, where the performance degradation is otherwise severe. The value proposition is clear: a small memory overhead for a large accuracy gain.

In summary, MiLo is a pivotal development for the practical deployment of large MoE models, offering a robust, efficient, and calibration-free method for extreme int3 quantization that significantly reduces memory footprint and boosts inference speed on single GPUs.

Key Takeaways

• Extreme Quantization for MoE Models: MiLo enables practical int3 quantization for large Mixture of Experts (MoE) models, which are otherwise challenging to fit on single GPUs due to their massive parameter counts.
• Low-Rank Compensators Mitigate Accuracy Loss: The core of the MiLo algorithm is the use of mixture of low-rank compensators that are added to quantized weights to restore precision, particularly in "insignificant weights" where most quantization errors occur.
• Adaptive Rank Selection Enhances Compensation: MiLo employs an adaptive rank selection strategy for compensators, dynamically assigning ranks based on layer type (dense vs. sparse), expert activation frequency, and kurtosis values, leading to superior accuracy.
• Calibration-Free and Fast Algorithm: MiLo is a calibration-free method, utilizing HQQ for initial quantization, which significantly reduces execution time (avoiding thousands of seconds for calibration) and eliminates the risk of calibration data bias.
• Optimized Int3 GEMM Kernel: The MiLo kernel provides a custom, highly efficient int3 quantized GEMM backend that overcomes hardware limitations (no native int3 support) through innovative bit-packing and MoE-specific optimizations, achieving higher TFLOPS than state-of-the-art int4 kernels.
• Significant End-to-End Performance: MiLo delivers a substantial 20% reduction in end-to-end inference latency for models like Mixture 8 * 7B, making large MoE models deployable and performant on single GPUs.

About the Speaker(s)

Beichen Huang and Yueming Yuan presented the MiLo work at MLSys 2025. They are part of the research team behind this innovative approach to efficient MoE quantization, alongside Zelei Shao and Minjia Zhang. The talk highlighted their contributions to both the MiLo algorithm and its optimized kernel, demonstrating a deep understanding of both the theoretical challenges of extreme quantization and the practicalities of high-performance system design for AI/ML inference.