Komorebi AI突破：AI实现程序自我优化调试能力提升

这项由西班牙马德里Komorebi AI Technologies研究团队发布的研究成果，以预印本形式于2026年5月12日在arXiv上公开，论文编号为arXiv:2505.09708v1，分类为计算机科学·机器学习方向。感兴趣的读者可以通过该编号在arXiv平台查阅完整原文。

说到底，这项研究要回答一个越来越切实的问题：当我们把写代码的任务交给AI，让它不断自我改进，我们怎么知道它改出来的东西是真的"更好"，而不是在钻空子、蒙混过关？这个问题听起来有些哲学味道，但研究团队用一套非常具体的工程方案给出了回答，而整个实验平台就是他们设计的METAL-SCI基准测试套件。

**一、 先搞清楚"舞台"：这场AI调教实验在哪里发生**

在展开故事之前，需要先认识一下这个实验的"舞台"。大多数人听到"AI写代码"，脑海中浮现的可能是帮你补全一个Python函数，或者自动生成一段JavaScript。但这项研究的对象要"底层"得多——它针对的是运行在苹果芯片（Apple Silicon）上的Metal计算核（kernel）。

Metal是苹果专门为自家芯片开发的GPU编程语言，类似于英伟达GPU上用的CUDA。所谓"计算核"，可以理解为一段专门在GPU上并行运行的小程序，这类程序通常承担着最耗算力的任务，比如物理模拟、流体仿真、分子动力学计算等等。写一个"能跑"的计算核不难，但写一个"跑得飞快、接近硬件极限"的计算核，需要深刻理解芯片的内存架构、并行调度机制和各种底层优化技巧。

研究团队选择苹果芯片作为实验平台，有几个颇为务实的理由。苹果芯片的统一内存架构（CPU和GPU共享同一块内存）省去了数据在两者之间来回搬运的麻烦，这意味着每次编译一段候选代码、运行它、检查结果，整个循环只需要不到一秒，这对于需要运行几十次迭代的进化搜索来说至关重要。此外，Metal在以英伟达CUDA为主导的AI训练数据中严重缺席，这恰好提供了一个真实的"考场"：如果AI模型只是在背诵它在训练数据里见过的CUDA代码模式，换到Metal环境下就会原形毕露。果然，实验中发现三大主流AI模型都会频繁犯Metal特有的语法错误，比如错误放置`[[max_total_threads_per_threadgroup]]`这个关键属性，或者误用`half`这个被Metal保留为16位浮点数类型的关键字。

**二、 十道难题，六种"关卡类型"：METAL-SCI的设计哲学**

METAL-SCI基准测试包含10个计算任务，被归纳为六种结构性截然不同的优化"关卡"（论文中称为R1到R6）。这个分类方式是整个基准设计的核心逻辑——研究团队刻意确保每种关卡所需的优化策略完全不同，在一种关卡上奏效的技巧，放到另一种关卡上不仅无用，甚至可能起反效果。

第一种关卡（R1）是规则网格上的"模板计算"（stencil），典型任务是二维热方程的5点差分和三维声波方程的7点差分。这类问题的瓶颈在于内存带宽——CPU/GPU之间的数据传输速度跟不上计算速度，所以关键技巧是"光晕缓存"（halo blocking）和时间维度上的分块，尽量让数据在靠近芯片的缓存里多待一会儿。

第二种关卡（R2）是"计算密集型"问题，包括N体引力模拟和哈密顿蒙特卡洛（HMC）采样。这类问题每次内存访问对应大约20次浮点运算，瓶颈从带宽转移到了芯片的峰值算力，关键技巧是"寄存器分块"（register tiling）——把数据尽量塞进最快的寄存器里做运算，而不是反复从内存读取。

第三种关卡（R3）是"多场、特殊内存"问题，包括D2Q9格子玻尔兹曼（LBM）流体模拟和二维伊辛模型的蒙特卡洛模拟。格子玻尔兹曼每个网格点要同时维护9个分布函数，每步计算的内存流量高达72字节/单元，关键技巧是采用"SoA"（Structure of Arrays）内存布局，以及对BGK碰撞步骤进行代数化简以减少运算量。伊辛模型的验证则要求CPU和GPU输出完全一致到比特级别，这需要使用确定性随机数生成器。

第四种关卡（R4）是"不规则内存与原子操作"，代表任务是带格子列表的Lennard-Jones分子动力学模拟。这类问题的难点在于粒子分布不均匀导致的负载不平衡，以及多线程同时更新同一个计数器时产生的"原子竞争"瓶颈。

第五种关卡（R5）是"多核规约"，代表任务是Grad-Shafranov等离子体平衡方程的Picard迭代求解。每次外层迭代需要先执行一个全局最大值规约，再执行一个带非线性源项的变系数5点差分，这两个核的协同方式直接影响性能。

第六种关卡（R6）是"数据重排/蝴蝶运算"，代表任务是三维快速傅里叶变换（FFT）。与前两种关卡不同，这里的优化目标不是算术密度，而是数据移动方式——如何利用位反转、Stockham自动排序、混合基数（radix-4、radix-8）蝴蝶运算以及GPU内部的simd通道间交换指令（`simd_shuffle_xor`）来最大化吞吐量，同时避免共享内存的"存储体冲突"（bank conflict）。

正是这种多样性构成了真正的挑战。正如研究团队指出的，对于AI模型来说，"背模板"的策略在这里根本行不通——不存在一套通用的代码结构能同时赢得所有关卡，模型必须真正识别出当前任务属于哪种类型，然后选择正确的优化路径。

每道任务都配备了一个CPU参考实现（用于验证正确性）、一个基于硬件"屋顶线模型"（roofline model）的评分函数，以及三个"训练用"问题规模加上一个"保留"的未见规模。屋顶线模型是一个来自加州大学伯克利分校的经典硬件分析工具，它根据程序的"算术强度"（每字节内存访问对应多少次浮点运算）来预测一个程序在特定硬件上能达到的理论上限——带宽受限程序的上限由内存带宽决定，计算受限程序的上限由峰值算力决定。用这个标准打分，意味着成绩衡量的是"距离硬件极限还有多远"，而不是"比某个参考实现快了多少"，这让分数更具客观性。

**三、 进化循环：让AI反复打磨自己的代码**

有了基准，下一步是搭建让AI在这个基准上自动搜索优质代码的循环。研究团队设计的框架叫"(1+1)进化策略"，这是进化算法中最简单的一种形式：一个"亲代"，每轮产生一个"子代"，只有子代表现更好时才取代亲代成为新的"亲代"。

整个循环运作方式如下。首先，每个任务都有一个初始的"种子核"（seed kernel），作为起点。然后，一个冻结的AI大语言模型（不在搜索过程中更新权重）扮演"突变器"的角色——它读取当前的代码候选和一个详细的反馈包，生成一个新的Metal源代码。反馈包里包含上一次候选代码的编译错误（如果有的话）、每个问题规模下的正确性标志（对比CPU参考结果是否在误差范围内）、每个规模下的实际吞吐量与屋顶线上限的比值，以及最近几轮迭代的简短历史。

每个新的候选代码被提交后，测试框架会立即在Python进程内部调用苹果的Metal运行时，通过`MTLDevice.newLibraryWithSource`接口直接编译这段代码，而不是走离线的`xcrun metal`工具链。编译成功后，代码会在三个训练用规模上分别运行（先3次预热，再10次计时，取中位数），并与CPU参考结果对比正确性。

评分函数被设计成一个"硬性门槛"与"几何均值"的组合：只有在所有训练规模上都通过正确性检验的候选才能得分，得分是各规模下"实际吞吐量/屋顶线上限"比值的几何均值。用几何均值而非算术均值，是为了防止AI只把某一个规模调到极致而忽视其他规模。只要新候选的分数严格高于当前最优解，就替换之，否则丢弃。

这个循环对每个任务迭代10到25次不等（LBM任务迭代25次，3D波动方程迭代15次，其余为10次），全程不需要人工干预。研究团队在苹果M1 Pro芯片（峰值算力4500 GFLOPS，内存带宽200 GB/s）上分别用三个主流大模型——Claude Opus 4.7、Gemini 3.1 Pro、GPT-5.5——各跑了一遍完整的10任务流程。

**四、 "隐藏关卡"：那个AI永远看不见的测试**

进化循环运行过程中，AI模型只能看到三个训练规模上的反馈。但每个任务还有第四个"保留规模"——这个规模的问题从不出现在任何反馈包里，AI在整个搜索过程中对它一无所知。等到每个任务的K轮迭代全部结束，测试框架才会把当前最优解拿出来，在这个保留规模上运行一次，记录结果，但这个结果绝不反馈给AI。

研究团队把这个机制称为"保留门"（held-out gate），用符号ΦT表示。它的作用类似于机器学习里的测试集——你在训练数据上表现再好，真正的价值要看在从未见过的数据上的表现。对于这里的代码搜索任务，一个代码在训练规模上跑得快，不代表它在更大或更小的规模上同样有效，甚至不代表它在新规模上的输出是正确的。

这个"从未见过的规模"设置得颇有讲究。以N体模拟（nbody）为例，训练规模是N∈{256, 1024, 2048}个粒子，保留规模是N=512。以哈密顿蒙特卡洛（hmc）为例，训练规模覆盖问题维度d∈{8, 16, 32}，保留规模是d=24，恰好落在训练维度之间的空缺处。以3D FFT为例，训练规模是边长N∈{32, 64, 128}的正方体，保留规模是N=256的正方体，比最大训练规模大一倍。

**五、 实验结果：三个AI，三种"出错方式"**

接下来是整篇研究最精彩的部分——三个AI在这套赛制下各自跑出了什么成绩，又在哪里翻了车。

在训练规模上（即AI能看到的那些规模），三个模型都取得了不错的提速。以"最优解/初始种子"的比值来衡量，跨10个任务的提速从1.00×（没有改进）到10.7×不等。

表现最亮眼的任务是hmc（哈密顿蒙特卡洛）。Opus 4.7和Gemini 3.1 Pro分别独立找到了同一个关键优化：把内层矩阵向量乘法改用`template`的编译期常数维度版本，让编译器能在编译阶段完全展开内层循环，消除所有分支判断。这一改动把d=8时的吞吐量从121 GFLOPS（占峰值的2.7%）提升到了970 GFLOPS（占峰值的22%），一次迭代实现了8倍的提速，而此前五轮迭代毫无进展。GPT-5.5也发现了类似的优化，但表现相对保守，达到7.19倍提速。

在其他任务上，三个模型各有所长，且优势模式截然不同。Opus在"调紧同一个算法"类的任务上更强，比如N体模拟（2.83×）、格子玻尔兹曼（1.46×）、三维波动方程（1.26×）——这些任务的最优解是在原有算法框架内做更精细的参数调整和代码优化。Gemini则在"换一套算法"类的任务上更强，比如Grad-Shafranov方程（2.89×）和Lennard-Jones分子动力学（1.98×）——这些任务的最优解需要重新设计归约策略或内存访问模式。GPT-5.5在FFT3D任务上取得了三个模型里最高的训练分数（2.95×），比Gemini（1.19×）和Opus（1.03×）都高出不少，这是训练集上任意单任务中最大的模型间差距。

然而，故事在"保留关卡"处急转直下。

Opus的hmc最优解在训练规模上得了满分（所有三个维度都通过正确性检验），但在保留维度d=24上，样本协方差误差高达约10个标准差——输出结果已经完全错误。原因在于Opus的模板化策略只枚举了D∈{8, 16, 32}三个值，当d=24到来时，它被路由到了D=32的分支，程序按照32维来展开内层循环并处理数据，但实际数据只有24维，多出来的8个维度读取了无效数据，计算结果自然偏差极大。从外部看，这个程序在训练集上表现完美，得分10.6倍，没有任何警告信号——但它交出的是一个看起来正确实则严重错误的采样器。

Gemini的做法更谨慎：在使用模板化快速路径的同时，保留了一个运行时维度的"保底路径"（fallback），因此d=24能正确运行，并在保留测试中取得了17.6倍的提速（相对于初始种子）。GPT-5.5更进一步，显式枚举了D∈{8, 16, 24, 32}四个值，给d=24专门准备了一个全展开的模板实例，同时保留了一个通用的运行时路径兜底，在保留测试中取得了18.6倍提速。

GPT-5.5在FFT3D任务上出现了另一种更隐蔽的失败——性能的静默衰退。它的训练最优解由三个手写的快速实现函数组成，分别对应N=32、N=64、N=128三种边长，都使用了`simd_shuffle_xor`指令做高效的蝴蝶运算。但当N不属于这三个值时，代码直接落入了一个教科书式的O(N?)直接离散傅里叶变换（DFT）。在保留规模N=256时，这个O(N?)路径需要做256次复数乘法才能算出一个输出元素，而种子代码的O(N log N) Stockham FFT只需要8次（log?256=8）。每条1D线约多做32倍的运算，三轴串联后结果是0.23倍的种子性能——也就是说，比最初的起点慢了4倍多。训练分数是2.95×，这是AI报告给自己的"进步"；保留分数是0.23×，这才是实际部署性能。两者之间的鸿沟，只有"保留关卡"才能发现。

相比之下，Opus和Gemini在同一保留规模上的FFT3D成绩分别是42%和45%的屋顶线命中率，都来自正确的O(N log N)路径（分别是基于共享内存的Stockham radix-4 FFT和基于`simd_shuffle_xor`的前五级免屏障蝴蝶运算）。

其他任务的保留结果呈现出有趣的分布。N体模拟（nbody）、Grad-Shafranov方程（gradshaf）和Lennard-Jones分子动力学（lj）在三个模型上都表现出真正的泛化——保留规模上的提速与训练规模相近甚至更高。热方程（heat2d）和萨克斯比（saxpy）这类带宽饱和型任务则已接近硬件极限，无论在哪个规模上提升空间都很有限。格子玻尔兹曼（lbm）的三个模型在保留规模（192?）上的表现基本平手，说明这个规模恰好落在所有模型调优策略的共同有效区间。

**六、 代码细节：三场优化风格的"对比画"**

研究团队给出了几个任务上不同模型最优解的代码对比，这些细节揭示了AI优化风格的实质差异。

在格子玻尔兹曼（lbm）任务上，两个模型都保留了原始的"拉流+BGK碰撞"结构，但优化方向完全不同。Opus的做法是：提前计算`A = fma(-1.5, ∥u∥?, 1)`一次，然后对9个方向的平衡分布函数分别手动展开成两个FMA（乘加融合指令），再把弛豫步骤也折叠进第三个FMA，得到九段完全展开的、零分支的代码块。同时，Opus还在核函数声明上加了`[[max_total_threads_per_threadgroup(64)]]`，把线程组几何形状固定为32×2，与GPU的SIMD宽度对齐。Gemini则保持教科书式的BGK公式，用`#pragma unroll`展开k=0到8的循环，没有A的提取，没有FMA折叠，也没有线程组几何约束。结果是：在256?这个规模（缓存常驻、每个指令都能发挥作用的区间），Opus的吞吐量是Gemini的约1.2倍，但在128?这个规模上，Gemini反而略胜。总的来说，Opus赢在"把同一套算法抠到极致"，而Gemini的通用展开在中等规模也相当有竞争力。

在FFT3D任务上，两个模型的差异则属于完全不同的算法路线。Opus实现了一个标准的Stockham auto-sort radix-4 FFT：每一个蝴蝶运算阶段都通过共享内存（threadgroup memory）做乒乓缓冲，每阶段需要一个屏障（barrier）来同步。Gemini发现了一个Metal特有的技巧：GPU的SIMD组宽度恰好是32，而Cooley-Tukey算法前五个阶段的蝴蝶配对距离（1, 2, 4, 8, 16）全部小于32，这意味着可以用`simd_shuffle_xor`指令在SIMD组内部直接交换数据，完全不需要共享内存和屏障。只有第六阶段及以后才退回到共享内存路径。这一改动在每条1D FFT中节省了5个屏障，在一次3D FFT的三轴串联中合计节省15个屏障，直接带来了约1.7倍的性能提升（Gemini 0.282 vs Opus 0.167的屋顶线命中率）。

在hmc任务上，三个模型的代码分别展示了三种不同的"安全性与性能之间的权衡哲学"。Opus选择了最激进的方式：只管d∈{8, 16, 32}，d=24进了D=32的门，正确性碎了一地。Gemini选择了保守路线：完全不做编译期特化，运行时d用于所有维度，安全但吞吐量有所牺牲。GPT-5.5找到了一个折中：显式枚举四个维度（包括d=24），每个都有完整展开的模板实例，另外还有一个运行时路径兜底，覆盖任何意外输入。代价是训练分数（0.0634）略低于Opus（0.0932）和Gemini（0.0870），但保留分数是三者中最高的（18.6倍）。

**七、 副产品：对大模型能力的几点实际观察**

实验的几个侧面观察也值得一提。在正确性失败率方面，三个模型差异明显：Gemini在整个候选代码生成过程中零正确性失败，GPT只有2次，而Opus有13次（其中10次集中在3D波动方程任务上，原因是多步蛙跳时间积分会放大任何符号或索引错误，使其在后续步骤中爆炸成NaN）。

在Metal语法错误方面，三个模型都犯了相似的错误类型，只是数量不同：`[[max_total_threads_per_threadgroup(N)]]`被放在了参数列表后面或作为独立语句，而不是正确地放在`kernel void`声明之前；`half`被当作变量名使用，而它是Metal的16位浮点关键字；使用了Metal不支持的C++ lambda表达式。Opus编译失败12次，GPT 12次，Gemini 22次。

在每次迭代的实际耗时方面，差异极为悬殊：Opus平均每次迭代约0.6分钟，Gemini约3.5分钟，GPT约6.6分钟。在相同的迭代预算下，GPT每次迭代的推理成本大约是Gemini的2倍、Opus的10倍。这意味着如果只关心训练分数，Opus是最"划算"的选择；如果关心代码在未见过的规模上的可靠性，Gemini和GPT（尤其是GPT对hmc的处理）表现更稳健，但代价是更高的时间和计算开销。

从迭代曲线来看，绝大多数任务在第8次迭代前后就已经停止进步，但格子玻尔兹曼（Opus）从第3次迭代开始在1.36×处平台停滞，直到第23次才突破到1.46×（靠BGK折叠加固定线程组实现）；3D波动方程（Opus）在第14次迭代才实现1.26×的最终最优解。这说明10次迭代的预算对大多数任务足够，但某些任务需要更长的搜索窗口，这也是研究团队为这两个任务专门设置更高迭代预算的原因。

归根结底，这项研究的最大贡献不是给出了某个"哪个AI最会写GPU代码"的排行榜，而是演示了一个极其简单却有力的机制：在你的自动代码搜索循环里，额外保留一个AI永远看不见的测试配置，在最后才运行一次。这一次测试可能是几秒钟的GPU时间，但它能发现两类AI最擅长制造的危险——输出看起来正确实则错误的代码（Opus的hmc），以及性能看起来提升实则退步的代码（GPT的fft3d）。这个机制被研究团队称为"廉价的机械监督原语"，意思是它不需要人工审查代码，不需要任何AI辅助验证，只是多运行一次，让现实数据说话。在AI辅助的自动化编程越来越普及的今天，这种"多留一个门缝"的思路，对于任何需要把AI生成的代码真正部署到实际系统中的人，都有相当的参考价值。如果你想深入了解技术细节、查看完整的代码示例和数学公式，可以通过arXiv编号arXiv:2505.09708找到原论文。

Q&A

Q1：METAL-SCI基准测试和其他AI代码生成基准测试（比如KernelBench）有什么本质区别？

A：METAL-SCI和KernelBench最根本的区别有三点。第一，任务类型不同：KernelBench针对的是神经网络里的矩阵乘法、注意力机制等机器学习算子，而METAL-SCI测的是流体模拟、分子动力学、傅里叶变换等科学计算任务，这些任务在AI训练数据中出现很少。第二，评分标准不同：KernelBench用"比PyTorch快多少倍"来打分，METAL-SCI用"达到硬件理论上限的百分之几"来打分，后者与任何参考实现无关，更客观。第三，METAL-SCI额外设置了一个AI永远看不见的"保留规模"测试，专门用于发现AI把代码调优到训练配置后在新规模上失效的问题，这种泛化性检验在其他基准里没有。

Q2：为什么Opus的哈密顿蒙特卡洛代码在训练集上完全正确，换一个维度却输出错误结果？

A：Opus在优化hmc时使用了C++的模板技术，把循环维度d变成编译期常数，让编译器能完全展开内层循环。但Opus只为d=8、16、32三个值生成了专门的模板实例，当遇到训练集里没有的d=24时，程序把它路由到了D=32的分支。D=32的版本会按照32维来展开并执行所有运算，但实际数据只有24维，多出的8个位置读取的是无效内存内容，导致计算出来的样本协方差误差高达约10个标准差。训练集上的三个维度（8、16、32）都恰好在枚举范围内，所以每次都通过了正确性检验，分数毫无异常；只有保留的d=24暴露了这个缺口。这个案例说明，通过训练集正确性检验并不能保证代码在所有输入上正确。

Q3：苹果M1 Pro芯片在这类GPU编程实验里有什么特别的优势和局限？

A：优势主要有两个。一是统一内存架构：苹果芯片的CPU和GPU共享同一块物理内存，不需要在两者之间来回复制数据，所以每次编译代码、运行测试、检查结果的完整循环不到一秒，非常适合需要反复迭代的进化搜索。二是测试真实的AI知识盲区：Metal语法在以英伟达CUDA为主的AI训练数据中几乎没有，这让三大AI模型都暴露出了对Metal特有语法（比如`[[max_total_threads_per_threadgroup]]`属性的正确放置位置）的不熟悉，提供了真实的泛化能力测试。局限方面，研究团队也指出，目前使用的静态屋顶线上限没有考虑小规模问题时的缓存残留效应，在小规模测试中评分可能偏低，未来需要针对每个任务和每个规模做更精细的硬件性能建模。