FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling

이 논문은 비대칭 하드웨어 확장 특성을 보이는 블랙웰 (Blackwell) 아키텍처의 병목 현상을 해결하기 위해 알고리즘과 커널 파이프라인을 공동 설계한 'FlashAttention-4'를 제안하며, CuTe-DSL 기반의 효율적인 구현을 통해 B200 GPU 에서 cuDNN 대비 최대 1.3 배, Triton 대비 2.7 배의 성능 향상을 달성함을 보여줍니다.

핵심

FlashAttention-4: 더 큰 뇌를 위한 '초고속' 학습법

이 논문은 인공지능 (AI) 이 더 길고 복잡한 내용을 이해하도록 돕기 위해, 최신 그래픽 카드 (NVIDIA Blackwell B200) 에 최적화된 새로운 기술인 FlashAttention-4를 소개합니다.

기존의 AI 모델은 긴 글을 읽거나 복잡한 코드를 분석할 때 '병목 현상'에 걸려 느려졌습니다. FlashAttention-4 는 이 문제를 해결하기 위해 하드웨어의 특성을 완벽하게 이해하고, 소프트웨어를 재설계한 혁신적인 방법입니다.

이 복잡한 기술을 일상적인 비유로 쉽게 설명해 드리겠습니다.

---

1. 문제 상황: "슈퍼카 엔진은 달렸는데, 도로가 좁아!"

과거의 AI 칩 (Hopper H100) 은 마치 슈퍼카 엔진처럼 계산을 엄청나게 빠르게 처리했습니다. 하지만 AI 가 긴 문서를 읽을 때 필요한 '기억 공간 (공유 메모리)'과 '특수 계산기 (지수 함수)'는 엔진만큼 빨라지지 않았습니다.

• 비유: 엔진은 제트기처럼 빨라졌는데, 차를 움직이는 **바퀴 (메모리)**와 **조향 장치 (특수 계산기)**는 여전히 자전거 수준입니다.
• 결과: 엔진이 아무리 빨라도 바퀴가 따라주지 못하면 차는 제자리걸음입니다. 최신 칩 (Blackwell B200) 은 엔진 성능이 2 배가 되었지만, 바퀴와 조향 장치는 그대로라 오히려 병목 현상이 더 심해졌습니다.

2. FlashAttention-4 의 해결책: "도로와 차를 함께 재설계하다"

FlashAttention-4 는 단순히 엔진을 더 빠르게 만드는 게 아니라, 도로 (하드웨어) 와 운전법 (소프트웨어) 을 함께 재설계했습니다.

① "동시 작업의 마법" (파이프라이닝)

기존 방식은 한 가지 일을 끝내고 다음 일을 시작했습니다. 하지만 FlashAttention-4 는 요리사처럼 행동합니다.

• 비유: 한 요리사가 국을 끓이는 동안 (계산), 다른 요리사가 채소를 썰고 (메모리 작업), 또 다른 사람이 접시를 준비합니다 (소프트웨어).
• 효과: 모든 작업이 겹쳐서 동시에 일어나므로, 기다리는 시간이 거의 사라집니다.

② "수학 대신 추측하기" (소프트웨어 시뮬레이션)

AI 가 '소프트맥스 (Softmax)'라는 계산을 할 때, 가장 느린 '지수 함수' 계산이 걸림돌이었습니다.

• 비유: 정확한 값을 구하는 데 10 분 걸리는 복잡한 수학 공식 대신, **경험칙 (다항식 근사)**을 써서 1 초 만에 "거의 정확한" 값을 추측해 내는 것입니다.
• 효과: AI 에게는 100% 정확한 값보다 '충분히 정확한' 값이 훨씬 빠릅니다. 이 방법으로 계산 속도를 획기적으로 높였습니다.

③ "작업장을 넓히고, 쓰레기를 줄이다" (메모리 최적화)

기존 방식은 계산 중간 결과를 계속 메모리에 저장하고 꺼내야 해서 시간이 걸렸습니다.

• 비유: 작업대 (공유 메모리) 가 좁아 물건을 옮기느라 바빴다면, 이제는 **작업대 위에 바로 보관할 수 있는 특수 선반 (텐서 메모리)**을 설치했습니다.
• 효과: 물건을 옮기는 횟수가 줄어든 덕분에, 계산하는 시간이 훨씬 늘어났습니다.

④ "두 팀이 한 팀이 되어 일하기" (2-CTA 모드)

최신 칩은 두 개의 작업 팀 (CTA) 이 협력하면 더 큰 일을 할 수 있게 해줍니다.

• 비유: 한 사람이 무거운 상자를 들면 힘들지만, 두 사람이 상자를 반반씩 나누어 들면 훨씬 가볍고 빠르게 이동할 수 있습니다.
• 효과: 메모리 이동량을 반으로 줄이고, 불필요한 대기 시간을 없앴습니다.

---

3. 왜 이것이 중요한가요? (성과)

이 기술을 적용한 FlashAttention-4 는 다음과 같은 놀라운 결과를 보여줍니다.

• 속도: 기존 최고 성능의 상용 소프트웨어 (cuDNN) 보다 최대 1.3 배, 다른 오픈소스 도구 (Triton) 보다 최대 2.7 배 빠릅니다.
• 효율: 칩이 이론적으로 낼 수 있는 성능의 **71%**까지 끌어올렸습니다. (기존에는 50% 정도만 활용됨)
• 개발자 친화적: 예전에는 복잡한 C++ 코드를 짜야 했지만, 이제는 파이썬으로 쉽게 작성할 수 있어 개발 속도가 20~30 배 빨라졌습니다.

4. 결론: AI 의 미래를 여는 열쇠

FlashAttention-4 는 단순히 "더 빠른 AI"를 만드는 것이 아닙니다. 하드웨어가 비대칭적으로 발전하는 (계산은 빨라지는데 메모리는 느린) 시대에, 알고리즘과 하드웨어가 서로 맞춰 춤추는 새로운 패러다임을 제시합니다.

이 기술 덕분에 AI 는 이제 수만 페이지의 책이나 긴 영상, 복잡한 코드 전체를 한 번에 읽고 이해할 수 있게 되었습니다. 마치 단숨에 도서관 전체를 훑어보는 슈퍼 독서광이 된 것과 같습니다.

이 기술은 오픈소스로 공개되어 전 세계 연구자와 개발자들이 더 빠르고 똑똑한 AI 를 만들 수 있는 길을 터주었습니다.

기술 요약

1. 문제 정의 (Problem)

• 비대칭 하드웨어 확장 (Asymmetric Hardware Scaling): NVIDIA Blackwell 아키텍처 (B200, GB200) 는 Hopper (H100) 에 비해 텐서 코어 (Matrix Multiply-Accumulate, MMA) 처리량이 2 배 증가했습니다 (FP16/BF16 기준 2.25 PFLOPS). 그러나 공유 메모리 (Shared Memory) 대역폭, 지수 함수 (Exponential) 유닛, 기타 기능 유닛의 성능은 동일하거나 느리게 확장되었습니다.
• 이동한 병목 현상 (Shifting Bottlenecks): 이전 버전 (FlashAttention-3) 이 Hopper 에서 MMA 연산에 최적화되었다면, Blackwell 에서는 MMA 연산이 너무 빨라져 공유 메모리 트래픽과 소프트맥스 (Softmax) 계산에 필요한 지수 연산이 새로운 병목 현상으로 작용합니다. 로프라인 (Roofline) 분석 결과, 이러한 비-MMA 자원의 사용 시간이 MMA 연산 시간을 25~60% 초과하는 것으로 나타났습니다.
• 기존 알고리즘의 한계: 기존 FlashAttention-3 을 Blackwell 에 단순히 이식하는 것은 새로운 아키텍처의 비동기적 특성과 더 큰 타일 크기를 활용하지 못해 성능을 제대로 끌어내지 못합니다.

2. 방법론 (Methodology)

FlashAttention-4 는 하드웨어의 비대칭적 특성을 해결하기 위해 알고리즘과 커널 구현을 공동 설계 (Co-design) 했습니다.

2.1. 전방향 (Forward Pass) 최적화

1. 재설계된 비동기 파이프라인:
   • Blackwell 의 완전 비동기 MMA 연산과 **더 큰 타일 크기 (128x128)**를 활용하여 텐서 코어, 소프트맥스 계산, 메모리 연산 간의 오버랩을 극대화합니다.
   • MMA 연산이 텐서 메모리 (TMEM) 에 결과를 직접 기록하므로, 레지스터 압력을 줄이고 더 큰 타일을 처리할 수 있습니다.
2. 지수 함수 (Exponential) 병목 완화:
   • 하드웨어 MUFU (Multi-Function Unit) 의 지수 연산 속도가 느리기 때문에, 소프트웨어로 시뮬레이션된 지수 함수를 도입했습니다.
   • FMA (Fused Multiply-Add) 유닛을 사용하여 다항식 근사 (Polynomial Approximation) 로 $2^x$를 계산하여 처리량을 높였습니다.
   • 정확도 손실을 최소화하기 위해 일부 행만 소프트웨어로 계산하고 나머지는 하드웨어 명령어를 사용하는 부분 시뮬레이션 (Partial Emulation) 전략을 채택했습니다.
3. 조건부 소프트맥스 리스케일링 (Conditional Rescaling):
   • 온라인 소프트맥스 알고리즘에서 새로운 블록이 이전 블록보다 더 큰 값을 가지지 않는 경우 (즉, $m_j - m_{j-1} \le \tau$), 불필요한 리스케일링 연산을 생략합니다. 이는 비-MMA 연산을 크게 줄여줍니다.

2.2. 역방향 (Backward Pass) 최적화

1. 텐서 메모리 (TMEM) 활용:
   • 중간 결과 (S, P, dS 등) 를 공유 메모리가 아닌 텐서 메모리에 저장하여 공유 메모리 트래픽을 줄였습니다.
2. 2-CTA MMA 모드 활용:
   • Blackwell 의 2-CTA (Cooperative Thread Array) 텐서 코어 모드를 활용합니다. 두 개의 CTA 가 하나의 큰 MMA 연산을 협력하여 수행하며, 피연산자 B 를 두 CTA 가 각각 절반씩 공유 메모리에 로드합니다.
   • 이를 통해 공유 메모리 트래픽을 약 50% 줄이고, dQ (Query Gradient) 계산 시 글로벌 원자적 덧셈 (Global Atomic Adds) 횟수를 절반으로 감소시킵니다.
3. 확정적 실행 모드 (Deterministic Execution):
   • 학습의 재현성을 위해 원자적 연산의 순서를 제어하는 잠금 (Lock) 메커니즘을 도입하되, CTA 스와즐링 (Swizzling) 과 SPT (Shortest-Processing-Time-First) 스케줄링을 통해 성능 저하를 최소화했습니다.

2.3. 구현 프레임워크 (CuTe-DSL)

• FlashAttention-4 는 전통적인 C++ 템플릿 메타프로그래밍 대신 Python 에 내장된 CuTe-DSL로 완전히 구현되었습니다.
• 이 접근 방식은 저수준 GPU 프로그래밍의 모든 표현력을 유지하면서, 컴파일 시간을 20~30 배 단축하여 개발 생산성을 극대화했습니다.

3. 주요 기여 (Key Contributions)

1. 비대칭 하드웨어 대응 알고리즘: 텐서 코어 속도가 다른 유닛보다 훨씬 빠른 Blackwell 아키텍처의 병목 현상을 해결하기 위해 지수 연산 시뮬레이션과 조건부 리스케일링을 도입했습니다.
2. 2-CTA 및 TMEM 활용: 공유 메모리 대역폭 병목을 해결하기 위해 2-CTA MMA 모드와 텐서 메모리를 적극 활용하여 데이터 이동량을 획기적으로 줄였습니다.
3. CuTe-DSL 기반 개발: 복잡한 커널 개발의 진입 장벽을 낮추고 빠른 반복 개발을 가능하게 하는 새로운 프로그래밍 모델을 제시했습니다.
4. 확정적 역방향 연산: 성능 저하를 최소화하면서 재현 가능한 학습을 지원하는 역방향 커널을 구현했습니다.

4. 실험 결과 (Results)

• 하드웨어: NVIDIA B200 GPU (BF16 정밀도).
• 성능 향상:
  • cuDNN 9.13 대비 최대 1.3 배 속도 향상.
  • Triton 구현체 대비 최대 2.7 배 속도 향상.
  • 최대 1613 TFLOPs/s 달성 (이론적 최대 성능의 약 71% 활용).
• 컴파일 시간: FlashAttention-3 (C++ 템플릿) 대비 22~32 배 빠른 컴파일 시간 (Forward: 55s → 2.5s, Backward: 45s → 1.4s).
• 시퀀스 길이: 중간~긴 시퀀스 (4k 이상) 에서 cuDNN 및 Triton 을 일관되게 능가했습니다.

5. 의의 및 결론 (Significance)

FlashAttention-4 는 단순히 하드웨어 성능을 끌어올리는 것을 넘어, 하드웨어의 비대칭적 진화 (Asymmetric Scaling) 에 맞춘 알고리즘적 재설계의 중요성을 보여줍니다.

• 하드웨어-소프트웨어 공동 설계: 하드웨어가 제공하는 비동기 연산과 새로운 메모리 계층 (TMEM) 을 알고리즘 단계에서 직접 반영하여 병목 현상을 해결했습니다.
• 생태계 확장: CuTe-DSL 을 통해 연구자들이 C++ 템플릿 메타프로그래밍의 복잡함 없이도 고성능 어텐션 변형을 빠르게 프로토타이핑하고 배포할 수 있는 기반을 마련했습니다.
• 미래 지향성: 컴퓨팅 성능이 비-MMA 유닛을 빠르게 추월하는 추세가 지속될 것으로 예상되므로, FlashAttention-4 의 접근 방식은 차세대 가속기 아키텍처에서도 중요한 패러다임이 될 것입니다.

이 논문은 대규모 언어 모델 (LLM) 과 긴 컨텍스트 애플리케이션의 성능 한계를 극복하기 위한 핵심 기술로서, Blackwell 기반 시스템의 효율성을 극대화하는 데 기여합니다.