What is FlashAttention?
Introduction
Project introduction
What is FlashAttention?
Attention compares every query with every key. For a sequence of n tokens that is a table of n × n scores, then a softmax over each row of the table, then a weighted sum of the value vectors. At 4096 tokens the table already has about 16.8 million entries for each head.
FlashAttention, introduced by Tri Dao and colleagues in 2022, computes exactly the same result without ever storing that table. It works on small tiles of queries and keys, and keeps a running maximum and a running sum for each query so the softmax still comes out exact.
Why is it faster if it does the same math?
A GPU has a large main memory and a tiny, much faster memory on the chip itself. On an A100, main memory moves about 1.5 to 2.0 TB per second, while each of its 108 processors has only 192 KB of on-chip memory, and together those on-chip memories move about 19 TB per second, roughly ten times faster. Softmax does very little arithmetic per number, so standard attention spends most of its time writing the score table to main memory and reading it back.
Tiled attention keeps each tile of scores on the chip and never sends the table to main memory. The paper measured attention for GPT-2 medium on an A100, forward and backward pass together, moving 40.3 GB in the standard version and 4.4 GB with FlashAttention, and running in 7.3 ms instead of 41.7 ms. This project has no GPU, so you will not time anything; you will count the reads and writes exactly instead.
Attention, one tile at a time
6 × 6 tilesqueries ↓ · keys →
This step
What you will build
01
Softmax in pieces
Safe softmax, online softmax with a running max and sum, and merging two partial results.
02
Tiled attention
Attention computed tile by tile, with causal masking that skips whole tiles.
03
Why it is fast
An exact count of memory reads and writes for both versions, then a Triton kernel.
Start with a safe softmax
Turn scores into weights without overflowing, including fully masked rows.