server: enforce prompt cache RAM limit (#25070)
Before this commit, --cache-ram was not a hard limit: - The cache always kept at least one entry, even if that entry exceeded the RAM/token limits. - Old entries were only evicted for the RAM/token limits after saving the new one, which could cause the cache to temporarily exceed the RAM/token limits even if individual entries were below the limit. Now, ensure that the RAM limit is strict with these changes: - Skip saving state to cache if by itself it exceeds the RAM limit. - Evict old entries as necessary to make the new entry fit. Additionally, token-limit cleanup may now evict the last remaining cache entry instead of always preserving one.
This commit is contained in:
@@ -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();) {
|
for (auto it = states.begin(); it != states.end();) {
|
||||||
const int len = it->tokens.get_common_prefix(prompt.tokens);
|
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<uint8_t> state_data_tgt;
|
std::vector<uint8_t> state_data_tgt;
|
||||||
std::vector<uint8_t> state_data_dft;
|
std::vector<uint8_t> state_data_dft;
|
||||||
|
|
||||||
@@ -1752,12 +1777,7 @@ bool server_prompt_cache::load(server_prompt & prompt, const server_tokens & tok
|
|||||||
|
|
||||||
void server_prompt_cache::update() {
|
void server_prompt_cache::update() {
|
||||||
if (limit_size > 0) {
|
if (limit_size > 0) {
|
||||||
// always keep at least one state, regardless of the limits
|
while (!states.empty() && size() > limit_size) {
|
||||||
while (states.size() > 1 && size() > limit_size) {
|
|
||||||
if (states.empty()) {
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
|
|
||||||
SRV_WRN(" - cache size limit reached, removing oldest entry (size = %.3f MiB)\n", states.front().size() / (1024.0 * 1024.0));
|
SRV_WRN(" - cache size limit reached, removing oldest entry (size = %.3f MiB)\n", states.front().size() / (1024.0 * 1024.0));
|
||||||
|
|
||||||
states.pop_front();
|
states.pop_front();
|
||||||
@@ -1771,11 +1791,7 @@ void server_prompt_cache::update() {
|
|||||||
const size_t limit_tokens_cur = limit_size > 0 ? std::max<size_t>(limit_tokens, limit_size/size_per_token) : limit_tokens;
|
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) {
|
if (limit_tokens > 0) {
|
||||||
while (states.size() > 1 && n_tokens() > limit_tokens_cur) {
|
while (!states.empty() && n_tokens() > limit_tokens_cur) {
|
||||||
if (states.empty()) {
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
|
|
||||||
SRV_WRN(" - cache token limit (%zu, est: %zu) reached, removing oldest entry (size = %.3f MiB)\n",
|
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));
|
limit_tokens, limit_tokens_cur, states.front().size() / (1024.0 * 1024.0));
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user