Lab 01 — Flash Attention (Tiled)

Phase: 08 — Inference Optimization at Depth | Difficulty: ⭐⭐⭐⭐⭐ | Time: 4–5 hours

Flash Attention (Dao et al., 2022) reduces attention memory from O(N²) to O(N) using online softmax tiling.

What you build

  • online_softmax — running max/sum trick for numerically stable softmax without materializing the full N×N matrix
  • flash_attention_forward — tiled attention: process Q, K, V in blocks; update running softmax statistics
  • Benchmark memory usage: standard O(N²) vs Flash O(N)
  • Verify numerical equivalence with PyTorch reference attention

Key concepts

ConceptWhat to understand
Online softmaxm_new = max(m, row_max); exp_sum update — avoids full row storage
SRAM tilingLoad block of K, V to fast SRAM; iterate over Q blocks
Memory savingsStandard: Q+K+V+attn_matrix = O(N²); Flash: only O(N) activations
RecomputationFlash recomputes attention during backward to save memory

Files

FilePurpose
lab.pyStubs with TODO markers
solution.pyComplete working implementation
test_lab.pypytest suite
requirements.txtDependencies

Run

pip install -r requirements.txt
python lab.py
pytest test_lab.py -v