Cohere, North Mini Code용 Megakernel 엔진 공개… vLLM 대비 최대 1.41배 속도 향상

Cohere가 North Mini Code(30B, active 3.3B) 모델을 위한 Megakernel serving engine을 공개했다. 이 엔진은 메모리 대역폭에 제한되는 Autoregressive decoding에서 GPU 유휴 시간을 줄이고 대역폭을 더 효율적으로 활용하기 위해 전체 forward pass를 하나의 persistent kernel로 실행하는 방식을 채택했다. H100 단일 BF16 환경 기준 vLLM 대비 end-to-end 1.25~1.41배의 속도 향상을 기록했으며, vLLM은 Speed-of-Light(SoL)의 약 39%인 185 tok/s를 기록한 반면 Megakernel은 더 높은 효율성을 보인다.

핵심 기술 및 아키텍처

기존 kernel-per-operation 방식에서 연산마다 발생하던 launch/sync 비용을 decode step당 한 번으로 줄이고, wave quantization으로 인한 유휴 시간을 감소시키는 것이 핵심이다. GPU의 100~150개 SM에 threadblock을 상주시키고 host가 global memory의 task list로 작업을 분배하며, kernel 경계 대신 global memory의 explicit counter와 fine-grained barrier로 의존성을 관리한다. Weight prefetching과 deterministic backfilling을 통해 idle SM에 ready work를 배치해 wave quantization을 감소시킨다.

구현은 Hazy Research의 선행 연구를 차용하되 Continuous batching, paged attention 등 실전 서버 기능을 추가했다. 복잡한 컴파일러 없이 단일 CUDA 파일로 작성되었으며, 기존 tiled GEMM과 paged attention을 재구성했다. 모든 연산은 3 warp groups와 고정 size task descriptor로 통일된 ABI를 준수하고, RoPE, SiLU-multiply 등을 융합했다.

스케줄링 및 성능 벤치마크

Scheduler는 Mostly static + dynamic work stealing 방식을 적용한다. Serving Engine은 Python thread(제어, prefill, batch 관리)와 C++ thread(decode 소유)가 interleaved되는 구조를 가진다.

성능 벤치마크는 Single H100(132 SMs)에서 vLLM v0.24를 baseline으로 진행했다. 8K context, 1K output token 기준 decode throughput을 측정했으며, Real expert distribution은 sparse MoE work로 bubble이 증가해 megakernel에 유리하다. BS=8 기준 Real checkpoint에서 1.32x, Uniform routing에서 1.14x faster를 기록했다. End-to-end serving (BS=8)에서는 Avg decode throughput 1.25~1.41x gain을 보였다. AIME 2025에서 1.41x (935 vs 661 tok/s), GPQA 1.25x, MMLU-Pro 1.33x, SciCode 1.37x, LiveCodeBench v6 1.28x의 성능 향상을 달성했으며, SciCode와 LiveCodeBench v6에서 정확도 손실 없이 품질이 유지됨이 확인되었다.

한계 및 향후 계획

현재는 Prefill과 decode 혼합 불가(prefill 시 decode 일시 정지), Max batch size 8 제한, Megakernel의 decode 전용 사용(prefill은 일반 PyTorch kernel 사용) 등의 한계가 있다. 향후 Prefill/mixed batch 지원, RTX Blackwell(RTX Pro 6000 등)용 FP8/FP4 megakernel 출시, 데이터센터 Blackwell 및 multi-GPU 확장을 계획하고 있다.