Skip to content

FlashAttention — A Systems / Kernel Course

FSKC

Deep Dive · Compiler Track

Specialized course identity for FlashAttention — A Systems / Kernel Course.

Artifact: compiler/inference optimization · Measure: op count, memory, latency

Parent: 02 — Kernel Engineering

Format: 10 lectures. Theory → repo / code reading → small lab. Every lecture ships an artifact (notebook, benchmark script, kernel sketch, correctness test, or patch).

Why this course exists. Most "FlashAttention" material online stops at "tiling makes it faster." That is not enough to ship attention kernels. By the end of this course you can read the FlashAttention repository, explain the IO math, debug correctness, benchmark kernels, and modify or integrate attention paths for real LLM inference and training.

What you will be able to do at the end

  • Derive the IO complexity of standard attention vs FlashAttention and predict where each one is compute- or memory-bound on a given GPU.
  • Implement the online-softmax recurrence in plain NumPy and prove numerical equivalence to a one-shot softmax over the full attention matrix.
  • Read flash-attention/csrc and trace a flash_attn_func(...) call from Python into the launched CUDA kernel.
  • Explain the warp / block / sequence / head partitioning differences between FA1, FA2, and FA3.
  • Stand up a correctness harness that compares your kernel against PyTorch's reference SDPA with controlled tolerance bands.
  • Patch one focused path (e.g. a varlen mask, a small numerics fix, an alternative dispatch) and produce a benchmark + report.

Prerequisites

  • Phase 4 Track B — Jetson, CUDA basics, TensorRT.
  • Phase 4 Track C Unit 01 — Graph and operator optimization.
  • Comfort with linear algebra at the level of multi-head attention, softmax, and matrix multiply.
  • A workstation with at least one NVIDIA GPU, CUDA 12.x, PyTorch built against that CUDA, and a clone of Dao-AILab/flash-attention.

If you do not have a GPU locally, a single L4 / A10 / 3090 in a cloud instance is enough for lectures 1–7. Lectures 8–10 benefit from H100/H200 access for the FA3 / FP8 lab.

Syllabus

# Lecture Lab artifact
1 Attention bottleneck + roofline Notebook: per-shape arithmetic intensity table; roofline plot for one GPU
2 Online softmax + numerical correctness NumPy script proving blockwise-LSE equals one-shot softmax bit-for-bit in fp64; tolerance plot for fp16 / bf16
3 FlashAttention-1 algorithm Pseudocode walkthrough + an HBM-byte counter that matches the FA1 paper's IO complexity bound
4 GPU kernel performance basics Nsight Compute capture of a coalesced vs uncoalesced copy kernel; warp-reduction microbenchmark
5 Repo anatomy + Python / CUDA API Annotated flash_attn_func call trace from Python to launched kernel; minimal setup.py-style build recipe
6 FlashAttention-2 Benchmark table: FA1 vs FA2 vs PyTorch SDPA across seqlen / head-dim sweep
7 Backward pass + validation Correctness harness that compares dQ/dK/dV against a reference with documented tolerances
8 Inference path (KV cache, decode) Microbenchmark of decode-step latency with and without paged KV; RoPE / GQA sanity check
9 Hopper / FA3 / FA4 FA3 vs FA2 benchmark on H100 / H200; identify the WGMMA + TMA + warp-specialization paths in the source
10 Capstone One focused kernel-path change + benchmark + correctness + write-up

How to use this course

  • Do the lectures in order. Each one builds primitives the next assumes.
  • Read the listed source files inside the repo before reading the lecture's "Build it" section. The lectures are tour guides, not replacements.
  • Keep every lab artifact in a single flash-attn-course/ working directory so you have a personal benchmark archive at the end.
  • Treat correctness as a first-class output: every benchmark must have a paired correctness test.

Core sources

Role mapping

  • MTS Kernels / DL Inference Optimization Engineer — direct skill match. The capstone artifact is in the form a hiring manager can read.
  • HPC / Distributed AI Engineer — needed for understanding the kernel layer that NCCL and friends sit on top of.
  • LLM Inference Engineer (vLLM, TensorRT-LLM, SGLang) — Lectures 8–10 map directly to the decode-path optimizations these stacks ship.

What this course is not

  • Not a CUDA-from-scratch course. Lecture 4 reviews the strict minimum; everything deeper is referenced to Track B and to NVIDIA's own docs.
  • Not a transformer-architecture course. We assume you know what Q, K, V, and softmax(QKᵀ/√d)·V are; we focus on how to make that line of math run fast on real hardware.
  • Not exhaustive. We intentionally skip FlashAttention forks for non-NVIDIA hardware and quantized-only variants. After this course you should be able to read those repos on your own.