LLM Wiki

Mixture-of-Kittens: Cursor가 새로 공개한 NVL72용 MoE 학습 메가커널

Properties2
descriptionPyTorchKR 9bow 게시글의 핵심 내용과 실무적 의미를 원문에 근거해 정리한다.
tagsai-models, open-source, research

6 min read

출처

Mixture-of-Kittens: Cursor가 공개한 NVL72용 MoE 학습 메가커널

개요

Mixture-of-Kittens(MoK)는 MoE 계층에서 발생하는 GPU 간 토큰 통신과 전문가 연산을 하나의 CUDA 메가커널로 융합한 Cursor의 오픈소스 구현입니다. NVIDIA GB300 NVL72 랙을 목표로 설계했으며, 통신을 전문가 FFN 연산과 겹쳐 통신 대기 시간을 줄이는 데 초점을 둡니다. MoK는 순전파·역전파의 디스패치와 결합 연산에 푸시와 풀 방식을 혼합해 사용하고, 미니배치 단위로 통신과 연산을 파이프라이닝합니다. Cursor는 자사 프로덕션 학습 스택에서 종단간 초당 토큰 처리량이 1.41배 향상됐다고 보고했지만, 성능은 워크로드와 학습 설정에 따라 달라질 수 있습니다.

핵심 포인트

  • MoE 통신 병목을 커널 설계에 포함: 전문가 병렬화에서는 토큰을 전문가가 위치한 GPU로 보내고 다시 돌려받아야 하므로 디스패치·전문가 FFN·결합이 병목이 될 수 있습니다. MoK는 이 세 과정을 별도 단계로 처리하지 않고 하나의 메가커널 안에서 조정합니다.
  • NVL72에 맞춘 하드웨어 특화: NVL72에서는 72개 GPU가 하나의 NVLink 도메인에 속하므로 통신과 연산을 세밀하게 겹칠 수 있습니다. 동시에 상대적으로 느린 Grace CPU의 개입과 CPU-GPU 동기화도 줄이는 것을 목표로 삼았습니다.
  • 통신 방향을 연산별로 분리: 순전파 디스패치와 역전파 역결합에는 풀 방식을, 순전파 결합과 역전파 역디스패치에는 푸시 방식을 사용합니다. 풀 기반 디스패치는 스케줄 테이블을 단순화하고 GPU 간 완료 신호를 줄이며, 원문 실험에서는 불균형 토큰 분포에서 푸시 대비 최대 29% 높은 NVLink 대역폭 활용률을 보였습니다.
  • 스케줄과 신호 비용 절감: 풀 기반 디스패치는 {src_rank, src_index} 형태의 스케줄을 사용해 정렬과 CPU-GPU 통신 없이 디바이스에서 처리합니다. Cursor는 스케줄 커널이 전체 MoE 실행 시간의 3% 미만을 차지하며, 푸시 기반 디스패치의 시그널링 지연이 풀 기반보다 약 5.8배 높았다고 보고했습니다.
  • 미니배치 크기를 튜닝 가능하게 구성: 통신을 너무 잘게 나누면 동기화가 잦아지고, 너무 크게 묶으면 초기 통신과 최종 결합을 기다리는 시간이 늘어납니다. Cursor는 전문가 그룹 GEMM이 GPU에서 최소 두 개의 완전한 웨이브를 만들도록 미니배치 크기를 정하는 휴리스틱을 제시했으며, Kimi 2.5 형태의 예시에서는 계산상 기준 2,368개에 가까운 2,560개 토큰에서 가장 좋은 결과가 측정됐습니다.
  • SM을 통신과 연산 그룹으로 분할: 일부 SM은 전문가 FFN을 담당하고, 나머지는 디스패치·결합을 담당하도록 나눕니다. TMA 로드·스토어를 활용해 통신용 SM 수를 제한하면서도 NVLink 대역폭을 활용하고, 두 그룹이 지역 카운터로 서로의 진행 상황을 알립니다.

왜 중요한가

대규모 MoE 모델에서는 전문가 연산량뿐 아니라 GPU 간 토큰 이동과 동기화가 전체 학습 시간에 큰 영향을 줄 수 있습니다. MoK는 통신 라이브러리와 연산 커널을 별도로 최적화하는 대신, 통신 방향·스케줄·미니배치·SM 배치를 MoE 계층 전체의 실행 흐름 안에서 함께 설계했다는 점에 의미가 있습니다. 특히 NVL72처럼 GPU 간 연결성이 높은 시스템에서는 어떤 통신 방식을 선택하고 어느 크기로 연산과 겹칠지에 따라 성능이 크게 달라질 수 있음을 구체적인 프로파일링과 휴리스틱으로 설명합니다.

후속으로 볼 링크와 키워드

관련 위키

원문 보존 위치

원문 전체는 raw source: 2026-08-07-pytorchkr-topic-11548-mixture-of-kittens-cursor-nvl72-moe에 source_url 및 HTML 원문과 함께 저장되어 있습니다.