Flash attention and efficient attention

The Standard Attention Problem

The computational flow:

Compute $S = QK^T \in \mathbb{R}^{N \times N}$ — the attention scores matrix
Apply $P = \text{softmax}(S / \sqrt{d_k})$ — the attention probability matrix
Compute output $O = PV$

Memory Hierarchy: Why Speed Isn't About FLOPs

Modern GPUs have a memory hierarchy:

SRAM (on-chip, fast): ~20 MB, ~19 TB/s bandwidth
HBM (GPU RAM, slow): ~40 GB, ~1.5 TB/s bandwidth

Key insight: Moving data from HBM to SRAM is ~12× slower (in bandwidth) than moving within on-chip SRAM, so memory movement—not arithmetic—dominates runtime!

Example calculation:

N = 2048, d = 64 (typical)
HBM reads/writes: $2048^2 = 4M$ floats = 16 MB (just for S)
Memory transfer time: $16\text{ MB} / 1.5\text{ TB/s} \approx 10\ \mu s$
Compute time for the softmax elementwise work: ~4 µs
The transfer alone takes about the same order as compute, and because standard attention repeatedly reads/writes these $N^2$ matrices, memory movement (not FLOPs) becomes the true bottleneck.

Flash Attention: Tiled Computation

Derivation: Incremental Softmax

Standard softmax requires the full row: $\text{softmax}(x_i) = \frac{e^{x_i}}{\sum_j e^{x_j}}$

Problem: We need $\sum_j e^{x_j}$ before computing any output!

Solution: Process in chunks and update statistics incrementally.

Step 1: Safe softmax (numerical stability)

Raw $e^{x}$ overflows. Standard trick: $\text{softmax}(x_i) = \frac{e^{x_i - m}}{\sum_j e^{x_j - m}}$ where $m = \max(x)$ .

Step 2: Split sum across blocks

Suppose we process attention scores in two blocks: $x^{(1)}$ and $x^{(2)}$ .

Block 1: $m^{(1)} = \max(x^{(1)}), \quad \ell^{(1)} = \sum_{j \in \text{block 1}} e^{x_j - m^{(1)}}$

Block 2: $m^{(2)} = \max(x^{(2)}), \quad \ell^{(2)} = \sum_{j \in \text{block 2}} e^{x_j - m^{(2)}}$

Step 3: Merge blocks

Global max: $m^{\text{new}} = \max(m^{(1)}, m^{(2)})$

Why rescale? Our sums used different "reference points" ( $m^{(1)}$ vs $m^{(2)}$ ). We need a common baseline.

Corrected sum: $\ell^{\text{new}} = e^{m^{(1)} - m^{\text{new}}}\, \ell^{(1)} + e^{m^{(2)} - m^{\text{new}}}\, \ell^{(2)}$

Why this works: $\ell^{\text{new}} = \sum_{j} e^{x_j - m^{\text{new}}}$ $= \sum_{j \in \text{block 1}} e^{x_j - m^{\text{new}}} + \sum_{j \in \text{block 2}} e^{x_j - m^{\text{new}}}$ $= \sum_{j \in \text{block 1}} e^{(x_j - m^{(1)}) + (m^{(1)} - m^{\text{new}})} + \sum_{j \in \text{block 2}} e^{(x_j - m^{(2)}) + (m^{(2)} - m^{\text{new}})}$ $= e^{m^{(1)} - m^{\text{new}}}\, \ell^{(1)} + e^{m^{(2)} - m^{\text{new}}}\, \ell^{(2)}$

Similarly, the weighted output accumulates. If $\tilde{O}^{(b)} = \sum_{j \in \text{block } b} e^{x_j - m^{(b)}}\, v_j$ is the unnormalized block output, then: $O^{\text{new}} = \frac{e^{m^{(1)} - m^{\text{new}}}\, \tilde{O}^{(1)} + e^{m^{(2)} - m^{\text{new}}}\, \tilde{O}^{(2)}}{\ell^{\text{new}}}$

This generalizes to any number of blocks!

Worked example Flash Attention Forward Pass (Simplified)

Setup: N = 4 tokens, d = 2 dimensions, block size B = 2. We follow query row 0 only.

Q = [[1, 0],    K = [[1, 0],    V = [[1, 1],
     [0, 1],         [0, 1],         [2, 0],
     [1, 1],         [1, 1],         [0, 2],
     [0, 0]]         [0, 0]]         [1, 1]]

Block 1: Keys/values for tokens 0-1

Load $q_0 = [1,0]$ , $K_{0:2}$ , $V_{0:2}$ into SRAM.
Scores $s = q_0 K_{0:2}^T / \sqrt{2} = [\,1{\cdot}1 + 0{\cdot}0,\ 1{\cdot}0 + 0{\cdot}1\,]/\sqrt 2 = [\,1/\sqrt2,\ 0\,] \approx [0.707,\ 0]$ . Why divide by $\sqrt{d_k}$ ? Keeps variance stable as dimensions grow.
$m^{(1)} = \max(s) = 1/\sqrt 2 \approx 0.707$ .
Subtract the max: exponents $= [\,e^{0},\ e^{-1/\sqrt2}\,] = [\,1,\ 0.493\,]$ , so $\ell^{(1)} = 1 + 0.493 = 1.493$ . Why this step? After the max-shift one entry is exactly $1$ and the other is $e^{-1/\sqrt2}\approx 0.493$ ; the normalizer is their sum, not $[1,1]$ .
Unnormalized output $\tilde{O}^{(1)} = 1\cdot v_0 + 0.493\cdot v_1 = 1\cdot[1,1] + 0.493\cdot[2,0] = [1.986,\ 1.0]$ . Why this step? The block output is the exponential-weighted sum of the values, not $V$ itself — token 0 contributes weight 1, token 1 contributes weight 0.493.

Block 2: Keys/values for tokens 2-3

Load $K_{2:4}, V_{2:4}$ , keep $q_0$ .
Scores $s = q_0 K_{2:4}^T/\sqrt 2 = [\,1{\cdot}1+0{\cdot}1,\ 1{\cdot}0+0{\cdot}0\,]/\sqrt2 = [\,1/\sqrt2,\ 0\,] \approx [0.707,\ 0]$ .
$m^{(2)} = 1/\sqrt2 \approx 0.707$ ; exponents $[1,\ 0.493]$ , so $\ell^{(2)} = 1.493$ .
Unnormalized output $\tilde{O}^{(2)} = 1\cdot v_2 + 0.493\cdot v_3 = 1\cdot[0,2] + 0.493\cdot[1,1] = [0.493,\ 2.493]$ .

Merge:

$m^{\text{new}} = \max(0.707, 0.707) = 0.707$ , so both rescale factors are $e^{0.707-0.707}=1$ here. (If the maxima differed, e.g. $m^{(1)}=1/\sqrt2$ vs a larger $m^{(2)}$ , the block-1 factor would be $e^{1/\sqrt2 - \sqrt2} \approx 0.493$ — this is where rescaling matters.)
$\ell^{\text{new}} = 1\cdot 1.493 + 1\cdot 1.493 = 2.986$ .
$O^{\text{new}} = \dfrac{1\cdot[1.986,1.0] + 1\cdot[0.493,2.493]}{2.986} = \dfrac{[2.479,\ 3.493]}{2.986} \approx [0.830,\ 1.170]$ .

Key: Never stored the full 4×4 attention matrix! Only 1×2 score blocks lived in SRAM, and the result equals exact standard attention for row 0.

Figure — Flash attention and efficient attention

Backward Pass: Recomputation

Why this works:

Forward pass: 1 read of Q,K,V + writes to HBM
Standard backward: Read stored S,P matrices ( $N^2$ elements)
Flash backward: Recompute S,P in SRAM (no HBM read), only load Q,K,V

Memory bandwidth: Flash reads $3Nd$ vs standard reads $Nd + 2N^2$ . For long sequences, Flash wins massively.

Other Efficient Attention Variants

Worked example Local Attention Window

Window size $w = 3$ : each token attends to itself + 1 neighbor on each side.

For N=6:

[1 1 0 0 0 0]  ← token 0 attends to 0,1
[1 1 1 0 0 0]  ← token 1 attends to 0,1,2
[0 1 1 1 0 0]  ← token 2 attends to 1,2,3
[0 0 1 1 1 0]
[0 0 0 1 1 1]
[0 0 0 0 1 1]

Complexity: $O(Nw)$ instead of $O(N^2)$ . For w=128 constant, this is linear in N!

Trade-off: Long-range dependencies require information to propagate through multiple layers (like a convolutional network).

Practical Impact

Recall Explain to a 12-Year-Old

Imagine you're solving a 1000×1000 Sudoku puzzle, and you need to check how every cell relates to every other cell—that's 1 million comparisons!

Bad way: Write down all 1 million results on paper, then use them. But you run out of paper (memory)!

Flash Attention way: Divide the puzzle into smaller 100×100 chunks. Check one chunk at a time, keep a running note of "what I've seen so far," then throw away the chunk and load the next one. You never need a million sheets of paper—just enough for one chunk!

The trick is the "running note"—you update your summary of "what's the biggest number I've seen" and "what's the sum so far" without keeping every single number. That's the incremental softmax math!

Why it's faster: Moving paper in and out of your backpack (memory) is slow. If you can keep everything in your hands (fast memory), you work much faster. Flash Attention keeps everything in your hands by working on smaller chunks.

Connections

4.1.1-Self-attention-mechanism — The core operation Flash optimizes
4.1.3-Multi-head-attention — Flash Attention works per-head
4.13-Scaling-laws-for-transformers — Efficient attention enables larger N, affecting scaling laws
5.27-Gradient-checkpointing — Similar recomputation trade-off
4.1.8-Positional-encoding — Flash Attention preserves positional information
4.2.1-BERT-architecture — BERT benefits from Flash for long documents
4.2.3-GPT-architecture — GPT-3/4 use Flash-like optimizations for long context

#flashcards/ai-ml

What is the memory complexity of standard attention vs Flash Attention?

Standard:

O(N^2)

for storing attention matrix. Flash:

O(N)

by tiling and never materializing full matrix.

Why is Flash Attention faster despite doing more FLOPs?

Memory bandwidth is the bottleneck, not compute. Flash reduces HBM accesses from

O(N^2)

O(N)

, which dominates the small increase in FLOPs from recomputation.

What is incremental softmax in Flash Attention?

Computing softmax in chunks by maintaining running max (

m

) and sum (

\ell

), then merging blocks by rescaling:

\ell^{\text{new}} = e^{m^{(1)} - m^{\text{new}}} \ell^{(1)} + e^{m^{(2)} - m^{\text{new}}} \ell^{(2)}

Does Flash Attention produce the same output as standard attention?

Yes, numerically equivalent in forward pass. Tiny differences in backward pass due to recomputation/rounding order, but negligible.

What is the complexity of local attention with window size

w

O(Nw)

. For constant

w

, this is linear in sequence length

N

, compared to

O(N^2)

for full attention.

How does linear attention achieve

O(N)

complexity?

By using kernel trick:

\phi(Q)(\phi(K)^T V)

instead of

(QK^T)V

. This changes order of operations from

O(N^2 d)

O(Nd^2)

—linear in

N

d

constant.

What is the main trade-off in Flash Attention's backward pass?

Recompute attention matrices (

S

P

) instead of storing them. Adds ~20% compute but saves

O(N^2)

memory, which is a huge win since memory bandwidth is the bottleneck.

Why does Flash Attention speedup increase with sequence length?

Standard attention's memory cost grows as

O(N^2)

, Flash's as

O(N)

. The absolute gap widens for longer sequences.

Concept Map

Hinglish (regional understanding)

Intuition Hinglish mein samjho

Hinglish (regional understanding)

Intuition Hinglish mein samjho

Chalo is topic ka core intuition samajhte hain. Transformer mein jab attention calculate hota hai, to woh ek bahut bada matrix banata hai jiska size sequence_length × sequence_length hota hai. Matlab agar tumhare paas 2048 tokens hai, to yeh matrix 2048×2048 ka ho jaata hai — bahut sara memory kha jaata hai! Normal attention pehle poora yeh matrix GPU ki memory mein store karta hai, phir use karta hai. Yeh waise hi hai jaise tum poori library ki photocopy nikalo bina yeh soche ki tumhe sirf kuch pages chahiye. Flash Attention ka core idea yeh hai ki poora matrix banao hi mat — chhote-chhote tiles (blocks) mein kaam karo, ek kitaab ki tarah page-by-page padhte jao.

Ab yeh "why it matters" wali baat samajhna zaroori hai. Log sochte hain ki GPU fast hone ka matlab hai zyada calculations (FLOPs) karna, lekin asli bottleneck compute nahi, balki memory movement hai. GPU mein do type ki memory hoti hai — SRAM (on-chip, super fast, chhoti) aur HBM (GPU RAM, badi lekin slow). Jab bar-bar yeh N×N matrix ko HBM mein likhna aur wapas padhna padta hai, to yeh data ka aana-jaana hi sabse zyada time leta hai, chahe calculation kitni bhi jaldi ho jaaye. Isliye Flash Attention in memory round-trips ko hataake attention ko 2-4× faster bana deta hai aur memory bhi O(N²) se O(N) tak kam kar deta hai.

Iska practical faayda yeh hai ki ab hum bahut lambi sequences ko process kar sakte hain jo pehle memory ki wajah se possible hi nahi thi. Flash Attention ek clever trick use karta hai jise "incremental softmax" kehte hain — softmax ke liye normally poori row chahiye hoti hai, lekin yeh running statistics maintain karke chunks mein kaam kar leta hai, aur backward pass mein values ko store karne ke bajaye on-the-fly recompute kar leta hai. Yeh yaad rakhna important hai kyunki modern LLMs jaise GPT aur LLaMA sab isi technique par chalte hain — yeh sirf theory nahi, real-world efficiency ka backbone hai jo AI ko scalable banata hai.

Flash attention and efficient attention

The Standard Attention Problem

Memory Hierarchy: Why Speed Isn't About FLOPs

Flash Attention: Tiled Computation

Derivation: Incremental Softmax

Backward Pass: Recomputation

Other Efficient Attention Variants

Practical Impact

Connections

Concept Map

Hinglish (regional understanding)

Hinglish (regional understanding)

Go deeper — visual, from zero

Test yourself — Transformer Architecture

Connections