llama: share SM75 MTP compute buffers on request

ggml/include/ggml-alloc.h

@@ -73,6 +73,9 @@ GGML_API bool ggml_gallocr_alloc_graph(ggml_gallocr_t galloc, struct ggml_cgraph
 
 GGML_API size_t ggml_gallocr_get_buffer_size(ggml_gallocr_t galloc, int buffer_id);
 
+// Discard allocated graphs before sharing device buffers. Callers must serialize all use of both allocators.
+GGML_API bool ggml_gallocr_share_buffers(ggml_gallocr_t dst, ggml_gallocr_t src);
+
 // Utils
 // Create a buffer and allocate all the tensors in a ggml_context
 // ggml_backend_alloc_ctx_tensors_from_buft_size returns the size of the buffer that would be allocated by ggml_backend_alloc_ctx_tensors_from_buft

ggml/include/ggml-backend.h

@@ -319,6 +319,9 @@ 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);
 
+    // Callers must serialize both schedulers, including input writes and output copies.
+    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

ggml/src/ggml-alloc.c

@@ -397,12 +397,16 @@ static size_t ggml_dyn_tallocr_max_size(struct ggml_dyn_tallocr * alloc, int chu
 
 struct vbuffer {
     ggml_backend_buffer_t chunks[GGML_VBUFFER_MAX_CHUNKS];
+    int refs;
 };
 
 static void ggml_vbuffer_free(struct vbuffer * buf) {
     if (buf == NULL) {
         return;
     }
+    if (--buf->refs > 0) {
+        return;
+    }
     for (int i = 0; i < GGML_VBUFFER_MAX_CHUNKS; ++i) {
         ggml_backend_buffer_free(buf->chunks[i]);
     }
@@ -426,6 +430,7 @@ static struct vbuffer * ggml_vbuffer_alloc(ggml_backend_buffer_type_t buft, cons
     if (buf == NULL) {
         return NULL;
     }
+    buf->refs = 1;
 
     for (int n = 0; n < talloc->n_chunks; n++) {
         size_t chunk_size = talloc->chunks[n]->max_size;
@@ -535,6 +540,59 @@ ggml_gallocr_t ggml_gallocr_new(ggml_backend_buffer_type_t buft) {
     return ggml_gallocr_new_n(&buft, 1);
 }
 
+bool ggml_gallocr_share_buffers(ggml_gallocr_t dst, ggml_gallocr_t src) {
+    if (!dst || !src || dst == src || dst->n_buffers != src->n_buffers) {
+        return false;
+    }
+    for (int i = 0; i < dst->n_buffers; ++i) {
+        if (dst->bufts[i] != src->bufts[i]) {
+            return false;
+        }
+        if (ggml_backend_buft_is_host(dst->bufts[i])) continue;
+        bool src_fits = true;
+        bool dst_fits = true;
+        for (int c = 0; c < GGML_VBUFFER_MAX_CHUNKS; ++c) {
+            const size_t a = src->buffers[i] ? ggml_vbuffer_chunk_size(src->buffers[i], c) : 0;
+            const size_t b = dst->buffers[i] ? ggml_vbuffer_chunk_size(dst->buffers[i], c) : 0;
+            src_fits &= a >= b;
+            dst_fits &= b >= a;
+        }
+        if (!src_fits && !dst_fits) return false;
+    }
+    bool shared = false;
+    for (int i = 0; i < dst->n_buffers; ++i) {
+        if (ggml_backend_buft_is_host(dst->bufts[i]) || (!src->buffers[i] && !dst->buffers[i])) {
+            continue;
+        }
+        bool duplicate = false;
+        for (int j = 0; j < i; ++j) {
+            if (dst->bufts[j] == dst->bufts[i]) {
+                dst->buffers[i] = dst->buffers[j];
+                src->buffers[i] = src->buffers[j];
+                duplicate = true;
+                break;
+            }
+        }
+        if (!duplicate) {
+            const size_t a = src->buffers[i] ? ggml_vbuffer_size(src->buffers[i]) : 0;
+            const size_t b = dst->buffers[i] ? ggml_vbuffer_size(dst->buffers[i]) : 0;
+            struct vbuffer * chosen = src->buffers[i] && (!dst->buffers[i] || a >= b) ? src->buffers[i] : dst->buffers[i];
+            if (dst->buffers[i] != chosen) {
+                ggml_vbuffer_free(dst->buffers[i]);
+                dst->buffers[i] = chosen;
+                ++chosen->refs;
+            }
+            if (src->buffers[i] != chosen) {
+                ggml_vbuffer_free(src->buffers[i]);
+                src->buffers[i] = chosen;
+                ++chosen->refs;
+            }
+        }
+        shared = true;
+    }
+    return shared;
+}
+
 void ggml_gallocr_free(ggml_gallocr_t galloc) {
     if (galloc == NULL) {
         return;

ggml/src/ggml-backend.cpp

@@ -1946,6 +1946,22 @@ 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 || !src || dst == src || dst->is_alloc || src->is_alloc || dst->n_copies != 1 || src->n_copies != 1 ||
+            dst->n_backends != src->n_backends) {
+        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;
+        }
+    }
+    ggml_backend_sched_synchronize(src);
+    ggml_backend_sched_synchronize(dst);
+    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

sm75/test-shared-compute.cpp

@@ -0,0 +1,94 @@
+#include "ggml.h"
+#include "ggml-alloc.h"
+#include "ggml-backend.h"
+
+#include <cstdio>
+#include <cstdlib>
+#include <vector>
+
+static void require(bool ok) {
+    if (!ok) std::exit(1);
+}
+
+struct graph {
+    ggml_context * ctx;
+    ggml_tensor * input;
+    ggml_tensor * output;
+    ggml_cgraph * gf;
+    graph(int size) {
+        ctx = ggml_init({1024 * 1024, nullptr, true});
+        input = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, size);
+        ggml_set_input(input);
+        output = ggml_scale(ctx, input, 2.0f);
+        ggml_set_output(output);
+        gf = ggml_new_graph(ctx);
+        ggml_build_forward_expand(gf, output);
+    }
+    ~graph() { ggml_free(ctx); }
+    void check(ggml_backend_t backend, float value) {
+        std::vector<float> data(ggml_nelements(input), value);
+        ggml_backend_tensor_set(input, data.data(), 0, ggml_nbytes(input));
+        require(ggml_backend_graph_compute(backend, gf) == GGML_STATUS_SUCCESS);
+        ggml_backend_tensor_get(output, data.data(), 0, ggml_nbytes(output));
+        for (float x : data) require(x == 2 * value);
+    }
+};
+
+int main() {
+    ggml_backend_load_all();
+    auto backend = ggml_backend_init_by_name("CUDA0", nullptr);
+    require(backend != nullptr);
+    auto buft = ggml_backend_get_default_buffer_type(backend);
+    for (int grow_dst : {0, 1}) {
+        auto src = ggml_gallocr_new(buft);
+        auto dst = ggml_gallocr_new(buft);
+        graph a(1024), b(1024), big(1024 * 1024);
+        require(ggml_gallocr_alloc_graph(src, a.gf));
+        require(ggml_gallocr_share_buffers(dst, src));
+        require(ggml_gallocr_alloc_graph(dst, b.gf));
+        require(a.input->buffer == b.input->buffer);
+        for (int i = 1; i <= 10; ++i) {
+            a.check(backend, i);
+            b.check(backend, -i);
+        }
+        require(ggml_gallocr_alloc_graph(grow_dst ? dst : src, big.gf));
+        require(big.input->buffer != (grow_dst ? a.input : b.input)->buffer);
+        big.check(backend, 0.5f);
+        (grow_dst ? a : b).check(backend, 4);
+        ggml_gallocr_free(grow_dst ? dst : src);
+        (grow_dst ? a : b).check(backend, 5);
+        ggml_gallocr_free(grow_dst ? src : dst);
+    }
+    for (int free_dst : {0, 1}) {
+        auto src = ggml_gallocr_new(buft);
+        auto dst = ggml_gallocr_new(buft);
+        graph a(1024), b(1024);
+        require(ggml_gallocr_alloc_graph(src, a.gf));
+        require(ggml_gallocr_share_buffers(dst, src));
+        require(ggml_gallocr_alloc_graph(dst, b.gf));
+        ggml_gallocr_free(free_dst ? dst : src);
+        (free_dst ? a : b).check(backend, 7);
+        ggml_gallocr_free(free_dst ? src : dst);
+    }
+    for (int larger_dst : {0, 1}) {
+        auto src = ggml_gallocr_new(buft);
+        auto dst = ggml_gallocr_new(buft);
+        {
+            graph a(larger_dst ? 1024 : 2048), b(larger_dst ? 2048 : 1024);
+            require(ggml_gallocr_alloc_graph(src, a.gf));
+            require(ggml_gallocr_alloc_graph(dst, b.gf));
+        }
+        require(ggml_gallocr_share_buffers(dst, src));
+        graph a(1024), b(1024);
+        require(ggml_gallocr_alloc_graph(src, a.gf));
+        require(ggml_gallocr_alloc_graph(dst, b.gf));
+        require(a.input->buffer == b.input->buffer);
+        a.check(backend, 9);
+        b.check(backend, 10);
+        ggml_gallocr_free(src);
+        b.check(backend, 11);
+        ggml_gallocr_free(dst);
+    }
+    ggml_backend_free(backend);
+    std::puts("PASS: shared buffers, alternating graphs, growth detachment, both free orders, larger-buffer selection");
+}

src/llama-context.cpp

@@ -147,6 +147,12 @@ llama_context::llama_context(
 
     cparams.ctx_other = nullptr;
 
+    const char * share_env = std::getenv("LLAMA_SM75_MTP_SHARED_COMPUTE");
+    if (share_env && strcmp(share_env, "1") == 0 && params.ctx_other &&
+            cparams.ctx_type == LLAMA_CONTEXT_TYPE_MTP && model.arch == LLM_ARCH_QWEN35 && cparams.n_seq_max == 1) {
+        compute_source = params.ctx_other;
+    }
+
     // TODO: more generic
     if (model.arch == LLM_ARCH_GEMMA4_ASSISTANT) {
         if (params.ctx_other == nullptr) {
@@ -484,6 +490,7 @@ llama_context::llama_context(
 }
 
 llama_context::~llama_context() {
+    auto compute_lock = lock_compute();
     // wait for any pending asynchronous copies into the output buffers before they are freed
     synchronize();
 
@@ -505,6 +512,25 @@ llama_context::~llama_context() {
         }
     }
     ggml_opt_free(opt_ctx);
+    if (compute_share) {
+        compute_share->active = nullptr;
+        for (auto & context : compute_share->contexts) {
+            if (context == this) context = nullptr;
+        }
+        sched.reset();
+    }
+}
+
+std::unique_lock<std::recursive_mutex> llama_context::lock_compute() {
+    if (!compute_share) {
+        return {};
+    }
+    std::unique_lock<std::recursive_mutex> lock(compute_share->mutex);
+    if (compute_share->active && compute_share->active != this) {
+        compute_share->active->synchronize();
+    }
+    compute_share->active = this;
+    return lock;
 }
 
 void llama_context::resolve_fused_ops(const llama_memory_context_i * mctx, uint32_t n_seqs) {
@@ -619,9 +645,15 @@ static int llama_graph_n_input_tensors(ggml_cgraph * gf) {
 }
 
 void llama_context::sched_reserve() {
+    auto compute_lock = lock_compute();
     if (!sched_need_reserve) {
         return;
     }
+    if (compute_share) {
+        for (auto * context : compute_share->contexts) {
+            if (context && context != this) compute_source = context;
+        }
+    }
 
     sched_need_reserve = false;
 
@@ -732,6 +764,45 @@ void llama_context::sched_reserve() {
         }
     }
 
+    if (compute_source) {
+        int devices = 0;
+        bool supported = (!compute_source->compute_share || compute_source->compute_share == compute_share) &&
+            compute_source->cparams.n_seq_max == 1;
+        for (auto backend : backend_ptrs) {
+            auto dev = ggml_backend_get_device(backend);
+            if (ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU) continue;
+            ++devices;
+            auto reg = ggml_backend_dev_backend_reg(dev);
+            auto supports = (bool (*)(ggml_backend_dev_t)) ggml_backend_reg_get_proc_address(reg, "ggml_cuda_gdn_txn_supported");
+            supported = supported && supports && supports(dev);
+        }
+        // Both graphs must be discarded before either reserved buffer can be replaced.
+        synchronize();
+        compute_source->synchronize();
+        gf_res_prev->reset();
+        gf_res_reserve->reset();
+        compute_source->gf_res_prev->reset();
+        compute_source->gf_res_reserve->reset();
+        ggml_backend_sched_reset(sched.get());
+        ggml_backend_sched_reset(compute_source->sched.get());
+        if (supported && devices == 1 && ggml_backend_sched_share_compute_buffers(sched.get(), compute_source->sched.get())) {
+            if (!compute_share) {
+                compute_share = std::make_shared<llama_compute_share>();
+                compute_share->contexts[0] = this;
+                compute_share->contexts[1] = compute_source;
+            }
+            compute_source->compute_share = compute_share;
+            compute_share->active = this;
+            for (size_t i = 0; i < compute_source->backend_ptrs.size(); ++i) {
+                compute_source->backend_buf_exp_size[i] = ggml_backend_sched_get_buffer_size(compute_source->sched.get(), compute_source->backend_ptrs[i]);
+            }
+            LLAMA_LOG_INFO("%s: shared MTP device compute buffers enabled\n", __func__);
+        } else {
+            LLAMA_LOG_INFO("%s: shared MTP compute unavailable, using separate buffers\n", __func__);
+        }
+        compute_source = nullptr;
+    }
+
     for (size_t i = 0; i < backend_ptrs.size(); ++i) {
         ggml_backend_t             backend = backend_ptrs[i];
         ggml_backend_buffer_type_t buft    = backend_buft[i];
@@ -1468,6 +1539,7 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll
 }
 
 int llama_context::encode(const llama_batch & batch_inp) {
+    auto compute_lock = lock_compute();
     // MTP hook batches carry both token (next-token id) and embd (h_nextn row),
     // so accept either present rather than requiring exactly one.
     GGML_ASSERT(batch_inp.token || batch_inp.embd);
@@ -1706,6 +1778,7 @@ static bool needs_raw_logits(const llama_ubatch & ubatch, const std::map<llama_s
 }
 
 int llama_context::decode(const llama_batch & batch_inp) {
+    auto compute_lock = lock_compute();
     // MTP hook batches carry both token (next-token id) and embd (h_nextn row),
     // so accept either present rather than requiring exactly one.
     GGML_ASSERT(batch_inp.token || batch_inp.embd);

src/llama-context.h

@@ -13,6 +13,7 @@
 
 #include <array>
 #include <map>
+#include <mutex>
 #include <vector>
 
 struct llama_model;
@@ -40,6 +41,12 @@ struct llama_memory_buffer {
 
 using llama_memory_buffers = std::map<ggml_backend_buffer_type_t, llama_memory_buffer>;
 
+struct llama_compute_share {
+    std::recursive_mutex mutex;
+    llama_context * active = nullptr;
+    llama_context * contexts[2] = {};
+};
+
 struct llama_context {
     // init scheduler and compute buffers, reserve worst-case graphs
     llama_context(
@@ -346,6 +353,9 @@ private:
     std::vector<swap_info> output_swaps;
 
     ggml_backend_sched_ptr sched;
+    llama_context * compute_source = nullptr;
+    std::shared_ptr<llama_compute_share> compute_share;
+    std::unique_lock<std::recursive_mutex> lock_compute();
 
     bool sched_need_reserve = true;