""" ============================================================================= CS349E · 第 11 讲 · 04_flash_decoding.py —— decode 时的注意力：沿 KV 的长度再切一刀 ============================================================================= 【这个脚本要回答的问题】 FlashAttention 沿 query 切块、每块一个 CTA。训练和 prefill 时 query 有几千个，块多得填满 GPU； decode 时每个序列只有 1 个新 query，CTA 数 = 批 × 头数——批 1、16 个头只有 16 个 CTA，108 个 SM 闲着九成。 flash-decoding（Dao、Haziza、Massa、Sizov，2023）沿 KV 的长度再切一刀，最后用 log-sum-exp 合并。 演示 1：分段 + LSE 合并是精确的；GQA 时同一组的 query 头打包在一起算。 演示 2：并行度：FlashAttention 式的 decode 能用上几个 SM。 演示 3：时间模型 vs flash-decoding 博客的实测（A100，CodeLlama-34B 在 4 卡上的形状：16 个 query 头、2 个 KV 头、d = 128）。 演示 4：段数怎么选；H100 上 Llama-3.1-8B 的 decode 注意力。 【运行】python3 04_flash_decoding.py （约 1 秒） 【依赖】numpy；本目录的 flash_lab；../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")) from flash_lab import A100_DEC, H100_DEC, best_splits, decode_attention_time, split_kv_decode, standard_attention # noqa: E402 from console import fmt_time, print_header, print_observations, print_table # noqa: E402 # flash-decoding 博客（PyTorch 博客，2023-10）的微基准：A100、FP16、批 B、序列长度 S，单位 µs BLOG = [(256, 256, 3058.6, 390.5, 63.4), (128, 512, 3151.4, 366.3, 67.7), (64, 1024, 3160.4, 364.8, 77.7), (32, 2048, 3158.3, 352.0, 58.5), (16, 4096, 3157.0, 401.7, 57.0), (8, 8192, 3173.1, 529.2, 56.4), (4, 16384, 3223.0, 582.7, 58.2), (2, 32768, 3224.1, 1156.1, 60.3), (1, 65536, 1335.6, 2300.6, 64.4), (1, 131072, 2664.0, 4592.2, 106.6)] def demo1_exact(): print_header("演示 1：分段 + log-sum-exp 合并是精确的（S = 10,000 个缓存的 key，d = 128）") rng = np.random.default_rng(0) S, d = 10000, 128 rows = [] for R in (1, 8): q = rng.standard_normal((R, d)) K, V = rng.standard_normal((S, d)), rng.standard_normal((S, d)) ref, _, _ = standard_attention(q, K, V) for ns in (1, 7, 64, 1000): o, _, parts = split_kv_decode(q, K, V, ns) rows.append([R, ns, len(parts), f"{np.abs(o - ref).max():.1e}"]) print_table(["打包的 query 头数 R", "段数", "实际的段数", "和直接计算的最大差"], rows) print_observations([ "每段算出本段归一化过的输出 o_s 和 lse_s = log Σ_段 exp(分数)；合并时 lse = log Σ exp(lse_s)，o = Σ exp(lse_s − lse)·o_s。" "和 01 脚本的 (m, l, o) 合并是同一件事，只是每段只多存一个数。切成 1000 段（每段 10 个 key）也精确。", "GQA（第 2 讲）：16 个 query 头共享 2 组 K、V，每组 8 个 query 头。把同一组的 8 个 query 当成一个 8 行的 Q 块，" "一次读 K、V 就服务 8 个头——flash-decoding 的实现就是这样把 GQA 的头\"打包\"进矩阵乘法的 M 维的。", ]) def demo2_parallelism(): print_header("演示 2：FlashAttention 式的 decode 能用上几个 SM？（A100：108 个 SM；16 个 query 头）") rows = [] for B in (1, 2, 4, 8, 16, 64): ctas = B * 16 rows.append([B, ctas, f"{min(ctas, 108) / 108:.0%}", f"{min(1.0, 1 / 128):.1%}"]) print_table(["批大小", "CTA 数 = 批 × 头", "同时在算的 SM 比例", "每个 CTA 的 Q 块（128 行）里有效的比例"], rows) print_observations([ "FlashAttention 沿 query 切块：decode 时每个序列只有 1 个 query，一个 CTA 负责一个 (序列, 头)。批 1 时 16 个 CTA，" "108 个 SM 只用了 15%——博客说\"批 1 时 FlashAttention 用不到 GPU 的 1%\"，指的是把这两个因素乘起来。", "第二个因素：Tensor Core 按块算，一个 CTA 的 Q 块是 128 行，只有 1 行是真的 query，其余 127 行是补的 0——" "这个 CTA 读 K、V 的速度被\"白算\"的那部分矩阵乘法拖住（下一个演示的模型里，批小、序列长时它受限于计算）。", "长上下文时批往往更小（KV cache 占满显存，第 1 讲 §4.4），偏偏这时注意力是 decode 的大头——问题最严重的地方就是最需要的地方。", ]) def demo3_vs_blog(): print_header("演示 3：时间模型 vs flash-decoding 博客的实测（A100，16 个 query 头 / 2 个 KV 头，d = 128，FP16）") print("模型参数（事先按 Little 定律和 FlashAttention-2 论文定的，没有按实测调）：" f"Tensor Core 效率 {A100_DEC.tc_eff}，每个 SM 最多拉 {A100_DEC.sm_bw_cap / 1e9:.0f} GB/s，每个内核固定开销 {A100_DEC.launch * 1e6:.0f} µs；" "FlashAttention-2 每个 CTA 的 Q 块按 128 行算，flash-decoding 把 8 个头打包、补到 64 行。") rows = [] for B, S, _, fa2_meas, fd_meas in BLOG: fa = decode_attention_time(B, 16, 2, S, 128, "fa2", A100_DEC) ns, fd = best_splits(B, 16, 2, S, 128, A100_DEC) rows.append([B, S, f"{fa2_meas:.0f}", f"{fa['t'] * 1e6:.0f}", fa["bound"], f"{fd_meas:.0f}", f"{fd['t'] * 1e6:.0f}", ns, fd["ctas"], f"{fa2_meas / fd_meas:.1f}", f"{fa['t'] / fd['t']:.1f}"]) print_table(["批", "序列长度", "FA2 实测 µs", "FA2 模型", "模型瓶颈", "FD 实测 µs", "FD 模型", "FD 段数", "FD 的 CTA 数", "实测 FA2/FD", "模型 FA2/FD"], rows) print_observations([ "模型对 FlashAttention-2 的估计在实测的 ±12% 以内：批大、序列短时受限于带宽（每个 query 头单独读一遍它那组的 K、V，读了 8 倍的数据）；" "批小、序列长时受限于计算——128 行的 Q 块只有 1 行有效，16 个 CTA 在 16 个 SM 上做了 128 倍的无用功。", "flash-decoding 的模型在 48–96 µs，实测 56–107 µs：几乎不随序列长度变——总的 KV 字节数（批 × 序列长度）在这张表里是常数，" "只要 CTA 填满了 GPU，时间就是\"读一遍 KV 的时间 + 两个内核的固定开销\"。最后一行 KV 翻倍，时间也翻倍。", "博客说 flash-decoding 的注意力最多比 FlashAttention-2 快 50 倍，表里最大的是 2300.6 / 64.4 = 36 倍（批 1、64K）；" "128K 时 43 倍。模型给出 40 倍、44 倍。", "\"PyTorch eager\"那一列（约 3 ms）我们没有建模：它把注意力写成几个独立的算子，还要为 GQA 把 K、V 复制 8 份。", ]) def demo4_splits(): print_header("演示 4：段数怎么选？以及 H100 上 Llama-3.1-8B（32 个 query 头 / 8 个 KV 头）的 decode 注意力") rows = [] for B, S in ((1, 65536), (8, 8192), (64, 2048)): for ns in (1, 2, 4, 8, 16, 32, 64, 128): r = decode_attention_time(B, 16, 2, S, 128, "fd", A100_DEC, splits=ns) rows.append([f"批 {B}，S = {S}", ns, r["ctas"], r["waves"], f"{r['util']:.0%}", fmt_time(r["t"])]) print_table(["设置（A100，16/2 头）", "段数", "CTA 数", "波数", "同时活跃的 SM", "模型时间"], rows) print() rows = [] for S in (1024, 8192, 32768, 131072): for B in (1, 16): fa = decode_attention_time(B, 32, 8, S, 128, "fa2", H100_DEC) ns, fd = best_splits(B, 32, 8, S, 128, H100_DEC) kv = 2 * B * S * 8 * 128 * 2 rows.append([S, B, f"{kv / 1e6:.1f} MB", fmt_time(kv / 3.35e12), fmt_time(fa["t"]), fmt_time(fd["t"]), ns, f"{fa['t'] / fd['t']:.1f}"]) print_table(["序列长度", "批", "一层的 KV（读一遍）", "按 3.35 TB/s", "FA2 式（模型）", "flash-decoding（模型）", "段数", "倍数"], rows) print_observations([ "段数太少：CTA 填不满 SM（批 1、64K 时 1 段只有 2 个 CTA）。段数太多：CTA 超过一波、每段太短，固定开销和合并的代价占比变大" "（这个模型里只体现为波数和第二个内核的开销）。好的段数让 CTA 数刚好是 SM 数的一到几倍——flash-decoding 的实现用类似的启发式。", "批 64、序列 2K：不切时有 64 × 2 = 128 个 CTA（每个 KV 头一个），比 108 个 SM 多一点，第二波几乎空着（151 µs）；" "切 4–16 段让波次排得更满，降到约 100 µs。波次量化（前置讲 P2 §7.3）在这里和并行度一样重要；真实的实现按批大小、长度和 SM 数自动选段数。", "Llama-3.1-8B 每组 4 个 query 头（GQA 4:1）：一层的 KV 在 128K 时 0.5 GB，按满带宽也要约 160 µs，32 层就是 5 ms——" "长上下文的 decode 注意力本身就是受限于带宽的大头。flash-decoding 让它接近\"读一遍 KV\"的下限；再往下只能减少 KV 的字节" "（GQA / MLA、KV 量化、稀疏注意力，第 2、3 讲）。", ]) if __name__ == "__main__": demo1_exact() demo2_parallelism() demo3_vs_blog() demo4_splits()