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

作者/机构: Elias Frantar, Roberto L. Castro, Jiale Chen, Torsten Hoefler, Dan Alistarh

A1 主要贡献

本文旨在解决大语言模型（LLM）推理中的一个核心问题：权重（weight）量化技术虽然在单用户推理场景下能通过减少内存移动带来显著加速，但在多用户并行的批处理（batched）场景下，由于算术强度大幅增加，其加速效果会大打折扣。批处理场景对于实际的大规模LLM服务至关重要，因此，如何在保持GPU内存受限（memory-bound）特性的同时，支持批处理带来的高计算需求，是一个亟待解决的开放问题。

本文提出了名为MARLIN（Mixed-precision AutoRegressive LINear kernels）的混合精度自回归线性核，正面回答了上述问题。MARLIN的核心贡献在于，它证明了即使在批处理设置下，也能实现接近理论极限的量化加速。具体贡献如下：

1. 设计并实现了MARLIN核：这是一种新型的GPU核，专门为混合精度（例如，4-bit权重）的批处理推理设计。它结合了多种先进技术，包括异步内存访问、复杂的任务调度与流水线（pipelining）、以及为量化特别设计的内存布局和计算优化。

2. 实现接近最优的批处理加速：实验表明，对于4-bit权重量化，MARLIN在批大小（batch size）高达16-32时，仍能提供接近理论最大值（4倍）的加速。即使批大小增加到64-128，虽然加速比逐渐下降，但仍然非常显著。这解决了现有混合精度实现中批处理性能迅速下降的痛点。下图1展示了MARLIN在单个大型LLM线性层上，随着批大小增加，其性能与理想值和其他开源核的对比。

   

3. 通过与vLLM集成验证端到端性能：将MARLIN集成到流行的LLM服务引擎vLLM中，在Llama和Falcon等主流开源LLM上进行了端到端测试。结果显示，在批大小为16时，与vLLM的标准精度核相比，MARLIN能带来高达2.8倍的端到端推理加速。

4. 扩展性设计（Sparse-MARLIN）：MARLIN的设计具有良好的扩展性，能够支持其他压缩技术。本文提出了Sparse-MARLIN，它在MARLIN的基础上集成了NVIDIA的2:4稀疏格式，相比于原始（密集）版本，可额外带来高达65%的加速，端到端加速比最高可达3.2倍。

总而言之，本文通过一种新颖的、充分利用硬件（特别是混合精度）能力的GPU核设计，证明了在批大小远大于1的情况下，也能实现权重化LLM推理的近乎最优加速。MARLIN及其稀疏变体的代码、以及vLLM的集成均已开源。

A3 背景知识

图形处理器（GPU）

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

张量核心单元（TCU）。TCU最早在Volta架构中引入，主要通过每个周期执行一次矩阵乘加（matrix multiply-and-accumulate, MMA）操作来针对机器学习工作负载。这降低了执行此类计算所需获取和解码多条指令的成本。在Ampere架构中，TCU在FP16上的性能比在FPU上运行的融合乘加（fused multiply-add, FMA）操作高出16倍。

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

现代张量核心单元

Ampere架构对TCU的扩展。Ampere GPU对其TCU进行了扩展，以支持1) 细粒度的结构化稀疏性，形成了稀疏张量核心单元（Sparse Tensor-Core Units, SPTCs），以及2) 异步复制操作。首先，结构化稀疏性通过一种新的2:4格式得到支持，其性能有望比原始TCU快2倍，比FPU快32倍。

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

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

LLM的混合精度推理

混合精度推理的优势。混合精度LLM推理有潜力减少模型巨大的内存占用，并通过在推理过程中按需动态解压静态压缩的预训练模型权重，从而加速内存受限的工作负载。

权重量化。一种标准的LLM压缩方法是仅权重量化（weight-only quantization），它降低权重$W$的存储精度，而保持层输入$X$不变。这种方法非常流行（例如，【10, Gptq: Accurate post-training quantization for generative pretrained transformers. 2022】，【15, Awq: Activation-aware weight quantization for llm compression and acceleration. 2023】，【4, The case for 4-bit precision: k-bit inference scaling laws. 2022】，【5, LLM.int8(): 8-bit matrix multiplication for transformers at scale. 2022】，【6, Spqr: A sparse-quantized representation for near-lossless llm weight compression. 2023】），因为它即使在相对较高的压缩率下也表现出卓越的准确性鲁棒性。

均匀量化。广义上，权重量化通过将浮点权重映射到有限的整数级别集合来进行有损压缩。本文关注均匀量化，即给定一个向量$v \in R^n$，我们定义：

$$Q(v, b)=\left|\frac{\boldsymbol{v}-\min (\boldsymbol{v})}{\max (\boldsymbol{v})-\min (\boldsymbol{v})}\left(2^{b}-1\right)\right|=\lfloor(v-z) / s\rfloor,$$

其中$\lfloor\cdot\rceil$表示四舍五入，$z = z(v) = \min(v)$映射到零，而$s = s(v) = (\max(v) - \min(v)) / (2^b - 1)$是缩放因子。

分组量化。重建误差可以计算为 $\epsilon_r = \|v - Q(v, b)\|_2$。我们可以通过将$v$划分为多个组并对每个组独立进行量化来改善误差，从而为每个组（例如，每128个连续值）存储$s$和$z$值。

GPTQ量化方法。本文中，实际的权重量化是通过GPTQ的一个变体【10, Gptq: Accurate post-training quantization for generative pretrained transformers. 2022】来执行的，该方法利用二阶信息来补偿量化误差，从而只产生微小的准确性下降。然而，我们强调我们的核技术独立于任何特定的量化算法。

A2 方法细节

The MARLIN Kernel

动机

高FLOPs/Bytes比率驱动量化。LLM权重量化的动机在于现代GPU具有很高的FLOPs/Bytes比率，这意味着它们执行浮点运算的速度远快于从内存中读取数据的速度。例如，一块A10 GPU的FLOPs/Bytes比率约为200【17, Nvidia a10 datasheet. 2022】。在单层矩阵乘法中，处理一个输入token每个权重需要2个FLOP，而GPU加载一个4-bit权重所需的时间内可以执行100个FLOP。因此，只要输入批大小小于$b_{opt} \approx 50$，内存加载将主导运行时。实际上，$b_{opt}$是延迟既不受内存限制也不受计算限制的批大小，即我们能以最大吞吐量实现最低延迟的点。理论上，这正是我们希望在实践中操作的批大小：更小不会带来进一步的加速，更大则不会提高吞吐量。（该分析在5.1节有更详细的阐述。）

实现挑战。然而，要实际实现一个能够同时最大化GPU所有资源（即计算和内存）的混合精度（FP16-INT4）矩阵乘法（matmul）核是一项重大挑战。在下文中，我们将通过设计MARLIN（一个极度优化的混合精度自回归线性核）来尽可能接近这个目标。

Ampere矩阵乘法

通用高性能矩阵乘法实现。我们首先描述在GPU上，特别是在Ampere级设备上实现峰值性能（均匀精度）矩阵乘法核所使用的通用概念。我们紧密遵循CUTLASS的分层并行化模型【19, Efficient GEMM in CUDA. 2024】。具体来说，我们考虑将一个$M \times K$的矩阵A与一个$K \times N$的矩阵B相乘以产生一个$M \times N$的输出矩阵C的问题。

SM级别并行化。第一步，将矩阵A划分为$M_{sm} \times K_{sm}$大小的块$A_{sm}[i_{sm}, k_{sm}]$，B划分为$K_{sm} \times N_{sm}$大小的块$B_{sm}[k_{sm}, j_{sm}]$，C划分为$M_{sm} \times N_{sm}$大小的块$C_{sm}[i_{sm}, j_{sm}]$。由于矩阵乘法的性质，所有的$C_{sm}[i_{sm}, j_{sm}]$都可以通过累加所有$k_{sm}$上的$A_{sm}[i_{sm}, k_{sm}]B_{sm}[k_{sm}, j_{sm}]$的结果来独立计算。因此，可以通过将这些$C_{sm}[i_{sm}, j_{sm}]$子问题分配给GPU的独立计算单元——即SMs——来轻松实现并行化。在这个阶段，$A_{sm}[i_{sm}, k_{sm}]$和$B_{sm}[k_{sm}, j_{sm}]$块必须从全局GPU内存中加载。同样，$C_{sm}[i_{sm}, j_{sm}]$最终必须写回全局存储，但中间的累加可以直接在寄存器中进行，我们接下来会讨论。

Warp级别并行化。在单个SM处理的子问题内部，会进行另一次等效的划分，这次的参数是$M_{wa}$, $K_{wa}$, 和$N_{wa}$。这是为了将独立的$C_{wa}[i_{sm}, j_{sm}][i_{wa}, j_{wa}]$输出累加任务分配给不同的warp。关键在于，SM块$A_{sm}[i_{sm}, k_{sm}]$和$B_{sm}[k_{sm}, j_{sm}]$可以临时存储在共享内存中，这样不同warp重复加载$A_{wa}[i_{sm}, k_{sm}][i_{wa}, k_{wa}]$和$B_{wa}[k_{sm}, j_{sm}][k_{wa}, j_{wa}]$的速度会快得多。同时，输出$C_{wa}[i_{sm}, j_{sm}][i_{wa}, j_{wa}]$保存在相应warp的寄存器中，消除了累加过程中的任何额外内存访问成本。

Tensor Core级别并行化。最终，每个warp将重复乘以$M_{wa} \times K_{wa}$和$K_{wa} \times N_{wa}$的矩阵。虽然在这个级别上相应的矩阵维度很小，但它们通常仍超过基本的Tensor Core (Mtc, Ktc, Ntc) 形状。因此，需要进行最后一步的划分。然而，与之前不同，$C_{tc}[i_{sm}, j_{sm}][i_{wa}, j_{wa}][i_{tc}, j_{tc}]$是由单个warp顺序累加的。虽然此时所有数据都在寄存器中，因此没有内存访问成本，但将$k_{tc}$的循环放在最外层仍然很重要。这是为了尽可能多地消除Tensor Core操作之间的顺序依赖，以最大化吞吞吐量。需要注意的是，实际利用Tensor Core还需要以非常特定的模式将矩阵元素分布到线程中。然而，这是一个由微架构强制规定的技术细节，而不是另一个灵活的并行化机会。

混合精度挑战

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

MARLIN的解决方案。我们的MARLIN核专门解决了上述所有挑战，最终使其能够在许多实际场景中实现接近峰值的性能。

核设计

问题设定。在下文中，我们假设矩阵A是全FP16精度，而$K \times N$的矩阵B已经被（对称地）量化为INT4，要么对N列中的每一列使用一个FP16缩放因子，要么对每列中每G个连续权重使用一个缩放因子，总共有$\lceil K/G \rceil N$个缩放因子。

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

$$(2M K_{\text{sm}} + 0.5K_{\text{sm}}N_{\text{sm}})/B_{12} < (0.5K_{\text{sm}}N_{\text{sm}})/B_{\text{gl}},$$

其中$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块的索引$i,j$, $(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$表示异或操作【20, CUTLASS convolution. 2024】。这个索引转换的另一个关键方面是，如果一个warp读取全局A瓦片的一个连续子瓦片（例如，前4行），那么它将被置换但总体上仍连续地写入共享内存。虽然没有文档说明，但这似乎是避免写入时bank冲突所必需的，正如我们分析NVIDIA profiler的输出时观察到的。这些索引计算有些复杂，动态处理可能很慢；然而，由于它们只影响相对少量的共享内存位置，并且在主循环中保持静态，我们可以将它们预计算到寄存器中，并伴随下面描述的适当展开。

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

Warp布局。在单个SM上计算$C_{sm}$必须进一步在warp之间进行细分：如果直接进行，每个warp将计算输出的一个$M \times (N_{sm} / #warps)$的瓦片。为了达到峰值计算吞吐量，我们希望至少使用四个（因为Ampere GPU有四个warp调度器），理想情况下是八个warp【26, Dissecting tensor cores via microbenchmarks: Latency, throughput and numeric behaviors. 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, Optimizing parallel reduction in cuda. 2007】来完成，这通常只会产生最小的开销。

去量化和Tensor Cores。从INT4到FP16的简单类型转换速度很慢；相反，我们遵循Kim等人【13, Who says elephants can’t run: Bringing large scale moe models into cloud scale production. 2022】的二进制操作的修改版本。我们现在以最简单的情况来说明这个过程：将位于INT16中12-15位的INT4转换为有符号的FP16值。首先，我们仅提取对应于我们INT4的位（通过与一个掩码进行AND操作），并将结果的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$的情况，我们只需要每隔一个瓦片从全局和共享内存加载一次新的缩放因子（并且在这里只在第一个子瓦片期间加载一次）。然而，编译器似乎对代码最关键部分的这种不规则性相当脆弱，导致在某些形状设置下指令排序不佳，整体速度减慢10-20%。相反，我们发现为每个子瓦片从共享内存重新加载缩放因子可以保持峰值性能。这样做增加了一些技术上不必要的共享内存加载，但有足够的额外带宽来支持这一点而没有开销，同时它也保留了编译器为非分组量化精心设计的流水线指令排序。

条带化分区。通过上述所有技术，我们可以达到近乎最优的计算和带宽性能，前提是矩阵很大并且可以完美地沿着N轴在所有SM上进行分区。在实践中，这种情况很少见。在这种情况下，标准的补救措施是也沿着K维度进行分区，但对于许多流行的层形状和GPU组合，我们需要大量的额外分割才能在没有显著波浪量化的情况下达到均匀分布。这反过来又增加了许多全局归约步骤，带来了额外的开销。相反，我们选择了一种条带化分区方案，其中一个SM处理的B的“条带”可以跨越多个$C_{sm}$瓦片（另见图5）。具体来说，我们首先确定每个SM要处理的$B_{sm}$瓦片的数量$T = \lceil #tiles / #SMs \rceil$，然后从左上角开始按列分配（最多）T个瓦片。关键是，如果我们到达一个瓦片列的底部，但当前的SM尚未拥有T个瓦片，我们就从下一个瓦片列的顶部开始分配瓦片；换句话说，条带可以跨越多列。这确保了瓦片在所有SM上的大致均匀分布，同时最小化了所需的全局归约步骤数量。这个策略类似于stream-k分区【21, Stream-k: Work-centric parallel decomposition for dense matrix-matrix multiplication on the GPU. 2023】。我们按顺序从下到上实现同一瓦片列的条带之间的全局归约。后一种方法是最有效的，因为在存在任何列溢出的情况下，最底部的SM将最快得到其结果，而最顶部的SM最慢。我们在输出缓冲区中直接以FP16进行归约，以最大化L2缓存命中率，从而最小化任何全局读取开销。这也使操作基本上是就地的，只需要一个小的额外锁缓冲区进行同步。最后，我们注意到，对于远大于64的批大小，我们可以为条带化索引计算虚拟地复制B，然后通过模运算移回原始矩阵，并将A指针前进到相应的大小为64的输入批次段。这导致在大型输入批大小（如LLM预填充期间发生）的情况下，全局归约显著减少，并在此设置下提高了计算吞吐量。

GPTQ 修改

MARLIN的量化格式。MARLIN使用的量化格式专为峰值推理效率而设计，与原始GPTQ实现【10, Gptq: Accurate post-training quantization for generative pretrained transformers. 2022】略有不同，但仍能产生高度准确的模型。我们还将两个小改进集成到GPTQ中：(a) 通过搜索最优的逐组裁剪阈值来选择组缩放因子，类似于【15, Awq: Activation-aware weight quantization for llm compression and acceleration. 2023】的方法，以及 (b) 支持可变长度的校准序列。我们发现这些修改比标准GPTQ带来了微小但一致的准确性改进，同时具有更高性能的优势。（我们也提供了一个简单的模型格式转换脚本。）

准确性结果。图6展示了我们的GPTQ变体与原始未压缩模型的困惑度（越低越好）与模型大小（以比特为单位）的对比。这表明MARLIN量化的模型在与未压缩基线相同的困惑度下，体积小了约3.33倍。虽然这不是无损的（在此位宽和组大小下，理想增益应为3.87倍），但这仍然是一个显著的改进，特别是考虑到MARLIN的高推理效率。

The Sparse-MARLIN Kernel

集成2:4稀疏性。为了进一步提高FLOPS/Byte比率，我们可以在4-bit量化权重表示的基础上集成2:4稀疏方案。作为背景，NVIDIA Ampere架构中的稀疏张量核（SPTCs）提供了一种在专门为稀疏计算设计的硬件单元上执行50%稀疏矩阵的有效方法。然而，要利用SPTCs，需要对之前描述的MARLIN核进行一些修改和扩展。

计算重构与数据布局。首先，为了适应mma.sp指令的约束——该指令能够利用SPTCs，并要求稀疏矩阵作为张量操作中的左侧（LHS）操作数【18, Nvidia instruction set. 2022】——我们设计了新的特定数据布局。具体来说，将A与B相乘的问题现在被内部重构为求解$((B^T A^T)^T)$以产生C。然而，这种重构保留了密集MARLIN核设计中描述的所有技术和优化。请注意，B可以根据需要进行离线预处理，而A可以通过ldmatrix指令及其.trans可选限定符在SMEM中即时转置，而不会导致性能下降。

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

量化的非零值。图7左侧展示了一个大小为$N \times K$的4-bit量化矩阵B，该矩阵已被修剪为2:4稀疏性。这个矩阵的压缩表示，如1a所示，其内部维度的大小将是原始矩阵的一半，即$N \times K/2$。然而，由于每个值是一个4-bit元素，我们可以在此基础上应用密集的MARLIN压缩方法，将8个元素进一步压缩到一个32-bit值中，如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列的前八个连续输出元素）。因此，这种乘法的重构不仅能在不降低性能的情况下转置存储结果，而且与基线密集设计相比，进一步提高了输出写入的效率。

元数据索引。为了从右侧（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 实验环境

• 模型:

  • Llama (Llama-2-7B, Llama-2-13B, Llama-2-70B)【30, Llama 2: Open foundation and fine-tuned chat models. 2023】
  • Falcon (Falcon-180B)【28, The Falcon family of large language models. 2023】
  • Yi-34B

• 硬件配置:

  • GPU:

    • NVIDIA A10 (推理优化型GPU，用于核心基准测试)
    • NVIDIA GeForce RTX 3090 (消费级GPU)
    • NVIDIA RTX A6000
    • NVIDIA A100 (旗舰级GPU)

  • 多GPU设置: 对大型模型（Llama-2-70B, Falcon-180B）测试了2、4、8个GPU的权重分片（sharding）场景。

• 软件配置:

  • 服务引擎: vLLM【14, Efficient memory management for large language model serving with pagedattention. 2023】，将MARLIN作为自定义核集成。

  • 对比基线:

    • PyTorch (调用CUTLASS)【22, Pytorch: An imperative style, high-performance deep learning library. 2019】
    • AWQ【15, Awq: Activation-aware weight quantization for llm compression and acceleration. 2023】
    • ExLlamaV2【8, Exllamav2: A memory efficient fork of hf transformers optimized for llama models. 2024】
    • bits-and-bytes【5, LLM.int8(): 8-bit matrix multiplication for transformers at scale. 2022】

  • 量化方法:

    • MARLIN使用本文修改版的GPTQ【10, Gptq: Accurate post-training quantization for generative pretrained transformers. 2022】进行4-bit量化，组大小（group size）为128。
    • Sparse-MARLIN使用SparseGPT【9, Sparsegpt: Massive language models can be accurately pruned in one-shot. 2023】生成INT4 + 2:4稀疏模型。

A4 实验结果

核基准测试 (Kernel Benchmarks)

• 峰值性能对比 (图1): 在NVIDIA A10上使用一个大型72k×18k矩阵进行测试，MARLIN在批大小为1时达到接近理想的3.87倍加速，并且在批大小增加到16-32时仍能保持该水平。相比之下，其他开源核（如PyTorch, AWQ, ExLlamaV2, bits-and-bytes）的性能随着批大小的增加而迅速下降。
• 真实层形状性能 (图9): 在批大小为16时，测试了流行模型（Llama, Falcon）中不同层形状的性能。MARLIN在各种GPU（A10, RTX 3090, A6000, A100）上均表现出强大的加速效果。在消费级GPU（如RTX 3090）上加速比更高，而在旗舰级A100上相对较低，因为后者的内存带宽和计算能力更强，使得开销占比更大。
• 持续性能 (图10): 在锁定GPU基础时钟频率的生产环境中，MARLIN的相对加速比几乎不受影响，而其他核的性能则显著下降。
• 大批次（Prefill）性能: 在A100上测试非常大的批次时，MARLIN的性能与未压缩的计算密集型矩阵乘法几乎相同，直到批大小1024，在更大尺寸下也只有约10%的性能下降。
• Roofline分析 (图11): 对MARLIN在A10上的分析显示，当批大小小于64时，核是内存受限的；当批大小更大时，则转为计算受限，这验证了设计的初衷。分析表明，MARLIN在各种矩阵尺寸和算术强度下都实现了很高的硬件利用率。
• Sparse-MARLIN性能 (图12, 13): 与密集版本相比，Sparse-MARLIN在峰值和持续性能测试中都展现了进一步的加速，验证了MARLIN设计对其他压缩格式（如2:4稀疏）的可扩展性。

端到端实验 (End-to-End Experiments)

• 准确性 (表1): 在Llama-2-7B上的测试表明，与基线模型相比，INT4量化（用于MARLIN）和INT4+2:4稀疏量化（用于Sparse-MARLIN）的模型准确性得到了很好的恢复。
• vLLM集成性能 (图14): 将MARLIN集成到vLLM后，在A10上对Llama-2-7B进行端到端生成测试。结果显示，MARLIN的加速比最高可达约3倍，而Sparse-MARLIN在MARLIN的基础上还能提供额外的1.2倍加速。
• 跨GPU和模型类型的性能 (表2): 在多种模型（Llama-2, Yi-34B, Falcon-180B）、GPU（A10, 3090, A6000, A100）和多GPU分片设置下，MARLIN均展现了性能提升。最大的加速出现在内存受限（批大小≤16）且GPU性能较弱或数量较少的情况下，表明MARLIN在资源受限的环境中尤其有益。
• 服务基准测试 (图15, 16): 在模拟的服务器-客户端设置中，使用NVIDIA RTX A6000测试Llama-2-7B。结果显示，在不同请求强度（QPS）下，MARLIN实现了约2.8倍的TPOT（每输出Token时间）延迟降低，而Sparse-MARLIN则达到约3.3倍。随着QPS增加，相对FP16的加速比保持稳定甚至略有增加。MARLIN同样也改善了TTFT（首个Token时间）。

A7 补充细节

相关工作

高效量化LLM推理支持。本文的工作与为量化LLM推理提供高效支持的领域密切相关。先前已有大量关于准确的LLM权重量化的研究，流行的方法包括GPTQ【10, Gptq: Accurate post-training quantization for generative pretrained transformers. 2022】和AWQ【15, Awq: Activation-aware weight quantization for llm compression and acceleration. 2023】，以及对四舍五入（round-to-nearest, RTN）量化【4, The case for 4-bit precision: k-bit inference scaling laws. 2022】的探索，后者通常准确性较低。MARLIN的并行化方法可以推广到这些量化方法。事实上，自从我们最初发布针对GPTQ格式的核以来，vLLM中已经独立引入了一个支持AWQ的MARLIN版本【31, vllm project pull request #6612: Add support for awq marlin. 2024】。

权重与激活同时量化。更广泛地看，LLM量化方法也可以考虑同时压缩权重和激活（activations）【5, LLM.int8(): 8-bit matrix multiplication for transformers at scale. 2022】，一些先进的方法如SmoothQuant【32, Smoothquant: Accurate and efficient post-training quantization for large language models. 2022】或QuaRot【1, Quarot: Outlier-free 4-bit inference in rotated llms. 2024】。然而，激活的量化往往更为复杂，因为会出现大的“离群值”【5, LLM.int8(): 8-bit matrix multiplication for transformers at scale. 2022】。因此，这些方法通常要么目标是较低的8-bit精度，要么需要更复杂的额外处理步骤，例如通过旋转矩阵【1, Quarot: Outlier-free 4-bit inference in rotated llms. 2024】。MARLIN方法可以扩展到这种情况，例如，最近一个独立的MARLIN后续工作已将我们的方法扩展到激活量化为8-bit而权重为4-bit的情况【34, Qqq: Quality quattuor-bit quantization for large language models. 2024】。

A5 结论

本文介绍了MARLIN，一种为LLM生成式推理实现混合精度核的通用方法。该方法通过利用新的GPU硬件指令和并行化技术，实现了近乎最优的效率。具体而言，我们已经证明MARLIN及其稀疏对应版本Sparse-MARLIN达到了接近最优的单层效率，并且在实际部署场景中，能够在精度影响适中的情况下，带来高达3倍的加速。

未来的工作方向包括：
1. 研究MARLIN对近期提出的、更复杂的“极端”压缩技术的支持，例如向量量化（vector quantization）【3, Quip: 2-bit quantization of large language models with guarantees. 2023】、【7, Extreme compression of large language models via additive quantization. 2024】，这些技术需要更复杂的解压过程。
2. 探索MARLIN对其他形式混合精度的支持，例如由激活压缩或稀疏性产生的混合精度。