FlashAttention-2 (KO)

From Systems analysis wiki
Jump to navigation Jump to search

FlashAttention-2 — 대형 언어 모델(LLM)에서 attention 메커니즘을 계산하기 위해 설계된 개선된 알고리즘입니다. 이 알고리즘은 Tri Dao와 스탠퍼드 대학교 연구진이 개발하여 2023년 7월에 발표되었습니다[1]. 핵심 목표는 GPU 하드웨어 자원을 보다 효율적으로 활용함으로써 transformer 모델의 학습 및 추론(inference) 속도를 크게 높이는 것이며, 동시에 표준 attention 메커니즘과 완전히 동일한 계산 결과를 유지합니다. 즉, 정확도 손실이 없습니다.

FlashAttention-2는 같은 팀이 2022년에 발표한 FlashAttention 알고리즘의 논리적 후속작입니다. 새 버전은 전작에서 관찰된 GPU 부하 미충족 문제를 해결하고, 1세대 대비 약 두 배의 속도 향상을 달성합니다.

배경: transformer에서의 attention 문제

표준 자기 주의(self-attention) 메커니즘은 transformer에서 긴 텍스트 시퀀스를 처리할 때 병목이 됩니다. 시퀀스 길이(N)에 따라 계산 복잡도와 메모리 사용량이 이차적으로(O(N²)) 증가하여 최대 컨텍스트 길이와 LLM의 확장성에 심각한 제약을 가합니다[1].

이 문제를 해결하기 위해 2022년에 FlashAttention 알고리즘이 발표되었습니다[2]. 핵심 아이디어는 다음과 같습니다:

  • GPU 메모리 계층 구조 고려(IO-awareness): 알고리즘은 GPU의 느린 메모리(HBM)와 칩 내의 빠른 정적 메모리(SRAM) 사이에서 발생하는 비용이 큰 읽기/쓰기 연산을 최소화합니다.
  • 블록 처리(tiling): 계산을 작은 블록(타일)으로 분할하여 빠른 SRAM에서 처리함으로써 전체 attention 행렬을 메모리에 구체화하는 것을 피합니다.

이를 통해 메모리 사용량의 선형 증가(O(N))와 표준 구현 대비 2–4배의 속도 향상을 달성했습니다[2]. FlashAttention은 널리 보급되었으며, 예를 들어 2천–4천 토큰(GPT-3)에서 128천 토큰(GPT-4) 이상으로 컨텍스트가 크게 확장된 모델의 등장에 기여했습니다[3]. Falcon-40B 모델에서는 FlashAttention을 사용함으로써 추론 속도가 3배, 전반적인 생성 성능이 GPT-3 대비 5배 향상되었습니다[4].

FlashAttention-2의 개발과 목표

성공에도 불구하고, FlashAttention 1세대는 GPU 연산 자원을 완전히 활용하지 못했습니다. NVIDIA A100 GPU에서 성능은 이론적 최대치(FLOPs/s)의 25–40%에 불과했습니다[1]. 주된 원인은 Streaming Multiprocessors의 비최적 부하와 공유 메모리의 과잉 연산이었습니다[5].

FlashAttention-2의 목표는 작업 병렬화를 보다 효율적으로 수행하고 보조 연산을 최소화하여 계산을 더욱 가속화하는 것이었습니다. 알고리즘은 최대 성능을 달성하기 위해 NVIDIA CUTLASS 3.x 라이브러리의 저수준 기본 요소를 사용하여 완전히 새로 작성되었습니다[6].

기술 아키텍처 및 동작 원리

FlashAttention-2는 병렬성과 효율성을 높이기 위한 세 가지 핵심 개선 사항을 도입합니다[1]:

1. 비행렬 연산 최소화

알고리즘은 행렬 곱셈이 아닌 보조 부동소수점 연산(non-matmul FLOPs)의 수를 줄입니다. GPU의 텐서 코어는 행렬 연산(GEMM)에 최적화되어 있으며 최대 16배 빠르게 실행하므로, 이 변경을 통해 가장 성능이 높은 GPU 블록을 더 많은 시간 동안 활용할 수 있습니다.

2. 향상된 병렬성

원래의 FlashAttention에서는 단일 attention 헤드에 대한 작업이 병렬화되지 않아, 긴 시퀀스와 작은 배치 크기에서 유휴 상태가 발생했습니다. FlashAttention-2는 블록 간 병렬화를 도입하여, 단일 attention 헤드에 대한 계산을 서로 다른 GPU Streaming Multiprocessors에 분산함으로써 부하율을 크게 향상시킵니다.

3. 블록 내 작업 분할 최적화

단일 연산 블록 수준에서 작업이 스레드 그룹(warp) 간에 재분배되어 공유 메모리(shared memory)를 통한 데이터 교환이 줄어들었습니다. 이를 통해 Softmax 정규화에 필요한 불필요한 읽기/쓰기 연산 수가 감소합니다.

성능과 효율성

아키텍처 개선 덕분에 FlashAttention-2는 상당한 성능 향상을 보여줍니다:

  • 두 배 속도 향상: 알고리즘은 FlashAttention 1세대에 비해 약 2배 빠르게 동작합니다[1].
  • 높은 GPU 활용률: NVIDIA A100 GPU에서 이론적 최대 처리량(TFLOPs)의 50–73%를 달성하며, 이는 최적화된 행렬 곱셈(GEMM) 연산의 효율에 근접합니다[1].
  • 기록적인 계산 속도:
    • A100 GPU에서 GPT 유형 모델의 end-to-end 학습 루프에서 최대 225 TFLOP/s의 속도를 달성하며, 이는 연산 블록 활용률 72%에 해당합니다. 비교하자면, 동일한 조건에서 표준 attention은 GPU에 100 TFLOP/s 미만의 부하만 걸었습니다[7].
    • H100 GPU에서는 성능이 335 TFLOP/s에 달합니다[7].

이러한 성능 향상 덕분에, 예를 들어 컨텍스트 윈도우 8k 토큰에 필요했던 시간과 동일한 시간 내에 16k 토큰 컨텍스트 윈도우를 가진 모델을 학습할 수 있습니다[5]. 중요한 점은 알고리즘이 정확하고 결정론적이므로, 이를 적용해도 모델의 예측 품질에 영향을 주지 않는다는 것입니다[8].

적용 및 생태계 통합

FlashAttention-2는 LLM 생태계에서 빠르게 표준 도구로 자리 잡았습니다. 다양한 인기 프레임워크 및 라이브러리에 통합되어 있습니다:

  • PyTorch: 네이티브 지원.
  • Hugging Face Transformers: 모델 로드 시 `attn_implementation="flash_attention_2"` 매개변수로 지원을 활성화합니다[9]. 수십 가지 아키텍처(GPT, Llama, Falcon, BERT 등)와 호환됩니다[10].
  • TensorRT-LLM, xFormers 및 Triton: 알고리즘이 이러한 플랫폼을 위해 구현되어 광범위한 적용이 가능합니다[7].

통합을 통해 FlashAttention-2를 양자화(GPTQ, QLoRA) 및 효율적인 파인튜닝(PEFT)과 같은 다른 최적화 방법과 쉽게 결합할 수 있습니다[9].

이후 버전과의 비교

FlashAttention-3

attention 최적화 연구는 계속되고 있습니다. 2024년 7월 Tri Dao는 NVIDIA Hopper 아키텍처(H100/H200) GPU의 기능 활용을 목표로 하는 FlashAttention-3를 발표했습니다. 주요 새 기능은 다음과 같습니다[3]:

  • FP8 지원: 추가 가속을 위해 8비트 부동소수점 계산을 사용합니다.
  • 비동기 연산: GPU의 비동기 기능을 보다 효율적으로 활용합니다.

FlashAttention-3는 H100 GPU에서 FlashAttention-2 대비 1.5–2배의 가속을 제공하며, 최대 740 TFLOP/s의 성능(이론적 최대치의 75%)을 달성합니다[11].

참고 문헌

  • Dao, T. (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. arXiv:2307.08691.
  • Dao, T. et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. arXiv:2205.14135.
  • 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.
  • Ye, Z. et al. (2025). FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving. arXiv:2501.01005.
  • Chen, Y. et al. (2023). FlashDecoding++: Faster Large Language Model Inference on GPUs. arXiv:2311.01282.
  • Liu, Y. et al. (2024). FastAttention: Extending FlashAttention-2 to NPUs and Low-Resource GPUs. OpenReview: 76NYyOrnfk.
  • 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. (2024). FlashMask: Efficient and Rich Mask Extension of FlashAttention. OpenReview: rog0J435OO.
  • Abbott, V.; Zardini, G. (2025). FlashAttention on a Napkin: A Diagrammatic Approach to Deep Learning IO-Awareness. arXiv:2412.03317.

주석

  1. ↑ 1.0 1.1 1.2 1.3 1.4 1.5 Дао, Три. «FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning». arXiv:2307.08691 [cs.LG], 17 июля 2023 г. [1]
  2. ↑ 2.0 2.1 «Optimizing LLMs for Speed and Memory». Hugging Face Documentation. [2]
  3. ↑ 3.0 3.1 Дао, Три. «FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision». Tri Dao's Blog. [3]
  4. ↑ «FlashAttention vs FlashAttention-2 - an Analysis». E2E Networks Blog. [4]
  5. ↑ 5.0 5.1 «FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning». OpenReview. [5]
  6. ↑ «FlashAttention-2». Hazy Research, Stanford University. [6]
  7. ↑ 7.0 7.1 7.2 Дао, Три. «FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning» (PDF). arXiv:2307.08691. [7]
  8. ↑ Рашка, Себастьян. «Llama 2 and FlashAttention 2». Ahead of AI Magazine. [8]
  9. ↑ 9.0 9.1 Белькада, Юнес. «Faster and more memory efficient models with Flash Attention 2!». LinkedIn. [9]
  10. ↑ «GPU inference». Hugging Face Documentation. [10]
  11. ↑ Дао, Три, и др. «FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision». arXiv:2407.08608 [cs.LG], 11 июля 2024 г. [11]
Retrieved from "https://systems-analysis.ru/int/index.php?title=FlashAttention-2_(KO)&oldid=22323"