""" ============================================================================= 第 5 讲 · 代码 5：FlashAttention —— 用分块 + 在线 softmax 精确计算注意力 ============================================================================= 对应讲义：lectures/lecture05_gpus_tpus.md 第 6 节 ----------------------------------------------------------------------------- 【这个文件要解决什么问题？】 ----------------------------------------------------------------------------- 标准注意力的实现： S = Q K^T / √d ← 写一个 n×n 矩阵到显存 P = softmax(S) ← 读 S，再写一个 n×n 矩阵 O = P V ← 读 P n = 32K 时，一个头的 S 就有 10 亿个元素（BF16 下 2 GB）。显存装不下，搬运也极慢。 FlashAttention（Dao et al. 2022）把前面学到的技巧组合起来： - **分块**：把 Q、K、V 切成小块，每次只把一小块搬进片上 SRAM（04 号文件）； - **在线 softmax**：不需要先看完整行，就能一块一块地算出 softmax（本文件第 1 部分）； - **融合**：QK^T、softmax、乘 V 在一个核函数里完成，n×n 矩阵从不写回显存（03 号文件）； - **重计算**：反向传播时不读存好的 P，而是重新算出来（03 号文件）。 结果是**精确**的注意力（不是近似），显存从 O(n²) 降到 O(n)，速度也快得多。 这个文件用纯 PyTorch 实现 FlashAttention 的前向和反向（与 CS336 作业 2 中 "PyTorch 版 FlashAttention-2" 的思路一致）， 并与标准注意力、PyTorch 自动求导逐一对比。它的目的是讲清算法，不是追求速度 （Python 循环很慢；真正的加速要靠第 6 讲的 Triton 核函数）。 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python 05_flash_attention.py """ import math import torch import torch.nn.functional as F def print_header(title): """打印一个带分隔线的标题。""" print() print("=" * 70) print(f" {title}") print("=" * 70) # ============================================================================= # 第 1 部分：在线 softmax # ============================================================================= def online_softmax_stats(x, block=4): """ 一块一块地读入 x，维护两个"运行统计量"： m = 目前见过的最大值 l = Σ exp(x_i - m) （以当前最大值为基准的指数和） 【关键一步】读入新的一块后，最大值可能变大（m → m_new），之前累加的指数和要"换基准"： l ← l × exp(m - m_new) + Σ_{新块} exp(x_i - m_new) 这样一趟扫描结束后，softmax(x)_i = exp(x_i - m) / l。 （标准做法需要三趟：求最大值、求指数和、归一化；在线算法把前两趟合成一趟。） 【参数】x : 一维张量；block : 每次读入多少个数 【返回】(m, l) """ m = torch.tensor(float("-inf")) l = torch.tensor(0.0) for s in range(0, len(x), block): chunk = x[s:s + block] m_new = torch.maximum(m, chunk.max()) l = l * torch.exp(m - m_new) + torch.exp(chunk - m_new).sum() m = m_new return m, l def demo_online_softmax(): """演示 1：在线 softmax 与标准 softmax 结果一致。""" print_header("演示 1：在线 softmax —— 一趟扫描，不需要先看完整行") torch.manual_seed(0) x = torch.randn(10) * 5 m, l = online_softmax_stats(x, block=3) online = torch.exp(x - m) / l print(f" x = {[round(v, 2) for v in x.tolist()]}") print(f" 在线算法：m = {m.item():.3f}，l = {l.item():.4f}") print(f" 与 torch.softmax 的最大差异：{(online - torch.softmax(x, 0)).abs().max().item():.1e}") print(f" logsumexp = m + log(l) = {(m + torch.log(l)).item():.5f}，torch.logsumexp = {torch.logsumexp(x, 0).item():.5f}") # ============================================================================= # 第 2 部分：FlashAttention 前向 # ============================================================================= def standard_attention(Q, K, V, causal=False): """标准注意力（会生成完整的 n×n 矩阵），作为参照。形状：Q (n, d)，K、V (m, d)。""" S = Q @ K.T / math.sqrt(Q.shape[-1]) if causal: n, m = S.shape S = S.masked_fill(torch.ones(n, m, dtype=torch.bool).triu(1), float("-inf")) return torch.softmax(S, dim=-1) @ V def flash_attention_forward(Q, K, V, causal=False, block_q=16, block_k=16): """ FlashAttention 前向（纯 PyTorch 模拟，一个头）。 【算法】对每个 query 块 Q_i（在 GPU 上，每个 query 块由一个线程块负责，数据留在 SRAM）： O_i = 0，m_i = -∞，l_i = 0 （每行一组运行统计量） 对每个 key/value 块 (K_j, V_j)： S_ij = Q_i K_j^T / √d （块大小 block_q × block_k，只在 SRAM 里） m_new = max(m_i, rowmax(S_ij)) P_ij = exp(S_ij - m_new) l_i = exp(m_i - m_new) × l_i + rowsum(P_ij) O_i = exp(m_i - m_new) × O_i + P_ij V_j （之前累加的结果也要"换基准"） m_i = m_new O_i = O_i / l_i （最后统一除以分母） L_i = m_i + log(l_i) （logsumexp，保存下来给反向传播用） 【因果掩码】位置 q 只能看 k ≤ q。完全在对角线右上方的 key 块直接跳过（省一半计算）。 【参数】 Q : (n, d)；K、V : (m, d) block_q, block_k : 分块大小 【返回】(O, L)：O 的形状 (n, d)，L 的形状 (n,) 【边界】必须按 key 块从左到右处理：因果情况下第 0 个 key 块对每个 query 都可见， 保证第一次更新时 m 就是有限值，不会出现 (-∞) - (-∞) = nan。 """ n, d = Q.shape m_total = K.shape[0] scale = 1.0 / math.sqrt(d) O = torch.zeros(n, d, dtype=Q.dtype) L = torch.zeros(n, dtype=Q.dtype) for qs in range(0, n, block_q): qe = min(qs + block_q, n) Qi = Q[qs:qe] Oi = torch.zeros(qe - qs, d, dtype=Q.dtype) mi = torch.full((qe - qs,), float("-inf"), dtype=Q.dtype) li = torch.zeros(qe - qs, dtype=Q.dtype) for ks in range(0, m_total, block_k): if causal and ks > qe - 1: break # 这个 key 块整个在"未来"，跳过 ke = min(ks + block_k, m_total) S = Qi @ K[ks:ke].T * scale if causal: q_idx = torch.arange(qs, qe)[:, None] k_idx = torch.arange(ks, ke)[None, :] S = S.masked_fill(k_idx > q_idx, float("-inf")) m_new = torch.maximum(mi, S.max(dim=-1).values) P = torch.exp(S - m_new[:, None]) correction = torch.exp(mi - m_new) li = correction * li + P.sum(dim=-1) Oi = correction[:, None] * Oi + P @ V[ks:ke] mi = m_new O[qs:qe] = Oi / li[:, None] L[qs:qe] = mi + torch.log(li) return O, L def demo_forward(): """演示 2：FlashAttention 前向与标准注意力结果一致。""" print_header("演示 2：FlashAttention 前向 —— 精确，不是近似") torch.manual_seed(0) n, d = 100, 32 # 故意让 n 不是块大小的整数倍 Q, K, V = (torch.randn(n, d, dtype=torch.float64) for _ in range(3)) for causal in [False, True]: O, L = flash_attention_forward(Q, K, V, causal=causal, block_q=16, block_k=24) ref = standard_attention(Q, K, V, causal=causal) ref2 = F.scaled_dot_product_attention(Q[None], K[None], V[None], is_causal=causal)[0] print(f" causal={causal!s:<5}：与标准实现的最大差异 {(O - ref).abs().max().item():.1e}，" f"与 F.scaled_dot_product_attention {(O - ref2).abs().max().item():.1e}") # ============================================================================= # 第 3 部分：FlashAttention 反向（重计算） # ============================================================================= def flash_attention_backward(Q, K, V, O, L, dO, causal=False, block_q=16, block_k=16): """ FlashAttention 反向（FlashAttention-2 的写法）：不存 n×n 的 P，而是用保存的 L 逐块重算。 【推导要点】（S = QK^T/√d，P = softmax(S)，O = PV） dV = P^T dO dP = dO V^T dS = P ⊙ (dP - D)，其中 D_i = Σ_j P_ij dP_ij = rowsum(dO ⊙ O) ← softmax 的求导 dQ = dS K / √d，dK = dS^T Q / √d 关键技巧： 1. P_ij 可以由保存的 logsumexp 重算：P_ij = exp(S_ij - L_i)，不需要 m、l 两个量； 2. D_i = rowsum(dO_i ⊙ O_i) 只要 O 和 dO 就能算，不需要完整的 P。 所以反向只需要 Q、K、V、O、L（都是 O(n) 的大小）和 dO。 【返回】(dQ, dK, dV) """ n, d = Q.shape scale = 1.0 / math.sqrt(d) D = (dO * O).sum(dim=-1) # (n,) dQ, dK, dV = torch.zeros_like(Q), torch.zeros_like(K), torch.zeros_like(V) for ks in range(0, K.shape[0], block_k): # 外层循环 key 块（FlashAttention-2 的顺序） ke = min(ks + block_k, K.shape[0]) Kj, Vj = K[ks:ke], V[ks:ke] for qs in range(0, n, block_q): qe = min(qs + block_q, n) if causal and ks > qe - 1: continue S = Q[qs:qe] @ Kj.T * scale if causal: q_idx = torch.arange(qs, qe)[:, None] k_idx = torch.arange(ks, ke)[None, :] S = S.masked_fill(k_idx > q_idx, float("-inf")) P = torch.exp(S - L[qs:qe, None]) # 重算注意力权重 dV[ks:ke] += P.T @ dO[qs:qe] dP = dO[qs:qe] @ Vj.T dS = P * (dP - D[qs:qe, None]) dQ[qs:qe] += dS @ Kj * scale dK[ks:ke] += dS.T @ Q[qs:qe] * scale return dQ, dK, dV class FlashAttention(torch.autograd.Function): """ 把前向和反向包装成一个 autograd.Function，这样就能像普通 PyTorch 运算一样用 .backward()。 前向只保存 Q、K、V、O、L —— 没有任何 n×n 的张量。 """ @staticmethod def forward(ctx, Q, K, V, causal=False): O, L = flash_attention_forward(Q, K, V, causal) ctx.save_for_backward(Q, K, V, O, L) ctx.causal = causal return O @staticmethod def backward(ctx, dO): Q, K, V, O, L = ctx.saved_tensors dQ, dK, dV = flash_attention_backward(Q, K, V, O, L, dO, ctx.causal) return dQ, dK, dV, None # causal 不是张量，不需要梯度 def demo_backward(): """演示 3：FlashAttention 反向与 PyTorch 自动求导结果一致。""" print_header("演示 3：FlashAttention 反向 —— 重算，不存 n×n 矩阵") torch.manual_seed(0) n, d = 70, 16 for causal in [False, True]: Q, K, V = (torch.randn(n, d, dtype=torch.float64, requires_grad=True) for _ in range(3)) Q2, K2, V2 = (t.detach().clone().requires_grad_(True) for t in (Q, K, V)) FlashAttention.apply(Q, K, V, causal).pow(2).sum().backward() standard_attention(Q2, K2, V2, causal).pow(2).sum().backward() diffs = [(a.grad - b.grad).abs().max().item() for a, b in [(Q, Q2), (K, K2), (V, V2)]] print(f" causal={causal!s:<5}：dQ、dK、dV 与自动求导的最大差异 {max(diffs):.1e}") # ============================================================================= # 第 4 部分：省了多少？ # ============================================================================= def demo_memory(): """演示 4：显存占用与显存访问量的对比。""" print_header("演示 4：省了多少显存和搬运？（一个头，BF16，d = 128）") d = 128 sram_elems = 100 * 1024 // 2 # 假设可用约 100 KB 的 SRAM（BF16 元素个数） print(f" {'序列长度 n':>10}{'标准：S 和 P 占用':>18}{'Flash：额外占用':>16}" f"{'标准：访问量(百万元素)':>22}{'Flash：访问量(百万元素)':>24}") for n in [1024, 4096, 16384, 65536]: std_mem = 2 * n * n * 2 # S 和 P 两个 n×n 的 BF16 矩阵 flash_mem = n * 4 # 只多存一个 FP32 的 logsumexp 向量 std_io = 4 * n * n # 写 S、读 S、写 P、读 P（数量级估计） flash_io = n * n * d * d / sram_elems # 论文中的 IO 复杂度 Θ(n²d²/M) print(f" {n:>10}{std_mem / 1e9:>15.3f} GB{flash_mem / 1e6:>13.3f} MB" f"{std_io / 1e6:>22,.0f}{flash_io / 1e6:>24,.0f}") print(""" - 显存：标准实现要存两个 n×n 矩阵（n = 64K 时一个头就要 17 GB）；FlashAttention 只多存一个长度为 n 的向量。 这也让训练时不需要为注意力矩阵保存激活值（第 3 讲 11.3 节：GPT-2 XL 一半以上的激活都是这两个矩阵）。 - 显存访问：标准实现要把 n×n 矩阵写、读各两次（约 4n²）；FlashAttention 的访问量约为 n²d²/M （M 是 SRAM 能放下的元素个数，这里 d²/M ≈ 0.32），按这个估算约为标准实现的 1/12。 （这是论文中的渐近分析，常数因子不精确；实际加速取决于硬件和实现。 FlashAttention 论文报告 GPT-2 的端到端训练约快 3 倍。） - 计算量（FLOPs）并没有减少，反向时还因为重算而增加了 —— 但注意力是内存受限的，少搬数据才是关键。""") def main(): demo_online_softmax() demo_forward() demo_backward() demo_memory() print("\n全部演示运行完毕。") if __name__ == "__main__": main()