피드로 돌아가기
Dev.toAI/ML
원문 읽기
HBM 병목 해결을 위한 Online Softmax 기반 Flash Attention 설계 분석
I Built Flash Attention From Scratch — Here's What Nobody Tells You About It
AI 요약
Context
표준 Attention 구조에서 N×N Score Matrix를 GPU HBM에 전체 생성함에 따른 메모리 대역폭 병목 발생. 연산량(FLOPs)보다 메모리 I/O 비용이 전체 성능을 결정하는 Memory Wall 문제 직면.
Technical Solution
- HBM 내 대규모 행렬 생성 생략을 통한 Memory Wall 문제 근본적 해결
- 전체 시퀀스를 블록 단위로 나누어 처리하는 Tiling 기법 도입으로 Fast on-chip SRAM 활용 극대화
- 부분적인 블록 처리 중에도 정규화를 유지하는 Online Softmax 알고리즘 적용
- 새로운 타일의 Max 값이 갱신될 때마다 이전 부분 출력값을 재조정하는 Rescaling 로직 구현
- 데이터 이동을 최소화하는 IO-aware 디자인을 통한 연산 효율성 확보
실천 포인트
1. 메모리 집약적 연산 설계 시 Compute-bound인지 Memory-bound인지 우선 판별
2. 대규모 데이터 처리 시 전체 구체화(Materialization) 대신 Tiling을 통한 스트리밍 구조 검토
3. 전역 정보가 필요한 연산(Softmax 등)을 부분 합산 및 보정 가능한 Incremental 로직으로 변환 가능성 확인