Cohere's Open-Source Megakernel Beats vLLM by 1.58x on H100

Cohere released an open-source serving engine that collapses an entire LLM forward pass into a single persistent CUDA kernel. Instead of launching dozens of small kernels per layer and paying synchronization costs at every boundary, all compute stays resident on the GPU from start to finish. The engine targets Cohere's North Mini Code model and reports up to 1.58x faster throughput than vLLM on an H100.

The code, released as cohere-megakernel under Apache 2.0, ships as an OpenAI-compatible server with streaming, tool calling, continuous batching, and paged KV cache. Cohere frames it as a research release rather than a general-purpose engine, but it is the first end-to-end serving system built around the megakernel approach.

Why decode leaves the GPU idle

Autoregressive decoding at low batch sizes is memory-bound. Every generated token forces the GPU to stream a large chunk of weights out of HBM to do relatively little arithmetic, so the bottleneck is bandwidth, not compute.

Cohere quantifies this with a Speed-of-Light (SoL) calculation. North Mini Code is a 30B model with 3.3B active parameters per token, which in BF16 means streaming 6.6 GB of weights per decode step plus roughly 0.5 GB of KV cache at 8K context. An H100 delivers 3.35 TB/s of HBM bandwidth, putting the theoretical ceiling at about 470 tok/s. vLLM serves the model at 185 tok/s, 39% of SoL. The remaining 61% is lost to the GPU waiting between kernel launches.

One kernel for the full decode step

A GPU contains roughly 100 to 150 independent processors called SMs that all run the same program on different data. A megakernel keeps one threadblock per SM resident for the entire decode step. Each block reads a task list in global memory rather than receiving work from the driver, and data dependencies are encoded as counters that tasks increment on completion and spin on when waiting for inputs.

For a memory-bound workload, that design yields three concrete gains:

• No wave quantization at kernel boundaries. A tile whose inputs are ready starts on whichever SM is free. North Mini Code's parallel transformer layers amplify this because attention and the MoE feed-forward branches are independent of each other.
• No false dependencies. O-proj for a given KV group starts as soon as that group's attention output is ready, rather than waiting for the slowest SM to finish attention across all heads.
• Weight prefetch across task boundaries. Because weights are immutable, a task can begin streaming its weight tiles from HBM into shared memory before its activation dependency resolves, consuming bandwidth that would otherwise sit idle.

The calling convention that makes it hand-writable

The design extends Hazy Research's megakernel work on Llama-3.2-1B. Every operation inside the kernel, whether GEMM, attention, RMSNorm, or MoE routing, follows a fixed ABI: exactly 3 warp groups (8 consumer warps, 1 controller, 1 producer, 1 storer) reading parameters from a 32-int32 task descriptor. The decode graph lowers to 16 opcodes covering QKV projection, attention decode and combine, router GEMM and top-k, MoE up/down/combine, and the LM head. All GEMM opcodes share one pipeline, differing only in which tensors they touch, whether they fuse an epilogue like SiLU-and-multiply, and which barrier they signal.

Synchronization is intentionally minimal. Barriers are counters in global memory:

// wait: spin until enough upstream tasks have arrived
while (*(volatile const uint32_t*)bar < target) {
    __nanosleep(20);
}
__threadfence();
// arrive: publish my tile, then signal downstream
fence.proxy.async;
__threadfence();
atomicAdd(bar, 1);

Each task waits on a single count, making both signaling and dependency checks O(1) regardless of fan-in or fan-out.

Scheduling: mostly static, selectively dynamic

The host builds named waves per layer and assigns task k to SM k mod 132. Attention and MoE, whose tile counts depend on live sequence lengths and expert routing, use claimer tasks that atomically pull from shared work queues.

Wave ordering has a measurable impact. An ablation at batch size 1 showed the tuned order reaching 291 tok/s, an interleaved variant at 282 tok/s, and an attention-first variant at only 236 tok/s. A greedy topology-aware scheduler and brute-force search over hundreds of candidates were both tried and abandoned once MoE was introduced, since routing makes work placement dynamic anyway.

Benchmark results

On a single H100 in BF16, the megakernel reaches 292 tok/s at batch size 1, or 62% of Speed-of-Light. End-to-end serving throughput, measured with real prompts through the API including prefill:

Benchmark	Megakernel	vLLM	Speedup
AIME 2025	935 tok/s	661 tok/s	1.41x
SciCode	711 tok/s	560 tok/s	1.37x
MMLU-Pro (CS)	948 tok/s	713 tok/s	1.33x
LiveCodeBench v6	803 tok/s	625 tok/s	1.28x
GPQA	787 tok/s	631 tok/s	1.25x

Accuracy holds: on SciCode the megakernel scored 38.9% against vLLM's 38.2%, and both tied at 70.3% on LiveCodeBench v6.

Speedup grows larger under real expert routing than under uniform simulated routing. Real requests tend to concentrate on the same experts, leaving the active-expert set sparse. The MoE does less total work, so pipeline bubbles represent a larger fraction of the step, and those are exactly the gaps the megakernel eliminates.

Current limitations

Straight from the repo, the release is narrow by design:

• H100 with SM90a and BF16 only. No Blackwell, no FP8 or FP4.
• Batch sizes 1 through 8 only.
• Decode only. Prefill runs as ordinary PyTorch kernels and pauses decode while it runs.
• Sampling is greedy or temperature; top_p must be 1.0, top_k must be 1 or unset.
• The task schedule is specific to North Mini Code.

Cohere has an RTX Blackwell megakernel with FP8 and FP4 quantization on the roadmap, along with datacenter Blackwell and multi-GPU inference with tensor and expert parallelism.

A recipe other teams can follow

The prevailing assumption about megakernels has been that they require a compiler, a new programming model, or an exotic abstraction layer. This release challenges that. Ordinary GEMM and attention kernels already handle the hard math; fitting them to a shared 12-warp threadblock ABI is sufficient to assemble a megakernel by hand, one operation at a time. Start from competitive standalone kernels, wire them to the task descriptor format, add counter barriers for data dependencies, and emit descriptors into the task list.

For teams serving MoE models at low-to-moderate batch sizes, a meaningful share of that missing 61% of HBM bandwidth is recoverable without switching frameworks or rewriting kernels from scratch. The sparser the expert routing under real traffic, the larger the payoff.