Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
120 changes: 26 additions & 94 deletions docs/getting_started/installation.md
Original file line number Diff line number Diff line change
Expand Up @@ -4,100 +4,32 @@ This guide provides instructions for installing and running `tpu-inference`.

There are three ways to install `tpu-inference`:

1. **[Install using pip via uv](#install-using-pip-via-uv)**
2. **[Run with Docker](#run-with-docker)**
3. **[Install from source](#install-from-source)**

## Install using pip via uv

We recommend using [uv](https://docs.astral.sh/uv/) (`uv pip install`) instead of standard `pip` as it improves installation speed.

1. Create a working directory:

```shell
mkdir ~/work-dir
cd ~/work-dir
```

1. Install `uv` and set up a Python virtual environment:

```shell
# If you prefer standard pip, simply use `python3.12 -m venv vllm_env`
curl -LsSf https://astral.sh/uv/install.sh | sh
source $HOME/.local/bin/env
uv venv vllm_env --python 3.12
source vllm_env/bin/activate
```

1. Use the following command to install vllm-tpu using `uv` or `pip`:

```shell
uv pip install vllm-tpu
# Or instead: pip install vllm-tpu
```

## Run with Docker

Include the `--privileged`, `--net=host`, and `--shm-size=150gb` options to enable TPU interaction and shared memory.

```shell
export DOCKER_URI=vllm/vllm-tpu:latest
sudo docker run -it --rm --name $USER-vllm --privileged --net=host \
-v /dev/shm:/dev/shm \
--shm-size 150gb \
-p 8000:8000 \
--entrypoint /bin/bash ${DOCKER_URI}
```

## Install from source

For debugging or development purposes, you can install `tpu-inference` from source. `tpu-inference` is a plugin for `vllm`, so you need to install both from source.

1. Install system dependencies:

```shell
sudo apt-get update && sudo apt-get install -y libopenblas-base libopenmpi-dev libomp-dev
```

1. Clone the `vllm` and `tpu-inference` repositories:

```shell
git clone https://github.com/vllm-project/tpu-inference.git
export VLLM_COMMIT_HASH="$(cat tpu-inference/.buildkite/vllm_lkg.version)"
git clone https://github.com/vllm-project/vllm.git
cd vllm
git checkout "${VLLM_COMMIT_HASH}"
cd ..
```

1. Install `uv` and set up a Python virtual environment:

```shell
curl -LsSf https://astral.sh/uv/install.sh | sh
source $HOME/.local/bin/env
uv venv vllm_env --python 3.12
source vllm_env/bin/activate
```

1. Install `vllm` from source, targeting the TPU device:

NOTE: `tpu-inference` repo pins `vllm` revision in `vllm_lkg.version` file,
make sure to checkout proper revision beforehand.

```shell
cd vllm
uv pip install -r requirements/tpu.txt --torch-backend=cpu
VLLM_TARGET_DEVICE="tpu" uv pip install -e . --no-build-isolation
cd ..
```

1. Install `tpu-inference` from source:

```shell
cd tpu-inference
uv pip install -e .
cd ..
```
- **uv pip** (Recommended): Fast installation of official stable releases into a Python virtual environment. Best for standard inference and serving.
- **Docker**: Pre-built containers with shared memory configured. Best for reproducible environments and Kubernetes (GKE).
- **Source**: Build directly from the repository. Best for debugging, contributors, and custom kernel development.

## Installation Commands

Select your preferred installation method to generate the exact setup command.

<div class="command-generator-container">
<div class="cg-options-group">
<span class="cg-label">Method</span>
<button class="cg-btn active" role="button" aria-pressed="true" data-group="method" data-val="uv_pip">uv pip</button>
<button class="cg-btn" role="button" aria-pressed="false" data-group="method" data-val="docker">Docker</button>
<button class="cg-btn" role="button" aria-pressed="false" data-group="method" data-val="source">Source</button>
</div>
<div class="cg-options-group" id="docker-image-group" style="display: none;">
<span class="cg-label">Image</span>
<button class="cg-btn active" role="button" aria-pressed="true" data-group="docker_img" data-val="latest">Release (latest)</button>
<button class="cg-btn" role="button" aria-pressed="false" data-group="docker_img" data-val="nightly">Nightly</button>
</div>

<div id="cg-output-instructions" class="cg-instructions"></div>
<div class="cg-output-container">
<pre><code id="cg-output-command" class="language-shell"></code></pre>
</div>
</div>

## Verify Installation

Expand Down
202 changes: 142 additions & 60 deletions docs/getting_started/tpu_setup.md
Original file line number Diff line number Diff line change
Expand Up @@ -2,33 +2,71 @@

This guide provides information on setting up and provisioning Google Cloud TPUs for use with `tpu-inference`.

## TPU Versions and Topologies
## TPU Generations

Tensor Processing Units (TPUs) are Google's custom-developed application-specific
integrated circuits (ASICs) used to accelerate machine learning workloads. TPUs
are available in different versions each with different hardware specifications.
For more information about TPUs, see [TPU System Architecture](https://cloud.google.com/tpu/docs/system-architecture-tpu-vm).
Tensor Processing Units (TPUs) are Google's custom-developed application-specific integrated circuits (ASICs) used to accelerate machine learning workloads. TPUs are available in different versions each with different hardware specifications. For more information about TPUs, see [TPU System Architecture](https://cloud.google.com/tpu/docs/system-architecture-tpu-vm).

The following TPU versions are compatible with `tpu-inference`:
The following TPU versions are compatible with `tpu-inference`. Selecting a topology allows you to configure the physical arrangements of the TPU chips, improving throughput and networking performance.

### Recommended
- [TPU v7x](https://cloud.google.com/tpu/docs/tpu7x)
- [TPU v6e](https://cloud.google.com/tpu/docs/v6e)
- [TPU v5e](https://cloud.google.com/tpu/docs/v5e)
**Recommended**

### Experimental
- [TPU v5p](https://cloud.google.com/tpu/docs/v5p)
- [TPU v4](https://cloud.google.com/tpu/docs/v4)
- [TPU v3](https://cloud.google.com/tpu/docs/v3)
<div class="grid cards" markdown>

These TPU versions allow you to configure the physical arrangements of the TPU
chips. This can improve throughput and networking performance. For more
information see:
- __TPU v7x (Ironwood)__

- [TPU v6e topologies](https://cloud.google.com/tpu/docs/v6e#configurations)
- [TPU v5e topologies](https://cloud.google.com/tpu/docs/v5e#tpu-v5e-config)
- [TPU v5p topologies](https://cloud.google.com/tpu/docs/v5p#tpu-v5p-config)
- [TPU v4 topologies](https://cloud.google.com/tpu/docs/v4#tpu-v4-config)
<span class="cg-badge cg-badge-recommended">Recommended</span>

Next-generation architecture for bleeding-edge research and ultra-large model training.

[View Topology Guidelines &rarr;](https://cloud.google.com/tpu/docs/tpu7x)

- __TPU v6e (Trillium)__

<span class="cg-badge cg-badge-recommended">Recommended</span>

Optimal performance for mainstream AI workloads with balanced memory and compute.

[View Topology Guidelines &rarr;](https://cloud.google.com/tpu/docs/v6e#configurations)

- __TPU v5e__

<span class="cg-badge cg-badge-recommended">Recommended</span>

Cost-effective performance for medium-to-large scale inference and training.

[View Topology Guidelines &rarr;](https://cloud.google.com/tpu/docs/v5e#tpu-v5e-config)

</div>

**Experimental**

<div class="grid cards" markdown>

- __TPU v5p__

<span class="cg-badge cg-badge-experimental">Experimental</span>

High-performance architecture optimized for peak compute and memory bandwidth.

[View Topology Guidelines &rarr;](https://cloud.google.com/tpu/docs/v5p#tpu-v5p-config)

- __TPU v4__

<span class="cg-badge cg-badge-experimental">Experimental</span>

Previous generation flagship architecture for large-scale training.

[View Topology Guidelines &rarr;](https://cloud.google.com/tpu/docs/v4#tpu-v4-config)

- __TPU v3__

<span class="cg-badge cg-badge-experimental">Experimental</span>

Legacy architecture suitable for smaller scale experimentation.

[View Topology Guidelines &rarr;](https://cloud.google.com/tpu/docs/v3)

</div>

## Quota and Pricing

Expand All @@ -37,47 +75,91 @@ Google Cloud project. For more information, see [TPU quota](https://cloud.google

For TPU pricing information, see [Cloud TPU pricing](https://cloud.google.com/tpu/pricing).

You may need additional persistent storage for your TPU VMs. For more
information, see [Storage options for Cloud TPU data](https://cloud.devsite.corp.google.com/tpu/docs/storage-options).

## Provisioning Cloud TPUs

You can provision Cloud TPUs using the [Cloud TPU API](https://cloud.google.com/tpu/docs/reference/rest)
or the [queued resources](https://cloud.google.com/tpu/docs/queued-resources)
API (preferred). This section shows how to create TPUs using the queued resource API.

### Provision a Cloud TPU with the queued resource API

Use the following command to provision a Cloud TPU. Replace the parameters in all caps with your own values.

```bash
gcloud alpha compute tpus queued-resources create QUEUED_RESOURCE_ID \
--node-id TPU_NAME \
--project PROJECT_ID \
--zone ZONE \
--accelerator-type ACCELERATOR_TYPE \
--runtime-version RUNTIME_VERSION \
--service-account SERVICE_ACCOUNT
```

| Parameter name | Description |
|--------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
| QUEUED_RESOURCE_ID | The user-assigned ID of the queued resource request. |
| TPU_NAME | The user-assigned name of the TPU which is created when the queued resource request is allocated. |
| PROJECT_ID | Your Google Cloud project |
| ZONE | The Google Cloud zone where you want to create your Cloud TPU. The value you use depends on the version of TPUs you are using. For more information, see [TPU regions and zones] |
| ACCELERATOR_TYPE | Specify the TPU version, for example `v5litepod-4` specifies a v5e TPU with 4 cores, `v6e-1` specifies a v6e TPU with 1 core. For more information, see [TPU versions]. |
| RUNTIME_VERSION | The TPU VM runtime version to use. For example, use `v2-alpha-tpuv6e` for a VM loaded with one or more v6e TPU(s). For more information, see [TPU software versions](https://docs.cloud.google.com/tpu/docs/runtimes) |
| SERVICE_ACCOUNT | The email address for your service account. You can find it in the IAM Cloud Console under *Service Accounts*. For example: `tpu-service-account@<your_project_ID>.iam.gserviceaccount.com` |

Connect to your TPU VM using SSH:

```bash
gcloud compute tpus tpu-vm ssh TPU_NAME --project PROJECT_ID --zone ZONE
```

!!! note
When configuring `RUNTIME_VERSION` ("TPU software version") for your TPU, ensure it matches the TPU generation you've selected by referencing the [TPU VM images] compatibility matrix. Using an incompatible version may prevent vLLM from running correctly.
Google Cloud supports two primary APIs for provisioning TPUs:

- **[Compute Engine API](https://cloud.google.com/tpu/docs/tpus-in-compute-engine) (`gcloud compute instances create`)**: Recommended for modern TPU generations starting with **TPU v6e (Trillium)** and **TPU v5p**.
- **[Cloud TPU API](https://cloud.google.com/tpu/docs/queued-resources) (`gcloud alpha compute tpus queued-resources create`)**: Legacy API used for earlier generations like **TPU v5e**. Note that the Cloud TPU API is no longer under active development.

!!! note "TPU v7x (Ironwood)"
TPU v7x (Ironwood) is in preview status and is provisioned via **Google Kubernetes Engine (GKE)** rather than standalone on-demand VM commands. Note that TPU v7x does not support Flex-start (DWS). For v7x provisioning and cluster orchestration, see [About TPUs in GKE](https://cloud.google.com/tpu/docs/tpus-in-gke) and the official [TPU v7x documentation](https://cloud.google.com/tpu/docs/tpu7x).

You can also choose between two capacity models:

- **Standard (On-Demand)**: Immediate allocation at standard pay-as-you-go rates.
- **Flex-start (DWS)**: Discounted capacity via Dynamic Workload Scheduler (DWS) that runs uninterrupted for up to 7 days. Note that Flex-start is supported on **TPU v5e**, **TPU v5p**, and **TPU v6e** (**TPU v7x**, as well as older generations like **v3** and **v4**, do not support Flex-start). For more details, see [About Flex-start VMs](https://cloud.google.com/compute/docs/instances/about-flex-start-vms) and [DWS Pricing](https://cloud.google.com/products/dws/pricing#flex-start-tpu-vm-pricing).

Select your desired TPU hardware, number of chips, and capacity model to generate the exact provisioning command. Be sure to replace placeholder variables (like `PROJECT_ID` and `SERVICE_ACCOUNT`) with your own values before running.

<div class="command-generator-container" id="prov-generator">
<div class="cg-options-group">
<span class="cg-label">Hardware</span>
<button class="cg-btn active" role="button" aria-pressed="true" data-group="prov_hw" data-val="v6e">TPU v6e</button>
<button class="cg-btn" role="button" aria-pressed="false" data-group="prov_hw" data-val="v5e">TPU v5e</button>
<button class="cg-btn" role="button" aria-pressed="false" data-group="prov_hw" data-val="v5p">TPU v5p</button>
<button class="cg-btn" role="button" aria-pressed="false" data-group="prov_hw" data-val="v4">TPU v4</button>
<button class="cg-btn" role="button" aria-pressed="false" data-group="prov_hw" data-val="v3">TPU v3</button>
</div>
<div class="cg-options-group">
<span class="cg-label">Chips</span>
<button class="cg-btn" role="button" aria-pressed="false" data-group="prov_chips" data-val="1">1</button>
<button class="cg-btn active" role="button" aria-pressed="true" data-group="prov_chips" data-val="4">4</button>
<button class="cg-btn" role="button" aria-pressed="false" data-group="prov_chips" data-val="8">8</button>
<button class="cg-btn" role="button" aria-pressed="false" data-group="prov_chips" data-val="16">16</button>
<button class="cg-btn" role="button" aria-pressed="false" data-group="prov_chips" data-val="32">32</button>
<button class="cg-btn" role="button" aria-pressed="false" data-group="prov_chips" data-val="64">64</button>
</div>
<div class="cg-options-group">
<span class="cg-label">Model</span>
<button class="cg-btn active" role="button" aria-pressed="true" data-group="prov_model" data-val="standard">Standard</button>
<button class="cg-btn" role="button" aria-pressed="false" data-group="prov_model" data-val="flex_start">Flex-start (DWS)</button>
</div>

<div id="prov-output-instructions" class="cg-instructions"></div>
<div class="cg-output-container">
<pre><code id="prov-output-command" class="language-shell"></code></pre>
</div>
</div>

| Parameter | Description |
|-----------|-------------|
| `PROJECT_ID` | Your Google Cloud project ID. |
| `ZONE` | The Google Cloud zone where you have TPU quota (e.g., `us-east5-a`, `europe-west4-a`, `us-central2-b`). See [TPU regions and zones](https://cloud.google.com/tpu/docs/regions-zones) for availability. |
| `SERVICE_ACCOUNT` | The email address for your service account, found in the Cloud Console under IAM Service Accounts (e.g., `tpu-service-account@<your_project_ID>.iam.gserviceaccount.com`). Required for legacy Cloud TPU API calls. |
| `RUNTIME_VERSION` | Automatically populated by the generator above based on your selected TPU hardware generation. |

### Connecting and Checking Status

**Connect to your TPU VM using SSH**:

- **For GCE VM Instances (TPU v6e, v5p)**:

```bash
gcloud compute ssh my-tpu-vm --zone ZONE
```

- **For Legacy TPU VMs (TPU v5e, v4, v3)**:

```bash
gcloud compute tpus tpu-vm ssh my-tpu-name --project PROJECT_ID --zone ZONE
```

**Check Provisioning Status**:

To check whether your TPU VM or Flex-start request has been allocated and is running:

- **For GCE VM Instances (TPU v6e, v5p)**:

```bash
gcloud compute instances describe my-tpu-vm --zone ZONE
```

- **For Queued Resources (TPU v5e, v4, v3)**:

```bash
gcloud alpha compute tpus queued-resources describe my-queued-resource --zone ZONE
```

[TPU versions]: https://cloud.google.com/tpu/docs/runtimes
[TPU VM images]: https://cloud.google.com/tpu/docs/runtimes
Expand Down
Loading
Loading