FlashAttention from Scratch — The Paradox of Doing More Math to Go Faster

FlashAttention deliberately recomputes the same math in the backward pass. It does strictly more arithmetic and still runs faster — because on a GPU, moving numbers costs more than multiplying them. From counting HBM round trips, through the softmax wall that blocks tiling, to the online-softmax recurrence that breaks it, and the silent backend fallback that bites people in production.

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

An optimization that adds arithmetic

When you hear "make it faster," you usually picture cutting wasted work. FlashAttention does the opposite. In the backward pass it recomputes almost everything the forward pass already computed. The arithmetic count strictly goes up. And it is still faster than the naive implementation.

The trick fits in one line: on a modern GPU, fetching a number from memory costs more than multiplying it. That sets an exchange rate — if recomputing saves you a trip to memory, recomputing is the cheaper option. And FlashAttention is not an approximation. Nothing is dropped or sparsified; it produces the same answer as the naive version. What changed is the route to the answer, not the answer.

The metaphor: never write the score sheet down

Picture a library where the documents live in a basement warehouse and your desk is tiny. Trips to the warehouse are slow; work on the desk is instant.

The job: score 1,000 applicants on 1,000 criteria, then blend their opinions in proportion to the scores. The naive approach builds a million-cell score sheet, can't fit it on the desk, files it in the warehouse, fetches it back to compute row sums, converts to proportions, files it again, and fetches it a third time to blend the opinions.

FlashAttention never writes the sheet down. It splits the applicants into batches of 50, scores each batch on the desk, and keeps only a running partial answer sitting on the desk. When the last batch is done, the finished answer is already there.

The warehouse is the GPU's HBM (large, slow, off-chip); the desk is SRAM (small, fast, on-chip). How many orders of magnitude separate them is covered in The GPU Memory Hierarchy.

Recap: only the middle blows up to N×N

As shown in Attention from Scratch, self-attention is three lines. Let $N$ be the sequence length (number of tokens) and $d$ the head dimension.

$Q, K, V$ are all $N \times d$ matrices whose rows are each token's "question," "name tag," and "content." $S$ is the $N \times N$ score table pairing everyone with everyone, $P$ is that table rescaled so each row sums to one, and $O$ is the output ($N \times d$). So input and output are both $N \times d$, but the middle swells to $N \times N$. That is the root of everything.

The slow part is the traffic, not the math

Numbers help. With $N = 4096$ and $d = 128$, $Q$ holds about 520K elements while $S$ holds about 16.8M — 32× more. Quadruple the sequence length and $Q$ grows 4× while $S$ grows 16×, widening the ratio to 128×.

The naive implementation shuttles that giant table through HBM repeatedly: write $S$, read it back for the softmax, write $P$, read it again for the product with $V$. Bytes moved scale as $N^2$ while arithmetic scales as $N^2 d$, so you get only about $d$ operations per byte moved. That is squarely bandwidth-bound territory.

And both $S$ and $P$ are intermediates that get thrown away. We are hauling something larger than the final output to the warehouse and back, only to discard it. That is the opening.

The softmax wall that blocks tiling

"Cut the big matrix into blocks and process whatever fits in fast memory" — tiling has been standard practice for matrix multiplication for decades. So why wasn't attention doing it?

Because softmax gets in the way.

$$\mathrm{softmax}(s)_j = \frac{e^{s_j - m}}{\sum_{k} e^{s_k - m}}, \qquad m = \max_k s_k$$

Here $s_j$ is one cell of the score table and $m$ is that row's maximum. Subtracting the maximum before exponentiating is the standard guard against $e^x$ overflowing; it multiplies numerator and denominator by the same constant, so the answer is unchanged.

The trouble is $m$ and that sum in the denominator. Neither is final until you have seen the whole row. After one block, a larger score might still show up later. So implementations produced the entire table before normalizing, which meant holding $N \times N$ — or so it seemed for a long time.