플래시 어텐션

FlashAttention; 플래시 어텐션
어텐션 결과는 그대로 두고, GPU의 느린 주 메모리와 빠른 온칩 메모리 사이의 데이터 이동을 줄여 정확한 어텐션을 빠르고 적은 메모리로 계산하는 알고리즘

플래시 어텐션(FlashAttention)은 2022년 스탠퍼드 대학교의 트리 다오(Tri Dao) 등이 발표한 어텐션 계산 알고리즘이다.^[1] 트랜스포머의 셀프 어텐션은 시퀀스 길이 N에 대해 N×N 점수 행렬을 만드는데, 표준 구현은 이 행렬을 GPU 주 메모리(HBM)에 썼다가 다시 읽느라 시간과 메모리를 많이 쓴다. 플래시 어텐션은 입력을 블록으로 잘라 빠른 온칩 메모리(SRAM) 안에서 계산을 끝내고 N×N 행렬을 아예 메모리에 만들지 않는다.

근사가 아니라 수학적으로 같은 결과를 내는 정확한(exact) 어텐션이라는 점이 그 이전의 효율적 어텐션 연구(희소·저랭크 근사)와 다르다. 연산량(FLOPs)은 줄이지 않고 오히려 조금 늘지만, 메모리 접근을 크게 줄여 실제 실행 시간을 줄인다. 저자들은 이를 "IO 인식"(IO-aware) 알고리즘이라 불렀다. 이후 버전 2(2023), 3(2024), 4(2026)가 나오며 GPU 세대마다 다시 최적화되었고, 파이토치의 표준 어텐션 함수에도 백엔드로 들어갔다.^[2] 긴 컨텍스트 윈도우를 가진 대형 언어 모델을 가능하게 한 핵심 기반 기술 중 하나다.

배경: GPU 메모리 계층

GPU에는 크기와 속도가 다른 메모리가 여러 층으로 있다. 플래시 어텐션 논문은 엔비디아 A100을 예로 다음 수치를 들었다.^[1]

메모리	위치	용량(A100)	대역폭(A100)
HBM(고대역폭 메모리)	GPU 칩 바깥, 같은 패키지	40~80GB	1.5~2.0TB/s
SRAM(공유 메모리·L1)	각 스트리밍 멀티프로세서(SM) 안	SM당 192KB, SM 108개	약 19TB/s(추정)

온칩 SRAM은 HBM보다 한 자릿수 빠르지만 크기는 몇 자릿수 작다. 연산 속도가 메모리 속도보다 빠르게 발전해 왔기 때문에, 오늘날 트랜스포머의 많은 연산은 계산이 아니라 HBM 접근이 병목인 메모리 제한(memory-bound) 연산이다.^[1] 소프트맥스, 마스킹, 드롭아웃처럼 원소별 연산은 특히 그렇다.

표준 어텐션의 문제

어텐션 $O = s o f t m a x (Q K^{⊤} / \sqrt{d}) V$ (Q, K, V는 N×d 행렬)를 단순히 구현하면 다음과 같이 계산한다.

Q, K를 HBM에서 읽어 $S = Q K^{⊤}$ 를 계산하고 N×N 행렬 S를 HBM에 쓴다.
S를 읽어 $P = s o f t m a x (S)$ 를 계산하고 N×N 행렬 P를 HBM에 쓴다.
P와 V를 읽어 $O = P V$ 를 계산하고 O를 HBM에 쓴다.

마스킹이나 드롭아웃이 끼면 N×N 행렬을 읽고 쓰는 단계가 더 늘어난다. 역전파를 위해 P를 저장해 두어야 하므로 메모리도 $O (N^{2})$ 가 든다. N이 수천~수만이면 이 왕복이 실행 시간과 메모리를 지배한다. 논문은 표준 어텐션의 HBM 접근량을 $Θ (N d + N^{2})$ 로 분석했다.^[1]

플래시 어텐션의 원리

타일링과 커널 융합

Q, K, V를 SRAM에 들어갈 크기의 블록으로 나눈다. K, V의 블록을 SRAM에 올리고, 그 블록에 대해 Q의 블록들을 차례로 불러 점수 계산, 소프트맥스, V와의 곱을 SRAM 안에서 한 번에 끝낸 뒤 출력 블록만 HBM에 쓴다. 모든 단계를 하나의 GPU 커널로 융합하므로 중간 결과인 S와 P는 HBM에 전혀 쓰이지 않는다.^[1]

온라인 소프트맥스

문제는 소프트맥스가 행 전체를 봐야 계산된다는 점이다. 분모 $\sum_{j} e^{s_{j}}$ 와, 지수 함수의 넘침을 막기 위해 빼는 최댓값 $m = \max_{j} s_{j}$ 가 모두 행 전체에 의존한다. 플래시 어텐션은 블록을 하나씩 볼 때마다 "지금까지의 최댓값과 분모"를 갱신하고 이미 계산한 부분 결과를 다시 보정하는 온라인 소프트맥스를 쓴다. 이 기법은 2018년 엔비디아의 밀라코프(Maxim Milakov)와 김멜셰인(Natalia Gimelshein)이 소프트맥스의 메모리 접근을 줄이는 방법으로 발표한 것이다.^[3]

한 행에 대해 지금까지의 최댓값 $m$ , 분모 $ℓ$ , 정규화 전 출력 누적값 $𝐨$ 를 들고 있다가, 새 점수 블록 $s^{(j)}$ 와 값 블록 $V^{(j)}$ 가 오면 다음처럼 갱신한다.

$m^{'} = \max (m, \max s^{(j)}), ℓ^{'} = e^{m - m^{'}} ℓ + \sum e^{s^{(j)} - m^{'}}$

$𝐨^{'} = e^{m - m^{'}} 𝐨 + e^{s^{(j)} - m^{'}} V^{(j)}$

모든 블록을 본 뒤 $𝐨 / ℓ$ 가 정확한 어텐션 출력이다. 최댓값이 바뀔 때마다 기존 누적값에 $e^{m - m^{'}}$ 를 곱해 기준을 맞추는 것이 핵심이다.

# 한 쿼리 행 q에 대한 플래시 어텐션 순전파의 개념 코드 (실제는 GPU 커널)
m, l, o = -inf, 0.0, zeros(d)
for Kj, Vj in blocks(K, V):            # K, V를 블록 단위로 SRAM에 적재
    s = (q @ Kj.T) / sqrt(d)           # 이 블록의 점수
    m_new = max(m, s.max())
    scale = exp(m - m_new)
    p = exp(s - m_new)
    l = scale * l + p.sum()
    o = scale * o + p @ Vj
    m = m_new
out = o / l                            # 정확한 softmax(qK^T/sqrt(d)) V

역전파의 재계산

역전파에는 소프트맥스 확률 P가 필요하다. 표준 구현은 순전파 때 N×N 행렬 P를 저장해 두지만, 플래시 어텐션은 행마다 소프트맥스 정규화 통계(최댓값과 분모)만 저장해 두고 역전파 때 Q, K, V 블록에서 P를 SRAM 안에서 다시 계산한다. 계산량은 늘지만 HBM에서 큰 행렬을 읽는 것보다 빠르다.^[1] 그 결과 메모리 사용량이 시퀀스 길이에 대해 제곱이 아니라 선형이 된다.

IO 복잡도

논문은 SRAM 크기를 M( $d \leq M \leq N d$ )이라 할 때 플래시 어텐션의 HBM 접근량이 $Θ (N^{2} d^{2} M^{- 1})$ 임을 증명했다. 헤드 차원 d가 64~128, M이 약 100KB인 일반적인 경우 $d^{2}$ 가 M보다 훨씬 작으므로 표준 구현의 $Θ (N d + N^{2})$ 보다 접근 횟수가 여러 배 적다. 또한 일정 범위의 SRAM 크기에 대해 어떤 정확한 어텐션 알고리즘도 이보다 점근적으로 적게 HBM에 접근할 수 없다는 하한도 보였다.^[1]

버전별 발전

버전(발표)	대상 GPU	주요 기법	보고된 성능
플래시 어텐션(2022년 5월)^[1]	A100 등	타일링, 온라인 소프트맥스, 커널 융합, 역전파 재계산, 블록 희소 확장	GPT-2 어텐션 연산 최대 7.6배 빠름. BERT-large 학습이 MLPerf 1.1 기록보다 15% 빠름, GPT-2(길이 1K) 학습 3배. 메모리는 정확한 어텐션 기준선 대비 최대 20배 효율
플래시 어텐션 2(2023년 7월)^[4]	A100	행렬 곱이 아닌 연산 감소, 헤드 하나도 시퀀스 길이 방향으로 여러 스레드 블록에 병렬화, 워프 사이 작업 분할 개선	버전 1 대비 약 2배. A100 이론 최대치의 50~73%. GPT 계열 학습에서 GPU당 225TFLOPs/s(모델 FLOPs 활용률 72%)
플래시 어텐션 3(2024년 7월)^[5]	H100(호퍼)	텐서 코어와 TMA의 비동기 실행, 워프 특화로 계산과 데이터 이동 겹치기, 행렬 곱과 소프트맥스 교차 실행, FP8 블록 양자화	FP16에서 버전 2 대비 1.5~2.0배, 최대 740TFLOPs/s(활용률 75%). FP8에서 1.2PFLOPs/s에 근접, 기준 FP8 어텐션보다 수치 오차 2.6배 작음
플래시 어텐션 4(2026년 3월)^[6]	B200·GB200(블랙웰)	완전 비동기 행렬 곱 파이프라인과 큰 타일, 지수 함수를 소프트웨어로 근사해 특수 함수 유닛 병목 회피, 조건부 소프트맥스 재조정, 텐서 메모리와 2-CTA 행렬 곱 모드로 역전파의 공유 메모리 트래픽 감소	B200 BF16에서 cuDNN 9.13 대비 최대 1.3배, Triton 대비 2.7배, 최대 1,613TFLOPs/s(활용률 71%). 파이썬 내장 CuTe-DSL로 구현해 컴파일 시간 20~30배 단축

버전 2 논문은 버전 1이 A100 이론 최대 FLOPs/s의 25~40%밖에 쓰지 못하는 원인을 스레드 블록과 워프 사이의 비효율적인 작업 분할로 보았다.^[4] 버전 3 논문은 버전 2가 H100에서 활용률 35%에 그친다고 지적하고 호퍼 세대의 비동기 기능을 활용했다.^[5] 버전 4 논문은 블랙웰 세대에서 텐서 코어 처리량은 두 배가 되었지만 공유 메모리 대역폭과 지수 연산 유닛은 그만큼 늘지 않은 "비대칭 하드웨어 확장"이 새 병목이 되었다고 분석했다.^[6] 하드웨어가 바뀔 때마다 병목이 옮겨 가므로 알고리즘도 그에 맞춰 다시 설계해야 한다는 것이 이 계열 연구의 공통된 교훈이다.

효과와 영향

긴 컨텍스트: 메모리가 길이에 선형이 되면서 같은 GPU에서 훨씬 긴 시퀀스로 학습할 수 있게 되었다. 버전 1 논문은 GPT-2 퍼플렉서티 0.7 개선, 긴 문서 분류 6.4점 향상을 보고했고, 길이 16K의 Path-X 과제에서 처음으로 우연 수준을 넘는 트랜스포머(정확도 61.4%), 블록 희소 버전으로 길이 64K의 Path-256(63.1%)을 달성했다.^[1]
근사 어텐션의 재평가: 희소·저랭크 근사 방법은 FLOPs를 줄여도 메모리 접근 비용을 무시해 실제 속도 향상이 없는 경우가 많았다. 플래시 어텐션은 정확한 어텐션을 더 빠르게 만들어, 품질을 희생하는 근사보다 구현 최적화가 실용적이라는 흐름을 만들었다.^[1]
프레임워크 통합: 파이토치의 scaled_dot_product_attention은 입력 조건에 따라 플래시 어텐션 2, 메모리 효율 어텐션, 일반 C++ 구현 중 하나를 자동으로 고른다.^[2]
IO 인식 설계의 확산: 같은 원리가 다른 구조로 퍼졌다. 선택적 상태 공간 모델 Mamba는 확장된 상태를 SRAM 안에서만 만드는 스캔 커널을 썼고, 선형 어텐션의 하드웨어 효율 구현인 FlashLinearAttention은 짧은 시퀀스에서도 플래시 어텐션 2보다 빠르다고 보고했다.^[7]

한계

연산량 자체는 여전히 $O (N^{2} d)$ 다. 시퀀스가 아주 길면 계산 시간이 제곱으로 늘어나는 것은 변하지 않으며, 이는 상태 공간 모델이나 희소 어텐션 같은 구조적 해법의 영역이다.
GPU 아키텍처에 깊게 맞춘 커널이라 새 하드웨어가 나올 때마다 다시 작성해야 하고, 지원하는 헤드 차원, 데이터 형식, 마스크 종류에 제약이 있다.
블록 순서에 따라 부동소수점 합산 순서가 달라지므로, 표준 구현과 비트 단위까지 같은 결과가 나오지는 않는다(수학적으로는 같다).

같이 보기

각주

↑ ^1.00 ^1.01 ^1.02 ^1.03 ^1.04 ^1.05 ^1.06 ^1.07 ^1.08 ^1.09 Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré, 〈FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness〉, NeurIPS 2022. arXiv:2205.14135
↑ ^2.0 ^2.1 PyTorch, 〈torch.nn.functional.scaled_dot_product_attention〉 문서. docs.pytorch.org
↑ ^3.0 ^3.1 Maxim Milakov, Natalia Gimelshein, 〈Online normalizer calculation for softmax〉, 2018. arXiv:1805.02867
↑ ^4.0 ^4.1 Tri Dao, 〈FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning〉, ICLR 2024. arXiv:2307.08691
↑ ^5.0 ^5.1 Jay Shah 외, 〈FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision〉, NeurIPS 2024. arXiv:2407.08608
↑ ^6.0 ^6.1 Ted Zadouri 외, 〈FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling〉, 2026. arXiv:2603.05451
↑ Songlin Yang 외, 〈Gated Linear Attention Transformers with Hardware-Efficient Training〉, ICML 2024. arXiv:2312.06635
↑ Markus N. Rabe, Charles Staats, 〈Self-attention Does Not Need O(n²) Memory〉, 2021. arXiv:2112.05682
↑ Rewon Child 외, 〈Generating Long Sequences with Sparse Transformers〉, 2019. arXiv:1904.10509
↑ Woosuk Kwon 외, 〈Efficient Memory Management for Large Language Model Serving with PagedAttention〉, SOSP 2023. arXiv:2309.06180
↑ Ke Hong 외, 〈FlashDecoding++: Faster Large Language Model Inference on GPUs〉, MLSys 2024. arXiv:2311.01282

참고 문헌

[fa1-1] 1.00 ^1.01 ^1.02 ^1.03 ^1.04 ^1.05 ^1.06 ^1.07 ^1.08 ^1.09 Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré, 〈FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness〉, NeurIPS 2022. arXiv:2205.14135

[sdpa-2] 2.0 ^2.1 PyTorch, 〈torch.nn.functional.scaled_dot_product_attention〉 문서. docs.pytorch.org

[online-3] 3.0 ^3.1 Maxim Milakov, Natalia Gimelshein, 〈Online normalizer calculation for softmax〉, 2018. arXiv:1805.02867

[fa2-4] 4.0 ^4.1 Tri Dao, 〈FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning〉, ICLR 2024. arXiv:2307.08691

[fa3-5] 5.0 ^5.1 Jay Shah 외, 〈FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision〉, NeurIPS 2024. arXiv:2407.08608

[fa4-6] 6.0 ^6.1 Ted Zadouri 외, 〈FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling〉, 2026. arXiv:2603.05451

[gla-7] Songlin Yang 외, 〈Gated Linear Attention Transformers with Hardware-Efficient Training〉, ICML 2024. arXiv:2312.06635

[rabe-8] Markus N. Rabe, Charles Staats, 〈Self-attention Does Not Need O(n²) Memory〉, 2021. arXiv:2112.05682

[sparse-9] Rewon Child 외, 〈Generating Long Sequences with Sparse Transformers〉, 2019. arXiv:1904.10509

[paged-10] Woosuk Kwon 외, 〈Efficient Memory Management for Large Language Model Serving with PagedAttention〉, SOSP 2023. arXiv:2309.06180

[fdpp-11] Ke Hong 외, 〈FlashDecoding++: Faster Large Language Model Inference on GPUs〉, MLSys 2024. arXiv:2311.01282

[1]

[2]

[3]

[4]

[5]

[6]

[7]

[8]

[9]

[10]

[11]

기법	관계
온라인 소프트맥스(2018)^[3]	플래시 어텐션의 블록별 소프트맥스 갱신의 토대. 소프트맥스 단독 최대 1.3배, Top-K와 융합 시 최대 5배 가속을 보고
메모리 효율 어텐션(라베·스타츠, 2021)^[8]	같은 시기 어텐션 메모리를 $O (\sqrt{n})$ 로 줄이는 구현을 제시. 길이 16,384에서 추론 메모리 59배, 미분 시 32배 절감. 다만 실행 시간은 표준 구현과 비슷
Sparse Transformer의 재계산(2019)	어텐션 행렬을 저장하지 않고 역전파 때 다시 계산해 메모리를 아끼는 아이디어를 먼저 사용^[9]
PagedAttention(2023)^[10]	추론 서빙에서 KV 캐시를 운영체제의 페이지처럼 나누어 관리. 플래시 어텐션이 계산 쪽 최적화라면 이쪽은 KV 캐시 메모리 관리 최적화로 vLLM의 기반
FlashDecoding++(2023)^[11]	토큰을 하나씩 생성하는 디코딩 단계에 맞춰, 부분 소프트맥스들의 동기화 비용을 통일된 최댓값으로 없애는 등 추론 특화 최적화
MQA·GQA	키·값 헤드를 공유해 디코딩 때 읽어야 할 KV 캐시 자체를 줄임. 플래시 어텐션과 함께 쓰임