The TPU serving stack

Chapter 7 derived what a serving system should do. This chapter is about the software that does it on TPU, which changed three times in 18 months. It ends with the Scaling Book's serving case study.

How TPU serving got here

• JetStream (2024–2025). Google's engine, with prefill and generate engines on separate slices and an interface of prefill, insert and generate [SB7]. Archived on 2026-02-01; its functions moved to vLLM TPU [R-jetstream].
• vLLM on PyTorch/XLA (to 2025). vLLM's first TPU backend ran PyTorch through PyTorch/XLA's lazy tensors. PyTorch/XLA is now deprecated, and its repository is archived on 2026-11-15 [T-xla].
• tpu-inference, a unified JAX backend (October 2025). One lowering path for all models: JAX-native models where they exist, and upstream vLLM PyTorch models through torchax otherwise. Against the PyTorch/XLA backend, Google reported 3.6× on Llama-3.1-8B on one v6e chip and 2.1× on 70B on eight. About 20% of the gain came from JAX lowering and about 10% from Ragged Paged Attention v3 over v2 [B-vllmtpu].
• The split (October 2026). tpu-inference is being renamed vLLM JAX and focuses on JAX (Flax NNX) models. PyTorch model definitions get vllm-torchtpu, described as "the recommended path for serving PyTorch model definitions on TPU". SGLang made the same split [R-tpuinf] [T-engines].

Which path for a new port.

• A JAX implementation of the model exists in vLLM JAX: use it. tpu-inference picks it by default (MODEL_IMPL_TYPE=auto resolves to flax_nnx) [R-tpuinf]. These implementations are written for TPU sharding and kernels.
• You are bringing PyTorch model code: vllm-torchtpu is now the recommended path. It runs upstream vLLM model definitions through torch.compile on TorchTPU, with TPU kernels inserted as custom ops [R-vllm-torchtpu].
• You have an existing torchax port: it continues to run (MODEL_IMPL_TYPE=vllm). Google's blog found torchax-lowered and hand-written JAX models performed about the same [B-vllmtpu], but new PyTorch work is going to vllm-torchtpu.

The scheduler is vLLM's

Neither TPU backend replaces vLLM's engine. The API server, the v1 scheduler that implements continuous batching, the KV block manager and prefix caching are upstream vLLM [R-tpuinf] [T-engines]. The TPU code supplies a platform plugin, a worker and a model runner. In tpu-inference the runner directory has a tpu_runner, a kv_cache_manager, a block_table, a persistent_batch_manager, a compilation_manager and managers for speculative and structured decoding [R-tpuinf].

torchax: PyTorch code, JAX execution

torchax lets PyTorch model code run as JAX [R-torchax]:

• A torchax.tensor.Tensor is a torch.Tensor subclass that holds a jax.Array.
• It overrides __torch_dispatch__. When PyTorch dispatches an ATen operation, such as the aten::mm inside torch.nn.Linear, torchax unwraps the JAX arrays, runs a JAX implementation from its op registry (a jnp or lax call), and wraps the result.
• Under jax.jit, those JAX calls are traced into a jaxpr like any other JAX code, lowered to StableHLO and compiled by XLA.

A Linear layer becomes a jnp.matmul in a jaxpr, then a dot_general in HLO, then an MXU fusion in the compiled program. JAX shardings, Pallas kernels and Qwix all apply to it.

tpu-inference wraps a PyTorch model in what its developer guide calls a "sandwich" [R-torchax-guide]:

• The middle is upstream vLLM's PyTorch model code.
• Around it is JAX: inputs, the sampler, the KV cache and the Ragged Paged Attention kernel.
• torch_view and jax_view convert at the boundaries.

The worker calls one pure function, model_fn(weights, kv_caches, inputs). PyTorch modules are made pure with torch.func.functional_call, and the KV cache is passed in and returned, functional-update style, rather than mutated in place. The guide's summary: "torchax is JAX".

Continuous batching with static shapes

Continuous batching means the set of sequences in each step changes every step: new requests join, finished ones leave, some are in prefill and some in decode. A TPU program is compiled for exact shapes. tpu-inference reconciles the two in three ways.

Flatten the batch into tokens. Every query token in the step, from every sequence, is packed into one array: queries are [max_num_tokens, num_q_heads, head_dim]. Per-sequence structure is carried in small metadata arrays [R-tpuinf, kernels/ragged_paged_attention/v3]:

• cu_q_lens, cumulative query lengths that mark where each sequence's tokens start;
• kv_lens, each sequence's cache length;
• page_indices, the block table;
• distribution, which counts how many sequences are decode-only, prefill-only and mixed.

A decode sequence contributes one token and a prefill chunk contributes many, in the same array.

Pad to buckets and compile ahead. The number of tokens and of sequences is padded up to one of a fixed set of bucket sizes. The compilation_manager compiles every bucket before serving starts, so no request waits on a compile [R-tpuinf]. The cost is padding waste, which the bucket spacing controls, and startup time.

Make attention handle raggedness inside the kernel. Ragged Paged Attention (RPA) is a Pallas kernel that reads each sequence's KV pages through the block table and attends over that sequence's true length [P-RPA]. Version 3 adds three things:

• It fuses the KV-cache update into the attention kernel, which "completely hides scatter latency".
• It compiles three specialized sub-kernels, for decode-only, prefill-only and mixed batches, chosen by the distribution metadata [B-vllmtpu].
• The paper reports up to 86% of memory bandwidth in decode and 73% of FLOPs in prefill on Llama-3-8B on Ironwood [P-RPA].

The KV-cache layout

Per layer, tpu-inference allocates the cache as [R-tpuinf, runner/kv_cache.py]

Pages of tokens; within a page, K and V heads interleaved and packed p to a 32-bit word.

K and V are interleaved along the head axis (the factor 2), then grouped in fours for 8-bit values or pairs for bf16. The packing dimension exists because of Chapter 2's tiling. XLA stores bf16 as (8,128)(2,1) and 8-bit types as (8,128)(4,1), two or four values per 32-bit sublane [X-tiled]. Laying the cache out that way lets the kernel load whole 32-bit words that already hold the right K and V values. The DMA then moves full tiles, and no relayout is needed between the cache and the MXU.

Sharding in tpu-inference

tpu-inference defines its mesh axes as data, attn_dp, attn_dp_expert, expert, model, dcp and pcp. The default is a 2D mesh of ('data', 'model'), in which attention heads, KV heads, the MLP's tensor dimension and experts all shard over model, and the batch over data [R-tpuinf, layers/common/sharding.py]. That is Chapter 7's decode sharding:

• model parallelism over model;
• replicas over data;
• dcp and pcp for decode and prefill context parallelism when sequences are long;
• expert for expert parallelism. The source does not document attn_dp; its name suggests data parallelism for attention, the usual answer when a model has fewer KV heads than the model axis has chips unverified.

Multi-slice deployments set NUM_SLICES.

Quantization and MoE

tpu-inference loads checkpoints already quantized in several formats: MXFP4, compressed-tensors FP8 and INT8, AWQ INT4 and NVIDIA ModelOpt NVFP4. It can also quantize JAX models at load time with Qwix rules (Chapter 7) [R-tpuinf]. For MoE it carries Pallas kernels for fused MoE and grouped matmul (megablox), and SparseCore kernels for the ragged gathers and scatters of token routing [R-tpuinf, kernels/].

vllm-torchtpu provides the same kernels as custom ops inside its compiled graph [R-vllm-torchtpu] [T-engines]:

• RPA, dense and sparse MLA, Gated DeltaNet, and fused MoE with SparseCore offload;
• FP8, INT8 and MXFP4 quantized matmuls and FP8 KV cache;
• TP, DP, PP, EP and context parallelism;
• disaggregated prefill and decode through TPU Sync, which reads the device buffers behind PyTorch tensors directly;
• KV offload to host DRAM or Managed Lustre.

Google also describes a Pallas "megakernel" that runs a whole decode step of Kimi K3 as a single kernel.

Case study: serving LLaMA-3-70B

The Scaling Book's Chapter 8 works through serving LLaMA-3-70B on v5e, chosen because it was the cheapest TPU per FLOP at the time [SB8]:

The book's steps, with int8 weights and bf16 math:

• KV cache: $2 \cdot 8 \cdot 128 \cdot 80 = 164$ kB per token in int8; the book rounds to 160 kB. A batch of 32 at 8k context needs 113 GB with the weights, so 7 chips minimum: a 4×2.
• Step time on a 4×2: 113 GB over 8 × 820 GB/s = 17 ms, 1,860 tokens/s, about 232 tokens/s per chip. On a 4×4: 8.6 ms, same per-chip throughput. The book gives 17 ms and 235.
• Sharding check: bf16 weights on a 4×8 at batch 64 give $T_\text{ICI}$ = 12 µs, $T_\text{HBM}$ = 18 µs and $T_\text{math}$ = 5 µs. HBM-bound, so 32-way model parallelism lowers latency without raising throughput.
• The smallest slice is not the cheapest. At 8k context the KV cache, not the weights, limits batch. Doubling the slice more than doubles the batch that fits and raises throughput per chip by about 1.6× in the book's table.
• Latency against cost. In the book's interactive plot, doubling per-token latency can cut cost per token by about 100×, and at any context above 2,048 tokens KV loading takes longer than FLOPs [SB8].