From 5c17a5d2a87be40cf3c38b033d6584f462c928c3 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Mon, 14 Sep 2026 19:48:04 +0000 Subject: [PATCH] [ROCm] Pair gfx1250 support with gfx1250-strict gfx1250-strict is a distinct LLVM/HIP offload target with the same CDNA5 ISA as gfx1250. Runtime equality checks against "gfx1250" alone miss it, and CMAKE HIP_ARCHITECTURES / CK GEMM filtering must treat both names as the same unsupported-CK family. Add gfx1250-strict next to gfx1250 in arch lists, hipBLASLt/hipSparseLt gates, shared-memory and workspace sizing, device preprocessor checks, wheel PYTORCH_ROCM_ARCH, and tests. Keep using the existing amd-gfx1250 AOTriton image, whose regex already matches gfx1250-strict. Test Plan: - rg gfx1250 to confirm remaining references include gfx1250-strict where they are architecture identifiers (AOTriton image name excluded). - No local ROCm gfx1250-strict hardware in this environment; unit tests that key off gcnArchName were updated in lockstep with CUDAHooks. Authored with an AI assistant. Co-authored-by: Bo Li --- .ci/docker/almalinux/build.sh | 2 +- .ci/docker/manywheel/build.sh | 2 +- aten/src/ATen/CMakeLists.txt | 10 +++++----- aten/src/ATen/cuda/CUDABlas.cpp | 1 + aten/src/ATen/cuda/CublasHandlePool.cpp | 4 ++-- aten/src/ATen/cuda/detail/CUDAHooks.cpp | 4 ++-- aten/src/ATen/native/cuda/GroupedBlas.cpp | 2 +- aten/src/ATen/native/cuda/KernelUtils.cuh | 2 +- aten/src/ATen/native/cuda/MemoryAccess.cuh | 8 ++++---- aten/src/ATen/native/cuda/ScaledBlas.cpp | 3 ++- aten/src/ATen/native/cuda/int4mm.cu | 10 +++++----- .../ATen/native/cuda/reduction_template.cuh | 2 +- .../ATen/native/sparse/cuda/cuSPARSELtOps.cpp | 2 +- cmake/Dependencies.cmake | 19 ++++++++++--------- cmake/External/aotriton.cmake | 2 ++ test/inductor/test_aot_inductor.py | 4 ++-- test/test_cuda.py | 6 +++--- test/test_linalg.py | 6 +++--- torch/cuda/_utils.py | 4 ++-- torch/testing/_internal/common_cuda.py | 18 ++++++++++-------- 20 files changed, 59 insertions(+), 52 deletions(-) diff --git a/.ci/docker/almalinux/build.sh b/.ci/docker/almalinux/build.sh index 01dfa8ae67bb2..5043419db0772 100755 --- a/.ci/docker/almalinux/build.sh +++ b/.ci/docker/almalinux/build.sh @@ -36,7 +36,7 @@ case ${DOCKER_TAG_PREFIX} in ;; rocm*) BASE_TARGET=rocm - PYTORCH_ROCM_ARCH="gfx900;gfx906;gfx908;gfx90a;gfx942;gfx1030;gfx1100;gfx1101;gfx1102;gfx1103;gfx1200;gfx1201;gfx950;gfx1150;gfx1151;gfx1250" + PYTORCH_ROCM_ARCH="gfx900;gfx906;gfx908;gfx90a;gfx942;gfx1030;gfx1100;gfx1101;gfx1102;gfx1103;gfx1200;gfx1201;gfx950;gfx1150;gfx1151;gfx1250;gfx1250-strict" EXTRA_BUILD_ARGS="${EXTRA_BUILD_ARGS} --build-arg PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH}" ;; *) diff --git a/.ci/docker/manywheel/build.sh b/.ci/docker/manywheel/build.sh index 4bcb1b6143c84..232ff9563a71b 100755 --- a/.ci/docker/manywheel/build.sh +++ b/.ci/docker/manywheel/build.sh @@ -95,7 +95,7 @@ case ${image} in MANY_LINUX_VERSION="2_28" DEVTOOLSET_VERSION="13" GPU_IMAGE=rocm/dev-almalinux-8:${GPU_ARCH_VERSION}-complete - PYTORCH_ROCM_ARCH="gfx900;gfx906;gfx908;gfx90a;gfx942;gfx1030;gfx1100;gfx1101;gfx1102;gfx1103;gfx1200;gfx1201;gfx950;gfx1150;gfx1151;gfx1250" + PYTORCH_ROCM_ARCH="gfx900;gfx906;gfx908;gfx90a;gfx942;gfx1030;gfx1100;gfx1101;gfx1102;gfx1103;gfx1200;gfx1201;gfx950;gfx1150;gfx1151;gfx1250;gfx1250-strict" DOCKER_GPU_BUILD_ARG="--build-arg ROCM_VERSION=${GPU_ARCH_VERSION} --build-arg PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} --build-arg DEVTOOLSET_VERSION=${DEVTOOLSET_VERSION}" ;; manylinux2_28-builder:xpu) diff --git a/aten/src/ATen/CMakeLists.txt b/aten/src/ATen/CMakeLists.txt index ef5a2291740ae..764d22435f598 100644 --- a/aten/src/ATen/CMakeLists.txt +++ b/aten/src/ATen/CMakeLists.txt @@ -667,14 +667,14 @@ if(USE_ROCM) exclude(ATen_HIP_SRCS "${ATen_HIP_SRCS}" ${native_hip_bgemm} ${native_hip_ck}) else() - # composable_kernel has no gfx1250 support yet, so its CK GEMM kernels fail to - # compile for that arch. Build CK GEMM as a separate library with gfx1250 - # filtered out of HIP_ARCHITECTURES so every other requested arch still gets - # CK GEMM. Remove this filtering once the CK submodule supports gfx1250. + # composable_kernel has no gfx1250/gfx1250-strict support yet, so its CK GEMM + # kernels fail to compile for those archs. Build CK GEMM as a separate library + # with them filtered out of HIP_ARCHITECTURES so every other requested arch + # still gets CK GEMM. Remove this filtering once the CK submodule supports them. exclude(ATen_HIP_SRCS "${ATen_HIP_SRCS}" ${native_hip_bgemm} ${native_hip_ck}) set(_ck_gemm_hip_arches ${PYTORCH_ROCM_ARCH}) - list(REMOVE_ITEM _ck_gemm_hip_arches "gfx1250") + list(REMOVE_ITEM _ck_gemm_hip_arches "gfx1250" "gfx1250-strict") set_source_files_properties(${native_hip_bgemm} ${native_hip_ck} PROPERTIES LANGUAGE HIP) add_library(ck_gemm STATIC ${native_hip_bgemm} ${native_hip_ck}) set_target_properties(ck_gemm PROPERTIES diff --git a/aten/src/ATen/cuda/CUDABlas.cpp b/aten/src/ATen/cuda/CUDABlas.cpp index 5befe801dc90c..ea4c7b642950f 100644 --- a/aten/src/ATen/cuda/CUDABlas.cpp +++ b/aten/src/ATen/cuda/CUDABlas.cpp @@ -2104,6 +2104,7 @@ void scaled_gemm( std::vector mx_archs{"gfx950"}; #if ROCM_VERSION >= 71400 mx_archs.push_back("gfx1250"); + mx_archs.push_back("gfx1250-strict"); #endif if (at::detail::getCUDAHooks().isGPUArch(mx_archs)) { // TODO: add constraints based on hipblaslt internals diff --git a/aten/src/ATen/cuda/CublasHandlePool.cpp b/aten/src/ATen/cuda/CublasHandlePool.cpp index ab64b28305f16..fe28271e7e67e 100644 --- a/aten/src/ATen/cuda/CublasHandlePool.cpp +++ b/aten/src/ATen/cuda/CublasHandlePool.cpp @@ -172,8 +172,8 @@ size_t parseChosenWorkspaceSize() { // for extra convenience val = c10::utils::get_env("ROCBLAS_WORKSPACE_CONFIG"); } - /* 32MiB default, 128MiB for gfx942/gfx950/gfx1250 */ - const bool gfx942_950_1250 = at::detail::getCUDAHooks().isGPUArch({"gfx942", "gfx950", "gfx1250"}); + /* 32MiB default, 128MiB for gfx942/gfx950/gfx1250/gfx1250-strict */ + const bool gfx942_950_1250 = at::detail::getCUDAHooks().isGPUArch({"gfx942", "gfx950", "gfx1250", "gfx1250-strict"}); const size_t default_size = gfx942_950_1250 ? 1024 * 128 * 1024 : 1024 * 32 * 1024; #else /* :4096:2:16:8 default, 32MiB for Hopper and Blackwell */ diff --git a/aten/src/ATen/cuda/detail/CUDAHooks.cpp b/aten/src/ATen/cuda/detail/CUDAHooks.cpp index 47e83d03e6b36..1f7de03d5bf99 100644 --- a/aten/src/ATen/cuda/detail/CUDAHooks.cpp +++ b/aten/src/ATen/cuda/detail/CUDAHooks.cpp @@ -560,7 +560,7 @@ const std::vector& CUDAHooks::getHipblasltPreferredArchs() const { "gfx1100", "gfx1101", "gfx1151", #endif #if ROCM_VERSION >= 71400 - "gfx1250", + "gfx1250", "gfx1250-strict", #endif }; return archs; @@ -576,7 +576,7 @@ const std::vector& CUDAHooks::getHipblasltSupportedArchs() const { "gfx950", "gfx1150", "gfx1151", #endif #if ROCM_VERSION >= 71400 - "gfx1250" + "gfx1250", "gfx1250-strict" #endif }; return archs; diff --git a/aten/src/ATen/native/cuda/GroupedBlas.cpp b/aten/src/ATen/native/cuda/GroupedBlas.cpp index 1a9b242c0cdcd..ff507e09c63c6 100644 --- a/aten/src/ATen/native/cuda/GroupedBlas.cpp +++ b/aten/src/ATen/native/cuda/GroupedBlas.cpp @@ -80,7 +80,7 @@ bool _scaled_mm_allowed_device(bool sm90_only=false, bool sm100_only=false) { "gfx950", #endif #if ROCM_VERSION >= 71400 - "gfx1250", + "gfx1250", "gfx1250-strict", #endif }; return at::detail::getCUDAHooks().isGPUArch(archs); diff --git a/aten/src/ATen/native/cuda/KernelUtils.cuh b/aten/src/ATen/native/cuda/KernelUtils.cuh index 0c899f5c6318c..fb0f5229d6b36 100644 --- a/aten/src/ATen/native/cuda/KernelUtils.cuh +++ b/aten/src/ATen/native/cuda/KernelUtils.cuh @@ -257,7 +257,7 @@ __device__ inline void cmtdStore(void* address, T value) { if constexpr (wait_for_commit) { __atomic_signal_fence(__ATOMIC_SEQ_CST); -#ifdef __gfx1250__ +#if defined(__gfx1250__) || defined(__gfx1250_strict__) asm volatile("s_wait_loadcnt(0)" ::: "memory"); #else asm volatile("s_waitcnt vmcnt(0)" ::: "memory"); diff --git a/aten/src/ATen/native/cuda/MemoryAccess.cuh b/aten/src/ATen/native/cuda/MemoryAccess.cuh index 66f871626e739..e83faf313754b 100644 --- a/aten/src/ATen/native/cuda/MemoryAccess.cuh +++ b/aten/src/ATen/native/cuda/MemoryAccess.cuh @@ -193,10 +193,10 @@ __device__ aligned_vector load_vector(const scalar_t *base_p auto *from = reinterpret_cast(base_ptr); #if defined(USE_ROCM) // NOTE: This nontemporal vectorized load is tuned for gfx942 (CDNA3, Wave64). - // Do NOT extend to gfx1250 (GFX12.5, Wave32) without first validating - // intrinsic codegen and re-benchmarking — the register layout and wave - // width differ. gfx1250 falls through to the generic path below, which - // is correctness-safe. + // Do NOT extend to gfx1250/gfx1250-strict (GFX12.5, Wave32) without first + // validating intrinsic codegen and re-benchmarking -- the register layout + // and wave width differ. Those archs fall through to the generic path below, + // which is correctness-safe. if(__builtin_amdgcn_processor_is("gfx942")) { using longx2 = __attribute__((__vector_size__(4*sizeof(int)))) int; if constexpr (sizeof(vec_t) == sizeof(int)) { diff --git a/aten/src/ATen/native/cuda/ScaledBlas.cpp b/aten/src/ATen/native/cuda/ScaledBlas.cpp index 87c96c56ffb40..aebe5fef15e2f 100644 --- a/aten/src/ATen/native/cuda/ScaledBlas.cpp +++ b/aten/src/ATen/native/cuda/ScaledBlas.cpp @@ -81,7 +81,7 @@ bool _scaled_mm_allowed_device(bool sm90_only=false, bool sm100_only=false) { "gfx950", #endif #if ROCM_VERSION >= 71400 - "gfx1250", + "gfx1250", "gfx1250-strict", #endif }; return at::detail::getCUDAHooks().isGPUArch(archs); @@ -106,6 +106,7 @@ static void check_blockwise_e8m0fnu_arch_supported() { std::vector mx_archs{"gfx950"}; #if ROCM_VERSION >= 71400 mx_archs.push_back("gfx1250"); + mx_archs.push_back("gfx1250-strict"); #endif TORCH_CHECK_NOT_IMPLEMENTED( at::detail::getCUDAHooks().isGPUArch(mx_archs), diff --git a/aten/src/ATen/native/cuda/int4mm.cu b/aten/src/ATen/native/cuda/int4mm.cu index d5d5b63f4bc1b..1c84ce3c31bd6 100644 --- a/aten/src/ATen/native/cuda/int4mm.cu +++ b/aten/src/ATen/native/cuda/int4mm.cu @@ -145,7 +145,7 @@ static bool isCDNA2orLater(int index) { } static bool isCDNA5orLater(int index) { - return at::detail::getCUDAHooks().isGPUArch({"gfx1250"}, index); + return at::detail::getCUDAHooks().isGPUArch({"gfx1250", "gfx1250-strict"}, index); } #else @@ -1111,8 +1111,8 @@ at::Tensor _weight_int4pack_mm_cuda( #if defined(USE_ROCM) if (isCDNA5orLater(A.device().index())) { TORCH_CHECK(false, - "_weight_int4pack_mm_cuda is not yet supported on gfx1250. " - "A WMMA-based implementation is required for gfx1250."); + "_weight_int4pack_mm_cuda is not yet supported on gfx1250/gfx1250-strict. " + "A WMMA-based implementation is required for gfx1250/gfx1250-strict."); } if (!isCDNA2orLater(A.device().index())) { TORCH_CHECK(false, "_weight_int4pack_mm_cuda is only supported on AMD gpu arch greater than or equal to CDNA2"); @@ -1310,8 +1310,8 @@ at::Tensor _convert_weight_to_int4pack_cuda( #if defined(USE_ROCM) if (isCDNA5orLater(in.device().index())) { TORCH_CHECK(false, - "_convert_weight_to_int4pack_cuda is not yet supported on gfx1250. " - "A WMMA-based implementation is required for gfx1250."); + "_convert_weight_to_int4pack_cuda is not yet supported on gfx1250/gfx1250-strict. " + "A WMMA-based implementation is required for gfx1250/gfx1250-strict."); } if (!isCDNA2orLater(in.device().index())) { TORCH_CHECK(false, "_convert_weight_to_int4pack_cuda is only supported on AMD gpu arch greater than or equal to CDNA2"); diff --git a/aten/src/ATen/native/cuda/reduction_template.cuh b/aten/src/ATen/native/cuda/reduction_template.cuh index 3be0cad8b67b1..04767e286564f 100644 --- a/aten/src/ATen/native/cuda/reduction_template.cuh +++ b/aten/src/ATen/native/cuda/reduction_template.cuh @@ -604,7 +604,7 @@ struct ReduceJitOp { if constexpr (wait_for_commit) { __atomic_signal_fence(__ATOMIC_SEQ_CST); - #ifdef __gfx1250__ + #if defined(__gfx1250__) || defined(__gfx1250_strict__) asm volatile("s_wait_loadcnt(0)" ::: "memory"); #else asm volatile("s_waitcnt vmcnt(0)" ::: "memory"); diff --git a/aten/src/ATen/native/sparse/cuda/cuSPARSELtOps.cpp b/aten/src/ATen/native/sparse/cuda/cuSPARSELtOps.cpp index f0f4172ccf6f4..69d0056d7dcef 100644 --- a/aten/src/ATen/native/sparse/cuda/cuSPARSELtOps.cpp +++ b/aten/src/ATen/native/sparse/cuda/cuSPARSELtOps.cpp @@ -27,7 +27,7 @@ static bool g_hipSparseLtSupported = false; static const std::vector& hipSparseLtSupportedArchs() { #if ROCM_VERSION >= 71400 - static const std::vector archs = {"gfx950", "gfx942", "gfx1250"}; + static const std::vector archs = {"gfx950", "gfx942", "gfx1250", "gfx1250-strict"}; #elif ROCM_VERSION >= 71200 static const std::vector archs = {"gfx950", "gfx942"}; #else diff --git a/cmake/Dependencies.cmake b/cmake/Dependencies.cmake index fc8c604adba57..710592c1cdc45 100644 --- a/cmake/Dependencies.cmake +++ b/cmake/Dependencies.cmake @@ -1040,18 +1040,19 @@ if(USE_ROCM) if(HIPBLASLT_VEC_EXT) list(APPEND HIP_CXX_FLAGS -DHIPBLASLT_VEC_EXT) endif() - # composable_kernel has no gfx1250 support, so aten/src/ATen/CMakeLists.txt - # filters gfx1250 out of the ck_gemm target's HIP_ARCHITECTURES. When gfx1250 - # is the only requested arch that list is empty and CMake fails at generate - # time with 'HIP_ARCHITECTURES is empty for target "ck_gemm"', so turn CK GEMM - # off entirely here -- before -DUSE_ROCM_CK_GEMM is added below, since the - # non-CK TUs guarded by that define (HIPBlas.cpp, HIPHooks.cpp, GroupedBlas.cpp) - # would otherwise reference symbols the unbuilt ck_gemm target never provides. + # composable_kernel has no gfx1250/gfx1250-strict support, so + # aten/src/ATen/CMakeLists.txt filters those archs out of the ck_gemm + # target's HIP_ARCHITECTURES. When they are the only requested archs that + # list is empty and CMake fails at generate time with 'HIP_ARCHITECTURES is + # empty for target "ck_gemm"', so turn CK GEMM off entirely here -- before + # -DUSE_ROCM_CK_GEMM is added below, since the non-CK TUs guarded by that + # define (HIPBlas.cpp, HIPHooks.cpp, GroupedBlas.cpp) would otherwise + # reference symbols the unbuilt ck_gemm target never provides. # Mirrors the USE_ROCM_CK_SDPA / USE_MSLK auto-disable pattern in - # aten/src/ATen/CMakeLists.txt. Remove once the CK submodule supports gfx1250. + # aten/src/ATen/CMakeLists.txt. Remove once the CK submodule supports these archs. if(USE_ROCM_CK_GEMM) set(_ck_gemm_supported_arches ${PYTORCH_ROCM_ARCH}) - list(REMOVE_ITEM _ck_gemm_supported_arches "gfx1250") + list(REMOVE_ITEM _ck_gemm_supported_arches "gfx1250" "gfx1250-strict") if(_ck_gemm_supported_arches STREQUAL "") message(STATUS "USE_ROCM_CK_GEMM disabled: PYTORCH_ROCM_ARCH (${PYTORCH_ROCM_ARCH}) has no arch supported by composable_kernel") caffe2_update_option(USE_ROCM_CK_GEMM OFF) diff --git a/cmake/External/aotriton.cmake b/cmake/External/aotriton.cmake index 4e2a163966cce..a1949b1e992ca 100644 --- a/cmake/External/aotriton.cmake +++ b/cmake/External/aotriton.cmake @@ -48,6 +48,8 @@ if(NOT __AOTRITON_INCLUDED) "amd-gfx120x" "amd-gfx1250" ) + # gfx1250-strict is not a separate AOTriton image; the amd-gfx1250 regex + # also matches gfx1250-strict in PYTORCH_ROCM_ARCH. set(__AOTRITON_IMAGE_SHA256_LIST "a3d1a6868ce290ba8118618207093e785252eff4e18a64f495752cb5a03ffed6" # amd-gfx90a "ccdbc7e3d96839be4895ee004f21531cc55d590c9018937b9e314bba363b3927" # amd-gfx942 diff --git a/test/inductor/test_aot_inductor.py b/test/inductor/test_aot_inductor.py index 6f2d9c943a0ac..257ca937106c1 100644 --- a/test/inductor/test_aot_inductor.py +++ b/test/inductor/test_aot_inductor.py @@ -7702,7 +7702,7 @@ def test__weight_int4pack_mm(self, m, n, q_group, num_groups): if TEST_WITH_ROCM: if CDNA5OrLater(): - self.skipTest("int4 mm not yet implemented for gfx1250 (needs WMMA)") + self.skipTest("int4 mm not yet implemented for gfx1250/gfx1250-strict (needs WMMA)") if not CDNA2OrLater(): self.skipTest("_int4_mm is supported only for CDNA2 or later") @@ -7741,7 +7741,7 @@ def test__weight_int4pack_mm_with_scales_and_zeros(self, m, n, q_group, num_grou if TEST_WITH_ROCM: if CDNA5OrLater(): - self.skipTest("int4 mm not yet implemented for gfx1250 (needs WMMA)") + self.skipTest("int4 mm not yet implemented for gfx1250/gfx1250-strict (needs WMMA)") if not CDNA2OrLater(): self.skipTest("_int4_mm is supported only for CDNA2 or later") diff --git a/test/test_cuda.py b/test/test_cuda.py index e4d7182370daa..41f5acc575b69 100644 --- a/test/test_cuda.py +++ b/test/test_cuda.py @@ -841,7 +841,7 @@ def _check_default(): if ROCM_VERSION >= (7, 13): archs.extend(["gfx1100", "gfx1101", "gfx1151"]) if ROCM_VERSION >= (7, 14): - archs.extend(["gfx1250"]) + archs.extend(["gfx1250", "gfx1250-strict"]) 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: @@ -931,7 +931,7 @@ def test_cublas_workspace_explicit_allocation(self): gcn_arch = str( torch.cuda.get_device_properties(0).gcnArchName.split(":", 1)[0] ) - if gcn_arch in ["gfx942", "gfx950", "gfx1250"]: + if gcn_arch in ["gfx942", "gfx950", "gfx1250", "gfx1250-strict"]: default_workspace_size = 1024 * 128 * 1024 # :1024:128 else: default_workspace_size = ( @@ -9729,7 +9729,7 @@ def test_compile_kernel_large_shared_memory(self): # Test error handling with more than supported shared memory size if torch.version.hip: gcn_arch = get_device_properties().gcnArchName.split(":", 1)[0] - if gcn_arch == "gfx1250": + if gcn_arch in ("gfx1250", "gfx1250-strict"): max_smem = 320 * 1024 elif gcn_arch == "gfx950": max_smem = 160 * 1024 diff --git a/test/test_linalg.py b/test/test_linalg.py index 4546c9a8065bf..5b4d79964f859 100644 --- a/test/test_linalg.py +++ b/test/test_linalg.py @@ -73,7 +73,7 @@ def blaslt_supported_device(): if ROCM_VERSION >= (6, 5): archs.append('gfx95') if ROCM_VERSION >= (7, 14): - archs.append('gfx1250') + archs.extend(['gfx1250', 'gfx1250-strict']) for arch in archs: if arch in torch.cuda.get_device_properties(0).gcnArchName: return True @@ -6611,7 +6611,7 @@ def test__int4_mm(self, device, m, k, n): self.skipTest("requires SM80 or later") if TEST_WITH_ROCM and self.device_type == 'cuda' and CDNA5OrLater(): - self.skipTest("int4 mm not yet implemented for gfx1250 (needs WMMA)") + self.skipTest("int4 mm not yet implemented for gfx1250/gfx1250-strict (needs WMMA)") if TEST_WITH_ROCM and self.device_type == 'cuda' and not CDNA2OrLater(): self.skipTest("_convert_weight_to_int4pack_cuda is supported only for CDNA2 or later") @@ -6683,7 +6683,7 @@ def test_compile_int4_mm(self, device, m, k, n): self.skipTest("requires SM80 or later") if TEST_WITH_ROCM and self.device_type == 'cuda' and CDNA5OrLater(): - self.skipTest("int4 mm not yet implemented for gfx1250 (needs WMMA)") + self.skipTest("int4 mm not yet implemented for gfx1250/gfx1250-strict (needs WMMA)") if TEST_WITH_ROCM and self.device_type == 'cuda' and not CDNA2OrLater(): self.skipTest("_convert_weight_to_int4pack_cuda supported only for CDNA2 or later") diff --git a/torch/cuda/_utils.py b/torch/cuda/_utils.py index bede63233f38e..30b47b3349719 100644 --- a/torch/cuda/_utils.py +++ b/torch/cuda/_utils.py @@ -476,11 +476,11 @@ def set_shared_memory_config(self, shared_mem_bytes: int) -> None: # HIP doesn't have shared_memory_per_block_optin in device properties, so we hard-code it here if torch.version.hip: # navi, CDNA1-CDNA3 allows a max of 64KB shared memory, - # CDNA4 (gfx950) 160KB, and CDNA5 (gfx1250) 320KB. + # CDNA4 (gfx950) 160KB, and CDNA5 (gfx1250/gfx1250-strict) 320KB. gcn_arch = device_props.gcnArchName.split(":", 1)[0] if gcn_arch == "gfx950": max_shared_mem = 160 * 1024 - elif gcn_arch == "gfx1250": + elif gcn_arch in ("gfx1250", "gfx1250-strict"): max_shared_mem = 320 * 1024 else: max_shared_mem = 65536 diff --git a/torch/testing/_internal/common_cuda.py b/torch/testing/_internal/common_cuda.py index 9818b3c4ab5ec..9b51284174876 100644 --- a/torch/testing/_internal/common_cuda.py +++ b/torch/testing/_internal/common_cuda.py @@ -70,7 +70,7 @@ def evaluate_gfx_arch_within(arch_list): _CDNA2_ARCHS = ["gfx90a"] _CDNA3_ARCHS = ["gfx942", "gfx950"] # GFX1250 (CDNA 5) -_CDNA5_ARCHS = ["gfx1250"] +_CDNA5_ARCHS = ["gfx1250", "gfx1250-strict"] def CDNA5OrLater(): return evaluate_gfx_arch_within(_CDNA5_ARCHS) @@ -83,13 +83,13 @@ def CDNA2OrLater(): # Archs that take the opportunistic_fastAtomicAdd path (packed 2x16 atomics + DPP # lane coalescing) in ScatterGatherKernel.cu. Keep in sync with that kernel's arch -# gate; this is intentionally not CDNA3OrLater (gfx1250 uses plain fastAtomicAdd). +# gate; this is intentionally not CDNA3OrLater (gfx1250/gfx1250-strict use plain fastAtomicAdd). def gfx_arch_supports_opportunistic_fastatomics(): return evaluate_gfx_arch_within(["gfx942", "gfx950"]) def evaluate_platform_supports_flash_attention(): if TEST_WITH_ROCM: - arch_list = ["gfx90a", "gfx942", "gfx1100", "gfx1201", "gfx950", "gfx1250"] + arch_list = ["gfx90a", "gfx942", "gfx1100", "gfx1201", "gfx950", "gfx1250", "gfx1250-strict"] if os.environ.get("TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL", "0") != "0": arch_list += ["gfx1101", "gfx1102", "gfx1150", "gfx1151", "gfx1200"] return evaluate_gfx_arch_within(arch_list) @@ -107,7 +107,7 @@ def evaluate_platform_supports_ck_sdpa(): def evaluate_platform_supports_efficient_attention(): if TEST_WITH_ROCM: - arch_list = ["gfx90a", "gfx942", "gfx1100", "gfx1201", "gfx950", "gfx1250"] + arch_list = ["gfx90a", "gfx942", "gfx1100", "gfx1201", "gfx950", "gfx1250", "gfx1250-strict"] if os.environ.get("TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL", "0") != "0": arch_list += ["gfx1101", "gfx1102", "gfx1150", "gfx1151", "gfx1200"] return evaluate_gfx_arch_within(arch_list) @@ -194,7 +194,7 @@ def evaluate_platform_supports_fp8(): if ROCM_VERSION >= (6, 5): archs.append('gfx95') if ROCM_VERSION >= (7, 14): - archs.append('gfx1250') + archs.extend(['gfx1250', 'gfx1250-strict']) for arch in archs: if arch in torch.cuda.get_device_properties(0).gcnArchName: return True @@ -211,8 +211,8 @@ def evaluate_platform_supports_fp8_grouped_gemm(): if torch.version.hip: if "USE_MSLK" not in torch.__config__.show(): return False - # gfx1250 omitted: MSLK only builds gfx942/gfx950 kernels (see the arch - # filter in aten/src/ATen/CMakeLists.txt). Add gfx1250 here once MSLK does. + # gfx1250/gfx1250-strict omitted: MSLK only builds gfx942/gfx950 kernels + # (see the arch filter in aten/src/ATen/CMakeLists.txt). Add them here once MSLK does. archs = ['gfx942', 'gfx950'] for arch in archs: if arch in torch.cuda.get_device_properties(0).gcnArchName: @@ -226,7 +226,9 @@ def evaluate_platform_supports_mx_gemm(): if torch.version.hip: if ROCM_VERSION >= (7, 0): gcn_name = torch.cuda.get_device_properties(0).gcnArchName - return 'gfx950' in gcn_name or ('gfx1250' in gcn_name and ROCM_VERSION >= (7, 14)) + return 'gfx950' in gcn_name or ( + ('gfx1250' in gcn_name or 'gfx1250-strict' in gcn_name) and ROCM_VERSION >= (7, 14) + ) else: return SM100OrLater return False