LLM领域首次实现量化推理自由，效果和性能双SOTA！字节开源ABQ-LLM_新浪财经

字节开源的 ABQ-LLM （项目地址：https://github.com/bytedance/ABQ-LLM）是面向 AIGC 领域的算法系统协同优化工作，旨在解决 LLM 量化中存在的两大主要挑战：

低位量化致使的效果严重降低；
主流 GPU 对整数矩阵乘支持有限所引发的计算效率低下问题。

在推理引擎层面，ABQ-LLM 基于二值化矩阵乘(BTC) 等价重构了任意精度组合的矩阵乘，从而突破了 INT4/INT8 计算单元的限制，开创性地实现了任意比特量化组合的直接加速，结合深度工程优化，实现了 decoding 阶段对 cutlass/cublas 加速库的大幅性能领先。在算法方面，ABQ-LLM 运用了基于 transformer block 的分布纠正和低比特位平衡策略，有效提升低位量化效果。在各类量化配置下均达成了 SOTA 的效果，综合模型效果优于 OmniQuant、AffineQuant 等前期工作，同时工程实现大幅超越 SmoothQuant，端到端综合性能超越 SmoothQuant 实现了 1.6x 的推理加速和 2.7x 的内存压缩。

图 1. (上) 算子级性能对比，ABQKernel 可以有效地将位宽收益转化为加速收益，以 W2A8 为例，相比 cutlass W8A8 具备 3.8~7.48 倍的单核加速；(下) 端到端性能对比，在 fastertransformer 框架中，ABQ-LLM 相比业界 SOTA 的 SmoothQuant 方案实现了 1.6x 的推理加速和 2.7x 的内存压缩。

概述

大型语言模型（LLM）为自然语言处理带来了革命性的变化，但在推理过程中却受到大量内存和计算需求的限制。PTQ 提供了一种很有前景的解决方案，但它也面临着一些挑战，例如低位量化导致的性能下降，以及主流 GPU 的整数支持有限导致的计算效率低下。

字节智能创作团队提出的任意位量化方案 ABQ-LLM 是一种新型的 PTQ 方法，可以解决上述问题。ABQ 支持任意位量化权重和激活，并使用二进制表示执行矩阵计算。这种方法有两大优势：它能在不损失精度的情况下实现任意位量化矩阵计算，并绕过 INT4/INT8 单元的限制，利用二值化矩阵乘单元实现高效计算和高吞吐量。

为了解决低位宽时的性能下降问题，他们采用了两种增强技术：一种是 transformer block 的分布校正方法，用于解决权重和激活全量化过程中的分布偏移问题；另一种是位平衡策略，用于缓解超低位宽（如 2 比特）时的非对称分布问题。实验结果表明，ABQ 在各种量化设置（包括 W2A8、W4A4 和 W6A6）下都具有出色的灵活性和具备竞争力的性能。

介绍

大规模语言模型（LLM）对计算/带宽/内存的需求日益增长，其在实际应用场景中的部署受到了极大限制。后训练量化（PTQ）成为解决计算和带宽瓶颈的一种有效手段。然而当前业界的量化方法，无论是仅针对权重的量化(weight-only)，还是权重和激活值的全量化(weight-activation)，都无法突破主流 GPU 整数运算单元的限制。在推理阶段，weight-only 量化方法仍然需要将数据反量化为 FP16 进行计算。对于 weight-activation 量化，如图 2 所示，主流推理引擎局限于 INT4/INT8 计算单元的能力，在实际计算时仍需要将 W/A 精度转换为 INT4/INT8，引入额外的反量化指令耗时，这导致 LLM 领域量化算法的探索局限在 W4A4 和 W8A8，极大地限制了算法设计空间。

图 2. 在 NVIDIA GPU 上，非 W4A4/W8A8 位宽组合需要离线或者在线反量化至 4/8Bit，以满足 INT4/8 TensorCore 的要求。

此外，现有的量化方案在低比特位精度下仍存在显著的性能劣化。例如，SoothQuant、OmniQuant、AffineQuant 等方法在 W4A4、W3A4、W2A16 和 W2A8 等低比特位设置下表现出较大的精度下降，从而影响了量化模型的实际使用体验。为了解决上述问题，他们提出了 ABQ-LLM。据知，这是首个在大规模语言模型（LLM）中实现量化推理自由的研究工作。主要贡献包括：

在 LLM 领域中首次实现了量化推理自由。巧妙地基于 BTC（Binary TensorCore，BTC）等效重构了任意精度组合的量化矩阵乘法，突破了 INT4/INT8 计算单元的限制，使每一位压缩的位宽收益都能够转化为实际的加速收益，并有效规避了 GEMV 问题，充分挖掘了量化模型在更多混合精度（如 W2A6、W2A8）下的综合优势。
引入了 Transformer 模块的分布纠正方法，以缓解全量化处理 weight 和 activation 时导致的分布差异，从而提升低比特位下的模型性能。针对极低比特位（如 2 比特）下出现的非对称分布问题，采用位平衡策略，以弥补 W2 量化带来的性能损失。

Method

分布矫正策略（Distribution Correction）

图 3. LLaMA-7B transformer block 各模块量化带来的精度损失（Wikitext2 PPL 越低越好）

LLM 模型在量化时不同层的敏感性存在显著差异，某些层对 LLM 的量化性能起着决定性作用。为验证这一观点，如上图 3 所示，他们在 weight-activation 全量化场景下对 LLaMA-7B 模型 Transformer Block 中的不同组件进行了量化。当量化 Attention 和 MLP 中的 gate_proj 层和 up_proj 层时，模型性能仅出现了轻微下降。然而，一旦对 down_proj linear 层进行量化时，模型性能出现了显著的下降，表明对 down_proj 量化的处理是性能恢复的关键难点。字节智能创作技术团队的同学还进一步观察到，down_proj 量化导致性能下降的主要原因在于对 down_proj 激活的量化。由于在 INT4、INT3 和 INT2 等低比特位下表示范围受限，量化后模型的分布相较于全精度表示发生了显著偏移。因此，在量化校验过程中，他们对 down_proj 的输出建立了双余弦相似度损失，以纠正量化后模型 down_proj 输出的分布。具体的损失函数如下：

其中，

代表第 i 个 Transformer Block量化后的输出，

代表第 i 个 transformer 块全精度的输出，

代表第 i 个transformer 块全精度的输出，但其输入来自第 i−1个 transformer 块量化后的输出。此外，我们进一步分析了LLaMA-7B模型中激活值经过decoder block后的输入和输出余弦相似度变化。观察结果表明，模型首尾几个块的激活输入和输出余弦相似度存在显著差异，这表明这些首尾块对模型推理性能有较大影响。针对这一现象，他们对首尾块的down_proj层施加了vector分布补偿，以弥补这些块之间的分布差异。

其中，

代表 round-to-nearest operation, n代表比特数,

代表 step size , z 是 zero point.

是量化后的权重， W是全精度的权重，a和b是分布补偿向量，b用于控制是否进行补偿，b=1表示进行补偿，b=0则不进行补偿。

图4. 第一个Transformer Block（上图）和最后一个Transformer Block（下图）的注意力分布图

为了进一步提升量化模型的性能，他们研究了量化前(a)后(b)Attention Map的分布变化，如图4 所示，在全精度模型中，大量注意力集中在序列的首token上，这表明在文本生成过程中，首token起到了关键的引导作用，这一现象与LLM-QAT中的结论一致。然而，量化后模型的注意力分布显著改变，文本序列对首token的关注度受到了干扰。为了解决这一问题，并在量化过程中恢复模型对首token的关注，他们引入了attention-aware KL散度，以重建attention map 注意力分布，经过AKL损失重建我们的ABQ量化后与全精度保持了相似的Attention Map分布。

其中，

代表第 i 个 Transformer Block 量化后attention map输出，

代表第 i 个 Transformer Block全精度的attention map输出。

图5. ABQ-LLM 算法概述。ABQ-LLM 使用 DLC loss 和 AKL loss 来更新可学习的量化参数

最后，将 DLC 损失和 AKL 损失结合起来，如图5 所示得到最终的优化目标如下：

其中，

为校验完成后第 i 个Transformer Block的量化参数。当量化输出和全精度输出的分布相匹配时，总体损失接近于 0，这将有效地指导量化过程。

位平衡策略（Bit Balance Strategy）

预训练的LLM模型权重分布趋近于正态分布，具有对称性。然而，在普通的INT2量化配置下，数值表示范围仅包含四个值，进行对称量化时可选的表示范围为{-2, -1, 0, 1}或{-1, 0, 1, 2}，这破坏了模型权重原有的对称分布属性。这种不对称性导致模型性能显著下降，从W4A16到W3A16性能下降了0.46，而从W3A16到W2A16性能下降了4.87，出现了大幅异常下跌。为了解决这种不对称性对LLM模型量化的影响，他们借鉴了BitNet1.58的位平衡策略，对INT2量化进行了改进，将INT2对称量化的表示空间扩展为{-2, -1, 0, 1, 2}。通过这一改进，模型性能恢复到7.50，与W3A16性能相比，属于一个合理的变化范围。

计算等价重构（Reconstructing Arbitrary Bit Computation）

NVIDIA 在图灵以上架构中引入了 INT1 TensorCore 用于支持业界 W1A1 量化算法的发展，由于完全基于位运算，因此其峰值算力高达 INT4/INT8 的 4/8 倍。然而，由于巨大的效果损失，1Bit 量化至今未实际落地，该硬件单元始终未得到业界充分利用。他们在数学层面对量化矩阵乘进行分析，发现任意量化组合的运算均可分解为 1-bit 的矩阵乘的特殊叠加。

假设神经网路的某一层的权重W被量化为p bit，输入激活值X被量化为q bit，W和X执行矩阵乘运算得到32 bit的输出Y=WX，关键观察是，权重W和激活值X的任意位置的标量数值都可以分解为一系列 1 bit标量数字，任意精度组合的标量运算可以被分解为 1 bit运算和移位操作。形式上，为了支持标量级别的任意精度组合运算，例如 1 bit权重 𝑤 和 2 bit 激活值的标量级任意精度计算 𝑤𝑥，我们可以首先进行形式变换，例如2 bit的x可以被表示为

其中

他们使用OP(a,b)来表示输入是1bit数据，输出是32-bit的计算操作，那么原来的标量级任意精度计算 𝑤𝑥可以被表示为

上述过程可以自然推广到具有任意位 𝑝 和 𝑞 的矩阵乘组合。形式上，给定一个𝑝位权重矩阵𝑊和一个𝑞位权重矩阵𝑋，我们可以首先分解为1位矩阵

和

假如以BMMA代表1 bit矩阵乘运算，那么可以调用BMMA运算p*q 次，计算一系列的1-bit矩阵乘分量

最终，他们将所有的1-bit矩阵乘分量进行位叠加的缩放系数处理，之后进行累加得到32-bit的输出矩阵

经过上述变换过程，他们将任意量化组合的运算分解为1-bit的矩阵乘的特殊叠加，使底层可以调用高算力的BMMA指令实现，主要收益如下

有效解决GEMV的计算效率问题，由于底层计算时真实M维度是BatchSize * X_BIT，X_BIT=8时BatchSize无论多大均是8的整数，有效提升Decoding阶段(BatchSize=1)的计算密度。
首次以一种统一的方式支持任意Bit的量化组合计算，不但可以根据模型规模合理选择不同的量化组合，同时底层实现高效。

推理引擎实现（Engine Implementation）

图6. (上) 推理引擎计算任务拆解图. p和q分别代表输入X和权重W的量化位宽，数据流动被精心设计以实现高效计算：GL->SMEM->FR->SMEM->GL. (下) ABQKernel构成transformer block. ReQuant 和 DeQuant 分别表示在线量化和反量化操作，BitPacking 表示激活的在线布局转换操作

NVIDIA GPU 拥有许多称为流多处理器 (SM) 的处理单元，并使用大量线程并行执行计算任务。线程被构造为线程块，线程块成为 SM 上最小的调度执行单元。因此，计算目标被分解和映射到每个线程块上(称为Thread Block Tile) ，以实现并行计算。如图6 (a) 所示，对于形状为 M×N×K 的 GEMM 任务，每个线程块负责计算一个 BM×BN 输出块，该输出块被分解为形状为 BM×BN×BK 的 K/BK 子 GEMM 任务。我们的推理引擎将位宽配置为{p，q} 的量化矩阵乘法转换为 pq 个二值化矩阵乘法的特殊累积，因此Thread Block Tile的真正计算任务是 pBM×q*BN，整体计算流水如下所示：

首先，为了提高内存访问的连续性，我们提出BitPacking策略，将量化张量分解为n个二进制矩阵，其中n为量化位宽。以输入X为例，这意味着它的bit视角的内存布局由[M，K，p]变为[p，M，K]。线程块内的所有线程共享相同的共享内存空间(shared memory, SMEM)。在每个线程块内，线程进一步组织成warp，每个warp由32个连续的线程组成。

接下来，warp协作地从全局内存(global memory, GL)，加载当前的Thread Block需要计算的Thread Block Tile块所需的A矩阵（pBM×BK）和B矩阵（BK×qBN）数据并缓存在SMEM中。得益于BitPacking，读取pBMBK单比特row-major的Tile块和写入pBM*BK bit到SMEM的过程是高效且连续的。

随后，由于thread block包含多个warp，因此可以进一步将Thread Block Tile分解为Warp Tile，实现warp级并行，每个warp的计算任务为WM×WN。在计算准备阶段，将A矩阵（WM×WK，row-major）和B矩阵（WK×WN，col-major）独立从SMEM加载到FR，随后将计算分解为WARP_M_TILES*WARP_N_TILES个 Tensor Core MMA(matrix-multiply-accumulate)运算。由于A和B是二值化矩阵，因此我们实际使用的是Binary TensorCoreMMA（BMMA），其计算能力分别比INT8和INT4 TensorCore提升8倍和4倍。

所有warp协同完成Thread BlockTile计算，结果存储在每个warp的c fragment中，因此每个warp需要独立将计算结果写回SMEM。

输出图块（pBM × qBN）在SMEM中执行全局缩减求和以获得最终的正确结果（BM × BN），其中每个 BM × BN 子图块需要乘以某个缩放因子。我们称此过程为 Bit Reduction。

最后一步，warp 协同地从 SMEM 加载最终结果并写回 GL 中的目标位置。

他们将上述计算过程实现为一个名为 ABQKernel 的 GPUKernel。如图6 (b) 所示，ABQKernel 用于替换解码器层中的所有 gemm 操作，结合必要的 BitPacking、量化和反量化辅助操作，实现了 LLaMA 模型的任意量化推理。我们通过将辅助运算符融合到现有运算符中来精心管控量化运算符的开销，并且权重 BitPacking 是离线实现的，以提升效率。此外，我们还运用了多种优化手段，包括 GEMV 消除、Bank Conflicts 消除、Computational and Pipeline 优化和 AutoKernelSearch。

GEMV 运算消除（Eliminate GEMV）

图7. GEMV实质性转化为GEMM问题

LLM 类模型中80%的计算量和参数访存集中在GEMM/GEMV运算，解码阶段的GEMV是核心性能瓶颈。不失一般性地说，GEMV/GEMM 操作都可以使用 M、N、K 表示，其中两个相乘矩阵的大小为 M × K 和 K × N。当M=1时，GEMM问题会转换为GEMV问题，同时算子会从计算密集型转变为访存密集型。为了使用Tensor Core进行加速计算，通常会沿M维按BM=8进行分块，若M<8则进行0填充，导致冗余计算。

基于BTC的任意精度量化推理，通过将矩阵分解成若干个单比特的矩阵分量，把多个单比特的GEMV问题转换回GEMM问题，在利用Tensor Core的同时，减少甚至避免padding导致的冗余计算。

如图7 所示的一个w2a8算子实现，相比于直接将BM填充至8再调用Tensor Core的方式，我们可以将X矩阵分解成8个子矩阵，若设M=1，BM=1，WM=MMA_M=8，padding的占比从7/8=87.5%降低至0%。在其他精度的情况下，亦可以根据M的大小，通过调整BM的大小，减少甚至避免padding导致的冗余计算，提高Tensor Core的利用率。

BankConflict 消除（BankConflict Free）

图8. 共享内存内发生4路bank冲突的访存模式，E表示int32的一个元素，B表示bank，T表示一个线程.

BankConflicts的存在是影响访存效率的一大因素，图8 展示了Bank冲突发生的场景(BM=8,BK=512)，其中E代表int32视角的元素（Element)，B代表Bank，T代表Thread。数据从全局内存拷贝至共享内存时，为了充分利用内存带宽，每个线程以16Bytes进行内存拷贝，此时8x512bit数据会分成4个阶段（Phase)，每个阶段读取128Bytes（即图中连续两行数据），所有阶段都不会发生Bank冲突。数据从共享内存拷贝至寄存器时，由于8x8x128的BMMA以32-bit寄存器来存储矩阵A（8x128 bit）和矩阵B（8x128 bit)，每个线程需要拷贝4Bytes数据，拷贝一个矩阵块（同颜色块）时，线程T0~3、T8~11、T16~19、T24~27均会访问B0~3，T4~7、T12~15、T20~23、T28~31均会访问B16~19，此时发生4路Bank冲突，降低了访问效率。

如图9所示，通过采用swizzle操作解决了这一问题。如上所述8x512bit数据会分成4个阶段（Phase)，每个阶段8个线程完全访问32个Bank(128 Bits)，每个线程整体负责16Bytes访存，因此将连续两行作为一个整体进行坐标标定，每个线程负责访存的数据块(同颜色的16Bytes块)坐标为(irow, icol)，例如T0负责访问的数据块为(0, 0)，T7负责访问的数据块为(0，7) ，同理T8和T15访问的数据块坐标也是(0, 0)和(0，7) ，在全局内存拷贝至共享内存时，对坐标进行swizzle，本质上是应用new_icol = irow ^ icol的坐标变换，使同一个矩阵的8x128bit数据错落在不同的32个Bank中。在共享内存拷贝至寄存器时，T0~31访问的数据分别来自不同的Bank，访问不同的矩阵块均不会发生Bank冲突，进而提高访存效率。

图9. 利用swizzle实现共享内存内无bankconflicts.

计算流水优化（Computational and Pipeline Optimization）

图10. 计算流水线的时空图.

上图展示了在SM80+架构上进行的计算流水优化。在共享内存级别，通过cp.async指令异步执行全局内存到共享内存的写入。在处理第一个循环前，发出同步来确保第一个循环所需要的数据TILE-0已经准备在共享内存中。在处理第一个循环时，第二个循环所需的数据TILE-1同时也会往共享内存当中写入。从而在计算第一个循环时掩盖掉第二个循环所需共享内存的访存耗时。

在寄存器级别，当k=0时，共享内存的TILE-0数据会被加载进第一组寄存器A0、B0中，同时k=1所需的数据也会同时进行预加载至寄存器A1、B1。寄存器A0、B0的数据准备就绪后，会调用bmma对A0、B0进行计算。当k=1时，bmma所需的数据A1、B1已经在k=0时预加载，此时对k=2所需的数据进行预加载至寄存器A0、B0，如此循环，从而实现在Tensor Core bmma计算的同时将数据从共享内存写入到寄存器，完成在bmma计算时掩盖掉寄存器的访存耗时。

自动算子搜索（Auto Kernel Search）

为了在所有层级上更高效地利用GPU内存，在计算过程中采用分块策略至关重要。块大小对GEMM性能至关重要。在启动任意精度推理操作之前，他们会对不同的块大小进行性能测试，以选择最佳的实现方法。

对于经典的GEMM任务，M × N × K问题通常在线程块级别上被分块为BM × BN × BK，每个线程块在warp级别上进一步被分块为WM × WN × WK。每个warp基于GPU模型支持的Tensor Core块大小获取MMA_M × MMA_N × MMA_K ，然后在获得这三个层级的块大小后，搜索不同的块形状。在任意精度操作中，还需要考虑权重比特数(q)和激活比特数(p)，因此搜索空间比经典GEMM更大。给定BTC块大小为MMA_M = 8, MMA_N = 8, MMA_K = 128，权重维度warps的数量为W_WARPS_NUM = BN × q / WN，激活维度warps的数量为X_WARPS_NUM = BM × p / WM，实际warps数量为1 ≤ X_WARPS_NUM × W_WARPS_NUM ≤ 32。为了减少搜索空间，他们将WK = MMA_K = 128固定值，并将BK的长度设置为{128, 256, 384, 512}。由于单个线程块的共享内存和寄存器使用受限，BM、BN、BK、WM和WN不能无限大。对于

配置，搜索空间中候选实例的设计过程总结如下：

基于专家经验，确定搜索空间中想要进行测速的每个配置的warps的数量和布局。例如，X_WARPS_NUM × W_WARPS_NUM = 1 × 4。

基于计算任务的M维度和激活的量化位宽p，确定在具有最小冗余填充的情况下，想要进行测速的每个配置的Thread Block Tile大小（)。

根据warps的布局和Thread Block Tile大小，计算warp块的大小（）。最终，我们在不同平台执行速度测试，并采用速度最优的的并行配置。

Experiments

效果

实验结果表明：ABQ-LLM 在 LLaMA、LLaMA-2 系列模型上取得了优异的效果，在 weight-only 和 weight-activation 量化方面始终优于先前的 LLM 量化工作。这些结果彰显了 ABQ-LLM 的健壮性，能够适应多种量化配置和量化场景。

此外，ABQ-LLM 天然与 per-group 量化正交，而 per-group 量化也是当前主流量化推理引擎所采用的主流方式。下表的结果显示，当对 ABQ-LLM 施加分组大小为 128 的 per-group 量化时，W4A4 取得了与全精度模型相当的结果，同时超越了其他 per-group 量化的工作，例如 Atom 等。再次凸显出ABQ-LLM的健壮性。

性能

Kernel Benchmark：在LLM模型典型GEMM计算shape下对比ABQKernel和cutlass/cublas库内核的性能。值得注意的是，CUTLASS 仅支持 W4A4 和 W8A8，而 cuBLAS 仅支持 W8A8 进行量化操作。图 11 显示了结果，表明ABQ-LLM的 ABQKernel 在所有矩阵大小上都实现了卓越的加速。具体来说，对于 W2A8 和 W2A4 等特殊位组合，ABQKernel 明显优于基线方法，因为 cuBLAS 和 CUTLASS 需要转换为 W8A8 和 W4A4 才能进行计算。在 W2A8 配置中，与使用 CUTLASS 和 cuBLAS 的 W8A8 相比，ABQ-LLM的吞吐量在尺寸 (1, 4096) × (4096, 4096) 上提高了 7.47 倍。

E2E Benchmark：将 ABQKernel 集成到 fastertransformer 中，并将其与 fastertransformer 的 FP16 实现和 SmoothQuant 的 INT8 实现进行比较。与 FP16 相比，字节智能创作团队的方案实现了 2.95x 的推理加速和 4.8x的内存压缩，同时在 LLaMA-30B 模型上仅需要 10GB 的内存进行推理，这低于使用 FP16 的 LLaMA-7B 所需的内存。此外，与 SmoothQuant 相比，他们的方案实现了1.6x 的加速和2.7x的内存压缩增益，这是一个了不起的成就，因为 SmoothQuant 已经实现了很多细致的优化。