From eb43748b0597aaf5ff5c493b788541a64efb9c49 Mon Sep 17 00:00:00 2001 From: Sascha Rogmann Date: Wed, 21 Jan 2026 22:46:28 +0100 Subject: [PATCH] common: add vector of speculative states --- common/CMakeLists.txt | 2 + common/arg.cpp | 34 +- common/ngram-map.cpp | 2 +- common/speculative.cpp | 456 ++++++++++++------ common/speculative.h | 4 +- .../speculative-simple/speculative-simple.cpp | 2 +- tools/server/server-context.cpp | 7 +- 7 files changed, 348 insertions(+), 159 deletions(-) diff --git a/common/CMakeLists.txt b/common/CMakeLists.txt index ae02c0bd77..3bc7bc6210 100644 --- a/common/CMakeLists.txt +++ b/common/CMakeLists.txt @@ -73,6 +73,8 @@ add_library(${TARGET} STATIC log.h ngram-cache.cpp ngram-cache.h + ngram-map.cpp + ngram-map.h peg-parser.cpp peg-parser.h preset.cpp diff --git a/common/arg.cpp b/common/arg.cpp index a258d1a494..811b95fa39 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -626,25 +626,33 @@ static bool common_params_parse_ex(int argc, char ** argv, common_params_context if (!params.speculative.tensor_buft_overrides.empty()) { params.speculative.tensor_buft_overrides.push_back({nullptr, nullptr}); } - if (!params.speculative.model.path.empty()) { - bool found_draft = false; - bool found_eagle3 = false; + { + bool has_draft =!params.speculative.model.path.empty(); + bool has_draft_eagle3 = false; // TODO PR-18039: if params.speculative.eagle3 + bool has_lookup_caches = !params.lookup_cache_static.empty() && !params.lookup_cache_dynamic.empty(); + bool found_config_draft = false; + bool found_config_eagle3 = false; + bool found_config_ngram_cache = false; for (const auto & config : params.speculative.configs) { if (config.type == COMMON_SPECULATIVE_TYPE_DRAFT) { - found_draft = true; + found_config_draft = true; } if (config.type == COMMON_SPECULATIVE_TYPE_EAGLE3) { - found_eagle3 = true; - break; + found_config_eagle3 = true; + } + if (config.type == COMMON_SPECULATIVE_TYPE_NGRAM_CACHE) { + found_config_ngram_cache = true; } } - if (!found_draft) { + if (has_draft && !found_config_draft) { params.speculative.configs.push_back(common_speculative_config(COMMON_SPECULATIVE_TYPE_DRAFT)); } - // TODO PR-18039: if params.speculative.eagle3 - //if (!found_eagle3) { - // params.speculative.configs.push_back(common_speculative_config(COMMON_SPECULATIVE_TYPE_DRAFT)); - //} + if (has_draft_eagle3 && !found_config_eagle3) { + params.speculative.configs.push_back(common_speculative_config(COMMON_SPECULATIVE_TYPE_EAGLE3)); + } + if (has_lookup_caches && !found_config_ngram_cache) { + params.speculative.configs.push_back(common_speculative_config(COMMON_SPECULATIVE_TYPE_NGRAM_CACHE)); + } } if (!params.chat_template.empty() && !common_chat_verify_template(params.chat_template, params.use_jinja)) { @@ -1239,14 +1247,14 @@ common_params_context common_params_parser_init(common_params & params, llama_ex [](common_params & params, const std::string & value) { params.lookup_cache_static = value; } - ).set_examples({LLAMA_EXAMPLE_LOOKUP})); + ).set_examples({LLAMA_EXAMPLE_LOOKUP, LLAMA_EXAMPLE_SERVER})); add_opt(common_arg( {"-lcd", "--lookup-cache-dynamic"}, "FNAME", "path to dynamic lookup cache to use for lookup decoding (updated by generation)", [](common_params & params, const std::string & value) { params.lookup_cache_dynamic = value; } - ).set_examples({LLAMA_EXAMPLE_LOOKUP})); + ).set_examples({LLAMA_EXAMPLE_LOOKUP, LLAMA_EXAMPLE_SERVER})); add_opt(common_arg( {"-c", "--ctx-size"}, "N", string_format("size of the prompt context (default: %d, 0 = loaded from model)", params.n_ctx), diff --git a/common/ngram-map.cpp b/common/ngram-map.cpp index f79066cae0..9c589ab6f7 100644 --- a/common/ngram-map.cpp +++ b/common/ngram-map.cpp @@ -1,6 +1,6 @@ -#include "ngram-map.h" #include "common.h" #include "log.h" +#include "ngram-map.h" #include #include diff --git a/common/speculative.cpp b/common/speculative.cpp index f52121df59..5791c328a0 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -1,15 +1,17 @@ #include "speculative.h" +#include +#include +#include +#include + #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 -#include -#include #define SPEC_VOCAB_MAX_SIZE_DIFFERENCE 128 #define SPEC_VOCAB_CHECK_START_TOKEN_ID 5 @@ -34,14 +36,7 @@ const std::map 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 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 impls; // list of implementations to use and their statistics - common_speculative_impl * curr_impl = nullptr; // current implementation in use (for stats) + std::vector> 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 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 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 & 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 & 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 & 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 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 implementations = {}; - LOG_INF("common_speculative_init: configs.size = %zu\n", configs.size()); - for (const auto & config : configs) { + std::vector> 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(config.type)); + break; + } + case COMMON_SPECULATIVE_TYPE_EAGLE3: { + implementations.push_back(std::make_unique(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( + /* .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( + (config.type), get_common_ngram_map(config) + )); + break; + } + case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V: { + implementations.push_back(std::make_unique( + (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(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(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(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(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(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(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; } diff --git a/common/speculative.h b/common/speculative.h index 7fdfd6116a..3cf3fdbf3b 100644 --- a/common/speculative.h +++ b/common/speculative.h @@ -22,9 +22,9 @@ enum common_speculative_type common_speculative_type_from_name(const std::string std::string common_speculative_type_to_str(enum common_speculative_type type); struct common_speculative * common_speculative_init( + struct common_params & params, struct llama_context * ctx_tgt, - struct llama_context * ctx_dft, - const std::vector configs = {} // incubator config (options not yet in common_params) + struct llama_context * ctx_dft ); void common_speculative_free(struct common_speculative * spec); diff --git a/examples/speculative-simple/speculative-simple.cpp b/examples/speculative-simple/speculative-simple.cpp index 8141052a22..0439209131 100644 --- a/examples/speculative-simple/speculative-simple.cpp +++ b/examples/speculative-simple/speculative-simple.cpp @@ -132,7 +132,7 @@ int main(int argc, char ** argv) { params_spec.n_reuse = llama_n_ctx(ctx_dft) - n_draft; params_spec.p_min = p_min; - struct common_speculative * spec = common_speculative_init(ctx_tgt, ctx_dft); + struct common_speculative * spec = common_speculative_init(params, ctx_tgt, ctx_dft); for (auto &pair : params.speculative.replacements) { common_speculative_add_replacement_tgt_dft(spec, pair.first.c_str(), pair.second.c_str()); } diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 46bb8194d1..20ec61940c 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -775,8 +775,7 @@ private: return false; } - slot.spec = common_speculative_init(slot.ctx, slot.ctx_dft, - params_base.speculative.configs); + slot.spec = common_speculative_init(params_base, slot.ctx, slot.ctx_dft); if (slot.spec == nullptr) { SRV_ERR("%s", "failed to create speculator\n"); return false; @@ -785,9 +784,7 @@ private: common_speculative_add_replacement_tgt_dft(slot.spec, pair.first.c_str(), pair.second.c_str()); } } else if (params_base.speculative.configs.size() > 0) { - SLT_INF(slot, "init spec for speculative decoding without draft model, slot %d\n", i); - slot.spec = common_speculative_init(nullptr, nullptr, - params_base.speculative.configs); + slot.spec = common_speculative_init(params_base, nullptr, nullptr); } SLT_INF(slot, "new slot, n_ctx = %d\n", slot.n_ctx);