diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index e068116f3..1122232a9 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -1409,13 +1409,13 @@ static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const using traits = batched_mul_mat_traits; using cuda_t = typename traits::cuda_type; - GGML_ASSERT(ggml_is_contiguous(dst)); - - // Byte offsets and tensor dimensions are currently used in an inconsistent way for dst. - // As long as dst is contiguous this does not matter though. + // F32 chunks can have padding between rows to preserve the original destination stride. + GGML_ASSERT(ggml_is_contiguous(dst) || + (compute_type == GGML_TYPE_F32 && ggml_is_contiguous_1(dst))); GGML_TENSOR_BINARY_OP_LOCALS + const int64_t ldc = nb1 / sizeof(float); const int64_t ne_dst = ggml_nelements(dst); cudaStream_t main_stream = ctx.stream(); cublasHandle_t cublas_h = ctx.cublas_handle(); @@ -1545,14 +1545,14 @@ static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const ne01, ne11, ne10, (const float *) alpha, (const float *) src0_ptr, s01, (const float *) src1_ptr, s11, - (const float *) beta, (float *) dst_ptr, ne0)); + (const float *) beta, (float *) dst_ptr, ldc)); } else if (ne12 == 1 && ne13 == 1) { CUBLAS_CHECK( 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, - beta, dst_ptr, cu_data_type, ne0, + beta, dst_ptr, cu_data_type, ldc, cu_compute_type, CUBLAS_GEMM_DEFAULT_TENSOR_OP)); } else if (r2 == 1 && r3 == 1 && is_src0_cont_2 && is_src1_cont_2) { @@ -1567,7 +1567,7 @@ static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const ne01, ne11, ne10, alpha, src0_ptr, cu_data_type_a, s01, sma, // strideA src1_ptr, cu_data_type_b, s11, smb, // strideB - beta, dst_ptr, cu_data_type, ne0, ne1*ne0, // strideC + beta, dst_ptr, cu_data_type, ldc, ne1*ldc, // strideC ne12*ne13, cu_compute_type, CUBLAS_GEMM_DEFAULT_TENSOR_OP)); @@ -1605,7 +1605,7 @@ static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const 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, - beta, ( void **) (ptrs_dst.get() + 0*ne23), cu_data_type, ne0, + beta, ( void **) (ptrs_dst.get() + 0*ne23), cu_data_type, ldc, ne23, cu_compute_type, CUBLAS_GEMM_DEFAULT_TENSOR_OP)); @@ -1658,6 +1658,32 @@ static void ggml_cuda_mul_mat_cublas(ggml_backend_cuda_context & ctx, const ggml } } + constexpr size_t max_src0_convert_size = 512ull * 1024 * 1024; + const size_t src0_f32_size = ggml_nelements(src0) * sizeof(float); + + if (compute_type == GGML_TYPE_F32 && + (src0->type == GGML_TYPE_F16 || src0->type == GGML_TYPE_BF16) && + src0_f32_size > max_src0_convert_size) { + const size_t f32_row_size = src0_f32_size / src0->ne[1]; + const int64_t rows_per_chunk = std::max(1, (int64_t) (max_src0_convert_size / f32_row_size)); + + if (rows_per_chunk < src0->ne[1]) { + ggml_tensor src0_chunk = *src0; + ggml_tensor dst_chunk = *dst; + + for (int64_t i01 = 0; i01 < src0->ne[1]; i01 += rows_per_chunk) { + src0_chunk.ne[1] = std::min(rows_per_chunk, src0->ne[1] - i01); + src0_chunk.data = (char *) src0->data + i01*src0->nb[1]; + + dst_chunk.ne[0] = src0_chunk.ne[1]; + dst_chunk.data = (char *) dst->data + i01*dst->nb[0]; + + ggml_cuda_mul_mat_cublas_impl(ctx, &src0_chunk, src1, &dst_chunk); + } + return; + } + } + switch (compute_type) { case GGML_TYPE_F32: ggml_cuda_mul_mat_cublas_impl(ctx, src0, src1, dst);