Add ROCM_VERSION to triton dependency

ROCm · Mar 22, 2024 · be72ec9 · be72ec9
1 parent 608d4f4
commit be72ec9
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/manywheel/build_rocm.sh b/manywheel/build_rocm.sh
@@ -282,9 +282,9 @@ if [ ${PYTORCH_VERSION%%\.*} -ge 2 ]; then
         TRITON_VERSION=$(cat $PYTORCH_ROOT/.ci/docker/triton_version.txt)
 
         if [[ -z "$PYTORCH_EXTRA_INSTALL_REQUIREMENTS" ]]; then
-            export PYTORCH_EXTRA_INSTALL_REQUIREMENTS="pytorch-triton-rocm==${TRITON_VERSION}+${TRITON_SHORTHASH}"
+            export PYTORCH_EXTRA_INSTALL_REQUIREMENTS="pytorch-triton-rocm==${TRITON_VERSION}+${ROCM_VERSION}_${TRITON_SHORTHASH}"
         else
-            export PYTORCH_EXTRA_INSTALL_REQUIREMENTS="${PYTORCH_EXTRA_INSTALL_REQUIREMENTS} | pytorch-triton-rocm==${TRITON_VERSION}+${TRITON_SHORTHASH}"
+            export PYTORCH_EXTRA_INSTALL_REQUIREMENTS="${PYTORCH_EXTRA_INSTALL_REQUIREMENTS} | pytorch-triton-rocm==${TRITON_VERSION}+${ROCM_VERSION}_${TRITON_SHORTHASH}"
         fi
     fi
 fi