Revert "[main] Conditionalize triton package name on rocm version (#91)"

This reverts commit 66b8d8d.

Author: pragupta

manywheel/build_rocm.sh

@@ -118,7 +118,7 @@ do_lightweight_build() {
     else
         BUILD_SCRIPT=build_libtorch.sh
     fi
-    export WHEELNAME_MARKER="${LIGHTWEIGHT_WHEELNAME_MARKER}" 
+    export WHEELNAME_MARKER="${LIGHTWEIGHT_WHEELNAME_MARKER}"
     source "$SCRIPTPATH/${BUILD_SCRIPT}"
 
     echo "=== Done building LIGHTWEIGHT variant ==="
@@ -152,7 +152,7 @@ do_heavyweight_build() {
         "libroctracer64.so"
         "libroctx64.so"
     )
-    
+
     # Adjust list based on ROCm version
     if [[ $ROCM_INT -ge 60100 ]]; then
         HEAVYWEIGHT_ROCM_SO_FILES+=("librocprofiler-register.so")
@@ -165,8 +165,8 @@ do_heavyweight_build() {
         HEAVYWEIGHT_ROCM_SO_FILES+=("libhipsparselt.so")
 	HEAVYWEIGHT_ROCM_SO_FILES+=("librocprofiler-sdk.so")
     fi
-    
-    
+
+
     OS_NAME=`awk -F= '/^NAME/{print $2}' /etc/os-release`
     if [[ "$OS_NAME" == *"CentOS Linux"* || "$OS_NAME" == *"AlmaLinux"* ]]; then
         LIBGOMP_PATH="/usr/lib64/libgomp.so.1"
@@ -246,15 +246,15 @@ do_heavyweight_build() {
         file_name="${lib##*/}" # Substring removal of path to get filename
         OS_SO_FILES[${#OS_SO_FILES[@]}]=$file_name # Append lib to array
     done
-    
+
     # rocBLAS library files
     ROCBLAS_LIB_SRC=$ROCM_HOME/lib/rocblas/library
     ROCBLAS_LIB_DST=lib/rocblas/library
     ARCH=$(echo $PYTORCH_ROCM_ARCH | sed 's/;/|/g') # Replace ; seperated arch list to bar for grep
     ARCH_SPECIFIC_FILES=$(ls $ROCBLAS_LIB_SRC | grep -E $ARCH)
     OTHER_FILES=$(ls $ROCBLAS_LIB_SRC | grep -v gfx)
     ROCBLAS_LIB_FILES=($ARCH_SPECIFIC_FILES $OTHER_FILES)
-    
+
     # hipblaslt library files
     HIPBLASLT_LIB_SRC=$ROCM_HOME/lib/hipblaslt/library
     HIPBLASLT_LIB_DST=lib/hipblaslt/library
@@ -354,11 +354,12 @@ if [ ${PYTORCH_VERSION%%\.*} -ge 2 ]; then
             TRITON_SHORTHASH=$(cut -c1-8 $PYTORCH_ROOT/.ci/docker/ci_commit_pins/triton-rocm.txt)
 	fi
         TRITON_VERSION=$(cat $PYTORCH_ROOT/.ci/docker/triton_version.txt)
-	# Only linux Python < 3.13 are supported wheels for triton
-	TRITON_CONSTRAINT="platform_system == 'Linux' and platform_machine == 'x86_64'$(if [[ $(ver "$PYTORCH_VERSION") -le $(ver "2.5") ]]; then echo " and python_version < '3.13'"; fi)"
-
-        if [[ -z "$PYTORCH_EXTRA_INSTALL_REQUIREMENTS" ]]; then
-            export PYTORCH_EXTRA_INSTALL_REQUIREMENTS="pytorch-triton-rocm==${TRITON_VERSION}+${ROCM_VERSION_WITH_PATCH}.git${TRITON_SHORTHASH}; ${TRITON_CONSTRAINT}"
+        # Only linux Python < 3.13 are supported wheels for triton
+        TRITON_CONSTRAINT="platform_system == 'Linux' and platform_machine == 'x86_64'$(if [[ $(ver "$PYTORCH_VERSION") -le $(ver "2.5") ]]; then echo " and python_version < '3.13'"; fi)"
+        # Use "triton" for dev builds, else "pytorch-triton-rocm"
+        # Temp: Currently enabling for rocm7.1_internal_testing branch only but plan to expand it to other branches
+        if [[ "$PYTORCH_VERSION_FULL" == *"2.9.0a0"* ]]; then
+            PKG="triton"
         else
             export PYTORCH_EXTRA_INSTALL_REQUIREMENTS="${PYTORCH_EXTRA_INSTALL_REQUIREMENTS} | pytorch-triton-rocm==${TRITON_VERSION}+${ROCM_VERSION_WITH_PATCH}.git${TRITON_SHORTHASH}; ${TRITON_CONSTRAINT}"
         fi
@@ -371,7 +372,7 @@ echo "PYTORCH_ROCM_ARCH: ${PYTORCH_ROCM_ARCH}"
 export LIGHTWEIGHT_WHEELNAME_MARKER="${LIGHTWEIGHT_WHEELNAME_MARKER}"
 SCRIPTPATH="$( cd "$(dirname "$0")" ; pwd -P )"
 if [[ -z "$BUILD_PYTHONLESS" ]]; then
-    BUILD_SCRIPT=build_torch_wheel.sh	
+    BUILD_SCRIPT=build_torch_wheel.sh
 fi
 source $SCRIPTPATH/${BUILD_SCRIPT}
 if [[ "$BUILD_LIGHTWEIGHT" == "1" ]]; then