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.

published · difficulty 4/5 · 18 min read

Intuition

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

PredictName the object in plain language, then predict what should change.Leave with one reusable mental picture before notation appears.

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.

Section prompt

Copy-only prompts — each action copies a page-grounded prompt to your clipboard. Nothing is sent by this site.

Math

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

InspectTrack the same object through the notation and check each symbol.Leave with the invariant the equations preserve.

Equation 1Inspect equation objectAttn(Q,K,V)=softmax ((QK^T)/(sqrt(d_k)))V.Equation 2Inspect equation objectsoftmax(s)_j = (\exp(s_j-m))/(\sum_{k=1)^T \exp(s_k-m)},\qquad m=\max_k s_k.Equation 3Inspect equation objectm_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)...

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$ .

Section prompt

Copy-only prompts — each action copies a page-grounded prompt to your clipboard. Nothing is sent by this site.

Code

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

TraceMatch variables to symbols before reading the implementation.Leave with a runnable witness for the math.

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))

Section prompt

Copy-only prompts — each action copies a page-grounded prompt to your clipboard. Nothing is sent by this site.

Interactive Demo

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

ManipulateChange one control and predict the visible response before reveal.Leave with the observed invariant or a repaired model.

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 inquiry checkpoint

Manipulate one control and predict the visible change.

01Choose lensTrace a quantity

02ObserveDemo state pending

03GroundName the equation, invariant, or control that explains it.

04CarryNext: LLM Serving at Scale: Prefill, Decode & Continuous Batching

Choose what to inspect in FlashAttention: IO-Aware Attention. This shared fallback is an observation guide, not evidence of learning.

Loading interactive demo...

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.

Section prompt

Copy-only prompts — each action copies a page-grounded prompt to your clipboard. Nothing is sent by this site.

Object

4/4 sections ready

Concept: FlashAttention: IO-Aware Attention

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

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 contextAttention & Transformers

StateAttention & Transformers

1 source attachedChecking local snapshot

concept:attention-transformers/flash-attention

Math3 equation objectsAttn(Q,K,V)=softmax ((QK^T)/(sqrt(d_k)))V.CodeFlashAttention: IO-Aware Attention code witness 1out_stream = float(attn_out_stream(s, v, block=2))DemoFlashAttention: IO-Aware Attention predictionManipulate one control and predict the visible change.

ConceptLearner lens

FlashAttention: IO-Aware Attention

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

Mode questionCan I say the mechanism back in one sentence before I reveal anything?

Start with the prediction checkpoint, then compare the reveal to the mental model.

Take this move

Study modes

Keep the object fixed; change the lens.

01PredictCommit before reveal.02ExplainName the invariant.03Ground1 source link nearby.04ExperimentChange one knob.05CarrySave the next question.

Route back through the notebook

Carry the same object through intuition, math, code, and demo.

4/4 sections ready

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.

After The First Pass

Turn the concept into an inspected object.

The lower panels are one second act: keep the object fixed, inspect it visually, check source boundaries, practice transfer, then attach the research question.

ConceptFlashAttention: IO-Aware AttentionAttention & Transformers

01MechanismInspect motionSee how the idea changes form before adding more explanation.02VisualChoose a cueBind one visible mark to one mathematical role.03SourcesCheck boundary1 source link attached to the object.04ClaimAudit evidenceTreat the page claim as something to verify.05PracticeAttempt transferWrite before feedback so the model has something to repair.06RoomAsk groundedAttach the question to the same exact object.

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.

Demo notes open01 / Intuition

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

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 readyDemo notes 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.

Object - ConceptFlashAttention: IO-Aware AttentionQuestion

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

concept:attention-transformers/flash-attention

Boundary

sources: dao-2022-flashattention

Check

Open the closest source note before trusting the local explanation.

Evidence

1 selected-object source shown first; 1 reference total.

Next move

Audit the claim boundary, then ask from the same selected object.

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

Located CF editorial boundary

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.

Used here as

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...

Caveat

"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...

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.

Object - ConceptFlashAttention: IO-Aware AttentionQuestion

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

concept:attention-transformers/flash-attention

Boundary

sources: dao-2022-flashattention

Check

Treat every claim as provisional until source support and a local witness agree.

Evidence

1 structured claim check on this concept.

Next move

Run the prediction or practice transfer before asking for a grounded review.

1 CF editorial source-scope review recorded

Publisher-side editorial review is not independent replication. Claims without it still need exact source-support review. 1 reference and 3 local witnesses are available for inspection.

FlashAttention 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.

Registered CF editorial boundary

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-AwarenessPrimary source for the dense FlashAttention mechanism: IO-aware HBM/SRAM tiling plus block-wise/incremental softmax statistics reduce HBM reads/wri...

Used here as

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 atte...

Local witness

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

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.

Caveat

"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...

Review stateCF editorial source-scope reviewClaim metadata: source checkedPublisher-side editorial review only; not independent replication. 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

Practice notebook

Use the idea, then test it somewhere new

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.

AttemptNo learning claim inferred

Object - ConceptFlashAttention: IO-Aware AttentionQuestion

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

concept:attention-transformers/flash-attention

Boundary

sources: dao-2022-flashattention

Check

Use one state from FlashAttention: IO-Aware Attention to explain what changes, why it changes, and which assumption the explanation needs.

Evidence

No learner move yet; no learning state is inferred.

Next move

Write first, use only the help you need, then try a new case without it.

Explain

Use one state from FlashAttention: IO-Aware Attention to explain what changes, why it changes, and which assumption the explanation needs.

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

Attempt context

Learner move

Write an attempt before asking the companion.

Assistance used

0 of 3 progressive hints opened.

Evidence boundary

This draft and any AI response do not establish mastery; a later unassisted case can.

Misconception check

Grounded object roomClose

Selected object routeAsk from this object; carry one invariant back.sources: dao-2022-flashattention

ObjectConceptFlashAttention: IO-Aware Attention
PredictBefore revealFlashAttention: IO-Aware Attention prediction
WitnessCompare codeFlashAttention: IO-Aware Attention code witness 1
RoomAsk groundedChecking local snapshot

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 predictionManipulate 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?Checking local snapshot

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?

Source boundaryInspect source ids: dao-2022-flashattentionStable content-object key attached

Role lenses for this object

These are fixed, deterministic perspectives derived from the selected object. They do not represent people, community contributions, or independent review.

Learner evidence requestAsk what would make "FlashAttention: IO-Aware Attention" feel predictable rather than familiar.

Assumption

Source ids dao-2022-flashattention must support the exact object, not just the surrounding topic.

Source-checking summary

Treat this as a mechanism object: connect the definition to one equation, code witness, or demo before broadening the discussion.

Proposed experiment

Ask the learner to perturb one representation, then check whether the same invariant survives in math, code, and demo.

Next action

The learner can state the mechanism in their own words

Evidence4 checks

PredictionChecking carried observation

ActionReady for one action

AILearner handoff ready

Open source object

01PredictionChecking browser-local route memory

02EvidenceChecking for a carried observation

03BoundaryInspect source ids: dao-2022-flashattention

04Next moveSave one next action

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

Object-attached 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 Deterministic role lenses for this object: - Boundary: fixed perspectives, not people, community contributions, or independent review - Source-checking summary: Treat this as a mechanism object: connect the definition to one equation, code witness, or demo before broadening the discussion. - Proposed experiment: Ask the learner to perturb one representation, then check whether the same invariant survives in math, code, and demo. - Teach/transfer move: Turn the mechanism into one sentence that predicts a neighboring concept. - Assumptions: - Source ids dao-2022-flashattention must support the exact object, not just the surrounding topic. - The stable content-object key lets local drafts, prompts, and route memory attach without changing the source page. - The concept explanation is local atlas prose until checked against its math, code, and source support. - Prerequisite gaps should become a repair route, not a reason to leave the object vague. - Role-lens requests: - Learner: ask for "Ask what would make "FlashAttention: IO-Aware Attention" feel predictable rather than familiar." | assumption: Source ids dao-2022-flashattention must support the exact object, not just the surrounding topic. | next action: The learner can state the mechanism in their own words - Researcher: ask for "Source ids to inspect: dao-2022-flashattention" | assumption: The stable content-object key lets local drafts, prompts, and route memory attach without changing the source page. | next action: The learner can name the prerequisite that would repair confusion - Experimenter: ask for "Choose one variable or condition to perturb before asking for an explanation." | assumption: The concept explanation is local atlas prose until checked against its math, code, and source support. | next action: The learner can predict how the mechanism changes under one perturbation - Professor: ask for "Find the smallest transferable rule a learner could reuse without the AI." | assumption: Prerequisite gaps should become a repair route, not a reason to leave the object vague. | next action: Teach or transfer: Turn the mechanism into one sentence that predicts a neighboring concept. 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. Current deterministic role lens for this object: - Role lens: Learner - Evidence request: Ask what would make "FlashAttention: IO-Aware Attention" feel predictable rather than familiar. - Assumption to keep visible: Source ids dao-2022-flashattention must support the exact object, not just the surrounding topic. - Proposed experiment: Ask the learner to perturb one representation, then check whether the same invariant survives in math, code, and demo. - Next action: The learner can state the mechanism in their own words

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

Learnerpicture first Researchersource check Experimentertest the demo Professorname invariant

01Predict 02Ground 03Code 04Carry

01Intuitionready
02Mathready
03Codeready
04Interactive 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.

Copy-only tool: this panel composes a prompt for an AI tool you already use. Nothing you enter here is sent from this 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

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

Online softmax (streaming trick)

Concept: FlashAttention: IO-Aware Attention

FlashAttention: IO-Aware Attention

Carry the same object through intuition, math, code, and demo.

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.

Attach the question to an exact object

FlashAttention: IO-Aware Attention

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