vllm-torchtpu-core-test 0.29.0rc1

This repository contains the integration of TorchTPU and vLLM. It is a vLLM platform plugin packaged as vllm_torchtpu, with TPU kernels and runtime code for TorchTPU.

For contribution guidelines, developer environment setup, and architecture details, please refer to CONTRIBUTING.md.

---

🛠️ Installation

Follow these installation steps:

We recommend using uv for installing dependencies as it is significantly faster than standard pip.

1. Google Cloud Authentication

We need to authenticate with Google Cloud to access the private Torch TPU Virtual Registry (https://us-python.pkg.dev/ml-oss-artifacts-transient/torch-tpu-virtual-registry/simple/). This registry contains packages like torch-tpu, torch, and other dependencies required by vllm-torchtpu.

Make sure you are logged into gcloud using your corporate account (i.e., one that has read permissions for the Torch TPU registry). If you need access to this registry, please reach out to the Torch TPU team.

You can check your active account by running:

gcloud auth list

If the active account is not your corporate account, switch to it or log in by running:

gcloud auth login

# For keyring authentication, see https://github.com/GoogleCloudPlatform/artifact-registry-python-tools?tab=readme-ov-file#authentication
gcloud auth application-default login

2. Install vLLM-TorchTPU and dependencies

First, clone the repo, for one viewing this doc from web.

git clone git@github.com:vllm-project/vllm-torchtpu.git && cd vllm-torchtpu

VLLM_TARGET_DEVICE="empty" is a switch that defined and recognized by vLLM's build process. We use it since that vllm-torchtpu is an out-of-tree platform of the vllm product. We pass this variable during the installation of vllm-torchtpu, as we are building vllm wheel from source during this process.

Package names: this repo is distributed as vllm-torchtpu-core (the import package is still vllm_torchtpu). vLLM is not a core dependency; the vllm dependency group (--group vllm) installs it from source for development. vllm-torchtpu is the name of the combined wheel (vLLM + a dependency on vllm-torchtpu-core).

Migrating an existing environment: if this repo was installed before the rename (distribution name vllm_torchtpu), uninstall it first:

pip uninstall -y vllm_torchtpu   # or: uv pip uninstall vllm_torchtpu

pip normalizes vllm_torchtpu and vllm-torchtpu to the same project, so the old plugin would be "upgraded" by the combined vllm-torchtpu wheel. For non-editable installs, that removal can delete vllm_torchtpu/ files now owned by vllm-torchtpu-core. If this already happened, repair it with pip install --force-reinstall --no-deps vllm-torchtpu-core.

Option A: Using uv (Recommended)

uv venv --python 3.12 ~/uv_venv
source ~/uv_venv/bin/activate

# Set the username for the index
export UV_INDEX_TORCH_TPU_REGISTRY_USERNAME="oauth2accesstoken"

# Install keyring and Google Artifact Registry plugin for persistent auth
uv tool install keyring --with keyrings.google-artifactregistry-auth

# Use CPU Torch when building vLLM's editable TPU package.
export UV_TORCH_BACKEND=cpu

# Install vLLM-TorchTPU and dependencies (--group vllm pulls vLLM from source)
VLLM_TARGET_DEVICE="empty" uv pip install --pre -e . --group vllm

Option B: Using pip

python3.12 -m venv ~/pip_venv --symlinks
source ~/pip_venv/bin/activate
pip install --upgrade pip

# Install keyring and Google Artifact Registry plugin for persistent auth
pip install keyring keyrings.google-artifactregistry-auth

# Set the PIP_INDEX_URL environment variable (using keyring)
export PIP_INDEX_URL="https://oauth2accesstoken@us-python.pkg.dev/ml-oss-artifacts-transient/torch-tpu-virtual-registry/simple/"

# Install vLLM-TorchTPU and dependencies (--group vllm pulls vLLM from source;
# needs pip >= 25.1)
VLLM_TARGET_DEVICE="empty" pip install --pre -e . --group vllm

3. [Optional] Install vLLM

The local installation of vllm is optional. Follow these steps when one need to develop with vllm directly.

Note: Currently, vllm==0.29.0 is supported.

Option A: Using uv (Recommended)

# Use the venv from last section
source ~/uv_venv/bin/activate

# Clone vLLM to allow making local patches for debugging
git clone --depth 1 --branch v0.29.0 https://github.com/vllm-project/vllm.git ../vllm

# Install vLLM in editable mode
VLLM_TARGET_DEVICE="empty" uv pip install  --no-deps -e ../vllm

Option B: Using pip

# Use the venv from last section
source ~/pip_venv/bin/activate

# Clone vLLM to allow making local patches for debugging
git clone --depth 1 --branch v0.29.0 https://github.com/vllm-project/vllm.git ../vllm

# Install vLLM in editable mode
VLLM_TARGET_DEVICE="empty" pip install --no-deps -e ../vllm

---

🌐 Online Serving

Start the server with the following command:

Note: Prioritize compile mode for better performance. The first startup may take several minutes while TPU graphs compile. Add --enforce-eager if you want eager mode. On TPUv7, Qwen3-Coder-30B can fit on a single device. On v6, use a smaller model like Qwen3-4B or test with TP/EP.

vllm serve "Qwen/Qwen3-Coder-30B-A3B-Instruct" \
  --tensor_parallel_size=2 \
  --max-model-len=256 \
  --max-num-batched-tokens=256 \
  --attention-backend CUSTOM

Send a Request

Once the server is running, you can verify it by sending a request:

curl http://localhost:8000/v1/completions \
  -H "Content-Type: application/json" \
  -d '{
        "model": "Qwen/Qwen3-Coder-30B-A3B-Instruct",
        "prompt": "Hello, my name is",
        "max_tokens": 20,
        "temperature": 0
    }'

📉 Running Offline Inference

You can also run a simple offline inference script to verify the setup without starting a full server.

Tensor Parallelism (TP)

python3 examples/offline_inference.py \
  --model Qwen/Qwen3-Coder-30B-A3B-Instruct \
  --max-model-len 256 \
  --max-num-batched-tokens 256 \
  --tensor_parallel_size=2

Expert Parallelism (EP)

python3 examples/offline_inference.py \
  --model Qwen/Qwen3-Coder-30B-A3B-Instruct \
  --max-model-len 256 \
  --max-num-batched-tokens 256 \
  --tensor_parallel_size=2 \
  --enable-expert-parallel

Pipeline Parallelism (PP)

Each pipeline stage runs a slice of the layers on one core and hands its hidden states to the next stage with a collective-permute over ICI. Tensor parallelism inside a stage is not supported yet, so the worker count is pipeline_parallel_size.

A stage may instead span a prefill context-parallel group (--prefill-context-parallel-size), which splits each chunk's attention over the stage's cores; MoE models then need --enable-expert-parallel. This needs a torch_tpu that compiles a stage's kernels and collectives for the stage's devices only (the 2026-09-28 build of google-pytorch/torch_tpu#4387); on a torch_tpu that builds world-size meshes only, the platform refuses such stages. A two-core stage must not be the two cores of one chip: set VLLM_TPU_DEVICE_ORDER to the device each local rank binds, e.g. 0,2,1,3,4,6,5,7 for four stages of two cores on one host.

python3 examples/offline_inference.py \
  --model Qwen/Qwen3-Coder-30B-A3B-Instruct \
  --max-model-len 256 \
  --max-num-batched-tokens 256 \
  --pipeline-parallel-size=2 \
  --tensor_parallel_size=1 \
  --no-async-scheduling