논문: FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness 저자: Dao et al. (Stanford) 링크: https://arxiv.org/abs/2205.14135

논문 핵심

FlashAttention은 Self-Attention의 계산 결과는 정확히 동일하면서, GPU 메모리 계층(HBM vs SRAM)을 최적화하여 2-4배 빠르고 메모리 효율적인 구현을 달성합니다.

핵심 아이디어 — IO-Awareness

"We propose FlashAttention, an IO-aware exact attention algorithm."

기존 방식의 문제: Attention 행렬(N×N)을 HBM에 쓰고 읽는 IO가 병목 FlashAttention: 타일링(Tiling)으로 SRAM에서만 연산, HBM 접근 최소화

성능 향상

  • BERT-large 학습: 15% 속도 향상
  • GPT-2 학습: 3배 속도 향상
  • 시퀀스 길이 4K → 16K로 확장 가능 (같은 메모리)

채택 현황

현재 사실상 모든 주요 LLM 프레임워크에서 표준:

  • PyTorch 2.0 내장
  • Hugging Face Transformers 기본 지원
  • vLLM, TensorRT-LLM 등에서 필수 사용

실무 시사점

FlashAttention은 "알고리즘 개선이 하드웨어 업그레이드보다 효과적일 수 있다"는 것을 보여줍니다. LLM 추론 서버를 운영하는 기업에서는 FlashAttention 적용만으로도 비용을 크게 절감할 수 있습니다.