FlashAttention-3 (KO)

FlashAttention-3 — 트랜스포머 신경망에서 어텐션(attention) 메커니즘을 최적화하기 위한 알고리즘으로, NVIDIA Hopper 아키텍처(H100) GPU의 하드웨어 기능을 최대한 활용하도록 개발되었습니다. 이 알고리즘은 2024년 Colfax Research, Meta, NVIDIA, Georgia Tech, 프린스턴 대학교, Together AI 소속 연구자들로 구성된 팀에 의해 발표되었습니다. 해당 논문은 NeurIPS 2024 컨퍼런스에 채택되었으며 spotlight 논문으로 선정되었습니다.

FlashAttention-3은 FlashAttention(2022)과 FlashAttention-2(2023)에 이은 알고리즘 시리즈의 세 번째 버전입니다. 주요 목표는 계산 정확도를 유지하면서 대형 언어 모델(LLM)의 학습 및 추론(inference)을 크게 가속화하는 것입니다.

도입 및 배경

어텐션 메커니즘의 문제점

트랜스포머의 핵심 구성 요소는 셀프 어텐션(self-attention) 메커니즘이지만, 이의 계산 복잡도와 메모리 소비는 입력 시퀀스 길이(n)가 증가함에 따라 이차적으로(O(n²)) 증가합니다. 현대 GPU는 빠른 행렬 곱셈에 최적화되어 있지만, 지수 함수(예: Softmax)의 계산은 수십 배 더 느리기 때문에 심각한 병목 현상이 발생합니다. 또한 단순한 구현 방식에서는 GPU 메모리에 대규모의 중간 어텐션 텐서를 저장해야 하므로 모델의 확장성이 제한됩니다.

FlashAttention 및 FlashAttention-2

이 문제를 해결하기 위해 2022년 FlashAttention이 제안되었으며, 두 가지 기법을 통해 느린 글로벌 메모리(HBM)에 대한 접근 횟수를 줄였습니다:

• 블록 처리(tiling): 계산을 블록(타일)으로 분할하여 빠른 온칩(on-chip) 메모리(SRAM)에서 처리합니다.
• 연산 융합: 모든 연산(행렬 곱셈, Softmax)을 중간 결과를 글로벌 메모리에 쓰지 않고 하나의 GPU 커널에서 수행합니다.

이를 통해 메모리 복잡도를 이차에서 선형으로 줄이고 계산 속도를 2–4배 향상시킬 수 있었습니다.

2023년에는 개선된 버전인 FlashAttention-2가 발표되어 계산 병렬화를 최적화했습니다. NVIDIA Ampere 아키텍처(A100) GPU에서 이론적 최대 성능의 약 70%를 달성했습니다. 그러나 더 새로운 NVIDIA Hopper 아키텍처(H100)에서는 효율이 훨씬 낮은 약 35% 수준에 머물렀습니다. 이는 알고리즘이 Hopper의 새로운 하드웨어 기능을 활용하지 못했기 때문이며, 이것이 FlashAttention-3 개발의 계기가 되었습니다.

GPU Hopper(H100)의 새로운 하드웨어 기능

NVIDIA Hopper 아키텍처는 FlashAttention-3이 최대 성능을 달성하기 위해 활용하는 여러 새로운 기능을 제공합니다:

• WGMMA (Warpgroup Matrix Multiply-Accumulate): 텐서 코어용 새로운 명령어 유형으로, Ampere 아키텍처 대비 거의 두 배의 성능 향상을 제공하는 행렬 곱셈을 수행합니다.
• TMA (Tensor Memory Accelerator): 글로벌 메모리(HBM)와 공유 메모리(shared memory) 간의 데이터 전송을 가속화하는 하드웨어 모듈입니다. TMA는 주소 계산을 자동으로 처리하여 연산 코어의 부담을 줄입니다.
• FP8 형식: 8비트 부동소수점 데이터 형식에 대한 하드웨어 지원으로, FP16 대비 이론적 성능을 두 배로 높이지만 제한된 동적 범위로 인해 정확도 손실의 위험이 있습니다.

FlashAttention-3의 기술적 혁신

이 알고리즘은 Hopper 아키텍처를 위해 특별히 설계된 세 가지 핵심 최적화 방법을 구현합니다:

1. 비동기 실행 및 워프 특화

FlashAttention-3은 warp-specialization 원칙을 활용하며, GPU의 서로 다른 스레드 그룹(warps)이 서로 다른 작업을 전담합니다:

• Producer warps: TMA를 사용하여 글로벌 메모리에서 데이터를 로드합니다.
• Consumer warps: 텐서 코어에서 행렬 곱셈을 수행합니다.

Hopper의 하드웨어 비동기 특성 덕분에 이 연산들은 시간적으로 중첩됩니다. 한 워프 그룹이 계산을 수행하는 동안, 다른 그룹은 병렬로 다음 블록의 데이터를 로드합니다. ping-pong scheduling 방식으로 구성된 이 파이프라인(pipeline) 접근 방식은 느린 연산(예: Softmax)의 지연을 숨기고 GPU의 모든 기능 모듈을 최대한 활용할 수 있게 합니다.

2. 메모리 연산 최소화

알고리즘은 이전 버전의 tiling 철학을 유지하되, TMA를 적극적으로 활용하여 현재 계산과 병렬로 다음 데이터 블록을 비동기적으로 로드합니다. 느린 HBM에서 빠른 SRAM으로의 데이터 전송은 사실상 주요 계산의 "그늘 속에서" 이루어지므로, GPU가 데이터를 기다리며 유휴 상태에 머무는 시간이 줄어듭니다.

3. 양자화 오류를 줄인 낮은 정밀도(FP8)

FP8으로의 전환은 속도를 두 배로 높이지만 양자화로 인해 심각한 정확도 손실을 초래할 수 있습니다. 이를 해결하기 위해 개발자들은 incoherent processing 기법을 도입했습니다. 그 내용은 다음과 같습니다:

1. 어텐션 계산 전에 특징 벡터(쿼리 Q와 키 K)를 무작위 직교 행렬(예: 아다마르 행렬)과 곱합니다.
2. 이 변환은 비정상적으로 큰 절댓값을 가진 값(이상값)을 모든 좌표에 "분산"시켜 분포를 균일하게 만듭니다.
3. 그 후 FP8로 양자화를 수행하면 오류가 더 적게 발생합니다.
4. 이 변환은 직교 변환이므로 행렬의 효과가 곱셈 시 상쇄되어 최종 어텐션 결과(QKᵀ)를 왜곡하지 않습니다.

이 기법을 통해 FP8 어텐션 계산 오류를 변환 없이 표준 FP8을 적용하는 것과 비교하여 약 2.6배 줄일 수 있었습니다.

성능 및 의의

앞서 언급한 기법들을 적용함으로써 FlashAttention-3은 H100 GPU에서 이전 버전 대비 상당한 우위를 달성했습니다:

• FlashAttention-2 대비 1.5–2배 속도 향상.
• 높은 GPU 활용률: H100 이론적 최대 성능의 약 75–85% 달성.
• 처리량:
  • 반정밀도(FP16/BF16)에서 최대 740–840 TFLOPS.
  • 8비트 정밀도(FP8) 사용 시 최대 1.2–1.3 PFLOPS(페타플롭스).

FlashAttention-3의 높은 효율성은 LLM의 개발 및 적용에 직접적인 영향을 미칩니다:

• 학습 시간 단축: 어텐션 속도가 75–100% 향상되면 수 주 또는 수 개월이 걸릴 수 있는 모델 학습 시간이 크게 단축됩니다.
• 컨텍스트 윈도우 확장: 모델이 더 긴 시퀀스(수십만 개의 token)를 효율적으로 처리할 수 있어 대규모 문서나 코드 분석에 유리합니다.
• 자원의 효율적 활용: 더 적은 수의 GPU로 동일한 성능을 달성하거나 동일한 하드웨어에서 더 높은 속도를 얻을 수 있어 모델 배포 비용이 절감됩니다.

가용성 및 통합

저자들은 FlashAttention-3의 소스 코드를 오픈 라이선스로 GitHub에 공개했습니다. PyTorch 및 Hugging Face Transformers 라이브러리와 같은 주요 딥러닝 프레임워크에 통합될 것으로 예상되며, 이를 통해 기술이 광범위한 개발자와 연구자들에게 접근 가능해질 것입니다. 이전 버전들은 이미 업계 사실상의 표준이 되었으며, FlashAttention-3도 이 추세를 이어갈 것으로 보입니다.

외부 링크

• GitHub의 FlashAttention 공식 저장소
• FlashAttention-3 발표를 담은 Together AI 블로그

참고 문헌

• Shah, J. et al. (2024). FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision. arXiv:2407.08608.
• 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.
• 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. arXiv:2410.01359.
• Abbott, V.; Zardini, G. (2025). FlashAttention on a Napkin: A Diagrammatic Approach to Deep Learning IO-Awareness. arXiv:2412.03317.