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, 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.

That 1.9x is one row of AMD’s Table 1, and the honest way to read a vendor table is to read the rows that go the other way too. The same table records Matrix FP64 falling from 256 to 128 FLOPS per clock per CU, 163.4 TF down to 78.6 TF, which AMD scores at roughly 0.5x: half the double-precision matrix throughput of the part it replaces. Vector FP64, FP32, and FP16 all land at about 0.96x, slightly slower than MI300X rather than faster. And TF32, which the prior generation supported in hardware, has been removed from hardware entirely; AMD writes that after extensive discussion with customers and the ecosystem it “has been moved out of hardware and is supported through software emulation utilizing the BF16 datatype.” The generational picture is therefore conditional on datatype rather than uniform: matrix FP16, BF16, FP8, and INT8 all gain 1.9x, while a kernel whose matmuls need FP64 loses half its peak and one built on TF32 now runs on an emulation over BF16. That is worth knowing before reading 1.9x as a property of the part.

The other change a kernel author will feel is in the memory the bank-conflict material of Foundations was about. AMD’s shared memory is the Local Data Share, and through CDNA 3 it was a directly addressed structure of 32 banks, each holding 512 entries of 32 bits, for 64 KB in total, with conflict detection and scheduling logic, a crossbar and swizzle unit, and atomic execution units in front of it. CDNA 4 raises the LDS to 160 KB by adding banks, and doubles its read bandwidth to 256 bytes per clock. It also lets the LDS load data directly from the L1 vector data cache, which AMD says reduces vector register usage and latency. The L1 those loads come from is itself mostly unchanged from the previous generation at 32 KB of capacity with 128-byte lines and 64-way associativity. AMD ties both LDS changes to matrix multiplication specifically, calling the extra capacity and bandwidth crucial to improving utilization of the vector and matrix execution resources because of the extensive data reuse in matrix multiply routines, and it pairs them with a doubling of transcendental rates aimed at the softmax in attention.

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.

A CDNA 4 kernel that wants those matrix numbers goes through the MFMA instruction family, and the instruction set architecture document describes a machine that differs from a tensor core in ways a port will notice. The matrix core is an extension to the CDNA shader instruction set with its own register file, the accumulation VGPRs, separate from the architectural VGPRs the rest of the SIMD uses. Which file an operand lives in is part of the instruction: the low and high bits of an MFMA’s ACC field select the file for the A and B matrices respectively, and a separate ACC_CD bit selects it for C and D, with V_ACCVGPR_READ and V_ACCVGPR_WRITE moving values between the two files. Because the accumulation file is separate, registers holding a running accumulation are not drawn from the same pool as the ones the addressing and loop code around the matmul is using.

Underneath the whole family is one operation. AMD states that “the core operation implemented inside the matrix core is the 4 x 1 times 1 x 4 outer matrix product, yielding 16 output values,” performed on dense inputs or on 2,4 sparse ones, and that the matrix core composes these outer products in parallel and in series to implement every MFMA instruction. The instruction names read off the composition: V_MFMA_[output type]_[M]X[N]X[K][_[B]B]_[input type], where M, N, and K are the multiplication dimensions of one block and B, taken as 1 when absent, is how many blocks the instruction multiplies. V_MFMA_F32_32x32x1_2B_F32 therefore performs two independent 32 x 32 block products in one instruction rather than one larger one, which is the distinction that decides how a tile maps onto it.

Register placement is a hard constraint rather than a preference. The registers holding an MFMA’s input or output “must be contiguous, and the first register must be aligned to the number of registers required as the input or output.” The document’s own example is the one to remember: an instruction needing four input registers for matrix A may use registers 4 through 7 but not registers 5 through 8. Sparsity is a separate instruction family rather than a modifier. The V_SMFMAC instructions multiply a 4:2 structurally sparse A against dense B, C, and D, meaning two of every four elements along K are zero; those zeros are not stored, and a separate VGPR carries pairs of 2-bit indices naming which two of each group of four survive, for 2:1 compression. Only A may be sparse, C and D are the same registers, and the C operand slot is repurposed to carry the index offset.

The constraint with no counterpart anywhere in the CUDA chapters of this book is the last one. Because matrix instructions “do not produce their output in a single cycle, and since their partially-written results may be observable, a certain amount of independent instructions must sometimes be present between the issuance of a matrix core instruction and accesses to its results or modification of the registers that hold its inputs.” The ISA prints a table of how many, and states that it is the user who inserts them, as NOPs or as independent VALU instructions. A VALU write that is not a dot-product op, followed by an MFMA reading that register, needs 2. An MFMA result read as the C operand of a following matrix instruction of exactly the same pass count and register offset can need none, because that path forwards. The same result read as the A or B operand of a following MFMA needs 5, 8, 12, or 20 depending on the instruction, since there is no forwarding path and the first instruction must commit its result to a VGPR before the second may read it. Getting this wrong does not cost performance, it produces wrong answers, and it turns the question of what independent work sits between two dependent matrix instructions from a scheduling preference into a correctness requirement.

That hardware is reached through HIP, described by AMD as “a C++ runtime API and kernel language for AMD GPUs” and as “a thin API with little or no performance impact over coding directly in AMD ROCm.” The claim that matters to a CUDA programmer is about the kernel language rather than the runtime. HIP marks a device function with __global__, indexes with threadIdx, blockIdx, and blockDim, and launches with the triple chevron, so the introductory kernel in its programming model documentation is the vecAdd of Foundations under a different name:

__global__ void AddKernel(float* a, const float* b)
{
  int global_idx = threadIdx.x + blockIdx.x * blockDim.x;

  a[global_idx] += b[global_idx];
}

AddKernel<<<number_of_blocks, threads_per_block>>>(a, b);

Nothing in that listing is AMD-specific, and that is the design. Inside the chevrons the same four things are configurable as in CUDA: the number of blocks, the threads per block, the shared memory the host allocates, and the stream to enqueue on. HIP also provides hipLaunchKernelGGL(), which “by default, is a macro expanding to the triple-chevron syntax,” for cases where the language extension is unwanted or the kernel is templated or overloaded. The host side is unified rather than merely mirrored: CUDA’s split between cu-prefixed driver calls and cuda-prefixed runtime calls mostly collapses into a single HIP API, with the driver-only module and context functionality kept under hipModule and hipCtx prefixes.

Because the correspondence is that close, most of the port is mechanical, and ROCm ships tools for the mechanical part. HIPIFY “is a collection of tools that automatically translate CUDA code to HIP code,” in two flavors with a real tradeoff. hipify-clang “parses code, translates it into an Abstract Syntax Tree, and generates the HIP source,” which means it needs source that actually compiles and a CUDA install with all necessary headers. hipify-perl “uses pattern matching,” needs no CUDA installation, tolerates source that is not syntactically correct, and “is therefore easier to set up and use, but is not as powerful as hipify-clang.” Before converting anything, --examine runs a scan “without changing the source files” that reports which files contain CUDA and how much of it converts automatically. On one of NVIDIA’s own samples the documented scan finds 15 convertible references across 512 lines of code and names every one, cudaMalloc to hipMalloc, cuda_runtime.h to hip/hip_runtime.h. Passing --inplace then rewrites the files, keeping the originals as .prehip backups.

That number is the argument for the tool and the warning about it at once: 15 references in 512 lines is most of the port and none of the performance work. AMD states the limit plainly, that HIP “is not intended to be a drop-in replacement for NVIDIA CUDA, and developers should expect to do some manual coding and performance tuning work to port existing projects to AMD GPUs.” The porting guide’s advice around the residue is unglamorous: start from a working CUDA codebase, port incrementally, and test each section as it converts.

The residue is where the assumptions of Foundations stop holding. Code “should not assume a warp size of 32 or 64, as AMD GPU architectures have different warp sizes”; the warpSize built-in is the device-side answer and the device properties the host-side one. The concrete failure is lane masks. GCN and CDNA have a warp size of 64, so a mask built as 1 << (threadIdx.x % warpSize) shifts a 32-bit constant by up to 63, and shifting by such values “would clear the 32-bit register to which the shift operation is applied,” with the documented fixes being a 1ull literal or an architecture-conditional lane_mask_t typedef. __launch_bounds__ survives with a changed second parameter: CUDA’s MIN_BLOCKS_PER_MULTIPROCESSOR becomes HIP’s MIN_WARPS_PER_EXECUTION_UNIT, obtained by multiplying by the maximum threads per block and dividing by warpSize * 2 in the compute-unit mode typical of CDNA, or by warpSize * 4 in RDNA’s workgroup-processor mode. And --maxregcount has no equivalent at all: amdclang++ does not support it, and the documentation steers toward __launch_bounds__ because its parameters are “more intuitive than micro-architecture details like registers.”

What does carry over is the memory reasoning, stated in AMD’s documentation in almost the words Foundations used. Global memory transactions “typically occur in chunks of 32, 64, or 128 bytes, which must be naturally aligned,” and coalescing means arranging accesses “so that multiple threads in a warp can combine their memory requests into the fewest possible transactions.” The prescriptions are the familiar four: use naturally aligned data types, arrange consecutive threads to touch consecutive locations, avoid strided access such as array[i * stride], and pad structures where alignment requires it. Shared memory is exempted from coalescing and subject to the other hazard instead, since accesses “do not require coalescing, but they can suffer from bank conflicts,” which serialize when threads in a warp hit different addresses in the same bank. The units are AMD’s, the failure modes are not.

Once ported, the reasoning is the reasoning of Foundations, and ROCm’s own documentation says so in the same vocabulary: a compute roof at the peak arithmetic rate, a memory roof whose slope is the memory bandwidth, and a ridge point whose horizontal coordinate “gives the minimum arithmetic intensity required to escape the memory bottleneck.” Its optimization workflow makes that the second step, after profiling: “determine if kernels are compute-bound or memory-bound. Check arithmetic intensity, memory bandwidth achieved vs peak, and compute throughput,” with rocprofv3 gathering achieved FLOPs, memory transactions, and operational intensity. The generational observation it adds is the one every specification in this chapter has been illustrating: “improvements in compute throughput have far outpaced memory growth,” pushing ridge points steadily to the right. CDNA 4’s doubled execution resources for 16-bit and 8-bit datatypes are one more instance, raising the compute roof so that a kernel which does not also raise its arithmetic intensity simply lands further below it. The model is deliberately blind to what that costs, since it “ignores latency entirely, focusing only on sustained throughput limits.”

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.

That hardware reference table is worth reading across rather than down, because it puts ten TPU variants against the same set of columns and so states the argument this chapter keeps making. Per TensorCore, HBM bandwidth goes from 358 GB/s on v2 to 3,700 GB/s on 7x, a little over 10x, while BF16 peak goes from 23 TFLOPs/s to 1,155, about 50x. On-chip memory does not move monotonically with either: VMEM sits at 16 MiB from v2 through v4, jumps to 128 MiB on v5e, drops to 64 MiB on v5p, returns to 128 MiB on v6e, and is back at 64 MiB on 7x, so a kernel tiled to fill VMEM on one generation is not portable to the next by arithmetic. SMEM goes 16 KiB to 1,024 KiB at v4 and then stops moving, and CMEM appears on v4 alone. The table also now carries two rows past 7x, an 8I and an 8T of a TPU 8 generation, with 8T listed at 128 MiB of VMEM, 231 GB of HBM at 6,400 GB/s, and 5,977 TFLOPs/s of FP8. Those rows are what the JAX hardware reference records; the Google Cloud documentation this section draws its TPU7x figures from does not describe that generation, so nothing else here is said about it.

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 is rated at 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.

A NeuronCore is not one processor. AWS documents its internals in the architecture guide for the second-generation NeuronCore, which it marks as relevant for Trainium2 and Trainium3, and the shape it describes is the thing a NKI kernel is written against. A NeuronCore-v2 holds four heterogeneous compute engines, Tensor, Vector, Scalar, and GpSimd, each meant to accelerate a different class of operator. Each has its own sequencer that fetches, decodes, and issues its instructions, so the four engines execute four independent instruction streams asynchronously in parallel. Nothing about that is implicitly ordered: synchronization to satisfy a data dependency between engines is carried out by atomic semaphores in hardware, which the Neuron Compiler inserts during compilation from the dependencies it finds in the kernel rather than the programmer writing them. Alongside the four sits a Sync Engine, a fifth sequencer that runs the same control instructions and is most commonly used to trigger DMA transfers without disturbing the instruction scheduling and ordering of the compute engines.

The division of labor is what makes the pipelining advice below mean something. The Tensor Engine accelerates matrix multiplications and the operators that reduce to them, such as 2D convolutions, and is built around a systolic array of 128 rows and 128 columns of processing elements that streams its inputs from SBUF and writes its outputs to PSUM. The Vector Engine handles operations where an output element depends on several input elements, reductions and elementwise operations between two tensors, across 128 parallel lanes each streaming from one memory partition. The Scalar Engine handles the case where each output element depends on exactly one input element, and carries hardware evaluation of non-linear functions such as Gelu and Sqrt, also across 128 lanes. The GpSimd Engine is the general-purpose remainder, for operators that do not lower efficiently onto any of the other three, AWS’s example being applying a triangular mask to a tensor; it is eight fully programmable processors that can execute arbitrary C and C++, each one a 512-bit vector machine. The engines do not run at the same speed either, and AWS publishes the data-path widths and frequencies on the grounds that they are often worth taking into account when optimizing a multi-engine operator: on Trainium/Inferentia2 the Tensor Engine takes 2x128 elements per cycle and emits 1x128 at 2.8 GHz, the Vector Engine moves 128 elements in and out at 1.12 GHz, and the Scalar and GpSimd engines run at 1.4 GHz.

The two on-chip SRAMs those engines share are software-managed and two-dimensional, which is the structural fact behind every tile shape in a NKI kernel. On a NeuronCore-v2, SBUF is 24 MiB of main data storage and PSUM is 2 MiB serving as the Tensor Engine’s dedicated accumulation buffer, and both are organized as 128 partitions: one SBUF partition holds 192 KiB and one PSUM partition holds 16 KiB. That partition count, not a software convention, is why a NKI tile’s first dimension is capped at 128 and why the getting-started example below asserts against nl.tile_size.pmax before doing anything else. The Trainium3 figures earlier in this chapter are the same structure with a larger SBUF. Two pairs of engines also contend for memory ports rather than running truly independently: the Vector and GpSimd engines cannot access SBUF in parallel, and the Vector and Scalar engines cannot access PSUM in parallel, so those instructions get serialized. The Neuron Compiler enforces both, but a pipeline drawn on the assumption that all four engines overlap freely will not match the profile.

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 to minimize intermediates spilling 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