네이버클라우드, B200 GPU용 Mamba-2 커널 최적화로 단일 레이어 1.33배 가속

네이버클라우드가 NVIDIA Blackwell(B200) GPU의 하드웨어 잠재력을 최대한 끌어내기 위해 Mamba-2 커널을 최적화했다. 기존 오픈소스 커널인 mamba-ssm은 A100 기반 휴리스틱 설정으로 인해 B200의 새로운 하드웨어 기능을 충분히 활용하지 못했으며, 이를 해결하기 위해 Helion DSL을 도입했다.

Helion 기반 3단계 최적화 공정

최적화 과정은 병목 분석, 수치 정밀도 확보, 로직 및 휴리스틱 개선의 3단계로 진행됐다. 먼저 전체 연산 시간의 82%를 차지하는 상위 10개 커널을 프로파일링해 병목을 선별했다. 이어 기존 Triton 구현체와 수학적 연산 메커니즘을 1:1로 재현해 수치 오차를 최소화했다. 마지막으로 B200의 아키텍처 특성을 반영해 연산 구조를 재구성했다. 특히 _state_passing 연산의 순차 루프를 Helion의 Associative Scan 명령어 기반으로 전환해 병렬성을 극대화했다.

병목 해결과 성능 향상

핵심 병목 중 하나인 _layer_norm_bwd 커널은 A100 기준의 고정 파라미터가 B200의 대역폭 증가와 지연 시간 불균형을 반영하지 못해 성능 저하를 유발했다. Helion의 Autotuner가 B200 구조에 맞춰 nrow_groups 파라미터를 기존 74에서 2601로 정밀 탐색해 조정함으로써, 해당 커널에서 3.18배의 속도 향상을 달성했다. 이는 새로운 알고리즘 도입이 아닌, 하드웨어 세대 전환에 따른 병렬성 파라미터 조정을 통해 기존 하드웨어의 잠재력을 되찾은 결과다.

실제 학습 환경 적용 및 효과

오토튜닝 오버헤드를 줄이기 위해 학습 환경별 최적 파라미터를 사전에 저장하고 로드하는 Automated Config Generator를 구축했다. 이를 통해 실제 Pre-training 단계에서 1~2시간의 대기 시간 없이 최적화된 성능을 즉시 발휘할 수 있다. Nemotron-3 Nano 30B 모델의 End-to-End 학습 테스트에서 전체 학습 스텝당 처리 속도가 1.12배 단축되었으며, Loss 수렴 곡선은 Baseline과 동일한 양상을 보여 안정성을 입증했다. 네이버클라우드는 향후 Context Parallelism 등 분산 학습 패러다임과 연계해 클러스터 전체의 연산 효율을 높일 계획이다.