[ROCm] Harden HIP version for native library naming and selection

.github/scripts/build-rocm.sh

@@ -1,26 +1,24 @@
 #!/bin/bash
 set -xeuo pipefail
 
-: "${RUNNER_OS:?RUNNER_OS must be set (Linux/Windows)}"
-: "${ROCM_VERSION:?ROCM_VERSION must be set}"
+if [[ "${RUNNER_OS:-}" != "Linux" && "${RUNNER_OS:-}" != "Windows" ]]; then
+    echo "Invalid RUNNER_OS '${RUNNER_OS:-}'; expected Linux or Windows." >&2
+    exit 1
+fi
 
-rocm_version_at_least() {
-    local required_version="$1"
-    local current_major current_minor required_major required_minor
+if [[ ! "${ROCM_VERSION:-}" =~ ^([0-9]+)\.([0-9]+)(\.[0-9]+)?$ ]]; then
+    echo "Invalid ROCM_VERSION '${ROCM_VERSION:-}'; expected a dotted ROCm release such as 7.14." >&2
+    exit 1
+fi
 
-    IFS=. read -r current_major current_minor _ <<< "${ROCM_VERSION}"
-    IFS=. read -r required_major required_minor _ <<< "${required_version}"
+rocm_version_major="$((10#${BASH_REMATCH[1]}))"
+rocm_version_minor="$((10#${BASH_REMATCH[2]}))"
 
-    if ((current_major > required_major)); then
-        return 0
-    fi
-    if ((current_major < required_major)); then
-        return 1
-    fi
-    if ((current_minor >= required_minor)); then
-        return 0
-    fi
-    return 1
+rocm_version_at_least() {
+    local required_major required_minor
+    IFS=. read -r required_major required_minor <<< "$1"
+    ((rocm_version_major > required_major ||
+        (rocm_version_major == required_major && rocm_version_minor >= required_minor)))
 }
 
 bnb_rocm_arch="gfx90a;gfx942;gfx1100;gfx1101;gfx1102;gfx1103"
@@ -94,4 +92,17 @@ fi
 
 output_dir="output/${RUNNER_OS}/X64"
 mkdir -p "${output_dir}"
-(shopt -s nullglob && cp bitsandbytes/*.{so,dylib,dll} "${output_dir}")
+libraries=()
+shopt -s nullglob
+for extension in so dylib dll; do
+    for library in bitsandbytes/libbitsandbytes_rocm*."${extension}"; do
+        libraries+=("${library}")
+    done
+done
+shopt -u nullglob
+
+if [ "${#libraries[@]}" -eq 0 ]; then
+    echo "No ROCm-backend library was built (expected bitsandbytes/libbitsandbytes_rocm*.{so,dylib,dll})." >&2
+    exit 1
+fi
+cp "${libraries[@]}" "${output_dir}/"

CMakeLists.txt

@@ -10,17 +10,22 @@
 #                        Separate by semicolons, i.e. `-DCOMPUTE_CAPABILITY=89;90;100;120`
 #                        Check your compute capability here: https://developer.nvidia.com/cuda-gpus
 #  - PTXAS_VERBOSE: Pass the `-v` option to the PTX Assembler
-#  - ROCM_VERSION: Override the ROCm version shortcode used in the output library name.
-#                  Useful when PyTorch was built against a different ROCm version than the
-#                  system install. For example, `-DROCM_VERSION=70` produces
-#                  libbitsandbytes_rocm70.so even if the system has ROCm 7.2.
+#  - HIP_VERSION: Override the HIP version used in the ROCm-backend output library name.
+#                  Accepts a dotted version or compact shortcode. For example,
+#                  `-DHIP_VERSION=7.14` and `-DHIP_VERSION=714` both produce
+#                  libbitsandbytes_rocm714.so.
 cmake_minimum_required(VERSION 3.22.1)
 
 # On Windows with HIP backend, auto-detect compilers from ROCM_PATH before project()
 if(WIN32 AND COMPUTE_BACKEND STREQUAL "hip")
-    if(DEFINED ENV{ROCM_PATH})
-        file(TO_CMAKE_PATH "$ENV{ROCM_PATH}" ROCM_PATH)
+    if(NOT DEFINED ENV{ROCM_PATH} OR "$ENV{ROCM_PATH}" STREQUAL "")
+        message(FATAL_ERROR
+            "ROCM_PATH must be set for HIP builds on Windows. "
+            "After 'rocm-sdk init', set it from 'rocm-sdk path --root'. "
+            "PowerShell: $env:ROCM_PATH = (rocm-sdk path --root)"
+        )
     endif()
+    file(TO_CMAKE_PATH "$ENV{ROCM_PATH}" ROCM_PATH)
     if(ROCM_PATH AND NOT DEFINED CMAKE_CXX_COMPILER)
         set(CMAKE_CXX_COMPILER "${ROCM_PATH}/lib/llvm/bin/clang++.exe")
     endif()
@@ -270,6 +275,11 @@ elseif(BUILD_HIP)
       set(CMAKE_HIP_ARCHITECTURES "gfx90a;gfx942;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201")
     endif()
 
+    find_program(HIPCONFIG_EXECUTABLE hipconfig)
+    if(NOT HIPCONFIG_EXECUTABLE)
+        message(FATAL_ERROR "hipconfig was not found on PATH; a ROCm/HIP SDK is required for HIP builds")
+    endif()
+
     enable_language(HIP)
     message(STATUS "HIP Compiler: ${CMAKE_HIP_COMPILER}")
     message(STATUS "HIP Targets: ${CMAKE_HIP_ARCHITECTURES}")
@@ -282,20 +292,32 @@ elseif(BUILD_HIP)
 
     string(APPEND BNB_OUTPUT_NAME "_rocm")
 
-    # get hip version
-    execute_process(COMMAND hipconfig --version OUTPUT_VARIABLE HIP_CONFIG_VERSION)
-    string(REGEX MATCH "[0-9]+\\.[0-9]+" HIP_VERSION "${HIP_CONFIG_VERSION}")
-    string(REPLACE "." "" HIP_VERSION_SHORT "${HIP_VERSION}")
+    execute_process(
+        COMMAND "${HIPCONFIG_EXECUTABLE}" --version
+        OUTPUT_VARIABLE HIPCONFIG_VERSION
+        OUTPUT_STRIP_TRAILING_WHITESPACE
+        COMMAND_ERROR_IS_FATAL ANY
+    )
 
-    # Expose a cache variable that the user can set to override the ROCm version in the library name
-    set(ROCM_VERSION "${HIP_VERSION_SHORT}" CACHE STRING "Expected ROCm Version Shortcode")
+    if(NOT HIPCONFIG_VERSION MATCHES "([0-9]+)\\.([0-9]+)")
+        message(FATAL_ERROR "Could not parse HIP version: ${HIPCONFIG_VERSION}")
+    endif()
+    set(DETECTED_HIP_VERSION_TAG "${CMAKE_MATCH_1}${CMAKE_MATCH_2}")
 
-    message(STATUS "ROCm Version: ${HIP_VERSION_SHORT} (from hipconfig)")
-    if(NOT ROCM_VERSION STREQUAL "${HIP_VERSION_SHORT}")
-        message(WARNING "Overriding ROCm version in library name: ${HIP_VERSION_SHORT} -> ${ROCM_VERSION}")
+    set(HIP_VERSION "${DETECTED_HIP_VERSION_TAG}" CACHE STRING
+        "HIP version used in the ROCm-backend library name"
+    )
+    if(HIP_VERSION MATCHES "^([0-9]+)\\.([0-9]+)(\\.[0-9]+)?$")
+        set(HIP_VERSION_TAG "${CMAKE_MATCH_1}${CMAKE_MATCH_2}")
+    elseif(HIP_VERSION MATCHES "^[0-9]+$")
+        set(HIP_VERSION_TAG "${HIP_VERSION}")
+    else()
+        message(FATAL_ERROR "HIP_VERSION must be dotted or compact, such as 7.14 or 714")
     endif()
 
-    string(APPEND BNB_OUTPUT_NAME "${ROCM_VERSION}")
+    message(STATUS "HIP ${HIPCONFIG_VERSION}; library suffix: rocm${HIP_VERSION_TAG}")
+
+    string(APPEND BNB_OUTPUT_NAME "${HIP_VERSION_TAG}")
     add_compile_definitions(__HIP_PLATFORM_AMD__)
     add_compile_definitions(__HIP_PLATFORM_HCC__)
     add_compile_definitions(BUILD_HIP)
@@ -416,11 +438,16 @@ if(BUILD_CUDA)
     )
 endif()
 if(BUILD_HIP)
-    # Determine ROCM_PATH from environment variable, fallback to /opt/rocm on Linux
-    if(DEFINED ENV{ROCM_PATH})
-      file(TO_CMAKE_PATH "$ENV{ROCM_PATH}" ROCM_PATH)
+    # Determine ROCM_PATH from an existing CMake value or environment variable.
+    if(ROCM_PATH)
+        file(TO_CMAKE_PATH "${ROCM_PATH}" ROCM_PATH)
+    elseif(DEFINED ENV{ROCM_PATH} AND NOT "$ENV{ROCM_PATH}" STREQUAL "")
+        file(TO_CMAKE_PATH "$ENV{ROCM_PATH}" ROCM_PATH)
+    elseif(WIN32)
+        message(FATAL_ERROR "ROCM_PATH must be set for HIP builds on Windows")
     else()
-      set(ROCM_PATH /opt/rocm)
+        message(WARNING "ROCM_PATH is not set; falling back to /opt/rocm")
+        set(ROCM_PATH /opt/rocm)
     endif()
     list(APPEND CMAKE_PREFIX_PATH ${ROCM_PATH})
     macro(find_package_and_print_version PACKAGE_NAME)
@@ -452,12 +479,8 @@ if(BUILD_HIP)
     set_source_files_properties(${GPU_FILES} PROPERTIES LANGUAGE HIP)
     set_target_properties(bitsandbytes PROPERTIES LINKER_LANGUAGE CXX)
 
-    if(HIP_VERSION VERSION_LESS "6.1")
-	target_compile_definitions(bitsandbytes PUBLIC NO_HIPBLASLT)
-    else()
-	find_package(hipblaslt)
-        target_link_libraries(bitsandbytes PUBLIC roc::hipblaslt)
-    endif()
+    find_package(hipblaslt REQUIRED)
+    target_link_libraries(bitsandbytes PUBLIC roc::hipblaslt)
 endif()
 if(BUILD_XPU)
     set(SYCL_LINK_FLAGS "-fsycl;--offload-compress;-fsycl-targets=spir64_gen,spir64;-Xs;-device pvc,xe-lpg,ats-m150 -options ' -cl-intel-enable-auto-large-GRF-mode -cl-poison-unsupported-fp64-kernels -cl-intel-greater-than-4GB-buffer-required'")

bitsandbytes/cextension.py

@@ -25,58 +25,89 @@ def get_cuda_bnb_library_path(cuda_specs: CUDASpecs) -> Path:
 
     When no override is set, selects from packaged libraries using the following priority:
     1. Exact version match.
-    2. Highest packaged version <= runtime version, same major (e.g. runtime 12.9, ship 12.8).
-    3. Lowest packaged version > runtime version, same major (e.g. runtime 12.0, ship 12.1).
-    No cross-major fallback: if no same-major library exists, returns the exact non-existent
-    path so the caller raises a clear "not found" error.
-    A warning is logged when falling back. Override env vars bypass selection entirely
-    and load the named version with no fallback. The returned path is not guaranteed to
-    exist when no packaged libs are found, or when an override names an absent version.
+    2. Highest packaged version <= runtime version, same major (e.g. runtime 12.9, packaged 12.8).
+    3. Lowest packaged version > runtime version, same major (e.g. runtime 12.0, packaged 12.1).
+    4. For the ROCm backend only, repeat the same older-first selection across HIP major versions.
+    ROCm-backend selection uses the version tuple reported by ``torch.version.hip``; it does
+    not infer a ROCm release from ``torch.version.rocm``. CUDA does not fall back across major
+    versions. A warning is logged when falling back.
+    Overrides select the requested filename directly. The returned path is not guaranteed
+    to exist when no packaged libraries are found or an override names an absent version.
     """
     is_hip = bool(torch.version.hip)
     prefix = "rocm" if is_hip else "cuda"
     override_var = "BNB_ROCM_VERSION" if is_hip else "BNB_CUDA_VERSION"
+    other_override_var = "BNB_CUDA_VERSION" if is_hip else "BNB_ROCM_VERSION"
+
+    if os.environ.get(other_override_var):
+        logger.warning(
+            "%s is ignored because PyTorch is using %s; use %s instead.",
+            other_override_var,
+            "ROCm" if is_hip else "CUDA",
+            override_var,
+        )
 
     override_value = os.environ.get(override_var)
 
     if override_value is not None:
-        if not override_value.isdigit():
-            raise RuntimeError(f"{override_var}={override_value!r}: value must be digits only (e.g. '124' for 12.4).")
-        library_name = f"libbitsandbytes_{prefix}{override_value}{DYNAMIC_LIBRARY_SUFFIX}"
+        try:
+            override_version = _parse_version_override(override_value, is_hip)
+        except ValueError as error:
+            example = "7.14 or 714" if is_hip else "12.8 or 128"
+            raise RuntimeError(
+                f"{override_var}={override_value!r}: expected a dotted version or shortcode ({example})."
+            ) from error
+
+        version_tag = _format_native_version_tag(override_version)
+        library_name = f"libbitsandbytes_{prefix}{version_tag}{DYNAMIC_LIBRARY_SUFFIX}"
+        override_path = PACKAGE_DIR / library_name
         logger.warning(
-            f"WARNING: {override_var}={override_value} environment variable detected; loading {library_name}.\n"
+            f"WARNING: {override_var}={override_value} environment variable detected; "
+            f"loading {override_path.name}.\n"
             f"This overrides automatic {'ROCm' if is_hip else 'CUDA'} version selection.\n"
             f"If this was unintended clear the variable and retry: unset {override_var}\n",
         )
-        return PACKAGE_DIR / library_name
+        return override_path
 
     available = _find_cuda_libs(prefix, is_hip)
     runtime_version = cuda_specs.cuda_version_tuple
+    runtime_tag = _format_native_version_tag(runtime_version)
 
     if not available:
-        return PACKAGE_DIR / f"libbitsandbytes_{prefix}{cuda_specs.cuda_version_string}{DYNAMIC_LIBRARY_SUFFIX}"
+        return PACKAGE_DIR / f"libbitsandbytes_{prefix}{runtime_tag}{DYNAMIC_LIBRARY_SUFFIX}"
 
     if runtime_version in available:
         return available[runtime_version]
 
-    lower = [v for v in available if v[0] == runtime_version[0] and v < runtime_version]
+    missing_path = PACKAGE_DIR / f"libbitsandbytes_{prefix}{runtime_tag}{DYNAMIC_LIBRARY_SUFFIX}"
+    same_major = [version for version in available if version[0] == runtime_version[0]]
+    cross_major = False
+    lower = [version for version in same_major if version < runtime_version]
     if lower:
         selected = max(lower)
+    elif same_major:
+        selected = min(same_major)
+    elif is_hip:
+        lower = [version for version in available if version < runtime_version]
+        selected = max(lower) if lower else min(available)
+        cross_major = True
     else:
-        higher_same = [v for v in available if v[0] == runtime_version[0] and v > runtime_version]
-        if higher_same:
-            selected = min(higher_same)
-        else:
-            # No same-major library available. Return the non-existent exact path so
-            # get_native_library() raises a clear "not found" error.
-            return PACKAGE_DIR / f"libbitsandbytes_{prefix}{cuda_specs.cuda_version_string}{DYNAMIC_LIBRARY_SUFFIX}"
-
-    logger.warning(
-        f"No prebuilt binary for {'ROCm' if is_hip else 'CUDA'} "
-        f"{runtime_version[0]}.{runtime_version[1]}, loading "
-        f"{'ROCm' if is_hip else 'CUDA'} {selected[0]}.{selected[1]} instead. "
-        f"Set {override_var} to override."
-    )
+        return missing_path
+
+    if cross_major:
+        logger.warning(
+            f"No prebuilt ROCm-backend binary for HIP {runtime_version[0]}.{runtime_version[1]}, loading "
+            f"HIP {selected[0]}.{selected[1]} across major versions. This binary may be incompatible "
+            "or may not contain code for your GPU architecture. "
+            f"Set {override_var} to override or compile from source."
+        )
+    else:
+        logger.warning(
+            f"No prebuilt binary for {'ROCm-backend HIP' if is_hip else 'CUDA'} "
+            f"{runtime_version[0]}.{runtime_version[1]}, loading "
+            f"{'HIP' if is_hip else 'CUDA'} {selected[0]}.{selected[1]} instead. "
+            f"Set {override_var} to override."
+        )
     return available[selected]
 
 
@@ -124,26 +155,47 @@ def __init__(self, lib: ct.CDLL):
             lib.cget_managed_ptr.restype = ct.c_void_p
 
 
-def _split_cuda_version(compact: str, is_hip: bool) -> tuple[int, int]:
-    """Split a compact CUDA/ROCm version string from a library filename into (major, minor).
+def _format_native_version_tag(version: tuple[int, int]) -> str:
+    major, minor = version
+    return f"{major}{minor}"
+
+
+def _split_cuda_version(version_tag: str, is_hip: bool) -> tuple[int, int]:
+    """Split a CUDA/ROCm library filename tag into (major, minor).
 
     CUDA: major is always 2 digits (11, 12, 13...), e.g. '118' -> (11, 8), '132' -> (13, 2).
-    ROCm: major is always 1 digit for now (6, 7...), e.g. '72' -> (7, 2), '713' -> (7, 13).
-    Note: revisit if ROCm major reaches 10.
+    ROCm: supported majors 6-9 use one digit, e.g. '72' -> (7, 2)
+    and '713' -> (7, 13). Tags starting with 1-5 reserve two major digits.
     """
     if is_hip:
-        return int(compact[:1]), int(compact[1:])
-    return int(compact[:2]), int(compact[2:])
+        if len(version_tag) >= 3 and version_tag[0] in "12345":
+            return int(version_tag[:2]), int(version_tag[2:])
+        return int(version_tag[:1]), int(version_tag[1:])
+    return int(version_tag[:2]), int(version_tag[2:])
+
+
+def _parse_version_override(value: str, is_hip: bool) -> tuple[int, int]:
+    dotted = re.fullmatch(r"(\d+)\.(\d+)(?:\.\d+(?:[A-Za-z0-9+_.-]*)?)?", value)
+    if dotted:
+        return int(dotted.group(1)), int(dotted.group(2))
+    if value.isdigit():
+        return _split_cuda_version(value, is_hip)
+    raise ValueError(f"Invalid version override: {value}")
 
 
 def _find_cuda_libs(prefix: str, is_hip: bool) -> dict[tuple[int, int], Path]:
     """Return a {(major, minor): Path} mapping for all packaged CUDA/ROCm library files."""
     result = {}
     for lib in PACKAGE_DIR.glob(f"libbitsandbytes_{prefix}*{DYNAMIC_LIBRARY_SUFFIX}"):
-        match = re.search(rf"{prefix}(\d+)", lib.name)
+        match = re.fullmatch(
+            rf"libbitsandbytes_{re.escape(prefix)}(\d+){re.escape(DYNAMIC_LIBRARY_SUFFIX)}",
+            lib.name,
+        )
         if match:
             try:
-                result[_split_cuda_version(match.group(1), is_hip)] = lib
+                version_tag = match.group(1)
+                version = _split_cuda_version(version_tag, is_hip)
+                result[version] = lib
             except (ValueError, IndexError):
                 continue
     return result
@@ -153,11 +205,11 @@ def get_available_cuda_binary_versions() -> list[str]:
     """Get formatted CUDA/ROCm versions from existing library files."""
     is_hip = bool(torch.version.hip)
     prefix = "rocm" if is_hip else "cuda"
-    return sorted(f"{major}.{minor}" for major, minor in _find_cuda_libs(prefix, is_hip))
+    return [f"{major}.{minor}" for major, minor in sorted(_find_cuda_libs(prefix, is_hip))]
 
 
 def parse_cuda_version(version_str: str) -> str:
-    """Convert a raw version code string (e.g. '118', '713') to a dotted version (e.g. '11.8', '7.13')."""
+    """Convert a compact version tag (e.g. '118', '714') to a dotted version."""
     if version_str.isdigit():
         is_hip = bool(torch.version.hip)
         try:
@@ -275,7 +327,7 @@ def _format_lib_error_message(
                 "You have two options:\n"
                 "1. COMPILE FROM SOURCE as mentioned here:\n"
                 "   https://huggingface.co/docs/bitsandbytes/main/en/installation?backend=AMD+ROCm#amd-gpu\n"
-                "2. Use BNB_ROCM_VERSION to specify a DIFFERENT ROCm version from the detected one, matching the version the library was built with.\n\n"
+                "2. Use BNB_ROCM_VERSION to select a DIFFERENT HIP-version suffix for the ROCm-backend binary, matching the version the library was built with.\n\n"
             )
         )