FlashAttention: Part 1

This is Part 1 of a two-part FlashAttention series. Check out Part 2 next. Part 2 covers the FlashAttention evolution (FA-1/2/3/4).

How IO-Aware Attention Makes Transformers Faster Without Approximating Attention

The reference handbook for this blog post was inspired by this tweet and is titled: Understanding FlashAttention: How IO-Aware Attention Makes Transformers Faster Without Approximating Attention.

The mechanism, in three ideas: Tiling + Online Softmax + Recomputation. Everything in this handbook¹ is elaboration on that summary. The handbook is a technical deep dive on exact tiled attention: GPU memory traffic, online softmax, forward and backward passes, IO complexity, the evolution from FlashAttention-1 through FlashAttention-4, and current framework behavior.

In this blog post, we cover the fundamental problem FlashAttention-1 (FA1) addresses, the mathematical tricks it uses, the GPU implementation of those tricks, and the architectural compatibility (MHA/MQA/GQA) of FlashAttention.

Table of Contents

0. Introduction

• 0.1 How to Read This Handbook
• 0.2 Notation

1. The Fundamental Problem

• 1.1 What FlashAttention actually optimizes
• 1.2 The attention equation is not the implementation
• 1.3 GPU memory hierarchy and why IO matters
• 1.4 Why materializing the score and probability matrices is expensive
• 1.5 Dense arithmetic is still quadratic
• 1.6 Memory-efficient exact attention predates FlashAttention

2. The Mathematical Trick

• 2.1 Tiling queries, keys, and values
• 2.2 Softmax is the difficult part of streaming
• 2.3 Online softmax from first principles
• 2.4 The blockwise merge recurrence
• 2.5 A complete numerical example

3. Putting the Mathematics onto the GPU

• 3.1 FlashAttention forward pass
• 3.2 Why dense FlashAttention is exact
• 3.3 IO complexity: what the theorem actually says
• 3.4 Memory complexity: linear auxiliary state, not linear compute
• 3.5 Causal masking and tile skipping
• 3.6 Backward pass: recompute instead of save
• 3.7 Why more FLOPs can still be faster

4. Architectural Compatibility

• 4.1 MHA, MQA, and GQA compatibility
• 4.2 Variable lengths, local attention, and dropout

5. Summary

Appendix

• Reference Map
• References
• Citation

0. Introduction

0.1 How to Read This Handbook

What stands out to me is the number of round trips to HBM in naive standard attention. Every intermediate value $(S$, $P$, $O)$ has to be written out and read back. That is the problem FlashAttention solves. The handbook frames the subject as easiest to understand when three different questions are kept separate:

1. What mathematical function is being computed? For dense attention, the target remains ordinary scaled dot-product attention.

2. How much arithmetic does that function require? Dense all-pairs query-key scoring remains quadratic in sequence length.

3. How does the implementation move data through the GPU memory hierarchy? This is where FlashAttention changes the execution of the algorithm dramatically.

The lesson I keep coming back to: FLOP count alone doesn’t decide wall-clock speed. An algorithm can do essentially the same math, or even recompute intermediates from scratch, and still finish faster because it shuttles far less data to and from high-bandwidth memory (HBM).

The word exact is doing real work here. Exactness is a statement about the mathematical function, not about bitwise reproducibility. The kernel is free to reorder floating-point operations. It is not free to change the function being computed.² $^,$ ³ $^,$ ⁴

0.2 Notation

For one attention head, let

For self-attention, typically $N_q = N_k = N$. The scaled score matrix is

where $B$ represents an optional additive mask or bias, and

Throughout, HBM refers to large off-chip high-bandwidth GPU memory. On-chip memory is a broad teaching term for much smaller, faster storage such as registers and shared memory/SRAM. Exact hardware details vary by GPU generation.

The contrast in one table:

That question, not the arithmetic, is what FlashAttention was built to answer.

1. The Fundamental Problem

1.1 What FlashAttention Actually Optimizes

Start with the ordinary attention function:

A textbook implementation boils this down to three big steps:

• form the score matrix $S$,
• apply row-wise softmax to get $P$, then
• multiply by $V$.

Mathematically, that’s perfectly fine. On a GPU, though, writing a giant intermediate to HBM and reading it back is often far more expensive than the equation lets on.

FlashAttention’s real contribution is an IO-aware implementation. It splits the computation into tiles that fit in fast on-chip memory, streams blocks of $K$ and $V$ past them, and keeps just enough row-wise softmax state around to produce the exact output. No full $N \times N$ attention matrix ever needs to land in HBM.²

My working definition of IO-awareness:

minimize data movement between the different levels of GPU memory (HBM and SRAM), rather than just trying to reduce the number of mathematical operations (FLOPs).

The reasoning is that the speed bottleneck in modern AI hardware is often not how fast the GPU can compute math, but how fast it can read and write data. This is the memory-compute tradeoff.

The problem, in one line:

compute → HBM write → HBM read → compute

The arithmetic often isn’t the bottleneck. The round trips are. Tiling is the fix. Instead of one big $QK^T$, process it in blocks:

Each score block gets:

1. computed,
2. softmax-processed,
3. immediately multiplied by its corresponding $V_j$, and
4. discarded.

For $N = 8192$, instead of an $8192 \times 8192$ intermediate matrix, you process blocks like $128 \times 128$. All the $QK$ interactions still happen, just in manageable pieces.

This framing matters. Calling FlashAttention “a faster kind of attention” misses the point. It’s better understood as an algorithm and kernel family for evaluating attention efficiently on accelerators. A model can use causal masking, RoPE, MQA/GQA, or any other attention feature and still run a FlashAttention kernel underneath it.

The original paper draws a line between this and approximate attention methods, which cut arithmetic by changing the math itself. Dense FlashAttention doesn’t do that. It keeps the math intact. The block-sparse extension from the same paper is a different story, since skipping blocks literally changes which query-key interactions get computed.²

1.2 The Attention Equation Is Not the Implementation

The equation does not tell you where tensors live. Let’s look at a standard naive attention implementation:

1. Calculate $S$, store $S$.
2. Read $S$, calculate $P$, store $P$.
3. Read $P$, calculate $O$.

The problem is the number of HBM round trips. A simple materializing implementation conceptually does:

Prior to FlashAttention, here is what the standard attention implementation looks like in more detail in terms of memory communication. Load $Q, K, V \in \mathbb{R}^{N \times d}$ from HBM, then:

1. Read $Q, K$ from HBM, compute $S$, write $S$ to HBM.
2. Read $S$ from HBM, compute $P$, write $P$ to HBM.
3. Read $P, V$ by blocks from HBM, compute $O$, write $O$ to HBM.
4. Return $O$.

The GPU spends significant time moving an $N^2$ object through memory. FlashAttention avoids repeatedly moving huge intermediates:

1. Calculate a small $S$ tile, apply softmax to the tile, use the tile with $V_j$, and discard the tile.
2. Calculate the next tile, and repeat step 1 until all tiles are completed.

The key questions I ask when looking at the attention equation:

1. How many operations are required?
   
2. What data must move between memory levels to perform those operations?

FlashAttention’s move is simple: keep the work close. Blocks of $Q$, $K$, and $V$ sit near the compute units, score tiles get built and consumed on the spot, and only the running row statistics plus the output survive across tiles. Nothing else needs to stick around.

The pattern isn’t unique to attention. Fused kernels, tiling, recomputation, and operator scheduling all do the same thing: spend a little extra arithmetic to avoid dragging huge intermediates through memory. On modern hardware, that’s usually a winning bet. Matmul throughput has raced ahead of the rest of the memory hierarchy, so the math is cheap and the data movement is what hurts. FlashAttention-2 (FA2), FlashAttention-3 (FA3), and FlashAttention-4 (FA4) take this further, and they’re upfront about it: the algorithm bends to the hardware, not the other way around.⁵ $^,$ ⁶ $^,$ ⁷

1.3 GPU Memory Hierarchy and Why IO Matters

GPUs expose a hierarchy rather than one uniform pool of equally fast memory. The names and capacities vary by architecture, but the mental model is:

Figure 1: GPU Memory Hierarchy.

FA1 models the HBM/SRAM asymmetry and explicitly optimizes the number of transfers between them.² My breakdown of each GPU memory level in the figure above:

1. HBM:

• Relatively large
• Relatively slower to access
• Physically farther from the compute units, so data must travel further to reach them
• Stores model weights, $Q/K/V$, activations, and other large tensors

2. SRAM / shared memory:

• Much smaller
• Faster
• Cheaper reuse cost

3. Registers:

• Even smaller
• More local

Why tiling helps: Suppose $Q_i$ (a query tile/block) needs to interact with many $K/V$ tiles/blocks prior to writing its output state. Instead of constantly moving $Q_i$ back and forth, we can keep it close to the compute units while processing:

Therefore, one loaded $Q_i$ can participate in lots of computation. This is called reuse. Tiling doesn’t only make tensors smaller, it also increases reuse while a tile is resident on chip. Conversely, $K/V$ blocks can be streamed through query blocks according to the chosen schedule.

This does not mean attention is “always memory-bound.” FA3 and FA4 exist partly because, as hardware changed, the dominant bottlenecks changed too.⁶ $^,$ ⁷ The bottleneck depends on:

• Sequence length, $N$
• Head dimension, $d$
• dtype
• Mask pattern
• GPU generation (Ampere/Hopper/Blackwell)
• Forward vs. backward pass
• Which kernel is running

1.4 Why Materializing the Score and Probability Matrices Is Expensive

The quadratic intermediate becomes concrete very quickly. Suppose a batch contains one sequence, with 32 attention heads, sequence length $N = 8192$, and a two-byte dtype such as FP16 or BF16. Concretely, my arithmetic:

At two bytes per element, the total memory is: \(32 \times 8192^2 \times 2 = 4{,}294{,}967{,}296 \text{ bytes} = 4 \text{ GiB}.\)

One dense tensor with shape $[1, 32, 8192, 8192]$ contains $32 \times 8192^2$ elements. At two bytes per element, that is $4 \text{ GiB}$. That is the size of one full score- or probability-like tensor for this example. Imagine reading and writing this much data.

Naive pipeline:

A naive decomposition can produce arrays of this size at several stages. That’s not to say every framework keeps $S$ and $P$ alive at once, or that compiler fusion hasn’t already cut some of the traffic. The example is there to make the problem concrete: a dense $N^2$ intermediate is big enough that rewriting and rereading it can dominate memory and bandwidth.

Figure 2: Naive Attention Implementation Pipeline.

FlashAttention pipeline:

FlashAttention skips the naive decomposition entirely. Score tiles get formed, softmaxed on chip, used to accumulate their contribution to the output, then thrown away.² The full matrix never touches HBM.

Figure 3: FlashAttention Implementation Pipeline.

This matters most during training, where naive autograd would want those large intermediates saved for the backward pass.

Important distinction: FlashAttention acknowledges that $N^2$ interactions exist and says: Do not store the entire result of these interactions as a giant intermediate if we can consume each piece immediately.

1.5 Dense Arithmetic Is Still Quadratic

FlashAttention changes the memory schedule, not the math. The function itself is still dense attention. That means forming all query-key scores for $N$ tokens with head dimension $d$ still costs work proportional to $N^2 d$. There’s no way around that.

Multiplying probabilities by values is another dense pairwise matrix multiplication of the same broad order. FlashAttention reshuffles these operations, but it doesn’t skip the dense set of query-key interactions. Those still happen.² $^,$ ⁵

So three things must not get conflated:

• Arithmetic complexity: dense attention stays at $O(N^2 d)$.
• Large intermediate storage: FlashAttention avoids an $O(N^2)$ materialized score/probability tensor in HBM.
• HBM traffic: the original paper proves a lower IO cost than standard materializing attention under its two-level memory model.²

This is how a method makes much longer sequences practical without making long context “free.” Double $N$, and you roughly quadruple the number of dense query-key pairs. FlashAttention just attacks the data-movement and memory-footprint side of that.

When you report speedups, always separate asymptotic arithmetic from measured runtime. A kernel can run several times faster at the same $O(N^2 d)$ complexity because the constant factors, occupancy, fusion, and memory traffic change dramatically.

Essentially, the key takeaways are summarized below:

1.6 Memory-Efficient Exact Attention Predates FlashAttention

It would be historically inaccurate to say FlashAttention first discovered that exact attention can avoid quadratic memory. Earlier work by Rabe and Staats demonstrated exact memory-efficient attention with quadratic computation but subquadratic memory. They showed that attention does not require $O(N^2)$ memory with respect to $N$.⁸ Milakov and Gimelshein also did earlier work on online softmax, showing a memory-efficient, online recurrent way of computing the classical stable softmax normalizer.⁹ FlashAttention builds the same kind of running-max/running-normalizer idea into tiled attention, while also accumulating the value-weighted output.²

The lineage from stable softmax to FlashAttention is shown in Figure 4 below:

Figure 4: Memory-Efficient Exact Attention Lineage.

FlashAttention’s key contribution was to bring together the techniques below into a practical, high-performance algorithm:

1. Tiling
2. Online softmax
3. Fused computation
4. GPU memory hierarchy awareness
5. IO complexity analysis

This distinction matters because “memory-efficient” does not automatically mean “IO-optimized for a particular hardware model.”

2. The Mathematical Trick

2.1 Tiling Queries, Keys, and Values

Most important. This is the section where the trick lives. Score tiles are small enough to stay near the compute units. Each tile’s contribution gets folded into the running row statistics and the output accumulator, then the tile is thrown away.

• Suppose $Q \in \mathbb{R}^{N_q \times d}$, $K \in \mathbb{R}^{N_k \times d}$, and $V \in \mathbb{R}^{N_k \times d_v}$. Instead of processing all $N_q, N_k$ rows at once, divide them into blocks. For example:

  \[Q: \boxed{Q_1} \boxed{Q_2} \boxed{Q_3} \quad K, V: \boxed{K_1, V_1} \boxed{K_2, V_2} \boxed{K_3, V_3}\]

  For one query-key/value pair ($Q_i, K_j/V_j$), we process that score tile $S_{ij}$ locally:

  \[S_{ij} \doteq \frac{Q_i K_j^T}{\sqrt{d}} + B_{ij},\]

  Then, for each query block:

    1. Calculate score tile, S_ij
    2. Row-wise softmax and update
    3. Multiply by V_j and accumulate
    4. Update the running output
    5. Discard the tile 
    6. Move to the next tile, never writing a global S or P matrix.
  

  Recall that softmax is defined as

  \[\mathrm{softmax}(x_i) = \frac{e^{x_i}}{\sum_j e^{x_j}}.\]

• Matrix multiplication is easy to tile because it is fundamentally accumulation:

  \[AB = \sum_j A_j B_j,\]

  so we can calculate pieces and add them.

• Softmax is harder because every element depends on the entire row. If we process the first block, we do not know the eventual denominator. Even worse, numerical stability requires knowing the maximum.

• In the kernel, we choose tile shapes and loop order based on on-chip capacity, head dimension, GPU generation, causal structure, and work partitioning. FA1 uses tile sizes derived from SRAM capacity $M$.²

2.2 Softmax Is the Difficult Part of Streaming

A numerically stable softmax for one row $x_1, \ldots, x_N$ uses

Subtracting $m$ keeps large logits from blowing up. However, it seems to create a streaming problem: how does one normalize an early block if a later block turns out to have a bigger maximum?

The answer is to keep running statistics and rescale them when the maximum changes. Say the state after the earlier elements is $(m_{\text{old}}, \ell_{\text{old}})$, and the new block’s maximum is $m_b$. Then

Anything accumulated under the old max converts to the new one by a single factor:

The new block is then evaluated against the same $m_{\text{new}}$.

None of this is an approximation. It’s just the identity

The online-normalizer recurrence predates FlashAttention and is the mathematical basis for streaming stable softmax.⁹

Worked example 1: why subtract the max

Start with two extreme values to see why the max subtraction matters:

The exponential values $1$ and $0.368$ are perfectly manageable.

Worked example 2: streaming two blocks

Now let’s work through streaming two blocks.

Putting it all together, with $x = [2, 1, 4, 3]$:

This matches the full softmax computed in one shot.

2.3 Online Softmax from First Principles

Process scalar logits $x_1, x_2, \ldots$ one at a time. Initialize

After observing $x_j$:

Milakov and Gimelshein showed this produces the same stable softmax normalizer as the conventional safe-softmax procedure, with fewer passes over the input.⁹

Attention needs a weighted value sum, so we also maintain a value-weighted, unnormalized accumulator,

When the maximum moves from $m_{j-1}$ to $m_j$, the accumulator is rescaled exactly like the normalizer before the new value contribution is folded in:

After the last element, $x_n$, the output is

\(o = \frac{a_n}{\ell_n}.\)

The same idea works row by row and block by block, which is what makes a tiled exact softmax-attention forward pass possible.

Key Insight. We don’t need the entire probability vector. We only need enough information to reconstruct its contribution to the final output.

2.4 The Blockwise Merge Recurrence

For one query row, suppose the running state after some key blocks is

where $m$ is the running maximum score, $\ell$ is the stable softmax normalizer under that maximum, and $a \in \mathbb{R}^{d_v}$ is the unnormalized value accumulator.

A new block lands with scores $s \in \mathbb{R}^b$ and values $V_b \in \mathbb{R}^{b \times d_v}$. First, find the block’s own maximum and update the running one:

The rescale factor and the block’s exponentials follow:

Then update:

At the end:

When you process a block of query rows rather than a single row, the running state changes shape. $m$ and $\ell$ each become vectors with one entry per row, and $a$ becomes a matrix. If you have masks, apply them to the score tile before the exponentials. Masked positions contribute zero probability, so they drop out of the update entirely. FA1 and FA2 express equivalent running-max/running-normalizer/output updates in block form.² $^,$ ⁵

This recurrence is what makes tiled exact attention possible.

2.5 A Complete Numerical Example

Take one already-scaled, unmasked attention row

with two-dimensional values

The global maximum is 4, so the stable unnormalized weights are

Their sum is

and full softmax gives

Now run it through in two blocks, the way a tiled kernel would: $[2, 1]$ and $[4, 3]$.

Block 1: $s = [2, 1]$. Running max $m_1 = 2$, stable exponentials $[e^{0}, e^{-1}] = [1, e^{-1}]$. So

and the value accumulator comes out to

Block 2: $s = [4, 3]$. The block max is $m_b = 4$, beating the running max, so $m’ = \max(2, 4) = 4$. The rescale factor:

Then

and

Output:

matching the full-row computation.

This rescale-instead-of-recompute is what makes streaming possible.

3. Putting the Mathematics onto the GPU

3.1 FlashAttention Forward Pass

The forward pass, conceptually:

1. Partition $Q$ into query-row tiles, and $K$ and $V$ into key/value tiles.
2. For each query tile, initialize row-wise running maxima, normalizers, and output accumulators: $m = -\infty$, $\ell = 0$, $a = 0$.
3. Load a $K_j, V_j$ tile and form the local score tile $S_{ij} = Q_i K_j^T / \sqrt{d}$.
4. Apply the masks or biases $B_{ij}$ belonging to this tile.
5. Find the tile maximum, update the running maximum $\max(m, m_b)$, and rescale the previous state.
6. Exponentiate the current tile relative to the updated maximum.
7. Update the normalizer $\ell$ and the value-weighted accumulator $a$.
8. Repeat for every required key tile.
9. Normalize the accumulator row-wise, $O = a / \ell$, then write the output.

The original algorithm picks block sizes so the tiles and running state fit in on-chip SRAM, cutting trips to HBM.² A quick summary of the forward pass is shown below:

for each K/V block:
    scores = Q @ K_block.T / sqrt(d)
    update the running max m and normalizer l
    update the output accumulator a
O = a / l   # normalize once after the last block

You can reproduce the algebra in a few lines of PyTorch, but that’s not the same as a high-performance FlashAttention kernel. Real production implementations additionally exploit GPU-specific tiling, registers, shared memory, warp scheduling, tensor cores, async copies, specialized intrinsics, and occupancy optimization.

Minimal educational PyTorch.

import math
import torch

# Single query block, tiled over K/V. Production kernels also
# tile Q; the outer loop is omitted here to focus on the recurrence.
def tiled_attention(q, k, v, block=128):
    scale = 1 / math.sqrt(q.shape[-1])
    m = q.new_full((q.shape[0],), -float("inf"))
    l = q.new_zeros(q.shape[0])
    a = q.new_zeros((q.shape[0], v.shape[-1]))

    for j in range(0, k.shape[0], block):
        kj, vj = k[j:j+block], v[j:j+block]
        s = (q @ kj.T) * scale
        m_new = torch.maximum(m, s.max(dim=-1).values)
        alpha = torch.exp(m - m_new)
        p = torch.exp(s - m_new[:, None])
        a = a * alpha[:, None] + p @ vj
        l = l * alpha + p.sum(dim=-1)
        m = m_new

    return a / l[:, None]

The gap between this educational code and a production kernel matters a lot for FA2, FA3, and FA4. Every generation keeps the same semantic structure, but schedules the work differently for newer hardware.⁵ $^,$ ⁶ $^,$ ⁷ Same math, different schedule, each tuned to the hardware of its generation.

3.2 Why Dense FlashAttention Is Exact

Does FlashAttention approximate attention? For dense FlashAttention, NO. The target function is untouched:

A short checklist for why this counts as exact:

• No Q-K pairs are intentionally removed.
• No low-rank approximation is introduced.
• No alternative kernelized attention function replaces softmax.
• Tiling just processes the same interactions in blocks.

Online softmax doesn’t swap the exponential or normalization for something else either. The block recurrence only changes the order in which sufficient statistics get accumulated.

But exact needs a qualification. Suppose two implementations compute $a + b + c$:

• A: $(a + b) + c$
• B: $a + (b + c)$

In exact mathematics,

In floating point, for some values,

Tiling changes the reduction order. Fusion can change how rounding shakes out. So:

Three caveats keep the word exact precise:

• Floating point has finite precision. Reorder the additions and reductions and you’ll get small numerical differences from another implementation. PyTorch says so directly: the output of SDPA can differ depending on which backend kernel is chosen, because each backend accumulates floating-point values in a different order.³ $^,$ ⁴

• Dropout is random during training. If you want two runs to match bit-for-bit, you have to line up the random behavior too. Attention semantics alone won’t be enough.

• Sparse variants aren’t dense. The block-sparse variant in the original paper skips blocks entirely, so it approximates full dense attention rather than reproducing it.²

This distinction becomes important in numerical testing. A sensible tolerance depends on dtype, accumulation order, sequence length, and backend, rather than requiring binary identity.

3.3 IO Complexity: What the Theorem Actually Says

The original FlashAttention paper works in a two-level memory model: HBM plus on-chip SRAM of size $M$. Under its assumptions, which include head dimension $d$ and an SRAM size in the range $d \leq M \leq Nd$, the two attention implementations require:

• Standard materializing attention: $\Theta(Nd + N^2)$ HBM accesses
• FlashAttention: $\Theta\left(\frac{N^2 d^2}{M}\right)$ HBM accesses

The second quantity is a communication-complexity bound, and a few things are worth keeping straight:

• It holds for a specific memory model, not for every GPU.
• It counts scalar elements moved between the modeled memory levels, up to asymptotic factors. It is not a byte count.
• It is not the arithmetic complexity. Dense attention still performs $O(N^2 d)$ work.
• More usable on-chip memory $M$ means more reuse and less modeled HBM traffic.

The intuition: bigger on-chip memory lets larger working tiles stay resident. Bigger tiles mean more reuse and fewer HBM reloads. In the theorem’s regime, that reuse is exactly why the HBM-access bound shrinks as $M$ grows. It’s also why GPU architecture matters so much to FlashAttention performance. For typical values ($d$ of 64 to 128, $M$ around 100KB), $d^2$ is many times smaller than $M$, so FlashAttention needs many times fewer HBM accesses.²

3.4 Memory Complexity: Linear Auxiliary State, Not Linear Compute

FlashAttention papers often say memory becomes linear instead of quadratic. That’s true, but it’s easy to misread, so it’s worth being precise about what’s linear and what isn’t:

A naive attention implementation creates $O(N^2)$ intermediates. FlashAttention doesn’t. It keeps only:

• $Q/K/V$ tiles
• row-wise $m$
• row-wise $\ell$
• the output accumulator $a$

The inputs and output already occupy $O(Nd)$ memory on their own. A materializing dense attention implementation additionally creates $O(N^2)$ score and probability state. FlashAttention skips those full matrices and holds only tile-sized working state on chip, plus $O(N)$ row-wise statistics. So the extra memory tied to the attention operation grows linearly with sequence length rather than quadratically. But linear memory is not linear compute. The computation is still $O(N^2 d)$.

Example. For $N = 10{,}000$, there are roughly $10{,}000^2 = 100{,}000{,}000$ Q-K interactions. FlashAttention doesn’t make those disappear. It just prevents you from needing a giant intermediate holding all of them at once.

Training is where this difference bites hardest. A straightforward backward pass would otherwise want the full probability matrix $P$ kept around. FlashAttention stores compact information instead, the output plus row-wise normalization statistics, and recomputes score and probability tiles during the backward pass.

Two more things to keep in mind. First, don’t claim that total model memory is $O(N)$. Other Transformer components consume activation memory, and autoregressive serving has KV-cache memory that grows with context length. FlashAttention changes how the attention computation manages its intermediates, not the whole model’s memory profile. Second, FlashAttention and activation checkpointing can coexist, since both trade recomputation for reduced stored state, just at different scopes of the training graph.

3.5 Causal Masking and Tile Skipping

In autoregressive attention, query token $i$ can’t attend to future key token $j > i$. The attention matrix looks like

with mask

Apply the mask before softmax and you get

A materialized mask is another $N \times N$ object. But a tiled kernel can reason about the geometry of each tile directly:

• Blocks strictly above the causal boundary are fully masked, so they can be skipped.
• Blocks strictly below the boundary are fully valid.
• Only blocks intersecting the diagonal need element-level causal masking.

That’s a lot of invalid tiles you never have to compute. FA2 explicitly exploits this structure,⁵ and current implementations also define alignment rules for unequal query/key lengths. The FlashAttention interface aligns the causal mask to the bottom-right corner of the attention matrix, whereas PyTorch’s is_causal=True assumes upper-left alignment. The two only differ when $N_q \neq N_k$, so check which convention your code path uses.¹⁰ $^,$ ³

Connection to local attention. Sliding-window attention works on the same principle. If each token only attends to 128 nearby tokens, most tiles can be skipped the same way. But this is a different move: you’ve changed the mathematical attention pattern itself. The model is now doing sparse or local attention, and FlashAttention is just the kernel executing that pattern. It’s no longer the same problem as unrestricted dense attention.

3.6 Backward Pass: Recompute Instead of Save

A naive training implementation saves $P = \mathrm{softmax}(S)$ for the backward pass. That’s an $N^2$ tensor sitting in HBM. FlashAttention doesn’t keep it. It stores compact row-wise normalization information instead, then recomputes score and probability tiles when gradients are needed.

Forward:

• Store the output and row-wise normalization statistics.
• Do not store $P \in \mathbb{R}^{N \times N}$.

Backward:

• For each tile: recompute $QK^T$, reconstruct the local probabilities, compute gradients, and discard the tile.
• Same trade as before: more arithmetic, less memory traffic.

For

the backward pass relies on a row-wise identity, where $D_i$ is a per-row scalar:

After recomputing a tile of $P$, the local derivatives fall out as

and then

Masks contribute zero probability and zero gradient. FA2 stores a row-wise log-sum-exp quantity so the probability tile can be reconstructed stably from recomputed scores.⁵

This is the memory-compute tradeoff applied to the backward pass.

3.7 Why More FLOPs Can Still Be Faster

This is one of the biggest lessons from FlashAttention. The usual assumption is that fewer FLOPs means faster. On a GPU, that’s not always true.

Different operations run at completely different throughputs. Tensor cores chew through matrix multiplication. Everything else is comparatively expensive: exponentials, reductions, synchronization, shared-memory operations, and HBM transfers.² $^,$ ⁵ Therefore, doing extra arithmetic can be the right move if it eliminates expensive memory traffic.

Consider two options:

• A: compute, write huge $P$ to HBM, read $P$ back, compute
• B: compute, discard, recompute later

B does more arithmetic, but it avoids pushing a giant tensor through HBM. On modern accelerators, that trade usually wins.

Every FlashAttention generation is built around this. FA2’s goal is partly to reduce non-matmul FLOPs, since those don’t have the same throughput as tensor-core GEMMs. It also reworks the work partitioning so more of the GPU is doing useful work at once.⁵ FA3 pushes further on Hopper, overlapping matrix multiplication, softmax, and data movement with asynchronous hardware features.⁶ FA4 responds to Blackwell, where tensor-core throughput grew faster than other resources. That shift makes exponentials and shared-memory traffic relatively more important, and FA4 is designed around it.⁷

This is the systems lesson that makes FlashAttention matter beyond attention itself. Hardware efficiency doesn’t come from the math alone. It comes from co-designing the mathematical scheduling with the memory/execution hierarchy.

4. Architectural Compatibility

4.1 MHA, MQA, and GQA Compatibility (Shared K/V Heads)

FlashAttention isn’t tied to standard Multi-Head Attention (MHA). In general, the inputs have shapes

Let’s look into FlashAttention’s compatibility with MHA, Multi-Query Attention (MQA), and Grouped-Query Attention (GQA). Three head-sharing regimes are distinguished by $H_q$ versus $H_{kv}$:

• MHA: Every query head has its own K/V head. \(H_q = H_{kv}\)
• MQA: All query heads share a single K/V head. \(H_{kv} = 1\)
• GQA: Several query heads share each K/V head. \(1 < H_{kv} < H_q\)

Here’s the distinction that matters:

1. MQA/GQA defines the architecture. It decides how heads share K/V projections.
2. RoPE defines the position transform. It decides how position gets baked into Q and K.
3. FlashAttention defines the execution. It decides how the resulting attention operation gets scheduled and computed efficiently.

Any combination of these three orthogonal concepts can be used together. A kernel doesn’t care how many query heads share a K/V head, or whether the Q/K inputs were RoPE-rotated. It just evaluates attention between each query head and its assigned K/V head. Head sharing changes the architecture and the KV-memory footprint. RoPE changes the input representation. FlashAttention changes how the work gets scheduled on the hardware.

Current implementations do impose one shape constraint. In the Dao-AILab kernels, the number of query heads must be divisible by the number of KV heads¹⁰:

PyTorch SDPA exposes enable_gqa, but its docs have described GQA as an experimental feature with its own shape constraints, so production code should follow the exact version’s documentation rather than assuming universal fused-kernel support.³

4.2 Variable Lengths, Local Attention, and Dropout

Real systems aren’t always the clean case of same sequence length, dense attention, and no dropout.

Production implementations support a broader feature set:

• variable-length sequences
• causal attention
• sliding-window attention (SWA)
• dropout, in training-oriented interfaces
• MQA/GQA
• ALiBi-style score bias
• KV-cache decoding, including optional RoPE handling

But these are implementation capabilities, not properties of the FlashAttention mathematical idea itself.¹⁰

What a given build supports depends on:

• GPU
• CUDA/ROCm backend (NVIDIA and AMD have separate implementations, with different support matrices)
• dtype
• head dimension
• mask
• library version
• kernel generation

One practical warning on dropout. PyTorch’s scaled_dot_product_attention applies dropout according to its dropout_p argument regardless of training or eval mode, so eval code needs to pass 0.0 explicitly when dropout should be off.³

5. Summary

Let’s walk through the whole FlashAttention story in one derivation, and hit the key points to remember.

Ordinary attention: $QK^T \in \mathbb{R}^{N \times N}$. Materializing it creates an enormous intermediate matrix.¹¹

Naive execution: This creates huge memory traffic.

FlashAttention: Tiling is used. Split $Q \to Q_i$ and $K, V \to K_j, V_j$, then calculate the score tile:

• The problem with tiling: Softmax needs the maximum and the denominator over the whole row.

• The solution: online softmax

  1. Maintain $(m, \ell, a)$ where:
     • $m$ = running maximum
     • $\ell$ = running sum of exponentials (the softmax denominator)
     • $a$ = running unnormalized value-weighted sum (accumulator)
  2. When a new block has a larger maximum $m’$, calculate a rescaling factor $\alpha$ and rescale the old state ($\ell$ and $a$) by it:
  \[\alpha = e^{m - m'}\]

• Therefore, we can process the full FlashAttention pipeline without ever materializing the $N \times N$ matrix:

  \[\text{tile} \to \text{score} \to \text{online softmax} \to V \text{ accumulation} \to \text{discard tile}.\]

• Result: The same exact dense attention mathematics, but much better memory behavior: less HBM traffic, much smaller intermediates. Arithmetic remains $O(N^2 d)$.

FlashAttention in one table:

Appendix

Reference Map

The table below shows which paper to refer to for each question.¹

References

Citation

If you found this blog post helpful, please consider citing it:

@article{obasi2026FlashAttentionPt1,
  title   = "FlashAttention: Part 1",
  author  = "Obasi, Chizoba",
  journal = "chizkidd.github.io",
  year    = "2026",
  month   = "Sep",
  url     = "https://chizkidd.github.io/2026/09/13/flashattention/"
}