Axe: A Simple Unified Layout Abstraction for Machine Learning Compilers
Not on the map. Why?
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.
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
17 of 40 references.
References · 17
- 2020Language Models are Few-Shot Learnerscited by 32,419
- 2019PyTorch: An Imperative Style, High-Performance Deep Learning Librarycited by 20,305
- 2024The Llama 3 Herd of Modelscited by 14,715
- 2025Qwen3 Technical Reportcited by 8,820
- 2023Efficient Memory Management for Large Language Model Serving with PagedAttentioncited by 5,944
- 2024Gemma 2: Improving Open Language Models at a Practical Sizecited by 1,998
- 2023SGLang: Efficient Execution of Structured Language Model Programscited by 1,141
- 2017In-Datacenter Performance Analysis of a Tensor Processing Unitcited by 618
- 2025FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Servingcited by 258
- 2018Beyond Data and Model Parallelism for Deep Neural Networkscited by 179
- 2018Mesh-TensorFlow: Deep Learning for Supercomputerscited by 169
- 2021GSPMD: General and Scalable Parallelization for ML Computation Graphscited by 111
5 more
- 2025TileLang: A Composable Tiled Programming Model for AI Systemscited by 43
- 2022TensorIR: An Abstraction for Automatic Tensorized Program Optimizationcited by 38
- 2025Triton-distributed: Programming Overlapping Kernels on Distributed AI Systems with the Triton Compilercited by 25
- 2025TileLink: Generating Efficient Compute-Communication Overlapping Kernels using Tile-Centric Primitivescited by 24
- 2025Linear Layouts: Robust Code Generation of Efficient Tensor Computation Using mathbbF2cited by 4
Built on by
Citing papers per year: 2026: 3.
3 papers.