""" ============================================================================= CS349E · 第 11 讲 · 01_online_softmax.py —— 从在线 softmax 到"一遍扫完 K、V"的注意力 ============================================================================= 【这个脚本要回答的问题】 前置讲 P3 §5.3 讲过在线 softmax：一遍扫数据同时维护最大值 m 和 Σ exp(x − m)。 FlashAttention 再往前走一步：连同乘 V 也放进这一遍里，分数向量从来不需要完整地存下来。 演示 1：一个 query 的注意力的三种算法：先存下整行分数（3 遍）→ 在线 softmax（2 遍）→ 连同 V 一起（1 遍）。 演示 2：局部状态 (m, l, o) 的合并：顺序、树形、乱序合并结果都一样——所以可以把 K、V 切开并行算。 演示 3：FlashAttention-2 的"推迟归一化"：只在最后除一次 l，省掉多少非矩阵乘法的运算。 演示 4：精度：m、l、累加器用 FP32，P 转成 FP16 / BF16 再乘 V，误差多大。 【运行】python3 01_online_softmax.py （约 1 秒） 【依赖】numpy；本目录的 flash_lab；../prereq3 的 numfmt（FP16、BF16 的舍入）；../lecture01 的 console。 """ 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")) sys.path.insert(0, os.path.join(HERE, "..", "prereq3")) from flash_lab import merge_states, partial_state, softmax_3pass # noqa: E402 from numfmt import BF16, FP16 # noqa: E402 from console import print_header, print_observations, print_table # noqa: E402 def attn_row_3pass(q, K, V): """存下整行分数：s = K·q（写一个长度 S 的向量）→ 求 max → 求和 → 归一化 → 乘 V。读 K 一遍、读 s 三遍。""" s = K @ q / np.sqrt(len(q)) p, _ = softmax_3pass(s) return p @ V def attn_row_2pass(q, K, V, block): """第一遍：分块算分数，在线维护 (m, l)；第二遍：重新算分数、除以 l、乘 V。不存整行分数，但 K 要读两遍。""" S = K.shape[0] m, l = -np.inf, 0.0 for a in range(0, S, block): s = K[a:a + block] @ q / np.sqrt(len(q)) m_new = max(m, s.max()) l = l * np.exp(m - m_new) + np.exp(s - m_new).sum() m = m_new out = np.zeros(V.shape[1]) for a in range(0, S, block): s = K[a:a + block] @ q / np.sqrt(len(q)) out += np.exp(s - m) / l @ V[a:a + block] return out def attn_row_1pass(q, K, V, block, trace=None): """ 一遍：手里维护 (m, l, acc)，每来一块 K、V： m_new = max(m, 块内最大值)；α = exp(m − m_new) l ← α·l + Σ exp(s − m_new)；acc ← α·acc + exp(s − m_new)·V_块 最后 out = acc / l。acc 是"按当前 m 缩放过的加权和"，每次 m 变大就整体乘 α 换算。 """ S = K.shape[0] m, l = -np.inf, 0.0 acc = np.zeros(V.shape[1]) for a in range(0, S, block): s = K[a:a + block] @ q / np.sqrt(len(q)) m_new = max(m, s.max()) alpha = np.exp(m - m_new) p = np.exp(s - m_new) l = alpha * l + p.sum() acc = alpha * acc + p @ V[a:a + block] if trace is not None: trace.append((a // block, s.max(), m_new, alpha, l)) m = m_new return acc / l def demo1_three_ways(): print_header("演示 1：一个 query 的注意力（S = 4096 个 key，d = 64，每块 256 个）的三种算法") rng = np.random.default_rng(0) S, d, blk = 4096, 64, 256 q = rng.standard_normal(d) K = rng.standard_normal((S, d)) V = rng.standard_normal((S, d)) ref = attn_row_3pass(q, K, V) rows = [ ["存下整行分数（3 遍）", "K 一遍", "S 个分数：写 1 次、读 3 次", "长度 S 的向量", 0.0], ["在线 softmax（2 遍）", "K 两遍、V 一遍", "不存", "常数", np.abs(attn_row_2pass(q, K, V, blk) - ref).max()], ["在线 softmax × V（1 遍）", "K、V 各一遍", "不存", "常数（m、l、一个 d 维累加器）", np.abs(attn_row_1pass(q, K, V, blk) - ref).max()], ] print_table(["算法", "读 K、V", "中间的分数向量", "额外存储", "和第一种的最大差"], [r[:4] + [f"{r[4]:.1e}"] for r in rows]) trace = [] attn_row_1pass(q, K[:1024], V[:1024], blk, trace) print("\n一遍算法的前 4 块（m 变大时，之前的 l 和累加器都乘 α 换算）：") print_table(["块", "块内最大分数", "新的 m", "α = exp(m_旧 − m_新)", "l"], [[b, f"{bm:.3f}", f"{mn:.3f}", f"{al:.4f}", f"{l:.3f}"] for b, bm, mn, al, l in trace]) print_observations([ "三种算法的结果相同（差在 1e-16 量级，只是舍入顺序不同）。区别全在\"中间的分数向量\"：标准写法要把 S 个分数写出来再读三遍；" "一遍算法从来不需要完整的分数向量——每块算完就扔。", "α 的作用：累加器 acc 和 l 一直是\"按当前见过的最大值 m 缩放\"的。新来的块如果有更大的分数，m 变大，之前累加的东西整体乘 α < 1 换算；" "m 没变时 α = 1。第 0 块的 α 是 exp(−∞) = 0（之前什么也没有）。", "对 N 个 query 来说，标准写法要存 N×S 的分数矩阵（N = S = 4096、FP16 时 32 MiB / 每个头），而且要在显存里进出好几次；" "一遍算法只需要每个 query 的 (m, l, acc)——这就是 FlashAttention 的全部数学，剩下的是怎样分块（02 脚本）。", ]) def demo2_merge_orders(): print_header("演示 2：局部状态 (m, l, o) 的合并与顺序无关") rng = np.random.default_rng(1) S, d = 3000, 32 x = rng.standard_normal(S) * 3 V = rng.standard_normal((S, d)) _, l_ref, o_ref = partial_state(x[None, :], V) ref = o_ref[0] / l_ref[0] cuts = np.sort(rng.choice(np.arange(1, S), size=11, replace=False)) segs = np.split(np.arange(S), cuts) states = [partial_state(x[None, s], V[s]) for s in segs] def finish(st): return st[2][0] / st[1][0] def seq(sts): acc = sts[0] for s in sts[1:]: acc = merge_states(acc, s) return acc def tree(sts): while len(sts) > 1: sts = [merge_states(sts[i], sts[i + 1]) if i + 1 < len(sts) else sts[i] for i in range(0, len(sts), 2)] return sts[0] perm = rng.permutation(len(states)) rows = [["从左到右依次合并", f"{np.abs(finish(seq(states)) - ref).max():.1e}"], ["两两合并（树形，像并行归约）", f"{np.abs(finish(tree(states)) - ref).max():.1e}"], ["打乱顺序再依次合并", f"{np.abs(finish(seq([states[i] for i in perm])) - ref).max():.1e}"]] print(f"把长度 {S} 的一行随机切成 {len(segs)} 段（段长 {min(len(s) for s in segs)}–{max(len(s) for s in segs)}），各算局部状态再合并：") print_table(["合并顺序", "和一次算完的最大差"], rows) print_observations([ "合并运算 merge((m_a, l_a, o_a), (m_b, l_b, o_b)) 满足结合律和交换律：先合并谁、怎么分组都得到同一个结果（差别只是浮点舍入，1e-16 量级）。", "这给了并行的自由：FlashAttention-2 让不同的 SM 算不同的 query 块（各自顺序扫 K、V）；" "flash-decoding 让不同的 SM 算同一个 query 的不同 K、V 段，最后合并（04 脚本）；" "张量并行时不同的 GPU 各算一部分头或一部分序列，也用同一个合并公式（ring attention）。", "flash-decoding 存的是 (o_s, lse_s)：o_s 已经除过本段的 l，lse_s = m_s + log l_s，只多存一个数；合并时按 exp(lse_s − lse) 加权。两种写法等价。", ]) def demo3_delayed_normalization(): print_header("演示 3：FlashAttention-2 的\"推迟归一化\"——非矩阵乘法的运算省了多少（一个头，N = 4096，d = 128，Br = Bc = 128）") N, d, Bc = 4096, 128, 128 Tc = N // Bc matmul = 4 * N * N * d # 每处理一个 (i, j) 块：求块内 max、减 max、exp、求和（softmax 本身，两种写法都要） softmax_ops = 4 * N * N # FA1 式：每一步都把 O_i 换算并重新归一化：O ← (l_old·α·O + β·PV) / l_new → 每个元素约 4 次运算 fa1_rescale = Tc * N * d * 4 # FA2 式：每一步只把 acc 乘 α（1 次）、加 PV（矩阵乘法里）；最后除一次 l fa2_rescale = Tc * N * d * 1 + N * d rows = [] for name, extra in (("论文算法 1（每步重新归一化 O）", fa1_rescale), ("FlashAttention-2（只在最后除以 l）", fa2_rescale)): nonmm = softmax_ops + extra rows.append([name, f"{matmul / 1e9:.2f}", f"{nonmm / 1e9:.3f}", f"{extra / 1e9:.3f}", f"{nonmm / matmul:.1%}", f"{16 * nonmm / matmul:.0%}"]) print_table(["写法", "矩阵乘法 GFLOP", "非矩阵乘法 GFLOP", "其中换算 / 归一化", "非矩阵乘法 / 矩阵乘法", "按 A100 上贵 16 倍折算的时间占比"], rows) print_observations([ "注意力里 exp、max、求和这些\"非矩阵乘法\"的运算只有矩阵乘法的几个百分点，但它们不能用 Tensor Core：" "A100 的 FP16 矩阵乘法峰值 312 TFLOPS，非矩阵乘法的 FP32 只有 19.5 TFLOPS，每次运算贵 16 倍（FlashAttention-2 论文的说法）。", "按 16 倍折算，非矩阵乘法的时间相当于矩阵乘法的 25%（每步重新归一化）和 16%（推迟归一化）——不可忽略。" "所以 FlashAttention-2 尽量削减它们：每一步只把累加器乘 α，最后才除一次 l；反向传播只存 logsumexp L 一个数而不是 m、l 两个。", "这也是第 1 讲 Blackwell Ultra 把特殊函数单元（SFU）的 exp 吞吐翻倍的原因：注意力的 softmax 在新硬件上开始成为瓶颈；" "FlashAttention-3 还把 softmax 和下一块的矩阵乘法在不同的 warpgroup 之间交错执行（ping-pong，第 8 讲 §8.5）。", ]) def demo4_precision(): print_header("演示 4：精度——P 用什么格式乘 V？（S = 8192 个 key，d = 64，FP64 的结果作参照）") rng = np.random.default_rng(2) S, d = 8192, 64 q = rng.standard_normal(d) K = rng.standard_normal((S, d)) V = rng.standard_normal((S, d)) s = K @ q / np.sqrt(d) m = s.max() p = np.exp(s - m) ref = (p / p.sum()) @ V rows = [] for name, fmt in (("FP32", None), ("BF16", BF16), ("FP16", FP16)): pq = p.astype(np.float32) if fmt is None else fmt.round(p.astype(np.float32)) out = (pq.astype(np.float64) @ V) / p.sum() rows.append([f"P 用 {name} 乘 V（l 用 FP32 精确累加）", f"{np.abs(out - ref).max() / np.abs(ref).max():.1e}"]) l16 = np.float16(0.0) for v in p.astype(np.float16): l16 = np.float16(l16 + v) lb = 0.0 for v in BF16.round(p.astype(np.float32)): lb = float(BF16.round(np.float32(lb + v))) for name, lv in (("FP16", float(l16)), ("BF16", lb)): rows.append([f"l（Σ exp）也用 {name} 逐个累加", f"{abs(lv - p.sum()) / p.sum():.1e}"]) with np.errstate(over="ignore"): big = np.float16(np.exp(np.float16(12.0))) print_table(["做法", "相对误差"], rows) print(f"\n另外：不减最大值时，FP16 的 exp(12) = {big}（FP16 最大约 65504，exp(11.09) 就溢出）。") print_observations([ "FlashAttention 的做法：m、l、输出累加器都用 FP32；P = exp(S − m) 转成 FP16 / BF16 后和 V 做第二个矩阵乘法（Tensor Core 的输入只能是低精度），" "累加仍是 FP32。表里第一组：P 用 BF16 的相对误差约 1e-3、FP16 约 1e-4——BF16 的尾数只有 8 位（FP16 是 11 位）。", "l 绝不能用低精度逐个累加：8192 个 0 到 1 之间的数加起来，部分和变大以后，后面加的小数被\"吞掉\"（前置讲 P3 §4.3 的停滞），" "误差比 FP32 大几个数量级。", "减最大值不只是为了防止 FP32 溢出：FP16 的 exp 在 11 附近就溢出了。减掉 m 之后所有的 exp ≤ 1，FP16 / BF16 都放得下。", ]) if __name__ == "__main__": demo1_three_ways() demo2_merge_orders() demo3_delayed_normalization() demo4_precision()