init implementation

example/32_batched_gemm_scale_softmax_gemm/CMakeLists.txt

@@ -12,6 +12,7 @@ add_example_executable(example_grouped_gemm_lower_triangle_scale_softmax_gemm_pe
 add_example_executable(example_batched_multihead_attention_backward_fp16 batched_multihead_attention_backward_fp16.cpp)
 add_example_executable(example_batched_multihead_attention_backward_pt1_fp16 batched_multihead_attention_backward_pt1_fp16.cpp)
 add_example_executable(example_grouped_multihead_attention_backward_fp16 grouped_multihead_attention_backward_fp16.cpp)
+add_example_executable(example_batched_multihead_attention_backward_fp16_dropout batched_multihead_attention_backward_fp16_dropout.cpp)
 
 add_custom_target(example_gemm_scale_softmax_gemm)
 add_dependencies(example_gemm_scale_softmax_gemm example_batched_gemm_scale_softmax_gemm_xdl_fp16)

example/32_batched_gemm_scale_softmax_gemm/batched_gemm_scale_softmax_gemm_permute_train_xdl_bf16.cpp

@@ -32,7 +32,7 @@ template <ck::index_t... Is>
 using S = ck::Sequence<Is...>;
 
 using BF16 = ck::bhalf_t;
-using F32 = float;
+using F32  = float;
 
 using PassThrough = ck::tensor_operation::element_wise::PassThrough;
 

example/32_batched_gemm_scale_softmax_gemm/batched_multihead_attention_backward_fp16_dropout.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2022, Advanced Micro Devices, Inc. All rights reserved.
+/*
+Backprop for Gemm + Softmax + Gemm fused operation, where forward prop is defined as:
+
+  Y_g_m_o = Softmax(alpha * Q_g_m_k * K_g_k_n) * V_g_n_o
+
+Computation graph:
+
+          K^T                   V
+          |                     |
+          |                     |
+    Q --- * ----- Softmax ----- * --> Y
+              S             P
+
+Kernel inputs:
+
+    Q, K, V, Y, dY, per-row softmax stats (LSE)
+
+Kernel outputs:
+
+    dQ, dK, dV
+
+*/
+
+#define PRINT_HOST 0
+#define USING_MASK 1
+
+#include <iostream>
+#include <numeric>
+#include <initializer_list>
+#include <cstdlib>
+#include <fstream>
+
+#include "ck/ck.hpp"
+#include "ck/tensor_operation/gpu/device/gemm_specialization.hpp"
+#include "ck/tensor_operation/gpu/device/tensor_specialization.hpp"
+#include "ck/tensor_operation/gpu/device/impl/device_batched_multihead_attention_backward_train_xdl_cshuffle.hpp"
+#include "ck/tensor_operation/gpu/element/element_wise_operation.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/reference_tensor_operation/cpu/reference_batched_gemm.hpp"
+#include "ck/library/reference_tensor_operation/cpu/reference_softmax.hpp"
+#include "ck/library/reference_tensor_operation/cpu/reference_dropout.hpp"
+
+template <ck::index_t... Is>
+using S = ck::Sequence<Is...>;
+
+using F16 = ck::half_t;
+using F32 = float;
+using U16 = unsigned short;
+
+using PassThrough = ck::tensor_operation::element_wise::PassThrough;
+using Scale       = ck::tensor_operation::element_wise::Scale;
+
+using QKVElementOp = PassThrough;
+using YElementOp   = PassThrough;
+using VElementOp   = Scale;
+
+using DataType         = F16;
+using AccDataType      = F32;
+using ShuffleDataType  = F32;
+using LSEDataType      = F32;
+using ZDataType        = U16;
+using Acc0BiasDataType = ck::Tuple<>;
+using Acc1BiasDataType = ck::Tuple<>;
+
+static constexpr ck::index_t NumDimG = 2;
+static constexpr ck::index_t NumDimM = 1;
+static constexpr ck::index_t NumDimN = 1;
+static constexpr ck::index_t NumDimK = 1;
+static constexpr ck::index_t NumDimO = 1;
+
+static constexpr auto GemmSpec = ck::tensor_operation::device::GemmSpecialization::MNKOPadding;
+#if USING_MASK
+static constexpr auto MaskingSpec =
+    ck::tensor_operation::device::MaskingSpecialization::MaskOutUpperTriangle;
+#else
+static constexpr auto MaskingSpec =
+    ck::tensor_operation::device::MaskingSpecialization::MaskDisabled;
+#endif
+
+static constexpr auto TensorSpecQ = ck::tensor_operation::device::TensorSpecialization::Default;
+static constexpr auto TensorSpecK = ck::tensor_operation::device::TensorSpecialization::Default;
+static constexpr auto TensorSpecV = ck::tensor_operation::device::TensorSpecialization::Default;
+static constexpr auto TensorSpecY = ck::tensor_operation::device::TensorSpecialization::Default;
+
+using DeviceGemmInstance =
+    ck::tensor_operation::device::DeviceBatchedMultiheadAttentionBackward_Train_Xdl_CShuffle<
+        NumDimG,
+        NumDimM,
+        NumDimN,
+        NumDimK,
+        NumDimO,
+        DataType,
+        ZDataType,
+        LSEDataType,
+        Acc0BiasDataType,
+        Acc1BiasDataType,
+        AccDataType,
+        ShuffleDataType,
+        QKVElementOp,
+        QKVElementOp,
+        Scale,
+        QKVElementOp,
+        YElementOp,
+        GemmSpec,
+        TensorSpecQ,
+        TensorSpecK,
+        TensorSpecV,
+        TensorSpecY,
+        1,
+        256,
+        128,         // MPerBlock
+        128,         // NPerBlock
+        32,          // KPerBlock
+        128,         // Gemm1NPerBlock
+        64,          // Gemm1KPerBlock
+        8,           // AK1
+        8,           // BK1
+        2,           // B1K1
+        32,          // MPerXDL
+        32,          // NPerXDL
+        1,           // MXdlPerWave
+        4,           // NXdlPerWave
+        4,           // Gemm1NXdlPerWave
+        S<4, 64, 1>, // ABlockTransfer
+        S<1, 0, 2>,
+        S<1, 0, 2>,
+        2,
+        8,
+        8,
+        true,
+        S<4, 64, 1>, // BBlockTransfer
+        S<1, 0, 2>,
+        S<1, 0, 2>,
+        2,
+        8,
+        8,
+        true,
+        S<8, 32, 1>, // B1BlockTransfer
+        S<0, 2, 1>,
+        S<0, 2, 1>,
+        1,
+        4,
+        2,
+        false,
+        1,              // CShuffleMXdlPerWavePerShuffle
+        4,              // CShuffleNXdlPerWavePerShuffle
+        S<1, 32, 1, 8>, // CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock
+        8,              // CShuffleBlockTransferScalarPerVector_NPerBlock
+        MaskingSpec>;   // MaskingSpecialization
+
+// Ref Gemm0: S = alpha * Q * K^T
+// fp16 in, fp32 out
+using ReferenceGemm0Instance = ck::tensor_operation::host::ReferenceBatchedGemm<DataType,
+                                                                                DataType,
+                                                                                AccDataType,
+                                                                                AccDataType,
+                                                                                PassThrough,
+                                                                                PassThrough,
+                                                                                Scale>;
+
+// Ref Softmax: P = Softmax(S)
+// fp32 in, fp16 out
+using ReferenceSoftmaxInstance =
+    ck::tensor_operation::host::ReferenceSoftmax<AccDataType, DataType, AccDataType>;
+
+// Ref Gemm1: Y = P * V
+// fp16 in, fp16 out
+using ReferenceGemm1Instance = ck::tensor_operation::host::ReferenceBatchedGemm<DataType,
+                                                                                DataType,
+                                                                                DataType,
+                                                                                AccDataType,
+                                                                                PassThrough,
+                                                                                PassThrough,
+                                                                                PassThrough>;
+
+// Ref Gemm for backward pass
+// fp16 in, fp16 out
+using ReferenceGemmGradInstance = ck::tensor_operation::host::ReferenceBatchedGemm<DataType,
+                                                                                   DataType,
+                                                                                   DataType,
+                                                                                   AccDataType,
+                                                                                   PassThrough,
+                                                                                   PassThrough,
+                                                                                   Scale>;
+// Ref dropout
+using ReferenceDropoutInstance =
+    ck::tensor_operation::host::ReferenceDropout<ushort, DataType, DataType>;
+
+template <typename TensorQ,
+          typename TensorK,
+          typename TensorV,
+          typename TensorS,
+          typename TensorP,
+          typename TensorZ,
+          typename TensorY,
+          typename TensorLSE = TensorP>
+void run_attention_fwd_host(const TensorQ& q_g_m_k,
+                            const TensorK& k_g_n_k,
+                            const TensorV& v_g_n_o,
+                            const float alpha,
+                            TensorS& s_g_m_n,
+                            TensorP& p_g_m_n,
+                            TensorY& y_g_m_o,
+                            TensorLSE& lse_g_m,
+                            TensorP& p_drop_g_m_n,
+                            TensorZ& z_g_m_n,
+                            ushort p_dropout_in_16bits,
+                            float rp_dropout)
+{
+    // S = alpha * Q * K^T
+    auto k_g_k_n            = k_g_n_k.Transpose({0, 2, 1});
+    auto ref_gemm0          = ReferenceGemm0Instance{};
+    auto ref_gemm0_invoker  = ref_gemm0.MakeInvoker();
+    auto ref_gemm0_argument = ref_gemm0.MakeArgument(
+        q_g_m_k, k_g_k_n, s_g_m_n, PassThrough{}, PassThrough{}, Scale{alpha});
+
+    ref_gemm0_invoker.Run(ref_gemm0_argument);
+
+    // masking
+#if USING_MASK
+    auto N          = s_g_m_n.GetLengths()[2];
+    const auto mask = DeviceGemmInstance::C0MatrixMask(N);
+    s_g_m_n.ForEach([&](auto& self, auto idx) {
+        if(mask.IsMaskedElement(idx[1], idx[2]))
+            self(idx) = -ck::NumericLimits<float>::Infinity();
+    });
+#endif
+
+    // P = Softmax(S)
+    auto ref_softmax          = ReferenceSoftmaxInstance{};
+    auto ref_softmax_invoker  = ref_softmax.MakeInvoker();
+    auto ref_softmax_argument = ref_softmax.MakeArgument(s_g_m_n, p_g_m_n, 1, 0, {2}, &lse_g_m);
+
+    ref_softmax_invoker.Run(ref_softmax_argument);
+
+    // P_dropped
+    auto ref_dropout         = ReferenceDropoutInstance{};
+    auto ref_dropout_invoker = ref_dropout.MakeInvoker();
+    auto ref_dropout_argment =
+        ref_dropout.MakeArgument(z_g_m_n, p_g_m_n, p_drop_g_m_n, p_dropout_in_16bits, rp_dropout);
+    ref_dropout_invoker.Run(ref_dropout_argment);
+
+    // Y = P_dropout * V
+    auto ref_gemm1          = ReferenceGemm1Instance{};
+    auto ref_gemm1_invoker  = ref_gemm1.MakeInvoker();
+    auto ref_gemm1_argument = ref_gemm1.MakeArgument(
+        p_drop_g_m_n, v_g_n_o, y_g_m_o, PassThrough{}, PassThrough{}, PassThrough{});
+
+    ref_gemm1_invoker.Run(ref_gemm1_argument);
+}
+
+int run(int argc, char* argv[])
+{
+    bool do_verification = true;
+    int init_method      = 2; // method 1 will have slightly higher error; TODO: to investigate
+    bool time_kernel     = true;
+
+    // Overall QKV matrices shape
+    // y_g_m_o = Softmax(alpha * Q_g_m_k * K_g_k_n) * V_g_n_o
+    // y_g0_g1_m_o = reshape(y_g_m_o, [G0, G1, M, O])
+    // y_g0_m_g1_o = permute(y_g0_g1_m_o, [0, 2, 1, 3])
+    ck::index_t M  = 512;
+    ck::index_t N  = 512;
+    ck::index_t K  = 128;
+    ck::index_t O  = 128;
+    ck::index_t G0 = 3;
+    ck::index_t G1 = 2;
+
+    float alpha = 1.f / std::sqrt(K);
+
+    bool input_permute  = false;
+    bool output_permute = false;
+
+    float p_drop                    = 0.2;
+    float p_dropout                 = 1 - p_drop;
+    uint16_t p_dropout_in_16bits    = uint16_t(std::floor(p_dropout * 65535.0));
+    float rp_dropout                = 1.0 / p_dropout;
+    const unsigned long long seed   = 1;
+    const unsigned long long offset = 0;
+
+    if(argc == 1)
+    {
+        // use default case
+    }
+    else if(argc == 4)
+    {
+        do_verification = std::stoi(argv[1]);
+        init_method     = std::stoi(argv[2]);
+        time_kernel     = std::stoi(argv[3]);
+    }
+    else if(argc == 13)
+    {
+        do_verification = std::stoi(argv[1]);
+        init_method     = std::stoi(argv[2]);
+        time_kernel     = std::stoi(argv[3]);
+
+        M  = std::stoi(argv[4]);
+        N  = std::stoi(argv[5]);
+        K  = std::stoi(argv[6]);
+        O  = std::stoi(argv[7]);
+        G0 = std::stoi(argv[8]);
+        G1 = std::stoi(argv[9]);
+
+        alpha = std::stof(argv[10]);
+
+        input_permute  = std::stoi(argv[11]);
+        output_permute = std::stoi(argv[12]);
+    }
+    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 11: M, N, K, O, G0, G1\n");
+        printf("arg10: scale (alpha)\n");
+        printf("arg11 to 12: input / output permute\n");
+        exit(0);
+    }
+
+    const ck::index_t BatchCount = G0 * G1;
+
+    std::vector<ck::index_t> q_gs_ms_ks_lengths{G0, G1, M, K};
+    std::vector<ck::index_t> q_gs_ms_ks_strides =
+        input_permute
+            ? std::vector<ck::index_t>{M * G1 * K, K, G1 * K, 1} // Q layout [G0, M, G1, K]
+            : std::vector<ck::index_t>{G1 * M * K, M * K, K, 1}; // Q layout [G0, G1, M, K]
+
+    std::vector<ck::index_t> k_gs_ns_ks_lengths{G0, G1, N, K};
+    std::vector<ck::index_t> k_gs_ns_ks_strides =
+        input_permute
+            ? std::vector<ck::index_t>{N * G1 * K, K, G1 * K, 1} // K layout [G0, N, G1, K]
+            : std::vector<ck::index_t>{G1 * N * K, N * K, K, 1}; // K layout [G0, G1, N, K]
+
+    std::vector<ck::index_t> v_gs_os_ns_lengths{G0, G1, O, N};
+    std::vector<ck::index_t> v_gs_os_ns_strides =
+        input_permute
+            ? std::vector<ck::index_t>{N * G1 * O, O, 1, G1 * O} // V layout [G0, N, G1, O]
+            : std::vector<ck::index_t>{G1 * N * O, N * O, 1, O}; // V layout [G0, G1, N, O]
+
+    std::vector<ck::index_t> y_gs_ms_os_lengths{G0, G1, M, O};
+    std::vector<ck::index_t> y_gs_ms_os_strides =
+        output_permute
+            ? std::vector<ck::index_t>{M * G1 * O, O, G1 * O, 1} // Y layout [G0, M, G1, O]
+            : std::vector<ck::index_t>{G1 * M * O, M * O, O, 1}; // Y layout [G0, G1, M, O]
+
+    std::vector<ck::index_t> z_gs_ms_ns_lengths{G0, G1, M, N};
+    std::vector<ck::index_t> z_gs_ms_ns_strides =
+        input_permute
+            ? std::vector<ck::index_t>{M * G1 * N, N, G1 * N, 1} // Z layout [G0, M, G1, N]
+            : std::vector<ck::index_t>{G1 * M * N, M * N, N, 1}; // Z layout [G0, G1, M, N]
+    // The softmax stat log-sum-exp (LSE) is used to speed up softmax calculation in backward pass
+    // Pi = exp(Si) / sum(exp(S0) + exp(S1) + ...)
+    //    = exp(Si) / exp(log(sum(exp() + ...)))
+    //    = exp(Si - log(sum(exp() + ...)))
+    //               ^^^^^^^^^^^^^^^^^^^^^
+    //                       LSE
+    std::vector<ck::index_t> lse_gs_ms_lengths{G0, G1, M};
+    std::vector<ck::index_t> lse_gs_ms_strides{G1 * M, M, 1}; // LSE layout [G0, G1, M]
+
+    Tensor<DataType> q_gs_ms_ks(q_gs_ms_ks_lengths, q_gs_ms_ks_strides);
+    Tensor<DataType> k_gs_ns_ks(k_gs_ns_ks_lengths, k_gs_ns_ks_strides);
+    Tensor<ZDataType> z_gs_ms_ns(z_gs_ms_ns_lengths, z_gs_ms_ns_strides);
+    Tensor<DataType> v_gs_os_ns(v_gs_os_ns_lengths, v_gs_os_ns_strides);
+    Tensor<DataType> y_gs_ms_os(y_gs_ms_os_lengths, y_gs_ms_os_strides);
+    Tensor<DataType> ygrad_gs_ms_os(y_gs_ms_os_lengths, y_gs_ms_os_strides);
+    Tensor<LSEDataType> lse_gs_ms(lse_gs_ms_lengths, lse_gs_ms_strides);
+
+    std::cout << "q_gs_ms_ks: " << q_gs_ms_ks.mDesc << std::endl;
+    std::cout << "k_gs_ns_ks: " << k_gs_ns_ks.mDesc << std::endl;
+    std::cout << "z_gs_ms_ks: " << z_gs_ms_ns.mDesc << std::endl;
+    std::cout << "v_gs_os_ns: " << v_gs_os_ns.mDesc << std::endl;
+    std::cout << "y_gs_ms_os: " << y_gs_ms_os.mDesc << std::endl;
+    std::cout << "lse_gs_ms_os: " << lse_gs_ms.mDesc << std::endl;
+
+    z_gs_ms_ns.GenerateTensorValue(GeneratorTensor_1<DataType>{0});
+    switch(init_method)
+    {
+    case 0: break;
+    case 1:
+        q_gs_ms_ks.GenerateTensorValue(GeneratorTensor_2<DataType>{-2, 2});
+        k_gs_ns_ks.GenerateTensorValue(GeneratorTensor_2<DataType>{-2, 2});
+        v_gs_os_ns.GenerateTensorValue(GeneratorTensor_2<DataType>{-2, 2});
+        ygrad_gs_ms_os.GenerateTensorValue(GeneratorTensor_2<DataType>{-2, 2});
+        break;
+    case 2:
+        q_gs_ms_ks.GenerateTensorValue(GeneratorTensor_3<DataType>{0.0, 1.0});
+        k_gs_ns_ks.GenerateTensorValue(GeneratorTensor_3<DataType>{0.0, 1.0});
+        v_gs_os_ns.GenerateTensorValue(GeneratorTensor_3<DataType>{-0.5, 0.5});
+        ygrad_gs_ms_os.GenerateTensorValue(GeneratorTensor_3<DataType>{-0.5, 0.5});
+        break;
+    case 3:
+        q_gs_ms_ks.GenerateTensorValue(GeneratorTensor_2<DataType>{-5, 5});
+        k_gs_ns_ks.GenerateTensorValue(GeneratorTensor_Diagonal<DataType>{});
+        v_gs_os_ns.GenerateTensorValue(GeneratorTensor_Diagonal<DataType>{});
+        ygrad_gs_ms_os.GenerateTensorValue(GeneratorTensor_Diagonal<DataType>{});
+        break;
+    case 4:
+        q_gs_ms_ks.GenerateTensorValue(GeneratorTensor_1<DataType>{1});
+        k_gs_ns_ks.GenerateTensorValue(GeneratorTensor_Diagonal<DataType>{});
+        v_gs_os_ns.GenerateTensorValue(GeneratorTensor_Diagonal<DataType>{});
+        ygrad_gs_ms_os.GenerateTensorValue(GeneratorTensor_1<DataType>{2});
+        break;
+    case 5:
+        q_gs_ms_ks.GenerateTensorValue(GeneratorTensor_1<DataType>{1});
+        k_gs_ns_ks.GenerateTensorValue(GeneratorTensor_Diagonal<DataType>{});
+        v_gs_os_ns.GenerateTensorValue(GeneratorTensor_Diagonal<DataType>{});
+        ygrad_gs_ms_os.GenerateTensorValue(GeneratorTensor_Sequential<2>{}); // dy[g0, g1, m, o]
+        // dO dot O = [0; 1; 2; ...]
+        break;
+    case 6:
+        q_gs_ms_ks.GenerateTensorValue(GeneratorTensor_1<DataType>{1});
+        k_gs_ns_ks.GenerateTensorValue(GeneratorTensor_Diagonal<DataType>{});
+        v_gs_os_ns.GenerateTensorValue(GeneratorTensor_Diagonal<DataType>{});
+        ygrad_gs_ms_os.GenerateTensorValue(GeneratorTensor_Sequential<3>{}); // dy[g0, g1, m, o]
+        // assume mnko = 256
+        // P = softmax(QK) = 0.0039 * ones
+        // O = P V = 0.0039 * ones
+        // dP = dO V = [0, 1, 2, ...; 0, 1, 2, ...; ...]
+        // dO dot O = [127.5; ...]
+        // dS = P * (dP - dO dot O)
+        //
+        break;
+    default:
+        q_gs_ms_ks.GenerateTensorValue(GeneratorTensor_1<DataType>{1});
+        k_gs_ns_ks.GenerateTensorValue(GeneratorTensor_Diagonal<DataType>{});
+        v_gs_os_ns.GenerateTensorValue(GeneratorTensor_Diagonal<DataType>{});
+        ygrad_gs_ms_os.GenerateTensorValue(GeneratorTensor_1<DataType>{1}); // dy[g0, g1, m, o]
+        // assume mnko = 256
+        // P = softmax(QK) = 0.0039 * ones
+        // O = P V = 0.0039 * ones
+        // dP = dO V = ones
+        // dS = P * (dP - (dO dot O))
+        //    = 0.0039 * ones * (ones - 0.0039*256)
+        //    = 0.0039 * ones * (ones - 1)
+        //    = 0
+    }
+
+    // calculate y & log-sum-exp beforehand
+    Tensor<DataType> q_g_m_k({BatchCount, M, K});
+    Tensor<DataType> k_g_n_k({BatchCount, N, K});
+    Tensor<ZDataType> z_g_m_n({BatchCount, M, N});
+    Tensor<DataType> v_g_n_o({BatchCount, N, O});
+    Tensor<AccDataType> s_g_m_n({BatchCount, M, N});
+    Tensor<DataType> p_g_m_n({BatchCount, M, N});
+    Tensor<DataType> p_drop_g_m_n({BatchCount, M, N});
+    Tensor<DataType> y_g_m_o({BatchCount, M, O});
+    Tensor<LSEDataType> lse_g_m({BatchCount, M});
+
+    q_gs_ms_ks.ForEach(
+        [&](auto& self, auto idx) { q_g_m_k(idx[0] * G1 + idx[1], idx[2], idx[3]) = self(idx); });
+    k_gs_ns_ks.ForEach(
+        [&](auto& self, auto idx) { k_g_n_k(idx[0] * G1 + idx[1], idx[2], idx[3]) = self(idx); });
+    z_gs_ms_ns.ForEach(
+        [&](auto& self, auto idx) { z_g_m_n(idx[0] * G1 + idx[1], idx[2], idx[3]) = self(idx); });
+    v_gs_os_ns.ForEach(
+        [&](auto& self, auto idx) { v_g_n_o(idx[0] * G1 + idx[1], idx[3], idx[2]) = self(idx); });
+    lse_gs_ms.ForEach(
+        [&](auto& self, auto idx) { lse_g_m(idx[0] * G1 + idx[1], idx[2]) = self(idx); });
+
+    run_attention_fwd_host(q_g_m_k,
+                           k_g_n_k,
+                           v_g_n_o,
+                           alpha,
+                           s_g_m_n,
+                           p_g_m_n,
+                           y_g_m_o,
+                           lse_g_m,
+                           p_drop_g_m_n,
+                           z_g_m_n,
+                           p_dropout_in_16bits,
+                           rp_dropout);
+
+    y_gs_ms_os.ForEach(
+        [&](auto& self, auto idx) { self(idx) = y_g_m_o(idx[0] * G1 + idx[1], idx[2], idx[3]); });
+    lse_gs_ms.ForEach(
+        [&](auto& self, auto idx) { self(idx) = lse_g_m(idx[0] * G1 + idx[1], idx[2]); });
+
+    // qkv gradients have the same descriptor as with qkv
+    DeviceMem q_device_buf(sizeof(DataType) * q_gs_ms_ks.mDesc.GetElementSpaceSize());
+    DeviceMem k_device_buf(sizeof(DataType) * k_gs_ns_ks.mDesc.GetElementSpaceSize());
+    DeviceMem z_device_buf(sizeof(ZDataType) * z_gs_ms_ns.mDesc.GetElementSpaceSize());
+    DeviceMem v_device_buf(sizeof(DataType) * v_gs_os_ns.mDesc.GetElementSpaceSize());
+    DeviceMem y_device_buf(sizeof(DataType) * y_gs_ms_os.mDesc.GetElementSpaceSize());
+    DeviceMem lse_device_buf(sizeof(LSEDataType) * lse_gs_ms.mDesc.GetElementSpaceSize());
+    DeviceMem qgrad_device_buf(sizeof(DataType) * q_gs_ms_ks.mDesc.GetElementSpaceSize());
+    DeviceMem kgrad_device_buf(sizeof(DataType) * k_gs_ns_ks.mDesc.GetElementSpaceSize());
+    DeviceMem vgrad_device_buf(sizeof(DataType) * v_gs_os_ns.mDesc.GetElementSpaceSize());
+    DeviceMem ygrad_device_buf(sizeof(DataType) * y_gs_ms_os.mDesc.GetElementSpaceSize());
+
+    q_device_buf.ToDevice(q_gs_ms_ks.mData.data());
+    k_device_buf.ToDevice(k_gs_ns_ks.mData.data());
+    z_device_buf.ToDevice(z_gs_ms_ns.mData.data());
+    v_device_buf.ToDevice(v_gs_os_ns.mData.data());
+    y_device_buf.ToDevice(y_gs_ms_os.mData.data());
+    lse_device_buf.ToDevice(lse_gs_ms.mData.data());
+    ygrad_device_buf.ToDevice(ygrad_gs_ms_os.mData.data());
+    kgrad_device_buf.SetZero();
+    vgrad_device_buf.SetZero();
+    // z_device_buf.SetZero();
+
+    auto gemm     = DeviceGemmInstance{};
+    auto invoker  = gemm.MakeInvoker();
+    auto argument = gemm.MakeArgument(
+        static_cast<DataType*>(q_device_buf.GetDeviceBuffer()),
+        static_cast<DataType*>(k_device_buf.GetDeviceBuffer()),
+        static_cast<ZDataType*>(z_device_buf.GetDeviceBuffer()),
+        static_cast<DataType*>(v_device_buf.GetDeviceBuffer()),
+        static_cast<DataType*>(y_device_buf.GetDeviceBuffer()),
+        static_cast<LSEDataType*>(lse_device_buf.GetDeviceBuffer()),
+        static_cast<DataType*>(ygrad_device_buf.GetDeviceBuffer()),
+        static_cast<DataType*>(qgrad_device_buf.GetDeviceBuffer()),
+        static_cast<DataType*>(kgrad_device_buf.GetDeviceBuffer()),
+        static_cast<DataType*>(vgrad_device_buf.GetDeviceBuffer()),
+        {}, // std::array<void*, 1> p_acc0_biases;
+        {}, // std::array<void*, 1> p_acc1_biases;
+        q_gs_ms_ks_lengths,
+        q_gs_ms_ks_strides,
+        k_gs_ns_ks_lengths,
+        k_gs_ns_ks_strides,
+        z_gs_ms_ns_lengths,
+        z_gs_ms_ns_strides,
+        v_gs_os_ns_lengths,
+        v_gs_os_ns_strides,
+        y_gs_ms_os_lengths,
+        y_gs_ms_os_strides,
+        lse_gs_ms_lengths,
+        {}, // std::array<std::vector<ck::index_t>, 1>{acc0_biases_gs_ms_ns_lengths},
+        {}, // std::array<std::vector<ck::index_t>, 1>{acc0_biases_gs_ms_ns_strides},
+        {}, // std::array<std::vector<ck::index_t>, 1>{acc1_biases_gs_ms_os_lengths},
+        {}, // std::array<std::vector<ck::index_t>, 1>{acc1_biases_gs_ms_os_strides},
+        QKVElementOp{},
+        QKVElementOp{},
+        Scale{alpha},
+        QKVElementOp{},
+        YElementOp{},
+        p_drop,
+        std::tuple<unsigned long long, unsigned long long>(seed, offset));
+
+    if(!gemm.IsSupportedArgument(argument))
+    {
+        std::cout << gemm.GetTypeString() << " does not support this problem" << std::endl;
+
+        return 0;
+    }
+
+    float ave_time = invoker.Run(argument, StreamConfig{nullptr, time_kernel});
+
+    // 5 GEMM ops in total:
+    // S_MNK / dP_MNO Gemm (Gemm0 rcr)
+    // dQ_MKN Gemm (Gemm1 rrr)
+    // dV_NOM / dK_NKM Gemm (Gemm2 crr)
+    // 3x MNK + 2x MNO
+    std::size_t flop = (size_t(3) * M * N * K + size_t(2) * M * N * O) * 2 * BatchCount;
+    // Q/K/V/Y, dQ/dK/dV/dY, LSE
+    std::size_t num_btype = (sizeof(DataType) * M * K + sizeof(DataType) * K * N +
+                             sizeof(DataType) * N * O + sizeof(DataType) * M * O) *
+                                size_t(2) * BatchCount +
+                            sizeof(LSEDataType) * M * BatchCount;
+
+    float tflops = static_cast<float>(flop) / 1.E9 / ave_time;
+
+    float gb_per_sec = num_btype / 1.E6 / ave_time;
+
+    std::cout << "Perf: " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec << " GB/s, "
+              << gemm.GetTypeString() << std::endl;
+
+    // copy z matirx data form device
+    std::ofstream file("./z_matrix_txt");
+    z_device_buf.FromDevice(z_g_m_n.mData.data());
+    file << z_g_m_n << std::endl;
+    //       std::cout << "z_g_m_n ref:\n" << z_g_m_n;
+    bool pass = true;
+    if(do_verification)
+    {
+        // run fowad again for y, cause z_g_m_n update
+        run_attention_fwd_host(q_g_m_k,
+                               k_g_n_k,
+                               v_g_n_o,
+                               alpha,
+                               s_g_m_n,
+                               p_g_m_n,
+                               y_g_m_o,
+                               lse_g_m,
+                               p_drop_g_m_n,
+                               z_g_m_n,
+                               p_dropout_in_16bits,
+                               rp_dropout);
+        y_gs_ms_os.ForEach([&](auto& self, auto idx) {
+            self(idx) = y_g_m_o(idx[0] * G1 + idx[1], idx[2], idx[3]);
+        });
+        y_device_buf.ToDevice(y_gs_ms_os.mData.data());
+
+        //
+        // call kernel again
+        //
+        // example set Z matrix to null, will not ouput z matrix data
+        argument = gemm.MakeArgument(
+            static_cast<DataType*>(q_device_buf.GetDeviceBuffer()),
+            static_cast<DataType*>(k_device_buf.GetDeviceBuffer()),
+            static_cast<ZDataType*>(nullptr), // set to nullptr
+            static_cast<DataType*>(v_device_buf.GetDeviceBuffer()),
+            static_cast<DataType*>(y_device_buf.GetDeviceBuffer()),
+            static_cast<LSEDataType*>(lse_device_buf.GetDeviceBuffer()),
+            static_cast<DataType*>(ygrad_device_buf.GetDeviceBuffer()),
+            static_cast<DataType*>(qgrad_device_buf.GetDeviceBuffer()),
+            static_cast<DataType*>(kgrad_device_buf.GetDeviceBuffer()),
+            static_cast<DataType*>(vgrad_device_buf.GetDeviceBuffer()),
+            {}, // std::array<void*, 1> p_acc0_biases;
+            {}, // std::array<void*, 1> p_acc1_biases;
+            q_gs_ms_ks_lengths,
+            q_gs_ms_ks_strides,
+            k_gs_ns_ks_lengths,
+            k_gs_ns_ks_strides,
+            z_gs_ms_ns_lengths,
+            z_gs_ms_ns_strides,
+            v_gs_os_ns_lengths,
+            v_gs_os_ns_strides,
+            y_gs_ms_os_lengths,
+            y_gs_ms_os_strides,
+            lse_gs_ms_lengths,
+            {}, // std::array<std::vector<ck::index_t>, 1>{acc0_biases_gs_ms_ns_lengths},
+            {}, // std::array<std::vector<ck::index_t>, 1>{acc0_biases_gs_ms_ns_strides},
+            {}, // std::array<std::vector<ck::index_t>, 1>{acc1_biases_gs_ms_os_lengths},
+            {}, // std::array<std::vector<ck::index_t>, 1>{acc1_biases_gs_ms_os_strides},
+            QKVElementOp{},
+            QKVElementOp{},
+            Scale{alpha},
+            QKVElementOp{},
+            YElementOp{},
+            p_drop,
+            std::tuple<unsigned long long, unsigned long long>(seed, offset));
+        kgrad_device_buf.SetZero(); // reset global accum buffer and rerun
+        vgrad_device_buf.SetZero();
+        invoker.Run(argument, StreamConfig{nullptr, false});
+
+        Tensor<DataType> qgrad_g_m_k({BatchCount, M, K});
+        Tensor<DataType> kgrad_g_n_k({BatchCount, N, K});
+        Tensor<DataType> vgrad_g_n_o({BatchCount, N, O});
+        Tensor<DataType> sgrad_g_m_n({BatchCount, M, N});
+        Tensor<DataType> pgrad_g_m_n({BatchCount, M, N});
+        Tensor<DataType> pgrad_drop_g_m_n({BatchCount, M, N});
+        Tensor<DataType> ygrad_g_m_o({BatchCount, M, O});
+        Tensor<DataType> ygrad_dot_y_g_m({BatchCount, M});
+
+        ygrad_gs_ms_os.ForEach([&](auto& self, auto idx) {
+            ygrad_g_m_o(idx[0] * G1 + idx[1], idx[2], idx[3]) = self(idx);
+        });
+
+#if PRINT_HOST
+        {
+            std::cout << "q_g_m_k ref:\n" << q_g_m_k;
+            std::cout << "k_g_n_k ref:\n" << k_g_n_k;
+            std::cout << "v_g_n_o ref:\n" << v_g_n_o;
+            std::cout << "ygrad_g_m_o ref:\n" << ygrad_g_m_o;
+        }
+#endif
+
+        // Gradients
+        auto ref_gemm_grad         = ReferenceGemmGradInstance{};
+        auto ref_gemm_grad_invoker = ref_gemm_grad.MakeInvoker();
+        using RefGemmGradArg       = ReferenceGemmGradInstance::Argument;
+
+        // dP_dropout = dY * V^T
+        auto v_g_o_n = v_g_n_o.Transpose({0, 2, 1});
+        ref_gemm_grad_invoker.Run(RefGemmGradArg{
+            ygrad_g_m_o, v_g_o_n, pgrad_drop_g_m_n, PassThrough{}, PassThrough{}, Scale{1.f}});
+#if PRINT_HOST
+        {
+            std::cout << "===== dP = dY * V^T\n";
+            std::cout << "ygrad_drop_g_m_o ref:\n" << ygrad_drop_g_m_n;
+            std::cout << "v_g_o_n ref:\n" << v_g_o_n;
+            std::cout << "pgrad_drop_g_m_n ref:\n" << pgrad_drop_g_m_n;
+        }
+#endif
+        // dP = dP_dropout x Z
+        auto ref_dropout         = ReferenceDropoutInstance{};
+        auto ref_dropout_invoker = ref_dropout.MakeInvoker();
+        auto ref_dropout_argment = ref_dropout.MakeArgument(
+            z_g_m_n, pgrad_drop_g_m_n, pgrad_g_m_n, p_dropout_in_16bits, rp_dropout);
+        ref_dropout_invoker.Run(ref_dropout_argment);
+
+        // dS_i_j = P_i_j .* (dP_i_j -  dY_i dot Y_i)
+        sgrad_g_m_n.ForEach([&](auto& self, auto idx_gmn) {
+            float ygrad_dot_y = 0;
+            for(int o = 0; o < O; o++)
+            {
+                auto idx_gmo = idx_gmn;
+                idx_gmo[2]   = o;
+                ygrad_dot_y += ygrad_g_m_o(idx_gmo) * y_g_m_o(idx_gmo);
+            }
+            self(idx_gmn) = p_g_m_n(idx_gmn) * (pgrad_g_m_n(idx_gmn) - ygrad_dot_y);
+        });
+#if PRINT_HOST
+        {
+            std::cout << "===== dS_i_j = P_i_j .* (dP_i_j - dY_i dot Y_i)\n";
+            std::cout << "p_g_m_n ref:\n" << p_g_m_n;
+            std::cout << "pgrad_g_m_n ref:\n" << pgrad_g_m_n;
+            std::cout << "y_g_m_o ref:\n" << y_g_m_o;
+            std::cout << "ygrad_g_m_o ref:\n" << ygrad_g_m_o;
+            std::cout << "sgrad_g_m_n ref:\n" << sgrad_g_m_n;
+        }
+#endif
+        // dV = P_drop^T * dY
+        auto p_drop_g_n_m = p_drop_g_m_n.Transpose({0, 2, 1});
+        ref_gemm_grad_invoker.Run(RefGemmGradArg{
+            p_drop_g_n_m, ygrad_g_m_o, vgrad_g_n_o, PassThrough{}, PassThrough{}, Scale{1.0f}});
+#if PRINT_HOST
+        {
+            std::cout << "===== dV = P^T * dY\n";
+            std::cout << "p_drop_g_n_m ref:\n" << p_drop_g_n_m;
+            std::cout << "ygrad_g_m_o ref:\n" << ygrad_g_m_o;
+            std::cout << "vgrad_g_n_o ref:\n" << vgrad_g_n_o;
+        }
+#endif
+
+        // dQ = alpha * dS * K
+        ref_gemm_grad_invoker.Run(RefGemmGradArg{
+            sgrad_g_m_n, k_g_n_k, qgrad_g_m_k, PassThrough{}, PassThrough{}, Scale{alpha}});
+#if PRINT_HOST
+        {
+            std::cout << "===== dQ = alpha * dS * K\n";
+            std::cout << "sgrad_g_m_n ref:\n" << sgrad_g_m_n;
+            std::cout << "k_g_n_k ref:\n" << k_g_n_k;
+            std::cout << "qgrad_g_m_k ref:\n" << qgrad_g_m_k;
+        }
+#endif
+
+        // dK = alpha * dS^T * Q
+        auto sgrad_g_n_m = sgrad_g_m_n.Transpose({0, 2, 1});
+        ref_gemm_grad_invoker.Run(RefGemmGradArg{
+            sgrad_g_n_m, q_g_m_k, kgrad_g_n_k, PassThrough{}, PassThrough{}, Scale{alpha}});
+#if PRINT_HOST
+        {
+            std::cout << "===== dK = alpha * dS^T * Q\n";
+            std::cout << "sgrad_g_n_m ref:\n" << sgrad_g_n_m;
+            std::cout << "q_g_m_k ref:\n" << q_g_m_k;
+            std::cout << "kgrad_g_n_k ref:\n" << kgrad_g_n_k;
+        }
+#endif
+
+        Tensor<DataType> qgrad_gs_ms_ks_host_result(q_gs_ms_ks_lengths, q_gs_ms_ks_strides);
+        Tensor<DataType> kgrad_gs_ns_ks_host_result(k_gs_ns_ks_lengths, k_gs_ns_ks_strides);
+        Tensor<DataType> vgrad_gs_os_ns_host_result(v_gs_os_ns_lengths, v_gs_os_ns_strides);
+
+        Tensor<DataType> qgrad_gs_ms_ks_device_result(q_gs_ms_ks_lengths, q_gs_ms_ks_strides);
+        Tensor<DataType> kgrad_gs_ns_ks_device_result(k_gs_ns_ks_lengths, k_gs_ns_ks_strides);
+        Tensor<DataType> vgrad_gs_os_ns_device_result(v_gs_os_ns_lengths, v_gs_os_ns_strides);
+
+        qgrad_device_buf.FromDevice(qgrad_gs_ms_ks_device_result.mData.data());
+        kgrad_device_buf.FromDevice(kgrad_gs_ns_ks_device_result.mData.data());
+        vgrad_device_buf.FromDevice(vgrad_gs_os_ns_device_result.mData.data());
+
+        // permute
+        qgrad_gs_ms_ks_host_result.ForEach([&](auto& self, auto idx) {
+            const size_t& g0 = idx[0];
+            const size_t& g1 = idx[1];
+
+            const size_t g = g0 * G1 + g1;
+
+            self(idx) = qgrad_g_m_k(g, idx[2], idx[3]);
+        });
+        kgrad_gs_ns_ks_host_result.ForEach([&](auto& self, auto idx) {
+            const size_t& g0 = idx[0];
+            const size_t& g1 = idx[1];
+
+            const size_t g = g0 * G1 + g1;
+
+            self(idx) = kgrad_g_n_k(g, idx[2], idx[3]);
+        });
+        vgrad_gs_os_ns_host_result.ForEach([&](auto& self, auto idx) {
+            const size_t& g0 = idx[0];
+            const size_t& g1 = idx[1];
+
+            const size_t g = g0 * G1 + g1;
+
+            self(idx) = vgrad_g_n_o(g, idx[3], idx[2]);
+        });
+
+        std::cout << "Checking qgrad:\n";
+        pass &= ck::utils::check_err(qgrad_gs_ms_ks_device_result.mData,
+                                     qgrad_gs_ms_ks_host_result.mData,
+                                     "error",
+                                     1e-2,
+                                     1e-2);
+        std::cout << "Checking kgrad:\n";
+        pass &= ck::utils::check_err(kgrad_gs_ns_ks_device_result.mData,
+                                     kgrad_gs_ns_ks_host_result.mData,
+                                     "error",
+                                     1e-2,
+                                     1e-2);
+        std::cout << "Checking vgrad:\n";
+        pass &= ck::utils::check_err(vgrad_gs_os_ns_device_result.mData,
+                                     vgrad_gs_os_ns_host_result.mData,
+                                     "error",
+                                     1e-2,
+                                     1e-2);
+    }
+
+    return pass ? ((void)(std::cout << "pass\n"), 0) : ((void)(std::cout << "fail\n"), 1);
+}
+
+int main(int argc, char* argv[]) { return run(argc, argv); }

example/32_batched_gemm_scale_softmax_gemm/grouped_multihead_attention_backward_fp16.cpp

@@ -44,6 +44,7 @@ Kernel outputs:
 #include "ck/library/utility/host_tensor_generator.hpp"
 #include "ck/library/reference_tensor_operation/cpu/reference_batched_gemm.hpp"
 #include "ck/library/reference_tensor_operation/cpu/reference_softmax.hpp"
+#include "ck/library/reference_tensor_operation/cpu/reference_dropout.hpp"
 
 #define FLASH_ATTN_IMPLENTATION 0
 
@@ -52,6 +53,7 @@ using S = ck::Sequence<Is...>;
 
 using F16 = ck::half_t;
 using F32 = float;
+using U16 = unsigned short;
 
 using PassThrough = ck::tensor_operation::element_wise::PassThrough;
 using Scale       = ck::tensor_operation::element_wise::Scale;
@@ -63,6 +65,7 @@ using DataType         = F16;
 using AccDataType      = F32;
 using ShuffleDataType  = F32;
 using LSEDataType      = F32;
+using ZDataType        = U16;
 using Acc0BiasDataType = ck::Tuple<>;
 using Acc1BiasDataType = ck::Tuple<>;
 
@@ -160,6 +163,7 @@ using DeviceGemmInstance =
         NumDimK,
         NumDimO,
         DataType,
+        ZDataType,
         LSEDataType,
         Acc0BiasDataType,
         Acc1BiasDataType,
@@ -179,9 +183,9 @@ using DeviceGemmInstance =
         256,
         128,         // MPerBlock
         128,         // NPerBlock
-        64,          // KPerBlock
-        64,          // Gemm1NPerBlock
-        32,          // Gemm1KPerBlock
+        32,          // KPerBlock
+        128,         // Gemm1NPerBlock
+        64,          // Gemm1KPerBlock
         8,           // AK1
         8,           // BK1
         2,           // B1K1
@@ -189,7 +193,7 @@ using DeviceGemmInstance =
         32,          // NPerXDL
         1,           // MXdlPerWave
         4,           // NXdlPerWave
-        2,           // Gemm1NXdlPerWave
+        4,           // Gemm1NXdlPerWave
         S<4, 64, 1>, // ABlockTransfer
         S<1, 0, 2>,
         S<1, 0, 2>,
@@ -212,7 +216,7 @@ using DeviceGemmInstance =
         2,
         false,
         1,              // CShuffleMXdlPerWavePerShuffle
-        2,              // CShuffleNXdlPerWavePerShuffle
+        4,              // CShuffleNXdlPerWavePerShuffle
         S<1, 32, 1, 8>, // CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock
         8,              // CShuffleBlockTransferScalarPerVector_NPerBlock
         MaskingSpec>;   // MaskingSpecialization
@@ -252,11 +256,15 @@ using ReferenceGemmGradInstance = ck::tensor_operation::host::ReferenceBatchedGe
                                                                                    PassThrough,
                                                                                    Scale>;
 
+using ReferenceDropoutInstance =
+    ck::tensor_operation::host::ReferenceDropout<ushort, DataType, DataType>;
+
 template <typename TensorQ,
           typename TensorK,
           typename TensorV,
           typename TensorS,
           typename TensorP,
+          typename TensorZ,
           typename TensorY,
           typename TensorLSE = TensorP>
 void run_attention_fwd_host(const TensorQ& q_g_m_k,
@@ -266,7 +274,11 @@ void run_attention_fwd_host(const TensorQ& q_g_m_k,
                             TensorS& s_g_m_n,
                             TensorP& p_g_m_n,
                             TensorY& y_g_m_o,
-                            TensorLSE& lse_g_m)
+                            TensorLSE& lse_g_m,
+                            TensorP& p_drop_g_m_n,
+                            TensorZ& z_g_m_n,
+                            ushort p_dropout_in_16bits,
+                            float rp_dropout)
 {
     // S = alpha * Q * K^T
     auto k_g_k_n            = k_g_n_k.Transpose({0, 2, 1});
@@ -294,6 +306,12 @@ void run_attention_fwd_host(const TensorQ& q_g_m_k,
 
     ref_softmax_invoker.Run(ref_softmax_argument);
 
+    auto ref_dropout         = ReferenceDropoutInstance{};
+    auto ref_dropout_invoker = ref_dropout.MakeInvoker();
+    auto ref_dropout_argment =
+        ref_dropout.MakeArgument(z_g_m_n, p_g_m_n, p_drop_g_m_n, p_dropout_in_16bits, rp_dropout);
+    ref_dropout_invoker.Run(ref_dropout_argment);
+
     // Y = P * V
     auto ref_gemm1          = ReferenceGemm1Instance{};
     auto ref_gemm1_invoker  = ref_gemm1.MakeInvoker();
@@ -313,8 +331,14 @@ int run(int argc, char* argv[])
     // y_g_m_o = Softmax(alpha * Q_g_m_k * K_g_k_n) * V_g_n_o
     // y_g0_g1_m_o = reshape(y_g_m_o, [G0, G1, M, O])
     // y_g0_m_g1_o = permute(y_g0_g1_m_o, [0, 2, 1, 3])
-    float K     = 128;
-    float alpha = 1.f / std::sqrt(K);
+    float K                         = 128;
+    float alpha                     = 1.f / std::sqrt(K);
+    float p_drop                    = 0.2;
+    float p_dropout                 = 1 - p_drop;
+    uint16_t p_dropout_in_16bits    = uint16_t(std::floor(p_dropout * 65535.0));
+    float rp_dropout                = 1.0 / p_dropout;
+    const unsigned long long seed   = 1;
+    const unsigned long long offset = 0;
 
     bool input_permute  = false;
     bool output_permute = false;
@@ -358,6 +382,7 @@ int run(int argc, char* argv[])
     using DeviceMemPtr = std::unique_ptr<DeviceMem>;
     std::vector<const void*> p_q;
     std::vector<const void*> p_k;
+    std::vector<void*> p_z;
     std::vector<const void*> p_v;
     std::vector<const void*> p_y;
     std::vector<const void*> p_lse;
@@ -368,6 +393,7 @@ int run(int argc, char* argv[])
 
     std::vector<Tensor<DataType>> q_g_m_ks;
     std::vector<Tensor<DataType>> k_g_n_ks;
+    std::vector<Tensor<ZDataType>> z_g_m_ns;
     std::vector<Tensor<DataType>> v_g_n_os;
     std::vector<Tensor<AccDataType>> s_g_m_ns;
     std::vector<Tensor<DataType>> p_g_m_ns;
@@ -376,6 +402,7 @@ int run(int argc, char* argv[])
     std::vector<Tensor<DataType>> k_tensors;
     std::vector<Tensor<DataType>> v_tensors;
     std::vector<Tensor<DataType>> y_tensors;
+    std::vector<Tensor<ZDataType>> z_tensors;
     std::vector<Tensor<LSEDataType>> lse_tensors;
     std::vector<Tensor<DataType>> qgrad_tensors;
     std::vector<Tensor<DataType>> kgrad_tensors;
@@ -384,6 +411,7 @@ int run(int argc, char* argv[])
 
     std::vector<DeviceMemPtr> q_tensors_device;
     std::vector<DeviceMemPtr> k_tensors_device;
+    std::vector<DeviceMemPtr> z_tensors_device;
     std::vector<DeviceMemPtr> v_tensors_device;
     std::vector<DeviceMemPtr> y_tensors_device;
     std::vector<DeviceMemPtr> lse_tensors_device;
@@ -425,6 +453,11 @@ int run(int argc, char* argv[])
                 ? std::vector<ck::index_t>{M * G1 * O, O, G1 * O, 1} // Y layout [G0, M, G1, O]
                 : std::vector<ck::index_t>{G1 * M * O, M * O, O, 1}; // Y layout [G0, G1, M, O]
 
+        std::vector<ck::index_t> z_gs_ms_ns_lengths{G0, G1, M, N};
+        std::vector<ck::index_t> z_gs_ms_ns_strides =
+            input_permute
+                ? std::vector<ck::index_t>{M * G1 * N, N, G1 * N, 1} // Z layout [G0, M, G1, N]
+                : std::vector<ck::index_t>{G1 * M * N, M * N, N, 1}; // Z layout [G0, G1, M, N]
         // The softmax stat log-sum-exp (LSE) is used to speed up softmax calculation in backward
         // pass Pi = exp(Si) / sum(exp(S0) + exp(S1) + ...)
         //    = exp(Si) / exp(log(sum(exp() + ...)))
@@ -438,6 +471,8 @@ int run(int argc, char* argv[])
             q_gs_ms_ks_strides,
             k_gs_ns_ks_lengths,
             k_gs_ns_ks_strides,
+            z_gs_ms_ns_lengths,
+            z_gs_ms_ns_strides,
             v_gs_os_ns_lengths,
             v_gs_os_ns_strides,
             y_gs_ms_os_lengths,
@@ -460,6 +495,7 @@ int run(int argc, char* argv[])
 
         Tensor<DataType> q_gs_ms_ks(q_gs_ms_ks_lengths, q_gs_ms_ks_strides);
         Tensor<DataType> k_gs_ns_ks(k_gs_ns_ks_lengths, k_gs_ns_ks_strides);
+        Tensor<ZDataType> z_gs_ms_ns(z_gs_ms_ns_lengths, z_gs_ms_ns_strides);
         Tensor<DataType> v_gs_os_ns(v_gs_os_ns_lengths, v_gs_os_ns_strides);
         Tensor<DataType> y_gs_ms_os(y_gs_ms_os_lengths, y_gs_ms_os_strides);
         Tensor<DataType> ygrad_gs_ms_os(y_gs_ms_os_lengths, y_gs_ms_os_strides);
@@ -472,6 +508,7 @@ int run(int argc, char* argv[])
             std::cout << "y_gs_ms_os: " << y_gs_ms_os.mDesc << std::endl;
             std::cout << "lse_gs_ms_os: " << lse_gs_ms.mDesc << std::endl;
         }
+        z_gs_ms_ns.GenerateTensorValue(GeneratorTensor_1<DataType>{0});
         switch(init_method)
         {
         case 0: break;
@@ -535,11 +572,13 @@ int run(int argc, char* argv[])
         }
         Tensor<DataType> q_g_m_k({BatchCount, M, K});
         Tensor<DataType> k_g_n_k({BatchCount, N, K});
+        Tensor<ZDataType> z_g_m_n({BatchCount, M, N});
         Tensor<DataType> v_g_n_o({BatchCount, N, O});
         Tensor<AccDataType> s_g_m_n({BatchCount, M, N});
         Tensor<DataType> p_g_m_n({BatchCount, M, N});
         Tensor<DataType> y_g_m_o({BatchCount, M, O});
         Tensor<LSEDataType> lse_g_m({BatchCount, M});
+        Tensor<DataType> p_drop_g_m_n({BatchCount, M, N});
 
         q_gs_ms_ks.ForEach([&](auto& self, auto idx) {
             q_g_m_k(idx[0] * G1 + idx[1], idx[2], idx[3]) = self(idx);
@@ -547,14 +586,27 @@ int run(int argc, char* argv[])
         k_gs_ns_ks.ForEach([&](auto& self, auto idx) {
             k_g_n_k(idx[0] * G1 + idx[1], idx[2], idx[3]) = self(idx);
         });
+        z_gs_ms_ns.ForEach([&](auto& self, auto idx) {
+            z_g_m_n(idx[0] * G1 + idx[1], idx[2], idx[3]) = self(idx);
+        });
         v_gs_os_ns.ForEach([&](auto& self, auto idx) {
             v_g_n_o(idx[0] * G1 + idx[1], idx[3], idx[2]) = self(idx);
         });
         lse_gs_ms.ForEach(
             [&](auto& self, auto idx) { lse_g_m(idx[0] * G1 + idx[1], idx[2]) = self(idx); });
 
-        run_attention_fwd_host(
-            q_g_m_k, k_g_n_k, v_g_n_o, alpha, s_g_m_n, p_g_m_n, y_g_m_o, lse_g_m);
+        run_attention_fwd_host(q_g_m_k,
+                               k_g_n_k,
+                               v_g_n_o,
+                               alpha,
+                               s_g_m_n,
+                               p_g_m_n,
+                               y_g_m_o,
+                               lse_g_m,
+                               p_drop_g_m_n,
+                               z_g_m_n,
+                               p_dropout_in_16bits,
+                               rp_dropout);
 
         y_gs_ms_os.ForEach([&](auto& self, auto idx) {
             self(idx) = y_g_m_o(idx[0] * G1 + idx[1], idx[2], idx[3]);
@@ -564,6 +616,7 @@ int run(int argc, char* argv[])
 
         q_g_m_ks.push_back(q_g_m_k);
         k_g_n_ks.push_back(k_g_n_k);
+        z_g_m_ns.push_back(z_g_m_n);
         v_g_n_os.push_back(v_g_n_o);
         s_g_m_ns.push_back(s_g_m_n);
         p_g_m_ns.push_back(p_g_m_n);
@@ -572,12 +625,15 @@ int run(int argc, char* argv[])
         k_tensors.push_back(k_gs_ns_ks);
         v_tensors.push_back(v_gs_os_ns);
         y_tensors.push_back(y_gs_ms_os);
+        z_tensors.push_back(z_gs_ms_ns);
         lse_tensors.push_back(lse_gs_ms);
         ygrad_tensors.push_back(ygrad_gs_ms_os);
         q_tensors_device.emplace_back(
             std::make_unique<DeviceMem>(sizeof(DataType) * q_gs_ms_ks.GetElementSpaceSize()));
         k_tensors_device.emplace_back(
             std::make_unique<DeviceMem>(sizeof(DataType) * k_gs_ns_ks.GetElementSpaceSize()));
+        z_tensors_device.emplace_back(
+            std::make_unique<DeviceMem>(sizeof(ZDataType) * z_gs_ms_ns.GetElementSpaceSize()));
         v_tensors_device.emplace_back(
             std::make_unique<DeviceMem>(sizeof(DataType) * v_gs_os_ns.GetElementSpaceSize()));
         y_tensors_device.emplace_back(
@@ -594,6 +650,7 @@ int run(int argc, char* argv[])
             std::make_unique<DeviceMem>(sizeof(DataType) * y_gs_ms_os.GetElementSpaceSize()));
         q_tensors_device.back()->ToDevice(q_gs_ms_ks.data());
         k_tensors_device.back()->ToDevice(k_gs_ns_ks.data());
+        z_tensors_device.back()->ToDevice(z_gs_ms_ns.data());
         v_tensors_device.back()->ToDevice(v_gs_os_ns.data());
         y_tensors_device.back()->ToDevice(y_gs_ms_os.data());
         lse_tensors_device.back()->ToDevice(lse_gs_ms.data());
@@ -603,6 +660,7 @@ int run(int argc, char* argv[])
         ygrad_tensors_device.back()->ToDevice(ygrad_gs_ms_os.data());
         p_q.push_back(q_tensors_device.back()->GetDeviceBuffer());
         p_k.push_back(k_tensors_device.back()->GetDeviceBuffer());
+        p_z.push_back(z_tensors_device.back()->GetDeviceBuffer());
         p_v.push_back(v_tensors_device.back()->GetDeviceBuffer());
         p_y.push_back(y_tensors_device.back()->GetDeviceBuffer());
         p_lse.push_back(lse_tensors_device.back()->GetDeviceBuffer());
@@ -611,23 +669,27 @@ int run(int argc, char* argv[])
         p_ygrad.push_back(ygrad_tensors_device.back()->GetDeviceBuffer());
         p_qgrad.push_back(qgrad_tensors_device.back()->GetDeviceBuffer());
     }
-    auto argument = gemm.MakeArgument(p_q,
-                                      p_k,
-                                      p_v,
-                                      p_y,
-                                      p_lse,
-                                      p_ygrad,
-                                      p_qgrad,
-                                      p_kgrad,
-                                      p_vgrad,
-                                      {}, // std::array<void*, 1> p_acc0_biases;
-                                      {}, // std::array<void*, 1> p_acc1_biases;
-                                      problem_descs,
-                                      QKVElementOp{},
-                                      QKVElementOp{},
-                                      Scale{alpha},
-                                      QKVElementOp{},
-                                      YElementOp{});
+    auto argument =
+        gemm.MakeArgument(p_q,
+                          p_k,
+                          p_z,
+                          p_v,
+                          p_y,
+                          p_lse,
+                          p_ygrad,
+                          p_qgrad,
+                          p_kgrad,
+                          p_vgrad,
+                          {}, // std::array<void*, 1> p_acc0_biases;
+                          {}, // std::array<void*, 1> p_acc1_biases;
+                          problem_descs,
+                          QKVElementOp{},
+                          QKVElementOp{},
+                          Scale{alpha},
+                          QKVElementOp{},
+                          YElementOp{},
+                          p_drop,
+                          std::tuple<unsigned long long, unsigned long long>(seed, offset));
 
     DeviceMem problem_desc_workspace(gemm.GetWorkSpaceSize(&argument));
 
@@ -640,13 +702,6 @@ int run(int argc, char* argv[])
         return 0;
     }
 
-    if(!gemm.IsSupportedArgument(argument))
-    {
-        std::cout << gemm.GetTypeString() << " does not support this problem" << std::endl;
-
-        return 0;
-    }
-
     float ave_time = invoker.Run(argument, StreamConfig{nullptr, time_kernel});
 
     float tflops = static_cast<float>(flop) / 1.E9 / ave_time;

include/ck/tensor_operation/gpu/block/blockwise_dropout.hpp

@@ -16,12 +16,15 @@ struct BlockwiseDropout
     static constexpr index_t MRepeat = ThreadSliceDesc_M_K{}.GetLength(I0);
     static constexpr index_t KRepeat = ThreadSliceDesc_M_K{}.GetLength(I1);
 
-    template <typename CThreadBuffer>
+    template <typename CThreadBuffer, bool using_sign_bit = false>
     __host__ __device__ void ApplyDropout(CThreadBuffer& in_thread_buf, ck::philox ph)
     {
 
         auto execute_dropout = [&](bool keep, DataType val) {
-            return keep ? val * p_dropout_rescale : float(0);
+            if constexpr(using_sign_bit)
+                return keep ? val : -val;
+            else
+                return keep ? val * p_dropout_rescale : float(0);
         };
 
         constexpr int tmp_size = MRepeat * KRepeat;
@@ -47,6 +50,42 @@ struct BlockwiseDropout
         });
     }
 
+    template <typename CThreadBuffer, typename ZThreadBuffer, bool using_sign_bit = false>
+    __host__ __device__ void
+    ApplyDropout(CThreadBuffer& in_thread_buf, ck::philox ph, ZThreadBuffer& z_thread_buf)
+    {
+
+        auto execute_dropout = [&](bool keep, DataType val) {
+            if constexpr(using_sign_bit)
+                return keep ? val : -val;
+            else
+                return keep ? val * p_dropout_rescale : float(0);
+        };
+
+        constexpr int tmp_size = MRepeat * KRepeat;
+
+        int philox_calls = tmp_size / 8;
+
+        ushort tmp[tmp_size];
+        for(int i = 0; i < philox_calls; i++)
+        {
+            ph.get_random_8x16((tmp + i * 8));
+        }
+
+        block_sync_lds();
+
+        int tmp_index = 0;
+        static_for<0, MRepeat, 1>{}([&](auto iM) {
+            static_for<0, KRepeat, 1>{}([&](auto iK) {
+                auto offset = Number<ThreadSliceDesc_M_K{}.CalculateOffset(make_tuple(iM, iK))>{};
+                in_thread_buf(offset) =
+                    execute_dropout(tmp[tmp_index] < p_dropout_16bits, in_thread_buf(offset));
+                z_thread_buf(offset) = tmp[tmp_index];
+                tmp_index            = tmp_index + 1;
+            });
+        });
+    }
+
     ushort p_dropout_16bits;
     DataType p_dropout_rescale;
 };

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

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2022, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include <iostream>
+#include <sstream>
+
+#include "ck/utility/common_header.hpp"
+#include "ck/utility/philox_rand.hpp"
+#include "ck/tensor_description/tensor_descriptor.hpp"
+#include "ck/tensor_description/tensor_descriptor_helper.hpp"
+// #include "ck/tensor_operation/gpu/device/device_batched_multihead_attention_backward.hpp" // TODO
+#include "ck/tensor_operation/gpu/device/device_base.hpp"
+#include "ck/tensor_operation/gpu/device/gemm_specialization.hpp"
+#include "ck/tensor_operation/gpu/device/masking_specialization.hpp"
+#include "ck/tensor_operation/gpu/device/matrix_padder.hpp"
+#include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
+#include "ck/tensor_operation/gpu/grid/gridwise_batched_multihead_attention_backward_xdl_cshuffle_v2.hpp"
+#include "ck/tensor_operation/operator_transform/transform_contraction_to_gemm.hpp"
+#include "ck/host_utility/device_prop.hpp"
+#include "ck/host_utility/kernel_launch.hpp"
+
+#include "ck/library/utility/host_tensor.hpp"
+
+namespace ck {
+namespace tensor_operation {
+namespace device {
+
+template <typename GridwiseGemm,
+          typename DataType,
+          typename ZDataType,
+          typename LSEDataType,
+          typename AElementwiseOperation,
+          typename BElementwiseOperation,
+          typename AccElementwiseOperation,
+          typename B1ElementwiseOperation,
+          typename CElementwiseOperation,
+          typename AGridDesc_AK0_M_AK1,
+          typename BGridDesc_BK0_N_BK1,
+          typename ZGridDescriptor_M0_N0_M1_N1_M2_N2_M3_N3_N4_N5,
+          typename B1GridDesc_BK0_N_BK1,
+          typename YGridDescriptor_MBlock_MPerBlock_OBlock_OPerBlock,
+          typename LSEGridDescriptor_M,
+          typename VGradGridDescriptor_N_O,
+          typename YGradGridDesc_M0_O_M1,
+          typename Block2CTileMap,
+          typename ComputeBasePtrOfStridedBatch,
+          typename C0MatrixMask,
+          bool HasMainKBlockLoop>
+__global__ void
+#if CK_USE_LAUNCH_BOUNDS
+    __launch_bounds__(CK_MAX_THREAD_PER_BLOCK, CK_MIN_BLOCK_PER_CU)
+#endif
+        kernel_batched_gemm_softmax_gemm_xdl_cshuffle_v2(
+            const DataType* __restrict__ p_a_grid,
+            const DataType* __restrict__ p_b_grid,
+            ZDataType* __restrict__ p_z_grid,
+            const DataType* __restrict__ p_b1_grid,
+            const DataType* __restrict__ p_c_grid,
+            const LSEDataType* __restrict__ p_lse_grid,
+            const DataType* __restrict__ p_ygrad_grid,
+            DataType* __restrict__ p_qgrad_grid,
+            DataType* __restrict__ p_kgrad_grid,
+            DataType* __restrict__ p_vgrad_grid,
+            const AElementwiseOperation a_element_op,
+            const BElementwiseOperation b_element_op,
+            const AccElementwiseOperation acc_element_op,
+            const B1ElementwiseOperation b1_element_op,
+            const CElementwiseOperation c_element_op,
+            const AGridDesc_AK0_M_AK1 a_grid_desc_ak0_m_ak1,
+            const BGridDesc_BK0_N_BK1 b_grid_desc_bk0_n_bk1,
+            const ZGridDescriptor_M0_N0_M1_N1_M2_N2_M3_N3_N4_N5
+                c_grid_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5,
+            const B1GridDesc_BK0_N_BK1 b1_grid_desc_bk0_n_bk1,
+            const YGridDescriptor_MBlock_MPerBlock_OBlock_OPerBlock
+                c_grid_desc_mblock_mperblock_nblock_nperblock,
+            const LSEGridDescriptor_M lse_grid_desc_m,
+            const VGradGridDescriptor_N_O vgrad_grid_desc_n_o,
+            const YGradGridDesc_M0_O_M1 ygrad_grid_desc_m0_o_m1,
+            const Block2CTileMap block_2_ctile_map,
+            const index_t batch_count,
+            const ComputeBasePtrOfStridedBatch compute_base_ptr_of_batch,
+            const C0MatrixMask c0_matrix_mask,
+            const float p_dropout,
+            const unsigned long long seed,
+            const unsigned long long offset)
+{
+#if(!defined(__HIP_DEVICE_COMPILE__) || defined(__gfx908__) || defined(__gfx90a__))
+    __shared__ char p_shared[GridwiseGemm::GetSharedMemoryNumberOfByte()];
+    const index_t num_blocks_per_batch =
+        __builtin_amdgcn_readfirstlane(get_grid_size() / batch_count);
+    const index_t g_idx = __builtin_amdgcn_readfirstlane(get_block_1d_id() / num_blocks_per_batch);
+
+    // NOTE: assumes QKVY has the same layout as dQ/dK/dV/dY therefore being able to reuse batch
+    // offsets
+    const long_index_t a_batch_offset = __builtin_amdgcn_readfirstlane(
+        static_cast<long_index_t>(compute_base_ptr_of_batch.GetABasePtr(g_idx)));
+    const long_index_t b_batch_offset = __builtin_amdgcn_readfirstlane(
+        static_cast<long_index_t>(compute_base_ptr_of_batch.GetBBasePtr(g_idx)));
+    const long_index_t z_batch_offset = __builtin_amdgcn_readfirstlane(
+        static_cast<long_index_t>(compute_base_ptr_of_batch.GetZBasePtr(g_idx)));
+    const long_index_t b1_batch_offset = __builtin_amdgcn_readfirstlane(
+        static_cast<long_index_t>(compute_base_ptr_of_batch.GetB1BasePtr(g_idx)));
+    const long_index_t c_batch_offset = __builtin_amdgcn_readfirstlane(
+        static_cast<long_index_t>(compute_base_ptr_of_batch.GetCBasePtr(g_idx)));
+    const long_index_t lse_batch_offset = __builtin_amdgcn_readfirstlane(
+        static_cast<long_index_t>(compute_base_ptr_of_batch.GetLSEBasePtr(g_idx)));
+
+    const index_t global_thread_id = get_thread_global_1d_id();
+    ck::philox ph(seed, global_thread_id, offset);
+    ZDataType* z_matrix_ptr = (p_z_grid == nullptr ? nullptr : p_z_grid + z_batch_offset);
+
+    GridwiseGemm::template Run<HasMainKBlockLoop>(p_a_grid + a_batch_offset,
+                                                  p_b_grid + b_batch_offset,
+                                                  z_matrix_ptr,
+                                                  p_b1_grid + b1_batch_offset,
+                                                  p_c_grid + c_batch_offset,
+                                                  p_lse_grid + lse_batch_offset,
+                                                  p_ygrad_grid + c_batch_offset,
+                                                  p_qgrad_grid + a_batch_offset,
+                                                  p_kgrad_grid + b_batch_offset,
+                                                  p_vgrad_grid + b1_batch_offset,
+                                                  p_shared,
+                                                  a_element_op,
+                                                  b_element_op,
+                                                  acc_element_op,
+                                                  b1_element_op,
+                                                  c_element_op,
+                                                  a_grid_desc_ak0_m_ak1,
+                                                  b_grid_desc_bk0_n_bk1,
+                                                  c_grid_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5,
+                                                  b1_grid_desc_bk0_n_bk1,
+                                                  c_grid_desc_mblock_mperblock_nblock_nperblock,
+                                                  lse_grid_desc_m,
+                                                  vgrad_grid_desc_n_o,
+                                                  ygrad_grid_desc_m0_o_m1,
+                                                  block_2_ctile_map,
+                                                  c0_matrix_mask,
+                                                  p_dropout,
+                                                  ph);
+#else
+    ignore = p_a_grid;
+    ignore = p_b_grid;
+    ignore = p_b1_grid;
+    ignore = p_c_grid;
+    ignore = a_element_op;
+    ignore = b_element_op;
+    ignore = acc_element_op;
+    ignore = b1_element_op;
+    ignore = c_element_op;
+    ignore = a_grid_desc_ak0_m_ak1;
+    ignore = b_grid_desc_bk0_n_bk1;
+    ignore = b1_grid_desc_bk0_n_bk1;
+    ignore = c_grid_desc_mblock_mperblock_nblock_nperblock;
+    ignore = block_2_ctile_map;
+    ignore = batch_count;
+    ignore = compute_base_ptr_of_batch;
+    ignore = c0_matrix_mask;
+#endif // end of if (defined(__gfx908__) || defined(__gfx90a__))
+}
+
+// Computes C = A * B0 * B1
+//              ^^^^^^ (Acc0)
+//              ^^^^^^^^^^^ (Acc1)
+template <index_t NumDimG,
+          index_t NumDimM,
+          index_t NumDimN,
+          index_t NumDimK,
+          index_t NumDimO, // NumDimGemm1N
+          typename DataType,
+          typename ZDataType,
+          typename LSEDataType,
+          typename Acc0BiasDataType,
+          typename Acc1BiasDataType,
+          typename GemmAccDataType,
+          typename CShuffleDataType,
+          typename AElementwiseOperation,
+          typename BElementwiseOperation,
+          typename AccElementwiseOperation,
+          typename B1ElementwiseOperation,
+          typename CElementwiseOperation,
+          GemmSpecialization GemmSpec,
+          TensorSpecialization ASpec,
+          TensorSpecialization BSpec,
+          TensorSpecialization B1Spec,
+          TensorSpecialization CSpec,
+          index_t NumGemmKPrefetchStage,
+          index_t BlockSize,
+          index_t MPerBlock,
+          index_t NPerBlock, // Gemm0NPerBlock
+          index_t KPerBlock, // Gemm0KPerBlock
+          index_t Gemm1NPerBlock,
+          index_t Gemm1KPerBlock,
+          index_t AK1,
+          index_t BK1,
+          index_t B1K1,
+          index_t MPerXDL,
+          index_t NPerXDL,
+          index_t MXdlPerWave,
+          index_t NXdlPerWave,
+          index_t Gemm1NXdlPerWave,
+          typename ABlockTransferThreadClusterLengths_AK0_M_AK1,
+          typename ABlockTransferThreadClusterArrangeOrder,
+          typename ABlockTransferSrcAccessOrder,
+          index_t ABlockTransferSrcVectorDim,
+          index_t ABlockTransferSrcScalarPerVector,
+          index_t ABlockTransferDstScalarPerVector_AK1,
+          bool ABlockLdsExtraM,
+          typename BBlockTransferThreadClusterLengths_BK0_N_BK1,
+          typename BBlockTransferThreadClusterArrangeOrder,
+          typename BBlockTransferSrcAccessOrder,
+          index_t BBlockTransferSrcVectorDim,
+          index_t BBlockTransferSrcScalarPerVector,
+          index_t BBlockTransferDstScalarPerVector_BK1,
+          bool BBlockLdsExtraN,
+          typename B1BlockTransferThreadClusterLengths_BK0_N_BK1,
+          typename B1BlockTransferThreadClusterArrangeOrder,
+          typename B1BlockTransferSrcAccessOrder,
+          index_t B1BlockTransferSrcVectorDim,
+          index_t B1BlockTransferSrcScalarPerVector,
+          index_t B1BlockTransferDstScalarPerVector_BK1,
+          bool B1BlockLdsExtraN,
+          index_t CShuffleMXdlPerWavePerShuffle,
+          index_t CShuffleNXdlPerWavePerShuffle,
+          typename CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock,
+          index_t CShuffleBlockTransferScalarPerVector_NPerBlock,
+          MaskingSpecialization MaskingSpec,
+          LoopScheduler LoopSched = LoopScheduler::Default>
+struct DeviceBatchedMultiheadAttentionBackward_Train_Xdl_CShuffle
+    : public BaseOperator // TODO inherit atten bwd op once API stablizes
+{
+    static_assert(NumDimG > 0 && NumDimM > 0 && NumDimN > 0 && NumDimK > 0 && NumDimO > 0,
+                  "Number of dimension must be greater than 0");
+
+    static constexpr index_t NumAcc0Bias = Acc0BiasDataType::Size();
+    static constexpr index_t NumAcc1Bias = Acc1BiasDataType::Size();
+
+    // TODO: implement bias combination
+    static_assert(NumAcc0Bias == 0 && NumAcc0Bias == 0, "Bias addition is unimplemented");
+
+#if 0
+    // TODO: use alias
+    static constexpr index_t NumDimGemm0M = NumDimM;
+    static constexpr index_t NumDimGemm0N = NumDimN;
+    static constexpr index_t NumDimGemm0K = NumDimK;
+    static constexpr index_t NumDimGemm1M = NumDimM;
+    static constexpr index_t NumDimGemm1N = NumDimO;
+    static constexpr index_t NumDimGemm1K = NumDimN;
+#endif
+
+    using DeviceOp = DeviceBatchedMultiheadAttentionBackward_Train_Xdl_CShuffle;
+
+    static constexpr auto I0 = Number<0>{};
+    static constexpr auto I1 = Number<1>{};
+    static constexpr auto I2 = Number<2>{};
+
+    static constexpr index_t Q_K1 = 8;
+    static constexpr index_t K_K1 = 8;
+    static constexpr index_t V_N1 = 2;
+
+    static constexpr index_t Q_M1 = 2;
+    static constexpr index_t K_N1 = 2;
+    static constexpr index_t V_O1 = 8;
+    static constexpr index_t Y_O1 = 8;
+    static constexpr index_t Y_M1 = 2;
+
+    static constexpr auto padder = GemmGemmPadder<GemmSpec,
+                                                  Number<MPerBlock>,
+                                                  Number<NPerBlock>,
+                                                  Number<KPerBlock>,
+                                                  Number<Gemm1NPerBlock>>{};
+
+    using Transform = TransformBatchedContractionContractionToBatchedGemmGemm<
+        Sequence<NumDimG, NumDimM, NumDimN, NumDimK, NumDimO>,
+        Sequence<MPerBlock, NPerBlock, KPerBlock, Gemm1NPerBlock>,
+        GemmSpec,
+        ASpec,
+        BSpec,
+        B1Spec,
+        CSpec>;
+
+    /*
+    Descriptors for inputs:
+
+      Q, K, V, Y, dY, per-row softmax stats
+
+    Descriptors for outputs:
+
+      dQ, dK, dV
+
+    */
+
+    // Q in Gemm A position
+    static auto MakeAGridDescriptor_AK0_M_AK1(const std::vector<index_t>& a_gs_ms_ks_lengths_vec,
+                                              const std::vector<index_t>& a_gs_ms_ks_strides_vec)
+    {
+        return Transform::MakeAGridDescriptor_AK0_M_AK1(
+            Transform::MakeAGridDescriptor_M_K(a_gs_ms_ks_lengths_vec, a_gs_ms_ks_strides_vec),
+            Number<AK1>{});
+    }
+
+    // K in Gemm B0 position
+    static auto MakeBGridDescriptor_BK0_N_BK1(const std::vector<index_t>& b_gs_ns_ks_lengths_vec,
+                                              const std::vector<index_t>& b_gs_ns_ks_strides_vec)
+    {
+        return Transform::MakeB0GridDescriptor_BK0_N_BK1(
+            Transform::MakeB0GridDescriptor_N_K(b_gs_ns_ks_lengths_vec, b_gs_ns_ks_strides_vec),
+            Number<BK1>{});
+    }
+
+    // V in Gemm B1 position
+    static auto
+    MakeB1GridDescriptor_BK0_N_BK1(const std::vector<index_t>& b1_gs_gemm1ns_gemm1ks_lengths_vec,
+                                   const std::vector<index_t>& b1_gs_gemm1ns_gemm1ks_strides_vec)
+    {
+        return Transform::MakeB1GridDescriptor_BK0_N_BK1(
+            Transform::MakeB1GridDescriptor_N_K(b1_gs_gemm1ns_gemm1ks_lengths_vec,
+                                                b1_gs_gemm1ns_gemm1ks_strides_vec),
+            Number<B1K1>{});
+    }
+
+    //
+    // dV = P^T * dY
+    //
+
+    // VGrad in Gemm C position
+    static auto MakeVGradGridDescriptor_N_O(const std::vector<index_t>& v_gs_os_ns_lengths_vec,
+                                            const std::vector<index_t>& v_gs_os_ns_strides_vec)
+    {
+        // v_gs_os_ns -> vgrad_gs_ns_os. O dims last because output is row-major.
+        // Here directly rearrange lengths/strides before constructing tensor descriptor to reduce
+        // transformation overhead
+        // TODO: This will be much easier when inputs are Gs, Ms, Ns, Os. So there's no need to
+        // extract subsequence and shuffle them.
+        const index_t num_dims = NumDimG + NumDimN + NumDimO;
+
+        // 0, 1, .. NumDimG - 1
+        std::vector<index_t> gs_ids(NumDimG);
+        std::iota(gs_ids.begin(), gs_ids.end(), 0);
+
+        // NumDimG, NumDimG + 1, ... NumDimG + NumDimO - 1
+        std::vector<index_t> os_ids(NumDimO);
+        std::iota(os_ids.begin(), os_ids.end(), NumDimG);
+
+        // NumDimG + NumDimO, NumDimG + NumDimO + 1, ... NumDimG + NumDimO + NumDimN - 1
+        std::vector<index_t> ns_ids(NumDimN);
+        std::iota(ns_ids.begin(), ns_ids.end(), NumDimG + NumDimO);
+
+        std::vector<index_t> ids_old2new;
+        ids_old2new.insert(ids_old2new.end(), gs_ids.begin(), gs_ids.end());
+        ids_old2new.insert(ids_old2new.end(), ns_ids.begin(), ns_ids.end());
+        ids_old2new.insert(ids_old2new.end(), os_ids.begin(), os_ids.end());
+
+        std::vector<index_t> v_gs_ns_os_lengths_vec(num_dims), v_gs_ns_os_strides_vec(num_dims);
+        for(int i = 0; i < num_dims; i++)
+        {
+            index_t id_new            = ids_old2new[i];
+            v_gs_ns_os_lengths_vec[i] = v_gs_os_ns_lengths_vec[id_new];
+            v_gs_ns_os_strides_vec[i] = v_gs_os_ns_strides_vec[id_new];
+        }
+
+        const auto vgrad_desc_nraw_oraw =
+            MakeGridDescriptorPair<NumDimG, NumDimN, NumDimO, TensorSpecialization::Default>(
+                v_gs_ns_os_lengths_vec, v_gs_ns_os_strides_vec)
+                .second;
+
+        return PadTensorDescriptor(vgrad_desc_nraw_oraw,
+                                   make_tuple(NPerBlock, Gemm1NPerBlock),
+                                   Sequence<padder.PadN, padder.PadO>{});
+    }
+
+    template <typename YGridDesc_M_O>
+    static auto MakeYGradGridDescriptor_M0_O_M1(const YGridDesc_M_O& ygrad_grid_desc_m_o)
+    {
+        const auto M = ygrad_grid_desc_m_o.GetLength(I0);
+        const auto O = ygrad_grid_desc_m_o.GetLength(I1);
+
+        const auto Y_M0 = M / Y_M1;
+
+        return transform_tensor_descriptor(
+            ygrad_grid_desc_m_o,
+            make_tuple(make_unmerge_transform(make_tuple(Y_M0, Y_M1)),
+                       make_pass_through_transform(O)),
+            make_tuple(Sequence<0>{}, Sequence<1>{}),
+            make_tuple(Sequence<0, 2>{}, Sequence<1>{}));
+    }
+
+    //
+    // dP = dY * V^T
+    //
+
+    // YGrad in Gemm A position
+    static auto MakeYGradGridDescriptor_O0_M_O1(const std::vector<index_t>& y_gs_ms_os_lengths_vec,
+                                                const std::vector<index_t>& y_gs_ms_os_strides_vec)
+    {
+        return Transform::MakeAGridDescriptor_AK0_M_AK1(
+            Transform::MakeAGridDescriptor_M_K(y_gs_ms_os_lengths_vec, y_gs_ms_os_strides_vec),
+            Number<Y_O1>{});
+    }
+
+    // V in Gemm B position
+    static auto MakeVGridDescriptor_O0_N_O1(const std::vector<index_t>& v_gs_os_ns_lengths_vec,
+                                            const std::vector<index_t>& v_gs_os_ns_strides_vec)
+    {
+        // v_gs_os_ns -> vgrad_gs_ns_os. O dims last because output is row-major.
+        // Here directly rearrange lengths/strides before constructing tensor descriptor to reduce
+        // transformation overhead
+        // TODO: This will be much easier when inputs are Gs, Ms, Ns, Os. So there's no need to
+        // extract subsequence and shuffle them.
+        const index_t num_dims = NumDimG + NumDimN + NumDimO;
+
+        // 0, 1, .. NumDimG - 1
+        std::vector<index_t> gs_ids(NumDimG);
+        std::iota(gs_ids.begin(), gs_ids.end(), 0);
+
+        // NumDimG, NumDimG + 1, ... NumDimG + NumDimO - 1
+        std::vector<index_t> os_ids(NumDimO);
+        std::iota(os_ids.begin(), os_ids.end(), NumDimG);
+
+        // NumDimG + NumDimO, NumDimG + NumDimO + 1, ... NumDimG + NumDimO + NumDimN - 1
+        std::vector<index_t> ns_ids(NumDimN);
+        std::iota(ns_ids.begin(), ns_ids.end(), NumDimG + NumDimO);
+
+        std::vector<index_t> ids_old2new;
+        ids_old2new.insert(ids_old2new.end(), gs_ids.begin(), gs_ids.end());
+        ids_old2new.insert(ids_old2new.end(), ns_ids.begin(), ns_ids.end());
+        ids_old2new.insert(ids_old2new.end(), os_ids.begin(), os_ids.end());
+
+        std::vector<index_t> v_gs_ns_os_lengths_vec(num_dims), v_gs_ns_os_strides_vec(num_dims);
+        for(int i = 0; i < num_dims; i++)
+        {
+            index_t id_new            = ids_old2new[i];
+            v_gs_ns_os_lengths_vec[i] = v_gs_os_ns_lengths_vec[id_new];
+            v_gs_ns_os_strides_vec[i] = v_gs_os_ns_strides_vec[id_new];
+        }
+
+        const auto v_grid_desc_nraw_oraw =
+            MakeGridDescriptorPair<NumDimG, NumDimN, NumDimO, TensorSpecialization::Default>(
+                v_gs_ns_os_lengths_vec, v_gs_ns_os_strides_vec)
+                .second;
+
+        const auto v_grid_desc_n_o = PadTensorDescriptor(v_grid_desc_nraw_oraw,
+                                                         make_tuple(NPerBlock, Gemm1NPerBlock),
+                                                         Sequence<padder.PadN, padder.PadO>{});
+
+        // N_O to O0_N_O1; to refactor
+        return Transform::MakeB0GridDescriptor_BK0_N_BK1(v_grid_desc_n_o, Number<V_O1>{});
+    }
+
+    // Z in Gemm0 C position
+    static auto MakeZGridDescriptor_M_N(const std::vector<index_t>& z_gs_ms_ns_lengths_vec,
+                                        const std::vector<index_t>& z_gs_ms_ns_strides_vec)
+    {
+        return Transform::MakeCGridDescriptor_M_N(z_gs_ms_ns_lengths_vec, z_gs_ms_ns_strides_vec);
+    }
+    //
+    // dS_i_j = P_i_j .* (dP_i_j - dY_i dot Y_i)
+    //
+
+    //
+    // dQ = alpha * dS * K
+    //
+
+    // QGrad in Gemm C position
+    static auto MakeQGradGridDescriptor_M_K(const std::vector<index_t>& q_gs_ms_ks_lengths_vec,
+                                            const std::vector<index_t>& q_gs_ms_ks_strides_vec)
+    {
+        return Transform::MakeCGridDescriptor_M_N(q_gs_ms_ks_lengths_vec, q_gs_ms_ks_strides_vec);
+    }
+
+    //
+    // dK = alpha * dS^T * Q
+    //
+
+    // KGrad in Gemm C position
+    static auto MakeKGradGridDescriptor_N_K(const std::vector<index_t>& k_gs_ns_ks_lengths_vec,
+                                            const std::vector<index_t>& k_gs_ns_ks_strides_vec)
+    {
+        return Transform::MakeCGridDescriptor_M_N(k_gs_ns_ks_lengths_vec, k_gs_ns_ks_strides_vec);
+    }
+
+    static auto MakeLSEGridDescriptor_M(index_t MRaw)
+    {
+        const auto lse_grid_desc_mraw = make_naive_tensor_descriptor_packed(make_tuple(MRaw));
+
+        const auto M    = math::integer_divide_ceil(MRaw, MPerBlock) * MPerBlock;
+        const auto MPad = M - MRaw;
+
+        if constexpr(GemmSpec == GemmSpecialization::MPadding ||
+                     GemmSpec == GemmSpecialization::MNPadding ||
+                     GemmSpec == GemmSpecialization::MKPadding ||
+                     GemmSpec == GemmSpecialization::MNKPadding)
+        {
+            // pad M
+            return transform_tensor_descriptor(lse_grid_desc_mraw,
+                                               make_tuple(make_right_pad_transform(MRaw, MPad)),
+                                               make_tuple(Sequence<0>{}),
+                                               make_tuple(Sequence<0>{}));
+        }
+        else
+        {
+            // not pad M
+            return lse_grid_desc_mraw;
+        }
+    }
+
+    using AGridDesc_AK0_M_AK1  = decltype(MakeAGridDescriptor_AK0_M_AK1({}, {}));
+    using BGridDesc_BK0_N_BK1  = decltype(MakeBGridDescriptor_BK0_N_BK1({}, {}));
+    using B1GridDesc_BK0_N_BK1 = decltype(MakeB1GridDescriptor_BK0_N_BK1({}, {}));
+    using YGridDesc_M_O        = decltype(Transform::MakeCGridDescriptor_M_N({}, {}));
+    using LSEGridDesc_M        = decltype(MakeLSEGridDescriptor_M(1));
+    using AGridDesc_G_M_K      = decltype(Transform::MakeAGridDescriptor_G_M_K({}, {}));
+    using BGridDesc_G_N_K      = decltype(Transform::MakeB0GridDescriptor_G_N_K({}, {}));
+    using B1GridDesc_G_N_K     = decltype(Transform::MakeB1GridDescriptor_G_N_K({}, {}));
+    using CGridDesc_G_M_N      = decltype(Transform::MakeCGridDescriptor_G_M_N({}, {}));
+    using ZGridDesc_G_M_N      = decltype(Transform::MakeCGridDescriptor_G_M_N({}, {}));
+
+    using VGradGridDesc_N_O     = decltype(MakeVGradGridDescriptor_N_O({}, {}));
+    using YGradGridDesc_M0_O_M1 = decltype(MakeYGradGridDescriptor_M0_O_M1(YGridDesc_M_O{}));
+    using ZGridDesc_M_N         = decltype(MakeZGridDescriptor_M_N({}, {}));
+
+    constexpr static auto make_MaskOutPredicate()
+    {
+        if constexpr(MaskingSpec == MaskingSpecialization::MaskDisabled)
+        {
+            return MaskDisabledPredicate{};
+        }
+        else if constexpr(MaskingSpec == MaskingSpecialization::MaskOutUpperTriangle)
+        {
+            return MaskOutUpperTrianglePredicate{};
+        }
+    }
+    using C0MatrixMask = C0MatrixMask_impl<decltype(make_MaskOutPredicate())>;
+
+    struct ComputeBasePtrOfStridedBatch
+    {
+        ComputeBasePtrOfStridedBatch(const AGridDesc_G_M_K& a_grid_desc_g_m_k,
+                                     const BGridDesc_G_N_K& b_grid_desc_g_n_k,
+                                     const ZGridDesc_G_M_N& z_grid_desc_g_m_n,
+                                     const B1GridDesc_G_N_K& b1_grid_desc_g_n_k,
+                                     const CGridDesc_G_M_N& c_grid_desc_g_m_n,
+                                     index_t BatchStrideLSE)
+            : a_grid_desc_g_m_k_(a_grid_desc_g_m_k),
+              b_grid_desc_g_n_k_(b_grid_desc_g_n_k),
+              z_grid_desc_g_m_n_(z_grid_desc_g_m_n),
+              b1_grid_desc_g_n_k_(b1_grid_desc_g_n_k),
+              c_grid_desc_g_m_n_(c_grid_desc_g_m_n),
+              BatchStrideLSE_(BatchStrideLSE)
+        {
+        }
+
+        __host__ __device__ constexpr long_index_t GetABasePtr(index_t g_idx) const
+        {
+            return a_grid_desc_g_m_k_.CalculateOffset(make_multi_index(g_idx, 0, 0));
+        }
+
+        __host__ __device__ constexpr long_index_t GetBBasePtr(index_t g_idx) const
+        {
+            return b_grid_desc_g_n_k_.CalculateOffset(make_multi_index(g_idx, 0, 0));
+        }
+
+        __host__ __device__ constexpr long_index_t GetZBasePtr(index_t g_idx) const
+        {
+            return z_grid_desc_g_m_n_.CalculateOffset(make_multi_index(g_idx, 0, 0));
+        }
+
+        __host__ __device__ constexpr long_index_t GetB1BasePtr(index_t g_idx) const
+        {
+            return b1_grid_desc_g_n_k_.CalculateOffset(make_multi_index(g_idx, 0, 0));
+        }
+
+        __host__ __device__ constexpr long_index_t GetCBasePtr(index_t g_idx) const
+        {
+            return c_grid_desc_g_m_n_.CalculateOffset(make_multi_index(g_idx, 0, 0));
+        }
+
+        __host__ __device__ constexpr long_index_t GetLSEBasePtr(index_t g_idx) const
+        {
+            return g_idx * static_cast<long_index_t>(BatchStrideLSE_);
+        }
+
+        private:
+        AGridDesc_G_M_K a_grid_desc_g_m_k_;
+        BGridDesc_G_N_K b_grid_desc_g_n_k_;
+        ZGridDesc_G_M_N z_grid_desc_g_m_n_;
+        B1GridDesc_G_N_K b1_grid_desc_g_n_k_;
+        CGridDesc_G_M_N c_grid_desc_g_m_n_;
+
+        index_t BatchStrideLSE_;
+    };
+
+    // GridwiseGemm
+    using GridwiseGemm = GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2<
+        DataType, // TODO: distinguish A/B datatype
+        LSEDataType,
+        GemmAccDataType,
+        CShuffleDataType,
+        AElementwiseOperation,
+        BElementwiseOperation,
+        AccElementwiseOperation,
+        B1ElementwiseOperation,
+        CElementwiseOperation,
+        InMemoryDataOperationEnum::Set,
+        AGridDesc_AK0_M_AK1,
+        BGridDesc_BK0_N_BK1,
+        ZGridDesc_M_N,
+        B1GridDesc_BK0_N_BK1,
+        YGridDesc_M_O,
+        LSEGridDesc_M,
+        NumGemmKPrefetchStage,
+        BlockSize,
+        MPerBlock,
+        NPerBlock,
+        KPerBlock,
+        Gemm1NPerBlock,
+        Gemm1KPerBlock,
+        AK1,
+        BK1,
+        B1K1,
+        MPerXDL,
+        NPerXDL,
+        MXdlPerWave,
+        NXdlPerWave,
+        Gemm1NXdlPerWave,
+        ABlockTransferThreadClusterLengths_AK0_M_AK1,
+        ABlockTransferThreadClusterArrangeOrder,
+        ABlockTransferSrcAccessOrder,
+        ABlockTransferSrcVectorDim,
+        ABlockTransferSrcScalarPerVector,
+        ABlockTransferDstScalarPerVector_AK1,
+        true,
+        ABlockLdsExtraM,
+        BBlockTransferThreadClusterLengths_BK0_N_BK1,
+        BBlockTransferThreadClusterArrangeOrder,
+        BBlockTransferSrcAccessOrder,
+        BBlockTransferSrcVectorDim,
+        BBlockTransferSrcScalarPerVector,
+        BBlockTransferDstScalarPerVector_BK1,
+        true,
+        BBlockLdsExtraN,
+        B1BlockTransferThreadClusterLengths_BK0_N_BK1,
+        B1BlockTransferThreadClusterArrangeOrder,
+        B1BlockTransferSrcAccessOrder,
+        B1BlockTransferSrcVectorDim,
+        B1BlockTransferSrcScalarPerVector,
+        B1BlockTransferDstScalarPerVector_BK1,
+        false,
+        B1BlockLdsExtraN,
+        CShuffleMXdlPerWavePerShuffle,
+        CShuffleNXdlPerWavePerShuffle,
+        CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock,
+        CShuffleBlockTransferScalarPerVector_NPerBlock,
+        LoopSched,
+        Transform::matrix_padder.PadN,
+        MaskingSpec == MaskingSpecialization::MaskOutUpperTriangle>;
+
+    // Argument
+    struct Argument : public BaseArgument
+    {
+        Argument(
+            const DataType* p_a_grid,
+            const DataType* p_b_grid,
+            ZDataType* p_z_grid,
+            const DataType* p_b1_grid,
+            const DataType* p_c_grid, // for dS
+            const LSEDataType* p_lse_grid,
+            const DataType* p_ygrad_grid,
+            DataType* p_qgrad_grid,
+            DataType* p_kgrad_grid,
+            DataType* p_vgrad_grid,
+            const std::array<void*, NumAcc0Bias> p_acc0_biases,
+            const std::array<void*, NumAcc1Bias> p_acc1_biases,
+            const std::vector<index_t>& a_gs_ms_ks_lengths,
+            const std::vector<index_t>& a_gs_ms_ks_strides,
+            const std::vector<index_t>& b_gs_ns_ks_lengths,
+            const std::vector<index_t>& b_gs_ns_ks_strides,
+            const std::vector<index_t>& z_gs_ms_ns_lengths,
+            const std::vector<index_t>& z_gs_ms_ns_strides,
+            const std::vector<index_t>& b1_gs_gemm1ns_gemm1ks_lengths, // b1_gs_os_ns_lengths
+            const std::vector<index_t>& b1_gs_gemm1ns_gemm1ks_strides, // b1_gs_os_ns_strides
+            const std::vector<index_t>& c_gs_ms_gemm1ns_lengths,       // c_gs_ms_os_lengths
+            const std::vector<index_t>& c_gs_ms_gemm1ns_strides,       // c_gs_ms_os_strides
+            const std::vector<index_t>& lse_gs_ms_lengths,
+            const std::array<std::vector<ck::index_t>, NumAcc0Bias> acc0_biases_gs_ms_ns_lengths,
+            const std::array<std::vector<ck::index_t>, NumAcc0Bias> acc0_biases_gs_ms_ns_strides,
+            const std::array<std::vector<ck::index_t>, NumAcc1Bias>
+                acc1_biases_gs_ms_gemm1ns_lengths, // acc1_biases_gs_ms_os_lengths
+            const std::array<std::vector<ck::index_t>, NumAcc1Bias>
+                acc1_biases_gs_ms_gemm1ns_strides, // acc1_biases_gs_ms_os_strides
+            AElementwiseOperation a_element_op,
+            BElementwiseOperation b_element_op,
+            AccElementwiseOperation acc_element_op,
+            B1ElementwiseOperation b1_element_op,
+            CElementwiseOperation c_element_op,
+            float p_drop,
+            std::tuple<unsigned long long, unsigned long long> seeds)
+            : p_a_grid_{p_a_grid},
+              p_b_grid_{p_b_grid},
+              p_z_grid_{p_z_grid},
+              p_b1_grid_{p_b1_grid},
+              p_c_grid_{p_c_grid},
+              p_lse_grid_{p_lse_grid},
+              p_ygrad_grid_{p_ygrad_grid},
+              p_qgrad_grid_{p_qgrad_grid},
+              p_kgrad_grid_{p_kgrad_grid},
+              p_vgrad_grid_{p_vgrad_grid},
+              a_grid_desc_ak0_m_ak1_{
+                  DeviceOp::MakeAGridDescriptor_AK0_M_AK1(a_gs_ms_ks_lengths, a_gs_ms_ks_strides)},
+              b_grid_desc_bk0_n_bk1_{
+                  DeviceOp::MakeBGridDescriptor_BK0_N_BK1(b_gs_ns_ks_lengths, b_gs_ns_ks_strides)},
+              z_grid_desc_m_n_{MakeZGridDescriptor_M_N(z_gs_ms_ns_lengths, z_gs_ms_ns_strides)},
+              b1_grid_desc_bk0_n_bk1_{DeviceOp::MakeB1GridDescriptor_BK0_N_BK1(
+                  b1_gs_gemm1ns_gemm1ks_lengths, b1_gs_gemm1ns_gemm1ks_strides)},
+              y_grid_desc_m_o_{Transform::MakeCGridDescriptor_M_N(c_gs_ms_gemm1ns_lengths,
+                                                                  c_gs_ms_gemm1ns_strides)},
+              lse_grid_desc_m_{DeviceOp::MakeLSEGridDescriptor_M(lse_gs_ms_lengths[NumDimG])},
+              vgrad_grid_desc_n_o_{DeviceOp::MakeVGradGridDescriptor_N_O(
+                  b1_gs_gemm1ns_gemm1ks_lengths, b1_gs_gemm1ns_gemm1ks_strides)},
+              ygrad_grid_desc_m0_o_m1_{DeviceOp::MakeYGradGridDescriptor_M0_O_M1(y_grid_desc_m_o_)},
+              // batch offsets
+              a_grid_desc_g_m_k_{
+                  Transform::MakeAGridDescriptor_G_M_K(a_gs_ms_ks_lengths, a_gs_ms_ks_strides)},
+              b_grid_desc_g_n_k_{
+                  Transform::MakeB0GridDescriptor_G_N_K(b_gs_ns_ks_lengths, b_gs_ns_ks_strides)},
+              b1_grid_desc_g_n_k_{Transform::MakeB1GridDescriptor_G_N_K(
+                  b1_gs_gemm1ns_gemm1ks_lengths, b1_gs_gemm1ns_gemm1ks_strides)},
+              c_grid_desc_g_m_n_{Transform::MakeCGridDescriptor_G_M_N(c_gs_ms_gemm1ns_lengths,
+                                                                      c_gs_ms_gemm1ns_strides)},
+              z_grid_desc_g_m_n_{
+                  Transform::MakeCGridDescriptor_G_M_N(z_gs_ms_ns_lengths, z_gs_ms_ns_strides)},
+              y_grid_desc_mblock_mperblock_oblock_operblock_{},
+              block_2_ctile_map_{GridwiseGemm::MakeDefaultBlock2CTileMap(y_grid_desc_m_o_)},
+              a_element_op_{a_element_op},
+              b_element_op_{b_element_op},
+              acc_element_op_{acc_element_op},
+              b1_element_op_{b1_element_op},
+              c_element_op_{c_element_op},
+              c0_matrix_mask_{b_grid_desc_g_n_k_.GetLength(I1)},
+              raw_lengths_mz_nz_kz_gemm1nz_{a_gs_ms_ks_lengths[NumDimG + NumDimM - 1],
+                                            b_gs_ns_ks_lengths[NumDimG + NumDimN - 1],
+                                            b_gs_ns_ks_lengths[NumDimG + NumDimN + NumDimK - 1],
+                                            b1_gs_gemm1ns_gemm1ks_lengths[NumDimG + NumDimO - 1]},
+              a_mz_kz_strides_{a_gs_ms_ks_strides[NumDimG + NumDimM - 1],
+                               a_gs_ms_ks_strides[NumDimG + NumDimM + NumDimK - 1]},
+              b_nz_kz_strides_{b_gs_ns_ks_strides[NumDimG + NumDimN - 1],
+                               b_gs_ns_ks_strides[NumDimG + NumDimN + NumDimK - 1]},
+              b1_nz_kz_strides_{b1_gs_gemm1ns_gemm1ks_strides[NumDimG + NumDimO - 1],
+                                b1_gs_gemm1ns_gemm1ks_strides[NumDimG + NumDimO + NumDimN - 1]},
+              c_mz_gemm1nz_strides_{c_gs_ms_gemm1ns_strides[NumDimG + NumDimM - 1],
+                                    c_gs_ms_gemm1ns_strides[NumDimG + NumDimM + NumDimO - 1]},
+              batch_count_{c_grid_desc_g_m_n_.GetLength(I0)},
+              compute_base_ptr_of_batch_{
+                  a_grid_desc_g_m_k_,
+                  b_grid_desc_g_n_k_,
+                  z_grid_desc_g_m_n_,
+                  b1_grid_desc_g_n_k_,
+                  c_grid_desc_g_m_n_,
+                  type_convert<index_t>(lse_grid_desc_m_.GetElementSpaceSize())}
+        {
+            // TODO: implement bias addition
+            ignore = p_acc0_biases;
+            ignore = p_acc1_biases;
+            ignore = acc0_biases_gs_ms_ns_lengths;
+            ignore = acc0_biases_gs_ms_ns_strides;
+            ignore = acc1_biases_gs_ms_gemm1ns_lengths;
+            ignore = acc1_biases_gs_ms_gemm1ns_strides;
+
+            if(GridwiseGemm::CheckValidity(a_grid_desc_ak0_m_ak1_,
+                                           b_grid_desc_bk0_n_bk1_,
+                                           b1_grid_desc_bk0_n_bk1_,
+                                           y_grid_desc_m_o_,
+                                           block_2_ctile_map_))
+            {
+                y_grid_desc_mblock_mperblock_oblock_operblock_ =
+                    GridwiseGemm::MakeCGridDescriptor_MBlock_MPerBlock_NBlock_NPerBlock(
+                        y_grid_desc_m_o_);
+            }
+
+            p_dropout_        = 1.f - p_drop;
+            float rp_dropout_ = 1.f / p_dropout_;
+            acc_element_op_.Append(rp_dropout_);
+
+            seed_   = std::get<0>(seeds);
+            offset_ = std::get<1>(seeds);
+
+            c_grid_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5_ =
+                GridwiseGemm::MakeCGridDescriptor_M0_N0_M1_N1_M2_N2_M3_N3_N4_N5(z_grid_desc_m_n_);
+            // Print();
+        }
+
+        void Print() const
+        {
+            std::cout << "a_grid_desc_g_m_k_: " << a_grid_desc_g_m_k_.GetLength(I0) << ", "
+                      << a_grid_desc_g_m_k_.GetLength(I1) << ", "
+                      << a_grid_desc_g_m_k_.GetLength(I2) << '\n';
+            // a_grid_desc_g_m_k_.Print();
+            std::cout << "b_grid_desc_g_n_k_: " << b_grid_desc_g_n_k_.GetLength(I0) << ", "
+                      << b_grid_desc_g_n_k_.GetLength(I1) << ", "
+                      << b_grid_desc_g_n_k_.GetLength(I2) << '\n';
+            // b_grid_desc_g_n_k_.Print();
+            std::cout << "b1_grid_desc_g_o_n_: " << b1_grid_desc_g_n_k_.GetLength(I0) << ", "
+                      << b1_grid_desc_g_n_k_.GetLength(I1) << ", "
+                      << b1_grid_desc_g_n_k_.GetLength(I2) << '\n';
+            // b1_grid_desc_g_n_k_.Print();
+            std::cout << "c_grid_desc_g_m_o_: " << c_grid_desc_g_m_n_.GetLength(I0) << ", "
+                      << c_grid_desc_g_m_n_.GetLength(I1) << ", "
+                      << c_grid_desc_g_m_n_.GetLength(I2) << '\n';
+            // c_grid_desc_g_m_n_.Print();
+            std::cout << "vgrad_grid_desc_n_o_: " << vgrad_grid_desc_n_o_.GetLength(I0) << ", "
+                      << vgrad_grid_desc_n_o_.GetLength(I1) << '\n';
+            std::cout << "ygrad_grid_desc_m0_o_m1_: " << ygrad_grid_desc_m0_o_m1_.GetLength(I0)
+                      << ", " << ygrad_grid_desc_m0_o_m1_.GetLength(I1) << ", "
+                      << ygrad_grid_desc_m0_o_m1_.GetLength(I2) << '\n';
+        }
+
+        // pointers
+        const DataType* p_a_grid_;
+        const DataType* p_b_grid_;
+        ZDataType* p_z_grid_;
+        const DataType* p_b1_grid_;
+        const DataType* p_c_grid_;
+        const LSEDataType* p_lse_grid_;
+        const DataType* p_ygrad_grid_;
+        DataType* p_qgrad_grid_;
+        DataType* p_kgrad_grid_;
+        DataType* p_vgrad_grid_;
+
+        // tensor descriptor
+        AGridDesc_AK0_M_AK1 a_grid_desc_ak0_m_ak1_;
+        BGridDesc_BK0_N_BK1 b_grid_desc_bk0_n_bk1_;
+        ZGridDesc_M_N z_grid_desc_m_n_;
+        B1GridDesc_BK0_N_BK1 b1_grid_desc_bk0_n_bk1_;
+        YGridDesc_M_O y_grid_desc_m_o_;
+        LSEGridDesc_M lse_grid_desc_m_;
+        VGradGridDesc_N_O vgrad_grid_desc_n_o_;
+        YGradGridDesc_M0_O_M1 ygrad_grid_desc_m0_o_m1_;
+
+        // batch offsets
+        AGridDesc_G_M_K a_grid_desc_g_m_k_;
+        BGridDesc_G_N_K b_grid_desc_g_n_k_;
+        B1GridDesc_G_N_K b1_grid_desc_g_n_k_;
+        CGridDesc_G_M_N c_grid_desc_g_m_n_;
+        ZGridDesc_G_M_N z_grid_desc_g_m_n_;
+        typename GridwiseGemm::YGridDescriptor_MBlock_MPerBlock_OBlock_OPerBlock
+            y_grid_desc_mblock_mperblock_oblock_operblock_;
+
+        typename GridwiseGemm::ZGridDescriptor_M0_N0_M1_N1_M2_N2_M3_N3_N4_N5
+            c_grid_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5_;
+
+        // block-to-c-tile map
+        typename GridwiseGemm::DefaultBlock2CTileMap block_2_ctile_map_;
+
+        // element-wise op
+        AElementwiseOperation a_element_op_;
+        BElementwiseOperation b_element_op_;
+        AccElementwiseOperation acc_element_op_;
+        B1ElementwiseOperation b1_element_op_;
+        CElementwiseOperation c_element_op_;
+
+        // check C0 masking and padding
+        C0MatrixMask c0_matrix_mask_;
+
+        // For robust IsSupportedArgument() check
+        std::vector<index_t> raw_lengths_mz_nz_kz_gemm1nz_;
+        std::vector<index_t> a_mz_kz_strides_;
+        std::vector<index_t> b_nz_kz_strides_;
+        std::vector<index_t> b1_nz_kz_strides_;
+        std::vector<index_t> c_mz_gemm1nz_strides_;
+
+        index_t batch_count_;
+        ComputeBasePtrOfStridedBatch compute_base_ptr_of_batch_;
+
+        float p_dropout_;
+        unsigned long long seed_;
+        unsigned long long offset_;
+    };
+
+    // Invoker
+    struct Invoker : public BaseInvoker
+    {
+        using Argument = DeviceOp::Argument;
+
+        float Run(const Argument& arg, const StreamConfig& stream_config = StreamConfig{})
+        {
+            if(!DeviceOp::IsSupportedArgument(arg))
+            {
+                throw std::runtime_error("wrong! unsupported argument");
+            }
+
+            const index_t grid_size =
+                arg.block_2_ctile_map_.CalculateGridSize(arg.y_grid_desc_m_o_) * arg.batch_count_;
+
+            // Gemm0_K
+            const auto K =
+                arg.a_grid_desc_ak0_m_ak1_.GetLength(I0) * arg.a_grid_desc_ak0_m_ak1_.GetLength(I2);
+
+            float ave_time = 0;
+
+            auto launch_kernel = [&](auto has_main_k_block_loop_) {
+                const auto kernel = kernel_batched_gemm_softmax_gemm_xdl_cshuffle_v2<
+                    GridwiseGemm,
+                    DataType,
+                    ZDataType,
+                    LSEDataType,
+                    AElementwiseOperation,
+                    BElementwiseOperation,
+                    AccElementwiseOperation,
+                    B1ElementwiseOperation,
+                    CElementwiseOperation,
+                    DeviceOp::AGridDesc_AK0_M_AK1,
+                    DeviceOp::BGridDesc_BK0_N_BK1,
+                    typename GridwiseGemm::ZGridDescriptor_M0_N0_M1_N1_M2_N2_M3_N3_N4_N5,
+                    DeviceOp::B1GridDesc_BK0_N_BK1,
+                    typename GridwiseGemm::YGridDescriptor_MBlock_MPerBlock_OBlock_OPerBlock,
+                    DeviceOp::LSEGridDesc_M,
+                    DeviceOp::VGradGridDesc_N_O,
+                    DeviceOp::YGradGridDesc_M0_O_M1,
+                    typename GridwiseGemm::DefaultBlock2CTileMap,
+                    ComputeBasePtrOfStridedBatch,
+                    C0MatrixMask,
+                    has_main_k_block_loop_>;
+
+                return launch_and_time_kernel(stream_config,
+                                              kernel,
+                                              dim3(grid_size),
+                                              dim3(BlockSize),
+                                              0,
+                                              arg.p_a_grid_,
+                                              arg.p_b_grid_,
+                                              arg.p_z_grid_,
+                                              arg.p_b1_grid_,
+                                              arg.p_c_grid_,
+                                              arg.p_lse_grid_,
+                                              arg.p_ygrad_grid_,
+                                              arg.p_qgrad_grid_,
+                                              arg.p_kgrad_grid_,
+                                              arg.p_vgrad_grid_,
+                                              arg.a_element_op_,
+                                              arg.b_element_op_,
+                                              arg.acc_element_op_,
+                                              arg.b1_element_op_,
+                                              arg.c_element_op_,
+                                              arg.a_grid_desc_ak0_m_ak1_,
+                                              arg.b_grid_desc_bk0_n_bk1_,
+                                              arg.c_grid_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5_,
+                                              arg.b1_grid_desc_bk0_n_bk1_,
+                                              arg.y_grid_desc_mblock_mperblock_oblock_operblock_,
+                                              arg.lse_grid_desc_m_,
+                                              arg.vgrad_grid_desc_n_o_,
+                                              arg.ygrad_grid_desc_m0_o_m1_,
+                                              arg.block_2_ctile_map_,
+                                              arg.batch_count_,
+                                              arg.compute_base_ptr_of_batch_,
+                                              arg.c0_matrix_mask_,
+                                              arg.p_dropout_,
+                                              arg.seed_,
+                                              arg.offset_);
+            };
+
+            // Gemm1_K is split into Gemm1_K0/K1 where K1 is known at compile time, so we only need
+            // to concern Gemm0's loop
+#if 1
+            if(GridwiseGemm::CalculateHasMainKBlockLoop(K))
+            {
+                ave_time = launch_kernel(integral_constant<bool, true>{});
+            }
+            else
+            {
+                ave_time = launch_kernel(integral_constant<bool, false>{});
+            }
+#endif
+            return ave_time;
+        }
+
+        // polymorphic
+        float Run(const BaseArgument* p_arg,
+                  const StreamConfig& stream_config = StreamConfig{}) override
+        {
+            return Run(*dynamic_cast<const Argument*>(p_arg), stream_config);
+        }
+    };
+
+    static constexpr bool IsValidCompilationParameter()
+    {
+        // TODO: properly implement this check
+        return true;
+    }
+
+    static bool IsSupportedArgument(const Argument& arg)
+    {
+#if 0
+        arg.Print();
+#endif
+
+        if(!(ck::get_device_name() == "gfx908" || ck::get_device_name() == "gfx90a"))
+        {
+            return false;
+        }
+
+        // TODO: Check if tensor specialization & strides mismatch
+
+        // Check if C permute dimension matches GEMM + GEMM shape
+        const index_t c_g       = arg.c_grid_desc_g_m_n_.GetLength(I0); // unpadded
+        const index_t c_m       = arg.y_grid_desc_m_o_.GetLength(I0);
+        const index_t c_gemm1n  = arg.y_grid_desc_m_o_.GetLength(I1);
+        const index_t a_m       = arg.a_grid_desc_ak0_m_ak1_.GetLength(I1);
+        const index_t b1_gemm1n = arg.b1_grid_desc_bk0_n_bk1_.GetLength(I1);
+
+        if(!(c_g == arg.batch_count_ && c_m == a_m && c_gemm1n == b1_gemm1n))
+        {
+            return false;
+        }
+
+        // Note: we need raw lengths since threadwise copy can not handle vector load when part of
+        // vector is out of bounds
+        // Note: need lowest dim in Ms/Ns/Ks/Os, not merged M/N/K/O
+        const auto MzRaw      = arg.raw_lengths_mz_nz_kz_gemm1nz_[0];
+        const auto NzRaw      = arg.raw_lengths_mz_nz_kz_gemm1nz_[1];
+        const auto KzRaw      = arg.raw_lengths_mz_nz_kz_gemm1nz_[2];
+        const auto Gemm1NzRaw = arg.raw_lengths_mz_nz_kz_gemm1nz_[3];
+
+        // Check scalar per vector requirement
+        const auto a_extent_lowest  = ABlockTransferSrcVectorDim == 2 ? KzRaw : MzRaw;
+        const auto b_extent_lowest  = BBlockTransferSrcVectorDim == 2 ? KzRaw : NzRaw;
+        const auto b1_extent_lowest = B1BlockTransferSrcVectorDim == 2 ? NzRaw : Gemm1NzRaw;
+        const auto c_extent_lowest  = Gemm1NzRaw;
+
+        if(!(a_extent_lowest % ABlockTransferSrcScalarPerVector == 0 &&
+             b_extent_lowest % BBlockTransferSrcScalarPerVector == 0 &&
+             b1_extent_lowest % B1BlockTransferSrcScalarPerVector == 0 &&
+             c_extent_lowest % CShuffleBlockTransferScalarPerVector_NPerBlock == 0))
+        {
+            return false;
+        }
+
+        // Check vector load/store requirement
+        const auto a_stride_lowest =
+            ABlockTransferSrcVectorDim == 2 ? arg.a_mz_kz_strides_[1] : arg.a_mz_kz_strides_[0];
+        const auto b_stride_lowest =
+            BBlockTransferSrcVectorDim == 2 ? arg.b_nz_kz_strides_[1] : arg.b_nz_kz_strides_[0];
+        const auto b1_stride_lowest =
+            B1BlockTransferSrcVectorDim == 2 ? arg.b1_nz_kz_strides_[1] : arg.b1_nz_kz_strides_[0];
+        const auto c_stride_lowest =
+            arg.c_mz_gemm1nz_strides_[1]; // cshuffle assumes lowest dim in Gemm1Ns to be contiguous
+
+        if(!(a_stride_lowest == 1 || b_stride_lowest == 1 || b1_stride_lowest == 1 ||
+             c_stride_lowest == 1))
+        {
+            return false;
+        }
+
+        return GridwiseGemm::CheckValidity(arg.a_grid_desc_ak0_m_ak1_,
+                                           arg.b_grid_desc_bk0_n_bk1_,
+                                           arg.b1_grid_desc_bk0_n_bk1_,
+                                           arg.y_grid_desc_m_o_,
+                                           arg.block_2_ctile_map_);
+    }
+
+    // polymorphic
+    bool IsSupportedArgument(const BaseArgument* p_arg) override
+    {
+        return IsSupportedArgument(*dynamic_cast<const Argument*>(p_arg));
+    }
+
+    static auto MakeArgument(
+        const DataType* p_a,
+        const DataType* p_b,
+        ZDataType* p_z,
+        const DataType* p_b1,
+        const DataType* p_c,
+        const LSEDataType* p_lse,
+        const DataType* p_ygrad_grid,
+        DataType* p_qgrad_grid,
+        DataType* p_kgrad_grid,
+        DataType* p_vgrad_grid,
+        const std::array<void*, NumAcc0Bias> p_acc0_biases,
+        const std::array<void*, NumAcc1Bias> p_acc1_biases,
+        const std::vector<index_t>& a_gs_ms_ks_lengths,
+        const std::vector<index_t>& a_gs_ms_ks_strides,
+        const std::vector<index_t>& b_gs_ns_ks_lengths,
+        const std::vector<index_t>& b_gs_ns_ks_strides,
+        const std::vector<index_t>& z_gs_ms_ns_lengths,
+        const std::vector<index_t>& z_gs_ms_ns_strides,
+        const std::vector<index_t>& b1_gs_gemm1ns_gemm1ks_lengths, // b1_gs_os_ns_lengths
+        const std::vector<index_t>& b1_gs_gemm1ns_gemm1ks_strides, // b1_gs_os_ns_strides
+        const std::vector<index_t>& c_gs_ms_gemm1ns_lengths,       // c_gs_ms_os_lengths
+        const std::vector<index_t>& c_gs_ms_gemm1ns_strides,       // c_gs_ms_os_strides
+        const std::vector<index_t>& lse_gs_ms_lengths,
+        const std::array<std::vector<ck::index_t>, NumAcc0Bias> acc0_biases_gs_ms_ns_lengths,
+        const std::array<std::vector<ck::index_t>, NumAcc0Bias> acc0_biases_gs_ms_ns_strides,
+        const std::array<std::vector<ck::index_t>, NumAcc1Bias>
+            acc1_biases_gs_ms_gemm1ns_lengths, // acc1_biases_gs_ms_os_lengths
+        const std::array<std::vector<ck::index_t>, NumAcc1Bias>
+            acc1_biases_gs_ms_gemm1ns_strides, // acc1_biases_gs_ms_os_strides
+        AElementwiseOperation a_element_op,
+        BElementwiseOperation b_element_op,
+        AccElementwiseOperation acc_element_op,
+        B1ElementwiseOperation b1_element_op,
+        CElementwiseOperation c_element_op,
+        float p_drop,
+        std::tuple<unsigned long long, unsigned long long> seeds)
+    {
+        return Argument{p_a,
+                        p_b,
+                        p_z,
+                        p_b1,
+                        p_c,
+                        p_lse,
+                        p_ygrad_grid,
+                        p_qgrad_grid,
+                        p_kgrad_grid,
+                        p_vgrad_grid,
+                        p_acc0_biases,
+                        p_acc1_biases,
+                        a_gs_ms_ks_lengths,
+                        a_gs_ms_ks_strides,
+                        b_gs_ns_ks_lengths,
+                        b_gs_ns_ks_strides,
+                        z_gs_ms_ns_lengths,
+                        z_gs_ms_ns_strides,
+                        b1_gs_gemm1ns_gemm1ks_lengths, // b1_gs_os_ns_lengths
+                        b1_gs_gemm1ns_gemm1ks_strides, // b1_gs_os_ns_strides
+                        c_gs_ms_gemm1ns_lengths,       // c_gs_ms_os_lengths
+                        c_gs_ms_gemm1ns_strides,       // c_gs_ms_os_strides
+                        lse_gs_ms_lengths,
+                        acc0_biases_gs_ms_ns_lengths,
+                        acc0_biases_gs_ms_ns_strides,
+                        acc1_biases_gs_ms_gemm1ns_lengths, // acc1_biases_gs_ms_os_lengths
+                        acc1_biases_gs_ms_gemm1ns_strides, // acc1_biases_gs_ms_os_strides
+                        a_element_op,
+                        b_element_op,
+                        acc_element_op,
+                        b1_element_op,
+                        c_element_op,
+                        p_drop,
+                        seeds};
+    }
+
+    static auto MakeInvoker() { return Invoker{}; }
+
+    // polymorphic
+    // FIXME: constness
+    std::unique_ptr<BaseArgument> MakeArgumentPointer(
+        const void* p_a,
+        const void* p_b,
+        void* p_z,
+        const void* p_b1,
+        const void* p_c,
+        const void* p_lse,
+        const void* p_ygrad_grid,
+        void* p_qgrad_grid,
+        void* p_kgrad_grid,
+        void* p_vgrad_grid,
+        const std::array<void*, NumAcc0Bias> p_acc0_biases,
+        const std::array<void*, NumAcc1Bias> p_acc1_biases,
+        const std::vector<index_t>& a_gs_ms_ks_lengths,
+        const std::vector<index_t>& a_gs_ms_ks_strides,
+        const std::vector<index_t>& b_gs_ns_ks_lengths,
+        const std::vector<index_t>& b_gs_ns_ks_strides,
+        const std::vector<index_t>& z_gs_ms_ns_lengths,
+        const std::vector<index_t>& z_gs_ms_ns_strides,
+        const std::vector<index_t>& b1_gs_gemm1ns_gemm1ks_lengths, // b1_gs_os_ns_lengths
+        const std::vector<index_t>& b1_gs_gemm1ns_gemm1ks_strides, // b1_gs_os_ns_strides
+        const std::vector<index_t>& c_gs_ms_gemm1ns_lengths,       // c_gs_ms_os_lengths
+        const std::vector<index_t>& c_gs_ms_gemm1ns_strides,       // c_gs_ms_os_strides
+        const std::vector<index_t>& lse_gs_ms_lengths,
+        const std::array<std::vector<ck::index_t>, NumAcc0Bias> acc0_biases_gs_ms_ns_lengths,
+        const std::array<std::vector<ck::index_t>, NumAcc0Bias> acc0_biases_gs_ms_ns_strides,
+        const std::array<std::vector<ck::index_t>, NumAcc1Bias>
+            acc1_biases_gs_ms_gemm1ns_lengths, // acc1_biases_gs_ms_os_lengths
+        const std::array<std::vector<ck::index_t>, NumAcc1Bias>
+            acc1_biases_gs_ms_gemm1ns_strides, // acc1_biases_gs_ms_os_strides
+        AElementwiseOperation a_element_op,
+        BElementwiseOperation b_element_op,
+        AccElementwiseOperation acc_element_op,
+        B1ElementwiseOperation b1_element_op,
+        CElementwiseOperation c_element_op,
+        float p_drop,
+        std::tuple<unsigned long long, unsigned long long> seeds) // override
+    {
+        return std::make_unique<Argument>(static_cast<const DataType*>(p_a),
+                                          static_cast<const DataType*>(p_b),
+                                          static_cast<ZDataType*>(p_z),
+                                          static_cast<const DataType*>(p_b1),
+                                          static_cast<const DataType*>(p_c),
+                                          static_cast<const LSEDataType*>(p_lse),
+                                          static_cast<const DataType*>(p_ygrad_grid),
+                                          static_cast<DataType*>(p_qgrad_grid),
+                                          static_cast<DataType*>(p_kgrad_grid),
+                                          static_cast<DataType*>(p_vgrad_grid),
+                                          p_acc0_biases, // cast in struct Argument
+                                          p_acc1_biases, // cast in struct Argument
+                                          a_gs_ms_ks_lengths,
+                                          a_gs_ms_ks_strides,
+                                          b_gs_ns_ks_lengths,
+                                          b_gs_ns_ks_strides,
+                                          z_gs_ms_ns_lengths,
+                                          z_gs_ms_ns_strides,
+                                          b1_gs_gemm1ns_gemm1ks_lengths, // b1_gs_os_ns_lengths
+                                          b1_gs_gemm1ns_gemm1ks_strides, // b1_gs_os_ns_strides
+                                          c_gs_ms_gemm1ns_lengths,       // c_gs_ms_os_lengths
+                                          c_gs_ms_gemm1ns_strides,       // c_gs_ms_os_strides
+                                          lse_gs_ms_lengths,
+                                          acc0_biases_gs_ms_ns_lengths,
+                                          acc0_biases_gs_ms_ns_strides,
+                                          acc1_biases_gs_ms_gemm1ns_lengths,
+                                          acc1_biases_gs_ms_gemm1ns_strides,
+                                          a_element_op,
+                                          b_element_op,
+                                          acc_element_op,
+                                          b1_element_op,
+                                          c_element_op,
+                                          p_drop,
+                                          seeds);
+    }
+
+    // polymorphic
+    std::unique_ptr<BaseInvoker> MakeInvokerPointer() // override
+    {
+        return std::make_unique<Invoker>(Invoker{});
+    }
+
+    // polymorphic
+    std::string GetTypeString() const override
+    {
+        auto str = std::stringstream();
+
+        // clang-format off
+        str << "DeviceBatchedMultiheadAttentionBackward_Train_Xdl_CShuffle"
+            << "<"
+            << BlockSize << ", "
+            << MPerBlock << ", "
+            << NPerBlock << ", "
+            << KPerBlock << ", "
+            << AK1 << ", "
+            << BK1 << ", "
+            << MPerBlock << ", "
+            << Gemm1NPerBlock << ", "
+            << Gemm1KPerBlock << ", "
+            << B1K1 << ", "
+            << getGemmSpecializationString(GemmSpec) << ", "
+            << "ASpec" << getTensorSpecializationString(ASpec) << ", "
+            << "B0Spec" << getTensorSpecializationString(BSpec) << ", "
+            << "B1Spec" << getTensorSpecializationString(B1Spec) << ", "
+            << "CSpec" << getTensorSpecializationString(CSpec) << ", "
+            << getMaskingSpecializationString(MaskingSpec) << ">";
+        // clang-format on
+
+        return str.str();
+    }
+};
+
+} // namespace device
+} // namespace tensor_operation
+} // namespace ck

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

@@ -7,6 +7,7 @@
 #include <sstream>
 
 #include "ck/utility/common_header.hpp"
+#include "ck/utility/philox_rand.hpp"
 #include "ck/tensor_description/tensor_descriptor.hpp"
 #include "ck/tensor_description/tensor_descriptor_helper.hpp"
 // #include "ck/tensor_operation/gpu/device/device_batched_multihead_attention_backward.hpp" // TODO

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

@@ -7,6 +7,7 @@
 #include <sstream>
 
 #include "ck/utility/common_header.hpp"
+#include "ck/utility/philox_rand.hpp"
 #include "ck/tensor_description/tensor_descriptor.hpp"
 #include "ck/tensor_description/tensor_descriptor_helper.hpp"
 // #include "ck/tensor_operation/gpu/device/device_batched_multihead_attention_backward.hpp" // TODO
@@ -45,7 +46,10 @@ __global__ void
             const BElementwiseOperation b_element_op,
             const AccElementwiseOperation acc_element_op,
             const B1ElementwiseOperation b1_element_op,
-            const CElementwiseOperation c_element_op)
+            const CElementwiseOperation c_element_op,
+            const float p_dropout,
+            const unsigned long long seed,
+            const unsigned long long offset)
 {
 #if(!defined(__HIP_DEVICE_COMPILE__) || defined(__gfx908__) || defined(__gfx90a__))
     __shared__ char p_shared[GridwiseGemm::GetSharedMemoryNumberOfByte()];
@@ -80,15 +84,25 @@ __global__ void
         static_cast<long_index_t>(arg_ptr[group_id].compute_base_ptr_of_batch_.GetABasePtr(g_idx)));
     const long_index_t b_batch_offset = __builtin_amdgcn_readfirstlane(
         static_cast<long_index_t>(arg_ptr[group_id].compute_base_ptr_of_batch_.GetBBasePtr(g_idx)));
+    const long_index_t z_batch_offset = __builtin_amdgcn_readfirstlane(
+        static_cast<long_index_t>(arg_ptr[group_id].compute_base_ptr_of_batch_.GetZBasePtr(g_idx)));
     const long_index_t b1_batch_offset = __builtin_amdgcn_readfirstlane(static_cast<long_index_t>(
         arg_ptr[group_id].compute_base_ptr_of_batch_.GetB1BasePtr(g_idx)));
     const long_index_t c_batch_offset  = __builtin_amdgcn_readfirstlane(
         static_cast<long_index_t>(arg_ptr[group_id].compute_base_ptr_of_batch_.GetCBasePtr(g_idx)));
     const long_index_t lse_batch_offset = __builtin_amdgcn_readfirstlane(static_cast<long_index_t>(
         arg_ptr[group_id].compute_base_ptr_of_batch_.GetLSEBasePtr(g_idx)));
+
+    const index_t global_thread_id = get_thread_global_1d_id();
+    ck::philox ph(seed, global_thread_id, offset);
+    unsigned short* z_matrix_ptr =
+        (arg_ptr[group_id].p_z_grid_ == nullptr ? nullptr
+                                                : arg_ptr[group_id].p_z_grid_ + z_batch_offset);
+
     GridwiseGemm::template Run<HasMainKBlockLoop>(
         arg_ptr[group_id].p_a_grid_ + a_batch_offset,
         arg_ptr[group_id].p_b_grid_ + b_batch_offset,
+        z_matrix_ptr,
         arg_ptr[group_id].p_b1_grid_ + b1_batch_offset,
         arg_ptr[group_id].p_c_grid_ + c_batch_offset,
         arg_ptr[group_id].p_lse_grid_ + lse_batch_offset,
@@ -104,13 +118,16 @@ __global__ void
         c_element_op,
         arg_ptr[group_id].a_grid_desc_ak0_m_ak1_,
         arg_ptr[group_id].b_grid_desc_bk0_n_bk1_,
+        arg_ptr[group_id].c_grid_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5_,
         arg_ptr[group_id].b1_grid_desc_bk0_n_bk1_,
         arg_ptr[group_id].y_grid_desc_mblock_mperblock_oblock_operblock_,
         arg_ptr[group_id].lse_grid_desc_m_,
         arg_ptr[group_id].vgrad_grid_desc_n_o_,
-        arg_ptr[group_id].ygrad_grid_desc_o0_m_o1_,
+        arg_ptr[group_id].ygrad_grid_desc_m0_o_m1_,
         arg_ptr[group_id].block_2_ctile_map_,
-        arg_ptr[group_id].c0_matrix_mask_);
+        arg_ptr[group_id].c0_matrix_mask_,
+        p_dropout,
+        ph);
 #else
     ignore = group_kernel_args;
     ignore = group_count;
@@ -131,6 +148,7 @@ template <index_t NumDimG,
           index_t NumDimK,
           index_t NumDimO, // NumDimGemm1N
           typename DataType,
+          typename ZDataType,
           typename LSEDataType,
           typename Acc0BiasDataType,
           typename Acc1BiasDataType,
@@ -209,6 +227,9 @@ struct DeviceGroupedMultiheadAttentionBackward_Xdl_CShuffle_V2
         std::vector<index_t> b_gs_ns_ks_lengths;
         std::vector<index_t> b_gs_ns_ks_strides;
 
+        std::vector<index_t> z_gs_ms_ns_lengths;
+        std::vector<index_t> z_gs_ms_ns_strides;
+
         std::vector<index_t> b1_gs_gemm1ns_gemm1ks_lengths;
         std::vector<index_t> b1_gs_gemm1ns_gemm1ks_strides;
 
@@ -359,19 +380,6 @@ struct DeviceGroupedMultiheadAttentionBackward_Xdl_CShuffle_V2
             make_tuple(Sequence<0, 2>{}, Sequence<1>{}));
     }
 
-    //
-    // dP = dY * V^T
-    //
-
-    // YGrad in Gemm A position
-    static auto MakeYGradGridDescriptor_O0_M_O1(const std::vector<index_t>& y_gs_ms_os_lengths_vec,
-                                                const std::vector<index_t>& y_gs_ms_os_strides_vec)
-    {
-        return Transform::MakeAGridDescriptor_AK0_M_AK1(
-            Transform::MakeAGridDescriptor_M_K(y_gs_ms_os_lengths_vec, y_gs_ms_os_strides_vec),
-            Number<Y_O1>{});
-    }
-
     // V in Gemm B position
     static auto MakeVGridDescriptor_O0_N_O1(const std::vector<index_t>& v_gs_os_ns_lengths_vec,
                                             const std::vector<index_t>& v_gs_os_ns_strides_vec)
@@ -447,6 +455,12 @@ struct DeviceGroupedMultiheadAttentionBackward_Xdl_CShuffle_V2
         return Transform::MakeCGridDescriptor_M_N(k_gs_ns_ks_lengths_vec, k_gs_ns_ks_strides_vec);
     }
 
+    static auto MakeZGridDescriptor_M_N(const std::vector<index_t>& z_gs_ms_ns_lengths_vec,
+                                        const std::vector<index_t>& z_gs_ms_ns_strides_vec)
+    {
+        return Transform::MakeCGridDescriptor_M_N(z_gs_ms_ns_lengths_vec, z_gs_ms_ns_strides_vec);
+    }
+
     static auto MakeLSEGridDescriptor_M(index_t MRaw)
     {
         const auto lse_grid_desc_mraw = make_naive_tensor_descriptor_packed(make_tuple(MRaw));
@@ -481,9 +495,11 @@ struct DeviceGroupedMultiheadAttentionBackward_Xdl_CShuffle_V2
     using BGridDesc_G_N_K      = decltype(Transform::MakeB0GridDescriptor_G_N_K({}, {}));
     using B1GridDesc_G_N_K     = decltype(Transform::MakeB1GridDescriptor_G_N_K({}, {}));
     using CGridDesc_G_M_N      = decltype(Transform::MakeCGridDescriptor_G_M_N({}, {}));
+    using ZGridDesc_G_M_N      = decltype(Transform::MakeCGridDescriptor_G_M_N({}, {}));
 
     using VGradGridDesc_N_O     = decltype(MakeVGradGridDescriptor_N_O({}, {}));
-    using YGradGridDesc_O0_M_O1 = decltype(MakeYGradGridDescriptor_O0_M_O1({}, {}));
+    using YGradGridDesc_M0_O_M1 = decltype(MakeYGradGridDescriptor_M0_O_M1(YGridDesc_M_O{}));
+    using ZGridDesc_M_N         = decltype(MakeZGridDescriptor_M_N({}, {}));
 
     constexpr static auto make_MaskOutPredicate()
     {
@@ -502,11 +518,13 @@ struct DeviceGroupedMultiheadAttentionBackward_Xdl_CShuffle_V2
     {
         ComputeBasePtrOfStridedBatch(const AGridDesc_G_M_K& a_grid_desc_g_m_k,
                                      const BGridDesc_G_N_K& b_grid_desc_g_n_k,
+                                     const ZGridDesc_G_M_N& z_grid_desc_g_m_n,
                                      const B1GridDesc_G_N_K& b1_grid_desc_g_n_k,
                                      const CGridDesc_G_M_N& c_grid_desc_g_m_n,
                                      index_t BatchStrideLSE)
             : a_grid_desc_g_m_k_(a_grid_desc_g_m_k),
               b_grid_desc_g_n_k_(b_grid_desc_g_n_k),
+              z_grid_desc_g_m_n_(z_grid_desc_g_m_n),
               b1_grid_desc_g_n_k_(b1_grid_desc_g_n_k),
               c_grid_desc_g_m_n_(c_grid_desc_g_m_n),
               BatchStrideLSE_(BatchStrideLSE)
@@ -523,6 +541,11 @@ struct DeviceGroupedMultiheadAttentionBackward_Xdl_CShuffle_V2
             return b_grid_desc_g_n_k_.CalculateOffset(make_multi_index(g_idx, 0, 0));
         }
 
+        __host__ __device__ constexpr long_index_t GetZBasePtr(index_t g_idx) const
+        {
+            return z_grid_desc_g_m_n_.CalculateOffset(make_multi_index(g_idx, 0, 0));
+        }
+
         __host__ __device__ constexpr long_index_t GetB1BasePtr(index_t g_idx) const
         {
             return b1_grid_desc_g_n_k_.CalculateOffset(make_multi_index(g_idx, 0, 0));
@@ -541,6 +564,7 @@ struct DeviceGroupedMultiheadAttentionBackward_Xdl_CShuffle_V2
         private:
         AGridDesc_G_M_K a_grid_desc_g_m_k_;
         BGridDesc_G_N_K b_grid_desc_g_n_k_;
+        ZGridDesc_G_M_N z_grid_desc_g_m_n_;
         B1GridDesc_G_N_K b1_grid_desc_g_n_k_;
         CGridDesc_G_M_N c_grid_desc_g_m_n_;
         index_t BatchStrideLSE_;
@@ -560,6 +584,7 @@ struct DeviceGroupedMultiheadAttentionBackward_Xdl_CShuffle_V2
         InMemoryDataOperationEnum::Set,
         AGridDesc_AK0_M_AK1,
         BGridDesc_BK0_N_BK1,
+        ZGridDesc_M_N,
         B1GridDesc_BK0_N_BK1,
         YGridDesc_M_O,
         LSEGridDesc_M,
@@ -617,6 +642,7 @@ struct DeviceGroupedMultiheadAttentionBackward_Xdl_CShuffle_V2
         // pointers
         const DataType* p_a_grid_;
         const DataType* p_b_grid_;
+        ZDataType* p_z_grid_;
         const DataType* p_b1_grid_;
         const DataType* p_c_grid_;
         const LSEDataType* p_lse_grid_;
@@ -628,14 +654,17 @@ struct DeviceGroupedMultiheadAttentionBackward_Xdl_CShuffle_V2
         // tensor descriptors for block/thread-wise copy
         AGridDesc_AK0_M_AK1 a_grid_desc_ak0_m_ak1_;
         BGridDesc_BK0_N_BK1 b_grid_desc_bk0_n_bk1_;
+        ZGridDesc_M_N z_grid_desc_m_n_;
         B1GridDesc_BK0_N_BK1 b1_grid_desc_bk0_n_bk1_;
         YGridDesc_M_O y_grid_desc_m_o_;
+
         typename GridwiseGemm::YGridDescriptor_MBlock_MPerBlock_OBlock_OPerBlock
             y_grid_desc_mblock_mperblock_oblock_operblock_;
-
+        typename GridwiseGemm::ZGridDescriptor_M0_N0_M1_N1_M2_N2_M3_N3_N4_N5
+            c_grid_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5_;
         LSEGridDesc_M lse_grid_desc_m_;
         VGradGridDesc_N_O vgrad_grid_desc_n_o_;
-        YGradGridDesc_O0_M_O1 ygrad_grid_desc_o0_m_o1_;
+        YGradGridDesc_M0_O_M1 ygrad_grid_desc_m0_o_m1_;
         // block-to-c-tile map
         Block2CTileMap block_2_ctile_map_;
         index_t num_blocks_per_batch_;
@@ -667,6 +696,7 @@ struct DeviceGroupedMultiheadAttentionBackward_Xdl_CShuffle_V2
     {
         Argument(const std::vector<const void*>& p_As,
                  const std::vector<const void*>& p_Bs,
+                 const std::vector<void*>& p_Zs,
                  const std::vector<const void*>& p_B1s,
                  const std::vector<const void*>& p_Cs, // for dS
                  const std::vector<const void*>& p_LSEs,
@@ -681,13 +711,21 @@ struct DeviceGroupedMultiheadAttentionBackward_Xdl_CShuffle_V2
                  BElementwiseOperation b_element_op,
                  AccElementwiseOperation acc_element_op,
                  B1ElementwiseOperation b1_element_op,
-                 CElementwiseOperation c_element_op)
+                 CElementwiseOperation c_element_op,
+                 float p_drop,
+                 std::tuple<unsigned long long, unsigned long long> seeds)
             : a_element_op_{a_element_op},
               b_element_op_{b_element_op},
               acc_element_op_{acc_element_op},
               b1_element_op_{b1_element_op},
               c_element_op_{c_element_op}
         {
+            p_dropout_        = 1.f - p_drop;
+            float rp_dropout_ = 1.f / p_dropout_;
+            acc_element_op_.Append(rp_dropout_);
+
+            seed_   = std::get<0>(seeds);
+            offset_ = std::get<1>(seeds);
 
             group_count_ = ck::type_convert<ck::index_t>(problem_desc_vec.size());
 
@@ -714,6 +752,7 @@ struct DeviceGroupedMultiheadAttentionBackward_Xdl_CShuffle_V2
             {
                 const auto p_a_grid     = static_cast<const DataType*>(p_As[i]);
                 const auto p_b_grid     = static_cast<const DataType*>(p_Bs[i]);
+                auto p_z_grid           = static_cast<ZDataType*>(p_Zs[i]);
                 const auto p_b1_grid    = static_cast<const DataType*>(p_B1s[i]);
                 const auto p_c_grid     = static_cast<const DataType*>(p_Cs[i]);
                 const auto p_lse_grid   = static_cast<const LSEDataType*>(p_LSEs[i]);
@@ -728,6 +767,8 @@ struct DeviceGroupedMultiheadAttentionBackward_Xdl_CShuffle_V2
                     problem_desc.a_gs_ms_ks_lengths, problem_desc.a_gs_ms_ks_strides);
                 const auto b_grid_desc_bk0_n_bk1 = DeviceOp::MakeBGridDescriptor_BK0_N_BK1(
                     problem_desc.b_gs_ns_ks_lengths, problem_desc.b_gs_ns_ks_strides);
+                const auto z_grid_desc_m_n = DeviceOp::MakeZGridDescriptor_M_N(
+                    problem_desc.z_gs_ms_ns_lengths, problem_desc.z_gs_ms_ns_strides);
                 const auto b1_grid_desc_bk0_n_bk1 = DeviceOp::MakeVGridDescriptor_O0_N_O1(
                     problem_desc.b1_gs_gemm1ns_gemm1ks_lengths,
                     problem_desc.b1_gs_gemm1ns_gemm1ks_strides);
@@ -739,13 +780,15 @@ struct DeviceGroupedMultiheadAttentionBackward_Xdl_CShuffle_V2
                 const auto vgrad_grid_desc_n_o = DeviceOp::MakeVGradGridDescriptor_N_O(
                     problem_desc.b1_gs_gemm1ns_gemm1ks_lengths,
                     problem_desc.b1_gs_gemm1ns_gemm1ks_strides);
-                const auto ygrad_grid_desc_o0_m_o1 = DeviceOp::MakeYGradGridDescriptor_O0_M_O1(
-                    problem_desc.c_gs_ms_gemm1ns_lengths, problem_desc.c_gs_ms_gemm1ns_strides);
+                const auto ygrad_grid_desc_m0_o_m1 =
+                    DeviceOp::MakeYGradGridDescriptor_M0_O_M1(y_grid_desc_m_o);
 
                 const auto a_grid_desc_g_m_k = Transform::MakeAGridDescriptor_G_M_K(
                     problem_desc.a_gs_ms_ks_lengths, problem_desc.a_gs_ms_ks_strides);
                 const auto b_grid_desc_g_n_k = Transform::MakeB0GridDescriptor_G_N_K(
                     problem_desc.b_gs_ns_ks_lengths, problem_desc.b_gs_ns_ks_strides);
+                const auto z_grid_desc_g_m_n = Transform::MakeCGridDescriptor_G_M_N(
+                    problem_desc.z_gs_ms_ns_lengths, problem_desc.z_gs_ms_ns_strides);
                 const auto b1_grid_desc_g_n_k = Transform::MakeB1GridDescriptor_G_N_K(
                     problem_desc.b1_gs_gemm1ns_gemm1ks_lengths,
                     problem_desc.b1_gs_gemm1ns_gemm1ks_strides);
@@ -753,6 +796,8 @@ struct DeviceGroupedMultiheadAttentionBackward_Xdl_CShuffle_V2
                     problem_desc.c_gs_ms_gemm1ns_lengths, problem_desc.c_gs_ms_gemm1ns_strides);
                 typename GridwiseGemm::YGridDescriptor_MBlock_MPerBlock_OBlock_OPerBlock
                     y_grid_desc_mblock_mperblock_oblock_operblock;
+                typename GridwiseGemm::ZGridDescriptor_M0_N0_M1_N1_M2_N2_M3_N3_N4_N5
+                    c_grid_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5;
                 const index_t BlockStart     = grid_size_;
                 const auto block_2_ctile_map = Block2CTileMap(y_grid_desc_m_o, BlockStart);
                 if(GridwiseGemm::CheckValidity(a_grid_desc_ak0_m_ak1,
@@ -765,15 +810,20 @@ struct DeviceGroupedMultiheadAttentionBackward_Xdl_CShuffle_V2
                         GridwiseGemm::MakeCGridDescriptor_MBlock_MPerBlock_NBlock_NPerBlock(
                             y_grid_desc_m_o);
                 }
+
+                c_grid_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5 =
+                    GridwiseGemm::MakeCGridDescriptor_M0_N0_M1_N1_M2_N2_M3_N3_N4_N5(
+                        z_grid_desc_m_n);
+
                 const index_t batch_count = c_grid_desc_g_m_n.GetLength(I0);
                 const index_t grid_size_grp =
                     block_2_ctile_map.CalculateGridSize(y_grid_desc_m_o) * batch_count;
                 const index_t BlockEnd = grid_size_ + grid_size_grp;
-
                 // batch stride
                 const auto compute_base_ptr_of_batch = ComputeBasePtrOfStridedBatch(
                     a_grid_desc_g_m_k,
                     b_grid_desc_g_n_k,
+                    z_grid_desc_g_m_n,
                     b1_grid_desc_g_n_k,
                     c_grid_desc_g_m_n,
                     type_convert<index_t>(lse_grid_desc_m.GetElementSpaceSize()));
@@ -797,6 +847,7 @@ struct DeviceGroupedMultiheadAttentionBackward_Xdl_CShuffle_V2
 
                 group_kernel_args_.push_back({p_a_grid,
                                               p_b_grid,
+                                              p_z_grid,
                                               p_b1_grid,
                                               p_c_grid,
                                               p_lse_grid,
@@ -806,12 +857,14 @@ struct DeviceGroupedMultiheadAttentionBackward_Xdl_CShuffle_V2
                                               p_vgrad_grid,
                                               a_grid_desc_ak0_m_ak1,
                                               b_grid_desc_bk0_n_bk1,
+                                              z_grid_desc_m_n,
                                               b1_grid_desc_bk0_n_bk1,
                                               y_grid_desc_m_o,
                                               y_grid_desc_mblock_mperblock_oblock_operblock,
+                                              c_grid_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5,
                                               lse_grid_desc_m,
                                               vgrad_grid_desc_n_o,
-                                              ygrad_grid_desc_o0_m_o1,
+                                              ygrad_grid_desc_m0_o_m1,
                                               block_2_ctile_map,
                                               block_2_ctile_map.CalculateGridSize(y_grid_desc_m_o),
                                               compute_base_ptr_of_batch,
@@ -851,6 +904,10 @@ struct DeviceGroupedMultiheadAttentionBackward_Xdl_CShuffle_V2
         B1ElementwiseOperation b1_element_op_;
         CElementwiseOperation c_element_op_;
 
+        float p_dropout_;
+        unsigned long long seed_;
+        unsigned long long offset_;
+
         index_t grid_size_;
         index_t group_count_;
 
@@ -911,7 +968,10 @@ struct DeviceGroupedMultiheadAttentionBackward_Xdl_CShuffle_V2
                     arg.b_element_op_,
                     arg.acc_element_op_,
                     arg.b1_element_op_,
-                    arg.c_element_op_);
+                    arg.c_element_op_,
+                    arg.p_dropout_,
+                    arg.seed_,
+                    arg.offset_);
             };
 
             // Gemm1_K is split into Gemm1_K0/K1 where K1 is known at compile time, so we only need
@@ -1037,6 +1097,7 @@ struct DeviceGroupedMultiheadAttentionBackward_Xdl_CShuffle_V2
 
     static auto MakeArgument(const std::vector<const void*>& p_As,
                              const std::vector<const void*>& p_Bs,
+                             const std::vector<void*>& p_Zs,
                              const std::vector<const void*>& p_B1s,
                              const std::vector<const void*>& p_Cs, // for dS
                              const std::vector<const void*>& p_LSEs,
@@ -1051,10 +1112,13 @@ struct DeviceGroupedMultiheadAttentionBackward_Xdl_CShuffle_V2
                              BElementwiseOperation b_element_op,
                              AccElementwiseOperation acc_element_op,
                              B1ElementwiseOperation b1_element_op,
-                             CElementwiseOperation c_element_op)
+                             CElementwiseOperation c_element_op,
+                             float p_drop,
+                             std::tuple<unsigned long long, unsigned long long> seeds)
     {
         return Argument{p_As,
                         p_Bs,
+                        p_Zs,
                         p_B1s,
                         p_Cs,
                         p_LSEs,
@@ -1069,7 +1133,9 @@ struct DeviceGroupedMultiheadAttentionBackward_Xdl_CShuffle_V2
                         b_element_op,
                         acc_element_op,
                         b1_element_op,
-                        c_element_op};
+                        c_element_op,
+                        p_drop,
+                        seeds};
     }
 
     static auto MakeInvoker() { return Invoker{}; }
@@ -1079,6 +1145,7 @@ struct DeviceGroupedMultiheadAttentionBackward_Xdl_CShuffle_V2
     std::unique_ptr<BaseArgument>
     MakeArgumentPointer(const std::vector<const void*>& p_As,
                         const std::vector<const void*>& p_Bs,
+                        const std::vector<void*>& p_Zs,
                         const std::vector<const void*>& p_B1s,
                         const std::vector<const void*>& p_Cs, // for dS
                         const std::vector<const void*>& p_LSEs,
@@ -1093,10 +1160,13 @@ struct DeviceGroupedMultiheadAttentionBackward_Xdl_CShuffle_V2
                         BElementwiseOperation b_element_op,
                         AccElementwiseOperation acc_element_op,
                         B1ElementwiseOperation b1_element_op,
-                        CElementwiseOperation c_element_op) // override
+                        CElementwiseOperation c_element_op,
+                        float p_drop,
+                        std::tuple<unsigned long long, unsigned long long> seeds) // override
     {
         return std::make_unique<Argument>(p_As,
                                           p_Bs,
+                                          p_Zs,
                                           p_B1s,
                                           p_Cs,
                                           p_LSEs,
@@ -1111,7 +1181,9 @@ struct DeviceGroupedMultiheadAttentionBackward_Xdl_CShuffle_V2
                                           b_element_op,
                                           acc_element_op,
                                           b1_element_op,
-                                          c_element_op);
+                                          c_element_op,
+                                          p_drop,
+                                          seeds);
     }
 
     // polymorphic

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

@@ -95,6 +95,8 @@ struct Scale
         y = scale_ * x;
     };
 
+    __host__ __device__ void Append(float scale) { scale_ = scale_ * scale; }
+
     float scale_;
 };
 

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

@@ -4,6 +4,7 @@
 #pragma once
 
 #include "ck/utility/common_header.hpp"
+#include "ck/utility/philox_rand.hpp"
 #include "ck/tensor_description/multi_index_transform_helper.hpp"
 #include "ck/tensor_description/tensor_descriptor.hpp"
 #include "ck/tensor_description/tensor_descriptor_helper.hpp"
@@ -15,6 +16,7 @@
 #include "ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer.hpp"
 #include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"
 #include "ck/tensor_operation/gpu/block/blockwise_softmax.hpp"
+#include "ck/tensor_operation/gpu/block/blockwise_dropout.hpp"
 
 namespace ck {
 
@@ -30,7 +32,8 @@ template <typename DataType,
           InMemoryDataOperationEnum CGlobalMemoryDataOperation,
           typename QGridDesc_K0_M_K1,
           typename KGridDesc_K0_N_K1,
-          typename VGridDesc_O0_N_O1,
+          typename ZGridDesc_M_N,
+          typename VGridDesc_N0_O_N1,
           typename CGridDesc_M_N,
           typename LSEGridDesc_M,
           index_t NumGemmKPrefetchStage,
@@ -93,7 +96,10 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
     static constexpr auto I5 = Number<5>{};
     static constexpr auto I6 = Number<6>{};
     static constexpr auto I7 = Number<7>{};
+    static constexpr auto I8 = Number<8>{};
+    static constexpr auto I9 = Number<9>{};
 
+    static constexpr auto WaveSize = 64;
     // K1 should be Number<...>
     // Gemm0
     static constexpr auto AK0 = Number<KPerBlock / AK1Value>{};
@@ -104,8 +110,90 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
     static constexpr auto Gemm0MWaves = MPerBlock / (MPerXdl * MXdlPerWave);
     static constexpr auto Gemm0NWaves = NPerBlock / (NPerXdl * NXdlPerWave);
 
+    // Gemm1
+    static constexpr auto B1K0 = Number<Gemm1KPerBlock / B1K1Value>{};
+    static constexpr auto B1K1 = Number<B1K1Value>{};
+
     using ThisThreadBlock = ThisThreadBlock<BlockSize>;
 
+    using GridwiseGemmPipe = remove_cvref_t<decltype(
+        GridwiseGemmPipeline_Selector<PipelineVer, NumGemmKPrefetchStage>())>;
+
+    // C desc for source in blockwise copy
+    __host__ __device__ static constexpr auto
+    MakeCGridDescriptor_M0_N0_M1_N1_M2_N2_M3_N3_N4_N5(const ZGridDesc_M_N& z_grid_desc_m_n)
+    {
+        const auto M = z_grid_desc_m_n.GetLength(I0);
+        const auto N = z_grid_desc_m_n.GetLength(I1);
+
+        constexpr auto mfma = MfmaSelector<DataType, MPerXdl, NPerXdl>::selected_mfma;
+        constexpr auto N3   = mfma.num_groups_per_blk;
+        constexpr auto N4   = mfma.num_input_blks;
+        constexpr auto N5   = mfma.group_size;
+        return transform_tensor_descriptor(
+            z_grid_desc_m_n,
+            make_tuple(make_unmerge_transform(
+                           make_tuple(M / MPerBlock, MXdlPerWave, Gemm0MWaves, MPerXdl)),
+                       make_unmerge_transform(
+                           make_tuple(N / NPerBlock, NXdlPerWave, Gemm0NWaves, N3, N4, N5))),
+            make_tuple(Sequence<0>{}, Sequence<1>{}),
+            make_tuple(Sequence<0, 2, 4, 6>{}, Sequence<1, 3, 5, 7, 8, 9>{}));
+    }
+    __host__ __device__ static constexpr auto
+    MakeZGridDescriptor_M0_N0_M1_N1_M2_N2_M3_N3_N4_N5(const index_t M, const index_t N)
+    {
+        constexpr auto mfma = MfmaSelector<DataType, MPerXdl, NPerXdl>::selected_mfma;
+        constexpr auto N3   = mfma.num_groups_per_blk;
+        constexpr auto N4   = mfma.num_input_blks;
+        constexpr auto N5   = mfma.group_size;
+        return transform_tensor_descriptor(
+            make_naive_tensor_descriptor_packed(make_tuple(M, N)),
+            make_tuple(make_unmerge_transform(
+                           make_tuple(M / MPerBlock, MXdlPerWave, Gemm0MWaves, MPerXdl)),
+                       make_unmerge_transform(
+                           make_tuple(N / NPerBlock, NXdlPerWave, Gemm0NWaves, N3, N4, N5))),
+            make_tuple(Sequence<0>{}, Sequence<1>{}),
+            make_tuple(Sequence<0, 2, 4, 6>{}, Sequence<1, 3, 5, 7, 8, 9>{}));
+    }
+
+    __device__ static auto GetGemm0WaveIdx()
+    {
+        const index_t thread_id = get_thread_local_1d_id();
+
+        constexpr auto threadid_to_wave_idx_adaptor = make_single_stage_tensor_adaptor(
+            make_tuple(make_merge_transform(make_tuple(Gemm0MWaves, Gemm0NWaves, WaveSize))),
+            make_tuple(Sequence<0, 1, 2>{}),
+            make_tuple(Sequence<0>{}));
+
+        return threadid_to_wave_idx_adaptor.CalculateBottomIndex(make_multi_index(thread_id));
+    }
+
+    __device__ static auto GetGemm0WaveMNIdx(const index_t thread_id)
+    {
+        constexpr auto wave_threadid_to_mn_idx_adaptor = make_single_stage_tensor_adaptor(
+            make_tuple(make_merge_transform(make_tuple(WaveSize / MPerXdl, MPerXdl))),
+            make_tuple(Sequence<0, 1>{}),
+            make_tuple(Sequence<0>{}));
+
+        return wave_threadid_to_mn_idx_adaptor.CalculateBottomIndex(make_multi_index(thread_id));
+    }
+
+    template <typename ABlockDesc_AK0_M_AK1>
+    __host__ __device__ static constexpr auto
+    MakeGemm1AMmaTileDescriptor_M0_M1_M2_K(const ABlockDesc_AK0_M_AK1&)
+    {
+        return MakeGemmMmaTileDescriptor_MN0_MN1_MN2_K<MXdlPerWave, 1, 1>(ABlockDesc_AK0_M_AK1{});
+    }
+
+    template <typename BBlockDesc_BK0_N_BK1>
+    __host__ __device__ static constexpr auto
+    MakeGemm1BMmaTileDescriptor_N0_N1_N2_K(const BBlockDesc_BK0_N_BK1&)
+    {
+        constexpr index_t Gemm1NWaves = Gemm1NPerBlock / (Gemm1NXdlPerWave * NPerXdl);
+        return MakeGemmMmaTileDescriptor_MN0_MN1_MN2_K<Gemm1NXdlPerWave, Gemm1NWaves, NPerXdl>(
+            BBlockDesc_BK0_N_BK1{});
+    }
+
     __host__ __device__ static constexpr auto GetABlockDescriptor_AK0PerBlock_MPerBlock_AK1()
     {
         // A matrix in LDS memory, dst of blockwise copy
@@ -149,6 +237,14 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
             make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}));
     }
 
+    __host__ __device__ static constexpr auto GetB1BlockDescriptor_BK0PerBlock_NPerBlock_BK1()
+    {
+        // B1 matrix in LDS memory, dst of blockwise copy
+        return make_naive_tensor_descriptor(
+            make_tuple(B1K0, Number<Gemm1NPerBlock>{}, B1K1),
+            make_tuple(Number<Gemm1NPerBlock + B1BlockLdsExtraN>{} * B1K1, B1K1, I1));
+    }
+
     __host__ __device__ static constexpr auto
     GetCShuffleBlockDescriptor_MBlock_MPerBlock_NBlock_NPerBlock()
     {
@@ -178,12 +274,25 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
                        I1));
     }
 
+    template <typename Gemm2Param>
+    __host__ __device__ static constexpr auto GetB2BlockDescriptor_M0_O_M1()
+    {
+        return make_naive_tensor_descriptor(
+            make_tuple(Number<Gemm2Param::B_M0>{},
+                       Number<Gemm2Param::Free1_O>{},
+                       Number<Gemm2Param::B_M1>{}),
+            make_tuple(Number<Gemm2Param::Free1_O + Gemm2Param::B_LdsPad>{} *
+                           Number<Gemm2Param::B_M1>{},
+                       Number<Gemm2Param::B_M1>{},
+                       I1));
+    }
+
     // block_id to matrix tile idx (m0, n0) mapping are controlled by {M01, N01}
     template <typename Block2CTileMap>
     __host__ __device__ static constexpr bool
     CheckValidity(const QGridDesc_K0_M_K1& q_grid_desc_k0_m_k1,
                   const KGridDesc_K0_N_K1& k_grid_desc_k0_n_k1,
-                  const VGridDesc_O0_N_O1& v_grid_desc_o0_n_o1,
+                  const VGridDesc_N0_O_N1& v_grid_desc_n0_o_n1,
                   const CGridDesc_M_N& c_grid_desc_m_n,
                   const Block2CTileMap& block_2_ctile_map)
     {
@@ -194,7 +303,7 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
         const auto M      = q_grid_desc_k0_m_k1.GetLength(I1);
         const auto N      = k_grid_desc_k0_n_k1.GetLength(I1);
         const auto K      = q_grid_desc_k0_m_k1.GetLength(I0) * q_grid_desc_k0_m_k1.GetLength(I2);
-        const auto Gemm1N = v_grid_desc_o0_n_o1.GetLength(I0) * v_grid_desc_o0_n_o1.GetLength(I2);
+        const auto Gemm1N = v_grid_desc_n0_o_n1.GetLength(I1);
 
         // This assumption reduces implemention complexity by categorizing 6 separate GEMMs into 3
         // types of GEMM operations, therefore some code body can be reused accordingly
@@ -218,12 +327,25 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
             return false;
         }
 
+        // check gemm0 gridwise gemm pipeline
+        const auto num_gemm0_k_loop = K / KPerBlock;
+        if(!GridwiseGemmPipe::IsSupported(num_gemm0_k_loop))
+        {
+            return false;
+        }
+
         // check gemm1 gridwise gemm pipeline
         if(!(NPerBlock % Gemm1KPerBlock == 0))
         {
             return false;
         }
 
+        const auto num_gemm1_k_inner_loop = NPerBlock / Gemm1KPerBlock;
+        if(!GridwiseGemmPipe::IsSupported(num_gemm1_k_inner_loop))
+        {
+            return false;
+        }
+
         if(!block_2_ctile_map.CheckValidity(c_grid_desc_m_n))
         {
             return false;
@@ -233,12 +355,12 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
         return true;
     }
 
-    // __host__ __device__ static constexpr bool CalculateHasMainKBlockLoop(index_t K)
-    // {
-    //     const index_t num_loop = K / KPerBlock;
+    __host__ __device__ static constexpr bool CalculateHasMainKBlockLoop(index_t K)
+    {
+        const index_t num_loop = K / KPerBlock;
 
-    //     return GridwiseGemmPipe::CalculateHasMainLoop(num_loop);
-    // }
+        return GridwiseGemmPipe::CalculateHasMainLoop(num_loop);
+    }
 
     __host__ __device__ static constexpr auto
     MakeCGridDescriptor_MBlock_MPerBlock_NBlock_NPerBlock(const CGridDesc_M_N& c_grid_desc_m_n)
@@ -290,25 +412,12 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
     using DefaultBlock2CTileMap =
         remove_cvref_t<decltype(MakeDefaultBlock2CTileMap(CGridDesc_M_N{}))>;
 
-    // Q / K / V / dY 
-    struct GemmBlockwiseCopy
-    {
-        // Q matrix in LDS memory, dst of blockwise copy
-        static constexpr auto q_block_desc_k0_m_k1 =
-            GetABlockDescriptor_AK0PerBlock_MPerBlock_AK1();
-
-        // K matrix in LDS memory, dst of blockwise copy
-        static constexpr auto k_block_desc_k0_n_k1 =
-            GetBBlockDescriptor_BK0PerBlock_NPerBlock_BK1();
-
-        // V matrix in LDS memory, dst of blockwise copy
-        static constexpr auto v_block_desc_k0_n_k1 =
-            GetBBlockDescriptor_BK0PerBlock_NPerBlock_BK1();
-
-        // dY matrix in LDS memory, dst of blockwise copy
-        static constexpr auto ygrad_block_desc_k0_m_k1 =
-            GetABlockDescriptor_AK0PerBlock_MPerBlock_AK1();
+    using ZGridDescriptor_M0_N0_M1_N1_M2_N2_M3_N3_N4_N5 = remove_cvref_t<decltype(
+        MakeCGridDescriptor_M0_N0_M1_N1_M2_N2_M3_N3_N4_N5(ZGridDesc_M_N{}))>;
 
+    // S / dP Gemm (type 1 rcr)
+    struct Gemm0
+    {
         // A matrix in LDS memory, dst of blockwise copy
         static constexpr auto a_block_desc_ak0_m_ak1 =
             GetABlockDescriptor_AK0PerBlock_MPerBlock_AK1();
@@ -317,10 +426,30 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
         static constexpr auto b_block_desc_bk0_n_bk1 =
             GetBBlockDescriptor_BK0PerBlock_NPerBlock_BK1();
 
+        template <typename ABlockDesc_AK0_M_AK1>
+        __host__ __device__ static constexpr auto
+        MakeGemm0AMmaTileDescriptor_M0_M1_M2_K(const ABlockDesc_AK0_M_AK1&)
+        {
+            constexpr index_t MWaves = MPerBlock / (MXdlPerWave * MPerXdl);
+
+            return MakeGemmMmaTileDescriptor_MN0_MN1_MN2_K<MXdlPerWave, MWaves, MPerXdl>(
+                ABlockDesc_AK0_M_AK1{});
+        }
+
+        template <typename BBlockDesc_BK0_N_BK1>
+        __host__ __device__ static constexpr auto
+        MakeGemm0BMmaTileDescriptor_N0_N1_N2_K(const BBlockDesc_BK0_N_BK1&)
+        {
+            constexpr index_t NWaves = NPerBlock / (NXdlPerWave * NPerXdl);
+
+            return MakeGemmMmaTileDescriptor_MN0_MN1_MN2_K<NXdlPerWave, NWaves, NPerXdl>(
+                BBlockDesc_BK0_N_BK1{});
+        }
+
         template <typename GridDesc_K0_M_K1>
-        using QBlockwiseCopy =
+        using ABlockwiseCopy =
             ThreadGroupTensorSliceTransfer_v4r1<ThisThreadBlock,
-                                                tensor_operation::element_wise::PassThrough,
+                                                AElementwiseOperation,
                                                 tensor_operation::element_wise::PassThrough,
                                                 InMemoryDataOperationEnum::Set,
                                                 Sequence<AK0, MPerBlock, AK1>,
@@ -329,7 +458,7 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
                                                 DataType,
                                                 DataType,
                                                 GridDesc_K0_M_K1,
-                                                decltype(q_block_desc_k0_m_k1),
+                                                decltype(a_block_desc_ak0_m_ak1),
                                                 ABlockTransferSrcAccessOrder,
                                                 Sequence<1, 0, 2>,
                                                 ABlockTransferSrcVectorDim,
@@ -343,34 +472,9 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
                                                 NumGemmKPrefetchStage>;
 
         template <typename GridDesc_K0_N_K1>
-        using KBlockwiseCopy =
-            ThreadGroupTensorSliceTransfer_v4r1<ThisThreadBlock,
-                                                tensor_operation::element_wise::PassThrough,
-                                                tensor_operation::element_wise::PassThrough,
-                                                InMemoryDataOperationEnum::Set,
-                                                Sequence<BK0, NPerBlock, BK1>,
-                                                BBlockTransferThreadClusterLengths_BK0_N_BK1,
-                                                BBlockTransferThreadClusterArrangeOrder,
-                                                DataType,
-                                                DataType,
-                                                GridDesc_K0_N_K1,
-                                                decltype(k_block_desc_k0_n_k1),
-                                                BBlockTransferSrcAccessOrder,
-                                                Sequence<1, 0, 2>,
-                                                BBlockTransferSrcVectorDim,
-                                                2,
-                                                BBlockTransferSrcScalarPerVector,
-                                                BBlockTransferDstScalarPerVector_BK1,
-                                                1,
-                                                1,
-                                                true, // SrcResetCoord
-                                                true, // DstResetCoord
-                                                NumGemmKPrefetchStage>;
-
-        template <typename GridDesc_K0_N_K1>
-        using VBlockwiseCopy =
+        using BBlockwiseCopy =
             ThreadGroupTensorSliceTransfer_v4r1<ThisThreadBlock,
-                                                tensor_operation::element_wise::PassThrough,
+                                                BElementwiseOperation,
                                                 tensor_operation::element_wise::PassThrough,
                                                 InMemoryDataOperationEnum::Set,
                                                 Sequence<BK0, NPerBlock, BK1>,
@@ -379,7 +483,7 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
                                                 DataType,
                                                 DataType,
                                                 GridDesc_K0_N_K1,
-                                                decltype(v_block_desc_k0_n_k1),
+                                                decltype(b_block_desc_bk0_n_bk1),
                                                 BBlockTransferSrcAccessOrder,
                                                 Sequence<1, 0, 2>,
                                                 BBlockTransferSrcVectorDim,
@@ -392,66 +496,6 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
                                                 true, // DstResetCoord
                                                 NumGemmKPrefetchStage>;
 
-        template <typename GridDesc_K0_M_K1>
-        using YGradBlockwiseCopy =
-            ThreadGroupTensorSliceTransfer_v4r1<ThisThreadBlock,
-                                                tensor_operation::element_wise::PassThrough,
-                                                tensor_operation::element_wise::PassThrough,
-                                                InMemoryDataOperationEnum::Set,
-                                                Sequence<AK0, MPerBlock, AK1>,
-                                                ABlockTransferThreadClusterLengths_AK0_M_AK1,
-                                                ABlockTransferThreadClusterArrangeOrder,
-                                                DataType,
-                                                DataType,
-                                                GridDesc_K0_M_K1,
-                                                decltype(ygrad_block_desc_k0_m_k1),
-                                                ABlockTransferSrcAccessOrder,
-                                                Sequence<1, 0, 2>,
-                                                ABlockTransferSrcVectorDim,
-                                                2,
-                                                ABlockTransferSrcScalarPerVector,
-                                                ABlockTransferDstScalarPerVector_AK1,
-                                                1,
-                                                1,
-                                                true, // SrcResetCoord
-                                                true, // DstResetCoord
-                                                NumGemmKPrefetchStage>;
-
-        static constexpr auto gemm_tile_k_block_slice_copy_step = make_multi_index(0, NPerBlock, 0);
-        static constexpr auto gemm_tile_v_block_slice_copy_step = make_multi_index(0, NPerBlock, 0);
-    };
-
-    // S / dP Gemm (type 1 rcr)
-    struct Gemm0
-    {
-        // A matrix in LDS memory, dst of blockwise copy
-        static constexpr auto a_block_desc_ak0_m_ak1 =
-            GetABlockDescriptor_AK0PerBlock_MPerBlock_AK1();
-
-        // B matrix in LDS memory, dst of blockwise copy
-        static constexpr auto b_block_desc_bk0_n_bk1 =
-            GetBBlockDescriptor_BK0PerBlock_NPerBlock_BK1();
-
-        template <typename ABlockDesc_AK0_M_AK1>
-        __host__ __device__ static constexpr auto
-        MakeGemm0AMmaTileDescriptor_M0_M1_M2_K(const ABlockDesc_AK0_M_AK1&)
-        {
-            constexpr index_t MWaves = MPerBlock / (MXdlPerWave * MPerXdl);
-
-            return MakeGemmMmaTileDescriptor_MN0_MN1_MN2_K<MXdlPerWave, MWaves, MPerXdl>(
-                ABlockDesc_AK0_M_AK1{});
-        }
-
-        template <typename BBlockDesc_BK0_N_BK1>
-        __host__ __device__ static constexpr auto
-        MakeGemm0BMmaTileDescriptor_N0_N1_N2_K(const BBlockDesc_BK0_N_BK1&)
-        {
-            constexpr index_t NWaves = NPerBlock / (NXdlPerWave * NPerXdl);
-
-            return MakeGemmMmaTileDescriptor_MN0_MN1_MN2_K<NXdlPerWave, NWaves, NPerXdl>(
-                BBlockDesc_BK0_N_BK1{});
-        }
-
         static constexpr index_t KPack = math::max(
             math::lcm(AK1, BK1), MfmaSelector<DataType, MPerXdl, NPerXdl>::selected_mfma.k_per_blk);
 
@@ -473,6 +517,9 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
             NXdlPerWave,
             KPack,
             true>; // TransposeC
+
+        static constexpr auto a_block_slice_copy_step = make_multi_index(KPerBlock / AK1, 0, 0);
+        static constexpr auto b_block_slice_copy_step = make_multi_index(KPerBlock / BK1, 0, 0);
     };
 
     // Y / dQ Gemm (type 2 rrr)
@@ -508,23 +555,8 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
             make_tuple(AThreadSliceLength_K0, AThreadSliceLength_M, AThreadSliceLength_K1));
 
         // B matrix in LDS memory, dst of blockwise copy
-        static constexpr auto b_block_desc_bn0_k_bn1 =
-            GetBBlockDescriptor_BK0PerBlock_NPerBlock_BK1();
-
-        template <typename ABlockDesc_AK0_M_AK1>
-        __host__ __device__ static constexpr auto
-        MakeGemm1AMmaTileDescriptor_M0_M1_M2_K(const ABlockDesc_AK0_M_AK1&)
-        {
-            return MakeGemmMmaTileDescriptor_MN0_MN1_MN2_K<MXdlPerWave, 1, 1>(ABlockDesc_AK0_M_AK1{});
-        }
-
-        template <typename BBlockDesc_BK0_N_BK1>
-        __host__ __device__ static constexpr auto
-        MakeGemm1BMmaTileDescriptor_N0_N1_N2_K(const BBlockDesc_BK0_N_BK1&)
-        {
-            return MakeGemmMmaTileDescriptor_MN0_MN1_MN2_K<Gemm1NXdlPerWave, 1, 1>(
-                BBlockDesc_BK0_N_BK1{});
-        }
+        static constexpr auto b_block_desc_bk0_n_bk1 =
+            GetB1BlockDescriptor_BK0PerBlock_NPerBlock_BK1();
 
         static constexpr auto ASrcScalarPerVector = n4;
 
@@ -541,9 +573,36 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
             2,
             ASrcScalarPerVector>;
 
+        template <typename GridDesc_K0_N_K1>
+        using BBlockwiseCopy =
+            ThreadGroupTensorSliceTransfer_v4r1<ThisThreadBlock,
+                                                BElementwiseOperation,
+                                                tensor_operation::element_wise::PassThrough,
+                                                InMemoryDataOperationEnum::Set,
+                                                Sequence<B1K0, Gemm1NPerBlock, B1K1>,
+                                                B1BlockTransferThreadClusterLengths_BK0_N_BK1,
+                                                B1BlockTransferThreadClusterArrangeOrder,
+                                                DataType,
+                                                DataType,
+                                                GridDesc_K0_N_K1,
+                                                decltype(b_block_desc_bk0_n_bk1),
+                                                B1BlockTransferSrcAccessOrder,
+                                                Sequence<1, 0, 2>,
+                                                B1BlockTransferSrcVectorDim,
+                                                2,
+                                                B1BlockTransferSrcScalarPerVector,
+                                                B1BlockTransferDstScalarPerVector_BK1,
+                                                1,
+                                                1,
+                                                B1ThreadTransferSrcResetCoordinateAfterRun,
+                                                true, // DstResetCoord
+                                                NumGemmKPrefetchStage>;
+
         // for a_block_slice_copy_step to be able to address static buffers, it MUST be a
         // tuple-based container as well as containing ONLY integral constants
         static constexpr auto a_block_slice_copy_step = make_tuple(AThreadSliceLength_K0, I0, I0);
+        static constexpr auto b_block_slice_copy_step =
+            make_multi_index(Gemm1KPerBlock / B1K1, 0, 0);
 
         // selected_mfma.group_size or B1K1 <= Gemm1KPack <= selected_mfma.group_size
         // selected_mfma.k_per_blk <= Gemm1KPack
@@ -557,101 +616,14 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
         static constexpr index_t GemmKPack =
             MfmaSelector<DataType, MPerXdl, NPerXdl>::selected_mfma.group_size;
 
-        static constexpr index_t GemmMWave   = Gemm0MWaves;
-        static constexpr index_t GemmNWave   = Gemm0NWaves;
-        static constexpr index_t GemmMRepeat = MXdlPerWave;
-        static constexpr index_t GemmNRepeat = Gemm1NXdlPerWave;
-        static constexpr index_t GemmKLoop   = NPerBlock / Gemm1KPerBlock; // 128/32=4
-        static constexpr index_t B_K3 = GemmKPack;                         // 4
-        static constexpr index_t B_K2 = N3;                                // 2
-        static constexpr index_t B_K1 = Gemm1KPerBlock / B_K2 / B_K3;      // 4
-        static constexpr index_t B_K0 = GemmKLoop;                         // 4
-
-        __host__ __device__ static constexpr auto MakeBBlockDesc_N0_N1_N2_K0_K1_K2_K3()
-        {
-            const auto N0_ = b_block_desc_bn0_k_bn1.GetLength(I0);
-            const auto K_  = b_block_desc_bn0_k_bn1.GetLength(I1);
-            const auto N1_ = b_block_desc_bn0_k_bn1.GetLength(I2);
-
-            constexpr auto b_block_desc_n_k = transform_tensor_descriptor( //(64, 128)
-                b_block_desc_bn0_k_bn1,
-                make_tuple(make_merge_transform_v3_division_mod(make_tuple(N0_, N1_)), //(8, 8)
-                           make_pass_through_transform(K_)), //128
-                make_tuple(Sequence<0, 2>{}, Sequence<1>{}),
-                make_tuple(Sequence<0>{}, Sequence<1>{}));
-
-            return transform_tensor_descriptor(
-                b_block_desc_n_k,
-                make_tuple(make_unmerge_transform(make_tuple(GemmNRepeat, GemmNWave, NPerXdl)), //(2, 1, 32)
-                           make_unmerge_transform(make_tuple(B_K0, B_K1, B_K2, B_K3))), //(4, 4, 2, 4)
-                make_tuple(Sequence<0>{}, Sequence<1>{}),
-                make_tuple(Sequence<0, 1, 2>{}, Sequence<3, 4, 5, 6>{}));
-        }
-
-        static constexpr auto b_block_desc_n0_n1_n2_k0_k1_k2_k3 = MakeBBlockDesc_N0_N1_N2_K0_K1_K2_K3();
-
-        using BThreadSlice_N0_N1_N2_K0_K1_K2_K3 = Sequence<GemmNRepeat, 
-                                                           1, 
-                                                           1, 
-                                                           1, 
-                                                           B_K1, 
-                                                           1, 
-                                                           B_K3>;
-
-        static constexpr auto b_thread_desc_n0_n1_n2_k0_k1_k2_k3 = make_naive_tensor_descriptor_packed(
-                                                                    make_tuple(Number<GemmNRepeat>{}, 
-                                                                               I1, 
-                                                                               I1, 
-                                                                               I1, 
-                                                                               Number<B_K1>{}, 
-                                                                               I1, 
-                                                                               Number<B_K3>{}));
-
-        __host__ __device__ static constexpr auto MakeBThreadDesc_K0_N_K1()
-        {
-            constexpr auto b_thread_desc_n_k = transform_tensor_descriptor(
-                b_thread_desc_n0_n1_n2_k0_k1_k2_k3,
-                make_tuple(make_merge_transform_v3_division_mod(
-                               make_tuple(Number<GemmNRepeat>{}, I1, I1)),
-                           make_merge_transform_v3_division_mod(
-                               make_tuple(I1, Number<B_K1>{}, I1, Number<B_K3>{}))),
-                make_tuple(Sequence<0, 1, 2>{}, Sequence<3, 4, 5, 6>{}),
-                make_tuple(Sequence<0>{}, Sequence<1>{}));
-
-            return transform_tensor_descriptor(
-                b_thread_desc_n_k,
-                make_tuple(make_pass_through_transform(Number<GemmNRepeat>{}),
-                           make_unmerge_transform(make_tuple(Number<B_K1>{}, Number<B_K3>{}))),
-                make_tuple(Sequence<0>{}, Sequence<1>{}),
-                make_tuple(Sequence<1>{}, Sequence<0, 2>{}));
-        }
-
-        static constexpr auto b_thread_desc_k0_n_k1 = MakeBThreadDesc_K0_N_K1();
-
-        using BBlockwiseCopy = 
-            ThreadwiseTensorSliceTransfer_v2<DataType,
-                                             DataType,
-                                             decltype(b_block_desc_n0_n1_n2_k0_k1_k2_k3),
-                                             decltype(b_thread_desc_n0_n1_n2_k0_k1_k2_k3),
-                                             BThreadSlice_N0_N1_N2_K0_K1_K2_K3,
-                                             Sequence<0, 1, 2, 3, 4, 5, 6>,
-                                             6,
-                                             1,
-                                             1,
-                                             true>;
-
-        static constexpr auto b_block_slice_copy_step = make_multi_index(0, 0, 0, 1, 0, 0, 0);
-        static constexpr auto b_block_reset_copy_step = make_multi_index(0, 0, 0, -B_K0, 0, 0, 0);
-
-        using BlockwiseGemm =
-            BlockwiseGemmXdlops_v2<
+        using BlockwiseGemm = BlockwiseGemmXdlops_v2<
             BlockSize,
             DataType,
             FloatGemmAcc,
             decltype(a_thread_desc_k0_m_k1),
-            decltype(b_thread_desc_k0_n_k1),
+            decltype(b_block_desc_bk0_n_bk1),
             decltype(MakeGemm1AMmaTileDescriptor_M0_M1_M2_K(a_thread_desc_k0_m_k1)),
-            decltype(MakeGemm1BMmaTileDescriptor_N0_N1_N2_K(b_thread_desc_k0_n_k1)),
+            decltype(MakeGemm1BMmaTileDescriptor_N0_N1_N2_K(b_block_desc_bk0_n_bk1)),
             MPerBlock,
             Gemm1NPerBlock,
             Gemm1KPerBlock,
@@ -662,7 +634,8 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
             GemmKPack,
             true,      // TransposeC
             GemmKPack, // AMmaKStride
-            GemmKPack>;
+            GemmKPack * XdlopsGemm<DataType, MPerXdl, NPerXdl, GemmKPack, false>{}
+                            .K0PerXdlops /* BMmaKStride */>;
     };
 
     // dV / dK Gemm (type 3 crr)
@@ -671,7 +644,6 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
     struct Gemm2Params_N_O_M_
     {
         static constexpr index_t Free0_N = NPerBlock;
-        static constexpr index_t Free1_M = MPerBlock;
         static constexpr index_t Free1_O = Gemm1NPerBlock;
         static constexpr index_t Sum_M   = Sum_M_;
 
@@ -679,19 +651,27 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
         static constexpr index_t A_M0     = Sum_M / A_M1;
         static constexpr index_t A_LdsPad = 0; // how many multiples of M1 per N * M1 elements
 
+        static constexpr index_t B_M1     = 2; // dY assumed row-major, typically =2 for fp16
+        static constexpr index_t B_M0     = Sum_M / B_M1;
+        static constexpr index_t B_LdsPad = 0; // how many multiples of M1 per N * M1 elements
+
         static_assert(Sum_M % MPerXdl == 0, "");
 
+        static constexpr index_t BSrcVectorDim       = 1; // Free1_O dimension
+        static constexpr index_t BSrcScalarPerVector = 4;
+
         static constexpr index_t GemmNWave   = 2;
         static constexpr index_t GemmOWave   = BlockSize / get_warp_size() / GemmNWave;
         static constexpr index_t GemmNRepeat = Free0_N / GemmNWave / MPerXdl;
         static constexpr index_t GemmORepeat = Free1_O / GemmOWave / NPerXdl;
-        static constexpr index_t GemmMLoop   = Free1_M / Sum_M;
         static constexpr index_t GemmMPack =
-            math::max(A_M1, MfmaSelector<DataType, MPerXdl, NPerXdl>::selected_mfma.k_per_blk);
-        static constexpr index_t B_M3 = GemmMPack;                                                              // 8
-        static constexpr index_t B_M2 = XdlopsGemm<DataType, MPerXdl, NPerXdl, GemmMPack, false>{}.K0PerXdlops; // 2
-        static constexpr index_t B_M1 = Sum_M / B_M2 / B_M3;                                                    // 4
-        static constexpr index_t B_M0 = GemmMLoop;                                                              // 2
+            math::max(math::lcm(A_M1, B_M1),
+                      MfmaSelector<DataType, MPerXdl, NPerXdl>::selected_mfma.k_per_blk);
+
+        using BBlockSliceLengths = Sequence<B_M0, Free1_O, B_M1>;
+        using BThreadClusterLengths =
+            Sequence<BlockSize / (Free1_O / BSrcScalarPerVector), Free1_O / BSrcScalarPerVector, 1>;
+        using BThreadClusterArrangeOrder = Sequence<0, 2, 1>;
 
         __host__ __device__ static constexpr auto GetABlockSliceLengths_M0_N0_M1_N1_M2_N2()
         {
@@ -749,23 +729,9 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
         static constexpr auto a_block_desc_m0_n_m1 =
             GetA2BlockDescriptor_M0_N_M1<Gemm2Params_N_O_M>();
 
-        // // B matrix in LDS memory, dst of blockwise copy
-        static constexpr auto b_block_desc_o0_m_o1 =
-            GetABlockDescriptor_AK0PerBlock_MPerBlock_AK1();
-
-        template <typename ABlockDesc_M0_N_M1>
-        __host__ __device__ static constexpr auto
-        MakeGemm2AMmaTileDescriptor_N0_N1_N2_M(const ABlockDesc_M0_N_M1&)
-        {
-            return MakeGemmMmaTileDescriptor_MN0_MN1_MN2_K<Gemm2Params_N_O_M::GemmNRepeat, Gemm2Params_N_O_M::GemmNWave, MPerXdl>(ABlockDesc_M0_N_M1{});
-        }
-
-        template <typename BBlockDesc_M0_O_M1>
-        __host__ __device__ static constexpr auto
-        MakeGemm2BMmaTileDescriptor_O0_O1_O2_M(const BBlockDesc_M0_O_M1&)
-        {
-            return MakeGemmMmaTileDescriptor_MN0_MN1_MN2_K<Gemm2Params_N_O_M::GemmORepeat, 1, 1>(BBlockDesc_M0_O_M1{});
-        }
+        // B matrix in LDS memory, dst of blockwise copy
+        static constexpr auto b_block_desc_m0_o_m1 =
+            GetB2BlockDescriptor_M0_O_M1<Gemm2Params_N_O_M>();
 
         __host__ __device__ static constexpr auto MakeABlockDesc_M0_N0_M1_N1_M2_N2_N3_N4()
         {
@@ -822,12 +788,13 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
                               typename Gemm2Params_N_O_M::ABlockSliceLengths_M0_N0_M1_N1,
                               false>;
 
+        template <typename ElementwiseOp = tensor_operation::element_wise::PassThrough>
         using ABlockwiseCopy = ThreadwiseTensorSliceTransfer_v1r3<
             FloatGemmAcc,
             DataType,
             decltype(a_src_thread_desc_m0_n0_m1_n1_m2_n2_n3_n4),
             decltype(a_block_desc_m0_n0_m1_n1_m2_n2_n3_n4),
-            tensor_operation::element_wise::PassThrough,
+            ElementwiseOp,
             Sequence<Gemm2Params_N_O_M::ABlockSliceLengths_M0_N0_M1_N1::At(
                          I0), // ThreadSliceLengths
                      Gemm2Params_N_O_M::ABlockSliceLengths_M0_N0_M1_N1::At(I1),
@@ -844,105 +811,50 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
             1, // DstScalarStrideInVector
             true>;
 
-        __host__ __device__ static constexpr auto MakeBBlockDesc_O0_O1_O2_M0_M1_M2_M3()
-        {
-            const auto O0_ = b_block_desc_o0_m_o1.GetLength(I0);
-            const auto M_  = b_block_desc_o0_m_o1.GetLength(I1);
-            const auto O1_ = b_block_desc_o0_m_o1.GetLength(I2);
-
-            constexpr auto b_block_desc_o_m = transform_tensor_descriptor( //(64, 128)
-                b_block_desc_o0_m_o1,
-                make_tuple(make_merge_transform_v3_division_mod(make_tuple(O0_, O1_)), //(8, 8)
-                           make_pass_through_transform(M_)), //128
-                make_tuple(Sequence<0, 2>{}, Sequence<1>{}),
-                make_tuple(Sequence<0>{}, Sequence<1>{}));
-
-            return transform_tensor_descriptor(
-                b_block_desc_o_m,
-                make_tuple(make_unmerge_transform(make_tuple(Gemm2Params_N_O_M::GemmORepeat, Gemm2Params_N_O_M::GemmOWave, NPerXdl)), //(1, 2, 32)
-                           make_unmerge_transform(make_tuple(Gemm2Params_N_O_M::B_M0, Gemm2Params_N_O_M::B_M1, Gemm2Params_N_O_M::B_M2, Gemm2Params_N_O_M::B_M3))), //(2, 4, 2, 8)
-                make_tuple(Sequence<0>{}, Sequence<1>{}),
-                make_tuple(Sequence<0, 1, 2>{}, Sequence<3, 4, 5, 6>{}));
-        }
-
-        static constexpr auto b_block_desc_o0_o1_o2_m0_m1_m2_m3 = MakeBBlockDesc_O0_O1_O2_M0_M1_M2_M3();
-
-        using BThreadSlice_O0_O1_O2_M0_M1_M2_M3 = Sequence<Gemm2Params_N_O_M::GemmORepeat, 
-                                                           1, 
-                                                           1, 
-                                                           1, 
-                                                           Gemm2Params_N_O_M::B_M1, 
-                                                           1, 
-                                                           Gemm2Params_N_O_M::B_M3>;
-
-        static constexpr auto b_thread_desc_o0_o1_o2_m0_m1_m2_m3 = make_naive_tensor_descriptor_packed(
-                                                                    make_tuple(Number<Gemm2Params_N_O_M::GemmORepeat>{}, 
-                                                                               I1, 
-                                                                               I1, 
-                                                                               I1, 
-                                                                               Number<Gemm2Params_N_O_M::B_M1>{}, 
-                                                                               I1, 
-                                                                               Number<Gemm2Params_N_O_M::B_M3>{}));
-
-        __host__ __device__ static constexpr auto MakeBThreadDesc_M0_O_M1()
-        {
-            constexpr auto b_thread_desc_o_m = transform_tensor_descriptor(
-                b_thread_desc_o0_o1_o2_m0_m1_m2_m3,
-                make_tuple(make_merge_transform_v3_division_mod(
-                               make_tuple(Number<Gemm2Params_N_O_M::GemmORepeat>{}, I1, I1)),
-                           make_merge_transform_v3_division_mod(
-                               make_tuple(I1, Number<Gemm2Params_N_O_M::B_M1>{}, I1, Number<Gemm2Params_N_O_M::B_M3>{}))),
-                make_tuple(Sequence<0, 1, 2>{}, Sequence<3, 4, 5, 6>{}),
-                make_tuple(Sequence<0>{}, Sequence<1>{}));
-
-            return transform_tensor_descriptor(
-                b_thread_desc_o_m,
-                make_tuple(make_pass_through_transform(Number<Gemm2Params_N_O_M::GemmORepeat>{}),
-                           make_unmerge_transform(make_tuple(Number<Gemm2Params_N_O_M::B_M1>{}, Number<Gemm2Params_N_O_M::B_M3>{}))),
-                make_tuple(Sequence<0>{}, Sequence<1>{}),
-                make_tuple(Sequence<1>{}, Sequence<0, 2>{}));
-        }
-
-        static constexpr auto b_thread_desc_m0_o_m1 = MakeBThreadDesc_M0_O_M1();
-
-        using BBlockwiseCopy = 
-            ThreadwiseTensorSliceTransfer_v2<DataType,
-                                             DataType,
-                                             decltype(b_block_desc_o0_o1_o2_m0_m1_m2_m3),
-                                             decltype(b_thread_desc_o0_o1_o2_m0_m1_m2_m3),
-                                             BThreadSlice_O0_O1_O2_M0_M1_M2_M3,
-                                             Sequence<0, 1, 2, 3, 4, 5, 6>,
-                                             6,
-                                             1,
-                                             1,
-                                             true>;
-
-        static constexpr auto b_block_slice_copy_step = make_multi_index(0, 0, 0, 1, 0, 0, 0);
-        static constexpr auto b_block_reset_copy_step = make_multi_index(0, 0, 0, -Gemm2Params_N_O_M::B_M0, 0, 0, 0);
-
-        using BlockwiseGemm =
-            BlockwiseGemmXdlops_v2<
-            BlockSize,
+        template <typename GridDesc_M0_O_M1>
+        using BBlockwiseCopy = ThreadGroupTensorSliceTransfer_v4r1<
+            ThisThreadBlock,
+            tensor_operation::element_wise::PassThrough,
+            tensor_operation::element_wise::PassThrough,
+            InMemoryDataOperationEnum::Set,
+            typename Gemm2Params_N_O_M::BBlockSliceLengths,
+            typename Gemm2Params_N_O_M::BThreadClusterLengths,
+            typename Gemm2Params_N_O_M::BThreadClusterArrangeOrder,
             DataType,
-            FloatGemmAcc,
-            decltype(a_block_desc_m0_n_m1),
-            decltype(b_thread_desc_m0_o_m1),
-            decltype(MakeGemm2AMmaTileDescriptor_N0_N1_N2_M(a_block_desc_m0_n_m1)),
-            decltype(MakeGemm2BMmaTileDescriptor_O0_O1_O2_M(b_thread_desc_m0_o_m1)),
-            NPerBlock,
-            Gemm1NPerBlock,
-            Gemm2Params_N_O_M::Sum_M,
-            MPerXdl,
-            NPerXdl,
-            Gemm2Params_N_O_M::GemmNRepeat,
-            Gemm2Params_N_O_M::GemmORepeat,
-            Gemm2Params_N_O_M::GemmMPack,
-            true,       // TransposeC
-            Gemm2Params_N_O_M::GemmMPack * XdlopsGemm<DataType, MPerXdl, NPerXdl, Gemm2Params_N_O_M::GemmMPack, false>{}.K0PerXdlops,
-            Gemm2Params_N_O_M::GemmMPack>;
+            DataType,
+            GridDesc_M0_O_M1,
+            decltype(b_block_desc_m0_o_m1),
+            typename Gemm2Params_N_O_M::BThreadClusterArrangeOrder, // access order == thread order
+            Sequence<1, 0, 2>,
+            Gemm2Params_N_O_M::BSrcVectorDim,
+            2, // DstVectorDim
+            Gemm2Params_N_O_M::BSrcScalarPerVector,
+            Gemm2Params_N_O_M::B_M1,
+            1,
+            1,
+            true,
+            true,
+            1>;
 
+        using BlockwiseGemm =
+            BlockwiseGemmXdlops_k0mk1_k0nk1_m0n0m1n1m2m3m4n2_v1<BlockSize,
+                                                                DataType,
+                                                                FloatGemmAcc,
+                                                                decltype(a_block_desc_m0_n_m1),
+                                                                decltype(b_block_desc_m0_o_m1),
+                                                                MPerXdl,
+                                                                NPerXdl,
+                                                                Gemm2Params_N_O_M::GemmNRepeat,
+                                                                Gemm2Params_N_O_M::GemmORepeat,
+                                                                Gemm2Params_N_O_M::GemmMPack,
+                                                                true>; // TranspossC
+
+        static constexpr auto b_block_slice_copy_step =
+            make_multi_index(Gemm2Params_N_O_M::B_M0, 0, 0);
         static constexpr auto c_block_slice_copy_step =
             make_multi_index(Gemm2Params_N_O_M::GemmNRepeat, 0, 0, 0, 0, 0, 0, 0);
+        static constexpr auto b_block_reset_copy_step =
+            make_multi_index(-MPerBlock / Gemm2Params_N_O_M::B_M1, 0, 0);
 
         template <typename CGradDesc_N_O>
         __host__ __device__ static auto
@@ -1001,31 +913,6 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
         static constexpr auto ThreadSliceLength_M =
             Number<BlockSliceLength_M_ * ThreadClusterLength_O / BlockSize_>{};
 
-        // dY matrix in LDS memory, dst of blockwise copy
-        static constexpr auto ygrad_block_desc_k0_m_k1 =
-            GetABlockDescriptor_AK0PerBlock_MPerBlock_AK1();
-
-        __host__ __device__ static constexpr auto MakeYGradBlockDesc_M0_K0_M1_K1()
-        {
-            const auto K0_ = ygrad_block_desc_k0_m_k1.GetLength(I0);
-            const auto M_  = ygrad_block_desc_k0_m_k1.GetLength(I1);
-            const auto K1_ = ygrad_block_desc_k0_m_k1.GetLength(I2);
-
-            constexpr auto ygrad_block_desc_k_m = transform_tensor_descriptor( //(64, 128)
-                ygrad_block_desc_k0_m_k1,
-                make_tuple(make_merge_transform_v3_division_mod(make_tuple(K0_, K1_)), //(8, 8)
-                           make_pass_through_transform(M_)), //128
-                make_tuple(Sequence<0, 2>{}, Sequence<1>{}),
-                make_tuple(Sequence<0>{}, Sequence<1>{}));
-
-            return transform_tensor_descriptor( //(32, 8, 4, 8)
-                ygrad_block_desc_k_m,
-                make_tuple(make_unmerge_transform(make_tuple(ThreadClusterLength_O, ThreadSliceLength_O)),  //(8, 8)
-                           make_unmerge_transform(make_tuple(ThreadClusterLength_M, ThreadSliceLength_M))), //(32, 4)
-                make_tuple(Sequence<0>{}, Sequence<1>{}),
-                make_tuple(Sequence<1, 3>{}, Sequence<0, 2>{}));
-        }
-
         static_assert(ThreadClusterLength_O * ThreadSliceLength_O == BlockSliceLength_O_, "");
         static_assert(ThreadClusterLength_M * ThreadSliceLength_M == BlockSliceLength_M_, "");
 
@@ -1039,6 +926,53 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
     };
     using YDotYGrad_M_O = YDotYGrad_M_O_<BlockSize, MPerBlock, Gemm1NPerBlock>;
 
+    // PGrad Gemm has the same layout as P = Q * K^T Gemm (A row-major B col-major)
+    struct PGradGemmTile_M_N_O
+    {
+        // TODO:
+        // Make all input tensors 2D and transform them into appropriate 3D form in kernel to make
+        // things more concise
+        template <typename YGradGridDesc_M0_O_M1_>
+        __device__ static auto
+        MakeYGradGridDesc_O0_M_O1(const YGradGridDesc_M0_O_M1_& ygrad_grid_desc_m0_o_m1)
+        {
+            const auto M0 = ygrad_grid_desc_m0_o_m1.GetLength(I0);
+            const auto O  = ygrad_grid_desc_m0_o_m1.GetLength(I1);
+            const auto M1 = ygrad_grid_desc_m0_o_m1.GetLength(I2);
+
+            constexpr auto Y_O1 = AK1;
+            const auto Y_O0     = O / Y_O1;
+
+            const auto ygrad_grid_desc_o0_m_o1 = transform_tensor_descriptor(
+                ygrad_grid_desc_m0_o_m1,
+                make_tuple(make_unmerge_transform(make_tuple(Y_O0, Y_O1)),
+                           make_merge_transform_v3_division_mod(make_tuple(M0, M1))),
+                make_tuple(Sequence<1>{}, Sequence<0, 2>{}),
+                make_tuple(Sequence<0, 2>{}, Sequence<1>{}));
+
+            return ygrad_grid_desc_o0_m_o1;
+        }
+
+        template <typename VGridDesc_N0_O_N1_>
+        __device__ static auto MakeVGridDesc_O0_N_O1(const VGridDesc_N0_O_N1_& v_grid_desc_n0_o_n1)
+        {
+            const auto N0 = v_grid_desc_n0_o_n1.GetLength(I0);
+            const auto O  = v_grid_desc_n0_o_n1.GetLength(I1);
+            const auto N1 = v_grid_desc_n0_o_n1.GetLength(I2);
+
+            constexpr auto V_O1 = BK1;
+            const auto V_O0     = O / V_O1;
+
+            const auto v_grid_desc_o0_n_o1 = transform_tensor_descriptor(
+                v_grid_desc_n0_o_n1,
+                make_tuple(make_unmerge_transform(make_tuple(V_O0, V_O1)),
+                           make_merge_transform_v3_division_mod(make_tuple(N0, N1))),
+                make_tuple(Sequence<1>{}, Sequence<0, 2>{}),
+                make_tuple(Sequence<0, 2>{}, Sequence<1>{}));
+
+            return v_grid_desc_o0_n_o1;
+        }
+    };
 
     // QGrad Gemm has the same layout as Y = P * V Gemm (A in acc B row-major)
     struct QGradGemmTile_M_K_N
@@ -1069,10 +1003,51 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
                 make_tuple(Sequence<0>{}, Sequence<1>{}),
                 make_tuple(Sequence<0, 1>{}, Sequence<2, 3>{}));
         }
+
+        template <typename KGridDesc_K0_N_K1_>
+        __device__ static auto MakeKGridDesc_N0_K_N1(const KGridDesc_K0_N_K1_& k_grid_desc_k0_n_k1)
+        {
+            const auto K_K0 = k_grid_desc_k0_n_k1.GetLength(I0);
+            const auto N    = k_grid_desc_k0_n_k1.GetLength(I1);
+            const auto K_K1 = k_grid_desc_k0_n_k1.GetLength(I2);
+
+            constexpr auto K_N1 = B1K1;
+            const auto K_N0     = N / K_N1;
+
+            const auto k_grid_desc_n0_k_n1 = transform_tensor_descriptor(
+                k_grid_desc_k0_n_k1,
+                make_tuple(make_unmerge_transform(make_tuple(K_N0, K_N1)),
+                           make_merge_transform_v3_division_mod(make_tuple(K_K0, K_K1))),
+                make_tuple(Sequence<1>{}, Sequence<0, 2>{}),
+                make_tuple(Sequence<0, 2>{}, Sequence<1>{}));
+
+            return k_grid_desc_n0_k_n1;
+        }
     };
 
     struct KGradGemmTile_N_K_M
     {
+        // B position
+        template <typename QGridDesc_K0_M_K1_>
+        __device__ static auto MakeQGridDesc_M0_K_M1(const QGridDesc_K0_M_K1_& q_grid_desc_k0_m_k1)
+        {
+            const auto Q_K0 = q_grid_desc_k0_m_k1.GetLength(I0);
+            const auto M    = q_grid_desc_k0_m_k1.GetLength(I1);
+            const auto Q_K1 = q_grid_desc_k0_m_k1.GetLength(I2);
+
+            constexpr auto Q_M1 = B1K1;
+            const auto Q_M0     = M / Q_M1;
+
+            const auto q_grid_desc_m0_k_m1 = transform_tensor_descriptor(
+                q_grid_desc_k0_m_k1,
+                make_tuple(make_unmerge_transform(make_tuple(Q_M0, Q_M1)),
+                           make_merge_transform_v3_division_mod(make_tuple(Q_K0, Q_K1))),
+                make_tuple(Sequence<1>{}, Sequence<0, 2>{}),
+                make_tuple(Sequence<0, 2>{}, Sequence<1>{}));
+
+            return q_grid_desc_m0_k_m1;
+        }
+
         // C position
         template <typename KGridDesc_K0_N_K1_>
         __device__ static auto MakeKGradGridDesc_N_K(const KGridDesc_K0_N_K1_& k_grid_desc_k0_n_k1)
@@ -1092,43 +1067,42 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
 
     struct SharedMemTrait
     {
-        // // LDS allocation for A and B: be careful of alignment
-        static constexpr auto q_block_desc_k0_m_k1 =
+        // LDS allocation for A and B: be careful of alignment
+        static constexpr auto a_block_desc_ak0_m_ak1 =
             GetABlockDescriptor_AK0PerBlock_MPerBlock_AK1();
-        static constexpr auto k_block_desc_k0_n_k1 =
-            GetBBlockDescriptor_BK0PerBlock_NPerBlock_BK1();
-        static constexpr auto v_block_desc_k0_n_k1 =
+        static constexpr auto b_block_desc_bk0_n_bk1 =
             GetBBlockDescriptor_BK0PerBlock_NPerBlock_BK1();
-        static constexpr auto ygrad_block_desc_k0_m_k1 =
-            GetABlockDescriptor_AK0PerBlock_MPerBlock_AK1();
-        static constexpr auto p_slash_sgrad_block_desc_m0_n_m1 =
+        static constexpr auto b1_block_desc_bk0_n_bk1 =
+            GetB1BlockDescriptor_BK0PerBlock_NPerBlock_BK1();
+        static constexpr auto a2_block_desc_m0_n_m1 =
             GetA2BlockDescriptor_M0_N_M1<Gemm2Params_N_O_M>();
+        static constexpr auto b2_block_desc_m0_o_m1 =
+            GetB2BlockDescriptor_M0_O_M1<Gemm2Params_N_O_M>();
 
         static constexpr auto max_lds_align = Number<16 / sizeof(DataType)>{};
 
-        static constexpr auto q_block_space_size_aligned = math::integer_least_multiple(
-            q_block_desc_k0_m_k1.GetElementSpaceSize(), max_lds_align);
-        static constexpr auto k_block_space_size_aligned = math::integer_least_multiple(
-            k_block_desc_k0_n_k1.GetElementSpaceSize(), max_lds_align);
-        static constexpr auto v_block_space_size_aligned = math::integer_least_multiple(
-            v_block_desc_k0_n_k1.GetElementSpaceSize(), max_lds_align);
+        static constexpr auto a_block_space_size_aligned = math::integer_least_multiple(
+            a_block_desc_ak0_m_ak1.GetElementSpaceSize(), max_lds_align);
+        static constexpr auto b_block_space_size_aligned = math::integer_least_multiple(
+            b_block_desc_bk0_n_bk1.GetElementSpaceSize(), max_lds_align);
+        static constexpr auto b1_block_space_size_aligned = math::integer_least_multiple(
+            b1_block_desc_bk0_n_bk1.GetElementSpaceSize(), max_lds_align);
+        static constexpr auto p_block_space_size_aligned = math::integer_least_multiple(
+            a2_block_desc_m0_n_m1.GetElementSpaceSize(), max_lds_align);
         static constexpr auto ygrad_block_space_size_aligned = math::integer_least_multiple(
-            ygrad_block_desc_k0_m_k1.GetElementSpaceSize(), max_lds_align);
-        static constexpr auto p_slash_sgrad_block_space_size_aligned = math::integer_least_multiple(
-            p_slash_sgrad_block_desc_m0_n_m1.GetElementSpaceSize(), max_lds_align);
+            b2_block_desc_m0_o_m1.GetElementSpaceSize(), max_lds_align);
 
-        static constexpr auto ygrad_block_space_offset   = 0;
-        static constexpr auto q_block_space_offset       = ygrad_block_space_size_aligned.value;
-        static constexpr auto k_block_space_offset       = ygrad_block_space_size_aligned.value + q_block_space_size_aligned.value;
-        static constexpr auto v_block_space_offset       = ygrad_block_space_size_aligned.value + q_block_space_size_aligned.value;
-        static constexpr auto p_slash_sgrad_block_space_offset = ygrad_block_space_size_aligned.value + q_block_space_size_aligned.value;
+        static constexpr auto a_block_space_offset  = 0;
+        static constexpr auto b_block_space_offset  = a_block_space_size_aligned.value;
+        static constexpr auto b1_block_space_offset = 0;
+        static constexpr auto a2_block_space_offset = 0;
+        static constexpr auto b2_block_space_offset = p_block_space_size_aligned.value;
 
         // LDS allocation for reduction
         static constexpr index_t reduction_space_size_aligned =
             math::integer_least_multiple(BlockSize, max_lds_align);
 
         static constexpr auto reduction_space_offset = 0;
-        // static constexpr auto reduction_space_offset = ygrad_block_space_size_aligned.value + q_block_space_size_aligned.value;
 
         // LDS allocation for C shuffle in LDS
         static constexpr auto c_shuffle_block_desc_mblock_mperblock_nblock_nperblock =
@@ -1139,24 +1113,25 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
 
     __host__ __device__ static constexpr index_t GetSharedMemoryNumberOfByte()
     {
-        const index_t k_bytes_end = (SharedMemTrait::k_block_space_offset +
-                                     SharedMemTrait::k_block_space_size_aligned) *
-                                    sizeof(DataType);
-        const index_t v_bytes_end = (SharedMemTrait::v_block_space_offset +
-                                     SharedMemTrait::v_block_space_size_aligned) *
-                                    sizeof(DataType);
-        const index_t p_slash_sgrad_bytes_end = (SharedMemTrait::p_slash_sgrad_block_space_offset +
-                                                 SharedMemTrait::p_slash_sgrad_block_space_size_aligned) *
-                                                sizeof(DataType);
+        const index_t gemm0_bytes_end = (SharedMemTrait::a_block_space_size_aligned +
+                                         SharedMemTrait::b_block_space_size_aligned) *
+                                        sizeof(DataType);
+        const index_t gemm1_bytes_end =
+            (SharedMemTrait::b1_block_space_offset + SharedMemTrait::b1_block_space_size_aligned) *
+            sizeof(DataType);
+        const index_t vgrad_gemm_bytes_end = (SharedMemTrait::p_block_space_size_aligned +
+                                              SharedMemTrait::ygrad_block_space_size_aligned) *
+                                             sizeof(DataType);
+
         const index_t softmax_bytes_end = (SharedMemTrait::reduction_space_offset +
                                            SharedMemTrait::reduction_space_size_aligned) *
                                           sizeof(FloatGemmAcc);
         const index_t c_block_bytes_end =
             SharedMemTrait::c_block_space_size * sizeof(FloatCShuffle);
 
-        return math::max(k_bytes_end,
-                         v_bytes_end,
-                         p_slash_sgrad_bytes_end,
+        return math::max(gemm0_bytes_end,
+                         gemm1_bytes_end,
+                         vgrad_gemm_bytes_end,
                          softmax_bytes_end,
                          c_block_bytes_end);
     }
@@ -1165,9 +1140,10 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
               typename Block2CTileMap,
               typename C0MatrixMask,
               typename VGradGridDescriptor_N_O,
-              typename YGradGridDesc_O0_M_O1>
+              typename YGradGridDesc_M0_O_M1>
     __device__ static void Run(const DataType* __restrict__ p_q_grid,
                                const DataType* __restrict__ p_k_grid,
+                               unsigned short* __restrict__ p_z_grid,
                                const DataType* __restrict__ p_v_grid,
                                const DataType* __restrict__ p_y_grid,
                                const FloatLSE* __restrict__ p_lse_grid,
@@ -1183,27 +1159,34 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
                                const CElementwiseOperation& c_element_op,
                                const QGridDesc_K0_M_K1& q_grid_desc_k0_m_k1,
                                const KGridDesc_K0_N_K1& k_grid_desc_k0_n_k1,
-                               const VGridDesc_O0_N_O1& v_grid_desc_o0_n_o1,
+                               const ZGridDescriptor_M0_N0_M1_N1_M2_N2_M3_N3_N4_N5&
+                                   z_grid_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5,
+                               const VGridDesc_N0_O_N1& v_grid_desc_n0_o_n1,
                                const YGridDescriptor_MBlock_MPerBlock_OBlock_OPerBlock&
                                    y_grid_desc_mblock_mperblock_oblock_operblock,
                                const LSEGridDesc_M& lse_grid_desc_m,
                                const VGradGridDescriptor_N_O& vgrad_grid_desc_n_o,
-                               const YGradGridDesc_O0_M_O1& ygrad_grid_desc_o0_m_o1,
+                               const YGradGridDesc_M0_O_M1& ygrad_grid_desc_m0_o_m1,
                                const Block2CTileMap& block_2_ctile_map,
-                               const C0MatrixMask& c0_matrix_mask)
+                               const C0MatrixMask& c0_matrix_mask,
+                               FloatGemmAcc p_dropout,
+                               ck::philox& ph)
     {
+        const ushort p_dropout_in_16bits = uint16_t(std::floor(p_dropout * 65535.0));
+        const FloatGemmAcc rp_dropout    = 1.0f / p_dropout;
+
         const auto q_grid_buf = make_dynamic_buffer<AddressSpaceEnum::Global>(
             p_q_grid, q_grid_desc_k0_m_k1.GetElementSpaceSize());
         const auto k_grid_buf = make_dynamic_buffer<AddressSpaceEnum::Global>(
             p_k_grid, k_grid_desc_k0_n_k1.GetElementSpaceSize());
         const auto v_grid_buf = make_dynamic_buffer<AddressSpaceEnum::Global>(
-            p_v_grid, v_grid_desc_o0_n_o1.GetElementSpaceSize());
+            p_v_grid, v_grid_desc_n0_o_n1.GetElementSpaceSize());
         const auto y_grid_buf = make_dynamic_buffer<AddressSpaceEnum::Global>(
             p_y_grid, y_grid_desc_mblock_mperblock_oblock_operblock.GetElementSpaceSize());
         const auto lse_grid_buf = make_dynamic_buffer<AddressSpaceEnum::Global>(
             p_lse_grid, lse_grid_desc_m.GetElementSpaceSize());
         const auto ygrad_grid_buf = make_dynamic_buffer<AddressSpaceEnum::Global>(
-            p_ygrad_grid, ygrad_grid_desc_o0_m_o1.GetElementSpaceSize());
+            p_ygrad_grid, ygrad_grid_desc_m0_o_m1.GetElementSpaceSize());
         auto vgrad_grid_buf = make_dynamic_buffer<AddressSpaceEnum::Global>(
             p_vgrad_grid, vgrad_grid_desc_n_o.GetElementSpaceSize());
         auto qgrad_grid_buf = make_dynamic_buffer<AddressSpaceEnum::Global>(
@@ -1227,81 +1210,91 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
         const index_t m_block_data_idx_on_grid =
             __builtin_amdgcn_readfirstlane(block_work_idx[I0] * MPerBlock);
 
-        // const index_t o_block_data_idx_on_grid =
-        //     __builtin_amdgcn_readfirstlane(block_work_idx[I1] * Gemm1NPerBlock);
+        const index_t o_block_data_idx_on_grid =
+            __builtin_amdgcn_readfirstlane(block_work_idx[I1] * Gemm1NPerBlock);
 
         // 6 GEMM operations are categorized into 3 buckets. SizeK == SizeO == head_dim
         // S_MNK / dP_MNO Gemm (Gemm0 rcr)
         // Y_MON / dQ_MKN Gemm (Gemm1 rrr)
         // dV_NOM / dK_NKM Gemm (Gemm2 crr)
 
-        // LDS allocation for Q / K / V / dY
-        auto q_block_buf = make_dynamic_buffer<AddressSpaceEnum::Lds>(
-            static_cast<DataType*>(p_shared) + SharedMemTrait::q_block_space_offset,
-            GemmBlockwiseCopy::q_block_desc_k0_m_k1.GetElementSpaceSize());
-
-        auto k_block_buf = make_dynamic_buffer<AddressSpaceEnum::Lds>(
-            static_cast<DataType*>(p_shared) + SharedMemTrait::k_block_space_offset,
-            GemmBlockwiseCopy::k_block_desc_k0_n_k1.GetElementSpaceSize());
-
-        auto v_block_buf = make_dynamic_buffer<AddressSpaceEnum::Lds>(
-            static_cast<DataType*>(p_shared) + SharedMemTrait::v_block_space_offset,
-            GemmBlockwiseCopy::v_block_desc_k0_n_k1.GetElementSpaceSize());
-
-        auto ygrad_block_buf = make_dynamic_buffer<AddressSpaceEnum::Lds>(
-            static_cast<DataType*>(p_shared) + SharedMemTrait::ygrad_block_space_offset,
-            GemmBlockwiseCopy::ygrad_block_desc_k0_m_k1.GetElementSpaceSize());
+        //
+        // set up S / dP Gemm (type 1 rcr)
+        //
 
-        // Q matrix blockwise copy
-        auto gemm_tile_q_blockwise_copy =
-            typename GemmBlockwiseCopy::template QBlockwiseCopy<decltype(q_grid_desc_k0_m_k1)>(
+        // Gemm0: LDS allocation for A and B: be careful of alignment
+        auto gemm0_a_block_buf = make_dynamic_buffer<AddressSpaceEnum::Lds>(
+            static_cast<DataType*>(p_shared) + SharedMemTrait::a_block_space_offset,
+            Gemm0::a_block_desc_ak0_m_ak1.GetElementSpaceSize());
+
+        auto gemm0_b_block_buf = make_dynamic_buffer<AddressSpaceEnum::Lds>(
+            static_cast<DataType*>(p_shared) + SharedMemTrait::b_block_space_offset,
+            Gemm0::b_block_desc_bk0_n_bk1.GetElementSpaceSize());
+
+        // Gemm0: gridwise GEMM pipeline
+        // Only supports LoopScheduler::Default
+        const auto gemm0_gridwise_gemm_pipeline =
+            GridwiseGemmPipeline_Selector<PipelineVer,
+                                          NumGemmKPrefetchStage,
+                                          LoopScheduler::Default>();
+        // S: A matrix blockwise copy
+        auto s_gemm_tile_q_blockwise_copy =
+            typename Gemm0::template ABlockwiseCopy<decltype(q_grid_desc_k0_m_k1)>(
                 q_grid_desc_k0_m_k1,
                 make_multi_index(0, m_block_data_idx_on_grid, 0),
                 a_element_op,
-                GemmBlockwiseCopy::q_block_desc_k0_m_k1,
+                Gemm0::a_block_desc_ak0_m_ak1,
                 make_multi_index(0, 0, 0),
                 tensor_operation::element_wise::PassThrough{});
 
-        // K matrix blockwise copy
-        auto gemm_tile_k_blockwise_copy =
-            typename GemmBlockwiseCopy::template KBlockwiseCopy<decltype(k_grid_desc_k0_n_k1)>(
+        // S: B matrix blockwise copy
+        auto s_gemm_tile_k_blockwise_copy =
+            typename Gemm0::template BBlockwiseCopy<decltype(k_grid_desc_k0_n_k1)>(
                 k_grid_desc_k0_n_k1,
                 make_multi_index(0, 0, 0), // will loop over GemmN dimension
                 b_element_op,
-                GemmBlockwiseCopy::k_block_desc_k0_n_k1,
+                Gemm0::b_block_desc_bk0_n_bk1,
                 make_multi_index(0, 0, 0),
                 tensor_operation::element_wise::PassThrough{});
 
-        // V matrix blockwise copy
-        auto gemm_tile_v_blockwise_copy =
-            typename GemmBlockwiseCopy::template VBlockwiseCopy<decltype(v_grid_desc_o0_n_o1)>(
-                v_grid_desc_o0_n_o1,
-                make_multi_index(0, 0, 0), // will loop over GemmN dimension
-                b1_element_op,
-                GemmBlockwiseCopy::v_block_desc_k0_n_k1,
-                make_multi_index(0, 0, 0),
-                tensor_operation::element_wise::PassThrough{});
-
-        // dY matrix blockwise copy
-        auto gemm_tile_ygrad_blockwise_copy =
-            typename GemmBlockwiseCopy::template YGradBlockwiseCopy<decltype(ygrad_grid_desc_o0_m_o1)>(
-                ygrad_grid_desc_o0_m_o1,
-                make_multi_index(0, m_block_data_idx_on_grid, 0),
-                a_element_op,
-                GemmBlockwiseCopy::ygrad_block_desc_k0_m_k1,
-                make_multi_index(0, 0, 0),
-                tensor_operation::element_wise::PassThrough{});
-
-        //
-        // set up S / dP Gemm (type 1 rcr)
-        //
-
         // S: blockwise gemm
         auto s_blockwise_gemm = typename Gemm0::BlockwiseGemm{}; // TransposeC
 
         auto s_slash_p_thread_buf = s_blockwise_gemm.GetCThreadBuffer();
 
+        const auto s_gemm_tile_a_block_reset_copy_step =
+            make_multi_index(-q_grid_desc_k0_m_k1.GetLength(I0), 0, 0);
+        const auto s_gemm_tile_b_block_reset_copy_step =
+            make_multi_index(-k_grid_desc_k0_n_k1.GetLength(I0), NPerBlock, 0);
+
+        const index_t num_k_block_main_loop = __builtin_amdgcn_readfirstlane(
+            (q_grid_desc_k0_m_k1.GetLength(I0) * q_grid_desc_k0_m_k1.GetLength(I2)) / KPerBlock);
 
+        // dP: transform input and output tensor descriptors
+        const auto ygrad_grid_desc_o0_m_o1 =
+            PGradGemmTile_M_N_O::MakeYGradGridDesc_O0_M_O1(ygrad_grid_desc_m0_o_m1);
+        const auto v_grid_desc_o0_n_o1 =
+            PGradGemmTile_M_N_O::MakeVGridDesc_O0_N_O1(v_grid_desc_n0_o_n1);
+
+        // dP: A matrix blockwise copy
+        auto pgrad_gemm_tile_ygrad_blockwise_copy =
+            typename Gemm0::template ABlockwiseCopy<decltype(ygrad_grid_desc_o0_m_o1)>(
+                ygrad_grid_desc_o0_m_o1,
+                make_multi_index(0, m_block_data_idx_on_grid, 0),
+                tensor_operation::element_wise::PassThrough{},
+                Gemm0::a_block_desc_ak0_m_ak1,
+                make_multi_index(0, 0, 0),
+                tensor_operation::element_wise::PassThrough{});
+
+        // dP: B matrix blockwise copy
+        auto pgrad_gemm_tile_v_blockwise_copy =
+            typename Gemm0::template BBlockwiseCopy<decltype(v_grid_desc_o0_n_o1)>(
+                v_grid_desc_o0_n_o1,
+                make_multi_index(0, 0, 0), // will loop over GemmN dimension
+                tensor_operation::element_wise::PassThrough{},
+                Gemm0::b_block_desc_bk0_n_bk1,
+                make_multi_index(0, 0, 0),
+                tensor_operation::element_wise::PassThrough{});
 
         // dP: blockwise gemm
         // we need separate blockwise gemm object because we need separate thread buffer
@@ -1309,6 +1302,14 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
 
         auto pgrad_thread_buf = pgrad_blockwise_gemm.GetCThreadBuffer();
 
+        const auto pgrad_gemm_tile_ygrad_block_reset_copy_step =
+            make_multi_index(-ygrad_grid_desc_o0_m_o1.GetLength(I0), 0, 0);
+        const auto pgrad_gemm_tile_v_block_reset_copy_step =
+            make_multi_index(-v_grid_desc_o0_n_o1.GetLength(I0), NPerBlock, 0);
+
+        const index_t num_o_block_main_loop = __builtin_amdgcn_readfirstlane(
+            (ygrad_grid_desc_o0_m_o1.GetLength(I0) * ygrad_grid_desc_o0_m_o1.GetLength(I2)) /
+            KPerBlock);
 
         //
         // set up Y / dQ Gemm (type 2 rrr)
@@ -1318,14 +1319,17 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
             Gemm1<decltype(s_blockwise_gemm.GetCThreadDescriptor_M0_N0_M1_N1_M2_N2_N3_N4()),
                   decltype(s_blockwise_gemm.GetCBlockDescriptor_M0_N0_M1_N1_M2_N2_N3_N4())>;
 
-        // Gemm1: VGPR allocation for A and B
+        // Gemm1: VGPR allocation for A and LDS allocation for B
         auto gemm1_a_thread_buf = make_static_buffer<AddressSpaceEnum::Vgpr, DataType>(
             Gemm1::a_thread_desc_k0_m_k1.GetElementSpaceSize());
 
-        auto gemm1_b_thread_buf = make_static_buffer<AddressSpaceEnum::Vgpr, DataType>(
-            Gemm1::b_thread_desc_n0_n1_n2_k0_k1_k2_k3.GetElementSpaceSize()); 
+        auto gemm1_b_block_buf = make_dynamic_buffer<AddressSpaceEnum::Lds>(
+            static_cast<DataType*>(p_shared) + SharedMemTrait::b1_block_space_offset,
+            Gemm1::b_block_desc_bk0_n_bk1.GetElementSpaceSize());
 
         // dQ: transform input and output tensor descriptors
+        const auto k_grid_desc_n0_k_n1 =
+            QGradGemmTile_M_K_N::MakeKGridDesc_N0_K_N1(k_grid_desc_k0_n_k1);
         auto qgrad_grid_desc_mblock_mperblock_kblock_kperblock =
             QGradGemmTile_M_K_N::MakeQGradGridDesc_MBlock_MPerBlock_KBlock_KPerBlock(
                 q_grid_desc_k0_m_k1);
@@ -1334,23 +1338,20 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
         auto qgrad_gemm_tile_sgrad_blockwise_copy =
             typename Gemm1::ABlockwiseCopy{tensor_operation::element_wise::PassThrough{}};
 
+        // dQ: B matrix blockwise copy
+        auto qgrad_gemm_tile_k_blockwise_copy =
+            typename Gemm1::template BBlockwiseCopy<decltype(k_grid_desc_n0_k_n1)>(
+                k_grid_desc_n0_k_n1,
+                make_multi_index(0, o_block_data_idx_on_grid, 0),
+                b1_element_op,
+                Gemm1::b_block_desc_bk0_n_bk1, // there n actually is k, k is N, so name can be
+                                               // b_block_desc_bn0_k_bn1
+                make_multi_index(0, 0, 0),
+                tensor_operation::element_wise::PassThrough{});
+
         // dQ: blockwise gemm
         auto qgrad_blockwise_gemm =
-            typename Gemm1::BlockwiseGemm{make_tuple(0, 0, 0, 0), make_tuple(0, 0, 0, 0)};
-
-        // dQ: B matrix blockwise copy        
-        auto k_thread_origin = qgrad_blockwise_gemm.CalculateBThreadOriginDataIndex();
-
-        auto qgrad_gemm_tile_k_blockwise_copy = 
-            typename Gemm1::BBlockwiseCopy{
-                Gemm1::b_block_desc_n0_n1_n2_k0_k1_k2_k3,
-                make_multi_index(0,                                     // nrepeat
-                                 k_thread_origin[I1],                   // nwave
-                                 k_thread_origin[I2],                   // nperxdl
-                                 0,                                     // k0
-                                 0,                                     // k1
-                                 k_thread_origin[I3] / Gemm1::GemmKPack,// k2
-                                 0)};                                   // k3
+            typename Gemm1::BlockwiseGemm{make_tuple(0, 0, 0, 0)}; // A_origin
 
         auto qgrad_thread_buf = qgrad_blockwise_gemm.GetCThreadBuffer();
 
@@ -1407,6 +1408,9 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
                                                   decltype(thread_cluster_desc_m_n),
                                                   decltype(thread_slice_desc_m_n)>{};
 
+        auto blockwise_dropout = BlockwiseDropout<FloatGemmAcc, decltype(thread_slice_desc_m_n)>{
+            p_dropout_in_16bits, rp_dropout};
+
         auto lse_grid_desc_mblock_mrepeat_mwave_mperxdl =
             MakeLSEGridDescriptor_MBlock_MRepeat_NWave_MPerXdl(lse_grid_desc_m);
 
@@ -1436,6 +1440,75 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
                                  acc0_thread_origin[I2],   // mwave
                                  acc0_thread_origin[I4])}; // mperxdl
 
+        //
+        // z vgpr copy to global
+        //
+        // z matrix threadwise desc
+        constexpr auto z_thread_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5 =
+            make_naive_tensor_descriptor_packed(make_tuple(I1,   // MBlockId
+                                                           I1,   // NBlockID
+                                                           m0,   // MRepeat
+                                                           n0,   // NRepeat
+                                                           m1,   // MWaveId
+                                                           n1,   // NWaveId
+                                                           m2,   // MPerXdl
+                                                           n2,   // NGroupNum
+                                                           n3,   // NInputNum
+                                                           n4)); // registerNum
+
+        StaticBuffer<AddressSpaceEnum::Vgpr,
+                     unsigned short,
+                     z_thread_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5.GetElementSpaceSize(),
+                     true>
+            z_tenor_buffer;
+        z_tenor_buffer.Clear();
+        // z matrix global desc
+        /*const auto M = q_grid_desc_k0_m_k1.GetLength(I1);
+        const auto N = k_grid_desc_k0_n_k1.GetLength(I1);
+
+        auto z_grid_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5 =
+            MakeZGridDescriptor_M0_N0_M1_N1_M2_N2_M3_N3_N4_N5(M, N);*/
+
+        auto z_grid_buf = make_dynamic_buffer<AddressSpaceEnum::Global>(
+            p_z_grid, z_grid_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5.GetElementSpaceSize());
+
+        const auto wave_id     = GetGemm0WaveIdx();
+        const auto wave_m_n_id = GetGemm0WaveMNIdx(wave_id[I2]); // I2: 0~63
+
+        auto z_thread_copy_vgpr_to_global = ThreadwiseTensorSliceTransfer_v1r3<
+            ushort,
+            ushort,
+            decltype(z_thread_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5),
+            decltype(z_grid_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5),
+            tensor_operation::element_wise::PassThrough,
+            Sequence<I1, // MBlockId
+                     I1, // NBlockID
+                     m0, // MRepeat
+                     n0, // NRepeat
+                     m1, // MWaveId
+                     n1, // NWaveId
+                     m2, // MPerXdl
+                     n2, // NGroupNum
+                     n3, // NInputNum
+                     n4>,
+            Sequence<0, 1, 2, 3, 4, 5, 6, 7, 8, 9>,
+            9,  // DstVectorDim
+            n4, // DstScalarPerVector
+            InMemoryDataOperationEnum::Set,
+            1, // DstScalarStrideInVector
+            true>{z_grid_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5,
+                  make_multi_index(block_work_idx[I0], // MBlockId
+                                   0,                  // NBlockId
+                                   0,                  // mrepeat
+                                   0,                  // nrepeat
+                                   wave_id[I0],        // MWaveId
+                                   wave_id[I1],        // NWaveId
+                                   wave_m_n_id[I1],    // MPerXdl
+                                   0,                  // group
+                                   wave_m_n_id[I0],    // NInputIndex
+                                   0),
+                  tensor_operation::element_wise::PassThrough{}};
+
         //
         // set up dV / dK Gemm (type 3 crr)
         //
@@ -1443,38 +1516,38 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
 
         // Gemm2: LDS allocation for A and B: be careful of alignment
         auto gemm2_a_block_buf = make_dynamic_buffer<AddressSpaceEnum::Lds>(
-            static_cast<DataType*>(p_shared) + SharedMemTrait::p_slash_sgrad_block_space_offset,
+            static_cast<DataType*>(p_shared) + SharedMemTrait::a2_block_space_offset,
             Gemm2::a_block_desc_m0_n_m1.GetElementSpaceSize());
 
-        auto gemm2_b_thread_buf = make_static_buffer<AddressSpaceEnum::Vgpr, DataType>(
-            Gemm2::b_thread_desc_o0_o1_o2_m0_m1_m2_m3.GetElementSpaceSize());
+        auto gemm2_b_block_buf = make_dynamic_buffer<AddressSpaceEnum::Lds>(
+            static_cast<DataType*>(p_shared) + SharedMemTrait::b2_block_space_offset,
+            Gemm2::b_block_desc_m0_o_m1.GetElementSpaceSize());
 
         // dV: transform input and output tensor descriptors
         const auto vgrad_grid_desc_n0_o0_n1_o1_n2_o2_o3_o4 =
             Gemm2::MakeCGridDesc_N0_O0_N1_O1_N2_O2_O3_O4(vgrad_grid_desc_n_o);
 
         // dV: A matrix VGPR-to-LDS blockwise copy
-        auto vgrad_gemm_tile_p_thread_copy_vgpr_to_lds = typename Gemm2::ABlockwiseCopy{
-            Gemm2::a_block_desc_m0_n0_m1_n1_m2_n2_n3_n4,
-            Gemm2::MakeAThreadOriginOnBlock_M0_N0_M1_N1_M2_N2_N3_N4(),
-            tensor_operation::element_wise::PassThrough{}};
+        auto vgrad_gemm_tile_p_thread_copy_vgpr_to_lds =
+            typename Gemm2::template ABlockwiseCopy<tensor_operation::element_wise::Relu>{
+                Gemm2::a_block_desc_m0_n0_m1_n1_m2_n2_n3_n4,
+                Gemm2::MakeAThreadOriginOnBlock_M0_N0_M1_N1_M2_N2_N3_N4(),
+                tensor_operation::element_wise::Relu{}}; // relu(P-dropped)
+
+        // dV: B matrix global-to-LDS blockwise copy
+        auto vgrad_gemm_tile_ygrad_blockwise_copy =
+            typename Gemm2::template BBlockwiseCopy<decltype(ygrad_grid_desc_m0_o_m1)>(
+                ygrad_grid_desc_m0_o_m1,
+                make_multi_index(m_block_data_idx_on_grid / Gemm2Params_N_O_M::B_M1,
+                                 o_block_data_idx_on_grid,
+                                 0),
+                tensor_operation::element_wise::PassThrough{},
+                Gemm2::b_block_desc_m0_o_m1,
+                make_multi_index(0, 0, 0),
+                tensor_operation::element_wise::PassThrough{});
 
         // dV: blockwise gemm
-        auto v_slash_k_grad_blockwise_gemm = typename Gemm2::BlockwiseGemm{1, make_tuple(0, 0, 0, 0)};
-
-        auto q_slash_ygrad_thread_origin = v_slash_k_grad_blockwise_gemm.CalculateBThreadOriginDataIndex();
-
-        // dV: B matrix LDS-to-VGPR blockwise copy
-        auto vgrad_gemm_tile_ygrad_blockwise_copy = 
-            typename Gemm2::BBlockwiseCopy{
-                Gemm2::b_block_desc_o0_o1_o2_m0_m1_m2_m3,
-                make_multi_index(0,                                                 // orepeat
-                                 q_slash_ygrad_thread_origin[I1],                   // owave
-                                 q_slash_ygrad_thread_origin[I2],                   // nperxdl
-                                 0,                                                 // m0
-                                 0,                                                 // m1
-                                 q_slash_ygrad_thread_origin[I3] / Gemm2Params_N_O_M::GemmMPack,// m2
-                                 0)};                                               // m3
+        auto v_slash_k_grad_blockwise_gemm = typename Gemm2::BlockwiseGemm{};
 
         auto v_slash_k_grad_thread_buf = v_slash_k_grad_blockwise_gemm.GetCThreadBuffer();
 
@@ -1484,35 +1557,40 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
             make_multi_index(
                 I0, block_work_idx[I1] * Gemm2Params_N_O_M::GemmORepeat, I0, I0, I0, I0, I0, I0);
 
-        auto vgrad_thread_copy_vgpr_to_global = typename Gemm2::template CBlockwiseCopy<decltype(
-            vgrad_grid_desc_n0_o0_n1_o1_n2_o2_o3_o4)>(
-            vgrad_grid_desc_n0_o0_n1_o1_n2_o2_o3_o4,
-            vgrad_thread_origin_on_grid_n0_o0_n1_o1_n2_o2_o3_o4,
-            tensor_operation::element_wise::PassThrough{});
-
-        // dK: transform output tensor descriptors
+        auto vgrad_thread_copy_vgpr_to_global =
+            typename Gemm2::template CBlockwiseCopy<decltype(
+                                                        vgrad_grid_desc_n0_o0_n1_o1_n2_o2_o3_o4),
+                                                    tensor_operation::element_wise::Scale>(
+                vgrad_grid_desc_n0_o0_n1_o1_n2_o2_o3_o4,
+                vgrad_thread_origin_on_grid_n0_o0_n1_o1_n2_o2_o3_o4,
+                tensor_operation::element_wise::Scale{rp_dropout});
+
+        // dK: transform input and output tensor descriptors
+        const auto q_grid_desc_m0_k_m1 =
+            KGradGemmTile_N_K_M::MakeQGridDesc_M0_K_M1(q_grid_desc_k0_m_k1);
         const auto kgrad_grid_desc_n_k =
             KGradGemmTile_N_K_M::MakeKGradGridDesc_N_K(k_grid_desc_k0_n_k1);
         const auto kgrad_grid_desc_n0_o0_n1_o1_n2_o2_o3_o4 =
             Gemm2::MakeCGridDesc_N0_O0_N1_O1_N2_O2_O3_O4(kgrad_grid_desc_n_k);
 
         // dK: A matrix VGPR-to-LDS blockwise copy
-        auto kgrad_gemm_tile_sgrad_thread_copy_vgpr_to_lds = typename Gemm2::ABlockwiseCopy{
-            Gemm2::a_block_desc_m0_n0_m1_n1_m2_n2_n3_n4,
-            Gemm2::MakeAThreadOriginOnBlock_M0_N0_M1_N1_M2_N2_N3_N4(),
-            tensor_operation::element_wise::PassThrough{}};
-
-        // dK: B matrix LDS-to-VGPR blockwise copy
-        auto kgrad_gemm_tile_q_blockwise_copy = 
-            typename Gemm2::BBlockwiseCopy{
-                Gemm2::b_block_desc_o0_o1_o2_m0_m1_m2_m3,
-                make_multi_index(0,                                                 // orepeat
-                                 q_slash_ygrad_thread_origin[I1],                   // owave
-                                 q_slash_ygrad_thread_origin[I2],                   // nperxdl
-                                 0,                                                 // m0
-                                 0,                                                 // m1
-                                 q_slash_ygrad_thread_origin[I3] / Gemm2Params_N_O_M::GemmMPack,// m2
-                                 0)};                                               // m3
+        auto kgrad_gemm_tile_sgrad_thread_copy_vgpr_to_lds =
+            typename Gemm2::template ABlockwiseCopy<tensor_operation::element_wise::PassThrough>{
+                Gemm2::a_block_desc_m0_n0_m1_n1_m2_n2_n3_n4,
+                Gemm2::MakeAThreadOriginOnBlock_M0_N0_M1_N1_M2_N2_N3_N4(),
+                tensor_operation::element_wise::PassThrough{}};
+
+        // dK: B matrix global-to-LDS blockwise copy
+        auto kgrad_gemm_tile_q_blockwise_copy =
+            typename Gemm2::template BBlockwiseCopy<decltype(q_grid_desc_m0_k_m1)>(
+                q_grid_desc_m0_k_m1,
+                make_multi_index(m_block_data_idx_on_grid / Gemm2Params_N_O_M::B_M1,
+                                 o_block_data_idx_on_grid,
+                                 0),
+                tensor_operation::element_wise::PassThrough{},
+                Gemm2::b_block_desc_m0_o_m1,
+                make_multi_index(0, 0, 0),
+                tensor_operation::element_wise::PassThrough{});
 
         // dK: blockwise gemm
         /* reuse v_slash_k_grad_blockwise_gemm, v_slash_k_grad_thread_buf */
@@ -1574,21 +1652,6 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
             true /* InvalidElementAsNaN */>(y_grid_desc_mblock_mperblock_oblock_operblock,
                                             y_thread_data_on_grid_idx);
 
-        // // performs for ygrad
-        // auto ygrad_threadwise_copy = ThreadwiseTensorSliceTransfer_v2<
-        //     DataType,
-        //     DataType,
-        //     YBlockDescriptor_MBlock_MPerBlock_OBlock_OPerBlock,
-        //     decltype(y_thread_desc_m0_m1_o0_o1),
-        //     decltype(y_thread_desc_m0_m1_o0_o1.GetLengths()),
-        //     Sequence<0, 1, 2, 3>,
-        //     3,                                 // SrcVectorDim
-        //     YDotYGrad_M_O::SrcScalarPerVector, // SrcScalarPerVector
-        //     1,                                 // SrcScalarStrideInVector
-        //     true /* ResetCoordAfterRun */,
-        //     true /* InvalidElementAsNaN */>(y_block_desc_mblock_mperblock_oblock_operblock,
-        //                                     y_thread_data_on_block_idx);
-
         auto y_thread_buf                 = typename YDotYGrad_M_O::SrcBufType{};
         auto ygrad_thread_buf             = typename YDotYGrad_M_O::SrcBufType{};
         auto y_dot_ygrad_thread_accum_buf = typename YDotYGrad_M_O::DstBufType{};
@@ -1664,7 +1727,7 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
         static_for<0, YDotYGrad_M_O::ThreadSliceLength_M, 1>{}([&](auto iM) {
             const auto idx_on_block = y_thread_data_on_block_idx[I1] + iM;
             y_dot_ygrad_block_accum_buf.AtomicAdd(
-                idx_on_block, true, y_dot_ygrad_thread_accum_buf[iM]);
+                idx_on_block, true, y_dot_ygrad_thread_accum_buf[iM] * p_dropout); // p_dropoutD1
         });
         block_sync_lds();
 
@@ -1685,15 +1748,11 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
         const index_t num_gemm1_k_block_outer_loop = k_grid_desc_k0_n_k1.GetLength(I1) / NPerBlock;
         constexpr index_t num_gemm1_k_block_inner_loop = NPerBlock / Gemm1KPerBlock;
 
+        const index_t K    = k_grid_desc_k0_n_k1.GetLength(I0) * k_grid_desc_k0_n_k1.GetLength(I2);
+        const float scalar = 1.0f / std::sqrt(K);
         // Initialize dQ
         qgrad_thread_buf.Clear();
 
-        // load q
-        gemm_tile_q_blockwise_copy.Run(q_grid_desc_k0_m_k1, q_grid_buf, GemmBlockwiseCopy::q_block_desc_k0_m_k1, q_block_buf, I0);
-
-        // load ygrad
-        gemm_tile_ygrad_blockwise_copy.Run(ygrad_grid_desc_o0_m_o1, ygrad_grid_buf, GemmBlockwiseCopy::ygrad_block_desc_k0_m_k1, ygrad_block_buf, I0);
-
         // gemm1 K loop
         index_t gemm1_k_block_outer_index = 0;
         do
@@ -1705,20 +1764,23 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
             {
                 continue;
             }
-            
-            // gemm dP
-            // dP = dY * V^T
-            pgrad_thread_buf.Clear();
-            gemm_tile_v_blockwise_copy.Run(v_grid_desc_o0_n_o1, v_grid_buf, GemmBlockwiseCopy::v_block_desc_k0_n_k1, v_block_buf, I0);
-            block_sync_lds();
-            pgrad_blockwise_gemm.Run(ygrad_block_buf, v_block_buf, pgrad_thread_buf);
-
-            // gemm S
             // S = Q * K^T
-            s_slash_p_thread_buf.Clear();
-            gemm_tile_k_blockwise_copy.Run(k_grid_desc_k0_n_k1, k_grid_buf, GemmBlockwiseCopy::k_block_desc_k0_n_k1, k_block_buf, I0);
-            block_sync_lds();
-            s_blockwise_gemm.Run(q_block_buf, k_block_buf, s_slash_p_thread_buf);
+            gemm0_gridwise_gemm_pipeline.template Run<HasMainKBlockLoop>(
+                q_grid_desc_k0_m_k1,
+                Gemm0::a_block_desc_ak0_m_ak1,
+                s_gemm_tile_q_blockwise_copy,
+                q_grid_buf,
+                gemm0_a_block_buf,
+                Gemm0::a_block_slice_copy_step,
+                k_grid_desc_k0_n_k1,
+                Gemm0::b_block_desc_bk0_n_bk1,
+                s_gemm_tile_k_blockwise_copy,
+                k_grid_buf,
+                gemm0_b_block_buf,
+                Gemm0::b_block_slice_copy_step,
+                s_blockwise_gemm,
+                s_slash_p_thread_buf,
+                num_k_block_main_loop);
 
             // do MNK padding or upper triangular masking
             if constexpr(MaskOutUpperTriangle || PadN)
@@ -1768,14 +1830,14 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
                     }
                     else
                     {
-                        s_element_op(s_slash_p_thread_buf(i), s_slash_p_thread_buf[i]);
+                        s_slash_p_thread_buf(i) = scalar * s_slash_p_thread_buf[i];
                     }
                 });
             }
             else
             {
                 static_for<0, s_slash_p_thread_buf.Size(), 1>{}(
-                    [&](auto i) { s_element_op(acc_thread_buf(i), s_slash_p_thread_buf[i]); });
+                    [&](auto i) { s_slash_p_thread_buf(i) = scalar * s_slash_p_thread_buf[i]; });
             }
 
             block_sync_lds(); // wait for lds read in gemm0 blockwise gemm
@@ -1784,59 +1846,29 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
             // scaling is already performed in the preceding statements with s_element_op
             blockwise_softmax.RunWithPreCalcStats(s_slash_p_thread_buf, lse_thread_buf);
 
-            block_sync_lds(); // wait for gemm1 LDS read
-
-            // dS = P * (dP - Y_dot_dY)
-            auto& sgrad_thread_buf = pgrad_thread_buf;
-            constexpr auto pgrad_thread_tile_iterator =
-                pgrad_blockwise_gemm.MakeCThreadTileIterator();
-            constexpr auto pgrad_thread_idx_to_m_n_adaptor =
-                pgrad_blockwise_gemm.MakeCThreadIndexAdaptor8DTo2D();
-            static_for<0, pgrad_thread_tile_iterator.GetNumOfAccess(), 1>{}([&](auto i) {
-                constexpr auto pgrad_thread_idx = pgrad_thread_tile_iterator.GetIndex(i);
-                constexpr auto m =
-                    pgrad_thread_idx_to_m_n_adaptor.CalculateBottomIndex(pgrad_thread_idx)[I0];
-                // dS and P has same thread buf layout
-                sgrad_thread_buf(i) = s_slash_p_thread_buf[i] *
-                                      (pgrad_thread_buf[i] - y_dot_ygrad_thread_buf[Number<m>{}]);
-            });
-
-            // gemm dQ
-            // dQ = scalar * dS * K
+            // save z to global
+            if(p_z_grid)
             {
-                // TODO: explore using dynamic buffer for a1 thread buffer
-                // For a1_blockwise_copy, the goal is to satisfy pipeline requirements RunRead(),
-                // RunWrite(), and MoveSliceWindow(). But it is impossible to implement given that
-                // the A1 source buffer is static buffer holding the output of first GEMM and
-                // requires constexpr offset by design. Therefore, we pass tensor coordinate offset
-                // explicitly in Run() below.
-
-                // main body
-                static_for<0, num_gemm1_k_block_inner_loop, 1>{}([&](auto i) {
-                    qgrad_gemm_tile_sgrad_blockwise_copy.Run(Gemm1::a_src_thread_desc_k0_m_k1,
-                                                                Gemm1::a_block_slice_copy_step * i,
-                                                                sgrad_thread_buf,
-                                                                Gemm1::a_thread_desc_k0_m_k1,
-                                                                make_tuple(I0, I0, I0),
-                                                                gemm1_a_thread_buf);
-
-                    qgrad_gemm_tile_k_blockwise_copy.Run(Gemm1::b_block_desc_n0_n1_n2_k0_k1_k2_k3,
-                                                            k_block_buf,
-                                                            Gemm1::b_thread_desc_n0_n1_n2_k0_k1_k2_k3,
-                                                            make_tuple(I0, I0, I0, I0, I0, I0, I0),
-                                                            gemm1_b_thread_buf);
-
-                    qgrad_gemm_tile_k_blockwise_copy.MoveSrcSliceWindow(Gemm1::b_block_desc_n0_n1_n2_k0_k1_k2_k3,
-                                                                        Gemm1::b_block_slice_copy_step);
-
-                    block_sync_lds();
-
-                    qgrad_blockwise_gemm.Run(
-                        gemm1_a_thread_buf, gemm1_b_thread_buf, qgrad_thread_buf);
+                // P_dropped
+                blockwise_dropout.template ApplyDropout<decltype(s_slash_p_thread_buf),
+                                                        decltype(z_tenor_buffer),
+                                                        true>(
+                    s_slash_p_thread_buf, ph, z_tenor_buffer);
+
+                z_thread_copy_vgpr_to_global.Run(z_thread_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5,
+                                                 make_tuple(I0, I0, I0, I0, I0, I0, I0, I0, I0, I0),
+                                                 z_tenor_buffer,
+                                                 z_grid_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5,
+                                                 z_grid_buf);
+            }
+            else
+            {
+                // P_dropped
+                blockwise_dropout.template ApplyDropout<decltype(s_slash_p_thread_buf), true>(
+                    s_slash_p_thread_buf, ph);
+            }
 
-                    // block_sync_lds();
-                });
-            } // end gemm dQ
+            block_sync_lds(); // wait for gemm1 LDS read
 
             SubThreadBlock<BlockSize> gemm2_a_copy_subgroup(s_blockwise_gemm.GetWaveIdx()[I0],
                                                             s_blockwise_gemm.GetWaveIdx()[I1]);
@@ -1846,9 +1878,12 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
                           "");
 
             // TODO: tune gemm2 pipeline
-            // dV = P^T * dY
+            // dV = P_drop^T * dY
             v_slash_k_grad_thread_buf.Clear();
             static_for<0, num_gemm2_loop, 1>{}([&](auto gemm2_loop_idx) { // gemm dV
+                // load VGrad Gemm B
+                vgrad_gemm_tile_ygrad_blockwise_copy.RunRead(ygrad_grid_desc_m0_o_m1,
+                                                             ygrad_grid_buf);
 
                 // load VGrad Gemm A
                 const auto p_slice_idx =
@@ -1870,20 +1905,18 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
                         gemm2_a_block_buf);
                 }
 
-                // block_sync_lds(); // sync before write
-
-                vgrad_gemm_tile_ygrad_blockwise_copy.Run(Gemm2::b_block_desc_o0_o1_o2_m0_m1_m2_m3,
-                                                         ygrad_block_buf,
-                                                         Gemm2::b_thread_desc_o0_o1_o2_m0_m1_m2_m3,
-                                                         make_tuple(I0, I0, I0, I0, I0, I0, I0),
-                                                         gemm2_b_thread_buf);
+                // ygrad slice window is moved with MoveSrcSliceWindow() since it is dynamic buffer
+                // p slice window is moved by loop index
+                vgrad_gemm_tile_ygrad_blockwise_copy.MoveSrcSliceWindow(
+                    ygrad_grid_desc_m0_o_m1, Gemm2::b_block_slice_copy_step);
 
-                vgrad_gemm_tile_ygrad_blockwise_copy.MoveSrcSliceWindow(Gemm2::b_block_desc_o0_o1_o2_m0_m1_m2_m3,
-                                                                        Gemm2::b_block_slice_copy_step);
+                block_sync_lds(); // sync before write
+                vgrad_gemm_tile_ygrad_blockwise_copy.RunWrite(Gemm2::b_block_desc_m0_o_m1,
+                                                              gemm2_b_block_buf);
 
                 block_sync_lds(); // sync before read
                 v_slash_k_grad_blockwise_gemm.Run(
-                    gemm2_a_block_buf, gemm2_b_thread_buf, v_slash_k_grad_thread_buf);
+                    gemm2_a_block_buf, gemm2_b_block_buf, v_slash_k_grad_thread_buf);
 
             }); // end gemm dV
             // atomic_add dV
@@ -1893,10 +1926,120 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
                                                  vgrad_grid_desc_n0_o0_n1_o1_n2_o2_o3_o4,
                                                  vgrad_grid_buf);
 
+            // gemm dP
+            block_sync_lds();
+            // dP = dY * V^T
+            // assume size K == size O so HasMainKBlockLoop is the same
+            gemm0_gridwise_gemm_pipeline.template Run<HasMainKBlockLoop>(
+                ygrad_grid_desc_o0_m_o1,
+                Gemm0::a_block_desc_ak0_m_ak1, // reuse
+                pgrad_gemm_tile_ygrad_blockwise_copy,
+                ygrad_grid_buf,
+                gemm0_a_block_buf,              // reuse
+                Gemm0::a_block_slice_copy_step, // reuse
+                v_grid_desc_o0_n_o1,
+                Gemm0::b_block_desc_bk0_n_bk1, // reuse
+                pgrad_gemm_tile_v_blockwise_copy,
+                v_grid_buf,
+                gemm0_b_block_buf,              // reuse
+                Gemm0::b_block_slice_copy_step, // reuse
+                pgrad_blockwise_gemm,
+                pgrad_thread_buf,
+                num_o_block_main_loop);
+
+            // dS = P * (dP - Y_dot_dY)
+            auto& sgrad_thread_buf = pgrad_thread_buf;
+            constexpr auto pgrad_thread_tile_iterator =
+                pgrad_blockwise_gemm.MakeCThreadTileIterator();
+            constexpr auto pgrad_thread_idx_to_m_n_adaptor =
+                pgrad_blockwise_gemm.MakeCThreadIndexAdaptor8DTo2D();
+            static_for<0, pgrad_thread_tile_iterator.GetNumOfAccess(), 1>{}([&](auto i) {
+                constexpr auto pgrad_thread_idx = pgrad_thread_tile_iterator.GetIndex(i);
+                constexpr auto m =
+                    pgrad_thread_idx_to_m_n_adaptor.CalculateBottomIndex(pgrad_thread_idx)[I0];
+                // dS and P has same thread buf layout
+                if(s_slash_p_thread_buf[i] >= 0)
+                {
+                    sgrad_thread_buf(i) =
+                        s_slash_p_thread_buf[i] *
+                        (pgrad_thread_buf[i] - y_dot_ygrad_thread_buf[Number<m>{}]);
+                }
+                else
+                {
+                    sgrad_thread_buf(i) =
+                        s_slash_p_thread_buf[i] * y_dot_ygrad_thread_buf[Number<m>{}];
+                }
+            });
+
+            // gemm dQ
+            // dQ = scalar * dS * K
+            {
+                // TODO: explore using dynamic buffer for a1 thread buffer
+                // For a1_blockwise_copy, the goal is to satisfy pipeline requirements RunRead(),
+                // RunWrite(), and MoveSliceWindow(). But it is impossible to implement given that
+                // the A1 source buffer is static buffer holding the output of first GEMM and
+                // requires constexpr offset by design. Therefore, we pass tensor coordinate offset
+                // explicitly in Run() below.
+
+                // preload data into LDS
+                qgrad_gemm_tile_k_blockwise_copy.RunRead(k_grid_desc_n0_k_n1, k_grid_buf);
+
+                qgrad_gemm_tile_k_blockwise_copy.MoveSrcSliceWindow(k_grid_desc_n0_k_n1,
+                                                                    Gemm1::b_block_slice_copy_step);
+
+                block_sync_lds(); // wait for previous LDS read
+
+                qgrad_gemm_tile_k_blockwise_copy.RunWrite(Gemm1::b_block_desc_bk0_n_bk1,
+                                                          gemm1_b_block_buf);
+
+                // main body
+                if constexpr(num_gemm1_k_block_inner_loop > 1)
+                {
+                    static_for<0, num_gemm1_k_block_inner_loop - 1, 1>{}([&](auto i) {
+                        qgrad_gemm_tile_sgrad_blockwise_copy.Run(Gemm1::a_src_thread_desc_k0_m_k1,
+                                                                 Gemm1::a_block_slice_copy_step * i,
+                                                                 sgrad_thread_buf,
+                                                                 Gemm1::a_thread_desc_k0_m_k1,
+                                                                 make_tuple(I0, I0, I0),
+                                                                 gemm1_a_thread_buf);
+                        qgrad_gemm_tile_k_blockwise_copy.RunRead(k_grid_desc_n0_k_n1, k_grid_buf);
+
+                        block_sync_lds();
+
+                        qgrad_blockwise_gemm.Run(
+                            gemm1_a_thread_buf, gemm1_b_block_buf, qgrad_thread_buf);
+
+                        block_sync_lds();
+
+                        qgrad_gemm_tile_k_blockwise_copy.MoveSrcSliceWindow(
+                            k_grid_desc_n0_k_n1, Gemm1::b_block_slice_copy_step);
+
+                        qgrad_gemm_tile_k_blockwise_copy.RunWrite(Gemm1::b_block_desc_bk0_n_bk1,
+                                                                  gemm1_b_block_buf);
+                    });
+                }
+                // tail
+                {
+                    qgrad_gemm_tile_sgrad_blockwise_copy.Run(
+                        Gemm1::a_src_thread_desc_k0_m_k1,
+                        Gemm1::a_block_slice_copy_step * Number<num_gemm1_k_block_inner_loop - 1>{},
+                        sgrad_thread_buf,
+                        Gemm1::a_thread_desc_k0_m_k1,
+                        make_tuple(I0, I0, I0),
+                        gemm1_a_thread_buf);
+
+                    block_sync_lds();
+
+                    qgrad_blockwise_gemm.Run(
+                        gemm1_a_thread_buf, gemm1_b_block_buf, qgrad_thread_buf);
+                }
+            } // end gemm dQ
+
             // dK = scalar * dS^T * dQ
             v_slash_k_grad_thread_buf.Clear();
-            // gemm2_b_thread_buf.Clear();
             static_for<0, num_gemm2_loop, 1>{}([&](auto gemm2_loop_idx) { // gemm dK
+                // load KGrad Gemm B
+                kgrad_gemm_tile_q_blockwise_copy.RunRead(q_grid_desc_m0_k_m1, q_grid_buf);
 
                 // load KGrad Gemm A
                 const auto sgrad_slice_idx =
@@ -1921,18 +2064,18 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
                         gemm2_a_block_buf);
                 }
 
-                kgrad_gemm_tile_q_blockwise_copy.Run(Gemm2::b_block_desc_o0_o1_o2_m0_m1_m2_m3,
-                                                     q_block_buf,
-                                                     Gemm2::b_thread_desc_o0_o1_o2_m0_m1_m2_m3,
-                                                     make_tuple(I0, I0, I0, I0, I0, I0, I0),
-                                                     gemm2_b_thread_buf);
-
-                kgrad_gemm_tile_q_blockwise_copy.MoveSrcSliceWindow(Gemm2::b_block_desc_o0_o1_o2_m0_m1_m2_m3,
+                // kgrad slice window is moved with MoveSrcSliceWindow() since it is dynamic buffer
+                // sgrad slice window is moved by loop index
+                kgrad_gemm_tile_q_blockwise_copy.MoveSrcSliceWindow(q_grid_desc_m0_k_m1,
                                                                     Gemm2::b_block_slice_copy_step);
 
+                block_sync_lds(); // sync before write
+                kgrad_gemm_tile_q_blockwise_copy.RunWrite(Gemm2::b_block_desc_m0_o_m1,
+                                                          gemm2_b_block_buf);
+
                 block_sync_lds(); // sync before read
                 v_slash_k_grad_blockwise_gemm.Run(
-                    gemm2_a_block_buf, gemm2_b_thread_buf, v_slash_k_grad_thread_buf);
+                    gemm2_a_block_buf, gemm2_b_block_buf, v_slash_k_grad_thread_buf);
 
             }); // end gemm dK
             // atomic_add dK
@@ -1943,26 +2086,32 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
                                                  kgrad_grid_buf);
 
             // move slice window
-            gemm_tile_k_blockwise_copy.MoveSrcSliceWindow(
+            s_gemm_tile_q_blockwise_copy.MoveSrcSliceWindow(
+                q_grid_desc_k0_m_k1,
+                s_gemm_tile_a_block_reset_copy_step); // rewind K
+            s_gemm_tile_k_blockwise_copy.MoveSrcSliceWindow(
                 k_grid_desc_k0_n_k1,
-                GemmBlockwiseCopy::gemm_tile_k_block_slice_copy_step); // step N
-            gemm_tile_v_blockwise_copy.MoveSrcSliceWindow(
-                v_grid_desc_o0_n_o1,
-                GemmBlockwiseCopy::gemm_tile_v_block_slice_copy_step); // step N
+                s_gemm_tile_b_block_reset_copy_step); // rewind K and step N
             vgrad_gemm_tile_ygrad_blockwise_copy.MoveSrcSliceWindow(
-                Gemm2::b_block_desc_o0_o1_o2_m0_m1_m2_m3,
+                ygrad_grid_desc_m0_o_m1,
                 Gemm2::b_block_reset_copy_step); // rewind M
             vgrad_thread_copy_vgpr_to_global.MoveDstSliceWindow(
                 vgrad_grid_desc_n0_o0_n1_o1_n2_o2_o3_o4, Gemm2::c_block_slice_copy_step); // step N
-            qgrad_gemm_tile_k_blockwise_copy.MoveSrcSliceWindow(
-                Gemm1::b_block_desc_n0_n1_n2_k0_k1_k2_k3,
-                Gemm1::b_block_reset_copy_step); // rewind K
+            pgrad_gemm_tile_ygrad_blockwise_copy.MoveSrcSliceWindow(
+                ygrad_grid_desc_o0_m_o1, pgrad_gemm_tile_ygrad_block_reset_copy_step); // rewind O
+            pgrad_gemm_tile_v_blockwise_copy.MoveSrcSliceWindow(
+                v_grid_desc_o0_n_o1,
+                pgrad_gemm_tile_v_block_reset_copy_step); // rewind O and step N
             kgrad_gemm_tile_q_blockwise_copy.MoveSrcSliceWindow(
-                Gemm2::b_block_desc_o0_o1_o2_m0_m1_m2_m3,
+                q_grid_desc_m0_k_m1,
                 Gemm2::b_block_reset_copy_step); // rewind M
             kgrad_thread_copy_vgpr_to_global.MoveDstSliceWindow(
                 kgrad_grid_desc_n0_o0_n1_o1_n2_o2_o3_o4, Gemm2::c_block_slice_copy_step); // step N
 
+            z_thread_copy_vgpr_to_global.MoveDstSliceWindow(
+                z_grid_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5,
+                make_multi_index(0, 1, 0, 0, 0, 0, 0, 0, 0, 0));
+
         } while(++gemm1_k_block_outer_index < num_gemm1_k_block_outer_loop); // end j loop
 
         // shuffle dQ and write

library/include/ck/library/reference_tensor_operation/cpu/reference_dropout.hpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2022, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include <iostream>
+#include <sstream>
+#include <vector>
+#include <algorithm>
+
+#include "ck/tensor_operation/gpu/device/device_base.hpp"
+#include "ck/library/utility/host_tensor.hpp"
+#include "ck/library/utility/host_tensor_generator.hpp"
+
+namespace ck {
+namespace tensor_operation {
+namespace host {
+
+template <typename RefDataType, typename InDataType, typename OutDataType>
+struct ReferenceDropout : public device::BaseOperator
+{
+    // Argument
+    struct Argument : public device::BaseArgument
+    {
+        Argument(const Tensor<RefDataType>& ref,
+                 const Tensor<InDataType>& in,
+                 Tensor<OutDataType>& out,
+                 RefDataType p_dropout_in_16bits,
+                 float rp_dropout)
+            : ref_(ref),
+              in_(in),
+              out_(out),
+              p_dropout_in_16bits_(p_dropout_in_16bits),
+              rp_dropout_(ck::type_convert<OutDataType>(rp_dropout))
+        {
+        }
+        const Tensor<RefDataType>& ref_;
+        const Tensor<InDataType>& in_;
+        Tensor<OutDataType>& out_;
+        RefDataType p_dropout_in_16bits_;
+        OutDataType rp_dropout_;
+    };
+
+    // Invoker
+    struct Invoker : public device::BaseInvoker
+    {
+        float Run(const Argument& arg)
+        {
+            arg.out_.ForEach([&](auto& self, auto idx) {
+                self(idx) =
+                    arg.ref_(idx) < arg.p_dropout_in_16bits_ ? arg.in_(idx) * arg.rp_dropout_ : 0;
+            });
+            return 0;
+        }
+
+        float Run(const device::BaseArgument* p_arg,
+                  const StreamConfig& /* stream_config */ = StreamConfig{}) override
+        {
+            return Run(*dynamic_cast<const Argument*>(p_arg));
+        }
+    };
+
+    static constexpr bool IsValidCompilationParameter()
+    {
+        // TODO: properly implement this check
+        return true;
+    }
+
+    bool IsSupportedArgument(const device::BaseArgument*) override { return true; }
+
+    static auto MakeArgument(const Tensor<RefDataType>& ref,
+                             const Tensor<InDataType>& in,
+                             Tensor<OutDataType>& out,
+                             RefDataType p_dropout_in_16bits,
+                             float rp_dropout)
+    {
+        return Argument{ref, in, out, p_dropout_in_16bits, rp_dropout};
+    }
+
+    static auto MakeInvoker() { return Invoker{}; }
+
+    virtual std::unique_ptr<device::BaseInvoker> MakeInvokerPointer()
+    {
+        return std::make_unique<Invoker>(Invoker{});
+    }
+
+    std::string GetTypeString() const override
+    {
+        auto str = std::stringstream();
+
+        // clang-format off
+        str << "ReferenceDropout"
+            << std::endl;
+        // clang-format on
+
+        return str.str();
+    }
+};
+
+} // namespace host
+} // namespace tensor_operation
+} // namespace ck

library/include/ck/library/reference_tensor_operation/cpu/reference_softmax.hpp

@@ -149,13 +149,6 @@ struct ReferenceSoftmax : public device::BaseOperator
                                                   ck::type_convert<AccDataType>(
                                                       arg.sm_stats_ptr_[0](to_sm_stats_idx(idx)))) +
                             arg.beta_ * self(idx);
-                // printf(
-                //     "exponent %f, exp() = %f\n",
-                //     ck::type_convert<AccDataType>(arg.in_(idx)) -
-                //         ck::type_convert<AccDataType>(arg.sm_stats_ptr_[0](to_sm_stats_idx(idx))),
-                //     std::exp(
-                //         ck::type_convert<AccDataType>(arg.in_(idx)) -
-                //         ck::type_convert<AccDataType>(arg.sm_stats_ptr_[0](to_sm_stats_idx(idx)))));
             });
 
             return 0;