AxCore: A Quantization-Aware Approximate GEMM Unit for LLM Inference

作者/机构：Jiaxiang Zou, Yonghao Chen, Xingyu Chen, Chenxi Xu, Xinyu Chen / The Hong Kong University of Science and Technology (Guangzhou)

A1 主要贡献

大型语言模型（LLMs）已成为现代自然语言处理的基础，但其巨大的计算和内存需求对高效推理构成了主要障碍。基于Transformer的LLM严重依赖浮点通用矩阵乘法（FP-GEMM），这在计算和带宽方面都占据主导地位。本文介绍了 AxCore，这是一个量化感知（quantization-aware）的近似GEMM单元，它结合了仅权重量化（weight-only quantization）与浮点乘法近似（FPMA），以提供高效且准确的LLM推理。

与传统的GEMM单元不同，AxCore完全消除了乘法器，取而代之的是新型脉动阵列中的低位整数加法。AxCore具有以下关键创新点：
1. 混合精度FPMA处理单元（PE）：支持直接在压缩权重和高精度激活值上进行计算，减少了数据通路宽度和PE复杂性。
2. 轻量级精度保持策略：包括子正规数（subnormal number）处理、误差补偿和格式感知量化，以减轻近似误差。
3. 脉动阵列优化：包括共享校正和归一化逻辑，以减少硬件资源消耗。

在开源LLM上的评估表明，AxCore实现了比传统FP GEMM单元高出 6.3×-12.5× 的计算密度。与最先进的基于INT4的加速器（FIGLUT和FIGNA）相比，AxCore的计算密度分别提高了53%和70%，同时提供了更低的困惑度（Perplexity）。

Figure 1: AxCore实现了显著更高的计算密度，并且在困惑度方面与传统FP GEMM核心（FPC）和最先进的INT4基加速器FIGNA相比具有相当或更好的表现。
Figure 1: AxCore实现了显著更高的计算密度，并且在困惑度方面与传统FP GEMM核心（FPC）和最先进的INT4基加速器FIGNA相比具有相当或更好的表现。

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

2.1 LLM推理中的GEMM

LLM通常由多个堆叠的Transformer解码器块组成，其中线性层（Linear layers）主导了推理的计算工作负载。
* 计算占比：如图2所示，尽管注意力机制的计算比例随序列长度增加，但在实际序列长度（10k-20k tokens）下，线性层中的GEMM操作仍占计算工作负载的主导地位（69% - 99%）。
* 优化关键：优化线性层GEMM对于提高大规模LLM推理效率至关重要。

Figure 2: OPT-175B和LLaMA-3.1-405B在不同序列长度（Batch Size 32）下，注意力机制与线性层的操作数（OPs）相对比例。
Figure 2: OPT-175B和LLaMA-3.1-405B在不同序列长度（Batch Size 32）下，注意力机制与线性层的操作数（OPs）相对比例。

2.2 仅权重量化（Weight-only Quantization）

为了减少内存和计算开销，量化技术被广泛使用。
* 仅权重量化：将模型权重压缩为低位格式（如INT4或FP4），同时保持激活值为较高精度（如FP16或BF16）。这种方法有效减少了内存占用和带宽需求，同时避免了激活值量化带来的精度损失。
* 量化公式：权重 $w$ 通过缩放因子 $s$ 映射为低位表示 $w_q$：

$$s = \frac{w_{\max}}{F_{\max}}, \quad w_q = \text{clamp}\left(\text{round}\left(\frac{w}{s}\right), -F_{\max}, F_{\max}\right)$$


其中 $F_{\max}$ 是目标格式的最大可表示值。
* 分组量化：为保持精度，通常采用分组量化（grouped quantization），将权重张量划分为较小的组（如128个元素），每组分配专用的缩放因子。

2.3 量化感知GEMM

在仅权重量化下，有两种GEMM执行策略（如图3所示）：
1. 间接GEMM（Indirect GEMM）：先将量化权重反量化为浮点值，再进行GEMM。
2. 直接混合精度GEMM（mpGEMM）：直接在低位权重和FP激活值之间执行GEMM，仅在累加输出后进行反量化。这种方法硬件效率更高，因为避免了逐权重的反量化开销。

Figure 3: GEMM计算模式的比较，使用4-bit量化权重和16-bit激活值。
Figure 3: GEMM计算模式的比较，使用4-bit量化权重和16-bit激活值。

2.4 基于整数加法的浮点乘法近似（FPMA）

浮点（FP）乘法因其复杂性而在硬件面积上成本高昂。FPMA通过用整数加法替换乘法来解决此问题。
* 浮点表示：标准浮点数 $x$ 表示为：

$$x = (-1)^{S_x} \cdot 2^{E_x - B} \cdot (1 + M_x), \quad 0 \leq M_x < 1$$


* 对数近似：基于Mitchell的对数近似 $\log_2(1+x) \approx x$，浮点数 $x$ 近似到对数域为：

$$\log_2(|x|) = E_x - B + \log_2(1 + M_x) \approx E_x - B + M_x$$
* 乘法近似：浮点乘法 $r = x \cdot y$ 近似为：
$$\log_2(|r|) = \log_2(|x \cdot y|) \approx (E_x + M_x) + (E_y + M_y) - 2B$$
* 整数实现：近似结果可以通过整数加法实现：
$$R = X + Y - B$$
其中 $X$、$Y$、$R$ 分别是 $x$、$y$、$r$ 的二进制表示。这消除了复杂的乘法器，且结果 $R$ 已经是标准的浮点值格式。

A2 方法细节

3.1 挑战与解决方案概述

尽管FPMA具有潜力，但在量化LLM推理中应用它面临挑战：
1. 硬件支持mpGEMM：传统的FPMA仅支持统一精度。扩展到混合精度（如FP16 x FP4）需要重新设计处理单元（PE）和数据通路，以解决格式对齐和偏置不匹配问题。
2. 精度保持：FPMA引入了系统误差（$\log_2(1+x) \approx x$），且在低位格式（如FP4）中，子正规数（subnormal numbers）的比例显著增加。由于子正规数没有隐含的“1”，传统的FPMA公式在数学上不再成立，导致严重误差。如图4所示，未处理子正规数的mpFPMA会导致显著的困惑度退化。

Figure 4: 不同OPT模型规模下各种计算方法的困惑度比较。未处理子正规数的FPMA和mpFPMA会导致显著的精度损失。
Figure 4: 不同OPT模型规模下各种计算方法的困惑度比较。未处理子正规数的FPMA和mpFPMA会导致显著的精度损失。

4.1 将FPMA扩展为mpFPMA

为了支持混合精度GEMM，AxCore重新设计了运算逻辑。假设激活值 $A$ 为FP16，量化权重 $W_q$ 为FP4。
* 对齐（Alignment）：首先将操作数对齐到共同的定点表示。由于FP4的尾数位数少于FP16，需要对FP4尾数进行左移（零填充）以匹配FP16的分辨率：

$$\text{Align}(w_q) = w_q \ll (\text{Mantissa}_{\text{FP16}} - \text{Mantissa}_{\text{FP4}})$$


* 偏置校正（Bias Correction）：由于指数偏置不同（例如FP16为15，FP4 E2M1为1），需要引入格式感知的偏置校正项 $B_1$：

$$B_1 = B_a + B_{w_q} - B_r$$
其中 $B_a, B_{w_q}, B_r$ 分别是激活值、量化权重和结果的指数偏置。
* 近似计算：结合对齐和偏置校正，混合精度乘积的近似结果 $R$ 为：
$$R = A + \text{Align}(W_q) - B_1$$
为了提高数值保真度，AxCore引入了一个常数补偿项 $C_1$（见4.3节），最终的mpFPMA表达式为：
$$R = A + \text{Align}(W_q) - B_1 + C_1$$
这一公式允许仅使用整数加法高效近似混合精度乘法。

4.2 mpFPMA中的子正规数处理

子正规数的问题：
子正规数表示为 $x_{\text{sub}} = (-1)^S \cdot 2^{1-B} \cdot M$。与规范化数不同，它去除了尾数中的隐含“1+”，导致FPMA的近似基础 $\log_2(1+x) \approx x$ 失效。在低位格式（如FP4）中，由于指数位极少，子正规数不再是边缘情况，而是频繁出现（例如可表示高达0.5的值）。

子正规数转换（SNC）：
AxCore提出了一种轻量级的子正规数转换（SNC）方法。该方法基于观察：子正规编码和某些规范化编码表示的值在数值上非常接近。
* 转换逻辑：如图5所示，FP4 (E1M2) 中的子正规数 "011" (数值0.75) 可以映射到规范化编码 "010" (数值0.75)，两者数值相等。SNC单元在运行时识别子正规编码，并将其替换为数值上最接近的规范化表示。
* 转换表：Table 1展示了不同尾数位宽下的映射关系。对于无法精确映射的情况（表中带下划线的值），选择最近的规范化值。
* 随机舍入：为了减轻固定方向舍入带来的系统偏差，AxCore在近似子正规数时采用随机选择策略（向上或向下舍入），平衡累积误差。

Table 1: 子正规数转换表
Table 1: 子正规数转换表

Figure 5: 子正规编码“011”及其在E1M2中等效的规范化编码“010”所表示的数值。
Figure 5: 子正规编码“011”及其在E1M2中等效的规范化编码“010”所表示的数值。

4.3 mpFPMA的误差补偿

FPMA引入了可预测的数值误差。AxCore通过分析发现误差分布高度非均匀（图6a），且针对每对尾数进行细粒度补偿成本过高。
* 基于均值的常数补偿：AxCore提出使用单一预计算的校正值 $C_1$ 来减少累积误差。该值是通过计算所有有效尾数组合的平均近似误差得出的：

$$C_1 = \frac{1}{2^{N_{M_w}} \cdot 2^{N_{M_a}}} \sum_{m_a, m_w} \epsilon(m_a, m_w)$$


这种方法没有运行时开销，因为补偿值是预先计算的，且显著恢复了精度（图6b）。

Figure 6: mpFPMA的平方误差分布。x轴为激活值尾数，y轴为权重尾数。(a) 补偿前 (b) 补偿后。
Figure 6: mpFPMA的平方误差分布。x轴为激活值尾数，y轴为权重尾数。(a) 补偿前 (b) 补偿后。

4.4 自适应格式感知量化

为了在激进量化下保持精度，AxCore支持统一框架内的多种低位FP格式。
* 块级格式选择：不同于全模型固定格式，AxCore采用块级自适应策略。对于每个权重块（如 $128 \times 64$），从候选格式（E3M0, E2M1, E1M2）中选择最小化均方误差（MSE）的格式（公式12）。如图7所示，不同层适合不同的格式（Layer 0适合E3M0，Layer 29适合E1M2/E2M1）。
* 与FPMA集成：量化和反量化过程被重新定义以适应FPMA风格的近似（公式13-16）。

$$w_q = \text{clamp}(\text{round}(w - S + B - C))$$

$$w_r = w_q + S - B + C_2$$
通过精心设计，FPMA补偿值 $C$ 和 $C_2$ 可以相互抵消，从而保持数值正确性。相比传统除法/乘法，这种基于加减法的量化减少了舍入偏差。

Figure 7: Llama2-7B选定层注意力输出张量的权重分布，展示了不同FP4格式的适用性。
Figure 7: Llama2-7B选定层注意力输出张量的权重分布，展示了不同FP4格式的适用性。

5.1 AxCore架构概览

AxCore采用权重固定（weight-stationary）数据流。
* 数据流：量化权重（如FP4）预加载并保持在PE中，高精度激活值（如FP16）水平传播。
* PreAdd单元：集中式PreAdd单元预计算中间值 $T$，通过应用校正项到激活值上：$T = A - B_1 + C_1$。该值沿行传播，最大限度减少了PE内部的逻辑重复。
* PE处理：PE内部包含子正规数转换（SNC）单元，将权重统一为共享内部格式（S1E3M2），然后与 $T$ 进行简单的整数加法。
* 后处理：包括归一化（Norm）、AxScale（FPMA基反量化）和累加器（Accumulator）。

Figure 8: AxCore脉动阵列空间架构。
Figure 8: AxCore脉动阵列空间架构。

5.2 mpFPMA处理单元（PE）细节

PE逻辑上分为近似乘法块和累加块（图9）。
1. 输入处理：PE接收量化权重 $W_q$ 和来自PreAdd的中间值 $T$。
2. SNC单元：如图10a所示，SNC单元根据FormatSel信号选择解码器，通过逻辑表将子正规编码映射为规范化值。它还包含零标志逻辑，用于处理零输入和辅助随机舍入（通过随机位控制零标志来实现交替舍入）。所有输出转换为统一的S1E3M2格式（图10c），支持自适应量化。
3. 近似乘法：SNC处理后的权重进行尾数对齐（零填充），然后通过低位整数加法器与 $T$ 相加，得到乘积 $R$。
4. 累加：Guard单元检查零输入并强制输出为零。结果 $P'_{sum}$ 传递给部分浮点加法器，与垂直传播的部分和进行累加。

Figure 9: AxCore中处理单元（PE）的架构。
Figure 9: AxCore中处理单元（PE）的架构。
Figure 10: 子正规数转换（SNC）单元。
Figure 10: 子正规数转换（SNC）单元。

5.3 脉动阵列优化

  1. 校正前移（Correction Advancing）：在mpFPMA中，校正项 $-B_1 + C_1$ 与权重值无关，仅取决于格式。AxCore将这些项的计算移出PE，放入集中式的PreAdd模块（图11b）。这使得PE只需执行轻量级加法 $R = T + \text{Align}(W_q)$。如图12所示，这避免了在每个PE中重复使用宽位宽加法器（例如15-bit），仅需一个7-bit加法器即可处理E3M2权重和FP16激活值的组合，显著节省了面积。
  2. 归一化推迟（Normalization Postponing）：为了避免PE内部昂贵的归一化逻辑（LZD、移位等），AxCore将归一化推迟到PE阵列之外的共享Norm模块。PE累加未归一化的中间和，Norm模块（图11c）处理最终的Abs、LZD、Cmp和Round操作。这在 $N \times N$ 阵列中减少了 $N$ 倍的逻辑重复。
  3. 基于FPMA的反量化（AxScale）：AxCore使用FPMA实现反量化缩放，而不是乘法器。输出 $O_q$ 通过两个整数加法进行缩放：

    $$O = O_q + S - B + C_2$$


    其中 $S$ 是缩放因子的二进制表示，$C_2$ 是补偿常数。

Figure 11: 跨PE资源共享的优化，包括PreAdd模块中的高级校正和Norm模块中的推迟归一化。
Figure 11: 跨PE资源共享的优化，包括PreAdd模块中的高级校正和Norm模块中的推迟归一化。
Figure 12: mpFPMA不同设计的比较。(b) 采用了校正前移技术。
Figure 12: mpFPMA不同设计的比较。(b) 采用了校正前移技术。

5.4 基于AxCore的LLM加速器

整个系统架构如图13所示，核心是AxCore GEMM单元，由多个处理Tile组成的2D阵列。系统包括权重缓冲区（Weight Buffer）、统一缓冲区（Unified Buffer）用于激活值、向量单元（Vector Unit）处理层级操作，以及控制单元（CTRL）。这种模块化和脉动阵列友好的设计使AxCore能够高效支持低位量化下的大规模Transformer推理。

Figure 13: 基于AxCore的LLM加速器架构。
Figure 13: 基于AxCore的LLM加速器架构。

A4 实验环境

  • 数据集：

    • WikiText-2（用于困惑度PPL评估，序列长度2048）。
    • Zero-shot基准：ARC-e, HellaSwag, PiQA, Winogrande。
    • 校准集：Pile数据集的子集。

  • 模型：

    • OPT系列 (2.7B, 6.7B, 13B, 30B)。
    • LLaMA2系列 (7B, 70B)。
    • 量化设置：4-bit权重量化，组大小128（OPT）或64（LLaMA2）。

  • 硬件配置：

    • 实现：SpinalHDL，生成Verilog RTL。
    • 工艺：Synopsys Design Compiler, TSMC 28nm。
    • 频率：1GHz。
    • 配置：64x64脉动阵列，4x4 Tiling。

  • 软件配置：

    • 仿真器：基于DNNWeaver开发的周期级仿真器。
    • 功耗模拟：CACTI (SRAM)。

  • 基线对比：

    • FPC (传统FP GEMM核心)。
    • FPMA (传统近似计算)。
    • FIGNA (SOTA FP-INT mpGEMM设计)。
    • FIGLUT (SOTA LUT-based FP-INT GEMM)。
    • Tender (SOTA Integer-based GEMM)。

A5 实验结果

6.2 面积效率

  • PE面积：如图14所示，AxCore在所有格式下均实现了最小的PE面积。在W4-FP16配置下，AxCore比FIGLUT减少了31%的面积，比FIGNA减少了32%-39%。SNC单元仅占PE总面积的3.5%。
  • GEMM单元总面积：如图15所示，AxCore的总面积始终最低。在4-bit权重场景下，AxCore比FIGLUT减少了26%-34%的总面积，比FIGNA减少了29%-37%。
Figure 14: 不同格式下处理单元（PE）的归一化面积分解。
Figure 14: 不同格式下处理单元（PE）的归一化面积分解。
Figure 15: 六种输入格式配置下GEMM单元的归一化面积分解，分解为PE阵列和共享模块（Others）。
Figure 15: 六种输入格式配置下GEMM单元的归一化面积分解，分解为PE阵列和共享模块（Others）。

6.3 计算密度

  • 结果：如图16所示，AxCore在所有格式下均提供最高的计算密度。
    • W4-FP16：AxCore比FPC高出 6.7×，比FIGNA高出 1.7× (6.7 vs 4.0)，比FIGLUT高出 1.6× (6.7 vs 4.3)。
    • W4-FP32：AxCore比FPC高出 12.5×。
Figure 16: 六种输入格式配置下GEMM阵列的归一化计算密度（TOPS/mm²）。
Figure 16: 六种输入格式配置下GEMM阵列的归一化计算密度（TOPS/mm²）。

6.4 能效

  • 结果：如图17所示，AxCore在所有评估配置中实现了最低的能耗和最高的TOPS/W。
  • 对比：AxCore平均比FPC、FPMA、FIGNA和FIGLUT分别实现了2.2×、1.5×、1.1×和1.3×的总能耗降低，以及6.4×、3.1×、1.4×和2.0×的TOPS/W提升。
Figure 17: AxCore和基线加速器在不同数据格式和模型配置下的归一化能耗。
Figure 17: AxCore和基线加速器在不同数据格式和模型配置下的归一化能耗。

6.5 准确性评估

  • 端到端困惑度（PPL）：如Table 2所示，AxCore在OPT和LLaMA 2模型上始终提供有竞争力或更优的PPL。例如，在OPT-30B上，AxCore的PPL为9.60，优于FIGNA (9.62) 和 FIGLUT (9.70)。

  • 消融研究：

    • 基础mpFPMA PPL较高（如OPT-6.7B为11.83）。
    • 添加SNC降低了PPL（11.45）。
    • 添加补偿进一步降低PPL（11.14）。
    • 完整AxCore（含格式感知量化）达到最佳（11.01）。

  • KV Cache量化：AxCore-KV（W4A16KV4）引入的精度损失极小。

  • Zero-shot性能：如Table 3所示，AxCore在LLaMA2 70B上的平均准确率为81.78%，与FP16基线（81.91%）相当，且优于INT4和FP4量化。
  • 数值精度（SNR）：图18证实SNC和补偿显著提高了信噪比（SNR）。
Table 2: OPT和LLaMA 2模型的困惑度比较。
Table 2: OPT和LLaMA 2模型的困惑度比较。
Table 3: 四个基准数据集上的Zero-shot性能。
Table 3: 四个基准数据集上的Zero-shot性能。
Figure 18: AxCore的信噪比（SNR）分析。
Figure 18: AxCore的信噪比（SNR）分析。

6.6 与非混合精度设计的比较

  • 对比Tender：如图19所示，AxCore (W4A16KV4) 在计算密度上比Tender (W8A8KV4) 高出1.72×-1.86×，且在准确性上显著更优（OPT-30B PPL: AxCore 9.79 vs Tender 14.49）。这证明了保留高精度激活值的混合精度设计的优势。
Figure 19: 与基于整数的非混合精度GEMM加速器Tender的比较。
Figure 19: 与基于整数的非混合精度GEMM加速器Tender的比较。

A6 结论

本文提出了AxCore，这是首个利用FPMA潜力进行LLM推理的架构。通过将FPMA与低位浮点量化相结合，AxCore消除了乘法器并简化了PE逻辑。通过引入子正规数转换、均值误差补偿和自适应格式感知量化，AxCore在保持高精度的同时实现了极高的硬件效率。评估表明，AxCore在计算密度上比FP基线高出达12.5倍，比INT4加速器节省50%-70%的面积。未来的工作方向包括扩展AxCore以支持自定义数据类型或基于块的格式。

A7 补充细节

A.1 Artifact Appendix 概述

本部分提供了复现论文关键结果所需的组件，包括：
1. AxCore硬件设计：SpinalHDL RTL代码。
2. 评估脚本：用于LLM准确性评估。
3. 仿真器：用于端到端性能评估的周期精确仿真器。

A.3 软件与硬件依赖

  • 硬件：评估使用了配备四个NVIDIA RTX 6000 Ada GPU的服务器。
  • 软件：依赖包括Ubuntu 22.04.5 LTS, Python 3.9, PyTorch 2.5.1, CUDA 12.4, Cacti 7.0等。

A.5 实验工作流

实验分为三个主要部分：
1. 硬件功能验证：验证AxCore硬件逻辑。
2. LLM准确性评估：复现Table 2和Table 3。通过脚本自动下载Hugging Face模型并执行AxCore评估。
3. 仿真器性能：复现Figure 17。运行提供的脚本进行所有模拟并绘制结果图。