kv-cache : replace struct graph_params with llama_context &

ggerganov · ggerganov · commit eb623f2fcd53 · 2025-04-30T11:30:29.000+03:00
ggml-ci
diff --git a/src/llama-context.cpp b/src/llama-context.cpp
@@ -396,6 +396,22 @@ const llama_model & llama_context::get_model() const {
     return model;
 }
 
+const llama_cparams & llama_context::get_cparams() const {
+    return cparams;
+}
+
+const ggml_backend_sched_ptr & llama_context::get_sched() const {
+    return sched;
+}
+
+const ggml_context_ptr & llama_context::get_ctx_compute() const {
+    return ctx_compute;
+}
+
+const std::vector<ggml_backend_ptr> & llama_context::get_backends() const {
+    return backends;
+}
+
 uint32_t llama_context::n_ctx() const {
     return cparams.n_ctx;
 }
@@ -439,15 +455,7 @@ void llama_context::kv_self_update() {
 
     llama_kv_cache * kv_self = static_cast<llama_kv_cache *>(memory.get());
 
-    need_reserve = kv_self->update({
-        /*.cparams         =*/ cparams,
-        /*.sched           =*/ sched.get(),
-        /*.backends        =*/ backends,
-        /*.n_max_nodes     =*/ graph_max_nodes(),
-        /*.get_ctx_compute =*/ [this]() { return ctx_compute.get(); },
-        /*.graph_init      =*/ [this]() { return graph_init(); },
-        /*.graph_compute   =*/ [this](ggml_cgraph * gf) { graph_compute(gf, false); },
-    });
+    need_reserve = kv_self->update(*this);
 
     // reserve a worst case graph if needed
     if (need_reserve) {
diff --git a/src/llama-context.h b/src/llama-context.h
@@ -27,7 +27,14 @@ struct llama_context {
 
     void synchronize();
 
-    const llama_model & get_model() const;
+    const llama_model   & get_model()   const;
+    const llama_cparams & get_cparams() const;
+
+    const ggml_backend_sched_ptr & get_sched() const;
+
+    const ggml_context_ptr & get_ctx_compute() const;
+
+    const std::vector<ggml_backend_ptr> & get_backends() const;
 
     uint32_t n_ctx()         const;
     uint32_t n_ctx_per_seq() const;
@@ -141,22 +148,24 @@ struct llama_context {
     // graph
     //
 
+public:
     int32_t graph_max_nodes() const;
 
     // zero-out inputs and create the ctx_compute for the compute graph
     ggml_cgraph * graph_init();
 
+    // returns the result of ggml_backend_sched_graph_compute_async execution
+    ggml_status graph_compute(
+            ggml_cgraph * gf,
+                   bool   batched);
+
+private:
     llm_graph_result_ptr graph_build(
             ggml_context * ctx,
              ggml_cgraph * gf,
       const llama_ubatch & ubatch,
           llm_graph_type   gtype);
 
-    // returns the result of ggml_backend_sched_graph_compute_async execution
-    ggml_status graph_compute(
-            ggml_cgraph * gf,
-                   bool   batched);
-
     llm_graph_cb graph_get_cb() const;
 
     // TODO: read/write lora adapters and cvec
diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp
@@ -4,6 +4,7 @@
 #include "llama-batch.h"
 #include "llama-cparams.h"
 #include "llama-model.h"
+#include "llama-context.h"
 
 #include <algorithm>
 #include <cassert>
@@ -367,10 +368,10 @@ void llama_kv_cache_unified::commit() {
     pending.ranges.clear();
 }
 
-bool llama_kv_cache_unified::update(const graph_params & params) {
+bool llama_kv_cache_unified::update(llama_context & lctx) {
     bool need_reserve = false;
 
-    const auto & sched = params.sched;
+    const auto & sched = lctx.get_sched();
 
     if (has_shift) {
         if (!get_can_shift()) {
@@ -381,17 +382,17 @@ bool llama_kv_cache_unified::update(const graph_params & params) {
 
         // apply K-shift if needed
         if (hparams.rope_type != LLAMA_ROPE_TYPE_NONE) {
-            ggml_backend_sched_reset(sched);
+            ggml_backend_sched_reset(sched.get());
 
-            auto * gf = params.graph_init();
+            auto * gf = lctx.graph_init();
 
-            auto res = build_graph_shift(params, gf);
+            auto res = build_graph_shift(lctx, gf);
 
-            ggml_backend_sched_alloc_graph(sched, gf);
+            ggml_backend_sched_alloc_graph(sched.get(), gf);
 
             res->set_inputs(nullptr);
 
-            params.graph_compute(gf);
+            lctx.graph_compute(gf, false);
 
             need_reserve = true;
         }
@@ -408,18 +409,18 @@ bool llama_kv_cache_unified::update(const graph_params & params) {
     if (do_defrag) {
         LLAMA_LOG_DEBUG("%s: defragmenting KV cache\n", __func__);
 
-        if (defrag_prepare(params.n_max_nodes)) {
-            ggml_backend_sched_reset(sched);
+        if (defrag_prepare(lctx.graph_max_nodes())) {
+            ggml_backend_sched_reset(sched.get());
 
-            auto * gf = params.graph_init();
+            auto * gf = lctx.graph_init();
 
-            auto res = build_graph_defrag(params, gf);
+            auto res = build_graph_defrag(lctx, gf);
 
-            ggml_backend_sched_alloc_graph(sched, gf);
+            ggml_backend_sched_alloc_graph(sched.get(), gf);
 
             res->set_inputs(nullptr);
 
-            params.graph_compute(gf);
+            lctx.graph_compute(gf, false);
 
             need_reserve = true;
         }
@@ -591,17 +592,17 @@ size_t llama_kv_cache_unified::size_v_bytes() const {
 }
 
 ggml_tensor * llama_kv_cache_unified::build_rope_shift(
-        const graph_params & params,
-              ggml_context * ctx,
-               ggml_tensor * cur,
-               ggml_tensor * shift,
-               ggml_tensor * factors,
-                     float   freq_base,
-                     float   freq_scale,
-       ggml_backend_buffer * bbuf) const {
-    const auto & cparams  = params.cparams;
-    const auto & backends = params.backends;
-    const auto & sched    = params.sched;
+        llama_context & lctx,
+         ggml_context * ctx,
+          ggml_tensor * cur,
+          ggml_tensor * shift,
+          ggml_tensor * factors,
+                float   freq_base,
+                float   freq_scale,
+  ggml_backend_buffer * bbuf) const {
+    const auto & cparams  = lctx.get_cparams();
+    const auto & backends = lctx.get_backends();
+    const auto & sched    = lctx.get_sched();
 
     const auto & n_ctx_orig = cparams.n_ctx_orig_yarn;
 
@@ -622,11 +623,12 @@ ggml_tensor * llama_kv_cache_unified::build_rope_shift(
         // dequantize to f32 -> RoPE -> quantize back
         tmp = ggml_cast(ctx, cur, GGML_TYPE_F32);
 
+        // TODO: can we simplify/avoid this?
         if (bbuf) {
             for (const auto & backend : backends) {
                 // Figure out which backend KV cache belongs to
                 if (ggml_backend_supports_buft(backend.get(), ggml_backend_buffer_get_type(bbuf))) {
-                    ggml_backend_sched_set_tensor_backend(sched, tmp, backend.get());
+                    ggml_backend_sched_set_tensor_backend(sched.get(), tmp, backend.get());
                     break;
                 }
             }
@@ -674,13 +676,13 @@ void llm_graph_input_k_shift::set_input(const llama_ubatch * ubatch) {
 }
 
 llm_graph_result_ptr llama_kv_cache_unified::build_graph_shift(
-            const graph_params & params,
-                   ggml_cgraph * gf) const {
+        llama_context & lctx,
+        ggml_cgraph * gf) const {
     auto res = std::make_unique<llm_graph_result>();
 
-    auto * ctx = params.get_ctx_compute();
+    auto * ctx = lctx.get_ctx_compute().get();
 
-    const auto & cparams = params.cparams;
+    const auto & cparams = lctx.get_cparams();
 
     const auto & n_layer = hparams.n_layer;
 
@@ -716,7 +718,7 @@ llm_graph_result_ptr llama_kv_cache_unified::build_graph_shift(
                 ggml_row_size(k_l[il]->type, n_embd_k_gqa),
                 0);
 
-        ggml_tensor * cur = build_rope_shift(params, ctx, k, inp->k_shift, rope_factors, freq_base_l, freq_scale_l, k_l[il]->buffer);
+        ggml_tensor * cur = build_rope_shift(lctx, ctx, k, inp->k_shift, rope_factors, freq_base_l, freq_scale_l, k_l[il]->buffer);
 
         ggml_build_forward_expand(gf, cur);
     }
@@ -727,15 +729,15 @@ llm_graph_result_ptr llama_kv_cache_unified::build_graph_shift(
 }
 
 llm_graph_result_ptr llama_kv_cache_unified::build_graph_defrag(
-            const graph_params & params,
-                   ggml_cgraph * gf) const {
+        llama_context & lctx,
+          ggml_cgraph * gf) const {
     auto res = std::make_unique<llm_graph_result>();
 
-    auto * ctx = params.get_ctx_compute();
+    auto * ctx = lctx.get_ctx_compute().get();
 
     const auto & ids = defrag_info.ids;
 
-    const auto & cparams = params.cparams;
+    const auto & cparams = lctx.get_cparams();
 
 #if 0
     // CPU defrag
@@ -1725,8 +1727,8 @@ void llama_kv_cache_recurrent::commit() {
     pending.ranges.clear();
 }
 
-bool llama_kv_cache_recurrent::update(const graph_params & params) {
-    GGML_UNUSED(params);
+bool llama_kv_cache_recurrent::update(llama_context & lctx) {
+    GGML_UNUSED(lctx);
     return false;
 }
 
diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h
@@ -7,7 +7,6 @@
 
 #include "ggml-cpp.h"
 
-#include <functional>
 #include <set>
 #include <vector>
 
@@ -16,26 +15,9 @@ struct llama_hparams;
 struct llama_ubatch;
 struct llama_sbatch;
 struct llama_model;
+struct llama_context;
 
 struct llama_kv_cache : public llama_memory_i {
-    struct graph_params {
-        const llama_cparams & cparams;
-
-        const ggml_backend_sched_t & sched;
-
-        const std::vector<ggml_backend_ptr> & backends;
-
-        int32_t n_max_nodes;
-
-        std::function<ggml_context * ()> get_ctx_compute;
-
-        // function for creating ggml graphs
-        std::function<ggml_cgraph * ()> graph_init;
-
-        // function for computing ggml graphs
-        std::function<void(ggml_cgraph * gf)> graph_compute;
-    };
-
     virtual ~llama_kv_cache() = default;
 
     // call if batch processing fails - restores the cache state
@@ -46,7 +28,7 @@ struct llama_kv_cache : public llama_memory_i {
 
     // process any pending defrag/shift/etc. operations
     // optionally call once before processing a new batch
-    virtual bool update(const graph_params & params) = 0;
+    virtual bool update(llama_context & lctx) = 0;
 
     // schedule a defrag if the fragmentation threshold is exceeded. otherwise, do nothing
     virtual void defrag_sched(float thold) = 0;
@@ -161,7 +143,7 @@ class llama_kv_cache_unified : public llama_kv_cache {
     void restore() override;
     void commit()  override;
 
-    bool update(const graph_params & params) override;
+    bool update(llama_context & ctx) override;
 
     void defrag_sched(float thold) override;
 
@@ -251,22 +233,22 @@ class llama_kv_cache_unified : public llama_kv_cache {
     size_t size_v_bytes() const;
 
     ggml_tensor * build_rope_shift(
-            const graph_params & params,
-                  ggml_context * ctx,
-                   ggml_tensor * cur,
-                   ggml_tensor * shift,
-                   ggml_tensor * factors,
-                         float   freq_base,
-                         float   freq_scale,
-           ggml_backend_buffer * bbuf) const;
+            llama_context & lctx,
+             ggml_context * ctx,
+              ggml_tensor * cur,
+              ggml_tensor * shift,
+              ggml_tensor * factors,
+                    float   freq_base,
+                    float   freq_scale,
+      ggml_backend_buffer * bbuf) const;
 
     llm_graph_result_ptr build_graph_shift(
-            const graph_params & params,
-                   ggml_cgraph * gf) const;
+            llama_context & lctx,
+              ggml_cgraph * gf) const;
 
     llm_graph_result_ptr build_graph_defrag(
-            const graph_params & params,
-                   ggml_cgraph * gf) const;
+            llama_context & lctx,
+              ggml_cgraph * gf) const;
 
     void state_write_meta(llama_io_write_i & io, const std::vector<std::pair<uint32_t, uint32_t>> & cell_ranges, llama_seq_id seq_id = -1) const;
     void state_write_data(llama_io_write_i & io, const std::vector<std::pair<uint32_t, uint32_t>> & cell_ranges) const;
@@ -331,7 +313,7 @@ class llama_kv_cache_recurrent : public llama_kv_cache {
     void restore() override;
     void commit()  override;
 
-    bool update(const graph_params & params) override;
+    bool update(llama_context & lctx) override;
 
     void defrag_sched(float thold) override;
 
