Title: QServe: W4A8KV4 Quantization and System Co-design for Efficient LLM Serving
Title: QServe: W4A8KV4 Quantization and System Co-design for Efficient LLM Serving
发表时间: 2025-05 · arXiv:2405.04532 (MLSys 2025)
原文: https://arxiv.org/abs/2405.04532
Yujun Lin * 1, Haotian Tang * 1, Shang Yang * 1, Zhekai Zhang 1, Guangxuan Xiao 1, Chuang Gan 2 3, Song Han 1 4
URL: https://hanlab.mit.edu/projects/qserve
速读
一句话结论
本文提出了一种名为 QoQ 的 W4A8KV4 量化算法及配套推理系统 QServe,通过将大语言模型的权重和 KV Cache 压缩至 4 比特、激活值保持 8 比特,在几乎不损失精度的前提下,使模型在 GPU 上的大批量服务吞吐量比 TensorRT-LLM 提升了 1.2 至 3.5 倍。
要解决什么问题
现有的 INT4 量化方法(如 W4A4)在小批量边缘设备推理上有效,但在云端大批量服务中却无法带来吞吐量提升,甚至比 W4A16 更慢。核心卡点在于 GPU 架构的底层执行机制:为了保住精度,W4A4 必须使用细粒度的分组量化(per-group),这导致 Tensor Core 计算出的 INT32 部分和必须在 GEMM 算子的主循环中反量化为浮点数。这一反量化过程只能依赖 GPU 上吞吐量极低的 CUDA Core 来执行。以 A100 为例,CUDA Core 的算力仅为 INT4 Tensor Core 的五十分之一,导致反量化操作占据了 20% 到 90% 的运行时间。此外,在处理长上下文时,将 KV Cache 压缩到 4 比特(KV4)同样会因为复杂的位操作和反量化计算,让原本受限于显存带宽的 Attention 算子退化成受限于 CUDA Core 算力的计算密集型任务,从而完全抵消了低比特带来的显存读取红利。
怎么做的
作者提出了 QoQ 算法与 QServe 推理系统,核心思路是让所有矩阵乘法都在高速的 INT8 Tensor Core 上运行,同时将反量化开销降到最低。第一步,为了绕开 CUDA Core 的计算瓶颈,算法端引入了渐进式分组量化。它先将权重用通道级缩放因子量化到 INT8,再用分组缩放因子进一步量化到 INT4。关键设计是为第一步的 INT8 设定了 [-119, 119] 的安全保护范围,而非用满整个 8 比特空间。第二步,在系统实现层,这个保护范围使得 QServe 能够在反量化时采用“先乘法后减法”的计算顺序:
这种顺序保证了中间结果不会溢出 INT8 的表示范围,从而完美利用 GPU 的寄存器级并行,用一条指令同时处理 4 个 INT4 到 INT8 的反量化,并将零点减法移出主循环,彻底消除了 CUDA Core 瓶颈。第三步,针对内存读取效率,QServe 引入了计算感知权重重排。由于 4 比特权重和 8 比特激活的数据类型不匹配,常规的内存加载指令会失效并引发大量指针算术运算。系统通过在离线阶段按线程实际计算时的读取顺序对权重进行重排,确保了运行时能以 128 比特的最高带宽满载读取。第四步,为了解决 4 比特 KV Cache 掉精度的问题,算法端提出了 SmoothAttention。由于 Value 矩阵没有明显的异常值,该方法只针对 Key 矩阵,通过一个通道级缩放因子 $\lambda_i$ 抹平 Key 的异常通道,并将量化难度转移到不被量化的 Query 上:
效果如何
实验在 A100(80GB)和 L40S(48GB)GPU 上进行,测试了 Llama-1/2/3、Mistral、Mixtral、Yi 和 Qwen1.5 等 7B 到 72B 规模的模型。对比基线包括代表工业界主流路线的 TensorRT-LLM(分别测试了 FP16、W8A8、W4A16 三种配置),以及代表 W4A4 激进量化路线的最新系统 Atom 和 QuaRot。在输入 1024 tokens、输出 512 tokens 的大批量生成任务下,QServe 在 A100 上比 TensorRT-LLM 表现最好的配置还要快 1.2 到 2.4 倍,比 Atom 和 QuaRot 快 2.5 到 2.9 倍。在 L40S 上,QServe 的吞吐量提升更为显著,达到了 TensorRT-LLM 的 1.4 到 3.5 倍。得益于 W4A8KV4 极小的显存占用,QServe 在显存较小的 L40S 上能跑满与 A100 相同的 Batch Size,导致在测试的 8 个模型中有 6 个,L40S 运行 QServe 的绝对吞吐量甚至超越了价格贵 3 倍的 A100 运行 TensorRT-LLM 的吞吐量。在精度方面,QoQ 算法在 WikiText2 困惑度和常识推理零样本准确率上均大幅超越了 W4A4 基线,仅比近乎无损的 W4A16 出现极微小的下降。代价方面,该方法在处理 KV Cache 时需要动态的逐头量化来维持精度,这在系统实现上要求将 FP16 的缩放因子与 KV 数据交错存储在显存页中以支持实时更新。
主要贡献
大型语言模型(LLM)的巨大规模使其部署极具挑战性,目前最先进的INT4量化技术仅能加速小批量、边缘LLM推理,未能在大批量、基于云的LLM服务中带来性能提升。当前4位整数量化方法在当前代GPU上对权重或部分和进行反量化时,会遭遇显著的运行时开销($20\% - 90\%$)。
为了解决这一挑战,本文引入了QoQ算法,这是一种具有4位权重、8位激活和4位KV缓存的W4A8KV4量化算法。QoQ代表quattuor-octo-quattuor,在拉丁语中表示4-8-4。QoQ由QServe推理系统实现,该系统取得了显著的实测加速。驱动QServe的核心见解是:GPU上LLM服务的效率受到低吞吐量CUDA核心操作的严重影响。基于此,在QoQ算法中,本文引入了渐进式量化(progressive quantization),允许在W4A8 GEMM中实现低反量化开销。此外,本文开发了SmoothAttention,以有效缓解由4位KV量化引起的精度下降。在QServe系统中,本文执行计算感知的权重重排序(compute-aware weight reordering),并利用寄存器级并行性(register-level parallelism)来减少反量化延迟。QServe还将KV4注意力带来的理论内存节省转化为实测加速。结果表明,与TensorRT-LLM相比,QServe将Llama-3-8B在A100上的最大可实现服务吞吐量提高了$1.2\times$,在L40S上提高了$1.4\times$;将Qwen1.5-72B在A100上的吞吐量提高了$2.4\times$,在L40S上提高了$3.5\times$。值得注意的是,L40S GPU上的QServe甚至可以实现比A100上的TensorRT-LLM更高的吞吐量。
背景知识与关键Observation
大型语言模型
大型语言模型(LLMs)是由多个结构相同的层组成的因果Transformer模型系列。每一层结合了注意力块、前馈网络(FFN)和归一化层。每层的输入$\mathbf{x}$是一个$N \times HD$的张量,其中$N$是输入Token的数量,$H$代表注意力头的数量,$D$是每个头的隐藏维度。服务LLMs涉及两个阶段:预填充阶段(所有提示Token同时呈现,$N > 1$)和解码阶段(模型每次只处理一个Token,$N = 1$)。在注意力块中,$\mathbf{x}$首先经过线性投影得到$\mathbf{q} \in \mathbb{R}^{N \times HD}$,$\mathbf{k}, \mathbf{v} \in \mathbb{R}^{N \times H_{KV}D}$,其中$H_{KV}$是键/值头的数量。将$\mathbf{k}, \mathbf{v}$与预计算的$S$个先前Token的KV缓存特征连接,得到$\mathbf{K}, \mathbf{V} \in \mathbb{R}^{(S+N) \times H_{KV}D}$,并计算注意力:
结果$\mathbf{o}$与输出投影矩阵$\mathbf{W}_O \in \mathbb{R}^{HD \times HD}$相乘,乘积加到$\mathbf{x}$上作为FFN的输入。
整数量化
整数量化将高精度数字映射到离散级别。该过程可以公式化为:
其中$\mathbf{X}$是浮点张量,$\mathbf{Q}_{\mathbf{X}}$是其$n$位量化对应物,$s$是缩放因子,$z$是零点。反量化张量可以表示为:
为什么选择W4A8KV4
W4A8KV4在Roofline模型上优于W8A8和W4A16:如图2a所示,在实际对话中,注意力和GEMM占据了LLM部署的大部分运行时间,且解码阶段的运行时间约为预填充阶段的6倍。对于$m \times n \times k$的GEMM问题,当$n, k$远大于$m$时,计算强度约为$m$。根据图3中的A100 Roofline模型,当$m < 78$时,W4A16具有更高的理论吞吐量,因为此时GEMM受内存限制,W4A16较小的内存占用带来了更好的性能;当$m > 78$时,问题变为计算受限,W8A8由于INT8张量核心的更高吞吐量而表现更好。W4A8能够结合两者的优势。此外,LLM解码中的注意力工作负载可表述为一系列批处理GEMV操作,计算强度为1 MAC/element,内存流量由KV缓存访问主导。将KV缓存量化为INT4相当于有效增加了内存带宽,为注意力提供了相比KV8 2倍的峰值性能。
为什么不选择W4A4KV4:如图2b所示,现有的W4A4服务系统Atom【1,Atom: Low-bit quantization for efficient and accurate llm serving + 2023 + MLSys】和QuaRot【2,Quarot: Outlier-free 4-bit inference in rotated llms + 2024 + arXiv】甚至比TensorRT-LLM的W16A16解决方案慢得多。这是因为GPU上映射逐组量化的W4A4 GEMM存在固有的困难。如图4所示,最先进的系统采用输出固定数据流实现张量核心GEMM。对于$m \times n \times k$问题,每个线程块通过顺序迭代归约维度$k$来计算$t_m \times t_n$的输出瓦片,这个顺序循环被称为主循环。在FP16和W8A8 GEMM中(图5a),主循环完全在张量核心上执行。而TensorRT-LLM-W4A16(图5b)和Atom-W4A4(图5c)都需要在主循环中进行反量化操作,这些操作在CUDA核心上运行。在A100和H100上,FP32 CUDA核心的峰值性能仅为其INT4张量核心对应物的$2\%$。因此,Atom中反量化单个部分和相当于50次张量核心MAC操作,导致主循环被缓慢的CUDA核心操作主导。其次,Atom创建两组寄存器(一组FP32,一组INT32)来保存部分和,高寄存器消耗限制了SM上同时执行的warp数量,进一步加剧了主循环开销。相比之下,QServe的W4A8设计(图5d)采用两级渐进式组量化,确保所有计算都在INT8张量核心上执行,并应用4路寄存器级并行性同时解码四个INT4权重,从而减少主循环开销。
方法细节
QoQ量化算法
为了在不影响LLM效果的情况下释放W4A8KV4的全部潜力,本文提出了包含渐进式组量化、SmoothAttention以及各种通用量化优化的QoQ算法。
渐进式组量化:为了解决系统实现中反量化开销抵消低比特量化精度提升的问题,作者引入了如图6所示的渐进式组量化方法。给定权重张量$\mathbf{W} \in \mathbb{R}^{k \times n}$,首先应用逐通道对称INT8量化:
其中$\mathbf{Q}_{\mathbf{W}_{\mathbf{s8}}}^{(0)} \in \mathbb{N}^{n \times k}$是中间8位量化权重张量,$\mathbf{s}_{\mathrm{fp16}}^{(0)} \in \mathbb{R}^{n \times 1}$是逐通道量化缩放因子。接着,对中间权重张量进一步使用逐组非对称INT4量化:
保护性量化范围:朴素地应用上述公式并不能保证中间反量化权重完美地落在8位整数表示范围内(即$[-128, 127]$)。如果简单地在反量化期间打开算术指令中的饱和选项,会严重损害计算吞吐量(降低高达$67\%$)。通过将公式代入推导,可以得到:
由于$s_{u8} \leq \frac{127 - (-128)}{15 - 0} = 17$,可以推导出$\hat{q}_{s8} \leq 127 - \frac{1}{2}s_{u8} \leq 119.5$。因此,作者将INT8对称量化范围从$[-127, 127]$缩小到保护范围$[-119, 119]$以避免反量化溢出。与先前直接应用目标精度的组量化然后对组级浮点缩放因子执行逐通道量化的方法(如QLoRA【3,Qlora: Efficient finetuning of quantized llms + 2023 + arXiv】)不同,QoQ引入了保护范围,允许将反量化操作融合到具有完全寄存器级并行性的W4A8 GEMM内核中,最小化CUDA核心开销,使得QServe的W4A8逐组GEMM实现了比W8A8 cuBLAS GEMM快$1.5\times$的加速。
SmoothAttention:直接将KV缓存降低到4位会显著降低LLM的准确性。通过可视化采样的Key和Value缓存激活的幅度分布(图7),作者观察到Value矩阵没有显著的异常值模式,而Key矩阵倾向于在每个头部具有固定的异常值通道,这些异常值比大多数激活值大$10\times$左右。受SmoothQuant【4,SmoothQuant: Accurate and efficient post-training quantization for large language models + 2023 + ICML】启发,作者提出了SmoothAttention,通过逐通道因子$\lambda$缩小Key缓存中的异常值通道:
由于不量化Queries,只需关注Keys并选择缩放因子$\lambda_i = \max(|\mathbf{K}_i|)^\alpha$,实践中$\alpha=0.5$即可。为了消除内核调用开销并使缩放在旋转位置嵌入(RoPE)方面可交换,作者添加了硬约束$\lambda_i = \lambda_{i+\frac{D}{2}} = \max(\max(|\mathbf{K}_i|), \max(|\mathbf{K}_{i+\frac{D}{2}}|))^\alpha$。最后将缩放$\mathbf{\Lambda}$融合到先前层的权重中。
通用LLM量化优化:
- 块输入模块旋转:受相关工作启发,通过乘以旋转矩阵来旋转块输入激活(如QKV投影层和FFN第一层),并在相反方向旋转对应的权重以保持数学等价性(图8)。旋转后,异常值通道被有效抑制,且旋转矩阵可被前一个线性层的权重吸收。
- 块输出模块平滑:对于生成块输出的层(如输出投影层和FFN第二层),通过将块中间激活除以逐通道平滑因子来进行平滑(图9)。作者发现迁移强度$\alpha$应该接近0,即平滑因子主要由权重决定。
- 激活感知通道重排序:使用$\max(|\mathbf{X}|)$来确定通道显著性,然后重新排序通道,使得具有相似显著性的通道在同一个量化组中(图10)。
- 权重裁剪:应用裁剪比例$\alpha$到动态范围。作者最小化所有线性层(除q_proj和k_proj外)的层输出误差,对于q_proj和k_proj则优化块输出均方误差:$\arg \min_{\boldsymbol{\alpha}} \|\mathrm{Block}(\mathbf{X}; \mathbf{W}) - \mathrm{Block}(\mathbf{X}; Q(\mathbf{W}; \boldsymbol{\alpha}))\|$。
QServe服务系统
QServe系统运行时:如图11所示,QServe中的所有GEMM层都在W4A8输入上运行,在INT8张量核心上执行计算,并生成FP16输出。所有注意力层都在CUDA核心上以FP16执行计算。为了确保每个GEMM接收INT8激活,激活量化被融合到前面的layernorm或激活内核中。在KV缓存管理方面,采用分页KV缓存,并在每个KV缓存页面中量化KV特征之后立即存储每头的FP16缩放因子和零点,允许动态更新。
计算感知权重重排序:在每次主循环迭代期间加载操作数时,张量核心GEMM内在函数要求跨步布局。当存储和计算数据类型相同时,ldmatrix指令可以解决此问题(图12a)。但对于W4A8,ldmatrix会导致线程获取的数据与计算所需数据不匹配(图12b),产生指针算术开销并降低带宽利用率。为此,作者应用计算感知权重重排序(图12c),将整个GEMM问题划分为$32 \times 32$瓦片,并将线程使用的32个通道连接成单一的128位字,按计算使用的顺序存储权重。这种离线重排序不引入运行时开销,消除了指针算术开销并保证了高带宽的内存事务。
逐通道W4A8 GEMM中的快速反量化:对于主循环中从ZINT4到SINT8的转换,首先进行权重解包(UINT4到UINT8)。如图13所示,将每32个UINT4权重重新排序,利用寄存器级并行性,仅用三个逻辑操作将其解包为UINT8。对于UINT8到SINT8的转换(零点减法),为了避免主循环中的整数减法开销,作者采用先乘后减的方法。通过重写公式,将零点减法从主循环移动到epilogue中:
第二项可以转化为$\mathbf{t}_{\mathbf{X}} \times (\mathbf{z}_{\mathbf{W}} \odot \mathbf{s}_{\mathbf{W}})$,其中$\mathbf{t}_{\mathbf{X}} = \mathbf{X} \mathbf{1}_k$。预计算$\mathbf{t}_{\mathbf{X}}$的内核被融合到前一个内存受限内核中,开销可忽略。
逐组W4A8 GEMM中的快速反量化:由于零点是逐组定义的,无法合并到epilogue中,且需要额外的INT8乘法。作者坚持采用先乘后减的方法以实现寄存器级并行性(RLP)。如图14所示,NVIDIA GPU提供vadd4指令执行四路INT8加法。在先减后乘方法中,比例乘法的结果会溢出,破坏RLP。而在先乘后减计算顺序中,渐进式组量化算法确保初始乘法步骤的结果永远不会超过INT8范围,从而允许在乘法和减法中充分利用RLP。
QServe中的KV4注意力:单纯地将KV缓存量化为INT4在A100上会导致$1.2\times$的减速(表1),因为融合的注意力内核计算强度极高(反量化单个INT4数字需要5个ALU操作),迅速饱和了FP32 CUDA核心的roofline转折点。为了缓解计算受限瓶颈,作者采用双向方法:首先,将原始内核中的所有FP32操作替换为FP16,使计算上限翻倍;其次,应用位技巧将反量化的算术强度降低到每个元素2个操作。接着,简化控制逻辑,并异步预取缩放因子和零值。这些优化使得在A100上相比KV8基线实现了$1.5\times$的加速。
实验环境
- 数据集:语言建模任务使用WikiText2(困惑度),零样本任务使用PIQA、ARC、HellaSwag和WinoGrande。长上下文评估使用LongBench。校准数据集主要使用Pile验证集,也报告了使用WikiText2的结果。
- 模型架构:Llama-1 (30B)、Llama-2 (7B, 13B, 70B)、Llama-3 (8B)、Mistral-7B、Mixtral-8x7B和Yi-34B。
- 硬件配置:评估在A100-80G-SXM4和L40S-48G GPU上进行。主机配备x86_64 CPU。
- 软件配置:QoQ算法基于PyTorch 2.2.0和HuggingFace实现。QServe系统使用CUDA 12.2和PTX汇编实现。基线系统包括TensorRT-LLM v0.9.0、QuaRot和Atom。除了QuaRot外,所有系统均启用了分页注意力。
实验结果
- 准确度评估 (WikiText2困惑度与零样本精度):如表2所示,QoQ在Llama-2-7B上的困惑度仅比W8A8和W4A16基线增加最多0.16,且一致优于Atom和QuaRot。如表3所示,在零样本任务中,QoQ显著优于其他4位量化方法(如在Winogrande任务上比QuaRot高4.82%),相比FP16仅引入了极小的精度损失(Llama-2 7B/13B/70B分别为1.03%、0.89%和0.40%)。表5证明QoQ能保持与BF16基线相当的长上下文性能。
- 系统吞吐量评估:如图15和表4所示,在输入序列长度1024、输出序列长度512的设置下,QServe在A100上相比TensorRT-LLM的最佳配置实现了$1.2\times - 2.4\times$的吞吐量提升。在L40S上,性能提升更为显著,达到了$1.47\times - 3.47\times$。值得注意的是,QServe在L40S上运行34B以下的五个模型时,实现了比在A100上运行TensorRT-LLM更高的吞吐量,有效降低了硬件成本。
- 消融实验(量化技术):如图16所示,以W8A8为起点,将权重降至4位会降低精度但提高吞吐量并节省内存。引入块输入旋转和权重裁剪后,W4A8配置达到了与W4A16相当的困惑度。将KV降至4位再次降低精度但大幅提升吞吐量。SmoothAttention、渐进式组量化和激活感知通道重排序逐步修复了精度损失,且几乎不增加系统开销。
- 消融实验(系统反量化开销):如图18所示,QServe的反量化开销与TRT-LLM-W4A16相当,但远小于Atom-W4A4(减少高达90%),且由于在INT8张量核心上计算,享受了更高的吞吐量。
- 同等批次大小比较:如图17所示,QServe的加速来源于同等批次下的单批次加速(如Llama-2-7B的$1.45\times$)以及由于内存节省带来的批次规模扩大。
- KV4注意力优化分解:基础KV4实现延迟为$0.48\mathrm{ms}$,应用位技巧降至$0.44\mathrm{ms}$,简化控制流降至$0.39\mathrm{ms}$,转换为FP16降至$0.36\mathrm{ms}$,异步预取最终将延迟降至$0.28\mathrm{ms}$,实现了端到端$1.7\times$的提升。
结论
本文介绍了QServe,这是一个算法与系统协同设计的框架,旨在将大型语言模型量化为W4A8KV4精度,以促进其在GPU上的高效部署。在算法方面,设计了QoQ量化方法,采用渐进式量化使W4A8 GEMM能在INT8张量核心上执行,并提出了SmoothAttention以减少KV4量化带来的精度损失。在QServe系统中,利用渐进式量化第一级建立的保护范围来实现INT4到INT8的反量化,利用完全的寄存器级并行性并采用先乘后减的计算顺序。此外,实现了计算感知的权重重排序以最小化指针算术开销。最终,QServe在A100和L40S上服务代表性LLM时,实现了比工业标准TensorRT-LLM高出$2.4\times - 3.5\times$的吞吐量。
附录细节
Artifact Appendix:附录提供了重现论文中关键实验所需的脚本和依赖项。要求主机系统配备x86_64 CPU,以及至少一个A100或L40S NVIDIA GPU。作者提供了一个预构建的Docker镜像来简化环境设置过程(基于NVIDIA Container Toolkit)。文档中详细列出了启动QServe基准测试和TensorRT-LLM基准测试的bash命令,并提供了预期的吞吐量结果参考。用户可以通过修改基准测试脚本来使用不同的模型和批次大小进行实验。精度评估结果可以通过开源库deepcompressor重现。
💬 评论讨论
欢迎在这里分享您的想法和见解!