The GPU Memory Hierarchy — HBM, SRAM, Registers, and Why Movement Wins

What sets a GPU's speed is not the arithmetic units but where the data sits — registers, shared memory, L2 or HBM — and how many times it is moved. Capacities and bandwidths by order of magnitude, arithmetic intensity and tiling, a roofline per level of the hierarchy, and finally FlashAttention: more FLOPs, less time.

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

An analogy: palm, bench, shelf, warehouse

A skilled worker is assembling parts. Their palm holds two or three. The bench in front of them holds a few dozen, the shelf behind them a few hundred, and the full stock sits in a warehouse across the yard.

What limits this person's output is not how fast their hands move but where the next part happens to be. From the palm it is instant; from the warehouse it is a one-minute walk. If every part means a walk, no amount of hand speed gets you past one unit per minute.

A GPU has exactly this shape. The arithmetic units are more than fast enough; runtime is decided by which drawer the data is in and how many round trips it takes. The memory wall from scratch covered why moving data is expensive, starting from the physics of a wire. This piece opens the drawers of one concrete machine.

The GPU has four drawers

The numbers below are for an NVIDIA A100.

Capacity grows by more than five orders of magnitude from top to bottom, and speed falls by about as much. The FlashAttention paper puts measured figures on the gap for the A100: roughly 20 MB of on-chip SRAM at about 19 TB/s, against 40 GB of HBM at about 1.5 TB/s. Two thousand times the capacity, more than ten times slower. Step off the chip to host DRAM and bandwidth drops another two orders of magnitude.

Two things about this table surprise people coming from CPUs.

First, the register file is enormous. A CPU core has a few hundred bytes of registers; a GPU has 256 KB per SM, and with 108 SMs an A100 carries roughly 27 MB in total — the same order as its L2. The reason is that a GPU keeps the state of tens of thousands of threads resident at once and switches between them to fill memory stalls. Registers are both the fastest scratch space and the inventory that hides latency — the flip side of the occupancy story in why GPUs are fast.

Second, shared memory is not a cache. A cache decides for itself what to keep; shared memory (__shared__ in CUDA) is SRAM whose contents you choose, with the guarantee that what you put there stays until you remove it. That guarantee is what makes optimisations like FlashAttention — where you design the number of round trips — possible at all.

Putting "movement wins" into a formula

A single kernel has a floor on its runtime that no amount of tuning moves.

$W$ is the operation count in FLOPs, $\pi$ the peak arithmetic throughput, $Q$ the bytes crossing some level of the hierarchy, and $\beta$ the bandwidth of that level. In words: you cannot finish before the slower of "done computing" and "done moving". The ratio of those two terms is arithmetic intensity.

Which says: how many operations do you get out of each byte you move — how many times you reused the material after carrying it in. If $I$ falls below the machine balance (peak throughput ÷ bandwidth), you are bandwidth bound; above it, compute bound. So far this is the same story as the memory wall.

The one thing this article adds is this: $Q$ is a different number at every level. For the same kernel, the bytes crossing HBM, the bytes crossing L2 and the bytes crossing shared memory are three separate quantities, so there are as many arithmetic intensities and as many rooflines as there are levels. "Is this kernel bandwidth bound?" is an incomplete question. The complete one is "at which level is it bandwidth bound?"