diff --git a/ggml/include/ggml-alloc.h b/ggml/include/ggml-alloc.h index a7926a21a..bb547d7c6 100644 --- a/ggml/include/ggml-alloc.h +++ b/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 diff --git a/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h index cc3f8cd36..bd26fb45f 100644 --- a/ggml/include/ggml-backend.h +++ b/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 diff --git a/ggml/src/ggml-alloc.c b/ggml/src/ggml-alloc.c index a71838eaf..ecb24fa34 100644 --- a/ggml/src/ggml-alloc.c +++ b/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; diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index 20bf96501..8624f1fe6 100644 --- a/ggml/src/ggml-backend.cpp +++ b/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 diff --git a/sm75/test-shared-compute.cpp b/sm75/test-shared-compute.cpp new file mode 100644 index 000000000..e8add2532 --- /dev/null +++ b/sm75/test-shared-compute.cpp @@ -0,0 +1,94 @@ +#include "ggml.h" +#include "ggml-alloc.h" +#include "ggml-backend.h" + +#include +#include +#include + +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 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"); +} diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 5893eb1c8..37e68d626 100644 --- a/src/llama-context.cpp +++ b/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 llama_context::lock_compute() { + if (!compute_share) { + return {}; + } + std::unique_lock 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(); + 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 #include +#include #include struct llama_model; @@ -40,6 +41,12 @@ struct llama_memory_buffer { using llama_memory_buffers = std::map; +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 output_swaps; ggml_backend_sched_ptr sched; + llama_context * compute_source = nullptr; + std::shared_ptr compute_share; + std::unique_lock lock_compute(); bool sched_need_reserve = true;