IBM：Krish Agarwal、Rishi Astra、Adnan Hoque、Mudhakar Srivatsa、Raghu Ganti
Meta：Less Wright、Sijia Chen
量化是一种通过压缩模型权重并在较低精度数据类型中执行（更快）计算来提高模型推理速度的方法。然而，由于异常值的存在，量化可能会导致精度损失。最近的工作，如 QuaRot、SpinQuant 和 FlashAttention-3，介绍了提高 LLM 中 INT4、INT8 和 FP8 量化数值精度的方法。这些方法依赖于 Hadamard 变换。在本博客中，我们介绍了 HadaCore，一个 Hadamard 变换 CUDA 内核，可在 NVIDIA A100 和 H100 GPU 上实现最先进的性能。与 Dao AI Lab 的 Fast Hadamard Transform Kernel 相比，我们的内核实现了 1.1–1.4 倍 和 1.0–1.3 倍 的加速，峰值增益分别为 3.5 倍 和 3.6 倍。我们利用硬件感知的工作分解，受益于张量核心加速，同时保持量化误差的减少。
图 1：HadaCore 与 Dao AI Hadamard CUDA 内核的速度提升。在使用 128 次旋转和 840 万个元素的情况下，在 A100 上实现了 3.46 倍的峰值增益。
背景
QuaRot 和 SpinQuant 都提出了提高 LLM 中 INT4 和 INT8 量化数值精度的方法。这两种方法都旋转模型激活，因为旋转在统计学上可能减少异常值的幅度，因为它将极端值“分布”到其他（不太极端的）维度，并且旋转也是一个使用旋转矩阵的逆矩阵容易反转的操作。这些方法还可以提高 FP8 推理精度，例如在 FlashAttention-3 中。
图 2. Transformer 模块，显示 QuaRot 中的在线（红色）和离线旋转（蓝色）
应用这些旋转矩阵会引入模型运行时开销，这是由于图 2 中显示的在线操作造成的。这些旋转可以通过矩阵乘法应用，但增加的开销会减少量化带来的好处。因此，QuaRot 和 SpinQuant 选择使用 Walsh-Hadamard 矩阵，这是一种特殊的旋转矩阵，可以使用 快速 Walsh-Hadamard 变换 算法比矩阵乘法更快地应用。HadaCore 是此算法针对支持张量核心的 NVIDIA GPU 的优化实现。
张量核心加速 Hadamard 变换
HadaCore 利用 NVIDIA 张量核心，它是 NVIDIA GPU 上专门用于矩阵乘法的计算单元。为了实现这一点，我们的内核对快速 Walsh-Hadamard 算法执行硬件感知的工作分解。这种工作分解确保我们可以利用在张量核心芯片上执行的 MMA PTX 指令。HadaCore 对输入数据的块应用 16×16 Hadamard 变换。然后可以将计算卸载到 FP16 张量核心，并使用 mma.m16n8k16 指令。HadaCore 的 warp 级并行化如下所示。
图 3：HadaCore 并行化，1x256 向量（行）被大小为 256 的 Hadamard 矩阵旋转。
我们使用 warp 级张量核心操作并行处理 256 个元素的片段，以实现高达 256 大小的 Hadamard 变换。对于更大的尺寸，我们在 warp 之间混洗数据并重复。
微基准测试
我们在 NVIDIA H100 和 A100 GPU 上，针对不同的 Hadamard 和输入张量大小，对 HadaCore 与 Dao AI Lab Hadamard Kernel 进行了基准测试。
图 4：NVIDIA A100 上 HadaCore 内核相对于 Dao AI Lab Fast Hadamard Kernel 的速度提升
NVIDIA A100 的颜色编码速度提升表，绿色 = 相对于基线的速度提升
图 5：NVIDIA H100 上 HadaCore 内核相对于 Dao AI Lab Fast Hadamard Kernel 的速度提升
NVIDIA H100 的颜色编码速度提升表，绿色 = 相对于基线的速度提升
我们在图表中展示了速度提升，随着输入张量大小（标记为元素计数）的增加而增加。元素计数是我们要旋转的目标矩阵中的元素数量。例如，在多头注意力机制中
查询 (Q)、键 (K) 和值 (V) 张量是大小为 4D 的张量
(batch_size, seq_len, n_heads, head_dim)
大小为 head_dim 的 Hadamard 矩阵应用于这些激活张量，因此我们将其称为使用大小为 head_dim 的 Hadamard 矩阵，元素计数为
batch_size*seq_len*n_heads*head_dim。
注意力模块中查询旋转的常用元素计数
| 模型 \ Tokens | 预填充 | 解码 |
| Llama-2 70b | 33,554,432 个元素 128 Hadamard 大小 （1 批次 * 64 头 * 4096 个 tokens * 每个头每个 token 128 维嵌入） |
8192 个元素 128 Hadamard 大小 （1 批次 * 64 头 * 1 个 token * 每个头每个 token 128 维嵌入） |
| Llama-3 8b | 33,554,432 个元素 128 Hadamard 大小 （1 批次 * 32 头 * 8192 个 tokens * 每个头每个 token 128 维嵌入） |
4,096 个元素 128 Hadamard 大小 （1 批次 * 32 头 * 1 个 token * 每个头每个 token 128 维嵌入） |
HadaCore 在 A100 上实现了 1.1–1.4 倍 的速度提升，在 H100 上实现了 1.0–1.3 倍 的速度提升，与 Dao AI Lab 的 Fast Hadamard 内核相比，峰值增益分别为 3.5 倍和 3.6 倍。对于 H100 上较小的尺寸，HadaCore 的增益有所降低。对于未来的工作，我们计划结合使用 Hopper 特定功能，如 TMA 和 WGMMA，以提高 H100 的性能。
MMLU 基准测试
我们评估了 Llama 3.1-8B 推理工作负载上的 MMLU 分数，其中 FlashAttention 计算以 FP8 执行。新一代 NVIDIA Hopper GPU 配备了 FP8 张量核心，与 FP16 相比，可提供显著的计算增益。
我们的结果表明，当与 FP8 FlashAttention 等优化相结合时，使用 HadaCore 有利于保持精度。
| 格式 | 方法 | Llama3.1-8B
平均 5-Shot MMLU 准确率 |
| Q、K、V：FP16
FlashAttention：FP16 |
N/A | 65.38 |
| Q、K、V：FP16
FlashAttention：FP8 |
无 Hadamard | 64.40 |
| Q、K、V：FP8
FlashAttention：FP8 |
HadaCore | 65.09 |
| Q、K、V：FP8
FlashAttention：FP8 |
Dao AI Fast Hadamard Kernel | 65.45 |
表 1：使用 Hadamard 变换的 FP16 基线和 FP8 注意力的 Llama3.1 8B 的 MMLU 分数，比较了使用显式 Hadamard 矩阵乘法的实现与 HadaCore 的实现（越高越好）
从上面的 MMLU 分数中，我们注意到对于使用 FP8 注意力的 Llama3.1-8B 推理，HadaCore 改善了以较低精度计算注意力机制引入的量化误差。
结论
我们展示了通过将 Fast-Walsh Hadamard 算法移入 CUDA 内核所实现的速度提升，该内核利用张量核心加速，并且在 NVIDIA A100 和 H100 上分别比 Dao AI Fast-Hadamard 内核实现了 3.5 倍 和 3.6 倍 的峰值速度提升。
此外，我们在 MMLU 基准测试中表明，使用 HadaCore 进行旋转与 Fast-Hadamard 内核保持了相似的量化误差减少，同时提供了计算加速。
未来工作
我们计划实现 Triton 版本的内核，并尝试更先进的技术，例如内核融合，以支持融合的 Hadamard 变换和量化。此外，我们计划扩展我们的内核以支持 BF16 张量核心计算。
