Overview of AI Hardware: GPUs, TPUs, and AWS Trainium #

Before we can talk about teaching an LLM to write system kernels, we have to be precise about what a kernel runs on. Every kernel targets a specific piece of silicon with a specific set of compute units, a specific on-chip memory layout, and a specific way of moving data around. This section builds a mental model of the hardware that the rest of the tutorial assumes. We will look at three representative tensor accelerators — NVIDIA GPUs, Google TPUs, and AWS Trainium. Despite different vendors and programming models, they share the same three architectural ingredients. Those three ingredients are exactly what a kernel has to reason about, which is why understanding them here sets up the kernel-programming discussion in Section 1.3.

Why specialized AI hardware, and why it is hard #

The bulk of the work in deep neural networks is dense linear algebra (matrix multiplications and convolutions) over a small, repetitive set of operators (matmul, elementwise activations, normalization, reductions), often in reduced precision (FP16, BF16, FP8, INT8). This regularity is what makes specialization pay off: instead of a general-purpose CPU that spends most of its transistors and energy on control logic, branch prediction, and caches to run arbitrary code fast, an AI accelerator can pour its silicon into wide arrays of multiply-accumulate units that do one thing — tensor math — extremely efficiently per watt.

This opportunity has produced a crowded AI landscape, and a recurring theme today is that most of the participants are “hardware rich, software poor”. A chip can be genuinely superior on paper — more FLOPs, better performance per watt, lower cost — and still fail to realize these gains because it does not have effective software and tooling. Along with the hardware, there should also be a way to give programmers a way to actually extract performance from the chip.

Writing that software is hard for reasons that are structural, not incidental:

Both software and hardware are evolving: The dominant model family is evolving: MLPs, then CNNs, then Transformers/LLMs; and each shift stresses the hardware differently. Moreover, vendors ship new architectures on a fast cadence (for NVIDIA, Ampere → Hopper → Blackwell), each adding new compute units, datatypes, and memory features.
Hardware/software co-design is common. Techniques like quantization and sparsity blur the line between the algorithm and the silicon, so the software cannot treat the hardware as an opaque black box.

As a consequence, high-level languages and AI frameworks (PyTorch, JAX, HuggingFace Transformers, and others) must be lowered onto a different low-level, performance-oriented programming interface for each backend, and code written for one rarely ports cleanly to another. That software-lowering story is the subject of Section 1.2, and the per-backend kernel-programming interfaces are the subject of Section 1.3. Here we stay focused on the silicon itself.

A unifying mental model of a tensor accelerator #

Under the hood, most tensor accelerators are built to exploit the same regular, tensor-based math over a limited operator set. Vendor-specific details become variations of the following:

1. Compute engines. At the heart of the chip is a matrix/tensor unit — frequently a systolic array — that performs matrix-multiply-accumulate on fixed-size tiles (common tile granularities are on the order of 16×16 or 128×128). Around it sit a vector unit for reductions and elementwise operations, and a scalar unit for the remaining scalar/control work and activation functions. Almost every DNN kernel is ultimately a schedule that keeps the matrix unit fed while the vector and scalar units handle the epilogue and glue.

2. An explicitly-managed memory hierarchy. Off-chip memory is large but far away (high-bandwidth memory, HBM, or DRAM). On-chip memory is small but fast: a software-managed scratchpad SRAM, plus an accumulator (or register file) where partial sums from the matrix unit are collected. Crucially, on accelerators this hierarchy is often explicitly managed rather than governed by a hardware cache: the program (or compiler, or kernel author) decides what to stage on-chip and when. The central performance concern is therefore minimizing data movement — reusing data while it is resident in fast memory — and overlapping computation with data movement so the compute engines never stall waiting for HBM.

3. Interconnects for multi-chip scaling. A single chip is rarely enough for large models, so accelerators are wired together with high-bandwidth interconnects into larger systems, letting many chips cooperate on one computation.

A schematic view of a single accelerator chip:

        ┌───────────────────────────────────────────────┐
        │                Off-chip HBM / DRAM            │   large, slow
        └───────────────────────────────────────────────┘
                     ▲                     │
        data movement│                     │ (explicitly managed)
                     │                     ▼
        ┌───────────────────────────────────────────────┐
        │            On-chip scratchpad SRAM            │   small, fast
        └───────────────────────────────────────────────┘
                     ▲            ▲             ▲
                     │            │             │
             ┌───────────┐ ┌────────────┐ ┌────────────┐
             │  Matrix / │ │   Vector   │ │   Scalar   │   compute
             │  Tensor   │ │    unit    │ │    unit    │   engines
             │   unit    │ │(reductions,│ │(activations│
             │ (GEMM /   │ │ elementwise│ │  control)  │
             │ systolic) │ │    ops)    │ │            │
             └───────────┘ └────────────┘ └────────────┘
                     │
                     ▼
             ┌───────────────┐
             │  Accumulator  │   partial sums
             └───────────────┘

  ... and chip ⇄ chip over a high-bandwidth interconnect (scale-out)

The explicitly-managed memory hierarchy is worth visualizing as a pyramid, because the trade-off it encodes is the same on every accelerator. Memories closer to the compute engines offer far higher bandwidth and lower latency but hold far less; memories farther away hold the whole working set but are slow to reach. The figure below draws this pyramid for a Trainium NeuronCore, but the shape is what matters and it generalizes to GPUs and TPUs alike. Since compute is usually cheap relative to moving data, keeping the working set high in the pyramid — through tiling, reuse, and fusion — is the single most important lever a kernel has.

The memory pyramid, drawn here for a Trainium NeuronCore but representative of the pattern on GPUs and TPUs alike: moving closer to the compute engines buys higher bandwidth and lower latency at the cost of much smaller capacity. The numbers are illustrative, not exact (credit: AWS Neuron Docs).

Designs differ enormously in size (edge NPU vs. datacenter GPU vs. wafer-scale part), in dataflow (how tiles stream through the matrix unit), and in programming model (hand-written kernels vs. whole-graph compilation). But the three ingredients — engines, memory hierarchy, interconnect — recur. We now instantiate this model for each of the three systems.

GPUs (NVIDIA) #

A modern NVIDIA GPU is built from many Streaming Multiprocessors (SMs). Each SM contains two kinds of compute engines that map onto our model:

CUDA cores: scalar/vector ALUs that execute general-purpose arithmetic — the “vector + scalar unit” side of the picture.
Tensor Cores: dedicated matrix-multiply-accumulate units for GEMM. These are the GPU’s “matrix unit,” operating on small tiles in reduced precision.

Execution follows the SIMT (Single Instruction, Multiple Threads) model. The programmer writes the code for a single thread — a scalar-looking program such as “compute one output element” — and the GPU runs that same program across a huge number of threads at once, each with its own registers and its own slice of the data. This gives the convenience of scalar per-thread code with the throughput of massive parallelism, and it is the key distinction from pure SIMD, where one instead writes vector instructions over vector registers explicitly. Those threads are not a flat pool, but instead, are organized into a three-level hierarchy, each level of which corresponds to a physical or scheduling reality on the chip.

Warps (32 threads, one instruction stream). The hardware does not schedule threads individually. It bundles them into fixed groups of 32 called warps, and the warp (not the thread) is the true unit of execution. All 32 threads in a warp share a single instruction stream: the ordered flow of machine instructions driven by one program counter. The GPU fetches and decodes one instruction and issues it to all 32 lanes together (“lockstep”), so each lane runs the same instruction on its own data. Sharing the stream is what makes SIMT efficient — the cost of fetching and decoding an instruction is paid once and amortized across 32 threads — and it is also the source of a classic performance pitfall. If threads in the same warp take different sides of an if/else, there is still only one program counter to go around, so the hardware must run both paths, masking off the non-participating threads on each. This warp divergence serializes what looked like parallel branches, which is why keeping the 32 threads of a warp on the same control-flow path (and, touching contiguous addresses) matters for performance.

Thread blocks / CTAs (cooperating warps on one SM). Warps are grouped into thread blocks, also called CTAs (cooperative thread arrays). The “cooperative” is the point: an entire block is guaranteed to be scheduled onto a single SM, and because its threads all live on the same SM they can genuinely work together — sharing that SM’s fast, software-managed on-chip shared memory and synchronizing at barriers. A block cannot span two SMs, which is exactly what makes the single-SM placement a hard guarantee. This is the level at which a kernel author chooses tile sizes: how large a chunk of the problem one block should stage into shared memory and process cooperatively.

Grid (independent blocks covering the problem). All the thread blocks together form a grid, sized so that the blocks collectively cover the entire problem. Different blocks are largely independent — they do not cheaply share memory or synchronize with one another — so the GPU is free to run them in any order, across whatever SMs are available, even sequentially if there are more blocks than SMs. That independence is precisely what lets the same kernel scale unchanged from a small GPU to a large one.

A kernel author reasons explicitly about this hierarchy when deciding how to partition work: how many threads per block and how many blocks to launch, how each thread’s index maps to the data it owns, what to stage into a block’s shared memory, and where to synchronize while avoiding divergence. Mapping the computation onto warps, blocks, and the grid is a large part of what writing a GPU kernel means.

The GPU memory hierarchy is a concrete instance of the explicitly-managed hierarchy above:

per-thread registers            (fastest, private)
        ↓
per-SM shared memory / L1       (software-managed, shared within a block)
        ↓
shared L2 cache                 (across SMs)
        ↓
global memory (HBM)             (largest, slowest, off-chip)

Getting performance means keeping working data in registers and shared memory, and issuing coalesced global-memory accesses so that the threads of a warp touch contiguous addresses. Shared memory in particular is software-managed: the kernel decides what tiles to stage there, mirroring the scratchpad in our general model. GPUs support a wide range of precisions — FP32, TF32, FP16, BF16, FP8, INT8, and more — which the Tensor Cores exploit for throughput.

For multi-GPU scaling, NVIDIA uses NVLink and NVSwitch to provide high-bandwidth GPU-to-GPU communication (with PCIe as the more general host/device interconnect). Across generations — Ampere → Hopper → Blackwell — the recurring story is more and faster Tensor Cores, new datatypes, larger and higher-bandwidth on-chip and off-chip memory, and faster interconnect. Each generation shifts the performance sweet spot, which is part of why kernels must be re-tuned over time.

TPUs (Google) #

Google’s Tensor Processing Unit (TPU) is a domain-specific accelerator built around a large systolic-array Matrix Multiply Unit (MXU). In a systolic array, data flows rhythmically through a grid of multiply-accumulate cells, so a large matrix multiply is performed with very high arithmetic density and minimal control overhead. Alongside the MXU sit a vector unit for elementwise and reduction work and a scalar unit for control — again matching the matrix/vector/scalar decomposition of our model. Explicitly-managed on-chip memory (a scratchpad, e.g. VMEM) is fed from off-chip HBM, and the MXU accumulates partial sums as tiles stream through.

Where the TPU differs most sharply from the GPU is the programming model. TPUs are programmed primarily through whole-graph compilation rather than by hand-writing individual kernels. A model expressed in JAX, TensorFlow, or PyTorch/XLA is captured as a computation graph and handed to the XLA compiler, which fuses operations, assigns layouts, and emits an optimized executable for the whole graph at once. The programmer usually reasons at the level of tensor operations and lets the compiler own tiling and scheduling — a very different experience from writing SIMT kernels. (Section 1.2 covers this XLA/HLO path in detail.)

For scale-out, TPUs are assembled into pods connected by a dedicated high-bandwidth inter-chip interconnect (ICI) arranged in a torus topology, so that many chips form a tightly-coupled machine for large training and inference jobs. TPUs have gone through multiple generations (v2, v3, v4, v5e/v5p, and newer), broadly trending toward more matrix throughput, larger memory, and larger, faster-connected pods.

AWS Trainium #

AWS Trainium is a family of tensor accelerators designed by Annapurna Labs (AWS). Trainium chips are exposed to users through Trn EC2 instances (for example, a trn1.2xlarge instance exposes two NeuronCore-v2 cores). Each NeuronCore contains several heterogeneous compute engines that map cleanly onto the matrix/vector/scalar decomposition — plus a fully programmable “catch-all” engine:

Tensor Engine: the primary high-throughput unit. It is built around a large systolic array (a 128×128 grid of processing elements) and performs matrix multiplies, convolutions, and transposes at very high arithmetic intensity, in mixed precision (BF16, FP16, TF32, FP8, FP32, MXFP8/MXFP4). Because attention projections, feed-forward layers, and most other transformer work reduce to GEMMs, the Tensor Engine carries the vast majority of a model’s FLOPs. It is the core’s matrix unit.
Vector Engine: elementwise and vectorized operations in which each output element depends on multiple input elements — axpy (Z = aX + Y), LayerNorm, reductions, pooling. Offloading this non-matmul work to dedicated hardware lets it overlap with the Tensor Engine (avoids stalls).
Scalar Engine: per-element operations in which each output element depends on a single input element (activations, absolute value, and similar), with lightweight per-element control logic. Isolating this work onto its own pipeline keeps it from stalling the Tensor or Vector engines.
GPSIMD Engine: the flexible catch-all. It is a set of fully programmable wide vector processors (eight 512-bit-wide units) that can run general-purpose C code and access on-chip SRAM directly, which is what makes it possible to implement custom operators that do not map cleanly onto the tensor, vector, or scalar engines.

A single NeuronCore: the Tensor, Vector, Scalar, and GPSIMD engines arranged around software-managed on-chip SRAM (credit: AWS Neuron Docs).

Feeding these engines is an explicitly-managed on-chip memory hierarchy with two named components: SBUF, a software-managed scratchpad SRAM fed from off-chip HBM via DMA, and PSUM, a dedicated accumulator buffer written by the Tensor Engine with matrix-multiply partial sums. The Vector, Scalar, and GPSIMD engines can read and write SBUF, while the Tensor Engine reads its inputs from SBUF and writes its outputs only to PSUM; inputs are explicitly loaded HBM→SBUF before compute and results written SBUF→HBM afterward, so managing tile lifetimes to avoid spills back to HBM is a central concern. PSUM is smaller and specialized: it lets the Tensor Engine accumulate partial results from many matmul tiles into the same output region, with completed tiles evicted to SBUF. A defining characteristic of the architecture is that it works efficiently on groups of 128 partitions: SBUF and PSUM are organized as 128-partition memories, so tiling computations in units of 128 is central to mapping work onto the hardware, analogous to (but distinct from) the tile sizes seen on GPUs and TPUs.

              HBM  (off-chip)
                │        ▲
    DMA load    ▼        │  store
              SBUF  (on-chip scratchpad SRAM, 128 partitions)
        ┌───────┼───────────────┬───────────────┐
        ▼       ▼               ▼               ▼
   Tensor    Vector          Scalar          GPSIMD
   Engine    Engine          Engine          Engine
   (matmul,  (reduce,        (activations)   (general SIMD)
    systolic  elementwise)
    array)
        │
        ▼
      PSUM  (accumulator buffer for matmul partial sums)

For multi-chip scaling, Trainium follows the same pattern as the GPU and TPU: chips are joined by a dedicated chip-to-chip interconnect (NeuronLink) so that the accelerators in and across instances can cooperate, and Trn instances scale out over the datacenter network using AWS’s Elastic Fabric Adapter (EFA). Because it is easy to underappreciate the importance of interchip communication from a single-chip view, we give it its own subsection below, using Trainium as a concrete case study.

Trainium is programmed through the AWS Neuron SDK, which — like the other stacks — offers two levels of access:

A high-level, XLA-based graph-compiler path, reached from PyTorch or JAX via the framework integrations (torch-neuronx / jax-neuronx) and compiled by the Neuron graph compiler (neuronx-cc). This traces the model graph and fuses operations automatically, much like the TPU’s whole-graph approach.
A lower-level kernel path, NKI (the Neuron Kernel Interface), a Python-embedded DSL for hand-writing high-performance kernels with explicit control over the engines and the SBUF/PSUM/HBM hierarchy.

Scaling out: from chip to server to datacenter #

The interconnect is the one that a single-chip diagram makes easy to forget, yet it shapes what a large-model kernel and its collective operations can assume. It is cleanest to see as a three-tier hierarchy: on-chip execution within one accelerator, a scale-up fabric that binds several accelerators inside a server, and a scale-out network that binds many servers into a datacenter. We use Trainium as the concrete case study here, but the same three tiers appear on GPUs (NVLink/NVSwitch inside a node, InfiniBand/Ethernet across nodes) and TPUs (an ICI torus within a pod).

The physical chip. A Trn3 chip packages eight NeuronCores together with four banks of HBM on a single substrate. The compute die sits at the center with the HBM stacks placed close around it for bandwidth; inside the die the cores are laid out in a regular grid and joined by a high-bandwidth on-chip interconnect so data and synchronization move efficiently across the chip. The whole package is deployed as a PCIe card/module in a server, exposing high-speed links both to the host and to its neighbor chips.

Schematic of a Trn3 chip: 8 NeuronCores and 4 HBM banks

A Trn3 chip: eight NeuronCores plus four banks of HBM on one package — photo (left) and schematic (right) (credit: SemiAnalysis & AWS Neuron Docs).

Scale-up: aggregating accelerators inside a server. One chip is rarely enough, so a server wires several accelerators together with high-bandwidth, low-latency links, pooling their compute and memory.

A Trn3 ultraserver with several connected chips

A Trn3 ultraserver aggregating several Trn3 chips over a high-bandwidth scale-up interconnect (credit: AWS re:Invent 2025).

There are two broad ways to wire that scale-up fabric, and the choice is a general trade-off between structured and irregular communication rather than anything specific to one workload:

Direct topology (torus / mesh). Each accelerator is wired to a small, fixed set of neighbors in a 2D or 3D grid with wrap-around links, and traffic to distant devices is routed hop by hop. This is hardware-efficient (no switches), cheap to scale, and gives predictable nearest-neighbor bandwidth, which suits structured, regular communication such as tensor- or pipeline-parallel exchanges between fixed groups of devices. The cost is that latency grows with distance and the aggregate bandwidth between arbitrary pairs of devices is limited by the topology.

Scaling up with a 3D torus: each accelerator links directly to a few neighbors with wrap-around edges (Trn2 ultraserver shown; credit: AWS).

Switched fabric. Each accelerator instead connects to one or more high-radix switches, so any device can reach any other in roughly a single hop. This delivers higher bisection bandwidth and lower, more uniform latency, and it handles irregular or global communication patterns far better — at the cost of extra switch hardware, cabling, and power, larger fault domains (one switch touches many devices), and more latency variability under contention. It is the same idea as NVSwitch on GPU servers.

Compute trays connected to switching trays in a switched fabric

A switched fabric: compute trays connected through dedicated switching trays, giving any-to-any communication in about one hop (credit: AWS).

Schematic of a Trn3 NL 32x2 server: two racks combining compute trays, switching trays, CPU, and power (credit: AWS Neuron Docs).

Scale-out: many servers into a datacenter. Once a model or a serving fleet outgrows a single server, servers are joined across the datacenter with high-speed Ethernet fabrics such as AWS’s Elastic Fabric Adapter (EFA). Latency here is orders of magnitude higher than on-chip or intra-server communication, so this tier is used for data-parallel replication, sharding, and coordination rather than for the tightest inner loops. At this scale, network latency, bandwidth, and tail behavior become first-class performance concerns, not just FLOPs and memory bandwidth.

Datacenter-scale cluster of accelerator servers

Scale-out: many accelerator-equipped servers joined into a single datacenter-scale training and inference fabric (credit: AWS, Project Rainier).

Taken together, the interconnect is best understood as this hierarchy — on-chip execution, a scale-up fabric inside the server, and a scale-out network across the datacenter. Kernels themselves live at the bottom tier, but the tiers above decide which collective operations are cheap and which are expensive, and that is exactly the kind of structure a kernel author (or an LLM writing kernels) has to keep in mind.

Synthesis: what a kernel actually targets #

Step back and the three systems tell one story. A GPU’s Tensor Cores + CUDA cores + registers/shared memory/L2/HBM + NVLink; a TPU’s MXU + vector/scalar units + VMEM/HBM + ICI; a Trainium NeuronCore’s Tensor/Vector/Scalar/GPSIMD engines + SBUF/PSUM/HBM + NeuronLink — these are similar ingredients:

Ingredient	NVIDIA GPU	Google TPU	AWS Trainium
Matrix unit	Tensor Cores	MXU (systolic array)	Tensor Engine
Vector / scalar	CUDA cores	vector + scalar unit	Vector/Scalar/GPSIMD
On-chip memory	registers, shared memory, L2	VMEM scratchpad	SBUF + PSUM
Off-chip memory	HBM (global)	HBM	HBM
Interconnect	NVLink / NVSwitch	ICI (torus, pods)	NeuronLink + EFA (scale-out)

A high-performance kernel — whatever the backend — is fundamentally an answer to the same set of questions: How do I tile the computation so the matrix unit stays busy? What do I stage in the on-chip scratchpad, and when? How do I map the elementwise/reduction epilogue onto the vector and scalar units? How do I overlap data movement with computation so the engines never stall? The differences between CUDA, XLA, and NKI are, to a large degree, differences in how much of that reasoning the programmer does explicitly versus how much the compiler does for them.

That is the bridge to the rest of Section 1. The hardware here defines the targets; Section 1.2 will build a mental model of the software stacks that lower models down toward this hardware; and Section 1.3 will turn to the kernel-programming interfaces — CUDA, Triton, and NKI — through which a person, a compiler, or an LLM actually writes the code that these engines run.

1.1 Overview of AI Hardware: GPUs, TPUs, and AWS Trainium