Foundations & Infrastructure

GPUs: How They Work and How to Make Them Fast

Foundations & Infrastructure GPUs
Topic
Foundations & Infrastructure
Venue
CS336 (Language Modeling from Scratch), Stanford · Lecture 5
Read
22 min
Source

In one line

A modern GPU is a factory that can compute far faster than it can fetch data, so almost every trick for making machine-learning code fast is really a trick for touching slow memory less often — and FlashAttention is the poster child of doing exactly that.

The breakdown

TL;DR

  • GPUs win by running thousands of tiny, dumb workers in lockstep instead of a few smart ones. That is great for the one operation deep learning lives on: matrix multiply.
  • On today’s hardware, compute got cheap faster than memory got fast. So the bottleneck for most ML code is not “how many multiplies” but “how many trips to slow memory.”
  • A single mental model, the roofline, tells you which regime you are in. If you are memory-bound, buying more compute does nothing; you have to move less data.
  • Every optimization in the lecture — low precision, operator fusion, recomputation, memory coalescing, and above all tiling — is a way to feed the compute units without stalling on memory.
  • FlashAttention is these tricks stacked together on the attention operation: tile the matrix multiplies, keep the intermediate scores in fast on-chip memory, and compute the softmax incrementally so you never write the giant attention matrix to slow memory at all.

Problem & Motivation

Language models improve predictably when you throw more compute at them (the scaling laws). For two decades, that compute came almost for free: chips got faster every year on their own (this was Dennard scaling — as transistors shrank, they also got faster and cooler at the same power). Around 2005 that free lunch ended. Clock speeds stopped climbing.

The industry’s answer was to stop making one worker faster and start adding more workers in parallel — that is the GPU. GPU throughput has grown more than 1000x in ten years, and there is no LLM progress without it.

But parallel hardware creates a new, sharper pain. You now have thousands of compute units that can multiply numbers absurdly fast — and a memory system that cannot hand them data quickly enough. A single square matrix multiply, the simplest thing in the world, can run at wildly different speeds depending on its exact size. The lecture’s motivating puzzle: why is multiplying two 1793×1793 matrices sometimes slower than multiplying two 1792×1792 matrices — bigger work, less time? If you cannot answer that, you cannot reason about why your model is slow. This lecture makes the GPU “less magic” so you can.

What’s New (Core Contribution)

This is a lecture, not a research paper, so the “contribution” is a way of seeing rather than a new algorithm. Three ideas do the heavy lifting:

  • Before: People treat GPU performance as a black box (“just use a bigger GPU”). Now: A single lens — the roofline model — predicts whether any given kernel is limited by compute or by memory, so you know which optimizations will actually help.
  • Before: Optimizations feel like a grab-bag of unrelated tricks. Now: They are unified under one goal — respect the memory hierarchy. Low precision, fusion, recomputation, coalescing, and tiling are all “touch slow memory less.”
  • Before: FlashAttention looks like clever, opaque CUDA wizardry. Now: It is shown to be nothing but the standard tricks (tiling + fusion + an online-softmax accounting trick) applied to attention. Once you see the parts, the “magic” disappears.

How It Works (Technically)

Part 1 — What a GPU actually is

A CPU is built for latency: a few powerful cores, lots of cache and branch-prediction machinery, so each individual thread finishes as fast as possible. A GPU is built for throughput: sacrifice per-thread speed and cleverness, pack in thousands of tiny arithmetic units (ALUs), and win by sheer parallel volume. Total work done per second is what matters, not how fast any one item finishes.

The physical structure, from big to small:

  • SM (Streaming Multiprocessor): the GPU has many of these; each independently runs a “block” of work. An NVIDIA A100 has 108 SMs — remember that number, it explains the puzzle later.
  • SP (Streaming Processor) / core: inside each SM, many of these run individual threads.
  • Tensor Cores: special circuits (introduced in the V/T generations) that do nothing but small matrix multiplies — and do them more than 10x faster than the general-purpose units. This is why “make it a matmul” is a performance strategy.

The parts and how they nest:

flowchart LR
  subgraph GPU
    direction LR
    subgraph SM1[SM · Streaming Multiprocessor]
      direction TB
      W[Warp = 32 threads in lockstep] --> SP[Many SP cores]
      TC[Tensor Core: fast matmul]
      SHM[(Shared memory / L1 · fast, tiny)]
      SP <--> SHM
    end
    L2[(L2 cache · on-die)]
    SM1 <--> L2
  end
  HBM[(Global memory / HBM · big, slow)]
  L2 <-->|the expensive trip| HBM

The execution model has three players:

  • Threads do the work. All threads run the same instruction on different data. NVIDIA calls this SIMT (Single Instruction, Multiple Threads).
  • Blocks are groups of threads. Each block lives on one SM and shares that SM’s fast scratchpad memory.
  • Warps are the unit that actually executes: 32 consecutively numbered threads that step together, in lockstep. This detail drives two later tricks (divergence and coalescing).

The memory hierarchy is the whole game. Closer to the compute unit = faster and smaller:

LevelWhereSpeedSizeCost
RegistersPer threadFastestTiny
Shared memory / L1 (SRAM)Inside the SM~8x faster than DRAMSmall (KBs)~100x more expensive per byte
L2 cacheOn the GPU dieMediumBigger
Global memory / HBM (DRAM)Chips beside the GPUSlowestBig (GBs)Cheap per byte

The killer fact: compute has scaled faster than memory bandwidth (“the memory wall”). You have more multiply-power than you can keep fed. So the job of a fast kernel is to pull data into fast on-chip memory once and reuse it, instead of repeatedly crossing to slow global memory.

The memory hierarchy as depth layers (schematic). Registers and shared memory sit inside the SM and are tiny but fast; global memory (HBM) is huge but far. Drag to orbit. Every optimization in Part 2 is about keeping work in the near, fast layers.

Part 2 — The roofline: the one model that tells you what to fix

Whether a kernel is limited by compute or by memory comes down to one number: arithmetic intensity = FLOPs performed ÷ bytes moved. “How much math do I do per byte I fetch?”

  • Low intensity (few operations per byte) → you spend your time waiting on memory → memory-bound. Adding compute does nothing.
  • High intensity (lots of operations per byte) → the compute units are the limit → compute-bound. Now, and only now, does faster/more compute help.

The roofline plots achievable performance against intensity. It rises along a slanted line (the memory-bandwidth “roof”) until it hits the flat line (the peak-compute “roof”). The corner where they meet is the intensity you need to reach to stop being memory-starved. Most naive ML kernels sit on the left, slanted part — memory-bound. The entire goal of Part 2 is to push right, off the memory roof.

Interactive roofline (schematic). Drag the slider to change a kernel's arithmetic intensity. Left of the ridge point it is memory-bound (buying compute is wasted); right of it, compute-bound. The optimizations below all move a kernel rightward.

Now the toolbox, each item explained as “why it touches memory less”:

1. Control divergence (the one non-memory issue). Because a warp’s 32 threads execute in lockstep, an if where some threads go one way and some the other forces the hardware to run both branches and mask out the threads that shouldn’t act. You pay for both paths. Keep branches aligned within a warp.

2. Low precision. Fewer bits per number = fewer bytes to move = higher arithmetic intensity for free. A ReLU (x = max(0, x)) over a float32 vector moves 8 bytes per FLOP; in float16 it moves 4 bytes per FLOP — half the memory traffic for the same math. And Tensor Cores run fastest in low/mixed precision, so low precision speeds up the multiply itself too.

3. Operator fusion. Picture the GPU as a factory and global memory as a warehouse across town. Every separate operation ships the data to the warehouse and back. Computing sin²(x) + cos²(x) naively launches 5 separate kernels — five round trips. Fusing them into one kernel does all the pointwise work while the data is on-site, writing to slow memory only once. Compilers like torch.compile do these “easy” fusions automatically.

4. Recomputation. During backprop you normally store every layer’s activations so you can reuse them on the backward pass. Storing and reloading them is a lot of slow-memory traffic. Counterintuitively, throwing them away and recomputing them from scratch when needed can be faster, because compute is cheap and memory is expensive — one example cuts memory accesses to 5/8ths. (This is activation checkpointing.)

5. Memory coalescing. DRAM is read in bursts — one access returns a whole contiguous chunk, not one number. Accesses are coalesced when the 32 threads of a warp read addresses that all fall inside the same burst: one burst serves the whole warp. If the threads read scattered addresses, you pay for many bursts. For a row-major matrix, threads walking down a column are scattered and uncoalesced; threads walking along a row are coalesced. Layout matters.

6. Tiling (the big one). In a naive matmul, every input element is read from global memory N times (once per output it contributes to). Tiling fixes this: cut the matrices into small tiles, load a pair of tiles into the SM’s shared memory once, and compute all the partial products that use them before moving on. Reads that used to hit slow global memory now hit fast shared memory, and they can be coalesced on the way in. With tile size T, each input is read from global memory a factor of T fewer times. That is the single biggest lever in dense linear algebra.

Solving the matmul mystery

Two tiling side-effects explain why bigger can be faster:

  • Tile quantization / alignment. If the matrix dimension does not divide evenly by the tile size, the edge tiles are partly empty — wasted work — and memory bursts stop aligning with tile boundaries, so you lose coalescing and have to pad. Sizes that align cleanly with tiles (and with the 32-thread warp / burst width) run much faster.
  • Wave quantization. Tiles are handed out to SMs in “waves.” At 1792 with a 256×128 tile you get 7×14 = 98 tiles, which fit in the A100’s 108 SMs in a single wave. Bump to 1793 and you need 8×15 = 120 tiles — more than 108 — so a second wave spins up for the leftover 12 tiles, and the whole op waits nearly twice as long. More work, worse time. That is the mystery, solved: performance is quantized by how tiles pack into a fixed number of SMs.

The algorithm, simplified — tiled matmul (the core idea)

# Tiled matrix multiply C = A @ B, the idea behind coalescing + shared-memory reuse.
# A: [M, K], B: [K, N], C: [M, N]. Tile size T (e.g. 128). "shared" = fast on-chip scratchpad.
def tiled_matmul(A, B, T):
    M, K = A.shape; K2, N = B.shape
    C = zeros((M, N))
    for i0 in range(0, M, T):            # each (i0, j0) block == one GPU block on one SM
        for j0 in range(0, N, T):
            acc = zeros((T, T))          # partial sums live in registers, never touch global mem
            for k0 in range(0, K, T):    # walk across the shared dimension, one tile-pair at a time
                a_tile = load_to_shared(A[i0:i0+T, k0:k0+T])   # ONE coalesced global read...
                b_tile = load_to_shared(B[k0:k0+T, j0:j0+T])   # ...then reused T times from fast mem
                acc += a_tile @ b_tile   # all reuse hits shared memory, not slow global memory
            C[i0:i0+T, j0:j0+T] = acc    # write each output exactly once
    return C

The whole point is the two load_to_shared calls: each slow global read is amortized over T fast reuses.

Part 3 — FlashAttention is just these tricks, on attention

Standard attention is three matrix multiplies with a softmax in the middle: score S = Q @ Kᵀ, normalize P = softmax(S), output O = P @ V. The problem is S: for a sequence of length n it is an n×n matrix. Materializing it — writing all numbers to global memory and reading them back for the softmax — is pure memory traffic and dominates the cost. Attention is memory-bound, not compute-bound.

Here is the data flow — standard attention writes the giant S to slow memory; FlashAttention keeps everything on-chip and streams:

flowchart TD
  Q[Q tile] --> S[S = Q·Kᵀ tile]
  K[K tile from global] --> S
  S --> E[exp, fused in same kernel]
  E --> ON{Online softmax: update running max m and running sum l}
  V[V tile from global] --> ACC[Accumulate O += P·V, rescaled]
  ON --> ACC
  ACC -->|next K,V tile| S
  ACC --> O[Output O written once]

FlashAttention never writes S to global memory. It stacks the Part-2 tricks:

  1. Tile the Q, K, V matmuls exactly like tiled matmul above — figure 1 of the paper is literally that.
  2. Fuse the exponential (the exp in softmax) into the same kernel, so scores are consumed on-chip the instant they are produced.
  3. Online (incremental) softmax — the one genuinely subtle piece. Softmax needs the max and the sum over a whole row for numerical stability, but you only ever hold one tile of that row at a time. The trick (from Milakov & Gimelshein 2018): keep a running max and a running sum, and when a new tile arrives with a bigger max, rescale the sum you already have (a telescoping correction) so the final result is exactly the true softmax. This lets you compute the softmax tile-by-tile without ever seeing the full row at once.

Result: attention runs with far less global-memory traffic, so it is dramatically faster and uses far less memory — with identical outputs (it is exact, not an approximation). The backward pass uses the same idea plus recomputation (trick 4): recompute the scores tile-by-tile instead of storing the n×n matrix.

Why attention is memory-bound and how tiling fixes it (schematic). Toggle between "materialize S" (write/read the whole n×n score matrix to slow global memory) and "FlashAttention" (stream K/V tiles through fast shared memory, keep only running max + sum). Watch the global-memory traffic counter.

Built on Prior Work

Prior ideaWhat it gaveWhat this lecture adds
Kaplan et al., Neural Scaling LawsCompute predictably buys model qualityFrames GPUs as the engine that makes that compute available
Dennard scaling → its end (~2005)Free single-thread speedupsExplains the pivot to parallel (GPU) scaling
Bill Dally, HotChips keynoteGPU throughput up >1000x in 10 yearsGrounds “no LLM scaling without GPU scaling”
Williams et al., Roofline modelCompute-vs-memory-bound diagnosisUses it as the organizing lens for all optimizations
Horace He, Making GPUs go BRRRFusion / factory-vs-warehouse intuitionTies fusion to the roofline
PyTorch AOTAutograd / min-cut recomputeOptimal activation checkpointingFramed as “trade compute for memory”
Milakov & Gimelshein 2018, Online softmaxStreaming, one-pass softmaxThe accounting trick that makes FlashAttention possible
Dao et al. 2022/2023, FlashAttentionExact, IO-aware fast attentionShown to be tiling + fusion + online softmax combined
thonking.ai “matmul shapes” postTile & wave quantizationThe worked explanation of the matmul mystery

Results & Evidence

As a lecture, the “evidence” is a set of demonstrations rather than a benchmark table, and they are convincing:

  • The matmul mystery is fully explained, quantitatively: the 1792→1793 slowdown falls straight out of 98 vs 120 tiles against 108 SMs. This is a real, reproducible A100 effect.
  • Arithmetic-intensity math checks out: float16 halves ReLU’s bytes-per-FLOP versus float32; the tiling analysis shows a clean factor-of-T reduction in global reads. These are first-principles, not hand-waving.
  • FlashAttention’s speedups are real and widely reproduced in production (it ships in PyTorch and every serious inference stack).

Caveats worth naming, because you sell this for a living:

  • The concrete numbers (108 SMs, 256×128 tiles, the 1792 boundary) are A100-specific. Newer GPUs (H100/Blackwell) move the boundaries; the reasoning transfers, the exact numbers do not.
  • “Compute is basically free, trade memory for it” holds today because of the memory wall. It is a regime, not a law — on very compute-heavy, low-precision kernels you can flip to compute-bound.
  • The lecture stays at the mechanism level and skips the messy CUDA reality (register pressure, occupancy, bank conflicts) that decides whether a hand-written kernel actually hits these ceilings.

How You’d Use It

You run an AI services company, so the payoff is diagnosis and cost control, not writing CUDA.

  • Serving-cost estimation and quoting. Knowing attention is memory-bound tells you why long-context requests get expensive fast (the n×n traffic) and why FlashAttention / paged-KV backends (vLLM, TensorRT-LLM) cut cost. You can size a client’s GPU bill from the shape of their workload instead of guessing.
  • The right first question for any “it’s slow” ticket: is this memory-bound or compute-bound? If memory-bound (most inference, most attention, most small-batch work), throwing a bigger GPU at it is wasted spend — you want fusion, quantization, and a FlashAttention-style backend instead. This alone saves clients real money.
  • Quantization decisions with eyes open. “Use FP16/FP8/INT8” is not just an accuracy trade — it directly raises arithmetic intensity and unlocks Tensor Cores. You can explain to a client what they gain in throughput and what they risk in accuracy.
  • Batching strategy. Small batches keep you memory-bound; larger batches raise intensity and push you toward compute-bound (better utilization). This is the lever behind continuous batching in modern serving.
  • Buy-vs-build clarity. The lecture shows FlashAttention is “just” known tricks — but also that hand-writing kernels is brutal. The honest client answer is almost always buy the optimized backend, and now you can say exactly why it is fast.

Build Your Own (Minimal Recipe)

You do not need to write a CUDA kernel to internalize this. Smallest version that captures ~80% of the value:

  1. Reproduce the roofline intuition in NumPy. Time an elementwise op (memory-bound) vs a large square matmul (compute-bound). Plot GFLOP/s vs arithmetic intensity. You will see the two regimes. (An afternoon.)
  2. Write a tiled matmul in Python (the pseudocode above), then compare it to a naive triple loop by counting global-memory reads (a counter on the “slow” array). Watch tiling cut reads by a factor of T. This teaches the core idea without a GPU.
  3. Reproduce the wave-quantization cliff. On any NVIDIA GPU, torch.matmul two square matrices sweeping size from ~1700 to ~1900 and time each. You will find the sawtooth. (Half a day; needs a GPU.)
  4. Implement online softmax in NumPy: process a long vector in chunks, maintaining running max + running sum with rescaling, and check it matches scipy.special.softmax. This is the one non-obvious FlashAttention piece — nail it in ~30 lines.
  5. The hard part(s): an actual fused, tiled GPU kernel. Do not start in raw CUDA. Reach for Triton (OpenAI’s Python-like kernel language) — its official tutorials walk you through a fused-softmax and a matmul kernel, and a working “flash-attention in Triton” is a well-trodden exercise. That is the realistic ceiling for a services shop.

Libraries/tools to reach for: PyTorch (torch.compile for automatic fusion, scaled_dot_product_attention for built-in FlashAttention), Triton for custom kernels, nsight-compute/torch.profiler to measure whether you are memory- or compute-bound, vLLM / TensorRT-LLM as the production backends that bake all of this in.

How to Improve It

Limitations of the techniques here, which are where the field is actively pushing — each a testable direction:

  • Attack the memory wall differently. Tiling reduces traffic; you can also increase effective bandwidth. Test grouped-query / multi-query attention and KV-cache quantization to shrink the very n×n/KV traffic that makes attention memory-bound.
  • Go lower precision, carefully. FP8 (and experiments in FP4) push arithmetic intensity further. The testable question is the accuracy floor per model/task — build a small harness that sweeps precision and plots throughput vs quality, so quantization is a measured decision, not a vibe.
  • Better recomputation policies. Activation checkpointing is a coarse “store vs recompute” switch. The min-cut framing suggests a per-operator optimal policy; you could test learned or profile-guided checkpointing that recomputes only the cheap-to-redo, expensive-to-store ops.
  • Auto-tiling for the target GPU. The wave-quantization cliff is hardware-specific. A small autotuner that picks tile sizes to align with the actual SM count and burst width (à la Triton autotune) can claw back the sawtooth losses automatically — a concrete, shippable tool.
  • Fuse more aggressively across the whole block. FlashAttention fuses within attention; the frontier is fusing attention + normalization + the MLP’s pointwise ops into fewer kernels to kill even more round trips. Test end-to-end fused transformer blocks and measure the memory-traffic reduction.

Glossary

  • SM (Streaming Multiprocessor) — an independent compute cluster on the GPU; the A100 has 108. Runs one block at a time (conceptually).
  • SP / core — a small execution unit inside an SM that runs one thread.
  • Tensor Core — dedicated circuitry that does small matrix multiplies >10x faster than general units, fastest in low precision.
  • Thread / Block / Warp — a thread does the work; a block is a group of threads sharing on-chip memory on one SM; a warp is the 32 threads that actually execute in lockstep.
  • SIMT (Single Instruction, Multiple Threads) — all threads in a warp run the same instruction on different data.
  • Global memory / HBM (DRAM) — the big, slow memory chips beside the GPU.
  • Shared memory / SRAM (L1) — small, ~8x faster scratchpad inside the SM; the target for tiling.
  • Arithmetic intensity — FLOPs ÷ bytes moved; the number that decides memory-bound vs compute-bound.
  • Roofline model — a plot of achievable performance vs arithmetic intensity that shows whether memory bandwidth or peak compute is your ceiling.
  • Memory-bound / compute-bound — limited by data movement vs limited by math throughput.
  • Coalescing — arranging a warp’s memory reads so they fall in one DRAM burst, so one access serves all 32 threads.
  • Burst mode — DRAM returns a contiguous chunk per access, not a single value.
  • Operator fusion — merging several ops into one kernel so intermediates stay on-chip instead of round-tripping to global memory.
  • Recomputation / activation checkpointing — discarding stored activations and recomputing them on the backward pass to save memory traffic.
  • Tiling — splitting matrices into small blocks loaded once into shared memory and reused, cutting global reads by a factor of the tile size.
  • Tile / wave quantization — performance jumps caused by tiles not dividing the matrix evenly, or by needing more tile-waves than the GPU has SMs.
  • Online softmax — computing softmax in a single streaming pass using a running max and running sum with rescaling; the trick that makes FlashAttention exact and tileable.
  • FlashAttention — an exact, IO-aware attention implementation that tiles the matmuls, fuses the exponential, and uses online softmax to avoid writing the n×n score matrix to global memory.
  • Dennard scaling — the historical trend (ended ~2005) where shrinking transistors got faster at constant power; its end forced the shift to parallelism.