Dao 2022 - Flash Attention

Paper Link

This paper argues that the attention mechanism is slow because of reading / writing between GPU High Bandwidth Memory and GPU on-chip SRAM. The authors hence create a block-wise attention algorithm that minimizes such IO read / writes and speeds up attention significantly especially when the sequence length is long.

Brief Overview of Attention

Suppose we have an input sequence of embeddings $X=(x_1,...,x_N)$ where $x_i\in {\mathbb{R}}^k$, such that $X\in {\mathbb{R}}^{k\times N}$. Naively, we can compute activations by $V=X^T\cdot W_v$, where $W_v\in {\mathbb{R}}^{k\times d}$, such that $V\in {\mathbb{R}}^{N\times d}$. However, this naive way of encoding our input sequence does not allow interaction between inputs at different positions (say $x_i$ with $x_j$). We can see this by observing that the first row of $V$ is only affected by the first column of $X$ (i.e. first encoding $x_1$), and likewise for all the other positions.

Attention addresses this problem by adding an interaction mechanism. Besides $W_v$, we also create weight parameters $W_q,W_k\in {\mathbb{R}}^{k\times d}$. Given an input $X$, we compute $Q,K,V\in {\mathbb{R}}^{N\times d}$ as follows:

• $Q=X^T\cdot W_q$
• $K=X^T\cdot W_k$
• $V=X^T\cdot W_v$

We then create an interaction matrix $S=Q{K}^T\in {\mathbb{R}}^{N\times N}$, and apply row-wise softmax to get $P=softmax(S)\in {\mathbb{R}}^{N\times N}$. $S$ can be thought of as a pairwise similarity matrix between the encoding at position $i$ and position $j$ that captures the degree of interaction. For example, in a sentence the economy has been in decline, the value of $S_{1,5}$ (assuming 0-index) measuring the interaction between economy and decline might be high.

Finally, we produce the output $O=PV\in {\mathbb{R}}^{N\times d}$, which is an activation output from the input sequence that has captured the interactions between tokens at different positions of the input. This simple mechanism has led to significant improvements in language modelling.

GPU Memory Hierarchy

The memory hierarchy is such that read/write speed is super fast on the SRAM but memory is highly limited. Hence, the N x N attention matrix is written/read repeatedly to/from HBM, resulting in IO being a bottleneck. The numbers are as such on an A100 GPU:

• SRAM: 19 TB/s (20 MB RAM)
• HBM: 1.5 TB/s (40 GB RAM)

Naive Attention Algorithm

The naive attention algorithm has many reads and writes to HBM. (ps: Not sure why we cannot persist the intermediate matrices on SRAM and complete the computations, but in any case the naive algorithm requires materializing the $N\times N$ matrices on SRAM which will quickly flood it. For example, a sequence length of 2,048 at float32 already takes up 33MB for the $S$ matrix).

1. Load $Q,K$ from HBM, compute $S=Q{K}^T$, write $S$ to HBM
2. Read $S$ from HBM, compute $P=softmax(S)$, write $P$ to HBM
3. Load $P,V$ from HBM, compute $O=PV$, write $O$ to HBM

Flash Attention

The main idea is quite simple: instead of computing the full attention matrix, we use block-wise tiling to compute parts of it at a time. This reduces the memory required for each block and allows the whole computation to be done on SRAM while minimizing the amount of IO read from HBM, leading to faster compute time and lower memory usage on SRAM. The difficulty is in devising a block-wise softmax algorithm that yields the exact same result as computing it all at once.

Consider the naive softmax algorithm on an arbitrary vector $x\in {\mathbb{R}}^B$.

$$\begin{array}{rl}m(x)& :=\underset{i}{\max }{x}_i\\ p(x)& :=[e^{x_i-m(x)}...e^{x_B-m(x)}]\\ l(x)& :=\sum _{i}p(x{)}_i\\ softmax(x)& :=\frac{p(x)}{l(x)}\end{array}$$

Note that the maximum value $m(x)$ is subtracted for numerical stability to avoid overflow (underflow is ok because $e^{-inf}=0$). $f(x)\in {\mathbb{R}}^B$ is the numerator and $l(x)\in \mathbb{R}$ is the sum of all elements in $p(x)$.

Now, the problem with the naive softmax algorithm in the context of attention is that we need an entire row of $S$ ($N$ elements) to perform the row-wise softmax computation. This will not be available if we are performing block-wise computation, since we are splitting $Q,K\in {\mathbb{R}}^{Nxd}$ row-wise into blocks of $Q_i,K_j\in {\mathbb{R}}^{Bxd}$. When we compute $S_{ij}:=Q_i\cdot K_j^T\in {\mathbb{R}}^{BxB}$, blocks of $S$ will be materialized in each pass, but not the entire row at a time.

Hence, we need a modified algorithm that allows us to compute chunks of the final output $O\in {\mathbb{R}}^{Nxd}$ at a time by iterating block-wise through $S$, such that the combination of the new chunk of $O$ at each step with the already written intermediate $O$ gives the correct result at the end. The key to realizing this algorithm is in decomposing the softmax step, as shown below.

Consider two vectors $x^a,x^b\in {\mathbb{R}}^B$. We can decompose the softmax of their concatenated vector $x=[x^a{x}^b]\in {\mathbb{R}}^{2B}$ as follows:

$$\begin{array}{rl}m(x)& =max(m(x^a),m(x^b))\\ p(x)& =[e^{m(x^a)-m(x)}\cdot p(x^a)e^{m(x^b)-m(x)}\cdot p(x^b)]\\ l(x)& =e^{m(x^a)-m(x)}\cdot l(x^a)+e^{m(x^b)-m(x)}\cdot l(x^b)\\ softmax(x)& =\frac{p(x)}{l(x)}\end{array}$$

The first line of the above simply notes that the maximum of $x$ is the maximum over each of the subvector maximums $x^a,x^b$. The second line notes that we previously multiplied each element of $p(x)$ by a factor, say $e^{-m(x^a)}\cdot e^{x_i}$ for those in $x_a$. To get the correct multiplier for the full vector $x$, we need to divide away the previous multiplier and apply the new multiplier, i.e. $e^{-m(x)}/e^{-m(x^a)}=e^{m(x^a)-m(x)}$. The third line notes that the new denominator is the sum over each of the subvector sums, after we apply the correct multiplier from line 2.

The decomposition is simple but powerful. It implies that so long as we keep track of intermediate statistics $m(x)$ and $l(x)$, we can compute the softmax of a long vector $x$ by splitting $x$ into subvectors and operate over each subvector at a time.

Now we are ready for Algorithm 1: Flash Attention of the paper.

$$\begin{array}{rlr}& \text{Require: }Q,K,V\in {\mathbb{R}}^{Nxd}& \\ & \text{01. }\text{Initialize }O=0^{Nxd},l=0^N,m={-inf}^N& \\ & \text{02. }\text{Divide }Q,K,V\text{ row-wise into B-sized blocks of }{Q}_i,K_j,V_j\in {\mathbb{R}}^{Bxd}& \\ & \text{03. }\text{Divide }O\text{ row-wise into B-sized blocks of }{O}_i\in {\mathbb{R}}^{Bxd}& \\ & \text{04. }\text{Divide }m,l\text{ into B-sized arrays of }{m}_i,l_i\in {\mathbb{R}}^B& \\ & \text{05. }\text{for j in 1,... :}& \\ & \text{06. }\text{Load Kj,Vj from HBM to SRAM}& \\ & \text{07. }\text{for i in 1,... :}& \\ & \text{08. }\text{Load Qi,Oi,li,mi from HBM to SRAM}& \\ & \text{09. }\text{Compute }{S}_{ij}=Q_i\cdot K_j^T\in {\mathbb{R}}^{BxB}& \\ & \text{10. }\text{Compute }{m}_{ij}=rowmax(S_{ij})\in {\mathbb{R}}^B& (1)\\ & \text{11. }\text{Compute }{P}_{ij}=e^{S_{ij}-m_{ij}}\in {\mathbb{R}}^{BxB}& (2)\\ & \text{12. }\text{Compute }{l}_{ij}=rowsum(S_{ij})\in {\mathbb{R}}^B& (3)\\ & \text{13. }\text{Compute }{m}_i^{new}=\text{elementwise-max}(m_i,m_{ij})\in {\mathbb{R}}^B& (5)\\ & \text{14. }\text{Compute }{l}_i^{new}=e^{m_i-m_i^{new}}\cdot l_i+e^{m_{ij}-m_i^{new}}\cdot l_{ij}\in {\mathbb{R}}^B& (7)\\ & \text{15. }\text{Write }{O}_i\leftarrow (l^{new}{)}^{-1}\times \left[l_i\times e^{m_i-m_i^{new}}\times O_i+e^{m_{ij}-m_i^{new}}\times P_{ij}{V}_j\right]\text{ to HBM}& \\ & \text{16. }\text{Write }{l}_i\leftarrow l_i^{new},m_i\leftarrow m_i^{new}\text{ to HBM}& \\ & \text{17. }\text{Return }O& \end{array}$$

Note that we use $0^{Nxd},0^N$ to denote a zero matrix of size $Nxd$ and a zero array of size $N$ respectively. For simplicity, we divide $Q,K,V$ into equal $B$-sized blocks but the paper allows different block sizes for $Q$ and $K,V$. The equation numbers on the right in parentheses show which equations the lines correspond to above. Equation line 15 is a bit confusing because it combines multiple steps together. The next few paras try to unpack this.

Firstly, note that we are using the $\times$ operator to denote an element-wise broadcasted multiplication. For a vector $l_i\in {\mathbb{R}}^B$ and a matrices $P_{ij}\in {\mathbb{R}}^{BxB},V_j\in {\mathbb{R}}^{Bxd}$, observe the associative property $(l_i\times P_{ij})\cdot V_j=l_i\times (P_{ij}\cdot V_j)$, since each element of $l_i$ only affects the corresponding row in the final matrix. This allows us to apply the scaling to either $O_i$ or $P_{ij}$ and the result will be the same.

Next, see that the term $e^{m_{ij}-m_i^{new}}\times P_{ij}{V}_j$ is simply the corrected numerator of the softmax dotted with $V_j$. Dividing this term by $l^{new}$ gives the output block for this particular $S_{ij},V_j$ pair.

Similarly, the other term $l_i\times e^{m_i-m_i^{new}}\times O_i$ is the existing output that has been accumulated from previous steps $\{S_{ij},V_j:j=1,...\}$. Due to the associative property, we can also directly apply the scaling correction to $O_i$. The $l_i/l^{new}\times e^{m_i-m_i^{new}}$ are scaling factors according to equations (6), (8) to correct the scaling of previous steps.

Finally, we should understand why there is a + in equation 15. I find it easier to visualize if we set $B=1$. If we trace the matrix multiplications, we will observe that $O_i\in {\mathbb{R}}^{1xd}$ is only affected by $Q_i\in {\mathbb{R}}^{1xd}$, i.e. it corresponds to only the query token in position $i$. Now, $O_i$ represents the weighted average over all $N$ positions of the $V$ matrix where the weights are determined by the softmax over the interaction between $Q_i$ (representing one token) and all $N$ positions on the $K$ matrix. This weighted average is why it is a $+$ symbol: we are accumulating the weighted sum over $V_j$ into $O_i$. The only complication is that we are applying the scaling corrections at each step.

Hopefully these explanations provide some intuition to the FlashAttention algorithm, which is quite a simple idea but makes a ton of difference practically. It should be easy to implement this algorithm in numpy if the reader wishes to understand it better.