제3부 딥러닝 라이브러리

3.6 어텐션 I — 메모리 문제·FlashAttention 원리

3.6.1 이 장에서 배우는 것 / 선수 지식

지금까지의 도구 — GEMM(3.1), 융합(3.3), 타일 계층(3.4), 레이아웃(3.5) — 이 모이는 첫 큰 응용이 어텐션이다. 어텐션은 두 개의 GEMM($QK^T$·PV) 사이에 softmax가 낀 구조인데, 순진하게 계산하면 $N\times N$ 중간 행렬($S=QK^T$)을 HBM에 구현해 시퀀스 길이의 제곱으로 메모리가 폭증한다. FlashAttention은 이것을 타일링 + online softmax로 HBM에 한 번도 쓰지 않고 푼다. 이 장은 그 원리와 수학을 익히고, 실전 라이브러리·적용은 3.7으로 잇는다.

배우는 것:

• 표준 어텐션이 왜 메모리 바운드가 되는지 — $S=QK^T$의 $O(N^2)$ 중간값을 연산-바이트로 분석한다.
• online(streaming) softmax 의 점화식을 직접 전개한다 — 전체 행을 보지 않고 running max·running sum으로 softmax를 누적하는 핵심 트릭.
• FlashAttention의 타일 루프가 $N\times N$을 구현하지 않고 출력을 만드는 과정을, 그리고 FLOP은 그대로인데 HBM 트래픽만 줄여 빨라지는 이유(1.10 루프라인)를 설명한다.
• 메모리가 $O(N^2)\to O(N)$으로 줄어 긴 컨텍스트가 가능해지는 이차 효과를 안다.

선수 지식: 3.1(GEMM·배치 GEMM), 1.10(루프라인·메모리 바운드 — 이 장의 핵심 프레임), 1.9(HBM 대역폭), 3.3(융합·HBM 왕복 제거), 1.7(공유메모리·타일), 1.13(정밀도·softmax는 FP32 누적). 트랜스포머·self-attention 수식은 안다고 가정. 본 장은 원리 중심 — 실전 커널·라이브러리 선택은 3.7, CuTe 구현 독해는 3.5.

3.6.2 표준 어텐션과 N² 문제

scaled dot-product attention(한 헤드, 시퀀스 길이 N, 헤드 차원 d):

두 GEMM($QK^T$, PV) 사이에 행별 softmax. 문제는 가운데 S·P가 $N\times N$ 이라는 것. N이 커지면:

배치 B·헤드 H까지 곱하면 S만으로 HBM이 폭발한다. 표준 구현은 이 S를 HBM에 쓰고(GEMM1), 다시 읽어 softmax 하고(쓰고), 또 읽어 PV(GEMM2) 한다.

왜 메모리 바운드인가(연산-바이트, 1.10). GEMM 두 개는 연산이 많지만($O(N^2\cdot d)$), 그 사이의 softmax는 $N\times N$ 데이터에 대한 element-wise exp·합 — 연산 강도가 낮은 순수 메모리 바운드. 그리고 S·P를 HBM에 구현·재독하는 트래픽이 전체를 지배한다:

표준 어텐션 HBM 트래픽(개략):

$N^2$ 텐서를 여러 번 왕복한다. N이 커지면 대역폭(1.9)이 천장.

핵심: FLOP은 $O(N^2d)$로 어쩔 수 없지만, $N\times N$ 중간값의 HBM 왕복은 본질이 아니다. 그 왕복을 없애는 것이 FlashAttention.

3.6.3 핵심 아이디어 — N×N을 구현하지 않는다

FlashAttention의 한 문장: $S=QK^T$를 타일(블록)로 분할해 계산하고, 각 블록의 기여를 출력에 누적하되, $N\times N$ 전체를 HBM에 절대 구현하지 않는다. S의 한 블록은 공유메모리·레지스터(1.7)에 잠깐 머물고 즉시 소비된다 — 3.3의 "가상 텐서" 융합을 어텐션에 적용한 셈.

문제: softmax는 행 전체를 봐야 한다(분모 = 행 전체 exp 합, 분자도 max 정규화 필요). 블록으로 분할하면 "아직 안 본 블록"이 분모를 바꾼다. 이 난관을 푸는 것이 online softmax — 전체를 보지 않고도 나중에 들어온 블록으로 이미 누적한 결과를 보정하는 점화식이다.

3.6.4 online softmax — 핵심 수학

먼저 안정적 softmax(복습). 한 행 x = [x_1..x_N]:

• $m = \max_j x_j$ (오버플로 방지용 최대값)
• $\mathrm{softmax}_j = \exp(x_j - m) / \sum_k \exp(x_k - m)$

m을 빼지 않으면 exp가 오버플로(1.13). 그래서 max를 알아야 한다 — 그런데 블록 스트리밍에선 전체 max를 아직 모른다. online softmax는 "지금까지 본 max" 로 누적하다가 더 큰 값이 나오면 이미 누적한 합·출력을 소급 보정한다.

점화식. 출력 한 행을 블록 단위로 누적한다. 상태 세 개를 유지:

• m = 지금까지 본 최대 점수(running max)
• ℓ = 지금까지의 정규화 합 $\Sigma\exp(x - m)$(running sum)
• O = 지금까지의 가중 출력 누적(running output, [d] 벡터)

새 K·V 블록 j가 들어와 그 블록의 점수 S_block(길이 Bk)과 블록 로컬 max $\tilde{m} = \max(S_{\mathrm{block}})$를 계산하면:

• $m_{\mathrm{new}} = \max(m_{\mathrm{old}}, \tilde{m})$ — 새 전역 max
• $\alpha = \exp(m_{\mathrm{old}} - m_{\mathrm{new}})$ — 옛 상태 보정 계수($\le 1$)
• $p = \exp(S_{\mathrm{block}} - m_{\mathrm{new}})$ [길이 Bk] — 이 블록의 비정규 가중치
• $\ell_{\mathrm{new}} = \alpha \cdot \ell_{\mathrm{old}} + \sum p$ — 합 보정 후 누적
• $O_{\mathrm{new}} = \alpha \cdot O_{\mathrm{old}} + p \cdot V_{\mathrm{block}}$ [d] — 출력도 같은 $\alpha$로 보정 후 누적

마지막에 한 번만 정규화: $O_{\mathrm{final}} = O_{\mathrm{new}} / \ell_{\mathrm{new}}$.

요지: 새 블록에서 더 큰 점수($\tilde{m} > m_{\mathrm{old}}$)가 나오면, 옛 누적은 "더 작은 max로 과대평가"돼 있다. $\alpha = \exp(m_{\mathrm{old}} - m_{\mathrm{new}}) < 1$ 로 옛 $\ell \cdot O$를 줄여(보정) 새 기준에 맞춘 뒤 새 블록을 더한다. max가 안 커지면 $\alpha=1$ (보정 없음). → 어떤 순서로 블록이 와도 정확히 같은 결과.

이 점화식이 FlashAttention의 핵심이다 — 전체 $N\times N$을 한 번에 안 보고도 정확한(근사 아님) softmax 가중 출력을 만든다. 블록 순서·개수와 무관하게 수학적으로 동일.

3.6.5 워크스루 A — online softmax 수계산

작게 검산한다. 한 행 점수 x = [1, 3, 2, 5], 값 V는 스칼라 v = [10, 20, 30, 40](d=1로 단순화). 정답부터: 표준 softmax 가중 평균.

표준(전체): m=5, exp(x-5)=[e⁻⁴,e⁻²,e⁻³,e⁰]≈[0.0183,0.1353,0.0498,1], 합 $\ell\approx1.2034$. 출력 $O=\Sigma p\cdot v / \ell\approx(0.183+2.707+1.494+40)/1.2034 \approx44.38/1.2034 \approx36.88$.

online(블록 2개: [1,3] then [2,5]):

블록1 [1,3], v=[10,20]:

• $\tilde{m} = 3$, $m_{\mathrm{new}} = \max(-\infty,3)=3$, $\alpha = \exp(-\infty-3)=0$
• $p = \exp([1,3]-3) = [e^{-2},1] \approx [0.1353, 1]$
• $\ell = 0 \cdot 0 + (0.1353+1) = 1.1353$
• $O = 0 \cdot 0 + (0.1353 \cdot 10 + 1 \cdot 20) = 1.353 + 20 = 21.353$

블록2 [2,5], v=[30,40]:

• $\tilde{m} = 5$, $m_{\mathrm{new}} = \max(3,5)=5$, $\alpha = \exp(3-5)=e^{-2} \approx 0.1353$
• $p = \exp([2,5]-5) = [e^{-3},1] \approx [0.0498, 1]$
• $\ell = 0.1353 \cdot 1.1353 + (0.0498+1) = 0.1536 + 1.0498 = 1.2034$ ✅ 표준과 일치
• $O = 0.1353 \cdot 21.353 + (0.0498 \cdot 30 + 1 \cdot 40) = 2.889 + 41.494 = 44.383$
• $O_{\mathrm{final}} = O / \ell = 44.383 / 1.2034 \approx 36.88$ ✅ 표준과 일치

블록을 분할해 누적했는데 전체로 계산한 것과 정확히 같다. $\alpha=\exp(3-5)$가 블록1의 누적(작은 max 3 기준)을 블록2의 새 기준(5)으로 끌어내려 맞춘 것이 핵심. 순서를 바꿔도(블록2 먼저) 같은 결과가 나온다 — 직접 해 보면 보정 계수가 반대로 적용돼 동일.

체크포인트 선행: 이 수계산을 막힘없이 재현할 수 있으면 FlashAttention의 80%를 이해한 것이다. 나머지는 이것을 GPU 타일·텐서코어 위에 구현하는 엔지니어링이다.

3.6.6 워크스루 B — FlashAttention 타일 루프

위 점화식을 2D(헤드 차원 d)·타일로 확장한 FlashAttention의 골격(의사코드, FA2 스타일 — Q를 바깥 루프로):

# Q,K,V: [N, d].  블록 크기 Bq, Bk.  O: [N, d] 출력
for i in range(0, N, Bq):                       # Q 블록 (각 블록이 한 threadblock, 병렬)
    Qi = Q[i:i+Bq]                              # [Bq, d]  → 공유메모리/레지스터
    Oi = zeros(Bq, d);  mi = -inf(Bq);  li = zeros(Bq)   # 행별 running 상태
    for j in range(0, N, Bk):                   # K,V 블록 (안쪽 루프, 누적)
        Kj, Vj = K[j:j+Bk], V[j:j+Bk]          # [Bk, d] → 공유메모리
        Sij = (Qi @ Kj.T) * scale              # [Bq, Bk]  ← 텐서코어 GEMM, 공유메모리에만!
        # (인과 마스크면 j>i 블록 skip / 부분 마스크)
        m_new = maximum(mi, rowmax(Sij))       # [Bq]
        P = exp(Sij - m_new[:,None])           # [Bq, Bk]  비정규 가중치
        alpha = exp(mi - m_new)                # [Bq] 옛 상태 보정
        li = alpha*li + rowsum(P)              # 합 누적
        Oi = alpha[:,None]*Oi + P @ Vj         # [Bq,d] 출력 누적 ← 둘째 텐서코어 GEMM
        mi = m_new
    Oi = Oi / li[:,None]                        # 최종 정규화
    O[i:i+Bq] = Oi                              # HBM에 O만 쓴다 (S·P는 안 씀)

읽는 법:

• Sij는 공유메모리에서만 산다 — [Bq,Bk]로 작고, 쓰이고 즉시 소비. $N\times N$이 HBM에 안 생긴다.
• 두 GEMM(Qi@Kjᵀ, P@Vj)이 텐서코어로(1.11). 그 사이 softmax는 레지스터·공유메모리에서. 이것이 3.3의 융합을 어텐션에 적용한 것 — 두 GEMM과 softmax를 한 커널로 융합.
• 안쪽 루프가 3.6.4 점화식 그대로. alpha로 보정, Oi·li 누적.
• 인과(causal) 마스크: j>i인 블록 전체를 건너뛰어(미래 토큰) 연산을 절반으로. 대각 블록만 부분 마스크.
• HBM에 쓰는 것은 O(N×d)뿐. 메모리 $O(N^2)\to O(N)$.

왜 빨라지나(루프라인, 1.10). FLOP은 표준과 같다(두 GEMM은 동일). 줄어든 것은 HBM 트래픽 — $N\times N$ 텐서를 안 쓰고·안 읽는다. 어텐션은 그 트래픽 때문에 메모리 바운드였으므로, 트래픽을 없애면 연산 바운드(텐서코어 포화)로 이동해 빨라진다. "연산을 줄여서"가 아니라 "메모리를 줄여서" 빠른 것 — 루프라인에서 점이 천장 아래(memory)에서 천장 위(compute)로 옮겨간다.

3.6.7 backward — 재계산 vs 저장

학습(backward)은 forward의 S·P가 필요하다. FlashAttention은 $N\times N$을 저장하지 않으므로 backward에서 블록별로 재계산한다(S 다시 $QK^T$).

   trade: 저장 안 함(메모리 O(N)) ↔ backward에서 S·P 재계산(FLOP 추가)
   그런데 어텐션은 메모리 바운드였으므로, 재계산(연산)을 더해도
   HBM 왕복을 없앤 이득이 보통 더 커서 전체로도 빠르다.  (1.10 참조)

forward에서 행 통계(m, ℓ 또는 logsumexp $L=m+\log\ell$)만 O(N)로 저장해 두면 backward 재계산이 한결 싸진다. "메모리 폭증을 막는 대신 약간의 재계산"이라는 이 교환이 FA의 학습 측 핵심이다.

3.6.8 중간 점검 (체크포인트)

• [ ] 표준 어텐션이 메모리 바운드인 이유를 $N\times N$ 중간 행렬의 HBM 왕복으로 설명한다.
• [ ] online softmax 점화식(m_new·$\alpha$·ℓ·O)을 외워 3.6.5 수계산을 재현한다.
• [ ] FlashAttention이 FLOP은 그대로 두고 HBM 트래픽을 줄여 빨라진다는 점을, 루프라인에서 점이 옮겨가는 그림으로 말한다.
• [ ] 메모리 $O(N^2)\to O(N)$이 왜 긴 컨텍스트를 가능케 하는지, backward 재계산 교환을 안다.

증상→의심: 긴 시퀀스에서 OOM → $N\times N$ 구현하는 표준 경로(eager attention) 사용 중, FA 커널로 전환. FA 켰는데 안 빨라짐 → 짧은 시퀀스(메모리 바운드 아니라 이득 작음)이거나 헤드 차원·정렬이 커널 지원 밖(폴백). 수치가 미세하게 다름 → softmax FP32 누적 여부·연산 순서(정상 범위) 또는 마스크 처리 차이.

3.6.9 트레이드오프 — 언제 FlashAttention인가

견해: 현대 트랜스포머에서 어텐션은 사실상 항상 FA 계열 융합 커널로 실행한다(PyTorch SDPA가 자동 선택). 직접 신경 쓸 일은 ① 헤드 차원·dtype·마스크가 커널 지원 범위 밖이라 느린 폴백으로 떨어질 때, ② 추론 디코딩처럼 다른 병렬화가 필요할 때(3.7)다.

비NVIDIA 정직: FlashAttention은 NVIDIA 독점 기술이 아니라 알고리즘(Dao et al.)이며, 핵심은 online softmax + 타일 융합이라는 수학·IO 인식이다. 같은 원리가 AMD·TPU·CPU에도 이식됐고, online softmax 자체는 그 전부터 알려진 streaming 기법이다. NVIDIA의 기여는 이를 텐서코어·cp.async·(Hopper) TMA·WGMMA로 극한 구현한 것(FA2·FA3, cuDNN SDPA, 3.7·4.2).

3.6.10 흔한 함정

• online softmax의 보정을 빠뜨림. 블록 누적 시 α=exp(m_old-m_new)로 옛 ℓ·O를 둘 다 보정해야 한다. 하나만 보정하면 결과가 틀린다(분모·분자 불일치). 3.6.5 수계산으로 검증.
• m 빼기 생략 → 오버플로. online이라도 exp는 항상 (S - m_new)로. 안 빼면 긴 시퀀스·큰 점수에서 inf(1.13).
• FA가 항상 빠를 거라는 기대. 짧은 시퀀스·작은 배치는 $N^2$ 트래픽이 작아 이득이 미미하고, 디코딩(쿼리 1)엔 다른 병렬화가 필요. 무조건 빠르지 않다 — 형상 의존.
• 헤드 차원·dtype 미지원 폴백. FA 커널은 보통 특정 d(64·128 등)·dtype·정렬에 최적화. 벗어나면 느린 일반 경로로 폴백하면서 "FA 켰는데 느리다". 지원 범위 확인(3.7).
• 마스크·바이어스 처리 불일치. 인과 마스크는 블록 skip으로 싸지만, 커스텀 마스크·ALiBi·sliding window는 커널이 지원해야 한다. 미지원이면 잘못된 결과나 폴백.
• 수치 재현성 기대. 블록 누적 순서·FP32 부분합 차이로 비트 단위는 표준과 다를 수 있다(수학적으론 동일, 부동소수 순서 차이). 정상.

3.6.11 연습문제

연습 1 (N² 비용). N=16384, d=128, FP16, 배치 B=8, 헤드 H=32. 표준 어텐션이 구현하는 S 한 텐서(한 배치·한 헤드)의 크기와, $B\cdot H$ 전체 S의 총 메모리를 구하라. 왜 이것이 문제인가?

연습 2 (online softmax). 행 x=[0,2,4], 값 v=[1,1,1](d=1). 블록 [0,2] then [4]로 online softmax를 갱신해 ℓ·O_final을 구하고, v가 전부 1이므로 O_final이 1이어야 함을 확인하라.

연습 3 (루프라인 이동). "FlashAttention은 FLOP을 줄이지 않는데 어떻게 빨라지나?"를 1.10 루프라인 용어(산술 강도·memory bound·compute bound)로 설명하라.

연습 4 (backward 교환). FlashAttention은 forward의 $S\cdot P$를 저장하지 않아 backward에서 재계산한다. 재계산(추가 FLOP)을 하는데도 전체가 더 빠를 수 있는 이유와, forward에서 무엇을 O(N)로 저장해 재계산을 줄이는지.

연습 5 (언제 안 쓰나). FlashAttention이 이득이 거의 없거나 다른 변형이 필요한 두 경우를 들고, 각각 무엇으로 대응하는지.

3.6.12 요약 · 다음 장

• 표준 어텐션은 $S=QK^T$의 $N\times N$ 중간 행렬을 HBM에 구현·왕복해 시퀀스 길이의 제곱으로 메모리·대역폭 요구가 폭증한다 — 메모리 바운드(1.10).
• FlashAttention은 Q·K·V를 타일링하고 online softmax(running max·보정 계수 $\alpha$·running sum)로 $N\times N$을 구현하지 않고 출력을 누적한다. 두 GEMM과 softmax를 한 커널로 융합(3.3의 어텐션판).
• FLOP은 그대로, HBM 트래픽만 줄여 메모리 바운드를 탈출해 빨라진다. 메모리는 $O(N^2)\to O(N)$ — 긴 컨텍스트가 가능. backward는 통계만 저장하고 블록 재계산.
• online softmax 수계산(3.6.5)을 재현할 수 있으면 원리의 핵심을 잡은 것.

다음 장: 3.7 어텐션 II — 융합 커널·실전 라이브러리. 이 원리를 실제로 어디서 어떻게 쓰는지 — PyTorch SDPA의 백엔드 자동선택, FlashAttention 2·3·cuDNN SDPA·xFormers의 차이, 추론 디코딩용 FlashDecoding, 헤드 차원·마스크·dtype 지원 범위와 폴백 진단을 다룬다.