[release/2.12] [ROCm] Add initial support for gfx1250 (pytorch#188597) (#3421)

## Summary

Validation theRock Run:
https://github.com/ROCm/TheRock/actions/runs/29045360944

Cherry-pick of upstream
[pytorch#188597](pytorch#188597)
onto `release/2.12` to add initial ROCm support for **gfx1250 (CDNA5)**.

Co-authored-by: @glen-amd

### What this PR enables
- Gates gfx1250-specific behavior on ROCm 7.14+ (the support floor),
applied consistently across CUDABlas.cpp, ScaledBlas.cpp, CUDAHooks.cpp,
and the test helpers.
- Adds gfx1250 to the hipBLASLt preferred/supported arch lists and the
hipSparseLt support check (ROCm 7.14+).
- Extends the scaled GEMM / MX-format paths (block-wise Float8_e8m0fnu,
mxfp8, mxfp4) to gfx1250.
- Sets the gfx1250 shared-memory limit to 320 KB (vs 160 KB on gfx950),
matching on the suffix-stripped arch name.
- Adds CDNA2/CDNA3/CDNA5 "or later" arch helpers and updates FP8 / MX
GEMM test gating in common_cuda.py.
- Builds Composable Kernel (CK) GEMM as a separate ck_gemm library with
gfx1250 filtered out of HIP_ARCHITECTURES, so a multi-arch build that
includes gfx1250 keeps CK GEMM for the other archs instead of failing to
compile for gfx1250 (composable_kernel has no gfx1250 support yet). The
filter is removed once CK supports gfx1250.
- Extends the gfx942 nontemporal vectorized-load path in
MemoryAccess.cuh to gfx1250.
- Refactors duplicated arch-check logic in CUDABlas.cpp, ScaledBlas.cpp,
and cuSPARSELtOps.cpp per review feedback.
- Writes the CublasHandlePool workspace archs in full
(gfx942/gfx950/gfx1250).

Enabled via the main merge:
- Flash attention and memory-efficient attention on gfx1250 via AOTriton
0.12.1b (from pytorch#188242), which this branch merges in.

What is not enabled yet:
- CK SDPA on gfx1250: composable_kernel has no gfx1250 support; the CK
SDPA target auto-filters to gfx942/gfx950.
- CK GEMM on gfx1250: filtered out (see above) until composable_kernel
supports it.
- FP8 grouped GEMM on gfx1250: MSLK builds only gfx942/gfx950.
- int4 mm on gfx1250: the tinygemm MFMA kernel needs a WMMA port; the
ops return a clear "not supported yet" error and the int4 unit tests
skip gfx1250.

CI note:
- gfx1250 is not added to the ROCm 7.2 (rocm-n) docker image, whose HIP
compiler cannot target gfx1250. gfx1250 CI belongs on the nightly ROCm
image and is added in a follow-up.

Co-authored-by: @glen-amd

## References

- Upstream PR: pytorch#188597
- Related AOTriton bump: pytorch#188242

---------

Co-authored-by: Prachi Gupta <prachi.gupta@amd.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: Jeff Daily <jeff.daily@amd.com>

Author: 4 people

aten/src/ATen/CMakeLists.txt

@@ -682,11 +682,33 @@ if(USE_ROCM)
     ${native_quantized_hip_hip}
     ${native_transformers_hip_hip} ${native_transformers_src_hip_hip}
   )
+  file(GLOB native_hip_bgemm CONFIGURE_DEPENDS "native/hip/bgemm_kernels/*.hip")
+  file(GLOB native_hip_ck CONFIGURE_DEPENDS "native/hip/ck*.hip")
   if(NOT USE_ROCM_CK_GEMM)
-    file(GLOB native_hip_bgemm "native/hip/bgemm_kernels/*.hip")
-    file(GLOB native_hip_ck "native/hip/ck*.hip")
     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.
+    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")
+    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
+      POSITION_INDEPENDENT_CODE ON
+      HIP_ARCHITECTURES "${_ck_gemm_hip_arches}")
+    target_compile_options(ck_gemm PRIVATE ${HIP_CXX_FLAGS})
+    target_include_directories(ck_gemm PRIVATE
+      ${CMAKE_CURRENT_SOURCE_DIR}/hip
+      ${CMAKE_CURRENT_SOURCE_DIR}/../../../third_party/composable_kernel/include
+      ${CMAKE_CURRENT_SOURCE_DIR}/../../../third_party/composable_kernel/library/include
+      ${CMAKE_CURRENT_BINARY_DIR}/composable_kernel)
+    # Ensure torch_cpu (and its codegen) is built before ck_gemm to avoid race conditions
+    add_dependencies(ck_gemm torch_cpu)
   endif()
 
   # TODO: Codegen separate files for HIP and use those (s/cuda_generated_sources/hip_generated_sources)

aten/src/ATen/cuda/CUDABlas.cpp

@@ -2012,7 +2012,11 @@ void scaled_gemm(
   }
   else if (mat1_scale_dtype == kFloat8_e8m0fnu && mat2_scale_dtype == kFloat8_e8m0fnu) {
   #if ROCM_VERSION >= 70000
-            if (at::detail::getCUDAHooks().isGPUArch({"gfx950"})) {
+            std::vector<std::string> mx_archs{"gfx950"};
+  #if ROCM_VERSION >= 71400
+            mx_archs.push_back("gfx1250");
+  #endif
+            if (at::detail::getCUDAHooks().isGPUArch(mx_archs)) {
                 // TODO: add constraints based on hipblaslt internals
                 TORCH_CHECK((m % 16 == 0) && (n % 16 == 0) && (k % 128 == 0),
                            "M, N must be multiples of 16 and K should be multiple of 128 for MX format. "

aten/src/ATen/cuda/CublasHandlePool.cpp

@@ -165,9 +165,9 @@ size_t parseChosenWorkspaceSize() {
     // for extra convenience
     val = c10::utils::get_env("ROCBLAS_WORKSPACE_CONFIG");
   }
-  /* 32MiB default, 128MiB for gfx94x/gfx95x */
-  const bool gfx94_95 = at::detail::getCUDAHooks().isGPUArch({"gfx94", "gfx95"});
-  const size_t default_size = gfx94_95 ? 1024 * 128 * 1024 : 1024 * 32 * 1024;
+  /* 32MiB default, 128MiB for gfx942/gfx950/gfx1250 */
+  const bool gfx942_950_1250 = at::detail::getCUDAHooks().isGPUArch({"gfx942", "gfx950", "gfx1250"});
+  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 */
   cudaDeviceProp* properties = at::cuda::getCurrentDeviceProperties();

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

@@ -557,7 +557,10 @@ const std::vector<std::string>& CUDAHooks::getHipblasltPreferredArchs() const {
     "gfx950",
 #endif
 #if ROCM_VERSION >= 71300
-    "gfx1100", "gfx1101", "gfx1151"
+    "gfx1100", "gfx1101", "gfx1151",
+#endif
+#if ROCM_VERSION >= 71400
+    "gfx1250",
 #endif
   };
   return archs;
@@ -572,7 +575,7 @@ const std::vector<std::string>& CUDAHooks::getHipblasltSupportedArchs() const {
 #if ROCM_VERSION >= 70000
     "gfx950", "gfx1150", "gfx1151",
 #endif
-#if ROCM_VERSION >= 70200
+#if ROCM_VERSION >= 71400
     "gfx1250"
 #endif
   };

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

@@ -187,28 +187,35 @@ template <int vec_size, typename scalar_t>
 __device__ aligned_vector<scalar_t, vec_size> load_vector(const scalar_t *base_ptr, uint32_t offset) {
   using vec_t = aligned_vector<scalar_t, vec_size>;
   auto *from = reinterpret_cast<const vec_t *>(base_ptr);
-#if defined(USE_ROCM) && defined(__gfx942__)
-  using longx2 = __attribute__((__vector_size__(4*sizeof(int)))) int;
-  if constexpr (sizeof(vec_t) == sizeof(int)) {
-   union {
-     vec_t v;
-     int   i;
-   } tmpt = { .i = __builtin_nontemporal_load(reinterpret_cast<const int *>(&(from[offset]))) };
-   return tmpt.v;
-  }
-  else if constexpr (sizeof(vec_t) == sizeof(long)) {
-   union {
-     vec_t v;
-     long   i;
-   } tmpt = { .i = __builtin_nontemporal_load(reinterpret_cast<const long *>(&(from[offset]))) };
-   return tmpt.v;
-  }
-  else if constexpr (sizeof(vec_t) == sizeof(longx2)) {
-   union {
-     vec_t v;
-     longx2  i;
-   } tmpt = { .i = __builtin_nontemporal_load(reinterpret_cast<const longx2 *>(&(from[offset]))) };
-   return tmpt.v;
+#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.
+  if(__builtin_amdgcn_processor_is("gfx942")) {
+    using longx2 = __attribute__((__vector_size__(4*sizeof(int)))) int;
+    if constexpr (sizeof(vec_t) == sizeof(int)) {
+    union {
+      vec_t v;
+      int   i;
+    } tmpt = { .i = __builtin_nontemporal_load(reinterpret_cast<const int *>(&(from[offset]))) };
+    return tmpt.v;
+    }
+    else if constexpr (sizeof(vec_t) == sizeof(long)) {
+    union {
+      vec_t v;
+      long   i;
+    } tmpt = { .i = __builtin_nontemporal_load(reinterpret_cast<const long *>(&(from[offset]))) };
+    return tmpt.v;
+    }
+    else if constexpr (sizeof(vec_t) == sizeof(longx2)) {
+    union {
+      vec_t v;
+      longx2  i;
+    } tmpt = { .i = __builtin_nontemporal_load(reinterpret_cast<const longx2 *>(&(from[offset]))) };
+    return tmpt.v;
+    }
   }
 #endif
   return from[offset];

aten/src/ATen/native/cuda/ScaledBlas.cpp

@@ -17,6 +17,7 @@
 #include <ATen/cuda/tunable/TunableGemm.h>
 #include <ATen/native/Resize.h>
 #include <c10/util/MaybeOwned.h>
+#include <c10/util/StringUtil.h>
 #include <ATen/native/GroupedMMUtils.h>
 #include <ATen/native/cuda/RowwiseScaledMM.h>
 #include <ATen/native/cuda/ScaledGroupMM.h>
@@ -78,7 +79,10 @@ bool _scaled_mm_allowed_device(bool sm90_only=false, bool sm100_only=false) {
         "gfx1200", "gfx1201",
 #endif
 #if ROCM_VERSION >= 60500
-        "gfx950"
+        "gfx950",
+#endif
+#if ROCM_VERSION >= 71400
+        "gfx1250",
 #endif
     };
     return at::detail::getCUDAHooks().isGPUArch(archs);
@@ -97,6 +101,19 @@ bool _scaled_mm_allowed_device(bool sm90_only=false, bool sm100_only=false) {
 bool _scaled_mm_is_fnuz() {
     return at::detail::getCUDAHooks().isGPUArch({"gfx942"});
 }
+
+#if ROCM_VERSION >= 70000
+static void check_blockwise_e8m0fnu_arch_supported() {
+  std::vector<std::string> mx_archs{"gfx950"};
+#if ROCM_VERSION >= 71400
+  mx_archs.push_back("gfx1250");
+#endif
+  TORCH_CHECK_NOT_IMPLEMENTED(
+      at::detail::getCUDAHooks().isGPUArch(mx_archs),
+      "Block-wise scaling for Float8_e8m0fnu is only supported on ",
+      c10::Join(",", mx_archs));
+}
+#endif
 #endif
 
 /*
@@ -621,9 +638,8 @@ _scaled_mm_out_cuda(const Tensor& mat1, const Tensor& mat2,
   }
   else if (scaling_choice_a == ScalingType::BlockWise1x32 && scaling_choice_b == ScalingType::BlockWise1x32) {
 #ifdef USE_ROCM
-    #if ROCM_VERSION >= 70000
-    TORCH_CHECK_NOT_IMPLEMENTED(at::detail::getCUDAHooks().isGPUArch({"gfx950"}),
-                "Block-wise scaling for Float8_e8m0fnu is only supported on gfx950");
+#if ROCM_VERSION >= 70000
+    check_blockwise_e8m0fnu_arch_supported();
 
     int packed_factor = 1;
     if (mat1.scalar_type() == ScalarType::Float4_e2m1fn_x2) {
@@ -1064,8 +1080,7 @@ _scaled_mxfp8_mxfp8(
 
 #ifdef USE_ROCM
 #if ROCM_VERSION >= 70000
-  TORCH_CHECK_NOT_IMPLEMENTED(at::detail::getCUDAHooks().isGPUArch({"gfx950"}),
-              "Block-wise scaling for Float8_e8m0fnu is only supported on gfx950");
+  check_blockwise_e8m0fnu_arch_supported();
 
   TORCH_CHECK_VALUE(mat_a.size(0) % 32 == 0 && mat_a.size(1) % 32 == 0 &&
               mat_b.size(0) % 32 == 0 && mat_b.size(1) % 32 == 0,
@@ -1150,8 +1165,7 @@ _scaled_mxfp4_mxfp4(
   auto scaling_choice_b = ScalingType::BlockWise1x32;
 
 #if ROCM_VERSION >= 70000
-  TORCH_CHECK_NOT_IMPLEMENTED(at::detail::getCUDAHooks().isGPUArch({"gfx950"}),
-              "Block-wise scaling for Float8_e8m0fnu is only supported on gfx950");
+  check_blockwise_e8m0fnu_arch_supported();
 
   TORCH_CHECK_VALUE(mat_a.size(0) % 32 == 0 && mat_a.size(1) % 32 == 0 &&
               mat_b.size(0) % 32 == 0 && mat_b.size(1) % 32 == 0,

aten/src/ATen/native/cuda/int4mm.cu

@@ -146,6 +146,10 @@ static bool isCDNA2orLater(int index) {
     return at::detail::getCUDAHooks().isGPUArch({"gfx90a", "gfx942", "gfx950"}, index);
 }
 
+static bool isCDNA5orLater(int index) {
+  return at::detail::getCUDAHooks().isGPUArch({"gfx1250"}, index);
+}
+
 #else
 constexpr int32_t kWarpSize = 32;
 #endif
@@ -1098,6 +1102,11 @@ at::Tensor _weight_int4pack_mm_cuda(
       A.device() == B.device() && A.device() == qScaleAndZeros.device());
 
 #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.");
+  }
   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");
   }
@@ -1293,6 +1302,11 @@ at::Tensor _convert_weight_to_int4pack_cuda(
   TORCH_CHECK(innerKTiles == 2 || innerKTiles == 4 || innerKTiles == 8);
 
 #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.");
+  }
   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/sparse/cuda/cuSPARSELtOps.cpp

@@ -2,6 +2,9 @@
 #include <unordered_map>
 #include <mutex>
 #include <string_view>
+#include <unordered_map>
+#include <vector>
+#include <c10/util/StringUtil.h>
 #if AT_CUSPARSELT_ENABLED()
 
 namespace at::native {
@@ -22,37 +25,51 @@ thread_local bool handle_initialized = false;
 c10::once_flag g_hipSparseLtSupportInitFlag;
 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"};
+#elif ROCM_VERSION >= 71200
+  static const std::vector<std::string> archs = {"gfx950", "gfx942"};
+#else
+  static const std::vector<std::string> archs = {};
+#endif
+  return archs;
+}
+
 // Initialize the hipSparseLt support status once for the platform
 static void initHipSparseLtSupport() {
-    // Default to not supported
-    g_hipSparseLtSupported = false;
-
-    // Check only the first available device
-    try {
-        if (at::cuda::device_count() > 0) {
-            g_hipSparseLtSupported = at::detail::getCUDAHooks().isGPUArch({"gfx950", "gfx942"}, 0);
-        }
-    } catch (const std::exception&) {
-        // If an exception occurs during device property check, we assume hipSparseLt is not supported
-        // This could happen due to driver issues, device access problems, or other runtime errors
-        g_hipSparseLtSupported = false;
-        TORCH_WARN("Exception occurred while checking hipSparseLt support. Assuming not supported.");
+  // Default to not supported
+  g_hipSparseLtSupported = false;
+
+  // Check only the first available device
+  try {
+    if (at::cuda::device_count() > 0) {
+      g_hipSparseLtSupported = at::detail::getCUDAHooks().isGPUArch(
+          hipSparseLtSupportedArchs(), 0);
     }
+  } catch (const std::exception&) {
+    // If an exception occurs during device property check, we assume
+    // hipSparseLt is not supported This could happen due to driver issues,
+    // device access problems, or other runtime errors
+     g_hipSparseLtSupported = false;
+    TORCH_WARN(
+        "Exception occurred while checking hipSparseLt support. Assuming not supported.");
+  }
 }
 
 static bool isHipSparseLtSupported() {
     // Initialize support check only once
     c10::call_once(g_hipSparseLtSupportInitFlag, initHipSparseLtSupport);
 
-    // Return cached result (platform-wide)
-    if (!g_hipSparseLtSupported) {
-        TORCH_CHECK(
-            false,
-            "hipSparseLt not supported on this device, supported architectures: "
-            "gfx950, gfx942. "
-            "required ROCM version: 6.4.0 or later.");
-    }
-    return g_hipSparseLtSupported;
+  // Return cached result (platform-wide)
+  if (!g_hipSparseLtSupported) {
+    TORCH_CHECK(
+        false,
+        "hipSparseLt not supported on this device. Supported architectures: ",
+        c10::Join(", ", hipSparseLtSupportedArchs()),
+        ". hipSparseLt on ROCm requires ROCm 7.12 or newer.");
+  }
+  return g_hipSparseLtSupported;
 }
 #endif
 

caffe2/CMakeLists.txt

@@ -1866,6 +1866,10 @@ if(USE_ROCM)
     target_link_libraries(torch_hip PRIVATE ck_sdpa)
   endif()
 
+  if(TARGET ck_gemm)
+    target_link_libraries(torch_hip PRIVATE ck_gemm)
+  endif()
+
   if(USE_MSLK)
     if(USE_ROCM)
       target_link_libraries(torch_hip PRIVATE mslk)

test/inductor/test_aot_inductor.py

@@ -46,6 +46,7 @@
 from torch.testing._internal import common_utils
 from torch.testing._internal.common_cuda import (
     CDNA2OrLater,
+    CDNA5OrLater,
     PLATFORM_SUPPORTS_FLASH_ATTENTION,
     PLATFORM_SUPPORTS_FP8,
     PLATFORM_SUPPORTS_FP8_GROUPED_GEMM,
@@ -7174,6 +7175,8 @@ def test__weight_int4pack_mm(self, m, n, q_group, num_groups):
             raise unittest.SkipTest("requires GPU")
 
         if TEST_WITH_ROCM:
+            if CDNA5OrLater():
+                self.skipTest("int4 mm not yet implemented for gfx1250 (needs WMMA)")
             if not CDNA2OrLater():
                 self.skipTest("_int4_mm is supported only for CDNA2 or later")
 
@@ -7211,6 +7214,8 @@ def test__weight_int4pack_mm_with_scales_and_zeros(self, m, n, q_group, num_grou
             raise unittest.SkipTest("requires Intel GPU")
 
         if TEST_WITH_ROCM:
+            if CDNA5OrLater():
+                self.skipTest("int4 mm not yet implemented for gfx1250 (needs WMMA)")
             if not CDNA2OrLater():
                 self.skipTest("_int4_mm is supported only for CDNA2 or later")