플래시어텐션: 메모리를 아끼는 정확한 어텐션

플래시어텐션: 메모리를 아끼는 정확한 어텐션

어텐션을 근사하지 않고 똑같이 계산하면서, GPU의 느린 메모리(HBM)와 빠른 메모리(SRAM) 사이를 오가는 횟수만 크게 줄여 속도와 메모리를 함께 잡은 논문입니다. 큰 $N \times N$ 어텐션 행렬을 아예 메모리에 쓰지 않는 "타일링 + 재계산"으로, GPT-2 학습을 최대 3배 빠르게 하고 더 긴 문맥을 가능하게 했습니다. 지금 대부분의 LLM 학습 · 추론 프레임워크에 들어간 기본 부품입니다.

들어가기 전에

이 해설은 아래 내용을 안다고 보고 씁니다. 처음이라면 먼저 읽고 오면 수식과 손계산을 훨씬 쉽게 따라올 수 있습니다.

• 추론 최적화: KV 캐시 · MQA/GQA · 배칭 · 추측 디코딩 · 서빙 지표 — 추론 비용이 어디서 생기는지 먼저 정리하는 수업
• 어텐션만 있으면 된다 (트랜스포머) — 순환 없이 어텐션만으로 번역 모델을 만든 트랜스포머 — 스케일드 닷-프로덕트 · 멀티헤드 어텐션 · 위치 인코딩의 수식과 의미

한눈에 보기

• 어텐션이 느린 진짜 이유가 계산량(FLOPs)이 아니라 메모리 읽기 · 쓰기(IO)라는 관점을 알게 됩니다.
• 소프트맥스를 블록 단위로 나눠 계산해도 정확히 같은 답이 나오는 온라인 소프트맥스(타일링) 요령을 손으로 계산해 보게 됩니다.
• 역전파 때 어텐션 행렬을 저장하지 않고 다시 계산하는 편이 오히려 빠른 이유를 알게 됩니다.
• BERT · GPT-2 · 긴 문서 과제에서 실제로 얼마나 빨라지고 좋아졌는지 숫자로 알게 됩니다.

초록 (번역)

트랜스포머는 긴 시퀀스에서 느리고 메모리를 많이 씁니다. 셀프 어텐션의 시간 · 메모리 복잡도가 시퀀스 길이의 제곱이기 때문입니다. 근사 어텐션 방법들은 모델 품질을 일부 내주고 계산 복잡도를 줄이려 했지만, 실제 걸리는 시간(wall-clock)은 빨라지지 않는 경우가 많았습니다. 우리는 빠진 원칙이 어텐션 알고리즘을 IO를 고려하도록(IO-aware) 만드는 것, 곧 GPU 메모리 계층 사이의 읽기 · 쓰기를 따지는 것이라고 주장합니다. 우리는 타일링으로 GPU 고대역폭 메모리(HBM)와 GPU 칩 위 SRAM 사이의 메모리 읽기 · 쓰기 횟수를 줄이는, IO를 고려한 정확한 어텐션 알고리즘 플래시어텐션(FlashAttention)을 제안합니다. 플래시어텐션의 IO 복잡도를 분석해, 표준 어텐션보다 HBM 접근이 적고 일정 범위의 SRAM 크기에서 최적임을 보입니다. 또한 플래시어텐션을 블록 희소 어텐션으로 확장해, 기존의 어떤 근사 어텐션 방법보다 빠른 근사 어텐션 알고리즘을 만듭니다. 플래시어텐션은 기존 기준선보다 트랜스포머를 빠르게 학습시킵니다. MLPerf 1.1 학습 속도 기록과 비교해 BERT-large(시퀀스 길이 512)에서 전체 학습 시간 15% 단축, GPT-2(시퀀스 길이 1K)에서 3배, long-range arena(시퀀스 길이 1K~4K)에서 2.4배 빨라집니다. 플래시어텐션과 블록 희소 플래시어텐션은 트랜스포머의 문맥을 더 길게 해 더 높은 품질의 모델(GPT-2에서 퍼플렉서티 0.7 개선, 긴 문서 분류에서 6.4점 향상)과 완전히 새로운 능력을 가능하게 합니다. Path-X 과제(시퀀스 길이 16K, 정확도 61.4%)와 Path-256(시퀀스 길이 64K, 정확도 63.1%)에서 우연보다 나은 성능을 낸 최초의 트랜스포머가 됐습니다.

무엇을 풀었나

셀프 어텐션은 시퀀스 길이 $N$에 대해 시간 · 메모리가 $N^2$으로 늘어납니다. 문맥을 두 배로 늘리면 어텐션 행렬은 네 배가 됩니다. 이를 줄이려고 희소 근사 · 저랭크 근사 같은 근사 어텐션이 많이 나왔지만, 대부분 실제 시간은 빨라지지 않아 널리 쓰이지 못했습니다.

저자들은 원인을 이렇게 봅니다. 근사 방법들은 연산량(FLOPs)을 줄이는 데 집중했지만, 현대 GPU에서는 계산보다 메모리 접근(IO)이 병목입니다.

• A100 GPU 기준: HBM은 40~80GB · 대역폭 1.5~2.0TB/s, 칩 위 SRAM은 스트리밍 멀티프로세서 108개에 각 192KB · 대역폭 약 19TB/s입니다.
• SRAM은 HBM보다 한 자릿수 빠르지만 크기는 여러 자릿수 작습니다.
• 표준 어텐션은 $N \times N$ 행렬 $S$와 $P$를 HBM에 쓰고 다시 읽습니다. 소프트맥스 · 마스킹 · 드롭아웃 같은 연산은 계산은 적고 메모리 접근이 많은 메모리 병목 연산입니다.

PyTorch 같은 파이썬 인터페이스로는 이런 메모리 접근을 세밀하게 제어할 수 없다는 점도 문제였습니다.

전체 흐름

12345678910flowchart LR
  A[Q · K · V<br/>HBM에 저장] --> B[K · V를 블록으로 나눠<br/>SRAM에 올리기 바깥 반복]
  B --> C[Q 블록을 SRAM에 올리기<br/>안쪽 반복]
  C --> D[SRAM 안에서<br/>S블록 = Q블록 · K블록ᵀ]
  D --> E[블록 최댓값 m · 합 ℓ로<br/>소프트맥스 부분 계산]
  E --> F[이전 결과를 m · ℓ로 다시 보정해<br/>출력 O 갱신 → HBM에 쓰기]
  F -->|다음 Q 블록| C
  F -->|다음 K · V 블록| B
  F --> G[최종 O = softmax QKᵀ V<br/>N×N 행렬은 HBM에 안 씀]
  G --> H[역전파: 저장한 m · ℓ로<br/>S · P를 SRAM에서 다시 계산]

핵심 아이디어

1. 표준 어텐션이 느린 이유 — HBM 왕복

한 문장 직관: 작업대(SRAM)는 좁고 창고(HBM)는 멀리 있는데, 표준 어텐션은 커다란 중간 결과물을 매번 창고에 내려놓았다가 다시 가져옵니다.

단계별 설명 (표준 구현)

1. $Q, K$를 HBM에서 읽어 $S = QK^{\top}$를 계산하고, $S$를 HBM에 씁니다.
2. $S$를 HBM에서 다시 읽어 $P = \mathrm{softmax}(S)$를 계산하고, $P$를 HBM에 씁니다.
3. $P$와 $V$를 HBM에서 읽어 $O = PV$를 계산하고, $O$를 HBM에 씁니다.

GPT-2처럼 $N = 1024$, $d = 64$이면 $N \gg d$라서 $N \times N$ 행렬이 입력 · 출력보다 훨씬 큽니다.

숫자로 해 보기 (해설용 계산): $N = 1024$, 헤드 16개, 배치 64라면 $S$의 원소 수는 $1024^2 \times 16 \times 64 \approx 10.7$억 개입니다. 16비트(2바이트)로 저장해도 약 2.1GB이고, $P$도 같은 크기입니다. 이것을 쓰고 읽는 데 드는 시간이 계산 시간보다 큽니다.

2. 타일링 — 소프트맥스를 블록 단위로 정확히 계산

한 문장 직관: 시험 점수 평균을 낼 때 반별로 "최고점 · 합계"만 적어 두면, 나중에 전체 평균을 정확히 다시 맞출 수 있습니다. 소프트맥스도 블록마다 최댓값과 합만 기억하면 블록별로 나눠 계산할 수 있습니다.

단계별 설명

1. 소프트맥스는 수치 안정성을 위해 최댓값을 빼고 계산합니다.
2. 벡터를 두 블록 $x^{(1)}, x^{(2)}$로 나눠 각 블록의 최댓값 $m$, 지수 값 $f$, 합 $\ell$을 따로 구합니다.
3. 전체 최댓값을 새로 정하고, 각 블록 값을 $e^{(\text{블록 최댓값} - \text{전체 최댓값})}$로 보정해 합칩니다.
4. 이렇게 하면 전체 벡터를 한 번에 보지 않고도 정확히 같은 소프트맥스를 얻습니다.

숫자로 해 보기 (해설용, 점수 4개를 2블록으로)

$x = [1, 2 \mid 3, 0]$이라고 합시다.

1. 블록 1: $m_1 = 2$, $f_1 = [e^{-1}, e^{0}] \approx [0.368, 1]$, $\ell_1 \approx 1.368$
2. 블록 2: $m_2 = 3$, $f_2 = [e^{0}, e^{-3}] \approx [1, 0.050]$, $\ell_2 \approx 1.050$
3. 합치기: $m = 3$, $\ell = e^{2-3} \times 1.368 + e^{3-3} \times 1.050 \approx 0.503 + 1.050 = 1.553$
4. 확인(한 번에 계산): $e^{-2} + e^{-1} + e^{0} + e^{-3} \approx 0.135 + 0.368 + 1 + 0.050 = 1.553$ — 똑같습니다.

출력 $O$도 같은 원리로, 이전 블록까지의 결과에 보정 계수를 곱하고 새 블록 기여를 더해 갱신합니다.

3. 플래시어텐션 알고리즘 — 하나의 커널로 묶기

한 문장 직관: 창고에서 재료를 한 번 가져오면 작업대에서 곱하기 · 소프트맥스 · 마스킹 · 드롭아웃 · 곱하기를 모두 끝내고, 완성품만 창고에 돌려놓습니다.

단계별 설명 (알고리즘 1 요약)

1. SRAM 크기 $M$에 맞춰 블록 크기를 정합니다: $B_c = \lceil M / 4d \rceil$, $B_r = \min(\lceil M / 4d \rceil, d)$.
2. $Q$를 $B_r$행씩, $K, V$를 $B_c$행씩 블록으로 나눕니다.
3. 바깥 반복: $K, V$ 블록 하나를 SRAM에 올립니다.
4. 안쪽 반복: $Q$ 블록과 지금까지의 $O, \ell, m$ 블록을 SRAM에 올려, 칩 위에서 $S$ 블록 · 행별 최댓값 · 지수 · 합을 계산합니다.
5. 새 최댓값 $m^{\text{new}}$과 합 $\ell^{\text{new}}$를 구하고, $O$ 블록을 보정 · 갱신해 HBM에 씁니다.
6. 이 모든 단계를 CUDA 커널 하나로 묶어(커널 융합) 중간값을 HBM에 쓰지 않습니다.

논문은 이 알고리즘이 정확히 $O = \mathrm{softmax}(QK^{\top})V$를 돌려주고, 연산은 $O(N^2 d)$, 추가 메모리는 $O(N)$만 필요함을 증명합니다(정리 1).

왜 근사가 아니라 정확한 계산인가: 근사 어텐션은 품질을 내주는데도 실제로는 느린 경우가 많았습니다. 플래시어텐션은 수학적으로 같은 결과를 내므로 모델 정의를 바꿀 필요가 없고, 기존 모델에 그대로 끼워 쓸 수 있습니다. 실제로 GPT-2 학습에서 기존 구현과 같은 퍼플렉서티가 나왔습니다.

4. 재계산 — 저장 대신 다시 계산이 더 빠르다

한 문장 직관: 큰 중간 결과물을 창고에 보관했다가 꺼내 오는 것보다, 작은 메모(최댓값 · 합)만 남겨 두고 필요할 때 작업대에서 다시 만드는 편이 빠릅니다.

단계별 설명

1. 역전파에는 보통 $S$와 $P$($N \times N$)가 필요해 HBM에 저장해 둡니다.
2. 플래시어텐션은 출력 $O$와 정규화 통계 $(m, \ell)$만 저장합니다($O(N)$ 크기).
3. 역전파 때 SRAM에 올린 $Q, K, V$ 블록과 $(m, \ell)$로 $S, P$를 다시 계산합니다.
4. 연산(FLOPs)은 늘지만 HBM 접근이 크게 줄어 전체가 빨라집니다.

이는 선택적 그래디언트 체크포인팅의 한 형태입니다. 보통 체크포인팅은 메모리를 아끼는 대신 느려지지만, 여기서는 메모리도 아끼고 속도도 빨라졌습니다.

논문 측정(GPT-2 medium, 길이 1,024, 헤드 차원 64, 헤드 16, 배치 64, A100, 순전파 + 역전파)

연산은 13% 늘었지만 HBM 접근은 약 9분의 1로 줄어 실행 시간이 약 5.7배 짧아졌습니다(표에서 계산한 비율).

5. IO 복잡도 — 얼마나 덜 오가는가

한 문장 직관: 작업대가 클수록(SRAM이 클수록) 창고를 덜 오가도 됩니다. 플래시어텐션은 그 관계를 식으로 보장합니다.

$d$가 64~128, $M$이 약 100KB 수준이면 $d^2$이 $M$보다 여러 배 작아서, 플래시어텐션의 HBM 접근이 훨씬 적습니다. 논문은 또 모든 SRAM 크기에 대해 이보다 점근적으로 더 적게 접근하는 정확한 어텐션 알고리즘은 없다는 하한도 증명합니다(명제 3).

숫자로 해 보기 (해설용, 단위를 "원소 수"로 단순화): $N = 1024$, $d = 64$, $M = 100{,}000$이라 하면 표준은 약 $N^2 \approx 105$만, 플래시는 약 $N^2 d^2 / M = 105\text{만} \times 4{,}096 / 100{,}000 \approx 4.3$만입니다. 대략 24배 차이입니다. 실제 측정(위 표)에서는 약 9배였는데, 상수와 구현 세부가 다르기 때문입니다.

6. 블록 희소 플래시어텐션

한 문장 직관: 어텐션 행렬에서 계산할 필요 없는 블록(0인 블록)을 미리 정해 두고 건너뛰면, 건너뛴 비율만큼 더 빨라집니다.

• 블록 단위 마스크 $\tilde{M}$에서 0인 블록은 계산하지 않습니다.
• HBM 접근은 $\Theta(Nd + N^2 d^2 M^{-1} s)$로, 큰 항이 0이 아닌 블록 비율 $s$만큼 줄어듭니다.
• 실험에서는 고정된 버터플라이 희소 패턴을 썼습니다. 플래시어텐션보다 2~4배 빠르고, 길이 64K까지 확장됐습니다.

결과

학습 속도

GPT-2의 퍼플렉서티는 구현과 상관없이 같았습니다(small 18.2, medium 14.2~14.3). 모델 정의를 바꾸지 않았기 때문입니다.

긴 문맥 → 더 좋은 모델

문맥을 4배(4K)로 늘려도 Megatron 1K보다 30% 빨랐고, 퍼플렉서티는 0.7 좋아졌습니다. 긴 문서 분류(MIMIC-III · ECtHR)에서는 긴 시퀀스로 6.4점 향상이 났습니다.

새로운 능력: Path-X(길이 16K)에서 61.4%, 블록 희소 버전으로 Path-256(길이 64K)에서 63.1%로, 우연(50%)보다 나은 성능을 낸 첫 트랜스포머가 됐습니다.

어텐션 자체 벤치마크: 길이 128~2K에서 표준 구현보다 최대 3배 빨랐고, 메모리는 시퀀스 길이에 선형으로 늘었습니다. 길이 512까지는 어떤 기존 방법보다 빠르고 메모리도 적게 썼으며, 1K를 넘으면 Linformer 같은 일부 근사 방법이 더 빨라지기 시작했습니다. 블록 희소 버전은 시험한 모든 근사 방법보다 빨랐습니다.

왜 중요한가 · 이후 영향

• IO를 보는 눈: "FLOPs를 줄이는 것"과 "실제로 빠른 것"은 다르다는 점을 분명히 보여 줬습니다. 이후 GPU 커널 최적화 연구 전반에 영향을 줬습니다.
• 긴 문맥 LLM의 기반: 메모리가 시퀀스 길이에 선형이 되면서, 수천~수만 토큰 문맥을 다루는 모델 학습이 현실적으로 쉬워졌습니다. 후속 플래시어텐션-2가 병렬화를 더 개선합니다.
• 표준 부품화: 정확한 계산이라 기존 모델을 바꾸지 않고 바로 교체할 수 있어, 주요 학습 · 추론 라이브러리에 빠르게 들어갔습니다.
• 근사 어텐션의 재평가: 블록 희소 확장처럼, IO를 고려하면 근사 방법도 실제 속도 이득을 낼 수 있음을 보였습니다.

한계 · 주의할 점

• CUDA로 직접 작성해야 합니다. 새 어텐션 변형마다 PyTorch보다 훨씬 낮은 수준의 CUDA 커널을 새로 짜야 하고, GPU 구조가 바뀌면 옮기기 어려울 수 있다고 논문이 밝힙니다.
• 한 GPU 안의 최적화: 단일 GPU에서는 상수 배 안에서 최적이지만, 여러 GPU에 나눠 계산할 때의 데이터 이동은 다루지 않았습니다.
• 연산량 자체는 그대로 $N^2$: 메모리 접근과 메모리 사용은 줄지만 계산량은 여전히 시퀀스 길이의 제곱입니다. 아주 긴 시퀀스에서는 Mamba 같은 다른 구조가 대안으로 연구됩니다.
• 블록 희소 버전은 근사입니다. 정확한 결과를 보장하는 것은 기본 플래시어텐션뿐입니다.
• 하드웨어 · 설정 의존: 속도 수치는 A100 GPU와 특정 배치 · 헤드 설정 기준입니다.

이어서 읽을 것

• 앞: 어텐션만 있으면 된다 (트랜스포머)
• 뒤: 플래시어텐션-2: 더 나은 병렬화 · PagedAttention: LLM 서빙의 메모리 관리 (vLLM) · Mamba: 선택적 상태 공간 모델

용어 정리

스스로 점검

1. 플래시어텐션은 FLOPs가 더 많은데도 왜 더 빠른가요?

   현대 GPU에서 어텐션은 계산보다 HBM 읽기 · 쓰기가 병목이기 때문입니다. $N \times N$ 행렬을 HBM에 쓰지 않고 SRAM 안에서 처리해, 측정 예에서 HBM 접근을 40.3GB에서 4.4GB로 줄였습니다.

2. 소프트맥스를 블록으로 나눠 계산해도 정확한 이유는 무엇인가요?

   블록마다 최댓값 $m$과 지수 합 $\ell$을 기억해 두면, 전체 최댓값 기준으로 $e^{m_{\text{블록}} - m_{\text{전체}}}$를 곱해 합치는 것만으로 한 번에 계산한 것과 똑같은 분모와 출력을 얻을 수 있기 때문입니다.

3. 역전파에서 무엇을 저장하고 무엇을 다시 계산하나요?

   출력 $O$와 정규화 통계 $(m, \ell)$만 저장하고, $S$와 $P$는 SRAM에 올린 $Q, K, V$ 블록으로 다시 계산합니다.

다음에 볼 것

이 글을 이해했다면 이어서 보면 좋은 항목입니다.

• 플래시어텐션-2: 더 나은 병렬화 — 일 나누기와 병렬화를 고쳐 플래시어텐션을 약 2배 빠르게 한 플래시어텐션-2

출처

• 원문: https://arxiv.org/abs/2205.14135 (2022 · Tri Dao 외)