What is the difference between SRAM and HBM?

HBM is the GPU's large main memory — many gigabytes, but relatively slow. SRAM is tiny on-chip memory — only kilobytes to megabytes, but roughly an order of magnitude faster. FlashAttention's whole strategy is to do as much work as possible inside fast SRAM and touch slow HBM as little as possible.

Is FlashAttention exact or approximate?

Exact. Unlike earlier efficient-attention methods that approximated the attention matrix (sparse or low-rank), FlashAttention computes the mathematically identical result as standard attention — just with a smarter memory schedule. You lose no accuracy; you only save time and memory.

How does FlashAttention enable long context windows?

Standard attention's memory grows with the square of the sequence length because it stores the full scores matrix. FlashAttention never materializes that matrix, so its memory grows only linearly with sequence length. That makes much longer context windows practical, which is why it became a default building block for modern long-context models.

Paper Breakdowns / FlashAttention

Paper 09~14 min readNeurIPS 2022worked math + runnable code

Paper Breakdown

FlashAttention,
explained.

Q: What is FlashAttention?

FlashAttention is a 2022 algorithm that computes exact attention faster and with far less memory by being IO-aware — minimizing data movement between a GPU's slow high-bandwidth memory (HBM) and its fast on-chip memory (SRAM). It never writes the full attention matrix to slow memory; instead it processes attention in tiles that fit in SRAM, using an online softmax to combine them.

Q: Why was standard attention slow?

Standard attention builds the full N-by-N scores matrix and writes it to the GPU's slow HBM, then reads it back for the softmax and again for the value multiply. For long sequences that matrix is huge, and all that reading and writing of slow memory — not the arithmetic — is the bottleneck. Attention was memory-bound, but everyone had been optimizing the math.

Everyone knew attention got painfully slow on long inputs, and everyone assumed the problem was too much math. This 2022 paper showed the arithmetic was never the bottleneck — the GPU was spending almost all its time shuffling data between fast and slow memory. Fix the shuffling, and attention gets several times faster while computing the exact same answer.

Video breakdown

The animated walkthrough is in production.

Read the full breakdown below in the meantime ↓

Prefer video? Browse published breakdowns →

The real bottleneck Fast memory, slow memory Tiling & online softmax The math, worked Exact, not approximate Run it yourself The results Why it still matters

The real bottleneck

Self-attention has every token compare itself to every other token, producing an N×N scores matrix. For a sequence of length N, that matrix grows with the square of N — at 8,000 tokens it's 64 million numbers. The standard implementation writes that whole matrix out to memory, reads it back to apply softmax, then reads it a third time to multiply by the values.

The paper's diagnosis: the GPU's arithmetic units were mostly idle, waiting on those reads and writes. Attention was memory-bound, not compute-bound. The field had been optimizing the wrong thing.

Fast memory, slow memory

A GPU has two very different kinds of memory. HBM is the big main memory — tens of gigabytes, but relatively slow. SRAM is tiny on-chip memory — only kilobytes to a few megabytes, but roughly an order of magnitude faster.

The GPU memory hierarchy

SRAM~20 MB, ~19 TB/s — tiny but blazing fast. Work here.

HBM~40 GB, ~1.5 TB/s — huge but ~10x slower. Touch as little as possible.

Standard attention parks that giant scores matrix in slow HBM and keeps trudging back to it. FlashAttention's entire strategy is to never put it there in the first place.

Tiling and the online softmax

Instead of computing the full matrix at once, FlashAttention breaks the work into tiles — small blocks of queries and keys that fit inside fast SRAM. It loads a block, computes its piece of attention entirely in SRAM, and moves on, never writing the intermediate scores back to HBM.

The tricky part is softmax, which normally needs to see a whole row at once to normalize it. FlashAttention uses an online softmax: it processes blocks one at a time and keeps a running total and running maximum, rescaling previous partial results as new blocks arrive — so it gets the exact normalized answer without ever holding the full row. All the attention operations are fused into a single pass over fast memory.

One query block walks the key/value blocks — only the amber tile is in SRAM

The math, worked with real numbers

The rescaling trick is the whole paper, so let's do it by hand. Softmax is computed "safely" by subtracting the row's max m before exponentiating (so nothing overflows), then dividing by the sum ℓ:

softmax(s_i) = e^{s_i − m} / ℓ where m = max_i(s_i), ℓ = Σ_i e^{s_i − m}

Both m and ℓ look like they need the whole row. They don't. Take a row of scores s = [2, 1, 3, 0] and process it in two blocks — [2, 1] first, then [3, 0]:

Step	Running max m	Running sum ℓ	What happened
Block [2, 1]	2	e⁰ + e⁻¹ = 1.3679	Everything measured against max-so-far = 2
Block [3, 0] arrives	3	—	New max! The old sum used the wrong max
Rescale old ℓ	3	1.3679 × e²⁻³ = 0.5033	One multiply fixes every earlier term at once
Absorb new block	3	0.5033 + e⁰ + e⁻³ = 1.5530	Exactly what the full row gives: e⁻¹+e⁻²+e⁰+e⁻³ = 1.5530 ✓

Why does one multiply fix everything? Every old term was e^s−2 but should have been e^s−3 — and e^s−3 = e^s−2·e²⁻³. The correction factor is the same for every term, so it can be applied to their sum. That's the update rule the kernel runs per block:

m_new = max(m_old, m_block)
ℓ_new = ℓ_old · e^{m_old − m_new} + Σ_i∈block e^{s_i − m_new}
O_new = O_old · e^{m_old − m_new} + Σ_i∈block e^{s_i − m_new} · V_i

The third line is the punchline: the output accumulator O (the weighted sum of values built so far) is corrected by the same factor. Final answer = O / ℓ. Finishing the example: with m = 3 and ℓ = 1.5530, the weights come out [0.2369, 0.0871, 0.6439, 0.0321] — identical to computing softmax on the whole row at once. No approximation anywhere; the N×N matrix simply never existed.

Exact, not approximate

This is what set FlashAttention apart. Earlier attempts to speed up attention approximated it — sparse patterns, low-rank tricks — trading accuracy for speed. FlashAttention changes only the memory schedule. The result is bit-for-bit the same attention; you give up nothing.

Worth knowing

Because it never stores the N×N matrix, FlashAttention's memory grows only linearly with sequence length instead of quadratically. That's the door to long context windows — the reason models could suddenly handle tens of thousands of tokens.

RUN IT YOURSELF

Online softmax & tiled attention, for real

This is the paper's core algorithm in ~40 lines of plain Python, running in your browser through WebAssembly. naive_softmax needs the whole row; online_softmax streams it in blocks with the rescale trick; flash_attention extends the same trick to the full attention output. Run it and watch the tiled version match the standard one to 16 decimal places — then change the row, the block size, or Q/K/V and check it still does.

CPython · WebAssembly

import math

def naive_softmax(row):
    """The textbook way - needs the WHOLE row before it can start."""
    m = max(row)                              # pass 1: global max
    exps = [math.exp(s - m) for s in row]
    total = sum(exps)                         # pass 2: global sum
    return [x / total for x in exps]

def online_softmax(row, block=2):
    """FlashAttention's way - one block at a time, never the whole row."""
    m, l = float("-inf"), 0.0                 # running max, running sum
    for b in range(0, len(row), block):
        chunk = row[b:b+block]
        m_new = max(m, max(chunk))
        # THE trick: one multiply re-bases every earlier term at once
        l = l * math.exp(m - m_new) + sum(math.exp(s - m_new) for s in chunk)
        m = m_new
    return [math.exp(s - m) / l for s in row]

row = [2.0, 1.0, 3.0, 0.0]                    # the row from the worked example
print("naive :", [round(x, 4) for x in naive_softmax(row)])
print("online:", [round(x, 4) for x in online_softmax(row)])

# ---- the same trick, extended to full attention ----
def attention(Q, K, V):
    """Standard attention - materializes every score."""
    d = len(Q[0]); out = []
    for q in Q:
        scores = [sum(qi*ki for qi, ki in zip(q, k)) / math.sqrt(d) for k in K]
        w = naive_softmax(scores)
        out.append([sum(wi*v[j] for wi, v in zip(w, V)) for j in range(len(V[0]))])
    return out

def flash_attention(Q, K, V, block=2):
    """Tiled - walks K/V in blocks, rescaling the accumulator O too."""
    d = len(Q[0]); out = []
    for q in Q:
        m, l = float("-inf"), 0.0
        acc = [0.0] * len(V[0])               # the running output O
        for b in range(0, len(K), block):
            idx = range(b, min(b+block, len(K)))
            scores = [(i, sum(qi*ki for qi, ki in zip(q, K[i])) / math.sqrt(d))
                      for i in idx]
            m_new = max([m] + [s for _, s in scores])
            corr = math.exp(m - m_new)        # same rescale, applied to l AND O
            l *= corr
            acc = [a * corr for a in acc]
            for i, s in scores:
                w = math.exp(s - m_new)
                l += w
                acc = [a + w*v for a, v in zip(acc, V[i])]
            m = m_new
        out.append([a / l for a in acc])      # final answer = O / l
    return out

Q = [[1, 0], [0, 1], [1, 1], [0.5, 0.5]]
K = [[1, 0], [0, 1], [1, 1], [1, 0.5]]
V = [[1, 2], [3, 4], [5, 6], [7, 8]]
std, fl = attention(Q, K, V), flash_attention(Q, K, V)
print("standard:", [[round(x, 4) for x in r] for r in std])
print("flash   :", [[round(x, 4) for x in r] for r in fl])
diff = max(abs(x-y) for ra, rf in zip(std, fl) for x, y in zip(ra, rf))
print("max difference:", diff)                # ~1e-16. Exact, not approximate.

# Try it: change block=2 to block=1 or 3 - the answer must not move.

What the paper actually measured

The headline numbers from the NeurIPS 2022 paper, not paraphrases:

Benchmark	Against	Result
BERT-large training (seq 512)	MLPerf 1.1 record	15% faster end-to-end wall clock
GPT-2 training (seq 1K)	HuggingFace / Megatron-LM	up to 3× / ~1.7× faster end-to-end
Long Range Arena	standard attention	2.4× speedup
Path-X (16,384 tokens)	all prior Transformers	first to beat chance (61.4%); block-sparse variant reached Path-256 at 64K tokens
Attention memory	quadratic baseline	linear in N — up to ~20× less at long lengths

And the structural comparison that explains why:

	Writes full N×N matrix?	Attention memory	Accuracy
Standard attention	Yes — to slow HBM	Quadratic in N	Exact
Sparse / low-rank	No	Lower	Approximate
FlashAttention	No	Linear in N	Exact

What it did not solve

FlashAttention makes attention IO-efficient, not sub-quadratic — the arithmetic is still O(N²); only the memory traffic became linear. The backward pass recomputes attention on the fly (trading FLOPs for memory), and each hardware generation needs its kernel rewritten — which is exactly what FlashAttention-2 and -3 turned out to be.

Why it still matters

FlashAttention became a default building block baked into every major training and inference stack — PyTorch's scaled_dot_product_attention, vLLM, TensorRT-LLM, xFormers. If you've used a model with a long context window, you've almost certainly used it. And it grew into a lineage:

Version	Year	What changed
FlashAttention	2022	Tiling + online softmax + recomputation; IO-aware exact attention
FlashAttention-2	2023	Better parallelism & work partitioning between GPU thread blocks and warps — roughly 2× again, reaching 50–73% of A100 peak FLOPs
FlashAttention-3	2024	Rebuilt for Hopper (H100): warp specialization, asynchronous TMA copies, FP8 support — another ~1.5–2×

Its deeper lesson is one every systems engineer learns eventually: on modern hardware, data movement is the cost, not arithmetic. The fastest code isn't the one that does the least math — it's the one that moves the least memory. FlashAttention is that principle, applied to the single most important operation in modern AI.

Frequently asked

Quick answers

What is FlashAttention?

An IO-aware algorithm that computes exact attention faster and with far less memory by minimizing data movement between a GPU's slow HBM and fast SRAM — processing attention in tiles with an online softmax.

Why was standard attention slow?

It writes the full N×N scores matrix to slow HBM and reads it back multiple times. The reads and writes — not the math — were the bottleneck; attention was memory-bound.

SRAM vs HBM?

HBM is large but slow main memory; SRAM is tiny but ~10x faster on-chip memory. FlashAttention does as much as possible in SRAM and touches HBM as little as possible.

Is it exact or approximate?

Exact — it computes the identical result as standard attention, just with a smarter memory schedule. No accuracy is lost.

How does it enable long context?

By never storing the full scores matrix, its memory grows linearly (not quadratically) with sequence length, making much longer context windows practical.

Finished this one? 0 / 101 Paper Breakdowns done

Explore the topic

See this alongside everything else on the same subject — handbooks, system designs, challenges and tools, in one place.

LLM Engineering

More Paper Breakdowns

Cite this page

Reference it in your work, paper or an AI's context window.

APASingh, S. (2026). FlashAttention. Vibe Engines. https://vibeengines.com/paper/flash-attention

MLASingh, Saurabh. “FlashAttention.” Vibe Engines, 2026, vibeengines.com/paper/flash-attention.

BibTeX

@online{vibeengines-flash-attention,
  author       = {Singh, Saurabh},
  title        = {FlashAttention},
  year         = {2026},
  organization = {Vibe Engines},
  url          = {https://vibeengines.com/paper/flash-attention}
}