FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling——针对非对称硬件

本文是FlashAttention注意力系列的第四篇论文，在前面我们已经系统学习了前三篇文章，如下所示：

《FlashAttention: Fast and Memory-Efficient Exact Attentionwith IO-Awareness——快速且内存高效的精确注意力机制，具有IO感知能力》

《FlashAttention-2:Faster Attention with Better Parallelism and Work Partitioning——更好的并行性和工作分区带来更快的注意力机制》

《FlashAttention-3:Fast and Accurate Attention with Asynchrony and Low-precision——利用异步和低精度实现快速准确的注意力机制》

感兴趣的话可以自行移步阅读学习回顾FlashAttention系列的前置内容。

《FLASHATTENTION-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling》主要研究了在NVIDIA Blackwell（B200/GB200）GPU架构上，针对非对称硬件扩展（即张量核心计算速度远快于其他单元）导致的性能瓶颈，如何通过算法与内核流水线的协同设计，对注意力机制（Attention）进行极致优化。

以下是该文章的核心研究内容全面总结：

1. 研究背景与核心问题

• 硬件变化： 随着从Hopper（H100）架构演进到Blackwell（B200）架构，硬件性能提升呈现非对称性。张量核心（Tensor Core）的FP16/BF16吞吐量翻倍，但共享内存带宽、指数单元（MUFU）等其它部件的提升几乎停滞。

• 新瓶颈： 这使得注意力机制的瓶颈从“矩阵乘法”转移到了“共享内存流量”和“非矩阵乘法运算（如Softmax中的指数计算）”。

• 目标： 针对Blackwell的新特性，设计新的算法（FlashAttention-4），以最大化硬件利用率。

2. 主要技术创新（解决方案）

文章提出了四大核心优化技术来应对上述瓶颈：

(1) 重新设计的并行流水线

• 利用完全异步性： Blackwell的MMA（矩阵乘累加）指令是完全异步的，且支持更大的图块尺寸（128×128）。

• 乒乓调度： 设计了新的前向和反向传播流水线，使得一个图块的张量核心计算与另一个图块的Softmax计算可以充分重叠，隐藏了Softmax的延迟。

(2) 指数单元瓶颈的缓解

• 软件模拟指数： 由于硬件指数单元（MUFU）吞吐量低，作者使用FMA单元通过多项式逼近（基于Cody-Waite算法）软件模拟了2x2x的计算。

• 部分模拟+精度控制： 仅对部分条目使用软件模拟以平衡性能和寄存器压力。实验表明，3阶多项式在转换为BF16后，精度与硬件单元相当，但吞吐量更高。

• 有条件的Softmax重缩放： 修改了在线Softmax算法，仅在最大值变化超过阈值时才进行重缩放，跳过了大部分不必要的向量乘法操作。

(3) 共享内存流量与原子操作的减少

• 利用TMEM： 利用Blackwell新增的张量内存（TMEM）存储中间结果，减少了对共享内存的读写压力。

• 2-CTA MMA模式： 利用2-CTA（协作线程数组）协同执行MMA的特性，使每个CTA只加载一半的操作数，大幅减少共享内存流量。

• 优化dQ计算： 在反向传播中，通过重构dQ的计算步骤，将全局内存的原子加（AtomicAdd）操作次数减半，既提升了速度，也配合了确定性执行的需求。

(4) 改进的调度策略（LPT调度）

• 最长处理时间优先： 针对因果掩码和可变序列长度导致的负载不均衡，采用了最长处理时间优先（LPT）调度算法。

• 缓存感知： 通过优化CTA的遍历顺序（如按批次、按头部、逆序处理块），提高了L2缓存的命中率，减少了停顿。

3. 实现框架的创新

• CuTe-DSL + Python： FlashAttention-4完全使用嵌入在Python中的CuTe-DSL实现，而不是传统的C++模板。

• 编译速度： 这一选择带来了巨大的开发效率提升：相比FlashAttention-3，编译速度快了20到30倍（前向2.5秒 vs 55秒），使得研究人员能快速迭代，而无需深谙C++元编程。

4. 性能表现

• 实测数据： 在B200 GPU上，FlashAttention-4相比cuDNN 9.13实现了最高1.3倍加速，相比Triton实现了最高2.7倍加速。

• 峰值性能： 达到了约1613 TFLOPS，占B200理论峰值的71%。

• 确定性模式： 提供的确定性反向传播模式性能损失极小，可达非确定性版本75%的速度，适用于需要严格可重现性的场景（如强化学习）。

FlashAttention-4是一项针对Blackwell GPU架构“量身定制”的注意力机制优化方案。 它通过软件模拟硬件短板（指数单元）、利用新硬件特性（TMEM、2-CTA）来消除新瓶颈（共享内存流量），并通过编译框架的重构（CuTe-DSL+Python）大幅降低了开发门槛和编译时间，为长上下文Transformer模型在最新硬件上的部署提供了显著的性能提升。这里是自己的论文阅读记录，感兴趣的话可以参考一下，如果需要阅读原文的话可以看这里，如下所示：

项目地址在这里，如下所示：

摘要

作为无处不在的Transformer架构的核心层，Attention（注意力机制）已成为大型语言模型和长上下文应用的瓶颈。虽然FLASHATTENTION-3通过异步执行和线程束专门化针对Hopper GPU优化了注意力机制，但它主要针对H100架构。AI行业已迅速转向部署基于Blackwell的系统，如B200和GB200，这些系统由于非对称硬件扩展而表现出根本不同的性能特征：张量核心吞吐量翻倍，而其他功能单元（共享内存带宽、指数单元）扩展较慢或保持不变。我们开发了几种技术来解决Blackwell GPU上这些变化的瓶颈：（1）重新设计的流水线，利用完全异步的MMA操作和更大的图块尺寸，（2）软件模拟的指数运算和有条件的Softmax重新缩放，减少了非矩阵乘法操作，（3）利用张量内存和2-CTA MMA模式减少反向传递中的共享内存流量和原子加操作。我们证明了我们的方法FLASHATTENTION-4在B200 GPU上使用BF16，相比cuDNN 9.13实现了高达1.3倍的加速，相比Triton实现了2.7倍的加速，达到了高达1613 TFLOPS/s（71%的利用率）。除了算法创新，我们还完全用嵌入在Python中的CuTe-DSL实现了FLASHATTENTION-4，与传统的C++模板方法相比，编译速度快了20-30倍，同时保持了完全的表达能力。

1 引言

Transformer架构[27]仍然是几乎所有AI应用的主要支柱，从大型语言模型[2]到视觉[8]和多模态系统。对于Transformer而言，注意力机制构成了主要的计算瓶颈，因为查询和键之间计算的自注意力分数在序列长度上呈现二次方扩展。将注意力扩展到更长的上下文解锁了新功能，例如对多个文档进行推理[10, 24]、对整个代码库进行建模[22]以及处理高分辨率视频[3, 11]。与此同时，加速器硬件持续快速发展[19]，每一代都提供显著更高的峰值计算吞吐量。然而，这种发展是非对称的：虽然矩阵乘法单元扩展迅猛，但其他功能单元（如内存带宽和专用计算单元）扩展较慢，造成了日益不平衡的硬件流水线，这需要精心的算法协同设计。

这引发了通过深度融合GPU硬件特性知识的算法创新来加速注意力的持续兴趣。Dao等人[6]引入了FLASHATTENTION，它通过以下方式消除了对慢速全局内存的中间读写：

除了算法创新，我们完全用嵌入在Python中的CuTe-DSL实现了FLASHATTENTION-4，与传统的C++模板方法相比，实现了20-30倍的编译加速，同时保持了完全的表达能力。这个框架显著提高了开发者的生产力。

通过实证验证我们的方法，我们在B200 GPU上对FLASHATTENTION-4进行了基准测试，结果显示：（1）BF16相比cuDNN实现了高达1.3倍的加速，相比Triton实现高达2.7倍的加速；（2）我们在变化的瓶颈资源上实现了接近峰值的利用率，达到了约1600 TFLOPS（理论最大值的71%）；（3）对于长序列长度，FLASHATTENTION-4优于其他注意力实现。

我们以宽松许可证开源FLASHATTENTION-4，并正在努力将其集成到流行的库中，以使尽可能多的研究人员和开发者受益。

2 背景

2.1 多头注意力

2.2 GPU 硬件特性与执行模型

我们描述与 FLASHATTENTION-4 相关的 GPU 执行模型方面，重点关注 NVIDIA Blackwell 架构（B200 和 GB200）。我们强调了与之前的 Hopper 架构的关键区别，这些区别促使了 FLASHATTENTION-4 中的优化。

• 内存层次结构： GPU 的内存组织为数据位置的层次结构，容量与带宽成反比。全局内存（GMEM），也称为 HBM，是片外 DRAM，所有流多处理器（SM）都可访问。来自 GMEM 的数据被透明地缓存在片上 L2 缓存中。接下来，每个 SM 包含一个小的、程序员管理的、高度分库的片上缓存，称为共享内存（SMEM）。最后是每个 SM 内部的寄存器文件。 Blackwell 引入了一个新的内存层级，称为张量内存（TMEM），每个 SM 有 256 KB 的片上内存，专门用于存储张量核心操作的中间结果。与共享内存不同，TMEM 是线程束同步的，并与张量核心紧密耦合，使得矩阵乘累加（MMA）单元可以直接将输出写入 TMEM，而无需消耗寄存器。这缓解了困扰 Hopper 内核的极端寄存器压力，并支持更大的图块尺寸。TMEM 以 32 列（16 KB）为粒度进行分配，需要程序员显式管理分配、释放和数据移动。

• 线程层次结构： GPU 的编程模型围绕称为线程的执行单元的逻辑分组进行组织。从最细到最粗的级别，线程层次结构包括线程、线程束（32 个线程）、线程束组（4 个连续的线程束）、线程块（即协作线程数组 CTA）、线程块簇和网格。同一 CTA 中的线程被共同调度到同一个 SM 上，同一簇中的 CTA 被共同调度到同一个 GPC 上。SMEM 可由 CTA 内的所有线程直接寻址，而每个线程最多拥有 256 个私有寄存器（RMEM）。

• 张量核心与增强的异步性： Blackwell 具有第五代张量核心，其操作的图块尺寸比以前的架构大得多。每个 MMA 张量核心指令处理 128×N128×N 的图块（通常 N=128 或 256），而 Hopper 上是 64×N。至关重要的是，Blackwell MMA 将其输出异步直接写入 TMEM，而 Hopper MMA 则写入寄存器。这种完全的异步性使得计算与其他操作之间能够更好地重叠，因为 MMA 单元不再阻塞于寄存器写回。 硬件对异步性的支持允许实现线程束专门化的内核，其中 CTA 的线程束被划分为生产者或消费者角色，仅负责发出数据移动或计算指令[1]。

• 2-CTA 张量核心： Blackwell 支持一种 2-CTA 张量核心 MMA 模式，其中同一线程块簇内的一对 CTA 协作执行单个 MMA，允许该操作从两个 CTA 读写张量内存。对中的一个线程发起 MMA，但伙伴 CTA 必须已启动并在操作进行期间保持活动状态。与将维度 M 限制为 128 的单 CTA MMA 相比，配对模式通过在对之间沿维度 M 分割 A 图块和累加器，并沿维度 N 在两个 CTA 之间分割 B 图块来支持 M=128 或 256，这样每个 CTA 仅在自己的共享内存中暂存一半的 B，而硬件在乘法期间消耗组合后的 B 图块。这减少了冗余的共享内存容量和带宽，但由于这些操作会触及跨 CTA 对的张量内存，内核必须以固定对的形式启动 CTA，并在此内核中始终使用一致的 2-CTA 模式进行张量内存和张量核心操作。

• 变化的瓶颈： Blackwell 反映的一个关键趋势是，张量核心吞吐量的增长速度快于其他功能单元。与 Hopper 相比，Blackwell 的 FP16/BF16 张量核心吞吐量翻倍（每个 GPU 2.25 PFLOPS [19] 对比 1 PFLOPS [17]），但共享内存带宽和指数单元吞吐量保持不变或增长较慢。这种不平衡使性能瓶颈从矩阵乘法转移到共享内存流量和像 softmax 这样的非矩阵乘法操作。正如我们在 3.1 和 3.2 节中的屋顶线分析所示，这需要精心设计内核，以最大化 MMA 操作与这些瓶颈资源之间的重叠。

B200（和 GB200）上几个硬件组件的吞吐量如下：

3 算法

3.1 注意力前向传递

我们首先进行屋顶线分析，以展示注意力前向传递的瓶颈，这促使我们设计新的流水线，并修改 FLASHATTENTION 算法以提高指数单元的吞吐量并避免大部分 softmax 重新缩放步骤。

3.1.1 速率分析

我们通过首先分析屋顶线来提供内核设计和优化的直觉，该分析基于矩阵乘法单元（张量核心）、共享内存（smem）和指数单元的吞吐量。我们注意到，这是一个简化的分析，未考虑 GPU 中的所有资源（例如，浮点运算、寄存器带宽、L2 带宽）。尽管如此，它可以识别瓶颈。

令沿 Q 和 K 序列长度维度的图块形状为 M×N，头维度为 d。我们分析计算和内存流量需求，以确定性能瓶颈。

表 1：前向传递的屋顶线分析（周期）。

3.1.2 用于重叠矩阵乘法和 Softmax 的新流水线

由于 Blackwell 架构再次将张量核心 FLOPS 翻倍，因此像在 Hopper 上一样小心地重叠 softmax 和张量核心操作变得更加关键。我们遵循类似于 FA-3 的乒乓调度，其中每个线程块计算两个输出图块。当一个图块的张量核心操作被执行时，另一个图块计算 softmax。Hopper 张量核心将累加器保存在寄存器中，每行有四个线程交织排列，而 Blackwell 张量核心将其累加器保存在张量内存中。此外，Blackwell 上的单个累加器图块大小为 128×128 元素，而 Hopper 的图块大小是 64×128。

然后，在这些图块之间分配工作的自然方式是拥有两个 128 线程的线程束组，每个线程处理一整行。这消除了用于归并行最大值的线程束间混洗的需要，也消除了每个线程拥有多个统计寄存器的需要。就像 FA-3 一样，我们显式同步两个 softmax 线程束组，使它们的关键部分（即指数计算部分）不重叠。每个 softmax 线程束组首先将整行加载到寄存器中，然后计算最大值，然后计算 softmax（即减去最大值、重新缩放、取指数、转换为输入精度），最后计算行和。

与 FA-3 的另一个区别是，由于我们通过张量内存而不是寄存器文件传输 P，我们可以将输出的重新缩放解耦到一个单独的“校正”线程束组中，从而将其移出关键路径。

有几种张量内存分区方式可以实现这种流水线重叠。所有分区都必须分配两个输出图块的空间，在头维度为 128 的情况下，留下一半的张量内存来存储 S 和 P。该内存可以存储两个 S 副本或四个 P 副本（假设是 FP16 或 BF16 张量核心的输入）。这给我们留下了两种剩余张量内存的分区选项：一个 S 图块和两个 P 图块，或者两个与 P 重叠的 S 图块。我们选择后者，因为它允许我们通过立即计算两个 S 图块来启动我们的软件流水线。它还留下了一些张量内存来将重新缩放统计信息传递给校正线程束组。

更大的 Blackwell 图块尺寸和所选的线程分配带来的一个问题是，除非我们从张量内存重新加载，否则我们必须将一整行 128 个元素保存在寄存器中。鉴于我们使用两个 softmax 线程束组、一个校正线程束组和一个驱动张量核心和 TMA 单元的线程束组，为 softmax 分配足够的寄存器并防止寄存器溢出至关重要。对于 BF16 输入数据类型，我们需要 128 个寄存器来保存输入，可能还需要 64 个寄存器来保存输出（加上其他杂项和临时寄存器）。为了减少寄存器压力，我们分阶段存储 P：前四分之三在存储后立即触发相应的 MMA 操作，最后四分之一单独存储。

3.1.3 指数函数的模拟

• 指数吞吐量瓶颈： 在现代 GPU 上，指数函数由多功能单元（MUFU）计算，其吞吐量远低于用于矩阵乘法的张量核心。在 B200 和 GB200 GPU 上，MUFU 提供每个时钟每 SM 16 次操作，而矩阵乘法为每个时钟每 SM 8192 次操作。由于 softmax 计算需要大量指数运算，这种差异使得指数函数成为注意力内核中的关键瓶颈。

• 通过多项式逼近进行软件模拟： 为了提高指数吞吐量，我们使用浮点 FMA 单元实现 2x2x 的软件模拟，这些单元可以与 MUFU 并行操作。我们使用经典的归约技术（Cody-Waite），然后进行多项式逼近[16]。关键思想是分解指数计算：

• 通过将指数计算分布到 MUFU 和 FMA 单元上，这种方法有效地提高了指数吞吐量，缓解了注意力计算中的一个关键瓶颈。

• 部分模拟： 尽管多项式模拟提高了指数吞吐量，但它也是有代价的：额外的寄存器（用于保存中间值和系数）、更高的寄存器带宽消耗以及与 MUFU 指令相比更长的延迟。对所有指数计算使用模拟会增加寄存器压力，并可能导致寄存器溢出，从而抵消吞吐量的提升。相反，我们仅对每个 softmax 行中的一部分条目（10-25%）应用模拟，其余条目通过硬件 MUFU.EX2 计算。具体比例是根据给定图块配置的 MMA 和指数吞吐量之比凭经验调整的。

3.1.4 跳过在线 Softmax 重新缩放

• 这种修改显著减少了重新缩放操作的次数，同时保持了数值精度，因为最终的归一化步骤校正了跳过中间重新缩放引入的任何微小偏差。

在实践中，为了避免线程束发散，当线程束中的任何线程需要重新缩放时，我们就进行重新缩放。

3.2 注意力反向传递

3.2.1 速率分析

与前向传递类似，我们首先通过基于矩阵乘法单元（张量核心）、共享内存（smem）和指数单元吞吐量的屋顶线分析，为内核设计和优化提供直觉。

3.2.2 用于重叠矩阵乘法和 Softmax 的新流水线

Flash Attention 中的反向传递执行五个 MMA 操作，分别对应于重新计算 S，以及由 QK（产生 dQ 和 dK）和 PVPV（产生 dP 和 dV）引起的两个梯度计算。在 FA-3 中，累加器存储在寄存器中，而寄存器是一种有限的资源。这施加了显著的有序性约束——实际上——将计算图串行化，即计算 S、dP、dV、dQ、dK，只有 TMA 加载可以显著超前运行。除此之外，算法是类似的：它沿着 KV 序列长度维度迭代，并计算相对于前向传递转置的值，因为这是 dV 和 dK 梯度计算所需的一种布局，以便从张量内存读取它们的一个操作数。dQ 通过原子操作累积。

在 FA-4 中，TMEM 使得与 FA-3 相比能够采用额外的调度，从而在 MMA 和非 MMA 操作之间提供显著的重叠。具体来说，与前向传递一样，我们试图隐藏 softmax 计算的延迟。在 FA-3 中，softmax 计算与 dP 的 MMA 重叠。从前一节我们知道，在 Blackwell 上，我们需要至少两个 MMA 操作同时运行。

我们通过使用前一次迭代的 dQ 和 dK MMA 来实现这一点。这需要在加载、MMA、计算和归约操作之间仔细管理共享内存和张量内存资源。特别要注意的是，我们没有足够的张量内存来容纳五个累加器图块。最多可以容纳四个 128×128 元素的图块，并且 dV 和 dK 会累积，因此它们不能共享空间。在我们的实现中，我们让 S 和 P 共享一个 tmem 块（在偏移量 0 处），并让 dP、dS 和 dQ 共享另一个块。我们在图 2 中展示了 FA-4 反向传递的计算图。

3.2.3 2-CTA 反向传递：减少共享内存流量和全局原子加操作

即使改进了流水线，并且十个 GEMM 操作数中有两个驻留在张量内存中，共享内存带宽仍然主导着反向传递。在五个 GEMM 中，剩余的八个 BF16 操作数从共享内存加载到张量核心，这种共享内存流量产生的周期数比张量核心计算多约 30%。为了进一步缓解这一瓶颈，我们使用了 Blackwell 引入的 2-CTA MMA 模式，其中输出累加器在 M 维度上分割。对于 M=256 和 N=K=128 的 MMA 图块形状，两个 CTA 就像一个更大的图块：每个 CTA 加载并暂存一半的操作数 B，并只保留自己的累加器切片。

3.2.4 确定性反向传递

我们的反向内核由于全局内存中的 CTA 间归约（通常影响 dQ，在 GQA 情况下也影响 dK/dV）而为梯度计算引入了非确定性。为了确保可重现性并促进训练期间的可靠调试，我们还提供了一种确定性执行模式。我们也采用的标准解决方案是使用信号量锁序列化全局归约。具体来说，每个写入公共 dQ 图块的 CTA 必须根据预定义的顺序获取锁，执行其归约，然后通过递增信号量计数器释放锁。

这种基于锁的方法会影响性能，主要原因有两个：（1）发出内存屏障以确保信号量写入的设备级可见性（需要正确的获取-释放语义），以及（2）引入停顿，因为每个 CTA 需要等待在其之前对公共 dQ 图块进行归约的 CTA 完成。在负载不均衡的情况下，一个朴素的 CTA 顺序选择会严重降低性能。通常，我们在头部和批次维度上进行 CTA 交织，以减少停顿，直到达到 L2 缓存容量，参见第 3.3 节。对于因果掩码，我们额外以降序启动 KV 块，以升序从对角线开始遍历查询块，并按降序的查询块索引排列 dQ 归约。这种“最短处理时间优先”（SPT）调度确保了没有 CTA 在其首次 dQ 写入时停顿。

3.3 调度

在许多情况下，例如使用因果掩码或可变序列长度（varlen）时，注意力内核自然存在负载不均衡——分配给 SM 的工作图块的主循环长度不同，因为某些工作图块需要比其他图块更多的加载和 MMA 操作。此外，我们可以选择 SM 处理图块的顺序，例如通过定义网格坐标的优选线性化。抽象掉注意力的任何特定特性，我们可以将适用于相同并行处理器的完工时间最小化的一般结果应用于我们的上下文。特别地，在 FlashAttention-4 中，我们使用了经典的最长处理时间优先（LPT）调度思想[9]。我们强调，我们应用这一思想的方式适用于所有 GPU 架构，并且也已被验证为在 Hopper GPU 上对 FlashAttention-3 的改进。

• 用于因果掩码的 LPT： 标准的注意力网格由（mblocks, heads, batches）给出，并按从左到右的递增顺序计算。但由于分数在对角线上方被掩蔽，对于固定的头部和批次，SM 最终会以从最短到最长的顺序低效地处理工作图块。另一方面，一个朴素的 LPT 顺序也是次优的，因为对于不同的批次，主循环 KV 加载不会在 L2 缓存中命中，并且如果超过 L2 缓存容量，首先加载所有 KV 头部可能会颠簸 L2 缓存。相反，我们始终将批次作为最外层维度处理，并在头部上扫描。这意味着我们将头部划分为不超过 L2 缓存容量的部分；然后图块调度器按每部分头部、逆序的 mblocks、部分、最后是批次的顺序遍历网格。特别地，对于 MQA 或 GQA，我们总是在改变 mblocks 之前，遍历每个 KV 头部的所有查询头部。经验上，我们验证了这种 LPT 顺序非常有效；例如，对于 BF16 和头维度 128，在 H200 GPU 上测量，我们对 MHA 获得了 4-8% 的 FLOPS 增益，对 MQA 8 获得了 7-14% 的 FLOPS 增益。

• 用于可变序列长度的 LPT： 对于 varlen，我们还必须应对由于批次间差异导致的负载不均衡。例如，在解码工作负载中，不同的批次可能关注不同数量的上下文，而在混合或连续批处理中，某些批次可能处于预填充阶段，而其他批次处于解码阶段。每个批次的查询和 KV 序列长度列表通常作为注意力元数据存储在设备上，标准的 varlen 注意力内核在按递增顺序处理批次时读取这些整数。然而，给定的批次顺序在负载平衡方面可能是任意次优的——例如，可能有短的正方形预填充，后跟长上下文的解码。为了改善这一点，我们可以通过启动一个预处理内核来强制实施 LPT 顺序，该内核根据每个工作图块的最大执行时间对批次进行排序，写出额外的元数据（虚拟到实际批次索引的映射），随后注意力内核会读回这些元数据，以便按排序后的顺序遍历批次。这个元数据可以被缓存，因此排序不会导致性能损失。

4 语言与框架

我们完全用嵌入在 Python 中的 CuTe-DSL [21] 编写 FLASHATTENTION-4，没有任何 CUDA C++ 组件。CuTe-DSL 编译器获取 Python 源代码，将其降级为 PTX，然后使用 PTX 编译器（ptxas）最终生成汇编代码（SASS）。

• 具有清晰抽象的完全表达力： CuTe-DSL 编程模型与 CUTLASS C++ 同构，确保 FLASHATTENTION-4 在保留底层 GPU 编程的完全表达力的同时，受益于用 Python 而非 C++ 进行元编程的生产力提升以及快速的 JIT 编译时间。CuTe-DSL 提供对 PTX 的直接访问，作为逃生舱口，允许开发人员实现他们需要的任何功能，而不受框架限制。例如，我们利用定制的 PTX 序列来实现尚未在 CuTe-DSL API 中完全公开的操作（尽管这些将在未来版本中集成），这表明我们的框架并未将开发人员限制在 GPU 功能的有限子集内。

• 通过 JIT 实现快速编译： 由于复杂的 C++ 模板元程序，编译时间一直是过去 FlashAttention 实现的一个瓶颈。通过将 CuTe-DSL 嵌入 Python 并进行即时（JIT）编译，与传统的基于 C++ 模板的方法相比，FlashAttention-4 实现了更快的构建时间。如表 4 所示，与 FlashAttention-3 相比，FlashAttention-4 将编译时间减少了 20-30 倍。这个快速的迭代周期显著提高了开发者的生产力，使得在内核开发期间能够进行更快的实验和调试。

表 4：单个内核的编译时间：FA3（C++ 模板）和 FA4（CuTe-DSL）。通常 FA2 和 FA3 需要为不同的注意力变体预编译数百个内核。

• 灵活性和可访问性： 基于 Python 的框架已经在实践中证明了其灵活性：开发者已经成功地在 FlashAttention-4 之上构建了 FlexAttention 和块稀疏注意力变体，而无需修改核心框架。通过降低准入门槛，我们的方法使得仅有几个月 GPU 编程经验的研究人员和工程师也能够贡献有意义的扩展，而无需深入了解 C++ 模板元编程。这种可访问性加速了创新，并允许注意力机制研究界更快速地探索新的算法变体。

我们的愿景是提供一个全面的框架，用于构建各种具有最佳性能的注意力变体。FlashAttention-4 并非从头实现每个注意力变体，而是将通用功能分解为独立、可组合的原语。诸如块稀疏模式、掩码策略、可变序列长度处理和工作调度等操作都被公开为正交原语，可以自由组合。这种模块化设计确保优化和新功能惠及在该框架上构建的所有注意力实现，同时通过编译为高效的 GPU 内核，仍然达到最高性能。

5 实证评估

我们评估了 FLASHATTENTION-4 相对于各种开源和闭源基线的效率。

5.1 前向传递

我们在图 4 和 5 中报告了前向传递的结果，显示 FLASHATTENTION-4 比 cuDNN 9.13 快 1.1−1.3 倍，比 Triton 快 2.1−2.7 倍。对于中长序列（4k 及以上），在不同的头维度和因果掩码设置下，FLASHATTENTION-4 持续优于所有基线。在因果情况下，增益更大，我们将其归因于最长处理时间优先（LPT）调度器。

图 4：B200 上前向传递 TFLOPS（FP16/BF16），头维度 128。左：非因果注意力。右：因果注意力。FA4 在各种序列长度上比 cuDNN 9.13.0 实现了 1.1−1.3× 加速，比 Triton 实现了 2.1−2.7× 加速。自我们实现首次发布以来，较新版本的 cuDNN 已经采纳了本文描述的许多技术，产生了与 FA4 相似的性能。

图 5：B200 上前向传递 TFLOPS 比较（FP16/BF16），头维度 (192, 128)，因果注意力（通常用于 DeepSeek V3 架构）。

5.2 反向传递

我们在图 6 中报告了反向传递的结果。FLASHATTENTION-4 在长序列长度和因果掩码下实现了持续的加速，证明了我们 2-CTA 反向传递的有效性。

我们还在图 7 中展示了确定性反向传递的性能。我们精心设计的交织和调度使得确定性反向传递快得多，达到了 1-CTA 反向传递非确定性版本速度的 75%。

图 6：B200 上反向传递 TFLOPS（FP16/BF16），头维度 128。左：非因果注意力。右：因果注意力。

图 7：B200 上确定性反向传递的消融研究（FP16/BF16），头维度 128。因果注意力 – SPT、带逆序 mblock 的 LPT、LPT，以及无批次/头部交织的朴素方法。

6 讨论与结论

FlashAttention-4 解决了非对称硬件扩展问题，其中张量核心速度如此之快，以至于主要瓶颈转移到共享内存流量和指数吞吐量，这促使我们进行算法和内核的协同设计以缓解这些限制。我们围绕完全异步的 MMA 重新设计了流水线，以将 softmax 与更大图块的矩阵乘法重叠，并引入了软件模拟的指数和有条件的 softmax 重新缩放，以减少非矩阵乘法操作。我们利用张量内存和 2-CTA MMA 模式来减少共享内存流量。此外，2-CTA 使得重构全局原子累积成为可能，将全局原子加操作的次数减半。

FlashAttention-4 完全用嵌入在 Python 中的 CuTe-DSL 实现，在保持底层控制的同时，实现了比 C++ 模板内核快 20-30 倍的编译时间。虽然针对 Blackwell GPU 进行了优化，但随着计算持续超越非矩阵乘法单元，其中一些算法可以扩展到其他加速器。

附录 A 实验和基准测试的补充细节

A.1 系统和库

我们在 B100 180GB SXM6 (1000W) 上对速度进行基准测试。我们预热 5 次，然后重复基准测试 10 次，并取平均时间。

我们通常使用撰写本文时（2025 年 3 月）的最新版本库。具体来说，我们使用：

• CUDA 13.1

• FLASHATTENTION 2.8.3

• Triton 3.6

• PyTorch 2.10.0

• CuTe-DSL 4.4.1

对于 cuDNN，在主要论文中，我们与 cuDNN 9.13 和最新版本 cuDNN 9.19.1.2 进行了比较。从版本 9.13 和 9.14 [20] 开始，我们与 cuDNN 团队合作，将 FLASHATTENTION-4 中的一些技术融入到 cuDNN 中，以便我们的工作能够惠及尽可能多的实践者。

A.2 确定性反向非因果传递

为了完整起见，我们在图 8 中还包含了无因果掩码的确定性反向内核的性能数据，与因果掩码的情况并列展示。

图 8：B200 上确定性反向传递的消融研究，头维度 128。左：非因果注意力 – 批次/头部交织与朴素方法对比。右：因果注意力 – SPT、带逆序 mblock 的 LPT、LPT，以及无批次/头部交织的朴素方法。