FlashAttention Analysis

FlashAttention 심층 분석

IO-aware Exact Attention · Tiling · Online Softmax · Recomputation · FA-2/FA-3 · Flash-Decoding

개요

FlashAttention은 어텐션을 근사하지 않으면서(exact) 더 빠르고 메모리 효율적으로 만드는 IO-aware 알고리즘입니다. 핵심 통찰은 어텐션이 연산(FLOPs)이 아니라 메모리 접근(HBM 읽기/쓰기)에 병목이 있다는 점입니다.

표준 어텐션은 거대한 N×N 행렬을 HBM에 만들고 여러 번 읽고 씁니다. FlashAttention은 Q·K·V를 SRAM에 맞는 블록으로 쪼개(tiling) online softmax로 점진적으로 계산해 N×N을 아예 만들지 않고, 역전파에서는 통계만 저장해 재계산(recomputation)합니다. 결과는 정확하면서 메모리 O(N)·속도 2~4배입니다. 본 문서는 문제(메모리 바운드) → GPU 메모리 계층 → 해법①(tiling+online softmax) → 해법②(recomputation)와 결과 → 진화(FA-2/FA-3·Flash-Decoding)와 위치의 순서로 분석합니다.

핵심 개념

• IO-aware: HBM 읽기/쓰기를 줄여 전체 속도를 끌어올리는 설계입니다.
• Tiling: Q·K·V를 SRAM에 맞는 블록으로 나누어 블록 단위로 계산합니다.
• Online softmax: running max와 running sum으로 블록 결과를 점진적으로 합칩니다.
• Recomputation: 역전파에서 필요한 값은 통계만 저장하고 다시 계산합니다.
• FA-2/FA-3: 병렬화와 Hopper 비동기·FP8으로 FlashAttention을 확장합니다.

비교/분석

동작 원리

FlashAttention의 동작 원리는 표준 어텐션이 어떤 지점에서 HBM 왕복을 반복하는지 파악한 뒤, 같은 수학적 결과를 유지한 채 데이터를 SRAM 안에서 최대한 오래 머무르게 만드는 데 있습니다. 아래 순서대로 보면 왜 tiling, online softmax, recomputation이 함께 필요해지는지 자연스럽게 연결됩니다.

1. 문제 — 표준 어텐션은 메모리 바운드(IO 병목)

표준 어텐션은 S=QK^T(N×N) → P=softmax(S)(N×N) → O=PV로 계산합니다. 여기서 N은 시퀀스 길이이며, 거대한 N×N 행렬 S·P를 HBM에 실체화(materialize)하고 여러 패스로 읽고 씁니다. 시간·메모리가 N²에 비례합니다.

그림 1. 표준 어텐션의 N×N 실체화와 메모리 바운드 특성

핵심 — 어텐션은 연산이 아니라 HBM 접근이 병목

• 어텐션의 진짜 병목은 FLOPs가 아니라 HBM 접근(데이터 이동)입니다. A100의 연산 처리량은 312 TFLOPS(FP16)로 매우 빠른 반면 HBM 대역폭은 ~2TB/s에 그쳐, 연산 유닛이 데이터를 기다리는 memory-bound 상태가 됩니다.

• 그래서 근사 어텐션(Linformer·Performer·sparse)이 FLOPs를 줄여도 실제 wall-clock 속도는 잘 빨라지지 않습니다 — 그들은 IO(메모리 읽고 쓰기)를 고려하지 않았기 때문입니다. 이것이 FlashAttention의 핵심 지적입니다.

• 해결의 열쇠는 'IO-aware' — N×N을 HBM에 만들지 않고 메모리 이동을 최소화하는 것입니다.

2. GPU 메모리 계층 — IO-awareness의 무대

FlashAttention을 이해하려면 GPU의 가파른 메모리 계층을 봐야 합니다. 작고 매우 빠른 on-chip SRAM과 크고 상대적으로 느린 HBM의 차이가 모든 것의 출발점입니다.

그림 2. SRAM↔HBM 메모리 계층과 IO 복잡도

메모리 계층과 IO 복잡도

• SRAM은 수십 MB에 불과하지만 대역폭이 HBM의 약 10배입니다(연산 유닛 바로 옆). HBM은 수십 GB로 크지만 상대적으로 느립니다.

• 표준 어텐션의 비효율은 N×N 행렬을 HBM에 만들고 scaling·masking·softmax·dropout 단계마다 다시 읽고 쓰는 데 있습니다. IO-aware 설계의 목표는 데이터를 SRAM에 올려 계산을 끝내고 HBM 왕복을 최소화하는 것입니다.

• FlashAttention의 IO 복잡도는 O(N²d²/M) HBM 접근(M=SRAM 크기, d=헤드 차원)으로, 표준 어텐션의 Ω(Nd+N²)보다 훨씬 적습니다. SRAM이 클수록 유리하며, 정확 어텐션에서 이보다 더 줄일 수 없다는 이론적 하한도 증명했습니다. 즉 '연산을 줄이는 것'이 아니라 '데이터 이동을 줄이는 것'이 핵심입니다.

3. 해법 ① — Tiling + Online Softmax

FlashAttention은 Q·K·V를 SRAM에 맞는 블록으로 쪼개(tiling) 블록 단위로 계산하고, softmax를 점진적으로(online) 계산해 N×N 행렬을 아예 만들지 않습니다.

그림 3. Tiling과 online softmax(running max/sum rescaling)

Online Softmax 누적 (개념)

Tiling + Online Softmax의 핵심

• Tiling — Q·K·V를 블록으로 쪼개 블록 쌍을 SRAM에 올려 계산하고 출력만 누적합니다. 거대한 N×N 행렬 S·P를 HBM에 전혀 만들지 않습니다(materialization 회피).

• 난점과 해결 — softmax는 행 전체의 max·합이 필요해 블록 단위로 쪼개기 어렵습니다(비결합적). Online softmax는 블록마다 running max(m)와 running sum(l)을 갱신하고, 새 블록의 max가 더 크면 이전까지 누적한 출력 O를 지수 인자로 rescale해 보정합니다. softmax를 '결합 가능'하게 만든 것이 FlashAttention의 핵심 수학입니다.

• 전부 한 GPU 커널로 융합(fuse)합니다 — scaling·masking·softmax·dropout·PV를 한 패스에 처리합니다. 결과는 표준 어텐션과 정확히 동일(exact)하며 근사가 아닙니다. 단지 IO 패턴만 바꿉니다.

# Q를 블록 단위로 고정하고, K·V 블록을 순회하며 출력을 누적
m = -inf
l = 0
O = 0
for j in range(num_kv_blocks):
    # K_j, V_j 를 SRAM에 로드
    S_ij = Q_i @ K_j.T
    m_new = max(m, rowmax(S_ij))
    P_ij = exp(S_ij - m_new)
    l = exp(m - m_new) * l + rowsum(P_ij)
    O = exp(m - m_new) * O + P_ij @ V_j
    m = m_new

O = O / l

4. 해법 ② — Recomputation, 그리고 결과

그림 4. Recomputation(통계만 저장해 역전파에서 재계산)과 결과

Recomputation

• 역전파에는 보통 순전파의 N×N 어텐션 행렬이 필요합니다 — 저장하면 O(N²) 메모리입니다. FlashAttention은 N×N을 저장하지 않고 softmax 정규화 통계(running max·sum)만 저장합니다.

• 역전파 때 그 통계로 S·P를 SRAM에서 즉석 재계산(recompute)하여 메모리를 선형으로 만듭니다. 핵심 통찰은 recompute로 FLOPs는 늘지만, memory-bound라 HBM 접근 감소가 더 커서 '순속도'가 오히려 빨라진다는 것입니다.

결과 (FlashAttention-1)

또한 정확 어텐션에서 HBM 접근을 점근적으로 더 줄일 수 없다는 IO 복잡도 하한을 증명했고(최적), block-sparse로 확장하면 IO 인식 + 희소 패턴으로 기존 근사 어텐션보다도 빠른 근사 버전이 됩니다.

장단점

관련 기술

그림 5. 세대별 진화, Flash-Decoding, 그리고 FlashAttention의 위치

세대별 진화

FlashAttention-2는 병렬화와 work 분할을 개선해 A100에서 약 70%의 이론적 최대 FLOPS를 달성합니다. FlashAttention-3는 Hopper(H100)의 비동기 기능(warp-specialization으로 연산·데이터 이동 중첩, TMA)과 FP8 저정밀(incoherent processing)을 활용해, FP16에서 740 TFLOPS(75% 활용)·FP8에서 약 1.2 PFLOPS에 이릅니다(FA-2는 H100에서 35%만 활용).

Flash-Decoding — 추론(디코딩)용 변형

추론의 디코딩 단계는 질의가 1개인데 KV는 수천 개라 memory-bound이고, 질의 길이로 병렬화하는 FlashAttention만으로는 작은 배치에서 병렬성이 부족합니다. Flash-Decoding은 KV 길이 축으로 분할해 부분 어텐션을 병렬 계산하고, log-sum-exp로 합쳐(reduction) 작은 배치에서도 GPU를 가득 채웁니다. xFormers·vLLM 등에 통합되어 있으며, FlashDecoding++는 디코딩에서 최대 2.02배 향상을 보고합니다.

내부 문서

• PagedAttention Analysis
• KV Cache Quantization Analysis
• KV Cache Offloading Analysis
• Continuous Batching Analysis
• LLM Inference Scheduler Analysis
• Memory Centric LLM Serving Survey

원문

성능 요약

세대별 비교

핵심 정리

FlashAttention은 IO-aware exact attention으로, N×N을 HBM에 만들지 않고(tiling + online softmax) 역전파는 통계로 재계산(recompute)한다. 어텐션이 memory-bound라는 통찰에서 출발해, 연산을 줄이는 대신 HBM 접근을 줄여 정확하면서 메모리 O(N)·속도 2~4배를 달성한다.

FA-1(원리)→FA-2(병렬화·work 분할)→FA-3(Hopper 비동기·FP8)로 하드웨어에 맞춰 진화했고, 추론 디코딩용 Flash-Decoding은 KV 길이 축으로 분할 병렬화한다. 중요한 점은 FlashAttention이 KV 캐시를 '줄이지' 않는다는 것 — 어텐션 계산의 IO 패턴을 최적화할 뿐이다.

따라서 KV 절감(양자화·MLA·오프로딩)과 직교적으로 함께 쓰이며, PagedAttention/vLLM도 내부적으로 Flash 커널을 사용한다. 메모리 시스템 관점에서는 SRAM↔HBM 계층에서 데이터를 어디에 두고 언제 다시 계산할지에 대한 일반적인 IO 트레이드오프를 잘 보여주는 사례다.

주의 — 본문 수치는 원논문·기술 자료(FlashAttention arXiv 2205.14135·NeurIPS 2022, FA-2 arXiv 2307.08691, FA-3 arXiv 2407.08608, Flash-Decoding/FlashDecoding++ 자료)의 보고값이다. 'GPT-2 7.6배', 'BERT 15%', '메모리 10~20배', 'A100 312 TFLOPS·~2TB/s', '740 TFLOPS·1.2 PFLOPS' 같은 값은 특정 모델·하드웨어·시퀀스 길이·정밀도 조건의 보고값으로 일반화에 주의가 필요하다. SRAM·HBM 용량/대역폭 수치는 GPU 세대에 따라 다르며, 속도 향상은 시퀀스가 길수록(더 memory-bound일수록) 커지는 경향이 있다.