CS336 Lecture 5: GPUs Win by Moving Data Less, Not by Making Each Thread Fast

🌏 中文版

This post covers CS336 Spring 2026 Lecture 5: GPUs, TPUs, taught by Tatsunori Hashimoto on April 13, 2026. Its primary source is the official lecture_05.pdf.

Lecture 2 used the roofline model to distinguish compute-bound from memory-bound operations. Lecture 5 opens the GPU to show where those limits come from. Its goal is not memorizing CUDA vocabulary but making questions such as “why is this shape slow?” and “why is FlashAttention fast?” stop looking like magic.

GPUs choose throughput over single-thread latency

CPUs use sophisticated control, caches, and a small number of powerful cores to minimize the completion time of individual threads. GPUs contain many simpler compute units running the same instructions over different data, optimizing total throughput.

A GPU contains many streaming multiprocessors (SMs). Threads form blocks, blocks are scheduled onto SMs, and threads execute in warps. When threads in a warp take different branches, the hardware must execute paths separately, causing divergence. Conditionals are legal, but they weaken the same-instruction, multiple-data advantage.

TPUs share the high-level principles: lightweight control, large matrix-multiplication units, and high-bandwidth memory. Their differences lie mainly in compute organization and device interconnects, not in one being capable of matrix multiplication and the other not.

Memory closer to an SM is smaller and faster

Threads have registers; a block can share shared memory; farther away are L2 cache and HBM or global memory. Data shared across blocks generally travels through slower global memory. Matrix-multiplication throughput has improved faster than memory bandwidth, so feeding tensor cores is now a central constraint.

This is arithmetic intensity in hardware form. Load data from HBM once and reuse it repeatedly in registers or shared memory, and substantial computation amortizes the transfer. Write every intermediate result back to HBM, and additional tensor cores wait for data.

Five optimizations all reshape the data path

Low precision reduces bytes per value and unlocks higher tensor-core throughput. FP8, MXFP8, and lower formats still need scale factors and selective precision; they are not unconditional casts of every tensor.

Operator fusion combines adjacent pointwise operations into one kernel. Intermediate values stay in registers or shared memory instead of returning to HBM after each step.

Recomputation discards selected activations and recalculates them later. It increases arithmetic while reducing memory traffic; in a memory-bound region, recomputing can be faster than loading.

Memory coalescing makes threads in a warp access contiguous addresses so a DRAM burst returns useful data. Identical arithmetic with a different traversal direction can perform very differently.

Tiling divides matrices into blocks that fit in shared memory. Each tile is loaded from HBM once, reused, and then written out. Tile size must also account for divisibility, alignment, register pressure, and the number of resident blocks.

Why matrix shapes create periodic performance cliffs

Adding a single row to a square matrix can create a sudden runtime jump. FLOPs may not be the cause. Tile boundaries, memory alignment, or wave quantization can be responsible: the number of SMs is fixed, so a small final group of blocks still requires another complete scheduling wave.

Benchmarks should therefore cover more than one convenient power-of-two shape, and advertised peak speed should not be treated as program speed. Sweep sizes, look for periodicity and discontinuities, then use a profiler to inspect kernels, memory transactions, and occupancy.

FlashAttention combines the preceding principles

Standard attention constructs a large score matrix, applies softmax, then multiplies by values. Writing each intermediate matrix to HBM makes data movement expensive as sequences grow.

FlashAttention tiles Q, K, and V and computes blocks in fast SRAM or shared memory. Online softmax maintains a running maximum and normalization sum, so the complete score matrix never needs to exist at once. The backward pass can recompute selected values to save additional storage.

The algorithm still computes exact attention. Its speedup comes mainly from IO-aware organization, not approximating away attention entries. It therefore combines the lecture's core tools: tiling, fusion, reduced HBM traffic, and deliberate recomputation.

How to debug performance after this lecture

Use roofline analysis to classify the operation, then inspect five questions: does precision actually use tensor cores, can operations be fused, are activations worth recomputing, are accesses coalesced, and do tiles fit the shape and hardware? Validate with several problem sizes and a profiler rather than inferring a cause from one wall-clock measurement.

GPU optimization is often presented as a bag of tricks. Lecture 5 reduces it to one question: where does data come from, where is it reused, and when is it written back?

Material fidelity

This lecture has a Spring 2026 schedule entry and a complete official PDF. This guide follows its hardware, performance-technique, and FlashAttention sections.

References

• CS336 Spring 2026 course and schedule
• Lecture 5 official slides
• FlashAttention
• CUDA Programming Guide