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?"

Reuse is the only way to raise intensity

$W$ is fixed by the problem you are solving, so the only way to raise $I$ is to shrink $Q$ — and in practice there is exactly one way to do that: keep data in a fast drawer and use it many times. Tiling.

Cut an $n \times n$ matrix multiply into $b \times b$ tiles. There are $(n/b)^3$ tile triples; each reads two tiles ($2b^2$ elements) and performs $2b^3$ operations. With $s$ bytes per element,

In words, arithmetic intensity is proportional to the edge length of the tile. In fp16 ($s = 2$) a 128-wide tile gives 64 FLOP/byte. The bigger the fast drawer, the bigger the tile it holds, and the less you carry.

The flip side: any computation that touches each element once — adds, ReLU, normalisation, embedding lookups — cannot raise $I$ by tiling harder. It is structurally bandwidth bound forever.

A roofline per level

Work an example on paper. An $n = 8192$ fp16 GEMM with $b = 128$: bytes from HBM are $2n^3/b$ elements × 2 bytes ≈ 17 GB, and the work is $2n^3 \approx 1.1$ TFLOP. Take 312 TFLOP/s of fp16 tensor-core peak and 1.5 TB/s of HBM bandwidth, and you get 3.5 ms of compute against 11 ms of transfer. The conclusion would be that even large GEMMs are bandwidth bound.

Real high-performance GEMMs run faster than that estimate. The faulty part is the assumption: the calculation above says every tile read goes all the way to HBM. In reality there are 40 MB of L2, and neighbouring SMs ask for the same tiles at nearly the same moment, so the second and later requests are served on chip. The threadblock swizzle used by libraries such as CUTLASS — reordering which block runs when — exists precisely to farm those L2 hits.

That is what a per-level roofline is for in practice. Once you have estimated HBM traffic on paper, always check it against measured traffic. If the measurement is far below your estimate, L2 is doing the work; if it matches, L2 is being bypassed and there is headroom left. In Nsight Compute, dram__bytes.sum is the number of bytes that actually crossed HBM.

Counting round trips in code

A more everyday example. Count the trips:

y = torch.relu(x)   # read x, write y   → 2 trips
z = y * scale       # read y, write z   → 2 trips
w = z + bias        # read z, write w   → 2 trips

If those three lines are three separate kernels, every element crosses HBM six times for three operations. That is $I = 3/(6 \times 2) = 0.25$ FLOP/byte, three orders of magnitude below the A100's fp16 machine balance of roughly 200 FLOP/byte. The arithmetic units spend over 99% of the time waiting.

Fuse the three into one kernel and you read $x$ once and write $w$ once: six trips become two. Not a single FLOP was removed, and it runs about three times faster. A large part of what torch.compile does is this fusion, and it is the first thing to check when hand-writing Triton.

Why FlashAttention works

Apply the same logic to attention. As covered in attention from scratch, self-attention forms an $N \times N$ score matrix with $QK^\top$, applies softmax, and multiplies by $V$. Implemented literally, that matrix is written out to HBM, read back for the softmax, and read again for the product with $V$.

Put a number on it. At sequence length $N = 4096$ in fp16, one head's score matrix is $4096^2 \times 2 = 33.5$ MB. With 32 heads and a batch of 8 that is 8.6 GB, moved repeatedly, layer after layer. And it is not even output: it is a discarded intermediate.

FlashAttention's idea is to never write that matrix to HBM at all. Split $Q$, $K$ and $V$ into blocks that fit in shared memory, compute partial scores block by block, and carry the running maximum and denominator of the softmax forward as you go (online softmax). The $N \times N$ matrix never exists; it is born and dies inside SRAM. The paper's analysis puts HBM accesses at $O(N^2 d^2 / M)$, where $M$ is the SRAM size, against $\Omega(Nd + N^2)$ for the standard implementation. In words: the bigger the fast drawer, the less you carry — the same statement as the tiling formula above.

The most important detail is in the backward pass. Because the score matrix was never saved, FlashAttention recomputes it during backpropagation. The operation count goes up. It is still faster. In a bandwidth-bound world, FLOPs are not the currency. Bytes are, and redoing arithmetic to avoid a transfer is a trade that pays.

How this shows up on the job

Who touches it, and when. Inference-platform engineers raising how many requests one GPU can serve; training engineers shaving step time; anyone writing a custom kernel in Triton or CUDA. The headroom that remains without changing a line of the model architecture usually lives at this level.

The first metrics to look at. In Nsight Compute's Memory Workload Analysis (ncu --set full): dram__bytes.sum for real HBM traffic, L1/L2 hit rates, and smsp__sass_average_data_bytes_per_sector_mem_global_op_ld for how well accesses coalesce. In PyTorch, get per-kernel time from torch.profiler first and only then descend into the top few. Compute $Q$ on paper before you read any profiler output — without an expectation, you cannot tell whether a number is good or bad.

The knobs you actually turn. Shared memory per block (going past 164 KB requires setting cudaFuncAttributeMaxDynamicSharedMemorySize explicitly), register caps via __launch_bounds__ or -maxrregcount, tile shapes (BLOCK_M / BLOCK_N / BLOCK_K in Triton), asynchronous copies with cp.async on Ampere and later, and TMA on Hopper. Every one of them decides what goes in the fast drawer and when the carrying starts.

Traps that turn into incidents.

• Spilled registers land in "local memory", which is HBM. The name suggests otherwise. This is the classic reason a kernel gets suddenly slower after a small change; -Xptxas -v reports the spill bytes.
• More shared memory means lower occupancy. SRAM per SM is fixed, so a block that claims more of it leaves room for fewer concurrent blocks — less inventory to hide latency with.
• Bank conflicts. Shared memory is split into 32 banks; when threads in a warp hit different addresses in the same bank, the accesses serialise. Padding a transpose buffer from [32][32] to [32][33] is the standard trick for avoiding exactly that.
• Uncoalesced access. A warp reading contiguous addresses needs a handful of transactions; a strided one can need up to 32. The byte count is unchanged, but effective bandwidth drops by an order of magnitude.
• Small benchmarks lie. With 40 MB of L2, a few megabytes of weights measures a world where everything fits on chip. Benchmark at production sizes.
• Host round trips. Calling .item() every step drags data down a path two orders of magnitude slower and inserts a synchronisation point, which makes unrelated kernels look slow in a profile.

A standard design-review question. "At sequence length 4096, how many bytes is one head's fp16 score matrix in HBM, and what does that imply?" — 33.5 MB; multiply by heads and batch and you are instantly in gigabytes, all of it a discarded intermediate. Being able to do that arithmetic on the spot is the same thing as being able to explain, in your own words, why FlashAttention exists.

Summary

• The GPU's drawers are registers, shared memory/L1, L2 and HBM: five-plus orders of magnitude in capacity, orders of magnitude in bandwidth
• The register file is huge because it is inventory for hiding latency; shared memory is not a cache but a bench whose contents you choose
• In $I = W/Q$, the $Q$ differs per level — you only know what to do once you can name which level is the bottleneck
• The only way to raise $I$ is reuse. Tiling gives $I = b/s$, so a bigger fast drawer wins
• FlashAttention adds FLOPs to remove transfers. When you are bandwidth bound, the currency is bytes, not FLOPs

Next we go after the bytes themselves — quantisation, read as an exchange rate against accuracy.