About this series

Transformers power modern language models, but explanations often jump from a colorful attention diagram to dense matrix notation. This series takes the middle route: clear intuition, enough mathematics to verify the mechanism, and practical implications for model architecture and GPU inference. It assumes familiarity with programming, arrays and basic algebra, not previous experience building Transformers.

Reading path. Parts 1–4 build one attention calculation. Parts 5–8 assemble a Transformer block, including positional information and feed-forward networks (FFNs). Parts 9–13 connect the math to memory-efficient GPU execution and text generation. Parts 14–15 close with training, implementation and evidence-based performance analysis.

Editorial conventions.  $T$ denotes token positions, $d_{model}$ the number of hidden features per token, $d_k$ Query/Key width, $B$ batch size and $H$ attention-head count. Dimensions are usually written as rows × columns. Approximate costs are analytic estimates, not measurements of a particular GPU.

1. Part 1 – Why a Word Needs Its Neighbors
2. Part 2 – Q, K and V: The Matching Engine
3. Part 3 – Why Divide Attention Scores by $\sqrt{d_k}$?
4. Part 4 – Softmax, Causal Masks and a Three-Token Walkthrough
5. Part 5 – Multiple Heads and Positional Information
6. Part 6 – Normalization and Residual Connections
7. Part 7 – Trace One Complete Transformer Block
8. Part 8 – The Feed-Forward Network Is Not Just Another Matrix
9. Part 9 – The GPU Cost of Attention
10. Part 10 – FlashAttention Without the Mystery ← You are here
11. Part 11 – Why Predicting the Next Token Is a Learning Problem
12. Part 12 – Prefill, Decode and the First Generated Token
13. Part 13 – KV Cache, Grouped-Query Attention and PagedAttention
14. Part 14 – What Gradients and Distributed GPUs Really Change
15. Part 15 – Build, Verify and Profile Your Own Attention

Overview

Imagine computing a weighted average over a dataset too large to hold in a small, fast workspace. You could materialize every intermediate weight, write it to slow storage, then read it again. Or you could process the data in blocks, keeping only the summary needed to update the final result.

FlashAttention uses this second strategy for dense attention. Its key achievement is reducing intermediate global-memory traffic by operating on smaller tiles and maintaining a numerically stable running softmax summary. It does not simply discard low-weight tokens or make exact dense attention linear in sequence length.

Why softmax complicates tiling

For a single Query, ordinary attention produces

where $s_j$ is the Query-Key score for Key $j$ and $v_j$ is its Value vector. You cannot finalise this expression from only the first few Keys, because later Keys affect both numerator and denominator.

But you can preserve a small running summary. After processing some valid scores, keep: $m$, the largest score seen so far; $\ell$, the sum of exponentials measured relative to that maximum; and $u$, the correspondingly weighted sum of Value vectors. Here $m$ and $\ell$ are scalars for one Query, while $u$ has the Value-vector width.

The current attention output is $u/\ell$ after all relevant entries have been processed. The numerical maximum matters because subtracting it keeps exponentials from unnecessarily overflowing.

The surprising part: an old summary can be rescaled

Suppose the next tile contains a score larger than the old maximum. All previous exponentials were measured relative to the old reference. Rather than reloading all old scores, we can convert the old totals to the new reference by multiplying them by one common factor.

If the new maximum is $m\prime=max\left(m,s\right)$ for one incoming score , then

The same rescaling factor $e^{m-m\prime}$ applies to both accumulated denominator and numerator. If we rescale only one, the output becomes incorrect. If the new score does not increase the maximum, the factor equals one and the old totals stay unchanged.

This update is a mathematical identity in exact real arithmetic. It is not an approximation to softmax. Hardware implementations may nevertheless differ slightly in rounded results because floating-point reductions occur in different orders.

Tile merging is the same idea at a larger scale

A GPU can process separate Key/Value tiles and compute one summary for each. To merge two disjoint tiles A and B, choose the larger maximum $m=max\left(m_A,m_B\right)$ and define $\alpha=e^{m_A-m}$ and $\beta=e^{m_B-m}$. Then

The denominator and Value numerator are now expressed in the same scale. The final output $u/\ell$ matches the ordinary weighted average in real arithmetic. The core insight here is rescale, then merge.

What changes on the GPU?

A traditional materialized implementation may store the entire $T\times T$ score or probability matrix in HBM. A tiled implementation loads a portion of Q/K/V, does calculations in registers and shared memory, updates running summaries and ultimately writes the output. It avoids some expensive intermediate transfers and reduces the footprint of those buffers.

FlashAttention remains dense attention: each eligible Query-Key interaction can still matter. Reducing intermediate I/O does not mean eliminating all pairwise arithmetic. Sparse and approximate approaches change a different aspect of the computation.

Numerical precision still matters

FP16 and BF16 each use two bytes per value but differ in representable range and resolution. BF16 has a wider exponent range, while FP16 generally resolves more closely spaced values near one. Fused kernels may accumulate selected reductions in FP32 or in another supported internal format. Therefore “same mathematical answer” does not always mean bit-for-bit identical floating-point output.

For a simple caution, directly exponentiating a score of 12 in FP16 can overflow, while subtracting the row maximum turns scores $\left[12,11\right]$ into $\left[0,-1\right]$, whose exponentials fit. The resulting stable softmax is approximately $\left[0.7311,0.2689\right]$.

Takeaway. FlashAttention is primarily an I/O-aware execution method for ordinary dense attention. Online softmax is the mathematical device that makes tiled, stable accumulation possible without keeping every intermediate score and probability in HBM.

Next: Part 11 changes perspective from the attention kernel to how language models learn to predict the next token.