Axe: A Simple Unified Layout Abstraction for Machine Learning Compilers

2026/01/27 by Bohan Hou, Hongyi Jin, Guanjie Wang +7 · 1 voice · 1 citation
Computer Science · #Abstraction #Abstraction layer #Advanced Neural Network Applications #Cloud Computing and Resource Management #Compiler #Digital subscriber line #Parallel Computing and Optimization Techniques #Polygon mesh #Scaling #Simple (philosophy) #Tensor (intrinsic definition) #cs.AI #cs.DC #cs.LG #cs.PL

paper · pdf · doi:10.48550/arxiv.2601.19092

openalex publication_date 2026/01/27 · arxiv published 2026/01/27 · arxiv updated 2026/01/28 · openalex created_date 2026/01/29 · openalex updated_date 2026/07/28

Abstract

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.

Citations

• Linear Layouts: Robust Code Generation of Efficient Tensor Computation Using \mathbbF₂
• Qwen3 Technical Report
• Triton-distributed: Programming Overlapping Kernels on Distributed AI Systems with the Triton Compiler
• TileLang: A Composable Tiled Programming Model for AI Systems
• TileLink: Generating Efficient Compute-Communication Overlapping Kernels using Tile-Centric Primitives
• FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving
• The Llama 3 Herd of Models
• Gemma 2: Improving Open Language Models at a Practical Size
• SGLang: Efficient Execution of Structured Language Model Programs
• Efficient Memory Management for Large Language Model Serving with PagedAttention
• TensorIR: An Abstraction for Automatic Tensorized Program Optimization
• GSPMD: General and Scalable Parallelization for ML Computation Graphs
• Array Programming with NumPy
• Language Models are Few-Shot Learners
• PyTorch: An Imperative Style, High-Performance Deep Learning Library
• Mesh-TensorFlow: Deep Learning for Supercomputers
• Beyond Data and Model Parallelism for Deep Neural Networks
• In-Datacenter Performance Analysis of a Tensor Processing Unit
• TensorFlow: A system for large-scale machine learning

Cited by

• Ada-MK: Adaptive MegaKernel Optimization via Automated DAG-based Search for LLM Inference

Discussions

• Axe: A Simple Unified Layout Abstraction for Machine Learning Compilers [hn, 1 points, 0 comments]

Related

• Stripe: Tensor Compilation via the Nested Polyhedral Model
• TileLink: Generating Efficient Compute-Communication Overlapping Kernels using Tile-Centric Primitives
• A High-Level Compiler Integration Approach for Deep Learning Accelerators Supporting Abstraction and Optimization
• FLUX: Fast Software-based Communication Overlap On GPUs Through Kernel Fusion
• ISA Mapper: A Compute and Hardware Agnostic Deep Learning Compiler