FlashAttention: Attention Is Limited by Memory, Not Math

Key takeaway: flash attention (FlashAttention) computes the same attention result as the textbook version but never writes the N×N score matrix to GPU memory. It works on small tiles in fast on-chip SRAM instead. Standard attention is slow mostly because of memory traffic, not arithmetic, so a kernel that moves fewer bytes wins even at identical FLOPs. After reading, you should know when it helps (training, prefill, long context), when it does not (decode), and which version your stack probably runs.

Previous post: Prefill and Decode: Why One Model Has Two Bottlenecks.

TLDR: FlashAttention gives the same exact result from one fused, tiled kernel. It cuts HBM reads and writes, not FLOPs. FA2 improved GPU utilisation, FA3 targets Hopper, and decode needs other tools.

What this post covers

• What FlashAttention is, where it came from, and where it sits in the LLM stack
• How tiling and online softmax stay exact without the N×N matrix
• What FlashAttention 1, 2 and 3 each changed
• Where it helps and where it does not
• A snippet you can run
• The tradeoffs and a next step

What is flash attention (FlashAttention)?

FlashAttention is a GPU kernel for self-attention, introduced by Tri Dao and colleagues in May 2022. The paper calls it "IO-aware exact attention": a kernel designed around how much data moves between memory levels, not only how many operations run.

It sits between your framework and the hardware: PyTorch's scaled_dot_product_attention can dispatch to FlashAttention-2 kernels, and vLLM and SGLang expose FlashAttention backends.

Standard attention computes S = QKᵀ/√d, P = softmax(S), then O = PV, each as its own kernel, so the N×N matrices S and P round-trip through HBM. Derived size in fp16: 32 MiB per head at N = 4,096, 2 GiB at N = 32,768. HBM is large (80 GB on an H100, about 3.35 TB/s) but slow next to on-chip SRAM, which holds only about 200 KB per streaming multiprocessor.

How tiling and online softmax stay exact

FlashAttention splits Q, K and V into blocks that fit in SRAM and fuses the computation into one kernel. It streams over key and value blocks and accumulates the output without holding the full score matrix.

Softmax is the obstacle, because each row needs a global max and sum. Online softmax gets around it by keeping a running row max m and row sum l. When a new block raises the max, the earlier partial output and sum are rescaled by exp(m_old − m_new). The idea predates FlashAttention (Milakov and Gimelshein, 2018).

• The result is exact, up to floating-point rounding from a different summation order. FLOPs stay O(N²d), and the backward pass recomputes attention on-chip instead of storing it.
• Memory is linear. The extra state is one softmax statistic per row, so memory grows as O(N) instead of O(N²), and HBM accesses drop from Θ(Nd + N²) to Θ(N²d²/M), where M is SRAM size.

FlashAttention 1, 2 and 3

Each version went after the next bottleneck. Numbers are author-reported.

• FlashAttention (2022): up to 7.6x faster attention than PyTorch on GPT-2, and 1.7x faster end-to-end GPT-2 training than Megatron-LM at 1K context. It still reached only 25-40% of A100 peak throughput.
• FlashAttention-2 (2023): fewer non-matmul operations, parallelism over sequence length as well as batch and heads, and a better work split across warps. Roughly 2x over FA1, up to 230 TFLOPs/s (73% of peak) on A100.
• FlashAttention-3 (2024): built for Hopper, with asynchronous copies and matmuls, overlapped softmax and FP8 support. FP16 forward runs 1.5-2.0x faster than FA2, up to 740 TFLOPs/s (75% of the 989 dense peak) on H100.

vLLM defaults to FA3 on Hopper, FA2 on earlier GPUs and FA4 on Blackwell, per its docs, so check which one your engine selected. We quote no numbers for FA4.

Where it helps, and where it does not

The saving scales with the N×N intermediate, so the wins are biggest on long sequences: training, prefill and long context, where big matmuls and N×N traffic dominate. Decode is different. One query token per step is a matrix-vector product against the whole KV cache, with no N×N matrix to avoid, so the step is bound by streaming K and V from HBM.

The PyTorch team found that at batch size 1, FlashAttention uses less than 1% of the GPU during decode. Flash-Decoding splits the K/V sequence across the GPU and merges partial results, reporting up to 8x faster generation on very long sequences. FlashAttention also does not shrink the KV cache or replace PagedAttention; engines combine them.

Try it

This runs on CPU and checks the tiled algorithm against reference attention.

import math
import torch
import torch.nn.functional as F


def tiled_attention(q, k, v, block=128):
    """Single head, non-causal. Online softmax over K/V blocks; never builds the NxN matrix."""
    n, d = q.shape
    scale = 1.0 / math.sqrt(d)
    m = torch.full((n, 1), float("-inf"), dtype=q.dtype)  # running row max
    l = torch.zeros((n, 1), dtype=q.dtype)                # running row sum of exp
    o = torch.zeros((n, d), dtype=q.dtype)                # unnormalised output
    for j in range(0, k.shape[0], block):
        s = (q @ k[j:j + block].T) * scale                # (n, block) tile only
        m_new = torch.maximum(m, s.max(dim=1, keepdim=True).values)
        p = torch.exp(s - m_new)
        alpha = torch.exp(m - m_new)                      # rescale old statistics
        l = l * alpha + p.sum(dim=1, keepdim=True)
        o = o * alpha + p @ v[j:j + block]
        m = m_new
    return o / l


torch.manual_seed(0)
q, k, v = (torch.randn(1000, 64, dtype=torch.float64) for _ in range(3))
ref = F.scaled_dot_product_attention(q[None, None], k[None, None], v[None, None])[0, 0]
print("max abs diff:", (tiled_attention(q, k, v) - ref).abs().max().item())  # ~6e-16

The float64 difference on our CPU run was about 6e-16, which is rounding. To force the real kernel on a GPU, wrap F.scaled_dot_product_attention in sdpa_kernel(SDPBackend.FLASH_ATTENTION); we did not run that path.

Tradeoffs

• The win is HBM traffic, not FLOPs, and short sequences gain less because the N×N traffic is small.
• The flash-attn package needs Ampere, Ada or Hopper GPUs with fp16 or bf16, and FA3 is Hopper-specific.
• PyTorch falls back to another backend for unsupported shapes unless you force the kernel.
• Our numbers come from the cited papers and docs, not our own benchmarks.

Where to go from here

Attention cost grows with the context you send, so retrieving only what the question needs is the cheapest optimisation before any kernel work. Inherent's managed ingestion and deterministic retrieval do that for your agents: get started in the docs.

Which attention backend is your engine really running, and did it match what you expected? DM me on X with the version and GPU.

Next read: Continuous Batching Is How One GPU Serves a Crowd.