Is FlashAttention a different attention mechanism?

No. FlashAttention computes the same scaled dot-product attention expression, but changes the order and location of the computation. The score blocks are produced, consumed by softmax, and combined with values while still on chip, so the full score or probability matrix is never written out to high-bandwidth memory.

Why is FlashAttention described as an IO optimisation?

Because its main win is avoiding unnecessary reads and writes of the quadratic intermediate matrices. The tensor-core matrix multiplications still happen, and the softmax still happens. The difference is that the kernel tiles the loop and carries enough running softmax state to avoid spilling the whole attention matrix to external memory.

What changed from FlashAttention to FA2 and FA3?

FA2 kept the same exact-attention idea but improved parallel work partitioning, occupancy, and the balance between tensor-core work and scalar softmax bookkeeping, especially on Ampere-class hardware. FA3 is more Hopper-specific: it pipelines loading, matrix multiply, and softmax around hardware features such as TMA, WGMMA, asynchronous execution, and FP8 support.

Does FlashAttention speed up every LLM inference path?

No. It helps most when dense attention would otherwise materialise large score or probability matrices, especially in prefill. In single-token decode, there is only one new query position, so the dominant cost is often streaming cached keys and values. In that case, KV-cache layout, paging, batching, and scheduling may matter more.

FlashAttention

FlashAttention is an exact GPU attention algorithm that reduces high-bandwidth-memory traffic by computing attention in tiles, keeping score and probability blocks on chip instead of materialising the full quadratic attention matrix. It computes the same softmax attention result, but reorganises the kernel around IO, online softmax, and hardware-aware scheduling.

The problem FlashAttention addresses is not that attention has the wrong formula, but that a naive implementation moves too many bytes. Standard scaled dot-product attention forms the full score matrix, writes it to high-bandwidth memory, reads it for softmax, then often handles a probability matrix before multiplying by values. For long prefills, that quadratic intermediate can dominate runtime even when the arithmetic is well suited to tensor cores.

FlashAttention streams the computation block by block. A tile of queries is combined with a tile of keys to produce a small score block in registers or shared memory. Instead of storing all scores, the kernel maintains each row’s running maximum, normalisation term, and partial output. When the next key and value tile arrives, it updates those quantities with a numerically stable online softmax and rescales the accumulated output.

The trade-off is specialised kernel complexity. FlashAttention saves memory traffic, but it adds online-softmax bookkeeping, tighter constraints on tile sizes, register pressure, shared-memory use, and architecture-specific scheduling choices. It is commonly misunderstood as an approximation like sparse attention. It is not: the usual result is exact attention, subject to normal floating-point differences. Whether it helps depends on sequence shape, masking, decode versus prefill, and the GPU generation.

Engineers usually meet FlashAttention through inference and training backends rather than by writing the recurrence themselves. In production inference it matters most for prefill and chunked prefill, where many query positions attend over a context. During single-token decode, the bottleneck is often reading the KV cache, so FlashAttention is not a substitute for cache paging, batching, or scheduling. FA2 improved work partitioning on Ampere-class GPUs; FA3 targets Hopper features such as asynchronous copies, WGMMA, TMA, and FP8 paths.

Go deeper

Inference EngineeringFlashAttentionExplain FlashAttention as an IO result rather than a maths one: tiling the computation so the full attention score matrix is never written to HBM, and what FA2 and FA3 changed on Ampere and Hopper.Open the lesson

Common questions

Is FlashAttention a different attention mechanism?: No. FlashAttention computes the same scaled dot-product attention expression, but changes the order and location of the computation. The score blocks are produced, consumed by softmax, and combined with values while still on chip, so the full score or probability matrix is never written out to high-bandwidth memory.
Why is FlashAttention described as an IO optimisation?: Because its main win is avoiding unnecessary reads and writes of the quadratic intermediate matrices. The tensor-core matrix multiplications still happen, and the softmax still happens. The difference is that the kernel tiles the loop and carries enough running softmax state to avoid spilling the whole attention matrix to external memory.
What changed from FlashAttention to FA2 and FA3?: FA2 kept the same exact-attention idea but improved parallel work partitioning, occupancy, and the balance between tensor-core work and scalar softmax bookkeeping, especially on Ampere-class hardware. FA3 is more Hopper-specific: it pipelines loading, matrix multiply, and softmax around hardware features such as TMA, WGMMA, asynchronous execution, and FP8 support.
Does FlashAttention speed up every LLM inference path?: No. It helps most when dense attention would otherwise materialise large score or probability matrices, especially in prefill. In single-token decode, there is only one new query position, so the dominant cost is often streaming cached keys and values. In that case, KV-cache layout, paging, batching, and scheduling may matter more.

Every term