""" ============================================================================= CS349E · 第 15 讲 · 01_agent_loop.py —— 智能体的工具循环：一个有 bug 的内核、六种工具、一条 ReAct 式的轨迹 ============================================================================= 【这个脚本要回答的问题】 一个"写内核"的智能体，交出去的分块矩阵乘法内核有 bug。它能调用的工具各告诉它什么？它怎样一步步找到 bug、修好、确认？ 演示 1：七种 bug × 六种工具：每种工具在三种形状上看到什么（哪个工具能"看见"哪个 bug）。 演示 2：一条 ReAct 式的轨迹（Thought / Action / Observation）：按"每秒信息量"挑工具，找到 bug、修好、复查。 演示 3：同一个 bug，在三种"反馈接口"下智能体看到的东西（只有对错 / 加上错误模式 / 加上 sanitizer）。 【关于数字】bug 的现象来自真实的 numpy 计算（agent_lab.tiled_gemm 逐块模拟 GPU 内核）；工具的耗时是假设值（见 agent_lab.TOOLS）； "智能体"的推理是一个按信息量挑工具的规则（不是语言模型），修 bug 的办法查表（agent_lab.BUGS）。 【运行】python3 01_agent_loop.py （约 2 秒） 【依赖】numpy；本目录的 agent_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 agent_lab import (ACTIONS, BUGS, BUG_NAMES, SHAPES, SHAPE_NOTES, TOOLS, info_gain, likelihood_table, # noqa: E402 observe, posterior) from console import print_header, print_observations, print_table # noqa: E402 COMPILE_S = 25.0 # 每次改代码后重新编译一次（假设） def short(obs, tool): """把观察值写成表格里的短字符串。""" if tool == "test": return "FAIL" if obs else "pass" if tool in ("test_verbose", "test_identity"): return "pass" if not obs[0] else f"{obs[1]}/{obs[2]}" if tool == "run_twice": return obs if tool == "memcheck": return "越界" + "、".join("读" if "read" in e else "写" for e in obs) if obs else "—" if tool == "racecheck": return "竞争" if obs else "—" return str(obs) def demo1_matrix(): print_header("演示 1：七种 bug × 六种工具——每种工具在三种形状上看到什么") print("形状：" + "；".join(f"{k} = {v}" for k, v in SHAPE_NOTES.items())) print("工具：" + "；".join(f"{k}（{v[0]:.0f} s）：{v[1]}" for k, v in TOOLS.items())) print() cols = [("test", s) for s in SHAPES] + [("test_verbose", s) for s in SHAPES] + \ [("test_identity", "nice"), ("run_twice", "edge"), ("memcheck", "edge"), ("racecheck", "nice")] rows = [] for b in BUG_NAMES: rows.append([b] + [short(observe(t, s, b, run=0)[0], t) for t, s in cols]) print_table(["bug"] + [f"{t.replace('test_verbose', 'verbose').replace('test_identity', 'identity')}·{s}" for t, s in cols], rows) print("verbose 的格式：错误模式/误差量级。模式：all（> 90% 的元素错）、tail（只有最后不满的块行 / 块列错）、" "head_cols（每行开头几列错）、full_cols（若干整列错）、scattered（其他）；量级：small / large（最大误差相对参考值的标准差）；" "identity 用的输入是 B = 单位矩阵") print("\nbug 说明：") for b in BUG_NAMES[1:]: print(f" {b}：{BUGS[b][0]}") print_observations([ "\"整齐\"的形状（128³，都是 32 的倍数）上，三种和\"尾巴\"有关的 bug（K 没加掩码、网格向下取整、写越界）完全看不出来——" "而很多人（和模型）写测试时正好只用这种形状。换成 100 × 70 × 90，它们都暴露了。", "只说 PASS / FAIL 的测试分不开大部分 bug：在 edge 形状上，除了 missing_sync 有时碰巧通过，其余 6 个都是 FAIL。" "\"详细测试\"给出错误的模式：只错最后几行几列（网格没盖住）、每行开头几列被写成 0（越界写进了下一行）、几乎全错（K 循环或下标的错）、" "零星的小误差（FP16 累加）。", "有的 bug 要专门的工具：跑两次结果不一样，说明有非确定性（竞争或读到了未初始化的内存）；memcheck 只看到越界读写；" "racecheck 一定能抓住 missing_sync，但它最贵（这里假设 120 秒）。", ]) def demo2_trace(table): print_header("演示 2：一条 ReAct 式的轨迹——bug 是 store_overflow，按\"每秒信息量\"挑工具") true_bug = "store_overflow" prior = {b: 1 / len(BUG_NAMES) for b in BUG_NAMES} rng = np.random.default_rng(3) t_total = COMPILE_S print(f"[0] Action: compile(kernel.cu) → Observation: 编译通过（{COMPILE_S:.0f} s）") step = 0 after = [] while max(prior.values()) < 0.99 and step < 8: step += 1 gains = {a: info_gain(prior, table, a) for a in ACTIONS} a = max(ACTIONS, key=lambda a: gains[a] / TOOLS[a[0]][0]) cands = [b for b, p in sorted(prior.items(), key=lambda x: -x[1]) if p > 0.01] print(f"[{step}] Thought: 还剩 {len(cands)} 个可能的 bug（{', '.join(cands)}）；{a[0]}·{a[1]} 的期望信息量 " f"{gains[a]:.2f} 比特、耗时 {TOOLS[a[0]][0]:.0f} s，每秒 {gains[a] / TOOLS[a[0]][0]:.3f} 比特，最划算。") o, text = observe(a[0], a[1], true_bug, int(rng.integers(1_000_000))) print(f" Action: {a[0]}(shape={a[1]})") print(f" Observation: {text}") prior = posterior(prior, table, a, o) after.append([b for b, p in prior.items() if p > 0.01]) t_total += TOOLS[a[0]][0] best = max(prior, key=prior.get) print(f"[{step + 1}] Thought: 后验里 {best} 的概率 {prior[best]:.3f}。修法：{BUGS[best][1]}。") print(f" Action: edit(kernel.cu) + compile → Observation: 编译通过（{COMPILE_S:.0f} s）") t_total += COMPILE_S checks = [] for a in (("test", "nice"), ("test", "edge"), ("test", "bigK"), ("memcheck", "edge")): o, text = observe(a[0], a[1], "none", 0) checks.append(f"{a[0]}·{a[1]}：{text}") t_total += TOOLS[a[0]][0] print(f"[{step + 2}] Action: 复查 → Observation: " + "；".join(checks)) print(f"\n一共调用了 {step} 次诊断工具 + 2 次编译 + 4 次复查，假设的工具时间合计 {t_total:.0f} 秒。") print_observations([ "ReAct（Yao et al. 2022）的格式：思考（为什么做下一步）和动作（调用工具）交替出现，环境的观察写回上下文。" "这里的\"思考\"是一个固定的规则：维护一个\"可能是哪个 bug\"的概率分布，每一步挑\"每秒能带来最多信息\"的工具。", "第一步它选了便宜的详细测试，而且选了不整齐的形状——在整齐的形状上一半的 bug 根本不出现。" f"这一个观察（\"每行开头几列错\"）就把候选从 8 个缩到 {len(after[0])} 个（{'、'.join(after[0])}）" + ("：竞争偶尔也会碰巧只弄错开头几列，所以它又花一次便宜的测试把它排除了。" if len(after[0]) > 1 else "。") + f"一共 {step} 次诊断。", "修完之后的复查同样重要：三种形状都通过、memcheck 也干净，才算修好（第 14 讲：测试要覆盖不同的形状和输入）。", ]) def demo3_interfaces(): print_header("演示 3：同一个 bug（grid_floor），三种反馈接口下智能体看到的东西") rows = [] b = "grid_floor" rows.append(["只有对错（KernelBench 的 \"G\"：只看自己上一次的代码）", observe("test", "nice", b)[1] + "（nice）；" + observe("test", "edge", b)[1] + "（edge）"]) rows.append(["加上错误模式（\"G + E\"：执行结果）", observe("test_verbose", "edge", b)[1]]) rows.append(["加上工具（run_twice、memcheck）", observe("run_twice", "edge", b)[1] + "；memcheck：" + observe("memcheck", "edge", b)[1]]) print_table(["接口", "智能体看到的"], rows, aligns=["<", "<"]) print_observations([ "只有对错时，智能体只知道\"在不整齐的形状上错了\"——七种 bug 里有六种都会这样。错误模式（只有第 96–99 行、第 64–69 列错，误差很大）" "几乎直接指向\"最后不满的块没人算\"；跑两次结果不同，说明那里是未初始化的内存。", "KernelBench 的迭代优化实验里，把执行结果（E）和 profiler（P）喂回去，DeepSeek-R1 在 Level 2 上的 fast_1 从 36% 升到 72%；" "剩下的失败几乎都是\"结果不对\"——因为它们拿到的正确性反馈只是\"不等于参考\"，远没有编译错误信息具体。" "把\"不对\"变成\"哪里不对、错多少、什么模式\"，是工具接口设计里最便宜的一步。", ]) if __name__ == "__main__": demo1_matrix() tab = likelihood_table() demo2_trace(tab) demo3_interfaces()