# Other hardware stacks · Programming models and profiling

<!-- https://learn-kernels.com/chapters/programming-models/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.

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** (A technique Composable Kernel pairs with its tile-based programming model to reduce the algorithmic complexity of complex machine learning operators.). It is layered from templated tile operators up through a client API, a decomposition recognizable from CUTLASS. 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 [chapter 1’s hardware](https://learn-kernels.com/chapters/hardware).

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** (NKI's execution contract: the logical order of operations follows the syntactic order of statements in the kernel, and the compiler may only reorder operations that have no data dependencies.) 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](https://learn-kernels.com/chapters/reading#other-hardware-stacks) for each stack’s documentation.

Source

Pallas framing, the Triton-on-TPU argument, and lowering targets from the [Pallas design document](https://docs.jax.dev/en/latest/pallas/design/design.html). Composable Kernel description and layering from the [ROCm/composable\_kernel](https://github.com/ROCm/composable_kernel) README. HipKittens primitives and scheduling patterns from the [HazyResearch/HipKittens](https://github.com/HazyResearch/HipKittens) README. AITER description, backends, and vLLM integration from the [ROCm/aiter](https://github.com/ROCm/aiter) README. NKI model and quotes from the [NKI language guide](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/nki/programming_model.html).
