Axe: A Simple Unified Layout Abstraction for Machine Learning Compilers
Fuente:
arXiv
Enregistré dans:
| Auteurs principaux: | , , , , , , , , , |
|---|---|
| Format: | Preprint |
| Publié: |
2026
|
| Sujets: | |
| Accès en ligne: | |
| Tags: |
Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
|
| _version_ | 1866910004164952064 |
|---|---|
| author | Hou, Bohan Jin, Hongyi Wang, Guanjie Chen, Jinqi Cai, Yaxing Yang, Lijie Ye, Zihao Ding, Yaoyao Lai, Ruihang Chen, Tianqi |
| author_facet | Hou, Bohan Jin, Hongyi Wang, Guanjie Chen, Jinqi Cai, Yaxing Yang, Lijie Ye, Zihao Ding, Yaoyao Lai, Ruihang Chen, Tianqi |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2601_19092 |
| institution | arXiv |
| publishDate | 2026 |
| record_format | arxiv |
| spellingShingle | Axe: A Simple Unified Layout Abstraction for Machine Learning Compilers Hou, Bohan Jin, Hongyi Wang, Guanjie Chen, Jinqi Cai, Yaxing Yang, Lijie Ye, Zihao Ding, Yaoyao Lai, Ruihang Chen, Tianqi Distributed, Parallel, and Cluster Computing Artificial Intelligence Machine Learning Programming Languages 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. |
| title | Axe: A Simple Unified Layout Abstraction for Machine Learning Compilers |
| topic | Distributed, Parallel, and Cluster Computing Artificial Intelligence Machine Learning Programming Languages |
| url | https://arxiv.org/abs/2601.19092 |