Diving into TPU architectures

Published

This post intends to dive into a TPUv6e chip from a first principles perspective.

I’ve been training models, and optimizing AlphaFold3 for TPU v6e chips, so these notes are things I document as I go so that I understand the lower level more. I hope it will be an accessible guide once it’s fully written. This article assumes you have a bit of a background in computer architecture.11 I’m glad that at CMU we had guest lectures about TPUs from Google which dived in very well on how to code for TPUs. However at the time of learning, I was pretty naive since I wasn’t actively using TPUs. But now, I have more motivation to revisit and learn this more carefully.

What is inside a v6e chip

Alright, if we think from first principles, first off, for any kind of processor, we need to have something that controls the flow, meaning that it orders everybody to stay in sync.

Alright, diving bit deeper, we have the high bandwidth memory (HBM) that we all deeply value22 similar to how we say H100s have 80GB of VRAM.. While this is “high bandwidth” memory, the actual parts that do the computation operate quite fast33 each MXU giving out 459 TFLOP/s in bfloat16. That’s why VMEM exists.

Next, the question is, why do we have separate kind of arithmetic units? The simple answer is that these specialize on specific kind of operations :

ONE TPU v6e CHIPhost CPU (PCIe Gen5 x8)ICIICIICIICITENSORCORE · ONE PER CHIP918 TFLOP/s BF16Scalar unitcontrol flow, issues instructionsSMEM1 MiBVector registers8 × 128 words each, shared by all unitsVMEM · 128 MiBon-chip scratchpad, no cache:the compiler or your kernel fills itvector load / storeVPUelementwise mathMXU 0256 × 256, bf16 → f32MXU 1256 × 256, bf16 → f32XLUXLUcross-lane unitsSparseCore × 2embedding lookupsHBM32 GiB1.64 TB/s(31.25 GiB usable)DMAfour ICI links to neighbouring chips: a 2D torus of up to 256 chips, 800 GB/s per chip
Figure 1: One TPU v6e chip. Redrawn after lecture 12, slide 27, with the v6e parts and numbers from the Cloud TPU v6e documentation and the Pallas hardware reference (VMEM). Lecture slides: Srinath Mandalapu (Google), LLM Systems course, Spring 2026.

What does a Pallas kernel control?

Pallas is an extension to JAX that lets you write kernels. Some of Anthropic’s bioml optimizations use Pallas: the FlashPairformer kernels for their JAX models are written in it44 These target NVIDIA GPUs only. Their PyTorch models use CUDA and Triton versions of the same kernels..

If Pallas can be used for both GPUs and TPUs, how does Pallas know such low level detail about both of them? It doesn’t: it hands your kernel to a compiler for that chip. On TPUs that is Mosaic, which turns it into TPU instructions; on GPUs it is Triton (what Anthropic’s kernels use) or Mosaic GPU.

So in summary, your Pallas kernel controls these parts of the TPU, and the compiler takes care of the more fine-grained components.

WHAT A PALLAS KERNEL CONTROLS ON ONE v6e CHIPhost CPUICIICITENSORCOREScalar unitruns your grid loopSMEMyour scalarsVector registersyour last two dimensions map hereVMEMyour x_ref / o_ref blocksand scratch buffersx = x_ref[...]VPUelementwise mathMXU 0jnp.dot → hereMXU 1jnp.dot → hereXLUXLUSparseCore × 2separate Pallas APIHBMyour kernel’sinputs and outputsBlockSpec copies (DMA)
Figure 2: What a Pallas kernel controls, on the chip from Figure 1.

Your Pallas kernel decides

  • The grid: the loop your kernel body runs in, one step at a time on a v6e
  • BlockSpecs: tile shapes and which tiles each step uses; the copies are then overlapped with compute for you
  • VMEM blocks and scratch buffers; SMEM buffers
  • Dimension order: the last two dimensions map onto the 8 × 128 registers
  • Operations and precision (preferred_element_type)
  • Optionally your own DMA copies: pltpu.async_copy

The compiler decides

  • Mosaic picks the instructions, usually: lax.dot → MXU, elementwise math → VPU, cross-lane transposes → XLU, exp → EUP
  • Register allocation and spills to VMEM
  • XLA’s LLO stage then packs everything into static VLIW bundles

This split isn’t strict. You can pick a specific MXU yourself with pltpu.matmul_acc_lhs, and not every transpose needs an XLU: some only relabel dimensions, and others happen inside a matmul.

To be continued…

KUDOSDon’t
MoveThanks!