AlignQuant: Efficient Mixed-Precision Quantization for LLMs

Paper: AlignQuant: Tile-Aligned Mixed-Precision Quantization for Efficient LLM Generation

Authors: Hanzhi Zhang, Qiao Zhang, Qinglei Cao, Heng Fan, Yan Huang, Kewei Sha, Yunhe Feng

Executive Summary

Recent advancements in artificial intelligence have highlighted the potential for fine-grained mixed-precision quantization to enhance the efficiency of large language model (LLM) inference. However, a challenge arises from the mismatch between local precision choices and the requirements of standard GPU storage and computation units. This misalignment can hinder the practical benefits of compressing models to accelerate their performance during inference.

To address this issue, researchers have developed AlignQuant, a post-training quantization method designed to optimize the use of GPU resources. AlignQuant employs two-dimensional weight tiles, which serve as a unified unit for precision allocation, storage, and execution. This approach allows precision to be allocated according to the sensitivity of various output channels, ensuring that the most critical components of the model receive the appropriate level of precision.

The method involves a calibration process that assesses precision reductions by evaluating the impact of projection-output perturbations, which are influenced by language model loss gradients under quantized activations. This process is further refined through phase-normalized scoring, which emphasizes higher precision for tiles that are significant during either the pre-fill or decoding phases of model execution, all while adhering to a model-wide weight-storage budget.

In practical terms, AlignQuant enables each tile to store a single selected representation. Additionally, it incorporates phase-specialized kernels that efficiently reuse the packed model, allowing for the expansion of lower-bit weights for INT8 computation with 8-bit activations. The research demonstrates that AlignQuant achieves up to 2.5 times faster generation speed compared to the BF16 format, while maintaining the quality of the model.

The evaluations conducted cover four different language models with parameter sizes ranging from 3 billion to 14 billion. These tests were performed across three different GPUs and with context lengths of up to 64,000 tokens. The results indicate that AlignQuant effectively reconciles the flexibility of local precision with the requirements of standard GPU execution, offering a viable solution for improving inference efficiency in large language models.

The implementation of AlignQuant is accessible to the public, allowing organizations interested in optimizing their AI models to explore its potential benefits. This development could pave the way for more efficient deployment of AI solutions across various applications, from natural language processing to other areas where large language models are utilized. The availability of the implementation may encourage further advancements and explorations in mixed-precision quantization techniques, ultimately enhancing the speed and performance of AI systems in real-world scenarios.

Academic Abstract

Fine-grained mixed-precision quantization promises efficient large language model inference, but local precision choices can conflict with regular GPU storage and computation units. This precision-boundary mismatch limits the translation of compression into practical acceleration. We introduce AlignQuant, a post-training quantization method that uses GPU-compatible two-dimensional weight tiles as the common unit of precision allocation, compact storage, and execution. This shared partition lets precision follow sensitivity within output channels. Joint prefill/decode calibration scores precision reductions using projection-output perturbations weighted by language-model loss gradients under quantized activations. Phase-normalized scores prioritize higher precision for tiles important to either phase under a model-wide weight-storage budget. Each tile stores one selected representation, while phase-specialized kernels reuse the packed model and expand lower-bit weights for INT8 computation with 8-bit activations. Across four LLMs spanning 3B to 14B parameters, AlignQuant achieves up to $2.50\times$ generation speedup over BF16 while preserving model quality. Evaluations further cover three GPUs and contexts up to 64K tokens. These results show that local precision flexibility and regular GPU execution can coexist through a shared tile unit. The implementation is available at https://github.com/HanzhiZhang-Ulrica/AlignQuant.