owen/longhaul.cpp
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

stuff

Browse Source
This commit is contained in:
Owen Qwen
2026-07-29 17:17:25 -05:00
parent fdd9fcf85e
commit a673afad85
17 changed files with 1071 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/CMakeLists.txt
+8
View File
@@ -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)
src/llama-model-loader.cpp
+33
View File
@@ -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 <cstdint>
#include <cstring>
#include <future>
#include <filesystem>
#include <regex>
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<uint64_t>(size)) + '\n' +
std::to_string(time_ec ? 0 : static_cast<int64_t>(mtime));
std::vector<std::string> 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()) {
src/llama-model-loader.h
+3
View File
@@ -107,6 +107,7 @@ struct llama_model_loader {
std::vector<ggml_context_ptr> 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;
src/llama-token-cache.cpp
+514
View File
@@ -0,0 +1,514 @@
#include "llama-token-cache.h"
#include "llama-impl.h"
#include <sqlite3.h>
#include <algorithm>
#include <array>
#include <atomic>
#include <condition_variable>
#include <chrono>
#include <cstdlib>
#include <cstring>
#include <deque>
#include <filesystem>
#include <limits>
#include <mutex>
#include <thread>
#if defined(_WIN32)
# include <windows.h>
#else
# include <sys/stat.h>
# include <unistd.h>
#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<const uint8_t *>(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<uint8_t, 32> 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<uint8_t, 32> 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<uint8_t> payload;
};
struct cache_counters {
std::atomic<uint64_t> tokenizer_hits{0};
std::atomic<uint64_t> tokenizer_misses{0};
std::atomic<uint64_t> text_hits{0};
std::atomic<uint64_t> text_misses{0};
std::atomic<uint64_t> writes{0};
std::atomic<uint64_t> dropped_writes{0};
std::atomic<uint64_t> evictions{0};
std::atomic<uint64_t> errors{0};
};
class token_cache {
public:
token_cache() {
params.capacity_bytes = DEFAULT_CAPACITY;
worker = std::thread([this] { run(); });
}
~token_cache() {
flush();
{
std::lock_guard<std::mutex> lock(mutex);
stopping = true;
}
ready.notify_all();
if (worker.joinable()) worker.join();
close();
}
bool configure(llama_token_cache_params value) {
std::lock_guard<std::mutex> 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<uint8_t> & out) {
if (!ensure_open()) return false;
std::lock_guard<std::mutex> 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<const uint8_t *>(p), static_cast<const uint8_t *>(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<uint8_t> & payload, bool sync) {
if (params.capacity_bytes == 0 || payload.size() > params.capacity_bytes ||
payload.size() > static_cast<size_t>(std::numeric_limits<int>::max())) return;
if (sync) {
write_one({kind, key, payload});
return;
}
std::lock_guard<std::mutex> 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<std::mutex> lock(mutex);
drained.wait(lock, [&] { return queue.empty() && !writing; });
if (db) {
lock.unlock();
std::lock_guard<std::mutex> 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<std::mutex> 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<std::mutex> 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::microseconds>(
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<std::mutex> 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<sqlite3_int64>(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<std::mutex> 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<std::mutex> lock(mutex);
writing = false;
if (queue.empty()) drained.notify_all();
}
}
}
void write_one(const queued_write & item) {
if (!ensure_open()) return;
std::lock_guard<std::mutex> 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<uint64_t>(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<const char *>(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<std::mutex> 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<queued_write> queue;
uint64_t queue_bytes = 0;
bool initialized = false;
bool stopping = false;
bool writing = false;
std::atomic<bool> 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<uint8_t> & payload) {
return cache().lookup(kind, key, payload);
}
void llama_token_cache_store(llama_token_cache_kind kind, const std::string & key, const std::vector<uint8_t> & 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();
}
src/llama-token-cache.h
+29
View File
@@ -0,0 +1,29 @@
#pragma once
#include "llama.h"
#include <cstdint>
#include <string>
#include <vector>
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<uint8_t> & payload);
void llama_token_cache_store(
llama_token_cache_kind kind,
const std::string & key,
const std::vector<uint8_t> & payload,
bool synchronous = false);
void llama_token_cache_note_tokenizer_hit(bool hit);
src/llama-vocab.cpp
+276 -1
View File
@@ -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<llama_token> suppress_tokens;
std::unique_ptr<llm_tokenizer> tokenizer;
std::string token_cache_id;
std::vector<char> 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<uint8_t> & 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<llama_token> 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<uint8_t> 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<const char *>(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<uint8_t> snapshot;
auto put = [&snapshot](const auto & value) {
const auto * p = reinterpret_cast<const uint8_t *>(&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<std::pair<std::string, int>> 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<std::pair<std::string, llama_token>> 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<uint8_t> 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<uint8_t> 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<uint8_t> & 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<const char *>(snapshot.data() + offset), static_cast<size_t>(n));
offset += static_cast<size_t>(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<size_t>(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_token> llama_vocab::impl::tokenize(
std::vector<llama_token> llama_vocab::impl::tokenize_uncached(
const std::string & raw_text,
bool add_special,
bool parse_special) const {
@@ -3516,6 +3745,52 @@ std::vector<llama_token> llama_vocab::impl::tokenize(
return output;
}
std::vector<llama_token> 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<uint8_t> 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<uint32_t>(std::numeric_limits<int32_t>::max())) {
std::vector<llama_token> 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<size_t>(token) < id_to_token.size();
})) {
return result;
}
}
}
auto result = tokenize_uncached(raw_text, add_special, parse_special);
if (result.size() < static_cast<size_t>(std::numeric_limits<uint32_t>::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;