Commit a591cbe17 for llama.cpp
commit a591cbe17c497a3d18de55bc0c4e521964f296ab
Author: Xuan-Son Nguyen <son@huggingface.co>
Date: Tue Nov 10 11:45:09 2026 +0100
mtmd: share compute buffer among subcontext (#31945)
* shared compute buffer
* add load-time alloc for audio gen
diff --git a/tools/mtmd/README-dev.md b/tools/mtmd/README-dev.md
index 4ca410c67..b24c393a2 100644
--- a/tools/mtmd/README-dev.md
+++ b/tools/mtmd/README-dev.md
@@ -23,18 +23,21 @@ Overview of the `libmtmd` structure:
- `mtmd_context` (`mtmd.cpp`): the public handle
- `mtmd_models`: one mmproj file, loaded by `mtmd_model_load()`
- `mtmd_model` `v` / `a` / `gen_a` (`mtmd-model.h`): one per modality, holds hparams, weights (`ctx_data` + `buf`) and requirements
- - `mtmd_subcontext` `ctx_v` / `ctx_a` / `ctx_gen_a` (`mtmd-subcontext.h`): one per model, holds a const ref to it
- - backends + `sched`: GPU/CPU backends, compute buffers, resolved flash attn
+ - `mtmd_backends` (`mtmd-subcontext.h`): GPU/CPU backend instances, CPU threadpool and one scheduler, shared by the subcontexts
+ - `compute_mutex`: held by every public function that runs a compute pass, so only one pass at a time runs on the backends
+ - `mtmd_subcontext` `ctx_v` / `ctx_a` / `ctx_gen_a` (`mtmd-subcontext.h`): one per model, holds a const ref to it and a ref to the backends
+ - `buf_compute_meta` and the resolved flash attn of its graphs
- `mtmd_graph` (`mtmd-graph.h`, `models/`): built for each `encode()` call, temporary
- `mtmd_image_preprocessor` / `mtmd_audio_preprocessor` (`mtmd-image.h`, `mtmd-audio.h`): read the hparams of their model
- `mtmd_gen_state` / `mtmd_audio_stream`: public, one per sequence / stream
- `mtmd_subctx_gen_state` / `mtmd_subctx_enc_state`: allocated on the backend of a subcontext, see `gen_state_init()` / `enc_state_init()`
Rules of thumb:
-- `mtmd_model` is read-only once loaded. Anything that depends on the backend (scheduler, compute buffers, flash attn resolution, warmup) belongs to `mtmd_subcontext`
+- `mtmd_model` is read-only once loaded. Anything that depends on the backend (flash attn resolution, warmup) belongs to `mtmd_subcontext`
+- Like the pp/tg graphs of `llama_context`, the subcontexts share one scheduler: each warmup reserves its worst-case graph on it, so the compute buffers fit the largest one, and switching modality re-plans inside them
- A model declares its runtime needs (e.g. `max_nodes`, `flash_attn_auto`) in `load_hparams`, the subcontext applies them
-- Models must outlive their subcontexts, `mtmd_context` declares `models` first for that reason
-- Backends are owned by each subcontext for now; they will move to `mtmd_context` so that subcontexts can share them
+- Models and backends must outlive the subcontexts, `mtmd_context` declares them first for that reason
+- Locking happens only at the public API boundary (`mtmd.cpp`): a new public function that runs a graph must take `compute_mutex`. The internals (subcontexts, preprocessors, `mtmd_audio_log_mel()`) assume it is held and never lock
## Pipeline
diff --git a/tools/mtmd/mtmd-subcontext.cpp b/tools/mtmd/mtmd-subcontext.cpp
index 768eab562..f363cba95 100644
--- a/tools/mtmd/mtmd-subcontext.cpp
+++ b/tools/mtmd/mtmd-subcontext.cpp
@@ -32,7 +32,7 @@ static void mtmd_backend_cpu_set_n_threads(ggml_backend_t backend_cpu, int n_thr
}
}
-std::shared_ptr<ggml_threadpool> mtmd_threadpool_init(int n_threads) {
+static std::shared_ptr<ggml_threadpool> mtmd_threadpool_init(int n_threads) {
ggml_backend_dev_t dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU);
ggml_backend_reg_t reg = dev ? ggml_backend_dev_backend_reg(dev) : nullptr;
if (!reg || n_threads < 2) {
@@ -63,14 +63,9 @@ ggml_backend_dev_t mtmd_select_device(bool use_gpu, ggml_backend_dev_t device) {
return dev ? dev : ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_IGPU);
}
-mtmd_subcontext::mtmd_subcontext(const mtmd_model & model, const mtmd_subctx_params & params) :
- model(model),
- threadpool(params.threadpool),
- max_nodes(model.max_nodes),
- flash_attn_type(params.flash_attn_type) {
- if (flash_attn_type == MTMD_FLASH_ATTN_TYPE_AUTO && !model.flash_attn_auto) {
- flash_attn_type = MTMD_FLASH_ATTN_TYPE_DISABLED;
- }
+mtmd_backends::mtmd_backends(ggml_backend_dev_t device, int n_threads, int max_nodes,
+ ggml_backend_sched_eval_callback cb_eval, void * cb_eval_user_data) {
+ threadpool = mtmd_threadpool_init(n_threads);
backend_cpu = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr);
if (!backend_cpu) {
@@ -83,57 +78,70 @@ mtmd_subcontext::mtmd_subcontext(const mtmd_model & model, const mtmd_subctx_par
set_threadpool_fn(backend_cpu, threadpool.get());
}
}
- if (params.device != nullptr) {
- backend = ggml_backend_dev_init(params.device, nullptr);
+ if (device != nullptr) {
+ backend = ggml_backend_dev_init(device, nullptr);
if (!backend) {
ggml_backend_free(backend_cpu);
throw std::runtime_error(string_format("%s: failed to initialize \"%s\" backend\n",
- __func__, ggml_backend_dev_name(params.device)));
+ __func__, ggml_backend_dev_name(device)));
}
}
if (backend) {
LOG_INF("%s: using %s backend\n", __func__, ggml_backend_name(backend));
- backend_ptrs.push_back(backend);
- backend_buft.push_back(ggml_backend_get_default_buffer_type(backend));
+ ptrs.push_back(backend);
+ bufts.push_back(ggml_backend_get_default_buffer_type(backend));
} else {
backend = backend_cpu;
LOG_INF("%s: using CPU backend\n", __func__);
}
- backend_ptrs.push_back(backend_cpu);
- backend_buft.push_back(ggml_backend_get_default_buffer_type(backend_cpu));
+ ptrs.push_back(backend_cpu);
+ bufts.push_back(ggml_backend_get_default_buffer_type(backend_cpu));
+
+ sched.reset(ggml_backend_sched_new(ptrs.data(), bufts.data(), ptrs.size(), max_nodes, false, true));
+ if (cb_eval != nullptr) {
+ ggml_backend_sched_set_eval_callback(sched.get(), cb_eval, cb_eval_user_data);
+ }
+}
- sched.reset(
- ggml_backend_sched_new(backend_ptrs.data(), backend_buft.data(), backend_ptrs.size(), max_nodes, false, true)
- );
+mtmd_backends::~mtmd_backends() {
+ // before the backends it uses
+ sched.reset();
+ ggml_backend_free(backend);
+ if (backend != backend_cpu) {
+ ggml_backend_free(backend_cpu);
+ }
+}
+
+std::map<ggml_backend_dev_t, size_t> mtmd_backends::get_mem_compute() const {
+ std::map<ggml_backend_dev_t, size_t> res;
+ for (ggml_backend_t b : ptrs) {
+ res[ggml_backend_get_device(b)] += ggml_backend_sched_get_buffer_size(sched.get(), b);
+ }
+ return res;
+}
- if (params.cb_eval != nullptr) {
- ggml_backend_sched_set_eval_callback(sched.get(), params.cb_eval, params.cb_eval_user_data);
+mtmd_subcontext::mtmd_subcontext(const mtmd_model & model, mtmd_backends & backends, const mtmd_subctx_params & params) :
+ model(model),
+ backends(backends),
+ max_nodes(model.max_nodes),
+ flash_attn_type(params.flash_attn_type) {
+ if (flash_attn_type == MTMD_FLASH_ATTN_TYPE_AUTO && !model.flash_attn_auto) {
+ flash_attn_type = MTMD_FLASH_ATTN_TYPE_DISABLED;
}
debug_output_embeddings = std::getenv("MTMD_DEBUG_EMBEDDINGS") != nullptr;
buf_compute_meta.resize(max_nodes * ggml_tensor_overhead() + ggml_graph_overhead_custom(max_nodes, false));
- // TODO: fix warmup for audio gen
+ auto batch = get_dummy_batch();
if (model.modality != MTMD_MODALITY_GEN_AUDIO) {
// check batching support
- auto batch = get_dummy_batch();
support_batch = mtmd_get_graph_builder(this, batch)->support_batch();
-
- if (params.warmup) {
- warmup(batch);
- }
}
-}
-
-mtmd_subcontext::~mtmd_subcontext() {
- // the scheduler references the backends
- sched.reset();
- ggml_backend_free(backend);
- if (backend != backend_cpu) {
- ggml_backend_free(backend_cpu);
+ if (params.warmup) {
+ warmup(batch);
}
}
@@ -146,6 +154,11 @@ mtmd_image_f32_batch mtmd_subcontext::get_dummy_batch() const {
const int sz = hparams.warmup_image_size;
img.set_size({sz, sz}, false, false);
LOG_INF("%s: warmup with image size = %d x %d\n", __func__, sz, sz);
+ } else if (model.modality == MTMD_MODALITY_GEN_AUDIO) {
+ // TODO: improve this, move to per-model impl
+ const int n_embd = model.n_mmproj_embd();
+ img.set_size({n_embd, 1}, false, true);
+ LOG_INF("%s: warmup with one gen-audio step\n", __func__);
} else {
// GEMMA4UA uses n_mel_bins as a raw-waveform frame size (640), not a mel-bin count,
// so the [1, 256] bound only applies to FFT-based models.
@@ -170,7 +183,7 @@ void mtmd_subcontext::warmup(const mtmd_image_f32_batch & batch) {
if (!info.fattn && info.fattn_op) {
auto op = info.fattn_op;
LOG_WRN("%s: *****************************************************************\n", __func__);
- LOG_WRN("%s: WARNING: flash attention not supported by %s, memory usage will increase\n", __func__, ggml_backend_name(backend));
+ LOG_WRN("%s: WARNING: flash attention not supported by %s, memory usage will increase\n", __func__, ggml_backend_name(backends.backend));
LOG_WRN("%s: op params: \n", __func__);
static auto print_shape = [](const char * fn, const char * name, ggml_tensor * t) {
LOG_WRN("%s: %s: type = %s, ne = [%d %d %d %d], nb = [%d %d %d %d]\n", fn,
@@ -200,7 +213,7 @@ void mtmd_subcontext::warmup(const mtmd_image_f32_batch & batch) {
(flash_attn_type == MTMD_FLASH_ATTN_TYPE_ENABLED) ? "enabled" : "disabled");
// print ops that are not supported by the GPU backend (if there is one)
- if (backend && backend != backend_cpu) {
+ if (backends.backend != backends.backend_cpu) {
std::vector<support_info_op> unsupported_ops;
for (const auto & op : info.ops) {
if (!op.is_accel) {
@@ -211,7 +224,7 @@ void mtmd_subcontext::warmup(const mtmd_image_f32_batch & batch) {
LOG_WRN("%s: *****************************************************************\n", __func__);
LOG_WRN("%s: WARNING: the mmproj graph uses unsupported operators by the backend\n", __func__);
LOG_WRN("%s: the performance will be suboptimal \n", __func__);
- LOG_WRN("%s: list of unsupported ops (backend=%s):\n", __func__, ggml_backend_name(backend));
+ LOG_WRN("%s: list of unsupported ops (backend=%s):\n", __func__, ggml_backend_name(backends.backend));
for (const auto & op : unsupported_ops) {
LOG_WRN("%s: %16s: type = %s, ne = [%d %d %d %d]\n", __func__,
ggml_op_name(op.op->op),
@@ -229,22 +242,9 @@ void mtmd_subcontext::warmup(const mtmd_image_f32_batch & batch) {
mtmd_subcontext::support_info_graph mtmd_subcontext::reserve_compute_meta(const mtmd_image_f32_batch & batch) {
ggml_cgraph * gf = mtmd_get_graph_builder(this, batch)->build();
- ggml_backend_sched_reserve(sched.get(), gf);
-
- mem_compute.clear();
- 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];
- size_t size = ggml_backend_sched_get_buffer_size(sched.get(), backend);
- if (size > 1) {
- LOG_INF("%s: %10s compute buffer size = %8.2f MiB\n", __func__,
- ggml_backend_buft_name(buft),
- size / 1024.0 / 1024.0);
- }
- mem_compute[ggml_backend_get_device(backend)] += size;
- }
+ ggml_backend_sched_reserve(backends.sched.get(), gf);
- const int n_splits = ggml_backend_sched_get_n_splits(sched.get());
+ const int n_splits = ggml_backend_sched_get_n_splits(backends.sched.get());
const int n_nodes = ggml_graph_n_nodes(gf);
LOG_INF("%s: graph splits = %d, nodes = %d\n", __func__, n_splits, n_nodes);
@@ -259,7 +259,7 @@ mtmd_subcontext::support_info_graph mtmd_subcontext::reserve_compute_meta(const
for (int i = 0; i < ggml_graph_n_nodes(gf); i++) {
ggml_tensor * node = ggml_graph_node(gf, i);
res.ops.push_back({node, true});
- if (!ggml_backend_supports_op(backend, node)) {
+ if (!ggml_backend_supports_op(backends.backend, node)) {
res.ops.back().is_accel = false;
if (node->op == GGML_OP_FLASH_ATTN_EXT) {
res.fattn = false;
@@ -475,7 +475,7 @@ mtmd_subctx_gen_state * mtmd_subcontext::gen_state_init(uint32_t seed) {
ggml_format_name(t, "gen_state_%s", slot.name.c_str());
st->slots.push_back(t);
}
- st->buf.reset(ggml_backend_alloc_ctx_tensors(st->ctx_data.get(), backend));
+ st->buf.reset(ggml_backend_alloc_ctx_tensors(st->ctx_data.get(), backends.backend));
if (!st->buf) {
LOG_ERR("%s: failed to allocate the gen state\n", __func__);
return nullptr;
@@ -503,7 +503,7 @@ mtmd_subctx_enc_state * mtmd_subcontext::enc_state_init(int32_t n_frames_max) {
ggml_format_name(t, "enc_state_%s", slot.name.c_str());
st->slots.push_back(t);
}
- st->buf.reset(ggml_backend_alloc_ctx_tensors(st->ctx_data.get(), backend));
+ st->buf.reset(ggml_backend_alloc_ctx_tensors(st->ctx_data.get(), backends.backend));
if (!st->buf) {
LOG_ERR("%s: failed to allocate the encoder state\n", __func__);
return nullptr;
@@ -543,7 +543,7 @@ mtmd_subctx_enc_state * mtmd_subcontext::enc_state_init(int32_t n_frames_max) {
st->k.push_back(k);
st->v.push_back(v);
}
- st->buf.reset(ggml_backend_alloc_ctx_tensors(st->ctx_data.get(), backend));
+ st->buf.reset(ggml_backend_alloc_ctx_tensors(st->ctx_data.get(), backends.backend));
if (!st->buf) {
LOG_ERR("%s: failed to allocate the encoder state\n", __func__);
return nullptr;
@@ -588,14 +588,14 @@ bool mtmd_subcontext::encode(mtmd_subctx_encode_params * params) {
}
// build the inference graph
- ggml_backend_sched_reset(sched.get());
+ ggml_backend_sched_reset(backends.sched.get());
ggml_cgraph * gf = mtmd_get_graph_builder(this, imgs, params)->build();
// the last node is the embedding tensor, code2wav has no out_embd
// read it before scheduling: backends may reorder the nodes of graphs with several outputs (Metal)
ggml_tensor * embeddings = params->out_embd ? ggml_graph_node(gf, -1) : nullptr;
- if (!ggml_backend_sched_alloc_graph(sched.get(), gf)) {
+ if (!ggml_backend_sched_alloc_graph(backends.sched.get(), gf)) {
LOG_ERR("%s: failed to allocate compute graph\n", __func__);
return false;
}
@@ -2309,9 +2309,9 @@ bool mtmd_subcontext::encode(mtmd_subctx_encode_params * params) {
GGML_ABORT("Unknown projector type");
}
- mtmd_backend_cpu_set_n_threads(backend_cpu, params->n_threads);
+ mtmd_backend_cpu_set_n_threads(backends.backend_cpu, params->n_threads);
- auto status = ggml_backend_sched_graph_compute(sched.get(), gf);
+ auto status = ggml_backend_sched_graph_compute(backends.sched.get(), gf);
if (status != GGML_STATUS_SUCCESS) {
LOG_ERR("%s: ggml_backend_sched_graph_compute failed with error %d\n", __func__, status);
return false;
@@ -2709,11 +2709,3 @@ bool mtmd_subcontext::encode(mtmd_subctx_encode_params * params) {
return true;
}
-
-std::map<ggml_backend_dev_t, size_t> mtmd_subcontext::get_mem_usage() const {
- std::map<ggml_backend_dev_t, size_t> result = model.mem_weights;
- for (auto & [dev, size] : mem_compute) {
- result[dev] += size;
- }
- return result;
-}
diff --git a/tools/mtmd/mtmd-subcontext.h b/tools/mtmd/mtmd-subcontext.h
index d757c9b6e..7e6d8232a 100644
--- a/tools/mtmd/mtmd-subcontext.h
+++ b/tools/mtmd/mtmd-subcontext.h
@@ -16,8 +16,7 @@
// !!! Internal header, to be used by mtmd only !!!
-// a subcontext runs the graphs of one mtmd_model (one modality)
-// it owns the backends and the compute buffers, the weights are owned by the model
+// runs the graphs of one mtmd_model on the shared backends and scheduler
enum mtmd_flash_attn_type {
MTMD_FLASH_ATTN_TYPE_AUTO = -1,
@@ -26,22 +25,40 @@ enum mtmd_flash_attn_type {
};
struct mtmd_subctx_params {
- ggml_backend_dev_t device; // nullptr means CPU, see mtmd_select_device()
- std::shared_ptr<ggml_threadpool> threadpool; // CPU threadpool shared by the subcontexts, can be nullptr
enum mtmd_flash_attn_type flash_attn_type;
bool warmup;
- ggml_backend_sched_eval_callback cb_eval;
- void * cb_eval_user_data;
};
// the device the weights and the graphs are placed on, nullptr means CPU
ggml_backend_dev_t mtmd_select_device(bool use_gpu, ggml_backend_dev_t device);
-// nullptr if n_threads < 2 or the CPU backend has no threadpool support
-std::shared_ptr<ggml_threadpool> mtmd_threadpool_init(int n_threads);
+// shared by the subcontexts, not thread-safe
+struct mtmd_backends {
+ ggml_backend_t backend = nullptr; // backend_cpu if no accelerator
+ ggml_backend_t backend_cpu = nullptr;
+
+ // for ggml_backend_sched_new(), backend_cpu last
+ std::vector<ggml_backend_t> ptrs;
+ std::vector<ggml_backend_buffer_type_t> bufts;
+
+ std::shared_ptr<ggml_threadpool> threadpool; // can be nullptr
+
+ // its buffers fit the largest graph reserved by the subcontexts
+ ggml_backend_sched_ptr sched;
+
+ // device nullptr means CPU
+ mtmd_backends(ggml_backend_dev_t device, int n_threads, int max_nodes,
+ ggml_backend_sched_eval_callback cb_eval, void * cb_eval_user_data);
+ ~mtmd_backends();
+
+ std::map<ggml_backend_dev_t, size_t> get_mem_compute() const;
+
+ mtmd_backends(const mtmd_backends &) = delete;
+ mtmd_backends & operator=(const mtmd_backends &) = delete;
+};
// compute the log-mel of n_frames frames, the waveform must already be padded and window must hold one whole frame
-// out has the layout [n_frames, n_mel]; like the encoders, must not run at the same time as another pass on config.sched
+// out has the layout [n_frames, n_mel]; must not run at the same time as another pass
struct mtmd_audio_mel_config;
bool mtmd_audio_log_mel(const mtmd_audio_mel_config & config, const float * samples, const float * window, int64_t n_frames, float * out);
@@ -134,38 +151,26 @@ struct mtmd_subctx_encode_params {
struct mtmd_subcontext {
const mtmd_model & model;
+ mtmd_backends & backends;
std::vector<uint8_t> buf_compute_meta;
- std::vector<ggml_backend_t> backend_ptrs;
- std::vector<ggml_backend_buffer_type_t> backend_buft;
-
- std::shared_ptr<ggml_threadpool> threadpool; // shared by the subcontexts of a model
-
- ggml_backend_t backend = nullptr;
- ggml_backend_t backend_cpu = nullptr;
-
int max_nodes = 8192;
- ggml_backend_sched_ptr sched;
mtmd_flash_attn_type flash_attn_type = MTMD_FLASH_ATTN_TYPE_AUTO;
bool is_allocated = false;
bool debug_output_embeddings = false;
- // for measuring memory usage
- std::map<ggml_backend_dev_t, size_t> mem_compute;
-
bool support_batch = false;
- // the model must outlive the subcontext
- mtmd_subcontext(const mtmd_model & model, const mtmd_subctx_params & params);
- ~mtmd_subcontext();
+ // model and backends must outlive the subcontext
+ mtmd_subcontext(const mtmd_model & model, mtmd_backends & backends, const mtmd_subctx_params & params);
+
+ bool encode(mtmd_subctx_encode_params * params);
// reserve the compute buffers for the given batch
void warmup(const mtmd_image_f32_batch & batch);
- bool encode(mtmd_subctx_encode_params * params);
-
// TODO: remove encode_image() and always use batched version
bool encode_image(int n_threads, const mtmd_image_f32 * img, std::vector<float> & out_vec);
bool encode_batch(int n_threads, const mtmd_image_f32_batch * imgs, std::vector<float> & out_batch_embd);
@@ -176,9 +181,6 @@ struct mtmd_subcontext {
// seed UINT32_MAX means random, free the result with delete
mtmd_subctx_gen_state * gen_state_init(uint32_t seed);
- // weights and compute buffers
- std::map<ggml_backend_dev_t, size_t> get_mem_usage() const;
-
private:
struct support_info_op {
ggml_tensor * op;
diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp
index aa886b6c7..d95cad4f7 100644
--- a/tools/mtmd/mtmd.cpp
+++ b/tools/mtmd/mtmd.cpp
@@ -28,6 +28,7 @@
#include <cstring>
#include <climits>
#include <map>
+#include <mutex>
#include <type_traits>
#include <vector>
@@ -593,6 +594,10 @@ static bool cbx_is_mtl(const mtmd_model & model) {
struct mtmd_context {
// must outlive the subcontexts below
mtmd_models models;
+ std::unique_ptr<mtmd_backends> backends; // shared by the subcontexts
+
+ // held by every public function that runs a graph, the internals never lock
+ std::mutex compute_mutex;
std::unique_ptr<mtmd_subcontext> ctx_v; // vision
std::unique_ptr<mtmd_subcontext> ctx_a; // audio
@@ -709,23 +714,37 @@ struct mtmd_context {
throw std::runtime_error(string_format("Failed to load mmproj model from %s\n", mmproj_fname));
}
+ int max_nodes = 0;
+ for (const auto * m : { models.v.get(), models.a.get(), models.gen_a.get() }) {
+ if (m) {
+ max_nodes = std::max(max_nodes, m->max_nodes);
+ }
+ }
+
+ // no_alloc only measures memory, nothing is computed
+ backends = std::make_unique<mtmd_backends>(device, no_alloc ? 1 : ctx_params.n_threads, max_nodes,
+ ctx_params.cb_eval, ctx_params.cb_eval_user_data);
+
+ // like llama_context::sched_reserve(), each warmup reserves its graph on the shared scheduler
mtmd_subctx_params sctx_params {
- /* device */ device,
- // no_alloc only measures memory, nothing is computed
- /* threadpool */ mtmd_threadpool_init(no_alloc ? 1 : ctx_params.n_threads),
/* flash_attn_type */ mtmd_to_flash_attn_type(ctx_params.flash_attn_type),
/* warmup */ ctx_params.warmup,
- /* cb_eval */ ctx_params.cb_eval,
- /* cb_eval_user_data */ ctx_params.cb_eval_user_data,
};
if (models.v) {
- ctx_v = std::make_unique<mtmd_subcontext>(*models.v, sctx_params);
+ ctx_v = std::make_unique<mtmd_subcontext>(*models.v, *backends, sctx_params);
}
if (models.a) {
- ctx_a = std::make_unique<mtmd_subcontext>(*models.a, sctx_params);
+ ctx_a = std::make_unique<mtmd_subcontext>(*models.a, *backends, sctx_params);
}
if (models.gen_a) {
- ctx_gen_a = std::make_unique<mtmd_subcontext>(*models.gen_a, sctx_params);
+ ctx_gen_a = std::make_unique<mtmd_subcontext>(*models.gen_a, *backends, sctx_params);
+ }
+ for (size_t i = 0; i < backends->ptrs.size(); ++i) {
+ const size_t size = ggml_backend_sched_get_buffer_size(backends->sched.get(), backends->ptrs[i]);
+ if (size > 1) {
+ LOG_INF("%s: %10s compute buffer size = %8.2f MiB\n", __func__,
+ ggml_backend_buft_name(backends->bufts[i]), size / 1024.0 / 1024.0);
+ }
}
// if both vision and audio mmproj are present, we need to validate their n_embd
@@ -1261,7 +1280,7 @@ struct mtmd_context {
// initialize audio preprocessor
if (audio_preproc) {
audio_preproc->mel_config.n_threads = n_threads;
- audio_preproc->mel_config.sched = ctx_a->sched.get();
+ audio_preproc->mel_config.sched = backends->sched.get();
audio_preproc->initialize();
}
}
@@ -2361,6 +2380,7 @@ static int32_t mtmd_encode_chunk_impl(mtmd_context * ctx, const mtmd_input_chunk
int32_t mtmd_encode_chunk(mtmd_context * ctx, const mtmd_input_chunk * chunk) {
// this is the non-batching version
+ std::lock_guard<std::mutex> lock(ctx->compute_mutex);
try {
return mtmd_encode_chunk_impl(ctx, chunk, ctx->out_embd);
} catch (const std::exception & e) {
@@ -2370,6 +2390,7 @@ int32_t mtmd_encode_chunk(mtmd_context * ctx, const mtmd_input_chunk * chunk) {
}
int32_t mtmd_encode(mtmd_context * ctx, const mtmd_image_tokens * image_tokens) {
+ std::lock_guard<std::mutex> lock(ctx->compute_mutex);
try {
return mtmd_encode_impl(ctx, image_tokens, ctx->out_embd);
} catch (const std::exception & e) {
@@ -3140,6 +3161,7 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, mtmd_gen_state *
}
int32_t mtmd_gen_audio_process(mtmd_context * ctx, mtmd_gen_state * state, const struct mtmd_gen_inp * inp, struct mtmd_gen_out * out) {
+ std::lock_guard<std::mutex> lock(ctx->compute_mutex);
try {
return mtmd_gen_audio_process_impl(ctx, state, inp, out);
} catch (const std::exception & e) {
@@ -3248,6 +3270,7 @@ static int32_t mtmd_batch_encode_impl(mtmd_batch * batch) {
}
int32_t mtmd_batch_encode(mtmd_batch * batch) {
+ std::lock_guard<std::mutex> lock(batch->ctx->compute_mutex);
try {
return mtmd_batch_encode_impl(batch);
} catch (const std::exception & e) {
@@ -3861,6 +3884,7 @@ static void mtmd_debug_encode_impl(mtmd_context * ctx, mtmd_subcontext * sctx, m
}
void mtmd_debug_encode_image(mtmd_context * ctx, const std::vector<std::vector<float>> & image) {
+ std::lock_guard<std::mutex> lock(ctx->compute_mutex);
if (!ctx->ctx_v) {
LOG_ERR("%s: model does not support vision input\n", __func__);
return;
@@ -3879,6 +3903,7 @@ void mtmd_debug_encode_image(mtmd_context * ctx, const std::vector<std::vector<f
}
void mtmd_debug_encode_audio(mtmd_context * ctx, const std::vector<float> & input) {
+ std::lock_guard<std::mutex> lock(ctx->compute_mutex);
if (!ctx->ctx_a) {
LOG_ERR("%s: model does not support audio input\n", __func__);
return;
@@ -3923,6 +3948,7 @@ void mtmd_debug_preprocess_image(mtmd_context * ctx, const std::vector<uint8_t>
}
void mtmd_debug_preprocess_audio(mtmd_context * ctx, const std::vector<float> & samples) {
+ std::lock_guard<std::mutex> lock(ctx->compute_mutex);
if (!ctx->ctx_a || !ctx->audio_preproc) {
LOG_ERR("%s: model does not support audio input\n", __func__);
return;
@@ -3963,20 +3989,14 @@ mtmd_memory_usage mtmd_get_memory_usage(const char * mmproj_fname,
mtmd_log_set(stub_log_callback, nullptr); // suppress logging
ctx.reset(new mtmd_context(mmproj_fname, nullptr, ctx_params, true));
mtmd_log_set(saved_log_callback, saved_log_user_data); // restore log callback
- std::map<ggml_backend_dev_t, size_t> total_mem;
- auto merge = [&](const struct mtmd_subcontext * c) {
- for (auto & [dev, size] : c->get_mem_usage()) {
- total_mem[dev] += size;
+ // compute buffers are shared, weights are per model
+ std::map<ggml_backend_dev_t, size_t> total_mem = ctx->backends->get_mem_compute();
+ for (const auto * m : { ctx->models.v.get(), ctx->models.a.get(), ctx->models.gen_a.get() }) {
+ if (m) {
+ for (auto & [dev, size] : m->mem_weights) {
+ total_mem[dev] += size;
+ }
}
- };
- if (ctx->ctx_v) {
- merge(ctx->ctx_v.get());
- }
- if (ctx->ctx_a) {
- merge(ctx->ctx_a.get());
- }
- if (ctx->ctx_gen_a) {
- merge(ctx->ctx_gen_a.get());
}
mtmd_memory_usage res;
res.backend_mem_usage = std::move(total_mem);
@@ -4200,10 +4220,12 @@ int32_t mtmd_audio_stream_flush(mtmd_audio_stream * s, bool end) {
if (s->ended) {
return 1;
}
+ std::lock_guard<std::mutex> lock(s->ctx->compute_mutex);
return s->flush(end);
}
int32_t mtmd_audio_stream_step(mtmd_audio_stream * s, llama_token last_token) {
+ std::lock_guard<std::mutex> lock(s->ctx->compute_mutex);
return s->step(last_token);
}
diff --git a/tools/mtmd/mtmd.h b/tools/mtmd/mtmd.h
index e023f2c31..e14f345a6 100644
--- a/tools/mtmd/mtmd.h
+++ b/tools/mtmd/mtmd.h
@@ -353,6 +353,7 @@ DEPRECATED(MTMD_API int32_t mtmd_encode(mtmd_context * ctx, const mtmd_image_tok
"use mtmd_encode_chunk() instead");
// text chunk will be ignored silently, only media chunk will be encoded
+// the output is stored in ctx: do not call it from several threads at once
// returns 0 on success
// returns 1 on generic error
MTMD_API int32_t mtmd_encode_chunk(mtmd_context * ctx,
@@ -377,6 +378,7 @@ MTMD_API void mtmd_batch_free(mtmd_batch * batch);
// returns 3 if it cannot be batched with the existing chunks in the batch
MTMD_API int32_t mtmd_batch_add_chunk(mtmd_batch * batch, const mtmd_input_chunk * chunk);
+// thread-safe, the output is stored in the batch
// returns 0 on success
// returns 1 on generic error
MTMD_API int32_t mtmd_batch_encode(mtmd_batch * batch);
@@ -431,11 +433,13 @@ MTMD_API int32_t mtmd_audio_stream_push(mtmd_audio_stream * stream, const float
// pad the stream so that all pending rows can be produced
// end = false: the stream stays open (use it on a pause, to get a partial transcript)
// end = true: no more push() is accepted
+// thread-safe across streams
// returns 0 on success, 1 if the stream has ended
MTMD_API int32_t mtmd_audio_stream_flush(mtmd_audio_stream * stream, bool end);
// produce the next decoder rows
// last_token is the token sampled from the previous row, LLAMA_TOKEN_NULL before the first sample
+// thread-safe across streams
// return value:
// > 0 number of rows ready, read them with mtmd_audio_stream_get_output_embd()
// 0 more audio is needed
@@ -557,6 +561,7 @@ MTMD_API struct mtmd_gen_inp mtmd_gen_inp_default(const mtmd_context * ctx);
// run one step and update the state in place
// the caller drives the backbone and collects the audio frames
+// thread-safe across gen states
MTMD_API int32_t mtmd_gen_audio_process(mtmd_context * ctx,
mtmd_gen_state * state,
const struct mtmd_gen_inp * inp,