diff --git a/common/arg.cpp b/common/arg.cpp index 6e5b59a..c442896 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -1234,6 +1234,15 @@ bool common_params_parse(int argc, char ** argv, common_params & params, llama_e exit(0); } params.lr.init(); + llama_token_cache_params token_cache_params = llama_token_cache_default_params(); + token_cache_params.path = params.token_cache_path.empty() ? nullptr : params.token_cache_path.c_str(); + token_cache_params.capacity_bytes = params.token_cache_size_mib * 1024ull * 1024ull; + if (!llama_token_cache_configure(token_cache_params) && params.token_cache_size_mib != 0) { + LOG_WRN("%s", "token cache was already initialized; keeping its existing configuration\n"); + } + if (params.token_cache_clear && !llama_token_cache_clear()) { + throw std::runtime_error("failed to clear token cache"); + } } catch (const std::invalid_argument & ex) { fprintf(stderr, "%s\n", ex.what()); ctx_arg.params = params_org; @@ -1636,6 +1645,34 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.cache_ram_mib = value; } ).set_env("LLAMA_ARG_CACHE_RAM").set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + add_opt(common_arg( + {"--token-cache-size"}, "N", + string_format("shared tokenizer and tokenized-text cache size in MiB (default: %llu, 0 = disable)", + (unsigned long long) params.token_cache_size_mib), + [](common_params & params, int value) { + if (value < 0) { + throw std::invalid_argument("token-cache-size must be non-negative"); + } + params.token_cache_size_mib = value; + } + ).set_env("LLAMA_ARG_TOKEN_CACHE_SIZE")); + add_opt(common_arg( + {"--token-cache-dir"}, "PATH", + "path to the shared tokenizer and tokenized-text cache database", + [](common_params & params, const std::string & value) { + if (value.empty()) { + throw std::invalid_argument("token-cache-dir must not be empty"); + } + params.token_cache_path = value; + } + ).set_env("LLAMA_ARG_TOKEN_CACHE_DIR")); + add_opt(common_arg( + {"--token-cache-clear"}, + "clear the shared tokenizer and tokenized-text cache before loading the model", + [](common_params & params) { + params.token_cache_clear = true; + } + ).set_env("LLAMA_ARG_TOKEN_CACHE_CLEAR")); add_opt(common_arg( {"-kvu", "--kv-unified"}, {"-no-kvu", "--no-kv-unified"}, diff --git a/common/common.h b/common/common.h index a2801de..c56699a 100644 --- a/common/common.h +++ b/common/common.h @@ -626,6 +626,9 @@ struct common_params { int32_t n_ctx_checkpoints = 32; // max number of context checkpoints per slot int32_t checkpoint_min_step = 8192; // minimum spacing between context checkpoints int32_t cache_ram_mib = 8192; // -1 = no limit, 0 - disable, 1 = 1 MiB, etc. + uint64_t token_cache_size_mib = 5120; // shared tokenizer/text cache, 0 = disable + std::string token_cache_path; + bool token_cache_clear = false; std::string hostname = "127.0.0.1"; std::string public_path = ""; // NOLINT diff --git a/docs/token-cache.md b/docs/token-cache.md new file mode 100644 index 0000000..af2d0c3 --- /dev/null +++ b/docs/token-cache.md @@ -0,0 +1,32 @@ +# Tokenizer and tokenized-text cache + +libllama maintains a persistent cache containing constructed tokenizer snapshots +and text-to-token results. The cache is enabled by default with a shared 5 GiB +logical entry budget and is available to direct libllama callers as well as the +llama.cpp tools. + +The default database is `token-cache.sqlite3` below the platform llama.cpp cache +directory. `LLAMA_CACHE` changes the base directory. llama.cpp tools also accept: + +```text +--token-cache-size N cache size in MiB (default: 5120, 0 disables) +--token-cache-dir PATH use a different SQLite database +--token-cache-clear clear the database before loading a model +``` + +The matching environment variables are `LLAMA_ARG_TOKEN_CACHE_SIZE`, +`LLAMA_ARG_TOKEN_CACHE_DIR`, and `LLAMA_ARG_TOKEN_CACHE_CLEAR`. + +Cache entries are shared by processes running as the same OS user. Tokenized +text keys are SHA-256 digests and the original text is not stored. Token IDs can +usually be detokenized, however, so the database must still be treated as +sensitive user data. The cache directory and files are created with owner-only +permissions where supported. + +Writes use a bounded asynchronous queue. Cache failures, unavailable storage, +or corrupt entries are treated as misses and do not prevent model loading or +tokenization. Applications that require all pending entries to be committed can +call `llama_token_cache_flush()`. + +The 5 GiB limit accounts for logical entry payloads. SQLite metadata and bounded +journal files are not included. diff --git a/include/llama.h b/include/llama.h index d749287..f961ab3 100644 --- a/include/llama.h +++ b/include/llama.h @@ -304,6 +304,26 @@ extern "C" { ggml_backend_buffer_type_t buft; }; + struct llama_token_cache_params { + // NULL selects the platform llama.cpp cache directory. + const char * path; + // Shared logical entry budget. 0 disables the cache for this process. + uint64_t capacity_bytes; + }; + + struct llama_token_cache_stats { + uint64_t tokenizer_hits; + uint64_t tokenizer_misses; + uint64_t text_hits; + uint64_t text_misses; + uint64_t writes; + uint64_t dropped_writes; + uint64_t evictions; + uint64_t errors; + uint64_t entries; + uint64_t bytes_used; + }; + struct llama_model_params { // NULL-terminated list of devices to use for offloading (if NULL, all available devices are used) ggml_backend_dev_t * devices; @@ -460,6 +480,12 @@ extern "C" { // Helpers for getting default parameters // TODO: update API to start accepting pointers to params structs (https://github.com/ggml-org/llama.cpp/discussions/9172) LLAMA_API struct llama_model_params llama_model_default_params(void); + LLAMA_API struct llama_token_cache_params llama_token_cache_default_params(void); + // Must be called before the first model load or tokenization operation. + LLAMA_API bool llama_token_cache_configure(struct llama_token_cache_params params); + LLAMA_API bool llama_token_cache_flush(void); + LLAMA_API bool llama_token_cache_clear(void); + LLAMA_API struct llama_token_cache_stats llama_token_cache_get_stats(void); LLAMA_API struct llama_context_params llama_context_default_params(void); LLAMA_API struct llama_sampler_chain_params llama_sampler_chain_default_params(void); LLAMA_API struct llama_model_quantize_params llama_model_quantize_default_params(void); diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 0669246..6fb9f12 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -37,6 +37,7 @@ add_library(llama llama-model.cpp llama-quant.cpp llama-sampler.cpp + llama-token-cache.cpp llama-vocab.cpp unicode-data.cpp unicode.cpp @@ -56,6 +57,13 @@ target_compile_features (llama PRIVATE cxx_std_17) # don't bump target_link_libraries(llama PUBLIC ggml) +find_package(SQLite3 REQUIRED) +if (TARGET SQLite3::SQLite3) + target_link_libraries(llama PRIVATE SQLite3::SQLite3) +else() + target_link_libraries(llama PRIVATE SQLite::SQLite3) +endif() + if (BUILD_SHARED_LIBS) set_target_properties(llama PROPERTIES POSITION_INDEPENDENT_CODE ON) target_compile_definitions(llama PRIVATE LLAMA_BUILD) diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp index 2e3f2d7..ffbd649 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -1,4 +1,5 @@ #include "llama-model-loader.h" +#include "llama-token-cache.h" #include "ggml-alloc.h" #include "ggml.h" @@ -12,6 +13,7 @@ #include #include #include +#include #include static const size_t kiB = 1024; @@ -540,6 +542,33 @@ llama_model_loader::llama_model_loader( } } + if (!fname.empty()) { + std::error_code path_ec; + std::error_code size_ec; + std::error_code time_ec; + const auto absolute = std::filesystem::absolute(fname, path_ec); + const auto size = std::filesystem::file_size(fname, size_ec); + const auto mtime = std::filesystem::last_write_time(fname, time_ec).time_since_epoch().count(); + std::string source = (path_ec ? fname : absolute.string()) + '\n' + + std::to_string(size_ec ? 0 : static_cast(size)) + '\n' + + std::to_string(time_ec ? 0 : static_cast(mtime)); + std::vector override_keys; + override_keys.reserve(kv_overrides.size()); + for (const auto & item : kv_overrides) override_keys.push_back(item.first); + std::sort(override_keys.begin(), override_keys.end()); + for (const auto & key : override_keys) { + const auto & value = kv_overrides.at(key); + source += '\n' + key + ':' + std::to_string(value.tag) + ':'; + switch (value.tag) { + case LLAMA_KV_OVERRIDE_TYPE_INT: source += std::to_string(value.val_i64); break; + case LLAMA_KV_OVERRIDE_TYPE_FLOAT: source += std::to_string(value.val_f64); break; + case LLAMA_KV_OVERRIDE_TYPE_BOOL: source += value.val_bool ? "1" : "0"; break; + case LLAMA_KV_OVERRIDE_TYPE_STR: source += value.val_str; break; + } + } + token_cache_source_key = "source:" + llama_token_cache_hash(source); + } + tensor_buft_overrides = param_tensor_buft_overrides_p; this->use_mmap = load_mode == LLAMA_LOAD_MODE_MMAP || load_mode == LLAMA_LOAD_MODE_MMAP_MLOCK; @@ -862,6 +891,10 @@ enum llm_arch llama_model_loader::get_arch() const { return llm_kv.arch; } +const std::string & llama_model_loader::get_token_cache_source_key() const { + return token_cache_source_key; +} + const llama_model_loader::llama_tensor_weight * llama_model_loader::get_weight(const char * name) const { auto pos = weights_map.find(name); if (pos != weights_map.end()) { diff --git a/src/llama-model-loader.h b/src/llama-model-loader.h index 557208a..6224c31 100644 --- a/src/llama-model-loader.h +++ b/src/llama-model-loader.h @@ -107,6 +107,7 @@ struct llama_model_loader { std::vector contexts; std::string arch_name; + std::string token_cache_source_key; LLM_KV llm_kv = LLM_KV(LLM_ARCH_UNKNOWN); size_t size_done = 0; @@ -179,6 +180,8 @@ struct llama_model_loader { enum llm_arch get_arch() const; + const std::string & get_token_cache_source_key() const; + void configure_longhaul(uint64_t cache_bytes, uint32_t n_expert, uint32_t n_layer); const llama_tensor_weight * get_weight(const char * name) const; diff --git a/src/llama-token-cache.cpp b/src/llama-token-cache.cpp new file mode 100644 index 0000000..17fea0c --- /dev/null +++ b/src/llama-token-cache.cpp @@ -0,0 +1,514 @@ +#include "llama-token-cache.h" + +#include "llama-impl.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if defined(_WIN32) +# include +#else +# include +# include +#endif + +namespace fs = std::filesystem; + +namespace { + +constexpr uint64_t DEFAULT_CAPACITY = 5ull * 1024 * 1024 * 1024; +constexpr uint64_t MAX_QUEUE_BYTES = 64ull * 1024 * 1024; +constexpr int SCHEMA_VERSION = 1; + +struct sha256 { + uint32_t h[8] = { + 0x6a09e667, 0xbb67ae85, 0x3c6ef372, 0xa54ff53a, + 0x510e527f, 0x9b05688c, 0x1f83d9ab, 0x5be0cd19, + }; + uint8_t block[64] = {}; + uint64_t bits = 0; + size_t used = 0; + + static uint32_t rotr(uint32_t x, uint32_t n) { return (x >> n) | (x << (32 - n)); } + + void compress(const uint8_t * p) { + static const uint32_t k[64] = { + 0x428a2f98,0x71374491,0xb5c0fbcf,0xe9b5dba5,0x3956c25b,0x59f111f1,0x923f82a4,0xab1c5ed5, + 0xd807aa98,0x12835b01,0x243185be,0x550c7dc3,0x72be5d74,0x80deb1fe,0x9bdc06a7,0xc19bf174, + 0xe49b69c1,0xefbe4786,0x0fc19dc6,0x240ca1cc,0x2de92c6f,0x4a7484aa,0x5cb0a9dc,0x76f988da, + 0x983e5152,0xa831c66d,0xb00327c8,0xbf597fc7,0xc6e00bf3,0xd5a79147,0x06ca6351,0x14292967, + 0x27b70a85,0x2e1b2138,0x4d2c6dfc,0x53380d13,0x650a7354,0x766a0abb,0x81c2c92e,0x92722c85, + 0xa2bfe8a1,0xa81a664b,0xc24b8b70,0xc76c51a3,0xd192e819,0xd6990624,0xf40e3585,0x106aa070, + 0x19a4c116,0x1e376c08,0x2748774c,0x34b0bcb5,0x391c0cb3,0x4ed8aa4a,0x5b9cca4f,0x682e6ff3, + 0x748f82ee,0x78a5636f,0x84c87814,0x8cc70208,0x90befffa,0xa4506ceb,0xbef9a3f7,0xc67178f2, + }; + uint32_t w[64]; + for (int i = 0; i < 16; ++i) { + w[i] = uint32_t(p[4*i]) << 24 | uint32_t(p[4*i+1]) << 16 | uint32_t(p[4*i+2]) << 8 | p[4*i+3]; + } + for (int i = 16; i < 64; ++i) { + const uint32_t s0 = rotr(w[i-15], 7) ^ rotr(w[i-15], 18) ^ (w[i-15] >> 3); + const uint32_t s1 = rotr(w[i-2], 17) ^ rotr(w[i-2], 19) ^ (w[i-2] >> 10); + w[i] = w[i-16] + s0 + w[i-7] + s1; + } + uint32_t a=h[0],b=h[1],c=h[2],d=h[3],e=h[4],f=h[5],g=h[6],hh=h[7]; + for (int i = 0; i < 64; ++i) { + const uint32_t s1 = rotr(e,6)^rotr(e,11)^rotr(e,25); + const uint32_t ch = (e&f)^((~e)&g); + const uint32_t t1 = hh+s1+ch+k[i]+w[i]; + const uint32_t s0 = rotr(a,2)^rotr(a,13)^rotr(a,22); + const uint32_t maj = (a&b)^(a&c)^(b&c); + const uint32_t t2 = s0+maj; + hh=g; g=f; f=e; e=d+t1; d=c; c=b; b=a; a=t1+t2; + } + h[0]+=a;h[1]+=b;h[2]+=c;h[3]+=d;h[4]+=e;h[5]+=f;h[6]+=g;h[7]+=hh; + } + + void update(const void * data, size_t n) { + const auto * p = static_cast(data); + bits += uint64_t(n) * 8; + while (n) { + const size_t take = std::min(n, sizeof(block) - used); + memcpy(block + used, p, take); + used += take; p += take; n -= take; + if (used == sizeof(block)) { compress(block); used = 0; } + } + } + + std::array finish() { + const uint64_t original_bits = bits; + const uint8_t one = 0x80; + update(&one, 1); + const uint8_t zero = 0; + while (used != 56) update(&zero, 1); + uint8_t len[8]; + for (int i = 0; i < 8; ++i) len[7-i] = uint8_t(original_bits >> (8*i)); + update(len, 8); + std::array out; + for (int i = 0; i < 8; ++i) { + out[4*i]=uint8_t(h[i]>>24); out[4*i+1]=uint8_t(h[i]>>16); + out[4*i+2]=uint8_t(h[i]>>8); out[4*i+3]=uint8_t(h[i]); + } + return out; + } +}; + +struct queued_write { + llama_token_cache_kind kind; + std::string key; + std::vector payload; +}; + +struct cache_counters { + std::atomic tokenizer_hits{0}; + std::atomic tokenizer_misses{0}; + std::atomic text_hits{0}; + std::atomic text_misses{0}; + std::atomic writes{0}; + std::atomic dropped_writes{0}; + std::atomic evictions{0}; + std::atomic errors{0}; +}; + +class token_cache { +public: + token_cache() { + params.capacity_bytes = DEFAULT_CAPACITY; + worker = std::thread([this] { run(); }); + } + + ~token_cache() { + flush(); + { + std::lock_guard lock(mutex); + stopping = true; + } + ready.notify_all(); + if (worker.joinable()) worker.join(); + close(); + } + + bool configure(llama_token_cache_params value) { + std::lock_guard lock(mutex); + if (initialized) return false; + params.capacity_bytes = value.capacity_bytes; + configured_path = value.path ? value.path : ""; + return true; + } + + bool lookup(llama_token_cache_kind kind, const std::string & key, std::vector & out) { + if (!ensure_open()) return false; + std::lock_guard db_lock(db_mutex); + sqlite3_stmt * stmt = nullptr; + const char * sql = "SELECT payload FROM entries WHERE kind=?1 AND key=?2"; + if (sqlite3_prepare_v2(db, sql, -1, &stmt, nullptr) != SQLITE_OK) return db_error(); + sqlite3_bind_int(stmt, 1, kind); + sqlite3_bind_text(stmt, 2, key.data(), int(key.size()), SQLITE_TRANSIENT); + bool found = false; + if (sqlite3_step(stmt) == SQLITE_ROW) { + const int n = sqlite3_column_bytes(stmt, 0); + const void * p = sqlite3_column_blob(stmt, 0); + if (n >= 0 && (n == 0 || p != nullptr)) { + out.assign(static_cast(p), static_cast(p) + n); + found = true; + } + } + sqlite3_finalize(stmt); + if (found) { + sqlite3_stmt * touch = nullptr; + if (sqlite3_prepare_v2(db, "UPDATE entries SET last_access=?1 WHERE kind=?2 AND key=?3", -1, &touch, nullptr) == SQLITE_OK) { + sqlite3_bind_int64(touch, 1, now_tick()); + sqlite3_bind_int(touch, 2, kind); + sqlite3_bind_text(touch, 3, key.data(), int(key.size()), SQLITE_TRANSIENT); + sqlite3_step(touch); + } + sqlite3_finalize(touch); + if (kind == LLAMA_TOKEN_CACHE_KIND_TEXT) ++stats.text_hits; + } else if (kind == LLAMA_TOKEN_CACHE_KIND_TEXT) { + ++stats.text_misses; + } + return found; + } + + void store(llama_token_cache_kind kind, const std::string & key, const std::vector & payload, bool sync) { + if (params.capacity_bytes == 0 || payload.size() > params.capacity_bytes || + payload.size() > static_cast(std::numeric_limits::max())) return; + if (sync) { + write_one({kind, key, payload}); + return; + } + std::lock_guard lock(mutex); + if (queue_bytes + payload.size() > MAX_QUEUE_BYTES) { + ++stats.dropped_writes; + return; + } + queue.push_back({kind, key, payload}); + queue_bytes += payload.size(); + ready.notify_one(); + } + + void flush() { + std::unique_lock lock(mutex); + drained.wait(lock, [&] { return queue.empty() && !writing; }); + if (db) { + lock.unlock(); + std::lock_guard db_lock(db_mutex); + sqlite3_wal_checkpoint_v2(db, nullptr, SQLITE_CHECKPOINT_PASSIVE, nullptr, nullptr); + } + } + + bool clear_all() { + flush(); + if (!ensure_open()) return params.capacity_bytes == 0; + std::lock_guard db_lock(db_mutex); + if (sqlite3_exec(db, "DELETE FROM entries; PRAGMA incremental_vacuum;", nullptr, nullptr, nullptr) != SQLITE_OK) return db_error(); + return true; + } + + llama_token_cache_stats get_stats() { + llama_token_cache_stats result = { + stats.tokenizer_hits.load(), + stats.tokenizer_misses.load(), + stats.text_hits.load(), + stats.text_misses.load(), + stats.writes.load(), + stats.dropped_writes.load(), + stats.evictions.load(), + stats.errors.load(), + 0, + 0, + }; + if (ensure_open()) { + std::lock_guard db_lock(db_mutex); + sqlite3_stmt * stmt = nullptr; + if (sqlite3_prepare_v2(db, "SELECT COUNT(*), COALESCE(SUM(size),0) FROM entries", -1, &stmt, nullptr) == SQLITE_OK && + sqlite3_step(stmt) == SQLITE_ROW) { + result.entries = sqlite3_column_int64(stmt, 0); + result.bytes_used = sqlite3_column_int64(stmt, 1); + } + sqlite3_finalize(stmt); + } + return result; + } + + void note_tokenizer(bool hit) { + if (hit) ++stats.tokenizer_hits; else ++stats.tokenizer_misses; + } + +private: + static int64_t now_tick() { + return std::chrono::duration_cast( + std::chrono::system_clock::now().time_since_epoch()).count(); + } + + std::string default_path() { + if (const char * p = std::getenv("LLAMA_CACHE"); p && *p) return (fs::path(p) / "token-cache.sqlite3").string(); +#if defined(_WIN32) + if (const char * p = std::getenv("LOCALAPPDATA"); p && *p) return (fs::path(p) / "llama.cpp" / "token-cache.sqlite3").string(); +#elif defined(__APPLE__) + if (const char * p = std::getenv("HOME"); p && *p) return (fs::path(p) / "Library" / "Caches" / "llama.cpp" / "token-cache.sqlite3").string(); +#else + if (const char * p = std::getenv("XDG_CACHE_HOME"); p && *p) return (fs::path(p) / "llama.cpp" / "token-cache.sqlite3").string(); + if (const char * p = std::getenv("HOME"); p && *p) return (fs::path(p) / ".cache" / "llama.cpp" / "token-cache.sqlite3").string(); +#endif + return {}; + } + + bool ensure_open() { + std::lock_guard init_lock(init_mutex); + if (initialized) return db != nullptr; + initialized = true; + if (params.capacity_bytes == 0) return false; + path = configured_path.empty() ? default_path() : configured_path; + if (path.empty()) return false; + std::error_code ec; + fs::create_directories(fs::path(path).parent_path(), ec); + if (ec) return false; +#if !defined(_WIN32) + chmod(fs::path(path).parent_path().c_str(), 0700); +#endif + if (sqlite3_open_v2(path.c_str(), &db, SQLITE_OPEN_READWRITE | SQLITE_OPEN_CREATE | SQLITE_OPEN_FULLMUTEX, nullptr) != SQLITE_OK) { + warn_once("failed to open token cache database"); + close(); + return false; + } +#if !defined(_WIN32) + chmod(path.c_str(), 0600); +#endif + sqlite3_busy_timeout(db, 5000); + const char * schema = + "PRAGMA journal_mode=WAL;" + "PRAGMA synchronous=NORMAL;" + "PRAGMA auto_vacuum=INCREMENTAL;" + "CREATE TABLE IF NOT EXISTS meta(key TEXT PRIMARY KEY,value INTEGER NOT NULL);" + "CREATE TABLE IF NOT EXISTS entries(" + " kind INTEGER NOT NULL,key TEXT NOT NULL,payload BLOB NOT NULL,size INTEGER NOT NULL,last_access INTEGER NOT NULL," + " PRIMARY KEY(kind,key));" + "CREATE INDEX IF NOT EXISTS entries_lru ON entries(last_access);"; + if (sqlite3_exec(db, schema, nullptr, nullptr, nullptr) != SQLITE_OK) { + warn_once("failed to initialize token cache database"); + close(); + return false; + } + sqlite3_stmt * stmt = nullptr; + sqlite3_prepare_v2(db, "INSERT OR REPLACE INTO meta(key,value) VALUES('schema',?1)", -1, &stmt, nullptr); + sqlite3_bind_int(stmt, 1, SCHEMA_VERSION); + sqlite3_step(stmt); + sqlite3_finalize(stmt); + sqlite3_prepare_v2(db, "INSERT OR REPLACE INTO meta(key,value) VALUES('capacity',?1)", -1, &stmt, nullptr); + sqlite3_bind_int64(stmt, 1, static_cast(params.capacity_bytes)); + sqlite3_step(stmt); + sqlite3_finalize(stmt); + sqlite3_exec(db, "BEGIN IMMEDIATE", nullptr, nullptr, nullptr); + while (true) { + sqlite3_prepare_v2(db, "SELECT COALESCE(SUM(size),0) FROM entries", -1, &stmt, nullptr); + uint64_t total = 0; + if (sqlite3_step(stmt) == SQLITE_ROW) total = sqlite3_column_int64(stmt, 0); + sqlite3_finalize(stmt); + if (total <= params.capacity_bytes) break; + if (sqlite3_exec(db, + "DELETE FROM entries WHERE rowid=(SELECT rowid FROM entries ORDER BY last_access LIMIT 1)", + nullptr, nullptr, nullptr) != SQLITE_OK) { + break; + } + ++stats.evictions; + } + sqlite3_exec(db, "COMMIT", nullptr, nullptr, nullptr); + return true; + } + + void run() { + while (true) { + queued_write item; + { + std::unique_lock lock(mutex); + ready.wait(lock, [&] { return stopping || !queue.empty(); }); + if (stopping && queue.empty()) return; + item = std::move(queue.front()); + queue.pop_front(); + queue_bytes -= item.payload.size(); + writing = true; + } + write_one(item); + { + std::lock_guard lock(mutex); + writing = false; + if (queue.empty()) drained.notify_all(); + } + } + } + + void write_one(const queued_write & item) { + if (!ensure_open()) return; + std::lock_guard db_lock(db_mutex); + if (sqlite3_exec(db, "BEGIN IMMEDIATE", nullptr, nullptr, nullptr) != SQLITE_OK) { db_error(); return; } + uint64_t capacity = params.capacity_bytes; + sqlite3_stmt * capacity_stmt = nullptr; + if (sqlite3_prepare_v2(db, "SELECT value FROM meta WHERE key='capacity'", -1, &capacity_stmt, nullptr) == SQLITE_OK && + sqlite3_step(capacity_stmt) == SQLITE_ROW) { + capacity = static_cast(sqlite3_column_int64(capacity_stmt, 0)); + } + sqlite3_finalize(capacity_stmt); + if (item.payload.size() > capacity) { + sqlite3_exec(db, "ROLLBACK", nullptr, nullptr, nullptr); + return; + } + sqlite3_stmt * existing = nullptr; + uint64_t old_size = 0; + sqlite3_prepare_v2(db, "SELECT size FROM entries WHERE kind=?1 AND key=?2", -1, &existing, nullptr); + sqlite3_bind_int(existing, 1, item.kind); + sqlite3_bind_text(existing, 2, item.key.data(), int(item.key.size()), SQLITE_TRANSIENT); + if (sqlite3_step(existing) == SQLITE_ROW) old_size = sqlite3_column_int64(existing, 0); + sqlite3_finalize(existing); + + sqlite3_stmt * total_stmt = nullptr; + uint64_t total = 0; + sqlite3_prepare_v2(db, "SELECT COALESCE(SUM(size),0) FROM entries", -1, &total_stmt, nullptr); + if (sqlite3_step(total_stmt) == SQLITE_ROW) total = sqlite3_column_int64(total_stmt, 0); + sqlite3_finalize(total_stmt); + total -= std::min(total, old_size); + + while (total + item.payload.size() > capacity) { + sqlite3_stmt * victim = nullptr; + sqlite3_prepare_v2(db, "SELECT kind,key,size FROM entries WHERE NOT(kind=?1 AND key=?2) ORDER BY last_access LIMIT 1", -1, &victim, nullptr); + sqlite3_bind_int(victim, 1, item.kind); + sqlite3_bind_text(victim, 2, item.key.data(), int(item.key.size()), SQLITE_TRANSIENT); + if (sqlite3_step(victim) != SQLITE_ROW) { sqlite3_finalize(victim); break; } + const int victim_kind = sqlite3_column_int(victim, 0); + const std::string victim_key(reinterpret_cast(sqlite3_column_text(victim, 1))); + const uint64_t victim_size = sqlite3_column_int64(victim, 2); + sqlite3_finalize(victim); + sqlite3_stmt * del = nullptr; + sqlite3_prepare_v2(db, "DELETE FROM entries WHERE kind=?1 AND key=?2", -1, &del, nullptr); + sqlite3_bind_int(del, 1, victim_kind); + sqlite3_bind_text(del, 2, victim_key.data(), int(victim_key.size()), SQLITE_TRANSIENT); + sqlite3_step(del); + sqlite3_finalize(del); + total -= std::min(total, victim_size); + ++stats.evictions; + } + + sqlite3_stmt * put = nullptr; + sqlite3_prepare_v2(db, "INSERT OR REPLACE INTO entries(kind,key,payload,size,last_access) VALUES(?1,?2,?3,?4,?5)", -1, &put, nullptr); + sqlite3_bind_int(put, 1, item.kind); + sqlite3_bind_text(put, 2, item.key.data(), int(item.key.size()), SQLITE_TRANSIENT); + sqlite3_bind_blob(put, 3, item.payload.data(), int(item.payload.size()), SQLITE_TRANSIENT); + sqlite3_bind_int64(put, 4, item.payload.size()); + sqlite3_bind_int64(put, 5, now_tick()); + const int rc = sqlite3_step(put); + sqlite3_finalize(put); + if (rc == SQLITE_DONE) { + sqlite3_exec(db, "COMMIT", nullptr, nullptr, nullptr); + ++stats.writes; + } else { + sqlite3_exec(db, "ROLLBACK", nullptr, nullptr, nullptr); + db_error(); + } + } + + bool db_error() { + ++stats.errors; + warn_once("token cache database operation failed; treating it as a cache miss"); + return false; + } + + void warn_once(const char * message) { + bool expected = false; + if (warned.compare_exchange_strong(expected, true)) { + LLAMA_LOG_WARN("%s\n", message); + } + } + + void close() { + std::lock_guard db_lock(db_mutex); + if (db) sqlite3_close(db); + db = nullptr; + } + + llama_token_cache_params params = {nullptr, DEFAULT_CAPACITY}; + std::string configured_path; + std::string path; + sqlite3 * db = nullptr; + std::mutex init_mutex; + std::mutex db_mutex; + std::mutex mutex; + std::condition_variable ready; + std::condition_variable drained; + std::deque queue; + uint64_t queue_bytes = 0; + bool initialized = false; + bool stopping = false; + bool writing = false; + std::atomic warned{false}; + std::thread worker; + cache_counters stats; +}; + +token_cache & cache() { + static token_cache instance; + return instance; +} + +} // namespace + +std::string llama_token_cache_hash(const void * data, size_t size) { + sha256 hash; + hash.update(data, size); + const auto bytes = hash.finish(); + static const char hex[] = "0123456789abcdef"; + std::string result(64, '0'); + for (size_t i = 0; i < bytes.size(); ++i) { + result[2*i] = hex[bytes[i] >> 4]; + result[2*i+1] = hex[bytes[i] & 15]; + } + return result; +} + +std::string llama_token_cache_hash(const std::string & data) { + return llama_token_cache_hash(data.data(), data.size()); +} + +bool llama_token_cache_lookup(llama_token_cache_kind kind, const std::string & key, std::vector & payload) { + return cache().lookup(kind, key, payload); +} + +void llama_token_cache_store(llama_token_cache_kind kind, const std::string & key, const std::vector & payload, bool synchronous) { + cache().store(kind, key, payload, synchronous); +} + +void llama_token_cache_note_tokenizer_hit(bool hit) { + cache().note_tokenizer(hit); +} + +llama_token_cache_params llama_token_cache_default_params() { + return {nullptr, DEFAULT_CAPACITY}; +} + +bool llama_token_cache_configure(llama_token_cache_params params) { + return cache().configure(params); +} + +bool llama_token_cache_flush() { + cache().flush(); + return true; +} + +bool llama_token_cache_clear() { + return cache().clear_all(); +} + +llama_token_cache_stats llama_token_cache_get_stats() { + return cache().get_stats(); +} diff --git a/src/llama-token-cache.h b/src/llama-token-cache.h new file mode 100644 index 0000000..b922033 --- /dev/null +++ b/src/llama-token-cache.h @@ -0,0 +1,29 @@ +#pragma once + +#include "llama.h" + +#include +#include +#include + +enum llama_token_cache_kind : int32_t { + LLAMA_TOKEN_CACHE_KIND_TOKENIZER = 1, + LLAMA_TOKEN_CACHE_KIND_TEXT = 2, +}; + +std::string llama_token_cache_hash(const void * data, size_t size); +std::string llama_token_cache_hash(const std::string & data); + +bool llama_token_cache_lookup( + llama_token_cache_kind kind, + const std::string & key, + std::vector & payload); + +void llama_token_cache_store( + llama_token_cache_kind kind, + const std::string & key, + const std::vector & payload, + bool synchronous = false); + +void llama_token_cache_note_tokenizer_hit(bool hit); + diff --git a/src/llama-vocab.cpp b/src/llama-vocab.cpp index 9164a4d..72dd0e1 100644 --- a/src/llama-vocab.cpp +++ b/src/llama-vocab.cpp @@ -1,4 +1,5 @@ #include "llama-vocab.h" +#include "llama-token-cache.h" #include "ggml.h" #include "gguf.h" @@ -1841,6 +1842,7 @@ struct llama_vocab::impl { std::vector suppress_tokens; std::unique_ptr tokenizer; + std::string token_cache_id; std::vector precompiled_charsmap; @@ -1851,6 +1853,9 @@ struct llama_vocab::impl { void load(llama_model_loader & ml, const LLM_KV & kv); + bool load_cache_snapshot(const std::vector & snapshot); + void reset_cache_snapshot_state(); + enum llama_vocab_type get_type() const; std::string type_name() const; @@ -1881,6 +1886,11 @@ struct llama_vocab::impl { bool add_special, bool parse_special = false) const; + std::vector tokenize_uncached( + const std::string & raw_text, + bool add_special, + bool parse_special = false) const; + int32_t tokenize( const char * text, int32_t text_len, @@ -1921,6 +1931,25 @@ private: void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { struct gguf_context * ctx = ml.metadata; + const std::string & source_cache_key = ml.get_token_cache_source_key(); + if (!source_cache_key.empty()) { + std::vector snapshot; + bool found = llama_token_cache_lookup(LLAMA_TOKEN_CACHE_KIND_TOKENIZER, source_cache_key, snapshot); + if (found && snapshot.size() == 65 && snapshot[0] == 'A') { + const std::string content_key(reinterpret_cast(snapshot.data() + 1), 64); + snapshot.clear(); + found = llama_token_cache_lookup(LLAMA_TOKEN_CACHE_KIND_TOKENIZER, content_key, snapshot); + } + if (found && + load_cache_snapshot(snapshot)) { + token_cache_id = llama_token_cache_hash(snapshot.data(), snapshot.size()); + llama_token_cache_note_tokenizer_hit(true); + return; + } + reset_cache_snapshot_state(); + llama_token_cache_note_tokenizer_hit(false); + } + // determine vocab type { ml.get_key(LLM_KV_TOKENIZER_MODEL, tokenizer_model); @@ -3033,6 +3062,206 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { } } } + + // Persist a versioned canonical snapshot of the effective tokenizer state. + // The snapshot also serves as the namespace for text-to-token cache entries. + { + std::vector snapshot; + auto put = [&snapshot](const auto & value) { + const auto * p = reinterpret_cast(&value); + snapshot.insert(snapshot.end(), p, p + sizeof(value)); + }; + auto put_string = [&snapshot, &put](const std::string & value) { + const uint64_t n = value.size(); + put(n); + snapshot.insert(snapshot.end(), value.begin(), value.end()); + }; + const uint32_t version = 2; + put(version); + put(type); + put(pre_type); + put(n_token_types); + put(max_token_len); + put_string(tokenizer_model); + put_string(tokenizer_pre); + put(special_bos_id); put(special_eos_id); put(special_eot_id); put(special_eom_id); + put(special_unk_id); put(special_sep_id); put(special_pad_id); put(special_mask_id); + put(special_fim_pre_id); put(special_fim_suf_id); put(special_fim_mid_id); + put(special_fim_pad_id); put(special_fim_rep_id); put(special_fim_sep_id); + put(linefeed_id); + put(add_space_prefix); put(add_bos); put(add_eos); put(add_sep); put(ignore_merges); + put(clean_spaces); put(remove_extra_whitespaces); put(escape_whitespaces); put(treat_whitespace_as_suffix); + put(normalizer_opts.lowercase); + put(normalizer_opts.strip_accents); + const uint64_t n_tokens = id_to_token.size(); + put(n_tokens); + for (const auto & item : id_to_token) { + put_string(item.text); + put(item.score); + put(item.attr); + } + std::vector> merges; + merges.reserve(bpe_ranks.size()); + for (const auto & item : bpe_ranks) { + merges.emplace_back(item.first.first + '\0' + item.first.second, item.second); + } + std::sort(merges.begin(), merges.end()); + const uint64_t n_merges = merges.size(); + put(n_merges); + for (const auto & item : merges) { + put_string(item.first); + put(item.second); + } + const uint64_t charsmap_size = precompiled_charsmap.size(); + put(charsmap_size); + snapshot.insert(snapshot.end(), precompiled_charsmap.begin(), precompiled_charsmap.end()); + const uint64_t n_suppress = suppress_tokens.size(); + put(n_suppress); + for (const auto token : suppress_tokens) put(token); + const uint64_t n_special = cache_special_tokens.size(); + put(n_special); + for (const auto token : cache_special_tokens) put(token); + const uint64_t n_pieces = cache_token_to_piece.size(); + put(n_pieces); + for (const auto & piece : cache_token_to_piece) put_string(piece); + const uint64_t n_eog = special_eog_ids.size(); + put(n_eog); + for (const auto token : special_eog_ids) put(token); + std::vector> text_ids(token_to_id.begin(), token_to_id.end()); + std::sort(text_ids.begin(), text_ids.end()); + const uint64_t n_text_ids = text_ids.size(); + put(n_text_ids); + for (const auto & item : text_ids) { + put_string(item.first); + put(item.second); + } + + token_cache_id = llama_token_cache_hash(snapshot.data(), snapshot.size()); + std::vector cached; + const bool hit = llama_token_cache_lookup(LLAMA_TOKEN_CACHE_KIND_TOKENIZER, token_cache_id, cached); + if (!hit) { + llama_token_cache_store(LLAMA_TOKEN_CACHE_KIND_TOKENIZER, token_cache_id, snapshot, true); + } + if (!source_cache_key.empty()) { + std::vector alias(1 + token_cache_id.size()); + alias[0] = 'A'; + memcpy(alias.data() + 1, token_cache_id.data(), token_cache_id.size()); + llama_token_cache_store(LLAMA_TOKEN_CACHE_KIND_TOKENIZER, source_cache_key, alias, true); + } + } +} + +void llama_vocab::impl::reset_cache_snapshot_state() { + n_token_types = 0; + tokenizer_model.clear(); + tokenizer_pre.clear(); + type = LLAMA_VOCAB_TYPE_SPM; + pre_type = LLAMA_VOCAB_PRE_TYPE_DEFAULT; + max_token_len = 0; + special_bos_id = 1; special_eos_id = 2; special_eot_id = LLAMA_TOKEN_NULL; special_eom_id = LLAMA_TOKEN_NULL; + special_unk_id = 0; special_sep_id = LLAMA_TOKEN_NULL; special_pad_id = LLAMA_TOKEN_NULL; special_mask_id = LLAMA_TOKEN_NULL; + linefeed_id = 13; + special_fim_pre_id = special_fim_suf_id = special_fim_mid_id = LLAMA_TOKEN_NULL; + special_fim_pad_id = special_fim_rep_id = special_fim_sep_id = LLAMA_TOKEN_NULL; + add_space_prefix = false; add_bos = false; add_eos = false; add_sep = false; + ignore_merges = false; clean_spaces = false; remove_extra_whitespaces = false; + escape_whitespaces = true; treat_whitespace_as_suffix = false; + normalizer_opts = {}; + token_to_id.clear(); + id_to_token.clear(); + cache_special_tokens.clear(); + cache_token_to_piece.clear(); + bpe_ranks.clear(); + special_eog_ids.clear(); + suppress_tokens.clear(); + tokenizer.reset(); + precompiled_charsmap.clear(); + token_cache_id.clear(); +} + +bool llama_vocab::impl::load_cache_snapshot(const std::vector & snapshot) { + size_t offset = 0; + auto get = [&snapshot, &offset](auto & value) { + if (sizeof(value) > snapshot.size() - std::min(offset, snapshot.size())) return false; + memcpy(&value, snapshot.data() + offset, sizeof(value)); + offset += sizeof(value); + return true; + }; + auto get_string = [&snapshot, &offset, &get](std::string & value) { + uint64_t n = 0; + if (!get(n) || n > snapshot.size() - std::min(offset, snapshot.size())) return false; + value.assign(reinterpret_cast(snapshot.data() + offset), static_cast(n)); + offset += static_cast(n); + return true; + }; + auto get_count = [&get](uint64_t & n) { + return get(n) && n <= 2000000; + }; + + uint32_t version = 0; + if (!get(version) || version != 2 || + !get(type) || !get(pre_type) || !get(n_token_types) || !get(max_token_len) || + !get_string(tokenizer_model) || !get_string(tokenizer_pre) || + !get(special_bos_id) || !get(special_eos_id) || !get(special_eot_id) || !get(special_eom_id) || + !get(special_unk_id) || !get(special_sep_id) || !get(special_pad_id) || !get(special_mask_id) || + !get(special_fim_pre_id) || !get(special_fim_suf_id) || !get(special_fim_mid_id) || + !get(special_fim_pad_id) || !get(special_fim_rep_id) || !get(special_fim_sep_id) || + !get(linefeed_id) || + !get(add_space_prefix) || !get(add_bos) || !get(add_eos) || !get(add_sep) || !get(ignore_merges) || + !get(clean_spaces) || !get(remove_extra_whitespaces) || !get(escape_whitespaces) || !get(treat_whitespace_as_suffix) || + !get(normalizer_opts.lowercase) || !get(normalizer_opts.strip_accents)) { + return false; + } + if (type <= LLAMA_VOCAB_TYPE_NONE || type > LLAMA_VOCAB_TYPE_PLAMO2) return false; + + uint64_t count = 0; + if (!get_count(count)) return false; + id_to_token.resize(count); + for (auto & item : id_to_token) { + if (!get_string(item.text) || !get(item.score) || !get(item.attr)) return false; + } + + if (!get_count(count)) return false; + for (uint64_t i = 0; i < count; ++i) { + std::string joined; + int rank = 0; + if (!get_string(joined) || !get(rank)) return false; + const auto split = joined.find('\0'); + if (split == std::string::npos) return false; + bpe_ranks.emplace(std::make_pair(joined.substr(0, split), joined.substr(split + 1)), rank); + } + + uint64_t charsmap_size = 0; + if (!get_count(charsmap_size) || charsmap_size > snapshot.size() - std::min(offset, snapshot.size())) return false; + precompiled_charsmap.assign(snapshot.begin() + offset, snapshot.begin() + offset + charsmap_size); + offset += charsmap_size; + + if (!get_count(count)) return false; + suppress_tokens.resize(count); + for (auto & token : suppress_tokens) if (!get(token)) return false; + if (!get_count(count)) return false; + cache_special_tokens.resize(count); + for (auto & token : cache_special_tokens) if (!get(token)) return false; + if (!get_count(count)) return false; + cache_token_to_piece.resize(count); + for (auto & piece : cache_token_to_piece) if (!get_string(piece)) return false; + if (!get_count(count)) return false; + for (uint64_t i = 0; i < count; ++i) { + llama_token token; + if (!get(token)) return false; + special_eog_ids.insert(token); + } + if (!get_count(count)) return false; + for (uint64_t i = 0; i < count; ++i) { + std::string text; + llama_token token; + if (!get_string(text) || !get(token) || token < 0 || static_cast(token) >= id_to_token.size()) return false; + token_to_id.emplace(std::move(text), token); + } + if (offset != snapshot.size() || cache_token_to_piece.size() != id_to_token.size()) return false; + + init_tokenizer(type); + return true; } enum llama_vocab_type llama_vocab::impl::get_type() const { @@ -3310,7 +3539,7 @@ static std::string llama_decode_text(const std::string & text) { return decoded_text; } -std::vector llama_vocab::impl::tokenize( +std::vector llama_vocab::impl::tokenize_uncached( const std::string & raw_text, bool add_special, bool parse_special) const { @@ -3516,6 +3745,52 @@ std::vector llama_vocab::impl::tokenize( return output; } +std::vector llama_vocab::impl::tokenize( + const std::string & raw_text, + bool add_special, + bool parse_special) const { + if (raw_text.empty() || token_cache_id.empty()) { + return tokenize_uncached(raw_text, add_special, parse_special); + } + + std::string key_input; + key_input.reserve(token_cache_id.size() + raw_text.size() + 3); + key_input += token_cache_id; + key_input.push_back('\0'); + key_input.push_back(add_special ? '\1' : '\0'); + key_input.push_back(parse_special ? '\1' : '\0'); + key_input += raw_text; + const std::string key = llama_token_cache_hash(key_input); + + std::vector payload; + if (llama_token_cache_lookup(LLAMA_TOKEN_CACHE_KIND_TEXT, key, payload) && + payload.size() >= sizeof(uint32_t) && + (payload.size() - sizeof(uint32_t)) % sizeof(llama_token) == 0) { + uint32_t count = 0; + memcpy(&count, payload.data(), sizeof(count)); + const size_t available = (payload.size() - sizeof(count)) / sizeof(llama_token); + if (count == available && count < static_cast(std::numeric_limits::max())) { + std::vector result(count); + if (count) memcpy(result.data(), payload.data() + sizeof(count), count * sizeof(llama_token)); + if (std::all_of(result.begin(), result.end(), [&](llama_token token) { + return token >= 0 && static_cast(token) < id_to_token.size(); + })) { + return result; + } + } + } + + auto result = tokenize_uncached(raw_text, add_special, parse_special); + if (result.size() < static_cast(std::numeric_limits::max())) { + const uint32_t count = result.size(); + payload.resize(sizeof(count) + result.size() * sizeof(llama_token)); + memcpy(payload.data(), &count, sizeof(count)); + if (!result.empty()) memcpy(payload.data() + sizeof(count), result.data(), result.size() * sizeof(llama_token)); + llama_token_cache_store(LLAMA_TOKEN_CACHE_KIND_TEXT, key, payload); + } + return result; +} + int32_t llama_vocab::impl::token_to_piece(llama_token token, char * buf, int32_t length, int32_t lstrip, bool special) const { // ref: https://github.com/ggml-org/llama.cpp/pull/7587#discussion_r1620983843 static const int attr_special = LLAMA_TOKEN_ATTR_UNKNOWN | LLAMA_TOKEN_ATTR_CONTROL; diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 66fd403..5520b83 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -259,6 +259,8 @@ set_tests_properties(test-thread-safety PROPERTIES FIXTURES_REQUIRED test-downlo llama_build_and_test(test-arg-parser.cpp) llama_build_and_test(test-longhaul.cpp) +llama_build_and_test(test-token-cache.cpp) +target_include_directories(test-token-cache PRIVATE ${PROJECT_SOURCE_DIR}/src) if (NOT LLAMA_SANITIZE_ADDRESS AND NOT GGML_SCHED_NO_REALLOC) # TODO: repair known memory leaks diff --git a/tests/test-arg-parser.cpp b/tests/test-arg-parser.cpp index c9abf39..b278341 100644 --- a/tests/test-arg-parser.cpp +++ b/tests/test-arg-parser.cpp @@ -160,6 +160,14 @@ static void test(void) { assert(params.load_mode == LLAMA_LOAD_MODE_LONGHAUL); assert(params.longhaul_cache_bytes == 2ULL * 1024 * 1024 * 1024); + argv = {"binary_name", "--token-cache-size", "256", "--token-cache-dir", "/tmp/llama-token-cache-test.sqlite3"}; + assert(true == common_params_parse(argv.size(), list_str_to_char(argv).data(), params, LLAMA_EXAMPLE_COMMON)); + assert(params.token_cache_size_mib == 256); + assert(params.token_cache_path == "/tmp/llama-token-cache-test.sqlite3"); + + argv = {"binary_name", "--token-cache-size", "-1"}; + assert(false == common_params_parse(argv.size(), list_str_to_char(argv).data(), params, LLAMA_EXAMPLE_COMMON)); + // multi-value args (CSV) argv = {"binary_name", "--lora", "file1.gguf,\"file2,2.gguf\",\"file3\"\"3\"\".gguf\",file4\".gguf"}; assert(true == common_params_parse(argv.size(), list_str_to_char(argv).data(), params, LLAMA_EXAMPLE_COMMON)); diff --git a/tests/test-token-cache.cpp b/tests/test-token-cache.cpp new file mode 100644 index 0000000..ace783d --- /dev/null +++ b/tests/test-token-cache.cpp @@ -0,0 +1,69 @@ +#include "llama.h" +#include "../src/llama-token-cache.h" + +#include +#include +#include +#include + +#undef NDEBUG +#include + +int main() { + namespace fs = std::filesystem; + + const auto defaults = llama_token_cache_default_params(); + assert(defaults.capacity_bytes == 5ull * 1024 * 1024 * 1024); + + const fs::path root = fs::temp_directory_path() / + ("llama-token-cache-test-" + std::to_string( + std::chrono::high_resolution_clock::now().time_since_epoch().count())); + const fs::path db = root / "cache.sqlite3"; + + llama_token_cache_params params = { + nullptr, + 128, + }; + // configure copies the path; it must not retain the temporary c_str(). + const std::string db_string = db.string(); + params.path = db_string.c_str(); + assert(llama_token_cache_configure(params)); + assert(llama_token_cache_clear()); + + assert(llama_token_cache_hash("abc") == + "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"); + + const std::vector first(80, 0x11); + const std::vector second(80, 0x22); + std::vector loaded; + + llama_token_cache_store(LLAMA_TOKEN_CACHE_KIND_TOKENIZER, "first", first, true); + assert(llama_token_cache_lookup(LLAMA_TOKEN_CACHE_KIND_TOKENIZER, "first", loaded)); + assert(loaded == first); + + llama_token_cache_store(LLAMA_TOKEN_CACHE_KIND_TEXT, "second", second, true); + loaded.clear(); + assert(!llama_token_cache_lookup(LLAMA_TOKEN_CACHE_KIND_TOKENIZER, "first", loaded)); + assert(llama_token_cache_lookup(LLAMA_TOKEN_CACHE_KIND_TEXT, "second", loaded)); + assert(loaded == second); + + const std::vector async_payload(16, 0x33); + llama_token_cache_store(LLAMA_TOKEN_CACHE_KIND_TEXT, "async", async_payload); + assert(llama_token_cache_flush()); + loaded.clear(); + assert(llama_token_cache_lookup(LLAMA_TOKEN_CACHE_KIND_TEXT, "async", loaded)); + assert(loaded == async_payload); + + const auto stats = llama_token_cache_get_stats(); + assert(stats.entries == 2); + assert(stats.bytes_used == second.size() + async_payload.size()); + assert(stats.evictions >= 1); + assert(stats.writes >= 3); + + assert(llama_token_cache_clear()); + assert(llama_token_cache_get_stats().entries == 0); + + std::error_code ec; + fs::remove_all(root, ec); + return 0; +} diff --git a/tools/cli/README.md b/tools/cli/README.md index bcddd05..ac147a0 100644 --- a/tools/cli/README.md +++ b/tools/cli/README.md @@ -32,6 +32,9 @@ | `-ub, --ubatch-size N` | physical maximum batch size (default: 512)
(env: LLAMA_ARG_UBATCH) | | `--keep N` | number of tokens to keep from the initial prompt (default: 0, -1 = all) | | `--swa-full` | use full-size SWA cache (default: false)
[(more info)](https://github.com/ggml-org/llama.cpp/pull/13194#issuecomment-2868343055)
(env: LLAMA_ARG_SWA_FULL) | +| `--token-cache-size N` | shared tokenizer and tokenized-text cache size in MiB (default: 5120, 0 = disable)
(env: LLAMA_ARG_TOKEN_CACHE_SIZE) | +| `--token-cache-dir PATH` | path to the shared tokenizer and tokenized-text cache database
(env: LLAMA_ARG_TOKEN_CACHE_DIR) | +| `--token-cache-clear` | clear the shared tokenizer and tokenized-text cache before loading the model
(env: LLAMA_ARG_TOKEN_CACHE_CLEAR) | | `-fa, --flash-attn [on\|off\|auto]` | set Flash Attention use ('on', 'off', or 'auto', default: 'auto')
(env: LLAMA_ARG_FLASH_ATTN) | | `-p, --prompt PROMPT` | prompt to start generation with; for system message, use -sys | | `--perf, --no-perf` | whether to enable internal libllama performance timings (default: false)
(env: LLAMA_ARG_PERF) | @@ -54,11 +57,12 @@ | `-ctv, --cache-type-v TYPE` | KV cache data type for V
allowed values: f32, f16, bf16, q8_0, q4_0, q4_1, iq4_nl, q5_0, q5_1
(default: f16)
(env: LLAMA_ARG_CACHE_TYPE_V) | | `-dt, --defrag-thold N` | KV cache defragmentation threshold (DEPRECATED)
(env: LLAMA_ARG_DEFRAG_THOLD) | | `-np, --parallel N` | number of parallel sequences to decode (default: 1)
(env: LLAMA_ARG_N_PARALLEL) | -| `--rpc SERVERS` | comma-separated list of RPC servers (host:port)
(env: LLAMA_ARG_RPC) | | `--mlock` | DEPRECATED in favor of `--load-mode`: force system to keep model in RAM rather than swapping or compressing
(env: LLAMA_ARG_MLOCK) | | `--mmap, --no-mmap` | DEPRECATED in favor of `--load-mode`: whether to memory-map model. (if mmap disabled, slower load but may reduce pageouts if not using mlock)
(env: LLAMA_ARG_MMAP) | | `-dio, --direct-io, -ndio, --no-direct-io` | DEPRECATED in favor of `--load-mode`: use DirectIO if available
(env: LLAMA_ARG_DIO) | -| `-lm, --load-mode MODE` | model loading mode (default: mmap)
- none: no special loading mode
- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)
- mlock: force system to keep model in RAM rather than swapping or compressing
- mmap+mlock: mmap + force system to keep model in RAM rather than swapping or compressing
- dio: use DirectIO if available

(env: LLAMA_ARG_LOAD_MODE) | +| `-lm, --load-mode MODE` | model loading mode (default: mmap)
- none: no special loading mode
- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)
- mlock: force system to keep model in RAM rather than swapping or compressing
- mmap+mlock: mmap + force system to keep model in RAM rather than swapping or compressing
- dio: use DirectIO if available
- longhaul: stream routed MoE experts through a bounded cache

(env: LLAMA_ARG_LOAD_MODE) | +| `--longhaul` | stream routed Qwen3.5 MoE experts through a bounded Metal cache
(env: LLAMA_ARG_LONGHAUL) | +| `--longhaul-cache N` | longhaul expert cache size in GiB
(env: LLAMA_ARG_LONGHAUL_CACHE) | | `--numa TYPE` | attempt optimizations that help on some NUMA systems
- distribute: spread execution evenly over all nodes
- isolate: only spawn threads on CPUs on the node that execution started on
- numactl: use the CPU map provided by numactl
if run without this previously, it is recommended to drop the system page cache before using this
see https://github.com/ggml-org/llama.cpp/issues/1437
(env: LLAMA_ARG_NUMA) | | `-dev, --device ` | comma-separated list of devices to use for offloading (none = don't offload)
use --list-devices to see a list of available devices
(env: LLAMA_ARG_DEVICE) | | `--list-devices` | print list of available devices and exit | diff --git a/tools/completion/README.md b/tools/completion/README.md index bce71d6..d2b2df3 100644 --- a/tools/completion/README.md +++ b/tools/completion/README.md @@ -115,6 +115,9 @@ llama-completion.exe -m models\gemma-1.1-7b-it.Q4_K_M.gguf --ignore-eos -n -1 | `-ub, --ubatch-size N` | physical maximum batch size (default: 512)
(env: LLAMA_ARG_UBATCH) | | `--keep N` | number of tokens to keep from the initial prompt (default: 0, -1 = all) | | `--swa-full` | use full-size SWA cache (default: false)
[(more info)](https://github.com/ggml-org/llama.cpp/pull/13194#issuecomment-2868343055)
(env: LLAMA_ARG_SWA_FULL) | +| `--token-cache-size N` | shared tokenizer and tokenized-text cache size in MiB (default: 5120, 0 = disable)
(env: LLAMA_ARG_TOKEN_CACHE_SIZE) | +| `--token-cache-dir PATH` | path to the shared tokenizer and tokenized-text cache database
(env: LLAMA_ARG_TOKEN_CACHE_DIR) | +| `--token-cache-clear` | clear the shared tokenizer and tokenized-text cache before loading the model
(env: LLAMA_ARG_TOKEN_CACHE_CLEAR) | | `-fa, --flash-attn [on\|off\|auto]` | set Flash Attention use ('on', 'off', or 'auto', default: 'auto')
(env: LLAMA_ARG_FLASH_ATTN) | | `-p, --prompt PROMPT` | prompt to start generation with; for system message, use -sys | | `--perf, --no-perf` | whether to enable internal libllama performance timings (default: false)
(env: LLAMA_ARG_PERF) | @@ -137,11 +140,12 @@ llama-completion.exe -m models\gemma-1.1-7b-it.Q4_K_M.gguf --ignore-eos -n -1 | `-ctv, --cache-type-v TYPE` | KV cache data type for V
allowed values: f32, f16, bf16, q8_0, q4_0, q4_1, iq4_nl, q5_0, q5_1
(default: f16)
(env: LLAMA_ARG_CACHE_TYPE_V) | | `-dt, --defrag-thold N` | KV cache defragmentation threshold (DEPRECATED)
(env: LLAMA_ARG_DEFRAG_THOLD) | | `-np, --parallel N` | number of parallel sequences to decode (default: 1)
(env: LLAMA_ARG_N_PARALLEL) | -| `--rpc SERVERS` | comma-separated list of RPC servers (host:port)
(env: LLAMA_ARG_RPC) | | `--mlock` | DEPRECATED in favor of `--load-mode`: force system to keep model in RAM rather than swapping or compressing
(env: LLAMA_ARG_MLOCK) | | `--mmap, --no-mmap` | DEPRECATED in favor of `--load-mode`: whether to memory-map model. (if mmap disabled, slower load but may reduce pageouts if not using mlock)
(env: LLAMA_ARG_MMAP) | | `-dio, --direct-io, -ndio, --no-direct-io` | DEPRECATED in favor of `--load-mode`: use DirectIO if available
(env: LLAMA_ARG_DIO) | -| `-lm, --load-mode MODE` | model loading mode (default: mmap)
- none: no special loading mode
- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)
- mlock: force system to keep model in RAM rather than swapping or compressing
- mmap+mlock: mmap + force system to keep model in RAM rather than swapping or compressing
- dio: use DirectIO if available

(env: LLAMA_ARG_LOAD_MODE) | +| `-lm, --load-mode MODE` | model loading mode (default: mmap)
- none: no special loading mode
- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)
- mlock: force system to keep model in RAM rather than swapping or compressing
- mmap+mlock: mmap + force system to keep model in RAM rather than swapping or compressing
- dio: use DirectIO if available
- longhaul: stream routed MoE experts through a bounded cache

(env: LLAMA_ARG_LOAD_MODE) | +| `--longhaul` | stream routed Qwen3.5 MoE experts through a bounded Metal cache
(env: LLAMA_ARG_LONGHAUL) | +| `--longhaul-cache N` | longhaul expert cache size in GiB
(env: LLAMA_ARG_LONGHAUL_CACHE) | | `--numa TYPE` | attempt optimizations that help on some NUMA systems
- distribute: spread execution evenly over all nodes
- isolate: only spawn threads on CPUs on the node that execution started on
- numactl: use the CPU map provided by numactl
if run without this previously, it is recommended to drop the system page cache before using this
see https://github.com/ggml-org/llama.cpp/issues/1437
(env: LLAMA_ARG_NUMA) | | `-dev, --device ` | comma-separated list of devices to use for offloading (none = don't offload)
use --list-devices to see a list of available devices
(env: LLAMA_ARG_DEVICE) | | `--list-devices` | print list of available devices and exit | diff --git a/tools/server/README.md b/tools/server/README.md index f45c018..2df8893 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -53,6 +53,9 @@ For the full list of features, please refer to [server's changelog](https://gith | `-ub, --ubatch-size N` | physical maximum batch size (default: 512)
(env: LLAMA_ARG_UBATCH) | | `--keep N` | number of tokens to keep from the initial prompt (default: 0, -1 = all) | | `--swa-full` | use full-size SWA cache (default: false)
[(more info)](https://github.com/ggml-org/llama.cpp/pull/13194#issuecomment-2868343055)
(env: LLAMA_ARG_SWA_FULL) | +| `--token-cache-size N` | shared tokenizer and tokenized-text cache size in MiB (default: 5120, 0 = disable)
(env: LLAMA_ARG_TOKEN_CACHE_SIZE) | +| `--token-cache-dir PATH` | path to the shared tokenizer and tokenized-text cache database
(env: LLAMA_ARG_TOKEN_CACHE_DIR) | +| `--token-cache-clear` | clear the shared tokenizer and tokenized-text cache before loading the model
(env: LLAMA_ARG_TOKEN_CACHE_CLEAR) | | `-fa, --flash-attn [on\|off\|auto]` | set Flash Attention use ('on', 'off', or 'auto', default: 'auto')
(env: LLAMA_ARG_FLASH_ATTN) | | `--perf, --no-perf` | whether to enable internal libllama performance timings (default: false)
(env: LLAMA_ARG_PERF) | | `-e, --escape, --no-escape` | whether to process escapes sequences (\n, \r, \t, \', \", \\) (default: true) | @@ -71,11 +74,12 @@ For the full list of features, please refer to [server's changelog](https://gith | `-ctk, --cache-type-k TYPE` | KV cache data type for K
allowed values: f32, f16, bf16, q8_0, q4_0, q4_1, iq4_nl, q5_0, q5_1
(default: f16)
(env: LLAMA_ARG_CACHE_TYPE_K) | | `-ctv, --cache-type-v TYPE` | KV cache data type for V
allowed values: f32, f16, bf16, q8_0, q4_0, q4_1, iq4_nl, q5_0, q5_1
(default: f16)
(env: LLAMA_ARG_CACHE_TYPE_V) | | `-dt, --defrag-thold N` | KV cache defragmentation threshold (DEPRECATED)
(env: LLAMA_ARG_DEFRAG_THOLD) | -| `--rpc SERVERS` | comma-separated list of RPC servers (host:port)
(env: LLAMA_ARG_RPC) | | `--mlock` | DEPRECATED in favor of `--load-mode`: force system to keep model in RAM rather than swapping or compressing
(env: LLAMA_ARG_MLOCK) | | `--mmap, --no-mmap` | DEPRECATED in favor of `--load-mode`: whether to memory-map model. (if mmap disabled, slower load but may reduce pageouts if not using mlock)
(env: LLAMA_ARG_MMAP) | | `-dio, --direct-io, -ndio, --no-direct-io` | DEPRECATED in favor of `--load-mode`: use DirectIO if available
(env: LLAMA_ARG_DIO) | -| `-lm, --load-mode MODE` | model loading mode (default: mmap)
- none: no special loading mode
- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)
- mlock: force system to keep model in RAM rather than swapping or compressing
- mmap+mlock: mmap + force system to keep model in RAM rather than swapping or compressing
- dio: use DirectIO if available

(env: LLAMA_ARG_LOAD_MODE) | +| `-lm, --load-mode MODE` | model loading mode (default: mmap)
- none: no special loading mode
- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)
- mlock: force system to keep model in RAM rather than swapping or compressing
- mmap+mlock: mmap + force system to keep model in RAM rather than swapping or compressing
- dio: use DirectIO if available
- longhaul: stream routed MoE experts through a bounded cache

(env: LLAMA_ARG_LOAD_MODE) | +| `--longhaul` | stream routed Qwen3.5 MoE experts through a bounded Metal cache
(env: LLAMA_ARG_LONGHAUL) | +| `--longhaul-cache N` | longhaul expert cache size in GiB
(env: LLAMA_ARG_LONGHAUL_CACHE) | | `--numa TYPE` | attempt optimizations that help on some NUMA systems
- distribute: spread execution evenly over all nodes
- isolate: only spawn threads on CPUs on the node that execution started on
- numactl: use the CPU map provided by numactl
if run without this previously, it is recommended to drop the system page cache before using this
see https://github.com/ggml-org/llama.cpp/issues/1437
(env: LLAMA_ARG_NUMA) | | `-dev, --device ` | comma-separated list of devices to use for offloading (none = don't offload)
use --list-devices to see a list of available devices
(env: LLAMA_ARG_DEVICE) | | `--list-devices` | print list of available devices and exit | @@ -199,6 +203,8 @@ For the full list of features, please refer to [server's changelog](https://gith | `--ui-config-file, --webui-config-file PATH` | JSON file that provides default UI settings (overrides UI defaults)
(env: LLAMA_ARG_UI_CONFIG_FILE) | | `--ui-mcp-proxy, --webui-mcp-proxy, --no-ui-mcp-proxy, --no-webui-mcp-proxy` | experimental: whether to enable MCP CORS proxy - do not enable in untrusted environments (default: disabled)
(env: LLAMA_ARG_UI_MCP_PROXY) | | `--tools TOOL1,TOOL2,...` | experimental: whether to enable built-in tools for AI agents - do not enable in untrusted environments (default: no tools)
specify "all" to enable all tools
available tools: read_file, file_glob_search, grep_search, exec_shell_command, write_file, edit_file, get_datetime
note: for security reasons, this will limit --cors-origins to localhost by default
(env: LLAMA_ARG_TOOLS) | +| `--mcp-servers-config PATH` | experimental: path to JSON file with MCP server definitions (Cursor-compatible format) - do not enable in untrusted environments (default: none)
note: for security reasons, this will limit --cors-origins to localhost by default
(env: LLAMA_ARG_MCP_SERVERS_CONFIG) | +| `--mcp-servers-json JSON` | experimental: inline JSON with MCP server definitions (Cursor-compatible format) - do not enable in untrusted environments (default: none)
note: for security reasons, this will limit --cors-origins to localhost by default
(env: LLAMA_ARG_MCP_SERVERS_JSON) | | `-ag, --agent, -no-ag, --no-agent` | whether to enable CORS proxy and all built-in tools - do not enable in untrusted environments (default: disabled)
note: for security reasons, this will limit --cors-origins to localhost by default
(env: LLAMA_ARG_AGENT) | | `--ui, --webui, --no-ui, --no-webui` | whether to enable the Web UI (default: enabled)
(env: LLAMA_ARG_UI) | | `--embedding, --embeddings` | restrict to only support embedding use case; use only with dedicated embedding models (default: disabled)
(env: LLAMA_ARG_EMBEDDINGS) | diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 1621eb3..b316705 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -38,6 +38,7 @@ using json = nlohmann::ordered_json; constexpr int HTTP_POLLING_SECONDS = 1; +constexpr int64_t PROMPT_TIMING_LOG_INTERVAL_US = 5*1000*1000; static uint32_t server_n_outputs_max(const common_params & params) { const uint32_t n_batch = params.n_batch; @@ -276,6 +277,7 @@ struct server_slot { // TODO @ngxson : move all metrics to a sub-struct for clarity int64_t t_start_process_prompt; int64_t t_start_generation; + int64_t t_prompt_print_last = 0; int64_t t_print_last = 0; int32_t n_decoded_last = 0; @@ -569,13 +571,17 @@ struct server_slot { SLT_INF(*this, "n_decoded = %6d, tg = %6.2f t/s, tg_3s = %6.2f t/s\n", n_decoded, n_gen_second, n_gen_second_win); } - void print_timings_pp() const { + void print_timings_pp() { + const int64_t t_now = ggml_time_us(); + + if (t_now - t_prompt_print_last < PROMPT_TIMING_LOG_INTERVAL_US) { + return; + } + const double n_prompt_second = 1e3 / t_prompt_processing * n_prompt_tokens_processed; const double f_progress = (float) prompt.n_tokens() / task->n_tokens(); - if (t_prompt_processing < 3000.0) { - return; - } + t_prompt_print_last = t_now; SLT_INF(*this, "prompt processing, n_tokens = %6d, progress = %.2f, t = %6.2f s / %.2f tokens per second\n", n_prompt_tokens_processed, f_progress, t_prompt_processing / 1e3, n_prompt_second); @@ -3072,6 +3078,7 @@ private: // TODO: maybe move branch to outside of this loop in the future if (slot.state == SLOT_STATE_STARTED) { slot.t_start_process_prompt = ggml_time_us(); + slot.t_prompt_print_last = slot.t_start_process_prompt; slot.t_start_generation = 0; slot.state = SLOT_STATE_PROCESSING_PROMPT;