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，这是一款在 NVIDIA A100 和 H100 GPU 上实现顶尖性能的 Hadamard 变换 CUDA 内核。与 Dao AI Lab 的 快速 Hadamard 变换内核 相比，我们的内核实现了 1.1–1.4 倍 和 1.0–1.3 倍 的加速，峰值增益分别达到了 3.5 倍 和 3.6 倍。我们利用了硬件感知的任务分解，在减少量化误差的同时，充分利用了 Tensor Core 加速。
图 1：HadaCore 与 Dao AI Hadamard CUDA 内核的加速比。在使用 128 次旋转处理 840 万个元素的情况下，A100 上实现了 3.46 倍的峰值增益。
背景
QuaRot 和 SpinQuant 都提出了提高 LLM 中 INT4 和 INT8 量化数值精度的方法。这两种方法都会对模型激活值进行旋转，因为从统计学角度看，旋转有助于降低异常值的幅度，它将极端值“分布”到其他（不那么极端的）维度中，并且通过旋转矩阵的逆矩阵，旋转操作也非常容易反转。这些方法同样可以提高 FP8 的推理精度，例如在 FlashAttention-3 中。
图 2：显示 QuaRot 中在线（红色）和离线旋转（蓝色）的 Transformer 模块
应用这些旋转矩阵会由于图 2 所示的在线操作而引入模型运行时的开销。这些旋转可以通过矩阵乘法来应用，但增加的开销会抵消量化带来的收益。因此，QuaRot 和 SpinQuant 选择使用 Walsh-Hadamard 矩阵，这是一种特殊的旋转矩阵，利用 快速 Walsh-Hadamard 变换 算法，其应用速度比矩阵乘法更快。HadaCore 是针对支持 Tensor Core 的 NVIDIA GPU 对该算法进行的优化实现。
Tensor Core 加速的 Hadamard 变换
HadaCore 利用了 NVIDIA Tensor Cores，这是 NVIDIA GPU 上专为矩阵乘法优化的专用计算单元。为了实现这一点，我们的内核对快速 Walsh-Hadamard 算法执行了硬件感知的任务分解。这种任务分解确保了我们可以利用在 Tensor Core 芯片上执行的 MMA PTX 指令。HadaCore 对输入数据的块应用 16×16 的 Hadamard 变换。计算随后可以通过使用 mma.m16n8k16 指令卸载到 FP16 Tensor Core 上。HadaCore 的线程束（Warp）级并行度如下所示。
图 3：HadaCore 并行化，1×256 向量（行）通过大小为 256 的 Hadamard 进行旋转。
我们使用 Warp 级 Tensor Core 操作并行处理 256 个元素的片段，以实现最大 256 大小的 Hadamard 变换。对于更大的尺寸，我们会在 Warp 之间洗牌数据并重复该过程。
微基准测试
我们在 NVIDIA H100 和 A100 GPU 上，针对不同 Hadamard 和输入张量大小，将 HadaCore 与 Dao AI Lab Hadamard 内核 进行了基准测试。
图 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)
的 4D 张量。一个大小为 head_dim 的 Hadamard 矩阵被应用于这些激活张量，因此我们将其称为使用大小为 head_dim 的 Hadamard，其元素计数为
batch_size*seq_len*n_heads*head_dim。
注意力块中查询旋转的常见元素计数
| 模型 \ Tokens | 预填充 (Prefill) | 解码 (Decoding) |
| 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 倍 的加速，峰值增益分别为 3.5 倍和 3.6 倍。对于 H100 上的较小尺寸，HadaCore 的增益会下降。作为未来工作，我们计划整合 Hopper 特有功能（如 TMA 和 WGMMA），以进一步提高 H100 的性能。
MMLU 基准测试
我们在 Llama 3.1-8B 推理工作负载上评估了 MMLU 分数，其中 FlashAttention 计算以 FP8 执行。新一代 NVIDIA Hopper GPU 配备了 FP8 Tensor Core，与 FP16 相比可提供巨大的计算增益。
我们的结果显示，当与 FP8 FlashAttention 等优化结合使用时，使用 HadaCore 有助于保持精度。
| 格式 | 方法 | Llama3.1-8B 平均 5-Shot 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 的 MMLU 分数，包含 FP16 基准和使用 Hadamard 变换的 FP8 注意力，比较了显式 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 计算。