Exercises — Flash attention and efficient attention

$S$ is $N\times N=2048\times 2048=4,194,304$ entries. Bytes $=4,194,304\times 4=16,777,216$ bytes $=16$ MB. What we did: counted rows$\times$columns (that is what "$N\times N$" means) and multiplied by 4 bytes. Why: this is the object Flash Attention refuses to store — the whole point.

Answer: (c). $Q,K,V$ and the output $O$ still live in HBM — you must read inputs and write the answer. What Flash Attention avoids is materialising the full $N\times N$ matrices $S$ and $P$; they only ever appear as small tiles inside fast on-chip SRAM.

Standard: $O(N^2)$ (it stores the full square scores matrix). Flash: $O(N)$ (it keeps only running statistics $m,\ell ,\tilde{O}$ per query row plus one block at a time).

$m=\max (x)=5$. Shifted exponents: $e^{3-5}=e^{-2}=0.1353,e^{5-5}=1,e^{5-5}=1,e^{2-5}=e^{-3}=\mathrm{0.0498.}$ $\ell =0.1353+1+1+0.0498=\mathrm{2.1851.}$ Why subtract $m$: $e^5\approx 148$ is fine, but for larger scores $e^x$ overflows; shifting so the biggest exponent is $e^0=1$ keeps every term in $[0,1]$ and never changes the final softmax (top and bottom both get multiplied by $e^{-m}$, which cancels).

Using $\ell =2.1851$ from 2.1: $p=[0.1353,1,1,0.0498]/2.1851=[0.06192,0.45764,0.45764,0.02279].$ Output $=0.06192(10)+0.45764(20)+0.45764(30)+0.02279(40)$ $=0.6192+9.1528+13.7292+0.9117=\mathrm{24.413.}$ What it means: the two large-score tokens (both $5$) split most of the weight, dragging the output toward their values $20$ and $30$; the weak tokens barely count.

Block 1: $m^{(1)}=5$, ${\ell }^{(1)}=e^{3-5}+e^{5-5}=0.1353+1=1.1353$. Block 2: $m^{(2)}=5$, ${\ell }^{(2)}=e^{5-5}+e^{2-5}=1+0.0498=1.0498$. Merge: $m^{\text{new}}=\max (5,5)=5$. Rescale factors $e^{m^{(1)}-m^{\text{new}}}=e^0=1$ and $e^{m^{(2)}-m^{\text{new}}}=1$. ${\ell }^{\text{new}}=1\cdot 1.1353+1\cdot 1.0498=\mathrm{2.1851.}$ ✓ Why the rescale: each block measured its exponents against its own local max; before adding sums we must re-express them against the common global max $m^{\text{new}}$ — the factor $e^{m^{(b)}-m^{\text{new}}}$ does exactly that shift.

$m^{\text{new}}=\max (6,2)=6$. Rescale factors: $e^{6-6}=1$ for block 1, $e^{2-6}=e^{-4}=0.018316$ for block 2. ${\ell }^{\text{new}}=1\cdot 1.8+0.018316\cdot 3.0=1.8+0.054948=\mathrm{1.854948.}$ What it looks like: block 2's whole sum ($3.0$) shrinks to $0.055$ because its scores were tiny compared with block 1's peak — exactly the desired behaviour: far-below-max tokens contribute almost nothing.

Numerator $=e^{m^{(1)}-m^{\text{new}}}{\tilde{O}}^{(1)}+e^{m^{(2)}-m^{\text{new}}}{\tilde{O}}^{(2)}$ $=1\cdot [1.8,0.9]+0.018316\cdot [6.0,3.0]=[1.8+0.10990,0.9+0.054948]=[1.90990,0.954948].$ $O^{\text{new}}=[1.90990,0.954948]/{\ell }^{\text{new}}=[1.90990,0.954948]/1.854948=[1.02962,0.514810].$ Why: the numerator ($\tilde{O}$) and denominator ($\ell$) are rescaled by the same factors, so dividing at the very end gives the exact softmax-weighted average — proving you can accumulate one block at a time and never store the full matrix.

$N^2=4096^2=16,777,216$. $Nd=4096\times 64=262,144$. Fraction $=\frac{N^2}{N^2+Nd}=\frac{16,777,216}{16,777,216+262,144}=\frac{16,777,216}{17,039,360}=\mathrm{0.98462.}$ So about 98.5% of memory traffic is the quadratic term. Insight: since $N\gg d$, killing the $N^2$ round-trips (what Flash Attention does) removes essentially all the memory cost — the linear $Q,K,V$ traffic is a rounding error.

Start with block 1: $m=4,\ell =1.5$. Fold in block 2 ($m^{(2)}=7,{\ell }^{(2)}=2.0$): $m^{\text{new}}=\max (4,7)=7$. $\ell =e^{4-7}(1.5)+e^{7-7}(2.0)=e^{-3}(1.5)+2.0=0.049787(1.5)+2.0=0.074681+2.0=\mathrm{2.074681.}$ Running: $m=7,\ell =2.074681$. Fold in block 3 ($m^{(3)}=5,{\ell }^{(3)}=1.2$): $m^{\text{new}}=\max (7,5)=7$. $\ell =e^{7-7}(2.074681)+e^{5-7}(1.2)=2.074681+e^{-2}(1.2)=2.074681+0.135335(1.2)$ $=2.074681+0.162402=\mathrm{2.237083.}$ Final: $m=7,\ell =2.237083$. Why online works: each merge only ever needs the current running $(m,\ell )$ and the new block $(m^{(b)},{\ell }^{(b)})$ — never the raw scores. That is why memory stays $O(N)$: the running pair is a fixed-size summary of everything seen so far.

$e^{-3}(1.5)+1(2.0)+e^{-2}(1.2)=0.074681+2.0+0.162402=\mathrm{2.237083.}$ ✓ Identical to the online result — order of folding does not matter, because merging is associative. This is the correctness proof of Flash Attention in one line.

Block 1: $m^{(1)}=2$. Exponents $e^{2-2}=1,e^{0-2}=e^{-2}=0.135335$. ${\ell }^{(1)}=1.135335$. ${\tilde{O}}^{(1)}=1\cdot [1,0]+0.135335\cdot [0,1]=[1,0.135335].$ Block 2: $m^{(2)}=3$. Exponents $e^{3-3}=1,e^{1-3}=e^{-2}=0.135335$. ${\ell }^{(2)}=1.135335$. ${\tilde{O}}^{(2)}=1\cdot [1,1]+0.135335\cdot [2,0]=[1.270670,1].$ Merge: $m^{\text{new}}=\max (2,3)=3$. Factors: block 1 $\to e^{2-3}=e^{-1}=0.367879$; block 2 $\to e^{3-3}=1$. ${\ell }^{\text{new}}=0.367879(1.135335)+1(1.135335)=0.417668+1.135335=\mathrm{1.553003.}$ Numerator $=0.367879\cdot [1,0.135335]+1\cdot [1.270670,1]$ $=[0.367879,0.049788]+[1.270670,1]=[1.638549,1.049788].$ $O=[1.638549,1.049788]/1.553003=[1.055083,0.675974].$

Direct check (plain softmax over $[2,0,3,1]$): $m=3$, exps $[e^{-1},e^{-3},e^0,e^{-2}]=[0.367879,0.049787,1,0.135335]$, sum $=1.553001$. Weights $[0.236882,0.032056,0.643897,0.087141]$. $O=0.236882[1,0]+0.032056[0,1]+0.643897[1,1]+0.087141[2,0]$ $=[0.236882+0.643897+0.174282,0.032056+0.643897]$ $=[1.055061,0.675953].$ Matches the tiled result to rounding ✓ — Flash Attention returns exactly standard attention, just without ever storing the full row.

With equal scores every $e^{s_j-m}=e^{c-c}=1$, so $\ell =4$ and each weight $=1/4$. Output $=\frac{1}{4}(v_0+v_1+v_2+v_3)$ — the plain average of the values, independent of $c$. Why: softmax cares only about score differences; a common shift $c$ cancels. This is the same cancellation that makes the safe-softmax max-subtraction harmless, and it is the "no information" limit — uniform attention.

Naive: $e^{800}$ and $e^{801}$ both overflow to $\infty$ in float; $\infty /\infty =$ NaN — the network breaks. Safe: $m=801$. $e^{800-801}=e^{-1}=0.367879$, $e^{801-801}=1$. $\ell =1.367879$. Weights $[0.367879/1.367879,1/1.367879]=[0.268941,0.731059]$. Output $=0.268941(10)+0.731059(20)=2.68941+14.62119=\mathrm{17.31060.}$ Why it works: subtracting the max guarantees the largest exponent is $e^0=1$ and all others are in $(0,1]$ — never overflow, identical result. Flash Attention must do this per block, which is exactly why it tracks $m^{(b)}$ at all.