由 IBM 和 Meta 开发

IBM：Krish Agarwal，Rishi Astra，Adnan Hoque，Mudhakar Srivatsa，Raghu Ganti
Meta：Less Wright，Sijia Chen

量化是一种通过压缩模型权重和在低精度数据类型中进行（更快）计算来提高模型推理速度的方法。然而，由于异常值的存在，量化可能会导致精度损失。最近的工作，如 QuaRot、SpinQuant 和 FlashAttention-3，引入了提高 INT4、INT8 和 FP8 量化数值精度的方法。这些方法依赖于 Hadamard 变换。在这篇博客中，我们介绍了 HadaCore，这是一个在 NVIDIA A100 和 H100 GPU 上实现最先进性能的 Hadamard 变换 CUDA 内核。我们的内核在 Dao AI Lab 的快速 Hadamard 变换内核上实现了 1.1-1.4 倍和 1.0-1.3 倍的速度提升，峰值增益分别为 3.5 倍和 3.6 倍。我们利用了硬件感知的工作分解，这得益于 Tensor Core 加速，同时保持了量化误差的降低。

Figure 1: Speedup of HadaCore vs Dao AI Hadamard CUDA kernel. A peak gain of 3.46x on the A100 is achieved using 128 rotation by 8.4M elements.

图 1：HadaCore 与 Dao AI Hadamard CUDA 内核的加速对比。在 A100 上通过 128 次旋转 8.4M 个元素实现了 3.46 倍的峰值增益。

HadaCore 内核已公开。

背景

QuaRot 和 SpinQuant 都提出了提高 INT4 和 INT8 量化在LLMs中的数值精度的方法。这两种方法都通过旋转模型激活来降低异常值的幅度，因为旋转可以将极端值“分布”到其他（不那么极端）维度上，并且旋转也是使用旋转矩阵的逆矩阵容易逆操作。这些方法还可以提高 FP8 推理精度，例如在 FlashAttention-3 中。

Figure 2. Transformer block showing online (red) and offline rotations (blue) in QuaRot

图 2. QuaRot 中的 Transformer 块，显示在线（红色）和离线旋转（蓝色）

应用这些旋转矩阵会引入模型运行时开销，因为图 2 中显示的在线操作。这些旋转可以通过矩阵乘法应用，但增加的开销会减少量化的好处。因此，QuaRot 和 SpinQuant 选择使用瓦舍-哈达玛矩阵，这是一种特殊的旋转矩阵，可以使用快速瓦舍-哈达玛变换算法比矩阵乘法更快地应用。HadaCore 是针对支持 Tensor Core 的 NVIDIA GPU 优化的该算法的优化实现。

张量核心加速哈达玛变换

HadaCore 利用 NVIDIA Tensor Core，这是 NVIDIA GPU 上优化的专门计算单元，用于矩阵乘法。为了实现这一点，我们的内核执行了硬件感知的工作分解，以实现快速瓦舍-哈达玛算法。这种工作分解确保我们可以利用在 Tensor Core 芯片上执行的 MMA PTX 指令。HadaCore 对输入数据的块应用 16×16 哈达玛变换。然后，可以使用 mma.m16n8k16 指令将计算卸载到 FP16 Tensor Core。以下显示了 HadaCore 的 warp 级并行性。

Figure 3: HadaCore Parallelization, 1x256 vectors (rows) being rotated by a size 256 Hadamard.

图 3：HadaCore 并行化，1x256 个向量（行）通过 256 大小的 Hadamard 进行旋转。

我们使用 warp 级别的 Tensor Core 操作并行处理 256 个元素的片段，以实现 256 大小的 Hadamard 变换。对于更大的尺寸，我们在 warp 之间交换数据并重复。

微基准测试

我们在 NVIDIA H100 和 A100 GPU 上，针对不同的 Hadamard 和输入张量大小，将 HadaCore 与 Dao AI Lab 的 Hadamard 内核进行基准测试。

Figure 4: HadaCore Kernel Speedup on NVIDIA A100 over Dao AI Lab Fast Hadamard Kernel

图 4：HadaCore 在 NVIDIA A100 上相对于 Dao AI Lab 快速 Hadamard 内核的内核加速。

Color coded Speedup Table for NVIDIA A100, Green = Speedup over Baseline

彩色编码的 NVIDIA A100 加速表，绿色表示相对于基线的加速

Figure 5: HadaCore Kernel Speedup on NVIDIA H100 over Dao AI Lab Fast Hadamard Kernel

图 5：HadaCore 内核在 NVIDIA H100 上的加速与 Dao AI Lab 快速 Hadamard 内核

Color coded Speedup Table for NVIDIA H100, Green = Speedup over Baseline

彩色编码的 NVIDIA H100 加速表，绿色表示相对于基线的加速

我们在图表中展示了随着输入张量大小（标记的元素数量）的增加，我们的加速效果。元素数量是我们正在旋转的目标矩阵中的元素数量。例如，在多头注意力中：

查询（Q）、键（K）和值（V）张量的大小为：

(batch_size, seq_len, n_heads, head_dim)

对这些激活张量应用大小为 head_dim 的 Hadamard 矩阵，因此我们称之为使用 Hadamard 大小为 head_dim ，元素个数为：

batch_size*seq_len*n_heads*head_dim.

注意力块中查询旋转的常见元素个数：

模型 \ 标记 预填充 解码
Llama-2 70b 33,554,432 个元素
128 雅可比矩阵大小
(1 批次 * 64 头 * 4096 令牌 * 每个令牌每个头的 128 维嵌入)		 8192 个元素
128 雅可比矩阵大小
(1 批次 * 64 个头 * 1 个标记 * 每个头每个标记 128 维嵌入)
Llama-3 8b 33,554,432 个元素
128 个哈达玛德大小
(1 批次 * 32 个头 * 8192 个标记 * 每个头每个标记 128 维嵌入)		 4,096 个元素
128 个哈达玛尺寸
(1 批次 * 32 个头 * 1 个标记 * 每个头每个标记 128 维嵌入)

HadaCore 在 A100 上实现了 1.1-1.4 倍的加速，在 H100 上实现了 1.0-1.3 倍的加速，分别达到峰值增益 3.5 倍和 3.6 倍。对于 H100 上的较小尺寸，HadaCore 的增益会降低。对于未来的工作，我们计划结合使用 Hopper 特定的功能，如 TMA 和 WGMMA，以提升 H100 的性能。

MMLU 基准测试

我们评估了 Llama 3.1-8B 推理工作负载上的 MMLU 分数，其中 FlashAttention 的计算是在 FP8 下进行的。新一代的 NVIDIA Hopper GPU 配备了 FP8 Tensor Cores，相比 FP16 提供了显著的计算增益。

我们的结果表明，结合 FP8 FlashAttention 等优化，使用 HadaCore 可以保持准确性的优势。

格式 方法 Llama3.1-8B
平均 5 次射击 MMLU 准确率
Q, K, V：FP16
FlashAttention: 半精度浮点数 		无 65.38
Q, K, V: 半精度浮点数
FlashAttention: 8 位浮点数 		无 Hadamard 64.40
Q, K, V: FP8
FlashAttention: FP8 		HadaCore 65.09
Q, K, V: FP8
FlashAttention: FP8 		Dao AI 快速哈达玛核 65.45

表 1：Llama3.1 8B 的 MMLU 分数，使用 FP16 基线和 FP8 注意力以及哈达玛变换，比较显式哈达玛矩阵乘法实现与 HadaCore（数值越高越好）

从上述 MMLU 分数中，我们注意到，对于 Llama3.1-8B 的 FP8 注意力推理，HadaCore 可以改善在低精度计算注意力时引入的量化误差。

结论

我们展示了通过将 Fast-Walsh Hadamard 算法移入 CUDA 内核并利用 Tensor Core 加速，实现了 3.5 倍和 3.6 倍的峰值加速，分别超过 NVIDIA A100 和 H100 上的 Dao AI Fast-Hadamard 内核。

此外，我们在 MMLU 基准测试中展示了使用 HadaCore 旋转可以保持与 Fast-Hadamard 内核相似的量化误差降低，同时提供计算加速。

未来工作

我们计划实现 Triton 版本的内核，并尝试更多高级技术，如内核融合，以支持融合 Hadamard 变换和量化。此外，我们计划扩展内核以支持 BF16 Tensor Core 计算。