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 的快速 Hadamard 变换核 的 **1.1–1.4 倍**和 **1.0–1.3 倍**加速，峰值增益分别为 **3.5 倍**和 **3.6 倍**。我们利用了硬件感知的任务分解，该分解受益于 Tensor Core 加速，同时保持了量化误差的减少。
图 1：HadaCore 与 Dao AI Hadamard CUDA 核的加速比。在 A100 上使用 128 旋转和 8.4M 元素实现了 3.46 倍的峰值增益。
背景
QuaRot 和 SpinQuant 都提出了提高 LLM 中 INT4 和 INT8 量化数值精度的方法。这两种方法都旋转模型激活，因为旋转在统计学上可能降低异常值的幅度，因为它将极端值“分布”到其他（不那么极端）维度，而且旋转也是一种使用旋转矩阵逆矩阵容易可逆的操作。这些方法还可以提高 FP8 推理精度，例如在 FlashAttention-3 中。
图 2. Transformer 块显示 QuaRot 中的在线（红色）和离线旋转（蓝色）
应用这些旋转矩阵会引入模型运行时开销，如上图 2 所示的在线操作。这些旋转可以通过矩阵乘法应用，但额外的开销会削弱量化带来的好处。因此，QuaRot 和 SpinQuant 选择使用 Walsh-Hadamard 矩阵，这是一种特殊类型的旋转矩阵，可以使用 快速 Walsh-Hadamard 变换 算法比矩阵乘法更快地应用。HadaCore 是该算法针对支持 Tensor Cores 的 NVIDIA GPU 的优化实现。
Tensor Core 加速的 Hadamard 变换
HadaCore 利用 NVIDIA Tensor Cores，它们是 NVIDIA GPU 上专门用于矩阵乘法的计算单元。为此，我们的内核对快速 Walsh-Hadamard 算法进行了硬件感知的任务分解。这种任务分解确保我们可以利用在 Tensor Core 芯片上执行的 MMA PTX 指令。HadaCore 对输入数据块应用 16x16 Hadamard 变换。然后可以使用 mma.m16n8k16 指令将计算卸载到 FP16 Tensor Core。HadaCore 的 warp 级并行化如下所示。
图 3：HadaCore 并行化，1x256 向量（行）通过大小为 256 的 Hadamard 进行旋转。
我们使用 warp 级 Tensor Core 操作并行处理 256 个元素的片段，以实现高达 256 大小的 Hadamard 变换。对于更大的尺寸，我们会在 warp 之间重新组织数据并重复此过程。
微基准测试
我们针对 Dao AI Lab Hadamard 核 在 NVIDIA H100 和 A100 GPU 上，针对不同的 Hadamard 和输入张量大小，对 HadaCore 进行了基准测试。
图 4：HadaCore 核在 NVIDIA A100 上相对于 Dao AI Lab 快速 Hadamard 核的加速比
NVIDIA A100 的颜色编码加速表，绿色 = 相对于基线的加速
图 5：HadaCore 核在 NVIDIA H100 上相对于 Dao AI Lab 快速 Hadamard 核的加速比
NVIDIA H100 的颜色编码加速表，绿色 = 相对于基线的加速
我们展示了随着图表中输入张量大小（标记为元素计数）的增加而实现的加速。元素计数是我们要旋转的目标矩阵中的元素数量。例如，在多头注意力中：
查询 (Q)、键 (K) 和值 (V) 张量是大小为
(batch_size, seq_len, n_heads, head_dim)
一个大小为 `head_dim` 的 Hadamard 矩阵被应用于这些激活张量，因此我们将其称为使用 `head_dim` 大小的 Hadamard，其元素计数为
batch_size*seq_len*n_heads*head_dim。
注意力块中查询旋转的常见元素计数
| 模型 \ 令牌 | 预填充 | 解码 |
| Llama-2 70b | 33,554,432 元素 128 Hadamard 大小 （1 批次 * 64 头 * 4096 令牌 * 每个头每个令牌 128 维嵌入） |
8192 元素 128 Hadamard 大小 （1 批次 * 64 头 * 1 令牌 * 每个头每个令牌 128 维嵌入） |
| Llama-3 8b | 33,554,432 元素 128 Hadamard 大小 （1 批次 * 32 头 * 8192 令牌 * 每个头每个令牌 128 维嵌入） |
4,096 元素 128 Hadamard 大小 （1 批次 * 32 头 * 1 令牌 * 每个头每个令牌 128 维嵌入） |
HadaCore 在 A100 上比 Dao AI Lab 的快速 Hadamard 核实现了 **1.1-1.4 倍**的加速，在 H100 上实现了 **1.0-1.3 倍**的加速，峰值增益分别为 **3.5 倍和 3.6 倍**。对于 H100 上的较小尺寸，HadaCore 的增益会降低。对于未来的工作，我们计划结合使用 Hopper 特有功能，如 TMA 和 WGMMA，以提高 H100 性能。
MMLU 基准测试
我们评估了在 Llama 3.1-8B 推理工作负载上执行 FlashAttention 计算（使用 FP8 精度）时的 MMLU 分数。新一代 NVIDIA Hopper GPU 配备了 FP8 Tensor Cores，与 FP16 相比，可提供显著的计算增益。
我们的结果表明，当与 FP8 FlashAttention 等优化结合使用时，HadaCore 在精度保持方面具有优势。
| 格式 | 方法 | Llama3.1-8B 平均 5 次 MMLU 精度 |
| Q、K、V：FP16 FlashAttention：FP16 |
不适用 | 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 快速 Hadamard 核 | 65.45 |
表 1：Llama3.1 8B 在 FP16 基线和使用 Hadamard 变换的 FP8 注意力下的 MMLU 分数，比较了显式 Hadamard 矩阵乘法实现与 HadaCore（**越高越好**）
从以上 MMLU 分数，我们注意到对于使用 FP8 注意力的 Llama3.1-8B 推理，HadaCore 改善了在较低精度下计算注意力引入的量化误差。
结论
我们展示了通过将快速 Walsh-Hadamard 算法转换为利用 Tensor Core 加速的 CUDA 核所实现的加速，在 NVIDIA A100 和 H100 上分别比 Dao AI 快速 Hadamard 核实现了 **3.5 倍**和 **3.6 倍**的峰值加速。
此外，我们在 MMLU 基准测试中表明，使用 HadaCore 进行旋转可以保持与快速 Hadamard 核相似的量化误差减少，同时提供计算加速。
未来工作
我们计划实现我们内核的 Triton 版本，并尝试更高级的技术，如内核融合，以支持融合 Hadamard 变换和量化。此外，我们计划扩展我们的内核以支持 BF16 Tensor Core 计算。