CTRL K

SonicMoE로 보는 GPU 커널 최적화

양원석,2026년 8월 19일•29 min read

Introduction

여러분은 머신 러닝에서 무엇이 중요하다고 생각하시나요? 질 좋은 데이터를 확보하는 것이 떠오를 수 있고, 데이터로부터 더 많은 시그널을 뽑아낼 수 있는 새로운 모델링 기법이나 학습 알고리즘이 생각날 수도 있겠습니다.

실제로 AI 연구에서는 모델 구조나 데이터 등이 가장 많은 관심을 받습니다. 그런데 트릴리온 랩스에서 여러 차례 LLM Pretraining을 진행하면서, 의외인 곳에서 가장 큰 병목을 마주하게 되었습니다. 같은 데이터나 모델 구조여도 모델을 GPU 위에서 얼마나 효율적으로 학습시키느냐가 전체 개발 속도와 비용을 결정하는 경우가 훨씬 많았습니다.

GPU는 AI 연구에서 가장 비싼 자원입니다. 학습 속도를 몇 퍼센트만 개선해도 수백 시간의 GPU 사용 시간을 절약할 수 있고, 이는 곧 개발 비용과 연구 속도의 차이로 이어집니다. 따라서 최근에는 모델 자체뿐 아니라 GPU 커널과 시스템 최적화가 AI 연구의 중요한 경쟁력이 되고 있습니다. GPU 커널 최적화는 머신러닝 시스템을 구성하는 중요한 요소이지만, 데이터나 모델링 등에 비해 상대적으로 주목받지 못했던 분야이기도 합니다.

이번 글에서는 SonicMoE를 예시로, 효율적인 GPU 커널을 설계한다는 것이 무엇을 의미하는지 이야기 해보려고 합니다.

초심자를 위한 GPU 커널

먼저 GPU 커널이 무엇인지 알아보겠습니다. GPU 커널은 GPU 위에서 실행되는 하나의 함수입니다. CPU 프로그램이 여러 함수를 호출해 작업을 수행하듯, 딥러닝 모델 역시 행렬 곱셈, 활성화 함수, 정규화 등 대부분의 연산이 각각 GPU 커널 형태로 실행됩니다.

모든 GPU 커널은 여타 함수와 마찬가지로 메모리에서 데이터를 읽고 (Prologue), 계산을 수행한 뒤 (Mainloop), 다시 메모리에 결과를 저장하는 (Epilogue) 과정으로 요약될 수 있습니다. 즉, 하나의 커널을 연산과 메모리 접근 두 요소로 나누어 생각할 수 있습니다. 어떤 커널은 계산량이 많아 연산 자체가 병목이 되고, 반대로 어떤 커널은 계산보다 메모리 접근에 더 많은 시간을 사용합니다. 따라서 어떤 부분이 병목인지에 따라 최적화 전략도 달라집니다.

실험 결과 시각화 — GPU 커널은 프롤로그, 메인루프, 에필로그로 나뉩니다.

뿐만 아니라 GPU 커널이 학습 목적인지 추론 목적인지에 따라 최적화 방향이 달라집니다. 학습과 추론은 이질적인 워크로드이기 때문에 최적화 방향도 차이를 반영해야 합니다.

학습과 추론과 가장 큰 차이는 역전파를 수행해야 한다는 점입니다. 역전파를 위해서는 순전파 과정에서 생성된 activation을 메모리에 보관해야 하는데, 대규모 언어 모델에서는 이 activation이 모델 파라미터보다도 훨씬 큰 메모리를 차지하는 경우가 많습니다. 결국 activation은 학습 시 OOM의 가장 큰 원인 중 하나가 되며, 이를 줄이기 위해 필요한 값만 저장하고 나머지는 필요할 때 다시 계산하는 activation checkpointing 기법이 널리 사용됩니다. 다시 계산하는 비용은 다소 증가하지만, 메모리 사용량을 크게 줄일 수 있기 때문에 현대 LLM 학습에서는 매우 중요한 최적화 기법입니다.

Activation이 어떻게 학습에 병목이 되는지, 그리고 이를 해결하기 위한 activation checkpointing이 무엇인지는 이전 포스트 들의 중요한 논점이기도 하기 때문에 같이 보시면 좋을 것 같습니다 🙂

이처럼 GPU 커널을 최적화한다는 것은 단순히 연산을 빠르게 만드는 것만을 의미하지 않습니다. 메모리 사용량을 줄이고, 데이터 이동을 최소화하며, 필요한 계산만 효율적으로 수행하도록 만드는 것까지 모두 포함됩니다. 이러한 관점은 특히 MoE처럼 메모리 접근과 GPU 간 데이터 이동이 단순하지 않은 구조에서 더욱 중요해집니다.

Interlude: Mixture of Experts

최근 대규모 언어 모델의 발전에서 가장 눈에 띄는 변화 중 하나는 Mixture of Experts(MoE) 구조의 확산입니다. MoE는 토큰마다 여러 전문가(Expert) 중 일부 전문가만 선택적으로 활성화하여 연산을 수행하는 구조입니다.

계산 그래프 다시 쓰기 시각화 — 각 단계가 별도 커널이 되면서 단계마다 중간 텐서가 남습니다.

MoE의 가장 큰 장점은 모델의 총 파라미터 수는 크게 늘리면서도, 실제 연산량은 상대적으로 작게 유지할 수 있다는 점입니다. 덕분에 동일한 FLOPs로도 성능 좋은 모델을 확보할 수 있기 때문에 Deepseek, GLM, Qwen, Kimi, GPT-OSS, Inkling 등 동서양을 막론하고 모든 프론티어 급 LLM은 이러한 구조를 적극적으로 채택하고 있습니다.

MoE 구조는 지난 몇 년간 빠르게 변화했습니다. 그중 가장 뚜렷한 변화는 granularity (전문가의 은닉 차원 크기 대비 LLM의 은닉 차원 크기) 와 sparsity (활성 전문가 수 대비 실제 전문가 수)의 증가입니다. 초창기 Mixtral 8×22B나 OLMoE 7B와 같은 모델은 소수의 큰 전문가 네트워크를 사용하는 구조가 일반적이었습니다. 이후 DeepSeek이 Fine-grained MoE 구조를 제안하면서, 최근 모델들은 소수의 큰 전문가를 사용하는 대신 다수의 작은 전문가를 사용하는 방향으로 발전하고 있습니다. 이러한 구조는 동일한 활성 파라미터 수에서도 더 높은 모델 용량과 표현력을 확보할 수 있기 때문입니다.

GEMM Overlapping 시각화 — 전문가를 잘게 쪼개면 활성 전문가 수 K가 늘어납니다. 출처: [6]

하지만 이러한 변화는 시스템 관점에서는 새로운 문제를 만들어냅니다.

이를 이해하기 위해 먼저 arithmetic intensity(산술 집약도)라는 개념을 알아보겠습니다. 앞서 말씀드린 것처럼 GPU 커널은 데이터 읽고 쓰는 IO 작업과 실제 연산으로 이루어져 있습니다. Arithmetic intensity란 연산 대비 메모리에 읽고 쓰는 바이트양으로 정의됩니다.

\text{산술 집약도} = \frac{\text{연산량 (FLOPs)}}{\text{메모리량 (Bytes)}}

산술 집약도는 커널의 성격을 정의하는 간단하면서도 효과적인 척도입니다. 산술 집약도가 크다는 것은 읽고 쓰는 데이터 양 대비 연산량이 많다는 것을 의미합니다. 이러한 커널은 연산 처리 속도가 전체 성능을 제한하는 compute-bound 특성을 보이는 경우가 많습니다. 반대로 산술 집약도가 작다는 것은 데이터 이동량 대비 연산량이 적다는 것을 의미합니다. 이러한 커널은 메모리 접근 속도가 전체 성능을 제한하는 memory-bound 특성을 보이는 경우가 많습니다.

그럼 MoE의 산술 집약도는 granularity와 sparsity에 따라 어떻게 변하게 될까요? 이를 위해 산술 집약도를 실제 계산해보면 아래와 같습니다.

\text{MoE 산술 집약도} = \frac{3}{\frac{2+2G}{d}+\frac{3s}{T}}

아래 식에서 $G=\frac{d}{n}$ 는 MoE의 granularity, $s=\frac{|E|}{K}$ 는 sparsity를 의미합니다. 즉, granularity와 sparsity가 증가하면 MoE의 산술 집약도가 감소하게 됩니다. 최근 모델들은 fine-grained MoE 구조를 택하면서 granularity와 sparsity가 함께 커지는 방향으로 가고 있고, 그 결과 산술 집약도는 지날수록 오히려 낮아지고 있습니다.

Memory Bound Kernel 최적화하기

Memory bound 커널을 어떻게 최적화 할 수 있을까요? 단순하게 생각하면 아래 방식들을 떠올릴 수 있습니다.

IO를 없앤다: 애초에 읽고 쓰지 않는다.
IO를 숨긴다: 계산을 하는 동안 데이터를 읽고 쓴다.

IO 없애기

IO 없애는 대표적인 방식으로 kernel fusion이 있습니다. Kernel fusion을 이해하기 위해 GPU 커널에 대한 이해도를 더 높힐 필요가 있습니다. 일반적으로 하나의 커널은 여러 서브 커널들로 이루어져 있습니다. 각 서브 커널도 결국 또 하나의 커널이기에 각 서브 커널의 경계마다 데이터를 메모리로부터 읽고 연산 결과물을 메모리에 다시 써야 합니다. 우리가 원하는 것은 원래 커널의 결과물 뿐인데 말이죠.

MoE 커널 또한 이 문제에서 자유롭지 않습니다. 결국 우리가 원하는 것은 활성화된 전문가들의 종합 결과물인데, 보통 그 과정 속에서 수많은 중간 결과물을 읽고 써야 합니다. 이런 모든 중간 IO는 순수 오버헤드입니다.

Kernel fusion은 커널 경계를 지워서 메모리 IO 왕복을 없애는 기법입니다. 여러 세부 커널을 하나의 커널로 결합하여, 중간 결과물을 메모리에 저장하고 다시 읽어오는 과정을 줄입니다.

학습 커널의 경우 kernel fusion으로 해결되지 않은 경우가 있습니다. 학습의 경우 순전파에서 계산된 값들을 역전파에서 사용하는 경우가 있습니다. 순전파와 역전파는 다른 시점에 실행되기 때문에 서로 합칠 방법은 없습니다.

이 때 사용되는 방식이 앞서 소개한 activation checkpointing입니다. 즉, kernel fusion으로 인해 저장하지 않은 값을 역전파에서 다시 계산하는 것이지요. 중간 결과물을 순전파 때 저장하고 역전파 때 메모리에서 불러내는 비용보다 다시 계산하는 비용이 더 저렴할 수록 activation checkpointing을 사용할 유인이 더 큽니다.

Kernel fusion을 더 극단적으로 밀어붙이면 아예 계산 그래프를 재설계 하는 것이 있을 것입니다.

지금까지 이야기한 두 기법은 모두 계산 그래프를 주어진 것으로 놓고 그 위에서 IO를 줄입니다. Kernel fusion은 노드 사이의 경계를 지우고, activation checkpointing은 저장을 재계산으로 바꿉니다. 하지만 어떤 중간 결과물이 존재해야 하는지 자체는 그대로입니다.

$A(BC)$ 와 $(AB)C$ 가¹ 결과적으로는 같은 값이면서도 중간 텐서 크기는 전혀 다른 것처럼, 역전파 수식에도 결과는 유지한 채 중간 결과물만 바꾸는 자유도가 남아있는 경우가 많습니다. 그런 경로를 찾으면 기존에 저장해야 할 큰 텐서는 사라지고, 다루기 쉬운 작은 텐서만 저장해도 되는 상황이 될 수 있습니다.

IO 숨기기

다음 전략으로는 계산과 동시에 데이터 IO를 하여 IO가 끝날 때가지 기다리는 시간을 줄이고 GPU 계산 역량을 최대한 짜내는 방식을 취할 수 있을 겁니다.

이 전략의 핵심은 비동기 실행입니다. Hopper 세대 이후 GPU마다 비동기 메모리 instruction에 대한 지원을 강화해온 이유가 여기에 있습니다. 일반적으로 비동기 IO를 위해 사용되는 기법들을 아래에 정리해보았습니다.

Software Pipelining - 하나의 커다란 행렬 계산을 여러 서브 행렬 (이하 타일)로 쪼갭니다. 다음 타일을 미리 불러오면서 현재 타일을 계산합니다.
Warp Specialization - 스레드 그룹²의 역할을 나눕니다. 일부는 데이터를 불러오는 일만 하고(producer), 일부는 계산만 합니다(consumer).

지금까지 memory bound 커널을 IO 효율적으로 만들기 위한 방법들에 대해서 알아보았습니다. 물론 이런 방법들이 꼭 memory bound 커널에만 적용되는 것은 아닙니다. 하지만, memory bound 커널에서는 효율적인 IO가 가능하도록 커널의 세부 동작을 튜닝하는 것이 계산 효율성을 높히는 방식으로 튜닝하는 것보다 더 효과적인 전략일 것입니다. 반대로, compute bound 커널에 대해서는 계산을 빠르게 하거나 아님 계산량을 줄이는 방식이 더 효과적일 것이구요. 중요한 것은 내가 최적화 하는 커널이 어떤 성격을 가지는지 이론적 분석과 코드로 구현하여 프로파일링 해보고, 커널의 성격에 맞는 전략을 선택하는 것입니다.

SonicMoE로 이해하는 Memory Bound Kernel 최적화

SonicMoE는 NVIDIA Hopper와 Blackwell GPU를 대상으로 만들어진 학습용 MoE 커널입니다. Fine-grained MoE 학습에서 발생하는 activation 메모리 증가와 IO 병목을 함께 해결하는 것을 목표로 합니다.

SonicMoE의 설계 원칙은 한 문장으로 요약됩니다.

O(TKd)

크기를 가지는 텐서는 저장하지 않는다.

토큰 수 $T$ 와 은닉 차원 $d$ 를 고정한 채 계산 FLOPs를 유지하기 위해서는 전문가 중간 차원과 활성 전문가 수의 곱 $nK$ 을 유지해야 하는데, granularity를 높히기 위해서는 전문가 중간 차원을 줄이고 활성 전문가 수를 늘린다는 것을 의미합니다. 결국 크기가 $O(TKd)$ 인 텐서는 granularity에 비례해서 커지게 됩니다.

MoE 내부를 살펴보면 $O(TKd)$ 인 텐서가 여러개 존재합니다. 일반적으로 MoE를 $\text{gather}$ , $\text{up-projection}$ , $\text{SwiGLU}$ , $\text{down-projection}$ , $\text{scatter}$ , $\text{expert aggregation}$ 총 6개의 서브커널로 나눌 수 있습니다. 이중 $\text{gather}$ 된 $X$ , $\text{down-projection}$ 의 결과물인 $Y$ , $\text{scatter}$ 의 결과물인 $\tilde{Y}$ 는 각각 $2TKd$ 바이트를 차지하면서 역전파를 위해 메모리에 저장되게 됩니다.

이런 텐서들을 최대한 없애거나, 없앨 수 없다면 최대한 효율적으로 계산되도록 해야 할 것입니다. 위에 설명드린 테크닉 적용해서 어떻게 이 문제를 해결했는지 보겠습니다.

Gather Fusion

MoE에서 $\text{gather}$ 는 각 전문가가 자신에게 배정된 토큰들의 은닉 상태를 모아오는 작업입니다. 별도 커널로 처리하면 $2TKd$ 바이트를 읽어서 다시 그만큼 써야 합니다. 이 과정은 은닉 상태를 추후 개별 전문가들로 계산하기 쉽게 재배열하는 과정으로 실질적인 계산은 아예 없는, 순수 IO 오버헤드 커널입니다.

SonicMoE은 이 $\text{gather}$ 를 뒤따르는 $\text{up-projection}$ 의 데이터를 메모리에서 가져오는 과정과 융합합니다. $\text{up-projection}$ 에서 데이터를 불러올 때 라우팅된 인덱스를 따라 흩어진 행을 바로 읽어오게 하면 재배열된 텐서를 만들어 저장할 필요가 없을 것입니다. 즉, 이 kernel fusion을 통해서 $\text{gather}$ 의 결과물이 쓰이고 다시 읽히는 과정을 없앱니다.

여기서 한가지 의문이 들 수 있습니다. 결국 IO 해야 할 데이터 총량은 동일한 것이 아닌가요? 한번에 $2TKd$ 바이트를 가져오는 거나 아님 $2Td$ 바이트를 $K$ 번 가져오는 거나 결과적으로 동일하지 않나라는 의문입니다.

수식상으로는 맞습니다. 하지만, 하드웨어 관점에서 이 두 과정은 아예 다릅니다. $\text{gather}$ 의 결과물은 $\text{gather}$ 전의 텐서보다 $K$ 배 큰 텐서입니다. granularity가 올라가면 $K$ 가 커지고, 이 경우 $\text{gather}$ 의 결과물은 GPU의 L2 캐시의 최대 용량을 넘어가게 됩니다. 캐시 미스가 나면서 GPU는 빠른 L2 캐시에서 데이터를 읽어오지 못하고 상대적으로 느린 HBM에서 데이터를 가져와야 할 것입니다 (참고로, Blackwell에서 L2 캐시와 HBM의 최대 IO bandwidth는 약 20TB/s와 7TB/s로 3배 이상 차이가 납니다). 커널을 재설계할때 수식 상의 알고리즘 이상으로 하드웨어에 대한 이해가 바탕이 되어야 한다는 것을 의미합니다.

추가적으로, $\text{gather}$ 의 결과는 역전파 시 gradient계산에 필요합니다. SonicMoE에서는 이 결과물들을 재계산하는 방식을 통해 IO 병목을 최대한 없앱니다. 앞서 이야기 했던 activation checkpointing을 적용했다고 볼 수 있습니다.

역전파 계산 그래프 재작성하기

역전파 알고리즘에서 라우터 기울기 $dS$ 는 $dS_{t,e}=\langle dO_t,Y_{e,t}\rangle$ 로 계산되는데, 이 식은 $\text{down-projection}$ 의 결과물인 $Y$ 를 캐시하도록 강제합니다. SonicMoE에서는 이 계산 과정을 바꾸어 $Y$ 를 저장할 필요 없도록 만듭니다.

$Y_e=A_eW_{2,e}$ 를 대입하면

\begin{aligned} dS_{t,e} &= \langle dO_t,Y_{e,t}\rangle = \langle dO_t,A_eW_{2,e}\rangle = \langle dO_tW_{2,e}^{T},A_e\rangle \end{aligned}

가 됩니다.

두 식이 같은 값을 주더라도 메모리로부터 읽어야 하는 데이터 양이 다릅니다. 기존 식에서 사용되는 $Y$ 는 $2TKd$ 바이트인 반면, 새로운 식에서 사용되는 $dA'$ 과 $A$ 는 $2TKn$ 바이트입니다. 같은 값을 구하는데 대략 granularity $G$ 배 적은 데이터를 읽는 셈입니다.

뿐만 아닙니다. $dA'$ 은 $dH$ 를 구하기 위해 어차피 계산되는 값이고 $A$ 는 $H$ 로부터 얻을 수 있습니다 (하단 SonicMoE 역전파 참고). 둘 다 $dS$ 때문에 새로 생기는 텐서가 아닙니다. 반대로 $Y$ 는 오직 $dS$ 를 위해 다시 계산해야 했고 역전파에서 $dY$ 까지 함께 만들어졌습니다. 계산 그래프에서 노드 하나를 지우자 순전파와 역전파 양쪽의 IO가 동시에 사라진 것입니다. 계산 그래프를 바꾸는 것이 가장 어려우면서도 가장 강력한 방식인 이유입니다.

GEMM Overlapping

$dH$ 커널은 epilogue에서 $H$ 를 다시 불러오고 여러 activation과 reduction연산을 수행한 뒤 $dH$ , $dS$ , $A'$ 3개의 텐서를 저장하기 때문에 IO 부담이 큽니다. 이 병목을 그대로 두면 데이터를 쓰는 동안 GPU가 대기하게 됩니다.

여기서 쓰이는 것이 warp specialization입니다. 스레드 그룹의 역할을 나눠서, 일부는 다음 타일을 불러오고 일부는 행렬 연산을 수행하며 또 일부는 epilogue만 담당하게 만듭니다. 각 그룹이 서로를 기다리지 않고 자기 일을 계속하기 때문에, 한 그룹이 무거운 epilogue를 처리하는 동안에도 Tensor Core는 다음 타일의 행렬 연산을 돌릴 수 있습니다.

Qwen3-235B에서 $dH$ 커널의 epilogue는 HBM 트래픽을 24% 늘리지만, Tensor Core 활용률은 98%에서 88%로, TFLOPs는 1213에서 1078로 11%만 떨어집니다. IO 비용 증가가 TFLOPs 감소로 반비례하게 옮겨가지 않는다는 것, 이것이 IO를 숨긴다는 말의 의미입니다.

실험 결과

Blackwell GPU에서 7B부터 685B까지 6개의 실제 MoE 설정으로 벤치마크했을 때, SonicMoE는 모든 설정에서 가장 빨랐습니다. 기존 라이브러리를 조합해 쓰는 경우와 비교하면 순전파 54%, 역전파 35% 높은 처리량이 나왔고, 기존 MoE 학습 커널인 ScatterMoE·MoMoE 대비로는 최대 4배까지 차이가 났습니다.

메모리도 함께 줄었습니다. Hopper에서 fine-grained 7B MoE 기준으로 activation 메모리는 45% 감소하면서 처리량은 1.86배가 되었습니다. 보통 메모리와 속도는 맞바꾸는 관계인데, $O(TKd)$ 텐서를 없앤 것이 저장할 양과 옮길 양을 동시에 줄였기 때문에 둘 다 얻을 수 있었습니다.

가장 와닿는 숫자는 실제 학습 처리량입니다. H100 64장에서 하루 213B 토큰을 처리했는데, 같은 모델을 ScatterMoE로 학습하면 H100 96장에서 하루 225B 토큰입니다. GPU를 2/3배 덜 사용하고도 거의 같은 속도를 낸 셈입니다.

트릴리온 랩스에서도 이번 Gravity 16B를 학습했을 때 SonicMoE를 사용했습니다. 기존 pytorch native한 방식으로 구현한 MoE 커널 대비 2.1배 학습 속도 향상이 있었습니다.

Conclusion

SonicMoE의 핵심은 단순히 연산을 더 빠르게 구현하는 것이 아니라, MoE가 어떤 방식으로 GPU의 메모리와 연산 자원을 사용하는지 이해하고 그에 맞게 계산 자체를 재구성하는 것에 있습니다. Kernel fusion으로 불필요한 IO를 없애고, 계산 그래프를 바꾸어 큰 중간 텐서를 제거하며, warp specialization으로 남은 IO를 계산과 병행하도록 하였습니다. 그 결과 메모리 사용량과 학습 속도를 동시에 개선할 수 있었습니다.

결국 GPU 커널 최적화에서 중요한 것은 병목이 어디에서 발생하는지 이해하고 가장 큰 병목 지점들을 효과적으로 해결하는 테크닉을 적용하는 것 입니다. 모델이 점점 더 복잡해질수록 좋은 알고리즘을 만드는 것만큼, 그 알고리즘이 실제 하드웨어 위에서 효율적으로 동작하도록 만드는 시스템 최적화의 중요성도 더욱 커질 것입니다.

Reference

행렬의 shape을 $A\in\mathbb{R}^{m\times n}$ , $B\in\mathbb{R}^{n\times k}$ , $C\in\mathbb{R}^{k\times p}$ 라고 하면, $A(BC)$ 에서는 먼저 $BC\in\mathbb{R}^{n\times p}$ 를 만들고 $(AB)C$ 에서는 먼저 $AB\in\mathbb{R}^{m\times k}$ 를 만듭니다. 두 결과는 같지만 $n\times p$ 와 $m\times k$ 중 어느 쪽을 중간 텐서로 저장하느냐에 따라 메모리 사용량이 달라집니다. ↩
GPU에서 여러 스레드를 묶어 함께 실행하고 작업을 분담하는 단위입니다. ↩