数据截至 (上游 commit 0251105a2fb1)
01 · 注意力为什么是 IO 瓶颈
这一章讲什么: 不讲任何 kernel 细节,只回答一个问题——标准注意力实现到底慢在哪。读完你会理解 FlashAttention 的每个设计决定为什么长那样。
1. 它要解决的小问题
Transformer 的注意力是 softmax(QKᵀ)·V。序列长度 N 翻倍,分数矩阵 QKᵀ 的元素数变 4 倍。
问题不在算不过来,而在这张 N×N 的表被反复写出显存、再读回来。先给结论:
标准实现的时间大头是 HBM(显存)读写,不是 FLOPs。优化它的正确姿势是少搬数据,而不是算得更快。
2. 先建立直觉:GPU 的两层存储差多远
GPU 不是一块均匀的大内存。它有两层常用存储,容量和速度差着数量级(数字为 A100 量级,出自 FA1 论文 arXiv:2205.14135,论文链接见 README.md:8):
| 层级 | 容量 | 带宽(量级) | 类比 |
|---|---|---|---|
| HBM(显存) | 40–80 GB | ~2 TB/s | 磁盘:大、慢 |
| SRAM(每 SM 的 shared memory) | 每 SM 192 KB,全卡 ~20 MB | ~19 TB/s(聚合) | 内存:小、快 |
| 寄存器 | 每 SM 256 KB | 比 SRAM 更快 | CPU 缓存 |
要点只有一条:SRAM 比 HBM 快约一个数量级,但小三个数量级。
所以 GPU 程序的性能秘诀从来不是「算得快」,而是「尽量让数据待在片上,少过 HBM 这座窄桥」。
3. 算术强度:判断一个算子「卡在哪」的尺子
算术强度 = FLOPs ÷ 从 HBM 读写的字节数。
- 强度高 → 每搬 1 字节要算很多次 → 算力先封顶(compute-bound), 矩阵乘大方阵属于这类。
- 强度低 → 每搬 1 字节只算几下 → 带宽先封顶(memory-bound),逐元素操作属于这类。
注意力的两个 matmul 本身强度不低,但中间夹着 softmax,而 softmax 要整行看完才能算。教科书实现于是拆成三个 kernel:
| 步骤 | 算什么 | HBM 往返(N×d 输入,N×N 中间量) |
|---|---|---|
| ① S = QKᵀ | matmul | 读 Q、K;写出 S(N²) |
| ② P = softmax(S) | 逐行归一化 | 读 S(N²);写出 P(N²) |
| ③ O = PV | matmul | 读 P(N²)、V;写 O |
三次 N² 级的显存读写,把 matmul 省下的时间全吃掉了。d(头维)通常是 64/128,不随 N 涨;N 一拉长,搬运量按 N² 涨,计算有用的部分却只跟着 N²d 里的常数因子走。
这就是「注意力是 memory-bound」的全部含义。
4. 图示:数据在两层存储间的往返
标准实现(三次过窄桥):
HBM ──读 Q,K──► SM 内算 S ──写 S──► HBM
HBM ──读 S────► SM 内 softmax ──写 P──► HBM
HBM ──读 P,V──► SM 内算 O ──写 O──► HBM
FlashAttention(一次过桥):
HBM ──读 Q,K,V──► SM 内融合算完(S/P 只在 SRAM 里活一瞬间)
│
└──写 O + LSE──► HBM
怎么读这张图: 上半部分每条横线都是一次 N² 字节的 HBM 传输;下半部分 S 和 P 从未离开 SM。kernel 数量从 3 变 1,HBM 流量从 O(N²) 降到 O(N·d)。
5. 原理演示:把「搬运」标注出来
下面这段是教科书注意力,注释标出每一趟 HBM 读写:
# 示意,非源码 —— 教科书注意力,标注 HBM 流量
def naive_attention(Q, K, V):
# Q, K, V: (N, d),已在 HBM
S = Q @ K.T # 读 Q,K (Nd);写 S (N²) ← 第 1 趟 N²
P = softmax(S, dim=-1) # 读 S (N²);写 P (N²) ← 第 2、3 趟 N²
O = P @ V # 读 P (N²)、V;写 O ← 第 4 趟 N²
return O
FlashAttention 的想法一句话:别让 S 和 P 离开片上。把 Q 切成行块、K/V 切成列块,在一个 kernel 里边算 S 的小块、边更新 O——这就是 tiling(分块):
# 示意,非源码 —— tiling 的骨架(正确性细节见第 02 章)
def flash_attention_skeleton(Q, K, V):
O = zeros_like(Q) # 输出累加器,在 SRAM/寄存器
for Q_tile in split_rows(Q): # 每个 Q 行块
for K_tile, V_tile in split_cols(K, V): # K/V 逐块流过
S_tile = Q_tile @ K_tile.T # S_tile 只活在 SRAM
O_tile_update(O, S_tile, V_tile) # 边算边更新,不落地
return O # 只有 O 写回 HBM
问题只剩一个:softmax 不是逐元素运算,分块之后行和、行最大值都看不全。这就是第 02 章要解决的事。
6. 真实实现:去代码里确认「S/P 不落地」
这个论断可以在 kernel 源码里直接核实——前向 kernel 全程只给 Q/K/V/O/LSE 建过 gmem tensor,没有任何 gmem 的 S/P:
compute_attn_1rowblock(csrc/flash_attn/src/flash_fwd_kernel.h:55)里,gmem 侧只有gQ、gK/gV、gO、gLSE;分数acc_s是partition_fragment_C分出来的寄存器 fragment,P 由它转 fp16 后同样在寄存器(csrc/flash_attn/src/flash_fwd_kernel.h:326-375一带)。- 唯一多出来的输出是
softmax_lse,每行一个 float32,O(N) 大小——分配在csrc/flash_attn/flash_api.cpp:464。 - 融合成一个 kernel 的发射点在
run_flash_fwd(csrc/flash_attn/src/flash_fwd_launch_template.h:91):kernel<<<grid, kNThreads, smem_size, stream>>>,一次 launch 算完全部。
至于省显存的效果,README 给了量化结论:序列 2K 时省 10 倍、4K 时省 20 倍,且省幅随序列长线性增长(README.md:515)。注意它的措辞——标准实现显存随 N² 涨,FlashAttention 随 N 线性涨。
7. 关键细节 / 坑
-
「精确」二字是本库的立身之本。 FlashAttention 不改数学,输出与教科书注意力在数值容差内一致;它和稀疏注意力、线性注意力(近似家族)完全不是一个物种。测试标准也按这个写:误差不超过 PyTorch 基线的两倍(
README.md:551)。 -
省显存还有第二个来源:反向重计算。 标准实现反向要读回 P(N²);FlashAttention 只存 O 和 LSE,反向时把 P 重算出来——「算一遍」比「存+读一遍」便宜,这本身就是 IO 权衡。详见第 04 章。
-
IO 复杂度是有下界证明的。 FA1 论文证明:SRAM 大小为 M 时,精确注意力的 HBM 访问量下界是 Ω(N²d²/M),FlashAttention 达到该量级;标准实现是 Θ(N²)(论文 arXiv:2205.14135 的定理部分,入口见
README.md:8)。这解释了为什么 tile 尺寸总是顶着 SRAM 容量选(见第 05 章的_get_block_size_n,flash_attn/flash_attn_interface.py:31)。 -
收益大小取决于硬件的「桥有多窄」。 README 明确说:显存带宽越低的 GPU,speedup 越大(
README.md:490附近)。换句话说,这是一个「IO 感知」的优化——它优化的是计算与搬运的相对关系,不是计算本身。
8. 本章小结
- 注意力慢在 N² 中间矩阵的 HBM 往返,不慢在 FLOPs。
- 对策:一个 kernel 融合 到底,S/P 只活在 SRAM。
- 代价:softmax 的分块正确性(第 02 章)、kernel 工程复杂度(第 03–05 章)。
- 这一切成立的前提是「精确」——所以它能无声地替换掉所有框架里的注意力。