From 2648a706e0eeadd7c5720aa31df2fbddc2f473d6 Mon Sep 17 00:00:00 2001 From: Alexander Heisler <126129661+heislera763@users.noreply.github.com> Date: Thu, 20 Aug 2026 03:27:51 -0400 Subject: [PATCH] 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 --- ggml/src/ggml-cuda/common.cuh | 29 ++++++++++++++++++----------- ggml/src/ggml-cuda/ggml-cuda.cu | 19 +++++++++++-------- ggml/src/ggml-cuda/out-prod.cu | 2 -- ggml/src/ggml-cuda/solve_tri.cu | 8 +++----- ggml/src/ggml-cuda/ssm-scan.cu | 1 - 5 files changed, 32 insertions(+), 27 deletions(-) diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh index d27d8acb1..14dd1098c 100644 --- a/ggml/src/ggml-cuda/common.cuh +++ b/ggml/src/ggml-cuda/common.cuh @@ -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 diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 3b2a0ea85..a8a1c09ca 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -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, diff --git a/ggml/src/ggml-cuda/out-prod.cu b/ggml/src/ggml-cuda/out-prod.cu index 46b9f3a67..c46e0455d 100644 --- a/ggml/src/ggml-cuda/out-prod.cu +++ b/ggml/src/ggml-cuda/out-prod.cu @@ -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); diff --git a/ggml/src/ggml-cuda/solve_tri.cu b/ggml/src/ggml-cuda/solve_tri.cu index 07ca33f51..d96783420 100644 --- a/ggml/src/ggml-cuda/solve_tri.cu +++ b/ggml/src/ggml-cuda/solve_tri.cu @@ -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); } diff --git a/ggml/src/ggml-cuda/ssm-scan.cu b/ggml/src/ggml-cuda/ssm-scan.cu index ef342f01f..40cb38dee 100644 --- a/ggml/src/ggml-cuda/ssm-scan.cu +++ b/ggml/src/ggml-cuda/ssm-scan.cu @@ -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;