Tiled matmul

Cut each matrix into tiles and multiply as if the tiles were numbers: row of tiles times column of tiles. The only rule is the one from Part 2, applied to tiles: inner sizes must match, and order matters.

Cut into $2 \times 2$ tiles. To get the top-left tile of $C$, pair the first tile-row of $A$ with the first tile-column of $B$.

Each tile product is an ordinary $2 \times 2$ matmul; the results add up. Compare with the top-left of the full product: $AB$ starts with rows $[2, 3]$ and $[2, 2]$.

Naively, each output entry reads a whole row and a whole column. A GPU has 80 GB of slow HBM and only tens of megabytes of fast on-chip memory. Load a $b \times b$ tile once into the fast memory and every number in it gets used $b$ times before it’s evicted. Reads fall by a factor of $b$, so arithmetic intensity rises by about $b$.

By what factor does tiling cut the reads from slow memory? How many tile products make up one tile of $C$?

Each output tile accumulates 32 products of $128 \times 128$ tiles, summed in FP32 (the number formats deck).

Compute $C_{11} = A_{11}B_{11} + A_{12}B_{21}$.

$A_{12}$ and $B_{21}$ both swap, and two swaps make the identity.

• Block multiplication: $C_{ij} = \sum_k A_{ik}B_{kj}$ on tiles.
• A tile loaded once into fast memory is reused $b$ times: intensity ≈ $b$.
• Kernels accumulate tile products in FP32.
• FlashAttention is the same idea plus the online softmax.