【第578期】FlashAttention-4：针对 Blackwell 架构的算法与内核协优设计

Seventy3：借助NotebookLM的能力进行论文解读，专注人工智能、大模型、机器人算法、crypto方向，让大家跟着AI一起进步。如果你想要解读自己的论文，获得更多曝光度。请联系小助手微信：seventy3_podcast 加群。合作邮箱：zhiwudazhanjiangshi#gmail.com今天的主题是：FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware ScalingSummary作为 Transformer 架构的核心层，注意力机制（Attention） 是大语言模型和长文本应用中的性能瓶颈。虽然 FlashAttention-3 通过异步执行和 Warp 特化（warp specialization）针对 Hopper GPU 优化了注意力计算，但其主要目标仍是 H100 架构。当前，AI 行业正迅速转向部署基于 Blackwell 的系统（如 B200 和 GB200）。由于硬件扩展的不对称性，这些系统表现出截然不同的性能特征：其 Tensor Core 吞吐量翻倍，而其他功能单元（如共享内存带宽、指数运算单元）的提升较慢或保持不变。为了应对 Blackwell GPU 上这些变化的瓶颈，我们开发了 FlashAttention-4，并引入以下核心技术：重新设计流水线：充分利用完全异步的矩阵乘累加（MMA）操作和更大的分块（tile）尺寸。软件模拟与重缩放：通过软件模拟指数运算和条件式 Softmax 重缩放，减少非矩阵乘法（non-matmul）操作。优化内存流量：利用 Tensor Memory 和 2-CTA MMA 模式，减少反向传播过程中的共享内存流量和原子加法（atomic adds）。性能表现： 在 BF16 精度的 B200 GPU 上，FlashAttention-4 相比 cuDNN 9.13 实现了高达 1.3 倍的加速，相比 Triton 实现了 2.7 倍的加速。其性能达到 1613 TFLOPs/s，硬件利用率高达 71%。工程创新： 除了算法创新，FlashAttention-4 完全使用嵌入 Python 的 CuTe-DSL 实现。与传统的基于 C++ 模板的方法相比，这种方式在保持完整表达能力的同时，将编译速度提升了 20-30 倍。原文链接：https://arxiv.org/abs/2603.05451前往小宇宙评论区与主播互动