ggml-org · mushang0 · Aug 21, 2026
@@ -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
@@ -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
@@ -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
@@ -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

@@ -160,6 +160,12 @@ llama_context::llama_context(
         }
     }
 
+    if (params.ctx_type == LLAMA_CONTEXT_TYPE_MTP && params.ctx_other != nullptr &&
+        &params.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__);

@@ -342,6 +342,7 @@ struct llama_context {
     std::vector<swap_info> output_swaps;
 
     ggml_backend_sched_ptr sched;
+    llama_context * ctx_compute = nullptr;
 
     bool sched_need_reserve = true;
 

@@ -88,6 +88,8 @@ static void dummy_backend_buffer_clear(ggml_backend_buffer_t, uint8_t) {}
 struct dummy_backend {
     std::unique_ptr<dummy_backend_context> 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;
 }