Inferact released open-source TPU megakernels for Kimi K3. The implementation hits over 700 tokens per second with speculative decoding. It beats a GB200 baseline by nearly two times at batch size one. The megakernel fuses the full decoder into one Pallas program. It prefetches weights across layer boundaries using TPU’s large on-chip VMEM. Asynchronous DMAs keep HBM busy throughout each decode step. The kernel shards the model across 16 TPU v7 chips. It compiles in under 90 seconds. End-to-end validation scored 0.944 on GPQA-Diamond and 0.972 on GSM8K. The code and reproduction scripts are public on GitHub.
What changed
Decode steps used separate kernels with bandwidth underutilization between phases.
What it unlocks
Running Kimi K3 at over 700 tokens/s on TPU v7 with speculative decoding.
- 700 tokens/s with speculative decoding
- 452 tokens/s on GB200
- 1.4 to 2× decode throughput vs GB200 at batch 1-8
- 8.5 ms per decode step
What you need to act on it
- TPU v7 hardware
- Kimi K3 model
- Pallas kernel language
- 16 TPU v7 chips for full throughput
Sources