diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index cbfbff6b..eae1becd 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -68,6 +68,9 @@ enum slot_state { struct server_slot; // forward declaration +// most tokens reprocessed after a system prefix so its state can be learned +static constexpr int32_t LEARN_REPLAY_MAX = 8192; + // where the first user message starts, -1 without one static int32_t first_user_pos(const common_chat_msg_spans & spans) { for (const auto & span : spans.spans) { @@ -4109,7 +4112,19 @@ private: slot.prompt.scr_depth = 0; } - if (slot.task->params.scr && n_past + 1 < slot.task->n_tokens() && slot.alora_invocation_start <= 0) { + // a prefix to learn lies inside the reused text: replay from it once, if what follows it is short + bool learn_replay = false; + if (slot.learn_at > 0 && n_past > slot.learn_at) { + learn_replay = slot.task->n_tokens() - slot.learn_at <= LEARN_REPLAY_MAX; + if (learn_replay) { + SLT_INF(slot, "cache: reprocessing from %d to store a learned prefix\n", slot.learn_at); + n_past = slot.learn_at; + } else { + slot.learn_at = -1; + } + } + + if (slot.task->params.scr && !learn_replay && n_past + 1 < slot.task->n_tokens() && slot.alora_invocation_start <= 0) { const int32_t n_keep = scr_begin(slot); if (n_keep >= 0) { n_past = n_keep;