Metadata-Version: 2.4
Name: vllm-torchtpu-core-test
Version: 0.29.0rc1
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.1.1+release.2026.9.22.20261003154636
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.47
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]; 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`.

> [!TIP]
> A common cause of authentication errors when using `keyring` with `uv` is an existing `~/.netrc` file containing stale credentials.
>
> To resolve this, open `~/.netrc` in your preferred text editor and delete the block associated with `us-python.pkg.dev`. It will look something like this:
>
> ```text
> machine us-python.pkg.dev
> login oauth2accesstoken
> password <your_expired_token>
> ```

### 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:

```bash
gcloud auth list
```

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

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

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

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

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

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