diff --git a/kernel-builder/src/pyproject/templates/torch/preamble.cmake b/kernel-builder/src/pyproject/templates/torch/preamble.cmake index cc4ab774..1189b238 100644 --- a/kernel-builder/src/pyproject/templates/torch/preamble.cmake +++ b/kernel-builder/src/pyproject/templates/torch/preamble.cmake @@ -279,7 +279,7 @@ endif() if(GPU_LANG STREQUAL "CUDA") generate_build_name(BUILD_VARIANT_NAME "${TORCH_VERSION}" "cuda" "${CUDA_VERSION}" ${_STABLE_ABI_ARG}) elseif(GPU_LANG STREQUAL "HIP") - run_python(ROCM_VERSION "import torch.version; print(torch.version.hip.split('.')[0] + '.' + torch.version.hip.split('.')[1])" "Failed to get ROCm version") + run_python(ROCM_VERSION "import torch.version; v = getattr(torch.version, 'rocm', None) or torch.version.hip; print(v.split('.')[0] + '.' + v.split('.')[1])" "Failed to get ROCm version") generate_build_name(BUILD_VARIANT_NAME "${TORCH_VERSION}" "rocm" "${ROCM_VERSION}" ${_STABLE_ABI_ARG}) elseif(GPU_LANG STREQUAL "SYCL") generate_build_name(BUILD_VARIANT_NAME "${TORCH_VERSION}" "xpu" "${DPCPP_VERSION}" ${_STABLE_ABI_ARG}) diff --git a/kernel-builder/src/pyproject/templates/tvm_ffi/preamble.cmake b/kernel-builder/src/pyproject/templates/tvm_ffi/preamble.cmake index 86a66535..57650aa6 100644 --- a/kernel-builder/src/pyproject/templates/tvm_ffi/preamble.cmake +++ b/kernel-builder/src/pyproject/templates/tvm_ffi/preamble.cmake @@ -129,7 +129,12 @@ elseif(GPU_LANG STREQUAL "HIP") enable_language(HIP) find_package(hip REQUIRED) - set(ROCM_VERSION "${hip_VERSION_MAJOR}.${hip_VERSION_MINOR}") + find_package(rocm-core CONFIG QUIET) + if(rocm-core_FOUND) + set(ROCM_VERSION "${rocm-core_VERSION_MAJOR}.${rocm-core_VERSION_MINOR}") + else() + set(ROCM_VERSION "${hip_VERSION_MAJOR}.${hip_VERSION_MINOR}") + endif() message(STATUS "Found ROCm version: ${ROCM_VERSION}") set(ROCM_ARCHS "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1030;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1200;gfx1201") diff --git a/kernels/src/kernels/backends.py b/kernels/src/kernels/backends.py index b8e20a8a..f77c4061 100644 --- a/kernels/src/kernels/backends.py +++ b/kernels/src/kernels/backends.py @@ -228,7 +228,12 @@ def _backend() -> Backend: cuda_version = Version(torch.version.cuda) return CUDA(version=cuda_version) elif torch.version.hip is not None: - rocm_version = Version(torch.version.hip.split("-")[0]) + # Note: the ROCm version does not always match the HIP version + # (e.g. ROCm 10.0 uses HIP 7.15). So, we should use + # `torch.version.rocm`. However, this attribute is only available + # since Torch 2.10, so use `torch.version.hip` as a fallback. + rocm = getattr(torch.version, "rocm", None) or torch.version.hip + rocm_version = Version(rocm.split("-")[0]) return ROCm(version=rocm_version) elif torch.backends.mps.is_available(): return Metal() diff --git a/kernels/src/kernels/cli/benchmark.py b/kernels/src/kernels/cli/benchmark.py index 55cb5040..e74ff860 100644 --- a/kernels/src/kernels/cli/benchmark.py +++ b/kernels/src/kernels/cli/benchmark.py @@ -403,7 +403,8 @@ def collect_machine_info() -> MachineInfo: gpu = torch.cuda.get_device_name(0) # ROCm uses the CUDA API but has torch.version.hip if hasattr(torch.version, "hip") and torch.version.hip: - backend_type = f"ROCm {torch.version.hip}" + rocm_version = getattr(torch.version, "rocm", None) or torch.version.hip + backend_type = f"ROCm {rocm_version}" else: backend_type = f"CUDA {torch.version.cuda}" if torch.version.cuda else "CUDA" elif backend_name == "xpu": diff --git a/nix-builder/pkgs/python-modules/torch/binary/generic.nix b/nix-builder/pkgs/python-modules/torch/binary/generic.nix index 47526789..808ece90 100644 --- a/nix-builder/pkgs/python-modules/torch/binary/generic.nix +++ b/nix-builder/pkgs/python-modules/torch/binary/generic.nix @@ -107,6 +107,9 @@ let amdrocm-rand amdrocm-rccl ] + ++ lib.optionals (versionAtLeast version "2.15") [ + amdrocm-amdsmi + ] else [ aotriton