COMET: Towards Partical W4A4KV4 LLMs Serving

TL;DR

This work proposes a novel fine-grained mixed-precision quantization algorithm (FMPQ) that compresses most activations into 4-bit with negligible accuracy loss, and develops a highly optimized W4Ax kernel to support mixed-precision matrix multiplication for W4A4 and W4A8.

Abstract

Quantization is a widely-used compression technology to reduce the overhead of serving large language models (LLMs) on terminal devices and in cloud data centers. However, prevalent quantization methods, such as 8-bit weight-activation or 4-bit weight-only quantization, achieve limited performance improvements due to poor support for low-precision (e.g., 4-bit) activation. This work, for the first time, realizes practical W4A4KV4 serving for LLMs, fully utilizing the INT4 tensor cores on modern GPUs and reducing the memory bottleneck caused by the KV cache. Specifically, we propose a novel fine-grained mixed-precision quantization algorithm (FMPQ) that compresses most activations into 4-bit with negligible accuracy loss. To support mixed-precision matrix multiplication for W4A4 and W4A8, we develop a highly optimized W4Ax kernel. Our approach introduces a novel mixed-precision data layout to facilitate access and fast dequantization for activation and weight tensors, utilizing the GPU's software pipeline to hide the overhead of data loading and conversion. Additionally, we propose fine-grained streaming multiprocessor (SM) scheduling to achieve load balance across different SMs. We integrate the optimized W4Ax kernel into our inference framework, COMET, and provide efficient management to support popular LLMs such as LLaMA-3-70B. Extensive evaluations demonstrate that, when running LLaMA family models on a single A100-80G-SMX4, COMET achieves a kernel-level speedup of \textbf{$2.88\times$} over cuBLAS and a \textbf{$2.02 \times$} throughput improvement compared to TensorRT-LLM from an end-to-end framework perspective.

Figures (12)

• Figure 1: The inference procedure of LLMs. The inference procedure can be divided into two phases: prompting and token generation.
• Figure 2: Roofline model analysis for activation-activation and weight-activation operators. The activation-activation operator is highly memory-bound, while the weight-activation operator becomes compute-bound with large-batch parallelism.
• Figure 3: Activation distributions in LLaMA-7B. Outliers are present in certain channels of activations, leading to the ineffectiveness of conventional quantization strategies.
• Figure 4: The algorithm design of FMPQ.
• Figure 5: The design overview of the COMET-W4Ax kernel. (a) illustrates the tile-based GEMM computing with mixed-precision encoding. (b) presents the computing procedure when issuing W4A4 and W4A8 GEMM instructions simultaneously. (c) shows that we can effectively hide the overhead by software pipeline.