AI 하드웨어·소프트웨어 스택
FlashAttention 등 고속 어텐션 구현의 주된 발상으로 올바른 것은 무엇인가?
- A계산을 타일로 나누어 칩 위의 메모리에서 진행하고, n×n의 큰 중간 표를 HBM에 쓰지 않음으로써 메모리 읽기/쓰기를 줄인다
- B어텐션 계산을 생략하고 근사한 답을 낸다
- C어텐션을 CPU에서 계산한다
정답과 해설 보기
정답: A. 계산을 타일로 나누어 칩 위의 메모리에서 진행하고, n×n의 큰 중간 표를 HBM에 쓰지 않음으로써 메모리 읽기/쓰기를 줄인다
단순한 구현에서는 n×n 표를 HBM에 쓰고, 소프트맥스를 위해 다시 읽고, 다시 Value와의 곱을 위해 또 읽으므로 메모리 읽기/쓰기가 성능을 결정해 버립니다. FlashAttention은 Query와 Key·Value를 블록으로 나누어 공유 메모리에 읽어 들이고, 소프트맥스를 조금씩 갱신하면서 계산을 한 번에 진행하여 표를 HBM에 쓰지 않습니다. 계산량은 같아도 메모리 읽기/쓰기가 크게 줄어 빨라집니다.
- A. 정답입니다. 결과는 일반 어텐션과 같습니다.
- B. 근사가 아니라 같은 결과를 정확하게 냅니다.
- C. GPU의 칩 위 메모리를 활용합니다.