5 / 9 GPU / TPU Architecture

A Skinny Matmul on a 128x128 MXU

A TPU MXU before v6e is a 128 x 128 systolic array of multiply-accumulators, 16K MACs per cycle, holding a tile of the weight matrix stationary while activations stream through. Pallas also requires a block's last two dimensions to be multiples of 8 and 128.

A layer computes X @ W with X of shape (256, 16) and W of shape (16, 512), so the contracting dimension is 16. What share of the array's multiply-accumulate slots does real data occupy, and what fixes it?

A

87.5%, since only 16 of the 128 rows go unused.

B

12.5%, and the fix is to raise the batch dimension from 256, which fills the rest of the array.

C

12.5%: the contracting dimension is padded from 16 up to 128, so 112 of every 128 rows of processing elements multiply zeros. Reshaping or fusing so the contraction is a multiple of 128 is what recovers the array.

D

100%: the hardware pads transparently and the padded lanes are gated off, so utilization tracks useful FLOPs.

E

About 1.6%, since both dimensions are padded to 128, giving 16/128 x 16/128.