Inferact releases K3 AgentX results on TPU

In our prior work we released our first TPU megakernels for Kimi K3 and Qwen 3.8, demonstrating how weight prefetching and memory-compute overlap can deliver high single-user interactivity. Real-world serving systems need to do more than this, they must serve many concurrent user's sessions efficiently and handle long multi-turn conversations spanning hundreds of thousands of tokens. In this post we show how TPU megakernels, integrated into vLLM TorchTPU, can efficiently serve real-world agentic traffic.

On AgentX, SemiAnalysis's public agentic benchmark, our TPU megakernel backend achieves high interactivity at low concurrency and competitive price performance across the throughput/interactivity frontier compared with the existing GB200 results.

Throughput–interactivity frontier for Kimi K3 on InferenceX as of Oct 19, 2026. Tokens per dollar uses InferenceX's external TCO: $1.21 per TPU v7 chip-hour and $1.86 per GB200 GPU-hour. Source: InferenceX by SemiAnalysis (Apache 2.0).

‍

A megakernel backend in vLLM-TorchTPU

Our first megakernels targeted low-concurrency decode at short context lengths. Benchmarking on this simple traffic pattern was useful when analyzing the peak performance of our kernels and for identifying new opportunities for optimization, but this was not a realistic representation of actual production traffic. Production systems increasingly serve multi-turn agentic sessions, where prompts are long but mostly cached, with each turn adding a relatively small number of new tokens on top of a growing shared prefix. Concurrency is bursty, because agents spawn subagents that all hit the server at once. And the large KV cache must be efficiently transferred between workers or offloaded to host RAM to free up device memory.

Our goal was to retain the megakernel's decode speed while serving real traffic through vLLM. We achieve this through a new megakernel backend in vLLM-TorchTPU. The backend is selected with --load-format megakernel. Megakernels replace the model's layer stack, reusing all of the key features that come with vLLM, like async scheduling, prefix caching, KV transfer, and the API server.

A single decode worker. MegakernelModelRunner sits alongside the standard TPUModelRunner and replaces only the model's forward path; vLLM's engine and the rest of the worker are unchanged.

For end-to-end serving, we run prefill and decode on separate TPU v7 slices, each with 16 chips in a 2×2×4 topology. A sequence-parallel prefill stack runs on one slice, a fused decoder on the other, and Raiden moves the KV cache between them. On each slice, a single unified KV pool holds attention pages, linear-attention state, and speculative-draft state.

Both slices shard the KV cache with decode context parallelism (DCP). Each sequence's cache is split into 256-token pages, so even a very long context is spread evenly across all 32 ranks. Because both roles use the same layout, Raiden can move each rank's pages directly to the same rank on the decode slice.

P/D disaggregation with decode context parallelism. Each sequence's KV cache is split into 256-token pages, and page k lives on rank k mod 32, so a long context spreads evenly across all 32 ranks of a slice. Raiden moves each rank's pages to the same rank on the decode slice.

In addition to DCP, the decoder applies three other kinds of parallelism across the 32 ranks of each slice.

Expert and tensor parallelism for mixture-of-experts

K3 routes each token to 16 of its 896 experts. The decoder splits the experts into eight groups of 112, and each group's four ranks shard its matmuls (EP8 × TP4), so a rank only streams weights for experts its group was asked for. Attention projections, shared experts and the vocabulary projection are tensor-parallel across all 32 ranks.

Sequence parallelism for token-wise operations

Between each MoE and attention layer, K3 executes attention residual and RMSNorm operations. In standard tensor-parallel deployments, tokens are fully replicated and these operations are duplicated across ranks. To improve the efficiency of these token-wise operations, we employ sequence parallelism in our decode megakernel, splitting token ownership evenly across a configurable number of ranks.

Attention DP for high-concurrency decode

At higher concurrency (48+ concurrenct requests), the ranks split into two groups of 16 that each hold a full copy of the attention weights and serve their own share of the requests, with the KV cache sharded DCP16 inside each group. Attention collectives span 16 ranks instead of 32, while the MoE layers still span all 32.

Performance results

The megakernel backend delivers strong interactivity at low and mid concurrency, and competitive price-performance for high-throughput serving. Below are the full performance metrics on the AgentX benchmark dataset across a range of concurrency levels.

Concurrency	Decode Recipe	Total tok/s	Interactivity p50 (tok/s/user)
1	MoEEP8 × TP4 AttentionDCP32 Sequence Parallel4 SpecDSpark	23,000	408
4	MoEEP8 × TP4 AttentionDCP32 Sequence Parallel4 SpecDSpark	35,032	388
8	MoEEP8 × TP4 AttentionDCP32 Sequence Parallel4 SpecDSpark	58,699	270
16	MoEEP8 × TP4 AttentionDCP32 Sequence Parallel4 SpecDSpark	102,727	162
32	MoEEP8 × TP4 AttentionDCP32 Sequence Parallel4 SpecDSpark	170,495	90
48	MoEEP8 × TP4 AttentionDP2 × DCP16 Sequence Parallel4 SpecDSpark	204,294	58
72	MoEEP8 × TP4 AttentionDP2 × DCP16 Sequence Parallel4 Specnone	210,159	31

TPU v7, one 16 + 16-chip P/D pair. Every run uses a TP32/EP32 prefill slice, FP4 weights, and Raiden KV offload to host DRAM; the recipe column shows how the decode slice is parallelised; sequence parallel 4 applies to decode steps with 8 or more rows, and smaller steps run without it.

How to run K3 with vLLM TorchTPU

With the launch of vLLM TorchTPU, our full Kimi K3 serving stack is open-source and ready to use today. See the instructions below for how to get started with the megakernel backend and how to reproduce our results.

The full step-by-step guide, including every flag we used, is in the K3 megakernel AgentX guide in the vLLM TorchTPU repo. Everything runs from one Docker image, vllm/vllm-torchtpu:nightly-20261019-k3-megakernel-c78529c. At a high level:

Pre-shard the K3 checkpoint and the DSpark draft model into the TP32 layout the megakernels expect. This step needs only the CPU.
Start a prefill server and a decode server with vllm serve, each across 4 TPU v7 hosts, using the megakernel load format and the Raiden KV transfer connector.
Put vllm-router in front of the two servers to route requests between prefill and decode.
Run the AgentX benchmark against the router with aiperf from the agentx-harness.

What's next

The K3 megakernels are now serving through vLLM TorchTPU, but there is more to do. A few directions we are taking from here:

A clean integration pattern. We want to define a clean, general pattern for integrating TPU megakernels into vLLM TorchTPU, so that new megakernels can reuse the same scheduling, KV cache management, and P/D plumbing.

Higher throughput. Our strongest results are at low and mid concurrency, and there is more performance to unlock at the high-throughput end. Directions include pipeline parallelism, aggregated (non-disaggregated) serving, and sweeping more prefill-to-decode ratios to find the best split of chips between the two roles.

New model architectures. Little of the approach is specific to Kimi K3. The same ideas, whole-model kernels with explicit control over data movement, carry over to other models, and we plan to expand our megakernel work to new architectures.

‍

Serving agentic workloads with TPU megakernels in vLLM