Merge branch 'attn-bwd-develop' into attn-bwd-dropout-pt1

example/32_batched_gemm_scale_softmax_gemm/run_batched_multihead_attention_forward.inc

@@ -5,7 +5,7 @@ int run(int argc, char* argv[])
 {
     bool do_verification = true;
     int init_method      = 1;
-    bool time_kernel     = false;
+    bool time_kernel     = true;
 
     // GEMM shape for A/B0/B1/C
     // C_g_m_o = A_g_m_k * B0_g_k_n * B1_g_n_o
@@ -175,7 +175,7 @@ int run(int argc, char* argv[])
         static_cast<B0DataType*>(b0_device_buf.GetDeviceBuffer()),
         static_cast<B1DataType*>(b1_device_buf.GetDeviceBuffer()),
         static_cast<CDataType*>(c_device_buf.GetDeviceBuffer()),
-        static_cast<ZDataType*>(z_device_buf.GetDeviceBuffer()),
+        static_cast<ZDataType*>(nullptr),
         static_cast<LSEDataType*>(lse_device_buf.GetDeviceBuffer()),
         {}, // std::array<void*, 1> p_acc0_biases;
         {}, // std::array<void*, 1> p_acc1_biases;
@@ -228,6 +228,44 @@ int run(int argc, char* argv[])
 
     if(do_verification)
     {
+
+        // run for storing z tensor
+        argument = gemm.MakeArgument(
+            static_cast<ADataType*>(a_device_buf.GetDeviceBuffer()),
+            static_cast<B0DataType*>(b0_device_buf.GetDeviceBuffer()),
+            static_cast<B1DataType*>(b1_device_buf.GetDeviceBuffer()),
+            static_cast<CDataType*>(c_device_buf.GetDeviceBuffer()),
+            static_cast<ZDataType*>(z_device_buf.GetDeviceBuffer()),
+            static_cast<LSEDataType*>(lse_device_buf.GetDeviceBuffer()),
+            {}, // std::array<void*, 1> p_acc0_biases;
+            {}, // std::array<void*, 1> p_acc1_biases;
+            a_gs_ms_ks_lengths,
+            a_gs_ms_ks_strides,
+            b0_gs_ns_ks_lengths,
+            b0_gs_ns_ks_strides,
+            b1_gs_os_ns_lengths,
+            b1_gs_os_ns_strides,
+            c_gs_ms_os_lengths,
+            c_gs_ms_os_strides,
+            z_gs_ms_ns_lengths,
+            z_gs_ms_ns_strides,
+            lse_gs_ms_lengths,
+            {}, // std::array<std::vector<ck::index_t>, 1>{acc0_biases_gs_ms_ns_lengths},
+            {}, // std::array<std::vector<ck::index_t>, 1>{acc0_biases_gs_ms_ns_strides},
+            {}, // std::array<std::vector<ck::index_t>, 1>{acc1_biases_gs_ms_os_lengths},
+            {}, // std::array<std::vector<ck::index_t>, 1>{acc1_biases_gs_ms_os_strides},
+            a_element_op,
+            b0_element_op,
+            acc0_element_op,
+            b1_element_op,
+            c_element_op,
+            p_drop,          // dropout ratio
+            {seed, offset}); // dropout random seed and offset, offset should be at least the number
+                             // of elements on a thread
+        c_device_buf.SetZero();
+        lse_device_buf.SetZero();
+        invoker.Run(argument, StreamConfig{nullptr, false});
+
         c_device_buf.FromDevice(c_gs_ms_os_device_result.mData.data());
         z_device_buf.FromDevice(z_gs_ms_ns.mData.data());
         lse_device_buf.FromDevice(lse_gs_ms_device_result.mData.data());

example/32_batched_gemm_scale_softmax_gemm/run_grouped_multihead_attention_forward.inc

@@ -5,7 +5,7 @@ int run(int argc, char* argv[])
 {
     bool do_verification = true;
     int init_method      = 1;
-    bool time_kernel     = false;
+    bool time_kernel     = true;
 
     bool input_permute  = false;
     bool output_permute = true;
@@ -56,7 +56,8 @@ 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<void*> p_z;
+    std::vector<void*> p_z;         // for result verification
+    std::vector<void*> p_z_nullptr; // for time test
     std::vector<void*> p_lse;
     std::vector<std::vector<int>> g0_g1_m_n_k_o;
 
@@ -221,6 +222,7 @@ int run(int argc, char* argv[])
         p_b1.push_back(b1_tensors_device[i]->GetDeviceBuffer());
         p_c.push_back(c_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());
     }
 
@@ -233,12 +235,13 @@ int run(int argc, char* argv[])
     // do GEMM
     auto gemm    = DeviceGemmInstance{};
     auto invoker = gemm.MakeInvoker();
+
     auto argument =
         gemm.MakeArgument(p_a,
                           p_b0,
                           p_b1,
                           p_c,
-                          p_z,
+                          p_z_nullptr,
                           p_lse,
                           {}, // p_acc0_biases
                           {}, // p_acc1_biases
@@ -252,7 +255,6 @@ int run(int argc, char* argv[])
                           {seed, offset}); // dropout random seed and offset, offset should be
                                            // at least the number of elements on a thread
 
-
     // specify workspace for problem_desc
     DeviceMem problem_desc_workspace(gemm.GetWorkSpaceSize(&argument));
 
@@ -277,6 +279,31 @@ int run(int argc, char* argv[])
     bool pass = true;
     if(do_verification)
     {
+        argument =
+            gemm.MakeArgument(p_a,
+                              p_b0,
+                              p_b1,
+                              p_c,
+                              p_z,
+                              p_lse,
+                              {}, // p_acc0_biases
+                              {}, // p_acc1_biases
+                              problem_descs,
+                              a_element_op,
+                              b0_element_op,
+                              acc0_element_op,
+                              b1_element_op,
+                              c_element_op,
+                              p_drop,          // dropout ratio
+                              {seed, offset}); // dropout random seed and offset, offset should be
+                                               // at least the number of elements on a thread
+        // specify workspace for problem_desc
+        DeviceMem problem_desc_workspace_verify(gemm.GetWorkSpaceSize(&argument));
+
+        gemm.SetWorkSpacePointer(&argument, problem_desc_workspace_verify.GetDeviceBuffer());
+
+        invoker.Run(argument, StreamConfig{nullptr, false});
+
         for(std::size_t i = 0; i < group_count; i++)
         {
             const int& G0 = g0_g1_m_n_k_o[i][0];

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

@@ -44,7 +44,8 @@ template <typename GridwiseGemm,
           typename ComputeBasePtrOfStridedBatch,
           typename C0MatrixMask,
           bool HasMainKBlockLoop,
-          bool IsDropout>
+          bool IsDropout,
+          bool IsLseStoring>
 __global__ void
 #if CK_USE_LAUNCH_BOUNDS
     __launch_bounds__(CK_MAX_THREAD_PER_BLOCK, CK_MIN_BLOCK_PER_CU)
@@ -100,13 +101,13 @@ __global__ void
     const index_t global_thread_id = get_thread_global_1d_id();
     ck::philox ph(seed, global_thread_id, offset);
 
-    GridwiseGemm::template Run<HasMainKBlockLoop, IsDropout>(
+    GridwiseGemm::template Run<HasMainKBlockLoop, IsDropout, IsLseStoring>(
         p_a_grid + a_batch_offset,
         p_b_grid + b_batch_offset,
         p_b1_grid + b1_batch_offset,
         p_c_grid + c_batch_offset,
         nullptr ? nullptr : p_z_grid + z_batch_offset,
-        p_lse_grid + lse_batch_offset,
+        nullptr ? nullptr : p_lse_grid + lse_batch_offset,
         p_shared,
         a_element_op,
         b_element_op,
@@ -596,6 +597,12 @@ struct DeviceBatchedMultiheadAttentionForward_Xdl_CShuffle
 
             z_grid_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5_ =
                 GridwiseGemm::MakeCGridDescriptor_M0_N0_M1_N1_M2_N2_M3_N3_N4_N5(z_grid_desc_m_n_);
+
+            if(p_lse_grid == nullptr)
+            {
+                is_lse_storing_ = false;
+            }
+
         }
 
         void Print() const
@@ -669,6 +676,8 @@ struct DeviceBatchedMultiheadAttentionForward_Xdl_CShuffle
         unsigned long long seed_;
         unsigned long long offset_;
         bool is_dropout_;
+
+        bool is_lse_storing_ = true;
     };
 
     // Invoker
@@ -692,7 +701,9 @@ struct DeviceBatchedMultiheadAttentionForward_Xdl_CShuffle
 
             float ave_time = 0;
 
-            auto launch_kernel = [&](auto has_main_k_block_loop_, auto is_dropout_) {
+            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<
                     GridwiseGemm,
                     ADataType, // TODO: distiguish A/B datatype
@@ -715,7 +726,8 @@ struct DeviceBatchedMultiheadAttentionForward_Xdl_CShuffle
                     ComputeBasePtrOfStridedBatch,
                     C0MatrixMask,
                     has_main_k_block_loop_,
-                    is_dropout_>;
+                    is_dropout_,
+                    is_lse_storing_>;
 
                 return launch_and_time_kernel(stream_config,
                                               kernel,
@@ -755,26 +767,66 @@ struct DeviceBatchedMultiheadAttentionForward_Xdl_CShuffle
             {
                 if(arg.is_dropout_)
                 {
-                    ave_time = launch_kernel(integral_constant<bool, true>{},
-                                             integral_constant<bool, true>{});
+                    if(arg.is_lse_storing_)
+                    {
+                        ave_time = launch_kernel(integral_constant<bool, true>{},
+                                                 integral_constant<bool, true>{},
+                                                 integral_constant<bool, true>{});
+                    }
+                    else
+                    {
+                        ave_time = launch_kernel(integral_constant<bool, true>{},
+                                                 integral_constant<bool, true>{},
+                                                 integral_constant<bool, false>{});
+                    }
                 }
                 else
                 {
-                    ave_time = launch_kernel(integral_constant<bool, true>{},
-                                             integral_constant<bool, false>{});
+                    if(arg.is_lse_storing_)
+                    {
+                        ave_time = launch_kernel(integral_constant<bool, true>{},
+                                                 integral_constant<bool, false>{},
+                                                 integral_constant<bool, true>{});
+                    }
+                    else
+                    {
+                        ave_time = launch_kernel(integral_constant<bool, true>{},
+                                                 integral_constant<bool, false>{},
+                                                 integral_constant<bool, false>{});
+                    }
                 }
             }
             else
             {
                 if(arg.is_dropout_)
                 {
-                    ave_time = launch_kernel(integral_constant<bool, false>{},
-                                             integral_constant<bool, true>{});
+                    if(arg.is_lse_storing_)
+                    {
+                        ave_time = launch_kernel(integral_constant<bool, false>{},
+                                                 integral_constant<bool, true>{},
+                                                 integral_constant<bool, true>{});
+                    }
+                    else
+                    {
+                        ave_time = launch_kernel(integral_constant<bool, false>{},
+                                                 integral_constant<bool, true>{},
+                                                 integral_constant<bool, false>{});
+                    }
                 }
                 else
                 {
-                    ave_time = launch_kernel(integral_constant<bool, false>{},
-                                             integral_constant<bool, false>{});
+                    if(arg.is_lse_storing_)
+                    {
+                        ave_time = launch_kernel(integral_constant<bool, false>{},
+                                                 integral_constant<bool, false>{},
+                                                 integral_constant<bool, true>{});
+                    }
+                    else
+                    {
+                        ave_time = launch_kernel(integral_constant<bool, false>{},
+                                                 integral_constant<bool, false>{},
+                                                 integral_constant<bool, false>{});
+                    }
                 }
             }
 

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

@@ -273,11 +273,11 @@ struct GridwiseBatchedMultiheadAttentionForward_Xdl_CShuffle
         const auto K = a_grid_desc_ak0_m_ak1.GetLength(I0) * a_grid_desc_ak0_m_ak1.GetLength(I2);
         const auto Gemm1N = b1_grid_desc_bk0_n_bk1.GetLength(I1);
 
-        if(Gemm1N != K)
-        {
-            std::cout << "SizeK must be equal to SizeO (equal attention head size)" << '\n';
-            return false;
-        }
+        // if(Gemm1N != K)
+        //{
+        //    std::cout << "SizeK must be equal to SizeO (equal attention head size)" << '\n';
+        //    return false;
+        //}
 
         if(!(M == c_grid_desc_m_n.GetLength(I0) && Gemm1N == c_grid_desc_m_n.GetLength(I1)))
         {