post : fused kernel

Author: woongjoonchoi

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

@@ -33,7 +33,9 @@ context를 전체 input이 아닌 부분으로 한정하는 Window Attention 이
 하지만, 이러한 efficient attention은 computation Complexity의 장점을 얻는 대신에 정확도에서 단점을 얻게 됩니다.
 따라서, Memory Efficient한 Attention이 주목을 받고 있습니다. Gpu Programming을 활용하여 정확도를 유지한채 연산속도와 메모리 사용량을 상당부분 향상시키기 때문입니다.
 
-## Online-Softmax
+## Softmax
+
+### Safe-Softmax
 
 $$ y_i = \frac{e^{x_i}}{\sum_{j=1}^V e^{x_j}} $$
 기존의 softmax는 exponential 함수를 사용하기에 overflow나 underflow 문제가 발생할 수 있다.
@@ -46,12 +48,16 @@ $$ y_i = \frac{e^{x_i - \max_{k=1}^V x_k}}{\sum_{j=1}^V e^{x_j - \max_{k=1}^V x_
 
 
 
+### Online-Softmax 
 하지만,softmax는 전체 row를 봐야한다.
 전체 row를 메모리에 매번 적재를 하는 것은 부담이 됩니다.
 따라서, input의 일부만을 확인하고 매번 분모와 분자를 update한다.
 
 ![online-softmax](\assets\images\DeepLearning\KernelFusion\online-softmax.png)
 
+
+#### Proof 
+
 **Base case**: \(V = 1\)
 
 1. $$m_1 \leftarrow x_1$$  
@@ -75,6 +81,7 @@ $$d_S \leftarrow d_{S-1} \, e^{\,m_{S-1} - m_S} + e^{\,x_S - m_S} $$
 $$= \left(\sum_{j=1}^{S-1} e^{\,x_j - m_{S-1}}\right) e^{\,m_{S-1} - m_S} + e^{\,x_S - m_S} $$  
 $$ = \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\}$$  
@@ -89,6 +96,28 @@ $$ = \sum_{j=1}^{S} e^{\,x_j - m_S} $$
   32-bit floating point의 range보다 적은 값까지 표현이 가능하다 하는 이유는 안전수치를 보수적으로 잡은거 같습니다. 
   만약에, 더 많은 벡터를 처리한다면 64-bit floating point 를 사용해야 할 것입니다.  
 
+#### Block-Update Proof
+위에서 증명한 것들은 Online-Softmax를 1개의 data를 차례대로 계산하여 update하는 과정을 증명한 것입니다.  
+실제로, data를 여러개 받아서 $$d_j $$ , $$m_j $$ 등을 계산하기 때문에 이를 기존에 구한 $$d_i $$ , $$m_i $$ 에 반영할 수 있어야 합니다.  
+
+$$
+\begin{bmatrix}
+m_i \\
+d_i
+\end{bmatrix}
+\;\oplus\;
+\begin{bmatrix}
+m_j \\
+d_j
+\end{bmatrix}
+=
+\begin{bmatrix}
+\max(m_i, m_j) \\
+d_i \, e^{\,m_i - \max(m_i, m_j)} \;+\; d_j \, e^{\,m_j - \max(m_i, m_j)}
+\end{bmatrix}
+$$
+
+
 ## Online-Self Attention
 attention에는 softmax 때문에 input 전체를 봐야한다.