Memory Hierarchy and IO-Awareness

Modern hardware is memory-bound, not compute-bound for most ML operations. Understanding the memory hierarchy is the key to writing fast code.

The hierarchy (GPU, e.g. A100):

The bandwidth wall: A100 has 312 TFLOPS (bf16) but only 2 TB/s HBM bandwidth. If an operation loads more bytes than it computes FLOPs, it's memory-bound — the cores sit idle waiting for data.

Arithmetic intensity = FLOPs / bytes loaded. Operations are:

• Compute-bound — large matmuls, convolutions (high arithmetic intensity)
• Memory-bound — elementwise ops, softmax, layer norm, attention score computation (low arithmetic intensity)

FlashAttention — the canonical IO-aware algorithm:

• Standard attention materializes the $n \times n$ score matrix to HBM → $O(n^2)$ memory reads/writes
• FlashAttention tiles the computation: load blocks of $Q, K, V$ into SRAM, compute partial attention, write only the final output to HBM
• Result: exact attention with $O(n)$ memory and 2–4x wall-clock speedup
• The algorithm is mathematically identical — the speedup is purely from reducing IO

Practical implications:

• Kernel fusion — combine elementwise ops to avoid writing intermediates to HBM
• Activation checkpointing — trade recomputation for memory: don't store intermediates, recompute in backward pass
• Batch size tuning — larger batches improve arithmetic intensity of matmuls
• Sequence packing — avoid wasting compute on padding tokens by packing multiple sequences into one batch

See also: GPU Architecture and CUDA, Computational Complexity of Attention, Big-O and Complexity Analysis