FP8注意力计算中,调整数据顺序和缩放因子可显著提升精度,避免数值坍塌。
P-Cast Precision in FP8 Attention: Sink-Induced Collapse and the Optimality of S=2^8
- 通过反向KV块遍历与S=256缩放,防止概率矩阵下溢
- 实验显示在中等聚类强度下误差降低3到10倍
- 适合关注大模型推理优化的工程师与研究者
FP8(E4M3)加速注意力计算虽能大幅提升吞吐量,但3位尾数带来的精度挑战在于:在执行P·V矩阵乘法前将软最大概率矩阵P转换为FP8。本文分析了两个影响输出精度的关键实现选择:(1) KV块迭代顺序,(2) P转换前的静态缩放因子。我们发现正向迭代会导致‘P坍塌’——近似有Φ(Δ + δ_k - 6.93 - ln S)比例的非聚类项下溢至零,其中δ_k ≈ 1(当k_sink=4时),即块内得分最大值的期望。反向迭代可消除此现象,并在S=256时保证无下溢。进一步证明,S=256=2^8是满足三项最优条件的静态缩放因子:(i) 精确符合IEEE 754缩放,(ii) 覆盖E4M3数轴上锯齿函数dp(S)的下包络(dp=2^{-4},最小最差量化步长),(iii) 在所有2^k形式的精确缩放中提供最大正规范围覆盖(非精确缩放如448略有更高覆盖,见第5节)。两项优化已在FlashAttention-3/4中工程部署;我们的贡献在于量化解释其有效性,并给出预测内核级精度损失的闭式阈值Δ_c = 6.93 + ln S - δ_k。内核忠实实验(Q,K,V为FP32以隔离P转换效应)显示,在中等聚类强度下均方误差改善3~10倍,配对测试确认两者结合后达到相同精度下限,由此推动hpc-ops内核从S=1更新至S=256。
原文摘要 · Abstract (English)
FP8 (E4M3) acceleration for attention computation offers significant throughput gains, but the 3-bit mantissa introduces precision challenges when the softmax probability matrix~$P$ is cast to FP8 before the $P \cdot V$ matrix multiplication. We analyze two implementation choices that affect output precision under the \emph{Attention Sink} phenomenon: (1)~the KV block iteration order, and (2) the static scaling factor applied to $P$ before casting. We show that forward KV iteration causes \emph{P-collapse} -- to leading order a fraction $Φ(Δ+ δ_k - 6.93 - \ln S)$ of non-sink $P$ values underflow to zero, where the small shift $δ_k \approx 1$ (for $k_{\text{sink}}{=}4$) is the expected within-sink-block score maximum -- and that reverse iteration removes it, with a zero-underflow guarantee when reverse is combined with $S{=}256$. We further give a constructive characterization of $S = 256 = 2^8$ as the static scale that simultaneously satisfies (i)~bit-exact IEEE 754 scaling, (ii) the lower envelope of a sawtooth function $dp(S)$ over the E4M3 number line ($dp = 2^{-4}$, the minimum worst-case quantization step), and (iii)~the maximum normal-range coverage \emph{among bit-exact ($2^k$) scales} (a non-bit-exact scale such as $448$ attains slightly higher coverage; sec.5}). Both optimizations are already deployed in FlashAttention-3/4 on engineering grounds; our contribution is a quantitative account of \emph{why} these choices are good and a closed-form threshold $Δ_c = 6.93 + \ln S - δ_k$ for predicting kernel-level precision loss. Kernel-faithful experiments ($Q, K, V$ in FP32 to isolate the P-cast effect) show $3$-$10\times$ MSE improvement at moderate sink strengths, and paired tests confirm both fixes saturate to the same precision floor when combined -- which motivated updating the hpc-ops kernel from $S{=}1$ to $S{=}256$.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。