Flash Attention: The Hidden Gem

# flash attention
y = F.scaled_dot_product_attention(q, k, v, is_causal=True)

It is a single line in PyTorch. One call. But the math sitting underneath it is enormous, and most people who use it have no idea what is actually happening. That gap is what makes it a hidden gem to me.

I will try to explain the math as closely as I have understood it. There is a depth limit I have reached, and I will be honest when we get there. But we will go far.

Background

Before we talk about Flash Attention, we need to talk about where computation actually happens and where it gets stuck. Because the problem Flash Attention solves is not about math. It is about memory.

The attention mechanism computes this:

$$\text{attn} = \text{softmax}\!\left(\frac{QK^T}{\sqrt{d_k}}\right) V$$

Before the attention formula even runs, the input $X$ is projected into three separate matrices using learned weight matrices $W_Q$, $W_K$, $W_V$:

$$Q = XW_Q, \quad K = XW_K, \quad V = XW_V$$

$Q$ is the query, what each token is looking for. $K$ is the key, what each token is advertising. $V$ is the value, what each token actually carries. These three projections are done by nn.Linear(), which is already well optimized. That part is not the concern. The concern is everything that happens after: the matrix multiply, the softmax, the final multiply with $V$. That chain is not quite optimized, and the reason is memory.

A GPU has two main places it stores data during computation. There is SRAM (shared memory), which lives on-chip right next to the compute cores. It is tiny, around 20 MB on an A100, but extremely fast at 19 TB/s.

Then there is HBM (High Bandwidth Memory), the large off-chip memory at around 40 GB, but much slower at 1.5 TB/s. That is a 12x bandwidth gap.

Why is Flash Attention needed?

SRAM is fast but tiny. HBM is large but slow. Ideally we want to compute attention entirely inside SRAM. The problem is that the intermediate matrices do not fit. Here is what a standard implementation actually does.

Step 1. Compute the score matrix.

$$S = \frac{QK^T}{\sqrt{d_k}}$$

$Q$ and $K$ are both shape $N \times d$, so $S$ comes out $N \times N$. Too large for SRAM, so it is written to HBM.

Step 2. Apply softmax row by row.

$$P = \text{softmax}(S)$$

$S$ is read back from HBM. The result $P$, also $N \times N$, is written back to HBM.

Step 3. Multiply with $V$.

$$O = PV$$

$P$ is read from HBM one more time. The final output $O$ is written back.

Count the HBM trips: write $S$, read $S$, write $P$, read $P$, write $O$. Five passes over an $N \times N$ matrix. For a sequence length of 1024 with float32, that is about 4 MB per pass, 20 MB of HBM traffic per head per forward pass. The GPU cores are not the bottleneck. They sit idle, waiting for data. This is what I/O bound means.

Enter Flash Attention

Flash Attention is the solution to exactly this problem. It does not change the math. It changes where and how the math is executed.

To understand what it does, you need to know what a kernel is. A GPU kernel is a function that runs directly on the GPU hardware. When you call something like torch.matmul or F.softmax, each of those is a separate kernel.

Each kernel loads its input from HBM, does its computation, and writes its output back to HBM. That is why standard attention burns so many HBM trips: it is three separate kernels, each one blind to what the others are doing.

Flash Attention replaces all three steps with a single fused kernel. One kernel handles the score matrix, the softmax, and the final multiply together. Because it controls the whole computation, it can keep the intermediates in SRAM and never write $S$ or $P$ to HBM at all. Only the final output $O$ touches HBM. This is what fused means, and it is the core idea.

There are three Flash Attention papers.

FlashAttention (2022) by Dao et al. introduced the IO-aware algorithm and proved it produces exact attention.

FlashAttention-2 (2023) rewrote the parallelism strategy to better utilize GPU cores and cut non-matrix-multiply operations.

FlashAttention-3 (2024) targeted the Hopper architecture (H100) specifically, adding asynchronous execution and FP8 support.

The differences between them will be covered below.

Road Map

The rest of this post is split into two phases. Each phase starts with the math you need to understand it, then walks through the algorithm with a worked example.

Phase 1: Forward Pass. How Flash Attention computes the attention output in a single pass over the input, without ever writing the full $N \times N$ matrix to HBM.

Phase 2: Backward Pass. How gradients flow back through the fused kernel, and why that requires a different trick than the forward pass.

If you get lost at any point, come back here and find where you are.

Phase 1: Forward Pass

Three things to understand before the algorithm makes sense.

Safe Softmax

What softmax basically does is it computes the exp of every value in a row, sums all those exp values, and divides each individual value by that sum.

$$\text{softmax}(x_i) = \frac{e^{x_i}}{\underbrace{\sum_{j=1}^{N} e^{x_j}}_{\text{norm factor}}}$$

The denominator $\sum_{j=1}^{N} e^{x_j}$ is the norm factor. It sums all the exponentials across the row, which forces every output to be between 0 and 1 and makes all outputs sum to 1. Without it you just have raw exponentials that carry no meaning as a distribution.

It looks innocent but has a problem. If the value $x_i$ in the array is large, then $e^{x_i}$ will be even larger. The exponents will explode, which leads to numerical instability: the values cannot be represented by float32 or float16.

So a safe way to do it is to multiply both the numerator and denominator by a constant $c$:

$$\frac{c \cdot e^{x_i}}{c \cdot \sum_{j=1}^{N} e^{x_j}} = \frac{e^{x_i + \log c}}{\sum_{j=1}^{N} e^{x_j + \log c}}$$

Let $-\log c = k$, so $\log c = -k$. Then:

$$\text{softmax}(x_i) = \frac{e^{x_i - k}}{\sum_{j=1}^{N} e^{x_j - k}}$$

Now we choose $k = \max(x)$, the maximum of that row. What happens is most of the values $(x_i - k)$ will be negative, and for $x_i = \max(x)$ it will be exactly 0. Looking at the graph of $e^x$: when the value is negative the result is between 0 and 1, when the value is 0 the result is exactly 1. So the exponents do not explode anymore. Numerically stable.

Online Softmax

Safe softmax solves the instability problem but it is not quite optimized. Standard safe softmax requires reading the input data multiple times:

(a) To find the max in each row.

(b) To compute the exp sum.

(c) To normalize the values.

That is three separate passes over the data. Online softmax solves this by keeping a running maximum $m$ and a running sum $l$ which change as new data is introduced. We never need to re-read the full row.

Initialize $m_0 = -\infty$, $l_0 = 0$. For $i = 1$ to $N$:

$$m_i = \max(m_{i-1},\ x_i)$$ $$l_i = l_{i-1} \cdot e^{m_{i-1} - m_i} + e^{x_i - m_i}$$

At the end, for each $k$:

$$\text{softmax}(x_k) = \frac{e^{x_k - m_N}}{l_N}$$

How does this solve the three pass problem? The running $m$ tracks the maximum seen so far, so we find it on the fly without a dedicated pass. The running $l$ accumulates the exp sum as values arrive. And the final division by $l_N$ is the normalization. All three happen in one sweep.

The term $l_{i-1} \cdot e^{m_{i-1} - m_i}$ is what makes this work. When a new larger max arrives, this rescales the old sum to match the new max before adding the new term. This exact correction idea will appear again inside the Flash Attention algorithm.

Block Matrix Multiplication

A normal matrix multiplication computes every element of the result in one shot. Block matrix multiplication does the same thing but splits the input matrices into smaller tiles and computes the result tile by tile. The final answer is identical. The difference is that each tile is small enough to fit in SRAM, which is exactly what Flash Attention needs.

We have $Q$ and $K$ both of shape $(8, 128)$, so $K^T$ is $(128, 8)$.

$Q$ is divided into 4 row blocks of 2 rows each. Each block $Q_i$ is shape $(2, 128)$.

$K^T$ is divided into 4 column blocks of 2 columns each. Each block $K_j^T$ is shape $(128, 2)$.

Each tile $S_{ij} = Q_i K_j^T$ comes out shape $(2, 2)$ and is computed independently.

Now here is the problem with applying softmax to these blocks.

Softmax of $S_{11}$ alone would use only the max of $S_{11}$. But the correct softmax for row 1 needs the max across all four blocks in that row: $S_{11}$, $S_{12}$, $S_{13}$, $S_{14}$. Using a local max gives the wrong normalization.

This is exactly why online softmax is needed. By maintaining a running max $m$ that updates as each column block is processed, Flash Attention computes the correct global softmax without ever seeing the full row at once.

The Algorithm

This is the actual FlashAttention-2 forward pass algorithm from the paper. It looks dense but do not be overwhelmed. We will explain every line.

Let's take an example: $N = 8$ tokens, $d = 128$ embedding dim, block size $B_r = B_c = 2$, one head. $T_r = T_c = 4$.

1. Divide $Q$, $K$, $V$ into row blocks. Each is shape $(8, 128)$, split token-wise into 4 blocks of 2 rows. $Q \to Q_1, Q_2, Q_3, Q_4$ each $(2 \times 128)$. Same for $K$ and $V$. The split is always along the token dimension. $K_j$ is a $(2 \times 128)$ row slice, transposed to $(128 \times 2)$ for the dot product.

2. $O$ is the output matrix, shape $(8, 128)$. Split the same way as $Q$: $O \to O_1, O_2, O_3, O_4$ each $(2 \times 128)$. This is where results land after each outer loop iteration.

3. The computation.

Outer loop over $i = 1$ to $4$, one $Q_i$ at a time. For each $i$, initialize in SRAM:

$$O_0 = \underbrace{\begin{bmatrix} 0 & \cdots & 0 \\ 0 & \cdots & 0 \end{bmatrix}}_{2 \times 128} \qquad l_0 = \begin{bmatrix} 0 \\ 0 \end{bmatrix} \qquad m_0 = \begin{bmatrix} -\infty \\ -\infty \end{bmatrix}$$

$m$ is a vector of shape $(2,)$: one running maximum per row. $l$ is the running normalization sum, same shape. Load $Q_i$ from HBM to SRAM and hold it for the full inner loop.

Inner loop over $j = 1$ to $4$, loading one $(K_j, V_j)$ pair each step.

4. Inner loop $j = 1$.

$$S_{11} = Q_1 K_1^T \quad (2 \times 128) \cdot (128 \times 2) \to (2 \times 2)$$ $$m_{11} = \max(\text{rowmax}(S_{11}),\; m_0) = \text{rowmax}(S_{11})$$ $$P_{11} = \exp(S_{11} - m_{11})$$ $$l_{11} = \underbrace{l_0 \cdot e^{m_0 - m_{11}}}_{=\;0} + \text{rowsum}(P_{11}) = \text{rowsum}(P_{11})$$ $$O_{11} = \underbrace{\text{diag}(e^{m_0 - m_{11}}) \cdot O_0}_{=\;0} + P_{11} V_1 = P_{11} V_1$$

$m_{11}$ is the largest score seen so far. It keeps the exponents from blowing up. $l_{11}$ is the running sum of the unnormalized weights, the partial denominator for softmax. We do not divide $P_{11}$ by $l_{11}$ right away because we have only seen $K_1$. The true normalizer sums over all four key blocks. Dividing now would give the wrong probabilities. We wait until step 6, when $l$ is complete.

5. Inner loop $j = 2$. Same five steps. $m_{12}$ may be larger than $m_{11}$, so the correction factor $e^{m_{11} - m_{12}}$ rescales the old $l$ and $O$ before the new contribution is added.

$$S_{12} = Q_1 K_2^T \quad (2 \times 2)$$ $$m_{12} = \max(\text{rowmax}(S_{12}),\; m_{11})$$ $$P_{12} = \exp(S_{12} - m_{12})$$ $$l_{12} = l_{11} \cdot e^{m_{11} - m_{12}} + \text{rowsum}(P_{12})$$ $$O_{12} = \text{diag}(e^{m_{11} - m_{12}}) \cdot O_{11} + P_{12} V_2$$

$j = 3$ and $j = 4$ follow the same pattern. Each step stays in SRAM.

6. After the inner loop ($j = 4$). Normalize and compute the log-sum-exp:

$$O_1 = \text{diag}(l_{14})^{-1} \cdot O_{14} \quad (2 \times 128)$$ $$L_1 = m_{14} + \log(l_{14}) \quad (2,)$$

Write $O_1$ and $L_1$ to HBM. Two writes for the entire row block. Advance $i \to 2$, reinitialize, repeat.

7. Return. After $i = 4$:

$$O = \begin{bmatrix} O_1 \\ O_2 \\ O_3 \\ O_4 \end{bmatrix} \;\; (8 \times 128) \qquad L = \begin{bmatrix} L_1 \\ L_2 \\ L_3 \\ L_4 \end{bmatrix} \;\; (8,)$$

The Correction Factor

Look at the $l$ and $O$ updates in step 5. Both have $e^{m_{11} - m_{12}}$ in them. That is the correction factor.

When a new block brings a larger max, everything we accumulated before is wrong. It was built under the old max, not the real one. We scale it down before adding the new piece. Since $m_\text{old} \leq m_\text{new}$, the exponent is zero or negative, so $e^{m_\text{old} - m_\text{new}} \leq 1$. The old output shrinks. The new one comes in at the right scale. End result is the same as if we knew the global max from the start.

One more thing. $m$ has shape $(2,)$, so each row in the block has its own running max. The correction is per-row:

$$\text{diag}\!\left(e^{m_\text{old} - m_\text{new}}\right) = \begin{bmatrix} e^{m_{\text{old},0} - m_{\text{new},0}} & 0 \\ 0 & e^{m_{\text{old},1} - m_{\text{new},1}} \end{bmatrix}$$

Row 0 and row 1 can have different maxima, so they get corrected by different amounts. This is online softmax, now applied block by block.

That is the forward pass. We return $O = [O_1, O_2, O_3, O_4]$ of shape $(8 \times 128)$, the full attention output, and $L = [L_1, L_2, L_3, L_4]$ of shape $(8,)$, the log-sum-exp values. $L$ is not part of the attention result. It is saved specifically for the backward pass.

Phase 2: Backward Pass

Four things to understand before the algorithm makes sense. This section also assumes you know how tensors are laid out in memory. If you do not, read this first.

Gradient and the Jacobian

The word derivative means different things depending on the shape of the output.

When $f: \mathbb{R}^N \to \mathbb{R}$ has a scalar output, the derivative is a gradient. It is a vector of partial derivatives, one per input. The update step is a dot product.

When $f: \mathbb{R}^N \to \mathbb{R}^M$ has a vector output, the derivative is a Jacobian. It is an $M \times N$ matrix $J$ where $J_{jk} = \partial y_j / \partial x_k$. The update step is a matrix multiply: $\Delta y = J \cdot \Delta x$.

$$J = \begin{bmatrix} \partial y_1/\partial x_1 & \cdots & \partial y_1/\partial x_N \\ \vdots & \ddots & \vdots \\ \partial y_M/\partial x_1 & \cdots & \partial y_M/\partial x_N \end{bmatrix}$$

Softmax takes a vector and outputs a vector. So its derivative is a Jacobian. That is why backward through softmax is the hard part.

Chain Rule Through a Matrix Multiply

Take a linear layer $y = xW$ with $x \in \mathbb{R}^{1 \times N}$ and $W \in \mathbb{R}^{N \times M}$. If we already know $dY$, the gradient of the loss with respect to the output, the chain rule gives us the gradients with respect to the input and the weights:

$$\frac{\partial \ell}{\partial x} = dY \cdot W^T \qquad (1 \times N)$$ $$\frac{\partial \ell}{\partial W} = x^T \cdot dY \qquad (N \times M)$$

The attention forward pass computes $O = PV$. Same pattern. Given $dO$:

$$dV = P^T \cdot dO \qquad (N \times d)$$ $$dP = dO \cdot V^T \qquad (N \times N)$$

We get $dV$ and $dP$ directly. The hard one is $dS$, because it has to go back through softmax.

The Softmax Jacobian

Softmax for row $i$ is $P_{ij} = e^{S_{ij}} / Z_i$ where $Z_i = \sum_l e^{S_{il}}$. Every output $P_{ij}$ depends on every input $S_{ik}$ through $Z_i$. So we need $\partial P_{ij} / \partial S_{ik}$ for all $j$ and $k$. That is the full row Jacobian.

When $j = k$, $P_{ij}$ shows up in both the numerator and $Z_i$:

$$\frac{\partial P_{ij}}{\partial S_{ij}} = \frac{e^{S_{ij}}}{Z_i} - \frac{e^{S_{ij}} \cdot e^{S_{ij}}}{Z_i^2} = P_{ij} - P_{ij}^2 = P_{ij}(1 - P_{ij})$$

When $j \neq k$, $P_{ij}$ only sees $S_{ik}$ through $Z_i$:

$$\frac{\partial P_{ij}}{\partial S_{ik}} = -\frac{e^{S_{ij}} \cdot e^{S_{ik}}}{Z_i^2} = -P_{ij}\,P_{ik}$$

Put it together and the full Jacobian for row $i$ is $J_i = \text{diag}(P_i) - P_i P_i^T$:

Now apply the chain rule. $dS_i = dP_i \cdot J_i$. Expand entry $j$ using the two cases above:

$$dS_{ij} = \sum_k dP_{ik}\,J_{kj} = dP_{ij} \cdot P_{ij}(1-P_{ij}) - P_{ij}\!\sum_{k \neq j} dP_{ik}\,P_{ik}$$ $$= P_{ij}\!\left[dP_{ij} - \underbrace{\sum_k dP_{ik}\,P_{ik}}_{D_i}\right]$$ $$\boxed{dS_{ij} = P_{ij}(dP_{ij} - D_i)}$$

where $D_i = \text{rowsum}(dP_i \odot P_i)$.

The D Shortcut

To compute $dS$ we need $dP$. That is the $N \times N$ matrix $dO \cdot V^T$. Storing it costs $O(N^2)$ memory. Same problem as standard attention all over again.

But look at $D_i$. It is just the dot product of two rows. We do not need the whole $dP$ matrix to get it. Plug in $dP_{ik} = \sum_m dO_{im}\,V_{km}$ where $m$ is the head dimension and watch what happens:

$$D_i = \sum_k dP_{ik}\,P_{ik} = \sum_k \left(\sum_m dO_{im}\,V_{km}\right) P_{ik}$$ $$= \sum_m dO_{im} \underbrace{\left(\sum_k P_{ik}\,V_{km}\right)}_{O_{im}}$$ $$\boxed{D_i = \text{rowsum}(dO_i \odot O_i)}$$

$dO$ comes from upstream. $O$ was already written to HBM during the forward pass. Both are sitting there. We never have to store $dP$ at all.

The Algorithm

This is the actual FlashAttention-2 backward pass algorithm from the paper. It looks dense but do not be overwhelmed. We will explain every line.

Let's take an example: $N = 8$ tokens, $d = 128$ embedding dim, block size $B_r = B_c = 2$, one head. $T_r = T_c = 4$.

At this point the forward pass is done. Sitting in HBM we have $O$ and $L$. PyTorch computed $dO$, the gradient of the loss with respect to the output, and handed it to the kernel. $dO$ is shape $(N \times d) = (8, 128)$.

1. Divide $Q$, $K$, $V$ into blocks the same way as the forward pass. $Q \to Q_1, Q_2, Q_3, Q_4$ each $(2 \times 128)$. Same for $K$ and $V$.

2. Divide $O$, $dO$, $L$ the same way. $O \to O_1 \ldots O_4$ each $(2 \times 128)$. $dO \to dO_1 \ldots dO_4$ each $(2 \times 128)$. $L \to L_1 \ldots L_4$ each $(2,)$.

3. Initialize $dQ = 0$ of shape $(8 \times 128)$ and split into $dQ_1 \ldots dQ_4$ each $(2 \times 128)$. $dK$ and $dV$ are initialized later, one block at a time inside the outer loop.

4. Compute $D = \text{rowsum}(dO \odot O)$ of shape $(8,)$. This is the $D$ we derived earlier, one scalar per token. Write it to HBM and split into $D_1 \ldots D_4$ each $(2,)$. This is the only step that touches the full matrix before the loops start.

5. Outer loop over $j = 1$ to $4$. This is the first big difference from the forward pass. The forward pass looped over $Q$ on the outside. The backward pass loops over $K$ and $V$ on the outside. The reason is that $dK_j$ and $dV_j$ accumulate over all $i$ and can be finalized before a single write to HBM. $dQ_i$ cannot, so it needs a read-modify-write on every inner step.

6. Load $K_j$ and $V_j$ from HBM to SRAM.

7. Initialize $dK_j = 0$ and $dV_j = 0$, both $(2 \times 128)$, in SRAM.

8. Inner loop over $i = 1$ to $4$.

9. Load $Q_i$, $O_i$, $dO_i$, $L_i$, $D_i$ from HBM to SRAM.

10. Recompute the score tile.

$$S_{ij} = Q_i K_j^T \quad (2 \times 128) \cdot (128 \times 2) \to (2 \times 2)$$

11. Recompute the attention weight tile using the saved log-sum-exp.

$$P_{ij} = \exp(S_{ij} - L_i) \quad (2 \times 2)$$

$S$ and $P$ were never stored during the forward pass. We get them back here from $Q_i$, $K_j$, and $L_i$. This is the recomputation trick that keeps memory $O(N)$ instead of $O(N^2)$.

12. Update $dV_j$. From $O = PV$, the gradient is $dV = P^T \cdot dO$.

$$dV_j = dV_j + P_{ij}^T \cdot dO_i \quad (2 \times 2)^T \cdot (2 \times 128) \to (2 \times 128)$$

This stays in SRAM and accumulates across all $i$.

13. Compute $dP_{ij}$. From $O = PV$, the gradient is $dP = dO \cdot V^T$.

$$dP_{ij} = dO_i \cdot V_j^T \quad (2 \times 128) \cdot (128 \times 2) \to (2 \times 2)$$

14. Compute $dS_{ij}$ using the softmax Jacobian result from earlier.

$$dS_{ij} = P_{ij} \odot (dP_{ij} - D_i) \quad (2 \times 2)$$

$D_i$ is a $(2,)$ vector. It broadcasts across the columns of $dP_{ij}$, subtracting a different value from each row.

15. Update $dQ_i$. From $S = QK^T$, the gradient is $dQ = dS \cdot K$.

$$dQ_i = dQ_i + dS_{ij} \cdot K_j \quad (2 \times 2) \cdot (2 \times 128) \to (2 \times 128)$$

$dQ_i$ accumulates across the outer $j$ loop, so it cannot stay in SRAM. Each inner step reads $dQ_i$ from HBM, adds the new contribution, and writes it back.

16. Update $dK_j$. From $S = QK^T$, the gradient is $dK = dS^T \cdot Q$.

$$dK_j = dK_j + dS_{ij}^T \cdot Q_i \quad (2 \times 2)^T \cdot (2 \times 128) \to (2 \times 128)$$

This also stays in SRAM. Steps 12 and 16 are symmetric: $dV_j$ accumulates $P_{ij}^T dO_i$ and $dK_j$ accumulates $dS_{ij}^T Q_i$, both across the inner $i$ loop, both staying in SRAM until the inner loop is done.

17. Inner loop ends. Repeat steps 9 to 16 for $i = 2, 3, 4$.

18. Write $dK_j$ and $dV_j$ to HBM. One write each, after all $i$ have contributed. Advance $j \to 2$ and repeat from step 6.

Why the Outer Loop Flips

In the forward pass the outer loop is over $i$ (Q blocks). Each $O_i$ accumulates across the inner $j$ loop entirely in SRAM, then gets one write to HBM after the inner loop ends. That works because $O_i$ is independent across $i$.

In the backward pass the same logic applies to $dK_j$ and $dV_j$. They accumulate across the inner $i$ loop in SRAM and get one write to HBM after the inner loop ends at step 18. So the outer loop goes over $j$.

$dQ_i$ is the problem. It accumulates across $j$, meaning each new outer iteration adds a new contribution to the same $dQ_i$. It cannot stay in SRAM because a different $K_j$ and $V_j$ are loaded each outer step and SRAM is overwritten. So $dQ_i$ lives in HBM and gets a read-modify-write at every inner step (step 15). That is the one unavoidable cost.

Count the HBM writes. $dK$ and $dV$ each do $T_c = 4$ writes total, one per outer step. $dQ$ does $T_r \times T_c = 16$ round-trips. There is no arrangement that avoids this. $dQ$ accumulates over the outer loop by definition, so it always needs more trips than $dK$ and $dV$. The outer-$j$ structure is the choice that keeps the other two gradients cheap.

19-20. Return. After $j = 4$:

PyTorch takes these and handles the rest of the backward pass through the projection layers.

That is the backward pass. We never stored $S$, $P$, or $dP$ at any point. $S$ and $P$ were recomputed tile by tile from $Q_i$, $K_j$, and the saved $L_i$. $D$ was computed in one pass before the loops. $dP$ was never stored at all. Its only role was inside the $dS$ formula, and that reduced down to $D_i$ which we already had. No $N \times N$ matrix ever exists.

Putting It Together

The forward pass is one kernel. It reads $Q$, $K$, $V$ from HBM, does everything in SRAM, and writes $O$ and $L$ back. Two HBM passes. Standard attention needed five for the same result, and an $N \times N$ matrix sitting in memory the whole time.

The backward pass is two kernels. The first sweeps through $dO$ and $O$ once to compute $D$ and writes it to HBM. The second does everything else and writes $dQ$, $dK$, $dV$. No $N \times N$ matrix at any point.

FA1, FA2, FA3

Everything we walked through is FlashAttention-2. It is worth being clear about what changed across the three versions.

FA1 had the same core idea: fused kernel, tiling, online softmax, no $N \times N$ matrix. But it looped over key blocks on the outside and query blocks on the inside, the opposite of FA2. That meant the GPU threads doing the work had to share and merge partial results with each other on every step. FA2 flipped the loop order so each thread can own a complete output row and never has to talk to the others. The math is the same. FA2 is roughly 2x faster on the same hardware just from this change.

FA3 targets the H100 specifically. The H100 can load the next tile from memory while it is still computing the current one. FA3 is built around this: it keeps the chip busy the whole time instead of having it wait between loads. It also adds FP8 support, which doubles throughput at the cost of some precision. If you are on an A100 or anything older, FA3 does nothing for you. It is the same algorithm, just written for hardware most people do not have yet.

The Code

Everything covered in this post is implemented in Triton. This was my first time writing Triton code, so I do not feel confident enough to walk through the implementation here. A tutorial might come later. For now, the code is on GitHub if you want to read it alongside this post.

ayyp1/Flash-Attention

Conclusion

This took me weeks. Not days. Weeks of reading the same paper again, going back to the math, filling pages of notes, getting confused, sleeping on it, and starting over. There were parts where I thought I understood something and then sat down to write it and realized I did not. Writing forced the gaps to show.

I wrote this for myself first. I wanted something I could come back to six months from now and actually follow. If I had to explain it, that meant I had to understand it. That is the real reason this exists.

But I also wrote it because I remember how it felt looking for a resource that explained the math without skipping the hard parts. Most explanations either stay at the surface or assume you already know the internals. I wanted something in between. Something honest. This is my attempt at that.

If you are somewhere in the middle of this and it is not clicking yet, that is normal. It did not click for me for a long time either. The notes in the appendix are proof. Keep going.

And if you find something wrong, I want to know. I have tried to be careful but I am not an expert. This is one person trying to understand something hard and writing it down. If you have feedback on the content, the writing, or anything else, I would love to hear it. Reach me at ashishkumar.pokharel@gmail.com.

References

The biggest source for understanding the math in this post was Umar Jamil's lecture. He derives and codes Flash Attention from first principles in Triton, and his explanation of the inside mathematics is the clearest I have seen. If you want to see all of this explained out loud and then implemented, watch this first.

A note on Umar Jamil: there are people who understand things deeply and keep it to themselves, and there are people who understand things deeply and spend hours making sure others can too. Umar is the second kind. The way he breaks down the math, builds it piece by piece, and then codes it from scratch is something I have not seen many people do at that level. I am genuinely grateful that he puts this out for free. This post would not exist without that lecture.

• Umar Jamil — Flash Attention derived and coded from first principles with Triton
• Dao et al. (2022) — FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
• Dao (2023) — FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning
• Shah et al. (2024) — FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision
• Milakov & Gimelshein (2018) — Online normalizer calculation for softmax
• Vaswani et al. (2017) — Attention Is All You Need

Appendix: The Notes

These are the actual notes I wrote while trying to understand Flash Attention. Every page is a moment where something clicked or something confused me. I am putting them here because I think there is something honest about showing the mess behind a finished post. If you are someone sitting with paper trying to figure this out, this is what it looked like on my end too.