Skip to content
AI Primer
release

Cohere open-sources fused LLM decode kernel with 1.58x vLLM claim

Cohere released an open-source serving system that fuses the LLM decode step into one GPU kernel launch. On North Mini Code with one H100, it reports up to 1.58x vLLM performance at the tested batch size.

5 min read
Cohere open-sources fused LLM decode kernel with 1.58x vLLM claim
Cohere open-sources fused LLM decode kernel with 1.58x vLLM claim

TL;DR

  • Cohere shipped an open-source serving engine that runs an LLM's whole decode forward pass inside one persistent GPU kernel, according to cohere's release.
  • The headline 1.58x result comes from BF16 North Mini Code at batch size one on a single H100, while cohere's benchmark post gives a 1.25x to 1.41x end-to-end range at batch size eight.
  • Production features remain in the stack, including continuous batching, paged attention, ragged sequence lengths, an OpenAI-compatible endpoint, and tool calling, which cohere's announcement lists alongside the kernel.
  • The engine breaks decode into 16 operations and schedules tile-sized tasks across SMs, as cohere's design note describes.

The technical post specifies a 32-int32 task descriptor and a fixed 12-warp execution shape. A 2025 Hazy Research megakernel already fused a Llama-1B forward pass, while Cohere has extended the pattern to a serving surface with continuous batching and paged KV management.

The decode megakernel

Cohere launches one threadblock per streaming multiprocessor, then leaves those blocks resident for the full decode step. Instead of the driver dispatching separate QKV, attention, MoE, normalization, and output kernels, each SM reads a host-prepared list of small tasks from GPU memory.

Dependencies move from kernel boundaries into counters in global memory. A task increments its output counter on completion and waits only for the inputs it needs, allowing another ready tile to occupy an SM.

The 16-opcode task graph

The usual decode graph becomes 16 task types in Cohere's implementation write-up:

  • Dense GEMMs: QKV_PROJ, O_PROJ, FFN_UPGATE_ACT, FFN_DOWN, LM_HEAD
  • Attention: ATTN_DECODE, ATTN_COMBINE, ATTN_DRAIN
  • MoE routing: ROUTER_GEMM, ROUTER_TOPK, ROUTE_FINALIZE, MOE_GATHER
  • MoE GEMMs: MOE_UPGATE_ACT_DRAIN, MOE_DOWN_DRAIN, MOE_COMBINE
  • Residual and normalization: ADD_RMSNORM

The host encodes each task as 32 int32 fields, including its opcode, layer, output tile, split-K coordinates, and the counters it waits on and signals. Attention and MoE tasks whose sizes depend on live sequence lengths or routing use work queues that any SM can claim from; most other work follows a host-built static schedule.

Memory bandwidth and bubbles

North Mini Code has 30B parameters but activates 3.3B per token. Cohere estimates that BF16 decode streams 6.6 GB of weights per step, plus roughly 0.5 GB of KV cache at 8K context; against an H100's 3.35 TB/s HBM bandwidth, it puts the theoretical ceiling near 470 tokens per second.

The company identifies four sources of lost bandwidth in conventional decode:

  1. Kernel launches and whole-grid synchronization.
  2. Wave quantization, where the last partial wave leaves SMs idle.
  3. False dependencies, where a task waits for a full operation rather than its actual producer.
  4. Weight prefetch blocked by kernel boundaries.

Its scheduler uses North Mini Code's parallel attention and MoE branches to backfill idle SMs, while counter barriers unlock per-group work and immutable weights begin loading before activations are ready. The Cohere post says these latter effects outweigh launch overhead for this model.

That distinction matches vikhyatk's reply, which says launch overhead is only a small part of the appeal and that CUDA graphs already address much of it.

North Mini Code on one H100

Cohere reports 292 tokens per second at batch size one, or 62% of its stated memory-bandwidth ceiling. Its benchmark report compares that with vLLM at 185 tokens per second, producing the 1.58x figure.

| Test | Reported result |
| --- | --- |
| BF16, 1 H100, batch size 1 | 292 tok/s, 1.58x vLLM |
| End-to-end serving, batch size 8 | 1.25x to 1.41x vLLM |
| Context range | Gains reported through 256K |

Cohere says the batch-size-one margin persists across larger batch sizes and context lengths, with no measurable accuracy loss. The company has published a vendor benchmark on one model and one GPU configuration, rather than an independent cross-model comparison.

Production serving features

Continuous batching is the operationally meaningful claim around the kernel. Cohere says the engine handles paged attention and ragged sequence lengths behind an OpenAI-compatible endpoint, with tool calling; the post also says OpenCode can point at it.

Users can run the model on their own hardware efficiently and privately, nickfrosst wrote.

Hazy Research's precedent

Cohere credits Hazy Research's "Look Ma, No Bubbles!" as a key predecessor. That project fused a Llama-1B forward pass on an H100, reported 78% of memory bandwidth, and claimed more than 1.5x over existing systems.

The open Hazy repository includes a low-latency Llama demo with H100 and B200 build targets. Cohere's contribution is its North Mini Code-specific task graph and a server design that accommodates dynamic batching and routing.

The 12-warp ABI

Cohere implements the engine in a single CUDA file around a common ABI that every task obeys. Each block has 12 warps: one controller that prefetches descriptors, one producer for TMA loads and input waits, one storer for writes and output signals, eight consumer warps for computation, and one unused warp.

The fixed shape lets an SM run a QKV tile, an attention slice, and an MoE task back to back without changing its warp roles. Adding a GEMM task means supplying its tensors, barriers, and epilogue behavior under that convention, as cohere's design note explains.

Further reading

Discussion across the web

Where this story is being discussed, in original context.

On X· 2 threads
Memory bandwidth and bubbles1 post
Production serving features1 post
Share on X