플래시어텐션-2: 더 나은 병렬화

플래시어텐션-2: 더 나은 병렬화

플래시어텐션은 메모리 접근을 줄여 어텐션을 빠르게 했지만, GPU 이론 최대 성능의 25~40%밖에 쓰지 못했습니다. 플래시어텐션-2는 계산 순서 · 병렬화 단위 · 워프 간 일 나누기를 바꿔 약 2배 더 빠르게 만들고, A100에서 이론 최대의 50~73%까지 끌어올렸습니다. 결과는 여전히 근사 없는 정확한 어텐션입니다.

들어가기 전에

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

• 플래시어텐션: 메모리를 아끼는 정확한 어텐션 — 타일링과 재계산으로 HBM 접근을 줄여 정확한 어텐션을 빠르고 메모리 적게 계산한 플래시어텐션

한눈에 보기

• 플래시어텐션이 메모리 문제는 풀었는데도 왜 GPU 성능을 다 못 썼는지 알게 됩니다.
• 행렬 곱이 아닌 연산(나눗셈 · 지수 등)이 왜 16배 비싼지, 그래서 무엇을 줄였는지 손으로 계산해 봅니다.
• 시퀀스 길이 방향으로도 병렬화해 GPU의 SM을 더 많이 채우는 방법을 알게 됩니다.
• 워프 사이의 일 나누기를 "split-K"에서 "split-Q"로 바꿔 공유 메모리 왕복을 없앤 이유를 알게 됩니다.

초록 (번역)

트랜스포머를 더 긴 시퀀스로 확장하는 것은 지난 몇 년간 큰 문제였습니다. 언어 모델링과 고해상도 이미지 이해의 성능을 높이고, 코드 · 오디오 · 비디오 생성에서 새로운 응용을 열 것으로 기대되기 때문입니다. 어텐션 층은 실행 시간과 메모리가 시퀀스 길이의 제곱으로 늘어나 긴 시퀀스로 확장하는 데 주된 병목입니다. 플래시어텐션은 비대칭적인 GPU 메모리 계층을 활용해 근사 없이 상당한 메모리 절약(제곱이 아닌 선형)과 실행 시간 단축(최적화된 기준선 대비 2~4배)을 이뤘습니다. 그러나 플래시어텐션은 여전히 최적화된 행렬 곱(GEMM) 연산만큼 빠르지 않아, 이론 최대 FLOPs/s의 25~40%에 그칩니다. 우리는 이 비효율이 GPU의 서로 다른 스레드 블록과 워프 사이에 일을 최적이 아니게 나눈 탓에 점유율이 낮거나 불필요한 공유 메모리 읽기 · 쓰기가 생기기 때문임을 관찰합니다. 우리는 이 문제들을 다루는, 일을 더 잘 나누는 플래시어텐션-2를 제안합니다. 구체적으로 (1) 알고리즘을 다듬어 행렬 곱이 아닌 FLOPs의 수를 줄이고, (2) 헤드 하나에 대해서도 어텐션 계산을 서로 다른 스레드 블록에 병렬화해 점유율을 높이며, (3) 각 스레드 블록 안에서 워프들 사이에 일을 나눠 공유 메모리를 통한 통신을 줄입니다. 이로써 플래시어텐션 대비 약 2배 빨라져 A100에서 이론 최대 FLOPs/s의 50~73%에 이르고, GEMM 연산의 효율에 가까워집니다. GPT 방식 모델을 처음부터 끝까지 학습하는 데 쓰면, 플래시어텐션-2가 A100 GPU 한 장당 최대 225 TFLOPs/s(모델 FLOPs 활용률 72%)의 학습 속도에 이른다는 것을 실험으로 확인합니다.

무엇을 풀었나

플래시어텐션(1편)은 큰 어텐션 행렬을 GPU 주 메모리(HBM)에 쓰지 않게 해, 메모리를 시퀀스 길이에 선형으로 줄이고 2~4배 빠르게 했습니다. 그래서 대규모 학습 · 추론에 널리 쓰였습니다.

그런데 문맥 길이가 계속 늘면서(논문은 GPT-4 32k, MPT 65k, Claude 100k를 예로 듭니다) 또 다른 한계가 드러났습니다.

• 최적화된 행렬 곱(GEMM)은 GPU 이론 최대 성능의 80~90%까지 씁니다.
• 플래시어텐션은 순전파가 30~50%, 역전파는 25~35%에 그쳤습니다(A100 기준).

저자의 분석으로는 원인이 일 나누기(work partitioning)에 있었습니다.

1. 행렬 곱이 아닌 연산이 너무 많았습니다. 소프트맥스 보정 같은 연산은 전체 FLOPs 중 작은 비율이지만, 텐서 코어를 못 써서 훨씬 느립니다.
2. GPU를 다 채우지 못했습니다. 시퀀스가 길면 배치 · 헤드 수가 작아져, 스레드 블록 수가 SM 수보다 적어집니다(낮은 점유율).
3. 워프끼리 공유 메모리로 자주 주고받았습니다. 한 블록 안에서 일을 나눈 방식 때문에 중간 결과를 모아 더하는 과정이 필요했습니다.

전체 흐름

12345678flowchart LR
  A[Q · K · V in HBM] --> B[바깥 반복을 Q 행 블록으로 바꿈<br/>행 블록마다 스레드 블록 하나]
  B --> C[배치 · 헤드 · 시퀀스 길이<br/>세 방향으로 병렬화]
  C --> D[블록 안: Q를 워프 4~8개로 나눔<br/>K · V는 모든 워프가 공유]
  D --> E[K · V 블록을 차례로 읽으며<br/>보정 안 한 출력 Õ 누적]
  E --> F[마지막에 한 번만 ℓ로 나눔<br/>logsumexp L = m + log ℓ 저장]
  F --> G[출력 O · L을 HBM에 쓰기]
  G --> H[역전파: L로 P 재계산<br/>열 블록 병렬 · dQ는 원자적 덧셈]

핵심 아이디어

1. 행렬 곱이 아닌 연산을 줄인다

한 문장 직관: 공장에 자동 기계(텐서 코어)와 수작업대가 있다면, 수작업이 전체 일의 5%뿐이어도 16배 느리면 시간을 크게 잡아먹습니다. 수작업을 줄여야 합니다.

왜 16배인가: A100은 FP16/BF16 행렬 곱 이론 최대가 312 TFLOPs/s이지만, 행렬 곱이 아닌 FP32 연산은 19.5 TFLOPs/s뿐입니다. $312 / 19.5 = 16$이므로 행렬 곱이 아닌 FLOP 하나는 행렬 곱 FLOP보다 16배 비쌉니다.

숫자로 해 보기 (해설용 계산): 전체 FLOPs의 95%가 행렬 곱, 5%가 행렬 곱이 아닌 연산이라고 해 봅시다. 시간은 $0.95 \times 1 + 0.05 \times 16 = 1.75$단위이고, 그중 $0.8 / 1.75 \approx 46\%$를 5%짜리 연산이 차지합니다. 그래서 이 5%를 줄이는 것이 중요합니다.

단계별 설명 (순전파의 두 가지 손질)

1. 보정을 마지막에 한 번만: 1편은 블록마다 출력에 $\mathrm{diag}(\ell)^{-1}$을 곱해 정규화했습니다. 2편은 정규화하지 않은 출력 $\tilde{O}$와 통계 $\ell$만 계속 갱신하고, 맨 마지막에 한 번만 $\ell$로 나눕니다.
2. 통계는 하나만 저장: 역전파용으로 최댓값 $m$과 지수 합 $\ell$을 둘 다 저장하던 것을, 로그합지수(logsumexp) $L = m + \log \ell$ 하나만 저장합니다.

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

점수 $x = [1, 2 \mid 3, 0]$, 값 $v = [10, 20 \mid 30, 0]$이라 합시다.

1. 블록 1: $m = 2$, 지수 $[e^{-1}, e^{0}] \approx [0.368, 1]$, $\ell \approx 1.368$, $\tilde{O} = 0.368 \times 10 + 1 \times 20 \approx 23.68$
2. 블록 2: 새 최댓값 $m = 3$. 이전 누적에 $e^{2-3} \approx 0.368$을 곱하면 $\tilde{O} \approx 8.71$, 여기에 $e^{0} \times 30 + e^{-3} \times 0 = 30$을 더해 $\tilde{O} \approx 38.71$. $\ell \approx 0.368 \times 1.368 + 1 + 0.050 \approx 1.553$
3. 마지막에 한 번만 나누기: $O = 38.71 / 1.553 \approx 24.93$
4. 확인(한 번에 계산): 가중치 $\approx (0.087, 0.237, 0.644, 0.032)$로 $0.087 \times 10 + 0.237 \times 20 + 0.644 \times 30 \approx 24.93$ — 같습니다.
5. 저장할 통계: $L = 3 + \ln 1.553 \approx 3.440$. 역전파에서 세 번째 원소의 가중치는 $e^{3 - 3.440} \approx 0.644$로 바로 복원됩니다.

블록마다 하던 나눗셈이 마지막 한 번으로 줄었고, 저장할 통계도 둘에서 하나로 줄었습니다.

인과 마스크 처리: 언어 모델에서 쓰는 인과 마스크(앞쪽만 보기)에서는 열 번호가 모두 행 번호보다 큰 블록(긴 시퀀스에서 약 절반)은 계산을 통째로 건너뜁니다. 그래서 마스크 없는 경우보다 약 1.7~1.8배 빨라집니다. 또 행 번호가 열 번호보다 확실히 큰 블록에는 마스크를 적용할 필요가 없어, 정사각 블록이라면 행마다 블록 하나에만 마스크를 씁니다.

2. 시퀀스 길이 방향으로도 병렬화

한 문장 직관: 작업자(SM)가 108명인데 일감 묶음이 32개뿐이면 76명이 놉니다. 일감을 더 잘게 나눠 모두에게 일을 줍니다.

단계별 설명

1. 1편은 배치 × 헤드마다 스레드 블록 하나를 만들었습니다. A100에는 SM이 108개 있습니다.
2. 이 수가 크면(예: 80 이상) GPU를 거의 다 쓰지만, 긴 시퀀스에서는 배치 · 헤드가 작아 블록이 모자랍니다.
3. 2편은 시퀀스 길이 방향으로도 나눕니다. 순전파에서는 바깥 반복을 $Q$의 행 블록으로 바꿔, 행 블록마다 독립된 스레드 블록이 맡습니다(서로 통신 필요 없음).
4. 역전파에서는 열 블록마다 스레드 블록 하나를 두고, 서로 겹치는 $dQ$ 갱신은 원자적 덧셈(atomic add)으로 합칩니다.

숫자로 해 보기 (해설용 계산): 배치 2 × 헤드 16 = 스레드 블록 32개라면 SM 108개 중 약 30%만 일합니다. 시퀀스 길이 8,192를 행 블록 128씩 나누면 행 블록 64개가 생겨, 스레드 블록이 $32 \times 64 = 2{,}048$개로 늘어나 모든 SM을 채울 수 있습니다.

(바깥 · 안쪽 반복 순서를 바꾸고 시퀀스 방향으로 병렬화하는 생각은 Triton 구현에서 Phil Tillet이 먼저 제안 · 구현했다고 논문이 밝힙니다.)

3. 워프 사이 일 나누기 — split-K에서 split-Q로

한 문장 직관: 네 사람이 각자 다른 페이지를 읽고 나중에 메모를 모아 합치는 것보다, 각자 다른 질문을 맡아 같은 책 전체를 보며 바로 자기 답을 쓰는 편이 회의(동기화)가 필요 없습니다.

단계별 설명

1. 한 스레드 블록 안에는 보통 워프 4개 또는 8개가 있습니다(워프 = 스레드 32개 묶음).
2. 1편(split-K): $K, V$를 워프 4개에 나누고 $Q$는 모두가 봅니다. 각 워프가 $QK^{\top}$의 일부와 $V$의 일부를 곱한 뒤, 결과를 공유 메모리에 써서 모아 더해야 합니다.
3. 2편(split-Q): $Q$를 워프 4개에 나누고 $K, V$는 모두가 봅니다. 각 워프가 자기 $Q$ 조각으로 $QK^{\top}$ 조각을 구해 공유된 $V$와 곱하면 바로 자기 몫의 출력이 나옵니다. 워프끼리 통신할 필요가 없습니다.
4. 역전파도 split-K를 피하도록 나눴지만, 입력 · 기울기 사이 의존성이 복잡해 일부 동기화는 남습니다.

블록 크기 조정: 블록이 클수록 공유 메모리 왕복은 줄지만 레지스터 · 공유 메모리가 더 필요합니다. 너무 크면 레지스터가 넘쳐 크게 느려지거나 아예 실행할 수 없습니다. 보통 $\{64, 128\} \times \{64, 128\}$ 중에서 헤드 차원과 장치에 맞게 손으로 골랐고, 자동 조정은 앞으로의 과제로 남겼습니다.

4. MQA · GQA 지원

여러 질의 헤드가 같은 키 · 값 헤드를 공유하는 다중 질의 어텐션(MQA)과 그룹 질의 어텐션(GQA)은 추론 때 KV 캐시를 줄이려고 씁니다. 2편은 키 · 값 헤드를 실제로 복사하지 않고 헤드 인덱스만 조정해 같은 계산을 하고, 역전파에서는 암묵적으로 복제된 헤드들의 $dK, dV$를 더합니다.

결과

어텐션 단독 벤치마크(A100 80GB SXM4, 시퀀스 512~16k, 배치는 총 토큰 16k가 되게, 은닉 2,048, 헤드 차원 64 또는 128)

• 플래시어텐션보다 1.7~3.0배, Triton 구현 플래시어텐션보다 1.3~2.5배, PyTorch 표준 구현보다 3~10배 빨랐습니다.
• 순전파 + 역전파 기준 약 2배 빨랐고, 최대 230 TFLOPs/s(이론 최대의 73%)에 이르렀습니다. 역전파는 최대 63%였습니다.
• 같은 코드를 H100에서 새 기능(TMA · 4세대 텐서 코어) 없이 돌려도 최대 335 TFLOPs/s가 나왔습니다. 새 명령어를 쓰면 1.5~2배 더 빨라질 것으로 기대한다고 밝힙니다.

(FLOPs 계산: 순전파 $4 \times \text{길이}^2 \times \text{헤드 차원} \times \text{헤드 수}$, 인과 마스크면 절반, 역전파는 순전파의 2.5배 — 순전파 행렬 곱 2개, 역전파 5개이기 때문)

숫자로 해 보기 (해설용 계산): 길이 2,048, 헤드 차원 128, 헤드 16이면 순전파 FLOPs는 $4 \times 2{,}048^2 \times 128 \times 16 \approx 3.4 \times 10^{10}$입니다. 이것을 230 TFLOPs/s로 처리하면 약 0.15밀리초가 걸리는 셈입니다.

GPT 방식 모델 학습(A100 8장, GPU당 TFLOPs/s)

플래시어텐션 없는 기준보다 최대 2.8배, 플래시어텐션보다 최대 1.3배 빨랐고, 최고 225 TFLOPs/s(모델 FLOPs 활용률 72%)였습니다. 학습 FLOPs는 Megatron-LM 관례대로 $6 \times \text{길이} \times \text{파라미터 수} + 12 \times \text{층 수} \times \text{은닉 차원} \times \text{길이}^2$로 셌습니다(인과 마스크로 어텐션 항을 절반으로 줄이지 않음).

왜 중요한가 · 이후 영향

• 같은 비용으로 두 배 긴 문맥: 논문은 2배 빨라졌다는 것이 곧 이전 8k 문맥 모델 비용으로 16k 문맥 모델을 학습할 수 있다는 뜻이라고 말합니다. 긴 문서 · 고해상도 이미지 · 오디오 · 비디오 모델로 가는 길을 넓혔습니다.
• 어텐션이 행렬 곱 수준의 효율로: 이론 최대의 73%까지 가면서, 어텐션이 더 이상 "특별히 비효율적인 연산"이 아니게 됐습니다.
• GPU 프로그래밍의 교훈: 같은 알고리즘이라도 스레드 블록 · 워프에 일을 어떻게 나누느냐에 따라 성능이 2배 차이 날 수 있음을 보여 줬습니다.
• 실사용: 공개 라이브러리로 배포돼 학습 · 미세조정 · 추론에 바로 쓰였고, MQA · GQA처럼 최신 모델이 쓰는 어텐션 변형도 지원합니다.

한계 · 주의할 점

• 하드웨어 특화: 수치는 주로 A100 기준이고, H100의 새 기능(TMA · 4세대 텐서 코어)은 아직 쓰지 않았습니다. 다른 장치(AMD GPU 등)는 앞으로의 과제로 남겼습니다.
• 수동 조정: 블록 크기를 헤드 차원마다 손으로 골랐고, 자동 조정은 하지 않았습니다.
• 계산량 자체는 그대로: 효율을 높였을 뿐 어텐션 계산량은 여전히 시퀀스 길이의 제곱입니다.
• 역전파는 아직 덜 효율적: 순전파 최대 73%에 비해 역전파는 최대 63%로, 의존성 때문에 동기화가 일부 남습니다.
• FLOPs 계산 관례: 학습 처리량은 인과 마스크의 절반 계산을 반영하지 않는 관례식으로 셌으므로, 다른 논문 수치와 비교할 때 기준을 맞춰야 합니다.

이어서 읽을 것

• 앞: 플래시어텐션: 메모리를 아끼는 정확한 어텐션 · 어텐션만 있으면 된다 (트랜스포머)
• 함께: PagedAttention: LLM 서빙의 메모리 관리 (vLLM) · Mamba: 선택적 상태 공간 모델

용어 정리

스스로 점검

1. 플래시어텐션이 메모리 문제를 풀었는데도 이론 최대 성능의 25~40%밖에 못 쓴 이유는 무엇인가요?

   스레드 블록과 워프 사이에 일을 나누는 방식이 최적이 아니어서, 긴 시퀀스에서 점유율이 낮거나 워프끼리 공유 메모리로 불필요하게 주고받았고, 느린 행렬 곱이 아닌 연산도 많았기 때문입니다.

2. 순전파에서 행렬 곱이 아닌 연산을 줄인 두 가지 손질은 무엇인가요?

   출력을 블록마다 정규화하지 않고 마지막에 한 번만 $\ell$로 나누는 것, 그리고 역전파용 통계로 $m$과 $\ell$ 대신 logsumexp $L = m + \log \ell$ 하나만 저장하는 것입니다.

3. split-Q가 split-K보다 빠른 이유는 무엇인가요?

   $Q$를 워프마다 나누면 각 워프가 공유된 $K, V$로 자기 몫의 출력을 바로 완성할 수 있어, 중간 결과를 공유 메모리에 써서 모아 더하는 통신 · 동기화가 필요 없기 때문입니다.

다음에 볼 것

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

• PagedAttention: LLM 서빙의 메모리 관리 (vLLM) — KV 캐시를 페이지처럼 블록으로 나눠 관리해 LLM 서빙 처리량을 2~4배 높인 vLLM

출처

• 원문: arxiv.org/abs/2307.08691 (2023 · Tri Dao)