NVIDIA가 대규모 AI 모델 훈련 효율을 높이는 고급 융합 커널(fusion kernel)을 공개했다. 이 커널은 MoE(Mixture-of-Experts, 전문가 혼합) 아키텍처의 학습 처리량을 끌어올리기 위해 CuTe DSL로 맞춤 제작됐다. NVIDIA에 따르면 DeepSeek-V3 사전학습 전체 구간에서 종단 간(end-to-end) 성능이 8% 개선됐으며, GPT-OSS 사전학습 환경에서는 93%의 속도 향상을 확인했다고 밝혔다.
MoE 모델은 각 토큰 처리 시 전체 파라미터 중 일부만 활성화해 실용적인 컴퓨팅 예산 안에서 거대한 모델 용량을 구현하는 방식이다. 최근 AI 모델 규모가 커질수록 이 블록의 연산 최적화가 훈련 처리량 극대화의 핵심 과제로 부상했다. NVIDIA 연구팀은 프로파일링 결과 MoE 블록에서 세 가지 주요 병목을 발견했다. 첫째, 활성화 함수 연산 시 텐서 읽기·쓰기가 과도해 메모리 바운드 상태가 발생한다. 둘째, 라우팅 전문가 방식에서 전문가별 토큰 수 계산이 CPU에서 이뤄져 GPU가 대기하는 CPU 의존성이 생긴다. 셋째, 고정밀에서 저정밀로의 양자화 과정에서 메모리 바운드 커널이 텐서 코어를 유휴 상태로 만든다.
NVIDIA는 이 세 가지 문제를 해결하기 위해 GroupGemm+양자화, GroupGemm+활성화+양자화/전치, GroupGemm+역활성화+양자화/전치로 구성된 커널 3종을 설계했다. 핵심 기술은 GLU(Gated Linear Unit) 활성화를 GEMM 커널 에필로그에 융합하는 방식으로, 중간 텐서의 전역 메모리 읽기·쓰기를 제거한다. 또한 GPU 메모리 내에서 그룹별 토큰 수를 추적하는 방식으로 CPU 동기화 포인트를 없애 전체 반복 구간을 CUDA 그래프로 캡처할 수 있도록 했다. 단위 수준 벤치마크에서 순전파 최대 1.3배, 역전파 최대 2.1배의 가속이 확인됐다.
이 커널은 cuDNN 프런트엔드(v1.23.0 이상), 트랜스포머 엔진(v2.15 이상), Megatron-Core(26.04-alpha.rc2 이상)를 통해 지금 바로 사용할 수 있다. NVIDIA는 더 많은 융합 패턴 지원, JAX 프레임워크 연동, AOT(사전) 컴파일 등을 지속적으로 개발할 예정이라고 밝혔다. MoE 구조를 채택한 대형 모델 학습 환경에서 훈련 시간 단축과 하드웨어 활용률 개선을 동시에 노리는 조직에 실질적인 선택지가 될 전망이다.
저작권자 © STORIUM 무단전재 및 재배포 금지
