Efficient Arbitrary Precision Acceleration for Large Language Models on GPU Tensor Cores

TL;DR

This work tackles the bottlenecks of deploying ultra-low-bit quantized LLMs on GPUs by introducing a bipolar-INT data format and a full pipeline for arbitrary-precision MatMul via bit-level decomposition and recovery. It combines matrix preprocessing and a recovery-oriented memory management strategy to maximize shared-memory usage and minimize global-memory traffic, achieving significant speedups over baseline tensor-core implementations. The approach delivers up to 2.4x MatMul speedup against NVIDIA CUTLASS and up to 6.7x end-to-end LLM inference acceleration, demonstrating practical viability for diverse models (e.g., Llama2-7B, OPT-6.7B, BLOOM-7B). Overall, the methods enable true arbitrary-precision acceleration on GPUs, broadening the deployment of ultra-low-bit LLMs in real-time applications.

Abstract

Large language models (LLMs) have been widely applied but face challenges in efficient inference. While quantization methods reduce computational demands, ultra-low bit quantization with arbitrary precision is hindered by limited GPU Tensor Core support and inefficient memory management, leading to suboptimal acceleration. To address these challenges, we propose a comprehensive acceleration scheme for arbitrary precision LLMs. At its core, we introduce a novel bipolar-INT data format that facilitates parallel computing and supports symmetric quantization, effectively reducing data redundancy. Building on this, we implement an arbitrary precision matrix multiplication scheme that decomposes and recovers matrices at the bit level, enabling flexible precision while maximizing GPU Tensor Core utilization. Furthermore, we develop an efficient matrix preprocessing method that optimizes data layout for subsequent computations. Finally, we design a data recovery-oriented memory management system that strategically utilizes fast shared memory, significantly enhancing kernel execution speed and minimizing memory access latency. Experimental results demonstrate our approach's effectiveness, with up to 2.4\times speedup in matrix multiplication compared to NVIDIA's CUTLASS. When integrated into LLMs, we achieve up to 6.7\times inference acceleration. These improvements significantly enhance LLM inference efficiency, enabling broader and more responsive applications of LLMs.

Figures (7)

• Figure 1: Comparison between bipolar-INT and traditional integers. Bipolar-INT is well-suited for TCs’ parallel computing due to its symmetric quantization and unified operations.
• Figure 2: Illustration of the decomposition and recovery process for bipolar-INT MatMuls. Here matrices W and X are both 2-bit, which can be extended to arbitrary bit widths.
• Figure 3: Procedure of matrix decomposition and reassembly to save GPU's memory and data transfer instructions.
• Figure 4: Recovery-oriented memory scheduling strategy for arbitrary precision MatMul on GPUs, leveraging shared memory and fragments to reduce global memory access and accelerate computation.
• Figure 5: Comparison of throughput between our work and other methods in the context of square MatMuls.