Cohere Releases Megakernel Engine for North Mini Code, Achieving Up to 1.41x Speedup Over vLLM

Cohere released a Megakernel serving engine for the North Mini Code (30B, active 3.3B) model. This engine adopts a method of executing the entire forward pass as a single persistent kernel to reduce GPU idle time and utilize bandwidth more efficiently in memory-bandwidth-bound Autoregressive decoding. On a single H100 in BF16, it recorded a 1.25~1.41x end-to-end speedup compared to vLLM, where vLLM achieved 185 tok/s (approximately 39% of Speed-of-Light), while Megakernel demonstrates higher efficiency.

Key Technologies and Architecture

The core approach reduces launch/sync costs per operation from the existing kernel-per-operation model to once per decode step, and reduces idle time caused by wave quantization. Threadblocks are resident on the GPU's 100~150 SMs, with the host distributing tasks via a task list in global memory, and dependencies are managed using explicit counters and fine-grained barriers in global memory instead of kernel boundaries. Weight prefetching and deterministic backfilling place ready work on idle SMs to reduce wave quantization.

The implementation borrows from prior research by Hazy Research but adds production server features such as Continuous batching and paged attention. It is written as a single CUDA file without a complex compiler, restructuring existing tiled GEMM and paged attention. All operations adhere to a unified ABI with 3 warp groups and fixed-size task descriptors, and fuse operations such as RoPE and SiLU-multiply.

Scheduling and Performance Benchmarks

The Scheduler applies a Mostly static + dynamic work stealing approach. The Serving Engine has a structure where a Python thread (control, prefill, batch management) and a C++ thread (decode ownership) are interleaved.

Performance benchmarks were conducted on a Single H100 (132 SMs) with vLLM v0.24 as the baseline. Decode throughput was measured at 8K context and 1K output tokens. The Real expert distribution, with increased bubbles due to sparse MoE work, is favorable for the megakernel. At BS=8, the Real checkpoint achieved 1.32x and Uniform routing achieved 1.14x faster. End-to-end serving (BS=8) showed an Avg decode throughput 1.25~1.41x gain. Performance improvements were achieved on AIME 2025 at 1.41x (935 vs 661 tok/s), GPQA at 1.25x, MMLU-Pro at 1.33x, SciCode at 1.37x, and LiveCodeBench v6 at 1.28x, with quality maintained without accuracy loss on SciCode and LiveCodeBench v6.

Limitations and Future Plans

Current limitations include the inability to mix Prefill and decode (decode pauses during prefill), a Max batch size limit of 8, and the use of Megakernel exclusively for decode (prefill uses standard PyTorch kernels). Future plans include support for Prefill/mixed batches, the release of FP8/FP4 megakernels for RTX Blackwell (such as RTX Pro 6000), and expansion to datacenter Blackwell and multi-GPU.