Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
pragupta merged 2 commits into from
Jul 9, 2026
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ac47375
[ROCm] Add initial support for gfx1250 (#188597)
rraminen Jul 9, 2026
6722f8e
Fix conflicts
rraminen Jul 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 24 additions & 2 deletions aten/src/ATen/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
6 changes: 5 additions & 1 deletion aten/src/ATen/cuda/CUDABlas.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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. "
Expand Down
6 changes: 3 additions & 3 deletions aten/src/ATen/cuda/CublasHandlePool.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down
7 changes: 5 additions & 2 deletions aten/src/ATen/cuda/detail/CUDAHooks.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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
};
Expand Down
51 changes: 29 additions & 22 deletions aten/src/ATen/native/cuda/MemoryAccess.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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];
Expand Down
30 changes: 22 additions & 8 deletions aten/src/ATen/native/cuda/ScaledBlas.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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>
Expand Down Expand Up @@ -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);
Expand All @@ -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

/*
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
14 changes: 14 additions & 0 deletions aten/src/ATen/native/cuda/int4mm.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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");
}
Expand Down Expand Up @@ -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");
}
Expand Down
61 changes: 39 additions & 22 deletions aten/src/ATen/native/sparse/cuda/cuSPARSELtOps.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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

Expand Down
4 changes: 4 additions & 0 deletions caffe2/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
5 changes: 5 additions & 0 deletions test/inductor/test_aot_inductor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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")

Expand Down Expand Up @@ -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")

Expand Down
12 changes: 8 additions & 4 deletions test/test_cuda.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -774,7 +774,7 @@ def test_cublas_workspace_explicit_allocation(self):
gcn_arch = str(
torch.cuda.get_device_properties(0).gcnArchName.split(":", 1)[0]
)
if "gfx94" in gcn_arch or "gfx95" in gcn_arch:
if gcn_arch in ["gfx942", "gfx950", "gfx1250"]:
default_workspace_size = 1024 * 128 * 1024 # :1024:128
else:
default_workspace_size = (
Expand Down Expand Up @@ -8456,9 +8456,13 @@ def test_compile_kernel_large_shared_memory(self):

# Test error handling with more than supported shared memory size
if torch.version.hip:
max_smem = (
65536 if get_device_properties().gcnArchName != "gfx950" else 160 * 1024
)
gcn_arch = get_device_properties().gcnArchName.split(":", 1)[0]
if gcn_arch == "gfx1250":
max_smem = 320 * 1024
elif gcn_arch == "gfx950":
max_smem = 160 * 1024
else:
max_smem = 65536
else:
max_smem = get_device_properties().shared_memory_per_block_optin
excessive_shared_mem = max_smem * 2
Expand Down
Loading