Softmax가 만든 병목, FlashAttention는 이를 어떻게 풀었나

AI 요약

LLM 추론 비용은 학습보다 더 큰 문제이고, 주요 LLM API 가격 하락 뒤에는 추론 효율을 극한까지 끌어올리려는 최적화 노력이 있으며 그 핵심 축 중 하나가 FlashAttention으로, Tri Dao 연구팀이 지난 3월 Blackwell GPU에 맞춘 FlashAttention-4를 공개했습니다. Transformer의 Attention에서 Softmax의 exp(x)는 수치적으로 오버플로우 문제를 일으켜(FP16에서는 x가 약 11만을 넘으면 Overflow) 계산 안정화를 위해 최댓값을 찾아 빼는 과정과 정규화로 데이터의 반복 읽기·쓰기가 필요해 병목을 만듭니다. FlashAttention은 실제 병목이 연산 자체가 아니라 HBM↔SRAM 간의 데이터 이동임을 지적하고, 블록 단위 처리로 기존 구현에서 세 번 발생하던 HBM↔SRAM 왕복을 줄여 메모리 이동을 최소화한다고 설명합니다.