Today, Cohere presents a serving engine for North Mini Code built around a decode megakernel: BF16 on a single H100, 1.25× - 1.41× faster than vLLM end-to-end. Explore the code behind the serving engine on GitHub.
Most LLM serving stacks still treat each forward pass as a sequence of kernels: launch QKV, wait; launch attention, wait; launch the MoE, wait. Each launch is fine on its own. The problem is the waiting in between. At small batch sizes, the GPU spends a surprising fraction of every decode step waiting rather than computing.
Autoregressive decoding, especially at lower batch sizes is fundamentally memory-bound. For every decode step, we move a large fraction of memory from HBM while relatively doing less compute. This means the correct question to ask is, how effectively can we use memory bandwidth and not the flops. Take North Mini Code which is a 30B model with 3.3B parameters active per token, which in BF16 means streaming 6.6 GB of weights during every decode step, plus roughly 0.5 GB of KV cache at 8K context. An H100 delivers 3.35 TB/s of bandwidth through HBM, putting the Speed-of-Light (SoL) at about 470 tok/s. vLLM serves this model at 185 tok/s, merely 39% of SoL.
Megakernels have been getting attention lately as the way to close that gap: instead of a hundred small kernels, run the entire forward pass as one persistent kernel. Starting from the pioneering work by Hazy Research's "Look Ma, No Bubbles!", whose design we recap below, numerous follow-up works have been released, achieving various levels of speedup. Existing work has gone mainly in two directions: compilers that generate megakernels automatically, and standalone demos that measure decode speed at batch size 1.
We take one step further. This post presents what we believe is the first fully fledged serving system built around a decode megakernel. It supports everything a real server needs: continuous batching, paged attention, and ragged sequence lengths, all behind an OpenAI-compatible endpoint with tool calling. Point OpenCode at it and you can code with it.
On batch size 1, our megakernel reaches 292 tok/s, or 62% of SoL — 1.58× faster than vLLM. That margin holds across batch sizes and out to 256K of context, with no measurable loss of accuracy.
Image 1: Decode throughput of various context lengths at batch size 1. Megakernel consistently outperforms vLLM.
We also found that megakernels are much easier to write than their reputation suggests, so we include a recipe for porting kernels you already have into one. Ours is a single CUDA file: no compiler, no new programming paradigm, no exotic abstractions — just ordinary tiled GEMMs and ordinary paged attention, restructured to fit a single calling convention.
What is a Megakernel?
A GPU is roughly 100–150 independent processors, called SMs (streaming multiprocessors), that all run the same program — a kernel — on different pieces of data. A megakernel is a single persistent kernel that runs an entire forward pass: we launch exactly one threadblock per SM, and it stays resident for the entire decode step. Instead of receiving work from the driver, each block reads a task list — a list of small pieces of work it should execute, prepared on the host and sitting in global memory. Instead of kernel boundaries encoding the data dependencies, dependencies are expressed as explicit counters in global memory that tasks increment when they finish and spin on when they need an input.
The result is that the unit of scheduling shrinks from an entire operation to one tile of one operation, and the unit of synchronization shrinks from the whole GPU to the specific producers a task depends on.
Image 2: One decode step, from the host to the device. Instead of one kernel per operation, the host decomposes the step into tasks — one tile of one operation each — and distributes them round-robin over the SMs, so every SM gets its own task list in global memory. The order of most tasks are determined on host by the scheduler (see the scheduler section). Full attention and the MoE are the exception: their task counts depend on the live sequence lengths and on routing, so they go into shared work queues that any SM can pull from. We describe the details in the following sections.
Where the speedup comes from
Inference engines typically launch one kernel per operation — RMSNorm, QKV, attention, MoE, and so on — and most optimization effort goes into making each of those kernels as fast as it can be. That works well for training and prefill, where the workload is large compute-bound GEMMs and every kernel has enough work to saturate the SMs.
Decode is the opposite regime: it is bound by memory bandwidth and by latency. A decode step is mostly GEMVs with low arithmetic intensity, so its speed is determined by how fast weights can be streamed from HBM into shared memory. Anything that keeps weights from moving is lost time, and a kernel-per-operation approach has several places where they stop moving. Together those stalls account for most of the 61% of bandwidth that a typical inference engine leaves unused.
The simplest benefit comes from reduced launch and synchronization overhead. Between two consecutive kernels, every SM must finish before any SM can start the next one, and the driver has to dispatch the next grid. For a decode step made of dozens of small kernels per layer, those gaps add up. A megakernel pays that cost once per step rather than once per operation. We list three more benefits that matter more for this model, in rough order of impact:
- Reduce wave quantization
Suppose a kernel has 200 tiles of work to do and the GPU has 132 SMs. The first 132 tiles run in parallel; the remaining 68 run in a second wave while 64 SMs sit idle. The kernel takes two waves' worth of time to do 1.5 waves' worth of work, and the smaller the kernel, the worse that rounding gets. This is not something we can fix by dividing the work more evenly. GEMM tile shapes are constrained by the matrix dimensions and the kernel design, so the total tile count rarely lands on an exact multiple of the SM count.
Image 3: Wave quantization reduces GPU utilization.
In a megakernel there is no boundary to round up to: a tile whose inputs are ready starts on whichever SM is free. North Mini Code gains more from this than most architectures because it uses parallel transformer layers: attention and the MoE feed-forward are computed from the same normalized input and are rejoined only by a fused residual add + RMSNorm at the end of the layer, so neither attention nor the MoE needs each other's output.
Image 4: A parallel transformer layer used by North Mini Code.
With a megakernel, we can "backfill" the idle SMs with ready work. The parallel transformer layers allow us to perform backfilling in a more aggressive way: whenever possible, we deterministically place the tasks that are likely ready to run on the idle SMs. The detail of the placement is described in the task scheduler section. In the figure below, we compare an MoE decode layer run by a conventional serving stack and a megakernel.
Image 5: One MoE decode layer, the same operations in the same order, scheduled two ways. Top: one kernel per operation. Attention does not finish together due to wave quantization and hardware jitter. The barrier at the kernel boundary makes every SM wait for the last one (hatched). The same pattern repeats at each kernel boundary. Bottom: the megakernel. As each SM finishes its attention work it starts an MoE tile, so the same interval carries work instead of waiting. Drawn with 16 SMs and a few dozen tiles per operation so tasks stay visible; the real kernel has 132 SMs and thousands of tasks. Durations are illustrative and the time axis is relative.
- Drop false dependencies
SMs do not always finish work at the same time, even when workload is identical. A kernel boundary is a full-grid barrier, so the slowest SM sets the pace for all of them. For example, if attention is split across 4 key/value groups and one group finishes early, that SM idles until the other three catch up, even though the data its next operation actually needs is already in memory. Fine-grained barriers drop the false dependency: O-proj for a given KV group starts as soon as that group's attention output lands. Similarly, an MoE down projection task can start whenever the up projection task of the corresponding expert is completed, without the need to wait for all experts’ up projection.
- Weight prefetch
Weights are immutable — they do not depend on this step's activations at all. A task can therefore start streaming its weight tiles from HBM into shared memory before its activation dependency has been satisfied, which a kernel boundary would forbid. We use this most aggressively for the router and QKV projections, which prefetch their weights during the tail of the previous layer's O-proj, before RMSNorm has even run, to harvest bandwidth that would otherwise go unused.
Where we started
Our design borrows a lot of insight from the pioneering post on Hazy Research referenced earlier , which fused a Llama-3.2-1B forward pass into a single kernel and reached 78% of an H100's memory bandwidth at batch size 1, against roughly half that for vLLM and SGLang. Three of their ideas are crucial to us:
The “task interpreter” pattern on GPU. Each SM walks a list of task descriptors prepared on the host and reused across forward passes. A controller warp reads the descriptor and dispatches the task to various on-device functions, each implementing one kind of operation. We describe our implementation of the idea in Figure 2 and 6.
Counter-based synchronization. Dependency barriers are plain integers in global memory, zeroed before each step. A task increments one when it finishes and spins on one before it starts. We describe our implementation in detail in the Barriers section and Figure 7.
Overlap across task boundaries. A task can start loading its weights while the previous task on the same SM is still storing its results.
What we have done differently
Our implementation differs from their megakernel in a few ways. Our GEMM implementation relies heavily on tensor core instruction, even for batch size 1 because we found that wgmma is slightly faster in our cases compared to CUDA core and reduces register pressure.
One other difference is that we do not use shared memory paging to implement weight prefetching. We initially experimented with shared-memory paging, which allows memory loads to begin before the previous task releases its buffers. In practice, though, the bookkeeping is complex, introduced a steady source of bugs, and had high overhead that outweighed the benefit. Every opcode instead gets its own warp-specialized pipeline with its shared memory layout statically at compile time, and we get the overlap from two cheaper places.
Between consecutive GEMM tasks of the same type. An MoE GEMM task walks a list of tiles, and the pipeline carries its stage phases across the whole list instead of draining and refilling the pipeline at every tile boundary. Within the MoE GEMM, the next tile's weights are already in flight while the current tile is still on the tensor cores or epilogue. Because it is the same operation with the same shared memory layout, shared memory paging reduces to a simple multi-stage pipeline. This kind of prefetch shares the same spirit as persistent grouped GEMM kernels.
Inside the GEMM pipeline. Weights are immutable, so the producer warp issues its weight-tile loads before the cross-SM wait for activations; in the pseudocode (Listings 1) shown later in the post, prefetch_weight_tiles sits above wait_input_bars. MoE down projection is one such example case: its expert weights start streaming from HBM while up/gate is still computing the hidden state that down will consume. Therefore a task blocked on its inputs is still moving bytes. The o
[truncated for AI cost control]