LiquidGEMM: Hardware-Efficient W4A8 GEMM Kernel for High-Performance LLM Serving

作者/机构: Huanqi Hu (Shanghai Jiao Tong University), Bowen Xiao (ByteDance Seed), Shixuan Sun (Shanghai Jiao Tong University), Jianian Yin (ByteDance Seed), Zhexi Zhang (ByteDance Seed), Xiang Luo (ByteDance Seed), Chengquan Jiang (ByteDance Seed), Weiqi Xu (ByteDance Seed), Xiaoying Jia (ByteDance Seed), Xin Liu (ByteDance Seed), Minyi Guo (Shanghai Jiao Tong University)

A1 主要贡献

本文旨在解决大型语言模型（LLM）服务中4位权重和8位激活（W4A8）量化方案的性能瓶颈问题。尽管W4A8在理论上能在准确性和性能之间取得良好平衡，但现有的W4A8通用矩阵乘法（GEMM）内核因在CUDA核心上进行低效的反量化操作而表现不佳，其速度无法跟上张量核心（Tensor Cores）的高吞吐量。

核心问题: 现有W4A8 GEMM内核的反量化步骤在CUDA核心上执行，由于存在潜在的溢出问题，需要数十条指令来解决，这给计算能力有限的CUDA核心带来了巨大压力，使其成为性能瓶颈，导致无法充分发挥张量核心的性能优势。这造成了理论上的性能潜力（如屋顶线分析所示）与实际表现之间的巨大差距。

研究目标: 提出一个硬件高效的W4A8 GEMM内核，通过优化量化算法和执行流水线，解决反量化开销大的问题，从而为高性能LLM服务提供支持。

主要贡献与创新点:
1. 深入分析: 本文深入分析了W4A8 GEMM的执行流水线，并建立了一个成本模型，准确定位了性能瓶颈在于硬件未感知的反量化步骤。
2. 提出 LiquidGEMM: 这是一个为高效LLM服务优化的高性能W4A8 GEMM内核。
3. 开发 LiquidQuant (LQQ): 这是一种硬件高效的量化算法，通过基于旋转的变换和利用二进制补码特性，实现了快速、无溢出的反量化。该方法仅需两条算术指令（IMAD和XOR）即可处理四个元素，极大地降低了CUDA核心上的反量化开销。
4. 引入隐式细粒度流水线 (ImFP): 该流水线机制能够充分重叠权重加载、反量化和矩阵乘法累加（MMA）操作。它采用单生产者-多消费者的执行模型，利用硬件调度来避免软件同步开销和冗余的内存流量，从而最大化硬件利用率。

实验结果表明，LiquidGEMM相比最先进的W4A8内核实现了高达2.90倍的加速，端到端系统级加速高达4.94倍。与NVIDIA TensorRT-LLM中的各种量化GEMM内核相比，LiquidGEMM也取得了1.12-1.63倍的性能提升。

图1：NVIDIA A100和H100 GPU的关键性能指标，以及LLM服务中GEMM层的屋顶线模型。

A3 背景知识与关键洞察

2 基础知识

整数环化。整数环化是一种关键技术，通过将高精度的浮点值（FP32或FP16）转换为低精度的整数表示（如INT8或INT4），来减少LLM的内存占用和计算成本。这种转换使得模型能够在支持整数运算的GPU上更高效地执行。形式上，量化将一个浮点张量 $W$ 映射到一个 $n$ 位整数张量 $Q$，具体如下：

$$Q=\left|\frac{W}{s}+z\right|, s=\frac{\max (W)-\min (W)}{\max (Q)-\min (Q)}, z=\left|\min (Q)-\frac{\min (W)}{s}\right| .$$

其中 $s$ 是缩放因子，$z$ 是零点。操作符 $⌊·⌉$ 表示四舍五入到最近的整数。由于 $Q$ 使用 $n$ 位表示，其动态范围对于无符号整数被限制在 $[0, 2^n − 1]$，对于有符号整数则被限制在 $[−2^{n−1}, 2^{n−1} − 1]$，具体取决于量化类型。相应的反量化过程从量化的整数张量 $Q$ 中重建一个近似的浮点值 $\hat{W}$：$\hat{W} = s(Q - z)$。在实践中，使用了两种常见的量化变体：非对称量化，其中 $z$ 非零以适应任意输入范围；对称量化，其范围以零为中心，且 $z$ 设置为0。在非对称量化中，整数范围由 $max(Q) − min(Q) = 2^n − 1$ 给出，而在对称量化中，范围变为 $2^n − 2$，因为 $|max(Q)| = |min(Q)|$。与对称量化相比，非对称量化可以充分利用可用的值范围，但在反量化过程中需要额外的减法操作。

GPU上的GEMM。图2概述了GPU上的GEMM执行过程。给定一个GEMM操作 $Y = XW^T$，其中 $X \in R^{M \times K}$ 是输入张量，$W^T \in R^{K \times N}$ 是权重矩阵，$Y \in R^{M \times N}$ 是输出，GPU将 $Y$ 分割成大小为 $M_t \times N_t$ 的瓦片（tile），每个瓦片由一个线程块处理。为了计算其分配的瓦片，一个线程块以 $K_t$ 的步长迭代 $K$ 维度，执行一系列大小为 $M_t \times N_t \times K_t$ 的较小GEMM。在每次迭代中，它加载相应的 $X$ 和 $W$ 的切片，执行乘法累加操作，并更新输出瓦片。这个对 $K$ 维度的迭代，称为主循环，主导了GEMM的整体计算成本。每个输出瓦片进一步划分为片段（fragment），每个warp使用张量核心上的MMA（矩阵乘法累加）指令计算一个片段。这些硬件加速的张量核心针对小型矩阵形状（例如，64×256×32）进行了优化，通过并行处理多个片段实现高吞吐量计算。为简单起见，本文中我们交替使用术语瓦片和片段，因为它们的区别不影响核心分析。
图2：GPU上的GEMM概览，其中$K_t$, $M_t$, $N_t$分别表示沿$K$, $M$, $N$维度的循环迭代。

张量核心支持的操作。张量核心原生支持对具有匹配对称精度的操作数进行运算，即权重和激活值具有相同的数据类型。根据输入矩阵的精度，GEMM可分为两种类型：对称GEMM，其中两个操作数共享相同类型；以及非对称GEMM，其中权重和激活值的精度不同（通常，权重的位宽较低）。在非对称GEMM中，权重必须在主循环期间被张量核心处理之前进行反量化。图3比较了W4A8（一种非对称GEMM）和W8A8（一种对称GEMM）。在W4A8中，反量化在主循环期间在CUDA核心上执行，之后才在张量核心上进行MMA。相比之下，W8A8完全在张量核心上执行主循环，反量化被推迟到结尾（epilogue）阶段。
图3：TensorRT-LLM中的W8A8 GEMM与QServe中的W4A8 GEMM的比较。

3 动机

性能评估。我们评估了W4A8 GEMM在LLM服务中的实际性能，并与代表性的量化方法进行比较。具体来说，我们基准测试了QServe 【15, Qserve: W4a8kv4 quantization and system co-design for efficient llm serving, 2024, arXiv】(W4A8)、TRT-W4A16 (W4A16)、TRT-W8A8 (W8A8)、TRT-FP8 (FP8)和TRT-FP16 (FP16)，其中TRT指的是NVIDIA开发的TensorRT-LLM【20, TensorRT-LLM: A TensorRT Toolbox for Optimized Large Language Model Inference, 2023, https://github.com/NVIDIA/TensorRT-LLM】。我们还考虑了Atom 【35, Atom: Lowbit quantization for efficient and accurate llm serving, 2024, Proceedings of Machine Learning and Systems】(W4A4)和QQQ【34, QQQ: Quality Quattuor-Bit Quantization for Large Language Models, 2024, arXiv】(W4A8)。然而，Atom在H800 GPU上性能较差，因为张量核心不支持INT4。QQQ的性能也不如QServe。因此，我们从进一步的评估中省略了Atom和QQQ。

3.1 屋顶线分析与实践之间的差距

基准测试设置。为了评估这些GEMM配置的实际性能，我们在H800 GPU上对LLaMA2-7B（一个密集模型）和Mixtral-8×7B（一个MoE模型）进行LLM服务基准测试，批量大小从4到256不等。我们为LLaMA2-7B选择W8A8量化，为Mixtral-8×7B选择FP8量化，因为W8A8量化目前不支持Mixtral-8×7B。我们考虑了两种输入输出长度设置：1) 1024个输入token和512个输出token；2) 128个输入token和128个输出token。需要注意的是，增加token长度不会影响解码期间FFN和投影（PROJ）层的GEMM工作负载，但会增加注意力计算。图4显示了端到端推理中GEMM延迟（来自FFN和PROJ层）的比例。我们观察到，在小批量大小时，GEMM主导了延迟，并且在LLaMA2-7B上，即使在大批量长序列的情况下，仍然占总延迟的20%以上。对于Mixtral-8×7B，由于需要为每个专家运行单独的GEMM，GEMM在所有测试案例中仍然是延迟的主要贡献者。这些结果突显了GEMM在LLM服务性能中的基础性作用。
图4：输入长度为128（左）和1024（右）的推理时间分解。由于内存不足，省略了长度1024、批次大小256的条形图。

性能差距观察。图5显示了解码期间平均每层的GEMM延迟。与基于屋顶线模型的预测相反，W4A8在小批量大小（$B \le 64$）时表现与W8A8相似，但在大批量大小（$B \ge 128$）时变得慢了近2倍，而此时它本应具有竞争力。值得注意的是，W4A8的性能甚至不如FP16和W4A16，这两种方法不涉及或仅涉及部分量化。在Mixtral-8×7B上仅报告了FP8和W4A16的结果，因为其他系统不支持该模型。Mixtral上的延迟也远高于LLaMA2-7B。总之，尽管理论上预期W4A8在内存受限的情况下应优于W8A8，在计算受限的情况下应与其性能相当，但我们的结果表明，现有的W4A8实现在实践中始终未能达到预期，尤其是在计算受限的条件下，这凸显了理论潜力与实际性能之间的明显差距。
图5：在批次大小从4到256的情况下，单个Transformer层的GEMM延迟。

3.2 深入探究GEMM处理过程

性能差距分析。为了理解性能差距，我们首先剖析了W4A8反量化的开销，然后开发了一个成本模型来捕捉关键性能因素。

QServe的反量化开销。我们关注QServe的主循环，因为$K$维度通常远大于瓦片大小$K_t$，并主导了整个GEMM成本。在每次迭代中，QServe使用寄存器级并行将权重从UINT4（$Q_{u4}$）反量化为INT8（$Q_{i8}$），每个32位寄存器处理四个元素。给定$Q_{i8}$和$Q_{u4}$，可以根据公式1离线计算出缩放因子$s_{u8}$和零点$z_{i8}$。为了避免在寄存器级算术中发生溢出，QServe应用了两种技术：1) 渐进式量化：它将$Q_{i8}$限制在[−119, 119]范围内，确保$Q_{u4} \cdot s_{u8}$保持在有效范围内；2) 乘法后减法：QServe推迟了减法操作，以避免乘以负值，计算$Q_{u4} \cdot s_{u8} - z_{i8} \cdot s_{u8}$，而不是在乘法前减去$z_{i8}$（如公式2所示）。尽管做出了这些努力，减法步骤仍然可能溢出。为了缓解这个问题，QServe依赖vadd指令来将封装在32位寄存器中的四个8位元素相加。然而，vadd不是原生硬件指令，它被降级为十几个低级操作，给CUDA核心带来了巨大的压力。由于CUDA核心和张量核心之间存在巨大的性能差距（见图1），这种开销成为瓶颈。NVIDIA Nsight对LLaMA2-7B的FFN层的剖析显示，涉及vadd的减法占了21%的warp停顿，这凸显了QServe反量化策略的性能成本。

成本模型。接下来，我们提出了一个成本模型，以捕捉第2节中描述的带有反量化的流水线GEMM执行的关键性能因素。假设瓦片大小为 $M_t \times N_t \times K_t$。那么，输出瓦片的数量为 $m \times n$，其中 $m = \lceil M/M_t \rceil, n = \lceil N/N_t \rceil$，并且每个瓦片需要 $k = \lceil K/K_t \rceil$ 次迭代来完成主循环。主循环的每次迭代包括两个阶段：数据加载和计算。我们首先分析单次迭代的成本，然后将分析扩展到完整的流水线执行。

数据加载。每次迭代的数据加载时间由公式3给出，其中$\phi_{BD}^x$表示加载类型为$x$的数据的块级吞吐量（元素/秒），基于线程块可用的有效内存带宽。在LLM服务中，激活矩阵通常很小，并从快速内存中重用，因此成本主要由从全局内存传输权重决定。

$$T_{\text{LD}} = \frac{(M_t \cdot K_t + N_t \cdot K_t)}{\phi_{\text{BD}}^{x}} \approx \frac{N_t \cdot K_t}{\phi_{\text{BD}}^{x}}$$

计算。计算阶段包括：1) 在CUDA核心上进行反量化；2) 在张量核心上进行MMA。因此，每次迭代的计算时间由以下公式给出：
$T_{\text{COMP}} = \frac{\alpha \cdot N_t \cdot K_t}{\phi_{\text{CUDA}}} + \frac{2 \cdot \min(M_t, M) \cdot N_t \cdot K_t}{\phi_{\text{TC}}^y},$
其中，$\alpha$是反量化一个权重元素所需的指令数，$\phi_{CUDA}$是块级CUDA核心吞吐量（OPs/s），$\phi_{TC}^y$是数据类型为$y$的块级张量核心吞吐量（OPs/s）。一个MAC（乘法累加）等于两个操作（一个乘法和一个加法）。在所有迭代之后，每个输出瓦片必须写回全局内存，这会产生一个结尾部分的成本。由于主循环通常占主导地位，我们省略了结尾部分的成本。

单瓦片执行。线程块计算一个输出瓦片的总时间$T_t$包括初始的流水线填充以及重复的重叠加载和计算。对于大的$k$，填充和排空开销可以忽略不计，因此$T_t$可以近似为公式5：

$$\begin{aligned} \begin{aligned} T_t &= T_{\text{LD}} + T_{\text{COMP}} + (k-1) \cdot \max(T_{\text{LD}}, T_{\text{COMP}}) \\ &\approx k \cdot \max(T_{\text{LD}}, T_{\text{COMP}}). \end{aligned} \end{aligned}$$

GPU级执行。假设一个设备有 $S$ 个流式多处理器，每个处理器最多可以同时运行 $L$ 个线程块。用 $\Phi_{BD}^x$（内存）、$\Phi_{CUDA}$（CUDA核心）和 $\Phi_{TC}^y$（张量核心）表示设备级吞吐量。由于 $M_t, N_t, K_t$ 很小，我们通常有 $K \gg K_t$ 和 $N \gg N_t$；$M$ 取决于批处理大小。给定 $m \times n$ 个总瓦片，总执行时间 $T$ 近似为：

$$T = \frac{m \cdot n}{S \cdot L} \cdot T_t = m \cdot \max\left(\frac{n \cdot k}{S \cdot L} \cdot T_{LD}, \frac{n \cdot k}{S \cdot L} \cdot T_{COMP}\right)$$

$$\approx m \cdot \max\left(\frac{N \cdot K}{\phi_{BD}^x \cdot S \cdot L}, \frac{\alpha \cdot N \cdot K}{\phi_{CUDA} \cdot S \cdot L} + \frac{2 \cdot \min(M_t, M) \cdot N \cdot K}{\phi_{TC}^y \cdot S \cdot L}\right)$$

$$= \left\lceil\frac{M}{M_t}\right\rceil \cdot \max\left(\underbrace{\frac{N \cdot K}{\Phi_{BD}^x}}_{T_{LD}}, \underbrace{\alpha \cdot \frac{N \cdot K}{\Phi_{CUDA}}}_{T_{DQ}} + \underbrace{\min(M_t, M) \cdot \frac{2 \cdot N \cdot K}{\Phi_{TC}^y}}_{T_{MMA}}\right),$$

其中$T_{LD}$、$T_{DQ}$和$T_{MMA}$分别表示数据加载、反量化和MMA的时间。为简洁起见，我们使用相同的符号$T_{LD}$表示每次迭代的数据加载时间。我们定义有效输出高度为$min(M_t, M)$，以考虑批处理大小小于瓦片大小的情况。该成本模型突出了GEMM性能如何受批处理大小$M$、硬件指标（$\Phi_{BD}$、$\Phi_{CUDA}$和张量核心吞吐量$\Phi_{TC}$）以及量化精度（权重位宽$x$和激活位宽$y$）的影响。

3.3 从剖析和分析中得出的见解

差距的根本原因。根据模型，在没有反量化开销的情况下，W4A8和W8A8在计算密集型场景中应表现出相似的性能，因为两者都使用INT8 MMA并共享相同的$T_{MMA}$。在内存密集型情况下，由于W4A8的内存加载量（$T_{LD}$）较低，预计其性能将优于W8A8。转换点发生在$T_{LD} = T_{MMA}$时，根据图1中的指标，H100上W4A8的批大小阈值为150，W8A8为300。这些结果与之前的基于屋顶线分析【34, QQQ: Quality Quattuor-Bit Quantization for Large Language Models, 2024, arXiv】【35, Atom: Lowbit quantization for efficient and accurate llm serving, 2024, Proceedings of Machine Learning and Systems】一致。然而，反量化改变了这一性能曲线。由权重矩阵大小决定的开销$T_{DQ}$由于CUDA核心（$\Phi_{CUDA}$）的计算能力有限以及处理溢出导致的高单位元素成本$\alpha$而变得显著。结果是，尽管W4A8的$T_{LD}$较低，但在内存密集型情况下其性能与W8A8相似，而在计算密集型场景中，其性能最高可慢2倍，如第3.1节所示。虽然人们可能期望通过增加批大小$M$来摊销$T_{DQ}$，但算术强度最终受到瓦片大小$M_t$的限制，而$M_t$又受共享内存的限制。这一限制阻止了$T_{DQ}$被有效隐藏，导致理论预期与观察到的性能之间存在显著差距。

对高效GEMM设计的启示。该成本模型为高效W4A8 GEMM提出了两个关键设计原则。首先，权重加载、反量化和MMA应在异构硬件单元（TMA、CUDA核心和张量核心）之间完全流水线化，以避免反量化造成的串行化瓶颈。其次，反量化必须是高度硬件高效的，以便能与其他阶段有效重叠。原则上，为了在内存受限场景中匹配权重加载的延迟（$T_{DQ} \le T_{LD}$），根据图1中的指标，H100上每个反量化元素的指令成本必须为$\alpha \le 5.07$。在计算受限设置中（$T_{DQ} \le T_{MMA}$），当$M = 150$时，此阈值变为$\alpha \le 5.05$。此外，CUDA核心还必须执行地址计算等辅助任务，这进一步增加了计算负载。这些约束共同凸显了在现代GPU上实现低开销反量化的挑战。

对LLM服务的影响。我们简要讨论硬件趋势如何影响LLM服务。在生产环境中，希望在小批量大小时达到计算密集型状态，以：1) 充分利用GPU计算能力；2) 减少请求延迟；3) 支持长序列；4) 最小化硬件故障等操作风险。此外，批量大小也受内存大小的限制。然而，如图1所示，张量核心的性能提升速度快于内存带宽，这使得内存到计算的转换点推向了更高的批量大小，根据我们的模型，A100上W8A8为156，H100上为300。相比之下，W4A8将这些阈值减半。这既凸显了量化在实现高效推理中的价值，也强调了高性能W4A8 GEMM内核的重要性。为此，我们提出了LiquidGEMM，一个用于高性能LLM服务的硬件高效W4A8 GEMM内核。在接下来的章节中，我们将介绍我们的量化算法，描述内核流水线设计和优化，并介绍一个用于评估的端到端LLM服务系统的实现。

A2 方法细节

4 量化算法

LiquidQuant (LQQ) 简介。为了解决反量化溢出问题，我们提出了LiquidQuant (LQQ)，这是一种硬件指令原生支持的硬件高效W4A8量化方案。

量化过程。为了提高低位量化的准确性，LQQ采用了一种分组量化策略【8, Gptq: Accurate post-training quantization for generative pre-trained transformers, 2022, arXiv】【14, AWQ: Activation-aware Weight Quantization for On-Device LLM Compression and Acceleration, 2024, Proceedings of Machine Learning and Systems】【15, Qserve: W4a8kv4 quantization and system co-design for efficient llm serving, 2024, arXiv】【34, QQQ: Quality Quattuor-Bit Quantization for Large Language Models, 2024, arXiv】【35, Atom: Lowbit quantization for efficient and accurate llm serving, 2024, Proceedings of Machine Learning and Systems】和一个将FP16权重转换为UINT4的两级量化框架。由于第一级反量化发生在GEMM的收尾阶段，开销可忽略不计，我们的重点是第二级量化。具体来说，遵循QServe【15, Qserve: W4a8kv4 quantization and system co-design for efficient llm serving, 2024, arXiv】的方法，第一级使用逐通道缩放将$W$量化为INT8张量$Q_{i8}$，如公式1所定义。我们还采用了第3.2节中的保护性量化范围，将$Q_{i8}$限制在[−119, 119]内，以防止反量化时缩放过程中的溢出（证明见【15, Qserve: W4a8kv4 quantization and system co-design for efficient llm serving, 2024, arXiv】）。第二级将INT8转换为UINT4。我们的关键思想是将$Q_{i8}$的对称范围移位到UINT8张量$Q_{u8}$的无符号域中，然后将$Q_{u8}$量化为$Q_{u4}$。这种设计与我们的反量化方法相匹配，以消除推理过程中潜在的溢出，我们将在本节末尾证明这一点。量化过程在公式7中定义。我们省略了零点$z_{u8}$，因为$min(Q_{u8})$和$min(Q_{u4})$都为零。
$Q_{u8} = Q_{i8} - \min(Q_{i8}), Q_{u4} = \left\lceil \frac{Q_{u8}}{s_{u8}} \right\rceil, s_{u8} = \frac{\max(Q_{u8})}{\max(Q_{u4})}$
与公式1中的标准量化相比，我们的方法引入了一个从$Q_{i8}$到$Q_{u8}$的简单移位，这完全是离线执行的。核心优化集中在在线反量化上，这对于高效的LLM服务至关重要。

反量化过程。根据公式7，我们在推理过程中将张量从UINT4反量化回INT8，如下所示：
$\widehat{Q}{i8} = \widehat{Q}).$} + \min(Q_{i8}) = Q_{u4} \cdot s_{u8} + \min(Q_{i8
为确保不发生溢出，我们必须保证此计算保持在有效的数值范围内。从公式7可知，缩放因子满足$s_{u8} \le \lfloor \frac{119 - (-119)}{15} \rceil = 16$。由于$Q_{u4} \in [0, 15]$，我们有$\hat{Q}{u8} = Q \le 15 \times 16 = 240$，这保持在UINT8范围内，避免了乘法过程中的溢出。} \cdot s_{u8

溢出问题与解决方案。然而，直接加上可能是负数的$min(Q_{i8})$可能导致环绕问题。我们用一个例子来说明。假设$Q_{u4} = 15$，$max(Q_{i8}) = 119$，以及$min(Q_{i8}) = -104$。那么，我们有$s_{u8} = \lfloor \frac{119 - (-104)}{15} \rceil = \lfloor \frac{223}{15} \rceil = 15$，期望的结果是：$\hat{Q}{i8} = Q} \cdot s_{u8} + min(Q_{i8}) = 15 \times 15 + (-104) = 225 - 104 = 121$。在二进制中，$\hat{Q{u8} = 225$表示为1110 0001，$min(Q$转换为INT8也是无效的，因为1110 0001在INT8中表示-31，而不是225。这个例子凸显了加法步骤需要超越标准硬件指令的仔细处理。}) = -104$以二进制补码形式表示为1001 1000。如果在没有类型提升的情况下在位级别执行加法，1110 0001 + 1001 1000 = 1 0111 1001，这会发生溢出。另外，在加法之前将$\hat{Q}_{u8

巧妙的反量化方法。LQQ引入了一种巧妙的反量化方法，结合移位量化，通过利用二进制补码的特性来消除溢出：一个INT8值$v$和一个UINT8值$u$如果$v \equiv u \pmod{2^8}$，则它们共享相同的二进制表示。例如，-3 $\equiv$ 253 (mod $2^8$)，两者都表示为1111 1101。利用这一特性，我们将公式8重写为：

$$\begin{aligned} \begin{aligned} \hat{Q}_{i8} &\equiv Q_{u4} \cdot s_{u8} + \min(Q_{i8}) + x \cdot 2^8 \pmod{2^8} \\ &\equiv Q_{u4} \cdot s_{u8} + (2^7 + \min(Q_{i8})) + (2x - 1) \cdot 2^7 \pmod{2^8}, \end{aligned} \end{aligned}$$

其中$x$是一个整数。我们接下来证明，通过适当地控制$x$的值，公式9中的计算可以避免溢出，即所有中间结果都保持在UINT8范围内。

无溢出证明。设$q_i$是第一级量化后$Q_{i8}$中的一个元素，设$q_u = q_i - min(Q_{i8})$是$Q_{u8}$中对应的元素。根据公式9，反量化计算过程可以表示为：

$$\widehat{q}_{i} \equiv \left\lfloor \frac{q_u}{s_{u8}} \right\rfloor \cdot s_{u8} + \underbrace{\left(2^7 + \min(Q_{i8})\right)}_{a \in [0,255]} + \underbrace{(2x-1) \cdot 2^7}_{b} \pmod{2^8}. $$

图6：显式粗粒度流水线（ExCP）和隐式细粒度流水线（ImFP）设计的比较。
我们首先证明$\hat{q}u + a$在UINT8范围内。因为$s) = 238$，我们有：} \le 16$且$q_u \le max(Q_{i8}) - min(Q_{i8

$$\begin{aligned} \widehat{q}_{u} + a = \left\lfloor \frac{q_{u}}{s_{u8}} \right\rfloor \cdot s_{u8} + a \le q_{u} + \frac{s_{u8}}{2} + a \\ \le (\max(Q_{i8}) - \min(Q_{i8})) + 8 + (2^7 + \min(Q_{i8})) \\ = \max(Q_{i8}) + 8 + 2^7 \le 119 + 8 + 128 = 255. \end{aligned}$$

接下来，为了确保最终结果$\hat{q}_u + a + b$也保持在[0, 255]范围内，我们如下控制$x$的值：如果$\hat{q}_u + a \ge 128$，则设置$x = 0$使得$b = -128$；否则，设置$x = 1$使得$b = 128$。这保证了公式10中的计算在UINT8内是无溢出的。 □

硬件高效计算。在运行时检查$\hat{q}u + a$并确定$x$会给GEMM的主循环带来显著开销。经过分析，我们观察到加$b$等同于翻转$\hat{q}_u + a$的最高有效位。因此，反量化可以执行为：
$\widehat{Q}80,$} = (Q_{u4} \cdot s_{u8} + a) \oplus 0\text{x
其中$a = 2^7 + min(Q_{i8})$是离线预计算的，$\oplus$表示异或操作。这种形式化使得所有中间值都保持在UINT8范围内，避免了溢出，并实现了高效的硬件执行（见第5.3节）。对于第一级反量化，LQQ遵循公式2中的标准过程。

5 高性能W4A8 GEMM内核

LiquidGEMM 简介。基于LiquidQuant (LQQ)，我们提出了LiquidGEMM，一个具有异步计算流水线的高性能W4A8 GEMM内核。我们使用当前云端主力GPU H800来说明该内核。为了优化执行，我们计算$C = (B A^T)^T$而不是$C = AB^T$，具体原因在5.4节中解释。

5.1 异步计算流水线设计

显式粗粒度流水线 (ExCP)。高性能GEMM库（如CUTLASS）使用warp专业化来重叠权重加载和计算。在此模型中，线程块内的warp被分为专门的角色，如加载warp和MMA warp，它们以生产者-消费者的方式异步操作。在H800上，warp被分组为warp组（WGs），每个组由四个warp（128个线程）组成，协同工作。在反量化上下文中，一个直接的想法是将其应用于W4A8计算。具体来说，如图6所示，我们设计了一个三阶段流水线，其中三个WG分别负责加载权重、执行反量化和执行MMA。每个阶段都映射到一个不同的硬件单元：通过TMA进行权重加载，通过CUDA核心进行反量化，通过张量核心进行MMA。这些阶段并发操作，使得$T_{LD}$、$T_{DQ}$和$T_{MMA}$能够重叠。我们将这种方法称为显式粗粒度流水线（ExCP）。然而，由于其粗粒度的执行和对warp组的显式调度，ExCP可能引入流水线气泡，从而降低GEMM效率。特别是，反量化WG从SMEM加载权重（这些权重之前由加载WG从GMEM加载），然后在RF中进行反量化。反量化后，它将权重写回SMEM，并通知MMA WG开始执行。这种在RF和SMEM之间的往返数据移动会产生不可忽视的开销，并增加反量化WG的工作负载，造成流水线停顿。此外，反量化WG和MMA WG之间的基于软件的同步增加了额外的开销。

隐式细粒度流水线 (ImFP)。为了解决ExCP的低效问题，我们提出了隐式细粒度流水线（ImFP）。与ExCP为反量化和MMA分配单独的WG不同，ImFP使用一个统一的计算WG（Compute WG）负责这两项任务。这消除了将反量化结果从RF写回SMEM的需要，减少了数据移动开销（图6）。为了重叠反量化和MMA，我们利用了不同Compute WG之间的流水线阶段。具体来说，ImFP采用单生产者-多消费者的细粒度流水线模型。Load WG作为生产者，将权重从GMEM加载到SMEM，并将其分割成细粒度的任务，每个任务是权重矩阵的一个片段。然后，这些任务由多个Compute WG动态获取和处理，每个WG都执行反量化和MMA。由于不同的计算WG操作于不同的任务，一个WG中的反量化自然地与另一个WG中的MMA重叠，从而在没有显式同步的情况下实现了隐式并行。在我们的实现中，每个线程块由一个Load WG和两个Compute WG组成，这有效地平衡了硬件利用率和任务吞吐量。实验结果表明，ImFP显著优于粗粒度的ExCP设计。接下来，我们将详细介绍数据加载和计算。

5.2 内存布局和数据加载

数据加载流程。在每个主循环迭代中，所需的权重瓦片由Load WG从GMEM加载到SMEM，然后由Compute WG加载到RF中进行反量化和MMA。在张量核心上的MMA需要一种结构化的数据布局，以便跨线程符合硬件内部指令的要求。为了满足这一需求，权重矩阵的内存布局至关重要，因为它直接影响数据加载的效率。

传统方法。现代GPU支持由硬件定义的固定矩阵形状上的MMA操作。对于INT8输入，H800提供了WGMMA.m64nNk32和WGMMA.m64nNk64等指令，其中$N$的范围从8到256。如图7a所示，WGMMA.m64nNk32在张量核心上执行一个64×$N$×32的MMA，需要一个来自矩阵$W$的64×32片段。一个WG中的每个warp加载一个16×32的段，每个线程使用跨步布局将16个元素取入寄存器：每组四个连续元素，间隔排列以匹配内部指令的瓦片模式。线程T0访问的元素在图7a中以深蓝色显示。为了从SMEM加载到RF，H800提供了ldmatrix指令。每个线程在一个事务中加载16个连续字节，并将每个4字节组散布到适当的线程——假设每个元素是1字节。这个假设对于W4A8不成立，因为元素被压缩到4位。结果，ldmatrix会错误地散布数据，例如，本应给T2和T3的元素可能会被传递给T1，如图7a所示。一种替代方法是使用LDS.32指令，它从指定地址加载32位。然而，每个线程只需要四个4位值，这意味着一半的数据是未使用的，从而降低了有效带宽。此外，这种方法需要更多的加载指令和额外的地址计算，增加了算术开销并给CUDA核心带来了额外的负担【15, Qserve: W4a8kv4 quantization and system co-design for efficient llm serving, 2024, arXiv】。

双MMA打包布局。受QServe【15, Qserve: W4a8kv4 quantization and system co-design for efficient llm serving, 2024, arXiv】中计算感知权重重排的启发，我们提出了双MMA打包布局来解决这个问题。在单次MMA操作中，每个线程需要16个UINT4元素，而粗粒度的LDS.128指令每次事务加载32个元素。为了利用这一差距，我们将每个线程连续两次MMA操作所需的数据打包并连续存储，如图7b所示。这使得每个线程可以使用单个LDS.128指令加载所有32个UINT4元素。为了满足WGMMA的片段布局，我们对权重进行重排，使得每个线程在两次MMA中所需的元素在内存中是相邻的。与QServe将权重存储在二维布局中不同，我们将这些元素排列成一维布局，以消除共享内存的银行冲突，并无需进行数据重组或复杂的数据打包。这种布局支持跨线程同时进行八次LDS.128操作，充分利用了共享内存带宽。此外，双MMA打包布局显著减少了加载指令的数量，并最小化了CUDA核心上的地址计算开销。GMEM中的权重矩阵遵循与SMEM中相同的布局，从而能够使用LDG.128（每个warp可用的最粗粒度加载指令）进行高效传输。由于布局转换是离线应用的，它不会引入运行时开销。

图8：使用位操作和IMAD指令的反量化过程，这些指令由硬件原生支持。$s = s_{u8}$和$z = a$是根据公式12离线计算的。

5.3 硬件高效的反量化

反量化流程。将权重从SMEM加载到RF后，每个线程在四个32位寄存器中持有32个UINT4元素，如图8所示。元素$w8-w15$对应于第一次MMA操作，$w16-w31$对应于第二次。我们使用LQQ（第4节）在CUDA核心上将这些权重从UINT4反量化为INT8。图8展示了反量化过程。我们首先应用QServe【15, Qserve: W4a8kv4 quantization and system co-design for efficient llm serving, 2024, arXiv】中的解包方法，将一个寄存器中的八个4位元素扩展到两个持有8位值的寄存器中。然后我们使用公式12执行反量化：乘以缩放因子$s_{u8}$，加上偏移量$a$，并应用最终的XOR。因为LQQ确保没有溢出，所有操作都可以使用原生的32位硬件指令执行，特别是用于乘加的IMAD和用于偏移校正的XOR。注意，$s_{u8}$和$a$都可以离线预计算。反量化后，得到的UINT8元素与目标INT8值共享相同的二进制表示，使得它们可以直接用于后续在张量核心上的MMA操作。

效率总结。总之，我们的方法仅用两条硬件算术指令就反量化了四个元素。包括解包步骤在内，八个元素仅用七条指令就完成了反量化，显著降低了CUDA核心的计算开销，远低于为实现与权重加载和MMA有效重叠所需的阈值（第3.3节）。第一级反量化被融合到GEMM的收尾阶段，开销可以忽略不计。

5.4 其他GEMM优化

硬件特定优化。如上所述，GPU MMA指令受限于硬件定义的固定矩阵形状。对于INT8，H800将$M$维度固定为64，而$N$可以在几种配置中从8变化到256。为了在小批量大小下更好地利用张量核心，我们通过将$C = AB^T$重写为$C = (BA^T)^T$来应用硬件特定的优化，这使我们能够根据批量大小选择WGMMA指令，并最大化计算效率。此外，我们采用了标准的GEMM优化，如持久化内核。由于这些技术被广泛使用，我们在此省略细节。

实现细节。利用CUTLASS和Cute的编程原语，我们将瓦片调度器、主循环和收尾阶段等组件集成并调整为一个warp专用的乒乓内核。具体来说，我们的反量化算法被融合到MMA主循环中，并在数据加载期间使用双MMA打包布局。我们用PTX实现了WGMMA指令、屏障同步和像TMA这样的通用组件，并由CUTLASS封装。相比之下，反量化逻辑是直接用CUDA实现的。

A7 补充细节

6 LLM服务系统和离线量化

系统集成。为了支持端到端性能评估，我们通过集成了关键系统级组件的开源技术，实现了一个LLM服务系统，这些组件包括注意力计算、KV缓存管理和量化方案。本节简要概述了它们的实现以及离线量化。

服务系统。图9展示了我们为LLaMA模型设计的LLM服务系统的数据流。Query、Key、Value、Output和FFN层都使用我们提出的LiquidGEMM，对权重和激活值进行W4A8量化，并产生FP16输出。我们遵循TensorRT-LLM【20, TensorRT-LLM: A TensorRT Toolbox for Optimized Large Language Model Inference, 2023, https://github.com/NVIDIA/TensorRT-LLM】的方法，使用逐通道静态量化将KV缓存量化为INT8，缩放因子离线计算。为了提高内存效率，我们采用PagedAttention 【12, Efficient memory management for large language model serving with pagedattention, 2023, SOSP】进行KV缓存管理，并使用FlashAttention-2【6, Flashattention-2: Faster attention with better parallelism and work partitioning, 2023, arXiv】进行运行时注意力计算。我们没有采用FlashAttention3【22, Flashattention-3: Fast and accurate attention with asynchrony and low-precision, 2024, NeurIPS】，因为它专为FP8设计。对于激活值量化，我们遵循SmoothQuant【29, Smoothquant: Accurate and efficient post-training quantization for large language models, 2023, ICML】，在除以平滑尺度后，通过逐token量化，即时将FP16激活值动态映射到INT8。由于激活张量的内存占用小且计算开销低，量化是轻量级的，并且通常融合到其他内核中。
图9：我们用于LLaMA模型的LLM服务系统中的数据流概览。

离线量化。我们采用SmoothQuant【29, Smoothquant: Accurate and efficient post-training quantization for large language models, 2023, ICML】的训练后量化方法来离线量化权重。具体来说，权重首先通过一个平滑因子进行缩放，然后使用第4节中描述的两级方法进行量化：从FP16到INT8的逐通道量化，然后是到UINT4的逐组量化。我们遵循OutlierSuppression+【28, Outlier suppression+: Accurate quantization of large language models by equivalent and optimal shifting and scaling, 2023, arXiv】，应用网格搜索来确定最佳的平滑尺度。需要注意的是，我们的重点是优化W4A8 GEMM的效率；我们的方法与提高量化准确性的技术是正交的，并且可以与这些方法无缝集成。

8 相关工作

LLM推理的量化。这些方法通常分为两类：仅权重量化和权重-激活量化。对于仅权重量化，GPTQ【8, Gptq: Accurate post-training quantization for generative pre-trained transformers, 2022, arXiv】开创了通过使用近似二阶信息将权重压缩到3或4位的亚8位量化。AWQ【14, AWQ: Activation-aware Weight Quantization for On-Device LLM Compression and Acceleration, 2024, Proceedings of Machine Learning and Systems】通过结合激活统计数据来识别和保留关键权重，进一步提高了准确性。对于权重-激活量化，GPT3.int8()【7, Gpt3. int8 (): 8-bit matrix multiplication for transformers at scale, 2022, NeurIPS】引入了混合精度量化，以在单独的16位乘法中隔离异常激活值。SmoothQuant【29, Smoothquant: Accurate and efficient post-training quantization for large language models, 2023, ICML】提出通过数学上等价的变换将量化挑战从激活迁移到权重，有效平滑了激活异常值。Atom【35, Atom: Lowbit quantization for efficient and accurate llm serving, 2024, Proceedings of Machine Learning and Systems】采用了混合精度细粒度组量化，实现了吞吐量和准确性之间的平衡。OmniQuant【23, Omniquant: Omnidirectionally calibrated quantization for large language models, 2023, arXiv】提供了自动学习最优量化参数的方法。QServe【15, Qserve: W4a8kv4 quantization and system co-design for efficient llm serving, 2024, arXiv】利用了一种为GPU张量核心优化的两阶段W4A8KV4量化方法，而QQQ【34, QQQ: Quality Quattuor-Bit Quantization for Large Language Models, 2024, arXiv】结合了自适应平滑和基于Hessian的补偿，开发了一种新的W4A8 GEMM内核。其他近期工作，包括QuaRot【1, Quarot: Outlier-free 4-bit inference in rotated llms, 2024, arXiv】、SpinQuant【18, SpinQuant–LLM quantization with learned rotations, 2024, arXiv】和DuQuant【13, Duquant: Distributing outliers via dual transformation makes stronger quantized llms, 2025, NeurIPS】，应用了变换（如旋转）来有效分布异常值。其中，DuQuant与SpinQuant相比简化了训练复杂性，并表现出优于QuaRot的性能。与这些工作不同，本文专注于高效LLM服务的W4A8 GEMM的效率。

LLM训练的量化。量化感知训练（QAT）比训练后量化（PTQ）能获得更高的准确性，但由于其高计算开销，使用有限。LLM-QAT【17, Llm-qat: Data-free quantization aware training for large language models, 2023, arXiv】引入了一种使用预训练模型生成的数据进行无数据蒸馏的方法，使得对LLM进行实用的QAT成为可能。EfficientQAT【4, Efficientqat: Efficient quantization-aware training for large language models, 2024, arXiv】通过两阶段训练策略加速了QAT。Bondarenko等人【3, Low-Rank Quantization-Aware Training for LLMs, 2024, arXiv】提出了一种专为LLM设计的轻量级、内存高效的低秩QAT方法。EdgeQAT【24, EdgeQAT: Entropy and Distribution Guided Quantization-Aware Training for the Acceleration of Lightweight LLMs on the Edge, 2024, arXiv】引入了一种熵引导的方法，具有自适应的token重要性，以减少QAT中的信息失真。

LLM服务。Orca【32, Orca: A distributed serving system for {Transformer-Based} generative models, 2022, OSDI】通过迭代级调度和选择性批处理优化了服务性能。vLLM【12, Efficient memory management for large language model serving with pagedattention, 2023, SOSP】通过受虚拟内存机制启发的PagedAttention提高了KV缓存管理的效率。NVIDIA的TensorRT-LLM【20, TensorRT-LLM: A TensorRT Toolbox for Optimized Large Language Model Inference, 2023, https://github.com/NVIDIA/TensorRT-LLM】提供了一个专门为加速GPU上的LLM推理而优化的开源库。DistServe 【36, Distserve: Disaggregating prefill and decoding for goodput-optimized large language model serving, 2024, arXiv】通过使用先进的放置算法解耦预填充和解码计算，增强了服务性能。COMET【16, COMET: Towards Partical W4A4KV4 LLMs Serving, 2024, arXiv】提出了一个混合精度推理框架，其中包含了新颖、高度优化的内核，以最大化LLM推理性能。

A4 实验环境

研究对象: 本文提出的W4A8内核LiquidGEMM，以及包含该内核的LLM服务系统LiquidServe。
基线系统:
- QServe【15, Qserve: W4a8kv4 quantization and system co-design for efficient llm serving, 2024, arXiv】: 一个先进的W4A8 LLM服务系统。
- TensorRT-LLM【20, TensorRT-LLM: A TensorRT Toolbox for Optimized Large Language Model Inference, 2023, https://github.com/NVIDIA/TensorRT-LLM 】 (v0.16.0): NVIDIA官方的LLM推理框架，测试了多种精度配置（FP16, W4A16, W8A8, FP8）。
模型架构: 测试了多种模型，包括 LLaMA (v1, v2, v3)【9, The llama 3 herd of models, 2024, arXiv】【25, Llama: Open and efficient foundation language models, 2023, arXiv】【26, Llama 2: Open foundation and fine-tuned chat models, 2023, arXiv】, Mistral-7B【10, Mistral 7B, 2023, arXiv】, Mixtral-8×7B【11, Mixtral of experts, 2024, arXiv】, และ Yi-34B【31, Yi: Open foundation models by 01. ai, 2024, arXiv】。
数据集 (用于准确性评估):
- WikiText2【19, Pointer sentinel mixture models, 2016, arXiv】: 用于评估困惑度 (perplexity)。
- PIQA【2, Piqa: Reasoning about physical commonsense in natural language, 2020, AAAI】, ARC【5, Think you have solved question answering? try arc, the ai2 reasoning challenge, 2018, arXiv】, HellaSwag【33, Hellaswag: Can a machine really finish your sentence?, 2019, arXiv】, WinoGrande【21, Winogrande: An adversarial winograd schema challenge at scale, 2021, Commun. ACM】: 用于评估零样本准确率。
硬件配置:
- GPU: NVIDIA H800 (80 GB 显存)。
- CPU: Intel Xeon Platinum 8457C。
- 内存: 2.9 TB RAM。
软件配置:
- OS: Linux。
- 框架/库: PyTorch 2.4.0, CUDA 12.4。
- 实现语言: LiquidGEMM 使用 CUDA 和 PTX 实现。
- 基准测试框架: 使用一个统一的基于CUDA的框架来隔离和公平地比较各个GEMM内核的性能。

A4 实验结果

LLM服务效率对比

内存约束下的吞吐量 (表1): 在80GB显存限制下，LiquidServe 实现了最高的峰值吞吐量，尤其是在大型模型上。与 QServe 相比，LiquidServe 能够随着批处理大小的增加持续扩展性能。例如，在LLaMA2-70B上，LiquidServe 通过支持更大的批处理大小，比 TRT-W8A8 实现了3.16倍的加速；由于其更高的INT8 MMA计算吞吐量，比 TRT-W4A16 实现了1.63倍的加速。为了隔离 LiquidGEMM 的贡献，将 LiquidServe 中的内核替换为 QServe 的内核（LiquidServe/wo），LiquidServe 仍能实现1.13-1.98倍的端到端加速，证明了 LiquidGEMM 的有效性。

表1：LiquidServe、QServe和TRT在H800上80GB内存限制下的峰值token生成吞吐量（tokens/s）。括号中的数字表示达到峰值吞吐量时的批次大小。加速比是相对于QServe和TRT中表现最佳的基线报告的。LiquidServe/wo使用了QServe的W4A8 GEMM内核。

端到端LLM服务时间分解 (图10): LiquidServe 在所有测试模型上的GEMM延迟都与基线相当或更优。在LLaMA2-7B上，其GEMM延迟最低，比 QServe 快1.90倍，比TRT快1.58倍。这证实了 LiquidGEMM 的硬件友好型反量化设计有效缓解了CUDA内核瓶颈。

图10：在表1指定的批次大小下，处理LLaMA2-7B、LLaMA2-70B、LLaMA3-8B和Mistral-7B一个解码层的时间分解。

固定批次大小下的吞吐量 (图11): 在固定的批次大小（内存受限的16和接近计算受限的128）下，LiquidServe 在LLaMA2-7B和LLaMA2-70B上始终优于所有基线系统，展示了其效率优势。

GEMM内核效率对比

效率比较 (图12): 在内核级别的基准测试中，随着批处理大小的增加，QServe 的性能显著下降，而 LiquidGEMM 保持了持续的低延迟。在批处理大小为256时，LiquidGEMM 在LLaMA2-7B, 13B, 和 70B上分别比 QServe 实现了2.75倍、2.87倍和2.90倍的加速。对于Mixtral-8×7B，在批次大小超过32后，LiquidGEMM 比 TRT-FP8 实现了1.41-1.84倍的加速，比 TRT-W4A16 实现了1.12-2.53倍的加速，表现出强大的可扩展性和鲁棒性。

消融研究 (图13): 实验结果验证了 LiquidGEMM 各个组件的有效性。
- LQQ: 在计算密集型的大批次场景下，LQQ算法带来了高达1.29倍的加速。
- ExCP vs ImFP: 显式粗粒度流水线（ExCP）在小批次时由于往返通信和同步开销导致性能下降。相比之下，隐式细粒度流水线（ImFP）在所有批次大小下都能持续提升性能。这证明了ImFP流水线策略的优越性。

图13：LiquidGEMM的消融研究，首先启用LQQ，然后分别启用显式粗粒度流水线（ExCP）或隐式细粒度流水线（ImFP）。

A5 结论

本文成功解决了在LLM服务中W4A8量化的反量化瓶颈问题。通过对现有W4A8 GEMM内核的性能进行剖析并建立成本模型，文章揭示了关键的性能制约因素。基于此分析，本文提出了LiquidGEMM，一个硬件高效的W4A8 GEMM内核。LiquidGEMM集成了两项协同设计的关键技术：

LiquidQuant: 一种无溢出的反量化算法，其硬件效率极高。
隐式细粒度流水线: 一种能够最大化GPU子系统间并行度的执行机制。

实验结果强有力地证明了LiquidGEMM的有效性。与先前的W4A8内核相比，LiquidGEMM在内核层面实现了高达2.90倍的加速，在系统层面实现了高达4.94倍的加速。与NVIDIA TensorRT-LLM相比，也取得了1.12-1.63倍的性能提升。这些成果表明，通过硬件感知的协同设计，W4A8 GEMM可以变得既高效又可扩展，为高性能LLM推理提供了强大的支持。

PaperCache

LiquidGEMM: Hardware-Efficient W4A8 GEMM Kernel for High-Performance LLM Serving

LiquidGEMM: Hardware-Efficient W4A8 GEMM Kernel for High-Performance LLM Serving

A1 主要贡献