Commit 429460dc0 for llama.cpp
commit 429460dc0e0c52fcf108cd72aadb10e308e4ffa9
Author: Xuan-Son Nguyen <son@huggingface.co>
Date: Sun Nov 15 09:37:17 2026 +0100
server: refactor, add server_slot_sched, prompt running hash (#32154)
* server: refactor, add server_slot_sched, prompt running hash
* use sha256 hash
diff --git a/tests/test-server-component.cpp b/tests/test-server-component.cpp
index c4ed1cacc..3494cdf04 100644
--- a/tests/test-server-component.cpp
+++ b/tests/test-server-component.cpp
@@ -156,6 +156,82 @@ static void test_server_tokens_clone(testing & t) {
t.assert_true("get_text_tokens", cl.get_text_tokens() == llama_tokens({1, 2, 3, 4}));
}
+static void test_server_tokens_prefix_hashes(testing & t) {
+ const size_t n_blk = server_tokens::hash_block_size;
+
+ auto make_tokens = [](size_t n) {
+ llama_tokens res(n);
+ for (size_t i = 0; i < n; i++) {
+ res[i] = 100 + (llama_token) i;
+ }
+ return res;
+ };
+
+ {
+ server_tokens a(make_tokens(3 * n_blk + 10), false);
+ t.assert_equal("full blocks only", (size_t) 3, a.get_prefix_hashes().size());
+ t.assert_equal("same tokens", 3 * n_blk, a.get_common_prefix_lower(a.clone().get_prefix_hashes()));
+ t.assert_equal("no hashes", (size_t) 0, a.get_common_prefix_lower({}));
+
+ // the exact value is 2 * n_blk + 5
+ llama_tokens toks = make_tokens(3 * n_blk + 10);
+ toks[2 * n_blk + 5] = 7;
+ server_tokens b(toks, false);
+ t.assert_equal("diverge in block 2", 2 * n_blk, a.get_common_prefix_lower(b.get_prefix_hashes()));
+ t.assert_equal("exact, diverge in block 2", 2 * n_blk + 5, a.get_common_prefix(b));
+ t.assert_equal("exact, same tokens", 3 * n_blk + 10, a.get_common_prefix(a.clone()));
+ t.assert_equal("b is shorter", 1 * n_blk, a.get_common_prefix_lower(server_tokens(make_tokens(n_blk + 3), false).get_prefix_hashes()));
+
+ toks[0] = 7;
+ t.assert_equal("diverge in block 0", (size_t) 0, a.get_common_prefix_lower(server_tokens(toks, false).get_prefix_hashes()));
+ }
+
+ {
+ server_tokens a(make_tokens(3 * n_blk), false);
+ const auto hashes = a.get_prefix_hashes();
+
+ server_tokens b = a.clone();
+ b.set_token(n_blk + 1, 7);
+ t.assert_equal("set_token", n_blk, b.get_common_prefix_lower(hashes));
+
+ b = a.clone();
+ b.keep_first(n_blk + 1);
+ for (size_t i = n_blk + 1; i < 3 * n_blk; i++) {
+ b.push_back(7);
+ }
+ t.assert_equal("keep_first + push_back", n_blk, b.get_common_prefix_lower(hashes));
+
+ b.clear();
+ t.assert_equal("clear", (size_t) 0, b.get_prefix_hashes().size());
+ }
+
+ {
+ // the image chunk starts in block 0 and ends in a later block
+ mtmd::input_chunks chunks(mtmd_test_create_input_chunks());
+ const mtmd_input_chunk * image = chunks[1];
+ const size_t n_image = mtmd_input_chunk_get_n_tokens(image);
+ const size_t i_image = n_blk - 4;
+
+ auto make_mtmd = [&](bool with_image, llama_token last) {
+ server_tokens res(make_tokens(i_image), true);
+ if (with_image) {
+ res.push_back(image);
+ }
+ while (res.size() < i_image + n_image) {
+ res.push_back(5);
+ }
+ res.push_back(last);
+ return res;
+ };
+
+ server_tokens a = make_mtmd(true, 1);
+ t.assert_true("image is longer than a block", n_image > n_blk);
+ t.assert_equal("exact value", i_image + n_image, a.get_common_prefix(make_mtmd(true, 2)));
+ t.assert_equal("block ends inside the image", i_image, a.get_common_prefix_lower(make_mtmd(true, 2).get_prefix_hashes()));
+ t.assert_equal("text in place of the image", (size_t) 0, a.get_common_prefix_lower(make_mtmd(false, 1).get_prefix_hashes()));
+ }
+}
+
static void test_server_tokens(testing & t) {
t.test("basic", test_server_tokens_basic);
t.test("pos", test_server_tokens_pos);
@@ -163,6 +239,7 @@ static void test_server_tokens(testing & t) {
t.test("common_prefix", test_server_tokens_common_prefix);
t.test("serialize", test_server_tokens_serialize);
t.test("clone", test_server_tokens_clone);
+ t.test("prefix_hashes", test_server_tokens_prefix_hashes);
}
//
diff --git a/tools/server/CMakeLists.txt b/tools/server/CMakeLists.txt
index 4f3fa2458..9f9393f98 100644
--- a/tools/server/CMakeLists.txt
+++ b/tools/server/CMakeLists.txt
@@ -27,6 +27,8 @@ add_library(${TARGET} STATIC
server-mcp.h
server-schema.cpp
server-schema.h
+ server-slot-sched.cpp
+ server-slot-sched.h
)
if (BUILD_SHARED_LIBS)
@@ -36,6 +38,7 @@ endif()
target_include_directories(${TARGET} PRIVATE ../mtmd)
target_include_directories(${TARGET} PRIVATE ${PROJECT_SOURCE_DIR})
target_link_libraries(${TARGET} PUBLIC llama-common mtmd ${CMAKE_THREAD_LIBS_INIT})
+target_link_libraries(${TARGET} PRIVATE vendor::hash)
# llama-server-impl: server logic, reusable by app
diff --git a/tools/server/server-common.cpp b/tools/server/server-common.cpp
index 0013c09cd..ea34f75ac 100644
--- a/tools/server/server-common.cpp
+++ b/tools/server/server-common.cpp
@@ -6,6 +6,7 @@
#include "mtmd-helper.h"
#include "chat.h"
#include "base64.hpp"
+#include "hash/hash.h"
#include "server-common.h"
@@ -525,6 +526,7 @@ server_tokens::server_tokens(mtmd::input_chunks & mtmd_chunks, bool has_mtmd) :
}
server_tokens::server_tokens(const llama_tokens & tokens, bool has_mtmd) : has_mtmd(has_mtmd), tokens(tokens) {
+ update_hashes();
}
llama_pos server_tokens::pos_next(int64_t n_tokens) const {
@@ -641,6 +643,7 @@ void server_tokens::push_back(llama_token tok) {
throw std::runtime_error("Invalid token");
}
tokens.emplace_back(tok);
+ update_hashes();
}
void server_tokens::push_back(const mtmd_input_chunk * chunk) {
@@ -654,6 +657,7 @@ void server_tokens::push_back(const mtmd_input_chunk * chunk) {
}
mtmd::input_chunk_ptr new_chunk(mtmd_input_chunk_copy(chunk));
map_idx_to_media[start_idx] = std::move(new_chunk);
+ update_hashes();
} else if (type == MTMD_INPUT_CHUNK_TYPE_TEXT) {
size_t n_tokens;
const auto * text_tokens = mtmd_input_chunk_get_tokens_text(chunk, &n_tokens);
@@ -677,6 +681,7 @@ void server_tokens::push_back_placeholder(const mtmd_input_chunk * chunk) {
tokens.emplace_back(LLAMA_TOKEN_NULL);
}
map_idx_to_media[start_idx] = std::move(new_chunk);
+ update_hashes();
} else {
push_back(chunk);
}
@@ -701,6 +706,7 @@ void server_tokens::push_back(server_tokens & tokens) {
void server_tokens::insert(const llama_tokens & inp_tokens) {
tokens.insert(tokens.end(), inp_tokens.begin(), inp_tokens.end());
+ update_hashes();
}
const llama_tokens & server_tokens::get_tokens() const {
@@ -768,6 +774,10 @@ server_tokens server_tokens::deserialize(const llama_tokens & packed, bool has_m
result.map_idx_to_media[start_idx] = std::move(chunk);
}
+ // the hashes from the constructor do not include the media chunks
+ result.hashes.clear();
+ result.update_hashes();
+
if (reader.remaining() >= sizeof(llama_token)) {
throw std::runtime_error("Trailing data in server tokens state");
}
@@ -789,6 +799,7 @@ llama_tokens server_tokens::get_text_tokens() const {
void server_tokens::set_token(llama_pos pos, llama_token id) {
GGML_ASSERT(!has_mtmd); // only allow this if mtmd is disabled
tokens[pos] = id;
+ hashes.resize(std::min(hashes.size(), (size_t) pos / hash_block_size));
}
void server_tokens::keep_first(size_t n) {
@@ -822,6 +833,7 @@ void server_tokens::keep_first(size_t n) {
}
}
tokens.resize(n);
+ hashes.resize(std::min(hashes.size(), n / hash_block_size));
}
std::string server_tokens::detokenize(const llama_context * ctx, bool special) const {
@@ -838,8 +850,11 @@ std::string server_tokens::detokenize(const llama_context * ctx, bool special) c
size_t server_tokens::get_common_prefix(const server_tokens & b) const {
const size_t max_idx = std::min(tokens.size(), b.tokens.size());
+ // the blocks with equal hashes are not compared again
+ const size_t min_idx = get_common_prefix_lower(b.get_prefix_hashes());
+
if (!has_mtmd) {
- for (size_t i = 0; i < max_idx; ++i) {
+ for (size_t i = min_idx; i < max_idx; ++i) {
if (tokens[i] == b.tokens[i]) {
continue;
}
@@ -850,7 +865,7 @@ size_t server_tokens::get_common_prefix(const server_tokens & b) const {
return max_idx;
}
- for (size_t i = 0; i < max_idx; ++i) {
+ for (size_t i = min_idx; i < max_idx; ++i) {
const llama_token ai = tokens[i];
const llama_token bi = b.tokens[i];
@@ -885,6 +900,66 @@ size_t server_tokens::get_common_prefix(const server_tokens & b) const {
return max_idx; // all tokens are equal
}
+void server_tokens::update_hashes() {
+ const size_t n_blocks = tokens.size() / hash_block_size;
+
+ std::string buf;
+ const auto append = [&buf](const void * data, size_t len) {
+ buf.append((const char *) data, len);
+ };
+
+ for (size_t k = hashes.size(); k < n_blocks; k++) {
+ const size_t i_begin = k * hash_block_size;
+ const size_t i_end = i_begin + hash_block_size;
+
+ // the previous hash comes first, zeros for the first block
+ const prefix_hash prev = k > 0 ? hashes[k - 1] : prefix_hash{};
+
+ buf.clear();
+ append(prev.data(), prev.size());
+ append(tokens.data() + i_begin, hash_block_size * sizeof(llama_token));
+
+ // a media chunk is hashed in the block of its first token, it can span over many blocks
+ for (auto it = map_idx_to_media.lower_bound(i_begin); it != map_idx_to_media.end() && it->first < i_end; ++it) {
+ const std::string_view id = mtmd_input_chunk_get_id(it->second.get());
+ const uint64_t fields[] = { it->first, mtmd_input_chunk_get_n_tokens(it->second.get()), id.size() };
+ append(fields, sizeof(fields));
+ append(id.data(), id.size());
+ }
+
+ hashes.push_back(hash_sha256(buf.data(), buf.size()));
+ }
+}
+
+size_t server_tokens::get_common_prefix_lower(const std::vector<prefix_hash> & b_hashes) const {
+ const auto & a_hashes = hashes;
+
+ // a hash covers all the tokens before it, so the equal entries come first
+ size_t lo = 0;
+ size_t hi = std::min(a_hashes.size(), b_hashes.size());
+ while (lo < hi) {
+ const size_t mid = lo + (hi - lo) / 2;
+ if (a_hashes[mid] == b_hashes[mid]) {
+ lo = mid + 1;
+ } else {
+ hi = mid;
+ }
+ }
+
+ size_t n = lo * hash_block_size;
+
+ // do not cut a media chunk, go back to its start
+ auto it = map_idx_to_media.upper_bound(n);
+ if (it != map_idx_to_media.begin()) {
+ --it;
+ if (it->first < n && n < it->first + mtmd_input_chunk_get_n_tokens(it->second.get())) {
+ n = it->first;
+ }
+ }
+
+ return n;
+}
+
common_chat_msg_spans server_tokens::find_message_spans(const common_chat_msg_delimiters & delims) const {
std::map<size_t, size_t> skips;
for (const auto & it : map_idx_to_media) {
@@ -933,6 +1008,7 @@ server_tokens server_tokens::clone() const {
server_tokens res;
res.has_mtmd = has_mtmd;
res.tokens = tokens;
+ res.hashes = hashes;
for (auto it = map_idx_to_media.begin(); it != map_idx_to_media.end(); ++it) {
size_t idx = it->first;
const mtmd::input_chunk_ptr & chunk = it->second;
diff --git a/tools/server/server-common.h b/tools/server/server-common.h
index 05622849a..82c424d47 100644
--- a/tools/server/server-common.h
+++ b/tools/server/server-common.h
@@ -10,6 +10,7 @@
#include "json.h"
+#include <array>
#include <atomic>
#include <chrono>
#include <condition_variable>
@@ -167,6 +168,10 @@ std::vector<size_t> lora_get_enabled_ids(const std::vector<common_adapter_lora_i
struct server_tokens {
bool has_mtmd = false;
+ static constexpr size_t hash_block_size = 64;
+
+ using prefix_hash = std::array<uint8_t, 32>; // SHA-256
+
private: // disallow accessing these members directly, risking out-of-sync
// map a **start** index in tokens to the image chunk
@@ -186,6 +191,12 @@ private: // disallow accessing these members directly, risking out-of-sync
// pos 0 1 2 3 4 5 5 5 7 7 7
// map_idx_to_media will contain: {5, img0}, {8, img1}
+ // chained hashes: entry k covers the first (k + 1) * hash_block_size tokens
+ // set_token() drops the entries after the change, the next append brings them back
+ std::vector<prefix_hash> hashes;
+
+ void update_hashes();
+
public:
server_tokens() = default;
~server_tokens() = default;
@@ -254,6 +265,7 @@ public:
void clear() {
map_idx_to_media.clear();
tokens.clear();
+ hashes.clear();
}
void keep_first(size_t n);
@@ -262,6 +274,12 @@ public:
size_t get_common_prefix(const server_tokens & b) const;
+ const std::vector<prefix_hash> & get_prefix_hashes() const { return hashes; }
+
+ // lower bound of get_common_prefix(), b is known only by its prefix hashes
+ // the result is less than one block plus one media chunk below the exact value
+ size_t get_common_prefix_lower(const std::vector<prefix_hash> & b_hashes) const;
+
// split the tokens into message spans, skipping over media chunks
common_chat_msg_spans find_message_spans(const common_chat_msg_delimiters & delims) const;
diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp
index 45693f4e8..4015e23ea 100644
--- a/tools/server/server-context.cpp
+++ b/tools/server/server-context.cpp
@@ -10,6 +10,7 @@
#include "server-queue.h"
#include "server-realtime.h"
#include "server-schema.h"
+#include "server-slot-sched.h"
#include "server-stream.h"
#include "build-info.h"
@@ -1124,8 +1125,7 @@ private:
json json_ui_settings = json::object();
- // Necessary similarity of prompt for slot selection
- float slot_prompt_similarity = 0.0f;
+ server_slot_sched slot_sched;
std::string model_name; // name of the loaded model, to be used by API
std::set<std::string> model_aliases; // additional names for the model
@@ -1441,7 +1441,7 @@ private:
n_swa = params_base.swa_full ? 0 : llama_model_n_swa(model_tgt);
// Necessary similarity of prompt for slot selection
- slot_prompt_similarity = params_base.slot_prompt_similarity;
+ slot_sched.slot_prompt_similarity = params_base.slot_prompt_similarity;
const int n_ctx_train = llama_model_n_ctx_train(model_tgt);
@@ -1800,113 +1800,29 @@ private:
return nullptr;
}
- server_slot * get_available_slot(const server_task & task) {
- server_slot * ret = nullptr;
-
- bool update_cache = false;
-
- // if a specific slot is requested, use it (still goes through cache update logic below)
- if (task.id_slot != -1) {
- ret = get_slot_by_id(task.id_slot);
- if (ret) {
- // a busy slot is returned untouched, the caller defers the task
- if (ret->is_processing()) {
- return ret;
- }
-
- SLT_INF(*ret, "selected slot by id (%d)\n", task.id_slot);
- }
- }
-
- // find the slot that has at least n% prompt similarity
- if (slot_prompt_similarity != 0.0f) {
- float f_sim_best = 0;
-
- for (server_slot & slot : slots) {
- if (task.id_slot != -1 && slot.id != task.id_slot) {
- continue;
- }
-
- // skip the slot if it is not available
- if (slot.is_processing()) {
- SLT_TRC(slot, " - skipping, is_processing = %d\n", slot.is_processing());
- continue;
- }
-
- const auto & tokens = slot.prompt.tokens;
-
- // skip the slot if it does not contains cached tokens
- if (tokens.empty()) {
- SLT_TRC(slot, "%s", " - skipping, slot is empty\n");
- continue;
- }
-
- // fraction of the Longest Common Prefix length with respect to the input prompt length
- const size_t lcp_len = tokens.get_common_prefix(task.tokens);
- const float f_sim_cur = float(lcp_len) / task.tokens.size();
-
- SLT_TRC(slot, " - checking sim = %.3f (%zu/%zu) > %.3f\n", f_sim_cur, lcp_len, task.tokens.size(), slot_prompt_similarity);
-
- // select the current slot if the criteria match
- if (f_sim_cur > f_sim_best && f_sim_cur > slot_prompt_similarity) {
- f_sim_best = f_sim_cur;
-
- ret = &slot;
- }
- }
-
- if (ret != nullptr) {
- const float f_keep = (f_sim_best*task.tokens.size()) / ret->prompt.tokens.size();
-
- if (task.id_slot == -1) {
- SLT_INF(*ret, "selected slot by LCP similarity, f_sim_best = %.3f (> %.3f thold), f_keep = %.3f\n",
- f_sim_best, slot_prompt_similarity, f_keep);
- }
-
- // if we are about to lose a large portion of the existing context - save it in the prompt cache
- if (f_keep < 0.5f) {
- update_cache = true;
- }
- }
+ std::vector<server_slot_info> get_slot_infos() const {
+ std::vector<server_slot_info> res;
+ for (const server_slot & slot : slots) {
+ res.push_back({ slot.id, slot.task.get(), slot.is_processing(), slot.t_last_used, &slot.prompt });
}
+ return res;
+ }
- // find an empty slot, or the one that has been least recently used
- if (ret == nullptr) {
- const auto free_slots = get_free_slots(1, -1);
+ // call only when the task is sure to launch on the slot
+ void update_prompt_cache(server_slot & slot, const server_task & task) {
+ SRV_TRC("%s", "updating prompt cache\n");
- if (!free_slots.empty()) {
- ret = free_slots[0];
+ const int64_t t_start = ggml_time_us();
- SLT_INF(*ret, "selected slot by LRU, n_tokens = %d, t_last = %" PRId64 "\n", ret->prompt.n_tokens(), ret->t_last_used);
+ slot.prompt_save(*prompt_cache);
- update_cache = true;
- }
+ if (!slot.prompt_load(*prompt_cache, task.tokens)) {
+ slot.prompt_clear();
}
- if (ret) {
- update_cache = update_cache && prompt_cache;
-
- // cache prompts only for the tasks that reuse them
- update_cache = update_cache && !task.is_stateless();
-
- if (update_cache) {
- SRV_TRC("%s", "updating prompt cache\n");
+ prompt_cache->update();
- const int64_t t_start = ggml_time_us();
-
- ret->prompt_save(*prompt_cache);
-
- if (!ret->prompt_load(*prompt_cache, task.tokens)) {
- ret->prompt_clear();
- }
-
- prompt_cache->update();
-
- SRV_TRC("prompt cache update took %.2f ms\n", (ggml_time_us() - t_start) / 1000.0);
- }
- }
-
- return ret;
+ SRV_TRC("prompt cache update took %.2f ms\n", (ggml_time_us() - t_start) / 1000.0);
}
// return true if at least one slot has been cleared
@@ -2651,24 +2567,11 @@ private:
return true;
}
- // idle slots that lose the least by being reused: empty slots first, then the least recently used ones
std::vector<server_slot *> get_free_slots(size_t n_slots_needed, int exclude_id_slot) {
std::vector<server_slot *> free_slots;
- for (auto & slot : slots) {
- if (!slot.is_processing() && slot.id != exclude_id_slot) {
- free_slots.push_back(&slot);
- }
+ for (int id : slot_sched.get_free_slots(get_slot_infos(), n_slots_needed, exclude_id_slot)) {
+ free_slots.push_back(get_slot_by_id(id));
}
-
- std::stable_sort(free_slots.begin(), free_slots.end(), [](const server_slot * a, const server_slot * b) {
- if ((a->prompt.n_tokens() == 0) != (b->prompt.n_tokens() == 0)) {
- return a->prompt.n_tokens() == 0;
- }
- return a->t_last_used < b->t_last_used;
- });
-
- free_slots.resize(std::min(free_slots.size(), n_slots_needed));
-
return free_slots;
}
@@ -3040,7 +2943,9 @@ private:
const int id_task = task.id;
- server_slot * slot = get_available_slot(task);
+ const auto plan = slot_sched.select(get_slot_infos(), task);
+
+ server_slot * slot = plan.id_slot == -1 ? nullptr : get_slot_by_id(plan.id_slot);
//
// slot scheduling logic
@@ -3060,15 +2965,24 @@ private:
break;
}
+ std::vector<server_slot *> child_slots;
if (task.is_parent()) {
// try getting free slots for all child tasks
size_t n_child_tasks = task.child_tasks.size();
- std::vector<server_slot *> child_slots = get_free_slots(n_child_tasks, slot->id);
+ child_slots = get_free_slots(n_child_tasks, slot->id);
if (child_slots.size() < n_child_tasks) {
SRV_DBG("not enough free slots for child tasks, n_free = %zu, n_children = %zu, defer task, id_task = %d\n", child_slots.size(), n_child_tasks, id_task);
queue_tasks.defer(std::move(task));
break;
}
+ }
+
+ // the task is not deferred past this point
+ if (plan.update_cache && prompt_cache) {
+ update_prompt_cache(*slot, task);
+ }
+
+ if (task.is_parent()) {
if (!launch_slots_with_parent_task(*slot, child_slots, std::move(task))) {
SRV_ERR("failed to launch slot with parent task, id_task = %d\n", id_task);
break; // drop the task
@@ -4209,13 +4123,6 @@ private:
// ref: https://github.com/ggml-org/llama.cpp/pull/24110
const bool has_new_tokens = (n_past < slot.task->n_tokens());
- // the first evaluated token attends to the positions > pos_eval - n_swa, so a SWA state
- // starting at pos_eval - n_swa + 1 is enough, while a state without SWA must end before pos_eval
- const llama_pos pos_eval = pos_next - (has_new_tokens ? 0 : 1);
-
- // a state or a checkpoint is useful when its pos_min is below this threshold
- const auto pos_min_thold = std::max(0, n_swa > 0 ? pos_eval - n_swa + 2 : pos_eval);
-
if (n_past > 0 && n_past <= slot.prompt.n_tokens()) {
const auto pos_min = llama_memory_seq_pos_min(llama_get_memory(ctx_tgt), slot.id);
if (pos_min == -1) {
@@ -4266,58 +4173,35 @@ private:
SLT_WRN(slot, "%s\n", st1.str().c_str());
}
- if (pos_min >= pos_min_thold) {
- // search for a context checkpoint
- const auto it = std::find_if(
- slot.prompt.checkpoints.rbegin(),
- slot.prompt.checkpoints.rend(),
- [&](const auto & ckpt) {
- const auto & cur = *ckpt;
- // guarantee that a checkpoint will result in at least one token being processed [TAG_PROMPT_LOGITS]
- SLT_TRC(slot, "checking checkpoint with [%d, %d] against %d...\n", cur.pos_min, cur.pos_max, pos_min_thold);
- // workaround for [TAG_CHECKPOINTS_FIX_POS_MIN]
- if (cur.pos_max > pos_next) {
- return false;
- }
- return cur.pos_min < pos_min_thold || cur.pos_min == 0;
- }
- );
-
- bool do_reset = it == slot.prompt.checkpoints.rend();
+ auto reuse = slot.prompt.get_reuse(n_past, has_new_tokens, n_swa, pos_min);
- if (!do_reset) {
- const auto & cur = **it;
+ if (reuse.ckpt) {
+ const auto & cur = *reuse.ckpt;
- // restore the context checkpoint
- if (!cur.load_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) ||
- !cur.load_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY)) {
- if (cur.id_task != -1) {
- GGML_ABORT("failed to restore context checkpoint\n");
- }
- // restored from a slot file, not guaranteed to load - fall back to full prompt re-processing
- SLT_WRN(slot, "%s", "failed to load context checkpoint restored from a slot file\n");
- do_reset = true;
+ // restore the context checkpoint
+ if (!cur.load_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) ||
+ !cur.load_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY)) {
+ if (cur.id_task != -1) {
+ GGML_ABORT("failed to restore context checkpoint\n");
}
- }
-
- if (!do_reset) {
- const auto & cur = **it;
-
+ // restored from a slot file, not guaranteed to load - fall back to full prompt re-processing
+ SLT_WRN(slot, "%s", "failed to load context checkpoint restored from a slot file\n");
+ reuse = {};
+ } else {
// restore the draft's speculative state
common_speculative_set_state(spec.get(), slot.id, cur.data_spec);
- pos_next = std::min(pos_next, std::max(cur.pos_min + 1, cur.pos_max));
- n_past = std::min(slot.prompt.tokens.size_up_to_pos(pos_next), (size_t) cur.n_tokens);
- SLT_TRC(slot, "restored context checkpoint (pos_min = %d, pos_max = %d, n_tokens = %" PRId64 ", n_past = %d, size = %.3f MiB)\n", cur.pos_min, cur.pos_max, cur.n_tokens, n_past, (float) cur.size() / 1024 / 1024);
+ SLT_TRC(slot, "restored context checkpoint (pos_min = %d, pos_max = %d, n_tokens = %" PRId64 ", n_past = %d, size = %.3f MiB)\n", cur.pos_min, cur.pos_max, cur.n_tokens, reuse.n_past, (float) cur.size() / 1024 / 1024);
}
+ }
- if (do_reset) {
- SLT_TRC(slot, "forcing full prompt re-processing due to lack of cache data (likely due to SWA or hybrid/recurrent memory, see %s)\n",
- "https://github.com/ggml-org/llama.cpp/pull/13194#issuecomment-2868343055");
- pos_next = 0;
- n_past = 0;
- }
+ if (reuse.n_past == 0 && !reuse.ckpt) {
+ SLT_TRC(slot, "forcing full prompt re-processing due to lack of cache data (likely due to SWA or hybrid/recurrent memory, see %s)\n",
+ "https://github.com/ggml-org/llama.cpp/pull/13194#issuecomment-2868343055");
}
+
+ pos_next = reuse.pos_next;
+ n_past = reuse.n_past;
}
{
diff --git a/tools/server/server-slot-sched.cpp b/tools/server/server-slot-sched.cpp
new file mode 100644
index 000000000..5b33292c2
--- /dev/null
+++ b/tools/server/server-slot-sched.cpp
@@ -0,0 +1,417 @@
+#include "server-slot-sched.h"
+
+#include "common.h"
+#include "llama.h"
+#include "log.h"
+
+#include <algorithm>
+
+//
+// server_prompt
+//
+
+server_prompt_reuse server_prompt::get_reuse(int32_t n_prefix, bool has_new_tokens, int32_t n_swa, llama_pos pos_min) const {
+ server_prompt_reuse res;
+ res.n_past = n_prefix;
+ res.pos_next = tokens.pos_next(n_prefix);
+
+ // the first evaluated token attends to the positions > pos_eval - n_swa, so a SWA state
+ // starting at pos_eval - n_swa + 1 is enough, while a state without SWA must end before pos_eval
+ const llama_pos pos_eval = res.pos_next - (has_new_tokens ? 0 : 1);
+
+ // a state or a checkpoint is useful when its pos_min is below this threshold
+ const auto pos_min_thold = std::max(0, n_swa > 0 ? pos_eval - n_swa + 2 : pos_eval);
+
+ if (pos_min < pos_min_thold) {
+ return res;
+ }
+
+ // search for a context checkpoint
+ const auto it = std::find_if(
+ checkpoints.rbegin(),
+ checkpoints.rend(),
+ [&](const auto & ckpt) {
+ const auto & cur = *ckpt;
+ // guarantee that a checkpoint will result in at least one token being processed [TAG_PROMPT_LOGITS]
+ SRV_TRC("checking checkpoint with [%d, %d] against %d...\n", cur.pos_min, cur.pos_max, pos_min_thold);
+ // workaround for [TAG_CHECKPOINTS_FIX_POS_MIN]
+ if (cur.pos_max > res.pos_next) {
+ return false;
+ }
+ return cur.pos_min < pos_min_thold || cur.pos_min == 0;
+ }
+ );
+
+ if (it == checkpoints.rend()) {
+ res.n_past = 0;
+ res.pos_next = 0;
+ return res;
+ }
+
+ const auto & cur = **it;
+
+ res.pos_next = std::min(res.pos_next, std::max(cur.pos_min + 1, cur.pos_max));
+ res.n_past = std::min(tokens.size_up_to_pos(res.pos_next), (size_t) cur.n_tokens);
+ res.ckpt = *it;
+
+ return res;
+}
+
+//
+// server_prompt_cache
+//
+
+size_t server_prompt_cache::size() const {
+ size_t res = 0;
+
+ for (const auto & state : states) {
+ res += state.size();
+ }
+
+ return res;
+}
+
+size_t server_prompt_cache::n_tokens() const {
+ size_t res = 0;
+
+ for (const auto & state : states) {
+ res += state.prompt.n_tokens();
+ }
+
+ return res;
+}
+
+server_prompt_cache_state * server_prompt_cache::alloc(const server_prompt & prompt, size_t state_size_tgt, size_t state_size_dft) {
+ // first check if the current state is contained fully in the cache
+ for (auto it = states.begin(); it != states.end(); ++it) {
+ const int cur_lcp_len = it->prompt.tokens.get_common_prefix(prompt.tokens);
+
+ if (cur_lcp_len == (int) prompt.tokens.size()) {
+ SRV_TRC("%s", " - prompt is already in the cache, skipping\n");
+ return nullptr;
+ }
+ }
+
+ // calculate checkpoints size to see if it will fit with the prompt
+ size_t checkpoints_size = 0;
+ for (const auto & ckpt : prompt.checkpoints) {
+ checkpoints_size += ckpt->size();
+ }
+
+ const size_t state_size_new = state_size_tgt + state_size_dft + checkpoints_size;
+
+ // skip over-limit entries to avoid disturbing the cache
+ if (limit_size > 0 && state_size_new > limit_size) {
+ SRV_WRN(" - prompt state size %.3f MiB exceeds cache size limit %.3f MiB, skipping\n",
+ state_size_new / (1024.0 * 1024.0), limit_size / (1024.0 * 1024.0));
+ return nullptr;
+ }
+
+ // remove any cached prompts that are fully contained in the current prompt
+ for (auto it = states.begin(); it != states.end();) {
+ const int len = it->prompt.tokens.get_common_prefix(prompt.tokens);
+
+ if (len == (int) it->prompt.tokens.size()) {
+ SRV_TRC(" - removing obsolete cached prompt with length %d\n", len);
+
+ it = states.erase(it);
+ } else {
+ ++it;
+ }
+ }
+
+ if (limit_size > 0) {
+ // make room before allocating the new vectors to avoid breaching the limit
+ while (!states.empty() && size() + state_size_new > limit_size) {
+ SRV_WRN(" - making room for prompt cache entry, removing oldest entry (size = %.3f MiB)\n",
+ states.front().size() / (1024.0 * 1024.0));
+
+ states.pop_front();
+ }
+ }
+
+ common_state_data state_data_tgt;
+ common_state_data state_data_dft;
+
+ // check if we can allocate enough memory for the new state
+ try {
+ state_data_tgt.resize(state_size_tgt);
+ state_data_dft.resize(state_size_dft);
+ } catch (const std::bad_alloc & e) {
+ SRV_ERR("failed to allocate memory for prompt cache state: %s\n", e.what());
+
+ limit_size = std::max<size_t>(1, 0.4*size());
+
+ SRV_WRN(" - cache size limit reduced to %.3f MiB\n", limit_size / (1024.0 * 1024.0));
+
+ update();
+
+ return nullptr;
+ }
+
+ states.push_back({
+ /*.prompt =*/ {
+ /*.tokens =*/ prompt.tokens.clone(),
+ /*.checkpoints =*/ prompt.checkpoints,
+ },
+ /*.data =*/ {
+ /*.main =*/ std::move(state_data_tgt),
+ /*.drft =*/ std::move(state_data_dft),
+ },
+ });
+
+ return &states.back();
+}
+
+bool server_prompt_cache::load(server_prompt & prompt, const server_tokens & tokens_new, llama_context * ctx_tgt, llama_context * ctx_dft, int32_t id_slot) {
+ int lcp_best = prompt.tokens.get_common_prefix(tokens_new);
+
+ float f_keep_best = prompt.tokens.size() > 0 ? float(lcp_best) / prompt.tokens.size() : -1.0f; // empty slot: any cache entry wins
+ float f_sim_best = float(lcp_best) / tokens_new.size();
+
+ SRV_TRC(" - looking for better prompt, base f_keep = %.3f, f_sim = %.3f\n", f_keep_best, f_sim_best);
+
+ auto it_best = states.end();
+
+ // find the most similar cached prompt, that would also preserve the most context
+ for (auto it = states.begin(); it != states.end(); ++it) {
+ const int lcp_cur = it->prompt.tokens.get_common_prefix(tokens_new);
+
+ const float f_keep_cur = float(lcp_cur) / it->prompt.tokens.size();
+ const float f_sim_cur = float(lcp_cur) / tokens_new.size();
+
+ SRV_TRC(" - prompt with length %7zu, lcp = %7d, f_keep = %.3f, f_sim = %.3f\n", it->prompt.tokens.size(), lcp_cur, f_keep_cur, f_sim_cur);
+
+ // don't trash large prompts
+ if (f_keep_cur < 0.25f) {
+ continue;
+ }
+
+ if (f_keep_best < f_keep_cur && f_sim_best < f_sim_cur) {
+ f_keep_best = f_keep_cur;
+ f_sim_best = f_sim_cur;
+ lcp_best = lcp_cur;
+
+ it_best = it;
+ }
+ }
+
+ if (it_best != states.end()) {
+ SRV_TRC(" - found better prompt with f_keep = %.3f, f_sim = %.3f\n", f_keep_best, f_sim_best);
+
+ {
+ auto & data = it_best->data.main;
+
+ const size_t size = data.size();
+ const size_t n = llama_state_seq_set_data_ext(ctx_tgt, data.data(), size, id_slot, 0);
+ if (n != size) {
+ SRV_ERR("failed to restore state with size %zu\n", size);
+
+ return false;
+ }
+
+ data.clear();
+ data.shrink_to_fit();
+ }
+
+ {
+ auto & data = it_best->data.drft;
+
+ if (!data.empty()) {
+ GGML_ASSERT(ctx_dft);
+
+ const size_t size = data.size();
+ const size_t n = llama_state_seq_set_data_ext(ctx_dft, data.data(), size, id_slot, 0);
+ if (n != size) {
+ SRV_WRN("failed to restore state with size %zu\n", size);
+
+ return false;
+ }
+
+ data.clear();
+ data.shrink_to_fit();
+ }
+ }
+
+ SRV_DBG("%s", "__TEST_TAG_PROMPT_CACHE_RESTORE__\n");
+
+ // the whole cached prompt is a prefix of the new one
+ if (lcp_best == (int) it_best->prompt.tokens.size()) {
+ SRV_DBG("%s", "__TEST_TAG_PROMPT_CACHE_FULL_RESTORE__\n");
+ }
+
+ prompt = std::move(it_best->prompt);
+
+ states.erase(it_best);
+ }
+
+ return true;
+}
+
+void server_prompt_cache::update() {
+ if (limit_size > 0) {
+ while (!states.empty() && size() > limit_size) {
+ SRV_WRN(" - cache size limit reached, removing oldest entry (size = %.3f MiB)\n", states.front().size() / (1024.0 * 1024.0));
+
+ states.pop_front();
+ }
+ }
+
+ // average size per token
+ const float size_per_token = std::max<float>(1.0f, float(size()) / (std::max<size_t>(1, n_tokens())));
+
+ // dynamically increase the token limit if it can fit in the memory limit
+ const size_t limit_tokens_cur = limit_size > 0 ? std::max<size_t>(limit_tokens, limit_size/size_per_token) : limit_tokens;
+
+ if (limit_tokens > 0) {
+ while (!states.empty() && n_tokens() > limit_tokens_cur) {
+ SRV_WRN(" - cache token limit (%zu, est: %zu) reached, removing oldest entry (size = %.3f MiB)\n",
+ limit_tokens, limit_tokens_cur, states.front().size() / (1024.0 * 1024.0));
+
+ states.pop_front();
+ }
+ }
+
+ SRV_TRC(" - cache state: %zu prompts, %.3f MiB (limits: %.3f MiB, %zu tokens, %zu est)\n",
+ states.size(), size() / (1024.0 * 1024.0), limit_size / (1024.0 * 1024.0), limit_tokens, limit_tokens_cur);
+
+ for (const auto & state : states) {
+ SRV_TRC(" - prompt %p: %7d tokens, checkpoints: %2zu, %9.3f MiB\n",
+ (const void *)&state, state.prompt.n_tokens(), state.prompt.checkpoints.size(), state.size() / (1024.0 * 1024.0));
+ }
+}
+
+//
+// server_slot_sched
+//
+
+static std::vector<const server_slot_info *> get_free_slot_infos(const std::vector<server_slot_info> & slots, size_t n_slots_needed, int exclude_id_slot) {
+ std::vector<const server_slot_info *> free_slots;
+ for (const auto & slot : slots) {
+ if (!slot.is_processing && slot.id != exclude_id_slot) {
+ free_slots.push_back(&slot);
+ }
+ }
+
+ std::stable_sort(free_slots.begin(), free_slots.end(), [](const server_slot_info * a, const server_slot_info * b) {
+ if ((a->prompt->n_tokens() == 0) != (b->prompt->n_tokens() == 0)) {
+ return a->prompt->n_tokens() == 0;
+ }
+ return a->t_last_used < b->t_last_used;
+ });
+
+ free_slots.resize(std::min(free_slots.size(), n_slots_needed));
+
+ return free_slots;
+}
+
+std::vector<int> server_slot_sched::get_free_slots(const std::vector<server_slot_info> & slots, size_t n_slots_needed, int exclude_id_slot) const {
+ std::vector<int> res;
+ for (const auto * slot : get_free_slot_infos(slots, n_slots_needed, exclude_id_slot)) {
+ res.push_back(slot->id);
+ }
+
+ return res;
+}
+
+server_slot_plan server_slot_sched::select(const std::vector<server_slot_info> & slots, const server_task & task) const {
+ const server_slot_info * ret = nullptr;
+
+ bool update_cache = false;
+
+ // if a specific slot is requested, use it (still goes through cache update logic below)
+ if (task.id_slot != -1 && !slots.empty()) {
+ // note: allow id_slot to be out of bounds (wrap around)
+ const int id_slot = task.id_slot % slots.size();
+
+ for (const auto & slot : slots) {
+ if (slot.id == id_slot) {
+ ret = &slot;
+ break;
+ }
+ }
+
+ if (ret) {
+ // a busy slot is returned untouched, the caller defers the task
+ if (ret->is_processing) {
+ return { ret->id, false };
+ }
+
+ SLT_INF(*ret, "selected slot by id (%d)\n", task.id_slot);
+ }
+ }
+
+ // find the slot that has at least n% prompt similarity
+ if (slot_prompt_similarity != 0.0f) {
+ float f_sim_best = 0;
+
+ for (const auto & slot : slots) {
+ if (task.id_slot != -1 && slot.id != task.id_slot) {
+ continue;
+ }
+
+ // skip the slot if it is not available
+ if (slot.is_processing) {
+ SLT_TRC(slot, " - skipping, is_processing = %d\n", slot.is_processing);
+ continue;
+ }
+
+ const auto & tokens = slot.prompt->tokens;
+
+ // skip the slot if it does not contains cached tokens
+ if (tokens.empty()) {
+ SLT_TRC(slot, "%s", " - skipping, slot is empty\n");
+ continue;
+ }
+
+ // fraction of the Longest Common Prefix length with respect to the input prompt length
+ const size_t lcp_len = tokens.get_common_prefix(task.tokens);
+ const float f_sim_cur = float(lcp_len) / task.tokens.size();
+
+ SLT_TRC(slot, " - checking sim = %.3f (%zu/%zu) > %.3f\n", f_sim_cur, lcp_len, task.tokens.size(), slot_prompt_similarity);
+
+ // select the current slot if the criteria match
+ if (f_sim_cur > f_sim_best && f_sim_cur > slot_prompt_similarity) {
+ f_sim_best = f_sim_cur;
+
+ ret = &slot;
+ }
+ }
+
+ if (ret != nullptr) {
+ const float f_keep = (f_sim_best*task.tokens.size()) / ret->prompt->tokens.size();
+
+ if (task.id_slot == -1) {
+ SLT_INF(*ret, "selected slot by LCP similarity, f_sim_best = %.3f (> %.3f thold), f_keep = %.3f\n",
+ f_sim_best, slot_prompt_similarity, f_keep);
+ }
+
+ // if we are about to lose a large portion of the existing context - save it in the prompt cache
+ if (f_keep < 0.5f) {
+ update_cache = true;
+ }
+ }
+ }
+
+ // find an empty slot, or the one that has been least recently used
+ if (ret == nullptr) {
+ const auto free_slots = get_free_slot_infos(slots, 1, -1);
+
+ if (!free_slots.empty()) {
+ ret = free_slots[0];
+
+ SLT_INF(*ret, "selected slot by LRU, n_tokens = %d, t_last = %" PRId64 "\n", ret->prompt->n_tokens(), ret->t_last_used);
+
+ update_cache = true;
+ }
+ }
+
+ if (ret == nullptr) {
+ return {};
+ }
+
+ // cache prompts only for the tasks that reuse them
+ update_cache = update_cache && !task.is_stateless();
+
+ return { ret->id, update_cache };
+}
diff --git a/tools/server/server-slot-sched.h b/tools/server/server-slot-sched.h
new file mode 100644
index 000000000..8389d05a5
--- /dev/null
+++ b/tools/server/server-slot-sched.h
@@ -0,0 +1,126 @@
+#pragma once
+
+#include "server-common.h"
+#include "server-task.h"
+
+#include <list>
+#include <memory>
+#include <vector>
+
+// which prompt goes to which slot: slot prompt state, prompt cache and slot selection
+
+// what a cached prompt can give to a new prompt
+struct server_prompt_reuse {
+ int32_t n_past = 0; // tokens to keep
+ llama_pos pos_next = 0;
+
+ // restore this checkpoint first, nullptr when the live state is used as it is
+ std::shared_ptr<const common_prompt_checkpoint> ckpt;
+};
+
+struct server_prompt {
+ server_tokens tokens;
+
+ // a checkpoint does not change after it is created, so the prompt cache shares it with the slot instead of a copy
+ std::list<std::shared_ptr<const common_prompt_checkpoint>> checkpoints;
+
+ void clear() {
+ tokens.clear();
+ checkpoints.clear();
+ }
+
+ int n_tokens() const {
+ return tokens.size();
+ }
+
+ server_prompt clone() const {
+ return server_prompt {
+ tokens.clone(),
+ checkpoints,
+ };
+ }
+
+ // n_prefix: tokens shared with the new prompt, must be > 0
+ // has_new_tokens: the new prompt is longer than n_prefix
+ // pos_min: min position of the live state
+ // n_past == 0 without a checkpoint means that nothing can be reused
+ server_prompt_reuse get_reuse(int32_t n_prefix, bool has_new_tokens, int32_t n_swa, llama_pos pos_min) const;
+};
+
+struct server_prompt_data {
+ common_state_data main;
+ common_state_data drft;
+
+ size_t size() const {
+ return main.size() + drft.size();
+ }
+};
+
+struct server_prompt_cache_state {
+ server_prompt prompt;
+ server_prompt_data data;
+
+ size_t size() const {
+ size_t res = data.size();
+
+ for (const auto & ckpt : prompt.checkpoints) {
+ res += ckpt->size();
+ }
+
+ return res;
+ }
+};
+
+struct server_prompt_cache {
+ server_prompt_cache(int32_t limit_size_mib, size_t limit_tokens) {
+ this->limit_size = 1024ull*1024ull*(limit_size_mib < 0 ? 0 : limit_size_mib);
+ this->limit_tokens = limit_tokens;
+ }
+
+ std::list<server_prompt_cache_state> states;
+
+ // in bytes, 0 = no limit
+ size_t limit_size = 0;
+
+ // in tokens, 0 = no limit
+ size_t limit_tokens = 0;
+
+ size_t size() const;
+
+ size_t n_tokens() const;
+
+ server_prompt_cache_state * alloc(const server_prompt & prompt, size_t state_size_main, size_t state_size_drft);
+
+ bool load(server_prompt & prompt, const server_tokens & tokens_new, llama_context * ctx_tgt, llama_context * ctx_dft, int32_t id_slot);
+
+ void update();
+};
+
+// what the scheduler needs to know about a slot
+struct server_slot_info {
+ int id = -1;
+ const server_task * task = nullptr; // only for logging
+
+ bool is_processing = false;
+ int64_t t_last_used = -1;
+
+ const server_prompt * prompt = nullptr;
+};
+
+struct server_slot_plan {
+ // -1 if no slot is available; a busy slot if the task asks for it by id, the caller defers the task
+ int id_slot = -1;
+
+ // save the slot state to the prompt cache, then load the best match for the task
+ bool update_cache = false;
+};
+
+// decides only, the caller applies the plan
+struct server_slot_sched {
+ float slot_prompt_similarity = 0.0f;
+
+ // idle slots that lose the least by being reused: empty slots first, then the least recently used ones
+ std::vector<int> get_free_slots(const std::vector<server_slot_info> & slots, size_t n_slots_needed, int exclude_id_slot) const;
+
+ server_slot_plan select(const std::vector<server_slot_info> & slots, const server_task & task) const;
+};
diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp
index be3bc539a..9b2ba5ac5 100644
--- a/tools/server/server-task.cpp
+++ b/tools/server/server-task.cpp
@@ -1729,226 +1729,3 @@ json server_task_result_get_lora::to_json() {
json server_task_result_apply_lora::to_json() {
return json {{ "success", true }};
}
-
-//
-// server_prompt_cache
-//
-size_t server_prompt_cache::size() const {
- size_t res = 0;
-
- for (const auto & state : states) {
- res += state.size();
- }
-
- return res;
-}
-
-size_t server_prompt_cache::n_tokens() const {
- size_t res = 0;
-
- for (const auto & state : states) {
- res += state.prompt.n_tokens();
- }
-
- return res;
-}
-
-server_prompt_cache_state * server_prompt_cache::alloc(const server_prompt & prompt, size_t state_size_tgt, size_t state_size_dft) {
- // first check if the current state is contained fully in the cache
- for (auto it = states.begin(); it != states.end(); ++it) {
- const int cur_lcp_len = it->prompt.tokens.get_common_prefix(prompt.tokens);
-
- if (cur_lcp_len == (int) prompt.tokens.size()) {
- SRV_TRC("%s", " - prompt is already in the cache, skipping\n");
- return nullptr;
- }
- }
-
- // calculate checkpoints size to see if it will fit with the prompt
- size_t checkpoints_size = 0;
- for (const auto & ckpt : prompt.checkpoints) {
- checkpoints_size += ckpt->size();
- }
-
- const size_t state_size_new = state_size_tgt + state_size_dft + checkpoints_size;
-
- // skip over-limit entries to avoid disturbing the cache
- if (limit_size > 0 && state_size_new > limit_size) {
- SRV_WRN(" - prompt state size %.3f MiB exceeds cache size limit %.3f MiB, skipping\n",
- state_size_new / (1024.0 * 1024.0), limit_size / (1024.0 * 1024.0));
- return nullptr;
- }
-
- // remove any cached prompts that are fully contained in the current prompt
- for (auto it = states.begin(); it != states.end();) {
- const int len = it->prompt.tokens.get_common_prefix(prompt.tokens);
-
- if (len == (int) it->prompt.tokens.size()) {
- SRV_TRC(" - removing obsolete cached prompt with length %d\n", len);
-
- it = states.erase(it);
- } else {
- ++it;
- }
- }
-
- if (limit_size > 0) {
- // make room before allocating the new vectors to avoid breaching the limit
- while (!states.empty() && size() + state_size_new > limit_size) {
- SRV_WRN(" - making room for prompt cache entry, removing oldest entry (size = %.3f MiB)\n",
- states.front().size() / (1024.0 * 1024.0));
-
- states.pop_front();
- }
- }
-
- common_state_data state_data_tgt;
- common_state_data state_data_dft;
-
- // check if we can allocate enough memory for the new state
- try {
- state_data_tgt.resize(state_size_tgt);
- state_data_dft.resize(state_size_dft);
- } catch (const std::bad_alloc & e) {
- SRV_ERR("failed to allocate memory for prompt cache state: %s\n", e.what());
-
- limit_size = std::max<size_t>(1, 0.4*size());
-
- SRV_WRN(" - cache size limit reduced to %.3f MiB\n", limit_size / (1024.0 * 1024.0));
-
- update();
-
- return nullptr;
- }
-
- states.push_back({
- /*.prompt =*/ {
- /*.tokens =*/ prompt.tokens.clone(),
- /*.checkpoints =*/ prompt.checkpoints,
- },
- /*.data =*/ {
- /*.main =*/ std::move(state_data_tgt),
- /*.drft =*/ std::move(state_data_dft),
- },
- });
-
- return &states.back();
-}
-
-bool server_prompt_cache::load(server_prompt & prompt, const server_tokens & tokens_new, llama_context * ctx_tgt, llama_context * ctx_dft, int32_t id_slot) {
- int lcp_best = prompt.tokens.get_common_prefix(tokens_new);
-
- float f_keep_best = prompt.tokens.size() > 0 ? float(lcp_best) / prompt.tokens.size() : -1.0f; // empty slot: any cache entry wins
- float f_sim_best = float(lcp_best) / tokens_new.size();
-
- SRV_TRC(" - looking for better prompt, base f_keep = %.3f, f_sim = %.3f\n", f_keep_best, f_sim_best);
-
- auto it_best = states.end();
-
- // find the most similar cached prompt, that would also preserve the most context
- for (auto it = states.begin(); it != states.end(); ++it) {
- const int lcp_cur = it->prompt.tokens.get_common_prefix(tokens_new);
-
- const float f_keep_cur = float(lcp_cur) / it->prompt.tokens.size();
- const float f_sim_cur = float(lcp_cur) / tokens_new.size();
-
- SRV_TRC(" - prompt with length %7zu, lcp = %7d, f_keep = %.3f, f_sim = %.3f\n", it->prompt.tokens.size(), lcp_cur, f_keep_cur, f_sim_cur);
-
- // don't trash large prompts
- if (f_keep_cur < 0.25f) {
- continue;
- }
-
- if (f_keep_best < f_keep_cur && f_sim_best < f_sim_cur) {
- f_keep_best = f_keep_cur;
- f_sim_best = f_sim_cur;
- lcp_best = lcp_cur;
-
- it_best = it;
- }
- }
-
- if (it_best != states.end()) {
- SRV_TRC(" - found better prompt with f_keep = %.3f, f_sim = %.3f\n", f_keep_best, f_sim_best);
-
- {
- auto & data = it_best->data.main;
-
- const size_t size = data.size();
- const size_t n = llama_state_seq_set_data_ext(ctx_tgt, data.data(), size, id_slot, 0);
- if (n != size) {
- SRV_ERR("failed to restore state with size %zu\n", size);
-
- return false;
- }
-
- data.clear();
- data.shrink_to_fit();
- }
-
- {
- auto & data = it_best->data.drft;
-
- if (!data.empty()) {
- GGML_ASSERT(ctx_dft);
-
- const size_t size = data.size();
- const size_t n = llama_state_seq_set_data_ext(ctx_dft, data.data(), size, id_slot, 0);
- if (n != size) {
- SRV_WRN("failed to restore state with size %zu\n", size);
-
- return false;
- }
-
- data.clear();
- data.shrink_to_fit();
- }
- }
-
- SRV_DBG("%s", "__TEST_TAG_PROMPT_CACHE_RESTORE__\n");
-
- // the whole cached prompt is a prefix of the new one
- if (lcp_best == (int) it_best->prompt.tokens.size()) {
- SRV_DBG("%s", "__TEST_TAG_PROMPT_CACHE_FULL_RESTORE__\n");
- }
-
- prompt = std::move(it_best->prompt);
-
- states.erase(it_best);
- }
-
- return true;
-}
-
-void server_prompt_cache::update() {
- if (limit_size > 0) {
- while (!states.empty() && size() > limit_size) {
- SRV_WRN(" - cache size limit reached, removing oldest entry (size = %.3f MiB)\n", states.front().size() / (1024.0 * 1024.0));
-
- states.pop_front();
- }
- }
-
- // average size per token
- const float size_per_token = std::max<float>(1.0f, float(size()) / (std::max<size_t>(1, n_tokens())));
-
- // dynamically increase the token limit if it can fit in the memory limit
- const size_t limit_tokens_cur = limit_size > 0 ? std::max<size_t>(limit_tokens, limit_size/size_per_token) : limit_tokens;
-
- if (limit_tokens > 0) {
- while (!states.empty() && n_tokens() > limit_tokens_cur) {
- SRV_WRN(" - cache token limit (%zu, est: %zu) reached, removing oldest entry (size = %.3f MiB)\n",
- limit_tokens, limit_tokens_cur, states.front().size() / (1024.0 * 1024.0));
-
- states.pop_front();
- }
- }
-
- SRV_TRC(" - cache state: %zu prompts, %.3f MiB (limits: %.3f MiB, %zu tokens, %zu est)\n",
- states.size(), size() / (1024.0 * 1024.0), limit_size / (1024.0 * 1024.0), limit_tokens, limit_tokens_cur);
-
- for (const auto & state : states) {
- SRV_TRC(" - prompt %p: %7d tokens, checkpoints: %2zu, %9.3f MiB\n",
- (const void *)&state, state.prompt.n_tokens(), state.prompt.checkpoints.size(), state.size() / (1024.0 * 1024.0));
- }
-}
diff --git a/tools/server/server-task.h b/tools/server/server-task.h
index 6d550d104..aa6b779ed 100644
--- a/tools/server/server-task.h
+++ b/tools/server/server-task.h
@@ -685,78 +685,6 @@ struct server_task_result_apply_lora : server_task_result {
virtual json to_json() override;
};
-struct server_prompt {
- server_tokens tokens;
-
- // a checkpoint does not change after it is created, so the prompt cache shares it with the slot instead of a copy
- std::list<std::shared_ptr<const common_prompt_checkpoint>> checkpoints;
-
- void clear() {
- tokens.clear();
- checkpoints.clear();
- }
-
- int n_tokens() const {
- return tokens.size();
- }
-
- server_prompt clone() const {
- return server_prompt {
- tokens.clone(),
- checkpoints,
- };
- }
-};
-
-struct server_prompt_data {
- common_state_data main;
- common_state_data drft;
-
- size_t size() const {
- return main.size() + drft.size();
- }
-};
-
-struct server_prompt_cache_state {
- server_prompt prompt;
- server_prompt_data data;
-
- size_t size() const {
- size_t res = data.size();
-
- for (const auto & ckpt : prompt.checkpoints) {
- res += ckpt->size();
- }
-
- return res;
- }
-};
-
-struct server_prompt_cache {
- server_prompt_cache(int32_t limit_size_mib, size_t limit_tokens) {
- this->limit_size = 1024ull*1024ull*(limit_size_mib < 0 ? 0 : limit_size_mib);
- this->limit_tokens = limit_tokens;
- }
-
- std::list<server_prompt_cache_state> states;
-
- // in bytes, 0 = no limit
- size_t limit_size = 0;
-
- // in tokens, 0 = no limit
- size_t limit_tokens = 0;
-
- size_t size() const;
-
- size_t n_tokens() const;
-
- server_prompt_cache_state * alloc(const server_prompt & prompt, size_t state_size_main, size_t state_size_drft);
-
- bool load(server_prompt & prompt, const server_tokens & tokens_new, llama_context * ctx_tgt, llama_context * ctx_dft, int32_t id_slot);
-
- void update();
-};
-
// used exclusively by router mode
struct server_task_result_router : server_task_result {
json data;
diff --git a/vendor/hash/hash.cpp b/vendor/hash/hash.cpp
index 6493716bf..227f867ca 100644
--- a/vendor/hash/hash.cpp
+++ b/vendor/hash/hash.cpp
@@ -16,8 +16,13 @@ static std::string to_hex(const unsigned char * digest, size_t len) {
return out;
}
+std::array<uint8_t, 32> hash_sha256(const void * data, size_t len) {
+ std::array<uint8_t, SHA256_DIGEST_SIZE> digest;
+ sha256_hash(digest.data(), (const unsigned char *) data, len);
+ return digest;
+}
+
std::string hash_sha256_hex(const void * data, size_t len) {
- unsigned char digest[SHA256_DIGEST_SIZE];
- sha256_hash(digest, (const unsigned char *) data, len);
- return to_hex(digest, SHA256_DIGEST_SIZE);
+ const auto digest = hash_sha256(data, len);
+ return to_hex(digest.data(), digest.size());
}
diff --git a/vendor/hash/hash.h b/vendor/hash/hash.h
index 1298b4196..c450196d7 100644
--- a/vendor/hash/hash.h
+++ b/vendor/hash/hash.h
@@ -2,8 +2,13 @@
// C++ wrapper for the vendored hash functions
+#include <array>
#include <cstddef>
+#include <cstdint>
#include <string>
+// returns the raw SHA-256 digest
+std::array<uint8_t, 32> hash_sha256(const void * data, size_t len);
+
// returns the SHA-256 digest as a lowercase hex string
std::string hash_sha256_hex(const void * data, size_t len);