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.

In words, attention scores every token against every other token to get an N×N table S. Each row of S is then turned into proportions that sum to one (P), and those proportions are used to blend the V rows into the output O. Input and output are N×d, but the middle table is N×N.

$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$$

In words, softmax subtracts the row's maximum from every score, exponentiates, and divides by the sum. Because the same constant scales numerator and denominator, the answer is unchanged, but e^x no longer overflows. Both m and the denominator are only final once the whole row has been seen.

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.

Online softmax: shrink the running total after the fact

The wall comes down with online softmax. The idea is simple: when the maximum gets updated, rescale everything you have accumulated so far, all at once.

Walk the row block by block. You carry exactly three things:

• $m$: the largest score seen so far
• $\ell$: the running sum of exponentials (the denominator so far)
• $\tilde{O}$: the running weighted sum (a $d$-dimensional vector, not yet divided)

From the new block, compute its local maximum $\tilde{m}$, local exponential sum $\tilde{\ell}$, and local weighted sum $\tilde{P}\tilde{V}$. Then update with three lines:

In words, when a new block arrives, take the larger of the old and new maxima. Then scale the old running totals by e^{m−m'} and the new block's by e^{m̃−m'} so both sit on the same footing, and add them. If the maximum did not change, the old totals pass through untouched.

In words: pick the new maximum, scale the old accumulation by $e^{m - m'}$ and the new block by $e^{\tilde{m} - m'}$ to put them on the same footing, then add. If the maximum didn't change, $e^{m-m'} = 1$ and the old total passes through untouched.

The crux is that the shrink factor $e^{m - m'}$ is a single scalar. You never revisit the thousands of terms you already added — one multiply against the accumulated vector covers all of them. That is why you can add from left to right without fearing a later update to the maximum. After the last block, divide once: $O = \tilde{O}/\ell$. The result is mathematically identical to the naive computation; no approximation enters anywhere.

The recurrence itself comes from the 2018 online-softmax paper. FlashAttention extended it to fold in the product with $V$ and redesigned it around the GPU's memory hierarchy.

In pseudocode

import numpy as np

def flash_attention(Q, K, V, Bc=128):        # Bc = block width that fits on the "desk"
    N, d = Q.shape
    O = np.zeros((N, d))                      # running weighted sum (unnormalized)
    m = np.full(N, -np.inf)                   # per-row maximum
    l = np.zeros(N)                           # per-row sum of exponentials

    for j in range(0, N, Bc):                 # bring K/V to the desk one block at a time
        Kj, Vj = K[j:j+Bc], V[j:j+Bc]
        S = Q @ Kj.T / np.sqrt(d)             # only N×Bc here; N×N never exists
        m_new = np.maximum(m, S.max(axis=1))
        alpha = np.exp(m - m_new)             # factor that shrinks the old accumulation
        P = np.exp(S - m_new[:, None])
        l = alpha * l + P.sum(axis=1)
        O = alpha[:, None] * O + P @ Vj       # shrink in bulk, then add
        m = m_new

    return O / l[:, None]                     # divide exactly once, at the end

Check that no N × N array appears anywhere. What lives in memory is an N × Bc block plus two scalars per row. A real CUDA kernel also tiles the $Q$ side and assigns each pair of blocks to one SM, so $S$ never leaves SRAM.

How much HBM traffic does this actually save?

The paper formalizes this as IO complexity. The naive implementation makes $\Theta(Nd + N^2)$ HBM accesses; FlashAttention makes $\Theta(N^2 d^2 M^{-1})$, where $M$ is the size of SRAM in elements.

So both scale as $N^2$, but FlashAttention carries a small coefficient $d^2/M$. With realistic values — $d$ of 64 to 128 and $M$ on the order of 100 KB — $d^2$ is far smaller than $M$, and the traffic drops by orders of magnitude. Same asymptotic order, coefficient worth orders of magnitude: a very typical shape for a real optimization.

Recomputing on purpose in the backward pass

Here is the paradox from the opening. The backward pass needs $P$. The naive implementation saves it to HBM during the forward pass and reads it back. FlashAttention doesn't save it. It saves the output $O$ and the per-row $m$ and $\ell$ — quantities that scale with $N$, not $N \times N$.

During the backward pass it rebuilds the needed blocks of $P$ on the fly in SRAM from $Q, K, V, m, \ell$. That adds the entire cost of recomputing $S$. The whole thing still comes out ahead, because what was added is desk work and what was removed is warehouse trips. In a bandwidth-bound regime, FLOPs are not the currency. Bytes are.

The paper reports figures such as roughly 3× on GPT-2 training (sequence length 1K) and a 15% wall-clock reduction on BERT-large (sequence length 512) against the MLPerf 1.1 training speed record. How much you gain depends heavily on sequence length, head dimension, and GPU generation, so measure on your own workload.

Later versions, and what changes at inference time

FlashAttention-2 (2023) cut the non-matmul work, extended parallelism along the sequence-length axis, and reduced shared-memory traffic between warps. FlashAttention-3 (2024) exploits Hopper-era asynchrony — overlapping data movement with matrix multiplies — and adds techniques for holding accuracy at low precision such as FP8.

Decoding changes the picture again. During generation the query is a single token, so parallelizing along sequence length no longer works and the cores go idle. The fix is to split the $K, V$ side, compute partial results in parallel, and combine them at the end. The combining step is exactly the three equations above — bring each partition's $m, \ell, \tilde{O}$ together, put them on a common footing, and add; splitting the work does not change the answer.

One clarification worth internalizing: what shrinks here is the score matrix, not the KV cache. When you run out of memory on long contexts, the culprit is usually the KV cache, and The KV Cache from Scratch covers how to size it. People conflate these two constantly.

How this shows up in practice

Who touches it, and when. Engineers running LLM training or inference infrastructure, teams shipping long-context products, and researchers implementing custom mask shapes or positional biases. The first two are "users of the kernel," the last is "the person whose change makes the fast kernel ineligible" — different failure modes.

Names you will actually type. In PyTorch the entry point is torch.nn.functional.scaled_dot_product_attention (SDPA), which dispatches among MATH, FLASH_ATTENTION, EFFICIENT_ATTENTION, and CUDNN_ATTENTION backends. To pin and verify which one runs, use torch.nn.attention.sdpa_kernel. More directly there is the flash-attn package with flash_attn_func and, for ragged batches, flash_attn_varlen_func. Serving stacks like vLLM and TensorRT-LLM call into these under the hood.

Pitfalls that turn into incidents.

• Passing a mask silently falls back. An arbitrary additive mask (a tensor in attn_mask) tends to disqualify the flash backend, and dispatch drops to MATH. Nothing raises — you are back to materializing $N \times N$, so you find out via a several-fold slowdown or an OOM. Express causal masking with is_causal=True, pin the backend with sdpa_kernel, and make the run fail loudly if the fast path is unavailable.
• Dtype and shape constraints. Flash kernels assume fp16/bf16 and do not run in fp32. Head dimension has an upper bound that varies by version and GPU generation. Debugging in fp32 produces the baffling "only production is fast" gap.
• Bitwise mismatch is not a bug. Addition order changes, so exact equality with the naive implementation is not something to expect. Regression tests written with tight tolerances will fail on a perfectly correct migration; compare relative error instead.
• Benchmarking short sequences shows nothing. The traffic being saved is small there, so a sequence length of 512 hides the effect. Measure at production sequence length, batch size, and head count.
• Don't feed padded batches. Padding variable-length sequences up to the longest one means paying for that compute and traffic in full. The variable-length API (passing boundaries via cu_seqlens) removes it.

What design reviews ask. "Why is it faster when it does more FLOPs?" — because the regime is bandwidth-bound, the currency is bytes, and recomputation in SRAM is cheaper than a round trip to HBM. "What does the forward pass save?" — the output $O$ plus per-row $m$ and $\ell$, a quantity that scales with $N$. Answering those two in your own words is the difference between "we use it because it's fast" and "we use it because we know why it works."

Summary

• FlashAttention is exact, not approximate: same answer as the naive version, without ever building the $N \times N$ table
• The softmax wall that blocked tiling falls to online softmax — a recurrence that shrinks the old accumulation with a single scalar before adding
• The backward pass rebuilds $P$ instead of storing it. More arithmetic, but in a bandwidth-bound world bytes are the expensive thing
• Production incidents come from silent fallback, not from the kernel itself. Pin the backend and benchmark at production scale