핵심 요약
기존 FlashAttention 커널은 LLM의 긴 시퀀스에 최적화되어 있어, 짧고 불규칙한 시퀀스가 많은 실제 추천 시스템(RecSys) 환경에서는 성능이 저하되는 문제가 있다. Meta는 이를 해결하기 위해 소프트맥스를 GELU, SiLU 등 범용 활성화 함수로 대체한 GDPA(Generalized Dot-Product Attention) 커널을 새롭게 설계했다. 새로운 커널은 소프트웨어 파이프라이닝 재설계, 불규칙한 텐서를 위한 지그재그 타일 스케줄링, SFU 병목을 줄이는 테일러 급수 근사 기법 등을 도입했다. 그 결과 NVIDIA B200 GPU에서 기존 Triton 기반 구현 대비 포워드 패스에서 최대 2배, 백워드 패스에서 1.6배의 속도 향상을 기록하며 전체 모델 학습 처리량을 30% 이상 개선했다.
배경
Transformer 아키텍처 및 Attention 메커니즘, GPU 아키텍처(SM, Warp, SFU)에 대한 이해, CUDA 커널 프로그래밍 및 Flash Attention 개념
대상 독자
GPU 커널 개발자, 대규모 추천 시스템 학습 인프라 엔지니어, ML 성능 최적화 연구자
의미 / 영향
이 연구는 LLM 중심의 최적화 기법이 추천 시스템과 같은 다른 도메인에 그대로 적용되지 않음을 보여주며, 실제 데이터 분포에 기반한 맞춤형 커널 설계의 중요성을 강조한다. 특히 Blackwell 아키텍처(B200)의 성능을 한계까지 끌어올리는 구체적인 방법론을 제시했다는 점에서 업계에 큰 의미를 갖는다.
섹션별 상세
def gelu_taylor_approximation(x):
# ALU-only Taylor expansion up to x^6
# GELU(x) ≈ 1/2*x + 1/sqrt(2*pi)*x^2 - 1/(6*sqrt(2*pi))*x^4 + 1/(40*sqrt(2*pi))*x^6
c1 = 0.5
c2 = 0.3989422804 # 1/sqrt(2*pi)
c4 = -0.0664903801 # 1/(6*sqrt(2*pi))
c6 = 0.0099735570 # 1/(40*sqrt(2*pi))
x2 = x * x
x4 = x2 * x2
x6 = x4 * x2
return c1 * x + c2 * x2 + c4 * x4 + c6 * x6SFU 병목을 피하기 위해 ALU 연산만으로 구성한 GELU의 6차 테일러 급수 근사 구현 예시
실무 Takeaway
- 실제 서비스 데이터의 불규칙한 특성(Jagged Tensors)을 고려한 커널 설계가 벤치마크 성능보다 실제 학습 효율 개선에 더 결정적인 영향을 미친다.
- 짧은 시퀀스 비중이 높은 워크로드에서는 루프 플래닝(Loop Flattening)과 외부 루프 파이프라이닝을 통해 하드웨어 활용률을 극대화할 수 있다.
- SFU 자원이 부족한 최신 GPU 아키텍처에서는 복잡한 활성화 함수를 ALU 기반 테일러 급수로 근사하여 연산 병목을 효과적으로 해결할 수 있다.
AI 요약 · 북마크 · 개인 피드 설정 — 무료