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 12, 18:30 UTC
Revert #28875 (#29159)
Cyrus Leung committed
Nov 21, 2025 at 17:40 UTC
4d7231e7743e80078bbc68ccc37b5ba5a1f28bf5
3 files changed
+25
-4
docker/Dockerfile
+17
@@ -56,6 +56,7 @@ ARG UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
56
57
# PyTorch provides its own indexes for standard and nightly builds
58
ARG PYTORCH_CUDA_INDEX_BASE_URL=https://download.pytorch.org/whl
59
+ARG PYTORCH_CUDA_NIGHTLY_INDEX_BASE_URL=https://download.pytorch.org/whl/nightly
60
61
# PIP supports multiple authentication schemes, including keyring
62
# By parameterizing the PIP_KEYRING_PROVIDER variable and setting it to
@@ -97,6 +98,7 @@ RUN echo 'tzdata tzdata/Areas select America' | debconf-set-selections \
98
ARG PIP_INDEX_URL UV_INDEX_URL
99
ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
100
ARG PYTORCH_CUDA_INDEX_BASE_URL
101
+ARG PYTORCH_CUDA_NIGHTLY_INDEX_BASE_URL
102
ARG PIP_KEYRING_PROVIDER UV_KEYRING_PROVIDER
103
104
# Activate virtual environment and add uv to PATH
@@ -315,6 +317,7 @@ RUN echo 'tzdata tzdata/Areas select America' | debconf-set-selections \
317
ARG PIP_INDEX_URL UV_INDEX_URL
318
ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
319
ARG PYTORCH_CUDA_INDEX_BASE_URL
320
+ARG PYTORCH_CUDA_NIGHTLY_INDEX_BASE_URL
321
ARG PIP_KEYRING_PROVIDER UV_KEYRING_PROVIDER
322
323
# Install uv for faster pip installs
@@ -334,6 +337,20 @@ ENV UV_LINK_MODE=copy
337
# or future versions of triton.
338
RUN ldconfig /usr/local/cuda-$(echo $CUDA_VERSION | cut -d. -f1,2)/compat/
339
340
+# arm64 (GH200) build follows the practice of "use existing pytorch" build,
341
+# we need to install torch and torchvision from the nightly builds first,
342
+# pytorch will not appear as a vLLM dependency in all of the following steps
343
+# after this step
344
+RUN --mount=type=cache,target=/root/.cache/uv \
345
+ if [ "$TARGETPLATFORM" = "linux/arm64" ]; then \
346
+ uv pip install --system \
347
+ --index-url ${PYTORCH_CUDA_NIGHTLY_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
348
+ "torch==2.8.0.dev20250318+cu128" "torchvision==0.22.0.dev20250319" ; \
349
+ uv pip install --system \
350
+ --index-url ${PYTORCH_CUDA_NIGHTLY_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
351
+ --pre pytorch_triton==3.3.0+gitab727c40 ; \
352
+ fi
353
+
354
# Install vllm wheel first, so that torch etc will be installed.
355
RUN --mount=type=bind,from=build,src=/workspace/dist,target=/vllm-workspace/dist \
356
--mount=type=cache,target=/root/.cache/uv \
docs/deployment/docker.md
+4
-3
@@ -82,7 +82,8 @@ DOCKER_BUILDKIT=1 docker build . \
82
83
## Building for Arm64/aarch64
84
85
-A docker container can be built for aarch64 systems such as the Nvidia Grace-Hopper. At time of this writing, this should be considered **experimental**. Using the flag `--platform "linux/arm64"` will attempt to build for arm64.
85
+A docker container can be built for aarch64 systems such as the Nvidia Grace-Hopper. At time of this writing, this requires the use
86
+of PyTorch Nightly and should be considered **experimental**. Using the flag `--platform "linux/arm64"` will attempt to build for arm64.
87
88
!!! note
89
Multiple modules must be compiled, so this process can take a while. Recommend using `--build-arg max_jobs=` & `--build-arg nvcc_threads=`
@@ -93,6 +94,7 @@ A docker container can be built for aarch64 systems such as the Nvidia Grace-Hop
94
95
```bash
96
# Example of building on Nvidia GH200 server. (Memory usage: ~15GB, Build time: ~1475s / ~25 min, Image size: 6.93GB)
97
+ python3 use_existing_torch.py
98
DOCKER_BUILDKIT=1 docker build . \
99
--file docker/Dockerfile \
100
--target vllm-openai \
@@ -100,8 +102,7 @@ A docker container can be built for aarch64 systems such as the Nvidia Grace-Hop
102
-t vllm/vllm-gh200-openai:latest \
103
--build-arg max_jobs=66 \
104
--build-arg nvcc_threads=2 \
103
- --build-arg torch_cuda_arch_list="9.0 10.0+PTX" \
104
- --build-arg RUN_WHEEL_CHECK=false
105
+ --build-arg torch_cuda_arch_list="9.0 10.0+PTX"
106
```
107
108
!!! note
docs/getting_started/installation/gpu.cuda.inc.md
+4
-1
@@ -158,7 +158,10 @@ uv pip install -e .
158
159
##### Use an existing PyTorch installation
160
161
-There are scenarios where the PyTorch dependency cannot be easily installed with `uv`, for example, when building vLLM with non-default PyTorch builds (like nightly or a custom build).
161
+There are scenarios where the PyTorch dependency cannot be easily installed with `uv`, e.g.:
162
+
163
+- Building vLLM with PyTorch nightly or a custom PyTorch build.
164
+- Building vLLM with aarch64 and CUDA (GH200), where the PyTorch wheels are not available on PyPI. Currently, only the PyTorch nightly has wheels for aarch64 with CUDA. You can run `uv pip install --index-url https://download.pytorch.org/whl/nightly/cu128 torch torchvision torchaudio` to [install PyTorch nightly](https://pytorch.org/get-started/locally/) and then build vLLM on top of it.
165
166
To build vLLM using an existing PyTorch installation:
167