One-page visual reference

FlashAttention Cheat Sheet

How a logical B×H×L×L attention problem becomes many independent CTA tiles, how online softmax stays globally exact without storing the score matrix, and exactly where Skip Softmax saves work.

HBM / data movement Tensor Core computation On-chip online state Optional sparse path

1. The whole operation

\begin{aligned} Q &\in \mathbb{R}^{B\times H_q\times L_q\times D}, & K,V &\in \mathbb{R}^{B\times H_{kv}\times L_k\times D}, \\ S &= \frac{QK^{\mathsf T}}{\sqrt D}, & P &= \operatorname{softmax}(S), & O &= PV. \end{aligned}

Logical score tensor: [B, Hq, Lq, Lk] — FlashAttention never materializes it in HBM.

2. Logical tensor versus one CTA tile

Logical collection

Q tiles: [B, Hq, BM, D]

This describes all independent batch/head/query-block slices together.

One CTA / Triton program

Q[b,h,m:m+BM,:] → [BM,D]

Batch, head and query block are fixed by the launch-grid coordinates.

Its matrix products

[BM,D]×[D,BN]→[BM,BN]

[BM,BN]×[BN,D]→[BM,D]

launch grid →

b0 · h0 · qblock0

b0 · h0 · qblock1

b0 · h1 · qblock0

b0 · h1 · qblock1

more CTAs →

b1 · h0 · qblock0

b1 · h0 · qblock1

b1 · h1 · qblock0

b1 · h1 · qblock1

The GPU processes the batch concurrently by scheduling many CTAs—not by mixing different sequences or heads into one matrix multiply.

3. What an attention tile means

One head's logical Lq×Lk score matrix, partitioned into BM×BN tiles.

Q tile[BM,D]

Kᵀ tile[D,BN]

→

Score tile[BM,BN]

P tile[BM,BN]

V tile[BN,D]

→

Output update[BM,D]

Typical values are BM=64/128, BN=64/128, and D=64/128. Actual choices depend on architecture, datatype and kernel tuning.

4. The complete FlashAttention flow

Q/K/V in HBMFull tensors, often paged for decode

→

Load one Q tileKeep it resident for the CTA

→

Stream K tileTensor Core QKᵀ

→

Online softmaxUpdate m, l and accumulator

→

Stream V tileTensor Core PV

↻

Fix one batch item, query head and query block.
Load Q[BM,D] and initialize per-query-row state.

Loop over KV blocks.
Load K[BN,D], compute a temporary S[BM,BN], and apply causal/padding masks.

Merge the tile into global online softmax.
Maintain running maximum m[BM], denominator l[BM], and output accumulator A[BM,D].

Load the matching V tile and accumulate PV.
Temporary scores and probabilities die on-chip; the CTA advances to the next KV tile.

Normalize and store one output tile.
O=A/l is written to HBM only after every valid KV tile has been merged.

5. Why online softmax is globally exact

State after processed tiles, for each query row i:

\begin{aligned} m_i &= \max_{j\in\mathcal J} s_{ij}, \\ \ell_i &= \sum_{j\in\mathcal J} e^{s_{ij}-m_i}, \\ A_i &= \sum_{j\in\mathcal J} e^{s_{ij}-m_i}v_j. \end{aligned}

Merge the next tile, whose local maximum is m tilde:

\begin{aligned} m_i' &= \max(m_i,\widetilde m_i), & \alpha_i &= e^{m_i-m_i'}, \\ \ell_i' &= \alpha_i\ell_i + \sum_{j\in\mathcal B}e^{s_{ij}-m_i'}, \\ A_i' &= \alpha_i A_i + \sum_{j\in\mathcal B}e^{s_{ij}-m_i'}v_j. \end{aligned}

O_i=\frac{A_i}{\ell_i}.

The rescaling is the trick

If a later tile reveals a larger maximum, multiplying old state by exp(m_old−m_new) moves it into the new coordinate system.

No approximation is introduced: tile order can change, but the final global-softmax result is the same, up to floating-point numerics.

6. Tiny numerical example

Scores [2, 1, 3, 0], scalar values [10, 20, 30, 40], processed as two 2-token tiles.

Tile 1: scores [2,1]

m_1=2

\ell_1=1+e^{-1}\approx1.3679

A_1=10+e^{-1}\cdot20\approx17.358

Tile 2: scores [3,0]

m_2=3,\quad\alpha=e^{-1}\approx0.3679

\ell_2=\alpha\ell_1+1+e^{-3}\approx1.5530

A_2=\alpha A_1+30+e^{-3}\cdot40\approx38.376

O=\frac{A_2}{\ell_2}=\frac{38.376}{1.5530}\approx24.71.

Identical to one global softmax over all four scores.

7. Why it is faster: avoid the quadratic HBM round-trip

Materialized attention

Compute QKᵀ

Write full S to HBM

Read S; write/read full P

Compute PV

Write O to HBM

FlashAttention

Compute one QK tile on-chip

Update m, l and A on-chip

Compute one PV tile on-chip

Discard temporary S/P; repeat

Write O once

FlashAttention does not make dense attention sparse. It still evaluates all valid QK and PV pairs; its speed comes primarily from IO-aware tiling and fusion.

8. Where Skip Softmax enters

Compute exact QK tileK is already loaded

→

Compare maximam̃_tile vs running m

→

Negligible for every valid Q row?

e^{\widetilde m-m}<\lambda

↙ ↘

Keepsoftmax update + load V + PV

Skipno softmax update, V load or PV

The decision scope is one batch item × head × Q block × KV block. During prefill, all valid rows in the Q block must agree. During decode, there is usually one valid query row.

Work	Normal FlashAttention tile	Skipped tile
Load K and compute QK	Yes	Yes—needed to decide
Exponentiation / online-softmax update	Yes	Avoided
Load V from HBM	Yes	Avoided
Compute PV / BMM2	Yes	Avoided

9. Prefill versus decode

Prefill

BM ≈ 64/128 real query rows

Many Q tiles run in parallel. The workload is often compute-heavy. Skip Softmax saves softmax and PV, but one important query row can force the whole tile to remain dense.

Decode

one real query row per request

The query scans a long KV cache and is often bandwidth-bound. Avoiding V loads is the main Skip Softmax benefit; padded rows do not vote against skipping.

10. Shape and terminology quick reference

Symbol / term	Meaning	Typical shape or scope
B	Inference/training batch	Parallel grid dimension, not mixed into one attention GEMM
Hq / Hkv	Query and KV heads	In GQA, several query heads share one KV head
BM	Query-token block	Rows of the score tile
BN	KV-token block	Columns of the score tile
D	Per-head dimension	Reduction axis of QK and output axis of PV
CTA / Triton program	Cooperating GPU work unit	Usually one b,h,qblock, looping over KV blocks
Online state	Enough information to merge softmax tiles	m[BM], l[BM], A[BM,D]