A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[CI/Build] Add tool to build vllm-tpu wheel (#19165)
Signed-off-by: mgoin <michael@neuralmagic.com> Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Oct 12, 2025 at 18:25 UTC
7ef6052804819b270e5a69e54dc52ae9a016cf2d
2 files changed
+72
setup.py
+5
@@ -540,6 +540,11 @@ def get_gaudi_sw_version():
540
541
542
def get_vllm_version() -> str:
543
+ # Allow overriding the version. This is useful to build platform-specific
544
+ # wheels (e.g. CPU, TPU) without modifying the source.
545
+ if env_version := os.getenv("VLLM_VERSION_OVERRIDE"):
546
+ return env_version
547
+
548
version = get_version(write_to="vllm/_version.py")
549
sep = "+" if "+" not in version else "." # dev versions might contain +
550
tools/vllm-tpu/build.sh
new
+67
@@ -0,0 +1,67 @@
1
+#!/bin/bash
2
+set -e # Exit immediately if a command exits with a non-zero status.
3
+# Script to build VLLM wheel for TPU with an optional version override.
4
+
5
+SCRIPT_PATH_PARAM="$0"
6
+TOOLS_DIR=$(cd "$(dirname "$SCRIPT_PATH_PARAM")" && pwd) # Absolute path to the script's directory
7
+REPO_ROOT=$(cd "$TOOLS_DIR/../../" && pwd) # Absolute path to the repo root
8
+VLLM_DIR="$REPO_ROOT/" # Path to the vllm sources
9
+
10
+# Ensure we are not running from within the vllm directory if SCRIPT_PATH_PARAM is relative like "."
11
+if [ "$TOOLS_DIR" = "$VLLM_DIR" ]; then
12
+ echo "Error: This script should not be run from the vllm directory directly if using relative paths."
13
+ echo "Place it in a subdirectory like 'tools/vllm-tpu' and run it from the repository root or via its full path."
14
+ exit 1
15
+fi
16
+
17
+# Optional version argument
18
+if [ -n "$1" ]; then
19
+ USER_VERSION="$1"
20
+ export VLLM_VERSION_OVERRIDE="$USER_VERSION"
21
+ echo "User defined version: $USER_VERSION"
22
+else
23
+ echo "No version override supplied. Using default version from source."
24
+fi
25
+
26
+PYPROJECT_FILE="$VLLM_DIR/pyproject.toml"
27
+
28
+# Backup and update the project name.
29
+if ! grep -q "name = \"vllm-tpu\"" "$PYPROJECT_FILE"; then
30
+ echo "Patching pyproject.toml project name to vllm-tpu..."
31
+ cp "$PYPROJECT_FILE" "${PYPROJECT_FILE}.bak"
32
+ sed -i '0,/^name = "vllm"/s//name = "vllm-tpu"/' "$PYPROJECT_FILE"
33
+ PATCHED=true
34
+else
35
+ PATCHED=false
36
+fi
37
+
38
+# Navigate to the vllm directory
39
+cd "$VLLM_DIR"
40
+
41
+# Cleanup function to be called on exit or error
42
+cleanup() {
43
+ echo "Cleaning up..."
44
+ if [ "$PATCHED" = true ]; then
45
+ echo "Restoring original pyproject.toml..."
46
+ cp "${PYPROJECT_FILE}.bak" "$PYPROJECT_FILE"
47
+ rm -f "${PYPROJECT_FILE}.bak"
48
+ fi
49
+}
50
+trap cleanup EXIT HUP INT QUIT PIPE TERM # Register cleanup function to run on script exit and various signals
51
+
52
+echo "Updating pyproject.toml completed. Proceeding with build..."
53
+
54
+echo "Building wheel for TPU..."
55
+rm -rf dist/
56
+mkdir -p dist/
57
+
58
+# User confirmed to use 'python -m build' directly
59
+if ! VLLM_TARGET_DEVICE=tpu python -m build; then
60
+ echo "Error: Python build command failed. Check if 'python -m build' works and the 'build' module is installed."
61
+ exit 1
62
+fi
63
+
64
+trap - EXIT HUP INT QUIT PIPE TERM
65
+cleanup
66
+
67
+exit 0
\ No newline at end of file