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…

162 ·GitHub
Tpu-megakernels
这是什么？一个配套研究仓库，包含 Google TPU 上两个模型的融合 megakernel 解码实现：使用 DSpark 投机解码的 Kimi K3，以及使用 DFlash2 投机解码的 Qwen3.8-27B，为配合博客文章《在 Kimi K3 上达到 700 TPS：论 TPU Megakernel 的价值》而写。
这是通用推理库，还是针对这两个模型的基准演示？它是专门围绕两条模型流水线构建的配套演示代码，而非通用服务库：每个模型都有各自的融合 TPU 内核，外加交互式终端模式和 OpenAI 兼容的 HTTP 服务器。
怎么让它跑起来？从仓库根目录配置好 Python/JAX 环境，确保 Hugging Face 模型快照已缓存（启动器以离线模式运行 HF），然后交互式运行 demokimidspark.py 或 demoqwendflash.py，或加 --serve 启动 HTTP API。
运行它需要真实的 TPU 硬件（以及 Slurm 集群）吗？需要——Kimi K3 需要四节点的 Slurm 分配、每台主机 8 个 TPU 设备（共 32 个），Qwen3.8-27B 需要一台配备 8 个 TPU 设备的主机；CPU 正确性测试虽然存在，但不会真正运行 TPU 执行。
CPU 测试能证明 TPU 内核在硬件上真正正确工作吗？不能——CPU/解释器测试套件只检查数值正确性；它不会验证 TPU 特有的 Mosaic lowering、内存布局、VMEM 容量、DMA 调度、集合通信或真实性能。
它能在 GPU 上运行，还是只能在 Google TPU 上运行？它只能在 Google TPU 设备上运行；GB200 GPU 仅作为 megakernel 基准对比的性能基线出现，并不是受支持的执行目标。
这些 megakernel 比 GPU 快多少？在投机解码下，Kimi K3 在 16x TPU v7 上达到每秒 709 个已接受 token，对比 16x GB200 基线（1.57 倍）；Qwen3.8-27B 在 4x TPU v7 上达到每秒 1,515 个 token，对比 4x GB200（2.18 倍）。不使用投机时，TPU 内核在批量大小 1-8 下仍能提供约 1.4-2 倍于 GB200 的解码吞吐。
DSpark 和 DFlash2 是什么？它们是随各模型捆绑的融合投机解码方案：DSpark 为 Kimi K3 配备一个锚点行和七个草稿行，而 DFlash2 为 Qwen3.8-27B 的稠密 BF16 解码器提供可选的融合草稿/投机循环。
OpenAI 兼容服务器提供哪些接口？每个服务器都提供 GET /health、GET /v1/models、POST /tokenize、POST /detokenize、POST /v1/completions 和 POST /v1/chat/completions，因此可以像标准 OpenAI API 端点一样查询。
为什么要构建自定义 TPU megakernel，而不是使用标准推理框架？把整个解码步骤——注意力、MoE、归一化、残差和集合通信——融合进单个 TPU 内核，可以削减逐算子分发开销，让 TPU v7 芯片在解码吞吐上追平甚至超过 GB200 GPU，详见配套博客文章。
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
162
Updated
CommitOct 15
Age
25d
Commits · 90 Days