diff --git a/src/llama-context.cpp b/src/llama-context.cpp index f21767601d..ef53728d1d 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -599,8 +599,11 @@ void llama_context::sched_reserve() { LLAMA_LOG_DEBUG("%s: max_nodes = %zu\n", __func__, max_nodes); - gf_res_prev.reset(new llm_graph_result(max_nodes)); + for (auto & res : gf_res_prev) { + res.reset(); + } gf_res_reserve.reset(new llm_graph_result(max_nodes)); + gf_res_prev_active = nullptr; sched.reset(ggml_backend_sched_new(backend_ptrs.data(), backend_buft.data(), backend_ptrs.size(), max_nodes, cparams.pipeline_parallel, cparams.op_offload)); @@ -816,10 +819,14 @@ bool llama_context::memory_update(bool optimize) { } } - // reset the previous graph result to make sure that it won't be reused - // TODO: change the mctx->apply() to return information if a graph reserve is needed - // reset the graph result only if the memory module did reset the scheduler - gf_res_prev->reset(); + // reset the previous graph results to make sure that they won't be reused + // TODO: make mctx->apply() report if a graph reserve is needed, then reset graph results only if the memory module reset the scheduler + for (auto & res : gf_res_prev) { + if (res) { + res->reset(); + } + } + gf_res_prev_active = nullptr; if (!mctx->apply()) { LLAMA_LOG_ERROR("%s: failed to apply memory update\n", __func__); @@ -1340,14 +1347,14 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll return nullptr; } - auto * res = gf_res_prev.get(); + auto * res = get_gf_res_prev(); auto * gf = res->get_gf(); // the new graph parameters // in order to correctly reuse a graph, it's full topology has to be uniquely determined by these parameters const auto gparams = graph_params(res, ubatch, mctx, gtype); - if (!graph_reuse_disable && res->can_reuse(gparams)) { + if (!graph_reuse_disable && gf_res_prev_active == res && res->can_reuse(gparams)) { //LLAMA_LOG_DEBUG("%s: reusing previous graph\n", __func__); // with pipeline parallelism, the previous graph_compute_async may still be running @@ -1359,6 +1366,7 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll n_reused++; } else { + gf_res_prev_active = nullptr; res->reset(); ggml_backend_sched_reset(sched.get()); @@ -1381,6 +1389,8 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll ret = GGML_STATUS_ALLOC_FAILED; return nullptr; } + + gf_res_prev_active = res; } // set the input data for the input tensors @@ -2361,6 +2371,14 @@ llm_graph_result * llama_context::get_gf_res_reserve() const { return static_cast(gf_res_reserve.get()); } +llm_graph_result * llama_context::get_gf_res_prev() { + auto & res = gf_res_prev[n_outputs > 0]; + if (!res) { + res.reset(new llm_graph_result(gf_res_reserve->get_max_nodes())); + } + return res.get(); +} + // pack sampler outputs into as few sequences as possible before using sequences without samplers static void ubatch_prepare_reserve( llama_ubatch & ubatch, @@ -2430,8 +2448,13 @@ ggml_cgraph * llama_context::graph_reserve( ggml_backend_sched_reset(sched.get()); - // when the scheduler is reset, we cannot reuse the old graph, so we reset the previous graph result to prevent that - gf_res_prev->reset(); + // when the scheduler is reset, we cannot reuse old graphs, so we reset the previous graph results + for (auto & res : gf_res_prev) { + if (res) { + res->reset(); + } + } + gf_res_prev_active = nullptr; // store the n_outputs as it is, and restore it afterwards // TODO: not sure if needed, might simplify in the future by removing this @@ -3521,10 +3544,12 @@ void llama_context::opt_epoch_iter( break; } - auto * res = gf_res_prev.get(); + auto * res = get_gf_res_prev(); const auto gparams = graph_params(res, ubatch, mctx.get(), ctx_type_to_graph_type(cparams.ctx_type)); + // the optimizer graph is allocated outside sched, so the next decode must rebuild + gf_res_prev_active = nullptr; res->reset(); auto * gf = model.build_graph(gparams); diff --git a/src/llama-context.h b/src/llama-context.h index bf91daa8b5..b7a9db5913 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -11,6 +11,7 @@ #include "ggml-cpp.h" #include "ggml-opt.h" +#include #include #include @@ -254,6 +255,8 @@ public: bool set_sampler(llama_seq_id seq_id, llama_sampler * sampler); private: + llm_graph_result * get_gf_res_prev(); + llm_graph_params graph_params( llm_graph_result * res, const llama_ubatch & ubatch, @@ -364,9 +367,12 @@ private: std::vector backend_buft; std::vector backend_buf_exp_size; // expected buffer sizes - llm_graph_result_ptr gf_res_prev; + // Separate arenas give batches with and without outputs distinct CUDA graph cache keys. + std::array gf_res_prev; llm_graph_result_ptr gf_res_reserve; + llm_graph_result * gf_res_prev_active = nullptr; + // host buffer for the model output (logits and embeddings) ggml_backend_buffer_ptr buf_output;