Use https://download.pytorch.org/whl/test/cu128

Signed-off-by: Huy Do <[email protected]>

Author: huydhn

.pre-commit-config.yaml

@@ -46,7 +46,7 @@ repos:
   rev: 0.6.2
   hooks:
     - id: pip-compile
-      args: [requirements/test.in, -o, requirements/test.txt]
+      args: [requirements/test.in, -o, requirements/test.txt, --index-strategy, unsafe-best-match]
       files: ^requirements/test\.(in|txt)$
 - repo: local
   hooks:

docker/Dockerfile

@@ -66,7 +66,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 COPY requirements/common.txt requirements/common.txt
 COPY requirements/cuda.txt requirements/cuda.txt
 RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system -r requirements/cuda.txt
+    uv pip install --system -r requirements/cuda.txt --index-strategy unsafe-best-match
 
 # cuda arch list used by torch
 # can be useful for both `dev` and `test`
@@ -96,7 +96,7 @@ COPY requirements/build.txt requirements/build.txt
 ENV UV_HTTP_TIMEOUT=500
 
 RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system -r requirements/build.txt
+    uv pip install --system -r requirements/build.txt --index-strategy unsafe-best-match
 
 COPY . .
 ARG GIT_REPO_CHECK=0
@@ -171,7 +171,7 @@ COPY requirements/dev.txt requirements/dev.txt
 RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install --system mamba-ssm==2.2.4 --no-build-isolation
 RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system -r requirements/dev.txt
+    uv pip install --system -r requirements/dev.txt --index-strategy unsafe-best-match
 #################### DEV IMAGE ####################
 
 #################### vLLM installation IMAGE ####################
@@ -228,7 +228,9 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 # Install vllm wheel first, so that torch etc will be installed.
 RUN --mount=type=bind,from=build,src=/workspace/dist,target=/vllm-workspace/dist \
     --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system dist/*.whl --verbose
+    uv pip install --system dist/*.whl --verbose \
+    --extra-index-url https://download.pytorch.org/whl/cu128 \
+    --index-strategy unsafe-best-match
 
 # TESTING: install xformers from source to test 2.7.0 final RC
 RUN --mount=type=cache,target=/root/.cache/uv \
@@ -268,7 +270,7 @@ uv pip list
 # TODO: Remove this once FlashInfer AOT wheel is fixed
 COPY requirements/build.txt requirements/build.txt
 RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system -r requirements/build.txt
+    uv pip install --system -r requirements/build.txt --index-strategy unsafe-best-match
 
 #################### vLLM installation IMAGE ####################
 
@@ -288,7 +290,7 @@ ENV UV_HTTP_TIMEOUT=500
 RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install --system mamba-ssm==2.2.4 --no-build-isolation
 RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system -r requirements/dev.txt
+    uv pip install --system -r requirements/dev.txt --index-strategy unsafe-best-match
 
 # install development dependencies (for testing)
 RUN --mount=type=cache,target=/root/.cache/uv \

requirements/build.txt

@@ -4,6 +4,7 @@ ninja
 packaging
 setuptools>=61
 setuptools-scm>=8
+--extra-index-url https://download.pytorch.org/whl/cu128
 torch==2.7.0
 wheel
 jinja2>=3.1.6

requirements/cpu.txt

@@ -2,6 +2,7 @@
 -r common.txt
 
 # Dependencies for CPUs
+--extra-index-url https://download.pytorch.org/whl/cpu
 torch==2.7.0+cpu; platform_machine == "x86_64"
 torch==2.7.0; platform_system == "Darwin"
 torch==2.7.0; platform_machine == "ppc64le" or platform_machine == "aarch64"

requirements/cuda.txt

@@ -6,6 +6,7 @@ numba == 0.61.2; python_version > '3.9'
 
 # Dependencies for NVIDIA GPUs
 ray[cgraph]>=2.43.0, !=2.44.* # Ray Compiled Graph, required for pipeline parallelism in V1.
+--extra-index-url https://download.pytorch.org/whl/cu128
 torch==2.7.0
 torchaudio==2.7.0
 # These must be updated alongside torch

requirements/rocm-build.txt

@@ -1,6 +1,7 @@
 # Common dependencies
 -r common.txt
 
+--extra-index-url https://download.pytorch.org/whl/rocm6.2.4
 torch==2.7.0
 torchvision==0.22.0
 torchaudio==2.7.0

requirements/test.in

@@ -23,6 +23,7 @@ sentence-transformers # required for embedding tests
 soundfile # required for audio tests
 jiwer # required for audio tests
 timm # required for internvl test
+--extra-index-url https://download.pytorch.org/whl/cu128
 torch==2.7.0
 torchaudio==2.7.0
 torchvision==0.22.0

requirements/test.txt

@@ -1,5 +1,5 @@
 # This file was autogenerated by uv via the following command:
-#    uv pip compile requirements/test.in -o requirements/test.txt
+#    uv pip compile requirements/test.in -o requirements/test.txt --index-strategy unsafe-best-match
 absl-py==2.1.0
     # via rouge-score
 accelerate==1.0.1
@@ -718,12 +718,12 @@ torch==2.7.0
     #   torchvision
     #   vector-quantize-pytorch
     #   vocos
-torchaudio==2.7.0
+torchaudio==2.7.0+cu128
     # via
     #   -r requirements/test.in
     #   encodec
     #   vocos
-torchvision==0.22.0
+torchvision==0.22.0+cu128
     # via
     #   -r requirements/test.in
     #   timm

tests/standalone_tests/python_only_compile.sh

@@ -18,7 +18,8 @@ apt autoremove -y
 
 echo 'import os; os.system("touch /tmp/changed.file")' >> vllm/__init__.py
 
-VLLM_TEST_USE_PRECOMPILED_NIGHTLY_WHEEL=1 VLLM_USE_PRECOMPILED=1 pip3 install -vvv -e .
+VLLM_TEST_USE_PRECOMPILED_NIGHTLY_WHEEL=1 VLLM_USE_PRECOMPILED=1 pip3 install -vvv -e . \
+  --extra-index-url https://download.pytorch.org/whl/cu128
 
 # Run the script
 python3 -c 'import vllm'