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:
Alexander Heisler
2026-08-20 10:27:51 +03:00
committed by GitHub
co-authored by Oliver Simons
parent 929d47a391
commit d9b6be07d0
5 changed files with 32 additions and 27 deletions
+11 -8
View File
@@ -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,