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:
parent
acfad32841
commit
2648a706e0
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
Loading…
Reference in New Issue