Flash Attention

Flash Attention（FlashAttention）是一种算法及一系列 GPU 内核，可为Transformer模型精确计算注意力，同时始终不将完整注意力矩阵写入 GPU 主存。对 N 个词元执行标准自注意力时，会构建 N × N 的分数矩阵，将其写入高带宽存储器（HBM），然后再读回以进行Softmax计算以及与值的乘积运算，因此注意力的大部分运行时间都花在数据搬运上。斯坦福大学的 Tri Dao 及其同事在 2022 年的论文《FlashAttention：具备 I/O 感知能力的快速、内存高效精确注意力》中介绍了该方法。它将计算拆分为适合快速片上 SRAM 的分块，因此内存使用量随序列长度线性增长，输出与标准注意力的结果一致，误差仅来自浮点舍入。其内核在 PyTorch 以及 vLLM 等大型语言模型（LLM）推理引擎中运行。

Flash Attention 的工作原理#

GPU的计算单元旁有少量极高速内存，芯片外则有大量速度较慢的内存。FlashAttention 论文以 NVIDIA A100 为例：它配有 40-80 GB、带宽为 1.5-2.0 TB/s 的 HBM；其 108 个流式多处理器各有 192 KB SRAM，估算带宽为 19 TB/s。注意力每读取一个字节执行的运算很少，因此其速度取决于数据在这两个层级之间传输的频率；该方法会计算这些传输量，这就是“具备 I/O 感知能力”的含义。

FlashAttention 会加载一块查询，然后将多块键和值流式传入 SRAM。对每个分块，它都会计算分数，为每个查询行更新运行中的最大值和总和，在最大值增大时重新缩放部分输出，并加上新分块的贡献。这种在线 Softmax 无需一次看到整行数据就能得到精确的 Softmax 结果，而且只有最终输出和每行的两个数值会返回 HBM。反向传播会根据 Q、K 和 V 重新计算分数分块，而不是存储 N × N 矩阵；虽然这增加了运算量，但由于读取的内存大幅减少，运行速度仍然更快。下面是使用PyTorch实现的相同循环，每次处理所有查询与一块键值：

import torch

torch.manual_seed(0)
n, d, block = 4096, 64, 256
q, k, v = torch.randn(3, n, d).unbind(0)

# Standard attention builds the full n x n score matrix in memory
reference = torch.softmax(q @ k.T / d**0.5, dim=-1) @ v

# FlashAttention-style: stream key/value tiles, keeping a running max and sum per query (online softmax)
row_max, row_sum, out = torch.full((n, 1), -torch.inf), torch.zeros(n, 1), torch.zeros(n, d)
for start in range(0, n, block):
    scores = q @ k[start : start + block].T / d**0.5  # only an n x block tile exists at once
    new_max = torch.maximum(row_max, scores.max(dim=-1, keepdim=True).values)
    rescale, probs = torch.exp(row_max - new_max), torch.exp(scores - new_max)
    row_sum = row_sum * rescale + probs.sum(dim=-1, keepdim=True)
    out = out * rescale + probs @ v[start : start + block]
    row_max = new_max
out = out / row_sum

print(f"Largest score tile: {tuple(scores.shape)} vs full matrix: ({n}, {n})")
print(f"Max difference from standard attention: {(out - reference).abs().max().item():.1e}")

最大的分数分块为 4,096 × 256，是 4,096 × 4,096 矩阵的十六分之一；结果与标准注意力的差异最大为 4.3e-07，相当于 FP32 舍入误差的大小。真实内核还会将查询拆分成多个分块，每个分块对应一个 GPU 线程块，并将整个循环融合到单个内核中，使分块不会离开 SRAM。

FlashAttention-2、-3 和 -4#

每个版本都面向更新一代的 GPU。FlashAttention-2（2023 年）重新安排了线程间的工作划分，运行速度约为原版的两倍，达到 A100 峰值吞吐量的 50-73%。FlashAttention-3（2024 年）使用 H100 的异步 Tensor Core，并支持FP8，FP16 下达到 740 TFLOPs/s，FP8 下接近 1.2 PFLOPs/s。FlashAttention-4（2026 年）面向 Blackwell GPU；据报告，它在 B200 上使用 BF16 时最高可达 1,613 TFLOPs/s，即峰值性能的 71%。这些内核包含在开源 flash-attention库中，支持 FP16和 BF16 输入；PyTorch 的 scaled_dot_product_attention会在 GPU 和输入支持时自动选择 FlashAttention-2 内核，否则回退到内存高效或普通实现。Meta 的 SAM 2在 Ultralytics 软件包中运行，会通过该函数在 Hiera 图像编码器和记忆注意力模块中计算注意力，因此只要硬件和输入精度允许，PyTorch 就会将这些层路由到 FlashAttention。

Flash Attention 与相关术语的区别#

• 标准注意力：通过将完整的分数矩阵和概率矩阵写入 HBM 来计算相同的数学运算。
• 稀疏注意力：只计算选定的查询—键对的分数。FlashAttention 论文还介绍了块稀疏变体，可跳过空分块。
• 线性注意力：用特征映射替代 Softmax，使计算成本随序列长度线性增长，但会改变模型输出。
• PagedAttention：一种内存管理器，在服务过程中将 KV 缓存存储到固定大小的页中。注意力内核（包括 FlashAttention）会从这些页中读取键和值。
• 环形注意力：将超长序列拆分到多个设备上，每个设备在其分片上计算分块注意力，同时键值分块沿环形结构传递。

对训练和长上下文的影响#

由于内存随序列长度线性增长，FlashAttention 让训练长上下文窗口变得切实可行。原论文报告称，与 MLPerf 1.1 训练纪录相比，序列长度为 512 的BERT-large端到端速度提升了 15%；序列长度为 1,024 的GPT-2速度提升了 3 倍；在 Long Range Arena 基准上则提升了 2.4 倍。同样的节省也适用于视觉领域，视觉 Transformer（ViT）或扩散 Transformer会将高分辨率图像和视频转换成长词元序列。更长的序列也提高了准确率：论文报告称，首次有 Transformer 在 Path-X 上超越随机猜测；该任务要求判断展平为包含 16,384 个像素的序列后，一张 128 × 128 图像中的两个点是否由路径连接。