B200에 맞춘 Helion Mamba-2 커널 최적화

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

TL;DR

B200 같은 최신 GPU를 도입해도 기존 GPU 세대에 맞춰진 커널 설정이 남아 있으면 새 하드웨어의 메모리 대역폭과 병렬 자원을 충분히 사용하지 못합니다. 네이버클라우드는 Helion의 Autotuner로 Mamba-2 기반 mamba-ssm 커널의 병목을 찾고, 수치 동등성을 확인한 뒤 B200에 맞춰 연산 구조와 휴리스틱을 다시 조정했습니다. 그 결과 _layer_norm_bwd 커널은 최대 3.18배, Mamba-2 단일 레이어는 1.33배, Nemotron-3 Nano 30B 전체 학습은 1.12배 빨라졌습니다. 오토튜닝 결과를 사전 저장하는 Automated Config Generator와 1:1 교체형 모듈을 함께 구축해 반복 학습에서 발생하는 탐색 대기 시간과 통합 부담도 줄였습니다.

빠른 이해

새로운 점

B200의 성능을 제한하던 A100 기준 휴리스틱을 Helion Autotuner로 재탐색해, 새로운 알고리즘 없이 Mamba-2 커널의 병렬성 설정만으로 최대 3.18배 가속을 확보했습니다.

핵심 메커니즘

Mamba-2 학습의 Forward·Backward 커널을 프로파일링해 전체 실행 시간의 82%를 차지하는 상위 10개를 선별하고, 기존 Triton 연산을 Helion에서 수치적으로 재현한 뒤 Associative Scan과 B200 맞춤 휴리스틱을 적용합니다. Helion Autotuner는 GPU 아키텍처, 입력 디멘션, 데이터 타입에 따라 타일링·블록 크기와 nrow_groups를 탐색하며, Automated Config Generator는 최적 설정을 사전 저장해 실제 Pre-training에서 재탐색 없이 로드합니다.

핵심 수치

상위 10개 커널의 전체 연산 시간 비중: 82.28%- Mamba-2 Forward 및 Backward 전체 커널 프로파일링 기준
_layer_norm_bwd 속도 향상: 최대 3.18배- B200, bf16, mamba-ssm v2.3.1 기준
Mamba-2 단일 레이어 속도 향상: 1.33배- Forward + Backward, 기존 오픈소스 mamba-ssm 대비
Nemotron-3 Nano 30B E2E 학습 속도 향상: 1.12배- 32노드, global batch 3072, seqlen 8192, tensor parallel 2, expert parallel 8
_layer_norm_bwd 지연 시간: 0.28ms -> 0.09ms- B200 한 장, bf16, batch 2, seqlen 8192 조건
오토튜닝 대기 시간: 1~2시간- Automated Config Generator로 사전 탐색 및 Config 재사용

섹션별 상세

Long-context와 Mamba-2의 역할

긴 문맥을 처리하는 기존 Transformer의 Self-attention은 시퀀스 길이가 늘어날수록 연산 복잡도가 O(N²)로 증가하고 KV 캐시 메모리도 커져 OOM이나 응답 지연을 일으킬 수 있습니다. Linear attention 계열인 Mamba-2는 상태를 선형적으로 처리해 연산 복잡도를 O(N)으로 유지하고, 시퀀스가 길어질 때 메모리 부담을 낮추는 방향의 구조입니다. NVIDIA Nemotron, Alibaba Qwen, Moonshot AI Kimi 등은 Self-attention의 언어 이해 및 추론 특성과 Mamba-2, GDN, KDA 같은 Linear attention 계열의 효율을 섞는 하이브리드 구조를 도입하고 있습니다. 따라서 Mamba-2를 실제 최신 GPU에서 빠르게 실행하는 커널 구현은 Long-context 모델의 구조적 장점을 학습 및 추론 처리량으로 연결하는 핵심 요소가 됩니다.

Helion 기반 B200 커널 재설계

Helion은 Triton 위에서 동작하는 Domain-specific Language로, PyTorch와 Triton 워크플로우를 유지하면서 GPU 아키텍처와 입력 디멘션, 데이터 타입에 맞는 커널 설정을 자동 탐색합니다. B200에서 지원되는 TMA, TMEM, Warp Specialization 같은 기능을 적용 후보에 넣고, 기존 Triton 커널에서 사람이 고정했던 타일 및 휴리스틱 파라미터까지 탐색 대상으로 확장했습니다. 최적화는 병목 프로파일링, 수치 정밀도 확보, 로직과 휴리스틱 개선의 세 단계로 진행됐으며, 기존 Triton의 수학적 연산을 Helion에서 1:1로 재현한 뒤 구조를 바꿨습니다. 특히 _state_passing_fwd와 _state_passing_bwd의 순차 루프를 Associative Scan 기반으로 전환하고, SM 개수에 고정된 파라미터를 오토튜닝에 포함해 B200의 병렬 자원을 더 많이 사용하도록 했습니다.

오토튜닝 결과의 학습 파이프라인 통합

Helion의 오토튜닝은 학습 환경마다 최적 설정을 찾지만, 모델 구조나 배치가 바뀔 때 1~2시간의 탐색 시간이 다시 필요하다는 운영 부담이 있습니다. 네이버클라우드는 Automated Config Generator로 환경별 타일링과 블록 크기를 사전에 탐색해 파일로 저장하고, 실제 Pre-training에서는 저장된 Config를 즉시 불러오도록 분리했습니다. 또한 기존 mamba-ssm과 1:1로 교체할 수 있는 모듈 구조를 설계해 PyTorch와 Triton 기반 학습 코드의 대규모 수정을 피했습니다. 단일 레이어 단위의 수치 오차 추적과 Context Parallelism 같은 분산 학습 방식으로의 확장을 함께 고려할 수 있어, 커널 성능을 실제 클러스터 학습으로 옮기는 통합 경로를 마련했습니다.

B200에서 확인된 가속과 정밀도

B200 한 장에서 bf16, batch 2, seqlen 8192, Nemotron-3 Nano 30B 구성으로 측정한 결과, 최적화 대상 상위 10개 커널은 기존 1.80ms에서 Helion 1.26ms로 줄어 1.43배 빨라졌고 레이어 전체는 2.18ms에서 1.64ms로 줄어 1.33배 빨라졌습니다. _layer_norm_bwd는 0.28ms에서 0.09ms로 줄어 3.18배의 가장 큰 개선을 기록했으며, 기존 휴리스틱의 nrow_groups 74를 B200에 맞는 2601로 조정해 동시 메모리 요청을 늘린 결과입니다. 공식 mamba-ssm 유닛 테스트와 출력을 비교한 결과 diff_mean은 5×10^-9에서 3×10^-4 수준이었고, E2E 학습에서는 32노드, global batch 3072, seqlen 8192 조건에서 Iteration Time이 1.12배 개선됐습니다. Nemotron-3 Nano 30B의 Loss 수렴 곡선은 기준 커널과 같은 양상을 유지했고, 장기 학습에서 NaN이나 Loss Spike 없이 동작해 속도 향상과 학습 안정성을 함께 확인했습니다.

용어 해설

선형 어텐션(Linear attention): — 입력 시퀀스 길이에 따라 연산량이 제곱으로 늘어나는 Self-attention과 달리, 상태를 누적해 연산 복잡도를 O(N)으로 낮추는 구조입니다. 긴 문맥에서 메모리 사용량과 처리 지연을 줄이는 데 중요합니다.
GPU 커널(GPU Kernel): — GPU에서 특정 연산을 실제로 실행하는 프로그램 단위입니다. 모델의 수학적 구조가 같아도 커널의 병렬성, 메모리 접근 방식, 하드웨어 기능 사용 여부에 따라 실행 속도가 크게 달라집니다.
오토튜닝(Autotuning): — GPU 아키텍처, 입력 크기, 데이터 타입에 맞는 타일 크기와 병렬성 파라미터를 여러 후보 중 자동으로 탐색하는 방법입니다. 사람이 고정한 휴리스틱보다 특정 실행 환경에 맞는 설정을 찾는 데 유리합니다.
연관 스캔(Associative Scan): — 순차적으로 처리하던 누적 연산을 결합 법칙이 성립하는 형태로 바꿔 여러 GPU 스레드가 병렬 처리하도록 만드는 연산 방식입니다. 순차 루프의 병목을 줄이고 GPU 병렬성을 높이는 데 사용됩니다.
메모리 병목 연산(Memory-bound): — 연산 유닛의 계산 능력보다 GPU 메모리에서 데이터를 가져오는 속도와 지연 시간이 성능을 제한하는 상태입니다. 동시 메모리 요청을 충분히 늘려야 높은 메모리 대역폭을 실제 처리량으로 전환할 수 있습니다.

기술

NVIDIA Blackwell
B200
A100
FlashAttention
FA2
FA3
FA4
Mamba-2
mamba-ssm
Helion
Triton
PyTorch
TMA
TMEM
Warp Specialization
Associative Scan
Automated Config Generator
Context Parallelism
Nemotron-3 Nano 30B

언급된 리소스

문서원문 링크