server : refactor prompt cache state ownership (#25649)

* server : clear checkpoints upon prompt clear

* server : move the prompt state data to the server_prompt_cache

Assisted-by: pi:llama.cpp/Qwen3.6-27B

* server : handle batched slot being cleared
This commit is contained in:
Georgi Gerganov
2026-07-14 18:25:52 +03:00
committed by GitHub
parent 6e52db5b72
commit bf2c86ddc0
3 changed files with 76 additions and 74 deletions
+14 -12
View File
@@ -1646,16 +1646,16 @@ size_t server_prompt_cache::n_tokens() const {
size_t res = 0;
for (const auto & state : states) {
res += state.n_tokens();
res += state.prompt.n_tokens();
}
return res;
}
server_prompt * server_prompt_cache::alloc(const server_prompt & prompt, size_t state_size_tgt, size_t state_size_dft) {
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->tokens.get_common_prefix(prompt.tokens);
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");
@@ -1680,9 +1680,9 @@ server_prompt * server_prompt_cache::alloc(const server_prompt & prompt, size_t
// remove any cached prompts that are fully contained in the current prompt
for (auto it = states.begin(); it != states.end();) {
const int len = it->tokens.get_common_prefix(prompt.tokens);
const int len = it->prompt.tokens.get_common_prefix(prompt.tokens);
if (len == (int) it->tokens.size()) {
if (len == (int) it->prompt.tokens.size()) {
SRV_TRC(" - removing obsolete cached prompt with length %d\n", len);
it = states.erase(it);
@@ -1721,12 +1721,14 @@ server_prompt * server_prompt_cache::alloc(const server_prompt & prompt, size_t
}
states.push_back({
/*.tokens =*/ prompt.tokens.clone(),
/*.data =*/ {
/*.prompt =*/ {
/*.tokens =*/ prompt.tokens.clone(),
/*.checkpoints =*/ prompt.checkpoints,
},
/*.data =*/ {
/*.main =*/ std::move(state_data_tgt),
/*.drft =*/ std::move(state_data_dft),
},
/*.checkpoints =*/ prompt.checkpoints,
});
return &states.back();
@@ -1744,9 +1746,9 @@ bool server_prompt_cache::load(server_prompt & prompt, const server_tokens & tok
// 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->tokens.get_common_prefix(tokens_new);
const int lcp_cur = it->prompt.tokens.get_common_prefix(tokens_new);
const float f_keep_cur = float(lcp_cur) / it->tokens.size();
const float f_keep_cur = float(lcp_cur) / it->prompt.tokens.size();
const float sim_cur = float(lcp_cur) / tokens_new.size();
// don't trash large prompts
@@ -1799,7 +1801,7 @@ bool server_prompt_cache::load(server_prompt & prompt, const server_tokens & tok
}
}
prompt = std::move(*it_best);
prompt = std::move(it_best->prompt);
states.erase(it_best);
}
@@ -1836,6 +1838,6 @@ void server_prompt_cache::update() {
for (const auto & state : states) {
SRV_TRC(" - prompt %p: %7d tokens, checkpoints: %2zu, %9.3f MiB\n",
(const void *)&state, state.n_tokens(), state.checkpoints.size(), state.size() / (1024.0 * 1024.0));
(const void *)&state, state.prompt.n_tokens(), state.prompt.checkpoints.size(), state.size() / (1024.0 * 1024.0));
}
}