Metadata-Version: 2.4
Name: vllm-torchtpu-core
Version: 0.29.0
Author: vLLM TorchTPU Contributors
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Education
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: Apache Software License
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: <3.13,>=3.12
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: tpu-info
Requires-Dist: portpicker
Requires-Dist: pathwaysutils
Requires-Dist: torch==2.13.0
Requires-Dist: torchvision==0.28.0
Requires-Dist: torch-tpu==0.2.0
Requires-Dist: tpu-sync-torch==0.0.1.dev20261008082921
Requires-Dist: transformers<5.15,>=5.10.4
Requires-Dist: numba==0.65.0
Requires-Dist: jax==0.11.2
Requires-Dist: jaxlib==0.11.2
Requires-Dist: libtpu==0.0.50
Requires-Dist: ray[default]
Requires-Dist: ray[data]
Provides-Extra: benchmarking
Requires-Dist: nltk; extra == "benchmarking"
Requires-Dist: evaluate; extra == "benchmarking"
Requires-Dist: datasets; extra == "benchmarking"
Requires-Dist: rouge-score; extra == "benchmarking"
Requires-Dist: scikit-learn; extra == "benchmarking"
Requires-Dist: pandas; extra == "benchmarking"
Requires-Dist: lm-eval[api]==0.5.0.dev1; extra == "benchmarking"
Requires-Dist: evalscope==1.10.0; extra == "benchmarking"
Requires-Dist: runai-model-streamer[gcs]<0.17; extra == "benchmarking"
Provides-Extra: omni
Requires-Dist: vllm-omni==0.29.0rc1; extra == "omni"
Dynamic: license-file

<p align="center">
   <!-- This image will ONLY show up in GitHub's dark mode -->
  <img src="docs/assets/torch_tpu_dark_mode_short.png#gh-dark-mode-only" alt="vLLM TorchTPU" style="width: 86%;">
    <!-- This image will ONLY show up in GitHub's light mode (and on other platforms) -->
  <img src="docs/assets/torch_tpu_light_mode_short.png#gh-light-mode-only" alt="vLLM TorchTPU" style="width: 86%;">
</p>

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](https://github.com/vllm-project/vllm-torchtpu/blob/main/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.

```sh
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:
>
> ```bash
> 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)

```bash
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`

```bash
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)

```bash
# 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`

```bash
# 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.

```bash
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:

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

> [!NOTE]
> For current TorchTPU builds, set the temporary workarounds below or append `PYTHONPATH=$(pwd)/src` to the `python3` command:
>
> ```bash
> export TORCH_TPU_INTERNAL_MATERIALIZE_COLLECTIVE_TENSORS=false
> export TORCHINDUCTOR_AUTOGRAD_CACHE=0
> ```

#### Tensor Parallelism (TP)

```bash
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)

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

```bash
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
```
