FlashAttention：为什么少搬数据，反而算得更快

01 · 先记住一句话

FlashAttention 的核心是重组 Attention 的数据流：分数只算一小块，立刻把这块用于更新输出，然后腾出空间处理下一块。它避免把整张巨大的分数表和概率表写进显存，再读出来。

你可以先理解“少搬什么数据”，再理解“为什么分块后结果仍然正确”。下面先讲前向，训练反向的重计算单独放在后面。

02 · 原来的 Attention，慢在哪里

标准 Attention 可以拆成三个动作。为了方便阅读，N 表示序列位置数，S 专门表示分数矩阵。下面的直观说明暂时省略缩放和因果掩码；它们在实际计算中仍然存在。

步骤	做什么	产生什么
S = QKᵀ	每个查询 Q 给各位置的 K 打分	N×N 的分数表 S。
P = softmax(S)	把每行分数变成总和为 1 的相对权重	N×N 的权重表 P。
O = PV	按这些权重汇总 V	每个查询的输出向量 O。

如果三个动作分别执行，前一步写出的大表要让后一步再读回来。序列长时，N×N 的表很快超过片上小工作台的容量，于是大量中间数据在 HBM 和计算单元附近来回搬运。

完整中间表的执行路径

读 Q/K→算 S→写 S 到 HBM→再读 S→算 P→写 P 再读 P→乘 V 得 O

FlashAttention 的分块路径

读一块 Q/K/V→算小块分数→在片上更新统计与输出累积→继续下一块

没有完整 S/P 的写出和读回；输入读取与最终输出保存仍然存在。

问题不只是“算得够不够快”：计算单元可能在等数据。FlashAttention 首先减少这种中间表的读写。这里说的是会完整写出中间矩阵的基线，不能据此断定所有 Attention、所有形状都只受内存带宽限制。

依据：FlashAttention 原始论文 ↗ · PyTorch · SDPA 后端与约束（main） ↗

03 · 分块：大表不用整张保存

固定一小块 Q，依次取来对应的 K/V 小块。计算当前的分数块后，立即把它用于累积这一小块 Q 的输出。当前分数块用完，就可以被下一块覆盖。

取一小块 Q 和一小块 K，得到一小块分数。
读取这一块对应的 V，把贡献加到输出累积量。
更新必要的归一化统计，然后处理下一个 K/V 块。
遍历完需要关注的 K/V 后，输出这一块 Q 的最终结果。

遍历所有 Q 块后，得到完整输出 O。要关注的位置仍然都被处理；变化在于中间结果只按块暂存，不会形成一张完整驻留在 HBM 中的 S 或 P。

04 · 在线 Softmax：没看到全部分数，怎么先累加

这是最容易卡住的一步。Softmax 的分母取决于整行，如果先对每块独立归一化，再把块输出直接平均，通常会算错。FlashAttention 会保留足够的统计量，等全部块处理完再完成全行归一化。

先看一个能手算的例子。假设一行分数取指数之后得到 [1,2,4,3]，对应 V 简化成四个数 [10,20,30,40]。完整计算的权重是 [0.1,0.2,0.4,0.3]，输出为 29。下面把它分成两块，每块两个位置。

过程	权重按当前最大值改尺度	指数和 ℓ	加权累积 U
第1块 [1,2]	都除以 2 → [0.5,1]	1.5	0.5×10+1×20 = 25
第2块带来更大最大值 4	旧值再乘 2/4 = 0.5	旧 ℓ 变成 0.75	旧 U 变成 12.5
加入第2块 [4,3]	都除以 4 → [1,0.75]	0.75+1+0.75 = 2.5	12.5+30+30 = 72.5
最后归一化	统一尺度，统一分母	2.5	O=72.5/2.5 = 29

关键在第二步：最大值变化后，旧分子和旧分母一起改尺度，旧块之间的相对关系仍保留。新块也使用相同尺度，再把它们累加。没有丢掉任何贡献，所以最后得到相同的归一化结果。

逐块算一遍：输出为什么仍是 29

单个查询、两个块、标量 V 的手算例子；实际输出是向量。先处理第1块，再统一尺度，最后累加第2块。

实现里会为每个查询行记录 m（已见分数的最大值）、ℓ（改尺度后的指数和）、U（尚未除以 ℓ 的输出累积）。m 是分数的最大值，例子里的 2 和 4 是取指数后的最大强度。最大值从 ln 2 变成 ln 4 时，修正系数就是 exp(ln 2−ln 4)=1/2。

理解算例后，再展开对应公式

公式可以后看，先理解下面的例子

m_new = max(m_old, max(scores_block))
a = exp(m_old - m_new)
p_block = exp(scores_block - m_new)
l_new = a * l_old + sum(p_block)
u_new = a * u_old + p_block @ V_block

# 全部块处理完
O = u_final / l_final

真实 U 是向量或矩阵，每个输出分量都用同一个行尺度修正。初始状态为 m=−∞、ℓ=0、U=0；因果掩码、全掩码行和 Dropout 还需要正确处理。

“结果等价”指遵循同一个 Attention 数学定义，不代表不同浮点执行顺序逐 bit 完全一致。

依据：FlashAttention 原始论文 ↗ · FlashAttention · 原作者实现 ↗

05 · 训练反向：需要 P 时，重算当前块

训练的反向求导也要用到 Attention 权重 P。如果前向不保存完整 P，反向就不能直接读出那张大表。FlashAttention 选择再次计算需要的分数块，并利用前向的归一化统计重建对应的概率块，用完后继续下一块。

阶段	完整保存中间表的基线	FlashAttention 的思路
前向	生成并保存完整 S/P，供后续使用	按块处理，保存反向需要的输入、输出与少量行统计。
反向	读取完整保存的 P 来求导	重算当前需要的分数 / 概率块，再计算对应梯度。

为什么增加计算反而可能更快？因为重新算一个小块的代价，可能小于把大量中间数据写出去、再读回来的代价。是否更快依赖硬件与负载；这是计算量、访存量和执行效率的取舍。

依据：FlashAttention · 原作者实现 ↗ · FlashAttention 原始论文 ↗

06 · 序列变长，中间表有多大

N=4096 时，一个头的分数表有 4096×4096=16,777,216 个元素。假设每元素 2 字节，一张表约 32 MiB；B=1、32 个头合计约 1 GiB。若 S 与 P 都各自完整驻留，就约 2 GiB，仅这是两类中间数据。

把 N 翻倍到 8192，元素数变成 4 倍。FlashAttention 避免的是这类二次增长的完整中间表存储，并不减少 Dense Attention 要处理的所有位置关系。

在批量、头数和每头维度固定时，输入、输出和每行统计的存储随 N 近似线性增长，可以把 Attention 的这部分存储从 O(N²) 理解为 O(N)。这里讲的是存储；Dense Attention 的主要计算量仍随位置配对数二次增长，不能说计算复杂度也变成 O(N)。

调整序列长度，观察一张分数矩阵的容量

序列长度 S

B=1，H=32，2 字节 / 元素；只估计中间矩阵，不估计加速倍数。

这里计算的是教学基线的一张完整中间矩阵容量。真实 dtype、复用、掩码和工作区都会影响峰值；图中容量变化不用于预测加速倍数。

07 · FA1、FA2、FA3 各自继续解决什么

理解了前面的数据流，再看这三篇经典论文的演进。版本差别可以按“少搬数据 → 更好分工 → 更好重叠”来理解。

版本	继续处理的问题	主要做法
FA1	完整分数 / 概率矩阵导致大量 IO	分块、在线归一化与反向重计算。
FA2	省了 IO，但 GPU 工作划分仍可能低效	减少非矩阵运算，在序列维度增加并行，调整线程块 / 线程束分工，减少协作交换。
FA3	如何发挥 Hopper 的异步硬件能力	利用 WGMMA / TMA 的异步执行，重叠搬运、矩阵乘法与 Softmax；另有 FP8 路径。

线程束可以先理解为协同执行的一组 GPU 线程。FA2 不只是“多切几个 Q 块”：它同时调整线程块间与线程束内的工作划分，减少不必要的共享内存交换。

WGMMA 是 Hopper 上的线程束组矩阵乘加指令，TMA 是搬运 Tensor 数据的硬件机制。FA3 的生产者负责准备下一批数据，消费者负责处理就绪数据，让部分搬运和计算在时间上重叠。FP8 是另一种低精度取舍，需要数值误差控制。

FA2 论文在 A100 的报告中，FA1 约达到理论峰值 FLOPs/s 的 25–40%，FA2 约 50–73%。这是论文条件下的 Attention 运算效率，不是所有模型都适用的利用率，也不是整个训练任务的固定提速比例。FA3 的核心实现针对 Hopper，不可直接当作昇腾也支持相同指令的证据。

依据：FlashAttention-2 · 原始论文 ↗ · FlashAttention-3 · 原始论文 ↗ · Tri Dao · FlashAttention-3 作者讲解 ↗

08 · 理解原理后，怎样实际使用

通常先使用框架或推理引擎已有的 Attention 后端。SDPA 是缩放点积注意力的框架接口，不是保证选中 FlashAttention 的按钮；形状、布局、精度和设备会影响后端。

PyTorch 接入片段 · 假设 q/k/v 已准备好，实际后端由环境决定

import torch.nn.functional as F
# q,k,v: [batch, heads, sequence, head_dim]
out = F.scaled_dot_product_attention(
    q, k, v, dropout_p=0.0, is_causal=True
)

下方纯 Python 小例子验证“逐块累积后与完整 Softmax 输出一致”，不依赖 GPU。它帮助理解公式；实际加速内核还需要实现片上存储与线程调度。

依据：PyTorch · SDPA 后端与约束（main） ↗

09 · 在推理和训练中怎样判断收益

场景	看什么
长序列 Prefill / 训练	中间表存储、Attention 前后向时间与端到端内存。
小批次 Decode	查询通常只有一个位置，分数形状约为 1×历史长度，不能直接套用 N×N 的容量故事。
已经启用融合后端	先确认基线实际内核，避免比较同一条优化路径。
与分页 KV 组合	分页管理 KV 的存放，FlashAttention 主要优化计算中的中间 IO，二者解决不同问题。

测量前先预热，保证设备计时正确；比较输出误差，训练还比较梯度与 Loss。单个 Attention 提速不能直接等同于整个模型提速，其他层、通信与数据输入仍占时间。

10 · NVIDIA 与昇腾怎样放在同一框架里理解

先比较相同 Attention 数学操作和数据流，再看硬件实现。NVIDIA 可检查 CUDA / FlashAttention / TensorRT-LLM 后端；昇腾可检查 CANN 融合 Attention 算子与框架实现。分块和减少中间读写的思路可以跨平台采用，FA3 的特定硬件指令则有平台条件。

依据：FlashAttention · 原作者实现 ↗ · TensorRT-LLM · Attention ↗ · MindSpeed Core · 特性清单 ↗

11 · 读完应该能解释这四件事

FlashAttention 少搬的主要是哪两张中间表？
为什么各块独立做 Softmax 后直接平均可能算错？
最大值变化时，为什么分子和分母都要改尺度？
训练反向重计算增加了计算，为什么仍有可能缩短时间？

流程位置与关联方法

这个方法可能出现在多个位置；计算线包含模型计算与算子实现，执行线包含运行时、内存和通信。

I-A2 · Transformer 各层计算 T-A1 · 前向计算 T-A3 · 反向传播得到梯度

资料与核查范围

原理依据以原始论文和官方文档为主；算例与时间线是教学推导。文档核查日期：2026-10-10。引用 main / latest 文档时，实际接口以安装版本为准。

代码为接入片段或教学伪代码；本页没有声称完成 NVIDIA / 昇腾硬件实测。

FlashAttention：为什么少搬数据，反而算得更快