The TPU serving stack

Tick each objective you can deliver out loud without notes. Then write the one-paragraph explanation you would give a room of GPU engineers. Saved attempts stay below, so you can compare them over the week.

• I can describe how TPU serving got from JetStream through PyTorch/XLA and tpu-inference to the October 2026 split, and say which path to use for a new PyTorch or JAX port.
• I can explain what torchax does to a torch.nn.Linear call, why tpu-inference needs the model as a pure function that returns its KV cache, and where the code ends up as HLO.
• I can explain how continuous batching survives a compiler that wants static shapes — padded shape buckets compiled ahead, a flattened token batch, and Ragged Paged Attention's three sub-kernels.
• I can read the tpu-inference KV-cache layout and explain its packing dimension from (8,128) tiling.
• I can map tpu-inference's mesh axes to the decode sharding rules of Chapter 7.
• I can size a deployment from the Scaling Book's LLaMA-3-70B serving analysis, and solve the book's 405B problem under a latency limit.

Prompts to answer out loud

1. Explain to a vLLM-on-CUDA engineer how continuous batching works when every shape must be known at compile time.
2. What does torchax do to a torch.nn.Linear call, and where does it end up?
3. Why does the KV-cache tensor in tpu-inference have a 'packing' dimension?
4. You are porting a new PyTorch model to TPU this month. Which serving path do you choose and why?