FlashAttention: GPU 메모리 계층을 활용한 어텐션 최적화 분석

1. 들어가며
2. Standard Attention의 메모리 문제
3. GPU 메모리 계층 구조
4. IO Complexity 분석
5. Tiling 기법: SRAM에 맞는 블록 단위 연산
6. Online Softmax (Safe Softmax) 알고리즘
7. Backward Pass의 Recomputation 전략
8. FlashAttention-2 개선점
9. FlashAttention-3 최신 발전
10. 벤치마크: 속도/메모리 비교
11. PyTorch torch.nn.functional.scaled_dot_product_attention 연동
12. 정리 및 핵심 교훈
References

1. 들어가며

Transformer 아키텍처의 핵심인 Self-Attention은 시퀀스 내 모든 토큰 쌍 간의 관계를 계산한다. 이 연산은 강력한 표현력을 제공하지만, 시퀀스 길이 $N$ 에 대해 시간 및 메모리 복잡도가 $O(N^2)$ 으로 증가하는 근본적인 한계를 갖는다. GPT-4, LLaMA, Gemini 등 최신 LLM이 128K 이상의 긴 컨텍스트를 처리하려면, 이 $O(N^2)$ 병목을 실질적으로 해결해야 한다.

FlashAttention(Dao et al., 2022)은 이 문제를 근사(approximation) 없이 해결한다. 핵심 아이디어는 단순하면서도 깊다: attention 연산 자체의 계산량을 줄이는 것이 아니라, GPU 메모리 계층 간 데이터 이동(IO)을 최소화하는 것이다. 이 글에서는 FlashAttention의 원리를 GPU 하드웨어 관점에서 체계적으로 분석하고, FlashAttention-2와 FlashAttention-3까지의 발전을 살펴본다.

2. Standard Attention의 메모리 문제

2.1 Standard Attention 연산 흐름

Standard Self-Attention은 다음과 같이 계산된다. 입력 $Q, K, V \in \mathbb{R}^{N \times d}$ 에 대해:

S = QK^T \in \mathbb{R}^{N \times N}

P = \text{softmax}(S) \in \mathbb{R}^{N \times N}

O = PV \in \mathbb{R}^{N \times d}

여기서 $N$ 은 시퀀스 길이, $d$ 는 head dimension이다.

2.2 메모리 복잡도 분석

문제의 핵심은 중간 행렬 $S$ 와 $P$ 에 있다. 이 행렬들의 크기는 $N \times N$ 이며, 시퀀스 길이에 대해 이차(quadratic) 메모리를 요구한다. 구체적인 수치를 계산하면:

시퀀스 길이 ( $N$ )	Attention 행렬 크기	FP16 메모리
1,024	1M 원소	2 MB
4,096	16.7M 원소	33 MB
16,384	268M 원소	536 MB
65,536	4.3B 원소	8.6 GB
131,072	17.2B 원소	34.4 GB

이 수치는 단일 head, 단일 배치에 대한 것이다. Multi-head attention에서 head 수 $h$ 를 곱하고, 배치 크기 $B$ 를 곱하면 실제 메모리 사용량은 훨씬 커진다. 시퀀스 길이 65,536에서 이미 단일 head만으로 A100 80GB GPU의 HBM 상당 부분을 소비하게 된다.

2.3 HBM 병목 현상

Standard attention의 구현에서는 이 $N \times N$ 행렬을 GPU HBM(High Bandwidth Memory)에 materialization한다. 즉, $S = QK^T$ 를 계산하여 HBM에 쓰고, softmax를 위해 다시 읽고, 결과 $P$ 를 HBM에 쓰고, $O = PV$ 를 위해 다시 읽는다. 이 과정에서 HBM에 대한 읽기/쓰기 횟수는 $\Omega(Nd + N^2)$ 이 된다.

실제 GPU에서 이 연산이 느린 이유는 계산(compute)이 아니라 메모리 접근(memory access)이 병목이기 때문이다. A100 GPU의 계산 처리량은 312 TFLOPS(FP16)인 반면, HBM 대역폭은 약 2 TB/s에 불과하다. Attention 연산은 arithmetic intensity(연산량/메모리 접근량 비율)가 낮아 전형적인 memory-bound 연산이다.

3. GPU 메모리 계층 구조

FlashAttention을 이해하려면 GPU의 메모리 계층을 정확히 알아야 한다.

3.1 HBM (High Bandwidth Memory)

용량: A100 기준 40GB 또는 80GB
대역폭: 약 1.5-2.0 TB/s (A100 80GB SXM: 2,039 GB/s)
접근 지연: 약 200-600 사이클
역할: GPU의 메인 메모리. 모델 파라미터, 입력 텐서, 출력 텐서 등 모든 데이터가 저장됨

3.2 SRAM (On-chip Shared Memory)

용량: A100 기준 SM당 약 192KB, 전체 약 20MB (108개 SM)
대역폭: 약 19 TB/s
접근 지연: 약 20-30 사이클
역할: 각 Streaming Multiprocessor(SM) 내의 고속 온칩 메모리

3.3 핵심 비대칭성

SRAM과 HBM 사이에는 극적인 비대칭이 존재한다:

특성	SRAM	HBM
대역폭	~19 TB/s	~2 TB/s
용량	~20 MB	40-80 GB
접근 지연	20-30 사이클	200-600 사이클

SRAM은 HBM보다 약 10배 빠르지만, 용량은 약 4000배 작다. FlashAttention의 핵심 통찰은 이 비대칭성을 적극 활용하는 것이다: $N \times N$ 행렬 전체를 HBM에 materialization하는 대신, SRAM에 들어가는 작은 블록 단위로 연산을 수행하면 HBM 접근을 극적으로 줄일 수 있다.

4. IO Complexity 분석

4.1 Standard Attention의 IO Complexity

Standard attention은 다음과 같은 HBM 접근 패턴을 보인다:

$Q, K$ 를 HBM에서 읽어 $S = QK^T$ 계산 -> $S$ 를 HBM에 쓰기: $\Theta(Nd + N^2)$ IO
$S$ 를 HBM에서 읽어 $P = \text{softmax}(S)$ 계산 -> $P$ 를 HBM에 쓰기: $\Theta(N^2)$ IO
$P, V$ 를 HBM에서 읽어 $O = PV$ 계산 -> $O$ 를 HBM에 쓰기: $\Theta(Nd + N^2)$ IO

총 HBM 접근량: $\Theta(Nd + N^2)$

시퀀스 길이 $N$ 이 head dimension $d$ (보통 64 또는 128)보다 훨씬 크므로, $N^2$ 항이 지배적이 된다.

4.2 FlashAttention의 IO Complexity

FlashAttention은 tiling을 통해 HBM 접근량을 다음으로 줄인다:

O\left(\frac{N^2 d^2}{M}\right)

여기서 $M$ 은 SRAM 크기이다. 직관적으로, SRAM이 클수록 더 큰 블록을 한 번에 처리할 수 있어 HBM 접근이 줄어든다.

4.3 최적성 증명 (Lower Bound)

논문은 더 나아가 다음의 하한(lower bound)을 증명한다:

Theorem: $d \leq M \leq Nd$ 인 모든 SRAM 크기 $M$ 에 대해, exact attention을 계산하는 어떤 알고리즘도 $\Omega(N^2 d^2 / M)$ 의 HBM 접근이 필요하다.

이는 FlashAttention이 **IO complexity 관점에서 최적(optimal)**임을 의미한다. 상수 인자나 다항 로그 인자를 제외하면, 더 적은 HBM 접근으로 exact attention을 계산하는 것은 불가능하다.

4.4 수치 예시

A100의 SRAM 크기 $M \approx 192$ KB, head dimension $d = 64$ , 시퀀스 길이 $N = 4096$ 일 때:

Standard attention IO: $\Theta(Nd + N^2) \approx 4096 \times 64 + 4096^2 \approx 17M$ 원소
FlashAttention IO: $\Theta(N^2 d^2 / M) \approx 4096^2 \times 64^2 / (192 \times 512) \approx 7M$ 원소 (블록 크기에 따라 달라짐)

실제로는 $N^2$ 크기의 중간 행렬이 HBM에 전혀 기록되지 않으므로, 절약 효과는 더 크다. 특히 시퀀스 길이가 길어질수록 효과가 극대화된다.

5. Tiling 기법: SRAM에 맞는 블록 단위 연산

5.1 알고리즘 개요

FlashAttention의 핵심 알고리즘은 다음과 같다:

$Q$ 를 $T_r = \lceil N / B_r \rceil$ 개의 블록으로 나눈다: $Q_1, Q_2, \ldots, Q_{T_r}$ , 각 블록 크기 $B_r \times d$
$K, V$ 를 $T_c = \lceil N / B_c \rceil$ 개의 블록으로 나눈다: $K_1, \ldots, K_{T_c}$ 및 $V_1, \ldots, V_{T_c}$ , 각 블록 크기 $B_c \times d$
블록 크기 $B_r, B_c$ 는 SRAM 크기 $M$ 에 맞게 설정: $B_c = \lceil M / (4d) \rceil$ , $B_r = \min(\lceil M / (4d) \rceil, d)$

5.2 Forward Pass 의사코드

Algorithm: FlashAttention Forward Pass
---------------------------------------
Input: Q, K, V in HBM, SRAM size M
Output: O in HBM

1. 블록 크기 설정: B_c = ceil(M / 4d), B_r = min(ceil(M / 4d), d)
2. O = zeros(N, d), l = zeros(N), m = -inf * ones(N) 를 HBM에 초기화

3. for j = 1 to T_c:                        # Outer loop: K, V 블록
     K_j, V_j 를 HBM에서 SRAM으로 로드

     for i = 1 to T_r:                      # Inner loop: Q 블록
       Q_i, O_i, l_i, m_i 를 HBM에서 SRAM으로 로드

       # SRAM에서 블록 단위 연산 수행
       S_ij = Q_i @ K_j^T                   # (B_r x B_c)
       m_ij = rowmax(S_ij)
       P_ij = exp(S_ij - m_ij)
       l_ij = rowsum(P_ij)

       # 이전 블록의 통계치와 결합 (Online Softmax)
       m_new = max(m_i, m_ij)
       l_new = exp(m_i - m_new) * l_i + exp(m_ij - m_new) * l_ij

       # 출력 업데이트 (rescaling 포함)
       O_i = diag(exp(m_i - m_new))^(-1) * (diag(l_i) * O_i)
             + diag(exp(m_ij - m_new))^(-1) * P_ij @ V_j
       O_i = diag(l_new)^(-1) * O_i

       # 통계치 업데이트
       m_i = m_new, l_i = l_new

       O_i, l_i, m_i 를 HBM에 다시 쓰기
     end for
   end for

4. return O

5.3 왜 이것이 작동하는가

핵심은 $N \times N$ 크기의 attention 행렬 $S$ 와 $P$ 가 HBM에 전혀 materialization되지 않는다는 점이다. 각 $B_r \times B_c$ 블록 $S_{ij}$ 는 SRAM 내에서 계산되고, 즉시 softmax 통계치 업데이트와 출력 누적에 사용된 후 폐기된다.

이를 가능하게 하는 수학적 기법이 바로 Online Softmax이다.

6. Online Softmax (Safe Softmax) 알고리즘

6.1 Standard Softmax의 문제

Softmax는 **전역 연산(global operation)**이다. 행 벡터 $x = [x_1, \ldots, x_N]$ 에 대해:

\text{softmax}(x_i) = \frac{e^{x_i}}{\sum_{j=1}^{N} e^{x_j}}

이를 계산하려면 분모의 합계를 위해 전체 행을 한 번에 봐야 한다. 이것이 tiling을 어렵게 만드는 근본적인 장벽이다 -- 블록 $S_{i1}$ 만 보고는 softmax를 완성할 수 없다. 나머지 블록 $S_{i2}, S_{i3}, \ldots$ 의 값에 따라 분모가 달라지기 때문이다.

또한 수치 안정성을 위해 "safe softmax"를 사용한다:

\text{softmax}(x_i) = \frac{e^{x_i - m}}{\sum_{j=1}^{N} e^{x_j - m}}, \quad m = \max_j x_j

이 역시 전역 최댓값 $m$ 이 필요하므로 전체 행을 먼저 스캔해야 한다.

6.2 Online Softmax 트릭

Online Softmax(Milakov & Gimelshein, 2018)의 핵심 아이디어는 running statistics를 유지하면서 블록 단위로 점진적(incremental)으로 softmax를 계산하는 것이다.

두 개의 스칼라를 행마다 유지한다:

$m$ : 지금까지 본 원소들의 최댓값 (running max)
$\ell$ : 지금까지의 정규화 상수 (running sum of exponentials)

새로운 블록 $S_{ij}$ 가 들어오면:

새 블록의 행별 최댓값 계산: $\tilde{m} = \text{rowmax}(S_{ij})$
전역 최댓값 업데이트: $m_{\text{new}} = \max(m, \tilde{m})$
이전 정규화 상수를 rescale: $\ell_{\text{new}} = e^{m - m_{\text{new}}} \cdot \ell + e^{\tilde{m} - m_{\text{new}}} \cdot \tilde{\ell}$
이전 출력도 rescale: $O_{\text{new}} = \frac{e^{m - m_{\text{new}}} \cdot \ell \cdot O + e^{\tilde{m} - m_{\text{new}}} \cdot \tilde{P}V_j}{\ell_{\text{new}}}$

이 과정은 **수학적으로 정확(exact)**하다. 근사가 아니다. 블록을 어떤 순서로 처리하든, 최종 결과는 standard attention과 비트 단위로 동일하다(부동소수점 연산 순서에 따른 미세한 수치 차이 제외).

6.3 수학적 정당성

증명의 핵심은 softmax의 rescaling property이다:

\frac{e^{x_i - m'}}{\sum_j e^{x_j - m'}} = \frac{e^{x_i - m} \cdot e^{m - m'}}{\sum_j e^{x_j - m} \cdot e^{m - m'}} = \frac{e^{x_i - m}}{\sum_j e^{x_j - m}}

최댓값이 $m$ 에서 $m'$ 으로 업데이트되더라도, 분자와 분모에 동일한 factor $e^{m - m'}$ 가 곱해지므로 비율은 변하지 않는다. 이 성질 덕분에 이전 블록의 결과를 새로운 최댓값 기준으로 안전하게 rescale할 수 있다.

7. Backward Pass의 Recomputation 전략

7.1 Standard Backward Pass의 문제

Standard attention의 backward pass에서는 gradient 계산을 위해 forward pass에서 저장한 중간 행렬 $S$ 와 $P$ 가 필요하다. 이들의 크기가 $N \times N$ 이므로, forward에서 저장하고 backward에서 다시 읽는 것은 $O(N^2)$ 메모리를 요구한다.

7.2 FlashAttention의 Recomputation

FlashAttention은 gradient checkpointing의 변형을 사용한다. Forward pass에서 $S$ 와 $P$ 를 저장하지 않고, 대신 다음만 저장한다:

최종 출력 $O \in \mathbb{R}^{N \times d}$
Softmax 정규화 통계치 $m, \ell \in \mathbb{R}^{N}$ (행별 최댓값과 합계)

Backward pass에서는 이 통계치와 원본 $Q, K, V$ 를 사용하여 ** $S$ 와 $P$ 의 필요한 블록을 SRAM에서 다시 계산(recompute)**한다. 이 recomputation은 추가적인 FLOP을 요구하지만, HBM 접근을 크게 줄인다.

7.3 Recomputation의 역설적 효과

일반적으로 gradient checkpointing은 메모리를 절약하는 대신 속도를 희생한다. 그러나 FlashAttention의 recomputation은 오히려 속도까지 향상시킨다. 이유는 다음과 같다:

FLOP은 증가한다: forward에서 한 번 계산한 것을 backward에서 다시 계산하므로, 총 FLOP은 약간 증가한다.
HBM IO는 감소한다: $N \times N$ 크기의 $S$ , $P$ 를 HBM에 쓰고 읽는 비용이 사라진다.

현대 GPU에서는 HBM 접근이 계산보다 훨씬 느리므로, FLOP 증가량보다 IO 감소의 이득이 더 크다. 실험 결과, recomputation으로 인한 추가 런타임 오버헤드는 5% 미만이면서, 메모리 사용량은 $O(N^2)$ 에서 $O(N)$ 으로 감소한다.

7.4 메모리 절약 효과

시퀀스 길이	Standard Attention 메모리	FlashAttention 메모리	절약 비율
1K	~2 MB	~0.13 MB	~15x
2K	~8 MB	~0.26 MB	~30x
4K	~33 MB	~0.52 MB	~63x
8K	~131 MB	~1.04 MB	~126x

이 절약 효과 덕분에 동일한 GPU 메모리로 더 긴 시퀀스를 처리하거나, 더 큰 배치 크기를 사용할 수 있다.

8. FlashAttention-2 개선점

Dao(2023)는 FlashAttention-2에서 세 가지 핵심 개선을 도입했다.

8.1 Non-matmul FLOP 최소화

A100 GPU의 Tensor Core는 행렬 곱셈(matmul)에 대해 312 TFLOPS(FP16)를 제공하지만, non-matmul 연산(softmax의 exp, max, sum 등)은 19.5 TFLOPS(FP32)로 약 16배 느리다. FlashAttention-1에서는 non-matmul 연산 비중이 상당했다.

FlashAttention-2는 알고리즘을 재구성하여 이러한 non-matmul FLOP을 최소화한다. 구체적으로, rescaling 연산의 횟수를 줄이고, softmax 통계치 업데이트를 더 효율적으로 수행한다. 최종 rescaling을 루프 마지막에 한 번만 수행하도록 변경한 것이 핵심이다.

8.2 Parallelism 향상: Sequence Length 차원 병렬화

FlashAttention-1은 batch 차원과 head 차원에서만 병렬화했다. 배치 크기가 작거나 head 수가 적으면 GPU의 SM(Streaming Multiprocessor)을 충분히 활용하지 못했다.

FlashAttention-2는 시퀀스 길이 차원에서도 병렬화한다. 외부 루프를 Q 블록 기준으로 변경하여 ( $K, V$ 블록이 아닌 $Q$ 블록을 외부 루프로), 각 Q 블록을 독립적인 thread block에서 처리할 수 있게 했다. 이 변경으로 forward pass에서의 occupancy가 크게 향상된다.

8.3 Work Partitioning 최적화

Thread block 내에서 warp 간의 작업 분배도 개선되었다:

FlashAttention-1: K, V를 4개의 warp에 분할, 각 warp가 독립적으로 $QK^T$ 계산 후 결과를 동기화. 이 방식은 shared memory를 통한 통신과 동기화 오버헤드가 발생한다.
FlashAttention-2: Q를 4개의 warp에 분할, K와 V는 모든 warp가 공유. 각 warp는 Q의 다른 부분에 대해 독립적으로 출력을 계산하므로, warp 간 통신이 불필요하다.

8.4 성능 결과

이 세 가지 개선을 합치면:

FlashAttention-1 대비 약 2배 speedup
A100에서 FP16/BF16 기준 230 TFLOPS 달성 (이론적 최대의 약 73%)
Standard PyTorch attention 대비 최대 9배 speedup
GEMM(행렬 곱셈) 연산의 효율에 근접

9. FlashAttention-3 최신 발전

FlashAttention-3(Shah et al., 2024)는 NVIDIA Hopper 아키텍처(H100)의 새로운 하드웨어 기능을 활용하여 한 단계 더 발전했다.

9.1 Hopper GPU의 새로운 기능

H100 GPU는 A100에 비해 다음의 핵심 기능을 제공한다:

WGMMA (Warpgroup Matrix Multiply-Accumulate): A100의 mma.sync보다 훨씬 높은 처리량을 가진 새로운 Tensor Core 명령어
TMA (Tensor Memory Accelerator): Global memory와 Shared memory 간 데이터 전송을 전담하는 하드웨어 유닛. 인덱스 계산과 경계 검사를 하드웨어에서 처리

9.2 세 가지 핵심 기법

1. Warp Specialization을 통한 비동기 실행

연산(WGMMA)과 데이터 이동(TMA)을 서로 다른 warp group에 할당하여 파이프라인 방식으로 중첩(overlap) 실행한다. 한 warp group이 현재 블록을 계산하는 동안, 다른 warp group이 다음 블록의 데이터를 prefetch한다.

2. Matmul과 Softmax의 Interleaving

기존에는 matmul이 끝난 후 softmax를 수행하고, 다시 matmul을 수행하는 순차적 방식이었다. FlashAttention-3는 이를 인터리빙하여, matmul과 softmax가 서로 다른 하드웨어 유닛에서 동시에 실행되도록 한다. Tensor Core가 다음 블록의 $QK^T$ 를 계산하는 동안, CUDA Core가 현재 블록의 softmax를 처리한다.

3. FP8 Low-precision 지원

H100의 FP8 Tensor Core를 활용하여 처리량을 2배 높인다. 단순히 FP8로 양자화하면 정확도가 떨어지지만, FlashAttention-3는 두 가지 기법으로 이를 해결한다:

Block quantization: 블록 단위로 별도의 스케일 팩터를 유지하여 동적 범위를 보존
Incoherent processing: 랜덤 직교 행렬을 곱하여 outlier를 분산시킨 후 양자화. 이를 통해 FP8 baseline 대비 2.6배 낮은 수치 오차 달성

9.3 성능 결과

H100에서의 FlashAttention-3 성능:

설정	TFLOPS	GPU 활용률
FP16 FlashAttention-2	~400	~50%
FP16 FlashAttention-3	~740	~75%
FP8 FlashAttention-3	~1,200	~75%

FP16에서 FlashAttention-2 대비 1.5-2.0배 speedup, FP8에서는 1.2 PFLOPS에 근접하는 성능을 달성했다.

10. 벤치마크: 속도/메모리 비교

10.1 Attention Forward Pass 속도 (A100 80GB, FP16)

FlashAttention 논문과 후속 벤치마크에서 보고된 주요 수치는 다음과 같다:

시퀀스 길이	Standard Attention	FlashAttention	FlashAttention-2	Speedup (FA2 vs Std)
512	12.2 ms	3.5 ms	1.9 ms	6.4x
1K	45.8 ms	7.8 ms	4.1 ms	11.2x
2K	178 ms	18.9 ms	9.8 ms	18.2x
4K	710 ms	52.3 ms	27.1 ms	26.2x
8K	OOM	145 ms	75 ms	-
16K	OOM	520 ms	270 ms	-

시퀀스 길이가 길어질수록 speedup이 더 극적으로 증가한다. 8K 이상에서는 standard attention이 OOM(Out of Memory)으로 실행 자체가 불가능하지만, FlashAttention은 문제없이 처리한다.

10.2 End-to-End 학습 성능

모델	Standard	FlashAttention	Speedup
BERT-large (seq 512)	100% (MLPerf 기준)	115%	1.15x
GPT-2 (seq 1K)	100%	300%	3.0x
Long-range Arena (seq 1K-4K)	100%	240%	2.4x

10.3 메모리 사용량 비교

FlashAttention의 attention 연산 메모리는 시퀀스 길이에 대해 **선형(linear)**으로, standard attention의 이차(quadratic) 대비 극적인 개선이다:

시퀀스 길이 2K: 약 10배 메모리 절약
시퀀스 길이 4K: 약 20배 메모리 절약
시퀀스 길이 64K: standard attention은 A100 80GB에서도 OOM, FlashAttention은 정상 동작

11. PyTorch `torch.nn.functional.scaled_dot_product_attention` 연동

11.1 네이티브 통합

PyTorch 2.0부터 FlashAttention이 torch.nn.functional.scaled_dot_product_attention (SDPA)에 네이티브로 통합되어 있다. PyTorch 2.2부터는 FlashAttention-2가 기본 backend로 사용된다.

import torch
import torch.nn.functional as F

# 기본 사용법 - 자동으로 FlashAttention backend 선택
query = torch.randn(batch_size, num_heads, seq_len, head_dim,
                    device='cuda', dtype=torch.float16)
key = torch.randn(batch_size, num_heads, seq_len, head_dim,
                  device='cuda', dtype=torch.float16)
value = torch.randn(batch_size, num_heads, seq_len, head_dim,
                    device='cuda', dtype=torch.float16)

# PyTorch가 자동으로 최적의 backend를 선택한다
output = F.scaled_dot_product_attention(query, key, value)

11.2 Backend 명시적 선택

특정 backend를 강제로 사용하거나 제외할 수 있다:

from torch.nn.attention import sdpa_kernel, SDPBackend

# FlashAttention backend만 사용
with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
    output = F.scaled_dot_product_attention(query, key, value)

# Memory-efficient attention backend만 사용
with sdpa_kernel(SDPBackend.EFFICIENT_ATTENTION):
    output = F.scaled_dot_product_attention(query, key, value)

# Math (naive) backend 사용 - 디버깅용
with sdpa_kernel(SDPBackend.MATH):
    output = F.scaled_dot_product_attention(query, key, value)

# CuDNN backend 사용 (PyTorch 2.2+)
with sdpa_kernel(SDPBackend.CUDNN_ATTENTION):
    output = F.scaled_dot_product_attention(query, key, value)

11.3 Causal Mask와 함께 사용

LLM의 autoregressive generation에 필수적인 causal mask도 지원된다:

# is_causal=True로 causal mask 적용
# FlashAttention은 causal mask를 fused kernel 내에서 처리하여 추가 메모리 불필요
output = F.scaled_dot_product_attention(
    query, key, value,
    is_causal=True
)

# 커스텀 attention mask 사용
attn_mask = torch.tril(torch.ones(seq_len, seq_len, device='cuda', dtype=torch.bool))
output = F.scaled_dot_product_attention(
    query, key, value,
    attn_mask=attn_mask
)

11.4 Backend 선택 조건

PyTorch SDPA가 FlashAttention backend를 선택하기 위한 조건:

dtype: float16 또는 bfloat16 (float32는 불가)
device: CUDA GPU (CPU 지원 불가)
head dimension: 최대 256 (FlashAttention-2 기준)
attention mask: boolean mask 또는 is_causal=True 지원, 임의의 float mask는 비지원

이 조건을 만족하지 않으면, PyTorch는 자동으로 memory-efficient attention 또는 math backend로 fallback한다.

11.5 실무 적용 팁

# 어떤 backend가 사용되는지 확인
import torch.backends.cuda

# 각 backend의 활성화 상태 확인
print(f"Flash SDP enabled: {torch.backends.cuda.flash_sdp_enabled()}")
print(f"Mem efficient SDP enabled: {torch.backends.cuda.mem_efficient_sdp_enabled()}")
print(f"Math SDP enabled: {torch.backends.cuda.math_sdp_enabled()}")

# 전역적으로 특정 backend 비활성화
torch.backends.cuda.enable_flash_sdp(False)  # FlashAttention 비활성화
torch.backends.cuda.enable_mem_efficient_sdp(True)

11.6 flash-attn 라이브러리 직접 사용

PyTorch 네이티브 SDPA 외에도, Tri Dao의 flash-attn 패키지를 직접 사용할 수 있다. 이 패키지는 PyTorch SDPA보다 더 많은 기능(예: sliding window attention, ALiBi, cross-attention 최적화)을 제공한다:

# pip install flash-attn
from flash_attn import flash_attn_func

# (batch, seqlen, nheads, headdim) 형태
output = flash_attn_func(q, k, v, causal=True)

12. 정리 및 핵심 교훈

FlashAttention의 핵심 교훈은 알고리즘의 FLOP 복잡도만이 성능을 결정하지 않는다는 것이다. 현대 GPU에서는 메모리 접근 패턴이 실제 실행 시간을 지배하며, IO-aware 알고리즘 설계가 실용적 성능에 결정적이다.

주요 기여를 요약하면:

IO-Aware 설계 원칙: GPU 메모리 계층(HBM vs SRAM)의 비대칭성을 활용한 알고리즘 설계
Tiling + Online Softmax: SRAM에 맞는 블록 단위 연산으로 $N \times N$ 행렬의 HBM materialization 제거
Recomputation 전략: Backward pass에서 중간값을 재계산하여 $O(N^2)$ -> $O(N)$ 메모리 절약, 동시에 속도 향상
최적성 증명: IO complexity 관점에서 하한을 증명하여 알고리즘의 최적성 입증
Exact Computation: 모든 최적화에도 불구하고 근사 없는 exact attention 유지

FlashAttention은 이론적 아름다움과 실용적 효과를 동시에 갖춘 드문 연구로, 현대 LLM 학습과 추론의 핵심 인프라가 되었다. PyTorch의 네이티브 통합으로 인해, 별도의 구현 없이도 F.scaled_dot_product_attention 호출만으로 그 혜택을 누릴 수 있다.

References

Dao, T., Fu, D.Y., Ermon, S., Rudra, A., & Re, C. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022. https://arxiv.org/abs/2205.14135
Dao, T. (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. ICLR 2024. https://arxiv.org/abs/2307.08691
Shah, J., Bikshandi, G., Zhang, Y., Thakkar, V., Ramani, P., & Dao, T. (2024). FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision. NeurIPS 2024 Spotlight. https://arxiv.org/abs/2407.08608
Dao-AILab. flash-attention GitHub Repository. https://github.com/Dao-AILab/flash-attention
PyTorch Documentation. torch.nn.functional.scaled_dot_product_attention. https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html
PyTorch Documentation. torch.nn.attention.sdpa_kernel. https://docs.pytorch.org/docs/stable/generated/torch.nn.attention.sdpa_kernel.html
PyTorch Blog. FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision. https://pytorch.org/blog/flashattention-3/
Milakov, M. & Gimelshein, N. (2018). Online Normalizer Calculation for Softmax. arXiv:1805.02867. https://arxiv.org/abs/1805.02867
NVIDIA. A100 Tensor Core GPU Architecture Whitepaper. https://images.nvidia.com/aem-dam/en-zz/Solutions/data-center/nvidia-ampere-architecture-whitepaper.pdf
NVIDIA. Hopper Architecture In-Depth. https://developer.nvidia.com/blog/nvidia-hopper-architecture-in-depth/