Mixture-of-Kittens: NVL72를 위한 오픈소스 MoE 메가커널

Cursor가 에이전트 코딩 모델 Composer의 학습에 사용해 온 **Mixture-of-Kittens(MoK)**을 오픈소스로 공개했다. MoK는 NVL72 환경을 위해 설계된 MoE 학습 메가커널로, MoE의 통신과 연산을 하나의 결정론적 커널에 통합한다.

MoE 계층은 Composer 학습 시간의 절반 이상을 차지할 수 있는 주요 병목이었다. 기존 구현은 GPU 간 토큰 통신과 전문가 FFN 연산을 별도로 처리했지만, 실제 운영 환경에서는 통신 비용이 연산 비용만큼 커지면서 전체 성능을 제한했다.

GB300 NVL72는 72개 GPU가 하나의 NVLink 도메인에 연결된 멀티노드 랙이다. MoK는 이 구조를 활용해 통신과 연산을 세밀하게 중첩하고, 통합 Grace CPU가 GPU보다 느려 발생하는 CPU 작업 및 CPU-GPU 동기화 비용도 줄인다.

MoK가 해결하는 DSV3 스타일 MoE 계층은 하나의 공유 전문가와 다수의 라우팅 전문가로 구성된다. 라우터가 각 토큰을 상위 k개 전문가에 배정하면, 토큰은 전문가가 배치된 GPU로 이동해 FFN 연산을 거친 뒤 원래 GPU로 결과가 돌아온다.

주요 설계 방향은 다음과 같다.

• 모든 MoE 통신과 연산을 단일 메가커널로 통합
• 통신과 FFN 계산의 중첩을 통한 GPU 활용률 향상
• MXFP8 학습 지원
• 결정론적 실행 보장
• 링 토큰 버퍼를 활용한 CPU-GPU 동기화 최소화

공개 구현체 기준 비교에서 MoK는 GB300 NVL72의 MXFP8 forward 처리량을 최대 2.37배 높였다. 여러 NVL72 랙을 사용하는 실제 학습 스택에서는 전체 초당 처리 토큰 수가 1.41배 증가했다.