""" ============================================================================= CS349E · 第 13 讲 · 04_autotuning.py —— 自动调优：搜索策略与学习的代价模型 ============================================================================= 【这个脚本要回答的问题】 03 脚本的配置空间里，好的配置很稀少。真实的调优每"测一次"都要编译、运行一个内核（几秒钟）， 预算有限。怎样用尽量少的测量找到接近最优的配置？ 演示 1：学习的代价模型：用测过的配置训练一个回归模型，它对没测过的配置排序有多准？换一个形状还准吗？ 演示 2：三种搜索策略（随机、进化、代价模型引导）在不同的测量预算下找到的配置有多好（20 次重复，测量有 2% 的噪声）。 演示 3：要多少次测量才能到最优的 5% 以内。 演示 4：调优在真实系统里的成本账，以及几种降低成本的办法。 【关于数字】"硬件"是 sched_lab.gemm_time 的解析模型（不是实测，见 03 脚本的说明）；测量噪声是人为加的。 这里研究的是"搜索"本身：同一个代价函数下，几种策略谁用的测量少。 【运行】python3 04_autotuning.py （约 7 秒） 【依赖】numpy；本目录的 sched_lab；../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 sched_lab import RidgeModel, all_configs, gemm_time, static_valid, tune # noqa: E402 from console import print_header, print_observations, print_table # noqa: E402 BIG = (4096, 4096, 4096) OTHER = [((512, 14336, 4096), "prefill 512 × up 投影"), ((16, 4096, 4096), "decode 批 16")] POOL = [c for c in all_configs() if static_valid(c)] SEEDS = 20 BUDGETS = (16, 32, 64, 128, 256) def spearman(a, b): """Spearman 秩相关：两组数的排序有多一致（1 = 完全一致，0 = 无关）。""" ra = np.argsort(np.argsort(a)) rb = np.argsort(np.argsort(b)) return float(np.corrcoef(ra, rb)[0, 1]) def demo1_cost_model(): print_header("演示 1：学习的代价模型（岭回归，特征 = 各参数的 log2 及其两两乘积）") truth = {s: {c: gemm_time(c, *s) for c in POOL} for s in [BIG] + [s for s, _ in OTHER]} rng = np.random.default_rng(0) perm = rng.permutation(len(POOL)) test = [POOL[i] for i in perm[:1000]] rows = [] for n in (16, 32, 64, 128, 256, 512): tr = [POOL[i] for i in perm[1000:1000 + n]] m = RidgeModel().fit(tr, [truth[BIG][c] for c in tr]) pred = m.predict(test) r = [n, f"{spearman(pred, [truth[BIG][c] for c in test]):.3f}"] top = [test[i] for i in np.argsort(pred)[:10]] r.append(f"{min(truth[BIG][c] for c in top) / min(truth[BIG].values()):.2f}") for s, _ in OTHER: r.append(f"{spearman(pred, [truth[s][c] for c in test]):.3f}") rows.append(r) print_table(["训练用的测量数", "排序相关（4096³，1000 个没见过的配置）", "模型挑的前 10 个里最好的 / 最优", f"同一个模型用在 {OTHER[0][1]}", f"用在 {OTHER[1][1]}"], rows) allc = all_configs() mixed = [] for n in (64, 256): tr = [allc[i] for i in np.random.default_rng(1).permutation(len(allc))[:n]] m = RidgeModel().fit(tr, [gemm_time(c, *BIG) for c in tr]) mixed.append(f"{n} 个 → {spearman(m.predict(test), [truth[BIG][c] for c in test]):.2f}") print("对照：训练集里混进不能启动的配置（时间 = ∞，目标值记为见过的最慢 × 2）时的排序相关：" + "，".join(mixed)) print_observations([ "几十个测量就能训练出一个排序相当准的模型（32 个时 Spearman 0.88，64 个以上 0.95 以上）：调度参数和时间之间的关系虽然不光滑，但有很强的结构" "（块越大越好、直到资源不够），对数特征加二次项就能抓住大部分。", "模型最有用的地方是\"挑候选\"：让它在几千个没测过的配置里选出最有希望的几个，再去真测——这正是 AutoTVM（XGBoost）和 Ansor（梯度提升树）的做法。", "换一个形状：相近的形状（prefill 的 up 投影）上模型还能用（0.93），换到受限于带宽的 decode 形状就几乎没用了（0.2 左右）——" "时间和参数的关系随形状而变。TVM 的 Ansor / MetaSchedule 用" "\"从程序里抽出来的特征\"（循环长度、访存量、向量化宽度……）训练一个跨形状、跨算子的模型，来缓解这个问题。", "我们只用能静态通过检查的配置训练；如果把不能启动的配置也混进来（上面的对照），同样的模型排序相关掉到很低——先剪枝再学习很重要。", ]) def run_strategies(shape): truth = {c: gemm_time(c, *shape) for c in POOL} opt = min(truth.values()) res = {} for strat in ("random", "evolution", "model"): curves = [] for seed in range(SEEDS): r = np.random.default_rng(1000 + seed) hist = tune(strat, lambda c: gemm_time(c, *shape, noise=0.02, rng=r), max(BUDGETS), np.random.default_rng(seed), pool=POOL) curves.append([truth[c] / opt for c in hist]) res[strat] = np.array(curves) return res def demo2_search(res): print_header(f"演示 2：搜索策略（4096³；{SEEDS} 次重复的中位数 / 最差的 10%；测量有 2% 的对数正态噪声）") names = {"random": "随机搜索", "evolution": "进化搜索（变异最好的几个）", "model": "代价模型引导（岭回归，每轮测 8 个）"} rows = [] for strat, curves in res.items(): rows.append([names[strat]] + [f"{np.median(curves[:, b - 1]):.2f} / {np.percentile(curves[:, b - 1], 90):.2f}" for b in BUDGETS]) print_table(["策略（找到的配置的真实时间 / 最优）"] + [f"{b} 次测量" for b in BUDGETS], rows) print_observations([ f"随机搜索：3744 个配置里只有十几个在最优的 5% 以内，随机抽 {BUDGETS[-1]} 个才大概率碰到一个；预算小时差得很远。", "进化搜索利用了\"好的配置附近还有好的配置\"：从目前最好的几个出发做小的变异，很快爬到最优附近。", "代价模型引导最省：前 16 次随机测量训练出的模型已经能把搜索引向好的区域，32 次测量时中位数就到了最优。" "代价是每轮都要训练模型、给上千个候选打分——和\"编译 + 运行一个内核要几秒\"相比，这点计算可以忽略。", "表里的 1.00 是保留两位小数的结果：最优附近有好几个只差 1%–2% 的配置，2% 的测量噪声足以让调优器在它们之间选错。" "真实的调优会重复测量、取中位数（第 14 讲的评测也要处理同一个问题）。", ]) def demo3_trials(res): print_header("演示 3：到最优的 5% 以内要几次测量（20 次重复）") rows = [] for strat, curves in res.items(): hits = [] for c in curves: idx = np.flatnonzero(c <= 1.05) hits.append(idx[0] + 1 if len(idx) else math.inf) hits = np.array(hits) fin = hits[np.isfinite(hits)] rows.append([strat, f"{np.median(hits):.0f}" if np.isfinite(np.median(hits)) else f"> {max(BUDGETS)}", f"{int(np.isfinite(hits).sum())} / {len(hits)}", f"{fin.min():.0f}–{fin.max():.0f}" if len(fin) else "—"]) print_table(["策略", "中位数", f"{max(BUDGETS)} 次以内做到的重复数", "做到的那些的范围"], rows) print_observations([ "同样到 5% 以内：代价模型引导大约要几十次，进化搜索多一些，随机搜索在 256 次以内常常做不到。", "如果一次测量（编译 + 预热 + 计时）要 3 秒，256 次就是 13 分钟——这只是一个形状、一个算子。差几倍的测量次数，就是几十分钟到几小时的差别。", ]) def demo4_cost(): print_header("演示 4：调优的成本账（一个推理服务，估算）") rows = [] shapes = 4 * 16 + 4 * 8 for name, trials in (("随机搜索（到 5% 以内）", 256), ("代价模型引导", 48), ("只用启发式 / 预先调好的表", 0)): t = shapes * trials * 3 rows.append([name, trials, shapes, f"{t / 3600:.1f} 小时" if t else "0"]) print("假设：每层 4 种矩阵乘法（QKV、O、gate/up、down），decode 的批大小按 16 个桶、prefill 的长度按 8 个桶；每次测量 3 秒。") print_table(["做法", "每个形状测量几次", "形状数", "总调优时间"], rows) print_observations([ "推理服务的形状很多：每层几种矩阵乘法 × 很多种批大小 / 长度。全部离线调一遍，随机搜索要将近一天，模型引导要几个小时。", "常见的降低成本的办法：(1) 缓存——调好的结果按 (算子, 形状, GPU) 存起来，Triton 的 @triton.autotune 按 key 缓存，TVM 有调优日志；" "(2) 按桶调优，实际形状向上补齐（第 12 讲 §8）；(3) 只在热点上调优——Ansor 的任务调度器按\"这个子图占总时间多少、还有多大提升空间\"分配预算；" "(4) 启发式和预先调好的表——cuBLAS / cuBLASLt 内置了按形状选内核的启发式，Triton 的 matmul 也有一组默认配置。", "第 14 讲会看到另一种思路：让 LLM 提出新的程序（不只是在固定的参数空间里选），用同样的\"测量 + 搜索\"闭环去评估（AlphaEvolve）。", ]) if __name__ == "__main__": demo1_cost_model() res = run_strategies(BIG) demo2_search(res) demo3_trials(res) demo4_cost()