MoE Fusion Kernels는 NVIDIA가 CuTe DSL로 구현한 dense/MoE 모델용 fused MLP 커널 묶음이다. MoE(Mixture-of-Experts) 학습에서 활성화 함수, 저정밀 양자화, expert별 동적 실행이 만드는 메모리·동기화 병목을 줄여 Tensor Core 활용률과 전체 학습 처리량을 높인다.
왜 MoE 블록이 병목이 되는가
MoE는 토큰마다 일부 expert만 활성화해 모델 용량을 키우지만, 학습 시스템 관점에서는 세 가지 병목을 만든다.
| 병목 | 문제 |
|---|---|
| 활성화 함수 | SwiGLU, GeGLU 같은 GLU 계열 연산이 중간 텐서를 메모리에 쓰고 다시 읽음 |
| CPU 동기화 | expert별 토큰 수가 런타임에 정해져 CPU가 커널 크기와 실행을 조정해야 함 |
| 양자화 비용 | MXFP8, NVFP4 등 저정밀 학습에서 activation 뒤 별도 quantize/transpose 패스 발생 |
이 병목은 GEMM 자체보다 주변 메모리 이동과 launch/sync 비용에서 발생한다. NVIDIA의 접근은 이 주변 연산을 GEMM epilogue와 grouped GEMM 커널 안으로 흡수하는 것이다.
핵심 커널 패턴
| 커널 | 역할 |
|---|---|
| GroupGemm + Quantize | expert별 GEMM과 저정밀 변환을 한 번에 처리 |
| GroupGemm + Activation + Quantize/Transpose | forward pass에서 FC1, GLU 활성화, 양자화/전치를 융합 |
| GroupGemm + dActivation + Quantize/Transpose | backward pass에서 dActivation과 양자화/전치를 융합 |
지원 활성화 함수는 SwiGLU, GeGLU, sReLU이며, feature scaling, clamp, bias addition도 epilogue 안에서 처리할 수 있다.
GLU 활성화 융합
SwiGLU 같은 GLU 계열 함수는 입력 경로와 gate 경로를 나눠 계산한 뒤 결합한다. 일반적인 구현에서는 두 출력을 전역 메모리에 쓴 뒤 다시 읽어야 한다.
NVIDIA는 weight를 input/gate가 같은 thread block에서 접근 가능한 형태로 repack한다. 이렇게 하면 GEMM epilogue 안에서 gate와 input을 바로 결합할 수 있고, 중간 텐서의 global memory round-trip을 제거한다. weight repack은 학습 시작 전 체크포인트 로딩 단계에서 수행할 수 있다.
CPU 동기화 제거
전통적인 grouped GEMM은 expert별 토큰 수를 CPU가 알아야 커널 launch를 구성할 수 있다. MoE에서는 이 토큰 수가 매 iteration마다 달라지므로 host-device synchronization이 생긴다.
CuTe DSL GroupGEMM 커널은 그룹별 토큰 수를 GPU 메모리에서 직접 추적한다. 그 결과 CPU 개입 없이 전체 iteration CUDA Graph를 구성할 수 있고, expert 수가 늘어날수록 커지는 launch overhead를 줄인다.
성능 효과
NVIDIA Technical Blog 기준 수치는 다음과 같다.
| 범위 | 개선 |
|---|---|
| kernel level forward | 최대 1.3배 |
| kernel level backward | 최대 2.1배 |
| DeepSeek-V3 pre-training setup | end-to-end 8% 개선 |
| GPT-OSS pre-training setup | end-to-end 93% 개선 |
GPT-OSS 사례처럼 전체 개선 폭이 큰 경우는 단일 커널 속도뿐 아니라 sync-free 실행, CUDA Graph, 통신 overlap이 함께 작동한 결과다.
사용 경로
| 레이어 | 접근 방식 |
|---|---|
| cuDNN Frontend 1.23.0+ | 커널 직접 호출 또는 wrapper API |
| Transformer Engine 2.15+ | transformer_engine.pytorch.ops와 Sequential 패턴 매칭 |
| Megatron Core 26.04-alpha.rc2+ | Megatron 설정 knob로 활성화 |
직접 커널을 작성하지 않는 팀도 Transformer Engine이나 Megatron Core를 쓰면 상위 추상화에서 활용할 수 있다.
관련 문서
- vllm — LLM 추론·서빙 엔진
- tiny-vllm — C++와 CUDA로 LLM 추론 엔진을 구현하는 교육 과정
- kvarn — KV 캐시 양자화 기반 vLLM 백엔드
- north-mini-code — 30B 총 파라미터, 3B 활성 파라미터 MoE 코딩 모델
참고 자료
- Boosting MoE Training Throughput with Advanced Fusion Kernels — NVIDIA Technical Blog (2026-06-15)