The
@inferact team open sourced a TPU megakernel for Kimi K3 achieving 709 tokens/s, against 450 tokens/s on GB200.
All 92 of K3's MoE layers run in a single Pallas kernel, with weight prefetching that reaches across layer boundaries so transfers for one layer overlap with computation in the previous one.
Shoutout to the team! Writeup and repo link in the thread.
Our first TPU megakernel for Kimi K3 reaches 709 tokens/s on low-concurrency decode, against 450 tokens/s for our GB200 baseline, both with DSpark speculative decoding.
To our knowledge, this is the first TPU inference megakernel. The whole model runs in a single Pallas kernel, and without spec decoding it is roughly 1.4 to 2x the GB200 baseline at batch sizes 1 through 8.
We are open sourcing it today.
1/2
もっと見る