ggml-cuda: provide static workspace for cuBLAS handles (llama/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 03:27:51 -04:00 committed by Georgi Gerganov
parent acfad32841
commit 2648a706e0
5 changed files with 32 additions and 27 deletions

View File

@ -1418,7 +1418,9 @@ struct ggml_backend_cuda_context {
cudaEvent_t copy_event = nullptr;
cudaStream_t streams[GGML_CUDA_MAX_DEVICES][GGML_CUDA_MAX_STREAMS] = { { nullptr } };
cublasHandle_t cublas_handles[GGML_CUDA_MAX_DEVICES] = {nullptr};
cublasHandle_t cublas_handles[GGML_CUDA_MAX_DEVICES][GGML_CUDA_MAX_STREAMS] = {nullptr};
void * cublas_workspaces[GGML_CUDA_MAX_DEVICES][GGML_CUDA_MAX_STREAMS] = {nullptr};
size_t cublas_workspace_sizes[GGML_CUDA_MAX_DEVICES] = {0};
int curr_stream_no = 0;
@ -1495,17 +1497,22 @@ struct ggml_backend_cuda_context {
ggml_cuda_stream_context & stream_context() { return concurrent_stream_context; }
cublasHandle_t cublas_handle(int device) {
if (cublas_handles[device] == nullptr) {
ggml_cuda_set_device(device);
CUBLAS_CHECK(cublasCreate(&cublas_handles[device]));
CUBLAS_CHECK(cublasSetMathMode(cublas_handles[device], CUBLAS_TF32_TENSOR_OP_MATH));
}
return cublas_handles[device];
}
cublasHandle_t cublas_handle() {
return cublas_handle(device);
if (cublas_handles[device][curr_stream_no] == nullptr) {
ggml_cuda_set_device(device);
CUBLAS_CHECK(cublasCreate(&cublas_handles[device][curr_stream_no]));
CUBLAS_CHECK(cublasSetMathMode(cublas_handles[device][curr_stream_no], CUBLAS_TF32_TENSOR_OP_MATH));
CUBLAS_CHECK(cublasSetStream(cublas_handles[device][curr_stream_no], stream()));
#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && (CUBLAS_VER_MAJOR > 11 || (CUBLAS_VER_MAJOR == 11 && CUBLAS_VER_MINOR >= 2))
if (cublas_workspace_sizes[device] == 0) {
const int cc = ggml_cuda_info().devices[device].cc;
cublas_workspace_sizes[device] = (cc >= GGML_CUDA_CC_HOPPER) ? 32 * 1024 * 1024 : 4 * 1024 * 1024;
}
CUDA_CHECK(cudaMalloc(&cublas_workspaces[device][curr_stream_no], cublas_workspace_sizes[device]));
CUBLAS_CHECK(cublasSetWorkspace(cublas_handles[device][curr_stream_no], cublas_workspaces[device][curr_stream_no], cublas_workspace_sizes[device]));
#endif
}
return cublas_handles[device][curr_stream_no];
}
// pool

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,

View File

@ -54,8 +54,6 @@ void ggml_cuda_out_prod(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
const float alpha = 1.0f;
const float beta = 0.0f;
CUBLAS_CHECK(cublasSetStream(handle, stream));
const int64_t lda = nb01 / sizeof(float);
const int64_t ldc = nb1 / sizeof(float);

View File

@ -65,15 +65,13 @@ static void solve_tri_f32_cublas(ggml_backend_cuda_context & ctx,
get_batch_pointers<<<(total_batches + 255) / 256, 256, 0, stream>>>(A, X, A_ptrs_dev, X_ptrs_dev, ne02,
total_batches, s02, s03, s2, s3);
CUBLAS_CHECK(cublasSetStream(ctx.cublas_handle(id), stream));
// Yes, this is necessary, without this we get RMSE errors
CUBLAS_CHECK(cublasSetMathMode(ctx.cublas_handle(id), CUBLAS_DEFAULT_MATH));
CUBLAS_CHECK(cublasStrsmBatched(ctx.cublas_handle(id), CUBLAS_SIDE_RIGHT, CUBLAS_FILL_MODE_UPPER, CUBLAS_OP_N,
CUBLAS_CHECK(cublasSetMathMode(ctx.cublas_handle(), CUBLAS_DEFAULT_MATH));
CUBLAS_CHECK(cublasStrsmBatched(ctx.cublas_handle(), CUBLAS_SIDE_RIGHT, CUBLAS_FILL_MODE_UPPER, CUBLAS_OP_N,
CUBLAS_DIAG_NON_UNIT, k, n, &alpha, A_ptrs_dev, n, X_ptrs_dev, k, total_batches));
// revert to standard mode from common.cuh
CUBLAS_CHECK(cublasSetMathMode(ctx.cublas_handle(id), CUBLAS_TF32_TENSOR_OP_MATH));
CUBLAS_CHECK(cublasSetMathMode(ctx.cublas_handle(), CUBLAS_TF32_TENSOR_OP_MATH));
GGML_UNUSED_VARS(s12, s13);
}

View File

@ -632,7 +632,6 @@ static void ssm_scan_ssd_f32_cuda(
// Step 3: chunked SSD loop
// Per chunk: pre_matmul (incl. M) + 4 cuBLAS (CB, Y, S@C, state update) + scale_state
cublasHandle_t handle = ctx.cublas_handle();
CUBLAS_CHECK(cublasSetStream(handle, stream));
const float alpha_one = 1.0f;
const float beta_zero = 0.0f;
const float beta_one = 1.0f;