Anatomy of a TensorCore

A TPU chip is one or two TensorCores plus HBM. A TensorCore is the TPU's unit of compute, the counterpart of a GPU's SM, not of a GPU's Tensor Core. The name collision causes confusion, so this course always writes TensorCore as one word for the TPU unit and "Tensor Core" as two for NVIDIA's matrix unit [SB12].

An H100 has 132 SMs. An Ironwood chip has two TensorCores [SB12] [D-7x]. Almost every difference in this chapter follows from that ratio.

The parts of a TensorCore

The scalar unit fetches very long instruction words (VLIW) and issues them. On v2 and v3 a bundle was 322 bits with slots for 2 scalar ops, 4 vector ops (2 of them loads and stores), 2 matrix pushes, 1 matrix pop and 1 miscellaneous op [P-Norrie]. It also runs loops and issues DMAs, one DMA request per cycle per core on current chips [SB2]. It has its own small memory, SMEM, 1 MiB from v4 onward [J-hw].

The VPU is an 8×128 grid of 32-bit lanes: 128 lanes across, 8 sublanes deep. On Ironwood each (lane, sublane) position has 4 ALUs, so the VPU has 4,096 ALUs per core. It runs the elementwise work, reductions within a lane, and the softmax and normalization arithmetic [SB2].

The vector registers hold 64 registers of 8×128 32-bit values each, about 256 KiB per core on Ironwood (32 registers on v4). Each cycle can do up to 3 register loads and 2 stores [SB2].

The MXUs are the matrix units: weight-stationary systolic arrays, 128×128 through v5 and 256×256 on v6e and Ironwood. They multiply in bf16 (fp8 on Ironwood) and accumulate in fp32 [D-arch]. Ironwood has two per TensorCore, four per chip.

The XLU handles transposes, permutations and cross-lane reductions [P-Norrie]. Moving data across lanes is not free on a TPU; it goes through the XLU or through memory.

VMEM is the on-chip memory every operand passes through on its way from HBM to the compute units. It is covered in its own section below.

One thread of control

A GPU SM keeps up to 64 warps resident and switches between them every cycle. When one warp waits on memory, another runs. Latency hiding is a hardware property, and the programmer's job is to supply enough parallelism to keep occupancy up [N-tune, §1.4.1.1].

A TensorCore has one instruction stream. Nothing switches to other work when a load stalls. Latency is hidden only if the instruction stream was arranged in advance so that the DMA for the next tile is already in flight while the current tile is being multiplied. XLA does that arrangement for the whole program. It knows every shape at compile time, so it can schedule every HBM-to-VMEM copy ahead of its use [SB2]. In a Pallas kernel, you do it by choosing block shapes, and the framework double-buffers for you (Chapter 10).

The MXU is a weight-stationary systolic array

A systolic array is a grid of multiply-accumulate cells that pass data only to their neighbours. In the weight-stationary arrangement the TPU uses, one tile of the weight matrix $W[K, N]$ is loaded first, one value per cell [P-Norrie]. Then rows of the activation matrix enter from the left edge, one row per cycle, each skewed one cycle behind the previous array row. Each cell multiplies the activation passing through it by its stored weight, adds the partial sum arriving from above, and passes the sum down. Finished outputs leave from the bottom edge [SB2].

No cell reads memory during the multiply. Data moves one hop per cycle, and each value loaded into the array is used many times as it crosses. That reuse is what makes a large array affordable in area and power.

Activation $x[m, k]$ reaches cell $(k, n)$ at cycle $m + k + n$, so a tile with $M$ streamed rows finishes after $M + K + N - 2$ cycles. That gives the tile's utilization:

Useful multiply-accumulates over the cells available during the time the tile occupies the array. S is the array size; K and N are the parts of it the weight tile fills.

Two consequences, computed for a 256×256 array in scripts/ch02_core.py derived:

• Rows per weight tile must be large. With $K = N = 256$, streaming 8 rows uses 1.5% of the array's capacity over the tile, 256 rows uses 33%, and 2,048 rows uses 80%. The number of rows per weight tile is the token batch. This is the HBM roofline from Chapter 1 seen from inside the chip: a decode step loads each weight tile and uses it for a handful of tokens.
• K and N must fill the array. A contraction depth of 64 on a 256-deep array caps utilization at 25% before any other loss. 128×128 tiles on a 256×256 array use a quarter of it.

Padding

XLA pads every matrix dimension that feeds the MXU up to a multiple of the array size. The Scaling Book's advice is to make weight dimensions multiples of 256 on v6e and Ironwood, and 128 before [SB2]. A dimension of 100 keeps 78% of its work useful on a 128-wide array and 39% on a 256-wide one. Most model dimensions are already multiples of 256. The exception that matters is head dimension: attention contracts over it, and many models use 64 or 128. With head dimension 64, the $QK^\top$ product fills a quarter of a 256-deep array unless the kernel packs several heads into one contraction. Chapter 8 covers how the attention kernels in vLLM TPU deal with it.

VMEM: a scratchpad, not a cache

Every byte an MXU or VPU touches comes from VMEM, through the vector registers. A DMA engine copies data from HBM into VMEM before it is used and copies results back [SB2]. VMEM has no tags, no eviction policy and no hit rate. What it holds at any moment is decided at compile time, by XLA's memory-space assignment pass or by a Pallas kernel's block specification.

Per-TensorCore figures from the JAX hardware table [J-hw], chip totals multiplied out. v4's CMEM is from the v4 paper [P-v4]. The 8i total matches Google's "384 MB" figure [G-8dd].

VMEM bandwidth is not on any spec sheet. The Scaling Book puts it at about 7× HBM on Ironwood [SB2] and roughly 40 TB/s in its GPU chapter [SB12]. One of its problems uses 22× HBM for v5e [SB2]. Treat VMEM bandwidth as "several times HBM", with the exact ratio unverified.

That ratio matters when weights are already resident. If a layer's weights fit in VMEM and were prefetched while the previous layer ran, the matmul's ridge point is set by VMEM bandwidth, not HBM. On Ironwood at 7× HBM that ridge is about 45 FLOP/B instead of 313 derived. The Scaling Book gives the example of prefetching the feed-forward weights into VMEM while attention runs [SB2].

Layout: everything is tiled in 8×128

The vector registers are 8 sublanes × 128 lanes, so XLA stores arrays in HBM and VMEM in tiles of that shape. A 32-bit array uses (8,128) tiles over its last two dimensions. bf16 packs two rows into each 32-bit sublane, written (8,128)(2,1). int8 and fp8 pack four, written (8,128)(4,1) [X-tiled]. A profile writes the layout next to every array shape, as in bf16[32,4096]{1,0:T(8,128)(2,1)}. Chapter 10 reads these.

Three practical consequences:

• The last two dimensions should be multiples of 8 and 128. Anything else is padded in memory, not only in the MXU. A Pallas block shape must satisfy the same rule [PL-tpu].
• Small second-minor dimensions waste memory. An array of shape [3, 5] in f32 occupies a padded (8,128) tile.
• Changing layout costs a copy. A transpose or a reshape that crosses tile boundaries makes XLA insert a relayout, which appears in profiles as a copy or transpose fusion.

The KV-cache tensor in tpu-inference has an explicit "packing" dimension for the same reason: K and V heads are packed so that each 32-bit sublane holds two bf16 or four int8 values (Chapter 8).

How many devices per chip

Sources: [SB2], [D-7x], [J-hw].

The Ironwood change matters for code. A 4-chip Ironwood VM presents 8 JAX devices [D-7x]. Because the die-to-die link between a chip's two chiplets is 6× faster than one ICI link, it pays to put the two chiplets of a chip on the same mesh axis, the one that carries the most traffic. Pallas's megacore annotation, dimension_semantics=("parallel", ...), splits a grid across the two cores only on v4 and v5p, where they form one device [PL-pipe].

SparseCores and the Collectives Acceleration Engine

Not all work is dense matrix multiplication. Embedding lookups in recommendation models are random gathers and scatters over huge tables, and the dense units handle them poorly. TPU v4 added SparseCores for this: 16 tiles per SparseCore, each with its own HBM channel, fetch unit and 8-wide SIMD unit, about 5% of the die. Embedding-heavy models ran 5–7× faster than with embeddings held in CPU memory [P-v4, §3]. v5p has 4 SparseCores per chip, v6e 2, Ironwood 4 [D-arch].

They are programmable from Pallas. The JAX documentation's example runs a gather in 4.05 ms on SparseCore against 18.1 ms with jnp.take on the TensorCore [PL-sc]. tpu-inference uses SparseCore kernels for the ragged gather and scatter in MoE routing (Chapter 8).

TPU 8i replaces SparseCores with a Collectives Acceleration Engine (CAE), which Google says cuts on-chip collective latency by up to 5× [G-8dd]. The JAX hardware table still lists one SparseCore per 8i TensorCore [J-hw]. The two sources disagree, and which is current is unverified. 8t keeps SparseCores and adds an "LLM Decoder Engine" whose function Google has not described.

Mapping a GPU onto a TPU

Counts from the Scaling Book's GPU chapter [SB12]. The pattern: a GPU spends its die on many small independent units and a large register file; a TPU spends it on a few large matrix units and a large scratchpad.