KLLM: Fast LLM Inference with K-Means Quantization

Wu, Xueying, Zhou, Baijun, Gao, Zhihui, Fu, Yuzhe, Zheng, Qilin, He, Yintao, Li, Hai

arXiv.org Artificial Intelligence 

Abstract--Large language model (LLM) inference poses significant challenges due to its intensive memory and computation demands. Weight and activation quantization (W AQ) offers a promising solution by reducing both memory footprint and arithmetic complexity. Traditional W AQ designs rely on uniform integer quantization for hardware efficiency, but often suffer from significant model performance degradation at low precision. In contrast, K-Means quantization, a non-uniform technique, achieves higher accuracy by aligning with the Gaussian-like distributions of weights and activations in LLMs. However, two key challenges prevent the efficient deployment of K-Means-based W AQ designs for LLM inference: (1) The non-uniform structure of K-Means-quantized data precludes direct execution on low-precision compute units, necessitating dequantization and floating-point matrix multiplications (MatMuls) during inference. Offline thresholding methods for outlier detection degrade model performance substantially, while existing online detection techniques introduce significant runtime overhead. T o address the aforementioned challenges and fully unleash the potential of K-Means-based W AQ for LLM inference, in this paper, we propose KLLM, an LLM inference accelerator for efficient execution with K-Means-quantized weights and activations. KLLM features an index-based computation scheme for efficient execution of MatMuls and nonlinear operations on K-Means-quantized data, which avoids most of the dequantization and full-precision computations. Moreover, KLLM incorporates a lightweight outlier detection engine, Orizuru, that efficiently identifies the top-k largest and smallest elements in the activation data stream during online inference. Extensive experiments show that, KLLM achieves speedups of 8.66, 2.12 and energy efficiency improvements of 214.76, 2.79 compared to the A100 GPU and SpAtten, respectively. Large language models (LLMs) have achieved impressive advancements, demonstrating versatile applications such as chatbots [1], [18], code generation [41], [46], and circuit design [14], [20].