FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

논문 해설

문제 및 동기

트랜스포머는 시퀀스 길이 $N$에 대해 2차(quadratic) 시간 및 메모리 복잡도를 가집니다. 근사 어텐션 방법들은 이론적인 FLOP 수를 줄이지만, 최신 GPU 연산이 산술 연산보다는 주로 메모리 전송(IO)에 의해 병목이 발생하기 때문에 실제 소요 시간(wall-clock) 단축을 달성하지 못하는 경우가 많습니다.

표준 어텐션에서는 $\mathbf{S}=\mathbf{Q}{\mathbf{K}}^{\top }$, $\mathbf{P}=\mathrm{softmax}(\mathbf{S})$, $\mathbf{O}=\mathbf{PV}$를 계산할 때 크기 $N\times N$의 중간 행렬을 느린 고대역폭 메모리(HBM)에 쓰고 이를 여러 번 다시 읽어옵니다. 이는 $\Theta (Nd+N^2)$의 HBM 접근과 $O(N^2)$의 메모리 사용량을 발생시켜 실질적인 컨텍스트 윈도우를 제한합니다.

핵심 아이디어

• IO-Awareness: GPU 메모리 계층 구조를 명시적으로 고려합니다: 빠르고 작은 온칩 SRAM(예: A100에서 SM당 192KB, 대역폭 약 19TB/s) 대 크고 느린 HBM(대역폭 1.5–2.0TB/s의 40–80GB).
• 타일링(블록 소프트맥스 계산): 입력 $\mathbf{Q},\mathbf{K},\mathbf{V}$를 SRAM 내에 완전히 들어가는 블록으로 나눕니다. 연속적인(running) 소프트맥스 행 최댓값과 정규화 통계량을 추적함으로써 HBM에 전체 $N\times N$ 행렬을 구체화하지 않고도 블록 전반에 걸쳐 소프트맥스 축약(reduction)을 점진적으로 계산합니다.
• 역전파에서의 재연산(Recomputation): 역전파를 위해 중간 $N\times N$ 어텐션 행렬을 저장하지 않습니다. 대신 출력 $\mathbf{O}$와 소프트맥스 통계량 $(\ell ,m)$만 저장하고, 역전파 과정에서 SRAM 내에서 블록 단위로 $\mathbf{S}$와 $\mathbf{P}$를 온칩 재연산합니다.
• 커널 융합(Kernel Fusion): 행렬 곱셈, 마스킹, 소프트맥스, 드롭아웃을 단일 CUDA 커널로 융합하여 불필요한 HBM 왕복을 제거합니다.

방법 (핵심 수식 포함)

$\mathbf{Q},\mathbf{K},\mathbf{V}\in {\mathbb{R}}^{N\times d}$가 주어졌을 때, FlashAttention은 블록 크기를 $B_c=\lceil M/(4d)\rceil$ 및 $B_r=\min (\lceil M/(4d)\rceil ,d)$로 설정합니다. 여기서 $M$은 SRAM 용량입니다.

점진적 소프트맥스 추적

하위 벡터 $x^{(1)}$과 $x^{(2)}$로 분해된 벡터에 대해, running 통계량은 다음과 같이 업데이트됩니다: $$m(x)=\max (m(x^{(1)}),m(x^{(2)})),$$ $$\ell (x)=e^{m(x^{(1)})-m(x)}\ell (x^{(1)})+e^{m(x^{(2)})-m(x)}\ell (x^{(2)}).$$

키/값의 블록 $j$와 쿼리의 블록 $i$에 대해:

1. SRAM 상에서 정규화되지 않은 점수를 계산: ${\mathbf{S}}_{ij}={\mathbf{Q}}_i{\mathbf{K}}_j^{\top }\in {\mathbb{R}}^{B_r\times B_c}$.
2. 블록의 로컬 통계량 계산: ${\tilde{m}}_{ij}=\mathrm{rowmax}({\mathbf{S}}_{ij})$, ${\tilde{\mathbf{P}}}_{ij}=\exp ({\mathbf{S}}_{ij}-{\tilde{m}}_{ij})$, 그리고 ${\tilde{\ell }}_{ij}=\mathrm{rowsum}({\tilde{\mathbf{P}}}_{ij})$.
3. 누적 전역 정규화 값 업데이트: $$m_i^{\mathrm{new}}=\max (m_i,{\tilde{m}}_{ij}),{\ell }_i^{\mathrm{new}}=e^{m_i-m_i^{\mathrm{new}}}{\ell }_i+e^{{\tilde{m}}_{ij}-m_i^{\mathrm{new}}}{\tilde{\ell }}_{ij}.$$
4. 누적 출력 블록 ${\mathbf{O}}_i\in {\mathbb{R}}^{B_r\times d}$ 업데이트: $${\mathbf{O}}_i\leftarrow \mathrm{diag}({\ell }_i^{\mathrm{new}}{)}^{-1}\left(\mathrm{diag}({\ell }_i)e^{m_i-m_i^{\mathrm{new}}}{\mathbf{O}}_i+e^{{\tilde{m}}_{ij}-m_i^{\mathrm{new}}}{\tilde{\mathbf{P}}}_{ij}{\mathbf{V}}_j\right).$$

복잡도 분석

• IO 복잡도: 표준 어텐션은 $\Theta (Nd+N^2)$의 HBM 접근을 필요로 합니다. FlashAttention은 $\Theta (N^2{d}^2{M}^{-1})$의 HBM 접근을 필요로 합니다. 저자들은 모든 $M\in [d,Nd]$에 걸쳐 $o(N^2{d}^2{M}^{-1})$ 접근을 점근적으로 달성할 수 있는 정확한 어텐션 알고리즘이 없음을 증명합니다.
• FLOPs 및 메모리: FlashAttention은 $O(N^{2}d)$ FLOPs를 수행하며(재연산으로 인해 역전파에서 FLOP 수가 약간 증가함), 입력 및 출력 외에 $O(N)$의 보조 메모리만 필요로 합니다.

확장: Block-Sparse FlashAttention

사전 정의된 희소성 마스크 $\mathbf{M}\in \{0,1{\}}^{N/B_r\times N/B_c}$에 따라 0이 아닌 블록으로 평가를 제한하면 $\Theta (Nd+N^2{d}^2{M}^{-1}s)$의 IO 복잡도를 달성합니다. 여기서 $s$는 0이 아닌 블록의 비율입니다.

결과 (가장 중요한 수치 포함)

• 학습 속도 향상:
  • BERT-large (시퀀스 길이 512): MLPerf 1.1 학습 속도 기록을 15% 능가($8\times$ A100에서 20.0분 대비 17.4분).
  • GPT-2 (시퀀스 길이 1K): HuggingFace 대비 최대 3.5배, Megatron-LM 대비 최대 1.8배 속도 향상 달성.
  • Long-Range Arena (LRA): 표준 어텐션 대비 2.4배 속도 향상, Block-sparse FlashAttention은 2.8배 속도 향상 달성.
• 더 긴 컨텍스트 지원 역량:
  • 4K 컨텍스트의 GPT-2: 1K 컨텍스트의 Megatron-LM보다 빠르게 학습하면서 검증 perplexity를 0.7 포인트 개선(18.2 대비 17.5).
  • 장문 문서 분류: 컨텍스트를 512에서 16K로 확장하면 MIMIC-III에서 micro-$F_1$ 4.3 증가(52.8에서 57.1), ECtHR에서 최대 8.5 증가(8K에서 72.2에서 80.7)를 달성.
  • Path-X 및 Path-256: FlashAttention은 61.4%의 정확도로 Path-X(시퀀스 길이 16K)를 해결한 최초의 트랜스포머입니다. Block-sparse FlashAttention은 Path-256(시퀀스 길이 64K)에서 63.1%의 정확도를 달성합니다.

한계점 및 열린 질문

• CUDA 엔지니어링 오버헤드: IO-aware 커널을 구현하려면 표준 PyTorch나 컴파일러 추상화가 아닌 커스텀 로우레벨 CUDA 프로그래밍이 필요하므로, 다양한 하드웨어 아키텍처 간의 이식성이 떨어집니다.
• 다중 GPU 어텐션: 최적화된 단일 노드 IO 스케줄이 상호 연결(NVLink, InfiniBand)을 통한 분산 다중 GPU 어텐션으로 자동으로 확장되지는 않습니다.
• 정확한 어텐션의 하한선: 일반적인 $M$에 걸쳐 최적이지만, $M$에 관한 정밀한 매개변수화된 복잡도 하한선을 증명하는 것은 여전히 미해결 과제로 남아 있습니다.