FlashAttention-2 — GPU attention의 사실상 표준

원제: FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning

Tri Dao가 제안한 FlashAttention-2는 GPU SRAM에 tile 단위로 attention을 올려 HBM I/O를 최소화. FlashAttention v1 대비 2x, A100/H100에서 표준 attention 대비 4-9x 빠름. PyTorch native 통합.

출처: NeurIPS 2023 · arXiv:2311.01831 원문 ↗

1. 한줄 요약

Attention 계산을 GPU SRAM에 tile 단위로 올려 메모리 I/O를 최소화한 FlashAttention의 후속. A100/H100에서 2-9x faster, 모든 트랜스포머의 사실상 표준.

2. 왜 등장했는가 (Background)

트랜스포머의 attention은 O(N²) 메모리 + I/O. 2023년 기준 context length가 100K+ 토큰으로 늘어나면서 attention이 training/inference의 병목.

Standard attention: QKV 행렬을 HBM(고용량 메모리)에서 매번 읽고 씀 → 메모리 bandwidth 한계
GPU의 SRAM은 작지만(20MB) 매우 빠름 (≈19TB/s), HBM은 크지만 느림 (≈2TB/s)
FlashAttention v1 (2022): SRAM + tiling 개념 입증, 2-4x 가속

v1의 한계: GPU warp/wave의 활용이 비효율적, 비선형 패턴. v2는 work partitioning 재설계로 추가 2x.

3. 핵심 아이디어 — Tiling + Recomputation

두 가지 트릭의 합:

1. Tiling (tile-wise attention):
   Q, K, V를 작은 block으로 분할 → SRAM에 올림
   softmax(QK^T/√d)V를 block 단위로 계산
   HBM ↔ SRAM 왕복 최소화

2. Recomputation:
   backward pass에서 attention score 재계산
   (HBM에 저장 안 함 → 메모리 ↓, 계산 ↑, 전체 속도 ↑)

3. Work partitioning (v2 핵심):
   Q block을 GPU warp에 균등 분배
   K/V block은 함께 순회 → warp 간 통신 최소화

비유: 1,000쪽 책을 매번 HBM(창고)에서 책상(SRAM)으로 가져와 한 줄씩 읽는 대신, 책을 10쪽씩 끊어 책상에 두고 연속 작업. 창고 왕복이 없어지면 작업 자체는 빨라짐.

4. Architecture

Tiling 전략: Q는 row-tile, K/V는 column-tile. SRAM 사용률 90%+
Softmax in SRAM: 통계량(max, sum)을 누적하면서 online softmax
Work partitioning: query block을 4개 warp에 균등 분배, key block은 cooperative 순회
Backward: QKV gradient를 recompute (저장 X), 메모리 footprint 절반

5. Training & Data

하드웨어: A100 (80GB, 108 SMs), H100 (SXM5)
모델: GPT-3 style, 1.3B → 70B까지 검증
Context: 512 → 8K → 64K까지 확장
비교 대상: standard attention, FlashAttention v1, xformers memory-efficient attention

6. Result

모든 트랜스포머의 표준 attention:

Forward speed: A100 standard 대비 4.5x, H100 대비 9x
Backward speed: A100 standard 대비 3x
메모리: O(N²) → O(N) (sequence length linear)
정확도: standard attention과 수치적으로 동일 (masking, dropout 모두 지원)
긴 컨텍스트: 64K 컨텍스트에서 메모리 10x 절감

7. 한국 독자 적용 사례

모든 LLM 학습/추론: Llama 3, Qwen3, Claude 3.5, GPT-4o 모두 FlashAttention 사용
긴 컨텍스트 서비스: RAG·문서 요약에서 100K+ 토큰 처리에 필수
한국어 모델: KoLLaMA·KoSOLAR 등 한국어 모델 학습/추론 기본
PyTorch native: torch.nn.functional.scaled_dot_product_attention 한 줄로 사용

8. 한계 + 후속 영향

한계:

NVIDIA GPU에 최적화 (AMD/Intel GPU는 미지원, 별도 kernel 필요)
FlashAttention-3 (2024) — H100의 WGMMA + async 적극 활용해 추가 2x
Multi-head attention 외 sparse/multi-query attention은 별도 구현 필요

후속 영향:

PyTorch native scaled_dot_product_attention의 기본 구현
긴 컨텍스트(100K-1M 토큰) 시대의 토대
모든 LLM 프레임워크(HF, vLLM, TGI)의 사실상 표준