推理优化 | FlashAttention 1-4 详解

本文受 Datawhale llm-algo-leetcode 活动启发而作，具体参考了其中的 FlashAttention 仿真实现与 Triton Flash Attention 示例两个章节。该项目系统整理了 LLM 推理/训练中的经典算法，推荐关注。

算法#

原算法：标准 Attention 与显存瓶颈#

Transformer 的标准缩放点积注意力定义为：

\text{Attention}(Q, K, V) = \text{softmax}\!\left(\frac{QK^T}{\sqrt{d_k}}\right) V

其中 $Q, K, V \in \mathbb{R}^{N \times d}$ ， $N$ 为序列长度， $d$ 为头维度。计算流程如下：

\begin{aligned} S &= Q K^T &&\in \mathbb{R}^{N \times N} \quad \text{(矩阵乘)} \\ P &= \text{softmax}\!\left(\frac{S}{\sqrt{d}}\right) &&\in \mathbb{R}^{N \times N} \quad \text{(逐行 softmax)} \\ O &= P V &&\in \mathbb{R}^{N \times d} \quad \text{(矩阵乘)} \end{aligned}

瓶颈分析：

$S$ 和 $P$ 均为 $O(N^2)$ 大小的中间矩阵，必须写入 HBM 再读出，带宽严重不足
GPU HBM 带宽（A100 ~2TB/s）远低于片上 SRAM（~19TB/s），而 Attention 的计算量仅为 $O(N^2 d)$ ，使得整体呈 Memory-Bound
对长序列（如 $N \ge 2048$ ），超过 90% 的时间耗费在显存读写上

1
HBM 数据流视角（标准实现）：
2
  Load Q,K → Compute S → Write S to HBM
3
  Read S → Softmax → Write P to HBM
4
  Read P,V → Compute O → Write O to HBM
5
每一步跨越慢速 HBM，瓶颈显著。

FlashAttention 1（2022）：Tiling + Online Softmax#

论文：Dao et al., FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, NeurIPS 2022.

核心思想：将注意力矩阵分解为 tile，在 SRAM 内完成关键计算，避免 $S$ 和 $P$ 写入 HBM。

Online Softmax 数学推导#

标准 softmax 需要两遍扫描——先求最大值做数值稳定，再求和做归一化：

\text{softmax}(x_i) = \frac{e^{x_i - m}}{\sum_j e^{x_j - m}}, \quad m = \max_k x_k

Online Softmax 将其融合为单次遍历，不断用新的局部最大值修正历史结果。设当前处理第 $j$ 个 tile，已累积的全局最大值为 $m$ 、归一化分母为 $\ell$ ：

\begin{aligned} m_{\text{new}} &= \max(m, \text{rowmax}(S_{ij})) \\[4pt] \widetilde{P}_{ij} &= \exp\!\left(S_{ij} - m_{\text{new}}\right) \\[4pt] \ell_{\text{new}} &= e^{m - m_{\text{new}}} \cdot \ell + \sum \widetilde{P}_{ij} \\[4pt] O_i^{\text{new}} &= e^{m - m_{\text{new}}} \cdot O_i + \widetilde{P}_{ij} V_j \end{aligned}

最终输出为 $O_i = O_i^{\text{new}} / \ell_{\text{new}}$ 。由于 $m_{\text{new}} \ge m$ ，修正因子 $e^{m - m_{\text{new}}} \le 1$ ，数值稳定。

PyTorch 实现#

1
def flash_attention_v1(Q, K, V, Br=128, Bc=128):
2
    N, d = Q.shape
3
    O = torch.zeros_like(Q)
4
    m = torch.full((N,), -float('inf'), device=Q.device)
5
    ell = torch.zeros((N,), device=Q.device)
6

7
    for i in range(0, N, Br):
8
        Qi = Q[i:i+Br]
9
        Oi, mi, li = O[i:i+Br], m[i:i+Br], ell[i:i+Br]
10

11
        for j in range(0, N, Bc):
12
            Kj, Vj = K[j:j+Bc], V[j:j+Bc]
13
            Sij = Qi @ Kj.T
14
            mi_new = torch.max(mi, Sij.max(dim=1).values)
15
            Pij = torch.exp(Sij - mi_new.unsqueeze(1))
16
            li = torch.exp(mi - mi_new) * li + Pij.sum(dim=1)
17
            Oi = torch.exp(mi - mi_new).unsqueeze(1) * Oi + Pij @ Vj
18
            mi = mi_new
19

20
        O[i:i+Br] = Oi / li.unsqueeze(1)
21

22
    return O

FlashAttention 2（2023）：并行性重构#

论文：Dao, T., FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning, 2023.

FlashAttention 1 已经逼近 IO 理论下界，但并行效率仍有提升空间。FlashAttention 2 的改进思路是”让矩阵乘法占据更多 GPU 时间，减少其他操作的开销”。

内外循环交换#

这是 FlashAttention 2 最关键的改动：

1
FlashAttention 1 循环结构:                  FlashAttention 2 循环结构:
2
  for Qi in Q_blocks:  (外循环)              for Kj,Vj in KV_blocks: (外循环,可沿序列并行)
3
    for Kj,Vj in KV_blocks: (内循环)           for Qi in Q_blocks:   (内循环,独立可并行)
4
      Sij = Qi @ Kj^T                           Sij = Qi @ Kj^T

为什么这样改？

FlashAttention 2 的外循环沿 KV 序列维度划分，每个 KV block 交给不同的 thread block 并行处理
内循环的 $Q_i$ 之间没有依赖，可完全并行
$K, V$ 被加载到 SRAM 后可供多个 $Q_i$ block 复用

减少非矩阵乘运算占比#

FlashAttention 1 中 softmax 相关的 rescale、mask、dropout 等 element-wise 操作虽然计算量小，但占据了 GPU 调度时间。FlashAttention 2 将这些操作尽可能后置、合并，让 Tensor Core 的 GEMM 时间占比接近 100%。

用公式表示，FlashAttention 2 将 softmax 的归一化步骤延迟到内循环全部结束之后，减少了每次迭代中的除法指令：

\begin{aligned} \text{FlashAttention 1:}&\quad O_i \leftarrow \text{rescale}(O_i) + P_{ij}V_j &&\text{(每次迭代都 rescale)} \\ \text{FlashAttention 2:}&\quad O_i \leftarrow O_i + P_{ij}V_j,\quad \text{最终 } O_i / \ell_i &&\text{(仅一次归一化)} \end{aligned}

注：FlashAttention 3（Shah et al., 2024）的算法架构与 FlashAttention 2 一致，主要针对 Hopper (H100) 架构做了硬件级优化（见”其他优化”部分），此处不单独讨论算法。

FlashAttention 4（2026）：面向非对称硬件扩展的算法-Kernel 流水线协同设计#

论文：Zadouri, Hoehnerbach, Shah et al., FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling, arXiv:2603.05451, March 2026.

背景：Blackwell 架构（B200/GB200）带来了”非对称硬件扩展”——Tensor Core 的算力翻倍，但共享内存带宽、指数函数单元等其他部件的提升不同步甚至停滞。这意味着之前在 Hopper 上不是瓶颈的组件，现在变成了新的瓶颈。

FlashAttention 4 就是围绕如何应对这种”不对称”，重新设计算法-系统的协同优化方案。

三项核心技术#

（1）全异步 MMA 流水线 + 更大 tile

Blackwell 支持完全异步的矩阵乘加（MMA）操作。FlashAttention 4 重新设计了计算流水线，使用更大的 tile 尺寸配合异步执行，让 Tensor Core 在等待数据时不空闲：

\text{Tile 大小从 } (B_r, B_c) \to (B_r', B_c'), \quad B_r' \approx 2\times B_r

（2）软件模拟指数函数 + 条件 softmax rescale

Blackwell 的指数函数单元带宽没有成比例增长。FlashAttention 4 用软件模拟的指数函数替代硬件 exp 指令，并结合”条件 rescale”策略——仅在确实需要时才执行 softmax 的修正步骤，减少非矩阵乘运算占比：

\text{rescale 仅在 } m_{\text{new}} > m \text{ 时触发，否则跳过}

（3）Tensor Memory + 2-CTA MMA 模式

利用 Blackwell 的 Tensor Memory 层级和 2-CTA MMA 模式，将反向传播中的 shared memory 流量和 atomic add 操作大幅减少。两个 CTA（Cooperative Thread Array）共用同一块 MMA 单元，减少了数据搬移开销。

性能#

相比 cuDNN 9.13 快 1.3×
相比 FlashAttention 3 快 2.7×（在同代 Blackwell 硬件上）

其他优化#

以下为各版本中超出核心算法的工程优化技术。

FlashAttention 2 的工程优化#

减少 warp 间通信：forward 中调整计算顺序（Q 与 K/V tile 的配对方式），减少 shared memory 读写，消除跨 warp 的同步等待
Bank conflict 消除：Q 按行优先、K 按列优先存储于 shared memory，避免同一 warp 内的线程访问同一 bank
Causal mask 下三角计算：利用下三角结构跳过被 mask 的 tile，将计算量减半至 $\approx N^2/2$

FlashAttention 3 的硬件特性利用（Hopper / H100）#

WGMMA 指令：替代传统 mma.sync，支持 warp group 级别的异步矩阵乘，在共享内存和寄存器之间直接传输
TMA 硬件拷贝：Tensor Memory Accelerator 异步从 HBM 加载数据到 SRAM，完全卸载数据传输开销
GEMM-Softmax 流水线重叠：Tensor Core 做下一 tile 的矩阵乘时，CUDA Core 同时做当前 tile 的 softmax
FP8 低精度：forward 中使用 FP8 Tensor Core 加速矩阵乘，反向传播保持 BF16 精度

FlashAttention 4 的工程优化（Blackwell / B200）#

全异步 MMA 流水线：利用 Blackwell 的异步 MMA 指令，将 GEMM、指数运算、归一化三级流水覆盖，消除等待延迟
软件指数函数模拟：绕过硬件 exp 单元瓶颈，使用多项式近似 + 分段线性拟合替代
条件 softmax rescale：仅在全局最大值变化时才执行修正，减少冗余的除法/指数运算
2-CTA MMA 模式：两个 CTA 共享 MMA 计算单元，减少反向传播中的 atomic add 和 shared memory 流量
Tensor Memory 层级利用：直接使用 Blackwell 的 Tensor Memory 存储中间结果，绕过 shared memory 瓶颈

总结#

	FlashAttention 1	FlashAttention 2	FlashAttention 3	FlashAttention 4
论文标题	FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness	FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning	FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-Precision	FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling
中文标题	具备 IO 感知能力的快速低显存精确注意力机制	通过更好并行性与工作划分实现更快的注意力计算	结合异步执行与低精度计算的快速精确注意力机制	面向非对称硬件扩展的算法与 Kernel 流水线协同设计
GPU 型号	A100 (Ampere)	A100 / H100	H100 (Hopper)	B200 / GB200 (Blackwell)
核心算法改进	IO 感知分块、增量 softmax	内外循环交换、序列维度并行	流水线异步执行、低精度计算	全异步流水线、软件指数模拟、条件 rescale、2-CTA MMA
关键工程优化	SRAM 驻留计算、反向传播重计算	减少 warp 通信、bank conflict 消除、causal 稀疏化	WGMMA 指令、TMA 硬件拷贝、FP8 低精度	异步 MMA 流水线、软件 exp 替代、Tensor Memory 层级利用

参考#

Datawhale llm-algo-leetcode：在线文档 | 代码仓库
FlashAttention 1：Dao, T., et al. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022. doi:10.48550/arXiv.2205.14135
FlashAttention 2：Dao, T. FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. 2023. doi:10.48550/arXiv.2307.08691
FlashAttention 3：Shah, J., et al. FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-Precision. 2024. doi:10.48550/arXiv.2407.08608
FlashAttention 4：Zadouri, T., Hoehnerbach, M., Shah, J., et al. FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling. arXiv:2603.05451, 2026. doi:10.48550/arXiv.2603.05451

音乐

音乐

算法#