add d vector load template parameters

example/52_flash_atten_bias/grouped_multihead_attention_bias_forward_v2.cpp

@@ -136,6 +136,7 @@ using DeviceGemmInstance =
         8,
         8,
         true,
+        1,
         S<16, 16, 1>, // B1BlockTransfer
         S<0, 2, 1>,
         S<0, 2, 1>,
@@ -147,7 +148,8 @@ using DeviceGemmInstance =
         1,              // CShuffleNXdlPerWavePerShuffle
         S<1, 64, 1, 4>, // CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock
         8,              // CShuffleBlockTransferScalarPerVector_NPerBlock
-        MaskingSpec,    // MaskingSpecialization
+        1,
+        MaskingSpec, // MaskingSpecialization
         Deterministic>;
 #elif(DIM <= 64)
 using DeviceGemmInstance =
@@ -207,6 +209,7 @@ using DeviceGemmInstance =
         8,
         8,
         true,
+        1,
         S<16, 16, 1>, // B1BlockTransfer
         S<0, 2, 1>,
         S<0, 2, 1>,
@@ -218,7 +221,8 @@ using DeviceGemmInstance =
         2,              // CShuffleNXdlPerWavePerShuffle
         S<1, 32, 1, 8>, // CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock
         8,              // CShuffleBlockTransferScalarPerVector_NPerBlock
-        MaskingSpec,    // MaskingSpecialization
+        1,
+        MaskingSpec, // MaskingSpecialization
         Deterministic>;
 #elif(DIM <= 128)
 using DeviceGemmInstance =
@@ -278,6 +282,7 @@ using DeviceGemmInstance =
         8,
         8,
         true,
+        1,
         S<8, 32, 1>, // B1BlockTransfer
         S<0, 2, 1>,
         S<0, 2, 1>,
@@ -289,7 +294,8 @@ using DeviceGemmInstance =
         2,              // CShuffleNXdlPerWavePerShuffle
         S<1, 32, 1, 8>, // CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock
         8,              // CShuffleBlockTransferScalarPerVector_NPerBlock
-        MaskingSpec,    // MaskingSpecialization
+        1,
+        MaskingSpec, // MaskingSpecialization
         Deterministic>;
 #endif
 

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

@@ -256,6 +256,7 @@ template <index_t NumDimG,
           index_t BBlockTransferSrcScalarPerVector,
           index_t BBlockTransferDstScalarPerVector_BK1,
           bool BBlockLdsExtraN,
+          index_t Acc0BiasTransferSrcScalarPerVector,
           typename B1BlockTransferThreadClusterLengths_BK0_N_BK1,
           typename B1BlockTransferThreadClusterArrangeOrder,
           typename B1BlockTransferSrcAccessOrder,
@@ -267,6 +268,7 @@ template <index_t NumDimG,
           index_t CShuffleNXdlPerWavePerShuffle,
           typename CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock,
           index_t CShuffleBlockTransferScalarPerVector_NPerBlock,
+          index_t Acc1BiasTransferSrcScalarPerVector,
           MaskingSpecialization MaskingSpec,
           bool Deterministic,
           LoopScheduler LoopSched = LoopScheduler::Default>
@@ -561,6 +563,7 @@ struct DeviceGroupedMultiheadAttentionForward_Xdl_CShuffle_V2R2
         BBlockTransferSrcScalarPerVector,
         BBlockTransferDstScalarPerVector_BK1,
         true,
+        Acc0BiasTransferSrcScalarPerVector,
         BBlockLdsExtraN,
         B1BlockTransferThreadClusterLengths_BK0_N_BK1,
         B1BlockTransferThreadClusterArrangeOrder,
@@ -574,6 +577,7 @@ struct DeviceGroupedMultiheadAttentionForward_Xdl_CShuffle_V2R2
         CShuffleNXdlPerWavePerShuffle,
         CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock,
         CShuffleBlockTransferScalarPerVector_NPerBlock,
+        Acc1BiasTransferSrcScalarPerVector,
         LoopSched,
         Transform::matrix_padder.PadN,
         MaskingSpec != MaskingSpecialization::MaskDisabled,

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

@@ -74,6 +74,7 @@ template <typename FloatAB,
           index_t BBlockTransferDstScalarPerVector_BK1,
           bool BThreadTransferSrcResetCoordinateAfterRun, // ignored
           index_t BBlockLdsExtraN,
+          index_t D0BlockTransferSrcScalarPerVector,
           typename B1BlockTransferThreadClusterLengths_BK0_N_BK1,
           typename B1BlockTransferThreadClusterArrangeOrder,
           typename B1BlockTransferSrcAccessOrder,
@@ -86,6 +87,7 @@ template <typename FloatAB,
           index_t CShuffleNXdlPerWavePerShuffle,
           typename CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock,
           index_t CShuffleBlockTransferScalarPerVector_NPerBlock,
+          index_t D1BlockTransferSrcScalarPerVector,
           LoopScheduler LoopSched,
           bool PadN,
           bool MaskOutUpperTriangle,
@@ -930,7 +932,7 @@ struct GridwiseBatchedMultiheadAttentionForward_Xdl_CShuffle_V2R2
                                                       n4>,
                                              Sequence<0, 1, 2, 3, 4, 5, 6, 7, 8, 9>,
                                              9,
-                                             4,
+                                             D0BlockTransferSrcScalarPerVector,
                                              1,
                                              false>(d_grid_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5,
                                                     make_multi_index(block_work_idx[I0], // MBlockId