FP6-LLM: Efficiently Serving Large Language Models Through FP6-Centric Algorithm-System Co-Design

作者/机构: Haojun Xia (University of Sydney), Zhen Zheng (Microsoft), Xiaoxia Wu (Microsoft), Shiyang Chen (Rutgers University), Zhewei Yao (Microsoft), Stephen Youn (Microsoft), Arash Bakhtiari (Microsoft), Michael Wyatt (Microsoft), Donglin Zhuang (University of Sydney), Zhongzhu Zhou (University of Sydney), Olatunji Ruwase (Microsoft), Yuxiong He (Microsoft), Shuaiwen Leon Song (Microsoft)

A1 主要贡献

本文旨在解决大型语言模型（LLM）部署中的显存和内存带宽瓶颈问题，特别是针对6比特浮点（FP6）量化在现有GPU系统上缺乏高效支持的现状。尽管FP6在模型质量和推理成本之间提供了良好的权衡，但其在GPU上的实现面临两大挑战：（1）不规则位宽导致的内存访问不友好问题；（2）权重反量化过程中的高昂运行时开销。

为应对这些挑战，本文提出并实现了以下核心贡献：
* 识别了在现代GPU上支持FP6量化的重要性及其关键挑战。论文论证了相比8比特和4比特量化，FP6在推理成本和模型质量上具有更优的权衡。
* 提出了TC-FPx，这是首个全栈GPU核心设计方案。该方案为不同位宽（如FP6）的浮点权重提供了统一的Tensor Core支持，有效缓解了LLM推理中的“内存墙”问题。
* 通过集成TC-FPx，提供了名为FP6-LLM的全新端到端量化LLM推理支持。该系统在推理成本和模型质量之间实现了更优的平衡，支持包括LLaMA和OPT在内的多种流行LLM。
* 通过实验验证了FP6-LLM的卓越性能。实验表明，FP6-LLM能够在单张GPU上运行LLaMA-70b模型，其标准化推理吞吐量比FP16基线高出1.69倍至2.65倍。

A3 背景知识/关键Observation/设计原则

背景知识

大型语言模型的量化。大型语言模型（LLM）因其巨大的模型尺寸给部署带来了挑战。模型量化是一种常用技术，它通过减少表示模型权重的位数来创建一个更紧凑的模型。仅权重（Weight-only）量化方法仅降低模型权重的精度（例如，使用8位整数INT8），而激活值仍保持FP16精度。量化的主要目标是线性层（即矩阵乘法）的权重，这部分占LLM总权重的99%以上。本文中，我们使用“WxAy”来描述权重（Weights）和激活值（Activations）的精度，其中x/y分别代表权重和激活值的位宽。此外，“反量化”（dequantization）过程指的是将量化后的权重转换回FP16格式。

IEEE浮点标准。IEEE 754浮点标准定义了一种表示实数的二进制格式。每个浮点数由三部分组成：符号位（S）、指数位（E）和尾数位（M）。一个浮点数对应的值f可以通过以下公式计算：

Tensor Cores与SIMT Cores。SIMT（单指令多线程）核心负责GPU中的通用处理任务，处理包括整数运算、浮点运算、加载/存储操作等多种指令。SIMT核心执行操作于单个（或向量）数据元素上的标量（或向量）指令。而Tensor Cores是为加速矩阵乘法而设计的专用硬件。在A100【索引20，Nvidia a100 tensor core gpu architecture，2020】和H100【索引21，Nvidia h100 tensor core gpu architecture，2022】GPU上，Tensor Cores的浮点运算性能（FLOPS）分别比SIMT Cores高出16.0倍和14.8倍。此外，Tensor Cores以粗粒度的方式工作，例如，通过一条mma（矩阵乘法和累加）指令即可执行两个形状为16×16和16×8的FP16矩阵之间的乘法。

动机

FP6量化相比8比特与4比特量化的优势。尽管8比特【索引4, 38】和4比特量化【索引7, 14, 42】是当前LLM训练后量化的主流方案，但近期的算法研究【索引30, 35】表明，与这两种方案相比，FP6量化能在推理成本和模型质量之间实现更优的权衡。

(I) 比8比特量化具有更低的推理成本。与8比特量化相比，通过更激进的6比特量化可以在没有明显精度下降的情况下进一步降低LLM的部署成本。一方面，LLM的权重大小可以显著减小，比FP16基线小近2.7倍，从而减少存储模型权重所需的GPU显存，进而减少所需GPU数量并降低服务成本。另一方面，6比特量化可以更有效地加速LLM推理。鉴于LLM推理在生成token时通常是内存带宽受限的，通过减少从GPU DRAM读取模型权重的次数可以实现更快的推理速度。如图1所示，使用我们新提出的6比特量化系统设计（TC-FPx_W6A16）实现的llama-65b模型【索引32，Llama: Open and efficient foundation language models，2023】中的线性层执行速度始终比最先进的8比特量化支持（TensorRT-LLM_W8A16【索引26，Tensorrt-llm，2023】）快（最高可达1.42倍）。由于线性层是LLM中耗时最长的部分，这种加速将直接转化为端到端推理场景的性能提升。

(II) 比4比特量化具有更好的模型质量。尽管4比特量化更积极地减少了内存占用和DRAM访问，但它不可避免地会导致模型质量下降。相比之下，6比特量化可以实现近乎无损的模型压缩。如表1和表2所示，FP6在包括代码生成和零样本困惑度在内的各种任务中表现出强大且一致的性能。它在各种模型尺寸（如1B、13B和65B的LLaMA模型【索引32，Llama: Open and efficient foundation language models，2023】）上也显示出高鲁棒性。我们还发现，INT4量化严重依赖细粒度量化（FGQ）方法来维持高模型质量，而我们的FP6量化在粗粒度量化上已经表现得很好。值得注意的是，表1和表2中的数据点选自【索引35，Zeroquant(4+2): Redefining llms quantization with a new fp6-centric strategy for diverse generative tasks，2023】。综上所述，在算法层面，FP6量化是一个实用的替代方案，可以在不显著牺牲复杂任务和不同模型尺寸上模型质量的情况下，进一步普及LLM的部署。
表1: 零样本评估，平均了包括PTB、Wikitext和C4在内的五个数据集。指标：困惑度，越低越好。

设计选择与挑战

设计选择

现有系统缺乏对FP6的高效支持。尽管对高性能FP6训练后量化的需求日益增长，但目前尚无高效的以FP6为中心的系统设计能够实现其相对于4比特和8比特量化的权衡。现有线性层支持主要针对位宽是2的幂（如4位、8位、16位）的数据类型。鉴于如何在现代GPU上高效支持FP6尚不明确，本节阐述了两个重要的设计选择。

启用Tensor Cores的必要性。我们发现在执行量化LLM推理时，支持Tensor Cores至关重要。例如，我们评估了AWQ【索引14, 15】纯SIMT核心执行在不同批量大小下的性能以测试其可扩展性。如图1所示，随着推理批量大小的增加，没有Tensor Core支持的线性层（AWQ_W4A16_SIMT）的运行时性能变得极低。这背后有两个原因：一方面，如2.3节所述，传统SIMT核心在线性层执行上比Tensor Cores慢一个数量级。另一方面，SIMT核心的大部分计算能力将被用于在运行时反量化模型权重，这进一步减少了SIMT核心用于计算矩阵乘法的可用计算能力。这促使我们启用Tensor Cores进行密集的矩阵乘法计算，同时利用多功能的SIMT核心进行权重反量化。

统一内核方案优于双内核方案。WxA16量化的独特之处在于，激活矩阵使用FP16，但权重矩阵以更窄的位宽存储。然而，Tensor Cores要求权重和激活矩阵都以相同的数据类型存储，例如FP16/INT8/INT4。直接的解决方案（即双内核方案）是增加一个额外的GPU内核，在调用常规FP16内核之前将权重反量化为FP16。然而，这样的推理速度甚至会比未量化的模型慢。如图2（左）所示，线性层的执行将启动两个GPU内核，反量化后的FP16权重将被写入GPU DRAM，然后被第二个GPU内核读取，导致2倍的DRAM访问。将反量化和矩阵乘法过程融合成一个单一的GPU内核更为高效，这样可以消除对反量化权重（FP16格式的W）的读/写操作。

设计挑战

实现统一内核的挑战。基于4.1节的设计选择，设计一个在现代GPU上支持FP6×FP16矩阵乘法的统一GPU内核是具有挑战性的。一方面，现代GPU内存系统天然不支持不规则位宽（不是2的幂），因为GPU全局/共享内存的最小访问大小是每个线程8/32位，并且访问的内存地址必须对齐。Tensor Cores复杂的数据布局要求使得对不规则位宽的支持更具挑战性。另一方面，反量化计算开销很大，因为它需要大量复杂的位级操作。因此，如何将反量化融入线性层计算而不损害整体性能也是一个不容忽视的问题。

硬件不友好的内存访问。在现代GPU上执行线性层期间，模型权重应在进行相应乘法计算前从DRAM加载到寄存器。通常，模型权重分两步加载，以隐藏DRAM的高访问延迟，从而实现高性能。具体来说，模型权重首先从GPU DRAM加载并缓冲到片上内存（如共享内存）中以实现数据复用。之后，缓冲的权重再从共享内存读入寄存器进行实际计算。

不规则位宽导致访问困难。由于每个GPU线程不能直接访问其他线程的寄存器，每个线程必须自行将其所需的模型权重放入其私有寄存器中。当权重以不规则位宽（非2^n，如6位）存储时，鉴于Tensor Cores严格的数据布局要求，这个过程变得极具挑战性。如图3a所示，现代GPU架构中FP16 Tensor Cores的最小输入是一个8×8的子矩阵，每个GPU线程应在其寄存器中持有一对权重。在正常情况下，每个权重以16位存储，每对权重可以自然地以32位字的粒度从共享内存中读取。然而，在我们的工作中，每个权重以x位存储，这使得内存访问对现代GPU内存层次结构极其不友好。

片上内存访问存在未使用的比特。我们以6比特量化为例，说明访问不规则位宽权重的低效性。如图3b所示，权重已经缓冲在共享内存中，每个GPU线程需要从共享内存中读取一对权重（12位，2 * 6位）。然而，共享内存有32个内存库（bank），每个内存库在每次内存请求中输出一个32位字。因此，从共享内存中读取的大部分比特将被闲置，导致共享内存带宽的巨大浪费。例如，图3b中的T0（线程#0）只需要12位，但会读取一个32位字（W1），导致32位中有20位（62.5%）未使用。由于现代GPU内存层次结构中对齐内存访问的要求，未用比特的浪费会更加严重。如图3b所示，T2（线程#2）所需的比特分布在W1和W2中。因此，T2需要读取W1和W2，即从共享内存中读取2 * 32位。然而，最终只有6 * 2位被使用，导致64位中有52位（81.25%）被闲置和浪费。值得注意的是，由于不规则位宽，GPU DRAM和寄存器上的内存管理和访问也存在类似问题。

反量化的高计算开销。FPx到FP16的反量化运行时开销可能极高，这很容易拖慢整体执行速度。一方面，运行时需要反量化大量的模型权重，例如，在LLaMA-70b【索引33，Llama 2: Open foundation and fine-tuned chat models，2023】推理的每个LLM解码步骤中，需要反量化700亿个FPx权重。另一方面，反量化每个FPx权重的运行时开销很高，需要复杂的位操作。根据公式2，为了获得与给定FPx等值的FP16，运行时需要计算新的符号、指数和尾数。

在公式2中，$bias_{fp16} = 15$ 且 $bias_{fpx} = 2^{len(E_{fpx})-1} - 1$。FP16的符号字段与FPx的相同，FP16的尾数也可以通过在FPx的尾数后填充零来计算。更重要的是，FP16的指数应该是$E_{fp16} = E_{fpx} + bias_{fp16} - bias_{fpx}$，这在计算上更为昂贵。总而言之，如何高效地反量化FPx值也成为一个主要挑战。

A2 方法细节

概述

TC-FPx设计概览。图4比较了TC-FPx（我们设计的x位仅权重（weight-only）量化线性层内核）与传统通用矩阵乘法（GEMM）的设计（其中两个输入矩阵均为FP16）。在TC-FPx中，模型权重以减少的位数存储。因此，在寄存器层面引入了一个额外的反量化阶段（Dequant W），其中FP6权重在每个线程内使用SIMT核局部地反量化为FP16。值得注意的是，这些FP16权重不会写回共享内存，而是存储在寄存器中供后续使用，从而消除了不必要的共享内存往返访问。另一个区别是，TC-FPx使用细粒度的lds（加载共享内存）指令将x位权重从共享内存加载到寄存器，而不是使用粗粒度的ldmatrix（加载矩阵）内在函数，后者具有严格的布局要求且灵活性较差。

Ahead-of-time Bit-level Pre-packing

通过预打包解决不规则位宽内存访问问题。正如4.2.1节所述，对具有不规则位宽的权重的内存访问对现代GPU内存层次结构不友好。为了解决这个问题，我们提出了一个洞见：可以将每32个x位权重的内存读取合并，从而每个GPU线程产生x个4字节字的请求。在这种情况下，所有的内存访问都将以32位字的粒度对齐，而不是不规则的位宽。

利用权重静态特性进行预打包。由于Tensor Cores严格的数据布局要求，合并权重的内存读取并非易事，因为每个GPU线程所需的权重并非存储在连续的内存空间中。为解决此问题，我们提出通过重新排序权重矩阵内的权重并提前进行预打包来优化运行时内存访问模式。由于模型权重在模型训练和量化后是静态确定的，因此可以提前对权重应用复杂的内存布局转换，从而不引入运行时开销。此外，我们只需预打包一次权重，因此权重预打包的开销可以被每次推理服务有效摊销，变得可以忽略不计。

权重预打包的两步法。通常，权重预打包包括两个步骤。第一步，我们收集每个GPU线程所需的所有权重，并在本地将这些权重组合起来。鉴于每个GPU线程所需的权重在权重矩阵内最初并非位于连续位置（见图3a），我们必须仔细为每个GPU线程挑选权重。为每个线程挑选的权重随后按照它们在运行时被Tensor Cores消耗的相对时间顺序在本地组合。第二步，我们将整个GPU WARP（由32个GPU线程组成）所需的所有权重组合成一个统一的线性内存空间，权重将按此顺序在运行前存储在GPU DRAM中。为完全消除共享内存库冲突，我们提出以锯齿状（jagged）顺序组合每个线程的32位字。

方案的通用性。值得注意的是，本小节讨论的所有技术都与模型权重的实际位宽（全程用x表示）无关。因此，我们的权重预打包可以自然地应用于任何位宽。

步骤1：逐线程权重收集。图5展示了T0（线程#0）挑选的权重以及组合它们的顺序。我们假设WARP级别的切片大小为64×64，这意味着每个权重矩阵被划分为64×64的数据块，并以该粒度为每个WARP加载到GPU的共享内存中。每个权重块随后被进一步划分为四个切片（slice），因为权重是逐一切片从共享内存加载并用于Tensor Core计算的。此外，每个切片又被划分为四个16×16的块（chunk），因为Tensor Core在每个指令中处理16×16的数据项。在每个16×16的块内，为T0挑选并组合了四对FPx权重。如图5所示，步骤1之后我们得到了32（即WARP大小）组FPx权重。在每组内，权重被组合并连续存储，每组权重将由一个特定的GPU线程消耗。总而言之，每个64×64的权重块最终被分配给32个线程（一个WARP），每个线程将消耗128个x位权重。

步骤2：按WARP进行比特级组装。在步骤2中，我们将不同组的所有权重组装到一个统一的内存空间中。在这个比特级预打包过程中，我们暂时忽略每个比特的含义，将组合后的权重视为连续的数据进行复制。具体来说，128个x位项被视为4x个32位项。我们提出以图5所示的锯齿状顺序组装所有组的权重。首先，将每个线程的第一个32位项连接在一起。之后，将每个线程的第二个32位项连接并附加到之前的结果之后。通过重复这个过程，所有权重可以连续地存储在一个线性内存空间中，并且良好对齐（128字节对齐）。这样，所有权重可以简单地以128字节块的粒度从DRAM复制到共享内存，无需任何更改，轻松实现最优的DRAM访问。此外，这些权重在运行时也可以从共享内存中以最优性能加载。具体来说，一个WARP的线程将在每次内存请求中读取共享内存中连续的32位项，从而完全避免了内存库冲突。

SIMT-Efficient GPU Runtime

并行反量化。为了减少FP-x权重反量化的运行时开销，我们使用优化的位操作SIMT核心指令实现了FP-x反量化。此外，我们提出并行反量化多个FPx权重，通过利用每个32位寄存器内的位级并行性，进一步将SIMT开销减少4倍。
* (1) 优化的位操作：如4.2.2节所述，将FPx转换为等效的FP16时，FP16的指数应为 $E_{fp16} = E_{fpx} + bias_{fp16} - bias_{fpx}$。为简化此过程，我们采纳了【索引35，Zeroquant(4+2): Redefining llms quantization with a new fp6-centric strategy for diverse generative tasks，2023】中的数学变换，改为用 $E_{fp16} = E_{fpx}$ 计算FP16的指数。为保持正确性，结果FP16随后会乘以FP16常数 $2^{bias_{fp16}−bias_{fpx}}$。图6a展示了优化的FP16到FP6的转换过程。尽管我们仅以FP6到FP16的转换为例，但该方法可应用于任何位宽。FP16的符号字段与FPx相同。此外，为了效率，指数字段的低位和尾数字段的高位可以从FPx一同复制到FP16。FP16的其他位应用零填充。通过精心设计，我们成功地仅用两个位“与”操作、一个“移位”操作和一个“或”操作就实现了从FP6到FP16的转换，如图6b的❶所示。符号字段通过第一个“与”操作从FP6复制到FP16，同时FP16的所有其他位被初始化为零，从而无需后续对指数和尾数字段进行零填充。然后，FP6的所有位通过位“右移”操作向右移动。之后，FP6中指数的低位和尾数的高位首先通过FP6与位掩码“0x1f1f1f1f”的“与”操作被选中，然后通过位“或”操作复制到FP16中。
* (2) 位级并行：基于利用32位字内位级并行性的洞见，我们提出并行反量化多个FPx权重，进一步减少反量化的运行时开销。图6b以FP6为例展示了详细设计。32位寄存器被视为四个处理槽，每个槽独立工作，使用相同的指令但输入不同的FP6数据。在反量化开始前，四个FP6应按图中初始数据布局存储在R1（寄存器#1）中。通过代码片段❶，这四个FP6可以同时被反量化为四个FP16，其中每个FP16的前8位存储在R2中。之后，通过代码片段❷，第一个和第二个FP16被提取到R1，其后8位用零填充。最后，通过代码片段❸和❹，第三个和第四个FP16被提取到R2。

权重拆分与拼接。接下来，我们将演示如何通过精心设计的内存布局，在GPU上高效地从2+4方案【索引35，Zeroquant(4+2): Redefining llms quantization with a new fp6-centric strategy for diverse generative tasks，2023】中重构6位权重，此方法同样适用于其他位宽。
* (1) 预先权重拆分：为了以对齐的方式将权重存储在GPU的32位寄存器中，我们将每个权重拆分成几个段，每个段的位宽为2^n，例如，每个6位权重可以拆分为2+4或4+2。基于此方案，后续设计的索引计算被显著简化。注意，5.2节描述的技术可应用于任何位宽，因此2位和4位的段可以根据5.2节分别进行预打包和高效处理。
* (2) 运行时权重拼接：在反量化之前，权重首先从共享内存加载到寄存器。由于每个权重被拆分成几个段，完整的权重需要在运行时于寄存器级别进行重构。为了减少此运行时开销，我们提出并行地提取和拼接权重。如图7所示，两组寄存器用于存储32个FP6权重，其中Frag1_PT R指向包含32个2位段的两个32位寄存器，而Frag2_PT R指向包含32个4位段的四个32位寄存器。通过我们的并行拼接，四个FP6权重被同时重构，将SIMT核心指令的数量减少了4倍。如图7所示，四个2位段首先被提取到寄存器#1（❶），然后四个4位段被提取到寄存器#2（❷）。之后，寄存器#2被右移（❸），其有效位被复制到寄存器#1（❹），从而得到完整的6位权重。

整体伪代码。算法1展示了包括并行反量化和权重拼接在内的伪代码（GPU代码）。伪代码中的所有输入和输出变量都存储在寄存器中。如图7所示，算法1总共反量化了32个FP6权重。在每个外层循环中，会生成四个FP16权重，并在代码末尾用两个寄存器存储。图7中的转换（❶, ❷, ❸, ❹）分别通过算法1的第6、7、9和10行的SIMT核心操作实现。输出的寄存器数组（OutputReg）随后直接作为Tensor Cores的输入。

算法 1 权重拼接与反量化

1: 输入: int Frag1_ptr[], int Frag2_ptr[], half Scales[]
2: 输出: int OutputReg[]
3: #pragma unroll
4: for int i = 0; i < 8; i++ do
5:   // 权重提取
6:   unsigned int R1 = (*Frag1_ptr) & 0xc0c0c0c0;      // ❶
7:   unsigned int R2 = (*Frag2_ptr) & 0xf0f0f0f0;      // ❷
8:   // 权重拼接
9:   R2 = R2 >> 2;                                     // ❸
10:  R1 = R1 | R2;                                     // ❹
11:  // 移动到下一个寄存器或移位当前寄存器
12:  if i % 4 == 3 then
13:    Frag1_PT R++;
14:  else
15:    (*Frag1_PT R) = (*Frag1_PT R) << 2;
16:  if i % 2 == 1 then
17:    Frag2_PT R++;
18:  else
19:    (*Frag2_PT R) = (*Frag2_PT R) << 4;
20:  // 4路并行反量化
21:  *R2 = *R1 & 0x80808080;
22:  *R1 = *R1 >> 2;
23:  *R1 = *R1 & 0x1f1f1f1f;
24:  *R2 = *R2 | *R1;
25:  *R1 = *R2 & 0x9f009f00;
26:  *R2 = *R2 & 0x009f009f;
27:  *R2 = *R2 << 8; // R1和R2现在各有两个FP16权重
28:  // 乘以量化尺度并输出到寄存器
29:  OutputReg[i * 2] = Multiply(R1, Scales[i/2 * 2]);
30:  OutputReg[i * 2 + 1] = Multiply(R1, Scales[i/2 * 2 + 1]);

软件流水线设计

逐片（slice-by-slice）反量化以减少寄存器使用。为了减少GPU寄存器的使用，我们逐片地反量化权重。此外，我们将反量化过程无缝地融入到线性层执行的传统软件流水线中，通过有效的指令并行性完全隐藏了反量化的运行时开销。我们不是一次性反量化所有权重，而是逐片反量化FPx权重。如图8a所示，我们假设一个FPx权重块和一个FP16激活块已经从DRAM复制到共享内存中。共享内存中的整个权重块随后分几步进行反量化。在每一步中，只有一个FPx权重的切片从共享内存加载到寄存器，通过SIMT高效的GPU运行时（5.3节）反量化为FP16权重，然后存储在寄存器缓冲A1或A2中作为Tensor Cores的输入。ASlice和BSlice随后使用Tensor Cores进行乘法运算。

相较于一次性反量化整个块，我们的逐片反量化方法将存储FP16权重所需的寄存器数量减少了4倍，显著降低了寄存器压力。此外，这也为指令级并行创造了更多机会，因为一旦一个权重的切片被反量化，Tensor Cores就可以立即用于计算，而无需等待整个块完成。

通过高效重叠隐藏开销。软件流水线通过图8b中的时空图来说明，其中SIMT核心（负责反量化）、Tensor Cores（负责矩阵乘法）和GPU内存层次结构协同工作，实现了高度的指令级并行性。首先，全局内存读取使用cp.async【索引20，Nvidia a100 tensor core gpu architecture，2020】内在函数异步执行，完全与其他操作重叠。在处理完第三个切片后（k=2结束时）发出内存屏障和线程块同步，确保下一个主循环的数据在共享内存中就绪，以便在k=3时可以开始“De-quant”（反量化）和“ldmatrix”操作。其次，共享内存读取也与Tensor Core操作重叠。当计算第i个切片时，第(i+1)个切片的数据通过“De-quant”和“ldmatrix”同时从共享内存中读取。最后，用于权重反量化的SIMT核心操作也与Tensor Core操作有效重叠。在第i个切片的“De-quant”过程中，FPx权重首先使用硬件内在函数load shared (LDS) 从共享内存加载到寄存器，然后立即用SIMT核心反量化为FP16权重。同时，Tensor Cores在计算第(i-1)个切片，两者之间没有数据依赖。

A4 实验环境与结果

实验环境

• 硬件配置：
  • 内核级基准测试：在NVIDIA A100-40GB平台上进行。
  • 端到端模型推理：在NVIDIA A100-SXM4-80GB DGX平台上进行。
• 软件配置：
  • 操作系统与CUDA：使用CUDA 11.8。
  • 代码实现：TC-FPx内核使用超过1200行CUDA代码实现，基于Flash-LLM【索引37，Flash-llm github，2023】的代码。该内核可被编译为动态链接库（.so），并提供C++ API调用。
  • 集成框架：为进行端到端评估，将TC-FPx内核集成到DeepSpeed【索引19，Deepspeed github，2023】推理框架中，形成FP6-LLM系统。
• 模型与数据集：
  • 模型：评估了多种尺寸的LLaMA【索引32, 33】模型（7b, 13b, 33b, 65b, 70b）和OPT【索引41，Opt: Open pre-trained transformer language models，2022】模型（30b, 65b, 175b）。
  • 工作负载：实验主要关注LLM在token生成阶段的性能。端到端测试中，每个请求的提示长度设置为512，生成长度为1536。
• 性能评测工具：
  • 使用NVIDIA Nsight Compute【索引23，Nsight compute profiling guide，2023】测量GPU硬件单元的运行时利用率。
  • 使用NVIDIA Nsight System【索引24，Nsight system，2023】测量端到端推理延迟和各部分延迟 breakdown。

实验结果

线性层相对于8/16比特的加速效果

• 实验内容：在LLaMA和OPT模型中不同形状的线性层上，评估TC-FPx（W6A16）内核与cuBLAS（W16A16）、TensorRT-LLM（W8A16）和BitsandBytes（W4A16）等基线的性能，测试了批量大小为8、16和32的场景。
• 实验结果：如图9所示，TC-FPx的性能显著优于所有基线。与cuBLAS（W16A16）相比，TC-FPx最高可达2.6倍加速；与TensorRT-LLM（W8A16）相比，最高可达1.9倍加速；与BitsandBytes（W4A16）相比，最高可达8.9倍加速。平均而言，在批量大小为8/16/32时，TC-FPx分别比cuBLAS快2.2倍/2.2倍/2.0倍，比TensorRT-LLM快1.3倍/1.3倍/1.2倍。
• 性能分析：如图10a所示，cuBLAS基线在小批量推理时受限于DRAM带宽（利用率>80%），而Tensor Core利用率不足（<50%）。TC-FPx通过6比特量化显著减少了DRAM访问（最多2.7倍），缓解了“内存墙”问题，从而更有效地利用了Tensor Core。这解释了其性能优势。同时，分析发现BitsandBytes性能差是因为其采用了双内核方案，引入了额外的反量化内核开销。图10b显示，尽管我们的反量化设计增加了SIMT核（ALU和FMA单元）的负载，但由于高效的软件流水线设计，这些开销被成功隐藏。
  

  

  
  

  

与4比特性能的比较

• 实验内容：将TC-FPx（W6A16）与当前最先进的TensorRT-LLM中的W4A16内核（包括粗粒度和细粒度量化）在LLaMA-65b模型的线性层上进行性能比较。
• 实验结果：如图11所示，TC-FPx（W6A16）的性能与细粒度W4A16相当（在不同批量大小下，速度比为1.06倍、1.04倍、0.94倍），仅比粗粒度W4A16慢16%-24%。
• 结论：考虑到6比特量化能提供显著优于4比特的模型质量，这种微小的性能差异是一个非常值得的权衡。
  

  

端到端推理性能

• 评估指标：使用“每GPU秒生成的token数”（tokens per GPU-second）作为标准化吞吐量指标，综合考虑了执行时间和硬件成本。
  

  

• LLaMA-70b推理：
  • 结果：FP6-LLM仅用单张A100 GPU即可运行LLaMA-70b，而FP16基线需要两张。在相同的最大批量大小（32）下，FP6-LLM的标准化推理吞吐量比FP16基线高1.69倍至2.65倍（图12a）。
  • 分析：性能提升主要源于两点：1）即使使用一半的GPU数量，TC-FPx内核的矩阵乘法（MatMul）速度平均比cuBLAS快1.20倍；2）单GPU方案完全避免了跨GPU通信（NCCL）的开销（图12b）。
    

    

• OPT-30b推理：
  • 结果：在单GPU上，FP6-LLM支持的最大批量是16，而FP16基线仅为4。FP6-LLM的最大吞吐量达到319.1 tokens/GPU-sec，比FP16基线的最高吞吐量（78.8 tokens/GPU-sec）高4.05倍。在相同批量大小（1/2/4）下，FP6-LLM的吞吐量也高出1.91倍/1.84倍/1.72倍（图13a）。
  • 分析：性能提升主要来自线性层执行时间的缩短，TC-FPx内核平均比FP16的cuBLAS内核快2.39倍（图13b）。
    

    

• LLaMA-13b推理：
  • 结果：在单GPU上，FP6-LLM与FP16基线最大批量大小均为32。FP6-LLM的平均吞吐量比FP16基线高1.23倍（图14a）。
  • 分析：虽然TC-FPx内核使线性层执行速度平均快了2.11倍，但整体性能增益不如大模型显著。这是因为在较小模型上，内核启动延迟和GPU同步等非内核开销（图14b中的Others）占总时间的比例更大，削弱了整体性能提升。
    

    

A5 结论

本文介绍了TC-FPx，这是首个为不同位宽的浮点权重提供统一Tensor Core支持的全栈GPU内核设计方案，有效缓解了LLM推理中的“内存墙”问题。通过将TC-FPx内核集成到先进的推理系统中，我们推出了FP6-LLM，一个为量化LLM提供端到端支持的新系统，它在推理成本和模型质量之间取得了更优的平衡。FP6-LLM通过一系列创新技术，成功解决了硬件不友好的内存访问和高昂的反量化计算开销两大挑战，实现了更快的推理速度和显著减少的GPU内存占用。实验评估表明，FP6-LLM能够使用单块GPU推理LLaMA-70b模型，其标准化推理吞吐量比FP16基线高出1.69倍至2.65倍。此外，FP6-LLM还将OPT-30b的推理吞吐量提升了1.72倍至4.05倍。