AI 모델 경쟁, 최신 GPU를 도입하는 것만으로 충분할까?
최신 GPU인 NVIDIA Blackwell(B200)을 투입하는 것만으로 성능을 효과적으로 끌어낼 수 있을까요?
아쉽게도 그렇지 않습니다. GPU 연산 소프트웨어가 GPU의 발전 속도에 맞춰 함께 새로 작성되어야만 가속기의 성능을 온전히 끌어낼 수 있습니다. Self-attention 연산 가속 프레임워크인 FlashAttention(FA)이 Ampere용 FA2에서 Hopper용 FA3, 그리고 Blackwell용 FA4로 지속적으로 커널 구조를 발전시켜 온 것과 동일한 원리입니다.
최신 B200 GPU를 도입했는데, 왜 커널은 기대만큼 성능을 내지 못했을까?
AI 모델을 학습하는 과정에서는 GPU 위에서 수많은 연산이 반복적으로 수행됩니다. 이때 GPU에서 각각의 연산을 실제로 수행하는 프로그램을 GPU 커널(Kernel)이라고 합니다.
기존 오픈소스 커널인 mamba-ssm은 최신 B200의 새로 추가된 하드웨어 기능을 활용하지 못하고 있었으며, 구현 당시 최신 GPU였던 A100 기반의 휴리스틱 설정이 오히려 B200 자원 배치의 병목으로 작용하고 있었습니다.
왜 Mamba-2를 최적화했을까? Long-context와 Hybrid Architecture
초거대 AI 모델의 성능이 고도화됨에 따라 ‘긴 문맥(Long-context)을 얼마나 빠르고 정확하게 다루는가’는 모델의 경쟁력을 가르는 핵심 지표가 되었습니다. 하지만 우리가 흔히 사용하는 기존 Transformer 아키텍처에는 치명적인 약점이 존재합니다. 핵심 연산인 Self-attention 구조 특성상, 입력되는 컨텍스트가 길어질수록 연산 복잡도와 KV(key-value) 캐시 메모리 비용이 시퀀스 길이에 비례해 급증한다는 점입니다.
실제 현장의 예시로 보면 어떨까요?
예를 들어 500페이지 분량의 복잡한 법률 문서 뭉치를 한 번에 요약하거나, 수 시간 분량의 고화질 영상 데이터 전체를 입력받아 이상 징후를 분석하는 상황을 가정해 보겠습니다. 이러한 Long-context 환경에서 기존 Self-attention 구조만 적용된 모델은 메모리 부족(Out of Memory, OOM)을 일으키거나, 답변을 내놓기까지 응답 속도가 급격히 느려지는 한계에 직면하게 됩니다. 이러한 ‘연산 병목’과 ‘메모리 폭증’이라는 두 마리 토끼를 잡기 위해, 최근 글로벌 AI 생태계는 큰 기술적 전환을 맞이하고 있습니다.
그림 1. 긴 문맥에서 Self-attention과 Linear attention의 계산 방식 비교
Self-attention vs. Linear attention: 시간, 공간 복잡도 비교
| Self-attention | Linear attention (e.g. Mamba-2) | |
|---|---|---|
| 연산 복잡도 | O(N²) (시퀀스 길이 제곱 비례) | O(N) (시퀀스 길이 선형 비례) |
| KV 캐시 메모리 | 시퀀스가 길어질수록 폭증 (OOM 유발) | 메모리 점유율 최소화 및 일정 유지 |
NVIDIA Nemotron, Alibaba Qwen, Moonshot AI Kimi 등 최첨단 LLM들은 기존 Self-attention 구조의 강점은 유지하면서, 연산 효율이 극대화된 Linear attention 계열(Mamba-2, GDN, KDA 등)을 일정 비율로 혼합하는 하이브리드 구조를 적극 도입하기 시작했습니다. 이를 통해 모델의 뛰어난 언어 이해 및 추론 성능은 그대로 유지하면서도, 연산 복잡도와 KV 캐시 메모리 부담을 낮추는 방향으로 활용되고 있습니다.
이처럼 Mamba-2와 같은 연산 구조의 활용이 확대되면서, 이를 최신 GPU에서 얼마나 효율적으로 실행할 수 있는지 역시 중요한 문제가 됩니다.
B200에 맞는 Mamba-2 커널 개발, 왜 Helion인가?
기존 오픈소스의 커널에서 B200 성능을 효과적으로 끌어내기 위해, 저희는 새로운 Helion DSL을 활용하는 전략을 선택했습니다. 저희가 선택한 이유는 하나로 요약됩니다. 기존 Triton의 Autotune 기능을 확장하여 GPU 아키텍처와 입출력 디멘션, 데이터 타입 등 실제 연산 조건에 맞는 설정값을 자동으로 탐색해 주기 때문입니다. 구체적으로는 다음 세 가지입니다.
1. 생태계 호환성 유지
Helion은 기존 Triton 위에서 추상화를 제공하는 Domain-specific Language(특정 목적에 맞춰 설계된 언어)입니다. 따라서 기존 PyTorch 및 Triton 개발 로직과 생태계 호환성을 그대로 유지하면서 커널 재설계의 자율성을 대폭 확보할 수 있습니다.
2. 강력한 Autotuner 탑재
A100 대비 B200에서 새로 지원하는 명령어와 하드웨어 기능(TMA, TMEM, Warp Specialization 등)을 최적의 연산 지점에 적절히 적용하도록 돕는 자동 튜닝 엔진을 제공합니다.
3. 휴리스틱 최적화의 자동화
기존 Triton 커널에서 엔지니어가 수동으로 고정해 두었던 휴리스틱 최적화 파라미터들을 오토튜닝 탐색 대상에 통합함으로써 하드웨어 타깃 맞춤형 최고 성능을 도출할 수 있습니다.
최적화 과정: 느린 곳을 찾고 → 정확성을 지키고 → B200에 맞게 바꾸다
B200 타깃의 mamba-ssm 커널 최적화는 단기적인 단편 개선이 아닌, 체계적인 3단계 공정으로 진행되었습니다. 연산 병목을 정확히 진단하고, 모델 학습의 수치 정밀도를 보장하며, 오토튜너가 하드웨어 기능을 적절하게 적용하는 것을 넘어 로직과 휴리스틱 개선을 수행하도록 변경했습니다.
그림 2. 세 단계 공정: 프로파일링 및 병목 분석 → 수치 정밀도 확보 → 로직·휴리스틱 개선
1단계. Profiling & Bottleneck (연산 비중 분석)
첫 번째 단계에서는 Mamba-2 모델 학습 과정의 Forward 및 Backward 연산 전체에서 수행되는 개별 커널들의 연산 비중을 정밀하게 프로파일링했습니다.
그 결과 상위 10개 커널이 전체 연산 시간의 82%를 차지하는 것을 확인했습니다. 프로파일링 결과를 바탕으로 병목이 가장 심한 핵심 커널을 먼저 선별하고 최적화 우선순위를 부여함으로써, 적은 작업량으로도 실행 지연 시간 단축 효과를 극대화하는 전략을 취했습니다.
커널별 연산 비중
상위 10개 커널이 전체 연산 시간의 82%를 차지합니다.
| Kernel Name | Phase + Step | % (Baseline) |
|---|---|---|
| _layer_norm_bwd | bwd_2 | 12.68 |
| _chunk_scan_fwd | fwd_5 | 12.32 |
| _chunk_scan_chunk_state_bwd_dx | bwd_10 | 10.81 |
| _chunk_scan_bwd_ddAcs_stable | bwd_16 | 9.43 |
| _chunk_state_fwd | fwd_2, bwd_5 | 9.25 |
| _chunk_scan_bwd_dC | bwd_12 | 6.09 |
| _chunk_state_bwd_db | bwd_11 | 6.04 |
| _state_passing_fwd | fwd_3, bwd_6 | 5.95 |
| _state_passing_bwd | bwd_9 | 5.54 |
| _chunk_scan_bwd_dstates | bwd_8 | 4.17 |
| 합계 | 82.28 |
이하 본문에 나오는 bwd_2, fwd_5 등의 표기는 위 표의 Phase + Step 열을 가리킵니다.
2단계. Numerical Parity (정밀도 보장)
아무리 연산 속도가 빨라지더라도 정밀도에 오차가 생겨 수렴성이 깨진다면 최적화의 의미가 퇴색됩니다.
따라서 가속화에 앞서 기존 Triton 구현체의 수학적 연산 메커니즘을 Helion상에 1:1로 동일하게 재현했습니다. 수치 오차를 최소화하여, 초거대 모델 학습 파이프라인 적용 시에도 안정성이 훼손되지 않도록 검증했습니다.
3단계. Logic & Heuristic Refinement (구조 및 휴리스틱 개선)
수치 정밀도가 확보된 후, B200의 성능을 극대화하기 위해 일부 커널에 연산 구조와 파라미터를 재구성했습니다.
Associative Scan 명령어 도입: 기존 _state_passing_fwd, _state_passing_bwd 연산의 순차 루프 구조를 그대로 Helion에서 사용했을 때 성능 향상이 미미했기 때문에 Helion의 Associative Scan 명령어 기반 알고리즘으로 전환하여 연산 병렬성을 극대화했습니다.
그림 3. 순차 루프와 Associative Scan의 실행 단계 비교
GPU SM Count 기반 휴리스틱 개선: 기존 SM 개수 기반으로 고정되어 있던 휴리스틱 파라미터는 B200에서 성능 저하를 유발할 가능성이 있습니다. 이를 오토튜닝 탐색 대상에 추가하여 새로운 환경에서 최적으로 동작하도록 개선했습니다. 이 변경이 가장 큰 효과를 냈습니다.
빠른 커널을 만드는 것에서 실제 학습에 사용하는 것까지
앞선 3단계 파이프라인(프로파일링 및 병목 분석, 수치 정밀도 확보, 로직 및 휴리스틱 개선)을 통해 개별 커널 연산의 수치적 정밀도와 실행 속도를 대폭 개선했습니다. 하지만 단지 커널 최적화만으로 엔지니어링이 완료되는 것은 아닙니다. 다음으로 해결해야 할 핵심 과제는 ‘이 최적화된 커널들을 대규모 학습 파이프라인에 어떻게 손실 없이 내재화하고 유연하게 통합할 것인가’였습니다.
커널 단에서 확보한 뛰어난 성능이 실제 대규모 Pre-training 환경까지 손실 없이 이어지기 위해서는 두 가지 핵심 엔지니어링 장벽을 극복해야 했습니다. 첫째는 모델 구조나 배치 설정이 바뀔 때마다 오토튜닝에 1~2시간씩 소요되는 시간 오버헤드이며, 둘째는 기존 학습 파이프라인과의 강력한 이식성 및 확장성 확보였습니다.
Automated Config Generator를 통한 오토튜닝 오버헤드 개선
Helion은 강력한 오토튜닝을 통해 B200에 최적화된 타일링 및 블록 하이퍼파라미터를 자동으로 탐색해 주는 장점이 있습니다. 그러나 학습 환경이 바뀔 때마다 매번 1~2시간의 탐색 오버헤드를 감수하는 것은 수시로 실험을 반복하고 대규모 클러스터를 운용해야 하는 실제 현장에서 큰 병목이 됩니다.
이를 해결하기 위해 오토튜닝 탐색과 실제 모델 학습 시점을 완벽히 분리하는 Automated Config Generator를 구축했습니다. 학습 환경별 최적의 타일링 및 블록 크기 파라미터를 사전에 자동 탐색하여 파일로 저장해 두고, 실제 Pre-training 단계에서는 저장된 Config를 즉시 로드하여 사용하도록 설계했습니다. 그 결과, 1~2시간에 달하는 오토튜닝 대기 시간 없는 최적화된 연산 성능을 즉각 발휘할 수 있게 되었습니다.
그림 4. 사전 오토튜닝과 실제 학습 시점의 Config 재사용 흐름
1:1 Drop-in 모듈 구조와 PyTorch/Triton 워크플로우 이식성
아무리 뛰어난 최적화 커널이라도 기존 학습 프레임워크와 호환되지 않거나 도입 시 대대적인 코드 수정이 필요하다면 실제 연구개발 현장에 바로 적용하기 어렵습니다. 이에 저희는 기존 mamba-ssm 오픈소스 패키지와 1:1로 직접 교체 가능한 모듈 구조를 설계하여 PyTorch/Triton 워크플로우와의 완벽한 이식성을 확보했습니다.
이러한 1:1 대응 모듈 구조는 두 가지 결정적인 엔지니어링 이점을 제공합니다.
신속한 수치 오차 추적
3단계 파이프라인에서 검증한 정밀도를 바탕으로, 대규모 연산 시 발생할 수 있는 수치 오차 누적 여부를 단일 레이어 단위에서 신속하게 추적하고 검증할 수 있습니다.
높은 모듈 재활용성 및 확장성
동일한 커널 모듈을 상위 프레임워크의 구조적 변경 없이 재활용할 수 있어, 향후 Long-context 데이터 처리를 위한 Context Parallelism 등 다양한 분산 학습 패러다임 연산과 연계하여 손쉽게 확장할 수 있는 기반을 완성했습니다.
그래서 얼마나 빨라졌을까?
앞서 구축한 3단계 커널 최적화와 자동화 파이프라인이 실제 하드웨어 상에서 어느 정도의 성능 격차를 만들어 냈을까요?
- 특정 병목 커널 _layer_norm_bwd: 최대 3.18배
- Mamba-2 단일 레이어(Forward + Backward): 약 1.33배
- 실제 Nemotron-3 Nano 30B Pre-training E2E: 약 1.12배
Mamba-2 단일 레이어(Forward + Backward) 기준으로 기존 오픈소스 mamba-ssm 베이스라인 대비 1.33배의 속도 향상을 달성했습니다.
커널별 Latency 및 Speedup
표의 Target은 최적화 대상으로 삼은 상위 10개 커널(전체의 82%), Total은 레이어 전체를 의미합니다. 대상 커널 전체에서 속도가 개선되었으며, 개선폭은 1.20배에서 3.18배까지 분포합니다.
| Phase + Step | % (Baseline) | Latency (Baseline, ms) | Latency (Helion, ms) | Speedup |
|---|---|---|---|---|
| bwd_2 | 12.68 | 0.28 | 0.09 | 3.18 |
| fwd_5 | 12.32 | 0.27 | 0.20 | 1.33 |
| bwd_10 | 10.81 | 0.24 | 0.19 | 1.22 |
| bwd_16 | 9.43 | 0.21 | 0.16 | 1.30 |
| fwd_2, bwd_5 | 9.25 | 0.20 | 0.15 | 1.33 |
| bwd_12 | 6.09 | 0.13 | 0.09 | 1.48 |
| bwd_11 | 6.04 | 0.13 | 0.10 | 1.31 |
| fwd_3, bwd_6 | 5.95 | 0.13 | 0.11 | 1.20 |
| bwd_9 | 5.54 | 0.12 | 0.09 | 1.32 |
| bwd_8 | 4.17 | 0.09 | 0.07 | 1.25 |
| Target 최적화한 커널 10개 |
82.28 | 1.80 | 1.26 | 1.43 |
| Total 레이어 전체 |
100.00 | 2.18 | 1.64 | 1.33 |
측정 조건: B200 1장, bf16, Nemotron-3 Nano 30B 구성, batch 2 / seqlen 8192. Warm-up 10회 후 100회 반복 측정 중앙값. mamba-ssm v2.3.1 기준.
단일 커널 단계별 성능을 추적한 결과, 특히 _layer_norm_bwd 연산에서 3.18배라는 독보적인 가속 성능을 기록했습니다.
결론부터 말하면, 새로운 알고리즘을 넣어서가 아닙니다. A100 시절 기준으로 굳어 있던 설정값 하나가 B200의 성능을 묶어두고 있었고, 그걸 바로잡은 결과입니다.
‘고속도로가 넓어졌는데 차를 안 늘렸다.’
고속도로를 4차선에서 16차선으로 넓혔다고 생각해 봅시다. 그런데 차 한 대가 서울에서 부산까지 가는 데 걸리는 시간은 그대로입니다. 도로를 넓힌 효과를 보려면 차를 그만큼 더 많이, 동시에 굴려야 합니다. 차를 4차선 시절만큼만 보낸다면 16차선은 텅 빈 채로 남습니다. B200이 이런 상황이라고 이해할 수 있습니다.
_layer_norm_bwd는 전형적인 Memory-bound(메모리 속도에 성능이 좌우되는) 연산입니다. A100에서 B200으로 넘어오며 HBM Bandwidth는 2TB/s에서 8TB/s로 약 4배 늘었지만, HBM Latency는 그만큼 줄지 않았습니다.
이 불균형이 핵심입니다. 대역폭만 4배가 되고 지연은 그대로라면, 파이프를 채우기 위해 SM당 동시 메모리 요청을 그만큼 더 많이 발행해야 합니다. 요청이 부족하면 커널은 대역폭이 아니라 지연 시간에 묶이고, 이 상태에서는 하드웨어 대역폭이 아무리 늘어도 그 혜택을 받지 못합니다.
어떻게 해결했을까요?
기존 Triton 휴리스틱은 ‘SM당 warp 4개’를 목표로 파라미터를 결정합니다. A100 기준으로는 타당한 값이지만, B200에서는 상황이 달라집니다. B200의 SM 148개를 기준으로 기존 휴리스틱을 적용하면 nrow_groups 파라미터가 74로 산출됩니다. Helion Autotuner가 B200 구조를 반영해 정밀 탐색한 최적값은 2601로, 약 35배 차이입니다. 동시 메모리 요청이 비약적으로 늘며 Latency가 감춰졌고, 3.18배의 성능 개선으로 이어졌습니다.
그림 5. B200 메모리 대역폭을 활용하기 위한 동시 요청 증가 원리
이 개선을 만든 것은 새로운 구현 기법이 아니라 병렬성 파라미터값 그 자체입니다. 기존 Triton 커널에서 이 값을 직접 조정해도 유사한 결과에 도달할 수 있을 것으로 보입니다.
그럼에도 오토튜닝이 결정적이었던 이유는, 이 파라미터가 SM 개수만으로 결정되지 않는다는 점에 있습니다. 입력 디멘션에 따라서도 최적값이 달라지기 때문에, 실험 조건이 바뀔 때마다 사람이 값을 다시 찾아야 합니다. 이 탐색 과정을 자동화한 것이 Helion을 도입해 얻은 실질적인 이득이었습니다.
이 사례가 말하는 것
이번 최적화는 새로운 알고리즘을 도입한 것이 아니라, B200이 원래 갖고 있던 성능을 되찾은 것에 가깝습니다. 커널이 하드웨어 세대 전환의 혜택을 온전히 받으려면 병렬성 파라미터도 함께 따라와야 한다는 사실을 보여주는 사례입니다. 이 개선 폭은 기존 커널이 B200이 새로 확보한 메모리 대역폭을 충분히 활용하지 못하고 있었다는 신호로 읽힙니다.
성능이 빨라져도 정확성과 학습 성능이 유지될까?
연산 속도가 아무리 빨라지더라도, 수치 오차가 누적되어 모델 학습의 수렴성이 깨진다면 가속화의 의미는 퇴색됩니다. 속도 향상이 정밀도 손실 없이 이루어졌음을 입증하기 위해 공식 mamba-ssm 유닛 테스트 스위트를 기반으로 엄격한 정밀도 검증을 진행했습니다.
단계별 정밀도 검증 데이터
기존 Triton 구현과 출력을 비교한 오차입니다.
diff_mean: 기존 구현과의 차이(절댓값)의 평균으로, 새 구현이 수학적으로 기존과 동치인지 확인하는 값입니다.
diff_max: 같은 차이의 최댓값으로, 연산 순서나 커널 내부 정밀도 변동 때문에 특정 지점에서 오차가 튀지 않는지 확인합니다.
| Phase + Step | diff_mean | diff_max |
|---|---|---|
| bwd_2 | 2.24E-08 | 1.56E-02 |
| fwd_5 | 5.59E-09 | 1.56E-02 |
| bwd_10 | 4.78E-06 | 1.80E-02 |
| bwd_16 | 3.43E-07 | 2.77E-05 |
| fwd_2, bwd_5 | 2.84E-04 | 2.84E-04 |
| bwd_12 | 1.08E-08 | 2.86E-06 |
| bwd_11 | 3.51E-08 | 6.20E-06 |
| fwd_3, bwd_6 | 2.98E-08 | 6.10E-05 |
| bwd_9 | 1.40E-08 | 3.13E-02 |
| bwd_8 | 2.29E-07 | 6.01E-02 |
검증 기준: |차이| ≤ atol + rtol × |기준값| (atol = 1×10⁻², rtol = 1×10⁻²)
검증 결과: Forward 및 Backward 전 단계에 걸쳐 출력을 비교한 결과, Absolute Difference 평균값이 5×10^-9 ~ 3×10^-4 수준을 기록했습니다.
Nemotron-3 Nano 30B E2E 학습 실증 과정
단일 레이어 벤치마크를 넘어, 실제 Pre-training 파이프라인인 Nemotron-3 Nano 30B (A3B) 모델 학습 환경에 최적화 커널을 1:1로 교체 투입하여 End-to-End 실증을 거쳤습니다.
Iteration Time 단축
전체 모델 단위의 학습 스텝당 처리 속도가 1.12배 단축되어 대규모 클러스터 운용 시간과 비용을 실질적으로 절감했습니다.
Loss 수렴성 보장
Baseline 커널 대비 Loss 수렴 곡선이 동일한 양상으로 나타나는 것을 확인했습니다. 이를 통해 수십만 스텝 이상의 장기 학습 과정에서도 NaN 발생이나 Loss Spike 없이 안정적으로 동작함을 실증했습니다.
그림 6. 학습 단계에 따른 Loss 곡선: 기존 커널과 Helion 최적화 커널 비교
E2E 실증 조건: 32노드, global batch 3072, seqlen 8192, tensor parallel 2, expert parallel 8
결론: GPU 확보와 성능 끌어올리기
B200과 같은 최첨단 GPU 가속기를 도입하는 것만으로는 초거대 AI 시대의 주도권을 확보하기 어렵습니다. 하드웨어의 잠재 성능을 끌어낼 수 있는 소프트웨어 코어 커널 기술의 독자적 내재화가 반드시 병행되어야 합니다. 네이버클라우드는 B200 특성에 맞춘 Helion 기반 Mamba-2 커널 최적화를 통해 단일 레이어 1.33배, End-to-End 학습 1.12배 가속이라는 정량적 결실을 증명해 냈습니다.
Nemotron-3 Nano 30B에서 검증된 가속 성과는 향후 좀 더 큰 모델 라인업으로 확대 적용할 예정입니다. 특히 시퀀스 길이가 더욱 극단적으로 길어지는 Long-context 환경에 대비하여, 이번 Helion Mamba-2 커널을 Context Parallelism 및 Multi-node GPU 클러스터 분산 환경과 연계하는 확장 연구를 진행하고 있습니다. 이를 통해 단일 노드를 넘어 클러스터 전체의 연산 효율을 수평적으로 끌어올릴 계획입니다.
앞으로도 네이버클라우드는 선제적인 AI 모델 아키텍처 연구와 차세대 가속기 소프트웨어 최적화 역량을 하나로 묶어, 연산 효율성을 갖춘 자주적 AI 인프라 생태계를 구축해 나갈 것입니다.