A Case Study in CUDA Kernel Fusion: Implementing FlashAttention-2 on NVIDIA Hopper Architecture using the CUTLASS Library

我们提供了FlashAttention-2前向传递的优化实现，它是一种流行的内存感知的缩放点积注意力算法，作为一个定制的融合CUDA内核，针对NVIDIA Hopper架构进行编写，使用开源的CUTLASS库。在此过程中，我们解释了将在线softmax与背靠背的GEMM内核融合的挑战和技术，利用了Hopper特定的张量内存加速器（TMA）和Warpgroup矩阵乘累加（WGMMA）指令，定义和转换CUTLASS布局和张量，重叠复制和GEMM操作，并选择最佳的瓷砖大小，以平衡寄存器压力和共享内存利用率。在一些常见的超参数选择上，我们在单个H100 PCIe GPU上进行了头对头的基准测试，观察到相对于为上一代NVIDIA Ampere架构优化的FlashAttention-2版本，FLOPs/s高出20-50%。

优化FlashAttention-2算法的前向传递，以适应NVIDIA Hopper架构。

使用CUTLASS库编写自定义融合CUDA内核，利用TMA和WGMMA指令，定义和转换CUTLASS布局和张量，重叠复制和GEMM操作，并选择最佳的瓷砖大小。

在单个H100 PCIe GPU上进行了头对头基准测试，观察到比优化过Ampere架构的版本高20-50%的FLOPs/s。

相关研究包括FlashAttention-2算法本身以及其他针对不同硬件架构的优化研究，如针对TPU的优化等。