AI 的最新优化 FlashAttention-4 在 NVIDIA 旗舰 Blackwell B200 GPU 上提供高达 1,605 TFLOPs/s 的速度，标志着 AI 工作负载的硬件效率和计算速度的显着飞跃。
Together AI 发布了 FlashAttention-4，在 NVIDIA Blackwell B200 GPU 上实现了高达 1,605 TFLOPs/s，代表着 71% 的硬件利用率，比 Triton 实现速度提高了 2.7 倍。该版本解决了现代人工智能硬件中的一个基本挑战：张量核心吞吐量的扩展速度远远快于其他关键资源。
就背景而言，截至 2026 年 3 月 4 日，NVIDIA 的市值为 4.49 万亿美元，股价为 179.86 美元。该公司昨天发布了自己的 Blackwell GPU Flash Attention 优化指南，表明注意力优化在生产 AI 工作负载中的重要性与日俱增。
不对称缩放问题
这就是为什么这很有趣。从 Hopper H100 到 Blackwell B200，BF16 张量核心吞吐量从 1 PFLOP 跃升至 2.25 PFLOP。但是用于指数运算和共享内存带宽的特殊功能单元呢？不变。这造成了没人预料到的瓶颈。
Together AI 团队发现，前向传递在 B200 上根本不受计算限制，它受到 Softmax 中指数计算的瓶颈。向后传球？共享内存流量占主导地位。传统的注意力优化侧重于错误的约束。
FA4 如何解决
前向传递使用乒乓调度来处理每个 CTA 的两个查询图块，并使用专用的 warpgroup 处理 softmax，而其他的则发出矩阵运算。巧妙之处在于：使用 FMA 单元和硬件 MUFU.EX2 对指数函数进行软件模拟，有效地将指数吞吐量翻倍。
有条件的在线 softmax 重新缩放完全跳过小修正。如果最大跳跃保持在阈值以下，内核会避免不必要的向量操作。最终标准化仍然会产生正确的结果，但关键路径大大缩短。
向后传递利用 Blackwell 的新 2-CTA MMA 模式，跨 CTA 对划分输出累加器。每个 CTA 暂存操作数 B 的一半，同时仅保留其累加器切片，将共享内存流量大致减半。 dQ 梯度的全局原子缩减也下降了一半。
性能数据
与 cuDNN 9.13 相比，FlashAttention-4 在前向传递方面提供了 1.1-1.3 倍的改进，并在大序列长度的后向传递方面提供了一致的增益。 Triton 的比较显示了最明显的差异 - 前进性能提高了 2.7 倍。
确定性模式将全局缩减序列化以进行可重复的训练，仍然可以实现 85-90% 的非确定性吞吐量。这对于需要在训练运行中精确再现的团队来说非常重要。
更广阔的前景
自 2022 年 5 月首次亮相以来，FlashAttention 发展迅速。版本 1 在 A100 上实现了 25-40% 的利用率。 FA2 在 2023 年 7 月将这一比例提高到 50-73%。FA3 特别针对 Hopper GPU，FP16 的利用率达到 75%，FP8 的利用率接近 1.2 PFLOPS。
FA4 代表了一种哲学转变——算法和内核协同设计，解决了不对称硬件演进的问题。通过与 NVIDIA 团队的合作，这些技术已部分纳入 cuDNN 9.13 和 9.14。
该实现使用 CuTe-DSL（CUTLASS 的 Python 内核 DSL），与 C++ 模板相比，编译时间缩短了 20-30 倍。对于在 Blackwell 硬件上运行大规模训练的团队来说，每天数百万次注意力操作的效率会得到复合提升。