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
+33 -38
View File
@@ -220,8 +220,6 @@ struct server_slot {
return false;
}
GGML_ASSERT(prompt.data.size() == 0);
const size_t cur_size_tgt = llama_state_seq_get_size_ext(ctx_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE);
const size_t cur_size_dft = ctx_dft ? llama_state_seq_get_size_ext(ctx_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE) : 0;
@@ -252,11 +250,7 @@ struct server_slot {
return res;
}
void prompt_clear(bool allow_processing) {
if (!allow_processing) {
GGML_ASSERT(!is_processing());
}
void prompt_clear() {
SLT_TRC(*this, "clearing prompt with %zu tokens\n", prompt.tokens.size());
common_context_seq_rm(ctx_tgt, id, -1, -1);
@@ -264,7 +258,7 @@ struct server_slot {
common_context_seq_rm(ctx_dft, id, -1, -1);
}
prompt.tokens.clear();
prompt.clear();
}
std::vector<common_adapter_lora_info> lora;
@@ -493,7 +487,7 @@ struct server_slot {
// do not keep context of the child slots - the parent's context is enough
if (task->is_child()) {
prompt_clear(false);
prompt_clear();
}
reset();
@@ -1626,7 +1620,7 @@ private:
ret->prompt_save(*prompt_cache);
if (!ret->prompt_load(*prompt_cache, task.tokens)) {
ret->prompt_clear(false);
ret->prompt_clear();
}
prompt_cache->update();
@@ -1658,7 +1652,7 @@ private:
if (slot.prompt.n_tokens() > 0) {
SRV_WRN("purging slot %d with %zu tokens\n", slot.id, slot.prompt.tokens.size());
slot.prompt_clear(false);
slot.prompt_clear();
res = true;
@@ -1691,7 +1685,7 @@ private:
// if lora has changed, check to see if the cache should be cleared
if (lora_should_clear_cache(slot.lora, task_loras)) {
SLT_TRC(slot, "clearing cache for lora change. %zu loras -> %zu loras\n", slot.lora.size(), task.params.lora.size());
slot.prompt.tokens.clear();
slot.prompt.clear();
} else {
SLT_TRC(slot, "keeping cache for alora. %zu target loras\n", task_loras.size());
}
@@ -2405,7 +2399,7 @@ private:
if (params_base.kv_unified) {
// [TAG_IDLE_SLOT_CLEAR]
slot.prompt_clear(false);
slot.prompt_clear();
}
}
}
@@ -2573,12 +2567,12 @@ private:
size_t token_count = 0;
size_t nread = llama_state_seq_load_file(ctx_tgt, filepath.c_str(), slot->id, tokens.data(), tokens.size(), &token_count);
if (nread == 0) {
slot->prompt.tokens.clear(); // KV may already been invalidated?
slot->prompt.clear(); // KV may already been invalidated?
send_error(task, "Unable to restore slot, no available space in KV cache or invalid slot save file", ERROR_TYPE_INVALID_REQUEST);
break;
}
tokens.resize(token_count);
slot->prompt.tokens.clear();
slot->prompt.clear();
slot->prompt.tokens.insert(tokens);
const int64_t t_end = ggml_time_us();
@@ -2615,7 +2609,7 @@ private:
// Erase token cache
const size_t n_erased = slot->prompt.tokens.size();
slot->prompt_clear(false);
slot->prompt_clear();
auto res = std::make_unique<server_task_result_slot_erase>();
res->id = task.id;
@@ -2775,6 +2769,27 @@ private:
abort_all_slots("pre_decode() failed: " + std::string(e.what()));
}
GGML_ASSERT(batch.slot_batched || batch.size() == 0);
if (batch.slot_batched) {
auto & slot_batched = batch.slot_batched;
auto & alora_scale = batch.alora_scale;
auto & alora_disabled_id = batch.alora_disabled_id;
// TODO @ngxson : alora handling is too messy, need to refactor it to be more clear and maintainable
// apply lora, only need to do it once per batch
common_set_adapter_lora(ctx_tgt, slot_batched->lora);
// if the lora is temporarily disabled for an alora, re-enable it
// for next time
if (alora_scale > 0.0f) {
SRV_DBG("re-enabling alora with scale %f\n", alora_scale);
slot_batched->lora[alora_disabled_id].scale = alora_scale;
}
llama_set_embeddings(ctx_tgt, slot_batched->need_embd());
}
llama_batch batch_view;
int32_t off_next = 0;
int32_t n_batch = llama_n_batch(ctx_tgt);
@@ -2814,7 +2829,6 @@ private:
abort_all_slots("post_decode() failed: " + std::string(e.what()));
break; // stop any further processing
}
}
}
@@ -2880,7 +2894,7 @@ private:
new_tokens.resize(slot.prompt.tokens.size() - n_discard);
slot.prompt.tokens.clear();
slot.prompt.clear();
slot.prompt.tokens.insert(new_tokens);
}
@@ -3556,25 +3570,6 @@ private:
bool decode(int32_t & n_batch, int32_t off, llama_batch & batch_view) {
SRV_DBG("n_batch (effective) = %d, off = %d\n", n_batch, off);
auto & slot_batched = batch.slot_batched;
auto & alora_scale = batch.alora_scale;
auto & alora_disabled_id = batch.alora_disabled_id;
// TODO @ngxson : alora handling is too messy, need to refactor it to be more clear and maintainable
if (slot_batched) {
// apply lora, only need to do it once per batch
common_set_adapter_lora(ctx_tgt, slot_batched->lora);
// if the lora is temporarily disabled for an alora, re-enable it
// for next time
if (alora_scale > 0.0f) {
SRV_DBG("re-enabling alora with scale %f\n", alora_scale);
slot_batched->lora[alora_disabled_id].scale = alora_scale;
}
llama_set_embeddings(ctx_tgt, slot_batched->need_embd());
}
if (batch.size() == 0) {
SRV_WRN("%s", "no tokens to decode\n");
@@ -3622,7 +3617,7 @@ private:
// note: it's complicated to keep track of how much of the current batch has been
// processed before the error occurred, so we simply clear the entire context
slot.prompt_clear(false);
slot.prompt_clear();
}
}