""" ============================================================================= CS349E · 第 11 讲 · 05_megakernel.py —— 一次 decode 一百多个内核：空隙从哪来，megakernel 怎样消灭它 ============================================================================= 【这个脚本要回答的问题】 批 1 的 decode 受限于带宽：理想时间 = 权重字节 / 显存带宽。Llama-3.2-1B 在 H100 上是 0.74 ms， 可 Hazy Research 的博客说现有的推理引擎最多用到一半的带宽。剩下的时间去哪了？ 演示 1：Llama-3.2-1B 一次 decode 的账本：各算子读多少字节；H100、B200 上的理想时间。 演示 2：时间线模型：每层 12 个内核 / 7 个内核（stream / CUDA Graphs）/ megakernel（有无权重预取），时间分解。 演示 3：尾部效应（波次量化）：Hazy 的例子"512 个块在 148 个 SM 上"，以及我们的 GEMV 内核。 演示 4：敏感性：模型里每个假设的参数变一变，结论还成立吗？ 【关于数字】这是一个**模型**：参数一部分来自 Hazy 的博客（启动间隙、每 SM 的共享内存页），一部分是假设 （序幕、同步、指令开销、带宽效率），都写在 flash_lab.TimelineParams 里。用它看"时间去哪了"和量级，不是预测具体的毫秒数。 【运行】python3 05_megakernel.py （约 1 秒） 【依赖】numpy；本目录的 flash_lab；../lecture01 的 hw_specs、console。 """ import os import sys from dataclasses import replace 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 TimelineParams, llama_1b_ops, timeline # noqa: E402 from hw_specs import B200, H100 # noqa: E402 from console import fmt_bytes, fmt_time, print_header, print_observations, print_table # noqa: E402 P_H100 = TimelineParams(n_sm=132, hbm=H100.mem_bw) P_B200 = TimelineParams(n_sm=148, hbm=B200.mem_bw) STREAM, GRAPHS = 2.1e-6, 1.3e-6 # Hazy：H100 上一个空内核在 stream 上约 2.1 µs，用 CUDA Graphs 约 1.3 µs def configs(p): """五种执行方式：(名字, 算子列表, 模式, 参数)。""" return [("每层 12 个内核，stream 启动", llama_1b_ops(kernels="12"), "kernels", replace(p, launch=STREAM)), ("每层 7 个内核，stream 启动", llama_1b_ops(kernels="7"), "kernels", replace(p, launch=STREAM)), ("每层 7 个内核，CUDA Graphs", llama_1b_ops(kernels="7"), "kernels", replace(p, launch=GRAPHS)), ("megakernel，不预取权重", llama_1b_ops(kernels="fused"), "megakernel", replace(p, smem_buf=0)), ("megakernel，预取权重", llama_1b_ops(kernels="fused"), "megakernel", p)] def demo1_ledger(): print_header("演示 1：Llama-3.2-1B（BF16）批 1 的一次 decode 读多少字节") ops = llama_1b_ops(kernels="fused") per_layer = {} for o in ops[:5]: per_layer[o.name] = o.bytes rows = [[name, fmt_bytes(b), fmt_bytes(b * 16), f"{b * 16 / sum(o.bytes for o in ops):.1%}"] for name, b in per_layer.items()] lm = ops[-1].bytes total = sum(o.bytes for o in ops) rows.append(["norm + lm_head（最后一次）", fmt_bytes(lm), fmt_bytes(lm), f"{lm / total:.1%}"]) print_table(["算子（megakernel 的指令粒度）", "每层", "16 层合计", "占总字节"], rows) rows = [] for chip in (H100, B200): t = total / chip.mem_bw rows.append([chip.name, f"{chip.mem_bw / 1e12:.2f} TB/s", fmt_time(t), f"{1 / t:,.0f}"]) print(f"\n总共 {fmt_bytes(total)}（注意力读的 KV cache 按 160 个 token 计，只有几百 KB）。") print_table(["GPU", "显存带宽", "理想时间（只读一遍权重）", "每秒最多几次前向"], rows) print_observations([ "批 1 的 decode 每个权重只用一次，强度约 1 FLOP/字节，完全受限于带宽（第 1 讲 §4）。理想时间 = 读一遍所有权重。" "H100 上约 0.74 ms、每秒约 1,350 次，和 Hazy 博客算的一样；B200 约 0.31 ms（博客说\"约 3,000 次/秒\"）。", "一个 1B 的模型被切成了大约一百个算子：每层的 QKV、O、gate/up、down 四个矩阵乘向量（GEMV），再加上归一化、RoPE、注意力、残差。" "最大的是 MLP（gate/up 每层 67 MB）和词表很大的 lm_head（525 MB，占五分之一）。", "每个算子只读几 MB 到几十 MB：按 3.35 TB/s，一个 8 MB 的 O 投影只要 2.5 µs——和一次内核启动的开销同一个量级。这就是问题所在。", ]) def demo2_timeline(): print_header("演示 2：时间线模型（H100）——时间去哪了") rows, parts_rows = [], [] res = {} for name, ops, mode, p in configs(P_H100): r = timeline(ops, mode, p) res[name] = r rows.append([name, len(ops), fmt_time(r["t"]), f"{1 / r['t']:,.0f}", f"{r['bw_util']:.0%}"]) parts_rows.append([name] + [fmt_time(v) if v > 1e-9 else "—" for v in r["parts"].values()]) print_table(["执行方式", "内核 / 指令数", "一次前向", "每秒", "带宽利用率（理想 / 实际）"], rows) print() keys = list(next(iter(res.values()))["parts"].keys()) print_table(["分解"] + keys, parts_rows) k7 = res["每层 7 个内核，CUDA Graphs"] mk = res["megakernel，预取权重"] print_observations([ f"多内核的写法：每层 7 个内核、用 CUDA Graphs 时模型给出 {fmt_time(k7['t'])}（带宽利用率 {k7['bw_util']:.0%}）。" "Hazy 的估计是\"每层 7 个内核 × 16 层 × 每个约 5 µs 的停顿 → 每秒约 770 次\"；我们的模型每个内核的停顿是" "启动间隙 + 序幕 + 尾部，合起来也是几微秒，结果在同一个量级。Hazy 测到 vLLM、SGLang 最多用到 50% 的带宽——" "真实的引擎还有采样、调度这些 CPU 上的开销，模型里没有。", "空隙由几部分组成：启动间隙（前一个内核全部结束、下一个才开始，CUDA Graphs 只能把它从 2.1 µs 压到 1.3 µs）；" "序幕（内核开始后要等第一批权重从显存到达）；尾部（最后一波块没占满 SM）；还有那些几乎不读数据、却要占一个完整内核时间的小内核。", f"megakernel：整个前向是一个内核，每个 SM 按顺序执行\"指令\"。权重不依赖上一个算子的结果，可以在共享内存的空页里提前读——" f"等上一个算子的激活时，权重还在源源不断地流进来。模型给出 {fmt_time(mk['t'])}、带宽利用率 {mk['bw_util']:.0%}；" "不预取时降到约八成。Hazy 在 H100 上实测：78% 的带宽、每次前向不到 1 ms，比 vLLM 快约 2.5 倍、比 SGLang 快 1.5 倍以上。", "模型比 Hazy 的实测乐观：它没有算矩阵乘向量本身的计算时间、每条指令读写激活和检查一致性的开销" "（Hazy 在 B200 上的分解：约 600 µs 里，存取激活和一致性检查约 250 µs、RMSNorm 和矩阵乘向量约 200 µs、设置等约 80 µs、warp 同步约 40 µs）。", ]) def demo3_tails(): print_header("演示 3：尾部效应（波次量化）") rows = [] for ctas, sms, what in ((512, 148, "Hazy 的例子：B200 上 512 个块"), (192, 132, "我们的 QKV 内核（3072 行，每块 16 行），H100"), (128, 132, "O 投影 / down（2048 行），H100"), (1024, 132, "gate/up（16384 行），H100"), (8016, 132, "lm_head（128,256 行），H100")): waves = -(-ctas // sms) last = ctas - (waves - 1) * sms rows.append([what, ctas, sms, waves, last, sms - last, f"{ctas / (waves * sms):.0%}"]) print_table(["内核", "块数", "SM 数", "波数", "最后一波的块", "最后一波空闲的 SM", "SM 的平均利用率"], rows) print_observations([ "一个内核的所有块结束，下一个内核才能开始（同一个 stream 上的语义）。最后一波如果只占了一部分 SM，其余的 SM 就在等：" "Hazy 的例子里 512 = 3 × 148 + 68，最后一波 80 个 SM 空闲。", "块数恰好比 SM 数多一点最糟（192 个块在 132 个 SM 上：利用率 73%）；块数是 SM 数的很多倍时尾部占比很小（lm_head 的 8016 个块：99%）。", "常见的缓解：持久化内核（块数 = SM 数，每个块循环处理多份工作）、split-K、Hopper 的 PDL（Programmatic Dependent Launch，" "让下一个内核在前一个结束前就开始执行序幕）。Hazy 指出 PDL 的同步太粗：注意力必须等**所有**的 Q、K、V 都算完，" "而 megakernel 用计数器做细粒度的依赖（例如 down 投影分成 4 块，每块只等自己那部分输入）。", ]) def demo4_sensitivity(): print_header("演示 4：敏感性——每个假设的参数变化时，结论还成立吗？（H100；带宽利用率）") base = P_H100 variants = [("基准参数", base), ("序幕 0.5 µs（原 1 µs）", replace(base, ramp=0.5e-6)), ("序幕 2 µs", replace(base, ramp=2e-6)), ("小内核 1 µs（原 2 µs）", replace(base, small=1e-6)), ("同步 2 µs（原 1 µs）", replace(base, sync=2e-6)), ("指令开销 1 µs（原 0.5 µs）", replace(base, instr_overhead=1e-6)), ("带宽效率 0.8（原 0.9）", replace(base, eff=0.8))] rows = [] for name, p in variants: r = {} for cname, ops, mode, pp in configs(p): if cname.startswith("每层 7 个内核，CUDA") or cname == "megakernel，预取权重": r[cname] = timeline(ops, mode, replace(pp, ramp=p.ramp, small=p.small, sync=p.sync, instr_overhead=p.instr_overhead, eff=p.eff)) k, m = r["每层 7 个内核，CUDA Graphs"], r["megakernel，预取权重"] rows.append([name, f"{k['bw_util']:.0%}", f"{m['bw_util']:.0%}", f"{k['t'] / m['t']:.2f}"]) for name, p in (("B200（148 SM、8 TB/s）", P_B200),): k = timeline(llama_1b_ops(kernels="7"), "kernels", replace(p, launch=GRAPHS)) m = timeline(llama_1b_ops(kernels="fused"), "megakernel", p) rows.append([name, f"{k['bw_util']:.0%}（{fmt_time(k['t'])}）", f"{m['bw_util']:.0%}（{fmt_time(m['t'])}）", f"{k['t'] / m['t']:.2f}"]) print_table(["参数", "7 个内核 + CUDA Graphs", "megakernel", "megakernel 快几倍"], rows) print_observations([ "每一种参数变化下，megakernel 都比多内核快 1.3–1.5 倍，带宽利用率高 20–29 个百分点：结论不依赖某一个假设的具体值。", "多内核的写法对\"序幕\"和\"小内核\"最敏感——这两项每个内核都要付一次；megakernel 对它们不敏感，" "但对\"带宽效率\"和\"指令开销\"敏感：空隙消除以后，剩下的就是读权重本身和每条指令的固定成本。", "B200 上带宽是 H100 的 2.4 倍、读权重的时间短了，固定的空隙占比更大，所以多内核的利用率更低、megakernel 的优势更大——" "Hazy 在 B200 上测到比 vLLM 快 3.5 倍以上（H100 上约 2.5 倍）、每次前向不到 680 µs。硬件越快，软件的空隙越显眼。" "（模型给 B200 的 megakernel 是 373 µs，比实测乐观得多：带宽变大以后，模型没算的计算和激活读写（演示 2 观察 4）占比更大。）", "这也是第 6 讲投机解码、第 5 讲批处理的另一面：批 1 的 decode 天生\"每个内核只有几微秒的活\"，要么把活变多（批、投机），" "要么把内核之间的空隙消灭（megakernel、CUDA Graphs、PDL）。", ]) if __name__ == "__main__": demo1_ledger() demo2_timeline() demo3_tails() demo4_sensitivity()