What is FlashAttention?

In this article

1. Before you start
2. Why interviewers ask this
3. FlashAttention from first principles
4. Worked example
5. Trade-offs and edge cases
6. How to answer it in an interview
7. Common mistakes
8. Follow-up questions
9. Sources

Before you start

Why interviewers ask this

FlashAttention is the cleanest example of an optimisation that changes speed without changing the answer, which is why it suits large language model (LLM) depth rounds. The question takes two forms: a conceptual one ("why is standard attention memory-bound and how does FlashAttention fix it") and an implementation one ("write single-head flash attention in NumPy").

A weak answer says "it is a faster attention" or "it approximates attention". A strong answer says the maths is unchanged, names the real bottleneck as memory traffic between HBM and SRAM, and explains the running max and sum that make a block-by-block softmax exact.

FlashAttention from first principles

In plain words

Tiling

Tiling is the technique of cutting a large matrix computation into small blocks that fit in fast memory, finishing each block's work before loading the next.

Memory-bound

Memory-bound is a property of a computation whose speed is limited by how fast data moves to and from memory, not by how fast the chip can do arithmetic.

Think of marking exams with a small desk and a large warehouse of papers. Carry each paper from the warehouse to the desk, mark it and carry it back, and you spend the day walking. Take one stack, finish all you can with it, then fetch the next, and you spend the day marking.

The analogy breaks here. Marking each paper is independent, but a softmax needs a total over a whole row, so a stack-by-stack method needs extra bookkeeping to get that total right.

In one precise sentence, FlashAttention reorders the same arithmetic so that the large intermediate matrix stays in on-chip memory.

Step 1: Why standard attention is memory-bound

Standard attention is memory-bound because it writes the $N\times N$ score matrix to HBM, reads it back for the softmax, writes the weights, and reads them again for the product with the values. The FlashAttention paper describes an A100 with 40 to 80 gigabytes (GB) of HBM at roughly 1.5 to 2.0 terabytes (TB) per second and 192 kilobytes (KB) of SRAM per streaming multiprocessor across 108 of them at roughly 19 TB per second. The SRAM is about ten times faster and many orders of magnitude smaller.

The input is the matrices of queries, keys and values for $N$ tokens. The matrix of scores has $N^2$ entries per head. At 8,192 tokens and 2 bytes per entry (fp16) that is 128 MiB for one head, and 4 GiB for one layer of 32 heads¹. Moving this to and from HBM costs more time than the multiplications that produced it. The paper's Theorem 2 puts standard attention at $\Theta (Nd+N^2)$ HBM accesses, with $d$ the head dimension.

Step 2: How online softmax keeps tiled attention exact

Online softmax keeps a running maximum $m$ and a running sum $\ell$ for each row; when a new block raises the maximum from $m$ to $m^{\prime }$, it multiplies the old sum by $e^{m-m^{\prime }}$. Milakov and Gimelshein introduced this online softmax in 2018 to cut memory accesses, and FlashAttention builds on it.

Take six illustrative scores, 1.0, 3.0, −2.0, 0.5, 2.5 and 4.0, as two blocks of three. After block one the running maximum is 3.0 and the running sum is 1.1421. Block two has a maximum of 4.0, so the maximum rises to 4.0, the old sum is multiplied by $e^{3-4}$, and the new block's sum is added. The final sum is 1.6735, which equals the sum computed on the whole row at once¹.

The same rescaling is applied to the partly built output. The input is a block of scores and the previous $m$, $\ell$ and output; the output is updated values that equal the exact softmax result after the last block. In the six-score example with value vectors, the largest difference from the full computation was $\approx 5.6\times {10}^{-17}$, which is floating-point rounding¹.

Step 3: Tiling keeps every block in SRAM

FlashAttention loops over blocks of keys and values in the outer loop and blocks of queries in the inner loop, loading each block into SRAM and computing its scores there. With keys and values in the outer loop, each query block's partial output and its running max $m$ and sum $\ell$ are read from HBM and written back after every key and value block.

That is $O(Nd)$ traffic per pass and no $N\times N$ score matrix. FlashAttention-2 swaps the loops, so the output and statistics stay on chip until one final write. Algorithm 1 in the paper sets the key block size to $\lceil M/(4d)\rceil$ and the query block size to the smaller of that and $d$, where $M$ is the SRAM size in elements.

With an illustrative $M=\mathrm{98,304}$ fp16 elements (192 KB) and $d=128$, the key block holds 192 rows and the query block 128 rows¹. Theorem 2 gives $\Theta (N^2{d}^2/M)$ HBM accesses. At 8,192 tokens the ratio of the two leading terms is about 6.1, so on that rough count, with constants ignored, the tiled version moves about a sixth of the data¹.

Step 4: The backward pass recomputes instead of storing

For training, FlashAttention stores the output and the statistics $m$ and $\ell$, then recomputes the score and weight tiles in on-chip memory during the backward pass. The paper says that even with more arithmetic, recomputation speeds up the backward pass because it reduces HBM accesses.

The trade is more floating-point work for less memory traffic. Because the chip has arithmetic to spare and little bandwidth, the trade wins.

The extra memory is the pair $m$ and $\ell$ per row, which is linear in $N$. At 8,192 tokens that is 65,536 bytes per head in fp32¹. The paper states the algorithm needs only $O(N)$ additional memory beyond inputs and output while still doing $O(N^{2}d)$ arithmetic.

1 · Memory-bound2 · Online softmax3 · Tiling4 · RecomputeReplay

Step-by-step description

1. Standard attention moves the N by N scores to HBM and back.
2. A running max and sum make block-wise softmax exact.
3. Tiles stay in SRAM; O, m and l are written back, never the scores.
4. The backward pass recomputes tiles instead of storing scores.

• Standard attention writes the N by N score matrix to HBM and reads it back.
• Online softmax keeps a running maximum and a running sum for each row.
• FlashAttention loads tiles of queries, keys and values into SRAM and writes the output and statistics back, never the score matrix.
• The backward pass recomputes the score and weight tiles from the stored output and statistics.

Worked example

Every row is $N^2$ entries at 2 bytes each¹. Quadrupling the tokens multiplies the size by 16, so at long contexts the matrix soon stops fitting in GPU memory. FlashAttention never builds it.

Text description of the diagram

• At 8,192 tokens one head's fp16 score matrix is 128 MiB, which is 0.125 GiB.
• At 131,072 tokens one head's fp16 score matrix is 32 GiB.
• Quadrupling the tokens multiplies the size by 16.

The exactness check is a full tiled run. With 256 tokens, head dimension 16 and 64-row blocks on random data, the tiled output differed from standard attention by at most $\approx 5.6\times {10}^{-16}$. Standard attention holds 65,536 scores at once; the tiled version holds at most 4,096 in a tile¹.

Text description of the diagram

• Standard attention holds 65,536 scores at once for 256 tokens.
• The tiled version holds at most 4,096 scores in a tile.
• The tiled output differed from standard attention by at most rounding error.

The sources are the FlashAttention, FlashAttention-2 and FlashAttention-3 abstracts. These are the authors' own benchmark claims on their hardware, not our measurements.

Trade-offs and edge cases

FlashAttention does not make attention cheaper in arithmetic. It does the same $O(N^{2}d)$ work, so compute still grows with the square of the sequence length; the saving is memory size and memory traffic. It also does not shrink the KV cache that inference stores for every past token, which is a separate memory cost covered in LLM inference.

The gains depend on hardware and data type. The FlashAttention-2 paper reports 25 to 40% of peak arithmetic throughput for version 1, and the FlashAttention-3 paper reports 35% utilisation of an H100 for version 2. The repository states version 2 supports fp16 and bf16, which is why the choice between them in fp16 versus bf16 matters here, and head dimensions up to 256.

Text description of the diagram

• Version 2 supports Ampere, Ada and Hopper GPUs with fp16 and bf16.
• Version 3 is in beta for Hopper GPUs with CUDA 12.3 or higher.
• Version 3 supports FP16 and BF16 forward and backward and an FP8 forward pass.
• ROCm builds cover AMD Instinct MI200x, MI250x, MI300x and MI355x accelerators and RDNA 3 and 4 GPUs.

Fused kernels have input limits. PyTorch's scaled_dot_product_attention selects among FlashAttention-2, a memory-efficient kernel and a plain C++ implementation, and when a fused kernel cannot run it emits a warning listing the reasons. Unusual masks or biases can force the slower path. The PyTorch default is enough unless you need a feature it lacks.

How to answer it in an interview

The spoken structure is: name the bottleneck as memory traffic, say tiling plus online softmax, then state that it is exact and that compute stays quadratic.

Common mistakes

• "FlashAttention approximates attention to go faster" (wrong). The output is the same up to floating-point rounding (right).
• "It reduces the number of floating-point operations" (wrong). Arithmetic stays $O(N^{2}d)$ and the backward pass adds recomputation (right: it reduces memory traffic).
• "It removes the KV cache" (wrong). The cache stores past keys and values for generation and is a separate cost (right).
• "Tiling alone is enough" (wrong). Without the running max and sum a block-wise softmax would normalise each block separately and give the wrong weights (right).
• "It works on any GPU with any data type" (wrong). Support depends on the GPU family, the vendor backend and the version (right).

Follow-up questions

What is FlashAttention in an LLM?

It is an exact attention kernel that cuts memory traffic. PyTorch's built-in scaled_dot_product_attention ships a FlashAttention-2 kernel (PyTorch 2.14 documentation, read October 2026), so models that call it can use FlashAttention with no separate install. It makes long contexts practical because it avoids storing the quadratic score matrix. See the transformers hub for where attention sits in the model.

How do I install FlashAttention?

For version 2 the repository gives pip install flash-attn --no-build-isolation, and it lists Compute Unified Device Architecture (CUDA) 12.0 or above, PyTorch 2.2 or later, and the packaging, psutil and ninja packages, with Linux as the supported platform as of October 2026. Often no install is needed, because PyTorch's scaled_dot_product_attention already includes a FlashAttention-2 kernel on CUDA.

Which GPUs support FlashAttention?

It depends on vendor and version. As of October 2026 the repository lists, for NVIDIA CUDA, version 2 on Ampere, Ada and Hopper GPUs such as the A100, GeForce 3090 and 4090 cards and the H100 (for Turing GPUs such as the T4 and 2080 it points to a separate flash-attention-turing project that supports a core subset of features), version 3 in beta on Hopper (H100 and H800) with CUDA 12.3 or higher, and version 4 for Hopper and Blackwell (B200); FlashAttention-3 supports FP16 and brain floating point 16 (BF16) forward and backward and an 8-bit floating point (FP8) forward pass.

For Advanced Micro Devices (AMD) ROCm, it lists AMD Instinct MI200x, MI250x, MI300x and MI355x accelerators and RDNA (AMD's Radeon graphics architecture) 3 and 4 GPUs on a Composable Kernel backend, and a Triton backend for CDNA (AMD's Instinct compute architecture) and RDNA with fp16, bf16 and fp32.

What is FlashAttention 2?

FlashAttention-2 keeps the exact tiled method but cuts non-matmul arithmetic, parallelises over sequence length as well as batch and heads, and splits work between warps to cut shared-memory traffic (Dao, July 2023). The author reports about twice the speed of version 1, reaching 50 to 73% of the theoretical maximum arithmetic throughput on A100 and up to 225 TFLOPs per second per GPU in GPT-style training.

What is FlashAttention 3?

FlashAttention-3 targets Hopper GPUs by overlapping data movement with computation through asynchrony and by using FP8 low-precision arithmetic. The authors report a 1.5 to 2.0 times speedup on H100 GPUs and about 2.6 times lower numerical error than a baseline FP8 attention.

Does FlashAttention make attention linear in compute?

No. The arithmetic is still quadratic in sequence length; only the extra memory becomes linear. Methods that cut compute, such as sparse or windowed attention, change the maths, which FlashAttention does not.

Sources

1. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (arXiv 2205.14135) — arXiv (Dao, Fu, Ermon, Rudra, Re), 27 May 2022
2. FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning (arXiv 2307.08691) — arXiv (Dao), 17 Jul 2023
3. FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision (arXiv 2407.08608) — arXiv (Shah, Bikshandi, Zhang, Thakkar, Ramani, Dao), 11 Jul 2024
4. Online normalizer calculation for softmax (arXiv 1805.02867) — arXiv (Milakov and Gimelshein, NVIDIA), 8 May 2018
5. Dao-AILab/flash-attention repository README (read 2026-10-24) — GitHub, Dao-AILab, 24 Oct 2026
6. torch.nn.functional.scaled_dot_product_attention, PyTorch 2.14 documentation — PyTorch (page shows version only; read 2026-10-24), 24 Oct 2026