Flash Attention forward pass

Assume that vectors are column vectors in notation

Computing safe softmax in online fashion

Online algorithm for the normalization factor

• $x_1,\dots ,x_n$

• Let $m_i$ be the max value at step $i$ and $s_i$ the running sum at step i

• At step 0

  • $m_0=x_0$, $s_0=e^{(}{x}_0-m_0)=1$

• At step 1

  • if $x_1>x_0$, $m_1=x_1$, $s_1=s^0\ast e^{(}{m}_0-m_1)+e^{(}{x}_1-m_1)={\sum }_{i=0}^1{e}^{(}{x}_i-m_1)$

• At step i,

  • By induction, we assume $s_{i-1}={\sum }_{j=0}^{i-1}{e}^{(}{x}_j-m_{i-1})$
  • then $m_i=max(x_i,m_{i-1})$ and $s_i=s^{i-1}\ast e(m_{i-1}-m_i)+e^{(}{x}_i-m_i)$

Computing $Q{K}^{T}V$ without materializing $Q{K}^T$

• Using Block Matrix Multiplication, the high-level idea is that you compute first a block of $Q{K}^T$ (ideally the shape of the block should be independent of the sequence length e.g. $(32,\dim )$), and then directly reuse the resulting $(32,32)$ block to compute with $V$ block $(32,\dim )$.

Online block softmax computation

• For the sake of simplicity, let’s say we have blocks of shape (c,dim) of $Q$ and shape (dim,c) for $K^T$

  • In this case, we’re selecting entire rows of Q as blocks (called $Q_i$), and entire columns of $K^T$ as blocks. (called $K_i^T$)

• We’ll call the score matrix $S=Q{K}^T$ of shape (T,T)

  • Assuming we do block-matrix MM, we get that $S_{ij}$ = $Q_i{K}_j^T$ of shape (c,c) (note that there’s no accumulation, because the blocks span entire rows/columns)

• Let’s call ${\text{softmax}}^{\ast }(x_i)=exp(x_i-x_{max})$, which is softmax without the normalization (still over a row vector)

• Let’s call $P={\text{softmax}}^{\ast }(S)$, where the softmax is applied over each block independently i.e. $P_{ij}=exp(S_{ij}-rowmax(S_{ij})$

  • 

• Then we do a block MM with $V$ with blocks $V_i$ of shape (c,dim), giving us block rows $O_j$ of shape (c,dim) for the output matrix $O$ i.e. $O_j={\sum }_{i=0}^{T/c}{P}_{ji}{V}_i$

  • 
  • 

• Now obviously, each $P_{ji}$ in $O_j={\sum }_{i=0}^{T/c}{P}_{ji}{V}_i$ uses a different local maximum. However, we can reuse our Computing safe softmax in online fashion idea, communicate the maximum of each row in a block, and just readjust the current maximum as we accumulate

  • We can compute the normalization factor in the same online fasion
  • Does this require the accumulation to be non-parallel i.e. one-by-one?
    • No, because $max$ is an associative operator, we can do a parallel-prefix sum

The pseudo-code with for-loops to make it simpler

For output block o_j = [0,0,...,0)$ # shape (c,dim)
	curr_max = eye(-infinity) # shape (c,1)
	normalization = [0,...0] # shape (c,1)
	for i in range(T/c):
		## compute the block
		S_ji = Q_iK_j^T
		## recomputing the max
		curr_max = max(curr_max, rowmax(S_ji))
		correction_factor = exp(rowmax(S_ji) - curr_max)
		P_ji = exp(S_ji - curr_max)
		## correcting the sums
		normalization = normalization*correction_factor + rowsum(P_ji)
		o_j = o_j*correction_factor + P_ji*V_i
	
	## normalize at the end
	o_j = o_j * normalization
 

Actual FA2 pseudocode (with more comments on loading between SRAM and HBM)

•