diff --git a/common/common.h b/common/common.h index 1535317008..7ed1a98275 100644 --- a/common/common.h +++ b/common/common.h @@ -1081,6 +1081,9 @@ enum ggml_opt_optimizer_type common_opt_get_optimizer(const char *); struct common_prompt_checkpoint { int64_t n_tokens; + // (optional) id of the task that created the checkpoint + int id_task = -1; + llama_pos pos_min; llama_pos pos_max; diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 4826ca5d03..5fadc1c8d9 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -2290,6 +2290,24 @@ private: // n_tokens_cur: the number of tokens added to the batch for the current slot void create_checkpoint(server_slot & slot, const int64_t n_tokens_cur, llama_pos pos_min, llama_pos pos_max) { + const int id_task = slot.task->id; + + // evict checkpoints within min-step of a previous checkpoint, unless they were + // created by the current task + int64_t last = -1; + for (auto it = slot.prompt.checkpoints.begin(); it != slot.prompt.checkpoints.end(); ) { + if (it->id_task != id_task && last >= 0 && it->n_tokens <= last + params_base.checkpoint_min_step) { + SLT_TRC(slot, "erasing context checkpoint too close to an earlier one (pos_min = %d, pos_max = %d, n_tokens = %" PRId64 ", size = %.3f MiB)\n", + it->pos_min, it->pos_max, it->n_tokens, (float) it->size() / 1024 / 1024); + + it = slot.prompt.checkpoints.erase(it); + continue; + } + + last = it->n_tokens; + ++it; + } + while (slot.prompt.checkpoints.size() >= (size_t) params_base.n_ctx_checkpoints) { // make room for the new checkpoint, if needed const auto & cur = slot.prompt.checkpoints.front(); @@ -2302,6 +2320,8 @@ private: auto & cur = slot.prompt.checkpoints.emplace_back(); + cur.id_task = id_task; + // [TAG_CHECKPOINTS_FIX_POS_MIN] // TODO: here we incorrectly deterimne that the saved checkpoint data covers the [pos_min, pos_max] range // this is not true for SWA models: https://github.com/ggml-org/llama.cpp/pull/24411#issuecomment-4677983225 @@ -3511,7 +3531,10 @@ private: do_checkpoint = do_checkpoint && !has_mtmd; // no need to create checkpoints that are too close together, unless it's the last user message - do_checkpoint = do_checkpoint && (slot.prompt.checkpoints.empty() || is_last_user_message || n_tokens_start > slot.prompt.checkpoints.back().n_tokens + params_base.checkpoint_min_step); + do_checkpoint = do_checkpoint && ( + slot.prompt.checkpoints.empty() || + is_last_user_message || near_prompt_end || + n_tokens_start > slot.prompt.checkpoints.back().n_tokens + params_base.checkpoint_min_step); SLT_DBG(slot, "main/do_checkpoint = %s, pos_min = %d, pos_max = %d\n", do_checkpoint ? "yes" : "no", pos_min, pos_max); // note: we create the checkpoint before calling llama_decode(), so the current batch is not