Flux perfermence improved by Flash attention3 + Triton

FlashAttention-3: 대규모 언어 모델을 위한 어텐션 메커니즘의 혁신

1. 서론

대규모 언어 모델(LLM)의 발전과 함께, 트랜스포머 구조의 핵심 요소인 어텐션 메커니즘의 효율성 개선이 중요한 과제로 대두되고 있습니다. 특히 긴 문맥(long-context)을 다루는 모델에서 어텐션 연산은 주요 병목 현상의 원인이 되어왔습니다. 이러한 맥락에서 FlashAttention 알고리즘의 발전, 특히 최근 발표된 FlashAttention-3는 주목할 만한 혁신을 가져왔습니다.

하지만 들어가기에 앞서 현재 FlashAttention3는 Hopper 기반(H100)에서만 사용이 가능하고 ADA 기반(RTX4090, L4) GPU는 사용을 할 수가 없음으로 추후에 업데이트를 기다려야 합니다.

2. FlashAttention의 원리

FlashAttention은 메모리 효율성과 연산 속도를 크게 개선한 어텐션 알고리즘입니다. 주요 특징은 다음과 같습니다:

1. 메모리 재정렬과 타일링: GPU의 주 메모리(HBM)에서 더 빠른 SRAM으로 데이터를 효율적으로 이동시켜 처리합니다.
2. 중간 결과 최소화: 대규모 어텐션 행렬을 HBM에 저장하지 않아 메모리 I/O를 줄입니다.
3. 블록 단위 처리: 타일링과 소프트맥스 재스케일링을 통해 정확도를 유지하면서 연산을 가속화합니다.

이러한 접근 방식으로 FlashAttention은 기존 방법 대비 2-4배의 속도 향상을 달성했습니다.

3. FlashAttention-3: H100 GPU에 최적화된 혁신

FlashAttention-3는 NVIDIA의 Hopper 아키텍처, 특히 H100 GPU의 새로운 기능들을 최대한 활용하여 개발되었습니다. 주요 특징과 개선 사항은 다음과 같습니다:

3.1 성능 향상

• FP16 모드에서 최대 740 TFLOPS (H100 이론적 최대의 75%)
• FP8 모드에서 1.2 PFLOPS에 근접한 성능
• FlashAttention-2 대비 1.5-2배 빠른 처리 속도

3.2 H100 GPU의 새로운 기능 활용

1. WGMMA(Warpgroup Matrix Multiply-Accumulate): 새로운 Tensor Core를 사용하여 높은 처리량 제공
2. TMA(Tensor Memory Accelerator): 글로벌 메모리와 공유 메모리 간 데이터 전송 가속화
3. 저정밀도 FP8: Tensor Core 처리량을 두 배로 증가시키는 대신 정확도 일부 희생

3.3 주요 개선 사항

• GPU 활용도 향상: H100 GPU 최대 용량의 75%까지 사용
• 저정밀도 성능 개선: FP8 연산에서도 높은 정확도 유지
• 긴 문맥 처리 능력 향상: LLM에서 더 긴 텍스트를 효율적으로 처리 가능

4. 핵심 기술: GEMM과 Softmax 연산의 최적화

FlashAttention-3의 주요 혁신 중 하나는 GEMM(General Matrix Multiply)과 Softmax 연산의 효율적인 처리입니다.

4.1 비동기적 처리 및 중첩

• GEMM 연산은 GPU에서 매우 빠르지만, Softmax와 같은 특수 함수는 상대적으로 느림
• 두 연산을 병렬로 처리하여 전체적인 효율성 향상

4.2 최적화 전략

1. 워프스케줄러 활용: 자동으로 워프그룹 간 중첩 수행
2. 핑퐁 스케줄링: 서로 다른 워프그룹에서 GEMM 연산을 번갈아가며 수행
3. 워프그룹 내 병렬 처리: 단일 워프그룹에서 GEMM과 Softmax를 동시에 처리

4.3 양자화 오류 감소

• 이상치 값(outlier) 처리를 위한 '통일성 없는 처리' 기법 도입
• 하다마드 변환을 활용한 이상치 분산 방식 적용

5. 벤치마크 결과

FlashAttention-3는 다양한 조건에서 우수한 성능을 보여주었습니다:

1. FP16 연산: FlashAttention-2 대비 약 1.6-1.8배 속도 개선
2. FP8 연산: 약 1.2 PFLOPS에 근접한 처리 속도 달성
3. 타 구현체 비교: cuDNN, Triton 등 Hopper GPU 최적화 라이브러리보다 우수한 성능

6. 결론 및 전망

FlashAttention-3는 대규모 언어 모델의 학습 및 추론 과정에서 중요한 돌파구를 제공합니다. GPU 하드웨어의 최신 기능을 최대한 활용하고, 연산 과정을 세밀하게 최적화함으로써 현재 트랜스포머 구조의 주요 병목이었던 어텐션 메커니즘의 효율성을 크게 향상시켰습니다.

이러한 발전은 더 큰 규모의 언어 모델 개발, 더 긴 문맥을 다룰 수 있는 모델의 등장, 그리고 실시간 응용에서의 LLM 활용 가능성을 높여줄 것으로 기대됩니다. 향후 새로운 GPU 아키텍처의 등장과 함께 FlashAttention의 추가적인 최적화 및 개선이 이루어질 것으로 예상되며, 이는 인공지능 분야의 지속적인 발전에 중요한 역할을 할 것입니다.