From 632d652ba42ce30908bd807e51b2b54f1397cb51 Mon Sep 17 00:00:00 2001 From: rraminen_amdeng Date: Mon, 15 Jun 2026 19:03:15 +0000 Subject: [PATCH 01/11] Cherry-pick bdbcbea8dbf09fb95685d499cd6b1de1e04fe4b0 (exclude CK submodule + triton pins) --- .ci/docker/build.sh | 4 +++- .ci/pytorch/binary_populate_env.sh | 2 +- .github/actionlint.yaml | 4 ++++ CMakeLists.txt | 10 +++++++--- aten/src/ATen/CMakeLists.txt | 4 ++-- aten/src/ATen/Context.cpp | 2 +- aten/src/ATen/cuda/CUDABlas.cpp | 8 ++++++++ aten/src/ATen/cuda/CublasHandlePool.cpp | 4 ++-- aten/src/ATen/cuda/detail/CUDAHooks.cpp | 5 ++++- aten/src/ATen/native/cuda/CUDALoops.cuh | 3 ++- aten/src/ATen/native/cuda/GroupedBlas.cpp | 2 +- aten/src/ATen/native/cuda/KernelUtils.cuh | 3 ++- aten/src/ATen/native/cuda/MemoryAccess.cuh | 3 ++- aten/src/ATen/native/cuda/ScaledBlas.cpp | 15 +++++++++------ aten/src/ATen/native/cuda/int4mm.cu | 19 ++++++++++++++++++- .../ATen/native/sparse/cuda/cuSPARSELtOps.cpp | 4 ++-- .../native/transformers/cuda/sdp_utils.cpp | 1 + .../hip/flash_attn/ck/launch_kernel_pt.hpp | 2 +- c10/core/AllocatorConfig.h | 8 +++++++- cmake/External/aotriton.cmake | 5 +++++ test/test_cuda.py | 8 +++++--- test/test_linalg.py | 6 ++++++ torch/_inductor/config.py | 3 ++- torch/cuda/_utils.py | 13 ++++++++----- torch/testing/_internal/common_cuda.py | 15 +++++++++++---- torch/testing/_internal/common_distributed.py | 2 +- 26 files changed, 115 insertions(+), 40 deletions(-) diff --git a/.ci/docker/build.sh b/.ci/docker/build.sh index 1eca06471a110..b7b74c3a0794f 100755 --- a/.ci/docker/build.sh +++ b/.ci/docker/build.sh @@ -163,7 +163,9 @@ case "$tag" in ROCM_VERSION=7.2 TRITON=yes KATEX=yes - PYTORCH_ROCM_ARCH="gfx90a;gfx942;gfx950;gfx1100" + UCX_COMMIT=${_UCX_COMMIT} + UCC_COMMIT=${_UCC_COMMIT} + PYTORCH_ROCM_ARCH="gfx90a;gfx942;gfx950;gfx1100;gfx1250" if [[ $tag =~ "benchmarks" ]]; then INDUCTOR_BENCHMARKS=yes fi diff --git a/.ci/pytorch/binary_populate_env.sh b/.ci/pytorch/binary_populate_env.sh index 53914914c8c93..4912b444bcca9 100755 --- a/.ci/pytorch/binary_populate_env.sh +++ b/.ci/pytorch/binary_populate_env.sh @@ -147,7 +147,7 @@ export PYTORCH_EXTRA_INSTALL_REQUIREMENTS="${PYTORCH_EXTRA_INSTALL_REQUIREMENTS: # TODO: We don't need this anymore IIUC export TORCH_PACKAGE_NAME='torch' -export USE_FBGEMM=1 +export USE_FBGEMM=0 export PIP_UPLOAD_FOLDER="$PIP_UPLOAD_FOLDER" export DOCKER_IMAGE="$DOCKER_IMAGE" diff --git a/.github/actionlint.yaml b/.github/actionlint.yaml index b894d7214fced..81690f184271a 100644 --- a/.github/actionlint.yaml +++ b/.github/actionlint.yaml @@ -62,9 +62,13 @@ self-hosted-runner: # gfx942 runners - linux.rocm.gpu.gfx942.1 - linux.rocm.gpu.gfx942.4 + - linux.rocm.gfx942.docker-cache # gfx950 runners - linux.rocm.gpu.gfx950.1 - linux.rocm.gpu.gfx950.4 + # gfx1250 runners + - linux.rocm.gpu.gfx1250.1 + - linux.rocm.gpu.gfx1250.4 # Org wise AWS `mac2.metal` runners (2020 Mac mini hardware powered by Apple silicon M1 processors) - macos-m1-stable - macos-m1-14 diff --git a/CMakeLists.txt b/CMakeLists.txt index d5941243d66db..7530948ed4d2e 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -275,7 +275,7 @@ cmake_dependent_option(USE_CUSPARSELT "Use cuSPARSELt" ON "USE_CUDA" OFF) cmake_dependent_option(USE_CUDSS "Use cuDSS" ON "USE_CUDA" OFF) # USE_ROCM is guarded against in Dependencies.cmake because USE_ROCM is not properly defined here cmake_dependent_option(USE_CUFILE "Use cuFile" ON "USE_CUDA AND NOT WIN32" OFF) -option(USE_FBGEMM "Use FBGEMM (quantized 8-bit server operators)" ON) +option(USE_FBGEMM "Use FBGEMM (quantized 8-bit server operators)" OFF) option(USE_KINETO "Use Kineto profiling library" ON) option(USE_CUPTI_SO "Use CUPTI as a shared library" ON) option(USE_GFLAGS "Use GFLAGS" OFF) @@ -955,9 +955,13 @@ cmake_dependent_option( OFF) +# TODO: +# MSLK related parts are missing that already exists upstream. +# gfx1250 for MSLK needs to be involved as well. + IF(USE_ROCM AND ("gfx942" IN_LIST PYTORCH_ROCM_ARCH OR "gfx950" IN_LIST PYTORCH_ROCM_ARCH)) - message(WARNING "Setting USE_MSLK for gfx942/gfx950 to ON by default, doing ROCM build") - set(USE_MSLK_DEFAULT ON) + message(WARNING "Setting USE_FBGEMM_GENAI for gfx942/gfx950 to ON by default, doing ROCM build") + set(USE_FBGEMM_GENAI_DEFAULT ON) elseif(USE_CUDA AND "$ENV{TORCH_CUDA_ARCH_LIST}" MATCHES "10.0" AND CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.8 AND NOT WIN32) message(STATUS "Setting USE_MSLK to ON by default , doing CUDA build for SM100a") set(USE_MSLK_DEFAULT ON) diff --git a/aten/src/ATen/CMakeLists.txt b/aten/src/ATen/CMakeLists.txt index 06522a9298893..fd4d1d64d70ca 100644 --- a/aten/src/ATen/CMakeLists.txt +++ b/aten/src/ATen/CMakeLists.txt @@ -430,10 +430,10 @@ IF(USE_MSLK) list(PREPEND MSLK_EXTRA_HIPCC_FLAGS -mllvm -amdgpu-coerce-illegal-types=1) endif() - # Only compile for gfx942 and gfx950. + # Only compile for gfx942, gfx950, and gfx1250. set(HIP_CLANG_FLAGS_ORIGINAL ${HIP_CLANG_FLAGS}) string(REGEX REPLACE "--offload-arch=[^ ]*" "" FILTERED_HIP_CLANG_FLAGS "${HIP_CLANG_FLAGS}") - foreach(ARCH gfx942 gfx950) + foreach(ARCH gfx942 gfx950 gfx1250) if(${ARCH} IN_LIST PYTORCH_ROCM_ARCH) list(APPEND FILTERED_HIP_CLANG_FLAGS --offload-arch=${ARCH}) endif() diff --git a/aten/src/ATen/Context.cpp b/aten/src/ATen/Context.cpp index c342590b58c42..ba63e73e1017c 100644 --- a/aten/src/ATen/Context.cpp +++ b/aten/src/ATen/Context.cpp @@ -505,7 +505,7 @@ at::BlasBackend Context::blasPreferredBackend() { bool Context::ckSupported() { #ifdef USE_ROCM static const std::vector supported_archs = { - "gfx90a", "gfx942", "gfx950" + "gfx90a", "gfx942", "gfx950", "gfx1250", }; for (auto index : c10::irange(detail::getCUDAHooks().deviceCount())) { if(!detail::getCUDAHooks().isGPUArch(supported_archs, index)) { diff --git a/aten/src/ATen/cuda/CUDABlas.cpp b/aten/src/ATen/cuda/CUDABlas.cpp index 748eecbc1572a..f0455e5e8061f 100644 --- a/aten/src/ATen/cuda/CUDABlas.cpp +++ b/aten/src/ATen/cuda/CUDABlas.cpp @@ -2019,6 +2019,14 @@ void scaled_gemm( "Got m=", m, ", n=", n, ", k=", k); } #endif + #if ROCM_VERSION >= 70200 + if (at::detail::getCUDAHooks().isGPUArch({"gfx1250"})) { + // 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. " + "Got m=", m, ", n=", n, ", k=", k); + } + #endif } #elif (CUDA_VERSION < 12090) && !defined(USE_ROCM) // hipblaslt supported row-wise before cublas, and did so their own way (via diff --git a/aten/src/ATen/cuda/CublasHandlePool.cpp b/aten/src/ATen/cuda/CublasHandlePool.cpp index 1fbe3cdedf474..6a79d1a794444 100644 --- a/aten/src/ATen/cuda/CublasHandlePool.cpp +++ b/aten/src/ATen/cuda/CublasHandlePool.cpp @@ -166,8 +166,8 @@ size_t parseChosenWorkspaceSize() { 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; + const bool gfx94_95_125 = at::detail::getCUDAHooks().isGPUArch({"gfx94", "gfx95", "gfx125"}); + const size_t default_size = gfx94_95_125 ? 1024 * 128 * 1024 : 1024 * 32 * 1024; #else /* :4096:2:16:8 default, 32MiB for Hopper and Blackwell */ cudaDeviceProp* properties = at::cuda::getCurrentDeviceProperties(); diff --git a/aten/src/ATen/cuda/detail/CUDAHooks.cpp b/aten/src/ATen/cuda/detail/CUDAHooks.cpp index 5f81407b1ac03..e957c398ebabf 100644 --- a/aten/src/ATen/cuda/detail/CUDAHooks.cpp +++ b/aten/src/ATen/cuda/detail/CUDAHooks.cpp @@ -554,7 +554,10 @@ const std::vector& CUDAHooks::getHipblasltPreferredArchs() const { "gfx1200", "gfx1201", #endif #if ROCM_VERSION >= 70000 - "gfx950" + "gfx950", +#endif +#if ROCM_VERSION >= 70200 + "gfx1250" #endif }; return archs; diff --git a/aten/src/ATen/native/cuda/CUDALoops.cuh b/aten/src/ATen/native/cuda/CUDALoops.cuh index 00f0aedf4578c..a23ff69ee722b 100644 --- a/aten/src/ATen/native/cuda/CUDALoops.cuh +++ b/aten/src/ATen/native/cuda/CUDALoops.cuh @@ -231,7 +231,8 @@ C10_LAUNCH_BOUNDS_1(num_threads()) __global__ void vectorized_elementwise_kernel(int N, func_t f, array_t data) { using traits = function_traits; constexpr auto io_size = calc_io_size(); -#if defined(USE_ROCM) && defined(__gfx942__) + // Extend the TWS (16) to GFX1250. +#if defined(USE_ROCM) && (defined(__gfx942__) || defined(__gfx1250__)) // Similar check in launch_vectorized_kernel() as well. Both should be in sync. constexpr int tws = 16; #else diff --git a/aten/src/ATen/native/cuda/GroupedBlas.cpp b/aten/src/ATen/native/cuda/GroupedBlas.cpp index 5b393401a29b4..7c779e6658a66 100644 --- a/aten/src/ATen/native/cuda/GroupedBlas.cpp +++ b/aten/src/ATen/native/cuda/GroupedBlas.cpp @@ -714,7 +714,7 @@ std::optional out_dtype) { bool use_fast_path = false; // ifdef USE_ROCM_CK_GEMM is required since ROCm systems w/o CK should not call ck path. #if defined(USE_ROCM_CK_GEMM) - if (at::globalContext().rocmAllowGroupGemmCk() && at::detail::getCUDAHooks().isGPUArch({"gfx942", "gfx950", "gfx90a"})) { + if (at::globalContext().rocmAllowGroupGemmCk() && at::detail::getCUDAHooks().isGPUArch({"gfx942", "gfx950", "gfx90a", "gfx1250"})) { use_fast_path = true; } #endif //USE_ROCM_CK_GEMM diff --git a/aten/src/ATen/native/cuda/KernelUtils.cuh b/aten/src/ATen/native/cuda/KernelUtils.cuh index 12feeb6d63af3..cf36105fb6e84 100644 --- a/aten/src/ATen/native/cuda/KernelUtils.cuh +++ b/aten/src/ATen/native/cuda/KernelUtils.cuh @@ -13,7 +13,8 @@ #if ROCM_VERSION < 60400 __device__ inline __hip_bfloat162 preview_unsafeAtomicAdd(__hip_bfloat162* address, __hip_bfloat162 value) { -#if (defined(__gfx942__)) && \ +// `__gfx1250__`-specific `s_wait_loadcnt(0)` path for committed store already there +#if (defined(__gfx942__) || defined(__gfx1250__)) && \ __has_builtin(__builtin_amdgcn_flat_atomic_fadd_v2bf16) typedef unsigned short __attribute__((ext_vector_type(2))) vec_short2; static_assert(sizeof(vec_short2) == sizeof(__hip_bfloat162_raw)); diff --git a/aten/src/ATen/native/cuda/MemoryAccess.cuh b/aten/src/ATen/native/cuda/MemoryAccess.cuh index 373b44cca7901..e96876c8ca149 100644 --- a/aten/src/ATen/native/cuda/MemoryAccess.cuh +++ b/aten/src/ATen/native/cuda/MemoryAccess.cuh @@ -187,7 +187,8 @@ template __device__ aligned_vector load_vector(const scalar_t *base_ptr, uint32_t offset) { using vec_t = aligned_vector; auto *from = reinterpret_cast(base_ptr); -#if defined(USE_ROCM) && defined(__gfx942__) + // Extend the non-temporal load optimization to GFX1250. +#if defined(USE_ROCM) && (defined(__gfx942__) || defined(__gfx1250__)) using longx2 = __attribute__((__vector_size__(4*sizeof(int)))) int; if constexpr (sizeof(vec_t) == sizeof(int)) { union { diff --git a/aten/src/ATen/native/cuda/ScaledBlas.cpp b/aten/src/ATen/native/cuda/ScaledBlas.cpp index 047aec2c8334a..37c470c6adf10 100644 --- a/aten/src/ATen/native/cuda/ScaledBlas.cpp +++ b/aten/src/ATen/native/cuda/ScaledBlas.cpp @@ -79,6 +79,9 @@ bool _scaled_mm_allowed_device(bool sm90_only=false, bool sm100_only=false) { #endif #if ROCM_VERSION >= 60500 "gfx950" +#endif +#if ROCM_VERSION >= 70200 + , "gfx1250" #endif }; return at::detail::getCUDAHooks().isGPUArch(archs); @@ -622,8 +625,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"); + TORCH_CHECK_NOT_IMPLEMENTED(at::detail::getCUDAHooks().isGPUArch({"gfx950", "gfx1250"}), + "Block-wise scaling for Float8_e8m0fnu is only supported on gfx950/gfx1250"); int packed_factor = 1; if (mat1.scalar_type() == ScalarType::Float4_e2m1fn_x2) { @@ -1064,8 +1067,8 @@ _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"); + TORCH_CHECK_NOT_IMPLEMENTED(at::detail::getCUDAHooks().isGPUArch({"gfx950", "gfx1250"}), + "Block-wise scaling for Float8_e8m0fnu is only supported on gfx950/gfx1250"); 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 +1153,8 @@ _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"); + TORCH_CHECK_NOT_IMPLEMENTED(at::detail::getCUDAHooks().isGPUArch({"gfx950", "gfx1250"}), + "Block-wise scaling for Float8_e8m0fnu is only supported on gfx950/gfx1250"); 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, diff --git a/aten/src/ATen/native/cuda/int4mm.cu b/aten/src/ATen/native/cuda/int4mm.cu index 8765bed83345a..07f3b9443ed40 100644 --- a/aten/src/ATen/native/cuda/int4mm.cu +++ b/aten/src/ATen/native/cuda/int4mm.cu @@ -127,7 +127,8 @@ inline __host__ __device__ uint32_t getAlignmentRoundUp(const void* p) { return diff == 0 ? 0 : uint32_t(Align) - diff; } -#if defined (__gfx90a__) || defined(__gfx942__) || defined(__gfx950__) +// CDNA arch with MFMA and Warp-32 support +#if defined(__gfx90a__) || defined(__gfx942__) || defined(__gfx950__) #define CDNA2_OR_LATER 1 #else #define CDNA2_OR_LATER 0 @@ -146,6 +147,12 @@ static bool isCDNA2orLater(int index) { return at::detail::getCUDAHooks().isGPUArch({"gfx90a", "gfx942", "gfx950"}, index); } +// Conceptual for now and subject to change +// gfx1250 (CDNA5 / CDNA-next / UDNA) +static bool isCDNA5orLater(int index) { + return at::detail::getCUDAHooks().isGPUArch({"gfx1250"}, index); +} + #else constexpr int32_t kWarpSize = 32; #endif @@ -1098,6 +1105,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 +1305,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"); } diff --git a/aten/src/ATen/native/sparse/cuda/cuSPARSELtOps.cpp b/aten/src/ATen/native/sparse/cuda/cuSPARSELtOps.cpp index 24893da98c6c3..07627b1587f97 100644 --- a/aten/src/ATen/native/sparse/cuda/cuSPARSELtOps.cpp +++ b/aten/src/ATen/native/sparse/cuda/cuSPARSELtOps.cpp @@ -30,7 +30,7 @@ static void initHipSparseLtSupport() { // Check only the first available device try { if (at::cuda::device_count() > 0) { - g_hipSparseLtSupported = at::detail::getCUDAHooks().isGPUArch({"gfx950", "gfx942"}, 0); + g_hipSparseLtSupported = at::detail::getCUDAHooks().isGPUArch({"gfx950", "gfx942", "gfx1250"}, 0); } } catch (const std::exception&) { // If an exception occurs during device property check, we assume hipSparseLt is not supported @@ -49,7 +49,7 @@ static bool isHipSparseLtSupported() { TORCH_CHECK( false, "hipSparseLt not supported on this device, supported architectures: " - "gfx950, gfx942. " + "gfx1250, gfx950, gfx942. " "required ROCM version: 6.4.0 or later."); } return g_hipSparseLtSupported; diff --git a/aten/src/ATen/native/transformers/cuda/sdp_utils.cpp b/aten/src/ATen/native/transformers/cuda/sdp_utils.cpp index 79b5df3f302bb..ff3eee28bc314 100644 --- a/aten/src/ATen/native/transformers/cuda/sdp_utils.cpp +++ b/aten/src/ATen/native/transformers/cuda/sdp_utils.cpp @@ -301,6 +301,7 @@ bool check_flash_attention_hardware_support(sdp_params const& params, bool debug using sm121 = SMVersion<12, 1>; #if USE_ROCM #if USE_ROCM_ATTENTION +// TODO: gfx1250 if(at::globalContext().getROCmFAPreferredBackend() == at::ROCmFABackend::Ck) { // User explicitly set CK as the flash attention backend. Return true for now // TODO: Flesh out sanity checks diff --git a/aten/src/ATen/native/transformers/hip/flash_attn/ck/launch_kernel_pt.hpp b/aten/src/ATen/native/transformers/hip/flash_attn/ck/launch_kernel_pt.hpp index 5f7a16cffa1c7..3838b3e2edf2b 100644 --- a/aten/src/ATen/native/transformers/hip/flash_attn/ck/launch_kernel_pt.hpp +++ b/aten/src/ATen/native/transformers/hip/flash_attn/ck/launch_kernel_pt.hpp @@ -27,7 +27,7 @@ __launch_bounds__(Kernel::kBlockSize, MinBlockPerCu) #endif __global__ void kentry_pt(Args... args) { -#if (defined(__gfx90a__) || defined(__gfx942__) || defined(__gfx950__)) +#if (defined(__gfx90a__) || defined(__gfx942__) || defined(__gfx950__) || defined(__gfx1250__)) Kernel{}(args...); #else CUDA_KERNEL_ASSERT(false && "Fatal! Attempting to call a CK SDPA kernel on unsupported hardware"); diff --git a/c10/core/AllocatorConfig.h b/c10/core/AllocatorConfig.h index d314c93f0494f..ef4e798939c5e 100644 --- a/c10/core/AllocatorConfig.h +++ b/c10/core/AllocatorConfig.h @@ -19,7 +19,13 @@ constexpr size_t kMinBlockSize = 512; // largest "small" allocation is 1 MiB constexpr size_t kSmallSize = 1048576; // allocations between 1 and 10 MiB may use kLargeBuffer -constexpr size_t kMinLargeAlloc = 10485760; +#if defined(USE_ROCM) && defined(__gfx1250__) +// Increase the buffer threshold for gfx1250 +// to avoid fragmentation on 432GB devices. +constexpr size_t kMinLargeAlloc = 20 * 1024 * 1024; // 20 MiB +#else +constexpr size_t kMinLargeAlloc = 10485760; // 10 MiB +#endif // round up large allocations to 2 MiB constexpr size_t kRoundLarge = 2097152; diff --git a/cmake/External/aotriton.cmake b/cmake/External/aotriton.cmake index d2f1e15ff1186..414dced29e297 100644 --- a/cmake/External/aotriton.cmake +++ b/cmake/External/aotriton.cmake @@ -45,6 +45,8 @@ if(NOT __AOTRITON_INCLUDED) "amd-gfx950" "amd-gfx11xx" "amd-gfx120x" + # TODO: Update on AOTriton integration + #"amd-gfx1250" ) set(__AOTRITON_IMAGE_SHA256_LIST "fe9f04b66bf52ac27cd025e1d89cfd04974dd3fb3ae076192f783641a4d80fdf" # amd-gfx90a @@ -52,6 +54,9 @@ if(NOT __AOTRITON_INCLUDED) "c1ba3bfe84217fd67df3dd1f8b67c80a7f7b33d0ad4d74b41d6567036e032ace" # amd-gfx950 "839299637fccb13fbe3e7823d57d1b2dcd0e0bed78abbcb7005ea5f4fd82b928" # amd-gfx11xx "0a4ff324bffdac0c2fde87a8a7f70563d3c84a80ad4e8f31345f2b40a1384e95" # amd-gfx120x + # TODO: Update when AOTriton publishes gfx1250 images. + # Until then, may need to set AOTRITON_INSTALL_FROM_SOURCE=1 to build from source. + #"0000000000000000000000000000000000000000000000000000000000000000" # amd-gfx1250 ) set(__AOTRITON_BASE_URL "https://github.com/ROCm/aotriton/releases/download/") # @lint-ignore set(__AOTRITON_Z "gz") diff --git a/test/test_cuda.py b/test/test_cuda.py index dc7b1150bae2c..b624fd7e70210 100644 --- a/test/test_cuda.py +++ b/test/test_cuda.py @@ -687,7 +687,7 @@ def _check_default(): gcn_arch = str( torch.cuda.get_device_properties(0).gcnArchName.split(":", 1)[0] ) - if gcn_arch in ["gfx90a", "gfx942", "gfx950", "gfx1200", "gfx1201"]: + if gcn_arch in ["gfx90a", "gfx942", "gfx950", "gfx1200", "gfx1201", "gfx1250"]: self.assertTrue(default == torch._C._BlasBackend.Cublaslt) else: self.assertTrue(default == torch._C._BlasBackend.Cublas) @@ -757,7 +757,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 "gfx94" in gcn_arch or "gfx95" in gcn_arch or "gfx1250" in gcn_arch: default_workspace_size = 1024 * 128 * 1024 # :1024:128 else: default_workspace_size = ( @@ -8434,7 +8434,9 @@ 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 + 65536 + if get_device_properties().gcnArchName not in ["gfx950", "gfx1250"] + else 160 * 1024 ) else: max_smem = get_device_properties().shared_memory_per_block_optin diff --git a/test/test_linalg.py b/test/test_linalg.py index 82c5ee64b51d8..7f8db8b28ff4f 100644 --- a/test/test_linalg.py +++ b/test/test_linalg.py @@ -65,6 +65,12 @@ def blaslt_supported_device(): archs.extend(['gfx110', 'gfx120']) if ROCM_VERSION >= (6, 5): archs.append('gfx95') + # We extend this in a way not assuming the exact ROCm version + # where gfx1250 lands, so we don't pin a ROCm_VERSION gate here + # because the exact landing version may shift. + # Prefix checks rather than exact matches treat MI-series / gfx1250 + # as BLASLt-capabile. + archs.append('gfx125') for arch in archs: if arch in torch.cuda.get_device_properties(0).gcnArchName: return True diff --git a/torch/_inductor/config.py b/torch/_inductor/config.py index f68171baef7c0..a3dbe5ae9d984 100644 --- a/torch/_inductor/config.py +++ b/torch/_inductor/config.py @@ -2376,10 +2376,11 @@ class rocm: # Enable the CK backend for CDNA2 and CDNA3 only (for now) # Processor name reference: https://llvm.org/docs/AMDGPUUsage.html#processors - ck_supported_arch: list[Literal["gfx90a", "gfx942", "gfx950"]] = [ + ck_supported_arch: list[Literal["gfx90a", "gfx942", "gfx950", "gfx1250"]] = [ "gfx90a", "gfx942", "gfx950", + "gfx1250", ] # Optimization level, use to balance compilation speed and runtime performance. diff --git a/torch/cuda/_utils.py b/torch/cuda/_utils.py index 8d9501afcbff3..3f49b0e5f3772 100644 --- a/torch/cuda/_utils.py +++ b/torch/cuda/_utils.py @@ -475,11 +475,14 @@ def set_shared_memory_config(self, shared_mem_bytes: int) -> None: device_props = torch.cuda.get_device_properties() # 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 allows a max of 160KB shared memory - max_shared_mem = ( - 65536 if device_props.gcnArchName != "gfx950" else 160 * 1024 - ) + # navi, CDNA1-CDNA3 allows a max of 64KB shared memory, + # CDNA4 (gfx950) 160KB, and CDNA5 (gfx1250) 320KB. + if device_props.gcnArchName == "gfx950": + max_shared_mem = 160 * 1024 + elif device_props.gcnArchName == "gfx1250": + max_shared_mem = 320 * 1024 + else: + max_shared_mem = 65536 else: max_shared_mem = getattr( device_props, "shared_memory_per_block_optin", 49152 diff --git a/torch/testing/_internal/common_cuda.py b/torch/testing/_internal/common_cuda.py index 51c2f1fb7e215..2d04f21366803 100644 --- a/torch/testing/_internal/common_cuda.py +++ b/torch/testing/_internal/common_cuda.py @@ -63,6 +63,10 @@ def evaluate_gfx_arch_within(arch_list): # Hence the matching should be done reversely return any(arch in effective_arch for arch in arch_list) +# CDNA 5 (CDNA-next / UDAN) arch helper +def CDNA5OrLater(): + return evaluate_gfx_arch_within(["gfx1250"]) + def CDNA3OrLater(): return evaluate_gfx_arch_within(["gfx942", "gfx950"]) @@ -71,7 +75,7 @@ def CDNA2OrLater(): def evaluate_platform_supports_flash_attention(): if TEST_WITH_ROCM: - arch_list = ["gfx90a", "gfx942", "gfx1100", "gfx1201", "gfx950"] + arch_list = ["gfx90a", "gfx942", "gfx1100", "gfx1201", "gfx950", "gfx1250"] 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) @@ -89,7 +93,7 @@ def evaluate_platform_supports_ck_sdpa(): def evaluate_platform_supports_efficient_attention(): if TEST_WITH_ROCM: - arch_list = ["gfx90a", "gfx942", "gfx1100", "gfx1201", "gfx950"] + arch_list = ["gfx90a", "gfx942", "gfx1100", "gfx1201", "gfx950", "gfx1250"] 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) @@ -175,6 +179,8 @@ def evaluate_platform_supports_fp8(): archs.extend(['gfx120']) if ROCM_VERSION >= (6, 5): archs.append('gfx95') + if ROCM_VERSION >= (7, 2): + archs.append('gfx1250') for arch in archs: if arch in torch.cuda.get_device_properties(0).gcnArchName: return True @@ -191,7 +197,7 @@ def evaluate_platform_supports_fp8_grouped_gemm(): if torch.version.hip: if "USE_MSLK" not in torch.__config__.show(): return False - archs = ['gfx942', 'gfx950'] + archs = ['gfx942', 'gfx950', 'gfx1250'] for arch in archs: if arch in torch.cuda.get_device_properties(0).gcnArchName: return True @@ -203,7 +209,8 @@ def evaluate_platform_supports_mx_gemm(): if torch.cuda.is_available(): if torch.version.hip: if ROCM_VERSION >= (7, 0): - return 'gfx950' in torch.cuda.get_device_properties(0).gcnArchName + gcn_name = torch.cuda.get_device_properties(0).gcnArchName + return 'gfx950' in gcn_name or ('gfx1250' in gcn_name and ROCM_VERSION >= (7, 2)) else: return SM100OrLater return False diff --git a/torch/testing/_internal/common_distributed.py b/torch/testing/_internal/common_distributed.py index 894acb6be19fe..ce0b6026e8425 100644 --- a/torch/testing/_internal/common_distributed.py +++ b/torch/testing/_internal/common_distributed.py @@ -493,7 +493,7 @@ def requires_multicast_support(): def evaluate_platform_supports_symm_mem(): if TEST_CUDA: if TEST_WITH_ROCM: - arch_list = ["gfx942", "gfx950"] + arch_list = ["gfx942", "gfx950", "gfx1250"] for arch in arch_list: if arch in torch.cuda.get_device_properties(0).gcnArchName: return True From a3dd7ed2932c897732712aaef1ec70ab215b2040 Mon Sep 17 00:00:00 2001 From: "Li, Bo" Date: Thu, 30 Apr 2026 12:03:31 -0500 Subject: [PATCH 02/11] CK - gfx1250 supports (#14) * CK - gfx1250 support (#5) * Enable ROCM_CK_SDPA build * [submodule] composable_kernel and aiter update (#172592) Summary: update ck to commit https://github.com/ROCm/composable_kernel/commit/fcc9372c009c8e0a23fece77b582da83b04a654f update aiter to commit https://github.com/ROCm/aiter/commit/9a469a608b2c10b7157df573a38d31e5bf4038b4 changes of caffe2/aten/src/ATen/CMakeLists.txt and caffe2/caffe2/CMakeLists.txt are adopted from https://github.com/pytorch/pytorch/pull/161759 updated caffe2/aten/src/ATen/native/transformers/hip/flash_attn/ck/launch_kernel_pt.hpp to match the ck version in https://github.com/ROCm/composable_kernel/blob/292df2719f28cd01464d5d059820684790c101da/include/ck_tile/host/kernel_launch.hpp update aiter fav3 bwd codegen according to changes in https://github.com/ROCm/aiter/pull/1573 update caffe2/aten/src/ATen/native/transformers/hip/flash_attn/ck mha fwd/bwd kernels according to the interfaces in https://github.com/ROCm/composable_kernel/tree/292df2719f28cd01464d5d059820684790c101da/example/ck_tile/01_fmha Differential Revision: D88991877 Pull Request resolved: https://github.com/pytorch/pytorch/pull/172592 Approved by: https://github.com/alugorey, https://github.com/izaitsevfb * Added MI450 supports and packages * Fix misalinged ck api * Replace aiter with ck for bwd * [ROCm] Bump AOTriton to 0.11.2b (#174105) Notable new features: * AOTriton 0.11.2b adds gfx1151/1152/1153 support. * Add precompiled AOTriton runtime for ROCM 7.2 * Match the sliding window attention behavior of `_flash_attention_forward/backward` with CUTLASS backend. Bug fixes: * Fixes #173204. Now all tests in `test/test_varlen_attention.py` are enabled on ROCm Notes: This replaces PR #173820 and #173469 Pull Request resolved: https://github.com/pytorch/pytorch/pull/174105 Approved by: https://github.com/jeffdaily * Fix philox data types for this version of ck * Update CK to use new gfx1250_pytorch branch * Add new gfx1250 compile flags for CK * add --targets to generate and a couple new compile flags * Remove default USE_ROCM_CK_SDPA --------- Co-authored-by: blorange-amd Co-authored-by: Yu Guo Co-authored-by: Xinya Zhang * Updated aiter module * Fixed merged error * Fixed additional merged error * Reset USE_ROCM_CK_SDPA config --------- Co-authored-by: LugoReyes, Andy Co-authored-by: Yu Guo Co-authored-by: Xinya Zhang --- aten/src/ATen/CMakeLists.txt | 11 +++- .../hip/flash_attn/ck/CMakeLists.txt | 16 +++--- .../hip/flash_attn/ck/fav_v3/CMakeLists.txt | 2 +- .../hip/flash_attn/ck/launch_kernel_pt.hpp | 2 +- .../hip/flash_attn/ck/mha_bwd_ck.hip | 53 +++++++++++++++---- .../hip/flash_attn/ck/mha_fwd_ck.hip | 34 ++++++++++-- .../hip/flash_attn/ck/mha_varlen_fwd_ck.hip | 23 +++++++- requirements.txt | 1 + third_party/aiter | 2 +- 9 files changed, 113 insertions(+), 31 deletions(-) diff --git a/aten/src/ATen/CMakeLists.txt b/aten/src/ATen/CMakeLists.txt index fd4d1d64d70ca..b030288ef5070 100644 --- a/aten/src/ATen/CMakeLists.txt +++ b/aten/src/ATen/CMakeLists.txt @@ -268,8 +268,15 @@ if(USE_FLASH_ATTENTION) CK_ENABLE_FP64 CK_ENABLE_FP8 CK_ENABLE_INT8 - CK_USE_FNUZ_FP8 - CK_USE_GFX94 + #CK_USE_FNUZ_FP8 + #CK_USE_GFX94 + CK_USE_GFX1250 + CK_USE_NATIVE_MX_SUPPORT + CK_GFX1250_SUPPORT + CK_GFX12_SUPPORT + CK_USE_OCP_FP8 + CK_USE_WMMA + CK_USE_WMMA_FP8 CK_USE_XDL __HIP_PLATFORM_AMD__=1 __HIP_PLATFORM_HCC__=1 diff --git a/aten/src/ATen/native/transformers/hip/flash_attn/ck/CMakeLists.txt b/aten/src/ATen/native/transformers/hip/flash_attn/ck/CMakeLists.txt index 819880cf3bc5c..3fb754cb8bf14 100644 --- a/aten/src/ATen/native/transformers/hip/flash_attn/ck/CMakeLists.txt +++ b/aten/src/ATen/native/transformers/hip/flash_attn/ck/CMakeLists.txt @@ -1,6 +1,6 @@ # generate a list of kernels, but not actually emit files at config stage execute_process( - COMMAND python3 ${CMAKE_SOURCE_DIR}/third_party/composable_kernel/example/ck_tile/01_fmha/generate.py + COMMAND python3 ${CMAKE_SOURCE_DIR}/third_party/composable_kernel/example/ck_tile/01_fmha/generate.py --targets gfx1250 --api fwd --receipt 4 --list_blobs ${CMAKE_CURRENT_LIST_DIR}/fwd_blob_list.txt RESULT_VARIABLE ret ) @@ -10,7 +10,7 @@ if(ret AND NOT ret EQUAL 0) endif() execute_process( - COMMAND python3 ${CMAKE_SOURCE_DIR}/third_party/composable_kernel/example/ck_tile/01_fmha/generate.py + COMMAND python3 ${CMAKE_SOURCE_DIR}/third_party/composable_kernel/example/ck_tile/01_fmha/generate.py --targets gfx1250 --api fwd_splitkv --receipt 4 --list_blobs ${CMAKE_CURRENT_LIST_DIR}/fwd_splitkv_blob_list.txt RESULT_VARIABLE ret ) @@ -20,7 +20,7 @@ if(ret AND NOT ret EQUAL 0) endif() execute_process( - COMMAND python3 ${CMAKE_SOURCE_DIR}/third_party/composable_kernel/example/ck_tile/01_fmha/generate.py + COMMAND python3 ${CMAKE_SOURCE_DIR}/third_party/composable_kernel/example/ck_tile/01_fmha/generate.py --targets gfx1250 --api fwd_appendkv --receipt 4 --list_blobs ${CMAKE_CURRENT_LIST_DIR}/fwd_appendkv_blob_list.txt RESULT_VARIABLE ret ) @@ -30,7 +30,7 @@ if(ret AND NOT ret EQUAL 0) endif() execute_process( - COMMAND python3 ${CMAKE_SOURCE_DIR}/third_party/composable_kernel/example/ck_tile/01_fmha/generate.py + COMMAND python3 ${CMAKE_SOURCE_DIR}/third_party/composable_kernel/example/ck_tile/01_fmha/generate.py --targets gfx1250 --api bwd --receipt 4 --list_blobs ${CMAKE_CURRENT_LIST_DIR}/bwd_blob_list.txt RESULT_VARIABLE ret ) @@ -40,28 +40,28 @@ if(ret AND NOT ret EQUAL 0) endif() # Generate the files for both fwd, fwd_splitkv, fwd_appendkv, and bwd -execute_process(COMMAND python3 ${CMAKE_SOURCE_DIR}/third_party/composable_kernel/example/ck_tile/01_fmha/generate.py --api fwd --receipt 4 --output_dir ${CMAKE_CURRENT_LIST_DIR} +execute_process(COMMAND python3 ${CMAKE_SOURCE_DIR}/third_party/composable_kernel/example/ck_tile/01_fmha/generate.py --targets gfx1250 --api fwd --receipt 4 --output_dir ${CMAKE_CURRENT_LIST_DIR} ) if(ret AND NOT ret EQUAL 0) message( FATAL_ERROR "CK Tile FMHA FAILED to generate FWD kernels.") endif() -execute_process(COMMAND python3 ${CMAKE_SOURCE_DIR}/third_party/composable_kernel/example/ck_tile/01_fmha/generate.py --api fwd_splitkv --receipt 4 --output_dir ${CMAKE_CURRENT_LIST_DIR} +execute_process(COMMAND python3 ${CMAKE_SOURCE_DIR}/third_party/composable_kernel/example/ck_tile/01_fmha/generate.py --targets gfx1250 --api fwd_splitkv --receipt 4 --output_dir ${CMAKE_CURRENT_LIST_DIR} ) if(ret AND NOT ret EQUAL 0) message( FATAL_ERROR "CK Tile FMHA FAILED to generate FWD_SPLITKV kernels.") endif() -execute_process(COMMAND python3 ${CMAKE_SOURCE_DIR}/third_party/composable_kernel/example/ck_tile/01_fmha/generate.py --api fwd_appendkv --receipt 4 --output_dir ${CMAKE_CURRENT_LIST_DIR} +execute_process(COMMAND python3 ${CMAKE_SOURCE_DIR}/third_party/composable_kernel/example/ck_tile/01_fmha/generate.py --targets gfx1250 --api fwd_appendkv --receipt 4 --output_dir ${CMAKE_CURRENT_LIST_DIR} ) if(ret AND NOT ret EQUAL 0) message( FATAL_ERROR "CK Tile FMHA FAILED to generate FWD_APPENDKV kernels.") endif() -execute_process(COMMAND python3 ${CMAKE_SOURCE_DIR}/third_party/composable_kernel/example/ck_tile/01_fmha/generate.py --api bwd --receipt 4 --output_dir ${CMAKE_CURRENT_LIST_DIR} +execute_process(COMMAND python3 ${CMAKE_SOURCE_DIR}/third_party/composable_kernel/example/ck_tile/01_fmha/generate.py --targets gfx1250 --api bwd --receipt 4 --output_dir ${CMAKE_CURRENT_LIST_DIR} RESULT_VARIABLE ret ) diff --git a/aten/src/ATen/native/transformers/hip/flash_attn/ck/fav_v3/CMakeLists.txt b/aten/src/ATen/native/transformers/hip/flash_attn/ck/fav_v3/CMakeLists.txt index e0eb28652c1e6..2bed9e565d200 100644 --- a/aten/src/ATen/native/transformers/hip/flash_attn/ck/fav_v3/CMakeLists.txt +++ b/aten/src/ATen/native/transformers/hip/flash_attn/ck/fav_v3/CMakeLists.txt @@ -2,7 +2,7 @@ include(CMakePrintHelpers) # Generate AITER/CK Asm code execute_process( - COMMAND ${CMAKE_COMMAND} -E env "AITER_GPU_ARCHS=gfx942;gfx950" + COMMAND ${CMAKE_COMMAND} -E env "AITER_GPU_ARCHS=gfx942;gfx950;gfx1250" python3 ${CMAKE_SOURCE_DIR}/third_party/aiter/hsa/codegen.py -m fmha_v3_bwd --output_dir ${CMAKE_CURRENT_LIST_DIR} RESULT_VARIABLE ret ) diff --git a/aten/src/ATen/native/transformers/hip/flash_attn/ck/launch_kernel_pt.hpp b/aten/src/ATen/native/transformers/hip/flash_attn/ck/launch_kernel_pt.hpp index 3838b3e2edf2b..6bbcf09d5c92d 100644 --- a/aten/src/ATen/native/transformers/hip/flash_attn/ck/launch_kernel_pt.hpp +++ b/aten/src/ATen/native/transformers/hip/flash_attn/ck/launch_kernel_pt.hpp @@ -14,7 +14,7 @@ __launch_bounds__(Kernel::kBlockSize, MinBlockPerCu) #endif __global__ void kentry_pt(Args... args) { -#if (defined(__gfx90a__) || defined(__gfx942__) || defined(__gfx950__)) +#if (defined(__gfx90a__) || defined(__gfx942__) || defined(__gfx950__) || defined(__gfx1250__)) Kernel{}(args...); #else CUDA_KERNEL_ASSERT(false && "Fatal! Attempting to call a CK SDPA kernel on unsupported hardware"); diff --git a/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_bwd_ck.hip b/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_bwd_ck.hip index d7633521d8eef..aec411e9978d7 100644 --- a/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_bwd_ck.hip +++ b/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_bwd_ck.hip @@ -3,13 +3,32 @@ ******************************************************************************/ #include -#include #include #include namespace pytorch_flash { -aiter::mha_bwd_args get_ck_fmha_bwd_args(const mask_info &mask, +fmha_bwd_traits get_ck_fmha_bwd_traits(const mask_info &mask, + std::string dtype, + int head_size, + bool has_dropout, + bool enable_bias, + bool deterministic, + bool bias_requires_grad) +{ + return fmha_bwd_traits{head_size, + head_size, + dtype, + false, // is_group_mode (batch mode) + mask.type, + enable_bias ? bias_enum::elementwise_bias : bias_enum::no_bias, + bias_requires_grad, + has_dropout, + false, // is_store_randval + deterministic}; +} + +fmha_bwd_args get_ck_fmha_bwd_args(const mask_info &mask, std::string dtype, bool has_dropout, bool enable_bias, @@ -167,16 +186,18 @@ aiter::mha_bwd_args get_ck_fmha_bwd_args(const mask_info &mask, nullptr, // cu_seqlen_k_ptr seqlen_q, seqlen_k, - b, // batch + b, seqlen_q, // max_seqlen_q seqlen_k, // max_seqlen_k - h, // nhead_q - h_k, // nhead_k - softmax_scale, // scale + hdim, // hdim_q + hdim, // hdim_v + h, // nhead_q + h_k, // nhead_k + softmax_scale, stride_q, stride_k, stride_v, - stride_attn_bias, // stride_bias + stride_attn_bias, stride_o, 0, // stride_randval stride_do, @@ -212,9 +233,10 @@ aiter::mha_bwd_args get_ck_fmha_bwd_args(const mask_info &mask, batch_stride_dv, batch_stride_dbias, split_stride_dq_acc, - mask.left, // window_size_left - mask.right, // window_size_right - p_dropout, // p_drop + mask.left, + mask.right, + static_cast(mask.type), + p_dropout, p_undrop, drop_seed_offset }; @@ -390,6 +412,15 @@ mha_bwd_ck(const at::Tensor &dout, // batch_size x seqlen_q x ck_tile::stream_config stream_config{stream}; dq.zero_(); // ck use atomic operation on dq + auto traits = + get_ck_fmha_bwd_traits(mask, + q_dtype_str, + head_size_8x, + is_dropout, + attn_bias_.has_value(), + deterministic, + bias_requires_grad); + auto args = get_ck_fmha_bwd_args( mask, @@ -421,7 +452,7 @@ mha_bwd_ck(const at::Tensor &dout, // batch_size x seqlen_q x p_dropout, drop_seed_offset); - float t = aiter::mha_bwd(args, stream_config); + float t = fmha_bwd(traits, args, stream_config); TORCH_CHECK(t >= 0, "invalid argument for fmha_bwd"); } else { diff --git a/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_fwd_ck.hip b/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_fwd_ck.hip index c5902e33afd36..7c33e7c3529de 100644 --- a/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_fwd_ck.hip +++ b/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_fwd_ck.hip @@ -87,6 +87,15 @@ fmha_fwd_args get_ck_fmha_fwd_args(bool has_lse, ck_tile::index_t batch_stride_bias = 0; ck_tile::index_t nhead_stride_bias = 0; + ck_tile::index_t nhead_stride_q_descale = 0; + ck_tile::index_t nhead_stride_k_descale = 0; + ck_tile::index_t nhead_stride_v_descale = 0; + ck_tile::index_t batch_stride_q_descale = 0; + ck_tile::index_t batch_stride_k_descale = 0; + ck_tile::index_t batch_stride_v_descale = 0; + ck_tile::index_t block_scale_size_q = 0; + ck_tile::index_t block_scale_size_kv = 0; + if (attn_bias_.has_value()) { auto a_b = attn_bias_.value(); CHECK_DEVICE(a_b); @@ -112,7 +121,8 @@ fmha_fwd_args get_ck_fmha_fwd_args(bool has_lse, nullptr, // seqlen_k_ptr nullptr, // cu_seqlen_q_ptr nullptr, // cu_seqlen_k_ptr - nullptr, // sink_ptr + nullptr, // block_scale_seqstart_q_ptr + nullptr, // block_scale_seqstart_k_ptr seqlen_q, seqlen_k, b, @@ -136,6 +146,9 @@ fmha_fwd_args get_ck_fmha_fwd_args(bool has_lse, nhead_stride_randval, nhead_stride_lse, nhead_stride_o, + nhead_stride_q_descale, + nhead_stride_k_descale, + nhead_stride_v_descale, batch_stride_q, batch_stride_k, batch_stride_v, @@ -143,6 +156,9 @@ fmha_fwd_args get_ck_fmha_fwd_args(bool has_lse, batch_stride_randval, batch_stride_lse, batch_stride_o, + batch_stride_q_descale, + batch_stride_k_descale, + batch_stride_v_descale, mask.left, mask.right, 0, // sink_size @@ -150,7 +166,9 @@ fmha_fwd_args get_ck_fmha_fwd_args(bool has_lse, -1, // min_seqlen_q p_dropout, has_dropout_randval, - drop_seed_offset}; + drop_seed_offset, + block_scale_size_q, + block_scale_size_kv}; } std::tuple @@ -305,9 +323,15 @@ mha_fwd_ck(const at::Tensor &q, // batch_size x seqlen_q x rng_state[0] = *(reinterpret_cast(&drop_seed)); rng_state[1] = *(reinterpret_cast(&drop_offset)); - } else { - seed_t = at::empty({}, at::dtype(at::kLong).device(at::kCUDA)); - offset_t = at::empty({}, at::dtype(at::kLong).device(at::kCUDA)); + hipLaunchKernelGGL( + flash::ParsePhiloxCudaState, dim3(1), dim3(64), 0, at::cuda::getCurrentCUDAStream(), philox_args, rng_state_ptr); + seed_t = at::scalar_tensor(at::Scalar(static_cast(rng_state_ptr[0])), at::dtype(at::kUInt64)); + offset_t = at::scalar_tensor(at::Scalar(static_cast(rng_state_ptr[1])), at::dtype(at::kUInt64)); + } + else + { + seed_t = at::empty({}, at::dtype(at::kUInt64).device(at::kCUDA)); + offset_t = at::empty({}, at::dtype(at::kUInt64).device(at::kCUDA)); } std::optional attn_bias; std::pair drop_seed_offset; diff --git a/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_varlen_fwd_ck.hip b/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_varlen_fwd_ck.hip index 39ce82ef1e700..e7d18a6028b6a 100644 --- a/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_varlen_fwd_ck.hip +++ b/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_varlen_fwd_ck.hip @@ -84,6 +84,16 @@ fmha_fwd_args get_ck_fmha_varlen_fwd_args(bool has_lse, ck_tile::index_t batch_stride_lse = 0; ck_tile::index_t batch_stride_randval = 0; + + ck_tile::index_t nhead_stride_q_descale = 0; + ck_tile::index_t nhead_stride_k_descale = 0; + ck_tile::index_t nhead_stride_v_descale = 0; + ck_tile::index_t batch_stride_q_descale = 0; + ck_tile::index_t batch_stride_k_descale = 0; + ck_tile::index_t batch_stride_v_descale = 0; + ck_tile::index_t block_scale_size_q = 0; + ck_tile::index_t block_scale_size_kv = 0; + void *attn_bias_ptr = nullptr; ck_tile::index_t stride_attn_bias = 0; @@ -113,7 +123,8 @@ fmha_fwd_args get_ck_fmha_varlen_fwd_args(bool has_lse, nullptr, // seqlen_k_ptr nullptr, // cu_seqlen_q_ptr nullptr, // cu_seqlen_k_ptr - nullptr, // sink_ptr + nullptr, // block_scale_seqstart_q_ptr + nullptr, // block_scale_seqstart_k_ptr total_q, total_k, b, @@ -137,6 +148,9 @@ fmha_fwd_args get_ck_fmha_varlen_fwd_args(bool has_lse, nhead_stride_randval, nhead_stride_lse, nhead_stride_o, + nhead_stride_q_descale, + nhead_stride_k_descale, + nhead_stride_v_descale, batch_stride_q, batch_stride_k, batch_stride_v, @@ -144,6 +158,9 @@ fmha_fwd_args get_ck_fmha_varlen_fwd_args(bool has_lse, batch_stride_randval, batch_stride_lse, batch_stride_o, + batch_stride_q_descale, + batch_stride_k_descale, + batch_stride_v_descale, mask.left, mask.right, 0, // sink_size @@ -151,7 +168,9 @@ fmha_fwd_args get_ck_fmha_varlen_fwd_args(bool has_lse, -1, // min_seqlen_q p_dropout, has_dropout_randval, - drop_seed_offset}; + drop_seed_offset, + block_scale_size_q, + block_scale_size_kv}; } std::tuple diff --git a/requirements.txt b/requirements.txt index d6552edad483b..13430c258d109 100644 --- a/requirements.txt +++ b/requirements.txt @@ -14,6 +14,7 @@ lintrunner==0.12.11 ; platform_machine != "s390x" networkx==2.8.8 optree==0.13.0 ; python_version < "3.14" optree==0.17.0 ; python_version >= "3.14" +pandas psutil==7.2.2 spin==0.17 sympy==1.13.3 diff --git a/third_party/aiter b/third_party/aiter index 9a469a608b2c1..bfa86a0b25652 160000 --- a/third_party/aiter +++ b/third_party/aiter @@ -1 +1 @@ -Subproject commit 9a469a608b2c10b7157df573a38d31e5bf4038b4 +Subproject commit bfa86a0b25652b9b0da5c3e3692136789dc4f984 From aec8828e295af22052301e6436bb8910decbad72 Mon Sep 17 00:00:00 2001 From: "Vasishta, Aaryaman (Jam)" Date: Fri, 8 May 2026 03:26:34 +0900 Subject: [PATCH 03/11] [ROCm] Fix large ROCm arange launch (#182657) (#16) Fix `torch.arange` (and the other range factories sharing this kernel) for very large outputs on ROCm. `torch.arange(N)` with `N >= 2^32` fails on ROCm because `hipLaunchKernel` does not support `gridDim.x * blockDim.x >= 2^32` for the per-thread kernel `aten/src/ATen/native/cuda/RangeFactories.cu` previously used. Depending on the ROCm version the launch returns `hipErrorInvalidConfiguration` or is accepted silently with the kernel never executing, leaving zero-initialized output. Concrete repro: `torch.arange(2 ** 32 + 1, device="cuda", dtype=torch.int32)`. The fix replaces the per-thread launch on the ROCm path with a grid-stride loop that fixes the grid at `sm_count * 4` blocks, so the launch limit is no longer load-bearing for correctness regardless of `N`. The non-ROCm path is untouched. On MI250X the grid-stride kernel matches the per-thread kernel within noise at `N=1024` and is 24-60% faster from `N=1M` up across `int32`, `int64`, and `float32`. On MI300X the grid-stride kernel matches within noise at `N=1024` and `N=1M`, and is 2-5x faster from `N=64M` up across `int32`, `int64`, and `float32`. The 64-bit-indexing test is extended to also cover `N = 2^32 + 1` and `N = 2^33 + 1` on ROCm when memory permits. Pull Request resolved: https://github.com/pytorch/pytorch/pull/182657 Approved by: https://github.com/jeffdaily Co-authored-by: Jeff Daily --- aten/src/ATen/native/cuda/RangeFactories.cu | 41 +++++++++++++++++++++ test/test_tensor_creation_ops.py | 18 +++++++++ 2 files changed, 59 insertions(+) diff --git a/aten/src/ATen/native/cuda/RangeFactories.cu b/aten/src/ATen/native/cuda/RangeFactories.cu index 9d7ead7e49892..bc8decb5469ec 100644 --- a/aten/src/ATen/native/cuda/RangeFactories.cu +++ b/aten/src/ATen/native/cuda/RangeFactories.cu @@ -8,6 +8,9 @@ #include #include #include +#if defined(USE_ROCM) +#include +#endif #ifndef AT_PER_OPERATOR_HEADERS #include @@ -48,12 +51,49 @@ __global__ void elementwise_kernel_with_index(index_t N, func_t f, typename func } } +#if defined(USE_ROCM) +// HIP does not support launches with gridDim.x * blockDim.x >= 2^32: +// depending on the ROCm version the launch returns +// hipErrorInvalidConfiguration or is accepted silently with the kernel +// never executing, leaving zero-initialized output. A grid-stride kernel +// with a fixed grid sized to device occupancy avoids the limit. +template +C10_LAUNCH_BOUNDS_1(num_threads()) +__global__ void elementwise_kernel_with_index_grid_stride( + index_t N, func_t f, + typename function_traits::result_type *data) { + index_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const index_t stride = static_cast(gridDim.x) * blockDim.x; + for (; idx < N; idx += stride) { + data[idx] = f(idx); + } +} +#endif + template void gpu_kernel_with_index(at::Tensor &output, func_t f) { int64_t N = output.numel(); if (N == 0) { return; } +#if defined(USE_ROCM) + constexpr int blocks_per_sm = 4; + const int sm_count = + at::cuda::getCurrentDeviceProperties()->multiProcessorCount; + const int64_t orig_grid = (N + block_work_size - 1) / block_work_size; + int64_t grid = std::min( + orig_grid, static_cast(sm_count) * blocks_per_sm); + grid = std::max(grid, 1); + auto stream = at::cuda::getCurrentCUDAStream(); + using scalar_t = typename function_traits::result_type; + if (N <= std::numeric_limits::max()) { + elementwise_kernel_with_index_grid_stride<<>>(N, f, output.mutable_data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } else { + elementwise_kernel_with_index_grid_stride<<>>(N, f, output.mutable_data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } +#else int64_t grid = (N + block_work_size - 1) / block_work_size; auto stream = at::cuda::getCurrentCUDAStream(); using scalar_t = typename function_traits::result_type; @@ -64,6 +104,7 @@ void gpu_kernel_with_index(at::Tensor &output, func_t f) { elementwise_kernel_with_index<<>>(N, f, output.mutable_data_ptr()); C10_CUDA_KERNEL_LAUNCH_CHECK(); } +#endif } } // namespace diff --git a/test/test_tensor_creation_ops.py b/test/test_tensor_creation_ops.py index 87e7db57318b6..71ebe5b41aa01 100644 --- a/test/test_tensor_creation_ops.py +++ b/test/test_tensor_creation_ops.py @@ -2789,6 +2789,24 @@ def test_range_factories_64bit_indexing(self, device): self.assertEqual(t[-1].item(), 2) del t + # On ROCm, launches with gridDim.x * blockDim.x >= 2^32 are not + # supported and either return hipErrorInvalidConfiguration or fail + # silently. Exercise the just-over case (~16 GB at int32) and a + # far-above case (~33 GB) when memory permits. arange computes + # int64 values then casts down, so for int32 the trailing + # 2^32 - 2, 2^32 - 1, 2^32 wrap to -2, -1, 0. + if TEST_WITH_ROCM: + for bigint in (2 ** 32 + 1, 2 ** 33 + 1): + free, _ = torch.cuda.mem_get_info(device) + if free < bigint * 4 + (3 << 30): + continue + t = torch.arange(bigint, dtype=torch.int32, device=device) + self.assertEqual(t.numel(), bigint) + self.assertEqual( + t[-3:].cpu(), torch.tensor([-2, -1, 0], dtype=torch.int32) + ) + del t + @expectedFailureMeta # RuntimeError: The tensor has a non-zero number of elements @onlyNativeDeviceTypes def test_tensor_ctor_device_inference(self, device): From 1e7198817c2bfb770a1a498df5a260b7883d3e2e Mon Sep 17 00:00:00 2001 From: "Cao, Glen" Date: Tue, 26 May 2026 13:35:23 -0700 Subject: [PATCH 04/11] Temp 2.11 1250 tdm (#20) * TDM on release/2.11 for bring-up based on careful selection * Triton commit: Upstream fe0c38b5262c0447fed6df0d37e02cb8ea75deb4 -> AMD-ROCm-Internal Triton 250bb5d5b821377f49dc2d83d87ded75b952f0f7; Consequence: Triton TDM support may miss. * Refinement according to reviewers' comments * Added/modified UT cases; NUM_STAGES issue of ineffectiveness * A couple of changes to related UTs * Got rid of configs like `waves_per_cu=2` --- .ci/pytorch/build.sh | 20 +- .github/workflows/inductor-rocm-gfx1250.yml | 77 ++++++++ test/inductor/test_flex_attention.py | 14 ++ test/inductor/test_max_autotune.py | 117 ++++++++++++ .../_inductor/codegen/rocm/compile_command.py | 4 + torch/_inductor/config.py | 23 +++ torch/_inductor/kernel/flex/common.py | 7 +- torch/_inductor/kernel/flex/flex_attention.py | 18 +- torch/_inductor/kernel/flex/flex_decoding.py | 12 +- torch/_inductor/kernel/mm.py | 71 ++++++-- torch/_inductor/template_heuristics/triton.py | 171 +++++++++++++++++- torch/_inductor/utils.py | 112 ++++++++++++ 12 files changed, 628 insertions(+), 18 deletions(-) create mode 100644 .github/workflows/inductor-rocm-gfx1250.yml diff --git a/.ci/pytorch/build.sh b/.ci/pytorch/build.sh index 265a6f579f427..0334a557634a7 100755 --- a/.ci/pytorch/build.sh +++ b/.ci/pytorch/build.sh @@ -110,6 +110,7 @@ if [[ "$BUILD_ENVIRONMENT" == *vulkan* ]]; then source /var/lib/jenkins/vulkansdk/setup-env.sh fi +# Example BUILD_ENVIRONMENT: linux-noble-rocm-py3.12-gfx1250 if [[ "$BUILD_ENVIRONMENT" == *rocm* ]]; then # hcc used to run out of memory, silently exiting without stopping # the build process, leaving undefined symbols in the shared lib, @@ -119,10 +120,23 @@ if [[ "$BUILD_ENVIRONMENT" == *rocm* ]]; then export MAX_JOBS=$(($(nproc) - 1)) fi + # Logic for multiple architectures based on the discriminator BUILD_ENVIRONMENT + # that is set by the workflow YAML and follows a consistent naming pattern. if [[ -n "$CI" && -z "$PYTORCH_ROCM_ARCH" ]]; then - # Set ROCM_ARCH to gfx906 for CI builds, if user doesn't override. - echo "Limiting PYTORCH_ROCM_ARCH to gfx906 for CI builds" - export PYTORCH_ROCM_ARCH="gfx906" + if [[ "$BUILD_ENVIRONMENT" == *gfx1250* ]]; then + echo "Setting PYTORCH_ROCM_ARCH to gfx1250 for CI builds" + export PYTORCH_ROCM_ARCH="gfx1250" + elif [[ "$BUILD_ENVIRONMENT" == *mi355* ]] || [[ "$BUILD_ENVIRONMENT" == *gfx950* ]]; then + echo "Setting PYTORCH_ROCM_ARCH to gfx950 for CI builds" + export PYTORCH_ROCM_ARCH="gfx950" + elif [[ "$BUILD_ENVIRONMENT" == *mi300* ]] || [[ "$BUILD_ENVIRONMENT" == *gfx942* ]]; then + echo "Setting PYTORCH_ROCM_ARCH to gfx942 for CI builds" + export PYTORCH_ROCM_ARCH="gfx942" + else + # Set ROCM_ARCH to gfx906 for CI builds, if user doesn't override. + echo "Limiting PYTORCH_ROCM_ARCH to gfx906 for CI builds" + export PYTORCH_ROCM_ARCH="gfx906" + fi fi # hipify sources diff --git a/.github/workflows/inductor-rocm-gfx1250.yml b/.github/workflows/inductor-rocm-gfx1250.yml new file mode 100644 index 0000000000000..33449a7d52bf1 --- /dev/null +++ b/.github/workflows/inductor-rocm-gfx1250.yml @@ -0,0 +1,77 @@ +# The name of this file is subject to change to stay consistent with other .yml files. +# +# The MI355 workflow (.github/workflows/inductor-rocm-mi355.yml) uses: +# - _linux-build.yml and _rocm-test.yml reusable workflows +# - Build environment linux-noble-rocm-py3.12-mi355 +# - Runner label linux.rocm.gpu.gfx950.1 +# - Docker image ci-image:pytorch-linux-noble-rocm-n-py3 +# - 2-shard test matrix for the inductor config +# +# The GFX1250 equivalent is following this exact pattern. + +name: inductor-rocm-gfx1250 + +on: + push: + branches: + - main + - release/* + tags: + - ciflow/inductor-rocm-gfx1250/* + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' }}-${{ github.event_name == 'schedule' }} + cancel-in-progress: true + +permissions: + id-token: write + contents: read + actions: read + +jobs: + target-determination: + if: github.repository_owner == 'pytorch' + name: before-test + uses: ./.github/workflows/target_determination.yml + + get-label-type: + name: get-label-type + uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.11 + if: ${{ (github.event_name != 'schedule' || github.repository == 'pytorch/pytorch') && github.repository_owner == 'pytorch' }} + with: + triggering_actor: ${{ github.triggering_actor }} + issue_owner: ${{ github.event.pull_request.user.login || github.event.issue.user.login }} + curr_branch: ${{ github.head_ref || github.ref_name }} + curr_ref_type: ${{ github.ref_type }} + opt_out_experiments: lf + + linux-noble-rocm-py3_12-inductor-build: + name: linux-noble-rocm-py3.12-gfx1250 + uses: ./.github/workflows/_linux-build.yml + needs: get-label-type + with: + runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" + build-environment: linux-noble-rocm-py3.12-gfx1250 + # Docker image stays the same as MI355 because ROCm image supports multiple arches. + docker-image-name: ci-image:pytorch-linux-noble-rocm-n-py3 + # Set PYTORCH_ROCM_ARCH directly in the workflow YAML as an env variable, + # so build.sh never needs to parse BUILD_ENVIRONMENT. + #env-var-script: | + # export PYTORCH_ROCM_ARCH=gfx1250 + test-matrix: | + { include: [ + { config: "inductor", shard: 1, num_shards: 2, runner: "linux.rocm.gpu.gfx1250.1" }, # It requires provisioning hardware. + { config: "inductor", shard: 2, num_shards: 2, runner: "linux.rocm.gpu.gfx1250.1" }, + ]} + secrets: inherit + + linux-noble-rocm-py3_12-inductor-test: + name: linux-noble-rocm-py3.12-gfx1250 + uses: ./.github/workflows/_rocm-test.yml + needs: linux-noble-rocm-py3_12-inductor-build + with: + build-environment: ${{ needs.linux-noble-rocm-py3_12-inductor-build.outputs.build-environment }} + docker-image: ${{ needs.linux-noble-rocm-py3_12-inductor-build.outputs.docker-image }} + test-matrix: ${{ needs.linux-noble-rocm-py3_12-inductor-build.outputs.test-matrix }} + secrets: inherit diff --git a/test/inductor/test_flex_attention.py b/test/inductor/test_flex_attention.py index d4bab42d80173..934e3b6668937 100644 --- a/test/inductor/test_flex_attention.py +++ b/test/inductor/test_flex_attention.py @@ -467,6 +467,20 @@ def batch_reserve(paged_attention: PagedAttention, target_seq_len: Tensor): ) +class TestFlexAttentionTDMOptions(InductorTestCase): + def test_apply_tdm_num_stages_uses_triton_launch_option(self): + from torch._inductor.kernel.flex.common import apply_tdm_num_stages + + kernel_options = {"num_stages": 1, "NUM_STAGES": 4} + + apply_tdm_num_stages(kernel_options) + + self.assertEqual( + kernel_options["num_stages"], config.tdm.max_outstanding_per_wave + ) + self.assertNotIn("NUM_STAGES", kernel_options) + + @large_tensor_test_class("2GB", device=test_device[0]) class TestFlexAttention(InductorTestCase): def setUp(self): diff --git a/test/inductor/test_max_autotune.py b/test/inductor/test_max_autotune.py index 58f1ae439abdf..9526272c2eeb2 100644 --- a/test/inductor/test_max_autotune.py +++ b/test/inductor/test_max_autotune.py @@ -544,6 +544,123 @@ def mock_get_tma_workspace_arg(*args, **kwargs): mm_tma_heuristic.mm_configs = original_tma_configs mm_heuristic.mm_configs = original_mm_configs + def test_tdm_arch_gate_accepts_only_gfx1250(self): + from torch._inductor.utils import is_gfx1250_arch + + self.assertTrue(is_gfx1250_arch("gfx1250")) + self.assertTrue(is_gfx1250_arch("gfx1250:sramecc+:xnack-")) + self.assertFalse(is_gfx1250_arch("gfx1251")) + self.assertFalse(is_gfx1250_arch("gfx1260")) + self.assertFalse(is_gfx1250_arch("amd-gfx1250")) + + def test_tdm_persistent_template_precedes_rocm_tma_fallback(self): + from torch._inductor.kernel import mm as mm_kernel + + for persistent_tma_enabled in (False, True): + templates = [] + with ( + config.patch( + { + "triton.enable_persistent_tma_matmul": persistent_tma_enabled + } + ), + mock.patch.object( + mm_kernel, + "use_triton_blackwell_tma_template", + return_value=False, + ), + mock.patch.object( + mm_kernel, + "use_triton_tdm_template", + return_value=True, + ), + mock.patch.object( + mm_kernel, + "use_triton_tma_template", + side_effect=AssertionError("TMA fallback should not run"), + ), + ): + selected = mm_kernel._append_persistent_mm_template( + templates, mock.Mock(), mock.Mock(), mock.Mock() + ) + + self.assertEqual(selected, "tdm") + self.assertEqual(templates, [mm_kernel.persistent_mm_template]) + + def test_tdm_template_add_guards_checks_compile_time_device(self): + from torch._inductor.utils import use_triton_tdm_template + + class FakeMatrix: + def __init__(self, device): + self.device = device + + def get_device(self): + return self.device + + def get_dtype(self): + return torch.float16 + + def get_stride(self): + return (128, 1) + + mat = FakeMatrix(torch.device("cuda", 0)) + with ( + config.patch({"enable_tdm_configs": True}), + mock.patch.object(torch.version, "hip", "7.2.0"), + mock.patch.object( + torch.cuda, + "get_device_properties", + return_value=mock.Mock(gcnArchName="gfx1250"), + ), + ): + self.assertTrue( + use_triton_tdm_template( + mat, + output_layout=mock.Mock(device=torch.device("cuda", 0)), + add_guards=True, + ) + ) + self.assertFalse( + use_triton_tdm_template( + mat, + output_layout=mock.Mock(device=torch.device("cuda", 1)), + add_guards=True, + ) + ) + + def test_tdm_block_k_filter_is_dtype_size_aware(self): + from torch._inductor.template_heuristics.triton import ( + _filter_tdm_block_k_configs, + ROCmGemmConfig, + ) + + configs = [ + ROCmGemmConfig(128, 64, 64, 4, 4, group_m=8), + ROCmGemmConfig(128, 64, 128, 4, 4, group_m=8), + ] + + self.assertEqual( + [c.block_k for c in _filter_tdm_block_k_configs(configs, 2)], + [64, 128], + ) + self.assertEqual( + [c.block_k for c in _filter_tdm_block_k_configs(configs, 1)], + [128], + ) + self.assertEqual( + [c.block_k for c in _filter_tdm_block_k_configs(configs, 4)], + [64, 128], + ) + + def test_shared_memory_estimation_counts_num_stages_once(self): + heuristic = ROCmMMTemplateConfigHeuristic() + gemm_config = GemmConfig(128, 64, 64, 4, 4, group_m=8) + + self.assertEqual( + heuristic.get_shared_memory_estimation(gemm_config, 2, False, 0), + (128 * 64 + 64 * 64) * 2 * 4 + 128, + ) + @unittest.skipIf( not has_triton_tma_device(), "Need device-side TMA support in Triton" ) diff --git a/torch/_inductor/codegen/rocm/compile_command.py b/torch/_inductor/codegen/rocm/compile_command.py index 2aa4494e0aa32..b44638879d512 100644 --- a/torch/_inductor/codegen/rocm/compile_command.py +++ b/torch/_inductor/codegen/rocm/compile_command.py @@ -74,6 +74,10 @@ def _rocm_lib_options(dst_file_ext: str) -> list[str]: def _rocm_compiler_options() -> list[str]: + # `config.rocm.arch`is populated from either: + # - The `PYTORCH_ROCM_ARCH` environment variable, or + # - Runtime device detection. + # The string "native" tells `hipcc` to compile for the current GPU. arch_list = config.rocm.arch or ["native"] gpu_arch_flags = [f"--offload-arch={arch}" for arch in arch_list] opts = [ diff --git a/torch/_inductor/config.py b/torch/_inductor/config.py index a3dbe5ae9d984..78db8b74715dc 100644 --- a/torch/_inductor/config.py +++ b/torch/_inductor/config.py @@ -822,6 +822,18 @@ def use_autoheuristic(name: str) -> bool: os.environ.get("TORCHINDUCTOR_CACHE_SDPA_CONSTRAINT", "0") == "1" ) +# AMD TDM config flag for TDM support for gfx1250, is a configuration gate +# rather than something that causes TDM instructions to be emitted directly. +# It's set to control if gfx1250-specific autotuning configs +# (with larger tile sizes, more pipeline stages) are included +# in the candidate set during `max_autotune`. +# When enabled and running on a gfx1250 device, Inductor may emit Triton kernels +# that leverage TDM for asynchronous global->LDS tensor tile copies. +# Requires Triton >= 3.6.0 with AMD TDM backend support. +# Disabled by default on ROCm; TDM is further gated by device arch (gfx1250) +# in `use_triton_tdm_template()` at codegen time. +# Set to "0" to disable TDM even on gfx1250 hardware. +enable_tdm_configs = os.environ.get("TORCHINDUCTOR_ENABLE_TDM_CONFIGS", "0") == "1" # Whether to keep the output strides the same as eager after layout optimization. keep_output_stride = os.environ.get("TORCHINDUCTOR_KEEP_OUTPUT_STRIDE", "1") == "1" @@ -2443,6 +2455,17 @@ class rocm: contiguous_threshold: int = 16 +class tdm: + # Maximum outstanding TDM address translations per wave is 4. + # For small tiles (e.g., 128x64 FP16), this is the binding constraint. + # For larger tiles with 2 waves/SIMD, the SIMD limit of 6 applies. + max_outstanding_per_wave: int = 4 + max_outstanding_per_simd: int = 6 + # TDM requires 128B/256B aligned contiguous regions + # in both global memory and LDS. + alignment_bytes: int = 128 + + # Backend to use for CPU codegen either "cpp" or "triton" (experimental) or "halide" (experimental) or "pallas" (experimental) cpu_backend: Literal["cpp", "triton", "halide", "pallas"] = "cpp" diff --git a/torch/_inductor/kernel/flex/common.py b/torch/_inductor/kernel/flex/common.py index 1ad6e50239170..6b37086e2e538 100644 --- a/torch/_inductor/kernel/flex/common.py +++ b/torch/_inductor/kernel/flex/common.py @@ -43,12 +43,17 @@ to_dtype, ) from ...select_algorithm import realize_inputs -from ...utils import load_template +from ...utils import config, load_template SubgraphResults = list[ComputedBuffer | None] | ComputedBuffer | None +def apply_tdm_num_stages(kernel_options: dict[str, Any]) -> None: + kernel_options["num_stages"] = config.tdm.max_outstanding_per_wave + kernel_options.pop("NUM_STAGES", None) + + def zeros_and_scatter_lowering(shape: list[int], indices, values): """To support backwards on captured buffers we register a specific lowering for our specific custom up""" # Always accumulate into fp32 then cast diff --git a/torch/_inductor/kernel/flex/flex_attention.py b/torch/_inductor/kernel/flex/flex_attention.py index 8c1523591c1ef..fb1c71c98a973 100644 --- a/torch/_inductor/kernel/flex/flex_attention.py +++ b/torch/_inductor/kernel/flex/flex_attention.py @@ -24,8 +24,9 @@ SymbolicGridFn, TritonTemplate, ) -from ...utils import can_use_tma +from ...utils import can_use_tma, use_triton_tdm_template from .common import ( + apply_tdm_num_stages, build_subgraph_buffer, create_indices_fake, create_num_blocks_fake_generator, @@ -418,6 +419,15 @@ def flex_attention( if cur_kernel_options["USE_TMA"] and not can_use_tma(query, key, value): cur_kernel_options["USE_TMA"] = False + # For gfx1250 TDM, ensure the non-TMA path uses enough pipeline stages + # to trigger TDM async copies in Triton's AMD backend (TTGIR pipelining pass). + # Standard attention block sizes (64, 128, 256) with FP16/BF16 + # produce 128B aligned tiles and are compatible. + if not cur_kernel_options["USE_TMA"] and use_triton_tdm_template( + query, key, value + ): + apply_tdm_num_stages(cur_kernel_options) + cur_kernel_options.setdefault("BLOCK_M", conf.block_m) cur_kernel_options.setdefault("BLOCK_N", conf.block_n) # Blocksparse options @@ -944,6 +954,12 @@ def flex_attention_backward(*args, **kwargs): if cur_kernel_options["USE_TMA"] and not can_use_tma(query, key, value): cur_kernel_options["USE_TMA"] = False + # See the comments at the corresponding place in function `flex_attention` above. + if not cur_kernel_options["USE_TMA"] and use_triton_tdm_template( + query, key, value + ): + apply_tdm_num_stages(cur_kernel_options) + cur_kernel_options.setdefault("BLOCK_M1", conf.block_m1) cur_kernel_options.setdefault("BLOCK_N1", conf.block_n1) cur_kernel_options.setdefault("BLOCK_M2", conf.block_m2) diff --git a/torch/_inductor/kernel/flex/flex_decoding.py b/torch/_inductor/kernel/flex/flex_decoding.py index 2cd25c5fa1063..5a7990ee7ba72 100644 --- a/torch/_inductor/kernel/flex/flex_decoding.py +++ b/torch/_inductor/kernel/flex/flex_decoding.py @@ -19,8 +19,9 @@ SymbolicGridFn, TritonTemplate, ) -from ...utils import can_use_tma +from ...utils import can_use_tma, use_triton_tdm_template from .common import ( + apply_tdm_num_stages, create_indices_fake, create_num_blocks_fake_generator, freeze_irnodes, @@ -356,6 +357,15 @@ def create_flex_decoding_kernel(*args, **kwargs): if cur_kernel_options["USE_TMA"] and not can_use_tma(query, key, value): cur_kernel_options["USE_TMA"] = False + # For gfx1250 TDM, ensure the non-TMA path uses enough pipeline stages + # to trigger TDM async copies in Triton's AMD backend (TTGIR pipelining pass). + # Standard attention block sizes (64, 128, 256) with FP16/BF16 + # produce 128B aligned tiles and are compatible. + if not cur_kernel_options["USE_TMA"] and use_triton_tdm_template( + query, key, value + ): + apply_tdm_num_stages(cur_kernel_options) + # Add ROCm-specific parameters if they exist in the config for attrib in ["kpack", "matrix_instr_nonkdim", "waves_per_eu"]: if hasattr(conf, attrib): diff --git a/torch/_inductor/kernel/mm.py b/torch/_inductor/kernel/mm.py index 13653d41e66d5..0eecc5524130e 100644 --- a/torch/_inductor/kernel/mm.py +++ b/torch/_inductor/kernel/mm.py @@ -54,6 +54,7 @@ use_triton_scaling_template, use_triton_template, use_triton_tma_template, + use_triton_tdm_template, ) from .mm_common import ( _is_static_problem, @@ -129,6 +130,59 @@ ) +def _append_persistent_mm_template( + templates_to_use: list[ExternKernelChoice | KernelTemplate], + mat1: Buffer, + mat2: Buffer, + layout: Layout, +) -> str | None: + if use_triton_blackwell_tma_template( + mat1, mat2, output_layout=layout, add_guards=True + ): + templates_to_use.append(blackwell_ws_persistent_device_tma_mm_template) + return "blackwell_tma" + if use_triton_tdm_template(mat1, mat2, output_layout=layout, add_guards=True): + # GFX1250 TDM: use the non-TMA persistent template. Triton's AMD backend + # automatically inserts TDM instructions when compiling with num_stages > 1. + # + # TDM hardware constraints (gfx1250): + # Documentation-only for now because Triton's pipeliner handles the + # wave assignment internally. If Triton exposes wave-specialization + # knobs in the future, these comments identify where to hook them. + # - 1 TDM unit per SIMD pair (2 SIMDs share 1 TDM unit) + # - Max 4 outstanding TDM address translations per wave + # - Max 6 outstanding TDM address translations per SIMD + # - Recommended: waves specialized to load A or B + # - 8 waves/WG (num_warps=8): 4 waves for A, 4 for B + # - 4 waves/WG (num_warps=4): 2 waves for A, 2 for B + # - All waves can load both A & B. This is preferable for uniform + # unrolling across CUs and helps multicast load latency. + # - Do not interleave different TDMs within a wave. Complete one TDM + # instruction's unrolling before switching. + # + # The persistent template heuristic supplies gfx1250 TDM configs. + # Key constraints: + # - TDM requests target 128B or 256B aligned contiguous regions in both + # global memory and LDS. + # - For FP16: BLOCK_K must be a multiple of 64 (64 * 2B = 128B). + # - For FP8/FP4: BLOCK_K must be a multiple of 128 (128 * 1B = 128B). + # - For FP32: BLOCK_K must be a multiple of 32 (32 * 4B = 128B). + # - BLOCK_M and BLOCK_N should also produce 128B-aligned LDS rows. + # TDM is most beneficial for small tiles such as 128x64/64x128 FP16 + # and 128x128/128x256 FP8 with num_stages=4 (matching the + # 4-outstanding-per-wave TDM address translation limit) and + # num_warps=4 or 8 (for wave-specialized A/B loading). + templates_to_use.append(persistent_mm_template) + return "tdm" + if use_triton_tma_template(mat1, mat2, output_layout=layout, add_guards=True): + if torch.version.hip is None: + templates_to_use.append(persistent_tma_mm_template) + else: + templates_to_use.append(persistent_mm_template) + return "tma" + return None + + # prevent duplication registration of extern functions @functools.cache def lazy_register_extern_choice(fn): @@ -426,8 +480,11 @@ def _to_dtype(x): if is_exhaustive or not use_decompose_k_choice(m, n, k, threshold_multiple=2): templates_to_use.append(mm_template) - if use_triton_blackwell_tma_template( - mat1, mat2, output_layout=layout, add_guards=True + _append_persistent_mm_template(templates_to_use, mat1, mat2, layout) + + if ( + inductor_config.is_fbcode() + and inductor_config.triton.enable_tlx_templates ): templates_to_use.append(blackwell_ws_persistent_device_tma_mm_template) elif use_triton_tma_template( @@ -685,15 +742,7 @@ def tuned_addmm(inp, mat1, mat2, *, alpha=1, beta=1, layout=None): if is_nonzero and use_triton_template(layout, check_max_autotune=False): templates_to_use.append(mm_template) - if use_triton_blackwell_tma_template( - mat1, mat2, output_layout=layout, add_guards=True - ): - templates_to_use.append(blackwell_ws_persistent_device_tma_mm_template) - elif use_triton_tma_template(mat1, mat2, output_layout=layout, add_guards=True): - if torch.version.hip is None: - templates_to_use.append(persistent_tma_mm_template) - else: - templates_to_use.append(persistent_mm_template) + _append_persistent_mm_template(templates_to_use, mat1, mat2, layout) templates_to_use.append(addmm_contiguous_subgraph_template) diff --git a/torch/_inductor/template_heuristics/triton.py b/torch/_inductor/template_heuristics/triton.py index 641c8568fb298..0eb18f5265968 100644 --- a/torch/_inductor/template_heuristics/triton.py +++ b/torch/_inductor/template_heuristics/triton.py @@ -35,6 +35,7 @@ get_default_kpack, get_num_sms, get_tma_workspace_arg, + is_gfx1250_arch, TMA_DESCRIPTOR_SIZE, using_b200, ) @@ -52,6 +53,42 @@ from torch._inductor.runtime.triton_compat import Config as TritonConfig +def _is_gfx1250_device() -> bool: + """Detect whether the current device is AMD gfx1250. + + gfx1250 has 320 KB LDS per CU, 512 B/clk/CU LDS bandwidth (R/W combined), + and TDM for async descriptor-based global->LDS copies. When num_stages > 1, + Triton's AMD backend automatically converts tl.load into TDM instructions. + + TDM hardware constraints: + - 1 TDM unit per SIMD pair (2 SIMDs share 1 TDM unit) + - Max 4 outstanding TDM address translations per wave + - Max 6 outstanding TDM address translations per SIMD + - Requests target 128B or 256B aligned contiguous regions + in both global memory and LDS + """ + if not torch.version.hip: + return False + if not inductor_config.enable_tdm_configs: + return False + try: + props = torch.cuda.get_device_properties(torch.cuda.current_device()) + arch = getattr(props, "gcnArchName", "") + return is_gfx1250_arch(arch) + except Exception: + return False + + +def _filter_tdm_block_k_configs( + configs: list[BaseConfig], + dtype_size: int, +) -> list[BaseConfig]: + if dtype_size <= 0: + return configs + block_k_multiple = max(inductor_config.tdm.alignment_bytes // dtype_size, 1) + return [c for c in configs if c.block_k % block_k_multiple == 0] + + # Gemm Configs @dataclasses.dataclass class BaseConfig: @@ -1469,6 +1506,7 @@ def __init__(self) -> None: super().__init__() self.default_num_stages = get_backend_num_stages() + self.uses_tdm_configs = False self.mm_configs: list[BaseConfig] = [ ROCmGemmConfig( @@ -1527,6 +1565,42 @@ def __init__(self) -> None: ROCmGemmConfig(256, 256, 64, self.default_num_stages, 8, group_m=4), ] + # TDM-optimized persistent MM configs for gfx1250. + # + # Hardware constraints on gfx1250: + # - 320 KB LDS per CU, 512 B/clk/CU bandwidth + # - Max 4 outstanding TDM translations per wave (bounds num_stages) + # - Max 6 outstanding TDM translations per SIMD + # - TDM requests target 128B/256B aligned contiguous regions + # - 1 TDM unit per SIMD pair -> recommended 1 wave/SIMD pair issuing + # TDM instructions (num_warps=4 gives 1 wave/SIMD on a 4-SIMD CU) + # + # BLOCK_K alignment requirements (128B / dtype_bytes): + # FP16/BF16 (2B): BLOCK_K multiple of 64 + # FP8/INT8 (1B): BLOCK_K multiple of 128 + # FP32 (4B): BLOCK_K multiple of 32 + # + # LDS usage estimate: (BM*BK + BK*BN)*dtype_size*num_stages + # Must stay under 320 KB. + self.tdm_persistent_mm_configs: list[BaseConfig] = [ + # Small tiles: TDM most beneficial here (LDS-BW-limited to 1 wave/SIMD). + # These are where TDM eliminates cluster load issue overhead. + # `waves_per_eu=0` means Triton compiler should figure out the best value. + ROCmGemmConfig(128, 64, 64, 4, 4, group_m=8, waves_per_eu=0), + ROCmGemmConfig( 64, 128, 64, 4, 4, group_m=8, waves_per_eu=0), + ROCmGemmConfig(128, 64, 128, 4, 4, group_m=8), + ROCmGemmConfig( 64, 128, 128, 4, 4, group_m=8), + # Medium tiles: benefit from TDM prologue overhead reduction. + ROCmGemmConfig(128, 128, 64, 4, 4, group_m=8), + ROCmGemmConfig(128, 128, 64, 4, 8, group_m=16), + ROCmGemmConfig(128, 128, 128, 4, 4, group_m=8), + ROCmGemmConfig(128, 128, 128, 3, 8, group_m=16), + # Larger tiles: exploit the 320 KB LDS with deep pipelines. + ROCmGemmConfig(256, 128, 64, 4, 8, group_m=16), + ROCmGemmConfig(128, 256, 64, 4, 8, group_m=16), + ROCmGemmConfig(256, 128, 128, 3, 8, group_m=16), + ] + # Exhaustive search for mm configs self.exhaustive_configs: list[BaseConfig] = [ ROCmGemmConfig( @@ -1662,12 +1736,57 @@ def _prune_exhaustive_configs( ] return pruned_configs + def preprocess_mm_configs( + self, + m: int, + n: int, + k: int, + configs: list[BaseConfig], + has_int8_tensor: bool = False, + scale: float = 1.0, + exclude: Callable[ + [sympy.Integer, sympy.Integer, sympy.Integer], bool + ] = lambda m, n, k: False, + dtype_size: int = 0, + op_name: str = "mm", + **kwargs, + ) -> Generator[TritonConfig, None, None]: + if self.uses_tdm_configs: + configs = _filter_tdm_block_k_configs(configs, dtype_size) + return super().preprocess_mm_configs( + m, + n, + k, + configs, + has_int8_tensor, + scale, + exclude, + dtype_size, + op_name, + **kwargs, + ) + def _filter_configs(self, configs: list[BaseConfig]) -> list[BaseConfig]: """ ROCm specific filtering + + Normally we force num_stages to `default_num_stages` because AMD's + Triton backend historically didn't benefit from multi-stage pipelining. + However, on gfx1250 with TDM enabled, num_stages > 1 is exactly + what triggers TDM async copies in the StreamPipeliner pass, so we + must preserve the requested num_stages for TDM configs. """ + tdm_enabled = _is_gfx1250_device() for c in configs: - c.num_stages = self.default_num_stages + # On gfx1250, preserve num_stages as provided (but cap at 4 to + # respect TDM's per-wave outstanding limit). Otherwise, fall + # back to the legacy behavior of forcing default_num_stages. + if tdm_enabled: + c.num_stages = ( + min(c.num_stages, 4) if c.num_stages > 1 else c.num_stages + ) + else: + c.num_stages = self.default_num_stages return super()._filter_configs(configs) def _finalize_mm_configs( @@ -3084,6 +3203,56 @@ def __init__(self) -> None: self.exhaustive_configs = self.mm_plus_mm_configs +# ROCm persistent MM template heuristic (non-TMA, standard pointer loads) + + +@register_template_heuristic( + persistent_mm_template.uid, + "cuda", + register=torch.version.hip is not None, +) +class PersistentMMTemplateConfigHeuristic( + MMTemplateConfigMixin, + ROCmConfigHeuristic, # type: ignore[misc] +): + """Persistent MM template heuristic (no TMA, standard pointer loads)""" + + def __init__(self) -> None: + super().__init__() + if _is_gfx1250_device(): + # On gfx1250, use TDM-optimized configs that exploit: + # - 320 KB LDS (larger tiles possible) + # - TDM async global->LDS copies (num_stages > 1) + # - Wave specialization across SIMD pairs + self.mm_configs = self.tdm_persistent_mm_configs + self.uses_tdm_configs = True + else: + self.mm_configs = self.persistent_mm_configs + + def _get_template_configs_impl( + self, + kernel_inputs: KernelInputs, + op_name: str, + **kwargs, + ) -> Generator[dict[str, Any], None, None]: + for template_kwargs in super()._get_template_configs_impl( + kernel_inputs, op_name, **kwargs + ): + yield {**template_kwargs, "NUM_SMS": get_num_sms()} + + +@register_template_heuristic( + persistent_mm_template.uid, + "cuda", + register=torch.version.hip is not None, + op_name="addmm", +) +class ROCmAddMMPersistentTemplateConfigHeuristic( + AddMMConfigMixin, PersistentMMTemplateConfigHeuristic +): + """Addmm specific mixin for persistent MM on ROCm""" + + # MTIA template-specific classes diff --git a/torch/_inductor/utils.py b/torch/_inductor/utils.py index 7dc7476f2c9ad..f7168a1101b75 100644 --- a/torch/_inductor/utils.py +++ b/torch/_inductor/utils.py @@ -2015,6 +2015,118 @@ def use_triton_tma_template( return can_use_tma(*matrices, output_layout=layout, add_guards=add_guards) +def is_gfx1250_arch(arch: str) -> bool: + """Return True only for the gfx1250 target, including feature-suffixed names.""" + return arch.split(":", 1)[0] == "gfx1250" + + +def use_triton_tdm_template(*matrices: IRNode, output_layout=None, add_guards=False): + """Coarsely check whether AMD TDM-optimized persistent template should be used. + Despite the fact that TDM has some restrictions on direct LDS loads ( + when unrolling requests, the requests are to 128B or 256B aligned and + continuous regions in global memory, and writing back to 128B or 256B + aligned and continuous regions in LDS), this gate doesn't check per-element + alignment, instead delegating to Triton compiler's AMD backend. + That said, for GeMM specifically, it checks the leading dimensions produce + 128B-aligned rows, which is the common case where TDM is beneficial. + + Args: + - output_layout: Output layout (used for device detection) + - add_guards: If True, add shape guards to the graph + Returns True when: + - Running on ROCm (torch.version.hip is not None) + - Device is gfx1250 + - config.enable_tdm_configs (defined elsewhere) is True + - Triton version supports gfx1250 TDM (>= 3.6.0) + """ + if not torch.version.hip: + return False + if not config.enable_tdm_configs: + return False + device = matrices[0].get_device() + if device.type != "cuda": + return False + try: + # props = torch.cuda.get_device_properties(torch.cuda.current_device()) + props = torch.cuda.get_device_properties(device) + arch = getattr(props, "gcnArchName", "") + if not is_gfx1250_arch(arch): + return False + except Exception: + return False + + # Check Triton version supports TDM (3.6.0+) + try: + from torch.torch_version import TorchVersion + import triton + if TorchVersion(triton.__version__) < "3.6.0": + return False + except (ImportError, Exception): + return False + + # TDM can lower the same matrix element types as standard Triton GeMM. + # The persistent-template config heuristic still filters BLOCK_K by dtype + # size so each TDM request covers an aligned 128B region. + # FIXME: torch.float4 + triton_tdm_dtypes = { + torch.float16, + torch.bfloat16, + torch.float32, + torch.float8_e4m3fn, + torch.float8_e5m2, + torch.float8_e4m3fnuz, + torch.float8_e5m2fnuz, + torch.int8, + } + # TDM unrolls requests to 128B or 256B aligned contiguous regions + # in both global memory and LDS. For GeMM, this means the leading + # dimension (stride) of each matrix should produce rows that start + # at 128B-aligned addresses. If strides are not 128B-aligned, TDM + # may still work (Triton's backend can fall back per-load), but + # performance benefit is reduced. We don't hard-gate on this because: + # 1. The autotuner will benchmark and reject TDM configs if they're slower. + # 2. Triton's pipeliner handles misaligned cases gracefully. + # 3. Most real GeMM workloads have contiguous row-major or col-major layouts. + # For now, only log for debugging purposes. May be upgraded to a hard gate later. + for mat in matrices: + mat_dtype = mat.get_dtype() + if mat_dtype not in triton_tdm_dtypes: + return False + strides = mat.get_stride() + dtype_bytes = mat_dtype.itemsize + # Check if the innermost stride is 1 (contiguous) and + # the outer stride produces 128B-aligned rows. + if len(strides) >= 2: + inner_stride = strides[-1] + outer_stride = strides[-2] + if hasattr(inner_stride, '__int__'): + inner_val = int(inner_stride) + if inner_val != 1: + # TODO: Should be a hard rejection, not just a warning. + log.debug( + "TDM: matrix has non-unit inner stride %d, " + "TDM requests may not be contiguous", + inner_val, + ) + if hasattr(outer_stride, '__int__'): + row_bytes = int(outer_stride) * dtype_bytes + if row_bytes % config.tdm.alignment_bytes != 0: + log.debug( + "TDM: matrix row stride %d bytes is not %dB-aligned, " + "TDM performance may be suboptimal", + row_bytes, + config.tdm.alignment_bytes, + ) + + if add_guards and output_layout is not None: + # `SizeVarAllocator` does not guard device properties. Keep this as a + # static consistency check for the layout/device pair used at compile time. + if (output_layout.device.index or 0) != (device.index or 0): + return False + + return True + + def use_triton_blackwell_tma_template( *matrices: IRNode, output_layout: Layout, add_guards: bool = False ) -> bool: From 07ec88f2f2542b9e782df480cdc6cfb78c473b8a Mon Sep 17 00:00:00 2001 From: rraminen_amdeng Date: Wed, 17 Jun 2026 03:54:47 +0000 Subject: [PATCH 05/11] Bug fix - Remove drop_seed, drop_offset as there are not used in the code --- .../ATen/native/transformers/hip/flash_attn/ck/mha_bwd_ck.hip | 1 - 1 file changed, 1 deletion(-) diff --git a/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_bwd_ck.hip b/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_bwd_ck.hip index aec411e9978d7..1fcf7a97c6cc7 100644 --- a/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_bwd_ck.hip +++ b/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_bwd_ck.hip @@ -400,7 +400,6 @@ mha_bwd_ck(const at::Tensor &dout, // batch_size x seqlen_q x auto gen = at::get_generator_or_default( std::nullopt, at::cuda::detail::getDefaultCUDAGenerator()); - uint64_t* drop_seed, drop_offset; int64_t counter_offset = batch_size * num_heads * ck_tile::get_warp_size(); std::pair drop_seed_offset = {nullptr,nullptr}; if(is_dropout) { From 873d4c5b100c9fc015841e42acabf01d4dc067e1 Mon Sep 17 00:00:00 2001 From: Prachi Gupta Date: Wed, 17 Jun 2026 12:06:23 -0500 Subject: [PATCH 06/11] [release/2.12] Fix all arch build (#3324) - Need to turn MSLK on for mi300 and mi350 - Need to turn CK off for gfx1250 ## Motivation ## Technical Details ## Test Plan ## Test Result ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. --- CMakeLists.txt | 4 ++-- aten/src/ATen/native/cuda/GroupedBlas.cpp | 2 ++ cmake/Dependencies.cmake | 8 ++++++++ 3 files changed, 12 insertions(+), 2 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 7530948ed4d2e..215009d868e9f 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -960,8 +960,8 @@ cmake_dependent_option( # gfx1250 for MSLK needs to be involved as well. IF(USE_ROCM AND ("gfx942" IN_LIST PYTORCH_ROCM_ARCH OR "gfx950" IN_LIST PYTORCH_ROCM_ARCH)) - message(WARNING "Setting USE_FBGEMM_GENAI for gfx942/gfx950 to ON by default, doing ROCM build") - set(USE_FBGEMM_GENAI_DEFAULT ON) + message(STATUS "Setting USE_MSLK for gfx942/gfx950 to ON by default, doing ROCM build") + set(USE_MSLK_DEFAULT ON) elseif(USE_CUDA AND "$ENV{TORCH_CUDA_ARCH_LIST}" MATCHES "10.0" AND CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.8 AND NOT WIN32) message(STATUS "Setting USE_MSLK to ON by default , doing CUDA build for SM100a") set(USE_MSLK_DEFAULT ON) diff --git a/aten/src/ATen/native/cuda/GroupedBlas.cpp b/aten/src/ATen/native/cuda/GroupedBlas.cpp index 7c779e6658a66..103363ebf7f43 100644 --- a/aten/src/ATen/native/cuda/GroupedBlas.cpp +++ b/aten/src/ATen/native/cuda/GroupedBlas.cpp @@ -721,7 +721,9 @@ std::optional out_dtype) { const auto out_dtype_ = _resolve_grouped_mm_out_dtype(mat_a, mat_b, out_dtype); Tensor out = create_grouped_gemm_output_tensor(mat_a, mat_b, offs, out_dtype_); if (use_fast_path) { +#if defined(USE_ROCM_CK_GEMM) at::hip::detail::group_gemm_ck(mat_a, mat_b, offs, bias, out); +#endif //USE_ROCM_CK_GEMM } else { _grouped_mm_fallback(mat_a, mat_b, offs, bias, out_dtype, out); } diff --git a/cmake/Dependencies.cmake b/cmake/Dependencies.cmake index 70c5d3de4cf23..266ca743ef662 100644 --- a/cmake/Dependencies.cmake +++ b/cmake/Dependencies.cmake @@ -985,6 +985,14 @@ if(USE_ROCM) message(INFO "Compiling with HIP for AMD.") caffe2_update_option(USE_ROCM ON) + # composable_kernel has no gfx1250 support, so its CK GEMM/SDPA kernels fail + # to compile for that arch. + if("gfx1250" IN_LIST PYTORCH_ROCM_ARCH) + message(WARNING "gfx1250 in PYTORCH_ROCM_ARCH: disabling USE_ROCM_CK_GEMM and USE_ROCM_CK_SDPA (composable_kernel lacks gfx1250 support)") + caffe2_update_option(USE_ROCM_CK_GEMM OFF) + caffe2_update_option(USE_ROCM_CK_SDPA OFF) + endif() + if(USE_NCCL AND NOT USE_SYSTEM_NCCL) message(INFO "Forcing USE_SYSTEM_NCCL to ON since it's required by using RCCL") caffe2_update_option(USE_SYSTEM_NCCL ON) From aeb64a7497d08b2da400801c9340834bd6bde3f1 Mon Sep 17 00:00:00 2001 From: Xinya Zhang Date: Wed, 17 Jun 2026 12:09:37 -0500 Subject: [PATCH 07/11] [release/2.12_gfx1250] Bump AOTriton to 0.12.50tp (#3313) # Bump to AOTriton 0.12.50tp Notable new features: Enable gfx1250 ## Features from AOTriton 0.12b Notable new features: * **BREAKING** Varlen LSE tensor shape changes to (H, Total_seqlen) * Support head_dim != head_dim_v * Support `use_deterministic_algorithims` * Support seqused_k in test/test_varlen_attention.py * gfx1100 and gfx1151 promoted out of experimental * Partial FAv3 support on gfx950 Bug Fixes: * GQA kernel failed to read bias tensor with the right offset. Known Issues * gfx950's Triton kernel has problem handling hdim=16's fwd, in addition to hdim=48/80's bwd. * Disables gfx90a's CK SDPA support due to GPU Segfault. --------- Co-authored-by: Claude Co-authored-by: Prachi Gupta --- aten/src/ATen/Context.cpp | 4 + .../native/transformers/cuda/attention.cu | 129 ++---- .../transformers/cuda/attention_backward.cu | 190 +++----- .../native/transformers/cuda/sdp_utils.cpp | 78 +++- .../transformers/hip/aotriton_adapter.h | 19 +- .../transformers/hip/aotriton_versions.h | 6 + .../hip/flash_attn/aot/mha_all_aot.hip | 417 +++++++----------- cmake/External/aotriton.cmake | 45 +- test/test_transformers.py | 8 - test/test_varlen_attention.py | 5 +- torch/nn/attention/varlen.py | 17 - 11 files changed, 365 insertions(+), 553 deletions(-) diff --git a/aten/src/ATen/Context.cpp b/aten/src/ATen/Context.cpp index ba63e73e1017c..e17352934f9ee 100644 --- a/aten/src/ATen/Context.cpp +++ b/aten/src/ATen/Context.cpp @@ -505,7 +505,11 @@ at::BlasBackend Context::blasPreferredBackend() { bool Context::ckSupported() { #ifdef USE_ROCM static const std::vector supported_archs = { +<<<<<<< xinyazhang/backport-aotriton-0.12b-2.12_gfx1250 + "gfx942", "gfx950" +======= "gfx90a", "gfx942", "gfx950", "gfx1250", +>>>>>>> release/2.12_gfx1250 }; for (auto index : c10::irange(detail::getCUDAHooks().deviceCount())) { if(!detail::getCUDAHooks().isGPUArch(supported_archs, index)) { diff --git a/aten/src/ATen/native/transformers/cuda/attention.cu b/aten/src/ATen/native/transformers/cuda/attention.cu index 69ecf31df0586..93466a28064f4 100644 --- a/aten/src/ATen/native/transformers/cuda/attention.cu +++ b/aten/src/ATen/native/transformers/cuda/attention.cu @@ -1563,6 +1563,10 @@ std::tuple _efficient_ // compute_logsumexp is false constexpr int kAlignLSE = 1; res = at::empty({B, M, num_heads, Kv}, query.options()); + // TODO: Use Compact Varlen LSE + // The current memory allocation is strictly larger than necessary + // (total_q <= max_seqlen_q * B) + // The problem is total_q is not available here. at::Tensor softmax_lse; logsumexp = at::empty( { B, num_heads, compute_logsumexp ? max_seqlen_q : 0}, @@ -1591,8 +1595,6 @@ std::tuple _efficient_ atomic_counter = at::zeros({1}, query.options().dtype(at::kInt)); } - using aotriton::v2::flash::attn_fwd; - using aotriton::v2::flash::attn_fwd_compact_varlen; using sdp::aotriton_adapter::mk_aotensor; using sdp::aotriton_adapter::mk_aoscalartensor; using sdp::aotriton_adapter::mk_philoxtensor; @@ -1608,92 +1610,47 @@ std::tuple _efficient_ auto offset_output = mk_philoxtensor(use_philox_state ? offset_t.data_ptr() : nullptr); auto persistent_counter = mk_atomictensor(is_causal ? atomic_counter.data_ptr() : nullptr); hipError_t err; // TODO: Error handling - if constexpr (AOTRITON_ALWAYS_V3_API) { // Better readability than nesting ifdef -#if AOTRITON_V3_API // if constexpr does not stop errors from undefined functions - using aotriton::v3::flash::CausalType; - using aotriton::v3::flash::VarlenType; - using aotriton::v3::flash::WindowValue; - aotriton::v3::flash::attn_fwd_params params; - params.Q = mk_aotensor(q_t, "q"); - params.K = mk_aotensor(k_t, "k"); - params.V = mk_aotensor(v_t, "v"); - params.Sm_scale = softmax_scale; - params.L = compute_logsumexp ? mk_aotensor<2>(softmax_lse, "M") : empty_t2; - params.Out = mk_aotensor(output_t, "Out"); - params.Max_seqlen_q = max_seqlen_q; // Unused if cu_seqlens_q is empty - params.Max_seqlen_k = max_seqlen_k; // Unused if cu_seqlens_k is empty - params.dropout_p = dropout_p; - params.philox_seed_ptr = seed; - params.philox_offset1 = offset1; - params.philox_offset2 = offset2; - params.philox_seed_output = seed_output; - params.philox_offset_output = offset_output; - params.encoded_softmax = mk_aotensor(softmax_fa_t, "encoded_softmax"); - params.persistent_atomic_counter = persistent_counter; - params.causal_type = is_causal ? CausalType::WindowedAttention : CausalType::None; - if (static_cast(sdp::CustomMaskType::CausalFromTopLeft) == custom_mask_type) { - params.window_left = WindowValue::TopLeftAligned; - params.window_right = WindowValue::TopLeftAligned; - } else if (static_cast(sdp::CustomMaskType::CausalFromBottomRight) == custom_mask_type) { - params.window_left = WindowValue::BottomRightAligned; - params.window_right = WindowValue::BottomRightAligned; - } - if (bias.has_value()) { - params.B = mk_aotensor(bias.value(), "bias"); - } - if (seqstart_q.has_value()) { - params.varlen_type = VarlenType::CompactVarlen; - params.cu_seqlens_q = mk_aotensor<1>(seqstart_q.value(), "cu_seqlens_q"); - params.cu_seqlens_k = mk_aotensor<1>(seqstart_k.value(), "cu_seqlens_k"); - } else { - params.varlen_type = VarlenType::None; - } - err = aotriton::v3::flash::attn_fwd(params, - aotriton::v3::flash::attn_fwd_params::kVersion, - stream); -#endif // AOTRITON_V3_API - } else if (seqstart_q.has_value()) { - // varlen aka nested tensor - err = attn_fwd_compact_varlen(mk_aotensor(q_t, "q"), - mk_aotensor(k_t, "k"), - mk_aotensor(v_t, "v"), - bias.has_value() ? mk_aotensor(bias.value(), "bias"): empty_t4, - mk_aotensor<1>(seqstart_q.value(), "cu_seqlens_q"), - mk_aotensor<1>(seqstart_k.value(), "cu_seqlens_k"), - max_seqlen_q, - max_seqlen_k, - softmax_scale, - compute_logsumexp ? mk_aotensor<2>(softmax_lse, "M") : empty_t2, - mk_aotensor(output_t, "Out"), - dropout_p, - seed, - offset1, - offset2, - seed_output, - offset_output, - mk_aotensor(softmax_fa_t, "encoded_softmax"), - is_causal, - persistent_counter, - stream); + using aotriton::v3::flash::CausalType; + using aotriton::v3::flash::VarlenType; + using aotriton::v3::flash::WindowValue; + aotriton::v3::flash::attn_fwd_params params; + params.Q = mk_aotensor(q_t, "q"); + params.K = mk_aotensor(k_t, "k"); + params.V = mk_aotensor(v_t, "v"); + params.Sm_scale = softmax_scale; + params.L = compute_logsumexp ? mk_aotensor<2>(softmax_lse, "M") : empty_t2; + params.Out = mk_aotensor(output_t, "Out"); + params.Max_seqlen_q = max_seqlen_q; // Unused if cu_seqlens_q is empty + params.Max_seqlen_k = max_seqlen_k; // Unused if cu_seqlens_k is empty + params.dropout_p = dropout_p; + params.philox_seed_ptr = seed; + params.philox_offset1 = offset1; + params.philox_offset2 = offset2; + params.philox_seed_output = seed_output; + params.philox_offset_output = offset_output; + params.encoded_softmax = mk_aotensor(softmax_fa_t, "encoded_softmax"); + params.persistent_atomic_counter = persistent_counter; + params.causal_type = is_causal ? CausalType::WindowedAttention : CausalType::None; + if (static_cast(sdp::CustomMaskType::CausalFromTopLeft) == custom_mask_type) { + params.window_left = WindowValue::TopLeftAligned; + params.window_right = WindowValue::TopLeftAligned; + } else if (static_cast(sdp::CustomMaskType::CausalFromBottomRight) == custom_mask_type) { + params.window_left = WindowValue::BottomRightAligned; + params.window_right = WindowValue::BottomRightAligned; + } + if (bias.has_value()) { + params.B = mk_aotensor(bias.value(), "bias"); + } + if (seqstart_q.has_value()) { + params.varlen_type = VarlenType::CompactVarlen; + params.cu_seqlens_q = mk_aotensor<1>(seqstart_q.value(), "cu_seqlens_q"); + params.cu_seqlens_k = mk_aotensor<1>(seqstart_k.value(), "cu_seqlens_k"); } else { - err = attn_fwd(mk_aotensor(q_t, "q"), - mk_aotensor(k_t, "k"), - mk_aotensor(v_t, "v"), - bias.has_value() ? mk_aotensor(bias.value(), "bias"): empty_t4, - softmax_scale, - compute_logsumexp ? mk_aotensor<2>(softmax_lse, "M") : empty_t2, - mk_aotensor(output_t, "Out"), - dropout_p, - seed, - offset1, - offset2, - seed_output, - offset_output, - mk_aotensor(softmax_fa_t, "encoded_softmax"), - is_causal, - persistent_counter, - stream); + params.varlen_type = VarlenType::None; } + err = aotriton::v3::flash::attn_fwd(params, + aotriton::v3::flash::attn_fwd_params::kVersion, + stream); #else TORCH_CHECK(false, "Attempting to use AOTriton mem_eff_forward backend in a build that has not built AOTriton"); #endif diff --git a/aten/src/ATen/native/transformers/cuda/attention_backward.cu b/aten/src/ATen/native/transformers/cuda/attention_backward.cu index 183f99e975cda..b7e52617a0ec6 100644 --- a/aten/src/ATen/native/transformers/cuda/attention_backward.cu +++ b/aten/src/ATen/native/transformers/cuda/attention_backward.cu @@ -547,6 +547,17 @@ _efficient_attention_backward( "[AOTriton] Accelerated SDPA only supports MI200/MI300X/7900XTX/9070XT GPUs" " (gfx90a/gfx942/gfx1100/gfx1201)") } + bool deterministic{false}; + auto& ctx = at::globalContext(); + if (ctx.deterministicAlgorithms()) { + if (ctx.deterministicAlgorithmsWarnOnly()) { + TORCH_WARN_ONCE( + "Memory Efficient attention defaults to a non-deterministic algorithm. ", + "To explicitly enable determinism call torch.use_deterministic_algorithms(True, warn_only=False)."); + } else { + deterministic = true; + } + } const auto softmax_scale = sdp::calculate_scale(query, scale).expect_float(); bool is_causal; if (static_cast(sdp::CustomMaskType::NoCustomMask) == custom_mask_type) { @@ -569,139 +580,60 @@ _efficient_attention_backward( at::Tensor dout_t = grad_out.permute({0,2,1,3}); at::Tensor softmax_lse = logsumexp.view({B * nH, max_seqlen_q}); hipError_t err; - using aotriton::v2::flash::attn_bwd; - using aotriton::v2::flash::attn_bwd_fused; - using aotriton::v2::flash::attn_bwd_compact_varlen; using sdp::aotriton_adapter::mk_aotensor; using sdp::aotriton_adapter::mk_aoscalartensor; using sdp::aotriton_adapter::cast_dtype; aotriton::TensorView<4> empty_t4(0, {0, 0, 0, 0}, {0, 0, 0, 0}, cast_dtype(query.dtype())); - if constexpr (AOTRITON_ALWAYS_V3_API) { // Better readability than nesting ifdef -#if AOTRITON_V3_API // if constexpr does not stop errors from undefined functions - using aotriton::v3::flash::CausalType; - using aotriton::v3::flash::VarlenType; - using aotriton::v3::flash::WindowValue; - aotriton::v3::flash::attn_bwd_params params; - params.Q = mk_aotensor(q_t, "q"); - params.K = mk_aotensor(k_t, "k"); - params.V = mk_aotensor(v_t, "v"); - params.B = bias.has_value() ? mk_aotensor(bias.value(), "bias") : empty_t4; - params.Sm_scale = softmax_scale; - params.Out = mk_aotensor(out_t, "out"); - params.DO = mk_aotensor(dout_t, "dout"); - params.DK = mk_aotensor(dk_t, "dk"); - params.DV = mk_aotensor(dv_t, "dv"); - params.DQ = mk_aotensor(dq_t, "dq"); - params.DB = bias_requires_grad ? mk_aotensor(grad_bias, "db") : empty_t4; - params.L = mk_aotensor<2>(softmax_lse, "L"); - params.Max_seqlen_q = max_seqlen_q; // Unused if cu_seqlens_q is empty - params.Max_seqlen_k = max_seqlen_k; // Unused if cu_seqlens_k is empty - params.dropout_p = float(dropout_p); - params.philox_seed_ptr = mk_aoscalartensor(philox_seed); - params.philox_offset1 = mk_aoscalartensor(philox_offset); - params.philox_offset2 = 0; - params.causal_type = is_causal ? CausalType::WindowedAttention : CausalType::None; - if (static_cast(sdp::CustomMaskType::CausalFromTopLeft) == custom_mask_type) { - params.window_left = WindowValue::TopLeftAligned; - params.window_right = WindowValue::TopLeftAligned; - } else if (static_cast(sdp::CustomMaskType::CausalFromBottomRight) == custom_mask_type) { - params.window_left = WindowValue::BottomRightAligned; - params.window_right = WindowValue::BottomRightAligned; - } -#if AOTRITON_ALWAYS_V3_API - using sdp::aotriton_adapter::mklazy_empty_like; - using sdp::aotriton_adapter::mklazy_fp32zeros; - using sdp::aotriton_adapter::LazyTensorContext; - LazyTensorContext lazy_delta { .like_tensor = softmax_lse, .tensor_name = "delta" }; - LazyTensorContext lazy_dq_acc { .like_tensor = dq_t, .tensor_name = "dq_acc" }; - params.D = mklazy_empty_like<2>(&lazy_delta); - params.DQ_ACC = mklazy_fp32zeros<4>(&lazy_dq_acc); -#else - at::Tensor delta = at::empty_like(softmax_lse).contiguous(); - params.D = mk_aotensor<2>(delta, "delta"); -#endif - if (cu_seqlens_q.has_value()) { - params.varlen_type = VarlenType::CompactVarlen; - params.cu_seqlens_q = mk_aotensor<1>(cu_seqlens_q.value(), "cu_seqlens_q"); - params.cu_seqlens_k = mk_aotensor<1>(cu_seqlens_k.value(), "cu_seqlens_k"); - } else { - params.varlen_type = VarlenType::None; - } - err = aotriton::v3::flash::attn_bwd(params, - aotriton::v3::flash::attn_bwd_params::kVersion, - stream); -#endif // AOTRITON_V3_API - } else if (cu_seqlens_q.has_value()) { - at::Tensor delta = at::empty_like(softmax_lse).contiguous(); - // varlen aka Nested tensor - err = attn_bwd_compact_varlen(mk_aotensor(q_t, "q"), - mk_aotensor(k_t, "k"), - mk_aotensor(v_t, "v"), - mk_aotensor<1>(cu_seqlens_q.value(), "cu_seqlens_q"), - mk_aotensor<1>(cu_seqlens_k.value(), "cu_seqlens_k"), - max_seqlen_q, - max_seqlen_k, - bias.has_value() ? mk_aotensor(bias.value(), "bias") : empty_t4, - softmax_scale, - mk_aotensor(out_t, "out"), - mk_aotensor(dout_t, "dout"), - mk_aotensor(dq_t, "dq"), - mk_aotensor(dk_t, "dk"), - mk_aotensor(dv_t, "dv"), - bias_requires_grad ? mk_aotensor(grad_bias, "db") : empty_t4, - mk_aotensor<2>(softmax_lse, "L"), - mk_aotensor<2>(delta, "delta"), - float(dropout_p), - mk_aoscalartensor(philox_seed), - mk_aoscalartensor(philox_offset), - 0, - is_causal, - stream); - } else { // cu_seqlens.has_value - auto d_head = Kv; - bool use_fused_bwd = d_head <= 192 && d_head * max_seqlen_q < 64 * 512; - if (use_fused_bwd) { - err = attn_bwd_fused(mk_aotensor(q_t, "q"), - mk_aotensor(k_t, "k"), - mk_aotensor(v_t, "v"), - bias.has_value() ? mk_aotensor(bias.value(), "bias") : empty_t4, - softmax_scale, - mk_aotensor(out_t, "out"), - mk_aotensor(dout_t, "dout"), - mk_aotensor(dq_t, "dq"), - mk_aotensor(dk_t, "dk"), - mk_aotensor(dv_t, "dv"), - bias_requires_grad ? mk_aotensor(grad_bias, "db") : empty_t4, - mk_aotensor<2>(softmax_lse, "L"), - float(dropout_p), - mk_aoscalartensor(philox_seed), - mk_aoscalartensor(philox_offset), - 0, - is_causal, - stream); - } else { - at::Tensor delta = at::empty_like(softmax_lse).contiguous(); - err = attn_bwd(mk_aotensor(q_t, "q"), - mk_aotensor(k_t, "k"), - mk_aotensor(v_t, "v"), - bias.has_value() ? mk_aotensor(bias.value(), "bias") : empty_t4, - softmax_scale, - mk_aotensor(out_t, "out"), - mk_aotensor(dout_t, "dout"), - mk_aotensor(dq_t, "dq"), - mk_aotensor(dk_t, "dk"), - mk_aotensor(dv_t, "dv"), - bias_requires_grad ? mk_aotensor(grad_bias, "db") : empty_t4, - mk_aotensor<2>(softmax_lse, "L"), - mk_aotensor<2>(delta, "delta"), - float(dropout_p), - mk_aoscalartensor(philox_seed), - mk_aoscalartensor(philox_offset), - 0, - is_causal, - stream); - } //used_fused_bwd - } // cuseqlen.has_value + using aotriton::v3::flash::CausalType; + using aotriton::v3::flash::VarlenType; + using aotriton::v3::flash::WindowValue; + aotriton::v3::flash::attn_bwd_params params; + params.Q = mk_aotensor(q_t, "q"); + params.K = mk_aotensor(k_t, "k"); + params.V = mk_aotensor(v_t, "v"); + params.B = bias.has_value() ? mk_aotensor(bias.value(), "bias") : empty_t4; + params.Sm_scale = softmax_scale; + params.Out = mk_aotensor(out_t, "out"); + params.DO = mk_aotensor(dout_t, "dout"); + params.DK = mk_aotensor(dk_t, "dk"); + params.DV = mk_aotensor(dv_t, "dv"); + params.DQ = mk_aotensor(dq_t, "dq"); + params.DB = bias_requires_grad ? mk_aotensor(grad_bias, "db") : empty_t4; + params.L = mk_aotensor<2>(softmax_lse, "L"); + params.Max_seqlen_q = max_seqlen_q; // Unused if cu_seqlens_q is empty + params.Max_seqlen_k = max_seqlen_k; // Unused if cu_seqlens_k is empty + params.dropout_p = float(dropout_p); + params.philox_seed_ptr = mk_aoscalartensor(philox_seed); + params.philox_offset1 = mk_aoscalartensor(philox_offset); + params.philox_offset2 = 0; + params.causal_type = is_causal ? CausalType::WindowedAttention : CausalType::None; + if (static_cast(sdp::CustomMaskType::CausalFromTopLeft) == custom_mask_type) { + params.window_left = WindowValue::TopLeftAligned; + params.window_right = WindowValue::TopLeftAligned; + } else if (static_cast(sdp::CustomMaskType::CausalFromBottomRight) == custom_mask_type) { + params.window_left = WindowValue::BottomRightAligned; + params.window_right = WindowValue::BottomRightAligned; + } + using sdp::aotriton_adapter::mklazy_empty_like; + using sdp::aotriton_adapter::mklazy_fp32zeros; + using sdp::aotriton_adapter::LazyTensorContext; + LazyTensorContext lazy_delta { .like_tensor = softmax_lse, .tensor_name = "delta" }; + LazyTensorContext lazy_dq_acc { .like_tensor = dq_t, .tensor_name = "dq_acc" }; + params.D = mklazy_empty_like<2>(&lazy_delta); + params.DQ_ACC = mklazy_fp32zeros<4>(&lazy_dq_acc); + if (cu_seqlens_q.has_value()) { + params.varlen_type = VarlenType::CompactVarlen; + params.cu_seqlens_q = mk_aotensor<1>(cu_seqlens_q.value(), "cu_seqlens_q"); + params.cu_seqlens_k = mk_aotensor<1>(cu_seqlens_k.value(), "cu_seqlens_k"); + } else { + params.varlen_type = VarlenType::None; + } + aotriton::v3::flash::attn_options opts; + opts.deterministic = deterministic; + err = aotriton::v3::flash::attn_bwd(params, + aotriton::v3::flash::attn_bwd_params::kVersion, + stream, + &opts); #else // DISABLE_AOTRITON TORCH_CHECK(false, "Attempting to use aotriton mem_eff_backward backend in a build that has not built AOTriton"); #endif diff --git a/aten/src/ATen/native/transformers/cuda/sdp_utils.cpp b/aten/src/ATen/native/transformers/cuda/sdp_utils.cpp index ff3eee28bc314..c7f91c134b657 100644 --- a/aten/src/ATen/native/transformers/cuda/sdp_utils.cpp +++ b/aten/src/ATen/native/transformers/cuda/sdp_utils.cpp @@ -140,19 +140,11 @@ int64_t minimum_gemm_alignment(sdp_params const& params) { return matmul_alignment_mn; } -// On ROCM, ME and FA share the backend, and hence they share the checking -// function for fundamental limitations by the GPU kernel -// caller_is_meff is added to make the TORCH_WARN message showing the correct result -template -bool check_head_dim_size_flash(sdp_params const& params, bool debug) { #if USE_ROCM_ATTENTION - if (at::cuda::device_count() == 0) { - return false; - } - // AOTriton 0.9+ supports head_dim up to 512 - const static auto max_hdim = []() { +inline int aotriton_max_hdim() { + static const int max_hdim = []() { #if AOTRITON_VERSION_CURRENT == AOTRITON_VERSION_INT(0, 11) - // gfx11xx only support hdim <= 256 on AOTriton 0.11 + // gfx11xx only support hdim <= 256 on AOTriton 0.11/0.12 auto dprops = at::cuda::getCurrentDeviceProperties(); const c10::basic_string_view arch(dprops->gcnArchName); if (arch.starts_with("gfx11")) { @@ -165,7 +157,28 @@ bool check_head_dim_size_flash(sdp_params const& params, bool debug) { return 256; #endif }(); - const auto max_size = c10::SymInt(max_hdim); + return max_hdim; +} +#endif // USE_ROCM_ATTENTION + +// For AOTriton <= 0.11: +// On ROCM, ME and FA share the backend, and hence they share the checking +// function for fundamental limitations by the GPU kernel +// caller_is_meff is added to make the TORCH_WARN message showing the correct result +// +// FIXME: revert this reuse when removing AOTriton <= 0.11 support +// +// AOTriton 0.12 supports hdim_qk != hdim_vo, but we cannot enable this in +// check_head_dim_size_flash because it changes the backend selection logic for +// FA, which can break certain workloads that rely on the behavior of rejecting +// FA for hdim_qk != hdim_vo +template +bool check_head_dim_size_flash(sdp_params const& params, bool debug) { +#if USE_ROCM_ATTENTION + if (at::cuda::device_count() == 0) { + return false; + } + const auto max_size = c10::SymInt(aotriton_max_hdim()); #else // All head_dim sizes must be equal and less than 256 const auto max_size = c10::SymInt(256); @@ -245,9 +258,20 @@ bool check_head_dim_size_flash_nested(sdp_params const& params, bool debug) { } bool check_head_dim_size_mem_efficient(sdp_params const& params, bool debug) { +#if USE_ROCM_ATTENTION +#if AOTRITON_VERSION_CURRENT < AOTRITON_VERSION_INT(0, 12) + return check_head_dim_size_flash_nested(params, debug); +#endif +#endif const auto query_size_last = params.query.sym_size(-1); const auto value_size_last = params.value.sym_size(-1); +#ifdef USE_ROCM + bool is_half = (params.query.dtype() == at::kHalf) || + (params.query.dtype() == at::kBFloat16); + const int64_t alignment = is_half ? 8 : 4; +#else const int64_t alignment = minimum_gemm_alignment(params); +#endif if (!(query_size_last == params.key.sym_size(-1) && query_size_last % alignment == 0 && query_size_last > 0 && value_size_last % alignment == 0 && value_size_last > 0)) { @@ -266,6 +290,27 @@ bool check_head_dim_size_mem_efficient(sdp_params const& params, bool debug) { } return false; } +#if USE_ROCM_ATTENTION +#if AOTRITON_VERSION_CURRENT >= AOTRITON_VERSION_INT(0, 12) + const auto max_size = c10::SymInt(aotriton_max_hdim()); + if (!(query_size_last <= max_size && value_size_last <= max_size)) { + if (debug) { + TORCH_WARN( + "Mem efficient attention on ROCM requires last dimension of inputs to less or equal than ", + max_size, + ". ", + "Got Query.size(-1): ", + query_size_last, + ", Key.size(-1): ", + params.key.sym_size(-1), + ", Value.size(-1): ", + params.value.sym_size(-1), + " instead. (Note this limit differs among architectures)"); + } + return false; + } +#endif // AOTRITON_VERSION_CURRENT >= AOTRITON_VERSION_INT(0, 12) +#endif // USE_ROCM_ATTENTION return true; } @@ -881,11 +926,7 @@ bool can_use_mem_efficient_attention(sdp_params const& params, bool debug) { check_all_tensors_on_device, check_mem_efficient_hardware_support, check_tensor_shapes, -#ifdef USE_ROCM - check_head_dim_size_flash -#else check_head_dim_size_mem_efficient -#endif ); for (auto& constraint : general_constraints) { if (!constraint(params, debug)) { @@ -895,11 +936,6 @@ bool can_use_mem_efficient_attention(sdp_params const& params, bool debug) { if (has_for_nested_inputs(params)) { constexpr auto nested_constraints = c10::array_of( -#ifndef USE_ROCM // ME and FA shares backend on ROCM and thus supports training - check_requires_grad_and_nested, -#else // Meanwhile ME on ROCM share the limits of FA about head dimensions - check_head_dim_size_flash_nested, -#endif check_batch_size_nested, check_for_seq_len_0_nested_tensor); for (auto& constraint : nested_constraints) { diff --git a/aten/src/ATen/native/transformers/hip/aotriton_adapter.h b/aten/src/ATen/native/transformers/hip/aotriton_adapter.h index e40376ae0c3a7..c0cac2a76a2a5 100644 --- a/aten/src/ATen/native/transformers/hip/aotriton_adapter.h +++ b/aten/src/ATen/native/transformers/hip/aotriton_adapter.h @@ -11,6 +11,10 @@ #include #include +#if AOTRITON_VERSION_CURRENT >= AOTRITON_VERSION_INT(0, 12) +#define AOTRITON_V2_API_FLASH_ATTN_H // Suppress the include of deprecated flash/v2.h +#endif + //////////////////////////////////////////////////////////////////////////////// // Common macros copied from cuda/mem_eff_attention/gemm_kernel_utils.h //////////////////////////////////////////////////////////////////////////////// @@ -127,8 +131,17 @@ struct LazyTensorContext { template struct LazyTensorFunctions : public LazyTensorContext { - static aotriton::TensorView acquire(void* cookie) { - auto ctx = (LazyTensorContext*)cookie; +#if AOTRITON_VERSION_CURRENT >= AOTRITON_VERSION_INT(0, 12) + using HolderType = aotriton::LazyTensor; +#else + using HolderType = void; +#endif + static aotriton::TensorView acquire(HolderType* self) { +#if AOTRITON_VERSION_CURRENT >= AOTRITON_VERSION_INT(0, 12) + auto ctx = (LazyTensorContext*)self->cookie; +#else + auto ctx = (LazyTensorContext*)self; +#endif if (!ctx->tensor.defined()) { auto q = ctx->like_tensor; if constexpr (kRequireZeros) { @@ -141,7 +154,7 @@ struct LazyTensorFunctions : public LazyTensorContext { return mk_aotensor(ctx->tensor, ctx->tensor_name); } - static void dispose(void* cookie) { + static void dispose(HolderType* cookie) { } }; diff --git a/aten/src/ATen/native/transformers/hip/aotriton_versions.h b/aten/src/ATen/native/transformers/hip/aotriton_versions.h index 2f5d3f0e12228..8b06cc1acc789 100644 --- a/aten/src/ATen/native/transformers/hip/aotriton_versions.h +++ b/aten/src/ATen/native/transformers/hip/aotriton_versions.h @@ -17,4 +17,10 @@ #define AOTRITON_V3_API 0 #endif +#if AOTRITON_VERSION_CURRENT >= AOTRITON_VERSION_INT(0, 12) +#define AOTRITON_COMPACT_VARLEN_LSE 1 +#else +#define AOTRITON_COMPACT_VARLEN_LSE 0 +#endif + #endif diff --git a/aten/src/ATen/native/transformers/hip/flash_attn/aot/mha_all_aot.hip b/aten/src/ATen/native/transformers/hip/flash_attn/aot/mha_all_aot.hip index f3816296a1d4f..c16f7d1aad233 100644 --- a/aten/src/ATen/native/transformers/hip/flash_attn/aot/mha_all_aot.hip +++ b/aten/src/ATen/native/transformers/hip/flash_attn/aot/mha_all_aot.hip @@ -257,7 +257,6 @@ mha_fwd_aot(const at::Tensor &q, // batch_size x seqlen_q x num_heads x } hipError_t err; // TODO: Error handling - using aotriton::v2::flash::attn_fwd; using sdp::aotriton_adapter::mk_aotensor; using sdp::aotriton_adapter::mk_aoscalartensor; using sdp::aotriton_adapter::mk_philoxtensor; @@ -270,54 +269,32 @@ mha_fwd_aot(const at::Tensor &q, // batch_size x seqlen_q x num_heads x auto seed_output = mk_philoxtensor(use_philox_state ? seed_t.data_ptr() : nullptr); auto offset_output = mk_philoxtensor(use_philox_state ? offset_t.data_ptr() : nullptr); auto persistent_counter = mk_atomictensor(is_causal ? atomic_counter.data_ptr() : nullptr); - if (uses_swa || AOTRITON_ALWAYS_V3_API) { -#if AOTRITON_V3_API - using aotriton::v3::flash::CausalType; - using aotriton::v3::flash::VarlenType; - aotriton::v3::flash::attn_fwd_params params; - params.Q = mk_aotensor(q_t, "q"); - params.K = mk_aotensor(k_t, "k"); - params.V = mk_aotensor(v_t, "v"); - params.Sm_scale = softmax_scale; - params.L = mk_aotensor<2>(M, "M"); - params.Out = mk_aotensor(output_t, "Out"); - params.Max_seqlen_q = seqlen_q; // Unused if cu_seqlens_q is empty - params.Max_seqlen_k = seqlen_k; // Unused if cu_seqlens_k is empty - params.dropout_p = p_dropout; - params.philox_seed_ptr = seed; - params.philox_offset1 = offset1; - params.philox_offset2 = offset2; - params.philox_seed_output = seed_output; - params.philox_offset_output = offset_output; - params.encoded_softmax = mk_aotensor(softmax_fa_t, "encoded_softmax"); - params.persistent_atomic_counter = persistent_counter; - params.causal_type = is_causal ? CausalType::WindowedAttention : CausalType::None; - params.varlen_type = VarlenType::None; - params.window_left = window_left; - params.window_right = window_right; - err = aotriton::v3::flash::attn_fwd(params, - aotriton::v3::flash::attn_fwd_params::kVersion, - stream); -#endif - } else { - err = attn_fwd(mk_aotensor(q_t, "q"), - mk_aotensor(k_t, "k"), - mk_aotensor(v_t, "v"), - empty_bias, - softmax_scale, - mk_aotensor<2>(M, "M"), - mk_aotensor(output_t, "Out"), - p_dropout, - seed, - offset1, - offset2, - seed_output, - offset_output, - mk_aotensor(softmax_fa_t, "encoded_softmax"), - is_causal, - persistent_counter, - stream); - } + using aotriton::v3::flash::CausalType; + using aotriton::v3::flash::VarlenType; + aotriton::v3::flash::attn_fwd_params params; + params.Q = mk_aotensor(q_t, "q"); + params.K = mk_aotensor(k_t, "k"); + params.V = mk_aotensor(v_t, "v"); + params.Sm_scale = softmax_scale; + params.L = mk_aotensor<2>(M, "M"); + params.Out = mk_aotensor(output_t, "Out"); + params.Max_seqlen_q = seqlen_q; // Unused if cu_seqlens_q is empty + params.Max_seqlen_k = seqlen_k; // Unused if cu_seqlens_k is empty + params.dropout_p = p_dropout; + params.philox_seed_ptr = seed; + params.philox_offset1 = offset1; + params.philox_offset2 = offset2; + params.philox_seed_output = seed_output; + params.philox_offset_output = offset_output; + params.encoded_softmax = mk_aotensor(softmax_fa_t, "encoded_softmax"); + params.persistent_atomic_counter = persistent_counter; + params.causal_type = is_causal ? CausalType::WindowedAttention : CausalType::None; + params.varlen_type = VarlenType::None; + params.window_left = window_left; + params.window_right = window_right; + err = aotriton::v3::flash::attn_fwd(params, + aotriton::v3::flash::attn_fwd_params::kVersion, + stream); // Note: These are propagated up to the return of mha_fwd(). comments // represent the assignments at that level return {out, // output @@ -350,7 +327,7 @@ mha_varlen_fwd_aot(const at::Tensor &q, // total_q x num_heads x head_size, tot std::optional window_size_right, const bool return_softmax, const std::optional& gen_) { - TORCH_CHECK(!seqused_k.has_value(), "[ROCm] mha_varlen_fwd: seqused_k must be nullopt"); + bool strided_varlen = seqused_k.has_value(); const bool paged_KV = block_table_.has_value(); TORCH_CHECK(!paged_KV, "[ROCm] mha_varlen_fwd: block_table_ must be nullopt"); TORCH_CHECK(!alibi_slopes_.has_value(), "[ROCm] mha_varlen_fwd: alibi_slopes_ must be nullopt"); @@ -425,8 +402,13 @@ mha_varlen_fwd_aot(const at::Tensor &q, // total_q x num_heads x head_size, tot auto opts = q.options(); +#if AOTRITON_COMPACT_VARLEN_LSE + auto softmax_lse = at::empty({num_heads, total_q}, opts.dtype(at::kFloat)); + at::Tensor M = softmax_lse; +#else auto softmax_lse = at::empty({batch_size, num_heads, max_seqlen_q}, opts.dtype(at::kFloat)); at::Tensor M = softmax_lse.view({batch_size * num_heads, max_seqlen_q}); +#endif at::Tensor softmax_fa_t; // Only return softmax if there's dropout to reduce compilation time if (return_softmax) { @@ -465,7 +447,6 @@ mha_varlen_fwd_aot(const at::Tensor &q, // total_q x num_heads x head_size, tot if (max_seqlen_k > 0) { hipError_t err; // TODO: Error handling - using aotriton::v2::flash::attn_fwd_compact_varlen; using sdp::aotriton_adapter::mk_aotensor; using sdp::aotriton_adapter::mk_aoscalartensor; using sdp::aotriton_adapter::mk_philoxtensor; @@ -483,60 +464,48 @@ mha_varlen_fwd_aot(const at::Tensor &q, // total_q x num_heads x head_size, tot auto seed_output = use_philox_state ? mk_philoxtensor(seed_t.data_ptr()) : nullscalar; auto offset_output = use_philox_state ? mk_philoxtensor(offset_t.data_ptr()) : nullscalar; auto persistent_counter = mk_atomictensor(is_causal ? atomic_counter.data_ptr() : nullptr); - if (uses_swa || AOTRITON_ALWAYS_V3_API) { -#if AOTRITON_V3_API - using aotriton::v3::flash::CausalType; - using aotriton::v3::flash::VarlenType; - aotriton::v3::flash::attn_fwd_params params; - params.Q = mk_aotensor(q_padded, "q"); - params.K = mk_aotensor(k_padded, "k"); - params.V = mk_aotensor(v_padded, "v"); - params.Sm_scale = softmax_scale; - params.L = mk_aotensor<2>(M, "M"); - params.Out = mk_aotensor(out_padded, "Out"); + using aotriton::v3::flash::CausalType; + using aotriton::v3::flash::VarlenType; + aotriton::v3::flash::attn_fwd_params params; + params.Q = mk_aotensor(q_padded, "q"); + params.K = mk_aotensor(k_padded, "k"); + params.V = mk_aotensor(v_padded, "v"); + params.Sm_scale = softmax_scale; + params.L = mk_aotensor<2>(M, "logsumexp"); + params.Out = mk_aotensor(out_padded, "Out"); + params.Max_seqlen_q = max_seqlen_q; // Unused if cu_seqlens_q is empty + params.Max_seqlen_k = max_seqlen_k; // Unused if cu_seqlens_k is empty + params.dropout_p = p_dropout; + params.philox_seed_ptr = seed; + params.philox_offset1 = offset1; + params.philox_offset2 = offset2; + params.philox_seed_output = seed_output; + params.philox_offset_output = offset_output; + params.encoded_softmax = mk_aotensor(softmax_fa_t, "encoded_softmax"); + params.persistent_atomic_counter = persistent_counter; + params.causal_type = is_causal ? CausalType::WindowedAttention : CausalType::None; + params.varlen_type = strided_varlen ? VarlenType::StridedVarlen : VarlenType::CompactVarlen; + if (strided_varlen) { + // seqused_k holds per-batch actual kv lengths; the kernel expects cumulative + // offsets for cu_seqlens_k so it can compute seqlen_k via differencing. + // seq_strides_k carries the real memory offsets into the KV cache. + const auto& seqused = seqused_k.value(); + const int num_seqs = seqused.size(0); + at::Tensor cu_seqlens_k_from_seqused = at::zeros({num_seqs + 1}, seqused.options()); + cu_seqlens_k_from_seqused.slice(0, 1).copy_(seqused.cumsum(0)); params.cu_seqlens_q = mk_aotensor<1>(cu_seqlens_q, "cu_seqlens_q"); - params.cu_seqlens_k = mk_aotensor<1>(cu_seqlens_k, "cu_seqlens_k"); - params.Max_seqlen_q = max_seqlen_q; // Unused if cu_seqlens_q is empty - params.Max_seqlen_k = max_seqlen_k; // Unused if cu_seqlens_k is empty - params.dropout_p = p_dropout; - params.philox_seed_ptr = seed; - params.philox_offset1 = offset1; - params.philox_offset2 = offset2; - params.philox_seed_output = seed_output; - params.philox_offset_output = offset_output; - params.encoded_softmax = mk_aotensor(softmax_fa_t, "encoded_softmax"); - params.persistent_atomic_counter = persistent_counter; - params.causal_type = is_causal ? CausalType::WindowedAttention : CausalType::None; - params.varlen_type = VarlenType::CompactVarlen; - params.window_left = window_left; - params.window_right = window_right; - err = aotriton::v3::flash::attn_fwd(params, - aotriton::v3::flash::attn_fwd_params::kVersion, - stream); -#endif + params.cu_seqlens_k = mk_aotensor<1>(cu_seqlens_k_from_seqused, "cu_seqlens_k"); + params.seq_strides_q = mk_aotensor<1>(cu_seqlens_q, "seq_strides_q"); + params.seq_strides_k = mk_aotensor<1>(cu_seqlens_k, "seq_strides_k"); } else { - err = attn_fwd_compact_varlen(mk_aotensor(q_padded, "q"), - mk_aotensor(k_padded, "k"), - mk_aotensor(v_padded, "v"), - empty_bias, - mk_aotensor<1>(cu_seqlens_q, "cu_seqlens_q"), - mk_aotensor<1>(cu_seqlens_k, "cu_seqlens_k"), - max_seqlen_q, - max_seqlen_k, - softmax_scale, - mk_aotensor<2>(M, "M"), - mk_aotensor(out_padded, "Out"), - p_dropout, - seed, - offset1, - offset2, - seed_output, - offset_output, - mk_aotensor(softmax_fa_t, "encoded_softmax"), - is_causal, - persistent_counter, - stream); + params.cu_seqlens_q = mk_aotensor<1>(cu_seqlens_q, "cu_seqlens_q"); + params.cu_seqlens_k = mk_aotensor<1>(cu_seqlens_k, "cu_seqlens_k"); } + params.window_left = window_left; + params.window_right = window_right; + err = aotriton::v3::flash::attn_fwd(params, + aotriton::v3::flash::attn_fwd_params::kVersion, + stream); } else { // If seqlen_k == 0, then we have an empty tensor. We need to set the output to 0. out.zero_(); @@ -678,96 +647,44 @@ mha_bwd_aot(const at::Tensor &dout, // batch_size x seqlen_q x num_heads, x hea hipError_t err; // TODO: Error handling using sdp::aotriton_adapter::mk_aotensor; using sdp::aotriton_adapter::mk_aoscalartensor; - if (uses_swa || AOTRITON_ALWAYS_V3_API) { -#if AOTRITON_V3_API - // Fused BWD does not support SWA - using aotriton::v3::flash::CausalType; - using aotriton::v3::flash::VarlenType; - aotriton::v3::flash::attn_bwd_params params; - params.Q = mk_aotensor(q_t, "q"); - params.K = mk_aotensor(k_t, "k"); - params.V = mk_aotensor(v_t, "v"); - params.Sm_scale = softmax_scale; - params.Out = mk_aotensor(out_t, "out"); - params.DO = mk_aotensor(dout_t, "dout"); - params.DQ = mk_aotensor(dq_t, "dq"); - params.DK = mk_aotensor(dk_t, "dk"); - params.DV = mk_aotensor(dv_t, "dv"); - params.L = mk_aotensor<2>(softmax_lse_cont, "L"); - params.Max_seqlen_q = seqlen_q; // Unused if cu_seqlens_q is empty - params.Max_seqlen_k = seqlen_k; // Unused if cu_seqlens_k is empty - params.dropout_p = p_dropout; - params.philox_seed_ptr = mk_aoscalartensor(philox_seed); - params.philox_offset1 = mk_aoscalartensor(philox_offset); - params.philox_offset2 = 0; - // SWA in AOTriton Kernels is treated as "Generalized Causal masks" - params.causal_type = is_causal || uses_swa ? CausalType::WindowedAttention : CausalType::None; - params.window_left = window_left; - params.window_right = window_right; - params.varlen_type = VarlenType::None; -#if AOTRITON_ALWAYS_V3_API - using sdp::aotriton_adapter::mklazy_empty_like; - using sdp::aotriton_adapter::mklazy_fp32zeros; - using sdp::aotriton_adapter::LazyTensorContext; - LazyTensorContext lazy_delta { .like_tensor = softmax_lse_cont, .tensor_name = "delta" }; - LazyTensorContext lazy_dq_acc { .like_tensor = dq_t, .tensor_name = "dq_acc" }; - params.D = mklazy_empty_like<2>(&lazy_delta); - params.DQ_ACC = mklazy_fp32zeros<4>(&lazy_dq_acc); -#else - at::Tensor delta = at::empty_like(softmax_lse_cont).contiguous(); - params.D = mk_aotensor<2>(delta, "delta"); -#endif - err = aotriton::v3::flash::attn_bwd(params, - aotriton::v3::flash::attn_bwd_params::kVersion, - stream); -#endif - } else if (use_fused_bwd) { - using aotriton::v2::flash::attn_bwd_fused; - using sdp::aotriton_adapter::cast_dtype; - aotriton::TensorView<4> empty_bias(0, {0,0,0,0}, {0,0,0,0}, cast_dtype(q.dtype())); - err = attn_bwd_fused(mk_aotensor(q_t, "q"), - mk_aotensor(k_t, "k"), - mk_aotensor(v_t, "v"), - empty_bias, - softmax_scale, - mk_aotensor(out_t, "out"), - mk_aotensor(dout_t, "dout"), - mk_aotensor(dq_t, "dq"), - mk_aotensor(dk_t, "dk"), - mk_aotensor(dv_t, "dv"), - empty_bias, // dbb - mk_aotensor<2>(softmax_lse_cont, "L"), - p_dropout, - mk_aoscalartensor(philox_seed), - mk_aoscalartensor(philox_offset), - 0, - is_causal, - stream); - } else { - at::Tensor delta = at::empty_like(softmax_lse_cont).contiguous(); - using aotriton::v2::flash::attn_bwd; - using sdp::aotriton_adapter::cast_dtype; - aotriton::TensorView<4> empty_bias(0, {0,0,0,0}, {0,0,0,0}, cast_dtype(q.dtype())); - err = attn_bwd(mk_aotensor(q_t, "q"), - mk_aotensor(k_t, "k"), - mk_aotensor(v_t, "v"), - empty_bias, - softmax_scale, - mk_aotensor(out_t, "out"), - mk_aotensor(dout_t, "dout"), - mk_aotensor(dq_t, "dq"), - mk_aotensor(dk_t, "dk"), - mk_aotensor(dv_t, "dv"), - empty_bias, // db - mk_aotensor<2>(softmax_lse_cont, "L"), - mk_aotensor<2>(delta, "delta"), - p_dropout, - mk_aoscalartensor(philox_seed), - mk_aoscalartensor(philox_offset), - 0, - is_causal, - stream); - } + // Fused BWD does not support SWA + using aotriton::v3::flash::CausalType; + using aotriton::v3::flash::VarlenType; + aotriton::v3::flash::attn_bwd_params params; + params.Q = mk_aotensor(q_t, "q"); + params.K = mk_aotensor(k_t, "k"); + params.V = mk_aotensor(v_t, "v"); + params.Sm_scale = softmax_scale; + params.Out = mk_aotensor(out_t, "out"); + params.DO = mk_aotensor(dout_t, "dout"); + params.DQ = mk_aotensor(dq_t, "dq"); + params.DK = mk_aotensor(dk_t, "dk"); + params.DV = mk_aotensor(dv_t, "dv"); + params.L = mk_aotensor<2>(softmax_lse_cont, "L"); + params.Max_seqlen_q = seqlen_q; // Unused if cu_seqlens_q is empty + params.Max_seqlen_k = seqlen_k; // Unused if cu_seqlens_k is empty + params.dropout_p = p_dropout; + params.philox_seed_ptr = mk_aoscalartensor(philox_seed); + params.philox_offset1 = mk_aoscalartensor(philox_offset); + params.philox_offset2 = 0; + // SWA in AOTriton Kernels is treated as "Generalized Causal masks" + params.causal_type = is_causal || uses_swa ? CausalType::WindowedAttention : CausalType::None; + params.window_left = window_left; + params.window_right = window_right; + params.varlen_type = VarlenType::None; + using sdp::aotriton_adapter::mklazy_empty_like; + using sdp::aotriton_adapter::mklazy_fp32zeros; + using sdp::aotriton_adapter::LazyTensorContext; + LazyTensorContext lazy_delta { .like_tensor = softmax_lse_cont, .tensor_name = "delta" }; + LazyTensorContext lazy_dq_acc { .like_tensor = dq_t, .tensor_name = "dq_acc" }; + params.D = mklazy_empty_like<2>(&lazy_delta); + params.DQ_ACC = mklazy_fp32zeros<4>(&lazy_dq_acc); + aotriton::v3::flash::attn_options options; + options.deterministic = deterministic; + err = aotriton::v3::flash::attn_bwd(params, + aotriton::v3::flash::attn_bwd_params::kVersion, + stream, + &options); return { dq, dk, dv, softmax_d }; } @@ -850,7 +767,11 @@ mha_varlen_bwd_aot(const at::Tensor &dout, // total_q x num_heads, x head_size CHECK_SHAPE(cu_seqlens_q, batch_size + 1); CHECK_SHAPE(cu_seqlens_k, batch_size + 1); +#if AOTRITON_COMPACT_VARLEN_LSE + at::Tensor softmax_lse_cont = softmax_lse.view({num_heads, total_q}).contiguous(); +#else at::Tensor softmax_lse_cont = softmax_lse.view({batch_size * num_heads, max_seqlen_q}).contiguous(); +#endif at::Tensor q_padded, k_padded, v_padded; q_padded = q.unsqueeze(0).transpose(1, 2); @@ -931,79 +852,45 @@ mha_varlen_bwd_aot(const at::Tensor &dout, // total_q x num_heads, x head_size hipError_t err; // TODO: Error handling using sdp::aotriton_adapter::mk_aotensor; using sdp::aotriton_adapter::mk_aoscalartensor; - if (uses_swa || AOTRITON_ALWAYS_V3_API) { -#if AOTRITON_V3_API - using aotriton::v3::flash::CausalType; - using aotriton::v3::flash::VarlenType; - aotriton::v3::flash::attn_bwd_params params; - params.Q = mk_aotensor(q_padded, "q"); - params.K = mk_aotensor(k_padded, "k"); - params.V = mk_aotensor(v_padded, "v"); - params.Sm_scale = softmax_scale; - params.Out = mk_aotensor(out_t, "out"); - params.DO = mk_aotensor(dout_t, "dout"); - params.DK = mk_aotensor(dk_padded, "dk"); - params.DV = mk_aotensor(dv_padded, "dv"); - params.DQ = mk_aotensor(dq_padded, "dq"); - params.L = mk_aotensor<2>(softmax_lse_cont, "L"); - params.cu_seqlens_q = mk_aotensor<1>(cu_seqlens_q, "cu_seqlens_q"); - params.cu_seqlens_k = mk_aotensor<1>(cu_seqlens_k, "cu_seqlens_k"); - params.Max_seqlen_q = max_seqlen_q; // Unused if cu_seqlens_q is empty - params.Max_seqlen_k = max_seqlen_k; // Unused if cu_seqlens_k is empty - params.dropout_p = p_dropout; - params.philox_seed_ptr = mk_aoscalartensor(philox_seed); - params.philox_offset1 = mk_aoscalartensor(philox_offset); - params.philox_offset2 = 0; - // SWA in AOTriton Kernels is treated as "Generalized Causal masks" - params.causal_type = is_causal || uses_swa ? CausalType::WindowedAttention : CausalType::None; - params.varlen_type = VarlenType::CompactVarlen; - params.window_left = window_left; - params.window_right = window_right; -#if AOTRITON_ALWAYS_V3_API - using sdp::aotriton_adapter::mklazy_empty_like; - using sdp::aotriton_adapter::mklazy_fp32zeros; - using sdp::aotriton_adapter::LazyTensorContext; - LazyTensorContext lazy_delta { .like_tensor = softmax_lse_cont, .tensor_name = "delta" }; - LazyTensorContext lazy_dq_acc { .like_tensor = dq_padded, .tensor_name = "dq_acc" }; - params.D = mklazy_empty_like<2>(&lazy_delta); - params.DQ_ACC = mklazy_fp32zeros<4>(&lazy_dq_acc); -#else - at::Tensor delta = at::empty_like(softmax_lse_cont).contiguous(); - params.D = mk_aotensor<2>(delta, "delta"); -#endif - err = aotriton::v3::flash::attn_bwd(params, - aotriton::v3::flash::attn_bwd_params::kVersion, - stream); -#endif // AOTRITON_ALWAYS_V3_API - } else { - using aotriton::v2::flash::attn_bwd_compact_varlen; - using sdp::aotriton_adapter::cast_dtype; - at::Tensor delta = at::empty_like(softmax_lse_cont).contiguous(); - aotriton::TensorView<4> empty_bias(0, {0,0,0,0}, {0,0,0,0}, cast_dtype(q.dtype())); - err = attn_bwd_compact_varlen(mk_aotensor(q_padded, "q"), - mk_aotensor(k_padded, "k"), - mk_aotensor(v_padded, "v"), - mk_aotensor<1>(cu_seqlens_q, "cu_seqlens_q"), - mk_aotensor<1>(cu_seqlens_k, "cu_seqlens_k"), - max_seqlen_q, - max_seqlen_k, - empty_bias, - softmax_scale, - mk_aotensor(out_t, "out"), - mk_aotensor(dout_t, "dout"), - mk_aotensor(dq_padded, "dq"), - mk_aotensor(dk_padded, "dk"), - mk_aotensor(dv_padded, "dv"), - empty_bias, - mk_aotensor<2>(softmax_lse_cont, "L"), - mk_aotensor<2>(delta, "delta"), - p_dropout, - mk_aoscalartensor(philox_seed), - mk_aoscalartensor(philox_offset), - 0, - is_causal, - stream); - } + using aotriton::v3::flash::CausalType; + using aotriton::v3::flash::VarlenType; + aotriton::v3::flash::attn_bwd_params params; + params.Q = mk_aotensor(q_padded, "q"); + params.K = mk_aotensor(k_padded, "k"); + params.V = mk_aotensor(v_padded, "v"); + params.Sm_scale = softmax_scale; + params.Out = mk_aotensor(out_t, "out"); + params.DO = mk_aotensor(dout_t, "dout"); + params.DK = mk_aotensor(dk_padded, "dk"); + params.DV = mk_aotensor(dv_padded, "dv"); + params.DQ = mk_aotensor(dq_padded, "dq"); + params.L = mk_aotensor<2>(softmax_lse_cont, "L"); + params.cu_seqlens_q = mk_aotensor<1>(cu_seqlens_q, "cu_seqlens_q"); + params.cu_seqlens_k = mk_aotensor<1>(cu_seqlens_k, "cu_seqlens_k"); + params.Max_seqlen_q = max_seqlen_q; // Unused if cu_seqlens_q is empty + params.Max_seqlen_k = max_seqlen_k; // Unused if cu_seqlens_k is empty + params.dropout_p = p_dropout; + params.philox_seed_ptr = mk_aoscalartensor(philox_seed); + params.philox_offset1 = mk_aoscalartensor(philox_offset); + params.philox_offset2 = 0; + // SWA in AOTriton Kernels is treated as "Generalized Causal masks" + params.causal_type = is_causal || uses_swa ? CausalType::WindowedAttention : CausalType::None; + params.varlen_type = VarlenType::CompactVarlen; + params.window_left = window_left; + params.window_right = window_right; + using sdp::aotriton_adapter::mklazy_empty_like; + using sdp::aotriton_adapter::mklazy_fp32zeros; + using sdp::aotriton_adapter::LazyTensorContext; + LazyTensorContext lazy_delta { .like_tensor = softmax_lse_cont, .tensor_name = "delta" }; + LazyTensorContext lazy_dq_acc { .like_tensor = dq_padded, .tensor_name = "dq_acc" }; + params.D = mklazy_empty_like<2>(&lazy_delta); + params.DQ_ACC = mklazy_fp32zeros<4>(&lazy_dq_acc); + aotriton::v3::flash::attn_options options; + options.deterministic = deterministic; + err = aotriton::v3::flash::attn_bwd(params, + aotriton::v3::flash::attn_bwd_params::kVersion, + stream, + &options); } else { // If seqlen_q == 0, then we have an empty tensor. We need to set the output to 0. dq.zero_(); diff --git a/cmake/External/aotriton.cmake b/cmake/External/aotriton.cmake index 414dced29e297..0feb798320b72 100644 --- a/cmake/External/aotriton.cmake +++ b/cmake/External/aotriton.cmake @@ -9,56 +9,57 @@ if(NOT __AOTRITON_INCLUDED) # Replaces .ci/docker/aotriton_version.txt # Note packages information may have versions skipped (due to no ABI breaks) # But they must be listed from lower version to higher version - set(__AOTRITON_VER "0.11.2b") + set(__AOTRITON_VER "0.12.50tp") set(__AOTRITON_MANYLINUX_LIST - "manylinux_2_28" # rocm6.2 - "manylinux_2_28" # rocm6.3 "manylinux_2_28" # rocm6.4 "manylinux_2_28" # rocm7.0 "manylinux_2_28" # rocm7.1 "manylinux_2_28" # rocm7.2 + "manylinux_2_28" # rocm7.14 ) set(__AOTRITON_ROCM_LIST - "rocm6.2" - "rocm6.3" "rocm6.4" "rocm7.0" "rocm7.1" "rocm7.2" + "rocm7.14" ) if(DEFINED ENV{PYTORCH_AOTRITON_COMMIT}) set(__AOTRITON_CI_COMMIT "$ENV{PYTORCH_AOTRITON_COMMIT}") else() - set(__AOTRITON_CI_COMMIT "dd1b68b604b5258ee7a9f7b66ad95e7a82c18065") + set(__AOTRITON_CI_COMMIT "098300d1b2509d0e3599eacae542510e3f00a3cb") endif() set(__AOTRITON_SHA256_LIST - "d784314849ba1911181dfc80cd845064ff6f0cdad10e2f4c53eb84a8b89245b9" # rocm6.2 - "f4b14dc111c334e967b28a1cf9ed4c63264c634dbdccbb5849aa9490022992f7" # rocm6.3 - "6b51d8479c85b902334e4f5518f404a8f5d563fd8d4732cb8b621ed4b45c2876" # rocm6.4 - "5501a0a3b300890001b6625f2a3539a7bad60f386f0a061ebe7d4ed5ca0fafb9" # rocm7.0 - "fee36beb3ea484ce18155bbafe026c577fd6705e4469e59405b260bd74b8cc10" # rocm7.1 - "cd8abf27bbb63cec45c94135e9b28745966074263a6b0555e5878ae1cb6a2349" # rocm7.2 + "b9bd9f9d4aec5e398edb26d1ff662ba2cc36b358c02eebacbcd7361a73ebb55e" # rocm6.4 + "a3f3284e864f95d4b3b3e20c78deb6ec4c0d4f176749ff166325bb3e74b2dbeb" # rocm7.0 + "ee4c11db1a07bfa7f6f946be1e2fb02d4709c675c478cf48bd80bec771b0962e" # rocm7.1 + "b00baf5271e44d5db4c9c58be930c604494432587c9d791b75423af80740aeeb" # rocm7.2 + "a1562984a053c33e34ca6615c84c34d3ef3092c84609bf41ce62d9868bea4c52" # rocm7.14 ) set(__AOTRITON_IMAGE_LIST "amd-gfx90a" "amd-gfx942" "amd-gfx950" - "amd-gfx11xx" + "amd-gfx110x" + "amd-gfx115x" "amd-gfx120x" # TODO: Update on AOTriton integration #"amd-gfx1250" ) set(__AOTRITON_IMAGE_SHA256_LIST - "fe9f04b66bf52ac27cd025e1d89cfd04974dd3fb3ae076192f783641a4d80fdf" # amd-gfx90a - "0a7bcee19d3bb6d548732248c3234f7b92736c2ab7a7aae65294b87a7fd64c06" # amd-gfx942 - "c1ba3bfe84217fd67df3dd1f8b67c80a7f7b33d0ad4d74b41d6567036e032ace" # amd-gfx950 - "839299637fccb13fbe3e7823d57d1b2dcd0e0bed78abbcb7005ea5f4fd82b928" # amd-gfx11xx - "0a4ff324bffdac0c2fde87a8a7f70563d3c84a80ad4e8f31345f2b40a1384e95" # amd-gfx120x - # TODO: Update when AOTriton publishes gfx1250 images. - # Until then, may need to set AOTRITON_INSTALL_FROM_SOURCE=1 to build from source. - #"0000000000000000000000000000000000000000000000000000000000000000" # amd-gfx1250 + "bb8bf2237b77fc503bc2967ea0d99d6ca419126c479e951ea42b712737128086" # amd-gfx90a + "f08edacf83c9ccf1c4bdcb51f1cab052d1680abea31c9e035f3f9fadb2f13ba4" # amd-gfx942 + "307a37d729cda3a2120449909e5192cd71c2badccbd37f0222786098e69c7a91" # amd-gfx950 + "c9cac7cf6f277168e1659ac2f04706f8823580b7c7e3e895f5a5503ed6bdd55f" # amd-gfx110x + "3177387a15c678b30057f4584d1fc1b8f8db56163890cb5c98f27450209f5a7b" # amd-gfx115x + "68572511ce6487a83f9014bd255bd69c8943f87d0c93bd57b2daac5fbc6c79c1" # amd-gfx120x + "f2415cc328430a84287cb3f545d6965811b2fa54fe9986d91357b3d455f50b12" # amd-gfx1250 + ) - set(__AOTRITON_BASE_URL "https://github.com/ROCm/aotriton/releases/download/") # @lint-ignore + set(__AOTRITON_BASE_URL "$ENV{PYTORCH_AOTRITON_BASE_URL}") + if(NOT __AOTRITON_BASE_URL) + set(__AOTRITON_BASE_URL "https://github.com/ROCm/aotriton/releases/download/") # @lint-ignore + endif() set(__AOTRITON_Z "gz") # Set the default __AOTRITON_LIB path if(NOT WIN32) diff --git a/test/test_transformers.py b/test/test_transformers.py index ced9b0133e11d..ea89876790944 100644 --- a/test/test_transformers.py +++ b/test/test_transformers.py @@ -3145,7 +3145,6 @@ def test_cudnn_attention_broken_166211(self): self.assertFalse(dk.isnan().any()) self.assertFalse(dv.isnan().any()) - @skipIfRocm @unittest.skipIf(not PLATFORM_SUPPORTS_CUDNN_ATTENTION, "cudnn Attention is not supported on this system") def test_cudnn_attention_mask_broken_177842(self): # https://github.com/pytorch/pytorch/issues/177842 @@ -3356,8 +3355,6 @@ def test_scaled_dot_product_attention_fused_kernels_packed(self, device, type: s @parametrize("type", ["nested"]) @parametrize("is_contiguous", [True, False]) def test_scaled_dot_product_attention_cudnn_nested(self, device, type: str, is_contiguous: bool): - if TEST_WITH_ROCM and type == 'nested': - self.skipTest("ROCM does not support efficient attention on nested tensors, for now") make_tensor = partial(rand_sdpa_tensor, type=type, device=device, dtype=torch.float16, packed=True) batch_size, seq_len, num_heads, head_dim = 8, 64, 16, 64 @@ -3655,7 +3652,6 @@ def compiled_func(order): reset_order = torch._C._get_sdp_priority_order() self.assertEqual(default_order, reset_order, "expected SDPA context manager to reset priority order.") - @skipIfRocm # Missing deterministic algo @unittest.skipIf(not PLATFORM_SUPPORTS_FUSED_ATTENTION, "Fused SDPA was not built for this system") @parametrize("fused_kernel", PLATFORM_SPECIFIC_SDPA) @parametrize("warn_only", [True, False]) @@ -4378,9 +4374,6 @@ def test_fused_kernels_nested_broadcasting( rand_nested_tensor = partial(rand_sdpa_tensor, type="nested", device=device, dtype=dtype) batch, num_heads, head_dim = 32, 8, 64 head_dim_v = 32 if is_efficient else head_dim - if TEST_WITH_ROCM and head_dim != head_dim_v: - self.skipTest("head_dim != head_dim_v unsupported on ROCm for now") - return seq_lens_q = (torch.randint(low=1, high=5, size=(1,)).item() if expand_q_batch else torch.randint(low=1, high=32, size=(batch,)).tolist()) @@ -4444,7 +4437,6 @@ def _broadcast(t, batch_broadcasted, num_heads_broadcasted): self.assertEqual(actual.contiguous(), math_ref.contiguous().to(dtype), atol=1.5e-3, rtol=1e-2) - @skipIfRocm(msg="Efficient Attention on ROCM does not support head_dim != head_dim_v for now.") @unittest.skipIf(not PLATFORM_SUPPORTS_MEM_EFF_ATTENTION, "Fused SDPA was not built for this system") def test_fused_kernels_nested_broadcasting_query_dense(self, device): rand_nested_tensor = partial(rand_sdpa_tensor, type="nested", device=device, dtype=torch.float32) diff --git a/test/test_varlen_attention.py b/test/test_varlen_attention.py index dd382aea0bf05..f1283857a4d83 100644 --- a/test/test_varlen_attention.py +++ b/test/test_varlen_attention.py @@ -831,7 +831,6 @@ def test_batch_invariance( @unittest.skipIf( not PLATFORM_SUPPORTS_FLASH_ATTENTION, "Flash Attention not supported" ) - @unittest.skipIf(TEST_WITH_ROCM, "ROCm does not support seqused_k") @decorateIf( unittest.expectedFailure, lambda params: params["backend"] != "fa2" @@ -857,6 +856,8 @@ def test_seqused_k_kv_cache( self, device, dtype, actual_kv_lens, backend, sdpa_backend=None ): if TEST_WITH_ROCM: + if sdpa_backend == "ck": + self.skipTest("CK backend does not support seqused_k") torch.backends.cuda.preferred_rocm_fa_library(sdpa_backend) torch.manual_seed(42) @@ -961,7 +962,7 @@ def test_seqused_k_kv_cache( @unittest.skipIf( not PLATFORM_SUPPORTS_FLASH_ATTENTION, "Flash Attention not supported" ) - @unittest.skipIf(TEST_WITH_ROCM, "ROCm does not support seqused_k") + @unittest.skipIf(TEST_WITH_ROCM, "ROCm does not support block_table") @parametrize("dtype", [torch.bfloat16, torch.float16]) @parametrize("page_size", [32, 64, 128, 256]) @parametrize("compile", [False, True]) diff --git a/torch/nn/attention/varlen.py b/torch/nn/attention/varlen.py index 96cfc332b8699..454bb82c45e2a 100644 --- a/torch/nn/attention/varlen.py +++ b/torch/nn/attention/varlen.py @@ -168,15 +168,6 @@ def _varlen_attn_fake( (num_heads, total_q), dtype=torch.float, device=query.device ) - if torch.version.hip: - preferred = torch._C._get_rocm_fa_preferred_backend() - if preferred == torch._C._ROCmFABackend.AOTriton: - # AOTriton ROCm path uses batched 3D - batch_size = cu_seq_q.size(0) - 1 - logsumexp = torch.empty( - (batch_size, num_heads, max_q), dtype=torch.float, device=query.device - ) - rng_state = torch.empty((2,), dtype=torch.uint64, device=query.device) return output, logsumexp, rng_state @@ -416,14 +407,6 @@ def _varlen_attn_out_fake( (num_heads, total_q), dtype=torch.float, device=query.device ) - if torch.version.hip: - preferred = torch._C._get_rocm_fa_preferred_backend() - if preferred == torch._C._ROCmFABackend.AOTriton: - batch_size = cu_seq_q.size(0) - 1 - logsumexp = torch.empty( - (batch_size, num_heads, max_q), dtype=torch.float, device=query.device - ) - return logsumexp From 55719a5a3b1ff08b18c1fe20cec2e468bacd19d3 Mon Sep 17 00:00:00 2001 From: Prachi Gupta Date: Wed, 17 Jun 2026 13:42:30 -0500 Subject: [PATCH 08/11] Fixes previous PR: [ROCm] Bump AOTriton to 0.12.50tp (#3328) Co-authored-by: Xinya Zhang --- aten/src/ATen/Context.cpp | 6 +----- cmake/External/aotriton.cmake | 4 +--- 2 files changed, 2 insertions(+), 8 deletions(-) diff --git a/aten/src/ATen/Context.cpp b/aten/src/ATen/Context.cpp index e17352934f9ee..2a0f9cfcae86b 100644 --- a/aten/src/ATen/Context.cpp +++ b/aten/src/ATen/Context.cpp @@ -505,11 +505,7 @@ at::BlasBackend Context::blasPreferredBackend() { bool Context::ckSupported() { #ifdef USE_ROCM static const std::vector supported_archs = { -<<<<<<< xinyazhang/backport-aotriton-0.12b-2.12_gfx1250 - "gfx942", "gfx950" -======= - "gfx90a", "gfx942", "gfx950", "gfx1250", ->>>>>>> release/2.12_gfx1250 + "gfx942", "gfx950", "gfx1250", }; for (auto index : c10::irange(detail::getCUDAHooks().deviceCount())) { if(!detail::getCUDAHooks().isGPUArch(supported_archs, index)) { diff --git a/cmake/External/aotriton.cmake b/cmake/External/aotriton.cmake index 0feb798320b72..e9d3cab22a1d6 100644 --- a/cmake/External/aotriton.cmake +++ b/cmake/External/aotriton.cmake @@ -43,8 +43,7 @@ if(NOT __AOTRITON_INCLUDED) "amd-gfx110x" "amd-gfx115x" "amd-gfx120x" - # TODO: Update on AOTriton integration - #"amd-gfx1250" + "amd-gfx1250" ) set(__AOTRITON_IMAGE_SHA256_LIST "bb8bf2237b77fc503bc2967ea0d99d6ca419126c479e951ea42b712737128086" # amd-gfx90a @@ -54,7 +53,6 @@ if(NOT __AOTRITON_INCLUDED) "3177387a15c678b30057f4584d1fc1b8f8db56163890cb5c98f27450209f5a7b" # amd-gfx115x "68572511ce6487a83f9014bd255bd69c8943f87d0c93bd57b2daac5fbc6c79c1" # amd-gfx120x "f2415cc328430a84287cb3f545d6965811b2fa54fe9986d91357b3d455f50b12" # amd-gfx1250 - ) set(__AOTRITON_BASE_URL "$ENV{PYTORCH_AOTRITON_BASE_URL}") if(NOT __AOTRITON_BASE_URL) From bbecc5657577e15f0c0aa057daf34b4e4be41c31 Mon Sep 17 00:00:00 2001 From: Prachi Gupta Date: Wed, 17 Jun 2026 15:22:51 -0500 Subject: [PATCH 09/11] Turn off MSLK's CK Kernels for gfx1250 (#3329) --- aten/src/ATen/CMakeLists.txt | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/aten/src/ATen/CMakeLists.txt b/aten/src/ATen/CMakeLists.txt index b030288ef5070..4b770af863d01 100644 --- a/aten/src/ATen/CMakeLists.txt +++ b/aten/src/ATen/CMakeLists.txt @@ -437,10 +437,10 @@ IF(USE_MSLK) list(PREPEND MSLK_EXTRA_HIPCC_FLAGS -mllvm -amdgpu-coerce-illegal-types=1) endif() - # Only compile for gfx942, gfx950, and gfx1250. + # Only compile for gfx942 and gfx950 (composable_kernel lacks gfx1250 support). set(HIP_CLANG_FLAGS_ORIGINAL ${HIP_CLANG_FLAGS}) string(REGEX REPLACE "--offload-arch=[^ ]*" "" FILTERED_HIP_CLANG_FLAGS "${HIP_CLANG_FLAGS}") - foreach(ARCH gfx942 gfx950 gfx1250) + foreach(ARCH gfx942 gfx950) if(${ARCH} IN_LIST PYTORCH_ROCM_ARCH) list(APPEND FILTERED_HIP_CLANG_FLAGS --offload-arch=${ARCH}) endif() From 04b54055439beb8a156f244c3fd3cdb9e31a1d3b Mon Sep 17 00:00:00 2001 From: Ramya Ramineni <62723901+rraminen@users.noreply.github.com> Date: Wed, 17 Jun 2026 16:31:46 -0500 Subject: [PATCH 10/11] [release/2.12_gfx1250] Support for gfx1250 - bug fixes (#3330) This PR is to address the reviewed comments on PR https://github.com/ROCm/pytorch/pull/3307 --- .ci/docker/build.sh | 2 - .ci/pytorch/binary_populate_env.sh | 2 +- .github/actionlint.yaml | 4 -- .github/workflows/inductor-rocm-gfx1250.yml | 77 --------------------- CMakeLists.txt | 2 +- requirements.txt | 1 - 6 files changed, 2 insertions(+), 86 deletions(-) delete mode 100644 .github/workflows/inductor-rocm-gfx1250.yml diff --git a/.ci/docker/build.sh b/.ci/docker/build.sh index b7b74c3a0794f..41589ef4c7031 100755 --- a/.ci/docker/build.sh +++ b/.ci/docker/build.sh @@ -163,8 +163,6 @@ case "$tag" in ROCM_VERSION=7.2 TRITON=yes KATEX=yes - UCX_COMMIT=${_UCX_COMMIT} - UCC_COMMIT=${_UCC_COMMIT} PYTORCH_ROCM_ARCH="gfx90a;gfx942;gfx950;gfx1100;gfx1250" if [[ $tag =~ "benchmarks" ]]; then INDUCTOR_BENCHMARKS=yes diff --git a/.ci/pytorch/binary_populate_env.sh b/.ci/pytorch/binary_populate_env.sh index 4912b444bcca9..53914914c8c93 100755 --- a/.ci/pytorch/binary_populate_env.sh +++ b/.ci/pytorch/binary_populate_env.sh @@ -147,7 +147,7 @@ export PYTORCH_EXTRA_INSTALL_REQUIREMENTS="${PYTORCH_EXTRA_INSTALL_REQUIREMENTS: # TODO: We don't need this anymore IIUC export TORCH_PACKAGE_NAME='torch' -export USE_FBGEMM=0 +export USE_FBGEMM=1 export PIP_UPLOAD_FOLDER="$PIP_UPLOAD_FOLDER" export DOCKER_IMAGE="$DOCKER_IMAGE" diff --git a/.github/actionlint.yaml b/.github/actionlint.yaml index 81690f184271a..b894d7214fced 100644 --- a/.github/actionlint.yaml +++ b/.github/actionlint.yaml @@ -62,13 +62,9 @@ self-hosted-runner: # gfx942 runners - linux.rocm.gpu.gfx942.1 - linux.rocm.gpu.gfx942.4 - - linux.rocm.gfx942.docker-cache # gfx950 runners - linux.rocm.gpu.gfx950.1 - linux.rocm.gpu.gfx950.4 - # gfx1250 runners - - linux.rocm.gpu.gfx1250.1 - - linux.rocm.gpu.gfx1250.4 # Org wise AWS `mac2.metal` runners (2020 Mac mini hardware powered by Apple silicon M1 processors) - macos-m1-stable - macos-m1-14 diff --git a/.github/workflows/inductor-rocm-gfx1250.yml b/.github/workflows/inductor-rocm-gfx1250.yml deleted file mode 100644 index 33449a7d52bf1..0000000000000 --- a/.github/workflows/inductor-rocm-gfx1250.yml +++ /dev/null @@ -1,77 +0,0 @@ -# The name of this file is subject to change to stay consistent with other .yml files. -# -# The MI355 workflow (.github/workflows/inductor-rocm-mi355.yml) uses: -# - _linux-build.yml and _rocm-test.yml reusable workflows -# - Build environment linux-noble-rocm-py3.12-mi355 -# - Runner label linux.rocm.gpu.gfx950.1 -# - Docker image ci-image:pytorch-linux-noble-rocm-n-py3 -# - 2-shard test matrix for the inductor config -# -# The GFX1250 equivalent is following this exact pattern. - -name: inductor-rocm-gfx1250 - -on: - push: - branches: - - main - - release/* - tags: - - ciflow/inductor-rocm-gfx1250/* - workflow_dispatch: - -concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' }}-${{ github.event_name == 'schedule' }} - cancel-in-progress: true - -permissions: - id-token: write - contents: read - actions: read - -jobs: - target-determination: - if: github.repository_owner == 'pytorch' - name: before-test - uses: ./.github/workflows/target_determination.yml - - get-label-type: - name: get-label-type - uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@release/2.11 - if: ${{ (github.event_name != 'schedule' || github.repository == 'pytorch/pytorch') && github.repository_owner == 'pytorch' }} - with: - triggering_actor: ${{ github.triggering_actor }} - issue_owner: ${{ github.event.pull_request.user.login || github.event.issue.user.login }} - curr_branch: ${{ github.head_ref || github.ref_name }} - curr_ref_type: ${{ github.ref_type }} - opt_out_experiments: lf - - linux-noble-rocm-py3_12-inductor-build: - name: linux-noble-rocm-py3.12-gfx1250 - uses: ./.github/workflows/_linux-build.yml - needs: get-label-type - with: - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - build-environment: linux-noble-rocm-py3.12-gfx1250 - # Docker image stays the same as MI355 because ROCm image supports multiple arches. - docker-image-name: ci-image:pytorch-linux-noble-rocm-n-py3 - # Set PYTORCH_ROCM_ARCH directly in the workflow YAML as an env variable, - # so build.sh never needs to parse BUILD_ENVIRONMENT. - #env-var-script: | - # export PYTORCH_ROCM_ARCH=gfx1250 - test-matrix: | - { include: [ - { config: "inductor", shard: 1, num_shards: 2, runner: "linux.rocm.gpu.gfx1250.1" }, # It requires provisioning hardware. - { config: "inductor", shard: 2, num_shards: 2, runner: "linux.rocm.gpu.gfx1250.1" }, - ]} - secrets: inherit - - linux-noble-rocm-py3_12-inductor-test: - name: linux-noble-rocm-py3.12-gfx1250 - uses: ./.github/workflows/_rocm-test.yml - needs: linux-noble-rocm-py3_12-inductor-build - with: - build-environment: ${{ needs.linux-noble-rocm-py3_12-inductor-build.outputs.build-environment }} - docker-image: ${{ needs.linux-noble-rocm-py3_12-inductor-build.outputs.docker-image }} - test-matrix: ${{ needs.linux-noble-rocm-py3_12-inductor-build.outputs.test-matrix }} - secrets: inherit diff --git a/CMakeLists.txt b/CMakeLists.txt index 215009d868e9f..8e66ebcb694df 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -275,7 +275,7 @@ cmake_dependent_option(USE_CUSPARSELT "Use cuSPARSELt" ON "USE_CUDA" OFF) cmake_dependent_option(USE_CUDSS "Use cuDSS" ON "USE_CUDA" OFF) # USE_ROCM is guarded against in Dependencies.cmake because USE_ROCM is not properly defined here cmake_dependent_option(USE_CUFILE "Use cuFile" ON "USE_CUDA AND NOT WIN32" OFF) -option(USE_FBGEMM "Use FBGEMM (quantized 8-bit server operators)" OFF) +option(USE_FBGEMM "Use FBGEMM (quantized 8-bit server operators)" ON) option(USE_KINETO "Use Kineto profiling library" ON) option(USE_CUPTI_SO "Use CUPTI as a shared library" ON) option(USE_GFLAGS "Use GFLAGS" OFF) diff --git a/requirements.txt b/requirements.txt index 13430c258d109..d6552edad483b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -14,7 +14,6 @@ lintrunner==0.12.11 ; platform_machine != "s390x" networkx==2.8.8 optree==0.13.0 ; python_version < "3.14" optree==0.17.0 ; python_version >= "3.14" -pandas psutil==7.2.2 spin==0.17 sympy==1.13.3 From 43a62ac6cc17a57487826c1b7b0f6e7cf96a43c1 Mon Sep 17 00:00:00 2001 From: Ramya Ramineni <62723901+rraminen@users.noreply.github.com> Date: Wed, 17 Jun 2026 18:58:23 -0500 Subject: [PATCH 11/11] [release/2.12_gfx1250] Fix typo (#3331) Fix typo --- test/inductor/test_max_autotune.py | 1 - 1 file changed, 1 deletion(-) diff --git a/test/inductor/test_max_autotune.py b/test/inductor/test_max_autotune.py index 9526272c2eeb2..c6f1dac05d280 100644 --- a/test/inductor/test_max_autotune.py +++ b/test/inductor/test_max_autotune.py @@ -550,7 +550,6 @@ def test_tdm_arch_gate_accepts_only_gfx1250(self): self.assertTrue(is_gfx1250_arch("gfx1250")) self.assertTrue(is_gfx1250_arch("gfx1250:sramecc+:xnack-")) self.assertFalse(is_gfx1250_arch("gfx1251")) - self.assertFalse(is_gfx1250_arch("gfx1260")) self.assertFalse(is_gfx1250_arch("amd-gfx1250")) def test_tdm_persistent_template_precedes_rocm_tma_fallback(self):