ForgeKernel with PyPTO

前两天看到了 OpenBMB 开源的 ForgeKernel，一个让 agent 自动优化算子的工作，在 FlashInfer-Bench（B200）、CANN-Bench (Ascend A3 & 950DT) 上都取得了第一名。这个进步还是很夸张的。笔者之前也参与过算子的竞赛活动，尝试过用 AI 自己迭代写算子，但是没能取得很好的成绩，看到这样的工作自然忍不住想学习一下。正好最近在看 PyPTO，就想能不能复现这个工作，使用 PyPTO 的后端。

需要说明的是，官方仓库目前还没有放出开源的 Ascend Agent 框架，本文中提到的实现都是笔者根据论文的描述以及已开源的 NVIDIA 实现模仿实现的，代码质量也是 vibe 质量，欢迎批评指正。

代码可以在这里查看：https://github.com/OpenBMB/ForgeKernel/pull/5

ForgeKernel 和普通 Agent 的区别

笔者之前一直在参加昇腾算子挑战赛，大概从 S8 赛季开始就全面投入了 AI 的怀抱，很少像之前一样手写算子了。不过，在最近结束的 S9 赛季上，笔者大概每道题花了 20+ 小时的迭代时间给 GPT-6 Astra，却并没有收获什么一骑绝尘的成绩，甚至在多数题上连前排水平都碰不到。一个很自然的想法是，目前的 Agent 也许还并不擅长写算子。

ForgeKernel 否定了这一点，他们认为，单个 /loop 或 /goal 会话对于算子性能的提升会在自己设定的上限处停止。agent 在最初的几次尝试中就确定了算子的架构，剩下的运行都是在既定架构上进行微调；而且，之后的每个新会话也都是在相同的知识基础上进行，上一个会话的经验会被丢弃，而不是继承到下一个会话。总的来说，就是“陷入局部最小值”和“无法积累长期知识”两个问题。

针对这两个问题，ForgeKernel 拿出了两套机制。

一套是架构层面的持续演化。要解决“局部最小值”的问题，ForgeKernel 提出的方案是定义两个 agent，一个 teacher 专门在顶层思考，用来跳出局部最值；另一个 student 专门在底层实现，用来接近局部最小值。具体到实现上，就是两个 claude code 实例。架构决策由一个专门的 Teacher 以里程碑的形式显式写出来，工作划分、数据流、流水组织这些牵一发动全身的决定，都在最早的里程碑里定下来。优化过程中一旦发现性能停滞、当前计划里的决策被用尽，系统就会停下来，让 Teacher 结合实现轨迹和实测数据重新评估架构，改动小就从相关的里程碑接着做，改动大就从最早受影响的决策开始重做。每一轮优化结束，通过验证的 kernel 会被封存为下一轮的参考实现，下一轮就在这个基础上继续调整架构。

另一套是跨会话的知识积累。ForgeKernel 准备了一个 hw_kb/ 目录，记录关于硬件机制的知识。这个目录下的每一条实现，都要先在目标硬件上做独立的探针实验，验证通过才会进知识库。这个知识库跨算子、跨输入形状共享，之后的优化里都能用。这样，上一轮跑出来的最好结果就是下一轮的起跑线，攒下来的经验也不会随着会话结束被丢掉。

Teacher 收到的信息和做出的决策

按照论文的描述，每一轮优化开始，Teacher 手里有四样东西:

• 算子定义和输入配置
• 目标硬件
• 知识库里的现有条目
• 一个密封的参考实现。这个参考实现是个黑盒，可以运行、计时、profile、反汇编，但是不给源码。Teacher 能用的是它的实测性能、资源占用和执行特征，这些构成后面定目标的参考证据。

拿到这些信息后，Teacher 的决策分两步。第一步是定架构，包括 kernel 分解、数据依赖、中间结果的存放、tiling 形状、线程角色，还有寄存器和共享内存的预算。第二步是把这些决策排成有序的里程碑序列 M0 到 Mk，M0 只要求把关键结构做对，不设性能门槛，后面的里程碑才在这个结构里逐个调优。

Student 收到的信息和做出的工作

轮到 Student，能看的东西就被刻意卡得很死。

• 里程碑计划里写的设计要求、接口约定和验收条件
• 计划中显式引用的知识库条目
• harness 对每个候选版本的检查反馈

参考实现的源码明确不给，harness 会强制守住这条边界，直接调用参考实现或者把计算外包给它的候选都会被拒掉。

Student 的工作其实和传统的 /goal 实现比较类似，但是它收到的约束比 /goal 更多。在计划的约束之内，它可以自由做局部优化，觉得非得动架构才有出路时，它没有权限自己改，只能把问题交回去。

每写完一个候选，harness 先查正确性，再把候选和参考实现放进同一个测量任务里配对计时，顺带收集机制验证要看的执行特征和资源占用，包括正确性、性能、机制验证。通过了就进下一个里程碑，通不过就拿着检查反馈继续改。实在改不动，或者验收怎么走都走不下去，系统会把整个实现轨迹和实测数据打包交回 Teacher，进入上一节说的修订流程。

适配 PyPTO harness 要做的工作

了解了两个 agent 主要的工作内容，适配新后端要做的事就比较清晰了，类似于实现一个算子 Judge 系统，对于 Student 实现的每个候选，运行时要返回：

• 正确性检查结果，过了还是挂了，挂在哪个 case 上
• 配对计时数据，候选和参考实现在同一个测量任务里跑出来的耗时
• profiling 数据，机制验证要看的执行特征和资源占用

不过这里，还是有很多细节要和论文对齐：

• 正确性判定不能随手写个 allclose 了事，要用算子来源对应的官方容差测试，FlashInfer-Bench 收录的算子用它的判定，FlashAttention 用它官方的判定。直接调用参考实现或者把计算外包给它的候选，harness 要直接拒掉
• 计时协议要统一。论文在 NVIDIA 上测的是 cold L2 下的 device time，重复多次取平均，候选和参考共用同一套协议和 cache 条件。搬到 Ascend 上，device time 要从 msprof 的 kernel_details 里拿
• 性能提升要超过测量噪声阈值才算通过
• harness 还要支持机制验证，检查计划里指定的执行特征或资源占用是否出现了预期的变化。这一条在普通 Judge 系统里一般没有。

一阶段实测结果

一阶段在两个算子上跑了完整的 campaign，平台是 Ascend 950PR、CANN 9.2、PyPTO 0.3.0，评测走 CANN-Bench 2.0.0 官方 evaluator，teacher 和 student 都用 deepseek-flash（1M context）。结果一成一败。

SwiGlu 的收益几乎全部来自两个尾部 case。20 个 case 取几何平均，多数 case 本来就在 30 到 150 微秒之间，swi_glu_12 一个就要约 484 微秒，swi_glu_20 要约 269 微秒，总成绩基本被这两个拽着走。agent 的处置有两手。一是给最后一维只有 2 的 shape 加了条专用路径，host 侧常数时间就能判断要不要走，路径里用 deinterleave 把交错的两列当连续块读入再拆开，循环次数从约 3907 降到约 123。二是把 c8 类的 tile 从 (32, 50) 放大到 (128, 50)，unroll 从 [1] 加到 [4]，trip 数从约 1632 降到约 408。两条路径的参数都是联合 probe 比出来的，比如 THIN_T0 在 1024、4096、8192、16384 里选中 8192。probe 只负责找方向，最终成绩由官方 20-case 的 3+3 确认决定。

GQA 的候选 20/20 通过官方正确性，数值和分组逻辑确实写对了，但几何平均 40.24 ms 对参照的 0.06116 ms，慢了约 658 倍，最慢的 gqa_8 差出一万三千多倍。这个量级换 tile 参数救不回来，要动的是 prefill 和 decode 分流、双缓冲、跨阶段重叠这一层，属于架构级改动。这轮 campaign 也不是全无收获，过程中认证了 pypto-operand-tile-construction 和 pypto-dynamic-execution 两条 NPU 接口 KB，还抓出过一次计时口径不一致造成的无效 gate。最重要的是系统认输了，没有把 20/20 正确性误报成优化成功。

PyPTO 的编译时长

最后记一个论文之外、实际接 PyPTO 时躲不开的问题，编译时长。同样的迭代在 CUDA 侧，nvcc 加 CUTLASS 头文件大概半分钟就见分晓。PyPTO 是 JIT，而且按 shape 特化，20 个 case 意味着 20 份编译。实测一次全量评测要几分钟到几十分钟，GQA 里有的 probe 单次编译就要几十分钟，逼近 7200 秒的 job 上限。几次 probe 或编译失败，一个外层循环的预算就没了。

这个成本直接改写了打法。probe 和正式评测必须分层，probe 只覆盖单个合同，负责缩小范围，判决一律交给官方 20-case。参数不能孤立地扫，tile 和 unroll 要放进同一轮联合 probe 里比，否则很容易在一个错误的循环结构上白烧几小时。架构探索和全量 gate 评测的成本要分开记账。回头看，GQA 停在 M0，一半原因是架构差距，另一半就是试不起。