Merge branch 'develop' into amd-develop

CMakeLists.txt

@@ -62,8 +62,14 @@ if (DTYPES)
     endif()
     message("DTYPES macro set to ${DTYPES}")
 else()
-    add_definitions(-DCK_ENABLE_INT8 -DCK_ENABLE_FP8 -DCK_ENABLE_BF8 -DCK_ENABLE_FP16 -DCK_ENABLE_FP32 -DCK_ENABLE_FP64 -DCK_ENABLE_BF16)
-    set(CK_ENABLE_ALL_DTYPES "ON")
+    add_definitions(-DCK_ENABLE_INT8 -DCK_ENABLE_FP16 -DCK_ENABLE_FP32 -DCK_ENABLE_FP64 -DCK_ENABLE_BF16 -DCK_ENABLE_FP8 -DCK_ENABLE_BF8)
+    set(CK_ENABLE_INT8 "ON")
+    set(CK_ENABLE_FP16 "ON")
+    set(CK_ENABLE_FP32 "ON")
+    set(CK_ENABLE_FP64 "ON")
+    set(CK_ENABLE_BF16 "ON")
+    set(CK_ENABLE_FP8 "ON")
+    set(CK_ENABLE_BF8 "ON")
 endif()
 
 #for f8/bf8_t type
@@ -182,12 +188,18 @@ endif()
 configure_file(include/ck/config.h.in ${CMAKE_CURRENT_BINARY_DIR}/include/ck/config.h)
 
 if(NOT WIN32 AND ${hip_VERSION_FLAT} GREATER 500723302)
-   message("Adding the fno-offload-uniform-block compiler flag")
-   add_compile_options(-fno-offload-uniform-block)
+  check_cxx_compiler_flag("-fno-offload-uniform-block" HAS_NO_OFFLOAD_UNIFORM_BLOCK)
+  if(HAS_NO_OFFLOAD_UNIFORM_BLOCK)
+    message("Adding the fno-offload-uniform-block compiler flag")
+    add_compile_options(-fno-offload-uniform-block)
+  endif()
 endif()
 if(NOT WIN32 AND ${hip_VERSION_FLAT} GREATER 600140090)
-   message("Adding the enable-post-misched=0 compiler flag")
-   add_compile_options("SHELL: -mllvm -enable-post-misched=0")
+  check_cxx_compiler_flag("-mllvm -enable-post-misched=0" HAS_ENABLE_POST_MISCHED)
+  if(HAS_ENABLE_POST_MISCHED)
+    message("Adding the enable-post-misched=0 compiler flag")
+    add_compile_options("SHELL: -mllvm -enable-post-misched=0")
+  endif()
 endif()
 set(check-coerce)
 check_cxx_compiler_flag(" -mllvm -amdgpu-coerce-illegal-types=1" check-coerce)
@@ -541,12 +553,7 @@ if(NOT DEFINED INSTANCES_ONLY)
         PACKAGE_NAME examples
    )
    add_subdirectory(example)
-   if(GPU_TARGETS MATCHES "gfx9" AND NOT INSTANCES_ONLY)
-      add_subdirectory(codegen)
-   endif()
-   if(BUILD_TESTING)
-      add_subdirectory(test)
-   endif()
+   add_subdirectory(test)
 
    rocm_package_setup_component(profiler
         LIBRARY_NAME composablekernel
@@ -563,6 +570,10 @@ if(NOT DEFINED INSTANCES_ONLY)
   endif()
 endif()
 
+if(NOT DEFINED PROFILER_ONLY AND (GPU_TARGETS MATCHES "gfx9" OR DEFINED INSTANCES_ONLY))
+  add_subdirectory(codegen)
+endif()
+
 #Create an interface target for the include only files and call it "composablekernels"
 include(CMakePackageConfigHelpers)
 

Jenkinsfile

@@ -426,8 +426,9 @@ def runCKProfiler(Map conf=[:]){
                             archiveArtifacts "perf_resnet50_N4.log"
                             archiveArtifacts "perf_batched_gemm.log"
                             archiveArtifacts "perf_grouped_gemm.log"
-                            archiveArtifacts "perf_conv_fwd.log"
-                            archiveArtifacts "perf_conv_bwd_data.log"
+                            archiveArtifacts "perf_grouped_conv_fwd.log"
+                            archiveArtifacts "perf_grouped_conv_bwd_data.log"
+                            archiveArtifacts "perf_grouped_conv_bwd_weight.log"
                             archiveArtifacts "perf_gemm_bilinear.log"
                             archiveArtifacts "perf_reduction.log"
                             archiveArtifacts "perf_splitK_gemm.log"
@@ -439,8 +440,9 @@ def runCKProfiler(Map conf=[:]){
                             stash name: "perf_resnet50_N4.log"
                             stash name: "perf_batched_gemm.log"
                             stash name: "perf_grouped_gemm.log"
-                            stash name: "perf_conv_fwd.log"
-                            stash name: "perf_conv_bwd_data.log"
+                            stash name: "perf_grouped_conv_fwd.log"
+                            stash name: "perf_grouped_conv_bwd_data.log"
+                            stash name: "perf_grouped_conv_bwd_weight.log"
                             stash name: "perf_gemm_bilinear.log"
                             stash name: "perf_reduction.log"
                             stash name: "perf_splitK_gemm.log"
@@ -648,8 +650,9 @@ def process_results(Map conf=[:]){
                         unstash "perf_resnet50_N4.log"
                         unstash "perf_batched_gemm.log"
                         unstash "perf_grouped_gemm.log"
-                        unstash "perf_conv_fwd.log"
-                        unstash "perf_conv_bwd_data.log"
+                        unstash "perf_grouped_conv_fwd.log"
+                        unstash "perf_grouped_conv_bwd_data.log"
+                        unstash "perf_grouped_conv_bwd_weight.log"
                         unstash "perf_gemm_bilinear.log"
                         unstash "perf_reduction.log"
                         unstash "perf_splitK_gemm.log"
@@ -746,6 +749,10 @@ pipeline {
             name: "RUN_PERFORMANCE_TESTS",
             defaultValue: true,
             description: "Run the performance tests (default: ON)")
+        booleanParam(
+            name: "RUN_GROUPED_CONV_LARGE_CASES_TESTS",
+            defaultValue: false,
+            description: "Run the grouped conv large cases tests (default: OFF)")
         booleanParam(
             name: "RUN_CK_TILE_TESTS",
             defaultValue: false,
@@ -837,6 +844,30 @@ pipeline {
                 }
             }
         }
+        stage("Run Grouped Conv Large Case Tests")
+        {
+            parallel
+            {
+                stage("Run Grouped Conv Large Case Tests on gfx90a")
+                {
+                    when {
+                        beforeAgent true
+                        expression { params.RUN_GROUPED_CONV_LARGE_CASES_TESTS.toBoolean() }
+                    }
+                    agent{ label rocmnode("gfx90a")}
+                    environment{
+                        setup_args = "NO_CK_BUILD"
+                        execute_args = """ ../script/cmake-ck-dev.sh  ../ gfx90a && \
+                                           make -j64 test_grouped_convnd_fwd_large_cases_xdl && \
+                                           ./bin/test_grouped_convnd_fwd_large_cases_xdl"""
+                   }
+                    steps{
+                        buildHipClangJobAndReboot(setup_args:setup_args, no_reboot:true, build_type: 'Release', execute_cmd: execute_args)
+                        cleanWs()
+                    }
+                }
+            }
+        }
         stage("Run CK_TILE Tests")
         {
             parallel

client_example/07_grouped_convnd_fwd/CMakeLists.txt

@@ -5,17 +5,17 @@ if(GPU_TARGETS MATCHES "gfx9")
     add_executable(client_grouped_conv1d_fwd grouped_conv1d_fwd.cpp)
     target_link_libraries(client_grouped_conv1d_fwd PRIVATE composable_kernel::device_conv_operations)
 
-    if((DTYPES MATCHES "fp8") OR NOT DEFINED DTYPES)
+    if((DTYPES MATCHES "fp8") OR (NOT DEFINED DTYPES AND GPU_TARGETS MATCHES "gfx94"))
         add_executable(client_grouped_conv3d_fwd_fp8 grouped_conv3d_fwd_fp8.cpp)
         target_link_libraries(client_grouped_conv3d_fwd_fp8 PRIVATE composable_kernel::device_conv_operations)
     endif()
 
-    if((DTYPES MATCHES "bf8") OR NOT DEFINED DTYPES)
+    if((DTYPES MATCHES "bf8") OR (NOT DEFINED DTYPES AND GPU_TARGETS MATCHES "gfx94"))
         add_executable(client_grouped_conv3d_fwd_bf8 grouped_conv3d_fwd_bf8.cpp)
         target_link_libraries(client_grouped_conv3d_fwd_bf8 PRIVATE composable_kernel::device_conv_operations)
     endif()
 
-    if((DTYPES MATCHES "fp8" AND DTYPES MATCHES "bf8") OR NOT DEFINED DTYPES)
+    if((DTYPES MATCHES "fp8" AND DTYPES MATCHES "bf8") OR (NOT DEFINED DTYPES AND GPU_TARGETS MATCHES "gfx94"))
         add_executable(client_grouped_conv3d_fwd_fp8_bf8 grouped_conv3d_fwd_fp8_bf8.cpp)
         target_link_libraries(client_grouped_conv3d_fwd_fp8_bf8 PRIVATE composable_kernel::device_conv_operations)
 

client_example/10_grouped_convnd_bwd_data/CMakeLists.txt

@@ -4,5 +4,7 @@ target_link_libraries(client_grouped_conv2d_bwd_data PRIVATE composable_kernel::
 add_executable(client_grouped_conv3d_bwd_data grouped_conv3d_bwd_data.cpp)
 target_link_libraries(client_grouped_conv3d_bwd_data PRIVATE composable_kernel::device_conv_operations)
 
-add_executable(client_grouped_conv3d_bwd_data_input_fp16_comp_bf8f8 grouped_conv3d_bwd_data_input_fp16_comp_bf8f8.cpp)
-target_link_libraries(client_grouped_conv3d_bwd_data_input_fp16_comp_bf8f8 PRIVATE composable_kernel::device_conv_operations)
+if((DTYPES MATCHES "fp8" AND DTYPES MATCHES "bf8") OR (NOT DEFINED DTYPES AND GPU_TARGETS MATCHES "gfx94"))
+    add_executable(client_grouped_conv3d_bwd_data_input_fp16_comp_bf8f8 grouped_conv3d_bwd_data_input_fp16_comp_bf8f8.cpp)
+    target_link_libraries(client_grouped_conv3d_bwd_data_input_fp16_comp_bf8f8 PRIVATE composable_kernel::device_conv_operations)
+endif()

client_example/11_grouped_conv_bwd_weight/CMakeLists.txt

@@ -2,10 +2,13 @@ add_executable(client_grouped_conv1d_bwd_weight_fp16 grouped_conv1d_bwd_weight_f
 add_executable(client_grouped_conv2d_bwd_weight_fp16 grouped_conv2d_bwd_weight_fp16.cpp)
 add_executable(client_grouped_conv3d_bwd_weight_fp16 grouped_conv3d_bwd_weight_fp16.cpp)
 add_executable(client_grouped_conv3d_bwd_weight_fp32 grouped_conv3d_bwd_weight_fp32.cpp)
-add_executable(client_grouped_conv3d_bwd_weight_fp16_comp_bf8_fp8 grouped_conv3d_bwd_weight_fp16_comp_bf8_fp8.cpp)
 
 target_link_libraries(client_grouped_conv1d_bwd_weight_fp16 PRIVATE composable_kernel::device_conv_operations)
 target_link_libraries(client_grouped_conv2d_bwd_weight_fp16 PRIVATE composable_kernel::device_conv_operations)
 target_link_libraries(client_grouped_conv3d_bwd_weight_fp16 PRIVATE composable_kernel::device_conv_operations)
 target_link_libraries(client_grouped_conv3d_bwd_weight_fp32 PRIVATE composable_kernel::device_conv_operations)
-target_link_libraries(client_grouped_conv3d_bwd_weight_fp16_comp_bf8_fp8 PRIVATE composable_kernel::device_conv_operations)
+
+if((DTYPES MATCHES "fp8" AND DTYPES MATCHES "bf8") OR (NOT DEFINED DTYPES AND GPU_TARGETS MATCHES "gfx94"))
+    add_executable(client_grouped_conv3d_bwd_weight_fp16_comp_bf8_fp8 grouped_conv3d_bwd_weight_fp16_comp_bf8_fp8.cpp)
+    target_link_libraries(client_grouped_conv3d_bwd_weight_fp16_comp_bf8_fp8 PRIVATE composable_kernel::device_conv_operations)
+endif()

client_example/16_convnd_fwd/CMakeLists.txt

@@ -4,7 +4,7 @@ if((DTYPES MATCHES "fp16") OR NOT DEFINED DTYPES)
 
 endif()
 
-if((DTYPES MATCHES "fp8") OR NOT DEFINED DTYPES)
+if((DTYPES MATCHES "fp8") OR (NOT DEFINED DTYPES AND GPU_TARGETS MATCHES "gfx94"))
     add_executable(client_conv3d_fwd_fp16_comp_fp8 conv3d_fwd_fp16_comp_fp8.cpp)
     target_link_libraries(client_conv3d_fwd_fp16_comp_fp8 PRIVATE composable_kernel::device_conv_operations)
 endif()

client_example/20_splitk_gemm/CMakeLists.txt

@@ -1,4 +1,4 @@
-if(GPU_TARGETS MATCHES "gfx9" AND ((DTYPES MATCHES "fp8" AND DTYPES MATCHES "fp16") OR NOT DEFINED DTYPES))
+if((DTYPES MATCHES "fp8" AND DTYPES MATCHES "fp16") OR (NOT DEFINED DTYPES AND GPU_TARGETS MATCHES "gfx94"))
   add_executable(client_splitK_gemm splitK_gemm_fp16_f8.cpp)
   target_link_libraries(client_splitK_gemm PRIVATE composable_kernel::device_gemm_operations)
 endif()

client_example/24_grouped_conv_activation/CMakeLists.txt

@@ -1,6 +1,6 @@
 if(GPU_TARGETS MATCHES "gfx9")
 # Fwd scaleadd scaleadd relu
-add_executable(client_grouped_convnd_fwd_scaleadd_scaleadd_relu_fp32 
+add_executable(client_grouped_convnd_fwd_scaleadd_scaleadd_relu_fp32
                grouped_convnd_fwd_scaleadd_scaleadd_relu/grouped_conv_fwd_scaleadd_scaleadd_relu_fp32.cpp)
 target_link_libraries(client_grouped_convnd_fwd_scaleadd_scaleadd_relu_fp32 PRIVATE composable_kernel::device_conv_operations)
 
@@ -36,7 +36,7 @@ add_executable(client_grouped_convnd_fwd_bilinear_residual_fp16
                grouped_convnd_fwd_bilinear/grouped_conv_fwd_bilinear_residual_fp16.cpp)
 target_link_libraries(client_grouped_convnd_fwd_bilinear_residual_fp16 PRIVATE composable_kernel::device_conv_operations)
 # Fwd convinvscale
-add_executable(client_conv3d_fwd_convinvscale_fp8 
+add_executable(client_conv3d_fwd_convinvscale_fp8
                grouped_convnd_fwd_convinvscale/conv3d_fwd_convinvscale_fp8.cpp)
 target_link_libraries(client_conv3d_fwd_convinvscale_fp8 PRIVATE composable_kernel::device_conv_operations)
 # Fwd convscale + Bias
@@ -47,6 +47,22 @@ target_link_libraries(client_conv3d_fwd_convscale_add_fp8 PRIVATE composable_ker
 add_executable(client_conv3d_fwd_convscale_relu_fp8
                grouped_convnd_fwd_convscale_relu/conv3d_fwd_convscale_relu_fp8.cpp)
 target_link_libraries(client_conv3d_fwd_convscale_relu_fp8 PRIVATE composable_kernel::device_conv_operations)
+# Fwd convscale + ReLU + AMAX
+add_executable(client_conv3d_fwd_convscale_relu_amax_fp8
+               grouped_convnd_fwd_convscale_reduce/conv3d_fwd_convscale_relu_amax_fp8.cpp)
+target_link_libraries(client_conv3d_fwd_convscale_relu_amax_fp8
+                      PRIVATE composable_kernel::device_conv_operations
+                              composable_kernel::device_other_operations
+                              composable_kernel::device_reduction_operations
+                              utility)
+# Fwd convscale + AMAX
+add_executable(client_conv3d_fwd_convscale_amax_fp8
+               grouped_convnd_fwd_convscale_reduce/conv3d_fwd_convscale_amax_fp8.cpp)
+target_link_libraries(client_conv3d_fwd_convscale_amax_fp8
+                      PRIVATE composable_kernel::device_conv_operations
+                              composable_kernel::device_other_operations
+                              composable_kernel::device_reduction_operations
+                              utility)
 # Fwd convscale
 add_executable(client_conv3d_fwd_convscale_fp8
                grouped_convnd_fwd_convscale/conv3d_fwd_convscale_fp8.cpp)
@@ -56,11 +72,11 @@ add_executable(client_conv3d_fwd_convscale_bf8
                grouped_convnd_fwd_convscale/conv3d_fwd_convscale_bf8.cpp)
 target_link_libraries(client_conv3d_fwd_convscale_bf8 PRIVATE composable_kernel::device_conv_operations)
 
-add_executable(client_conv3d_fwd_convscale_fp8_bf8 
+add_executable(client_conv3d_fwd_convscale_fp8_bf8
                grouped_convnd_fwd_convscale/conv3d_fwd_convscale_fp8_bf8.cpp)
 target_link_libraries(client_conv3d_fwd_convscale_fp8_bf8 PRIVATE composable_kernel::device_conv_operations)
 
-add_executable(client_conv3d_fwd_convscale_bf8_fp8 
+add_executable(client_conv3d_fwd_convscale_bf8_fp8
                grouped_convnd_fwd_convscale/conv3d_fwd_convscale_bf8_fp8.cpp)
 target_link_libraries(client_conv3d_fwd_convscale_bf8_fp8 PRIVATE composable_kernel::device_conv_operations)
 # Bwd data bilinear

client_example/24_grouped_conv_activation/grouped_convnd_fwd_convscale_reduce/common.hpp

@@ -0,0 +1,834 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#include <algorithm>
+#include <cstdlib>
+#include <iomanip>
+#include <iostream>
+#include <iterator>
+#include <numeric>
+#include <string>
+#include <vector>
+
+#include "ck/ck.hpp"
+#include "ck/library/utility/algorithm.hpp"
+#include "ck/tensor_operation/gpu/device/device_elementwise.hpp"
+#include "ck/tensor_operation/gpu/device/device_reduce.hpp"
+#include "ck/tensor_operation/gpu/device/reduction_operator_mapping.hpp"
+#include "ck/utility/tuple.hpp"
+#include "ck/utility/type.hpp"
+#include "ck/library/tensor_operation_instance/gpu/grouped_convolution_forward_convscale_relu.hpp"
+#include "ck/library/tensor_operation_instance/gpu/grouped_convolution_forward_convscale.hpp"
+#include "ck/utility/reduction_enums.hpp"
+#include "ck/library/tensor_operation_instance/gpu/permute_scale.hpp"
+#include "ck/library/tensor_operation_instance/gpu/reduce/reduce.hpp"
+#include "ck/library/utility/host_tensor.hpp"
+
+using PassThrough   = ck::tensor_operation::element_wise::PassThrough;
+using ConvScaleRelu = ck::tensor_operation::element_wise::ScaleScaleRelu;
+using ConvScale     = ck::tensor_operation::element_wise::ScaleScalePass;
+
+struct SimpleDeviceMem
+{
+    SimpleDeviceMem() = delete;
+
+    SimpleDeviceMem(std::size_t mem_size) : p_mem_{}
+    {
+        (void)hipMalloc(static_cast<void**>(&p_mem_), mem_size);
+    }
+
+    void* GetDeviceBuffer() { return p_mem_; }
+
+    ~SimpleDeviceMem() { (void)hipFree(p_mem_); }
+
+    void* p_mem_;
+};
+
+template <ck::index_t NumDimSpatial, ck::index_t NumNonSpatialDim = 3>
+std::size_t
+GetFlops(const std::array<ck::index_t, NumDimSpatial + NumNonSpatialDim>& output_lengths,
+         const std::array<ck::index_t, NumDimSpatial + NumNonSpatialDim>& weights_lengths,
+         const std::size_t& ds_size)
+{
+    // 2 * G * N * K * C * <output spatial lengths product> * <filter spatial lengths product> +
+    // + ds_size * <output tensor size> =>
+    // => <output tensor size> * ( 2 * C * <filter spatial lengths product> + ds_size) =>
+    // => G * N * K * <output spatial lengths product> * (2 * C * <filter spatial lengths product> +
+    // ds_size)
+    ck::index_t G = weights_lengths[0];
+    ck::index_t N = output_lengths[1];
+    ck::index_t K = weights_lengths[1];
+    ck::index_t C = weights_lengths[2];
+
+    return G * N * K *
+           std::accumulate(std::next(std::begin(output_lengths), NumNonSpatialDim),
+                           std::end(output_lengths),
+                           static_cast<std::size_t>(1),
+                           std::multiplies<>()) *
+           (ds_size + static_cast<std::size_t>(2) * C *
+                          std::accumulate(std::next(std::begin(weights_lengths), NumNonSpatialDim),
+                                          std::end(weights_lengths),
+                                          static_cast<std::size_t>(1),
+                                          std::multiplies<>()));
+}
+
+template <ck::index_t NumDimSpatial, ck::index_t NumNonSpatialDim = 3>
+std::size_t GetTensorSize(const std::array<ck::index_t, NumDimSpatial + NumNonSpatialDim>& lengths)
+{
+
+    return std::accumulate(std::begin(lengths),
+                           std::end(lengths),
+                           static_cast<std::size_t>(1),
+                           std::multiplies<std::size_t>());
+}
+
+template <typename InDataType, ck::index_t NumDimSpatial, ck::index_t NumNonSpatialDim = 3>
+std::size_t
+GetInputByte(const std::array<ck::index_t, NumDimSpatial + NumNonSpatialDim>& input_lengths)
+{
+    // sizeof(InDataType) * (G * N * C * <input spatial lengths product>) +
+    return sizeof(InDataType) * GetTensorSize<NumDimSpatial>(input_lengths);
+}
+
+template <typename WeiDataType, ck::index_t NumDimSpatial, ck::index_t NumNonSpatialDim = 3>
+std::size_t
+GetWeightByte(const std::array<ck::index_t, NumDimSpatial + NumNonSpatialDim>& weights_lengths)
+{
+    // sizeof(WeiDataType) * (G * K * C * <filter spatial lengths product>) +
+    return sizeof(WeiDataType) * GetTensorSize<NumDimSpatial>(weights_lengths);
+}
+
+template <typename OutDataType, ck::index_t NumDimSpatial, ck::index_t NumNonSpatialDim = 3>
+std::size_t
+GetOutputByte(const std::array<ck::index_t, NumDimSpatial + NumNonSpatialDim>& output_lengths)
+{
+    // sizeof(OutDataType) * (G * N * K * <output spatial lengths product>);
+    return sizeof(OutDataType) * GetTensorSize<NumDimSpatial>(output_lengths);
+}
+
+template <typename InDataType,
+          typename WeiDataType,
+          typename OutDataType,
+          typename ConvElementOp,
+          typename InLayout,
+          typename WeiLayout,
+          typename OutLayout,
+          ck::index_t NumDimSpatial,
+          ck::index_t NumNonSpatialDim = 3,
+          typename AComputeType        = InDataType,
+          typename BComputeType        = AComputeType>
+bool ConvolutionScale(SimpleDeviceMem& in,
+                      SimpleDeviceMem& wei,
+                      SimpleDeviceMem& out,
+                      ConvElementOp elementwise_op,
+                      const std::array<ck::index_t, NumDimSpatial + NumNonSpatialDim>& in_lengths,
+                      const std::array<ck::index_t, NumDimSpatial + NumNonSpatialDim>& in_strides,
+                      const std::array<ck::index_t, NumDimSpatial + NumNonSpatialDim>& wei_lengths,
+                      const std::array<ck::index_t, NumDimSpatial + NumNonSpatialDim>& wei_strides,
+                      const std::array<ck::index_t, NumDimSpatial + NumNonSpatialDim>& out_lengths,
+                      const std::array<ck::index_t, NumDimSpatial + NumNonSpatialDim>& out_strides);
+
+template <typename InDataType,
+          typename OutDataType,
+          ck::index_t NumDimSpatial,
+          ck::index_t NumNonSpatialDim = 3>
+bool TensorScaleConvert(SimpleDeviceMem& in,
+                        SimpleDeviceMem& out,
+                        float scale_out,
+                        const std::array<ck::index_t, NumDimSpatial + NumNonSpatialDim>& lengths,
+                        const std::array<ck::index_t, NumDimSpatial + NumNonSpatialDim>& strides);
+
+template <typename InDataType,
+          typename OutDataType,
+          ck::ReduceTensorOp ReduceOpId,
+          ck::index_t NumDimSpatial,
+          ck::index_t NumNonSpatialDim = 3>
+bool TensorFullReduction(SimpleDeviceMem& tensor,
+                         SimpleDeviceMem& out_amax,
+                         const std::array<ck::index_t, NumDimSpatial + NumNonSpatialDim>& lengths,
+                         const std::array<ck::index_t, NumDimSpatial + NumNonSpatialDim>& strides);
+
+template <ck::index_t NumDimSpatial,
+          typename InDataType,
+          typename WeiDataType,
+          typename ConvOutDataType,
+          typename OutDataType,
+          typename ConvElementOp,
+          ck::ReduceTensorOp ReduceOp,
+          typename InLayout,
+          typename WeiLayout,
+          typename OutLayout,
+          ck::index_t NumNonSpatialDim = 3,
+          typename AComputeType        = InDataType,
+          typename BComputeType        = AComputeType>
+bool run_grouped_conv_fwd_convscale_reduce(
+    std::array<ck::index_t, NumDimSpatial + NumNonSpatialDim> in_lengths,
+    std::array<ck::index_t, NumDimSpatial + NumNonSpatialDim> wei_lengths,
+    std::array<ck::index_t, NumDimSpatial + NumNonSpatialDim> out_lengths)
+{
+
+    namespace ctc = ck::tensor_layout::convolution;
+    static_assert(NumDimSpatial == 3 && ck::is_same_v<InLayout, ctc::NDHWGC> &&
+                      ck::is_same_v<WeiLayout, ctc::GKZYXC> &&
+                      ck::is_same_v<OutLayout, ctc::NDHWGK>,
+                  "Unsupported configuration");
+
+    const ck::index_t G  = in_lengths[4];
+    const ck::index_t N  = in_lengths[0];
+    const ck::index_t K  = wei_lengths[1];
+    const ck::index_t C  = in_lengths[5];
+    const ck::index_t Z  = wei_lengths[2];
+    const ck::index_t Y  = wei_lengths[3];
+    const ck::index_t X  = wei_lengths[4];
+    const ck::index_t Di = in_lengths[1];
+    const ck::index_t Hi = in_lengths[2];
+    const ck::index_t Wi = in_lengths[3];
+    const ck::index_t Do = out_lengths[1];
+    const ck::index_t Ho = out_lengths[2];
+    const ck::index_t Wo = out_lengths[3];
+
+    const std::size_t in_mem_size       = sizeof(InDataType) * N * Di * Hi * Wi * G * C;
+    const std::size_t wei_mem_size      = sizeof(WeiDataType) * G * K * Z * Y * X * C;
+    const std::size_t conv_out_mem_size = sizeof(ConvOutDataType) * N * Do * Ho * Wo * G * K;
+    const std::size_t out_mem_size      = sizeof(OutDataType) * N * Do * Ho * Wo * G * K;
+
+    SimpleDeviceMem in(in_mem_size);
+    SimpleDeviceMem wei(wei_mem_size);
+    SimpleDeviceMem conv_out(conv_out_mem_size);
+    SimpleDeviceMem out(out_mem_size);
+
+    float scale_in  = float(std::rand()) / float(RAND_MAX);
+    float scale_wei = float(std::rand()) / float(RAND_MAX);
+    float scale_out = float(std::rand()) / float(RAND_MAX);
+
+    // We have NDHWGC/GKZYXC/NDHWGK (x, weight, y) in memory space.
+    // However, CK's API only accepts lengths and strides with order of GNCDHW/GKCZYX/GNKDHW.
+    // Hence, we need to adjust the order of strides.
+    const std::array<ck::index_t, NumDimSpatial + 3> input_lengths{G, N, C, Di, Hi, Wi};
+    const std::array<ck::index_t, NumDimSpatial + 3> input_strides{
+        C, Di * Hi * Wi * G * C, 1, Hi * Wi * G * C, Wi * G * C, G * C};
+    const std::array<ck::index_t, NumDimSpatial + 3> weights_lengths{G, K, C, Z, Y, X};
+    const std::array<ck::index_t, NumDimSpatial + 3> weights_strides{
+        K * Z * Y * X * C, Z * Y * X * C, 1, Y * X * C, X * C, C};
+    const std::array<ck::index_t, NumDimSpatial + 3> output_lengths{G, N, K, Do, Ho, Wo};
+    const std::array<ck::index_t, NumDimSpatial + 3> output_strides{
+        K, Do * Ho * Wo * G * K, 1, Ho * Wo * G * K, Wo * G * K, G * K};
+
+    /*
+     * FP8 Convolution with Scaling
+     */
+    std::cout << "\n\nConvolution with scale Benchmarking:" << std::endl;
+    auto elementwise_op = ConvElementOp{ck::tensor_operation::element_wise::Scale{scale_in},
+                                        ck::tensor_operation::element_wise::Scale{scale_wei},
+                                        {}};
+    auto conv_ok        = ConvolutionScale<InDataType,
+                                    WeiDataType,
+                                    ConvOutDataType,
+                                    ConvElementOp,
+                                    InLayout,
+                                    WeiLayout,
+                                    OutLayout,
+                                    NumDimSpatial>(in,
+                                                   wei,
+                                                   conv_out,
+                                                   elementwise_op,
+                                                   input_lengths,
+                                                   input_strides,
+                                                   weights_lengths,
+                                                   weights_strides,
+                                                   output_lengths,
+                                                   output_strides);
+
+    if(!conv_ok)
+        return false;
+
+    /*
+     *  Scale with output weight and convert to FP8
+     */
+    std::cout << "\n\nElement-wise scale + convert Benchmarking:" << std::endl;
+    auto elem_wise_ok = TensorScaleConvert<ConvOutDataType, OutDataType, NumDimSpatial>(
+        conv_out, out, scale_out, output_lengths, output_strides);
+
+    if(!elem_wise_ok)
+        return false;
+
+    /*
+     *  Compute AMAX
+     */
+    std::cout << "\n\nAMAX Benchmarking:" << std::endl;
+    SimpleDeviceMem amax_device(sizeof(ConvOutDataType));
+    auto reduction_ok =
+        TensorFullReduction<ConvOutDataType,
+                            ConvOutDataType,
+                            ck::ReduceTensorOp::AMAX,
+                            NumDimSpatial>(conv_out, amax_device, output_lengths, output_strides);
+
+    if(!reduction_ok)
+        return false;
+
+    return true;
+}
+
+template <typename InDataType,
+          typename WeiDataType,
+          typename OutDataType,
+          typename ConvElementOp,
+          typename InLayout,
+          typename WeiLayout,
+          typename OutLayout,
+          ck::index_t NumDimSpatial,
+          ck::index_t NumNonSpatialDim,
+          typename AComputeType,
+          typename BComputeType>
+bool ConvolutionScale(SimpleDeviceMem& in,
+                      SimpleDeviceMem& wei,
+                      SimpleDeviceMem& out,
+                      ConvElementOp elementwise_op,
+                      const std::array<ck::index_t, NumDimSpatial + NumNonSpatialDim>& in_lengths,
+                      const std::array<ck::index_t, NumDimSpatial + NumNonSpatialDim>& in_strides,
+                      const std::array<ck::index_t, NumDimSpatial + NumNonSpatialDim>& wei_lengths,
+                      const std::array<ck::index_t, NumDimSpatial + NumNonSpatialDim>& wei_strides,
+                      const std::array<ck::index_t, NumDimSpatial + NumNonSpatialDim>& out_lengths,
+                      const std::array<ck::index_t, NumDimSpatial + NumNonSpatialDim>& out_strides)
+{
+
+    const std::array<ck::index_t, NumDimSpatial> conv_filter_strides{1, 1, 1};
+    const std::array<ck::index_t, NumDimSpatial> conv_filter_dilations{1, 1, 1};
+    const std::array<ck::index_t, NumDimSpatial> input_left_pads{1, 1, 1};
+    const std::array<ck::index_t, NumDimSpatial> input_right_pads{1, 1, 1};
+
+    const auto in_mem_size  = GetInputByte<InDataType, NumDimSpatial>(in_lengths);
+    const auto wei_mem_size = GetWeightByte<WeiDataType, NumDimSpatial>(wei_lengths);
+    const auto out_mem_size = GetOutputByte<OutDataType, NumDimSpatial>(out_lengths);
+
+    std::size_t ds_size = 2; // 2 element-wise scale multipliers
+    if constexpr(ck::is_same_v<ConvElementOp, ConvScaleRelu>)
+    {
+        ds_size += 1; // +1 element-wise relu
+    }
+    std::size_t flop = GetFlops<NumDimSpatial>(out_lengths, wei_lengths, ds_size);
+    std::size_t num_bytes =
+        in_mem_size + wei_mem_size + sizeof(float) + sizeof(float) + out_mem_size;
+
+    using ConvDeviceOp =
+        ck::tensor_operation::device::DeviceGroupedConvFwdMultipleABD<NumDimSpatial,
+                                                                      InLayout,
+                                                                      WeiLayout,
+                                                                      ck::Tuple<>,
+                                                                      OutLayout,
+                                                                      InDataType,
+                                                                      WeiDataType,
+                                                                      ck::Tuple<>,
+                                                                      OutDataType,
+                                                                      PassThrough,
+                                                                      PassThrough,
+                                                                      ConvElementOp,
+                                                                      AComputeType,
+                                                                      BComputeType>;
+    // get device op instances
+    const auto conv_ptrs = ck::tensor_operation::device::instance::DeviceOperationInstanceFactory<
+        ConvDeviceOp>::GetInstances();
+
+    std::cout << "found " << conv_ptrs.size() << " instances" << std::endl;
+
+    std::string conv_best_op_name;
+    int conv_best_op_id        = -1;
+    float conv_best_avg_time   = std::numeric_limits<float>::max();
+    float conv_best_gb_per_sec = 0;
+    float conv_best_tflops     = 0;
+
+    // profile device operation instances
+    std::cout << "Run all convolution instances and do timing" << std::endl;
+
+    for(int i = 0; i < conv_ptrs.size(); ++i)
+    {
+        auto& op_ptr      = conv_ptrs[i];
+        auto argument_ptr = op_ptr->MakeArgumentPointer(
+            in.GetDeviceBuffer(),
+            wei.GetDeviceBuffer(),
+            std::array<const void*, 0>{},
+            out.GetDeviceBuffer(),
+            in_lengths,
+            in_strides,
+            wei_lengths,
+            wei_strides,
+            std::array<std::array<ck::index_t, NumDimSpatial + NumNonSpatialDim>, 0>{},
+            std::array<std::array<ck::index_t, NumDimSpatial + NumNonSpatialDim>, 0>{},
+            out_lengths,
+            out_strides,
+            conv_filter_strides,
+            conv_filter_dilations,
+            input_left_pads,
+            input_right_pads,
+            PassThrough{},
+            PassThrough{},
+            elementwise_op);
+
+        auto invoker_ptr    = op_ptr->MakeInvokerPointer();
+        std::string op_name = op_ptr->GetTypeString();
+
+        if(op_ptr->IsSupportedArgument(argument_ptr.get()))
+        {
+            float avg_time = invoker_ptr->Run(argument_ptr.get(), StreamConfig{nullptr, true});
+
+            float tflops     = static_cast<float>(flop) / 1.E9 / avg_time;
+            float gb_per_sec = num_bytes / 1.E6 / avg_time;
+
+            std::cout << "Perf: " << std::setw(10) << avg_time << " ms, " << tflops << " TFlops, "
+                      << gb_per_sec << " GB/s, " << op_name << std::endl;
+
+            if(tflops > conv_best_tflops)
+            {
+                conv_best_op_id      = i;
+                conv_best_op_name    = op_name;
+                conv_best_avg_time   = avg_time;
+                conv_best_gb_per_sec = gb_per_sec;
+                conv_best_tflops     = tflops;
+            }
+        }
+        else
+        {
+            std::cerr << op_name << " does not support this problem" << std::endl;
+        }
+    }
+
+    if(conv_best_op_id < 0)
+    {
+        std::cerr << "no suitable instance" << std::endl;
+        return false;
+    }
+
+    std::cout << "Best Perf: " << std::setw(10) << conv_best_avg_time << " ms, " << conv_best_tflops
+              << " TFlops, " << conv_best_gb_per_sec << " GB/s, " << conv_best_op_name << std::endl;
+
+    // run the best instance
+    {
+        auto& op_ptr = conv_ptrs[conv_best_op_id];
+        std::cout << "Run the best instance without timing: " << op_ptr->GetTypeString()
+                  << std::endl;
+        auto argument_ptr = op_ptr->MakeArgumentPointer(
+            in.GetDeviceBuffer(),
+            wei.GetDeviceBuffer(),
+            std::array<const void*, 0>{},
+            out.GetDeviceBuffer(),
+            in_lengths,
+            in_strides,
+            wei_lengths,
+            wei_strides,
+            std::array<std::array<ck::index_t, NumDimSpatial + NumNonSpatialDim>, 0>{},
+            std::array<std::array<ck::index_t, NumDimSpatial + NumNonSpatialDim>, 0>{},
+            out_lengths,
+            out_strides,
+            conv_filter_strides,
+            conv_filter_dilations,
+            input_left_pads,
+            input_right_pads,
+            PassThrough{},
+            PassThrough{},
+            elementwise_op);
+
+        auto invoker_ptr = op_ptr->MakeInvokerPointer();
+
+        if(op_ptr->IsSupportedArgument(argument_ptr.get()))
+        {
+            invoker_ptr->Run(argument_ptr.get(), StreamConfig{nullptr, false});
+        }
+
+        std::cout << "Done" << std::endl;
+    }
+
+    return true;
+}
+
+template <typename InDataType,
+          typename OutDataType,
+          ck::index_t NumDimSpatial,
+          ck::index_t NumNonSpatialDim>
+bool TensorScaleConvert(SimpleDeviceMem& in,
+                        SimpleDeviceMem& out,
+                        float scale_out,
+                        const std::array<ck::index_t, NumDimSpatial + NumNonSpatialDim>& lengths,
+                        const std::array<ck::index_t, NumDimSpatial + NumNonSpatialDim>& strides)
+{
+
+    const auto tensor_size = GetTensorSize<NumDimSpatial>(lengths);
+
+    const std::size_t in_mem_size  = sizeof(InDataType) * tensor_size;
+    const std::size_t out_mem_size = sizeof(OutDataType) * tensor_size;
+
+    std::size_t flop = 2 * tensor_size; // element-wise scale + convert
+
+    std::size_t bytes =
+        in_mem_size + sizeof(float) + out_mem_size; // read from in, scale, write to out
+
+    using DeviceScaleConvert =
+        ck::tensor_operation::device::DeviceElementwise<ck::Tuple<InDataType>,
+                                                        ck::Tuple<OutDataType>,
+                                                        ck::tensor_operation::element_wise::Scale,
+                                                        NumDimSpatial + NumNonSpatialDim>;
+
+    // get device op instances
+    const auto op_ptrs = ck::tensor_operation::device::instance::DeviceOperationInstanceFactory<
+        DeviceScaleConvert>::GetInstances();
+
+    std::cout << "found " << op_ptrs.size() << " instances" << std::endl;
+
+    std::string best_op_name;
+    int best_op_id        = -1;
+    float best_avg_time   = std::numeric_limits<float>::max();
+    float best_gb_per_sec = 0;
+    float best_tflops     = 0;
+
+    // profile device operation instances
+    std::cout << "Run all DeviceScaleConvert instances and do timing" << std::endl;
+
+    auto scale_convert = ck::tensor_operation::element_wise::Scale{scale_out};
+
+    for(int i = 0; i < op_ptrs.size(); ++i)
+    {
+        auto& op_ptr      = op_ptrs[i];
+        auto argument_ptr = op_ptr->MakeArgumentPointer(lengths,
+                                                        {strides},
+                                                        {strides},
+                                                        {in.GetDeviceBuffer()},
+                                                        {out.GetDeviceBuffer()},
+                                                        scale_convert);
+
+        auto invoker_ptr    = op_ptr->MakeInvokerPointer();
+        std::string op_name = op_ptr->GetTypeString();
+
+        if(op_ptr->IsSupportedArgument(argument_ptr.get()))
+        {
+            float avg_time = invoker_ptr->Run(argument_ptr.get(), StreamConfig{nullptr, true});
+
+            float tflops     = static_cast<float>(flop) / 1.E9 / avg_time;
+            float gb_per_sec = bytes / 1.E6 / avg_time;
+
+            std::cout << "Perf: " << std::setw(10) << avg_time << " ms, " << tflops << " TFlops, "
+                      << gb_per_sec << " GB/s, " << op_name << std::endl;
+
+            if(tflops > best_tflops)
+            {
+                best_op_id      = i;
+                best_op_name    = op_name;
+                best_avg_time   = avg_time;
+                best_gb_per_sec = gb_per_sec;
+                best_tflops     = tflops;
+            }
+        }
+        else
+        {
+            std::cerr << op_name << " does not support this problem" << std::endl;
+        }
+    }
+
+    if(best_op_id < 0)
+    {
+        std::cerr << "no suitable instance found." << std::endl;
+        return false;
+    }
+    else
+    {
+        std::cout << "Best Perf: " << std::setw(10) << best_avg_time << " ms, " << best_tflops
+                  << " TFlops, " << best_gb_per_sec << " GB/s, " << best_op_name << std::endl;
+
+        // run the best intance
+        auto& op_ptr = op_ptrs[best_op_id];
+        std::cout << "Run the best instance without timing: " << op_ptr->GetTypeString()
+                  << std::endl;
+        auto argument_ptr = op_ptr->MakeArgumentPointer(lengths,
+                                                        {strides},
+                                                        {strides},
+                                                        {in.GetDeviceBuffer()},
+                                                        {out.GetDeviceBuffer()},
+                                                        scale_convert);
+
+        auto invoker_ptr = op_ptr->MakeInvokerPointer();
+
+        if(op_ptr->IsSupportedArgument(argument_ptr.get()))
+        {
+            invoker_ptr->Run(argument_ptr.get(), StreamConfig{nullptr, false});
+        }
+
+        std::cout << "Done" << std::endl;
+    }
+
+    return true;
+}
+
+template <typename InDataType,
+          typename OutDataType,
+          ck::ReduceTensorOp ReduceOpId,
+          ck::index_t NumDimSpatial,
+          ck::index_t NumNonSpatialDim>
+bool TensorFullReduction(SimpleDeviceMem& tensor,
+                         SimpleDeviceMem& out_amax,
+                         const std::array<ck::index_t, NumDimSpatial + NumNonSpatialDim>& lengths,
+                         const std::array<ck::index_t, NumDimSpatial + NumNonSpatialDim>& strides)
+{
+    const auto spatial_dim_size = std::accumulate(std::next(std::begin(lengths), NumNonSpatialDim),
+                                                  std::end(lengths),
+                                                  static_cast<std::size_t>(1),
+                                                  std::multiplies<>());
+    const auto tensor_size      = GetTensorSize<NumDimSpatial>(lengths);
+
+    auto copy = [](const auto& x, auto& y) { ck::ranges::copy(x, y.begin()); };
+
+    // Get the reduction operation
+    using ReduceOperation = typename ck::reduce_binary_operator<ReduceOpId>::opType;
+    using InElementwiseOperation =
+        typename ck::reduce_unary_operator<ReduceOpId, true, true>::InElementwiseOperation;
+    using AccElementwiseOperation =
+        typename ck::reduce_unary_operator<ReduceOpId, true, true>::AccElementwiseOperation;
+
+    InElementwiseOperation in_elementwise_op;
+    AccElementwiseOperation acc_elementwise_op;
+    std::tie(in_elementwise_op, acc_elementwise_op) =
+        ck::reduce_unary_operator<ReduceOpId, true, true>::GetElementwiseOperator(
+            static_cast<int32_t>(tensor_size));
+
+    std::array<ck::index_t, 1> reduce_out_lengths{1};
+    std::array<ck::index_t, 1> reduce_out_strides{1};
+
+    SimpleDeviceMem partial_reduce_tensor(sizeof(OutDataType) * spatial_dim_size);
+    std::array<ck::index_t, NumDimSpatial> reduce_part_lengths;
+    std::copy(std::next(std::begin(lengths), NumNonSpatialDim),
+              std::end(lengths),
+              std::begin(reduce_part_lengths));
+    std::array<ck::index_t, NumDimSpatial> reduce_part_strides;
+    copy(HostTensorDescriptor(reduce_part_lengths).GetStrides(), reduce_part_strides);
+
+    {
+        std::cout << "\nReduction of nonspatial dimensions:" << std::endl;
+        using DeviceOp =
+            ck::tensor_operation::device::DeviceReduce<InDataType,
+                                                       OutDataType,
+                                                       OutDataType,
+                                                       NumDimSpatial + NumNonSpatialDim,
+                                                       NumNonSpatialDim,
+                                                       ReduceOperation,
+                                                       InElementwiseOperation,
+                                                       PassThrough,
+                                                       true,   // PropagateNan
+                                                       false>; // OutputIndex
+        const auto op_ptrs = ck::tensor_operation::device::instance::DeviceOperationInstanceFactory<
+            DeviceOp>::GetInstances();
+
+        std::cout << "found " << op_ptrs.size() << " instances" << std::endl;
+
+        std::string best_op_name;
+        int best_op_id        = -1;
+        float best_ave_time   = std::numeric_limits<float>::max();
+        float best_gb_per_sec = 0;
+
+        std::array<int, NumNonSpatialDim> reduce_dims;
+        std::iota(reduce_dims.begin(), reduce_dims.end(), 0); // 0,..., NumNonSpatialDim-1
+
+        ck::index_t num_in_elements  = tensor_size;
+        ck::index_t num_out_elements = spatial_dim_size;
+
+        // profile device operation instances
+        std::cout << "Run partial reduction and do timing" << std::endl;
+
+        for(int i = 0; i < op_ptrs.size(); ++i)
+        {
+            auto& op_ptr = op_ptrs[i];
+
+            auto argument_ptr   = op_ptr->MakeArgumentPointer(lengths,
+                                                            strides,
+                                                            reduce_part_lengths,
+                                                            reduce_part_strides,
+                                                            reduce_dims,
+                                                            1.0,
+                                                            0.0,
+                                                            tensor.GetDeviceBuffer(),
+                                                            nullptr,
+                                                            partial_reduce_tensor.GetDeviceBuffer(),
+                                                            nullptr,
+                                                            in_elementwise_op,
+                                                            PassThrough{});
+            auto invoker_ptr    = op_ptr->MakeInvokerPointer();
+            std::string op_name = op_ptr->GetTypeString();
+
+            if(op_ptr->IsSupportedArgument(argument_ptr.get()))
+            {
+                float ave_time = invoker_ptr->Run(argument_ptr.get(), StreamConfig{nullptr, true});
+                std::size_t num_bytes =
+                    num_in_elements * sizeof(InDataType) + num_out_elements * sizeof(OutDataType);
+
+                float gb_per_sec = num_bytes / 1.E6 / ave_time;
+
+                std::cout << "Perf: " << std::setw(10) << ave_time << " ms, " << gb_per_sec
+                          << " GB/s, " << op_name << std::endl;
+
+                if(ave_time < best_ave_time)
+                {
+                    best_op_id      = i;
+                    best_op_name    = op_name;
+                    best_ave_time   = ave_time;
+                    best_gb_per_sec = gb_per_sec;
+                }
+            }
+            else
+            {
+                std::cout << op_name << " does not support this problem" << std::endl;
+            }
+        }
+
+        if(best_op_id < 0)
+        {
+            std::cerr << "no suitable instance found." << std::endl;
+            return false;
+        }
+        else
+        {
+            std::cout << "Best Perf: " << best_ave_time << " ms, " << best_gb_per_sec << " GB/s, "
+                      << best_op_name << std::endl;
+
+            // run the best instance
+            auto& op_ptr = op_ptrs[best_op_id];
+            std::cout << "Run the best instance without timing: " << op_ptr->GetTypeString()
+                      << std::endl;
+            auto argument_ptr = op_ptr->MakeArgumentPointer(lengths,
+                                                            strides,
+                                                            reduce_part_lengths,
+                                                            reduce_part_strides,
+                                                            reduce_dims,
+                                                            1.0,
+                                                            0.0,
+                                                            tensor.GetDeviceBuffer(),
+                                                            nullptr,
+                                                            partial_reduce_tensor.GetDeviceBuffer(),
+                                                            nullptr,
+                                                            in_elementwise_op,
+                                                            PassThrough{});
+
+            auto invoker_ptr = op_ptr->MakeInvokerPointer();
+
+            if(op_ptr->IsSupportedArgument(argument_ptr.get()))
+            {
+                invoker_ptr->Run(argument_ptr.get(), StreamConfig{nullptr, false});
+            }
+
+            std::cout << "Done" << std::endl;
+        }
+    }
+
+    {
+        std::cout << "\nReduction of spatial dimensions:" << std::endl;
+        using DeviceOp     = ck::tensor_operation::device::DeviceReduce<OutDataType,
+                                                                    OutDataType,
+                                                                    OutDataType,
+                                                                    NumDimSpatial,
+                                                                    NumDimSpatial,
+                                                                    ReduceOperation,
+                                                                    PassThrough,
+                                                                    AccElementwiseOperation,
+                                                                    true,   // PropagateNan
+                                                                    false>; // OutputIndex
+        const auto op_ptrs = ck::tensor_operation::device::instance::DeviceOperationInstanceFactory<
+            DeviceOp>::GetInstances();
+
+        std::cout << "found " << op_ptrs.size() << " instances" << std::endl;
+
+        std::string best_op_name;
+        int best_op_id        = -1;
+        float best_ave_time   = std::numeric_limits<float>::max();
+        float best_gb_per_sec = 0;
+
+        std::array<int, NumDimSpatial> reduce_dims;
+        std::iota(reduce_dims.begin(), reduce_dims.end(), 0); // 0,..., NumDimSpatial-1
+
+        ck::index_t num_in_elements  = spatial_dim_size;
+        ck::index_t num_out_elements = 1;
+
+        // profile device operation instances
+        std::cout << "Run final reduction and do timing" << std::endl;
+
+        for(int i = 0; i < op_ptrs.size(); ++i)
+        {
+            auto& op_ptr = op_ptrs[i];
+
+            auto argument_ptr   = op_ptr->MakeArgumentPointer(reduce_part_lengths,
+                                                            reduce_part_strides,
+                                                            reduce_out_lengths,
+                                                            reduce_out_strides,
+                                                            reduce_dims,
+                                                            1.0,
+                                                            0.0,
+                                                            partial_reduce_tensor.GetDeviceBuffer(),
+                                                            nullptr,
+                                                            out_amax.GetDeviceBuffer(),
+                                                            nullptr,
+                                                            PassThrough{},
+                                                            acc_elementwise_op);
+            auto invoker_ptr    = op_ptr->MakeInvokerPointer();
+            std::string op_name = op_ptr->GetTypeString();
+
+            if(op_ptr->IsSupportedArgument(argument_ptr.get()))
+            {
+                float ave_time = invoker_ptr->Run(argument_ptr.get(), StreamConfig{nullptr, true});
+
+                std::size_t num_bytes =
+                    num_in_elements * sizeof(OutDataType) + num_out_elements * sizeof(OutDataType);
+
+                float gb_per_sec = num_bytes / 1.E6 / ave_time;
+
+                std::cout << "Perf: " << std::setw(10) << ave_time << " ms, " << gb_per_sec
+                          << " GB/s, " << op_name << std::endl;
+
+                if(ave_time < best_ave_time)
+                {
+                    best_op_id      = i;
+                    best_op_name    = op_name;
+                    best_ave_time   = ave_time;
+                    best_gb_per_sec = gb_per_sec;
+                }
+            }
+            else
+            {
+                std::cout << op_name << " does not support this problem" << std::endl;
+            }
+        }
+
+        if(best_op_id < 0)
+        {
+            std::cerr << "no suitable instance found." << std::endl;
+            return false;
+        }
+        else
+        {
+            std::cout << "Best Perf: " << best_ave_time << " ms, " << best_gb_per_sec << " GB/s, "
+                      << best_op_name << std::endl;
+
+            // run the best instance
+            auto& op_ptr = op_ptrs[best_op_id];
+            std::cout << "Run the best instance without timing: " << op_ptr->GetTypeString()
+                      << std::endl;
+            auto argument_ptr = op_ptr->MakeArgumentPointer(reduce_part_lengths,
+                                                            reduce_part_strides,
+                                                            reduce_out_lengths,
+                                                            reduce_out_strides,
+                                                            reduce_dims,
+                                                            1.0,
+                                                            0.0,
+                                                            partial_reduce_tensor.GetDeviceBuffer(),
+                                                            nullptr,
+                                                            out_amax.GetDeviceBuffer(),
+                                                            nullptr,
+                                                            PassThrough{},
+                                                            acc_elementwise_op);
+
+            auto invoker_ptr = op_ptr->MakeInvokerPointer();
+
+            if(op_ptr->IsSupportedArgument(argument_ptr.get()))
+            {
+                invoker_ptr->Run(argument_ptr.get(), StreamConfig{nullptr, false});
+            }
+
+            std::cout << "Done" << std::endl;
+        }
+    }
+
+    return true;
+}

client_example/24_grouped_conv_activation/grouped_convnd_fwd_convscale_reduce/conv3d_fwd_convscale_amax_fp8.cpp

@@ -0,0 +1,58 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "common.hpp"
+
+#include "ck/ck.hpp"
+#include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
+
+using InDataType       = ck::f8_t;
+using WeiDataType      = ck::f8_t;
+using CShuffleDataType = float;
+using ConvOutDataType  = float;    // data type of convolution result
+using OutDataType      = ck::f8_t; // data type of final result
+using AComputeDataType = ck::f8_t;
+using BComputeDataType = ck::f8_t;
+
+using ConvElementOp = ConvScale;
+
+using InLayout  = ck::tensor_layout::convolution::NDHWGC;
+using WeiLayout = ck::tensor_layout::convolution::GKZYXC;
+using OutLayout = ck::tensor_layout::convolution::NDHWGK;
+
+constexpr auto ReduceOpId = ck::ReduceTensorOp::AMAX;
+
+static constexpr ck::index_t NumDimSpatial = 3;
+static constexpr ck::index_t G             = 1;
+static constexpr ck::index_t N             = 64;
+static constexpr ck::index_t K             = 128;
+static constexpr ck::index_t C             = 64;
+static constexpr ck::index_t Z             = 3;
+static constexpr ck::index_t Y             = 3;
+static constexpr ck::index_t X             = 3;
+static constexpr ck::index_t Di            = 28;
+static constexpr ck::index_t Hi            = 28;
+static constexpr ck::index_t Wi            = 3;
+static constexpr ck::index_t Do            = 28;
+static constexpr ck::index_t Ho            = 28;
+static constexpr ck::index_t Wo            = 3;
+
+int main()
+{
+    return run_grouped_conv_fwd_convscale_reduce<NumDimSpatial,
+                                                 InDataType,
+                                                 WeiDataType,
+                                                 ConvOutDataType,
+                                                 OutDataType,
+                                                 ConvElementOp,
+                                                 ReduceOpId,
+                                                 InLayout,
+                                                 WeiLayout,
+                                                 OutLayout,
+                                                 3,
+                                                 AComputeDataType,
+                                                 BComputeDataType>(
+               {N, Di, Hi, Wi, G, C}, {G, K, Z, Y, X, C}, {N, Do, Ho, Wo, G, K})
+               ? EXIT_SUCCESS
+               : EXIT_FAILURE;
+}

client_example/24_grouped_conv_activation/grouped_convnd_fwd_convscale_reduce/conv3d_fwd_convscale_relu_amax_fp8.cpp

@@ -0,0 +1,58 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "common.hpp"
+
+#include "ck/ck.hpp"
+#include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
+
+using InDataType       = ck::f8_t;
+using WeiDataType      = ck::f8_t;
+using CShuffleDataType = float;
+using ConvOutDataType  = float;    // data type of convolution result
+using OutDataType      = ck::f8_t; // data type of final result
+using AComputeDataType = ck::f8_t;
+using BComputeDataType = ck::f8_t;
+
+using ConvElementOp = ConvScaleRelu;
+
+using InLayout  = ck::tensor_layout::convolution::NDHWGC;
+using WeiLayout = ck::tensor_layout::convolution::GKZYXC;
+using OutLayout = ck::tensor_layout::convolution::NDHWGK;
+
+constexpr auto ReduceOpId = ck::ReduceTensorOp::AMAX;
+
+static constexpr ck::index_t NumDimSpatial = 3;
+static constexpr ck::index_t G             = 1;
+static constexpr ck::index_t N             = 64;
+static constexpr ck::index_t K             = 128;
+static constexpr ck::index_t C             = 64;
+static constexpr ck::index_t Z             = 3;
+static constexpr ck::index_t Y             = 3;
+static constexpr ck::index_t X             = 3;
+static constexpr ck::index_t Di            = 28;
+static constexpr ck::index_t Hi            = 28;
+static constexpr ck::index_t Wi            = 3;
+static constexpr ck::index_t Do            = 28;
+static constexpr ck::index_t Ho            = 28;
+static constexpr ck::index_t Wo            = 3;
+
+int main()
+{
+    return run_grouped_conv_fwd_convscale_reduce<NumDimSpatial,
+                                                 InDataType,
+                                                 WeiDataType,
+                                                 ConvOutDataType,
+                                                 OutDataType,
+                                                 ConvElementOp,
+                                                 ReduceOpId,
+                                                 InLayout,
+                                                 WeiLayout,
+                                                 OutLayout,
+                                                 3,
+                                                 AComputeDataType,
+                                                 BComputeDataType>(
+               {N, Di, Hi, Wi, G, C}, {G, K, Z, Y, X, C}, {N, Do, Ho, Wo, G, K})
+               ? EXIT_SUCCESS
+               : EXIT_FAILURE;
+}

client_example/CMakeLists.txt

@@ -34,8 +34,17 @@ if (DTYPES)
     endif()
     message("DTYPES macro set to ${DTYPES}")
 else()
-    add_definitions(-DCK_ENABLE_INT8 -DCK_ENABLE_FP8 -DCK_ENABLE_BF8 -DCK_ENABLE_FP16 -DCK_ENABLE_FP32 -DCK_ENABLE_FP64 -DCK_ENABLE_BF16)
-    set(CK_ENABLE_ALL_DTYPES "ON")
+    add_definitions(-DCK_ENABLE_INT8 -DCK_ENABLE_FP16 -DCK_ENABLE_FP32 -DCK_ENABLE_FP64 -DCK_ENABLE_BF16)
+    set(CK_ENABLE_INT8 "ON")
+    set(CK_ENABLE_FP16 "ON")
+    set(CK_ENABLE_FP32 "ON")
+    set(CK_ENABLE_FP64 "ON")
+    set(CK_ENABLE_BF16 "ON")
+    if (GPU_TARGETS MATCHES "gfx94")
+        add_definitions(-DCK_ENABLE_FP8 -DCK_ENABLE_BF8)
+        set(CK_ENABLE_FP8 "ON")
+        set(CK_ENABLE_BF8 "ON")
+    endif()
 endif()
 
 if (GPU_TARGETS)

codegen/CMakeLists.txt

@@ -27,6 +27,8 @@ file(GLOB_RECURSE KERNEL_FILES CONFIGURE_DEPENDS
 add_embed_library(ck_headers ${KERNEL_FILES} RELATIVE ${CK_ROOT}/include)
 
 file(GLOB SOURCES CONFIGURE_DEPENDS src/*.cpp)
+
+##message(STATUS "SOURCE_FILES: ${SOURCES}")
 # TODO: Use object library
 add_library(ck_host STATIC ${SOURCES})
 target_link_libraries(ck_host PRIVATE ck_headers)
@@ -48,6 +50,4 @@ rocm_install(
 )
 rocm_install(DIRECTORY include/ck DESTINATION ${CMAKE_INSTALL_INCLUDEDIR})
 
-if(BUILD_TESTING)
 add_subdirectory(test)
-endif()

codegen/test/CMakeLists.txt

@@ -1,15 +1,19 @@
 list(APPEND CMAKE_PREFIX_PATH /opt/rocm)
 add_subdirectory(rtc)
 file(GLOB TEST_SRCS CONFIGURE_DEPENDS *.cpp)
-foreach(TEST_SRC ${TEST_SRCS})
-  set_source_files_properties(${TEST_SRC} PROPERTIES LANGUAGE HIP)
-  get_filename_component(BASE_NAME ${TEST_SRC} NAME_WE)
-  add_executable(test_host_${BASE_NAME} ${TEST_SRC})
-  add_dependencies(codegen test_host_${BASE_NAME})
-  add_test(NAME codegen_test_${BASE_NAME} COMMAND test_host_${BASE_NAME})
-  target_link_libraries(test_host_${BASE_NAME} ck_rtc ck_host)
-  # target_link_libraries(test_host_${BASE_NAME} ${CK_ROOT}/build/lib/libutility.a)
-  target_include_directories(test_host_${BASE_NAME} PUBLIC include())
-  target_include_directories(test_host_${BASE_NAME} PUBLIC ${CK_ROOT}/include)
-  target_include_directories(test_host_${BASE_NAME} PUBLIC ${CK_ROOT}/library/include)
-endforeach()
+if(NOT INSTANCES_ONLY)
+  foreach(TEST_SRC ${TEST_SRCS})
+    set_source_files_properties(${TEST_SRC} PROPERTIES LANGUAGE HIP)
+    get_filename_component(BASE_NAME ${TEST_SRC} NAME_WE)
+    add_executable(codegen_test_${BASE_NAME} ${TEST_SRC})
+    add_dependencies(codegen codegen_test_${BASE_NAME})
+    add_dependencies(tests codegen_test_${BASE_NAME})
+    add_dependencies(check codegen_test_${BASE_NAME})
+    add_test(NAME codegen_test_${BASE_NAME} COMMAND codegen_test_${BASE_NAME})
+    message("adding test codegen_test_${BASE_NAME}")
+    target_link_libraries(codegen_test_${BASE_NAME} ck_rtc ck_host)
+    target_include_directories(codegen_test_${BASE_NAME} PUBLIC ${CK_ROOT}/codegen/test/include)
+    target_include_directories(codegen_test_${BASE_NAME} PUBLIC ${CK_ROOT}/include)
+    target_include_directories(codegen_test_${BASE_NAME} PUBLIC ${CK_ROOT}/library/include)
+  endforeach()
+endif()

codegen/test/rtc/CMakeLists.txt

@@ -1,5 +1,3 @@
-
-find_package(hip)
 file(GLOB RTC_SOURCES CONFIGURE_DEPENDS src/*.cpp)
 add_library(ck_rtc ${RTC_SOURCES})
 target_include_directories(ck_rtc PUBLIC include)

docs/sphinx/requirements.in

@@ -1,2 +1,2 @@
-rocm-docs-core==1.6.2
+rocm-docs-core==1.7.2
 sphinxcontrib-bibtex==2.6.2

docs/sphinx/requirements.txt

@@ -103,7 +103,7 @@ requests==2.32.3
     # via
     #   pygithub
     #   sphinx
-rocm-docs-core==1.6.2
+rocm-docs-core==1.7.2
     # via -r requirements.in
 six==1.16.0
     # via pybtex

example/01_gemm/gemm_xdl_fp8.cpp

@@ -1,5 +1,5 @@
 // SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
 
 #include "common.hpp"
 
@@ -7,7 +7,7 @@
 
 using ADataType        = ck::f8_t;
 using BDataType        = ck::f8_t;
-using CDataType        = ck::half_t;
+using CDataType        = ck::f8_t;
 using AccDataType      = float;
 using CShuffleDataType = float;
 

example/01_gemm/run_gemm_example.inc

@@ -1,5 +1,5 @@
 // SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
 
 #pragma once
 
@@ -34,11 +34,11 @@ inline __host__ __device__ constexpr double get_rtol()
     }
     else if constexpr(std::is_same_v<DataType, ck::f8_t>)
     {
-        return 1e-1; // 240 and 224 are acceptable
+        return 2e-1;
     }
     else if constexpr(std::is_same_v<DataType, ck::bf8_t>)
     {
-        return 1.5e-1; // 57344 and 49152 are acceptable
+        return 2e-1;
     }
     else
     {
@@ -75,11 +75,11 @@ inline __host__ __device__ constexpr double get_atol()
     }
     else if constexpr(std::is_same_v<DataType, ck::f8_t>)
     {
-        return 16.1; // 240 and 224 are acceptable
+        return 2e-1;
     }
     else if constexpr(std::is_same_v<DataType, ck::bf8_t>)
     {
-        return 8192.1; // 57344 and 49152 are acceptable
+        return 2e-1;
     }
     else
     {

example/12_reduce/reduce_blockwise.cpp

@@ -1,5 +1,5 @@
 // SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
 
 #include <iostream>
 #include <initializer_list>
@@ -255,34 +255,61 @@ int main(int argc, char* argv[])
     else
     {
         // for testing half_t
+        pass =
+            pass && reduce_blockwise_test<ck::half_t, float, ReduceOpId, PropagateNan, OutputIndex>(
+                        true, 2, true, {3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3}, {0, 1, 2}, 1.0f, 0.0f);
         pass =
             pass && reduce_blockwise_test<ck::half_t, float, ReduceOpId, PropagateNan, OutputIndex>(
                         true, 2, true, {16, 64, 32, 960}, {0, 1, 2}, 1.0f, 0.0f);
 
         // for testing float
+        pass =
+            pass && reduce_blockwise_test<float, float, ReduceOpId, PropagateNan, OutputIndex>(
+                        true, 2, true, {3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3}, {0, 1, 2}, 1.0f, 0.0f);
+
         pass = pass && reduce_blockwise_test<float, float, ReduceOpId, PropagateNan, OutputIndex>(
                            true, 2, true, {16, 64, 32, 960}, {0, 1, 2}, 1.0f, 0.0f);
 
         // for testing double
+        pass =
+            pass && reduce_blockwise_test<float, float, ReduceOpId, PropagateNan, OutputIndex>(
+                        true, 2, true, {3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3}, {0, 1, 2}, 1.0f, 0.0f);
+
         pass = pass && reduce_blockwise_test<float, float, ReduceOpId, PropagateNan, OutputIndex>(
                            true, 2, true, {16, 64, 32, 960}, {0, 1, 2}, 1.0f, 0.0f);
 
         // for testing bhalf_t
+        pass = pass &&
+               reduce_blockwise_test<ck::bhalf_t, float, ReduceOpId, PropagateNan, OutputIndex>(
+                   true, 2, true, {3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3}, {0, 1, 2}, 1.0f, 0.0f);
+
         pass = pass &&
                reduce_blockwise_test<ck::bhalf_t, float, ReduceOpId, PropagateNan, OutputIndex>(
                    true, 2, true, {16, 64, 32, 960}, {0, 1, 2}, 1.0f, 0.0f);
 
         // for testing int8_t
+        pass =
+            pass && reduce_blockwise_test<int8_t, int32_t, ReduceOpId, PropagateNan, OutputIndex>(
+                        true, 2, true, {3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3}, {0, 1, 2}, 1.0f, 0.0f);
+
         pass =
             pass && reduce_blockwise_test<int8_t, int32_t, ReduceOpId, PropagateNan, OutputIndex>(
                         true, 2, true, {16, 64, 32, 960}, {0, 1, 2}, 1.0f, 0.0f);
 
 #ifdef CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4
         // for testing int4_t using AVG operation
+        pass =
+            pass && reduce_blockwise_test<int4_t, int32_t, ReduceTensorOp::AVG, false, false>(
+                        true, 2, true, {3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3}, {0, 1, 2}, 1.0f, 0.0f);
+
         pass = pass && reduce_blockwise_test<int4_t, int32_t, ReduceTensorOp::AVG, false, false>(
                            true, 2, true, {16, 64, 32, 960}, {0, 1, 2}, 1.0f, 0.0f);
 
         // for testing int4_t using MAX operation
+        pass =
+            pass && reduce_blockwise_test<int4_t, int8_t, ReduceTensorOp::MAX, false, false>(
+                        true, 2, true, {3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3}, {0, 1, 2}, 1.0f, 0.0f);
+
         pass = pass && reduce_blockwise_test<int4_t, int8_t, ReduceTensorOp::MAX, false, false>(
                            true, 2, true, {16, 64, 32, 960}, {0, 1, 2}, 1.0f, 0.0f);
 #endif

example/12_reduce/reduce_example_common.hpp

@@ -1,5 +1,5 @@
 // SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
 
 #pragma once
 
@@ -38,7 +38,8 @@ struct ReduceShape
     static constexpr ck::index_t NumReduceDim_ = NumReduceDim;
 };
 
-using reduce_shape_instances = std::tuple<ReduceShape<3, 1>,
+using reduce_shape_instances = std::tuple<ReduceShape<12, 3>,
+                                          ReduceShape<3, 1>,
                                           ReduceShape<3, 2>,
                                           ReduceShape<4, 1>,
                                           ReduceShape<4, 2>,

example/20_grouped_conv_bwd_weight/common.hpp

@@ -23,12 +23,8 @@
 using BF16 = ck::bhalf_t;
 using F16  = ck::half_t;
 using F32  = float;
-#ifdef CK_ENABLE_FP8
-using F8 = ck::f8_t;
-#endif
-#ifdef CK_ENABLE_BF8
-using BF8 = ck::bf8_t;
-#endif
+using F8   = ck::f8_t;
+using BF8  = ck::bf8_t;
 
 template <ck::index_t... Is>
 using S = ck::Sequence<Is...>;

example/62_convnd_activ/CMakeLists.txt

@@ -3,6 +3,7 @@ add_subdirectory(convinvscale)
 add_subdirectory(convscale)
 add_subdirectory(convscale_relu)
 add_subdirectory(convscale_add)
+add_subdirectory(convscale_reduce)
 add_subdirectory(multi_AB)
 add_subdirectory(unary)
 

example/62_convnd_activ/convscale_reduce/CMakeLists.txt

@@ -0,0 +1,14 @@
+list(APPEND gpu_list gfx908 gfx90a gfx940 gfx941 gfx942)
+set(target 0)
+foreach(gpu IN LISTS GPU_TARGETS)
+ if(gpu IN_LIST gpu_list AND target EQUAL 0)
+     add_custom_target(example_convnd_activ_xdl_convscale_reduce)
+     add_example_executable(example_convnd_fwd_xdl_convscale_relu_amax_fp8 convnd_fwd_xdl_convscale_relu_amax_fp8.cpp)
+     add_example_dependencies(example_convnd_activ_xdl_convscale_reduce example_convnd_fwd_xdl_convscale_relu_amax_fp8)
+
+     add_example_executable(example_convnd_fwd_xdl_convscale_amax_fp8 convnd_fwd_xdl_convscale_amax_fp8.cpp)
+     add_example_dependencies(example_convnd_activ_xdl_convscale_reduce example_convnd_fwd_xdl_convscale_amax_fp8)
+
+     set(target 1)
+ endif()
+endforeach()

example/62_convnd_activ/convscale_reduce/convnd_fwd_convscale_reduce_common.hpp

@@ -0,0 +1,502 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#include <cstdlib>
+#include <iostream>
+
+#include "ck/ck.hpp"
+
+#include "ck/library/utility/algorithm.hpp"
+#include "ck/library/utility/check_err.hpp"
+#include "ck/library/utility/device_memory.hpp"
+#include "ck/library/utility/host_tensor.hpp"
+#include "ck/library/utility/host_tensor_generator.hpp"
+#include "ck/library/utility/convolution_parameter.hpp"
+#include "ck/library/reference_tensor_operation/cpu/reference_conv_fwd.hpp"
+#include "ck/library/reference_tensor_operation/cpu/reference_reduce.hpp"
+#include "ck/tensor_operation/gpu/element/combined_element_wise_operation.hpp"
+#include "ck/tensor_operation/gpu/element/unary_element_wise_operation.hpp"
+#include "ck/tensor_operation/gpu/device/impl/device_elementwise_dynamic_vector_dims_impl.hpp"
+#include "ck/tensor_operation/gpu/device/impl/device_reduce_multiblock.hpp"
+#include "ck/utility/reduction_operator.hpp"
+#include "ck/utility/reduction_enums.hpp"
+#include "ck/tensor_operation/gpu/device/reduction_operator_mapping.hpp"
+#include "ck/utility/type.hpp"
+
+namespace ew = ck::tensor_operation::element_wise;
+
+using PassThrough   = ew::PassThrough;
+using ConvScaleRelu = ew::UnaryCombinedOp<ew::Scale, ew::Scale, ew::Relu>;
+using ConvScale     = ew::UnaryCombinedOp<ew::Scale, ew::Scale, PassThrough>;
+
+using UnaryScaleConvert = ew::Scale;
+
+void print_helper_msg()
+{
+    std::cout << "arg1: verification (0=no, 1=yes)\n"
+              << "arg2: initialization (0=no init, 1=integer value, 2=decimal value)\n"
+              << "arg3: time kernel (0=no, 1=yes)\n"
+              << ck::utils::conv::get_conv_param_parser_helper_msg() << std::endl;
+}
+
+template <typename DataType>
+inline __host__ __device__ constexpr double get_rtol()
+{
+    if constexpr(std::is_same_v<DataType, float>)
+    {
+        return 1e-3;
+    }
+    else if constexpr(std::is_same_v<DataType, double>)
+    {
+        return 1e-6;
+    }
+    else if constexpr(std::is_same_v<DataType, ck::half_t>)
+    {
+        return 1e-3;
+    }
+    else if constexpr(std::is_same_v<DataType, ck::bhalf_t>)
+    {
+        return 5e-2;
+    }
+    else if constexpr(std::is_same_v<DataType, int32_t>)
+    {
+        return 1e-1;
+    }
+    else if constexpr(std::is_same_v<DataType, int8_t>)
+    {
+        return 1e-1;
+    }
+    else if constexpr(std::is_same_v<DataType, ck::f8_t>)
+    {
+        return 1e-1; // 240 and 224 are acceptable
+    }
+    else if constexpr(std::is_same_v<DataType, ck::bf8_t>)
+    {
+        return 1.5e-1; // 57344 and 49152 are acceptable
+    }
+    else
+    {
+        return 1e-3;
+    }
+}
+
+template <typename DataType>
+inline __host__ __device__ constexpr double get_atol()
+{
+    if constexpr(std::is_same_v<DataType, float>)
+    {
+        return 1e-3;
+    }
+    else if constexpr(std::is_same_v<DataType, double>)
+    {
+        return 1e-6;
+    }
+    else if constexpr(std::is_same_v<DataType, ck::half_t>)
+    {
+        return 1e-3;
+    }
+    else if constexpr(std::is_same_v<DataType, ck::bhalf_t>)
+    {
+        return 5e-2;
+    }
+    else if constexpr(std::is_same_v<DataType, int32_t>)
+    {
+        return 1e-1;
+    }
+    else if constexpr(std::is_same_v<DataType, int8_t>)
+    {
+        return 1e-1;
+    }
+    else if constexpr(std::is_same_v<DataType, ck::f8_t>)
+    {
+        return 16.1; // 240 and 224 are acceptable
+    }
+    else if constexpr(std::is_same_v<DataType, ck::bf8_t>)
+    {
+        return 8192.1; // 57344 and 49152 are acceptable
+    }
+    else
+    {
+        return 1e-3;
+    }
+}
+
+template <ck::index_t NDimSpatial,
+          typename InDataType,
+          typename WeiDataType,
+          typename ConvOutDataType,
+          typename OutDataType,
+          typename InElementOp,
+          typename WeiElementOp,
+          typename ConvElementOp,
+          typename DeviceConvNDFwdInstance>
+bool run_grouped_conv_fwd(bool do_verification,
+                          int init_method,
+                          bool time_kernel,
+                          const ck::utils::conv::ConvParam& conv_param,
+                          const HostTensorDescriptor& in_g_n_c_wis_desc,
+                          const HostTensorDescriptor& wei_g_k_c_xs_desc,
+                          const HostTensorDescriptor& out_g_n_k_wos_desc,
+                          const InElementOp& in_element_op,
+                          const WeiElementOp& wei_element_op)
+{
+    Tensor<InDataType> in(in_g_n_c_wis_desc);
+    Tensor<WeiDataType> wei(wei_g_k_c_xs_desc);
+    Tensor<ConvOutDataType> host_conv(out_g_n_k_wos_desc);
+    Tensor<ConvOutDataType> device_conv(out_g_n_k_wos_desc);
+    Tensor<OutDataType> out_host(out_g_n_k_wos_desc);
+    Tensor<OutDataType> out_device(out_g_n_k_wos_desc);
+
+    std::cout << "in: " << in.mDesc << std::endl;
+    std::cout << "wei: " << wei.mDesc << std::endl;
+    std::cout << "out: " << out_host.mDesc << std::endl;
+
+    switch(init_method)
+    {
+    case 0: break;
+    case 1:
+        in.GenerateTensorValue(GeneratorTensor_2<InDataType>{-5, 5});
+        wei.GenerateTensorValue(GeneratorTensor_2<WeiDataType>{-5, 5});
+        break;
+    case 11: // used for debugging
+        in.GenerateTensorValue(GeneratorTensor_1<InDataType>{1});
+        wei.GenerateTensorValue(GeneratorTensor_1<WeiDataType>{1});
+        break;
+    default:
+        in.GenerateTensorValue(GeneratorTensor_3<InDataType>{-1.0, 1.0});
+        wei.GenerateTensorValue(GeneratorTensor_3<WeiDataType>{-0.5, 0.5});
+    }
+
+    DeviceMem in_device_buf(sizeof(InDataType) * in.mDesc.GetElementSpaceSize());
+    DeviceMem wei_device_buf(sizeof(WeiDataType) * wei.mDesc.GetElementSpaceSize());
+    DeviceMem conv_device_buf(conv_param.GetOutputByte<ConvOutDataType>());
+    DeviceMem out_device_buf(conv_param.GetOutputByte<OutDataType>());
+
+    in_device_buf.ToDevice(in.mData.data());
+    wei_device_buf.ToDevice(wei.mData.data());
+
+    std::array<ck::index_t, NDimSpatial + 3> a_g_n_c_wis_lengths{};
+    std::array<ck::index_t, NDimSpatial + 3> a_g_n_c_wis_strides{};
+    std::array<ck::index_t, NDimSpatial + 3> b_g_k_c_xs_lengths{};
+    std::array<ck::index_t, NDimSpatial + 3> b_g_k_c_xs_strides{};
+    std::array<ck::index_t, NDimSpatial + 3> e_g_n_k_wos_lengths{};
+    std::array<ck::index_t, NDimSpatial + 3> e_g_n_k_wos_strides{};
+    std::array<ck::index_t, NDimSpatial> conv_filter_strides{};
+    std::array<ck::index_t, NDimSpatial> conv_filter_dilations{};
+    std::array<ck::index_t, NDimSpatial> input_left_pads{};
+    std::array<ck::index_t, NDimSpatial> input_right_pads{};
+
+    auto copy = [](const auto& x, auto& y) { ck::ranges::copy(x, y.begin()); };
+
+    copy(in_g_n_c_wis_desc.GetLengths(), a_g_n_c_wis_lengths);
+    copy(in_g_n_c_wis_desc.GetStrides(), a_g_n_c_wis_strides);
+    copy(wei_g_k_c_xs_desc.GetLengths(), b_g_k_c_xs_lengths);
+    copy(wei_g_k_c_xs_desc.GetStrides(), b_g_k_c_xs_strides);
+    copy(out_g_n_k_wos_desc.GetLengths(), e_g_n_k_wos_lengths);
+    copy(out_g_n_k_wos_desc.GetStrides(), e_g_n_k_wos_strides);
+    copy(conv_param.conv_filter_strides_, conv_filter_strides);
+    copy(conv_param.conv_filter_dilations_, conv_filter_dilations);
+    copy(conv_param.input_left_pads_, input_left_pads);
+    copy(conv_param.input_right_pads_, input_right_pads);
+
+    // random scale values
+    float scale_in  = float(std::rand()) / float(RAND_MAX);
+    float scale_wei = float(std::rand()) / float(RAND_MAX);
+    float scale_out = float(std::rand()) / float(RAND_MAX);
+
+    std::cout << std::endl;
+    std::cout << "scale_in: " << scale_in << std::endl;
+    std::cout << "scale_wei: " << scale_wei << std::endl;
+    std::cout << "scale_out: " << scale_out << std::endl;
+
+    // convolution elementwise operation
+    auto conv_element_op = ConvElementOp{ew::Scale{scale_in}, ew::Scale{scale_wei}, {}};
+    auto scale_convert   = UnaryScaleConvert{scale_out}; // elementwise scale and type cast
+
+    // do Conv
+    auto conv         = DeviceConvNDFwdInstance{};
+    auto conv_invoker = conv.MakeInvoker();
+    auto conv_argument =
+        conv.MakeArgument(in_device_buf.GetDeviceBuffer(),
+                          wei_device_buf.GetDeviceBuffer(),
+                          std::array<const void*, 0>{},
+                          conv_device_buf.GetDeviceBuffer(),
+                          a_g_n_c_wis_lengths,
+                          a_g_n_c_wis_strides,
+                          b_g_k_c_xs_lengths,
+                          b_g_k_c_xs_strides,
+                          std::array<std::array<ck::index_t, NDimSpatial + 3>, 0>{},
+                          std::array<std::array<ck::index_t, NDimSpatial + 3>, 0>{},
+                          e_g_n_k_wos_lengths,
+                          e_g_n_k_wos_strides,
+                          conv_filter_strides,
+                          conv_filter_dilations,
+                          input_left_pads,
+                          input_right_pads,
+                          in_element_op,
+                          wei_element_op,
+                          conv_element_op);
+
+    if(!conv.IsSupportedArgument(conv_argument))
+    {
+        throw std::runtime_error(
+            "wrong! device_conv with the specified compilation parameters does "
+            "not support this Conv problem");
+    }
+
+    std::string kernels = conv.GetTypeString();
+
+    float avg_time = conv_invoker.Run(conv_argument, StreamConfig{nullptr, time_kernel});
+
+    using DeviceElementwiseScale = ck::tensor_operation::device::DeviceElementwiseImpl<
+        ck::Tuple<ConvOutDataType>, // InDataTypeTuple
+        ck::Tuple<OutDataType>,     // OutDataTypeTuple
+        UnaryScaleConvert,          // UnaryScaleConvert
+        NDimSpatial + 3,            // NumDim
+        256,                        // BlockSize
+        128,                        // M0PerBlock
+        128,                        // M1PerBlock
+        8,                          // M0PerThread
+        8,                          // M1PerThread
+        ck::Sequence<1, 0>,         // ThreadClusterArrangeOrder
+        ck::Sequence<8>,            // InScalarPerVectorSeq
+        ck::Sequence<8>>;           // OutScalarPerVectorSeq
+
+    auto device_ew_scale = DeviceElementwiseScale{};
+    auto scale_invoker   = device_ew_scale.MakeInvoker();
+    auto scale_argument  = device_ew_scale.MakeArgument(e_g_n_k_wos_lengths,
+                                                       {e_g_n_k_wos_strides},
+                                                       {e_g_n_k_wos_strides},
+                                                       {conv_device_buf.GetDeviceBuffer()},
+                                                       {out_device_buf.GetDeviceBuffer()},
+                                                       scale_convert);
+
+    if(!device_ew_scale.IsSupportedArgument(scale_argument))
+    {
+        throw std::runtime_error(
+            "wrong! DeviceElementwiseScale with the specified compilation parameters does "
+            "not support this problem");
+    }
+
+    kernels += std::string("\n\t\t ") + device_ew_scale.GetTypeString();
+
+    avg_time += scale_invoker.Run(scale_argument, StreamConfig{nullptr, time_kernel});
+
+    constexpr auto ReduceOpId = ck::ReduceTensorOp::AMAX;
+    using ReduceOperation     = typename ck::reduce_binary_operator<ReduceOpId>::opType;
+    using InElementwiseOperation =
+        typename ck::reduce_unary_operator<ReduceOpId, true, true>::InElementwiseOperation;
+    using AccElementwiseOperation =
+        typename ck::reduce_unary_operator<ReduceOpId, true, true>::AccElementwiseOperation;
+    using DeviceReduceInstance =
+        ck::tensor_operation::device::DeviceReduceMultiBlock<ConvOutDataType,
+                                                             ConvOutDataType,
+                                                             ConvOutDataType,
+                                                             NDimSpatial + 3,
+                                                             NDimSpatial + 3,
+                                                             ReduceOperation,
+                                                             InElementwiseOperation,
+                                                             AccElementwiseOperation,
+                                                             ck::InMemoryDataOperationEnum::Set,
+                                                             true,  // PropagateNan
+                                                             false, // OutputIndex
+                                                             false, // HaveIndexInputIfOutputIndex
+                                                             256,   // BlockSize
+                                                             4,     // MThreadClusterSize
+                                                             64,    // KThreadClusterSize
+                                                             1,     // MThreadSliceSize
+                                                             1,     // KThreadSliceSize
+                                                             1,     // InSrcVectorDim
+                                                             1,     // InSrceVectorSize
+                                                             1>;    // OutDstVectorSize
+
+    std::vector<size_t> outLengths = {1};
+    Tensor<ConvOutDataType> amax_host(outLengths);
+    Tensor<ConvOutDataType> amax_from_device(outLengths);
+    auto amax_host_strides = amax_host.mDesc.GetStrides();
+
+    std::array<int, NDimSpatial + 3> reduce_dims;
+    std::iota(reduce_dims.begin(), reduce_dims.end(), 0); // 0,..., NDimSpatial+3-1
+
+    std::array<ck::index_t, 1> reduce_out_lengths{1};
+    std::array<ck::index_t, 1> reduce_out_strides{static_cast<ck::index_t>(amax_host_strides[0])};
+
+    DeviceMem amax_device(sizeof(ConvOutDataType) * amax_host.mDesc.GetElementSpaceSize());
+    DeviceMem index_device;
+
+    InElementwiseOperation in_elementwise_op;
+    AccElementwiseOperation acc_elementwise_op;
+    std::tie(in_elementwise_op, acc_elementwise_op) =
+        ck::reduce_unary_operator<ReduceOpId, true, true>::GetElementwiseOperator(
+            static_cast<int32_t>(host_conv.mDesc.GetElementSize()));
+
+    // Hack convolution output strides for reduction as kernel expects stride 1 for the last
+    // dimension. It only works because the reduction is done on the whole tensor and result is
+    // independent of the order of elements.
+    std::array<ck::index_t, NDimSpatial + 3> reduction_strides{};
+    copy(HostTensorDescriptor(e_g_n_k_wos_lengths).GetStrides(), reduction_strides);
+
+    auto device_reduce   = DeviceReduceInstance{};
+    auto reduce_invoker  = device_reduce.MakeInvokerPointer();
+    auto reduce_argument = device_reduce.MakeArgumentPointer(e_g_n_k_wos_lengths,
+                                                             reduction_strides,
+                                                             reduce_out_lengths,
+                                                             reduce_out_strides,
+                                                             reduce_dims,
+                                                             1.0,
+                                                             0.0,
+                                                             conv_device_buf.GetDeviceBuffer(),
+                                                             nullptr,
+                                                             amax_device.GetDeviceBuffer(),
+                                                             nullptr,
+                                                             in_elementwise_op,
+                                                             acc_elementwise_op);
+
+    if(!device_reduce.IsSupportedArgument(reduce_argument.get()))
+    {
+        throw std::runtime_error(
+            "wrong! DeviceReduceInstance with the specified compilation parameters does "
+            "not support this runtime parameters!");
+    };
+
+    kernels += std::string("\n\t\t ") + device_reduce.GetTypeString();
+
+    float reduce_time =
+        reduce_invoker->Run(reduce_argument.get(), StreamConfig{nullptr, time_kernel});
+
+    if(time_kernel)
+        std::cout << "\nReduce time: " << reduce_time << " ms" << std::endl;
+
+    avg_time += reduce_time;
+
+    std::size_t flop    = conv_param.GetFlops();      // convolution FLOPs
+    auto conv_out_elems = host_conv.GetElementSize(); // number of elements in conv result tensor
+
+    // 3 element-wise scale multipliers + 1 AMAX
+    std::size_t elementwise_ops = 3 + 1;
+    if constexpr(ck::is_same_v<ConvElementOp, ConvScaleRelu>)
+    {
+        elementwise_ops += 1; // +1 element-wise relu
+    }
+
+    flop += elementwise_ops * conv_out_elems;
+
+    // convolution + elementwise scaling (in + wei + output byte count)
+    std::size_t num_btype = conv_param.GetByte<InDataType, WeiDataType, ConvOutDataType>();
+    num_btype += sizeof(float) + sizeof(float); //  + 2 scales
+
+    // elementwise scaling + F8 conversion
+    num_btype += conv_param.GetOutputByte<ConvOutDataType>() + sizeof(float) +
+                 conv_param.GetOutputByte<OutDataType>();
+
+    // AMAX
+    num_btype += conv_param.GetOutputByte<ConvOutDataType>() + sizeof(float);
+
+    if(time_kernel)
+    {
+        float tflops     = static_cast<float>(flop) / 1.E9 / avg_time;
+        float gb_per_sec = num_btype / 1.E6 / avg_time;
+        std::cout << "Perf: " << avg_time << " ms, " << tflops << " TFlops, " << gb_per_sec
+                  << " GB/s, " << std::endl;
+    }
+
+    std::cout << "\nKernels: " << kernels << std::endl;
+
+    if(do_verification)
+    {
+        auto ref_conv = ck::tensor_operation::host::ReferenceConvFwd<NDimSpatial,
+                                                                     InDataType,
+                                                                     WeiDataType,
+                                                                     ConvOutDataType,
+                                                                     InElementOp,
+                                                                     WeiElementOp,
+                                                                     ConvElementOp>();
+
+        auto ref_invoker  = ref_conv.MakeInvoker();
+        auto ref_argument = ref_conv.MakeArgument(in,
+                                                  wei,
+                                                  host_conv,
+                                                  conv_param.conv_filter_strides_,
+                                                  conv_param.conv_filter_dilations_,
+                                                  conv_param.input_left_pads_,
+                                                  conv_param.input_right_pads_,
+                                                  in_element_op,
+                                                  wei_element_op,
+                                                  conv_element_op);
+
+        ref_invoker.Run(ref_argument);
+
+        conv_device_buf.FromDevice(device_conv.mData.data());
+
+        out_device_buf.FromDevice(out_device.mData.data());
+
+        out_host.ForEach([&](auto&, auto idx) { scale_convert(out_host(idx), host_conv(idx)); });
+
+        std::cout << "\nComparing output to reference: " << std::endl;
+        auto tight_tol_check = ck::utils::check_err(out_device, out_host, "Error: ");
+        if(!tight_tol_check)
+        {
+            std::cout << "\n\tRecompare applying tolerances...\n";
+            std::cout << "\t\trtol = " << get_rtol<OutDataType>() << std::endl;
+            std::cout << "\t\tatol = " << get_atol<OutDataType>() << std::endl;
+            auto loose_tol_check = ck::utils::check_err(out_device,
+                                                        out_host,
+                                                        "Error: incorrect convolution results!",
+                                                        get_rtol<OutDataType>(),
+                                                        get_atol<OutDataType>());
+            if(!loose_tol_check)
+            {
+                return false;
+            }
+        }
+        std::cout << "Success!" << std::endl;
+
+        /// Verify AMAX
+
+        using RefReduceInstance =
+            ck::tensor_operation::host::ReferenceReduce<ConvOutDataType,
+                                                        ConvOutDataType,
+                                                        ConvOutDataType,
+                                                        NDimSpatial + 3,
+                                                        NDimSpatial + 3,
+                                                        ReduceOperation,
+                                                        InElementwiseOperation,
+                                                        AccElementwiseOperation,
+                                                        true,
+                                                        false>;
+
+        auto ref_reduce          = RefReduceInstance{};
+        auto ref_reduce_invoker  = ref_reduce.MakeInvokerPointer();
+        auto ref_reduce_argument = ref_reduce.MakeArgumentPointer(e_g_n_k_wos_lengths,
+                                                                  e_g_n_k_wos_strides,
+                                                                  reduce_out_lengths,
+                                                                  reduce_out_strides,
+                                                                  reduce_dims,
+                                                                  1.0,
+                                                                  0.0,
+                                                                  host_conv.mData.data(),
+                                                                  nullptr,
+                                                                  amax_host.mData.data(),
+                                                                  nullptr,
+                                                                  in_elementwise_op,
+                                                                  acc_elementwise_op);
+
+        if(!ref_reduce.IsSupportedArgument(ref_reduce_argument.get()))
+        {
+            throw std::runtime_error(
+                "wrong! RefReduceInstance with the specified compilation parameters does "
+                "not support this runtime parameters!");
+        };
+
+        ref_reduce_invoker->Run(ref_reduce_argument.get());
+
+        amax_device.FromDevice(amax_from_device.mData.data());
+
+        std::cout << "\namax: " << amax_from_device.mData[0] << std::endl;
+        std::cout << "amax_ref: " << amax_host.mData[0] << std::endl;
+
+        return ck::utils::check_err(amax_from_device, amax_host, "Error: incorrect AMAX results!");
+    }
+
+    return true;
+}

example/62_convnd_activ/convscale_reduce/convnd_fwd_xdl_convscale_amax_fp8.cpp

@@ -0,0 +1,82 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "convnd_fwd_convscale_reduce_common.hpp"
+
+#include "ck/tensor_operation/gpu/device/impl/device_grouped_conv_fwd_multiple_abd_xdl_cshuffle.hpp"
+
+using InDataType       = ck::f8_t;
+using WeiDataType      = ck::f8_t;
+using AccDataType      = float;
+using CShuffleDataType = float;
+using ConvOutDataType  = float;    // data type of convolution result
+using OutDataType      = ck::f8_t; // data type of final result
+using AComputeDataType = ck::f8_t;
+using BComputeDataType = ck::f8_t;
+
+template <ck::index_t... Is>
+using S = ck::Sequence<Is...>;
+
+using InElementOp  = PassThrough;
+using WeiElementOp = PassThrough;
+using OutElementOp = ConvScale;
+
+static constexpr auto ConvSpec =
+    ck::tensor_operation::device::ConvolutionForwardSpecialization::Default;
+
+static constexpr auto GemmSpec = ck::tensor_operation::device::GemmSpecialization::MNKPadding;
+
+template <ck::index_t NDimSpatial, typename InLayout, typename WeiLayout, typename OutLayout>
+using DeviceGroupedConvNDFwdInstance =
+    ck::tensor_operation::device::DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle<
+        NDimSpatial,
+        InLayout,
+        WeiLayout,
+        ck::Tuple<>,
+        OutLayout,
+        InDataType,
+        WeiDataType,
+        AccDataType,
+        CShuffleDataType,
+        ck::Tuple<>,
+        ConvOutDataType,
+        InElementOp,
+        WeiElementOp,
+        OutElementOp,
+        ConvSpec,    // ConvForwardSpecialization
+        GemmSpec,    // GemmSpecialization
+        1,           //
+        256,         // BlockSize
+        128,         // MPerBlock
+        256,         // NPerBlock
+        32,          // KPerBlock
+        8,           // AK1
+        8,           // BK1
+        32,          // MPerXdl
+        32,          // NPerXdl
+        2,           // MXdlPerWave
+        4,           // NXdlPerWave
+        S<4, 64, 1>, // ABlockTransferThreadClusterLengths_AK0_M_AK1
+        S<1, 0, 2>,  // ABlockTransferThreadClusterArrangeOrder
+        S<1, 0, 2>,  // ABlockTransferSrcAccessOrder
+        2,           // ABlockTransferSrcVectorDim
+        8,           // ABlockTransferSrcScalarPerVector
+        8,           // ABlockTransferDstScalarPerVector_AK1
+        1,           // ABlockLdsExtraM
+        S<4, 64, 1>, // BBlockTransferThreadClusterLengths_BK0_N_BK1
+        S<1, 0, 2>,  // BBlockTransferThreadClusterArrangeOrder
+        S<1, 0, 2>,  // BBlockTransferSrcAccessOrder
+        2,           // BBlockTransferSrcVectorDim
+        8,           // BBlockTransferSrcScalarPerVector
+        8,           // BBlockTransferDstScalarPerVector_BK1
+        1,           // BBlockLdsExtraN
+        1,
+        1,
+        S<1, 32, 1, 8>,
+        8,
+        AComputeDataType,
+        BComputeDataType>;
+
+#include "run_convnd_fwd_example.inc"
+
+int main(int argc, char* argv[]) { return run_convnd_fwd_example(argc, argv) ? 0 : 1; }

example/62_convnd_activ/convscale_reduce/convnd_fwd_xdl_convscale_relu_amax_fp8.cpp

@@ -0,0 +1,82 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "convnd_fwd_convscale_reduce_common.hpp"
+
+#include "ck/tensor_operation/gpu/device/impl/device_grouped_conv_fwd_multiple_abd_xdl_cshuffle.hpp"
+
+using InDataType       = ck::f8_t;
+using WeiDataType      = ck::f8_t;
+using AccDataType      = float;
+using CShuffleDataType = float;
+using ConvOutDataType  = float;    // data type of convolution result
+using OutDataType      = ck::f8_t; // data type of final result
+using AComputeDataType = ck::f8_t;
+using BComputeDataType = ck::f8_t;
+
+template <ck::index_t... Is>
+using S = ck::Sequence<Is...>;
+
+using InElementOp  = PassThrough;
+using WeiElementOp = PassThrough;
+using OutElementOp = ConvScaleRelu;
+
+static constexpr auto ConvSpec =
+    ck::tensor_operation::device::ConvolutionForwardSpecialization::Default;
+
+static constexpr auto GemmSpec = ck::tensor_operation::device::GemmSpecialization::MNKPadding;
+
+template <ck::index_t NDimSpatial, typename InLayout, typename WeiLayout, typename OutLayout>
+using DeviceGroupedConvNDFwdInstance =
+    ck::tensor_operation::device::DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle<
+        NDimSpatial,
+        InLayout,
+        WeiLayout,
+        ck::Tuple<>,
+        OutLayout,
+        InDataType,
+        WeiDataType,
+        AccDataType,
+        CShuffleDataType,
+        ck::Tuple<>,
+        ConvOutDataType,
+        InElementOp,
+        WeiElementOp,
+        OutElementOp,
+        ConvSpec,    // ConvForwardSpecialization
+        GemmSpec,    // GemmSpecialization
+        1,           //
+        256,         // BlockSize
+        128,         // MPerBlock
+        256,         // NPerBlock
+        32,          // KPerBlock
+        8,           // AK1
+        8,           // BK1
+        32,          // MPerXdl
+        32,          // NPerXdl
+        2,           // MXdlPerWave
+        4,           // NXdlPerWave
+        S<4, 64, 1>, // ABlockTransferThreadClusterLengths_AK0_M_AK1
+        S<1, 0, 2>,  // ABlockTransferThreadClusterArrangeOrder
+        S<1, 0, 2>,  // ABlockTransferSrcAccessOrder
+        2,           // ABlockTransferSrcVectorDim
+        8,           // ABlockTransferSrcScalarPerVector
+        8,           // ABlockTransferDstScalarPerVector_AK1
+        1,           // ABlockLdsExtraM
+        S<4, 64, 1>, // BBlockTransferThreadClusterLengths_BK0_N_BK1
+        S<1, 0, 2>,  // BBlockTransferThreadClusterArrangeOrder
+        S<1, 0, 2>,  // BBlockTransferSrcAccessOrder
+        2,           // BBlockTransferSrcVectorDim
+        8,           // BBlockTransferSrcScalarPerVector
+        8,           // BBlockTransferDstScalarPerVector_BK1
+        1,           // BBlockLdsExtraN
+        1,
+        1,
+        S<1, 32, 1, 8>,
+        8,
+        AComputeDataType,
+        BComputeDataType>;
+
+#include "run_convnd_fwd_example.inc"
+
+int main(int argc, char* argv[]) { return run_convnd_fwd_example(argc, argv) ? 0 : 1; }

example/62_convnd_activ/convscale_reduce/run_convnd_fwd_example.inc

@@ -0,0 +1,98 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+bool run_convnd_fwd_example(int argc, char* argv[])
+{
+    print_helper_msg();
+
+    bool do_verification = true;
+    int init_method      = 1;
+    bool time_kernel     = false;
+
+    ck::utils::conv::ConvParam conv_param{
+        2, 1, 128, 256, 192, {3, 3}, {71, 71}, {2, 2}, {1, 1}, {1, 1}, {1, 1}};
+
+    if(argc == 1)
+    {
+        // use default
+    }
+    else if(argc == 4)
+    {
+        do_verification = std::stoi(argv[1]);
+        init_method     = std::stoi(argv[2]);
+        time_kernel     = std::stoi(argv[3]);
+    }
+    else
+    {
+        do_verification                   = std::stoi(argv[1]);
+        init_method                       = std::stoi(argv[2]);
+        time_kernel                       = std::stoi(argv[3]);
+        const ck::index_t num_dim_spatial = std::stoi(argv[4]);
+
+        conv_param = ck::utils::conv::parse_conv_param(num_dim_spatial, 5, argv);
+    }
+
+    // instantiate in and wei element ops, will
+    // instantiate out_element_op below for every iteration
+    const auto in_element_op  = InElementOp{};
+    const auto wei_element_op = WeiElementOp{};
+
+    const auto run = [&](auto ndim_spatial, auto in_layout, auto wei_layout, auto out_layout) {
+        constexpr ck::index_t ndim_spatial_value = ndim_spatial.value;
+
+        using InLayout  = decltype(in_layout);
+        using WeiLayout = decltype(wei_layout);
+        using OutLayout = decltype(out_layout);
+
+        const auto in_g_n_c_wis_desc =
+            ck::utils::conv::make_input_host_tensor_descriptor_g_n_c_wis_packed<InLayout>(
+                conv_param);
+
+        const auto wei_g_k_c_xs_desc =
+            ck::utils::conv::make_weight_host_tensor_descriptor_g_k_c_xs_packed<WeiLayout>(
+                conv_param);
+
+        const auto out_g_n_k_wos_desc =
+            ck::utils::conv::make_output_host_tensor_descriptor_g_n_k_wos_packed<OutLayout>(
+                conv_param);
+
+        return run_grouped_conv_fwd<
+            ndim_spatial_value,
+            InDataType,
+            WeiDataType,
+            ConvOutDataType,
+            OutDataType,
+            InElementOp,
+            WeiElementOp,
+            OutElementOp,
+            DeviceGroupedConvNDFwdInstance<ndim_spatial_value, InLayout, WeiLayout, OutLayout>>(
+            do_verification,
+            init_method,
+            time_kernel,
+            conv_param,
+            in_g_n_c_wis_desc,
+            wei_g_k_c_xs_desc,
+            out_g_n_k_wos_desc,
+            in_element_op,
+            wei_element_op);
+    };
+
+    namespace ctc = ck::tensor_layout::convolution;
+
+    if(conv_param.num_dim_spatial_ == 1)
+    {
+        return run(ck::Number<1>{}, ctc::GNWC{}, ctc::GKXC{}, ctc::GNWK{});
+    }
+    else if(conv_param.num_dim_spatial_ == 2)
+    {
+        return run(ck::Number<2>{}, ctc::GNHWC{}, ctc::GKYXC{}, ctc::GNHWK{});
+    }
+    else if(conv_param.num_dim_spatial_ == 3)
+    {
+        return run(ck::Number<3>{}, ctc::GNDHWC{}, ctc::GKZYXC{}, ctc::GNDHWK{});
+    }
+
+    return true;
+}

example/65_gemm_multiply_multiply/gemm_add_add_xdl_fp16.cpp

@@ -208,6 +208,7 @@ int main(int argc, char* argv[])
                                StrideB,
                                std::array<ck::index_t, NumDTensor>{StrideD, StrideD},
                                StrideE,
+                               1,
                                a_element_op,
                                b_element_op,
                                cde_element_op);

example/65_gemm_multiply_multiply/gemm_multiply_multiply_xdl_fp8.cpp

@@ -69,7 +69,7 @@ using AElementOp   = PassThrough;
 using BElementOp   = PassThrough;
 using CDEElementOp = MultiplyMultiply;
 
-static constexpr auto GemmSpec = ck::tensor_operation::device::GemmSpecialization::Default;
+static constexpr auto GemmSpec = ck::tensor_operation::device::GemmSpecialization::MNPadding;
 
 using DeviceOpInstance = ck::tensor_operation::device::DeviceGemmMultiD_Xdl_CShuffle_V3
     // clang-format off
@@ -99,6 +99,8 @@ int main(int argc, char* argv[])
     ck::index_t StrideD = 0;
     ck::index_t StrideE = N;
 
+    ck::index_t KBatch = 1;
+
     if(argc == 1)
     {
         // use default case
@@ -109,7 +111,7 @@ int main(int argc, char* argv[])
         init_method     = std::stoi(argv[2]);
         time_kernel     = std::stoi(argv[3]);
     }
-    else if(argc == 11)
+    else if(argc == 12)
     {
         do_verification = std::stoi(argv[1]);
         init_method     = std::stoi(argv[2]);
@@ -123,13 +125,16 @@ int main(int argc, char* argv[])
         StrideB = std::stoi(argv[8]);
         StrideD = std::stoi(argv[9]);
         StrideE = std::stoi(argv[10]);
+
+        KBatch = std::stoi(argv[11]);
     }
     else
     {
         printf("arg1: verification (0=no, 1=yes)\n");
         printf("arg2: initialization (0=no init, 1=integer value, 2=decimal value)\n");
         printf("arg3: time kernel (0=no, 1=yes)\n");
-        printf("arg4 to 9: M (256x), N(128x), K(32x), StrideA, StrideB, StrideD, StrideE\n");
+        printf(
+            "arg4 to 9: M (256x), N(128x), K(32x), StrideA, StrideB, StrideD, StrideE, KBatch\n");
         exit(0);
     }
 
@@ -212,6 +217,7 @@ int main(int argc, char* argv[])
                                StrideB,
                                std::array<ck::index_t, NumDTensor>{I0, I0},
                                StrideE,
+                               KBatch,
                                a_element_op,
                                b_element_op,
                                cde_element_op);
@@ -236,10 +242,12 @@ int main(int argc, char* argv[])
     std::cout << "Perf: " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec << " GB/s"
               << std::endl;
 
-    e_device_buf.FromDevice(e_m_n_device_result.mData.data());
-
     if(do_verification)
     {
+        invoker.Run(argument, StreamConfig{nullptr, false});
+
+        e_device_buf.FromDevice(e_m_n_device_result.mData.data());
+
         Tensor<CShuffleDataType> c_m_n({M, N});
 
         using ReferenceGemmInstance = ck::tensor_operation::host::ReferenceGemm<A0DataType,

example/CMakeLists.txt

@@ -72,6 +72,20 @@ function(add_example_executable EXAMPLE_NAME FILE_NAME)
             list(REMOVE_ITEM FILE_NAME "${source}")
         endif()
     endforeach()
+    #Do not build any FP8 examples if CK_ENABLE_FP8 not set
+    foreach(source IN LISTS FILE_NAME)
+        if(NOT DEFINED CK_ENABLE_FP8 AND source MATCHES "_fp8")
+            message("removing fp8 example ${source} ")
+            list(REMOVE_ITEM FILE_NAME "${source}")
+        endif()
+    endforeach()
+    #Do not build any BF8 examples if CK_ENABLE_BF8 not set
+    foreach(source IN LISTS FILE_NAME)
+        if(NOT DEFINED CK_ENABLE_BF8 AND source MATCHES "_bf8")
+            message("removing bf8 example ${source} ")
+            list(REMOVE_ITEM FILE_NAME "${source}")
+        endif()
+    endforeach()
     #only continue if there are some source files left on the list
     if(FILE_NAME)
         if(FILE_NAME MATCHES "_xdl")

example/ck_tile/01_fmha/CMakeLists.txt

@@ -6,7 +6,7 @@ execute_process(
 
 execute_process(
   COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_LIST_DIR}/generate.py
-  --api bwd --list_blobs ${CMAKE_CURRENT_BINARY_DIR}/bwd_blob_list.txt
+  --api bwd --list_blobs ${CMAKE_CURRENT_BINARY_DIR}/bwd_blob_list.txt --receipt 3
 )
 
 # NOTE: for cmake, the FMHA_FWD_GEN_BLOBS/FMHA_BWD_GEN_BLOBS files must be in the same directory
@@ -23,7 +23,7 @@ add_custom_command(
 add_custom_command(
   OUTPUT ${FMHA_BWD_GEN_BLOBS}
   COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_LIST_DIR}/generate.py
-  --api bwd --output_dir ${CMAKE_CURRENT_BINARY_DIR}
+  --api bwd --output_dir ${CMAKE_CURRENT_BINARY_DIR} --receipt 3
 )
 
 set(EXAMPLE_FMHA_FWD "tile_example_fmha_fwd")
@@ -55,11 +55,10 @@ set(EXAMPLE_FMHA_BWD_COMPILE_OPTIONS)
 #       ... because they are auto-generated
 if(FMHA_FWD_FAST_EXP2)
 	list(APPEND EXAMPLE_FMHA_FWD_COMPILE_OPTIONS -Wno-undefined-func-template -DCK_TILE_FMHA_FWD_FAST_EXP2=1 -fgpu-flush-denormals-to-zero)
-  list(APPEND EXAMPLE_FMHA_BWD_COMPILE_OPTIONS -Wno-undefined-func-template -DCK_TILE_FMHA_FWD_FAST_EXP2=1 -fgpu-flush-denormals-to-zero)
 else()
 	list(APPEND EXAMPLE_FMHA_FWD_COMPILE_OPTIONS -Wno-undefined-func-template -DCK_TILE_FMHA_FWD_FAST_EXP2=0)
-  list(APPEND EXAMPLE_FMHA_BWD_COMPILE_OPTIONS -Wno-undefined-func-template -DCK_TILE_FMHA_FWD_FAST_EXP2=0)
 endif()
+list(APPEND EXAMPLE_FMHA_BWD_COMPILE_OPTIONS -Wno-undefined-func-template -fgpu-flush-denormals-to-zero)
 
 # Allow comparing floating points directly in order to check sentinel values
 list(APPEND EXAMPLE_FMHA_FWD_COMPILE_OPTIONS -Wno-float-equal)

example/ck_tile/01_fmha/codegen/cpp_symbol_map.py

@@ -66,6 +66,22 @@ BIAS_CHECK_MAP = {
     "alibi" : "bias_enum::alibi"
 }
 
+DROPOUT_MAP = {
+    "no"                        : "ck_tile::BlockDropoutBwd<false, true,  false>",
+    "dropout_wg32"              : "ck_tile::BlockDropoutBwd<true,  true,  false>",
+    "dropout_wg32_storerandval" : "ck_tile::BlockDropoutBwd<true,  true,  true >",
+    "dropout_wg16"              : "ck_tile::BlockDropoutBwd<true,  false, false>",
+    "dropout_wg16_storerandval" : "ck_tile::BlockDropoutBwd<true,  false, true >"
+}
+
+DROPOUT_CHECK_MAP = {
+    "no"                        : "t.has_dropout == false",
+    "dropout_wg32"              : "t.has_dropout == true && t.is_store_randval == false",
+    "dropout_wg32_storerandval" : "t.has_dropout == true && t.is_store_randval == true",
+    "dropout_wg16"              : "t.has_dropout == true && t.is_store_randval == false",
+    "dropout_wg16_storerandval" : "t.has_dropout == true && t.is_store_randval == true",
+}
+
 MODE_MAP = {
     "batch" : "false",
     "group" : "true"

example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py

@@ -14,15 +14,13 @@ from codegen.cpp_symbol_map import *
 
 
 BWD_DQDKDV_PIPELINE_MAP = {
-    "ks_kts_vr" : "ck_tile::BlockFmhaBwdDQDKDVPipelineKSKTSVR",
-    "qs_ks_vr_dos" : "ck_tile::BlockFmhaBwdDQDKDVPipelineQSKSVROGradS",
-    "ks_vr" : "ck_tile::BlockFmhaBwdDQDKDVPipelineKSVR",
+    "kr_ktr_vr_iglp" : "ck_tile::BlockFmhaBwdDQDKDVPipelineKRKTRVRIGLP",
+    "kr_ktr_vr"      : "ck_tile::BlockFmhaBwdDQDKDVPipelineKRKTRVR",
 }
 
 BWD_DQDKDV_PIPELINE_ENUM_MAP = {
-    "ks_kts_vr" : "ck_tile::BlockFmhaBwdPipelineEnum::KSKTSVR",
-    "qs_ks_vr_dos" : "ck_tile::BlockFmhaBwdPipelineEnum::QSKSVROGradS",
-    "ks_vr" : "ck_tile::BlockFmhaBwdPipelineEnum::KSVR",
+    "kr_ktr_vr_iglp" : "ck_tile::BlockFmhaBwdPipelineEnum::KRKTRVR_IGLP",
+    "kr_ktr_vr"      : "ck_tile::BlockFmhaBwdPipelineEnum::KRKTRVR",
 }
 
 FMHA_BWD_KERNEL_HEADER = """// SPDX-License-Identifier: MIT
@@ -34,39 +32,42 @@ FMHA_BWD_KERNEL_HEADER = """// SPDX-License-Identifier: MIT
 FMHA_BWD_DQ_DK_DV_KERNEL_BODY="""
 using fmha_dtype_{F_idx} = {F_dtype};
 
-using fmha_block_tile_{F_idx} = ck_tile::sequence<{F_bm0}, {F_bn0}, {F_bk0}, {F_bk1}, {F_bk2}, {F_bk3}, {F_bk4}, {F_bhdq}, {F_bhdv}>;
+using fmha_block_tile_{F_idx} = ck_tile::
+    sequence<{F_bm0}, {F_bn0}, {F_bk0}, {F_bk1}, {F_bk2}, {F_bk3}, {F_bk4}, {F_bhdq}, {F_bhdv}>;
 using fmha_block_warps0_{F_idx} = ck_tile::sequence<{F_rm0}, {F_rn0}, {F_rk0}>;
 using fmha_block_warps1_{F_idx} = ck_tile::sequence<{F_rm1}, {F_rn1}, {F_rk1}>;
 using fmha_block_warps2_{F_idx} = ck_tile::sequence<{F_rm2}, {F_rn2}, {F_rk2}>;
-using fmha_warp_tile_{F_idx} = ck_tile::sequence<{F_wm}, {F_wn}, {F_wk}>;
+using fmha_warp_tile0_{F_idx}   = ck_tile::sequence<{F_wm0}, {F_wn0}, {F_wk0}>;
+using fmha_warp_tile1_{F_idx}   = ck_tile::sequence<{F_wm1}, {F_wn1}, {F_wk1}>;
 
 // TODO: simplify Gemm0~4BlockWarps in TileFmhaBwdShape
 //       G0&G2 -> GSdP
 //       G1&G3 -> GdKV
 //       G4    -> GdQ
 using fmha_bwd_shape_{F_idx} = ck_tile::TileFmhaBwdShape<fmha_block_tile_{F_idx},
-                                      fmha_block_warps0_{F_idx},
-                                      fmha_warp_tile_{F_idx},
-                                      fmha_block_warps1_{F_idx},
-                                      fmha_warp_tile_{F_idx},
-                                      fmha_block_warps0_{F_idx},
-                                      fmha_warp_tile_{F_idx},
-                                      fmha_block_warps1_{F_idx},
-                                      fmha_warp_tile_{F_idx},
-                                      fmha_block_warps2_{F_idx},
-                                      fmha_warp_tile_{F_idx}>;
+                                                         fmha_block_warps0_{F_idx},
+                                                         fmha_warp_tile0_{F_idx},
+                                                         fmha_block_warps1_{F_idx},
+                                                         fmha_warp_tile1_{F_idx},
+                                                         fmha_block_warps0_{F_idx},
+                                                         fmha_warp_tile0_{F_idx},
+                                                         fmha_block_warps1_{F_idx},
+                                                         fmha_warp_tile1_{F_idx},
+                                                         fmha_block_warps2_{F_idx},
+                                                         fmha_warp_tile0_{F_idx}>;
 
 using fmha_bwd_trait_{F_idx} = ck_tile::TileFmhaTraits<{F_spad},
-                                                    {F_skpad},
-                                                    {F_dpad},
-                                                    {F_dvpad},
-                                                    {F_bias},
-                                                    {F_dbias},
-                                                    false,
-                                                    {F_dropout},
-                                                    false,
-                                                    {F_occupancy}>;
-using fmha_mask_{F_idx} = {F_mask};
+                                                       {F_skpad},
+                                                       {F_dpad},
+                                                       {F_dvpad},
+                                                       {F_bias},
+                                                       {F_dbias},
+                                                       false,
+                                                       false,
+                                                       false,
+                                                       {F_occupancy}>;
+using fmha_mask_{F_idx}      = {F_mask};
+using fmha_dropout_{F_idx}   = {F_dropout};
 
 using fmha_bwd_pipeline_problem_{F_idx} = ck_tile::BlockFmhaBwdPipelineProblem<
     typename FmhaBwdTypeConfig<fmha_dtype_{F_idx}>::QDataType,
@@ -86,55 +87,72 @@ using fmha_bwd_pipeline_problem_{F_idx} = ck_tile::BlockFmhaBwdPipelineProblem<
     typename FmhaBwdTypeConfig<fmha_dtype_{F_idx}>::BiasGradDataType,
     fmha_bwd_shape_{F_idx},
     {F_mode},
+    {F_deterministic},
     fmha_mask_{F_idx},
+    fmha_dropout_{F_idx},
     fmha_bwd_trait_{F_idx}>;
 
-using fmha_bwd_pipeline_{F_idx} = {F_pipeline}<
-    fmha_bwd_pipeline_problem_{F_idx}>;
+using fmha_bwd_pipeline_{F_idx} = {F_pipeline}<fmha_bwd_pipeline_problem_{F_idx}>;
 
-using fmha_bwd_dk_epilogue_{F_idx} =
-    ck_tile::Default2DEpilogue<ck_tile::Default2DEpilogueProblem<typename FmhaBwdTypeConfig<{F_dtype}>::AccDataType,
-                               typename FmhaBwdTypeConfig<{F_dtype}>::KGradDataType,
-                               false, false>>;
+using fmha_bwd_dk_epilogue_{F_idx} = ck_tile::Default2DEpilogue<
+    ck_tile::Default2DEpilogueProblem<typename FmhaBwdTypeConfig<{F_dtype}>::AccDataType,
+                                      typename FmhaBwdTypeConfig<{F_dtype}>::KGradDataType,
+                                      {F_skpad},
+                                      {F_dpad}>>;
 
-using fmha_bwd_dv_epilogue_{F_idx} =
-    ck_tile::Default2DEpilogue<ck_tile::Default2DEpilogueProblem<typename FmhaBwdTypeConfig<{F_dtype}>::AccDataType,
-                               typename FmhaBwdTypeConfig<{F_dtype}>::VGradDataType,
-                               false, false>>;
+using fmha_bwd_dv_epilogue_{F_idx} = ck_tile::Default2DEpilogue<
+    ck_tile::Default2DEpilogueProblem<typename FmhaBwdTypeConfig<{F_dtype}>::AccDataType,
+                                      typename FmhaBwdTypeConfig<{F_dtype}>::VGradDataType,
+                                      {F_skpad},
+                                      {F_dvpad}>>;
 
 using fmha_bwd_dq_dk_dv_kernel_{F_idx} =
-    ck_tile::FmhaBwdDQDKDVKernel<ck_tile::FmhaBwdTilePartitioner<fmha_bwd_shape_{F_idx}>,
-                        fmha_bwd_pipeline_{F_idx},
-                        fmha_bwd_dk_epilogue_{F_idx},
-                        fmha_bwd_dv_epilogue_{F_idx}>;
+    ck_tile::FmhaBwdDQDKDVKernel<fmha_bwd_pipeline_{F_idx},
+                                 fmha_bwd_dk_epilogue_{F_idx},
+                                 fmha_bwd_dv_epilogue_{F_idx}>;
 
-using dq_dk_dv_trait_{F_idx} = fmha_bwd_dq_dk_dv_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_pipeline_enum}, fmha_mask_{F_idx}, {F_bias}, {F_dbias}, {F_dropout}, {F_spad}, {F_skpad}, {F_dpad}, {F_dvpad}>;
+using dq_dk_dv_trait_{F_idx} = fmha_bwd_dq_dk_dv_traits_<{F_hdim},
+                                                         {F_dtype},
+                                                         {F_mode},
+                                                         {F_pipeline_enum},
+                                                         fmha_mask_{F_idx},
+                                                         fmha_dropout_{F_idx},
+                                                         {F_bias},
+                                                         {F_dbias},
+                                                         {F_spad},
+                                                         {F_skpad},
+                                                         {F_dpad},
+                                                         {F_dvpad},
+                                                         {F_deterministic}>;
 
 #include <iostream>
 
-template<>
+template <>
 float fmha_bwd_dq_dk_dv_<dq_dk_dv_trait_{F_idx}>(const ck_tile::stream_config& s, fmha_bwd_args a)
 {{
     using k_ = fmha_bwd_dq_dk_dv_kernel_{F_idx};
     if(s.log_level_ > 0)
         std::cout << ", " << k_::GetName() << std::flush;
-    auto [kargs, grids] = fmha_bwd_dq_dk_dv_create_kargs_and_grids<k_>(a);
-    constexpr dim3 blocks             = k_::BlockSize();
+    auto [kargs, grids]                    = fmha_bwd_dq_dk_dv_create_kargs_and_grids<k_>(a);
+    constexpr dim3 blocks                  = k_::BlockSize();
     constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu;
-    return ck_tile::launch_kernel(s, ck_tile::make_kernel<blocks.x, kBlockPerCu>(k_{{}}, grids, blocks, 0, kargs));
+    return ck_tile::launch_kernel(
+        s, ck_tile::make_kernel<blocks.x, kBlockPerCu>(k_{{}}, grids, blocks, 0, kargs));
 }}
 
-template<>
-void fmha_bwd_dq_dk_dv_oneshot_<dq_dk_dv_trait_{F_idx}>(const ck_tile::stream_config& s, fmha_bwd_args a)
+template <>
+void fmha_bwd_dq_dk_dv_oneshot_<dq_dk_dv_trait_{F_idx}>(const ck_tile::stream_config& s,
+                                                        fmha_bwd_args a)
 {{
-    using k_ = fmha_bwd_dq_dk_dv_kernel_{F_idx};
-    auto [kargs, grids] = fmha_bwd_dq_dk_dv_create_kargs_and_grids<k_>(a);
-    constexpr dim3 blocks             = k_::BlockSize();
+    using k_                               = fmha_bwd_dq_dk_dv_kernel_{F_idx};
+    auto [kargs, grids]                    = fmha_bwd_dq_dk_dv_create_kargs_and_grids<k_>(a);
+    constexpr dim3 blocks                  = k_::BlockSize();
     constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu;
-    ck_tile::make_kernel<blocks.x, kBlockPerCu>(k_{{}}, grids, blocks, 0, kargs)(ck_tile::stream_config{{s.stream_id_}});
+    ck_tile::make_kernel<blocks.x, kBlockPerCu>(k_{{}}, grids, blocks, 0, kargs)(
+        ck_tile::stream_config{{s.stream_id_}});
 }}
 
-template<>
+template <>
 std::string fmha_bwd_dq_dk_dv_get_name_<dq_dk_dv_trait_{F_idx}>()
 {{
     using k_ = fmha_bwd_dq_dk_dv_kernel_{F_idx};
@@ -146,14 +164,15 @@ FMHA_BWD_API_FILENAME="fmha_bwd_api.cpp"
 FMHA_BWD_API="""
 #include <iostream>
 
-template<typename dot_do_o_trait_, typename dq_dk_dv_trait_>
+template <typename dot_do_o_trait_, typename dq_dk_dv_trait_, typename convert_dq_trait_>
 float fmha_bwd_(const ck_tile::stream_config& s, fmha_bwd_args a)
 {{
     if(s.log_level_ > 0)
-        std::cout << ", " << fmha_bwd_dot_do_o_get_name_<dot_do_o_trait_>() << ", " << fmha_bwd_dq_dk_dv_get_name_<dq_dk_dv_trait_>() << std::flush;
+        std::cout << ", " << fmha_bwd_dot_do_o_get_name_<dot_do_o_trait_>() << ", " << fmha_bwd_dq_dk_dv_get_name_<dq_dk_dv_trait_>() << ", " << fmha_bwd_convert_dq_get_name_<convert_dq_trait_>() << std::flush;
     return ck_tile::launch_kernel(s,
-            [=](const ck_tile::stream_config& s_){{ fmha_bwd_dot_do_o_oneshot_<dot_do_o_trait_>(s_, a); }},
-            [=](const ck_tile::stream_config& s_){{ fmha_bwd_dq_dk_dv_oneshot_<dq_dk_dv_trait_>(s_, a); }}
+        [=](const ck_tile::stream_config& s_){{ fmha_bwd_dot_do_o_oneshot_<dot_do_o_trait_>(s_, a); }},
+        [=](const ck_tile::stream_config& s_){{ fmha_bwd_dq_dk_dv_oneshot_<dq_dk_dv_trait_>(s_, a); }},
+        [=](const ck_tile::stream_config& s_){{ fmha_bwd_convert_dq_oneshot_<convert_dq_trait_>(s_, a); }}
     );
 }}
 
@@ -173,38 +192,36 @@ FMHA_BWD_API_PER_HDIM_CASE="""        {F_if} (t.hdim_q <= {F_hdim} && t.hdim_v <
         }}
 """
 
-FMHA_BWD_API_INNER_DISPATCH="""            {F_if}((t.is_group_mode == {F_mode}) && ({F_mask_check}) && (t.bias_type == {F_bias_check}) && (t.has_dbias == {F_dbias}) && (t.has_dropout == {F_dropout}) &&
-                        ({F_scheck}) && ({F_skcheck}) && ({F_dcheck}) && ({F_dvcheck})) {{
-                using dq_dk_dv_trait_ = fmha_bwd_dq_dk_dv_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_pipeline_enum}, {F_mask}, {F_bias}, {F_dbias}, {F_dropout}, {F_spad0}, {F_skpad}, {F_dpad}, {F_dvpad}>;
+FMHA_BWD_API_INNER_DISPATCH="""            {F_if}((t.is_group_mode == {F_mode}) && ({F_mask_check}) && (t.bias_type == {F_bias_check}) && (t.has_dbias == {F_dbias}) && ({F_dropout_check}) &&
+                        ({F_scheck}) && ({F_skcheck}) && ({F_dcheck}) && ({F_dvcheck}) && (t.is_deterministic == {F_deterministic})) {{
                 using dot_do_o_trait_ = fmha_bwd_dot_do_o_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_spad1}, {F_dvpad}>;
-                r = fmha_bwd_<dot_do_o_trait_, dq_dk_dv_trait_>(s, a);
+                using dq_dk_dv_trait_ = fmha_bwd_dq_dk_dv_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_pipeline_enum}, {F_mask}, {F_dropout}, {F_bias}, {F_dbias}, {F_spad0}, {F_skpad}, {F_dpad}, {F_dvpad}, {F_deterministic}>;
+                using convert_dq_trait_ = fmha_bwd_convert_dq_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_spad1}, {F_dpad}, {F_deterministic}>;
+                r = fmha_bwd_<dot_do_o_trait_, dq_dk_dv_trait_, convert_dq_trait_>(s, a);
                 return r;
             }}
 """
 
 @dataclass
 class FmhaBwdDQDKDVApiTrait:
-    pipeline  : str
+    pipeline      : str
     # sync with fmha_bwd_traits<>, to generate fallback calls
-    hdim      : str
-    dtype     : str  # data type
-    mode      : str  # value from MODE_MAP
-    bm0       : int  # tile size along q seqlen (block size)
-    bn0       : int  # tile size along k seqlen
-    bhdq      : int  # q head_dim
-    bhdv      : int  # v head_dim
-    mask      : str
-    bias      : str
-    dbias     : str
-    dropout   : str
-    spad      : str
-    skpad     : str
-    dpad      : str
-    dvpad     : str
-
-    @property
-    def name(self) -> str:
-        return f'{self.pipeline}-{self.hdim}-{self.dtype}-{self.mode}-{self.mask}-{self.bias}-{self.dbias}-{self.dropout}-{self.spad}-{self.skpad}-{self.dpad}-{self.dvpad}'
+    hdim          : str
+    dtype         : str  # data type
+    mode          : str  # value from MODE_MAP
+    bm0           : int  # tile size along q seqlen (block size)
+    bn0           : int  # tile size along k seqlen
+    bhdq          : int  # q head_dim
+    bhdv          : int  # v head_dim
+    mask          : str
+    bias          : str
+    dbias         : str
+    dropout       : str
+    spad          : str
+    skpad         : str
+    dpad          : str
+    dvpad         : str
+    deterministic : str
 
     def scheck(self, spad1 : str) -> str:
         if self.mode == 'group':
@@ -212,9 +229,9 @@ class FmhaBwdDQDKDVApiTrait:
         elif self.spad == 't' and spad1 == 't':
             return f'a.seqlen_q % {self.bm0} != 0'
         elif self.spad == 'f' and spad1 == 't':
-            return f'a.seqlen_q % {self.bm0} == 0 and a.seqlen_q % 256 != 0' # BlockSize
+            return f'a.seqlen_q % {self.bm0} == 0 and a.seqlen_q % 64 != 0'
         else: # self.skpad == 'f' and skpad1 == 'f'
-            return f'a.seqlen_q % 256 == 0' # BlockSize
+            return f'a.seqlen_q % 64 == 0'
 
     @property
     def skcheck(self) -> str:
@@ -256,16 +273,19 @@ class FmhaBwdApiPool:
             per_hdim_case=str()
             for j, hdim in enumerate(self.dq_dk_dv_pool[dtype].keys()):
                 traits=self.dq_dk_dv_pool[dtype][hdim]
+                hdim_int = int(hdim)
                 inners=str()
                 for k, trait in enumerate(traits):
                     if_k = 'if' if k == 0 else 'else if'
                     for spad1 in ["t", "f"]:
-                        if ((spad1 == "f" and trait.spad == "t") or (trait.mode == "group" and spad1 == "f")):
+                        if (spad1 == "f" and (trait.spad == "t" or trait.mode == "group")):
                             continue
-                        inners = inners + FMHA_BWD_API_INNER_DISPATCH.format(F_if=if_k, F_mode=MODE_MAP[trait.mode], F_mask=get_mask_map(self.mask_impl)[trait.mask], F_pipeline_enum=BWD_DQDKDV_PIPELINE_ENUM_MAP[trait.pipeline],
-                                    F_mask_check=get_mask_check_map(self.mask_impl)[trait.mask], F_bias_check=BIAS_CHECK_MAP[trait.bias], F_bias=BIAS_MAP[trait.bias], F_dbias=BOOL_MAP[trait.dbias], F_dropout=BOOL_MAP[trait.dropout],
+                        inners = inners + FMHA_BWD_API_INNER_DISPATCH.format(F_if=if_k, F_mode=MODE_MAP[trait.mode], F_pipeline_enum=BWD_DQDKDV_PIPELINE_ENUM_MAP[trait.pipeline],
+                                    F_mask_check=get_mask_check_map(self.mask_impl)[trait.mask], F_mask=get_mask_map(self.mask_impl)[trait.mask], F_bias_check=BIAS_CHECK_MAP[trait.bias],
+                                    F_bias=BIAS_MAP[trait.bias], F_dbias=BOOL_MAP[trait.dbias], F_dropout_check=DROPOUT_CHECK_MAP[trait.dropout], F_dropout=DROPOUT_MAP[trait.dropout],
                                     F_scheck=trait.scheck(spad1=spad1), F_skcheck=trait.skcheck, F_dcheck=trait.dcheck, F_dvcheck=trait.dvcheck, F_hdim=hdim, F_dtype=DTYPE_MAP[dtype],
-                                    F_spad0=BOOL_MAP[trait.spad], F_spad1=BOOL_MAP[spad1], F_skpad=BOOL_MAP[trait.skpad], F_dpad=BOOL_MAP[trait.dpad], F_dvpad=BOOL_MAP[trait.dvpad])
+                                    F_spad0=BOOL_MAP[trait.spad], F_spad1=BOOL_MAP[spad1], F_skpad=BOOL_MAP[trait.skpad], F_dpad=BOOL_MAP[trait.dpad], F_dvpad=BOOL_MAP[trait.dvpad],
+                                    F_deterministic=BOOL_MAP[trait.deterministic])
 
                 if_j = 'if' if j == 0 else 'else if'
                 per_hdim_case = per_hdim_case + FMHA_BWD_API_PER_HDIM_CASE.format(F_if=if_j, F_hdim=hdim, F_inner_dispatch=inners)
@@ -295,81 +315,89 @@ class FmhaBwdDQDKDVTileSize:
     F_bhdv      : int  # v head_dim
     F_rm0       : int  # number of warps along q seqlen (block warps) in gemm0/gemm2
     F_rn0       : int  # number of warps along k seqlen (block warps) in gemm0/gemm2
-    F_rk0       : int  # number of warps along gemm-k (not used) in gemm0/gemm2
+    F_rk0       : int  # number of warps along headdim_qk/v (not used) in gemm0/gemm2
     F_rm1       : int  # number of warps along k seqlen (block warps) in gemm1/gemm3
-    F_rn1       : int  # number of warps along q seqlen (block warps) in gemm1/gemm3
-    F_rk1       : int  # number of warps along gemm-k (not used) in gemm1/gemm3
-    F_rm2       : int  # number of warps along k seqlen (block warps) in gemm4
-    F_rn2       : int  # number of warps along q seqlen (block warps) in gemm4
-    F_rk2       : int  # number of warps along gemm-k (not used) in gemm4
-    F_wm        : int  # warp size along m (warp size)
-    F_wn        : int  # warp size along n
-    F_wk        : int  # warp size along k
+    F_rn1       : int  # number of warps along headdim_qk/v (block warps) in gemm1/gemm3
+    F_rk1       : int  # number of warps along q seqlen (not used) in gemm1/gemm3
+    F_rm2       : int  # number of warps along q seqlen (block warps) in gemm4
+    F_rn2       : int  # number of warps along headdim_qk (block warps) in gemm4
+    F_rk2       : int  # number of warps along k seqlen (not used) in gemm4
+    F_wm0       : int  # warp size along m in gemm0/gemm2/gemm4
+    F_wn0       : int  # warp size along n in gemm0/gemm2/gemm4
+    F_wk0       : int  # warp size along k in gemm0/gemm2/gemm4
+    F_wm1       : int  # warp size along m in gemm1/gemm3
+    F_wn1       : int  # warp size along n in gemm1/gemm3
+    F_wk1       : int  # warp size along k in gemm1/gemm3
     F_occupancy : int  # occupancy
     @property
     def name(self) -> str:
         return f"b{self.F_bm0}x{self.F_bn0}x{self.F_bk0}x{self.F_bk1}x{self.F_bk2}x{self.F_bk3}x{self.F_bk4}x{self.F_bhdq}x{self.F_bhdv}" +\
         f"_r{self.F_rm0}x{self.F_rn0}x{self.F_rk0}_r{self.F_rm1}x{self.F_rn1}x{self.F_rk1}_r{self.F_rm2}x{self.F_rn2}x{self.F_rk2}" +\
-        f"_w{self.F_wm}x{self.F_wn}x{self.F_wk}_o{self.F_occupancy}"
+        f"_w{self.F_wm0}x{self.F_wn0}x{self.F_wk0}_w{self.F_wm1}x{self.F_wn1}x{self.F_wk1}_o{self.F_occupancy}"
 
 @dataclass
 class FmhaBwdDQDKDVKernel:
-    F_idx       : int  # this is not a tunable, but a counter to differentiate symbol
-    F_hdim      : int  # hdim
-    F_dtype     : str  # data type
-    F_tile      : FmhaBwdDQDKDVTileSize
-    F_spad      : str  # true/false
-    F_skpad     : str  #
-    F_dpad      : str  #
-    F_dvpad     : str  #
-    F_bias      : str  #
-    F_dbias     : str  #
-    F_dropout   : str  #
-    F_mask      : str  # value from MASK_MAP
-    F_mode      : str  # value from MODE_MAP
-    F_pipeline  : str
-    mask_impl   : str
+    F_idx           : int  # this is not a tunable, but a counter to differentiate symbol
+    F_hdim          : int  # hdim
+    F_dtype         : str  # data type
+    F_tile          : FmhaBwdDQDKDVTileSize
+    F_spad          : str  # true/false
+    F_skpad         : str  #
+    F_dpad          : str  #
+    F_dvpad         : str  #
+    F_bias          : str  #
+    F_dbias         : str  #
+    F_dropout       : str  #
+    F_mask          : str  # value from MASK_MAP
+    F_mode          : str  # value from MODE_MAP
+    F_deterministic : str  #
+    F_pipeline      : str  #
+    mask_impl       : str  #
 
     @property
     def template(self) -> str:
         return FMHA_BWD_KERNEL_HEADER + \
             FMHA_BWD_DQ_DK_DV_KERNEL_BODY.format(
-                F_idx       = self.F_idx,
-                F_hdim      = self.F_hdim,
-                F_dtype     = DTYPE_MAP[self.F_dtype],
-                F_bm0       = self.F_tile.F_bm0,
-                F_bn0       = self.F_tile.F_bn0,
-                F_bk0       = self.F_tile.F_bk0,
-                F_bk1       = self.F_tile.F_bk1,
-                F_bk2       = self.F_tile.F_bk2,
-                F_bk3       = self.F_tile.F_bk3,
-                F_bk4       = self.F_tile.F_bk4,
-                F_bhdq      = self.F_tile.F_bhdq,
-                F_bhdv      = self.F_tile.F_bhdv,
-                F_rm0       = self.F_tile.F_rm0,
-                F_rn0       = self.F_tile.F_rn0,
-                F_rk0       = self.F_tile.F_rk0,
-                F_rm1       = self.F_tile.F_rm1,
-                F_rn1       = self.F_tile.F_rn1,
-                F_rk1       = self.F_tile.F_rk1,
-                F_rm2       = self.F_tile.F_rm2,
-                F_rn2       = self.F_tile.F_rn2,
-                F_rk2       = self.F_tile.F_rk2,
-                F_wm        = self.F_tile.F_wm,
-                F_wn        = self.F_tile.F_wn,
-                F_wk        = self.F_tile.F_wk,
-                F_spad      = BOOL_MAP[self.F_spad],
-                F_skpad     = BOOL_MAP[self.F_skpad],
-                F_dpad      = BOOL_MAP[self.F_dpad],
-                F_dvpad     = BOOL_MAP[self.F_dvpad],
-                F_bias      = BIAS_MAP[self.F_bias],
-                F_dbias     = BOOL_MAP[self.F_dbias],
-                F_dropout   = BOOL_MAP[self.F_dropout],
-                F_occupancy = self.F_tile.F_occupancy,
-                F_mask      = get_mask_map(self.mask_impl)[self.F_mask],
-                F_mode      = MODE_MAP[self.F_mode],
+                F_idx           = self.F_idx,
+                F_hdim          = self.F_hdim,
+                F_dtype         = DTYPE_MAP[self.F_dtype],
+                F_bm0           = self.F_tile.F_bm0,
+                F_bn0           = self.F_tile.F_bn0,
+                F_bk0           = self.F_tile.F_bk0,
+                F_bk1           = self.F_tile.F_bk1,
+                F_bk2           = self.F_tile.F_bk2,
+                F_bk3           = self.F_tile.F_bk3,
+                F_bk4           = self.F_tile.F_bk4,
+                F_bhdq          = self.F_tile.F_bhdq,
+                F_bhdv          = self.F_tile.F_bhdv,
+                F_rm0           = self.F_tile.F_rm0,
+                F_rn0           = self.F_tile.F_rn0,
+                F_rk0           = self.F_tile.F_rk0,
+                F_rm1           = self.F_tile.F_rm1,
+                F_rn1           = self.F_tile.F_rn1,
+                F_rk1           = self.F_tile.F_rk1,
+                F_rm2           = self.F_tile.F_rm2,
+                F_rn2           = self.F_tile.F_rn2,
+                F_rk2           = self.F_tile.F_rk2,
+                F_wm0           = self.F_tile.F_wm0,
+                F_wn0           = self.F_tile.F_wn0,
+                F_wk0           = self.F_tile.F_wk0,
+                F_wm1           = self.F_tile.F_wm1,
+                F_wn1           = self.F_tile.F_wn1,
+                F_wk1           = self.F_tile.F_wk1,
+                F_spad          = BOOL_MAP[self.F_spad],
+                F_skpad         = BOOL_MAP[self.F_skpad],
+                F_dpad          = BOOL_MAP[self.F_dpad],
+                F_dvpad         = BOOL_MAP[self.F_dvpad],
+                F_bias          = BIAS_MAP[self.F_bias],
+                F_dbias         = BOOL_MAP[self.F_dbias],
+                F_dropout       = DROPOUT_MAP[self.F_dropout],
+                F_occupancy     = self.F_tile.F_occupancy,
+                F_mask          = get_mask_map(self.mask_impl)[self.F_mask],
+                F_mode          = MODE_MAP[self.F_mode],
+                F_deterministic = BOOL_MAP[self.F_deterministic],
                 F_pipeline_enum = BWD_DQDKDV_PIPELINE_ENUM_MAP[self.F_pipeline],
-                F_pipeline  = BWD_DQDKDV_PIPELINE_MAP[self.F_pipeline])
+                F_pipeline      = BWD_DQDKDV_PIPELINE_MAP[self.F_pipeline])
 
     @property
     def name(self) -> str:
@@ -382,7 +410,7 @@ class FmhaBwdDQDKDVKernel:
             if n != '' : n = 'p' + n
             return n
         pn = pad_name()
-        n = f"fmha_bwd_d{self.F_hdim}_{self.F_dtype}_{self.F_mode}_" + self.F_tile.name
+        n = f"fmha_bwd_d{self.F_hdim}_{self.F_dtype}_{self.F_mode}_" + self.F_tile.name + f'_{self.F_pipeline}'
         if pn != '' : n += f'_{pn}'
         if self.F_bias != 'no' : n += f'_{self.F_bias}'
         if self.F_dbias == 't' : n += '_dbias'
@@ -390,7 +418,8 @@ class FmhaBwdDQDKDVKernel:
             if self.F_mask == 's_mask': n += f'_mask'
         else:
             if self.F_mask != 'no' : n += f'_m{self.F_mask[0]}'
-        if self.F_dropout == 't' : n += '_dropout'
+        if self.F_dropout != 'no' : n += f'_{self.F_dropout}'
+        if self.F_deterministic == 't' : n += '_deterministic'
         return n
 
     @property
@@ -413,19 +442,23 @@ class FmhaBwdDQDKDVKernel:
                 spad=self.F_spad,
                 skpad=self.F_skpad,
                 dpad=self.F_dpad,
-                dvpad=self.F_dvpad)
+                dvpad=self.F_dvpad,
+                deterministic=self.F_deterministic
+                )
 
 # TODO: design a more practical way to do it
 # this is current supported tile size & pipeline.
 def get_fmha_bwd_dq_dk_dv_tile_ppl_dict_from_dtype(dtype : str) -> Optional[dict]:
     if dtype == 'fp16' or dtype == 'bf16':
         return {
-            '32'  : [FmhaBwdDQDKDVTileSize(128, 128, 32, 32, 32, 32, 32,  32,  32, 1, 4, 1, 4, 1, 1, 4, 1, 1, 32, 32, 16, 1),
-                        "qs_ks_vr_dos"],
-            '64'  : [FmhaBwdDQDKDVTileSize( 64, 128, 32, 32, 32, 32, 32,  64,  64, 1, 4, 1, 4, 1, 1, 2, 2, 1, 32, 32, 16, 1),
-                        "qs_ks_vr_dos"],
-            '128' : [FmhaBwdDQDKDVTileSize( 64, 128, 32, 32, 32, 32, 32, 128, 128, 1, 4, 1, 4, 1, 1, 2, 2, 1, 32, 32, 16, 1),
-                        "ks_vr"]
+            '32'  : [FmhaBwdDQDKDVTileSize( 32, 128,  32, 32,  32, 32, 64,  32,  32, 1, 4, 1, 4, 1, 1, 2, 2, 1, 16, 16, 32, 16, 16, 16, 1),
+                        "kr_ktr_vr_iglp", "kr_ktr_vr"],
+            '64'  : [FmhaBwdDQDKDVTileSize( 32, 128,  64, 32,  64, 32, 32,  64,  64, 1, 4, 1, 4, 1, 1, 1, 4, 1, 16, 16, 32, 16, 16, 16, 1),
+                        "kr_ktr_vr_iglp", "kr_ktr_vr"],
+            '128' : [FmhaBwdDQDKDVTileSize( 16, 128, 128, 16, 128, 16, 32, 128, 128, 1, 4, 1, 4, 1, 1, 1, 4, 1, 16, 16, 32, 16, 16, 16, 1),
+                        "kr_ktr_vr_iglp", "kr_ktr_vr"],
+            '256' : [FmhaBwdDQDKDVTileSize( 16,  64, 256, 16, 256, 16, 32, 256, 256, 1, 4, 1, 4, 1, 1, 1, 4, 1, 16, 16, 32, 16, 16, 16, 1),
+                        "kr_ktr_vr_iglp", "kr_ktr_vr"]
         }
     else:
         return None
@@ -440,7 +473,7 @@ def get_bwd_dq_dk_dv_blobs(kernel_filter : Optional[str], receipt, mask_impl) ->
         d = get_fmha_bwd_dq_dk_dv_tile_ppl_dict_from_dtype(dtype)
         if d == None:
             continue
-        for hdim_str, mode, mask, bias, dbias, dropout, spad, skpad, dpad, dvpad in itertools.product(d.keys(), MODE_MAP.keys(), get_mask_map(mask_impl).keys(), BIAS_MAP.keys(), ["t", "f"], ["t", "f"], ["t", "f"], ["t", "f"], ["t", "f"], ["t", "f"]):
+        for hdim_str, mode, mask, bias, dbias, dropout, spad, skpad, dpad, dvpad, deterministic in itertools.product(d.keys(), MODE_MAP.keys(), get_mask_map(mask_impl).keys(), BIAS_MAP.keys(), ["t", "f"], DROPOUT_MAP.keys(), ["t", "f"], ["t", "f"], ["t", "f"], ["t", "f"], ["t", "f"]):
             tile = d[hdim_str][0]
             ppl = d[hdim_str][1]
             hdim = int(hdim_str)
@@ -448,16 +481,29 @@ def get_bwd_dq_dk_dv_blobs(kernel_filter : Optional[str], receipt, mask_impl) ->
                 continue
             if ((bias == "no" or bias == "alibi") and dbias == "t"):
                 continue
+            if ("wg32" in dropout):
+                continue
+            if (dpad == "t" or dvpad == "t"):
+                ppl = d[hdim_str][2]
             k = FmhaBwdDQDKDVKernel(F_idx=0, F_hdim=hdim, F_dtype=dtype, F_tile=tile,
                                 F_spad=spad, F_skpad=skpad, F_dpad=dpad, F_dvpad=dvpad,
                                 F_bias=bias, F_dbias=dbias, F_dropout=dropout, F_mask=mask, F_mode=mode,
-                                F_pipeline=ppl, mask_impl=mask_impl)
+                                F_pipeline=ppl, mask_impl=mask_impl, F_deterministic=deterministic)
             if kernel_filter != None:
                 if not fnmatch.fnmatch(k.name, kernel_filter):
                     continue
             if receipt == 2:
                     cond = dtype in ['fp16', 'bf16']
                     cond &= bias in ['no', 'alibi']
+                    cond &= dropout in ['no', 'dropout_wg32',  'dropout_wg16']
+                    cond &= dpad == dvpad
+                    if not cond:
+                        continue
+            if receipt == 3:
+                    cond = dtype in ['fp16', 'bf16']
+                    cond &= bias in ['no', 'alibi']
+                    cond &= dpad == dvpad
+                    cond &= deterministic == "f"
                     if not cond:
                         continue
             api_pool.register_dq_dk_dv_traits(k.api_trait())
@@ -468,53 +514,54 @@ def get_bwd_dq_dk_dv_blobs(kernel_filter : Optional[str], receipt, mask_impl) ->
 FMHA_BWD_DOT_DO_O_KERNEL_BODY="""
 using fmha_dtype_{F_idx} = {F_dtype};
 
-using fmha_bwd_dot_do_o_trait_{F_idx} = ck_tile::TileFmhaBwdOGradDotOTraits<{F_spad},
-                                                    {F_dvpad},
-                                                    {F_occupancy}>;
+using fmha_bwd_dot_do_o_trait_{F_idx} =
+    ck_tile::TileFmhaBwdOGradDotOTraits<{F_spad}, {F_dvpad}, {F_occupancy}>;
 
 using fmha_bwd_dot_do_o_pipeline_problem_{F_idx} = ck_tile::BlockFmhaBwdOGradDotOPipelineProblem<
     typename FmhaBwdTypeConfig<fmha_dtype_{F_idx}>::ODataType,
     typename FmhaBwdTypeConfig<fmha_dtype_{F_idx}>::OGradDataType,
     typename FmhaBwdTypeConfig<fmha_dtype_{F_idx}>::DDataType,
-    /* BlockSize = */ 256,
+    /* BlockSize = */ 64,
     {F_hdim},
     {F_mode},
     fmha_bwd_dot_do_o_trait_{F_idx}>;
 
-using fmha_bwd_dot_do_o_{F_idx} = typename ck_tile::BlockFmhaBwdOGradDotO<
-    fmha_bwd_dot_do_o_pipeline_problem_{F_idx}>;
+using fmha_bwd_dot_do_o_{F_idx} =
+    typename ck_tile::BlockFmhaBwdOGradDotO<fmha_bwd_dot_do_o_pipeline_problem_{F_idx}>;
 
 using fmha_bwd_dot_do_o_kernel_{F_idx} =
-    ck_tile::FmhaBwdOGradDotOKernel<ck_tile::FmhaBwdOGradDotOTilePartitioner</* BlockSize = */ 256>,
-                                    fmha_bwd_dot_do_o_{F_idx}>;
+    ck_tile::FmhaBwdOGradDotOKernel<fmha_bwd_dot_do_o_{F_idx}>;
 
-using dot_do_o_trait_{F_idx} = fmha_bwd_dot_do_o_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_spad}, {F_dvpad}>;
+using dot_do_o_trait_{F_idx} =
+    fmha_bwd_dot_do_o_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_spad}, {F_dvpad}>;
 
 #include <iostream>
 
-template<>
+template <>
 float fmha_bwd_dot_do_o_<dot_do_o_trait_{F_idx}>(const ck_tile::stream_config& s, fmha_bwd_args a)
 {{
     using k_ = fmha_bwd_dot_do_o_kernel_{F_idx};
     if(s.log_level_ > 0)
         std::cout << ", " << k_::GetName() << std::flush;
-    auto [kargs, grids] = fmha_bwd_dot_do_o_create_kargs_and_grids<k_>(a);
-    constexpr dim3 blocks             = k_::BlockSize();
+    auto [kargs, grids]                    = fmha_bwd_dot_do_o_create_kargs_and_grids<k_>(a);
+    constexpr dim3 blocks                  = k_::BlockSize();
     constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu;
-    return ck_tile::launch_kernel(s, ck_tile::make_kernel<blocks.x, kBlockPerCu>(k_{{}}, grids, blocks, 0, kargs));
+    return ck_tile::launch_kernel(
+        s, ck_tile::make_kernel<blocks.x, kBlockPerCu>(k_{{}}, grids, blocks, 0, kargs));
 }}
 
-template<>
+template <>
 void fmha_bwd_dot_do_o_oneshot_<dot_do_o_trait_{F_idx}>(const ck_tile::stream_config& s, fmha_bwd_args a)
 {{
-    using k_ = fmha_bwd_dot_do_o_kernel_{F_idx};
-    auto [kargs, grids] = fmha_bwd_dot_do_o_create_kargs_and_grids<k_>(a);
-    constexpr dim3 blocks             = k_::BlockSize();
+    using k_                               = fmha_bwd_dot_do_o_kernel_{F_idx};
+    auto [kargs, grids]                    = fmha_bwd_dot_do_o_create_kargs_and_grids<k_>(a);
+    constexpr dim3 blocks                  = k_::BlockSize();
     constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu;
-    ck_tile::make_kernel<blocks.x, kBlockPerCu>(k_{{}}, grids, blocks, 0, kargs)(ck_tile::stream_config{{s.stream_id_}});
+    ck_tile::make_kernel<blocks.x, kBlockPerCu>(k_{{}}, grids, blocks, 0, kargs)(
+        ck_tile::stream_config{{s.stream_id_}});
 }}
 
-template<>
+template <>
 std::string fmha_bwd_dot_do_o_get_name_<dot_do_o_trait_{F_idx}>()
 {{
     using k_ = fmha_bwd_dot_do_o_kernel_{F_idx};
@@ -584,12 +631,150 @@ def get_bwd_dot_do_o_blobs() -> List[FmhaBwdOGradDotOKernel]:
 
     return gen
 
+FMHA_BWD_CONVERT_DQ_KERNEL_BODY="""
+using fmha_dtype_{F_idx} = {F_dtype};
+
+using fmha_bwd_convert_dq_trait_{F_idx} =
+    ck_tile::TileFmhaBwdConvertQGradTraits<{F_spad}, {F_dpad}, {F_occupancy}>;
+
+using fmha_bwd_convert_dq_pipeline_problem_{F_idx} =
+    ck_tile::BlockFmhaBwdConvertQGradPipelineProblem<
+        typename FmhaBwdTypeConfig<fmha_dtype_{F_idx}>::AccDataType,
+        typename FmhaBwdTypeConfig<fmha_dtype_{F_idx}>::QGradDataType,
+        /* BlockSize = */ 256,
+        {F_bm0},
+        {F_bn0},
+        {F_hdim},
+        {F_mode},
+        {F_deterministic},
+        fmha_bwd_convert_dq_trait_{F_idx}>;
+
+using fmha_bwd_convert_dq_{F_idx} =
+    typename ck_tile::BlockFmhaBwdConvertQGrad<fmha_bwd_convert_dq_pipeline_problem_{F_idx}>;
+
+using fmha_bwd_convert_dq_kernel_{F_idx} =
+    ck_tile::FmhaBwdConvertQGradKernel<fmha_bwd_convert_dq_{F_idx}>;
+
+using convert_dq_trait_{F_idx} = fmha_bwd_convert_dq_traits_<{F_hdim},
+                                                             {F_dtype},
+                                                             {F_mode},
+                                                             {F_spad},
+                                                             {F_dpad},
+                                                             {F_deterministic}>;
+
+#include <iostream>
+
+template <>
+float fmha_bwd_convert_dq_<convert_dq_trait_{F_idx}>(const ck_tile::stream_config& s, fmha_bwd_args a)
+{{
+    using k_ = fmha_bwd_convert_dq_kernel_{F_idx};
+    if(s.log_level_ > 0)
+        std::cout << ", " << k_::GetName() << std::flush;
+    auto [kargs, grids]                    = fmha_bwd_convert_dq_create_kargs_and_grids<k_>(a);
+    constexpr dim3 blocks                  = k_::BlockSize();
+    constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu;
+    return ck_tile::launch_kernel(
+        s, ck_tile::make_kernel<blocks.x, kBlockPerCu>(k_{{}}, grids, blocks, 0, kargs));
+}}
+
+template <>
+void fmha_bwd_convert_dq_oneshot_<convert_dq_trait_{F_idx}>(const ck_tile::stream_config& s,
+                                                            fmha_bwd_args a)
+{{
+    using k_                               = fmha_bwd_convert_dq_kernel_{F_idx};
+    auto [kargs, grids]                    = fmha_bwd_convert_dq_create_kargs_and_grids<k_>(a);
+    constexpr dim3 blocks                  = k_::BlockSize();
+    constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu;
+    ck_tile::make_kernel<blocks.x, kBlockPerCu>(k_{{}}, grids, blocks, 0, kargs)(
+        ck_tile::stream_config{{s.stream_id_}});
+}}
+
+template <>
+std::string fmha_bwd_convert_dq_get_name_<convert_dq_trait_{F_idx}>()
+{{
+    using k_ = fmha_bwd_convert_dq_kernel_{F_idx};
+    return k_::GetName();
+}}
+"""
+
+@dataclass
+class FmhaBwdConvertQGradKernel:
+    F_idx           : int  # this is not a tunable, but a counter to differentiate symbol
+    F_hdim          : int  # hdim
+    F_dtype         : str  # data type
+    F_bm0           : int  # tile size along q seqlen (block size)
+    F_bn0           : int  # tile size along k seqlen
+    F_spad          : str  # true/false
+    F_dpad          : str  #
+    F_mode          : str  # value from MODE_MAP
+    F_occupancy     : int  #
+    F_deterministic : str  #
+
+    @property
+    def template(self) -> str:
+        return FMHA_BWD_KERNEL_HEADER + \
+            FMHA_BWD_CONVERT_DQ_KERNEL_BODY.format(
+                F_idx           = self.F_idx,
+                F_hdim          = self.F_hdim,
+                F_dtype         = DTYPE_MAP[self.F_dtype],
+                F_bm0           = self.F_bm0,
+                F_bn0           = self.F_bn0,
+                F_spad          = BOOL_MAP[self.F_spad],
+                F_dpad          = BOOL_MAP[self.F_dpad],
+                F_mode          = MODE_MAP[self.F_mode],
+                F_occupancy     = self.F_occupancy,
+                F_deterministic = BOOL_MAP[self.F_deterministic])
+
+    @property
+    def name(self) -> str:
+        def pad_name() -> str:
+            n = ''
+            if self.F_spad == 't': n += 's'
+            if self.F_dpad == 't' : n += 'd'
+            if n != '' : n = 'p' + n
+            return n
+        pn = pad_name()
+        n = f"fmha_bwd_convert_dq_d{self.F_hdim}_{self.F_dtype}_b{self.F_bm0}x{self.F_bn0}_{self.F_mode}_o{self.F_occupancy}"
+        if pn != '' : n += f'_{pn}'
+        if self.F_deterministic == 't' : n += f'_deterministic'
+        return n
+
+    @property
+    def filename(self) -> str:
+        return self.name + ".cpp"
+
+def get_bwd_convert_dq_blobs() -> List[FmhaBwdConvertQGradKernel]:
+    # TODO: we don't support tuning yet, so pick up one value for pad/occupancy
+    #       support this in future
+    def get_occupancy(dtype, hdim):
+        return 2
+
+    gen = list()
+
+    for dtype in DTYPE_MAP.keys():
+        d = get_fmha_bwd_dq_dk_dv_tile_ppl_dict_from_dtype(dtype)
+        if d == None:
+            continue
+        for hdim_str, mode, spad, dpad, deterministic in itertools.product(d.keys(), MODE_MAP.keys(), ["t", "f"], ["t", "f"], ["t", "f"]):
+            hdim = int(hdim_str)
+            tile = d[hdim_str][0]
+            if (mode == "group" and spad == "f"):
+                continue
+            k = FmhaBwdConvertQGradKernel(F_idx=0, F_hdim=hdim, F_dtype=dtype, F_bm0=64, F_bn0=tile.F_bn0,
+                                F_spad=spad, F_dpad=dpad, F_mode=mode, F_occupancy=get_occupancy(dtype, hdim), F_deterministic=deterministic)
+            gen.append(k)
+
+    return gen
+
 def write_single_bwd_dq_dk_dv_kernel(kernel: FmhaBwdDQDKDVKernel, autogen_dir: Path) -> None:
     (autogen_dir / kernel.filename).write_text(kernel.template)
 
 def write_single_bwd_dot_do_o_kernel(kernel: FmhaBwdOGradDotOKernel, autogen_dir: Path) -> None:
     (autogen_dir / kernel.filename).write_text(kernel.template)
 
+def write_single_bwd_convert_dq_kernel(kernel: FmhaBwdConvertQGradKernel, autogen_dir: Path) -> None:
+    (autogen_dir / kernel.filename).write_text(kernel.template)
+
 def write_bwd_api(api_pool : FmhaBwdApiPool, autogen_dir: Path) -> None:
     (autogen_dir / FMHA_BWD_API_FILENAME).write_text(api_pool.api)
 
@@ -597,6 +782,9 @@ def write_blobs(output_dir : Path, kernel_filter : Optional[str], receipt, mask_
     kernels = get_bwd_dot_do_o_blobs()
     for kernel in kernels:
         write_single_bwd_dot_do_o_kernel(kernel, output_dir)
+    kernels = get_bwd_convert_dq_blobs()
+    for kernel in kernels:
+        write_single_bwd_convert_dq_kernel(kernel, output_dir)
     api_pool, kernels = get_bwd_dq_dk_dv_blobs(kernel_filter, receipt, mask_impl)
     for kernel in kernels:
         write_single_bwd_dq_dk_dv_kernel(kernel, output_dir)
@@ -605,6 +793,9 @@ def write_blobs(output_dir : Path, kernel_filter : Optional[str], receipt, mask_
 def list_blobs(file_path : Path, kernel_filter : Optional[str], receipt, mask_impl) -> None:
     with file_path.open('a') as f:
         kernels = get_bwd_dot_do_o_blobs()
+        for kernel in kernels:
+            f.write(str(file_path.parent / GEN_DIR / kernel.filename) + "\n")
+        kernels = get_bwd_convert_dq_blobs()
         for kernel in kernels:
             f.write(str(file_path.parent / GEN_DIR / kernel.filename) + "\n")
         _, kernels = get_bwd_dq_dk_dv_blobs(kernel_filter, receipt, mask_impl)

example/ck_tile/01_fmha/fmha_bwd.cpp

@@ -87,7 +87,11 @@ auto create_args(int argc, char* argv[])
         .insert("drop_offset", "0", "offset for random number generator")
         .insert("timer", "gpu", "gpu:gpu timer, cpu:cpu timer")
         .insert("warmup", "5", "number of iterations before benchmark the kernel")
-        .insert("repeat", "20", "number of iterations to benchmark the kernel");
+        .insert("repeat", "20", "number of iterations to benchmark the kernel")
+        .insert("deterministic",
+                "0",
+                "if set to 1 will use multi-buffer reduction strategy for dq, atomic opeartion "
+                "will not be used");
 
     bool result = arg_parser.parse(argc, argv);
     return std::make_tuple(result, arg_parser);
@@ -128,11 +132,6 @@ bool run(const ck_tile::ArgParser& arg_parser)
     ck_tile::index_t hdim_v = arg_parser.get_int("d_v");
     if(hdim_v < 0)
         hdim_v = hdim_q;
-    if(hdim_q % 2 != 0 || hdim_v % 2 != 0)
-    {
-        std::cerr << "FMHA Bwd kernel currently only supports even headdim" << std::endl;
-        return false;
-    }
 
     bool i_perm = arg_parser.get_bool("iperm"); // if true, will be batch * nhead * seqlen * hdim
     bool o_perm = arg_parser.get_bool("operm"); // if false, will be batch * seqlen * nhead * hdim
@@ -177,9 +176,10 @@ bool run(const ck_tile::ArgParser& arg_parser)
         seed.reset();
     }
 
-    int stream_warmup = arg_parser.get_int("warmup");
-    int stream_repeat = arg_parser.get_int("repeat");
-    bool kname        = arg_parser.get_bool("kname");
+    int stream_warmup  = arg_parser.get_int("warmup");
+    int stream_repeat  = arg_parser.get_int("repeat");
+    bool kname         = arg_parser.get_bool("kname");
+    bool deterministic = arg_parser.get_bool("deterministic");
 
     ck_tile::stream_config stream_config{nullptr,
                                          true,
@@ -265,6 +265,9 @@ bool run(const ck_tile::ArgParser& arg_parser)
         (mode == mode_enum::batch ? seqlen_q : seqstart_q_host.back());
     const ck_tile::index_t shape_seqlen_k =
         (mode == mode_enum::batch ? seqlen_k : seqstart_k_host.back());
+    const ck_tile::index_t kN0 = (hdim_q <= 128) ? 128 : 64;
+    const ck_tile::index_t nsplits =
+        deterministic ? ck_tile::integer_divide_ceil(max_seqlen_k, kN0) : 1;
 
     ck_tile::HostTensor<QDataType> q_host(
         get_lengths(i_perm, shape_batch, nhead, shape_seqlen_q, hdim_q));
@@ -284,9 +287,9 @@ bool run(const ck_tile::ArgParser& arg_parser)
     ck_tile::HostTensor<ODataType> o_host(
         get_lengths(o_perm, shape_batch, nhead, shape_seqlen_q, hdim_v));
     ck_tile::HostTensor<LSEDataType> lse_host(
-        std::array<ck_tile::index_t, 3>{batch, nhead, max_seqlen_q});
+        std::array<ck_tile::index_t, 3>{shape_batch, nhead, shape_seqlen_q});
     ck_tile::HostTensor<DDataType> d_host(
-        std::array<ck_tile::index_t, 3>{batch, nhead, max_seqlen_q});
+        std::array<ck_tile::index_t, 3>{shape_batch, nhead, shape_seqlen_q});
     ck_tile::HostTensor<RandValOutputDataType> randval_host(
         p_drop > 0 ? get_lengths(true, shape_batch, nhead, shape_seqlen_q, max_seqlen_k)
                    : std::array<ck_tile::index_t, 4>{1, 1, 1, 1});
@@ -302,6 +305,10 @@ bool run(const ck_tile::ArgParser& arg_parser)
         use_dbias
             ? get_lengths(i_perm, shape_batch, nhead, shape_seqlen_q, max_seqlen_k)
             : std::array<ck_tile::index_t, 4>{1, 1, 1, 1} /* dummy shape for simplifying code */);
+    ck_tile::HostTensor<AccDataType> dq_acc_host(
+        i_perm
+            ? std::array<ck_tile::index_t, 5>{nsplits, shape_batch, nhead, shape_seqlen_q, hdim_q}
+            : std::array<ck_tile::index_t, 5>{nsplits, shape_batch, shape_seqlen_q, nhead, hdim_q});
 
     if(init_method == 0)
     {
@@ -362,6 +369,7 @@ bool run(const ck_tile::ArgParser& arg_parser)
     ck_tile::DeviceMem seqstart_q(seqstart_q_host.size() * sizeof(int32_t));
     ck_tile::DeviceMem seqstart_k(seqstart_k_host.size() * sizeof(int32_t));
     ck_tile::DeviceMem alibi_slope_buf(alibi_slope_host.get_element_space_size_in_bytes());
+    ck_tile::DeviceMem dq_acc_buf(dq_acc_host.get_element_space_size_in_bytes());
 
     q_buf.ToDevice(q_host.data());
     k_buf.ToDevice(k_host.data());
@@ -387,8 +395,17 @@ bool run(const ck_tile::ArgParser& arg_parser)
     std::cout << "[" << prec << "|" << mode << "|" << io_layout(i_perm, o_perm) << "] b:" << batch
               << ", h:" << nhead << "/" << nhead_k << ", s:" << seqlen_q << "/" << seqlen_k
               << ", d:" << hdim_q << "/" << hdim_v << ", scale:" << scale << ", bias:" << bias
-              << ", dbias:" << use_dbias << ", p_drop:" << p_drop << ", mask:" << mask
-              << std::flush;
+              << ", dbias:" << use_dbias << ", p_drop:" << p_drop << ", s_randval:" << s_randval
+              << ", deterministic:" << deterministic << ", mask:" << mask << std::flush;
+
+    std::size_t workspace_size =
+        dq_acc_host.get_element_space_size_in_bytes() * sizeof(AccDataType) / (1024 * 1024);
+
+    if(deterministic == 1)
+    {
+        std::cout << "\nDeterministic mode ON: " << workspace_size
+                  << " MByte memory workspace allocated" << std::endl;
+    }
 
     auto fmha_traits = fmha_bwd_traits{hdim_q,
                                        hdim_v,
@@ -397,7 +414,9 @@ bool run(const ck_tile::ArgParser& arg_parser)
                                        mask.type,
                                        bias.type,
                                        use_dbias,
-                                       p_drop > 0.0f};
+                                       p_drop > 0.0f,
+                                       s_randval,
+                                       deterministic};
     auto fmha_args   = [&]() {
         assert(nhead % nhead_k == 0);
         /// NOTE: we broadcast bias from [1, 1, seqlen_q, seqlen_k] to [batch, nhead, seqlen_q,
@@ -422,7 +441,7 @@ bool run(const ck_tile::ArgParser& arg_parser)
         const ck_tile::index_t nhead_stride_o       = (o_perm ? shape_seqlen_q * hdim_v : hdim_v);
         const ck_tile::index_t nhead_stride_randval = (shape_seqlen_q * max_seqlen_k);
         const ck_tile::index_t nhead_stride_do      = (o_perm ? shape_seqlen_q * hdim_v : hdim_v);
-        const ck_tile::index_t nhead_stride_lsed    = max_seqlen_q;
+        const ck_tile::index_t nhead_stride_lsed    = shape_seqlen_q;
         const ck_tile::index_t nhead_stride_dbias =
             (i_perm ? shape_seqlen_q * max_seqlen_k : max_seqlen_k);
         // setup batch_stride_* arguments
@@ -433,10 +452,12 @@ bool run(const ck_tile::ArgParser& arg_parser)
         const ck_tile::index_t batch_stride_o       = (nhead * shape_seqlen_q * hdim_v);
         const ck_tile::index_t batch_stride_randval = (nhead * shape_seqlen_q * max_seqlen_k);
         const ck_tile::index_t batch_stride_do      = (nhead * shape_seqlen_q * hdim_v);
-        const ck_tile::index_t batch_stride_lsed    = (nhead * max_seqlen_q);
+        const ck_tile::index_t batch_stride_lsed    = (nhead * shape_seqlen_q);
         const ck_tile::index_t batch_stride_dk      = (nhead * shape_seqlen_k * hdim_q);
         const ck_tile::index_t batch_stride_dv      = (nhead * shape_seqlen_k * hdim_v);
         const ck_tile::index_t batch_stride_dbias   = (nhead * shape_seqlen_q * max_seqlen_k);
+        const ck_tile::index_t split_stride_dq_acc =
+            (shape_batch * nhead * shape_seqlen_q * hdim_q);
 
         return fmha_bwd_args{q_buf.GetDeviceBuffer(),
                              k_buf.GetDeviceBuffer(),
@@ -452,6 +473,7 @@ bool run(const ck_tile::ArgParser& arg_parser)
                              dk_buf.GetDeviceBuffer(),
                              dv_buf.GetDeviceBuffer(),
                              dbias_buf.GetDeviceBuffer(),
+                             dq_acc_buf.GetDeviceBuffer(),
                              seqstart_q.GetDeviceBuffer(),
                              seqstart_k.GetDeviceBuffer(),
                              nullptr,
@@ -473,6 +495,8 @@ bool run(const ck_tile::ArgParser& arg_parser)
                              stride_o,
                              stride_randval,
                              stride_do,
+                             stride_q, // stride_dq_acc
+                             stride_q, // stride_dq
                              stride_dk,
                              stride_dv,
                              stride_dbias,
@@ -484,6 +508,10 @@ bool run(const ck_tile::ArgParser& arg_parser)
                              nhead_stride_randval,
                              nhead_stride_do,
                              nhead_stride_lsed,
+                             nhead_stride_q, // nhead_stride_dq_acc
+                             nhead_stride_q, // nhead_stride_dq
+                             nhead_stride_k, // nhead_stride_dk
+                             nhead_stride_v, // nhead_stride_dv
                              nhead_stride_dbias,
                              batch_stride_q,
                              batch_stride_k,
@@ -493,15 +521,17 @@ bool run(const ck_tile::ArgParser& arg_parser)
                              batch_stride_randval,
                              batch_stride_do,
                              batch_stride_lsed,
+                             batch_stride_q, // batch_stride_dq_acc
+                             batch_stride_q, // batch_stride_dq
                              batch_stride_dk,
                              batch_stride_dv,
                              batch_stride_dbias,
+                             split_stride_dq_acc,
                              mask.left,
                              mask.right,
                              static_cast<ck_tile::index_t>(mask.type),
                              p_drop,
                              p_undrop,
-                             s_randval,
                              {drop_seed, drop_offset}};
     }();
 
@@ -719,7 +749,7 @@ bool run(const ck_tile::ArgParser& arg_parser)
         if(o_perm) o_host_ref.ForEach([&](auto& self, auto idx) { o_host(b, idx[0], idx[1] + query_offset, idx[2]) = self(idx); });
         else       o_host_ref.ForEach([&](auto& self, auto idx) { o_host(b, idx[1] + query_offset, idx[0], idx[2]) = self(idx); });
 
-        lse_host_ref.ForEach([&](auto& self, auto idx) { lse_host(wb, idx[0], idx[1]) = self(idx); });
+        lse_host_ref.ForEach([&](auto& self, auto idx) { lse_host(b, idx[0], idx[1] + query_offset) = self(idx); });
         // clang-format on
 
         q_host_refs.push_back(q_host_ref);
@@ -738,6 +768,7 @@ bool run(const ck_tile::ArgParser& arg_parser)
     lse_buf.ToDevice(lse_host.data());
     dq_buf.SetZero();
     dbias_buf.SetZero();
+    dq_acc_buf.SetZero();
 
     ck_tile::stream_config stream_config_v{
         nullptr, true, 0, 0, 1, arg_parser.get_str("timer") == std::string("gpu")};

example/ck_tile/01_fmha/fmha_bwd.hpp

@@ -77,6 +77,7 @@ struct fmha_bwd_args
     void* dk_ptr;
     void* dv_ptr;
     void* dbias_ptr;
+    void* dq_acc_ptr;
     const void* seqstart_q_ptr;
     const void* seqstart_k_ptr;
     const void* seqlen_k_ptr;
@@ -97,6 +98,8 @@ struct fmha_bwd_args
     ck_tile::index_t stride_o;
     ck_tile::index_t stride_randval;
     ck_tile::index_t stride_do;
+    ck_tile::index_t stride_dq_acc;
+    ck_tile::index_t stride_dq;
     ck_tile::index_t stride_dk;
     ck_tile::index_t stride_dv;
     ck_tile::index_t stride_dbias;
@@ -108,6 +111,10 @@ struct fmha_bwd_args
     ck_tile::index_t nhead_stride_randval;
     ck_tile::index_t nhead_stride_do;
     ck_tile::index_t nhead_stride_lsed;
+    ck_tile::index_t nhead_stride_dq_acc;
+    ck_tile::index_t nhead_stride_dq;
+    ck_tile::index_t nhead_stride_dk;
+    ck_tile::index_t nhead_stride_dv;
     ck_tile::index_t nhead_stride_dbias;
     ck_tile::index_t batch_stride_q;
     ck_tile::index_t batch_stride_k;
@@ -117,15 +124,17 @@ struct fmha_bwd_args
     ck_tile::index_t batch_stride_randval;
     ck_tile::index_t batch_stride_do;
     ck_tile::index_t batch_stride_lsed;
+    ck_tile::index_t batch_stride_dq_acc;
+    ck_tile::index_t batch_stride_dq;
     ck_tile::index_t batch_stride_dk;
     ck_tile::index_t batch_stride_dv;
     ck_tile::index_t batch_stride_dbias;
+    ck_tile::index_t split_stride_dq_acc;
     ck_tile::index_t window_size_left;
     ck_tile::index_t window_size_right;
     ck_tile::index_t mask_type;
     float p_drop;
     float p_undrop;
-    bool s_randval;
     std::tuple<uint64_t, uint64_t> drop_seed_offset;
 };
 
@@ -145,10 +154,10 @@ auto fmha_bwd_dq_dk_dv_create_kargs_and_grids(fmha_bwd_args args)
                                                   args.do_ptr,
                                                   args.d_ptr,
                                                   args.rand_val_ptr,
-                                                  args.dq_ptr,
                                                   args.dk_ptr,
                                                   args.dv_ptr,
                                                   args.dbias_ptr,
+                                                  args.dq_acc_ptr,
                                                   args.seqstart_q_ptr,
                                                   args.seqstart_k_ptr,
                                                   args.seqlen_k_ptr,
@@ -163,6 +172,7 @@ auto fmha_bwd_dq_dk_dv_create_kargs_and_grids(fmha_bwd_args args)
                                                   args.stride_bias,
                                                   args.stride_randval,
                                                   args.stride_do,
+                                                  args.stride_dq_acc,
                                                   args.stride_dk,
                                                   args.stride_dv,
                                                   args.stride_dbias,
@@ -173,13 +183,15 @@ auto fmha_bwd_dq_dk_dv_create_kargs_and_grids(fmha_bwd_args args)
                                                   args.nhead_stride_randval,
                                                   args.nhead_stride_do,
                                                   args.nhead_stride_lsed,
+                                                  args.nhead_stride_dq_acc,
+                                                  args.nhead_stride_dk,
+                                                  args.nhead_stride_dv,
                                                   args.nhead_stride_dbias,
-                                                  args.batch_stride_lsed,
+                                                  args.split_stride_dq_acc,
                                                   args.window_size_left,
                                                   args.window_size_right,
                                                   args.mask_type,
                                                   args.p_drop,
-                                                  args.s_randval,
                                                   args.drop_seed_offset);
         }
         else
@@ -192,10 +204,10 @@ auto fmha_bwd_dq_dk_dv_create_kargs_and_grids(fmha_bwd_args args)
                                                   args.do_ptr,
                                                   args.d_ptr,
                                                   args.rand_val_ptr,
-                                                  args.dq_ptr,
                                                   args.dk_ptr,
                                                   args.dv_ptr,
                                                   args.dbias_ptr,
+                                                  args.dq_acc_ptr,
                                                   args.seqlen_q,
                                                   args.seqlen_k,
                                                   args.hdim_q,
@@ -209,6 +221,7 @@ auto fmha_bwd_dq_dk_dv_create_kargs_and_grids(fmha_bwd_args args)
                                                   args.stride_bias,
                                                   args.stride_randval,
                                                   args.stride_do,
+                                                  args.stride_dq_acc,
                                                   args.stride_dk,
                                                   args.stride_dv,
                                                   args.stride_dbias,
@@ -219,6 +232,9 @@ auto fmha_bwd_dq_dk_dv_create_kargs_and_grids(fmha_bwd_args args)
                                                   args.nhead_stride_randval,
                                                   args.nhead_stride_do,
                                                   args.nhead_stride_lsed,
+                                                  args.nhead_stride_dq_acc,
+                                                  args.nhead_stride_dk,
+                                                  args.nhead_stride_dv,
                                                   args.nhead_stride_dbias,
                                                   args.batch_stride_q,
                                                   args.batch_stride_k,
@@ -227,14 +243,15 @@ auto fmha_bwd_dq_dk_dv_create_kargs_and_grids(fmha_bwd_args args)
                                                   args.batch_stride_randval,
                                                   args.batch_stride_do,
                                                   args.batch_stride_lsed,
+                                                  args.batch_stride_dq_acc,
                                                   args.batch_stride_dk,
                                                   args.batch_stride_dv,
                                                   args.batch_stride_dbias,
+                                                  args.split_stride_dq_acc,
                                                   args.window_size_left,
                                                   args.window_size_right,
                                                   args.mask_type,
                                                   args.p_drop,
-                                                  args.s_randval,
                                                   args.drop_seed_offset);
         }
     }();
@@ -260,8 +277,7 @@ auto fmha_bwd_dot_do_o_create_kargs_and_grids(fmha_bwd_args args)
                                                      args.stride_o,
                                                      args.nhead_stride_do,
                                                      args.nhead_stride_o,
-                                                     args.nhead_stride_lsed,
-                                                     args.batch_stride_lsed);
+                                                     args.nhead_stride_lsed);
         }
         else
         { // create batch mode kernel arguments
@@ -286,19 +302,59 @@ auto fmha_bwd_dot_do_o_create_kargs_and_grids(fmha_bwd_args args)
     return ck_tile::make_tuple(kargs, grids);
 }
 
+template <typename FmhaBwdConvertQGradKernel>
+auto fmha_bwd_convert_dq_create_kargs_and_grids(fmha_bwd_args args)
+{
+    auto kargs = [&] {
+        // create group mode kernel arguments
+        if constexpr(FmhaBwdConvertQGradKernel::kIsGroupMode)
+        {
+            return FmhaBwdConvertQGradKernel::MakeKargs(args.dq_acc_ptr,
+                                                        args.dq_ptr,
+                                                        args.seqstart_q_ptr,
+                                                        args.seqstart_k_ptr,
+                                                        args.hdim_q,
+                                                        args.stride_dq,
+                                                        args.stride_dq_acc,
+                                                        args.nhead_stride_dq,
+                                                        args.nhead_stride_dq_acc,
+                                                        args.split_stride_dq_acc);
+        }
+        else
+        { // create batch mode kernel arguments
+            return FmhaBwdConvertQGradKernel::MakeKargs(args.dq_acc_ptr,
+                                                        args.dq_ptr,
+                                                        args.seqlen_q,
+                                                        args.seqlen_k,
+                                                        args.hdim_q,
+                                                        args.stride_dq,
+                                                        args.stride_dq_acc,
+                                                        args.nhead_stride_dq,
+                                                        args.nhead_stride_dq_acc,
+                                                        args.batch_stride_dq,
+                                                        args.batch_stride_dq_acc,
+                                                        args.split_stride_dq_acc);
+        }
+    }();
+
+    dim3 grids = FmhaBwdConvertQGradKernel::GridSize(args.batch, args.nhead_q, args.max_seqlen_q);
+    return ck_tile::make_tuple(kargs, grids);
+}
+
 // this is used to pattern-match internl kernel implementation, not to instantiate kernel
 template <ck_tile::index_t HDim_,
           typename DataType_,
           bool kIsGroupMode_,
           ck_tile::BlockFmhaBwdPipelineEnum FmhaBwdPipelineEnum_,
           typename FmhaMask_,
+          typename FmhaDropout_,
           ck_tile::BlockAttentionBiasEnum BiasEnum_,
           bool kHasBiasGrad_,
-          bool kHasDropout_,
           bool kPadS_,
           bool kPadSK_,
           bool kPadD_,
-          bool kPadDv_>
+          bool kPadDv_,
+          bool kIsDeterministic_>
 struct fmha_bwd_dq_dk_dv_traits_
 {
     static constexpr ck_tile::index_t HDim    = HDim_;
@@ -306,13 +362,14 @@ struct fmha_bwd_dq_dk_dv_traits_
     static constexpr bool kIsGroupMode        = kIsGroupMode_;
     static constexpr auto FmhaBwdPipelineEnum = FmhaBwdPipelineEnum_;
     using FmhaMask                            = ck_tile::remove_cvref_t<FmhaMask_>;
+    using FmhaDropout                         = ck_tile::remove_cvref_t<FmhaDropout_>;
     static constexpr auto BiasEnum            = BiasEnum_;
     static constexpr bool kHasBiasGrad        = kHasBiasGrad_;
-    static constexpr bool kHasDropout         = kHasDropout_;
     static constexpr bool kPadS               = kPadS_;
     static constexpr bool kPadSK              = kPadSK_;
     static constexpr bool kPadD               = kPadD_;
     static constexpr bool kPadDv              = kPadDv_;
+    static constexpr bool kIsDeterministic    = kIsDeterministic_;
 };
 
 template <typename Traits_>
@@ -343,6 +400,31 @@ void fmha_bwd_dot_do_o_oneshot_(const ck_tile::stream_config&, fmha_bwd_args);
 template <typename Traits_>
 std::string fmha_bwd_dot_do_o_get_name_();
 
+template <ck_tile::index_t HDim_,
+          typename DataType_,
+          bool kIsGroupMode_,
+          bool kPadS_,
+          bool kPadD_,
+          bool kIsDeterministic_>
+struct fmha_bwd_convert_dq_traits_
+{
+    static constexpr ck_tile::index_t HDim = HDim_;
+    using DataType                         = ck_tile::remove_cvref_t<DataType_>;
+    static constexpr bool kIsGroupMode     = kIsGroupMode_;
+    static constexpr bool kPadS            = kPadS_;
+    static constexpr bool kPadD            = kPadD_;
+    static constexpr bool kIsDeterministic = kIsDeterministic_;
+};
+
+template <typename Traits_>
+float fmha_bwd_convert_dq_(const ck_tile::stream_config&, fmha_bwd_args);
+
+template <typename Traits_>
+void fmha_bwd_convert_dq_oneshot_(const ck_tile::stream_config&, fmha_bwd_args);
+
+template <typename Traits_>
+std::string fmha_bwd_convert_dq_get_name_();
+
 // This is the public API, will be generated by script
 struct fmha_bwd_traits
 {
@@ -354,6 +436,8 @@ struct fmha_bwd_traits
     bias_enum bias_type; // 0:no bias, 1:elementwise bias, 2:alibi. sync with BlockAttentionBiasEnum
     bool has_dbias;
     bool has_dropout;
+    bool is_store_randval;
+    bool is_deterministic;
     // TODO: padding check is inside this api
 };
 float fmha_bwd(fmha_bwd_traits, fmha_bwd_args, const ck_tile::stream_config&);

example/ck_tile/01_fmha/fmha_fwd.cpp

@@ -479,16 +479,18 @@ bool run(const ck_tile::ArgParser& arg_parser)
             : std::array<ck_tile::index_t, 2>{1, 1});
 
     ck_tile::HostTensor<LSEDataType> lse_acc_host(
-        1 < num_splits ? std::array<ck_tile::index_t, 4>{num_splits, batch, nhead, max_seqlen_q}
-                       : std::array<ck_tile::index_t, 4>{1, 1, 1, 1});
+        1 < num_splits
+            ? std::array<ck_tile::index_t, 4>{num_splits, shape_batch, nhead, shape_seqlen_q}
+            : std::array<ck_tile::index_t, 4>{1, 1, 1, 1});
     ck_tile::HostTensor<OaccDataType> o_acc_host(
         1 < num_splits
             ? std::array<ck_tile::index_t, 5>{num_splits, batch, nhead, max_seqlen_q, hdim_v}
             : std::array<ck_tile::index_t, 5>{1, 1, 1, 1, 1});
 
-    // self define lse data layout as [batch, nhead, max_seqlen_q]
+    // batch mode of lse data layout is [batch, nhead, seqlen_q]
+    // group mode of lse data layout is [nhead, total_seqlen_q]
     ck_tile::HostTensor<LSEDataType> lse_host(
-        lse ? std::array<ck_tile::index_t, 3>{batch, nhead, max_seqlen_q}
+        lse ? std::array<ck_tile::index_t, 3>{shape_batch, nhead, shape_seqlen_q}
             : std::array<ck_tile::index_t, 3>{1, 1, 1} /* dummy shape for simplifying code */);
 
     ck_tile::HostTensor<ODataType> o_host(
@@ -669,8 +671,8 @@ bool run(const ck_tile::ArgParser& arg_parser)
         const ck_tile::index_t nhead_stride_bias =
             (i_perm ? 0 * shape_seqlen_q * shape_seqlen_k : 0 * shape_seqlen_k);
         const ck_tile::index_t nhead_stride_randval = (shape_seqlen_q * max_seqlen_k);
-        const ck_tile::index_t nhead_stride_lse     = max_seqlen_q;
-        const ck_tile::index_t nhead_stride_lse_acc = max_seqlen_q;
+        const ck_tile::index_t nhead_stride_lse     = shape_seqlen_q;
+        const ck_tile::index_t nhead_stride_lse_acc = shape_seqlen_q;
         const ck_tile::index_t nhead_stride_o_acc   = (max_seqlen_q * hdim_v);
         const ck_tile::index_t nhead_stride_o       = (o_perm ? shape_seqlen_q * hdim_v : hdim_v);
         // setup batch_stride_* arguments
@@ -679,12 +681,12 @@ bool run(const ck_tile::ArgParser& arg_parser)
         const ck_tile::index_t batch_stride_v       = (nhead_k * hdim_v * shape_seqlen_k);
         const ck_tile::index_t batch_stride_bias    = (0 * nhead * shape_seqlen_q * shape_seqlen_k);
         const ck_tile::index_t batch_stride_randval = (nhead * shape_seqlen_q * max_seqlen_k);
-        const ck_tile::index_t batch_stride_lse     = (nhead * max_seqlen_q);
-        const ck_tile::index_t batch_stride_lse_acc = (nhead * max_seqlen_q);
+        const ck_tile::index_t batch_stride_lse     = (nhead * shape_seqlen_q);
+        const ck_tile::index_t batch_stride_lse_acc = (nhead * shape_seqlen_q);
         const ck_tile::index_t batch_stride_o_acc   = (nhead * max_seqlen_q * hdim_v);
         const ck_tile::index_t batch_stride_o       = (nhead * shape_seqlen_q * hdim_v);
         // setup split_stride_* arguments (only used in split-kv kernel)
-        const ck_tile::index_t split_stride_lse_acc = (batch * nhead * max_seqlen_q);
+        const ck_tile::index_t split_stride_lse_acc = (shape_batch * nhead * shape_seqlen_q);
         const ck_tile::index_t split_stride_o_acc   = (batch * nhead * max_seqlen_q * hdim_v);
 
         return fmha_fwd_args{q_buf.GetDeviceBuffer(),
@@ -996,8 +998,9 @@ bool run(const ck_tile::ArgParser& arg_parser)
         if(lse)
         {
             ck_tile::HostTensor<SMPLComputeDataType> lse_host_result({nhead, real_seqlen_q});
-            lse_host_result.ForEach(
-                [&](auto& self, auto idx) { self(idx) = lse_host(wb, idx[0], idx[1]); });
+            lse_host_result.ForEach([&](auto& self, auto idx) {
+                self(idx) = lse_host(b, idx[0], idx[1] + query_offset);
+            });
 
             cur_pass = ck_tile::check_err(lse_host_result,
                                           lse_host_ref,

example/ck_tile/01_fmha/fmha_fwd.hpp

@@ -185,7 +185,6 @@ auto fmha_fwd_create_kargs_and_grids(fmha_fwd_args args)
                                          args.nhead_stride_randval,
                                          args.nhead_stride_lse,
                                          args.nhead_stride_o,
-                                         args.batch_stride_lse,
                                          args.window_size_left,
                                          args.window_size_right,
                                          args.mask_type,
@@ -284,7 +283,6 @@ auto fmha_fwd_splitkv_create_kargs_and_grids(fmha_fwd_args args)
                                      args.nhead_stride_randval,
                                      args.nhead_stride_lse_acc,
                                      args.nhead_stride_o_acc,
-                                     args.batch_stride_lse_acc,
                                      args.batch_stride_o_acc,
                                      args.split_stride_lse_acc,
                                      args.split_stride_o_acc,
@@ -376,9 +374,7 @@ auto fmha_fwd_splitkv_combine_create_kargs_and_grids(fmha_fwd_args args)
                                      args.nhead_stride_o_acc,
                                      args.nhead_stride_lse,
                                      args.nhead_stride_o,
-                                     args.batch_stride_lse_acc,
                                      args.batch_stride_o_acc,
-                                     args.batch_stride_lse,
                                      args.split_stride_lse_acc,
                                      args.split_stride_o_acc);
         }

example/ck_tile/01_fmha/script/smoke_test_bwd.sh

@@ -11,18 +11,19 @@ COMMON_ARGS='-v=1'
 set -x
 for prec in "fp16" "bf16" ; do
 for perm in 0 1 ; do
-for hdim in 32 64 128 ; do
+for hdim in 32 64 128 256 ; do
 for mode in 0 1 ; do
-for bias in "n" "e" "a"; do
-for dbias in 0 1 ; do
-for p_drop in 0.0 0.2; do
+for bias in "n" "a" ; do
+for dbias in 0 ; do
+for p_drop in 0.0 0.2 ; do
+for deterministic in 0 ; do
 
-$EXE -prec=$prec -b=1 -h=4 -h_k=2 -d=$hdim -s=259 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS
-$EXE -prec=$prec -b=2 -h=2 -d=$hdim -s=516 -s_k=253 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS
-$EXE -prec=$prec -b=1 -h=4 -h_k=1 -d=$hdim -s=500 -s_k=251 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -mask=1 -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS
-$EXE -prec=$prec -b=1 -h=2 -d=$hdim -s=900 -s_k=258 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -mask=2 -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS
-$EXE -prec=$prec -b=2 -h=1 -d=$hdim -s=987 -s_k=219 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -mask=t:128,30 -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS
-$EXE -prec=$prec -b=2 -h=3 -h_k=1 -d=$hdim -s=244 -s_k=499 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -mask=b:4,35 -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS
+$EXE -prec=$prec -b=1 -h=4 -h_k=2 -d=$hdim -s=259 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -deterministic=$deterministic -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS
+$EXE -prec=$prec -b=2 -h=2 -d=$hdim -s=516 -s_k=253 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -deterministic=$deterministic -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS
+$EXE -prec=$prec -b=1 -h=4 -h_k=1 -d=$hdim -s=500 -s_k=251 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -mask=1 -deterministic=$deterministic -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS
+$EXE -prec=$prec -b=1 -h=2 -d=$hdim -s=900 -s_k=258 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -mask=2 -v=1 -deterministic=$deterministic -mode=$mode -kname=$KNAME $COMMON_ARGS
+$EXE -prec=$prec -b=2 -h=1 -d=$hdim -s=987 -s_k=219 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -mask=t:128,30 -deterministic=$deterministic -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS
+$EXE -prec=$prec -b=2 -h=3 -h_k=1 -d=$hdim -s=244 -s_k=499 -bias=$bias -dbias=$dbias -p_drop=$p_drop -iperm=$perm -operm=$perm -mask=b:4,35 -deterministic=$deterministic -v=1 -mode=$mode -kname=$KNAME $COMMON_ARGS
 
 done
 done
@@ -31,4 +32,5 @@ done
 done
 done
 done
+done
 set +x

include/ck/ck.hpp

@@ -1,5 +1,5 @@
 // SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
 
 #pragma once
 
@@ -153,8 +153,8 @@ CK_DECLARE_ENV_VAR_BOOL(CK_LOGGING)
 // LDS direct loads using inline assembly
 #define CK_USE_AMD_LDS_DIRECT_LOAD_INLINE_ASM 0
 
-// set stochastic rounding as default for f8 conversions
-#define CK_USE_SR_F8_CONVERSION 1
+// set rounding to nearest even as default for f8 conversions
+#define CK_USE_SR_F8_CONVERSION 0
 
 // block synchronization only s_wait lgkmcnt(0), not vmcnt(0)
 #define CK_EXPERIMENTAL_BLOCK_SYNC_LDS_WITHOUT_SYNC_VMEM 1

include/ck/host_utility/device_prop.hpp

@@ -65,6 +65,12 @@ inline bool is_lds_direct_load_supported()
            ck::get_device_name() == "gfx941" || ck::get_device_name() == "gfx942";
 }
 
+inline bool is_bf16_atomic_supported()
+{
+    return ck::get_device_name() == "gfx940" || ck::get_device_name() == "gfx941" ||
+           ck::get_device_name() == "gfx942";
+}
+
 inline bool is_gfx101_supported()
 {
     return ck::get_device_name() == "gfx1010" || ck::get_device_name() == "gfx1011" ||

include/ck/tensor_operation/gpu/device/device_gemm_multiple_d.hpp

@@ -53,6 +53,49 @@ struct DeviceGemmMultipleD : public BaseOperator
     virtual std::unique_ptr<BaseInvoker> MakeInvokerPointer() = 0;
 };
 
+// GEMM:
+//   input : A[M, K], B[K, N],
+//   input : D0[M, N], D1[M, N], ...
+//   output : E[M, N]
+//   C = a_op(A) * b_op(B)
+//   E = cde_op(C, D0, D1, ...)
+// Assume:
+//   D0, D1, ... and E have the same layout
+template <typename ALayout,
+          typename BLayout,
+          typename DsLayout,
+          typename ELayout,
+          typename ADataType,
+          typename BDataType,
+          typename DsDataType,
+          typename EDataType,
+          typename AElementwiseOperation,
+          typename BElementwiseOperation,
+          typename CDEElementwiseOperation>
+struct DeviceGemmMultipleDSplitK : public BaseOperator
+{
+    static constexpr index_t NumDTensor = DsDataType::Size();
+
+    virtual std::unique_ptr<BaseArgument>
+    MakeArgumentPointer(const void* p_a,
+                        const void* p_b,
+                        std::array<const void*, NumDTensor> p_ds,
+                        void* p_e,
+                        ck::index_t M,
+                        ck::index_t N,
+                        ck::index_t K,
+                        ck::index_t StrideA,
+                        ck::index_t StrideB,
+                        std::array<ck::index_t, NumDTensor> StrideDs,
+                        ck::index_t StrideE,
+                        ck::index_t KBatch,
+                        AElementwiseOperation a_element_op,
+                        BElementwiseOperation b_element_op,
+                        CDEElementwiseOperation cde_element_op) = 0;
+
+    virtual std::unique_ptr<BaseInvoker> MakeInvokerPointer() = 0;
+};
+
 } // namespace device
 } // namespace tensor_operation
 } // namespace ck

include/ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_xdl_cshuffle_v3.hpp

@@ -69,17 +69,17 @@ template <typename ALayout,
           typename ComputeTypeB                       = ComputeTypeA,
           typename LDSTypeA                           = ComputeTypeA,
           typename LDSTypeB                           = ComputeTypeB>
-struct DeviceGemmMultiD_Xdl_CShuffle_V3 : public DeviceGemmMultipleD<ALayout,
-                                                                     BLayout,
-                                                                     DsLayout,
-                                                                     CLayout,
-                                                                     ADataType,
-                                                                     BDataType,
-                                                                     DsDataType,
-                                                                     CDataType,
-                                                                     AElementwiseOperation,
-                                                                     BElementwiseOperation,
-                                                                     CElementwiseOperation>
+struct DeviceGemmMultiD_Xdl_CShuffle_V3 : public DeviceGemmMultipleDSplitK<ALayout,
+                                                                           BLayout,
+                                                                           DsLayout,
+                                                                           CLayout,
+                                                                           ADataType,
+                                                                           BDataType,
+                                                                           DsDataType,
+                                                                           CDataType,
+                                                                           AElementwiseOperation,
+                                                                           BElementwiseOperation,
+                                                                           CElementwiseOperation>
 {
     static constexpr index_t NumDTensor = DsDataType::Size();
 
@@ -192,15 +192,11 @@ struct DeviceGemmMultiD_Xdl_CShuffle_V3 : public DeviceGemmMultipleD<ALayout,
                         // rotating mem
                         rotating_mem.Next();
                         // clear c mem
-                        if constexpr(!is_same<remove_cvref_t<CDataType>, bhalf_t>::value)
-                        {
-                            if(arg_.KBatch > 1)
-                                hipGetErrorString(
-                                    hipMemsetAsync(arg_.p_c_grid,
-                                                   0,
-                                                   arg_.M * arg_.N * sizeof(CDataType),
-                                                   stream_config.stream_id_));
-                        }
+                        if(arg_.KBatch > 1)
+                            hipGetErrorString(hipMemsetAsync(arg_.p_c_grid,
+                                                             0,
+                                                             arg_.M * arg_.N * sizeof(CDataType),
+                                                             stream_config.stream_id_));
                     };
 
                     ave_time = ck::utility::launch_and_time_kernel_with_preprocess<false>(
@@ -234,38 +230,49 @@ struct DeviceGemmMultiD_Xdl_CShuffle_V3 : public DeviceGemmMultipleD<ALayout,
                 if constexpr(BlkGemmPipelineVer == BlockGemmPipelineVersion::v1 ||
                              BlkGemmPipelineVer == BlockGemmPipelineVersion::v3)
                 {
+                    if(arg.KBatch > 1)
+                    {
+                        const auto kernel = kernel_gemm_xdl_cshuffle_v3_multi_d<
+                            GridwiseGemm,
+                            true,
+                            InMemoryDataOperationEnum::AtomicAdd,
+                            minimum_occupancy>;
+                        Run(kernel);
+                    }
+                    else
                     {
                         const auto kernel =
-                            kernel_gemm_xdl_cshuffle_v3<GridwiseGemm,
-                                                        true,
-                                                        InMemoryDataOperationEnum::Set,
-                                                        minimum_occupancy>;
+                            kernel_gemm_xdl_cshuffle_v3_multi_d<GridwiseGemm,
+                                                                true,
+                                                                InMemoryDataOperationEnum::Set,
+                                                                minimum_occupancy>;
                         Run(kernel);
                     }
                 }
                 // Tail number could be One to Seven
                 else if constexpr(BlkGemmPipelineVer == BlockGemmPipelineVersion::v2)
                 {
+                    if(arg.KBatch > 1)
                     {
                         if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::One)
                         {
-                            const auto kernel =
-                                kernel_gemm_xdl_cshuffle_v3<GridwiseGemm,
-                                                            true,
-                                                            InMemoryDataOperationEnum::Set,
-                                                            minimum_occupancy,
-                                                            TailNumber::One>;
+                            const auto kernel = kernel_gemm_xdl_cshuffle_v3_multi_d<
+                                GridwiseGemm,
+                                true,
+                                InMemoryDataOperationEnum::AtomicAdd,
+                                minimum_occupancy,
+                                TailNumber::One>;
                             Run(kernel);
                         }
                         else if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) ==
                                 TailNumber::Full)
                         {
-                            const auto kernel =
-                                kernel_gemm_xdl_cshuffle_v3<GridwiseGemm,
-                                                            true,
-                                                            InMemoryDataOperationEnum::Set,
-                                                            minimum_occupancy,
-                                                            TailNumber::Full>;
+                            const auto kernel = kernel_gemm_xdl_cshuffle_v3_multi_d<
+                                GridwiseGemm,
+                                true,
+                                InMemoryDataOperationEnum::AtomicAdd,
+                                minimum_occupancy,
+                                TailNumber::Full>;
                             Run(kernel);
                         }
 
@@ -273,12 +280,12 @@ struct DeviceGemmMultiD_Xdl_CShuffle_V3 : public DeviceGemmMultipleD<ALayout,
                         {
                             if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Two)
                             {
-                                const auto kernel =
-                                    kernel_gemm_xdl_cshuffle_v3<GridwiseGemm,
-                                                                true,
-                                                                InMemoryDataOperationEnum::Set,
-                                                                minimum_occupancy,
-                                                                TailNumber::Two>;
+                                const auto kernel = kernel_gemm_xdl_cshuffle_v3_multi_d<
+                                    GridwiseGemm,
+                                    true,
+                                    InMemoryDataOperationEnum::AtomicAdd,
+                                    minimum_occupancy,
+                                    TailNumber::Two>;
                                 Run(kernel);
                             }
                         }
@@ -288,12 +295,12 @@ struct DeviceGemmMultiD_Xdl_CShuffle_V3 : public DeviceGemmMultipleD<ALayout,
                             if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) ==
                                TailNumber::Three)
                             {
-                                const auto kernel =
-                                    kernel_gemm_xdl_cshuffle_v3<GridwiseGemm,
-                                                                true,
-                                                                InMemoryDataOperationEnum::Set,
-                                                                minimum_occupancy,
-                                                                TailNumber::Three>;
+                                const auto kernel = kernel_gemm_xdl_cshuffle_v3_multi_d<
+                                    GridwiseGemm,
+                                    true,
+                                    InMemoryDataOperationEnum::AtomicAdd,
+                                    minimum_occupancy,
+                                    TailNumber::Three>;
                                 Run(kernel);
                             }
                         }
@@ -303,12 +310,12 @@ struct DeviceGemmMultiD_Xdl_CShuffle_V3 : public DeviceGemmMultipleD<ALayout,
                             if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) ==
                                TailNumber::Four)
                             {
-                                const auto kernel =
-                                    kernel_gemm_xdl_cshuffle_v3<GridwiseGemm,
-                                                                true,
-                                                                InMemoryDataOperationEnum::Set,
-                                                                minimum_occupancy,
-                                                                TailNumber::Four>;
+                                const auto kernel = kernel_gemm_xdl_cshuffle_v3_multi_d<
+                                    GridwiseGemm,
+                                    true,
+                                    InMemoryDataOperationEnum::AtomicAdd,
+                                    minimum_occupancy,
+                                    TailNumber::Four>;
                                 Run(kernel);
                             }
                         }
@@ -318,12 +325,12 @@ struct DeviceGemmMultiD_Xdl_CShuffle_V3 : public DeviceGemmMultipleD<ALayout,
                             if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) ==
                                TailNumber::Five)
                             {
-                                const auto kernel =
-                                    kernel_gemm_xdl_cshuffle_v3<GridwiseGemm,
-                                                                true,
-                                                                InMemoryDataOperationEnum::Set,
-                                                                minimum_occupancy,
-                                                                TailNumber::Five>;
+                                const auto kernel = kernel_gemm_xdl_cshuffle_v3_multi_d<
+                                    GridwiseGemm,
+                                    true,
+                                    InMemoryDataOperationEnum::AtomicAdd,
+                                    minimum_occupancy,
+                                    TailNumber::Five>;
                                 Run(kernel);
                             }
                         }
@@ -332,12 +339,12 @@ struct DeviceGemmMultiD_Xdl_CShuffle_V3 : public DeviceGemmMultipleD<ALayout,
                         {
                             if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Six)
                             {
-                                const auto kernel =
-                                    kernel_gemm_xdl_cshuffle_v3<GridwiseGemm,
-                                                                true,
-                                                                InMemoryDataOperationEnum::Set,
-                                                                minimum_occupancy,
-                                                                TailNumber::Six>;
+                                const auto kernel = kernel_gemm_xdl_cshuffle_v3_multi_d<
+                                    GridwiseGemm,
+                                    true,
+                                    InMemoryDataOperationEnum::AtomicAdd,
+                                    minimum_occupancy,
+                                    TailNumber::Six>;
                                 Run(kernel);
                             }
                         }
@@ -347,12 +354,124 @@ struct DeviceGemmMultiD_Xdl_CShuffle_V3 : public DeviceGemmMultipleD<ALayout,
                             if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) ==
                                TailNumber::Seven)
                             {
-                                const auto kernel =
-                                    kernel_gemm_xdl_cshuffle_v3<GridwiseGemm,
-                                                                true,
-                                                                InMemoryDataOperationEnum::Set,
-                                                                minimum_occupancy,
-                                                                TailNumber::Seven>;
+                                const auto kernel = kernel_gemm_xdl_cshuffle_v3_multi_d<
+                                    GridwiseGemm,
+                                    true,
+                                    InMemoryDataOperationEnum::AtomicAdd,
+                                    minimum_occupancy,
+                                    TailNumber::Seven>;
+                                Run(kernel);
+                            }
+                        }
+                    }
+                    else
+                    {
+                        if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::One)
+                        {
+                            const auto kernel =
+                                kernel_gemm_xdl_cshuffle_v3_multi_d<GridwiseGemm,
+                                                                    true,
+                                                                    InMemoryDataOperationEnum::Set,
+                                                                    minimum_occupancy,
+                                                                    TailNumber::One>;
+                            Run(kernel);
+                        }
+                        else if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) ==
+                                TailNumber::Full)
+                        {
+                            const auto kernel =
+                                kernel_gemm_xdl_cshuffle_v3_multi_d<GridwiseGemm,
+                                                                    true,
+                                                                    InMemoryDataOperationEnum::Set,
+                                                                    minimum_occupancy,
+                                                                    TailNumber::Full>;
+                            Run(kernel);
+                        }
+
+                        if constexpr(GridwiseGemm::BlockwiseGemmPipe::PrefetchStages > 2)
+                        {
+                            if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Two)
+                            {
+                                const auto kernel = kernel_gemm_xdl_cshuffle_v3_multi_d<
+                                    GridwiseGemm,
+                                    true,
+                                    InMemoryDataOperationEnum::Set,
+                                    minimum_occupancy,
+                                    TailNumber::Two>;
+                                Run(kernel);
+                            }
+                        }
+
+                        if constexpr(GridwiseGemm::BlockwiseGemmPipe::PrefetchStages > 3)
+                        {
+                            if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) ==
+                               TailNumber::Three)
+                            {
+                                const auto kernel = kernel_gemm_xdl_cshuffle_v3_multi_d<
+                                    GridwiseGemm,
+                                    true,
+                                    InMemoryDataOperationEnum::Set,
+                                    minimum_occupancy,
+                                    TailNumber::Three>;
+                                Run(kernel);
+                            }
+                        }
+
+                        if constexpr(GridwiseGemm::BlockwiseGemmPipe::PrefetchStages > 4)
+                        {
+                            if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) ==
+                               TailNumber::Four)
+                            {
+                                const auto kernel = kernel_gemm_xdl_cshuffle_v3_multi_d<
+                                    GridwiseGemm,
+                                    true,
+                                    InMemoryDataOperationEnum::Set,
+                                    minimum_occupancy,
+                                    TailNumber::Four>;
+                                Run(kernel);
+                            }
+                        }
+
+                        if constexpr(GridwiseGemm::BlockwiseGemmPipe::PrefetchStages > 5)
+                        {
+                            if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) ==
+                               TailNumber::Five)
+                            {
+                                const auto kernel = kernel_gemm_xdl_cshuffle_v3_multi_d<
+                                    GridwiseGemm,
+                                    true,
+                                    InMemoryDataOperationEnum::Set,
+                                    minimum_occupancy,
+                                    TailNumber::Five>;
+                                Run(kernel);
+                            }
+                        }
+
+                        if constexpr(GridwiseGemm::BlockwiseGemmPipe::PrefetchStages > 6)
+                        {
+                            if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Six)
+                            {
+                                const auto kernel = kernel_gemm_xdl_cshuffle_v3_multi_d<
+                                    GridwiseGemm,
+                                    true,
+                                    InMemoryDataOperationEnum::Set,
+                                    minimum_occupancy,
+                                    TailNumber::Six>;
+                                Run(kernel);
+                            }
+                        }
+
+                        if constexpr(GridwiseGemm::BlockwiseGemmPipe::PrefetchStages > 7)
+                        {
+                            if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) ==
+                               TailNumber::Seven)
+                            {
+                                const auto kernel = kernel_gemm_xdl_cshuffle_v3_multi_d<
+                                    GridwiseGemm,
+                                    true,
+                                    InMemoryDataOperationEnum::Set,
+                                    minimum_occupancy,
+                                    TailNumber::Seven>;
                                 Run(kernel);
                             }
                         }
@@ -361,51 +480,98 @@ struct DeviceGemmMultiD_Xdl_CShuffle_V3 : public DeviceGemmMultipleD<ALayout,
                 // Tail number could be Odd or Even
                 else if constexpr(BlkGemmPipelineVer == BlockGemmPipelineVersion::v4)
                 {
-
+                    if(arg.KBatch > 1)
                     {
                         if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Odd)
                         {
-                            const auto kernel =
-                                kernel_gemm_xdl_cshuffle_v3_2lds<GridwiseGemm,
-                                                                 true,
-                                                                 InMemoryDataOperationEnum::Set,
-                                                                 minimum_occupancy,
-                                                                 TailNumber::Odd>;
+                            const auto kernel = kernel_gemm_xdl_cshuffle_v3_multi_d_2lds<
+                                GridwiseGemm,
+                                true,
+                                InMemoryDataOperationEnum::AtomicAdd,
+                                minimum_occupancy,
+                                TailNumber::Odd>;
                             Run(kernel);
                         }
                         else
                         {
-                            const auto kernel =
-                                kernel_gemm_xdl_cshuffle_v3_2lds<GridwiseGemm,
-                                                                 true,
-                                                                 InMemoryDataOperationEnum::Set,
-                                                                 minimum_occupancy,
-                                                                 TailNumber::Even>;
+                            const auto kernel = kernel_gemm_xdl_cshuffle_v3_multi_d_2lds<
+                                GridwiseGemm,
+                                true,
+                                InMemoryDataOperationEnum::AtomicAdd,
+                                minimum_occupancy,
+                                TailNumber::Even>;
+                            Run(kernel);
+                        }
+                    }
+                    else
+                    {
+                        if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Odd)
+                        {
+                            const auto kernel = kernel_gemm_xdl_cshuffle_v3_multi_d_2lds<
+                                GridwiseGemm,
+                                true,
+                                InMemoryDataOperationEnum::Set,
+                                minimum_occupancy,
+                                TailNumber::Odd>;
+                            Run(kernel);
+                        }
+                        else
+                        {
+                            const auto kernel = kernel_gemm_xdl_cshuffle_v3_multi_d_2lds<
+                                GridwiseGemm,
+                                true,
+                                InMemoryDataOperationEnum::Set,
+                                minimum_occupancy,
+                                TailNumber::Even>;
                             Run(kernel);
                         }
                     }
                 }
                 else
                 {
+                    if(arg.KBatch > 1)
+                    {
+                        if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Odd)
+                        {
+                            const auto kernel = kernel_gemm_xdl_cshuffle_v3_multi_d<
+                                GridwiseGemm,
+                                true,
+                                InMemoryDataOperationEnum::AtomicAdd,
+                                minimum_occupancy,
+                                TailNumber::Odd>;
+                            Run(kernel);
+                        }
+                        else
+                        {
+                            const auto kernel = kernel_gemm_xdl_cshuffle_v3_multi_d<
+                                GridwiseGemm,
+                                true,
+                                InMemoryDataOperationEnum::AtomicAdd,
+                                minimum_occupancy,
+                                TailNumber::Even>;
+                            Run(kernel);
+                        }
+                    }
+                    else
                     {
                         if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Odd)
                         {
                             const auto kernel =
-                                kernel_gemm_xdl_cshuffle_v3<GridwiseGemm,
-                                                            true,
-                                                            InMemoryDataOperationEnum::Set,
-                                                            minimum_occupancy,
-                                                            TailNumber::Odd>;
+                                kernel_gemm_xdl_cshuffle_v3_multi_d<GridwiseGemm,
+                                                                    true,
+                                                                    InMemoryDataOperationEnum::Set,
+                                                                    minimum_occupancy,
+                                                                    TailNumber::Odd>;
                             Run(kernel);
                         }
                         else
                         {
                             const auto kernel =
-                                kernel_gemm_xdl_cshuffle_v3<GridwiseGemm,
-                                                            true,
-                                                            InMemoryDataOperationEnum::Set,
-                                                            minimum_occupancy,
-                                                            TailNumber::Even>;
+                                kernel_gemm_xdl_cshuffle_v3_multi_d<GridwiseGemm,
+                                                                    true,
+                                                                    InMemoryDataOperationEnum::Set,
+                                                                    minimum_occupancy,
+                                                                    TailNumber::Even>;
                             Run(kernel);
                         }
                     }
@@ -416,12 +582,22 @@ struct DeviceGemmMultiD_Xdl_CShuffle_V3 : public DeviceGemmMultipleD<ALayout,
                 // Tail number always 1
                 if constexpr(BlkGemmPipelineVer == BlockGemmPipelineVersion::v1)
                 {
+                    if(arg.KBatch > 1)
+                    {
+                        const auto kernel = kernel_gemm_xdl_cshuffle_v3_multi_d<
+                            GridwiseGemm,
+                            false,
+                            InMemoryDataOperationEnum::AtomicAdd,
+                            minimum_occupancy>;
+                        Run(kernel);
+                    }
+                    else
                     {
                         const auto kernel =
-                            kernel_gemm_xdl_cshuffle_v3<GridwiseGemm,
-                                                        false,
-                                                        InMemoryDataOperationEnum::Set,
-                                                        minimum_occupancy>;
+                            kernel_gemm_xdl_cshuffle_v3_multi_d<GridwiseGemm,
+                                                                false,
+                                                                InMemoryDataOperationEnum::Set,
+                                                                minimum_occupancy>;
                         Run(kernel);
                     }
                 }
@@ -451,6 +627,11 @@ struct DeviceGemmMultiD_Xdl_CShuffle_V3 : public DeviceGemmMultipleD<ALayout,
             return false;
         }
 
+        if(!is_bf16_atomic_supported() && std::is_same_v<CDataType, ck::bhalf_t> && arg.KBatch > 1)
+        {
+            return false;
+        }
+
         if((arg.K % AK1 != 0 || arg.K % BK1 != 0) && !(GemmSpec == GemmSpecialization::MKPadding ||
                                                        GemmSpec == GemmSpecialization::NKPadding ||
                                                        GemmSpec == GemmSpecialization::MNKPadding ||
@@ -479,6 +660,7 @@ struct DeviceGemmMultiD_Xdl_CShuffle_V3 : public DeviceGemmMultipleD<ALayout,
                              index_t StrideB,
                              std::array<index_t, NumDTensor> StrideDs,
                              index_t StrideC,
+                             index_t KBatch,
                              AElementwiseOperation a_element_op,
                              BElementwiseOperation b_element_op,
                              CElementwiseOperation c_element_op)
@@ -494,7 +676,7 @@ struct DeviceGemmMultiD_Xdl_CShuffle_V3 : public DeviceGemmMultipleD<ALayout,
                         StrideB,
                         StrideDs,
                         StrideC,
-                        1,
+                        KBatch,
                         a_element_op,
                         b_element_op,
                         c_element_op};
@@ -514,6 +696,7 @@ struct DeviceGemmMultiD_Xdl_CShuffle_V3 : public DeviceGemmMultipleD<ALayout,
                                                       index_t StrideB,
                                                       std::array<ck::index_t, NumDTensor> StrideDs,
                                                       index_t StrideC,
+                                                      index_t KBatch,
                                                       AElementwiseOperation a_element_op,
                                                       BElementwiseOperation b_element_op,
                                                       CElementwiseOperation c_element_op) override
@@ -529,7 +712,7 @@ struct DeviceGemmMultiD_Xdl_CShuffle_V3 : public DeviceGemmMultipleD<ALayout,
                                           StrideB,
                                           StrideDs,
                                           StrideC,
-                                          1,
+                                          KBatch,
                                           a_element_op,
                                           b_element_op,
                                           c_element_op);

include/ck/tensor_operation/gpu/device/impl/device_gemm_xdl_cshuffle_v3.hpp

@@ -168,15 +168,11 @@ struct DeviceGemm_Xdl_CShuffleV3 : public DeviceGemmV2<ALayout,
                         // rotating mem
                         rotating_mem.Next();
                         // clear c mem
-                        if constexpr(!is_same<remove_cvref_t<CDataType>, bhalf_t>::value)
-                        {
-                            if(arg_.KBatch > 1)
-                                hipGetErrorString(
-                                    hipMemsetAsync(arg_.p_c_grid,
-                                                   0,
-                                                   arg_.M * arg_.N * sizeof(CDataType),
-                                                   stream_config.stream_id_));
-                        }
+                        if(arg_.KBatch > 1)
+                            hipGetErrorString(hipMemsetAsync(arg_.p_c_grid,
+                                                             0,
+                                                             arg_.M * arg_.N * sizeof(CDataType),
+                                                             stream_config.stream_id_));
                     };
 
                     ave_time = ck::utility::launch_and_time_kernel_with_preprocess<false>(
@@ -190,14 +186,11 @@ struct DeviceGemm_Xdl_CShuffleV3 : public DeviceGemmV2<ALayout,
                 }
                 else
                 {
-                    if constexpr(!is_same<remove_cvref_t<CDataType>, bhalf_t>::value)
-                    {
-                        if(arg.KBatch > 1)
-                            hipGetErrorString(hipMemsetAsync(arg.p_c_grid,
-                                                             0,
-                                                             arg.M * arg.N * sizeof(CDataType),
-                                                             stream_config.stream_id_));
-                    }
+                    if(arg.KBatch > 1)
+                        hipGetErrorString(hipMemsetAsync(arg.p_c_grid,
+                                                         0,
+                                                         arg.M * arg.N * sizeof(CDataType),
+                                                         stream_config.stream_id_));
 
                     ave_time = launch_and_time_kernel(
                         stream_config, kernel, dim3(gdx, gdy, gdz), dim3(BlockSize), 0, arg);
@@ -215,15 +208,12 @@ struct DeviceGemm_Xdl_CShuffleV3 : public DeviceGemmV2<ALayout,
                 {
                     if(arg.KBatch > 1)
                     {
-                        if constexpr(!is_same<remove_cvref_t<CDataType>, bhalf_t>::value)
-                        {
-                            const auto kernel =
-                                kernel_gemm_xdl_cshuffle_v3<GridwiseGemm,
-                                                            true,
-                                                            InMemoryDataOperationEnum::AtomicAdd,
-                                                            minimum_occupancy>;
-                            Run(kernel);
-                        }
+                        const auto kernel =
+                            kernel_gemm_xdl_cshuffle_v3<GridwiseGemm,
+                                                        true,
+                                                        InMemoryDataOperationEnum::AtomicAdd,
+                                                        minimum_occupancy>;
+                        Run(kernel);
                     }
                     else
                     {
@@ -240,118 +230,113 @@ struct DeviceGemm_Xdl_CShuffleV3 : public DeviceGemmV2<ALayout,
                 {
                     if(arg.KBatch > 1)
                     {
-                        if constexpr(!is_same<remove_cvref_t<CDataType>, bhalf_t>::value)
+                        if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::One)
                         {
-                            if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::One)
+                            const auto kernel =
+                                kernel_gemm_xdl_cshuffle_v3<GridwiseGemm,
+                                                            true,
+                                                            InMemoryDataOperationEnum::AtomicAdd,
+                                                            minimum_occupancy,
+                                                            TailNumber::One>;
+                            Run(kernel);
+                        }
+                        else if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) ==
+                                TailNumber::Full)
+                        {
+                            const auto kernel =
+                                kernel_gemm_xdl_cshuffle_v3<GridwiseGemm,
+                                                            true,
+                                                            InMemoryDataOperationEnum::AtomicAdd,
+                                                            minimum_occupancy,
+                                                            TailNumber::Full>;
+                            Run(kernel);
+                        }
+
+                        if constexpr(GridwiseGemm::BlockwiseGemmPipe::PrefetchStages > 2)
+                        {
+                            if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Two)
                             {
                                 const auto kernel = kernel_gemm_xdl_cshuffle_v3<
                                     GridwiseGemm,
                                     true,
                                     InMemoryDataOperationEnum::AtomicAdd,
                                     minimum_occupancy,
-                                    TailNumber::One>;
+                                    TailNumber::Two>;
                                 Run(kernel);
                             }
-                            else if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) ==
-                                    TailNumber::Full)
+                        }
+
+                        if constexpr(GridwiseGemm::BlockwiseGemmPipe::PrefetchStages > 3)
+                        {
+                            if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) ==
+                               TailNumber::Three)
                             {
                                 const auto kernel = kernel_gemm_xdl_cshuffle_v3<
                                     GridwiseGemm,
                                     true,
                                     InMemoryDataOperationEnum::AtomicAdd,
                                     minimum_occupancy,
-                                    TailNumber::Full>;
+                                    TailNumber::Three>;
                                 Run(kernel);
                             }
+                        }
 
-                            if constexpr(GridwiseGemm::BlockwiseGemmPipe::PrefetchStages > 2)
+                        if constexpr(GridwiseGemm::BlockwiseGemmPipe::PrefetchStages > 4)
+                        {
+                            if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) ==
+                               TailNumber::Four)
                             {
-                                if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) ==
-                                   TailNumber::Two)
-                                {
-                                    const auto kernel = kernel_gemm_xdl_cshuffle_v3<
-                                        GridwiseGemm,
-                                        true,
-                                        InMemoryDataOperationEnum::AtomicAdd,
-                                        minimum_occupancy,
-                                        TailNumber::Two>;
-                                    Run(kernel);
-                                }
+                                const auto kernel = kernel_gemm_xdl_cshuffle_v3<
+                                    GridwiseGemm,
+                                    true,
+                                    InMemoryDataOperationEnum::AtomicAdd,
+                                    minimum_occupancy,
+                                    TailNumber::Four>;
+                                Run(kernel);
                             }
+                        }
 
-                            if constexpr(GridwiseGemm::BlockwiseGemmPipe::PrefetchStages > 3)
+                        if constexpr(GridwiseGemm::BlockwiseGemmPipe::PrefetchStages > 5)
+                        {
+                            if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) ==
+                               TailNumber::Five)
                             {
-                                if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) ==
-                                   TailNumber::Three)
-                                {
-                                    const auto kernel = kernel_gemm_xdl_cshuffle_v3<
-                                        GridwiseGemm,
-                                        true,
-                                        InMemoryDataOperationEnum::AtomicAdd,
-                                        minimum_occupancy,
-                                        TailNumber::Three>;
-                                    Run(kernel);
-                                }
+                                const auto kernel = kernel_gemm_xdl_cshuffle_v3<
+                                    GridwiseGemm,
+                                    true,
+                                    InMemoryDataOperationEnum::AtomicAdd,
+                                    minimum_occupancy,
+                                    TailNumber::Five>;
+                                Run(kernel);
                             }
+                        }
 
-                            if constexpr(GridwiseGemm::BlockwiseGemmPipe::PrefetchStages > 4)
+                        if constexpr(GridwiseGemm::BlockwiseGemmPipe::PrefetchStages > 6)
+                        {
+                            if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Six)
                             {
-                                if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) ==
-                                   TailNumber::Four)
-                                {
-                                    const auto kernel = kernel_gemm_xdl_cshuffle_v3<
-                                        GridwiseGemm,
-                                        true,
-                                        InMemoryDataOperationEnum::AtomicAdd,
-                                        minimum_occupancy,
-                                        TailNumber::Four>;
-                                    Run(kernel);
-                                }
+                                const auto kernel = kernel_gemm_xdl_cshuffle_v3<
+                                    GridwiseGemm,
+                                    true,
+                                    InMemoryDataOperationEnum::AtomicAdd,
+                                    minimum_occupancy,
+                                    TailNumber::Six>;
+                                Run(kernel);
                             }
+                        }
 
-                            if constexpr(GridwiseGemm::BlockwiseGemmPipe::PrefetchStages > 5)
+                        if constexpr(GridwiseGemm::BlockwiseGemmPipe::PrefetchStages > 7)
+                        {
+                            if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) ==
+                               TailNumber::Seven)
                             {
-                                if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) ==
-                                   TailNumber::Five)
-                                {
-                                    const auto kernel = kernel_gemm_xdl_cshuffle_v3<
-                                        GridwiseGemm,
-                                        true,
-                                        InMemoryDataOperationEnum::AtomicAdd,
-                                        minimum_occupancy,
-                                        TailNumber::Five>;
-                                    Run(kernel);
-                                }
-                            }
-
-                            if constexpr(GridwiseGemm::BlockwiseGemmPipe::PrefetchStages > 6)
-                            {
-                                if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) ==
-                                   TailNumber::Six)
-                                {
-                                    const auto kernel = kernel_gemm_xdl_cshuffle_v3<
-                                        GridwiseGemm,
-                                        true,
-                                        InMemoryDataOperationEnum::AtomicAdd,
-                                        minimum_occupancy,
-                                        TailNumber::Six>;
-                                    Run(kernel);
-                                }
-                            }
-
-                            if constexpr(GridwiseGemm::BlockwiseGemmPipe::PrefetchStages > 7)
-                            {
-                                if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) ==
-                                   TailNumber::Seven)
-                                {
-                                    const auto kernel = kernel_gemm_xdl_cshuffle_v3<
-                                        GridwiseGemm,
-                                        true,
-                                        InMemoryDataOperationEnum::AtomicAdd,
-                                        minimum_occupancy,
-                                        TailNumber::Seven>;
-                                    Run(kernel);
-                                }
+                                const auto kernel = kernel_gemm_xdl_cshuffle_v3<
+                                    GridwiseGemm,
+                                    true,
+                                    InMemoryDataOperationEnum::AtomicAdd,
+                                    minimum_occupancy,
+                                    TailNumber::Seven>;
+                                Run(kernel);
                             }
                         }
                     }
@@ -473,28 +458,25 @@ struct DeviceGemm_Xdl_CShuffleV3 : public DeviceGemmV2<ALayout,
                 {
                     if(arg.KBatch > 1)
                     {
-                        if constexpr(!is_same<remove_cvref_t<CDataType>, bhalf_t>::value)
+                        if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Odd)
                         {
-                            if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Odd)
-                            {
-                                const auto kernel = kernel_gemm_xdl_cshuffle_v3_2lds<
-                                    GridwiseGemm,
-                                    true,
-                                    InMemoryDataOperationEnum::AtomicAdd,
-                                    minimum_occupancy,
-                                    TailNumber::Odd>;
-                                Run(kernel);
-                            }
-                            else
-                            {
-                                const auto kernel = kernel_gemm_xdl_cshuffle_v3_2lds<
-                                    GridwiseGemm,
-                                    true,
-                                    InMemoryDataOperationEnum::AtomicAdd,
-                                    minimum_occupancy,
-                                    TailNumber::Even>;
-                                Run(kernel);
-                            }
+                            const auto kernel = kernel_gemm_xdl_cshuffle_v3_2lds<
+                                GridwiseGemm,
+                                true,
+                                InMemoryDataOperationEnum::AtomicAdd,
+                                minimum_occupancy,
+                                TailNumber::Odd>;
+                            Run(kernel);
+                        }
+                        else
+                        {
+                            const auto kernel = kernel_gemm_xdl_cshuffle_v3_2lds<
+                                GridwiseGemm,
+                                true,
+                                InMemoryDataOperationEnum::AtomicAdd,
+                                minimum_occupancy,
+                                TailNumber::Even>;
+                            Run(kernel);
                         }
                     }
                     else
@@ -525,28 +507,25 @@ struct DeviceGemm_Xdl_CShuffleV3 : public DeviceGemmV2<ALayout,
                 {
                     if(arg.KBatch > 1)
                     {
-                        if constexpr(!is_same<remove_cvref_t<CDataType>, bhalf_t>::value)
+                        if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Odd)
                         {
-                            if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Odd)
-                            {
-                                const auto kernel = kernel_gemm_xdl_cshuffle_v3<
-                                    GridwiseGemm,
-                                    true,
-                                    InMemoryDataOperationEnum::AtomicAdd,
-                                    minimum_occupancy,
-                                    TailNumber::Odd>;
-                                Run(kernel);
-                            }
-                            else
-                            {
-                                const auto kernel = kernel_gemm_xdl_cshuffle_v3<
-                                    GridwiseGemm,
-                                    true,
-                                    InMemoryDataOperationEnum::AtomicAdd,
-                                    minimum_occupancy,
-                                    TailNumber::Even>;
-                                Run(kernel);
-                            }
+                            const auto kernel =
+                                kernel_gemm_xdl_cshuffle_v3<GridwiseGemm,
+                                                            true,
+                                                            InMemoryDataOperationEnum::AtomicAdd,
+                                                            minimum_occupancy,
+                                                            TailNumber::Odd>;
+                            Run(kernel);
+                        }
+                        else
+                        {
+                            const auto kernel =
+                                kernel_gemm_xdl_cshuffle_v3<GridwiseGemm,
+                                                            true,
+                                                            InMemoryDataOperationEnum::AtomicAdd,
+                                                            minimum_occupancy,
+                                                            TailNumber::Even>;
+                            Run(kernel);
                         }
                     }
                     else
@@ -579,18 +558,14 @@ struct DeviceGemm_Xdl_CShuffleV3 : public DeviceGemmV2<ALayout,
                 // Tail number always 1
                 if constexpr(BlkGemmPipelineVer == BlockGemmPipelineVersion::v1)
                 {
-
                     if(arg.KBatch > 1)
                     {
-                        if constexpr(!is_same<remove_cvref_t<CDataType>, bhalf_t>::value)
-                        {
-                            const auto kernel =
-                                kernel_gemm_xdl_cshuffle_v3<GridwiseGemm,
-                                                            false,
-                                                            InMemoryDataOperationEnum::AtomicAdd,
-                                                            minimum_occupancy>;
-                            Run(kernel);
-                        }
+                        const auto kernel =
+                            kernel_gemm_xdl_cshuffle_v3<GridwiseGemm,
+                                                        false,
+                                                        InMemoryDataOperationEnum::AtomicAdd,
+                                                        minimum_occupancy>;
+                        Run(kernel);
                     }
                     else
                     {
@@ -628,6 +603,11 @@ struct DeviceGemm_Xdl_CShuffleV3 : public DeviceGemmV2<ALayout,
             return false;
         }
 
+        if(!is_bf16_atomic_supported() && std::is_same_v<CDataType, ck::bhalf_t> && arg.KBatch > 1)
+        {
+            return false;
+        }
+
         if((arg.K % AK1 != 0 || arg.K % BK1 != 0) && !(GemmSpec == GemmSpecialization::MKPadding ||
                                                        GemmSpec == GemmSpecialization::NKPadding ||
                                                        GemmSpec == GemmSpecialization::MNKPadding ||

include/ck/tensor_operation/gpu/device/impl/device_reduce_common.hpp

@@ -1,5 +1,5 @@
 // SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
 
 #pragma once
 
@@ -19,7 +19,7 @@ namespace device {
 template <index_t Rank, int NumReduceDim>
 std::pair<long_index_t, long_index_t> get_2d_lengths(const std::vector<index_t>& inLengths)
 {
-    static_assert(Rank <= 6, "bigger Rank size not supported!");
+    static_assert(Rank <= 12, "bigger Rank size not supported!");
 
     long_index_t invariant_total_length = 1;
     long_index_t reduce_total_length    = 1;
@@ -38,7 +38,7 @@ std::pair<long_index_t, long_index_t> get_2d_lengths(const std::vector<index_t>&
 template <index_t Rank, int NumReduceDim>
 std::pair<long_index_t, long_index_t> get_2d_lengths(const std::array<index_t, Rank>& inLengths)
 {
-    static_assert(Rank <= 6, "bigger Rank size not supported!");
+    static_assert(Rank <= 12, "bigger Rank size not supported!");
 
     long_index_t invariant_total_length = 1;
     long_index_t reduce_total_length    = 1;

include/ck/tensor_operation/gpu/device/impl/device_reduce_multiblock.hpp

@@ -1,5 +1,5 @@
 // SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
 
 #pragma once
 
@@ -51,7 +51,7 @@ struct DeviceReduceMultiBlock : public DeviceReduce<InDataType,
                                                     PropagateNan,
                                                     OutputIndex>
 {
-    static_assert(Rank <= 6, "Bigger Rank size is not supported!");
+    static_assert(Rank <= 12, "Bigger Rank size is not supported!");
     static_assert(BlockSize == MThreadClusterSize * KThreadClusterSize,
                   "Invalid thread cluster size assignments!");
 

include/ck/tensor_operation/gpu/device/impl/device_reduce_threadwise.hpp

@@ -1,5 +1,5 @@
 // SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
 
 #pragma once
 
@@ -47,7 +47,7 @@ struct DeviceReduceThreadWise : public DeviceReduce<InDataType,
                                                     OutputIndex>
 
 {
-    static_assert(Rank <= 6, "Bigger Rank size is not supported!");
+    static_assert(Rank <= 12, "Bigger Rank size is not supported!");
 
     static_assert(((InSrcVectorDim == 0 && MThreadSliceSize % InSrcVectorSize == 0) ||
                    (InSrcVectorDim == 1 && KThreadSliceSize % InSrcVectorSize == 0)) &&

include/ck/tensor_operation/gpu/device/impl/device_reduce_threadwise_multi_d.hpp

@@ -45,7 +45,7 @@ struct DeviceReduceThreadWiseMultiD : public DeviceReduceMultiD<InDataType,
                                                                 OutElementwiseOperation>
 
 {
-    static_assert(Rank <= 6, "Bigger Rank size is not supported!");
+    static_assert(Rank <= 12, "Bigger Rank size is not supported!");
 
     static_assert(((InSrcVectorDim == 0 && MThreadSliceSize % InSrcVectorSize == 0) ||
                    (InSrcVectorDim == 1 && KThreadSliceSize % InSrcVectorSize == 0)) &&

include/ck/tensor_operation/gpu/element/combined_element_wise_operation.hpp

@@ -3,7 +3,6 @@
 
 #pragma once
 
-#include "ck/utility/data_type.hpp"
 #include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"
 
 namespace ck {
@@ -107,6 +106,9 @@ struct TrinaryWithUnaryCombinedOp
     UnaryOp2 unary_op2_{};
 };
 
+using ScaleScalePass = UnaryCombinedOp<Scale, Scale, PassThrough>;
+using ScaleScaleRelu = UnaryCombinedOp<Scale, Scale, Relu>;
+
 } // namespace element_wise
 } // namespace tensor_operation
 } // namespace ck

include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3.hpp

@@ -417,6 +417,13 @@ struct GridwiseGemm_xdl_cshuffle_v3
             }
         }();
 
+        // pad M and N
+        return transform_tensor_descriptor(c_grid_desc_mraw_nraw,
+                                           make_tuple(make_right_pad_transform(M, MPad - M),
+                                                      make_right_pad_transform(N, NPad - N)),
+                                           make_tuple(Sequence<0>{}, Sequence<1>{}),
+                                           make_tuple(Sequence<0>{}, Sequence<1>{}));
+#if 0
         using GemmSpecialization = tensor_operation::device::GemmSpecialization;
 
         if constexpr(GemmSpec == GemmSpecialization::MNPadding ||
@@ -454,6 +461,7 @@ struct GridwiseGemm_xdl_cshuffle_v3
             // not pad M or N
             return c_grid_desc_mraw_nraw;
         }
+#endif
     }
 
     struct Problem
@@ -953,7 +961,8 @@ struct GridwiseGemm_xdl_cshuffle_v3
         if constexpr(!(GemmSpec == tensor_operation::device::GemmSpecialization::MPadding ||
                        GemmSpec == tensor_operation::device::GemmSpecialization::MNPadding ||
                        GemmSpec == tensor_operation::device::GemmSpecialization::MKPadding ||
-                       GemmSpec == tensor_operation::device::GemmSpecialization::MNKPadding))
+                       GemmSpec == tensor_operation::device::GemmSpecialization::MNKPadding) &&
+                     !(is_same<tensor_layout::gemm::RowMajor, ALayout>::value))
         {
             if(!(karg.M % MPerBlock == 0))
             {
@@ -970,7 +979,8 @@ struct GridwiseGemm_xdl_cshuffle_v3
         if constexpr(!(GemmSpec == tensor_operation::device::GemmSpecialization::NPadding ||
                        GemmSpec == tensor_operation::device::GemmSpecialization::MNPadding ||
                        GemmSpec == tensor_operation::device::GemmSpecialization::NKPadding ||
-                       GemmSpec == tensor_operation::device::GemmSpecialization::MNKPadding))
+                       GemmSpec == tensor_operation::device::GemmSpecialization::MNKPadding) &&
+                     (is_same<tensor_layout::gemm::RowMajor, BLayout>::value))
         {
             if(!(karg.N % NPerBlock == 0))
             {
@@ -1105,7 +1115,9 @@ struct GridwiseGemm_xdl_cshuffle_v3
         }
 
         if constexpr(!(is_same<remove_cvref_t<CDataType>, half_t>::value ||
-                       is_same<remove_cvref_t<CDataType>, float>::value))
+                       is_same<remove_cvref_t<CDataType>, float>::value ||
+                       is_same<remove_cvref_t<CDataType>, bhalf_t>::value ||
+                       is_same<remove_cvref_t<CDataType>, int32_t>::value))
         {
             if(!karg.IsReduceAdd())
             {

include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3_multi_d.hpp

@@ -36,10 +36,9 @@ __global__ void
     __launch_bounds__(CK_MAX_THREAD_PER_BLOCK, MinimumOccupancy)
 #endif
     // __attribute__((amdgpu_waves_per_eu(1, 1)))
-    kernel_gemm_xdl_cshuffle_v3(typename GridwiseGemm::Argument karg)
+    kernel_gemm_xdl_cshuffle_v3_multi_d(typename GridwiseGemm::Argument karg)
 {
-#if(!defined(__HIP_DEVICE_COMPILE__) || defined(__gfx908__) || defined(__gfx90a__) || \
-    defined(__gfx940__) || defined(__gfx941__) || defined(__gfx942__))
+#if(!defined(__HIP_DEVICE_COMPILE__) || defined(__gfx9__))
     __shared__ char p_shared[GridwiseGemm::GetSharedMemoryNumberOfByte()];
 
     auto splitk_batch_offset = typename GridwiseGemm::SplitKBatchOffset(karg);
@@ -56,7 +55,7 @@ __global__ void
         karg.c_element_op);
 #else
     ignore = karg;
-#endif // end of if (defined(__gfx908__) || defined(__gfx90a__))
+#endif // end of if (defined(__gfx9__))
 }
 
 template <typename GridwiseGemm,
@@ -69,10 +68,9 @@ __global__ void
     __launch_bounds__(CK_MAX_THREAD_PER_BLOCK, MinimumOccupancy)
 #endif
     // __attribute__((amdgpu_waves_per_eu(1, 1)))
-    kernel_gemm_xdl_cshuffle_v3_2lds(typename GridwiseGemm::Argument karg)
+    kernel_gemm_xdl_cshuffle_v3_multi_d_2lds(typename GridwiseGemm::Argument karg)
 {
-#if(!defined(__HIP_DEVICE_COMPILE__) || defined(__gfx908__) || defined(__gfx90a__) || \
-    defined(__gfx940__) || defined(__gfx941__) || defined(__gfx942__))
+#if(!defined(__HIP_DEVICE_COMPILE__) || defined(__gfx9__))
     // Pass two lds pointer is the key to tell compiler that ds_read/write
     // operate on different lds chunk at same time without order dependecy
     __shared__ char p_shared_0[GridwiseGemm::GetSharedMemoryNumberOfByte()];
@@ -93,7 +91,7 @@ __global__ void
         karg.c_element_op);
 #else
     ignore = karg;
-#endif // end of if (defined(__gfx908__) || defined(__gfx90a__))
+#endif // end of if (defined(__gfx9__))
 }
 
 template <typename ALayout,
@@ -454,6 +452,13 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3
             }
         }();
 
+        // pad M and N
+        return transform_tensor_descriptor(c_grid_desc_mraw_nraw,
+                                           make_tuple(make_right_pad_transform(M, MPad - M),
+                                                      make_right_pad_transform(N, NPad - N)),
+                                           make_tuple(Sequence<0>{}, Sequence<1>{}),
+                                           make_tuple(Sequence<0>{}, Sequence<1>{}));
+#if 0
         using GemmSpecialization = tensor_operation::device::GemmSpecialization;
 
         if constexpr(GemmSpec == GemmSpecialization::MNPadding ||
@@ -491,6 +496,7 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3
             // not pad M or N
             return c_grid_desc_mraw_nraw;
         }
+#endif
     }
 
     __host__ __device__ static auto MakeDsGridDescriptor_M_N(
@@ -1016,7 +1022,8 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3
         if constexpr(!(GemmSpec == tensor_operation::device::GemmSpecialization::MPadding ||
                        GemmSpec == tensor_operation::device::GemmSpecialization::MNPadding ||
                        GemmSpec == tensor_operation::device::GemmSpecialization::MKPadding ||
-                       GemmSpec == tensor_operation::device::GemmSpecialization::MNKPadding))
+                       GemmSpec == tensor_operation::device::GemmSpecialization::MNKPadding) &&
+                     !(is_same<tensor_layout::gemm::RowMajor, ALayout>::value))
         {
             if(!(karg.M % MPerBlock == 0))
             {
@@ -1033,7 +1040,8 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3
         if constexpr(!(GemmSpec == tensor_operation::device::GemmSpecialization::NPadding ||
                        GemmSpec == tensor_operation::device::GemmSpecialization::MNPadding ||
                        GemmSpec == tensor_operation::device::GemmSpecialization::NKPadding ||
-                       GemmSpec == tensor_operation::device::GemmSpecialization::MNKPadding))
+                       GemmSpec == tensor_operation::device::GemmSpecialization::MNKPadding) &&
+                     (is_same<tensor_layout::gemm::RowMajor, BLayout>::value))
         {
             if(!(karg.N % NPerBlock == 0))
             {

include/ck/utility/amd_buffer_addressing.hpp

@@ -562,6 +562,34 @@ __device__ void amd_buffer_store_impl(const typename vector_type<T, N>::type src
                                                         dst_wave_addr_offset);
 }
 
+template <typename T, index_t N>
+__device__ void amd_global_atomic_add_impl(const typename vector_type<T, N>::type src_thread_data,
+                                           T* addr)
+{
+    static_assert((is_same<T, bhalf_t>::value && (N == 2 || N == 4 || N == 8)) ||
+                      (is_same<T, half_t>::value && (N == 2 || N == 4 || N == 8)),
+                  "wrong! not implemented");
+
+    if constexpr(is_same<T, half_t>::value)
+    {
+        vector_type<half_t, N> tmp{src_thread_data};
+        static_for<0, N / 2, 1>{}([&](auto i) {
+            __builtin_amdgcn_global_atomic_fadd_v2f16(bit_cast<half2_t*>(addr) + i,
+                                                      tmp.template AsType<half2_t>()[i]);
+        });
+    }
+#if defined(__gfx942__)
+    else if constexpr(is_same<T, bhalf_t>::value)
+    {
+        vector_type<bhalf_t, N> tmp{src_thread_data};
+        static_for<0, N / 2, 1>{}([&](auto i) {
+            __builtin_amdgcn_global_atomic_fadd_v2bf16(bit_cast<bhalf2_t*>(addr) + i,
+                                                       tmp.template AsType<bhalf2_t>()[i]);
+        });
+    }
+#endif
+}
+
 template <typename T, index_t N>
 __device__ void amd_buffer_atomic_add_impl(const typename vector_type<T, N>::type src_thread_data,
                                            int32x4_t dst_wave_buffer_resource,
@@ -907,18 +935,29 @@ amd_buffer_atomic_add(const typename vector_type_maker<T, N>::type::type src_thr
     using scalar_t                = typename scalar_type<vector_t>::type;
     constexpr index_t vector_size = scalar_type<vector_t>::vector_size;
 
-#if CK_EXPERIMENTAL_USE_BUFFER_ATOMIC_ADD_OOB_CHECK_OFFSET_TRICK
-    uint32_t dst_addr_shift = dst_thread_element_valid ? 0 : 0x80000000;
-
-    amd_buffer_atomic_add_impl<scalar_t, vector_size>(
-        src_thread_data, dst_wave_buffer_resource, dst_addr_shift + dst_thread_addr_offset, 0);
-#else
-    if(dst_thread_element_valid)
+    if constexpr(is_same<T, bhalf_t>::value)
     {
-        amd_buffer_atomic_add_impl<scalar_t, vector_size>(
-            src_thread_data, dst_wave_buffer_resource, dst_thread_addr_offset, 0);
+        if(dst_thread_element_valid)
+        {
+            amd_global_atomic_add_impl<scalar_t, vector_size>(
+                src_thread_data, p_dst_wave + dst_thread_element_offset);
+        }
     }
+    else
+    {
+#if CK_EXPERIMENTAL_USE_BUFFER_ATOMIC_ADD_OOB_CHECK_OFFSET_TRICK
+        uint32_t dst_addr_shift = dst_thread_element_valid ? 0 : 0x80000000;
+
+        amd_buffer_atomic_add_impl<scalar_t, vector_size>(
+            src_thread_data, dst_wave_buffer_resource, dst_addr_shift + dst_thread_addr_offset, 0);
+#else
+        if(dst_thread_element_valid)
+        {
+            amd_buffer_atomic_add_impl<scalar_t, vector_size>(
+                src_thread_data, dst_wave_buffer_resource, dst_thread_addr_offset, 0);
+        }
 #endif
+    }
 }
 
 // buffer_atomic_max requires:

include/ck/utility/dynamic_buffer.hpp

@@ -358,13 +358,15 @@ struct DynamicBuffer
         bool constexpr use_amd_buffer_addressing =
             is_same_v<remove_cvref_t<scalar_t>, int32_t> ||
             is_same_v<remove_cvref_t<scalar_t>, float> ||
-            (is_same_v<remove_cvref_t<scalar_t>, half_t> && scalar_per_x_vector % 2 == 0);
+            (is_same_v<remove_cvref_t<scalar_t>, half_t> && scalar_per_x_vector % 2 == 0) ||
+            (is_same_v<remove_cvref_t<scalar_t>, bhalf_t> && scalar_per_x_vector % 2 == 0);
 #elif CK_USE_AMD_BUFFER_ATOMIC_ADD_INTEGER && (!CK_USE_AMD_BUFFER_ATOMIC_ADD_FLOAT)
         bool constexpr use_amd_buffer_addressing = is_same_v<remove_cvref_t<scalar_t>, int32_t>;
 #elif(!CK_USE_AMD_BUFFER_ATOMIC_ADD_INTEGER) && CK_USE_AMD_BUFFER_ATOMIC_ADD_FLOAT
         bool constexpr use_amd_buffer_addressing =
             is_same_v<remove_cvref_t<scalar_t>, float> ||
-            (is_same_v<remove_cvref_t<scalar_t>, half_t> && scalar_per_x_vector % 2 == 0);
+            (is_same_v<remove_cvref_t<scalar_t>, half_t> && scalar_per_x_vector % 2 == 0) ||
+            (is_same_v<remove_cvref_t<scalar_t>, bhalf_t> && scalar_per_x_vector % 2 == 0);
 #else
         bool constexpr use_amd_buffer_addressing = false;
 #endif

include/ck_tile/core/algorithm/coordinate_transform.hpp

@@ -1341,7 +1341,7 @@ struct modulo : public base_transform<1, 1>
 };
 
 // 2D XOR, NOTE: "xor" is a keyword
-template <typename LowLengths, typename RightShift>
+template <typename LowLengths>
 struct xor_t : public base_transform<2, 2>
 {
     static constexpr auto type_enum = coord_transform_enum::xor_t;
@@ -1352,15 +1352,10 @@ struct xor_t : public base_transform<2, 2>
     using UpLengths = LowLengths;
 
     UpLengths up_lengths_;
-    RightShift right_shift_;
 
-    CK_TILE_HOST_DEVICE constexpr xor_t() : up_lengths_{}, right_shift_{} {}
+    CK_TILE_HOST_DEVICE constexpr xor_t() : up_lengths_{} {}
 
-    CK_TILE_HOST_DEVICE constexpr xor_t(const LowLengths& low_lengths,
-                                        const RightShift& right_shift)
-        : up_lengths_{low_lengths}, right_shift_{right_shift}
-    {
-    }
+    CK_TILE_HOST_DEVICE constexpr xor_t(const LowLengths& low_lengths) : up_lengths_{low_lengths} {}
 
     CK_TILE_HOST_DEVICE static constexpr auto get_type_enum()
     {
@@ -1378,13 +1373,8 @@ struct xor_t : public base_transform<2, 2>
 
         idx_low(number<0>{}) = idx_up[number<0>{}];
 
-        const auto idx_low_1_tmp =
-            (idx_up[number<1>{}] - idx_up[number<0>{}] * right_shift_) % up_lengths_[number<1>{}];
-
-        const auto idx_low_1 =
-            (idx_low_1_tmp >= 0) ? idx_low_1_tmp : up_lengths_[number<1>{}] + idx_low_1_tmp;
-
-        idx_low(number<1>{}) = idx_low_1;
+        idx_low(number<1>{}) =
+            idx_up[number<1>{}] ^ (idx_up[number<0>{}] % up_lengths_[number<1>{}]);
     }
 
     template <typename LowIdxDiff, typename UpIdxDiff, typename LowIdx, typename UpIdx>
@@ -1419,8 +1409,7 @@ struct xor_t : public base_transform<2, 2>
 
     CK_TILE_HOST_DEVICE static constexpr bool is_known_at_compile_time()
     {
-        return ck_tile::is_known_at_compile_time<UpLengths>::value &&
-               ck_tile::is_known_at_compile_time<RightShift>::value;
+        return ck_tile::is_known_at_compile_time<UpLengths>::value;
     }
 
     // MUST be static function
@@ -1432,14 +1421,6 @@ struct xor_t : public base_transform<2, 2>
         array<index_t, 2> up_vector_lengths = low_vector_lengths;
         array<index_t, 2> up_vector_strides = low_vector_strides;
 
-        if constexpr(ck_tile::is_known_at_compile_time<RightShift>::value)
-        {
-            if(low_vector_lengths[1] != -1)
-            {
-                up_vector_lengths(1) = gcd(low_vector_lengths[1], abs(right_shift_));
-            }
-        }
-
         return make_tuple(up_vector_lengths, up_vector_strides);
     }
 
@@ -1452,10 +1433,6 @@ struct xor_t : public base_transform<2, 2>
         print(up_lengths_);
         printf(", ");
 
-        //
-        printf("right_shift_: ");
-        print(right_shift_);
-
         printf("}");
     }
 };
@@ -1655,11 +1632,10 @@ CK_TILE_HOST_DEVICE constexpr auto make_modulo_transform(const Modulus& modulus,
     return modulo<Modulus, UpLength>{modulus, up_length};
 }
 
-template <typename LowLengths, typename RightShift>
-CK_TILE_HOST_DEVICE constexpr auto make_xor_transform(const LowLengths& low_lengths,
-                                                      const RightShift& right_shift)
+template <typename LowLengths>
+CK_TILE_HOST_DEVICE constexpr auto make_xor_transform(const LowLengths& low_lengths)
 {
-    return xor_t<LowLengths, RightShift>{low_lengths, right_shift};
+    return xor_t<LowLengths>{low_lengths};
 }
 
 template <typename LowLength, typename OffsetLength>

include/ck_tile/core/numeric/vector_type.hpp

@@ -117,6 +117,15 @@ using int32x16_t = int32_t __attribute__((ext_vector_type(16)));
 using int32x32_t = int32_t __attribute__((ext_vector_type(32)));
 using int32x64_t = int32_t __attribute__((ext_vector_type(64)));
 
+// u32
+// using uint32_t = ...
+using uint32x2_t  = uint32_t __attribute__((ext_vector_type(2)));
+using uint32x4_t  = uint32_t __attribute__((ext_vector_type(4)));
+using uint32x8_t  = uint32_t __attribute__((ext_vector_type(8)));
+using uint32x16_t = uint32_t __attribute__((ext_vector_type(16)));
+using uint32x32_t = uint32_t __attribute__((ext_vector_type(32)));
+using uint32x64_t = uint32_t __attribute__((ext_vector_type(64)));
+
 // i16
 // using int16_t = ...
 using int16x2_t  = int16_t __attribute__((ext_vector_type(2)));

include/ck_tile/core/tensor/tile_distribution.hpp

@@ -746,8 +746,9 @@ CK_TILE_HOST_DEVICE constexpr auto slice_distribution_from_x(
     return make_tuple(
         make_static_tile_distribution(
             tile_distribution_encoding<typename Encoding::RsLengths,
-                                       decltype(sliced_h_lengths), // only need to change the
-                                                                   // h_lengths type
+                                       remove_cvref_t<decltype(sliced_h_lengths)>, // only need to
+                                                                                   // change the
+                                                                                   // h_lengths type
                                        typename Encoding::Ps2RHssMajor,
                                        typename Encoding::Ps2RHssMinor,
                                        typename Encoding::Ys2RHsMajor,

include/ck_tile/core/utility/philox_rand.hpp

@@ -53,6 +53,39 @@ class philox
         out_tmp[3] = tmp_ph.w;
     }
 
+    CK_TILE_HOST_DEVICE void get_random_8x8(uint8_t* out,
+                                            const unsigned long long subsequence,
+                                            const index_t start_idx) const
+    {
+        uint4 tmp_ph;
+        tmp_ph = get_philox_4x32(subsequence);
+
+        uint32x4_t tmp;
+        tmp[0]            = tmp_ph.x;
+        tmp[1]            = tmp_ph.y;
+        tmp[2]            = tmp_ph.z;
+        tmp[3]            = tmp_ph.w;
+        uint32_t* out_tmp = reinterpret_cast<uint32_t*>(&out[0]);
+        out_tmp[0]        = tmp[start_idx];
+        out_tmp[1]        = tmp[start_idx + 2];
+    }
+
+    CK_TILE_HOST_DEVICE void get_random_4x8(uint8_t* out,
+                                            const unsigned long long subsequence,
+                                            const index_t start_idx) const
+    {
+        uint4 tmp_ph;
+        tmp_ph = get_philox_4x32(subsequence);
+
+        uint32x4_t tmp;
+        tmp[0]            = tmp_ph.x;
+        tmp[1]            = tmp_ph.y;
+        tmp[2]            = tmp_ph.z;
+        tmp[3]            = tmp_ph.w;
+        uint32_t* out_tmp = reinterpret_cast<uint32_t*>(&out[0]);
+        out_tmp[0]        = tmp[start_idx];
+    }
+
     private:
     struct ull2
     {

include/ck_tile/ops/fmha.hpp

@@ -8,21 +8,16 @@
 #include "ck_tile/ops/fmha/block/block_masking.hpp"
 #include "ck_tile/ops/fmha/block/block_position_encoding.hpp"
 #include "ck_tile/ops/fmha/kernel/fmha_bwd_kernel.hpp"
-#include "ck_tile/ops/fmha/kernel/fmha_bwd_tile_partitioner.hpp"
 #include "ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp"
 #include "ck_tile/ops/fmha/kernel/fmha_fwd_splitkv_combine_kernel.hpp"
 #include "ck_tile/ops/fmha/kernel/fmha_fwd_splitkv_combine_tile_partitioner.hpp"
 #include "ck_tile/ops/fmha/kernel/fmha_fwd_splitkv_kernel.hpp"
 #include "ck_tile/ops/fmha/kernel/fmha_fwd_splitkv_tile_partitioner.hpp"
 #include "ck_tile/ops/fmha/kernel/fmha_fwd_tile_partitioner.hpp"
+#include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_convert_dq.hpp"
 #include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_dot_do_o.hpp"
-#include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_dot_do_o_default_policy.hpp"
-#include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_ks_kts_vr.hpp"
-#include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_ks_kts_vr_default_policy.hpp"
-#include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_ks_vr.hpp"
-#include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_ks_vr_default_policy.hpp"
-#include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_qs_ks_vr_dos.hpp"
-#include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_qs_ks_vr_dos_default_policy.hpp"
+#include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_kr_ktr_vr.hpp"
+#include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_kr_ktr_vr_iglp.hpp"
 #include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_default_policy.hpp"
 #include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_enum.hpp"
 #include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_problem.hpp"

include/ck_tile/ops/fmha/block/block_dropout.hpp

@@ -286,11 +286,226 @@ struct BlockDropout
         });
     }
 
+    ck_tile::philox ph;
+    const float rp_undrop;
+    const uint8_t p_undrop_in_uint8_t;
+    const bool is_store_randval;
+};
+
+template <bool IsDropout_, bool IsWG32_, bool IsStoreRandval_>
+struct BlockDropoutBwd;
+
+template <bool IsWG32_, bool IsStoreRandval_>
+struct BlockDropoutBwd<false, IsWG32_, IsStoreRandval_>
+{
+    static constexpr bool IsDropout      = false;
+    static constexpr bool IsStoreRandval = IsStoreRandval_;
+
+    template <typename BlockGemm, bool IsFwd = true, typename RandValDramBlockWindowTmp>
+    __host__ __device__ static constexpr auto
+    MakeRandvalDramWindow(RandValDramBlockWindowTmp& randval_dram_block_window_tmp,
+                          index_t seqlen_qk_start)
+    {
+        (void)randval_dram_block_window_tmp;
+        (void)seqlen_qk_start;
+
+        return make_null_tile_window(make_tuple(number<0>{}, number<0>{}));
+    }
+};
+
+template <bool IsWG32_, bool IsStoreRandval_>
+struct BlockDropoutBwd<true, IsWG32_, IsStoreRandval_>
+{
+    static constexpr bool IsDropout = true;
+    // true:  32*32 warp gemm
+    // false: 16*16 warp gemm
+    static constexpr bool IsWG32         = IsWG32_;
+    static constexpr bool IsStoreRandval = IsStoreRandval_;
+
+    CK_TILE_HOST_DEVICE BlockDropoutBwd(index_t i_batch,
+                                        index_t i_head,
+                                        index_t nheads,
+                                        unsigned long long seed,
+                                        unsigned long long offset,
+                                        float rp_undrop_,
+                                        uint8_t p_undrop_in_uint8_t_)
+        : ph(seed,
+             offset + (i_batch * nheads + i_head) * get_warp_size() +
+                 (IsWG32 ? get_lane_id() : ((get_lane_id() & 47) + ((get_warp_id() & 1) << 4)))),
+          rp_undrop(rp_undrop_),
+          p_undrop_in_uint8_t(p_undrop_in_uint8_t_)
+    {
+    }
+
+    template <typename BlockGemm, bool IsFwd = true, typename RandValDramBlockWindowTmp>
+    CK_TILE_HOST_DEVICE static constexpr auto
+    MakeRandvalDramWindow(RandValDramBlockWindowTmp& randval_dram_block_window_tmp,
+                          index_t seqlen_qk_start)
+    {
+        constexpr auto config =
+            BlockGemm::Policy::template GetWarpGemmMWarpNWarp<typename BlockGemm::Problem>();
+        using BlockGemmShape                  = remove_cvref_t<typename BlockGemm::BlockGemmShape>;
+        using WG                              = remove_cvref_t<decltype(config.template at<0>())>;
+        constexpr index_t kMPerBlock          = BlockGemmShape::kM;
+        constexpr index_t MWarp               = config.template at<1>();
+        constexpr index_t NWarp               = config.template at<2>();
+        constexpr bool MBwdWG16MultiIterCheck = (!IsFwd) && (!IsWG32) && (kMPerBlock > 16);
+        constexpr index_t kMPerStep           = [&]() {
+            if constexpr(MBwdWG16MultiIterCheck)
+            {
+                return MWarp * WG::kM * 2;
+            }
+            else
+            {
+                return MWarp * WG::kM;
+            }
+        }();
+        constexpr index_t kNPerStep = NWarp * WG::kN;
+
+        const auto block_origin  = randval_dram_block_window_tmp.get_window_origin();
+        auto randval_dram_window = [&]() {
+            if constexpr(IsFwd)
+            {
+                return make_tile_window(
+                    randval_dram_block_window_tmp.get_bottom_tensor_view(),
+                    ck_tile::make_tuple(number<kMPerStep>{}, number<kNPerStep>{}),
+                    {block_origin.at(number<0>{}), seqlen_qk_start}); // M/N
+            }
+            else
+            {
+                return make_tile_window(
+                    randval_dram_block_window_tmp.get_bottom_tensor_view(),
+                    ck_tile::make_tuple(number<kMPerStep>{}, number<kNPerStep>{}),
+                    {seqlen_qk_start, block_origin.at(number<1>{})}); // M/N
+            }
+        }();
+
+        return randval_dram_window;
+    }
+
+    template <typename BlockGemm>
+    CK_TILE_HOST_DEVICE static constexpr auto MakeRandValLdsBlockDescriptor()
+    {
+        constexpr auto config =
+            BlockGemm::Policy::template GetWarpGemmMWarpNWarp<typename BlockGemm::Problem>();
+        using WG                    = remove_cvref_t<decltype(config.template at<0>())>;
+        constexpr index_t MWarp     = config.template at<1>();
+        constexpr index_t kMPerStep = MWarp * WG::kM;
+        constexpr index_t kNPerStep = WG::kN;
+        constexpr index_t kN1       = 8;
+        constexpr index_t kN0       = kNPerStep / kN1;
+
+        constexpr auto randval_lds_block_desc_0 = make_naive_tensor_descriptor(
+            ck_tile::make_tuple(number<kN0>{}, number<kMPerStep>{}, number<kN1>{}),
+            ck_tile::make_tuple(number<(kMPerStep + 1) * kN1>{}, number<kN1>{}, number<1>{}),
+            number<kN1>{},
+            number<1>{});
+
+        constexpr auto randval_lds_block_desc = transform_tensor_descriptor(
+            randval_lds_block_desc_0,
+            ck_tile::make_tuple(
+                make_pass_through_transform(number<kMPerStep>{}),
+                make_merge_transform(ck_tile::make_tuple(number<kN0>{}, number<kN1>{}))),
+            ck_tile::make_tuple(sequence<1>{}, sequence<0, 2>{}),
+            ck_tile::make_tuple(sequence<0>{}, sequence<1>{}));
+
+        return randval_lds_block_desc;
+    }
+
+    template <typename BlockGemm, bool IsFwd = true>
+    CK_TILE_HOST_DEVICE static constexpr auto MakeRandValTileDistribution()
+    {
+        constexpr auto config =
+            BlockGemm::Policy::template GetWarpGemmMWarpNWarp<typename BlockGemm::Problem>();
+        using BlockGemmShape                  = remove_cvref_t<typename BlockGemm::BlockGemmShape>;
+        constexpr index_t kMPerBlock          = BlockGemmShape::kM;
+        constexpr index_t MWarp               = config.template at<1>();
+        constexpr index_t NWarp               = config.template at<2>();
+        constexpr bool MBwdWG16MultiIterCheck = (!IsFwd) && (!IsWG32) && (kMPerBlock > 16);
+
+        constexpr index_t MIterPerWarp = [&]() {
+            if constexpr(MBwdWG16MultiIterCheck)
+            {
+                return 2;
+            }
+            else
+            {
+                return 1;
+            }
+        }();
+        constexpr index_t NIterPerWarp = 1;
+
+        constexpr auto randval_block_outer_part_dstr_encoding = tile_distribution_encoding<
+            sequence<>,
+            tuple<sequence<MIterPerWarp, MWarp>, sequence<NIterPerWarp, NWarp>>,
+            tuple<sequence<1, 2>>,
+            tuple<sequence<1, 1>>,
+            sequence<1, 2>,
+            sequence<0, 0>>{};
+
+        // Use Bwd WarpGemm to ensure that Fwd's random values ​​are consistent with Bwd.
+        // except headdim256.
+        constexpr auto randval_block_inner_part_dstr_encoding = []() {
+            if constexpr(std::is_same_v<typename BlockGemm::ADataType, half_t> &&
+                         std::is_same_v<typename BlockGemm::BDataType, half_t> &&
+                         std::is_same_v<typename BlockGemm::CDataType, float>)
+            {
+                if constexpr(IsWG32)
+                    return typename WarpGemmMfmaF16F16F32M32N32K16SwizzleA::CWarpDstrEncoding{};
+                else
+                    return typename WarpGemmMfmaF16F16F32M16N16K16::CWarpDstrEncoding{};
+            }
+            else
+            {
+                if constexpr(IsWG32)
+                    return typename WarpGemmMfmaBf16Bf16F32M32N32K16SwizzleA::CWarpDstrEncoding{};
+                else
+                    return typename WarpGemmMfmaBf16Bf16F32M16N16K16::CWarpDstrEncoding{};
+            }
+        }();
+
+        constexpr auto randval_block_part_dstr_encode =
+            detail::make_embed_tile_distribution_encoding(randval_block_outer_part_dstr_encoding,
+                                                          randval_block_inner_part_dstr_encoding);
+
+        return make_static_tile_distribution(randval_block_part_dstr_encode);
+    }
+
+    template <typename BlockGemm>
+    CK_TILE_HOST_DEVICE static constexpr auto MakeRandValLdsShuffleTileDistribution()
+    {
+        constexpr auto config =
+            BlockGemm::Policy::template GetWarpGemmMWarpNWarp<typename BlockGemm::Problem>();
+        using WG                = remove_cvref_t<decltype(config.template at<0>())>;
+        constexpr index_t MWarp = config.template at<1>();
+        constexpr index_t NWarp = config.template at<2>();
+
+        constexpr index_t MIterPerWarp = 1;
+        constexpr index_t NIterPerWarp = 1;
+
+        constexpr auto randval_block_outer_part_dstr_encoding = tile_distribution_encoding<
+            sequence<>,
+            tuple<sequence<MIterPerWarp, MWarp>, sequence<NIterPerWarp, NWarp>>,
+            tuple<sequence<1, 2>>,
+            tuple<sequence<1, 1>>,
+            sequence<1, 2>,
+            sequence<0, 0>>{};
+
+        constexpr auto randval_block_part_dstr_encode =
+            detail::make_embed_tile_distribution_encoding(randval_block_outer_part_dstr_encoding,
+                                                          typename WG::CWarpDstrEncoding{});
+
+        return make_static_tile_distribution(randval_block_part_dstr_encode);
+    }
+
     template <typename BlockGemm,
+              typename PComputeDataType,
               typename RandValOutputDataType,
               typename PComputeWindow,
               typename RandValDramWindow>
-    CK_TILE_HOST_DEVICE void Run(const index_t start_m0_idx,
+    CK_TILE_HOST_DEVICE void Run(void* randval_ptr,
+                                 const index_t start_m0_idx,
+                                 const index_t start_n0_idx,
                                  PComputeWindow& p_compute,
                                  RandValDramWindow& randval_dram_window) const
     {
@@ -305,30 +520,177 @@ struct BlockDropout
         constexpr index_t kMPerStep  = MWarp * WG::kM;
         constexpr index_t kNPerStep  = NWarp * WG::kN;
 
-        // register distribute
-        auto randval =
-            make_static_distributed_tensor<uint8_t>(MakeRandValTileDistribution<BlockGemm>());
-        static_assert(randval.kThreadElementSpaceSize == 16);
+        // randval tile in LDS
+        auto randval_lds = make_tensor_view<address_space_enum::lds>(
+            reinterpret_cast<uint8_t*>(randval_ptr), MakeRandValLdsBlockDescriptor<BlockGemm>());
 
-        const int start_n0_idx = randval_dram_window.get_window_origin().at(number<1>{});
-        static_for<0, kNPerBlock / kNPerStep, 1>{}([&](auto i_n0) {
-            static_for<0, kMPerBlock / kMPerStep, 1>{}([&](auto i_m0) {
-                int block_row_start = (start_m0_idx / WG::kM) + i_m0;
-                int block_col_start = (start_n0_idx / WG::kN) + (i_n0 * NWarp) + get_warp_id();
+        auto randval_lds_window = make_tile_window(
+            randval_lds, MakeRandValLdsBlockDescriptor<BlockGemm>().get_lengths(), {0, 0});
+
+        // register distribute
+        auto randval_dist_generated =
+            make_static_distributed_tensor<uint8_t>(MakeRandValTileDistribution<BlockGemm>());
+        static_assert(randval_dist_generated.kThreadElementSpaceSize == 16);
+
+        auto randval_lds_read_window =
+            make_tile_window(randval_lds_window.get_bottom_tensor_view(),
+                             randval_lds_window.get_window_lengths(),
+                             randval_lds_window.get_window_origin(),
+                             MakeRandValLdsShuffleTileDistribution<BlockGemm>());
+
+        static_for<0, kMPerBlock / kMPerStep, 1>{}([&](auto i_m0) {
+            static_for<0, kNPerBlock / kNPerStep, 1>{}([&](auto i_n0) {
+                int block_row_start = (start_m0_idx / WG::kM) + (i_m0 * MWarp) + get_warp_id();
+                int block_col_start = (start_n0_idx / WG::kN) + i_n0;
                 uint2 rowcol        = make_uint2(block_row_start, block_col_start);
 
                 // generate random number
                 uint8_t random_uint8_t[16];
                 ph.get_random_16x8(random_uint8_t, reinterpret_cast<unsigned long long&>(rowcol));
 
+                constexpr auto randval_dist_generated_spans =
+                    decltype(randval_dist_generated)::get_distributed_spans();
+                int i_random_idx = 0;
+                sweep_tile_span(randval_dist_generated_spans[number<0>{}], [&](auto idx0) {
+                    sweep_tile_span(randval_dist_generated_spans[number<1>{}], [&](auto idx1) {
+                        constexpr auto i_j_idx          = ck_tile::make_tuple(idx0, idx1);
+                        randval_dist_generated(i_j_idx) = random_uint8_t[i_random_idx++];
+                    });
+                });
+                // save to LDS
+                store_tile(randval_lds_window, randval_dist_generated);
+                block_sync_lds();
+                // read from LDS to register
+                auto randval                 = load_tile(randval_lds_read_window);
+                constexpr auto randval_spans = decltype(randval)::get_distributed_spans();
+                sweep_tile_span(randval_spans[number<0>{}], [&](auto idx0) {
+                    sweep_tile_span(randval_spans[number<1>{}], [&](auto idx1) {
+                        constexpr auto p_idx0 = tile_distributed_index<i_m0>{};
+                        constexpr auto p_idx1 =
+                            tile_distributed_index<i_n0, idx1.impl_.at(1), idx1.impl_.at(2)>{};
+                        constexpr auto p_idx = ck_tile::make_tuple(p_idx0, p_idx1);
+                        constexpr auto r_idx = ck_tile::make_tuple(idx0, idx1);
+                        p_compute(p_idx)     = randval[r_idx] <= p_undrop_in_uint8_t
+                                                   ? p_compute[p_idx] * rp_undrop
+                                                   : PComputeDataType(0);
+                    });
+                });
+                // save to Global
+                if constexpr(IsStoreRandval)
+                {
+                    const auto randval_store = cast_tile<RandValOutputDataType>(randval);
+                    store_tile(randval_dram_window, randval_store);
+                    move_tile_window(randval_dram_window, {0, kNPerStep});
+                }
+            });
+            if constexpr(IsStoreRandval)
+            {
+                move_tile_window(randval_dram_window, {kMPerStep, -kNPerBlock});
+            }
+        });
+        if constexpr(IsStoreRandval)
+        {
+            move_tile_window(randval_dram_window, {-kMPerBlock, kNPerBlock});
+        }
+    }
+
+    template <typename BlockGemm,
+              typename RandValOutputDataType,
+              typename PComputeWindow,
+              typename RandValDramWindow>
+    CK_TILE_HOST_DEVICE void Run(const index_t start_m0_idx,
+                                 const index_t start_n0_idx,
+                                 PComputeWindow& p_compute,
+                                 RandValDramWindow& randval_dram_window) const
+    {
+        constexpr auto config =
+            BlockGemm::Policy::template GetWarpGemmMWarpNWarp<typename BlockGemm::Problem>();
+        using WG                               = remove_cvref_t<decltype(config.template at<0>())>;
+        constexpr index_t MWarp                = config.template at<1>();
+        constexpr index_t NWarp                = config.template at<2>();
+        using BlockGemmShape                   = remove_cvref_t<typename BlockGemm::BlockGemmShape>;
+        constexpr index_t kMPerBlock           = BlockGemmShape::kM;
+        constexpr index_t kNPerBlock           = BlockGemmShape::kN;
+        constexpr bool MBwdWG16MultiIterCheck  = (!IsWG32) && (kMPerBlock > 16);
+        constexpr bool MBwdWG16SingleIterCheck = (!IsWG32) && (kMPerBlock == 16);
+        constexpr index_t kMPerStep            = [&]() {
+            if constexpr(MBwdWG16MultiIterCheck)
+            {
+                return MWarp * WG::kM * 2;
+            }
+            else
+            {
+                return MWarp * WG::kM;
+            }
+        }();
+        constexpr index_t kNPerStep = NWarp * WG::kN;
+
+        // register distribute
+        auto randval = make_static_distributed_tensor<uint8_t>(
+            MakeRandValTileDistribution<BlockGemm, false>());
+        if constexpr(IsWG32)
+            static_assert(randval.kThreadElementSpaceSize == 16);
+        else
+            static_assert(randval.kThreadElementSpaceSize == 4 ||
+                          randval.kThreadElementSpaceSize == 8);
+
+        static_for<0, kNPerBlock / kNPerStep, 1>{}([&](auto i_n0) {
+            static_for<0, kMPerBlock / kMPerStep, 1>{}([&](auto i_m0) {
+                int block_row_start, block_col_start;
+                if constexpr(IsWG32)
+                {
+                    block_row_start = (start_m0_idx / WG::kM) + i_m0;
+                    block_col_start = (start_n0_idx / WG::kN) + (i_n0 * NWarp) + get_warp_id();
+                }
+                else
+                {
+                    block_row_start = start_m0_idx / 32 + i_m0;
+                    block_col_start = (start_n0_idx / 32) + get_warp_id() / 2 + i_n0 * 2;
+                }
+                uint2 rowcol = make_uint2(block_row_start, block_col_start);
+
+                // generate random number
+                uint8_t* random_uint8_t_;
+                if constexpr(MBwdWG16SingleIterCheck)
+                {
+                    uint8_t random_uint8_t[4];
+                    // m0t0 ~m0t15/m0t32~m0t47: 0
+                    // m0t16~m0t31/m0t48~m0t63: 1
+                    // m1t0 ~m1t15/m1t32~m1t47: 2
+                    // m1t16~m1t31/m1t48~m1t63: 3
+                    const index_t start_idx =
+                        ((get_lane_id() >> 4) & 1) + (((start_m0_idx >> 4) & 1) << 1);
+                    ph.get_random_4x8(
+                        random_uint8_t, reinterpret_cast<unsigned long long&>(rowcol), start_idx);
+                    random_uint8_t_ = random_uint8_t;
+                }
+                else if constexpr(MBwdWG16MultiIterCheck)
+                {
+                    uint8_t random_uint8_t[8];
+                    // t0 ~t15/t32~t47: 0
+                    // t16~t31/t48~t63: 1
+                    const index_t start_idx = (get_lane_id() >> 4) & 1;
+                    ph.get_random_8x8(
+                        random_uint8_t, reinterpret_cast<unsigned long long&>(rowcol), start_idx);
+                    random_uint8_t_ = random_uint8_t;
+                }
+                else
+                {
+                    uint8_t random_uint8_t[16];
+                    ph.get_random_16x8(random_uint8_t,
+                                       reinterpret_cast<unsigned long long&>(rowcol));
+                    random_uint8_t_ = random_uint8_t;
+                }
+
                 constexpr auto randval_spans = decltype(randval)::get_distributed_spans();
                 int i_random_idx             = 0;
                 sweep_tile_span(randval_spans[number<0>{}], [&](auto idx0) {
                     sweep_tile_span(randval_spans[number<1>{}], [&](auto idx1) {
-                        constexpr auto r_idx = ck_tile::make_tuple(idx0, idx1);
-                        randval(r_idx)       = random_uint8_t[i_random_idx++];
-                        constexpr auto p_idx0 =
-                            tile_distributed_index<i_m0, idx0.impl_.at(1), idx0.impl_.at(2)>{};
+                        constexpr auto r_idx  = ck_tile::make_tuple(idx0, idx1);
+                        randval(r_idx)        = random_uint8_t_[i_random_idx++];
+                        constexpr auto p_idx0 = tile_distributed_index<i_m0 + idx0.impl_.at(0),
+                                                                       idx0.impl_.at(1),
+                                                                       idx0.impl_.at(2)>{};
                         constexpr auto p_idx1 = tile_distributed_index<i_n0>{};
                         constexpr auto p_idx  = ck_tile::make_tuple(p_idx0, p_idx1);
                         p_compute(p_idx)      = randval[r_idx] <= p_undrop_in_uint8_t
@@ -337,19 +699,19 @@ struct BlockDropout
                     });
                 });
                 // save to Global
-                if(is_store_randval)
+                if constexpr(IsStoreRandval)
                 {
                     const auto randval_store = cast_tile<RandValOutputDataType>(randval);
                     store_tile(randval_dram_window, randval_store);
                     move_tile_window(randval_dram_window, {kMPerStep, 0});
                 }
             });
-            if(is_store_randval)
+            if constexpr(IsStoreRandval)
             {
                 move_tile_window(randval_dram_window, {-kMPerBlock, kNPerStep});
             }
         });
-        if(is_store_randval)
+        if constexpr(IsStoreRandval)
         {
             move_tile_window(randval_dram_window, {kMPerBlock, -kNPerBlock});
         }
@@ -358,7 +720,6 @@ struct BlockDropout
     ck_tile::philox ph;
     const float rp_undrop;
     const uint8_t p_undrop_in_uint8_t;
-    const bool is_store_randval;
 };
 
 } // namespace ck_tile

include/ck_tile/ops/fmha/kernel/fmha_bwd_tile_partitioner.hpp

@@ -1,54 +0,0 @@
-// SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
-
-#pragma once
-
-#include "ck_tile/core.hpp"
-
-namespace ck_tile {
-
-template <typename BlockFmhaShape_>
-struct FmhaBwdTilePartitioner
-{
-    using BlockFmhaShape = ck_tile::remove_cvref_t<BlockFmhaShape_>;
-
-    static constexpr ck_tile::index_t kN0 = BlockFmhaShape::kN0;
-
-    CK_TILE_HOST static constexpr auto
-    GridSize(ck_tile::index_t batch_size_, ck_tile::index_t nhead_, ck_tile::index_t seqlen_k_)
-    {
-        // TODO: this may need tuning
-        return dim3(ck_tile::integer_divide_ceil(seqlen_k_, kN0), nhead_, batch_size_);
-    }
-
-    CK_TILE_DEVICE auto operator()(ck_tile::index_t /*seqlen_k*/)
-    {
-        const index_t i_block = blockIdx.x;
-        const index_t i_nhead = blockIdx.y;
-        const index_t i_batch = blockIdx.z;
-
-        return ck_tile::make_tuple(i_block, i_nhead, i_batch);
-    }
-};
-
-template <ck_tile::index_t kBlockSize>
-struct FmhaBwdOGradDotOTilePartitioner
-{
-    CK_TILE_HOST static constexpr auto
-    GridSize(ck_tile::index_t batch_size_, ck_tile::index_t nhead_, ck_tile::index_t seqlen_q_)
-    {
-        // TODO: this may need tuning
-        return dim3(ck_tile::integer_divide_ceil(seqlen_q_, kBlockSize), nhead_, batch_size_);
-    }
-
-    CK_TILE_DEVICE auto operator()(ck_tile::index_t /*seqlen_q*/)
-    {
-        const index_t i_block = blockIdx.x;
-        const index_t i_nhead = blockIdx.y;
-        const index_t i_batch = blockIdx.z;
-
-        return ck_tile::make_tuple(i_block, i_nhead, i_batch);
-    }
-};
-
-} // namespace ck_tile

include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp

@@ -86,7 +86,7 @@ struct FmhaFwdKernel
             "w" + _TS_(gwt::at(ck_tile::number<0>{})) + "x" + _TS_(gwt::at(ck_tile::number<1>{})) + "x" + _TS_(gwt::at(ck_tile::number<2>{})) + "_" +
             (kBlockPerCuInput == -1 ? "" : ("o" + _TS_(kBlockPerCu) + "_")) + _SS_(FmhaPipeline::name) + "_" +
             "v" + (std::is_same_v<VLayout, ck_tile::tensor_layout::gemm::RowMajor> ? "r" : "c") + (pn.empty() ? "" : "_" + pn) +
-            (BiasEnum == BlockAttentionBiasEnum::NO_BIAS ? _SS_("") : (_SS_("_") + BlockAttentionBiasEnumToStr<BiasEnum>::name)) + 
+            (BiasEnum == BlockAttentionBiasEnum::NO_BIAS ? _SS_("") : (_SS_("_") + BlockAttentionBiasEnumToStr<BiasEnum>::name)) +
             (kHasMask ? "_" + _SS_(FmhaMask::name) : "") + (kStoreLSE ? "_lse" : "" ) + (kHasDropout ? "_dropout" : "" ) + (kDoFp8StaticQuant ? "_squant" : "" );
         #undef _SS_
         #undef _TS_
@@ -387,7 +387,6 @@ struct FmhaFwdKernel
               ck_tile::index_t nhead_stride_randval,
               ck_tile::index_t nhead_stride_lse,
               ck_tile::index_t nhead_stride_o,
-              ck_tile::index_t batch_stride_lse,
               ck_tile::index_t window_size_left,
               ck_tile::index_t window_size_right,
               ck_tile::index_t mask_type,
@@ -448,7 +447,6 @@ struct FmhaFwdKernel
         {
             kargs.lse_ptr          = lse_ptr;
             kargs.nhead_stride_lse = nhead_stride_lse;
-            kargs.batch_stride_lse = batch_stride_lse;
         }
         if constexpr(kDoFp8StaticQuant)
         {
@@ -524,7 +522,7 @@ struct FmhaFwdKernel
             }
             if constexpr(kStoreLSE)
             {
-                batch_offset_lse = static_cast<long_index_t>(i_batch) * kargs.batch_stride_lse;
+                batch_offset_lse = query_start;
             }
             if constexpr(kHasDropout)
             {

include/ck_tile/ops/fmha/kernel/fmha_fwd_splitkv_combine_kernel.hpp

@@ -55,7 +55,7 @@ struct FmhaFwdSplitKVCombineKernel
             (kBlockPerCuInput == -1 ? "" : ("o" + _TS_(kBlockPerCu) + "_")) +
             _SS_(FmhaPipeline::name) +
             (pn.empty() ? "" : "_" + pn) +
-            (kStoreLSE ? "_lse" : "" ) + 
+            (kStoreLSE ? "_lse" : "" ) +
             (kDoFp8StaticQuant ? "_squant" : "" );
         #undef _SS_
         #undef _TS_
@@ -91,7 +91,6 @@ struct FmhaFwdSplitKVCombineKernel
         ck_tile::index_t nhead_stride_o_acc;
         ck_tile::index_t nhead_stride_o;
 
-        ck_tile::index_t batch_stride_lse_acc;
         ck_tile::index_t batch_stride_o_acc;
 
         ck_tile::index_t split_stride_lse_acc;
@@ -116,6 +115,7 @@ struct FmhaFwdSplitKVCombineKernel
           std::conditional_t<kDoFp8StaticQuant, Fp8StaticQuantKargs, EmptyKargs<1>>
     {
         ck_tile::index_t batch_stride_o;
+        ck_tile::index_t batch_stride_lse_acc;
     };
 
     struct GroupModeKargs
@@ -166,13 +166,13 @@ struct FmhaFwdSplitKVCombineKernel
                      nhead_stride_lse_acc,
                      nhead_stride_o_acc,
                      nhead_stride_o,
-                     batch_stride_lse_acc,
                      batch_stride_o_acc,
                      split_stride_lse_acc,
                      split_stride_o_acc}, // args for common karg
                     {},                   // placeholder for lse
                     {},                   // placeholder for fp8_static_quant args
-                    batch_stride_o};
+                    batch_stride_o,
+                    batch_stride_lse_acc};
 
         if constexpr(kStoreLSE)
         {
@@ -206,9 +206,7 @@ struct FmhaFwdSplitKVCombineKernel
               ck_tile::index_t nhead_stride_o_acc,
               ck_tile::index_t nhead_stride_lse,
               ck_tile::index_t nhead_stride_o,
-              ck_tile::index_t batch_stride_lse_acc,
               ck_tile::index_t batch_stride_o_acc,
-              ck_tile::index_t batch_stride_lse,
               ck_tile::index_t split_stride_lse_acc,
               ck_tile::index_t split_stride_o_acc)
     {
@@ -225,7 +223,6 @@ struct FmhaFwdSplitKVCombineKernel
                      nhead_stride_lse_acc,
                      nhead_stride_o_acc,
                      nhead_stride_o,
-                     batch_stride_lse_acc,
                      batch_stride_o_acc,
                      split_stride_lse_acc,
                      split_stride_o_acc}, // args for common karg
@@ -237,7 +234,6 @@ struct FmhaFwdSplitKVCombineKernel
         {
             kargs.lse_ptr          = lse_ptr;
             kargs.nhead_stride_lse = nhead_stride_lse;
-            kargs.batch_stride_lse = batch_stride_lse;
         }
         if constexpr(kDoFp8StaticQuant)
         {
@@ -274,24 +270,25 @@ struct FmhaFwdSplitKVCombineKernel
         const index_t i_m0 = __builtin_amdgcn_readfirstlane(i_tile_m * FmhaPipeline::kM0);
         const index_t i_n1 = __builtin_amdgcn_readfirstlane(i_tile_n * FmhaPipeline::kN1);
 
-        const long_index_t batch_offset_lse_acc =
-            static_cast<long_index_t>(i_batch) * kargs.batch_stride_lse_acc;
         const long_index_t batch_offset_o_acc =
             static_cast<long_index_t>(i_batch) * kargs.batch_stride_o_acc;
-        long_index_t batch_offset_lse = 0;
-        long_index_t batch_offset_o   = 0;
 
-        if constexpr(kStoreLSE)
-        {
-            batch_offset_lse = static_cast<long_index_t>(i_batch) * kargs.batch_stride_lse;
-        }
+        long_index_t batch_offset_lse_acc = 0;
+        long_index_t batch_offset_lse     = 0;
+        long_index_t batch_offset_o       = 0;
 
         if constexpr(kIsGroupMode)
         {
             // get starting offset for each batch
             const long_index_t query_start = kargs.seqstart_q_ptr[i_batch];
 
-            batch_offset_o = query_start * kargs.row_stride_o;
+            batch_offset_o       = query_start * kargs.row_stride_o;
+            batch_offset_lse_acc = query_start;
+
+            if constexpr(kStoreLSE)
+            {
+                batch_offset_lse = query_start;
+            }
 
             // get real # queries & # keys under group mode
             const auto adjusted_seqstart_q_ptr = kargs.seqstart_q_ptr + i_batch;
@@ -306,7 +303,13 @@ struct FmhaFwdSplitKVCombineKernel
         }
         else
         {
-            batch_offset_o = static_cast<long_index_t>(i_batch) * kargs.batch_stride_o;
+            batch_offset_o       = static_cast<long_index_t>(i_batch) * kargs.batch_stride_o;
+            batch_offset_lse_acc = static_cast<long_index_t>(i_batch) * kargs.batch_stride_lse_acc;
+
+            if constexpr(kStoreLSE)
+            {
+                batch_offset_lse = static_cast<long_index_t>(i_batch) * kargs.batch_stride_lse;
+            }
         }
 
         // for simplicity, batch stride we just modify the pointer

include/ck_tile/ops/fmha/kernel/fmha_fwd_splitkv_kernel.hpp

@@ -85,7 +85,7 @@ struct FmhaFwdSplitKVKernel
             "w" + _TS_(gwt::at(ck_tile::number<0>{})) + "x" + _TS_(gwt::at(ck_tile::number<1>{})) + "x" + _TS_(gwt::at(ck_tile::number<2>{})) + "_" +
             (kBlockPerCuInput == -1 ? "" : ("o" + _TS_(kBlockPerCu) + "_")) + _SS_(FmhaPipeline::name) + "_" +
             "v" + (std::is_same_v<VLayout, ck_tile::tensor_layout::gemm::RowMajor> ? "r" : "c") + (pn.empty() ? "" : "_" + pn) +
-            (BiasEnum == BlockAttentionBiasEnum::NO_BIAS ? _SS_("") : (_SS_("_") + BlockAttentionBiasEnumToStr<BiasEnum>::name)) + 
+            (BiasEnum == BlockAttentionBiasEnum::NO_BIAS ? _SS_("") : (_SS_("_") + BlockAttentionBiasEnumToStr<BiasEnum>::name)) +
             (kHasMask ? "_" + _SS_(FmhaMask::name) : "") + (kHasDropout ? "_dropout" : "" ) + (kDoFp8StaticQuant ? "_squant" : "" );
         #undef _SS_
         #undef _TS_
@@ -136,7 +136,6 @@ struct FmhaFwdSplitKVKernel
         ck_tile::index_t nhead_stride_lse_acc;
         ck_tile::index_t nhead_stride_o_acc;
 
-        ck_tile::index_t batch_stride_lse_acc;
         ck_tile::index_t batch_stride_o_acc;
 
         ck_tile::index_t split_stride_lse_acc;
@@ -216,6 +215,7 @@ struct FmhaFwdSplitKVKernel
         ck_tile::index_t batch_stride_q;
         ck_tile::index_t batch_stride_k;
         ck_tile::index_t batch_stride_v;
+        ck_tile::index_t batch_stride_lse_acc;
     };
 
     struct GroupModeKargs
@@ -313,7 +313,6 @@ struct FmhaFwdSplitKVKernel
                      nhead_stride_v,
                      nhead_stride_lse_acc,
                      nhead_stride_o_acc,
-                     batch_stride_lse_acc,
                      batch_stride_o_acc,
                      split_stride_lse_acc,
                      split_stride_o_acc}, // args for common karg
@@ -323,7 +322,8 @@ struct FmhaFwdSplitKVKernel
                     {},                   // placeholder for dropout
                     batch_stride_q,
                     batch_stride_k,
-                    batch_stride_v};
+                    batch_stride_v,
+                    batch_stride_lse_acc};
 
         if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS)
         {
@@ -394,7 +394,6 @@ struct FmhaFwdSplitKVKernel
               ck_tile::index_t nhead_stride_randval,
               ck_tile::index_t nhead_stride_lse_acc,
               ck_tile::index_t nhead_stride_o_acc,
-              ck_tile::index_t batch_stride_lse_acc,
               ck_tile::index_t batch_stride_o_acc,
               ck_tile::index_t split_stride_lse_acc,
               ck_tile::index_t split_stride_o_acc,
@@ -433,7 +432,6 @@ struct FmhaFwdSplitKVKernel
                      nhead_stride_v,
                      nhead_stride_lse_acc,
                      nhead_stride_o_acc,
-                     batch_stride_lse_acc,
                      batch_stride_o_acc,
                      split_stride_lse_acc,
                      split_stride_o_acc}, // args for common karg
@@ -511,8 +509,7 @@ struct FmhaFwdSplitKVKernel
         long_index_t batch_offset_v       = 0;
         long_index_t batch_offset_bias    = 0;
         long_index_t batch_offset_randval = 0;
-        const long_index_t batch_offset_lse_acc =
-            static_cast<long_index_t>(i_batch) * kargs.batch_stride_lse_acc;
+        long_index_t batch_offset_lse_acc = 0;
         const long_index_t batch_offset_o_acc =
             static_cast<long_index_t>(i_batch) * kargs.batch_stride_o_acc;
 
@@ -522,8 +519,9 @@ struct FmhaFwdSplitKVKernel
             const long_index_t query_start = kargs.seqstart_q_ptr[i_batch];
             const long_index_t key_start   = kargs.seqstart_k_ptr[i_batch];
 
-            batch_offset_q = query_start * kargs.stride_q;
-            batch_offset_k = key_start * kargs.stride_k;
+            batch_offset_q       = query_start * kargs.stride_q;
+            batch_offset_k       = key_start * kargs.stride_k;
+            batch_offset_lse_acc = query_start;
             if constexpr(std::is_same_v<VLayout, ck_tile::tensor_layout::gemm::RowMajor>)
             {
                 batch_offset_v = key_start * kargs.stride_v;
@@ -564,9 +562,10 @@ struct FmhaFwdSplitKVKernel
         }
         else
         {
-            batch_offset_q = static_cast<long_index_t>(i_batch) * kargs.batch_stride_q;
-            batch_offset_k = static_cast<long_index_t>(i_batch) * kargs.batch_stride_k;
-            batch_offset_v = static_cast<long_index_t>(i_batch) * kargs.batch_stride_v;
+            batch_offset_q       = static_cast<long_index_t>(i_batch) * kargs.batch_stride_q;
+            batch_offset_k       = static_cast<long_index_t>(i_batch) * kargs.batch_stride_k;
+            batch_offset_v       = static_cast<long_index_t>(i_batch) * kargs.batch_stride_v;
+            batch_offset_lse_acc = static_cast<long_index_t>(i_batch) * kargs.batch_stride_lse_acc;
             if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS)
             {
                 batch_offset_bias = static_cast<long_index_t>(i_batch) * kargs.batch_stride_bias;

include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_convert_dq.hpp

@@ -0,0 +1,141 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include "ck_tile/core.hpp"
+#include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_default_policy.hpp"
+
+namespace ck_tile {
+
+template <typename Problem, typename Policy = BlockFmhaBwdPipelineDefaultPolicy>
+struct BlockFmhaBwdConvertQGrad
+{
+    using AccDataType   = remove_cvref_t<typename Problem::AccDataType>;
+    using QGradDataType = remove_cvref_t<typename Problem::QGradDataType>;
+
+    static constexpr index_t kM0 = Problem::kM0;
+    static constexpr index_t kN0 = Problem::kN0;
+
+    static constexpr index_t kBlockPerCu = Problem::kBlockPerCu;
+    static constexpr index_t kBlockSize  = Problem::kBlockSize;
+    static constexpr index_t kQKHeaddim  = Problem::kQKHeaddim;
+
+    static constexpr bool kIsGroupMode     = Problem::kIsGroupMode;
+    static constexpr bool kPadSeqLenQ      = Problem::kPadSeqLenQ;
+    static constexpr bool kPadHeadDimQ     = Problem::kPadHeadDimQ;
+    static constexpr bool kIsDeterministic = Problem::kIsDeterministic;
+
+    static constexpr index_t kAlignmentQGradAcc =
+        kPadHeadDimQ ? 1 : Policy::template GetAlignmentPostQGradAcc<Problem>();
+    static constexpr index_t kAlignmentQGrad =
+        kPadHeadDimQ ? 1 : Policy::template GetAlignmentPostQGrad<Problem>();
+
+    CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetSmemSize() { return 0; }
+
+    // Convert only
+    template <typename QGradAccDramBlockWindowTmp, typename QGradDramBlockWindowTmp>
+    CK_TILE_HOST_DEVICE void
+    operator()(const QGradAccDramBlockWindowTmp& dq_acc_dram_block_window_tmp,
+               QGradDramBlockWindowTmp& dq_dram_block_window_tmp) const
+    {
+        static_assert(
+            std::is_same_v<AccDataType,
+                           remove_cvref_t<typename QGradAccDramBlockWindowTmp::DataType>> &&
+                std::is_same_v<QGradDataType,
+                               remove_cvref_t<typename QGradDramBlockWindowTmp::DataType>>,
+            "wrong!");
+
+        static_assert(kM0 == QGradDramBlockWindowTmp{}.get_window_lengths()[number<0>{}], "wrong!");
+
+        auto dq_acc_dram_window =
+            make_tile_window(dq_acc_dram_block_window_tmp.get_bottom_tensor_view(),
+                             dq_acc_dram_block_window_tmp.get_window_lengths(),
+                             dq_acc_dram_block_window_tmp.get_window_origin(),
+                             Policy::template MakePostQGradDramTileDistribution<Problem>());
+
+        auto dq_acc   = load_tile(dq_acc_dram_window);
+        const auto dq = cast_tile<QGradDataType>(dq_acc);
+
+        store_tile(dq_dram_block_window_tmp, dq);
+    }
+
+    // Reduce + Convert
+    template <typename QGradAccDramBlockWindowTmp, typename QGradDramBlockWindowTmp>
+    CK_TILE_HOST_DEVICE void
+    operator()(const QGradAccDramBlockWindowTmp& dq_acc_dram_block_window_tmp,
+               QGradDramBlockWindowTmp& dq_dram_block_window_tmp,
+               index_t nsplits) const
+    {
+        static_assert(
+            std::is_same_v<AccDataType,
+                           remove_cvref_t<typename QGradAccDramBlockWindowTmp::DataType>> &&
+                std::is_same_v<QGradDataType,
+                               remove_cvref_t<typename QGradDramBlockWindowTmp::DataType>>,
+            "wrong!");
+
+        static_assert(kM0 == QGradDramBlockWindowTmp{}.get_window_lengths()[number<0>{}], "wrong!");
+
+        auto dq_acc_dram_window =
+            make_tile_window(dq_acc_dram_block_window_tmp.get_bottom_tensor_view(),
+                             dq_acc_dram_block_window_tmp.get_window_lengths(),
+                             dq_acc_dram_block_window_tmp.get_window_origin(),
+                             Policy::template MakePostQGradAccDramTileDistribution<Problem>());
+
+        auto dq_acc = decltype(load_tile(dq_acc_dram_window)){};
+        clear_tile(dq_acc);
+
+        constexpr auto dq_acc_spans = decltype(dq_acc)::get_distributed_spans();
+        index_t i_total_loops       = 0;
+        auto dq_acc_buf             = load_tile(dq_acc_dram_window);
+        move_tile_window(dq_acc_dram_window, {1, 0, 0});
+
+        do
+        {
+            sweep_tile_span(dq_acc_spans[number<0>{}], [&](auto idx0) {
+                sweep_tile_span(dq_acc_spans[number<1>{}], [&](auto idx1) {
+                    sweep_tile_span(dq_acc_spans[number<2>{}], [&](auto idx2) {
+                        constexpr auto n_i_j_idx = make_tuple(idx0, idx1, idx2);
+                        dq_acc(n_i_j_idx) += dq_acc_buf(n_i_j_idx);
+                    });
+                });
+            });
+
+            dq_acc_buf = load_tile(dq_acc_dram_window);
+            move_tile_window(dq_acc_dram_window, {1, 0, 0});
+
+            i_total_loops += 1;
+        } while(i_total_loops < (nsplits - 1));
+
+        sweep_tile_span(dq_acc_spans[number<0>{}], [&](auto idx0) {
+            sweep_tile_span(dq_acc_spans[number<1>{}], [&](auto idx1) {
+                sweep_tile_span(dq_acc_spans[number<2>{}], [&](auto idx2) {
+                    constexpr auto n_i_j_idx = make_tuple(idx0, idx1, idx2);
+                    dq_acc(n_i_j_idx) += dq_acc_buf(n_i_j_idx);
+                });
+            });
+        });
+
+        // declare dq
+        constexpr auto dq_converted_dstr =
+            Policy::template MakePostQGradAccDramTileDistribution<Problem>();
+        auto dq_converted = make_static_distributed_tensor<QGradDataType>(dq_converted_dstr);
+
+        sweep_tile_span(dq_acc_spans[number<0>{}], [&](auto idx0) {
+            sweep_tile_span(dq_acc_spans[number<1>{}], [&](auto idx1) {
+                sweep_tile_span(dq_acc_spans[number<2>{}], [&](auto idx2) {
+                    constexpr auto n_i_j_idx = make_tuple(idx0, idx1, idx2);
+                    dq_converted(n_i_j_idx)  = type_convert<QGradDataType>(dq_acc[n_i_j_idx]);
+                });
+            });
+        });
+
+        constexpr auto dq_dstr = Policy::template MakePostQGradDramTileDistribution<Problem>();
+        auto dq                = make_static_distributed_tensor<QGradDataType>(dq_dstr);
+        dq.get_thread_buffer() = dq_converted.get_thread_buffer();
+
+        store_tile(dq_dram_block_window_tmp, dq);
+    }
+};
+
+} // namespace ck_tile

include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dot_do_o.hpp

@@ -4,11 +4,11 @@
 #pragma once
 
 #include "ck_tile/core.hpp"
-#include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_dot_do_o_default_policy.hpp"
+#include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_default_policy.hpp"
 
 namespace ck_tile {
 
-template <typename Problem, typename Policy = BlockFmhaBwdOGradDotODefaultPolicy>
+template <typename Problem, typename Policy = BlockFmhaBwdPipelineDefaultPolicy>
 struct BlockFmhaBwdOGradDotO
 {
     using ODataType     = remove_cvref_t<typename Problem::ODataType>;
@@ -26,7 +26,7 @@ struct BlockFmhaBwdOGradDotO
     static constexpr index_t kAlignmentO =
         kPadHeadDimV ? 1 : Policy::template GetAlignmentO<Problem>();
     static constexpr index_t kAlignmentOGrad =
-        kPadHeadDimV ? 1 : Policy::template GetAlignmentOGrad<Problem>();
+        kPadHeadDimV ? 1 : Policy::template GetAlignmentO<Problem>();
 
     CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetSmemSize() { return 0; }
 

include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dot_do_o_default_policy.hpp

@@ -1,20 +0,0 @@
-// SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
-
-#pragma once
-
-#include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_default_policy.hpp"
-
-namespace ck_tile {
-
-// These templates are not used here.
-using BlockFmhaBwdOGradDotODefaultPolicy =
-    BlockFmhaBwdPipelineDefaultPolicy</* QLoadOnce_      = */ false,
-                                      /* QTLoadOnce_     = */ false,
-                                      /* KLoadOnce_      = */ false,
-                                      /* KTLoadOnce_     = */ false,
-                                      /* VLoadOnce_      = */ false,
-                                      /* OGradLoadOnce_  = */ false,
-                                      /* OGradTLoadOnce_ = */ false>;
-
-} // namespace ck_tile

include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_kr_ktr_vr.hpp

@@ -0,0 +1,782 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include "ck_tile/core.hpp"
+#include "ck_tile/ops/fmha/block/block_attention_bias_enum.hpp"
+#include "ck_tile/ops/fmha/block/block_dropout.hpp"
+#include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_default_policy.hpp"
+#include "ck_tile/ops/reduce/block/block_reduce.hpp"
+
+namespace ck_tile {
+
+template <typename Problem, typename Policy = BlockFmhaBwdPipelineDefaultPolicy>
+struct BlockFmhaBwdDQDKDVPipelineKRKTRVR
+{
+    using QDataType             = remove_cvref_t<typename Problem::QDataType>;
+    using KDataType             = remove_cvref_t<typename Problem::KDataType>;
+    using VDataType             = remove_cvref_t<typename Problem::VDataType>;
+    using GemmDataType          = remove_cvref_t<typename Problem::GemmDataType>;
+    using BiasDataType          = remove_cvref_t<typename Problem::BiasDataType>;
+    using LSEDataType           = remove_cvref_t<typename Problem::LSEDataType>;
+    using AccDataType           = remove_cvref_t<typename Problem::AccDataType>;
+    using DDataType             = remove_cvref_t<typename Problem::DDataType>;
+    using RandValOutputDataType = remove_cvref_t<typename Problem::RandValOutputDataType>;
+    using ODataType             = remove_cvref_t<typename Problem::ODataType>;
+    using OGradDataType         = remove_cvref_t<typename Problem::OGradDataType>;
+    using QGradDataType         = remove_cvref_t<typename Problem::QGradDataType>;
+    using KGradDataType         = remove_cvref_t<typename Problem::KGradDataType>;
+    using VGradDataType         = remove_cvref_t<typename Problem::VGradDataType>;
+    using BiasGradDataType      = remove_cvref_t<typename Problem::BiasGradDataType>;
+    using FmhaMask              = remove_cvref_t<typename Problem::FmhaMask>;
+    using FmhaDropout           = remove_cvref_t<typename Problem::FmhaDropout>;
+    using HotLoopScheduler      = typename Policy::template HotLoopScheduler<Problem>;
+
+    using BlockFmhaShape = remove_cvref_t<typename Problem::BlockFmhaShape>;
+
+    static constexpr index_t kBlockPerCu = Problem::kBlockPerCu;
+    static constexpr index_t kBlockSize  = Problem::kBlockSize;
+
+    static constexpr index_t kM0        = BlockFmhaShape::kM0;
+    static constexpr index_t kN0        = BlockFmhaShape::kN0;
+    static constexpr index_t kK0        = BlockFmhaShape::kK0;
+    static constexpr index_t kK1        = BlockFmhaShape::kK1;
+    static constexpr index_t kK2        = BlockFmhaShape::kK2;
+    static constexpr index_t kK3        = BlockFmhaShape::kK3;
+    static constexpr index_t kK4        = BlockFmhaShape::kK4;
+    static constexpr index_t kQKHeaddim = BlockFmhaShape::kQKHeaddim;
+    static constexpr index_t kVHeaddim  = BlockFmhaShape::kVHeaddim;
+
+    static constexpr bool kIsGroupMode     = Problem::kIsGroupMode;
+    static constexpr bool kPadSeqLenQ      = Problem::kPadSeqLenQ;
+    static constexpr bool kPadSeqLenK      = Problem::kPadSeqLenK;
+    static constexpr bool kPadHeadDimQ     = Problem::kPadHeadDimQ;
+    static constexpr bool kPadHeadDimV     = Problem::kPadHeadDimV;
+    static constexpr auto BiasEnum         = Problem::BiasEnum;
+    static constexpr bool kHasBiasGrad     = Problem::kHasBiasGrad;
+    static constexpr bool kIsDeterministic = Problem::kIsDeterministic;
+
+    // last dimension vector length used to create tensor view(and decide buffer_load vector length)
+    // ... together with tensor distribution. tensor dist should able to overwrite this
+    static constexpr index_t kAlignmentQ =
+        kPadHeadDimQ ? 1 : Policy::template GetAlignmentQ<Problem>();
+    static constexpr index_t kAlignmentK =
+        kPadHeadDimQ ? 1 : Policy::template GetAlignmentK<Problem>();
+    static constexpr index_t kAlignmentV =
+        kPadHeadDimV ? 1 : Policy::template GetAlignmentV<Problem>();
+    static constexpr index_t kAlignmentOGrad =
+        kPadHeadDimV ? 1 : Policy::template GetAlignmentOGrad<Problem>();
+    static constexpr index_t kAlignmentQGrad = 1;
+    static constexpr index_t kAlignmentKGrad =
+        kPadHeadDimQ ? 1 : Policy::template GetAlignmentKGrad<Problem>();
+    static constexpr index_t kAlignmentVGrad =
+        kPadHeadDimV ? 1 : Policy::template GetAlignmentVGrad<Problem>();
+    static constexpr index_t kAlignmentBias =
+        kPadSeqLenK ? 1 : Policy::template GetTransposedAlignmentBias<Problem>();
+
+    static constexpr const char* name = "kr_ktr_vr";
+
+    CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetSmemSize()
+    {
+        return Policy::template GetSmemSize<Problem>();
+    }
+
+    template <typename QDramBlockWindowTmp,
+              typename KDramBlockWindowTmp,
+              typename VDramBlockWindowTmp,
+              typename BiasDramBlockWindowTmp,
+              typename RandValDramBlockWindowTmp,
+              typename OGradDramBlockWindowTmp,
+              typename LSEDramBlockWindowTmp,
+              typename DDramBlockWindowTmp,
+              typename QGradDramBlockWindowTmp,
+              typename BiasGradDramBlockWindowTmp,
+              typename PositionEncoding>
+    CK_TILE_HOST_DEVICE auto
+    operator()(const QDramBlockWindowTmp& q_dram_block_window_tmp,
+               const KDramBlockWindowTmp& k_dram_block_window_tmp,
+               const VDramBlockWindowTmp& v_dram_block_window_tmp,
+               const BiasDramBlockWindowTmp& bias_dram_block_window_tmp,
+               const RandValDramBlockWindowTmp& randval_dram_block_window_tmp,
+               const OGradDramBlockWindowTmp& do_dram_block_window_tmp,
+               const LSEDramBlockWindowTmp& lse_dram_block_window_tmp,
+               const DDramBlockWindowTmp& d_dram_block_window_tmp,
+               const QGradDramBlockWindowTmp& dq_dram_block_window_tmp,
+               const BiasGradDramBlockWindowTmp& dbias_dram_block_window_tmp,
+               FmhaMask mask,
+               PositionEncoding position_encoding,
+               float raw_scale,
+               float scale,
+               float rp_undrop,
+               float scale_rp_undrop,
+               void* smem_ptr,
+               FmhaDropout& dropout) const
+    {
+        static_assert(
+            std::is_same_v<QDataType, remove_cvref_t<typename QDramBlockWindowTmp::DataType>> &&
+                std::is_same_v<KDataType, remove_cvref_t<typename KDramBlockWindowTmp::DataType>> &&
+                std::is_same_v<VDataType, remove_cvref_t<typename VDramBlockWindowTmp::DataType>> &&
+                std::is_same_v<OGradDataType,
+                               remove_cvref_t<typename OGradDramBlockWindowTmp::DataType>> &&
+                std::is_same_v<LSEDataType,
+                               remove_cvref_t<typename LSEDramBlockWindowTmp::DataType>> &&
+                std::is_same_v<DDataType, remove_cvref_t<typename DDramBlockWindowTmp::DataType>>,
+            "wrong!");
+
+        static_assert(kM0 == QDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
+                          kN0 == KDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
+                          kN0 == VDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
+                          kM0 == BiasDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
+                          kN0 == BiasDramBlockWindowTmp{}.get_window_lengths()[number<1>{}] &&
+                          kM0 == OGradDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
+                          kM0 == LSEDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
+                          kM0 == DDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
+                          kM0 == QGradDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
+                          kM0 == BiasGradDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
+                          kN0 == BiasGradDramBlockWindowTmp{}.get_window_lengths()[number<1>{}],
+                      "wrong!");
+
+        // Block GEMM
+        constexpr auto gemm_0 = Policy::template GetQKBlockGemm<Problem>();
+        constexpr auto gemm_1 = Policy::template GetPTOGradTBlockGemm<Problem>();
+        constexpr auto gemm_2 = Policy::template GetOGradVBlockGemm<Problem>();
+        constexpr auto gemm_3 = Policy::template GetSGradTQTBlockGemm<Problem>();
+        constexpr auto gemm_4 = Policy::template GetSGradKTBlockGemm<Problem>();
+
+        // init VGrad & KGrad
+        auto dv_acc = decltype(gemm_1.MakeCBlockTile()){};
+        auto dk_acc = decltype(gemm_3.MakeCBlockTile()){};
+
+        // K, HBM ->LDS ->Reg
+        auto k_dram_window =
+            make_tile_window(k_dram_block_window_tmp.get_bottom_tensor_view(),
+                             k_dram_block_window_tmp.get_window_lengths(),
+                             k_dram_block_window_tmp.get_window_origin(),
+                             Policy::template MakeKDramTileDistribution<Problem>());
+
+        const auto k_origin = k_dram_window.get_window_origin();
+        // Early termination
+        const auto [seqlen_q_start, seqlen_q_end] =
+            mask.GetTileRangeAlongY(k_origin.at(number<0>{}), number<kM0>{}, number<kN0>{});
+
+        const auto num_total_loop = integer_divide_ceil(seqlen_q_end - seqlen_q_start, kM0);
+
+        // check early exit if masked and no work to do.
+        if constexpr(FmhaMask::IsMasking)
+        {
+            if(num_total_loop <= 0)
+            {
+                // Note: here dk_acc&dv_acc are all cleard, return it
+                // Note: v loaded but no fence, ignore it.
+                return make_tuple(dk_acc, dv_acc);
+            }
+        }
+        KDataType* k_lds_ptr =
+            static_cast<KDataType*>(static_cast<void*>(static_cast<char*>(smem_ptr)));
+        auto k_lds = make_tensor_view<address_space_enum::lds>(
+            k_lds_ptr, Policy::template MakeKLdsWriteBlockDescriptor<Problem>());
+
+        auto k_lds_write_window =
+            make_tile_window(k_lds, make_tuple(number<kN0>{}, number<kK0>{}), {0, 0});
+
+        auto k_lds_read_window =
+            make_tile_window(k_lds_write_window.get_bottom_tensor_view(),
+                             make_tuple(number<kN0>{}, number<kK0>{}),
+                             k_lds_write_window.get_window_origin(),
+                             Policy::template MakeKRegSliceBlockDescriptor<Problem>());
+
+        auto k_reg_tensor = make_static_distributed_tensor<KDataType>(
+            Policy::template MakeKRegBlockDescriptor<Problem>());
+
+        //------------------------------------------------------------------
+        // V, HBM ->LDS ->Reg
+        auto v_dram_window =
+            make_tile_window(v_dram_block_window_tmp.get_bottom_tensor_view(),
+                             v_dram_block_window_tmp.get_window_lengths(),
+                             v_dram_block_window_tmp.get_window_origin(),
+                             Policy::template MakeVDramTileDistribution<Problem>());
+
+        VDataType* v_lds_ptr =
+            static_cast<VDataType*>(static_cast<void*>(static_cast<char*>(smem_ptr)));
+
+        auto v_lds = make_tensor_view<address_space_enum::lds>(
+            v_lds_ptr, Policy::template MakeVLdsWriteBlockDescriptor<Problem>());
+
+        auto v_lds_write_window =
+            make_tile_window(v_lds, make_tuple(number<kN0>{}, number<kK2>{}), {0, 0});
+
+        auto v_lds_read_window =
+            make_tile_window(v_lds_write_window.get_bottom_tensor_view(),
+                             make_tuple(number<kN0>{}, number<kK2>{}),
+                             v_lds_write_window.get_window_origin(),
+                             Policy::template MakeVRegSliceBlockDescriptor<Problem>());
+
+        auto v_reg_tensor = make_static_distributed_tensor<VDataType>(
+            Policy::template MakeVRegBlockDescriptor<Problem>());
+
+        //------------------------------------------------------------------
+        // KT, Reg ->LDS ->Reg
+        auto shuffled_k_block_tile = make_static_distributed_tensor<KDataType>(
+            Policy::template MakeShuffledKRegWriteBlockDescriptor<Problem>());
+
+        KDataType* kt_lds_ptr = static_cast<KDataType*>(static_cast<void*>(
+            static_cast<char*>(smem_ptr) + Policy::template GetSmemSizeK<Problem>()));
+
+        auto shuffled_k_lds_write = make_tensor_view<address_space_enum::lds>(
+            kt_lds_ptr, Policy::template MakeShuffledKLdsWriteBlockDescriptor<Problem>());
+
+        auto shuffled_k_lds_write_window = make_tile_window(
+            shuffled_k_lds_write, make_tuple(number<kN0>{}, number<kK0>{}), {0, 0});
+
+        auto kt_lds_read = make_tensor_view<address_space_enum::lds>(
+            kt_lds_ptr, Policy::template MakeKTLdsReadBlockDescriptor<Problem>());
+
+        auto kt_lds_read_window =
+            make_tile_window(kt_lds_read,
+                             make_tuple(number<kQKHeaddim>{}, number<kN0>{}),
+                             {0, 0},
+                             Policy::template MakeKTRegBlockDescriptor<Problem>());
+
+        //------------------------------------------------------------------
+        // Pre-Load KV into Registers
+        auto k_block_tile = load_tile(k_dram_window);
+        auto v_block_tile = load_tile(v_dram_window);
+
+        store_tile(k_lds_write_window, k_block_tile);
+        shuffle_tile(shuffled_k_block_tile, k_block_tile);
+        store_tile(shuffled_k_lds_write_window, shuffled_k_block_tile);
+
+        block_sync_lds();
+        k_reg_tensor = load_tile(k_lds_read_window);
+        block_sync_lds();
+
+        auto kt_reg_tensor = load_tile(kt_lds_read_window);
+
+        store_tile(v_lds_write_window, v_block_tile);
+
+        block_sync_lds();
+
+        v_reg_tensor = load_tile(v_lds_read_window);
+        block_sync_lds();
+        //---------------------------- Loop Load in ----------------------------//
+        // Q: HBM ->Reg ->LDS
+        auto q_dram_window =
+            make_tile_window(q_dram_block_window_tmp.get_bottom_tensor_view(),
+                             q_dram_block_window_tmp.get_window_lengths(),
+                             {seqlen_q_start, 0},
+                             Policy::template MakeQDramTileDistribution<Problem>());
+
+        QDataType* q_lds_ptr = static_cast<QDataType*>(static_cast<void*>(
+            static_cast<char*>(smem_ptr) + Policy::template GetSmemSizeQT<Problem>() +
+            Policy::template GetSmemSizeOGrad<Problem>() +
+            Policy::template GetSmemSizeOGradT<Problem>()));
+
+        auto q_lds = make_tensor_view<address_space_enum::lds>(
+            q_lds_ptr, Policy::template MakeQLdsBlockDescriptor<Problem>());
+
+        auto q_lds_window =
+            make_tile_window(q_lds, make_tuple(number<kM0>{}, number<kK0>{}), {0, 0});
+
+        auto q_lds_read_window =
+            make_tile_window(q_lds_window.get_bottom_tensor_view(),
+                             make_tuple(number<kM0>{}, number<kK0>{}),
+                             q_lds_window.get_window_origin(),
+                             Policy::template MakeQRegSliceBlockDescriptor<Problem>());
+
+        auto pt_reg_tensor = make_static_distributed_tensor<GemmDataType>(
+            Policy::template MakePTRegSliceBlockDescriptor<Problem>());
+        // QT: Reg -> Reg-> LDS
+        auto shuffled_q_block_tile = make_static_distributed_tensor<QDataType>(
+            Policy::template MakeShuffledQRegWriteBlockDescriptor<Problem>());
+
+        QDataType* qt_lds_ptr =
+            static_cast<QDataType*>(static_cast<void*>(static_cast<char*>(smem_ptr)));
+
+        auto shuffled_q_lds_write = make_tensor_view<address_space_enum::lds>(
+            qt_lds_ptr, Policy::template MakeShuffledQLdsWriteBlockDescriptor<Problem>());
+
+        auto shuffled_q_lds_write_window = make_tile_window(
+            shuffled_q_lds_write, make_tuple(number<kM0>{}, number<kK0>{}), {0, 0});
+
+        auto qt_lds_read = make_tensor_view<address_space_enum::lds>(
+            qt_lds_ptr, Policy::template MakeQTLdsReadBlockDescriptor<Problem>());
+
+        auto qt_lds_read_window =
+            make_tile_window(qt_lds_read,
+                             make_tuple(number<kQKHeaddim>{}, number<kM0>{}),
+                             {0, 0},
+                             Policy::template MakeQTRegSliceBlockDescriptor<Problem>());
+
+        // dO: HBM ->Reg ->LDS
+        auto do_dram_window =
+            make_tile_window(do_dram_block_window_tmp.get_bottom_tensor_view(),
+                             do_dram_block_window_tmp.get_window_lengths(),
+                             {seqlen_q_start, 0},
+                             Policy::template MakeOGradDramTileDistribution<Problem>());
+
+        OGradDataType* do_lds_ptr = static_cast<OGradDataType*>(static_cast<void*>(
+            static_cast<char*>(smem_ptr) + Policy::template GetSmemSizeQT<Problem>()));
+
+        auto do_lds = make_tensor_view<address_space_enum::lds>(
+            do_lds_ptr, Policy::template MakeOGradLdsBlockDescriptor<Problem>());
+
+        auto do_lds_window =
+            make_tile_window(do_lds, make_tuple(number<kM0>{}, number<kK2>{}), {0, 0});
+
+        auto do_lds_read_window =
+            make_tile_window(do_lds_window.get_bottom_tensor_view(),
+                             make_tuple(number<kM0>{}, number<kK2>{}),
+                             do_lds_window.get_window_origin(),
+                             Policy::template MakeOGradRegSliceBlockDescriptor<Problem>());
+        // dOT: Reg ->Reg ->LDS
+        auto shuffled_do_block_tile = make_static_distributed_tensor<OGradDataType>(
+            Policy::template MakeShuffledOGradRegWriteBlockDescriptor<Problem>());
+
+        OGradDataType* dot_lds_ptr = static_cast<OGradDataType*>(static_cast<void*>(
+            static_cast<char*>(smem_ptr) + Policy::template GetSmemSizeQT<Problem>() +
+            Policy::template GetSmemSizeOGrad<Problem>()));
+
+        auto shuffled_do_lds_write = make_tensor_view<address_space_enum::lds>(
+            dot_lds_ptr, Policy::template MakeShuffledOGradLdsWriteBlockDescriptor<Problem>());
+
+        auto shuffled_do_lds_write_window = make_tile_window(
+            shuffled_do_lds_write, make_tuple(number<kM0>{}, number<kK2>{}), {0, 0});
+
+        auto dot_read_lds = make_tensor_view<address_space_enum::lds>(
+            dot_lds_ptr, Policy::template MakeOGradTLdsReadBlockDescriptor<Problem>());
+
+        auto dot_lds_read_window =
+            make_tile_window(dot_read_lds,
+                             make_tuple(number<kVHeaddim>{}, number<kM0>{}),
+                             {0, 0},
+                             Policy::template MakeOGradTRegSliceBlockDescriptor<Problem>());
+
+        // dS: Reg -> Reg -> LDS
+        GemmDataType* ds_lds_ptr = static_cast<GemmDataType*>(static_cast<void*>(
+            static_cast<char*>(smem_ptr) + Policy::template GetSmemSizeQT<Problem>() +
+            Policy::template GetSmemSizeOGrad<Problem>() +
+            Policy::template GetSmemSizeOGradT<Problem>() +
+            Policy::template GetSmemSizeQ<Problem>() + Policy::template GetSmemSizeLSE<Problem>() +
+            Policy::template GetSmemSizeD<Problem>()));
+
+        auto ds_lds = make_tensor_view<address_space_enum::lds>(
+            ds_lds_ptr, Policy::template MakeSGradLdsBlockDescriptor<Problem>());
+
+        auto ds_lds_window =
+            make_tile_window(ds_lds, make_tuple(number<kM0>{}, number<kN0>{}), {0, 0});
+
+        auto ds_lds_read_window =
+            make_tile_window(ds_lds_window.get_bottom_tensor_view(),
+                             make_tuple(number<kM0>{}, number<kK4>{}),
+                             ds_lds_window.get_window_origin(),
+                             Policy::template MakeSGradRegSliceBlockDescriptor<Problem>());
+
+        auto dst_reg_tensor = make_static_distributed_tensor<GemmDataType>(
+            Policy::template MakeSGradTRegSliceBlockDescriptor<Problem>());
+        // Bias: HBM ->Reg ->Reg ->LDS
+        const auto bias_origin = bias_dram_block_window_tmp.get_window_origin();
+
+        auto bias_dram_window =
+            make_tile_window(bias_dram_block_window_tmp.get_bottom_tensor_view(),
+                             bias_dram_block_window_tmp.get_window_lengths(),
+                             {seqlen_q_start, bias_origin.at(number<1>{})},
+                             Policy::template MakeBiasTileDistribution<Problem>());
+
+        BiasDataType* bias_lds_ptr = static_cast<BiasDataType*>(static_cast<void*>(
+            static_cast<char*>(smem_ptr) + Policy::template GetSmemSizeQT<Problem>() +
+            Policy::template GetSmemSizeOGrad<Problem>() +
+            Policy::template GetSmemSizeOGradT<Problem>() +
+            Policy::template GetSmemSizeQ<Problem>() + Policy::template GetSmemSizeLSE<Problem>() +
+            Policy::template GetSmemSizeD<Problem>()));
+
+        auto bias_lds = make_tensor_view<address_space_enum::lds>(
+            bias_lds_ptr, Policy::template MakeBiasLdsBlockDescriptor<Problem>());
+
+        auto bias_lds_write_window =
+            make_tile_window(bias_lds, make_tuple(number<kM0>{}, number<kN0>{}), {0, 0});
+
+        auto bias_s_lds_read_window =
+            make_tile_window(bias_lds_write_window.get_bottom_tensor_view(),
+                             bias_lds_write_window.get_window_lengths(),
+                             bias_lds_write_window.get_window_origin(),
+                             Policy::template MakeBiasSTileDistribution<decltype(gemm_0)>());
+
+        static_assert(std::is_same_v<BiasDataType, BiasGradDataType>,
+                      "BiasDataType and BiasGradDataType should be the same!");
+
+        // LSE: HBM -> LDS ->Reg
+        auto lse_dram_window = make_tile_window(
+            lse_dram_block_window_tmp.get_bottom_tensor_view(),
+            lse_dram_block_window_tmp.get_window_lengths(),
+            {seqlen_q_start},
+            Policy::template MakeLSEDDramTileDistribution<Problem, decltype(gemm_0)>());
+
+        LSEDataType* lse_lds_ptr = static_cast<LSEDataType*>(static_cast<void*>(
+            static_cast<char*>(smem_ptr) + Policy::template GetSmemSizeQT<Problem>() +
+            Policy::template GetSmemSizeOGrad<Problem>() +
+            Policy::template GetSmemSizeOGradT<Problem>() +
+            Policy::template GetSmemSizeQ<Problem>()));
+
+        auto lse_lds = make_tensor_view<address_space_enum::lds>(
+            lse_lds_ptr, Policy::template MakeLSEDLdsWriteBlockDescriptor<Problem>());
+
+        auto lse_lds_write_window = make_tile_window(lse_lds, make_tuple(number<kM0>{}), {0});
+
+        auto lse_lds_read_window = make_tile_window(
+            lse_lds,
+            make_tuple(number<kM0>{}),
+            {0},
+            Policy::template MakeLSEDLdsReadBlockDescriptor<Problem, decltype(gemm_0)>());
+
+        // D: HBM ->Reg
+        auto d_dram_window = make_tile_window(
+            d_dram_block_window_tmp.get_bottom_tensor_view(),
+            d_dram_block_window_tmp.get_window_lengths(),
+            {seqlen_q_start},
+            Policy::template MakeLSEDDramTileDistribution<Problem, decltype(gemm_0)>());
+
+        DDataType* d_lds_ptr = static_cast<DDataType*>(static_cast<void*>(
+            static_cast<char*>(smem_ptr) + Policy::template GetSmemSizeQT<Problem>() +
+            Policy::template GetSmemSizeOGrad<Problem>() +
+            Policy::template GetSmemSizeOGradT<Problem>() +
+            Policy::template GetSmemSizeQ<Problem>() + Policy::template GetSmemSizeLSE<Problem>()));
+
+        auto d_lds = make_tensor_view<address_space_enum::lds>(
+            d_lds_ptr, Policy::template MakeLSEDLdsWriteBlockDescriptor<Problem>());
+
+        auto d_lds_write_window = make_tile_window(d_lds, make_tuple(number<kM0>{}), {0});
+
+        auto d_lds_read_window = make_tile_window(
+            d_lds,
+            make_tuple(number<kM0>{}),
+            {0},
+            Policy::template MakeLSEDLdsReadBlockDescriptor<Problem, decltype(gemm_0)>());
+
+        // RandVal: HBM ->Reg
+        auto randval_dram_window = dropout.template MakeRandvalDramWindow<decltype(gemm_0), false>(
+            randval_dram_block_window_tmp, seqlen_q_start);
+
+        // BiasGrad
+        // Reg ->LDS ->Reg ->HBM
+        const auto dbias_origin = dbias_dram_block_window_tmp.get_window_origin();
+
+        auto dbias_dram_window =
+            make_tile_window(dbias_dram_block_window_tmp.get_bottom_tensor_view(),
+                             dbias_dram_block_window_tmp.get_window_lengths(),
+                             {seqlen_q_start, dbias_origin.at(number<1>{})}); // M/N
+
+        auto dbias_lds_read_window =
+            make_tile_window(bias_lds,
+                             make_tuple(number<kM0>{}, number<kN0>{}),
+                             {0, 0},
+                             Policy::template MakeShuffledBiasTileDistribution<Problem>());
+
+        // ----------------------------Loop write out------------------------------//
+        auto dq_dram_window = make_tile_window(dq_dram_block_window_tmp.get_bottom_tensor_view(),
+                                               dq_dram_block_window_tmp.get_window_lengths(),
+                                               {seqlen_q_start, 0});
+
+        using SPBlockTileType     = decltype(gemm_0.MakeCBlockTile());
+        using SPGradBlockTileType = decltype(gemm_2.MakeCBlockTile());
+        using QGradBlockTileType  = decltype(gemm_4.MakeCBlockTile());
+
+        index_t i_total_loops = 0;
+        index_t seqlen_q_step = seqlen_q_start;
+        static_assert(kQKHeaddim == kK0, "kQKHeaddim should equal to kK0");
+        static_assert(kM0 == kK1, "kM0 should equal to kK1");
+        static_assert(kVHeaddim == kK2, "kVHeaddim should equal to kK2");
+        static_assert(kM0 == kK3, "kM0 should equal to kK3");
+        constexpr index_t k4_loops = kN0 / kK4;
+
+        clear_tile(dv_acc);
+        clear_tile(dk_acc);
+
+        __builtin_amdgcn_sched_barrier(0);
+        // Hot loop
+        while(i_total_loops < num_total_loop)
+        {
+            auto q_block_tile = load_tile(q_dram_window);
+            move_tile_window(q_dram_window, {kM0, 0});
+
+            auto lse_block_tile = load_tile(lse_dram_window);
+            move_tile_window(lse_dram_window, {kM0});
+
+            store_tile(q_lds_window, q_block_tile);
+            shuffle_tile(shuffled_q_block_tile, q_block_tile);
+            store_tile(shuffled_q_lds_write_window, shuffled_q_block_tile);
+
+            store_tile(lse_lds_write_window, lse_block_tile);
+
+            block_sync_lds();
+
+            auto q_reg_tensor = load_tile(q_lds_read_window);
+            auto lse          = load_tile(lse_lds_read_window);
+
+            block_sync_lds();
+
+            // STAGE 1, Q@K Gemm0
+            auto s_acc = SPBlockTileType{};
+
+            s_acc = gemm_0(q_reg_tensor, k_reg_tensor);
+
+            // STAGE 2, Scale, Add bias, Mask, Softmax, Dropout
+            if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS)
+            {
+                const auto bias_tile    = load_tile(bias_dram_window);
+                auto shuffled_bias_tile = make_static_distributed_tensor<BiasDataType>(
+                    Policy::template MakeShuffledBiasTileDistribution<Problem>());
+                shuffle_tile(shuffled_bias_tile, bias_tile);
+                store_tile(bias_lds_write_window, shuffled_bias_tile);
+                block_sync_lds();
+                auto bias_s_tile = load_tile(bias_s_lds_read_window);
+                tile_elementwise_inout(
+                    [&](auto& x, const auto& y) {
+                        x = scale * x + log2e_v<AccDataType> * type_convert<AccDataType>(y);
+                    },
+                    s_acc,
+                    bias_s_tile);
+                move_tile_window(bias_dram_window, {kM0, 0});
+                __builtin_amdgcn_sched_barrier(0);
+            }
+            else if constexpr(BiasEnum == BlockAttentionBiasEnum::ALIBI)
+            {
+                constexpr auto s_spans = decltype(s_acc)::get_distributed_spans();
+                sweep_tile_span(s_spans[number<0>{}], [&](auto idx0) {
+                    sweep_tile_span(s_spans[number<1>{}], [&](auto idx1) {
+                        const auto tile_idx = get_x_indices_from_distributed_indices(
+                            s_acc.get_tile_distribution(), make_tuple(idx0, idx1));
+
+                        const auto row = seqlen_q_step + tile_idx.at(number<0>{});
+                        const auto col = k_origin.at(number<0>{}) + tile_idx.at(number<1>{});
+                        constexpr auto i_j_idx = make_tuple(idx0, idx1);
+
+                        s_acc(i_j_idx) *= scale;
+                        position_encoding.update(s_acc(i_j_idx), row, col);
+                    });
+                });
+            }
+
+            if constexpr(kPadSeqLenK || FmhaMask::IsMasking)
+            {
+                bool need_perpixel_check = mask.IsEdgeTile(
+                    seqlen_q_step, k_origin.at(number<0>{}), number<kM0>{}, number<kN0>{});
+                if(need_perpixel_check)
+                {
+                    set_tile_if(s_acc, -numeric<AccDataType>::infinity(), [&](auto tile_idx) {
+                        const auto row = seqlen_q_step + tile_idx.at(number<0>{});
+                        const auto col = k_origin.at(number<0>{}) + tile_idx.at(number<1>{});
+                        return mask.IsOutOfBound(row, col);
+                    });
+                }
+            }
+
+            static const auto get_validated_lse = [](LSEDataType raw_lse) {
+                if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS ||
+                             FmhaMask::IsMasking)
+                {
+                    return raw_lse == -numeric<LSEDataType>::infinity()
+                               ? type_convert<LSEDataType>(0.f)
+                               : raw_lse;
+                }
+                else
+                {
+                    return raw_lse;
+                }
+            };
+
+            auto p                 = SPBlockTileType{};
+            constexpr auto p_spans = decltype(p)::get_distributed_spans();
+            sweep_tile_span(p_spans[number<0>{}], [&](auto idx0) {
+                constexpr auto i_idx = make_tuple(idx0);
+                auto row_lse         = log2e_v<LSEDataType> * get_validated_lse(lse[i_idx]);
+
+                sweep_tile_span(p_spans[number<1>{}], [&](auto idx1) {
+                    constexpr auto i_j_idx = make_tuple(idx0, idx1);
+
+                    if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS ||
+                                 BiasEnum == BlockAttentionBiasEnum::ALIBI)
+                    {
+                        p(i_j_idx) = exp2(s_acc[i_j_idx] - row_lse);
+                    }
+                    else
+                    {
+                        p(i_j_idx) = exp2(scale * s_acc[i_j_idx] - row_lse);
+                    }
+                });
+            });
+
+            if constexpr(FmhaDropout::IsDropout)
+            {
+                dropout.template Run<decltype(gemm_0), RandValOutputDataType>(
+                    seqlen_q_step, k_origin.at(number<0>{}), p, randval_dram_window);
+            }
+            const auto p_gemm = [&]() {
+                if constexpr(FmhaDropout::IsDropout)
+                {
+                    return tile_elementwise_in(
+                        [](const auto& x) { return type_convert<GemmDataType>(x > 0.f ? x : 0.f); },
+                        p);
+                }
+                else
+                {
+                    return cast_tile<GemmDataType>(p);
+                }
+            }();
+
+            // STAGE 3, P^T@OGrad^T Gemm1
+            auto do_block_tile = load_tile(do_dram_window);
+            move_tile_window(do_dram_window, {kM0, 0});
+
+            auto d_block_tile = load_tile(d_dram_window);
+            move_tile_window(d_dram_window, {kM0});
+
+            store_tile(do_lds_window, do_block_tile);
+            shuffle_tile(shuffled_do_block_tile, do_block_tile);
+            store_tile(shuffled_do_lds_write_window, shuffled_do_block_tile);
+
+            store_tile(d_lds_write_window, d_block_tile);
+
+            block_sync_lds();
+
+            auto dot_reg_tensor = load_tile(dot_lds_read_window);
+
+            block_sync_lds();
+
+            Policy::template PTFromGemm0CToGemm1A<Problem,
+                                                  decltype(pt_reg_tensor),
+                                                  decltype(p_gemm)>(pt_reg_tensor, p_gemm);
+            gemm_1(dv_acc, pt_reg_tensor, dot_reg_tensor);
+
+            // STAGE 4, OGrad@V Gemm2
+            auto do_reg_tensor = load_tile(do_lds_read_window);
+            auto d             = load_tile(d_lds_read_window);
+            block_sync_lds();
+
+            auto dp_acc = SPGradBlockTileType{};
+
+            dp_acc = gemm_2(do_reg_tensor, v_reg_tensor);
+
+            // STAGE 5, P^T(PGrad^T - D)
+            auto ds                 = SPGradBlockTileType{};
+            constexpr auto ds_spans = decltype(ds)::get_distributed_spans();
+            sweep_tile_span(ds_spans[number<0>{}], [&](auto idx0) {
+                constexpr auto i_idx = make_tuple(idx0);
+                sweep_tile_span(ds_spans[number<1>{}], [&](auto idx1) {
+                    constexpr auto i_j_idx = make_tuple(idx0, idx1);
+                    bool undrop_flag       = p[i_j_idx] >= 0;
+                    ds(i_j_idx)            = p[i_j_idx] * (!FmhaDropout::IsDropout || undrop_flag
+                                                               ? (dp_acc[i_j_idx] - d[i_idx])
+                                                               : d[i_idx]);
+                });
+            });
+
+            if constexpr(kHasBiasGrad)
+            {
+                const auto dbias = [&]() {
+                    if constexpr(FmhaDropout::IsDropout)
+                    {
+                        return tile_elementwise_in(
+                            [&rp_undrop](const auto& x) {
+                                return type_convert<BiasGradDataType>(x * rp_undrop);
+                            },
+                            ds);
+                    }
+                    else
+                    {
+                        return cast_tile<BiasGradDataType>(ds);
+                    }
+                }();
+                store_tile(bias_lds_write_window, dbias);
+                block_sync_lds();
+                auto shuffled_dbias_tile = load_tile(dbias_lds_read_window);
+                auto dbias_tile          = make_static_distributed_tensor<BiasGradDataType>(
+                    Policy::template MakeBiasTileDistribution<Problem>());
+                shuffle_tile(dbias_tile, shuffled_dbias_tile);
+                store_tile(dbias_dram_window, dbias_tile);
+                move_tile_window(dbias_dram_window, {kM0, 0});
+                __builtin_amdgcn_sched_barrier(0);
+            }
+
+            // STAGE 6, SGrad^T@Q^T Gemm3
+            auto qt_reg_tensor = load_tile(qt_lds_read_window);
+            block_sync_lds();
+
+            const auto ds_gemm = cast_tile<GemmDataType>(ds);
+
+            Policy::template SGradTFromGemm2CToGemm3A<Problem,
+                                                      decltype(dst_reg_tensor),
+                                                      decltype(ds_gemm)>(dst_reg_tensor, ds_gemm);
+
+            gemm_3(dk_acc, dst_reg_tensor, qt_reg_tensor);
+
+            store_tile(ds_lds_window, ds_gemm);
+
+            block_sync_lds();
+
+            auto ds_reg_tensor      = load_tile(ds_lds_read_window);
+            auto ds_reg_tensor_next = decltype(ds_reg_tensor){};
+            move_tile_window(ds_lds_read_window, {0, kK4});
+
+            // STAGE7 SGrad@K^T Gemm4
+            auto dq_acc = QGradBlockTileType{};
+            clear_tile(dq_acc);
+
+            static_for<0, k4_loops, 1>{}([&](auto i_k4) {
+                if constexpr(i_k4 < k4_loops - 1)
+                {
+                    ds_reg_tensor_next = load_tile(ds_lds_read_window);
+                    move_tile_window(ds_lds_read_window, {0, kK4});
+                }
+                auto kt_reg_tensor_slice = get_slice_tile(kt_reg_tensor,
+                                                          sequence<0, i_k4 * kK4>{},
+                                                          sequence<kQKHeaddim, (i_k4 + 1) * kK4>{});
+                gemm_4(dq_acc, ds_reg_tensor, kt_reg_tensor_slice);
+
+                if constexpr(i_k4 < k4_loops - 1)
+                {
+                    ds_reg_tensor.get_thread_buffer() = ds_reg_tensor_next.get_thread_buffer();
+                }
+            });
+            move_tile_window(ds_lds_read_window, {0, -kN0});
+            // QGrad Scale
+            if constexpr(FmhaDropout::IsDropout)
+            {
+                tile_elementwise_inout([&scale_rp_undrop](auto& x) { x = x * scale_rp_undrop; },
+                                       dq_acc);
+            }
+            else
+            {
+                tile_elementwise_inout([&raw_scale](auto& x) { x = x * raw_scale; }, dq_acc);
+            }
+            if constexpr(kIsDeterministic)
+            {
+                store_tile(dq_dram_window, dq_acc);
+            }
+            else
+            {
+                update_tile(dq_dram_window, dq_acc);
+            }
+            move_tile_window(dq_dram_window, {kM0, 0});
+
+            i_total_loops += 1;
+            seqlen_q_step += kM0;
+        }
+
+        // Results Scale
+        if constexpr(FmhaDropout::IsDropout)
+        {
+            tile_elementwise_inout([&scale_rp_undrop](auto& x) { x = x * scale_rp_undrop; },
+                                   dk_acc);
+            tile_elementwise_inout([&rp_undrop](auto& x) { x = x * rp_undrop; }, dv_acc);
+        }
+        else
+        {
+            tile_elementwise_inout([&raw_scale](auto& x) { x = x * raw_scale; }, dk_acc);
+        }
+
+        return make_tuple(dk_acc, dv_acc);
+    }
+};
+
+} // namespace ck_tile

include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_ks_kts_vr.hpp

@@ -1,848 +0,0 @@
-// SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
-
-#pragma once
-
-#include "ck_tile/core.hpp"
-#include "ck_tile/ops/fmha/block/block_attention_bias_enum.hpp"
-#include "ck_tile/ops/fmha/block/block_dropout.hpp"
-#include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_ks_kts_vr_default_policy.hpp"
-#include "ck_tile/ops/reduce/block/block_reduce.hpp"
-
-namespace ck_tile {
-
-template <typename Problem, typename Policy = BlockFmhaBwdDQDKDVPipelineKSKTSVRDefaultPolicy>
-struct BlockFmhaBwdDQDKDVPipelineKSKTSVR
-{
-    using QDataType             = remove_cvref_t<typename Problem::QDataType>;
-    using KDataType             = remove_cvref_t<typename Problem::KDataType>;
-    using VDataType             = remove_cvref_t<typename Problem::VDataType>;
-    using GemmDataType          = remove_cvref_t<typename Problem::GemmDataType>;
-    using BiasDataType          = remove_cvref_t<typename Problem::BiasDataType>;
-    using LSEDataType           = remove_cvref_t<typename Problem::LSEDataType>;
-    using AccDataType           = remove_cvref_t<typename Problem::AccDataType>;
-    using DDataType             = remove_cvref_t<typename Problem::DDataType>;
-    using RandValOutputDataType = remove_cvref_t<typename Problem::RandValOutputDataType>;
-    using ODataType             = remove_cvref_t<typename Problem::ODataType>;
-    using OGradDataType         = remove_cvref_t<typename Problem::OGradDataType>;
-    using QGradDataType         = remove_cvref_t<typename Problem::QGradDataType>;
-    using KGradDataType         = remove_cvref_t<typename Problem::KGradDataType>;
-    using VGradDataType         = remove_cvref_t<typename Problem::VGradDataType>;
-    using BiasGradDataType      = remove_cvref_t<typename Problem::BiasGradDataType>;
-    using FmhaMask              = remove_cvref_t<typename Problem::FmhaMask>;
-
-    using BlockFmhaShape = remove_cvref_t<typename Problem::BlockFmhaShape>;
-
-    static constexpr index_t kBlockPerCu = Problem::kBlockPerCu;
-    static constexpr index_t kBlockSize  = Problem::kBlockSize;
-
-    static constexpr index_t kM0        = BlockFmhaShape::kM0;
-    static constexpr index_t kN0        = BlockFmhaShape::kN0;
-    static constexpr index_t kK0        = BlockFmhaShape::kK0;
-    static constexpr index_t kK1        = BlockFmhaShape::kK1;
-    static constexpr index_t kK2        = BlockFmhaShape::kK2;
-    static constexpr index_t kK3        = BlockFmhaShape::kK3;
-    static constexpr index_t kK4        = BlockFmhaShape::kK4;
-    static constexpr index_t kQKHeaddim = BlockFmhaShape::kQKHeaddim;
-    static constexpr index_t kVHeaddim  = BlockFmhaShape::kVHeaddim;
-
-    static constexpr bool kQLoadOnce      = false;
-    static constexpr bool kQTLoadOnce     = false;
-    static constexpr bool kKLoadOnce      = true;
-    static constexpr bool kKTLoadOnce     = true;
-    static constexpr bool kVLoadOnce      = true;
-    static constexpr bool kOGradLoadOnce  = false;
-    static constexpr bool kOGradTLoadOnce = false;
-
-    static constexpr bool kIsGroupMode = Problem::kIsGroupMode;
-    static constexpr bool kPadSeqLenQ  = Problem::kPadSeqLenQ;
-    static constexpr bool kPadSeqLenK  = Problem::kPadSeqLenK;
-    static constexpr bool kPadHeadDimQ = Problem::kPadHeadDimQ;
-    static constexpr bool kPadHeadDimV = Problem::kPadHeadDimV;
-    static constexpr auto BiasEnum     = Problem::BiasEnum;
-    static constexpr bool kHasBiasGrad = Problem::kHasBiasGrad;
-    static constexpr bool kHasDropout  = Problem::kHasDropout;
-
-    // last dimension vector length used to create tensor view(and decide buffer_load vector length)
-    // ... together with tensor distribution. tensor dist should able to overwrite this
-    static constexpr index_t kAlignmentQ =
-        kPadHeadDimQ ? 1 : Policy::template GetAlignmentQ<Problem>();
-    static constexpr index_t kAlignmentK =
-        kPadHeadDimQ ? 1 : Policy::template GetAlignmentK<Problem>();
-    static constexpr index_t kAlignmentV =
-        kPadHeadDimV ? 1 : Policy::template GetAlignmentV<Problem>();
-    static constexpr index_t kAlignmentO =
-        kPadHeadDimV ? 1 : Policy::template GetAlignmentO<Problem>();
-    static constexpr index_t kAlignmentOGrad =
-        kPadHeadDimV ? 1 : Policy::template GetAlignmentOGrad<Problem>();
-    static constexpr index_t kAlignmentQGrad =
-        kPadHeadDimQ ? 2 : Policy::template GetAlignmentQGrad<Problem>();
-    static constexpr index_t kAlignmentKGrad =
-        kPadHeadDimQ ? 1 : Policy::template GetAlignmentKGrad<Problem>();
-    static constexpr index_t kAlignmentVGrad =
-        kPadHeadDimV ? 1 : Policy::template GetAlignmentVGrad<Problem>();
-    static constexpr index_t kAlignmentBias =
-        kPadSeqLenK ? 1 : Policy::template GetTransposedAlignmentBias<Problem>();
-
-    static constexpr const char* name = "ks_kts_vr";
-
-    CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetSmemSize()
-    {
-        return Policy::template GetSmemSize<Problem>();
-    }
-
-    template <typename QDramBlockWindowTmp,
-              typename QTDramBlockWindowTmp,
-              typename KDramBlockWindowTmp,
-              typename KTDramBlockWindowTmp,
-              typename VDramBlockWindowTmp,
-              typename BiasDramBlockWindowTmp,
-              typename RandValDramBlockWindowTmp,
-              typename OGradDramBlockWindowTmp,
-              typename OGradTDramBlockWindowTmp,
-              typename LSEDramBlockWindowTmp,
-              typename DDramBlockWindowTmp,
-              typename QGradDramBlockWindowTmp,
-              typename BiasGradDramBlockWindowTmp,
-              typename PositionEncoding>
-    CK_TILE_HOST_DEVICE auto
-    operator()(const QDramBlockWindowTmp& q_dram_block_window_tmp,
-               const QTDramBlockWindowTmp& qt_dram_block_window_tmp,
-               const KDramBlockWindowTmp& k_dram_block_window_tmp,
-               const KTDramBlockWindowTmp& kt_dram_block_window_tmp,
-               const VDramBlockWindowTmp& v_dram_block_window_tmp,
-               const BiasDramBlockWindowTmp& bias_dram_block_window_tmp,
-               const RandValDramBlockWindowTmp& randval_dram_block_window_tmp,
-               const OGradDramBlockWindowTmp& do_dram_block_window_tmp,
-               const OGradTDramBlockWindowTmp& dot_dram_block_window_tmp,
-               const LSEDramBlockWindowTmp& lse_dram_block_window_tmp,
-               const DDramBlockWindowTmp& d_dram_block_window_tmp,
-               const QGradDramBlockWindowTmp& dq_dram_block_window_tmp,
-               const BiasGradDramBlockWindowTmp& dbias_dram_block_window_tmp,
-               FmhaMask mask,
-               PositionEncoding position_encoding,
-               float raw_scale,
-#if CK_TILE_FMHA_FWD_FAST_EXP2
-               float scale,
-#endif
-               float rp_undrop,
-               float scale_rp_undrop,
-               void* smem_ptr,
-               BlockDropout& dropout) const
-    {
-        static_assert(
-            std::is_same_v<QDataType, remove_cvref_t<typename QDramBlockWindowTmp::DataType>> &&
-                std::is_same_v<QDataType,
-                               remove_cvref_t<typename QTDramBlockWindowTmp::DataType>> &&
-                std::is_same_v<KDataType, remove_cvref_t<typename KDramBlockWindowTmp::DataType>> &&
-                std::is_same_v<KDataType,
-                               remove_cvref_t<typename KTDramBlockWindowTmp::DataType>> &&
-                std::is_same_v<VDataType, remove_cvref_t<typename VDramBlockWindowTmp::DataType>> &&
-                std::is_same_v<OGradDataType,
-                               remove_cvref_t<typename OGradDramBlockWindowTmp::DataType>> &&
-                std::is_same_v<OGradDataType,
-                               remove_cvref_t<typename OGradTDramBlockWindowTmp::DataType>> &&
-                std::is_same_v<LSEDataType,
-                               remove_cvref_t<typename LSEDramBlockWindowTmp::DataType>> &&
-                std::is_same_v<DDataType, remove_cvref_t<typename DDramBlockWindowTmp::DataType>> &&
-                std::is_same_v<QGradDataType,
-                               remove_cvref_t<typename QGradDramBlockWindowTmp::DataType>>,
-            "wrong!");
-
-        static_assert(kM0 == QDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
-                          kQKHeaddim == QTDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
-                          kN0 == KDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
-                          kQKHeaddim == KTDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
-                          kN0 == VDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
-                          kM0 == BiasDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
-                          kN0 == BiasDramBlockWindowTmp{}.get_window_lengths()[number<1>{}] &&
-                          kM0 == OGradDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
-                          kVHeaddim ==
-                              OGradTDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
-                          kM0 == LSEDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
-                          kM0 == DDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
-                          kM0 == QGradDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
-                          kM0 == BiasGradDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
-                          kN0 == BiasGradDramBlockWindowTmp{}.get_window_lengths()[number<1>{}],
-                      "wrong!");
-
-        // Q tile in LDS
-        QDataType* q_lds_ptr = static_cast<QDataType*>(static_cast<void*>(
-            static_cast<char*>(smem_ptr) + Policy::template GetSmemSizeK<Problem>() +
-            Policy::template GetSmemSizeKT<Problem>()));
-        auto q_lds           = make_tensor_view<address_space_enum::lds>(
-            q_lds_ptr, Policy::template MakeQLdsBlockDescriptor<Problem>());
-        auto q_lds_window =
-            make_tile_window(q_lds, make_tuple(number<kM0>{}, number<kK0>{}), {0, 0});
-
-        // QT tile in LDS
-        QDataType* qt_lds_ptr = static_cast<QDataType*>(static_cast<void*>(
-            static_cast<char*>(smem_ptr) + Policy::template GetSmemSizeK<Problem>() +
-            Policy::template GetSmemSizeKT<Problem>()));
-        auto qt_lds           = make_tensor_view<address_space_enum::lds>(
-            qt_lds_ptr, Policy::template MakeQTLdsBlockDescriptor<Problem>());
-        auto qt_lds_window =
-            make_tile_window(qt_lds, make_tuple(number<kQKHeaddim>{}, number<kK3>{}), {0, 0});
-
-        // K tile in LDS
-        auto k_lds = make_tensor_view<address_space_enum::lds>(
-            reinterpret_cast<KDataType*>(smem_ptr),
-            Policy::template MakeKLdsBlockDescriptor<Problem>());
-        auto k_lds_window =
-            make_tile_window(k_lds, make_tuple(number<kN0>{}, number<kQKHeaddim>{}), {0, 0});
-
-        // KT tile in LDS
-        KDataType* kt_lds_ptr = static_cast<KDataType*>(static_cast<void*>(
-            static_cast<char*>(smem_ptr) + Policy::template GetSmemSizeK<Problem>()));
-        auto kt_lds           = make_tensor_view<address_space_enum::lds>(
-            kt_lds_ptr, Policy::template MakeKTLdsBlockDescriptor<Problem>());
-        auto kt_lds_window =
-            make_tile_window(kt_lds, make_tuple(number<kQKHeaddim>{}, number<kN0>{}), {0, 0});
-
-        // OGrad tile in LDS
-        OGradDataType* do_lds_ptr = static_cast<OGradDataType*>(static_cast<void*>(
-            static_cast<char*>(smem_ptr) + Policy::template GetSmemSizeK<Problem>() +
-            Policy::template GetSmemSizeKT<Problem>()));
-        auto do_lds               = make_tensor_view<address_space_enum::lds>(
-            do_lds_ptr, Policy::template MakeOGradLdsBlockDescriptor<Problem>());
-        auto do_lds_window =
-            make_tile_window(do_lds, make_tuple(number<kM0>{}, number<kK2>{}), {0, 0});
-
-        // OGradT tile in LDS
-        OGradDataType* dot_lds_ptr = static_cast<OGradDataType*>(static_cast<void*>(
-            static_cast<char*>(smem_ptr) + Policy::template GetSmemSizeK<Problem>() +
-            Policy::template GetSmemSizeKT<Problem>()));
-        auto dot_lds               = make_tensor_view<address_space_enum::lds>(
-            dot_lds_ptr, Policy::template MakeOGradTLdsBlockDescriptor<Problem>());
-        auto dot_lds_window =
-            make_tile_window(dot_lds, make_tuple(number<kVHeaddim>{}, number<kK1>{}), {0, 0});
-
-        // SGrad tile in LDS
-        GemmDataType* ds_lds_ptr = static_cast<GemmDataType*>(static_cast<void*>(
-            static_cast<char*>(smem_ptr) + Policy::template GetSmemSizeK<Problem>() +
-            Policy::template GetSmemSizeKT<Problem>()));
-        auto ds_lds              = make_tensor_view<address_space_enum::lds>(
-            ds_lds_ptr, Policy::template MakeSGradLdsBlockDescriptor<Problem>());
-        auto ds_lds_window =
-            make_tile_window(ds_lds, make_tuple(number<kM0>{}, number<kN0>{}), {0, 0});
-
-        // BiasT/BiasGradT tile in LDS, use the same size and layout
-        BiasDataType* biast_lds_ptr = static_cast<BiasDataType*>(static_cast<void*>(
-            static_cast<char*>(smem_ptr) + Policy::template GetSmemSizeK<Problem>() +
-            Policy::template GetSmemSizeKT<Problem>()));
-        auto biast_lds              = make_tensor_view<address_space_enum::lds>(
-            biast_lds_ptr, Policy::template MakeBiasTLdsBlockDescriptor<Problem>());
-        auto biast_lds_shuffle_window =
-            make_tile_window(biast_lds, make_tuple(number<kM0>{}, number<kN0>{}), {0, 0});
-        auto dbiast_lds_shuffle_window =
-            make_tile_window(biast_lds,
-                             make_tuple(number<kM0>{}, number<kN0>{}),
-                             {0, 0},
-                             Policy::template MakeShuffledBiasTileDistribution<Problem>());
-
-        static_assert(std::is_same_v<BiasDataType, BiasGradDataType>,
-                      "BiasDataType and BiasGradDataType should be the same!");
-
-        // Block GEMM
-        constexpr auto gemm_0 = Policy::template GetQKBlockGemm<Problem>();
-        constexpr auto gemm_1 = Policy::template GetPTOGradTBlockGemm<Problem>();
-        constexpr auto gemm_2 = Policy::template GetOGradVBlockGemm<Problem>();
-        constexpr auto gemm_3 = Policy::template GetSGradTQTBlockGemm<Problem>();
-        constexpr auto gemm_4 = Policy::template GetSGradKTBlockGemm<Problem>();
-
-        auto v_dram_window = make_tile_window(
-            v_dram_block_window_tmp.get_bottom_tensor_view(),
-            v_dram_block_window_tmp.get_window_lengths(),
-            v_dram_block_window_tmp.get_window_origin(),
-            Policy::template MakeVInRegDramTileDistribution<Problem, decltype(gemm_2)>());
-
-        auto v = load_tile(v_dram_window); // persistent V register tile
-
-        using SPTBlockTileType     = decltype(gemm_0.MakeCBlockTile());
-        using SPGradTBlockTileType = decltype(gemm_2.MakeCBlockTile());
-        using QGradBlockTileType   = decltype(gemm_4.MakeCBlockTile());
-
-        // init VGrad & KGrad
-        auto dv_acc = decltype(gemm_1.MakeCBlockTile()){};
-        auto dk_acc = decltype(gemm_3.MakeCBlockTile()){};
-
-        clear_tile(dv_acc);
-        clear_tile(dk_acc);
-
-        auto k_dram_window = make_tile_window(
-            k_dram_block_window_tmp.get_bottom_tensor_view(),
-            k_dram_block_window_tmp.get_window_lengths(),
-            k_dram_block_window_tmp.get_window_origin(),
-            Policy::template MakeKDramTileDistribution<Problem>()); // K DRAM tile window for
-                                                                    // load
-
-        __builtin_amdgcn_sched_barrier(0);
-        const auto k_origin = k_dram_window.get_window_origin();
-        const auto [seqlen_q_start, seqlen_q_end] =
-            mask.GetTileRangeAlongY(k_origin.at(number<0>{}), number<kM0>{}, number<kN0>{});
-
-        const auto num_total_loop = integer_divide_ceil(seqlen_q_end - seqlen_q_start, kM0);
-
-        // check early exit if masked and no work to do.
-        if constexpr(FmhaMask::IsMasking)
-        {
-            if(num_total_loop <= 0)
-            {
-                // Note: here dk_acc&dv_acc are all cleard, return it
-                // Note: v loaded but no fence, ignore it.
-                return ck_tile::make_tuple(dk_acc, dv_acc);
-            }
-        }
-
-        auto k_block_tile = load_tile(k_dram_window);
-
-        store_tile(k_lds_window, k_block_tile); // // persistent K in LDS
-
-        auto kt_dram_block_window = kt_dram_block_window_tmp;
-
-        auto kt_dram_window = make_tile_window(
-            kt_dram_block_window.get_bottom_tensor_view(),
-            kt_dram_block_window.get_window_lengths(),
-            kt_dram_block_window.get_window_origin(),
-            Policy::template MakeKTDramTileDistribution<Problem>()); // K^T DRAM tile window for
-                                                                     // load
-
-        auto kt_block_tile = load_tile(kt_dram_window);
-
-        auto kt_shuffle_tmp = make_static_distributed_tensor<KDataType>(
-            Policy::template MakeShuffledKTRegBlockDescriptor<Problem>());
-        shuffle_tile(kt_shuffle_tmp, kt_block_tile);
-
-        store_tile(kt_lds_window, kt_shuffle_tmp); // persistent K^T in LDS
-
-        auto q_dram_block_window =
-            make_tile_window(q_dram_block_window_tmp.get_bottom_tensor_view(),
-                             q_dram_block_window_tmp.get_window_lengths(),
-                             {seqlen_q_start, 0});
-
-        auto qt_dram_block_window =
-            make_tile_window(qt_dram_block_window_tmp.get_bottom_tensor_view(),
-                             qt_dram_block_window_tmp.get_window_lengths(),
-                             {0, seqlen_q_start});
-
-        auto do_dram_block_window =
-            make_tile_window(do_dram_block_window_tmp.get_bottom_tensor_view(),
-                             do_dram_block_window_tmp.get_window_lengths(),
-                             {seqlen_q_start, 0});
-
-        auto dot_dram_block_window =
-            make_tile_window(dot_dram_block_window_tmp.get_bottom_tensor_view(),
-                             dot_dram_block_window_tmp.get_window_lengths(),
-                             {0, seqlen_q_start});
-
-        auto dq_dram_block_window =
-            make_tile_window(dq_dram_block_window_tmp.get_bottom_tensor_view(),
-                             dq_dram_block_window_tmp.get_window_lengths(),
-                             {seqlen_q_start, 0});
-
-        auto lse_dram_block_window =
-            make_tile_window(lse_dram_block_window_tmp.get_bottom_tensor_view(),
-                             lse_dram_block_window_tmp.get_window_lengths(),
-                             {seqlen_q_start});
-
-        auto d_dram_block_window =
-            make_tile_window(d_dram_block_window_tmp.get_bottom_tensor_view(),
-                             d_dram_block_window_tmp.get_window_lengths(),
-                             {seqlen_q_start});
-
-        const auto bias_origin = bias_dram_block_window_tmp.get_window_origin();
-        auto bias_dram_block_window =
-            make_tile_window(bias_dram_block_window_tmp.get_bottom_tensor_view(),
-                             bias_dram_block_window_tmp.get_window_lengths(),
-                             {seqlen_q_start, bias_origin.at(number<1>{})}); // M/N
-
-        const auto dbias_origin = dbias_dram_block_window_tmp.get_window_origin();
-        auto dbias_dram_block_window =
-            make_tile_window(dbias_dram_block_window_tmp.get_bottom_tensor_view(),
-                             dbias_dram_block_window_tmp.get_window_lengths(),
-                             {seqlen_q_start, dbias_origin.at(number<1>{})}); // M/N
-
-        auto qt_dram_window =
-            make_tile_window(qt_dram_block_window.get_bottom_tensor_view(),
-                             qt_dram_block_window.get_window_lengths(),
-                             qt_dram_block_window.get_window_origin(),
-                             Policy::template MakeQTDramTileDistribution<Problem>());
-
-        auto dot_dram_window =
-            make_tile_window(dot_dram_block_window.get_bottom_tensor_view(),
-                             dot_dram_block_window.get_window_lengths(),
-                             dot_dram_block_window.get_window_origin(),
-                             Policy::template MakeOGradTDramTileDistribution<Problem>());
-
-        auto lse_dram_window = make_tile_window(
-            lse_dram_block_window.get_bottom_tensor_view(),
-            lse_dram_block_window.get_window_lengths(),
-            lse_dram_block_window.get_window_origin(),
-            Policy::template MakeLSEDDramTileDistribution<Problem, decltype(gemm_0)>());
-
-        auto d_dram_window = make_tile_window(
-            d_dram_block_window.get_bottom_tensor_view(),
-            d_dram_block_window.get_window_lengths(),
-            d_dram_block_window.get_window_origin(),
-            Policy::template MakeLSEDDramTileDistribution<Problem, decltype(gemm_0)>());
-
-        auto bias_dram_window =
-            make_tile_window(bias_dram_block_window.get_bottom_tensor_view(),
-                             bias_dram_block_window.get_window_lengths(),
-                             bias_dram_block_window.get_window_origin(),
-                             Policy::template MakeBiasTileDistribution<Problem>());
-
-        auto biast_lds_window =
-            make_tile_window(biast_lds_shuffle_window.get_bottom_tensor_view(),
-                             biast_lds_shuffle_window.get_window_lengths(),
-                             biast_lds_shuffle_window.get_window_origin(),
-                             Policy::template MakeBiasTTileDistribution<decltype(gemm_0)>());
-
-        auto randval_dram_window = dropout.MakeRandvalDramWindow<decltype(gemm_0), false>(
-            randval_dram_block_window_tmp, seqlen_q_start);
-
-        index_t i_total_loops      = 0;
-        constexpr index_t k0_loops = kQKHeaddim / kK0;
-        constexpr index_t k1_loops = kM0 / kK1;
-        constexpr index_t k2_loops = kVHeaddim / kK2;
-        constexpr index_t k3_loops = kM0 / kK3;
-        constexpr index_t k4_loops = kN0 / kK4;
-        do
-        {
-            auto q_dram_window = make_tile_window(
-                q_dram_block_window.get_bottom_tensor_view(),
-                q_dram_block_window.get_window_lengths(),
-                q_dram_block_window.get_window_origin(),
-                Policy::template MakeQDramTileDistribution<Problem>()); // Q DRAM tile window for
-                                                                        // load
-
-            auto do_dram_window = make_tile_window(
-                do_dram_block_window.get_bottom_tensor_view(),
-                do_dram_block_window.get_window_lengths(),
-                do_dram_block_window.get_window_origin(),
-                Policy::template MakeOGradDramTileDistribution<Problem>()); // OGrad DRAM tile
-                                                                            // window for load
-
-            // STAGE 1, Q@K Gemm0
-            auto st_acc = SPTBlockTileType{};
-
-            auto q_block_tile = load_tile(q_dram_window);
-            {
-                move_tile_window(q_dram_window, {0, kK0});
-
-                clear_tile(st_acc); // Initialize S^T
-
-                store_tile(q_lds_window, q_block_tile);  // LDS write 0
-                q_block_tile = load_tile(q_dram_window); // global read 1
-            }
-
-            if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS)
-            {
-                __builtin_amdgcn_sched_barrier(
-                    0); // prevent from messing up the order of global loads
-            }
-            const auto bias_tile = load_tile(bias_dram_window); // load bias tile
-            if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS)
-            {
-                __builtin_amdgcn_sched_barrier(
-                    0); // prevent from messing up the order of global loads
-            }
-
-            if constexpr(k0_loops > 2)
-            {
-                static_for<0, k0_loops - 2, 1>{}([&](auto i_k0) {
-                    block_sync_lds();
-                    gemm_0(st_acc,
-                           q_lds_window,
-                           get_slice_tile(k_lds_window,
-                                          sequence<0, i_k0 * kK0>{},
-                                          sequence<kN0, (i_k0 + 1) * kK0>{}));
-                    block_sync_lds();
-                    move_tile_window(q_dram_window, {0, kK0});
-
-                    store_tile(q_lds_window,
-                               q_block_tile);                // LDS write i + 1
-                    q_block_tile = load_tile(q_dram_window); // global read i + 2
-                });
-            }
-
-            const auto dot_prefetch = load_tile(dot_dram_window); // prefetch load OGrad^T tile
-            {                                                     // tail
-                block_sync_lds();
-                gemm_0(st_acc,
-                       q_lds_window,
-                       get_slice_tile(k_lds_window,
-                                      sequence<0, (k0_loops - 2) * kK0>{},
-                                      sequence<kN0, (k0_loops - 1) * kK0>{}));
-                block_sync_lds();
-
-                store_tile(q_lds_window, q_block_tile);
-                block_sync_lds();
-
-                gemm_0(st_acc,
-                       q_lds_window,
-                       get_slice_tile(k_lds_window,
-                                      sequence<0, (k0_loops - 1) * kK0>{},
-                                      sequence<kN0, k0_loops * kK0>{}));
-            }
-
-            // STAGE 2, Scale, Add bias, Mask, Softmax, Dropout
-            if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS)
-            {
-                block_sync_lds();
-                auto bias_shuffle_tmp = make_static_distributed_tensor<BiasDataType>(
-                    Policy::template MakeShuffledBiasTileDistribution<Problem>());
-                shuffle_tile(bias_shuffle_tmp, bias_tile);
-                store_tile(biast_lds_shuffle_window, bias_shuffle_tmp);
-                block_sync_lds();
-                auto biast_tile = load_tile(biast_lds_window);
-                tile_elementwise_inout(
-                    [&](auto& x, const auto& y) {
-#if !CK_TILE_FMHA_FWD_FAST_EXP2
-                        x = raw_scale * x + type_convert<AccDataType>(y);
-#else
-                        x = scale * x + log2e_v<AccDataType> * type_convert<AccDataType>(y);
-#endif
-                    },
-                    st_acc,
-                    biast_tile);
-                move_tile_window(bias_dram_window, {kM0, 0});
-            }
-            else if constexpr(BiasEnum == BlockAttentionBiasEnum::ALIBI)
-            {
-                const auto q_origin     = q_dram_block_window.get_window_origin();
-                constexpr auto st_spans = decltype(st_acc)::get_distributed_spans();
-                sweep_tile_span(st_spans[number<0>{}], [&](auto idx0) {
-                    sweep_tile_span(st_spans[number<1>{}], [&](auto idx1) {
-                        const auto tile_idx = get_x_indices_from_distributed_indices(
-                            st_acc.get_tile_distribution(), make_tuple(idx0, idx1));
-
-                        const auto row = q_origin.at(number<0>{}) + tile_idx.at(number<0>{});
-                        const auto col = k_origin.at(number<0>{}) + tile_idx.at(number<1>{});
-                        constexpr auto i_j_idx = make_tuple(idx0, idx1);
-
-#if !CK_TILE_FMHA_FWD_FAST_EXP2
-                        st_acc(i_j_idx) *= raw_scale;
-#else
-                        st_acc(i_j_idx) *= scale;
-#endif
-                        position_encoding.update(st_acc(i_j_idx), row, col);
-                    });
-                });
-            }
-            else
-            {
-#if !CK_TILE_FMHA_FWD_FAST_EXP2
-                tile_elementwise_inout([&raw_scale](auto& x) { x = x * raw_scale; }, st_acc);
-#endif
-            }
-
-            if constexpr(kPadSeqLenK || FmhaMask::IsMasking)
-            {
-                const auto q_origin      = q_dram_block_window.get_window_origin();
-                bool need_perpixel_check = mask.IsEdgeTile(q_origin.at(number<0>{}),
-                                                           k_origin.at(number<0>{}),
-                                                           number<kM0>{},
-                                                           number<kN0>{});
-                if(need_perpixel_check)
-                {
-                    set_tile_if(st_acc, -numeric<AccDataType>::infinity(), [&](auto tile_idx) {
-                        const auto row = q_origin.at(number<0>{}) + tile_idx.at(number<0>{});
-                        const auto col = k_origin.at(number<0>{}) + tile_idx.at(number<1>{});
-                        return mask.IsOutOfBound(row, col);
-                    });
-                }
-            }
-
-            const auto lse = load_tile(lse_dram_window);
-
-            static const auto get_validated_lse = [](LSEDataType raw_lse) {
-                if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS ||
-                             FmhaMask::IsMasking)
-                {
-                    return raw_lse == -numeric<LSEDataType>::infinity()
-                               ? type_convert<LSEDataType>(0.f)
-                               : raw_lse;
-                }
-                else
-                {
-                    return raw_lse;
-                }
-            };
-
-            auto pt                 = SPTBlockTileType{};
-            constexpr auto pt_spans = decltype(pt)::get_distributed_spans();
-            sweep_tile_span(pt_spans[number<0>{}], [&](auto idx0) {
-                constexpr auto i_idx = make_tuple(idx0);
-#if CK_TILE_FMHA_FWD_FAST_EXP2
-                auto row_lse = log2e_v<LSEDataType> * get_validated_lse(lse[i_idx]);
-#endif
-                sweep_tile_span(pt_spans[number<1>{}], [&](auto idx1) {
-                    constexpr auto i_j_idx = make_tuple(idx0, idx1);
-#if CK_TILE_FMHA_FWD_FAST_EXP2
-                    if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS ||
-                                 BiasEnum == BlockAttentionBiasEnum::ALIBI)
-                    {
-                        pt(i_j_idx) = exp2(st_acc[i_j_idx] - row_lse);
-                    }
-                    else
-                    {
-                        pt(i_j_idx) = exp2(scale * st_acc[i_j_idx] - row_lse);
-                    }
-#else
-                    pt(i_j_idx) = exp(st_acc[i_j_idx] - get_validated_lse(lse[i_idx]));
-#endif
-                });
-            });
-
-            auto dot_shuffle_tmp = make_static_distributed_tensor<OGradDataType>(
-                Policy::template MakeShuffledOGradTRegBlockDescriptor<Problem>());
-            block_sync_lds();
-            {
-                shuffle_tile(dot_shuffle_tmp, dot_prefetch);
-                store_tile(dot_lds_window,
-                           dot_shuffle_tmp); // store the prefetch
-            }
-            move_tile_window(dot_dram_window, {0, kK1});
-
-            if constexpr(kHasDropout)
-            {
-                dropout.Run<decltype(gemm_0), RandValOutputDataType>(
-                    seqlen_q_start + i_total_loops * kM0, pt, randval_dram_window);
-            }
-
-            // STAGE 3, P^T@OGrad^T Gemm1
-            const auto pt_gemm = [&]() {
-                if constexpr(kHasDropout)
-                {
-                    return tile_elementwise_in(
-                        [](const auto& x) { return type_convert<GemmDataType>(x > 0.f ? x : 0.f); },
-                        pt);
-                }
-                else
-                {
-                    return cast_tile<GemmDataType>(pt);
-                }
-            }();
-
-            if constexpr(k1_loops > 1)
-            {
-                static_for<0, k1_loops - 1, 1>{}([&](auto i_k1) {
-                    const auto dot = load_tile(dot_dram_window); // load next OGrad^T
-                    block_sync_lds();
-                    gemm_1(dv_acc,
-                           get_slice_tile(pt_gemm,
-                                          sequence<i_k1 * kK1, 0>{},
-                                          sequence<(i_k1 + 1) * kK1, kN0>{}),
-                           dot_lds_window);
-                    block_sync_lds();
-                    shuffle_tile(dot_shuffle_tmp, dot);
-                    store_tile(dot_lds_window,
-                               dot_shuffle_tmp); // store the prefetch
-
-                    move_tile_window(dot_dram_window, {0, kK1});
-                });
-            }
-            auto do_block_tile = load_tile(do_dram_window); // prefetch load OGrad tile
-            // tail
-            {
-                block_sync_lds();
-                gemm_1(dv_acc,
-                       get_slice_tile(
-                           pt_gemm, sequence<(k1_loops - 1) * kK1, 0>{}, sequence<kM0, kN0>{}),
-                       dot_lds_window);
-                block_sync_lds();
-            }
-
-            // STAGE 4, OGrad@V Gemm2
-            auto dpt_acc = SPGradTBlockTileType{};
-
-            {
-                move_tile_window(do_dram_window, {0, kK2});
-
-                clear_tile(dpt_acc); // Initialize PGrad^T
-
-                store_tile(do_lds_window, do_block_tile);  // LDS write 0
-                do_block_tile = load_tile(do_dram_window); // global read 1
-            }
-
-            if constexpr(k2_loops > 2)
-            {
-                static_for<0, k2_loops - 2, 1>{}([&](auto i_k2) {
-                    block_sync_lds();
-                    gemm_2(dpt_acc,
-                           do_lds_window,
-                           get_slice_tile(
-                               v, sequence<0, i_k2 * kK2>{}, sequence<kN0, (i_k2 + 1) * kK2>{}));
-                    block_sync_lds();
-                    move_tile_window(do_dram_window, {0, kK2});
-
-                    store_tile(do_lds_window,
-                               do_block_tile);                 // LDS write i + 1
-                    do_block_tile = load_tile(do_dram_window); // global read i + 2
-                });
-            }
-
-            const auto qt_prefetch = load_tile(qt_dram_window); // prefetch load Q^T tile
-            {                                                   // tail
-                block_sync_lds();
-                gemm_2(dpt_acc,
-                       do_lds_window,
-                       get_slice_tile(v,
-                                      sequence<0, (k2_loops - 2) * kK2>{},
-                                      sequence<kN0, (k2_loops - 1) * kK2>{}));
-                block_sync_lds();
-
-                store_tile(do_lds_window, do_block_tile);
-                block_sync_lds();
-
-                gemm_2(dpt_acc,
-                       do_lds_window,
-                       get_slice_tile(v,
-                                      sequence<0, (k2_loops - 1) * kK2>{},
-                                      sequence<kN0, k2_loops * kK2>{}));
-            }
-
-            // STAGE 5, P^T(PGrad^T - D)
-            const auto d = load_tile(d_dram_window);
-
-            auto dst                 = SPGradTBlockTileType{};
-            constexpr auto dst_spans = decltype(dst)::get_distributed_spans();
-            sweep_tile_span(dst_spans[number<0>{}], [&](auto idx0) {
-                constexpr auto i_idx = make_tuple(idx0);
-                sweep_tile_span(dst_spans[number<1>{}], [&](auto idx1) {
-                    constexpr auto i_j_idx = make_tuple(idx0, idx1);
-                    bool undrop_flag       = pt[i_j_idx] >= 0;
-                    dst(i_j_idx) =
-                        pt[i_j_idx] *
-                        (!kHasDropout || undrop_flag ? (dpt_acc[i_j_idx] - d[i_idx]) : d[i_idx]);
-                });
-            });
-
-            if constexpr(kHasBiasGrad)
-            {
-                const auto dbiast = [&]() {
-                    if constexpr(kHasDropout)
-                    {
-                        return tile_elementwise_in(
-                            [&rp_undrop](const auto& x) {
-                                return type_convert<BiasGradDataType>(x * rp_undrop);
-                            },
-                            dst);
-                    }
-                    else
-                    {
-                        return cast_tile<BiasGradDataType>(dst);
-                    }
-                }();
-                store_tile(biast_lds_shuffle_window, dbiast);
-                block_sync_lds();
-                auto dbiast_tile        = load_tile(dbiast_lds_shuffle_window);
-                auto dbiast_shuffle_tmp = make_static_distributed_tensor<BiasGradDataType>(
-                    Policy::template MakeBiasTileDistribution<Problem>());
-                shuffle_tile(dbiast_shuffle_tmp, dbiast_tile);
-                store_tile(dbias_dram_block_window, dbiast_shuffle_tmp);
-                move_tile_window(dbias_dram_block_window, {kM0, 0});
-            }
-
-            // STAGE 6, SGrad^T@Q^T Gemm3
-            auto qt_shuffle_tmp = make_static_distributed_tensor<QDataType>(
-                Policy::template MakeShuffledQTRegBlockDescriptor<Problem>());
-            block_sync_lds();
-            {
-                shuffle_tile(qt_shuffle_tmp, qt_prefetch);
-                store_tile(qt_lds_window,
-                           qt_shuffle_tmp); // store the prefetch
-            }
-            move_tile_window(qt_dram_window, {0, kK3});
-
-            const auto dst_gemm = cast_tile<GemmDataType>(dst);
-
-            if constexpr(k3_loops > 1)
-            {
-                static_for<0, k3_loops - 1, 1>{}([&](auto i_k3) {
-                    const auto qt = load_tile(qt_dram_window); // load next Q^T
-                    block_sync_lds();
-                    gemm_3(dk_acc,
-                           get_slice_tile(dst_gemm,
-                                          sequence<i_k3 * kK3, 0>{},
-                                          sequence<(i_k3 + 1) * kK3, kN0>{}),
-                           qt_lds_window);
-                    block_sync_lds();
-                    shuffle_tile(qt_shuffle_tmp, qt);
-                    store_tile(qt_lds_window,
-                               qt_shuffle_tmp); // store the prefetch
-
-                    move_tile_window(qt_dram_window, {0, kK3});
-                });
-            }
-            // tail
-            {
-                block_sync_lds();
-                gemm_3(dk_acc,
-                       get_slice_tile(
-                           dst_gemm, sequence<(k3_loops - 1) * kK3, 0>{}, sequence<kM0, kN0>{}),
-                       qt_lds_window);
-                block_sync_lds();
-            }
-
-            // STAGE 7, SGrad@K^T Gemm4
-            store_tile(ds_lds_window, dst_gemm);
-
-            auto dq_acc = QGradBlockTileType{};
-            clear_tile(dq_acc); // Initialize QGrad
-
-            block_sync_lds();
-
-            static_for<0, k4_loops, 1>{}([&](auto i_k4) {
-                gemm_4(dq_acc,
-                       get_slice_tile(ds_lds_window,
-                                      sequence<0, i_k4 * kK4>{},
-                                      sequence<kM0, (i_k4 + 1) * kK4>{}),
-                       get_slice_tile(kt_lds_window,
-                                      sequence<0, i_k4 * kK4>{},
-                                      sequence<kQKHeaddim, (i_k4 + 1) * kK4>{}));
-            });
-
-            // QGrad Scale
-            if constexpr(kHasDropout)
-            {
-                tile_elementwise_inout([&scale_rp_undrop](auto& x) { x = x * scale_rp_undrop; },
-                                       dq_acc);
-            }
-            else
-            {
-                tile_elementwise_inout([&raw_scale](auto& x) { x = x * raw_scale; }, dq_acc);
-            }
-            const auto dq = cast_tile<QGradDataType>(dq_acc);
-            update_tile(dq_dram_block_window, dq);
-
-            // move tile windows
-            move_tile_window(q_dram_block_window, {kM0, 0});
-            move_tile_window(dq_dram_block_window, {kM0, 0});
-            move_tile_window(do_dram_block_window, {kM0, 0});
-            move_tile_window(lse_dram_window, {kM0});
-            move_tile_window(d_dram_window, {kM0});
-        } while(++i_total_loops < num_total_loop);
-
-        // KGrad Scale
-        if constexpr(kHasDropout)
-        {
-            tile_elementwise_inout([&scale_rp_undrop](auto& x) { x = x * scale_rp_undrop; },
-                                   dk_acc);
-        }
-        else
-        {
-            tile_elementwise_inout([&raw_scale](auto& x) { x = x * raw_scale; }, dk_acc);
-        }
-        // VGrad Scale
-        if constexpr(kHasDropout)
-        {
-            tile_elementwise_inout([&rp_undrop](auto& x) { x = x * rp_undrop; }, dv_acc);
-        }
-
-        return ck_tile::make_tuple(dk_acc, dv_acc);
-    }
-};
-
-} // namespace ck_tile

include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_ks_kts_vr_default_policy.hpp

@@ -1,20 +0,0 @@
-// SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
-
-#pragma once
-
-#include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_default_policy.hpp"
-
-namespace ck_tile {
-
-// This pipeline is v located in regs, k & k^t located in lds.
-using BlockFmhaBwdDQDKDVPipelineKSKTSVRDefaultPolicy =
-    BlockFmhaBwdPipelineDefaultPolicy</* QLoadOnce_      = */ false,
-                                      /* QTLoadOnce_     = */ false,
-                                      /* KLoadOnce_      = */ true,
-                                      /* KTLoadOnce_     = */ true,
-                                      /* VLoadOnce_      = */ true,
-                                      /* OGradLoadOnce_  = */ false,
-                                      /* OGradTLoadOnce_ = */ false>;
-
-} // namespace ck_tile

include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_ks_vr.hpp

@@ -1,821 +0,0 @@
-// SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
-
-#pragma once
-
-#include "ck_tile/core.hpp"
-#include "ck_tile/ops/fmha/block/block_attention_bias_enum.hpp"
-#include "ck_tile/ops/fmha/block/block_dropout.hpp"
-#include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_ks_vr_default_policy.hpp"
-#include "ck_tile/ops/reduce/block/block_reduce.hpp"
-
-namespace ck_tile {
-
-template <typename Problem, typename Policy = BlockFmhaBwdDQDKDVPipelineKSVRDefaultPolicy>
-struct BlockFmhaBwdDQDKDVPipelineKSVR
-{
-    using QDataType             = remove_cvref_t<typename Problem::QDataType>;
-    using KDataType             = remove_cvref_t<typename Problem::KDataType>;
-    using VDataType             = remove_cvref_t<typename Problem::VDataType>;
-    using GemmDataType          = remove_cvref_t<typename Problem::GemmDataType>;
-    using BiasDataType          = remove_cvref_t<typename Problem::BiasDataType>;
-    using LSEDataType           = remove_cvref_t<typename Problem::LSEDataType>;
-    using AccDataType           = remove_cvref_t<typename Problem::AccDataType>;
-    using DDataType             = remove_cvref_t<typename Problem::DDataType>;
-    using RandValOutputDataType = remove_cvref_t<typename Problem::RandValOutputDataType>;
-    using ODataType             = remove_cvref_t<typename Problem::ODataType>;
-    using OGradDataType         = remove_cvref_t<typename Problem::OGradDataType>;
-    using QGradDataType         = remove_cvref_t<typename Problem::QGradDataType>;
-    using KGradDataType         = remove_cvref_t<typename Problem::KGradDataType>;
-    using VGradDataType         = remove_cvref_t<typename Problem::VGradDataType>;
-    using BiasGradDataType      = remove_cvref_t<typename Problem::BiasGradDataType>;
-    using FmhaMask              = remove_cvref_t<typename Problem::FmhaMask>;
-
-    using BlockFmhaShape = remove_cvref_t<typename Problem::BlockFmhaShape>;
-
-    static constexpr index_t kBlockPerCu = Problem::kBlockPerCu;
-    static constexpr index_t kBlockSize  = Problem::kBlockSize;
-
-    static constexpr index_t kM0        = BlockFmhaShape::kM0;
-    static constexpr index_t kN0        = BlockFmhaShape::kN0;
-    static constexpr index_t kK0        = BlockFmhaShape::kK0;
-    static constexpr index_t kK1        = BlockFmhaShape::kK1;
-    static constexpr index_t kK2        = BlockFmhaShape::kK2;
-    static constexpr index_t kK3        = BlockFmhaShape::kK3;
-    static constexpr index_t kK4        = BlockFmhaShape::kK4;
-    static constexpr index_t kQKHeaddim = BlockFmhaShape::kQKHeaddim;
-    static constexpr index_t kVHeaddim  = BlockFmhaShape::kVHeaddim;
-
-    static constexpr bool kQLoadOnce      = false;
-    static constexpr bool kQTLoadOnce     = false;
-    static constexpr bool kKLoadOnce      = true;
-    static constexpr bool kKTLoadOnce     = false;
-    static constexpr bool kVLoadOnce      = true;
-    static constexpr bool kOGradLoadOnce  = false;
-    static constexpr bool kOGradTLoadOnce = false;
-
-    static constexpr bool kIsGroupMode = Problem::kIsGroupMode;
-    static constexpr bool kPadSeqLenQ  = Problem::kPadSeqLenQ;
-    static constexpr bool kPadSeqLenK  = Problem::kPadSeqLenK;
-    static constexpr bool kPadHeadDimQ = Problem::kPadHeadDimQ;
-    static constexpr bool kPadHeadDimV = Problem::kPadHeadDimV;
-    static constexpr auto BiasEnum     = Problem::BiasEnum;
-    static constexpr bool kHasBiasGrad = Problem::kHasBiasGrad;
-    static constexpr bool kHasDropout  = Problem::kHasDropout;
-
-    // last dimension vector length used to create tensor view(and decide buffer_load vector length)
-    // ... together with tensor distribution. tensor dist should able to overwrite this
-    static constexpr index_t kAlignmentQ =
-        kPadHeadDimQ ? 1 : Policy::template GetAlignmentQ<Problem>();
-    static constexpr index_t kAlignmentK =
-        kPadHeadDimQ ? 1 : Policy::template GetAlignmentK<Problem>();
-    static constexpr index_t kAlignmentV =
-        kPadHeadDimV ? 1 : Policy::template GetAlignmentV<Problem>();
-    static constexpr index_t kAlignmentO =
-        kPadHeadDimV ? 1 : Policy::template GetAlignmentO<Problem>();
-    static constexpr index_t kAlignmentOGrad =
-        kPadHeadDimV ? 1 : Policy::template GetAlignmentOGrad<Problem>();
-    static constexpr index_t kAlignmentQGrad =
-        kPadHeadDimQ ? 2 : Policy::template GetAlignmentQGrad<Problem>();
-    static constexpr index_t kAlignmentKGrad =
-        kPadHeadDimQ ? 1 : Policy::template GetAlignmentKGrad<Problem>();
-    static constexpr index_t kAlignmentVGrad =
-        kPadHeadDimV ? 1 : Policy::template GetAlignmentVGrad<Problem>();
-    static constexpr index_t kAlignmentBias =
-        kPadSeqLenK ? 1 : Policy::template GetTransposedAlignmentBias<Problem>();
-
-    static constexpr const char* name = "ks_vr";
-
-    CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetSmemSize()
-    {
-        return Policy::template GetSmemSize<Problem>();
-    }
-
-    template <typename QDramBlockWindowTmp,
-              typename QTDramBlockWindowTmp,
-              typename KDramBlockWindowTmp,
-              typename KTDramBlockWindowTmp,
-              typename VDramBlockWindowTmp,
-              typename BiasDramBlockWindowTmp,
-              typename RandValDramBlockWindowTmp,
-              typename OGradDramBlockWindowTmp,
-              typename OGradTDramBlockWindowTmp,
-              typename LSEDramBlockWindowTmp,
-              typename DDramBlockWindowTmp,
-              typename QGradDramBlockWindowTmp,
-              typename BiasGradDramBlockWindowTmp,
-              typename PositionEncoding>
-    CK_TILE_HOST_DEVICE auto
-    operator()(const QDramBlockWindowTmp& q_dram_block_window_tmp,
-               const QTDramBlockWindowTmp& qt_dram_block_window_tmp,
-               const KDramBlockWindowTmp& k_dram_block_window_tmp,
-               const KTDramBlockWindowTmp& /*kt_dram_block_window_tmp*/,
-               const VDramBlockWindowTmp& v_dram_block_window_tmp,
-               const BiasDramBlockWindowTmp& bias_dram_block_window_tmp,
-               const RandValDramBlockWindowTmp& randval_dram_block_window_tmp,
-               const OGradDramBlockWindowTmp& do_dram_block_window_tmp,
-               const OGradTDramBlockWindowTmp& dot_dram_block_window_tmp,
-               const LSEDramBlockWindowTmp& lse_dram_block_window_tmp,
-               const DDramBlockWindowTmp& d_dram_block_window_tmp,
-               const QGradDramBlockWindowTmp& dq_dram_block_window_tmp,
-               const BiasGradDramBlockWindowTmp& dbias_dram_block_window_tmp,
-               FmhaMask mask,
-               PositionEncoding position_encoding,
-               float raw_scale,
-#if CK_TILE_FMHA_FWD_FAST_EXP2
-               float scale,
-#endif
-               float rp_undrop,
-               float scale_rp_undrop,
-               void* smem_ptr,
-               BlockDropout& dropout) const
-    {
-        static_assert(
-            std::is_same_v<QDataType, remove_cvref_t<typename QDramBlockWindowTmp::DataType>> &&
-                std::is_same_v<QDataType,
-                               remove_cvref_t<typename QTDramBlockWindowTmp::DataType>> &&
-                std::is_same_v<KDataType, remove_cvref_t<typename KDramBlockWindowTmp::DataType>> &&
-                std::is_same_v<VDataType, remove_cvref_t<typename VDramBlockWindowTmp::DataType>> &&
-                std::is_same_v<OGradDataType,
-                               remove_cvref_t<typename OGradDramBlockWindowTmp::DataType>> &&
-                std::is_same_v<OGradDataType,
-                               remove_cvref_t<typename OGradTDramBlockWindowTmp::DataType>> &&
-                std::is_same_v<LSEDataType,
-                               remove_cvref_t<typename LSEDramBlockWindowTmp::DataType>> &&
-                std::is_same_v<DDataType, remove_cvref_t<typename DDramBlockWindowTmp::DataType>> &&
-                std::is_same_v<QGradDataType,
-                               remove_cvref_t<typename QGradDramBlockWindowTmp::DataType>>,
-            "wrong!");
-
-        static_assert(kM0 == QDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
-                          kQKHeaddim == QTDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
-                          kN0 == KDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
-                          kN0 == VDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
-                          kM0 == BiasDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
-                          kN0 == BiasDramBlockWindowTmp{}.get_window_lengths()[number<1>{}] &&
-                          kM0 == OGradDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
-                          kVHeaddim ==
-                              OGradTDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
-                          kM0 == LSEDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
-                          kM0 == DDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
-                          kM0 == QGradDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
-                          kM0 == BiasGradDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
-                          kN0 == BiasGradDramBlockWindowTmp{}.get_window_lengths()[number<1>{}],
-                      "wrong!");
-
-        // Q tile in LDS
-        QDataType* q_lds_ptr = static_cast<QDataType*>(static_cast<void*>(
-            static_cast<char*>(smem_ptr) + Policy::template GetSmemSizeK<Problem>()));
-        auto q_lds           = make_tensor_view<address_space_enum::lds>(
-            q_lds_ptr, Policy::template MakeQLdsBlockDescriptor<Problem>());
-        auto q_lds_window =
-            make_tile_window(q_lds, make_tuple(number<kM0>{}, number<kK0>{}), {0, 0});
-
-        // QT tile in LDS
-        QDataType* qt_lds_ptr = static_cast<QDataType*>(static_cast<void*>(
-            static_cast<char*>(smem_ptr) + Policy::template GetSmemSizeK<Problem>()));
-        auto qt_lds           = make_tensor_view<address_space_enum::lds>(
-            qt_lds_ptr, Policy::template MakeQTLdsBlockDescriptor<Problem>());
-        auto qt_lds_window =
-            make_tile_window(qt_lds, make_tuple(number<kQKHeaddim>{}, number<kK3>{}), {0, 0});
-
-        // K tile in LDS
-        auto k_lds = make_tensor_view<address_space_enum::lds>(
-            reinterpret_cast<KDataType*>(smem_ptr),
-            Policy::template MakeKLdsBlockDescriptor<Problem>());
-        auto k_lds_window =
-            make_tile_window(k_lds, make_tuple(number<kN0>{}, number<kQKHeaddim>{}), {0, 0});
-
-        // KT tile in LDS
-        auto kt_lds = make_tensor_view<address_space_enum::lds>(
-            reinterpret_cast<KDataType*>(smem_ptr),
-            Policy::template MakeKLdsBlockDescriptorAsKT<Problem>());
-        auto kt_lds_window =
-            make_tile_window(kt_lds, make_tuple(number<kQKHeaddim>{}, number<kN0>{}), {0, 0});
-
-        // OGrad tile in LDS
-        OGradDataType* do_lds_ptr = static_cast<OGradDataType*>(static_cast<void*>(
-            static_cast<char*>(smem_ptr) + Policy::template GetSmemSizeK<Problem>()));
-        auto do_lds               = make_tensor_view<address_space_enum::lds>(
-            do_lds_ptr, Policy::template MakeOGradLdsBlockDescriptor<Problem>());
-        auto do_lds_window =
-            make_tile_window(do_lds, make_tuple(number<kM0>{}, number<kK2>{}), {0, 0});
-
-        // OGradT tile in LDS
-        OGradDataType* dot_lds_ptr = static_cast<OGradDataType*>(static_cast<void*>(
-            static_cast<char*>(smem_ptr) + Policy::template GetSmemSizeK<Problem>()));
-        auto dot_lds               = make_tensor_view<address_space_enum::lds>(
-            dot_lds_ptr, Policy::template MakeOGradTLdsBlockDescriptor<Problem>());
-        auto dot_lds_window =
-            make_tile_window(dot_lds, make_tuple(number<kVHeaddim>{}, number<kK1>{}), {0, 0});
-
-        // SGrad tile in LDS
-        GemmDataType* ds_lds_ptr = static_cast<GemmDataType*>(static_cast<void*>(
-            static_cast<char*>(smem_ptr) + Policy::template GetSmemSizeK<Problem>()));
-        auto ds_lds              = make_tensor_view<address_space_enum::lds>(
-            ds_lds_ptr, Policy::template MakeSGradLdsBlockDescriptor<Problem>());
-        auto ds_lds_window =
-            make_tile_window(ds_lds, make_tuple(number<kM0>{}, number<kN0>{}), {0, 0});
-
-        // BiasT/BiasGradT tile in LDS, use the same size and layout
-        BiasDataType* biast_lds_ptr = static_cast<BiasDataType*>(static_cast<void*>(
-            static_cast<char*>(smem_ptr) + Policy::template GetSmemSizeK<Problem>()));
-        auto biast_lds              = make_tensor_view<address_space_enum::lds>(
-            biast_lds_ptr, Policy::template MakeBiasTLdsBlockDescriptor<Problem>());
-        auto biast_lds_shuffle_window =
-            make_tile_window(biast_lds, make_tuple(number<kM0>{}, number<kN0>{}), {0, 0});
-        auto dbiast_lds_shuffle_window =
-            make_tile_window(biast_lds,
-                             make_tuple(number<kM0>{}, number<kN0>{}),
-                             {0, 0},
-                             Policy::template MakeShuffledBiasTileDistribution<Problem>());
-
-        static_assert(std::is_same_v<BiasDataType, BiasGradDataType>,
-                      "BiasDataType and BiasGradDataType should be the same!");
-
-        // Block GEMM
-        constexpr auto gemm_0 = Policy::template GetQKBlockGemm<Problem>();
-        constexpr auto gemm_1 = Policy::template GetPTOGradTBlockGemm<Problem>();
-        constexpr auto gemm_2 = Policy::template GetOGradVBlockGemm<Problem>();
-        constexpr auto gemm_3 = Policy::template GetSGradTQTBlockGemm<Problem>();
-        constexpr auto gemm_4 = Policy::template GetSGradKTBlockGemm<Problem>();
-
-        auto v_dram_window = make_tile_window(
-            v_dram_block_window_tmp.get_bottom_tensor_view(),
-            v_dram_block_window_tmp.get_window_lengths(),
-            v_dram_block_window_tmp.get_window_origin(),
-            Policy::template MakeVInRegDramTileDistribution<Problem, decltype(gemm_2)>());
-
-        auto v = load_tile(v_dram_window); // persistent V register tile
-
-        using SPTBlockTileType     = decltype(gemm_0.MakeCBlockTile());
-        using SPGradTBlockTileType = decltype(gemm_2.MakeCBlockTile());
-        using QGradBlockTileType   = decltype(gemm_4.MakeCBlockTile());
-
-        // init VGrad & KGrad
-        auto dv_acc = decltype(gemm_1.MakeCBlockTile()){};
-        auto dk_acc = decltype(gemm_3.MakeCBlockTile()){};
-
-        clear_tile(dv_acc);
-        clear_tile(dk_acc);
-
-        auto k_dram_window = make_tile_window(
-            k_dram_block_window_tmp.get_bottom_tensor_view(),
-            k_dram_block_window_tmp.get_window_lengths(),
-            k_dram_block_window_tmp.get_window_origin(),
-            Policy::template MakeKDramTileDistribution<Problem>()); // K DRAM tile window for
-                                                                    // load
-
-        __builtin_amdgcn_sched_barrier(0);
-        const auto k_origin = k_dram_window.get_window_origin();
-        const auto [seqlen_q_start, seqlen_q_end] =
-            mask.GetTileRangeAlongY(k_origin.at(number<0>{}), number<kM0>{}, number<kN0>{});
-
-        const auto num_total_loop = integer_divide_ceil(seqlen_q_end - seqlen_q_start, kM0);
-
-        // check early exit if masked and no work to do.
-        if constexpr(FmhaMask::IsMasking)
-        {
-            if(num_total_loop <= 0)
-            {
-                // Note: here dk_acc&dv_acc are all cleard, return it
-                // Note: v loaded but no fence, ignore it.
-                return ck_tile::make_tuple(dk_acc, dv_acc);
-            }
-        }
-
-        auto k_block_tile = load_tile(k_dram_window);
-
-        store_tile(k_lds_window, k_block_tile); // // persistent K in LDS
-
-        auto q_dram_block_window =
-            make_tile_window(q_dram_block_window_tmp.get_bottom_tensor_view(),
-                             q_dram_block_window_tmp.get_window_lengths(),
-                             {seqlen_q_start, 0});
-
-        auto qt_dram_block_window =
-            make_tile_window(qt_dram_block_window_tmp.get_bottom_tensor_view(),
-                             qt_dram_block_window_tmp.get_window_lengths(),
-                             {0, seqlen_q_start});
-
-        auto do_dram_block_window =
-            make_tile_window(do_dram_block_window_tmp.get_bottom_tensor_view(),
-                             do_dram_block_window_tmp.get_window_lengths(),
-                             {seqlen_q_start, 0});
-
-        auto dot_dram_block_window =
-            make_tile_window(dot_dram_block_window_tmp.get_bottom_tensor_view(),
-                             dot_dram_block_window_tmp.get_window_lengths(),
-                             {0, seqlen_q_start});
-
-        auto dq_dram_block_window =
-            make_tile_window(dq_dram_block_window_tmp.get_bottom_tensor_view(),
-                             dq_dram_block_window_tmp.get_window_lengths(),
-                             {seqlen_q_start, 0});
-
-        auto lse_dram_block_window =
-            make_tile_window(lse_dram_block_window_tmp.get_bottom_tensor_view(),
-                             lse_dram_block_window_tmp.get_window_lengths(),
-                             {seqlen_q_start});
-
-        auto d_dram_block_window =
-            make_tile_window(d_dram_block_window_tmp.get_bottom_tensor_view(),
-                             d_dram_block_window_tmp.get_window_lengths(),
-                             {seqlen_q_start});
-
-        const auto bias_origin = bias_dram_block_window_tmp.get_window_origin();
-        auto bias_dram_block_window =
-            make_tile_window(bias_dram_block_window_tmp.get_bottom_tensor_view(),
-                             bias_dram_block_window_tmp.get_window_lengths(),
-                             {seqlen_q_start, bias_origin.at(number<1>{})}); // M/N
-
-        const auto dbias_origin = dbias_dram_block_window_tmp.get_window_origin();
-        auto dbias_dram_block_window =
-            make_tile_window(dbias_dram_block_window_tmp.get_bottom_tensor_view(),
-                             dbias_dram_block_window_tmp.get_window_lengths(),
-                             {seqlen_q_start, dbias_origin.at(number<1>{})}); // M/N
-
-        auto qt_dram_window =
-            make_tile_window(qt_dram_block_window.get_bottom_tensor_view(),
-                             qt_dram_block_window.get_window_lengths(),
-                             qt_dram_block_window.get_window_origin(),
-                             Policy::template MakeQTDramTileDistribution<Problem>());
-
-        auto dot_dram_window =
-            make_tile_window(dot_dram_block_window.get_bottom_tensor_view(),
-                             dot_dram_block_window.get_window_lengths(),
-                             dot_dram_block_window.get_window_origin(),
-                             Policy::template MakeOGradTDramTileDistribution<Problem>());
-
-        auto lse_dram_window = make_tile_window(
-            lse_dram_block_window.get_bottom_tensor_view(),
-            lse_dram_block_window.get_window_lengths(),
-            lse_dram_block_window.get_window_origin(),
-            Policy::template MakeLSEDDramTileDistribution<Problem, decltype(gemm_0)>());
-
-        auto d_dram_window = make_tile_window(
-            d_dram_block_window.get_bottom_tensor_view(),
-            d_dram_block_window.get_window_lengths(),
-            d_dram_block_window.get_window_origin(),
-            Policy::template MakeLSEDDramTileDistribution<Problem, decltype(gemm_0)>());
-
-        auto bias_dram_window =
-            make_tile_window(bias_dram_block_window.get_bottom_tensor_view(),
-                             bias_dram_block_window.get_window_lengths(),
-                             bias_dram_block_window.get_window_origin(),
-                             Policy::template MakeBiasTileDistribution<Problem>());
-
-        auto biast_lds_window =
-            make_tile_window(biast_lds_shuffle_window.get_bottom_tensor_view(),
-                             biast_lds_shuffle_window.get_window_lengths(),
-                             biast_lds_shuffle_window.get_window_origin(),
-                             Policy::template MakeBiasTTileDistribution<decltype(gemm_0)>());
-
-        auto randval_dram_window = dropout.MakeRandvalDramWindow<decltype(gemm_0), false>(
-            randval_dram_block_window_tmp, seqlen_q_start);
-
-        index_t i_total_loops      = 0;
-        constexpr index_t k0_loops = kQKHeaddim / kK0;
-        constexpr index_t k1_loops = kM0 / kK1;
-        constexpr index_t k2_loops = kVHeaddim / kK2;
-        constexpr index_t k3_loops = kM0 / kK3;
-        constexpr index_t k4_loops = kN0 / kK4;
-        do
-        {
-            auto q_dram_window = make_tile_window(
-                q_dram_block_window.get_bottom_tensor_view(),
-                q_dram_block_window.get_window_lengths(),
-                q_dram_block_window.get_window_origin(),
-                Policy::template MakeQDramTileDistribution<Problem>()); // Q DRAM tile window for
-                                                                        // load
-
-            auto do_dram_window = make_tile_window(
-                do_dram_block_window.get_bottom_tensor_view(),
-                do_dram_block_window.get_window_lengths(),
-                do_dram_block_window.get_window_origin(),
-                Policy::template MakeOGradDramTileDistribution<Problem>()); // OGrad DRAM tile
-                                                                            // window for load
-
-            // STAGE 1, Q@K Gemm0
-            auto st_acc = SPTBlockTileType{};
-
-            auto q_block_tile = load_tile(q_dram_window);
-            {
-                move_tile_window(q_dram_window, {0, kK0});
-
-                clear_tile(st_acc); // Initialize S^T
-
-                store_tile(q_lds_window, q_block_tile);  // LDS write 0
-                q_block_tile = load_tile(q_dram_window); // global read 1
-            }
-
-            if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS)
-            {
-                __builtin_amdgcn_sched_barrier(
-                    0); // prevent from messing up the order of global loads
-            }
-            const auto bias_tile = load_tile(bias_dram_window); // load bias tile
-            if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS)
-            {
-                __builtin_amdgcn_sched_barrier(
-                    0); // prevent from messing up the order of global loads
-            }
-
-            if constexpr(k0_loops > 2)
-            {
-                static_for<0, k0_loops - 2, 1>{}([&](auto i_k0) {
-                    block_sync_lds();
-                    gemm_0(st_acc,
-                           q_lds_window,
-                           get_slice_tile(k_lds_window,
-                                          sequence<0, i_k0 * kK0>{},
-                                          sequence<kN0, (i_k0 + 1) * kK0>{}));
-                    block_sync_lds();
-                    move_tile_window(q_dram_window, {0, kK0});
-
-                    store_tile(q_lds_window,
-                               q_block_tile);                // LDS write i + 1
-                    q_block_tile = load_tile(q_dram_window); // global read i + 2
-                });
-            }
-
-            const auto dot_prefetch = load_tile(dot_dram_window); // prefetch load OGrad^T tile
-            {                                                     // tail
-                block_sync_lds();
-                gemm_0(st_acc,
-                       q_lds_window,
-                       get_slice_tile(k_lds_window,
-                                      sequence<0, (k0_loops - 2) * kK0>{},
-                                      sequence<kN0, (k0_loops - 1) * kK0>{}));
-                block_sync_lds();
-
-                store_tile(q_lds_window, q_block_tile);
-                block_sync_lds();
-
-                gemm_0(st_acc,
-                       q_lds_window,
-                       get_slice_tile(k_lds_window,
-                                      sequence<0, (k0_loops - 1) * kK0>{},
-                                      sequence<kN0, k0_loops * kK0>{}));
-            }
-
-            // STAGE 2, Scale, Add bias, Mask, Softmax, Dropout
-            if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS)
-            {
-                block_sync_lds();
-                auto bias_shuffle_tmp = make_static_distributed_tensor<BiasDataType>(
-                    Policy::template MakeShuffledBiasTileDistribution<Problem>());
-                shuffle_tile(bias_shuffle_tmp, bias_tile);
-                store_tile(biast_lds_shuffle_window, bias_shuffle_tmp);
-                block_sync_lds();
-                auto biast_tile = load_tile(biast_lds_window);
-                tile_elementwise_inout(
-                    [&](auto& x, const auto& y) {
-#if !CK_TILE_FMHA_FWD_FAST_EXP2
-                        x = raw_scale * x + type_convert<AccDataType>(y);
-#else
-                        x = scale * x + log2e_v<AccDataType> * type_convert<AccDataType>(y);
-#endif
-                    },
-                    st_acc,
-                    biast_tile);
-                move_tile_window(bias_dram_window, {kM0, 0});
-            }
-            else if constexpr(BiasEnum == BlockAttentionBiasEnum::ALIBI)
-            {
-                const auto q_origin     = q_dram_block_window.get_window_origin();
-                constexpr auto st_spans = decltype(st_acc)::get_distributed_spans();
-                sweep_tile_span(st_spans[number<0>{}], [&](auto idx0) {
-                    sweep_tile_span(st_spans[number<1>{}], [&](auto idx1) {
-                        const auto tile_idx = get_x_indices_from_distributed_indices(
-                            st_acc.get_tile_distribution(), make_tuple(idx0, idx1));
-
-                        const auto row = q_origin.at(number<0>{}) + tile_idx.at(number<0>{});
-                        const auto col = k_origin.at(number<0>{}) + tile_idx.at(number<1>{});
-                        constexpr auto i_j_idx = make_tuple(idx0, idx1);
-
-#if !CK_TILE_FMHA_FWD_FAST_EXP2
-                        st_acc(i_j_idx) *= raw_scale;
-#else
-                        st_acc(i_j_idx) *= scale;
-#endif
-                        position_encoding.update(st_acc(i_j_idx), row, col);
-                    });
-                });
-            }
-            else
-            {
-#if !CK_TILE_FMHA_FWD_FAST_EXP2
-                tile_elementwise_inout([&raw_scale](auto& x) { x = x * raw_scale; }, st_acc);
-#endif
-            }
-
-            if constexpr(kPadSeqLenK || FmhaMask::IsMasking)
-            {
-                const auto q_origin      = q_dram_block_window.get_window_origin();
-                bool need_perpixel_check = mask.IsEdgeTile(q_origin.at(number<0>{}),
-                                                           k_origin.at(number<0>{}),
-                                                           number<kM0>{},
-                                                           number<kN0>{});
-                if(need_perpixel_check)
-                {
-                    set_tile_if(st_acc, -numeric<AccDataType>::infinity(), [&](auto tile_idx) {
-                        const auto row = q_origin.at(number<0>{}) + tile_idx.at(number<0>{});
-                        const auto col = k_origin.at(number<0>{}) + tile_idx.at(number<1>{});
-                        return mask.IsOutOfBound(row, col);
-                    });
-                }
-            }
-
-            const auto lse = load_tile(lse_dram_window);
-
-            static const auto get_validated_lse = [](LSEDataType raw_lse) {
-                if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS ||
-                             FmhaMask::IsMasking)
-                {
-                    return raw_lse == -numeric<LSEDataType>::infinity()
-                               ? type_convert<LSEDataType>(0.f)
-                               : raw_lse;
-                }
-                else
-                {
-                    return raw_lse;
-                }
-            };
-
-            auto pt                 = SPTBlockTileType{};
-            constexpr auto pt_spans = decltype(pt)::get_distributed_spans();
-            sweep_tile_span(pt_spans[number<0>{}], [&](auto idx0) {
-                constexpr auto i_idx = make_tuple(idx0);
-#if CK_TILE_FMHA_FWD_FAST_EXP2
-                auto row_lse = log2e_v<LSEDataType> * get_validated_lse(lse[i_idx]);
-#endif
-                sweep_tile_span(pt_spans[number<1>{}], [&](auto idx1) {
-                    constexpr auto i_j_idx = make_tuple(idx0, idx1);
-#if CK_TILE_FMHA_FWD_FAST_EXP2
-                    if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS ||
-                                 BiasEnum == BlockAttentionBiasEnum::ALIBI)
-                    {
-                        pt(i_j_idx) = exp2(st_acc[i_j_idx] - row_lse);
-                    }
-                    else
-                    {
-                        pt(i_j_idx) = exp2(scale * st_acc[i_j_idx] - row_lse);
-                    }
-#else
-                    pt(i_j_idx) = exp(st_acc[i_j_idx] - get_validated_lse(lse[i_idx]));
-#endif
-                });
-            });
-
-            auto dot_shuffle_tmp = make_static_distributed_tensor<OGradDataType>(
-                Policy::template MakeShuffledOGradTRegBlockDescriptor<Problem>());
-            block_sync_lds();
-            {
-                shuffle_tile(dot_shuffle_tmp, dot_prefetch);
-                store_tile(dot_lds_window,
-                           dot_shuffle_tmp); // store the prefetch
-            }
-            move_tile_window(dot_dram_window, {0, kK1});
-
-            if constexpr(kHasDropout)
-            {
-                dropout.Run<decltype(gemm_0), RandValOutputDataType>(
-                    seqlen_q_start + i_total_loops * kM0, pt, randval_dram_window);
-            }
-
-            // STAGE 3, P^T@OGrad^T Gemm1
-            const auto pt_gemm = [&]() {
-                if constexpr(kHasDropout)
-                {
-                    return tile_elementwise_in(
-                        [](const auto& x) { return type_convert<GemmDataType>(x > 0.f ? x : 0.f); },
-                        pt);
-                }
-                else
-                {
-                    return cast_tile<GemmDataType>(pt);
-                }
-            }();
-
-            if constexpr(k1_loops > 1)
-            {
-                static_for<0, k1_loops - 1, 1>{}([&](auto i_k1) {
-                    const auto dot = load_tile(dot_dram_window); // load next OGrad^T
-                    block_sync_lds();
-                    gemm_1(dv_acc,
-                           get_slice_tile(pt_gemm,
-                                          sequence<i_k1 * kK1, 0>{},
-                                          sequence<(i_k1 + 1) * kK1, kN0>{}),
-                           dot_lds_window);
-                    block_sync_lds();
-                    shuffle_tile(dot_shuffle_tmp, dot);
-                    store_tile(dot_lds_window,
-                               dot_shuffle_tmp); // store the prefetch
-
-                    move_tile_window(dot_dram_window, {0, kK1});
-                });
-            }
-            auto do_block_tile = load_tile(do_dram_window); // prefetch load OGrad tile
-            // tail
-            {
-                block_sync_lds();
-                gemm_1(dv_acc,
-                       get_slice_tile(
-                           pt_gemm, sequence<(k1_loops - 1) * kK1, 0>{}, sequence<kM0, kN0>{}),
-                       dot_lds_window);
-                block_sync_lds();
-            }
-
-            // STAGE 4, OGrad@V Gemm2
-            auto dpt_acc = SPGradTBlockTileType{};
-
-            {
-                move_tile_window(do_dram_window, {0, kK2});
-
-                clear_tile(dpt_acc); // Initialize PGrad^T
-
-                store_tile(do_lds_window, do_block_tile);  // LDS write 0
-                do_block_tile = load_tile(do_dram_window); // global read 1
-            }
-
-            if constexpr(k2_loops > 2)
-            {
-                static_for<0, k2_loops - 2, 1>{}([&](auto i_k2) {
-                    block_sync_lds();
-                    gemm_2(dpt_acc,
-                           do_lds_window,
-                           get_slice_tile(
-                               v, sequence<0, i_k2 * kK2>{}, sequence<kN0, (i_k2 + 1) * kK2>{}));
-                    block_sync_lds();
-                    move_tile_window(do_dram_window, {0, kK2});
-
-                    store_tile(do_lds_window,
-                               do_block_tile);                 // LDS write i + 1
-                    do_block_tile = load_tile(do_dram_window); // global read i + 2
-                });
-            }
-
-            const auto qt_prefetch = load_tile(qt_dram_window); // prefetch load Q^T tile
-            {                                                   // tail
-                block_sync_lds();
-                gemm_2(dpt_acc,
-                       do_lds_window,
-                       get_slice_tile(v,
-                                      sequence<0, (k2_loops - 2) * kK2>{},
-                                      sequence<kN0, (k2_loops - 1) * kK2>{}));
-                block_sync_lds();
-
-                store_tile(do_lds_window, do_block_tile);
-                block_sync_lds();
-
-                gemm_2(dpt_acc,
-                       do_lds_window,
-                       get_slice_tile(v,
-                                      sequence<0, (k2_loops - 1) * kK2>{},
-                                      sequence<kN0, k2_loops * kK2>{}));
-            }
-
-            // STAGE 5, P^T(PGrad^T - D)
-            const auto d = load_tile(d_dram_window);
-
-            auto dst                 = SPGradTBlockTileType{};
-            constexpr auto dst_spans = decltype(dst)::get_distributed_spans();
-            sweep_tile_span(dst_spans[number<0>{}], [&](auto idx0) {
-                constexpr auto i_idx = make_tuple(idx0);
-                sweep_tile_span(dst_spans[number<1>{}], [&](auto idx1) {
-                    constexpr auto i_j_idx = make_tuple(idx0, idx1);
-                    bool undrop_flag       = pt[i_j_idx] >= 0;
-                    dst(i_j_idx) =
-                        pt[i_j_idx] *
-                        (!kHasDropout || undrop_flag ? (dpt_acc[i_j_idx] - d[i_idx]) : d[i_idx]);
-                });
-            });
-
-            if constexpr(kHasBiasGrad)
-            {
-                const auto dbiast = [&]() {
-                    if constexpr(kHasDropout)
-                    {
-                        return tile_elementwise_in(
-                            [&rp_undrop](const auto& x) {
-                                return type_convert<BiasGradDataType>(x * rp_undrop);
-                            },
-                            dst);
-                    }
-                    else
-                    {
-                        return cast_tile<BiasGradDataType>(dst);
-                    }
-                }();
-                store_tile(biast_lds_shuffle_window, dbiast);
-                block_sync_lds();
-                auto dbiast_tile        = load_tile(dbiast_lds_shuffle_window);
-                auto dbiast_shuffle_tmp = make_static_distributed_tensor<BiasGradDataType>(
-                    Policy::template MakeBiasTileDistribution<Problem>());
-                shuffle_tile(dbiast_shuffle_tmp, dbiast_tile);
-                store_tile(dbias_dram_block_window, dbiast_shuffle_tmp);
-                move_tile_window(dbias_dram_block_window, {kM0, 0});
-            }
-
-            // STAGE 6, SGrad^T@Q^T Gemm3
-            auto qt_shuffle_tmp = make_static_distributed_tensor<QDataType>(
-                Policy::template MakeShuffledQTRegBlockDescriptor<Problem>());
-            block_sync_lds();
-            {
-                shuffle_tile(qt_shuffle_tmp, qt_prefetch);
-                store_tile(qt_lds_window,
-                           qt_shuffle_tmp); // store the prefetch
-            }
-            move_tile_window(qt_dram_window, {0, kK3});
-
-            const auto dst_gemm = cast_tile<GemmDataType>(dst);
-
-            if constexpr(k3_loops > 1)
-            {
-                static_for<0, k3_loops - 1, 1>{}([&](auto i_k3) {
-                    const auto qt = load_tile(qt_dram_window); // load next Q^T
-                    block_sync_lds();
-                    gemm_3(dk_acc,
-                           get_slice_tile(dst_gemm,
-                                          sequence<i_k3 * kK3, 0>{},
-                                          sequence<(i_k3 + 1) * kK3, kN0>{}),
-                           qt_lds_window);
-                    block_sync_lds();
-                    shuffle_tile(qt_shuffle_tmp, qt);
-                    store_tile(qt_lds_window,
-                               qt_shuffle_tmp); // store the prefetch
-
-                    move_tile_window(qt_dram_window, {0, kK3});
-                });
-            }
-            // tail
-            {
-                block_sync_lds();
-                gemm_3(dk_acc,
-                       get_slice_tile(
-                           dst_gemm, sequence<(k3_loops - 1) * kK3, 0>{}, sequence<kM0, kN0>{}),
-                       qt_lds_window);
-                block_sync_lds();
-            }
-
-            // STAGE 7, SGrad@K^T Gemm4
-            store_tile(ds_lds_window, dst_gemm);
-
-            auto dq_acc = QGradBlockTileType{};
-            clear_tile(dq_acc); // Initialize QGrad
-
-            block_sync_lds();
-
-            static_for<0, k4_loops, 1>{}([&](auto i_k4) {
-                gemm_4(dq_acc,
-                       get_slice_tile(ds_lds_window,
-                                      sequence<0, i_k4 * kK4>{},
-                                      sequence<kM0, (i_k4 + 1) * kK4>{}),
-                       get_slice_tile(kt_lds_window,
-                                      sequence<0, i_k4 * kK4>{},
-                                      sequence<kQKHeaddim, (i_k4 + 1) * kK4>{}));
-            });
-
-            // QGrad Scale
-            if constexpr(kHasDropout)
-            {
-                tile_elementwise_inout([&scale_rp_undrop](auto& x) { x = x * scale_rp_undrop; },
-                                       dq_acc);
-            }
-            else
-            {
-                tile_elementwise_inout([&raw_scale](auto& x) { x = x * raw_scale; }, dq_acc);
-            }
-            const auto dq = cast_tile<QGradDataType>(dq_acc);
-            update_tile(dq_dram_block_window, dq);
-
-            // move tile windows
-            move_tile_window(q_dram_block_window, {kM0, 0});
-            move_tile_window(dq_dram_block_window, {kM0, 0});
-            move_tile_window(do_dram_block_window, {kM0, 0});
-            move_tile_window(lse_dram_window, {kM0});
-            move_tile_window(d_dram_window, {kM0});
-        } while(++i_total_loops < num_total_loop);
-
-        // KGrad Scale
-        if constexpr(kHasDropout)
-        {
-            tile_elementwise_inout([&scale_rp_undrop](auto& x) { x = x * scale_rp_undrop; },
-                                   dk_acc);
-        }
-        else
-        {
-            tile_elementwise_inout([&raw_scale](auto& x) { x = x * raw_scale; }, dk_acc);
-        }
-        // VGrad Scale
-        if constexpr(kHasDropout)
-        {
-            tile_elementwise_inout([&rp_undrop](auto& x) { x = x * rp_undrop; }, dv_acc);
-        }
-
-        return ck_tile::make_tuple(dk_acc, dv_acc);
-    }
-};
-
-} // namespace ck_tile

include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_ks_vr_default_policy.hpp

@@ -1,20 +0,0 @@
-// SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
-
-#pragma once
-
-#include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_default_policy.hpp"
-
-namespace ck_tile {
-
-// This pipeline is v located in regs, k located in lds.
-using BlockFmhaBwdDQDKDVPipelineKSVRDefaultPolicy =
-    BlockFmhaBwdPipelineDefaultPolicy</* QLoadOnce_      = */ false,
-                                      /* QTLoadOnce_     = */ false,
-                                      /* KLoadOnce_      = */ true,
-                                      /* KTLoadOnce_     = */ false,
-                                      /* VLoadOnce_      = */ true,
-                                      /* OGradLoadOnce_  = */ false,
-                                      /* OGradTLoadOnce_ = */ false>;
-
-} // namespace ck_tile

include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_qs_ks_vr_dos.hpp

@@ -1,692 +0,0 @@
-// SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
-
-#pragma once
-
-#include "ck_tile/core.hpp"
-#include "ck_tile/ops/fmha/block/block_attention_bias_enum.hpp"
-#include "ck_tile/ops/fmha/block/block_dropout.hpp"
-#include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_qs_ks_vr_dos_default_policy.hpp"
-#include "ck_tile/ops/reduce/block/block_reduce.hpp"
-
-namespace ck_tile {
-
-template <typename Problem, typename Policy = BlockFmhaBwdDQDKDVPipelineQSKSVROGradSDefaultPolicy>
-struct BlockFmhaBwdDQDKDVPipelineQSKSVROGradS
-{
-    using QDataType             = remove_cvref_t<typename Problem::QDataType>;
-    using KDataType             = remove_cvref_t<typename Problem::KDataType>;
-    using VDataType             = remove_cvref_t<typename Problem::VDataType>;
-    using GemmDataType          = remove_cvref_t<typename Problem::GemmDataType>;
-    using BiasDataType          = remove_cvref_t<typename Problem::BiasDataType>;
-    using LSEDataType           = remove_cvref_t<typename Problem::LSEDataType>;
-    using AccDataType           = remove_cvref_t<typename Problem::AccDataType>;
-    using DDataType             = remove_cvref_t<typename Problem::DDataType>;
-    using RandValOutputDataType = remove_cvref_t<typename Problem::RandValOutputDataType>;
-    using ODataType             = remove_cvref_t<typename Problem::ODataType>;
-    using OGradDataType         = remove_cvref_t<typename Problem::OGradDataType>;
-    using QGradDataType         = remove_cvref_t<typename Problem::QGradDataType>;
-    using KGradDataType         = remove_cvref_t<typename Problem::KGradDataType>;
-    using VGradDataType         = remove_cvref_t<typename Problem::VGradDataType>;
-    using BiasGradDataType      = remove_cvref_t<typename Problem::BiasGradDataType>;
-    using FmhaMask              = remove_cvref_t<typename Problem::FmhaMask>;
-
-    using BlockFmhaShape = remove_cvref_t<typename Problem::BlockFmhaShape>;
-
-    static constexpr index_t kBlockPerCu = Problem::kBlockPerCu;
-    static constexpr index_t kBlockSize  = Problem::kBlockSize;
-
-    static constexpr index_t kM0        = BlockFmhaShape::kM0;
-    static constexpr index_t kN0        = BlockFmhaShape::kN0;
-    static constexpr index_t kK0        = BlockFmhaShape::kK0;
-    static constexpr index_t kK1        = BlockFmhaShape::kK1;
-    static constexpr index_t kK2        = BlockFmhaShape::kK2;
-    static constexpr index_t kK3        = BlockFmhaShape::kK3;
-    static constexpr index_t kK4        = BlockFmhaShape::kK4;
-    static constexpr index_t kQKHeaddim = BlockFmhaShape::kQKHeaddim;
-    static constexpr index_t kVHeaddim  = BlockFmhaShape::kVHeaddim;
-
-    static constexpr bool kQLoadOnce      = true;
-    static constexpr bool kQTLoadOnce     = false;
-    static constexpr bool kKLoadOnce      = true;
-    static constexpr bool kKTLoadOnce     = false;
-    static constexpr bool kVLoadOnce      = true;
-    static constexpr bool kOGradLoadOnce  = true;
-    static constexpr bool kOGradTLoadOnce = false;
-
-    static constexpr bool kIsGroupMode = Problem::kIsGroupMode;
-    static constexpr bool kPadSeqLenQ  = Problem::kPadSeqLenQ;
-    static constexpr bool kPadSeqLenK  = Problem::kPadSeqLenK;
-    static constexpr bool kPadHeadDimQ = Problem::kPadHeadDimQ;
-    static constexpr bool kPadHeadDimV = Problem::kPadHeadDimV;
-    static constexpr auto BiasEnum     = Problem::BiasEnum;
-    static constexpr bool kHasBiasGrad = Problem::kHasBiasGrad;
-    static constexpr bool kHasDropout  = Problem::kHasDropout;
-
-    // last dimension vector length used to create tensor view(and decide buffer_load vector length)
-    // ... together with tensor distribution. tensor dist should able to overwrite this
-    static constexpr index_t kAlignmentQ =
-        kPadHeadDimQ ? 1 : Policy::template GetAlignmentQ<Problem>();
-    static constexpr index_t kAlignmentK =
-        kPadHeadDimQ ? 1 : Policy::template GetAlignmentK<Problem>();
-    static constexpr index_t kAlignmentV =
-        kPadHeadDimV ? 1 : Policy::template GetAlignmentV<Problem>();
-    static constexpr index_t kAlignmentO =
-        kPadHeadDimV ? 1 : Policy::template GetAlignmentO<Problem>();
-    static constexpr index_t kAlignmentOGrad =
-        kPadHeadDimV ? 1 : Policy::template GetAlignmentOGrad<Problem>();
-    static constexpr index_t kAlignmentQGrad =
-        kPadHeadDimQ ? 2 : Policy::template GetAlignmentQGrad<Problem>();
-    static constexpr index_t kAlignmentKGrad =
-        kPadHeadDimQ ? 1 : Policy::template GetAlignmentKGrad<Problem>();
-    static constexpr index_t kAlignmentVGrad =
-        kPadHeadDimV ? 1 : Policy::template GetAlignmentVGrad<Problem>();
-    static constexpr index_t kAlignmentBias =
-        kPadSeqLenK ? 1 : Policy::template GetTransposedAlignmentBias<Problem>();
-
-    static constexpr const char* name = "qs_ks_vr_dos";
-
-    CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetSmemSize()
-    {
-        return Policy::template GetSmemSize<Problem>();
-    }
-
-    template <typename QDramBlockWindowTmp,
-              typename QTDramBlockWindowTmp,
-              typename KDramBlockWindowTmp,
-              typename KTDramBlockWindowTmp,
-              typename VDramBlockWindowTmp,
-              typename BiasDramBlockWindowTmp,
-              typename RandValDramBlockWindowTmp,
-              typename OGradDramBlockWindowTmp,
-              typename OGradTDramBlockWindowTmp,
-              typename LSEDramBlockWindowTmp,
-              typename DDramBlockWindowTmp,
-              typename QGradDramBlockWindowTmp,
-              typename BiasGradDramBlockWindowTmp,
-              typename PositionEncoding>
-    CK_TILE_HOST_DEVICE auto
-    operator()(const QDramBlockWindowTmp& q_dram_block_window_tmp,
-               const QTDramBlockWindowTmp& /*qt_dram_block_window_tmp*/,
-               const KDramBlockWindowTmp& k_dram_block_window_tmp,
-               const KTDramBlockWindowTmp& /*kt_dram_block_window_tmp*/,
-               const VDramBlockWindowTmp& v_dram_block_window_tmp,
-               const BiasDramBlockWindowTmp& bias_dram_block_window_tmp,
-               const RandValDramBlockWindowTmp& randval_dram_block_window_tmp,
-               const OGradDramBlockWindowTmp& do_dram_block_window_tmp,
-               const OGradTDramBlockWindowTmp& /*dot_dram_block_window_tmp*/,
-               const LSEDramBlockWindowTmp& lse_dram_block_window_tmp,
-               const DDramBlockWindowTmp& d_dram_block_window_tmp,
-               const QGradDramBlockWindowTmp& dq_dram_block_window_tmp,
-               const BiasGradDramBlockWindowTmp& dbias_dram_block_window_tmp,
-               FmhaMask mask,
-               PositionEncoding position_encoding,
-               float raw_scale,
-#if CK_TILE_FMHA_FWD_FAST_EXP2
-               float scale,
-#endif
-               float rp_undrop,
-               float scale_rp_undrop,
-               void* smem_ptr,
-               BlockDropout& dropout) const
-    {
-        static_assert(
-            std::is_same_v<QDataType, remove_cvref_t<typename QDramBlockWindowTmp::DataType>> &&
-                std::is_same_v<KDataType, remove_cvref_t<typename KDramBlockWindowTmp::DataType>> &&
-                std::is_same_v<VDataType, remove_cvref_t<typename VDramBlockWindowTmp::DataType>> &&
-                std::is_same_v<OGradDataType,
-                               remove_cvref_t<typename OGradDramBlockWindowTmp::DataType>> &&
-                std::is_same_v<LSEDataType,
-                               remove_cvref_t<typename LSEDramBlockWindowTmp::DataType>> &&
-                std::is_same_v<DDataType, remove_cvref_t<typename DDramBlockWindowTmp::DataType>> &&
-                std::is_same_v<QGradDataType,
-                               remove_cvref_t<typename QGradDramBlockWindowTmp::DataType>>,
-            "wrong!");
-
-        static_assert(kM0 == QDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
-                          kN0 == KDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
-                          kN0 == VDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
-                          kM0 == BiasDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
-                          kN0 == BiasDramBlockWindowTmp{}.get_window_lengths()[number<1>{}] &&
-                          kM0 == OGradDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
-                          kM0 == LSEDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
-                          kM0 == DDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
-                          kM0 == QGradDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
-                          kM0 == BiasGradDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
-                          kN0 == BiasGradDramBlockWindowTmp{}.get_window_lengths()[number<1>{}],
-                      "wrong!");
-
-        // Q tile in LDS
-        QDataType* q_lds_ptr = static_cast<QDataType*>(static_cast<void*>(
-            static_cast<char*>(smem_ptr) + Policy::template GetSmemSizeK<Problem>()));
-        auto q_lds           = make_tensor_view<address_space_enum::lds>(
-            q_lds_ptr, Policy::template MakeQLdsBlockDescriptor<Problem>());
-        auto q_lds_window =
-            make_tile_window(q_lds, make_tuple(number<kM0>{}, number<kQKHeaddim>{}), {0, 0});
-
-        // QT tile in LDS
-        auto qt_lds = make_tensor_view<address_space_enum::lds>(
-            q_lds_ptr, Policy::template MakeQLdsBlockDescriptorAsQT<Problem>());
-        auto qt_lds_window =
-            make_tile_window(qt_lds, make_tuple(number<kQKHeaddim>{}, number<kM0>{}), {0, 0});
-
-        // K tile in LDS
-        auto k_lds = make_tensor_view<address_space_enum::lds>(
-            reinterpret_cast<KDataType*>(smem_ptr),
-            Policy::template MakeKLdsBlockDescriptor<Problem>());
-        auto k_lds_window =
-            make_tile_window(k_lds, make_tuple(number<kN0>{}, number<kQKHeaddim>{}), {0, 0});
-
-        // KT tile in LDS
-        auto kt_lds = make_tensor_view<address_space_enum::lds>(
-            reinterpret_cast<KDataType*>(smem_ptr),
-            Policy::template MakeKLdsBlockDescriptorAsKT<Problem>());
-        auto kt_lds_window =
-            make_tile_window(kt_lds, make_tuple(number<kQKHeaddim>{}, number<kN0>{}), {0, 0});
-
-        // OGrad tile in LDS
-        OGradDataType* do_lds_ptr = static_cast<OGradDataType*>(static_cast<void*>(
-            static_cast<char*>(smem_ptr) + Policy::template GetSmemSizeK<Problem>() +
-            Policy::template GetSmemSizeQ<Problem>()));
-        auto do_lds               = make_tensor_view<address_space_enum::lds>(
-            do_lds_ptr, Policy::template MakeOGradLdsBlockDescriptor<Problem>());
-        auto do_lds_window =
-            make_tile_window(do_lds, make_tuple(number<kM0>{}, number<kVHeaddim>{}), {0, 0});
-
-        // OGradT tile in LDS
-        auto dot_lds = make_tensor_view<address_space_enum::lds>(
-            do_lds_ptr, Policy::template MakeOGradLdsBlockDescriptorAsOGradT<Problem>());
-        auto dot_lds_window =
-            make_tile_window(dot_lds, make_tuple(number<kVHeaddim>{}, number<kM0>{}), {0, 0});
-
-        // SGrad tile in LDS
-        GemmDataType* ds_lds_ptr = static_cast<GemmDataType*>(static_cast<void*>(
-            static_cast<char*>(smem_ptr) + Policy::template GetSmemSizeK<Problem>() +
-            Policy::template GetSmemSizeQ<Problem>() +
-            Policy::template GetSmemSizeOGrad<Problem>()));
-        auto ds_lds              = make_tensor_view<address_space_enum::lds>(
-            ds_lds_ptr, Policy::template MakeSGradLdsBlockDescriptor<Problem>());
-        auto ds_lds_window =
-            make_tile_window(ds_lds, make_tuple(number<kM0>{}, number<kN0>{}), {0, 0});
-
-        // BiasT/BiasGradT tile in LDS, use the same size and layout
-        BiasDataType* biast_lds_ptr = static_cast<BiasDataType*>(static_cast<void*>(
-            static_cast<char*>(smem_ptr) + Policy::template GetSmemSizeK<Problem>() +
-            Policy::template GetSmemSizeQ<Problem>() +
-            Policy::template GetSmemSizeOGrad<Problem>()));
-        auto biast_lds              = make_tensor_view<address_space_enum::lds>(
-            biast_lds_ptr, Policy::template MakeBiasTLdsBlockDescriptor<Problem>());
-        auto biast_lds_shuffle_window =
-            make_tile_window(biast_lds, make_tuple(number<kM0>{}, number<kN0>{}), {0, 0});
-        auto dbiast_lds_shuffle_window =
-            make_tile_window(biast_lds,
-                             make_tuple(number<kM0>{}, number<kN0>{}),
-                             {0, 0},
-                             Policy::template MakeShuffledBiasTileDistribution<Problem>());
-
-        static_assert(std::is_same_v<BiasDataType, BiasGradDataType>,
-                      "BiasDataType and BiasGradDataType should be the same!");
-
-        // Block GEMM
-        constexpr auto gemm_0 = Policy::template GetQKBlockGemm<Problem>();
-        constexpr auto gemm_1 = Policy::template GetPTOGradTBlockGemm<Problem>();
-        constexpr auto gemm_2 = Policy::template GetOGradVBlockGemm<Problem>();
-        constexpr auto gemm_3 = Policy::template GetSGradTQTBlockGemm<Problem>();
-        constexpr auto gemm_4 = Policy::template GetSGradKTBlockGemm<Problem>();
-
-        auto v_dram_window = make_tile_window(
-            v_dram_block_window_tmp.get_bottom_tensor_view(),
-            v_dram_block_window_tmp.get_window_lengths(),
-            v_dram_block_window_tmp.get_window_origin(),
-            Policy::template MakeVInRegDramTileDistribution<Problem, decltype(gemm_2)>());
-
-        auto v = load_tile(v_dram_window); // persistent V register tile
-
-        using SPTBlockTileType     = decltype(gemm_0.MakeCBlockTile());
-        using SPGradTBlockTileType = decltype(gemm_2.MakeCBlockTile());
-        using QGradBlockTileType   = decltype(gemm_4.MakeCBlockTile());
-
-        // init VGrad & KGrad
-        auto dv_acc = decltype(gemm_1.MakeCBlockTile()){};
-        auto dk_acc = decltype(gemm_3.MakeCBlockTile()){};
-
-        clear_tile(dv_acc);
-        clear_tile(dk_acc);
-
-        auto k_dram_window = make_tile_window(
-            k_dram_block_window_tmp.get_bottom_tensor_view(),
-            k_dram_block_window_tmp.get_window_lengths(),
-            k_dram_block_window_tmp.get_window_origin(),
-            Policy::template MakeKDramTileDistribution<Problem>()); // K DRAM tile window for
-                                                                    // load
-
-        __builtin_amdgcn_sched_barrier(0);
-        const auto k_origin = k_dram_window.get_window_origin();
-        const auto [seqlen_q_start, seqlen_q_end] =
-            mask.GetTileRangeAlongY(k_origin.at(number<0>{}), number<kM0>{}, number<kN0>{});
-
-        const auto num_total_loop = integer_divide_ceil(seqlen_q_end - seqlen_q_start, kM0);
-
-        // check early exit if masked and no work to do.
-        if constexpr(FmhaMask::IsMasking)
-        {
-            if(num_total_loop <= 0)
-            {
-                // Note: here dk_acc&dv_acc are all cleard, return it
-                // Note: v loaded but no fence, ignore it.
-                return ck_tile::make_tuple(dk_acc, dv_acc);
-            }
-        }
-
-        auto k_block_tile = load_tile(k_dram_window);
-
-        store_tile(k_lds_window, k_block_tile); // // persistent K in LDS
-
-        auto q_dram_block_window =
-            make_tile_window(q_dram_block_window_tmp.get_bottom_tensor_view(),
-                             q_dram_block_window_tmp.get_window_lengths(),
-                             {seqlen_q_start, 0});
-
-        auto do_dram_block_window =
-            make_tile_window(do_dram_block_window_tmp.get_bottom_tensor_view(),
-                             do_dram_block_window_tmp.get_window_lengths(),
-                             {seqlen_q_start, 0});
-
-        auto dq_dram_block_window =
-            make_tile_window(dq_dram_block_window_tmp.get_bottom_tensor_view(),
-                             dq_dram_block_window_tmp.get_window_lengths(),
-                             {seqlen_q_start, 0});
-
-        auto lse_dram_block_window =
-            make_tile_window(lse_dram_block_window_tmp.get_bottom_tensor_view(),
-                             lse_dram_block_window_tmp.get_window_lengths(),
-                             {seqlen_q_start});
-
-        auto d_dram_block_window =
-            make_tile_window(d_dram_block_window_tmp.get_bottom_tensor_view(),
-                             d_dram_block_window_tmp.get_window_lengths(),
-                             {seqlen_q_start});
-
-        const auto bias_origin = bias_dram_block_window_tmp.get_window_origin();
-        auto bias_dram_block_window =
-            make_tile_window(bias_dram_block_window_tmp.get_bottom_tensor_view(),
-                             bias_dram_block_window_tmp.get_window_lengths(),
-                             {seqlen_q_start, bias_origin.at(number<1>{})}); // M/N
-
-        const auto dbias_origin = dbias_dram_block_window_tmp.get_window_origin();
-        auto dbias_dram_block_window =
-            make_tile_window(dbias_dram_block_window_tmp.get_bottom_tensor_view(),
-                             dbias_dram_block_window_tmp.get_window_lengths(),
-                             {seqlen_q_start, dbias_origin.at(number<1>{})}); // M/N
-
-        auto lse_dram_window = make_tile_window(
-            lse_dram_block_window.get_bottom_tensor_view(),
-            lse_dram_block_window.get_window_lengths(),
-            lse_dram_block_window.get_window_origin(),
-            Policy::template MakeLSEDDramTileDistribution<Problem, decltype(gemm_0)>());
-
-        auto d_dram_window = make_tile_window(
-            d_dram_block_window.get_bottom_tensor_view(),
-            d_dram_block_window.get_window_lengths(),
-            d_dram_block_window.get_window_origin(),
-            Policy::template MakeLSEDDramTileDistribution<Problem, decltype(gemm_0)>());
-
-        auto bias_dram_window =
-            make_tile_window(bias_dram_block_window.get_bottom_tensor_view(),
-                             bias_dram_block_window.get_window_lengths(),
-                             bias_dram_block_window.get_window_origin(),
-                             Policy::template MakeBiasTileDistribution<Problem>());
-
-        auto biast_lds_window =
-            make_tile_window(biast_lds_shuffle_window.get_bottom_tensor_view(),
-                             biast_lds_shuffle_window.get_window_lengths(),
-                             biast_lds_shuffle_window.get_window_origin(),
-                             Policy::template MakeBiasTTileDistribution<decltype(gemm_0)>());
-
-        auto randval_dram_window = dropout.MakeRandvalDramWindow<decltype(gemm_0), false>(
-            randval_dram_block_window_tmp, seqlen_q_start);
-
-        index_t i_total_loops      = 0;
-        constexpr index_t k0_loops = kQKHeaddim / kK0;
-        constexpr index_t k1_loops = kM0 / kK1;
-        constexpr index_t k2_loops = kVHeaddim / kK2;
-        constexpr index_t k3_loops = kM0 / kK3;
-        constexpr index_t k4_loops = kN0 / kK4;
-        do
-        {
-            auto q_dram_window = make_tile_window(
-                q_dram_block_window.get_bottom_tensor_view(),
-                q_dram_block_window.get_window_lengths(),
-                q_dram_block_window.get_window_origin(),
-                Policy::template MakeQDramTileDistribution<Problem>()); // Q DRAM tile window for
-                                                                        // load
-
-            auto do_dram_window = make_tile_window(
-                do_dram_block_window.get_bottom_tensor_view(),
-                do_dram_block_window.get_window_lengths(),
-                do_dram_block_window.get_window_origin(),
-                Policy::template MakeOGradDramTileDistribution<Problem>()); // OGrad DRAM tile
-                                                                            // window for load
-
-            // STAGE 1, Q@K Gemm0
-            auto st_acc = SPTBlockTileType{};
-
-            auto q_block_tile = load_tile(q_dram_window);
-            clear_tile(st_acc);                     // Initialize S^T
-            store_tile(q_lds_window, q_block_tile); // LDS write
-
-            if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS)
-            {
-                __builtin_amdgcn_sched_barrier(
-                    0); // prevent from messing up the order of global loads
-            }
-            const auto bias_tile = load_tile(bias_dram_window); // load bias tile
-            if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS)
-            {
-                __builtin_amdgcn_sched_barrier(
-                    0); // prevent from messing up the order of global loads
-            }
-
-            if constexpr(k0_loops > 1)
-            {
-                static_for<0, k0_loops - 1, 1>{}([&](auto i_k0) {
-                    block_sync_lds();
-                    gemm_0(st_acc,
-                           get_slice_tile(q_lds_window,
-                                          sequence<0, i_k0 * kK0>{},
-                                          sequence<kM0, (i_k0 + 1) * kK0>{}),
-                           get_slice_tile(k_lds_window,
-                                          sequence<0, i_k0 * kK0>{},
-                                          sequence<kN0, (i_k0 + 1) * kK0>{}));
-                    block_sync_lds();
-                });
-            }
-
-            auto do_block_tile = load_tile(do_dram_window); // prefetch load OGrad tile
-            {                                               // tail
-                block_sync_lds();
-                gemm_0(st_acc,
-                       get_slice_tile(q_lds_window,
-                                      sequence<0, (k0_loops - 1) * kK0>{},
-                                      sequence<kM0, k0_loops * kK0>{}),
-                       get_slice_tile(k_lds_window,
-                                      sequence<0, (k0_loops - 1) * kK0>{},
-                                      sequence<kN0, k0_loops * kK0>{}));
-                block_sync_lds();
-            }
-
-            // STAGE 2, Scale, Add bias, Mask, Softmax, Dropout
-            if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS)
-            {
-                block_sync_lds();
-                auto bias_shuffle_tmp = make_static_distributed_tensor<BiasDataType>(
-                    Policy::template MakeShuffledBiasTileDistribution<Problem>());
-                shuffle_tile(bias_shuffle_tmp, bias_tile);
-                store_tile(biast_lds_shuffle_window, bias_shuffle_tmp);
-                block_sync_lds();
-                auto biast_tile = load_tile(biast_lds_window);
-                tile_elementwise_inout(
-                    [&](auto& x, const auto& y) {
-#if !CK_TILE_FMHA_FWD_FAST_EXP2
-                        x = raw_scale * x + type_convert<AccDataType>(y);
-#else
-                        x = scale * x + log2e_v<AccDataType> * type_convert<AccDataType>(y);
-#endif
-                    },
-                    st_acc,
-                    biast_tile);
-                move_tile_window(bias_dram_window, {kM0, 0});
-            }
-            else if constexpr(BiasEnum == BlockAttentionBiasEnum::ALIBI)
-            {
-                const auto q_origin     = q_dram_block_window.get_window_origin();
-                constexpr auto st_spans = decltype(st_acc)::get_distributed_spans();
-                sweep_tile_span(st_spans[number<0>{}], [&](auto idx0) {
-                    sweep_tile_span(st_spans[number<1>{}], [&](auto idx1) {
-                        const auto tile_idx = get_x_indices_from_distributed_indices(
-                            st_acc.get_tile_distribution(), make_tuple(idx0, idx1));
-
-                        const auto row = q_origin.at(number<0>{}) + tile_idx.at(number<0>{});
-                        const auto col = k_origin.at(number<0>{}) + tile_idx.at(number<1>{});
-                        constexpr auto i_j_idx = make_tuple(idx0, idx1);
-
-#if !CK_TILE_FMHA_FWD_FAST_EXP2
-                        st_acc(i_j_idx) *= raw_scale;
-#else
-                        st_acc(i_j_idx) *= scale;
-#endif
-                        position_encoding.update(st_acc(i_j_idx), row, col);
-                    });
-                });
-            }
-            else
-            {
-#if !CK_TILE_FMHA_FWD_FAST_EXP2
-                tile_elementwise_inout([&raw_scale](auto& x) { x = x * raw_scale; }, st_acc);
-#endif
-            }
-
-            if constexpr(kPadSeqLenK || FmhaMask::IsMasking)
-            {
-                const auto q_origin      = q_dram_block_window.get_window_origin();
-                bool need_perpixel_check = mask.IsEdgeTile(q_origin.at(number<0>{}),
-                                                           k_origin.at(number<0>{}),
-                                                           number<kM0>{},
-                                                           number<kN0>{});
-                if(need_perpixel_check)
-                {
-                    set_tile_if(st_acc, -numeric<AccDataType>::infinity(), [&](auto tile_idx) {
-                        const auto row = q_origin.at(number<0>{}) + tile_idx.at(number<0>{});
-                        const auto col = k_origin.at(number<0>{}) + tile_idx.at(number<1>{});
-                        return mask.IsOutOfBound(row, col);
-                    });
-                }
-            }
-
-            const auto lse = load_tile(lse_dram_window);
-
-            static const auto get_validated_lse = [](LSEDataType raw_lse) {
-                if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS ||
-                             FmhaMask::IsMasking)
-                {
-                    return raw_lse == -numeric<LSEDataType>::infinity()
-                               ? type_convert<LSEDataType>(0.f)
-                               : raw_lse;
-                }
-                else
-                {
-                    return raw_lse;
-                }
-            };
-
-            auto pt                 = SPTBlockTileType{};
-            constexpr auto pt_spans = decltype(pt)::get_distributed_spans();
-            sweep_tile_span(pt_spans[number<0>{}], [&](auto idx0) {
-                constexpr auto i_idx = make_tuple(idx0);
-#if CK_TILE_FMHA_FWD_FAST_EXP2
-                auto row_lse = log2e_v<LSEDataType> * get_validated_lse(lse[i_idx]);
-#endif
-                sweep_tile_span(pt_spans[number<1>{}], [&](auto idx1) {
-                    constexpr auto i_j_idx = make_tuple(idx0, idx1);
-#if CK_TILE_FMHA_FWD_FAST_EXP2
-                    if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS ||
-                                 BiasEnum == BlockAttentionBiasEnum::ALIBI)
-                    {
-                        pt(i_j_idx) = exp2(st_acc[i_j_idx] - row_lse);
-                    }
-                    else
-                    {
-                        pt(i_j_idx) = exp2(scale * st_acc[i_j_idx] - row_lse);
-                    }
-#else
-                    pt(i_j_idx) = exp(st_acc[i_j_idx] - get_validated_lse(lse[i_idx]));
-#endif
-                });
-            });
-
-            if constexpr(kHasDropout)
-            {
-                dropout.Run<decltype(gemm_0), RandValOutputDataType>(
-                    seqlen_q_start + i_total_loops * kM0, pt, randval_dram_window);
-            }
-
-            // STAGE 3, P^T@OGrad^T Gemm1
-            block_sync_lds();
-            store_tile(do_lds_window, do_block_tile); // store the prefetch
-
-            const auto pt_gemm = [&]() {
-                if constexpr(kHasDropout)
-                {
-                    return tile_elementwise_in(
-                        [](const auto& x) { return type_convert<GemmDataType>(x > 0.f ? x : 0.f); },
-                        pt);
-                }
-                else
-                {
-                    return cast_tile<GemmDataType>(pt);
-                }
-            }();
-
-            static_for<0, k1_loops, 1>{}([&](auto i_k1) {
-                block_sync_lds();
-                gemm_1(dv_acc,
-                       get_slice_tile(
-                           pt_gemm, sequence<i_k1 * kK1, 0>{}, sequence<(i_k1 + 1) * kK1, kN0>{}),
-                       get_slice_tile(dot_lds_window,
-                                      sequence<0, i_k1 * kK1>{},
-                                      sequence<kVHeaddim, (i_k1 + 1) * kK1>{}));
-                block_sync_lds();
-            });
-
-            // STAGE 4, OGrad@V Gemm2
-            auto dpt_acc = SPGradTBlockTileType{};
-            clear_tile(dpt_acc); // Initialize PGrad^T
-
-            static_for<0, k2_loops, 1>{}([&](auto i_k2) {
-                block_sync_lds();
-                gemm_2(dpt_acc,
-                       get_slice_tile(do_lds_window,
-                                      sequence<0, i_k2 * kK2>{},
-                                      sequence<kM0, (i_k2 + 1) * kK2>{}),
-                       get_slice_tile(
-                           v, sequence<0, i_k2 * kK2>{}, sequence<kN0, (i_k2 + 1) * kK2>{}));
-                block_sync_lds();
-            });
-
-            // STAGE 5, P^T(PGrad^T - D)
-            const auto d = load_tile(d_dram_window);
-
-            auto dst                 = SPGradTBlockTileType{};
-            constexpr auto dst_spans = decltype(dst)::get_distributed_spans();
-            sweep_tile_span(dst_spans[number<0>{}], [&](auto idx0) {
-                constexpr auto i_idx = make_tuple(idx0);
-                sweep_tile_span(dst_spans[number<1>{}], [&](auto idx1) {
-                    constexpr auto i_j_idx = make_tuple(idx0, idx1);
-                    bool undrop_flag       = pt[i_j_idx] >= 0;
-                    dst(i_j_idx) =
-                        pt[i_j_idx] *
-                        (!kHasDropout || undrop_flag ? (dpt_acc[i_j_idx] - d[i_idx]) : d[i_idx]);
-                });
-            });
-
-            if constexpr(kHasBiasGrad)
-            {
-                const auto dbiast = [&]() {
-                    if constexpr(kHasDropout)
-                    {
-                        return tile_elementwise_in(
-                            [&rp_undrop](const auto& x) {
-                                return type_convert<BiasGradDataType>(x * rp_undrop);
-                            },
-                            dst);
-                    }
-                    else
-                    {
-                        return cast_tile<BiasGradDataType>(dst);
-                    }
-                }();
-                store_tile(biast_lds_shuffle_window, dbiast);
-                block_sync_lds();
-                auto dbiast_tile        = load_tile(dbiast_lds_shuffle_window);
-                auto dbiast_shuffle_tmp = make_static_distributed_tensor<BiasGradDataType>(
-                    Policy::template MakeBiasTileDistribution<Problem>());
-                shuffle_tile(dbiast_shuffle_tmp, dbiast_tile);
-                store_tile(dbias_dram_block_window, dbiast_shuffle_tmp);
-                move_tile_window(dbias_dram_block_window, {kM0, 0});
-            }
-
-            // STAGE 6, SGrad^T@Q^T Gemm3
-            block_sync_lds();
-            const auto dst_gemm = cast_tile<GemmDataType>(dst);
-
-            static_for<0, k3_loops, 1>{}([&](auto i_k3) {
-                block_sync_lds();
-                gemm_3(dk_acc,
-                       get_slice_tile(
-                           dst_gemm, sequence<i_k3 * kK3, 0>{}, sequence<(i_k3 + 1) * kK3, kN0>{}),
-                       get_slice_tile(qt_lds_window,
-                                      sequence<0, i_k3 * kK3>{},
-                                      sequence<kQKHeaddim, (i_k3 + 1) * kK3>{}));
-                block_sync_lds();
-            });
-
-            // STAGE 7, SGrad@K^T Gemm4
-            store_tile(ds_lds_window, dst_gemm);
-
-            auto dq_acc = QGradBlockTileType{};
-            clear_tile(dq_acc); // Initialize QGrad
-
-            block_sync_lds();
-
-            static_for<0, k4_loops, 1>{}([&](auto i_k4) {
-                gemm_4(dq_acc,
-                       get_slice_tile(ds_lds_window,
-                                      sequence<0, i_k4 * kK4>{},
-                                      sequence<kM0, (i_k4 + 1) * kK4>{}),
-                       get_slice_tile(kt_lds_window,
-                                      sequence<0, i_k4 * kK4>{},
-                                      sequence<kQKHeaddim, (i_k4 + 1) * kK4>{}));
-            });
-
-            // QGrad Scale
-            if constexpr(kHasDropout)
-            {
-                tile_elementwise_inout([&scale_rp_undrop](auto& x) { x = x * scale_rp_undrop; },
-                                       dq_acc);
-            }
-            else
-            {
-                tile_elementwise_inout([&raw_scale](auto& x) { x = x * raw_scale; }, dq_acc);
-            }
-            const auto dq = cast_tile<QGradDataType>(dq_acc);
-            update_tile(dq_dram_block_window, dq);
-
-            // move tile windows
-            move_tile_window(q_dram_block_window, {kM0, 0});
-            move_tile_window(dq_dram_block_window, {kM0, 0});
-            move_tile_window(do_dram_block_window, {kM0, 0});
-            move_tile_window(lse_dram_window, {kM0});
-            move_tile_window(d_dram_window, {kM0});
-        } while(++i_total_loops < num_total_loop);
-
-        // KGrad Scale
-        if constexpr(kHasDropout)
-        {
-            tile_elementwise_inout([&scale_rp_undrop](auto& x) { x = x * scale_rp_undrop; },
-                                   dk_acc);
-        }
-        else
-        {
-            tile_elementwise_inout([&raw_scale](auto& x) { x = x * raw_scale; }, dk_acc);
-        }
-        // VGrad Scale
-        if constexpr(kHasDropout)
-        {
-            tile_elementwise_inout([&rp_undrop](auto& x) { x = x * rp_undrop; }, dv_acc);
-        }
-
-        return ck_tile::make_tuple(dk_acc, dv_acc);
-    }
-};
-
-} // namespace ck_tile

include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_dq_dk_dv_pipeline_qs_ks_vr_dos_default_policy.hpp

@@ -1,20 +0,0 @@
-// SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
-
-#pragma once
-
-#include "ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_default_policy.hpp"
-
-namespace ck_tile {
-
-// This pipeline is v located in regs, q & k & do located in lds.
-using BlockFmhaBwdDQDKDVPipelineQSKSVROGradSDefaultPolicy =
-    BlockFmhaBwdPipelineDefaultPolicy</* QLoadOnce_      = */ true,
-                                      /* QTLoadOnce_     = */ false,
-                                      /* KLoadOnce_      = */ true,
-                                      /* KTLoadOnce_     = */ false,
-                                      /* VLoadOnce_      = */ true,
-                                      /* OGradLoadOnce_  = */ true,
-                                      /* OGradTLoadOnce_ = */ false>;
-
-} // namespace ck_tile

include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_enum.hpp

@@ -8,9 +8,8 @@ namespace ck_tile {
 // This class is used for codegen pattern matching
 enum class BlockFmhaBwdPipelineEnum
 {
-    KSKTSVR = 0,
-    QSKSVROGradS,
-    KSVR,
+    KRKTRVR_IGLP = 0,
+    KRKTRVR,
 };
 
 } // namespace ck_tile

include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_problem.hpp

@@ -24,7 +24,9 @@ template <typename QDataType_,
           typename BiasGradDataType_,
           typename BlockFmhaShape_,
           bool kIsGroupMode_,
+          bool kIsDeterministic_,
           typename FmhaMask_,
+          typename FmhaDropout_,
           typename Traits_>
 struct BlockFmhaBwdPipelineProblem
 {
@@ -45,10 +47,12 @@ struct BlockFmhaBwdPipelineProblem
     using BiasGradDataType      = remove_cvref_t<BiasGradDataType_>;
     using BlockFmhaShape        = remove_cvref_t<BlockFmhaShape_>;
     using FmhaMask              = remove_cvref_t<FmhaMask_>;
+    using FmhaDropout           = remove_cvref_t<FmhaDropout_>;
     using Traits                = remove_cvref_t<Traits_>;
 
-    static constexpr index_t kBlockSize = BlockFmhaShape::NumWarps * get_warp_size();
-    static constexpr bool kIsGroupMode  = kIsGroupMode_;
+    static constexpr index_t kBlockSize    = BlockFmhaShape::NumWarps * get_warp_size();
+    static constexpr bool kIsGroupMode     = kIsGroupMode_;
+    static constexpr bool kIsDeterministic = kIsDeterministic_;
 
     // attributes from traits
     static constexpr bool kPadSeqLenQ    = Traits::kPadSeqLenQ;
@@ -57,7 +61,6 @@ struct BlockFmhaBwdPipelineProblem
     static constexpr bool kPadHeadDimV   = Traits::kPadHeadDimV;
     static constexpr auto BiasEnum       = Traits::BiasEnum;
     static constexpr bool kHasBiasGrad   = Traits::kHasBiasGrad;
-    static constexpr bool kHasDropout    = Traits::kHasDropout;
     static constexpr index_t kBlockPerCu = Traits::kBlockPerCu;
 };
 
@@ -88,4 +91,35 @@ struct BlockFmhaBwdOGradDotOPipelineProblem
     static constexpr index_t kBlockPerCu = Traits::kBlockPerCu;
 };
 
+template <typename AccDataType_,
+          typename QGradDataType_,
+          index_t kBlockSize_,
+          index_t kM0_,
+          index_t kN0_,
+          index_t kQKHeaddim_,
+          bool kIsGroupMode_,
+          bool kIsDeterministic_,
+          typename Traits_>
+struct BlockFmhaBwdConvertQGradPipelineProblem
+{
+    using AccDataType   = remove_cvref_t<AccDataType_>;
+    using QGradDataType = remove_cvref_t<QGradDataType_>;
+    using Traits        = remove_cvref_t<Traits_>;
+
+    static_assert(0 < kBlockSize_ && kBlockSize_ % get_warp_size() == 0,
+                  "kBlockSize should be divisible by get_warp_size()");
+
+    static constexpr index_t kBlockSize    = kBlockSize_;
+    static constexpr index_t kM0           = kM0_;
+    static constexpr index_t kN0           = kN0_;
+    static constexpr index_t kQKHeaddim    = kQKHeaddim_;
+    static constexpr bool kIsGroupMode     = kIsGroupMode_;
+    static constexpr bool kIsDeterministic = kIsDeterministic_;
+
+    // attributes from traits
+    static constexpr bool kPadSeqLenQ    = Traits::kPadSeqLenQ;
+    static constexpr bool kPadHeadDimQ   = Traits::kPadHeadDimQ;
+    static constexpr index_t kBlockPerCu = Traits::kBlockPerCu;
+};
+
 } // namespace ck_tile

include/ck_tile/ops/fmha/pipeline/tile_fmha_traits.hpp

@@ -86,4 +86,14 @@ struct TileFmhaBwdOGradDotOTraits
     static constexpr index_t kBlockPerCu = kBlockPerCu_;
 };
 
+template <bool kPadSeqLenQ_ /* padding for seqlen_q */,
+          bool kPadHeadDimQ_ /* paddding for hdim_q */,
+          index_t kBlockPerCu_ = 2 /* hint to occupancy */>
+struct TileFmhaBwdConvertQGradTraits
+{
+    static constexpr bool kPadSeqLenQ    = kPadSeqLenQ_;
+    static constexpr bool kPadHeadDimQ   = kPadHeadDimQ_;
+    static constexpr index_t kBlockPerCu = kBlockPerCu_;
+};
+
 } // namespace ck_tile

include/ck_tile/ops/gemm.hpp

@@ -5,6 +5,9 @@
 
 #include "ck_tile/ops/gemm/block/block_gemm_areg_bgmem_creg_v1.hpp"
 #include "ck_tile/ops/gemm/block/block_gemm_areg_bgmem_creg_v1_default_policy.hpp"
+#include "ck_tile/ops/gemm/block/block_gemm_areg_breg_creg_v1.hpp"
+#include "ck_tile/ops/gemm/block/block_gemm_areg_breg_creg_v1_custom_policy.hpp"
+#include "ck_tile/ops/gemm/block/block_gemm_areg_breg_creg_v1_default_policy.hpp"
 #include "ck_tile/ops/gemm/block/block_gemm_areg_bsmem_creg_v1.hpp"
 #include "ck_tile/ops/gemm/block/block_gemm_areg_bsmem_creg_v1_custom_policy.hpp"
 #include "ck_tile/ops/gemm/block/block_gemm_areg_bsmem_creg_v1_default_policy.hpp"

include/ck_tile/ops/gemm/block/block_gemm_areg_breg_creg_v1.hpp

@@ -0,0 +1,202 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include "ck_tile/core.hpp"
+#include "ck_tile/ops/gemm/block/block_gemm_areg_breg_creg_v1_default_policy.hpp"
+
+namespace ck_tile {
+
+// A is block distributed tensor
+// B is block distributed tensor
+// C is block distributed tensor
+template <typename Problem_, typename Policy_ = BlockGemmARegBRegCRegV1DefaultPolicy>
+struct BlockGemmARegBRegCRegV1
+{
+    using Problem        = remove_cvref_t<Problem_>;
+    using Policy         = remove_cvref_t<Policy_>;
+    using ADataType      = remove_cvref_t<typename Problem::ADataType>;
+    using BDataType      = remove_cvref_t<typename Problem::BDataType>;
+    using CDataType      = remove_cvref_t<typename Problem::CDataType>;
+    using BlockGemmShape = remove_cvref_t<typename Problem::BlockGemmShape>;
+
+    static constexpr index_t kBlockSize = Problem::kBlockSize;
+
+    // C += A * B
+    template <typename CBlockTensor, typename ABlockTensor, typename BBlockTensor>
+    CK_TILE_DEVICE void operator()(CBlockTensor& c_block_tensor,
+                                   const ABlockTensor& a_block_tensor,
+                                   const BBlockTensor& b_block_tensor) const
+    {
+        static_assert(std::is_same_v<ADataType, remove_cv_t<typename ABlockTensor::DataType>> &&
+                          std::is_same_v<BDataType, remove_cv_t<typename BBlockTensor::DataType>> &&
+                          std::is_same_v<CDataType, remove_cv_t<typename CBlockTensor::DataType>>,
+                      "wrong!");
+
+        constexpr index_t MPerBlock = BlockGemmShape::kM;
+        constexpr index_t NPerBlock = BlockGemmShape::kN;
+        constexpr index_t KPerBlock = BlockGemmShape::kK;
+
+        constexpr auto config = Policy::template GetWarpGemmMWarpNWarp<Problem>();
+
+        using WG = remove_cvref_t<decltype(config.template at<0>())>;
+
+        constexpr index_t MWarp = config.template at<1>();
+        constexpr index_t NWarp = config.template at<2>();
+
+        constexpr index_t MIterPerWarp = MPerBlock / (MWarp * WG::kM);
+        constexpr index_t NIterPerWarp = NPerBlock / (NWarp * WG::kN);
+        constexpr index_t KIterPerWarp = KPerBlock / WG::kK;
+
+        // M->N Warp
+        constexpr auto a_block_outer_dstr_encoding =
+            tile_distribution_encoding<sequence<NWarp>,
+                                       tuple<sequence<MIterPerWarp, MWarp>, sequence<KIterPerWarp>>,
+                                       tuple<sequence<1, 0>>,
+                                       tuple<sequence<1, 0>>,
+                                       sequence<1, 2>,
+                                       sequence<0, 0>>{};
+
+        constexpr auto b_block_outer_dstr_encoding =
+            tile_distribution_encoding<sequence<MWarp>,
+                                       tuple<sequence<NIterPerWarp, NWarp>, sequence<KIterPerWarp>>,
+                                       tuple<sequence<0, 1>>,
+                                       tuple<sequence<0, 1>>,
+                                       sequence<1, 2>,
+                                       sequence<0, 0>>{};
+
+        constexpr auto c_block_outer_dstr_encoding = tile_distribution_encoding<
+            sequence<>,
+            tuple<sequence<MIterPerWarp, MWarp>, sequence<NIterPerWarp, NWarp>>,
+            tuple<sequence<1, 2>>,
+            tuple<sequence<1, 1>>,
+            sequence<1, 2>,
+            sequence<0, 0>>{};
+
+        constexpr auto a_block_dstr_encode = detail::make_embed_tile_distribution_encoding(
+            a_block_outer_dstr_encoding, typename WG::AWarpDstrEncoding{});
+
+        constexpr auto b_block_dstr_encode = detail::make_embed_tile_distribution_encoding(
+            b_block_outer_dstr_encoding, typename WG::BWarpDstrEncoding{});
+
+        constexpr auto c_block_dstr_encode = detail::make_embed_tile_distribution_encoding(
+            c_block_outer_dstr_encoding, typename WG::CWarpDstrEncoding{});
+
+        // check ABC-block-distribution
+        static_assert(
+            std::is_same_v<remove_cvref_t<decltype(a_block_dstr_encode)>,
+                           remove_cvref_t<decltype(ABlockTensor::get_tile_distribution()
+                                                       .get_static_tile_distribution_encoding())>>,
+            "A distribution is wrong!");
+        static_assert(
+            std::is_same_v<remove_cvref_t<decltype(b_block_dstr_encode)>,
+                           remove_cvref_t<decltype(BBlockTensor::get_tile_distribution()
+                                                       .get_static_tile_distribution_encoding())>>,
+            "B distribution is wrong!");
+        static_assert(
+            std::is_same_v<remove_cvref_t<decltype(c_block_dstr_encode)>,
+                           remove_cvref_t<decltype(CBlockTensor::get_tile_distribution()
+                                                       .get_static_tile_distribution_encoding())>>,
+            "C distribution is wrong!");
+
+        using AWarpDstr = typename WG::AWarpDstr;
+        using BWarpDstr = typename WG::BWarpDstr;
+        using CWarpDstr = typename WG::CWarpDstr;
+
+        using AWarpTensor = typename WG::AWarpTensor;
+        using BWarpTensor = typename WG::BWarpTensor;
+        using CWarpTensor = typename WG::CWarpTensor;
+
+        constexpr auto a_warp_y_lengths =
+            to_sequence(AWarpDstr{}.get_ys_to_d_descriptor().get_lengths());
+        constexpr auto b_warp_y_lengths =
+            to_sequence(BWarpDstr{}.get_ys_to_d_descriptor().get_lengths());
+        constexpr auto c_warp_y_lengths =
+            to_sequence(CWarpDstr{}.get_ys_to_d_descriptor().get_lengths());
+
+        constexpr auto a_warp_y_index_zeros = uniform_sequence_gen_t<AWarpDstr::NDimY, 0>{};
+        constexpr auto b_warp_y_index_zeros = uniform_sequence_gen_t<BWarpDstr::NDimY, 0>{};
+        constexpr auto c_warp_y_index_zeros = uniform_sequence_gen_t<CWarpDstr::NDimY, 0>{};
+
+        // hot loop:
+        static_for<0, KIterPerWarp, 1>{}([&](auto kIter) {
+            static_for<0, MIterPerWarp, 1>{}([&](auto mIter) {
+                // read A warp tensor from A Block window
+                AWarpTensor a_warp_tensor;
+
+                a_warp_tensor.get_thread_buffer() = a_block_tensor.get_y_sliced_thread_data(
+                    merge_sequences(sequence<mIter, kIter>{}, a_warp_y_index_zeros),
+                    merge_sequences(sequence<1, 1>{}, a_warp_y_lengths));
+
+                static_for<0, NIterPerWarp, 1>{}([&](auto nIter) {
+                    // read B warp tensor from B block tensor
+                    BWarpTensor b_warp_tensor;
+
+                    b_warp_tensor.get_thread_buffer() = b_block_tensor.get_y_sliced_thread_data(
+                        merge_sequences(sequence<nIter, kIter>{}, b_warp_y_index_zeros),
+                        merge_sequences(sequence<1, 1>{}, b_warp_y_lengths));
+
+                    // read C warp tensor from C block tensor
+                    CWarpTensor c_warp_tensor;
+
+                    c_warp_tensor.get_thread_buffer() = c_block_tensor.get_y_sliced_thread_data(
+                        merge_sequences(sequence<mIter, nIter>{}, c_warp_y_index_zeros),
+                        merge_sequences(sequence<1, 1>{}, c_warp_y_lengths));
+
+                    // warp GEMM
+                    WG{}(c_warp_tensor, a_warp_tensor, b_warp_tensor);
+
+                    // write C warp tensor into C block tensor
+                    c_block_tensor.set_y_sliced_thread_data(
+                        merge_sequences(sequence<mIter, nIter>{}, c_warp_y_index_zeros),
+                        merge_sequences(sequence<1, 1>{}, c_warp_y_lengths),
+                        c_warp_tensor.get_thread_buffer());
+                });
+            });
+        });
+    }
+
+    CK_TILE_DEVICE constexpr auto MakeCBlockTile() const
+    {
+        constexpr index_t MPerBlock = BlockGemmShape::kM;
+        constexpr index_t NPerBlock = BlockGemmShape::kN;
+
+        constexpr auto config = Policy::template GetWarpGemmMWarpNWarp<Problem>();
+
+        using WG = remove_cvref_t<decltype(config.template at<0>())>;
+
+        constexpr index_t MWarp = config.template at<1>();
+        constexpr index_t NWarp = config.template at<2>();
+
+        constexpr index_t MIterPerWarp = MPerBlock / (MWarp * WG::kM);
+        constexpr index_t NIterPerWarp = NPerBlock / (NWarp * WG::kN);
+        // constexpr index_t KIterPerWarp = KPerBlock / WG::kK;
+
+        constexpr auto c_block_outer_dstr_encoding = tile_distribution_encoding<
+            sequence<>,
+            tuple<sequence<MIterPerWarp, MWarp>, sequence<NIterPerWarp, NWarp>>,
+            tuple<sequence<1, 2>>,
+            tuple<sequence<1, 1>>,
+            sequence<1, 2>,
+            sequence<0, 0>>{};
+
+        constexpr auto c_block_dstr_encode = detail::make_embed_tile_distribution_encoding(
+            c_block_outer_dstr_encoding, typename WG::CWarpDstrEncoding{});
+        constexpr auto c_block_dstr = make_static_tile_distribution(c_block_dstr_encode);
+        auto c_block_tensor         = make_static_distributed_tensor<CDataType>(c_block_dstr);
+        return c_block_tensor;
+    }
+
+    // C = A * B
+    template <typename ABlockTensor, typename BBlockTensor>
+    CK_TILE_DEVICE auto operator()(const ABlockTensor& a_block_tensor,
+                                   const BBlockTensor& b_block_tensor) const
+    {
+        auto c_block_tensor = MakeCBlockTile();
+        operator()(c_block_tensor, a_block_tensor, b_block_tensor);
+        return c_block_tensor;
+    }
+};
+
+} // namespace ck_tile

include/ck_tile/ops/gemm/block/block_gemm_areg_breg_creg_v1_custom_policy.hpp

@@ -0,0 +1,36 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include "ck_tile/core.hpp"
+
+namespace ck_tile {
+
+template <typename AType_,
+          typename BType_,
+          typename CType_,
+          typename BlockWarps_,
+          typename WarpGemm_>
+struct BlockGemmARegBRegCRegV1CustomPolicy
+{
+    using AType = remove_cvref_t<AType_>;
+    using BType = remove_cvref_t<BType_>;
+    using CType = remove_cvref_t<CType_>;
+
+    using BlockWarps = remove_cvref_t<BlockWarps_>;
+
+    static constexpr index_t kMWarps = BlockWarps::at(number<0>{});
+    static constexpr index_t kNWarps = BlockWarps::at(number<1>{});
+    static constexpr index_t kKWarps = BlockWarps::at(number<2>{});
+
+    using WarpGemm = remove_cvref_t<WarpGemm_>;
+
+    template <typename Problem>
+    CK_TILE_HOST_DEVICE static constexpr auto GetWarpGemmMWarpNWarp()
+    {
+        return make_tuple(WarpGemm{}, kMWarps, kNWarps);
+    }
+};
+
+} // namespace ck_tile

include/ck_tile/ops/gemm/block/block_gemm_areg_breg_creg_v1_default_policy.hpp

@@ -0,0 +1,33 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include "ck_tile/core.hpp"
+#include "ck_tile/ops/gemm/warp/warp_gemm.hpp"
+
+namespace ck_tile {
+
+// Default policy for BlockGemmARegBRegCRegV1
+// Default policy class should not be templated, put template on member functions instead
+struct BlockGemmARegBRegCRegV1DefaultPolicy
+{
+    template <typename Problem>
+    CK_TILE_HOST_DEVICE static constexpr auto GetWarpGemmMWarpNWarp()
+    {
+        if constexpr(std::is_same_v<typename Problem::ADataType, half_t> &&
+                     std::is_same_v<typename Problem::BDataType, half_t> &&
+                     std::is_same_v<typename Problem::CDataType, float>)
+        {
+            return make_tuple(WarpGemmMfmaF16F16F32M32N32K8TransposedCDistribution{}, 4, 1);
+        }
+        else if constexpr(std::is_same_v<typename Problem::ADataType, bf16_t> &&
+                          std::is_same_v<typename Problem::BDataType, bf16_t> &&
+                          std::is_same_v<typename Problem::CDataType, float>)
+        {
+            return make_tuple(WarpGemmMfmaBf16Bf16F32M32N32K8TransposedCDistribution{}, 4, 1);
+        }
+    }
+};
+
+} // namespace ck_tile

include/ck_tile/ops/gemm/block/block_gemm_areg_bsmem_creg_v1.hpp

@@ -35,16 +35,13 @@ struct BlockGemmARegBSmemCRegV1
                 std::is_same_v<CDataType, remove_cv_t<typename CBlockTensor::DataType>>,
             "wrong!");
 
-        // constexpr index_t MPerBlock = ABlockTensorTmp{}.get_lengths()[number<0>{}];
-        // constexpr index_t NPerBlock = BBlockWindowTmp{}.get_window_lengths()[number<0>{}];
-        // constexpr index_t KPerBlock = ABlockTensorTmp{}.get_lengths()[number<1>{}];
-        constexpr index_t MPerBlock = BlockGemmShape::kM;
-        constexpr index_t NPerBlock = BlockGemmShape::kN;
-        constexpr index_t KPerBlock = BlockGemmShape::kK;
+        constexpr index_t MPerBlock = ABlockTensorTmp{}.get_lengths()[number<0>{}];
+        constexpr index_t NPerBlock = BBlockWindowTmp{}.get_window_lengths()[number<0>{}];
+        constexpr index_t KPerBlock = ABlockTensorTmp{}.get_lengths()[number<1>{}];
 
-        // static_assert(MPerBlock == BlockGemmShape::kM && NPerBlock == BlockGemmShape::kN &&
-        //                   KPerBlock == BlockGemmShape::kK,
-        //               "wrong!");
+        static_assert(MPerBlock == BlockGemmShape::kM && NPerBlock == BlockGemmShape::kN &&
+                          KPerBlock == BlockGemmShape::kK,
+                      "wrong!");
 
         constexpr auto config = Policy::template GetWarpGemmMWarpNWarp<Problem>();
 

include/ck_tile/ops/gemm/block/block_gemm_asmem_breg_creg_v1.hpp

@@ -35,16 +35,13 @@ struct BlockGemmASmemBRegCRegV1
                 std::is_same_v<CDataType, remove_cv_t<typename CBlockTensor::DataType>>,
             "wrong!");
 
-        // constexpr index_t MPerBlock = ABlockWindowTmp{}.get_window_lengths()[number<0>{}];
-        // constexpr index_t NPerBlock = BBlockTensorTmp{}.get_lengths()[number<0>{}];
-        // constexpr index_t KPerBlock = ABlockWindowTmp{}.get_window_lengths()[number<1>{}];
-        constexpr index_t MPerBlock = BlockGemmShape::kM;
-        constexpr index_t NPerBlock = BlockGemmShape::kN;
-        constexpr index_t KPerBlock = BlockGemmShape::kK;
+        constexpr index_t MPerBlock = ABlockWindowTmp{}.get_window_lengths()[number<0>{}];
+        constexpr index_t NPerBlock = BBlockTensorTmp{}.get_lengths()[number<0>{}];
+        constexpr index_t KPerBlock = ABlockWindowTmp{}.get_window_lengths()[number<1>{}];
 
-        // static_assert(MPerBlock == BlockGemmShape::kM && NPerBlock == BlockGemmShape::kN &&
-        //                   KPerBlock == BlockGemmShape::kK,
-        //               "wrong!");
+        static_assert(MPerBlock == BlockGemmShape::kM && NPerBlock == BlockGemmShape::kN &&
+                          KPerBlock == BlockGemmShape::kK,
+                      "wrong!");
 
         constexpr auto config = Policy::template GetWarpGemmMWarpNWarp<Problem>();
 

include/ck_tile/ops/gemm/warp/warp_gemm.hpp

@@ -22,6 +22,9 @@ using WarpGemmMfmaF16F16F32M32N32K16 =
 using WarpGemmMfmaF16F16F32M16N16K32 =
     WarpGemmImpl<WarpGemmAtrributeMfmaIterateK<WarpGemmAttributeMfmaImplF16F16F32M16N16K16, 2>>;
 
+using WarpGemmMfmaF16F16F32M32N32K8SwizzleA = WarpGemmImpl<
+    WarpGemmAtrributeMfmaIterateK_SwizzleA<WarpGemmAttributeMfmaImplF16F16F32M32N32K8, 1>>;
+
 using WarpGemmMfmaF16F16F32M32N32K16SwizzleA = WarpGemmImpl<
     WarpGemmAtrributeMfmaIterateK_SwizzleA<WarpGemmAttributeMfmaImplF16F16F32M32N32K8, 2>>;
 
@@ -59,6 +62,9 @@ using WarpGemmMfmaBf16Bf16F32M32N32K16 =
 using WarpGemmMfmaBf16Bf16F32M16N16K32 =
     WarpGemmImpl<WarpGemmAtrributeMfmaIterateK<WarpGemmAttributeMfmaImplBf16Bf16F32M16N16K16, 2>>;
 
+using WarpGemmMfmaBf16Bf16F32M32N32K8SwizzleA = WarpGemmImpl<
+    WarpGemmAtrributeMfmaIterateK_SwizzleA<WarpGemmAttributeMfmaImplBf16Bf16F32M32N32K8, 1>>;
+
 using WarpGemmMfmaBf16Bf16F32M32N32K16SwizzleA = WarpGemmImpl<
     WarpGemmAtrributeMfmaIterateK_SwizzleA<WarpGemmAttributeMfmaImplBf16Bf16F32M32N32K8, 2>>;
 

include/ck_tile/ops/gemm/warp/warp_gemm_attribute_mfma.hpp

@@ -119,9 +119,9 @@ struct WarpGemmAtrributeMfmaIterateK
 
         static_for<0, kKIter, 1>{}([&](auto iKIter) {
             Impl{}(c_vec,
-                   reinterpret_cast<const buf_a>(a_vec)
+                   reinterpret_cast<const buf_a&>(a_vec)
                        .template get_as<typename Impl::AVecType>()[iKIter],
-                   reinterpret_cast<const buf_b>(b_vec)
+                   reinterpret_cast<const buf_b&>(b_vec)
                        .template get_as<typename Impl::BVecType>()[iKIter]);
         });
     }
@@ -135,15 +135,15 @@ struct WarpGemmAtrributeMfmaIterateK
 
         // c = a * b
         auto c_vec = Impl{}(
-            reinterpret_cast<const buf_a>(a_vec).template get_as<typename Impl::AVecType>()[I0],
-            reinterpret_cast<const buf_b>(b_vec).template get_as<typename Impl::BVecType>()[I0]);
+            reinterpret_cast<const buf_a&>(a_vec).template get_as<typename Impl::AVecType>()[I0],
+            reinterpret_cast<const buf_b&>(b_vec).template get_as<typename Impl::BVecType>()[I0]);
 
         // c += a * b
         static_for<1, kKIter, 1>{}([&](auto iKIter) {
             Impl{}(c_vec,
-                   reinterpret_cast<const buf_a>(a_vec)
+                   reinterpret_cast<const buf_a&>(a_vec)
                        .template get_as<typename Impl::AVecType>()[iKIter],
-                   reinterpret_cast<const buf_b>(b_vec)
+                   reinterpret_cast<const buf_b&>(b_vec)
                        .template get_as<typename Impl::BVecType>()[iKIter]);
         });
 

include/ck_tile/ops/gemm/warp/warp_gemm_dispatcher.hpp

@@ -15,7 +15,8 @@ template <typename AType,
           index_t MPerWave,
           index_t NPerWave,
           index_t KPerWave,
-          bool TransposeC>
+          bool TransposeC,
+          bool SwizzleA = false>
 struct WarpGemmMfmaDispatcher;
 
 // clang-format off
@@ -29,6 +30,9 @@ template<> struct WarpGemmMfmaDispatcher<ck_tile::half_t, ck_tile::half_t, float
 template<> struct WarpGemmMfmaDispatcher<ck_tile::half_t, ck_tile::half_t, float, 16, 16, 32, false> { using Type = WarpGemmMfmaF16F16F32M16N16K32; };
 template<> struct WarpGemmMfmaDispatcher<ck_tile::half_t, ck_tile::half_t, float, 16, 16, 32, true> { using Type = WarpGemmMfmaF16F16F32M16N16K32TransposedCDistribution; };
 
+template<> struct WarpGemmMfmaDispatcher<ck_tile::half_t, ck_tile::half_t, float, 32, 32,  8, false, true> { using Type = WarpGemmMfmaF16F16F32M32N32K8SwizzleA; };
+template<> struct WarpGemmMfmaDispatcher<ck_tile::half_t, ck_tile::half_t, float, 32, 32, 16, false, true> { using Type = WarpGemmMfmaF16F16F32M32N32K16SwizzleA; };
+
 // bf16
 template<> struct WarpGemmMfmaDispatcher<ck_tile::bf16_t, ck_tile::bf16_t, float, 32, 32,  8, false> { using Type = WarpGemmMfmaBf16Bf16F32M32N32K8; };
 template<> struct WarpGemmMfmaDispatcher<ck_tile::bf16_t, ck_tile::bf16_t, float, 32, 32,  8, true> { using Type = WarpGemmMfmaBf16Bf16F32M32N32K8TransposedCDistribution; };
@@ -39,6 +43,9 @@ template<> struct WarpGemmMfmaDispatcher<ck_tile::bf16_t, ck_tile::bf16_t, float
 template<> struct WarpGemmMfmaDispatcher<ck_tile::bf16_t, ck_tile::bf16_t, float, 16, 16, 32, false> { using Type = WarpGemmMfmaBf16Bf16F32M16N16K32; };
 template<> struct WarpGemmMfmaDispatcher<ck_tile::bf16_t, ck_tile::bf16_t, float, 16, 16, 32, true> { using Type = WarpGemmMfmaBf16Bf16F32M16N16K32TransposedCDistribution; };
 
+template<> struct WarpGemmMfmaDispatcher<ck_tile::bf16_t, ck_tile::bf16_t, float, 32, 32,  8, false, true> { using Type = WarpGemmMfmaBf16Bf16F32M32N32K8SwizzleA; };
+template<> struct WarpGemmMfmaDispatcher<ck_tile::bf16_t, ck_tile::bf16_t, float, 32, 32, 16, false, true> { using Type = WarpGemmMfmaBf16Bf16F32M32N32K16SwizzleA; };
+
 // fp8
 template<> struct WarpGemmMfmaDispatcher<ck_tile::fp8_t, ck_tile::fp8_t, float, 32, 32,  16, false> { using Type = WarpGemmMfma_f32_32x32x16_fp8_fp8; };
 template<> struct WarpGemmMfmaDispatcher<ck_tile::fp8_t, ck_tile::fp8_t, float, 32, 32,  16, true> { using Type = WarpGemmMfma_f32_32x32x16_fp8_fp8_CTransposed; };
@@ -58,8 +65,15 @@ template <typename AType,
           index_t MPerWave,
           index_t NPerWave,
           index_t KPerWave,
-          bool TransposeC>
-using WarpGemmMfmaDispatcher = typename impl::
-    WarpGemmMfmaDispatcher<AType, BType, CType, MPerWave, NPerWave, KPerWave, TransposeC>::Type;
+          bool TransposeC,
+          bool SwizzleA = false>
+using WarpGemmMfmaDispatcher = typename impl::WarpGemmMfmaDispatcher<AType,
+                                                                     BType,
+                                                                     CType,
+                                                                     MPerWave,
+                                                                     NPerWave,
+                                                                     KPerWave,
+                                                                     TransposeC,
+                                                                     SwizzleA>::Type;
 
 } // namespace ck_tile

library/include/ck/library/tensor_operation_instance/gpu/gemm_multiply_multiply.hpp

@@ -18,134 +18,82 @@ namespace device {
 namespace instance {
 #if(defined(CK_ENABLE_BF16) || defined(CK_ENABLE_FP8))
 void add_device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_comp_default_instances(
-    std::vector<std::unique_ptr<DeviceGemmMultipleD<Row,
-                                                    Col,
-                                                    Tuple<Row, Col>,
-                                                    Row,
-                                                    F8,
-                                                    F8,
-                                                    Tuple<F32, F32>,
-                                                    BF16,
-                                                    PassThrough,
-                                                    PassThrough,
-                                                    MultiplyMultiply>>>& instances);
+    std::vector<std::unique_ptr<DeviceGemmMultipleDSplitK<Row,
+                                                          Col,
+                                                          Tuple<Row, Col>,
+                                                          Row,
+                                                          F8,
+                                                          F8,
+                                                          Tuple<F32, F32>,
+                                                          BF16,
+                                                          PassThrough,
+                                                          PassThrough,
+                                                          MultiplyMultiply>>>& instances);
 
 void add_device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_comp_kpadding_instances(
-    std::vector<std::unique_ptr<DeviceGemmMultipleD<Row,
-                                                    Col,
-                                                    Tuple<Row, Col>,
-                                                    Row,
-                                                    F8,
-                                                    F8,
-                                                    Tuple<F32, F32>,
-                                                    BF16,
-                                                    PassThrough,
-                                                    PassThrough,
-                                                    MultiplyMultiply>>>& instances);
-
-void add_device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_comp_mnpadding_instances(
-    std::vector<std::unique_ptr<DeviceGemmMultipleD<Row,
-                                                    Col,
-                                                    Tuple<Row, Col>,
-                                                    Row,
-                                                    F8,
-                                                    F8,
-                                                    Tuple<F32, F32>,
-                                                    BF16,
-                                                    PassThrough,
-                                                    PassThrough,
-                                                    MultiplyMultiply>>>& instances);
-
-void add_device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_comp_mnkpadding_instances(
-    std::vector<std::unique_ptr<DeviceGemmMultipleD<Row,
-                                                    Col,
-                                                    Tuple<Row, Col>,
-                                                    Row,
-                                                    F8,
-                                                    F8,
-                                                    Tuple<F32, F32>,
-                                                    BF16,
-                                                    PassThrough,
-                                                    PassThrough,
-                                                    MultiplyMultiply>>>& instances);
+    std::vector<std::unique_ptr<DeviceGemmMultipleDSplitK<Row,
+                                                          Col,
+                                                          Tuple<Row, Col>,
+                                                          Row,
+                                                          F8,
+                                                          F8,
+                                                          Tuple<F32, F32>,
+                                                          BF16,
+                                                          PassThrough,
+                                                          PassThrough,
+                                                          MultiplyMultiply>>>& instances);
 
 void add_device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v1_default_instances(
-    std::vector<std::unique_ptr<DeviceGemmMultipleD<Row,
-                                                    Col,
-                                                    Tuple<Row, Col>,
-                                                    Row,
-                                                    F8,
-                                                    F8,
-                                                    Tuple<F32, F32>,
-                                                    BF16,
-                                                    PassThrough,
-                                                    PassThrough,
-                                                    MultiplyMultiply>>>& instances);
+    std::vector<std::unique_ptr<DeviceGemmMultipleDSplitK<Row,
+                                                          Col,
+                                                          Tuple<Row, Col>,
+                                                          Row,
+                                                          F8,
+                                                          F8,
+                                                          Tuple<F32, F32>,
+                                                          BF16,
+                                                          PassThrough,
+                                                          PassThrough,
+                                                          MultiplyMultiply>>>& instances);
 
 void add_device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v1_kpadding_instances(
-    std::vector<std::unique_ptr<DeviceGemmMultipleD<Row,
-                                                    Col,
-                                                    Tuple<Row, Col>,
-                                                    Row,
-                                                    F8,
-                                                    F8,
-                                                    Tuple<F32, F32>,
-                                                    BF16,
-                                                    PassThrough,
-                                                    PassThrough,
-                                                    MultiplyMultiply>>>& instances);
-
-void add_device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v1_mnkpadding_instances(
-    std::vector<std::unique_ptr<DeviceGemmMultipleD<Row,
-                                                    Col,
-                                                    Tuple<Row, Col>,
-                                                    Row,
-                                                    F8,
-                                                    F8,
-                                                    Tuple<F32, F32>,
-                                                    BF16,
-                                                    PassThrough,
-                                                    PassThrough,
-                                                    MultiplyMultiply>>>& instances);
+    std::vector<std::unique_ptr<DeviceGemmMultipleDSplitK<Row,
+                                                          Col,
+                                                          Tuple<Row, Col>,
+                                                          Row,
+                                                          F8,
+                                                          F8,
+                                                          Tuple<F32, F32>,
+                                                          BF16,
+                                                          PassThrough,
+                                                          PassThrough,
+                                                          MultiplyMultiply>>>& instances);
 
 void add_device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v2_default_instances(
-    std::vector<std::unique_ptr<DeviceGemmMultipleD<Row,
-                                                    Col,
-                                                    Tuple<Row, Col>,
-                                                    Row,
-                                                    F8,
-                                                    F8,
-                                                    Tuple<F32, F32>,
-                                                    BF16,
-                                                    PassThrough,
-                                                    PassThrough,
-                                                    MultiplyMultiply>>>& instances);
+    std::vector<std::unique_ptr<DeviceGemmMultipleDSplitK<Row,
+                                                          Col,
+                                                          Tuple<Row, Col>,
+                                                          Row,
+                                                          F8,
+                                                          F8,
+                                                          Tuple<F32, F32>,
+                                                          BF16,
+                                                          PassThrough,
+                                                          PassThrough,
+                                                          MultiplyMultiply>>>& instances);
 
 void add_device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v2_kpadding_instances(
-    std::vector<std::unique_ptr<DeviceGemmMultipleD<Row,
-                                                    Col,
-                                                    Tuple<Row, Col>,
-                                                    Row,
-                                                    F8,
-                                                    F8,
-                                                    Tuple<F32, F32>,
-                                                    BF16,
-                                                    PassThrough,
-                                                    PassThrough,
-                                                    MultiplyMultiply>>>& instances);
-
-void add_device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v2_mnkpadding_instances(
-    std::vector<std::unique_ptr<DeviceGemmMultipleD<Row,
-                                                    Col,
-                                                    Tuple<Row, Col>,
-                                                    Row,
-                                                    F8,
-                                                    F8,
-                                                    Tuple<F32, F32>,
-                                                    BF16,
-                                                    PassThrough,
-                                                    PassThrough,
-                                                    MultiplyMultiply>>>& instances);
+    std::vector<std::unique_ptr<DeviceGemmMultipleDSplitK<Row,
+                                                          Col,
+                                                          Tuple<Row, Col>,
+                                                          Row,
+                                                          F8,
+                                                          F8,
+                                                          Tuple<F32, F32>,
+                                                          BF16,
+                                                          PassThrough,
+                                                          PassThrough,
+                                                          MultiplyMultiply>>>& instances);
 #endif
 
 template <typename ADataType,
@@ -154,7 +102,7 @@ template <typename ADataType,
           typename ALayout,
           typename BLayout,
           typename CLayout>
-struct DeviceOperationInstanceFactory<ck::tensor_operation::device::DeviceGemmMultipleD<
+struct DeviceOperationInstanceFactory<ck::tensor_operation::device::DeviceGemmMultipleDSplitK<
     ALayout,
     BLayout,
     Tuple<Row, Col>,
@@ -167,17 +115,18 @@ struct DeviceOperationInstanceFactory<ck::tensor_operation::device::DeviceGemmMu
     ck::tensor_operation::element_wise::PassThrough,
     ck::tensor_operation::element_wise::MultiplyMultiply>>
 {
-    using DeviceOp = DeviceGemmMultipleD<ALayout,
-                                         BLayout,
-                                         Tuple<Row, Col>,
-                                         CLayout,
-                                         ADataType,
-                                         BDataType,
-                                         Tuple<F32, F32>,
-                                         CDataType,
-                                         ck::tensor_operation::element_wise::PassThrough,
-                                         ck::tensor_operation::element_wise::PassThrough,
-                                         ck::tensor_operation::element_wise::MultiplyMultiply>;
+    using DeviceOp =
+        DeviceGemmMultipleDSplitK<ALayout,
+                                  BLayout,
+                                  Tuple<Row, Col>,
+                                  CLayout,
+                                  ADataType,
+                                  BDataType,
+                                  Tuple<F32, F32>,
+                                  CDataType,
+                                  ck::tensor_operation::element_wise::PassThrough,
+                                  ck::tensor_operation::element_wise::PassThrough,
+                                  ck::tensor_operation::element_wise::MultiplyMultiply>;
 
     static auto GetInstances()
     {
@@ -194,24 +143,16 @@ struct DeviceOperationInstanceFactory<ck::tensor_operation::device::DeviceGemmMu
                     op_ptrs);
                 add_device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_comp_kpadding_instances(
                     op_ptrs);
-                add_device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_comp_mnpadding_instances(
-                    op_ptrs);
-                add_device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_comp_mnkpadding_instances(
-                    op_ptrs);
 
                 add_device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v1_default_instances(
                     op_ptrs);
                 add_device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v1_kpadding_instances(
                     op_ptrs);
-                add_device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v1_mnkpadding_instances(
-                    op_ptrs);
 
                 add_device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v2_default_instances(
                     op_ptrs);
                 add_device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v2_kpadding_instances(
                     op_ptrs);
-                add_device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v2_mnkpadding_instances(
-                    op_ptrs);
             }
         }
 #endif

library/include/ck/library/tensor_operation_instance/gpu/gemm_universal.hpp

@@ -77,16 +77,6 @@ void add_device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_comp_kpadding_instances(
         DeviceGemmV2<Row, Col, Row, F16, F16, F16, PassThrough, PassThrough, PassThrough>>>&
         instances);
 
-void add_device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_comp_mnpadding_instances(
-    std::vector<std::unique_ptr<
-        DeviceGemmV2<Row, Col, Row, F16, F16, F16, PassThrough, PassThrough, PassThrough>>>&
-        instances);
-
-void add_device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_comp_mnkpadding_instances(
-    std::vector<std::unique_ptr<
-        DeviceGemmV2<Row, Col, Row, F16, F16, F16, PassThrough, PassThrough, PassThrough>>>&
-        instances);
-
 void add_device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_mem_v1_default_instances(
     std::vector<std::unique_ptr<
         DeviceGemmV2<Row, Col, Row, F16, F16, F16, PassThrough, PassThrough, PassThrough>>>&
@@ -97,11 +87,6 @@ void add_device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_mem_v1_kpadding_instance
         DeviceGemmV2<Row, Col, Row, F16, F16, F16, PassThrough, PassThrough, PassThrough>>>&
         instances);
 
-void add_device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_mem_v1_mnkpadding_instances(
-    std::vector<std::unique_ptr<
-        DeviceGemmV2<Row, Col, Row, F16, F16, F16, PassThrough, PassThrough, PassThrough>>>&
-        instances);
-
 void add_device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_mem_v2_default_instances(
     std::vector<std::unique_ptr<
         DeviceGemmV2<Row, Col, Row, F16, F16, F16, PassThrough, PassThrough, PassThrough>>>&
@@ -111,13 +96,8 @@ void add_device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_mem_v2_kpadding_instance
     std::vector<std::unique_ptr<
         DeviceGemmV2<Row, Col, Row, F16, F16, F16, PassThrough, PassThrough, PassThrough>>>&
         instances);
-
-void add_device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_mem_v2_mnkpadding_instances(
-    std::vector<std::unique_ptr<
-        DeviceGemmV2<Row, Col, Row, F16, F16, F16, PassThrough, PassThrough, PassThrough>>>&
-        instances);
 #endif
-#if(defined(CK_ENABLE_FP16) || defined(CK_ENABLE_FP8))
+#if(defined(CK_ENABLE_FP16) && defined(CK_ENABLE_FP8))
 void add_device_gemm_xdl_universal_f16_f8_f16_mk_kn_mn_comp_default_instances(
     std::vector<std::unique_ptr<
         DeviceGemmV2<Row, Row, Row, F16, F8, F16, PassThrough, PassThrough, PassThrough>>>&
@@ -177,16 +157,6 @@ void add_device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_comp_kpadding_instances(
         DeviceGemmV2<Row, Col, Row, F16, F8, F16, PassThrough, PassThrough, PassThrough>>>&
         instances);
 
-void add_device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_comp_mnpadding_instances(
-    std::vector<std::unique_ptr<
-        DeviceGemmV2<Row, Col, Row, F16, F8, F16, PassThrough, PassThrough, PassThrough>>>&
-        instances);
-
-void add_device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_comp_mnkpadding_instances(
-    std::vector<std::unique_ptr<
-        DeviceGemmV2<Row, Col, Row, F16, F8, F16, PassThrough, PassThrough, PassThrough>>>&
-        instances);
-
 void add_device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_mem_v1_default_instances(
     std::vector<std::unique_ptr<
         DeviceGemmV2<Row, Col, Row, F16, F8, F16, PassThrough, PassThrough, PassThrough>>>&
@@ -196,12 +166,6 @@ void add_device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_mem_v1_kpadding_instances
     std::vector<std::unique_ptr<
         DeviceGemmV2<Row, Col, Row, F16, F8, F16, PassThrough, PassThrough, PassThrough>>>&
         instances);
-
-void add_device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_mem_v1_mnkpadding_instances(
-    std::vector<std::unique_ptr<
-        DeviceGemmV2<Row, Col, Row, F16, F8, F16, PassThrough, PassThrough, PassThrough>>>&
-        instances);
-
 void add_device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_mem_v2_default_instances(
     std::vector<std::unique_ptr<
         DeviceGemmV2<Row, Col, Row, F16, F8, F16, PassThrough, PassThrough, PassThrough>>>&
@@ -212,10 +176,6 @@ void add_device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_mem_v2_kpadding_instances
         DeviceGemmV2<Row, Col, Row, F16, F8, F16, PassThrough, PassThrough, PassThrough>>>&
         instances);
 
-void add_device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_mem_v2_mnkpadding_instances(
-    std::vector<std::unique_ptr<
-        DeviceGemmV2<Row, Col, Row, F16, F8, F16, PassThrough, PassThrough, PassThrough>>>&
-        instances);
 void add_device_gemm_xdl_universal_f8_f16_f16_mk_kn_mn_comp_default_instances(
     std::vector<std::unique_ptr<
         DeviceGemmV2<Row, Row, Row, F8, F16, F16, PassThrough, PassThrough, PassThrough>>>&
@@ -275,16 +235,6 @@ void add_device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_comp_kpadding_instances(
         DeviceGemmV2<Row, Col, Row, F8, F16, F16, PassThrough, PassThrough, PassThrough>>>&
         instances);
 
-void add_device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_comp_mnpadding_instances(
-    std::vector<std::unique_ptr<
-        DeviceGemmV2<Row, Col, Row, F8, F16, F16, PassThrough, PassThrough, PassThrough>>>&
-        instances);
-
-void add_device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_comp_mnkpadding_instances(
-    std::vector<std::unique_ptr<
-        DeviceGemmV2<Row, Col, Row, F8, F16, F16, PassThrough, PassThrough, PassThrough>>>&
-        instances);
-
 void add_device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_mem_v1_default_instances(
     std::vector<std::unique_ptr<
         DeviceGemmV2<Row, Col, Row, F8, F16, F16, PassThrough, PassThrough, PassThrough>>>&
@@ -295,11 +245,6 @@ void add_device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_mem_v1_kpadding_instances
         DeviceGemmV2<Row, Col, Row, F8, F16, F16, PassThrough, PassThrough, PassThrough>>>&
         instances);
 
-void add_device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_mem_v1_mnkpadding_instances(
-    std::vector<std::unique_ptr<
-        DeviceGemmV2<Row, Col, Row, F8, F16, F16, PassThrough, PassThrough, PassThrough>>>&
-        instances);
-
 void add_device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_mem_v2_default_instances(
     std::vector<std::unique_ptr<
         DeviceGemmV2<Row, Col, Row, F8, F16, F16, PassThrough, PassThrough, PassThrough>>>&
@@ -309,11 +254,6 @@ void add_device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_mem_v2_kpadding_instances
     std::vector<std::unique_ptr<
         DeviceGemmV2<Row, Col, Row, F8, F16, F16, PassThrough, PassThrough, PassThrough>>>&
         instances);
-
-void add_device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_mem_v2_mnkpadding_instances(
-    std::vector<std::unique_ptr<
-        DeviceGemmV2<Row, Col, Row, F8, F16, F16, PassThrough, PassThrough, PassThrough>>>&
-        instances);
 #endif
 #ifdef CK_ENABLE_BF16
 void add_device_gemm_xdl_universal_bf16_bf16_bf16_mk_kn_mn_comp_default_instances(
@@ -376,16 +316,6 @@ void add_device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_comp_kpadding_instanc
         DeviceGemmV2<Row, Col, Row, BF16, BF16, BF16, PassThrough, PassThrough, PassThrough>>>&
         instances);
 
-void add_device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_comp_mnpadding_instances(
-    std::vector<std::unique_ptr<
-        DeviceGemmV2<Row, Col, Row, BF16, BF16, BF16, PassThrough, PassThrough, PassThrough>>>&
-        instances);
-
-void add_device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_comp_mnkpadding_instances(
-    std::vector<std::unique_ptr<
-        DeviceGemmV2<Row, Col, Row, BF16, BF16, BF16, PassThrough, PassThrough, PassThrough>>>&
-        instances);
-
 void add_device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_mem_v1_default_instances(
     std::vector<std::unique_ptr<
         DeviceGemmV2<Row, Col, Row, BF16, BF16, BF16, PassThrough, PassThrough, PassThrough>>>&
@@ -396,11 +326,6 @@ void add_device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_mem_v1_kpadding_insta
         DeviceGemmV2<Row, Col, Row, BF16, BF16, BF16, PassThrough, PassThrough, PassThrough>>>&
         instances);
 
-void add_device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_mem_v1_mnkpadding_instances(
-    std::vector<std::unique_ptr<
-        DeviceGemmV2<Row, Col, Row, BF16, BF16, BF16, PassThrough, PassThrough, PassThrough>>>&
-        instances);
-
 void add_device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_mem_v2_default_instances(
     std::vector<std::unique_ptr<
         DeviceGemmV2<Row, Col, Row, BF16, BF16, BF16, PassThrough, PassThrough, PassThrough>>>&
@@ -410,13 +335,53 @@ void add_device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_mem_v2_kpadding_insta
     std::vector<std::unique_ptr<
         DeviceGemmV2<Row, Col, Row, BF16, BF16, BF16, PassThrough, PassThrough, PassThrough>>>&
         instances);
-
-void add_device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_mem_v2_mnkpadding_instances(
-    std::vector<std::unique_ptr<
-        DeviceGemmV2<Row, Col, Row, BF16, BF16, BF16, PassThrough, PassThrough, PassThrough>>>&
-        instances);
 #endif
-#if(defined(CK_ENABLE_BF16) || defined(CK_ENABLE_FP8))
+#if(defined(CK_ENABLE_BF16) && defined(CK_ENABLE_FP8))
+void add_device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_comp_default_instances(
+    std::vector<std::unique_ptr<
+        DeviceGemmV2<Row, Row, Row, F8, F8, BF16, PassThrough, PassThrough, PassThrough>>>&
+        instances);
+
+void add_device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_comp_kpadding_instances(
+    std::vector<std::unique_ptr<
+        DeviceGemmV2<Row, Row, Row, F8, F8, BF16, PassThrough, PassThrough, PassThrough>>>&
+        instances);
+
+void add_device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_comp_nkpadding_instances(
+    std::vector<std::unique_ptr<
+        DeviceGemmV2<Row, Row, Row, F8, F8, BF16, PassThrough, PassThrough, PassThrough>>>&
+        instances);
+
+void add_device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v1_default_instances(
+    std::vector<std::unique_ptr<
+        DeviceGemmV2<Row, Row, Row, F8, F8, BF16, PassThrough, PassThrough, PassThrough>>>&
+        instances);
+
+void add_device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v1_kpadding_instances(
+    std::vector<std::unique_ptr<
+        DeviceGemmV2<Row, Row, Row, F8, F8, BF16, PassThrough, PassThrough, PassThrough>>>&
+        instances);
+
+void add_device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v1_nkpadding_instances(
+    std::vector<std::unique_ptr<
+        DeviceGemmV2<Row, Row, Row, F8, F8, BF16, PassThrough, PassThrough, PassThrough>>>&
+        instances);
+
+void add_device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v2_default_instances(
+    std::vector<std::unique_ptr<
+        DeviceGemmV2<Row, Row, Row, F8, F8, BF16, PassThrough, PassThrough, PassThrough>>>&
+        instances);
+
+void add_device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v2_kpadding_instances(
+    std::vector<std::unique_ptr<
+        DeviceGemmV2<Row, Row, Row, F8, F8, BF16, PassThrough, PassThrough, PassThrough>>>&
+        instances);
+
+void add_device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v2_nkpadding_instances(
+    std::vector<std::unique_ptr<
+        DeviceGemmV2<Row, Row, Row, F8, F8, BF16, PassThrough, PassThrough, PassThrough>>>&
+        instances);
+
 void add_device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_comp_default_instances(
     std::vector<std::unique_ptr<
         DeviceGemmV2<Row, Col, Row, F8, F8, BF16, PassThrough, PassThrough, PassThrough>>>&
@@ -427,16 +392,6 @@ void add_device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_comp_kpadding_instances(
         DeviceGemmV2<Row, Col, Row, F8, F8, BF16, PassThrough, PassThrough, PassThrough>>>&
         instances);
 
-void add_device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_comp_mnpadding_instances(
-    std::vector<std::unique_ptr<
-        DeviceGemmV2<Row, Col, Row, F8, F8, BF16, PassThrough, PassThrough, PassThrough>>>&
-        instances);
-
-void add_device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_comp_mnkpadding_instances(
-    std::vector<std::unique_ptr<
-        DeviceGemmV2<Row, Col, Row, F8, F8, BF16, PassThrough, PassThrough, PassThrough>>>&
-        instances);
-
 void add_device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_v1_default_instances(
     std::vector<std::unique_ptr<
         DeviceGemmV2<Row, Col, Row, F8, F8, BF16, PassThrough, PassThrough, PassThrough>>>&
@@ -447,11 +402,6 @@ void add_device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_v1_kpadding_instances
         DeviceGemmV2<Row, Col, Row, F8, F8, BF16, PassThrough, PassThrough, PassThrough>>>&
         instances);
 
-void add_device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_v1_mnkpadding_instances(
-    std::vector<std::unique_ptr<
-        DeviceGemmV2<Row, Col, Row, F8, F8, BF16, PassThrough, PassThrough, PassThrough>>>&
-        instances);
-
 void add_device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_v2_default_instances(
     std::vector<std::unique_ptr<
         DeviceGemmV2<Row, Col, Row, F8, F8, BF16, PassThrough, PassThrough, PassThrough>>>&
@@ -461,11 +411,6 @@ void add_device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_v2_kpadding_instances
     std::vector<std::unique_ptr<
         DeviceGemmV2<Row, Col, Row, F8, F8, BF16, PassThrough, PassThrough, PassThrough>>>&
         instances);
-
-void add_device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_v2_mnkpadding_instances(
-    std::vector<std::unique_ptr<
-        DeviceGemmV2<Row, Col, Row, F8, F8, BF16, PassThrough, PassThrough, PassThrough>>>&
-        instances);
 #endif
 
 template <typename ADataType,
@@ -532,28 +477,20 @@ struct DeviceOperationInstanceFactory<
             {
                 add_device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_comp_default_instances(op_ptrs);
                 add_device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_comp_kpadding_instances(op_ptrs);
-                add_device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_comp_mnpadding_instances(
-                    op_ptrs);
-                add_device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_comp_mnkpadding_instances(
-                    op_ptrs);
 
                 add_device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_mem_v1_default_instances(
                     op_ptrs);
                 add_device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_mem_v1_kpadding_instances(
                     op_ptrs);
-                add_device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_mem_v1_mnkpadding_instances(
-                    op_ptrs);
 
                 add_device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_mem_v2_default_instances(
                     op_ptrs);
                 add_device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_mem_v2_kpadding_instances(
                     op_ptrs);
-                add_device_gemm_xdl_universal_f16_f16_f16_mk_nk_mn_mem_v2_mnkpadding_instances(
-                    op_ptrs);
             }
         }
 #endif
-#if(defined(CK_ENABLE_FP16) || defined(CK_ENABLE_FP8))
+#if(defined(CK_ENABLE_FP16) && defined(CK_ENABLE_FP8))
         if constexpr(is_same_v<ADataType, half_t> && is_same_v<BDataType, f8_t> &&
                      is_same_v<CDataType, half_t>)
         {
@@ -562,21 +499,14 @@ struct DeviceOperationInstanceFactory<
             {
                 add_device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_comp_default_instances(op_ptrs);
                 add_device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_comp_kpadding_instances(op_ptrs);
-                add_device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_comp_mnpadding_instances(op_ptrs);
-                add_device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_comp_mnkpadding_instances(
-                    op_ptrs);
 
                 add_device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_mem_v1_default_instances(op_ptrs);
                 add_device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_mem_v1_kpadding_instances(
                     op_ptrs);
-                add_device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_mem_v1_mnkpadding_instances(
-                    op_ptrs);
 
                 add_device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_mem_v2_default_instances(op_ptrs);
                 add_device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_mem_v2_kpadding_instances(
                     op_ptrs);
-                add_device_gemm_xdl_universal_f16_f8_f16_mk_nk_mn_mem_v2_mnkpadding_instances(
-                    op_ptrs);
             }
             else if constexpr(is_same_v<ALayout, Row> && is_same_v<BLayout, Row> &&
                               is_same_v<CLayout, Row>)
@@ -608,21 +538,14 @@ struct DeviceOperationInstanceFactory<
             {
                 add_device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_comp_default_instances(op_ptrs);
                 add_device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_comp_kpadding_instances(op_ptrs);
-                add_device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_comp_mnpadding_instances(op_ptrs);
-                add_device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_comp_mnkpadding_instances(
-                    op_ptrs);
 
                 add_device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_mem_v1_default_instances(op_ptrs);
                 add_device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_mem_v1_kpadding_instances(
                     op_ptrs);
-                add_device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_mem_v1_mnkpadding_instances(
-                    op_ptrs);
 
                 add_device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_mem_v2_default_instances(op_ptrs);
                 add_device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_mem_v2_kpadding_instances(
                     op_ptrs);
-                add_device_gemm_xdl_universal_f8_f16_f16_mk_nk_mn_mem_v2_mnkpadding_instances(
-                    op_ptrs);
             }
             else if constexpr(is_same_v<ALayout, Row> && is_same_v<BLayout, Row> &&
                               is_same_v<CLayout, Row>)
@@ -684,51 +607,55 @@ struct DeviceOperationInstanceFactory<
                     op_ptrs);
                 add_device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_comp_kpadding_instances(
                     op_ptrs);
-                add_device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_comp_mnpadding_instances(
-                    op_ptrs);
-                add_device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_comp_mnkpadding_instances(
-                    op_ptrs);
 
                 add_device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_mem_v1_default_instances(
                     op_ptrs);
                 add_device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_mem_v1_kpadding_instances(
                     op_ptrs);
-                add_device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_mem_v1_mnkpadding_instances(
-                    op_ptrs);
 
                 add_device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_mem_v2_default_instances(
                     op_ptrs);
                 add_device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_mem_v2_kpadding_instances(
                     op_ptrs);
-                add_device_gemm_xdl_universal_bf16_bf16_bf16_mk_nk_mn_mem_v2_mnkpadding_instances(
-                    op_ptrs);
             }
         }
 #endif
-#if(defined(CK_ENABLE_BF16) || defined(CK_ENABLE_FP8))
+#if(defined(CK_ENABLE_BF16) && defined(CK_ENABLE_FP8))
         if constexpr(is_same_v<ADataType, f8_t> && is_same_v<BDataType, f8_t> &&
                      is_same_v<CDataType, bhalf_t>)
         {
-            if constexpr(is_same_v<ALayout, Row> && is_same_v<BLayout, Col> &&
+            if constexpr(is_same_v<ALayout, Row> && is_same_v<BLayout, Row> &&
                          is_same_v<CLayout, Row>)
+            {
+                add_device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_comp_default_instances(op_ptrs);
+                add_device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_comp_kpadding_instances(op_ptrs);
+                add_device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_comp_nkpadding_instances(op_ptrs);
+
+                add_device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v1_default_instances(op_ptrs);
+                add_device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v1_kpadding_instances(
+                    op_ptrs);
+                add_device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v1_nkpadding_instances(
+                    op_ptrs);
+
+                add_device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v2_default_instances(op_ptrs);
+                add_device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v2_kpadding_instances(
+                    op_ptrs);
+                add_device_gemm_xdl_universal_f8_f8_bf16_mk_kn_mn_mem_v2_nkpadding_instances(
+                    op_ptrs);
+            }
+            else if constexpr(is_same_v<ALayout, Row> && is_same_v<BLayout, Col> &&
+                              is_same_v<CLayout, Row>)
             {
                 add_device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_comp_default_instances(op_ptrs);
                 add_device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_comp_kpadding_instances(op_ptrs);
-                add_device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_comp_mnpadding_instances(op_ptrs);
-                add_device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_comp_mnkpadding_instances(
-                    op_ptrs);
 
                 add_device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_v1_default_instances(op_ptrs);
                 add_device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_v1_kpadding_instances(
                     op_ptrs);
-                add_device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_v1_mnkpadding_instances(
-                    op_ptrs);
 
                 add_device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_v2_default_instances(op_ptrs);
                 add_device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_v2_kpadding_instances(
                     op_ptrs);
-                add_device_gemm_xdl_universal_f8_f8_bf16_mk_nk_mn_mem_v2_mnkpadding_instances(
-                    op_ptrs);
             }
         }
 #endif

library/include/ck/library/tensor_operation_instance/gpu/grouped_conv_fwd/device_grouped_conv_fwd_xdl_outelementop_instance.hpp

@@ -184,6 +184,43 @@ using device_grouped_conv_fwd_xdl_outelementop_bf8_f8_instances = std::tuple<
     // clang-format on
     >;
 
+template <index_t NDimSpatial,
+          typename ALayout,
+          typename BLayout,
+          typename DsLayout,
+          typename ELayout,
+          ConvolutionForwardSpecialization ConvSpec,
+          typename OutElementOp>
+using device_grouped_conv_fwd_xdl_outelementop_f8_f8_f32_instances = std::tuple<
+// clang-format off
+        //########################################|     NumDim|      A|      B|          Ds|      E| AData| BData| AccData| CShuffle|       Ds| EData|           A|           B|           CDE|    ConvForward|           GEMM| NumGemmK| Block|  MPer|  NPer|  KPer| AK1| BK1| MPer| NPer| MXdl| NXdl|  ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockLds|  BBlockTransfer| BBlockTransfer| BBlockTransfer| BlockTransfer| BBlockTransfer| BBlockTransfer| BBlockLds|    CShuffle|    CShuffle| CBlockTransferClusterLengths|  CBlockTransfer| Compute| Compute|
+        //########################################|    Spatial| Layout| Layout|      Layout| Layout|  Type|  Type|    Type| DataType| DataType|  Type| Elementwise| Elementwise|   Elementwise| Specialization| Specialization| Prefetch|  Size| Block| Block| Block|    |    |  XDL|  XDL|  Per|  Per|   ThreadCluster|  ThreadCluster| SrcAccessOrder|   SrcVectorDim|      SrcScalar|      DstScalar| AddExtraM|   ThreadCluster|  ThreadCluster| SrcAccessOrder|  SrcVectorDim|      SrcScalar|      DstScalar| AddExtraN| MXdlPerWave| NXdlPerWave|         _MBlock_MWaveMPerXdl| ScalarPerVector|   TypeA|   TypeB|
+        //########################################|           |       |       |            |       |      |      |        |         |         |      |   Operation|   Operation|     Operation|               |               |    Stage|      |      |      |      |    |    |     |     | Wave| Wave| Lengths_K0_M_K1|   ArrangeOrder|               |               |      PerVector|   PerVector_K1|          | Lengths_K0_N_K1|   ArrangeOrder|               |              |      PerVector|   PerVector_K1|          |  PerShuffle|  PerShuffle|         _NBlock_NWaveNPerXdl|   _NWaveNPerXdl|        |        |
+        //########################################|           |       |       |            |       |      |      |        |         |         |      |            |            |              |               |               |         |      |      |      |      |    |    |     |     |     |     |                |               |               |               |               |               |          |                |               |               |              |               |               |          |            |            |                             |                |        |        |
+#ifdef CK_ENABLE_FP8
+        // generic instance
+        DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle<NDimSpatial,ALayout,BLayout,  DsLayout,ELayout,    F8,    F8,     F32,      F32,  Tuple<>,   F32, PassThrough, PassThrough,  OutElementOp,       ConvSpec, GemmMNKPadding,        1,    64,    64,    64,    32,   8,   8,   32,   32,    2,    2,     S<4, 16, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              1,              8,         1,     S<4, 16, 1>,     S<1, 0, 2>,     S<1, 0, 2>,             2,              1,              8,         1,           1,           1,               S<1, 16, 1, 4>,               1,     F8,      F8>,
+        // instances for small conv.K and conv.C
+        DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle<NDimSpatial,ALayout,BLayout,  DsLayout,ELayout,    F8,    F8,     F32,      F32,  Tuple<>,   F32, PassThrough, PassThrough,  OutElementOp,       ConvSpec, GemmMNKPadding,        1,    64,    64,    32,    32,   8,   8,   32,   32,    2,    1,     S<4, 16, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,         1,     S<4, 16, 1>,     S<1, 0, 2>,     S<1, 0, 2>,             2,              8,              8,         1,           1,           1,               S<1, 16, 1, 4>,               1,     F8,      F8>,
+        DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle<NDimSpatial,ALayout,BLayout,  DsLayout,ELayout,    F8,    F8,     F32,      F32,  Tuple<>,   F32, PassThrough, PassThrough,  OutElementOp,       ConvSpec, GemmMNKPadding,        1,   256,   128,   128,    32,   8,   8,   32,   32,    2,    2,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              1,              8,         1,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,             2,              1,              8,         1,           1,           1,               S<1, 32, 1, 8>,               8,     F8,      F8>,
+
+        DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle<NDimSpatial,ALayout,BLayout,  DsLayout,ELayout,    F8,    F8,     F32,      F32,  Tuple<>,   F32, PassThrough, PassThrough,  OutElementOp,       ConvSpec, GemmMNKPadding,        1,   256,   256,   128,    32,   8,   8,   32,   32,    4,    2,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,         1,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,             2,              8,              8,         1,           1,           1,               S<1, 32, 1, 8>,               8,     F8,      F8>,
+        DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle<NDimSpatial,ALayout,BLayout,  DsLayout,ELayout,    F8,    F8,     F32,      F32,  Tuple<>,   F32, PassThrough, PassThrough,  OutElementOp,       ConvSpec, GemmMNKPadding,        1,   256,   128,   256,    32,   8,   8,   32,   32,    2,    4,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,         1,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,             2,              8,              8,         1,           1,           1,               S<1, 32, 1, 8>,               8,     F8,      F8>,
+        DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle<NDimSpatial,ALayout,BLayout,  DsLayout,ELayout,    F8,    F8,     F32,      F32,  Tuple<>,   F32, PassThrough, PassThrough,  OutElementOp,       ConvSpec, GemmMNKPadding,        1,   128,   128,   128,    32,   8,   8,   32,   32,    4,    2,     S<4, 32, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,         1,     S<4, 32, 1>,     S<1, 0, 2>,     S<1, 0, 2>,             2,              8,              8,         1,           1,           1,               S<1, 16, 1, 8>,               8,     F8,      F8>,
+        DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle<NDimSpatial,ALayout,BLayout,  DsLayout,ELayout,    F8,    F8,     F32,      F32,  Tuple<>,   F32, PassThrough, PassThrough,  OutElementOp,       ConvSpec, GemmMNKPadding,        1,   256,   128,   128,    32,   8,   8,   32,   32,    2,    2,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,         1,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,             2,              8,              8,         1,           1,           1,               S<1, 32, 1, 8>,               8,     F8,      F8>,
+        DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle<NDimSpatial,ALayout,BLayout,  DsLayout,ELayout,    F8,    F8,     F32,      F32,  Tuple<>,   F32, PassThrough, PassThrough,  OutElementOp,       ConvSpec, GemmMNKPadding,        1,   128,   128,    64,    32,   8,   8,   32,   32,    2,    2,     S<4, 32, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,         1,     S<4, 32, 1>,     S<1, 0, 2>,     S<1, 0, 2>,             2,              8,              8,         1,           1,           1,               S<1, 32, 1, 4>,               8,     F8,      F8>,
+        DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle<NDimSpatial,ALayout,BLayout,  DsLayout,ELayout,    F8,    F8,     F32,      F32,  Tuple<>,   F32, PassThrough, PassThrough,  OutElementOp,       ConvSpec, GemmMNKPadding,        1,   128,    64,   128,    32,   8,   8,   32,   32,    2,    2,     S<4, 32, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,         1,     S<4, 32, 1>,     S<1, 0, 2>,     S<1, 0, 2>,             2,              8,              8,         1,           1,           1,               S<1, 16, 1, 8>,               8,     F8,      F8>,
+        DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle<NDimSpatial,ALayout,BLayout,  DsLayout,ELayout,    F8,    F8,     F32,      F32,  Tuple<>,   F32, PassThrough, PassThrough,  OutElementOp,       ConvSpec, GemmMNKPadding,        1,    64,    64,    64,    32,   8,   8,   32,   32,    2,    2,     S<4, 16, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,         1,     S<4, 16, 1>,     S<1, 0, 2>,     S<1, 0, 2>,             2,              8,              8,         1,           1,           1,               S<1, 16, 1, 4>,               8,     F8,      F8>,
+        DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle<NDimSpatial,ALayout,BLayout,  DsLayout,ELayout,    F8,    F8,     F32,      F32,  Tuple<>,   F32, PassThrough, PassThrough,  OutElementOp,       ConvSpec, GemmMNKPadding,        1,   256,   128,    64,    32,   8,   8,   32,   32,    2,    1,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,         1,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,             2,              8,              8,         1,           1,           1,               S<1, 32, 1, 8>,               8,     F8,      F8>,
+        DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle<NDimSpatial,ALayout,BLayout,  DsLayout,ELayout,    F8,    F8,     F32,      F32,  Tuple<>,   F32, PassThrough, PassThrough,  OutElementOp,       ConvSpec, GemmMNKPadding,        1,   256,    64,   128,    32,   8,   8,   32,   32,    1,    2,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,         1,     S<4, 64, 1>,     S<1, 0, 2>,     S<1, 0, 2>,             2,              8,              8,         1,           1,           1,               S<1, 32, 1, 8>,               8,     F8,      F8>,
+        DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle<NDimSpatial,ALayout,BLayout,  DsLayout,ELayout,    F8,    F8,     F32,      F32,  Tuple<>,   F32, PassThrough, PassThrough,  OutElementOp,       ConvSpec, GemmMNKPadding,        1,   128,   128,    32,    32,   8,   8,   32,   32,    2,    1,     S<4, 32, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,         1,     S<4, 32, 1>,     S<1, 0, 2>,     S<1, 0, 2>,             2,              8,              8,         1,           1,           1,               S<1, 32, 1, 4>,               8,     F8,      F8>,
+        DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle<NDimSpatial,ALayout,BLayout,  DsLayout,ELayout,    F8,    F8,     F32,      F32,  Tuple<>,   F32, PassThrough, PassThrough,  OutElementOp,       ConvSpec, GemmMNKPadding,        1,   128,    32,   128,    32,   8,   8,   32,   32,    1,    2,     S<4, 32, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,         1,     S<4, 32, 1>,     S<1, 0, 2>,     S<1, 0, 2>,             2,              8,              8,         1,           1,           1,               S<1, 16, 1, 8>,               8,     F8,      F8>,
+        DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle<NDimSpatial,ALayout,BLayout,  DsLayout,ELayout,    F8,    F8,     F32,      F32,  Tuple<>,   F32, PassThrough, PassThrough,  OutElementOp,       ConvSpec, GemmMNKPadding,        1,    64,    64,    32,    32,   8,   8,   32,   32,    2,    1,     S<4, 16, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,         1,     S<4, 16, 1>,     S<1, 0, 2>,     S<1, 0, 2>,             2,              8,              8,         1,           1,           1,               S<1, 16, 1, 4>,               8,     F8,      F8>,
+        DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle<NDimSpatial,ALayout,BLayout,  DsLayout,ELayout,    F8,    F8,     F32,      F32,  Tuple<>,   F32, PassThrough, PassThrough,  OutElementOp,       ConvSpec, GemmMNKPadding,        1,    64,    32,    64,    32,   8,   8,   32,   32,    1,    2,     S<4, 16, 1>,     S<1, 0, 2>,     S<1, 0, 2>,              2,              8,              8,         1,     S<4, 16, 1>,     S<1, 0, 2>,     S<1, 0, 2>,             2,              8,              8,         1,           1,           1,               S<1, 16, 1, 4>,               8,     F8,      F8>
+#endif
+    // clang-format on
+    >;
+
 } // namespace instance
 } // namespace device
 } // namespace tensor_operation

library/include/ck/library/tensor_operation_instance/gpu/grouped_convolution_forward_convscale.hpp

@@ -8,9 +8,7 @@
 
 #include "ck/ck.hpp"
 #include "ck/tensor_operation/gpu/device/device_grouped_conv_fwd_multiple_abd.hpp"
-#include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
-#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"
-
+#include "ck/tensor_operation/gpu/element/combined_element_wise_operation.hpp"
 #include "ck/library/tensor_operation_instance/device_operation_instance_factory.hpp"
 
 namespace ck {
@@ -177,6 +175,88 @@ struct DeviceOperationInstanceFactory<ck::tensor_operation::device::DeviceGroupe
     }
 };
 
+using CombConvScale = ck::tensor_operation::element_wise::ScaleScalePass;
+
+#ifdef CK_ENABLE_FP8
+void add_device_grouped_conv3d_fwd_xdl_combconvscale_ndhwgc_gkzyxc_ndhwgk_f8_f8_f32_instances(
+    std::vector<std::unique_ptr<DeviceGroupedConvFwdMultipleABD<3,
+                                                                NDHWGC,
+                                                                GKZYXC,
+                                                                ck::Tuple<>,
+                                                                NDHWGK,
+                                                                F8,
+                                                                F8,
+                                                                ck::Tuple<>,
+                                                                F32,
+                                                                PassThrough,
+                                                                PassThrough,
+                                                                CombConvScale,
+                                                                F8,
+                                                                F8>>>& instances);
+#endif
+
+template <ck::index_t NumDimSpatial,
+          typename InLayout,
+          typename WeiLayout,
+          typename DLayouts,
+          typename OutLayout,
+          typename InDataType,
+          typename WeiDataType,
+          typename DDataTypes,
+          typename OutDataType,
+          typename AComputeType,
+          typename BComputeType>
+struct DeviceOperationInstanceFactory<
+    ck::tensor_operation::device::DeviceGroupedConvFwdMultipleABD<NumDimSpatial,
+                                                                  InLayout,
+                                                                  WeiLayout,
+                                                                  DLayouts,
+                                                                  OutLayout,
+                                                                  InDataType,
+                                                                  WeiDataType,
+                                                                  DDataTypes,
+                                                                  OutDataType,
+                                                                  PassThrough,
+                                                                  PassThrough,
+                                                                  CombConvScale,
+                                                                  AComputeType,
+                                                                  BComputeType>>
+{
+    using DeviceOp = DeviceGroupedConvFwdMultipleABD<NumDimSpatial,
+                                                     InLayout,
+                                                     WeiLayout,
+                                                     DLayouts,
+                                                     OutLayout,
+                                                     InDataType,
+                                                     WeiDataType,
+                                                     DDataTypes,
+                                                     OutDataType,
+                                                     PassThrough,
+                                                     PassThrough,
+                                                     CombConvScale,
+                                                     AComputeType,
+                                                     BComputeType>;
+
+    static auto GetInstances()
+    {
+        std::vector<std::unique_ptr<DeviceOp>> op_ptrs;
+        if constexpr(NumDimSpatial == 3 && is_same_v<InLayout, NDHWGC> &&
+                     is_same_v<WeiLayout, GKZYXC> && is_same_v<OutLayout, NDHWGK>)
+        {
+#ifdef CK_ENABLE_FP8
+            if constexpr(is_same_v<InDataType, f8_t> && is_same_v<WeiDataType, f8_t> &&
+                         is_same_v<OutDataType, F32> && is_same_v<AComputeType, f8_t> &&
+                         is_same_v<BComputeType, f8_t>)
+            {
+                add_device_grouped_conv3d_fwd_xdl_combconvscale_ndhwgc_gkzyxc_ndhwgk_f8_f8_f32_instances(
+                    op_ptrs);
+            }
+#endif
+        }
+        return op_ptrs;
+    }
+};
+
 } // namespace instance
 } // namespace device
 } // namespace tensor_operation

library/include/ck/library/tensor_operation_instance/gpu/grouped_convolution_forward_convscale_relu.hpp

@@ -8,7 +8,7 @@
 
 #include "ck/ck.hpp"
 #include "ck/tensor_operation/gpu/device/device_grouped_conv_fwd_multiple_abd.hpp"
-#include "ck/tensor_operation/gpu/element/unary_element_wise_operation.hpp"
+#include "ck/tensor_operation/gpu/element/combined_element_wise_operation.hpp"
 #include "ck/library/tensor_operation_instance/device_operation_instance_factory.hpp"
 
 namespace ck {
@@ -99,6 +99,88 @@ struct DeviceOperationInstanceFactory<
     }
 };
 
+using CombConvScaleRelu = ck::tensor_operation::element_wise::ScaleScaleRelu;
+
+#ifdef CK_ENABLE_FP8
+void add_device_grouped_conv3d_fwd_xdl_combconvscale_relu_ndhwgc_gkzyxc_ndhwgk_f8_f8_f32_instances(
+    std::vector<std::unique_ptr<DeviceGroupedConvFwdMultipleABD<3,
+                                                                NDHWGC,
+                                                                GKZYXC,
+                                                                ck::Tuple<>,
+                                                                NDHWGK,
+                                                                F8,
+                                                                F8,
+                                                                ck::Tuple<>,
+                                                                F32,
+                                                                PassThrough,
+                                                                PassThrough,
+                                                                CombConvScaleRelu,
+                                                                F8,
+                                                                F8>>>& instances);
+#endif
+
+template <ck::index_t NumDimSpatial,
+          typename InLayout,
+          typename WeiLayout,
+          typename DLayouts,
+          typename OutLayout,
+          typename InDataType,
+          typename WeiDataType,
+          typename DDataTypes,
+          typename OutDataType,
+          typename AComputeType,
+          typename BComputeType>
+struct DeviceOperationInstanceFactory<
+    ck::tensor_operation::device::DeviceGroupedConvFwdMultipleABD<NumDimSpatial,
+                                                                  InLayout,
+                                                                  WeiLayout,
+                                                                  DLayouts,
+                                                                  OutLayout,
+                                                                  InDataType,
+                                                                  WeiDataType,
+                                                                  DDataTypes,
+                                                                  OutDataType,
+                                                                  PassThrough,
+                                                                  PassThrough,
+                                                                  CombConvScaleRelu,
+                                                                  AComputeType,
+                                                                  BComputeType>>
+{
+    using DeviceOp = DeviceGroupedConvFwdMultipleABD<NumDimSpatial,
+                                                     InLayout,
+                                                     WeiLayout,
+                                                     DLayouts,
+                                                     OutLayout,
+                                                     InDataType,
+                                                     WeiDataType,
+                                                     DDataTypes,
+                                                     OutDataType,
+                                                     PassThrough,
+                                                     PassThrough,
+                                                     CombConvScaleRelu,
+                                                     AComputeType,
+                                                     BComputeType>;
+
+    static auto GetInstances()
+    {
+        std::vector<std::unique_ptr<DeviceOp>> op_ptrs;
+        if constexpr(NumDimSpatial == 3 && is_same_v<InLayout, NDHWGC> &&
+                     is_same_v<WeiLayout, GKZYXC> && is_same_v<OutLayout, NDHWGK>)
+        {
+#ifdef CK_ENABLE_FP8
+            if constexpr(is_same_v<InDataType, f8_t> && is_same_v<WeiDataType, f8_t> &&
+                         is_same_v<OutDataType, F32> && is_same_v<AComputeType, f8_t> &&
+                         is_same_v<BComputeType, f8_t>)
+            {
+                add_device_grouped_conv3d_fwd_xdl_combconvscale_relu_ndhwgc_gkzyxc_ndhwgk_f8_f8_f32_instances(
+                    op_ptrs);
+            }
+#endif
+        }
+        return op_ptrs;
+    }
+};
+
 } // namespace instance
 } // namespace device
 } // namespace tensor_operation

library/include/ck/library/tensor_operation_instance/gpu/permute_scale.hpp

@@ -70,6 +70,12 @@ void add_device_permute_scale_6d_f32_instances(
         DeviceElementwise<ck::Tuple<F32>, ck::Tuple<F32>, element_wise::Scale, 6>>>&);
 #endif
 
+#ifdef CK_ENABLE_FP8
+void add_device_permute_scale_6d_f32_f8_instances(
+    std::vector<std::unique_ptr<
+        DeviceElementwise<ck::Tuple<F32>, ck::Tuple<F8>, element_wise::Scale, 6>>>&);
+#endif
+
 template <typename InDataTypeTuple,
           typename OutDataTypeTuple,
           typename ElementwiseOperation,
@@ -184,6 +190,13 @@ struct DeviceOperationInstanceFactory<
             {
                 add_device_permute_scale_6d_f16_instances(op_ptrs);
             }
+#endif
+#ifdef CK_ENABLE_FP8
+            if constexpr(is_same_v<InDataTypeTuple, ck::Tuple<F32>> &&
+                         is_same_v<OutDataTypeTuple, ck::Tuple<F8>>)
+            {
+                add_device_permute_scale_6d_f32_f8_instances(op_ptrs);
+            }
 #endif
         }
         return op_ptrs;

library/include/ck/library/tensor_operation_instance/gpu/permute_scale/device_permute_scale_instances.hpp

@@ -10,6 +10,7 @@ namespace tensor_operation {
 namespace device {
 namespace instance {
 
+using F8  = ck::f8_t;
 using F16 = ck::half_t;
 using F32 = float;
 
@@ -46,7 +47,7 @@ using device_permute_scale_f16_instances =
 
 #if 0
         // Disabled instances to improve compilation time
-        // They listed here to show other possible combinations of parameters 
+        // They listed here to show other possible combinations of parameters
          DeviceElementwiseImpl<ck::Tuple<F16>, ck::Tuple<F16>, ElementwiseOp,  NDims, 256, 256, 256, 16, 16, ck::Sequence<1, 0>, ck::Sequence<16>, ck::Sequence<16>>,
          DeviceElementwiseImpl<ck::Tuple<F16>, ck::Tuple<F16>, ElementwiseOp,  NDims, 128, 256, 128, 16, 16, ck::Sequence<1, 0>, ck::Sequence<16>, ck::Sequence<16>>,
          DeviceElementwiseImpl<ck::Tuple<F16>, ck::Tuple<F16>, ElementwiseOp,  NDims, 128, 128, 256, 16, 16, ck::Sequence<1, 0>, ck::Sequence<16>, ck::Sequence<16>>,
@@ -57,7 +58,7 @@ using device_permute_scale_f16_instances =
          DeviceElementwiseImpl<ck::Tuple<F16>, ck::Tuple<F16>, ElementwiseOp,  NDims, 64,  128, 128, 16, 16, ck::Sequence<1, 0>, ck::Sequence<16>, ck::Sequence<16>>,
          DeviceElementwiseImpl<ck::Tuple<F16>, ck::Tuple<F16>, ElementwiseOp,  NDims, 32,  128,  64, 16, 16, ck::Sequence<1, 0>, ck::Sequence<16>, ck::Sequence<16>>,
          DeviceElementwiseImpl<ck::Tuple<F16>, ck::Tuple<F16>, ElementwiseOp,  NDims, 32,   64, 128, 16, 16, ck::Sequence<1, 0>, ck::Sequence<16>, ck::Sequence<16>>,
-   
+
          DeviceElementwiseImpl<ck::Tuple<F16>, ck::Tuple<F16>, ElementwiseOp,  NDims, 256,  64, 128, 4, 8, ck::Sequence<1, 0>, ck::Sequence<4>, ck::Sequence<4>>,
          DeviceElementwiseImpl<ck::Tuple<F16>, ck::Tuple<F16>, ElementwiseOp,  NDims, 256, 128,  64, 4, 8, ck::Sequence<1, 0>, ck::Sequence<4>, ck::Sequence<4>>,
          DeviceElementwiseImpl<ck::Tuple<F16>, ck::Tuple<F16>, ElementwiseOp,  NDims, 128,  64,  64, 4, 8, ck::Sequence<1, 0>, ck::Sequence<4>, ck::Sequence<4>>,
@@ -97,7 +98,7 @@ using device_permute_scale_f16_instances =
          DeviceElementwiseImpl<ck::Tuple<F16>, ck::Tuple<F16>, ElementwiseOp,  NDims, 64,   64,  16, 4, 4, ck::Sequence<1, 0>, ck::Sequence<1>, ck::Sequence<1>>,
          DeviceElementwiseImpl<ck::Tuple<F16>, ck::Tuple<F16>, ElementwiseOp,  NDims, 32,   32,  16, 4, 4, ck::Sequence<1, 0>, ck::Sequence<1>, ck::Sequence<1>>,
          DeviceElementwiseImpl<ck::Tuple<F16>, ck::Tuple<F16>, ElementwiseOp,  NDims, 32,   16,  32, 4, 4, ck::Sequence<1, 0>, ck::Sequence<1>, ck::Sequence<1>>
-   
+
     >;
 
 template <index_t NDims,
@@ -131,7 +132,7 @@ using device_permute_scale_f32_instances = std::tuple<
 
 #if 0
         // Disabled instances to improve compilation time
-        // They listed here to show other possible combinations of parameters 
+        // They listed here to show other possible combinations of parameters
          DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F32>, ElementwiseOp,  NDims, 256, 256, 256, 16, 16, ck::Sequence<1, 0>, ck::Sequence<16>, ck::Sequence<16>>,
          DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F32>, ElementwiseOp,  NDims, 128, 256, 128, 16, 16, ck::Sequence<1, 0>, ck::Sequence<16>, ck::Sequence<16>>,
          DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F32>, ElementwiseOp,  NDims, 128, 128, 256, 16, 16, ck::Sequence<1, 0>, ck::Sequence<16>, ck::Sequence<16>>,
@@ -142,7 +143,7 @@ using device_permute_scale_f32_instances = std::tuple<
          DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F32>, ElementwiseOp,  NDims, 64,  128, 128, 16, 16, ck::Sequence<1, 0>, ck::Sequence<16>, ck::Sequence<16>>,
          DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F32>, ElementwiseOp,  NDims, 32,  128,  64, 16, 16, ck::Sequence<1, 0>, ck::Sequence<16>, ck::Sequence<16>>,
          DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F32>, ElementwiseOp,  NDims, 32,   64, 128, 16, 16, ck::Sequence<1, 0>, ck::Sequence<16>, ck::Sequence<16>>,
-   
+
          DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F32>, ElementwiseOp,  NDims, 256,  64, 128, 4, 8, ck::Sequence<1, 0>, ck::Sequence<4>, ck::Sequence<4>>,
          DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F32>, ElementwiseOp,  NDims, 256, 128,  64, 4, 8, ck::Sequence<1, 0>, ck::Sequence<4>, ck::Sequence<4>>,
          DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F32>, ElementwiseOp,  NDims, 128,  64,  64, 4, 8, ck::Sequence<1, 0>, ck::Sequence<4>, ck::Sequence<4>>,
@@ -168,7 +169,7 @@ using device_permute_scale_f32_instances = std::tuple<
          DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F32>, ElementwiseOp,  NDims, 64,  128,  16, 8, 4, ck::Sequence<1, 0>, ck::Sequence<4>, ck::Sequence<4>>,
          DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F32>, ElementwiseOp,  NDims, 32,   64,  16, 8, 4, ck::Sequence<1, 0>, ck::Sequence<4>, ck::Sequence<4>>,
          DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F32>, ElementwiseOp,  NDims, 32,   32,  32, 8, 4, ck::Sequence<1, 0>, ck::Sequence<4>, ck::Sequence<4>>,
-#endif         
+#endif
 
          DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F32>, ElementwiseOp,  NDims, 256,  64,  64, 4, 4, ck::Sequence<1, 0>, ck::Sequence<1>, ck::Sequence<1>>,
          DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F32>, ElementwiseOp,  NDims, 256, 128,  32, 4, 4, ck::Sequence<1, 0>, ck::Sequence<1>, ck::Sequence<1>>,
@@ -183,6 +184,51 @@ using device_permute_scale_f32_instances = std::tuple<
          DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F32>, ElementwiseOp,  NDims, 32,   32,  16, 4, 4, ck::Sequence<1, 0>, ck::Sequence<1>, ck::Sequence<1>>,
          DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F32>, ElementwiseOp,  NDims, 32,   16,  32, 4, 4, ck::Sequence<1, 0>, ck::Sequence<1>, ck::Sequence<1>>
     >;
+
+#ifdef CK_ENABLE_FP8
+template <index_t NDims,
+          typename ElementwiseOp>
+using device_permute_scale_f32_f8_instances = std::tuple<
+         DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F8>, ElementwiseOp,  NDims, 256,  64,  64, 4, 4, ck::Sequence<1, 0>, ck::Sequence<4>, ck::Sequence<4>>,
+         DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F8>, ElementwiseOp,  NDims, 256, 128,  32, 4, 4, ck::Sequence<1, 0>, ck::Sequence<4>, ck::Sequence<4>>,
+         DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F8>, ElementwiseOp,  NDims, 256,  32, 128, 4, 4,  ck::Sequence<1, 0>, ck::Sequence<4>, ck::Sequence<4>>,
+         DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F8>, ElementwiseOp,  NDims, 128,  64,  32, 4, 4, ck::Sequence<1, 0>, ck::Sequence<4>, ck::Sequence<4>>,
+         DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F8>, ElementwiseOp,  NDims, 128,  32,  64, 4, 4, ck::Sequence<1, 0>, ck::Sequence<4>, ck::Sequence<4>>,
+         DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F8>, ElementwiseOp,  NDims, 128,  16, 128, 4, 4, ck::Sequence<1, 0>, ck::Sequence<4>, ck::Sequence<4>>,
+         DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F8>, ElementwiseOp,  NDims, 128, 128,  16, 4, 4, ck::Sequence<1, 0>, ck::Sequence<4>, ck::Sequence<4>>,
+         DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F8>, ElementwiseOp,  NDims, 64,   32,  32, 4, 4, ck::Sequence<1, 0>, ck::Sequence<4>, ck::Sequence<4>>,
+         DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F8>, ElementwiseOp,  NDims, 64,   16,  64, 4, 4, ck::Sequence<1, 0>, ck::Sequence<4>, ck::Sequence<4>>,
+         DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F8>, ElementwiseOp,  NDims, 64,   64,  16, 4, 4, ck::Sequence<1, 0>, ck::Sequence<4>, ck::Sequence<4>>,
+         DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F8>, ElementwiseOp,  NDims, 32,   32,  16, 4, 4, ck::Sequence<1, 0>, ck::Sequence<4>, ck::Sequence<4>>,
+         DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F8>, ElementwiseOp,  NDims, 32,   16,  32, 4, 4, ck::Sequence<1, 0>, ck::Sequence<4>, ck::Sequence<4>>,
+
+         DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F8>, ElementwiseOp,  NDims, 256, 128, 128, 8, 8, ck::Sequence<1, 0>, ck::Sequence<8>, ck::Sequence<8>>,
+         DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F8>, ElementwiseOp,  NDims, 256, 256,  64, 8, 8, ck::Sequence<1, 0>, ck::Sequence<8>, ck::Sequence<8>>,
+         DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F8>, ElementwiseOp,  NDims, 256,  64, 256, 8, 8, ck::Sequence<1, 0>, ck::Sequence<8>, ck::Sequence<8>>,
+         DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F8>, ElementwiseOp,  NDims, 128, 128,  64, 8, 8, ck::Sequence<1, 0>, ck::Sequence<8>, ck::Sequence<8>>,
+         DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F8>, ElementwiseOp,  NDims, 128,  64, 128, 8, 8, ck::Sequence<1, 0>, ck::Sequence<8>, ck::Sequence<8>>,
+         DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F8>, ElementwiseOp,  NDims, 128,  32, 256, 8, 8, ck::Sequence<1, 0>, ck::Sequence<8>, ck::Sequence<8>>,
+         DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F8>, ElementwiseOp,  NDims, 128, 256,  32, 8, 8, ck::Sequence<1, 0>, ck::Sequence<8>, ck::Sequence<8>>,
+         DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F8>, ElementwiseOp,  NDims, 64,   64,  64, 8, 8, ck::Sequence<1, 0>, ck::Sequence<8>, ck::Sequence<8>>,
+         DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F8>, ElementwiseOp,  NDims, 64,   32, 128, 8, 8, ck::Sequence<1, 0>, ck::Sequence<8>, ck::Sequence<8>>,
+         DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F8>, ElementwiseOp,  NDims, 64,  128,  32, 8, 8, ck::Sequence<1, 0>, ck::Sequence<8>, ck::Sequence<8>>,
+         DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F8>, ElementwiseOp,  NDims, 32,   64,  32, 8, 8, ck::Sequence<1, 0>, ck::Sequence<8>, ck::Sequence<8>>,
+         DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F8>, ElementwiseOp,  NDims, 32,   32,  64, 8, 8, ck::Sequence<1, 0>, ck::Sequence<8>, ck::Sequence<8>>,
+
+         DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F8>, ElementwiseOp,  NDims, 256,  64,  64, 4, 4, ck::Sequence<1, 0>, ck::Sequence<1>, ck::Sequence<1>>,
+         DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F8>, ElementwiseOp,  NDims, 256, 128,  32, 4, 4, ck::Sequence<1, 0>, ck::Sequence<1>, ck::Sequence<1>>,
+         DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F8>, ElementwiseOp,  NDims, 256,  32, 128, 4, 4,  ck::Sequence<1, 0>, ck::Sequence<1>, ck::Sequence<1>>,
+         DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F8>, ElementwiseOp,  NDims, 128,  64,  32, 4, 4, ck::Sequence<1, 0>, ck::Sequence<1>, ck::Sequence<1>>,
+         DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F8>, ElementwiseOp,  NDims, 128,  32,  64, 4, 4, ck::Sequence<1, 0>, ck::Sequence<1>, ck::Sequence<1>>,
+         DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F8>, ElementwiseOp,  NDims, 128,  16, 128, 4, 4, ck::Sequence<1, 0>, ck::Sequence<1>, ck::Sequence<1>>,
+         DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F8>, ElementwiseOp,  NDims, 128, 128,  16, 4, 4, ck::Sequence<1, 0>, ck::Sequence<1>, ck::Sequence<1>>,
+         DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F8>, ElementwiseOp,  NDims, 64,   32,  32, 4, 4, ck::Sequence<1, 0>, ck::Sequence<1>, ck::Sequence<1>>,
+         DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F8>, ElementwiseOp,  NDims, 64,   16,  64, 4, 4, ck::Sequence<1, 0>, ck::Sequence<1>, ck::Sequence<1>>,
+         DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F8>, ElementwiseOp,  NDims, 64,   64,  16, 4, 4, ck::Sequence<1, 0>, ck::Sequence<1>, ck::Sequence<1>>,
+         DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F8>, ElementwiseOp,  NDims, 32,   32,  16, 4, 4, ck::Sequence<1, 0>, ck::Sequence<1>, ck::Sequence<1>>,
+         DeviceElementwiseImpl<ck::Tuple<F32>, ck::Tuple<F8>, ElementwiseOp,  NDims, 32,   16,  32, 4, 4, ck::Sequence<1, 0>, ck::Sequence<1>, ck::Sequence<1>>
+    >;
+#endif
 // clang-format on
 
 } // namespace instance

library/include/ck/library/tensor_operation_instance/gpu/reduce/device_reduce_instance_blockwise_f32_f32_f32_amax.hpp

@@ -14,15 +14,24 @@ namespace device {
 namespace instance {
 
 // clang-format off
-// InDataType | AccDataType | OutDataType | Rank | NumReduceDim | ReduceOperation | InElementwiseOp | AccElementwiseOp | PropagateNan | UseIndex 
-extern template void add_device_reduce_instance_blockwise<F32, F32, F32, 4, 3, ReduceAMax, UnaryAbs, PassThrough, false, false>(std::vector<DeviceReducePtr<F32, F32, F32, 4, 3, ReduceAMax, UnaryAbs, PassThrough, false, false>>&);
-extern template void add_device_reduce_instance_blockwise<F32, F32, F32, 4, 4, ReduceAMax, UnaryAbs, PassThrough, false, false>(std::vector<DeviceReducePtr<F32, F32, F32, 4, 4, ReduceAMax, UnaryAbs, PassThrough, false, false>>&);
-extern template void add_device_reduce_instance_blockwise<F32, F32, F32, 4, 1, ReduceAMax, UnaryAbs, PassThrough, false, false>(std::vector<DeviceReducePtr<F32, F32, F32, 4, 1, ReduceAMax, UnaryAbs, PassThrough, false, false>>&);
-extern template void add_device_reduce_instance_blockwise<F32, F32, F32, 2, 1, ReduceAMax, UnaryAbs, PassThrough, false, false>(std::vector<DeviceReducePtr<F32, F32, F32, 2, 1, ReduceAMax, UnaryAbs, PassThrough, false, false>>&);
-extern template void add_device_reduce_instance_blockwise<F32, F32, F32, 4, 3, ReduceAMax, UnaryAbs, PassThrough, false, true>(std::vector<DeviceReducePtr<F32, F32, F32, 4, 3, ReduceAMax, UnaryAbs, PassThrough, false, true>>&);
-extern template void add_device_reduce_instance_blockwise<F32, F32, F32, 4, 4, ReduceAMax, UnaryAbs, PassThrough, false, true>(std::vector<DeviceReducePtr<F32, F32, F32, 4, 4, ReduceAMax, UnaryAbs, PassThrough, false, true>>&);
-extern template void add_device_reduce_instance_blockwise<F32, F32, F32, 4, 1, ReduceAMax, UnaryAbs, PassThrough, false, true>(std::vector<DeviceReducePtr<F32, F32, F32, 4, 1, ReduceAMax, UnaryAbs, PassThrough, false, true>>&);
-extern template void add_device_reduce_instance_blockwise<F32, F32, F32, 2, 1, ReduceAMax, UnaryAbs, PassThrough, false, true>(std::vector<DeviceReducePtr<F32, F32, F32, 2, 1, ReduceAMax, UnaryAbs, PassThrough, false, true>>&);
+//                                                        InDataType | AccDataType | OutDataType | Rank | NumReduceDim | ReduceOperation | InElementwiseOp | AccElementwiseOp | PropagateNan | UseIndex 
+extern template void add_device_reduce_instance_blockwise<        F32,          F32,          F32,     4,             3,       ReduceAMax,         UnaryAbs,       PassThrough,        false,     false>(std::vector<DeviceReducePtr<F32, F32, F32, 4, 3, ReduceAMax, UnaryAbs, PassThrough, false, false>>&);
+extern template void add_device_reduce_instance_blockwise<        F32,          F32,          F32,     4,             4,       ReduceAMax,         UnaryAbs,       PassThrough,        false,     false>(std::vector<DeviceReducePtr<F32, F32, F32, 4, 4, ReduceAMax, UnaryAbs, PassThrough, false, false>>&);
+extern template void add_device_reduce_instance_blockwise<        F32,          F32,          F32,     4,             1,       ReduceAMax,         UnaryAbs,       PassThrough,        false,     false>(std::vector<DeviceReducePtr<F32, F32, F32, 4, 1, ReduceAMax, UnaryAbs, PassThrough, false, false>>&);
+extern template void add_device_reduce_instance_blockwise<        F32,          F32,          F32,     2,             1,       ReduceAMax,         UnaryAbs,       PassThrough,        false,     false>(std::vector<DeviceReducePtr<F32, F32, F32, 2, 1, ReduceAMax, UnaryAbs, PassThrough, false, false>>&);
+extern template void add_device_reduce_instance_blockwise<        F32,          F32,          F32,     4,             3,       ReduceAMax,         UnaryAbs,       PassThrough,        false,      true>(std::vector<DeviceReducePtr<F32, F32, F32, 4, 3, ReduceAMax, UnaryAbs, PassThrough, false, true>>&);
+extern template void add_device_reduce_instance_blockwise<        F32,          F32,          F32,     4,             4,       ReduceAMax,         UnaryAbs,       PassThrough,        false,      true>(std::vector<DeviceReducePtr<F32, F32, F32, 4, 4, ReduceAMax, UnaryAbs, PassThrough, false, true>>&);
+extern template void add_device_reduce_instance_blockwise<        F32,          F32,          F32,     4,             1,       ReduceAMax,         UnaryAbs,       PassThrough,        false,      true>(std::vector<DeviceReducePtr<F32, F32, F32, 4, 1, ReduceAMax, UnaryAbs, PassThrough, false, true>>&);
+extern template void add_device_reduce_instance_blockwise<        F32,          F32,          F32,     2,             1,       ReduceAMax,         UnaryAbs,       PassThrough,        false,      true>(std::vector<DeviceReducePtr<F32, F32, F32, 2, 1, ReduceAMax, UnaryAbs, PassThrough, false, true>>&);
+extern template void add_device_reduce_instance_blockwise<        F32,          F32,          F32,     6,             6,       ReduceAMax,         UnaryAbs,       PassThrough,         true,     false>(std::vector<DeviceReducePtr<F32, F32, F32, 6, 6, ReduceAMax, UnaryAbs, PassThrough, true, false>>&);
+extern template void add_device_reduce_instance_blockwise<        F32,          F32,          F32,     5,             5,       ReduceAMax,         UnaryAbs,       PassThrough,         true,     false>(std::vector<DeviceReducePtr<F32, F32, F32, 5, 5, ReduceAMax, UnaryAbs, PassThrough, true, false>>&);
+extern template void add_device_reduce_instance_blockwise<        F32,          F32,          F32,     4,             4,       ReduceAMax,         UnaryAbs,       PassThrough,         true,     false>(std::vector<DeviceReducePtr<F32, F32, F32, 4, 4, ReduceAMax, UnaryAbs, PassThrough, true, false>>&);
+extern template void add_device_reduce_instance_blockwise<        F32,          F32,          F32,     6,             3,       ReduceAMax,         UnaryAbs,       PassThrough,         true,     false>(std::vector<DeviceReducePtr<F32, F32, F32, 6, 3, ReduceAMax, UnaryAbs, PassThrough, true, false>>&);
+extern template void add_device_reduce_instance_blockwise<        F32,          F32,          F32,     5,             3,       ReduceAMax,         UnaryAbs,       PassThrough,         true,     false>(std::vector<DeviceReducePtr<F32, F32, F32, 5, 3, ReduceAMax, UnaryAbs, PassThrough, true, false>>&);
+extern template void add_device_reduce_instance_blockwise<        F32,          F32,          F32,     4,             3,       ReduceAMax,         UnaryAbs,       PassThrough,         true,     false>(std::vector<DeviceReducePtr<F32, F32, F32, 4, 3, ReduceAMax, UnaryAbs, PassThrough, true, false>>&);
+extern template void add_device_reduce_instance_blockwise<        F32,          F32,          F32,     3,             3,       ReduceAMax,      PassThrough,       PassThrough,         true,     false>(std::vector<DeviceReducePtr<F32, F32, F32, 3, 3, ReduceAMax, PassThrough, PassThrough, true, false>>&);
+extern template void add_device_reduce_instance_blockwise<        F32,          F32,          F32,     2,             2,       ReduceAMax,      PassThrough,       PassThrough,         true,     false>(std::vector<DeviceReducePtr<F32, F32, F32, 2, 2, ReduceAMax, PassThrough, PassThrough, true, false>>&);
+extern template void add_device_reduce_instance_blockwise<        F32,          F32,          F32,     1,             1,       ReduceAMax,      PassThrough,       PassThrough,         true,     false>(std::vector<DeviceReducePtr<F32, F32, F32, 1, 1, ReduceAMax, PassThrough, PassThrough, true, false>>&);
 // clang-format on
 
 } // namespace instance

library/include/ck/library/utility/check_err.hpp

@@ -1,5 +1,5 @@
 // SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
 
 #pragma once
 
@@ -146,7 +146,7 @@ check_err(const Range& out,
     bool res{true};
     int err_count  = 0;
     double err     = 0;
-    double max_err = std::numeric_limits<ranges::range_value_t<Range>>::min();
+    double max_err = NumericLimits<ranges::range_value_t<Range>>::Min();
     for(std::size_t i = 0; i < ref.size(); ++i)
     {
         const double o = type_convert<float>(*std::next(std::begin(out), i));
@@ -178,7 +178,9 @@ check_err(const Range& out,
 template <typename Range, typename RefRange>
 std::enable_if_t<(std::is_same_v<ranges::range_value_t<Range>, ranges::range_value_t<RefRange>> &&
                   std::is_integral_v<ranges::range_value_t<Range>> &&
-                  !std::is_same_v<ranges::range_value_t<Range>, bhalf_t>)
+                  !std::is_same_v<ranges::range_value_t<Range>, bhalf_t> &&
+                  !std::is_same_v<ranges::range_value_t<Range>, f8_t> &&
+                  !std::is_same_v<ranges::range_value_t<Range>, bf8_t>)
 #ifdef CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4
                      || std::is_same_v<ranges::range_value_t<Range>, int4_t>
 #endif
@@ -270,7 +272,8 @@ check_err(const Range& out,
     }
     if(!res)
     {
-        std::cerr << std::setw(12) << std::setprecision(7) << "max err: " << max_err << std::endl;
+        std::cerr << std::setw(12) << std::setprecision(7) << "max err: " << max_err
+                  << " number of errors: " << err_count << std::endl;
     }
     return res;
 }

library/src/tensor_operation_instance/gpu/gemm/CMakeLists.txt

@@ -111,6 +111,7 @@ list(APPEND GEMM_INSTANCES
     device_gemm_xdl_c_shuffle_fp8_fp8_fp8_km_kn_mn_instance.cpp
     device_gemm_xdl_c_shuffle_fp8_fp8_fp8_km_nk_mn_instance.cpp)
 
+
 list(APPEND GEMM_INSTANCES
     device_gemm_wmma_f16_f16_f16_mk_kn_mn_instance.cpp
     device_gemm_wmma_f16_f16_f16_mk_nk_mn_instance.cpp

library/src/tensor_operation_instance/gpu/gemm_ab_scale/CMakeLists.txt

@@ -11,4 +11,9 @@ list(APPEND GEMM_AB_SCALE_INSTANCES
         device_gemm_ab_scale_xdl_f8_f8_bf16/device_gemm_ab_scale_xdl_f8_f8_bf16_mk_nk_mn_128_128_128_mem_v1_mnkpadding_instance.cpp
         )
 
+set_source_files_properties(device_gemm_ab_scale_xdl_f8_f8_bf16/device_gemm_ab_scale_xdl_f8_f8_bf16_mk_nk_mn_128_128_128_comp_default_instance.cpp PROPERTIES COMPILE_OPTIONS ";-mllvm;-greedy-reverse-local-assignment=1")
+set_source_files_properties(device_gemm_ab_scale_xdl_f8_f8_bf16/device_gemm_ab_scale_xdl_f8_f8_bf16_mk_nk_mn_128_128_128_comp_kpadding_instance.cpp PROPERTIES COMPILE_OPTIONS ";-mllvm;-greedy-reverse-local-assignment=1")
+set_source_files_properties(device_gemm_ab_scale_xdl_f8_f8_bf16/device_gemm_ab_scale_xdl_f8_f8_bf16_mk_nk_mn_128_128_128_comp_mnpadding_instance.cpp PROPERTIES COMPILE_OPTIONS ";-mllvm;-greedy-reverse-local-assignment=1")
+set_source_files_properties(device_gemm_ab_scale_xdl_f8_f8_bf16/device_gemm_ab_scale_xdl_f8_f8_bf16_mk_nk_mn_128_128_128_comp_mnkpadding_instance.cpp PROPERTIES COMPILE_OPTIONS ";-mllvm;-greedy-reverse-local-assignment=1")
+
 add_instance_library(device_gemm_ab_scale_instance ${GEMM_AB_SCALE_INSTANCES})

library/src/tensor_operation_instance/gpu/gemm_multiply_multiply/CMakeLists.txt

@@ -4,14 +4,13 @@ set(GEMM_MULTIPLY_MULTIPLY_INSTANCES)
 list(APPEND GEMM_MULTIPLY_MULTIPLY_INSTANCES 
         device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_comp_default_instance.cpp
         device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_comp_kpadding_instance.cpp
-        device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_comp_mnpadding_instance.cpp
-        device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_comp_mnkpadding_instance.cpp
         device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v1_default_instance.cpp
         device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v1_kpadding_instance.cpp
-        device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v1_mnkpadding_instance.cpp
         device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v2_default_instance.cpp
         device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v2_kpadding_instance.cpp
-        device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_mem_v2_mnkpadding_instance.cpp
         )
 
+set_source_files_properties(device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_comp_default_instance.cpp PROPERTIES COMPILE_OPTIONS ";-mllvm;-greedy-reverse-local-assignment=1")
+set_source_files_properties(device_gemm_multiply_multiply_xdl_f8_f8_bf16/device_gemm_multiply_multiply_xdl_f8_f8_bf16_mk_nk_mn_comp_kpadding_instance.cpp PROPERTIES COMPILE_OPTIONS ";-mllvm;-greedy-reverse-local-assignment=1")
+
 add_instance_library(device_gemm_multiply_multiply_instance ${GEMM_MULTIPLY_MULTIPLY_INSTANCES})