cuda : fix tensor size calculation for non-split buffer (llama/5145)

This commit is contained in:
slaren 2024-01-26 18:59:43 +01:00 committed by Georgi Gerganov
parent c65edd5b64
commit 0878ab7c15
No known key found for this signature in database
GPG Key ID: 449E073F9DC10735
2 changed files with 8 additions and 15 deletions

View File

@ -30,7 +30,9 @@ size_t ggml_backend_buft_get_alignment(ggml_backend_buffer_type_t buft) {
GGML_CALL size_t ggml_backend_buft_get_alloc_size(ggml_backend_buffer_type_t buft, struct ggml_tensor * tensor) { GGML_CALL size_t ggml_backend_buft_get_alloc_size(ggml_backend_buffer_type_t buft, struct ggml_tensor * tensor) {
// get_alloc_size is optional, defaults to ggml_nbytes // get_alloc_size is optional, defaults to ggml_nbytes
if (buft->iface.get_alloc_size) { if (buft->iface.get_alloc_size) {
return buft->iface.get_alloc_size(buft, tensor); size_t size = buft->iface.get_alloc_size(buft, tensor);
assert(size >= ggml_nbytes(tensor));
return size;
} }
return ggml_nbytes(tensor); return ggml_nbytes(tensor);
} }

View File

@ -9790,8 +9790,8 @@ static void ggml_cuda_mul_mat_id(const ggml_tensor * src0, const ggml_tensor * s
// TODO: mmq/mmv support // TODO: mmq/mmv support
#endif #endif
const int64_t nb11 = src1->nb[1]; const size_t nb11 = src1->nb[1];
const int64_t nb1 = dst->nb[1]; const size_t nb1 = dst->nb[1];
const struct ggml_tensor * ids = src0; const struct ggml_tensor * ids = src0;
const int32_t id = ((int32_t *) dst->op_params)[0]; const int32_t id = ((int32_t *) dst->op_params)[0];
@ -10304,15 +10304,11 @@ GGML_CALL static void ggml_backend_cuda_buffer_init_tensor(ggml_backend_buffer_t
if (ggml_is_quantized(tensor->type)) { if (ggml_is_quantized(tensor->type)) {
// initialize padding to 0 to avoid possible NaN values // initialize padding to 0 to avoid possible NaN values
int64_t row_low = 0; size_t original_size = ggml_nbytes(tensor);
int64_t row_high = ggml_nrows(tensor);
int64_t nrows_split = row_high - row_low;
size_t original_size = ggml_nbytes_split(tensor, nrows_split);
size_t padded_size = ggml_backend_buft_get_alloc_size(buffer->buft, tensor); size_t padded_size = ggml_backend_buft_get_alloc_size(buffer->buft, tensor);
if (padded_size > original_size && tensor->view_src == nullptr) { if (padded_size > original_size && tensor->view_src == nullptr) {
CUDA_CHECK(cudaMemsetAsync((char *)tensor->data + original_size, 0, padded_size - original_size, g_cudaStreams[ctx->device][0])); CUDA_CHECK(cudaMemset((char *)tensor->data + original_size, 0, padded_size - original_size));
} }
} }
} }
@ -10415,12 +10411,7 @@ GGML_CALL static size_t ggml_backend_cuda_buffer_type_get_alignment(ggml_backend
} }
GGML_CALL static size_t ggml_backend_cuda_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor) { GGML_CALL static size_t ggml_backend_cuda_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor) {
int64_t row_low = 0; size_t size = ggml_nbytes(tensor);
int64_t row_high = ggml_nrows(tensor);
int64_t nrows_split = row_high - row_low;
size_t size = ggml_nbytes_split(tensor, nrows_split);
int64_t ne0 = tensor->ne[0]; int64_t ne0 = tensor->ne[0];
if (ggml_is_quantized(tensor->type)) { if (ggml_is_quantized(tensor->type)) {