Megakernels para TPU atingem 700 tokens por segundo no Kimi K3

O resultado

A Inferact lançou o projeto open-source inferact/tpu-megakernels, uma coleção de megakernels para TPU v7. A implementação para o Kimi K3 atinge mais de 700 tokens por segundo com speculative decoding (709 tok/s em 16 TPU v7, contra 452 tok/s em 16 GB200 com vLLM). Sem speculative decoding, os megakernels para Kimi K3 e Qwen 3.8 27B entregam cerca de 1,4× a 2× o throughput de decode do baseline GB200 em batch sizes de 1 a 8.

Por que TPUs e megakernels combinam

A velocidade de decode é limitada pela largura de banda de memória do hardware. No papel, o TPU v7 Ironwood do Google é comparável à GPU GB200 da NVIDIA: specs como compute BF16 (2,31 vs. 2,5 PFLOPS), capacidade e largura de banda de HBM são parelhos. A diferença relevante está na hierarquia de memória on-chip. A TPU tem um grande pool de SRAM chamado VMEM (64 MiB por TensorCore, 128 MiB por chip), totalmente gerenciado por software: os dados permanecem lá até o programa sobrescrevê-los explicitamente. A GPU, em contraste, divide ~38 MiB de Tensor Memory entre 152 SMs, com apenas 256 KB por SM - pouco espaço para pré-carregar pesos.

Kernels tradicionais alternam fases de carga, computação e escrita, deixando a banda de HBM subutilizada entre elas; um passo de decode dispara centenas desses kernels, e os "bubbles" de banda se acumulam. Recursos de GPU como CUDA Graphs, streams e programmatic dependent launch recuperam parte do overlap, mas mantêm fronteiras entre kernels. O megakernel elimina essas fronteiras: todo o passo de decode vive em um único programa, e o carregamento de pesos pode ser iniciado assim que houver espaço no armazenamento on-chip. As restrições passam a ser tempo de vida de buffers e capacidade. No TPU, isso permite pré-carregar os pesos da próxima camada enquanto a atual ainda computa ou se comunica.

Implementação com Pallas

Os autores usaram Pallas, a linguagem de kernels para TPUs, e afirmam que este é o primeiro megakernel de inferência open-source escrito nela. Três recursos foram decisivos:

Lançamento sem grid: em vez de um grid de programas, o megakernel roda como um único programa sequencial, aproveitando o modelo single-threaded da TPU.
Alocações de VMEM com escopo: o primitivo run_scoped controla o tempo de vida das alocações (ativações residem durante todo o forward; o KV cache do MLA vive pouco), simplificando o gerenciamento de memória.
DMAs assíncronas: cópias explícitas entre HBM e VMEM especificam quando iniciar o prefetch e quando esperar - e, dentro de um único kernel, o XLA não pode reordenar essas dependências.

O megakernel do Kimi K3 combina as 92 camadas Mixture-of-Experts do modelo em uma única chamada Pallas, escolhendo condicionalmente entre Kimi Delta Attention (KDA) e Multi-head Latent Attention (MLA) conforme o índice da camada. O modelo é sharded em 16 chips TPU v7 (32 TensorCores): cabeças de atenção divididas entre os 32 ranks, experts roteados com TP4 × EP8. Para batch size 8, foi necessário combinar tensor parallelism com parallelismo de sequência do residual stream. Um efeito inesperado foi a compilação rápida: enquanto modelos com milhares de ops XLA podem levar mais de 30 minutos para compilar, o megakernel inteiro compila do zero em menos de 90 segundos.

Speculative decoding e resultados

Com batch size 8, o megakernel serve como verificador no DSpark speculative decoding: cada lançamento avalia 1 âncora e 7 continuações propostas, com aceitação típica de 3 a 6 tokens por passo. A cerca de 8,5 ms por passo, o sistema entrega mais de 700 tokens/s a um único usuário. Nos testes de precisão com decodificação greedy, a implementação pontuou 0,944 no GPQA-Diamond e 0,972 no GSM8K, com scripts de reprodução disponíveis.

Próximos passos

Os autores listam direções futuras: suportar workloads de maior concorrência (onde trabalho vetorial e tráfego entre dispositivos crescem) e cargas agentic (que tornam o movimento do KV cache crucial); ampliar topologias além da configuração 2×2×4 usada, o que exige novos collectives; e criar melhores ferramentas de análise dos bundles de instruções compilados, incluindo para agentes de código que escreverão kernels como esses. O trabalho faz parte da plataforma da Inferact, que otimiza inferência sobre a infraestrutura dos clientes, construída sobre o vLLM open-source.