1 episodes
Paper: FlashAttention
FlashAttention (Dao et al. 2022) computes exact attention while never writing the N x N score matrix to GPU main memory. The module derives the online softmax (Milakov & Gimelshein 2018), shows that partial softmax states merge associatively, builds the tiled forward and recomputing backward passes in NumPy (tested equal to naive attention), counts memory traffic to probe the paper's IO claim, and ports the kernel to Triton (GPU) and WGSL (browser).
Progress — loading…
Objectives
Language
- Explain why long contexts make attention's memory traffic, not its arithmetic, the bottleneck
- State what the paper demonstrates (exactness, measured speed-ups on 2022 hardware) versus proves (an IO bound in a cost model)
Math & data structures
- Online softmax: running max, running sum, rescaling; associativity and the identity element (a monoid)
- IO complexity O(N^2 d^2 / M) vs O(Nd + N^2); block sizes from SRAM size
- Backward pass via recomputation and the rowsum(dP*P) = rowsum(dO*O) identity
Code
- nlpcore.kernels.flash_np (forward, backward, IO counter) and nlpcore.kernels.flash_triton (GPU, unverified)
- Hardened masking (padding, empty rows), a WGSL workgroup kernel with a CPU mirror, five-bug break-it
Episodes
Dao 2022; memory hierarchy; tiling; online softmax derivation.
Exercises
| Exercise | Tier | Ladder | Time |
|---|---|---|---|
| Online softmax, tiled exact attention, and FlashAttention's IO cost model in NumPy | browser | Predict · Scratch · Compose · Probe a claim | 90 min |
| FlashAttention backward with recomputation: masks, empty rows, ragged blocks and fp16 | browser | Harden · Scratch | 75 min |
| FlashAttention forward in Triton: port the NumPy tiles to a GPU kernel and benchmark it | gpu | Port · Compose | 120 min |
| Systems lane: tiled exact attention as a WebGPU compute shader, with a CPU mirror | systems · wgsl | Scratch · Harden | 120 min |
| Break it: five one-line bugs in an online-softmax kernel, and the tests that miss them | browser | Predict · Break it | 50 min |
Labs
- FlashAttention: Exact Attention That Respects the Memory Hierarchy browser F08.01 · 60 min
Library contributions
Files in nlpcore this module teaches you to write.
Sources
- Dao, Fu, Ermon, Rudra, Ré. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, 2022 · arXiv:2205.14135
- Dao. FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning, 2023 · arXiv:2307.08691
- Curriculum Cost Models: Big-O, Recurrences, and Why Vectorization Wins Cost Models (memory hierarchy, roofline)
- Milakov & Gimelshein (NVIDIA). Online normalizer calculation for softmax, 2018 · arXiv:1805.02867
- PyTorch documentation, torch.nn.functional.scaled_dot_product_attention (fused backends) · link
- Rabe & Staats. Self-attention Does Not Need O(n^2) Memory, 2021 · arXiv:2112.05682
- WebGPU specification: maxComputeWorkgroupStorageSize default limit (16384 bytes) · link