spec : save the dynamic ngram cache file (#22055)
Copilot Setup Steps / copilot-setup-steps (push) Has been cancelled
Check Pre-Tokenizer Hashes / pre-tokenizer-hashes (push) Has been cancelled
Python check requirements.txt / check-requirements (push) Has been cancelled
Python Type-Check / python type-check (push) Has been cancelled
Update Operations Documentation / update-ops-docs (push) Has been cancelled

* spec : save the dynamic ngram cache file

* save the cache at the common_speculative_impl_ngram_cache destruction
* remove save_static/save_dynamic flags as well as the relevant TODO
* if the dynamic cache file does not exist, start with an empty cache
  that is written on exit. if it exists but cannot be read, abort
  instead of overwriting it.

* spec : fix draft_one not updating the ngram cache history correctly

* common_ngram_cache_update received only the new tokens with none of
  the tokens before them causing ngrams ending at the new token to not
  being formed, leading the context cache to only ever hold the prompt

* fix: pass the full history to common_ngram_cache_update with nnew set
  to the number of new tokens the way llama-lookup does and also reset
  cache_size to 0 on each new request

---------

Co-authored-by: Petros Sideris <petros.sideris@nokia.com>
This commit is contained in:
Peter Sideris
2026-10-20 14:31:10 +03:00
committed by GitHub
parent 13786fbb3a
commit 83c2d012a2
+53 -55
View File
@@ -16,6 +16,7 @@
#include <cassert>
#include <cmath>
#include <cstring>
#include <fstream>
#include <iomanip>
#include <map>
#include <cinttypes>
@@ -2152,9 +2153,6 @@ struct common_speculative_impl_ngram_cache : public common_speculative_impl {
uint16_t n_draft;
bool save_dynamic;
bool save_static;
struct seq_info {
size_t cache_size = 0; // number of tokens in n-gram cache
@@ -2168,54 +2166,72 @@ struct common_speculative_impl_ngram_cache : public common_speculative_impl {
common_speculative_impl_ngram_cache(
const common_params_speculative & params,
uint32_t n_seq,
uint16_t n_draft,
const std::string & path_static,
const std::string & path_dynamic,
bool save_dynamic,
bool save_static)
uint16_t n_draft)
: common_speculative_impl(COMMON_SPECULATIVE_TYPE_NGRAM_CACHE, n_seq, n_draft)
, params(params.ngram_cache)
, n_draft(n_draft)
, save_dynamic(save_dynamic)
, save_static(save_static)
{
const auto & ngram = params.ngram_cache;
SPC_TRC("%s", "adding speculative implementation 'ngram-cache'\n");
SPC_TRC("- n_draft=%d, cache_static=%s, cache_dynamic=%s\n",
n_draft,
path_static.empty() ? "none" : path_static.c_str(),
path_dynamic.empty() ? "none" : path_dynamic.c_str());
ngram.lookup_cache_static.empty() ? "none" : ngram.lookup_cache_static.c_str(),
ngram.lookup_cache_dynamic.empty() ? "none" : ngram.lookup_cache_dynamic.c_str());
sinfos.resize(n_seq);
if (!path_static.empty()) {
if (!ngram.lookup_cache_static.empty()) {
try {
auto ngram_cache_static = common_ngram_cache_load(path_static);
auto ngram_cache_static = common_ngram_cache_load(ngram.lookup_cache_static);
for (auto & sinfo : sinfos) {
sinfo.ngram_cache_static = ngram_cache_static;
}
} catch (...) {
SPC_ERR("failed to open static lookup cache: %s", path_static.c_str());
} catch (const std::ifstream::failure &) {
SPC_ERR("failed to open static lookup cache: %s\n", ngram.lookup_cache_static.c_str());
GGML_ABORT("Couldn't read static lookup cache");
}
}
if (!path_dynamic.empty()) {
if (!ngram.lookup_cache_dynamic.empty()) {
try {
auto ngram_cache_dynamic = common_ngram_cache_load(path_dynamic);
auto ngram_cache_dynamic = common_ngram_cache_load(ngram.lookup_cache_dynamic);
for (auto & sinfo : sinfos) {
sinfo.ngram_cache_dynamic = ngram_cache_dynamic;
}
} catch (...) {
SPC_ERR("failed to open dynamic lookup cache: %s", path_dynamic.c_str());
GGML_ABORT("Couldn't read dynamic lookup cache");
} catch (const std::ifstream::failure &) {
std::error_code ec;
const bool exists = std::filesystem::exists(ngram.lookup_cache_dynamic, ec);
if (!exists && !ec) {
// missing file is created on destruction
SPC_WRN("dynamic lookup cache file does not exist: %s, starting with empty cache\n",
ngram.lookup_cache_dynamic.c_str());
} else {
SPC_ERR("failed to open dynamic lookup cache: %s\n", ngram.lookup_cache_dynamic.c_str());
GGML_ABORT("Couldn't read dynamic lookup cache");
}
}
}
}
void begin(llama_seq_id /*seq_id*/, const llama_tokens & /*prompt*/) override {
// noop
~common_speculative_impl_ngram_cache() override {
if (params.lookup_cache_dynamic.empty() || sinfos.empty()) {
return;
}
auto & dynamic_cache = sinfos.front().ngram_cache_dynamic;
for (auto & sinfo : sinfos) {
common_ngram_cache_merge(dynamic_cache, sinfo.ngram_cache_context);
}
SPC_INF("saving dynamic lookup cache at %s\n", params.lookup_cache_dynamic.c_str());
common_ngram_cache_save(dynamic_cache, params.lookup_cache_dynamic);
}
void begin(llama_seq_id seq_id, const llama_tokens & /*prompt*/) override {
sinfos[seq_id].cache_size = 0;
}
void draft_one(
@@ -2226,22 +2242,6 @@ struct common_speculative_impl_ngram_cache : public common_speculative_impl {
const auto & prompt = *dparams.prompt;
if (sinfo.cache_size < prompt.size() + 1) {
llama_tokens tokens_new;
tokens_new.reserve(prompt.size() + 1 - sinfo.cache_size);
for (size_t j = sinfo.cache_size; j < prompt.size(); ++j) {
tokens_new.push_back(prompt[j]);
}
tokens_new.push_back(dparams.id_last); // add the last token
// Update context ngram cache with new dparams.prompt:
common_ngram_cache_update(
sinfo.ngram_cache_context,
LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX,
tokens_new, tokens_new.size(), false);
sinfo.cache_size = prompt.size() + 1;
}
llama_tokens inp;
inp.reserve(prompt.size() + 1);
for (size_t j = 0; j < prompt.size(); ++j) {
@@ -2249,6 +2249,16 @@ struct common_speculative_impl_ngram_cache : public common_speculative_impl {
}
inp.push_back(dparams.id_last);
if (sinfo.cache_size < inp.size()) {
// Update context ngram cache with the full history so n-grams
// ending at the new tokens have their preceding context
common_ngram_cache_update(
sinfo.ngram_cache_context,
LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX,
inp, inp.size() - sinfo.cache_size, false);
sinfo.cache_size = inp.size();
}
result.push_back(dparams.id_last);
common_ngram_cache_draft(
@@ -2309,20 +2319,12 @@ static common_ngram_map get_common_ngram_map(
return common_ngram_map(size_key, size_value, key_only, min_hits);
}
static common_speculative_impl_ngram_cache create_state_ngram_cache(
static std::unique_ptr<common_speculative_impl_ngram_cache> create_state_ngram_cache(
const common_speculative_config & config,
uint32_t n_seq,
const std::string & path_static,
const std::string & path_dynamic) {
uint32_t n_seq) {
uint16_t n_draft = 8; // TODO get from config?
// TODO bool param in common/common.h to set save_static/save_dynamic?
bool save_static = false;
bool save_dynamic = false;
common_speculative_impl_ngram_cache state(config.params, n_seq, n_draft, path_static, path_dynamic, save_static, save_dynamic);
return state;
return std::make_unique<common_speculative_impl_ngram_cache>(config.params, n_seq, n_draft);
}
std::string common_speculative_type_name_str(const std::vector<common_speculative_type> & types) {
@@ -2822,11 +2824,7 @@ common_speculative * common_speculative_init(common_params_speculative & params,
break;
}
case COMMON_SPECULATIVE_TYPE_NGRAM_CACHE: {
auto state = create_state_ngram_cache(
config, n_seq,
params.ngram_cache.lookup_cache_static,
params.ngram_cache.lookup_cache_dynamic);
impls.push_back(std::make_unique<common_speculative_impl_ngram_cache>(state));
impls.push_back(create_state_ngram_cache(config, n_seq));
break;
}
default: