[ROCm] Pair gfx1250 support with gfx1250-strict

.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}"
     ;;
   *)

.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)

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

aten/src/ATen/cuda/CUDABlas.cpp

@@ -2104,6 +2104,7 @@ void scaled_gemm(
             std::vector<std::string> 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

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 */

aten/src/ATen/cuda/detail/CUDAHooks.cpp

@@ -560,7 +560,7 @@ const std::vector<std::string>& 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<std::string>& CUDAHooks::getHipblasltSupportedArchs() const {
     "gfx950", "gfx1150", "gfx1151",
 #endif
 #if ROCM_VERSION >= 71400
-    "gfx1250"
+    "gfx1250", "gfx1250-strict"
 #endif
   };
   return archs;

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);

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");

aten/src/ATen/native/cuda/MemoryAccess.cuh

@@ -193,10 +193,10 @@ __device__ aligned_vector<scalar_t, vec_size> load_vector(const scalar_t *base_p
   auto *from = reinterpret_cast<const vec_t *>(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)) {

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<std::string> 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),

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");

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");

aten/src/ATen/native/sparse/cuda/cuSPARSELtOps.cpp

@@ -27,7 +27,7 @@ static bool g_hipSparseLtSupported = false;
 
 static const std::vector<std::string>& hipSparseLtSupportedArchs() {
 #if ROCM_VERSION >= 71400
-  static const std::vector<std::string> archs = {"gfx950", "gfx942", "gfx1250"};
+  static const std::vector<std::string> archs = {"gfx950", "gfx942", "gfx1250", "gfx1250-strict"};
 #elif ROCM_VERSION >= 71200
   static const std::vector<std::string> archs = {"gfx950", "gfx942"};
 #else

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)

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

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")
 

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

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")

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

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