hipMemcpy -> hipMemcpyWithStream

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

@@ -651,11 +651,12 @@ struct DeviceGroupedContractionMultipleD_Xdl_CShuffle
                 }
             }
 
-            hipGetErrorString(hipMemcpy(arg.p_workspace_,
+            hipGetErrorString(hipMemcpyWithStream(arg.p_workspace_,
                                                   arg.contraction_multi_d_kernel_args_.data(),
                                                   arg.contraction_multi_d_kernel_args_.size() *
                                                       sizeof(ContractionMultiDKernelArg),
-                                        hipMemcpyHostToDevice));
+                                                  hipMemcpyHostToDevice,
+                                                  stream_config.stream_id_));
 
             float ave_time = 0;
 

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

@@ -610,10 +610,11 @@ struct DeviceGroupedGemmSoftmaxGemmPermute_Xdl_CShuffle
                 some_has_main_k_block_loop |= y;
             }
 
-            hipGetErrorString(hipMemcpy(arg.p_workspace_,
+            hipGetErrorString(hipMemcpyWithStream(arg.p_workspace_,
                                                   arg.group_kernel_args_.data(),
                                                   arg.group_kernel_args_.size() * sizeof(GroupKernelArg),
-                                        hipMemcpyHostToDevice));
+                                                  hipMemcpyHostToDevice,
+                                                  stream_config.stream_id_));
 
             float ave_time = 0;
 

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

@@ -543,10 +543,11 @@ struct DeviceGroupedGemm_Xdl : public DeviceGroupedGemm<ALayout,
             }
 
             hipGetErrorString(
-                hipMemcpy(arg.p_workspace_,
+                hipMemcpyWithStream(arg.p_workspace_,
                                     arg.gemm_desc_kernel_arg_.data(),
                                     arg.gemm_desc_kernel_arg_.size() * sizeof(GemmBiasTransKernelArg),
-                          hipMemcpyHostToDevice));
+                                    hipMemcpyHostToDevice,
+                                    stream_config.stream_id_));
 
             float ave_time = 0;
 

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

@@ -934,10 +934,11 @@ struct DeviceGroupedMultiheadAttentionBackward_Kloop_Xdl_CShuffle_V1
             //     some_has_main_k_block_loop |= y;
             // }
 
-            hipGetErrorString(hipMemcpy(arg.p_workspace_,
+            hipGetErrorString(hipMemcpyAsync(arg.p_workspace_,
                                              arg.group_kernel_args_.data(),
                                              arg.group_kernel_args_.size() * sizeof(GroupKernelArg),
-                                        hipMemcpyHostToDevice));
+                                             hipMemcpyHostToDevice,
+                                             stream_config.stream_id_));
 
             float ave_time = 0;
 
@@ -954,6 +955,7 @@ struct DeviceGroupedMultiheadAttentionBackward_Kloop_Xdl_CShuffle_V1
                         has_main_k_block_loop_,
                         Deterministic>;
 
+                std::cerr << "Calling kernel kernel_grouped_multihead_attention_backward_kloop_xdl_cshuffle_v1 LINE: " << __LINE__ << " arg.p_workspace_ = " << arg.p_workspace_ << std::endl;
                 return launch_and_time_kernel(
                     stream_config,
                     kernel,

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

@@ -941,10 +941,11 @@ struct DeviceGroupedMultiheadAttentionBackward_Kloop_Xdl_CShuffle_V2
                 some_has_main_k_block_loop |= y;
             }
 
-            hipGetErrorString(hipMemcpy(arg.p_workspace_,
+            hipGetErrorString(hipMemcpyWithStream(arg.p_workspace_,
                                                   arg.group_kernel_args_.data(),
                                                   arg.group_kernel_args_.size() * sizeof(GroupKernelArg),
-                                        hipMemcpyHostToDevice));
+                                                  hipMemcpyHostToDevice,
+                                                  stream_config.stream_id_));
 
             float ave_time = 0;
 

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

@@ -955,10 +955,11 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V1
             //     some_has_main_k_block_loop |= y;
             // }
 
-            hipGetErrorString(hipMemcpy(arg.p_workspace_,
+            hipGetErrorString(hipMemcpyWithStream(arg.p_workspace_,
                                                   arg.group_kernel_args_.data(),
                                                   arg.group_kernel_args_.size() * sizeof(GroupKernelArg),
-                                        hipMemcpyHostToDevice));
+                                                  hipMemcpyHostToDevice,
+                                                  stream_config.stream_id_));
 
             float ave_time = 0;
 

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

@@ -962,10 +962,11 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
                 some_has_main_k_block_loop |= y;
             }
 
-            hipGetErrorString(hipMemcpy(arg.p_workspace_,
+            hipGetErrorString(hipMemcpyWithStream(arg.p_workspace_,
                                                   arg.group_kernel_args_.data(),
                                                   arg.group_kernel_args_.size() * sizeof(GroupKernelArg),
-                                        hipMemcpyHostToDevice));
+                                                  hipMemcpyHostToDevice,
+                                                  stream_config.stream_id_));
 
             float ave_time = 0;
 

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

@@ -804,10 +804,11 @@ struct DeviceGroupedMultiheadAttentionForward_Xdl_CShuffle_V1
                 some_has_main_k_block_loop |= y;
             }
 
-            hipGetErrorString(hipMemcpy(arg.p_workspace_,
+            hipGetErrorString(hipMemcpyWithStream(arg.p_workspace_,
                                                   arg.group_kernel_args_.data(),
                                                   arg.group_kernel_args_.size() * sizeof(GroupKernelArg),
-                                        hipMemcpyHostToDevice));
+                                                  hipMemcpyHostToDevice,
+                                                  stream_config.stream_id_));
 
             float ave_time = 0;
 

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

@@ -826,10 +826,11 @@ struct DeviceGroupedMultiheadAttentionForward_Xdl_CShuffle_V2
                 some_has_main_k_block_loop |= y;
             }
 
-            hipGetErrorString(hipMemcpy(arg.p_workspace_,
+            hipGetErrorString(hipMemcpyWithStream(arg.p_workspace_,
                                                   arg.group_kernel_args_.data(),
                                                   arg.group_kernel_args_.size() * sizeof(GroupKernelArg),
-                                        hipMemcpyHostToDevice));
+                                                  hipMemcpyHostToDevice,
+                                                  stream_config.stream_id_));
 
             float ave_time = 0;