ABQ-LLM: Arbitrary-Bit Quantized Inference Acceleration for Large Language Models

where $W$ and $X$ are full-precision weight and activation, $Q(\cdot)$ denotes the quantizer of weight and activation, $clip(\cdot)$ denotes the clipping operation, $s$ is the scale factor, and letting $W_{max}=\alpha max(W)$ and $W_{min}=\beta min(W)$ to control the clipping range of the weight.

We observed significant variations in sensitivity across different layers of LLM models during quantization, with some layers having a critical impact on quantization performance. To validate this, as shown in Figure 1, we quantified various components of the LLaMA-7B model under weight-activation full quantization. While quantizing the gate_proj and up_proj layers in mlp and attention resulted in only minor performance degradation, quantizing the down_proj linear layer caused a substantial performance drop. This indicates that addressing down_proj quantization is crucial for performance recovery. Further analysis revealed that the primary cause of performance degradation due to down_proj quantization is the quantization of down_proj activation. At low bit-widths such as INT4, INT3, and INT2, the limited representation range causes a significant shift in the model distribution compared to full precision. As illustrated in Figure 3, during the block-wise quantization calibration process, we apply a double logarithm of cosine similarity loss on the output of down_proj to correct the distribution of the quantized model. The loss function called DLC loss $\mathcal{L}^{i}_{DLC}$:

where $d^{i}_{q}$ represent the quantized output of the $i$-th transformer block, $d^{i}_{fp}$ represent the full-precision output of the $i$-th transformer block, and $d^{i}_{fp^{*}}$ represent the full-precision output of the $i$-th transformer block, with its input originating from the quantized output of the $(i-1)$-th transformer block.

where $\lceil\cdot\rfloor$ denotes round operation, $n$ represents the target bit-width, $\bigtriangleup$ denotes the step-size, and $z$ is the zero-point. $W_{q}$ and $W$ denote the quantized and full-presicion weight, respectively. The vectors $a$ and $b$ are distribution compensation vectors, where $\gamma=1$ indicates compensation is performed, and $\gamma=0$ indicates no compensation.

To enhance the performance of the quantized model, we analyzed the changes in Attention Map distribution before and after quantization, as shown in Figure 2. In the full-precision model, attention is heavily focused on the first token, highlighting its key role in guiding text generation, consistent with LLM-QAT(Liu et al. 2023) findings. However, quantization disrupts this attention pattern, diminishing focus on the first token. To address this and restore the model’s attention during quantization, we introduced attention-aware KL divergence to reconstruct the attention map.

where $attn^{i}_{q}$ denotes the quantized attention map output of the $i$-th transformer block, while $attn^{i}_{fp}$ refers to the full-precision attention map output of the same block.

In the end, we combined DLC loss and AKL loss, and our final optimization goal is:

where $s^{i}_{*}$, $\alpha^{i}_{*}$ and $\beta^{*}_{i}$ are the parameters of the $i$-th transformer block after calibration. When the distributions of the quantized output and the full-precision output match, we have a loss close to 0, which effectively guides the quantization process.

Reconstructing Arbitrary Bit Computation. To support W1A1 quantization, NVIDIA introduced INT1 TensorCore in Turing and later architectures to provide hardware support. However, W1A1 quantization has not been widely applied due to significant performance degradation. Through mathematical analysis of quantized matrix multiplication, we find that any combination of quantization can be decomposed into a superposition of 1-bit matrix multiplications. Assuming the weight $W$ of a particular neural network layer are quantized to q bits and the input activation value $X$ is quantized to p bits, the matrix multiplication of $W$ and $X$ results in a 32-bit output $Y=WX$. The key is to observe that the scalar values at any position in W and X can be decomposed into a series of 1-bit scalar numbers. Scalar operations with any combination of precision can be decomposed into 1-bit operations and shift operations. For example, a 2-bit x can be expressed as:

where $x^{1}=(x\gg 1)\&1$, $x^{0}=(x\gg 0)\&1$. We use $OP(a,b)$ to denote a computational operation where the input is 1-bit data and the output is 32-bit. Thus, the original scalar-level arbitrary precision computation $wx$ can be represented as:

Engine Implementation. NVIDIA GPU has many processing elements called Streaming Multiprocessors (SMs) and uses a large number of threads to perform computing tasks in parallel. Threads are structured into thread blocks, which become the smallest scheduling execution unit on SMs. Therefore, the computation target is decomposed and mapped to each thread block, called Thread Block Tile, to achieve parallel computing. As shown in Figure 4(a), for a GEMM task of shape M×N×K, each thread block is responsible for computing a BM×BN output block, which is decomposed into $\frac{K}{BK}$ sub-GEMM tasks of shape BM×BN×BK. Our engine converts quantized matrix multiplications with bit widths configured as {p, q} into special accumulations of p*q binarized matrix multiplications, so the true calculation task of thread block tile is p*BM × q*BN. First, in order to improve memory access continuity, we propose BitPacking strategy to decompose the quantized tensor into $n$ binary matrices, where $n$ is the quantization bit width. Taking input X as an example, this means that its bit perspective layout changes from [M, K, p] to [p, M, K]. All threads within a thread block share the same shared memory space. Within each thread block, threads are further organized into a set of warps, with each warp consisting of 32 consecutive threads. Next, warps collaborates to load the A matrix (p*BM × BK) and B matrix (BK × q*BN) data required for thread block tile calculation from GL and caches them in SMEM. Thanks to BitPacking, the process of reading p BM*BK single-bit row-major tiles and writing p*BM*BK bits SMEM is efficient and continuous. Subsequently, thread block contains multiple warps, so thread block tile can be further decomposed into Warp Tile to achieve warp-level parallelism, and the computing tasks of each warp are WM× WN. In the calculation preparation stage, the A matrix (WM×WK, row-major) and B matrix (WK×WN, col-major) are independently loaded from SMEM to FR. Then, the calculation is decomposed into WM_TILES*WARP_N_TILES TensorCore MMA(matrix-multiply-accumulate). Since A and B are binarized matrices, we actually use Binary TensorCore MMA (BMMA), which has a computing power 8 times and 4 times higher than INT8 and INT4 TensorCore respectively. All warps collaboratively complete the Thread Block Tile calculation, and the results are stored in the c fragment of each warp. Therefore, each warp needs to independently write the calculation results back to SMEM. As shown in Eq. (10), output tile(p*BM × q*BN) is globally reduced to obtain a final result(BM × BN), where each BM × BN sub-Tile needs to be multiplied by a certain scaling factor. We call this process Bit Reduction. As the final step, warps collaboratively load the final result from SMEM and write back to the target location in GL.

GPU Kernel Optimization. When M=1, the GEMM problem of shape MxNxK transforms into a GEMV problem, shifting from computation-intensive to memory-intensive, which becomes a performance bottleneck for model inference. When using ordinary TensorCore for accelerated computation, the dimensions of M are usually chunked in groups of 8, requiring padding if M$<$8, leading to 87.5% redundant computation. Thanks to the revolutionary reconstruction of computing and BitPacking strategy, for the $W_{q}A_{p}$ configuration, the actual computing task undertaken by ABQKernel is p*M × q*N × K. The expansion of the M dimension can effectively reduce the redundant calculations when calling TensorCore, and even when p*M $>=$ 8 and p*M $\%$ 8 = 0, padding can be completely avoided. We call the above optimization strategy GEMV Elimination. In addition, Computational and Pipeline Optimization, Auto Kernel Search, and Bank Conflicts Elimination are also applied. For details, see Appendix D.

Baseline. For weight-only quantization, we compare our approach with GPTQ(Frantar et al. 2022), AWQ(Lin et al. 2024a), OmniQuant(Shao et al. 2023), and AffineQuant(Ma et al. 2024b). For weight-activation quantization, we benchmark our method against SmoothQuant(Xiao et al. 2023), OmniQuant(Shao et al. 2023), and I-LLM(Hu et al. 2024b).

Models and Datasets. We primarily evaluate our method using LLaMA (7B-13B) (Touvron et al. 2023a) and LLaMA-2 (7B-13B) (Touvron et al. 2023b)in this paper. Following previous work(Shao et al. 2023; Ma et al. 2024b), we evaluate the quantized models by reporting the perplexity of language generation experiments on WikiText2(Merity et al. 2016) and C4(Raffel et al. 2020). To assess performance on zero-shot tasks, we select several popular benchmarks including PIQA(Bisk et al. 2020), ARC(Clark et al. 2018), BoolQ(Clark et al. 2019), HellaSwag(Zellers et al. 2019), and Winogrande(Sakaguchi et al. 2021) using the lm-evaluation-harness(Gao et al. 2021).

Calibration We initialize the balance vectors for weights and activations following (Xiao et al. 2023), with the learnable clipping parameter for weights set to 1. For distribution correction compensation vectors, we set $a$ as an all-ones vector and $b$ as an all-zeros vector, ensuring $ab^{\top}$ starts at 0. Using the AdamW optimizer (Loshchilov and Hutter 2017) with no weight decay, we set learning rates of 5e-3 for balance vectors and 1e-2 for the clipping parameter and vector compensation vector. Calibration data includes 128 randomly selected 2048-token segments from WikiText2. The calibration process, conducted on an NVIDIA A800-40G GPU, utilized a batch size of 1 and spanned 20 epochs. For activation and KV Cache we perform per-token quantization, and for weight we perform per-channel quantization. By default, activation and KV cache use the same quantization bit.

Kernel Benchmark. We evaluated the GEMV speedup of our ABQKernel across three matrix dimensions in LLaMA-7B and compared it with the quantization kernel provided by cuBLAS and CUTLASS. It is important to note that CUTLASS only supports W4A4 and W8A8, while cuBLAS supports only W8A8 for quantization operations. Our experiments were conducted on two different GPUs: the RTX 4080 and the RTX 3070. Figure 5 presents the results, showing that our ABQKernel achieve superior acceleration across all matrix sizes. Specifically, for special bit combinations such as W2A8 and W2A4, our ABQKernel significantly outperforms the baseline approaches, as cuBLAS and CUTLASS require conversion to W8A8 and W4A4 for computation. In the W2A8 configuration, our throughput is improved by 7.47$\times$ compared to W8A8 with CUTLASS and cuBLAS on dimensions (1, 4096) $\times$ (4096, 4096).

End-to-end throughput. As shown in Figure 6, we integrate our ABQKernel into FastTransformer and compare it with the FP16 implementation of FastTransformer and the INT8 implementation of SmoothQuant. Compared to FP16, our scheme achieves 2.95$\times$ inference acceleration and 4.8$\times$ memory compression gain, while requiring only 10GB of memory for inference on the LLaMA-30B model, which is less than the memory required for LLaMA-7B with FP16. Additionally, our scheme achieves 1.6$\times$ speedup and 2.7$\times$ memory compression gain over SmoothQuant, significantly outperforming current mainstream inference methods. This substantial improvement reduces the cost of LLM services and facilitates their practical deployment.

Kernel Optimization Ablation. Table 4 presents the impact of various optimization techniques on the inference latency of the kernel when performing the GEMV operation with dimensions (1, 4096) $\times$ (4096, 4096). Our ABQKernel significantly outperforms the CUTLASS W8A8 kernel when unoptimized. Additionally, by employing pipeline optimization, GEMV elimination and auto kernel search, we achieve a latency reduction by 7.47$\times$ and a corresponding increase in throughput by 7.47$\times$. These results significantly outperform CUTLASS.