GPU Memory Hierarchy

The memory closest to the arithmetic is tiny and enormously fast; the memory large enough to hold a model is far away and comparatively slow. On an A100 that is roughly 20 MB of on-chip SRAM at 19 TB/s against 40 GB of HBM at 1.5 TB/s. Most performance questions on this hardware turn out to be a question about which of those two a piece of data is sitting in.

Registers and shared memory per streaming multiprocessor, then L2, then HBM off-die. Capacity and bandwidth run in opposite directions at every step, so a kernel's cost is often predicted better by how many bytes it moves across that boundary than by how many operations it performs. That is the observation FlashAttention is built on, and the reason a rewrite with identical arithmetic can be several times faster.

The two levels differ by about three orders of magnitude in capacity and one in bandwidth, in opposite directions: 20 MB at 19 TB/s against 40 GB at 1.5 TB/s, a ratio of roughly thirteen to one per byte. A kernel that must read b bytes from HBM cannot finish faster than b/1.5·10¹² seconds however few operations it performs — a floor no amount of arithmetic optimisation moves.