전문가 혼합(MoE, Mixture-of-Experts) 모델은 현대의 대규모 AI 시스템에서 빠르게 핵심 구성 요소로 자리 잡았습니다. MoE는 토큰마다 파라미터의 일부만 활성화하면서도 모델 용량을 크게 키울 수 있어 폭넓게 채택되고 있으며, 실용적인 연산 예산 안에서 성능을 확장하는 데 비할 데 없는 접근 방식을 제공합니다. 모델 규모가 계속 커짐에 따라, 이러한 블록을 최적화하는 일은 학습 처리량을 극대화하는 데 매우 중요해지고 있습니다.
이 한계를 더 밀어붙이기 위해, 우리는 NVIDIA CuTe DSL로 직접 구현한, 밀집(dense) 모델과 MoE 모델을 위한 고급 융합 MLP 커널을 선보입니다. 이 새로운 커널들은 내재된 메모리 및 동기화 병목을 해결함으로써 비융합(unfused) 경로 대비 커널 수준에서 인상적인 1.3~2배 속도 향상을 제공하는 동시에, 전체 이터레이션을 포괄하는 NVIDIA CUDA 그래프를 위한 동기화 없는(sync-free) MoE 실행을 가능하게 합니다.
NVIDIA 풀스택 DeepSeek-V3 사전 학습 설정에서 이 최적화는 종단 간(end-to-end) 성능을 8% 개선하는 데 기여합니다. 마찬가지로 GPT-OSS 사전 학습 설정에서는 이 최적화가 종단 간 성능을 93% 개선하는 데 기여합니다. 학습 시간을 단축하고 싶든 하드웨어 활용도를 최적화하고 싶든, 이 커널들은 오늘날 NVIDIA cuDNN Frontend에서 바로 사용할 수 있으며 NVIDIA Transformer Engine과 NVIDIA Megatron-Core를 통해 매끄럽게 접근할 수 있습니다.
그 원리를 이해하려면, 현대 MoE 블록을 괴롭히는 가장 큰 세 가지 병목을 해체하는 과정과, 텐서 코어(Tensor Core)에 작업을 끊김 없이 공급하기 위해 하드웨어를 고려한 소프트웨어 코디자인으로 스택을 어떻게 재설계했는지를 체계적으로 살펴봐야 합니다.
MoE 블록의 학습 병목 극복하기
MoE 모델의 처리량을 극대화하기 위해, 우리는 먼저 연산 사이클이 정확히 어디에서 소모되는지를 파악해야 했습니다. 표준 학습 이터레이션의 실행 타임라인을 MoE 블록 내부에서 프로파일링하자, 시스템 수준의 병목 세 가지가 두드러졌습니다.
- 활성화(Activation) 병목: 활성화 함수는 일반적으로 메모리 바운드 커널과 대규모 텐서 읽기/쓰기 연산을 유발하여, 이 구간 동안 텐서 코어가 충분히 활용되지 못한 채 남게 됩니다.
- CPU 의존성/오버헤드: 라우팅된 전문가(routed experts)에서는 전문가별 토큰 수가 런타임에 계산되며 보통 CPU에서 처리됩니다. CPU가 GPU를 따라가지 못하면 CPU 연산이 노출됩니다. 이는 CPU 동기화나 개입이 필요 없는 커널을 구축해야 할 필요성을 부각합니다.
- 양자화(Quantization) 비용: 활성화 함수와 마찬가지로, 텐서를 고정밀도에서 저정밀도로 양자화하는 작업 역시 메모리 바운드 커널을 유발해 텐서 코어를 유휴 상태로 둡니다.
우리는 CuTe DSL로 작성한 커스텀 커널로 MoE 블록을 재설계해 이러한 과제를 해결하고, 동기화 없는 MoE를 위해 작성한 세 가지 커널 계열을 도입합니다.
- GroupGemm + Quantize
- GroupGemm + Activation + Quantize/Transpose
- GroupGemm + dActivation + Quantize/Transpose
지원되는 활성화 함수는 SwiGLU, GeGLU, sReLU이며, 클램핑(clamping)과 스케일링(scaling)을 추가하는 옵션도 함께 제공합니다.
융합 GEMM 에필로그로 GLU 활성화 함수 최적화하기
게이트형 선형 함수는 최근 매우 널리 쓰이게 되었으며, 대부분의 현대 모델은 SwiGLU, GeGLU 등 게이트형 선형 유닛(GLU, Gated Linear Unit) 활성화 함수의 어떤 변형을 사용합니다. 이러한 활성화 함수는 FC1 레이어의 출력을 청크(chunk)로 나눈 뒤 이를 결합해 최종 GLU 출력을 기록합니다. 우리는 순방향 전파와 역방향 전파 모두에서 GEMM과 그에 대응하는 GLU 연산을 매끄럽게 병합하는 융합 커널을 구현합니다.
SwiGLU(x, W, V, b, cβ) = Swish_β(xW + b) ⊗ (xV + c)
GLU 활성화 함수는 GEMM의 에필로그 안에서 융합하기가 간단하지 않습니다. GLU가 텐서의 서로 다른 두 청크, 즉 입력(input)과 게이트(gate)에 접근해야 하기 때문입니다. 일반적으로 이 두 청크는 서로 다른 스레드 블록에서 계산되며, 두 출력을 결합하려면 커널이 두 출력을 모두 전역 메모리에 기록해야 합니다. 이 융합을 달성하기 위해 우리는 가중치를 입력과 게이트의 열(column)로 다시 패킹(repack)합니다. 이렇게 하면 동일한 스레드 블록이 입력 텐서의 절반 타일 폭과 게이트 텐서의 절반 타일 폭에 모두 접근할 수 있습니다. 그 결과 입력과 게이트를 전역 메모리에 다녀오지 않고도 에필로그에서 결합할 수 있습니다. 이 리패킹은 학습이 시작되기 전, 체크포인트를 로딩하는 동안 수행될 수 있습니다.
마찬가지로 역방향 전파에서는 에필로그가 GEMM 출력을 읽어 dSwiGLU를 계산하고, 이를 양자화한 뒤 전역 메모리에 다시 기록합니다.
특히 이러한 융합 패턴은 중간 텐서의 읽기·쓰기를 제거할 뿐만 아니라, 남아 있는 메모리 연산을 GEMM 자체와 직접 겹치게(overlap) 하여 활용도를 극대화합니다.
SwiGLU, GeGLU, sReLU 같은 핵심 활성화 함수를 넘어, 이 커널들은 특징 스케일링(feature scaling), 텐서 클램핑(tensor clamping), 바이어스 벡터 덧셈을 포함한 융합 에필로그 연산을 기본(native)으로 처리합니다.
호스트-디바이스 동기화와 CPU 실행 오버헤드 제거하기
전통적으로 커널이 수행하는 작업량은 실행 시점의 블록 수로 정의되며, 이는 형상(shape) 정보가 호스트에 미리 존재해야 함을 의미합니다. 예를 들어 멀티 스트림 그룹 GEMM은 G개의 서로 다른 GEMM을 별도의 스트림에서 실행하는데, 여기서 G는 그룹 수입니다. 그룹별 토큰 수가 런타임에 결정되기 때문에, CPU는 자원 활용을 극대화하기 위해 이처럼 동적으로 크기가 정해지는 GEMM들을 별도의 스트림에서 실행해야 합니다.
이는 두 가지 주요 문제로 이어집니다. 첫째, 실행해야 할 커널 수가 로컬 전문가 수에 비례해 늘어납니다. 둘째, 커널 실행 전에 호스트에서 형상 정보를 가져오기 위한 동기화 지점이 필수적으로 발생합니다. 이러한 과제를 해결하기 위해, CuTe DSL GroupGEMM 커널은 그룹별 토큰 수를 GPU 메모리 안에서 직접 추적합니다. 이는 이터레이션 동안의 CPU 의존성을 제거하고 전체 이터레이션에 걸친 CUDA 그래프를 가능하게 하여, CPU 병목을 사실상 없앱니다.
MXFP8·NVFP4 양자화를 융합해 노출된 메모리 오버헤드 줄이기
사전 학습을 위한 MXFP8, NVFP4 같은 저정밀도 레시피의 인기가 높아지고 있으며, 이러한 정밀도는 정확도에 미치는 영향을 최소화하면서도 상당한 속도 향상을 제공합니다. 이러한 저정밀도 레시피에서는 활성화 함수 뒤에 좁은 정밀도(narrow precision) GEMM 연산을 위한 양자화와 전치(transpose)가 이어집니다.
MXFP8의 경우, 양자화 커널은 활성화 함수의 출력(BF16)을 읽어 MXFP8 출력과 역방향 전파를 위한 전치 버전을 기록합니다. 새로 설계한 우리 커널은 이 양자화 단계를 GEMM 커널 자체에 융합하여, BF16 텐서의 추가적인 읽기·쓰기를 제거합니다. NVFP4의 경우에도 마찬가지로, 커널은 순방향 전파를 위해 BF16 출력과 텐서별 amax(array-maximum, 배열 최댓값)를 생성하고, 역방향 전파를 위해서는 출력의 전치 아다마르(Hadamard) 회전에 대한 amax를 계산합니다. 이는 텐서별 amax 계산을 위한 추가 메모리 패스의 필요성을 제거합니다.
커널 수준 이득에서 사전 학습 속도 향상으로
단위 수준 마이크로벤치마크 전반에서 이 융합 커널들은 상당한 속도 향상을 제공합니다. 전통적인 비융합 실행 경로 대비 순방향 패스를 최대 1.3배, 역방향 패스를 최대 2.1배 가속합니다.
이러한 속도 향상을 종단 간 학습 처리량 향상으로 이어가기 위해, 이 커널들은 다음과 같은 기능도 지원합니다.
- 동적 스케줄링(Dynamic Scheduling): 전문가 병렬(expert parallelism), 데이터 병렬(data parallelism) 등에서 발생하는 통신과 같은 다른 커널들과 효율적으로 겹치도록 지원합니다.
- 구성 가능한 클러스터 마진(Configurable Cluster Margin): 커널을 더 적은 SM으로 제한하여 사용자가 SM 자원의 마진을 구성 가능하게 예약할 수 있도록 하며, 이를 통해 다른 커널이 GPU에서 동시에 실행될 수 있는 여유를 남깁니다.
커널별 속도 향상에 더해, 이 동기화 없는 커널들은 종단 간 CUDA 그래프와 통신 커널과의 효율적인 겹침을 가능하게 하므로, 전체 애플리케이션 수준에서는 훨씬 더 큰 속도 향상이 나타납니다. 내부 테스트에서 우리는 이러한 최적화로 DeepSeek-V3에서 최대 8% 종단 간 속도 향상을, GPT-OSS 사전 학습 실행에서 최대 93% 종단 간 속도 향상을 확인했습니다.
우리는 새로운 커널을 지속적으로 추가하고 이 커널들에 새로운 기능을 더해 가고 있습니다.
CuTe DSL 융합 커널을 활용하는 방법
이 커널들은 여러 추상화 수준에서 사용할 수 있습니다.
- cuDNN Frontend(v1.23.0+): 커널들은 cuDNN Frontend 라이브러리에 들어 있습니다. 사용자는 이 라이브러리를 자신의 소프트웨어 스택에 설치하고 거기서 직접 커널을 호출할 수 있습니다. cuDNN Frontend는 이 커널들을 위한 래퍼(wrapper)도 제공하는데, 이 래퍼는 첫 호출 시 커널을 컴파일한 뒤 이후 호출에서는 캐시된 객체를 재사용합니다. 사용자는 커널을 직접 호출하거나 래퍼 API를 통해 접근하는 옵션을 가집니다. 또한 우리는 이 커널들에 대한 AOT(Ahead-of-Time, 사전) 컴파일 지원을 라이브러리에 도입하는 작업도 적극적으로 진행하고 있어, 커널을 cubin으로 컴파일해 디스크에 캐시할 수 있게 될 것입니다.
- Transformer Engine(v2.15+): 사용자는 Transformer Engine을 통해서도 이 커널들을 사용할 수 있습니다. Transformer Engine은
transformer_engine.pytorch.ops구문을 통해 이러한 연산을 노출합니다. 이러한 연산들은transformer_engine.pytorch.ops.Sequential블록을 사용해 결합할 수 있으며, 이 블록은 내부적으로 연산들을 패턴 매칭하여 cuDNN Frontend 라이브러리의 융합 커널을 호출합니다. - Megatron Core(26.04-alpha.rc2+): 사용자는 Megatron Core를 통해서도 이 커널들을 사용할 수 있으며, 알맞은 노브(knob) 집합을 지정하기만 하면 해당 기능을 호출할 수 있습니다.
다음 단계는?
우리는 더 많은 융합 패턴 지원과 JAX 같은 더 많은 프레임워크 지원 등 여러 새로운 기능을 적극적으로 개발하고 있습니다.
또한 활성화 재계산(activation recompute), 컴파일할 최적의 커널을 고르는 휴리스틱, 컴파일 비용을 줄이는 AOT(Ahead-of-Time) 컴파일, CPU 오버헤드 감소 등 다수의 커널 최적화가 진행 중입니다.
원하는 활성화 함수가 있다면, 사용자가 직접 cuDNN 커널을 수정해 PR로 기여하는 것을 권장합니다. 또는 cuDNN Frontend에 이슈를 등록해 주시면 해당 기능을 추적하겠습니다.
커뮤니티의 피드백을 언제든 환영합니다!
시작하기
GitHub의 단계를 따라 이 커널들을 실행하는 방법을 확인하세요.
