post :Efficient Attention

FlashAttention

Author: woongjoonchoi

_posts/DeepLearning/Kernel Fusion/2025-03-07-fused.md

@@ -29,8 +29,8 @@ self-attention을 수학적으로 동일하지만 Memory Efficient하게 구현
 Attention은 기존의 RNN,GRU,LSTM 스타일의 encoder-decoder[link](https://arxiv.org/pdf/1406.1078) 의 한계점을 보완하기 위해 나왔습니다.
 Attention[[link]](https://arxiv.org/pdf/1409.0473)은 Dzmitry Bahdanau라는 분과 뉴옥대의 조경현 교수님께서 처음 제안하셨습니다.  
 
-조경현 교수님께서는 fixed size vector에 sequence vector를 압축하는것이 불가능하다고 결론을 내리셨습니다.  
-[[link]](https://www.boostcourse.org/ai331/lecture/540234?isDesc=false)
+조경현 교수님께서는 fixed size vector에 sequence vector를 압축하는것이 불가능하다고 결론을 내리셨습니다.[[link]](https://www.boostcourse.org/ai331/lecture/540234?isDesc=false)  
+
 ![rnn-attention](\assets\images\DeepLearning\KernelFusion\rnn-attention.png)  
 따라서, sequence vector를 fixed size로 압축해서  decoder의 input에 주지 않고 token length에  비례한 vector를 decoder에 input으로 주었습니다.
 input sequence를 tokenize한 embedding vector들로 정보를 만들어야 합니다.
@@ -113,7 +113,6 @@ $$ = \sum_{j=1}^{S} e^{\,x_j - m_S} $$
 
 #### SafeNess 
 
-**SafeNess** :  
 - $$m_j \in \Bigl[\min_{k=1}^V m_k,\; \max_{k=1}^V m_k\Bigr], \quad \forall j \in \{1, \dots, V\}$$  
   따라서,, $$m_j$$ 는 overflow나 underflow가 될 수 없습니다.
 
@@ -166,7 +165,7 @@ attention에는 softmax 때문에 input 전체를 봐야한다.
 하지만, softmax는 Online-softmax로 전체를 안봐도 계산이 가능하다
 
 self-attention에도 이를 적용한다.[[link]](https://arxiv.org/pdf/2112.05682) 
-$$ v^* \in \mathbb{R}^d $$ ,$$ s^* \in \mathbb{R}^d $$ 를 0으로 초기화하고 , $$m^*$$를 -inf로 초기화합니다.
+$$ v^* \in \mathbb{R}^d $$ ,$$ s^* \in \mathbb{R} $$ 를 0으로 초기화하고 , $$m^*$$를 -inf로 초기화합니다.
 query $$q$$, keys $$k_1, \dots, k_n$$ 와 values $$v_1, \dots, v_n$$ 가 주어질 때 , keys와 values들을 순서대로 사용합니다.$$k_i$$ , $$v_i$$ 가 주어지면 $$ s_i = \mathrm{dot}(q, k_i) $$ 를 계산합니다.그리고 나서,$$m^* = \max(m^*,s_i)$$, $$ v^* \leftarrow v^* e^{m^* - m_i} + v_i e^{s_i - m_i} $$ , $$ s^* \leftarrow s^* e^{m^* - m_i} + e^{s_i - m_i} $$ 를  update합니다.그리고,$$ \frac{v^*}{s^*} $$를 계산해줍니다.
 
 
@@ -187,20 +186,35 @@ $$m^* = \max(m^*,s_i)$$, ,$$ s^* \leftarrow s^* e^{m^* - m_i} + e^{s_i - m_i} $$
 
 
 1. $$ Assume \quad v^* = \sum_{j=1}^{i-1} v_je^{\,x_j - m^*} $$  
-$$ m_i = \max(m^*,s_i)$$ , $$ s_i = \mathrm{dot}(q, k_i) $$
+$$ m_i = \max(m^*,s_i)$$ , $$ s_i = \mathrm{dot}(q, k_i) $$  
 $$ v^* \leftarrow v^* e^{m^* - m_i} + v_i e^{s_i - m_i} $$  
 $$= \left(\sum_{j=1}^{i-1} v_j e^{\,x_j - m^*}\right) e^{\,m^* - m_i} + v_i e^{\,s_i - m_i} $$  
 $$ = \sum_{j=1}^{i} v_j  e^{\,s_j - m_i} $$  
+그리고, $$ \frac{v^*}{s^*} $$ 를 계산해줍니다.  
 $$m_i = \max(m^*,s_i)$$ 가 현재까지 계산한 $$ \max $$ 이기 때문에 online self-attention 수식이 정확함이 증명이 됩니다.
 
 
-### FlashAttention
+## FlashAttention
+
+
+
+Online-Self Attention을 실제로 구현체를 Open-source로 제공한 것이 FlashAttention이다.
+FlashAttention은 Tri Dao라는 사람이 제안을 하였습니다.FlashAttention의 알고리즘에서 idea가 새로 나온것은 없지만 Online-Self Attention을 실제로 여러 LLM에 적용할 수 있도록 구현체를 제공한 최초의 사례입니다. 
+
+### Forward
+![FlashAttention](\assets\images\DeepLearning\KernelFusion\FlashAttention_forward.png)  
+
+FlashAttention의 forward는 이와 같습니다. algorithm의 line벼로 설명을 하도록 하겠습니다.
+
+1. 
+
+2. 위의 online-self attention 에서 보았듯이 $$ v^* \in \mathbb{R}^d $$ ,$$ s^* \in \mathbb{R} $$ 를 0으로 초기화 했는데, FlashAttention에서는 이를 $$ O , \ell $$ 로 각각 표기법을 바꾸었습니다. Batch dimension으로 확장했기에 $$ O \in \mathbb{R}^{N \times d} $$ ,$$ \ell \in \mathbb{R}^N $$ ,$$m = (-\infty)_N \in \mathbb{R}^N $$ 로 초기화합니다.
+
 
 
-![FlashAttention](\assets\images\DeepLearning\KernelFusion\FlashAttention_forward.png)
 
+#### Backward
 
-Online-Self Attention을 실제로 알고리즘 관점에서 자세하게 분석한게 FlashAttention이다.
 
 ## Conclusion