Pallas

A Pallas kernel is a Python function over refs: windows into arrays that the runtime has already staged into VMEM. You never load from HBMThe chip’s main memory: large, far, and the resource memory-bound ops spend. 8.2e11 bytes per second on v5e, 1.6e12 on v6e.taught in /l/tpu → yourself. The BlockSpec for each array says how it is carved into blocks and which block a given grid position sees, via an index map. The grid says how many steps run. Three objects, and the separation between them is the discipline this whole site color-codes: the kernel body is algorithm (copper), the specs and grid are schedule (steel).

def matmul_kernel(a_ref, b_ref, o_ref):      # algorithm
    k = pl.program_id(2)
    @pl.when(k == 0)
    def _init():
        o_ref[...] = jnp.zeros_like(o_ref)
    o_ref[...] += jnp.dot(a_ref[...], b_ref[...])

pl.pallas_call(                               # schedule
    matmul_kernel,
    grid=(M // bm, N // bn, K // bk),
    in_specs=[pl.BlockSpec((bm, bk), lambda i, j, k: (i, k)),
              pl.BlockSpec((bk, bn), lambda i, j, k: (k, j))],
    out_specs=pl.BlockSpec((bm, bn), lambda i, j, k: (i, j)),
    out_shape=jax.ShapeDtypeStruct((M, N), dtype),
)

On TPU the grid is a software pipeline: while step k computes, the runtime streams step k+1's blocks into the other half of a double buffer. You wrote no DMA, yet transfers overlap compute. The condition: your blocks must be sized so transfer time hides under compute time, and the whole working set (roughly three blocks in flight per array) must fit VMEM. Stage 1's EX·02 and EX·07 instruments play the pipeline on and off so the difference is visible rather than asserted.

Two practical notes that save days: interpret=True runs kernel logic anywhere, including your laptop, which is how every lab in this track verifies before touching a chip; and the (8, 128) tiling lattice from the chip layer governs your block shapes, so violations fail at compile time with errors that make sense only after you know the lattice exists.