One language, three backends

Everything in this chapter lowers through MosaicThe MLIR dialect Pallas lowers to, and the last layer of the TPU stack you can read; only LLO below it is closed.taught in /l/mosaic → on TPU, the compiler pass that turns the algebra plus the schedule into the vector-register operations and DMAsAn asynchronous copy between memories that runs while compute continues; the grid pipeline is DMAs the runtime writes for you.taught in /l/pallas → a v5e or v6e chip actually runs. The same Pallas frontend, the same refs, the same grid and BlockSpecHow one array is carved for the grid: a block shape plus an index map saying which block each grid step sees.taught in /l/pallas → grammar, also targets GPU, through two different backends: a Triton backend, and Mosaic GPU, which is warp32 GPU threads scheduled as one unit; the GPU hides latency by switching among resident warps rather than by pipelining a scratchpad.taught in /l/tpu →-specialized for H100-class hardware.

What transfers across all three backends is the part of the mental model that was actually about the algorithm, not about any one chip. The split between what runs as algebra inside the kernel body and what the schedule, the grid, the blocking, controls outside it: that split is a way of thinking about a kernel, and it holds regardless of where the kernel eventually runs. The ref-and-block abstraction itself transfers too, and so does interpret=TrueRuns kernel logic on any machine for correctness work. It ignores memory spaces and never sees compile-time errors like lattice or VMEM violations.taught in /l/pallas →, since checking that a kernel computes the right values never depended on which chip it would later run on.

What does not transfer is the execution model underneath, and the difference is worth naming precisely rather than waving at. TPU runs the grid as a sequential software pipeline, one step overlapped with the next, because that is how a single large core with a big VMEMThe TPU’s software-managed vector scratchpad, about 128 MiB. Blocks must be staged here before compute touches them; what is resident is what your schedule staged.taught in /l/tpu → wants to be fed: few large steps, each one doing a lot of work, staged carefully so the next step's data is ready before the current one finishes. A GPU is built from thousands of smaller cores, so the same schedule maps instead onto thousands of concurrently scheduled blocks, a different shape of parallelism entirely rather than a faster version of the same one.

VMEMThe TPU’s software-managed vector scratchpad, about 128 MiB. Blocks must be staged here before compute touches them; what is resident is what your schedule staged.taught in /l/tpu → has no direct GPU equivalent either. The closest structure is shared memory, allocated and managed differently: smaller per unit of compute, and shared among the threads running on one GPU core rather than serving a single large TPU core the way VMEM does. The automatic pipelining MosaicThe MLIR dialect Pallas lowers to, and the last layer of the TPU stack you can read; only LLO below it is closed.taught in /l/mosaic → builds for you on TPU becomes, on Mosaic GPU, warp32 GPU threads scheduled as one unit; the GPU hides latency by switching among resident warps rather than by pipelining a scratchpad.taught in /l/tpu → choreography you specify yourself: which groups of threads load, which compute, and when they hand off to each other, a decision the TPU backend makes for you by walking the grid in the order you already know.

A warp32 GPU threads scheduled as one unit; the GPU hides latency by switching among resident warps rather than by pipelining a scratchpad.taught in /l/tpu → is a fixed-size group of threads that execute the same instruction together, the unit MosaicThe MLIR dialect Pallas lowers to, and the last layer of the TPU stack you can read; only LLO below it is closed.taught in /l/mosaic → GPU's specialization schedules around. TPU has no equivalent unit at all, because a TPU core was never built out of many small threads in the first place; it is one large core reading and writing a big VMEMThe TPU’s software-managed vector scratchpad, about 128 MiB. Blocks must be staged here before compute touches them; what is resident is what your schedule staged.taught in /l/tpu →, not a crowd of small ones coordinating over shared memory.

None of this means starting over on GPU. It means the algorithm and the correctness story you built in Pallas on TPU were never chip-specific to begin with. The scheduling code is what changes, and knowing exactly which part that is, the part this chapter has spent eight sections on, is what makes moving between backends a matter of rewriting a schedule rather than relearning what a kernel even is.