diff --git a/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h index cc3f8cd36e..27375bd0a5 100644 --- a/ggml/include/ggml-backend.h +++ b/ggml/include/ggml-backend.h @@ -424,6 +424,10 @@ extern "C" { // Compare the output of two backends GGML_API bool ggml_backend_compare_graph_backend(ggml_backend_t backend1, ggml_backend_t backend2, struct ggml_cgraph * graph, ggml_backend_eval_callback callback, void * user_data, struct ggml_tensor const * const * test_nodes, size_t num_test_nodes); + // returns true for ops that may require additional memory for fleeting data on some backends, + // i.e. the backend's get_alloc_size may return more than ggml_nbytes for the output tensor + GGML_API bool ggml_backend_op_alloc_size_may_expand(enum ggml_op op); + // Tensor initialization GGML_API enum ggml_status ggml_backend_tensor_alloc(ggml_backend_buffer_t buffer, struct ggml_tensor * tensor, void * addr); GGML_API enum ggml_status ggml_backend_view_init(struct ggml_tensor * tensor); diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index 78eb10dfe9..fec7d7c92b 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -65,6 +65,13 @@ size_t ggml_backend_buft_get_alloc_size(ggml_backend_buffer_type_t buft, const s if (buft->iface.get_alloc_size) { size_t size = buft->iface.get_alloc_size(buft, tensor); assert(size >= ggml_nbytes(tensor)); + + // [TAG_ALLOC_SIZE_EXPAND] + // if you hit this assert, update ggml_backend_op_alloc_size_may_expand() accordingly + GGML_ASSERT(size <= ggml_nbytes(tensor) || + ggml_op_is_empty(tensor->op) || + ggml_backend_op_alloc_size_may_expand(tensor->op)); + return size; } return ggml_nbytes(tensor); @@ -2101,6 +2108,22 @@ ggml_backend_t ggml_backend_sched_get_tensor_backend(ggml_backend_sched_t sched, // utils +// [TAG_ALLOC_SIZE_EXPAND] +// returns true for ops that may require additional memory for fleeting data on some backends, +// i.e. the backend's get_alloc_size may return more than ggml_nbytes for the output tensor +bool ggml_backend_op_alloc_size_may_expand(enum ggml_op op) { + switch (op) { + case GGML_OP_FLASH_ATTN_EXT: + case GGML_OP_MUL_MAT_ID: + case GGML_OP_CUMSUM: + case GGML_OP_ARGSORT: + case GGML_OP_TOP_K: + return true; + default: + return false; + } +} + enum ggml_status ggml_backend_view_init(struct ggml_tensor * tensor) { GGML_ASSERT(tensor); GGML_ASSERT(tensor->buffer == NULL); diff --git a/ggml/src/ggml-rpc/ggml-rpc.cpp b/ggml/src/ggml-rpc/ggml-rpc.cpp index 9aa5883d80..58a8a030cf 100644 --- a/ggml/src/ggml-rpc/ggml-rpc.cpp +++ b/ggml/src/ggml-rpc/ggml-rpc.cpp @@ -826,10 +826,10 @@ static size_t ggml_backend_rpc_buffer_type_get_alloc_size(ggml_backend_buffer_ty // See comments in init_tensor. rpc_get |= ggml_is_quantized(tensor->type) && (tensor->ne[0] % 512 != 0) && (tensor->view_src == nullptr); - // ops that require additional memory for fleeting data on certain backends + // [TAG_ALLOC_SIZE_EXPAND] + // ops that may require additional memory for fleeting data on certain backends // ref: https://github.com/ggml-org/llama.cpp/pull/15966 - rpc_get |= tensor->op == GGML_OP_FLASH_ATTN_EXT; - rpc_get |= tensor->op == GGML_OP_MUL_MAT_ID; + rpc_get |= ggml_backend_op_alloc_size_may_expand(tensor->op); if (rpc_get) { ggml_backend_rpc_buffer_type_context * buft_ctx = (ggml_backend_rpc_buffer_type_context *)buft->context;