Other hardware stacks

Nothing about the tile abstraction is NVIDIA-specific, and the clearest evidence comes from Pallas, JAX’s extension for “kernel programming for both GPUs and TPUs using a Triton-like model.” Its design document makes a pointed observation about Triton itself: Triton “exposes a TPU-like programming model to users, i.e. writing programs for tiles of arrays in L1-cache,” and yet is specialized enough to GPU that it cannot be compiled directly for TPU; Triton’s atomic operations, built for parallel writes, “don’t necessarily make sense on TPU.” So Pallas keeps only the tile-based programming model, abstracts the platform details behind it, and lowers the same kernel per backend: to Mosaic GPU (formerly Triton) on GPUs, and to Mosaic on TPUs. The tile is the portable part; what surrounds it is not.

What Pallas adds around that tile is exactly three things, and the design document is emphatic that it is otherwise “just JAX.” First, reference types: a kernel receives Refs rather than immutable arrays, which “gives users more precise control over memory access and layout,” read and written in place with NumPy-style indexing, and output Refs are handed in alongside the input ones. Reading from a Ref “corresponds to loading an array into the lowest level of the memory hierarchy,” L1 cache on GPU and vector registers on TPU, which is this chapter’s tile wearing a JAX name. Second, a restricted subset of JAX primitives plus Pallas-specific ones, among them a program_id the document admits is borrowed terminology from Triton. Third, pallas_call, the higher-order function that runs the kernel, “analogous to pmap or shard_map, except with references to shared memory,” which loops over a grid and slices each iteration’s inputs through a BlockSpec carrying a block shape and an index map from loop indices to block indices.

The third one is the answer to why a kernel language would be embedded in an array framework rather than built beside it. The stated motivation is that “JAX transformations are key to its success,” so the same vmap and jvp that transform ordinary JAX programs can be applied to a kernel the user wrote by hand. A Triton kernel is a leaf the framework calls into; a Pallas kernel is a JAX program, and the transformations do not stop at its edge.

AMD’s stack arrived at the same shape from two directions. Composable Kernel, in AMD’s own description, “provides a programming model for writing performance-critical kernels for machine learning workloads across multiple architectures,” written in general purpose kernel languages such as HIP C++ and resting on two ideas: “a tile-based programming model” and a complexity-reduction technique it calls Tensor Coordinate Transformation. The library names four layers, ordered lowest to highest: Templated Tile Operators, Templated Kernel and Invoker, Instantiated Kernel and Invoker, and Client API. That is the CUTLASS decomposition again under different names, an atom-like tile operator at the bottom and a ready-made handle at the top, and the repository documents two operation sets against it, one for CK and a separate one for CK Tile. HipKittens, a research library of C++ tile primitives for AMD GPUs, states this section’s through-line as an experimental finding: porting from its NVIDIA sibling ThunderKittens, the core tile and bulk compute interfaces carry over, while the decisions around memory access patterns, compute and memory scheduling, and thread block ordering within the chiplet architecture differ. Its tiles are sized to the tensor core units, and its two core scheduling patterns, 8-wave ping pong and 4-wave interleave, are organized around CDNA’s waves rather than the warps of Foundations.

Where those AMD kernels end up is its own data point. AITER, the AI Tensor Engine for ROCm, is “AMD’s high-performance AI operator library, providing optimized GPU kernels for inference and training workloads on ROCm,” positioned as “a unified collection of production-ready operators that framework developers can integrate directly into their stacks.” What makes it relevant here is its backend list: the same operator can be served by a Triton kernel, a Composable Kernel implementation, or hand-tuned assembly, behind one C++ or Python API. The programming models this chapter has covered are not competing endpoints so much as interchangeable suppliers to an operator library, and AITER’s README notes it is the default attention backend for vLLM on ROCm, which is where a kernel chosen by that machinery actually meets production traffic.

AWS Trainium is the strongest test of the claim, because the accelerator underneath is not a GPU. NKI, the Neuron Kernel Interface for writing kernels that run on Trainium devices, still hands the programmer a tile: a kernel allocates tiles in on-chip SBUF memory, DMA-copies inputs into them from HBM, and checks that a tile’s first dimension fits within the on-chip tile size limit before operating on whole tiles at a time. What changes is everything around the tile. NKI kernels use a sequential programming model rather than a grid of parallel blocks, and its nki.isa functions are “designed to expose the underlying hardware capabilities in as direct a way as possible,” each call running one operation on one of the device’s compute engines while the compiler unrolls, inlines, and resolves everything else ahead of time. Across all four stacks the tile survives even where warps, blocks, and threads do not; what each stack builds around the tile tracks what its particular silicon makes cheap or expensive. See further reading for each stack’s documentation.