common: add vector of speculative states

common/speculative.cpp

@@ -1,15 +1,17 @@
 #include "speculative.h"
 
+#include <cstring>
+#include <algorithm>
+#include <map>
+#include <fstream>
+
 #include "ggml.h"
 #include "llama.h"
 #include "log.h"
 #include "common.h"
+#include "ngram-cache.h"
+#include "ngram-map.h"
 #include "sampling.h"
-#include "ngram-map.cpp"
-
-#include <cstring>
-#include <algorithm>
-#include <map>
 
 #define SPEC_VOCAB_MAX_SIZE_DIFFERENCE  128
 #define SPEC_VOCAB_CHECK_START_TOKEN_ID 5
@@ -34,14 +36,7 @@ const std::map<std::string, enum common_speculative_type> common_speculative_typ
     {"ngram_cache",   COMMON_SPECULATIVE_TYPE_NGRAM_CACHE}
 };
 
-struct common_speculative_self {
-    uint16_t         size_ngram = 12;      // size of n-grams to lookup in self-mode
-    uint16_t         size_mgram = 48;      // size of m-grams to draft in self-mode
-    const uint16_t   check_rate = 3;       // check for speculative decoding without draft model for each check_rate token
-    size_t           idx_last_check  =  0; // index of last check in context history
-};
-
-struct common_speculative_impl {
+struct common_speculative_state {
     const enum common_speculative_type type;
 
     size_t drafts_call_count       = 0; // number of times this implementation was called.
@@ -49,6 +44,95 @@ struct common_speculative_impl {
     size_t drafts_accepted_count   = 0; // number of times a draft or part was accepted by the target model.
     size_t drafts_generated_tokens = 0; // number of tokens generated by this implementation.
     size_t drafts_accepted_tokens  = 0; // number of tokens accepted by the target model.
+
+    virtual ~common_speculative_state() = default;
+
+    common_speculative_state(enum common_speculative_type type) : type(type) {}
+};
+
+struct common_speculative_state_draft : public common_speculative_state {
+    common_speculative_state_draft(enum common_speculative_type type) : common_speculative_state(type) {}
+};
+struct common_speculative_state_eagle3 : public common_speculative_state {
+    common_speculative_state_eagle3(enum common_speculative_type type) : common_speculative_state(type) {}
+};
+
+// state of self-speculation (simple implementation, not ngram-map)
+struct common_speculative_state_ngram_simple : public common_speculative_state {
+    uint16_t         size_ngram;      // size of n-grams to lookup in self-mode
+    uint16_t         size_mgram;      // size of m-grams to draft in self-mode
+    const uint16_t   check_rate;      // check for speculative decoding without draft model for each check_rate token
+
+    size_t           idx_last_check  =  0; // index of last check in context history (mutable)
+
+    common_speculative_state_ngram_simple(
+            enum common_speculative_type type,
+            uint16_t size_ngram,
+            uint16_t size_mgram,
+            uint16_t check_rate)
+        : common_speculative_state(type)
+        , size_ngram(size_ngram)
+        , size_mgram(size_mgram)
+        , check_rate(check_rate) {}
+};
+
+struct common_speculative_state_ngram_map_k : public common_speculative_state {
+    common_ngram_map        map;        // draft ngram map for speculative decoding without draft model
+
+    common_speculative_state_ngram_map_k(
+            enum common_speculative_type type,
+            common_ngram_map map)
+        : common_speculative_state(type), map(map) {}
+};
+struct common_speculative_state_ngram_map_k4v : public common_speculative_state_ngram_map_k {
+    common_speculative_state_ngram_map_k4v(
+            enum common_speculative_type type,
+            common_ngram_map map)
+        : common_speculative_state_ngram_map_k(type, std::move(map)) {}
+};
+
+struct common_speculative_state_ngram_cache : public common_speculative_state {
+    uint16_t n_draft;
+    bool save_dynamic;
+    bool save_static;
+
+    common_ngram_cache ngram_cache_context;
+    common_ngram_cache ngram_cache_dynamic;
+    common_ngram_cache ngram_cache_static;
+
+    size_t cache_size = 0; // number of tokens in n-gram cache
+
+    common_speculative_state_ngram_cache(
+            const enum common_speculative_type type,
+            std::string & path_static,
+            std::string & path_dynamic,
+            uint16_t      n_draft,
+            bool          save_dynamic,
+            bool          save_static)
+        : common_speculative_state(type)
+        , n_draft(n_draft)
+        , save_dynamic(save_dynamic)
+        , save_static(save_static)
+    {
+        if (!path_static.empty()) {
+            try {
+                ngram_cache_static = common_ngram_cache_load(path_static);
+            } catch (std::ifstream::failure const &) {
+                LOG_ERR("failed to open static lookup cache: %s", path_static.c_str());
+                GGML_ABORT("Couldn't read static lookup cache");
+            }
+        }
+
+        if (!path_dynamic.empty()) {
+            try {
+                ngram_cache_dynamic = common_ngram_cache_load(path_dynamic);
+            } catch (std::ifstream::failure const &) {
+                LOG_ERR("failed to open dynamic lookup cache: %s", path_dynamic.c_str());
+                GGML_ABORT("Couldn't read dynamic lookup cache");
+            }
+        }
+
+    }
 };
 
 struct common_speculative {
@@ -61,78 +145,77 @@ struct common_speculative {
     bool vocab_dft_compatible = true; // whether retokenization is needed
     std::map<std::string, std::string> tgt_dft_replacements = {};
 
-    common_ngram_map        map;        // draft ngram map for speculative decoding without draft model
-    common_speculative_self self_state; // state of self-speculation (simple implementation, not ngram-map)
-
-    std::vector<common_speculative_impl> impls; // list of implementations to use and their statistics
-    common_speculative_impl * curr_impl = nullptr;    // current implementation in use (for stats)
+    std::vector<std::unique_ptr<common_speculative_state>> impls; // list of implementations to use and their states
+    common_speculative_state *                             curr_impl = nullptr; // current implementation in use (for stats)
 };
 
-common_ngram_map get_common_ngram_map(std::vector<common_speculative_config> configs);
+common_ngram_map get_common_ngram_map(const common_speculative_config config);
+struct common_speculative_state_ngram_cache create_state_ngram_cache(
+        std::string path_static, std::string path_dynamic,
+        common_speculative_config config);
 
-common_ngram_map get_common_ngram_map(std::vector<common_speculative_config> configs) {
-    for (const auto & config : configs) {
-        switch (config.type) {
-            case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE:
-            case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K:
-            case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V: {
-                // create common_ngram_map from config.config
-                // compute size_key, size_value, key_only, check_rate, min_hits from config.config
-                uint16_t size_key   = 12;
-                uint16_t size_value = 48;
-                bool     key_only   = false;
-                uint16_t check_rate = 3;
-                uint16_t min_hits   = 1;
-                const std::map<std::string, std::string> & cfg = config.config;
-                // check for key "size_ngram" in cfg
-                if (cfg.find("size_ngram") != cfg.end()) {
-                    size_key = std::stoi(cfg.at("size_ngram"));
-                    if (size_key < 1 || size_key > 1024) {
-                        throw std::invalid_argument("size_ngram must be between 1 and 1024");
-                    }
-                }
-                // check for key "size_mgram" in cfg
-                if (cfg.find("size_mgram") != cfg.end()) {
-                    size_value = std::stoi(cfg.at("size_mgram"));
-                    if (size_value < 1 || size_value > 1024) {
-                        throw std::invalid_argument("size_mgram must be between 1 and 1024");
-                    }
-                }
-                // check for key "key_only" in cfg
-                if (cfg.find("key_only") != cfg.end()) {
-                    // key_onle == true, if cfg.at("key_only") has value "true".
-                    key_only = (cfg.at("key_only") == "true");
-                }
-                // check for key "check_rate" in cfg
-                if (cfg.find("check_rate") != cfg.end()) {
-                    check_rate = std::stoi(cfg.at("check_rate"));
-                    if (check_rate < 1 || check_rate > 1024) {
-                        throw std::invalid_argument("check_rate must be between 1 and 1024");
-                    }
-                }
-                // check for key "min_hits" in cfg
-                if (cfg.find("min_hits") != cfg.end()) {
-                    min_hits = std::stoi(cfg.at("min_hits"));
-                    if (min_hits < 1 || min_hits > 1024) {
-                        throw std::invalid_argument("min_hits must be between 1 and 1024");
-                    }
-                }
-                return common_ngram_map(size_key, size_value, key_only, check_rate, min_hits);
-                break;
-            }
-            case COMMON_SPECULATIVE_TYPE_NONE:
-            case COMMON_SPECULATIVE_TYPE_DRAFT:
-            case COMMON_SPECULATIVE_TYPE_EAGLE3:
-            case COMMON_SPECULATIVE_TYPE_NGRAM_CACHE:
-                break;
-            case COMMON_SPECULATIVE_TYPE_COUNT:
-                break;
+common_ngram_map get_common_ngram_map(const common_speculative_config config) {
+    uint16_t size_key   = 12;
+    uint16_t size_value = 48;
+    bool     key_only   = false;
+    uint16_t check_rate = 3;
+    uint16_t min_hits   = 1;
+    const std::map<std::string, std::string> & cfg = config.config;
+    if (cfg.find("size_ngram") != cfg.end()) {
+        size_key = std::stoi(cfg.at("size_ngram"));
+        if (size_key < 1 || size_key > 1024) {
+            throw std::invalid_argument("size_ngram must be between 1 and 1024");
         }
     }
-
-    return common_ngram_map(12, 48, false, 3, 1); // default fallback
+    if (cfg.find("size_mgram") != cfg.end()) {
+        size_value = std::stoi(cfg.at("size_mgram"));
+        if (size_value < 1 || size_value > 1024) {
+            throw std::invalid_argument("size_mgram must be between 1 and 1024");
+        }
+    }
+    if (cfg.find("key_only") != cfg.end()) {
+        key_only = (cfg.at("key_only") == "true");
+    }
+    if (cfg.find("check_rate") != cfg.end()) {
+        check_rate = std::stoi(cfg.at("check_rate"));
+        if (check_rate < 1 || check_rate > 1024) {
+            throw std::invalid_argument("check_rate must be between 1 and 1024");
+        }
+    }
+    if (cfg.find("min_hits") != cfg.end()) {
+        min_hits = std::stoi(cfg.at("min_hits"));
+        if (min_hits < 1 || min_hits > 1024) {
+            throw std::invalid_argument("min_hits must be between 1 and 1024");
+        }
+    }
+    return common_ngram_map(size_key, size_value, key_only, check_rate, min_hits);
 }
 
+struct common_speculative_state_ngram_cache create_state_ngram_cache(
+        std::string path_static, std::string path_dynamic,
+        common_speculative_config config) {
+    uint16_t n_draft = 8;
+    bool save_static = false;
+    bool save_dynamic = false;
+    const std::map<std::string, std::string> & cfg = config.config;
+    if (cfg.find("n_draft") != cfg.end()) {
+        n_draft = std::stoi(cfg.at("n_draft"));
+        if (n_draft < 1 || n_draft > 1024) {
+            throw std::invalid_argument("ngram-cache: n_draft must be between 1 and 1024");
+        }
+    }
+    if (cfg.find("save_static") != cfg.end()) {
+        save_static = (cfg.at("save_static") == "true");
+    }
+    if (cfg.find("save_dynamic") != cfg.end()) {
+        save_dynamic = (cfg.at("save_dynamic") == "true");
+    }
+    common_speculative_state_ngram_cache state(config.type,
+            path_static, path_dynamic, n_draft, save_static, save_dynamic);
+    return state;
+}
+
+
 std::string common_speculative_type_name_str() {
     std::string result = "";
     for (size_t i = 0; i < common_speculative_types.size(); i++) {
@@ -168,25 +251,59 @@ enum common_speculative_type common_speculative_type_from_name(const std::string
 
 
 struct common_speculative * common_speculative_init(
+        struct common_params & params,
         struct llama_context * ctx_tgt,
-        struct llama_context * ctx_dft,
-        const std::vector<common_speculative_config> configs
+        struct llama_context * ctx_dft
     ) {
-    common_ngram_map ngram_map = get_common_ngram_map(configs);
-    uint16_t ngram_size_key   = ngram_map.size_key;
-    uint16_t mgram_size_value = ngram_map.size_value;
-    uint16_t check_rate = ngram_map.check_rate;
-    common_speculative_self self_state = common_speculative_self{
-        /* .size_ngram      = */ ngram_size_key,
-        /* .size_mgram      = */ mgram_size_value,
-        /* .check_rate      = */ check_rate,
-        /* .idx_last_check  = */ 0,
-    };
-    std::vector<common_speculative_impl> implementations = {};
-    LOG_INF("common_speculative_init: configs.size = %zu\n", configs.size());
-    for (const auto & config : configs) {
+    std::vector<std::unique_ptr<common_speculative_state>> implementations = {};
+    for (const common_speculative_config & config : params.speculative.configs) {
         LOG_INF("common_speculative_init: adding implementation %s\n", common_speculative_type_to_str(config.type).c_str());
-        implementations.push_back(common_speculative_impl{config.type});
+        switch (config.type) {
+            case COMMON_SPECULATIVE_TYPE_NONE:
+                break;
+            case COMMON_SPECULATIVE_TYPE_DRAFT: {
+                implementations.push_back(std::make_unique<common_speculative_state_draft>(config.type));
+                break;
+            }
+            case COMMON_SPECULATIVE_TYPE_EAGLE3: {
+                implementations.push_back(std::make_unique<common_speculative_state_eagle3>(config.type));
+                break;
+            }
+            case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE: {
+                common_ngram_map ngram_map = get_common_ngram_map(config);
+                uint16_t ngram_size_key   = ngram_map.size_key;
+                uint16_t mgram_size_value = ngram_map.size_value;
+                uint16_t check_rate = ngram_map.check_rate;
+                auto state = std::make_unique<common_speculative_state_ngram_simple>(
+                    /* .type            = */ config.type,
+                    /* .size_ngram      = */ ngram_size_key,
+                    /* .size_mgram      = */ mgram_size_value,
+                    /* .check_rate      = */ check_rate
+                );
+                implementations.push_back(std::move(state));
+                break;
+            }
+            case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K: {
+                implementations.push_back(std::make_unique<common_speculative_state_ngram_map_k>(
+                    (config.type), get_common_ngram_map(config)
+                ));
+                break;
+            }
+            case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V: {
+                implementations.push_back(std::make_unique<common_speculative_state_ngram_map_k4v>(
+                            (config.type), get_common_ngram_map(config)));
+                break;
+            }
+            case COMMON_SPECULATIVE_TYPE_NGRAM_CACHE: {
+                auto state = create_state_ngram_cache(
+                        params.lookup_cache_static, params.lookup_cache_dynamic, config);
+                implementations.push_back(std::make_unique<common_speculative_state_ngram_cache>(state));
+
+                break;
+            }
+            default:
+                break;
+        }
     }
     auto * result = new common_speculative {
         /* .ctx_tgt    = */ ctx_tgt,
@@ -196,12 +313,9 @@ struct common_speculative * common_speculative_init(
         /* .prompt_dft = */ {},
         /* .vocab_dft_compatible = */ false,
         /* .tgt_dft_replacements = */ {},
-        /* .map             = */ ngram_map,
-        /* .self_state      = */ self_state,
-        /* .impls          = */ implementations
+        /* .impls          = */ std::move(implementations)
     };
 
-    LOG_INF("common_speculative_init: created speculative decoder, map.n = %d\n", result->map.size_key);
     // TODO: optimize or pass from outside?
 #if 0
     {
@@ -359,7 +473,11 @@ llama_tokens common_speculative_use_draft_model(
         llama_token id_last);
 
 llama_tokens common_speculative_gen_self_draft(
-        common_speculative * spec,
+        common_speculative_state_ngram_simple & state,
+        const llama_tokens & tokens, llama_token sampled);
+
+llama_tokens common_speculative_gen_ngram_cache(
+        common_speculative_state_ngram_cache & state,
         const llama_tokens & tokens, llama_token sampled);
 
 llama_tokens common_speculative_gen_draft(
@@ -371,9 +489,9 @@ llama_tokens common_speculative_gen_draft(
 
     spec->curr_impl = nullptr; // reset current implementation
     for (auto & impl : spec->impls) {
-        impl.drafts_call_count++;
+        impl->drafts_call_count++;
         // LOG name and call_count
-        switch (impl.type) {
+        switch (impl->type) {
             case COMMON_SPECULATIVE_TYPE_NONE:
             {
                 break;
@@ -392,24 +510,36 @@ llama_tokens common_speculative_gen_draft(
             case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE:
             {
                 // Use common_ngram_map_draft to generate a draft from the current context.
-                result = common_speculative_gen_self_draft(spec, prompt_tgt_main_model, id_last);
+                auto * state = dynamic_cast<struct common_speculative_state_ngram_simple *>(impl.get());
+                if (state) {
+                    result = common_speculative_gen_self_draft(*state, prompt_tgt_main_model, id_last);
+                }
                 break;
             }
             case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K:
             {
                 // Use common_ngram_map_draft to generate a draft from the current context.
-                common_ngram_map_draft(spec->map, prompt_tgt_main_model, id_last, result);
+                auto state = dynamic_cast<common_speculative_state_ngram_map_k *>(impl.get());
+                if (state) {
+                    common_ngram_map_draft(state->map, prompt_tgt_main_model, id_last, result);
+                }
                 break;
             }
             case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V:
             {
                 // Use common_ngram_map_draft to generate a draft from the current context.
-                common_ngram_map_draft(spec->map, prompt_tgt_main_model, id_last, result);
+                auto state = dynamic_cast<common_speculative_state_ngram_map_k *>(impl.get());
+                if (state) {
+                    common_ngram_map_draft(state->map, prompt_tgt_main_model, id_last, result);
+                }
                 break;
             }
             case COMMON_SPECULATIVE_TYPE_NGRAM_CACHE:
             {
-                // TODO call common/ngram-cache.cpp
+                auto * state= dynamic_cast<common_speculative_state_ngram_cache *>(impl.get());
+                if (state) {
+                    result = common_speculative_gen_ngram_cache(*state, prompt_tgt_main_model, id_last);
+                }
                 break;
             }
             case COMMON_SPECULATIVE_TYPE_COUNT:
@@ -420,12 +550,12 @@ llama_tokens common_speculative_gen_draft(
         }
         if (!result.empty()) {
             LOG_DBG("%s: called impl %s, hist size = %zu, call_count = %zu, gen = %zu\n", __func__,
-                    common_speculative_type_to_str(impl.type).c_str(),
+                    common_speculative_type_to_str(impl.get()->type).c_str(),
                     prompt_tgt_main_model.size(),
-                    impl.drafts_call_count, result.size());
-            spec->curr_impl = &impl; // set current implementation for stats
-            impl.drafts_generated_count++;
-            impl.drafts_generated_tokens += result.size();
+                    impl.get()->drafts_call_count, result.size());
+            spec->curr_impl = impl.get(); // set current implementation for stats
+            impl->drafts_generated_count++;
+            impl->drafts_generated_tokens += result.size();
 
             break; // We have a draft, so break out of the loop and return it.
         }
@@ -614,7 +744,7 @@ llama_tokens common_speculative_use_draft_model(
 }
 
 void common_speculative_send_accepted(struct common_speculative * spec, const uint16_t n_accepted) {
-    common_speculative_impl * impl = spec->curr_impl;
+    common_speculative_state * impl = spec->curr_impl;
     if (impl != nullptr) {
         if (n_accepted > 0) {
             impl->drafts_accepted_count++;
@@ -622,11 +752,28 @@ void common_speculative_send_accepted(struct common_speculative * spec, const ui
         }
         if (impl->type == COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K ||
             impl->type == COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V) {
-            common_ngram_map_send_accepted(spec->map, n_accepted);
+            auto state = static_cast<struct common_speculative_state_ngram_map_k *>(impl);
+            common_ngram_map_send_accepted(state->map, n_accepted);
         }
     }
 }
 
+void common_speculative_print_stats(const struct common_speculative * spec) {
+    if (spec == nullptr) {
+        return;
+    }
+    for (const auto & impl : spec->impls) {
+        LOG_INF("statistics %s: #calls = %zu, #gen drafts = %zu, #acc drafts = %zu, #gen tokens = %zu, #acc tokens = %zu\n",
+                common_speculative_type_to_str(impl->type).c_str(),
+                impl->drafts_call_count,
+                impl->drafts_generated_count,
+                impl->drafts_accepted_count,
+                impl->drafts_generated_tokens,
+                impl->drafts_accepted_tokens);
+    }
+}
+
+
 // self-speculative decoding
 //
 
@@ -634,28 +781,27 @@ void common_speculative_send_accepted(struct common_speculative * spec, const ui
  * Perform speculative generation using the model's own token history.
  * Searches for a matching pattern in the token history and returns draft tokens.
  *
- * @param spec      configuration of speculative drafts
+ * @param state     Current state of this implementation
  * @param tokens    Token history to search in
  * @param sampled   Last sampled token
  * @return Vector of draft tokens, empty if no matching pattern is found
  */
 llama_tokens common_speculative_gen_self_draft(
-        common_speculative * spec,
+        common_speculative_state_ngram_simple & state,
         const llama_tokens & tokens, llama_token sampled) {
 
     // Simple implementation of self-speculative decoding without draft model, without ngram-map.
     //
-    common_speculative_self & self_state = spec->self_state;
     const size_t cur_len = tokens.size();
     // Only check every check_rate tokens to save compute
     // i.e., perform check if (cur_len - idx_last_check) >= check_rate
-    if (self_state.idx_last_check + self_state.check_rate > cur_len) {
+    if (state.idx_last_check + state.check_rate > cur_len) {
         llama_tokens draft_tokens;
         return draft_tokens;
     }
 
-    size_t n_draft_min = self_state.size_ngram; // size of n-gram to lookup in token history
-    size_t n_draft_max = self_state.size_mgram; // the m-gram following the found n-gram is used for draft
+    size_t n_draft_min = state.size_ngram; // size of n-gram to lookup in token history
+    size_t n_draft_max = state.size_mgram; // the m-gram following the found n-gram is used for draft
 
     // vector for tokens we want to verify.
     // return empty vector if there is no match.
@@ -675,10 +821,10 @@ llama_tokens common_speculative_gen_self_draft(
     pattern.push_back(sampled); // add the last token to the pattern
 
     // We do a search in the token history.
-    self_state.idx_last_check = tokens.size();
+    state.idx_last_check = tokens.size();
 
     size_t match_pos = 0; // we ignore position 0, position 0 == no match
-    // search backwards, but skip the current match (we are currently there)
+                          // search backwards, but skip the current match (we are currently there)
     for (size_t j = cur_len - n_draft_min - 1; j > 0; --j) {
         bool match = true;
         for (size_t k = 0; k < pattern.size(); ++k) {
@@ -697,15 +843,15 @@ llama_tokens common_speculative_gen_self_draft(
     }
 
     const size_t copy_max = std::min(
-        n_draft_max,
-        cur_len - (match_pos + n_draft_min)
-    );
+            n_draft_max,
+            cur_len - (match_pos + n_draft_min)
+            );
     if (copy_max < n_draft_min) {
         return draft_tokens;
     }
     LOG_DBG("%s: #tokens = %zu: found matching pattern at pos %zu, length %zu, draft length %zu\n",
-        __func__, cur_len,
-        match_pos, pattern.size(), copy_max);
+            __func__, cur_len,
+            match_pos, pattern.size(), copy_max);
 
     draft_tokens.reserve(copy_max);
     for (size_t j = 0; j < copy_max; ++j) {
@@ -714,17 +860,53 @@ llama_tokens common_speculative_gen_self_draft(
     return draft_tokens;
 }
 
-void common_speculative_print_stats(const struct common_speculative * spec) {
-    if (spec == nullptr) {
-        return;
+// n-gram cache
+//
+
+/**
+ * Perform speculative generation using a 3-tier n-gram cache.
+ *
+ * @param state     Current state of this implementation
+ * @param tokens    Token history to search in
+ * @param sampled   Last sampled token
+ * @return Vector of draft tokens, empty if draft is found
+ */
+llama_tokens common_speculative_gen_ngram_cache(
+        common_speculative_state_ngram_cache & state,
+        const llama_tokens & tokens, llama_token sampled) {
+    if (state.cache_size < tokens.size() + 1) {
+        llama_tokens tokens_new;
+        tokens_new.reserve(tokens.size() + 1 - state.cache_size);
+        for (size_t j = state.cache_size; j < tokens.size(); ++j) {
+            tokens_new.push_back(tokens[j]);
+        }
+        tokens_new.push_back(sampled); // add the last token
+
+        // Update context ngram cache with new tokens:
+        common_ngram_cache_update(state.ngram_cache_context, LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX,
+                tokens_new, tokens_new.size(), false);
+        state.cache_size = tokens.size() + 1;
     }
-    for (const auto & impl : spec->impls) {
-        LOG_INF("statistics %s: #calls = %zu, #gen drafts = %zu, #acc drafts = %zu, #gen tokens = %zu, #acc tokens = %zu\n",
-                common_speculative_type_to_str(impl.type).c_str(),
-                impl.drafts_call_count,
-                impl.drafts_generated_count,
-                impl.drafts_accepted_count,
-                impl.drafts_generated_tokens,
-                impl.drafts_accepted_tokens);
-     }
+
+    llama_tokens inp;
+    inp.reserve(tokens.size() + 1);
+    for (size_t j = 0; j < tokens.size(); ++j) {
+        inp.push_back(tokens[j]);
+    }
+    inp.push_back(sampled);
+
+    llama_tokens draft;
+    draft.push_back(sampled);
+
+    common_ngram_cache_draft(inp, draft, state.n_draft, LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX,
+            state.ngram_cache_context,
+            state.ngram_cache_dynamic,
+            state.ngram_cache_static);
+
+    if (draft.size() > 0) {
+        // delete first token in draft (which is the sampled token)
+        draft.erase(draft.begin());
+    }
+
+    return draft;
 }