{"slug": "understanding-flashattention-pt-1-personal-notes", "title": "Understanding FlashAttention Pt 1: Personal Notes", "summary": "A technical handbook on FlashAttention explains that the algorithm achieves wall-clock speedups by reducing data movement between GPU high-bandwidth memory (HBM) and on-chip SRAM rather than by approximating attention, using tiling, online softmax, and recomputation. The handbook traces the evolution from FlashAttention-1 through FlashAttention-4 and covers forward and backward passes, IO complexity, and current framework behavior. It stresses that dense FlashAttention remains an exact attention algorithm, with only small floating-point rounding differences from reordered operations.", "body_md": "# Understanding FlashAttention Pt 1: Personal Notes\n\n## 0. Introduction\n\n**How IO-Aware Attention Makes Transformers Faster Without Approximating Attention**\n\nThe mechanism, in three words: **Tiling + Online Softmax + Recomputation**. Everything in this handbook is elaboration on that summary.\n\nA technical handbook 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.\n\n### 0.1 How to Read This Handbook\n\nThis [handbook](https://drive.google.com/file/d/1CLyK-9Cflcvi3fRl3qAHyzYvwJFjCyVg/view) was inspired by this [tweet](https://x.com/techNmak/status/2098057360908685358). Before the fix, here is what the standard attention implementation looks like. Load $Q, K, V \\in \\mathbb{R}^{N \\times d}$ in HBM, then:\n\n1. Read $Q, K$ from HBM, compute $S$, write $S$ to HBM.\n2. Read $S$ from HBM, compute $P$, write $P$ to HBM.\n3. Read $P, V$ by blocks from HBM, compute $O$, write $O$ to HBM.\n4. Return $O$.\n\nWhat stands out to me is the number of round trips to HBM. Every intermediate value — $S$, $P$, $O$ — has to be written out and read back. That is the problem FlashAttention is solving.\n\nThe handbook itself frames the subject as easiest to understand when three different questions are kept separate:\n\n- What mathematical function is being computed? For dense attention, the target remains ordinary scaled dot-product attention.\n- How much arithmetic does that function require? Dense all-pairs query-key scoring remains quadratic in sequence length.\n- How does the implementation move data through the GPU memory hierarchy? This is where FlashAttention changes the algorithmic execution dramatically.\n\nThe central lesson I take from this framing is that wall-clock speed is not determined by FLOP count alone. An algorithm can perform essentially the same mathematical work, or even recompute intermediate values, and still run faster because it moves far less data to and from high-bandwidth memory.\n\n**Core distinction.** Dense FlashAttention is an *exact* attention algorithm: it does not replace softmax attention with a low-rank, sparse, kernelized, or approximate formula. “Exact” refers to the mathematical attention computation. Floating-point kernels can still differ by small rounding effects because operations are reordered.\n\nThe 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.\n\n### 0.2 Notation\n\nFor one attention head, let\n\n$$\nQ \\in \\mathbb{R}^{N_q \\times d}, \\quad K \\in \\mathbb{R}^{N_k \\times d}, \\quad V \\in \\mathbb{R}^{N_k \\times d_v}.\n$$\nFor self-attention, typically $N_q = N_k = N$. The scaled score matrix is\n\n$$\nS = \\frac{QK^T}{\\sqrt{d}} + B,\n$$\nwhere $B$ represents an optional additive mask or bias, and\n\n$$\nP = \\mathrm{softmax}\\_{\\mathrm{row}}(S), \\quad O = PV.\n$$\nThroughout, 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.\n\nThe practical difference I keep coming back to:\n\n| HBM | SRAM | \n|---|---|\n| slow | faster | \n| large | smaller | \n| off-chip | on-chip | \n\n**System problem:** where do all those intermediate values $(S, P, O)$ live while the GPU computes them?\n\nThat question — not the arithmetic — is what FlashAttention was built to answer.\n\n## Contents\n\n### Part 1: The Fundamental Problem\n\n### Part 2: The Mathematical Trick\n\n### Part 3: Putting the Mathematics onto the GPU\n\n### Part 4: Architectural Compatibility\n\n### Part 5: FlashAttention Evolution\n\n### Part 6: Using FlashAttention in Frameworks\n\n### Part 7: FlashAttention vs. Other Techniques\n\n### Part 8: Training vs. Inference\n\n### Part 9: Practical Engineering\n\n## Part 1: The Fundamental Problem\n\n### 1.1 What FlashAttention Actually Optimizes\n\nStart with the ordinary attention function:\n\n$$\nO = \\mathrm{softmax}\\left(\\frac{QK^T}{\\sqrt{d}} + B\\right) V.\n$$\nFlashAttention is **IO-aware**. My working definition:\n\nMinimize data movement between the different levels of GPU memory, rather than just trying to reduce the number of mathematical operations (FLOPs).\n\nThe 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.\n\nA textbook implementation often makes this look like three large operations:\n\nForm the score matrix $S$, apply row-wise softmax to obtain $P$, then multiply by $V$. Mathematically that is fine. On a GPU, however, writing a huge intermediate matrix to HBM and reading it back can be far more expensive than the equation suggests.\n\nFlashAttention’s central contribution is to make the algorithm IO-aware. It partitions the computation into tiles that fit in fast on-chip memory, streams blocks of $K$ and $V$, and maintains enough row-wise softmax state to produce the exact output without materializing the full $N \\times N$ attention matrix in HBM.\n\n**What changes:** the execution schedule, memory traffic, and stored intermediates.\n\n**What does not change:** the dense scaled-dot-product attention function being evaluated.\n\nThis distinction is why “FlashAttention is a faster kind of attention” can be misleading. It is better thought of as an algorithm and kernel family for evaluating attention efficiently on accelerators. A model can use causal masking, RoPE, MQA/GQA, or other attention features and still use a FlashAttention implementation underneath.\n\nThe original paper contrasts this approach with approximate attention methods that reduce arithmetic by changing the mathematical problem. Dense FlashAttention does not make that trade. The same paper also introduced a block-sparse extension, but that sparse extension is a different case because omitting blocks changes which interactions are computed.\n\n### 1.2 The Attention Equation Is Not the Implementation\n\nThe equation does not tell you where tensors live.\n\nA **standard naive attention implementation**:\n\n1. Calculate $S$, store $S$.\n2. Read $S$, calculate $P$, store $P$.\n3. Read $P$, calculate $O$.\n\nThe problem is the number of HBM round trips. A simple materializing implementation does:\n\n$$\nS \\leftarrow QK^T / \\sqrt{d}, \\quad P \\leftarrow \\mathrm{softmax}(S), \\quad O \\leftarrow PV.\n$$\nIf $S$ is written to HBM after the first matrix multiplication, read for softmax, $P$ is written back, and then $P$ is read again for the $PV$ multiplication, the GPU spends significant time moving an $N^2$ object through memory.\n\n**FlashAttention avoids repeatedly moving huge intermediates:**\n\n1. Calculate small $S$ tile, softmax tile, use tile with $V_j$, discard tile.\n2. Calculate next tile.\n\nThe key questions I ask when looking at any equation:\n\n1. How many operations are required?\n2. What data must move between memory levels to perform those operations?\n\nFlashAttention reuses the same dependencies. Blocks of $Q$, $K$, and $V$ are brought near the compute units, score tiles are produced and consumed locally, and only compact row-wise statistics plus the output need to persist across tiles.\n\n**Algorithmic lesson.** A computational graph is not a memory schedule. Writing $P = \\mathrm{softmax}(QK^T)$ on paper does not require an implementation to store all of $QK^T$ or $P$ in off-chip memory at once.\n\nThis idea generalizes beyond attention. Fused kernels, tiling, recomputation, and operator scheduling often trade a small amount of extra arithmetic for much less movement of large intermediates. On modern accelerators, that can be the right trade because matrix-multiply throughput has grown much faster than many other parts of the memory and execution hierarchy. FlashAttention-3 and -4 make that hardware dependence increasingly explicit.\n\n### 1.3 GPU Memory Hierarchy and Why IO Matters\n\nGPUs expose a hierarchy rather than one uniform pool of equally fast memory. The names and capacities vary by architecture, but the mental model is:\n\n- Registers / very local state\n- On-chip shared memory / SRAM\n- HBM / device memory\n\nHBM is large, but data must travel to the compute units. On-chip storage is much smaller, but reuse there is much cheaper. The original FlashAttention analysis models this asymmetry using HBM and SRAM and explicitly optimizes the number of transfers between them.\n\nMy breakdown of each level:\n\n**HBM:**\n\n- Relatively large\n- Relatively slower to access\n- Physically farther from individual compute operations\n- Stores model weights, $Q/K/V$, activations, large tensors\n\n**SRAM / shared memory:**\n\n- Much smaller\n- Faster\n- Cheaper reuse cost\n\n**Registers:**\n\n- Even smaller\n- More local\n\n**Why tiling helps.** Suppose $Q_i$ (a query tile) needs to interact with many $K/V$ tiles. Instead of constantly moving $Q_i$ back and forth, we can keep it close to the compute units while processing:\n\nSo one loaded $Q_i$ can participate in lots of computation. This is called **reuse**.\n\nThe purpose of tiling is not simply “make tensors smaller.” It is to **increase reuse** while a tile is resident on chip. A block of $Q$ can interact with multiple $K/V$ blocks before its partial softmax/output state is written back. Conversely, $K/V$ blocks can be streamed through query blocks according to the chosen schedule.\n\nA kernel becomes IO-aware when the placement and movement of data are part of the algorithm, rather than an afterthought left to a sequence of separately launched tensor operations.\n\nDo not turn this into a universal slogan that attention is always “memory-bound.” The bottleneck depends on sequence length, head dimension, dtype, mask pattern, GPU generation, forward vs. backward, and which kernel is running. FA3 and FA4 exist partly because, as hardware changed, the dominant bottlenecks changed too.\n\n### 1.4 Why Materializing $S$ and $P$ Is Expensive\n\nThe 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. One dense tensor with shape\n\n$$\n[1, 32, 8192, 8192]\n$$\ncontains $32 \\times 8192^2$ elements. At two bytes per element, that is\n\n$$\n32 \\times 8192^2 \\times 2 = 4{,}294{,}967{,}296 \\text{ bytes} = 4 \\text{ GiB}.\n$$\nThat is the size of one full score- or probability-like tensor for this example.\n\nConcretely, my arithmetic:\n\n- Batch = 1\n- Heads = 32\n- seq_len = 8192\n- dtype = FP16/BF16\n- 2 bytes/elem\n- tensor shape = $[1, 32, 8192, 8192]$\n- $(32)(8192)(8192)$ elem $\\times$ (2 bytes/elem) $= 4{,}294{,}967{,}296$ bytes $= 4$ GiB\n\nImagine reading and writing this much data.\n\nA naive decomposition may produce arrays of this scale at multiple stages. This does not mean every modern framework keeps both $S$ and $P$ alive simultaneously, and compiler fusion can already avoid some traffic. The example illustrates the fundamental problem: a dense $N^2$ intermediate is large enough that repeatedly writing and rereading it can dominate memory use and bandwidth.\n\n**Naive pipeline:**\n\n**FlashAttention pipeline:**\n\nFlashAttention avoids storing the full matrix in HBM. It forms score tiles, applies the softmax update while the tile is on chip, immediately uses those probabilities to accumulate the corresponding contribution from $V$, then discards the tile.\n\n**Important wording:** FlashAttention removes the need to materialize the full attention matrix as an off-chip intermediate. It does not remove the logical pairwise interactions required by dense attention.\n\n**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.*\n\nThis is also why the memory benefit is especially important during training, where naive autograd would otherwise want large intermediates for the backward pass.\n\n### 1.5 Dense Arithmetic Is Still Quadratic\n\nFlashAttention changes the memory schedule, not the mathematical function. The function is still dense attention. That means:\n\nFor self-attention with $N$ tokens and head dimension $d$, forming all query-key scores requires work proportional to $N^2 d$.\n\nMultiplying the probabilities by values adds another dense pairwise matrix multiplication of the same broad order. FlashAttention reorganizes these operations, but it does not stop evaluating the dense set of query-key interactions.\n\nSo three statements must not be confused:\n\n**Arithmetic complexity:** dense attention remains $\\mathcal{O}(N^2 d)$.\n**Large intermediate storage:** FlashAttention avoids an $\\mathcal{O}(N^2)$ materialized score/probability tensor in HBM.\n**HBM traffic:** the original paper proves a lower IO cost under its two-level memory model than standard materializing attention.\n\nThis is how a method can make much longer sequences practical without making long context “free.” Doubling $N$ still roughly quadruples the number of dense query-key pairs. FlashAttention mainly attacks the data-movement and memory-footprint side of that computation.\n\nIf you want to reduce the *number* of query-key pairs themselves, you need a different mathematical structure: for example sparsity, a local pattern, or a different attention formulation. Those choices can change model behavior and are conceptually separate from dense FlashAttention.\n\nWhen reporting speedups, always distinguish asymptotic arithmetic from measured runtime. A kernel can become several times faster at the same $\\mathcal{O}(N^2 d)$ complexity because the constant factors, occupancy, fusion, and memory traffic change dramatically.\n\nMy side-by-side summary:\n\n| Quality | Dense FlashAttention | \n|---|---|\n| Arithmetic | $\\mathcal{O}(N^2 d)$ | \n| Full attention intermediates | Avoid $\\mathcal{O}(N^2)$ storage | \n| HBM traffic | Reduced substantially | \n\n### 1.6 Memory-Efficient Exact Attention Predates FlashAttention\n\nIt would be historically inaccurate to say FlashAttention first discovered that exact attention can avoid quadratic memory.\n\nRabe and Staats showed before FlashAttention that attention need not require $\\mathcal{O}(N^2)$ memory with respect to sequence length. Their work gave exact memory-efficient algorithms while retaining quadratic time, and a practical accelerator implementation with subquadratic memory.\n\nFlashAttention’s distinct contribution was to turn memory efficiency into an explicit IO-aware GPU algorithm: tile the computation against the accelerator memory hierarchy, analyze HBM accesses, fuse the relevant operations, and show substantial wall-clock gains.\n\nOnline softmax also has an earlier lineage. Milakov and Gimelshein described an online recurrence that computes the classical stable softmax normalizer with fewer memory accesses. FlashAttention builds the same kind of running-max/running-normalizer idea into tiled attention, while also accumulating the value-weighted output.\n\nA precise lineage is therefore:\n\n- stable / online softmax provides a streaming normalization tool,\n- earlier memory-efficient attention shows exact attention need not store $N^2$ state,\n- FlashAttention co-designs tiling, softmax state, and GPU IO to make the approach fast in practice.\n\nThe lineage, in my shorthand:\n\n$$\n\\text{Stable softmax} \\to \\text{Online normalization} \\to \\text{Exact memory-efficient attention} \\to \\text{FlashAttention} \\to \\text{FlashAttention-2} \\to \\text{FlashAttention-3} \\to \\text{FlashAttention-4}\n$$\nRabe and Staats demonstrated exact memory-efficient attention with quadratic computation but subquadratic memory (the earlier work). There is also earlier work on online softmax by Milakov and Gimelshein.\n\nFlashAttention’s key contribution was to bring together:\n\n1. Tiling\n2. Online softmax\n3. Fused computation\n4. GPU memory hierarchy awareness\n5. IO complexity analysis\n\ninto a practical high-performance algorithm.\n\nThis distinction matters because “memory-efficient” does not automatically mean “IO-optimized for a particular hardware model.”\n\n## Part 2: The Mathematical Trick\n\n### 2.1 Tiling Queries, Keys, and Values\n\n**Most important.** This is the section where the trick lives.\n\nSuppose $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 everything at once, divide them into blocks. For example:\n\n$$\nQ: [Q\\_1, Q\\_2, Q\\_3], \\quad K, V: [K\\_1, V\\_1], [K\\_2, V\\_2], [K\\_3, V\\_3]\n$$\nFor one pair,\n\n$$\nS\\_{ij} \\doteq \\frac{Q\\_i K\\_j^T}{\\sqrt{d}} + B\\_{ij},\n$$\nwe process that score tile locally.\n\nThen, for each query block:\n\n1. Calculate scores\n2. Softmax them\n3. Multiply by $V_j$\n4. Update the running output\n5. Discard the tile\n\nThe softmax used at each tile is\n\n$$\n\\mathrm{softmax}(x\\_i) = \\frac{e^{x\\_i}}{\\sum\\_j e^{x\\_j}}.\n$$\n**Matrix multiplication is easy to tile** because it is fundamentally accumulation:\n\nso we can calculate pieces and add them.\n\n**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.\n\nInstead of forming all $N_q N_k$ scores at once, partition the matrices into blocks. For one query block $Q_i$ and one key/value block $(K_j, V_j)$, compute\n\n$$\nS\\_{ij} = \\frac{Q\\_i K\\_j^T}{\\sqrt{d}} + B\\_{ij}.\n$$\nThe score tile $S_{ij}$ is small enough to be processed near the compute units. Its contribution is folded into running row statistics and an output accumulator, then the tile can be discarded.\n\nConceptually, for each query block:\n\n- load a tile of $Q_i$,\n- stream compatible $K/V$ tiles,\n- compute one score tile,\n- update row-wise softmax state,\n- accumulate the corresponding $V$ contribution,\n- move to the next tile without writing a global $S$ or $P$ matrix.\n\nActual kernels choose tile shapes and loop order based on on-chip capacity, head dimension, GPU generation, causal structure, and work partitioning. The original FlashAttention IO analysis uses tile sizes derived from SRAM capacity $M$.\n\nTiling alone is not enough. Matrix multiplication tiles compose naturally because sums can be accumulated. Softmax couples every score in a row through a shared maximum and denominator, so we need a way to merge blocks without seeing the full row at once.\n\nThat is the key mathematical trick in the next chapters.\n\n### 2.2 Softmax Is the Difficult Part of Streaming\n\nA numerically stable softmax for one row $x_1, \\ldots, x_N$ uses\n\n$$\nm = \\max\\_j x\\_j, \\qquad \\ell = \\sum\\_j e^{x\\_j - m}, \\qquad p\\_j = \\frac{e^{x\\_j - m}}{\\ell}.\n$$\nThe subtraction by $m$ prevents overflow from large positive logits. But it seems to create a streaming problem: how can an early block be normalized if a later block may contain a larger maximum?\n\nThe answer is to retain sufficient statistics that can be **rescaled** when the maximum changes. Suppose the running state after earlier elements is $(m_{\\text{old}}, \\ell_{\\text{old}})$ and a new block has maximum $m_b$. Define\n\nEvery contribution accumulated under the old maximum can be converted to the new reference by multiplying it by\n\n$$\n\\alpha \\doteq e^{m\\_{\\text{old}} - m\\_{\\text{new}}}.\n$$\nThe new block is evaluated relative to the same $m_{\\text{new}}$.\n\nThis is not an approximation. It is the identity\n\n$$\ne^{x\\_j - m\\_{\\text{old}}}\\, e^{m\\_{\\text{old}} - m\\_{\\text{new}}} = e^{x\\_j - m\\_{\\text{new}}}.\n$$\n**Mathematical insight.** The running maximum is a change of numerical reference point. When that reference changes, previously accumulated exponentials can be rescaled exactly in real arithmetic rather than recomputed from scratch.\n\nThe online-normalizer recurrence predates FlashAttention and provides the mathematical basis for streaming stable softmax.\n\n**Worked example.** Start with two extreme values to see why the max subtraction matters:\n\nDirectly computing $e^{1000}$ overflows. Instead, set $m = 1000$, so\n\n$$\nx - m = [1000 - 1000, \\; 999 - 1000] = [0, -1],\n$$\nand $e^{0} = 1$, $e^{-1} \\approx 0.368$ are perfectly manageable.\n\nNow stream two blocks. Let Block 1 $= [2, 1]$ and Block 2 $= [4, 3]$. After Block 1, $m_{\\text{old}} = 2$. Block 1 is evaluated at $m_{\\text{old}}$ as $[e^{0}, e^{-1}]$. When Block 2 arrives with $m_b = 4$,\n\n$$\nm\\_{\\text{new}} = \\max(2, 4) = 4, \\qquad \\alpha = e^{2 - 4} = e^{-2}.\n$$\nRescaling Block 1 under the new reference:\n\n$$\n[\\alpha e^{0}, \\; \\alpha e^{-1}] = [e^{-2}, \\; e^{-3}].\n$$\nBlock 2 at the new reference:\n\n$$\n[e^{4 - 4}, \\; e^{3 - 4}] = [e^{0}, \\; e^{-1}].\n$$\nPutting it all together, with $x = [2, 1, 4, 3]$:\n\n$$\n\\mathrm{softmax}(x) = \\frac{[e^{-2}, \\; e^{-3}, \\; e^{0}, \\; e^{-1}]}{e^{-2} + e^{-3} + e^{0} + e^{-1}}.\n$$\nThe streaming computation for Block 1 is\n\n$$\n\\ell = \\sum\\_j e^{x\\_j - m} = e^{0} + e^{-1} = 1 + e^{-1},\n$$\ngiving normalized outputs $[e^{0}, e^{-1}] / (1 + e^{-1})$. For Block 2, the running normalizer is\n\n$$\n\\ell\\_{\\text{new}} = \\alpha\\, \\ell\\_{\\text{old}} + \\sum\\_{j \\in \\text{Block 2}} e^{x\\_j - m\\_{\\text{new}}} = e^{-2}(1 + e^{-1}) + e^{0} + e^{-1},\n$$\nand the combined numerator is $[\\alpha e^{0}, \\alpha e^{-1}, e^{0}, e^{-1}] = [e^{-2}, e^{-3}, e^{0}, e^{-1}]$, matching the full softmax computed in one shot.\n\n### 2.3 Online Softmax from First Principles\n\nProcess scalar logits $x_1, x_2, \\ldots$ one at a time. Initialize\n\n$$\nm\\_0 = -\\infty, \\qquad \\ell\\_0 = 0.\n$$\nAfter observing $x_j$, update\n\n$$\nm\\_j = \\max(m\\_{j-1}, x\\_j), \\qquad \\ell\\_j = \\ell\\_{j-1} e^{m\\_{j-1} - m\\_j} + e^{x\\_j - m\\_j}.\n$$\nMilakov and Gimelshein show that this produces the same stable softmax normalizer while requiring fewer passes over the input than the conventional safe-softmax procedure.\n\nFor attention, we also need the weighted value sum. Introduce an unnormalized accumulator $a$:\n\n$$\na = \\sum\\_j e^{x\\_j - m} v\\_j.\n$$\nWhen the maximum changes from $m$ to $m’$, rescale both $\\ell$ and $a$ by $e^{m - m’}$. Then add the new exponentials and value contributions under the new reference. At the end,\n\n$$\no = \\frac{a}{\\ell}.\n$$\nMy working version of the update:\n\n$$\na \\doteq \\sum\\_j e^{x\\_j - m} v\\_j\n$$\nWhen the maximum changes:\n\n$$\na\\_{\\text{old}} \\to a\\_{\\text{old}}\\, e^{m\\_{\\text{old}} - m\\_{\\text{new}}}\n$$\n \n$$\n\\ell\\_j \\to \\ell\\_{j-1}\\, e^{m\\_{j-1} - m\\_j} + e^{x\\_j - m\\_j}\n$$\nThen add the new value contributions. At the end:\n\n$$\no = \\frac{a}{\\ell}.\n$$\nFor attention, $x_j$ is not a fixed input vector stored in advance. Each block of logits is generated on demand from a matrix product $Q K_j^T / \\sqrt{d}$ plus mask/bias terms. The online recurrence lets the kernel consume that block immediately.\n\nThe same idea works row by row and block by block, which is what makes a tiled exact softmax-attention forward pass possible.\n\n**Mathematical insight II.** We do not need the entire probability vector. We only need enough information to reconstruct its contribution to the final output.\n\n### 2.4 The Blockwise Merge Recurrence\n\nFor one query row, suppose the running state after some key blocks is\n\n$$\n(m, \\ell, a),\n$$\nwhere $m$ is the maximum score seen so far, $\\ell$ is the stable softmax denominator under that maximum, and $a \\in \\mathbb{R}^{d_v}$ is the unnormalized value accumulator.\n\nFor a new score block $s \\in \\mathbb{R}^b$ with matching values $V_b \\in \\mathbb{R}^{b \\times d_v}$, let\n\n$$\nm\\_b = \\max(s), \\qquad m' = \\max(m, m\\_b),\n$$\n \n$$\n\\alpha = e^{m - m'}, \\qquad p = e^{s - m'}.\n$$\nThen update\n\n$$\n\\ell' = \\alpha \\ell + \\sum\\_j p\\_j,\n$$\n \n$$\na' = \\alpha a + p^T V\\_b,\n$$\n \n$$\nm \\leftarrow m', \\qquad \\ell \\leftarrow \\ell', \\qquad a \\leftarrow a'.\n$$\nFinally,\n\n$$\no = a / \\ell.\n$$\nThe mean of the running state:\n\n- $m$ = max score seen so far\n- $\\ell$ = stable softmax denominator, exponential sum (normalizer)\n- $a$ = unnormalized value accumulator\n\nThe blockwise recurrence in my own notation:\n\n- For one query row, maintain $(m, \\ell, a)$.\n- Suppose the next score block is $s = [s_1, s_2, \\ldots, s_b]$ with corresponding values $V_b \\in \\mathbb{R}^{b \\times d_v}$.\n- First calculate the block maximum: $m_b = \\max(s)$, then update the global maximum: $m’ = \\max(m, m_b)$.\n- Define: $\\alpha = e^{m - m’}$, $p = e^{s - m’}$.\n- Then update: $\\ell’ = \\alpha \\ell + \\sum_j p_j$, $a’ = \\alpha a + p^T V_b$.\n- Then set $m \\leftarrow m’$, $\\ell \\leftarrow \\ell’$, $a \\leftarrow a’$.\n- Finally at the end: $o = a / \\ell$.\n\nFor a block of query rows, $m$ and $\\ell$ become row-wise vectors and $a$ becomes a matrix. Masks can be applied to the score tile before the exponentials, with masked positions contributing zero probability.\n\nThe recurrence is the algebraic reason tile boundaries do not change the dense softmax result. A different tiling changes the order of floating-point operations, but not the intended real-arithmetic function.\n\nFlashAttention’s published algorithms express equivalent running-max / running-normalizer / output updates in block form.\n\n**This recurrence is the heart of tiled exact attention.**\n\n### 2.5 A Complete Numerical Example\n\nConsider one already-scaled, unmasked attention row\n\n$$\ns = [2, 1, 4, 3]\n$$\nwith two-dimensional values\n\n$$\nv\\_1 = [1, 0], \\quad v\\_2 = [0, 1], \\quad v\\_3 = [2, 0], \\quad v\\_4 = [0, 2].\n$$\nThe global maximum is 4, so stable unnormalized weights are\n\n$$\n[e^{-2}, e^{-3}, e^{-1}, e^{-1}] \\approx [0.135335, 0.049787, 1, 0.367879].\n$$\nTheir sum is\n\n$$\n\\ell \\approx 1.55300179,\n$$\nand full softmax gives\n\n$$\nO \\approx [1.37497284, 0.50582424].\n$$\nNow process two blocks. For $[2, 1]$:\n\n$$\nm\\_1 = 2, \\qquad \\ell\\_1 = 1 + e^{-1} = 1.36787944, \\qquad a\\_1 = [1, e^{-1}].\n$$\nFor the second block $[4, 3]$, the new maximum is 4, so\n\n$$\n\\alpha = e^{2 - 4} = e^{-2}.\n$$\nThen\n\n$$\n\\ell\\_2 = \\alpha \\ell\\_1 + 1 + e^{-1} = 1.55300179,\n$$\n \n$$\na\\_2 = \\alpha a\\_1 + [2, 2 e^{-1}] \\approx [2.13533528, 0.78554595].\n$$\nTherefore\n\n$$\na\\_2 / \\ell\\_2 \\approx [1.37497284, 0.50582424],\n$$\nmatching the full-row computation.\n\nThe old block was not revisited. Its contribution was merely rescaled when a larger maximum appeared.\n\n**Tying everything together.** Block 1 was never recomputed when block 2 revealed a larger maximum; we rescaled the statistics from block 1. This is what makes streaming possible.\n\nMy step-by-step computation:\n\nThe score row: $s = [2, 1, 4, 3]$. Values: $v_1 = [1, 0]$, $v_2 = [0, 1]$, $v_3 = [2, 0]$, $v_4 = [0, 2]$.\n\n**Blocks:** $[2, 1]$ and $[4, 3]$.\n\n**Block 1.** Scores: $[2, 1]$. Maximum: $m_1 = 2$. Stable exponentials:\n\nSo\n\n$$\n\\ell\\_1 = 1 + e^{-1} \\approx 1.36788.\n$$\nValue accumulator:\n\n$$\na\\_1 = 1 \\cdot [1, 0] + e^{-1} \\cdot [0, 1] = [1, e^{-1}].\n$$\n**Block 2.** Scores: $[4, 3]$. $m_b = 4$. $m’ = \\max(2, 4) = 4$. $\\alpha = e^{m - m’} = e^{2 - 4} = e^{-2}$.\n\n**Output.**\n\n## Part 3: Putting the Mathematics onto the GPU\n\n### 3.1 FlashAttention Forward Pass\n\nA useful conceptual forward pass is:\n\n1. Partition $Q$ into query-row tiles and $K, V$ into key/value tiles.\n2. For each query tile, initialize row-wise running maxima, normalizers, and output accumulators.\n3. Load a key/value tile and form the local score tile $Q_i K_j^T / \\sqrt{d}$.\n4. Apply causal/local masks or additive biases that belong to this tile.\n5. Compute the tile maximum, update the running maximum, and rescale previous state.\n6. Exponentiate the current tile relative to the updated maximum.\n7. Update the denominator and the value-weighted output accumulator.\n8. Continue until every required key tile has been contributed.\n9. Normalize the accumulator row-wise and write the output.\n\nThe original FlashAttention algorithm chooses block sizes so the relevant tiles and state fit in on-chip SRAM, reducing trips to HBM.\n\nAn educational implementation can reproduce the algebra in a few lines of PyTorch, but such code is not a high-performance FlashAttention kernel. Production implementations depend on GPU-specific tiling, thread/warp scheduling, asynchronous copies, tensor-core instructions, and other low-level details.\n\nThe essential algorithmic idea is independent of one CUDA kernel: generate a score tile, consume it immediately through online softmax and $V$ accumulation, and never materialize the complete score/probability matrix in HBM.\n\n**Minimal educational PyTorch.**\n\n``` python\ndef tiled_attention(q, k, v, block=128):\n    scale = 1 / math.sqrt(q.shape[-1])\n    n_q, n_kv = q.shape[0], k.shape[0]\n    O = torch.zeros_like(q)\n\n    for i in range(0, n_q, block):\n        q_blk = q[i:i+block]\n        m = torch.full((q_blk.shape[0],), -float('inf'))\n        l = torch.zeros((q_blk.shape[0],))\n        a = torch.zeros((q_blk.shape[0], v.shape[-1]))\n\n        for j in range(0, n_kv, block):\n            k_blk = k[j:j+block]\n            v_blk = v[j:j+block]\n            s = (q_blk @ k_blk.T) * scale\n            m_new = torch.maximum(m, s.max(dim=-1).values)\n            p = torch.exp(s - m_new[:, None])\n            alpha = torch.exp(m - m_new)\n            l = alpha * l + p.sum(dim=-1)\n            a = alpha[:, None] * a + p @ v_blk\n            m = m_new\n\n        O[i:i+block] = a / l[:, None]\n\n    return O\n```\n\nConceptually, for each KV block:\n\n- scores = $Q \\times K_{block}^T$\n- update running softmax, $\\ell$\n- update output accumulator, $a$\n\nThis isn’t a high-performance FlashAttention kernel. Real implementations additionally exploit:\n\n- GPU-specific tiling\n- Registers\n- Shared memory\n- Warp scheduling\n- Tensor cores\n- Async copies\n- Specialized intrinsics\n- Occupancy optimization\n\nThis distinction will become very important for FA-2/3/4.\n\nLater FlashAttention generations keep this semantic structure while changing how work is scheduled on newer hardware.\n\nMy step list for the forward pass:\n\n1. Load $Q_i$.\n2. Initialize $(m, \\ell, a) = (-\\infty, 0, 0)$ — $m = -\\infty$, $\\ell = 0$, $a = 0$.\n3. Load a $K_j, V_j$ tile.\n4. Compute the local score tile $S_{ij} = Q_i K_j^T / \\sqrt{d}$.\n5. Apply masks/biases $B_{ij} / M_{ij}$.\n6. Find the tile maximum and update the running maximum $\\max(m, m_b)$.\n7. Compute exponentials relative to the new maximum.\n8. Update $\\ell$ and $a$.\n9. Move to the next $K, V$ tile until every tile is completed/covered.\n10. Normalize the accumulator row-wise $O = a / \\ell$, then write $O$.\n\n### 3.2 Why Dense FlashAttention Is Exact\n\nDoes FlashAttention approximate attention? For dense FlashAttention, **no**. The target remains\n\nMy checklist for why this is exact:\n\n- No Q-K pairs are intentionally removed.\n- No low-rank approximation is introduced.\n- No alternative kernelized attention function replaces softmax.\n- This algorithm simply processes the same interactions in blocks.\n\nBut *exact* has a qualification. Let me illustrate.\n\nSuppose two implementations compute $a + b + c$:\n\n- **A:** $(a + b) + c$\n- **B:** $a + (b + c)$\n\nIn exact mathematics,\n\n$$\n(a + b) + c = a + (b + c).\n$$\nBut in floating points, in some cases,\n\n$$\n(a + b) + c \\neq a + (b + c).\n$$\nTiling changes the reduction order. Fusion can change the rounding behavior. Therefore:\n\n$$\n\\text{exact algorithm} \\neq \\text{bitwise identical output}.\n$$\nThis distinction becomes important in numerical testing.\n\nDense FlashAttention is called exact because the target function remains\n\n$$\n\\mathrm{softmax}(QK^T / \\sqrt{d} + B) V.\n$$\nTiling does not delete query-key pairs. Online softmax does not replace the exponential or normalization with another function. The block recurrence simply changes the order in which sufficient statistics are accumulated.\n\nThree caveats keep the word *exact* precise:\n\n**Floating point is finite precision.** Reordering additions and reductions can produce small numerical differences from another implementation. PyTorch explicitly warns that SDPA backends can differ because floating-point operations are fused and ordered differently.\n**Dropout is stochastic during training.** Comparing two runs bit-for-bit requires matching random behavior in addition to mathematical attention semantics.\n**Sparse variants are different.** The original FlashAttention paper also presents block-sparse FlashAttention, which omits blocks and is therefore an approximate/sparse variant relative to full dense attention.\n\n“Exact” does not mean “bitwise identical to every reference kernel.” It means the algorithm is not intentionally changing dense softmax attention to reduce the mathematical work.\n\nThis distinction matters when evaluating numerical tests. A sensible tolerance depends on dtype, accumulation order, sequence length, and backend rather than requiring binary identity.\n\n### 3.3 IO Complexity: What the Theorem Actually Says\n\nLet me get more theoretical.\n\nThe original FlashAttention paper analyzes a two-level memory model with HBM and on-chip SRAM of size $M$. Under the paper’s assumptions, including head dimension $d$ and\n\n$$\nd \\leq M \\leq Nd,\n$$\nstandard materializing attention requires approximately\n\n$$\n\\Theta(Nd + N^2)\n$$\nHBM accesses, whereas FlashAttention requires\n\n$$\n\\Theta\\left(\\frac{N^2 d^2}{M}\\right)\n$$\nHBM accesses under the specified regime.\n\nThe exact theorem has assumptions, so do not interpret this as “FlashAttention always moves exactly this many bytes.” It is an asymptotic result for a particular memory model.\n\n**Why does larger SRAM help?** Suppose you have more on-chip memory. You can fit larger tiles. Larger tiles mean:\n\n- more data stays resident\n- more reuse\n- fewer HBM reloads\n\nSo increasing $M$ can decrease the amount of HBM traffic. This is one reason GPU architecture matters so much to FlashAttention performance.\n\nThe original FlashAttention paper analyzes a two-level memory model with HBM and on-chip SRAM of size $M$. Under the paper’s assumptions, including head dimension $d$ and $d \\leq M \\leq Nd$, standard materializing attention requires $\\Theta(Nd + N^2)$ HBM accesses, whereas FlashAttention requires $\\Theta(N^2 d^2 / M)$ HBM accesses. The paper also proves an optimality result over a range of SRAM sizes in this model.\n\nSeveral details matter:\n\n- These are IO-complexity results in a particular memory model, not a universal byte count for every GPU.\n- The quantity counts movement of scalar elements/words between the modeled memory levels, up to asymptotic factors.\n- It is not the arithmetic complexity. Dense attention still performs $O(N^2 d)$ work.\n- Increasing usable on-chip memory $M$ enables more reuse and reduces modeled HBM traffic.\n\nA common incorrect summary is “FlashAttention reduces attention IO from $O(N^2)$ to $O(N)$.” The linear quantity is the large auxiliary memory footprint with respect to sequence length, not the general HBM-access expression above.\n\nIntuitively, more usable on-chip memory lets larger working tiles stay resident and be reused for more attention work before data must be reloaded from HBM. In the theorem’s regime, that increased reuse is why the HBM-access bound decreases as $M$ grows.\n\n### 3.4 Memory Complexity: Linear Auxiliary State, Not Linear Compute\n\nAnother important distinction: linear auxiliary state is not linear compute.\n\nA naive attention implementation may create $\\mathcal{O}(N^2)$ attention intermediates. FlashAttention does not. Instead it maintains:\n\n- $Q/K/V$ tiles\n- row-wise $m$\n- row-wise $\\ell$\n- output accumulator $a$\n\nThe auxiliary attention state scales roughly linearly with sequence length, but linear memory is not linear computation. The computation is still\n\n$$\n\\mathcal{O}(N^2 d).\n$$\n**Example.** For $N = 10{,}000$, there are approximately\n\nQ-K interactions. FlashAttention does not make those disappear. However, it prevents you from needing a gigantic intermediate containing all of them.\n\nWhy do FlashAttention papers often say memory becomes linear instead of quadratic?\n\nThe inputs and output already contain $O(Nd)$ elements. A materializing dense attention implementation additionally creates $O(N^2)$ score/probability state. FlashAttention avoids storing those full matrices and retains only tiled working state plus row-wise statistics. As a result, the extra memory associated with the attention operation scales linearly with sequence length rather than quadratically.\n\nFor training, the difference is especially important because a straightforward backward pass might otherwise save a full probability matrix $P$. FlashAttention instead saves compact information such as output and row-wise normalization statistics, then recomputes score/probability tiles in backward.\n\n**Linear memory does not imply linear runtime.** The number of dense query-key interactions is still quadratic in $N$.\n\nAlso avoid claiming 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 is specifically changing how the attention computation manages its intermediates.\n\nOne reason FlashAttention and activation checkpointing can coexist: both trade recomputation for reduced stored state, but at different scopes of the training graph.\n\n### 3.5 Causal Masking and Tile Skipping\n\nConsider autoregressive attention. Token $i$ cannot attend to future token $j > i$.\n\nThe attention matrix looks like\n\n$$\n\\begin{bmatrix} \\checkmark & \\times & \\times \\\\ \\checkmark & \\checkmark & \\times \\\\ \\checkmark & \\checkmark & \\checkmark \\end{bmatrix}\n$$\nwith mask\n\n$$\nM = \\begin{bmatrix} 0 & -\\infty & -\\infty \\\\ 0 & 0 & -\\infty \\\\ 0 & 0 & 0 \\end{bmatrix}.\n$$\nApplying the mask before softmax gives\n\n$$\n\\mathrm{softmax}(S') = \\begin{bmatrix} \\text{value} & 0 & 0 \\\\ \\text{value} & \\text{value} & 0 \\\\ \\text{value} & \\text{value} & \\text{value} \\end{bmatrix}.\n$$\nWith tiled attention, we can skip entire tiles:\n\n1. Entire blocks above the diagonal can be entirely skipped.\n2. Blocks below the diagonal are fully valid.\n3. Diagonal blocks require element-level masking.\n\nThis is important because we are no longer doing useless work for obviously invalid future positions.\n\nIn causal self-attention, query position $i$ may not attend to future key positions $j > i$. A materialized mask would be another $N \\times N$ object, but an efficient tiled kernel can reason about the geometry of each tile.\n\nFor square self-attention:\n\n- blocks strictly above the causal boundary are fully masked and need not contribute,\n- blocks strictly below the boundary are fully valid,\n- only blocks intersecting the diagonal require element-level causal masking.\n\nThis avoids computing many invalid tiles in a causal kernel. FlashAttention-2 explicitly exploits causal structure, while current implementations also define precise alignment rules for unequal query/key lengths.\n\nMask semantics are an API detail that must not be guessed. For example, current PyTorch SDPA treats a Boolean `attn_mask` value of True as a position that *participates* in attention, while other PyTorch mask APIs use different conventions.\n\n**Connection to local attention.** Suppose each token can only attend to 128 nearby tokens. Then many tiles can also be skipped. But now we have changed the mathematical attention pattern: that is no longer undistributed dense attention. FlashAttention can be the kernel executing the local pattern, but the model itself is now doing sparse/local attention.\n\nLocal/sliding-window attention can similarly skip tiles outside the permitted window. But once the model intentionally restricts which pairs are attended, the *model’s attention pattern* is sparse/local. FlashAttention can be the kernel used to execute that pattern, but it is no longer the same mathematical problem as unrestricted dense attention.\n\n### 3.6 Backward Pass: Recompute Instead of Save\n\nThis is vital for training.\n\nDuring the forward pass, we do not save the entire $P$ matrix ($N \\times N$). We recompute it.\n\n**Forward:**\n\n1. Store compact info such as:\n    \n  - output\n  - row-wise normalization statistics\n2. Do not store $P \\in \\mathbb{R}^{N \\times N}$.\n\n**Backward:**\n\n1. For each tile:\n    \n  - a) Recompute $QK^T$\n  - b) Reconstruct the local probability values\n  - c) Calculate gradients\n  - d) Discard the tile\n2. This trades more computation for less memory traffic/storage.\n\nConceptually: forward: don’t store $P$. Backward: recompute $P$ tile-by-tile.\n\nThe gradient rehashing for $P$:\n\n$$\ndV \\mathrel{+}= P^T dO\n$$\n \n$$\ndP = dO \\cdot V^T\n$$\n \n$$\ndS = P \\odot (dP - D\\_i[:, \\text{None}])\n$$\nand then:\n\n$$\ndQ \\mathrel{+}= dS \\cdot K / \\sqrt{d}\n$$\n \n$$\ndK \\mathrel{+}= dS^T \\cdot Q / \\sqrt{d}\n$$\nwith appropriate masking. Where\n\n$$\nD\\_i[:, \\text{None}] = \\sum\\_j P\\_{ij}\\, dP\\_{ij}\n$$\nis the column/vector broadcast across each row.\n\nA naive training implementation can save $P = \\mathrm{softmax}(S)$ for backward. FlashAttention avoids keeping that $N^2$ tensor in HBM. Instead, it stores compact row-wise normalization information and recomputes score/probability tiles when gradients are needed.\n\nFor\n\n$$\nS = QK^T / \\sqrt{d}, \\quad P = \\mathrm{softmax}(S), \\quad O = PV,\n$$\na useful row-wise identity is\n\n$$\nD\\_i = \\sum\\_r (dO\\_i)\\, O\\_{ir} = \\sum\\_j P\\_{ij}\\, dP\\_{ij}.\n$$\nAfter recomputing a tile of $P$, the local derivatives can be expressed as\n\n$$\ndV \\mathrel{+}= P^T dO, \\quad dP = dO\\, V^T, \\quad dS = P \\odot (dP - D\\_i[\\text{:}, \\text{None}]),\n$$\nthen\n\n$$\ndQ \\mathrel{+}= dS\\, K / \\sqrt{d}, \\quad dK \\mathrel{+}= dS^T\\, Q / \\sqrt{d}.\n$$\nMasks imply zero probability/gradient contribution for masked entries. FA2 stores a row-wise log-sum-exp quantity that allows the probability tile to be reconstructed stably from recomputed scores.\n\nBackward recomputation is intentional. It spends extra matrix-multiply work to avoid reading and writing a giant probability tensor, which can be a favorable trade on GPUs.\n\nThis is the **memory-compute tradeoff** applied to the backward pass.\n\n### 3.7 Why More FLOPs Can Still Be Faster\n\nThis is one of the biggest lessons from FlashAttention.\n\nNormally we think: fewer FLOPs $\\Rightarrow$ faster. But on GPUs, that’s incomplete.\n\nDifferent operations have radically different throughput. Tensor cores are exceptionally good at matrix multiplication. Other operations have different performance profiles. These operations include:\n\n1. exponentials\n2. reductions\n3. synchronization\n4. shared-memory operations\n5. memory transfers\n\nTherefore, sometimes doing **extra arithmetic** is worthwhile if it eliminates expensive memory traffic.\n\nFor example, consider two options:\n\n- **A:** Compute $\\to$ write huge $P$ to HBM $\\to$ read $P$ $\\to$ compute.\n- **B:** Compute $\\to$ discard $\\to$ recompute later.\n\nOption B performs more arithmetic. But it might be faster because it avoids moving a giant tensor through HBM.\n\nThis is a fundamental ML systems principle:\n\nThe cost of a FLOP depends on what hardware executes it and what data movement surrounds it.\n\nIt is tempting to assume that fewer arithmetic operations always imply lower latency. Accelerator performance breaks that intuition regularly.\n\nMatrix multiplication maps exceptionally well to tensor cores. HBM traffic, synchronization, shared-memory traffic, exponentials, reductions, and kernel launch boundaries can be comparatively expensive.\n\nFA2 makes this contrast explicit: one of its goals is to reduce non-matmul FLOPs, because those operations do not enjoy the same throughput as tensor-core GEMMs. The paper also improves work partitioning so more of the GPU is occupied.\n\nFA3 goes further on Hopper by overlapping matrix multiplication, softmax, and data movement using asynchronous hardware features. FA4 responds to Blackwell, where tensor-core throughput increased faster than some other resources, making exponentials and shared-memory traffic relatively more important.\n\nA better performance question is not merely “How many FLOPs?” but “Which operations, on which units, with what data movement, reuse, parallelism, and synchronization?”\n\nThis is the systems lesson that makes FlashAttention important beyond attention itself: hardware efficiency often comes from co-designing mathematical scheduling with the memory/execution hierarchy.\n\n## Part 4: Architectural Compatibility\n\n### 4.1 MHA, MQA, and GQA Compatibility (Shared K/V Heads)\n\nFlashAttention is not tied to standard MHA.\n\nLet\n\n$$\nQ \\in \\mathbb{R}^{B \\times N\\_q \\times H\\_q \\times d}, \\quad K, V \\in \\mathbb{R}^{B \\times N\\_k \\times H\\_{kv} \\times d}.\n$$\nThree head-sharing regimes:\n\n- **MHA:** $H_q = H_{kv}$ $\\Rightarrow$ every query head has its own K/V head.\n- **MQA:** $H_{kv} = 1$ $\\Rightarrow$ all query heads share 1 K/V head.\n- **GQA:** $1 < H_{kv} < H_q$ $\\Rightarrow$ several query heads share K/V heads.\n\nVital distinction:\n\n1. MHA/MQA/GQA defines the architecture and head sharing.\n2. FlashAttention defines efficient execution of the attention computation.\n\nThus MHA/MQA/GQA and FlashAttention can coexist. Current implementations impose shape constraints: the number of query heads must be divisible by the number of KV heads:\n\n$$\nH\\_q \\bmod H\\_{kv} = 0.\n$$\nFlashAttention does not require every architecture to have the same number of query and key/value heads.\n\nFor ordinary MHA, $H_q = H_{kv}$. In MQA, $H_{kv} = 1$. In GQA, $1 < H_{kv} < H_q$. Current Dao-AILab kernels support MQA/GQA by passing fewer KV heads than query heads, with the requirement that the number of query heads be divisible by the number of KV heads.\n\nThe kernel still evaluates attention between each query head and its assigned KV head. Head sharing changes the architecture and KV-memory footprint. FlashAttention changes how the resulting attention operation is executed.\n\n**Orthogonal concepts:**\n\n- MQA/GQA: how heads share K/V projections,\n- RoPE: how position transforms Q/K,\n- FlashAttention: how attention is scheduled and computed efficiently.\nThese can be used together.\n\nCurrent PyTorch SDPA also exposes `enable_gqa`. Its documentation labels GQA support experimental and imposes backend- and tensor-shape constraints, so production code should follow the exact version’s documentation rather than assuming universal fused-kernel support.\n\n### 4.2 Variable Lengths, Local Attention, and Dropout\n\nReal systems aren’t always: same sequence length + dense attention + no dropout.\n\nProduction implementations may support:\n\n- variable-length sequences\n- causal attention\n- sliding-window attention (SWA)\n- dropout\n- MQA/GQA\n- ALiBi-style bias\n- KV-cache decoding (including optional RoPE handling)\n\nBut an important distinction: these are implementation capabilities, not fundamental properties of the FlashAttention mathematical idea.\n\nFeature support depends on:\n\n- GPU\n- CUDA/ROCm backend\n- dtype\n- head dimension\n- mask\n- library version\n- kernel generation\n\nPyTorch issue: `scaled_dot_product_attention` applies dropout according to the supplied `dropout_p`, so eval code should explicitly use `0.0` when dropout should be disabled.\n\nProduction attention rarely consists only of equal-length dense sequences with no dropout. Current FlashAttention implementations support a broader feature set, but these are implementation capabilities, not properties of the mathematical idea itself.\n\nThe current Dao-AILab repository documents kernels/interfaces for features including:\n\n- variable-length sequences,\n- causal attention,\n- local/sliding-window attention,\n- dropout in training-oriented interfaces,\n- MQA/GQA,\n- ALiBi-style score bias in relevant interfaces,\n- specialized incremental-decoding paths with KV cache, including optional RoPE handling.\n\nFeature support differs by CUDA/ROCm backend and evolves over time. For example, the repository documents separate NVIDIA and AMD backends with different implementation details and support matrices.\n\nDropout deserves a practical warning. Current PyTorch `scaled_dot_product_attention` always applies dropout according to its `dropout_p` argument, so callers must pass `0.0` during evaluation when dropout should be disabled.\n\nDo not infer feature support from the name “FlashAttention.” Check the exact library, kernel generation, device, dtype, head dimension, mask/bias, and training/inference path that will actually run.\n\n## Appendix: Source Section Mapping\n\n| Source § | Hierarchical | \n|---|---|\n| 1-6 | 1.1-1.6 | \n| 7-11 | 2.1-2.5 | \n| 12-18 | 3.1-3.7 | \n| 19-20 | 4.1-4.2 | \n| 21-30 | 5.1-5.10 | \n| 31-32 | 6.1-6.2 | \n| 33-34 | 7.1-7.2 | \n| 35 | 8.1 | \n| 36-37 | 9.1-9.2 |", "url": "https://wpnews.pro/news/understanding-flashattention-pt-1-personal-notes", "canonical_source": "https://chizkidd.github.io//2026/09/13/flashattention/", "published_at": "2026-09-14 14:01:04+00:00", "updated_at": "2026-09-14 14:18:49.195432+00:00", "lang": "en", "topics": ["artificial-intelligence", "machine-learning", "ai-research", "ai-infrastructure"], "entities": ["FlashAttention", "FlashAttention-1", "FlashAttention-4", "HBM", "SRAM"], "alternates": {"html": "https://wpnews.pro/news/understanding-flashattention-pt-1-personal-notes", "markdown": "https://wpnews.pro/news/understanding-flashattention-pt-1-personal-notes.md", "text": "https://wpnews.pro/news/understanding-flashattention-pt-1-personal-notes.txt", "jsonld": "https://wpnews.pro/news/understanding-flashattention-pt-1-personal-notes.jsonld"}}