Flash Attention

Flash Attention

표준 Attention의 메모리 병목을 해결하기 위해 IO-aware 알고리즘을 적용한 효율적 Attention 구현으로, 타일링(Tiling)을 통해 HBM↔SRAM 간 데이터 이동을 최소화하여 O(N) 메모리와 2~4배 속도 향상 달성

핵심 원리: 타일링(입력을 블록 단위로 분할하여 SRAM에서 연산), IO-aware(GPU 메모리 계층 고려/HBM 접근 최소화), Online Softmax(블록별 점진적 softmax 계산), Recomputation(역전파 시 어텐션 재계산으로 메모리 절약)

효과: 메모리 O(N²)→O(N)(시퀀스 길이에 선형), 속도 2~4배(wall-clock), 긴 시퀀스 처리 가능(4K→64K+), 학습·추론 모두 적용, 정확도 동일(수학적 등가)

비교: 표준 Attention(O(N²) 메모리/느림/간단) vs Flash Attention(O(N) 메모리/빠름/IO 최적화) vs Sparse Attention(O(N√N)/근사/패턴 기반)

연관: Transformer, Self-Attention, GPU 메모리 최적화, LLM 학습