From Tiles to a Kernel
Introduction
Project introduction
From Tiles to a Kernel
You have written a softmax that cannot overflow, a running max and sum that are corrected as chunks arrive, a rule for merging two partial results, a tiled forward pass, the causal rule for skipping tiles, and an exact count of memory traffic. That is every piece FlashAttention needs.
In the notebook these pieces become one loop. Each tile of queries walks through the tiles of keys, merging each one into its running state, and divides once at the end. You check every version against PyTorch's own attention.
What you will see
First you will meet a version of online softmax that forgets to rescale the old sum, which looks right and passes easy tests. You will build the input that exposes it, and see how often it is wrong on real attention scores.
Then you will count how many numbers each version moves through main memory, compare with the numbers published for a real GPU, and write the core update inside a Triton kernel. Triton runs here in interpreter mode on the CPU, which checks that the kernel is correct; the speed claims come from the counts and the papers, not from timing.
Attention, one tile at a time
6 × 6 tilesqueries ↓ · keys →
This step
How the notebook is organized
01
Softmax in pieces
Safe and online softmax, the rescale bug, and merging partial results.
02
Tiles and traffic
Tiled and causal attention, then an exact count of reads and writes.
03
The kernel
The same update inside a Triton kernel, checked against PyTorch.
Write FlashAttention in the notebook
Build it step by step in the browser, about 50 minutes.