Attention & Transformers

FlashAttention: IO-Aware Attention

A fused, tiled attention implementation that avoids materializing the full T x T matrix by using an online softmax, reducing memory traffic and speeding up long-context training/inference.

status: publishedimportance: importantdifficulty 4/5math: graduateread: 18mlive demo

Back to Attention & Transformers Next: LLM Serving at Scale: Prefill, Decode & Continuous Batching

Editorial transformer-systems illustration of tiled attention streaming through a compact on-chip scratchpad without materializing the full matrix.

Concept Structure

FlashAttention: IO-Aware Attention

01Intuition

Start with the picture, metaphor, or geometric mechanism.

02Math

Make the objects explicit and connect them with notation.

03Code

Mirror the equations with runnable implementation details.

04Interactive Demo

Manipulate the mechanism and watch the idea respond.

2prerequisites

1next concepts

2related links

Learning map

FlashAttention: IO-Aware Attention

BeforeEfficient Attention at Scale: KV Cache, GQA & FlashAttentionNow4/4 sections readyTryManipulate one control and predict the visible change.NextLLM Serving at Scale: Prefill, Decode & Continuous Batching

Object flow

4/4 sections readyAsk about this Research room

ConceptFlashAttention: IO-Aware AttentionAttention & Transformers EquationFlashAttention: IO-Aware Attention equation 1Exact equation object CodeFlashAttention: IO-Aware Attention code witness 1Exact code witness DemoFlashAttention: IO-Aware Attention interactive demoVisualization object ClaimFlashAttention computes exact attention while using tiling and online...Exact claim check SourceFlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awa...Exact source object

ConceptFlashAttention: IO-Aware AttentionAttention & Transformers

1 source attachedLocal snapshot ready

concept:attention-transformers/flash-attention

Codewitness nearby Predictbefore reveal Roomobject handoff

Conceptual Bridge

What should feel connected as you move through this page.

Carry inEfficient Attention at Scale: KV Cache, GQA & FlashAttention

Bring the mental model from Efficient Attention at Scale: KV Cache, GQA & FlashAttention; this page will reuse it instead of restarting from zero.

Work hereFlashAttention: IO-Aware Attention

A fused, tiled attention implementation that avoids materializing the full T x T matrix by using an online softmax, reducing memory traffic and speeding up long-context training/inference.

Carry outLLM Serving at Scale: Prefill, Decode & Continuous Batching

The next edge should feel earned: use the demo prediction here before following LLM Serving at Scale: Prefill, Decode & Continuous Batching.

Test the linkManipulate one control and predict the visible change.Then continue to LLM Serving at Scale: Prefill, Decode & Continuous Batching

01IntuitionStart with the picture, metaphor, or geometric mechanism.02MathMake the objects explicit and connect them with notation.03CodeMirror the equations with runnable implementation details.04Interactive DemoManipulate the mechanism and watch the idea respond.

Intuition

Build the mental picture first so the rest of the page has something to attach to.

Section prompt

Standard attention is conceptually simple but computationally awkward on GPUs:

Compute the $T\times T$ score matrix $S=QK^\top/\sqrt{d_k}$ .
Apply a row-wise softmax to get attention weights $P=\mathrm{softmax}(S)$ .
Multiply by values: $PV$ .

The problem is step (1): materializing a $T\times T$ matrix is huge, and moving it to and from GPU memory (HBM) is slow. For long sequences, attention is often memory-bandwidth bound rather than FLOPs bound.

FlashAttention is the idea: keep computations in fast on-chip SRAM by tiling, and compute the softmax in a streaming way so you never need to store the full $T\times T$ matrix.

Math

Translate the story into symbols, assumptions, and a derivation you can inspect.

Section prompt

Equation 1\mathrm{Attn}(Q,K,V)=\mathrm{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right)V.Equation 2\mathrm{softmax}(s)_j = \frac{\exp(s_j-m)}{\sum_{k=1}^T \exp(s_k-m)},\qquad m=\max_k s_k.

Standard attention (and the $T\times T$ issue)

For one head, with $Q,K,V\in\mathbb R^{T\times d_k}$ :

\mathrm{Attn}(Q,K,V)=\mathrm{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right)V.

The attention weights matrix is $P\in\mathbb R^{T\times T}$ . If you store it explicitly, memory scales like $O(T^2)$ .

Online softmax (streaming trick)

For a single row of scores $s\in\mathbb R^T$ , the softmax uses:

\mathrm{softmax}(s)_j = \frac{\exp(s_j-m)}{\sum_{k=1}^T \exp(s_k-m)},\qquad m=\max_k s_k.

You can compute this in blocks without ever storing all scores at once by maintaining a running max $m$ and a running normalizer $\ell$ .

For a block $b$ of scores with:

m_b=\max(s_b),\qquad \ell_b=\sum_{j\in b}\exp(s_j-m_b),\qquad o_b=\sum_{j\in b}\exp(s_j-m_b)\,v_j,

you can merge block statistics with the running state via:

m\leftarrow \max(m,m_b),

\ell\leftarrow e^{m_{old}-m}\,\ell_{old}+e^{m_b-m}\,\ell_b,

o\leftarrow e^{m_{old}-m}\,o_{old}+e^{m_b-m}\,o_b.

At the end, the attention output for that row is:

\mathrm{out}=\frac{o}{\ell}.

This is the core idea: you can compute $PV$ without ever materializing $P$ .

Code

Keep the implementation aligned with the notation so the algorithm is legible.

Section prompt

Code witness 1import numpy as np def softmax(x): m = float(x.max()) e = np.exp(x - m) return e / e.sum() de...python

import numpy as np

def softmax(x):
    m = float(x.max())
    e = np.exp(x - m)
    return e / e.sum()

def attn_out_stream(scores, values, block=2):
    m, l, o = -np.inf, 0.0, 0.0
    for i in range(0, len(scores), block):
        sb = scores[i:i+block]
        vb = values[i:i+block]
        mb = float(sb.max())
        eb = np.exp(sb - mb)
        lb = float(eb.sum())
        ob = float((eb * vb).sum())
        m_new = max(m, mb)
        l = np.exp(m - m_new) * l + np.exp(mb - m_new) * lb
        o = np.exp(m - m_new) * o + np.exp(mb - m_new) * ob
        m = m_new
    return o / l

s = np.array([2.0, 1.0, -1.0, 0.5, 3.2, -0.3])
v = np.array([0.1, 0.2, -0.4, 0.0, 0.7, 0.3])

out_full = float((softmax(s) * v).sum())
out_stream = float(attn_out_stream(s, v, block=2))
print("full:", round(out_full, 6), "stream:", round(out_stream, 6), "diff:", round(abs(out_full - out_stream), 9))

Interactive Demo

Use direct manipulation to connect the explanation to a moving system.

Section prompt

The demo below asks you to predict the memory bottleneck first, then reveal how the online softmax state $(m,\ell,o)$ is merged across tiles. The key invariant is that FlashAttention computes exact attention while avoiding a stored $T\times T$ probability matrix.

Live Concept Demo

Explore FlashAttention: IO-Aware Attention

The stage is code-native and interactive. Use it to test the explanation against the mechanism.

difficulty 4/5graduatecode-aligned

Demo Prediction Checkpoint

Manipulate one control and predict the visible change.

Commit to what FlashAttention: IO-Aware Attention should make visible before reading the result.

After The First Pass

Turn the concept into an inspected object.

Once the invariant is visible in the intuition, math, code, and demo, use these panels to inspect the mechanism visually, check source support, practice the idea, and attach a grounded research question.

Mechanism Storyboard

See the idea move before the page explains it

A fused, tiled attention implementation that avoids materializing the full T x T matrix by using an online softmax, reducing memory traffic and speeding up long-context training/inference.

Prediction open01 / Intuition

Prediction lens

Start with the picture, metaphor, or geometric mechanism.

Commit first

Before reading further, choose the kind of change FlashAttention: IO-Aware Attention should make visible.

Visual Inquiry

Make the image answer a mathematical question

A fused, tiled attention implementation that avoids materializing the full T x T matrix by using an online softmax, reducing memory traffic and speeding up long-context training/inference.

4/4 stages readyLive demo connected

Visual cueWhich visible object should carry the first intuition?

Inspection depth2/4

Prediction

Which visible object should carry the first intuition?

Commit first

Pick the cue that should make FlashAttention: IO-Aware Attention easier to reason about before the page gives the answer.

Source Grounding

Canonical references for the mechanism on this page.

paper · 2022FlashAttention: Fast and Memory-Efficient Exact Attention with IO-AwarenessDao et al.

Primary source for the dense FlashAttention mechanism: IO-aware HBM/SRAM tiling plus block-wise/incremental softmax statistics reduce HBM reads/writes and avoid materializing the full N x N attention matrix.

Open source

Claim Review

A fused, tiled attention implementation that avoids materializing the full T x T matrix by using an online softmax, reducing memory traffic and speeding up long-context training/inference.

Status1 substantive review recorded

Claims without a substantive review badge still need exact source-support review.

Sources1 reference

dao-2022-flashattention

Witnesses4 local objects

Use equation, code, and demo objects to check whether the source support is operational.

Substantively reviewedFlashAttention computes exact attention while using tiling and online softmax to reduce high-bandwidth-memory traffic instead of materializing the full T x T attention matrix.Claim metadata: source checked

Dao et al. frame FlashAttention as IO-aware exact attention using HBM/SRAM tiling; Algorithm 1 updates block softmax statistics and returns O=softmax(QK^T)V while avoiding reads/writes of the full N x N attention matrix to HBM.

Sources: FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness"Online softmax" means the paper's block-wise/incremental softmax statistics. This checks the dense FlashAttention mechanism, not block-sparse variants, CUDA numeric equivalence, backward-pass implementation, or universal speedups.A bounded review summary is present; still check caveats and exact source scope.

Checked Dao et al. against the exact claim: the paper presents FlashAttention as exact attention, uses HBM/SRAM tiling, maintains block softmax statistics (m,l,O) to compute the same softmax(QK^T)V result, and avoids reading/writing the full N x N attention matrix to HBM; local math/code/demo illustrate the same finite-row invariant.

Reviewer: codex+oracle; reviewed 2026-05-07

source-span-dao-2022-flashattention math-object-1 math-object-2 code-witness-1 interactive-demo

Source support candidates

paper 2022FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

Mechanism witnesses

Equation 1

\mathrm{Attn}(Q,K,V)=\mathrm{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right)V.

Equation 2

\mathrm{softmax}(s)_j = \frac{\exp(s_j-m)}{\sum_{k=1}^T \exp(s_k-m)},\qquad m=\max_k s_k.

Code witness 1import numpy as np def softmax(x): m = float(x.max()) e = np.exp(x - m) return e / e.sum() de...Demo stateLive mechanism probe

Practice Loop

Try the idea before it explains itself

A fused, tiled attention implementation that avoids materializing the full T x T matrix by using an online softmax, reducing memory traffic and speeding up long-context training/inference.

Readiness0/3 checks ready

Predict

Before touching the demo, predict one visible change that should happen in FlashAttention: IO-Aware Attention.

Hint 1

Reveal when your model needs a nudge.

Hint 2

Reveal when your model needs a nudge.

Hint 3

Reveal when your model needs a nudge.

Your answer canvas

Local checks

Claim

A concrete answer is on the canvas.

Mechanism

The answer names why the claim should hold.

Bridge

It touches the page context or a neighboring idea.

Misconception check

Object research drawerClose

ConceptFlashAttention: IO-Aware AttentionAttention & Transformers

Code witness comparisonFlashAttention: IO-Aware Attention code witness 1out_stream = float(attn_out_stream(s, v, block=2))Prediction before revealFlashAttention: IO-Aware Attention interactive demoManipulate one control and predict the visible change.

Grounded room questionWhat is the smallest example that makes FlashAttention: IO-Aware Attention click without losing the math?Local snapshot ready

Research Room

Attach the question to an exact object

Pick the concept, equation, source, code witness, claim, misconception, or demo state before asking for help. The handoff stays grounded to that object.

Next local actionNo local draft saved yet

Open the draft below to save one note and next action in this browser.

conceptAttention & Transformers

FlashAttention: IO-Aware Attention

Anchored question

What is the smallest example that makes FlashAttention: IO-Aware Attention click without losing the math?

Local action draftNo local draft saved yetExpand only when ready to capture one local next action

Local action draft

This draft stays locally in this browser for concept:attention-transformers/flash-attention.

Draft noteNext action

No local draft saved.

Evidence to inspect

Source ids to inspect: dao-2022-flashattention
Definition, prerequisite, and contrast concept links
The equation or code witness that makes the concept operational
One demo state that shows the invariant instead of a slogan

What would resolve this

The learner can state the mechanism in their own words
The learner can name the prerequisite that would repair confusion
The learner can predict how the mechanism changes under one perturbation

Grounded AI handoff

I am working in Continuous Function's research reading room. Object: concept - FlashAttention: IO-Aware Attention Object key: concept:attention-transformers/flash-attention Context: Attention & Transformers Anchor id: concept/concept-notebook/attention-transformers/flash-attention Open question: What is the smallest example that makes FlashAttention: IO-Aware Attention click without losing the math? Evidence to inspect: - Source ids to inspect: dao-2022-flashattention - Definition, prerequisite, and contrast concept links - The equation or code witness that makes the concept operational - One demo state that shows the invariant instead of a slogan What would resolve this: - The learner can state the mechanism in their own words - The learner can name the prerequisite that would repair confusion - The learner can predict how the mechanism changes under one perturbation Answer as a careful research tutor: stay source-grounded, separate verified evidence from assumptions, name the relevant math objects, and end with one next action.

Open source object

concept/concept-notebook/attention-transformers/flash-attention
concept:attention-transformers/flash-attention

Learning Map

Before / Now / Try / Next

BeforeEfficient Attention at Scale: KV Cache, GQA & FlashAttention

NowIntuition → Math → Code → Demo

TryManipulate one control and predict the visible change.

NextLLM Serving at Scale: Prefill, Decode & Continuous Batching

Intuitionready
Mathready
Codeready
Interactive Demoready

Object Companion

Ask beside the selected object

A fused, tiled attention implementation that avoids materializing the full T x T matrix by using an online softmax, reducing memory traffic and speeding up long-context training/inference.

Your question

GoalComfortStyleStuck on

Context prompt

You are my AI learning companion for Continuous Function. Current context: Attention & Transformers concept. Learning surface: FlashAttention: IO-Aware Attention. What this page says: A fused, tiled attention implementation that avoids materializing the full T x T matrix by using an online softmax, reducing memory traffic and speeding up long-context training/inference. Current section: Intuition, math, code, and interactive demo. Suggested next step: Manipulate one control and predict the visible change.. Learner goal: Understand the idea. Learner comfort level: New to this. Preferred explanation style: Visual first. Task: Explain the central idea in plain language, then restate it with the exact math objects from the page. Answer in a way that helps me learn: ask one clarifying question only if needed, use intuition before notation, and end with one thing I should try on the page.

Domain

Attention & Transformers

attention-transformersflash-attentionefficiencymemoryio-aware

Prerequisites

Efficient Attention at Scale: KV Cache, GQA & FlashAttention Long Context Engineering: RoPE Scaling, KV Compression & Memory Optimization

Leads To

LLM Serving at Scale: Prefill, Decode & Continuous Batching

Efficiency: Quantization, Distillation, LoRA & Sparse MoE Scaled Dot-Product Attention & Transformer Layers

Within this domain

FlashAttention: IO-Aware Attention

FlashAttention: IO-Aware Attention

What should feel connected as you move through this page.

Standard attention (and the T×TT\times TT×T issue)

Online softmax (streaming trick)

Turn the concept into an inspected object.

Start with the picture, metaphor, or geometric mechanism.

Which visible object should carry the first intuition?

Canonical references for the mechanism on this page.

A fused, tiled attention implementation that avoids materializing the full T x T matrix by using an online softmax, reducing memory traffic and speeding up long-context training/inference.

Source support candidates

Mechanism witnesses

Attach the question to an exact object

FlashAttention: IO-Aware Attention

Standard attention (and the $T\times T$ issue)