FP8 分块量化注意力:Hopper FP8 TensorCore 适配,块缩放量化,控制低精度误差

发布时间:2026/7/31 18:24:46
FP8 分块量化注意力:Hopper FP8 TensorCore 适配,块缩放量化,控制低精度误差 在 LLM 推理与训练中引入 FP88 位浮点数 是将 GPU 计算吞吐翻倍例如在 Hopper H100 上从 FP16 的 750 TFLOPS 飙升至 1.2~1.5 PFLOPS并使显存带宽占用减半的最直接途径。然而FP8 的数值动态范围极窄E4M3 格式仅有 4 位尾数E5M2 格式动态范围虽大但精度极低。在 FlashAttention 的QKTQ K^TQKT点积与 Softmax 指数运算中直接使用 FP8 极易引发严重的数值溢出Overflow/Underflow或精度大幅崩塌。为了在 FP8 下保持与 FP16 几乎一致的 Attention 准确率FlashAttention-3FA3设计了一套专门适配 Hopper FP8 Tensor Core 的分块量化Block-wise Scaling与精度保护机制。一、 Hopper 架构的 FP8 Tensor Core 硬件特性Hopper 架构sm_90的 Tensor Core 提供了对两种 FP8 格式的原生支持FP8 E4M31 位符号, 4 位指数, 3 位尾数 精度相对较高最大表示值为±448\pm 448±448。适合表达张量数值如Q,K,VQ, K, VQ,K,V矩阵和 Softmax 概率PPP。FP8 E5M21 位符号, 5 位指数, 2 位尾数 动态范围极大与 FP16 相同但精度较差。通常用于梯度Gradients或特定中间量。在 FA3 中前向传播Forward Pass主要采用 FP8 E4M3 格式来驱动 wgmma 矩阵乘法。E4M3 格式布局[ S (1 bit) | E E E E (4 bits) | M M M (3 bits) ]最大值 (Max Normal): 448最小值 (Min Normal): 2^-6 ≈ 0.0156二、 核心挑战为什么传统 Per-Tensor 量化在 Attention 中会失效在传统的神经网络量化中通常采用 Per-Tensor张量级量化对整个QQQ或KKK矩阵使用单一的缩放因子sssScale Factor。但在 Attention 计算中Per-Tensor 量化存在两大致命痛点1. 长文本Long-Context离群值Outliers在注意力机制中某些 Token如 Initial Tokens 或标点符号的Q,KQ, KQ,K激活值可能会出现极高的峰值Outliers。如果使用统一的 Scale为了照顾极少数的大值绝大多数正常数值会在 FP8 的有限量化阶梯中被压缩为 0Underflow导致注意力注意力分布彻底失效。2. Softmax 指数放大效应Attention 得分SQKTS Q K^TSQKT随后需要经过eS−me^{S - m}eS−m进行指数运算。FP8 E4M3 的最大值仅为 448如果QKTQ K^TQKT计算出的局部得分超出该范围或者量化噪声被指数函数放大会导致 Softmax 的归一化结果严重偏离真实值。三、 FlashAttention-3 的核心解决方案Block-wise 分块量化FA3 放弃了全局缩放将量化粒度与 FlashAttention 本身的分块Tiling架构深度绑定实现了 Block-wise Quantization块级量化。 Q 矩阵 K 矩阵┌──────────────────────────────┐ ┌──────────────────────────────┐│ Tile Q_0 (计算独立 Scale s_q0) │ │ Tile K_0 (计算独立 Scale s_k0)│├──────────────────────────────┤ ├──────────────────────────────┤│ Tile Q_1 (计算独立 Scale s_q1) │ │ Tile K_1 (计算独立 Scale s_k1)│└──────────────────────────────┘ └──────────────────────────────┘│▼FP8 WGMMA 矩阵乘法: S_block Q_block * K_block^T还原真实浮点值: S_real S_block * (s_q * s_k)动态 Block-wise 缩放因子计算对于每一个在 SRAM 中加载的QQQ块Br×dB_r \times dBr​×d和KKK块Bc×dB_c \times dBc​×dFA3 在硬件局部层面独立计算其最大绝对值Amax并导出该 Block 专属的 ScalesQblockFP8_Max_Valmax⁡(∣Qblock∣)448max⁡(∣Qblock∣)s_{Q_{\text{block}}} \frac{\text{FP8\_Max\_Val}}{\max(\vert{}Q_{\text{block}}\vert{})} \frac{448}{\max(\vert{}Q_{\text{block}}\vert{})}sQblock​​max(∣Qblock​∣)FP8_Max_Val​max(∣Qblock​∣)448​收益 即使整个序列中存在 Outliers影响也仅局限在包含该 Outlier 的单个 Block 内其余 99% 的 Block 依然能使用极小的 Scaling 占满 FP8 的所有量化阶梯Bit-width极大提升了量化精度。四、 Hopper FP8 矩阵乘法的解量化与 Softmax 融合在利用 Hopper 的 WGMMAWarp Group MMA 执行 FP8 矩阵乘法时FP8 Tensor Core 输出的结果通常直接累加为 FP32 梯度的累加器Accumulator。FA3 将 Block-wise 解量化Dequantization无缝融合到了 Online Softmax 的流水线中[ FP8 Q_block ] × [ FP8 K_block ] ──(FP8 WGMMA)──► [ FP32 Accumulator S_block ]│▼乘以 Block 缩放系数修正:S_scaled S_block * (s_q * s_k)│▼Vector Core 实时执行 Online Softmax:- 动态更新 max(m_old, m_new)- 计算指数 exp(S_scaled - m_new)FP8 计算FP32 累加Qblock⋅KblockTQ_{\text{block}} \cdot K_{\text{block}}^TQblock​⋅KblockT​的计算过程由 FP8 Tensor Core 极速完成但中间累加结果保存在 FP32 寄存器 中防止点积求和时的精度损失。延迟 Scale 乘法 不在输入时做复杂的反向缩放而是在 FP32 累加器结果出来后仅需将整个 FP32 矩阵乘以标量因子(sq⋅sk)(s_q \cdot s_k)(sq​⋅sk​)。高精度 Softmax 转换 缩放修正后的SscaledS_{\text{scaled}}Sscaled​在 Vector Core 中以 FP32 精度 计算 Softmax 归一化与修正因子α,β\alpha, \betaα,β完全避免了 FP8 在 Softmax 指数运算中的数值下溢与溢出。FP8 量化PPP阵 Softmax 输出的概率矩阵PPP在与VVV矩阵执行第二次 GEMMP⋅VP \cdot VP⋅V之前再动态量化为 FP8 E4M3驱动第二次 FP8 WGMMA。五、 不相干处理Incoherent Processing消除离群值除了 Block-wise 量化FA3 还引入了一项来自 SmoothQuant / QuIP 领域的数学技术——Incoherent Processing不相干变换以进一步控制 FP8 误差。原理 某些 LLM 模型在特定维度上存在极高的离群特征Outlier Channels。FA3 利用正交矩阵HHH如随机 Hadamard 矩阵对Q,KQ, KQ,K进行不相干正交旋转Q~Q⋅H,K~K⋅H\tilde{Q} Q \cdot H, \quad \tilde{K} K \cdot HQ~​Q⋅H,K~K⋅H效果 根据正交变换性质(Q~)(K~)TQHHTKTQKT(\tilde{Q})(\tilde{K})^T Q H H^T K^T Q K^T(Q~​)(K~)TQHHTKTQKT注意力得分保持完全不变但 Hadamard 变换将原本集中在少数通道的峰值离群值“平摊”到了所有维度上使数据分布变得平滑极大地降低了 FP8 的量化截断误差。六、 总结与性能对比维度FP16 / BF16 FlashAttention-3FP8 Block-wise FlashAttention-3硬件指令WGMMA (FP16/BF16)WGMMA (FP8 E4M3)理论算力 (H100)~750 TFLOPS~1.2 PFLOPS (1200 TFLOPS)量化粒度N/A (高精度)Block-wise (按Br×dB_r \times dBr​×d动态计算 Scale)SRAM Bandwidth 需求100% (16-bit 加载)降至 50% (8-bit 加载Tile 尺寸可扩至 256)精度损失基准与 FP16 相比困惑度Perplexity损失 0.1%通过 Hopper FP8 Tensor Core 原生驱动 Block-wise 动态缩放 高精度 FP32 累加器与 Softmax 融合 不相干变换FlashAttention-3 完美解决了低精度下的“内存墙”与“数值崩塌”矛盾真正实现了 PFLOPS 级别的大模型注意力计算。