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:
@@ -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();
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user