cleanup cuda arm64 build

Signed-off-by: Daniele Trifirò <[email protected]>
vllm-project · Feb 19, 2025 · 1a2bd8d · 1a2bd8d
1 parent 6c4a7d0
commit 1a2bd8d
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 7 deletions.
diff --git a/Dockerfile b/Dockerfile
@@ -102,8 +102,7 @@ RUN --mount=type=cache,target=/root/.cache/pip \
         && export SCCACHE_IDLE_TIMEOUT=0 \
         && export CMAKE_BUILD_TYPE=Release \
         && sccache --show-stats \
-        && PIP_EXTRA_INDEX_URL="$([ "${TARGETPLATFORM}" = "linux/arm64" ] && echo "https://download.pytorch.org/whl/nightly/cu124")" \
-            python3 -m build --wheel --outdir=dist --no-isolation -v --config-setting py-limited-api=38  \
+        && python3 -m build --wheel --outdir=dist --no-isolation -v --config-setting py-limited-api=38  \
         && sccache --show-stats; \
     fi
 
@@ -112,8 +111,7 @@ RUN --mount=type=cache,target=/root/.cache/ccache \
     --mount=type=cache,target=/root/.cache/pip \
     --mount=type=bind,source=.git,target=.git  \
     if [ "$USE_SCCACHE" != "1" ]; then \
-        PIP_EXTRA_INDEX_URL="$([ "${TARGETPLATFORM}" = "linux/arm64" ] && echo "https://download.pytorch.org/whl/nightly/cu124")" \
-            python3 -m build --wheel --outdir=dist --no-isolation -v --config-setting py-limited-api=38 -v ;\
+        python3 -m build --wheel --outdir=dist --no-isolation -v --config-setting py-limited-api=38 -v ;\
     fi
 
 # Check the size of the wheel if RUN_WHEEL_CHECK is true

diff --git a/requirements-cuda.txt b/requirements-cuda.txt
@@ -9,6 +9,5 @@ torchaudio==2.5.1
 torchvision == 0.20.1; platform_machine != "aarch64" # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version
 xformers == 0.0.28.post3; platform_system == 'Linux' and platform_machine == 'x86_64'  # Requires PyTorch 2.5.1
 
-# FIXME: add torchaudio dependency
-# aarch64 cuda deps
---pre --extra-index-url https://download.pytorch.org/whl/nightly/cu126 torch==2.7.0.dev20250121+cu126 torchvision==0.22.0.dev20250121; platform_system == "Linux" and platform_machine == "aarch64"
+# aarch64 deps
+--pre --extra-index-url https://download.pytorch.org/whl/nightly/cu126 torch==2.7.0.dev20250121+cu126 torchvision==0.22.0.dev20250121 torchaudio==2.6.0.dev20250121+cu126; platform_system == "Linux" and platform_machine == "aarch64"