FlashAttention은 데이터 이동을 줄여 정확한 어텐션을 더 빠르게 만들었다

긴 시퀀스 트랜스포머는 표준 셀프 어텐션이 큰 중간 행렬을 물질화하고 이를 GPU 고대역폭 메모리와 연산 장치 사이에서 반복적으로 이동시켰기 때문에 느리고 메모리를 많이 소모했다.

매력적인 해결책은 근사 어텐션이었다. 모델의 수학을 바꿔 복잡도를 줄이는 것이다. FlashAttention은 반대 경로를 택했다. 정확한 어텐션을 유지하되, 블록을 SRAM에서 처리하고 실행 중인 소프트맥스 통계를 재계산하여 전체 어텐션 행렬을 저장할 필요가 없게 했다.

논문은 MLPerf 1.1 훈련 기록 대비 BERT-large에서 15%의 종단 간 실제 시간 단축, 시퀀스 길이 1K의 GPT-2에서 3배 단축, Long Range Arena 작업에서 2.4배 단축을 보고했다.

수학이 아니라 메모리 트래픽을 최적화하라

핵심 제약은 부동소수점 연산만이 아니라 HBM과 SRAM 사이의 IO였다.

타일링을 통해 각 블록은 데이터가 이미 칩에 있는 동안 유용한 작업을 수행할 수 있었다.

알고리즘을 정확하게 유지함으로써 많은 근사 방법을 매력적이지 않게 만든 품질 저하를 피했다.

동일한 기본 요소가 트랜스포머 추상화를 변경하지 않고도 더 긴 문맥과 새로운 시퀀스 작업을 가능하게 했다.

하드웨어 이동이 병목일 때, 더 나은 알고리즘은 동일한 답을 계산하면서 메모리 계층을 더 적게 오가는 알고리즘일 수 있다.

FlashAttention은 효율적인 트랜스포머 훈련과 추론을 위한 표준 구성 요소가 되었다. 그 더 큰 교훈은 어텐션을 넘어 퍼졌다. 현대 가속기 알고리즘은 종종 새로운 모델 수학을 발명하기 전에 메모리 계층을 중심으로 설계함으로써 승리한다.