Stanford CS336 Language Modeling from Scratch | Spring 2026 | Lecture 5: GPUs, TPUs — Annotated Lecture Notes & Key Takeaways

Course Information

Institution: Stanford

Original Course: Stanford CS336 Language Modeling from Scratch | Spring 2026 | Lecture 5: GPUs, TPUs

Instructor Bio: This lecture is delivered by **Percy Liang**, Associate Professor of Computer Science at Stanford University and core faculty of the Stanford Institute for Human-Centered Artificial Intelligence (HAI). Percy Liang leads the Stanford Natural Language Processing Group, with research spanning the theoretical foundations and practical systems of language modeling, machine learning alignment, compositional semantics, and NLP evaluation. He received his PhD in Computer Science from the University of California, Berkeley and his BA in Mathematics from Harvard University. His work has been recognized with the NSF CAREER Award, Google Faculty Research Award, and multiple best paper awards at top-tier machine learning and NLP conferences. He is also widely known for creating influential benchmarks and open-source frameworks for language model research and assessment.

Course Description: This lecture introduces the hardware accelerators that enable large-scale language model training and inference. It breaks down the internal architecture of GPUs and TPUs, including their compute cores, memory hierarchies, interconnects, and specialized tensor processing primitives. The lecture explains how accelerator architectures are optimized for the matrix operations at the heart of transformers, compares performance characteristics across hardware generations, and discusses key hardware-level bottlenecks and practical considerations for running language model workloads at different scales.

The pitch at the top of Lecture 5 was refreshingly concrete: by the end of the hour you should be able to look at a benchmark plot of matrix-multiply throughput against matrix size — one of those jagged charts where the line falls off a cliff at certain dimensions for no visible reason — and account for every wiggle. That promise is the spine of the whole session. Everything before the payoff is setup: what a GPU actually is, why it behaves nothing like the CPU you're used to, and a short list of tricks that separates a kernel that saturates the hardware from one that leaves most of the chip sitting idle. The instructor (a student addresses him as Tatsu) keeps reminding the room that he is not a systems person, which turns out to be a feature — he explains the parts that confused him when he started.
 

Why the lecture is about compute, not clock speed

 

The bridge out of the neural-network half of the course is one sentence: compute is the currency of this field. Faster hardware, better utilization, more chips, smarter parallelization — all of it converts into better models. If this were the 1990s, the conversation would be about clock frequency. Dennard scaling, the regime where shrinking transistors also made them faster, more or less ran out in the 2000s. Transistor counts kept climbing, but smaller no longer meant quicker, for physical reasons he gestures at without deriving.
 

So the field went horizontal. Instead of one instruction stream running faster and faster, you throw many instruction streams at the problem at once. That's the GPU story in one line, and it's the same story as the parallelism lectures that follow this one. The flop chart he shows makes the point: K20s and M40s look almost quaint, and then around the P100/V100 generation the curve bends upward into something close to super-exponential year over year. Tensor cores arrive with the V100 in 2017; structured sparsity and low-precision formats push further after that.
 

He also opens by pointing at the reading list he borrowed from: Horace He's explainers, the GPU Mode community (a student interrupts to note the group used to be called CUDA Mode), and the Google-authored TPU book that has since grown into a GPU book, whose exercises apparently resemble what the class assignment will look like.
 

The hardware mental model: SMs, a memory ladder, and warps

 

The CPU/GPU contrast is drawn in terms of design philosophy rather than specs (7:34). A CPU is built for fast serial execution with messy branching logic, so it spends its silicon on big control units and a handful of arithmetic units, and it cares obsessively about the latency of a single instruction. A GPU cares about aggregate throughput. Individual tasks may take a long time to finish, and the scheduler freely swaps between them, but the total work completed per unit time is enormous because there are hundreds of lightweight compute units on the die.
 

The unit that matters is the streaming multiprocessor, or SM. He describes an A100-class part as having well over a hundred of them, each an independent compute unit with its own sub-components and a view of global memory. On the die diagram, the SMs are the green blocks, L2 sits in the middle, and each SM carries its own L1 — physically close to the compute, which is precisely why it's quick.
 

Then the memory ladder, which he correctly flags as the thing that will dominate the rest of the hour. Registers are fastest and most local. L1 and shared memory land in the tens of cycles. L2 is notably slower. Global memory — the high-bandwidth memory your framework reports when it says a card has 144 GB — is roughly an order of magnitude slower than L1. The physical split explains it: HBM is separate memory silicon sitting off the compute die, while the caches are on-die.
 

A student asks the obvious question: if shared memory is so much better, why not build the whole chip out of it? Because it's far more expensive, power-hungry, and has to sit physically close to be useful. He mentions an all-SRAM design from Groq as attracting attention for inference workloads, but says most accelerators will keep the hierarchy and force you to respect it. (More on that remark below — it's one of two things I'd verify before repeating.)
 

The follow-up question is sharper: what's the difference between shared memory and L1? The answer is about control. A cache decides on its own what to hold, based on recent access. Shared memory is a scratchpad you explicitly write into and read out of. That distinction is what makes tiling possible later.
 

The programming model comes as three nested objects. A thread is a lightweight unit of execution, and threads follow the single-instruction-multiple-thread rule: same instruction, different data, which is the tradeoff that buys efficiency at the cost of flexibility. A block is a group of threads guaranteed to land on one SM, which is why blocks can share scratchpad memory. A warp is 32 consecutive threads and is the scheduler's real unit of work — when he says a warp is executing, that's what he means. Asked whether the same-instruction rule applies to a whole block or just a warp, he says warp.
 

The memory model that goes with this is a list of places you can put things: registers, per-thread local memory, shared memory visible to the block, global memory, constant memory (which he admits he almost never sees used), and host memory if you need to spill past the device. The takeaway he repeats is blunt: once you leave shared memory, you pay, so structuring work to minimize global reads is the whole game.
 

The TPU detour is short and framed as convergent evolution. Build an energy-efficient ML accelerator and you end up in roughly the same place: a matrix-multiply unit, vector units, a control path, slow HBM, and fast local memory. The differences are sizing and flexibility — lighter control logic, much bigger matmul units. He claims the genuinely large divergence is in networking rather than the chips themselves, and then explicitly declines to cover it. There's a nice caution about vocabulary: on a TPU, "tensor core" means a processor; on a GPU it means the matmul unit. Same phrase, opposite referents. In rough numbers, an H100 has something like 132 SMs against a TPU's two, and hundreds of small matmul units against the TPU's eight. Which produces a consequence he clearly finds funny: a batch-size sweep in one of his own papers stops at 64 because the hardware refuses smaller matmul dimensions.
 

He closes the hardware half with three reasons GPUs won, and a warning. Scaling is easy — add SMs. Programming is deceptively approachable because it's SIMD, closer to a functional map than to hand-managing threads. Threads are cheap to start, stop, and swap, so stalled work gets covered by other work. The warning is a plot of growth rates: compute climbs steeply, memory bandwidth climbs slowly, interconnect climbs slowly. The gap widens every generation, so more and more of the optimization burden is memory and communication. In the Q&A he notes this is even more extreme for inference, mentioning prefill/decode disaggregation and an open-source model that routes attention and MLP layers to different accelerators entirely.
 

Six tricks for making the thing go fast

 

The organizing frame is the roofline model, which Percy had already introduced. Plot throughput against arithmetic intensity and you get a diagonal region where you're memory-bound and a flat region where the compute units are saturated. You want to live on the flat part, which means raising how much computation you do per byte moved.
 

Control divergence

 

The one trick that isn't about memory. Because threads in a warp share an instruction, a branch isn't like a CPU branch: both sides get executed, and the threads on the wrong side sit masked and idle while the other side runs. That's why GPU code tends to express things like ReLU as a multiplication by a mask rather than as a conditional.
 

Low precision, or the black art of choosing which operation gets how many bits

 

He gives this the most time of any single topic, and it's where the lecture gets genuinely interesting. A non-trivial share of that super-exponential flops curve comes from number formats: FP32 to BF16 to INT8, halving bits and therefore halving bytes moved.
 

But the real practice is not "cast everything down." In a modern low-precision matmul you downcast the inputs, accumulate in higher precision, and often emit results in FP32. The hard-won knowledge is per-operation: matmul inputs can go low; softmax and exponentials may need full precision; some of it you can get away with, some you can't. He describes it as years of slow empirical work.
 

At FP8 there's no longer one canonical format — you pick between variants like E4M3 and E5M2 depending on whether you need exponent range or mantissa precision. And with so few exponent bits you need scaling factors to keep values in range. The block-scaled formats (MXFP8) take this further with many small scale factors, one per block of 32 elements, themselves stored as power-of-two values in eight bits.
 

Then the detail that got the biggest reaction in the room: transposing a block-scaled matrix breaks the scaling pattern, so a transpose is no longer a cheap view — it may require re-quantizing. The workaround implementations use is to keep two quantized copies of every matrix, one in each orientation. I found this both absurd and delightful, and it's a good illustration of how far the format wars have pushed things.
 

The payoff is real but sobering: FP8 matmuls might save you 20–30%, not the 2x the bit math suggests, because quantization itself costs work. First and last layers are the hardest to quantize; for the last layer he offers the intuition that it feeds the loss directly, and freely admits he has no intuition for the first. MXFP4 exists, with so few representable values that they fit on one slide, and there's a paper on FP4 training, but he hasn't heard of anyone successfully training a serious large model in it. He expects the next generation to try.
 

On sparsity he's more dismissive: mixtures of experts are the one structured-sparse idea that clearly worked, and for structured matrices generally the compute gains versus representation loss has, in his phrasing, washed out. Asked for the state of the art in quantization, he lands on "train bigger, then quantize," mixing quantization-aware training with post-training methods — and says plainly that industry teams are still doing the science of quantization.
 

Fusion, recomputation, and coalesced reads

 

Fusion gets the factory analogy: a warehouse of memory, a compute factory, and a conveyor belt between them. Every separate operation is another round trip on the belt. Computing sin²x + cos²x naively in PyTorch gives you a chain of kernels, each reading and writing global memory. Fuse them and you read once, do everything inside the SM, and write once. Compilers like torch.compile and JAX's handle the easy cases for you.
 

Recomputation is the trick that feels wrong until you accept the economics. Storing activations for three stacked sigmoids costs eight memory operations across forward and backward. Throw the intermediates away and re-derive them during the backward pass and you're down to five. Identical math, extra flops, less traffic — a good trade only because flops are cheap relative to bytes.
 

Coalescing (53:05) is DRAM trivia with real consequences. Memory delivers data in bursts rather than single values; if the threads in a warp happen to request addresses inside the same burst, you get the rest essentially free. In a row-major matrix, threads walking along the major axis are not coalesced, and a traversal in the other direction can read an entire block in one go.
 

Tiling, the big one

 

Tiling is where the memory hierarchy stops being theory (58:07). Cut your operands into tiles, pull each tile into shared memory once, do all the reuse there, and write results back when the tile is finished. In a naive n-by-n matmul, every input element gets read n times from global memory. With tiles of size T, each element makes n/T global trips and T cheap shared-memory reads — a T-fold cut in expensive traffic.
 

The complications pile up fast. Tile size interacts with matrix shape: a 256×256 matrix cuts cleanly into 128×128 tiles, but add a single row and column and you spawn thin, mostly-empty tiles. Shared memory capacity, coalescing behavior, and problem shape all constrain the choice, which is why PyTorch's max-autotune flag spends what feels like forever benchmarking tile configurations. His practical rule for the rest of us: keep dimensions at powers of two and ideally divisible by 32, not because powers of two are magic but because that's what makes the reads line up.
 

Where the payoff lands: the plot, and flash attention

 

The matmul mystery gets dismantled in two steps. First, color the throughput curves by how divisible the matrix dimension is: the odd-sized cases sit far below the rest, divisibility by 2 is still poor, and by the time you reach divisibility by 16 or 32 the curves converge — not because those numbers are special, but because the tiles then align with the memory burst windows.
 

Second, the periodic collapses. With 256×128 tiles, a dimension of 1792 yields 98 tiles, which fits inside the A100's 108 SMs in one wave. Bump to 1793 and you get 120 tiles, so a dozen of them wait for a second wave while nearly the whole chip idles (1:09:20). That's wave quantization, and it explains why adding one row can cost you a large fraction of your throughput. He cites Karpathy's nanoGPT speedrun, where padding the vocabulary size by a few dozen entries produced something like a 25% speedup. Percy allegedly calls this class of knowledge "GPU trivia" and dislikes it; Tatsu is obviously fond of it.
 

Flash attention is the victory lap (1:11:53), and the point is that by now nothing in it should look like magic. The paper's own summary — tiling plus recomputation, done so that memory traffic grows sub-quadratically — is now vocabulary you own. The matmuls tile trivially; the paper's figure one is just a tiled matmul. The obstacle is the softmax, which is global and therefore seems to weld all the tiles together.
 

The unlock is the online softmax. Keep a running maximum and a running normalizer; each time a new maximum appears, rescale what you've accumulated and carry on. Now you can process tile by tile, holding only a couple of accumulators in shared memory or registers, apply the exponential per tile, keep the running sums, multiply by V in tiled fashion, and divide once at the end. In the backward pass you don't store the quadratic attention matrix at all — you recompute it tile by tile.
 

The closing minute is an anti-superstition plea: don't cargo-cult the "make it divisible by 32" rule, understand why the rule exists. Matmuls are the arithmetically dense operation worth designing around, the compute-versus-memory gap means data movement is where the engineering lives, and architecture design that ignores the hardware is leaving performance on the table.
 

Where the lecture left me hanging

 

A few things I'd want cleaned up before relying on them.
 

The SM count is inconsistent. Early on, the A100 is described as having around 128 SMs; the wave-quantization argument an hour later depends on 108. The second number is the one the argument needs, and since the whole point of that example is the arithmetic, the slip is more than cosmetic. I'd assume the early figure was misspoken, but in the room it's genuinely confusing.
 

The Groq remark — an all-SRAM design, described as recently acquired by Nvidia — is the kind of offhand parenthetical that drifts into lectures. It isn't load-bearing for anything else he says, but I'd verify it before repeating it.
 

Several admissions of ignorance are more interesting than the errors, because they map where the field actually is. Quantization is explicitly unsettled: he doesn't know why the first layer is hard to quantize, FP4 training exists only in a paper as far as he knows, and the gap between the theoretical savings of lower precision and the realized 20–30% is hand-waved as quantization overhead. He also declines to say how much low precision helps outside matmuls beyond "usually not worth it," deflects a question about systolic array dimensions, and hands the wafer-scale question to someone else in the room entirely. Structured sparsity gets written off as empirically washed out with no numbers attached.
 

The weakest structural choice is the TPU segment. Two slides, a concept-mapping table, and then the explicit statement that the real difference — networking — is out of scope. If you came for the TPUs in the title, you leave with a translation glossary, not a mental model. And the claim that tensor cores make matmuls more than ten times faster than any other floating-point operation is stated flatly, with no source and no caveat about which generation or dtype. It's load-bearing for his argument that future architectures will keep containing matmuls, so I'd have liked a citation.
 

Still, the promise at the top is kept. If you've followed the tiling and wave quantization sections, you can look at that jagged throughput plot and know exactly why the line falls off a cliff at 1793.
 

Content Disclaimer:

This article is for general reference only and does not constitute professional R&D guidance, production process advice or quality certification. All material performance data has specific test premises; readers should verify parameters against actual equipment and working conditions.