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

📅 2026/7/31 18:33:58 👁️ 阅读次数 📝 编程学习
FP8 分块量化注意力:Hopper FP8 TensorCore 适配,块缩放量化,控制低精度误差

在 LLM 推理与训练中,引入 FP8(8 位浮点数) 是将 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-3(FA3)设计了一套专门适配 Hopper FP8 Tensor Core 的分块量化(Block-wise Scaling)与精度保护机制。一、 Hopper 架构的 FP8 Tensor Core 硬件特性Hopper 架构(sm_90)的 Tensor Core 提供了对两种 FP8 格式的原生支持:FP8 E4M3(1 位符号, 4 位指数, 3 位尾数): 精度相对较高,最大表示值为±448\pm 448±448。适合表达张量数值(如Q,K,VQ, K, VQ,K,V矩阵和 Softmax 概率PPP)。FP8 E5M2(1 位符号, 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(张量级)量化:对整个QQQKKK矩阵使用单一的缩放因子sss(Scale Factor)。但在 Attention 计算中,Per-Tensor 量化存在两大致命痛点:1. 长文本(Long-Context)离群值(Outliers)在注意力机制中,某些 Token(如 Initial Tokens 或标点符号)的Q,KQ, KQ,K激活值可能会出现极高的峰值(Outliers)。如果使用统一的 Scale,为了照顾极少数的大值,绝大多数正常数值会在 FP8 的有限量化阶梯中被压缩为 0(Underflow),导致注意力注意力分布彻底失效。2. Softmax 指数放大效应Attention 得分S=QKTS = Q K^TS=QKT随后需要经过eS−me^{S - m}eSm进行指数运算。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)

  1. 动态 Block-wise 缩放因子计算对于每一个在 SRAM 中加载的QQQ块(Br×dB_r \times dBr×d)和KKK块(Bc×dB_c \times dBc×d),FA3 在硬件局部层面独立计算其最大绝对值(Amax),并导出该 Block 专属的 Scale:sQblock=FP8_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 的 WGMMA(Warp 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}}^TQblockKblockT的计算过程由 FP8 Tensor Core 极速完成,但中间累加结果保存在 FP32 寄存器 中,防止点积求和时的精度损失。延迟 Scale 乘法: 不在输入时做复杂的反向缩放,而是在 FP32 累加器结果出来后,仅需将整个 FP32 矩阵乘以标量因子(sq⋅sk)(s_q \cdot s_k)(sqsk)。高精度 Softmax 转换: 缩放修正后的SscaledS_{\text{scaled}}Sscaled在 Vector Core 中以 FP32 精度 计算 Softmax 归一化与修正因子α,β\alpha, \betaα,β,完全避免了 FP8 在 Softmax 指数运算中的数值下溢与溢出。FP8 量化PPP阵: Softmax 输出的概率矩阵PPP在与VVV矩阵执行第二次 GEMM(P⋅VP \cdot VPV)之前,再动态量化为 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~=QH,K~=KH效果: 根据正交变换性质,(Q~)(K~)T=QHHTKT=QKT(\tilde{Q})(\tilde{K})^T = Q H H^T K^T = Q K^T(Q~)(K~)T=QHHTKT=QKT,注意力得分保持完全不变!但 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 级别的大模型注意力计算。