From a4699400e3ccca756297ef4433eb1d7998fca525 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Mon, 26 Oct 2026 16:21:25 +0200 Subject: [PATCH] llama : move no_alloc out of llama_hparams (#31141) * llama : move no_alloc out of llama_hparams Assisted-by: dsh:llama.cpp/ggml-org/GLM-5.3-Flash-GGUF (2x DGX Spark (RDMA, TP=2)) * cont : clean-up * cont : clean-up Assisted-by: dsh:llama.cpp/ggml-org/GLM-5.3-Flash-GGUF (2x DGX Spark (RDMA, TP=2)) * tests : add a no_alloc case to test-llama-context Assisted-by: dsh:llama.cpp/ggml-org/GLM-5.3-Flash-GGUF (2x DGX Spark (RDMA, TP=2)) --- src/llama-context.cpp | 17 ++++---- src/llama-hparams.h | 1 - src/llama-kv-cache-dsa-iswa.cpp | 5 ++- src/llama-kv-cache-dsa-iswa.h | 1 + src/llama-kv-cache-dsa.cpp | 5 ++- src/llama-kv-cache-dsa.h | 1 + src/llama-kv-cache-dsv4.cpp | 68 +++++++++++--------------------- src/llama-kv-cache-dsv4.h | 22 +++++------ src/llama-kv-cache-iswa.cpp | 10 +++-- src/llama-kv-cache-iswa.h | 4 +- src/llama-kv-cache-msa.cpp | 5 ++- src/llama-kv-cache-msa.h | 1 + src/llama-kv-cache.cpp | 39 ++++-------------- src/llama-kv-cache.h | 3 +- src/llama-memory-hybrid-idx.cpp | 10 ++--- src/llama-memory-hybrid-idx.h | 1 + src/llama-memory-hybrid-iswa.cpp | 3 ++ src/llama-memory-hybrid-iswa.h | 1 + src/llama-memory-hybrid.cpp | 3 ++ src/llama-memory-hybrid.h | 1 + src/llama-memory-recurrent.cpp | 34 +++------------- src/llama-memory-recurrent.h | 5 ++- src/llama-memory.cpp | 52 ++++++++++++++++++++++++ src/llama-memory.h | 31 +++++++++++++++ src/llama-model.cpp | 26 +++++++++++- src/llama-model.h | 1 + src/llama-moe-cache.cpp | 43 ++++++-------------- src/llama-moe-cache.h | 2 +- src/llama.cpp | 1 - tests/llama-context/cases-kv.cpp | 49 +++++++++++++++++++++++ 30 files changed, 266 insertions(+), 179 deletions(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index d5a66a499..455f5198e 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -489,6 +489,7 @@ llama_context::llama_context( /*.swa_full =*/ params.swa_full, /*.ctx_type =*/ cparams.ctx_type, /*.mem_other =*/ llama_get_memory(cparams.ctx_other), + /*.no_alloc =*/ model.no_alloc(), }; memory.reset(model.create_memory(params_mem, cparams)); @@ -559,7 +560,7 @@ llama_context::llama_context( } if (cparams.moe_cache_size > 0) { - moe_cache = std::make_unique(model, backend_ptrs, backend_buft, cparams.moe_cache_size); + moe_cache = std::make_unique(model, model.no_alloc(), backend_ptrs, backend_buft, cparams.moe_cache_size); } sched_reserve(); @@ -587,7 +588,7 @@ llama_context::~llama_context() { synchronize(); // when training, ggml_opt allocates extra buffers through the scheduler, so the sizes no longer match the expectation - if (!model.hparams.no_alloc && !opt_ctx) { + if (!model.no_alloc() && !opt_ctx) { 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]; @@ -788,7 +789,7 @@ 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); + model.no_alloc(), model.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__); @@ -815,7 +816,7 @@ void llama_context::sched_reserve() { n_splits_tg = n_splits_pp; n_nodes_tg = n_nodes_pp; } else { - 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(), model.no_alloc()); if (!gf) { throw std::runtime_error("failed to allocate compute tg buffers"); } @@ -837,10 +838,10 @@ void llama_context::sched_reserve() { // [TAG_RESERVE_DIAG_DECAY] // the `inp_diag_decay` tensor size scales with `n_seq_tokens^2` which // makes `n_seqs == 1` use more memory for the compute graph compared to `n_seqs > 1` - gf = graph_reserve(n_tokens, 1, n_outputs_pp, mctx.get(), model.hparams.no_alloc); + gf = graph_reserve(n_tokens, 1, n_outputs_pp, mctx.get(), model.no_alloc()); break; default: - gf = graph_reserve(n_tokens, n_seqs, n_outputs_pp, mctx.get(), model.hparams.no_alloc); + gf = graph_reserve(n_tokens, n_seqs, n_outputs_pp, mctx.get(), model.no_alloc()); }; if (!gf) { @@ -851,7 +852,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 (!model.no_alloc()) { backend_buf_exp_size[i] = ggml_backend_sched_get_buffer_size(sched.get(), backend); } if (backend_buf_exp_size[i] > 1) { @@ -3730,7 +3731,7 @@ llama_memory_breakdown llama_context::memory_breakdown() const { ret[buft].context += size; } } - if (model.hparams.no_alloc) { + if (model.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); diff --git a/src/llama-hparams.h b/src/llama-hparams.h index 760bb8ac6..7abb372ce 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -61,7 +61,6 @@ struct llama_hparams { // for example: n_embd_out() vs n_embd_out_impl bool vocab_only; - bool no_alloc; // TODO: no_alloc is not a model hyperparameter; move it out of llama_hparams bool rope_finetuned; bool use_par_res; bool swin_norm; diff --git a/src/llama-kv-cache-dsa-iswa.cpp b/src/llama-kv-cache-dsa-iswa.cpp index 784c78aa7..02d14a344 100644 --- a/src/llama-kv-cache-dsa-iswa.cpp +++ b/src/llama-kv-cache-dsa-iswa.cpp @@ -13,6 +13,7 @@ llama_kv_cache_dsa_iswa::llama_kv_cache_dsa_iswa( const llama_model & model, + bool no_alloc, ggml_type type_k, ggml_type type_v, bool v_trans, @@ -63,14 +64,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, + model, no_alloc, 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, + model, no_alloc, 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 28cf95bf0..1aae269da 100644 --- a/src/llama-kv-cache-dsa-iswa.h +++ b/src/llama-kv-cache-dsa-iswa.h @@ -14,6 +14,7 @@ class llama_kv_cache_dsa_iswa : public llama_memory_i { public: llama_kv_cache_dsa_iswa( const llama_model & model, + bool no_alloc, 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 6e5f7d484..307ff31dd 100644 --- a/src/llama-kv-cache-dsa.cpp +++ b/src/llama-kv-cache-dsa.cpp @@ -13,6 +13,7 @@ llama_kv_cache_dsa::llama_kv_cache_dsa( const llama_model & model, + bool no_alloc, ggml_type type_k, ggml_type type_v, bool v_trans, @@ -31,7 +32,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, + model, no_alloc, 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 +49,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, + model, no_alloc, 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 e74fc4d91..584ab40b9 100644 --- a/src/llama-kv-cache-dsa.h +++ b/src/llama-kv-cache-dsa.h @@ -16,6 +16,7 @@ class llama_kv_cache_dsa : public llama_memory_i { public: llama_kv_cache_dsa( const llama_model & model, + bool no_alloc, 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 6186f9d34..70bfee8f9 100644 --- a/src/llama-kv-cache-dsv4.cpp +++ b/src/llama-kv-cache-dsv4.cpp @@ -931,21 +931,24 @@ static void dsv4_make_k_only(llama_hparams & hparams) { llama_dsv4_comp_state::llama_dsv4_comp_state( const llama_model & model, - bool offload, - bool unified, - uint32_t n_seq_max, - uint32_t ratio, - uint32_t state_size, - uint32_t n_embd_state, - uint32_t n_rs_seq, - const char * name, + bool no_alloc, + bool offload, + bool unified, + uint32_t n_seq_max, + uint32_t ratio, + uint32_t state_size, + uint32_t n_embd_state, + uint32_t n_rs_seq, + const char * name, const llama_memory_i::layer_filter_cb & filter) : - hparams(model.hparams), ratio(ratio), state_size(state_size), n_embd_state(n_embd_state), n_stream(unified ? 1 : n_seq_max), - n_rs_seq(n_rs_seq) { + n_rs_seq(n_rs_seq), + ctxs_bufs(no_alloc) { + const llama_hparams & hparams = model.hparams; + struct ggml_backend_buft_comparator { bool operator()(const ggml_backend_buffer_type_t & lhs, const ggml_backend_buffer_type_t & rhs) const { return strcmp(ggml_backend_buft_name(lhs), ggml_backend_buft_name(rhs)) < 0; @@ -1020,25 +1023,10 @@ llama_dsv4_comp_state::llama_dsv4_comp_state( } for (auto & [buft, ctx] : ctx_map) { - ggml_backend_buffer_t buf; - if (hparams.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 DSV4 state so that the backend scheduler won't try to allocate it - } - } else { - buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx.get(), buft); // real buffer - } - if (!buf) { - throw std::runtime_error("failed to allocate buffer for DSV4 compressor state"); - } - - ggml_backend_buffer_clear(buf, 0); + ggml_backend_buffer_t buf = ctxs_bufs.emplace(std::move(ctx), buft, "DSV4 compressor state"); LLAMA_LOG_INFO("%s: %10s DSV4 %s state buffer size = %8.2f MiB\n", __func__, ggml_backend_buffer_name(buf), name, ggml_backend_buffer_get_size(buf)/1024.0/1024.0); - - ctxs_bufs.emplace_back(std::move(ctx), buf); } LLAMA_LOG_INFO("%s: %s ratio = %u, state = %u x %u, streams = %u, rs_seq = %u, layers = %zu, size = %7.2f MiB\n", @@ -1119,18 +1107,7 @@ 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 & [ctx, buf] : ctxs_bufs) { - ggml_backend_buffer_type_t buft = ggml_backend_buffer_get_type(buf.get()); - - if (hparams.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 { - ret[buft] += ggml_backend_buffer_get_size(buf.get()); - } - } - return ret; + return ctxs_bufs.memory_breakdown(); } void llama_dsv4_comp_state::state_write( @@ -1265,6 +1242,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, ggml_type type_k, ggml_type type_v, bool v_trans, @@ -1309,7 +1287,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, + model, no_alloc, 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); @@ -1413,7 +1391,7 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4( __func__, dsv4_comp_size(kv_size, ratio_csa)); kv_csa = std::make_unique( - model, hparams_csa, type_k, type_v, + model, no_alloc, hparams_csa, type_k, type_v, v_trans, offload, unified_compressed, GGML_PAD(dsv4_comp_size(kv_size, ratio_csa), 256u), n_seq_max, n_pad, 0, LLAMA_SWA_TYPE_NONE, nullptr, filter_csa, nullptr, nullptr); @@ -1421,7 +1399,7 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4( __func__, dsv4_comp_size(kv_size, ratio_hca)); kv_hca = std::make_unique( - model, hparams_hca, type_k, type_v, + model, no_alloc, hparams_hca, type_k, type_v, v_trans, offload, unified_compressed, GGML_PAD(dsv4_comp_size(kv_size, ratio_hca), 256u), n_seq_max, n_pad, 0, LLAMA_SWA_TYPE_NONE, nullptr, filter_hca, nullptr, nullptr); @@ -1429,7 +1407,7 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4( __func__, dsv4_comp_size(kv_size, ratio_csa)); kv_lid = std::make_unique( - model, hparams_lid, type_k, type_v, + model, no_alloc, hparams_lid, type_k, type_v, v_trans, offload, unified_compressed, GGML_PAD(dsv4_comp_size(kv_size, ratio_csa), 256u), n_seq_max, n_pad, 0, LLAMA_SWA_TYPE_NONE, nullptr, filter_lid, nullptr, nullptr); @@ -1439,19 +1417,19 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4( const uint32_t csa_state_embd = is_v41 ? model.hparams.n_embd_head_k() : 2*model.hparams.n_embd_head_k(); csa_state = std::make_unique( - model, offload, unified_compressed, n_seq_max, ratio_csa, csa_state_size, + model, no_alloc, offload, unified_compressed, n_seq_max, ratio_csa, csa_state_size, csa_state_embd, n_rs_seq, "csa", filter_csa_state); LLAMA_LOG_INFO("%s: creating DSV4 HCA compressor state\n", __func__); hca_state = std::make_unique( - model, offload, unified_compressed, n_seq_max, ratio_hca, ratio_hca, + model, no_alloc, offload, unified_compressed, n_seq_max, ratio_hca, ratio_hca, model.hparams.n_embd_head_k(), n_rs_seq, "hca", filter_hca_state); 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, + model, no_alloc, offload, unified_compressed, n_seq_max, DSV4_CSA_RATIO, 2*DSV4_CSA_RATIO, 2*model.hparams.indexer_head_size, n_rs_seq, "lid", filter_lid); // 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 bfc52947f..3e97c46e2 100644 --- a/src/llama-kv-cache-dsv4.h +++ b/src/llama-kv-cache-dsv4.h @@ -16,14 +16,15 @@ public: llama_dsv4_comp_state( const llama_model & model, - bool offload, - bool unified, - uint32_t n_seq_max, - uint32_t ratio, - uint32_t state_size, - uint32_t n_embd_state, - uint32_t n_rs_seq, - const char * name, + bool no_alloc, + bool offload, + bool unified, + uint32_t n_seq_max, + uint32_t ratio, + uint32_t state_size, + uint32_t n_embd_state, + uint32_t n_rs_seq, + const char * name, const llama_memory_i::layer_filter_cb & filter); void clear(llama_seq_id seq_id, bool data); @@ -61,15 +62,13 @@ private: std::vector score_stream; }; - const llama_hparams & hparams; - const uint32_t ratio; const uint32_t state_size; const uint32_t n_embd_state; const uint32_t n_stream; const uint32_t n_rs_seq; - std::vector> ctxs_bufs; + llama_memory_bufs ctxs_bufs; std::vector layers; @@ -92,6 +91,7 @@ class llama_kv_cache_dsv4 : public llama_memory_i { public: llama_kv_cache_dsv4( const llama_model & model, + bool no_alloc, 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 da440f6a2..4606bcf98 100644 --- a/src/llama-kv-cache-iswa.cpp +++ b/src/llama-kv-cache-iswa.cpp @@ -13,6 +13,7 @@ llama_kv_cache_iswa::llama_kv_cache_iswa( const llama_model & model, + bool no_alloc, ggml_type type_k, ggml_type type_v, bool v_trans, @@ -27,13 +28,14 @@ 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(model, no_alloc, 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, - const llama_hparams & hparams, + bool no_alloc, + const llama_hparams & hparams, ggml_type type_k, ggml_type type_v, bool v_trans, @@ -93,14 +95,14 @@ llama_kv_cache_iswa::llama_kv_cache_iswa( } kv_base = std::make_unique( - model, hparams, type_k, type_v, + model, no_alloc, 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, + model, no_alloc, 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 7dab6eaa8..8352669dd 100644 --- a/src/llama-kv-cache-iswa.h +++ b/src/llama-kv-cache-iswa.h @@ -15,6 +15,7 @@ class llama_kv_cache_iswa : public llama_memory_i { public: llama_kv_cache_iswa( const llama_model & model, + bool no_alloc, ggml_type type_k, ggml_type type_v, bool v_trans, @@ -32,7 +33,8 @@ public: llama_kv_cache_iswa( const llama_model & model, - const llama_hparams & hparams, + bool no_alloc, + const llama_hparams & hparams, ggml_type type_k, ggml_type type_v, bool v_trans, diff --git a/src/llama-kv-cache-msa.cpp b/src/llama-kv-cache-msa.cpp index 55ef286ca..a57937ff4 100644 --- a/src/llama-kv-cache-msa.cpp +++ b/src/llama-kv-cache-msa.cpp @@ -12,6 +12,7 @@ llama_kv_cache_msa::llama_kv_cache_msa( const llama_model & model, + bool no_alloc, ggml_type type_k, ggml_type type_v, bool v_trans, @@ -32,7 +33,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, + model, no_alloc, 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 +45,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, + model, no_alloc, 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 f09b6d32b..3b1477c83 100644 --- a/src/llama-kv-cache-msa.h +++ b/src/llama-kv-cache-msa.h @@ -15,6 +15,7 @@ class llama_kv_cache_msa : public llama_memory_i { public: llama_kv_cache_msa( const llama_model & model, + bool no_alloc, 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 4ab280a61..61b5cb370 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -64,7 +64,8 @@ static void ggml_gen_hadamard(ggml_tensor * tensor) { llama_kv_cache::llama_kv_cache( const llama_model & model, - const llama_hparams & hparams, + bool no_alloc, + const llama_hparams & hparams, ggml_type type_k, ggml_type type_v, bool v_trans, @@ -79,9 +80,10 @@ llama_kv_cache::llama_kv_cache( const layer_filter_cb & filter, const layer_reuse_cb & reuse, const layer_share_cb & share, - const char * name_tag) : + const char * name_tag) : 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), + ctxs_bufs(no_alloc), other(static_cast(mem_other)), v_cells_impl(other ? other->v_cells_impl : std::make_shared()), v_cells(*v_cells_impl) { @@ -273,25 +275,11 @@ llama_kv_cache::llama_kv_cache( } } - // allocate tensors and initialize the buffers to avoid NaNs in the padding + // allocate the 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) { - 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 - } - } else { - buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx.get(), buft); // real buffer - } - if (!buf) { - throw std::runtime_error("failed to allocate buffer for kv cache"); - } + ggml_backend_buffer_t buf = ctxs_bufs.emplace(std::move(ctx), buft, "kv cache"); LLAMA_LOG_INFO("%s: %10s KV buffer size = %8.2f MiB\n", __func__, ggml_backend_buffer_name(buf), ggml_backend_buffer_get_size(buf)/1024.0/1024.0); - - ggml_backend_buffer_clear(buf, 0); - ctxs_bufs.emplace_back(std::move(ctx), buf); } { @@ -705,20 +693,7 @@ llama_pos llama_kv_cache::seq_pos_max(llama_seq_id seq_id) const { } std::map llama_kv_cache::memory_breakdown() const { - std::map ret; - for (const auto & [ctx, buf] : ctxs_bufs) { - ggml_backend_buffer_type_t buft = ggml_backend_buffer_get_type(buf.get()); - - if (hparams.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 { - // GGML_ASSERT(ggml_backend_buffer_get_base(buf.get()) != nullptr); // multi_buffer does not have a defined base - ret[buft] += ggml_backend_buffer_get_size(buf.get()); - } - } - - return ret; + return ctxs_bufs.memory_breakdown(); } llama_memory_context_ptr llama_kv_cache::init_batch( diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index 9bfb8f231..b31cc226d 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -98,6 +98,7 @@ public: // likely through `struct llama_memory_params` llama_kv_cache( const llama_model & model, + bool no_alloc, const llama_hparams & hparams, ggml_type type_k, ggml_type type_v, @@ -296,7 +297,7 @@ private: const llama_swa_type swa_type = LLAMA_SWA_TYPE_NONE; // ggml contexts for the KV cache along with the allocated backend buffers: - std::vector> ctxs_bufs; + llama_memory_bufs ctxs_bufs; // the current index from where we start searching for a free slot in the ring buffer of KV cells (see find_slot()) // note: this is not part of the KV state and it's only used to speed-up the find_slot() method diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index 8113f5c21..03f0c342d 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/src/llama-memory-hybrid-idx.cpp @@ -1,20 +1,16 @@ #include "llama-memory-hybrid-idx.h" -#include -#include -#include - #include "llama-impl.h" #include "llama-batch.h" #include "llama-io.h" #include "llama-model.h" - #include #include #include #include #include +#include // // llama_memory_hybrid_idx @@ -22,6 +18,7 @@ llama_memory_hybrid_idx::llama_memory_hybrid_idx( const llama_model & model, + bool no_alloc, /* attn */ ggml_type type_k, ggml_type type_v, @@ -45,6 +42,7 @@ llama_memory_hybrid_idx::llama_memory_hybrid_idx( const layer_filter_cb & filter_idx) : llama_memory_hybrid( model, + no_alloc, 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, @@ -68,7 +66,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, + model, no_alloc, 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 0ac4265df..faa1724b0 100644 --- a/src/llama-memory-hybrid-idx.h +++ b/src/llama-memory-hybrid-idx.h @@ -18,6 +18,7 @@ class llama_memory_hybrid_idx : public llama_memory_hybrid { public: llama_memory_hybrid_idx( const llama_model & model, + bool no_alloc, /* 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 06f7fd542..4c0f14c22 100644 --- a/src/llama-memory-hybrid-iswa.cpp +++ b/src/llama-memory-hybrid-iswa.cpp @@ -10,6 +10,7 @@ llama_memory_hybrid_iswa::llama_memory_hybrid_iswa( const llama_model & model, + bool no_alloc, /* attn */ ggml_type type_k, ggml_type type_v, @@ -33,6 +34,7 @@ llama_memory_hybrid_iswa::llama_memory_hybrid_iswa( hparams(model.hparams), mem_attn(new llama_kv_cache_iswa( model, + no_alloc, type_k, type_v, v_trans, @@ -52,6 +54,7 @@ llama_memory_hybrid_iswa::llama_memory_hybrid_iswa( )), mem_recr(new llama_memory_recurrent( model, + no_alloc, type_r, type_s, offload, diff --git a/src/llama-memory-hybrid-iswa.h b/src/llama-memory-hybrid-iswa.h index c9d3f9f57..6e68f3f44 100644 --- a/src/llama-memory-hybrid-iswa.h +++ b/src/llama-memory-hybrid-iswa.h @@ -20,6 +20,7 @@ class llama_memory_hybrid_iswa : public llama_memory_i { public: llama_memory_hybrid_iswa( const llama_model & model, + bool no_alloc, /* attn */ ggml_type type_k, ggml_type type_v, diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp index bbd4bee0b..0361b23ca 100644 --- a/src/llama-memory-hybrid.cpp +++ b/src/llama-memory-hybrid.cpp @@ -10,6 +10,7 @@ llama_memory_hybrid::llama_memory_hybrid( const llama_model & model, + bool no_alloc, /* attn */ ggml_type type_k, ggml_type type_v, @@ -33,6 +34,7 @@ llama_memory_hybrid::llama_memory_hybrid( hparams(model.hparams), mem_attn(new llama_kv_cache( model, + no_alloc, model.hparams, type_k, type_v, @@ -53,6 +55,7 @@ llama_memory_hybrid::llama_memory_hybrid( )), mem_recr(new llama_memory_recurrent( model, + no_alloc, type_r, type_s, offload, diff --git a/src/llama-memory-hybrid.h b/src/llama-memory-hybrid.h index 484eafb74..74fa89ea6 100644 --- a/src/llama-memory-hybrid.h +++ b/src/llama-memory-hybrid.h @@ -20,6 +20,7 @@ class llama_memory_hybrid : public llama_memory_i { public: llama_memory_hybrid( const llama_model & model, + bool no_alloc, /* attn */ ggml_type type_k, ggml_type type_v, diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp index d00630d4c..5fd644645 100644 --- a/src/llama-memory-recurrent.cpp +++ b/src/llama-memory-recurrent.cpp @@ -19,13 +19,14 @@ llama_memory_recurrent::llama_memory_recurrent( const llama_model & model, + bool no_alloc, 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) : hparams(model.hparams), n_seq_max(n_seq_max), ctxs_bufs(no_alloc) { const int32_t n_layer = hparams.n_layer(); head = 0; @@ -114,23 +115,11 @@ llama_memory_recurrent::llama_memory_recurrent( } } - // allocate tensors and initialize the buffers to avoid NaNs in the padding + // allocate the 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) { - 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 RS cache so that the backend scheduler won't try to allocate it - } - } else { - buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx.get(), buft); // real buffer - } - if (!buf) { - throw std::runtime_error("failed to allocate buffer for rs cache"); - } - ggml_backend_buffer_clear(buf, 0); + ggml_backend_buffer_t buf = ctxs_bufs.emplace(std::move(ctx), buft, "rs cache"); + LLAMA_LOG_INFO("%s: %10s RS buffer size = %8.2f MiB\n", __func__, ggml_backend_buffer_name(buf), ggml_backend_buffer_get_size(buf)/1024.0/1024.0); - ctxs_bufs.emplace_back(std::move(ctx), buf); } if (is_empty()) { @@ -442,18 +431,7 @@ 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 & [ctx, buf] : ctxs_bufs) { - ggml_backend_buffer_type_t buft = ggml_backend_buffer_get_type(buf.get()); - - if (hparams.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 { - ret[buft] += ggml_backend_buffer_get_size(buf.get()); - } - } - return ret; + return ctxs_bufs.memory_breakdown(); } llama_memory_context_ptr llama_memory_recurrent::init_batch(llama_batch_allocr & balloc, uint32_t n_ubatch, bool embd_all) { diff --git a/src/llama-memory-recurrent.h b/src/llama-memory-recurrent.h index 08489a4be..d0fc14779 100644 --- a/src/llama-memory-recurrent.h +++ b/src/llama-memory-recurrent.h @@ -18,6 +18,7 @@ class llama_memory_recurrent : public llama_memory_i { public: llama_memory_recurrent( const llama_model & model, + bool no_alloc, ggml_type type_r, ggml_type type_s, bool offload, @@ -120,8 +121,8 @@ private: const uint32_t n_seq_max = 1; - // ggml contexts for the KV cache along with the allocated backend buffers: - std::vector> ctxs_bufs; + // ggml contexts for the RS cache along with the allocated backend buffers: + llama_memory_bufs ctxs_bufs; // true if no layers - can happen if the layer filter removes all layers bool is_empty() const; diff --git a/src/llama-memory.cpp b/src/llama-memory.cpp index 961dd74ef..042aab608 100644 --- a/src/llama-memory.cpp +++ b/src/llama-memory.cpp @@ -4,6 +4,8 @@ #include "llama-io.h" #include +#include +#include #include llama_memory_status llama_memory_status_combine(llama_memory_status s0, llama_memory_status s1) { @@ -64,6 +66,52 @@ bool llama_memory_status_is_fail(llama_memory_status status) { return false; } +// +// llama_memory_bufs +// + +ggml_backend_buffer_ptr llama_memory_bufs::alloc_buf(ggml_context * ctx, ggml_backend_buffer_type_t buft, bool no_alloc, const char * name) { + ggml_backend_buffer_ptr buf; + if (no_alloc) { + buf.reset(ggml_backend_buft_alloc_buffer(buft, /*size =*/ 0)); // dummy buffer + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { + t->buffer = buf.get(); // set dummy buffer so that the backend scheduler won't try to allocate the tensors + } + } else { + buf.reset(ggml_backend_alloc_ctx_tensors_from_buft(ctx, buft)); // real buffer + } + if (!buf) { + throw std::runtime_error(std::string("failed to allocate buffer for ") + name); + } + return buf; +} + +ggml_backend_buffer_t llama_memory_bufs::emplace(ggml_context_ptr ctx, ggml_backend_buffer_type_t buft, const char * name) { + ggml_backend_buffer_ptr buf = alloc_buf(ctx.get(), buft, no_alloc, name); + + // initialize the buffer to avoid NaNs in the padding + ggml_backend_buffer_clear(buf.get(), 0); + + bufs.emplace_back(std::move(ctx), std::move(buf)); + + return bufs.back().second.get(); +} + +std::map llama_memory_bufs::memory_breakdown() const { + std::map ret; + for (const auto & [ctx, buf] : bufs) { + ggml_backend_buffer_type_t buft = ggml_backend_buffer_get_type(buf.get()); + + if (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 { + ret[buft] += ggml_backend_buffer_get_size(buf.get()); + } + } + return ret; +} + // collects the buffers of the tensors that a state write visits, without keeping the data struct llama_memory_buffers_writer : public llama_io_write_i { std::set buffers; @@ -84,6 +132,10 @@ struct llama_memory_buffers_writer : public llama_io_write_i { }; std::vector llama_memory_get_buffers(const llama_memory_t mem) { + if (mem == nullptr) { + return {}; + } + llama_memory_buffers_writer writer; mem->state_write(writer); diff --git a/src/llama-memory.h b/src/llama-memory.h index b33c7aee4..37d01ed11 100644 --- a/src/llama-memory.h +++ b/src/llama-memory.h @@ -6,6 +6,7 @@ #include #include #include +#include struct llama_ubatch; @@ -25,6 +26,36 @@ struct llama_memory_params { llama_context_type ctx_type; llama_memory_t mem_other; + + // only simulate the memory allocations + bool no_alloc; +}; + +// the tensor contexts and backend buffers of a memory instance +// with no_alloc, dummy buffers are attached instead of allocating the tensor data, so the memory use can still be measured +struct llama_memory_bufs { + explicit llama_memory_bufs(bool no_alloc) : no_alloc(no_alloc) {} + + // allocate the tensors of ctx on buft, or attach a dummy buffer if no_alloc is set + // returns the buffer, for logging + ggml_backend_buffer_t emplace(ggml_context_ptr ctx, ggml_backend_buffer_type_t buft, const char * name); + + // the memory used by the tensors: the buffer sizes, or the simulated sizes if no_alloc is set + std::map memory_breakdown() const; + + // allocate a single ctx buffer - for the memory instances that keep their own buffer layout + static ggml_backend_buffer_ptr alloc_buf(ggml_context * ctx, ggml_backend_buffer_type_t buft, bool no_alloc, const char * name); + + bool empty() const { return bufs.empty(); } + + auto begin() { return bufs.begin(); } + auto end() { return bufs.end(); } + auto begin() const { return bufs.begin(); } + auto end() const { return bufs.end(); } + + std::vector> bufs; + + const bool no_alloc; }; enum llama_memory_status { diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 1de65c728..78a387e5d 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2448,10 +2448,14 @@ llama_split_mode llama_model::split_mode() const { return params.split_mode; } +bool llama_model::no_alloc() const { + return params.no_alloc; +} + std::map llama_model::memory_breakdown() const { std::map ret; for (const auto & [ctx, bufs] : pimpl->ctxs_bufs) { - if (hparams.no_alloc) { + if (no_alloc()) { GGML_ASSERT(bufs.size() == 1); ggml_backend_buffer_t buf = bufs[0].get(); GGML_ASSERT(ggml_backend_buffer_get_base(buf) == nullptr); @@ -2506,7 +2510,7 @@ void llama_model::print_info() const { // hparams LLAMA_LOG_INFO("%s: arch = %s\n", __func__, arch_name().c_str()); LLAMA_LOG_INFO("%s: vocab_only = %d\n", __func__, hparams.vocab_only); - LLAMA_LOG_INFO("%s: no_alloc = %d\n", __func__, hparams.no_alloc); + LLAMA_LOG_INFO("%s: no_alloc = %d\n", __func__, no_alloc()); if (!hparams.vocab_only) { LLAMA_LOG_INFO("%s: n_ctx_train = %u\n", __func__, hparams.n_ctx_train); @@ -2830,6 +2834,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, res = new llama_kv_cache_msa( *this, + params.no_alloc, params.type_k, params.type_v, !cparams.flash_attn, @@ -2856,6 +2861,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, res = new llama_kv_cache( *this, + params.no_alloc, hparams, params.type_k, params.type_v, @@ -2882,6 +2888,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, res = new llama_kv_cache_dsa( *this, + params.no_alloc, params.type_k, params.type_v, !cparams.flash_attn, @@ -2923,6 +2930,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, res = new llama_memory_hybrid_idx( /* model */ *this, + /* no_alloc */ params.no_alloc, /* attn_type_k */ params.type_k, /* attn_type_v */ params.type_v, /* attn_v_trans */ !cparams.flash_attn, @@ -2947,6 +2955,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, // full-attention checkpoint: no indexer, so no indexer key cache res = new llama_kv_cache( *this, + params.no_alloc, hparams, params.type_k, params.type_v, @@ -2968,6 +2977,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, res = new llama_kv_cache_dsa( *this, + params.no_alloc, params.type_k, params.type_v, !cparams.flash_attn, @@ -2994,6 +3004,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, res = new llama_kv_cache( *this, + params.no_alloc, hparams, params.type_k, params.type_v, @@ -3019,6 +3030,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, res = new llama_kv_cache_dsa_iswa( *this, + params.no_alloc, params.type_k, params.type_v, !cparams.flash_attn, @@ -3046,6 +3058,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, res = new llama_kv_cache_iswa( *this, + params.no_alloc, params.type_k, params.type_v, !cparams.flash_attn, @@ -3063,6 +3076,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, } else { res = new llama_kv_cache_dsv4( *this, + params.no_alloc, params.type_k, params.type_v, !cparams.flash_attn, @@ -3086,6 +3100,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, res = new llama_kv_cache_iswa( *this, + params.no_alloc, params.type_k, params.type_v, !cparams.flash_attn, @@ -3120,6 +3135,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, GGML_TYPE_F32, GGML_TYPE_F32, cparams.offload_kqv, @@ -3175,6 +3191,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, // Use hybrid-iswa for hybrid models with SWA res = new llama_memory_hybrid_iswa( /* model */ *this, + /* no_alloc */ params.no_alloc, /* attn_type_k */ params.type_k, /* attn_type_v */ params.type_v, /* attn_v_trans */ !cparams.flash_attn, @@ -3195,6 +3212,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, // sparse attention over a per-token indexer cache, in its own memory type res = new llama_memory_hybrid_idx( /* model */ *this, + /* no_alloc */ params.no_alloc, /* attn_type_k */ params.type_k, /* attn_type_v */ params.type_v, /* attn_v_trans */ !cparams.flash_attn, @@ -3215,6 +3233,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, } else { res = new llama_memory_hybrid( /* model */ *this, + /* no_alloc */ params.no_alloc, /* attn_type_k */ params.type_k, /* attn_type_v */ params.type_v, /* attn_v_trans */ !cparams.flash_attn, @@ -3281,6 +3300,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, res = new llama_kv_cache_iswa( *this, + params.no_alloc, params.type_k, params.type_v, !cparams.flash_attn, @@ -3298,6 +3318,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, } else { res = new llama_kv_cache_iswa( *this, + params.no_alloc, params.type_k, params.type_v, !cparams.flash_attn, @@ -3318,6 +3339,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, res = new llama_kv_cache( *this, + params.no_alloc, hparams, params.type_k, params.type_v, diff --git a/src/llama-model.h b/src/llama-model.h index ccc2602e5..222238efe 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -886,6 +886,7 @@ struct llama_model { uint32_t n_gpu_layers() const; llama_split_mode split_mode() const; + bool no_alloc() const; std::map memory_breakdown() const; diff --git a/src/llama-moe-cache.cpp b/src/llama-moe-cache.cpp index 4d3596b24..6a343a760 100644 --- a/src/llama-moe-cache.cpp +++ b/src/llama-moe-cache.cpp @@ -1,6 +1,7 @@ #include "llama-moe-cache.h" #include "llama-impl.h" +#include "llama-memory.h" #include "llama-model.h" #include "ggml-cpp.h" @@ -252,7 +253,7 @@ struct llama_moe_cache::impl { // views used by copy_experts ggml_context_ptr ctx_views; - impl(const llama_model & model, const std::vector & backends, const std::vector & bufts, size_t size) : + impl(const llama_model & model, bool no_alloc, const std::vector & backends, const std::vector & bufts, size_t size) : n_expert_used(model.hparams.n_expert_used_max()), layers(model.layers.size()) { for (size_t i = 0; i < backends.size(); ++i) { const auto dev_type = ggml_backend_dev_type(ggml_backend_get_device(backends[i])); @@ -418,37 +419,19 @@ struct llama_moe_cache::impl { devices[g.id].buf_size += alloc_size(g, g.n_slots); } - if (model.hparams.no_alloc) { - // only used to measure the memory use, see llama_context::memory_breakdown - for (device & d : devices) { - if (!d.ctx) { - continue; - } - d.buf.reset(ggml_backend_buft_alloc_buffer(d.buft, 0)); - for (ggml_tensor * t = ggml_get_first_tensor(d.ctx.get()); t != nullptr; t = ggml_get_next_tensor(d.ctx.get(), t)) { - t->buffer = d.buf.get(); - } + // with no_alloc, dummy buffers are attached so the memory use can still be measured, see llama_context::memory_breakdown + for (device & d : devices) { + if (!d.ctx) { + continue; } - buf_host.reset(ggml_backend_buft_alloc_buffer(buft_host, 0)); - for (ggml_tensor * t = ggml_get_first_tensor(ctx_host.get()); t != nullptr; t = ggml_get_next_tensor(ctx_host.get(), t)) { - t->buffer = buf_host.get(); - } - } else { - for (device & d : devices) { - if (!d.ctx) { - continue; - } - d.buf.reset(ggml_backend_alloc_ctx_tensors_from_buft(d.ctx.get(), d.buft)); - if (!d.buf) { - throw std::runtime_error("failed to allocate the MoE cache buffers"); - } + d.buf = llama_memory_bufs::alloc_buf(d.ctx.get(), d.buft, no_alloc, "MoE cache"); + if (!no_alloc) { ggml_backend_buffer_clear(d.buf.get(), 0); d.buf_size = ggml_backend_buffer_get_size(d.buf.get()); } - buf_host.reset(ggml_backend_alloc_ctx_tensors_from_buft(ctx_host.get(), buft_host)); - if (!buf_host) { - throw std::runtime_error("failed to allocate the MoE cache buffers"); - } + } + buf_host = llama_memory_bufs::alloc_buf(ctx_host.get(), buft_host, no_alloc, "MoE cache slot maps"); + if (!no_alloc) { ggml_backend_buffer_clear(buf_host.get(), 0xff); // all slots are -1 buf_host_size = ggml_backend_buffer_get_size(buf_host.get()); @@ -633,8 +616,8 @@ struct llama_moe_cache::impl { } }; -llama_moe_cache::llama_moe_cache(const llama_model & model, const std::vector & backends, const std::vector & bufts, size_t size) : - pimpl(new impl(model, backends, bufts, size)) { +llama_moe_cache::llama_moe_cache(const llama_model & model, bool no_alloc, const std::vector & backends, const std::vector & bufts, size_t size) : + pimpl(new impl(model, no_alloc, backends, bufts, size)) { } llama_moe_cache::~llama_moe_cache() = default; diff --git a/src/llama-moe-cache.h b/src/llama-moe-cache.h index 9baf06390..6c8e36582 100644 --- a/src/llama-moe-cache.h +++ b/src/llama-moe-cache.h @@ -13,7 +13,7 @@ struct llama_model; class llama_moe_cache { public: // backends are all the backends of the context, each GPU gets its own cache of the given size for the layers assigned to it - llama_moe_cache(const llama_model & model, const std::vector & backends, const std::vector & bufts, size_t size); + llama_moe_cache(const llama_model & model, bool no_alloc, const std::vector & backends, const std::vector & bufts, size_t size); ~llama_moe_cache(); // the device that caches layer il diff --git a/src/llama.cpp b/src/llama.cpp index 9beb14841..e3699ea1d 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -342,7 +342,6 @@ static std::pair llama_model_load(struct gguf_context * meta model->t_start_us = tm.t_start_us; model->hparams.vocab_only = params.vocab_only; - model->hparams.no_alloc = params.no_alloc; try { model->load_hparams(ml); diff --git a/tests/llama-context/cases-kv.cpp b/tests/llama-context/cases-kv.cpp index 27e6493e4..5086e0fb8 100644 --- a/tests/llama-context/cases-kv.cpp +++ b/tests/llama-context/cases-kv.cpp @@ -718,6 +718,54 @@ struct case_state_restore_fragmented : ctx_test { } }; +// +// no_alloc +// + +// - load a second instance of the model with no_alloc, so the tensor data stays on disk +// - create a context from it, which reserves the compute graphs over the dummy buffers +// - the memory use must still be reported, and nothing real may be allocated for the memory +struct case_noalloc : ctx_test { + const char * name() const override { return "noalloc"; } + + const char * description() const override { + return "with no_alloc the memory use is simulated: non-zero breakdown, null buffer bases"; + } + + case_result run(model_run & mr) const override { + const std::string & path = mr.params.model.path; + CTX_REQUIRE(!path.empty(), "the model path is not known"); + + llama_model_params mparams = llama_model_default_params(); + mparams.no_alloc = true; + mparams.load_mode = LLAMA_LOAD_MODE_NONE; + + llama_model_ptr model(llama_model_load_from_file(path.c_str(), mparams)); + CTX_REQUIRE(model, "failed to load the model with no_alloc"); + + auto ctx = init_ctx(model.get(), common_context_params_to_llama(mr.params), 0); + CTX_REQUIRE(ctx, "failed to create the context with no_alloc"); + + size_t n_ctx_mem = 0; + for (const auto & [_, data] : llama_get_memory_breakdown(ctx.get())) { + n_ctx_mem += data.context; + } + + // models without a memory module (BERT, diffusion) have a null memory + if (llama_get_memory(ctx.get()) != nullptr) { + CTX_REQUIRE(n_ctx_mem > 0, "the simulated memory breakdown is zero"); + } else { + CTX_REQUIRE(n_ctx_mem == 0, "a model without memory reported context memory"); + } + + for (ggml_backend_buffer_t buf : llama_memory_get_buffers(llama_get_memory(ctx.get()))) { + CTX_REQUIRE(ggml_backend_buffer_get_base(buf) == nullptr, "a memory buffer was really allocated"); + } + + return case_result::pass(); + } +}; + std::vector make_kv_cases() { std::vector cases; @@ -733,6 +781,7 @@ std::vector make_kv_cases() { cases.push_back(std::make_unique()); cases.push_back(std::make_unique()); cases.push_back(std::make_unique()); + cases.push_back(std::make_unique()); return cases; }