TPU Megakernel Achieves 709 tokens/s on Kimi K3, Outpacing GB200

An Open-Source Megakernel That's Turning Heads

The @inferact team recently open-sourced a TPU megakernel for Kimi K3, claiming an inference throughput of 709 tokens/s — significantly higher than the 450 tokens/s achieved on the GB200 platform. That number alone is striking: it suggests that, under certain workloads, a deeply optimized TPU kernel can match or even surpass NVIDIA's latest high-end GPU.

For teams focused on large model inference costs and deployment efficiency, the value of this kind of engineering work lies less in any single benchmark number and more in what it reveals about pushing hardware to its absolute limits. The core idea behind a megakernel is to merge many separate operator calls into one large kernel, thereby reducing launch overhead, memory round-trips, and scheduling latency.

The Key Innovation: 92 MoE Layers in a Single Pallas Kernel

According to information shared by the team, all 92 MoE (Mixture of Experts) layers in Kimi K3 run inside a single Pallas kernel. Pallas is a programming interface within the JAX ecosystem for writing custom TPU/GPU kernels, giving developers low-level control over memory layout and compute scheduling.

Packing all 92 layers into one kernel eliminates per-layer kernel launch overhead and gives the compiler and runtime more room to apply aggressive global optimizations. For MoE architectures, expert routing and sparse computation naturally generate a large number of fragmented small operators. Scheduling these one by one incurs substantial overhead — fusing them into a megakernel is a targeted fix for exactly this problem.

Cross-Layer Weight Prefetching

Even more noteworthy is the weight prefetching mechanism. The implementation prefetches weights across layer boundaries — meaning that while the current layer is being computed, the weights needed for the next layer are already being transferred from slow storage to fast storage. This overlaps data movement with computation, hiding what would otherwise be a sequential, blocking memory transfer.

For MoE models with massive parameter counts but only a fraction of experts activated per forward pass, weight loading is often a key bottleneck. By overlapping compute and transfer, the kernel keeps compute units continuously busy rather than stalling on data fetches — and this is a major engineering detail behind hitting 709 tokens/s.

---

MoE (Mixture of Experts) is a sparsely activated neural network architecture: the overall model has a large parameter count, but for each input token, only a few "expert" sub-networks are selected by a router and participate in computation — the rest remain dormant during that forward pass. This design allows models to scale parameter counts dramatically without a proportional increase in inference compute. Kimi K3 uses this architecture; with 92 layers, each layer has its own set of experts with independent routing decisions. The challenge of sparsity is that activated experts are distributed irregularly, which tends to generate many small, fragmented matrix multiply calls. Under traditional per-layer scheduling, this means frequent kernel launches and significant compute unit idle time.

Weight prefetching is essentially a "double-buffering" or pipelining technique with a long history in computer architecture. The idea is to split the memory hierarchy into two tiers — slow (large capacity, low bandwidth, e.g., HBM or off-chip DRAM) and fast (small capacity, high bandwidth, e.g., SRAM or on-chip cache) — and asynchronously move the next batch of data into fast storage while the compute units are consuming the current batch, running the compute and transfer pipelines in parallel. For large MoE models, a single expert's weight matrix can be several GB, and the full model's weights far exceed on-chip storage. Without prefetching, load latency is fully exposed on the critical path, causing a sharp drop in compute utilization. Cross-layer prefetching extends this idea beyond individual layer boundaries: as long as the compiler or runtime can statically or dynamically anticipate which weights the next layer needs, the transfer time can be almost entirely hidden.

---

Why This Comparison Matters

The 709 vs. 450 tokens/s comparison has attracted attention because it touches on a core question in today's inference hardware landscape: outside of NVIDIA GPUs, can alternative hardware like TPUs deliver competitive cost-efficiency on real workloads?

That said, a single throughput number is influenced by many factors — batch size, precision settings, context length, whether latency is measured end-to-end, and more. This result is best understood as "what a heavily optimized TPU kernel can achieve under specific conditions," not as a general verdict on hardware superiority. The original information comes from a single source (@inferact's tweets and repository); the full writeup and code repository link are in their thread for developers who want to reproduce and verify the results independently.

---

The GB200 refers to NVIDIA's Blackwell-architecture Grace Blackwell Superchip — currently NVIDIA's flagship product for datacenter inference, integrating a Grace CPU and B200 GPU via NVLink-C2C interconnect. The 450 tokens/s baseline represents throughput for running the same model (Kimi K3) on that platform. However, the two systems differ fundamentally in hardware configuration, memory bandwidth, interconnect topology, and software stack, so this comparison is best read as "the relative position TPUs can reach with extreme optimization" rather than a direct apples-to-apples competitive benchmark.

---

Takeaways for Developers

This kind of open-source work offers several direct insights for the inference optimization community:

• Operator fusion is a high-value direction: Merging many small operators into a megakernel significantly reduces scheduling and launch overhead — especially well-suited to fragmented, compute-intensive architectures like MoE.
• Memory movement optimization can't be overlooked: Cross-layer weight prefetching demonstrates how to hide I/O latency behind computation — a general-purpose optimization pattern for large-parameter model inference.
• Hardware diversity is growing: The TPU + Pallas combination shows that high-performance inference stacks around non-NVIDIA hardware are maturing, giving deployment teams more options.

For teams evaluating inference infrastructure, this open-source implementation is both a directly reusable engineering reference and a reminder: hardware capability is often constrained by the maturity of the software stack, and the real performance gap is frequently closed — or opened — through precise kernel-level optimization.