MoonshotAI/MoonEP 딥다이브
— MoE의 "쏠림"을 통신 계층에서 없애는 전문가 병렬 라이브러리

1프로젝트 한 줄 요약

먼저 문제 상황부터 그리자. 요즘 초거대 LLM은 대부분 MoE 구조다. 모델 안에 전문가(작은 신경망)가 수백 명 있고, 토큰마다 그중 몇 명(top-k)만 골라 계산한다. 전문가가 너무 많아 GPU 한 장에 다 안 들어가므로, 전문가들을 여러 GPU에 나눠서 배치한다. 이걸 전문가 병렬(Expert Parallelism, EP)이라고 부른다.

여기서 진짜 골칫거리는 쏠림(imbalance)이다. 라우터가 토큰을 전문가에게 배정할 때, 현실에서는 특정 전문가(예: "코드 담당")에게 토큰이 몰린다. 그러면 그 전문가가 사는 GPU만 일감이 폭증하고, 나머지 GPU는 놀면서 가장 바쁜 GPU가 끝날 때까지 전부 기다린다. 8장을 붙여도 1장 속도밖에 안 나는 셈이다.

즉 MoonEP는 모델을 바꾸지 않는다. 라우팅 결과(어느 토큰이 어느 전문가로 가는지)는 그대로 두고, GPU 사이에 토큰을 실어 나르는 "통신 계층"만 갈아 끼운다. 학습 프레임워크 입장에서는 dispatch → 전문가 계산 → combine이라는 익숙한 3단계를 그대로 쓰되, 그 안에서 부하가 마법처럼 평평해진다.

2왜 지금 주목받는가

2025년 DeepSeek가 DeepEP를 공개하면서 "MoE 통신 라이브러리"라는 카테고리가 사실상 표준이 됐다. MoonEP는 그 후속 경쟁자로, README가 처음부터 끝까지 DeepEP v2와의 비교로 채워져 있다. 주목 포인트는 세 가지다.

① 쏠림에 면역이라는 강한 주장

DeepEP는 부하가 몰리면(maxvio 증가) 반복 시간이 계속 올라가고, 활성 텐서 모양이 매 스텝 바뀌어 GPU 메모리가 조각나다가 결국 OOM(메모리 초과)으로 학습이 죽는다. MoonEP는 모든 GPU가 항상 정확히 S × K 토큰만 계산하므로 쏠림이 커져도 반복 시간이 평평하고, 메모리 모양이 완전히 고정이라 조각화도 OOM도 없다. H20 8장(EP=8) 실험에서 이 "평평함"을 그래프로 보여준다.

② 제로 카피(zero copy)로 통신 자체가 더 빠름

기존 방식은 토큰을 통신 버퍼로 받은 뒤, 전문가별로 다시 정렬(permute)해 사용자 버퍼로 복사한다. 이 복사가 뒷단(epilogue)의 시간을 잡아먹는다. MoonEP는 토큰을 원격 GPU의 "최종 전문가 위치"에 곧바로 써넣고, 그 통신 버퍼의 뷰(view)를 그대로 전문가 계산에 넘긴다. permute-in도 permute-out도, 버퍼 복사도 없다. 그래서 복제 전문가·가중치 프리페치 같은 추가 커널을 다 얹고도 총 dispatch 시간이 DeepEP의 dispatch만큼이고, combine은 모든 구간에서 더 빠르다고 한다.

③ 정적 모양(static shapes)으로 호스트 동기화 제거

토큰 개수가 스텝마다 바뀌면 CPU(호스트)가 매 MoE 레이어에서 "이번엔 몇 개지?"를 GPU에 물어보며 동기화해야 한다(느림). MoonEP는 버퍼 크기가 항상 고정 S × K라 모양이 컴파일 시점에 정해져, 레이어마다 호스트 동기화가 사라진다. 이게 실제 학습 처리량에 큰 영향을 준다.

덧붙여 이 레포가 MIT 라이선스에 코드가 완전 공개라는 점, 그리고 Moonshot이 방금 Kimi-K3라는 초거대 MoE를 낸 직후라는 맥락이 관심을 키웠다. "Kimi를 실제로 어떻게 그렇게 싸게 돌리는가"에 대한 실물 증거이기 때문이다.

3기술 스택 전체 지도

MoonEP는 보통의 파이썬 라이브러리처럼 from moonep import Buffer로 쓰지만, 실제 일은 전부 GPU 위에서 도는 커널이 한다. 스택은 크게 3층이다.

① 최하층 — CUDA C++ 확장 (csrc/)

bindings.cu + nvl_shared_buffer.cuh 딱 두 파일. 하는 일은 오직 하나, 여러 GPU의 메모리를 하나의 주소 공간처럼 잇는 것이다. CUDA의 저수준 VMM(Virtual Memory Management) API로 각 GPU 메모리를 POSIX 파일 디스크립터(fd)로 내보내고(nvl_dist_alloc), 다른 GPU가 그 fd를 받아 자기 주소 공간에 매핑한다(nvl_dist_map). 여기에 NVSwitch의 멀티캐스트(SHARP) 지원까지 얹는다.

② 중간층 — CuTe DSL 커널 (moonep/planning.py, dispatch.py 등)

놀랍게도 GPU 커널의 대부분이 파이썬으로 쓰여 있다. NVIDIA의 CuTe DSL(nvidia-cutlass-dsl)을 쓰기 때문이다. 파이썬 문법으로 GPU 커널을 기술하면 컴파일러가 MLIR을 거쳐 실제 CUDA 커널로 낮춘다. planning.py(1,316줄)가 가장 크고, 부하 계획·중복 전문가 선정 같은 핵심 알고리즘이 전부 여기 들어 있다.

③ 최상층 — 파이썬 API (moonep/api.py, buffer.py)

프레임워크가 실제로 부르는 얼굴. Buffer 클래스 하나에 dispatch / combine / prefetch_weight / reduce_grad 네 메서드가 걸려 있다. api.py(1,083줄)가 이들을 오케스트레이션하고, buffer.py(287줄)가 대칭 메모리 할당·정렬을 담당한다.

주목할 점: 외부 의존성이 극단적으로 적다. setup.py의 install_requires는 nvidia-cutlass-dsl==4.4.2 하나뿐이고, 나머지는 PyTorch와 CUDA 툴체인이 전부다. 프레임워크에 종속되지 않은 순수 통신 부품으로 설계됐다는 뜻이다.

4아키텍처 심화 분석 — 쏠림을 없애는 4단계

MoonEP의 한 MoE 레이어는 아래 흐름으로 돈다. 왼쪽이 순전파(forward), 오른쪽이 역전파(backward)다.

설계 패턴 1 — "먼저 계획하고, 그 계획을 재사용한다"

가장 비싼 연산(어디에 무엇을 복제할지 결정)은 순전파에서 딱 한 번 planning 커널이 계산한다. 그 결과 plan(MoonEPCommPlan)을 저장해 두고, prefetch·combine·양쪽 backward가 전부 이 plan을 그대로 재사용한다. backward에서는 planning을 다시 돌리지 않으므로 공짜다.

설계 패턴 2 — "고정 크기 버퍼 + 정적 모양"

버퍼는 항상 NvS(= S × K + 정렬 패딩) 크기로 잡힌다. 쏠림이 아무리 심해도 이 숫자는 안 변한다. 그래서 텐서 모양이 컴파일 시점에 확정되고, 메모리 조각화·호스트 동기화·OOM이 원천 차단된다. 대가는 약간의 패딩 낭비뿐이다 — MoonEP는 "예측 가능성"을 위해 이 낭비를 기꺼이 받아들인다.

설계 패턴 3 — "복제 전문가의 gradient는 반드시 본점으로 합산"

학습(training)에서는 복제 전문가가 만든 gradient가 임시라서, 프레임워크의 정상 gradient와 섞이면 안 된다. MoonEP는 복제분을 별도 reduce buffer에 모아 두고, reduce_grad에서 각 rank가 자기 전문가의 조각들을 모든 rank의 reduce buffer에서 NVLink로 읽어와(remote read) 로컬 gradient에 더한다. 그래서 학습 결과는 복제가 전혀 없던 것과 수학적으로 동일하다.

5디렉토리 구조 해부

파일 구성이 곧 파이프라인 순서다: planning → dispatch(+epilogue) → prefetch → (전문가 계산) → combine(+prologue) → grad_reduce. 코드를 읽을 때 이 순서대로 따라가면 데이터가 어떻게 흐르는지 자연스럽게 잡힌다. planning_reference.py가 커널의 "정답"을 순수 파이썬으로 다시 구현해 둔 것도 학습자에게 큰 선물이다 — 복잡한 GPU 커널이 이해 안 되면 이 참조본과 대조하면 된다.

6학습 포인트

MoonEP는 규모가 작지만(1만2천 줄) 밀도가 높다. 여기서 건질 수 있는 개념을 난이도 순으로 정리한다.

배울 것 1 — MoE 부하 불균형 문제의 실체

논문에서 "load balancing loss" 한 줄로 넘어가던 개념이, 여기서는 "가장 바쁜 rank가 전체를 잡아끈다"는 물리적 병목으로 구체화된다. maxvio 정의와 벤치마크 그래프를 보면 쏠림이 성능에 얼마나 치명적인지 수치로 감이 잡힌다.

배울 것 2 — dispatch / combine 이라는 통신 원형(primitive)

MoE 병렬화의 심장은 all-to-all 통신이다. 토큰을 흩뿌리고(dispatch) 다시 모으는(combine) 이 왕복은 분산 학습의 핵심 패턴이며, MoonEP의 API가 이를 가장 깔끔하게 보여준다. zero_copy=True 옵션으로 "복사를 없앤다"는 최적화의 실체도 코드로 확인할 수 있다.

배울 것 3 — NVLink 대칭 메모리 & VMM

csrc/는 "GPU들이 서로 메모리를 직접 읽는다"는 개념을 CUDA VMM API로 구현한 가장 짧고 명료한 예제다. cuMemGetAllocationGranularity, POSIX fd IPC, NVSwitch 멀티캐스트(SHARP)까지 실무 코드로 만난다.

배울 것 4 — CuTe DSL로 GPU 커널을 파이썬으로 쓰기

"커널 = C++"이라는 고정관념을 깨는 자료. 배리어·원자적 연산·비트 조작(_common.py)까지 파이썬 함수로 표현돼 있어, 최신 GPU 프로그래밍의 흐름을 엿볼 수 있다.

7하드웨어 / 시스템 요구사항

MoonEP는 태생이 데이터센터급 인프라다. README와 테스트가 요구하는 최소 환경은 다음과 같다.

8직접 해볼 수 있는 실습 과제

과제 A — README 벤치마크 그래프 읽고 재해석하기 난이도 ★☆☆

figure/comm_vs_deepep.png와 e2e_vs_deepep.png를 열어, x축(maxvio)이 커질 때 두 라이브러리의 곡선이 어떻게 갈라지는지 말로 설명해 보자. "왜 MoonEP는 평평한가"를 3문장으로 요약하면 핵심을 잡은 것이다.

과제 B — planning_reference.py를 CPU에서 돌려보기 난이도 ★★☆

tests/generate_topk_routing.py로 인위적으로 쏠린 라우팅을 만들고, planning_reference.py에 넣어 "어떤 전문가가 어디로 복제됐는지"를 출력해 보자. 쏠림 강도(maxvio)를 바꿔가며 복제 개수가 어떻게 변하는지 관찰한다. GPU 불필요.

과제 C — dispatch/combine의 왕복 항등성 증명하기 난이도 ★★☆

개념 실습: 임의의 [S, H] 텐서를 dispatch → (아무 계산 안 함) → combine 했을 때 원본과 같아야 한다. 이 "왕복 = 항등"을 NumPy로 흉내 내는 미니 버전을 짜 보면, permute/unpermute의 대칭성이 손에 잡힌다.

과제 D — zero_copy on/off 성능 차이 측정 난이도 ★★★

멀티 GPU 서버가 있다면 benchmarks/bench_comm.py를 zero_copy=True/False로 각각 돌려 combine epilogue의 복사 비용이 실제로 얼마나 되는지 재 보자.

과제 E — 미니 프레임워크에 MoonEP 붙이기 난이도 ★★★

작은 MoE 레이어(전문가 몇 개)를 직접 만들고 README의 API 워크스루대로 dispatch→FFN→combine과 backward(reduce_grad 포함)를 연결해, gradient가 복제 없이 학습한 것과 일치하는지 수치로 검증한다. 이 레포를 진짜로 "쓸 수 있게" 되는 관문이다.

9관련 기술 심화 학습 로드맵

10핵심 키워드 사전

11참고 링크

• MoonEP 저장소 — github.com/MoonshotAI/MoonEP (README + API 워크스루가 최고의 1차 자료)
• DeepEP (비교 기준) — github.com/deepseek-ai/DeepEP
• UltraEP — github.com/Dots-Infra/UltraEP (MoonEP가 영감으로 언급)
• Echo 논문 — arxiv.org/abs/2603.07685 (README Acknowledgments)
• CUTLASS / CuTe DSL — github.com/NVIDIA/cutlass
• CUDA VMM API 문서 — NVIDIA CUDA Driver API의 cuMemCreate/cuMemMap 계열
• Moonshot AI — huggingface.co/moonshotai (Kimi-K3 등 실제 MoE 모델)