Inferact open-sources TPU inference megakernel, with Kimi K3 exceeding 700 tokens per second in low-concurrency decoding

On September 23, Inferact released an inference megakernel targeting TPU v7, implementing Kimi K3 decoding in Pallas. The company measured more than 700 output tokens per second for a single user on 16 TPU chips, compared with 452 tokens per second on 16 GB200 GPUs using a vLLM recipe. The chart uses speculative decoding with an acceptance length of 6, a condition that should be considered when interpreting the results. [Official technical article](https://inferact.ai/blog/tpu-megakernels)

The implementation fuses 92 MoE layers into a single Pallas call, prefetching the next layer’s attention projection weights while executing expert computation in the current layer. The aim is to overlap computation, communication, and weight transfers, reducing idle time on the high-bandwidth memory interface during decoding. This is particularly relevant for small-batch serving: even when generating only a few tokens at a time, large amounts of model weights must still be loaded repeatedly. [Kernel design overview](https://inferact.ai/blog/tpu-megakernels)

The control Pallas provides is key to understanding this optimization. JAX documentation explains that TPUs primarily execute instructions sequentially, but can perform DMA transfers and matrix operations in the background. Kernels use on-chip buffers such as VMEM to reduce time spent waiting directly on HBM. Engineers can therefore control data lifetimes and prefetch order, but compilation will still fail if buffer allocations exceed capacity. Fusion alone does not guarantee a speedup. [Pallas documentation](https://docs.jax.dev/en/latest/pallas/tpu/details.html)

This also offers a concrete case study for compiler research: the ordering of data transfers across layers may have a greater impact on results than local optimization of individual operators. The design suggests that switching to another model would change weight sizes, attention state, and temporary storage requirements, so the existing schedule may not transfer directly. Anyone reproducing the results needs to check both numerical correctness and memory scheduling to identify where the gains originate.

Reproducing the results also requires distinguishing physical chips from devices exposed by the framework. Google’s documentation states that each Ironwood chip exposes two devices to JAX, and that a 16-chip configuration spans four hosts. The repository’s reference to 32 TPU devices therefore matches the article’s 16 chips. Deployment teams should verify the topology rather than compare device counts alone. [Google architecture documentation](https://docs.cloud.google.com/tpu/docs/tpu7x)

The public code includes an interactive demo, an HTTP server, and correctness tests. However, the README explicitly states that CPU tests do not cover TPU compilation, VMEM capacity, DMA scheduling, or performance. The implementation also targets a 2×2×4 topology; communication and vector operations could become new bottlenecks at high concurrency. For serving teams, the next step is to hold the model version and draft acceptance rate constant, then measure performance with long contexts, time to first token, and workloads with multiple requests. The current figures are insufficient to estimate total serving costs. [Code and test coverage](https://github.com/Inferact/tpu-megakernels), [Applicability and limitations](https://inferact.ai/blog/tpu-megakernels)