3 / 9 GPU / TPU Architecture

The Block That Will Not Fit in VMEM

A Pallas TPU kernel is written with pl.BlockSpec(block_shape=(1024, 4096), ...) for three float32 operands: two inputs and one output. On this TPU generation VMEM is on the order of 16 MB; HBM holds tens of GB.

What happens when this kernel is lowered, and why?

A

It is rejected because 1024 and 4096 are not multiples of 8 and 128, which the last two block dimensions must satisfy.

B

It is rejected because a BlockSpec may never be larger than one quarter of the full array's shape.

C

It is rejected. Each block is 1024 x 4096 x 4 bytes = 16.8 MB and three of them need about 50 MB, but block windows have to live in VMEM, which is around 16 MB. A smaller block, say (256, 4096) at 4 MB each, fits.

D

It compiles and runs, spilling whichever blocks do not fit back to HBM on demand, at a bandwidth cost.

E

It runs fine: VMEM only ever holds the 8 x 128 vector-register tile, and the rest of the block streams from HBM element by element.