让 AI 参加“流水线提速考试”,只在精心挑选的小题上拿高分还不够:真实工厂里常做哪道工序、处理多大一批零件,才决定它有没有用。Atrex-Bench 因此从实际推理集群的生产轨迹——服务运行时留下的计算记录——中出题,检验 AI 编写的 GPU 内核,也就是在 GPU 上执行具体计算的小程序。
基准收录 30 种算子和 440 种张量形状,来自超过一万块加速器、1,303 份性能记录。算子是矩阵乘法等计算动作,形状则是输入数据的尺寸;同一动作换个尺寸,最快写法可能完全不同。它还按每道题在真实服务中占用的 GPU 时间加权,并用 Roofline 上限——芯片理论上能达到的速度——衡量内核离硬件极限还有多远。
作者测试六款前沿编程 Agent 后发现,单看“答案正确”会高估能力:不少通过项其实调用了 PyTorch 的备用实现,并非模型自己写出的目标内核。配套的 Atrex-Kernel-Agent 会反复测量、修改,并在搜索停滞时局部重启;作者称,在受控案例中,它把备用调用改成了真正的内核,性能达到或超过人工调优的生产基线。