Cohere funde decode inteiro em um único kernel de GPU

Um kernel persistente por passo de decode

A Cohere apresentou um motor de serving para o North Mini Code construído em torno de um megakernel de decode, em BF16 e numa única H100, com ganho de 1,25× a 1,41× sobre o vLLM de ponta a ponta. O código está no GitHub. A empresa afirma acreditar ser o primeiro sistema de serving completo baseado em megakernel de decode, com batching contínuo, paged attention e endpoint compatível com OpenAI. Em batch 1, o megakernel chega a 292 tok/s, 62% do limite teórico de banda e 1,58× o vLLM; a margem se mantém em vários batch sizes e até 256K de contexto, sem perda mensurável de acurácia.

A maioria das stacks de serving trata cada forward pass como uma sequência de kernels: lança QKV e espera; a atenção e espera; o MoE e espera. Cada lançamento é adequado isoladamente; o problema é a espera entre eles. Em batch pequeno, a GPU passa uma fração surpreendente do passo esperando em vez de computar.

O decode autoregressivo, em batch baixo, é limitado por memória: move grande parte do HBM com relativamente pouca computação; a pergunta correta é quanta banda se aproveita, não quantos flops. O North Mini Code tem 30 bilhões de parâmetros, 3,3 bilhões ativos por token - em BF16, 6,6 GB de pesos por passo, mais cerca de 0,5 GB de KV cache em 8K de contexto. A H100 entrega 3,35 TB/s via HBM, pondo o limite físico (Speed-of-Light) em cerca de 470 tok/s. O vLLM serve o modelo a 185 tok/s, 39% desse limite.

De onde vem o ganho

O desenho parte do "Look Ma, No Bubbles!", da Hazy Research, que fundiu o forward pass de um Llama-3.2-1B em um único kernel e atingiu 78% da banda de uma H100 em batch 1, contra cerca de metade para vLLM e SGLang. Três ideias são centrais: o padrão de interpretador de tarefas na GPU, sincronização por contadores e sobreposição entre tarefas.

Um megakernel roda um threadblock por SM, residente durante todo o passo. Cada bloco lê uma lista de tarefas preparada no host, e as dependências viram contadores em memória global, incrementados quando uma tarefa termina e consultados por quem precisa do dado. A unidade de escalonamento cai de uma operação inteira para um tile; a de sincronização, do GPU inteiro para os produtores específicos.

O ganho tem quatro fontes: menos overhead de lançamento e sincronização; menos wave quantization (200 tiles em 132 SMs deixam 64 SMs ociosos na segunda onda); eliminação de dependências falsas, já que barreiras de kernel são barreiras de grid inteiro; e prefetch de pesos, imutáveis e transmissíveis antes das ativações - usado sobretudo no router e nas projeções QKV, na cauda do O-proj da camada anterior.

ABI, barreiras e escalonamento

O modelo usa camadas de transformer paralelas: atenção e MoE partem da mesma entrada normalizada e só se reúnem num residual add + RMSNorm fundido, o que permite preencher SMs ociosos. As 16 operações do passo obedecem a uma convenção comparada a uma ABI: 12 warps em 3 warpgroups - controller, producer, storer e consumidores - lendo um descritor de 32 inteiros. A Cohere diz que o megakernel é mais fácil de escrever do que a reputação sugere: é um único arquivo CUDA, sem compilador nem paradigma novo. A parte difícil são as barreiras, que liberam quando N threads chegam sem verificar quais; um opcode que chega o número errado de vezes pode causar deadlock mais tarde.

O escalonamento é majoritariamente estático, com work stealing local para atenção e MoE. O host monta waves nomeadas e distribui a tarefa k à SM k mod 132. Ablações de ordem de waves mostram diferenças de 1% a 19% em throughput; o posicionamento por afinidade de dependência atrasou o kernel em 1% a 2%, então a versão publicada usa round-robin. Um modelo causal completo segue em aberto.

Servidor, limitações e resultados

O servidor tem uma thread Python de plano de controle, que aceita requisições, roda prefill e gerencia KV, e uma thread C++ que detém o decode. Como a lista de tarefas é construída de antemão, o Python estaciona o decode antes de admitir, retirar ou remodelar o batch.

As limitações declaradas são da implementação atual, não do desenho: prefill e decode não se misturam (um prefill pausa toda a decodificação), o batch máximo é 8 por configuração e o megakernel cobre só o decode. Nos testes, o baseline é o vLLM v0.24 com atenção FA3 e MoE Triton, sem prefill, medindo 1K tokens de saída. Em 8K de contexto e batch 8, o megakernel é 1,32× mais rápido com a distribuição real de experts e 1,14× sob roteamento uniforme; a Cohere diz que o roteamento uniforme sintético subestima o ganho no tráfego real.