범용 내적 어텐션(GDPA): 실제 추천 시스템 학습 부하를 위한 커널 최적화

이 요약은 AI가 원문을 분석해 생성했습니다. 정확한 내용은 원문 기준으로 확인하세요.

핵심 요약

기존 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)의 성능을 한계까지 끌어올리는 구체적인 방법론을 제시했다는 점에서 업계에 큰 의미를 갖는다.

섹션별 상세

GDPA는 표준 어텐션의 소프트맥스를 GELU, SiLU 등 범용 활성화 함수로 대체하여 추천 시스템의 다양한 상호작용 패턴을 하나의 고성능 커널로 통합한다.

Self-attention, PMA, pFFN 구조가 GDPA로 통합되는 과정을 보여주는 다이어그램 — Diagram추천 시스템에서 사용되는 다양한 어텐션 변형들이 공통적으로 두 개의 행렬 곱셈과 중간 활성화 함수 패턴을 공유함을 시각화한다. 이를 통해 GDPA 커널 하나로 여러 모듈을 통합 최적화할 수 있는 근거를 제시한다.

실제 운영 환경의 데이터는 짧고 비대칭적인 시퀀스와 가변적인 길이를 가지며, 이는 기존 LLM 중심 커널 설계와 성능 격차를 발생시키는 주요 원인이다.

실제 운영 데이터와 벤치마크 데이터 간의 커널 성능 격차를 보여주는 차트 — Chart합성 데이터 기반 벤치마크와 달리 실제 데이터에서는 포워드 패스에서 2.6배, 백워드 패스에서 1.6배의 성능 저하가 발생함을 보여준다. 이는 기존 커널이 실제 데이터의 불규칙성을 제대로 처리하지 못하고 있음을 증명한다.

소프트맥스 보정 단계를 제거하고 워프(Warp) 구성을 최적화하여 레지스터 자원을 추가로 확보하고 파이프라인 효율을 높였다.

짧은 K/V 시퀀스에서 발생하는 오버헤드를 줄이기 위해 내부 루프가 아닌 외부 루프 수준에서 소프트웨어 파이프라이닝(SWP)을 적용하여 연산 중첩을 극대화했다.

루프 플래닝(Loop Flattening) 적용 전후의 워프 실행 타임라인 비교 — Diagram내부 루프 파이프라이닝에서 발생하던 버블(유휴 시간)이 외부 루프 수준의 SWP 적용 후 크게 줄어드는 것을 보여준다. 특히 짧은 K/V 시퀀스에서 연산 효율이 어떻게 개선되는지 시각적으로 설명한다.

불규칙한(Jagged) 입력을 처리하기 위해 CPU에서 유효 타일을 미리 계산하고 지그재그 패턴으로 SM에 할당하는 새로운 로드 밸런싱 알고리즘을 도입했다.

지그재그 타일 스케줄링을 통한 SM 부하 분산 과정 설명도 — Diagram불규칙한 입력으로 인해 발생한 SM 간의 작업량 불균형을 라운드 로빈과 지그재그 정렬 스케줄링을 통해 해결하는 과정을 보여준다. 최대/최소 부하 격차가 12:1에서 5:4 수준으로 좁혀지는 효과를 수치로 제시한다.

SFU(Special Function Unit) 병목을 해결하기 위해 tanh 기반 GELU 대신 ALU만 사용하는 6차 테일러 급수 근사법을 사용하여 연산 밀도를 높였다.

python

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 * x6

SFU 병목을 피하기 위해 ALU 연산만으로 구성한 GELU의 6차 테일러 급수 근사 구현 예시

표준 tanh 기반 GELU와 6차 테일러 급수 근사 모델의 입력 범위별 출력 비교 그래프 — Chart테일러 급수 근사가 특정 입력 범위 내에서 매우 정확함을 보여준다. 실제 모델에서는 RMSNorm을 통해 입력값이 이 범위 내로 제한되므로, 정확도 손실 없이 SFU 병목을 해결할 수 있음을 뒷받침한다.

실무 Takeaway

실제 서비스 데이터의 불규칙한 특성(Jagged Tensors)을 고려한 커널 설계가 벤치마크 성능보다 실제 학습 효율 개선에 더 결정적인 영향을 미친다.
짧은 시퀀스 비중이 높은 워크로드에서는 루프 플래닝(Loop Flattening)과 외부 루프 파이프라이닝을 통해 하드웨어 활용률을 극대화할 수 있다.
SFU 자원이 부족한 최신 GPU 아키텍처에서는 복잡한 활성화 함수를 ALU 기반 테일러 급수로 근사하여 연산 병목을 효과적으로 해결할 수 있다.

언급된 리소스

GitHubAds Model Kernel Library (GDPA)

논문Kunlun: Establishing Scaling Laws for Massive-Scale Recommendation Systems