Flash Attention: CPU-GPU의 BottleNeck

flash attention에 대해 알아보자.

original transformer

GPU 내부의 이야기

Attention의 과정

Attention은 Q, K, V 세 가지 행렬로 계산한다.

1. 먼저 Query와 Key를 행렬곱 한다.
   이 결과가 Attention Score인데,
2. 여기에 softmax를 씌워서 각 위치별 가중치를 만든다.
3. 그리고, 이 가중치를 Value에 가중합(Weighted Sum)하여 최종 Attention 출력을 얻는다.

$\text{Attention}(Q,K,V) = \mathrm{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$

- Q (Query) : [N, d_k]
- K (Key) : [M, d_k]
- V (Value) : [M, d_v]

• N : query 개수
• M : key/value 개수
• $d_k$ : Embedding Dimension

(자세한 Attention 계산 과정은 생략하도록 하겠다. )

그럼, 크게 보았을 때 Attention 계산은
(QK Inner Product -> Softmax -> Weighted sum:V)

---

GPU 메모리 구조

cv2.absdiff에서 나온 아이디어

HBM <-> SRAM 사이의 bottleneck 처럼 **OpenCV는 Numpy (CPU 동작)과 PyTorch 라이브러리 (GPU 연산) 사이의 병목을 생각하며 작업을 하자

참고자료

https://youtu.be/4Tw_ytMYHLI?si=D447MDPMzpoWyMTs