FlashAttention-1부터 4까지의 발전 과정을 PyTorch로 구현한 교육용 저장소

사용자들은 복잡한 CUDA 코드 없이 알고리즘의 핵심 논리를 이해할 수 있다는 점에 대해 매우 긍정적인 반응을 보였다.

• FlashAttention의 내부 작동 원리를 깊이 있게 이해하고 싶다면 공식 CUDA 코드 대신 이 PyTorch 구현체의 버전별 차이점을 먼저 분석하라.
• FP8 연산이나 파이프라인 최적화가 실제 알고리즘 단계에서 어떻게 구현되는지 확인하려면 FA3와 FA4의 구현부를 참고하라.

FlashAttention

— GPU의 메모리 계층 구조를 활용하여 어텐션 연산의 속도를 높이고 메모리 사용량을 줄이는 알고리즘이다. 중간 결과물을 HBM에 저장하지 않고 타일링 기법을 통해 SRAM 내에서 연산하여 메모리 대역폭 병목 현상을 해결한다.

Online Softmax

— 전체 데이터를 한 번에 보지 않고 데이터를 순차적으로 읽으면서 소프트맥스 값을 계산하는 기법이다. FlashAttention에서 메모리 효율적인 어텐션 계산을 가능하게 하는 핵심 수학적 토대가 된다.

Tiling

— 큰 행렬 연산을 작은 블록(타일) 단위로 나누어 처리하는 기법이다. GPU의 빠른 메모리인 SRAM에 데이터를 올려 연산함으로써 느린 메인 메모리(HBM) 접근을 최소화하여 성능을 최적화한다.

FP8

— 데이터를 8비트로 표현하는 저정밀도 부동소수점 형식이다. FlashAttention-3부터 본격적으로 도입되었으며, 연산 속도를 높이고 메모리 사용량을 절반으로 줄여 대규모 모델 학습 및 추론 효율을 극대화한다.