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*V matrix multiplication. We analyze two implementation choices that affect output precision under the 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 "P-collapse" -- to leading order, a fraction Phi(Delta + delta_k - 6.93 - ln S) of non-sink P values underflow to zero, where the small shift delta_k ~ 1 (for k_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 among bit-exact (2^k) scales (a non-bit-exact scale such as 448 attains slightly higher coverage). Both optimizations are already deployed in FlashAttention-3/4 on engineering grounds; our contribution is a quantitative account of why these choices are good and a closed-form threshold Delta_c = 6.93 + ln S - delta_k for predicting kernel-level precision loss. Kernel-faithful experiments (Q, K, V in FP32 to isolate the P-cast effect) show 3-10x MSE improvement at moderate sink strengths, and paired tests confirm both fixes saturate to the same precision floor when combined.
翻译:FP8(E4M3)加速注意力计算能显著提升吞吐量,但3位尾数在P*V矩阵乘法前将softmax概率矩阵P转换为FP8格式时会引入精度挑战。针对注意力汇聚现象影响输出精度的两种实现选择——(1) KV块迭代顺序,(2) 转换前应用于P的静态缩放因子——我们进行了分析。研究表明,前向KV迭代会导致“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可实现略高覆盖度)。两项优化已基于工程考量部署于FlashAttention-3/4中;本研究的贡献在于定量阐明这些选择的优越性,并给出预测核级精度损失的闭式阈值Δ_c = 6.93 + ln S - δ_k。基于核函数的实验(在FP32下处理Q、K、V以隔离P转换效应)表明,在中等汇聚强度下均方误差可提升3-10倍,配对实验证实两项修正结合使用时精度会收敛至相同下限。