FlashAttention, Explained

Why attention was slow for the wrong reason, how tiling and online softmax fixed it without approximating anything, and the full maths worked out with small numbers. No CUDA required.

In 2020, GPT-3 could read about 2,000 tokens at a time. Roughly 1,500 words. Four pages of a book.

Today you can paste an entire novel into a chat window and ask questions about chapter forty. The models got bigger, sure, and the GPUs got faster. But neither of those explains a 500× jump in how much text a model can hold in its head at once. Something in the middle had to change.

That something is a GPU kernel published in May 2022 by Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra and Christopher Ré. It is called FlashAttention. It does not change the transformer architecture. It does not approximate anything. It produces the same answer the old code produced, and it is two to four times faster while using a fraction of the memory.

The reason it works is the interesting part, and it is not what most people assume. Attention was not slow because it did too much arithmetic. It was slow because it kept putting things down and picking them up again.

This article walks the whole thing: the problem, the hardware fact that everyone had been ignoring, the maths trick at the centre, the full algorithm, and the four versions that followed as the chips kept changing underneath it. Where there is maths, it is worked out with actual numbers you can check on a calculator.

Part 1 - The Problem

1. The 4× problem nobody could solve

Attention has a cost that scales with the square of the input length.

Double the context from 2,000 to 4,000 tokens and the cost does not double. It quadruples. Go from 2,000 to 100,000 and the cost goes up by a factor of 2,500. This is the single sentence behind every “context window” limit you have ever hit.

Between 2019 and 2022 this was one of the most attacked problems in machine learning. Dozens of papers proposed ways around it, and they nearly all took the same approach: if the exact computation is too expensive, do an approximate one. Attend to a sliding window of nearby tokens. Hash tokens into buckets and only compare within a bucket. Project the whole thing down to a lower rank. Replace softmax with a kernel function that factorises nicely.

Many of these methods genuinely reduced the number of arithmetic operations, often dramatically. And yet, in practice, models trained with them were rarely faster on the clock, and often slightly worse at their jobs. Something was clearly wrong with the diagnosis.

FlashAttention’s contribution begins with a different question. Instead of asking “how do we do less maths?”, it asks: when attention runs on a real GPU, where does the time actually go?

The answer turned out to be: not in the maths.

2. A mini dictionary before we start

Nine terms carry the rest of this article. If a sentence later stops making sense, the missing piece is almost always one of these.

Two of them do most of the work: HBM and SRAM. Every GPU has both. HBM is the big memory advertised on the box — the “80 GB” of an A100 or H100. SRAM is a tiny scratchpad built into the chip itself, next to the units that do the arithmetic. HBM is large and comparatively slow. SRAM is minuscule and extremely fast.

Hold on to that distinction. The entire article is a story about the gap between them.

3. What attention actually does

Before we can talk about making attention fast, we need to be precise about what it computes.

Every token in a sequence produces three vectors:

• a query (Q), which represents what this token is looking for,
• a key (K), which represents what this token offers to others,
• a value (V), which is the information this token actually passes along.

The library metaphor is worn out but accurate. You walk in with a question (your query). Every book has a spine label (its key). You compare your question to every spine, decide how relevant each book is, and then read from the relevant books in proportion to their relevance (their values).

The formula is:

Attention(Q, K, V) = softmax( Q Kᵀ / √d ) V

Read left to right, that is four steps:

1. Q Kᵀ — compare every query with every key. This produces an N × N matrix of scores, where N is the sequence length. Entry (i, j) says how much token i cares about token j.
2. ÷ √d — divide by the square root of the head dimension. Without this the dot products grow with dimension and push softmax into an extreme, nearly one-hot distribution where gradients vanish.
3. softmax — turn each row of scores into weights that are positive and sum to 1.
4. × V — take a weighted average of the value vectors using those weights.

Here is a single row of that, with numbers.

Figure 1 — One query attends over four keys. Scores become weights; weights mix the values.

That is the entire operation. There is nothing conceptually hard about it. The difficulty is entirely in step 1 and the size of the thing it creates.

4. The quadratic problem, measured

When people say attention is quadratic, this is what they mean concretely: Q Kᵀ produces a matrix with N rows and N columns. For 1,000 tokens that is a million numbers. For 100,000 tokens it is ten billion.

And a standard implementation does not just produce that matrix. It stores it in HBM, reads it back, applies softmax, stores the result, and reads that back again. Each of those copies is a full N × N matrix.

Put real numbers on it.

Look at the fourth column. That is one layer, of one model, for one request, at 16-bit precision. A model has dozens of layers. Training needs to keep these around for the backward pass. And none of this includes the model’s own weights, the gradients, or the optimiser state, which are what the GPU memory was supposed to be for.

This is why context windows were small. Not because nobody wanted longer ones. Because the attention matrix physically did not fit.

Part 2 - The Insight

5. The plot twist: it was never about the maths

Here is the fact that the approximate-attention papers had underweighted.

A modern GPU can do arithmetic far faster than it can fetch the numbers to do arithmetic on. Not slightly faster. Enormously faster. An A100 can perform roughly 312 trillion 16-bit floating point operations per second. In that same second it can only pull about 1.5 terabytes out of its main memory.

Work through what that ratio implies. If each number is 2 bytes, 1.5 TB/s is about 750 billion numbers per second delivered, against 312 trillion operations per second available. The chip can do roughly 400 operations in the time it takes to fetch one number. If your algorithm does fewer than a few hundred operations per number fetched, the arithmetic units spend most of their life idle, waiting for data to arrive.

The gap gets worse the closer you look, because there is more than one kind of memory.

Figure 2 — The two memories that matter, with A100 numbers from the FlashAttention paper.

SRAM delivers data about 13× faster than HBM. It is also about 2,000× smaller: 192 KB per streaming multiprocessor, around 20 MB across the whole chip. You cannot put an attention matrix in there. You can put a small piece of one.

That is the hinge the whole algorithm turns on.

6. A two-minute GPU primer: SMs, warps and why kernels exist

The rest of this article talks about warps, thread blocks and shared memory. If those are unfamiliar, here is the minimum needed, and it is genuinely the minimum.

A GPU is not one processor. An A100 contains 108 streaming multiprocessors (SMs), and an H100 has 132. Each SM is an independent little machine with its own arithmetic units, its own registers and its own private scratchpad. That scratchpad is the SRAM from Figure 2: 192 KB per SM on an A100.

Work is handed to the GPU as a kernel, which is a program plus a grid of thread blocks to run it on. Each thread block is assigned to one SM and stays there. Threads within a block can see the same shared memory and can synchronise with each other. Threads in different blocks cannot — they are separate islands, and the only thing they share is HBM.

Inside a block, threads are grouped into warps of 32. A warp is the real unit of execution: all 32 threads run the same instruction at the same time on different data. When you read “the warps are re-partitioned” in the FlashAttention-2 section, this is what is being rearranged — which warp is responsible for which slice of the tile.

Three more facts, each of which explains a design decision later:

Registers are faster than shared memory, which is faster than HBM. A thread’s registers are the fastest storage on the chip. Keeping a block of the output in registers, rather than writing it to shared memory between steps, is worth real time. This is part of why FlashAttention-2 swapped the loop order.

Tensor cores are separate units with their own instructions. They do one thing — small matrix multiply-accumulates — extremely fast. Everything else (exponentials, divisions, comparisons, masking) runs on the general-purpose CUDA cores or the special-function units. On an A100, a non-matmul floating-point operation costs roughly 16× a matmul one. Shifting work off the tensor cores is expensive; shifting work onto them is cheap.

Occupancy matters. If your kernel only creates 16 thread blocks, 92 of the 108 SMs have nothing to do. A perfectly efficient kernel that only fills a seventh of the chip is a seventh as fast as it looks. This is the problem FlashAttention-2 fixed by parallelising along the sequence.

That is the whole vocabulary. SM, thread block, warp, register, shared memory, tensor core, occupancy.

7. Arithmetic intensity, or why your GPU spends its life waiting

The standard way to formalise this is arithmetic intensity: how many arithmetic operations an operation performs per byte it reads from memory.

• A large matrix multiplication reads two matrices and does an enormous number of multiply-accumulates with them. High intensity. The chip is busy.
• Adding two tensors reads two numbers and does one addition. Intensity near zero. The chip fetches, adds, fetches, adds, and is idle almost the entire time.

Every operation sits somewhere on this line, and the line has a knee in it, usually called the ridge point: the intensity at which the chip’s compute capability and its memory bandwidth are exactly balanced. For an A100 at 312 TFLOPs/s and 1.5 TB/s, that knee sits at roughly 208 operations per byte.

Figure 3 — Below the ridge point, the tensor cores wait. Above it, they work.

Here is what this means for attention specifically. The Q Kᵀ and PV steps are matrix multiplications, and they are compute-bound and efficient. But everything between and around them — scaling by √d, applying the causal mask, computing softmax, applying dropout — touches all N² entries and does a handful of operations on each.

Those steps are wildly memory-bound. And in a naive implementation they dominate the runtime, even though they account for a tiny fraction of the FLOPs.

This is the diagnosis the approximate-attention papers missed. They were optimising the numerator of a fraction whose denominator was the actual problem.

8. What standard attention does on real hardware

To make it concrete, here is what actually happens when you write the attention formula in PyTorch and run it.

Each operation in the formula becomes a separate GPU kernel. A kernel is a program that runs on the GPU. Kernels cannot hand results to each other directly — one kernel’s output goes to HBM, and the next kernel reads it back from HBM. That is the only channel they have.

Figure 4 — Eight steps, four kernel launches, and the N × N matrix crossing the memory bus four times.

Read steps 3, 4, 6 and 7 again. The full N × N score matrix is written out to HBM and read back, twice, purely because the next kernel needs somewhere to find it. Nobody wants those numbers. They are scaffolding. But the scaffolding is the size of a small dataset and it is being dragged across the slowest link in the system.

When you measure a real workload, the picture is stark.

That table is the whole article in five rows, so it is worth pausing on the first one. FlashAttention does more arithmetic than the standard implementation. About 13% more. And it finishes in a sixth of the time, because it moves a ninth as much data.

Once you have internalised that trade, the rest of the design follows almost inevitably. The goal is no longer “compute attention with fewer operations”. It is: compute attention while touching HBM as few times as possible.

Ideally, once.

Part 3 - The Maths, Slowly

The goal is now clear: load a piece of Q, K and V into SRAM, compute everything we can with it, and write out only the final answer. Never let the N × N matrix touch HBM.

There is exactly one thing standing in the way, and it is softmax.

9. Softmax, and the problem with softmax

Softmax turns a row of arbitrary numbers into a set of weights that are all positive and add up to 1. The textbook definition is:

softmax(x)ᵢ = exp(xᵢ) / Σⱼ exp(xⱼ)

Two things about this formula matter enormously.

First, it has a denominator that depends on the whole row. You cannot finish computing the weight for element 1 until you have seen element 8,000, because the sum in the denominator includes it. Softmax is a global operation over the row.

Second, computed literally, it explodes. exp grows terrifyingly fast. A 16-bit float tops out at 65,504, and exp(12) is already 162,754. Attention scores routinely reach that range. The result is inf, and once you divide inf by inf you get NaN, and a NaN propagates through the model like ink in water. Your loss becomes NaN and your training run is over.

The standard fix, used everywhere, is to subtract the row maximum first:

Figure 5 — Subtracting the maximum changes nothing mathematically and everything numerically.

Why is this allowed? Because multiplying the top and bottom of a fraction by the same constant does not change it:

exp(xᵢ − m) / Σⱼ exp(xⱼ − m)
  = [ exp(xᵢ) · exp(−m) ] / [ exp(−m) · Σⱼ exp(xⱼ) ]
  = exp(xᵢ) / Σⱼ exp(xⱼ)

The exp(−m) cancels. The answer is identical. But now every exponent is zero or negative, so every exponential lands between 0 and 1 and nothing overflows.

This is called safe softmax, and essentially every implementation uses it.

But notice what it costs us conceptually. We have just added a second global dependency. Before you can exponentiate anything, you need the maximum of the entire row. Before you can normalise anything, you need the sum over the entire row.

So a block-based algorithm appears to be impossible. If you only have four of the eight scores in SRAM, you do not know the maximum, you do not know the sum, and therefore you cannot compute a single final weight.

That is the knot. The next section unties it.

10. Online softmax: the trick the whole algorithm rests on

The insight is this: do not wait for the true maximum. Use the best maximum you have so far, and correct yourself when a better one turns up.

Keep two running numbers per row:

• m — the largest score seen so far,
• ℓ — the running sum of exponentials, measured relative to m.

Now suppose a new block arrives containing a bigger score. Everything you accumulated so far was computed relative to the old, smaller maximum. Is it ruined?

No. It is off by exactly one known factor. If every earlier term was exp(xᵢ − m_old) and it should have been exp(xᵢ − m_new), then:

exp(xᵢ − m_new) = exp(xᵢ − m_old) · exp(m_old − m_new)

Every single earlier term is wrong by the same multiplier, exp(m_old − m_new). Call it α. One multiplication fixes all of them at once, no matter how many there were.

That is the entire idea. Here it is written out properly:

Figure 6 — Four lines of bookkeeping that make an exact softmax computable block by block.

Two details in that figure deserve attention.

α is always ≤ 1. Since m_new ≥ m_old, the exponent m_old − m_new is zero or negative. So the correction shrinks old contributions, never inflates them. Nothing overflows during the correction, which is the whole reason we subtract maxima in the first place.

The division happens once, at the very end. During the loop we accumulate an unnormalised output, written Õ, alongside the running sum ℓ. Only after the last block do we compute O = Õ / ℓ. Dividing inside the loop would be both slower and harder to keep exact.

One more thing worth saying plainly, because it is the reason FlashAttention won and its rivals did not: this is not an approximation. The rescaling is an algebraic identity. There is no error term, no bound, no “close enough for training”. If you run the blocked version and the ordinary version in exact arithmetic, they agree completely.

Which is easier to believe once you have watched it happen.

11. A worked example with real numbers

Take eight scores and eight corresponding values, and process them four at a time. The second block deliberately contains a larger score than the first, because that is the awkward case.

Scores: 1, 3, 2, 5, 4, 7, 6, 2 Values: 10, 20, 30, 40, 50, 60, 70, 80

Figure 7 — The same weighted sum, computed in two halves, agreeing to every digit shown.

Follow it through by hand.

Block 1 contains scores 1, 3, 2, 5. The largest is 5, so m₁ = 5. Exponentiating relative to it:

exp(1−5)=0.0183   exp(3−5)=0.1353   exp(2−5)=0.0498   exp(5−5)=1.0000
ℓ₁ = 0.0183 + 0.1353 + 0.0498 + 1.0000 = 1.2034
Õ₁ = 0.0183·10 + 0.1353·20 + 0.0498·30 + 1.0000·40 = 44.383

If we stopped here we would report 44.383 / 1.2034 = 36.88. That is the correct answer to a different question — the softmax over only the first four scores.

Block 2 contains 4, 7, 6, 2. Its maximum is 7, which is larger than our running maximum of 5. So the running maximum becomes 7, and the correction factor is:

α = exp(5 − 7) = 0.1353

Compute the new block relative to 7:

exp(4−7)=0.0498   exp(7−7)=1.0000   exp(6−7)=0.3679   exp(2−7)=0.0067
block sum = 1.4244
block × V = 0.0498·50 + 1.0000·60 + 0.3679·70 + 0.0067·80 = 88.780

Merge. Scale the old accumulators by α and add the new ones:

ℓ = 0.1353 × 1.2034 + 1.4244 = 0.1629 + 1.4244 = 1.5873
Õ = 0.1353 × 44.383  + 88.780 = 6.006  + 88.780 = 94.786

Normalise once, at the end:

O = 94.786 / 1.5873 = 59.717

Now compute it the ordinary way. Softmax over all eight scores with maximum 7 gives weights 0.0016, 0.0115, 0.0042, 0.0853, 0.0314, 0.6300, 0.2318, 0.0042, and the weighted sum of the values is 59.717.

Identical. And we never held more than four scores at a time.

That is FlashAttention in miniature. Everything that follows is the engineering required to do this at scale, on a GPU, for thousands of rows at once.

12. Tiling: doing this for the whole matrix

Scaling up means cutting both dimensions of the score matrix into blocks.

Split the queries into blocks Q₁ … Q_t, and the keys and values into blocks K₁ … K_e and V₁ … V_e. The block sizes are not chosen for mathematical reasons. They are chosen so that a block of Q, a block of K, a block of V and a block of the output all fit together inside one streaming multiprocessor’s 192 KB of SRAM. In practice this means 64 or 128 rows at a time.

Now the algorithm walks the grid.

Figure 8 — One tile of the score matrix exists at a time, inside SRAM, and is destroyed before the next one is made.

For each block of keys and values, the kernel sweeps down every block of queries. For each (query block, key block) pair, it computes that tile of the score matrix directly in SRAM, does the softmax bookkeeping on it, folds the result into that query block’s running output, and throws the tile away.

The tile never receives an address in HBM. It is computed, used and discarded within microseconds, entirely inside the chip.

The only things that travel back and forth are Q, K, V, the output O, and the two running statistics per row. For a sequence of 100,000 tokens that is 200,000 extra numbers instead of ten billion.

This is why the memory footprint collapses from quadratic to linear. Not because the N × N matrix has been made smaller or sparser or approximated, but because it is never materialised in the first place.

13. The full forward algorithm

Here it is in its entirety. Read the middle column as pseudocode and the right column as the reason each line exists.

Figure 9 — The complete FlashAttention forward pass.

Written out as plain pseudocode:

# Q is N×d, split into t blocks of rows
# K, V are N×d, split into e blocks of rows

initialise, for every row:
    O = 0            # running output, unnormalised
    ℓ = 0            # running sum of exponentials
    m = −∞           # running maximum
for j = 1 … e:                          # outer loop over key/value blocks
    load K_j, V_j from HBM into SRAM
    for i = 1 … t:                      # inner loop over query blocks
        load Q_i, O_i, m_i, ℓ_i into SRAM
        S      = Q_i K_jᵀ / √d          # this tile of scores, in SRAM only
        m_tile = rowmax(S)
        P_tile = exp(S − m_tile)
        ℓ_tile = rowsum(P_tile)
        m_new  = max(m_i, m_tile)       # merge, per Figure 6
        α      = exp(m_i    − m_new)
        β      = exp(m_tile − m_new)
        ℓ_i    = α·ℓ_i + β·ℓ_tile
        O_i    = α·O_i + β·(P_tile V_j)
        m_i    = m_new
        write O_i, m_i, ℓ_i back to HBM
O = O / ℓ                               # one division, at the very end
save m and ℓ for the backward pass

Three observations about this code.

It is one kernel. Everything above happens inside a single GPU program. There are no intermediate results handed between kernels, which means no intermediate results parked in HBM. This is called kernel fusion, and it is the mechanism by which the eight steps of Figure 4 become one.

Each K,V block is read from HBM exactly once for the entire pass. Each Q block is read once per key block, which is where the residual HBM traffic comes from.

The matrix multiplications still run on the tensor cores. Nothing here gives up the hardware’s fast path. Q_i K_jᵀ and P_tile V_j are ordinary matrix multiplies on small matrices, exactly what the tensor cores are built for. The softmax work happens on the general-purpose units, in between, on data that is already in SRAM.

14. What about the causal mask?

Language models are causal: token 50 may attend to tokens 1 through 50, but not to token 51. This is usually implemented by adding −∞ to the upper triangle of the score matrix before softmax, so those positions receive zero weight.

In a standard implementation, this mask is a separate kernel that reads and writes the entire N × N matrix in order to zero out half of it. It is pure overhead, and it is exactly the sort of memory-bound step Figure 3 warns about.

Tiling turns it into something better. Because the kernel walks the score matrix block by block, it knows where each tile sits, and there are three cases:

for each tile (query block i, key block j):

    if every key in block j comes after every query in block i:
        skip the tile entirely          # nothing to compute
    else if the blocks overlap on the diagonal:
        compute the tile, apply the mask element-wise in SRAM
    else:
        compute the tile with no mask at all

Roughly half the tiles fall into the first case and are never computed. The mask stops being an extra pass over memory and becomes a branch that removes work. For causal attention this is close to a 2× saving on the whole operation, and FlashAttention-2 onwards exploits it fully.

Get Abinesh Siva Manikandan J R’s stories in your inbox

Join Medium for free to get updates from this writer.

Remember me for faster sign in

The same reasoning extends to other structured masks. A sliding-window mask skips every tile outside the band. Document masking in packed training batches skips every tile that crosses a document boundary. This is precisely the generalisation FlexAttention automates: you describe the mask as a small function, and the compiler works out which tiles can be skipped.

Part 4 - Why It Is Not a Trade-off

15. A complete walkthrough: one 4,096-token prompt

Abstractions are easier to trust once you have watched one concrete case go through. Take a single attention layer with:

• N = 4,096 tokens
• d = 64 head dimension
• 32 heads, batch size 1
• bfloat16, so every number is 2 bytes

Choosing the block size. The SRAM budget decides this. An SM has 192 KB, which holds about 98,000 numbers at 2 bytes each. The kernel needs four tiles resident at once — a block of Q, a block of K, a block of V and a block of O — so the block size B must satisfy roughly 4 · B · d ≤ M. With d = 64 that allows B up to about 384. In practice register pressure and occupancy pull it down, and real kernels use 64 or 128. Take B = 128.

The grid. 4,096 / 128 gives 32 query blocks and 32 key/value blocks, so 32 × 32 = 1,024 tiles per head, and 32,768 tiles across the layer. Each tile is a 128 × 128 patch of the score matrix: 16,384 numbers, 32 KB, which sits comfortably in SRAM alongside its inputs.

What one tile costs. Computing it is two small matrix multiplies on the tensor cores, plus a row-max, an exponential, a row-sum and the rescaling from Figure 6 on the general-purpose units. Then it is gone. Those 32 KB are overwritten by the next tile.

The traffic, counted. Q, K, V and O are each 4,096 × 64 × 2 bytes = 512 KB per head.

Standard implementation, per head: the score matrix is 4,09⁶² × 2 bytes = 33.5 MB. It is written once, read once, the softmax result is written once and read once. That is about 134 MB, plus 2 MB for the inputs and output.

FlashAttention, per head: with the query blocks as the outer loop, each Q block is loaded once and each K,V pair is loaded once per query block. That is 32 query blocks × 32 KB of keys and values per inner step × 32 steps, roughly 32 MB, plus the 1 MB for Q and O.

So the same layer moves about 33 MB instead of 134 MB, a little over 4× less - at this block size, with no mask.

Now add causality. This is a language model, so roughly half the tiles lie entirely in the future and are skipped without being computed at all. Traffic falls to around 17 MB, and the arithmetic falls by about half too. The gap is now close to 8×, which is the neighbourhood of the 9× reduction the original paper measured on GPT-2.

Two footnotes on that arithmetic. Larger blocks reduce the repeated loading further, which is exactly why the FlashAttention-4 work on Blackwell cares so much about bigger tiles and tensor memory. And on real hardware the L2 cache absorbs some of the repeated reads, so the measured improvement is usually better than this back-of-the-envelope count suggests.

And the memory held: the standard version keeps that 33.5 MB matrix alive per head for the backward pass — over a gigabyte for the layer. FlashAttention keeps m and ℓ, which is 4,096 × 2 numbers per head: 16 KB per head, 512 KB for the whole layer.

16. The backward pass and the recomputation bet

Everything so far concerns inference. Training adds a complication: the backward pass needs the attention weights again, to push gradients back through the softmax.

The standard implementation handles this the obvious way. It saved the N × N matrix during the forward pass, so it simply reads it back. That is the reason the matrix was being stored at all.

FlashAttention does not have it. It threw every tile away.

So it makes them again.

Figure 10 — Keep the matrix, or make it twice. On memory-bound hardware, making it twice is faster.

During the backward pass, for each tile, the kernel recomputes S = Q_i K_jᵀ / √d and P = exp(S − m_i) / ℓ_i from the inputs it still has — and crucially, from the saved statistics m and ℓ.

This is the quiet reason those two statistics are kept. Without them, recomputing the tile would require a second pass over the row to find its maximum and sum. With them, each tile can be reconstructed exactly, in isolation, in one shot.

The technique is recomputation, also known as gradient checkpointing, and it is not new. What is new is the arithmetic of the trade. The usual version of checkpointing accepts a slowdown in exchange for memory. Here the recomputation is so cheap compared with the HBM traffic it avoids that the backward pass gets both smaller and faster.

This is the clearest illustration of the whole principle. On memory-bound hardware, arithmetic is nearly free and memory movement is nearly everything.

17. The backward pass, in pseudocode

The forward pass is the famous half. The backward pass is where the memory saving is actually banked during training, and it has one trick of its own worth seeing.

Gradients flow backwards through the four steps of attention, so from dO (the gradient arriving at the output) the kernel needs to produce dQ, dK and dV. Differentiating a weighted sum gives dV and dP immediately. The awkward one is pushing the gradient back through softmax, whose derivative is:

dS = P ∘ ( dP − rowsum(dP ∘ P) )

where ∘ is element-wise multiplication. That rowsum(dP ∘ P) term is a global operation over the row again, and dP is an N × N matrix we are trying very hard not to materialise.

The escape is a small identity: rowsum(dP ∘ P) equals rowsum(dO ∘ O). The right-hand side involves only the output and its gradient, both of which are N × d, not N × N. So it can be computed once, cheaply, before the tile loop begins.

Call that vector D. The backward pass then looks like this:

D = rowsum(dO ∘ O)                     # one cheap pass, N×d data only

for j = 1 … e:                         # over key/value blocks
    load K_j, V_j into SRAM
    dK_j = 0, dV_j = 0
    for i = 1 … t:                     # over query blocks
        load Q_i, O_i, dO_i, m_i, ℓ_i, D_i into SRAM
        S = Q_i K_jᵀ / √d              # recomputed, not loaded
        P = exp(S − m_i) / ℓ_i         # exact, thanks to the saved m and ℓ
        dV_j += Pᵀ dO_i
        dP    = dO_i V_jᵀ
        dS    = P ∘ (dP − D_i)         # the identity above
        dQ_i += dS K_j / √d            # accumulated in HBM
        dK_j += dSᵀ Q_i / √d
    write dK_j, dV_j to HBM

Three things to notice.

S and P are recomputed, not read back. This is the recomputation bet from the previous section, executed. The inputs are Q, K and V, which we still have.

The saved m and ℓ make each tile self-sufficient. Without them, reconstructing P for a tile would require finding the row maximum and row sum again, which would mean a second pass over the entire row. With them, one tile can be rebuilt in isolation, correctly normalised, in a single step. This is the real reason those two vectors are kept.

dQ is accumulated across iterations of the outer loop, so it lives in HBM and is updated with atomic additions, while dK and dV are accumulated in SRAM within one outer iteration and written once. This asymmetry is a consequence of the loop order, and it is one of the things FlashAttention-2 revisited.

The upshot: the backward pass holds O(N) extra memory instead of O(N²), does noticeably more arithmetic than the standard version, and still finishes faster.

18. The IO complexity result

Reducing memory traffic is a good engineering story. Proving something about it is what made this a paper.

The authors analysed attention in terms of HBM accesses rather than FLOPs, and showed:

• Standard attention requires Θ(N d + N²) accesses. The N² term dominates whenever the sequence is longer than the head dimension, which is always.
• FlashAttention requires Θ(N² d² / M) accesses, where M is the size of SRAM.

Figure 11 — The same arithmetic, a different number of trips to memory.

The ratio between them is d² / M. Put in real numbers: with a head dimension of 64 and roughly 100 KB of usable SRAM, that is 4,096 / 50,000, or about 12× fewer accesses. Measured on GPT-2, the observed reduction was about 9×.

Two things in that formula are worth noticing.

The bound depends on the hardware. M is a property of the chip, not the model. A bigger scratchpad means fewer trips to HBM, automatically. The algorithm improves when the silicon does, without anyone rewriting it. This is what “IO-aware” means in the paper’s title.

They proved it is optimal. For exact attention, over a range of SRAM sizes, no algorithm can achieve asymptotically fewer HBM accesses. You cannot do better than this without changing the problem or changing the hardware.

19. Exact, not approximate

This is the section to remember if you remember nothing else.

FlashAttention produces the same output as standard attention. Not similar. Not within a tolerance. The same computation, reordered.

In practice you will see tiny differences from a reference implementation, at the level of the last bits of a floating-point number, because summing numbers in a different order gives slightly different rounding. Every fused kernel has this property, including the ones already inside PyTorch. It is not an approximation in any meaningful sense.

Compare that with what the competition was offering.

Look at the last two columns. An approximate method asks you to accept some quality loss, retrain your model, and hope the speedup materialises on real hardware. FlashAttention asks you to upgrade a library.

That asymmetry is the entire adoption story. There is no debate to have, no ablation study to run, no risk to weigh. A maintainer can switch the default and every user benefits without knowing it happened.

It also carries a lesson beyond attention: an optimisation that requires no behaviour change from its users will beat a better optimisation that does.

Part 5 - The Family, and the Impact

The algorithm has not changed since 2022. The hardware has changed three times, and each time someone had to work out which resource had become scarce and rearrange the schedule around it.

20. FlashAttention-2: fixing the work division

A year of watching version 1 run revealed something uncomfortable. It had solved the memory problem and was still only using 25–40% of the chip’s theoretical peak — while a plain matrix multiply, on the same hardware, reaches 80–90%.

Three changes closed most of the gap.

Figure 12 — Three scheduling fixes, no change to the maths.

Fewer non-matmul operations. Version 1 rescaled and normalised the running output at every block. Those divisions and exponentials run on the general-purpose units, not the tensor cores, and on an A100 a non-matmul FLOP costs roughly 16× a matmul FLOP. Version 2 keeps the output unnormalised throughout and divides exactly once at the end. Same result, far less work on the slow units.

Parallelism along the sequence. Version 1 assigned one thread block per (batch, head) pair. With a batch of 1 and 16 heads you get 16 thread blocks on a chip with 108 streaming multiprocessors, and most of the GPU sits idle. Version 2 also splits the sequence itself across thread blocks. This matters precisely in the case people care about most: long sequences, which almost always come with small batches.

Swapped loops and re-partitioned warps. Version 2 makes the query blocks the outer loop and keys/values the inner loop, so the output block stays resident in registers. Inside a thread block, it splits the queries across warps rather than the keys, so each warp owns its slice of the output outright. In version 1, splitting keys across warps meant partial results had to meet in shared memory and be summed. That synchronisation disappears.

The result: about 2× faster than version 1, reaching 230 TFLOPs/s on an A100 and 50–73% of theoretical peak. End-to-end GPT training hit 225 TFLOPs/s, or 72% model FLOPs utilisation.

There is a general lesson in this. Version 1 fixed the dominant bottleneck, and a completely different bottleneck was waiting underneath it. Performance work is layered, and you only see the next layer after you remove the one above it.

21. FlashAttention-3: Hopper and asynchrony

The H100 introduced hardware that version 2 could not use: a dedicated asynchronous copy engine (the Tensor Memory Accelerator), fourth-generation tensor cores that run asynchronously, and FP8 support.

Figure 13 — Loading, multiplying and exponentiating, all at once.

Warp specialisation. Instead of every warp doing a bit of everything, warps are given jobs. Producer warps do nothing but fetch the next tiles from HBM using the copy engine. Consumer warps do nothing but compute. They pass work through a queue, and neither has to wait for the other’s memory operations.

Interleaved softmax and matmul. In version 2 there is a strict dependency chain: compute QKᵀ, then softmax it, then multiply by V. The tensor cores idle during softmax and the special-function units idle during the matmuls. Version 3 runs them in a ping-pong: while one tile's softmax is on the special-function units, the next tile's matrix multiply is on the tensor cores.

FP8 without the accuracy cost. Eight-bit precision doubles throughput again, but naive FP8 attention loses accuracy on outlier values. Block quantisation and incoherent processing keep the numerical error roughly 2.6× lower than a straightforward FP8 implementation.

Result: 1.5–2× faster than version 2 on an H100, around 740 TFLOPs/s in 16-bit — about 75% of theoretical peak — and close to 1.2 PFLOPs/s in FP8.

22. FlashAttention-4: when softmax became the bottleneck

Then Blackwell arrived and broke the assumptions again, in a way that is genuinely surprising.

Tensor core throughput roughly doubled, from about 1 PFLOPS to 2.25 PFLOPS in FP16/BF16. Shared memory bandwidth and the exponential unit did not keep pace. The paper calls this asymmetric hardware scaling, and the consequence is remarkable: on a B200, for typical attention workloads, the non-matmul work — shared memory traffic and exponentials — now exceeds the matmul work by 25–60%.

Softmax, the awkward thing in the middle, has become the bottleneck.

Figure 14 — Blackwell made one part of the chip twice as fast and left the rest behind.

FlashAttention-4’s answers are unusually inventive:

Exponentials computed in software. Rather than queueing for the one special-function unit that Blackwell did not speed up, FA-4 approximates 2ˣ with a polynomial evaluated on the fused multiply-add units, which Blackwell has in abundance. It moves work from a scarce resource to a plentiful one.

Conditional rescaling. Recall that online softmax rescales the accumulator every time the running maximum moves. FA-4 skips that rescaling unless the shift is large enough to threaten numerical stability. According to Tri Dao’s Hot Chips presentation, this removes roughly 90% of rescaling operations.

A dedicated correction warpgroup, whose only job is performing those rescalings, keeping them off the critical path. Two query tiles are alternated in a ping-pong schedule so something is always in flight.

Written in CuTe-DSL, a Python-embedded kernel language in NVIDIA’s CUTLASS library, so compiling takes seconds instead of hours. That sounds like a footnote, but it changes how quickly the next optimisation can be found.

The result is roughly 1,605 TFLOPs/s on a B200 — up to 2.7× faster than a Triton implementation and 1.1–1.3× faster than NVIDIA’s own cuDNN kernels on the forward pass. Attention now runs at close to the speed of plain matrix multiplication, which is about as good as this problem gets.

23. The four versions side by side

The utilisation column is the honest measure of the engineering. Raw throughput rises largely because the chips do. The share of the chip actually kept busy is what the algorithm earns.

24. What it unlocked

Figure 15 — Published context limits of well-known models, on a logarithmic scale.

FlashAttention is not the sole cause of this line moving. Position encodings had to be extended, KV caches had to be managed, serving systems had to be rebuilt, and training data had to contain long documents in the first place. But the quadratic memory wall was the first obstacle, and it was the one that an exact, drop-in kernel removed without asking anyone to retrain a model.

Concretely, it bought four things:

Longer context, affordably. Memory linear in sequence length means a 128K-token prompt stores like a 128K-token prompt, not like a 16-billion-entry matrix.

Cheaper training. A 15% end-to-end speedup on BERT-large against the MLPerf 1.1 record, 3× on GPT-2 at 1K context, 2.4× on long-range arena. Those are wall-clock numbers on the whole training run, not microbenchmarks.

Better models, not just faster ones. The original paper reports 0.7 better perplexity on GPT-2 from the longer context that became affordable, and 6.4 points of lift on long-document classification. It also produced the first transformers to beat chance on Path-X (16K tokens) and Path-256 (64K), tasks that were previously out of reach.

Faster inference for everyone. Which is why you are almost certainly using it already.

Figure 16 — The default path inside libraries you already have.

PyTorch’s scaled_dot_product_attention dispatches to it. Hugging Face exposes it with attn_implementation="flash_attention_2". vLLM, SGLang, TensorRT-LLM and TGI all call into the family, and modern serving stacks pick the right version for the GPU they find: FA-2 on Ampere, FA-3 on Hopper, FA-4 on Blackwell data-centre parts.

25. The neighbours

Once the field accepted that memory movement was the thing to optimise, the same reasoning was applied up and down the stack.

Flash-Decoding deserves a note, because generation is a different problem from training. When a model produces one token at a time, each step is a single query attending over a long cached context. There is almost no parallelism in the query dimension, so the GPU sits half-empty. Flash-Decoding splits the keys and values into chunks processed in parallel, then combines the partial results using exactly the rescaling identity from Figure 6. The same trick, pointed at a different axis.

PagedAttention attacks the KV cache rather than attention itself, borrowing virtual memory from operating systems to stop the cache fragmenting. It is the foundation vLLM is built on.

FlexAttention attacks a maintenance problem. Every attention variant — ALiBi, sliding windows, document masking, soft-capping — used to need its own handwritten kernel. FlexAttention lets you write a small scoring or masking function in Python and compiles it into a fused FlashAttention-style kernel, now with a FlashAttention-4 backend on Hopper and Blackwell.

26. How to actually use it, and how to tell whether you are

You almost never call FlashAttention directly. You call something that calls it, and the interesting question becomes whether you landed on the fast path or quietly fell off it.

In PyTorch. torch.nn.functional.scaled_dot_product_attention(q, k, v, is_causal=True) is the front door. PyTorch inspects your tensors and dispatches to one of several backends: a FlashAttention kernel, a memory-efficient kernel, or a plain maths fallback that materialises the full matrix. The fallback is correct and slow, and it does not warn you.

In Hugging Face Transformers. Load a model with attn_implementation="flash_attention_2" (which requires the flash-attn package) or attn_implementation="sdpa" to go through PyTorch's dispatcher.

In serving. vLLM, SGLang and TensorRT-LLM select a backend automatically from the detected GPU architecture, generally without configuration.

The more useful knowledge is the list of things that silently disqualify you from the fast kernel:

• Float32 inputs. The flash kernels are written for FP16 and BF16. Running in full precision drops you to a different backend. This catches people who debug in float32 and never switch back.
• An arbitrary additive attention mask. Passing is_causal=True is a structured mask the kernel understands. Passing a custom float mask tensor often forces a fallback, because an arbitrary mask cannot be turned into skipped tiles. If your mask is a known shape — causal, sliding window, document boundaries — express it as one rather than as a matrix of numbers.
• Unsupported head dimensions. Support has widened over the versions, but unusual head sizes can still miss the fast path.
• Non-contiguous or oddly strided tensors. A stray transpose or a view can be enough.
• Padding instead of packing. If you pad a batch to a common length and mask the padding, you pay for the padded tokens. The variable-length entry points (flash_attn_varlen_func and the equivalents inside serving engines) take packed sequences with an offsets array and skip the padding entirely. On batches with uneven lengths this matters more than any other item on this list.

If you want to know what actually ran, PyTorch provides a context manager for restricting the backends, and a profiler that names the kernel it chose. Forcing the flash backend and watching whether it errors is a blunt but effective diagnostic.

27. Limits and honest caveats

It is worth being clear about what FlashAttention does not do.

It does not make attention sub-quadratic in compute. Memory is linear in N. Arithmetic is still proportional to N². Doubling the context still roughly quadruples the multiplications. This is why sparse attention, sliding windows, linear attention and state-space models remain active research areas: they attack the compute term, which FlashAttention leaves untouched.

It does not solve the KV cache. During generation, the cached keys and values grow linearly with context and must be kept for the whole session. At long contexts on a busy server, that cache, not attention, is the memory problem.

It is hardware-specific. Each version targets particular architectures. Consumer Blackwell cards (RTX 5090, RTX PRO 6000) use SM120 and lack the tensor-memory subsystem FA-4 depends on, so they fall back to FA-2 via Triton. AMD and Apple silicon have their own separate implementations at their own maturity levels.

Head dimension limits exist. Early versions supported head dimensions up to 128; support widened over time, but unusual configurations can still fall off the fast path silently and land on a slower fallback.

It does not help you use long context well. Fitting 200,000 tokens in memory is not the same as the model reliably finding the relevant sentence among them. That is a modelling problem, and it is still open.

Part 6 - Closing

28. Six things people get wrong

The linear-versus-quadratic confusion in the first row is the one worth carrying with you. It explains why a million-token prompt is affordable to store but still expensive to process, and why the research into approximate attention did not stop in 2022.

29. Five ideas worth stealing

The reason this paper is taught outside kernel-engineering circles is that its reasoning transfers.

Profile the system, not the formula. The field spent three years optimising FLOPs because FLOPs are what the notation shows. The clock was measuring something else entirely. Whenever you optimise a proxy, check occasionally that the proxy still correlates with the thing you care about.

Know which resource is scarce. Compute-bound and memory-bound problems want opposite things. Recomputing a value is idiotic in one regime and free in the other. Almost every optimisation is only correct relative to a bottleneck.

Trading the cheap resource for the expensive one is usually right. FlashAttention does 13% more arithmetic to avoid 90% of the memory traffic. Stated abstractly that sounds obviously good; stated as “let’s make the algorithm do more work” it sounds obviously bad. The framing matters more than it should.

Exact beats approximate when adoption is the goal. An approximate method has to win an argument with every potential user. An exact one only has to be merged.

Fuse across boundaries. Most of the waste was not inside any single operation. It was in the gaps between them, where intermediate results were parked in slow memory because that is how the abstraction worked. That pattern is everywhere in software, not just on GPUs.

30. Key takeaways

• Attention’s quadratic cost was a real problem, but the reason it was slow was memory traffic, not arithmetic.
• GPUs have a small, fast SRAM and a large, slow HBM. Standard attention wrote the N × N score matrix to HBM and read it back twice, for numbers nobody wanted.
• Online softmax lets you compute an exact softmax block by block, by carrying a running maximum and a running sum and rescaling earlier work by exp(m_old − m_new) when a larger value appears.
• Tiling applies this to the whole matrix: each tile is computed in SRAM, folded into a running output, and discarded. The full matrix is never materialised.
• Recomputation reconstructs the tiles during the backward pass instead of storing them, which is faster as well as smaller because the operation is memory-bound.
• Memory drops from quadratic to linear in sequence length; HBM accesses drop by roughly an order of magnitude; the output is exact.
• Four versions have followed, each re-tuning the schedule for the resource that the newest GPU made scarce: parallelism on Ampere, asynchrony on Hopper, softmax throughput on Blackwell.
• It is already the default in PyTorch, Hugging Face and every major serving engine. You have used it.

31. References and further reading

The papers

• Dao, Fu, Ermon, Rudra and Ré, FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, arXiv:2205.14135 (2022). The original. Section 3 contains the algorithm and the IO complexity proof.
• Dao, FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning, arXiv:2307.08691 (2023).
• Shah, Bikshandi, Zhang, Thakkar, Ramani and Dao, FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-Precision, arXiv:2407.08608 (2024).
• Zadouri, Hoehnerbach, Shah, Liu, Thakkar and Dao, FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling, arXiv:2603.05451 (2026).
• Milakov and Gimelshein, Online normalizer calculation for softmax, arXiv:1805.02867 (2018). The online softmax trick predates FlashAttention; this is where it comes from.
• Rabe and Staats, Self-attention Does Not Need O(n²) Memory, arXiv:2112.05682 (2021). Independent, closely related work on memory-efficient attention.
• Vaswani et al., Attention Is All You Need, arXiv:1706.03762 (2017). Where the formula comes from.

The neighbours

• Kwon et al., Efficient Memory Management for Large Language Model Serving with PagedAttention, arXiv:2309.06180 (2023). The paper behind vLLM.
• Dao, Haziza, Massa and Sizov, Flash-Decoding for long-context inference (2023). Blog post from the same group on making generation fast.
• Ye et al., FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving (2025).
• PyTorch team, FlexAttention and FlexAttention + FlashAttention-4 (pytorch.org/blog). How to get custom attention variants without writing kernels.

Blogs and talks worth your time

• Tri Dao’s blog (tridao.me) — the FlashAttention-2, 3 and 4 announcement posts, written for people who want the reasoning rather than the citations.
• The Stanford Hazy Research blog (hazyresearch.stanford.edu) — the original lab’s posts on all versions.
• Horace He, Making Deep Learning Go Brrrr From First Principles — the clearest short explanation anywhere of compute-bound versus memory-bound versus overhead-bound.
• GPU MODE (formerly CUDA MODE) lecture series on YouTube — several sessions walk through FlashAttention kernels line by line.

Code

• Dao-AILab/flash-attention on GitHub — the reference implementation.
• PyTorch’s torch.nn.functional.scaled_dot_product_attention documentation — how to use it without thinking about any of this.

If this explainer helped, consider following for more breakdowns of how modern AI systems actually work.