# TPU Agent Suite in Claude Code: Setup, Benchmarking and Common Traps

Google's preview adds TPU-specific workflows to coding assistants, but a successful install does not mean your compute target is ready.

By Paras Madan | 2026-11-10


Google's TPU agent suite gives Claude Code and Antigravity specialized workflows for getting models onto TPUs, diagnosing slow runs and improving kernels. It is an open-source preview, and the useful starting point is a controlled onboarding experiment with a verified compute target rather than an unrestricted request to make a model faster.

The distinction matters because installing the assistant package and preparing the accelerator are separate jobs. The installer can succeed while the remote machine still has the wrong Python environment, missing dependencies or no usable model cache.

## What is available in the preview

Google's [Cloud TPU documentation](https://docs.cloud.google.com/tpu/docs/pytorch-on-tpu) describes the suite as focused on transformer inference, including dense and mixture-of-experts models, on TPU v6e and TPU7x. Its router delegates onboarding, diagnostics, Pallas kernel work and documentation questions to specialists. It is not a replacement for the inference server or the TPU runtime.

There is also a packaging detail worth checking before installation. The [repository](https://github.com/google-pytorch/tpu-agent-suite) describes Claude Code and Antigravity support, with Codex and Cursor coming soon. Its build system can generate distributions for all four, but the installation guide says prebuilt release archives are published for Claude Code and Antigravity. **A generated distribution is not the same thing as a supported release.**

## Install the package, then verify the target

Start with the [installation guide](https://github.com/google-pytorch/tpu-agent-suite/blob/main/docs/install.md) and the release archive for your assistant. The host running that assistant needs Python 3.12 or newer and the archive's requirements-client.txt dependencies in an activated virtual environment. Installing only PyYAML is enough for the compiler and installer, but does not supply the complete host-side agent environment, including XProf.

For an extracted release, the documented project-scoped installation is:

```bash
python3 tpu-agent-suite-<harness>-v<version>/framework/installer.py install --destination /path/to/project
```

Replace the placeholders with the archive you actually downloaded and your project directory. Review the files the installer manages before choosing a user-wide installation. For Claude Code, the package installs skills, specialist definitions and managed configuration; it also adds a managed instruction block to the project's CLAUDE.md.

The first session then configures a compute profile. Choose the backend based on where the workload will execute:

| Backend | Use it when |
| --- | --- |
| local | The coding assistant is running on the TPU VM itself |
| gce_tpu_vm | The assistant reaches a Cloud TPU VM through gcloud |
| gke_pod | You want to execute inside an existing TPU pod |
| gke_job | You want to submit a new TPU job to a cluster |

A local profile does not give an ordinary laptop a TPU. For a remote VM, verify the project, zone, machine and Python interpreter. For GKE, verify the cluster context and namespace rather than relying on whichever context happens to be active in a terminal.

On the TPU VM, follow the guide's separate workload-environment instructions. It specifically tells users to install torch_tpu before manually installing PyTorch, because the TPU package supplies its compatible build. Keep the host environment and the workload environment distinct in your troubleshooting notes.

## Give the agent an experiment it can finish

The [usage guide](https://github.com/google-pytorch/tpu-agent-suite/blob/main/docs/using.md) recommends specifying the task, environment, resource limits and desired output. Here is an example prompt for a team with an existing model and test set. It is an experiment template, not a report of a test performed for this article:

```text
Use the verified TPU profile for this project. Onboard the model at
<existing model path> for inference using the pinned serving image.
Before changing kernels, check numerical outputs against our CPU reference
on <test set path> and report any differences against our stated tolerance.

Then benchmark the unchanged baseline using <input lengths>,
<output lengths> and <concurrency>. Separate download and compilation time
from warmed-up TTFT, inter-token latency and throughput.

Propose at most two optimization attempts. Ask before executing commands.
Keep the original configuration, record every change and report failures
alongside improvements. Do not leave a serving job running afterward.
```

The practical benefit is a comparison you can inspect. If accuracy changes or the request lengths change between runs, a faster number tells you very little about whether the optimization belongs in production.

## Why the first run is slow, or every run downloads again

The usage guide describes two separate startup waits: model download and compilation. Look for the weight-download progress before deciding that the compiler has stalled. Also leave enough disk space on the compute target, where the weights are stored.

Cache persistence depends on the backend. A VM retains its Hugging Face cache between runs. A pod retains its local cache only while that pod lives. A fresh gke_job starts without a cache unless you configure the shared model-cache claim, hf_cache_pvc. Repeated downloads can therefore be an infrastructure setup issue rather than a model or agent failure.

There is a second cost trap after benchmarking: serving jobs submitted with --serve are exempt from the suite's ordinary session-end cancellation. They keep their accelerators until explicitly stopped, subject to cluster-level bounds. Ending the coding session is not a reliable shutdown procedure for a service.

## Keep the evidence, not just the answer

The suite records artifacts and an experiment journal for each run. Review those records before accepting an optimization, particularly if the assistant changed a kernel, compiler setting or serving configuration. Trajectory export is optional, and the usage guide warns that exported transcripts and artifacts are not redacted, so inspect the destination and contents before sharing them.

A TPU-specific assistant can reduce the work needed to investigate a migration. The useful outcome is a reproducible model, environment and benchmark combination that another engineer can rerun, with the accelerator bill ending when the experiment does.

I’m Paras. I work at the intersection of AI, engineering and distribution.

I help technical teams turn infrastructure workflows and engineering decisions into clear developer stories.

[Book a 30-minute call](https://cal.com/paras-madan/30min).