FlashAttention-4震撼发布：性能再攀新高

在备受瞩目的半导体行业盛会 Hot Chips 2025 上，TogetherAI 首席科学家 Tri Dao 震撼发布了FlashAttention-4。

据官方介绍，在 Backwell 上，FlashAttention-4 的性能表现极为出色，其速度比英伟达 cuDNN 库中的注意力核实现快达 22%！

此次发布的 FlashAttention 新版本，Tri Dao 团队实现了两项关键算法改进。

首先，他们采用了一种全新的在线 softmax 算法，成功跳过了 90% 的输出 rescaling，大幅提升了效率。

其次，为了更高效地重叠 softmax 计算与张量核计算，他们引入了指数 (MUFU.EX2) 的软件模拟技术，从而显著提高了吞吐量。

此外，FlashAttention-4 采用了 CUTLASS CuTe Python DSL 进行开发，尽管其移植到 ROCm HIP 的难度比 CUDA C++ 高出 10 倍，但 CUDA C++ 到 ROCm HIP 的移植却相对容易。

值得一提的是，Tri Dao 还宣布，在执行 A@B+C 计算时，对于 Blackwell 上归约维度 K 较小的计算场景，他使用 CUTLASS CuTe-DSL 编写的核（kernel）比英伟达最新的 cuBLAS 13.0 库更快。而在标准矩阵算法 A@B 下，两者速度总体相当。

据悉，他的核通过两个累积缓冲区来重叠 epilogue，从而超越了 cuBLAS 的性能。

据 Semi Analysis 分析，像 Tri Dao 这样的开发者是 CUDA 护城河的核心优势之一，因为他们专注于使用英伟达 GPU，并大量开源其核代码，以造福其他英伟达开发者。目前，Tri Dao 等研究者尚未采用 ROCm AMD GPU 或 Trainium 芯片。

这对 AMD 等厂商来说无疑是个挑战。如果 AMD 希望 Tri Dao 及其团队在 ROCm 上实现算法突破，那么它应该为 TogetherAI GPU 云服务上的 AMD GPU 提供优惠支持。Semi Analysis 建议：「谷歌为 Noam Shazeer 投入了 27 亿美元，Zucc 为 OpenAI 工程师支付了 1 亿美元，AMD 拥有足够的资金，可以为 TogetherAI/Tri Dao 提供 5000 万美元来启动 ROCm 生态系统。」

FlashAttention 最初由 Tri Dao 等人在 2022 年提出，论文标题为《FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness》。

论文地址：https://arxiv.org/pdf/2205.14135

FlashAttention 针对传统注意力机制需生成 N×N 的注意力矩阵导致的二次时间和内存开销问题，提出了「IO-awareness」策略。通过「tiling+softmax rescaling」策略，将数据块临时存入高速缓存（SRAM），从而避免大量读写开销，将内存复杂度从 O(N²) 降低到 O(N)。

如图所示（左图），FlashAttention 利用 tiling 技术防止在相对较慢的 GPU HBM 上执行大规模的 N×N 注意力矩阵（虚线框）。在外层循环（红色箭头）中，FlashAttention 循环遍历 K 和 V 矩阵的块，并将其加载到快速片上 SRAM 中。在每个块中，FlashAttention 循环遍历 Q 矩阵的块（蓝色箭头），将其加载到 SRAM 中，并将注意力计算的输出写回 HBM。

右图显示，相比 GPT-2 上 PyTorch 注意力实现，FlashAttention 速度更快——无需将大型 N×N 注意力矩阵读写到 HBM，从而将注意力计算速度提升了 7.6 倍。

总体来看，初代 FlashAttention 的改进效果显著：在 BERT-large（序列长度 512）中相比 MLPerf 基线提升训练速度约 15%；GPT-2（序列长度 1K）提升约 3 倍；在 Long-Range Arena（序列长度 1K–4K）提升约 2.4 倍。

一年后，《FlashAttention-2》问世。这一次，作者仅 Tri Dao 一人。值得一提的是，他还在这一年的晚些时候与 Albert Gu 共同提出了 Mamba。

论文地址：https://arxiv.org/pdf/2307.08691

FlashAttention-2 的改进焦点在于：尽管性能已显著提升，但在 GPU 上仍存在低吞吐率问题，仅能达到理论峰值的很低比例（约 25–40%）。

为此，Tri Dao 提出了一系列解决策略：

• 工作划分优化：重新设计分块策略与线程分配，提升并行效率，增加硬件利用率；
• 减少非矩阵运算，加快整体执行；
• 支持更大 head size（至 256）及多查询注意力（MQA）和分组查询注意力（GQA），以满足更多模型架构需求。