diff --git a/aten/src/ATen/cuda/detail/CUDAHooks.cpp b/aten/src/ATen/cuda/detail/CUDAHooks.cpp index 4dfc3003e0cdb..7b7ed89c4adab 100644 --- a/aten/src/ATen/cuda/detail/CUDAHooks.cpp +++ b/aten/src/ATen/cuda/detail/CUDAHooks.cpp @@ -560,6 +560,9 @@ const std::vector& CUDAHooks::getHipblasltPreferredArchs() const { #if ROCM_VERSION >= 70000 "gfx950", #endif +#if ROCM_VERSION >= 71300 + "gfx1100", "gfx1101", "gfx1151", +#endif #if ROCM_VERSION >= 70200 "gfx1250" #endif diff --git a/test/test_cuda.py b/test/test_cuda.py index cc209e5f03b28..890fb7622f351 100644 --- a/test/test_cuda.py +++ b/test/test_cuda.py @@ -681,10 +681,18 @@ def _check_default(): self.assertTrue(default == torch._C._BlasBackend.Cublas) else: # ROCm logic is less so, it's cublaslt for some Instinct, cublas for all else - gcn_arch = str( - torch.cuda.get_device_properties(0).gcnArchName.split(":", 1)[0] - ) - if gcn_arch in ["gfx90a", "gfx942", "gfx950", "gfx1200", "gfx1201", "gfx1250"]: + # Mirror CUDAHooks::getHipblasltPreferredArchs in CUDAHooks.cpp + ROCM_VERSION = tuple(int(v) for v in torch.version.hip.split(".")[:2]) + archs = ["gfx90a", "gfx942"] + if ROCM_VERSION >= (6, 4): + archs.extend(["gfx1200", "gfx1201"]) + if ROCM_VERSION >= (7, 0): + archs.append("gfx950") + if ROCM_VERSION >= (7, 13): + archs.extend(["gfx1100", "gfx1101", "gfx1151"]) + gcn_arch_name = torch.cuda.get_device_properties(0).gcnArchName + hipblaslt_preferred = any(arch in gcn_arch_name for arch in archs) + if hipblaslt_preferred: self.assertTrue(default == torch._C._BlasBackend.Cublaslt) else: self.assertTrue(default == torch._C._BlasBackend.Cublas)