实现4比特权重激活与键值缓存的高效推理,显著提升大模型部署速度。
COMET: Towards Partical W4A4KV4 LLMs Serving
- 提出细粒度混合精度量化算法,将多数激活压缩至4比特
- 自研W4Ax内核与数据布局,在A100上实现2.88倍算子加速
- 支持主流大模型如LLaMA-3-70B,适合追求低延迟推理的部署场景
量化是降低大语言模型在终端设备和云数据中心服务开销的常用压缩技术。然而,当前主流方法如8比特权值激活或4比特仅权值量化,受限于对低精度(如4比特)激活的支持不足,性能提升有限。本文首次实现适用于大型语言模型的实用化W4A4KV4推理,充分运用现代GPU的INT4张量核心,缓解由键值缓存引发的内存瓶颈。我们提出一种新型细粒度混合精度量化算法(FMPQ),将大部分激活压缩至4比特,且精度损失可忽略。为支持W4A4和W4A8的混合精度矩阵乘法,开发了高度优化的W4Ax内核。设计新型混合精度数据布局,促进激活与权值张量的快速访问与解量化,利用GPU软件流水线隐藏数据加载与转换开销。此外,提出细粒度流式多处理器(SM)调度策略,实现不同SM间负载均衡。将优化后的W4Ax内核集成至推理框架COMET,并提供对LLaMA-3-70B等主流大模型的高效支持。大量实验表明,在单块A100-80G-SMX4上运行LLaMA系列模型时,COMET相比cuBLAS实现2.88倍的内核级加速,从端到端框架角度看,相比TensorRT-LLM提升2.02倍吞吐量。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。