💬 NLP

Axon: A Synthesizing Superoptimizer for Tensor Programs

Axon 是一个合成型超优化器，它通过使用程序合成和基于 SMT 的语义验证来自动为 AI 加速器生成高性能张量程序，从而在不依赖手工编写的重写规则的情况下发现代数变换、优化分块（tiling）与指令选择，并进行算子融合。

原作者： Akash Kothari, Shaowei Zhu, Daniel Kroening, Chungha Sung

发布于 2026-06-26

📖 1 分钟阅读☕ 轻松阅读

原作者： Akash Kothari, Shaowei Zhu, Daniel Kroening, Chungha Sung

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正试图烤出一个完美的蛋糕，但你是在一个拥有非常特定、极其挑剔的烤箱和搅拌机的超高科技厨房里工作的。你知道你想烤什么（食谱），但要让机器尽可能快地完成，你需要确切知道如何装载食材、何时搅拌以及在烤箱上按下哪些特定的按钮。

长期以来，为 AI 芯片（即“烤箱”）编写代码就像是手工一点一滴地制作蛋糕。你需要成为一名精通烤箱具体机械原理的大厨，才能获得理想的结果。如果你在装载食材时出了小错，蛋糕就会变得又慢又干。

Axon 是一个全新的工具，它扮演着一个自动化的、超级聪明的副厨的角色——它不仅能遵循你的食谱，还能发明出更好的烹饪方法。

以下是 Axos 的工作原理，通过简单的步骤进行拆解：

1. “做什么”与“怎么做” (The "What" vs. The "How")

通常，程序员会写下一份食谱，比如：“将面粉和鸡蛋混合，然后烘烤。”他们不会去担心烤箱是如何混合它们的。

问题所在： 烤箱（AI 芯片）非常挑剔。它有不同的隔间（比如一个大搅拌碗、一个小搅拌器和一个加热元件）。如果你告诉它把所有东西都放在大碗里混合，而它其实可以用小搅拌器，那么就会浪费时间。
Axon 的职责： Axons 接过你的简单食谱并说：“等等，我可以做得更快。”它不仅仅是遵循食谱，它还会根据烤箱的具体优势来重写指令。

2. “神奇的重组” (代数变换 - Algebraic Transformations)

想象一下你的食谱写着：“首先切碎所有的胡萝卜，然后将它们混入汤中，最后加入盐。”

传统方式： 你先切碎，然后混合，最后加盐。你必须在切菜板和炉灶之间来回搬运碗。
Axon 的方式： Axon 查看了数学逻辑并意识到：“嘿，我可以在切胡萝卜的同时加入盐，或者我可以在切碎之前就把胡萝卜混入汤中，效果其实是一模一样的！”
结果： 通过重新排列操作顺序，Axon 让烤箱可以同时做两件事（比如同时进行切碎和加热），而不是一个接一个地进行。它在数学上证明了即使步骤不同，蛋糕的味道依然保持不变。

3. “打包拼图” (分块/平铺 - Tiling)

想象你有一大堆食材，但你的搅拌碗很小。你必须把食材从储藏室（慢速）移动到搅拌碗（快速）中，混合一点，再移回去，如此循环。

问题所在： 如果一次移动太多，碗会溢出来；如果一次移动太少，你就会把时间都花在往返于储藏室的路上。
Axon 的方式： Axon 是一个打包大师。它尝试数百万种不同的方式来将食材分组成“批次”（tiles）。它会问：“如果我一次移动 100 个胡萝卜呢？如果我一次移动 50 个呢？”它测试这些不同的分组大小，以找到完美的尺寸，既能让碗保持满载而不溢出，又能最大限度地减少往返储藏室的次数。

4. “说明书” (指令集合成 - ISA Synthesis)

烤箱有一个特定的控制面板，上面有“高温”、“搅拌”和“折叠”等按钮。

问题所在： 有时食谱说“混合”，但烤箱并没有“混合”按钮。它只有“搅拌”和“折叠”。
Axon 的方式： Axon 将“混合”翻译成该烤箱能理解的精确“搅拌然后折叠”的序列。它甚至能找到捷径，比如意识到“在移动到炉灶的过程中进行折叠”在这个特定的烤箱上是一个单一的按键操作，从而节省时间。

5. “品尝测试” (经验选择 - Empirical Selection)

这是最重要的一部分。Axon 不仅仅是靠猜测哪种方法最好。

过程： 它生成了数千个不同版本的食谱（有些具有不同的顺序，有些具有不同的批次大小，有些具有不同的按键操作）。
测试： 它实际上在真实的烤箱（AI 芯片）上运行所有这些版本，并记录时间。它会挑选出完成速度最快的那一个。这就像是烤了 1000 个微型蛋糕，看看哪一个能在 5 分钟内出炉，而不是 10 分钟。

结果：有多快？

论文在 Amazon Trainium 芯片（一种 AI 烤箱）上测试了 Axon，测试了 20 种不同类型的“蛋糕”（AI 任务）。

对比标准烤箱： 与标准的烘焙方式（Amazon 的 Neuron 编译器）相比，Axon 使单项任务的速度提升了高达 3.7 倍，复杂任务的速度提升了高达 19 倍。
对比专家厨师： 即使是与精通烤箱细节的人类专家编写的食谱相比，Axon 有时也比他们快了 1.35 倍。
对比其他机器人： 它比另一个名为“Mirage”的自动化工具平均快了约 10%，并且能够处理另外 16 种该工具完全无法处理的“蛋糕”类型。

总结

Axon 是一个工具，它接收关于 AI 需要做什么的简单描述，然后自动找出在特定计算机芯片上执行该任务的最快方式。它通过在数学上证明自己可以重新排列步骤、测试数百万种数据打包方式，并实际运行代码来挑选优胜者来实现这一点。它消除了人类需要成为硬件细微、枯燥细节专家的需求，让人们只需编写食谱，而由 Axon 来处理烘焙工作。

技术摘要：Axon：一种用于张量程序的合成超优化器

问题陈述

编写针对 AI 加速器（如 Amazon 的 Trainium）的高性能算子（kernel）需要深厚的硬件特定细节知识，包括内存层级、分块策略（tiling strategies）、指令选择、数据布局以及算子融合。虽然基于分块的算子语言（如 NKI，Neuron Kernel Interface）抽象掉了一些底层细节，但对于大多数机器学习工程师而言，实现峰值硬件利用率仍然具有挑战性。现有的优化方法存在以下局限性：

手工编写的规则： 许多系统依赖于手动编写的代数变换重写规则，这些规则容易出错，且难以扩展到新的算子组合。
基于模板的生成： 一些方法使用手写的代码生成模板，限制了灵活性。
作用范围有限： 许多优化器仅针对 CPU 或 GPU，未能解决现代 AI 加速器独特的多引擎、DMA 驱动架构问题。
缺乏可证明的正确性： 基于概率的测试或启发式的重排序可能会遗漏最优配置或引入语义错误。

方法论

Axon 是一个专为基于分块（tile-based）的 AI 加速器程序设计的合成超优化器。它接收高层张量程序（以类 NumPy 接口编写），并通过一个多阶段流水线自动生成优化的目标指令（NKI）。其核心创新是 $\nu$ Graph（nu-graph），这是一个持久的内部表示，在整个合成过程中同时追踪所有语义等价的程序变体，将优化决策推迟到经验性选择阶段。

该流水线由以下阶段组成：

通过算子传播进行代数变换：
- Axon 通过在计算图中传播算子来发现语义等价的计算图。
- 它尝试重新排列相邻算子（例如，将元素级乘法移过矩阵乘法），以暴露并行性和融合机会。
- 正确性保证： Axon 并非使用手工编写的重写规则，而是将张量操作分解为细粒度的计算，并使用针对无界张量的 SMT（可满足性模理论）求解器 来证明重排序保持了语义一致性。这可以处理任意张量形状，且无需手动定义属性。
符号分块（Symbolic Tiling）：
- 系统使用用于 strip、block 和 tile 维度的符号参数对每个张量操作进行分块。
- 这些参数受硬件描述（例如片上缓冲区容量和特定指令限制）的约束。
- Axon 生成具有不同分块配置的多个版本，允许独立优化并行分支（例如，将 RMSNorm 和 MatMul 分支分离，以便在不同的计算引擎上运行）。
ISA 合成（指令选择）：
- 利用基于 Sketch 的程序合成方法，Axon 将高层张量操作降低（lower）为目标 ISA 指令（例如，Trainium 的 nc_matmul、dma_copy）。
- 它探索有效的指令序列和特定硬件指令所需的数据布局变换（如转置）。
- 高层操作与合成的指令序列之间的等价性同样通过 SMT 进行验证。
ISA 与算子融合：
- 在合成后的 ISA 图上进行第二轮代数变换，以暴露由合成过程产生的新的融合机会（例如，将转置与矩阵乘法融合）。
- 节点融合（Node Fusion）： 合并具有兼容分块维度的连续操作，使中间结果保留在片上缓冲区（SBUF/PSUM）中，从而减少 HBM 流量。
- 指令融合（Instruction Fusion）： 用融合后的变体替换 ISA 指令序列（例如，dma_copy + nc_transpose $\rightarrow$ dma_transpose）。
代码发射与经验选择：
- Axon 使用具体数值实例化符号参数，以发射具体的 NKI 程序。
- 每个候选程序都由 Neuron 编译器编译，并在目标硬件上执行。
- 根据经验延迟测量值和针对参考实现的数值正确性验证，选择性能最佳的算子。

核心贡献

本文提出了三个主要贡献：

可证明的代数变换： 一种通过传播算子并使用针对无界张量的 SMT 验证等价性，从而自动发现语义等效计算图的方法。这消除了对手工编写重写规则的需求。
基于合成的流水线： 一个完整的基于分块的 AI 加速器流水线，共同探索代数变换、指令选择、分块配置和融合。它在 $\nu$ Graph 中维护所有有效的变体，并进行经验性选择。
全面的评估： 在 Amazon 的 Trainium 上进行了评估，涵盖了 20 个基准测试，包括单个算子和复杂的多算子 LLM 算子（如 Group Query Attention、Gated MLP）。

实验结果

Axon 与三个基准进行了对比：Amazon 的 Neuron 编译器、手工优化的 NKI 算子以及 Mirage（一种最先进的基于搜索的编译器）。

对比 Neuron 编译器：
- 单个算子： Axon 实现了高达 3.7 倍 的加速（例如在大维度下的 SiLU），在各种算子上的几何平均加速比为 1.23x–1.90x。
- 多算子算子（Kernels）： Axon 在 Transpose+MatMul 等算子上实现了高达 19 倍 的加速，在 RMSNorm+MatMul 上实现了高达 1.47 倍的加速。几何平均加速比范围为 1.07x 至 1.32x。
- 原因分析： 收益归功于更好的分块策略、高效的 DMA 利用，以及跨不同计算引擎融合操作的能力（例如，在 MatMul 在张量引擎运行的同时，在标量/向量引擎上并行运行 RMSNorm）。
对比手工优化的 NKI：
- 在特定配置下（如 Cumsum 和 RoPE），Axon 能够匹配或超越专家编写的算子，最高达 1.35 倍，证明了自动化探索可以在不需要深厚硬件专业知识的情况下，达到甚至超过手动调优的效果。
对比 Mirage：
- 在两者都支持的基准测试中，Axon 实现了 2–10% 的几何平均加速。
- 至关重要的是，Axon 成功合成了 16 个额外的基准测试，而 Mirage 由于依赖手工模板而无法处理这些测试。
编译开销：
- 对于复杂的算子，搜索空间可能会生成约 30,000 个候选方案。
- 第一阶段（符号合成）很快（<82 秒）。第二阶段（编译与执行）占据了主要时间（对于复杂的 LLM 层，最高可达约 1.7 小时）。论文指出这是一次性的成本，与其它基于调度（schedule-based）的编译器相当，且优化的算子可以缓存以供重复使用。

重要性与主张

本文声称 Axon 代表了从基于启发式或模板驱动的优化向基于语义的、合成式方法的转变。通过利用 SMT 求解器来保证语义等价性，Axon 可以自动发现复杂的代数变换（例如，跨不同计算引擎重新排列操作），而这些变换很难通过手动编码实现。

其重要性在于：

自动化： 减轻了程序员为特定硬件手动调优算子的负担。
正确性： 与基于概率的测试方法不同，它提供了可证明的变换保证。
性能： 通过在统一的搜索空间内联合优化分块、指令选择和融合，在现代 AI 加速器上实现了显著的加速。
通用性： 能够处理现有基于搜索的编译器（如 Mirage）无法合成的复杂、多算子 LLM 算子。

作者强调，虽然编译开销不容忽视，但这是优化编译器标准的单次成本，且由此带来的性能提升（特别是对于大规模 LLM 工作负载）证明了这一投入是值得的。