From 23ac1168a8febb159be42faff4bd576c4cc82022 Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Tue, 10 Nov 2026 08:17:34 -0600 Subject: [PATCH] context: measure host and device allocations on the loaded model Assisted-by: Codex --- include/llama.h | 11 ++++++++ src/CMakeLists.txt | 1 + src/llama-context-memory.cpp | 33 +++++++++++++++++++++++ src/llama-context.cpp | 45 +++++++++++++++++++++----------- src/llama-context.h | 8 +++++- src/llama-kv-cache-dsa-iswa.cpp | 6 ++--- src/llama-kv-cache-dsa-iswa.h | 2 +- src/llama-kv-cache-dsa.cpp | 6 ++--- src/llama-kv-cache-dsa.h | 2 +- src/llama-kv-cache-dsv4.cpp | 34 +++++++++++++++--------- src/llama-kv-cache-dsv4.h | 5 ++-- src/llama-kv-cache-iswa.cpp | 10 +++---- src/llama-kv-cache-iswa.h | 4 +-- src/llama-kv-cache-msa.cpp | 6 ++--- src/llama-kv-cache-msa.h | 2 +- src/llama-kv-cache.cpp | 8 +++--- src/llama-kv-cache.h | 3 ++- src/llama-memory-hybrid-idx.cpp | 6 ++--- src/llama-memory-hybrid-idx.h | 2 +- src/llama-memory-hybrid-iswa.cpp | 6 ++--- src/llama-memory-hybrid-iswa.h | 2 +- src/llama-memory-hybrid.cpp | 6 ++--- src/llama-memory-hybrid.h | 2 +- src/llama-memory-recurrent.cpp | 22 +++++++++++----- src/llama-memory-recurrent.h | 3 ++- src/llama-memory.h | 1 + src/llama-model.cpp | 30 ++++++++++----------- 27 files changed, 177 insertions(+), 89 deletions(-) create mode 100644 src/llama-context-memory.cpp diff --git a/include/llama.h b/include/llama.h index 84bdebaf771a..d6bb8b5654a4 100644 --- a/include/llama.h +++ b/include/llama.h @@ -660,6 +660,17 @@ extern "C" { // Get the model file type (quantization), e.g. LLAMA_FTYPE_MOSTLY_Q8_0 LLAMA_API enum llama_ftype llama_model_ftype(const struct llama_model * model); + // Context-owned allocations only. Host buffers (including pinned host memory) are + // reported separately from device buffers. Buffer names are borrowed for the callback. + typedef void (*llama_context_memory_callback)(const char * buffer, bool host, uint64_t bytes, void * data); + // Build the same context layout and graph without allocating tensor payloads. + // Uses the loaded model's placement; never reloads or mutates its weights. + LLAMA_API bool llama_context_measure_memory( + struct llama_model * model, struct llama_context_params params, + llama_context_memory_callback callback, void * data); + LLAMA_API void llama_context_visit_memory( + const struct llama_context * ctx, llama_context_memory_callback callback, void * data); + // Returns the total size of all the tensors in the model in bytes LLAMA_API uint64_t llama_model_size(const struct llama_model * model); diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 7dcd7a5a3e23..fb49c88fa748 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -16,6 +16,7 @@ add_library(llama llama-batch.cpp llama-chat.cpp llama-context.cpp + llama-context-memory.cpp llama-context-opt.cpp llama-context-opt-measure.cpp llama-context-learn.cpp diff --git a/src/llama-context-memory.cpp b/src/llama-context-memory.cpp new file mode 100644 index 000000000000..e3773d4fcf75 --- /dev/null +++ b/src/llama-context-memory.cpp @@ -0,0 +1,33 @@ +#include "llama-context.h" +#include "llama-model.h" +#include + +void llama_context::visit_context_memory(llama_context_memory_callback callback, void * data) const { + for (const auto & [buft, size] : memory_breakdown()) { + callback(ggml_backend_buft_name(buft), ggml_backend_buft_is_host(buft), + size.context + size.compute, data); + } + if (output_buft) { + callback(ggml_backend_buft_name(output_buft), ggml_backend_buft_is_host(output_buft), output_bytes, data); + } +} + +bool llama_context_measure_memory(llama_model * model, llama_context_params params, + llama_context_memory_callback callback, void * data) { + if (!model || !callback || params.n_ctx == 0 || params.n_batch == 0 || params.n_ubatch == 0) { + return false; + } + try { + std::unique_ptr measured(llama_context_create(model, params, true)); + if (!measured) { return false; } + measured->visit_context_memory(callback, data); + return true; + } catch (const std::exception & error) { + LLAMA_LOG_ERROR("%s: %s\n", __func__, error.what()); + return false; + } +} + +void llama_context_visit_memory(const llama_context * ctx, llama_context_memory_callback callback, void * data) { + if (ctx && callback) { ctx->visit_context_memory(callback, data); } +} diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 83f195043c6c..50cd11c17255 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -78,8 +78,8 @@ static const llm_fused_op_probe llm_fused_op_dsv4_hc_post_probe = { llama_context::llama_context( const llama_model & model, - llama_context_params params) : - model(model), + llama_context_params params, bool measure_only) : + no_alloc(measure_only || model.hparams.no_alloc), model(model), cvec(std::make_unique()), loras(std::make_unique()), balloc(std::make_unique(model.hparams.n_pos_per_embd())) { @@ -368,7 +368,8 @@ llama_context::llama_context( // graph outputs buffer { - if (output_reserve(params.n_seq_max) < params.n_seq_max) { + const uint32_t output_count = no_alloc ? cparams.n_batch : params.n_seq_max; + if (output_reserve(output_count) < output_count) { throw std::runtime_error("failed to reserve initial output buffer"); } @@ -381,6 +382,7 @@ llama_context::llama_context( // init the memory module if (!hparams.vocab_only) { llama_memory_params params_mem = { + /*.no_alloc =*/ no_alloc, /*.type_k =*/ params.type_k, /*.type_v =*/ params.type_v, /*.swa_full =*/ params.swa_full, @@ -478,7 +480,7 @@ llama_context::~llama_context() { // wait for any pending asynchronous copies into the output buffers before they are freed synchronize(); - if (!model.hparams.no_alloc) { + if (!no_alloc) { 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]; @@ -639,13 +641,13 @@ void llama_context::sched_reserve() { // reserve pp (prompt processing) graph first so that buffers are only allocated once { auto * gf = graph_reserve(n_tokens, n_seqs, n_outputs_pp, mctx.get(), - model.hparams.no_alloc, model.hparams.no_alloc ? backend_buf_exp_size.data() : nullptr); + no_alloc, no_alloc ? backend_buf_exp_size.data() : nullptr); if (!gf) { if (cparams.pipeline_parallel) { LLAMA_LOG_WARN("%s: compute buffer allocation failed, retrying without pipeline parallelism\n", __func__); cparams.pipeline_parallel = false; sched.reset(ggml_backend_sched_new(backend_ptrs.data(), backend_buft.data(), backend_ptrs.size(), max_nodes, false, cparams.op_offload)); - gf = graph_reserve(n_tokens, n_seqs, n_outputs_pp, mctx.get()); + gf = graph_reserve(n_tokens, n_seqs, n_outputs_pp, mctx.get(), no_alloc, no_alloc ? backend_buf_exp_size.data() : nullptr); } if (!gf) { throw std::runtime_error("failed to allocate compute pp buffers"); @@ -658,7 +660,7 @@ void llama_context::sched_reserve() { // reserve with tg (token generation) graph to get the number of splits and nodes { - auto * gf = graph_reserve(n_seqs, n_seqs, n_seqs, mctx.get(), model.hparams.no_alloc); + auto * gf = graph_reserve(n_seqs, n_seqs, n_seqs, mctx.get(), no_alloc, no_alloc ? backend_buf_exp_size.data() : nullptr); if (!gf) { throw std::runtime_error("failed to allocate compute tg buffers"); } @@ -673,7 +675,7 @@ void llama_context::sched_reserve() { // // auto * gf = graph_reserve(n_tokens, 1, n_tokens, mctx.get()); // - auto * gf = graph_reserve(n_tokens, n_seqs, n_outputs_pp, mctx.get(), model.hparams.no_alloc); + auto * gf = graph_reserve(n_tokens, n_seqs, n_outputs_pp, mctx.get(), no_alloc, no_alloc ? backend_buf_exp_size.data() : nullptr); if (!gf) { throw std::runtime_error("failed to allocate compute pp buffers"); } @@ -682,7 +684,7 @@ void llama_context::sched_reserve() { 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]; - if (!model.hparams.no_alloc) { + if (!no_alloc) { backend_buf_exp_size[i] = ggml_backend_sched_get_buffer_size(sched.get(), backend); } if (backend_buf_exp_size[i] > 1) { @@ -2211,6 +2213,13 @@ uint32_t llama_context::output_reserve(int32_t n_outputs) { if (output_dev_host_buft) { buft = output_dev_host_buft; } + output_buft = buft; + output_bytes = new_size; + if (no_alloc) { + buf_output.reset(ggml_backend_buft_alloc_buffer(buft, 0)); + if (!buf_output) { return 0; } + return n_outputs_max; + } buf_output.reset(ggml_backend_buft_alloc_buffer(buft, new_size)); if (buf_output == nullptr) { LLAMA_LOG_ERROR("%s: failed to allocate output buffer of size %.2f MiB\n", __func__, new_size / (1024.0 * 1024.0)); @@ -2546,7 +2555,9 @@ ggml_cgraph * llama_context::graph_reserve( // initialize scheduler with the specified graph if (split_only) { if (sizes) { - ggml_backend_sched_reserve_size(sched.get(), gf, sizes); + std::vector measured(backend_ptrs.size(), 0); + ggml_backend_sched_reserve_size(sched.get(), gf, measured.data()); + for (size_t i = 0; i < measured.size(); ++i) { sizes[i] = std::max(sizes[i], measured[i]); } } else { ggml_backend_sched_split_graph(sched.get(), gf); } @@ -3401,11 +3412,11 @@ llama_memory_breakdown llama_context::memory_breakdown() const { ret[buft].context += size; } } - if (model.hparams.no_alloc) { + if (no_alloc) { for (size_t i = 0; i < backends.size(); ++i) { ggml_backend_t backend = backends[i].get(); ggml_backend_buffer_type_t buft = ggml_backend_sched_get_buffer_type(sched.get(), backend); - ret[buft].compute += backend_buf_exp_size[i]; + ret[buft].compute = std::max(ret[buft].compute, backend_buf_exp_size[i]); } } else { std::vector reported; @@ -3475,9 +3486,9 @@ llama_context_params llama_context_default_params() { return result; } -llama_context * llama_init_from_model( +llama_context * llama_context_create( llama_model * model, - llama_context_params params) { + llama_context_params params, bool measure_only) { if (!model) { LLAMA_LOG_ERROR("%s: model cannot be NULL\n", __func__); return nullptr; @@ -3562,7 +3573,7 @@ llama_context * llama_init_from_model( } try { - auto * ctx = new llama_context(*model, params); + auto * ctx = new llama_context(*model, params, measure_only); return ctx; } catch (const std::exception & err) { LLAMA_LOG_ERROR("%s: failed to initialize the context: %s\n", __func__, err.what()); @@ -3571,6 +3582,10 @@ llama_context * llama_init_from_model( return nullptr; } +llama_context * llama_init_from_model(llama_model * model, llama_context_params params) { + return llama_context_create(model, params, false); +} + // deprecated llama_context * llama_new_context_with_model( llama_model * model, diff --git a/src/llama-context.h b/src/llama-context.h index 645be5cf9e69..d7c690cc1802 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -44,7 +44,7 @@ struct llama_context { // init scheduler and compute buffers, reserve worst-case graphs llama_context( const llama_model & model, - llama_context_params params); + llama_context_params params, bool measure_only = false); ~llama_context(); @@ -198,6 +198,7 @@ struct llama_context { void perf_reset(); llama_memory_breakdown memory_breakdown() const; + void visit_context_memory(llama_context_memory_callback callback, void * data) const; // // training @@ -380,6 +381,7 @@ struct llama_context { // members // + const bool no_alloc; const llama_model & model; llama_cparams cparams; @@ -506,6 +508,8 @@ struct llama_context { // host buffer for the model output (logits and embeddings) ggml_backend_buffer_ptr buf_output; + size_t output_bytes = 0; + ggml_backend_buffer_type_t output_buft = nullptr; // keep copies of the per-sequence memory on the device std::map mem_storage; @@ -529,3 +533,5 @@ struct llama_context { mutable int32_t n_reused = 0; // number of times the previous graph was reused }; + +llama_context * llama_context_create(llama_model * model, llama_context_params params, bool measure_only); diff --git a/src/llama-kv-cache-dsa-iswa.cpp b/src/llama-kv-cache-dsa-iswa.cpp index 90842dcd3236..d7a6d1056e66 100644 --- a/src/llama-kv-cache-dsa-iswa.cpp +++ b/src/llama-kv-cache-dsa-iswa.cpp @@ -12,7 +12,7 @@ // llama_kv_cache_dsa_iswa::llama_kv_cache_dsa_iswa( - const llama_model & model, + bool no_alloc, const llama_model & model, ggml_type type_k, ggml_type type_v, bool v_trans, @@ -63,14 +63,14 @@ llama_kv_cache_dsa_iswa::llama_kv_cache_dsa_iswa( LLAMA_LOG_INFO("%s: creating DSA KV cache, size = %u cells\n", __func__, size_dsa); kv_dsa = std::make_unique( - model, type_k, type_v, + no_alloc, model, type_k, type_v, v_trans, offload, unified, size_dsa, n_seq_max, n_pad, 0, LLAMA_SWA_TYPE_NONE, filter_dsa, filter_lid, reuse); LLAMA_LOG_INFO("%s: creating SWA KV cache, size = %u cells\n", __func__, size_swa); kv_swa = std::make_unique( - model, hparams, type_k, type_v, + no_alloc, model, hparams, type_k, type_v, v_trans, offload, unified, size_swa, n_seq_max, n_pad, hparams.n_swa, hparams.swa_type, nullptr, filter_swa, reuse, nullptr); } diff --git a/src/llama-kv-cache-dsa-iswa.h b/src/llama-kv-cache-dsa-iswa.h index 890b324e9477..cc31333f48ea 100644 --- a/src/llama-kv-cache-dsa-iswa.h +++ b/src/llama-kv-cache-dsa-iswa.h @@ -13,7 +13,7 @@ class llama_kv_cache_dsa_iswa : public llama_memory_i { public: llama_kv_cache_dsa_iswa( - const llama_model & model, + bool no_alloc, const llama_model & model, ggml_type type_k, ggml_type type_v, bool v_trans, diff --git a/src/llama-kv-cache-dsa.cpp b/src/llama-kv-cache-dsa.cpp index 163f21c45619..fef65f663a99 100644 --- a/src/llama-kv-cache-dsa.cpp +++ b/src/llama-kv-cache-dsa.cpp @@ -12,7 +12,7 @@ // llama_kv_cache_dsa::llama_kv_cache_dsa( - const llama_model & model, + bool no_alloc, const llama_model & model, ggml_type type_k, ggml_type type_v, bool v_trans, @@ -31,7 +31,7 @@ llama_kv_cache_dsa::llama_kv_cache_dsa( LLAMA_LOG_INFO("%s: creating main KV cache, size = %u cells\n", __func__, kv_size); kv_mla = std::make_unique( - model, model.hparams, type_k, type_v, + no_alloc, model, model.hparams, type_k, type_v, v_trans, offload, unified, kv_size, n_seq_max, n_pad, n_swa, swa_type, nullptr, filter_mla, reuse, nullptr); @@ -48,7 +48,7 @@ llama_kv_cache_dsa::llama_kv_cache_dsa( LLAMA_LOG_INFO("%s: creating indexer KV cache, size = %u cells\n", __func__, kv_size); kv_lid = std::make_unique( - model, hparams_lid, type_k, type_v, + no_alloc, model, hparams_lid, type_k, type_v, v_trans, offload, unified, kv_size, n_seq_max, n_pad, n_swa, swa_type, nullptr, filter_lid, reuse, nullptr); } diff --git a/src/llama-kv-cache-dsa.h b/src/llama-kv-cache-dsa.h index e6e21eb4f73b..f91f728901c8 100644 --- a/src/llama-kv-cache-dsa.h +++ b/src/llama-kv-cache-dsa.h @@ -15,7 +15,7 @@ class llama_kv_cache_dsa : public llama_memory_i { public: llama_kv_cache_dsa( - const llama_model & model, + bool no_alloc, const llama_model & model, ggml_type type_k, ggml_type type_v, bool v_trans, diff --git a/src/llama-kv-cache-dsv4.cpp b/src/llama-kv-cache-dsv4.cpp index d9ff820f7871..64eebbe1d4e4 100644 --- a/src/llama-kv-cache-dsv4.cpp +++ b/src/llama-kv-cache-dsv4.cpp @@ -893,7 +893,7 @@ static void dsv4_make_k_only(llama_hparams & hparams) { // llama_dsv4_comp_state::llama_dsv4_comp_state( - const llama_model & model, + bool no_alloc, const llama_model & model, bool offload, bool unified, uint32_t n_seq_max, @@ -903,6 +903,7 @@ llama_dsv4_comp_state::llama_dsv4_comp_state( uint32_t n_rs_seq, const char * name, const llama_memory_i::layer_filter_cb & filter) : + no_alloc(no_alloc || model.hparams.no_alloc), ratio(ratio), state_size(state_size), n_embd_state(n_embd_state), @@ -984,7 +985,14 @@ llama_dsv4_comp_state::llama_dsv4_comp_state( } for (auto & [buft, ctx] : ctx_map) { - ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx.get(), buft); + ggml_backend_buffer_t buf = this->no_alloc + ? ggml_backend_buft_alloc_buffer(buft, 0) + : ggml_backend_alloc_ctx_tensors_from_buft(ctx.get(), buft); + if (this->no_alloc && buf) { + for (auto * t = ggml_get_first_tensor(ctx.get()); t; t = ggml_get_next_tensor(ctx.get(), t)) { + t->buffer = buf; + } + } if (!buf) { throw std::runtime_error("failed to allocate buffer for DSV4 compressor state"); } @@ -1072,9 +1080,9 @@ uint32_t llama_dsv4_comp_state::get_n_rows() const { std::map llama_dsv4_comp_state::memory_breakdown() const { std::map ret; - for (const auto & [_, buf] : ctxs_bufs) { + for (const auto & [ctx, buf] : ctxs_bufs) { ggml_backend_buffer_type_t buft = ggml_backend_buffer_get_type(buf.get()); - ret[buft] += ggml_backend_buffer_get_size(buf.get()); + ret[buft] += no_alloc ? ggml_backend_alloc_ctx_tensors_from_buft_size(ctx.get(), buft) : ggml_backend_buffer_get_size(buf.get()); } return ret; } @@ -1198,7 +1206,7 @@ ggml_tensor * llama_dsv4_comp_state::cpy_score(ggml_context * ctx, ggml_tensor * size_t llama_dsv4_comp_state::total_size() const { size_t size = 0; - for (const auto & [_, buf] : ctxs_bufs) { + for (const auto & [ctx, buf] : ctxs_bufs) { size += ggml_backend_buffer_get_size(buf.get()); } @@ -1210,7 +1218,7 @@ size_t llama_dsv4_comp_state::total_size() const { // llama_kv_cache_dsv4::llama_kv_cache_dsv4( - const llama_model & model, + bool no_alloc, const llama_model & model, ggml_type type_k, ggml_type type_v, bool v_trans, @@ -1255,7 +1263,7 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4( dsv4_make_k_only(hparams_raw); kv_raw = std::make_unique( - model, hparams_raw, type_k, type_v, + no_alloc, model, hparams_raw, type_k, type_v, v_trans, offload, swa_full, unified_raw, kv_size, n_seq_max, n_ubatch, n_pad, nullptr, filter_raw, reuse, nullptr); @@ -1292,7 +1300,7 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4( __func__, dsv4_comp_size(kv_size, DSV4_CSA_RATIO)); kv_csa = std::make_unique( - model, hparams_csa, type_k, type_v, + no_alloc, model, hparams_csa, type_k, type_v, v_trans, offload, unified_compressed, GGML_PAD(dsv4_comp_size(kv_size, DSV4_CSA_RATIO), 256u), n_seq_max, n_pad, 0, LLAMA_SWA_TYPE_NONE, nullptr, filter_csa, nullptr, nullptr); @@ -1300,7 +1308,7 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4( __func__, dsv4_comp_size(kv_size, DSV4_HCA_RATIO)); kv_hca = std::make_unique( - model, hparams_hca, type_k, type_v, + no_alloc, model, hparams_hca, type_k, type_v, v_trans, offload, unified_compressed, GGML_PAD(dsv4_comp_size(kv_size, DSV4_HCA_RATIO), 256u), n_seq_max, n_pad, 0, LLAMA_SWA_TYPE_NONE, nullptr, filter_hca, nullptr, nullptr); @@ -1308,26 +1316,26 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4( __func__, dsv4_comp_size(kv_size, DSV4_CSA_RATIO)); kv_lid = std::make_unique( - model, hparams_lid, type_k, type_v, + no_alloc, model, hparams_lid, type_k, type_v, v_trans, offload, unified_compressed, GGML_PAD(dsv4_comp_size(kv_size, DSV4_CSA_RATIO), 256u), n_seq_max, n_pad, 0, LLAMA_SWA_TYPE_NONE, nullptr, filter_csa, nullptr, nullptr); LLAMA_LOG_INFO("%s: creating DSV4 CSA compressor state\n", __func__); csa_state = std::make_unique( - model, offload, unified_compressed, n_seq_max, DSV4_CSA_RATIO, 2*DSV4_CSA_RATIO, + no_alloc, model, offload, unified_compressed, n_seq_max, DSV4_CSA_RATIO, 2*DSV4_CSA_RATIO, 2*model.hparams.n_embd_head_k(), n_rs_seq, "csa", filter_csa); LLAMA_LOG_INFO("%s: creating DSV4 HCA compressor state\n", __func__); hca_state = std::make_unique( - model, offload, unified_compressed, n_seq_max, DSV4_HCA_RATIO, DSV4_HCA_RATIO, + no_alloc, model, offload, unified_compressed, n_seq_max, DSV4_HCA_RATIO, DSV4_HCA_RATIO, model.hparams.n_embd_head_k(), n_rs_seq, "hca", filter_hca); LLAMA_LOG_INFO("%s: creating DSV4 lightning-indexer compressor state\n", __func__); lid_state = std::make_unique( - model, offload, unified_compressed, n_seq_max, DSV4_CSA_RATIO, 2*DSV4_CSA_RATIO, + no_alloc, model, offload, unified_compressed, n_seq_max, DSV4_CSA_RATIO, 2*DSV4_CSA_RATIO, 2*model.hparams.indexer_head_size, n_rs_seq, "lid", filter_csa); // DSV4 attention reads compressed-K / compressor-state rows that the current diff --git a/src/llama-kv-cache-dsv4.h b/src/llama-kv-cache-dsv4.h index 293307434f13..77c7ab8fdfd1 100644 --- a/src/llama-kv-cache-dsv4.h +++ b/src/llama-kv-cache-dsv4.h @@ -15,7 +15,7 @@ class llama_dsv4_comp_state { stream_copy_info sc_info; llama_dsv4_comp_state( - const llama_model & model, + bool no_alloc, const llama_model & model, bool offload, bool unified, uint32_t n_seq_max, @@ -50,6 +50,7 @@ class llama_dsv4_comp_state { ggml_tensor * cpy_score(ggml_context * ctx, ggml_tensor * cur, ggml_tensor * idxs, int32_t il) const; private: + const bool no_alloc; struct layer { uint32_t il; @@ -88,7 +89,7 @@ class llama_dsv4_comp_state { class llama_kv_cache_dsv4 : public llama_memory_i { public: llama_kv_cache_dsv4( - const llama_model & model, + bool no_alloc, const llama_model & model, ggml_type type_k, ggml_type type_v, bool v_trans, diff --git a/src/llama-kv-cache-iswa.cpp b/src/llama-kv-cache-iswa.cpp index 4ac04ae50ce7..f5042af1e474 100644 --- a/src/llama-kv-cache-iswa.cpp +++ b/src/llama-kv-cache-iswa.cpp @@ -12,7 +12,7 @@ // llama_kv_cache_iswa::llama_kv_cache_iswa( - const llama_model & model, + bool no_alloc, const llama_model & model, ggml_type type_k, ggml_type type_v, bool v_trans, @@ -27,12 +27,12 @@ llama_kv_cache_iswa::llama_kv_cache_iswa( const layer_filter_cb & filter, const layer_reuse_cb & reuse, const layer_share_cb & share) : - llama_kv_cache_iswa(model, model.hparams, type_k, type_v, v_trans, offload, swa_full, unified, + llama_kv_cache_iswa(no_alloc, model, model.hparams, type_k, type_v, v_trans, offload, swa_full, unified, kv_size, n_seq_max, n_ubatch, n_pad, mem_other, filter, reuse, share) { } llama_kv_cache_iswa::llama_kv_cache_iswa( - const llama_model & model, + bool no_alloc, const llama_model & model, const llama_hparams & hparams, ggml_type type_k, ggml_type type_v, @@ -93,14 +93,14 @@ llama_kv_cache_iswa::llama_kv_cache_iswa( } kv_base = std::make_unique( - model, hparams, type_k, type_v, + no_alloc, model, hparams, type_k, type_v, v_trans, offload, unified, size_base, n_seq_max, n_pad, 0, LLAMA_SWA_TYPE_NONE, mem_other_base, filter_base, reuse, share); LLAMA_LOG_INFO("%s: creating SWA KV cache, size = %u cells\n", __func__, size_swa); kv_swa = std::make_unique( - model, hparams, type_k, type_v, + no_alloc, model, hparams, type_k, type_v, v_trans, offload, unified, size_swa, n_seq_max, n_pad, hparams.n_swa, hparams.swa_type, mem_other_swa, filter_swa, reuse, share); } diff --git a/src/llama-kv-cache-iswa.h b/src/llama-kv-cache-iswa.h index 546b8325d8ad..af12089903cc 100644 --- a/src/llama-kv-cache-iswa.h +++ b/src/llama-kv-cache-iswa.h @@ -14,7 +14,7 @@ class llama_kv_cache_iswa : public llama_memory_i { public: llama_kv_cache_iswa( - const llama_model & model, + bool no_alloc, const llama_model & model, ggml_type type_k, ggml_type type_v, bool v_trans, @@ -31,7 +31,7 @@ class llama_kv_cache_iswa : public llama_memory_i { const layer_share_cb & share); llama_kv_cache_iswa( - const llama_model & model, + bool no_alloc, const llama_model & model, const llama_hparams & hparams, ggml_type type_k, ggml_type type_v, diff --git a/src/llama-kv-cache-msa.cpp b/src/llama-kv-cache-msa.cpp index d4de229a9b81..f83cfffce2f0 100644 --- a/src/llama-kv-cache-msa.cpp +++ b/src/llama-kv-cache-msa.cpp @@ -11,7 +11,7 @@ // llama_kv_cache_msa llama_kv_cache_msa::llama_kv_cache_msa( - const llama_model & model, + bool no_alloc, const llama_model & model, ggml_type type_k, ggml_type type_v, bool v_trans, @@ -32,7 +32,7 @@ llama_kv_cache_msa::llama_kv_cache_msa( LLAMA_LOG_INFO("%s: creating main KV cache, size = %u cells\n", __func__, kv_size); kv_base = std::make_unique( - model, model.hparams, type_k, type_v, + no_alloc, model, model.hparams, type_k, type_v, v_trans, offload, unified, kv_size, n_seq_max, n_pad, n_swa, swa_type, nullptr, filter, reuse, nullptr); @@ -44,7 +44,7 @@ llama_kv_cache_msa::llama_kv_cache_msa( LLAMA_LOG_INFO("%s: creating indexer KV cache, size = %u cells\n", __func__, kv_size); kv_idx = std::make_unique( - model, hparams_idx, type_k, type_v, + no_alloc, model, hparams_idx, type_k, type_v, v_trans, offload, unified, kv_size, n_seq_max, n_pad, n_swa, swa_type, nullptr, filter_idx, reuse, nullptr); } diff --git a/src/llama-kv-cache-msa.h b/src/llama-kv-cache-msa.h index 04fce2d66671..0da457e5371b 100644 --- a/src/llama-kv-cache-msa.h +++ b/src/llama-kv-cache-msa.h @@ -14,7 +14,7 @@ class llama_kv_cache_msa : public llama_memory_i { public: llama_kv_cache_msa( - const llama_model & model, + bool no_alloc, const llama_model & model, ggml_type type_k, ggml_type type_v, bool v_trans, diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 76aa5692442e..d5daa3141a77 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -62,7 +62,7 @@ static void ggml_gen_hadamard(ggml_tensor * tensor) { // llama_kv_cache::llama_kv_cache( - const llama_model & model, + bool no_alloc, const llama_model & model, const llama_hparams & hparams, ggml_type type_k, ggml_type type_v, @@ -79,7 +79,7 @@ llama_kv_cache::llama_kv_cache( const layer_reuse_cb & reuse, const layer_share_cb & share, const char * name_tag) : - model(model), hparams(hparams), v_trans(v_trans), + no_alloc(no_alloc || hparams.no_alloc), model(model), hparams(hparams), v_trans(v_trans), n_seq_max(n_seq_max), n_stream(unified ? 1 : n_seq_max), n_pad(n_pad), n_swa(n_swa), swa_type(swa_type), other(static_cast(mem_other)), v_cells_impl(other ? other->v_cells_impl : std::make_shared()), @@ -275,7 +275,7 @@ llama_kv_cache::llama_kv_cache( // allocate tensors and initialize the buffers to avoid NaNs in the padding for (auto & [buft, ctx] : ctx_map) { ggml_backend_buffer_t buf; - if (hparams.no_alloc) { + if (this->no_alloc) { buf = ggml_backend_buft_alloc_buffer(buft, /*size =*/ 0); // dummy buffer for (ggml_tensor * t = ggml_get_first_tensor(ctx.get()); t != nullptr; t = ggml_get_next_tensor(ctx.get(), t)) { t->buffer = buf; // set dummy buffer for KV cache so that the backend scheduler won't try to allocate it @@ -686,7 +686,7 @@ std::map llama_kv_cache::memory_breakdown() for (const auto & [ctx, buf] : ctxs_bufs) { ggml_backend_buffer_type_t buft = ggml_backend_buffer_get_type(buf.get()); - if (hparams.no_alloc) { + if (this->no_alloc) { GGML_ASSERT(ggml_backend_buffer_get_base(buf.get()) == nullptr); ret[buft] += ggml_backend_alloc_ctx_tensors_from_buft_size(ctx.get(), buft); } else { diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index cba1f663b2e4..00ce5bfd4261 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -97,7 +97,7 @@ class llama_kv_cache : public llama_memory_i { // instead pass all necessary info (e.g. hparams, dev layers, arch, etc.) directly // likely through `struct llama_memory_params` llama_kv_cache( - const llama_model & model, + bool no_alloc, const llama_model & model, const llama_hparams & hparams, ggml_type type_k, ggml_type type_v, @@ -229,6 +229,7 @@ class llama_kv_cache : public llama_memory_i { void set_input_v_rot(ggml_tensor * dst) const; private: + const bool no_alloc; const llama_model & model; const llama_hparams & hparams; diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index bfba14c56454..46442d3e5659 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/src/llama-memory-hybrid-idx.cpp @@ -16,7 +16,7 @@ // llama_memory_hybrid_idx::llama_memory_hybrid_idx( - const llama_model & model, + bool no_alloc, const llama_model & model, /* attn */ ggml_type type_k, ggml_type type_v, @@ -39,7 +39,7 @@ llama_memory_hybrid_idx::llama_memory_hybrid_idx( const layer_filter_cb & filter_recr, const layer_filter_cb & filter_idx) : llama_memory_hybrid( - model, + no_alloc, model, type_k, type_v, v_trans, kv_size, n_pad, n_swa, swa_type, type_r, type_s, rs_size, n_seq_max, n_rs_seq, offload, unified, @@ -53,7 +53,7 @@ llama_memory_hybrid_idx::llama_memory_hybrid_idx( LLAMA_LOG_INFO("%s: creating indexer KV cache, size = %u cells\n", __func__, kv_size); return new llama_kv_cache( - model, hparams_idx, type_k, type_v, v_trans, offload, unified, + no_alloc, model, hparams_idx, type_k, type_v, v_trans, offload, unified, kv_size, n_seq_max, n_pad, n_swa, swa_type, nullptr, filter_idx, nullptr, nullptr, "idx_"); }()) {} diff --git a/src/llama-memory-hybrid-idx.h b/src/llama-memory-hybrid-idx.h index f8a3f2dad905..c0d2d349f328 100644 --- a/src/llama-memory-hybrid-idx.h +++ b/src/llama-memory-hybrid-idx.h @@ -18,7 +18,7 @@ class llama_memory_hybrid_idx : public llama_memory_hybrid { public: llama_memory_hybrid_idx( - const llama_model & model, + bool no_alloc, const llama_model & model, /* attn */ ggml_type type_k, ggml_type type_v, diff --git a/src/llama-memory-hybrid-iswa.cpp b/src/llama-memory-hybrid-iswa.cpp index 9ab3d6f70b18..443beb40766d 100644 --- a/src/llama-memory-hybrid-iswa.cpp +++ b/src/llama-memory-hybrid-iswa.cpp @@ -9,7 +9,7 @@ // llama_memory_hybrid_iswa::llama_memory_hybrid_iswa( - const llama_model & model, + bool no_alloc, const llama_model & model, /* attn */ ggml_type type_k, ggml_type type_v, @@ -32,7 +32,7 @@ llama_memory_hybrid_iswa::llama_memory_hybrid_iswa( const layer_filter_cb & filter_recr) : hparams(model.hparams), mem_attn(new llama_kv_cache_iswa( - model, + no_alloc, model, type_k, type_v, v_trans, @@ -51,7 +51,7 @@ llama_memory_hybrid_iswa::llama_memory_hybrid_iswa( nullptr )), mem_recr(new llama_memory_recurrent( - model, + no_alloc, model, type_r, type_s, offload, diff --git a/src/llama-memory-hybrid-iswa.h b/src/llama-memory-hybrid-iswa.h index a6b9d84c1ecb..13ff7de0fee2 100644 --- a/src/llama-memory-hybrid-iswa.h +++ b/src/llama-memory-hybrid-iswa.h @@ -19,7 +19,7 @@ class llama_memory_hybrid_iswa : public llama_memory_i { public: llama_memory_hybrid_iswa( - const llama_model & model, + bool no_alloc, const llama_model & model, /* attn */ ggml_type type_k, ggml_type type_v, diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp index 6cb6faee3cf9..efd7fbda71b7 100644 --- a/src/llama-memory-hybrid.cpp +++ b/src/llama-memory-hybrid.cpp @@ -9,7 +9,7 @@ // llama_memory_hybrid::llama_memory_hybrid( - const llama_model & model, + bool no_alloc, const llama_model & model, /* attn */ ggml_type type_k, ggml_type type_v, @@ -32,7 +32,7 @@ llama_memory_hybrid::llama_memory_hybrid( const layer_filter_cb & filter_recr) : hparams(model.hparams), mem_attn(new llama_kv_cache( - model, + no_alloc, model, model.hparams, type_k, type_v, @@ -52,7 +52,7 @@ llama_memory_hybrid::llama_memory_hybrid( nullptr )), mem_recr(new llama_memory_recurrent( - model, + no_alloc, model, type_r, type_s, offload, diff --git a/src/llama-memory-hybrid.h b/src/llama-memory-hybrid.h index 8cd1dfc80982..4c6a94ec5f83 100644 --- a/src/llama-memory-hybrid.h +++ b/src/llama-memory-hybrid.h @@ -19,7 +19,7 @@ class llama_memory_hybrid : public llama_memory_i { public: llama_memory_hybrid( - const llama_model & model, + bool no_alloc, const llama_model & model, /* attn */ ggml_type type_k, ggml_type type_v, diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp index 208c9d0d557b..afff0d215bae 100644 --- a/src/llama-memory-recurrent.cpp +++ b/src/llama-memory-recurrent.cpp @@ -18,14 +18,15 @@ // llama_memory_recurrent::llama_memory_recurrent( - const llama_model & model, + bool no_alloc, const llama_model & model, ggml_type type_r, ggml_type type_s, bool offload, uint32_t mem_size, uint32_t n_seq_max, uint32_t n_rs_seq, - const layer_filter_cb & filter) : hparams(model.hparams), n_seq_max(n_seq_max) { + const layer_filter_cb & filter) : + no_alloc(no_alloc || model.hparams.no_alloc), hparams(model.hparams), n_seq_max(n_seq_max) { const int32_t n_layer = hparams.n_layer(); head = 0; @@ -107,7 +108,14 @@ llama_memory_recurrent::llama_memory_recurrent( // allocate tensors and initialize the buffers to avoid NaNs in the padding for (auto & [buft, ctx] : ctx_map) { - ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx.get(), buft); + ggml_backend_buffer_t buf = this->no_alloc + ? ggml_backend_buft_alloc_buffer(buft, 0) + : ggml_backend_alloc_ctx_tensors_from_buft(ctx.get(), buft); + if (this->no_alloc && buf) { + for (auto * t = ggml_get_first_tensor(ctx.get()); t; t = ggml_get_next_tensor(ctx.get(), t)) { + t->buffer = buf; + } + } if (!buf) { throw std::runtime_error("failed to allocate buffer for rs cache"); } @@ -408,8 +416,10 @@ void llama_memory_recurrent::set_rs_idx(llama_seq_id seq_id, uint32_t idx) { std::map llama_memory_recurrent::memory_breakdown() const { std::map ret; - for (const auto & [_, buf] : ctxs_bufs) { - ret[ggml_backend_buffer_get_type(buf.get())] += ggml_backend_buffer_get_size(buf.get()); + for (const auto & [ctx, buf] : ctxs_bufs) { + auto * buft = ggml_backend_buffer_get_type(buf.get()); + ret[buft] += no_alloc ? ggml_backend_alloc_ctx_tensors_from_buft_size(ctx.get(), buft) + : ggml_backend_buffer_get_size(buf.get()); } return ret; } @@ -729,7 +739,7 @@ bool llama_memory_recurrent::get_can_shift() const { size_t llama_memory_recurrent::total_size() const { size_t size = 0; - for (const auto & [_, buf] : ctxs_bufs) { + for (const auto & [ctx, buf] : ctxs_bufs) { size += ggml_backend_buffer_get_size(buf.get()); } diff --git a/src/llama-memory-recurrent.h b/src/llama-memory-recurrent.h index 3cbc682a88de..aa35c23fc22d 100644 --- a/src/llama-memory-recurrent.h +++ b/src/llama-memory-recurrent.h @@ -17,7 +17,7 @@ class llama_memory_recurrent : public llama_memory_i { public: llama_memory_recurrent( - const llama_model & model, + bool no_alloc, const llama_model & model, ggml_type type_r, ggml_type type_s, bool offload, @@ -126,6 +126,7 @@ class llama_memory_recurrent : public llama_memory_i { std::vector s_l; private: + const bool no_alloc; uint32_t protected_first = 0, protected_count = 0; //const llama_model & model; const llama_hparams & hparams; diff --git a/src/llama-memory.h b/src/llama-memory.h index ceb327db46f1..621c7e756bf4 100644 --- a/src/llama-memory.h +++ b/src/llama-memory.h @@ -17,6 +17,7 @@ class llama_io_write_i; class llama_io_read_i; struct llama_memory_params { + bool no_alloc; // kv cache ggml_type type_k; ggml_type type_v; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 63fa565e7843..1c12fad178de 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2227,7 +2227,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, [&](int32_t il) { return (uint32_t) il >= hparams.n_layer_dense_lead; }; res = new llama_kv_cache_msa( - *this, + params.no_alloc, *this, params.type_k, params.type_v, !cparams.flash_attn, @@ -2253,7 +2253,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, [&](uint32_t il) { return il >= hparams.n_layer(); }; res = new llama_kv_cache( - *this, + params.no_alloc, *this, hparams, params.type_k, params.type_v, @@ -2279,7 +2279,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, llama_kv_cache::layer_filter_cb filter_lid = [&](uint32_t il) { return il < hparams.n_layer() && (arch != LLM_ARCH_GLM_DSA || hparams.is_indexer_full(il)); }; res = new llama_kv_cache_dsa( - *this, + params.no_alloc, *this, params.type_k, params.type_v, !cparams.flash_attn, @@ -2305,7 +2305,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, [&](uint32_t il) { return il >= hparams.n_layer(); }; res = new llama_kv_cache( - *this, + params.no_alloc, *this, hparams, params.type_k, params.type_v, @@ -2330,7 +2330,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, llama_kv_cache::layer_filter_cb filter_lid = [&](uint32_t il) { return il < hparams.n_layer() && hparams.is_indexer_full(il); }; res = new llama_kv_cache_dsa_iswa( - *this, + params.no_alloc, *this, params.type_k, params.type_v, !cparams.flash_attn, @@ -2356,7 +2356,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, }; res = new llama_kv_cache_iswa( - *this, + params.no_alloc, *this, params.type_k, params.type_v, !cparams.flash_attn, @@ -2373,7 +2373,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, nullptr); } else { res = new llama_kv_cache_dsv4( - *this, + params.no_alloc, *this, params.type_k, params.type_v, !cparams.flash_attn, @@ -2396,7 +2396,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, GGML_ASSERT(hparams.swa_type != LLAMA_SWA_TYPE_NONE); res = new llama_kv_cache_iswa( - *this, + params.no_alloc, *this, params.type_k, params.type_v, !cparams.flash_attn, @@ -2430,7 +2430,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, if (llm_arch_is_recurrent(arch)) { res = new llama_memory_recurrent( - *this, + params.no_alloc, *this, GGML_TYPE_F32, GGML_TYPE_F32, cparams.offload_kqv, @@ -2476,7 +2476,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, if (hparams.swa_type != LLAMA_SWA_TYPE_NONE) { // Use hybrid-iswa for hybrid models with SWA res = new llama_memory_hybrid_iswa( - /* model */ *this, + params.no_alloc, /* model */ *this, /* attn_type_k */ params.type_k, /* attn_type_v */ params.type_v, /* attn_v_trans */ !cparams.flash_attn, @@ -2496,7 +2496,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, } else if (needs_mem_idx) { // sparse attention over a per-token indexer cache, in its own memory type res = new llama_memory_hybrid_idx( - /* model */ *this, + params.no_alloc, /* model */ *this, /* attn_type_k */ params.type_k, /* attn_type_v */ params.type_v, /* attn_v_trans */ !cparams.flash_attn, @@ -2516,7 +2516,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, /* filter_idx */ std::move(filter_idx)); } else { res = new llama_memory_hybrid( - /* model */ *this, + params.no_alloc, /* model */ *this, /* attn_type_k */ params.type_k, /* attn_type_v */ params.type_v, /* attn_v_trans */ !cparams.flash_attn, @@ -2582,7 +2582,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, }; res = new llama_kv_cache_iswa( - *this, + params.no_alloc, *this, params.type_k, params.type_v, !cparams.flash_attn, @@ -2599,7 +2599,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, share); } else { res = new llama_kv_cache_iswa( - *this, + params.no_alloc, *this, params.type_k, params.type_v, !cparams.flash_attn, @@ -2619,7 +2619,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, GGML_ASSERT(!hparams.is_swa_any()); res = new llama_kv_cache( - *this, + params.no_alloc, *this, hparams, params.type_k, params.type_v,