Cruxly › AI engineering › Inference Serving

FlashAttention and Why Attention Is a Memory-Traffic Problem

a 134 MB score matrix per head, an online softmax you can do on paper, and which of FlashAttention, PagedAttention and GQA fixes which bottleneck

Prompt: "Attention does the same maths in every implementation, yet FlashAttention is two to four times faster and lets us run 8K contexts without running out of memory. Explain why, show me online softmax on a tiny example, tell me what FlashAttention does for decode, and tell me how it relates to PagedAttention and grouped-query attention."

40 minutes first. Size the score matrix for one head at 8,192 tokens, estimate the memory traffic of the naive version, and do the streaming softmax on four numbers by hand. The trap is answering "it is a faster algorithm"; the answer is that it is the same algorithm moved so that it touches slow memory far less.

The bottleneck is traffic: 134 MB per head, and a GPU that waits

Attention is slow in practice because of how many bytes it moves, not how many multiplications it does. Transformers and Attention, From the Whiteboard defines the operation: scores S = Q K^T / sqrt(d), probabilities P = softmax(S), output O = P V. The score matrix S has one entry per pair of tokens, n x n, and that is the problem.

Size one score matrix. At n = 8,192 and fp16 (2 bytes per entry): 8,192 x 8,192 x 2 B = 134,217,728 B, about 134 MB per head per layer. A 70B-class model with 64 query heads and 80 layers (the shape used in GPU Memory Math: Weights, KV Cache, and What Actually Fits) would, if all were held at once, need 134 MB x 64 x 80 = 687 GB. Frameworks free each layer's matrix as they go, so the peak is smaller, but the matrix still has to exist, and it quadruples every time context doubles: 16K tokens is 537 MB per head, 32K is 2.1 GB.

Now count the traffic. A GPU has a large, slow memory (HBM, the "VRAM" figure on the spec sheet) and a small, very fast on-chip memory (SRAM, the shared memory and registers inside each streaming multiprocessor, on the order of tens to a couple of hundred kilobytes per multiprocessor). Naive attention runs as separate kernels, each reading its inputs from HBM and writing results back. Simplified, per head: write S, read S for softmax, write P, read P for the value multiply, so four passes over the matrix, 4 x 134 MB = 537 MB. For the 64 heads of one layer that is 34.4 GB of traffic; over 80 layers, 2.75 TB for the attention scores of a single 8K-token prompt.

Compare with the arithmetic. The two matrix multiplies cost 4 x n^2 x d floating-point operations per head, with head dimension d = 128: 4 x 8,192^2 x 128 = 34.4 GFLOP per head, and 34.4 GFLOP x 64 x 80 = 176 TFLOP for the prompt. Take two assumed round figures for a modern data-centre GPU, 3 TB/s of HBM bandwidth and 500 TFLOP/s of achieved half-precision throughput (check your GPU's datasheet; the exact values vary):

• memory time for the score traffic: 2.75 TB / 3 TB/s = 0.92 s
• arithmetic time: 176 TFLOP / 500 TFLOP/s = 0.35 s

The traffic costs more than the maths, so the GPU's expensive tensor cores sit idle waiting for memory. (These figures ignore causal masking, which halves useful work, and assume one GPU does everything; the ratio is what matters.) This is the roofline idea: an operation with few FLOPs per byte moved is bandwidth-bound, and the cure is to move fewer bytes, not to compute faster.

The FlashAttention fix is to never write S or P to HBM at all. The paper, Dao et al., FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (May 2022), frames attention as an IO problem and uses tiling to cut reads and writes between HBM and SRAM. Reading Q, K, V and writing O for one head is 4 x n x d x 2 B = 4 x 8,192 x 128 x 2 B = 8.4 MB, against 537 MB naive: 64 times less traffic in this simplified count. Memory also stops growing as n^2: the extra memory is linear in n (a few statistics per row), which is why 32K-token prefill fits at all.

Say what it is not. It is not an approximation or sparse attention; it computes the same function exactly.

Tiling and online softmax: exact attention without ever holding all the scores

The obstacle to tiling is softmax, which needs every score in a row before it can normalise. softmax(s_i) = e^(s_i - m) / sum_j e^(s_j - m), where m is the row maximum, subtracted for numerical stability. Seeing the scores in blocks means the maximum and the sum are not known until the last block. Online softmax solves this by carrying two running numbers per row, the maximum m seen so far and the sum l of e^(s - m), and rescaling the old sum whenever a bigger maximum appears.

Worked example. One query row has four scores, arriving as two blocks: [2, 1 | 3, 0]. The ordinary softmax subtracts nothing in this textbook version: e^2, e^1, e^3, e^0 = 7.389, 2.718, 20.086, 1.000, sum 31.193, so the probabilities are [0.237, 0.087, 0.644, 0.032].

@@ 1 | Block 1: scores 2 and 1 Running max m = 2. Running sum l = e^(2-2) + e^(1-2) = 1 + 0.368 = 1.368.

@@ 2 | Block 2: scores 3 and 0 arrive, and 3 is a new maximum New max m' = 3. Rescale the old sum by e^(m - m') = e^(2-3) = 0.368: 1.368 x 0.368 = 0.503. Add the new block relative to the new max: e^(3-3) + e^(0-3) = 1 + 0.050 = 1.050. So l = 0.503 + 1.050 = 1.553.

@@ 3 | Normalise Each probability is e^(s - 3) / 1.553: for score 2, 0.368 / 1.553 = 0.237; for 1, 0.135 / 1.553 = 0.087; for 3, 1 / 1.553 = 0.644; for 0, 0.050 / 1.553 = 0.032. The same four numbers, and no step needed all four scores at once.

The same trick carries the output. FlashAttention also keeps a running output vector O, a weighted sum of value vectors. When a new maximum arrives, it multiplies the old O by the same rescale factor before adding the new block's contribution, then divides by l at the end. Take values V = [1, 2, 3, 4] (one number each, to keep it small). The exact result is (7.389 x 1 + 2.718 x 2 + 20.086 x 3 + 1.000 x 4) / 31.193 = 2.471. Streaming: after block 1, un-normalised output relative to max 2 is 1 + 0.368 x 2 = 1.736; rescale by 0.368 to get 0.639; add block 2, 1 x 3 + 0.050 x 4 = 3.199; total 3.838; divide by l = 1.553: 2.471. Identical.

Putting it together in the kernel. For each tile of queries (a block of rows, for example 128), loop over blocks of keys and values: load a K tile and a V tile into SRAM, compute the tile of scores, update m, l and O, and move to the next block. Only when the loop ends does the kernel write O (and a per-row log-sum-exp statistic) to HBM. An illustrative footprint at 128-wide tiles with d = 128: Q, K, V tiles in fp16 are 32 KB each (96 KB), and a 128 x 128 fp32 score tile is 64 KB, 160 KB in total. That is why tile sizes are tuned per GPU, and why the head dimension constrains which kernels apply. The backward pass does not store the attention matrix either; it recomputes tiles from Q, K, V and the saved statistics, trading extra FLOPs (cheap) for traffic (expensive). That is also why training memory drops.

Causal masks and windows come almost free. For decoder models, blocks that lie entirely above the diagonal are fully masked and can be skipped, which roughly halves the work in prefill. A sliding-window mask skips every block outside the window, and block-sparse variants skip blocks listed in a pattern. Skipping whole blocks is what makes masks cheap; masking inside a block is ordinary.

Versions 1 to 4, and the knobs an engineer actually controls

Each FlashAttention release re-fits the same idea to a new GPU generation. The core is the 2022 algorithm; the later versions change the work partitioning, the use of asynchronous hardware units and the precision.

:::key Verified October 2026 Figures are from the authors' own abstracts: arXiv 2205.14135 (May 2022), 2307.08691 (July 2023), 2407.08608 (July 2024) and 2603.05451 (5 March 2026). They are one group's measurements on specific kernels and shapes, and FlashAttention-4's headline comparison is against cuDNN 9.13 and Triton kernels, not against earlier FlashAttention versions. The project README (retrieved October 2026) lists FlashAttention-2 for Ampere, Ada and Hopper GPUs with fp16 and bf16 and head dimensions up to 256, FlashAttention-3 as a beta for Hopper (fp16/bf16 forward and backward, FP8 forward only), and FlashAttention-4 for Hopper and Blackwell. Flash-Decoding (Stanford CRFM, October 2023) splits the keys and values into chunks, runs attention on each chunk in parallel, and merges with the per-chunk log-sum-exp; its authors report up to 8x faster decoding for very long sequences on CodeLlama 34B. Check the README for your installed version before trusting this table. :::

What you control in practice is rarely the kernel itself; it is which backend your stack selects. Frameworks and serving engines expose attention backends (FlashAttention-family kernels, vendor kernels such as cuDNN or FlashInfer, a plain fallback), and the right one depends on GPU generation, dtype and head dimension. The controls to know:

• dtype. Half-precision (fp16 or bf16) is the fast path; fp32 or an unusual dtype may silently fall back to a slower kernel. FP8 attention exists on Hopper and later but changes numerics, so it needs the same evaluation gate as any quantisation (Quantization Methods Compared).
• Head dimension and GPU generation. An unsupported head dimension or an older GPU falls back. If throughput is far below expectation, check which kernel was selected in the logs before profiling anything else.
• Mask type. Causal and sliding-window masks have fast paths; arbitrary dense masks may not.
• Version drift. Upgrading the library can change the backend chosen and, as the last section explains, the last bits of the output.
• Build and install. Source builds of these kernels can take a long time and need a matching CUDA toolchain; prefer prebuilt wheels pinned in your image.

A reading rule for the speedup figures. The 2x or 3x are kernel or training-throughput numbers on particular shapes. In an end-to-end serving system, attention is one slice of step time; weights, communication and sampling are the rest. If attention is 30 percent of a step and you make it 2x faster, the step gets 1 / (0.7 + 0.15) = 1.18x faster, not 2x. Quote the end-to-end figure in the interview.

Prefill versus decode, the three bottlenecks, and exactness

FlashAttention helps prefill most, because that is where the n-by-n matrix exists. In prefill, all prompt tokens attend to each other at once, so S is n x n. In decode, each step has a query of length one: S is 1 x n, a single row. There is no big matrix to avoid, so decode is bandwidth-bound for a different reason: each step must read the entire KV cache. For the 70B-shaped model with 327,680 bytes of KV per token (from the memory lesson), a 128,000-token request holds 327,680 x 128,000 = 41.9 GB, and every decode step reads all of it. At an assumed 3.3 TB/s per GPU across 8 GPUs, that read alone is 41.9 GB / 26.4 TB/s = 1.6 ms per token for that one request. Decode kernels differ: with one query and few sequences the GPU has too little parallelism, so split-KV schemes (often called Flash-Decoding) divide the KV sequence across thread blocks and merge the partial results using the same running-max-and-sum rescale. Same algebra, different loop order.

Three different tools fix three different bottlenecks, and an interviewer will test whether you can separate them.

Continuous Batching and the Scheduler covers paged allocation, and the vLLM paper reports 2 to 4 times higher throughput than earlier systems at the same latency, which is the authors' own evaluation. They compose: a modern engine runs a FlashAttention-style kernel over paged KV blocks of a GQA model, possibly with an 8-bit cache. DeepSeek-V2 (May 2024) reports that its multi-head latent attention cuts the KV cache by 93.3 percent against its predecessor, the strongest case of the third row. Mixture-of-Experts and Long-Context Serving uses these numbers at 128K tokens.

Exactness has a precise meaning: the same function, not the same bits. Floating-point addition is not associative, so computing a sum in a different order can change the last digits. A small demonstration, in fp16: (1000 + 0.2) + 0.2 gives 1000.0 (each 0.2 is lost, since the spacing near 1000 is 0.5), while 1000 + (0.2 + 0.2) gives 1000.5. FlashAttention sums scores in blocks and rescales running totals, an order different from the naive kernel, so outputs agree to rounding error rather than bit-for-bit. Consequences to say aloud:

• Comparing a FlashAttention model to a reference with exact equality will fail; compare with a tolerance, then confirm with task metrics.
• Results can change between library versions, GPU generations or batch compositions, which bears on reproducibility claims (Sampling and Decoding: Temperature, Top-p, Beam Search, and Determinism discusses why greedy decoding is not guaranteed identical across runs).
• Differences are normally far below the noise of an evaluation, but when two near-tied tokens flip, a long generation can diverge. Do not file that as a bug in the model.

The interview answer, compressed. "Standard attention is memory-bound because it round-trips an n-squared matrix through HBM. FlashAttention tiles Q, K and V into SRAM and uses online softmax to merge blocks exactly, so traffic falls by roughly the factor of n over d and memory goes linear. It helps prefill and training most; decode is bound by KV reads, where PagedAttention, GQA, MLA and KV quantisation are the levers. It is exact up to floating-point reordering."

How to say it in the interview

"Standard attention builds the n by n score matrix, writes it to GPU high-bandwidth memory, reads it back for softmax, writes the probabilities and reads them again for the value multiply. At 8,192 tokens that matrix is 134 MB per head per layer, so the cost is memory traffic, not arithmetic. FlashAttention computes exactly the same result in tiles that stay in on-chip SRAM, using an online softmax that keeps a running maximum and running sum, so the full matrix is never written. It is exact, not an approximation. It helps prefill most; decode has a query of length one and is bound by reading the KV cache, which is what PagedAttention, grouped-query attention and KV quantisation address."

Checklist

• Computed the score-matrix size (n x n x bytes) for the target context length
• Can state that standard attention is bandwidth-bound and FlashAttention reduces HBM traffic, not FLOPs
• Did online softmax by hand on a small example, including the rescale when a new maximum arrives
• Said that FlashAttention is exact up to floating-point reordering, with the fp16 non-associativity example
• Separated FlashAttention (score traffic), PagedAttention (KV layout) and GQA/MLA/KV quantisation (cache size)
• Stated that decode has a 1 x n score row and is bound by reading the KV cache, with split-KV kernels as the fix
• Checked the attention backend, dtype, head dimension and GPU generation the stack actually selected
• Converted a kernel speedup into an end-to-end speedup with the attention share of the step
• Pinned the kernel library and used tolerance-based comparisons in regression tests
• Dated and sourced version claims (FA1 to FA4) and kept vendor and author figures labelled as such

Related lessons

• GPU Memory Math: Weights, KV Cache, and What Actually Fits
• Quantization Methods Compared: INT8, GPTQ, AWQ, FP8, FP4 and KV-Cache Quantization
• Continuous Batching and the Scheduler
• TTFT, TPOT, and the Latency Budget
• Mixture-of-Experts and Long-Context Serving
• Autoscaling a GPU Fleet Without Burning the Budget