针对黑格尔芯片特性优化注意力计算,速度提升超2.7倍。
FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling
- 设计全异步流水线与大块矩阵运算,适配黑格尔硬件特点。
- 在B200上达1613 TFLOPs/s,比cuDNN快1.3倍,比Triton快2.7倍。
- 用Python嵌入式DSL实现,编译快20-30倍,适合快速迭代开发。
注意力机制作为广泛使用的Transformer架构的核心层,是大语言模型和长序列应用的性能瓶颈。尽管FlashAttention-3通过异步执行和线程束特化优化了Hopper GPU上的注意力计算,但主要面向H100架构。随着AI行业快速转向基于黑格尔(Blackwell)架构的系统(如B200和GB200),其性能特征发生根本变化:张量核心吞吐量翻倍,而其他功能单元(共享内存带宽、指数单元)增长缓慢或保持不变。为此,我们提出多项技术应对新瓶颈:(1) 重新设计利用完全异步矩阵乘法累加(MMA)操作和更大分块大小的流水线;(2) 软件模拟指数和条件softmax归一化,减少非矩阵乘法操作开销;(3) 利用张量内存与2-CTA MMA模式,降低反向传播中的共享内存流量和原子操作。实验表明,FlashAttention-4在B200上以BF16精度达到最高1613 TFLOPs/s(71%利用率),相比cuDNN 9.13快1.3倍,比Triton快2.7倍。除算法创新外,我们使用嵌入于Python的CuTe-DSL完整实现,相比传统C++模板方法,编译速度提升20-30倍,同时保持全表达能力。
原文摘要 · Abstract (English)
Attention, as a core layer of the ubiquitous Transformer architecture, is the bottleneck for large language models and long-context applications. While FlashAttention-3 optimized attention for Hopper GPUs through asynchronous execution and warp specialization, it primarily targets the H100 architecture. The AI industry has rapidly transitioned to deploying Blackwell-based systems such as the B200 and GB200, which exhibit fundamentally different performance characteristics due to asymmetric hardware scaling: tensor core throughput doubles while other functional units (shared memory bandwidth, exponential units) scale more slowly or remain unchanged. We develop several techniques to address these shifting bottlenecks on Blackwell GPUs: (1) redesigned pipelines that exploit fully asynchronous MMA operations and larger tile sizes, (2) software-emulated exponential and conditional softmax rescaling that reduces non-matmul operations, and (3) leveraging tensor memory and the 2-CTA MMA mode to reduce shared memory traffic and atomic adds in the backward pass. We demonstrate that our method, FlashAttention-4, achieves up to 1.3$\times$ speedup over cuDNN 9.13 and 2.7$\times$ over Triton on B200 GPUs with BF16, reaching up to 1613 TFLOPs/s (71% utilization). Beyond algorithmic innovations, we implement FlashAttention-4 entirely in CuTe-DSL embedded in Python, achieving 20-30$\times$ faster compile times compared to traditional C++ template-based approaches while maintaining full expressivity.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。