700 TPS on Kimi K3: A Case for TPU Megakernels

Written byGeorge Novack, Xuting Liu, Jeff Ma, and Woosuk Kwon

Today we released inferact/tpu-megakernels, a collection of megakernels for TPU v7. Our Kimi K3 implementation delivers over 700 tokens/s with speculative decoding, compared with 452 tokens/s on GB200. Without speculative decoding, our megakernels for K3 and Qwen 3.8 27B deliver roughly 1.4 to 2× the decode throughput of the GB200 baseline at batch sizes 1 through 8.

Decode speed is bound by the hardware's peak memory bandwidth. Our megakernels push end-to-end performance closer to this theoretical peak through effective use of TPU's large on-chip VMEM, sequential programming model, and explicit asynchronous pipelines to prefetch weights across layer boundaries. This post explains our design and why megakernels and TPUs are a good fit.

Background — TPUs vs. GPUs

Considering only the top-level specs, Google's TPU v7 Ironwood chip is comparable to NVIDIA's GB200 GPU.

Sources: Google Cloud TPU7x docs, NVIDIA GB200 NVL72.

This apples-to-apples comparison glosses over one of the TPU's most interesting features: the shape and size of its on-chip memory. Decode speed depends mostly on moving data efficiently from HBM to fast on-chip memory, which makes those dimensions a key design consideration.

Compared to a GPU's multi-level memory hierarchy of hardware-managed L1/L2 caches and program-managed Shared Memory and Tensor Memory spaces, a TPU's memory hierarchy is simple. Data from HBM is transferred to a large pool of on-chip SRAM called Vector Memory (VMEM); the TPU's vector registers (VREGs) are populated directly from VMEM. Importantly, the lifespan of all data in VMEM is managed entirely by our TPU program: when we move some data to VMEM, it stays there until we explicitly overwrite it with something else.

With 64 MiB of VMEM per TensorCore, we can prefetch most model weight tensors, ensuring they are staged in fast on-chip SRAM by the time any operation needs to use them. The GPU's Tensor Memory (TMEM) is organized differently. An entire GPU holds ~38 MiB of it, but that capacity is divided among all 152 SMs, and 256 KB per SM leaves little room to stage weights ahead of time.

Why megakernels?

Most kernels have the same rough outline: load weights and activations from HBM, perform some computation, then store the result back. Only the load phase pulls a large amount of data from HBM, leaving bandwidth underutilized during the other phases. Well-designed kernels mitigate these bandwidth bubbles by subdividing and overlapping each phase, but only to a degree. A single decode step might dispatch hundreds of kernels like this, one after another, and these bubbles add up.

GPUs already provide ways to address this. CUDA Graphs reduce host launch overhead, streams allow independent branches to execute concurrently, and programmatic dependent launch (PDL) lets a dependent kernel begin preparation and weight loading before its producer finishes. These mechanisms recover some overlap while retaining separate kernels.

A megakernel removes the boundary. With the whole decoder step inside one program, a weight load is no longer tied to the kernel that uses it — it can be issued as early as on-chip storage allows. The binding constraints become buffer lifetimes and capacity. On TPU v7, large VMEM and explicit asynchronous transfers let us stage the next layer's weights while the current layer is still computing or communicating. The goal is to keep the HBM interface busy throughout a decode step whose cost is dominated by moving weights.

Existing TPU implementations rely on XLA to perform fusion and scheduling, but finding an optimal cross-layer pipeline requires joint decisions about storage reuse, transfer timing, and synchronization. Automatic optimization may not discover the schedule we want. Pallas lets us express those decisions directly while leaving instruction scheduling to the compiler.

Writing TPU megakernels with Pallas

Pallas is the kernel language for TPUs. To our knowledge this is the first open-source inference megakernel written in Pallas, and we found the language surprisingly well suited to it. Below are the Pallas features that mattered most in this work.

Grid-less kernel launch

The Pallas language can be used to write kernels for both TPU and GPU. As such, it is possible to write Pallas kernels in a way that will feel familiar to any GPU kernel engineer: we can launch our kernels as a grid of programs, with each program receiving and operating on a specific tile of the input.

TPUs, though, follow an inherently single-threaded execution model. Even if we launch our kernel as a grid, each individual program is still run sequentially, one at a time. While the grid-based programming model has its advantages, we found that for megakernel development, embracing TPU's single-threaded nature is the way to go for both readability and fine-grained control. Our megakernel is launched as a single program with no explicit grid.

Scoped VMEM allocations

Our megakernel must balance aggressive prefetching of model weights against the limited pool of 64 MiB VMEM per TensorCore. Some data, like activations, remain resident in VMEM throughout the entire forward pass, while others, like the KV cache used in MLA, have a very short lifetime. The Pallas run_scoped primitive enables fine-grained control over the lifetime of specific VMEM allocations. This simplifies the task of the programmer: rather than explicitly managing which byte ranges in VMEM are safe to read and write at any point in time, we can write self-contained "mini-kernels" that define exactly the shape and size of the VMEM they need.

Asynchronous DMAs

Asynchronous DMAs between HBM and VMEM are what make the megakernel possible. Our implementation makes heavy use of make_async_copy(...).start() and make_async_copy(...).wait() to specify when weight prefetching should be initiated and when to pause execution until it completes. Managing asynchronous dependencies between operations is notoriously awkward in JAX: XLA schedules chains of operations as a unit and reserves the right to reorder them or reclaim their buffers. Inside a single kernel the problem does not arise: a copy we start and wait on never crosses a boundary XLA is free to schedule across.

Our Kimi K3 megakernel

The Kimi K3 megakernel combines K3's 92 Mixture-of-Experts (MoE) layers into a single Pallas call, conditionally choosing between Kimi Delta Attention (KDA) and Multi-head Latent Attention (MLA) operations based on the layer index. Inside the kernel, every weight tensor is staged into VMEM by a transfer issued during an earlier phase so that it is, ideally, already resident when its consumer runs. Because the layer loop lives inside the kernel, asynchronous copies can reach across layer boundaries. For example, layer N+1's attention projection weights start moving while layer N is still in its MoE phase. A simplified view of the cross-layer prefetching strategy is shown below.

The kernel shards the model across 16 TPU v7 chips, a total of 32 TensorCores. Attention heads are split across all 32 ranks, while routed experts use TP4 × EP8. The MoE latent-to-hidden projection is sharded within each host (8 TensorCores) and replicated across hosts, trading additional weight reads for cheaper collective communications.

Tuning megakernels for a specific target request concurrency requires exploring more bespoke means of parallelizing work across chips beyond the familiar tensor, expert, and data parallel strategies. For example, even the seemingly modest increase in decode batch size from 1 to 8 exposed new bottlenecks in operations that primarily perform vector arithmetic on the residual stream (e.g., attention residuals and RMSNorm); to address this, we combine tensor parallelism with residual stream sequence parallelism to split vector arithmetic in addition to matrix multiplication across TPUs.

One benefit of the megakernel approach that we did not foresee is fast compilation. In hindsight it makes sense: the kernel author pays by hand, and only once, for the ordering and overlap that XLA must otherwise rediscover across every operation on every new compile. On TPU, compiling a large model made up of hundreds or thousands of XLA ops can regularly take 30+ minutes; our entire megakernel, on the other hand, compiles from scratch in less than 90 seconds, enabling much more rapid development and experimentation cycles.

DSpark

With support for batch size 8, the megakernel can also serve as the verifier when using DSpark speculative decoding. Each kernel launch evaluates 1 anchor and 7 proposed continuations in a single step. Depending on the exact workload, the draft model typically yields between 3 and 6 accepted tokens per step. At ~8.5 ms per decode step, the megakernel can deliver over 700 tokens per second to a single user.

Results

We evaluate the megakernel on 16 TPU v7 Ironwood chips and take vLLM's published Kimi K3 recipe on 16 GB200 GPUs as our baseline. With no speculative decoding, the megakernel leads at every batch size from 1 through 8, delivering nearly twice the throughput at batch size 1.

We validated the accuracy of our implementation end to end by exposing it through a simple API server. With greedy decoding and max reasoning effort, the megakernel scores 0.944 on GPQA-Diamond and 0.972 on GSM8K. Scripts to reproduce these results are included alongside the implementation.

Future work

This kernel is a first step. The ideas it is built from — whole-model visibility, explicit VMEM staging, transfers issued layers ahead of the phase that consumes them — are not specific to Kimi K3, and most of the work ahead is about making them reusable and pushing them further. A few directions we are taking from here:

More workloads. This work shows that megakernels can be effective for small-batch decode, but production traffic can take many more shapes. Higher concurrency moves the bottleneck, since vector work and cross-device traffic grow with the batch size, and agentic workloads make KV cache movement and management crucial for performance.

More topologies. Our implementation is tailored to a 2×2×4 TPU chip topology. Supporting other shapes means writing new collectives, since our communication and parallelization strategies are built around how those 16 chips are wired. On NVIDIA's rack-scale systems every GPU in an NVLink domain is equally close, so the only decision is how many to use; on TPU we choose the topology as well, which adds a dimension to the search.

Better tooling. Pallas worked well for writing the kernel. Some bottlenecks, though, only become visible through direct analysis of compiled instruction bundles. Tooling that points at those inefficiencies, for people and for the coding agents that will write more of these kernels, would matter a lot as this work grows.

What Inferact does

This work comes out of what we do at Inferact. We optimize inference on our customers' infrastructure, which means optimizing against the specific hardware, model, and workload shape in front of us. Kernel work is one example of the inference optimizations we provide in the platform and to our customers, and because the platform builds on open-source vLLM, customers pick these up without having to change how they already serve.