AMD, TPU, and Trainium

AMD’s fourth-generation CDNA architecture powers the Instinct MI350 series. Each MI350-series GPU integrates 8 vertically stacked accelerator complex dies (XCDs) and 2 I/O dies, tied together with AMD’s on-package Infinity Fabric and connected to 8 stacks of 12-Hi HBM3E memory. The eight compute chiplets are built on TSMC’s N3P process, while the two I/O dies, which hold the AMD Infinity Cache and the memory controllers, use TSMC’s N6 process; AMD splits the design this way because the memory and communication logic in the I/O dies does not benefit as much from the newer, more expensive node as compute logic does.

The family ships as two SKUs: the air-cooled MI350X at 1000 W, built to be drop-in compatible with the prior-generation MI325X platform, and the liquid-cooled MI355X at 1400 W for higher power and cooling budgets. The MI355X carries 288 GB of HBM3E, 8.0 TB/s of memory bandwidth, and over 1 TB/s of Infinity Fabric communication bandwidth between GPUs. Across the package its 8 XCDs total 256 compute units and 1,024 matrix cores, with each XCD contributing a 4 MB L2 cache in front of the shared Infinity Cache in the I/O dies below.

CDNA 4’s biggest generational jump is in reduced-precision matrix throughput. Building on the OCP microscaling standard, the CDNA 4 compute units add hardware support for MXFP8, MXFP6, and MXFP4 formats alongside doubled execution resources for the existing 16-bit and 8-bit datatypes. Comparing peak theoretical throughput per GPU, the MI355X reaches 5.0 PFLOPS of dense FP8 matrix compute (10 PFLOPS with structured sparsity) and 10 PFLOPS of dense FP6/FP4 (20 PFLOPS with sparsity), against 2.6 PFLOPS of dense FP8 on the prior-generation MI300X, a 1.9x generational gain on the formats the two share. See AMD ROCm and CDNA in the reading list for the ISA reference and ROCm software stack.

Whether a kernel is actually reaching those numbers is a measurable question, because ROCm documents the full hardware performance counter set for the MI350 series, organized by IP block and exposed through ROCprofiler-SDK and ROCm Compute Profiler. The command processor and shader pipe interpolator blocks count thread groups launched and waves in flight; the compute unit (SQ) counters break instructions and FLOPs out by datatype, down to a dedicated counter for F6 and F4 matrix instructions (SQ_INSTS_VALU_MFMA_F6F4); and the LDS counters report load, store, and atomic traffic in 64-byte units alongside the cycles their FIFOs spent full. The cache blocks carry their own stall accounting, including cycles stalled on data pending from L2 and stalls inside the UTCL1 address-translation unit, so the waiting half of a kernel’s roofline story is as countable on this hardware as the arithmetic half.

Google’s TPU7x, the first chip in the Ironwood generation, is a dual-chiplet design: each chiplet is a self-contained unit with one TensorCore, two SparseCores, and 96 GB of HBM, and the two chiplets are exposed to frameworks like JAX as two separate devices connected by a die-to-die interface. Counted at the full-chip level, that is 2 TensorCores, 4 SparseCores, and 192 GiB of HBM per chip, with roughly 7.38 TB/s of HBM bandwidth. A full pod scales to 9,216 chips connected in a 3D torus topology, with 200 GB/s of bidirectional inter-chip bandwidth per axis between neighboring chips.

Per chip, TPU7x reaches 2,307 TFLOPS of peak bf16 compute and 4,614 TFLOPS of peak FP8 compute. Each TPU7x virtual machine bundles 4 chips together with 224 vCPUs and 960 GB of host RAM, connected to its host over PCIe. See TPU architecture in the reading list for the Pallas TPU programming model that targets this hardware.

The design has a lineage worth knowing, because its core idea has not changed since the first TPU. Google’s 2017 ISCA paper describes the original chip, deployed in its datacenters since 2015 to accelerate neural network inference: at its heart sat a matrix multiply unit built from 65,536 8-bit multiply-accumulate units, delivering a peak of 92 TeraOps/second and fed from a large 28 MiB software-managed on-chip memory. The paper argues that the TPU’s deterministic execution model was a better match for the 99th-percentile response-time requirements of production inference than the time-varying optimizations of contemporary CPUs and GPUs (caches, out-of-order execution, multithreading, prefetching), and measures the chip at roughly 15x to 30x the speed of its contemporary Haswell CPU and K80 GPU on Google’s production workloads, at 30x to 80x their TeraOps per watt.

TPU v4, which Google describes as its fifth domain-specific architecture and third supercomputer for machine learning, shows how far that lineage had scaled by its 2020 deployment. Its supercomputer grew 4x larger than v3’s, to 4,096 chips, joined through optical circuit switches that dynamically reconfigure the interconnect topology (users can pick a twisted 3D torus if desired) at under 5 percent of system cost and under 3 percent of system power. Each TPU v4 also includes SparseCores, dataflow processors that accelerate embedding-reliant models by 5x to 7x while using only 5 percent of die area and power, and the paper reports the chip outperforming TPU v3 by 2.1x while improving performance per watt by 2.7x. The SparseCores and the torus interconnect in TPU7x above are direct descendants of both decisions.

To a kernel author working through Pallas, this hardware looks nothing like a GPU. The JAX documentation describes TPUs as sequential machines with a very wide vector register: the grid of a Pallas TPU kernel is generally processed not in parallel but sequentially, in lexicographic order, and HBM cannot be accessed directly by compute instructions; data has to be prefetched into lower levels of the memory hierarchy by DMA subunits, with matrix multiplies executed by the MXU and the bulk of remaining computation performed on 2D vector registers, typically 8x128 for 32-bit values. The references a kernel body receives point at buffers in VMEM, a vector memory the docs describe as fairly large for its level of the hierarchy at 16 MB or more (the hardware reference table lists 64 MiB per TensorCore for TPU 7x), and in SMEM, a low-latency scalar memory serving the separate scalar unit that handles control flow. The quickstart’s first kernel shows how little of that machinery the author touches directly:

def add_vectors_kernel(x_ref, y_ref, o_ref):
  x, y = x_ref[...], y_ref[...]
  o_ref[...] = x + y

@jax.jit
def add_vectors(x: jax.Array, y: jax.Array) -> jax.Array:
  return pl.pallas_call(
      add_vectors_kernel,
      out_shape=jax.ShapeDtypeStruct.like(x)
  )(x, y)

The quickstart notes that on TPU the references already live in on-chip memory by the time the kernel body runs: values are fetched from HBM before execution, the body moves them from SRAM into registers and back, and results return to HBM only after the kernel completes. The compiler, not the kernel author, schedules those transfers and overlaps them with compute, the same overlap that CUDA kernels in earlier chapters arranged by hand with producer and consumer pipelines.

AWS’s Trainium3 device is built from 8 NeuronCores (v4), 4 HBM stacks totaling 144 GiB of capacity at 4.7 TB/s of bandwidth, 128 DMA engines for moving data within and across devices, 20 CC-Cores dedicated to collective communication, and 4 NeuronLink-v4 links for device-to-device traffic. On-chip, each NeuronCore-v4’s SBUF grew to 32 MiB, up from 28 MiB in the prior NeuronCore-v3, while PSUM stayed at 2 MiB.

The Tensor Engine is where the generational jump shows up most: it runs at 2.4 GHz and delivers 315 TFLOPS of MXFP8/MXFP4 compute, 79 TFLOPS of BF16/FP16/TF32, and 20 TFLOPS of FP32. To hit the MXFP8/ MXFP4 rate it quadruples the matmul contraction dimension from 128 to 512 elements, presenting what the architecture guide calls a 512x128 systolic array to the programmer, even though the underlying grid of processing elements is still 128x128. See Trainium and NKI in the reading list for the full architecture guide and the tile-level programming model built on top of it.

AWS’s prescription for programming this hardware is spelled out in the NKI performance guide, and it reads like this book’s roofline chapters translated into Neuron vocabulary. Optimization work should end with a kernel that is either compute-bound, meaning at least one compute engine is active close to 100 percent of the execution time (the guide treats 90 percent or more as good in practice), or memory-bound, with achieved memory bandwidth utilization close to 100 percent (60 percent or more is considered good). Getting there is framed as a fight against data movement: keep inputs resident in SBUF instead of reloading them over DMA, fuse consecutive operators through explicit loop fusion so intermediates never spill to device memory (the profiler’s spill_save_bytes and spill_reload_bytes metrics expose how much traffic is spill), pipeline tiles across the Tensor, Scalar, Vector, and GpSimd engines so no engine idles waiting on another, and size instruction tiles large enough to amortize instruction overhead but small enough not to wreck pipelining or SBUF pressure. The language those optimizations are written in is small; the getting-started guide names the three phases every NKI kernel has, load from device memory into SBUF, compute, store back, and its first example is the whole model in one function:

from neuronxcc import nki
import neuronxcc.nki.language as nl


@nki.jit
def nki_tensor_add_kernel(a_input, b_input):

    """NKI kernel to compute element-wise addition of two input tensors
    """

    # Check all input/output tensor shapes are the same for element-wise operation
    assert a_input.shape == b_input.shape

    # Check size of the first dimension does not exceed on-chip memory tile size limit,
    # so that we don't need to tile the input to keep this example simple
    assert a_input.shape[0] <= nl.tile_size.pmax

    # Load the inputs from device memory to on-chip memory
    a_tile = nl.load(a_input)
    b_tile = nl.load(b_input)

    # Specify the computation (in our case: a + b)
    c_tile = nl.add(a_tile, b_tile)

    # Create a HBM tensor as the kernel output
    c_output = nl.ndarray(a_input.shape, dtype=a_input.dtype, buffer=nl.shared_hbm)

    # Store the result to c_output from on-chip memory to device memory
    nl.store(c_output, value=c_tile)

    # Return kernel output as function output
    return c_output