post : Efficient Attention

FlashAttention forward

Author: woongjoonchoi

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

@@ -227,18 +227,27 @@ FlashAttention은 Tri Dao라는 사람이 제안을 하였습니다.FlashAttenti
 
 FlashAttention의 forward는 이와 같습니다. algorithm의 line별로 설명을 하도록 하겠습니다.
 
-1. 
+1.  $$B_r$$는 $$Query$$의 block size를 의미하고 , $$B_c$$ 는 $$Key,Value$$의 block size를 의미합니다.
 
-2. 위의 online-self attention 에서 보았듯이 $$ v^* \in \mathbb{R}^d $$ ,$$ s^* \in \mathbb{R} $$ 를 0으로 초기화 했는데, FlashAttention에서는 이를 $$ O , \ell $$ 로 각각 표기법을 바꾸었습니다. Batch dimension으로 확장했기에 $$ O = \mathbf{0}_{N\times d} \in \mathbb{R}^{N \times d} $$ ,$$ \ell = \mathbf{0}_N \in \mathbb{R}^N $$ ,$$m = (-\infty)_N \in \mathbb{R}^N $$ 로 초기화합니다.
+2. 위의 online-self attention 에서는 $$ v^* \in \mathbb{R}^d $$ ,$$ s^* \in \mathbb{R} $$ 를 0으로 초기화 했는데, FlashAttention에서는 이를 $$ O , \ell $$ 로 각각 표기법을 바꾸었습니다. Batch dimension으로 확장했기에 $$ O = \mathbf{0}_{N\times d} \in \mathbb{R}^{N \times d} $$ ,$$ \ell = \mathbf{0}_N \in \mathbb{R}^N $$ ,$$m = (-\infty)_N \in \mathbb{R}^N $$ 로 초기화합니다.
 
+3. $$Query ,key , value$$ 를 여러 block으로 쪼갭니다.
 
+4. $$O$$는 $$Query$$ 의 weighted sum이므로 $$Query$$ 와 같은 Block size인 $$B_r$$ 으로 쪼개고, $$ \ell , m $$는 softmax를 계산할 때 사용되는 중간계산결과 이므로 $$Key,Value$$와 같은 Block size인 $$B_c$$ 를 사용합니다.
 
+5. 모든 $$Key,Value$$ Block에 대해서 반복합니다.
 
-#### Backward
+6. $Key,Value$의 block을 먼저 Cache에 load합니다. 
 
+7. 모든 $$Query $$ Block에 대해서 반복합니다.
 
-## Conclusion
+8. $$ Query,O, \ell , m $$의 block을 Cache에 load합니다.
+
+### Backward
 
 
+## Conclusion
+
 
+$$ \mathbf{S}_{ij} = \mathbf{Q}_i \mathbf{K}_j^{T} \in \mathbb{R}^{B_r \times B_c} $$
 ## References