diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp index 775f50baf..54fa949df 100644 --- a/tools/server/server-task.cpp +++ b/tools/server/server-task.cpp @@ -1631,7 +1631,22 @@ server_prompt * server_prompt_cache::alloc(const server_prompt & prompt, size_t } } - // next, remove any cached prompts that are fully contained in the current prompt + // 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->tokens.get_common_prefix(prompt.tokens); @@ -1644,6 +1659,16 @@ server_prompt * server_prompt_cache::alloc(const server_prompt & prompt, size_t } } + 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(); + } + } + std::vector state_data_tgt; std::vector state_data_dft; @@ -1752,12 +1777,7 @@ bool server_prompt_cache::load(server_prompt & prompt, const server_tokens & tok void server_prompt_cache::update() { if (limit_size > 0) { - // always keep at least one state, regardless of the limits - while (states.size() > 1 && size() > limit_size) { - if (states.empty()) { - break; - } - + 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(); @@ -1771,11 +1791,7 @@ void server_prompt_cache::update() { const size_t limit_tokens_cur = limit_size > 0 ? std::max(limit_tokens, limit_size/size_per_token) : limit_tokens; if (limit_tokens > 0) { - while (states.size() > 1 && n_tokens() > limit_tokens_cur) { - if (states.empty()) { - break; - } - + 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));