Shivam PandeyAI Research EngineerIM
Writing / KERNELIZED

FlashAttention · Transformers · Attention · Compute

FlashAttention Kernel: Forward Pass (MATH)

A first-principles derivation of the online, numerically stable FlashAttention forward pass.

Shivam Pandey3 min read
On this page

  • A∈Rs2s1A \in \mathbb{R}^{s_2 s_1} where s2s_2 and s1s_1 are the index set, for example in a 5D tensor A∈RijklmA \in \mathbb{R}^{ijklm} a possible index set could be s2={i,j}s_2 = \{ i, j \} and s1={k,l,m}s_1 = \{k, l, m \}.

  • Bs2s1=flashs1(Qs2s1,Ks3s1,Vs3s1)B_{s_2 s_1} = flash_{s_1}(Q_{s_2 s_1}, K_{s_3 s_1}, V_{s_3 s_1}) is a flash attention operation over s1 index set of tensor set <Q,K,V><Q, K, V>, and the resulting index set is the index set of QQ s.t. B∈Rs2s1B \in \mathbb{R}^{s_2 s_1}

Attention Operation

For this exercise we will simplify the target index sets to match the most common setup, i.e. Qs2s1∈R[M×d]Q_{s_2 s_1} \in \mathbb{R}^{[M \times d]}, Ks3s1∈R[N×d]K_{s_3 s_1} \in \mathbb{R}^{[N \times d]}, Vs3s1∈R[N×d]V_{s_3 s_1} \in \mathbb{R}^{[N \times d]}

Note: In this doc, whenever we have to denote the exact dimensions instead of index set it will be denoted as .[...×...×… ].^{[... \times ... \times \dots]} where ×\times symbol separates across different index sets.

Thus our operation becomes:

Bs2s1=attentions1(Qs2s1,Ks3s1,Vs3s1)B[M×d]=attentiond(Q[M×d],K[N×d],V[N×d])\begin{aligned}B_{s_2 s_1} = attention_{s_1}(Q_{s_2 s_1}, K_{s_3 s_1}, V_{s_3 s_1}) \\ B^{[M \times d]} = attention_{d}(Q^{[M \times d]}, K^{[N \times d]}, V^{[N \times d]}) \end{aligned}
S[M×N]=Q@KT=Q∗(M×d,N×d,M×N)K∈R[M×N]SrmaxM=max⁡NS[M×N]∈RMSrm[m,n]=S[m,n]−Srmaxm∀[m,n]∈[M×N]P[M×N]=softmaxN(Srm[M×N])O[M×d]=P@V=P∗(M×N,N×d,M×d)V∈R[M×d]\begin{aligned}S^{[M \times N]} = Q @ K^T = Q *_{(M \times d, N \times d, M \times N)} K \in \mathbb{R}^{[M \times N]} \\ S_{rmax}^{M} = \max_{N} S^{[M \times N]} \in \mathbb{R}^{M} \\ S_{rm}^{[m, n]} = S^{[m, n]} - S_{rmax}^{m} \forall [m, n] \in [M \times N] \\ P^{[M \times N]} = softmax_{N}(S_{rm}^{[M \times N]}) \\ O^{[M \times d]} = P @ V = P *_{(M \times N, N \times d, M \times d)} V \in \mathbb{R}^{[M \times d]} \end{aligned}

Note: Notation abuse -> O  ⟺  BO \iff B

Simplifying Forward Pass

P[m,n]=exp⁡(S[m,n]−max⁡n(S[m,n]))∑nexp⁡(S[m,n]−max⁡n(S[m,n]))=exp⁡(S[m,n])∑nexp⁡(S[m,n])\begin{aligned}P^{[m, n]} = \frac{\exp(S^{[m, n]} - \max_{n}(S^{[m, n]}))}{\sum_n \exp(S^{[m, n]} - \max_{n}(S^{[m, n]}))} = \frac{\exp(S^{[m, n]})}{\sum_n \exp(S^{[m, n]})} \end{aligned}

Note: The independence over M in softmax the only aggregation is required over N dimension

Computation in chunk

Here we will first look at what is reuired to generate the output for a single query i.e. O[m,d]O^{[m, d]}

O[m,d]=Pm∗(N,N×d,d)V=∑nP[m,n]⋅V[n,d]P[m,n]=softmaxn(S[m,n])=exp⁡(S[m,n])∑nexp⁡(S[m,n])O[m,d]=∑nP[m,n]⋅V[n,d]=∑nexp⁡(S[m,n])∑nexp⁡(S[m,n])⋅V[n,d]=1∑nexp⁡(S[m,n])∑nexp⁡(S[m,n])⋅V[n,d]\begin{aligned}O^{[m, d]} = P_m *_{(N, N \times d, d)} V = \sum_n P^{[m, n]} \cdot V^{[n, d]} \\ P^{[m, n]} = softmax_{n}(S^{[m, n]}) = \frac{\exp(S^{[m, n]})}{\sum_n \exp(S^{[m, n]})} \\ O^{[m, d]} = \sum_n P^{[m, n]} \cdot V^{[n, d]} = \sum_n \frac{\exp(S^{[m, n]})}{\sum_n \exp(S^{[m, n]})} \cdot V^{[n, d]} \\ = \frac{1}{\sum_n \exp(S^{[m, n]})}\sum_n \exp(S^{[m, n]}) \cdot V^{[n, d]} \end{aligned}

We want to process O[m,d]=∑N…O^{[m, d]} = \sum_N \dots over nn sequentially to avoid whole sequence loading.

for the sequence just processed till n=jn = j we can write:

Oj[m,d]=1∑n∈[0…j]exp⁡(S[m,n])∑n∈[0…j]exp⁡(S[m,n])⋅V[n,d]=1ljuj\begin{aligned}O^{[m, d]}_j = \frac{1}{\sum_{n\in[0 \dots j]} \exp(S^{[m, n]})}\sum_{n\in[0 \dots j]} \exp(S^{[m, n]}) \cdot V^{[n, d]} = \frac{1}{l_j} u_j \end{aligned}

Let’s say we proceed by a single setp n=j+1n = j+1:

Oj+1[m,d]=1∑n∈[0…j,j+1]exp⁡(S[m,n])∑n∈[0…j,j+1]exp⁡(S[m,n])⋅V[n,d]=1lj+exp⁡(S[m,n=j+1])(uj+exp⁡(S[m,n=j+1])⋅V[n=j+1,d])=Oj[m,d]∗ljlj+1+exp⁡(S[m,n=j+1])⋅V[n=j+1,d]lj+1\begin{aligned}O^{[m, d]}_{j+1} = \frac{1}{\sum_{n\in[0 \dots j, j+1]} \exp(S^{[m, n]})}\sum_{n\in[0 \dots j, j+1]} \exp(S^{[m, n]}) \cdot V^{[n, d]}\\ = \frac{1}{l_j + \exp(S^{[m, n=j+1]})} (u_j + \exp(S^{[m, n=j+1]}) \cdot V^{[n=j+1, d]}) \\ = \frac{O^{[m, d]}_{j} * l_j}{l_{j+1}} + \frac{\exp(S^{[m, n=j+1]}) \cdot V^{[n=j+1, d]}}{l_{j+1}} \end{aligned}

Thus we can compute the output simply by iterating over the NN dimension for Om,dO^{m, d} the final expression

O0[m,d]=exp⁡(S[m,n=0])⋅V[n=0,d]l0Oj+1[m,d]=Oj[m,d]∗ljlj+1+exp⁡(S[m,n=j+1])⋅V[n=j+1,d]lj+1l0=exp(S[m,n=0])lj+1=lj+exp(S[m,n=j+1])\begin{aligned}O^{[m, d]}_{0} = \frac{\exp(S^{[m, n=0]}) \cdot V^{[n=0, d]}}{l_{0}}\\ O^{[m, d]}_{j+1} = \frac{O^{[m, d]}_{j} * l_j}{l_{j+1}} + \frac{\exp(S^{[m, n=j+1]}) \cdot V^{[n=j+1, d]}}{l_{j+1}} \\ l_0 = exp(S^{[m, n=0]}) \\ l_{j+1} = l_j + exp(S^{[m, n=j+1]}) \end{aligned}

WTF: exp⁡\exp can explode coz of high multiplication values

max operation is used for numerical stability of the softmax especially keeping exp from exploding.

Here next we will try to incorporated this stabilization technique in the above derived framework.

m0=S[m,n=0]mj+1=max⁡(mj,S[m,n=j+1])l0=exp⁡(S[m,n=0]−m0)lj+1=∑n∈[0…j,j+1]exp⁡(S[m,n]−mj+1)lj+1=exp(−mj)exp⁡(mj+1−mj)∑n∈[0…j]exp⁡(S[m,n])+exp⁡(S[m,n=j+1]−mj+1)lj+1=ljexp⁡(mj−mj+1)+exp⁡(S[m,n=j+1]−mj+1)O0[m,d]=exp⁡(S[m,n=0]−m0)⋅V[n=0,d]l0Oj+1[m,d]=Oj[m,d]∗ljlj+1+exp⁡(S[m,n=j+1]−mj+1)⋅V[n=j+1,d]lj+1\begin{aligned}m_0 = S^{[m, n=0]} \\ m_{j+1} = \max(m_j, S^{[m, n=j+1]}) \\ l_0 = \exp(S^{[m, n=0]} - m_0) \\ l_{j+1} = \sum_{n\in[0 \dots j, j+1]}\exp(S^{[m, n]} - m_{j+1}) \\ l_{j+1} = \frac{exp(-m_{j})}{\exp(m_{j+1} - m_{j})}\sum_{n\in[0 \dots j]}\exp(S^{[m, n]}) + \exp(S^{[m, n=j+1]} - m_{j+1}) \\ l_{j+1} = l_{j}\exp(m_{j} - m_{j+1}) + \exp(S^{[m, n=j+1]} - m_{j+1}) \\ O^{[m, d]}_{0} = \frac{\exp(S^{[m, n=0]} - m_{0}) \cdot V^{[n=0, d]}}{l_{0}} \\ O^{[m, d]}_{j+1} = \frac{O^{[m, d]}_{j} * l_j}{l_{j+1}} + \frac{\exp(S^{[m, n=j+1]} - m_{j+1}) \cdot V^{[n=j+1, d]}}{l_{j+1}} \\ \end{aligned}