FlashAttention：IO 感知的精确注意力加速算法

FreeGuideOnline 最新 2026-06-14

什么是 FlashAttention？

Transformer 架构已经在 NLP、计算机视觉等领域无处不在，其核心模块“缩放点积注意力”的计算复杂度与输入序列长度成二次方关系。对于长序列，标准的注意力实现会在 GPU 高带宽内存（HBM）与片上 SRAM 之间产生大量读写，内存带宽成为了真正瓶颈。FlashAttention 正是针对这一问题的精确加速算法——它通过 IO 感知（IO‑aware）的分块策略与重计算技术，在保持数学等价性的前提下，大幅减少 HBM 访问次数，从而将注意力计算速度提升数倍并显著节省内存。

背景：为什么标准注意力会慢？

标准注意力计算与内存瓶颈

给定查询 Q、键 K、值 V，形状均为 [N, d]，其中 N 为序列长度，d 为头维度。标准缩放点积注意力的计算过程为：

S = Q × Kᵀ   (形状 [N, N])
P = softmax(S / √d)  (按行归一化)
O = P × V

在实际 GPU 实现中，中间矩阵 S 和 P 都必须在 HBM 中显式存储。N 增大时，这两个 N×N 矩阵会迅速消耗内存，并且每一步都要将张量从 HBM 写入 DRAM 再读回，带来巨大的内存带宽压力。大多数 GPU 上的注意力计算属于内存限制型，而非计算限制型，即处理单元大部分时间在等待数据移动。

内存层次结构速览

现代 GPU 的内存层次大致分为：

  • HBM（高带宽内存）：容量大（数十 GB），带宽相对低（例如 A100 为 1.5 TB/s），片外。
  • SRAM（片上共享内存）：容量小（A100 共享内存约 192 KB / SM），带宽极高（10–20 TB/s），片上。

标准实现中，S 和 P 存储在 HBM，反复的加载/存储会浪费大量时间；而 SRAM 高速但太小，无法完整放下 N×N 矩阵。

FlashAttention 的核心思想：IO 感知

FlashAttention 的标题就给出了答案：IO 感知的精确注意力加速。关键洞察是：我们不必在 HBM 中显式构造完整的注意力矩阵 S 和 P，而只需通过分块（tiling）将计算分解为许多小型、可融合的操作，使得中间结果在高速 SRAM 中产生并使用后即丢弃，最终只将输出 O 写回 HBM。

两大支柱：分块与重计算

1. 分块前向传播

将 Q、K、V 沿着序列维度切成多个块，每次只加载一个 Q 块和一个 K、V 块到 SRAM。在块内计算局部注意力，同时通过维护归一化统计量（行和与最大值）来增量地更新全局 softmax 结果。

这样，我们永远不需要在 HBM 中分配完整的 N×N 矩阵，内存占用从 O(N²) 降为 O(N)。

经典 softmax 的块式增量计算：

  • 普通 softmax：softmax(x_i) = exp(x_i - max) / Σ exp(x_j - max)
  • 当分块时，可以先计算块内的最大值与和，再根据后续块出现的更大值对之前的结果进行重新缩放修正。FlashAttention 通过精心设计的标量 l（分母和）和 m（最大值）来维护此状态。

2. 重计算反向传播（Recomputation）

在反向传播中，通常需要前向的注意力矩阵 P 或 S。为了不保存这些大矩阵，FlashAttention 在反向时重新计算需要的中间值。具体做法是：在反向过程中，以相同的方式分块加载 Q、K、V 并重新计算所需的前向部分，这实际是将内存节省与额外计算进行交换。由于重计算仅引入可忽略的 FLOPs 增加，而减少了大量 HBM 访问，总体仍然大幅加速。

算法细节：逐步拆解

以下以前向算法为例，展示 FlashAttention 是如何在数学等价的情况下分块计算的。假设使用 scale 因子 τ = 1/√d。

初始化

将 Q、K、V 划分为 Tᵣ 个块（沿 N 轴），块大小为 Bᵣ（Q）和 B_c（K, V）。输出 O 初始化为零向量，并维护两个辅助向量：

  • m：每行的当前最大值，初始化为 -inf。
  • l：每行的 softmax 分母累加和，初始化为 0。

外层循环（遍历 K, V 块）

对于 j = 1 到 T_c：

  1. 加载 Kⱼ, Vⱼ 块从 HBM 到 SRAM。
  2. 内层循环（遍历 Q 块）： 对于 i = 1 到 Tᵣ：
    1. 加载 Qᵢ 块以及当前的 Oᵢ, mᵢ, lᵢ 到 SRAM。
    2. 在芯片上计算局部注意力分数： Sᵢⱼ = Qᵢ × Kⱼᵀ × τ （形状 [Bᵣ, B_c]）
    3. 计算当前块的局部最大值与 softmax： m̃ = rowmax(Sᵢⱼ) P̃ = exp(Sᵢⱼ - m̃) （局部未归一化概率） l̃ = rowsum(P̃)
    4. 根据从前一个 K 块累积得到的 mᵢ 和 lᵢ，更新全局统计量： m_new = max(mᵢ, m̃) l_new = exp(mᵢ - m_new) * lᵢ + exp(m̃ - m_new) * l̃
    5. 以修正的方式更新输出块： Oᵢ = diag(exp(mᵢ - m_new)) × Oᵢ + exp(m̃ - m_new) × P̃ × Vⱼ
    6. 将更新后的 Oᵢ, mᵢ = m_new, lᵢ = l_new 写回 HBM。

经过所有 K 块循环后，最终 O 矩阵即为精确的注意力输出。

反向传播类似地通过分块实现，重计算 S 和 P，并使用同样的 m、l 统计量，无需存储前向中间矩阵。

FlashAttention 的优势与性能表现

加速效果

  • 训练速度：在 GPT-2 等模型中，将序列长度从 1k 扩展到 8k 时，FlashAttention 相比标准 PyTorch 实现可实现 3–5 倍的前向/反向加速。
  • 内存节省：内存占用从 O(N²) 降至 O(N)，使得在相同硬件上训练更长序列成为可能。例如，原本只能在 A100 上训练 1k 长度的注意力，现在可以训练 4k 甚至更长。

精确性

  • FlashAttention 与原始注意力数学上完全相同，不是近似算法。块内 softmax 的重缩放公式保证了浮点运算的等价（忽略结合律带来的微小差异），因此零精度损失。

易用性

  • 多数深度学习框架已原生集成 FlashAttention：PyTorch 的 torch.nn.functional.scaled_dot_product_attention 在满足条件时自动调用；Hugging Face Transformers 中也通过 use_flash_attention_2=True 启用。
  • 只需少许环境配置（安装 CUDA 版 FlashAttention 包或使用最新 PyTorch 2+），即可替换原有注意力层。

从一个例子开始使用 FlashAttention

以下是最小化的 PyTorch 例子，展示如何利用内置 API 自动获得加速：

import torch
import torch.nn.functional as F

device = "cuda"
N, d = 4096, 64
Q = torch.randn(N, d, device=device, requires_grad=True)
K = torch.randn(N, d, device=device, requires_grad=True)
V = torch.randn(N, d, device=device, requires_grad=True)

# PyTorch 2.0+ 会自动为支持的情况选择 FlashAttention 内核
output = F.scaled_dot_product_attention(Q, K, V)

loss = output.sum()
loss.backward()

如果需要强制使用，可显式指定后端：

with torch.backends.cuda.sdp_kernel(enable_flash=True):
    output = F.scaled_dot_product_attention(Q, K, V)

对于 Transformer 模型，例如 Hugging Face 的 GPT-2：

from transformers import AutoModel
model = AutoModel.from_pretrained("gpt2", use_flash_attention_2=True, torch_dtype=torch.float16)

这样即可享受到长序列训练时的显著速度提升和内存节省。

总结

  • FlashAttention 通过 IO 感知 的分块计算和重计算，在不改变数学结果的前提下，大幅减少了 HBM 访问，缓解了注意力机制的内存带宽瓶颈。
  • 它将注意力内存复杂度从 O(N²) 降到 O(N)，使得长上下文训练成为现实。
  • 随着框架的集成，开发者几乎不用修改模型代码就可立即获得 2–4 倍的训练加速。

对于任何需要处理长序列的 Transformer 任务，FlashAttention 已经是事实上的标配组件。理解其背后的分块 softmax 原理和重计算思想，有助于更深入地优化其他内存受限的深度学习算子。