发表机构
Meta; Colfax Research（Meta; Colfax Research）
机构由 AI 辅助整理，请以论文原文为准。AI 中文总结
本文在Blackwell上通过固定移位Softmax和持久调度实现高效FlashAttention，在B200上平均吞吐量提升7.3%，前向和反向分别提升13.0%和10.4%。
AI 中文摘要
在Blackwell架构上，归一化和操作数移动可能限制注意力内核，其矩阵乘法已经深度流水线化。我们在Triton TLX中实现了一个FA4风格的流水线，采用固定移位密集Softmax和用于反向传播的保存的逆分母。固定移位消除了循环累加器校正；保存的倒数将行归一化从二次反向循环移到线性预处理。受保护的因果路径使用带有选择性恢复的第一块锚点，并跳过完全掩码的对角线收缩。打包的BF16指数近似和流水线化的dQ发布减少了逐元素工作和片上移动。在B200上，在1000瓦功耗下，使用准备好的启动和原始FA4计时器，TLX在FA4的BF16可变批量网格中，通过头维度128、最多32768个token的24个案例中，将几何平均吞吐量提高了7.3%。密集前向和反向分别提高了13.0%和10.4%。因果反向提高了8.0%，而因果前向慢了1.5%。与TritonBench的比较中，前向的几何平均增益为13.4%，反向为24.4%。
英文摘要
On Blackwell, normalization and operand movement can limit attention kernels whose matrix multiplications are already deeply pipelined. We implement an FA4-style pipeline in Triton TLX with fixed-shift dense softmax and a saved inverse denominator for backward. The fixed shift removes recurrent accumulator corrections; the saved reciprocal moves row normalization from the quadratic backward loop into linear preprocessing. The guarded causal path uses a first-block anchor with selective recovery, and fully masked diagonal contractions are skipped. A packed BF16 exponential approximation and pipelined dQ publication reduce per-element work and on-chip movement. On B200 at 1,000 W, using prepared launches and the original FA4 timer, TLX improves geometric-mean throughput by 7.3 percent across the 24 cases in FA4's BF16 variable-batch grid through 32,768 tokens at head dimension 128. Dense forward and backward improve by 13.0 and 10.4 percent, respectively. Causal backward improves by 8.0 percent, while causal forward is 1.5 percent slower. The TritonBench comparison has geometric- mean gains of 13.4 percent in forward and 24.4 percent in backward
