Systolic designs

TPUs and Trainium: large weight-stationary arrays fed from HBM by a compiler, and why decode needs a batch as big as the array to fill it.

Loading the animation…

Concept

The second family keeps the GPU's memory system (weights in HBM) but builds its matrix engine differently: a few large systolic arrays in which each weight sits in one processing element while activations flow past it, so a weight read once from memory is used by every row that streams through. Two shipping lines say so in their own documentation. Google's TPUs: "an MXU is composed of either 256 x 256 (TPU v6e and TPU7x) or 128 x 128 (TPU versions prior to v6e) multiply-accumulators in a systolic array". AWS Trainium: the NeuronCore's "Tensor engines are based on a power-optimized systolic array", organised "as a grid of 128x128 processing elements". How such an array works cycle by cycle is the whole of Systolic Arrays Explained; this chapter is about what the design means for the running workload.

The array is a fixed pipeline, so the compiler, not the hardware, plans every transfer: Google's TPUs take their programs through XLA, and VMEM, the on-chip scratchpad in front of the MXUs, is sized and filled by the compiler and by custom kernels rather than by a cache.

The animation runs one 1,024 × 1,024 matmul. Each tile of weights holds the array for $\max(m, R)$ cycles: the $m$ input rows of the batch take $m$ cycles to stream through, and the next tile's $R$ rows of weights take $R$ cycles to shift in behind it. With a batch of 1 the array spends 1 cycle in 256 multiplying, so one decoder layer of the 70B model uses 0.39% of a 256 × 256 array. At a batch of 16 it uses 6.2%; at a batch of 256, 100%. A 128 × 128 array fills at half the batch: 100% at 128. The big array's lesson for decode is the GPU's, sharpened: it wants a batch at least as tall as itself.

	Google TPU v6e (Trillium)	Google TPU7x (Ironwood)	AWS Trainium2
BF16, dense	918 TFLOP/s[tpu-v6e]	2.31 PFLOP/s[tpu-7x]	667 TFLOP/s[trn2-arch]
FP8, dense	918 TFLOP/s[tpu-7x]	4.61 PFLOP/s[tpu-7x]	1.3 PFLOP/s[trn2-arch]
Off-chip memory	32 GB[tpu-v6e]	192 GiB[tpu-7x]	96 GiB[trn2-arch]
Off-chip bandwidth	1.638 TB/s[tpu-v6e]	7.38 TB/s[tpu-7x]	2.9 TB/s[trn2-arch]
On-chip SRAM	not disclosed[tpu-v6e]	not disclosed[tpu-7x]	224 MB (derived)[nc3-arch]
Chip-to-chip (both directions)	800 GB/s[tpu-v6e]	1.2 TB/s[tpu-7x]	1.28 TB/s[trn2-arch]
Cores / SMs / TensorCores	2[tpu-v6e]	2[tpu-7x]	8[trn2-arch]
Systolic array rows	256[tpu-arch]	256[tpu-arch]	128[nki-trn2]
Power	not disclosed[tpu-v6e]	not disclosed[tpu-7x]	not disclosed[trn2-arch]

The running workload on systolic chips

Because the weights live in HBM, decode on these chips is bounded the way it is on a GPU: by bytes over bandwidth. TPU7x reads its 192 GiB of HBM at 7.38 TB/s, so Llama 3.1 70B in FP8 fits on one chip and a batch-1 step is bounded at 9.44 ms. Trainium2 holds it in 96 GiB and reads at 2.9 TB/s: 24 ms. TPU v6e has 32 GB per chip, so the 16-bit model needs 5 chips. The ridge points are GPU-like too: 625 flops per byte for TPU7x in FP8, 448 for Trainium2. Its successor, Trainium3, is not in this site's model: AWS gives it "144 GB HBM3e per chip, 4.9 TB/s bandwidth - 1.7x higher than Trainium2" and "up to 2x MXFP8 compute throughput compared to Trainium2", in UltraServers of up to 144 chips (AWS Trainium, read 2026-10-16; Racks, chapter 10 uses those figures).

Part	Model, format	Chips	Step	Bound	Tokens/s per user
Google TPU v6e (Trillium)	llama3-70b, BF16	5	17 ms	memory	58.8
Google TPU7x (Ironwood)	llama3-70b, FP8	1	9.44 ms	memory	106
Google TPU7x (Ironwood)	llama3-70b, FP8	1	15.2 ms	memory	65.6
AWS Trainium2	llama3-70b, FP8	1	24 ms	memory	41.6
AWS Trainium2	llama3-70b, FP8	2	19.4 ms	memory	51.6

What differs from a GPU is the machine around the array: one large compiled program instead of many kernels, a scratchpad instead of a cache, and, at scale, the TPU pods' torus of chip-to-chip links (AI Racks Explained, chapter 10). Google and AWS do not state their chips' power in the sources used, so this chapter gives no energy per token for them.

Maths

A weight-stationary array of $R \times R$ processing elements running an $m \times k$ by $k \times n$ matmul works through $\lceil k/R \rceil \lceil n/R \rceil$ tiles. With the next tile's weights shifted in behind the current one, a tile holds the array for $\max(m, R)$ cycles, and the first fill and last drain add $2R - 2$ :

C = \left\lceil \frac{k}{R} \right\rceil \left\lceil \frac{n}{R} \right\rceil \max(m, R) + 2R - 2, \qquad U = \frac{m k n}{R^2 C} \;\approx\; \frac{\min(m, R)}{R} \text{ for large } k, n.

So a batch below $R$ uses about $m/R$ of the array, whatever the matrix.

Code

One matmul on the array, cut from src/lib/accel/model.ts:

const tiles = Math.ceil(k / rows) * Math.ceil(n / cols);
const perTile = Math.max(m, rows);
const cycles = tiles * perTile + rows + cols - 2;
const macs = m * k * n;