if z matrix ptr is null, will not save

example/32_batched_gemm_scale_softmax_gemm/batched_multihead_attention_backward_fp16_dropout.cpp

@@ -597,6 +597,45 @@ int run(int argc, char* argv[])
         });
         y_device_buf.ToDevice(y_gs_ms_os.mData.data());
 
+        //
+        // call kernel again
+        //
+        // example set Z matrix to null, will not ouput z matrix data
+        argument = gemm.MakeArgument(
+            static_cast<DataType*>(q_device_buf.GetDeviceBuffer()),
+            static_cast<DataType*>(k_device_buf.GetDeviceBuffer()),
+            static_cast<ZDataType*>(nullptr), // set to nullptr
+            static_cast<DataType*>(v_device_buf.GetDeviceBuffer()),
+            static_cast<DataType*>(y_device_buf.GetDeviceBuffer()),
+            static_cast<LSEDataType*>(lse_device_buf.GetDeviceBuffer()),
+            static_cast<DataType*>(ygrad_device_buf.GetDeviceBuffer()),
+            static_cast<DataType*>(qgrad_device_buf.GetDeviceBuffer()),
+            static_cast<DataType*>(kgrad_device_buf.GetDeviceBuffer()),
+            static_cast<DataType*>(vgrad_device_buf.GetDeviceBuffer()),
+            {}, // std::array<void*, 1> p_acc0_biases;
+            {}, // std::array<void*, 1> p_acc1_biases;
+            q_gs_ms_ks_lengths,
+            q_gs_ms_ks_strides,
+            k_gs_ns_ks_lengths,
+            k_gs_ns_ks_strides,
+            z_gs_ms_ns_lengths,
+            z_gs_ms_ns_strides,
+            v_gs_os_ns_lengths,
+            v_gs_os_ns_strides,
+            y_gs_ms_os_lengths,
+            y_gs_ms_os_strides,
+            lse_gs_ms_lengths,
+            {}, // std::array<std::vector<ck::index_t>, 1>{acc0_biases_gs_ms_ns_lengths},
+            {}, // std::array<std::vector<ck::index_t>, 1>{acc0_biases_gs_ms_ns_strides},
+            {}, // std::array<std::vector<ck::index_t>, 1>{acc1_biases_gs_ms_os_lengths},
+            {}, // std::array<std::vector<ck::index_t>, 1>{acc1_biases_gs_ms_os_strides},
+            QKVElementOp{},
+            QKVElementOp{},
+            Scale{alpha},
+            QKVElementOp{},
+            YElementOp{},
+            p_drop,
+            std::tuple<unsigned long long, unsigned long long>(seed, offset));
         kgrad_device_buf.SetZero(); // reset global accum buffer and rerun
         vgrad_device_buf.SetZero();
         invoker.Run(argument, StreamConfig{nullptr, false});

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

@@ -109,10 +109,11 @@ __global__ void
 
     const index_t global_thread_id = get_thread_global_1d_id();
     ck::philox ph(seed, global_thread_id, offset);
+    ZDataType* z_matrix_ptr = (p_z_grid == nullptr ? nullptr : p_z_grid + z_batch_offset);
 
     GridwiseGemm::template Run<HasMainKBlockLoop>(p_a_grid + a_batch_offset,
                                                   p_b_grid + b_batch_offset,
-                                                  p_z_grid + z_batch_offset,
+                                                  z_matrix_ptr,
                                                   p_b1_grid + b1_batch_offset,
                                                   p_c_grid + c_batch_offset,
                                                   p_lse_grid + lse_batch_offset,

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

@@ -1509,8 +1509,6 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
                                    0),
                   tensor_operation::element_wise::PassThrough{}};
 
-        ignore = z_thread_copy_vgpr_to_global;
-
         //
         // set up dV / dK Gemm (type 3 crr)
         //
@@ -1854,11 +1852,14 @@ struct GridwiseBatchedMultiheadAttentionBackward_Xdl_CShuffle_V2
                                                     true>(s_slash_p_thread_buf, ph, z_tenor_buffer);
 
             // save z to global
-            z_thread_copy_vgpr_to_global.Run(z_thread_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5,
-                                             make_tuple(I0, I0, I0, I0, I0, I0, I0, I0, I0, I0),
-                                             z_tenor_buffer,
-                                             z_grid_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5,
-                                             z_grid_buf);
+            if(p_z_grid)
+            {
+                z_thread_copy_vgpr_to_global.Run(z_thread_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5,
+                                                 make_tuple(I0, I0, I0, I0, I0, I0, I0, I0, I0, I0),
+                                                 z_tenor_buffer,
+                                                 z_grid_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5,
+                                                 z_grid_buf);
+            }
 
             block_sync_lds(); // wait for gemm1 LDS read