vLLM, Tunix and MaxText: Google TPU setup and demos

This Google I/O 2026 session covers Google's open source TPU software stack for inference, post-training and pre-training. The examples use Google Cloud infrastructure rather than a desktop local LLM setup. The speakers also point to limited free Colab and Kaggle TPU access for smaller inference and post-training experiments; Google Cloud TPU infrastructure incurs hosting charges.

The vLLM demo serves Gemma 4's 31-billion-parameter model through a pre-built Docker container, with tensor parallelism set to eight. Rob Mulla says the model could fit on four chips; the demo uses eight to leave room for the KV cache. The explanation covers paged attention, prefix caching and continuous batching, followed by profiling and a food-logging AI agent that calls a simulated nutrition database.

Tunix then handles GRPO post-training of a four-billion-parameter model for the same task. The walkthrough defines actor and reference models, loads Food 101 through Grain, and sets rewards for tool use and output formatting. Mulla says training takes 30-45 minutes, so he presents metrics and results instead of completing the run live. He reports that the smaller trained model runs on one Trillium chip.

The final sections explain MaxText training recipes, JAX transformations and OpenXLA compilation. The speakers describe Torch TPU as work underway to bring PyTorch workloads to TPUs with minimal code changes, and introduce Kinetic for cluster setup through a decorator.