TE building over TheRock

.github/scripts/aiter_prebuild_upload.sh

@@ -12,7 +12,14 @@ set -euo pipefail
 
 # Derive ROCm version and aiter commit -> cache key
 ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
-ROCM_PATH="${ROCM_PATH:-/opt/rocm}"
+if [ -n "${ROCM_PATH:-}" ]; then
+    true # Use provided ROCM_PATH
+elif [ -d "/opt/rocm/core" ]; then
+    ROCM_PATH="/opt/rocm/core"
+else
+    ROCM_PATH="/opt/rocm"
+fi
+export ROCM_PATH
 ROCM_VER=`head -n1 "${ROCM_PATH}/.info/version" | cut -d. -f1`
 
 AITER_DIR="${ROOT_DIR}/3rdparty/aiter"

.github/workflows/rocm-ci.yml

@@ -70,7 +70,7 @@ jobs:
           docker ps -a
 
           echo ">>> ROCm Installation:"
-          ls -d /opt/rocm* || echo "No /opt/rocm found"
+          (ls -d /opt/rocm/core-* || ls -d /opt/rocm-* || echo "No default ROCm path found") 2>/dev/null || true
           echo ">>> GPU info:"
           ls -l /dev/dri
           ls -l /dev/kfd
@@ -247,11 +247,6 @@ jobs:
           set -x -o pipefail
           ulimit -c 0 # Disable core dumps
 
-          # debug output
-          ls -d /opt/rocm*
-          python --version
-          pip list | egrep "transformer_e|torch|jax|numpy|ml_dtypes|typing_ext"
-
           HIP_VISIBLE_DEVICES=1 ci/pytorch.sh > /workspace/torch_sgpu.log 2>&1 &
           torch_pid=$!; echo Pytorch test pid $!
 

build_tools/rocm_utils.cmake

@@ -0,0 +1,20 @@
+# Copyright (c) 2022-2026, Advanced Micro Devices, Inc. All rights reserved.
+# License for AMD contributions = MIT. See LICENSE for more information
+
+#Determine ROCM_PATH
+if(NOT "$ENV{ROCM_PATH}" STREQUAL "")
+    set(ROCM_PATH "$ENV{ROCM_PATH}")
+elseif(EXISTS "/opt/rocm/core")
+    set(ROCM_PATH "/opt/rocm/core")
+else()
+    set(ROCM_PATH "/opt/rocm")
+endif()
+
+#Configure target GPU architectures
+if(NOT DEFINED ENV{NVTE_ROCM_ARCH})
+    SET(CMAKE_HIP_ARCHITECTURES gfx942 gfx950)
+else()
+    # Accept comma separated list for NVTE_ROCM_ARCH
+    string(REPLACE "," ";" HIP_ARCH_LIST "$ENV{NVTE_ROCM_ARCH}")
+    SET(CMAKE_HIP_ARCHITECTURES ${HIP_ARCH_LIST})
+endif()

build_tools/utils.py

@@ -184,42 +184,30 @@ def rocm_build() -> bool:
     Determines which build platform to use:
 
     - If `NVTE_USE_ROCM` is set:
-        - Non-zero value: Use ROCm, if hipcc is detected.
+        - Any value except "0": Use ROCm, if hipcc is detected.
         - Zero value: Use CUDA, if nvcc is detected.
-    - If `NVTE_USE_ROCM` is not set:
-        - Attempt to auto-detect: Check for ROCm first, then CUDA.
+    - If `NVTE_USE_ROCM` is not set, guess if from `HIP_PLATFORM` if it is set.
+    - Otherwise auto-detect trying ROCm first.
 
     Returns:
         bool: `True` for ROCm, `False` for CUDA.
 
     Raises:
-        ValueError: If NVTE_USE_ROCM is set to invalid value.
         FileNotFoundError: If required tools (hipcc or nvcc) are not found.
     """
-    nvte_use_rocm = os.getenv("NVTE_USE_ROCM")
-    if nvte_use_rocm:
+    nvte_use_rocm = os.getenv("NVTE_USE_ROCM", "")
+    if not nvte_use_rocm:
+        match os.getenv("HIP_PLATFORM", ""):
+            case "amd": nvte_use_rocm = "1"
+            case "nvidia": nvte_use_rocm = "0"
+    if nvte_use_rocm != "0":
         try:
-            nvte_use_rocm = bool(int(nvte_use_rocm))
-        except ValueError:
-            raise ValueError(
-                f"Invalid value for NVTE_USE_ROCM: '{nvte_use_rocm}'.")
-
-        if nvte_use_rocm:
-            _, hipcc_bin = rocm_path()
-            if hipcc_bin.is_file():
-                return True
-            else:
-                raise FileNotFoundError(f"Could not find hipcc at {hipcc_bin}")
-        else:
-            nvcc_path()
-            return False
-
-    # Try to detect ROCm
-    _, hipcc_bin = rocm_path()
-    if hipcc_bin.is_file():
-        return True
-
-    # Try to detect CUDA
+            rocm_path()
+            return True
+        except FileNotFoundError:
+            if nvte_use_rocm:
+                raise FileNotFoundError("Could not find ROCm installation.")
+    # Try to detect CUDA if NVTE_USE_ROCM is set to "0" or ROCm is not found
     try:
         nvcc_path()
         return False
@@ -230,20 +218,26 @@ def rocm_build() -> bool:
 
 @functools.lru_cache(maxsize=None)
 def rocm_path() -> Tuple[str, str]:
-    """ROCm root path and HIPCC binary path as a tuple"""
-    """If ROCm installation is not specified, use default /opt/rocm path"""
+    """
+    ROCm root path and HIPCC binary path as a tuple
+    If ROCm installation is not specified, use default ROCm path
+    """
     hipcc_bin = None
     if os.getenv("ROCM_PATH"):
         rocm_home = Path(os.getenv("ROCM_PATH"))
         hipcc_bin = rocm_home / "bin" / "hipcc"
     if hipcc_bin is None:
         hipcc_bin = shutil.which("hipcc")
         if hipcc_bin is not None:
-            hipcc_bin = Path(hipcc_bin)
+            hipcc_bin = Path(hipcc_bin).resolve()
             rocm_home = hipcc_bin.parent.parent
     if hipcc_bin is None:
-        rocm_home = Path("/opt/rocm/")
+        rocm_home = Path("/opt/rocm/core")
+        if not rocm_home.is_dir():
+            rocm_home = Path("/opt/rocm/")
         hipcc_bin = rocm_home / "bin" / "hipcc"
+    if not hipcc_bin.is_file():
+        raise FileNotFoundError(f"Could not find hipcc at {hipcc_bin}")
     return rocm_home, hipcc_bin
 
 

build_tools/wheel_utils/Dockerfile.rocm.manylinux.x86

@@ -7,7 +7,7 @@
 
 # Build args:
 # BASE_IMAGE - Base manylinux image to use. Default: quay.io/pypa/manylinux_2_28_x86_64
-# ROCM_REPO_URL - ROCm repository URL. Default: https://repo.radeon.com/rocm/rhel8/latest/main/
+# ROCM_REPO_URL - ROCm TheRock repository URL. Default: https://repo.amd.com/rocm/packages/rhel8/x86_64/
 # GPU_TARGETS - Semicolon separated list of target GPU architectures. Default: "gfx942;gfx950"
 # TARGET_BRANCH - Target branch for TransformerEngine. Default: none (use git default)
 # GPU_TARGETS and TARGET_BRANCH can be overriden when start a container with NVTE_ROCM_ARCH and TARGET_BRANCH environment variables.
@@ -16,22 +16,25 @@
 ARG BASE_IMAGE=quay.io/pypa/manylinux_2_28_x86_64
 FROM $BASE_IMAGE
 
-ARG ROCM_REPO_URL=https://repo.radeon.com/rocm/rhel8/latest/main/
+ARG ROCM_REPO_URL=https://repo.amd.com/rocm/packages/rhel8/x86_64/
 
 # Set up ROCm repo
 RUN echo -e "[rocm]\nname=ROCm\nbaseurl=${ROCM_REPO_URL}\nenabled=1\ngpgcheck=0" > /etc/yum.repos.d/rocm.repo
 
 # Setup packages
-RUN dnf install -y --disablerepo=epel rocm-dev hipblaslt hipblaslt-devel hipcub hipcub-devel hiprand-devel rocrand-devel
-RUN dnf group install -y "Development Tools" && dnf install -y git cmake llvm-toolset gcc-toolset-12
+# ROCm TheRock packages are built for single GPU family. However, nothing from HW libraries is statically linked to TE,
+# so they can be used to build for multiple GPU targets. Using gfx950 here.
+RUN dnf install -y --disablerepo=epel amdrocm-core-devel-gfx950
 
-#Uncomment the next line for ROCm 6.4 cmake workaround: remove newer incomnpatible cmake preinstalled on base image
-#RUN rm /usr/local/bin/cmake || true
+# xz-devel installs lzma needed by AOTriton
+RUN dnf install -y gcc-toolset-13 xz-devel
 
 RUN dnf clean all
 RUN rm -rf /var/cache/dnf/*
 
-ENV CMAKE_PREFIX_PATH=/opt/rocm/lib/cmake
+ENV HIP_PLATFORM=amd
+ENV ROCM_PATH=/opt/rocm/core
+
 ENV NVTE_RELEASE_BUILD=1
 
 ARG GPU_TARGETS="gfx942;gfx950"

build_tools/wheel_utils/build_wheels.sh

@@ -22,17 +22,21 @@ mkdir -p /wheelhouse/logs
 git config --global --add safe.directory /TransformerEngine
 cd /TransformerEngine
 
-#If there is default Python installation, use it
-PYTHON=`which python || true`
-if [ -z "$PYTHON" ]; then
-        PYBINDIR=/opt/python/cp310-cp310/bin/
-        #hipify expects python in PATH, also ninja may be installed to python bindir
-        PATH="$PYBINDIR:$PATH"
-else
-        PYBINDIR="" #python bindir is already in PATHs
-fi
-
-ROCM_BUILD=$(${PYBINDIR}python -c "import build_tools.utils as u; print('true' if u.rocm_build() else 'false')")
+#hipify and aiter expect python in PATH, also ninja may be installed to python bindir
+#set it first because system python may be too old
+PATH="/opt/python/cp310-cp310/bin/:$PATH"
+
+case "$HIP_PLATFORM" in
+  amd)
+    ROCM_BUILD=true
+    ;;
+  nvidia)
+    ROCM_BUILD=false
+    ;;
+  *)
+    ROCM_BUILD=$(python -c "import build_tools.utils as u; print('true' if u.rocm_build() else 'false')")
+    ;;
+esac
 
 if [ "$LOCAL_TREE_BUILD" != "1" ]; then
         if $ROCM_BUILD ; then
@@ -45,22 +49,23 @@ else
 fi
 
 if $ROCM_BUILD ; then
-  #dataclasses, psutil are needed for AITER
-  ${PYBINDIR}pip install pybind11[global] ninja setuptools dataclasses psutil
+  pip install pybind11[global] ninja setuptools wheel
+  #modules needed to build AITER
+  pip install dataclasses psutil numpy pandas
+  export PATH=$PATH:$ROCM_PATH/bin
 else
-  PYBINDIR=/opt/python/cp310-cp310/bin/
   /opt/python/cp310-cp310/bin/pip install cmake pybind11[global] ninja setuptools wheel
 fi
 
 if $BUILD_METAPACKAGE ; then
         cd /TransformerEngine
-        NVTE_BUILD_METAPACKAGE=1 ${PYBINDIR}python setup.py bdist_wheel 2>&1 | tee /wheelhouse/logs/metapackage.txt
+        NVTE_BUILD_METAPACKAGE=1 python setup.py bdist_wheel 2>&1 | tee /wheelhouse/logs/metapackage.txt
         mv dist/* /wheelhouse/
 fi
 
 if $BUILD_COMMON -a $ROCM_BUILD ; then
         # Create the wheel.
-        ${PYBINDIR}python setup.py bdist_wheel --verbose --plat-name=$PLATFORM 2>&1 | tee /wheelhouse/logs/common.txt
+        python setup.py bdist_wheel --verbose --plat-name=$PLATFORM 2>&1 | tee /wheelhouse/logs/common.txt
 
         # Rename the wheel to make it python version agnostic.
         whl_name=$(basename dist/*)
@@ -93,8 +98,8 @@ fi
 if $BUILD_PYTORCH -a $ROCM_BUILD ; then
         cd /TransformerEngine/transformer_engine/pytorch
         #Only need torch for creating sdist, install CPU version to avoid installing CUDA/ROCm dependencies
-        ${PYBINDIR}pip install torch --index-url https://download.pytorch.org/whl/cpu
-        ${PYBINDIR}python setup.py sdist 2>&1 | tee /wheelhouse/logs/torch.txt
+        pip install torch --index-url https://download.pytorch.org/whl/cpu
+        python setup.py sdist 2>&1 | tee /wheelhouse/logs/torch.txt
         mv dist/* /wheelhouse/
 elif $BUILD_PYTORCH ; then
 	cd /TransformerEngine/transformer_engine/pytorch
@@ -105,8 +110,8 @@ fi
 
 if $BUILD_JAX -a $ROCM_BUILD ; then
         cd /TransformerEngine/transformer_engine/jax
-        ${PYBINDIR}pip install jax
-        ${PYBINDIR}python setup.py sdist 2>&1 | tee /wheelhouse/logs/jax.txt
+        pip install jax
+        python setup.py sdist 2>&1 | tee /wheelhouse/logs/jax.txt
         mv dist/* /wheelhouse/
 elif $BUILD_JAX ; then
 	cd /TransformerEngine/transformer_engine/jax

ci/_utils.sh

@@ -225,7 +225,15 @@ check_test_filter() {
 
 start_message() {
     echo "Started with TEST_LEVEL=$TEST_LEVEL sGPU='$TEST_SGPU' mGPU='$TEST_MGPU' at `date`"
-    echo "ROCm: `ls -d /opt/rocm-*`"
+    if [ -n "$ROCM_PATH" ]; then
+        _rocm_path="$ROCM_PATH"
+    elif [ -d "/opt/rocm/core" ]; then
+        _rocm_path="/opt/rocm/core"
+    else
+        _rocm_path="/opt/rocm"
+    fi
+    _rocm_path=`$REALPATH "$_rocm_path"`
+    test -d "$_rocm_path" && echo "ROCm: $_rocm_path" || echo "ROCm path not found"
     python --version
 }
 

tests/cpp/CMakeLists.txt

@@ -33,11 +33,7 @@ if(USE_ROCM)
   # Disable Asserts In Code (Can't use asserts on HIP stack.)
   add_definitions(-DNDEBUG)
   add_definitions(-DUSE_ROCM)
-  if(NOT DEFINED ENV{NVTE_ROCM_ARCH})
-    SET(CMAKE_HIP_ARCHITECTURES gfx942 gfx950)
-  else()
-    SET(CMAKE_HIP_ARCHITECTURES $ENV{NVTE_ROCM_ARCH})
-  endif()
+  include("${CMAKE_CURRENT_SOURCE_DIR}/../../build_tools/rocm_utils.cmake")
 else()
   if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES)
     if (CUDAToolkit_VERSION VERSION_GREATER_EQUAL 12.8)

transformer_engine/common/CMakeLists.txt

@@ -133,23 +133,17 @@ else()
     message(FATAL_ERROR "This project requires the Ninja build system. Install it using 'pip install ninja'.")
   endif()
 
+  include("${CMAKE_CURRENT_SOURCE_DIR}/../../build_tools/rocm_utils.cmake")
+
   # Disable Asserts In Code (Can't use asserts on HIP stack.)
   add_definitions(-DNDEBUG)
   add_definitions(-DUSE_ROCM)
 
-  if(NOT DEFINED ENV{NVTE_ROCM_ARCH})
-    SET(CMAKE_HIP_ARCHITECTURES gfx942 gfx950)
-  else()
-    SET(CMAKE_HIP_ARCHITECTURES $ENV{NVTE_ROCM_ARCH})
-  endif()
-
   # build error will be dup-ed parallel-jobs times
   # set(CMAKE_HIP_FLAGS "${CMAKE_HIP_FLAGS} -parallel-jobs=4")
   if(CMAKE_BUILD_TYPE STREQUAL "Debug")
     set(CMAKE_HIP_FLAGS "${CMAKE_HIP_FLAGS} -g")
   endif()
-
-  list(APPEND CMAKE_MODULE_PATH "/opt/rocm")
 endif()
 
 set(message_line "-------------------------------------------------------------")
@@ -323,6 +317,10 @@ else()
   message(STATUS "nvte hipified sources: ${te_hip_sources}")
 
   add_library(transformer_engine SHARED ${te_hip_sources})
+
+  # Workaround for TheRock installation that moved some headers from system-wide location
+  # to rocm_sysdeps but missing it in the default include path.
+  target_include_directories(transformer_engine SYSTEM PRIVATE "${ROCM_PATH}/lib/rocm_sysdeps/include")
 endif()
 
 target_include_directories(transformer_engine PUBLIC
@@ -562,11 +560,6 @@ endif()
 install(TARGETS transformer_engine DESTINATION .)
 if (USE_ROCM)
   set_target_properties(transformer_engine PROPERTIES INSTALL_RPATH "$ORIGIN/lib;$ORIGIN/transformer_engine/lib")
-  if("$ENV{ROCM_PATH}" STREQUAL "")
-    set(ROCM_PATH "/opt/rocm")
-  else()
-    set(ROCM_PATH "$ENV{ROCM_PATH}")
-  endif()
   file(READ "${ROCM_PATH}/.info/version" ROCM_VER)
   string(STRIP "${ROCM_VER}" ROCM_VER)
   string(REGEX MATCH "^[0-9]+\\.[0-9]+" ROCM_VER "${ROCM_VER}")

transformer_engine/common/__init__.py

@@ -421,7 +421,10 @@ def _load_core_library():
     if te_rocm_build:
         try:
             # Get installed ROCm version
-            with open(os.getenv("ROCM_PATH", "/opt/rocm") + "/.info/version", "r") as f:
+            for rocm_path in (os.getenv("ROCM_PATH"), "/opt/rocm/core", "/opt/rocm"):
+                if rocm_path and os.path.exists(os.path.join(rocm_path, ".info/version")):
+                    break
+            with open(os.path.join(rocm_path, ".info/version"), "r") as f:
                 rocm_version= f.read().strip().split('.')[:2]
 
             # Get ROCm version from the build info file

transformer_engine/common/ck_fused_attn/aiter_prebuilt.cmake

@@ -10,11 +10,9 @@ if(POLICY CMP0135)
   cmake_policy(SET CMP0135 NEW)
 endif()
 
+include("${CMAKE_CURRENT_SOURCE_DIR}/../../../build_tools/rocm_utils.cmake")
+
 # Extract ROCm version
-set(ROCM_PATH "$ENV{ROCM_PATH}")
-if("${ROCM_PATH}" STREQUAL "")
-  set(ROCM_PATH "/opt/rocm")
-endif()
 file(READ "${ROCM_PATH}/.info/version" ROCM_VER_CONTENT)
 string(STRIP "${ROCM_VER_CONTENT}" ROCM_VER_CONTENT)
 string(REGEX MATCH "^[0-9]+\\.[0-9]+" ROCM_VER "${ROCM_VER_CONTENT}")

transformer_engine/common/rocshmem_api/CMakeLists.txt

@@ -19,10 +19,12 @@ set_target_properties(rocshmemapi PROPERTIES
     POSITION_INDEPENDENT_CODE ON
 )
 
+include("${CMAKE_CURRENT_SOURCE_DIR}/../../../build_tools/rocm_utils.cmake")
+
 if(DEFINED ENV{ROCSHMEM_HOME})
   set(ROCSHMEM_HOME "$ENV{ROCSHMEM_HOME}" CACHE STRING "Location of ROCSHMEM installation")
 else()
-  set(ROCSHMEM_HOME "/opt/rocm" CACHE STRING "Location of ROCSHMEM installation (default)")
+  set(ROCSHMEM_HOME "${ROCM_PATH}" CACHE STRING "Location of ROCSHMEM installation (default)")
 endif()
 set(ROCSHMEM_LIBRARY_PATH "${ROCSHMEM_HOME}/lib/librocshmem.a")
 if (EXISTS ${ROCSHMEM_LIBRARY_PATH})

transformer_engine/common/util/cuda_runtime.cpp

@@ -178,6 +178,7 @@ const std::string &include_directory(bool required) {
 #ifdef __HIP_PLATFORM_AMD__
     std::vector<std::pair<std::string, Path>> search_paths = {{"ROCM_PATH", ""},
                                                               {"HIP_PATH", ""},
+                                                              {"", "/opt/rocm/core"},
                                                               {"", "/opt/rocm"}};
 #else
     std::vector<std::pair<std::string, Path>> search_paths = {{"NVTE_CUDA_INCLUDE_DIR", ""},