grouped gemm remove multiple D

example/52_flash_atten_bias/grouped_multihead_attention_bias_forward_v2.cpp

@@ -53,8 +53,8 @@ using CShuffleDataType = F32;
 using CDataType        = DataType;
 using ZDataType        = U16; // INT32
 using LSEDataType      = F32;
-using Acc0BiasDataType = ck::Tuple<DDataType>;
-using Acc1BiasDataType = ck::Tuple<>;
+using Acc0BiasDataType = DDataType;
+using Acc1BiasDataType = void;
 
 static constexpr ck::index_t NumDimG = 2;
 static constexpr ck::index_t NumDimM = 1;

example/52_flash_atten_bias/run_batched_multihead_attention_bias_forward.inc

@@ -137,7 +137,7 @@ int run(int argc, char* argv[])
         a_gs_ms_ks.GenerateTensorValue(GeneratorTensor_2<ADataType>{-2, 2});
         b0_gs_ns_ks.GenerateTensorValue(GeneratorTensor_2<B0DataType>{-2, 2});
         b1_gs_os_ns.GenerateTensorValue(GeneratorTensor_2<B1DataType>{-2, 2});
-        d_gs_ms_ns.GenerateTensorValue(GeneratorTensor_2<DDataType>{-2, 2});
+        d_gs_ms_ns.GenerateTensorValue(GeneratorTensor_2<DDataType>{-1, 1});
         break;
     case 2:
         a_gs_ms_ks.GenerateTensorValue(GeneratorTensor_3<ADataType>{0.0, 1.0});

example/52_flash_atten_bias/run_grouped_multihead_attention_bias_forward.inc

@@ -57,7 +57,7 @@ int run(int argc, char* argv[])
     std::vector<const void*> p_b0;
     std::vector<const void*> p_b1;
     std::vector<void*> p_c;
-    std::vector<std::vector<const void*>> p_d;
+    std::vector<const void*> p_d;
     std::vector<void*> p_z;         // for result verification
     std::vector<void*> p_z_nullptr; // for time test
     std::vector<void*> p_lse;
@@ -147,10 +147,8 @@ int run(int argc, char* argv[])
                                  z_gs_ms_ns_strides,
                                  lse_gs_ms_lengths,
                                  lse_gs_ms_strides,
-                                 std::vector<std::vector<ck::index_t>>{
-                                     d_gs_ms_ns_lengths}, // acc0_biases_gs_ms_ns_lengths
-                                 std::vector<std::vector<ck::index_t>>{
-                                     d_gs_ms_ns_strides}, // acc0_biases_gs_ms_ns_strides
+                                 d_gs_ms_ns_lengths, // acc0_biases_gs_ms_ns_lengths
+                                 d_gs_ms_ns_strides, // acc0_biases_gs_ms_ns_strides
                                  {},                      // acc1_biases_gs_ms_os_lengths
                                  {}});                    // acc1_biases_gs_ms_os_strides
 
@@ -167,7 +165,7 @@ int run(int argc, char* argv[])
         flop += (size_t(M) * N * K * 2 + size_t(M) * N * O * 2) * Batch;
         num_byte += (sizeof(ADataType) * M * K + sizeof(B0DataType) * K * N +
                      sizeof(B1DataType) * N * O + sizeof(CDataType) * M * O +
-                     sizeof(DDataType) * M * N * (Acc0BiasDataType::Size() ? 0 : 1)) *
+                     sizeof(DDataType) * M * N * (std::is_void<Acc0BiasDataType>::value ? 0 : 1)) *
                     Batch;
 
         if(i < 4)
@@ -244,9 +242,7 @@ int run(int argc, char* argv[])
         p_b0.push_back(b0_tensors_device[i]->GetDeviceBuffer());
         p_b1.push_back(b1_tensors_device[i]->GetDeviceBuffer());
         p_c.push_back(c_tensors_device[i]->GetDeviceBuffer());
-        p_d.push_back({d_tensors_device[i]->GetDeviceBuffer()});
-        // std::cout << "from host  group id: " << i << " d address: " <<
-        // d_tensors_device[i]->GetDeviceBuffer() << std::endl;
+        p_d.push_back(d_tensors_device[i]->GetDeviceBuffer());
         p_z.push_back(z_tensors_device[i]->GetDeviceBuffer());
         p_z_nullptr.push_back(nullptr);
         p_lse.push_back(lse_tensors_device[i]->GetDeviceBuffer());

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

@@ -111,11 +111,11 @@ struct DeviceGroupedMultiheadAttentionForward : public BaseOperator
         std::vector<index_t> lse_gs_ms_lengths;
         std::vector<index_t> lse_gs_ms_strides;
 
-        std::vector<std::vector<index_t>> acc0_biases_gs_ms_ns_lengths;
-        std::vector<std::vector<index_t>> acc0_biases_gs_ms_ns_strides;
+        std::vector<index_t> acc0_biases_gs_ms_ns_lengths;
+        std::vector<index_t> acc0_biases_gs_ms_ns_strides;
 
-        std::vector<std::vector<index_t>> acc1_biases_gs_ms_os_lengths;
-        std::vector<std::vector<index_t>> acc1_biases_gs_ms_os_strides;
+        std::vector<index_t> acc1_biases_gs_ms_os_lengths;
+        std::vector<index_t> acc1_biases_gs_ms_os_strides;
     };
 
     virtual std::unique_ptr<BaseArgument>
@@ -125,9 +125,9 @@ struct DeviceGroupedMultiheadAttentionForward : public BaseOperator
                         std::vector<void*> p_c_vec,
                         std::vector<void*> p_z_vec,
                         std::vector<void*> p_lse_vec,
-                        std::vector<std::vector<const void*>> p_acc0_biases_vec,
-                        std::vector<std::vector<const void*>> p_acc1_biases_vec,
-                        std::vector<ProblemDesc> problem_desc_vec,
+                        std::vector<const void*> p_acc0_biases_vec,
+                        std::vector<const void*> p_acc1_biases_vec,
+                        std::vector<ProblemDesc>& problem_desc_vec,
                         AElementwiseOperation a_element_op,
                         B0ElementwiseOperation b0_element_op,
                         Acc0ElementwiseOperation acc0_element_op,

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

@@ -685,11 +685,6 @@ struct DeviceBatchedMultiheadAttentionForward_Xdl_CShuffle_V2
                 c_grid_desc_mblock_mperblock_nblock_nperblock_ =
                     GridwiseGemm::MakeCGridDescriptor_MBlock_MPerBlock_NBlock_NPerBlock(
                         c_grid_desc_m_n_);
-
-                D0GridDesc_M_N d0_grid_desc_m_n{};
-                d0_grid_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5_ =
-                    GridwiseGemm::MakeD0GridDescriptor_M0_N0_M1_N1_M2_N2_M3_N3_N4_N5(
-                        d0_grid_desc_m_n);
             }
 
             d0_n_length_stride_.push_back(acc0_biases_gs_ms_ns_lengths[NumDimG + NumDimM]);
@@ -724,6 +719,13 @@ struct DeviceBatchedMultiheadAttentionForward_Xdl_CShuffle_V2
             std::cout << "b_grid_desc_g_n_k_: " << b_grid_desc_g_n_k_.GetLength(I0) << ", "
                       << b_grid_desc_g_n_k_.GetLength(I1) << ", "
                       << b_grid_desc_g_n_k_.GetLength(I2) << '\n';
+            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';

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

@@ -99,14 +99,9 @@ __global__ void
         static_cast<long_index_t>(arg_ptr[group_id].compute_base_ptr_of_batch_.GetZBasePtr(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 d0_batch_offset  = __builtin_amdgcn_readfirstlane(static_cast<long_index_t>(
+        arg_ptr[group_id].compute_base_ptr_of_batch_.GetD0BasePtr(g_idx)));
 
-    typename GridwiseGemm::D0sGridPointer p_d0s_grid = arg_ptr[group_id].p_d0s_grid_;
-    static_for<0, p_d0s_grid.Size(), 1>{}([&](auto In) {
-        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, In)));
-        p_d0s_grid(In) = p_d0s_grid(In) + d0_batch_offset;
-    });
     if constexpr(Deterministic)
     {
         for(index_t i = 0; i < num_blocks_per_batch; i++)
@@ -114,7 +109,9 @@ __global__ void
             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,
-                p_d0s_grid,
+                arg_ptr[group_id].p_d0_grid_ == nullptr
+                    ? nullptr
+                    : arg_ptr[group_id].p_d0_grid_ + d0_batch_offset,
                 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
@@ -132,7 +129,7 @@ __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].d0s_grid_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5_,
+                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_m4_n4_n5_n6_,
@@ -153,7 +150,9 @@ __global__ void
         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,
-            p_d0s_grid,
+            arg_ptr[group_id].p_d0_grid_ == nullptr
+                ? nullptr
+                : arg_ptr[group_id].p_d0_grid_ + d0_batch_offset,
             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
@@ -170,7 +169,7 @@ __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].d0s_grid_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5_,
+            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_m4_n4_n5_n6_,
@@ -299,11 +298,10 @@ struct DeviceGroupedMultiheadAttentionForward_Xdl_CShuffle_V2
     static_assert(NumDimG > 0 && NumDimM > 0 && NumDimN > 0 && NumDimK > 0 && NumDimO > 0,
                   "Number of dimension must be greater than 0");
 
-    static constexpr index_t NumD0Tensor = Acc0BiasDataType::Size();
-    static constexpr index_t NumD1Tensor = Acc1BiasDataType::Size();
-
+    using D0DataType = Acc0BiasDataType;
+    using D1DataType = Acc1BiasDataType;
     // TODO ANT: implement bias combination
-    static_assert(NumD1Tensor == 0, "Acc1 Bias addition is unimplemented");
+    static_assert(std::is_void<Acc1BiasDataType>::value, "Acc1 Bias addition is unimplemented");
 
 #if 0
     // TODO ANT: use alias
@@ -406,44 +404,38 @@ struct DeviceGroupedMultiheadAttentionForward_Xdl_CShuffle_V2
         }
     }
 
-    static auto MakeD0sGridDescriptor_M_N(
-        const std::vector<std::vector<ck::index_t>>& acc0_biases_gs_ms_ns_lengths,
-        const std::vector<std::vector<ck::index_t>>& acc0_biases_gs_ms_ns_strides)
+    static auto
+    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 generate_tuple(
-            [&](auto i) {
-                return Transform::MakeCGridDescriptor_M_N(acc0_biases_gs_ms_ns_lengths[i],
-                                                          acc0_biases_gs_ms_ns_strides[i]);
-            },
-            Number<NumD0Tensor>{});
+
+        return Transform::MakeCGridDescriptor_M_N(acc0_biases_gs_ms_ns_lengths,
+                                                  acc0_biases_gs_ms_ns_strides);
     }
 
-    static auto MakeD0sGridDescriptor_G_M_N(
-        const std::vector<std::vector<ck::index_t>>& acc0_biases_gs_ms_ns_lengths,
-        const std::vector<std::vector<ck::index_t>>& acc0_biases_gs_ms_ns_strides)
+    static auto
+    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 generate_tuple(
-            [&](auto i) {
-                return Transform::MakeCGridDescriptor_G_M_N(acc0_biases_gs_ms_ns_lengths[i],
-                                                            acc0_biases_gs_ms_ns_strides[i]);
-            },
-            Number<NumD0Tensor>{});
+
+        return Transform::MakeCGridDescriptor_G_M_N(acc0_biases_gs_ms_ns_lengths,
+                                                    acc0_biases_gs_ms_ns_strides);
     }
 
     using AGridDesc_AK0_M_AK1  = decltype(MakeAGridDescriptor_AK0_M_AK1({}, {}));
     using BGridDesc_BK0_N_BK1  = decltype(MakeBGridDescriptor_BK0_N_BK1({}, {}));
-    using D0sGridDesc_M_N      = decltype(MakeD0sGridDescriptor_M_N({}, {}));
+    using D0GridDesc_M_N       = decltype(MakeD0GridDescriptor_M_N({}, {}));
     using B1GridDesc_BK0_N_BK1 = decltype(MakeB1GridDescriptor_BK0_N_BK1({}, {}));
     using CGridDesc_M_N        = decltype(Transform::MakeCGridDescriptor_M_N({}, {}));
     using LSEGridDesc_M        = decltype(MakeLSEGridDescriptor_M(1));
     using ZGridDesc_M_N        = decltype(MakeZGridDescriptor_M_N({}, {}));
 
-    using AGridDesc_G_M_K   = decltype(Transform::MakeAGridDescriptor_G_M_K({}, {}));
-    using BGridDesc_G_N_K   = decltype(Transform::MakeB0GridDescriptor_G_N_K({}, {}));
-    using D0sGridDesc_G_M_N = decltype(MakeD0sGridDescriptor_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 AGridDesc_G_M_K  = decltype(Transform::MakeAGridDescriptor_G_M_K({}, {}));
+    using BGridDesc_G_N_K  = decltype(Transform::MakeB0GridDescriptor_G_N_K({}, {}));
+    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({}, {}));
 
     constexpr static auto make_MaskOutPredicate()
     {
@@ -466,14 +458,14 @@ struct DeviceGroupedMultiheadAttentionForward_Xdl_CShuffle_V2
     {
         ComputeBasePtrOfStridedBatch(const AGridDesc_G_M_K& a_grid_desc_g_m_k,
                                      const BGridDesc_G_N_K& b_grid_desc_g_n_k,
-                                     const D0sGridDesc_G_M_N& d0s_grid_desc_g_m_n,
+                                     const D0GridDesc_G_M_N& d0_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 ZGridDesc_G_M_N& z_grid_desc_g_m_n,
                                      index_t BatchStrideLSE)
             : a_grid_desc_g_m_k_(a_grid_desc_g_m_k),
               b_grid_desc_g_n_k_(b_grid_desc_g_n_k),
-              d0s_grid_desc_g_m_n_(d0s_grid_desc_g_m_n),
+              d0_grid_desc_g_m_n_(d0_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),
               z_grid_desc_g_m_n_(z_grid_desc_g_m_n),
@@ -491,11 +483,9 @@ struct DeviceGroupedMultiheadAttentionForward_Xdl_CShuffle_V2
             return b_grid_desc_g_n_k_.CalculateOffset(make_multi_index(g_idx, 0, 0));
         }
 
-        template <index_t I>
-        __host__ __device__ constexpr long_index_t GetD0BasePtr(index_t g_idx,
-                                                                Number<I> d0_idx) const
+        __host__ __device__ constexpr long_index_t GetD0BasePtr(index_t g_idx) const
         {
-            return d0s_grid_desc_g_m_n_[d0_idx].CalculateOffset(make_multi_index(g_idx, 0, 0));
+            return d0_grid_desc_g_m_n_.CalculateOffset(make_multi_index(g_idx, 0, 0));
         }
 
         __host__ __device__ constexpr long_index_t GetB1BasePtr(index_t g_idx) const
@@ -521,7 +511,7 @@ struct DeviceGroupedMultiheadAttentionForward_Xdl_CShuffle_V2
         private:
         AGridDesc_G_M_K a_grid_desc_g_m_k_;
         BGridDesc_G_N_K b_grid_desc_g_n_k_;
-        D0sGridDesc_G_M_N d0s_grid_desc_g_m_n_;
+        D0GridDesc_G_M_N d0_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_;
         ZGridDesc_G_M_N z_grid_desc_g_m_n_;
@@ -547,7 +537,7 @@ struct DeviceGroupedMultiheadAttentionForward_Xdl_CShuffle_V2
         InMemoryDataOperationEnum::Set,
         AGridDesc_AK0_M_AK1,
         BGridDesc_BK0_N_BK1,
-        D0sGridDesc_M_N,
+        D0GridDesc_M_N,
         B1GridDesc_BK0_N_BK1,
         CGridDesc_M_N,
         ZGridDesc_M_N,
@@ -610,7 +600,7 @@ struct DeviceGroupedMultiheadAttentionForward_Xdl_CShuffle_V2
         // pointers
         const ADataType* p_a_grid_;
         const BDataType* p_b_grid_;
-        typename GridwiseGemm::D0sGridPointer p_d0s_grid_;
+        const D0DataType* p_d0_grid_;
         const B1DataType* p_b1_grid_;
         CDataType* p_c_grid_;
         ZDataType* p_z_grid_;
@@ -619,8 +609,8 @@ struct DeviceGroupedMultiheadAttentionForward_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_;
-        typename GridwiseGemm::D0sGridDescriptor_M0_N0_M1_N1_M2_N2_M3_N3_N4_N5
-            d0s_grid_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5_;
+        typename GridwiseGemm::D0GridDescriptor_M0_N0_M1_N1_M2_N2_M3_N3_N4_N5
+            d0_grid_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5_;
         B1GridDesc_BK0_N_BK1 b1_grid_desc_bk0_n_bk1_;
         typename GridwiseGemm::CGridDescriptor_MBlock_MPerBlock_NBlock_NPerBlock
             c_grid_desc_mblock_mperblock_nblock_nperblock_;
@@ -660,7 +650,7 @@ struct DeviceGroupedMultiheadAttentionForward_Xdl_CShuffle_V2
         CGridDesc_M_N c_grid_desc_m_n_;
 
         // raw data
-        std::array<std::vector<ck::index_t>, NumD0Tensor> d0s_n_length_stride_;
+        std::vector<ck::index_t> d0_n_length_stride_;
     };
 
     // Argument
@@ -673,9 +663,9 @@ struct DeviceGroupedMultiheadAttentionForward_Xdl_CShuffle_V2
                  std::vector<void*> p_c_vec,
                  std::vector<void*> p_z_vec,
                  std::vector<void*> p_lse_vec,
-                 std::vector<std::vector<const void*>> p_acc0_biases_vec,
-                 std::vector<std::vector<const void*>> p_acc1_biases_vec,
-                 std::vector<ProblemDesc> problem_desc_vec,
+                 std::vector<const void*> p_acc0_biases_vec,
+                 std::vector<const void*> p_acc1_biases_vec,
+                 std::vector<ProblemDesc>& problem_desc_vec,
                  AElementwiseOperation a_element_op,
                  BElementwiseOperation b_element_op,
                  AccElementwiseOperation acc_element_op,
@@ -706,23 +696,11 @@ struct DeviceGroupedMultiheadAttentionForward_Xdl_CShuffle_V2
 
             for(std::size_t i = 0; i < group_count_; i++)
             {
-                const auto p_a_grid = static_cast<const ADataType*>(p_a_vec[i]);
-                const auto p_b_grid = static_cast<const BDataType*>(p_b_vec[i]);
-
-                const auto& problem_desc = problem_desc_vec[i];
-                std::array<std::vector<ck::index_t>, NumD0Tensor> d0s_n_length_stride;
-                typename GridwiseGemm::D0sGridPointer p_d0s_grid;
-                static_for<0, NumD0Tensor, 1>{}([&](auto j) {
-                    using D0DataType = remove_cvref_t<tuple_element_t<j.value, Acc0BiasDataType>>;
-                    // D0 pointer
-                    p_d0s_grid(j) = static_cast<const D0DataType*>(p_acc0_biases_vec[i][j]);
-                    // for  check
-                    d0s_n_length_stride[j].push_back(
-                        problem_desc.acc0_biases_gs_ms_ns_lengths[j][NumDimG + NumDimM]);
-                    d0s_n_length_stride[j].push_back(
-                        problem_desc.acc0_biases_gs_ms_ns_strides[j][NumDimG + NumDimM]);
-                });
-
+                const auto p_a_grid  = static_cast<const ADataType*>(p_a_vec[i]);
+                const auto p_b_grid  = static_cast<const BDataType*>(p_b_vec[i]);
+                const auto p_d0_grid = p_acc0_biases_vec.size()
+                                           ? static_cast<const D0DataType*>(p_acc0_biases_vec[i])
+                                           : nullptr;
                 const auto p_b1_grid  = static_cast<const B1DataType*>(p_b1_vec[i]);
                 const auto p_c_grid   = static_cast<CDataType*>(p_c_vec[i]);
                 const auto p_z_grid   = static_cast<ZDataType*>(p_z_vec[i]);
@@ -733,16 +711,17 @@ struct DeviceGroupedMultiheadAttentionForward_Xdl_CShuffle_V2
                     is_lse_storing_ = false;
                 }
 
+                const auto& problem_desc         = problem_desc_vec[i];
                 const auto a_grid_desc_ak0_m_ak1 = MakeAGridDescriptor_AK0_M_AK1(
                     problem_desc.a_gs_ms_ks_lengths, problem_desc.a_gs_ms_ks_strides);
                 const auto b_grid_desc_bk0_n_bk1 = MakeBGridDescriptor_BK0_N_BK1(
                     problem_desc.b0_gs_ns_ks_lengths, problem_desc.b0_gs_ns_ks_strides);
-                const D0sGridDesc_M_N d0s_grid_desc_m_n{
-                    DeviceOp::MakeD0sGridDescriptor_M_N(problem_desc.acc0_biases_gs_ms_ns_lengths,
-                                                        problem_desc.acc0_biases_gs_ms_ns_strides)};
-                const auto d0s_grid_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5 =
-                    GridwiseGemm::MakeD0sGridDescriptor_M0_N0_M1_N1_M2_N2_M3_N3_N4_N5(
-                        d0s_grid_desc_m_n);
+                const D0GridDesc_M_N d0_grid_desc_m_n{
+                    DeviceOp::MakeD0GridDescriptor_M_N(problem_desc.acc0_biases_gs_ms_ns_lengths,
+                                                       problem_desc.acc0_biases_gs_ms_ns_strides)};
+                const auto d0_grid_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5 =
+                    GridwiseGemm::MakeD0GridDescriptor_M0_N0_M1_N1_M2_N2_M3_N3_N4_N5(
+                        d0_grid_desc_m_n);
                 const auto b1_grid_desc_bk0_n_bk1 = MakeB1GridDescriptor_BK0_N_BK1(
                     problem_desc.b1_gs_os_ns_lengths, problem_desc.b1_gs_os_ns_strides);
                 const auto c_grid_desc_m_n = Transform::MakeCGridDescriptor_M_N(
@@ -756,9 +735,9 @@ struct DeviceGroupedMultiheadAttentionForward_Xdl_CShuffle_V2
                     problem_desc.a_gs_ms_ks_lengths, problem_desc.a_gs_ms_ks_strides);
                 const auto b_grid_desc_g_n_k = Transform::MakeB0GridDescriptor_G_N_K(
                     problem_desc.b0_gs_ns_ks_lengths, problem_desc.b0_gs_ns_ks_strides);
-                const auto d0s_grid_desc_g_m_n = DeviceOp::MakeD0sGridDescriptor_G_M_N(
-                    problem_desc.acc0_biases_gs_ms_ns_lengths,
-                    problem_desc.acc0_biases_gs_ms_ns_strides);
+                const auto d0_grid_desc_g_m_n =
+                    DeviceOp::MakeD0GridDescriptor_G_M_N(problem_desc.acc0_biases_gs_ms_ns_lengths,
+                                                         problem_desc.acc0_biases_gs_ms_ns_strides);
                 const auto b1_grid_desc_g_n_k = Transform::MakeB1GridDescriptor_G_N_K(
                     problem_desc.b1_gs_os_ns_lengths, problem_desc.b1_gs_os_ns_strides);
                 const auto c_grid_desc_g_m_n = Transform::MakeCGridDescriptor_G_M_N(
@@ -786,7 +765,7 @@ struct DeviceGroupedMultiheadAttentionForward_Xdl_CShuffle_V2
                 const auto compute_base_ptr_of_batch = ComputeBasePtrOfStridedBatch(
                     a_grid_desc_g_m_k,
                     b_grid_desc_g_n_k,
-                    d0s_grid_desc_g_m_n,
+                    d0_grid_desc_g_m_n,
                     b1_grid_desc_g_n_k,
                     c_grid_desc_g_m_n,
                     z_grid_desc_g_m_n,
@@ -798,18 +777,6 @@ struct DeviceGroupedMultiheadAttentionForward_Xdl_CShuffle_V2
 
                 grid_size_ += grid_size_grp;
 
-                // for each group, make sure acc0_biases_gs_ms_ns_lengths.size() == NumD0Tensor and
-                // so on
-                if(!(problem_desc.acc0_biases_gs_ms_ns_lengths.size() == NumD0Tensor &&
-                     problem_desc.acc0_biases_gs_ms_ns_strides.size() == NumD0Tensor &&
-                     problem_desc.acc1_biases_gs_ms_os_lengths.size() == NumD1Tensor &&
-                     problem_desc.acc1_biases_gs_ms_os_strides.size() == NumD1Tensor))
-                {
-                    throw std::runtime_error(
-                        "wrong! number of biases in function argument does not "
-                        "match that in template argument");
-                }
-
                 const auto raw_m_padded = GridwiseGemm::GetPaddedSize(
                     problem_desc.a_gs_ms_ks_lengths[NumDimG + NumDimM - 1]);
                 const auto raw_n_padded = GridwiseGemm::GetPaddedSize(
@@ -817,14 +784,14 @@ struct DeviceGroupedMultiheadAttentionForward_Xdl_CShuffle_V2
 
                 group_kernel_args_.push_back({p_a_grid,
                                               p_b_grid,
-                                              p_d0s_grid,
+                                              p_d0_grid,
                                               p_b1_grid,
                                               p_c_grid,
                                               p_z_grid,
                                               p_lse_grid,
                                               a_grid_desc_ak0_m_ak1,
                                               b_grid_desc_bk0_n_bk1,
-                                              d0s_grid_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5,
+                                              d0_grid_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5,
                                               b1_grid_desc_bk0_n_bk1,
                                               c_grid_desc_mblock_mperblock_nblock_nperblock,
                                               z_grid_desc_m0_n0_m1_n1_m2_n2_m3_n3_m4_n4_n5_n6,
@@ -843,6 +810,13 @@ struct DeviceGroupedMultiheadAttentionForward_Xdl_CShuffle_V2
                 z_random_matrix_offset =
                     z_random_matrix_offset + raw_m_padded * raw_n_padded * batch_count;
 
+                // for  check
+                std::vector<ck::index_t> d0_n_length_stride;
+                d0_n_length_stride.push_back(
+                    problem_desc.acc0_biases_gs_ms_ns_lengths[NumDimG + NumDimM]);
+                d0_n_length_stride.push_back(
+                    problem_desc.acc0_biases_gs_ms_ns_strides[NumDimG + NumDimM]);
+
                 group_device_args_.push_back(
                     {{problem_desc.a_gs_ms_ks_lengths[NumDimG + NumDimM - 1],
                       problem_desc.b0_gs_ns_ks_lengths[NumDimG + NumDimN - 1],
@@ -857,7 +831,7 @@ 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,
-                     d0s_n_length_stride});
+                     d0_n_length_stride});
             }
 
             is_dropout_          = p_dropout > 0.0; //
@@ -1077,18 +1051,14 @@ struct DeviceGroupedMultiheadAttentionForward_Xdl_CShuffle_V2
                 return false;
             }
 
-            for(int In = 0; In < NumD0Tensor; In++)
+            if(device_arg.d0_n_length_stride_[1] == 1 &&
+               device_arg.d0_n_length_stride_[0] % Acc0BiasTransferSrcScalarPerVector != 0)
             {
-                if(device_arg.d0s_n_length_stride_[In][1] == 1 &&
-                   device_arg.d0s_n_length_stride_[In][0] % Acc0BiasTransferSrcScalarPerVector != 0)
-                {
-                    return false;
-                }
-                if(device_arg.d0s_n_length_stride_[In][1] != 1 &&
-                   Acc0BiasTransferSrcScalarPerVector != 1)
-                {
-                    return false;
-                }
+                return false;
+            }
+            if(device_arg.d0_n_length_stride_[1] != 1 && Acc0BiasTransferSrcScalarPerVector != 1)
+            {
+                return false;
             }
             // Check if having main loop
             const auto K = kernel_arg.a_grid_desc_ak0_m_ak1_.GetLength(I0) *
@@ -1170,9 +1140,9 @@ struct DeviceGroupedMultiheadAttentionForward_Xdl_CShuffle_V2
                              std::vector<void*> p_c_vec,
                              std::vector<void*> p_z_vec,
                              std::vector<void*> p_lse_vec,
-                             std::vector<std::vector<const void*>> p_acc0_biases_vec,
-                             std::vector<std::vector<const void*>> p_acc1_biases_vec,
-                             std::vector<ProblemDesc> problem_desc_vec,
+                             std::vector<const void*> p_acc0_biases_vec,
+                             std::vector<const void*> p_acc1_biases_vec,
+                             std::vector<ProblemDesc>& problem_desc_vec,
                              AElementwiseOperation a_element_op,
                              BElementwiseOperation b_element_op,
                              AccElementwiseOperation acc_element_op,
@@ -1209,9 +1179,9 @@ struct DeviceGroupedMultiheadAttentionForward_Xdl_CShuffle_V2
                         std::vector<void*> p_c_vec,
                         std::vector<void*> p_z_vec,
                         std::vector<void*> p_lse_vec,
-                        std::vector<std::vector<const void*>> p_acc0_biases_vec,
-                        std::vector<std::vector<const void*>> p_acc1_biases_vec,
-                        std::vector<ProblemDesc> problem_desc_vec,
+                        std::vector<const void*> p_acc0_biases_vec,
+                        std::vector<const void*> p_acc1_biases_vec,
+                        std::vector<ProblemDesc>& problem_desc_vec,
                         AElementwiseOperation a_element_op,
                         BElementwiseOperation b_element_op,
                         AccElementwiseOperation acc_element_op,

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

@@ -1291,7 +1291,7 @@ struct GridwiseBatchedMultiheadAttentionForward_Xdl_CShuffle_V2
             block_sync_lds(); // wait for lds read in gemm0 blockwise gemm
 
             // add bias
-            if constexpr(std::is_void<D0DataType>::value)
+            if constexpr(!std::is_void<D0DataType>::value)
             {
                 // get register
                 StaticBuffer<AddressSpaceEnum::Vgpr,