在 HotChips 上展示 Jalapeño MLA 内核 以及随后 SemiAnalysis 的评论之后，引发了大量讨论。作为 OpenAI 的硬件团队，我们只是 略微 触及了这个小金块：即 AI 正在编写我们的内核，而且当它这样做时，我们并不 需要 逐行理解内核的功能。我们明显遗漏了：这样的事情是如何成为可能的？与更传统的代码生成方法相比，正确的思考方式是什么？优化后的内核是否与未优化的内核一样可靠？
就我的背景而言，我在加速器编译器领域工作了十多年。我创立了 XLA，这是一个优秀的编译器基础设施，拥有跨公司的杰出团队和努力。在过去两年多的时间里，我在 OpenAI 一直试图重新构想，在 AI 时代编译器应该如何工作。新的编译器公式将借鉴现有优势，但不可否认的是，工具包中有一个强大的新工具可以利用。
这将是一段旅程，但我希望阐明 AI 如何被用于自动化 计算机程序改进；即优化编译。我确实认为，通过 AI，我们可能会体验到我们称之为“编译器 2.0”的东西。AI 在其可以提出的方案上受到的根本限制更少，而其提出的方案是模型训练和上下文的结果，这让我将其归类为“随机优化器”——这可能会带来挑战，但正如我们将看到的，这也是巨大优势的来源……
大量的学术研究和工业应用已经朝着这个方向发展，并迅速揭示了 AI 在优化编译器领域的潜力，但我们正处于一个需要对其进行宏观解释的时刻。
背景
编译器接收程序，并输出这些程序的翻译或改进版本。
在输入和输出两侧，程序都具有 语义，这些语义告诉我们程序的含义、它们可能做什么，以及如何推理这些可能做的事情。
我们这些从事编译器工作的人，将编译器视为纯函数——它们接收一个数据结构，并输出一个应具有对应语义的数据结构。
有时，我们的编译器专注于“降低”或“翻译”。例如，它们可能接收 C 语言并输出 x86-64 汇编，我们通常认为这是“更低级别”的。但通常，作为该过程的一部分，它们所做的不仅仅是翻译……
在实践中，我们的编译器专注于“优化”。它们可能接收一个代表程序的数据结构——用我们的行话来说就是“中间表示”（IR）——并尝试生成该程序的更好版本。有时“更好”意味着运行所需的周期更少，有时意味着不必要的代码更少，有时意味着针对我们可以证明“必须为真”的程序特性进行特化（部分求值）。
现在，简要地考虑一下，LLM 最初是为了将人类文本从一种语言翻译成另一种语言而创建的。显然，翻译是它们的强项。通过我们在日常任务中使用 LLM，我们可以看到它们也能编写新的解决方案并改进现有解决方案。我们许多程序员也有过让 LLM“优化这段代码”的经历，它们确实能出色地完成。（然而，我们需要知道它们是否正确优化了代码，我们稍后会谈到！）这只是为了强调 LLM 具备我们在优化编译器中寻找的能力。
优化与最优性
优化编译器，顾名思义，试图根据某个目标（通常是执行时间）提高 它们正在处理的程序的 最优性。在一般情况下，对于任意程序，这非常困难，以至于有一个被称为编译器工程师充分就业定理的定理。（我是在选择成为编译器工程师之后才发现这一点的，但这仍然让我感到安慰！）
“超级优化器”是优化编译器的一个令人惊叹的小众子领域。想象有一个给定的程序，我们可以通过语义来描述它的功能。那么，具有相同语义的 最 优程序是什么？这正是超级优化器试图解决的问题，它本质上是一个搜索问题……
想象一下，我试图找到具有相同语义的最短程序，并且我有办法询问一个候选程序是否具有相同的语义。我可以假设地按目标顺序枚举每个程序，并选择具有相同语义的最小程序。
然而，按目标顺序枚举每个程序听起来非常棘手。我最喜欢的学术论文之一，发表于 2013 年，题为“STOKE”（随机超级优化），提出了一个问题：“如果我们只是反复 随机调整 程序，最终能否观察到最佳程序？”他们提出，通过随机游走（以及我们 OG 机器学习朋友马尔可夫链蒙特卡洛 / 梅特罗波利斯-黑斯廷斯算法），最终你会看到那个最优程序。
蒙特卡洛调整通常是笨拙的（你随机选择一个调整），但速度很快。LLM 非常智能（许多推理 token），但相对较慢。
如果，不是采用笨拙/快速的蒙特卡洛调整，而是让 LLM 找出程序应该前进的方向呢？我们将拥有一个 非常 智能的随机优化器，引导我们的程序穿越优化后的程序空间。
优化的直觉
让我们退一步。想想你认识的最能体现“极致优化代码片段”的人。简而言之，我们称他为“优化狂 Ollie”。Ollie 可能对哪些代码调整可能产生效果有一种直觉。Ollie 可能会尝试一些方法看看是否有效，如果无效，就回滚并尝试其他方法。但他们对于哪些事情是可能的，以及他们如何 可能 击败编译器，有一种直觉。
Ollie 拥有的这些直觉通常超出了编译器的能力范围。尽管现代优化编译器的结果令人印象深刻，但它们基于相当简单的规则和启发式方法。用技术术语来说，它们基于一个运行到不动点的局部数据流变换的思想。我们还对考虑因素进行阶段排序；即，我们构建编译器流水线来先考虑 A 然后 考虑 B，而不是考虑复合的 AB 问题。调度器和寄存器分配器是一个臭名昭著的例子，许多博士论文都尝试过复合调度器-寄存器分配器（以获得消除阶段排序的好处），但在实践中很难使其工作。
这就是为什么 Ollie 的专业知识如此宝贵。Ollie 通常知道如何用针对特定情况定制的启发式方法来平衡几个 NP 完全问题。因此，它具有更多的定制化上下文感知和敏感性。Ollie 还能够采用优化编译器可能无法有效应用的技术，尤其是组合使用，例如提取函数、定制 ABI 或变换以实现向量化，以及其他让我们抱怨“我希望编译器有办法做到这一点……”的事情。
现在考虑一下，AI 通过其拥有的任何推理能力，可能能够扮演一个迷你 Ollie 的角色。它可能没有那种关于什么会成功的匹配直觉，但它对什么可能有利可图有一种模糊的概念，并且它可以进行许多、许多 次尝试。
采用这种方法，与 STOKE 论文不同，我们无法 保证 随着时间趋于无穷，我们能看到 最优 程序，但因为 AI 具有“更接近人类”的推理能力，它实际上可以在单位时间内获得显著的、类似人类的进展。
联系回：MLA 内核
首先让我说明：我 不 知道 AI 为我们的 Jalapeño MLA 内核生成了什么底层代码，但我 确实 知道如何为 MLA 输入 numpy 代码。
在我之前参与的 XLA 编译器中，我们会将这些 numpy 操作融合成块，然后使用一个称为“发射器”的元程序将其降低为循环、指令和更底层的原语。
当 XLA 编译器/发射器程序这样做时，我不需要关心后端输出了什么汇编代码。对于我们的随机优化器，AI 在概念上取代了发射器元程序——它既负责降低也负责优化，我们可以要求它进一步优化，直至接近理论峰值性能。
我希望这能清楚地说明 AI 的位置，以及它如何类似于现有优化编译器系统中的一个组件。同样有帮助的是思考：我们现在认为的“汇编代码”的层级正在 上移。当你用普通的 C++ 编写代码并以 -O3（最高典型优化级别）编译时，你不会期望理解输出的汇编代码，即使你理解了你输入的 C++ 代码。我们在这里做的事情是类似的，但使用的是更高级、更数学化的输入规范。
现在，一个关键问题是，我们 如何 检查 AI 输出的程序是否确实等价于高级描述 / numpy 代码。这个检查机制确立了 AI 随机优化过程的可靠性。我预计未来的博客文章可能会更详细地讨论这一点，但目前，只需说测试语义等价性是可能的，并且我们正在这样做就足够了。加速器程序特别适合强健、完整的契约，我们可以验证这些契约“正是 AI 优化程序所做的”，因为它们在广泛的上下文中非常数学化和数据流导向。
请注意，该领域的许多相关技术是由程序合成子领域的努力开创的。优化编译器说，“这里有一个具有语义的程序，让它变得更好但保持语义等价！”，而程序合成则说，“我相信存在一个具有这些语义的程序，请尝试找到你能找到的最好的一个”。程序合成是一个比优化编译更困难的问题，但它受到的根本限制也更少。这实际上就是像 Ollie 这样的人在做得比优化编译器更好时所做的事情，也是 AI 现在可以帮助我们自动化的事情。AI 可以从原始程序中汲取“灵感”，但它不必仅仅对其执行微小的局部变换。经典的优化编译器 不会 看到“哦，你写了一个冒泡排序”，并通过理解契约将其切换为快速排序，但 Ollie 和 AI 都能够做到这一点。这就是为什么我们更多地处于程序合成加随机优化的范畴，而不是经典的优化编译器范畴。
所有这些结合在一起，就形成了这样一个事实：你可以从一个“与 numpy 代码相差不远”的东西开始，等待 48 小时，然后得到一个具有相同语义的优化内核，正如我们在 HotChips 演讲中展示的那样：
正如幻灯片中提到的，在我们的机器上，我们经常能够观察到 AI 的性能超越我们的人类专家，甚至 在我们认为已经调优得相当好的内核上也是如此。通常，由于可能需要探索的各种组合/排列的多样性，仍然存在相当可观的可实现百分比提升。这对于人类性能工程师来说通常是难以忍受的繁琐工作。
回顾与总结
归根结底，编译器只是一个函数。我们将程序交给这个函数，然后得到一个更好的程序版本。我们期望输出的程序和输入的程序具有相同的语义。
传统的优化编译器通过数据流规则和启发式方法来改进程序。这些方法的来源完全可理解，但它们所能做出的变换也可能更受限制。
相比之下，AI 作为一个 随机 优化器，只需要“深入思考”并输出一些东西。它的变换不受根本性限制，使其更类似于我们的人类专家优化器。我们需要方法来检查它输出的程序是否可靠，并实现了我们输入的相同语义，而我们确实已经具备了这些方法。这种 AI 优化特别适用于具有非常强契约的数学运算。这些契约避免了逐行理解内核功能的需要。
这就是我们获得 AI 生成的 MLA 内核的方式！
