""" ============================================================================= 第 8 讲 · 代码 3：在线 softmax 与 FlashAttention——换一种计算顺序，结果逐数字相同 ============================================================================= 对应讲义：lectures/lecture08_gpu_flashattention.md 第 5 节 ----------------------------------------------------------------------------- 【这个文件要回答的问题 What this file answers】 ----------------------------------------------------------------------------- FlashAttention 不是"近似"注意力：它和标准注意力算的是**同一个函数**，只是换了计算顺序， 让 N×N 的中间矩阵不必完整地存在。这里一步一步验证这件事。 演示 1：为什么 softmax 要"减去最大值"：朴素公式在 float64 下 exp(1000) 溢出，在 fp16 下 exp(12) 就溢出。 演示 2：在线 softmax：分块读一个向量，只记两个数（当前最大值 m、当前分母 ℓ），逐步打印它们的变化。 演示 3：FlashAttention 前向（论文算法 1）与标准注意力对照：不同块大小、有无因果掩码，误差 1e-16 量级； 片上同时存在的数最多多少个，和 N×N 比一比。 演示 4：反向传播：不存 P、反向时重算——梯度与标准反向、与 PyTorch 自动求导都一致；代价是多一次矩阵乘法。 演示 5：把分块注意力装进第 1 讲的 base model：验证集 bit/字符不变（精确算法）； 对照一个真正的"近似注意力"（只看最近 w 个字符的滑动窗口）：省了计算，但模型变差了。 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python3 03_online_softmax_and_flash.py # 约 10 秒；需要 numpy、PyTorch 依赖本目录的 flash_tools.py，以及 code/lecture01 的 fm_world.py、tiny_gpt.py、console.py （第 1 讲的 base model 只读加载，并且先 copy.deepcopy 一份再改装，缓存本身不会被改动）。 【Python 小知识】 - types.MethodType(函数, 对象)：把一个普通函数"绑定"成某个对象的方法。演示 5 用它把 base model 里 每个注意力层的 forward 换成我们自己的版本——只换这一个对象，不影响类本身和别的模型。 - copy.deepcopy(model)：把整个模型（包括所有参数）完整复制一份。 """ import copy import math import os import sys import types import numpy as np import torch import torch.nn.functional as F 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 (IOCounter, attention_flops, flash_attention_backward, flash_attention_forward, # noqa: E402 online_softmax, softmax_naive, softmax_safe, standard_attention, standard_attention_backward, standard_io) from fm_world import ALPHABET, build_world # noqa: E402 from tiny_gpt import CharTokenizer, bits_per_char, load_or_train_base_model # noqa: E402 # ============================================================================= # 演示 1：朴素 softmax 为什么不能用 # ============================================================================= def demo1_overflow(): print_header("演示 1：朴素 softmax 的溢出") x = np.array([1000.0, 1001.0, 1002.0]) with np.errstate(over="ignore", invalid="ignore"): naive = softmax_naive(x) print(f"x = {x}") print(f" 朴素公式 exp(x)/Σexp(x)：{naive} ← exp(1000) 超过了 float64 的最大值 ≈ 1.8e308") print(f" 先减最大值： {np.round(softmax_safe(x), 6)}") for v in [11.0, 12.0]: with np.errstate(over="ignore"): e16 = np.exp(np.float16(v)) print(f" fp16 下 exp({v:.0f}) = {e16}（真实值 {math.exp(v):,.0f}；fp16 最大只能表示 65504）") print("\n观察 Observations：") print(" 1. 注意力分数 QKᵀ/√d 很容易到几十；在 16 位浮点数里，exp(12) 就已经溢出。") print(" 2. 所以所有实现都用'安全 softmax'：先找出这一行的最大值，减掉它再取 exp。") print(" 但'先找最大值'意味着必须先把一整行读完——这就是为什么标准实现要把 N×N 的分数矩阵存下来。") # ============================================================================= # 演示 2：在线 softmax # ============================================================================= def demo2_online(): print_header("演示 2：在线 softmax——每次读 3 个数，只记住 m（当前最大值）和 ℓ（当前分母）") x = np.array([2.0, -1.0, 3.0, 0.5, 6.0, 1.0, -2.0, 5.5]) out, steps = online_softmax(x, block=3, trace=True) print(f"x = {x.tolist()}") print(f"{pad('读到第几个数', 14)}{pad('新读入的块', 22)}{pad('m（最大值）', 14, '>')}{pad('ℓ（分母）', 14, '>')}") start = 0 for end, m, l in steps: blk = str(x[start:end].tolist()) print(f"{pad(end, 14)}{pad(blk, 22)}{pad(f'{m:.1f}', 14, '>')}{pad(f'{l:.6f}', 14, '>')}") start = end ref = softmax_safe(x) print(f"最后：softmax = exp(x − m) / ℓ = {np.round(out, 4).tolist()}") print(f"与'先读完整行再算'的安全 softmax 最大差 {np.abs(out - ref).max():.1e}") l_check = float(np.exp(x - x.max()).sum()) print(f"核对：Σ exp(x − 6) = {l_check:.6f}，与最后的 ℓ 相同") print("\n观察 Observations：") print(" 1. 读到第一块时 m = 3，ℓ 是相对 3 算的；读到第二块出现了 6，旧的 ℓ 乘以 exp(3 − 6) 换算到新基准，") print(" 再加上新块的贡献。第三块的最大值 5.5 < 6，m 不变，只累加。") print(" 2. 任何时候只需要两个数——不需要先读完整行。这让注意力可以一块一块地算：") print(" 读一块 K，算一小块分数，更新 m、ℓ 和输出，然后把这一小块分数扔掉。") # ============================================================================= # 演示 3：FlashAttention 前向 # ============================================================================= def demo3_forward(rng): print_header("演示 3：FlashAttention 前向（论文算法 1）与标准注意力对照（float64，N = 256，d = 32）") N, d = 256, 32 Q, K, V = (rng.normal(size=(N, d)) for _ in range(3)) print(f"{pad('块大小 (Br, Bc)', 18)}{pad('因果掩码', 10)}{pad('与标准注意力最大差', 22, '>')}" f"{pad('片上最多同时存在', 18, '>')}{pad('HBM 读写（元素）', 18, '>')}") for causal in (False, True): O_ref, _ = standard_attention(Q, K, V, causal) for Br, Bc in [(16, 16), (32, 64), (64, 32), (128, 128)]: c = IOCounter() O, _, _ = flash_attention_forward(Q, K, V, Br, Bc, causal, counter=c) err = f"{np.abs(O - O_ref).max():.1e}" print(f"{pad(f'({Br}, {Bc})', 18)}{pad('是' if causal else '否', 10)}{pad(err, 22, '>')}" f"{pad(f'{c.peak_sram:,}', 18, '>')}{pad(f'{c.total:,}', 18, '>')}") std = standard_io(N, d)["forward"] print(f"\n对照：N×N 的分数矩阵有 {N * N:,} 个数；标准注意力前向读写 4N² + 4Nd = {std:,} 个数") print("\n观察 Observations：") print(" 1. 所有块大小、有无掩码，结果与标准注意力的差都在 1e-16 量级（float64 的舍入误差）：") print(" FlashAttention 是**精确**算法，不是近似。") print(" 2. 片上同时存在的数最多两千多到三万多个（取决于块大小），而 N×N 矩阵有 65,536 个——N 越大差距越悬殊，") print(" 片上需要的存储只和块大小有关，和 N 无关。") print(" 3. 读写次数随块变大而减少：块越大，K、V 被重复读的遍数（Tc）越少。块大小受 SRAM 容量限制——") print(" 这就是论文 IO 复杂度 Θ(N²d²/M) 里的 M（代码 04 详细算）。因果掩码下整块被遮住的块直接跳过，读写更少。") # ============================================================================= # 演示 4：反向传播 # ============================================================================= def demo4_backward(rng): print_header("演示 4：反向传播——不存 P，从 Q、K、V 和 (m, ℓ) 重算") N, d, Br, Bc = 128, 16, 32, 32 Q, K, V, dO = (rng.normal(size=(N, d)) for _ in range(4)) for causal in (False, True): O_ref, P = standard_attention(Q, K, V, causal) g_std = standard_attention_backward(Q, K, V, P, dO) O, l, m = flash_attention_forward(Q, K, V, Br, Bc, causal) g_flash = flash_attention_backward(Q, K, V, O, dO, l, m, Br, Bc, causal) # PyTorch 自动求导作第三方裁判：把标准注意力写成 torch 运算，让 autograd 求梯度 tq, tk, tv = (torch.tensor(a, requires_grad=True) for a in (Q, K, V)) s = tq @ tk.T / math.sqrt(d) if causal: s = s.masked_fill(~torch.tril(torch.ones(N, N, dtype=torch.bool)), float("-inf")) out = torch.softmax(s, dim=-1) @ tv out.backward(torch.tensor(dO)) g_auto = (tq.grad.numpy(), tk.grad.numpy(), tv.grad.numpy()) e1 = max(np.abs(a - b).max() for a, b in zip(g_flash, g_std)) e2 = max(np.abs(a - b).max() for a, b in zip(g_flash, g_auto)) tag = "有因果掩码" if causal else "无掩码" print(f"{pad(tag, 12)}dQ、dK、dV：FlashAttention vs 标准反向 最大差 {e1:.1e}；vs PyTorch 自动求导 最大差 {e2:.1e}") N2, d2 = 1024, 64 f_std, f_flash = attention_flops(N2, d2), attention_flops(N2, d2, flash=True) print(f"\n一个头（N = {N2}，d = {d2}）前向 + 反向的矩阵乘法 FLOPs：标准 12N²d = {f_std / 1e6:.1f} M，" f"FlashAttention 14N²d = {f_flash / 1e6:.1f} M（多 {f_flash / f_std - 1:.0%}）") print("论文图 2（GPT-2 medium，N = 1024，A100）：GFLOPs 66.6 → 75.2（多 13%），HBM 读写 40.3 GB → 4.4 GB，用时 41.7 → 7.3 ms") print("\n观察 Observations：") print(" 1. 三种算法给出同样的梯度（差在 1e-15 量级）：反向时重算 P 不改变结果。") print(" 2. 标准反向需要前向存下的 N×N 矩阵 P；FlashAttention 只存 O 和每行两个数 (m, ℓ)，") print(" 反向时用它们一块一块地重算 S 和 P——多做一次 QKᵀ（FLOPs 多约 17%），") print(" 换来不读写 N×N 矩阵。论文实测 FLOPs 多 13%，用时反而少了 5.7 倍：在 GPU 上，多算比多搬便宜。") # ============================================================================= # 演示 5：装进 base model # ============================================================================= def flash_attention_torch(q, k, v, block=32): """ PyTorch 版的分块因果注意力（FlashAttention-2 风格：输出先不归一化，最后统一除以 ℓ）。 【思路】对每一块查询 Q_i，按顺序扫过它能看见的键块 K_j（因果：j 块的起点不超过 i 块的终点）， 维护每一行的 m、ℓ 和"未归一化的输出" acc： m_new = max(m, rowmax(S_ij)) acc = acc · e^{m − m_new} + exp(S_ij − m_new) V_j ℓ = ℓ · e^{m − m_new} + rowsum(exp(S_ij − m_new)) 扫完之后 O_i = acc / ℓ。比算法 1 少了每一步的"除以 ℓ 再乘回来"——这正是 FlashAttention-2 减少非矩阵乘法运算的改动之一。 【参数】q, k, v：(B, H, T, dh)；block：块大小（Br = Bc = block） 【返回】(B, H, T, dh) 【边界】T 不是 block 的整数倍时最后一块小一些；对角线上的块需要块内的因果掩码，其余可见块不需要。 这只是为了验证"结果相同"：Python 循环在 CPU 上比一次大矩阵乘法慢，速度不是这里的重点。 """ B, H, T, dh = q.shape scale = 1.0 / math.sqrt(dh) out = torch.empty_like(q) for i0 in range(0, T, block): qi = q[:, :, i0:i0 + block] n_i = qi.shape[2] m = torch.full((B, H, n_i, 1), float("-inf")) l = torch.zeros(B, H, n_i, 1) acc = torch.zeros(B, H, n_i, dh) for j0 in range(0, i0 + n_i, block): kj, vj = k[:, :, j0:j0 + block], v[:, :, j0:j0 + block] s = qi @ kj.transpose(-2, -1) * scale if j0 + kj.shape[2] > i0: # 和对角线相交的块：块内加因果掩码 rows = torch.arange(i0, i0 + n_i)[:, None] cols = torch.arange(j0, j0 + kj.shape[2])[None, :] s = s.masked_fill(cols > rows, float("-inf")) m_new = torch.maximum(m, s.amax(dim=-1, keepdim=True)) p = torch.exp(s - m_new) corr = torch.exp(m - m_new) # 第一块时 m = −∞，corr = 0，正好清掉初始值 acc = acc * corr + p @ vj l = l * corr + p.sum(dim=-1, keepdim=True) m = m_new out[:, :, i0:i0 + n_i] = acc / l return out def window_attention_torch(q, k, v, window): """ 滑动窗口注意力（近似）：每个位置只看自己和前面 window − 1 个位置。 【思路】Longformer、Sparse Transformer 一类"近似注意力"的最简单代表：把 N×N 变成 N×w，省计算和内存， 代价是看不见远处。这里仍然先算完整的分数再遮住（为了简单），只用来比较"模型质量"。 【参数】q, k, v：(B, H, T, dh)；window：窗口大小 【返回】(B, H, T, dh) """ T = q.shape[2] s = q @ k.transpose(-2, -1) / math.sqrt(q.shape[-1]) rows, cols = torch.arange(T)[:, None], torch.arange(T)[None, :] s = s.masked_fill((cols > rows) | (cols <= rows - window), float("-inf")) return torch.softmax(s, dim=-1) @ v def patched_model(base, attention_fn): """ 复制一份 base model，把每一层注意力的"核心计算"换成 attention_fn(q, k, v)。 【思路】q、k、v 的投影和输出投影（proj）仍用原来的参数，只换中间那一段—— 这正是 FlashAttention 在真实代码里的位置：一个可以直接替换的 kernel，模型参数一个都不用改。 【参数】base：第 1 讲的 TinyGPT（不会被修改）；attention_fn：函数 (q, k, v) → 输出，形状 (B, H, T, dh) 【返回】新的模型 """ model = copy.deepcopy(base) def forward(self, x): B, T, D = x.shape H = self.n_head q, k, v = self.qkv(x).split(D, dim=2) q, k, v = (t.view(B, T, H, D // H).transpose(1, 2) for t in (q, k, v)) y = attention_fn(q, k, v) return self.proj(y.transpose(1, 2).reshape(B, T, D)) for blk in model.blocks: blk.attn.forward = types.MethodType(forward, blk.attn) return model def demo5_in_base_model(): print_header("演示 5：把分块注意力装进第 1 讲的 base model（验证集 bit/字符）") world, tok = build_world(), CharTokenizer(ALPHABET) base, _ = load_or_train_base_model(world, tok) val = torch.tensor(tok.encode(world.val_text()), dtype=torch.long) ref = bits_per_char(base, val) print(f"{pad('注意力的算法', 44)}{pad('验证集 bit/字符', 18, '>')}{pad('与原模型之差', 16, '>')}") print(f"{pad('原模型（标准注意力，N×N 全部算出）', 44)}{pad(f'{ref:.6f}', 18, '>')}{pad('—', 16, '>')}") variants = [(f"分块精确注意力，块大小 {b}", (lambda b: lambda q, k, v: flash_attention_torch(q, k, v, b))(b)) for b in (8, 32, 96)] variants += [(f"滑动窗口（近似），只看最近 {w} 个字符", (lambda w: lambda q, k, v: window_attention_torch(q, k, v, w))(w)) for w in (32, 16, 8, 4)] for name, fn in variants: bpc = bits_per_char(patched_model(base, fn), val) print(f"{pad(name, 44)}{pad(f'{bpc:.6f}', 18, '>')}{pad(f'{bpc - ref:+.1e}', 16, '>')}") x = torch.tensor([tok.encode("kari lives in varno.\n12+30=42.\n")]) with torch.no_grad(): diff = (patched_model(base, lambda q, k, v: flash_attention_torch(q, k, v, 8))(x) - base(x)).abs().max().item() print(f"\n同一段输入的 logits：分块（块大小 8）与原模型最大差 {diff:.1e}（float32 的舍入误差）") print("\n观察 Observations：") print(" 1. 换成分块注意力，验证集 bit/字符的差在 1e-9 量级（小数点后 6 位完全相同），logits 差 1e-5 量级——") print(" 这是 float32 下求和顺序不同带来的舍入误差。模型参数一个没改，就可以直接换上：精确算法。") print(" 2. 滑动窗口是真正的近似：窗口 32 时几乎无损（这个迷你世界里一行只有几十个字符，远处的字符本来就不太重要），") print(" 窗口越小模型越差：窗口 8 时多了 0.19，窗口 4 时 bit/字符变成原来的 4.7 倍。2020–2021 年很多'高效 Transformer'走的是这条路：") print(" 省了 FLOPs，但改变了模型，在真实任务上常常掉点，墙上时间也未必更快。") print(" FlashAttention 走了另一条路：FLOPs 一个不少（甚至更多），只减少读写——质量完全不变。") if __name__ == "__main__": rng = np.random.default_rng(0) demo1_overflow() demo2_online() demo3_forward(rng) demo4_backward(rng) demo5_in_base_model()