FlashAttention Explained: IO-Aware Attention, Tiling and Online Softmax

A transformer’s attention layer looks like three matrix multiplications and a softmax. On paper it is compute-bound math that a modern GPU should devour. In practice, the naive implementation spends most of its wall-clock time not multiplying, but shuttling a gigantic intermediate matrix back and forth between GPU memory and the chip. FlashAttention is the family of kernels that fixed this by changing where the data lives, not what the math computes. The output is bit-for-bit the same exact attention; only the memory traffic changes.

That matters more every year. Context windows have grown from 2K tokens to hundreds of thousands, the attention score matrix grows with the square of sequence length, and each new GPU generation adds tensor-core throughput much faster than it adds memory bandwidth or special-function throughput. The bottleneck keeps moving, and each FlashAttention version is a response to where it moved to.

You will leave with a working mental model of the GPU memory hierarchy, the tiling and online softmax algorithm step by step, a worked example of the traffic savings, and a clear picture of what FlashAttention-2, FlashAttention-3 (Hopper) and FlashAttention-4 (Blackwell) each changed.

What this covers: the memory-bound nature of attention, tiling and the online softmax recurrence, the IO complexity argument, the FA2 to FA3 to FA4 evolution, failure modes, and a practical checklist for using these kernels in production.

Context and Background

Self-attention computes O = softmax(Q K^T / sqrt(d)) V for query, key and value matrices of shape N by d, where N is sequence length and d is the head dimension (commonly 64 or 128). The textbook implementation materializes the N by N score matrix S, applies softmax row by row to get P, then multiplies P by V. For a single head with N = 8,192 and 16-bit values, S alone is 8,192 squared times 2 bytes, which is 128 MiB. Multiply by the number of heads and the batch size and you are moving gigabytes per layer, per step. This is illustrative arithmetic, not a benchmark, but it shows why the intermediate dominates.

Before 2022, the field mostly attacked this as a FLOP problem. Sparse attention, low-rank approximations (Linformer, Performer) and other approximate methods reduced the arithmetic. Many of them did not deliver wall-clock speedups, because the arithmetic was never the bottleneck. The 2022 paper FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness by Tri Dao and collaborators reframed the problem: count reads and writes between slow and fast memory, not FLOPs. Their kernel is exact, and the authors report a 15% end-to-end speedup on BERT-large at sequence length 512 against the MLPerf 1.1 training record, and a 3x speedup on GPT-2 at sequence length 1K over existing baselines.

The follow-ups track hardware. FlashAttention-2 reworked parallelism and work partitioning to reach roughly 2x the speed of the original, with 50 to 73 percent of A100 theoretical peak FLOPs. FlashAttention-3 exploited Hopper’s asynchronous hardware to reach up to 740 TFLOPs/s in FP16 (about 75 percent utilization on H100) and close to 1.2 PFLOPs/s in FP8. FlashAttention-4, written for Blackwell and described in a 2026 paper and Tri Dao’s write-up, attacks the new bottlenecks that appear once matrix units get faster than everything around them.

Today these kernels sit under nearly every serious training and inference stack. PyTorch’s scaled dot-product attention dispatches to FlashAttention-style kernels, inference engines use them for prefill, and libraries such as cuDNN have absorbed several of the ideas. If you serve or train transformers, you are almost certainly running a descendant of this work already. This post explains what it is doing, so you can reason about when it helps and when it will not. For the memory-physics background, see our companion piece on GPU memory physics, DRAM versus SRAM and kernel fusion, and for an earlier survey of the family, FlashAttention 2, 3 and 4 explained.

The GPU Memory Hierarchy and Why Attention Is Memory-Bound

The short answer: attention is memory-bound because the N by N score matrix must be written to and read from high-bandwidth memory (HBM) several times, and HBM is one to two orders of magnitude slower than the on-chip SRAM that actually feeds the tensor cores. FlashAttention removes those round trips by never materializing S or P in HBM.

Figure 1: The GPU memory hierarchy. Standard attention writes the N by N matrix to HBM; FlashAttention keeps tiles in on-chip shared memory.

The figure shows the levels a kernel can use. Registers are private to each thread and are the fastest and smallest. Shared memory (SRAM, software-managed) lives in each streaming multiprocessor (SM). L2 cache is shared across the chip. HBM is the large off-chip pool. Approximate public numbers for an H100 SXM are 80 GB of HBM3 at about 3.35 TB/s and roughly 228 KB of combined shared memory and L1 per SM; for an A100 80 GB, about 2 TB/s of HBM bandwidth and 192 KB of combined shared memory and L1 per SM. Treat these as ballpark figures and check NVIDIA’s datasheets for your exact SKU.

Arithmetic intensity: the roofline view

The roofline model asks how many FLOPs you perform per byte moved. A matrix multiply of large square matrices has high arithmetic intensity because each loaded element is reused many times. Softmax has almost none: for each element it reads one value, does an exponential and a couple of adds, and writes one value. It does a handful of operations per byte moved, which sits far left on the roofline, deep in memory-bound territory.

Standard attention is three kernels: a GEMM producing S, a softmax kernel producing P, and a GEMM producing O. Each kernel launch reads its input from HBM and writes its output back. The softmax step in particular forces a full read and write of an N by N matrix for almost no arithmetic. Dropout and masking, if unfused, add more passes. The FLOP count is the same as a fused kernel’s, but the bytes moved are far higher, so the time is set by memory bandwidth.

Counting HBM traffic

Take one head, sequence length N, head dimension d, 2-byte elements. The inputs Q, K, V and output O each cost N times d elements. The standard algorithm additionally writes S (N squared), reads S and writes P (2 N squared), then reads P (N squared). That is roughly four N-squared passes on top of the O(N d) linear traffic. The paper states the standard algorithm needs Theta(N d + N squared) HBM accesses.

For N = 8,192 and d = 128, the linear terms are 8,192 times 128 times 2 bytes, about 2 MiB per tensor. The N-squared term is 128 MiB per pass, so the quadratic traffic is roughly 64 times larger per pass than a single input tensor. This is why sequence length hurts so badly: the useful data is tiny and the intermediate is huge. It is also why the memory footprint, not just the time, goes quadratic in the naive version.

The key move: fuse and never write S

A fused kernel computes a tile of S, converts it to probabilities, multiplies by V and accumulates into O, all while the tile sits in SRAM. If it can do that, HBM only ever sees Q, K, V going in and O coming out. The difficulty is softmax: it normalizes across an entire row, and a row spans all N keys, which will not fit in SRAM at once. The solution is the online softmax recurrence, which lets you process a row in pieces.

Tiling and Online Softmax: The Core Algorithm

The direct answer: FlashAttention splits Q, K and V into blocks that fit in SRAM, loops over key and value blocks for each query block, and maintains a running row maximum and running sum so that softmax can be finalized incrementally. The result equals exact softmax attention, with the N by N matrix never stored.

Figure 2: The FlashAttention inner loop. Each key and value tile updates running statistics and the output accumulator, and only the final output and a log-sum-exp value return to HBM.

The loop in the figure is the whole algorithm. For each block of query rows, the kernel iterates over blocks of keys and values. It computes the tile of scores, updates the running softmax statistics, rescales the output accumulator, adds the new contribution, and moves on. When the key and value blocks are exhausted, it divides by the final sum and writes O. It also stores the log-sum-exp of each row, a small N-length vector, which the backward pass needs.

Numerically stable softmax in three passes

Plain softmax of a vector x is exp(x_i) / sum_j exp(x_j). To avoid overflow, implementations subtract the maximum m first: exp(x_i - m) / sum_j exp(x_j - m). This “safe softmax” needs three passes over the data: find the max, compute the sum of shifted exponentials, then normalize. Three passes over HBM-resident data means three memory reads.

Milakov and Gimelshein’s 2018 paper, “Online normalizer calculation for softmax”, showed that the max and the sum can be computed in a single pass. They report softmax speedups of up to 1.3x and, when fused with TopK, up to 5x. FlashAttention extends the same idea to cover the matrix multiplication with V, collapsing the whole pipeline into one pass.

The online softmax recurrence

Process the scores of one query row in blocks. Keep three running quantities: the row maximum m, the sum l of shifted exponentials, and the unnormalized output vector o. Initially m is negative infinity, l is 0 and o is the zero vector. For each new block with scores s (a short vector) and corresponding value rows V_blk:

Compute the block maximum, and set m_new = max(m, max(s)).
Compute the correction factor alpha = exp(m - m_new). This is the factor by which everything accumulated so far must shrink, because the reference maximum increased.
Compute p = exp(s - m_new) for the block.
Update the sum: l = alpha * l + sum(p).
Update the output: o = alpha * o + p @ V_blk.
Set m = m_new.

After the last block, the final output row is o / l. Because each earlier contribution was scaled by exactly the right power of the old and new maxima, the result is algebraically identical to the three-pass computation. Floating-point rounding differs slightly from a reference implementation because the order of operations differs, but the method is exact in the sense that it makes no approximation of the softmax itself.

A small worked example makes it concrete. Suppose a row has four scores [1, 3, 2, 5], processed two at a time. Block one is [1, 3]: m becomes 3, p = [exp(-2), exp(0)] = [0.135, 1.0], l = 1.135. Block two is [2, 5]: m_new = 5, alpha = exp(3 – 5) = 0.135, p = [exp(-3), exp(0)] = [0.050, 1.0], l = 0.135 times 1.135 plus 1.050 = 1.203. The direct computation with max 5 gives exp(-4) + exp(-2) + exp(-3) + exp(0) = 0.018 + 0.135 + 0.050 + 1 = 1.203. They match, as they must.

Block sizes and the SRAM budget

The tile sizes are set by the SRAM budget. If query blocks have B_r rows, key blocks have B_c rows and head dimension is d, SRAM must hold the Q block, K block, V block, the B_r by B_c score tile and the output accumulator. In the original paper the block size for keys is roughly M divided by 4d, where M is the SRAM size in elements, capped by d for the query side. In practice modern kernels choose tiles such as 128 by 128 for d = 128 on Hopper, balancing register pressure, shared memory capacity and tensor-core instruction shapes.

Larger tiles mean more reuse and less HBM traffic, but they cost registers and shared memory, which lowers occupancy (how many thread blocks fit on an SM at once). Tuning this trade-off per head dimension and per GPU generation is much of what the kernel-engineering effort is.

The IO complexity argument

The paper’s central theoretical result is that FlashAttention needs Theta(N squared d squared divided by M) HBM accesses, where M is the SRAM size, against Theta(N d + N squared) for standard attention. For typical values of d (64 to 128) and M (on the order of 100 KB), the denominator makes the new count much smaller than N squared. The authors also prove a lower bound showing that, for the range of M considered, no exact attention algorithm can asymptotically beat this. That is the sense in which the paper calls it “optimal for a range of SRAM sizes.”

Note what is not claimed. The arithmetic remains O(N squared d); FlashAttention does not make attention sub-quadratic in compute. What goes away is the quadratic memory footprint (it becomes linear in N) and the quadratic HBM traffic. At very long sequences the compute still grows quadratically, which is why approaches such as ring attention for context parallelism distribute the sequence across devices.

The backward pass: recompute instead of store

Training needs gradients, which classically require the stored P matrix. FlashAttention instead saves only O and the per-row log-sum-exp, then recomputes S and P tile by tile during the backward pass. Recomputation costs extra FLOPs, but those FLOPs are cheap relative to reading an N-squared matrix from HBM. This is the same trade as activation checkpointing, applied inside the kernel. The reported result is faster training despite doing more arithmetic, which is only possible because the baseline was bandwidth-limited.

The backward pass is heavier than the forward: FlashAttention-4’s write-up notes that it performs about 2.5 times the tensor-core work of the forward, using five matrix multiplications rather than two. That asymmetry matters later, because the backward pass hits different hardware limits than the forward.

From FlashAttention-2 to FlashAttention-3 and 4: Following the Bottleneck

Each generation of FlashAttention exists because the previous kernel hit a new wall. FlashAttention-2 fixed wasted non-matmul work and poor parallelism on Ampere. FlashAttention-3 fixed the failure to overlap asynchronous hardware on Hopper. FlashAttention-4 fixes the exponential unit and shared memory limits on Blackwell. Seeing the sequence as “move the bottleneck, chase it” is the most useful way to remember the whole family.

Figure 3: The same attention algorithm, three GPU generations, three different bottlenecks and fixes.

FlashAttention-2: fewer non-matmul FLOPs, better parallelism

Tensor cores are fast, but the GPU also has ordinary scalar and special-function units that execute the exponentials, rescales and masks. On an A100, half-precision tensor-core throughput is far higher than that of the general-purpose units, so each non-matmul FLOP costs much more than a matmul FLOP. FlashAttention-2, per its paper, reduced non-matrix-multiply operations by restructuring the algorithm. One representative trick is to keep the output unnormalized during the loop and divide by the sum only once at the end, rather than rescaling by the running denominator on every iteration.

The second change was parallelism. The original kernel parallelized over batch and heads, which leaves the GPU underfilled when batch times heads is small, as with long sequences and small batches. FlashAttention-2 also splits across the sequence dimension, even for a single head, so more thread blocks run concurrently. The third change was work partitioning within a thread block, distributing work between warps to reduce shared-memory communication. The paper reports about a 2x speedup over the first FlashAttention, 50 to 73 percent of theoretical peak on A100 (up from 25 to 40 percent), and end-to-end GPT training up to 225 TFLOPs/s per A100, which it states is 72 percent model FLOPs utilization.

Why FlashAttention-2 underperforms on Hopper

When teams moved to H100 they found FlashAttention-2 reached only about 35 percent of the GPU’s peak, according to the FlashAttention-3 paper. The reason is architectural. Hopper introduced three features that FlashAttention-2 does not use: asynchronous warpgroup-level matrix multiply (WGMMA), the Tensor Memory Accelerator (TMA) for asynchronous bulk copies between global and shared memory, and hardware FP8 support. A kernel written in the Ampere style issues loads and math in lockstep, so tensor cores sit idle while the same warps do softmax, and memory units sit idle during math.

The deeper issue is a throughput mismatch. H100’s FP16 dense tensor-core peak is on the order of 1 PFLOP/s, yet the special-function unit that computes exponentials handles far fewer operations per cycle. When the matrix units get this fast, the softmax exponentials become a visible fraction of runtime even though they are a small fraction of FLOPs. FlashAttention-3 is largely about hiding that softmax time behind the matrix math.

FlashAttention-3: warp specialization, pingpong and FP8

The FlashAttention-3 paper by Shah, Bikshandi, Zhang, Thakkar, Ramani and Dao reports 1.5 to 2.0x speedup over FlashAttention-2 on H100, up to 740 TFLOPs/s in FP16 (about 75 percent utilization) and close to 1.2 PFLOPs/s in FP8. It names three techniques.

Warp-specialized producers and consumers. Instead of every warp doing everything, a small producer warpgroup issues TMA loads that stream K and V tiles into a shared-memory pipeline, while consumer warpgroups run the WGMMA instructions and softmax. Because TMA and tensor cores operate asynchronously, data movement for tile j+1 overlaps with compute on tile j. The producer can also hand its register budget to consumers, since it needs few.

Interleaving matmul and softmax. Even with loads hidden, a single consumer’s softmax stalls its own tensor-core use. FlashAttention-3 uses two consumer warpgroups that “pingpong”: while one runs its GEMMs, the other does its softmax, and then they swap. A synchronization barrier schedules the order so the tensor cores always have a GEMM queued from one warpgroup while the other computes exponentials. Within a single warpgroup it also pipelines, issuing the GEMM for the next score tile before finishing softmax of the current one, at the cost of holding extra registers.

Figure 4: Simplified FlashAttention-3 schedule. A producer streams tiles with TMA while two consumer warpgroups alternate GEMM and softmax so tensor cores stay busy.

FP8 with block quantization and incoherent processing. FP8 doubles tensor-core throughput relative to FP16 but is much less precise. The paper addresses this with per-block quantization scaling and with “incoherent processing”, which multiplies Q and K by a random orthogonal matrix (built from a Hadamard transform with random signs) before quantization to spread outlier values across dimensions. The orthogonal transform cancels in Q K^T, so the exact result is unchanged, but the quantization error shrinks. The authors report that FP8 FlashAttention-3 has 2.6x lower numerical error than a baseline FP8 attention. FP8 also requires a layout fix: the WGMMA instruction expects operand layouts that differ from what the first GEMM’s accumulator produces, so the kernel rearranges data in registers rather than going through shared memory.

FlashAttention-4: Blackwell and asymmetric scaling

According to Tri Dao’s write-up, going from H100 to B200 BF16 tensor-core throughput rises from about 1 to 2.25 PFLOPs/s, while the number of special-function units and the shared-memory bandwidth stay the same. This asymmetry is the central problem. If matmul gets 2.25x faster and exponentials do not, the exponential becomes a co-bottleneck.

The write-up gives a per-SM tile analysis (for a 128 by 128 by 128 tile). On the forward pass, tensor cores need 1,024 cycles, the exponential unit needs 1,024 cycles, and shared memory needs 768, so the forward is limited by compute and exponentials equally. On the backward pass with a single CTA, tensor cores need 2,560 cycles, exponentials 1,024 and shared memory 3,328, so the backward is limited by shared memory bandwidth. Different phases, different walls.

FlashAttention-4’s reported answers are:

Software-emulated exponential. Part of the exponentials are computed on the ordinary FMA (fused multiply-add) units using a Cody-Waite range reduction and a low-degree polynomial, instead of the special-function unit. This spreads the exponential load over more hardware.
Conditional rescaling. The online-softmax rescale (the alpha multiply) is only performed when the running maximum increases by more than a threshold; otherwise the stale maximum is kept. Final normalization still applies. This removes many redundant rescale operations, since the maximum typically stabilizes quickly.
Larger ping-pong with a correction warpgroup. Two query tiles per CTA, two softmax warpgroups, and a dedicated correction warpgroup for rescaling.
Tensor memory (TMEM). Blackwell adds 256 KB of tensor memory per SM, so accumulators and backward-pass intermediates stay out of registers and shared memory.
2-CTA MMA. A pair of CTAs cooperate on a 256 by 256 tile, which roughly halves shared-memory traffic for one operand and halves the global atomic reductions of the query gradient in the backward pass.

The same write-up reports up to 1,605 TFLOPs/s for BF16 forward on B200, about 71 percent utilization, and speedups up to 1.3x over cuDNN 9.13 and 2.7x over Triton, with the note that the latest cuDNN offers similar performance since the cuDNN team absorbed some FlashAttention-4 techniques. Those are the authors’ numbers on specific shapes; treat them as an upper bound for your own workloads. The kernels are written in CuTe-DSL, a Python-embedded DSL for CUTLASS, which the authors say cuts compile times by roughly 20 to 30x relative to C++ templates.

Deeper Analysis: A Worked Traffic Model and Where Time Goes

Qualitative stories are nice, but you can estimate savings on the back of an envelope. The numbers below are illustrative arithmetic under stated assumptions, not measurements.

A worked HBM traffic example

Assume one attention head, N = 8,192, d = 128, FP16. Each of Q, K, V, O is 8,192 times 128 times 2 bytes, which is 2 MiB, so 8 MiB total of unavoidable traffic. Standard attention adds S written (128 MiB), S read (128 MiB), P written (128 MiB) and P read (128 MiB): 512 MiB of intermediate traffic. Total is about 520 MiB per head. FlashAttention’s forward reads Q once and re-reads K and V once per query block. With query blocks of 128 rows there are 64 query blocks, so K and V (4 MiB together) are streamed up to 64 times from HBM or L2, about 256 MiB, though much of that is served by L2 cache since many thread blocks read the same K and V tiles concurrently. Even charging all of it to HBM, the total is roughly 270 MiB, and with realistic L2 hit rates it is well below that.

The lesson is subtler than “512 versus 8.” Savings depend on cache behavior and tile size, and the advantage grows with sequence length because the standard version’s traffic scales with N squared while the fused version’s HBM traffic scales closer to N squared divided by the tile height. Larger query tiles cut re-reads of K and V proportionally.

The storage story is cleaner. The standard approach needs 128 MiB per head for S (and the same again for P if not fused), so 32 heads is 4 GiB for scores in one layer for one sequence. FlashAttention stores O plus an N-length log-sum-exp, which is linear. This is what lets long contexts fit in memory at all.

Where time goes inside a fused kernel

Once S and P stay on chip, the remaining time splits across four consumers: tensor-core matmul, exponentials, softmax bookkeeping (max, sum, rescale) and data movement from HBM or L2 into shared memory. FlashAttention-2 left them poorly overlapped. FlashAttention-3 overlaps them with asynchrony. FlashAttention-4 rebalances which unit does what.

A useful per-tile heuristic comes from the FlashAttention-4 analysis: count cycles each resource needs for one tile, and the largest is your bound. With M = N = D = 128 the matmul needs 2 times 128 cubed multiply-accumulate operations per GEMM; the exponential unit needs 128 times 128 exponentials. Because the former scales with d and the latter does not, smaller head dimensions are more exponential-bound. This is why a model with d = 64 can see a lower fraction of peak than d = 128 on the same GPU, and why tricks that offload exponentials matter most for narrow heads.

Causal masking, variable lengths and scheduling

Decoder models use causal masks, so roughly half the tiles are entirely masked and skipped. That halves FLOPs but creates load imbalance: query blocks near the end of the sequence process many more key tiles than those at the start. Good kernels schedule the longest blocks first so shorter ones fill the gaps. FlashAttention-4’s write-up mentions longest-processing-time-first scheduling and grid swizzling for exactly this reason, and its deterministic backward mode uses a shortest-first ordering with a lock to fix accumulation order, reaching about 85 to 90 percent of non-deterministic throughput.

Variable-length batches, as in inference serving, are handled by packing sequences with cumulative-length indices (“varlen” APIs) so no compute is wasted on padding. During decoding, where one query attends to a long cached key and value sequence, the standard tiling parallelizes poorly (one query row gives one query block). The “flash-decoding” approach splits the key and value sequence across thread blocks and combines partial results using the stored log-sum-exp, which is precisely the online softmax merge rule applied across blocks. For the serving side of this story, see our posts on speculative decoding architecture and Multi-head Latent Attention versus GQA and MQA.

Comparison of the three generations

Aspect	FlashAttention-2 (Ampere)	FlashAttention-3 (Hopper)	FlashAttention-4 (Blackwell)
Primary target	A100	H100	B200
Main problem solved	Non-matmul FLOPs, low occupancy	Idle async units, softmax stalls	Exponential unit and shared-memory limits
Key hardware used	mma.sync tensor cores	WGMMA, TMA, FP8	TMEM, 2-CTA MMA, asymmetric scaling
Headline figure (authors)	50 to 73% of peak; up to 225 TFLOPs/s training	Up to 740 TFLOPs/s FP16, near 1.2 PFLOPs/s FP8	Up to 1,605 TFLOPs/s BF16 forward, 71%
Low precision	FP16 and BF16	FP8 with incoherent processing	BF16 focus in the write-up; check release notes for others

Each row is from the authors’ papers or write-ups; hardware and shapes differ between columns, so do not read it as a direct cross-GPU benchmark.

Other attention kernels you will meet

FlashAttention is not the only fused attention. FlexAttention in PyTorch lets you express masks and score modifications (ALiBi, sliding windows, document masks) in Python and compiles them to a fused kernel. Triton and cuDNN offer their own fused attention. Paged attention, popularized by vLLM, solves a different problem: managing KV cache memory in blocks to avoid fragmentation, and typically calls a FlashAttention-style kernel underneath. Our earlier piece on FlashAttention-3 and PagedAttention covers that interplay in more depth.

Trade-offs, Gotchas, and What Goes Wrong

FlashAttention is not free of cost, and most production surprises come from the places where its assumptions break.

Hardware coupling. Each version is tuned to an architecture. FlashAttention-3 requires Hopper-class features; FlashAttention-4 targets Blackwell. On older GPUs you get the previous version, so a fleet with mixed GPU generations will have mixed performance. Check which kernel your framework actually selects rather than assuming.

Numerical differences. The result is mathematically exact but not bitwise identical to a naive implementation, because accumulation order differs. Models trained with one kernel and evaluated with another may show tiny output drift. Backward passes that use atomic adds (for dQ) are non-deterministic run to run, which complicates reproducibility; deterministic modes exist but cost throughput, per the FlashAttention-4 notes about 10 to 15 percent.

FP8 is not a free lunch. Even with incoherent processing, FP8 attention loses precision relative to BF16, and the 2.6x error reduction is relative to a naive FP8 baseline, not to BF16. Outlier-heavy activations, long-context retrieval tasks and training stability can all be sensitive. Validate on your own evaluation set before shipping.

Head dimension and shape restrictions. Kernels support specific head dimensions (commonly 64, 96, 128, 192 and 256 depending on version), and unusual sizes may fall back to slower paths. Arbitrary attention masks, per-element biases and some custom score modifications historically were not supported, which is the gap FlexAttention fills.

Short sequences see small gains. At sequence lengths of a few hundred tokens, the N-squared intermediate is small, launch overhead matters, and the speedup shrinks. The benefit is concentrated where the intermediate was huge.

Inference decoding is a different problem. Autoregressive decoding is bound by reading the KV cache, not by the attention matrix. FlashAttention helps with prefill; decode needs flash-decoding style parallelism and a smaller or compressed KV cache.

Build and portability pain. Compiling the CUDA kernels is slow and sensitive to CUDA, PyTorch and GPU-architecture combinations; mismatched versions are a frequent source of import errors and silent fallbacks. AMD and other accelerators have separate ports with their own feature coverage.

Opinion, clearly labeled: the biggest long-term risk is not any single bug but ecosystem concentration. When most of the industry’s attention compute runs through a handful of hand-tuned kernels, performance cliffs on new hardware show up as a lag of months until the new version lands. Compiler-based approaches (Triton, FlexAttention, CuTe-DSL) are the hedge, and FlashAttention-4’s choice of a Python DSL reads as a nod in that direction.

Practical Recommendations

If you train or serve transformers, the practical stance is: use the fused kernel your framework already provides, verify it is actually being selected, and spend your tuning effort on the places it does not cover.

In PyTorch, torch.nn.functional.scaled_dot_product_attention selects among FlashAttention, memory-efficient and math backends depending on dtype, shape, mask and hardware. Use the profiler or the backend-selection context manager to confirm a flash backend runs, because an arbitrary mask or an unsupported head dimension will quietly route you to a slower one. For Hopper and newer, install the FlashAttention-3 or later package where your stack supports it and benchmark on your real shapes, since headline numbers use large batches and long sequences.

Choose precision deliberately. BF16 is the safe default. Move to FP8 attention only with an evaluation gate that includes long-context and retrieval tasks. For serving, pair a fused prefill kernel with a paged KV cache, GQA or latent-attention models, and flash-decoding style kernels for decode.

Checklist:

Confirm the active attention backend with a profiler trace, not an assumption.
Benchmark at your real sequence-length and batch distribution, including variable lengths.
Use varlen or packed APIs to avoid padding compute.
Keep head dimension in the supported set; prefer 128 where you control model design.
Treat FP8 attention as an experiment gated by quality evals.
Pin CUDA, PyTorch and kernel package versions together, and test upgrades on a canary.
Use deterministic backward only when reproducibility is worth the throughput.
Revisit kernel versions whenever you change GPU generation.

Frequently Asked Questions

What is FlashAttention in simple terms?

FlashAttention is a GPU kernel that computes exact transformer attention without ever writing the large N by N score matrix to main GPU memory. It processes the sequence in small tiles that fit in fast on-chip memory and uses an online softmax to combine them. The output matches standard attention, but memory use becomes linear in sequence length and memory traffic drops sharply, which makes training and inference faster.

Is FlashAttention an approximation?

No. FlashAttention computes exact softmax attention. It reorders the computation using tiling and the online softmax recurrence, so results match the standard formula up to normal floating-point rounding differences. This separates it from sparse, low-rank or linear attention methods that change the function being computed. The exception is the optional FP8 mode in FlashAttention-3, which trades some precision for speed, though with lower error than a naive FP8 baseline.

How does online softmax work?

Online softmax processes a row of scores in blocks while tracking a running maximum and a running sum of exponentials. When a new block raises the maximum, previous partial sums and the output accumulator are multiplied by exp(old max minus new max) to rescale them. After the last block, dividing the accumulated output by the running sum gives the exact softmax-weighted result in a single pass over the data.

What is new in FlashAttention-3 compared with FlashAttention-2?

FlashAttention-3 targets Hopper GPUs. It uses warp specialization so a producer streams data with the Tensor Memory Accelerator while consumers run asynchronous matrix multiplies, overlaps softmax of one warpgroup with the GEMMs of another, and adds FP8 with block quantization and incoherent processing. The authors report 1.5 to 2.0x faster than FlashAttention-2 on H100, reaching up to 740 TFLOPs/s in FP16.

What does FlashAttention-4 change on Blackwell?

FlashAttention-4 responds to Blackwell’s uneven scaling, where tensor-core throughput grows faster than exponential-unit throughput and shared-memory bandwidth. It emulates part of the exponential in software on FMA units, skips unnecessary softmax rescaling, uses tensor memory and 2-CTA matrix instructions, and is written in CuTe-DSL. The authors report up to 1,605 TFLOPs/s BF16 on B200, around 71 percent utilization.

Does FlashAttention help with long-context decoding?

Only partly. FlashAttention greatly speeds up prefill and training, where full attention matrices are computed. During token-by-token decoding, the cost is dominated by reading the cached keys and values, so techniques such as flash-decoding (splitting the KV sequence across thread blocks), paged KV caches, grouped-query attention and KV compression matter more. Production inference stacks combine these with a FlashAttention-style kernel for the prefill phase.

FlashAttention Explained: IO-Aware Attention, Tiling, Online Softmax and FlashAttention-3/4 on Modern GPUs (2026)