커서, 블랙웰 서버 전용 MoE 커널 'MoK' 공개…학습 처리량 41% 향상

대규모 AI 모델의 학습 속도를 혁신할 새로운 기술이 등장했다. 커서(Cursor)는 4일(현지시간) 엔비디아 블랙웰(Blackwell) 기반의 서버 환경에서 대규모 전문가 혼합(MoE) 모델의 학습을 가속화하는 오픈소스 커널 ‘MoK(Mixture-of-Kittens)’를 공개했다. 이 새로운 기술은 학습 처리량을 41%까지 향상시킬 수 있는 가능성을 보여주고 있어, AI 연구자들에게 큰 기대를 모으고 있다.

MoK의 혁신적인 접근 방식

MoK는 대규모 AI 모델의 학습 과정에서 가장 큰 병목 현상으로 지목되는 MoE 계층의 통신 문제를 해결하기 위해 설계되었다. 기존의 MoE 구조에서는 전문가 간 토큰을 GPU 사이에서 주고받는 과정이 필요하여, 이로 인해 통신 비용이 전체 학습 시간의 절반 이상을 차지하는 경우가 많았다. 커서는 이 문제점을 인식하고, 통신과 연산을 하나의 결정론적 커널로 통합하는 방식을 채택했다.

이 기술은 특히 엔비디아의 GB300 NVL72와 GB200 NVL72와 같은 대규모 GPU 환경에서 최적화되어 있으며, GPU 내부에서 모든 스케줄링을 처리하여 CPU와의 동기화를 최소화하는 데 주력했다. 커서는 기존의 푸시(Push) 방식 대신 상황에 따라 풀(Pull)과 푸시 방식을 혼합하여 새로운 통신 구조를 도입함으로써, NV링크 대역폭 활용률을 최대 29% 높이는 성과를 거두었다.

성능 개선과 기대되는 효과

MoK의 성능은 이미 검증된 바 있다. 커서는 DeepEP, NCCL+PyTorch와 같은 기존 구현들과 비교했을 때, 다양한 조건에서 최대 2.37배의 처리량 향상을 기록했다고 밝혔다. 예를 들어, 512개의 GPU를 사용한 학습 환경에서 GPU당 초당 처리 토큰 수가 760.9개에서 1070.2개로 증가하여 전체 학습 처리량이 약 41% 개선되었다.

이러한 성능 향상은 대규모 AI 연구소와 GPU 네오클라우드 사업자, 국가 슈퍼컴퓨팅 센터 등 대규모 블랙웰 인프라를 보유한 기관에 큰 도움이 될 것으로 기대된다. MoK는 BF16과 MXFP8 정밀도를 모두 지원하며, 결정을 내리는 방식이 고정되어 있어 다양한 학습 작업에서도 안정적인 성능을 제공할 수 있다.

MoK의 미래와 AI 연구에 미치는 영향

커서는 AI 기술이 발전하면서 단일 커널 최적화는 대부분 자동화할 수 있는 수준에 이르렀다고 설명했다. 이제는 소규모 팀도 MoK와 같은 복잡한 시스템을 AI의 도움을 받아 개발할 수 있는 시대가 열렸다. 이러한 변화를 통해 더 많은 연구자들이 대규모 MoE 모델을 효율적으로 학습할 수 있을 것으로 보인다.

결국, MoK의 출현은 AI 모델 학습의 패러다임을 바꿀 수 있는 잠재력을 지닌 혁신으로 평가된다. 이제 연구자들은 더 나은 성과를 위해 자원을 효율적으로 활용할 수 있는 기회를 얻게 되었고, 이는 AI 기술 발전에 긍정적인 영향을 미칠 것으로 기대된다. 커서의 MoK가 앞으로 어떤 새로운 가능성을 열어줄지 귀추가 주목된다.