CUDA 内核融合研究案例：在 NVIDIA Hopper 架构上使用 CUTLASS 库实现 FlashAttention-2

本文介绍了FlashAttention-2前向传递的优化实现，使用了自定义融合的CUDA内核，适应NVIDIA Hopper架构，并使用开源的CUTLASS库编写。通过解释在线softmax和连续的GEMM内核融合的挑战和技术，利用Hopper特定的Tensor Memory Accelerator（TMA）和Warpgroup Matrix-Multiply-Accumulate（WGMMA）指令，定义和转换CUTLASS布局和张量，重叠复制和GEMM操作，并选择最优瓦片大小，平衡寄存器压力和共享内存利用率。在单个H100 PCIe GPU上的对比性测试中，与针对上一代NVIDIA Ampere架构进行优化的FlashAttention-2版本相比，FLOPs/s高出20-50%。