MARLIN: Mixed-Precision Auto-Regressive Parallel Inference on Large Language Models

发表时间: 2024-08 · arXiv:2408.11743

原文: https://arxiv.org/abs/2408.11743

文章标题：MARLIN：大型语言模型上的混合精度自回归并行推理
作者/机构：Elias Frantar¹, Roberto L. Castro², Jiale Chen¹, Torsten Hoefler³, Dan Alistarh¹,⁴

速读

一句话结论 本文提出了一种名为 MARLIN 的混合精度自回归并行推理核心，通过针对批处理场景深度优化权重定点化与内存加载机制，在多客户端并发时实现了接近理论极限的推理加速，并成功扩展至支持结构化稀疏硬件。

要解决什么问题 现有的混合精度推理技术（如激活值保持16位浮点精度、权重压缩为4位定点化）在单用户推理时，能够通过大幅减少内存读取量来获得显著加速。这是因为现代图形处理器（GPU）执行浮点运算的速度远快于从内存读取数据的速度，单用户场景下的算术强度极低，整个过程被内存带宽所限制。然而，在多用户并发的批处理场景下，随着输入批量大小的增加，矩阵乘法的算术强度会大幅上升。此时计算开销急剧增加，导致高昂的计算成本难以被减少的内存移动开销所掩盖，原有的混合精度实现在这种情况下往往无法提供显著加速，性能会迅速衰减。因此，核心卡点在于如何设计一种新型的底层计算核心，使其在承担批处理（如批量大小为16到32）带来的大量计算任务的同时，依然能将性能瓶颈维持在内存读取上，从而在多客户端场景下充分兑现权重定点化带来的理论加速红利。

怎么做的 核心思路是通过极致的异步内存访问、多级流水线和任务调度，确保在中等批量规模下，全局内存加载依然是整个计算的绝对瓶颈。MARLIN 的关键设计由四个核心部件构成。首先是内存加载流水线，它利用 GPU 的 L2 缓存，将全局内存到 L2 缓存、L2 缓存到共享内存的加载过程完全流水线化。通过设定深度为4的流水线并配合双缓冲技术，MARLIN 能够完美隐藏数据加载延迟。只要满足以下条件，全局加载权重矩阵 $B$ 就会成为瓶颈：$$ \frac{M \cdot K_{sm}}{B_{l2}} + \frac{K_{sm} \cdot N_{sm}}{B_{gl}} \le \frac{M \cdot N_{sm}}{B_{l2}} $$ 其中 $M$ 是批量大小，$B_{l2}$ 和 $B_{gl}$ 分别是 L2 缓存和全局内存带宽。其次是高效的去定点化与张量核心（Tensor Core）调度，MARLIN 摒弃了缓慢的常规类型转换，采用修改版的二进制位运算，将封装在32位寄存器中的4位定点化权重直接转换为16位浮点数。同时，它将每个扭曲（warp）的子瓦片宽度固定为64，精心设计线程块内存布局以避免共享内存 bank 冲突，使去定点化操作能与张量核心的矩阵乘法完美重叠。第三是条带化分区方案，当矩阵维度无法被流式多处理器（SM）完美整除时，该方案允许一个 SM 处理的计算任务跨越多个输出瓦片，从而在所有 SM 之间均匀分配负载，将同步等待和全局归约的代价降到最低。最后，针对支持 2:4 结构化稀疏的硬件，作者扩展出了 Sparse-MARLIN。它将矩阵乘法重构为 $(B^T A^T)^T$，并定制了非零值与元数据的内存索引结构，在不破坏原有密集计算流水线的前提下，进一步吃透了稀疏张量核心的算力。

效果如何 实验在 NVIDIA A10、A100、RTX 3090 等多款 GPU 上进行，测试了 Llama-2-7B 和 Falcon 等主流开源大模型。对比基线包括 PyTorch 核心、AWQ 核心、ExLlamaV2 核心以及 bits-and-bytes 核心，它们代表了当前社区中最主流的4位定点化推理路线。在核心层面的基准测试中，当批量大小在16到32之间时，MARLIN 在 A10 GPU 上相比16位浮点基线实现了约3.9倍的加速，几乎达到了4位定点化的理论极限；而其他基线方法在批量大小增加时性能会迅速衰减。在端到端测试中，作者将 MARLIN 集成到 vLLM 推理引擎中，在批量大小为16的设置下，多用户令牌生成速度相比 vLLM 的标准精度核心提升了高达2.8倍。Sparse-MARLIN 在此基础上还能带来最高65%的额外加速，端到端加速比最高可达3.2倍。该方法的局限在于，当批量大小进一步增加（如达到128或在处理长上下文的预填充阶段）时，计算任务会彻底转变为计算密集型，此时 MARLIN 的加速比会逐渐降至1.5倍左右；此外，定点化和稀疏化本身仍会带来微小的模型精度损失，尽管作者通过引入改进版的 GPTQ 和 SparseGPT 算法在很大程度上弥补了这一代价。

A1 主要贡献

本文研究了在多客户端并行（即批处理）场景下，如何通过权重定点化技术加速大型语言模型（LLM）的推理。现有技术在单用户推理时能通过减少内存移动获得显著加速，但在批处理场景下，由于算术强度大幅增加，加速效果不明显。本文旨在解决这一问题，即设计一种GPU核心，使其在支持批处理工作负载所需的大量计算的同时，仍能保持为内存密集型（memory-bound）。

核心问题与研究目标:

创新点与主要贡献:

  1. MARLIN核心的设计与实现: 提出了一种名为MARLIN（Mixed-precision Auto-Regressive LINear kernels）的混合精度推理核心。MARLIN结合了多种先进技术，包括异步内存访问、复杂的任务调度与流水线技术，以及定制化的定点化支持。
  2. 实现接近最优的批处理加速: 实验表明，在单个LLM层上，MARLIN能够在NVIDIA Ampere GPU上实现接近最优的性能。具体而言，对于4位权重推理，在NVIDIA A10 GPU上，当批量大小最高达到16-32时，MARLIN相对于FP16可获得约3.9倍的加速；当批量大小增加到128时，加速比逐渐降至1.5倍，此时问题转为计算密集型（compute-bound）。
  3. 端到端性能提升: 将MARLIN集成到流行的LLM服务引擎vLLM中，在Llama和Falcon等主流开源LLM上进行了端到端测试。结果显示，在批量大小为16时，MARLIN相比vLLM的标准精度核心，可将多用户令牌生成速度提升高达2.8倍。
  4. 对稀疏性的扩展: 提出了MARLIN的扩展版本Sparse-MARLIN，它支持NVIDIA的2:4稀疏Tensor Core格式。相对于原始的密集（dense）MARLIN版本，Sparse-MARLIN可带来高达65%的额外加速，端到端加速比最高可达3.2倍。
  5. 开源贡献: MARLIN、Sparse-MARLIN及其vLLM集成的代码均已开源，为社区提供了在批处理场景下高效利用权重定点化LLM推理的实用解决方案。

图1：MARLIN在增加批量大小时的峰值性能图示，针对单个大型LLM线性层，与其他流行的开源核心进行比较，显示了我们在该场景下能够实现接近最优的性能。
图1：MARLIN在增加批量大小时的峰值性能图示，针对单个大型LLM线性层，与其他流行的开源核心进行比较，显示了我们在该场景下能够实现接近最优的性能。

A3 背景知识

2.1 图形处理器（GPU）

GPU架构概述。NVIDIA GPU由一组流式多处理器（Streaming Multiprocessor, SM）阵列构成，这些SM共享一个称为全局内存（Global Memory, GMEM）的DRAM内存和一个L2缓存。每个SM被划分为多个分区，每个分区包含各种处理块，如扭曲调度器（warp scheduler）、寄存器文件（Register File, RF）和L0指令缓存。SM内的处理块共享一个L1缓存，该缓存的一部分可配置为称为共享内存（Shared Memory, SMEM）的高速暂存存储器。每个处理块内有四种类型的计算单元：整数单元、特殊函数单元、浮点单元（FPU）/CUDA核心和张量核心单元（Tensor-Core Units, TCU）。

张量核心单元（TCU）的角色。TCU最早在Volta架构中引入，主要通过每个周期执行一次矩阵乘法累加（MMA）操作来针对机器学习工作负载，从而降低了获取和解码多条指令的成本。在Ampere架构中，TCU在FP16上的性能比在FPU上运行的融合乘加（FMA）操作高出多达16倍。

CUDA编程与执行模型。CUDA模型定义了三个粒度级别：线程块（thread blocks）、扭曲（warps）和线程（threads）。扭曲是CUDA中的基本调度单位，由32个并发执行的线程组成。线程块是一组扭曲的集合，被调度到同一个SM上执行。每个SM上同时运行的扭曲数量和线程块数量取决于硬件限制，例如扭曲调度器的数量、每个线程的寄存器数量或可用的SMEM大小。

2.1.1 现代张量核心单元（TCU）

图2：异步复制操作（有/无L1旁路，右侧）与标准操作（左侧）的图示。
图2：异步复制操作（有/无L1旁路，右侧）与标准操作（左侧）的图示。

Ampere架构TCU的扩展。Ampere GPU对TCU进行了扩展，以处理两个方面：1) 细粒度的结构化稀疏性，形成了稀疏张量核心单元（Sparse Tensor-Core Units, SPTC）；2) 异步复制操作。首先，通过新的2:4格式支持结构化稀疏性，其性能承诺比原始TCU快2倍，比FPU快32倍。

2:4稀疏格式。2:4格式将左侧（LHS）矩阵划分为长度为4的向量，并对每个向量中的两个元素进行清零，从而形成一个50%稀疏但结构化的矩阵。图2展示了一个SPTC的简化表示。稀疏矩阵由两个数据结构表示：（1）一个包含非零值的“值”结构（蓝色表示）；（2）一个包含每个非零值在每4个元素组中位置的“元数据”结构（紫色表示）。元数据结构被SPTC上的新硬件组件用来仅选择右侧（RHS）矩阵中计算所需的元素，从而跳过被清零的值。

异步复制指令。NVIDIA的Ampere微架构引入了数据获取的改进以增强Tensor Core性能，具体是通过一个新的异步复制指令，允许数据直接从GMEM加载到SMEM。如图2所示，在之前的架构中（图2左），数据需要先通过L1缓存加载到RF（全局加载指令），然后传输到SMEM（共享存储指令），最后再移回RF（共享加载指令）。Ampere的新异步复制指令通过避免中间的RF访问节省了SM内部带宽。该指令有两个变体：“access”变体将数据存入L1以供后续访问和重用（图2中），而“bypass”变体则跳过L1缓存（图2右）。

2.2 LLM上的混合精度推理

混合精度推理的动机。混合精度LLM推理通过静态压缩预训练模型权重，并在推理过程中根据需要动态解压，从而有潜力减少模型巨大的内存占用，并相应地加速内存密集型工作负载。

权重定点化。一种标准的LLM压缩方法是仅权重定点化（weight-only quantization），它降低了权重W的存储精度，而保持层输入X不变。这种方法非常流行，因为它即使在相对较高的压缩率下也表现出卓越的准确性鲁棒性。

均匀定点化。本文关注均匀定点化，即给定一个向量 $v \in R^n$，我们定义：

公式1
公式1

其中，$\lfloor·\rceil$表示四舍五入到最近的整数，$z = z(v) = \min(v)$ 映射到零，而 $s = s(v) = (\max(v) - \min(v)) / (2^b - 1)$ 是缩放因子（scale）。重构误差可以计算为 $\epsilon_r = \|v - Q(v, b)\|_2$。通过将v划分为多个组并对每个组独立进行定点化（例如，每组128个连续值），可以存储每个组的s和z值，从而改善误差。

定点化算法。本文使用GPTQ【【10】，Frantar, E., Ashkboos, S., Hoefler, T., and Alistarh, D. Gptq: Accurate post-training quantization for generative pretrained transformers. arXiv preprint arXiv:2210.17323, 2022】的一个变体来执行实际的权重定点化。该方法利用二阶信息来补偿定点化误差，从而仅导致微小的准确性下降。但需要强调的是，本文提出的核心（kernel）技术独立于任何特定的定点化算法。

A2 方法细节

3. The MARLIN Kernel

3.1 动机

理论基础。LLM权重定点化的动机在于现代GPU具有很高的浮点运算/字节比率（FLOPs/Bytes ratio），这意味着它们执行浮点运算的速度远快于从内存读取数据的速度。例如，A10 GPU的FLOPs/Bytes比率约为200【【17】，NVIDIA. Nvidia a10 datasheet. https: //www.nvidia.com/content/dam/en-zz/ Solutions/Data-Center/a10/pdf/ datasheet-new/nvidia-a10-datasheet. pdf, 2022a】。在单层矩阵乘法中，处理一个输入token每个权重需要2个FLOPs，而GPU在加载一个4位权重的时间内可以执行100个FLOPs。因此，只要输入批量大小小于理论最优值 $b_{opt} \approx 50$，内存加载将主导运行时。$b_{opt}$是延迟既不受内存限制也不受计算限制的批量大小，此时我们能以最大吞吐量实现最低延迟。在实践中，我们希望恰好在这个批量大小下运行：更小不会带来进一步的加速，更大则不会提高吞吐量。

实现挑战。然而，要实现这样一个能同时最大限度利用GPU所有资源（即计算和内存）的混合精度（FP16-INT4）矩阵乘法（matmul）核心是一个巨大的挑战。本文的目标是通过设计MARLIN（一个高度优化的混合精度自回归线性核心）来尽可能接近这个目标。

3.2 Ampere矩阵乘法

通用高性能矩阵乘法实现。本节描述在GPU上实现峰值性能（均匀精度）矩阵乘法核心的通用概念，特别是在Ampere级设备上。我们紧密遵循CUTLASS的分层并行化模型【【19】，NVIDIA. Efficient GEMM in CUDA. https: //github.com/NVIDIA/cutlass/blob/ main/media/docs/efficient_gemm.md, 2024a】。具体来说，我们考虑将一个 $M \times K$ 的矩阵A与一个 $K \times N$ 的矩阵B相乘，得到一个 $M \times N$ 的输出矩阵C。

3.3 混合精度挑战

混合精度实现面临的挑战。将上述均匀精度矩阵乘法模型应用于混合精度场景，同时保持峰值性能（特别是在M中等大小，操作接近内存密集型时）是具有挑战性的，原因如下：

  1. 并行化配置：必须非常仔细地配置各个并行化级别，以确保加载量化操作数B确实是核心的主要运行时瓶颈，而不是例如重复重新加载全精度 $A_{sm}$ 块。
  2. 内存加载效率：由于运行时主要由内存加载决定，尽管B的表示被显著压缩，这方面必须达到峰值效率。
  3. 计算与内存的重叠：对于中等大小的M，矩阵乘法计算的成本可能接近于总内存加载成本，因此需要极其小心的重叠才能接近理论性能。此外，还需要管理量化元数据，这使得这部分更加棘手。
  4. 并行化限制：由挑战1带来的划分约束，加上M不是很大这一事实，显著限制了并行化选项。这使得在SM和warp级别上都难以实现峰值的内存加载和计算。现有模型的矩阵形状可能对特定GPU不利，进一步放大了这种效应。

MARLIN核心专门解决了以上所有挑战，最终使其在许多实际环境中能够实现接近峰值的性能。

3.4 核心设计

核心设计的前提。我们假设矩阵A为全FP16精度，而 $K \times N$ 的矩阵B已被（对称地）量化为INT4，可以为N列中的每一列使用一个FP16缩放因子，或者为每列中每G个连续权重使用一个缩放因子，总共有 $\lceil K/G \rceil N$ 个缩放因子。

受限于权重加载。理论上，执行我们的目标矩阵乘法需要接触 $16MK + 4KN + 16MN$ 位内存（读取两个操作数并写入结果），同时执行 $MKN$ 次乘法累加操作（每次计为2个FLOPs）。如果M相对较小，我们的问题算术强度较低。因此，它应该受限于从全局GPU内存中读取量化权重B的成本。为了在实践中保持这一理论，与之前研究的M=1的情况（其中A和C都非常小）不同，现在输入和输出具有不可忽略的大小，因为这些操作数的位宽比我们的权重要高4倍。因此，我们需要选择一个足够大的 $N_{SM}$ 来最小化昂贵的 $A_{sm}$ 块的重载。同时，这减少了 $C_{sm}[i_{sm}, j_{sm}]$ 子问题的数量，使得难以充分利用所有SM。解决这些问题的关键是利用GPU的L2缓存，它通常比全局内存快得多。此外，GPU可以同时从L2加载到L1和从全局加载到L2。因此，我们可以将这些加载流水线化，只要总内存流量不超过L2带宽，就可以完全隐藏 $A_{sm}$ 块加载的带宽成本。因此，我们通过将C划分为大小为 $M \times N_{sm}$ （其中 $N_{sm} \in \{64, 128, 256\}$）的块，即中等宽度的完整输入批次大小的瓦片，然后将每个相应的独立矩阵乘法子问题分配给一个SM。在 $N_{sm}=256$ 时，即使批次大小 $M=64$ 仍然受限于全局权重加载。更精确地说，只要从L2读取 $A_{sm}$ 和 $B_{sm}$ 块的速度更快，全局加载 $A_{sm}$ 块仍然是瓶颈，即：

公式2
公式2

其中 $B_{l2}$ 和 $B_{gl}$ 分别表示L2和全局带宽。

最大化加载带宽。为了最大化实际加载带宽，我们旨在利用尽可能宽的加载；在当前GPU上，每个线程16字节（128位）。这意味着一个warp可以通过一条指令加载 $32 \times 32 = 1024$ 个INT4权重。为达到峰值效率，我们需要一个warp中的8个线程各自从GMEM读取128字节的连续块（假设是128字节对齐的地址），即一个完整的缓存行。对于形状为 $M \times K_{sm}$ 的A块，这要求 $K_{sm}$ 至少为64。由于权重在推理期间是静态的，因此可以离线预处理，我们通过重新排列 $16 \times 64$ 的瓦片，使其在内存中连续布局，从而实现最优加载，并简化了相应的索引。尽管我们不断从L2缓存中重新加载 $A_{sm}$ 块，但B的每个元素只被访问一次。然而，每次读取都会被放入L2缓存，可能会驱逐一些SM仍然需要的A的部分。为了避免这种缓存污染，我们使用带有“evict first”缓存提示的 cp.async 指令，确保不必要存储的B数据在任何其他缓存行之前被丢弃。

共享内存布局。我们总是通过Ampere的 cp.async 指令从全局（或L2）异步加载到共享内存；这不需要临时寄存器，也使得将这些加载与计算重叠变得更加容易。由于我们对B进行了离线预处理，我们可以简单地以连续的方式复制到共享内存，避免了bank冲突。相比之下，处理A的片段需要更加小心：具体来说，我们需要确保每个 $16 \times 16$ FP16 A块的索引 $ij, (i+8)j, i(j+1)$ 和 $(i+8)(j+1)$ 对应的16字节向量存储在不同的内存bank中。只有这样 ldmatrix.sync 指令才能无冲突地执行。（这些指令加载A操作数数据并将其分布到warp线程中，为Tensor Core的使用做准备。）这可以通过将16字节元素ij存储在激活瓦片中对应共享内存瓦片的 $i(i \oplus j)$ 位置来实现，其中 $\oplus$ 表示XOR操作【【20】，NVIDIA. CUTLASS convolution. https://github. com/NVIDIA/cutlass/blob/main/media/ docs/implicit_gemm_convolution.md, 2024b】。此索引转换的另一个关键方面是，如果一个warp读取全局A瓦片的一个连续子瓦片（例如，前4行），那么它将被置换后仍然整体上连续地写入共享内存。虽然没有文档说明，但这似乎是避免写入时bank冲突所必需的，正如我们在分析NVIDIA性能分析器输出时观察到的。这些索引计算有些复杂，动态处理可能很慢；然而，由于它们只影响相对少量的共享内存位置，并且在主循环中保持静态，我们可以在寄存器中预计算它们，并配合适当的展开。

内存加载流水线。同时达到接近最大带宽和最大计算能力的关键是完全重叠内存加载和Tensor Core数学运算。对于全局到共享内存的加载，这可以通过 cp.async 操作实现，在每次迭代中预取将在 P-1 步后使用的 $A_{sm}$ 和 $B_{sm}$ 块，其中 P 是流水线深度（我们还需要一个缓冲区用于当前瓦片）。此外，我们可以在累加当前部分矩阵乘法之前，从共享内存中预取下一个子瓦片（大多数GPU操作在遇到依赖之前不会阻塞），该部分矩阵乘法的操作数已在上次迭代中被取到寄存器中——这种技术也称为双缓冲（double buffering）【【19】，NVIDIA. Efficient GEMM in CUDA. https: //github.com/NVIDIA/cutlass/blob/ main/media/docs/efficient_gemm.md, 2024a】。我们选择流水线深度 P=4 有两个原因：(a) 在我们所有的测试中，这足以完全隐藏延迟，并且即使在 M=64 时也能容纳在共享内存中；(b) 因为它可以被2整除。后者至关重要，因为它允许我们在整个流水线中平滑地展开，因为在 P 次迭代后，流水线和寄存器缓冲区的索引都将始终为0。这种展开使得所有共享内存寻址完全静态，通过使用我们可用的额外寄存器避免了缓慢的变换索引计算。最后，这也似乎是让CUDA编译器正确排序指令以实现真正双缓冲的最可靠方法。图3展示了MARLIN核心中使用的多层流水线。

图3：MARLIN核心中的流水线层次。
图3：MARLIN核心中的流水线层次。

图4：MARLIN的warp布局图示。多个warp累加同一输出瓦片的部分结果；相应的伪代码见算法1。
图4：MARLIN的warp布局图示。多个warp累加同一输出瓦片的部分结果；相应的伪代码见算法1。

Warp布局。在单个SM上计算 $C_{sm}$ 必须进一步在warp之间进行细分：如果直接进行，每个warp将计算输出的一个 $M \times (N_{sm}/\#warps)$ 的瓦片。为了达到峰值计算吞吐量，我们希望至少使用四个（因为Ampere GPU有四个warp调度器），理想情况下是八个warp【【26】，Sun, W., Li, A., Geng, T., Stuijk, S., and Corporaal, H. Dissecting tensor cores via microbenchmarks: Latency, throughput and numeric behaviors. IEEE Transactions on Parallel and Distributed Systems, 34(1):246–261, 2022】，以获得额外的延迟隐藏。然而，这会导致瓦片尺寸变小，尤其是在较小的 $N_{sm}$ 下。这不仅对我们之前讨论的内存重排有问题，也阻碍了Tensor Core的吞吐量，因为小瓦片宽度会给tensorcore操作带来更多的顺序依赖（因为这些连续操作必须使用相同的累加器），可能导致停顿。因此，我们将每个warp的子瓦片宽度固定为64，并进一步在 $K_{sm}$ 上分割计算；图4说明了这样的warp布局，算法1提供了相应的伪代码。因此，多个warp将在寄存器中累加相同 $C_{wa}[i_{sm}, j_{sm}][i_{wa}, j_{wa}]$ 的部分结果。这些结果最终必须在共享内存中进行归约，然后才能最终写出。但这可以通过对数并行归约（logarithmic parallel reduction）【【12】，Harris, M. et al. Optimizing parallel reduction in cuda. Nvidia developer technology, 2007】来完成，这通常只会带来最小的开销。

算法1
算法1

去量化与Tensor Cores。从INT4到FP16的简单类型转换速度很慢；因此，我们采用了Kim等人【【13】，Kim, Y. J., Henry, R., Fahim, R., and Awadalla, H. H. Who says elephants can’t run: Bringing large scale moe models into cloud scale production. arXiv preprint arXiv:2211.10017, 2022】二进制操作的修改版本。以最简单的情况为例：将一个INT16中位于12-15位的INT4转换为一个有符号的FP16值。首先，我们通过AND掩码仅提取对应于我们INT4的位，并将结果的1-7位变为0110010（通过OR操作）；这可以在单个lop3指令中完成。现在，我们得到了一个指数为50的FP16数，其尾数的最后4位对应我们的转换目标。因此，减去一个指数为50、尾数为0的FP16值，将得到我们目标4位的无符号浮点表示。为了使该值有符号，我们还需减去8，但这可以直接融合到我们减去的总值的最后3位中。现代GPU可以同时用一个32位寄存器中封装的两个独立的16位操作数进行计算。因此，我们可以使用上述过程同时高效地去量化一个INT32中的两个INT4。最后，我们希望直接去量化到后续Tensor Core调用所需的正确寄存器布局。为此，我们再次利用B可以离线预处理的优势，重新组织权重，使得每个线程读取的16字节向量精确包含其所需的4个独立的 $16 \times 16$ Tensor Core块的8个量化权重。此外，在一个INT32内，权重按模式64207531交错存储，以支持前面提到的并行解码。在最内层，我们累加一个 $M \times 16$ 乘以 $16 \times 64$ 矩阵乘法的结果。我们按列执行这个累加，发出 $16 \times 16$ 乘以 $16 \times 8$ 的Tensor Core mma.sync 指令。这比按行执行的优势在于，我们可以将下一个B操作数的去量化与当前列的Tensor Core数学运算流水线化。

分组与指令排序。对于逐输出量化，我们可以在全局写出前简单地对最终输出进行一次缩放。一个有趣的观察是，尽管这些加载与任何计算都不是异步的，但通过 cp.async 后紧跟一个即时的wait group指令来执行它们仍然至关重要，以避免编译器对主循环进行不利的指令重排。对于分组量化（这对于保持良好准确性至关重要），我们必须在主循环期间加载并应用缩放。首先，我们以与量化权重类似的方式重组缩放因子的存储，使得同一类型线程所需的不同 $16 \times 16$ 块的缩放因子被打包在一起，并可以作为一个16字节向量从共享内存加载。原则上，对于组大小为128和 $B_{sm}$ 瓦片形状为 $64 \times 256$ 的情况，我们只需要每隔一个瓦片从全局和共享内存加载一次新的缩放因子（并且在这里只在第一个子瓦片期间加载一次）。然而，编译器似乎对代码最关键部分中的这种不规则性相当脆弱，导致在某些形状设置下指令排序不佳，整体 slowdown 10-20%。相反，我们发现为每个子瓦片从共享内存重新加载缩放因子可以保持峰值性能。这样做增加了一些技术上不必要的共享内存加载，但有足够的额外带宽来支持这一点而没有开销，同时它也保留了编译器为非分组量化精心设计的流水线指令排序。

条带化分区。通过上述所有技术，我们可以在矩阵较大且可以在N轴上完美地在所有SM之间划分时，达到接近最优的计算和带宽性能。在实践中，这种情况很少见。在这种情况下，标准的补救措施是在K维度上也进行分区，但对于许多流行的层形状和GPU组合，我们需要大量额外的分割才能达到均匀分布而没有显著的波浪量化。这反过来又增加了许多全局归约步骤，带来了额外的开销。相反，我们选择了一种条带化分区方案（striped partitioning scheme），其中一个SM处理的B的“条带”可以跨越多个 $C_{sm}$ 瓦片（参见图5）。具体来说，我们首先确定每个SM要处理的 $B_{sm}$ 瓦片数量 $T = \lceil \#tiles / \#SMs \rceil$，然后从左上角开始按列分配（最多）T个瓦片。关键是，如果我们到达一个瓦片列的底部，但当前SM尚未拥有T个瓦片，我们将继续从下一个瓦片列的顶部开始分配瓦片；换句话说，条带可以跨越多列。这确保了瓦片在所有SM之间的分布大致均匀，同时最小化了所需全局归约步骤的数量。该策略类似于stream-k分区【【21】，Osama, M., Merrill, D., Cecka, C., Garland, M., and Owens, J. D. Stream-k: Work-centric parallel decomposition for dense matrix-matrix multiplication on the GPU. In ACM SIGPLAN Annual Symposium on Principles and Practice of Parallel Programming, 2023】。

图5：MARLIN的条带化分区方案。
图5：MARLIN的条带化分区方案。

全局归约实现。我们以串行方式实现同一瓦片列中条带之间的全局归约，从下到上。后一种方法最有效，因为在存在任何列溢出的情况下，最底部的SM将最快得到其结果，而最顶部的SM最慢。我们直接在输出缓冲区中以FP16执行归约，以最大化L2缓存命中率，从而最小化任何全局读取开销。这也使操作基本上是就地的，只需要一个小的额外锁缓冲区用于同步。最后，我们注意到，对于批次大小远大于64的情况，我们可以为条带索引计算虚拟地复制B，然后通过模运算移回原始矩阵，并将A指针前进到相应的大小为64的输入批次段。这为大型输入批次大小（如LLM预填充期间发生的情况）显著减少了全局归约，并在此设置下提高了计算吞吐量。

3.5 GPTQ 修改

对GPTQ的改进。MARLIN使用的量化格式，专为峰值推理效率而设计，与原始GPTQ实现【【10】，Frantar, E., Ashkboos, S., Hoefler, T., and Alistarh, D. Gptq: Accurate post-training quantization for generative pretrained transformers. arXiv preprint arXiv:2210.17323, 2022】略有不同，但仍能产生高精度模型。我们还向GPTQ集成了两个小改进：(a) 通过搜索最优的组级裁剪阈值来选择组缩放因子，类似于【【15】，Lin, J., Tang, J., Tang, H., Yang, S., Dang, X., and Han, S. Awq: Activation-aware weight quantization for llm compression and acceleration. arXiv preprint arXiv:2306.00978, 2023】；(b) 支持可变长度的校准序列。我们发现这些修改相比标准GPTQ带来了微小但一致的准确性提升，同时具有更高性能的优势。（我们也提供了一个在模型格式之间进行简单转换的脚本。）图6展示了我们的GPTQ变体与原始未压缩模型在困惑度（越低越好）与模型大小（以比特为单位）方面的对比。这表明，MARLIN量化的模型在与未压缩基线相同的困惑度下，大小减少了约3.33倍。虽然这不是无损的（在此位宽和组大小下，理想增益为3.87倍），但这是一个显著的改进，特别是考虑到MARLIN的高推理效率。

图6：通过GPTQ量化到MARLIN格式的Llama2模型的帕累托曲线。
图6：通过GPTQ量化到MARLIN格式的Llama2模型的帕累托曲线。

4. The Sparse-MARLIN Kernel

设计动机。为了进一步提高FLOPS/Byte比率，我们可以在4位量化权重表示的基础上集成2:4稀疏方案。NVIDIA Ampere架构中的稀疏Tensor Cores (SPTCs)为在专门硬件单元上执行50%稀疏矩阵提供了有效手段。然而，要利用SPTCs，需要对先前描述的MARLIN核心进行一些修改和扩展。

矩阵乘法重构。首先，为了适应mma.sp指令的约束（该指令启用SPTCs，并要求稀疏矩阵作为张量操作中的左侧(LHS)操作数【【18】，NVIDIA. Nvidia instruction set. https://docs.nvidia.com/cuda/ parallel-thread-execution/index.html# warp-level-matrix-instructions-for-sparse-mma, 2022b】），我们设计了新的特定数据布局。具体来说，将A与B相乘的问题在底层被重构为求解 $(B^T A^T)^T$ 来产生C。然而，这种重构保留了先前描述的密集MARLIN核心设计中的所有技术和优化。请注意，B可以根据需要离线预处理，而A可以通过ldmatrix指令及其.trans可选限定符在SMEM中即时转置，而不会导致性能下降。

新的数据结构。接下来，我们描述在Sparse-MARLIN中编码2:4稀疏矩阵所需的两个新数据结构，以及为解决此特定问题而进行的调整：（1）非零值结构，和（2）元数据索引结构。

图7：Sparse-MARLIN非零值结构布局
图7：Sparse-MARLIN非零值结构布局

量化的非零值。图7左侧展示了一个大小为 $N \times K$ 的4位量化矩阵B，该矩阵已被修剪为2:4稀疏。该矩阵的压缩表示（图7中1）在内部维度上的大小将是原始矩阵的一半，即 $N \times K/2$。然而，由于每个值是4位元素，我们可以在此基础上应用密集的MARLIN压缩方法，将8个元素进一步压缩到一个32位值中（图7中1b），最终大小为 $N/8 \times K/2$。为最大化内存效率，由于权重在推理期间保持静态，每个 $64 \times 16$ 的瓦片都被重新排列，以便每个线程在连续的内存位置加载和存储元素，类似于密集的MARLIN。继续看图7，非零值结构中的彩色元素表示线程T0为一个 $64 \times 16$ 的B块提供的元素示例。成对的颜色表示在同一个mma.sp指令内处理的元素，需要4次迭代才能计算所有元素。这种布局在从GMEM加载非零值结构时，确保了最宽的128位指令（每个线程4个连续的32位元素，如1b所示）。此外，由于乘法的重新定义，并且输出C将是一个形状为 $M \times N$ 的FP16矩阵，这种布局在将结果从RF转置存储到GMEM时，也确保了128位指令（例如，T0寄存器中存储的第0列中前八个连续输出元素）。因此，这种乘法的重构不仅在不降低性能的情况下转置存储结果，而且与基线密集设计相比，进一步提高了输出写入的效率。

图8：Sparse-MARLIN元数据索引布局。
图8：Sparse-MARLIN元数据索引布局。

元数据索引。为了选择稀疏计算中将需要的右侧（RHS）操作数A中的元素，需要一个包含原始矩阵中非零元素索引的元数据结构。图8左侧展示了B的元数据索引结构。由于这是2:4稀疏性，索引将在0~3范围内，用2位编码。基于图7中描述的数据布局，并考虑到mma.sp指令的稀疏性选择器约束，我们为元数据结构提出了一种新的数据布局。稀疏性选择器指示四个连续线程组中的哪个线程（们）将为整个组提供元数据。在图8所示的示例中，稀疏性选择器可以是0（线程T0, T1）或1（线程T2, T3）。首先，我们必须根据先前描述的非零值结构重新排序行，如2a所示。这使我们能够使用128位加载指令从GMEM到SMEM。然后，为了无bank冲突地将数据从SMEM加载到RF，我们执行单个ldmatrix指令，该指令将包含要执行的下四个mma.sp操作的所有信息，并根据需要将信息分布到线程T0~T3。然而，需要进行先前的数据重排，如2b所示。这样，线程T0, T1将包含前两次迭代的信息，而线程T2, T3将包含剩余两次迭代的信息。需要注意的是，所有这些预处理都是一次性离线完成的，没有运行时开销。

A4 实验环境

A4 实验结果

5.1 核心基准测试

5.2 端到端实验

图9：MARLIN在流行模型的真实层形状上的性能。
图9：MARLIN在流行模型的真实层形状上的性能。

图10：MARLIN与其他流行开源核心的持续性能对比。
图10：MARLIN与其他流行开源核心的持续性能对比。

图11：MARLIN核心在四种不同矩阵形状下的Roofline分析。
图11：MARLIN核心在四种不同矩阵形状下的Roofline分析。

图12：MARLIN和Sparse-MARLIN与其他流行开源核心的峰值性能对比。
图12：MARLIN和Sparse-MARLIN与其他流行开源核心的峰值性能对比。

图13：MARLIN和Sparse-MARLIN与其他流行开源核心的持续性能对比。
图13：MARLIN和Sparse-MARLIN与其他流行开源核心的持续性能对比。

表1：Llama-2-7B原始模型及量化/稀疏模型的准确性。INT4模型用于MARLIN，由我们的GPTQ版本生成。INT4+2:4模型用于Sparse-MARLIN，由SparseGPT生成；其指标高于原始模型是因为通过知识蒸馏在合成数据上进行了进一步微调。
表1：Llama-2-7B原始模型及量化/稀疏模型的准确性。INT4模型用于MARLIN，由我们的GPTQ版本生成。INT4+2:4模型用于Sparse-MARLIN，由SparseGPT生成；其指标高于原始模型是因为通过知识蒸馏在合成数据上进行了进一步微调。

图14：MARLIN和Sparse-MARLIN与vLLM FP16基线的端到端生成时间对比。
图14：MARLIN和Sparse-MARLIN与vLLM FP16基线的端到端生成时间对比。

表2：MARLIN与vLLM的FP16基线相比的端到端生成加速比。
表2：MARLIN与vLLM的FP16基线相比的端到端生成加速比。

图15：MARLIN和Sparse-MARLIN与vLLM FP16基线的服务基准测试（TPOT，每个输出令牌的时间）。
图15：MARLIN和Sparse-MARLIN与vLLM FP16基线的服务基准测试（TPOT，每个输出令牌的时间）。

图16：MARLIN和Sparse-MARLIN与vLLM FP16基线的服务基准测试（TTFT，首个令牌的时间）。
图16：MARLIN和Sparse-MARLIN与vLLM FP16基线的服务基准测试（TTFT，首个令牌的时间）。

A7 补充细节

6. 相关工作

LLM定点化推理支持。本文主要关注为量化LLM推理提供高效支持的相关工作。先前已有大量关于LLM权重精确量化的研究，流行的方法包括GPTQ【【10】，Frantar, E., Ashkboos, S., Hoefler, T., and Alistarh, D. Gptq: Accurate post-training quantization for generative pretrained transformers. arXiv preprint arXiv:2210.17323, 2022】和AWQ【【15】，Lin, J., Tang, J., Tang, H., Yang, S., Dang, X., and Han, S. Awq: Activation-aware weight quantization for llm compression and acceleration. arXiv preprint arXiv:2306.00978, 2023】，以及对四舍五入（RTN）量化的探索【【4】，Dettmers, T. and Zettlemoyer, L. The case for 4-bit precision: k-bit inference scaling laws. arXiv preprint arXiv:2212.09720, 2022】，后者通常精度较低。MARLIN的并行化方法可以推广到这些量化方法。事实上，自本文核心发布GPTQ格式版本以来，vLLM中已独立引入了支持AWQ的MARLIN版本【【30】，vLLM Project Contributors. vllm project pull request #6612: Add support for awq marlin. https://github.com/ vllm-project/vllm/pull/6612, 2024】。

权重与激活值同时量化。更广泛地看，LLM量化方法也可以考虑同时压缩权重和激活值【【5】，Dettmers, T., Lewis, M., Belkada, Y., and Zettlemoyer, L. LLM.int8(): 8-bit matrix multiplication for transformers at scale. Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurIPS 2022, 2022】，高级方法如SmoothQuant【【31】，Xiao, G., Lin, J., Seznec, M., Demouth, J., and Han, S. Smoothquant: Accurate and efficient post-training quantization for large language models. arXiv preprint arXiv:2211.10438, 2022】或QuaRot【【1】，Ashkboos, S., Mohtashami, A., Croci, M. L., Li, B., Jaggi, M., Alistarh, D., Hoefler, T., and Hensman, J. Quarot: Outlier-free 4-bit inference in rotated llms. arXiv preprint arXiv:2404.00456, 2024】。然而，激活值的量化往往更复杂，因为会出现大的“离群值”【【5】，Dettmers, T., Lewis, M., Belkada, Y., and Zettlemoyer, L. LLM.int8(): 8-bit matrix multiplication for transformers at scale. Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurIPS 2022, 2022】。因此，这些方法要么针对较低的8位精度，要么需要更复杂的额外处理步骤，例如通过旋转矩阵【【1】，Ashkboos, S., Mohtashami, A., Croci, M. L., Li, B., Jaggi, M., Alistarh, D., Hoefler, T., and Hensman, J. Quarot: Outlier-free 4-bit inference in rotated llms. arXiv preprint arXiv:2404.00456, 2024】。MARLIN的方法可以扩展到这种情况，例如，最近一个独立的后续工作将MARLIN的方法扩展到了激活值量化为8位，而权重量化为4位的情况【【33】，Zhang, Y., Zhang, P., Huang, M., Xiang, J., Wang, Y., Wang, C., Zhang, Y., Yu, L., Liu, C., and Lin, W. Qqq: Quality quattuor-bit quantization for large language models, 2024】。

A5 结论

本文提出了MARLIN，一种用于实现LLM生成式推理的混合精度核心的通用方法。该方法通过利用新的GPU硬件指令和并行化技术，实现了接近最优的效率。具体而言，MARLIN及其稀疏版本在单层上达到了接近最优的效率，并在实际部署场景中，以适度的精度影响为代价，带来了高达3倍的加速。

未来工作展望：

  1. 支持更极端的压缩技术: 一个自然的方向是研究MARLIN对近期提出的更复杂的“极端”压缩技术的支持，例如向量量化【【3】，Chee, J., Cai, Y., Kuleshov, V., and Sa, C. D. Quip: 2-bit quantization of large language models with guarantees, 2023】和【【7】，Egiazarian, V., Panferov, A., Kuznedelev, D., Frantar, E., Babenko, A., and Alistarh, D. Extreme compression of large language models via additive quantization. arXiv preprint arXiv:2401.06118, 2024】，这些技术需要更复杂的解压过程。
  2. 支持其他形式的混合精度: 另一个方向是研究MARLIN对其他形式混合精度的支持，例如由激活值压缩或稀疏性产生的混合精度。

方法细节中引用的参考文献汇总