Use hipMemcpyAsync() to replace hipMemcpy() in grouped fwd device operator

e29a9111 · Qianfeng Zhang · 299ac0a3 · e29a9111
Commit e29a9111 authored Sep 13, 2023 by Qianfeng Zhang
Hide whitespace changes
Inline Side-by-side

Showing with 5 additions and 4 deletions

include/ck/tensor_operation/gpu/device/impl/device_grouped_mha_fwd_xdl_cshuffle_v2.hpp ...pu/device/impl/device_grouped_mha_fwd_xdl_cshuffle_v2.hpp +5 -4

No files found.
--- a/include/ck/tensor_operation/gpu/device/impl/device_grouped_mha_fwd_xdl_cshuffle_v2.hpp
+++ b/include/ck/tensor_operation/gpu/device/impl/device_grouped_mha_fwd_xdl_cshuffle_v2.hpp
@@ -912,10 +912,11 @@ struct DeviceGroupedMultiheadAttentionForward_Xdl_CShuffle_V2
                some_has_main_k_block_loop |= y;
            }
-            hipGetErrorString(hipMemcpy(arg.p_workspace_,
+            HIP_CHECK_ERROR(hipMemcpyAsync(arg.p_workspace_,
-                                        arg.group_kernel_args_.data(),
+                                           arg.group_kernel_args_.data(),
-                                        arg.group_kernel_args_.size() * sizeof(GroupKernelArg),
+                                           arg.group_kernel_args_.size() * sizeof(GroupKernelArg),
-                                        hipMemcpyHostToDevice));
+                                           hipMemcpyHostToDevice,
+                                           stream_config.stream_id_));
            float ave_time = 0;