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.
@article{arxiv.2601.19092,
title = {Axe: A Simple Unified Layout Abstraction for Machine Learning Compilers},
author = {Bohan Hou and Hongyi Jin and Guanjie Wang and Jinqi Chen and Yaxing Cai and Lijie Yang and Zihao Ye and Yaoyao Ding and Ruihang Lai and Tianqi Chen},
journal= {arXiv preprint arXiv:2601.19092},
year = {2026}
}