Merge branch 'mha-train-develop' into mha-train-develop-bias-shfl

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

@@ -74,16 +74,19 @@ __global__ void
             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 BGridDesc_BK0_N_BK1 bgrad_grid_desc_bk0_n_bk1,
             const D0GridDescriptor_M0_N0_M1_M2_N1_M3 d0_grid_desc_m0_n0_m1_m2_n1_m3,
             const ZGridDescriptor_M0_N0_M1_N1_M2_N2_M3_M4_M5_N3
                 c_grid_desc_m0_n0_m1_n1_m2_n2_m3_m4_m5_n3,
             const B1GridDesc_BK0_N_BK1 b1_grid_desc_bk0_n_bk1,
+            const B1GridDesc_BK0_N_BK1 b1grad_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 YGradGridDesc_M0_O_M1 ygrad_grid_desc_m0_o_m1,
             const Block2CTileMap block_2_ctile_map,
             const index_t batch_count,
+            const index_t h_ratio,
             const index_t nblock,
             const ComputeBasePtrOfStridedBatch compute_base_ptr_of_batch,
             const C0MatrixMask c0_matrix_mask,
@@ -99,21 +102,26 @@ __global__ void
     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);
+    const index_t gkv_idx = __builtin_amdgcn_readfirstlane(g_idx / h_ratio);
 
     // 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)));
+        static_cast<long_index_t>(compute_base_ptr_of_batch.GetBBasePtr(gkv_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)));
+        static_cast<long_index_t>(compute_base_ptr_of_batch.GetB1BasePtr(gkv_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 long_index_t bgrad_batch_offset = __builtin_amdgcn_readfirstlane(
+        static_cast<long_index_t>(compute_base_ptr_of_batch.GetBGradBasePtr(g_idx)));
+    const long_index_t b1grad_batch_offset = __builtin_amdgcn_readfirstlane(
+        static_cast<long_index_t>(compute_base_ptr_of_batch.GetB1GradBasePtr(g_idx)));
 
     ck::philox ph(seed, 0, offset);
     ZDataType* z_matrix_ptr = (p_z_grid == nullptr ? nullptr : p_z_grid + z_batch_offset);
@@ -124,7 +132,6 @@ __global__ void
     D0DataType* tmp_p_d0grad_grid   = nullptr;
     if constexpr(!is_same<D0DataType, void>::value)
     {
-
         const long_index_t d0_batch_offset = __builtin_amdgcn_readfirstlane(
             static_cast<long_index_t>(compute_base_ptr_of_batch.GetD0BasePtr(g_idx)));
         if(p_d0_grid != nullptr)
@@ -136,6 +143,7 @@ __global__ void
             tmp_p_d0grad_grid = p_d0grad_grid + d0_batch_offset;
         }
     }
+
     if constexpr(Deterministic)
     {
         for(index_t i = 0; i < nblock; i++)
@@ -150,9 +158,9 @@ __global__ void
                 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_kgrad_grid + bgrad_batch_offset,
                 tmp_p_d0grad_grid,
-                p_vgrad_grid + b1_batch_offset,
+                p_vgrad_grid + b1grad_batch_offset,
                 p_shared,
                 a_element_op,
                 b_element_op,
@@ -161,9 +169,11 @@ __global__ void
                 c_element_op,
                 a_grid_desc_ak0_m_ak1,
                 b_grid_desc_bk0_n_bk1,
+                bgrad_grid_desc_bk0_n_bk1,
                 d0_grid_desc_m0_n0_m1_m2_n1_m3,
                 c_grid_desc_m0_n0_m1_n1_m2_n2_m3_m4_m5_n3,
                 b1_grid_desc_bk0_n_bk1,
+                b1grad_grid_desc_bk0_n_bk1,
                 c_grid_desc_mblock_mperblock_nblock_nperblock,
                 lse_grid_desc_m,
                 ygrad_grid_desc_m0_o_m1,
@@ -188,9 +198,9 @@ __global__ void
             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_kgrad_grid + bgrad_batch_offset,
             tmp_p_d0grad_grid,
-            p_vgrad_grid + b1_batch_offset,
+            p_vgrad_grid + b1grad_batch_offset,
             p_shared,
             a_element_op,
             b_element_op,
@@ -199,9 +209,11 @@ __global__ void
             c_element_op,
             a_grid_desc_ak0_m_ak1,
             b_grid_desc_bk0_n_bk1,
+            bgrad_grid_desc_bk0_n_bk1,
             d0_grid_desc_m0_n0_m1_m2_n1_m3,
             c_grid_desc_m0_n0_m1_n1_m2_n2_m3_m4_m5_n3,
             b1_grid_desc_bk0_n_bk1,
+            b1grad_grid_desc_bk0_n_bk1,
             c_grid_desc_mblock_mperblock_nblock_nperblock,
             lse_grid_desc_m,
             ygrad_grid_desc_m0_o_m1,
@@ -233,14 +245,17 @@ __global__ void
     ignore = c_element_op;
     ignore = a_grid_desc_ak0_m_ak1;
     ignore = b_grid_desc_bk0_n_bk1;
+    ignore = bgrad_grid_desc_bk0_n_bk1;
     ignore = d0_grid_desc_m0_n0_m1_m2_n1_m3;
     ignore = c_grid_desc_m0_n0_m1_n1_m2_n2_m3_m4_m5_n3;
     ignore = b1_grid_desc_bk0_n_bk1;
+    ignore = b1grad_grid_desc_bk0_n_bk1;
     ignore = c_grid_desc_mblock_mperblock_nblock_nperblock;
     ignore = lse_grid_desc_m;
     ignore = ygrad_grid_desc_m0_o_m1;
     ignore = block_2_ctile_map;
     ignore = batch_count;
+    ignore = h_ratio;
     ignore = nblock;
     ignore = compute_base_ptr_of_batch;
     ignore = c0_matrix_mask;
@@ -612,6 +627,8 @@ struct DeviceBatchedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
                                      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,
+                                     const BGridDesc_G_N_K& bgrad_grid_desc_g_n_k,
+                                     const B1GridDesc_G_N_K& b1grad_grid_desc_g_n_k,
                                      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),
@@ -619,6 +636,8 @@ struct DeviceBatchedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
               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),
+              bgrad_grid_desc_g_n_k_(bgrad_grid_desc_g_n_k),
+              b1grad_grid_desc_g_n_k_(b1grad_grid_desc_g_n_k),
               BatchStrideLSE_(BatchStrideLSE)
         {
         }
@@ -637,6 +656,7 @@ struct DeviceBatchedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
         {
             return d0_grid_desc_g_m_n_.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));
@@ -657,6 +677,16 @@ struct DeviceBatchedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
             return g_idx * static_cast<long_index_t>(BatchStrideLSE_);
         }
 
+        __host__ __device__ constexpr long_index_t GetBGradBasePtr(index_t g_idx) const
+        {
+            return bgrad_grid_desc_g_n_k_.CalculateOffset(make_multi_index(g_idx, 0, 0));
+        }
+
+        __host__ __device__ constexpr long_index_t GetB1GradBasePtr(index_t g_idx) const
+        {
+            return b1grad_grid_desc_g_n_k_.CalculateOffset(make_multi_index(g_idx, 0, 0));
+        }
+
         private:
         AGridDesc_G_M_K a_grid_desc_g_m_k_;
         BGridDesc_G_N_K b_grid_desc_g_n_k_;
@@ -664,6 +694,8 @@ struct DeviceBatchedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
         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_;
+        BGridDesc_G_N_K bgrad_grid_desc_g_n_k_;
+        B1GridDesc_G_N_K b1grad_grid_desc_g_n_k_;
 
         index_t BatchStrideLSE_;
     };
@@ -771,6 +803,10 @@ struct DeviceBatchedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
                  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::vector<index_t>& bgrad_gs_ns_ks_lengths,
+                 const std::vector<index_t>& bgrad_gs_ns_ks_strides,
+                 const std::vector<index_t>& b1grad_gs_gemm1ns_gemm1ks_lengths,
+                 const std::vector<index_t>& b1grad_gs_gemm1ns_gemm1ks_strides,
                  const std::vector<ck::index_t>& acc0_bias_gs_ms_ns_lengths,
                  const std::vector<ck::index_t>& acc0_bias_gs_ms_ns_strides,
                  const std::vector<ck::index_t>&
@@ -800,9 +836,13 @@ struct DeviceBatchedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
                   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)},
+              bgrad_grid_desc_bk0_n_bk1_{DeviceOp::MakeBGridDescriptor_BK0_N_BK1(
+                  bgrad_gs_ns_ks_lengths, bgrad_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::MakeVGridDescriptor_O0_N_O1(
                   b1_gs_gemm1ns_gemm1ks_lengths, b1_gs_gemm1ns_gemm1ks_strides)},
+              b1grad_grid_desc_bk0_n_bk1_{DeviceOp::MakeVGridDescriptor_O0_N_O1(
+                  b1grad_gs_gemm1ns_gemm1ks_lengths, b1grad_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])},
@@ -820,6 +860,10 @@ struct DeviceBatchedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
                                                                       c_gs_ms_gemm1ns_strides)},
               z_grid_desc_g_m_n_{
                   Transform::MakeC0GridDescriptor_G_M_N(z_gs_ms_ns_lengths, z_gs_ms_ns_strides)},
+              bgrad_grid_desc_g_n_k_{Transform::MakeB0GridDescriptor_G_N_K(bgrad_gs_ns_ks_lengths,
+                                                                           bgrad_gs_ns_ks_strides)},
+              b1grad_grid_desc_g_n_k_{Transform::MakeB1GridDescriptor_G_N_K(
+                  b1grad_gs_gemm1ns_gemm1ks_lengths, b1grad_gs_gemm1ns_gemm1ks_strides)},
               y_grid_desc_mblock_mperblock_oblock_operblock_{},
               block_2_ctile_map_{GridwiseGemm::MakeDefaultBlock2CTileMap(k_grid_desc_n_k_)},
               a_element_op_{a_element_op},
@@ -841,6 +885,7 @@ struct DeviceBatchedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
               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)},
+              h_ratio_{c_grid_desc_g_m_n_.GetLength(I0) / b_grid_desc_g_n_k_.GetLength(I0)},
               p_drop_{p_drop}
         {
             // TODO: implement bias addition
@@ -880,6 +925,8 @@ struct DeviceBatchedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
                 z_grid_desc_g_m_n_,
                 b1_grid_desc_g_n_k_,
                 c_grid_desc_g_m_n_,
+                bgrad_grid_desc_g_n_k_,
+                b1grad_grid_desc_g_n_k_,
                 type_convert<index_t>(lse_grid_desc_m_.GetElementSpaceSize()));
 
             seed_   = std::get<0>(seeds);
@@ -903,7 +950,7 @@ struct DeviceBatchedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
                       << 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) << ", "
+            std::cout << "b1_grid_desc_g_n_k_: " << 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();
@@ -916,10 +963,17 @@ struct DeviceBatchedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
             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';
-
             std::cout << "d0_grid_desc_g_m_n_: " << d0_grid_desc_g_m_n_.GetLength(I0) << ", "
                       << d0_grid_desc_g_m_n_.GetLength(I1) << ", "
                       << d0_grid_desc_g_m_n_.GetLength(I2) << '\n';
+            std::cout << "bgrad_grid_desc_g_n_k_: " << bgrad_grid_desc_g_n_k_.GetLength(I0) << ", "
+                      << bgrad_grid_desc_g_n_k_.GetLength(I1) << ", "
+                      << bgrad_grid_desc_g_n_k_.GetLength(I2) << '\n';
+            // bgrad_grid_desc_g_n_k_.Print();
+            std::cout << "b1grad_grid_desc_g_n_k_: " << b1grad_grid_desc_g_n_k_.GetLength(I0)
+                      << ", " << b1grad_grid_desc_g_n_k_.GetLength(I1) << ", "
+                      << b1grad_grid_desc_g_n_k_.GetLength(I2) << '\n';
+            // b1grad_grid_desc_g_n_k_.Print();
         }
 
         // pointers
@@ -939,9 +993,11 @@ struct DeviceBatchedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
         // tensor descriptor
         AGridDesc_AK0_M_AK1 a_grid_desc_ak0_m_ak1_;
         BGridDesc_BK0_N_BK1 b_grid_desc_bk0_n_bk1_;
+        BGridDesc_BK0_N_BK1 bgrad_grid_desc_bk0_n_bk1_;
         typename GridwiseGemm::D0GridDescriptor_M0_N0_M1_M2_N1_M3 d0_grid_desc_m0_n0_m1_m2_n1_m3_;
         ZGridDesc_M_N z_grid_desc_m_n_;
         B1GridDesc_BK0_N_BK1 b1_grid_desc_bk0_n_bk1_;
+        B1GridDesc_BK0_N_BK1 b1grad_grid_desc_bk0_n_bk1_;
         YGridDesc_M_O y_grid_desc_m_o_;
         LSEGridDesc_M lse_grid_desc_m_;
         KGridDesc_N_K k_grid_desc_n_k_;
@@ -954,6 +1010,8 @@ struct DeviceBatchedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
         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_;
+        BGridDesc_G_N_K bgrad_grid_desc_g_n_k_;
+        B1GridDesc_G_N_K b1grad_grid_desc_g_n_k_;
         typename GridwiseGemm::YGridDescriptor_MBlock_MPerBlock_OBlock_OPerBlock
             y_grid_desc_mblock_mperblock_oblock_operblock_;
 
@@ -981,6 +1039,7 @@ struct DeviceBatchedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
         std::vector<index_t> c_mz_gemm1nz_strides_;
 
         index_t batch_count_;
+        index_t h_ratio_;
         ComputeBasePtrOfStridedBatch compute_base_ptr_of_batch_;
 
         float p_drop_;
@@ -1071,14 +1130,17 @@ struct DeviceBatchedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
                     arg.c_element_op_,
                     arg.a_grid_desc_ak0_m_ak1_,
                     arg.b_grid_desc_bk0_n_bk1_,
+                    arg.bgrad_grid_desc_bk0_n_bk1_,
                     arg.d0_grid_desc_m0_n0_m1_m2_n1_m3_,
                     arg.c_grid_desc_m0_n0_m1_n1_m2_n2_m3_m4_m5_n3_,
                     arg.b1_grid_desc_bk0_n_bk1_,
+                    arg.b1grad_grid_desc_bk0_n_bk1_,
                     arg.y_grid_desc_mblock_mperblock_oblock_operblock_,
                     arg.lse_grid_desc_m_,
                     arg.ygrad_grid_desc_m0_o_m1_,
                     arg.block_2_ctile_map_,
                     arg.batch_count_,
+                    arg.h_ratio_,
                     arg.block_2_ctile_map_.CalculateGridSize(arg.k_grid_desc_n_k_),
                     arg.compute_base_ptr_of_batch_,
                     arg.c0_matrix_mask_,
@@ -1144,13 +1206,14 @@ struct DeviceBatchedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
 
         // 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 b_g      = arg.b_grid_desc_g_n_k_.GetLength(I0);
         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(I0) * arg.b1_grid_desc_bk0_n_bk1_.GetLength(I2);
 
-        if(!(c_g == arg.batch_count_ && c_m == a_m && c_gemm1n == b1_gemm1n))
+        if(!(c_g == arg.batch_count_ && c_m == a_m && c_gemm1n == b1_gemm1n && c_g % b_g == 0))
         {
             return false;
         }
@@ -1189,6 +1252,17 @@ struct DeviceBatchedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
             return false;
         }
 
+        // saving dQ data with atomic_add instruction, so KzRaw must be a multiple of 2
+        if constexpr(is_same<OutputDataType, half_t>::value ||
+                     is_same<OutputDataType, bhalf_t>::value)
+        {
+            if(KzRaw % 2 != 0)
+            {
+                std::cout << "K_q must be a multiple of 2" << std::endl;
+                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];
@@ -1243,6 +1317,10 @@ struct DeviceBatchedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
                  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::vector<index_t>& bgrad_gs_ns_ks_lengths,
+                 const std::vector<index_t>& bgrad_gs_ns_ks_strides,
+                 const std::vector<index_t>& b1grad_gs_gemm1ns_gemm1ks_lengths,
+                 const std::vector<index_t>& b1grad_gs_gemm1ns_gemm1ks_strides,
                  const std::vector<ck::index_t>& acc0_bias_gs_ms_ns_lengths,
                  const std::vector<ck::index_t>& acc0_bias_gs_ms_ns_strides,
                  const std::vector<ck::index_t>&
@@ -1282,6 +1360,10 @@ struct DeviceBatchedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
                         c_gs_ms_gemm1ns_lengths,       // c_gs_ms_os_lengths
                         c_gs_ms_gemm1ns_strides,       // c_gs_ms_os_strides
                         lse_gs_ms_lengths,
+                        bgrad_gs_ns_ks_lengths,
+                        bgrad_gs_ns_ks_strides,
+                        b1grad_gs_gemm1ns_gemm1ks_lengths,
+                        b1grad_gs_gemm1ns_gemm1ks_strides,
                         acc0_bias_gs_ms_ns_lengths,
                         acc0_bias_gs_ms_ns_strides,
                         acc1_bias_gs_ms_gemm1ns_lengths, // acc1_bias_gs_ms_os_lengths
@@ -1325,6 +1407,10 @@ struct DeviceBatchedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
         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::vector<index_t>& bgrad_gs_ns_ks_lengths,
+        const std::vector<index_t>& bgrad_gs_ns_ks_strides,
+        const std::vector<index_t>& b1grad_gs_gemm1ns_gemm1ks_lengths,
+        const std::vector<index_t>& b1grad_gs_gemm1ns_gemm1ks_strides,
         const std::vector<ck::index_t>& acc0_bias_gs_ms_ns_lengths,
         const std::vector<ck::index_t>& acc0_bias_gs_ms_ns_strides,
         const std::vector<ck::index_t>&
@@ -1365,6 +1451,10 @@ struct DeviceBatchedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
             c_gs_ms_gemm1ns_lengths,       // c_gs_ms_os_lengths
             c_gs_ms_gemm1ns_strides,       // c_gs_ms_os_strides
             lse_gs_ms_lengths,
+            bgrad_gs_ns_ks_lengths,
+            bgrad_gs_ns_ks_strides,
+            b1grad_gs_gemm1ns_gemm1ks_lengths,
+            b1grad_gs_gemm1ns_gemm1ks_strides,
             acc0_bias_gs_ms_ns_lengths,
             acc0_bias_gs_ms_ns_strides,
             acc1_bias_gs_ms_gemm1ns_lengths,

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

@@ -47,8 +47,7 @@ template <typename GridwiseGemm,
           typename C0MatrixMask,
           bool HasMainKBlockLoop,
           bool IsDropout,
-          bool IsLseStoring,
-          bool Deterministic>
+          bool IsLseStoring>
 __global__ void
 #if CK_USE_LAUNCH_BOUNDS
     __launch_bounds__(CK_MAX_THREAD_PER_BLOCK, CK_MIN_BLOCK_PER_CU)
@@ -78,7 +77,7 @@ __global__ void
             const LSEGridDescriptor_M lse_grid_desc_m,
             const Block2CTileMap block_2_ctile_map,
             const index_t batch_count,
-            const index_t mblock,
+            const index_t h_ratio,
             const ComputeBasePtrOfStridedBatch compute_base_ptr_of_batch,
             const C0MatrixMask c0_matrix_mask,
             const uint8_t p_dropout_in_uint8_t,
@@ -94,13 +93,14 @@ __global__ void
     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);
+    const index_t gkv_idx = __builtin_amdgcn_readfirstlane(g_idx / h_ratio);
 
     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)));
+        static_cast<long_index_t>(compute_base_ptr_of_batch.GetBBasePtr(gkv_idx)));
     const long_index_t b1_batch_offset = __builtin_amdgcn_readfirstlane(
-        static_cast<long_index_t>(compute_base_ptr_of_batch.GetB1BasePtr(g_idx)));
+        static_cast<long_index_t>(compute_base_ptr_of_batch.GetB1BasePtr(gkv_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 z_batch_offset = __builtin_amdgcn_readfirstlane(
@@ -122,73 +122,34 @@ __global__ void
 
     const index_t z_random_matrix_offset = g_idx * raw_m_padded * raw_n_padded;
 
-    if constexpr(Deterministic)
-    {
-        for(index_t i = 0; i < mblock; i++)
-        {
-            GridwiseGemm::template Run<HasMainKBlockLoop, IsDropout, IsLseStoring>(
-                p_a_grid + a_batch_offset,
-                p_b_grid + b_batch_offset,
-                tmp_p_d0_grid,
-                p_b1_grid + b1_batch_offset,
-                p_c_grid + c_batch_offset,
-                p_z_grid == nullptr ? nullptr : p_z_grid + z_batch_offset,
-                p_lse_grid == nullptr ? nullptr : p_lse_grid + lse_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,
-                d0_griddesc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5,
-                b1_grid_desc_bk0_n_bk1,
-                c_grid_desc_mblock_mperblock_nblock_nperblock,
-                z_grid_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5,
-                lse_grid_desc_m,
-                block_2_ctile_map,
-                c0_matrix_mask,
-                p_dropout_in_uint8_t,
-                p_dropout_rescale,
-                ph,
-                z_random_matrix_offset,
-                raw_n_padded,
-                i);
-        }
-    }
-    else
-    {
-        GridwiseGemm::template Run<HasMainKBlockLoop, IsDropout, IsLseStoring>(
-            p_a_grid + a_batch_offset,
-            p_b_grid + b_batch_offset,
-            tmp_p_d0_grid,
-            p_b1_grid + b1_batch_offset,
-            p_c_grid + c_batch_offset,
-            p_z_grid == nullptr ? nullptr : p_z_grid + z_batch_offset,
-            p_lse_grid == nullptr ? nullptr : p_lse_grid + lse_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,
-            d0_griddesc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5,
-            b1_grid_desc_bk0_n_bk1,
-            c_grid_desc_mblock_mperblock_nblock_nperblock,
-            z_grid_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5,
-            lse_grid_desc_m,
-            block_2_ctile_map,
-            c0_matrix_mask,
-            p_dropout_in_uint8_t,
-            p_dropout_rescale,
-            ph,
-            z_random_matrix_offset,
-            raw_n_padded,
-            0);
-    }
+    GridwiseGemm::template Run<HasMainKBlockLoop, IsDropout, IsLseStoring>(
+        p_a_grid + a_batch_offset,
+        p_b_grid + b_batch_offset,
+        tmp_p_d0_grid,
+        p_b1_grid + b1_batch_offset,
+        p_c_grid + c_batch_offset,
+        p_z_grid == nullptr ? nullptr : p_z_grid + z_batch_offset,
+        p_lse_grid == nullptr ? nullptr : p_lse_grid + lse_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,
+        d0_griddesc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5,
+        b1_grid_desc_bk0_n_bk1,
+        c_grid_desc_mblock_mperblock_nblock_nperblock,
+        z_grid_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5,
+        lse_grid_desc_m,
+        block_2_ctile_map,
+        c0_matrix_mask,
+        p_dropout_in_uint8_t,
+        p_dropout_rescale,
+        ph,
+        z_random_matrix_offset,
+        raw_n_padded);
 #else
     ignore = p_a_grid;
     ignore = p_b_grid;
@@ -211,7 +172,7 @@ __global__ void
     ignore = lse_grid_desc_m;
     ignore = block_2_ctile_map;
     ignore = batch_count;
-    ignore = mblock;
+    ignore = h_ratio;
     ignore = compute_base_ptr_of_batch;
     ignore = c0_matrix_mask;
     ignore = p_dropout_in_uint8_t;
@@ -296,7 +257,6 @@ template <index_t NumDimG,
           index_t CShuffleBlockTransferScalarPerVector_NPerBlock,
           index_t Acc1BiasTransferSrcScalarPerVector,
           MaskingSpecialization MaskingSpec,
-          bool Deterministic,
           LoopScheduler LoopSched = LoopScheduler::Default>
 struct DeviceBatchedMultiheadAttentionForward_Xdl_CShuffle_V2
     : public DeviceBatchedMultiheadAttentionForward<NumDimG,
@@ -576,8 +536,7 @@ struct DeviceBatchedMultiheadAttentionForward_Xdl_CShuffle_V2
         Acc1BiasTransferSrcScalarPerVector,
         LoopSched,
         Transform::matrix_padder.PadN,
-        MaskingSpec != MaskingSpecialization::MaskDisabled,
-        Deterministic>;
+        MaskingSpec != MaskingSpecialization::MaskDisabled>;
 
     // Argument
     // FIXME: constness
@@ -662,7 +621,8 @@ struct DeviceBatchedMultiheadAttentionForward_Xdl_CShuffle_V2
                                 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)}
+              batch_count_{c_grid_desc_g_m_n_.GetLength(I0)},
+              h_ratio_{c_grid_desc_g_m_n_.GetLength(I0) / b_grid_desc_g_n_k_.GetLength(I0)}
         {
             // TODO ANT: implement bias addition
             ignore = p_acc1_biases;
@@ -736,10 +696,8 @@ struct DeviceBatchedMultiheadAttentionForward_Xdl_CShuffle_V2
             std::cout << "d0_grid_desc_g_m_n_: " << d0_grid_desc_g_m_n_.GetLength(I0) << ", "
                       << d0_grid_desc_g_m_n_.GetLength(I1) << ", "
                       << d0_grid_desc_g_m_n_.GetLength(I2) << '\n';
-
             std::cout << "d0_grid_desc_m_n_: " << d0_grid_desc_m_n_.GetLength(I0) << ", "
                       << d0_grid_desc_m_n_.GetLength(I1) << '\n';
-
             std::cout << "b1_grid_desc_g_n_k_: " << 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';
@@ -802,6 +760,7 @@ struct DeviceBatchedMultiheadAttentionForward_Xdl_CShuffle_V2
         std::vector<index_t> c_mz_gemm1nz_strides_;
 
         index_t batch_count_;
+        index_t h_ratio_;
         ComputeBasePtrOfStridedBatch compute_base_ptr_of_batch_;
 
         float p_dropout_;
@@ -833,9 +792,7 @@ struct DeviceBatchedMultiheadAttentionForward_Xdl_CShuffle_V2
             }
 
             const index_t grid_size =
-                (Deterministic ? 1
-                               : arg.block_2_ctile_map_.CalculateGridSize(arg.c_grid_desc_m_n_)) *
-                arg.batch_count_;
+                arg.block_2_ctile_map_.CalculateGridSize(arg.c_grid_desc_m_n_) * arg.batch_count_;
 
             // Gemm0_K
             const auto K =
@@ -843,73 +800,72 @@ struct DeviceBatchedMultiheadAttentionForward_Xdl_CShuffle_V2
 
             float ave_time = 0;
 
-            auto launch_kernel =
-                [&](auto has_main_k_block_loop_, auto is_dropout_, auto is_lse_storing_) {
-                    const auto kernel = kernel_batched_multiheadattention_forward_xdl_cshuffle_v2<
-                        GridwiseGemm,
-                        ADataType, // TODO: distiguish A/B datatype
-                        D0DataType,
-                        CDataType,
-                        ZDataType,
-                        LSEDataType,
-                        GemmAccDataType,
-                        AElementwiseOperation,
-                        BElementwiseOperation,
-                        AccElementwiseOperation,
-                        B1ElementwiseOperation,
-                        CElementwiseOperation,
-                        DeviceOp::AGridDesc_AK0_M_AK1,
-                        DeviceOp::BGridDesc_BK0_N_BK1,
-                        typename GridwiseGemm::D0GridDescriptor_M0_N0_M1_N1_M2_N2_M3_N3_N4_N5,
-                        DeviceOp::B1GridDesc_BK0_N_BK1,
-                        typename GridwiseGemm::CGridDescriptor_MBlock_MPerBlock_NBlock_NPerBlock,
-                        typename GridwiseGemm::ZGridDescriptor_M0_N0_M1_N1_M2_N2_M3_N3_N4_N5,
-                        DeviceOp::LSEGridDesc_M,
-                        typename GridwiseGemm::DefaultBlock2CTileMap,
-                        ComputeBasePtrOfStridedBatch,
-                        C0MatrixMask,
-                        has_main_k_block_loop_,
-                        is_dropout_,
-                        is_lse_storing_,
-                        Deterministic>;
-
-                    return launch_and_time_kernel(
-                        stream_config,
-                        kernel,
-                        dim3(grid_size),
-                        dim3(BlockSize),
-                        0,
-                        arg.p_a_grid_,
-                        arg.p_b_grid_,
-                        arg.p_d0_grid_,
-                        arg.p_b1_grid_,
-                        arg.p_c_grid_,
-                        arg.p_z_grid_,
-                        arg.p_lse_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.d0_grid_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5_,
-                        arg.b1_grid_desc_bk0_n_bk1_,
-                        arg.c_grid_desc_mblock_mperblock_nblock_nperblock_,
-                        arg.z_grid_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5_,
-                        arg.lse_grid_desc_m_,
-                        arg.block_2_ctile_map_,
-                        arg.batch_count_,
-                        arg.block_2_ctile_map_.CalculateGridSize(arg.c_grid_desc_m_n_),
-                        arg.compute_base_ptr_of_batch_,
-                        arg.c0_matrix_mask_,
-                        arg.p_dropout_in_uint8_t_,
-                        arg.p_dropout_rescale_,
-                        arg.seed_,
-                        arg.offset_,
-                        arg.m_raw_padded_,
-                        arg.n_raw_padded_);
-                };
+            auto launch_kernel = [&](auto has_main_k_block_loop_,
+                                     auto is_dropout_,
+                                     auto is_lse_storing_) {
+                const auto kernel = kernel_batched_multiheadattention_forward_xdl_cshuffle_v2<
+                    GridwiseGemm,
+                    ADataType, // TODO: distiguish A/B datatype
+                    D0DataType,
+                    CDataType,
+                    ZDataType,
+                    LSEDataType,
+                    GemmAccDataType,
+                    AElementwiseOperation,
+                    BElementwiseOperation,
+                    AccElementwiseOperation,
+                    B1ElementwiseOperation,
+                    CElementwiseOperation,
+                    DeviceOp::AGridDesc_AK0_M_AK1,
+                    DeviceOp::BGridDesc_BK0_N_BK1,
+                    typename GridwiseGemm::D0GridDescriptor_M0_N0_M1_N1_M2_N2_M3_N3_N4_N5,
+                    DeviceOp::B1GridDesc_BK0_N_BK1,
+                    typename GridwiseGemm::CGridDescriptor_MBlock_MPerBlock_NBlock_NPerBlock,
+                    typename GridwiseGemm::ZGridDescriptor_M0_N0_M1_N1_M2_N2_M3_N3_N4_N5,
+                    DeviceOp::LSEGridDesc_M,
+                    typename GridwiseGemm::DefaultBlock2CTileMap,
+                    ComputeBasePtrOfStridedBatch,
+                    C0MatrixMask,
+                    has_main_k_block_loop_,
+                    is_dropout_,
+                    is_lse_storing_>;
+
+                return launch_and_time_kernel(stream_config,
+                                              kernel,
+                                              dim3(grid_size),
+                                              dim3(BlockSize),
+                                              0,
+                                              arg.p_a_grid_,
+                                              arg.p_b_grid_,
+                                              arg.p_d0_grid_,
+                                              arg.p_b1_grid_,
+                                              arg.p_c_grid_,
+                                              arg.p_z_grid_,
+                                              arg.p_lse_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.d0_grid_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5_,
+                                              arg.b1_grid_desc_bk0_n_bk1_,
+                                              arg.c_grid_desc_mblock_mperblock_nblock_nperblock_,
+                                              arg.z_grid_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5_,
+                                              arg.lse_grid_desc_m_,
+                                              arg.block_2_ctile_map_,
+                                              arg.batch_count_,
+                                              arg.h_ratio_,
+                                              arg.compute_base_ptr_of_batch_,
+                                              arg.c0_matrix_mask_,
+                                              arg.p_dropout_in_uint8_t_,
+                                              arg.p_dropout_rescale_,
+                                              arg.seed_,
+                                              arg.offset_,
+                                              arg.m_raw_padded_,
+                                              arg.n_raw_padded_);
+            };
 
             // Gemm1_K is split into Gemm1_K0/K1 where K1 is known at compile time, so we only need
             // to concern Gemm0's loop
@@ -1014,12 +970,13 @@ struct DeviceBatchedMultiheadAttentionForward_Xdl_CShuffle_V2
 
         // 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 b_g       = arg.b_grid_desc_g_n_k_.GetLength(I0);
         const index_t c_m       = arg.c_grid_desc_m_n_.GetLength(I0);
         const index_t c_gemm1n  = arg.c_grid_desc_m_n_.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))
+        if(!(c_g == arg.batch_count_ && c_m == a_m && c_gemm1n == b1_gemm1n && c_g % b_g == 0))
         {
             return false;
         }

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

@@ -103,6 +103,7 @@ __global__ void
         kernel_grouped_multihead_attention_backward_qloop_xdl_cshuffle_light_v1(
             const void CK_CONSTANT_ADDRESS_SPACE* group_kernel_args,
             const index_t group_count,
+            const index_t h_ratio,
             const AElementwiseOperation a_element_op,
             const BElementwiseOperation b_element_op,
             const AccElementwiseOperation acc_element_op,
@@ -141,19 +142,26 @@ __global__ void
     const index_t num_blocks_per_batch = arg_ptr[group_id].num_blocks_per_batch_;
     const index_t g_idx                = __builtin_amdgcn_readfirstlane(
         (block_id - arg_ptr[group_id].block_start_) / (Deterministic ? 1 : num_blocks_per_batch));
+    const index_t gkv_idx = __builtin_amdgcn_readfirstlane(g_idx / h_ratio);
 
     const long_index_t a_batch_offset = __builtin_amdgcn_readfirstlane(
         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 b_batch_offset = __builtin_amdgcn_readfirstlane(static_cast<long_index_t>(
+        arg_ptr[group_id].compute_base_ptr_of_batch_.GetBBasePtr(gkv_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)));
+        arg_ptr[group_id].compute_base_ptr_of_batch_.GetB1BasePtr(gkv_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 long_index_t bgrad_batch_offset =
+        __builtin_amdgcn_readfirstlane(static_cast<long_index_t>(
+            arg_ptr[group_id].compute_base_ptr_of_batch_.GetBGradBasePtr(g_idx)));
+    const long_index_t b1grad_batch_offset =
+        __builtin_amdgcn_readfirstlane(static_cast<long_index_t>(
+            arg_ptr[group_id].compute_base_ptr_of_batch_.GetB1GradBasePtr(g_idx)));
 
     const index_t global_thread_id = get_thread_global_1d_id();
     ck::philox ph(seed, global_thread_id, offset);
@@ -168,6 +176,7 @@ __global__ void
         const long_index_t d0_batch_offset =
             __builtin_amdgcn_readfirstlane(static_cast<long_index_t>(
                 arg_ptr[group_id].compute_base_ptr_of_batch_.GetD0BasePtr(g_idx)));
+
         if(arg_ptr[group_id].p_d0_grid_ != nullptr)
             tmp_p_d0_grid = arg_ptr[group_id].p_d0_grid_ + d0_batch_offset;
         if(arg_ptr[group_id].p_d0grad_grid_)
@@ -187,9 +196,9 @@ __global__ void
                 arg_ptr[group_id].p_d_grid_ + lse_batch_offset,
                 arg_ptr[group_id].p_ygrad_grid_ + c_batch_offset,
                 arg_ptr[group_id].p_qgrad_grid_ + a_batch_offset,
-                arg_ptr[group_id].p_kgrad_grid_ + b_batch_offset,
+                arg_ptr[group_id].p_kgrad_grid_ + bgrad_batch_offset,
                 tmp_p_d0grad_grid,
-                arg_ptr[group_id].p_vgrad_grid_ + b1_batch_offset,
+                arg_ptr[group_id].p_vgrad_grid_ + b1grad_batch_offset,
                 p_shared,
                 a_element_op,
                 b_element_op,
@@ -198,9 +207,11 @@ __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].bgrad_grid_desc_bk0_n_bk1_,
                 arg_ptr[group_id].d0_grid_desc_m0_n0_m1_m2_n1_m3_,
                 arg_ptr[group_id].c_grid_desc_m0_n0_m1_n1_m2_n2_m3_m4_m5_n3_,
                 arg_ptr[group_id].b1_grid_desc_bk0_n_bk1_,
+                arg_ptr[group_id].b1grad_grid_desc_bk0_n_bk1_,
                 arg_ptr[group_id].lse_grid_desc_m_,
                 arg_ptr[group_id].ygrad_grid_desc_o0_m_o1_,
                 arg_ptr[group_id].block_2_ctile_map_,
@@ -225,9 +236,9 @@ __global__ void
             arg_ptr[group_id].p_d_grid_ + lse_batch_offset,
             arg_ptr[group_id].p_ygrad_grid_ + c_batch_offset,
             arg_ptr[group_id].p_qgrad_grid_ + a_batch_offset,
-            arg_ptr[group_id].p_kgrad_grid_ + b_batch_offset,
+            arg_ptr[group_id].p_kgrad_grid_ + bgrad_batch_offset,
             tmp_p_d0grad_grid,
-            arg_ptr[group_id].p_vgrad_grid_ + b1_batch_offset,
+            arg_ptr[group_id].p_vgrad_grid_ + b1grad_batch_offset,
             p_shared,
             a_element_op,
             b_element_op,
@@ -236,9 +247,11 @@ __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].bgrad_grid_desc_bk0_n_bk1_,
             arg_ptr[group_id].d0_grid_desc_m0_n0_m1_m2_n1_m3_,
             arg_ptr[group_id].c_grid_desc_m0_n0_m1_n1_m2_n2_m3_m4_m5_n3_,
             arg_ptr[group_id].b1_grid_desc_bk0_n_bk1_,
+            arg_ptr[group_id].b1grad_grid_desc_bk0_n_bk1_,
             arg_ptr[group_id].lse_grid_desc_m_,
             arg_ptr[group_id].ygrad_grid_desc_o0_m_o1_,
             arg_ptr[group_id].block_2_ctile_map_,
@@ -253,6 +266,7 @@ __global__ void
 #else
     ignore = group_kernel_args;
     ignore = group_count;
+    ignore = h_ratio;
     ignore = a_element_op;
     ignore = b_element_op;
     ignore = acc_element_op;
@@ -366,6 +380,12 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V1
         std::vector<index_t> lse_gs_ms_lengths;
         std::vector<index_t> lse_gs_ms_strides;
 
+        std::vector<index_t> bgrad_gs_ns_ks_lengths;
+        std::vector<index_t> bgrad_gs_ns_ks_strides;
+
+        std::vector<index_t> b1grad_gs_gemm1ns_gemm1ks_lengths;
+        std::vector<index_t> b1grad_gs_gemm1ns_gemm1ks_strides;
+
         std::vector<index_t> acc0_bias_gs_ms_ns_lengths;
         std::vector<index_t> acc0_bias_gs_ms_ns_strides;
 
@@ -576,7 +596,6 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V1
     static auto MakeD0GridDescriptor_M_N(const std::vector<ck::index_t>& acc0_bias_gs_ms_ns_lengths,
                                          const std::vector<ck::index_t>& acc0_bias_gs_ms_ns_strides)
     {
-
         return Transform::MakeC0GridDescriptor_M_N(acc0_bias_gs_ms_ns_lengths,
                                                    acc0_bias_gs_ms_ns_strides);
     }
@@ -585,7 +604,6 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V1
     MakeD0GridDescriptor_G_M_N(const std::vector<ck::index_t>& acc0_bias_gs_ms_ns_lengths,
                                const std::vector<ck::index_t>& acc0_bias_gs_ms_ns_strides)
     {
-
         return Transform::MakeC0GridDescriptor_G_M_N(acc0_bias_gs_ms_ns_lengths,
                                                      acc0_bias_gs_ms_ns_strides);
     }
@@ -625,7 +643,7 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V1
     using D0GridDesc_G_M_N     = decltype(MakeD0GridDescriptor_G_M_N({}, {}));
     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 ZGridDesc_G_M_N      = decltype(Transform::MakeC0GridDescriptor_G_M_N({}, {}));
 
     using KGridDesc_N_K         = decltype(Transform::MakeB0GridDescriptor_N_K({}, {}));
     using D0GridDesc_M_N        = decltype(MakeD0GridDescriptor_M_N({}, {}));
@@ -660,6 +678,8 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V1
                                      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,
+                                     const BGridDesc_G_N_K& bgrad_grid_desc_g_n_k,
+                                     const B1GridDesc_G_N_K& b1grad_grid_desc_g_n_k,
                                      index_t batch_stride_lse)
             : a_grid_desc_g_m_k_(a_grid_desc_g_m_k),
               b_grid_desc_g_n_k_(b_grid_desc_g_n_k),
@@ -667,6 +687,8 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V1
               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),
+              bgrad_grid_desc_g_n_k_(bgrad_grid_desc_g_n_k),
+              b1grad_grid_desc_g_n_k_(b1grad_grid_desc_g_n_k),
               batch_stride_lse_(batch_stride_lse)
         {
         }
@@ -706,6 +728,16 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V1
             return g_idx * static_cast<long_index_t>(batch_stride_lse_);
         }
 
+        __host__ __device__ constexpr long_index_t GetBGradBasePtr(index_t g_idx) const
+        {
+            return bgrad_grid_desc_g_n_k_.CalculateOffset(make_multi_index(g_idx, 0, 0));
+        }
+
+        __host__ __device__ constexpr long_index_t GetB1GradBasePtr(index_t g_idx) const
+        {
+            return b1grad_grid_desc_g_n_k_.CalculateOffset(make_multi_index(g_idx, 0, 0));
+        }
+
         private:
         AGridDesc_G_M_K a_grid_desc_g_m_k_;
         BGridDesc_G_N_K b_grid_desc_g_n_k_;
@@ -713,6 +745,8 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V1
         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_;
+        BGridDesc_G_N_K bgrad_grid_desc_g_n_k_;
+        B1GridDesc_G_N_K b1grad_grid_desc_g_n_k_;
         index_t batch_stride_lse_;
     };
 
@@ -817,9 +851,11 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V1
         // 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_;
+        BGridDesc_BK0_N_BK1 bgrad_grid_desc_bk0_n_bk1_;
         typename GridwiseGemm::D0GridDescriptor_M0_N0_M1_M2_N1_M3 d0_grid_desc_m0_n0_m1_m2_n1_m3_;
         ZGridDesc_M_N z_grid_desc_m_n_;
         B1GridDesc_BK0_N_BK1 b1_grid_desc_bk0_n_bk1_;
+        B1GridDesc_BK0_N_BK1 b1grad_grid_desc_bk0_n_bk1_;
         YGridDesc_M_O y_grid_desc_m_o_;
 
         typename GridwiseGemm::ZGridDescriptor_M0_N0_M1_N1_M2_N2_M3_M4_M5_N3
@@ -861,6 +897,7 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V1
         std::vector<index_t> c_mz_gemm1nz_strides_;
 
         // for gridwise gemm check
+        BGridDesc_G_N_K b_grid_desc_g_n_k_;
         CGridDesc_G_M_N c_grid_desc_g_m_n_;
 
         index_t batch_count_;
@@ -933,6 +970,9 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V1
 
             d_grid_size_ = 0;
 
+            h_ratio_ = problem_desc_vec[0].a_gs_ms_ks_lengths[NumDimG - 1] /
+                       problem_desc_vec[0].b_gs_ns_ks_lengths[NumDimG - 1];
+
             for(index_t i = 0; i < group_count_; i++)
             {
                 const auto p_a_grid = static_cast<const InputDataType*>(p_As[i]);
@@ -960,6 +1000,8 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V1
                     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 bgrad_grid_desc_bk0_n_bk1 = DeviceOp::MakeBGridDescriptor_BK0_N_BK1(
+                    problem_desc.bgrad_gs_ns_ks_lengths, problem_desc.bgrad_gs_ns_ks_strides);
 
                 std::vector<index_t> tmp_d0_gs_ms_ns_lengths;
                 std::vector<index_t> tmp_d0_gs_ms_ns_strides;
@@ -982,6 +1024,9 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V1
                 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);
+                const auto b1grad_grid_desc_bk0_n_bk1 = DeviceOp::MakeVGridDescriptor_O0_N_O1(
+                    problem_desc.b1grad_gs_gemm1ns_gemm1ks_lengths,
+                    problem_desc.b1grad_gs_gemm1ns_gemm1ks_strides);
                 const auto y_grid_desc_m_o = Transform::MakeCGridDescriptor_M_N(
                     problem_desc.c_gs_ms_gemm1ns_lengths, problem_desc.c_gs_ms_gemm1ns_strides);
 
@@ -1005,6 +1050,11 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V1
                     problem_desc.b1_gs_gemm1ns_gemm1ks_strides);
                 const auto c_grid_desc_g_m_n = Transform::MakeCGridDescriptor_G_M_N(
                     problem_desc.c_gs_ms_gemm1ns_lengths, problem_desc.c_gs_ms_gemm1ns_strides);
+                const auto bgrad_grid_desc_g_n_k = Transform::MakeB0GridDescriptor_G_N_K(
+                    problem_desc.bgrad_gs_ns_ks_lengths, problem_desc.bgrad_gs_ns_ks_strides);
+                const auto b1grad_grid_desc_g_n_k = Transform::MakeB1GridDescriptor_G_N_K(
+                    problem_desc.b1grad_gs_gemm1ns_gemm1ks_lengths,
+                    problem_desc.b1grad_gs_gemm1ns_gemm1ks_strides);
                 typename GridwiseGemm::ZGridDescriptor_M0_N0_M1_N1_M2_N2_M3_M4_M5_N3
                     c_grid_desc_m0_n0_m1_n1_m2_n2_m3_m4_m5_n3;
                 const index_t BlockStart     = grid_size_;
@@ -1027,7 +1077,9 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V1
                     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()));
+                    bgrad_grid_desc_g_n_k,
+                    b1grad_grid_desc_g_n_k,
+                    type_convert<index_t>(problem_desc.lse_gs_ms_strides[NumDimG - 1]));
 
                 // C0 mask
                 const auto c0_matrix_mask =
@@ -1073,9 +1125,11 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V1
                                               p_vgrad_grid,
                                               a_grid_desc_ak0_m_ak1,
                                               b_grid_desc_bk0_n_bk1,
+                                              bgrad_grid_desc_bk0_n_bk1,
                                               d0_grid_desc_m0_n0_m1_m2_n1_m3,
                                               z_grid_desc_m_n,
                                               b1_grid_desc_bk0_n_bk1,
+                                              b1grad_grid_desc_bk0_n_bk1,
                                               y_grid_desc_m_o,
                                               c_grid_desc_m0_n0_m1_n1_m2_n2_m3_m4_m5_n3,
                                               lse_grid_desc_m,
@@ -1119,6 +1173,7 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V1
                       problem_desc.b1_gs_gemm1ns_gemm1ks_strides[NumDimG + NumDimO + NumDimN - 1]},
                      {problem_desc.c_gs_ms_gemm1ns_strides[NumDimG + NumDimM - 1],
                       problem_desc.c_gs_ms_gemm1ns_strides[NumDimG + NumDimM + NumDimO - 1]},
+                     b_grid_desc_g_n_k,
                      c_grid_desc_g_m_n,
                      batch_count,
                      d0_n_length_stride});
@@ -1145,6 +1200,7 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V1
 
         index_t grid_size_;
         index_t group_count_;
+        index_t h_ratio_;
 
         std::vector<GroupKernelArg> group_kernel_args_;
         std::vector<GroupDeviceArg> group_device_args_;
@@ -1224,6 +1280,7 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V1
                     0,
                     cast_pointer_to_constant_address_space(arg.p_workspace_),
                     arg.group_count_,
+                    arg.h_ratio_,
                     arg.a_element_op_,
                     arg.b_element_op_,
                     arg.acc_element_op_,
@@ -1292,13 +1349,15 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V1
             const auto& device_arg = arg.group_device_args_[i];
             // Check if C permute dimension matches GEMM + GEMM shape
             const index_t c_g       = device_arg.c_grid_desc_g_m_n_.GetLength(I0); // unpadded
+            const index_t b_g       = device_arg.b_grid_desc_g_n_k_.GetLength(I0);
             const index_t c_m       = kernel_arg.y_grid_desc_m_o_.GetLength(I0);
             const index_t c_gemm1n  = kernel_arg.y_grid_desc_m_o_.GetLength(I1);
             const index_t a_m       = kernel_arg.a_grid_desc_ak0_m_ak1_.GetLength(I1);
             const index_t b1_gemm1n = kernel_arg.b1_grid_desc_bk0_n_bk1_.GetLength(I0) *
                                       kernel_arg.b1_grid_desc_bk0_n_bk1_.GetLength(I2);
 
-            if(!(c_g == device_arg.batch_count_ && c_m == a_m && c_gemm1n == b1_gemm1n))
+            if(!(c_g == device_arg.batch_count_ && c_m == a_m && c_gemm1n == b1_gemm1n &&
+                 c_g % b_g == 0 && c_g / b_g == arg.h_ratio_))
             {
                 return false;
             }
@@ -1335,6 +1394,17 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V1
                 return false;
             }
 
+            // saving dQ data with atomic_add instruction, so KzRaw must be a multiple of 2
+            if constexpr(is_same<OutputDataType, half_t>::value ||
+                         is_same<OutputDataType, bhalf_t>::value)
+            {
+                if(KzRaw % 2 != 0)
+                {
+                    std::cout << "K_q must be a multiple of 2" << std::endl;
+                    return false;
+                }
+            }
+
             // Check vector load/store requirement
             const auto a_stride_lowest = ABlockTransferSrcVectorDim == 2
                                              ? device_arg.a_mz_kz_strides_[1]

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

@@ -102,6 +102,7 @@ __global__ void
         kernel_grouped_multihead_attention_backward_qloop_xdl_cshuffle_light_v2(
             const void CK_CONSTANT_ADDRESS_SPACE* group_kernel_args,
             const index_t group_count,
+            const index_t h_ratio,
             const AElementwiseOperation a_element_op,
             const BElementwiseOperation b_element_op,
             const AccElementwiseOperation acc_element_op,
@@ -140,19 +141,26 @@ __global__ void
     const index_t num_blocks_per_batch = arg_ptr[group_id].num_blocks_per_batch_;
     const index_t g_idx                = __builtin_amdgcn_readfirstlane(
         (block_id - arg_ptr[group_id].block_start_) / (Deterministic ? 1 : num_blocks_per_batch));
+    const index_t gkv_idx = __builtin_amdgcn_readfirstlane(g_idx / h_ratio);
 
     const long_index_t a_batch_offset = __builtin_amdgcn_readfirstlane(
         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 b_batch_offset = __builtin_amdgcn_readfirstlane(static_cast<long_index_t>(
+        arg_ptr[group_id].compute_base_ptr_of_batch_.GetBBasePtr(gkv_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)));
+        arg_ptr[group_id].compute_base_ptr_of_batch_.GetB1BasePtr(gkv_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 long_index_t bgrad_batch_offset =
+        __builtin_amdgcn_readfirstlane(static_cast<long_index_t>(
+            arg_ptr[group_id].compute_base_ptr_of_batch_.GetBGradBasePtr(g_idx)));
+    const long_index_t b1grad_batch_offset =
+        __builtin_amdgcn_readfirstlane(static_cast<long_index_t>(
+            arg_ptr[group_id].compute_base_ptr_of_batch_.GetB1GradBasePtr(g_idx)));
 
     const index_t global_thread_id = get_thread_global_1d_id();
     ck::philox ph(seed, global_thread_id, offset);
@@ -166,7 +174,6 @@ __global__ void
         const long_index_t d0_batch_offset =
             __builtin_amdgcn_readfirstlane(static_cast<long_index_t>(
                 arg_ptr[group_id].compute_base_ptr_of_batch_.GetD0BasePtr(g_idx)));
-
         if(arg_ptr[group_id].p_d0_grid_ != nullptr)
             tmp_p_d0_grid = arg_ptr[group_id].p_d0_grid_ + d0_batch_offset;
         if(arg_ptr[group_id].p_d0grad_grid_)
@@ -187,9 +194,9 @@ __global__ void
                 arg_ptr[group_id].p_d_grid_ + lse_batch_offset,
                 arg_ptr[group_id].p_ygrad_grid_ + c_batch_offset,
                 arg_ptr[group_id].p_qgrad_grid_ + a_batch_offset,
-                arg_ptr[group_id].p_kgrad_grid_ + b_batch_offset,
+                arg_ptr[group_id].p_kgrad_grid_ + bgrad_batch_offset,
                 tmp_p_d0grad_grid,
-                arg_ptr[group_id].p_vgrad_grid_ + b1_batch_offset,
+                arg_ptr[group_id].p_vgrad_grid_ + b1grad_batch_offset,
                 p_shared,
                 a_element_op,
                 b_element_op,
@@ -198,9 +205,11 @@ __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].bgrad_grid_desc_bk0_n_bk1_,
                 arg_ptr[group_id].d0_grid_desc_m0_n0_m1_m2_n1_m3_,
                 arg_ptr[group_id].c_grid_desc_m0_n0_m1_n1_m2_n2_m3_m4_m5_n3_,
                 arg_ptr[group_id].b1_grid_desc_bk0_n_bk1_,
+                arg_ptr[group_id].b1grad_grid_desc_bk0_n_bk1_,
                 arg_ptr[group_id].lse_grid_desc_m_,
                 arg_ptr[group_id].ygrad_grid_desc_m0_o_m1_,
                 arg_ptr[group_id].block_2_ctile_map_,
@@ -225,9 +234,9 @@ __global__ void
             arg_ptr[group_id].p_d_grid_ + lse_batch_offset,
             arg_ptr[group_id].p_ygrad_grid_ + c_batch_offset,
             arg_ptr[group_id].p_qgrad_grid_ + a_batch_offset,
-            arg_ptr[group_id].p_kgrad_grid_ + b_batch_offset,
+            arg_ptr[group_id].p_kgrad_grid_ + bgrad_batch_offset,
             tmp_p_d0grad_grid,
-            arg_ptr[group_id].p_vgrad_grid_ + b1_batch_offset,
+            arg_ptr[group_id].p_vgrad_grid_ + b1grad_batch_offset,
             p_shared,
             a_element_op,
             b_element_op,
@@ -236,9 +245,11 @@ __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].bgrad_grid_desc_bk0_n_bk1_,
             arg_ptr[group_id].d0_grid_desc_m0_n0_m1_m2_n1_m3_,
             arg_ptr[group_id].c_grid_desc_m0_n0_m1_n1_m2_n2_m3_m4_m5_n3_,
             arg_ptr[group_id].b1_grid_desc_bk0_n_bk1_,
+            arg_ptr[group_id].b1grad_grid_desc_bk0_n_bk1_,
             arg_ptr[group_id].lse_grid_desc_m_,
             arg_ptr[group_id].ygrad_grid_desc_m0_o_m1_,
             arg_ptr[group_id].block_2_ctile_map_,
@@ -253,6 +264,7 @@ __global__ void
 #else
     ignore = group_kernel_args;
     ignore = group_count;
+    ignore = h_ratio;
     ignore = a_element_op;
     ignore = b_element_op;
     ignore = acc_element_op;
@@ -373,6 +385,12 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V2
         std::vector<index_t> lse_gs_ms_lengths;
         std::vector<index_t> lse_gs_ms_strides;
 
+        std::vector<index_t> bgrad_gs_ns_ks_lengths;
+        std::vector<index_t> bgrad_gs_ns_ks_strides;
+
+        std::vector<index_t> b1grad_gs_gemm1ns_gemm1ks_lengths;
+        std::vector<index_t> b1grad_gs_gemm1ns_gemm1ks_strides;
+
         std::vector<index_t> acc0_bias_gs_ms_ns_lengths;
         std::vector<index_t> acc0_bias_gs_ms_ns_strides;
 
@@ -639,7 +657,6 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V2
     static auto MakeD0GridDescriptor_M_N(const std::vector<ck::index_t>& acc0_bias_gs_ms_ns_lengths,
                                          const std::vector<ck::index_t>& acc0_bias_gs_ms_ns_strides)
     {
-
         return Transform::MakeC0GridDescriptor_M_N(acc0_bias_gs_ms_ns_lengths,
                                                    acc0_bias_gs_ms_ns_strides);
     }
@@ -648,7 +665,6 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V2
     MakeD0GridDescriptor_G_M_N(const std::vector<ck::index_t>& acc0_bias_gs_ms_ns_lengths,
                                const std::vector<ck::index_t>& acc0_bias_gs_ms_ns_strides)
     {
-
         return Transform::MakeC0GridDescriptor_G_M_N(acc0_bias_gs_ms_ns_lengths,
                                                      acc0_bias_gs_ms_ns_strides);
     }
@@ -723,6 +739,8 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V2
                                      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,
+                                     const BGridDesc_G_N_K& bgrad_grid_desc_g_n_k,
+                                     const B1GridDesc_G_N_K& b1grad_grid_desc_g_n_k,
                                      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),
@@ -730,6 +748,8 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V2
               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),
+              bgrad_grid_desc_g_n_k_(bgrad_grid_desc_g_n_k),
+              b1grad_grid_desc_g_n_k_(b1grad_grid_desc_g_n_k),
               BatchStrideLSE_(BatchStrideLSE)
         {
         }
@@ -769,6 +789,16 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V2
             return g_idx * static_cast<long_index_t>(BatchStrideLSE_);
         }
 
+        __host__ __device__ constexpr long_index_t GetBGradBasePtr(index_t g_idx) const
+        {
+            return bgrad_grid_desc_g_n_k_.CalculateOffset(make_multi_index(g_idx, 0, 0));
+        }
+
+        __host__ __device__ constexpr long_index_t GetB1GradBasePtr(index_t g_idx) const
+        {
+            return b1grad_grid_desc_g_n_k_.CalculateOffset(make_multi_index(g_idx, 0, 0));
+        }
+
         private:
         AGridDesc_G_M_K a_grid_desc_g_m_k_;
         BGridDesc_G_N_K b_grid_desc_g_n_k_;
@@ -776,6 +806,8 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V2
         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_;
+        BGridDesc_G_N_K bgrad_grid_desc_g_n_k_;
+        B1GridDesc_G_N_K b1grad_grid_desc_g_n_k_;
         index_t BatchStrideLSE_;
     };
 
@@ -888,9 +920,11 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_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_;
+        BGridDesc_BK0_N_BK1 bgrad_grid_desc_bk0_n_bk1_;
         typename GridwiseGemm::D0GridDescriptor_M0_N0_M1_M2_N1_M3 d0_grid_desc_m0_n0_m1_m2_n1_m3_;
         ZGridDesc_M_N z_grid_desc_m_n_;
         B1GridDesc_BK0_N_BK1 b1_grid_desc_bk0_n_bk1_;
+        B1GridDesc_BK0_N_BK1 b1grad_grid_desc_bk0_n_bk1_;
         YGridDesc_M_O y_grid_desc_m_o_;
 
         typename GridwiseGemm::ZGridDescriptor_M0_N0_M1_N1_M2_N2_M3_M4_M5_N3
@@ -932,6 +966,7 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V2
         std::vector<index_t> c_mz_gemm1nz_strides_;
 
         // for gridwise gemm check
+        BGridDesc_G_N_K b_grid_desc_g_n_k_;
         CGridDesc_G_M_N c_grid_desc_g_m_n_;
 
         index_t batch_count_;
@@ -1004,6 +1039,9 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V2
 
             d_grid_size_ = 0;
 
+            h_ratio_ = problem_desc_vec[0].a_gs_ms_ks_lengths[NumDimG - 1] /
+                       problem_desc_vec[0].b_gs_ns_ks_lengths[NumDimG - 1];
+
             for(index_t i = 0; i < group_count_; i++)
             {
                 const auto p_a_grid = static_cast<const InputDataType*>(p_As[i]);
@@ -1031,6 +1069,8 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_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 bgrad_grid_desc_bk0_n_bk1 = DeviceOp::MakeBGridDescriptor_BK0_N_BK1(
+                    problem_desc.bgrad_gs_ns_ks_lengths, problem_desc.bgrad_gs_ns_ks_strides);
 
                 std::vector<index_t> tmp_d0_gs_ms_ns_lengths;
                 std::vector<index_t> tmp_d0_gs_ms_ns_strides;
@@ -1053,6 +1093,9 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V2
                 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);
+                const auto b1grad_grid_desc_bk0_n_bk1 = DeviceOp::MakeVGridDescriptor_O0_N_O1(
+                    problem_desc.b1grad_gs_gemm1ns_gemm1ks_lengths,
+                    problem_desc.b1grad_gs_gemm1ns_gemm1ks_strides);
                 const auto y_grid_desc_m_o = Transform::MakeCGridDescriptor_M_N(
                     problem_desc.c_gs_ms_gemm1ns_lengths, problem_desc.c_gs_ms_gemm1ns_strides);
 
@@ -1076,6 +1119,11 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V2
                     problem_desc.b1_gs_gemm1ns_gemm1ks_strides);
                 const auto c_grid_desc_g_m_n = Transform::MakeCGridDescriptor_G_M_N(
                     problem_desc.c_gs_ms_gemm1ns_lengths, problem_desc.c_gs_ms_gemm1ns_strides);
+                const auto bgrad_grid_desc_g_n_k = Transform::MakeB0GridDescriptor_G_N_K(
+                    problem_desc.bgrad_gs_ns_ks_lengths, problem_desc.bgrad_gs_ns_ks_strides);
+                const auto b1grad_grid_desc_g_n_k = Transform::MakeB1GridDescriptor_G_N_K(
+                    problem_desc.b1grad_gs_gemm1ns_gemm1ks_lengths,
+                    problem_desc.b1grad_gs_gemm1ns_gemm1ks_strides);
                 typename GridwiseGemm::ZGridDescriptor_M0_N0_M1_N1_M2_N2_M3_M4_M5_N3
                     c_grid_desc_m0_n0_m1_n1_m2_n2_m3_m4_m5_n3;
                 const index_t BlockStart     = grid_size_;
@@ -1098,7 +1146,9 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V2
                     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()));
+                    bgrad_grid_desc_g_n_k,
+                    b1grad_grid_desc_g_n_k,
+                    type_convert<index_t>(problem_desc.lse_gs_ms_strides[NumDimG - 1]));
 
                 // C0 mask
                 const auto c0_matrix_mask =
@@ -1144,9 +1194,11 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V2
                                               p_vgrad_grid,
                                               a_grid_desc_ak0_m_ak1,
                                               b_grid_desc_bk0_n_bk1,
+                                              bgrad_grid_desc_bk0_n_bk1,
                                               d0_grid_desc_m0_n0_m1_m2_n1_m3,
                                               z_grid_desc_m_n,
                                               b1_grid_desc_bk0_n_bk1,
+                                              b1grad_grid_desc_bk0_n_bk1,
                                               y_grid_desc_m_o,
                                               c_grid_desc_m0_n0_m1_n1_m2_n2_m3_m4_m5_n3,
                                               lse_grid_desc_m,
@@ -1190,6 +1242,7 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V2
                       problem_desc.b1_gs_gemm1ns_gemm1ks_strides[NumDimG + NumDimO + NumDimN - 1]},
                      {problem_desc.c_gs_ms_gemm1ns_strides[NumDimG + NumDimM - 1],
                       problem_desc.c_gs_ms_gemm1ns_strides[NumDimG + NumDimM + NumDimO - 1]},
+                     b_grid_desc_g_n_k,
                      c_grid_desc_g_m_n,
                      batch_count,
                      d0_n_length_stride});
@@ -1216,6 +1269,7 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V2
 
         index_t grid_size_;
         index_t group_count_;
+        index_t h_ratio_;
 
         std::vector<GroupKernelArg> group_kernel_args_;
         std::vector<GroupDeviceArg> group_device_args_;
@@ -1294,6 +1348,7 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V2
                     0,
                     cast_pointer_to_constant_address_space(arg.p_workspace_),
                     arg.group_count_,
+                    arg.h_ratio_,
                     arg.a_element_op_,
                     arg.b_element_op_,
                     arg.acc_element_op_,
@@ -1362,13 +1417,15 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V2
             const auto& device_arg = arg.group_device_args_[i];
             // Check if C permute dimension matches GEMM + GEMM shape
             const index_t c_g       = device_arg.c_grid_desc_g_m_n_.GetLength(I0); // unpadded
+            const index_t b_g       = device_arg.b_grid_desc_g_n_k_.GetLength(I0);
             const index_t c_m       = kernel_arg.y_grid_desc_m_o_.GetLength(I0);
             const index_t c_gemm1n  = kernel_arg.y_grid_desc_m_o_.GetLength(I1);
             const index_t a_m       = kernel_arg.a_grid_desc_ak0_m_ak1_.GetLength(I1);
             const index_t b1_gemm1n = kernel_arg.b1_grid_desc_bk0_n_bk1_.GetLength(I0) *
                                       kernel_arg.b1_grid_desc_bk0_n_bk1_.GetLength(I2);
 
-            if(!(c_g == device_arg.batch_count_ && c_m == a_m && c_gemm1n == b1_gemm1n))
+            if(!(c_g == device_arg.batch_count_ && c_m == a_m && c_gemm1n == b1_gemm1n &&
+                 c_g % b_g == 0 && c_g / b_g == arg.h_ratio_))
             {
                 return false;
             }
@@ -1407,6 +1464,17 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V2
                 return false;
             }
 
+            // saving dQ data with atomic_add instruction, so KzRaw must be a multiple of 2
+            if constexpr(is_same<OutputDataType, half_t>::value ||
+                         is_same<OutputDataType, bhalf_t>::value)
+            {
+                if(KzRaw % 2 != 0)
+                {
+                    std::cout << "K_q must be a multiple of 2" << std::endl;
+                    return false;
+                }
+            }
+
             // Check vector load/store requirement
             const auto a_stride_lowest = ABlockTransferSrcVectorDim == 2
                                              ? device_arg.a_mz_kz_strides_[1]

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

@@ -44,6 +44,7 @@ __global__ void
         kernel_grouped_multihead_attention_backward_qloop_xdl_cshuffle_v1(
             const void CK_CONSTANT_ADDRESS_SPACE* group_kernel_args,
             const index_t group_count,
+            const index_t h_ratio,
             const AElementwiseOperation a_element_op,
             const BElementwiseOperation b_element_op,
             const AccElementwiseOperation acc_element_op,
@@ -82,19 +83,26 @@ __global__ void
     const index_t num_blocks_per_batch = arg_ptr[group_id].num_blocks_per_batch_;
     const index_t g_idx                = __builtin_amdgcn_readfirstlane(
         (block_id - arg_ptr[group_id].block_start_) / (Deterministic ? 1 : num_blocks_per_batch));
+    const index_t gkv_idx = __builtin_amdgcn_readfirstlane(g_idx / h_ratio);
 
     const long_index_t a_batch_offset = __builtin_amdgcn_readfirstlane(
         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 b_batch_offset = __builtin_amdgcn_readfirstlane(static_cast<long_index_t>(
+        arg_ptr[group_id].compute_base_ptr_of_batch_.GetBBasePtr(gkv_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)));
+        arg_ptr[group_id].compute_base_ptr_of_batch_.GetB1BasePtr(gkv_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 long_index_t bgrad_batch_offset =
+        __builtin_amdgcn_readfirstlane(static_cast<long_index_t>(
+            arg_ptr[group_id].compute_base_ptr_of_batch_.GetBGradBasePtr(g_idx)));
+    const long_index_t b1grad_batch_offset =
+        __builtin_amdgcn_readfirstlane(static_cast<long_index_t>(
+            arg_ptr[group_id].compute_base_ptr_of_batch_.GetB1GradBasePtr(g_idx)));
 
     const index_t global_thread_id = get_thread_global_1d_id();
     ck::philox ph(seed, global_thread_id, offset);
@@ -129,9 +137,9 @@ __global__ void
                 arg_ptr[group_id].p_lse_grid_ + lse_batch_offset,
                 arg_ptr[group_id].p_ygrad_grid_ + c_batch_offset,
                 arg_ptr[group_id].p_qgrad_grid_ + a_batch_offset,
-                arg_ptr[group_id].p_kgrad_grid_ + b_batch_offset,
+                arg_ptr[group_id].p_kgrad_grid_ + bgrad_batch_offset,
                 tmp_p_d0grad_grid,
-                arg_ptr[group_id].p_vgrad_grid_ + b1_batch_offset,
+                arg_ptr[group_id].p_vgrad_grid_ + b1grad_batch_offset,
                 p_shared,
                 a_element_op,
                 b_element_op,
@@ -140,9 +148,11 @@ __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].bgrad_grid_desc_bk0_n_bk1_,
                 arg_ptr[group_id].d0_grid_desc_m0_n0_m1_m2_n1_m3_,
                 arg_ptr[group_id].c_grid_desc_m0_n0_m1_n1_m2_n2_m3_m4_m5_n3_,
                 arg_ptr[group_id].b1_grid_desc_bk0_n_bk1_,
+                arg_ptr[group_id].b1grad_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].ygrad_grid_desc_o0_m_o1_,
@@ -168,9 +178,9 @@ __global__ void
             arg_ptr[group_id].p_lse_grid_ + lse_batch_offset,
             arg_ptr[group_id].p_ygrad_grid_ + c_batch_offset,
             arg_ptr[group_id].p_qgrad_grid_ + a_batch_offset,
-            arg_ptr[group_id].p_kgrad_grid_ + b_batch_offset,
+            arg_ptr[group_id].p_kgrad_grid_ + bgrad_batch_offset,
             tmp_p_d0grad_grid,
-            arg_ptr[group_id].p_vgrad_grid_ + b1_batch_offset,
+            arg_ptr[group_id].p_vgrad_grid_ + b1grad_batch_offset,
             p_shared,
             a_element_op,
             b_element_op,
@@ -179,9 +189,11 @@ __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].bgrad_grid_desc_bk0_n_bk1_,
             arg_ptr[group_id].d0_grid_desc_m0_n0_m1_m2_n1_m3_,
             arg_ptr[group_id].c_grid_desc_m0_n0_m1_n1_m2_n2_m3_m4_m5_n3_,
             arg_ptr[group_id].b1_grid_desc_bk0_n_bk1_,
+            arg_ptr[group_id].b1grad_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].ygrad_grid_desc_o0_m_o1_,
@@ -197,6 +209,7 @@ __global__ void
 #else
     ignore = group_kernel_args;
     ignore = group_count;
+    ignore = h_ratio;
     ignore = a_element_op;
     ignore = b_element_op;
     ignore = acc_element_op;
@@ -307,6 +320,12 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V1
         std::vector<index_t> lse_gs_ms_lengths;
         std::vector<index_t> lse_gs_ms_strides;
 
+        std::vector<index_t> bgrad_gs_ns_ks_lengths;
+        std::vector<index_t> bgrad_gs_ns_ks_strides;
+
+        std::vector<index_t> b1grad_gs_gemm1ns_gemm1ks_lengths;
+        std::vector<index_t> b1grad_gs_gemm1ns_gemm1ks_strides;
+
         std::vector<index_t> acc0_bias_gs_ms_ns_lengths;
         std::vector<index_t> acc0_bias_gs_ms_ns_strides;
 
@@ -508,7 +527,6 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V1
     static auto MakeD0GridDescriptor_M_N(const std::vector<ck::index_t>& acc0_bias_gs_ms_ns_lengths,
                                          const std::vector<ck::index_t>& acc0_bias_gs_ms_ns_strides)
     {
-
         return Transform::MakeC0GridDescriptor_M_N(acc0_bias_gs_ms_ns_lengths,
                                                    acc0_bias_gs_ms_ns_strides);
     }
@@ -517,7 +535,6 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V1
     MakeD0GridDescriptor_G_M_N(const std::vector<ck::index_t>& acc0_bias_gs_ms_ns_lengths,
                                const std::vector<ck::index_t>& acc0_bias_gs_ms_ns_strides)
     {
-
         return Transform::MakeC0GridDescriptor_G_M_N(acc0_bias_gs_ms_ns_lengths,
                                                      acc0_bias_gs_ms_ns_strides);
     }
@@ -564,6 +581,8 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V1
                                      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,
+                                     const BGridDesc_G_N_K& bgrad_grid_desc_g_n_k,
+                                     const B1GridDesc_G_N_K& b1grad_grid_desc_g_n_k,
                                      index_t batch_stride_lse)
             : a_grid_desc_g_m_k_(a_grid_desc_g_m_k),
               b_grid_desc_g_n_k_(b_grid_desc_g_n_k),
@@ -571,6 +590,8 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V1
               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),
+              bgrad_grid_desc_g_n_k_(bgrad_grid_desc_g_n_k),
+              b1grad_grid_desc_g_n_k_(b1grad_grid_desc_g_n_k),
               batch_stride_lse_(batch_stride_lse)
         {
         }
@@ -610,6 +631,16 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V1
             return g_idx * static_cast<long_index_t>(batch_stride_lse_);
         }
 
+        __host__ __device__ constexpr long_index_t GetBGradBasePtr(index_t g_idx) const
+        {
+            return bgrad_grid_desc_g_n_k_.CalculateOffset(make_multi_index(g_idx, 0, 0));
+        }
+
+        __host__ __device__ constexpr long_index_t GetB1GradBasePtr(index_t g_idx) const
+        {
+            return b1grad_grid_desc_g_n_k_.CalculateOffset(make_multi_index(g_idx, 0, 0));
+        }
+
         private:
         AGridDesc_G_M_K a_grid_desc_g_m_k_;
         BGridDesc_G_N_K b_grid_desc_g_n_k_;
@@ -617,6 +648,8 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V1
         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_;
+        BGridDesc_G_N_K bgrad_grid_desc_g_n_k_;
+        B1GridDesc_G_N_K b1grad_grid_desc_g_n_k_;
         index_t batch_stride_lse_;
     };
 
@@ -708,9 +741,11 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V1
         // 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_;
+        BGridDesc_BK0_N_BK1 bgrad_grid_desc_bk0_n_bk1_;
         typename GridwiseGemm::D0GridDescriptor_M0_N0_M1_M2_N1_M3 d0_grid_desc_m0_n0_m1_m2_n1_m3_;
         ZGridDesc_M_N z_grid_desc_m_n_;
         B1GridDesc_BK0_N_BK1 b1_grid_desc_bk0_n_bk1_;
+        B1GridDesc_BK0_N_BK1 b1grad_grid_desc_bk0_n_bk1_;
         YGridDesc_M_O y_grid_desc_m_o_;
 
         typename GridwiseGemm::YGridDescriptor_MBlock_MPerBlock_OBlock_OPerBlock
@@ -745,6 +780,7 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V1
         std::vector<index_t> c_mz_gemm1nz_strides_;
 
         // for gridwise gemm check
+        BGridDesc_G_N_K b_grid_desc_g_n_k_;
         CGridDesc_G_M_N c_grid_desc_g_m_n_;
 
         index_t batch_count_;
@@ -813,6 +849,9 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V1
 
             index_t z_random_matrix_offset = 0;
 
+            h_ratio_ = problem_desc_vec[0].a_gs_ms_ks_lengths[NumDimG - 1] /
+                       problem_desc_vec[0].b_gs_ns_ks_lengths[NumDimG - 1];
+
             for(index_t i = 0; i < group_count_; i++)
             {
                 const auto p_a_grid = static_cast<const InputDataType*>(p_As[i]);
@@ -840,6 +879,8 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V1
                     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 bgrad_grid_desc_bk0_n_bk1 = DeviceOp::MakeBGridDescriptor_BK0_N_BK1(
+                    problem_desc.bgrad_gs_ns_ks_lengths, problem_desc.bgrad_gs_ns_ks_strides);
 
                 std::vector<index_t> tmp_d0_gs_ms_ns_lengths;
                 std::vector<index_t> tmp_d0_gs_ms_ns_strides;
@@ -862,6 +903,9 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V1
                 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);
+                const auto b1grad_grid_desc_bk0_n_bk1 = DeviceOp::MakeVGridDescriptor_O0_N_O1(
+                    problem_desc.b1grad_gs_gemm1ns_gemm1ks_lengths,
+                    problem_desc.b1grad_gs_gemm1ns_gemm1ks_strides);
                 const auto y_grid_desc_m_o = Transform::MakeCGridDescriptor_M_N(
                     problem_desc.c_gs_ms_gemm1ns_lengths, problem_desc.c_gs_ms_gemm1ns_strides);
 
@@ -885,6 +929,11 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V1
                     problem_desc.b1_gs_gemm1ns_gemm1ks_strides);
                 const auto c_grid_desc_g_m_n = Transform::MakeCGridDescriptor_G_M_N(
                     problem_desc.c_gs_ms_gemm1ns_lengths, problem_desc.c_gs_ms_gemm1ns_strides);
+                const auto bgrad_grid_desc_g_n_k = Transform::MakeB0GridDescriptor_G_N_K(
+                    problem_desc.bgrad_gs_ns_ks_lengths, problem_desc.bgrad_gs_ns_ks_strides);
+                const auto b1grad_grid_desc_g_n_k = Transform::MakeB1GridDescriptor_G_N_K(
+                    problem_desc.b1grad_gs_gemm1ns_gemm1ks_lengths,
+                    problem_desc.b1grad_gs_gemm1ns_gemm1ks_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_M4_M5_N3
@@ -918,7 +967,9 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V1
                     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()));
+                    bgrad_grid_desc_g_n_k,
+                    b1grad_grid_desc_g_n_k,
+                    type_convert<index_t>(problem_desc.lse_gs_ms_strides[NumDimG - 1]));
 
                 // C0 mask
                 const auto c0_matrix_mask =
@@ -945,9 +996,11 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V1
                                               p_vgrad_grid,
                                               a_grid_desc_ak0_m_ak1,
                                               b_grid_desc_bk0_n_bk1,
+                                              bgrad_grid_desc_bk0_n_bk1,
                                               d0_grid_desc_m0_n0_m1_m2_n1_m3,
                                               z_grid_desc_m_n,
                                               b1_grid_desc_bk0_n_bk1,
+                                              b1grad_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_m4_m5_n3,
@@ -985,6 +1038,7 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V1
                       problem_desc.b1_gs_gemm1ns_gemm1ks_strides[NumDimG + NumDimO + NumDimN - 1]},
                      {problem_desc.c_gs_ms_gemm1ns_strides[NumDimG + NumDimM - 1],
                       problem_desc.c_gs_ms_gemm1ns_strides[NumDimG + NumDimM + NumDimO - 1]},
+                     b_grid_desc_g_n_k,
                      c_grid_desc_g_m_n,
                      batch_count,
                      d0_n_length_stride});
@@ -1011,6 +1065,7 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V1
 
         index_t grid_size_;
         index_t group_count_;
+        index_t h_ratio_;
 
         std::vector<GroupKernelArg> group_kernel_args_;
         std::vector<GroupDeviceArg> group_device_args_;
@@ -1070,6 +1125,7 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V1
                     0,
                     cast_pointer_to_constant_address_space(arg.p_workspace_),
                     arg.group_count_,
+                    arg.h_ratio_,
                     arg.a_element_op_,
                     arg.b_element_op_,
                     arg.acc_element_op_,
@@ -1138,13 +1194,15 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V1
             const auto& device_arg = arg.group_device_args_[i];
             // Check if C permute dimension matches GEMM + GEMM shape
             const index_t c_g       = device_arg.c_grid_desc_g_m_n_.GetLength(I0); // unpadded
+            const index_t b_g       = device_arg.b_grid_desc_g_n_k_.GetLength(I0);
             const index_t c_m       = kernel_arg.y_grid_desc_m_o_.GetLength(I0);
             const index_t c_gemm1n  = kernel_arg.y_grid_desc_m_o_.GetLength(I1);
             const index_t a_m       = kernel_arg.a_grid_desc_ak0_m_ak1_.GetLength(I1);
             const index_t b1_gemm1n = kernel_arg.b1_grid_desc_bk0_n_bk1_.GetLength(I0) *
                                       kernel_arg.b1_grid_desc_bk0_n_bk1_.GetLength(I2);
 
-            if(!(c_g == device_arg.batch_count_ && c_m == a_m && c_gemm1n == b1_gemm1n))
+            if(!(c_g == device_arg.batch_count_ && c_m == a_m && c_gemm1n == b1_gemm1n &&
+                 c_g % b_g == 0 && c_g / b_g == arg.h_ratio_))
             {
                 return false;
             }
@@ -1181,6 +1239,17 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V1
                 return false;
             }
 
+            // saving dQ data with atomic_add instruction, so KzRaw must be a multiple of 2
+            if constexpr(is_same<OutputDataType, half_t>::value ||
+                         is_same<OutputDataType, bhalf_t>::value)
+            {
+                if(KzRaw % 2 != 0)
+                {
+                    std::cout << "K_q must be a multiple of 2" << std::endl;
+                    return false;
+                }
+            }
+
             // Check vector load/store requirement
             const auto a_stride_lowest = ABlockTransferSrcVectorDim == 2
                                              ? device_arg.a_mz_kz_strides_[1]

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

@@ -44,6 +44,7 @@ __global__ void
         kernel_grouped_multihead_attention_backward_qloop_xdl_cshuffle_v2(
             const void CK_CONSTANT_ADDRESS_SPACE* group_kernel_args,
             const index_t group_count,
+            const index_t h_ratio,
             const AElementwiseOperation a_element_op,
             const BElementwiseOperation b_element_op,
             const AccElementwiseOperation acc_element_op,
@@ -82,19 +83,26 @@ __global__ void
     const index_t num_blocks_per_batch = arg_ptr[group_id].num_blocks_per_batch_;
     const index_t g_idx                = __builtin_amdgcn_readfirstlane(
         (block_id - arg_ptr[group_id].block_start_) / (Deterministic ? 1 : num_blocks_per_batch));
+    const index_t gkv_idx = __builtin_amdgcn_readfirstlane(g_idx / h_ratio);
 
     const long_index_t a_batch_offset = __builtin_amdgcn_readfirstlane(
         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 b_batch_offset = __builtin_amdgcn_readfirstlane(static_cast<long_index_t>(
+        arg_ptr[group_id].compute_base_ptr_of_batch_.GetBBasePtr(gkv_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)));
+        arg_ptr[group_id].compute_base_ptr_of_batch_.GetB1BasePtr(gkv_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 long_index_t bgrad_batch_offset =
+        __builtin_amdgcn_readfirstlane(static_cast<long_index_t>(
+            arg_ptr[group_id].compute_base_ptr_of_batch_.GetBGradBasePtr(g_idx)));
+    const long_index_t b1grad_batch_offset =
+        __builtin_amdgcn_readfirstlane(static_cast<long_index_t>(
+            arg_ptr[group_id].compute_base_ptr_of_batch_.GetB1GradBasePtr(g_idx)));
 
     const index_t global_thread_id = get_thread_global_1d_id();
     ck::philox ph(seed, global_thread_id, offset);
@@ -128,9 +136,9 @@ __global__ void
                 arg_ptr[group_id].p_lse_grid_ + lse_batch_offset,
                 arg_ptr[group_id].p_ygrad_grid_ + c_batch_offset,
                 arg_ptr[group_id].p_qgrad_grid_ + a_batch_offset,
-                arg_ptr[group_id].p_kgrad_grid_ + b_batch_offset,
+                arg_ptr[group_id].p_kgrad_grid_ + bgrad_batch_offset,
                 tmp_p_d0grad_grid,
-                arg_ptr[group_id].p_vgrad_grid_ + b1_batch_offset,
+                arg_ptr[group_id].p_vgrad_grid_ + b1grad_batch_offset,
                 p_shared,
                 a_element_op,
                 b_element_op,
@@ -139,9 +147,11 @@ __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].bgrad_grid_desc_bk0_n_bk1_,
                 arg_ptr[group_id].d0_grid_desc_m0_n0_m1_m2_n1_m3_,
                 arg_ptr[group_id].c_grid_desc_m0_n0_m1_n1_m2_n2_m3_m4_m5_n3_,
                 arg_ptr[group_id].b1_grid_desc_bk0_n_bk1_,
+                arg_ptr[group_id].b1grad_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].ygrad_grid_desc_m0_o_m1_,
@@ -167,9 +177,9 @@ __global__ void
             arg_ptr[group_id].p_lse_grid_ + lse_batch_offset,
             arg_ptr[group_id].p_ygrad_grid_ + c_batch_offset,
             arg_ptr[group_id].p_qgrad_grid_ + a_batch_offset,
-            arg_ptr[group_id].p_kgrad_grid_ + b_batch_offset,
+            arg_ptr[group_id].p_kgrad_grid_ + bgrad_batch_offset,
             tmp_p_d0grad_grid,
-            arg_ptr[group_id].p_vgrad_grid_ + b1_batch_offset,
+            arg_ptr[group_id].p_vgrad_grid_ + b1grad_batch_offset,
             p_shared,
             a_element_op,
             b_element_op,
@@ -178,9 +188,11 @@ __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].bgrad_grid_desc_bk0_n_bk1_,
             arg_ptr[group_id].d0_grid_desc_m0_n0_m1_m2_n1_m3_,
             arg_ptr[group_id].c_grid_desc_m0_n0_m1_n1_m2_n2_m3_m4_m5_n3_,
             arg_ptr[group_id].b1_grid_desc_bk0_n_bk1_,
+            arg_ptr[group_id].b1grad_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].ygrad_grid_desc_m0_o_m1_,
@@ -196,6 +208,7 @@ __global__ void
 #else
     ignore = group_kernel_args;
     ignore = group_count;
+    ignore = h_ratio;
     ignore = a_element_op;
     ignore = b_element_op;
     ignore = acc_element_op;
@@ -313,6 +326,12 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
         std::vector<index_t> lse_gs_ms_lengths;
         std::vector<index_t> lse_gs_ms_strides;
 
+        std::vector<index_t> bgrad_gs_ns_ks_lengths;
+        std::vector<index_t> bgrad_gs_ns_ks_strides;
+
+        std::vector<index_t> b1grad_gs_gemm1ns_gemm1ks_lengths;
+        std::vector<index_t> b1grad_gs_gemm1ns_gemm1ks_strides;
+
         std::vector<index_t> acc0_bias_gs_ms_ns_lengths;
         std::vector<index_t> acc0_bias_gs_ms_ns_strides;
 
@@ -448,19 +467,6 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_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,
-                                                const std::vector<index_t>& y_gs_ms_os_strides)
-    {
-        return Transform::MakeAGridDescriptor_AK0_M_AK1(
-            Transform::MakeAGridDescriptor_M_K(y_gs_ms_os_lengths, y_gs_ms_os_strides),
-            Number<Y_O1>{});
-    }
-
     // V in Gemm B position
     static auto MakeVGridDescriptor_O0_N_O1(const std::vector<index_t>& v_gs_os_ns_lengths,
                                             const std::vector<index_t>& v_gs_os_ns_strides)
@@ -570,7 +576,6 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
     static auto MakeD0GridDescriptor_M_N(const std::vector<ck::index_t>& acc0_bias_gs_ms_ns_lengths,
                                          const std::vector<ck::index_t>& acc0_bias_gs_ms_ns_strides)
     {
-
         return Transform::MakeC0GridDescriptor_M_N(acc0_bias_gs_ms_ns_lengths,
                                                    acc0_bias_gs_ms_ns_strides);
     }
@@ -579,7 +584,6 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
     MakeD0GridDescriptor_G_M_N(const std::vector<ck::index_t>& acc0_bias_gs_ms_ns_lengths,
                                const std::vector<ck::index_t>& acc0_bias_gs_ms_ns_strides)
     {
-
         return Transform::MakeC0GridDescriptor_G_M_N(acc0_bias_gs_ms_ns_lengths,
                                                      acc0_bias_gs_ms_ns_strides);
     }
@@ -626,6 +630,8 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
                                      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,
+                                     const BGridDesc_G_N_K& bgrad_grid_desc_g_n_k,
+                                     const B1GridDesc_G_N_K& b1grad_grid_desc_g_n_k,
                                      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),
@@ -633,6 +639,8 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
               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),
+              bgrad_grid_desc_g_n_k_(bgrad_grid_desc_g_n_k),
+              b1grad_grid_desc_g_n_k_(b1grad_grid_desc_g_n_k),
               BatchStrideLSE_(BatchStrideLSE)
         {
         }
@@ -672,6 +680,16 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
             return g_idx * static_cast<long_index_t>(BatchStrideLSE_);
         }
 
+        __host__ __device__ constexpr long_index_t GetBGradBasePtr(index_t g_idx) const
+        {
+            return bgrad_grid_desc_g_n_k_.CalculateOffset(make_multi_index(g_idx, 0, 0));
+        }
+
+        __host__ __device__ constexpr long_index_t GetB1GradBasePtr(index_t g_idx) const
+        {
+            return b1grad_grid_desc_g_n_k_.CalculateOffset(make_multi_index(g_idx, 0, 0));
+        }
+
         private:
         AGridDesc_G_M_K a_grid_desc_g_m_k_;
         BGridDesc_G_N_K b_grid_desc_g_n_k_;
@@ -679,6 +697,8 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
         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_;
+        BGridDesc_G_N_K bgrad_grid_desc_g_n_k_;
+        B1GridDesc_G_N_K b1grad_grid_desc_g_n_k_;
         index_t BatchStrideLSE_;
     };
 
@@ -778,9 +798,11 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_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_;
+        BGridDesc_BK0_N_BK1 bgrad_grid_desc_bk0_n_bk1_;
         typename GridwiseGemm::D0GridDescriptor_M0_N0_M1_M2_N1_M3 d0_grid_desc_m0_n0_m1_m2_n1_m3_;
         ZGridDesc_M_N z_grid_desc_m_n_;
         B1GridDesc_BK0_N_BK1 b1_grid_desc_bk0_n_bk1_;
+        B1GridDesc_BK0_N_BK1 b1grad_grid_desc_bk0_n_bk1_;
         YGridDesc_M_O y_grid_desc_m_o_;
 
         typename GridwiseGemm::YGridDescriptor_MBlock_MPerBlock_OBlock_OPerBlock
@@ -815,6 +837,7 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
         std::vector<index_t> c_mz_gemm1nz_strides_;
 
         // for gridwise gemm check
+        BGridDesc_G_N_K b_grid_desc_g_n_k_;
         CGridDesc_G_M_N c_grid_desc_g_m_n_;
 
         index_t batch_count_;
@@ -883,6 +906,9 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
 
             index_t z_random_matrix_offset = 0;
 
+            h_ratio_ = problem_desc_vec[0].a_gs_ms_ks_lengths[NumDimG - 1] /
+                       problem_desc_vec[0].b_gs_ns_ks_lengths[NumDimG - 1];
+
             for(index_t i = 0; i < group_count_; i++)
             {
                 const auto p_a_grid = static_cast<const InputDataType*>(p_As[i]);
@@ -910,6 +936,8 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_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 bgrad_grid_desc_bk0_n_bk1 = DeviceOp::MakeBGridDescriptor_BK0_N_BK1(
+                    problem_desc.bgrad_gs_ns_ks_lengths, problem_desc.bgrad_gs_ns_ks_strides);
 
                 std::vector<index_t> tmp_d0_gs_ms_ns_lengths;
                 std::vector<index_t> tmp_d0_gs_ms_ns_strides;
@@ -932,6 +960,9 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
                 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);
+                const auto b1grad_grid_desc_bk0_n_bk1 = DeviceOp::MakeVGridDescriptor_O0_N_O1(
+                    problem_desc.b1grad_gs_gemm1ns_gemm1ks_lengths,
+                    problem_desc.b1grad_gs_gemm1ns_gemm1ks_strides);
                 const auto y_grid_desc_m_o = Transform::MakeCGridDescriptor_M_N(
                     problem_desc.c_gs_ms_gemm1ns_lengths, problem_desc.c_gs_ms_gemm1ns_strides);
 
@@ -955,6 +986,11 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
                     problem_desc.b1_gs_gemm1ns_gemm1ks_strides);
                 const auto c_grid_desc_g_m_n = Transform::MakeCGridDescriptor_G_M_N(
                     problem_desc.c_gs_ms_gemm1ns_lengths, problem_desc.c_gs_ms_gemm1ns_strides);
+                const auto bgrad_grid_desc_g_n_k = Transform::MakeB0GridDescriptor_G_N_K(
+                    problem_desc.bgrad_gs_ns_ks_lengths, problem_desc.bgrad_gs_ns_ks_strides);
+                const auto b1grad_grid_desc_g_n_k = Transform::MakeB1GridDescriptor_G_N_K(
+                    problem_desc.b1grad_gs_gemm1ns_gemm1ks_lengths,
+                    problem_desc.b1grad_gs_gemm1ns_gemm1ks_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_M4_M5_N3
@@ -988,7 +1024,9 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
                     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()));
+                    bgrad_grid_desc_g_n_k,
+                    b1grad_grid_desc_g_n_k,
+                    type_convert<index_t>(problem_desc.lse_gs_ms_strides[NumDimG - 1]));
 
                 // C0 mask
                 const auto c0_matrix_mask =
@@ -1015,9 +1053,11 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
                                               p_vgrad_grid,
                                               a_grid_desc_ak0_m_ak1,
                                               b_grid_desc_bk0_n_bk1,
+                                              bgrad_grid_desc_bk0_n_bk1,
                                               d0_grid_desc_m0_n0_m1_m2_n1_m3,
                                               z_grid_desc_m_n,
                                               b1_grid_desc_bk0_n_bk1,
+                                              b1grad_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_m4_m5_n3,
@@ -1055,6 +1095,7 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
                       problem_desc.b1_gs_gemm1ns_gemm1ks_strides[NumDimG + NumDimO + NumDimN - 1]},
                      {problem_desc.c_gs_ms_gemm1ns_strides[NumDimG + NumDimM - 1],
                       problem_desc.c_gs_ms_gemm1ns_strides[NumDimG + NumDimM + NumDimO - 1]},
+                     b_grid_desc_g_n_k,
                      c_grid_desc_g_m_n,
                      batch_count,
                      d0_n_length_stride});
@@ -1081,6 +1122,7 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
 
         index_t grid_size_;
         index_t group_count_;
+        index_t h_ratio_;
 
         std::vector<GroupKernelArg> group_kernel_args_;
         std::vector<GroupDeviceArg> group_device_args_;
@@ -1139,6 +1181,7 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
                     0,
                     cast_pointer_to_constant_address_space(arg.p_workspace_),
                     arg.group_count_,
+                    arg.h_ratio_,
                     arg.a_element_op_,
                     arg.b_element_op_,
                     arg.acc_element_op_,
@@ -1207,13 +1250,15 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
             const auto& device_arg = arg.group_device_args_[i];
             // Check if C permute dimension matches GEMM + GEMM shape
             const index_t c_g       = device_arg.c_grid_desc_g_m_n_.GetLength(I0); // unpadded
+            const index_t b_g       = device_arg.b_grid_desc_g_n_k_.GetLength(I0);
             const index_t c_m       = kernel_arg.y_grid_desc_m_o_.GetLength(I0);
             const index_t c_gemm1n  = kernel_arg.y_grid_desc_m_o_.GetLength(I1);
             const index_t a_m       = kernel_arg.a_grid_desc_ak0_m_ak1_.GetLength(I1);
             const index_t b1_gemm1n = kernel_arg.b1_grid_desc_bk0_n_bk1_.GetLength(I0) *
                                       kernel_arg.b1_grid_desc_bk0_n_bk1_.GetLength(I2);
 
-            if(!(c_g == device_arg.batch_count_ && c_m == a_m && c_gemm1n == b1_gemm1n))
+            if(!(c_g == device_arg.batch_count_ && c_m == a_m && c_gemm1n == b1_gemm1n &&
+                 c_g % b_g == 0 && c_g / b_g == arg.h_ratio_))
             {
                 return false;
             }
@@ -1252,6 +1297,17 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
                 return false;
             }
 
+            // saving dQ data with atomic_add instruction, so KzRaw must be a multiple of 2
+            if constexpr(is_same<OutputDataType, half_t>::value ||
+                         is_same<OutputDataType, bhalf_t>::value)
+            {
+                if(KzRaw % 2 != 0)
+                {
+                    std::cout << "K_q must be a multiple of 2" << std::endl;
+                    return false;
+                }
+            }
+
             // Check vector load/store requirement
             const auto a_stride_lowest = ABlockTransferSrcVectorDim == 2
                                              ? device_arg.a_mz_kz_strides_[1]

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

@@ -694,7 +694,7 @@ struct DeviceGroupedMultiheadAttentionForward_Xdl_CShuffle_V1
                     b1_grid_desc_g_n_k,
                     c_grid_desc_g_m_n,
                     z_grid_desc_g_m_n,
-                    type_convert<index_t>(lse_grid_desc_m.GetElementSpaceSize()));
+                    type_convert<index_t>(lse_gs_ms_strides[NumDimG - 1]));
 
                 // C0 mask
                 const auto c0_matrix_mask =

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

@@ -35,8 +35,7 @@ template <typename GridwiseGemm,
           typename CElementwiseOperation,
           bool HasMainKBlockLoop,
           bool IsDropout,
-          bool IsLseStoring,
-          bool Deterministic>
+          bool IsLseStoring>
 __global__ void
 #if CK_USE_LAUNCH_BOUNDS
     __launch_bounds__(CK_MAX_THREAD_PER_BLOCK, CK_MIN_BLOCK_PER_CU)
@@ -44,6 +43,7 @@ __global__ void
         kernel_grouped_gemm_softmax_gemm_xdl_cshuffle_v2(
             const void CK_CONSTANT_ADDRESS_SPACE* group_kernel_args,
             const index_t group_count,
+            const index_t h_ratio,
             const AElementwiseOperation a_element_op,
             const BElementwiseOperation b_element_op,
             const AccElementwiseOperation acc_element_op,
@@ -87,14 +87,15 @@ __global__ void
     // per-group batch offset
     const index_t num_blocks_per_batch = arg_ptr[group_id].num_blocks_per_batch_;
     const index_t g_idx                = __builtin_amdgcn_readfirstlane(
-        (block_id - arg_ptr[group_id].block_start_) / (Deterministic ? 1 : num_blocks_per_batch));
+        (block_id - arg_ptr[group_id].block_start_) / num_blocks_per_batch);
+    const index_t gkv_idx = __builtin_amdgcn_readfirstlane(g_idx / h_ratio);
 
     const long_index_t a_batch_offset = __builtin_amdgcn_readfirstlane(
         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 b_batch_offset  = __builtin_amdgcn_readfirstlane(static_cast<long_index_t>(
+        arg_ptr[group_id].compute_base_ptr_of_batch_.GetBBasePtr(gkv_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)));
+        arg_ptr[group_id].compute_base_ptr_of_batch_.GetB1BasePtr(gkv_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 z_batch_offset = __builtin_amdgcn_readfirstlane(
@@ -113,87 +114,42 @@ __global__ void
         tmp_p_d0_grid = arg_ptr[group_id].p_d0_grid_ + d0_batch_offset;
     }
 
-    if constexpr(Deterministic)
-    {
-        for(index_t i = 0; i < num_blocks_per_batch; i++)
-        {
-            GridwiseGemm::template Run<HasMainKBlockLoop, IsDropout, IsLseStoring>(
-                arg_ptr[group_id].p_a_grid_ + a_batch_offset,
-                arg_ptr[group_id].p_b_grid_ + b_batch_offset,
-                tmp_p_d0_grid,
-                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_z_grid_ == nullptr
-                    ? nullptr
-                    : arg_ptr[group_id].p_z_grid_ + z_batch_offset,
-                arg_ptr[group_id].p_lse_grid_ == nullptr
-                    ? nullptr
-                    : arg_ptr[group_id].p_lse_grid_ + lse_batch_offset,
-                // arg_ptr[group_id].p_lse_grid_ + lse_batch_offset,
-                p_shared,
-                a_element_op,
-                b_element_op,
-                acc_element_op,
-                b1_element_op,
-                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].d0_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].c_grid_desc_mblock_mperblock_nblock_nperblock_,
-                arg_ptr[group_id].z_grid_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5_,
-                arg_ptr[group_id].lse_grid_desc_m_,
-                arg_ptr[group_id].block_2_ctile_map_,
-                arg_ptr[group_id].c0_matrix_mask_,
-                p_dropout_in_uint8_t,
-                p_dropout_rescale,
-                ph,
-                arg_ptr[group_id].z_random_matrix_offset_ +
-                    g_idx * arg_ptr[group_id].raw_m_padded_ * arg_ptr[group_id].raw_n_padded_,
-                arg_ptr[group_id].raw_n_padded_,
-                i);
-        }
-    }
-    else
-    {
-        GridwiseGemm::template Run<HasMainKBlockLoop, IsDropout, IsLseStoring>(
-            arg_ptr[group_id].p_a_grid_ + a_batch_offset,
-            arg_ptr[group_id].p_b_grid_ + b_batch_offset,
-            tmp_p_d0_grid,
-            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_z_grid_ == nullptr ? nullptr
-                                                   : arg_ptr[group_id].p_z_grid_ + z_batch_offset,
-            arg_ptr[group_id].p_lse_grid_ == nullptr
-                ? nullptr
-                : arg_ptr[group_id].p_lse_grid_ + lse_batch_offset,
-            // arg_ptr[group_id].p_lse_grid_ + lse_batch_offset,
-            p_shared,
-            a_element_op,
-            b_element_op,
-            acc_element_op,
-            b1_element_op,
-            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].d0_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].c_grid_desc_mblock_mperblock_nblock_nperblock_,
-            arg_ptr[group_id].z_grid_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5_,
-            arg_ptr[group_id].lse_grid_desc_m_,
-            arg_ptr[group_id].block_2_ctile_map_,
-            arg_ptr[group_id].c0_matrix_mask_,
-            p_dropout_in_uint8_t,
-            p_dropout_rescale,
-            ph,
-            arg_ptr[group_id].z_random_matrix_offset_ +
-                g_idx * arg_ptr[group_id].raw_m_padded_ * arg_ptr[group_id].raw_n_padded_,
-            arg_ptr[group_id].raw_n_padded_,
-            0);
-    }
+    GridwiseGemm::template Run<HasMainKBlockLoop, IsDropout, IsLseStoring>(
+        arg_ptr[group_id].p_a_grid_ + a_batch_offset,
+        arg_ptr[group_id].p_b_grid_ + b_batch_offset,
+        tmp_p_d0_grid,
+        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_z_grid_ == nullptr ? nullptr
+                                               : arg_ptr[group_id].p_z_grid_ + z_batch_offset,
+        arg_ptr[group_id].p_lse_grid_ == nullptr ? nullptr
+                                                 : arg_ptr[group_id].p_lse_grid_ + lse_batch_offset,
+        // arg_ptr[group_id].p_lse_grid_ + lse_batch_offset,
+        p_shared,
+        a_element_op,
+        b_element_op,
+        acc_element_op,
+        b1_element_op,
+        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].d0_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].c_grid_desc_mblock_mperblock_nblock_nperblock_,
+        arg_ptr[group_id].z_grid_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5_,
+        arg_ptr[group_id].lse_grid_desc_m_,
+        arg_ptr[group_id].block_2_ctile_map_,
+        arg_ptr[group_id].c0_matrix_mask_,
+        p_dropout_in_uint8_t,
+        p_dropout_rescale,
+        ph,
+        arg_ptr[group_id].z_random_matrix_offset_ +
+            g_idx * arg_ptr[group_id].raw_m_padded_ * arg_ptr[group_id].raw_n_padded_,
+        arg_ptr[group_id].raw_n_padded_);
 #else
     ignore = group_kernel_args;
     ignore = group_count;
+    ignore = h_ratio;
     ignore = a_element_op;
     ignore = b_element_op;
     ignore = acc_element_op;
@@ -279,7 +235,6 @@ template <index_t NumDimG,
           index_t CShuffleBlockTransferScalarPerVector_NPerBlock,
           index_t Acc1BiasTransferSrcScalarPerVector,
           MaskingSpecialization MaskingSpec,
-          bool Deterministic,
           LoopScheduler LoopSched = LoopScheduler::Default>
 struct DeviceGroupedMultiheadAttentionForward_Xdl_CShuffle_V2
     : public DeviceGroupedMultiheadAttentionForward<NumDimG,
@@ -415,7 +370,6 @@ struct DeviceGroupedMultiheadAttentionForward_Xdl_CShuffle_V2
     MakeD0GridDescriptor_M_N(const std::vector<ck::index_t>& acc0_biases_gs_ms_ns_lengths,
                              const std::vector<ck::index_t>& acc0_biases_gs_ms_ns_strides)
     {
-
         return Transform::MakeC0GridDescriptor_M_N(acc0_biases_gs_ms_ns_lengths,
                                                    acc0_biases_gs_ms_ns_strides);
     }
@@ -424,7 +378,6 @@ struct DeviceGroupedMultiheadAttentionForward_Xdl_CShuffle_V2
     MakeD0GridDescriptor_G_M_N(const std::vector<ck::index_t>& acc0_biases_gs_ms_ns_lengths,
                                const std::vector<ck::index_t>& acc0_biases_gs_ms_ns_strides)
     {
-
         return Transform::MakeC0GridDescriptor_G_M_N(acc0_biases_gs_ms_ns_lengths,
                                                      acc0_biases_gs_ms_ns_strides);
     }
@@ -597,8 +550,7 @@ struct DeviceGroupedMultiheadAttentionForward_Xdl_CShuffle_V2
         Acc1BiasTransferSrcScalarPerVector,
         LoopSched,
         Transform::matrix_padder.PadN,
-        MaskingSpec != MaskingSpecialization::MaskDisabled,
-        Deterministic>;
+        MaskingSpec != MaskingSpecialization::MaskDisabled>;
 
     using Block2CTileMap = OffsettedBlockToCTileMap<typename GridwiseGemm::DefaultBlock2CTileMap>;
 
@@ -655,6 +607,8 @@ struct DeviceGroupedMultiheadAttentionForward_Xdl_CShuffle_V2
 
         // for gridwise gemm check
         CGridDesc_M_N c_grid_desc_m_n_;
+        BGridDesc_G_N_K b_grid_desc_g_n_k_;
+        CGridDesc_G_M_N c_grid_desc_g_m_n_;
 
         // raw data
         std::vector<ck::index_t> d0_n_length_stride_;
@@ -703,6 +657,9 @@ struct DeviceGroupedMultiheadAttentionForward_Xdl_CShuffle_V2
 
             index_t z_random_matrix_offset = 0;
 
+            h_ratio_ = problem_desc_vec[0].a_gs_ms_ks_lengths[NumDimG - 1] /
+                       problem_desc_vec[0].b0_gs_ns_ks_lengths[NumDimG - 1];
+
             for(std::size_t i = 0; i < group_count_; i++)
             {
                 const auto p_a_grid  = static_cast<const ADataType*>(p_a_vec[i]);
@@ -783,8 +740,7 @@ struct DeviceGroupedMultiheadAttentionForward_Xdl_CShuffle_V2
                 const auto block_2_ctile_map = Block2CTileMap(c_grid_desc_m_n, BlockStart);
                 const index_t batch_count    = c_grid_desc_g_m_n.GetLength(I0);
                 const index_t grid_size_grp =
-                    (Deterministic ? 1 : block_2_ctile_map.CalculateGridSize(c_grid_desc_m_n)) *
-                    batch_count;
+                    block_2_ctile_map.CalculateGridSize(c_grid_desc_m_n) * batch_count;
                 const index_t BlockEnd = grid_size_ + grid_size_grp;
 
                 // batch stride
@@ -795,7 +751,7 @@ struct DeviceGroupedMultiheadAttentionForward_Xdl_CShuffle_V2
                     b1_grid_desc_g_n_k,
                     c_grid_desc_g_m_n,
                     z_grid_desc_g_m_n,
-                    type_convert<index_t>(lse_grid_desc_m.GetElementSpaceSize()));
+                    type_convert<index_t>(problem_desc.lse_gs_ms_strides[NumDimG - 1]));
 
                 // C0 mask
                 const auto c0_matrix_mask =
@@ -855,6 +811,8 @@ struct DeviceGroupedMultiheadAttentionForward_Xdl_CShuffle_V2
                      {problem_desc.c_gs_ms_os_strides[NumDimG + NumDimM - 1],
                       problem_desc.c_gs_ms_os_strides[NumDimG + NumDimM + NumDimO - 1]},
                      c_grid_desc_m_n,
+                     b_grid_desc_g_n_k,
+                     c_grid_desc_g_m_n,
                      d0_n_length_stride});
             }
 
@@ -880,6 +838,7 @@ struct DeviceGroupedMultiheadAttentionForward_Xdl_CShuffle_V2
         B1ElementwiseOperation b1_element_op_;
         CElementwiseOperation c_element_op_;
 
+        index_t h_ratio_;
         float p_dropout_;
         uint8_t p_dropout_in_uint8_t_;
         unsigned long long seed_;
@@ -958,8 +917,7 @@ struct DeviceGroupedMultiheadAttentionForward_Xdl_CShuffle_V2
                                                                          CElementwiseOperation,
                                                                          has_main_k_block_loop_,
                                                                          use_dropout_,
-                                                                         is_lse_storing_,
-                                                                         Deterministic>;
+                                                                         is_lse_storing_>;
 
                     return launch_and_time_kernel(
                         stream_config,
@@ -969,6 +927,7 @@ struct DeviceGroupedMultiheadAttentionForward_Xdl_CShuffle_V2
                         0,
                         cast_pointer_to_constant_address_space(arg.p_workspace_),
                         arg.group_count_,
+                        arg.h_ratio_,
                         arg.a_element_op_,
                         arg.b_element_op_,
                         arg.acc_element_op_,
@@ -1091,11 +1050,14 @@ struct DeviceGroupedMultiheadAttentionForward_Xdl_CShuffle_V2
             const auto& device_arg = arg.group_device_args_[i];
 
             // Check if C permute dimension matches GEMM + GEMM shape
+            const index_t c_g       = device_arg.c_grid_desc_g_m_n_.GetLength(I0); // unpadded
+            const index_t b_g       = device_arg.b_grid_desc_g_n_k_.GetLength(I0);
             const index_t c_m       = device_arg.c_grid_desc_m_n_.GetLength(I0);
             const index_t c_gemm1n  = device_arg.c_grid_desc_m_n_.GetLength(I1);
             const index_t a_m       = kernel_arg.a_grid_desc_ak0_m_ak1_.GetLength(I1);
             const index_t b1_gemm1n = kernel_arg.b1_grid_desc_bk0_n_bk1_.GetLength(I1);
-            if(!(c_m == a_m && c_gemm1n == b1_gemm1n))
+            if(!(c_m == a_m && c_gemm1n == b1_gemm1n && c_g % b_g == 0 &&
+                 c_g / b_g == arg.h_ratio_))
             {
                 return false;
             }

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

@@ -5,6 +5,7 @@
 
 #include <iostream>
 #include <sstream>
+#include <cstring>
 
 #include "ck/utility/common_header.hpp"
 #include "ck/tensor_description/tensor_descriptor.hpp"
@@ -685,12 +686,34 @@ struct DeviceGroupedMultiheadAttentionInfer_Xdl_CShuffle
                 some_has_main_k_block_loop |= y;
             }
 
-            hipGetErrorString(
-                hipMemcpyWithStream(arg.p_workspace_,
-                                    arg.group_kernel_args_.data(),
-                                    arg.group_kernel_args_.size() * sizeof(GroupKernelArg),
-                                    hipMemcpyHostToDevice,
-                                    stream_config.stream_id_));
+            hipStreamCaptureStatus status = hipStreamCaptureStatusNone;
+
+            HIP_CHECK_ERROR(hipStreamIsCapturing(stream_config.stream_id_, &status));
+
+            if(status == hipStreamCaptureStatusActive)
+            {
+                size_t copy_size = arg.group_kernel_args_.size() * sizeof(GroupKernelArg);
+
+                // ToDO: when to release this memory buffer?
+                char* persistent_ptr = new char[copy_size];
+
+                (void)std::memcpy(persistent_ptr, arg.group_kernel_args_.data(), copy_size);
+
+                HIP_CHECK_ERROR(hipMemcpyAsync(arg.p_workspace_,
+                                               persistent_ptr,
+                                               copy_size,
+                                               hipMemcpyHostToDevice,
+                                               stream_config.stream_id_));
+            }
+            else
+            {
+                HIP_CHECK_ERROR(
+                    hipMemcpyAsync(arg.p_workspace_,
+                                   arg.group_kernel_args_.data(),
+                                   arg.group_kernel_args_.size() * sizeof(GroupKernelArg),
+                                   hipMemcpyHostToDevice,
+                                   stream_config.stream_id_));
+            }
 
             float ave_time = 0;
 

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

@@ -88,6 +88,9 @@ template <typename InputDataType,
           PipelineVersion PipelineVer = PipelineVersion::v1>
 struct GridwiseBatchedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V1
 {
+    static_assert(AK1Value % ABlockTransferDstScalarPerVector_AK1 == 0);
+    static_assert(BK1Value % BBlockTransferDstScalarPerVector_BK1 == 0);
+
     static_assert(KPerBlock == Gemm1NPerBlock);
     static_assert(MPerBlock % Gemm1KPerBlock == 0);
     static_assert(NPerBlock % Gemm2KPerBlock == 0);
@@ -1440,10 +1443,12 @@ struct GridwiseBatchedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V1
         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 KGridDesc_K0_N_K1& kgrad_grid_desc_k0_n_k1,
         const D0GridDescriptor_M0_N0_M1_M2_N1_M3& d0_grid_desc_m0_n0_m1_m2_n1_m3,
         const ZGridDescriptor_M0_N0_M1_N1_M2_N2_M3_M4_M5_N3&
             z_grid_desc_m0_n0_m1_n1_m2_n2_m3_m4_m5_n3,
         const VGridDesc_O0_N_O1& v_grid_desc_o0_n_o1,
+        const VGridDesc_O0_N_O1& vgrad_grid_desc_o0_n_o1,
         const LSEGridDesc_M& lse_grid_desc_m,
         const YGradGridDesc_O0_M_O1& ygrad_grid_desc_o0_m_o1,
         const Block2CTileMap& block_2_ctile_map,
@@ -1474,11 +1479,11 @@ struct GridwiseBatchedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V1
         const auto ygrad_grid_buf = make_dynamic_buffer<AddressSpaceEnum::Global>(
             p_ygrad_grid, ygrad_grid_desc_o0_m_o1.GetElementSpaceSize());
         auto vgrad_grid_buf = make_dynamic_buffer<AddressSpaceEnum::Global>(
-            p_vgrad_grid, v_grid_desc_o0_n_o1.GetElementSpaceSize());
+            p_vgrad_grid, vgrad_grid_desc_o0_n_o1.GetElementSpaceSize());
         auto qgrad_grid_buf = make_dynamic_buffer<AddressSpaceEnum::Global>(
             p_qgrad_grid, q_grid_desc_k0_m_k1.GetElementSpaceSize());
         auto kgrad_grid_buf = make_dynamic_buffer<AddressSpaceEnum::Global>(
-            p_kgrad_grid, k_grid_desc_k0_n_k1.GetElementSpaceSize());
+            p_kgrad_grid, kgrad_grid_desc_k0_n_k1.GetElementSpaceSize());
 
         // divide block work by [N, K]
         const auto block_work_idx =
@@ -1628,7 +1633,7 @@ struct GridwiseBatchedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V1
 
         // dV: transform input and output tensor descriptors
         auto vgrad_grid_desc_nblock_nperblock_oblock_operblock =
-            MakeVGradGridDesc_NBlock_NPerBlock_OBlock_OPerBlock(v_grid_desc_o0_n_o1);
+            MakeVGradGridDesc_NBlock_NPerBlock_OBlock_OPerBlock(vgrad_grid_desc_o0_n_o1);
 
         // dK: A matrix blockwise copy
         auto kgrad_gemm_tile_sgrad_blockwise_copy =
@@ -1657,7 +1662,7 @@ struct GridwiseBatchedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V1
 
         // dK: transform input and output tensor descriptors
         auto kgrad_grid_desc_nblock_nperblock_oblock_operblock =
-            MakeKGradGridDesc_NBlock_NPerBlock_OBlock_OPerBlock(k_grid_desc_k0_n_k1);
+            MakeKGradGridDesc_NBlock_NPerBlock_OBlock_OPerBlock(kgrad_grid_desc_k0_n_k1);
 
         //
         // set up dQ Gemm (type 3 crr)

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

@@ -96,6 +96,10 @@ template <typename InputDataType,
           PipelineVersion PipelineVer = PipelineVersion::v1>
 struct GridwiseBatchedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V2
 {
+    static_assert(AK1Value % ABlockTransferDstScalarPerVector_AK1 == 0);
+    static_assert(BK1Value % BBlockTransferDstScalarPerVector_BK1 == 0);
+    static_assert(B1K1Value % B1BlockTransferDstScalarPerVector_BK1 == 0);
+
     static_assert(Gemm1NPerBlock % KPerBlock == 0);
     static_assert(MPerBlock % Gemm1KPerBlock == 0);
     static_assert(NPerBlock % Gemm2KPerBlock == 0);
@@ -1531,10 +1535,12 @@ struct GridwiseBatchedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_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 KGridDesc_K0_N_K1& kgrad_grid_desc_k0_n_k1,
         const D0GridDescriptor_M0_N0_M1_M2_N1_M3& d0_grid_desc_m0_n0_m1_m2_n1_m3,
         const ZGridDescriptor_M0_N0_M1_N1_M2_N2_M3_M4_M5_N3&
             z_grid_desc_m0_n0_m1_n1_m2_n2_m3_m4_m5_n3,
         const VGridDesc_O0_N_O1& v_grid_desc_o0_n_o1,
+        const VGridDesc_O0_N_O1& vgrad_grid_desc_o0_n_o1,
         const LSEGridDesc_M& lse_grid_desc_m,
         const YGradGridDesc_M0_O_M1& ygrad_grid_desc_m0_o_m1,
         const Block2CTileMap& block_2_ctile_map,
@@ -1565,11 +1571,11 @@ struct GridwiseBatchedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V2
         const auto ygrad_grid_buf = make_dynamic_buffer<AddressSpaceEnum::Global>(
             p_ygrad_grid, ygrad_grid_desc_m0_o_m1.GetElementSpaceSize());
         auto vgrad_grid_buf = make_dynamic_buffer<AddressSpaceEnum::Global>(
-            p_vgrad_grid, v_grid_desc_o0_n_o1.GetElementSpaceSize());
+            p_vgrad_grid, vgrad_grid_desc_o0_n_o1.GetElementSpaceSize());
         auto qgrad_grid_buf = make_dynamic_buffer<AddressSpaceEnum::Global>(
             p_qgrad_grid, q_grid_desc_k0_m_k1.GetElementSpaceSize());
         auto kgrad_grid_buf = make_dynamic_buffer<AddressSpaceEnum::Global>(
-            p_kgrad_grid, k_grid_desc_k0_n_k1.GetElementSpaceSize());
+            p_kgrad_grid, kgrad_grid_desc_k0_n_k1.GetElementSpaceSize());
 
         // divide block work by [N, K]
         const auto block_work_idx =
@@ -1742,7 +1748,7 @@ struct GridwiseBatchedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V2
 
         // dV: transform input and output tensor descriptors
         auto vgrad_grid_desc_nblock_nperblock_oblock_operblock =
-            MakeVGradGridDesc_NBlock_NPerBlock_OBlock_OPerBlock(v_grid_desc_o0_n_o1);
+            MakeVGradGridDesc_NBlock_NPerBlock_OBlock_OPerBlock(vgrad_grid_desc_o0_n_o1);
 
         // dK: transform input and output tensor descriptors
         const auto q_grid_desc_m0_k_m1 =
@@ -1775,7 +1781,7 @@ struct GridwiseBatchedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V2
 
         // dK: transform input and output tensor descriptors
         auto kgrad_grid_desc_nblock_nperblock_oblock_operblock =
-            MakeKGradGridDesc_NBlock_NPerBlock_OBlock_OPerBlock(k_grid_desc_k0_n_k1);
+            MakeKGradGridDesc_NBlock_NPerBlock_OBlock_OPerBlock(kgrad_grid_desc_k0_n_k1);
 
         //
         // set up dQ Gemm (type 3 crr)

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

@@ -87,6 +87,9 @@ template <typename InputDataType,
           PipelineVersion PipelineVer = PipelineVersion::v1>
 struct GridwiseBatchedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V1
 {
+    static_assert(AK1Value % ABlockTransferDstScalarPerVector_AK1 == 0);
+    static_assert(BK1Value % BBlockTransferDstScalarPerVector_BK1 == 0);
+
     static_assert(KPerBlock == Gemm1NPerBlock);
     static_assert(MPerBlock % Gemm1KPerBlock == 0);
     static_assert(NPerBlock % Gemm2KPerBlock == 0);
@@ -1521,10 +1524,12 @@ struct GridwiseBatchedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V1
         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 KGridDesc_K0_N_K1& kgrad_grid_desc_k0_n_k1,
         const D0GridDescriptor_M0_N0_M1_M2_N1_M3& d0_grid_desc_m0_n0_m1_m2_n1_m3,
         const ZGridDescriptor_M0_N0_M1_N1_M2_N2_M3_M4_M5_N3&
             z_grid_desc_m0_n0_m1_n1_m2_n2_m3_m4_m5_n3,
         const VGridDesc_O0_N_O1& v_grid_desc_o0_n_o1,
+        const VGridDesc_O0_N_O1& vgrad_grid_desc_o0_n_o1,
         const YGridDescriptor_MBlock_MPerBlock_OBlock_OPerBlock&
             y_grid_desc_mblock_mperblock_oblock_operblock,
         const LSEGridDesc_M& lse_grid_desc_m,
@@ -1557,11 +1562,11 @@ struct GridwiseBatchedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V1
         const auto ygrad_grid_buf = make_dynamic_buffer<AddressSpaceEnum::Global>(
             p_ygrad_grid, ygrad_grid_desc_o0_m_o1.GetElementSpaceSize());
         auto vgrad_grid_buf = make_dynamic_buffer<AddressSpaceEnum::Global>(
-            p_vgrad_grid, v_grid_desc_o0_n_o1.GetElementSpaceSize());
+            p_vgrad_grid, vgrad_grid_desc_o0_n_o1.GetElementSpaceSize());
         auto qgrad_grid_buf = make_dynamic_buffer<AddressSpaceEnum::Global>(
             p_qgrad_grid, q_grid_desc_k0_m_k1.GetElementSpaceSize());
         auto kgrad_grid_buf = make_dynamic_buffer<AddressSpaceEnum::Global>(
-            p_kgrad_grid, k_grid_desc_k0_n_k1.GetElementSpaceSize());
+            p_kgrad_grid, kgrad_grid_desc_k0_n_k1.GetElementSpaceSize());
 
         // divide block work by [N, K]
         const auto block_work_idx =
@@ -1711,7 +1716,7 @@ struct GridwiseBatchedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V1
 
         // dV: transform input and output tensor descriptors
         auto vgrad_grid_desc_nblock_nperblock_oblock_operblock =
-            MakeVGradGridDesc_NBlock_NPerBlock_OBlock_OPerBlock(v_grid_desc_o0_n_o1);
+            MakeVGradGridDesc_NBlock_NPerBlock_OBlock_OPerBlock(vgrad_grid_desc_o0_n_o1);
 
         // dK: A matrix blockwise copy
         auto kgrad_gemm_tile_sgrad_blockwise_copy =
@@ -1740,7 +1745,7 @@ struct GridwiseBatchedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V1
 
         // dK: transform input and output tensor descriptors
         auto kgrad_grid_desc_nblock_nperblock_oblock_operblock =
-            MakeKGradGridDesc_NBlock_NPerBlock_OBlock_OPerBlock(k_grid_desc_k0_n_k1);
+            MakeKGradGridDesc_NBlock_NPerBlock_OBlock_OPerBlock(kgrad_grid_desc_k0_n_k1);
 
         //
         // set up dQ Gemm (type 3 crr)

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

@@ -95,6 +95,10 @@ template <typename InputDataType,
           PipelineVersion PipelineVer = PipelineVersion::v1>
 struct GridwiseBatchedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
 {
+    static_assert(AK1Value % ABlockTransferDstScalarPerVector_AK1 == 0);
+    static_assert(BK1Value % BBlockTransferDstScalarPerVector_BK1 == 0);
+    static_assert(B1K1Value % B1BlockTransferDstScalarPerVector_BK1 == 0);
+
     static_assert(Gemm1NPerBlock % KPerBlock == 0);
     static_assert(MPerBlock % Gemm1KPerBlock == 0);
     static_assert(NPerBlock % Gemm2KPerBlock == 0);
@@ -320,18 +324,27 @@ struct GridwiseBatchedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
         // dV_NOM / dK_NKM Gemm (Gemm2 crr)
         if(O != K)
         {
+            std::cerr << "O = " << O << " K = " << K << std::endl;
             std::cerr << "SizeK must be equal to SizeO (equal attention head size)" << '\n';
             return false;
         }
 
         if(!(M == y_grid_desc_m_o.GetLength(I0) && O == y_grid_desc_m_o.GetLength(I1)))
         {
+            std::cerr << "M = " << M << " O = " << O
+                      << " y_grid_desc_m_o = " << y_grid_desc_m_o.GetLength(I0) << " , "
+                      << y_grid_desc_m_o.GetLength(I1) << std::endl;
+            std::cerr << "Un-matched sizes!" << std::endl;
             return false;
         }
 
         if(!(M % MPerBlock == 0 && N % NPerBlock == 0 && K % KPerBlock == 0 &&
              O % Gemm1NPerBlock == 0))
         {
+            std::cerr << "M = " << M << " N = " << N << " O = " << O << std::endl;
+            std::cerr << "MPerBlock = " << MPerBlock << " NPerBlock = " << NPerBlock
+                      << " KPerBlock = " << KPerBlock << std::endl;
+            std::cerr << "Un-aligned sizes!" << std::endl;
             return false;
         }
 
@@ -1587,10 +1600,12 @@ struct GridwiseBatchedMultiheadAttentionBackward_Qloop_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 KGridDesc_K0_N_K1& kgrad_grid_desc_k0_n_k1,
         const D0GridDescriptor_M0_N0_M1_M2_N1_M3& d0_grid_desc_m0_n0_m1_m2_n1_m3,
         const ZGridDescriptor_M0_N0_M1_N1_M2_N2_M3_M4_M5_N3&
             z_grid_desc_m0_n0_m1_n1_m2_n2_m3_m4_m5_n3,
         const VGridDesc_O0_N_O1& v_grid_desc_o0_n_o1,
+        const VGridDesc_O0_N_O1& vgrad_grid_desc_o0_n_o1,
         const YGridDescriptor_MBlock_MPerBlock_OBlock_OPerBlock&
             y_grid_desc_mblock_mperblock_oblock_operblock,
         const LSEGridDesc_M& lse_grid_desc_m,
@@ -1623,11 +1638,11 @@ struct GridwiseBatchedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
         const auto ygrad_grid_buf = make_dynamic_buffer<AddressSpaceEnum::Global>(
             p_ygrad_grid, ygrad_grid_desc_m0_o_m1.GetElementSpaceSize());
         auto vgrad_grid_buf = make_dynamic_buffer<AddressSpaceEnum::Global>(
-            p_vgrad_grid, v_grid_desc_o0_n_o1.GetElementSpaceSize());
+            p_vgrad_grid, vgrad_grid_desc_o0_n_o1.GetElementSpaceSize());
         auto qgrad_grid_buf = make_dynamic_buffer<AddressSpaceEnum::Global>(
             p_qgrad_grid, q_grid_desc_k0_m_k1.GetElementSpaceSize());
         auto kgrad_grid_buf = make_dynamic_buffer<AddressSpaceEnum::Global>(
-            p_kgrad_grid, k_grid_desc_k0_n_k1.GetElementSpaceSize());
+            p_kgrad_grid, kgrad_grid_desc_k0_n_k1.GetElementSpaceSize());
 
         // divide block work by [N, K]
         const auto block_work_idx =
@@ -1800,7 +1815,7 @@ struct GridwiseBatchedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
 
         // dV: transform input and output tensor descriptors
         auto vgrad_grid_desc_nblock_nperblock_oblock_operblock =
-            MakeVGradGridDesc_NBlock_NPerBlock_OBlock_OPerBlock(v_grid_desc_o0_n_o1);
+            MakeVGradGridDesc_NBlock_NPerBlock_OBlock_OPerBlock(vgrad_grid_desc_o0_n_o1);
 
         // dK: transform input and output tensor descriptors
         const auto q_grid_desc_m0_k_m1 =
@@ -1833,7 +1848,7 @@ struct GridwiseBatchedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
 
         // dK: transform input and output tensor descriptors
         auto kgrad_grid_desc_nblock_nperblock_oblock_operblock =
-            MakeKGradGridDesc_NBlock_NPerBlock_OBlock_OPerBlock(k_grid_desc_k0_n_k1);
+            MakeKGradGridDesc_NBlock_NPerBlock_OBlock_OPerBlock(kgrad_grid_desc_k0_n_k1);
 
         //
         // set up dQ Gemm (type 3 crr)

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

@@ -94,10 +94,13 @@ template <typename FloatAB,
           LoopScheduler LoopSched,
           bool PadN,
           bool MaskOutUpperTriangle,
-          bool Deterministic,
           PipelineVersion PipelineVer = PipelineVersion::v1>
 struct GridwiseBatchedMultiheadAttentionForward_Xdl_CShuffle_V2
 {
+    static_assert(AK1Value % ABlockTransferDstScalarPerVector_AK1 == 0);
+    static_assert(BK1Value % BBlockTransferDstScalarPerVector_BK1 == 0);
+    static_assert(B1K1Value % B1BlockTransferDstScalarPerVector_BK1 == 0);
+
     static_assert(D0BlockTransferSrcScalarPerVector == 1 ||
                       D0BlockTransferSrcScalarPerVector == 2 ||
                       D0BlockTransferSrcScalarPerVector == 4,
@@ -531,8 +534,7 @@ struct GridwiseBatchedMultiheadAttentionForward_Xdl_CShuffle_V2
                                FloatGemmAcc p_dropout_rescale,
                                ck::philox& ph,
                                const index_t z_random_matrix_offset,
-                               const index_t raw_n_padded,
-                               const index_t block_idx_m)
+                               const index_t raw_n_padded)
     {
         const auto a_grid_buf = make_dynamic_buffer<AddressSpaceEnum::Global>(
             p_a_grid, a_grid_desc_ak0_m_ak1.GetElementSpaceSize());
@@ -557,7 +559,7 @@ struct GridwiseBatchedMultiheadAttentionForward_Xdl_CShuffle_V2
             return;
         }
 
-        const index_t block_work_idx_m = Deterministic ? block_idx_m : block_work_idx[I0];
+        const index_t block_work_idx_m = block_work_idx[I0];
 
         // HACK: this force m/gemm1_n_block_data_idx_on_grid into SGPR
         const index_t m_block_data_idx_on_grid =
@@ -1145,11 +1147,6 @@ struct GridwiseBatchedMultiheadAttentionForward_Xdl_CShuffle_V2
                                                                       0),
                                                      tensor_operation::element_wise::PassThrough{}};
 
-        if constexpr(Deterministic)
-        {
-            block_sync_lds();
-        }
-
         do
         {
             auto n_block_data_idx_on_grid =

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

@@ -88,6 +88,10 @@ template <typename FloatAB,
           PipelineVersion PipelineVer = PipelineVersion::v1>
 struct GridwiseMultiHeadFlashAttentionInfer_Xdl_CShuffle
 {
+    static_assert(AK1Value % ABlockTransferDstScalarPerVector_AK1 == 0);
+    static_assert(BK1Value % BBlockTransferDstScalarPerVector_BK1 == 0);
+    static_assert(B1K1Value % B1BlockTransferDstScalarPerVector_BK1 == 0);
+
     static_assert(LoopSched == LoopScheduler::Default,
                   "Non-default loop scheduler is currently not supported");
 

include/ck/utility/type_convert.hpp

@@ -31,6 +31,51 @@ inline __host__ __device__ constexpr float type_convert<float, bhalf_t>(bhalf_t
     return u.fp32;
 }
 
+#ifdef USE_RTN_BF16_CONVERT
+// Convert fp32 to bf16 with RTN if higher precision is needed
+template <>
+inline __host__ __device__ constexpr bhalf_t type_convert<bhalf_t, float>(float x)
+{
+    union
+    {
+        float fp32;
+        uint32_t int32;
+    } u = {x};
+
+    // When the exponent bits are not all 1s, then the value is zero, normal,
+    // or subnormal. We round the bfloat16 mantissa up by adding 0x7FFF, plus
+    // 1 if the least significant bit of the bfloat16 mantissa is 1 (odd).
+    // This causes the bfloat16's mantissa to be incremented by 1 if the 16
+    // least significant bits of the float mantissa are greater than 0x8000,
+    // or if they are equal to 0x8000 and the least significant bit of the
+    // bfloat16 mantissa is 1 (odd). This causes it to be rounded to even when
+    // the lower 16 bits are exactly 0x8000. If the bfloat16 mantissa already
+    // has the value 0x7f, then incrementing it causes it to become 0x00 and
+    // the exponent is incremented by one, which is the next higher FP value
+    // to the unrounded bfloat16 value. When the bfloat16 value is subnormal
+    // with an exponent of 0x00 and a mantissa of 0x7f, it may be rounded up
+    // to a normal value with an exponent of 0x01 and a mantissa of 0x00.
+    // When the bfloat16 value has an exponent of 0xFE and a mantissa of 0x7F,
+    // incrementing it causes it to become an exponent of 0xFF and a mantissa
+    // of 0x00, which is Inf, the next higher value to the unrounded value.
+    bool flag0 = ~u.int32 & 0x7f800000;
+
+    // When all of the exponent bits are 1, the value is Inf or NaN.
+    // Inf is indicated by a zero mantissa. NaN is indicated by any nonzero
+    // mantissa bit. Quiet NaN is indicated by the most significant mantissa
+    // bit being 1. Signaling NaN is indicated by the most significant
+    // mantissa bit being 0 but some other bit(s) being 1. If any of the
+    // lower 16 bits of the mantissa are 1, we set the least significant bit
+    // of the bfloat16 mantissa, in order to preserve signaling NaN in case
+    // the bfloat16's mantissa bits are all 0.
+    bool flag1 = !flag0 && (u.int32 & 0xffff);
+
+    u.int32 += flag0 ? 0x7fff + ((u.int32 >> 16) & 1) : 0; // Round to nearest, round to even
+    u.int32 |= flag1 ? 0x10000 : 0x0;                      // Preserve signaling NaN
+
+    return uint16_t(u.int32 >> 16);
+}
+#else
 // convert fp32 to bfp16
 template <>
 inline __host__ __device__ constexpr bhalf_t type_convert<bhalf_t, float>(float x)
@@ -43,6 +88,7 @@ inline __host__ __device__ constexpr bhalf_t type_convert<bhalf_t, float>(float
 
     return uint16_t(u.int32 >> 16);
 }
+#endif
 
 // convert bfp16 to fp16 via fp32
 template <>

library/include/ck/library/utility/host_common_util.hpp

@@ -22,7 +22,7 @@ static inline void dumpBufferToFile(const char* fileName, T* data, size_t dataNu
     std::ofstream outFile(fileName, std::ios::binary);
     if(outFile)
     {
-        outFile.write(reinterpret_cast<char*>(data), dataNumItems * sizeof(T));
+        outFile.write(reinterpret_cast<const char*>(data), dataNumItems * sizeof(T));
         outFile.close();
         std::cout << "Write output to file " << fileName << std::endl;
     }

library/include/ck/library/utility/host_tensor_generator.hpp

@@ -130,10 +130,11 @@ struct GeneratorTensor_3<ck::bhalf_t>
 template <typename T>
 struct GeneratorTensor_4
 {
-    std::default_random_engine generator;
+    std::mt19937 generator;
     std::normal_distribution<float> distribution;
 
-    GeneratorTensor_4(float mean, float stddev) : generator(1), distribution(mean, stddev){};
+    GeneratorTensor_4(float mean, float stddev, unsigned int seed = 1)
+        : generator(seed), distribution(mean, stddev){};
 
     template <typename... Is>
     T operator()(Is...)