Flash attention and efficient attention — Hinglish notes

4.1.14 · AI-ML › Transformer Architecture

The Standard Attention Problem

Computational flow:

$S = Q K^{T} \in R^{N \times N}$ compute karo — attention scores matrix
$P = softmax (S / d_{k})$ apply karo — attention probability matrix
Output $O = P V$ compute karo

Memory Hierarchy: Why Speed Isn't About FLOPs

Modern GPUs ki ek memory hierarchy hoti hai:

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

Key insight: HBM se SRAM mein data move karna on-chip SRAM mein move karne se ~12× slower (bandwidth mein) hai, isliye memory movement — arithmetic nahi — runtime dominate karta hai!

Example calculation:

N = 2048, d = 64 (typical)
HBM reads/writes: $204 8^{2} = 4 M$ floats = 16 MB (sirf S ke liye)
Memory transfer time: $16 MB /1.5 TB/s \approx 10 μ s$
Softmax elementwise work ka compute time: ~4 µs
Transfer akela compute ke same order ka time leta hai, aur kyunki standard attention repeatedly $N^{2}$ matrices read/write karta hai, memory movement (FLOPs nahi) true bottleneck ban jaata hai.

Flash Attention: Tiled Computation

Derivation: Incremental Softmax

Standard softmax ko poori row chahiye: $softmax (x_{i}) = \frac{e ^{x_{i}}}{\sum _{j} e ^{x_{j}}}$

Problem: Koi bhi output compute karne se pehle humein $\sum_{j} e^{x_{j}}$ chahiye!

Solution: Chunks mein process karo aur statistics incrementally update karo.

Step 1: Safe softmax (numerical stability)

Raw $e^{x}$ overflow karta hai. Standard trick: $softmax (x_{i}) = \frac{e ^{x_{i} - m}}{\sum _{j} e ^{x_{j} - m}}$ jahan $m = max (x)$ .

Step 2: Sum ko blocks mein split karo

Maano hum attention scores do blocks mein process karte hain: $x^{(1)}$ aur $x^{(2)}$ .

Block 1: $m^{(1)} = max (x^{(1)}), ℓ^{(1)} = \sum_{j \in block 1} e^{x_{j} - m^{(1)}}$

Block 2: $m^{(2)} = max (x^{(2)}), ℓ^{(2)} = \sum_{j \in block 2} e^{x_{j} - m^{(2)}}$

Step 3: Blocks merge karo

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

Rescale kyun? Humare sums alag "reference points" use karte the ( $m^{(1)}$ vs $m^{(2)}$ ). Humein ek common baseline chahiye.

Corrected sum: $ℓ^{new} = e^{m^{(1)} - m^{new}} ℓ^{(1)} + e^{m^{(2)} - m^{new}} ℓ^{(2)}$

Yeh kyun kaam karta hai: $ℓ^{new} = \sum_{j} e^{x_{j} - m^{new}}$ $= \sum_{j \in block 1} e^{x_{j} - m^{new}} + \sum_{j \in block 2} e^{x_{j} - m^{new}}$ $= \sum_{j \in block 1} e^{(x_{j} - m^{(1)}) + (m^{(1)} - m^{new})} + \sum_{j \in block 2} e^{(x_{j} - m^{(2)}) + (m^{(2)} - m^{new})}$ $= e^{m^{(1)} - m^{new}} ℓ^{(1)} + e^{m^{(2)} - m^{new}} ℓ^{(2)}$

Isi tarah, weighted output accumulate hota hai. Agar $\tilde{O}^{(b)} = \sum_{j \in block b} e^{x_{j} - m^{(b)}} v_{j}$ unnormalized block output hai, to: $O^{new} = \frac{e ^{m^{(1)} - m^{new}} O ~ ^{(1)} + e ^{m^{(2)} - m^{new}} O ~ ^{(2)}}{ℓ ^{new}}$

Yeh kisi bhi number of blocks tak generalize hota hai!

Worked example Flash Attention Forward Pass (Simplified)

Setup: N = 4 tokens, d = 2 dimensions, block size B = 2. Hum sirf query row 0 follow karte hain.

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: Tokens 0-1 ke liye keys/values

$q_{0} = [1, 0]$ , $K_{0 : 2}$ , $V_{0 : 2}$ SRAM mein load karo.
Scores $s = q_{0} K_{0 : 2}^{T} / 2 = [1 \cdot 1 + 0 \cdot 0, 1 \cdot 0 + 0 \cdot 1] / 2 = [1/ 2, 0] \approx [0.707, 0]$ . $d_{k}$ se divide kyun? Dimensions badhne par variance stable rakhta hai.
$m^{(1)} = max (s) = 1/ 2 \approx 0.707$ .
Max subtract karo: exponents $= [e^{0}, e^{- 1/ 2}] = [1, 0.493]$ , to $ℓ^{(1)} = 1 + 0.493 = 1.493$ . Yeh step kyun? Max-shift ke baad ek entry exactly $1$ hai aur doosri $e^{- 1/ 2} \approx 0.493$ hai; normalizer unka sum hai, nahi $[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]$ . Yeh step kyun? Block output values ka exponential-weighted sum hai, $V$ khud ka nahi — token 0 weight 1 contribute karta hai, token 1 weight 0.493 contribute karta hai.

Block 2: Tokens 2-3 ke liye keys/values

$K_{2 : 4}, V_{2 : 4}$ load karo, $q_{0}$ rakho.
Scores $s = q_{0} K_{2 : 4}^{T} / 2 = [1 \cdot 1 + 0 \cdot 1, 1 \cdot 0 + 0 \cdot 0] / 2 = [1/ 2, 0] \approx [0.707, 0]$ .
$m^{(2)} = 1/ 2 \approx 0.707$ ; exponents $[1, 0.493]$ , to $ℓ^{(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^{new} = max (0.707, 0.707) = 0.707$ , to dono rescale factors $e^{0.707 - 0.707} = 1$ hain yahaan. (Agar maxima alag hote, e.g. $m^{(1)} = 1/ 2$ vs ek bada $m^{(2)}$ , to block-1 factor $e^{1/ 2 - 2} \approx 0.493$ hota — yahaan rescaling matter karti hai.)
$ℓ^{new} = 1 \cdot 1.493 + 1 \cdot 1.493 = 2.986$ .
$O^{new} = \frac{1 \cdot [ 1.986 , 1.0 ] + 1 \cdot [ 0.493 , 2.493 ]}{2.986} = \frac{[ 2.479 , 3.493 ]}{2.986} \approx [0.830, 1.170]$ .

Key: Full 4×4 attention matrix kabhi store nahi kiya! Sirf 1×2 score blocks SRAM mein the, aur result row 0 ke liye exact standard attention ke barabar hai.

Figure — Flash attention and efficient attention

Backward Pass: Recomputation

Yeh kyun kaam karta hai:

Forward pass: Q,K,V ka 1 read + HBM mein writes
Standard backward: Stored S,P matrices read karo ( $N^{2}$ elements)
Flash backward: S,P SRAM mein recompute karo (koi HBM read nahi), sirf Q,K,V load karo

Memory bandwidth: Flash $3 N d$ reads karta hai vs standard $N d + 2 N^{2}$ reads. Long sequences ke liye, Flash massively jeetta hai.

Other Efficient Attention Variants

Worked example Local Attention Window

Window size $w = 3$ : har token khud ko + har taraf 1 neighbor ko attend karta hai.

N=6 ke liye:

[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 (N^{2})$ ki jagah $O (N w)$ . w=128 constant ke liye, yeh N mein linear hai!

Trade-off: Long-range dependencies ko information multiple layers ke through propagate karna padta hai (jaise ek convolutional network).

Practical Impact

Common mistake "Flash Attention Model Change Karta Hai"

Galat belief: "Flash Attention ek approximation hai jo speed ke liye accuracy sacrifice karta hai."

Kyun sahi lagta hai: Bahut saare optimizations (quantization, pruning, knowledge distillation) accuracy sacrifice karte hain.

Sachchi baat: Flash Attention standard attention ke numerically equivalent outputs compute karta hai! Yeh ek pure systems optimization hai — alag memory access pattern, same math. Sirf backward pass mein difference hai (recomputation), jisme floating-point rounding order ki wajah se tiny numerical differences ho sakte hain, lekin yeh negligible hain.

Proof: Dono $softmax (Q K^{T} / d_{k}) V$ exactly compute karte hain. Flash sirf computation ko chunk karta hai, aur incremental-softmax merge exact normalizer reproduce karta hai.

Recall 12-Saal-Ke-Bachche Ko Explain Karo

Socho tum ek 1000×1000 Sudoku puzzle solve kar rahe ho, aur tumhe check karna hai ki har cell kaise har doosri cell se relate karti hai — yeh 1 million comparisons hain!

Bura tarika: Saare 1 million results kagaz par likh lo, phir use karo. Lekin tumhara kagaz khatam ho jaata hai (memory)!

Flash Attention tarika: Puzzle ko chhote 100×100 chunks mein divide karo. Ek time mein ek chunk check karo, "maine ab tak kya dekha" ka ek running note rakho, phir chunk phenk do aur agla load karo. Tumhe kabhi ek million sheets of paper ki zaroorat nahi — sirf ek chunk ke liye kaafi!

Trick "running note" hai — tum apna summary update karte ho "maine ab tak sabse bada number kya dekha" aur "ab tak sum kya hai" bina har ek number keep kiye. Yahi incremental softmax math hai!

Kyun faster hai: Apne backpack mein paper andar bahar karna (memory) slow hai. Agar tum sab kuch apne haath mein rakh sako (fast memory), tum bahut faster kaam karte ho. Flash Attention chhote chunks par kaam karke sab kuch haath mein rakhta hai.

Connections

4.1.1-Self-attention-mechanism — Woh core operation jise Flash optimize karta hai
4.1.3-Multi-head-attention — Flash Attention per-head kaam karta hai
4.13-Scaling-laws-for-transformers — Efficient attention larger N enable karta hai, scaling laws affect karta hai
5.27-Gradient-checkpointing — Similar recomputation trade-off
4.1.8-Positional-encoding — Flash Attention positional information preserve karta hai
4.2.1-BERT-architecture — BERT long documents ke liye Flash se benefit karta hai
4.2.3-GPT-architecture — GPT-3/4 long context ke liye Flash-like optimizations use karte hain

#flashcards/ai-ml

Standard attention vs Flash Attention ki memory complexity kya hai?

Standard:

O (N^{2})

attention matrix store karne ke liye. Flash:

O (N)

tiling se aur full matrix materialize kiye bina.

Flash Attention zyada FLOPs karne ke bawajood faster kyun hai?

Memory bandwidth bottleneck hai, compute nahi. Flash HBM accesses

O (N^{2})

O (N)

tak reduce karta hai, jo recomputation se thoda FLOPs badhne par dominate karta hai.

Flash Attention mein incremental softmax kya hai?

Running max (

m

) aur sum (

ℓ

) maintain karke chunks mein softmax compute karna, phir blocks merge karna rescaling se:

ℓ^{new} = e^{m^{(1)} - m^{new}} ℓ^{(1)} + e^{m^{(2)} - m^{new}} ℓ^{(2)}

Kya Flash Attention standard attention ke jaisa hi output produce karta hai?

Haan, forward pass mein numerically equivalent. Backward pass mein recomputation/rounding order ki wajah se tiny differences, lekin negligible hain.

Window size

w

ke saath local attention ki complexity kya hai?

O (N w)

. Constant

w

ke liye, yeh sequence length

N

mein linear hai, full attention ke

O (N^{2})

ke comparison mein.

Linear attention

O (N)

complexity kaise achieve karta hai?

Kernel trick use karke:

(Q K^{T}) V

ki jagah

ϕ (Q) (ϕ (K)^{T} V)

. Yeh operations ka order

O (N^{2} d)

O (N d^{2})

mein change karta hai — N mein linear agar

d

constant ho.

Flash Attention ke backward pass mein main trade-off kya hai?

Attention matrices (

S

P

) store karne ki jagah recompute karo. ~20% compute badhta hai lekin

O (N^{2})

memory bachti hai, jo ek huge win hai kyunki memory bandwidth bottleneck hai.

Flash Attention speedup sequence length ke saath kyun badhta hai?

Standard attention ki memory cost

O (N^{2})

badhti hai, Flash ki

O (N)

. Absolute gap longer sequences ke liye aur bada hota hai.