tpu-megakernels README

The project's README from github.com/Inferact/tpu-megakernels. As one markdown file: README.md.

TPU megakernels for Kimi K3 and Qwen3.8-27B

TPU megakernel: 709 tokens per second on Kimi K3, 16x TPU v7 vs 452 on 16x GB200

Companion code for the blog post 700 TPS on Kimi K3: A Case for TPU Megakernels.

This repository contains fused decode implementations for two models:

Both demos support terminal interaction and an OpenAI-compatible HTTP server.

Results

With speculative decoding at an acceptance length of 6, the megakernels reach 709 accepted tokens per second on Kimi K3 (16× TPU v7 vs 16× GB200, 1.57×) and 1,515 on Qwen3.8-27B (4× TPU v7 vs 4× GB200, 2.18×):

Speculative decode throughput at acceptance length 6: Kimi K3 709 vs 452 tokens/s; Qwen3.8-27B 1,515 vs 695 tokens/s

DSpark keeps a large lead at shorter acceptance lengths as well:

Accepted tokens per second with DSpark speculative decoding at acceptance lengths 3 and 6

Without speculative decoding, the megakernels deliver roughly 1.4–2× the decode throughput of the GB200 baseline at batch sizes 1 through 8:

Aggregate decode throughput at batch sizes 1, 2, 4 and 8, TPU megakernel vs GB200 baseline

See the blog post for the full analysis and measurement setup.

Setup

Create the environment from the repository root:

uv sync

The launchers run Hugging Face in offline mode, so the model snapshots must already be present in the cache selected by HF_HOME, or supplied as explicit local paths. All scripts automatically load an optional, untracked .env file from the repository root. A typical configuration is:

HF_HOME=/path/to/huggingface
KIMI_PRESHARDED_WEIGHTS=/path/to/kimi-k3-tp32
KIMI_DRAFT_PRESHARDED_WEIGHTS=/path/to/kimi-k3-tp32/dspark

The Kimi pre-sharded variables are optional. Without them, the demo prepares its kernel layouts from moonshotai/Kimi-K3 and RedHatAI/Kimi-K3-speculator.dspark. The Qwen demo loads Qwen/Qwen3.8-27B and z-lab/Qwen3.8-27B-DFlash2; packed weight containers under checkpoints/qwen38-tp8 (or $QWEN_PACKED_WEIGHTS when set) are used instead when present. The optional KIMI_CHECKPOINT_REVISION, KIMI_DRAFT_REVISION, QWEN_CHECKPOINT_REVISION, and QWEN_DRAFT_REVISION variables pin the commits resolved from the HF_HOME cache for the default repositories. Explicit --checkpoint, --draft, --weights, and --draft-weights arguments override these defaults.

Kimi requires a four-node Slurm allocation with eight TPU devices per host. Qwen uses one host with eight TPU devices and can run on the first node of any allocation containing at least one host.

CPU correctness tests

Run the fast reference and CPU-interpreter tests with:

JAX_PLATFORMS=cpu uv run pytest -m "not cpu32"

Run the complete suite, including 32-device integration tests, with 32 virtual CPU devices:

JAX_PLATFORMS=cpu \
XLA_FLAGS=--xla_force_host_platform_device_count=32 \
uv run pytest

These tests validate numerical behavior, not TPU-specific Mosaic lowering, layouts, VMEM capacity, DMA scheduling, collectives, or performance.

OpenAI-compatible servers

Start the Kimi K3 + DSpark server on a four-node allocation:

JOBID=<four-node-job> bash scripts/demo_kimi_dspark.sh \
  --serve 0.0.0.0:8000 --context 4096 --max-tokens 512

Start the Qwen3.8-27B + DFlash2 server on the first host of an allocation:

JOBID=<slurm-job> bash scripts/demo_qwen_dflash.sh \
  --serve 0.0.0.0:8000 --context 2048 --max-tokens 512

Both servers expose GET /health, GET /v1/models, POST /tokenize, POST /detokenize, POST /v1/completions, and POST /v1/chat/completions. Kimi uses DSpark by default and supports greedy or temperature/top-p generation; --target-only serves the target megakernel without speculation. Qwen uses DFlash2 by default with greedy generation; --no-spec selects its plain single-token decoder. Omit --serve to use either launcher interactively.

Source layout

Generated weights, checkpoints, compilation caches, evaluation output, benchmark output, and Slurm logs are excluded from version control.

Citation

If you use this codebase or build on our results, please cite the companion blog post:

@misc{novack2026tpumegakernels,
  author       = {Novack, George and Liu, Xuting and Ma, Jeff and Kwon, Woosuk},
  title        = {700 TPS on Kimi K3: A Case for TPU Megakernels},
  howpublished = {\url{https://inferact.ai/blog/tpu-megakernels}},
  year         = {2026},
  month        = sep,
  note         = {Inferact blog post, September 23, 2026},
}

tpu-megakernels (Inferact/tpu-megakernels)

Searching Database…

163 ·GitHub
Tpu-megakernels
Что это?Сопроводительный исследовательский репозиторий со слитными megakernel-реализациями декодирования для двух моделей на Google TPU: Kimi K3 со спекулятивным декодированием DSpark и Qwen3.8-27B со спекулятивным декодированием DFlash2; написан в дополнение к блогпосту '700 TPS on Kimi K3: A Case for TPU Megakernels'.
Это универсальная библиотека инференса или демо-бенчмарк для этих двух моделей?Это демо-код, построенный специально вокруг двух модельных пайплайнов, а не универсальная библиотека для сервинга: у каждой модели своё слитое TPU-ядро плюс интерактивный режим в терминале и HTTP-сервер, совместимый с OpenAI API.
Как это запустить?Разверните окружение Python/JAX из корня репозитория, убедитесь, что снапшоты моделей Hugging Face закэшированы (лаунчеры запускают HF в офлайн-режиме), затем запустите demokimidspark.py или demoqwendflash.py интерактивно либо с --serve, чтобы поднять HTTP API.
Нужно ли настоящее TPU-железо (и кластер Slurm), чтобы это запустить?Да — Kimi K3 требует выделения Slurm из четырёх узлов с восемью TPU-устройствами на хост (всего 32), а Qwen3.8-27B — одного хоста с восемью TPU-устройствами; тесты корректности на CPU есть, но реальное исполнение на TPU они не проверяют.
Доказывают ли CPU-тесты, что TPU-ядра реально корректно работают на железе?Нет — набор тестов на CPU/интерпретаторе проверяет только численную корректность; он не проверяет специфичные для TPU понижение (lowering) в Mosaic, раскладки памяти, ёмкость VMEM, планирование DMA, коллективные операции и реальную производительность.
Это работает на GPU или только на Google TPU?Работает только на устройствах Google TPU; GPU GB200 фигурируют лишь как базовая линия производительности, с которой сравниваются мегаядра, а не как поддерживаемая цель исполнения.
Насколько эти мегаядра быстрее GPU?Со спекулятивным декодированием Kimi K3 выдаёт 709 принятых токенов/с на 16x TPU v7 против базовой линии из 16x GB200 (ускорение 1.57x), а Qwen3.8-27B — 1 515 токенов/с на 4x TPU v7 против 4x GB200 (2.18x). Без спекуляции TPU-ядра всё равно дают примерно в 1.4–2 раза большую пропускную способность декодирования, чем GB200, при размерах батча 1–8.
Что такое DSpark и DFlash2?Это слитые схемы спекулятивного декодирования, идущие в комплекте с каждой моделью: DSpark сочетает Kimi K3 с одной якорной и семью черновыми строками, а DFlash2 добавляет плотному BF16-декодеру Qwen3.8-27B опциональный слитый цикл черновика и спекуляции.
Что предоставляет сервер, совместимый с OpenAI?Каждый сервер предоставляет GET /health, GET /v1/models, POST /tokenize, POST /detokenize, POST /v1/completions и POST /v1/chat/completions, то есть к нему можно обращаться как к стандартному эндпоинту OpenAI API.
Зачем писать собственные TPU-мегаядра вместо стандартного фреймворка инференса?Слияние всего шага декодирования — внимания (attention), MoE, нормализации, остаточных связей и коллективных операций — в единое TPU-ядро снижает накладные расходы на диспетчеризацию каждой операции, позволяя чипам TPU v7 догонять или обгонять GPU GB200 по пропускной способности декодирования; подробности — в сопроводительном блогпосте.
What are similar projects?vllm, deltafin, DMI, jeeves, v100-skinny, kaggle-tpu-lab, TensorFold, sglang
Which topics is it in?qwen3, jax, http-server, llm-serving, kimi, speculative-decoding, inference-engine, attention-mechanism, moe, llm-inference, high-performance, google
Creator
Inferact
Size
1.4 MB
Stars
163
Updated
CommitOct 15
Age
26d
Commits · 90 Days