Axe: A Simple Unified Layout Abstraction for Machine Learning Compilers

TL;DR

Axe introduces a simple, hardware-aware layout abstraction that maps logical tensor coordinates to a multi-axis physical space using named axes for sharding, replication, and offsets. Building on this, the authors develop a multi-granularity, distribution-aware DSL and compiler that seamlessly combines thread-local control with collective operators to generate efficient code across GPUs and AI accelerators. The approach achieves near hand-tuned performance on modern devices and demonstrates clear advantages in multi-GPU and heterogeneous backends, with substantial reductions in development effort. Overall, Axe provides a practical foundation for unifying layout semantics across distributed and heterogeneous hardware, enabling productive and high-performance ML compilation.

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.

Figures (13)

• Figure 1: Elements of Axe Layout. An iter specifies a triple $(\text{extent}, \text{stride}, \text{axis})$ and defines a linear, strided access on that axis. A list of iters forms the shard part $\mathbf{D}$, a set of iters forms the replica part $\mathbf{R}$, and $\mathbf{O}$ is a fixed offset.
• Figure 2: Examples of Axe layouts across various scenarios. Each column shows the logical tensor shape and the mapped physical axis values. Axes are color-coded. Left. Mapping a logical $8\times 16$ tile to 4 GPU warps with 32 lanes and 2 registers each; 2 warps sharded and 2 warps replicated. Middle. Distributed sharding of a $64\times 128$ matrix across 4 GPUs; the top uses full sharding across 4 GPUs, the bottom uses a $2\times 2$ mesh with shards and replicas. Right. Native hardware memories; the top depicts an AI accelerator 2D-partitioned scratchpad SRAM, the bottom shows NVIDIA Blackwell 2D tensor memory.
• Figure 3: Axe compiler overview.Left: a GEMM kernel written in the Axe compiler DSL. Execution scopes, tensors with Axe layouts, and operators are highlighted. The program uses load and gemm_warp macros (expanded when parsing into IR) and a three-stage pipeline with prologue, main, and epilogue. We omit several lines (__syncthreads() and tensor allocations) for brevity. Right, top: tensors carry Axe layouts in shared memory and registers. Right, middle: Use tile to compose a register tile with lane to form a warp view. Right, bottom: the copy.async operator is lowered to a thread-bound loop that issues cp.async.cg.shared.global, with addresses derived from the layouts. Together, these steps show how Axe couples multi-granularity programming with layout-driven code generation.
• Figure 4: An example Axe DSL snippet showing thread-level loop transformation and thread binding. In actual CuTe programs, C_slice is derived by combinations of partition APIs.
• Figure 5: An example Axe DSL snippet showing thread-block collective semantics.

Theorems & Definitions (46)

• Definition 2.1: Iter
• Definition 2.2: Layout
• Definition 2.3: Induced map
• Proposition 1.1
• Theorem 1.2: Canonical form uniqueness under GC
• Lemma 1.3: Exact digit identity
• proof
• Lemma 1.4: Axis–wise coefficient expansion
• proof
• Lemma 1.5: First–difference periodicity