leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 24 Packages Projects Releases Wiki Activity

Use cudaMemcpyAsync in gemm grouped with kRequiresPrecomputation schedule. (#2256)

Browse Source 
Co-authored-by: Yuhang Qi <qiyuhang@bytedance.com>
This commit is contained in:
Qi Yuhang
2025-04-30 15:28:05 -04:00
committed by GitHub
co-authored by Yuhang Qi
parent 2b78c2fe31
commit e5b810bed1
1 changed files with 7 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/cutlass/gemm/device/base_grouped.h
+7 -7
View File
@@ -127,8 +127,8 @@ private:
} }
/// Copy from `data` to `workspace` /// Copy from `data` to `workspace`
Status copy_to_workspace(void* workspace, void* data, size_t bytes) { Status copy_to_workspace(void* workspace, void* data, size_t bytes, cudaStream_t stream = nullptr) {
cudaError_t cuda_error = cudaMemcpy(workspace, data, bytes, cudaMemcpyHostToDevice); cudaError_t cuda_error = cudaMemcpyAsync(workspace, data, bytes, cudaMemcpyHostToDevice, stream);
if (cuda_error != cudaSuccess) { if (cuda_error != cudaSuccess) {
// Call cudaGetLastError() to clear the error bit // Call cudaGetLastError() to clear the error bit
cuda_error = cudaGetLastError(); cuda_error = cudaGetLastError();
@@ -142,14 +142,14 @@ private:
} }
/// Precomputes scheduling information for the grouped GEMM /// 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); size_t workspace_bytes = get_workspace_size(args);
std::vector<uint8_t> host_workspace(workspace_bytes); std::vector<uint8_t> host_workspace(workspace_bytes);
BaseKernel::ProblemVisitor::host_precompute(args.host_problem_sizes, BaseKernel::ProblemVisitor::host_precompute(args.host_problem_sizes,
args.problem_count, args.problem_count,
args.threadblock_count, args.threadblock_count,
(void*)host_workspace.data()); (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` /// Reorder `data` according to `indices`
@@ -361,7 +361,7 @@ public:
if (BaseKernel::ProblemVisitor::kRequiresPrecomputation) { if (BaseKernel::ProblemVisitor::kRequiresPrecomputation) {
int32_t tile_count = group_tile_count(args); 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) { if (status != Status::kSuccess) {
return status; return status;
} }
@@ -388,7 +388,7 @@ public:
} }
/// Lightweight update given a subset of arguments /// 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); size_t workspace_bytes = get_workspace_size(args);
@@ -398,7 +398,7 @@ public:
if (BaseKernel::ProblemVisitor::kRequiresPrecomputation) { if (BaseKernel::ProblemVisitor::kRequiresPrecomputation) {
int32_t tile_count = group_tile_count(args); 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) { if (status != Status::kSuccess) {
return status; return status;
} }