""" ============================================================================= CS349E · 第 14 讲 · 05_alphaevolve_ablations.py —— 迷你 AlphaEvolve 的消融：每个部件值多少 ============================================================================= 【这个脚本要回答的问题】 AlphaEvolve 论文的消融实验（图 8）去掉一个部件看效果掉多少：No evolution、No context in the prompt、 Small base LLM only、No full-file evolution、No meta prompt evolution。我们在迷你版上做同样的事： 演示 1：9 个变体 × 5 个随机种子，分数随"提出的程序数"怎样上升，测试形状上的分数，评估成本。 演示 2：反馈怎么用——同样是"看评估反馈"，挑哪个形状下手、拿什么当参照，结果差别很大。 演示 3：不检查所有分支会怎样——数据库里藏着多少"测试没走到的、根本不能启动的分支"。 【变体】（"LLM"都是 evolve.MockLLM，没有真正的语言模型） full ：完整的方法（内核常识的定向修改 + 按"离见过的最快还差多少"加权挑形状 + 借规则 + 4 个岛 + 级联 + 分支检查） fb_roofline ：定向修改按"离 roofline 还差多少"挑形状 fb_worst3 ：定向修改只盯着"差得最多的 3 个形状" no_feedback ：定向修改随便挑形状（不看评估结果） small_llm ：没有内核常识，只会单字段的小改动（"Small base LLM only"） no_evolution ：每次都从初始程序出发改（"No evolution"） no_borrow ：不从别的程序借规则 one_island ：只有 1 个岛 no_cascade ：不做级联，每个程序都完整评估 【关于数字】"硬件"是第 13 讲的解析模型；所有结果确定（固定随机种子），每次运行完全相同。 【运行】python3 05_alphaevolve_ablations.py （约 50 秒） 【依赖】numpy；本目录的 evolve；../lecture01 的 console。 """ import math 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 evolve import (Evaluator, MockLLM, Shape, LAYERS, TimeTable, all_shapes, choose, evolve, # noqa: E402 geomean_speedup, oracle_speedup, render, simplify, split_shapes, static_check_program) from console import print_header, print_observations, print_table # noqa: E402 SEEDS = 5 BUDGET = 1000 MARKS = (50, 100, 300, 1000) VARIANTS = { "full": {}, "fb_roofline": {"llm": MockLLM(feedback="roofline")}, "fb_worst3": {"llm": MockLLM(focus="worst3")}, "no_feedback": {"llm": MockLLM(feedback=None)}, "small_llm": {"llm": MockLLM(knowledge=False)}, "no_evolution": {"evolution": False}, "no_borrow": {"llm": MockLLM(use_inspirations=False)}, "one_island": {"n_islands": 1}, "no_cascade": {"cascade": False}, } def run_variant(kw, train, test, table, check=True): """跑 SEEDS 次，返回每次的 (历史分数, 测试分数, 跑内核次数, 简化后的规则数, evolve 的输出)。""" outs = [] for seed in range(SEEDS): ev = Evaluator(train, table, check_all_branches=check) out = evolve(ev, BUDGET, seed=seed, **kw) bp, _ = out["best"] outs.append((out["history"], geomean_speedup(bp, test, table), ev.bench_runs, simplify(bp).n_rules(), out)) return outs def demo1_ablations(table, train, test): print_header(f"演示 1：消融——{len(VARIANTS)} 个变体 × {SEEDS} 个种子（中位数），分数 = 训练形状上相对专家的几何平均加速比") res = {} rows = [] for name, kw in VARIANTS.items(): outs = run_variant(kw, train, test, table) res[name] = outs hist = np.array([o[0] for o in outs]) cells = [f"{np.median(hist[:, m - 1]):.3f}" for m in MARKS] cells.append(f"{hist[:, -1].min():.3f}–{hist[:, -1].max():.3f}") rows.append([name] + cells + [f"{np.median([o[1] for o in outs]):.3f}", f"{int(np.median([o[2] for o in outs])):,}", f"{np.median([o[3] for o in outs]):.0f}"]) print_table(["变体"] + [f"{m} 个程序后" for m in MARKS] + ["1000 个后（最小–最大）", "测试形状", "跑内核次数", "规则数（化简后）"], rows) print(f"上限（逐形状最优）：训练 {oracle_speedup(train, table):.3f}，测试 {oracle_speedup(test, table):.3f}") med = {k: float(np.median([o[0][-1] for o in v])) for k, v in res.items()} med100 = {k: float(np.median([o[0][99] for o in v])) for k, v in res.items()} print_observations([ f"去掉\"进化\"（每次都从初始程序出发）掉得最多：{BUDGET} 个程序后只有 {med['no_evolution']:.3f}，完整方法 {med['full']:.3f}。" "一次修改只能走一小步，好程序要靠在好程序上继续改才能积累出来——AlphaEvolve 的消融里 \"No evolution\" 同样是最差的之一。", f"去掉\"内核常识\"（small_llm，只会单字段小改动）：{med['small_llm']:.3f}。很多好配置要几个字段一起改才到得了（04 脚本演示 2），" "小改动走不过去。AlphaEvolve 的 \"Small base LLM only\" 消融也明显变差，论文的解释是更强的模型能提出更好的修改。", f"反馈怎么用影响很大：按\"离见过的最快还差多少\"加权挑形状（full）{med['full']:.3f}，按\"离 roofline 还差多少\"加权挑 {med['fb_roofline']:.3f}，" f"不看反馈随便挑 {med['no_feedback']:.3f}，只盯最差的 3 个 {med['fb_worst3']:.3f}" + ("——只盯最差的 3 个比随便挑还差。" if med['fb_worst3'] < med['no_feedback'] else "。") + "演示 2 看它们各自指向哪些形状。", f"借规则（no_borrow {med['no_borrow']:.3f}）、岛屿数（one_island {med['one_island']:.3f}）、级联（no_cascade {med['no_cascade']:.3f}）" f"和完整方法（{med['full']:.3f}）的差别在种子之间的波动以内（看\"最小–最大\"那一列），有的变体的中位数还略高。" "级联只省下一小部分评估（这个问题上坏程序不多）。这些部件的价值在更大、更难、评估更贵的问题上（多样性、并行、省钱）才显出来——" "一个小实验得不出\"它们没用\"的结论；反过来，5 个种子的中位数差 0.02 也不能说明谁更好。", f"100 个程序时的分数：full {med100['full']:.3f}、no_feedback {med100['no_feedback']:.3f}、small_llm {med100['small_llm']:.3f}、" f"no_evolution {med100['no_evolution']:.3f}——早期差别小、后期差别大；比较搜索方法一定要看\"随预算变化的曲线\"，不能只看一个点。", ]) return res def demo2_feedback(table, train): print_header("演示 2：反馈怎么用——\"离 roofline 多远\"和\"离见过的最快多远\"指向的形状不一样") ev = Evaluator(train, table) out = evolve(ev, 300, seed=0) bp, _ = out["best"] fb = ev.per_shape(bp) by_roof = sorted(fb, key=lambda x: x["vs_roofline"])[:5] by_best = sorted(fb, key=lambda x: x["vs_best"])[:5] rows = [] for a, b in zip(by_roof, by_best): rows.append([a["shape"].label(), f"{a['vs_roofline']:.2f}", f"{a['vs_best']:.2f}", b["shape"].label(), f"{b['vs_roofline']:.2f}", f"{b['vs_best']:.2f}"]) print("（300 个程序之后的最好程序；vs_roofline = roofline 下界 / 时间，vs_best = 见过的最快 / 时间，都是越小越\"该改\"）") print_table(["按 roofline 最该改的", "vs_roofline", "vs_best", "按\"见过的最快\"最该改的", "vs_roofline", "vs_best"], rows) stuck = [x for x in by_roof if x["vs_best"] > 0.95] overlap = len({x["shape"] for x in by_roof} & {x["shape"] for x in by_best}) def mass_on_stuck(key): """按 MockLLM 的加权方式（权重 = 1/比值 − 1 + 0.05），提案落在"没有已知改进"的形状（vs_best > 0.95）上的比例。""" w = np.array([1.0 / x[key] - 1.0 + 0.05 for x in fb]) return float(w[[x["vs_best"] > 0.95 for x in fb]].sum() / w.sum()) n_stuck = sum(1 for x in fb if x["vs_best"] > 0.95) m_roof, m_best = mass_on_stuck("vs_roofline"), mass_on_stuck("vs_best") print(f"训练形状里有 {n_stuck}/{len(fb)} 个\"没有已知改进\"（vs_best > 0.95）。按权重抽形状时，提案落在它们身上的比例：" f"按 roofline {m_roof:.0%}，按见过的最快 {m_best:.0%}") print_observations([ f"两种排序的前 5 个有 {overlap} 个相同——差得多的形状，往往两种看法都差。区别在\"已经做到见过的最好\"的那些形状：按 roofline 加权时，" f"{m_roof:.0%} 的提案落在它们身上（例如 " + ("、".join(x["shape"].label() for x in stuck) if stuck else "M 很小的 decode 形状") + "：离 roofline 远是因为启动开销、补齐、块太少占不满 132 个 SM，换配置也改不了多少）；按见过的最快加权时只有 " f"{m_best:.0%}。提案花在\"还不知道怎么改、甚至根本改不动\"的地方，命中率就低。", "按\"见过的最快\"排，指向的是**确实有更快的配置、只是当前程序没用上**的形状，而且直接给出了那个配置——这个信息来自评估器已经做过的测量，" "不花额外的成本。KernelBench 的迭代优化里，profiler 的输出比单纯的\"对 / 错 + 时间\"更有用，是同一个道理：" "反馈要**可操作**，而不只是\"离理想有多远\"。", "只盯\"最差的 3 个\"会反复改同几个形状、失去多样性；按差距**加权随机**挑，既偏向该改的，也给其他形状机会。" "真正的 LLM 读反馈时没有这样一条固定的规则，它怎么用反馈是 AlphaEvolve 能否成功的关键，也是我们的模拟最不像的地方。", ]) def demo3_branches(table, train, test): print_header("演示 3：不检查所有分支——数据库里藏着多少\"测试没走到的、不能启动的分支\"") extra = [Shape(n, M, N, K) for n, (N, K) in list(LAYERS.items()) + [("405B/8 qkv", (2560, 16384)), ("405B/8 down", (16384, 6656))] for M in (2, 8, 32, 96, 192, 384, 768, 1536, 3072, 6144)] rows = [] for check in (True, False): outs = run_variant({}, train, test, table, check=check) for seed, o in enumerate(outs): out = o[4] progs = [p for pop in out["db"].islands for p, _ in pop] latent = sum(1 for p in progs if not static_check_program(render(p))[0]) crash = sum(1 for p in progs if any(not math.isfinite(table.time(choose(p, s), s)) for s in extra)) bp, _ = out["best"] rows.append(["开" if check else "关", seed, out["n_added"], out["n_added_bad_branch"], len(progs), latent, crash, "是" if not static_check_program(render(bp))[0] else "否"]) print_table(["分支检查", "种子", "进数据库的程序", "其中带不合法分支的", "最终数据库里的程序", "其中带不合法分支的", "在 100 个新形状上出错的", "最好的程序带不合法分支？"], rows) off = [r for r in rows if r[0] == "关"] print_observations([ f"不做分支检查时，进数据库的程序里有 {sum(r[3] for r in off)} 个（5 个种子合计）至少有一个分支返回不能启动的配置——" "训练形状从来没走到那个分支，所以分数看不出来。最终的数据库里仍有 " f"{sum(r[5] for r in off)} 个这样的程序，{sum(1 for r in off if r[7] == '是')} 个种子的\"最好程序\"本身就带着一个。", f"在 100 个新形状上（新的 M、405B 分片的 N、K）真正出错的有 {sum(r[6] for r in off)} 个：隐患大多潜伏着，" "直到某个没见过的形状碰巧走进那个分支。这正是\"测试覆盖不到的分支\"在生产里出事的方式。", "对策很便宜：对程序的**每个**分支做静态检查（这里是 static_valid：共享内存、寄存器、warp 的块）。" "AlphaEvolve 的 TPU 电路优化要求\"通过严格的验证\"，FlashAttention 的 IR 优化最后由人确认\"对所有输入都正确\"——" "评估器测的是样本，部署面对的是全体。", ]) if __name__ == "__main__": shapes = all_shapes() table = TimeTable() for s in shapes: table.oracle(s) train, test = split_shapes(shapes) demo1_ablations(table, train, test) demo2_feedback(table, train) demo3_branches(table, train, test)