TPU Agent Suite: Model Porting, Profiling and Evidence

TPU Agent Suite is a preview toolkit for assisted TPU development. It routes model onboarding, performance investigation and kernel tasks through specialist workflows, using an explicitly configured compute environment. Its public source is google-pytorch/tpu-agent-suite. It is useful for turning a model-porting request into code changes and measurements, provided the request defines what success means.

The Cloud TPU documentation labels the suite Preview. This review checked the repository on November 10, 2026. We inspected its source and documentation; we did not install it into a working project or execute a TPU port.

What the suite actually organizes

The repository separates a framework, a TPU workflow package and evaluation configurations. The framework handles installation, assistant integration and compute access. The package contains specialist workflows. Evaluations exercise selected tasks; they do not certify every model that a user might request.

Ask for an artifact that supports the decision

Comparison

Request	Workflow output to ask for
Port a model	Minimal code diff, environment and numerical comparison
Diagnose low performance	Profile artifact, suspected cause and a controlled follow-up
Optimize a kernel	Correctness cases, matched timings and retained failures
Check hardware facts	Exact configuration, units and primary source

The usage guide describes a single tpu-agent entry point that selects the appropriate specialist. This can shorten the path to an experiment. The person running it still needs to understand which result supports the deployment decision.

The host and the TPU target are different environments

The installation guide distinguishes local execution on a TPU machine, access to a Cloud TPU VM, execution inside an existing GKE pod, and dispatch through a GKE job. GKE is Google Kubernetes Engine. The selected backend determines where code runs and which credentials and tools are needed.

For example, a local assistant may inspect an XProf trace while the model executes on a remote TPU VM. Installing host-side profiling packages says nothing about whether the remote environment has a compatible TorchTPU build. Record both environments in the result. If the target is a Kubernetes cluster, identify its context and namespace so a later shell configuration change cannot silently move the experiment.

From a request to a reviewable experiment

Step 1 of 5: Define the task

Previous stepNext stepStart again

1. Define the task

   Name the model revision, target topology, precision, constraints and success criteria.

2. Verify the compute target

   Check the selected environment and matching dependencies before running model work.

3. Establish a baseline

   Save fixed inputs, correctness results, cold-start time and warm timing.

4. Make one bounded change

   Collect the new artifact and compare it against the same baseline conditions.

5. Keep the decision evidence

   Retain the diff, profiles, failures and the reason to accept or reject the change.

Write the porting request around a decision

A useful request might read:

Evaluate this pinned model on the existing TPU target. Keep its tokenizer, precision and generation settings fixed. First identify CUDA-specific dependencies, then run a numerical comparison on the supplied inputs. Allow at most three optimization attempts. Save each diff, the profile and all timing samples. Report cold startup separately from synchronized warm execution. Do not change the model or the evaluation inputs to improve the score.

This is an example task description, not an executed run. It gives the assistant a bounded problem and gives a reviewer enough information to reject a misleading success. Replace “this model” and “existing target” with actual identifiers before using it.

For inference serving, add offered request rate, prompt/output-length distributions and latency objectives. Require failures and cancellations in the report. For training, include gradients, several optimizer updates and checkpoint restoration. The TorchTPU migration guide explains why a correct forward pass alone is insufficient.

Budget the experiment, including unsuccessful runs

Suppose an illustrative experiment has 12 minutes of setup, five minutes of warmup and three six-minute trials. Its elapsed allocation budget is at least 35 minutes before analysis or cleanup. Eight reserved chips held for that interval consume 280 chip-minutes, or about 4.67 chip-hours. Multiply by the actual quoted rate only after identifying the billing unit.

If each trial independently recompiles for five more minutes, the total rises to 50 minutes and 6.67 chip-hours. That difference is why the run record should say whether warmup and compilation were reused. These are planning assumptions, not suite performance claims or TPU prices.

A useful optimization report retains the full timing distribution. Selecting the fastest trial from many attempts rewards noise. Repeat a promising change under matched conditions, alternate ordering when practical, and inspect whether host contention or compilation explains the result. The free serving benchmarks concept provides the measurement vocabulary.

Review an optimization as a change to one experiment

Suppose the suite proposes replacing an attention implementation to speed up a port. Before accepting the timing result, compare the changed code with the workload contract. Did the new path preserve the attention mask, cache layout and output length? Did it remove an unsupported case by excluding that case from the input set? A clean run on a smaller problem does not answer the original request.

Keep a small evidence directory for each attempt. These are suggested review artifacts, not filenames or outputs guaranteed by the suite:

Evidence that lets another engineer review the port

Comparison

Record	What it establishes
Input and model identities	Both runs received the same tokenized inputs, weights and generation limits.
Code diff and package versions	The intended implementation changed, without an unrelated dependency update.
Numerical results	Ordinary and boundary cases still meet the chosen quality criteria.
Raw timing samples and trace	The measured interval includes completed work and identifies compilation.
Restart result	A fresh process can load the intended checkpoint and reproduce the accepted behavior.

A useful rejected attempt is specific: “The replacement fails the masked-attention case at this shape; retain the baseline and attach the reproducer.” This is an example review outcome, not a reported suite failure. It is more actionable than a generic failed-port label, and it prevents the next attempt from rediscovering the same constraint.

Keep optimization inputs separate from a final evaluation set. Repeatedly tuning against one fixed batch can reward a shape-specific change that does little for production traffic. After selecting a candidate, evaluate it on the reserved inputs and the expected length distribution, including uncommon long requests. For a kernel, include tail dimensions that do not divide evenly into its tiles.

Resume a failed experiment from recorded evidence. If the connection drops after a job is submitted, inspect that job's state before launching another copy. An assistant timeout does not establish that remote compute stopped. Record the job identifier and cleanup outcome so a retry neither hides a continuing allocation nor contaminates the next timing run.

Preview support needs a version check

The pinned repository's introduction says Claude Code and Antigravity are available while Codex and Cursor are upcoming. Its build paragraph lists targets for all four. That inconsistency prevents a blanket statement that every integration is equally ready. Select a tagged distribution, read its installation notes and verify the actual assistant workflow before depending on it.

Similarly, a suite that can generate a kernel change does not prove the kernel supports every dtype or tail dimension. Keep boundary-shape tests and numerical tolerances next to timing results. An optimization that removes a required precision mode has changed the problem.

Keep the related products straight

TorchTPU executes PyTorch operations on the accelerator. TPU Agent Suite assists the developer. TPU Sync moves model state across memory and host boundaries. Agent Compute Engine is a separate service for agent execution environments; its suspend/resume lifecycle is not a model-porting workflow.

Use the TPU hardware comparison to choose a candidate configuration and the accelerator question category to practice explaining the choice. Accept a port when the evidence meets the workload contract. If the required operator or engine feature is unsupported, keep the working deployment and record the missing capability precisely enough to reevaluate a later release.