# Axe: A Simple Unified Layout Abstraction for Machine Learning Compilers

arXiv: https://arxiv.org/abs/2601.19092 · HTML: https://mlsanity.com/paper/2601.19092/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 17 of its 40 references that are in the map.
- First posted on arXiv 27 Jan 2026.

## Builds on (most cited first, 17 of 17)

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Qwen3 Technical Report](https://mlsanity.com/paper/2505.09388/) (2025)
- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [Gemma 2: Improving Open Language Models at a Practical Size](https://mlsanity.com/paper/2408.00118/) (2024)
- [SGLang: Efficient Execution of Structured Language Model Programs](https://mlsanity.com/paper/2312.07104/) (2023)
- [In-Datacenter Performance Analysis of a Tensor Processing Unit](https://mlsanity.com/paper/1704.04760/) (2017)
- [FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving](https://mlsanity.com/paper/2501.01005/) (2025)
- [Beyond Data and Model Parallelism for Deep Neural Networks](https://mlsanity.com/paper/1807.05358/) (2018)
- [Mesh-TensorFlow: Deep Learning for Supercomputers](https://mlsanity.com/paper/1811.02084/) (2018)
- [GSPMD: General and Scalable Parallelization for ML Computation Graphs](https://mlsanity.com/paper/2105.04663/) (2021)
- [TileLang: A Composable Tiled Programming Model for AI Systems](https://mlsanity.com/paper/2504.17577/) (2025)
- [TensorIR: An Abstraction for Automatic Tensorized Program Optimization](https://mlsanity.com/paper/2207.04296/) (2022)
- [Triton-distributed: Programming Overlapping Kernels on Distributed AI Systems with the Triton Compiler](https://mlsanity.com/paper/2504.19442/) (2025)
- [TileLink: Generating Efficient Compute-Communication Overlapping Kernels using Tile-Centric Primitives](https://mlsanity.com/paper/2503.20313/) (2025)
- [Linear Layouts: Robust Code Generation of Efficient Tensor Computation Using mathbbF2](https://mlsanity.com/paper/2505.23819/) (2025)

## Built on by (most cited first, 3 of 3)

- [CAKE: Compiler-Agent Co-Design for Frontier Kernel Evolution](https://mlsanity.com/paper/2608.12629/) (2026)
- [Event Tensor: A Unified Abstraction for Compiling Dynamic Megakernel](https://mlsanity.com/paper/2604.13327/) (2026)
- [Ada-MK: Adaptive MegaKernel Optimization via Automated DAG-based Search for LLM Inference](https://mlsanity.com/paper/2605.11581/) (2026)

## Abstract (arXiv)

Scaling modern deep learning workloads demands coordinated placement of data and compute across device meshes, memory hierarchies, and heterogeneous accelerators. We present Axe Layout, a hardware-aware abstraction that maps logical tensor coordinates to a multi-axis physical space via named axes. Axe unifies tiling, sharding, replication, and offsets across inter-device distribution and on-device layouts, enabling collective primitives to be expressed consistently from device meshes to threads. Building on Axe, we design a multi-granularity, distribution-aware DSL and compiler that composes thread-local control with collective operators in a single kernel. Experiments show that our unified approach can bring performance close to hand-tuned kernels on across latest GPU devices and multi-device environments and accelerator backends.