Skip to content

Commit bf2c86d

Browse files
ggerganovggerganov
authored
server : refactor prompt cache state ownership (ggml-org#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
1 parent 6e52db5 commit bf2c86d

3 files changed

Lines changed: 75 additions & 73 deletions

File tree

‎tools/server/server-context.cpp‎

Lines changed: 33 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -220,8 +220,6 @@ struct server_slot {
220220
return false;
221221
}
222222

223-
GGML_ASSERT(prompt.data.size() == 0);
224-
225223
const size_t cur_size_tgt = llama_state_seq_get_size_ext(ctx_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE);
226224
const size_t cur_size_dft = ctx_dft ? llama_state_seq_get_size_ext(ctx_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE) : 0;
227225

@@ -252,19 +250,15 @@ struct server_slot {
252250
return res;
253251
}
254252

255-
void prompt_clear(bool allow_processing) {
256-
if (!allow_processing) {
257-
GGML_ASSERT(!is_processing());
258-
}
259-
253+
void prompt_clear() {
260254
SLT_TRC(*this, "clearing prompt with %zu tokens\n", prompt.tokens.size());
261255

262256
common_context_seq_rm(ctx_tgt, id, -1, -1);
263257
if (ctx_dft) {
264258
common_context_seq_rm(ctx_dft, id, -1, -1);
265259
}
266260

267-
prompt.tokens.clear();
261+
prompt.clear();
268262
}
269263

270264
std::vector<common_adapter_lora_info> lora;
@@ -493,7 +487,7 @@ struct server_slot {
493487

494488
// do not keep context of the child slots - the parent's context is enough
495489
if (task->is_child()) {
496-
prompt_clear(false);
490+
prompt_clear();
497491
}
498492

499493
reset();
@@ -1626,7 +1620,7 @@ struct server_context_impl {
16261620
ret->prompt_save(*prompt_cache);
16271621

16281622
if (!ret->prompt_load(*prompt_cache, task.tokens)) {
1629-
ret->prompt_clear(false);
1623+
ret->prompt_clear();
16301624
}
16311625

16321626
prompt_cache->update();
@@ -1658,7 +1652,7 @@ struct server_context_impl {
16581652
if (slot.prompt.n_tokens() > 0) {
16591653
SRV_WRN("purging slot %d with %zu tokens\n", slot.id, slot.prompt.tokens.size());
16601654

1661-
slot.prompt_clear(false);
1655+
slot.prompt_clear();
16621656

16631657
res = true;
16641658

@@ -1691,7 +1685,7 @@ struct server_context_impl {
16911685
// if lora has changed, check to see if the cache should be cleared
16921686
if (lora_should_clear_cache(slot.lora, task_loras)) {
16931687
SLT_TRC(slot, "clearing cache for lora change. %zu loras -> %zu loras\n", slot.lora.size(), task.params.lora.size());
1694-
slot.prompt.tokens.clear();
1688+
slot.prompt.clear();
16951689
} else {
16961690
SLT_TRC(slot, "keeping cache for alora. %zu target loras\n", task_loras.size());
16971691
}
@@ -2405,7 +2399,7 @@ struct server_context_impl {
24052399

24062400
if (params_base.kv_unified) {
24072401
// [TAG_IDLE_SLOT_CLEAR]
2408-
slot.prompt_clear(false);
2402+
slot.prompt_clear();
24092403
}
24102404
}
24112405
}
@@ -2573,12 +2567,12 @@ struct server_context_impl {
25732567
size_t token_count = 0;
25742568
size_t nread = llama_state_seq_load_file(ctx_tgt, filepath.c_str(), slot->id, tokens.data(), tokens.size(), &token_count);
25752569
if (nread == 0) {
2576-
slot->prompt.tokens.clear(); // KV may already been invalidated?
2570+
slot->prompt.clear(); // KV may already been invalidated?
25772571
send_error(task, "Unable to restore slot, no available space in KV cache or invalid slot save file", ERROR_TYPE_INVALID_REQUEST);
25782572
break;
25792573
}
25802574
tokens.resize(token_count);
2581-
slot->prompt.tokens.clear();
2575+
slot->prompt.clear();
25822576
slot->prompt.tokens.insert(tokens);
25832577

25842578
const int64_t t_end = ggml_time_us();
@@ -2615,7 +2609,7 @@ struct server_context_impl {
26152609
// Erase token cache
26162610
const size_t n_erased = slot->prompt.tokens.size();
26172611

2618-
slot->prompt_clear(false);
2612+
slot->prompt_clear();
26192613

26202614
auto res = std::make_unique<server_task_result_slot_erase>();
26212615
res->id = task.id;
@@ -2775,6 +2769,27 @@ struct server_context_impl {
27752769
abort_all_slots("pre_decode() failed: " + std::string(e.what()));
27762770
}
27772771

2772+
GGML_ASSERT(batch.slot_batched || batch.size() == 0);
2773+
2774+
if (batch.slot_batched) {
2775+
auto & slot_batched = batch.slot_batched;
2776+
auto & alora_scale = batch.alora_scale;
2777+
auto & alora_disabled_id = batch.alora_disabled_id;
2778+
2779+
// TODO @ngxson : alora handling is too messy, need to refactor it to be more clear and maintainable
2780+
// apply lora, only need to do it once per batch
2781+
common_set_adapter_lora(ctx_tgt, slot_batched->lora);
2782+
2783+
// if the lora is temporarily disabled for an alora, re-enable it
2784+
// for next time
2785+
if (alora_scale > 0.0f) {
2786+
SRV_DBG("re-enabling alora with scale %f\n", alora_scale);
2787+
slot_batched->lora[alora_disabled_id].scale = alora_scale;
2788+
}
2789+
2790+
llama_set_embeddings(ctx_tgt, slot_batched->need_embd());
2791+
}
2792+
27782793
llama_batch batch_view;
27792794
int32_t off_next = 0;
27802795
int32_t n_batch = llama_n_batch(ctx_tgt);
@@ -2814,7 +2829,6 @@ struct server_context_impl {
28142829
abort_all_slots("post_decode() failed: " + std::string(e.what()));
28152830
break; // stop any further processing
28162831
}
2817-
28182832
}
28192833
}
28202834

@@ -2880,7 +2894,7 @@ struct server_context_impl {
28802894

28812895
new_tokens.resize(slot.prompt.tokens.size() - n_discard);
28822896

2883-
slot.prompt.tokens.clear();
2897+
slot.prompt.clear();
28842898
slot.prompt.tokens.insert(new_tokens);
28852899
}
28862900

@@ -3556,25 +3570,6 @@ struct server_context_impl {
35563570
bool decode(int32_t & n_batch, int32_t off, llama_batch & batch_view) {
35573571
SRV_DBG("n_batch (effective) = %d, off = %d\n", n_batch, off);
35583572

3559-
auto & slot_batched = batch.slot_batched;
3560-
auto & alora_scale = batch.alora_scale;
3561-
auto & alora_disabled_id = batch.alora_disabled_id;
3562-
3563-
// TODO @ngxson : alora handling is too messy, need to refactor it to be more clear and maintainable
3564-
if (slot_batched) {
3565-
// apply lora, only need to do it once per batch
3566-
common_set_adapter_lora(ctx_tgt, slot_batched->lora);
3567-
3568-
// if the lora is temporarily disabled for an alora, re-enable it
3569-
// for next time
3570-
if (alora_scale > 0.0f) {
3571-
SRV_DBG("re-enabling alora with scale %f\n", alora_scale);
3572-
slot_batched->lora[alora_disabled_id].scale = alora_scale;
3573-
}
3574-
3575-
llama_set_embeddings(ctx_tgt, slot_batched->need_embd());
3576-
}
3577-
35783573
if (batch.size() == 0) {
35793574
SRV_WRN("%s", "no tokens to decode\n");
35803575

@@ -3622,7 +3617,7 @@ struct server_context_impl {
36223617

36233618
// note: it's complicated to keep track of how much of the current batch has been
36243619
// processed before the error occurred, so we simply clear the entire context
3625-
slot.prompt_clear(false);
3620+
slot.prompt_clear();
36263621
}
36273622
}
36283623

‎tools/server/server-task.cpp‎

Lines changed: 14 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1646,16 +1646,16 @@ size_t server_prompt_cache::n_tokens() const {
16461646
size_t res = 0;
16471647

16481648
for (const auto & state : states) {
1649-
res += state.n_tokens();
1649+
res += state.prompt.n_tokens();
16501650
}
16511651

16521652
return res;
16531653
}
16541654

1655-
server_prompt * server_prompt_cache::alloc(const server_prompt & prompt, size_t state_size_tgt, size_t state_size_dft) {
1655+
server_prompt_cache_state * server_prompt_cache::alloc(const server_prompt & prompt, size_t state_size_tgt, size_t state_size_dft) {
16561656
// first check if the current state is contained fully in the cache
16571657
for (auto it = states.begin(); it != states.end(); ++it) {
1658-
const int cur_lcp_len = it->tokens.get_common_prefix(prompt.tokens);
1658+
const int cur_lcp_len = it->prompt.tokens.get_common_prefix(prompt.tokens);
16591659

16601660
if (cur_lcp_len == (int) prompt.tokens.size()) {
16611661
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
16801680

16811681
// remove any cached prompts that are fully contained in the current prompt
16821682
for (auto it = states.begin(); it != states.end();) {
1683-
const int len = it->tokens.get_common_prefix(prompt.tokens);
1683+
const int len = it->prompt.tokens.get_common_prefix(prompt.tokens);
16841684

1685-
if (len == (int) it->tokens.size()) {
1685+
if (len == (int) it->prompt.tokens.size()) {
16861686
SRV_TRC(" - removing obsolete cached prompt with length %d\n", len);
16871687

16881688
it = states.erase(it);
@@ -1721,12 +1721,14 @@ server_prompt * server_prompt_cache::alloc(const server_prompt & prompt, size_t
17211721
}
17221722

17231723
states.push_back({
1724-
/*.tokens =*/ prompt.tokens.clone(),
1725-
/*.data =*/ {
1724+
/*.prompt =*/ {
1725+
/*.tokens =*/ prompt.tokens.clone(),
1726+
/*.checkpoints =*/ prompt.checkpoints,
1727+
},
1728+
/*.data =*/ {
17261729
/*.main =*/ std::move(state_data_tgt),
17271730
/*.drft =*/ std::move(state_data_dft),
17281731
},
1729-
/*.checkpoints =*/ prompt.checkpoints,
17301732
});
17311733

17321734
return &states.back();
@@ -1744,9 +1746,9 @@ bool server_prompt_cache::load(server_prompt & prompt, const server_tokens & tok
17441746

17451747
// find the most similar cached prompt, that would also preserve the most context
17461748
for (auto it = states.begin(); it != states.end(); ++it) {
1747-
const int lcp_cur = it->tokens.get_common_prefix(tokens_new);
1749+
const int lcp_cur = it->prompt.tokens.get_common_prefix(tokens_new);
17481750

1749-
const float f_keep_cur = float(lcp_cur) / it->tokens.size();
1751+
const float f_keep_cur = float(lcp_cur) / it->prompt.tokens.size();
17501752
const float sim_cur = float(lcp_cur) / tokens_new.size();
17511753

17521754
// don't trash large prompts
@@ -1799,7 +1801,7 @@ bool server_prompt_cache::load(server_prompt & prompt, const server_tokens & tok
17991801
}
18001802
}
18011803

1802-
prompt = std::move(*it_best);
1804+
prompt = std::move(it_best->prompt);
18031805

18041806
states.erase(it_best);
18051807
}
@@ -1836,6 +1838,6 @@ void server_prompt_cache::update() {
18361838

18371839
for (const auto & state : states) {
18381840
SRV_TRC(" - prompt %p: %7d tokens, checkpoints: %2zu, %9.3f MiB\n",
1839-
(const void *)&state, state.n_tokens(), state.checkpoints.size(), state.size() / (1024.0 * 1024.0));
1841+
(const void *)&state, state.prompt.n_tokens(), state.prompt.checkpoints.size(), state.size() / (1024.0 * 1024.0));
18401842
}
18411843
}

‎tools/server/server-task.h‎

Lines changed: 28 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -584,6 +584,28 @@ struct server_task_result_apply_lora : server_task_result {
584584
virtual json to_json() override;
585585
};
586586

587+
struct server_prompt {
588+
server_tokens tokens;
589+
590+
std::list<common_prompt_checkpoint> checkpoints;
591+
592+
void clear() {
593+
tokens.clear();
594+
checkpoints.clear();
595+
}
596+
597+
int n_tokens() const {
598+
return tokens.size();
599+
}
600+
601+
server_prompt clone() const {
602+
return server_prompt {
603+
tokens.clone(),
604+
checkpoints,
605+
};
606+
}
607+
};
608+
587609
struct server_prompt_data {
588610
std::vector<uint8_t> main;
589611
std::vector<uint8_t> drft;
@@ -593,36 +615,19 @@ struct server_prompt_data {
593615
}
594616
};
595617

596-
struct server_prompt {
597-
server_tokens tokens;
598-
618+
struct server_prompt_cache_state {
619+
server_prompt prompt;
599620
server_prompt_data data;
600621

601-
std::list<common_prompt_checkpoint> checkpoints;
602-
603622
size_t size() const {
604-
size_t res = 0;
605-
606-
res += data.size();
623+
size_t res = data.size();
607624

608-
for (const auto & ckpt : checkpoints) {
625+
for (const auto & ckpt : prompt.checkpoints) {
609626
res += ckpt.size();
610627
}
611628

612629
return res;
613630
}
614-
615-
int n_tokens() const {
616-
return tokens.size();
617-
}
618-
619-
server_prompt clone() const {
620-
return server_prompt {
621-
tokens.clone(),
622-
data,
623-
checkpoints,
624-
};
625-
}
626631
};
627632

628633
struct server_prompt_cache {
@@ -631,7 +636,7 @@ struct server_prompt_cache {
631636
this->limit_tokens = limit_tokens;
632637
}
633638

634-
std::list<server_prompt> states;
639+
std::list<server_prompt_cache_state> states;
635640

636641
// in bytes, 0 = no limit
637642
size_t limit_size = 0;
@@ -643,7 +648,7 @@ struct server_prompt_cache {
643648

644649
size_t n_tokens() const;
645650

646-
server_prompt * alloc(const server_prompt & prompt, size_t state_size_main, size_t state_size_drft);
651+
server_prompt_cache_state * alloc(const server_prompt & prompt, size_t state_size_main, size_t state_size_drft);
647652

648653
bool load(server_prompt & prompt, const server_tokens & tokens_new, llama_context * ctx_main, llama_context * ctx_drft, int32_t id_slot);
649654

0 commit comments

Comments
 (0)