ggml-cuda: provide static workspace for cuBLAS handles (#26574)
* provide static workspace for cuBLAS handles * account for concurrent streams when using GGML_CUDA_GRAPH_OPT * drop cublas_handle overloads and remove direct cublasSetStream calls * Update ggml/src/ggml-cuda/common.cuh --------- Co-authored-by: Oliver Simons <osimons@nvidia.com>
This commit is contained in:
co-authored by
Oliver Simons
parent
929d47a391
commit
d9b6be07d0
@@ -711,9 +711,12 @@ ggml_backend_cuda_context::~ggml_backend_cuda_context() {
|
||||
if (streams[i][j] != nullptr) {
|
||||
CUDA_CHECK(cudaStreamDestroy(streams[i][j]));
|
||||
}
|
||||
}
|
||||
if (cublas_handles[i] != nullptr) {
|
||||
CUBLAS_CHECK(cublasDestroy(cublas_handles[i]));
|
||||
if (cublas_handles[i][j] != nullptr) {
|
||||
CUBLAS_CHECK(cublasDestroy(cublas_handles[i][j]));
|
||||
}
|
||||
if (cublas_workspaces[i][j] != nullptr) {
|
||||
CUDA_CHECK(cudaFree(cublas_workspaces[i][j]));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1416,7 +1419,7 @@ static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const
|
||||
|
||||
const int64_t ne_dst = ggml_nelements(dst);
|
||||
cudaStream_t main_stream = ctx.stream();
|
||||
CUBLAS_CHECK(cublasSetStream(ctx.cublas_handle(), main_stream));
|
||||
cublasHandle_t cublas_h = ctx.cublas_handle();
|
||||
|
||||
const size_t src0_ts = ggml_type_size(src0->type);
|
||||
GGML_ASSERT(nb00 == src0_ts);
|
||||
@@ -1539,14 +1542,14 @@ static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const
|
||||
// probably because the internal kernel selection logic is suboptimal.
|
||||
if (compute_type == GGML_TYPE_F32 && ne12 == 1 && ne13 == 1) {
|
||||
CUBLAS_CHECK(
|
||||
cublasSgemm(ctx.cublas_handle(), CUBLAS_OP_T, CUBLAS_OP_N,
|
||||
cublasSgemm(cublas_h, CUBLAS_OP_T, CUBLAS_OP_N,
|
||||
ne01, ne11, ne10,
|
||||
(const float *) alpha, (const float *) src0_ptr, s01,
|
||||
(const float *) src1_ptr, s11,
|
||||
(const float *) beta, (float *) dst_ptr, ne0));
|
||||
} else if (ne12 == 1 && ne13 == 1) {
|
||||
CUBLAS_CHECK(
|
||||
cublasGemmEx(ctx.cublas_handle(), CUBLAS_OP_T, CUBLAS_OP_N,
|
||||
cublasGemmEx(cublas_h, CUBLAS_OP_T, CUBLAS_OP_N,
|
||||
ne01, ne11, ne10,
|
||||
alpha, src0_ptr, cu_data_type_a, s01,
|
||||
src1_ptr, cu_data_type_b, s11,
|
||||
@@ -1561,7 +1564,7 @@ static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const
|
||||
// there is no broadcast and src0, src1 are contiguous across dims 2, 3
|
||||
// use cublasGemmStridedBatchedEx
|
||||
CUBLAS_CHECK(
|
||||
cublasGemmStridedBatchedEx(ctx.cublas_handle(), CUBLAS_OP_T, CUBLAS_OP_N,
|
||||
cublasGemmStridedBatchedEx(cublas_h, CUBLAS_OP_T, CUBLAS_OP_N,
|
||||
ne01, ne11, ne10,
|
||||
alpha, src0_ptr, cu_data_type_a, s01, sma, // strideA
|
||||
src1_ptr, cu_data_type_b, s11, smb, // strideB
|
||||
@@ -1599,7 +1602,7 @@ static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const
|
||||
CUDA_CHECK(cudaGetLastError());
|
||||
|
||||
CUBLAS_CHECK(
|
||||
cublasGemmBatchedEx(ctx.cublas_handle(), CUBLAS_OP_T, CUBLAS_OP_N,
|
||||
cublasGemmBatchedEx(cublas_h, CUBLAS_OP_T, CUBLAS_OP_N,
|
||||
ne01, ne11, ne10,
|
||||
alpha, (const void **) (ptrs_src.get() + 0*ne23), cu_data_type_a, s01,
|
||||
(const void **) (ptrs_src.get() + 1*ne23), cu_data_type_b, s11,
|
||||
|
||||
Reference in New Issue
Block a user