""" ============================================================================= 第 10 讲 · 02：写又对又快的 kernel —— fast_p 指标与"有信息量的反馈" ============================================================================= 对应讲义：lectures/lecture10_swe_agents.md 第 3、4 节（KernelBench；Agent–System Interface） ----------------------------------------------------------------------------- 【toy 任务：调一个 GPU kernel 的配置】 ----------------------------------------------------------------------------- 每个任务是一个 kernel，有 5 个旋钮： tile ∈ {8, 16, 32, 64} 分块大小 fuse ∈ {0, 1} 是否把相邻的运算融合成一个 kernel（少读写一次显存） vec ∈ {1, 2, 4} 向量化宽度（一次读写几个数） shared ∈ {0, 1} 是否使用共享内存 unroll ∈ {1, 2, 4} 循环展开 每个任务有隐藏的特性（agent 看不到）： · 访存密集的程度 m（0~1）：m 大时，融合、向量化、共享内存更有用；m 小时，tile 和 unroll 更重要； · 最合适的 tile（16 或 32）； · 问题规模 n：vec 不能整除 n 时，结果会算错（只有在规模不整除的测试输入上才看得出来）； · 这个任务的运算能不能融合：不能融合时强行融合会算错； · tile = 64 且用共享内存时，超出共享内存的容量 → 编译 / 启动失败。 PyTorch 的基线 = 一个默认配置（tile 32，其余不开）的运行时间 × 0.9（库函数本身已经调得不错）。 KernelBench 的指标 fast_p：**正确、并且比 PyTorch 基线快 p 倍以上**的任务所占的比例。 ----------------------------------------------------------------------------- 【三种反馈下的迭代优化（每个任务 N 轮）】 ----------------------------------------------------------------------------- 只有标量反馈 ：每轮只知道"运行时间是多少"或"失败了"。提议者随机改一个旋钮，更快就保留（像 OpenTuner）。 + 执行反馈 ：失败时知道**为什么**失败（"共享内存超出容量""在规模 n = 1003 的输入上结果不对"）， 提议者据此做有针对性的修改（调小 tile / 关掉共享内存；把 vec 改成能整除 n 的值；关掉 fuse）。 + 性能分析建议 ：成功时还知道瓶颈在哪里（"访存瓶颈：试试融合 / 向量化 / 共享内存""计算瓶颈：调 tile / unroll"）， 提议者优先改相关的旋钮。（这对应 Agent–System Interface 的 AutoGuide：把原始输出翻译成"解释 + 建议"。） 对照：重复采样 —— 每轮独立地随机写一个新配置，留最快的那个（不利用任何反馈）。 【重要】三种反馈下"提议者怎样利用信息"是**设定**的，演示的是机制：反馈的信息量越大，搜索越省。 运行时间：约 1 秒（只用标准库）。 """ import math import random KNOBS = {"tile": (8, 16, 32, 64), "fuse": (0, 1), "vec": (1, 2, 4), "shared": (0, 1), "unroll": (1, 2, 4)} DEFAULT = {"tile": 32, "fuse": 0, "vec": 1, "shared": 0, "unroll": 1} def make_task(rng): """随机生成一个任务的隐藏特性。""" return {"m": rng.random(), "best_tile": rng.choice((16, 32)), "n": rng.choice((1024, 1000, 1002, 1003, 4096)), "fusable": rng.random() < 0.7} def run_kernel(task, cfg): """ "编译并运行"一个配置：返回 (状态, 运行时间或错误信息, 瓶颈)。 【返回】状态 ∈ {"ok", "compile_error", "wrong"}；ok 时第二项是运行时间（越小越快），否则是错误信息； 瓶颈 ∈ {"memory", "compute"}（只有 ok 时有意义） 【思路】运行时间 = (1 − m) × 计算部分 + m × 访存部分： 计算部分：tile 离最合适的值越远越慢；unroll 有帮助，但 tile 64 + unroll 4 会溢出寄存器； 访存部分：融合 × 0.55、共享内存 × 0.7（tile ≥ 16 时）、向量化 × (1 − 0.15 log2(vec))。 """ if cfg["tile"] == 64 and cfg["shared"]: return "compile_error", "共享内存超出容量（tile 64 时每个 block 需要的共享内存太大）", None if task["n"] % cfg["vec"]: return "wrong", f"在规模 n = {task['n']} 的输入上结果不对（{task['n']} 不能被向量宽度 {cfg['vec']} 整除）", None if cfg["fuse"] and not task["fusable"]: return "wrong", "融合之后结果不对（中间结果需要单独保存）", None compute = 1 + 0.15 * abs(math.log2(cfg["tile"] / task["best_tile"])) compute *= 1 - 0.08 * math.log2(cfg["unroll"]) if cfg["tile"] == 64 and cfg["unroll"] == 4: compute *= 1.25 memory = 1.0 if cfg["fuse"]: memory *= 0.55 if cfg["shared"] and cfg["tile"] >= 16: memory *= 0.7 memory *= 1 - 0.15 * math.log2(cfg["vec"]) t = (1 - task["m"]) * compute + task["m"] * memory bottleneck = "memory" if task["m"] * memory > (1 - task["m"]) * compute else "compute" return "ok", t, bottleneck def baseline_time(task): return 0.9 * run_kernel(task, DEFAULT)[1] def random_config(rng): """模型"一次写出来"的配置：偏好激进的优化（这是 KernelBench 观察到的现象：模型爱用花哨的技巧，常常出错）。""" return {"tile": rng.choice((16, 32, 64, 64)), "fuse": int(rng.random() < 0.6), "vec": rng.choice((1, 2, 4, 4)), "shared": int(rng.random() < 0.5), "unroll": rng.choice((1, 2, 4))} def mutate(cfg, rng, knob=None): """改一个旋钮（随机挑一个，或者指定一个）为另一个取值。""" new = dict(cfg) knob = knob or rng.choice(list(KNOBS)) new[knob] = rng.choice([v for v in KNOBS[knob] if v != cfg[knob]]) return new def propose(task_feedback, cfg, rng, mode): """ 根据上一轮的反馈提出下一个配置。 【参数】task_feedback：(状态, 信息, 瓶颈)；cfg：上一轮的配置；mode："scalar" / "explain" / "suggest" 【思路】 scalar ：不管发生了什么，随机改一个旋钮。 explain：失败时读错误信息，做针对性的修改；成功时随机改一个旋钮。 suggest：在 explain 的基础上，成功时按瓶颈挑旋钮（访存瓶颈 → fuse / vec / shared；计算瓶颈 → tile / unroll）。 """ status, info, bottleneck = task_feedback if mode in ("explain", "suggest") and status != "ok": new = dict(cfg) if "共享内存" in info: new["shared"] = 0 if rng.random() < 0.5 else new["shared"] new["tile"] = 32 if new["shared"] else new["tile"] elif "整除" in info: n = int(info.split("n = ")[1].split(" ")[0]) new["vec"] = max(v for v in KNOBS["vec"] if n % v == 0) elif "融合" in info: new["fuse"] = 0 return new if mode == "suggest" and status == "ok": knobs = ("fuse", "vec", "shared") if bottleneck == "memory" else ("tile", "unroll") return mutate(cfg, rng, rng.choice(knobs)) return mutate(cfg, rng) def optimize(task, rng, iterations, mode): """ 迭代优化 iterations 轮，返回找到的最快的正确配置的加速比（相对 PyTorch 基线；没有正确的配置时返回 0）。 【思路】第 1 轮：模型一次写出一个配置（random_config）。之后每轮：从"当前最好的配置"出发提议修改 （如果还没有成功过，就从上一轮的配置出发，这样才能根据错误信息修复它）。 mode = "resample"：每轮都独立地重新写一个配置（重复采样，不利用反馈）。 """ best_cfg, best_time = None, math.inf cfg = random_config(rng) for _ in range(iterations): fb = run_kernel(task, cfg) if fb[0] == "ok" and fb[1] < best_time: best_cfg, best_time = cfg, fb[1] if mode == "resample": cfg = random_config(rng) else: base = best_cfg if (best_cfg is not None and fb[0] == "ok") else cfg fb_for_base = run_kernel(task, base) if base is not cfg else fb cfg = propose(fb_for_base, base, rng, mode) return baseline_time(task) / best_time if best_cfg else 0.0 def fast_p(speedups, p): """fast_p：正确（加速比 > 0）并且加速比 > p 的任务所占的比例。""" return sum(1 for s in speedups if s > p) / len(speedups) def main(): rng = random.Random(230) tasks = [make_task(rng) for _ in range(200)] print("=" * 78) print("演示 1：一次写出来的 kernel（200 个任务，每个任务写一次）") print("=" * 78) outcomes, one_shot = [], [] for i, t in enumerate(tasks): status, info, _ = run_kernel(t, random_config(random.Random(231 + i))) outcomes.append(status) one_shot.append(baseline_time(t) / info if status == "ok" else 0.0) print(f" 编译 / 启动失败 {outcomes.count('compile_error') / len(tasks):.1%}，结果不对 {outcomes.count('wrong') / len(tasks):.1%}，" f"正确 {outcomes.count('ok') / len(tasks):.1%}") print(" " + "，".join(f"fast_{p} = {fast_p(one_shot, p):.3f}" for p in (0.0, 1.0, 1.5))) print(" （fast_0 = 正确率；fast_1 = 正确且比 PyTorch 快）") print("\n" + "=" * 78) print("演示 2：迭代优化 —— 不同的反馈，同样的轮数（fast_1 / fast_1.5）") print("=" * 78) modes = (("重复采样（不用反馈）", "resample"), ("只有标量反馈", "scalar"), ("+ 执行反馈（解释错误）", "explain"), ("+ 性能分析建议", "suggest")) rounds = (5, 10, 20, 50) print(pad_cell("反馈", 26) + "".join(pad_cell(f"{r} 轮", 16, right=True) for r in rounds)) for name, mode in modes: row = pad_cell(name, 26) for r in rounds: sp = [optimize(t, random.Random(1000 * r + i), r, mode) for i, t in enumerate(tasks)] row += pad_cell(f"{fast_p(sp, 1.0):.2f} / {fast_p(sp, 1.5):.2f}", 16, right=True) print(row) print("\n观察：") print(" · 一次写出来的 kernel 大多不对：模型爱用激进的优化（大 tile + 共享内存、宽向量化、融合），") print(" 每一种都可能在某些任务上出错。KernelBench 发现，前沿的 reasoning model 也只在不到 20% 的任务上") print(" 达到\"正确且不比 PyTorch 慢\"。") print(" · 只有标量反馈时，失败的配置只告诉你\"失败了\"。从一个失败的配置出发随机改一个旋钮，") print(" 常常还是失败 —— 在少数几轮里，它甚至不如每轮重新写一个（重复采样）。") print(" · 知道**为什么**失败（执行反馈），就能直接修好它：例如把向量宽度改成能整除问题规模的值。") print(" 同样的轮数下，它的 fast_1 始终明显高于标量反馈和重复采样；fast_1.5 在 10 轮以后也是最好的。") print(" · 性能分析的建议（瓶颈在访存还是计算）在前几轮最有用：10 轮时 fast_1 最高。但它把修改限制在") print(" 和瓶颈相关的旋钮上，轮数多了以后，在更高的加速要求（fast_1.5）上反而不如\"只解释错误\"——") print(" 有用的建议加快了早期的进步，也缩窄了搜索。好的反馈应该是\"建议\"，而不是\"命令\"。") print(" · 这是 Agent–System Interface 的核心发现：把原始的执行输出翻译成\"解释 + 建议\"，") print(" LLM 优化器用 10 轮就胜过了 OpenTuner 1000 轮的结果；KernelBench 里，用执行和性能分析的反馈做") print(" 10 轮迭代，DeepSeek-R1 在第 2 级任务上的 fast_1 从 36% 提高到 72%。") print(" · p 越大（要求越快），所有方法都越难 —— fast_p 随 p 下降的速度反映了\"优化的深度\"。\n") def pad_cell(text, width, right=False): """按显示宽度（汉字算 2 格）补空格。""" w = sum(2 if ord(ch) > 127 else 1 for ch in text) gap = " " * max(0, width - w) return gap + text if right else text + gap if __name__ == "__main__": main()