给高速公路扩了车道,收费站却没变快,整体通行时间就不会按比例下降。Blackwell GPU上的注意力也遇到类似问题:FP4——一种只用4位表示数值的低精度格式——能大幅加速矩阵乘法,但注意力还要经过Softmax,把相关性分数转换成总和为1的权重。乘法耗时缩短后,这段转换和数据交接反而占据了关键路径。
Robert Hu提出面向Blackwell的FlashAttention-4实现Direct-P。FlashAttention是一类通过重排计算、减少显存往返的方法;Direct-P进一步把归一化后的分数直接编码成FP4概率,并让后续乘法使用的舍入值参与归一化,省去一部分转换与等待。作者报告,在有利的计算形状下,NVIDIA GB200上的前向吞吐最高达到BF16版本的2.13倍。但论文也指出,更深的限制来自片上Tensor Memory:分数与输出已占满可用空间,前一步不释放位置,下一次乘法就无法提前开始。换句话说,FP4算得快只是起点;要把速度真正兑现,还得重新安排Softmax和片上数据的接力。