FlashAttention (KO)

From Systems analysis Wiki
Jump to navigation Jump to search

FlashAttention — 대형 언어 모델(LLM)의 학습 및 추론을 완전한 계산 정확도를 유지하면서 크게 가속화하기 위해 개발된 혁신적인 attention 메커니즘 계산 알고리즘입니다. 이 알고리즘은 2022년 스탠퍼드 대학교의 연구팀이 트리 다오(Tri Dao)의 지도하에 처음 발표했습니다[1].

FlashAttention의 핵심 아이디어는 GPU 메모리 계층 구조를 고려하여 계산을 재구성하는 것으로, 느린 메모리에 대한 접근 횟수를 최소화하고 표준 attention 메커니즘의 주요 병목 현상을 제거합니다.

표준 Attention의 문제점

트랜스포머에서 표준 자기 주의(self-attention) 메커니즘은 다음 공식으로 계산됩니다: Attention(Q,K,V)=softmax(QKTdk)V 여기서 Q, K, V는 각각 쿼리, 키, 값 행렬입니다.

이 접근 방식의 주요 문제는 시퀀스 길이 N에 대해 시간 및 메모리 복잡도가 이차(O(N²))라는 점입니다[1]. 단순한 구현에서는 N×N 크기의 전체 attention 행렬 S를 GPU 메모리에 계산하고 저장해야 하며, 이로 인해 두 가지 심각한 문제가 발생합니다:

  1. 높은 메모리 소비: 긴 컨텍스트를 다룰 때 N×N 행렬을 저장하는 것이 불가능해집니다.
  2. 입출력(IO) 연산: 주요 병목 현상은 산술 연산의 수가 아니라, 느린 GPU 메모리에 대한 지속적인 접근입니다.

GPU 메모리 계층 구조

문제를 이해하려면 GPU의 두 가지 메모리 유형을 구별하는 것이 중요합니다(NVIDIA A100 기준):

  • SRAM (정적 메모리): 소용량(~20MB)이지만 매우 빠른 온칩 메모리로, 엄청난 대역폭(최대 19TB/s)을 제공합니다.
  • HBM (고대역폭 메모리): 대용량(40~80GB)이지만 상대적으로 느린 메모리로, 대역폭이 훨씬 낮습니다(약 1.5TB/s)[2].

이러한 비대칭성 때문에 표준 attention 알고리즘은 메모리 대역폭에 의해 제한(memory-bound)되며, 느린 HBM에서 대형 행렬을 지속적으로 읽고 쓰는 것이 주요 지연 원인이 됩니다.

FlashAttention의 핵심 혁신

FlashAttention은 HBM에 대한 접근을 최소화하여 문제를 해결하는 IO-인식(IO-aware) 알고리즘입니다. 이는 세 가지 주요 기법을 통해 달성됩니다.

타일링 및 블록 처리

전체 행렬을 한꺼번에 처리하는 대신, FlashAttention은 입력 행렬 Q, K, V를 빠른 SRAM에 맞는 작은 블록(타일)으로 분할합니다. 알고리즘은 이 블록들을 순차적으로 로드하고, 각 블록에 대해 모든 attention 계산을 수행한 후, 느린 HBM에 전체 attention 행렬을 저장하지 않고 최종 결과를 업데이트합니다[1].

온라인 Softmax 계산

핵심적인 기술적 돌파구는 Softmax의 "온라인" 계산입니다. 표준 Softmax는 정규화를 위해 입력 벡터의 모든 요소를 알아야 합니다. FlashAttention은 Softmax를 부분적으로 계산할 수 있는 수정된 알고리즘을 사용합니다. 새로운 블록이 처리될 때마다 업데이트되는 두 개의 중간값(현재 최댓값과 지수의 합)을 유지함으로써, 전체 행렬에 한 번에 접근하지 않고도 정확한 결과를 얻을 수 있습니다[2].

단일 CUDA 커널로의 연산 융합

모든 attention 연산(행렬 곱 QKᵀ, 마스킹, Softmax, V와의 곱)이 단일 융합 CUDA 커널(fused kernel)로 통합됩니다. 이는 HBM에 대한 읽기/쓰기 연산 수를 대폭 줄입니다. 전체 행렬을 여러 번 순회하는 대신, 알고리즘은 블록을 SRAM에 한 번 로드하고 모든 계산을 수행한 후 최종 결과만 저장합니다.

이론적 및 실용적 효율성

복잡도 및 최적성

FlashAttention은 메모리 소비를 O(N²)에서 O(N)으로 줄여 선형 확장성을 제공합니다. 이 알고리즘의 IO 복잡도는 2단계 메모리 계층 구조에서 attention 계산을 위해 이론적으로 최적임이 증명되었습니다. 즉, 하드웨어를 변경하지 않고는 정확한 attention을 더 빠르게 실행할 수 없습니다[3].

실증적 결과

FlashAttention의 첫 번째 버전은 다음과 같은 상당한 개선을 보여주었습니다:

  • 속도 향상:
    • BERT-large (길이 512): 학습 속도 15% 향상.
    • GPT-2 (길이 1K): 3배 속도 향상.
    • Long-Range Arena 작업 (1K-4K): 2.4배 속도 향상[1].
  • 메모리 절약: 정확한 기준 구현 대비 최대 20배의 메모리 절약.
  • 모델 품질 향상: 더 긴 컨텍스트로 작업할 수 있게 됨으로써, FlashAttention은 모델 품질을 저하시키지 않을 뿐만 아니라 오히려 향상시킵니다. 예를 들어, GPT-2의 퍼플렉시티가 0.7포인트 향상되었고, 긴 문서 분류 작업의 정확도가 6.4포인트 향상되었습니다[1].

발전 및 후속 연구

FlashAttention의 성공은 하드웨어 지향 알고리즘의 새로운 흐름을 만들어 냈습니다.

FlashAttention-2 (2023)

두 번째 버전은 GPU 리소스를 보다 완전하게 활용하는 것을 목표로 했습니다. 원래의 FlashAttention은 NVIDIA A100에서 최대 성능의 25~40%만 달성했습니다. FlashAttention-2는 계산 병렬화의 개선을 도입하여 다음을 가능하게 했습니다[4]:

  • 첫 번째 버전 대비 2배의 속도 향상 달성.
  • GPU 활용도를 이론적 최댓값의 50~73%까지 증가.
  • 크기 256의 attention 헤드 및 Multi-Query Attention (MQA) 아키텍처에 대한 지원 확장.

FlashAttention-3 (2024)

세 번째 버전은 NVIDIA Hopper (H100) GPU 아키텍처에 특화하여 최적화되었습니다[5]. Tensor Core 비동기성 및 FP8 지원과 같은 새로운 하드웨어 기능을 활용하여 다음을 달성했습니다:

  • FlashAttention-2 대비 추가적인 1.5~2배 속도 향상.
  • FP16에서 최대 740 TFLOPS, FP8에서 1.2 PFLOPS에 근접하는 성능.

특화된 솔루션

FlashAttention의 아이디어는 다른 프로젝트에서도 발전되었습니다:

  • FlashInfer (2025): LLM 추론 작업에 특화하여 최적화된 맞춤형 attention 엔진입니다. 스트리밍 생성 모드에서 KV 캐시를 효율적으로 처리하는 데 중점을 둡니다[6].
  • FlashMLA (2024): 컨텍스트 캐시를 압축하는 attention(latent attention) 구현으로, 최소한의 정보 손실로 매우 긴 시퀀스에서 메모리를 절약할 수 있게 해줍니다[7].

산업 및 생태계에 미친 영향

FlashAttention은 근본적인 돌파구가 되었으며, LLM의 효율적인 학습 및 추론을 위한 산업 표준으로 빠르게 자리잡았습니다. PyTorch와 Hugging Face와 같은 주요 라이브러리에 통합되었으며, LLaMA, MPT, Falcon, Claude 등 대부분의 주요 언어 모델에서 사용됩니다.

FlashAttention과 그 후속 버전들은 언어 모델의 컨텍스트 윈도우를 확장하는 데 결정적인 역할을 했습니다. 2~4천 토큰(GPT-3)에서 128천 토큰(GPT-4)으로, 심지어 실험적 모델에서는 수백만 토큰까지 확장되었습니다[8]. 이 알고리즘은 트랜스포머 확장의 주요 장애물 중 하나를 제거하여, 긴 문서 분석부터 멀티모달 이해에 이르기까지 AI 응용 프로그램의 새로운 가능성을 열었습니다.

외부 링크

  • GitHub의 공식 FlashAttention 저장소

참고 문헌

  • Dao, T. et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. arXiv:2205.14135.
  • Dao, T. (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. arXiv:2307.08691.
  • Shah, J. et al. (2024). FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-Precision. arXiv:2407.08608.
  • Kwon, W. et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention. arXiv:2309.06180.
  • Hong, K. et al. (2023). FlashDecoding++: Faster Large Language Model Inference on GPUs. arXiv:2311.01282.
  • Ye, Z. et al. (2025). FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving. arXiv:2501.01005.
  • Dege, P. et al. (2025). FlashMLA-ETAP: Efficient Transpose Attention Pipeline for Accelerating MLA Inference on NVIDIA H20 GPUs. arXiv:2506.01969.
  • Wang, G. et al. (2025). FlashMask: Efficient and Rich Mask Extension of FlashAttention. OpenReview wUtXB43Chi.
  • Dao, T. et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (OpenReview version). OpenReview H4DqfPSibmx.
  • Gholami, A. et al. (2024). FlashAttention on a Napkin: A Diagrammatic Approach to Deep Learning IO-Awareness. OpenReview pF2ukh7HxA.

주석

  1. ↑ 1.0 1.1 1.2 1.3 1.4 Дао, Три, и др. «FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness». arXiv:2205.14135 [cs.LG], 28 мая 2022 г. [1]
  2. ↑ 2.0 2.1 Дао, Три, и др. «FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness». OpenReview. [2]
  3. ↑ «We're Training AI Twice as Fast This Year as Last». IEEE Spectrum. [3]
  4. ↑ Дао, Три. «FlashAttention-2». tridao.me. [4]
  5. ↑ «FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision». PyTorch Blog. [5]
  6. ↑ «[2501.01005] FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving». arXiv. [6]
  7. ↑ «GitHub - deepseek-ai/FlashMLA: FlashMLA: Efficient MLA decoding kernels». GitHub. [7]
  8. ↑ «The Evolution of Flash Attention: Revolutionizing Transformer Efficiency». Medium. [8]
Retrieved from "https://systems-analysis.info/int/index.php?title=FlashAttention_(KO)&oldid=2259"