Cohere Megakernel

// Overview

Cohere Megakernel is a research release of a single-H100 inference engine that serves Cohere's North Mini Code model behind an OpenAI-compatible API. Its defining idea is the decode megakernel: instead of launching one CUDA kernel per operation (RMSNorm, QKV, attention, MoE, output projection) and paying a full-grid synchronization at every boundary, the engine launches one thread block per streaming multiprocessor and keeps it resident for the entire decode step.

Each resident block walks a host-built task list where every task is a single GEMM tile or a single split of attention, and dependencies are explicit counters in global memory. Work therefore starts the moment its own inputs are ready rather than at a kernel boundary. Cohere reports three concrete wins from that finer schedule: ready mixture-of-experts tiles backfill the SMs left idle by the tail of attention, consumers stop waiting on false whole-grid dependencies, and immutable weights are prefetched from HBM while other work runs.

The project is explicitly an early research release rather than a general-purpose inference engine. Cohere tested it on a single NVIDIA H100 (sm_90a) with CUDA 13 or newer, CPython 3.12 or newer, Linux, and batch sizes up to eight; other configurations are not built or tuned. Within that envelope it is a complete server, with OpenAI-compatible completions and chat-completions endpoints, streaming, tool calling, continuous batching, ragged sequence lengths, a paged KV cache, sliding-window attention, prefix caching and preemption.

// What it does

• One persistent CUDA kernel executes the complete decode forward pass, removing per-op launch overhead and full-grid barriers
• 292 tokens/second at batch size 1 on an H100 — 62% of theoretical speed-of-light, and 1.58x vLLM v0.24 decode
• 1.25x-1.41x faster end to end than vLLM across AIME 2025, SciCode, MMLU-Pro (CS), LiveCodeBench v6 and GPQA
• OpenAI-compatible /v1/completions, /v1/chat/completions and /v1/models endpoints with streaming and tool calling
• Continuous batching, ragged sequence lengths, paged KV cache, sliding-window attention, prefix caching and preemption
• Speedup holds out to 256K context at batch size 1, with no measurable accuracy loss versus vLLM

// Getting started

You need one NVIDIA H100, CUDA 13+, Python 3.12+, Linux, and a local North Mini Code checkpoint. The repository's BUILD_AND_RUN.md covers detailed install, server flags, benchmarking and profiling.

git clone --recurse-submodules https://github.com/cohere-ai/cohere-megakernel.git
cd cohere-megakernel

uv venv && source .venv/bin/activate
uv pip install -r requirements.txt
uv pip install flash-attn-3 --index-url https://download.pytorch.org/whl/cu130
sudo apt install -y libssl-dev ninja-build

cmake -S . -B build -G Ninja \
  -DPython_EXECUTABLE="$(which python)" -DCMAKE_BUILD_TYPE=Release
cmake --build build

python src/serving/server.py \
  --host 127.0.0.1 --port 8000 \
  --lib "$PWD/build/libmk_release.so" \
  --ckpt <checkpoint-path> \
  --device cuda:0 --bs 8 --frac-vram-utilization 0.7

curl http://127.0.0.1:8000/v1/chat/completions \
  -H 'Content-Type: application/json' \
  -d '{
    "model": "North-Mini-Code-1.0",
    "messages": [{"role": "user", "content": "Write a Python Fibonacci function."}],
    "max_tokens": 128,
    "temperature": 0
  }'

Commands and code are distilled from the project's own documentation — always check the official repo for the latest.

// When to use it

• Serving North Mini Code on a single H100 at low batch sizes, where per-kernel launch overhead dominates decode time
• Dropping a faster backend behind an existing OpenAI-compatible client without changing application code
• Studying a worked, single-file example of a persistent megakernel before writing your own
• Benchmarking decode-bound mixture-of-experts serving against vLLM on identical hardware

// How Cohere Megakernel compares

Cohere Megakernel alongside other open-source gpu kernels & compilers tools AI/TLDR tracks, ranked by GitHub stars.