diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp index 42c7381a9e6f..27d9c5cf76ce 100644 --- a/src/llama-memory-hybrid.cpp +++ b/src/llama-memory-hybrid.cpp @@ -170,7 +170,9 @@ void llama_memory_hybrid::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p } llama_pos llama_memory_hybrid::seq_pos_min(llama_seq_id seq_id) const { - // the min of the total cache is the max of the two caches' min values + // the min of the total cache is the max of the two caches' min values. + // the recurrent state is valid only at its latest position, so the combined min must + // not report positions that the recurrent state cannot serve return std::max(mem_attn->seq_pos_min(seq_id), mem_recr->seq_pos_min(seq_id)); } diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp index 3d6c6db876b4..b1d9c208bb3d 100644 --- a/src/llama-memory-recurrent.cpp +++ b/src/llama-memory-recurrent.cpp @@ -186,6 +186,10 @@ bool llama_memory_recurrent::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos cell.pos = p0 - 1; return true; } + // cannot roll back beyond the available snapshots - the caller has to + // restore a checkpoint or reprocess the sequence + LLAMA_LOG_DEBUG("%s: cannot roll back recurrent state of seq %d by %d tokens (n_rs_seq = %u)\n", + __func__, (int) seq_id, (int) rollback, n_rs_seq); return false; } // invalidate tails which will be cleared diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 7564ad4e9cfb..ed7e6c1c93e7 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -2282,10 +2282,26 @@ struct server_context_impl { return true; } + // whether the memory state is valid only at its exact final position (hybrid/recurrent), + // as opposed to a range of positions (SWA) + bool ctx_tgt_state_exact() const { + return ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_FULL || + ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_RS; + } + // 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; + // an equivalent checkpoint already exists (e.g. it was just restored) + if (!slot.prompt.checkpoints.empty() && + slot.prompt.checkpoints.back().n_tokens == slot.prompt.n_tokens() - n_tokens_cur && + slot.prompt.checkpoints.back().pos_max == pos_max) { + // adopt the checkpoint so the min-step eviction below does not erase it + slot.prompt.checkpoints.back().id_task = id_task; + return; + } + // evict checkpoints within min-step of a previous checkpoint, unless they were // created by the current task int64_t last = -1; @@ -2317,8 +2333,13 @@ struct server_context_impl { 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 + // the state of hybrid/recurrent memory is valid only at its exact final position + // TODO: for SWA models the saved range can still claim more than it actually covers: + // https://github.com/ggml-org/llama.cpp/pull/24411#issuecomment-4677983225 + if (ctx_tgt_state_exact()) { + pos_min = pos_max; + } + cur.update_pos(slot.prompt.n_tokens() - n_tokens_cur, pos_min, pos_max); cur.update_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); @@ -2326,7 +2347,7 @@ struct server_context_impl { // stash the draft's speculative state with the checkpoint common_speculative_get_state(spec.get(), slot.id, cur.data_spec); - SLT_TRC(slot, + SLT_INF(slot, "created context checkpoint %d of %d (pos_min = %d, pos_max = %d, n_tokens = %" PRId64 ", size = %.3f MiB)\n", (int) slot.prompt.checkpoints.size(), params_base.n_ctx_checkpoints, cur.pos_min, cur.pos_max, cur.n_tokens, (float) cur.size() / 1024 / 1024); @@ -3223,6 +3244,10 @@ struct server_context_impl { llama_pos pos_next = slot.prompt.tokens.pos_next(n_past); + // pos_next can be reduced below by a checkpoint restore - remember the + // divergence point for the checkpoint invalidation + const llama_pos pos_next_lcp = pos_next; + // ref: https://github.com/ggml-org/llama.cpp/pull/24110 const bool has_new_tokens = (n_past < slot.task->n_tokens()); @@ -3280,6 +3305,10 @@ struct server_context_impl { } if (pos_min >= pos_min_thold) { + // whether the checkpoints hold a state that is valid only at its exact + // final position (hybrid/recurrent memory), as opposed to a range (SWA) + const bool ckpt_exact = ctx_tgt_state_exact(); + // search for a context checkpoint const auto it = std::find_if( slot.prompt.checkpoints.rbegin(), @@ -3287,6 +3316,13 @@ struct server_context_impl { [&](const auto & cur) { // guarantee that a checkpoint will result in at least one token being processed [TAG_PROMPT_LOGITS] SLT_TRC(slot, "checking checkpoint with [%d, %d] against %d...\n", cur.pos_min, cur.pos_max, pos_min_thold); + if (ckpt_exact) { + // usable only if the tokens up to and including its position are + // a prefix of the new prompt, with at least one token left to + // process [TAG_PROMPT_LOGITS]. the state is self-contained, so + // the SWA slack in pos_min_thold does not apply + return cur.pos_max < pos_next - (has_new_tokens ? 0 : 1); + } // workaround for [TAG_CHECKPOINTS_FIX_POS_MIN] if (cur.pos_max > pos_next) { return false; @@ -3306,11 +3342,11 @@ struct server_context_impl { pos_next = std::min(pos_next, std::max(it->pos_min + 1, it->pos_max)); n_past = std::min(slot.prompt.tokens.size_up_to_pos(pos_next), (size_t) it->n_tokens); - SLT_TRC(slot, "restored context checkpoint (pos_min = %d, pos_max = %d, n_tokens = %" PRId64 ", n_past = %d, size = %.3f MiB)\n", it->pos_min, it->pos_max, it->n_tokens, n_past, (float) it->size() / 1024 / 1024); + SLT_INF(slot, "restored context checkpoint (pos_min = %d, pos_max = %d, n_tokens = %" PRId64 ", n_past = %d, size = %.3f MiB)\n", it->pos_min, it->pos_max, it->n_tokens, n_past, (float) it->size() / 1024 / 1024); } if (do_reset) { - SLT_TRC(slot, "forcing full prompt re-processing due to lack of cache data (likely due to SWA or hybrid/recurrent memory, see %s)\n", + SLT_INF(slot, "forcing full prompt re-processing due to lack of cache data (likely due to SWA or hybrid/recurrent memory, see %s)\n", "https://github.com/ggml-org/llama.cpp/pull/13194#issuecomment-2868343055"); pos_next = 0; n_past = 0; @@ -3319,11 +3355,19 @@ struct server_context_impl { } { - // erase any checkpoints with pos_max > pos_next + // erase any checkpoints that cover diverged content - once the new + // tokens are decoded, their staleness would become undetectable + const llama_pos pos_stale = std::min(pos_next_lcp, slot.task->tokens.pos_next()); + + // an exact checkpoint at the divergence position irreversibly contains + // the diverged token, while a range (SWA) checkpoint gets that entry + // overwritten when decoding resumes from it + const llama_pos pos_stale_min = ctx_tgt_state_exact() ? pos_stale : pos_stale + 1; + for (auto it = slot.prompt.checkpoints.begin(); it != slot.prompt.checkpoints.end();) { const auto & cur = *it; - if (cur.pos_max > pos_next) { - SLT_TRC(slot, "erased invalidated context checkpoint (pos_min = %d, pos_max = %d, n_tokens = %" PRId64 ", n_swa = %d, pos_next = %d, size = %.3f MiB)\n", cur.pos_min, cur.pos_max, cur.n_tokens, n_swa, pos_next, (float) cur.size() / 1024 / 1024); + if (cur.pos_max > pos_next || cur.pos_max >= pos_stale_min) { + SLT_INF(slot, "erased invalidated context checkpoint (pos_min = %d, pos_max = %d, n_tokens = %" PRId64 ", n_swa = %d, pos_next = %d, size = %.3f MiB)\n", cur.pos_min, cur.pos_max, cur.n_tokens, n_swa, pos_next, (float) cur.size() / 1024 / 1024); it = slot.prompt.checkpoints.erase(it); } else { ++it;