flash attention

Motivation

Transformers are slow and memory-hungry on long sequences due to quadratic time and memory complexity of self-attention in sequence length. The flash attention paper propose an IO-aware and memory efficient algorithm that implements exact attention.

Hardware prerequisites

• GPU SRAM(Static Random-Access Memory): 19TB/s(20MB)
• GPU HBM(High Bandwidth Memory): 1.5TB/s(40GB)
• CPU DRAM: 12.8GB/s(>1TB)

Compute/memory bandwidths in attention calculation

finish the note

Consider A100-40GB SXM, which has a compute bandwidth of 312TFLOPS and memory bandwidth of 1555GB/s and mixed precision training, the operational intensity for the particular hardware is

$$\begin{array}{r}\frac{\pi }{\beta }=\frac{312\ast 1e12}{1555\ast 1e9}=201\text{FLOPS/Bytes}\end{array}$$

Suppose we are to calculate attention for $S=Q{K}^T$,

$$\begin{array}{r}\frac{{\pi }_t}{{\beta }_t}=\frac{2N^{2}d}{2Nd+2Nd+2N^2}=\frac{N^{2}d}{2Nd+N^2}\end{array}$$

For more details refer to this wonderful blog. The bottleneck may vary for different values of $N$ and $d$.

$N$	$d$	operations/bytes	bottleneck
256	128	64	memory-bound
1024	128	102	memory-bound
4096	128	120	memory-bound
256	256	85	memory-bound
1024	256	171	memory-bound
4096	256	228	compute-bound
256	512	102	memory-bound
1024	512	256	compute-bound
4096	512	410	compute-bound

Online softmax

Safe softmax

Vanilla softmax could encounter numerical instability, hence a safe version of softmax is used in flash attention where the maximum of the vector is subtracted before exponentiation.

$$\begin{array}{r}\text{softmax}(x_i)=\frac{\exp \{x_i-max(x)\}}{{\sum }_{j=1}^n\exp \{x_j-max(x)\}}\end{array}$$

3-pass softmax

For $i=1,\cdots ,N$,

$$\begin{array}{r}{m}_i:=\max (m_{i-1},x_i)\end{array}$$

For $i=1,\cdots ,N$,

$$\begin{array}{r}{d}_i:=d_{i-1}+e^{x_i-m_N}\end{array}$$

For $i=1,\cdots ,N$,

$$\begin{array}{r}{a}_i:=e^{x_i-m_N}/d_N\end{array}$$

2-pass softmax

Reference

• 图解大模型计算加速系列：Flash Attention V1，从硬件到计算逻辑