
DeepSeek
FlashAttention3的两位共同第一作者就职于Colfax Research。年初的时候,他们就有了一个初步版本的fp8 FlashAttention - 2,这个版本利用了Hopper架构的wgmma和tma指令以及pipeline进行改造,能够实现1 PFLOP/s的性能(Delivering 1 PFLOP/s of Performance with FP8 FlashAttention - 2)。此次有诸多亮点,由于wgmma的fp8指令仅支持k - major,所以实现了QKV中的V在核内转置(得到了cudnn团队的帮助),并且不需要写回到HBM(对于
DeepSeekV2的MLA来说,因为KV共享,所以非常需要这一实现)。另外一个亮点是对fp8量化误差的控制。当前github上的代码只有fp16的实现,fp8的实现还需要等几天或者几周的时间。非常期待能够从源代码
中学习到V的在核内转置。