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

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:

• Kimi K3 + DSpark runs on 32 TPU devices across four hosts. One megakernel executes the KDA, MLA, MoE, residual, normalization, and collective work for the target model; DSpark provides fused speculative decoding with one anchor and seven draft rows.
• Qwen3.8-27B + DFlash2 runs on eight TPU devices on one host. Its dense BF16 megakernel covers the hybrid gated-delta-net and grouped-query-attention stack, with an optional fused DFlash2 speculative-decoding loop.

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

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×):

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

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

See the blog post for the full analysis and measurement 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.

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.

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.

• kimi/decode_megakernel.py: fused Kimi K3 target decoder, including KDA, MLA, MoE, residual/normalization, and collective communication.
• kimi/dspark.py: DSpark draft model, fused speculation step, cache management, and draft-weight loading.
• kimi/load.py: Kimi tokenizer, checkpoint streaming, and TP32 weight-layout preparation.
• qwen/__init__.py: Qwen model configuration and reference equations.
• qwen/load.py: Qwen tokenizer, checkpoint loading, TP weight packing, packed containers, and device placement.
• qwen/decode_megakernel.py: Qwen TP-sharded prefill, fused decode, and block verification.
• qwen/dflash.py: DFlash2 reference and fused draft/speculation kernels.
• demo_kimi_dspark.py and demo_qwen_dflash.py: model loading, compilation, prefill, generation, metrics, and server orchestration.
• collectives32.py and pool_alias.py: TPU collective primitives and Pallas VMEM aliasing support used by Kimi.
• openai_server.py and distributed_entry.py: shared HTTP protocol support and Kimi multi-host process startup.
• model_paths.py: local-path and Hugging Face repository resolution through HF_HOME.
• scripts/: Slurm launchers, correctness evaluations, and performance benchmarks for both model families.
• tests/: CPU reference, packing, Pallas-interpreter, and integration tests.

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

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},
}