커서, MoE 커널 MoK 공개…처리량 2.37배 향상됐다

작성일

커서, MoE 훈련 최대 2.37배 빠른 메가커널 ‘MoK’ 오픈소스 공개
블랙웰 NVL72 랙 전용…GPU 수만 대로 컴포저 모델 훈련에 이미 적용

커서, MoE 커널 MoK 공개…처리량 2.37배 향상됐다 / AI 생성 일러스트(삽화)
커서, MoE 커널 MoK 공개…처리량 2.37배 향상됐다 / AI 생성 일러스트(삽화)

AI 코딩 도구로 알려진 커서(Cursor)의 연구 조직 커서 리서치가 자사 에이전트형 코딩 모델 ‘컴포저(Composer)’ 훈련에 쓰이는 MoE(전문가 혼합, Mixture-of-Experts) 훈련용 메가커널 ‘믹스처 오브 키튼즈(Mixture-of-Kittens, MoK)’를 오픈소스로 공개했다. MoK는 MoE 계층에서 일어나는 모든 통신과 연산 과정을 하나의 완전한 결정론적(deterministic) 커널로 통합한 것이 특징이다. 커서 측은 공개된 가장 빠른 퍼블릭 베이스라인 대비 최대 2.37배 높은 처리량을 냈다고 밝혔다. MoK는 이미 수만 대 규모의 GPU에서 컴포저 모델 훈련을 구동하고 있다. 다만 엔비디아 블랙웰(Blackwell) 계열 고성능 GPU 랙이 있어야만 쓸 수 있어 실제로 도입 가능한 조직은 제한적이다.

MoE 계층이 훈련의 병목이 된 배경

커서 리서치는 이전에도 MoE 연산 속도를 높이는 작업을 해왔다. 자체적으로 MXFP8·NVFP4 훈련 커널을 개발했고, MoE 추론을 위한 ‘워프 디코드(warp decode)’ 방식도 만들었다. 하지만 이 기술들은 연산 부분만 최적화했고, GPU 간 통신은 별도로 처리된다고 가정했다.

실제 서비스 환경에서는 통신이 오히려 속도를 좌우하는 요인이 됐다. 워크로드와 훈련 설정에 따라 MoE 계층이 전체 훈련 시간의 절반 이상을 차지하는 경우도 있었다. 여기에 GB300 NVL72 랙으로 이전하면서 문제가 새로운 양상을 띠었다. NVL72 랙은 72개 GPU가 하나의 NVLink 도메인 안에 묶여 있어 연산과 통신을 세밀하게 겹쳐 처리할 수 있다. 그런데 랙에 내장된 그레이스(Grace) CPU는 GPU보다 처리 속도가 느리다. GPU 작업 스트림이 CPU 작업을 금세 따라잡아버려 그 사이 GPU가 그대로 놀게 되는 문제가 생겼다. 이 때문에 CPU 작업과 CPU-GPU 간 동기화를 최대한 줄여야 했다.

통신 방향과 오버랩 방식을 재설계한 세 가지 핵심 아이디어 / AI 생성 이미지
통신 방향과 오버랩 방식을 재설계한 세 가지 핵심 아이디어 / AI 생성 이미지

통신 방향과 오버랩 방식을 재설계한 세 가지 핵심 아이디어

MoK 설계의 첫 번째 축은 통신 방향을 연산 단계별로 다르게 정한 것이다. 딥EP(DeepEP) 같은 기존 방식은 대체로 푸시(push) 기반 전송을 쓴다. 커서의 자체 벤치마크에 따르면 푸시 방식은 한 방향으로만 많은 바이트를 옮기다 보니 반대쪽 NVLink 레인이 거의 비어 있게 된다. 풀(pull) 기반 디스패치를 쓰면 전문가 배분이 불균형한 상황에서 NVLink 대역폭 활용률이 최대 29% 높아졌다. GPU 간 완료 신호를 주고받을 필요도 없어졌다. 실제로 푸시 방식 디스패치 신호 처리 시간은 103마이크로초였지만 풀 방식은 18마이크로초로, 약 5.8배 차이가 났다. 이에 따라 MoK는 순전파에서 풀 기반 디스패치와 푸시 기반 컴바인을 쓴다. 역전파는 이를 거꾸로 적용해 풀 기반 역방향 컴바인과 푸시 기반 역방향 디스패치를 쓴다. 이 네 가지 연산 모두를 하나의 스케줄 테이블로 처리하는데, 여기 드는 비용은 전체 MoE 실행 시간의 3% 미만이다.

두 번째 축은 연산과 통신을 겹치는 단위(granularity)를 조정한 것이다. 기존 기법인 코멧(Comet)은 세밀한 단위로, 딥EP는 큼직한 단위로 오버랩을 처리한다. 커서 리서치는 최적점이 그 중간 어딘가에 있고, 워크로드에 따라 달라진다고 설명한다. MoK가 쓰는 휴리스틱은 전문가별 그룹 GEMM(행렬 곱) 연산마다 최소 두 번의 완전한 SM(스트리밍 멀티프로세서) 웨이브를 확보하는 것을 목표로 한다. 컴포저 2.5의 기반 모델인 키미(Kimi) 2.5 형태의 경우 이 기준을 만족하는 최소 토큰 수는 2,368개였고, 실측 지연 시간도 이 추정치에 근접했다.

세 번째 축은 링 형태의 토큰 버퍼로 CPU를 연산 루프에서 완전히 빼낸 것이다. 기존 대안은 토큰을 일부 누락시키거나 CPU가 버퍼 크기를 매번 계산하도록 하는 방식이었다. MoK는 대신 수백 메가바이트 규모의 고정된 링 버퍼를 순환시킨다. 미니배치 단위로 작동하며 매크로배치 경계에서 디스패치와 컴바인을 서로 겹쳐 처리한다. 역전파 때는 순전파 활성값을 다시 계산하는 부담을 줄이기 위해 링을 역방향으로 순회한다.

벤치마크 결과와 도입 가능한 조직

성능 측정은 단일 NVL72 랙에서 EP(전문가 병렬화) 차수 64로 진행했다. 라우팅 전 각 GPU는 2,048개 토큰을 보유한 상태였다. 비교 대상 베이스라인은 NCCL+PyTorch, 딥EP+PyTorch, 딥EP+TransformerEngine, HybridEP+Megatron이었고, 테스트에는 키미 K2.7 코드, GLM-5.2, 큐원3.5-397B-A17B, 딥시크-V4-프로 형태의 모델 구조가 쓰였다.

가장 빠른 베이스라인과 비교했을 때 MoK는 MXFP8 순전파에서 최대 2.37배, MXFP8 역전파에서 1.78배 빨랐다. BF16 기준으로는 순전파 1.92배, 역전파 1.58배였다. 512개 GPU를 여러 GB300 NVL72 랙에 걸쳐 사용한 종단 간(end-to-end) 테스트에서는 GPU당 초당 처리 토큰 수가 760.9개에서 1,070.2개로 늘어 1.41배(약 41%) 증가했다.

다만 MoK를 실제로 쓸 수 있는 문턱은 상당히 높다. GitHub에 아파치 2.0(Apache-2.0) 라이선스로 공개됐지만, 엔비디아 블랙웰 SM100 또는 SM103 GPU가 필요하다. 이는 곧 GB200 NVL72나 GB300 NVL72 랙을 갖춰야 한다는 뜻이다. 파이썬 3.12 이상, PyTorch 2.10 이상, CUDA 툴킷 13.0 이상도 필요하고 GPU 간 버퍼는 PyTorch의 심메트릭 메모리(symmetric memory) 기능에 의존한다. 결과적으로 실질적인 도입 대상은 NVL72급 컴퓨팅 자원을 보유하거나 임대할 수 있는 프런티어 AI 연구소, 자금력 있는 모델 스타트업, GPU 네오클라우드, 국가급 컴퓨팅센터 정도로 좁혀진다. 단일 노드나 8개 GPU 수준의 소규모 팀은 사실상 해당되지 않는다.

활용 분야와 남은 과제

MoK가 겨냥하는 영역은 좁지만 부가가치가 크다. 딥시크-V3 스타일의 MoE 모델 사전학습·후속학습이 대표적이고, 결정론적 특성 덕분에 온폴리시(on-policy) 강화학습 기반 후속학습이나 내부 실험(ablation)에도 유용하다. 관련 산업으로는 AI 모델 개발, 클라우드 GPU 인프라, 코드 생성 도구, 퀀트 연구 분야가 꼽힌다.

기술적으로 MoK는 완전한 메가커널 구조로 설계됐고 완전한 결정론성을 갖췄다. BF16과 MXFP8 정밀도 모드를 모두 지원한다. 스케줄링은 블랙웰의 클러스터 런치 컨트롤(Cluster Launch Control) 기능을 거쳐 이뤄지기 때문에 랙 간 RDMA 통신이 이 과정에 밀려 대기하지 않는다. 라우터 가중치의 그래디언트 계산은 소닉MoE(SonicMoE) 방식을 SwiGLU 역전파 안에 통합해 처리한다.

하드웨어 문턱이 높은 만큼 MoK의 확산 속도는 결국 블랙웰 기반 NVL72 랙의 보급 속도에 달려 있다. 그럼에도 커서 리서치가 밝힌 처리량 개선폭과 오픈소스 공개 자체는 MoE 아키텍처를 대규모로 훈련하는 연구소들에게 참고할 만한 설계 사례를 제공한다는 점에서 주목받고 있다.