From e5b810bed18697ac4708f8a7ee6911752b8eb655 Mon Sep 17 00:00:00 2001 From: Qi Yuhang <45795032+HydraQYH@users.noreply.github.com> Date: Thu, 1 May 2025 03:28:05 +0800 Subject: [PATCH] Use cudaMemcpyAsync in gemm grouped with kRequiresPrecomputation schedule. (#2256) Co-authored-by: Yuhang Qi --- include/cutlass/gemm/device/base_grouped.h | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/include/cutlass/gemm/device/base_grouped.h b/include/cutlass/gemm/device/base_grouped.h index fc59d7ff..d9c2423b 100644 --- a/include/cutlass/gemm/device/base_grouped.h +++ b/include/cutlass/gemm/device/base_grouped.h @@ -127,8 +127,8 @@ private: } /// Copy from `data` to `workspace` - Status copy_to_workspace(void* workspace, void* data, size_t bytes) { - cudaError_t cuda_error = cudaMemcpy(workspace, data, bytes, cudaMemcpyHostToDevice); + Status copy_to_workspace(void* workspace, void* data, size_t bytes, cudaStream_t stream = nullptr) { + cudaError_t cuda_error = cudaMemcpyAsync(workspace, data, bytes, cudaMemcpyHostToDevice, stream); if (cuda_error != cudaSuccess) { // Call cudaGetLastError() to clear the error bit cuda_error = cudaGetLastError(); @@ -142,14 +142,14 @@ private: } /// Precomputes scheduling information for the grouped GEMM - Status precompute(Arguments const &args, int32_t tile_count, void* workspace) { + Status precompute(Arguments const &args, int32_t tile_count, void* workspace, cudaStream_t stream = nullptr) { size_t workspace_bytes = get_workspace_size(args); std::vector host_workspace(workspace_bytes); BaseKernel::ProblemVisitor::host_precompute(args.host_problem_sizes, args.problem_count, args.threadblock_count, (void*)host_workspace.data()); - return copy_to_workspace(workspace, host_workspace.data(), workspace_bytes); + return copy_to_workspace(workspace, host_workspace.data(), workspace_bytes, stream); } /// Reorder `data` according to `indices` @@ -361,7 +361,7 @@ public: if (BaseKernel::ProblemVisitor::kRequiresPrecomputation) { int32_t tile_count = group_tile_count(args); - Status status = precompute(args, tile_count, workspace); + Status status = precompute(args, tile_count, workspace, stream); if (status != Status::kSuccess) { return status; } @@ -388,7 +388,7 @@ public: } /// Lightweight update given a subset of arguments - Status update(Arguments const &args, void *workspace = nullptr) { + Status update(Arguments const &args, void *workspace = nullptr, cudaStream_t stream = nullptr) { size_t workspace_bytes = get_workspace_size(args); @@ -398,7 +398,7 @@ public: if (BaseKernel::ProblemVisitor::kRequiresPrecomputation) { int32_t tile_count = group_tile_count(args); - Status status = precompute(args, tile_count, workspace); + Status status = precompute(args, tile_count, workspace, stream); if (status != Status::kSuccess) { return status; }