💬 NLP

Axon: A Synthesizing Superoptimizer for Tensor Programs

Axon is a synthesizing superoptimizer that automatically generates high-performance tensor programs for AI accelerators by using program synthesis and SMT-based semantic verification to discover algebraic transformations, optimize tiling and instruction selection, and fuse operators without relying on hand-crafted rewrite rules.

Original authors: Akash Kothari, Shaowei Zhu, Daniel Kroening, Chungha Sung

Published 2026-06-26

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Akash Kothari, Shaowei Zhu, Daniel Kroening, Chungha Sung

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to bake the perfect cake, but you are working in a high-tech kitchen with very specific, finicky ovens and mixers. You know what you want to bake (the recipe), but getting the machine to do it as fast as possible requires you to know exactly how to load the ingredients, when to stir, and which specific buttons to press on the oven.

For a long time, writing code for AI chips (the "ovens") has been like trying to bake that cake by hand, one tiny crumb at a time. You need to be a master chef who understands the exact mechanics of the oven to get a good result. If you make a small mistake in how you load the ingredients, the cake comes out slow and dry.

Axon is a new tool that acts like an automated, super-smart sous-chef that doesn't just follow your recipe—it invents better ways to bake it.

Here is how Axon works, broken down into simple steps:

1. The "What" vs. The "How"

Usually, a programmer writes a recipe saying, "Mix flour and eggs, then bake." They don't worry about how the oven mixes them.

The Problem: The oven (the AI chip) is picky. It has different compartments (like a big mixing bowl, a small whisk, and a heating element). If you tell it to mix everything in the big bowl when it could have used the whisk, it wastes time.
Axon's Job: Axon takes your simple recipe and says, "Wait, I can do this faster." It doesn't just follow the recipe; it rewrites the instructions to match the oven's specific strengths.

2. The "Magic Rearranging" (Algebraic Transformations)

Imagine your recipe says: "First, chop all the carrots, then mix them into the soup, then add salt."

The Old Way: You chop, then mix, then add salt. You have to carry the bowl back and forth between the cutting board and the stove.
Axon's Way: Axon looks at the math and realizes, "Hey, I can add the salt while the carrots are being chopped, or I can mix the carrots into the soup before I chop them, and it will taste exactly the same!"
The Result: By rearranging the order of operations, Axon lets the oven do two things at the same time (like chopping and heating simultaneously) instead of one after the other. It proves mathematically that the cake will still taste the same, even though the steps are different.

3. The "Packing Puzzle" (Tiling)

Imagine you have a huge pile of ingredients, but your mixing bowl is small. You have to move ingredients from the pantry (slow) to the bowl (fast), mix a little, move them back, and repeat.

The Problem: If you move too much at once, the bowl overflows. If you move too little, you spend all your time walking back and forth to the pantry.
Axon's Way: Axon is a master packer. It tries millions of different ways to group the ingredients into "batches" (tiles). It asks, "What if I move 100 carrots at a time? What if I move 50?" It tests these different groupings to find the perfect size that keeps the bowl full without overflowing, minimizing the trips to the pantry.

4. The "Instruction Manual" (ISA Synthesis)

The oven has a specific control panel with buttons like "High Heat," "Whisk," and "Fold."

The Problem: Sometimes the recipe says "Mix," but the oven doesn't have a "Mix" button. It only has "Whisk" and "Fold."
Axon's Way: Axon translates "Mix" into the exact sequence of "Whisk then Fold" that the oven understands. It even finds shortcuts, like realizing that "Fold while moving to the stove" is a single button press on this specific oven, saving time.

5. The "Taste Test" (Empirical Selection)

This is the most important part. Axon doesn't just guess which way is best.

The Process: It generates thousands of different versions of the recipe (some with different orders, some with different batch sizes, some with different button presses).
The Test: It actually runs all these versions on the real oven (the AI chip) and times them. It picks the one that finishes the fastest. It's like baking 1,000 tiny cakes to see which one comes out of the oven in 5 minutes instead of 10.

The Results: How Much Faster?

The paper tested Axon on Amazon's Trainium chips (a type of AI oven) with 20 different types of "cakes" (AI tasks).

Vs. The Standard Oven: Compared to the standard way of baking (Amazon's Neuron compiler), Axon made single tasks up to 3.7 times faster and complex tasks up to 19 times faster.
Vs. Expert Chefs: Even when compared to recipes written by human experts who know the oven inside and out, Axon was sometimes 1.35 times faster.
Vs. Other Robots: It beat another automated tool called "Mirage" by about 10% on average, and it could bake 16 types of cakes that the other tool couldn't handle at all.

Summary

Axon is a tool that takes a simple description of what an AI needs to do and automatically figures out the absolute fastest way to do it on a specific computer chip. It does this by mathematically proving it can rearrange steps, testing millions of ways to pack data, and then actually running the code to pick the winner. It removes the need for humans to be experts in the tiny, boring details of the hardware, letting them just write the recipe while Axon handles the baking.

Technical Summary: Axon: A Synthesizing Superoptimizer for Tensor Programs

Problem Statement

Writing high-performance kernels for AI accelerators (such as Amazon's Trainium) requires deep expertise in hardware-specific details, including memory hierarchies, tiling strategies, instruction selection, data layout, and operator fusion. While tile-based kernel languages like NKI (Neuron Kernel Interface) abstract away some low-level details, achieving peak hardware utilization remains challenging for most ML engineers. Existing optimization approaches suffer from several limitations:

Hand-crafted rules: Many systems rely on manually written rewrite rules for algebraic transformations, which are error-prone and difficult to scale to new operator combinations.
Template-based generation: Some approaches use hand-written templates for code generation, limiting flexibility.
Limited scope: Many optimizers target only CPUs or GPUs, failing to address the unique multi-engine, DMA-driven architectures of modern AI accelerators.
Lack of provable correctness: Probabilistic testing or heuristic-based reordering may miss optimal configurations or introduce semantic errors.

Methodology

Axon is a synthesizing superoptimizer designed for tile-based AI accelerator programs. It accepts high-level tensor programs (written in a NumPy-like interface) and automatically generates optimized target instructions (NKI) through a multi-stage pipeline. The core innovation is the $\nu$ Graph ("nu-graph"), a persistent internal representation that tracks all semantically equivalent program variants simultaneously throughout the synthesis process, deferring optimization choices to empirical selection.

The pipeline consists of the following stages:

Algebraic Transformations via Operator Propagation:
- Axon discovers semantically equivalent computation graphs by propagating operators through the computation graph.
- It attempts to reorder adjacent operators (e.g., moving an element-wise multiply past a matrix multiplication) to expose parallelism and fusion opportunities.
- Correctness Guarantee: Instead of using hand-crafted rewrite rules, Axon decomposes tensor operations into granular computations and uses an SMT (Satisfiability Modulo Theories) solver over unbounded tensors to provably check that reordering preserves semantics. This handles arbitrary tensor shapes and requires no manual property definitions.
Symbolic Tiling:
- The system tiles each tensor operation using symbolic parameters for strip, block, and tile dimensions.
- These parameters are constrained by hardware descriptions (e.g., on-chip buffer capacities and instruction-specific limits).
- Axon generates multiple versions with different tiling configurations, allowing independent optimization of parallel branches (e.g., separating RMSNorm and MatMul branches to run on different compute engines).
ISA Synthesis (Instruction Selection):
- Using a sketch-driven program synthesis approach, Axon lowers high-level tensor operations to target ISA instructions (e.g., Trainium's nc_matmul, dma_copy).
- It explores valid instruction sequences and data layout transformations (like transposes) required by specific hardware instructions.
- Equivalence between the high-level operation and the synthesized instruction sequence is again verified via SMT.
ISA and Operator Fusion:
- A second round of algebraic transformations is performed on the synthesized ISA graph to expose new fusion opportunities created by the synthesis process (e.g., fusing a transpose with a matrix multiplication).
- Node Fusion: Merges consecutive operations with compatible tiling dimensions to keep intermediate results in on-chip buffers (SBUF/PSUM), reducing HBM traffic.
- Instruction Fusion: Replaces sequences of ISA instructions with fused variants (e.g., dma_copy + nc_transpose $\rightarrow$ dma_transpose).
Code Emission and Empirical Selection:
- Axon instantiates symbolic parameters with concrete values to emit concrete NKI programs.
- Each candidate is compiled by the Neuron compiler and executed on the target hardware.
- The best-performing kernel is selected based on empirical latency measurements and numerical correctness validation against a reference implementation.

Key Contributions

The paper presents three primary contributions:

Provable Algebraic Transformations: A methodology for automatically discovering semantically equivalent computation graphs by propagating operators and verifying equivalence using SMT over unbounded tensors. This eliminates the need for hand-crafted rewrite rules.
Synthesis-Based Pipeline: A complete pipeline for tile-based AI accelerators that jointly explores algebraic transformations, instruction selection, tiling configurations, and fusion. It maintains all valid variants in a $\nu$ Graph and selects the optimal kernel empirically.
Comprehensive Evaluation: An evaluation on Amazon's Trainium across 20 benchmarks, including individual operators and complex multi-operator LLM kernels (e.g., Group Query Attention, Gated MLP).

Experimental Results

Axon was evaluated against three baselines: Amazon's Neuron compiler, hand-optimized NKI kernels, and Mirage (a state-of-the-art search-based compiler).

Vs. Neuron Compiler:
- Individual Operators: Axon achieves speedups of up to 3.7x (e.g., SiLU at large dimensions) with a geometric mean speedup of 1.23x–1.90x across various operators.
- Multi-Operator Kernels: Axon achieves speedups of up to 19x on kernels like Transpose+MatMul and up to 1.47x on RMSNorm+MatMul. The geometric mean speedups range from 1.07x to 1.32x.
- Reasoning: Gains are attributed to better tiling strategies, efficient DMA utilization, and the ability to fuse operations across different compute engines (e.g., running RMSNorm on Scalar/Vector engines while MatMul runs on the Tensor Engine in parallel).
Vs. Hand-Optimized NKI:
- Axon matches or outperforms expert-written kernels by up to 1.35x on specific configurations (e.g., Cumsum and RoPE), demonstrating that automated exploration can rival manual tuning without requiring deep hardware expertise.
Vs. Mirage:
- On benchmarks supported by both, Axon achieves geometric mean speedups of 2–10%.
- Crucially, Axon successfully synthesizes kernels for 16 additional benchmarks that Mirage cannot handle due to its reliance on hand-crafted templates.
Compilation Overhead:
- The search space can generate up to ~30,000 candidates for complex kernels.
- Phase 1 (symbolic synthesis) is fast (<82 seconds). Phase 2 (compilation and execution) dominates the time (up to ~1.7 hours for complex LLM layers). The paper notes this is a one-time cost, comparable to other schedule-based compilers, and optimized kernels can be cached for reuse.

Significance and Claims

The paper claims that Axon represents a shift from heuristic-based or template-driven optimization to a semantics-driven, synthesizing approach. By leveraging SMT solvers to guarantee semantic equivalence, Axon can automatically discover complex algebraic transformations (such as reordering operations across different compute engines) that are difficult to encode manually.

The significance lies in:

Automation: Reducing the burden on programmers to manually tune kernels for specific hardware.
Correctness: Providing provable guarantees for transformations, unlike probabilistic testing methods.
Performance: Achieving significant speedups on modern AI accelerators by jointly optimizing tiling, instruction selection, and fusion in a unified search space.
Generality: The ability to handle complex, multi-operator LLM kernels that existing search-based compilers (like Mirage) cannot synthesize.

The authors emphasize that while the compilation overhead is non-trivial, it is a standard one-time cost for optimizing compilers, and the resulting performance gains justify the investment, particularly for large-scale LLM workloads.