Modal工程师逆向解析Flash Attention 4的工作原理

Modal工程师阅读了已发布的Flash Attention 4源代码，逆向解析了这个针对Blackwell优化的注意力内核的工作原理，因为目前尚无官方技术报告。文章追踪了一个tile在内核的warp特化异步流水线中的流转过程，并详述了两个数学技巧：一个可替代基于SFU的指数运算的三次多项式软件指数，以及一个更聪明的在线softmax重缩放规则——Tri Dao称其将修正次数减少了10倍。据报道，FA4比Nvidia闭源的cudnn注意力内核快约20%，而在撰写本文时反向版本尚不可用。