GN⁺: FlashAttention-3: 비동기 및 저정밀도로 빠르고 정확한 Attention 기술

FlashAttention-3: 비동기 및 저정밀도로 빠르고 정확한 Attention

• Attention의 중요성

  • Attention은 Transformer 구조의 핵심 계층으로, 대형 언어 모델과 긴 문맥 응용 프로그램에서 병목 현상을 일으킴.
  • FlashAttention과 FlashAttention-2는 GPU에서 메모리 읽기/쓰기를 최소화하여 Attention을 가속화하는 접근 방식을 개척함.
  • 이로 인해 LLM의 문맥 길이가 크게 증가함.

• FlashAttention-3의 주요 기술

  • 비동기성 활용: Tensor Cores와 TMA의 비동기성을 활용하여 전체 계산과 데이터 이동을 겹침.
  • 블록 단위 연산: 블록 단위의 행렬 곱셈과 softmax 연산을 교차 수행.
  • 저정밀도 처리: FP8 저정밀도 지원을 활용하여 성능을 향상시킴.

• FlashAttention-3의 성능 향상

  • GPU 활용 효율성: H100 GPU의 최대 성능을 75%까지 활용하여 이전 버전보다 1.5-2배 빠름.
  • 저정밀도 성능: FP8을 사용하여 처리 속도를 높이고 메모리 사용량을 줄임.
  • 긴 문맥 처리: Attention 메커니즘을 가속화하여 더 긴 텍스트를 효율적으로 처리 가능.

• FlashAttention 요약

  • FlashAttention은 Attention 계산을 재정렬하고 타일링과 재계산을 활용하여 속도를 크게 높이고 메모리 사용량을 줄임.
  • 타일링을 통해 입력 블록을 로드하고, 해당 블록에 대해 Attention을 수행한 후 출력을 업데이트함.
  • 중간 Attention 행렬을 메모리에 쓰지 않음으로써 메모리 읽기/쓰기 양을 줄임.

• Hopper GPU의 새로운 하드웨어 기능

  • WGMMA: 새로운 Tensor Cores를 활용하여 높은 처리량을 제공.
  • TMA: 글로벌 메모리와 공유 메모리 간 데이터 전송을 가속화하는 하드웨어 유닛.
  • FP8 저정밀도: FP8을 사용하여 Tensor Core 처리량을 두 배로 늘림.

• 비동기성: GEMM과 Softmax 겹치기

  • 겹치기의 필요성: GEMM과 softmax를 병렬로 수행하여 성능을 극대화함.
  • 핑퐁 스케줄링: 두 워프 그룹이 번갈아 가며 GEMM과 softmax를 수행하여 성능을 향상시킴.
  • 워프 그룹 내 겹치기: 동일한 워프 그룹 내에서 GEMM과 softmax를 병렬로 수행하여 처리량을 증가시킴.

• 저정밀도: 비일관 처리로 양자화 오류 감소

  • 비일관 처리: Hadamard 변환을 사용하여 양자화 오류를 줄임.
  • 실험 결과: 비일관 처리를 통해 양자화 오류를 2.6배 감소시킴.

• Attention 벤치마크

  • FP16: FlashAttention-2보다 약 1.6-1.8배 빠름.
  • FP8: 최대 1.2 PFLOPS에 도달.

GN⁺의 정리

• FlashAttention-3는 GPU의 새로운 하드웨어 기능을 활용하여 Attention 메커니즘의 성능을 크게 향상시킴.
• 긴 문맥을 효율적으로 처리할 수 있어 대형 언어 모델의 성능을 극대화함.
• PyTorch와 같은 주요 프레임워크에 통합될 가능성이 높아 향후 AI 연구와 응용에 큰 영향을 미칠 것임.
• 유사한 기능을 제공하는 프로젝트로는 Triton과 cuDNN이 있음.