The TPU has less public hardware documentation than a GPU. There is no vendor ISA manual, no die-level memory-model spec; you get peak FLOPs, HBM capacity, and a block diagram. But if you write JAX or Pallas kernels for TPUs, you are not actually flying blind, because Google ships, openly, Apache-2.0 licensed, the compilers that target the chip, and a compiler has to name the parts of the machine it generates code for. The memory spaces, the register geometry, the instruction set, the synchronization primitives: all of it is spelled out in code you can read.
The exact numbers, lane counts, memory sizes, latencies, stay in the closed libtpu.so. The dialect is parameterized rather than hardcoded: the layout-inference entry point takes a hardware_generation and a target_shape as arguments. That split is itself informative. The structure, spaces, tilings, op set, is stable enough to design around. The constants are the compiler's job, supplied per generation.
Eight memory spaces, and why your Ref's space is the first thing to get right
Start with memory, because on a TPU it is the thing you manage explicitly and the thing that most determines performance. Mosaic's MemorySpace enum names eight:
- kVmem (0): vector memory, the on-chip scratchpad for the vector unit.
- kSmem (1): scalar memory, on-chip, for the scalar core.
- kHbm (2): high-bandwidth memory, off-chip DRAM.
- kCmem (3): a further on-chip space.
- kSemaphoreMem (4): where sync primitives live.
- kVmemShared (5): VMEM shared across subcores.
- kHost (6): host DRAM.
- kAny: unconstrained, the compiler's choice.
Compared with CUDA's three that matter (global, shared, local), the TPU exposes a richer, explicitly typed hierarchy. How this helps you write better TPU code: in a Pallas TPU kernel every buffer is a Ref tagged with one of these spaces, and that tag is the single most consequential choice you make. The core performance idiom of every good TPU kernel follows directly: stage the tile you are about to compute on from HBM into VMEM, work on it there, and write it back, never compute directly against HBM. Getting the space wrong is not a micro-optimization; it is the difference between a kernel that runs at bandwidth and one that stalls or fails to lower.
Two kinds of core, and they do not have the same register shape
The CoreType enum names three things: TensorCore (the dense matmul/vector engine, what people mean by "TPU"), and SparseCore split into scalar and vector subcores. SparseCore is the hardware embedding engine for recommendation and ranking models, the gather/scatter-over-huge-tables workloads that are not dense matmul. The two engines do not share a register geometry.
The TPU's compiler is its manual. Google does not document the chip, but it documents the compiler, and the compiler has to describe the chip exactly to generate code for it. Read the dialect and you read the machine.
VMEM is the fast on-chip scratchpad your vector unit reads at speed. Never compute directly against HBM.