MixPE: Quantization and Hardware Co-design for Efficient LLM Inference

[ysh329]

基于 Transformer 的大型语言模型（LLMs）随着模型规模的不断扩大取得了显著的成功，但由于其对计算和内存的巨大需求，部署仍然具有挑战性。量化已成为一种很有前途的解决方案，而针对大型语言模型的最先进量化算法引入了对混合精度矩阵乘法（mpGEMM）的需求，即低精度权重与高精度激活相乘。尽管有其优势，但目前的硬件加速器如 GPU 和 TPU 缺乏对高效 mpGEMM 的原生支持，导致在主要顺序循环中的反量化操作效率低下。
为了解决这一限制，我们引入了 MixPE，这是一种专门的混合精度处理元件，专为在大语言模型推理中进行高效的低位量化而设计。MixPE 利用两项关键创新来最大限度地减少反量化开销并释放低位量化的全部潜力。首先，认识到在每个量化组内比例因子和零点是共享的，我们建议在每组矩阵乘法后进行反量化，显著减少反量化开销。其次，MixPE 不依赖传统的乘法器，而是利用高效的移位和加法操作进行乘法运算，优化了计算和能源效率。我们的实验结果表明，MixPE 比最先进的量化加速器速度快 2.6 倍，能源消耗减少 1.4 倍。

[ysh329]

人工智能行业正在应用许多技术来降低模型成本。量化是这些技术中的一个关键实例，其中单个张量值从全精度 FP32 转换为成本更低的数字标准。鉴于其广泛的硬件支持，最流行的量化格式包括 INT4、INT8 和 FP16。

基于这些数据格式，整数量化算法可分为三类：

• 8 位权重和 8 位激活（W8A8）
• 4 位权重和 16 位激活（W4A16）
• 4 位权重和 8 位激活（W4A8）

后两种混合精度方法，即用低精度权重和高精度激活量化大型语言模型，变得特别有吸引力，因为它们在节省内存和计算成本的同时保持了模型精度[1,2]。这是因为激活比权重更难量化。更具体地说，权重分布非常均匀和平坦，相对容易量化。先前的工作表明，用 8 位、4 位甚至 1 位量化大型语言模型的权重不会显著降低精度[3–5]。相反，激活是动态生成的，具有高方差，明显表现为动态异常值[6–8]。这些异常值可能导致显著的精度下降。

因此，减少反量化开销对于在大语言模型推理中实现最佳吞吐量至关重要。受此启发，我们引入了一种混合精度处理单元 MixPE，通过专用的混合精度矩阵乘法（mpGEMM）解决方案来加速低比特大语言模型推理。与传统的基于反量化的方法不同，MixPE 直接执行混合精度矩阵乘法，并将反量化推迟到每组矩阵乘法计算之后。

这种方法利用了每个量化组内比例因子和零点是共享的这一事实，从而可以更有效地处理低精度权重。此外，为了充分利用低精度计算的优势，MixPE 用高效、低功耗的移位加法操作取代传统乘法器，用于将低比特权重与高比特激活相乘。通过整体协同设计软件（在混合精度矩阵乘法后进行反量化）和硬件，即基于移位加法的处理单元（PE），MixPE 不仅减少了反量化开销，还利用低精度算术实现了高吞吐量、高能效的大语言模型推理。与传统的基于 INT8 的张量处理单元（TPU）相比，MixPE 在大语言模型推理上可以实现超过 4 倍的加速（如图 1 右图所示）。此外，我们提出了一个参数化的设计空间探索（DSE）框架，该框架能够评估各种矩阵乘法加速器。这个框架使我们能够绘制出帕累托前沿，突出数值保真度和硬件效率之间的最佳权衡。我们的主要贡献如下：

1. 设计了一种新颖的硬件加速器 MixPE，通过协同优化量化方案和处理单元（PE）设计，以最小的反量化开销有效地处理混合精度矩阵乘法（GEMM）。
2. 提出了一个 DSE 框架，该框架将设计空间形式化为数值保真度和硬件成本之间的权衡。通过对设计变量的详尽扫描，我们确定 MixPE 为数值精度和性能的帕累托前沿。

帕累托前沿（Pareto Frontier）是一个在多目标优化问题中非常重要的概念。它源自于经济学中的帕累托效率或帕累托最优（Pareto Efficiency 或 Pareto Optimality）的概念，指的是在一个分配系统中，没有可能通过重新分配使某个个体更好而不使任何其他个体变得更坏的状态。简而言之，这是一种“最优”的状态，其中任何个体的利益增加都将以另一个个体的损失为代价。
在多目标优化问题中，帕累托前沿是一个区域或一组点，代表了在考虑两个或更多目标时可能的最优解的集合。在这个边界上的每一点，一个目标不能够被进一步改善而不牺牲另一个目标。换句话说，它代表了在多个目标之间取得的最佳平衡点。
对于两个目标的问题，其帕累托最优前沿通常是条线。而对于多个目标，其帕累托最优前沿通常是一个超曲面。在实际应用中，帕累托前沿帮助决策者理解不同目标之间的权衡关系，并找到最佳的解决方案集合，从而在多个目标之间实现最佳的平衡。

3. 实验结果表明，在 W4A8 量化下，MixPE 与最先进的量化加速器相比，实现了 2.6 倍的加速和 1.4 倍的节能。此外，对于 W4A16，MixPE 实现了 2.44 倍的加速，同时与传统的 FP16 乘法处理单元相比，能耗降低了 68% 。

[ysh329]

相关工作

2.1 大语言模型中的整数量化

如今，大型语言模型的规模迅速增长，这反过来又需要大量的硬件资源。例如，Llama-3-70B [9]仅存储其模型权重（FP16 格式）就需要大约 148GB 的内存，远远超出了像 NVIDIA A100 或 H100 这样的现代 GPU 的容量。这给大语言模型的部署带来了相当大的挑战。为了降低大语言模型部署中的推理成本，将浮点张量映射到离散电平的低比特整数量化已成为一种流行的方法[1,7,10]。给定 n 位来表示整数，量化过程可以公式化为：

$Q_{x}=\left\lceil\frac{x}{s}+z\right\rfloor, s=\frac{x_{max }-x_{min }}{2^{n}-1}, z=\left\lceil-2^{n-1}-\frac{x_{min }}{s}\right\rfloor, (1)$

其中𝑥是浮点值，𝑞𝑥是 n 位量化后的对应值，𝑠是比例因子，𝑧是零点。因此，反量化过程可以表示为：

$\hat{x}=\left(Q_{x}-z\right) \cdot s . \quad(2)$

对于整数量化，INT4 和 INT8 格式在低比特 LLM 推理中被广泛使用。在本文中，我们将 LMs 中的 x 位权重和 y 位激活量化表示为 WxAy 以作缩写。除了位精度之外，量化可以通过不同的粒度实现，从而在准确性和效率之间产生不同的权衡。

• 对于逐张量量化，整个张量共享一组缩放因子和零点[11]。
• 将通道表示为输入矩阵的最后一个维度，权重的逐通道量化或激活的逐标记量化意味着张量的每一行内共享𝑠和𝑧[6]。
• 将每个通道划分为几个子组（g 列），逐组量化通过在每行的每个组中使用不同的𝑠和𝑧进一步降低参数共享的程度[8]。

粒度越细，量化越精确，但开销越高。此外，分组量化通常用于提高 4 位量化的准确性。一个典型的组大小是 128[1,8,10]，其中每 128 个连续元素块被量化为一个单独的组，允许在每个组内更准确地表示数据。

2.2 低比特语言模型推理中的 mpGEMM

在仅解码器的大型语言模型中，对权重和激活使用不同的位宽会对混合精度通用矩阵乘法产生独特需求，其中权重矩阵使用较低精度，而激活矩阵保持较高精度。这种混合精度要求适用于多头注意力和前馈块中的主要通用矩阵乘法操作，当权重和激活量化位数不同时（如图 2 所示）。

然而，当前的商业硬件如 GPU 和 TPU 主要支持标准通用矩阵乘法，其中两个输入都具有相同的格式和位宽，并且缺乏对混合精度通用矩阵乘法的原生支持。现有的混合精度通用矩阵乘法方法通常分为两大类：

1. 间接混合精度通用矩阵乘法，主要是基于反量化的混合精度通用矩阵乘法。
2. 直接混合精度通用矩阵乘法，其中包括基于查找表的混合精度通用矩阵乘法和基于处理单元的混合精度通用矩阵乘法方法。

[ysh329]

基于间接量化的矩阵乘法（Indirect mpGEMM）基于量化的矩阵乘法首先将低精度权重放大到与高精度激活相匹配，然后执行传统的矩阵乘法（GEMM）[1,2,8]。图展示了一个基于量化的矩阵乘法示例，其中 INT4 权重是乘以 FP16 激活值。虽然这种方法支持各种精度组合，但反量化会引入额外的操作，这些操作可能会成为性能瓶颈。此外，由于通用矩阵乘法（GEMM）是在高精度下执行的，基于反量化的多精度通用矩阵乘法（mpGEMM）不能充分利用低精度计算的优势。

基于查找表的直接 mpGEMM 是一种直接的 mpGEMM 方法，它使用查找表（LUT）来实现 mpGEMM[12-14]。具体来说，它预先计算高精度激活与有限的低精度权重集的点积，并在生成的表中用简单的查找来替代计算。然而，由于现有 LLM 推理硬件（如 GPU）对查找表的支持有限，基于查找表的 mpGEMM 方法通常不如基于去量化的方法有效。此外，为每个潜在的权重-激活组合存储一个查找表可能需要大量内存，尤其是随着数据格式范围的增加。另一种直接的 mpGEMM 方法涉及设计专门的处理元件（PE），可以直接执行混合精度 GEMM 计算。BitFusion[15]提出利用张量位宽自适应来为更高精度的计算（例如 8 位整数）重用低精度组件（例如 4 位整数）而无需额外开销。然而，BitFusion 仅限于整数类型，这限制了其量化效率，并经常导致增加的内存和计算需求。OLAccel[16]对离群值进行稀疏高精度计算，同时对常规张量进行密集低精度计算。然而，它将张量存储在片外存储器中，这导致更长的内存访问延迟，从而降低吞吐量。ANT[17]和 Olive[7]都提出了新颖的数据格式，用于有效地在张量中存储大值。然而，这些方法需要专门的数据解码器，引入了额外的面积和计算开销。相比之下，我们的工作 MixPE 引入了一种灵活的 mpGEMM 计算设计，支持整数和浮点格式，而不会引入硬件开销。

[ysh329]

3 方法

3.1 动机

对权重和激活进行量化可以显著减少内存使用并提高计算吞吐量。尽管模型权重是预先已知的，并且通常呈现均匀分布，但由于 4 位整数格式的精度有限，实现准确表示仍然具有挑战性。为了进一步提高精度，分组量化被广泛使用。这种方法将权重矩阵划分为子组，并在每个组内进行量化。

为了将分组量化纳入传统的 GEMM 流水线，包括 TensorRT-LLM-W4A16 [2] 和 QServe-W4A8 [1] 在内的最先进量化算法需要在主循环中进行反量化操作，如图 3 的左图所示。对于一个量化的 GEMM 问题，𝑚表示序列数量（或批量大小），而𝑛和 k 对应于通道维度。m 和 n 都是可并行化的维度，而𝑘作为归约维度，需要一个顺序的主循环。主循环包括超过 100 次迭代，这在很大程度上主导了 GEMM 操作的运行时间。

在主循环中的去量化过程引入了两个显著的效率瓶颈。首先，去量化需要额外的对大型的 $n ×k$ 权重矩阵进行乘法和减法运算（参见公式（2））。在大语言模型推理中，批量大小 m 通常较小（例如，8、16、32），而 n 和 k 则大得多（例如，2048 或 4096）。在每次迭代中放大这样一个大型矩阵会增加大量开销，很快就会成为性能瓶颈。其次，图 3 左图中的通用矩阵乘法（GEMM）运算仍然以高精度执行。这种高精度计算不仅降低了低位量化带来的潜在性能提升，还导致功耗和内存带宽使用量增加。

我们在图 3 的右图中预览了我们的组量化 mpGEMM 设计。观察到每个组共享一组共同的缩放和零点值，我们建议先执行混合精度 GEMM（步骤 1），然后进行每组反量化（步骤 2）。为实现这一点，我们设计了一个针对高效低位计算进行优化的混合精度处理元件，充分发挥低精度算术的全部潜力。然后对每个组的点积结果进行反量化并累加以产生最终输出（步骤 3），从而显著降低主循环中的反量化开销。

[ysh329]

3.2 GEMM 后反量化

考虑$k$维权重和激活向量：$w=\left[w_{0}, w_{1}, ..., w_{k-1}\right], x=\left[x_{0}, x_{1}, ..., x_{k-1}\right]$, (3)其中 $w$ 是权重矩阵 $w_{n ×k}$ 的行向量，$x$是激活矩阵 $X_{m ×k}$ 的行向量。基于公式（2），它们的内积可以表示为：$x w^{\top} \approx \sum_{i=0}^{k-1} x_{i} \cdot s_{w_{i}}\left(Q_{w_{i}}-z_{w_{i}}\right)$，其中 $s_{w_{i}}$ 和 $z_{w_{i}}$ 分别是 $w_{i}$ 的比例因子和零点。

这里 $x_{i}$ 和 $Q_{w_{i}}$ 具有不同的精度，$Q_{w}$ 为 INT4，$x_{i}$ 为 INT8 或 FP16。为了在 GPU 上执行 GEMM，$Q_{w_{i}}$ 在与 $x_{i}$ 相乘之前必须进行去量化，对应于图 3 中的左图。我们特别观察到，当采用分组量化时，比例因子和零点在每个组内共享。设组数为 $N_{group }=k / g$，其中$g$是组大小，并将第$n$组中的元素索引集表示为 $G_{n}$。因此，我们可以将公式（4）中的点积表示如下：

$\begin{aligned} x w^{\top} & \approx \sum_{n=0}^{N_{group }-1} s_{G_{n}}\left(\sum_{j \in G_{n}} Q_{w_{j}} x_{j}-z_{G_{n}} \sum_{j \in G_{n}} x_{j}\right), \ & \approx \sum_{n=0}^{N_{group }-1}\left(s_{G_{n}} \sum_{j \in G_{n}} Q_{w_{j}} x_{j}-s_{G_{n}} z_{G_{n}} \sum_{j \in G_{n}} x_{j}\right) . \end{aligned}$。

Here 𝑠𝐺𝑛 和 7e9659ae - a4bb - 4583 - 8ac1 - 1f4c1cd70472 分别代表 e06b9af6 - a732 - 4bbd - 9edb - ada94ac30477 的比例因子和零点。从公式（5）中，我们观察到两个关键点。首先，如果 ba20d128 - bec5 - 44be - 8461 - 1981bd64eaaa 和.8302373a - c4b2 - 4318 - b17e - 0968884326ab 之间的乘法可以直接进行，那么可以在组内点积之后进行反量化，而不会损失精度。这种方法将反量化操作的频率降低到原始频率的 e1654746 - 7b02 - 4673 - bf4c - a59dedfbcaf1，显著减少了主循环中的反量化开销。其次，公式（5）中的第一项的计算复杂度为 31a3e855 - 0781 - 430f - a58b - ad46dc0a5981，而第二项为 3637634f - 52f4 - 421c - bb9b - 68b58800583a。因此，大多数现有的加速器，对于 Transformers [18, 19]，优先加速内积项，将剩余的求和项委托给软件或其他专门的计算单元。因此，为混合精度点积设计一个高效的计算引擎对于在 LLM 推理中实现最佳的 mpGEMM 性能至关重要。

[ysh329]

3.3 混合精度处理单元

最先进的 LLM 推理量化算法通常将模型权重量化为低比特整数格式，实现高理论吞吐量并减少内存占用。在计算机系统中，整数以固定位宽的二进制表示，其中每个位可以是 0 或 1。对于 4 位整数（INT4），这允许有 $2^{4}=16$ 种可能的值，使 INT4 能够表示多达 16 个不同的数字。对于 W4A8 量化，我们在权重上采用无符号 4 位（UINT4）量化，如 QServe [1]，并引入专门为量化 GEMM 设计的特殊混合精度处理元件设计：$\begin{aligned} w & (UINT4 ) * x(I N T 8) \ = & \left(w_{0} * 2^{0}+w_{1} * 2^{1}+w_{2} * 2^{2}+w_{3} * 2^{3}\right) * x \ = & w_{0} * x<<0+w_{1} * x<<1+w_{2} * x<<2 \ & +w_{3} * x<<3 \ = & \sum_{i=0}^{3} 1\left(w_{i}\right)(x<<i), \end{aligned}$ 葡萄酒 $1(w_{i}):= \begin{cases}1, & , 0, & ,\end{cases}$ 如果 w 的第 i 位是 1，否则，其中 w 和𝑥分别表示 UINT4 权重和 INT8 激活中的值，并且<<表示整数数据格式的左移运算符。由于整数表示的结构，原始乘法可以有效地被位移动所取代。具体来说，向左移动 n 位等同于乘以 $2^{n}$。基于等式（6），我们使用移位和加法操作实现 MixPE，实现具有 INT4 权重和 INT8 激活的高效混合精度 GEMM。理论上，与张量核心中的当前 INT8 乘法模块相比，W4A8 的 MixPE 可以实现 2 倍的加速、25%-50%的内存成本节省和 25%-50%的通信节省，这对于扩展下一代硬件加速器非常有前景。然后，对于 W4A16 操作，尽管按位移动不能替代乘法器，因为它们不适用于浮点值，但仍然存在有效的技术来实现快速的 2 的幂次方缩放而无需对于传统的高功耗乘法器。

W4A16 的 MixPE 设计如下：$\begin{aligned} & w(U I N T 4) * x(F P 16) \ & =\left(w_{0} * 2^{0}+w_{1} * 2^{1}+w_{2} * 2^{2}+w_{3} * 2^{3}\right) * x \ & =\sum_{i=0}^{3} 1\left(w_{i}\right)\left(x \otimes 2^{i}\right), \end{aligned}$ whuc $1(w_{i}):= \begin{cases}1, 0, & \end{cases}$ ih，其中 ∞ 代表一个专门的 RTL 模块，用于高效地进行浮点值和 2 的幂的乘法。FP16 遵循 IEEE 754 [20]标准，采用 1-5-10 位格式：1 个符号位、5 个指数位和 10 个尾数（小数）位，实际值为：value $(FP 16)=(-1)^{sign } ×2^{E- bias } \times($ 1.mantissa。采用这种格式，将 FP16 值按 2 的幂进行缩放仅涉及调整指数位，而符号位和尾数保持不变。具体来说，我们首先对指数位执行移位操作，然后进行按位加法，从而实现高效的 2 的幂缩放。这个过程与 INT4×INT8 过程类似，利用相同的移位和加法方法进行高效计算。通过直接在 FP16 中修改指数，我们可以在不使用传统高功耗乘法器的情况下按 2 的幂缩放激活值。W4A16 的 MixPE 中的这种设计选择显著降低了计算开销，最大限度地减少了操作，并避免了传统乘法器的延迟和能量成本。