---
title: "TPU"
description: "11 posts about TPU, summarised, each linking to the original."
---

# TPU
> 11 posts about TPU, summarised, each linking to the original.

## Articles

### [Reproducing OLMo 3 7B Pre-training in MaxText: case study of large scale training on TPUs](https://yomu.fyi/post/reproducing-olmo-3-7b-pre-training-in-maxtext-case-study-of-large-scal.md)
- Company: [Google](https://yomu.fyi/company/google.md)
- Author: Gagik Amirkhanyan, Ran Ran, Aireen Mei, Matt Davidow, TPU Inference Software Engineering Team, Google Cloud, AI2 Team
- Published: Sep 24, 2026
- Type: Problem & solution

The post describes the reproduction of AI2’s OLMo 3 7B pre-training using MaxText on Google Cloud TPUs, covering both stages, matching the reference loss curve and held-out metrics. PyTorch weights were converted to JAX with logit parity, KL divergence 1.5e-3 and 98.75% top-1 agreement, and held-out evaluations caught a data loader bug that made incorrect gains appear. The run demonstrated reliability via exact checkpoint-resume with zero loss delta, resilience to a 75% capacity loss via resize with 1% throughput preservation, and seamless TPU generation change. Performance optimizations achieved 44.5% MFU on Ironwood, and head reshaping from 32 heads by 128 to 16 by 256 yielded 12.4% faster throughput at identical parameters and FLOPs. The post emphasizes that faithful reproduction requires verification on generalization metrics and provides instructions for readers to reproduce the run with open sourced code.


### [Autonomous LLM post-training with Tunix on TPUs](https://yomu.fyi/post/autonomous-llm-post-training-with-tunix-on-tpus.md)
- Company: [Google](https://yomu.fyi/company/google.md)
- Author: Wei Wei
- Published: Sep 11, 2026
- Type: Benchmark

autofinetune applies autonomous research loops to LLM post-training, covering supervised fine-tuning (SFT) and reinforcement learning with GRPO through Tunix, Gemma, Cloud TPUs, Antigravity CLI, and Gemini Flash 3.7. A human defines the arena in program.md and supplies a self-contained run.py, while an agent edits hyperparameters, launches jobs, evaluates a target metric, keeps winning Git commits, reverts regressions, and logs results in results.tsv. On FunctionGemma 270M with google/mobile-actions, 20 Cloud TPU v5e-1 experiments tested LoRA settings, optimization choices, schedules, batch size, and seeds without changing the dataset, epochs, or architecture. A second study used Gemma 3 1B, GSM8K, and Cloud TPU v6e-1 for 40 GRPO experiments over two to three days, improving the combined numerical-and-format reward by about 10%.


### [Expanding our use of Google Cloud TPUs](https://yomu.fyi/post/expanding-our-use-of-google-cloud-tpus.md)
- Company: [Anthropic](https://yomu.fyi/company/anthropic.md)
- Published: Sep 10, 2026
- Type: Announcement

Anthropic announces plans to expand its use of Google Cloud technologies, including up to one million TPUs, in an expansion worth tens of billions of dollars. The company expects the deployment to bring well over one gigawatt of capacity online in 2026, increasing resources for customer demand, testing, alignment research, and responsible deployment at scale. Anthropic says it serves more than 300,000 business customers and that large accounts, defined as customers representing more than $100,000 in run-rate revenue, have grown nearly sevenfold in the past year. The company describes a diversified compute strategy using Google TPUs, Amazon Trainium, and NVIDIA GPUs, while remaining committed to Amazon as its primary training partner and cloud provider and continuing Project Rainier.


### [Google for Developers Blog - News about Web, Mobile, AI and Cloud](https://yomu.fyi/post/google-for-developers-blog-news-about-web-mobile-ai-and-cloud.md)
- Company: [Google](https://yomu.fyi/company/google.md)
- Published: Sep 2, 2026
- Type: Explainer

The page presents six featured Google developer articles spanning multi-agent systems, content credentials, TPU optimization, agent plugins, Go, and Genkit Go. One article identifies bidirectional MCP, async event buses, unified validation, and tiered routing as recurring patterns in successful AI Agents Challenge submissions. Another introduces Credentio, an open-source C++ library for local-first C2PA validation, while a HeyGen case study reports a 1.86x speedup for Avatar IV on Trillium TPUs using torchax, XLA, FSDP, Ulysses sequence parallelism, and custom kernel/compiler optimizations. The collection also covers Agent Plugins 1.0.0’s plugin.json directory specification and Genkit Go’s progressive-disclosure Agent Skills based on SKILL.md bundles. A final article argues that Go’s compiler, toolchain, readability, and backward compatibility provide guardrails for reviewing and maintaining AI-generated software.


### [Enterprise-Grade Precision for Long-Context Multimodal Embedding Inference on Cloud TPU](https://yomu.fyi/post/enterprise-grade-precision-for-long-context-multimodal-embedding-infer.md)
- Company: [Google](https://yomu.fyi/company/google.md)
- Author: Anthony Su, Injae Kwak
- Published: Aug 26, 2026
- Type: Benchmark

Embedding models convert text, images, and audio into dense vectors used for semantic search, recommendation, intent classification, content discovery, and clustering. The post describes production support for long-context Qwen3 Embedding models on Cloud TPUs through native vLLM TPU integration and GKE Custom Compute Classes, allowing prioritized autoscaling across TPU and GPU capacity pools. Its implementation pads vocabulary tensors for hardware-safe sharding, hardens lazy loading, pre-warms sharding-aware JAX/XLA compilation caches, and uses a hybrid StepPool with CachedRequestState to preserve pooling state during chunked prefill and preemption. Parity evaluations target cosine similarity of at least 0.999 for text and 0.995 for multimodal inputs; on TPU Ironwood, Qwen3-Embedding-8B reached 83,996 total tokens per second and 5.13 requests per second at 16K+ sequence length with bf16 and TP=4.


### [HeyGen x Google Cloud: Bringing Avatar IV to TPUs](https://yomu.fyi/post/heygen-x-google-cloud-bringing-avatar-iv-to-tpus.md)
- Company: [Google](https://yomu.fyi/company/google.md)
- Published: Aug 13, 2026
- Type: Benchmark

HeyGen and Google Cloud ported Avatar IV, an 18B-plus-parameter diffusion stack for streamed talking-head video, to an eight-chip Trillium (v6e) host. The pipeline uses two diffusion transformers and a VAE decoder, with FSDP weight sharding and Ulysses sequence parallelism across the mesh, and runs through torchax, JAX, XLA, and shape-specific Pallas kernels. Optimization targeted exposed all-to-all collectives, frame-misaligned sparse-attention blocks, serial softmax work, layout copies, scheduler choices, and custom-kernel cost estimates; changes were measured per video chunk against output-quality gates. The resulting pipeline was 1.86× faster than the first working TPU version, comparable in performance to the 8×H100 production setup and up to 25% more cost efficient per generated-video minute, while shipped changes preserved byte-identical output or passed defined similarity review.


### [How to use Google microbenchmarks for evaluating TPU performance](https://yomu.fyi/post/how-to-use-google-microbenchmarks-for-evaluating-tpu-performance.md)
- Company: [Google](https://yomu.fyi/company/google.md)
- Author: Junjie Qian, Chi Shuen Lee, Yu-Hsuan (Amy) Lin, Haixiong (Sean) Wang
- Published: Jul 30, 2026
- Type: Explainer

Evaluating TPU performance requires measurements beyond product specifications, because real-world behavior depends on architecture and workload constraints. The suite measures network collectives, compute, HBM bandwidth, host transfers, and Ragged-Paged Attention primitives, using metrics such as GB/s, latency, TFLOPs, MFU, and TTFT-related throughput. Its src and configs directories provide implementations and tunable parameters, while sample Kubernetes recipes connect measurements to TPU architecture and runtime behavior. Roofline analysis uses peak MXU throughput and HBM bandwidth to distinguish compute-, memory-, and network-bound workloads and guide kernel, sharding, or rematerialization choices. In a 110B MoE workload on a 4x4x4 TPU 7x configuration, microbenchmark-guided changes addressed ICI stalls and memory-bound attention, reducing training step time by 21.2%.


### [Run Ray on TPU, Part 2: Ray AI libraries](https://yomu.fyi/post/run-ray-on-tpu-part-2-ray-ai-libraries.md)
- Company: [Google](https://yomu.fyi/company/google.md)
- Author: Ivan Nardini, Spencer Peterson
- Published: Jul 24, 2026
- Type: Tutorial

Running Ray on TPUs requires multi-host workloads to remain on one intact TPU slice whose hosts share a high-speed ICI link. This guide presents Ray's AI libraries on top of GKE and Ray Core: Ray Serve for serving, Ray Data's iter\_jax\_batches() for input pipelines and batch inference, and Ray Train's JaxTrainer for distributed JAX training. Serve uses accelerator\_config.topology to defer slice placement to the replica, while iter\_jax\_batches() returns device-sharded JAX arrays and handles ragged final batches through drop, pad, or raise choices. JaxTrainer takes a training function and slice topology, launches one worker per host, and provides checkpointing, fault-tolerant restarts, and multi-slice scale-out; official TPU images and dashboard metrics round out the workflow.


### [Run Ray on TPU, Part 1: The foundations](https://yomu.fyi/post/run-ray-on-tpu-part-1-the-foundations.md)
- Company: [Google](https://yomu.fyi/company/google.md)
- Author: Ivan Nardini, Spencer Peterson
- Published: Jul 20, 2026
- Type: Explainer

Ray 2.55 makes Google Cloud TPUs a first-class accelerator in Ray, with official APIs, pre-built images, and core-library support rather than the earlier experimental path. The central constraint is TPU slicing: a multi-host workload must land on one intact slice because its hosts communicate through the slice’s ICI (Inter-Chip Interconnect). On GKE, the Ray Operator installs KubeRay and a TPU webhook; the webhook labels hosts, while Ray Core uses those labels to reserve a whole slice atomically through the public slice\_placement\_group() API, with topology and accelerator version as inputs. The guide shows node selection by TPU generation and topology, a multi-host numOfHosts field, and a RayCluster manifest that brings up one worker per host. Ray AI libraries such as Data, Train, and Serve generally invoke this placement primitive themselves, while direct use targets custom distributed workloads; the API is public but alpha.


### [Systems Engineering Playbook: Optimizing Qwen 3.5-397B MoE on Ironwood (TPU7x)](https://yomu.fyi/post/systems-engineering-playbook-optimizing-qwen-3-5-397b-moe-on-ironwood.md)
- Company: [Google](https://yomu.fyi/company/google.md)
- Author: Google for Developers
- Published: Jul 14, 2026
- Type: Benchmark

Serving Qwen3.5-397B, a 397-billion-parameter sparse Mixture-of-Experts model, on Ironwood (TPU v7x) requires fitting roughly 400 GB of weights into HBM while sustaining utilization across contrasting workloads. The report applies a modular, model-agnostic playbook built from reusable JAX/Pallas kernels and hardware-aware cost models, focusing new work on Gated DeltaNet linear attention and Attention Data Parallelism while integrating changes into vLLM and SGLang. The architecture activates exactly 17 billion parameters per token and combines 15 repeating blocks of three GDN layers and one GQA layer. Benchmarks use a single host with four Ironwood chips, eight logical devices, 8K/1K prefill-heavy and 1K/8K decode-heavy workloads, and concurrency tiers from 64 to 512. Between April and June 2026, stated optimizations improved inference performance by approximately 4.7x for 512-concurrency prefill-heavy workloads and 3.1x for decode-heavy workloads.


### [Unlocking the Power of the TPU Stack: Introducing our new Developer Hub](https://yomu.fyi/post/unlocking-the-power-of-the-tpu-stack-introducing-our-new-developer-hub.md)
- Company: [Google](https://yomu.fyi/company/google.md)
- Author: Keelin McDonell
- Published: Jun 16, 2026
- Type: Announcement

Google Cloud has launched the TPU Developer Hub, an educational resource for model builders, optimizers, and developers working with TPU infrastructure and its supporting software stack. It organizes guidance across the developer lifecycle, covering pre-training, post-training, and inference, with material on hardware architecture, infrastructure consumption modes, and software capabilities including specialized compiler technology and XLA. Resources also address PyTorch migration and deployment, XProf-based tracing and observability, multi-chip parallelism, Pallas kernels, KV cache offloading, networking, and security for distributed workloads. The hub combines interactive Colabs, open-source code recipes, and deep-dive documentation, with content intended for both manual browsing and AI-assisted development. Its stated purpose is to help users select appropriate TPU infrastructure, diagnose bottlenecks, optimize model performance, and move from concepts toward production through practical, regularly updated guidance.