A model that does not fit on one GPU has a different problem from a model that fits but decodes slowly. Distributed strategies partition data, parameters, and states. FlashAttention improves attention memory traffic; KV caching avoids recomputing historical keys and values during decoding.
Long-context evaluation asks whether a model can process, retrieve, and use distant information. An SSM specifies sequence processing through state updates. SSM efficiency does not establish reliable document retrieval, and a large attention window does not guarantee relevant evidence is used.
Standard attention writes the N×N score matrix out to HBM and reads it back, and most of its time goes to that traffic. FlashAttention uses tiling plus softmax rescaling so each block finishes inside SRAM, and the backward pass recomputes instead of storing. Tri Dao's guest slides for 11-868 give one set of numbers: the backward pass does 13% more FLOPs, 9x less HBM traffic, and runs 6x faster. FA3 and FA4 follow the same theme: when the hardware changes, the bottleneck moves, and the algorithm has to move with it.
The last programming assignment in CS149 Fall 2025 is open-ended. Pick at least one of five kernels (Histogram, a 1D occupancy decoder, FlashAttention, a 3D heat equation with RK4, SwiGLU) and make it faster than its PyTorch baseline on an H100. You can write CUDA, Triton, or TileLang, and you may use LLMs. There is no speed threshold. The grade depends on a work log that shows what you measured at each step, what hypothesis you formed, and why you stopped. The H100 job queue and leaderboard need a SUNet ID; outside Stanford you can only run eval.py on your own NVIDIA GPU.
This post covers Fall 2026 material, not the Fall 2024 edition the rest of the series follows. Fall 2026 replaced the pruning lab with "Efficient AI Fundamentals" (lab1_gpu_basics.zip). Part 1 has you hand-write a triple-loop GEMM and compute MAC, FLOPs, and I/O. Part 2 plots GEMM and GEMV rooflines. Part 3 works through a gemma-3-270m-it decoder layer, computing attention and MLP costs and comparing prefill with decode. Part 4 uses the PyTorch Profiler to inspect kernels, has you write GeLU to feel kernel fusion, then tries torch.compile and CUDA Graphs. Part 5 compares SDPA with FlashAttention. The core is 80 points plus 20 bonus, and all of Part 5 became bonus because Colab's T4 can't run it.
In week 3 of ML 2026, Hung-yi Lee spends the first half of the inference lecture on one technique: Flash Attention. A GPU's execution units are fast, but their workbench (on-chip SRAM) is tiny, so data has to be carried to and from the warehouse (HBM). The carrying is the bottleneck. A naive softmax makes several round trips to the warehouse. Flash Attention assumes the current maximum is Amax, then multiplies by a correction factor when a larger value shows up. That lets it find the maximum, build the denominator, and compute the weighted sum in one pass, without ever materializing the attention weights. The output is identical to standard attention, no retraining is needed, and the cost is a little extra compute and a little brain strain.
HW3 is 20 multiple-choice questions at 0.5 points each. No code is submitted; students answer a quiz on NTU COOL. The first 10 questions come from reading papers: four on speculative decoding (Leviathan et al., DeepMind's Speculative Sampling, Inference with Reference, SpecInfer) plus FlashAttention 1–3. The last 10 require filling TODOs in the Colab and analyzing the results: acceptance rate of a hand-written speculative decoder, speed-up curves for an assistant model vs n-gram under two prompt regimes, HBM reads and theoretical FlashAttention speed-up from T4 specs, vLLM prefix caching across turns and a cache invalidation test, and the effect of CPU offload on throughput. All questions are printed in both Mandarin and English in the homework PDF, so outsiders can do the whole thing; they just cannot get the official answers.
Lecture 5 explains GPUs through SMs, warps, and the memory hierarchy, then unifies common optimization under low precision, fusion, recomputation, coalescing, and tiling. FlashAttention combines those principles for attention.