TPU agent suite

This repository builds and distributes TPU agent suite packages : declarative bundles of multi-agent skills, routing protocols, and tools that compile and install natively across major agent harnesses—including Claude Code and Antigravity , with Codex and Cursor coming soon.

Users talk to the tpu-agent skill. It routes each request to a specialist: onboarding, profiling, a hardware check, cited TPU facts, and kernels. See the design docs for how the framework and the suite are built.

Python 3.12+ and the host-side packages the agents run with, XProf among them, in a virtualenv ( pip install -r package/requirements-client.txt ; see docs/install.md ). PyYAML alone builds and installs the package, but leaves the agents without XProf.

Standalone CLI tools (backend-dependent): gcloud ( Google Cloud CLI ) for Cloud TPU VMs, GCS staging, and GCP credentials ( gcloud auth login && gcloud auth application-default login ) kubectl ( Kubernetes CLI ) and gke-gcloud-auth-plugin ( gcloud components install gke-gcloud-auth-plugin ) for GKE Pod, Job and XPK workloads ( gcloud container clusters get-credentials --region= ) xpk ( Accelerated Processing Kit ) for multi-node GKE cluster workloads

gcloud ( Google Cloud CLI ) for Cloud TPU VMs, GCS staging, and GCP credentials ( gcloud auth login && gcloud auth application-default login )

kubectl ( Kubernetes CLI ) and gke-gcloud-auth-plugin ( gcloud components install gke-gcloud-auth-plugin ) for GKE Pod, Job and XPK workloads ( gcloud container clusters get-credentials --region= )

xpk ( Accelerated Processing Kit ) for multi-node GKE cluster workloads

The installer checks only the host-side Python packages. The first session sets up your backend and checks its tools, credentials and target (see docs/install.md ).

Compile the canonical agent package into native distributions for every supported harness:

This writes dist/ /tpu-agent-suite/ for claude-code and antigravity .

See docs/install.md ( headless / CI ).

Ensure your virtual environment is activated whenever you use this agentic suite. Once inside the session, the tpu-agent skill is activated by any relevant natural language request, so no explicit invocation is needed. For example:

Onboarding & Benchmarking : "Onboard Qwen/Qwen3.8-27B to run as a vLLM inference server on Cloud TPU using the torch_tpu backend. Send it a test completion request, then benchmark it and report TTFT, ITL and throughput in a table."

Onboarding & Benchmarking :

"Onboard Qwen/Qwen3.8-27B to run as a vLLM inference server on Cloud TPU using the torch_tpu backend. Send it a test completion request, then benchmark it and report TTFT, ITL and throughput in a table."

"Onboard Qwen/Qwen3.8-27B to run as a vLLM inference server on Cloud TPU using the torch_tpu backend. Send it a test completion request, then benchmark it and report TTFT, ITL and throughput in a table."

Profiling : "Analyze the XProf profile at /path/to/profile and identify memory bottlenecks and TPU compute utilization."

"Analyze the XProf profile at /path/to/profile and identify memory bottlenecks and TPU compute utilization."

"Analyze the XProf profile at /path/to/profile and identify memory bottlenecks and TPU compute utilization."

Kernel Optimization : "Optimize this custom Pallas flash-attention kernel for TPU v6e and run a self-refinement benchmark loop."

Kernel Optimization :

"Optimize this custom Pallas flash-attention kernel for TPU v6e and run a self-refinement benchmark loop."

"Optimize this custom Pallas flash-attention kernel for TPU v6e and run a self-refinement benchmark loop."

What each request routes to, where workloads run, and where outputs land are in docs/using.md .

Evaluations are orchestrated using Harbor and configured under eval/ :

Automated PR Evaluations : Evaluations are automatically triggered via GitHub Actions whenever a pull request is opened or updated that modifies relevant parts of the agent packages or evaluation datasets.

Automated PR Evaluations : Evaluations are automatically triggered via GitHub Actions whenever a pull request is opened or updated that modifies relevant parts of the agent packages or evaluation datasets.

For the full list of configs, credentials, and the interactive sandbox, see eval/README.md .

Contribution & Code Review Policy

All code contributions, tests, and agent packages live in this repository. Every pull request runs the following GitHub Actions workflows (defined in .github/workflows/ ):

make verify : manifest, delegation graph, schema, and knowledge base checks.

make test : framework and package test suites (translation parity per harness).

make build : distributions for all harnesses compile cleanly.

make lint : advisory content checks, reported as annotations.

pre-commit : formatting and static checks ( .pre-commit-config.yaml ).

Local setup for these checks, and the short path for common contributions, is in docs/developing.md . Note that external contributions are not being accepted at this time (see CONTRIBUTING.md ).

This project is eligible for the Google Open Source Software Vulnerability Rewards Program (OSS VRP) .

If you suspect a security vulnerability in this project, please report it via g.co/vulnz . Please do not report security vulnerabilities through public GitHub issues.