vllm-torchtpu-core 0.29.0

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. 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

# 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 -e . --group vllm

Option B: Using pip

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

# Install vLLM-TorchTPU and dependencies (--group vllm pulls vLLM from source;
# needs pip >= 25.1)
VLLM_TARGET_DEVICE="empty" pip install \
  --extra-index-url https://download.pytorch.org/whl/cpu \
  -e . --group vllm

2. [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