What is this?
A companion research repo with fused 'megakernel' decode implementations for two models on Google TPUs: Kimi K3 with DSpark speculative decoding, and Qwen3.8-27B with DFlash2 speculative decoding, written to accompany the blog post '700 TPS on Kimi K3: A Case for TPU Megakernels'.
A companion research repo with fused 'megakernel' decode implementations for two models on Google TPUs: Kimi K3 with DSpark speculative decoding, and Qwen3.8-27B with DFlash2 speculative decoding, written to accompany the blog post '700 TPS on Kimi K3: A Case for TPU Megakernels'.Is this a general-purpose inference library or a benchmark demo for these two models?It is companion demo code built specifically around two model pipelines rather than a general-purpose serving library: each model gets its own fused TPU kernel plus an interactive terminal mode and an OpenAI-compatible HTTP server.
It is companion demo code built specifically around two model pipelines rather than a general-purpose serving library: each model gets its own fused TPU kernel plus an interactive terminal mode and an OpenAI-compatible HTTP server.How do I get this running?Set up the Python/JAX environment from the repository root, make sure Hugging Face model snapshots are cached (launchers run HF in offline mode), then run demokimidspark.py or demoqwendflash.py interactively or with --serve to start the HTTP API.
Set up the Python/JAX environment from the repository root, make sure Hugging Face model snapshots are cached (launchers run HF in offline mode), then run demokimidspark.py or demoqwendflash.py interactively or with --serve to start the HTTP API.Do I need real TPU hardware (and a Slurm cluster) to run this?Yes — Kimi K3 needs a four-node Slurm allocation with eight TPU devices per host (32 total), and Qwen3.8-27B needs one host with eight TPU devices; the CPU correctness tests exist but don't exercise real TPU execution.
Yes — Kimi K3 needs a four-node Slurm allocation with eight TPU devices per host (32 total), and Qwen3.8-27B needs one host with eight TPU devices; the CPU correctness tests exist but don't exercise real TPU execution.Do the CPU tests prove the TPU kernels actually work correctly on hardware?No — the CPU/interpreter test suite only checks numerical correctness; it does not validate TPU-specific Mosaic lowering, memory layouts, VMEM capacity, DMA scheduling, collectives, or real performance.
No — the CPU/interpreter test suite only checks numerical correctness; it does not validate TPU-specific Mosaic lowering, memory layouts, VMEM capacity, DMA scheduling, collectives, or real performance.Does this run on GPUs, or only on Google TPUs?It only runs on Google TPU devices; GB200 GPUs appear solely as the performance baseline the megakernels are benchmarked against, not as a supported execution target.
It only runs on Google TPU devices; GB200 GPUs appear solely as the performance baseline the megakernels are benchmarked against, not as a supported execution target.How much faster are these megakernels than GPUs?With speculative decoding, Kimi K3 reaches 709 accepted tokens/sec on 16x TPU v7 vs a 16x GB200 baseline (1.57x), and Qwen3.8-27B reaches 1,515 tokens/sec on 4x TPU v7 vs 4x GB200 (2.18x). Without speculation, the TPU kernels still deliver roughly 1.4-2x GB200 decode throughput at batch sizes 1-8.
With speculative decoding, Kimi K3 reaches 709 accepted tokens/sec on 16x TPU v7 vs a 16x GB200 baseline (1.57x), and Qwen3.8-27B reaches 1,515 tokens/sec on 4x TPU v7 vs 4x GB200 (2.18x). Without speculation, the TPU kernels still deliver roughly 1.4-2x GB200 decode throughput at batch sizes 1-8.What are DSpark and DFlash2?They are the fused speculative-decoding schemes bundled with each model: DSpark pairs Kimi K3 with one anchor and seven draft rows, while DFlash2 gives Qwen3.8-27B's dense BF16 decoder an optional fused draft/speculation loop.
They are the fused speculative-decoding schemes bundled with each model: DSpark pairs Kimi K3 with one anchor and seven draft rows, while DFlash2 gives Qwen3.8-27B's dense BF16 decoder an optional fused draft/speculation loop.What does the OpenAI-compatible server expose?Each server exposes GET /health, GET /v1/models, POST /tokenize, POST /detokenize, POST /v1/completions, and POST /v1/chat/completions, so it can be queried like a standard OpenAI API endpoint.
Each server exposes GET /health, GET /v1/models, POST /tokenize, POST /detokenize, POST /v1/completions, and POST /v1/chat/completions, so it can be queried like a standard OpenAI API endpoint.Why build custom TPU megakernels instead of using a standard inference framework?Fusing the whole decode step — attention, MoE, normalization, residuals, and collectives — into a single TPU kernel cuts per-op dispatch overhead, letting TPU v7 chips match or beat GB200 GPUs on decode throughput, as detailed in the companion blog post.
Fusing the whole decode step — attention, MoE, normalization, residuals, and collectives — into a single TPU kernel cuts per-op dispatch overhead, letting TPU v7 chips match or beat GB200 GPUs on decode throughput, as detailed in the companion blog post.What are similar projects?vllm, deltafin, DMI, jeeves, v100-skinny, kaggle-tpu-lab, TensorFold, sglang
vllm, deltafin, DMI, jeeves, v100-skinny, kaggle-tpu-lab, TensorFold, sglangWhich topics is it in?qwen3, jax, http-server, llm-serving, kimi, speculative-decoding, inference-engine, attention-mechanism, moe, llm-inference, high-performance, google
qwen3, jax, http-server, llm-serving, kimi, speculative-decoding, inference-engine, attention-mechanism, moe, llm-inference, high-performance, google