AI
Four Stages of LLM Development
Building a large language model has four phases: design, pre-training, post-training, and inference. In pre-training, the model learns to predict the next token. Post-training teaches it to be useful through supervised fine-tuning and reinforcement learning for reasoning. Inference now consumes more compute as thinking models generate many tokens. Google designed TPU hardware with two distinct families: V8T for training throughput and V8I for inference latency.
Inference on TPUs with vLLM
vLLM solves inference bottlenecks like KV cache memory fragmentation and irregular request scheduling. It uses paged attention to virtualize KV cache into fixed-size blocks, continuous batching for token-level scheduling, and hardware-aware execution on TPUs. Rob Mulla demoed serving Gemma 4 (31B parameters) on eight Trillium chips. The same API works on GPUs. A PR from UCSD added diffusion-style speculative decoding, predicting all tokens at once instead of one at a time, yielding 3x speedup on several tasks.
Post-Training with Tunix and GRPO
Tunix is Google’s lightweight post-training and reinforcement learning framework for TPUs. It supports supervised fine-tuning, RL methods like GRPO, and knowledge distillation. In a Kaggle hackathon with thousands of participants, Tunix taught an older Gemma model to reason through math problems. Rob Mulla showed GRPO on a 4B model: actor and reference models, a dataset of 100K food images from Hugging Face, and custom reward functions for tool calling and output formatting. The fine-tuned model ran on a single Trillium chip, matching the 31B model’s food-logging behavior.
Pre-Training with MaxText
MaxText is an open-source reference implementation for large-scale pre-training on TPUs, built on JAX and XLA. It provides battle-tested configurations for model families like Gemma, Mistral, DeepSeek, and Qwen. Developers can start from a recipe instead of from scratch. MaxText scales from single-host experiments to multi-pod production runs without code changes. Features include multi-token prediction, custom kernel support, and multi-tier checkpointing.
Frameworks: JAX, Torch TPU, and Kinetic
JAX is a minimal framework that lets developers write Python functions and add transformations like grad, jit, and vmap. Its ecosystem includes Flax (neural network library) and Orbax (checkpointing). Gemini is built on JAX. Torch TPU, developed with Meta, offers a native PyTorch experience on TPUs – no new APIs. Running existing PyTorch code often requires minimal changes. Kinetic, from the Keras team, lets users deploy on TPUs with one decorator, handling cluster setup automatically.
Notable Quotes
vLLM uses paged attention. It virtualizes KV cache into fixed-size blocks. Rob Mulla · ▶ Watch (4:47)
diffusion-style speculative decoding … we got about a 3x speed up on a bunch of tasks. Josh Gordon · ▶ Watch (16:03)
Gemini is built on top of JAX Josh Gordon · ▶ Watch (32:21)
if you know PyTorch, you already know TPUs. Girija Sathyamurthy · ▶ Watch (35:00)
Kinetic … makes running your code on TPUs a lot easier just using a decorator. Josh Gordon · ▶ Watch (36:18)
Key Takeaways
- vLLM on TPUs provides a unified inference backend for GPUs and TPUs with paged attention and continuous batching.
- Tunix enables GRPO-based reinforcement learning on TPUs, using reward functions instead of labeled data.
- MaxText offers reproducible, battle-tested pre-training recipes that scale from one VM to thousands of chips.
- Torch TPU runs PyTorch code on TPUs with no new APIs, often with minimal code changes.