""" ============================================================================= 第 8 讲 · 代码 4：IO 记账——标准注意力 vs FlashAttention 读写了多少显存 ============================================================================= 对应讲义：lectures/lecture08_gpu_flashattention.md 第 5.5、5.6、6 节 ----------------------------------------------------------------------------- 【这个文件要回答的问题 What this file answers】 ----------------------------------------------------------------------------- FlashAttention 论文的核心论断是 IO 复杂度：标准注意力读写显存 Θ(Nd + N²) 次，FlashAttention 是 Θ(N²d²/M)。 这个文件把它变成具体的数字。 演示 1：公式 vs 实测计数：用 flash_tools 里带计数器的算法实际跑一遍，读写次数与公式逐个相等。 演示 2：随序列长度 N 变化：两种算法前向 + 反向的读写量（元素个数），以及倍数。 演示 3：随片上存储 M 变化：M 越大块越大，读写越少——Θ(N²d²/M) 的"1/M"。 演示 4：roofline 估计：GPT-2 medium 的注意力（N = 1024）在 A100 上的理论用时，并和论文图 2 的实测对照。 演示 5：显存占用：标准注意力反向要存 N×N 的 P；N 到几万时，单是它就放不进一张 80 GB 的卡。 演示 6：块稀疏（block-sparse）：只算一部分块，读写量按非零块的比例下降（因果掩码是一个天然的例子）。 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python3 04_io_accounting.py # 约 3 秒；需要 numpy 【Python 小知识】 - 2 ** 20 是 1,048,576（1 Mi）；1 GiB = 2 ** 30 字节。显存容量常用 GiB（二进制），带宽常用 GB/s（十进制）。 - f"{x:,}"：给整数加千位分隔符。 """ import math import os import sys import numpy as np HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, HERE) sys.path.insert(0, os.path.join(HERE, "..", "lecture01")) from console import pad, print_header # noqa: E402 from flash_tools import (HARDWARE, IOCounter, attention_flops, flash_attention_backward, # noqa: E402 flash_attention_forward, flash_block_sizes, flash_io, roofline_time, standard_attention, standard_io) A100 = HARDWARE["A100 40GB (2020)"] BYTES = 2 M_A100 = A100["sram_kb"] * 1024 // BYTES # 一个 SM 的 SRAM 放得下多少个 16 位数：98,304 def total(io): """前向 + 反向的读写元素个数。""" return io["forward"] + io["backward"] # ============================================================================= # 演示 1：公式 vs 实测计数 # ============================================================================= def demo1_check(rng): print_header("演示 1：读写次数的公式与带计数器的实际算法逐个相等") print(f"{pad('N', 6, '>')}{pad('d', 5, '>')}{pad('(Br, Bc)', 12, '>')}{pad('前向：实测', 14, '>')}{pad('公式', 12, '>')}" f"{pad('反向：实测', 14, '>')}{pad('公式', 12, '>')}") for N, d, Br, Bc in [(64, 16, 16, 16), (100, 8, 7, 33), (256, 32, 32, 128), (512, 64, 64, 384)]: Q, K, V, dO = (rng.normal(size=(N, d)) for _ in range(4)) cf, cb = IOCounter(), IOCounter() O, l, m = flash_attention_forward(Q, K, V, Br, Bc, counter=cf) flash_attention_backward(Q, K, V, O, dO, l, m, Br, Bc, counter=cb) io = flash_io(N, d, Br, Bc) assert cf.total == io["forward"] and cb.total == io["backward"] fwd, bwd = f"{io['forward']:,}", f"{io['backward']:,}" # 先算成字符串，避免 f-string 里嵌套同种引号 print(f"{pad(N, 6, '>')}{pad(d, 5, '>')}{pad(f'({Br}, {Bc})', 12, '>')}{pad(f'{cf.total:,}', 14, '>')}" f"{pad(fwd, 12, '>')}{pad(f'{cb.total:,}', 14, '>')}{pad(bwd, 12, '>')}") Br, Bc = flash_block_sizes(M_A100, 64) print(f"\nA100 的一个 SM 有 {A100['sram_kb']} KB SRAM，放得下 M = {M_A100:,} 个 16 位数；" f"d = 64 时论文的块大小公式给出 Br = {Br}，Bc = {Bc}") print("\n观察 Observations：") print(" 1. 对每一种 N、d、块大小（包括除不尽的情况），实际算法的读写计数和 flash_io 的公式完全相等——") print(" 后面几个演示可以放心地只用公式，算到几万的序列长度。") # ============================================================================= # 演示 2：随 N 变化 # ============================================================================= def demo2_vs_N(): print_header("演示 2：一个头前向 + 反向读写显存的元素个数（d = 64，M = 98,304）") d = 64 print(f"{pad('N', 8, '>')}{pad('标准 11N² + 12Nd', 20, '>')}{pad('FlashAttention', 18, '>')}{pad('倍数', 8, '>')}" f"{pad('Tc（K/V 被扫几遍）', 20, '>')}{pad('32N²d²/M（主项）', 18, '>')}") for N in [512, 1024, 2048, 4096, 8192, 16384, 65536]: Br, Bc = flash_block_sizes(M_A100, d, N) s, f = total(standard_io(N, d)), total(flash_io(N, d, Br, Bc)) lead = 32 * N * N * d * d / M_A100 # 每扫一遍 K/V：前向 3Nd + 反向 5Nd；Tc ≈ 4Nd/M 遍 print(f"{pad(N, 8, '>')}{pad(f'{s:,}', 20, '>')}{pad(f'{f:,}', 18, '>')}{pad(f'{s / f:.1f}', 8, '>')}" f"{pad(math.ceil(N / Bc), 20, '>')}{pad(f'{lead:,.0f}', 18, '>')}") print("\n观察 Observations：") print(" 1. 标准注意力的读写量由 N² 主导（前向 4N²、反向 7N²）；FlashAttention 也随 N² 增长：") print(" 每扫一遍 K/V 要读写约 8Nd 个数（前向 3Nd、反向 5Nd），一共扫 Tc ≈ 4Nd/M 遍 → 32N²d²/M。") print(" d = 64、M = 98,304 时系数是 32 × 64² / 98,304 ≈ 1.33，标准算法是 11——差的那个倍数，正是由 d²/M 决定的。") print(" 2. N = 512 时只差 4.1 倍（K/V 只扫 2 遍，O(Nd) 的项还占不小的比例）；N 越大，倍数越接近 11 ÷ 1.33 ≈ 8.3，") print(" N = 65,536 时 8.1 倍。") print(" 3. 'FlashAttention 让注意力的显存读写从二次方变成线性'是常见的误解：读写仍是 N² 量级，") print(" 变成线性的是**显存占用**（演示 5）——不用存 N×N 矩阵。") # ============================================================================= # 演示 3：随 M 变化 # ============================================================================= def demo3_vs_M(): print_header("演示 3：片上存储 M 越大，块越大，读写越少（N = 4096，d = 64）") N, d = 4096, 64 s = total(standard_io(N, d)) print(f"{pad('M（元素）', 12, '>')}{pad('相当于', 12, '>')}{pad('(Br, Bc)', 14, '>')}{pad('FlashAttention 读写', 22, '>')}" f"{pad('标准 ÷ Flash', 14, '>')}") for p in range(10, 20): M = 2 ** p Br, Bc = flash_block_sizes(M, d, N) f = total(flash_io(N, d, Br, Bc)) kb = f"{M * BYTES / 1024:.0f} KB" print(f"{pad(f'{M:,}', 12, '>')}{pad(kb, 12, '>')}{pad(f'({Br}, {Bc})', 14, '>')}{pad(f'{f:,}', 22, '>')}" f"{pad(f'{s / f:.1f}', 14, '>')}") print(f"（标准注意力不分块，与 M 无关：{s:,}）") print("\n观察 Observations：") print(" 1. M 每翻一倍，Bc 翻一倍，K/V 少扫一半遍数，读写量近似减半——这就是 Θ(N²d²/M) 里的 1/M。") print(" 2. M ≤ 16 KB（块 ≤ 32 行）时，FlashAttention 反而比标准注意力读写**更多**（K/V 要扫上百遍）：") print(" 分块的好处完全取决于片上存储够不够大。A100 每个 SM 192 KB，介于表里 128 KB 和 256 KB 两档之间，") print(" 读写量比标准注意力少 5–10 倍。") print(" 3. M 大到能一次放下整个 K、V（Bc = N）时，K/V 只读一遍，曲线到底——论文的下界（命题 3）说的就是：") print(" 在所有 M 的范围内，没有一个精确注意力算法能在渐近意义上做得比 Θ(N²d²/M) 更好。") # ============================================================================= # 演示 4：roofline 估计 vs 论文图 2 # ============================================================================= def demo4_roofline(): print_header("演示 4：GPT-2 medium 的注意力（N = 1024，d = 64，16 头，batch 8）在 A100 上的理论用时（前向 + 反向）") N, d, heads = 1024, 64, 8 * 16 Br, Bc = flash_block_sizes(M_A100, d, N) # 标准注意力有两种记账：论文算法 0/3 的最简实现；以及更接近 PyTorch 朴素实现的写法—— # 缩放、掩码、dropout 各是单独的 kernel：前向 3 趟 × 读写 2N²，反向 dropout、缩放各 2N² eager_extra = 3 * 2 * N * N + 2 * 2 * N * N rows = [("标准（算法 0 + 3，最简）", attention_flops(N, d), BYTES * total(standard_io(N, d))), ("标准（朴素 PyTorch：+缩放、掩码、dropout）", attention_flops(N, d), BYTES * (total(standard_io(N, d)) + eager_extra)), ("FlashAttention", attention_flops(N, d, flash=True), BYTES * total(flash_io(N, d, Br, Bc)))] print(f"{pad('实现', 44)}{pad('GFLOPs', 10, '>')}{pad('读写 GB', 10, '>')}{pad('理论下界', 12, '>')}{pad('受限于', 8, '>')}") res = [] for name, f, b in rows: r = roofline_time(f * heads, b * heads, A100) res.append((f, b, r["t"])) t_str = f"{r['t'] * 1e3:.2f} ms" print(f"{pad(name, 44)}{pad(f'{f * heads / 1e9:.1f}', 10, '>')}{pad(f'{b * heads / 1e9:.2f}', 10, '>')}" f"{pad(t_str, 12, '>')}{pad('算力' if r['bound'] == 'compute' else '带宽', 8, '>')}") (f0, b0, t0), (_, b1, t1), (f2, b2, t2) = res print(f"\n{pad('比值（标准 ÷ Flash）', 30)}{pad('读写量', 16, '>')}{pad('FLOPs', 10, '>')}{pad('理论用时', 12, '>')}") print(f"{pad('最简标准', 30)}{pad(f'{b0 / b2:.1f}', 16, '>')}{pad(f'{f0 / f2:.2f}', 10, '>')}{pad(f'{t0 / t2:.1f}', 12, '>')}") print(f"{pad('朴素 PyTorch 标准', 30)}{pad(f'{b1 / b2:.1f}', 16, '>')}{pad(f'{f0 / f2:.2f}', 10, '>')}{pad(f'{t1 / t2:.1f}', 12, '>')}") print(f"{pad('论文图 2 实测', 30)}{pad(f'{40.3 / 4.4:.1f}', 16, '>')}{pad(f'{66.6 / 75.2:.2f}', 10, '>')}{pad(f'{41.7 / 7.3:.1f}', 12, '>')}") print("（论文图 2 左：A100，GPT-2 medium，N = 1024；标准 66.6 GFLOPs、40.3 GB、41.7 ms，FlashAttention 75.2 GFLOPs、4.4 GB、7.3 ms）") print("\n观察 Observations：") print(" 1. 论文没有写这张图的 batch 和实现细节，所以这里只比较**比值**：读写量少 5.8–10.8 倍（论文 9.2 倍），") print(" FLOPs 多 17%（论文多 13%，比值 0.86 vs 0.89），理论用时少 5.3–9.7 倍（论文实测 5.7 倍）——") print(" 读写量和用时的比值都落在我们两种记账之间，FLOPs 的比值也接近。") print(" 2. 标准注意力受带宽限制，FlashAttention 受算力限制：省下来的是搬数据的时间，多花的那点计算几乎不影响用时。") # ============================================================================= # 演示 5：显存占用 # ============================================================================= def demo5_memory(): print_header("演示 5：为了反向传播要存下的东西（batch 1，16 个头，16 位）") H, d = 16, 64 print(f"{pad('N', 8, '>')}{pad('标准：存 P（N×N）', 20, '>')}{pad('Flash：存 m、ℓ', 18, '>')}{pad('一张 80 GB 的卡放得下吗', 26, '>')}") for N in [1024, 2048, 8192, 16384, 32768, 65536, 131072]: p_bytes = H * N * N * BYTES stat_bytes = H * N * 2 * 4 # m、ℓ 每行两个 fp32 数 fits = "放得下" if p_bytes < 80e9 else "放不下（只算 P 就超了）" print(f"{pad(N, 8, '>')}{pad(f'{p_bytes / 2 ** 30:,.2f} GiB', 20, '>')}{pad(f'{stat_bytes / 2 ** 20:.2f} MiB', 18, '>')}" f"{pad(fits, 26, '>')}") print("（每一层都要存一份；一个 24 层的模型要再乘以 24。Q、K、V、O 本身是 N×d，两种算法都要存。）") print("\n观察 Observations：") print(" 1. P 的大小随 N² 增长：N = 16K 时一层就要 8 GiB，N = 64K 时 128 GiB——一张卡放不下一层。") print(" 2022 年以前，'上下文最多 2K 个 token'很大程度上是这个限制。") print(" 2. FlashAttention 只存每行两个数，显存占用随 N **线性**增长——这才是'线性'的那一半。") print(" 论文因此第一次在 Path-X（16K）和 Path-256（64K）上训练出比随机猜测好的 Transformer（61.4%、63.1%）。") # ============================================================================= # 演示 6：块稀疏 # ============================================================================= def demo6_block_sparse(rng): print_header("演示 6：只算一部分块——读写量按非零块的比例下降（N = 1024，d = 32，块 64）") N, d, Br, Bc = 1024, 32, 64, 64 Q, K, V = (rng.normal(size=(N, d)) for _ in range(3)) c_full, c_causal = IOCounter(), IOCounter() O_full, _, _ = flash_attention_forward(Q, K, V, Br, Bc, causal=False, counter=c_full) O_causal, _, _ = flash_attention_forward(Q, K, V, Br, Bc, causal=True, counter=c_causal) err = np.abs(O_causal - standard_attention(Q, K, V, causal=True)[0]).max() n_blk = N // Br visible = n_blk * (n_blk + 1) // 2 print(f"块的总数 {n_blk} × {n_blk} = {n_blk * n_blk}；因果掩码下需要算的块 {visible}（比例 {visible / n_blk ** 2:.3f}）") print(f"前向读写：全部块 {c_full.total:,}；因果（跳过整块被遮住的块）{c_causal.total:,}，比例 {c_causal.total / c_full.total:.3f}" f"；结果与标准因果注意力最大差 {err:.1e}") print("\n观察 Observations：") print(" 1. 因果掩码让一半的块完全不用算、不用读——读写量和计算量都降到约一半多一点（对角线上的块还得算）。") print(" 2. 这就是论文的块稀疏 FlashAttention（命题 4：IO 复杂度乘以非零块的比例 s）的最简单情形；") print(" 更一般的稀疏模式（局部窗口 + 少数全局块）是近似注意力；论文报告块稀疏 FlashAttention 比当时已有的近似注意力方法都快。") if __name__ == "__main__": rng = np.random.default_rng(0) demo1_check(rng) demo2_vs_N() demo3_vs_M() demo4_roofline() demo5_memory() demo6_block_sparse(rng)