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;