diff --git a/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h index cc3f8cd36e35..8c5916cd553a 100644 --- a/ggml/include/ggml-backend.h +++ b/ggml/include/ggml-backend.h @@ -319,6 +319,11 @@ extern "C" { GGML_API ggml_backend_sched_t ggml_backend_sched_new(ggml_backend_t * backends, ggml_backend_buffer_type_t * bufts, int n_backends, size_t graph_size, bool parallel, bool op_offload); GGML_API void ggml_backend_sched_free(ggml_backend_sched_t sched); + // Share physical compute buffers while keeping scheduler and graph allocation state independent. + // The caller must ensure the schedulers do not execute concurrently while the buffers are shared. + // Returns false for incompatible or unsupported scheduler layouts, or if dst is already allocated. + GGML_API bool ggml_backend_sched_share_compute_buffers(ggml_backend_sched_t dst, ggml_backend_sched_t src); + // Initialize backend buffers from a measure graph GGML_API void ggml_backend_sched_reserve_size(ggml_backend_sched_t sched, struct ggml_cgraph * measure_graph, size_t * sizes); GGML_API bool ggml_backend_sched_reserve(ggml_backend_sched_t sched, struct ggml_cgraph * measure_graph); // returns success diff --git a/ggml/src/ggml-alloc.c b/ggml/src/ggml-alloc.c index 3bda9abbe03d..5cb3445b06e5 100644 --- a/ggml/src/ggml-alloc.c +++ b/ggml/src/ggml-alloc.c @@ -478,8 +478,15 @@ struct node_alloc { struct tensor_alloc src[GGML_MAX_SRC]; }; +struct ggml_gallocr_shared_buffers { + size_t refs; + int n_buffers; + struct vbuffer ** buffers; +}; + struct ggml_gallocr { ggml_backend_buffer_type_t * bufts; // [n_buffers] + struct ggml_gallocr_shared_buffers * shared_buffers; struct vbuffer ** buffers; // [n_buffers] struct ggml_dyn_tallocr ** buf_tallocs; // [n_buffers] int n_buffers; @@ -534,11 +541,39 @@ ggml_gallocr_t ggml_gallocr_new(ggml_backend_buffer_type_t buft) { return ggml_gallocr_new_n(&buft, 1); } +static void ggml_gallocr_shared_buffers_unref(struct ggml_gallocr_shared_buffers * shared) { + GGML_ASSERT(shared != NULL && shared->refs > 0); + if (--shared->refs > 0) { + return; + } + + for (int i = 0; i < shared->n_buffers; ++i) { + bool freed = false; + for (int j = 0; j < i; ++j) { + if (shared->buffers[j] == shared->buffers[i]) { + freed = true; + break; + } + } + if (!freed) { + ggml_vbuffer_free(shared->buffers[i]); + } + } + + free(shared->buffers); + free(shared); +} + void ggml_gallocr_free(ggml_gallocr_t galloc) { if (galloc == NULL) { return; } + if (galloc->shared_buffers != NULL) { + ggml_gallocr_shared_buffers_unref(galloc->shared_buffers); + galloc->buffers = NULL; + } + for (int i = 0; i < galloc->n_buffers; i++) { if (galloc->buffers != NULL) { // skip if already freed @@ -578,6 +613,40 @@ void ggml_gallocr_free(ggml_gallocr_t galloc) { free(galloc); } +bool ggml_gallocr_share_buffers(ggml_gallocr_t dst, ggml_gallocr_t src) { + if (dst == NULL || src == NULL || dst == src || + dst->n_buffers != src->n_buffers || dst->shared_buffers != NULL) { + return false; + } + + for (int i = 0; i < dst->n_buffers; ++i) { + if (dst->bufts[i] != src->bufts[i] || dst->buffers[i] != NULL || src->buffers[i] == NULL) { + return false; + } + } + + struct ggml_gallocr_shared_buffers * shared = src->shared_buffers; + if (shared != NULL && + (shared->n_buffers != src->n_buffers || shared->buffers != src->buffers)) { + return false; + } + + if (shared == NULL) { + shared = calloc(1, sizeof(*shared)); + GGML_ASSERT(shared != NULL); + shared->refs = 1; + shared->n_buffers = src->n_buffers; + shared->buffers = src->buffers; + src->shared_buffers = shared; + } + + free(dst->buffers); + shared->refs++; + dst->shared_buffers = shared; + dst->buffers = shared->buffers; + return true; +} + typedef struct ggml_gallocr * ggml_gallocr_t; static struct hash_node * ggml_gallocr_hash_get(ggml_gallocr_t galloc, struct ggml_tensor * t) { @@ -821,6 +890,18 @@ static void ggml_gallocr_alloc_graph_impl(ggml_gallocr_t galloc, struct ggml_cgr } } +static bool ggml_gallocr_shared_buffers_need_grow(ggml_gallocr_t galloc) { + for (int i = 0; i < galloc->n_buffers; ++i) { + for (int c = 0; c < galloc->buf_tallocs[i]->n_chunks; ++c) { + const size_t current = galloc->buffers[i] ? ggml_vbuffer_chunk_size(galloc->buffers[i], c) : 0; + if (ggml_dyn_tallocr_max_size(galloc->buf_tallocs[i], c) > current) { + return true; + } + } + } + return false; +} + static bool ggml_gallocr_reserve_n_impl( ggml_gallocr_t galloc, struct ggml_cgraph * graph, const int * node_buffer_ids, const int * leaf_buffer_ids, bool no_alloc) { size_t min_hash_size = graph->n_nodes + graph->n_leafs; @@ -900,6 +981,18 @@ static bool ggml_gallocr_reserve_n_impl( } } + // Never resize storage still referenced by another gallocr. + if (galloc->shared_buffers != NULL && ggml_gallocr_shared_buffers_need_grow(galloc)) { + GGML_LOG_DEBUG("%s: detaching shared compute buffers for a larger reservation\n", __func__); + struct vbuffer ** buffers = calloc(galloc->n_buffers, sizeof(*buffers)); + GGML_ASSERT(buffers != NULL); + + struct ggml_gallocr_shared_buffers * shared = galloc->shared_buffers; + galloc->shared_buffers = NULL; + galloc->buffers = buffers; + ggml_gallocr_shared_buffers_unref(shared); + } + // reallocate buffers if needed for (int i = 0; i < galloc->n_buffers; i++) { // if the buffer type is used multiple times, we reuse the same buffer diff --git a/ggml/src/ggml-backend-impl.h b/ggml/src/ggml-backend-impl.h index 9c56ec30c5f1..b82fdb911ecf 100644 --- a/ggml/src/ggml-backend-impl.h +++ b/ggml/src/ggml-backend-impl.h @@ -84,6 +84,13 @@ extern "C" { GGML_API bool ggml_backend_buffer_is_multi_buffer(ggml_backend_buffer_t buffer); GGML_API void ggml_backend_multi_buffer_set_usage(ggml_backend_buffer_t buffer, enum ggml_backend_buffer_usage usage); + // graph allocator internals + struct ggml_gallocr; +#if defined(__GNUC__) || defined(__clang__) + __attribute__((visibility("hidden"))) +#endif + bool ggml_gallocr_share_buffers(struct ggml_gallocr * dst, struct ggml_gallocr * src); + // // Backend (meta) // diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index 3d6310f3ffea..f5b037133587 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -1890,6 +1890,23 @@ void ggml_backend_sched_free(ggml_backend_sched_t sched) { free(sched); } +bool ggml_backend_sched_share_compute_buffers(ggml_backend_sched_t dst, ggml_backend_sched_t src) { + if (dst == nullptr || src == nullptr || dst == src || + dst->n_copies != 1 || src->n_copies != 1 || + dst->n_backends != src->n_backends || dst->is_alloc) { + return false; + } + + for (int i = 0; i < dst->n_backends; ++i) { + if (dst->bufts[i] != src->bufts[i] || + ggml_backend_get_device(dst->backends[i]) != ggml_backend_get_device(src->backends[i])) { + return false; + } + } + + return ggml_gallocr_share_buffers(dst->galloc, src->galloc); +} + void ggml_backend_sched_reset(ggml_backend_sched_t sched) { GGML_ASSERT(sched); // reset state for the next run diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 52f8d53672a3..c0d4ca4f255f 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -160,6 +160,12 @@ llama_context::llama_context( } } + if (params.ctx_type == LLAMA_CONTEXT_TYPE_MTP && params.ctx_other != nullptr && + ¶ms.ctx_other->get_model() == &model && + cparams.n_seq_max == 1 && params.ctx_other->n_seq_max() == 1) { + ctx_compute = params.ctx_other; + } + auto rope_scaling_type = params.rope_scaling_type; if (rope_scaling_type == LLAMA_ROPE_SCALING_TYPE_UNSPECIFIED) { rope_scaling_type = hparams.rope_scaling_type_train; @@ -603,6 +609,28 @@ void llama_context::sched_reserve() { sched.reset(ggml_backend_sched_new(backend_ptrs.data(), backend_buft.data(), backend_ptrs.size(), max_nodes, cparams.pipeline_parallel, cparams.op_offload)); + if (ctx_compute != nullptr) { + int n_devices = 0; + bool is_cuda = true; + for (ggml_backend_t backend : backend_ptrs) { + ggml_backend_dev_t device = ggml_backend_get_device(backend); + if (device == nullptr || ggml_backend_dev_type(device) == GGML_BACKEND_DEVICE_TYPE_CPU) { + continue; + } + + n_devices++; + ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(device); + is_cuda = is_cuda && reg != nullptr && strcmp(ggml_backend_reg_name(reg), "CUDA") == 0; + } + + if (n_devices == 1 && is_cuda && + ggml_backend_sched_share_compute_buffers(sched.get(), ctx_compute->get_sched())) { + LLAMA_LOG_INFO("%s: sharing compute buffers with the target context\n", __func__); + } else { + LLAMA_LOG_INFO("%s: compute buffer sharing unavailable; using independent buffers\n", __func__); + } + } + llama_memory_context_ptr mctx; if (memory) { LLAMA_LOG_DEBUG("%s: reserving full memory module\n", __func__); diff --git a/src/llama-context.h b/src/llama-context.h index bf91daa8b562..ef153405b03d 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -342,6 +342,7 @@ struct llama_context { std::vector output_swaps; ggml_backend_sched_ptr sched; + llama_context * ctx_compute = nullptr; bool sched_need_reserve = true; diff --git a/tests/test-alloc.cpp b/tests/test-alloc.cpp index 6d5428493e70..1a6f9b6f75df 100644 --- a/tests/test-alloc.cpp +++ b/tests/test-alloc.cpp @@ -88,6 +88,8 @@ static void dummy_backend_buffer_clear(ggml_backend_buffer_t, uint8_t) {} struct dummy_backend { std::unique_ptr context; ggml_backend_buffer_type buffer_type; + ggml_backend_device device; + ggml_backend backend; }; static dummy_backend dummy_backend_init(size_t max_buffer_size, size_t alignment = 8) { @@ -113,6 +115,40 @@ static dummy_backend dummy_backend_init(size_t max_buffer_size, size_t alignment return b; } +// +// ggml_backend / device interface +// + +static const char * dummy_backend_get_name(ggml_backend_t) { + return "dummy_backend"; +} + +static enum ggml_backend_dev_type dummy_backend_device_get_type(ggml_backend_dev_t) { + return GGML_BACKEND_DEVICE_TYPE_CPU; +} + +static bool dummy_backend_device_supports_op(ggml_backend_dev_t, const ggml_tensor *) { + return true; +} + +static bool dummy_backend_device_supports_buft(ggml_backend_dev_t device, ggml_backend_buffer_type_t buft) { + return buft == device->context; +} + +static ggml_backend_sched_ptr dummy_backend_sched_new(dummy_backend & b) { + b.device.iface.get_type = dummy_backend_device_get_type; + b.device.iface.supports_op = dummy_backend_device_supports_op; + b.device.iface.supports_buft = dummy_backend_device_supports_buft; + b.device.context = &b.buffer_type; + + b.backend.iface.get_name = dummy_backend_get_name; + b.backend.device = &b.device; + + ggml_backend_t backends[] = { &b.backend }; + ggml_backend_buffer_type_t bufts[] = { &b.buffer_type }; + return ggml_backend_sched_ptr(ggml_backend_sched_new(backends, bufts, 1, 64, false, false)); +} + // // test utilities @@ -144,6 +180,13 @@ static ggml_tensor * make_input_with_size(ggml_context * ctx, size_t size_bytes) return make_input_1d(ctx, size_bytes / 4); } +static ggml_tensor * make_scale_graph(ggml_context * ctx, ggml_cgraph * graph, size_t size_bytes) { + ggml_tensor * out = ggml_scale(ctx, make_input_with_size(ctx, size_bytes), 2.0f); + ggml_set_output(out); + ggml_build_forward_expand(graph, out); + return out; +} + static void assign_names(ggml_context * ctx, const char * prefix = "x") { int i = 0; for (ggml_tensor * t = ggml_get_first_tensor(ctx); t; t = ggml_get_next_tensor(ctx, t)) { @@ -583,6 +626,55 @@ static void test_reallocation() { } } +static void test_shared_buffers() { + dummy_backend backend = dummy_backend_init(SIZE_MAX); + + auto [ctx_src, graph_src, ctx_src_ptr] = make_context(); + ggml_tensor * src_out = make_scale_graph(ctx_src, graph_src, 16); + + ggml_backend_sched_ptr src = dummy_backend_sched_new(backend); + GGML_ASSERT(ggml_backend_sched_reserve(src.get(), graph_src)); + GGML_ASSERT(ggml_backend_sched_alloc_graph(src.get(), graph_src)); + const size_t shared_size = backend.context->allocated_total(); + GGML_ASSERT(shared_size > 0); + + auto [ctx_dst, graph_dst, ctx_dst_ptr] = make_context(); + ggml_tensor * dst_out = make_scale_graph(ctx_dst, graph_dst, 8); + + ggml_backend_sched_ptr dst = dummy_backend_sched_new(backend); + GGML_ASSERT(ggml_backend_sched_share_compute_buffers(dst.get(), src.get())); + GGML_ASSERT(ggml_backend_sched_reserve(dst.get(), graph_dst)); + GGML_ASSERT(ggml_backend_sched_alloc_graph(dst.get(), graph_dst)); + GGML_ASSERT(backend.context->allocated_total() == shared_size); + GGML_ASSERT(src_out->buffer == dst_out->buffer); + + // Schedulers keep independent allocation plans. + ggml_backend_sched_reset(src.get()); + GGML_ASSERT(ggml_backend_sched_alloc_graph(src.get(), graph_src)); + + auto [ctx_big, graph_big, ctx_big_ptr] = make_context(); + ggml_tensor * big_out = make_scale_graph(ctx_big, graph_big, 64); + + // A larger reservation detaches from shared buffers. + ggml_backend_sched_ptr big = dummy_backend_sched_new(backend); + GGML_ASSERT(ggml_backend_sched_share_compute_buffers(big.get(), src.get())); + GGML_ASSERT(ggml_backend_sched_reserve(big.get(), graph_big)); + GGML_ASSERT(ggml_backend_sched_alloc_graph(big.get(), graph_big)); + GGML_ASSERT(big_out->buffer != src_out->buffer); + GGML_ASSERT(backend.context->allocated_total() > shared_size); + + big.reset(); + GGML_ASSERT(backend.context->allocated_total() == shared_size); + // Freeing the source does not invalidate the destination. + src.reset(); + GGML_ASSERT(backend.context->allocated_total() == shared_size); + ggml_backend_sched_reset(dst.get()); + GGML_ASSERT(ggml_backend_sched_alloc_graph(dst.get(), graph_dst)); + GGML_ASSERT(backend.context->allocated_total() == shared_size); + dst.reset(); + GGML_ASSERT(backend.context->allocated_total() == 0); +} + static void run(const char * name, void (*f)()) { printf("%s ", name); fflush(stdout); @@ -604,5 +696,6 @@ int main() { run("test_multiple_buffer_types", test_multiple_buffer_types); run("test_buffer_size_zero", test_buffer_size_zero); run("test_reallocation", test_reallocation); + run("test_shared_buffers", test_shared_buffers); return 0; }