How to run LLM inference on Google TPUs with vLLM

By Samagra Sharma · Reviewed September 6, 2026 A CUDA serving image will not run unchanged on a Google TPU. To evaluate TPU inference, start with a compatible TPU VM, install the matched vllm-tpu package, verify that the runtime sees TPU devices, and test a small serving configuration before comparing its cost with a GPU endpoint. This guide uses a single-host Cloud TPU v6e setup and Llama 3.1 8B Instruct. It follows upstream documentation reviewed September 6, 2026; the commands have not been executed on a TPU for this guide. These are Google Cloud instructions. Tensorfuse’s documented AWS deployment configuration does not provide a TPU option.

1. Prepare a TPU VM and model access

Use the upstream TPU setup instructions to obtain a single-host v6e VM with the required quota, region capacity, permissions, and TPU-compatible system image. Start with one chip for this bounded example. Run the commands below inside that VM, where Python 3.12 and its virtual-environment support are installed. TPU provisioning can incur charges before the model server is ready. Record the resource name and region so you can release the evaluation resources afterwards. For the meaning of chip count and slice layout, see TPU vs GPU inference. Request access to Meta’s Llama 3.1 8B Instruct checkpoint and accept its terms with the Hugging Face account used by the VM. This model appears in the TPU backend’s recommended model list. Check both the model and required features; support for a model does not establish support for every serving option.

2. Install the matched TPU packages

python3.12 -m venv .venv-tpu
source .venv-tpu/bin/activate
python -m pip install --upgrade pip
python -m pip install "vllm-tpu==0.28.0"
python -m pip check
python -m pip freeze > tpu-environment.txt
hf auth login

The published vllm-tpu 0.28.0 package includes a Python 3.12 Linux x86-64 wheel and requires the matching tpu-inference 0.28.0 backend. Use a clean environment so an existing CUDA installation does not interfere. Enter the model-access token when prompted; do not place it in the model command or a committed file.

3. Confirm TPU discovery before loading weights

import importlib.metadata
import jax
from vllm.platforms import current_platform

print("vllm-tpu:", importlib.metadata.version("vllm-tpu"))
print("tpu-inference:", importlib.metadata.version("tpu-inference"))
print("platform:", current_platform.get_device_name())
devices = jax.devices()
print("devices:", devices)
assert devices and all(device.platform == "tpu" for device in devices)

Run this in the activated environment. Expect TPU devices and a TPU platform name. Stop here if JAX reports CPU devices or cannot initialize the backend; downloading model weights will not fix device discovery. The installation guide describes the supported installation paths and runtime verification.

4. Start the API server

export MODEL_ID=meta-llama/Llama-3.1-8B-Instruct
export MODEL_REVISION="$(python -c 'from huggingface_hub import HfApi; print(HfApi().model_info("meta-llama/Llama-3.1-8B-Instruct").sha)')"
mkdir -p "$HOME/tpu-model-cache"

vllm serve "$MODEL_ID" \
  --revision "$MODEL_REVISION" \
  --served-model-name tpu-llama \
  --download-dir "$HOME/tpu-model-cache" \
  --tensor-parallel-size 1 \
  --max-model-len 2048 \
  --max-num-seqs 1 \
  --host 127.0.0.1 \
  --port 8000

Record the resolved model revision with the environment file. The context and concurrency limits make the first test easier to diagnose; they are not a recommended production capacity. Expect model download and initialization before the endpoint becomes ready. The TPU quickstart uses the same vLLM serving interface with its TPU backend.

5. Test from a second shell on the VM

curl --fail-with-body http://127.0.0.1:8000/health
curl --fail-with-body http://127.0.0.1:8000/v1/models
curl --fail-with-body http://127.0.0.1:8000/v1/chat/completions \
  -H 'Content-Type: application/json' \
  -d '{
    "model": "tpu-llama",
    "messages": [{"role": "user", "content": "Name two reasons an inference request can be slow."}],
    "max_tokens": 128,
    "temperature": 0
  }'

The model list should include tpu-llama, and the chat response should contain a generated assistant message. The example listens only on the VM’s loopback interface. Use an SSH tunnel for a local evaluation; use an authenticated gateway and controlled network access before serving remote application traffic.

6. Compare the service, including startup

Run the inference benchmark at several request rates after initialization. Keep the model revision, task-quality checks, input lengths, output limits, and latency targets equal to the GPU baseline. Record compilation/startup separately from warm requests and test the prompt-length range the application actually uses. Include TPU and host allocation, idle time, storage, and cross-cloud traffic in the cost window. A TPU endpoint serving an AWS application introduces a network path that a benchmark client on the TPU VM does not measure. Use the cost-reduction workflow to compare accepted requests per dollar.

Symptom	First check
JAX only sees CPUs	Verify the TPU VM, compatible system image, environment and device initialization
Model download returns 401/403	Check checkpoint access and the account used by `hf auth login`
Memory or compilation failure	Start with the stated one-request context limit; inspect the first backend error before changing parallelism
Fast local test, slow application requests	Measure from the application’s region and include network/queueing time

After collecting results, stop the server and release the evaluation TPU resources through the provisioning method you used. Stopping the Python process alone does not release billed infrastructure.

Get Started

LLM Inference

Inference How-to Guides

Inference Hardware

Concepts

Operations

Troubleshooting

Enterprise Setup

Architecture

How to run LLM inference on Google TPUs with vLLM

1. Prepare a TPU VM and model access

2. Install the matched TPU packages

3. Confirm TPU discovery before loading weights

4. Start the API server

5. Test from a second shell on the VM

6. Compare the service, including startup

​1. Prepare a TPU VM and model access

​2. Install the matched TPU packages

​3. Confirm TPU discovery before loading weights

​4. Start the API server

​5. Test from a second shell on the VM

​6. Compare the service, including startup

1. Prepare a TPU VM and model access

2. Install the matched TPU packages

3. Confirm TPU discovery before loading weights

4. Start the API server

5. Test from a second shell on the VM

6. Compare the service, including startup