Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

ggml-cuda: provide static workspace for cuBLAS handles (#26574)

Browse Source 
* 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
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-cuda/common.cuh
+18 -11
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