Update ck and enable test (#16383)

Apply the fix in https://github.com/ROCmSoftwarePlatform/composable_kernel/issues/728 Introduce more kernel instances and allow the introduction of streamk and splitk.
2026-07-17 18:40:28 +00:00 · 2023-08-22 11:08:55 +08:00 · 2023-08-22 11:08:55 +08:00 · 4e6cec4d09
commit 4e6cec4d09
parent aae9a52e8b
14 changed files with 54 additions and 44 deletions
--- a/cmake/external/composable_kernel.cmake
+++ b/cmake/external/composable_kernel.cmake
@ -1,5 +1,5 @@
 set(composable_kernel_URL https://github.com/ROCmSoftwarePlatform/composable_kernel.git)
-set(composable_kernel_TAG ed3a2e52265e11daa366f47b082141a652b67c58) # 2023-04-10 21:02:17 +0800
+set(composable_kernel_TAG d52ec01652b7d620386251db92455968d8d90bdc) # 2023-08-18 11:14:59 +0800

 set(PATCH ${PROJECT_SOURCE_DIR}/patches/composable_kernel/Fix_Clang_Build.patch)

@ -14,10 +14,14 @@ FetchContent_GetProperties(composable_kernel)
 if(NOT composable_kernel_POPULATED)
  FetchContent_Populate(composable_kernel)
  set(BUILD_DEV OFF CACHE BOOL "Disable -Weverything, otherwise, error: 'constexpr' specifier is incompatible with C++98 [-Werror,-Wc++98-compat]" FORCE)
+  # Exclude i8 device gemm instances due to excessive long compilation time and not being used
+  set(DTYPES fp32 fp16 bf16)
+  set(INSTANCES_ONLY ON)
  add_subdirectory(${composable_kernel_SOURCE_DIR} ${composable_kernel_BINARY_DIR} EXCLUDE_FROM_ALL)

  add_library(onnxruntime_composable_kernel_includes INTERFACE)
  target_include_directories(onnxruntime_composable_kernel_includes INTERFACE
    ${composable_kernel_SOURCE_DIR}/include
    ${composable_kernel_SOURCE_DIR}/library/include)
+  target_compile_definitions(onnxruntime_composable_kernel_includes INTERFACE __fp32__ __fp16__ __bf16__)
 endif()
--- a/cmake/patches/composable_kernel/Fix_Clang_Build.patch
+++ b/cmake/patches/composable_kernel/Fix_Clang_Build.patch
@ -1,5 +1,5 @@
 diff --git a/CMakeLists.txt b/CMakeLists.txt
-index f861e3020..f0b6bceae 100644
+index 514b98fde..59c8a568a 100644
 --- a/CMakeLists.txt
 +++ b/CMakeLists.txt
@@ -1,7 +1,7 @@
@ -11,7 +11,7 @@ index f861e3020..f0b6bceae 100644
 
 list(APPEND CMAKE_MODULE_PATH "${PROJECT_SOURCE_DIR}/cmake")
 
-@@ -41,27 +41,6 @@ set(CMAKE_CXX_STANDARD_REQUIRED ON)
+@@ -94,27 +94,6 @@ set(CMAKE_CXX_STANDARD_REQUIRED ON)
 set(CMAKE_CXX_EXTENSIONS OFF)
 message("CMAKE_CXX_COMPILER_ID: ${CMAKE_CXX_COMPILER_ID}")
 
@ -39,7 +39,7 @@ index f861e3020..f0b6bceae 100644
 ## HIP
 find_package(HIP REQUIRED)
 # Override HIP version in config.h, if necessary.
-@@ -83,8 +62,6 @@ if( DEFINED CK_OVERRIDE_HIP_VERSION_PATCH )
+@@ -136,8 +115,6 @@ if( DEFINED CK_OVERRIDE_HIP_VERSION_PATCH )
     message(STATUS "CK_HIP_VERSION_PATCH overriden with ${CK_OVERRIDE_HIP_VERSION_PATCH}")
 endif()
 message(STATUS "Build with HIP ${HIP_VERSION}")
@ -48,17 +48,7 @@ index f861e3020..f0b6bceae 100644
 
 ## tidy
 include(EnableCompilerWarnings)
-@@ -273,9 +250,6 @@ rocm_package_setup_component(profiler
- )
- 
- add_subdirectory(library)
-add_subdirectory(example)
-add_subdirectory(test)
-add_subdirectory(profiler)
- 
- #Create an interface target for the include only files and call it "composablekernels"
- include(CMakePackageConfigHelpers)
-@@ -301,11 +275,3 @@ rocm_install(FILES
+@@ -391,11 +368,3 @@ rocm_install(FILES
 
 set(CPACK_RESOURCE_FILE_LICENSE "${CMAKE_CURRENT_SOURCE_DIR}/LICENSE")
 set(CPACK_RPM_PACKAGE_LICENSE "MIT")
@ -71,7 +61,7 @@ index f861e3020..f0b6bceae 100644
 -    HEADER_ONLY
 -)
 diff --git a/library/src/tensor_operation_instance/gpu/CMakeLists.txt b/library/src/tensor_operation_instance/gpu/CMakeLists.txt
-index c206c4dc0..b283eeb64 100644
+index 1d54a141b..4edd7dbfb 100644
 --- a/library/src/tensor_operation_instance/gpu/CMakeLists.txt
 +++ b/library/src/tensor_operation_instance/gpu/CMakeLists.txt
@@ -1,7 +1,13 @@
--- a/onnxruntime/contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_ck_impl/impl.cuh
+++ b/onnxruntime/contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_ck_impl/impl.cuh
@ -52,10 +52,11 @@ template <ck::index_t NumDimG,
 using device_batched_gemm_softmax_gemm_permute_instances =
    std::tuple<
        // clang-format off
-        // #############################################|  NumDimG| NumDimM| NumDimN| NumDimK| NumDimO| AData| B0Data| B1Data| CData| Acc0BiasData| Acc1BiasData| AccData| CShuffle|           A|          B0|        Acc0|          B1|           C|           GEMM|   ATensorSpec|  B0TensorSpec|  B1TensorSpec|   CTensorSpec| NumGemmK| Block| Gemm01| Gemm0| Gemm0| Gemm1| Gemm1| AK1| BK1| B1K1| MPer| NPer| Gemm0| Gemm0| Gemm1|  ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockLds|  B0BlockTransfer| B0BlockTransfer| B0BlockTransfer| B0BlockTransfer| B0BlockTransfer| B0BlockTransfer| B0BlockLds|  B1BlockTransfer| B1BlockTransfer| B1BlockTransfer| B1BlockTransfer| B1BlockTransfer| B1BlockTransfer| B1BlockLds|    CShuffle|    CShuffle| CBlockTransferClusterLengths|  CBlockTransfer| MaskingSpec|
-        // #############################################|         |        |        |        |        |  Type|   Type|   Type|  Type|         Type|         Type|    Type| DataType| Elementwise| Elementwise| Elementwise| Elementwise| Elementwise| Specialization|              |              |              |              | Prefetch|  Size|   MPer|  NPer|  KPer|  NPer|  KPer|    |    |     |  XDL|  XDL|  MXdl|  NXdl|  NXdl|   ThreadCluster|  ThreadCluster| SrcAccessOrder|   SrcVectorDim|      SrcScalar|      DstScalar| AddExtraM|    ThreadCluster|   ThreadCluster|  SrcAccessOrder|    SrcVectorDim|       SrcScalar|       DstScalar|  AddExtraN|    ThreadCluster|   ThreadCluster|  SrcAccessOrder|    SrcVectorDim|       SrcScalar|       DstScalar|  AddExtraN| MXdlPerWave| NXdlPerWave|         _MBlock_MWaveMPerXdl| ScalarPerVector|            |
-        // #############################################|         |        |        |        |        |      |       |       |      |             |             |        |         |   Operation|   Operation|   Operation|   Operation|   Operation|               |              |              |              |              |    Stage|      |  Block| Block| Block| Block| Block|    |    |     |     |     |   Per|   Per|   Per| Lengths_K0_M_K1|   ArrangeOrder|               |               |      PerVector|   PerVector_K1|          |  Lengths_K0_N_K1|    ArrangeOrder|                |                |       PerVector|    PerVector_K1|           |  Lengths_K0_N_K1|    ArrangeOrder|                |                |       PerVector|    PerVector_K1|           |  PerShuffle|  PerShuffle|         _NBlock_NWaveNPerXdl|   _NWaveNPerXdl|            |
-        // #############################################|         |        |        |        |        |      |       |       |      |             |             |        |         |            |            |            |            |            |               |              |              |              |              |         |      |       |      |      |      |      |    |    |     |     |     |  Wave|  Wave|  Wave|                |               |               |               |               |               |          |                 |                |                |                |                |                |           |                 |                |                |                |                |                |           |            |            |                             |                |            |
+        // #############################################|  NumDimG| NumDimM| NumDimN| NumDimK| NumDimO| AData| B0Data| B1Data| CData| Acc0BiasData| Acc1BiasData| AccData| CShuffle|           A|          B0|        Acc0|          B1|           C|           GEMM|   ATensorSpec|  B0TensorSpec|  B1TensorSpec|   CTensorSpec| NumGemmK| Block| Gemm01| Gemm0| Gemm0| Gemm1| Gemm1| AK1| BK1| B1K1| MPer| NPer| Gemm0| Gemm0| Gemm1|  ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockLds|  B0BlockTransfer| B0BlockTransfer| B0BlockTransfer| B0BlockTransfer| B0BlockTransfer| B0BlockTransfer| B0BlockLds|  B1BlockTransfer| B1BlockTransfer| B1BlockTransfer| B1BlockTransfer| B1BlockTransfer| B1BlockTransfer| B1BlockLds|    CShuffle|    CShuffle| CBlockTransferClusterLengths|  CBlockTransfer| MaskingSpec|  D0s Bias|
+        // #############################################|         |        |        |        |        |  Type|   Type|   Type|  Type|         Type|         Type|    Type| DataType| Elementwise| Elementwise| Elementwise| Elementwise| Elementwise| Specialization|              |              |              |              | Prefetch|  Size|   MPer|  NPer|  KPer|  NPer|  KPer|    |    |     |  XDL|  XDL|  MXdl|  NXdl|  NXdl|   ThreadCluster|  ThreadCluster| SrcAccessOrder|   SrcVectorDim|      SrcScalar|      DstScalar| AddExtraM|    ThreadCluster|   ThreadCluster|  SrcAccessOrder|    SrcVectorDim|       SrcScalar|       DstScalar|  AddExtraN|    ThreadCluster|   ThreadCluster|  SrcAccessOrder|    SrcVectorDim|       SrcScalar|       DstScalar|  AddExtraN| MXdlPerWave| NXdlPerWave|         _MBlock_MWaveMPerXdl| ScalarPerVector|            | SrcScalar|
+        // #############################################|         |        |        |        |        |      |       |       |      |             |             |        |         |   Operation|   Operation|   Operation|   Operation|   Operation|               |              |              |              |              |    Stage|      |  Block| Block| Block| Block| Block|    |    |     |     |     |   Per|   Per|   Per| Lengths_K0_M_K1|   ArrangeOrder|               |               |      PerVector|   PerVector_K1|          |  Lengths_K0_N_K1|    ArrangeOrder|                |                |       PerVector|    PerVector_K1|           |  Lengths_K0_N_K1|    ArrangeOrder|                |                |       PerVector|    PerVector_K1|           |  PerShuffle|  PerShuffle|         _NBlock_NWaveNPerXdl|   _NWaveNPerXdl|            | PerVector|
+        // #############################################|         |        |        |        |        |      |       |       |      |             |             |        |         |            |            |            |            |            |               |              |              |              |              |         |      |       |      |      |      |      |    |    |     |     |     |  Wave|  Wave|  Wave|                |               |               |               |               |               |          |                 |                |                |                |                |                |           |                 |                |                |                |                |                |           |            |            |                             |                |            |          |
+        DeviceBatchedGemmSoftmaxGemmPermute_Xdl_CShuffle<  NumDimG, NumDimM, NumDimN, NumDimK, NumDimO,    DT,     DT,     DT,    DT,        D0sDT,  ck::Tuple<>,   AccDT,       DT, PassThrough, PassThrough,        D0Op, PassThrough, PassThrough,     GemmPadded, TensorDefault, TensorDefault, TensorDefault, TensorDefault,        1,   256,    128,    64,    32,   128,    32,   8,   8,    2,   32,   32,     1,     2,     4,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,      true,      S<4, 64, 1>,      S<1, 0, 2>,      S<1, 0, 2>,               2,               8,               8,       true,     S< 8, 32, 1>,      S<0, 2, 1>,      S<0, 2, 1>,               1,               4,               2,      false,           1,           2,               S<1, 32, 1, 8>,               8, MaskingSpec,       1>,
        DeviceBatchedGemmSoftmaxGemmPermute_Xdl_CShuffle<  NumDimG, NumDimM, NumDimN, NumDimK, NumDimO,    DT,     DT,     DT,    DT,        D0sDT,  ck::Tuple<>,   AccDT,       DT, PassThrough, PassThrough,        D0Op, PassThrough, PassThrough,    GemmDefault, TensorDefault, TensorDefault, TensorDefault, TensorDefault,        1,   256,    256,   128,    32,    64,    32,   8,   8,    2,   32,   32,     2,     4,     2,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,      true,      S<4, 64, 1>,      S<1, 0, 2>,      S<1, 0, 2>,               2,               8,               8,       true,     S<16, 16, 1>,      S<0, 2, 1>,      S<0, 2, 1>,               1,               4,               2,      false,           1,           2,               S<1, 32, 1, 8>,               8, MaskingSpec>,
        DeviceBatchedGemmSoftmaxGemmPermute_Xdl_CShuffle<  NumDimG, NumDimM, NumDimN, NumDimK, NumDimO,    DT,     DT,     DT,    DT,        D0sDT,  ck::Tuple<>,   AccDT,       DT, PassThrough, PassThrough,        D0Op, PassThrough, PassThrough,    GemmDefault, TensorDefault, TensorDefault, TensorDefault, TensorDefault,        1,   256,    256,   128,    32,   128,    32,   8,   8,    2,   32,   32,     2,     4,     4,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,      true,      S<4, 64, 1>,      S<1, 0, 2>,      S<1, 0, 2>,               2,               8,               8,       true,     S< 8, 32, 1>,      S<0, 2, 1>,      S<0, 2, 1>,               1,               4,               2,      false,           1,           2,               S<1, 32, 1, 8>,               8, MaskingSpec>,
 #if ROCM_VERSION >= 50500
@ -71,6 +72,7 @@ using device_batched_gemm_softmax_gemm_permute_instances =
        DeviceBatchedGemmSoftmaxGemmPermute_Xdl_CShuffle<  NumDimG, NumDimM, NumDimN, NumDimK, NumDimO,    DT,     DT,     DT,    DT,        D0sDT,  ck::Tuple<>,   AccDT,       DT, PassThrough, PassThrough,        D0Op, PassThrough, PassThrough,    GemmDefault, TensorDefault, TensorDefault, TensorDefault, TensorDefault,        1,   256,     64,   256,    64,   128,    32,   8,   8,    2,   16,   16,     1,    16,     8,     S<8, 32, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,      true,      S<8, 32, 1>,      S<1, 0, 2>,      S<1, 0, 2>,               2,               8,               8,       true,     S< 8, 32, 1>,      S<0, 2, 1>,      S<0, 2, 1>,               1,               4,               2,      false,           1,           8,               S<1, 16, 1,16>,               8, MaskingSpec>,
        DeviceBatchedGemmSoftmaxGemmPermute_Xdl_CShuffle<  NumDimG, NumDimM, NumDimN, NumDimK, NumDimO,    DT,     DT,     DT,    DT,        D0sDT,  ck::Tuple<>,   AccDT,       DT, PassThrough, PassThrough,        D0Op, PassThrough, PassThrough,    GemmDefault, TensorDefault, TensorDefault, TensorDefault, TensorDefault,        1,   256,     64,   256,    64,    64,    32,   8,   8,    2,   16,   16,     1,    16,     4,     S<8, 32, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,      true,      S<8, 32, 1>,      S<1, 0, 2>,      S<1, 0, 2>,               2,               8,               8,       true,     S<16, 16, 1>,      S<0, 2, 1>,      S<0, 2, 1>,               1,               4,               2,      false,           1,           4,               S<1, 32, 1, 8>,               8, MaskingSpec>,
        // Padded fallback kernel
+        DeviceBatchedGemmSoftmaxGemmPermute_Xdl_CShuffle<  NumDimG, NumDimM, NumDimN, NumDimK, NumDimO,    DT,     DT,     DT,    DT,        D0sDT,  ck::Tuple<>,   AccDT,       DT, PassThrough, PassThrough,        D0Op, PassThrough, PassThrough,     GemmPadded, TensorDefault, TensorDefault, TensorDefault, TensorDefault,        1,   256,    128,   128,    64,   128,    32,   8,   8,    2,   32,   32,     1,     4,     4,     S<8, 32, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,     false,      S<8, 32, 1>,      S<1, 0, 2>,      S<1, 0, 2>,               2,               8,               8,      false,     S< 8, 32, 1>,      S<0, 2, 1>,      S<0, 2, 1>,               1,               4,               2,      false,           1,           2,               S<1, 32, 1, 8>,               8, MaskingSpec,       1>,
        DeviceBatchedGemmSoftmaxGemmPermute_Xdl_CShuffle<  NumDimG, NumDimM, NumDimN, NumDimK, NumDimO,    DT,     DT,     DT,    DT,        D0sDT,  ck::Tuple<>,   AccDT,       DT, PassThrough, PassThrough,        D0Op, PassThrough, PassThrough,     GemmPadded, TensorDefault, TensorDefault, TensorDefault, TensorDefault,        1,   256,    128,   128,    64,   128,    32,   8,   8,    2,   32,   32,     1,     4,     4,     S<8, 32, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,     false,      S<8, 32, 1>,      S<1, 0, 2>,      S<1, 0, 2>,               2,               8,               8,      false,     S< 8, 32, 1>,      S<0, 2, 1>,      S<0, 2, 1>,               1,               4,               2,      false,           1,           2,               S<1, 32, 1, 8>,               8, MaskingSpec>,
        DeviceBatchedGemmSoftmaxGemmPermute_Xdl_CShuffle<  NumDimG, NumDimM, NumDimN, NumDimK, NumDimO,    DT,     DT,     DT,    DT,        D0sDT,  ck::Tuple<>,   AccDT,       DT, PassThrough, PassThrough,        D0Op, PassThrough, PassThrough,     GemmPadded, TensorDefault, TensorDefault, TensorDefault, TensorDefault,        1,   256,    128,    64,    32,   128,    32,   8,   8,    2,   32,   32,     1,     2,     4,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,      true,      S<4, 64, 1>,      S<1, 0, 2>,      S<1, 0, 2>,               2,               8,               8,       true,     S< 8, 32, 1>,      S<0, 2, 1>,      S<0, 2, 1>,               1,               4,               2,      false,           1,           2,               S<1, 32, 1, 8>,               8, MaskingSpec>
        // clang-format on
--- a/onnxruntime/core/framework/tunable.h
+++ b/onnxruntime/core/framework/tunable.h
@ -268,6 +268,10 @@ class TunableOp {
      WarmUp(candidate, params);

      auto approx_duration = Profile(candidate, params, approx_num_iter);
+      if (approx_duration > 2 * min_time) {
+        LOGS_DEFAULT(VERBOSE) << "FindFastestImpl skip slow instance " << op_sig << '(' << param_sig << ") id=" << i;
+        continue;
+      }
      int tuning_iter = std::max(1, int(std::min(double(max_tuning_iter), ctx->GetMaxTuningDurationMs() / approx_duration)));

      LOGS_DEFAULT(VERBOSE) << "FindFastestImpl run instance " << op_sig << '(' << param_sig << ") id=" << i << " " << tuning_iter << " times.";
@ -278,7 +282,7 @@ class TunableOp {
        id = static_cast<int>(i);
      }
    }
-    ORT_ENFORCE(id >= 0, "Cannot found viable op");
+    ORT_ENFORCE(id >= 0, "Could not find viable op");
    LOGS_DEFAULT(VERBOSE) << "FindFastestImpl for " << op_sig << '(' << param_sig << ") found fastest with id=" << id;
    std::this_thread::sleep_for(std::chrono::milliseconds(50));
    return id;
--- a/onnxruntime/core/providers/rocm/math/softmax_ck.cuh
+++ b/onnxruntime/core/providers/rocm/math/softmax_ck.cuh
@ -34,7 +34,7 @@ auto GetCKSoftmaxTypeStringAndOps() {
  using OutDataType = typename CKDataTypeAdaptor<OutputT>::type;
  using AccDataType = typename CKDataTypeAdaptor<AccT>::type;
  using DeviceSoftmax = ck::tensor_operation::device::
-      DeviceSoftmax<InDataType, AccDataType, OutDataType, Nop, Nop, Rank>;
+      DeviceSoftmax<InDataType, AccDataType, OutDataType, Nop, Nop, Rank, NumReduceDim>;
  using InstanceFactory = ck::tensor_operation::device::instance::DeviceOperationInstanceFactory<DeviceSoftmax>;

  std::vector<std::pair<std::string, tunable::Op<SoftmaxParams<InputT, OutputT>>>> ret;
@ -49,9 +49,6 @@ auto GetCKSoftmaxTypeStringAndOps() {
      TUNABLE_OP_RETURN_UNSUPPORTED_ARGUMENT_IF(
          params->is_log_softmax,
          impl->GetTypeString(), " does not support log softmax");
-      TUNABLE_OP_RETURN_UNSUPPORTED_ARGUMENT_IF(
-          impl->GetRank() != Rank || impl->GetNumReduceDim() != NumReduceDim,
-          impl->GetTypeString(), " does not support current Rank or NumReduceDim ", params->Signature());

      std::vector<ck::index_t> in_lengths{1, 1, params->batch_count, params->softmax_elements};
      std::vector<ck::index_t> in_strides{params->batch_count * params->input_stride, params->batch_count * params->input_stride, params->input_stride, 1};
--- a/onnxruntime/core/providers/rocm/tunable/gemm_ck.cuh
+++ b/onnxruntime/core/providers/rocm/tunable/gemm_ck.cuh
@ -46,17 +46,6 @@ auto GetCKGemmTypeStringAndOps() {
  std::vector<std::pair<std::string, Op<GemmParams<T>>>> ret;
  for (auto&& impl : InstanceFactory::GetInstances()) {
    auto type_string = impl->GetTypeString();
-
-    // FIXME: ck upstream have bugs in some input shapes coupled with specific impls. The `IsSupportedArgument` is not
-    // sound, we exclude those implementation here for now. Check back later when AMD fixed them.
-    //
-    // The DeviceGemmXdl<256, 128, 144, 8, 8, 16, 16, 2, 9> and DeviceGemmXdl<256, 128, 144, 4, 8, 16, 16, 2, 9> only
-    // occurs in DeviceGemm<Row, Col> for FP16. When k < 8, the result is wrong.
-    if (type_string == "DeviceGemmXdl<256, 128, 144, 8, 8, 16, 16, 2, 9>" ||
-        type_string == "DeviceGemmXdl<256, 128, 144, 4, 8, 16, 16, 2, 9>") {
-      continue;
-    }
-
    auto invoker = impl->MakeInvokerPointer();
    auto ck_gemm_op = [impl = std::move(impl), invoker = std::move(invoker)](const GemmParams<T>* params) -> Status {
      auto one = ToHipType<T>::FromFloat(1.0f);
--- a/onnxruntime/python/tools/kernel_explorer/kernel_explorer_interface.h
+++ b/onnxruntime/python/tools/kernel_explorer/kernel_explorer_interface.h
@ -94,6 +94,22 @@ class IKernelExplorer {
  int repeats_{100};
 };

+class WithMaxTuningDurationMs {
+ public:
+  WithMaxTuningDurationMs(TuningContextT* ctx, int ms) : ctx_(ctx) {
+    original_tuning_duration_ = ctx_->GetMaxTuningDurationMs();
+    ctx_->SetMaxTuningDurationMs(ms);
+  }
+
+  ~WithMaxTuningDurationMs() {
+    ctx_->SetMaxTuningDurationMs(original_tuning_duration_);
+  }
+
+ private:
+  TuningContextT* ctx_;
+  int original_tuning_duration_;
+};
+
 pybind11::module GetKernelExplorerModule();

 class KernelExplorerInit {
--- a/onnxruntime/python/tools/kernel_explorer/kernels/gemm_softmax_gemm_permute_test.py
+++ b/onnxruntime/python/tools/kernel_explorer/kernels/gemm_softmax_gemm_permute_test.py
@ -14,6 +14,8 @@ import numpy as np
 import pytest
 from utils import dtype_to_suffix, matmul, softmax

+max_batch_size = int(os.environ.get("KERNEL_EXPLORER_BATCHED_GEMM_MAX_BATCH_SIZE", 64))
+

 def multinormal_distribution(num_distribution, num_element_per_dist):
    arrays = []
@ -36,7 +38,7 @@ def get_ck_binding_name(dtype, biased: bool, masked: bool):


 dtypes = ["float16"]
-batches = [1, 64]
+batches = [1, max_batch_size]
 seqlens = [128, 512]
 total_seqlens = [128, 512]
 num_heads = [8, 12]
--- a/onnxruntime/python/tools/kernel_explorer/kernels/rocm/elementwise.cu
+++ b/onnxruntime/python/tools/kernel_explorer/kernels/rocm/elementwise.cu
@ -64,6 +64,7 @@ class ElementwiseTunable : public IKernelExplorer {
  }

  void Run() override {
+    WithMaxTuningDurationMs max_duration(TuningContext(), 250);
    ORT_THROW_IF_ERROR(op_(&params_));
  }

--- a/onnxruntime/python/tools/kernel_explorer/kernels/rocm/gemm_fast_gelu_tunable.cu
+++ b/onnxruntime/python/tools/kernel_explorer/kernels/rocm/gemm_fast_gelu_tunable.cu
@ -56,6 +56,7 @@ class GemmFastGeluTunable : public IKernelExplorer {
  }

  void Run() override {
+    WithMaxTuningDurationMs max_duration(TuningContext(), 250);
    ORT_THROW_IF_ERROR((op_(&params_)));
  }

--- a/onnxruntime/python/tools/kernel_explorer/kernels/rocm/gemm_tunable.cu
+++ b/onnxruntime/python/tools/kernel_explorer/kernels/rocm/gemm_tunable.cu
@ -56,6 +56,7 @@ class GemmTunable : public IKernelExplorer {
  }

  void Run() override {
+    WithMaxTuningDurationMs max_duration(TuningContext(), 250);
    ORT_THROW_IF_ERROR(op_(&params_));
  }

@ -117,6 +118,7 @@ class BatchedGemmTunable : public IBatchedGemmKernelExplorer<T> {
  }

  void Run() override {
+    WithMaxTuningDurationMs max_duration(params_.TuningContext(), 250);
    ORT_THROW_IF_ERROR(op_(&params_));
  }

@ -179,6 +181,7 @@ class StridedBatchedGemmTunable : public IKernelExplorer {
  }

  void Run() override {
+    WithMaxTuningDurationMs max_duration(params_.TuningContext(), 250);
    ORT_THROW_IF_ERROR(op_(&params_));
  }

--- a/onnxruntime/python/tools/kernel_explorer/kernels/rocm/softmax.cu
+++ b/onnxruntime/python/tools/kernel_explorer/kernels/rocm/softmax.cu
@ -111,6 +111,7 @@ class SoftmaxTunable : public IKernelExplorer {
  }

  void Run() override {
+    WithMaxTuningDurationMs max_duration(TuningContext(), 250);
    ORT_THROW_IF_ERROR(op_(&params_));
  }

--- a/onnxruntime/test/contrib_ops/multihead_attention_op_test.cc
+++ b/onnxruntime/test/contrib_ops/multihead_attention_op_test.cc
@ -478,7 +478,6 @@ TEST(MultiHeadAttentionTest, CrossAttention_Batch2_HeadSize32_RightSidePadding_M
 }

 TEST(MultiHeadAttentionTest, CrossAttention_Batch2_HeadSize32_RightSidePadding_Mask2D) {
-  ROCM_GTEST_SKIP("ROCm MHA expect failure due to ck bug");
  AttentionTestData data;
  GetCrossAttentionData_Batch2_HeadSize32_RightSidePadding(data, false);
  RunMultiHeadAttentionTests(data, true);
--- a/tools/ci_build/github/azure-pipelines/orttraining-pai-ci-pipeline.yml
+++ b/tools/ci_build/github/azure-pipelines/orttraining-pai-ci-pipeline.yml
@ -186,13 +186,14 @@ jobs:
          --user onnxruntimedev \
          --volume $(Build.SourcesDirectory):/onnxruntime_src \
          --volume $(Build.BinariesDirectory):/build \
+          -e OPENBLAS_NUM_THREADS=1 \
+          -e OPENMP_NUM_THREADS=1 \
+          -e MKL_NUM_THREADS=1 \
+          -e KERNEL_EXPLORER_BUILD_DIR=/build/$(BuildConfig) \
+          -e KERNEL_EXPLORER_BATCHED_GEMM_MAX_BATCH_SIZE=8 \
+          -e KERNEL_EXPLORER_TEST_USE_CUPY=1 \
          onnxruntimetrainingrocm-cibuild-rocm$(RocmVersion)-test \
-            /bin/bash -c "
-              set -ex; \
-              export KERNEL_EXPLORER_BUILD_DIR=/build/$(BuildConfig); \
-              export KERNEL_EXPLORER_BATCHED_GEMM_MAX_BATCH_SIZE=8; \
-              export KERNEL_EXPLORER_TEST_USE_CUPY=1; \
-              pytest /onnxruntime_src/onnxruntime/python/tools/kernel_explorer/ -n 8 --reruns 1 --durations=100"
+            pytest /onnxruntime_src/onnxruntime/python/tools/kernel_explorer/ -n 4 --reruns 1 --durations=100
      workingDirectory: $(Build.SourcesDirectory)
    displayName: 'Run kernel explorer tests'
    condition: succeededOrFailed()