> ## Documentation Index
> Fetch the complete documentation index at: https://tensorfuse.io/docs/llms.txt
> Use this file to discover all available pages before exploring further.

# How to run LLM inference on Google TPUs with vLLM

> Install vLLM's TPU backend, verify JAX devices, serve Llama on a Cloud TPU VM, test the API, and compare compilation, latency, and cost with GPUs.

By [Samagra Sharma](/docs/authors/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](https://docs.vllm.ai/projects/tpu/en/stable/getting_started/tpu_setup/) 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](/docs/guides/inference/hardware/google-tpu).

Request access to [Meta's Llama 3.1 8B Instruct checkpoint](https://huggingface.co/meta-llama/Llama-3.1-8B-Instruct) and accept its terms with the Hugging Face account used by the VM. This model appears in the TPU backend's [recommended model list](https://docs.vllm.ai/projects/tpu/en/stable/recommended_models_features/). 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

```bash theme={null}
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](https://pypi.org/project/vllm-tpu/0.28.0/) 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

```python theme={null}
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](https://docs.vllm.ai/projects/tpu/en/stable/getting_started/installation/) describes the supported installation paths and runtime verification.

## 4. Start the API server

```bash theme={null}
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](https://docs.vllm.ai/projects/tpu/en/stable/getting_started/quickstart/) uses the same vLLM serving interface with its TPU backend.

## 5. Test from a second shell on the VM

```bash theme={null}
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](/docs/guides/inference/benchmarking) 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](/docs/guides/inference/how-to/reduce-cost) 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.