MatGPTQ: Accurate and Efficient Post-Training Matryoshka Quantization

TL;DR

This work tackles the challenge of deploying large language models across diverse memory and latency budgets by enabling a single checkpoint to serve multiple bit-widths through Matryoshka quantization. By reframing MatQuant as a multi-precision GPTQ problem with MSB slicing and cross-width error sharing, MatGPTQ produces a jointly optimized master model for a user-defined set of widths $R$ (e.g., $\{2,3,4,6,8\}$) at master width $c$, computed in one shot from a small calibration set. The authors introduce efficient CUDA kernels, a budget-aware non-uniform per-layer width search via EvoPress (Mix-and-Match), and demonstrate near-GPTQ accuracy at high bit-widths with substantial gains at 3–4 bits, along with strong interpolation to intermediate widths. The practical impact is a open-source, end-to-end pipeline for sliceable, multi-precision deployment that significantly reduces memory and latency while preserving accuracy, enabling flexible serving on heterogeneous hardware. This combination of a scalable PTQ objective, fast kernels, and an autonomous width-search strategy establishes MatGPTQ as a new standard for Matryoshka-style post-training quantization.

Abstract

Matryoshka Quantization (MatQuant) is a recent quantization approach showing that a single integer-quantized model can be served across multiple precisions, by slicing the most significant bits (MSB) at inference time. This enables a single checkpoint to cover a wide range of memory and latency budgets, but renders quantization much more challenging. In particular, the initial MatQuant relies on expensive quantization-aware training (QAT) variants, rather than fast one-shot post training quantization (PTQ), and lacks open-source and kernel support. We address all of these limitations by introducing Post-Training Matryoshka Quantization (MatGPTQ), a new PTQ pipeline that produces a single parent model jointly optimized for multiple target precisions in one-shot, based on a small calibration set. MatGPTQ casts Matryoshka quantization as a multi-precision objective with bit-slicing and cross-bit error compensation, resulting in an algorithm that produces a multi-bit-width, "sliceable" model in a single pass. We also incorporate a new budget-aware search for heterogeneous per-layer bit-witdhs and provide efficient kernels that implement slicing and mixed-precision execution. Across standard LLMs and benchmarks, MatGPTQ preserves high-bit accuracy while substantially improving performance at low-bit-witdh settings. Overall, we establish a new state of the art for Matryoshka-style post-training quantization and make single-checkpoint, multi-precision deployment open and practical. Code is available at https://github.com/IST-DASLab/MatGPTQ.

Figures (4)

• Figure 1: Task average performance of Llama3.1-8B-Instruct quantized using MatGPTQ optimizing for 3, 4 and 8 bits. The results demonstrate that even under extreme compression, MatGPTQ maintains strong performance, with models between 2.5 and 4 bits achieving near-baseline accuracy for optimized bit-witdhs and even exceeding it for interpolated bit-witdhs.
• Figure 2: Heatmap of best-performing quantization configurations across transformer blocks of Llama-3.1-8B Instruct. We enumerated all configurations for bitwidths 2, 3 and 4 under an average bitwidth of 3 and measured per-token MSE against the FP16 baseline. Configurations without a single token are omitted. A single static configuration dominates across most blocks, with certain layers showing higher sensitivity.
• Figure 3: Benchmark results comparison across muliple matrix dimensions with torch.matmul as the baseline. Our kernels achieve speedups between 3x and 5.6x in the memory-bound regime.
• Figure 4: End-to-end speedups in vLLM vllm for Llama-3.1-8B-Instruct. Our kernels achieve speedups between 1.5x and 3.5x in the memory-bound regime. We evaluated on an RTX A6000 using Prompt/Decode 32/128.