FlexAttention + FlashAttention-4：既快又灵活

要点速览

在 Hopper 和 Blackwell GPU 上，FlexAttention 现在拥有了 FlashAttention-4 后端。

我们在 PyTorch 中增加了对自动生成 CuTeDSL 分数/掩码修改函数的支持，并实现了针对自定义注意力变体的 FlashAttention-4 即时（JIT）实例化。

在计算密集型工作负载上，这使得性能相比现有的 Triton 实现提升了 1.2 倍至 3.2 倍。

FlexAttention 回顾

FlexAttention 是一个 PyTorch API，让你只需几行 Python 代码即可实现自定义注意力变体，无需编写 CUDA。你只需编写一个 score_mod 或 mask_mod 函数来修改注意力分数，编译器会处理剩下的工作：ALiBi、滑动窗口、文档掩码、软截断（soft-capping）及其组合均可通过同一接口实现。

从底层来看，它在原生 FlashAttention 之上进行了两个扩展：

1. 对 Softmax 前分数的逐点修改，支持从全局内存进行任意加载。
2. 前向和反向传播的块稀疏迭代，并使用简单的数据结构在运行时编码与数据相关的稀疏性。

就是这样。当然，魔鬼藏在细节中，但正如我们在原始 FlexAttention 博客和推理版 FlexAttention 博客中所展示的，这两个扩展涵盖了广泛的主流注意力变体。

此次发布后，FlexAttention 现在拥有了 FlashAttention-4 (FA4) 后端。使用方法如下：

import torch
from functools import partial

from torch.nn.attention.flex_attention import flex_attention

flex_flash = torch.compile(
    partial(flex_attention, kernel_options={"BACKEND": "FLASH"}), dynamic=False
)

def local_boost(score, b_idx, h_idx, q_idx, kv_idx):
    return torch.where(torch.abs(q_idx - kv_idx) <= 8, score * 2, score)

B, H, S, D = 2, 8, 2048, 128
q = torch.randn(B, H, S, D, device="cuda", dtype=torch.bfloat16)
k = torch.randn(B, H, S, D, device="cuda", dtype=torch.bfloat16)
v = torch.randn(B, H, S, D, device="cuda", dtype=torch.bfloat16)
out = flex_flash(q, k, v, score_mod=local_boost)

将 BACKEND 设置为 "FLASH" 以使用 FA4 后端。你需要安装较新的 PyTorch nightly 版本和较新的 FlashAttention 代码库；请查看安装文档以了解版本兼容性。这是处于快速开发中的代码，在趋于稳定前可能会有一些破坏性变更。

用 FlexAttention 普及注意力研究

FlexAttention 最初设计（并命名）的初衷是为 AI 研究人员在原型设计和尝试新的注意力变体时提供灵活性。实践证明确实如此：已有数十篇论文引用了 FlexAttention，且超过一千个代码库采用了它：

尽管 Flex 在赋能研究人员方面取得了成功，但用户不断反馈的一个问题是，他们最终会遇到难以突破的性能上限。在最初发布博客文章时，我们将它与 Hopper GPU 上的 FlashAttention-3 (FA3) 进行了比较，性能大约是后者的 80%。

如果今天再进行测量，尽管两种实现都有改进，但 FlexAttention 的吞吐量仅约为 FlashAttention-3 的 60%！

一个常见的模式出现了：研究人员用 Flex 进行实验，找到可行方案后，一旦性能成为关键，就会撞上“墙”。此时，专家必须将其移植到底层实现中。FlashAttention-3 不断增加新参数以扩展功能，但每一个新参数/模式都需要底层重写。我们将负担从研究人员转移到了 ML 工程师身上！

在 Hopper 上，完全优化版本之间的性能差异或许值得用灵活性来交换，但在更新的 Blackwell GPU 上情况则完全不同。

让我们看看在 Blackwell GB200 GPU（1000W 功耗）上，使用现有 Triton 实现（开启所有自动调优开关）的 FlexAttention 与通过 PyTorch SDPA 可用的高度优化实现（如 cuDNN attention）之间的对比：

曾经微小的差距如今已变成一道鸿沟！

Blackwell：更大的张量核心，更大的挑战

在 Blackwell 上，高性能注意力需要深度流水线化、针对 Warp 优化的内核。这些技术在基于 Triton 的实现中无法表达。我们推荐阅读 逆向工程 FlashAttention-4 中的出色解释，它详细介绍了该实现如何利用 Blackwell 的新硬件功能，以及 softmax 计算方式的更新。游戏的规则一如既往：保持张量核心的繁忙；由于它们变得更快了，这需要大量使用深度异步流水线。

Blackwell 引入了张量内存（TMEM），这是一种程序员可管理的、靠近张量核心的暂存器，用于存储中间结果。更重要的是，数据移动和矩阵乘法现在都是完全异步的。Warp 可以启动矩阵乘法或加载操作并立即继续执行后续任务。

Warp 特化将工作拆分为多个阶段：一些 Warp 处理同步工作（如需要寄存器的 softmax），而另一些 Warp 通过发出加载和矩阵乘法请求来编排异步流水线，并协调同步。由于编排 Warp 的寄存器压力较低，更多的操作可以同时保持在运行状态，从而隐藏延迟。

张量核心变大变快了，但负责指数运算等特殊函数单元（SFU）的进步速度却没跟上。对于前向注意力，这改变了瓶颈：softmax 的 exp() 运算现在和矩阵乘法一样昂贵。为了让 GPU 保持完全饱和，你需要在这两个 tile 之间进行乒乓切换，将一个 tile 的矩阵乘法与另一个 tile 的指数运算重叠。下面的时间轴显示了这些阶段如何交替以隐藏延迟。

反向传播更加棘手。TMEM 中没有足够的空间同时容纳所有的累加器，因此内核需要精细的流水线化，以便在共享内存、寄存器和张量内存处于沉重压力下时，将计算与数据移动重叠。

这正是通用编译器难以轻易发现的底层编排。正如 Gluon 简介 所言：“虽然 Triton 编译器在为各种内核生成高效代码方面做得很好，但它仍会被手动调优的底层代码击败。在这种情况下，用户几乎无法做任何事来显著提高性能，因为所有的细节都被隐藏了。”对于 FlexAttention 来说这更难，因为它是一种元注意力实现，当模式由用户定义时，硬编码针对特定模式的编译器优化非常困难。正因如此，我们开始研究底层实现，以寻找提高 Blackwell 性能的最佳方式。

以 FlashAttention-4 为基础

在针对 Blackwell 新硬件的注意力实现方面，存在着很大的变数。cuDNN 较早地添加了高性能注意力支持，但 FA3（Hopper 上的现有 SOTA 实现）无法在 Blackwell 上工作。WGMMA 在 SM100 上已不再存在：它已被 TCGEN05 张量核心指令取代，且张量核心操作需要不同的内存空间（张量内存）。

Tri Dao 等人开始研究 FlashAttention-4，这是一个能够充分利用硬件的更新版本实现。

从 FA3 到 FA4 的一个重大变化是 CuTeDSL，这是 NVIDIA CUTLASS 团队最近发布的一个 Python DSL，用于使用 CUTLASS 抽象编写高性能 CUDA 内核。SOTA 注意力实现大量使用了 CUTLASS 抽象（如 cute.Layouts），但任何尝试过安装 FlashAttention 的人都知道这有多痛苦，且编译时间很长。因此，虽然之前有人提出过用 CUTLASS C++ 重写 Flex 的想法，但 FlexAttention 的动态特性（以及编译开销）使得这一前提不再具有吸引力。CuTeDSL 允许用 Python 编写过去需要 CUTLASS C++ 才能实现的功能，这使得 JIT 风格的工作流对 FlexAttention 来说更具实用性。

在与 Tri Dao 对此路径进行初步讨论后，我们决定联手为 FlexAttention 和 FlashAttention-4 共同推进这一实现。

我们没有构建一个单独的实现，而是通过协作直接扩展了 FA4，共享相同的异步流水线基础设施，并在 FlexAttention 需要注入分数修改和稀疏性的地方增加了扩展点。

这意味着在前向和反向传播中都增加了分数修改支持（通过将分数修改内联到 FlashAttention 实现中），并增加了对 FlexAttention 中使用的块稀疏元数据的支持。

这项工作大致分为两部分：FA4 的变更以生成 FlexAttention 模板实例化，以及 Inductor 的更新以从 PyTorch 表示生成所需的 CuTeDSL 代码。

Inductor → CuTeDSL：胶水层

那么我们需要为 FlexAttention 生成什么呢？逐点修改和任意加载。幸运的是，Inductor 并不是第一次做这件事，这种扩展已有现成的机制。例如，让我们看看一个可用于实现 ALiBi 的分数修改（有趣的是：这正是 FlexAttention 项目的推动示例）。

粗略地说，torch.compile 获取用户代码并将其通过多个 IR（中间表示）进行转换。这些转换会产生越来越底层的表示。在 FX IR 中，你仍然可以看到熟悉的 PyTorch 算子，以及在使用后立即被置为 None 的变量。AOTAutograd 传递会自动生成反向传播：由于 (X + A) 的导数等于 1，链式法则将梯度直接传递。

值得注意的是，在到达 Inductor 并最终生成待执行的内核代码之前，该堆栈的任何部分都不需要“知道”什么是 CuTeDSL 代码。

点击下面的标签页，查看 ALiBi 如何从用户代码演变为最终的 CuTeDSL 内核。

原始 FX IR AOTAutograd CuTeDSL

def alibi_mod(score, b, h, q_idx, kv_idx):
    scale = torch.exp2(-((h + 1) * 8.0 / H))
    bias = (kv_idx - q_idx) * scale
    return score + bias

class score_mod_0(torch.nn.Module):
    def forward(self, child, child_1, child_2, child_3, child_4):
        add: "i32[][]cuda:0" = child_2 + 1;  child_2 = None
        mul: "f32[][]cuda:0" = add * 8.0;  add = None
        truediv: "f32[][]cuda:0" = mul / 32;  mul = None
        neg: "f32[][]cuda:0" = -truediv;  truediv = None
        scale: "f32[][]cuda:0" = torch.exp2(neg);  neg = None
        sub: "i32[][]cuda:0" = child_4 - child_3
        bias: "f32[][]cuda:0" = sub * scale
        add_1: "f32[][]cuda:0" = child + bias
        return add_1

class fw_graph0(torch.nn.Module):
    def forward(self, arg0_1, arg1_1, arg2_1, arg3_1, arg4_1):
        add = torch.ops.aten.add.Tensor(arg2_1, 1)
        mul = torch.ops.aten.mul.Tensor(add, 8.0)
        div = torch.ops.aten.div.Tensor(mul, 32)
        neg = torch.ops.aten.neg.default(div)
        exp2 = torch.ops.aten.exp2.default(neg)
        sub = torch.ops.aten.sub.Tensor(arg4_1, arg3_1)
        mul_1 = torch.ops.aten.mul.Tensor(sub, exp2)
        add_1 = torch.ops.aten.add.Tensor(arg0_1, mul_1)
        return add_1

class joint_graph0(torch.nn.Module):  # backward pass
    def forward(self, arg0_1, ..., arg5_1):
        convert = torch.ops.prims.convert_element_type.default(arg5_1, torch.bfloat16)
        return [convert, None, None, None, None]

@cute.jit
def score_mod(tSrS_ssa, b_idx, h_idx, q_idx, kv_idx, seqlen_info, aux_tensors):
    tmp1 = tSrS_ssa.to(cutlass.Float32)
    tmp4 = (kv_idx - q_idx)
    tmp5 = tmp4.to(cutlass.Float32)
    tmp7 = (h_idx + cute.full_like(h_idx, 1))
    tmp8 = tmp7.to(cutlass.Float32)
    tmp9 = (tmp8 * cute.full_like(tmp8, 8.0))
    tmp10 = (tmp9 * cute.full_like(tmp9, 0.03125))
    tmp11 = -tmp10
    tmp12 = (tmp11 * cute.full_like(tmp11, 0.6931471805599453))
    tmp13 = cute.math.exp2(tmp12 * 1.4426950408889634)
    tmp14 = (tmp5 * tmp13)
    tSrS_ssa = (tmp1 + tmp14)

@cute.jit
def score_mod_bwd(grad_score_mod_ssa, tSrS_ssa, ...):
    grad_score_mod_ssa_out = grad_score_mod_ssa

def tanh_softcap(score, b, h, q_idx, kv_idx):
    return soft_cap * tanh(score / soft_cap)

bias = torch.randn(seq_q, seq_kv, device='cuda', requires_grad=True)
def bias_func(score, b, h, q_idx, kv_idx):
    return score + bias[q_idx, kv_idx]