정규화 오버헤드 제거: 정규화를 GEMM과 어텐션 커널에 융합하기

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

TL;DR

정규화 연산이 LLM과 추천 시스템에서 널리 사용되지만 TensorCore 활용이 불가하고 메모리 바운드 특성으로 전체 학습 지연의 상당 비중을 차지하기에 본문은 정규화 커널을 GEMM과 어텐션 커널에 융합하는 기술을 통해 이 병목을 완화했다. Lazy Pre-Norm과 Multi-CTA Norm Fusion은 정규화 연산을 인접한 행렬곱과 결합해 중간 메모리 IO를 제거하거나 겹치게 하여 정규화 지연의 최대 90%를 숨기는 효과를 보였고 FlashNormAttention은 LayerNorm과 RMSNorm을 어텐션 내부로 융합해 특정 커널에서 최대 35%의 속도 향상을 기록했다. 구현과 자동튜닝은 TLX와 Helion DSL을 기반으로 이루어졌으며 벤치마크는 bfloat16 데이터 타입과 NVIDIA B200 GPU, 750W 전력 제한 환경에서 수행되어 실제 데이터센터 환경 성능을 반영했다. 이러한 융합은 정규화가 차지하는 비중을 줄여 하드웨어 연산 자원을 더 잘 포화시키는 한편, 구현 난이도와 자동튜닝 요구가 트레이드오프로 남는다.

섹션별 상세

정규화 연산은 LayerNorm과 RMSNorm처럼 현대의 LLM과 추천 모델에서 널리 사용되며 훈련 안정화와 수렴 가속에 핵심 역할을 한다. 그러나 이러한 정규화는 TensorCore를 활용하지 못하는 메모리 바운드 연산이어서 전체 훈련 지연에서 상당 부분을 차지하며 Kunlun 기반의 Generative Ads Model 사례에서는 약 20%의 전체 학습 지연을 차지한 것으로 나타났다. 이 때문에 별도 커널로 실행되는 정규화 연산을 단순히 빠르게 실행하는 것만으로는 시스템 전체의 연산 포화도를 회복하기 어려웠다.

메모리 IO를 줄이고 정규화의 별도 로드를 제거하기 위해 Lazy Pre-Norm과 Multi-CTA Norm Fusion 같은 커널 융합 전략을 적용했다. 이러한 기법은 정규화 연산을 인접한 GEMM 연산과 연산적으로 결합하여 중간 결과를 메모리에 쓰고 읽는 비용을 제거하거나 겹치게 만들어 실행 흐름을 재구성한다. 블로그에서는 이 방식으로 정규화 커널의 지연을 최대 90%까지 숨길 수 있다고 제시하며, 메모리 중심 병목을 완화해 하드웨어 연산 자원을 더 잘 포화시킬 수 있다고 주장했다.

어텐션 주변의 다중 정규화를 대상으로 한 FlashNormAttention 알고리즘은 LayerNorm과 RMSNorm을 GDPA와 같은 어텐션 커널 안으로 융합하여 연산을 통합했다. 본문에서는 이 접근으로 특정 커널에서 최대 35%의 속도 향상이 관찰되었고, 실험은 bfloat16 데이터 타입과 NVIDIA B200 GPU, 750W 전력 한계 환경에서 수행되었다. 개발 생산성과 이식성 면에서는 TLX(하드웨어 인지형 Triton 확장)와 Helion(고수준 DSL 및 자동튜닝)이 주요 도구로 사용되어 성능·이식성의 균형을 맞추었다.

용어 해설

행렬곱 연산(GEMM): — 대형 신경망의 핵심 연산으로 입력 행렬과 가중치 행렬의 곱셈을 수행하며 하드웨어에서 높은 연산 집약도를 보이지만 메모리 IO와의 결합 방식에 따라 전체 처리량이 크게 달라진다.
레이어 정규화(LayerNorm): — 내부 임베딩 차원에 대한 정규화로 입력의 분산과 평균을 이용해 스케일·바이어스 보정을 수행하며 주로 훈련 안정화와 수렴 가속에 사용되지만 구현 시 메모리 바운드 성격으로 성능 병목이 된다.
RMS 정규화(RMSNorm): — 평균 대신 루트평균제곱(RMS)을 이용해 스케일링을 수행하는 정규화 방식으로 연산적 차이가 크지 않지만 메모리 접근 패턴이 정규화 성능 저하에 영향을 준다.
커널 융합(Kernel Fusion): — 별도 커널로 실행되던 연산들을 단일 커널로 통합해 메모리 IO와 중간 결과 저장을 줄이고 연속 연산을 연산 집약적으로 재배치하여 GPU 자원 활용도를 높이는 기법이다.
GDPA (어텐션 변형)(GDPA): — 어텐션 연산의 특정 변형으로 본문에서는 어텐션 주변의 다중 정규화를 융합하는 문맥에서 참조되며 해당 구조에 융합을 적용하면 추가적인 병목 해소가 가능하다.
bfloat16: — 훈련과 추론에서 사용되는 16비트 부동소수점 형식으로 낮은 비트폭에도 불구하고 동적 범위를 넓게 유지해 성능을 높일 수 있어 메모리·대역폭 제약 환경에서 유리하다.
Triton DSL: — GPU 실행 제어에 하드웨어 인지적인 저수준 확장을 더한 도메인 특화 언어로 효율적인 커널 구현과 성능 튜닝을 위해 사용되며 본문에서는 TLX의 기반 기술로 언급되었다.

기술

TLX
Helion
bfloat16
NVIDIA B200
GEMM

활용 사례

대형 언어 모델의 훈련 지연 감소
추천 시스템(ads) 모델의 정규화 오버헤드 완화
어텐션 기반 연산의 커널 효율성 향상

언급된 리소스

GitHubads_model_kernel_library - multi_cta_norm_fusion

GitHubads_model_kernel_library - gdpa_megakernel