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

Add tokenizer longhaul cache

Browse Source
This commit is contained in:
Owen Qwen
2026-07-30 23:32:52 -05:00
parent a673afad85
commit e0665b1042
21 changed files with 1078 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/CMakeLists.txt
+1
View File
@@ -38,6 +38,7 @@ add_library(llama
llama-quant.cpp
llama-sampler.cpp
llama-token-cache.cpp
llama-tokenizer-longhaul.cpp
llama-vocab.cpp
unicode-data.cpp
unicode.cpp
src/llama-model.cpp
+3 -1
View File
@@ -1249,7 +1249,7 @@ void llama_model_base::load_hparams(llama_model_loader & ml) {
void llama_model_base::load_vocab(llama_model_loader & ml) {
const auto kv = LLM_KV(arch);
vocab.load(ml, kv);
vocab.load(ml, kv, params.tokenizer_longhaul, params.tokenizer_longhaul_cache_bytes);
}
bool llama_model_base::load_tensors(llama_model_loader & ml) {
@@ -2380,6 +2380,8 @@ llama_model_params llama_model_default_params() {
/*.no_host =*/ false,
/*.no_alloc =*/ false,
/*.longhaul_cache_bytes =*/ 0,
/*.tokenizer_longhaul =*/ false,
/*.tokenizer_longhaul_cache_bytes =*/ 0,
};
return result;
src/llama-token-cache.cpp
+77 -1
View File
@@ -214,6 +214,10 @@ public:
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();
std::error_code ec;
fs::remove_all(fs::path(path).parent_path() / "tokenizer-indexes", ec);
if (ec) return false;
sqlite3_exec(db, "INSERT OR REPLACE INTO meta(key,value) VALUES('external_bytes',0)", nullptr, nullptr, nullptr);
return true;
}
@@ -239,6 +243,11 @@ public:
result.bytes_used = sqlite3_column_int64(stmt, 1);
}
sqlite3_finalize(stmt);
if (sqlite3_prepare_v2(db, "SELECT value FROM meta WHERE key='external_bytes'", -1, &stmt, nullptr) == SQLITE_OK &&
sqlite3_step(stmt) == SQLITE_ROW) {
result.bytes_used += static_cast<uint64_t>(sqlite3_column_int64(stmt, 0));
}
sqlite3_finalize(stmt);
}
return result;
}
@@ -247,6 +256,45 @@ public:
if (hit) ++stats.tokenizer_hits; else ++stats.tokenizer_misses;
}
std::string database_path() {
return ensure_open() ? path : std::string();
}
uint64_t capacity() const {
return params.capacity_bytes;
}
void set_external_bytes(uint64_t bytes) {
if (!ensure_open()) return;
std::lock_guard<std::mutex> db_lock(db_mutex);
sqlite3_exec(db, "BEGIN IMMEDIATE", nullptr, nullptr, nullptr);
sqlite3_stmt * stmt = nullptr;
sqlite3_prepare_v2(db, "INSERT OR REPLACE INTO meta(key,value) VALUES('external_bytes',?1)", -1, &stmt, nullptr);
sqlite3_bind_int64(stmt, 1, static_cast<sqlite3_int64>(bytes));
sqlite3_step(stmt);
sqlite3_finalize(stmt);
uint64_t total = 0;
sqlite3_prepare_v2(db, "SELECT COALESCE(SUM(size),0) FROM entries", -1, &stmt, nullptr);
if (sqlite3_step(stmt) == SQLITE_ROW) total = sqlite3_column_int64(stmt, 0);
sqlite3_finalize(stmt);
const uint64_t entry_capacity = bytes >= params.capacity_bytes ? 0 : params.capacity_bytes - bytes;
while (total > entry_capacity) {
uint64_t victim_size = 0;
sqlite3_prepare_v2(db, "SELECT size FROM entries ORDER BY last_access LIMIT 1", -1, &stmt, nullptr);
if (sqlite3_step(stmt) == SQLITE_ROW) victim_size = sqlite3_column_int64(stmt, 0);
sqlite3_finalize(stmt);
if (victim_size == 0 || sqlite3_exec(db,
"DELETE FROM entries WHERE rowid=(SELECT rowid FROM entries ORDER BY last_access LIMIT 1)",
nullptr, nullptr, nullptr) != SQLITE_OK) {
break;
}
total -= std::min(total, victim_size);
++stats.evictions;
}
sqlite3_exec(db, "COMMIT", nullptr, nullptr, nullptr);
}
private:
static int64_t now_tick() {
return std::chrono::duration_cast<std::chrono::microseconds>(
@@ -311,13 +359,22 @@ private:
sqlite3_bind_int64(stmt, 1, static_cast<sqlite3_int64>(params.capacity_bytes));
sqlite3_step(stmt);
sqlite3_finalize(stmt);
uint64_t external_bytes = 0;
sqlite3_prepare_v2(db, "SELECT value FROM meta WHERE key='external_bytes'", -1, &stmt, nullptr);
if (sqlite3_step(stmt) == SQLITE_ROW) {
external_bytes = static_cast<uint64_t>(sqlite3_column_int64(stmt, 0));
}
sqlite3_finalize(stmt);
const uint64_t entry_capacity = external_bytes >= params.capacity_bytes
? 0
: params.capacity_bytes - external_bytes;
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 (total <= entry_capacity) 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) {
@@ -355,12 +412,19 @@ private:
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;
uint64_t external_bytes = 0;
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 (sqlite3_prepare_v2(db, "SELECT value FROM meta WHERE key='external_bytes'", -1, &capacity_stmt, nullptr) == SQLITE_OK &&
sqlite3_step(capacity_stmt) == SQLITE_ROW) {
external_bytes = static_cast<uint64_t>(sqlite3_column_int64(capacity_stmt, 0));
}
sqlite3_finalize(capacity_stmt);
capacity = external_bytes >= capacity ? 0 : capacity - external_bytes;
if (item.payload.size() > capacity) {
sqlite3_exec(db, "ROLLBACK", nullptr, nullptr, nullptr);
return;
@@ -492,6 +556,18 @@ void llama_token_cache_note_tokenizer_hit(bool hit) {
cache().note_tokenizer(hit);
}
std::string llama_token_cache_database_path() {
return cache().database_path();
}
uint64_t llama_token_cache_capacity() {
return cache().capacity();
}
void llama_token_cache_set_external_bytes(uint64_t bytes) {
cache().set_external_bytes(bytes);
}
llama_token_cache_params llama_token_cache_default_params() {
return {nullptr, DEFAULT_CAPACITY};
}
src/llama-token-cache.h
+5
View File
@@ -27,3 +27,8 @@ void llama_token_cache_store(
void llama_token_cache_note_tokenizer_hit(bool hit);
// Internal helpers used by tokenizer-longhaul artifacts.
// Opening the cache here also fixes the configured path for the process.
std::string llama_token_cache_database_path();
uint64_t llama_token_cache_capacity();
void llama_token_cache_set_external_bytes(uint64_t bytes);
src/llama-tokenizer-longhaul.cpp
+471
View File
@@ -0,0 +1,471 @@
#include "llama-tokenizer-longhaul.h"
#include "llama-token-cache.h"
#include <sqlite3.h>
#include <algorithm>
#include <chrono>
#include <filesystem>
#include <list>
#include <mutex>
#include <thread>
#include <unordered_map>
#if defined(_WIN32)
# include <process.h>
#else
# include <unistd.h>
#endif
namespace fs = std::filesystem;
namespace {
constexpr int INDEX_VERSION = 1;
int process_id() {
#if defined(_WIN32)
return _getpid();
#else
return getpid();
#endif
}
bool exec_sql(sqlite3 * db, const char * sql, std::string & error) {
char * message = nullptr;
if (sqlite3_exec(db, sql, nullptr, nullptr, &message) == SQLITE_OK) {
return true;
}
error = message ? message : sqlite3_errmsg(db);
sqlite3_free(message);
return false;
}
struct lock_directory {
fs::path path;
bool owned = false;
~lock_directory() {
if (owned) {
std::error_code ec;
fs::remove(path, ec);
}
}
};
uint64_t tokenizer_artifact_bytes(const fs::path & directory) {
uint64_t total = 0;
std::error_code ec;
for (const auto & entry : fs::directory_iterator(directory, ec)) {
if (ec) break;
if (!entry.is_regular_file(ec) || entry.path().extension() != ".sqlite3") continue;
total += entry.file_size(ec);
if (ec) {
ec.clear();
}
}
return total;
}
} // namespace
struct llama_tokenizer_longhaul_index::impl {
struct cache_entry {
int value;
uint64_t bytes;
std::list<std::string>::iterator position;
};
std::string source_key;
uint64_t capacity;
fs::path artifact;
sqlite3 * db = nullptr;
std::string id;
uint64_t merge_count = 0;
mutable std::mutex mutex;
mutable uint64_t used = 0;
mutable std::list<std::string> lru;
mutable std::unordered_map<std::string, cache_entry> values;
impl(std::string source_key, uint64_t capacity)
: source_key(std::move(source_key)), capacity(capacity) {
}
~impl() {
if (db) {
sqlite3_close(db);
}
}
bool open_existing(std::string & error) {
if (!fs::exists(artifact)) {
return false;
}
sqlite3 * opened = nullptr;
if (sqlite3_open_v2(artifact.string().c_str(), &opened,
SQLITE_OPEN_READONLY | SQLITE_OPEN_FULLMUTEX, nullptr) != SQLITE_OK) {
error = opened ? sqlite3_errmsg(opened) : "failed to open tokenizer index";
if (opened) sqlite3_close(opened);
return false;
}
sqlite3_stmt * stmt = nullptr;
const char * sql =
"SELECT version, tokenizer_id, n_merges FROM metadata WHERE complete=1 LIMIT 1";
bool valid = sqlite3_prepare_v2(opened, sql, -1, &stmt, nullptr) == SQLITE_OK &&
sqlite3_step(stmt) == SQLITE_ROW &&
sqlite3_column_int(stmt, 0) == INDEX_VERSION;
if (valid) {
const unsigned char * value = sqlite3_column_text(stmt, 1);
id = value ? reinterpret_cast<const char *>(value) : "";
merge_count = static_cast<uint64_t>(sqlite3_column_int64(stmt, 2));
valid = !id.empty();
}
sqlite3_finalize(stmt);
if (!valid) {
error = "tokenizer index is incomplete or has an unsupported version";
sqlite3_close(opened);
return false;
}
db = opened;
sqlite3_busy_timeout(db, 5000);
return true;
}
bool build(
const std::vector<llama_tokenizer_longhaul_reverse> & reverse,
const std::vector<llama_tokenizer_longhaul_merge> & merges,
const std::string & tokenizer_id,
std::string & error) {
const fs::path temp = artifact.string() + ".tmp." + std::to_string(process_id());
std::error_code ec;
fs::remove(temp, ec);
sqlite3 * output = nullptr;
if (sqlite3_open_v2(temp.string().c_str(), &output,
SQLITE_OPEN_READWRITE | SQLITE_OPEN_CREATE | SQLITE_OPEN_EXCLUSIVE, nullptr) != SQLITE_OK) {
error = output ? sqlite3_errmsg(output) : "failed to create tokenizer index";
if (output) sqlite3_close(output);
return false;
}
auto fail = [&](const std::string & message) {
error = message;
sqlite3_close(output);
output = nullptr;
std::error_code remove_ec;
fs::remove(temp, remove_ec);
return false;
};
if (!exec_sql(output,
"PRAGMA journal_mode=OFF;"
"PRAGMA synchronous=OFF;"
"PRAGMA locking_mode=EXCLUSIVE;"
"CREATE TABLE metadata("
" version INTEGER NOT NULL,complete INTEGER NOT NULL,tokenizer_id TEXT NOT NULL,"
" n_reverse INTEGER NOT NULL,n_merges INTEGER NOT NULL);"
"CREATE TABLE reverse_tokens(text BLOB PRIMARY KEY,id INTEGER NOT NULL) WITHOUT ROWID;"
"CREATE TABLE merges(left_text BLOB NOT NULL,right_text BLOB NOT NULL,rank INTEGER NOT NULL,"
" PRIMARY KEY(left_text,right_text)) WITHOUT ROWID;"
"BEGIN IMMEDIATE;", error)) {
return fail(error);
}
sqlite3_stmt * insert_reverse = nullptr;
sqlite3_stmt * insert_merge = nullptr;
if (sqlite3_prepare_v2(output,
"INSERT INTO reverse_tokens(text,id) VALUES(?1,?2)", -1, &insert_reverse, nullptr) != SQLITE_OK ||
sqlite3_prepare_v2(output,
"INSERT INTO merges(left_text,right_text,rank) VALUES(?1,?2,?3)",
-1, &insert_merge, nullptr) != SQLITE_OK) {
if (insert_reverse) sqlite3_finalize(insert_reverse);
if (insert_merge) sqlite3_finalize(insert_merge);
return fail(sqlite3_errmsg(output));
}
for (const auto & item : reverse) {
sqlite3_bind_blob(insert_reverse, 1, item.text.data(), int(item.text.size()), SQLITE_STATIC);
sqlite3_bind_int(insert_reverse, 2, item.id);
if (sqlite3_step(insert_reverse) != SQLITE_DONE) {
const std::string message = sqlite3_errmsg(output);
sqlite3_finalize(insert_reverse);
sqlite3_finalize(insert_merge);
return fail(message);
}
sqlite3_reset(insert_reverse);
sqlite3_clear_bindings(insert_reverse);
}
for (const auto & item : merges) {
sqlite3_bind_blob(insert_merge, 1, item.left.data(), int(item.left.size()), SQLITE_STATIC);
sqlite3_bind_blob(insert_merge, 2, item.right.data(), int(item.right.size()), SQLITE_STATIC);
sqlite3_bind_int(insert_merge, 3, item.rank);
if (sqlite3_step(insert_merge) != SQLITE_DONE) {
const std::string message = sqlite3_errmsg(output);
sqlite3_finalize(insert_reverse);
sqlite3_finalize(insert_merge);
return fail(message);
}
sqlite3_reset(insert_merge);
sqlite3_clear_bindings(insert_merge);
}
sqlite3_finalize(insert_reverse);
sqlite3_finalize(insert_merge);
sqlite3_stmt * metadata = nullptr;
if (sqlite3_prepare_v2(output,
"INSERT INTO metadata VALUES(?1,1,?2,?3,?4)", -1, &metadata, nullptr) != SQLITE_OK) {
return fail(sqlite3_errmsg(output));
}
sqlite3_bind_int(metadata, 1, INDEX_VERSION);
sqlite3_bind_text(metadata, 2, tokenizer_id.c_str(), -1, SQLITE_STATIC);
sqlite3_bind_int64(metadata, 3, static_cast<sqlite3_int64>(reverse.size()));
sqlite3_bind_int64(metadata, 4, static_cast<sqlite3_int64>(merges.size()));
const bool metadata_ok = sqlite3_step(metadata) == SQLITE_DONE;
sqlite3_finalize(metadata);
if (!metadata_ok || !exec_sql(output, "COMMIT;", error)) {
return fail(metadata_ok ? error : sqlite3_errmsg(output));
}
sqlite3_close(output);
output = nullptr;
const uint64_t persistent_capacity = llama_token_cache_capacity();
const uint64_t artifact_size = fs::file_size(temp, ec);
if (ec || artifact_size > persistent_capacity) {
fs::remove(temp, ec);
error = "completed tokenizer index does not fit in --token-cache-size";
return false;
}
struct artifact_candidate {
fs::path path;
fs::file_time_type modified;
uint64_t size;
};
std::vector<artifact_candidate> candidates;
uint64_t existing_bytes = 0;
for (const auto & entry : fs::directory_iterator(artifact.parent_path(), ec)) {
if (ec) break;
if (!entry.is_regular_file(ec) || entry.path().extension() != ".sqlite3" ||
entry.path() == artifact || entry.path() == temp) {
continue;
}
const uint64_t size = entry.file_size(ec);
if (ec) {
ec.clear();
continue;
}
existing_bytes += size;
candidates.push_back({entry.path(), entry.last_write_time(ec), size});
if (ec) ec.clear();
}
std::sort(candidates.begin(), candidates.end(),
[](const artifact_candidate & a, const artifact_candidate & b) {
return a.modified < b.modified;
});
for (const auto & candidate : candidates) {
if (existing_bytes + artifact_size <= persistent_capacity) break;
fs::remove(candidate.path, ec);
if (!ec) existing_bytes -= std::min(existing_bytes, candidate.size);
ec.clear();
}
if (existing_bytes + artifact_size > persistent_capacity) {
fs::remove(temp, ec);
error = "tokenizer indexes do not fit in --token-cache-size";
return false;
}
fs::rename(temp, artifact, ec);
if (ec) {
const std::string rename_error = ec.message();
// A competing process may have published the same immutable artifact.
std::error_code remove_ec;
fs::remove(temp, remove_ec);
std::string ignored;
if (open_existing(ignored)) {
return true;
}
error = "failed to publish tokenizer index: " + rename_error;
return false;
}
const bool opened = open_existing(error);
if (opened) {
llama_token_cache_set_external_bytes(tokenizer_artifact_bytes(artifact.parent_path()));
}
return opened;
}
int lookup(const std::string & cache_key, const char * sql,
const std::string & first, const std::string * second, int missing) const {
std::lock_guard<std::mutex> lock(mutex);
auto found = values.find(cache_key);
if (found != values.end()) {
lru.splice(lru.begin(), lru, found->second.position);
return found->second.value;
}
sqlite3_stmt * stmt = nullptr;
int value = missing;
if (sqlite3_prepare_v2(db, sql, -1, &stmt, nullptr) == SQLITE_OK) {
sqlite3_bind_blob(stmt, 1, first.data(), int(first.size()), SQLITE_STATIC);
if (second) {
sqlite3_bind_blob(stmt, 2, second->data(), int(second->size()), SQLITE_STATIC);
}
if (sqlite3_step(stmt) == SQLITE_ROW) {
value = sqlite3_column_int(stmt, 0);
}
}
sqlite3_finalize(stmt);
const uint64_t bytes = cache_key.size() + sizeof(cache_entry) + sizeof(std::string);
if (bytes <= capacity) {
while (!lru.empty() && used + bytes > capacity) {
auto victim = values.find(lru.back());
if (victim != values.end()) {
used -= victim->second.bytes;
values.erase(victim);
}
lru.pop_back();
}
lru.push_front(cache_key);
values.emplace(cache_key, cache_entry{value, bytes, lru.begin()});
used += bytes;
}
return value;
}
};
llama_tokenizer_longhaul_index::llama_tokenizer_longhaul_index(
std::string source_key, uint64_t cache_bytes)
: pimpl(std::make_unique<impl>(std::move(source_key), cache_bytes)) {
}
llama_tokenizer_longhaul_index::~llama_tokenizer_longhaul_index() = default;
bool llama_tokenizer_longhaul_index::prepare(
const std::vector<llama_tokenizer_longhaul_reverse> & reverse,
const std::vector<llama_tokenizer_longhaul_merge> & merges,
const std::string & tokenizer_id,
std::string & error) {
const std::string cache_db = llama_token_cache_database_path();
if (cache_db.empty()) {
error = "persistent token cache is disabled or unavailable";
return false;
}
const fs::path directory = fs::path(cache_db).parent_path() / "tokenizer-indexes";
std::error_code ec;
fs::create_directories(directory, ec);
if (ec) {
error = "failed to create tokenizer index directory: " + ec.message();
return false;
}
pimpl->artifact = directory / (llama_token_cache_hash(pimpl->source_key) + ".sqlite3");
if (pimpl->open_existing(error)) {
if (pimpl->id == tokenizer_id) {
llama_token_cache_set_external_bytes(tokenizer_artifact_bytes(directory));
return true;
}
sqlite3_close(pimpl->db);
pimpl->db = nullptr;
std::error_code stale_ec;
fs::remove(pimpl->artifact, stale_ec);
}
lock_directory lock{pimpl->artifact.string() + ".lock", false};
lock.owned = fs::create_directory(lock.path, ec);
if (!lock.owned) {
// Another process is building. Wait for its atomic publication.
for (int i = 0; i < 24000; ++i) {
std::this_thread::sleep_for(std::chrono::milliseconds(25));
std::string ignored;
if (pimpl->open_existing(ignored)) {
if (pimpl->id == tokenizer_id) {
llama_token_cache_set_external_bytes(tokenizer_artifact_bytes(directory));
return true;
}
sqlite3_close(pimpl->db);
pimpl->db = nullptr;
}
if (!fs::exists(lock.path)) {
lock.owned = fs::create_directory(lock.path, ec);
if (lock.owned) break;
}
}
if (!lock.owned) {
error = "timed out waiting for another tokenizer index builder";
return false;
}
}
error.clear();
if (pimpl->open_existing(error)) {
if (pimpl->id == tokenizer_id) {
llama_token_cache_set_external_bytes(tokenizer_artifact_bytes(directory));
return true;
}
sqlite3_close(pimpl->db);
pimpl->db = nullptr;
}
std::error_code remove_ec;
fs::remove(pimpl->artifact, remove_ec);
return pimpl->build(reverse, merges, tokenizer_id, error);
}
llama_token llama_tokenizer_longhaul_index::text_to_token(const std::string & text) const {
const std::string key = "t" + text;
return pimpl->lookup(key,
"SELECT id FROM reverse_tokens WHERE text=?1", text, nullptr, LLAMA_TOKEN_NULL);
}
int llama_tokenizer_longhaul_index::find_bpe_rank(
const std::string & left, const std::string & right) const {
std::string key;
key.reserve(left.size() + right.size() + 2);
key.push_back('m');
key += left;
key.push_back('\0');
key += right;
return pimpl->lookup(key,
"SELECT rank FROM merges WHERE left_text=?1 AND right_text=?2", left, &right, -1);
}
std::vector<std::string> llama_tokenizer_longhaul_index::get_bpe_merges() const {
std::lock_guard<std::mutex> lock(pimpl->mutex);
std::vector<std::string> result;
result.reserve(pimpl->merge_count);
sqlite3_stmt * stmt = nullptr;
if (sqlite3_prepare_v2(pimpl->db,
"SELECT left_text,right_text FROM merges ORDER BY rank", -1, &stmt, nullptr) != SQLITE_OK) {
return result;
}
while (sqlite3_step(stmt) == SQLITE_ROW) {
const char * left = static_cast<const char *>(sqlite3_column_blob(stmt, 0));
const int left_size = sqlite3_column_bytes(stmt, 0);
const char * right = static_cast<const char *>(sqlite3_column_blob(stmt, 1));
const int right_size = sqlite3_column_bytes(stmt, 1);
std::string item(left, left_size);
item.push_back(' ');
item.append(right, right_size);
result.push_back(std::move(item));
}
sqlite3_finalize(stmt);
return result;
}
uint64_t llama_tokenizer_longhaul_index::cache_bytes_used() const {
std::lock_guard<std::mutex> lock(pimpl->mutex);
return pimpl->used;
}
uint64_t llama_tokenizer_longhaul_index::cache_capacity() const {
return pimpl->capacity;
}
uint64_t llama_tokenizer_longhaul_index::n_merges() const {
return pimpl->merge_count;
}
const std::string & llama_tokenizer_longhaul_index::tokenizer_id() const {
return pimpl->id;
}
src/llama-tokenizer-longhaul.h
+47
View File
@@ -0,0 +1,47 @@
#pragma once
#include "llama.h"
#include <cstdint>
#include <memory>
#include <string>
#include <vector>
struct llama_tokenizer_longhaul_reverse {
std::string text;
llama_token id;
};
struct llama_tokenizer_longhaul_merge {
std::string left;
std::string right;
int rank;
};
class llama_tokenizer_longhaul_index {
public:
llama_tokenizer_longhaul_index(std::string source_key, uint64_t cache_bytes);
~llama_tokenizer_longhaul_index();
llama_tokenizer_longhaul_index(const llama_tokenizer_longhaul_index &) = delete;
llama_tokenizer_longhaul_index & operator=(const llama_tokenizer_longhaul_index &) = delete;
bool prepare(
const std::vector<llama_tokenizer_longhaul_reverse> & reverse,
const std::vector<llama_tokenizer_longhaul_merge> & merges,
const std::string & tokenizer_id,
std::string & error);
llama_token text_to_token(const std::string & text) const;
int find_bpe_rank(const std::string & left, const std::string & right) const;
std::vector<std::string> get_bpe_merges() const;
uint64_t cache_bytes_used() const;
uint64_t cache_capacity() const;
uint64_t n_merges() const;
const std::string & tokenizer_id() const;
private:
struct impl;
std::unique_ptr<impl> pimpl;
};
src/llama-vocab.cpp
+193 -16
View File
@@ -1,5 +1,6 @@
#include "llama-vocab.h"
#include "llama-token-cache.h"
#include "llama-tokenizer-longhaul.h"
#include "ggml.h"
#include "gguf.h"
@@ -13,13 +14,16 @@
#include <cctype>
#include <cfloat>
#include <cmath>
#include <condition_variable>
#include <cstdarg>
#include <cstring>
#include <forward_list>
#include <limits>
#include <map>
#include <mutex>
#include <queue>
#include <set>
#include <thread>
#include <unordered_map>
//
@@ -1846,12 +1850,37 @@ struct llama_vocab::impl {
std::vector<char> precompiled_charsmap;
enum class longhaul_state {
disabled,
building,
ready,
failed,
};
bool tokenizer_longhaul = false;
uint64_t tokenizer_longhaul_cache_bytes = 0;
uint64_t tokenizer_longhaul_n_merges = 0;
mutable std::mutex tokenizer_longhaul_mutex;
mutable std::condition_variable tokenizer_longhaul_ready;
mutable longhaul_state tokenizer_longhaul_state = longhaul_state::disabled;
mutable std::string tokenizer_longhaul_error;
mutable std::unique_ptr<llama_tokenizer_longhaul_index> tokenizer_longhaul_index;
std::thread tokenizer_longhaul_worker;
impl(const llama_vocab & vocab) : vocab(vocab) {
}
~impl() = default;
~impl() {
if (tokenizer_longhaul_worker.joinable()) {
tokenizer_longhaul_worker.join();
}
}
void load(llama_model_loader & ml, const LLM_KV & kv);
void load(
llama_model_loader & ml,
const LLM_KV & kv,
bool tokenizer_longhaul,
uint64_t tokenizer_longhaul_cache_bytes);
void wait_for_tokenizer_longhaul() const;
bool load_cache_snapshot(const std::vector<uint8_t> & snapshot);
void reset_cache_snapshot_state();
@@ -1908,7 +1937,7 @@ struct llama_vocab::impl {
bool special) const;
// use cached data
const std::string & token_to_piece(llama_token token) const;
std::string token_to_piece(llama_token token) const;
int32_t detokenize(
const llama_token * tokens,
@@ -1928,11 +1957,29 @@ private:
const llama_vocab & vocab;
};
void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
void llama_vocab::impl::load(
llama_model_loader & ml,
const LLM_KV & kv,
bool tokenizer_longhaul_value,
uint64_t tokenizer_longhaul_cache_bytes_value) {
struct gguf_context * ctx = ml.metadata;
tokenizer_longhaul = tokenizer_longhaul_value;
tokenizer_longhaul_cache_bytes = tokenizer_longhaul_cache_bytes_value;
if (tokenizer_longhaul) {
if (tokenizer_longhaul_cache_bytes == 0) {
throw std::runtime_error("tokenizer-longhaul cache budget must be positive");
}
if (ml.get_arch() != LLM_ARCH_QWEN35MOE && ml.get_arch() != LLM_ARCH_LAGUNA) {
throw std::runtime_error("tokenizer-longhaul currently requires a qwen35moe or laguna model");
}
if (ml.get_token_cache_source_key().empty()) {
throw std::runtime_error("tokenizer-longhaul requires a file-backed GGUF model");
}
}
const std::string & source_cache_key = ml.get_token_cache_source_key();
if (!source_cache_key.empty()) {
if (!tokenizer_longhaul && !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') {
@@ -1958,6 +2005,9 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
ml.get_key(LLM_KV_TOKENIZER_TOKEN_TYPE_COUNT, n_token_types, false);
if (tokenizer_model == "no_vocab" || tokenizer_model == "none") {
if (tokenizer_longhaul) {
throw std::runtime_error("tokenizer-longhaul requires a BPE tokenizer");
}
type = LLAMA_VOCAB_TYPE_NONE;
// default special tokens
@@ -3063,6 +3113,69 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
}
}
if (tokenizer_longhaul) {
if (type != LLAMA_VOCAB_TYPE_BPE ||
(tokenizer_pre != "qwen35" && tokenizer_pre != "laguna")) {
throw std::runtime_error(
"tokenizer-longhaul currently supports only qwen35 and laguna BPE tokenizers");
}
std::vector<llama_tokenizer_longhaul_reverse> reverse;
reverse.reserve(token_to_id.size());
while (!token_to_id.empty()) {
auto node = token_to_id.extract(token_to_id.begin());
reverse.push_back({std::move(node.key()), node.mapped()});
}
token_to_id.rehash(0);
std::vector<llama_tokenizer_longhaul_merge> merges;
merges.reserve(bpe_ranks.size());
while (!bpe_ranks.empty()) {
auto node = bpe_ranks.extract(bpe_ranks.begin());
auto key = std::move(node.key());
merges.push_back({std::move(key.first), std::move(key.second), node.mapped()});
}
bpe_ranks.rehash(0);
tokenizer_longhaul_n_merges = merges.size();
cache_token_to_piece.clear();
cache_token_to_piece.shrink_to_fit();
token_cache_id = llama_token_cache_hash(
std::string("tokenizer-longhaul-v1\n") + source_cache_key);
{
std::lock_guard<std::mutex> lock(tokenizer_longhaul_mutex);
tokenizer_longhaul_state = longhaul_state::building;
}
const uint64_t cache_bytes = tokenizer_longhaul_cache_bytes;
const std::string source_key = source_cache_key;
const std::string tokenizer_id = token_cache_id;
tokenizer_longhaul_worker = std::thread(
[this, source_key, cache_bytes, tokenizer_id,
reverse = std::move(reverse), merges = std::move(merges)]() mutable {
auto index = std::make_unique<llama_tokenizer_longhaul_index>(source_key, cache_bytes);
std::string error;
const bool ok = index->prepare(reverse, merges, tokenizer_id, error);
{
std::lock_guard<std::mutex> lock(tokenizer_longhaul_mutex);
if (ok) {
tokenizer_longhaul_index = std::move(index);
tokenizer_longhaul_state = longhaul_state::ready;
LLAMA_LOG_INFO("%s: tokenizer-longhaul index is ready\n", "llama_vocab");
} else {
tokenizer_longhaul_error = error.empty()
? "unknown tokenizer-longhaul build error"
: std::move(error);
tokenizer_longhaul_state = longhaul_state::failed;
LLAMA_LOG_ERROR("%s: tokenizer-longhaul index failed: %s\n",
"llama_vocab", tokenizer_longhaul_error.c_str());
}
}
tokenizer_longhaul_ready.notify_all();
});
LLAMA_LOG_INFO("%s: tokenizer-longhaul index is building in the background\n", __func__);
return;
}
// Persist a versioned canonical snapshot of the effective tokenizer state.
// The snapshot also serves as the namespace for text-to-token cache entries.
{
@@ -3151,6 +3264,23 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
}
}
void llama_vocab::impl::wait_for_tokenizer_longhaul() const {
if (!tokenizer_longhaul) {
return;
}
std::unique_lock<std::mutex> lock(tokenizer_longhaul_mutex);
if (tokenizer_longhaul_state == longhaul_state::disabled) {
return;
}
tokenizer_longhaul_ready.wait(lock, [this] {
return tokenizer_longhaul_state == longhaul_state::ready ||
tokenizer_longhaul_state == longhaul_state::failed;
});
if (tokenizer_longhaul_state == longhaul_state::failed) {
throw std::runtime_error("tokenizer-longhaul failed: " + tokenizer_longhaul_error);
}
}
void llama_vocab::impl::reset_cache_snapshot_state() {
n_token_types = 0;
tokenizer_model.clear();
@@ -3282,36 +3412,43 @@ std::string llama_vocab::impl::type_name() const{
}
bool llama_vocab::impl::is_normal(llama_token id) const {
wait_for_tokenizer_longhaul();
GGML_ASSERT(type != LLAMA_VOCAB_TYPE_NONE);
return id_to_token[id].attr & LLAMA_TOKEN_ATTR_NORMAL;
}
bool llama_vocab::impl::is_unknown(llama_token id) const {
wait_for_tokenizer_longhaul();
GGML_ASSERT(type != LLAMA_VOCAB_TYPE_NONE);
return id_to_token[id].attr & LLAMA_TOKEN_ATTR_UNKNOWN;
}
bool llama_vocab::impl::is_control(llama_token id) const {
wait_for_tokenizer_longhaul();
GGML_ASSERT(type != LLAMA_VOCAB_TYPE_NONE);
return id_to_token[id].attr & LLAMA_TOKEN_ATTR_CONTROL;
}
bool llama_vocab::impl::is_byte(llama_token id) const {
wait_for_tokenizer_longhaul();
GGML_ASSERT(type != LLAMA_VOCAB_TYPE_NONE);
return id_to_token[id].attr & LLAMA_TOKEN_ATTR_BYTE;
}
bool llama_vocab::impl::is_user_defined(llama_token id) const {
wait_for_tokenizer_longhaul();
GGML_ASSERT(type != LLAMA_VOCAB_TYPE_NONE);
return id_to_token[id].attr & LLAMA_TOKEN_ATTR_USER_DEFINED;
}
bool llama_vocab::impl::is_unused(llama_token id) const {
wait_for_tokenizer_longhaul();
GGML_ASSERT(type != LLAMA_VOCAB_TYPE_NONE);
return id_to_token[id].attr & LLAMA_TOKEN_ATTR_UNUSED;
}
bool llama_vocab::impl::is_eog(llama_token id) const {
wait_for_tokenizer_longhaul();
return id != LLAMA_TOKEN_NULL && special_eog_ids.count(id) > 0;
}
@@ -3339,6 +3476,7 @@ uint8_t llama_vocab::impl::token_to_byte(llama_token id) const {
}
llama_token_attr llama_vocab::impl::token_get_attr(llama_token id) const {
wait_for_tokenizer_longhaul();
GGML_ASSERT(type != LLAMA_VOCAB_TYPE_NONE);
return id_to_token.at(id).attr;
}
@@ -3749,6 +3887,7 @@ std::vector<llama_token> llama_vocab::impl::tokenize(
const std::string & raw_text,
bool add_special,
bool parse_special) const {
wait_for_tokenizer_longhaul();
if (raw_text.empty() || token_cache_id.empty()) {
return tokenize_uncached(raw_text, add_special, parse_special);
}
@@ -3792,6 +3931,7 @@ std::vector<llama_token> llama_vocab::impl::tokenize(
}
int32_t llama_vocab::impl::token_to_piece(llama_token token, char * buf, int32_t length, int32_t lstrip, bool special) const {
wait_for_tokenizer_longhaul();
// 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;
const llama_token_attr attr = token_get_attr(token);
@@ -3908,7 +4048,11 @@ int32_t llama_vocab::impl::token_to_piece(llama_token token, char * buf, int32_t
return 0;
}
const std::string & llama_vocab::impl::token_to_piece(llama_token token) const {
std::string llama_vocab::impl::token_to_piece(llama_token token) const {
wait_for_tokenizer_longhaul();
if (tokenizer_longhaul) {
return token_to_piece_for_cache(token, true);
}
return cache_token_to_piece.at(token);
}
@@ -4027,7 +4171,15 @@ int32_t llama_vocab::impl::detokenize(
void llama_vocab::impl::print_info() const {
LLAMA_LOG_INFO("%s: vocab type = %s\n", __func__, type_name().c_str());
LLAMA_LOG_INFO("%s: n_vocab = %u\n", __func__, vocab.n_tokens());
LLAMA_LOG_INFO("%s: n_merges = %u\n", __func__, (uint32_t) bpe_ranks.size());
LLAMA_LOG_INFO("%s: n_merges = %u\n", __func__,
(uint32_t) (tokenizer_longhaul ? tokenizer_longhaul_n_merges : bpe_ranks.size()));
if (tokenizer_longhaul) {
std::lock_guard<std::mutex> lock(tokenizer_longhaul_mutex);
LLAMA_LOG_INFO("%s: tokenizer longhaul = %s, cache = %.2f MiB\n", __func__,
tokenizer_longhaul_state == longhaul_state::ready ? "ready" :
tokenizer_longhaul_state == longhaul_state::failed ? "failed" : "building",
tokenizer_longhaul_cache_bytes / 1024.0 / 1024.0);
}
// special tokens
if (special_bos_id != LLAMA_TOKEN_NULL) { LLAMA_LOG_INFO( "%s: BOS token = %d '%s'\n", __func__, special_bos_id, id_to_token.at(special_bos_id).text.c_str() ); }
@@ -4060,8 +4212,12 @@ llama_vocab::llama_vocab() : pimpl(new impl(*this)) {
llama_vocab::~llama_vocab() = default;
void llama_vocab::load(llama_model_loader & ml, const LLM_KV & kv) {
pimpl->load(ml, kv);
void llama_vocab::load(
llama_model_loader & ml,
const LLM_KV & kv,
bool tokenizer_longhaul,
uint64_t tokenizer_longhaul_cache_bytes) {
pimpl->load(ml, kv, tokenizer_longhaul, tokenizer_longhaul_cache_bytes);
}
std::string llama_vocab::get_tokenizer_model() const {
@@ -4131,23 +4287,29 @@ llama_token llama_vocab::byte_to_token(uint8_t ch) const {
case LLAMA_VOCAB_TYPE_SPM:
case LLAMA_VOCAB_TYPE_UGM: {
const char buf[7] = { '<', '0', 'x', hex[ch >> 4], hex[ch & 15], '>', 0 };
auto token = pimpl->token_to_id.find(buf);
if (token != pimpl->token_to_id.end()) {
return (*token).second;
const auto token = text_to_token(buf);
if (token != LLAMA_TOKEN_NULL) {
return token;
}
// Try to fall back to just the byte as a string
const char buf2[2] = { (char)ch, 0 };
return pimpl->token_to_id.at(buf2);
const auto fallback = text_to_token(buf2);
if (fallback == LLAMA_TOKEN_NULL) throw std::out_of_range("byte token not found");
return fallback;
}
case LLAMA_VOCAB_TYPE_WPM:
case LLAMA_VOCAB_TYPE_BPE: {
return pimpl->token_to_id.at(unicode_byte_to_utf8(ch));
const auto token = text_to_token(unicode_byte_to_utf8(ch));
if (token == LLAMA_TOKEN_NULL) throw std::out_of_range("byte token not found");
return token;
}
case LLAMA_VOCAB_TYPE_PLAMO2: {
// PLaMo-2 uses byte tokens in format <0xXX>
char hex_str[8];
snprintf(hex_str, sizeof(hex_str), "<0x%02X>", ch);
return pimpl->token_to_id.at(hex_str);
const auto token = text_to_token(hex_str);
if (token == LLAMA_TOKEN_NULL) throw std::out_of_range("byte token not found");
return token;
}
default:
GGML_ABORT("fatal error");
@@ -4156,6 +4318,10 @@ llama_token llama_vocab::byte_to_token(uint8_t ch) const {
llama_token llama_vocab::text_to_token(const std::string & text) const {
GGML_ASSERT(pimpl->type != LLAMA_VOCAB_TYPE_NONE);
pimpl->wait_for_tokenizer_longhaul();
if (pimpl->tokenizer_longhaul_index) {
return pimpl->tokenizer_longhaul_index->text_to_token(text);
}
auto it = pimpl->token_to_id.find(text);
if (it != pimpl->token_to_id.end()) {
return (*it).second;
@@ -4165,16 +4331,19 @@ llama_token llama_vocab::text_to_token(const std::string & text) const {
const llama_vocab::token_data & llama_vocab::get_token_data(llama_token id) const {
GGML_ASSERT(pimpl->type != LLAMA_VOCAB_TYPE_NONE);
pimpl->wait_for_tokenizer_longhaul();
return pimpl->id_to_token.at(id);
}
const char * llama_vocab::token_get_text(llama_token id) const {
GGML_ASSERT(pimpl->type != LLAMA_VOCAB_TYPE_NONE);
pimpl->wait_for_tokenizer_longhaul();
return pimpl->id_to_token.at(id).text.c_str();
}
float llama_vocab::token_get_score(llama_token id) const {
GGML_ASSERT(pimpl->type != LLAMA_VOCAB_TYPE_NONE);
pimpl->wait_for_tokenizer_longhaul();
return pimpl->id_to_token.at(id).score;
}
@@ -4306,6 +4475,10 @@ int llama_vocab::find_bpe_rank(const std::string & token_left, const std::string
GGML_ASSERT(token_left.find(' ') == std::string::npos);
GGML_ASSERT(token_right.find(' ') == std::string::npos);
pimpl->wait_for_tokenizer_longhaul();
if (pimpl->tokenizer_longhaul_index) {
return pimpl->tokenizer_longhaul_index->find_bpe_rank(token_left, token_right);
}
auto it = pimpl->bpe_ranks.find(std::make_pair(token_left, token_right));
if (it == pimpl->bpe_ranks.end()) {
return -1;
@@ -4315,6 +4488,10 @@ int llama_vocab::find_bpe_rank(const std::string & token_left, const std::string
}
std::vector<std::string> llama_vocab::get_bpe_merges() const {
pimpl->wait_for_tokenizer_longhaul();
if (pimpl->tokenizer_longhaul_index) {
return pimpl->tokenizer_longhaul_index->get_bpe_merges();
}
int max_rank = -1;
for (const auto & pair : pimpl->bpe_ranks) {
max_rank = std::max(max_rank, pair.second);
@@ -4364,7 +4541,7 @@ std::vector<llama_token> llama_vocab::tokenize(
return pimpl->tokenize(raw_text, add_special, parse_special);
}
const std::string & llama_vocab::token_to_piece(llama_token token) const {
std::string llama_vocab::token_to_piece(llama_token token) const {
return pimpl->token_to_piece(token);
}
src/llama-vocab.h
+6 -2
View File
@@ -86,7 +86,11 @@ struct llama_vocab {
llama_vocab();
~llama_vocab();
void load(llama_model_loader & ml, const LLM_KV & kv);
void load(
llama_model_loader & ml,
const LLM_KV & kv,
bool tokenizer_longhaul = false,
uint64_t tokenizer_longhaul_cache_bytes = 0);
std::string get_tokenizer_model() const;
std::string get_tokenizer_pre() const;
@@ -181,7 +185,7 @@ struct llama_vocab {
bool special) const;
// use cached data
const std::string & token_to_piece(llama_token token) const;
std::string token_to_piece(llama_token token) const;
int32_t detokenize(
const llama_token * tokens,