FlashAttention-4: Algorithm and Kernel Pipelining
9 days ago
- #Attention Mechanism
- #Deep Learning
- #GPU Optimization
- 现代GPU如Blackwell上的非对称硬件扩展导致张量核心吞吐量的增长速度快于共享内存带宽和特殊功能单元（SFU）等资源。
- FlashAttention-4是为Blackwell架构协同设计的算法与内核，在B200上使用BF16精度实现高达1605 TFLOPs/s，性能超越cuDNN和Triton。
- 关键创新包括用于重叠的新流水线技术、通过多项式近似进行指数运算的软件模拟以缓解SFU瓶颈，以及使用TMEM减少共享内存流量。
- Blackwell的硬件特性如Tensor Memory（TMEM）、完全异步的第五代张量核心和2-CTA MMA模式支持更大的分块、降低流量并提升性能。
- 反向传播受限于共享内存带宽；优化措施包括采用2-CTA MMA将流量减半并减少原子加法操作，以及确定性模式以实现可复现的训练。
- 调度改进通过网格线性化和最长处理时间优先（LPT）调度，解决了因果掩码和可变序列长度导致的负载不均衡问题。
- FlashAttention-4基于CuTe-DSL实现，显著缩短了编译时间，并推动了cuDNN 9.13及后续版本的优化。