Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 10 additions & 2 deletions server/src/server/http_server.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -3392,11 +3392,17 @@ HttpServer::GenerationCacheState HttpServer::prepare_generation_cache(
// requests prefer the reusable system/tool boundary; otherwise an
// enabled exact full-prompt cache retains its existing priority.
auto prepare_inline = [&]() {
// Never let the new snapshot land in the slot this request restores
// from: the guard below would cancel it, pinning the restore point at
// the deepest slot on linearly-growing conversations.
const int restore_source_slot =
cache.using_restore ? cache.cache_slot : -1;
const auto prepared_snapshot = prefix_cache_.prepare_inline_snap(
effective_prompt,
cache.using_restore ? logical_prefix_len : 0,
prefer_tools_boundary,
forced_cut);
forced_cut,
restore_source_slot);
Comment thread
cubic-dev-ai[bot] marked this conversation as resolved.
cache.snap_slot = prepared_snapshot.first;
cache.snap_cut = prepared_snapshot.second;
};
Expand Down Expand Up @@ -3660,7 +3666,9 @@ void HttpServer::remember_agent_turn(

const int canonical_end = (int) canonical_tokens.size();
const auto pending = prefix_cache_.prepare_inline_snap(
canonical_tokens, source_pos, false, canonical_end);
canonical_tokens, source_pos, false, canonical_end, source_slot);
// No safe victim (only the restore source and/or protected pins remain)
// or no useful boundary: nothing to replay into.
if (pending.first < 0 || pending.second != canonical_end) return;

const int slot = pending.first;
Expand Down
117 changes: 102 additions & 15 deletions server/src/server/prefix_cache.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
#include "common/sha1.h"

#include <algorithm>
#include <atomic>
#include <cstdio>
#include <cstring>
#include <chrono>
Expand Down Expand Up @@ -119,6 +120,7 @@ std::vector<int> find_all_boundaries(const std::vector<int32_t> & ids,
if (sys_idx < 0) return out;

int cursor = sys_idx + (int)markers.sys_role_prefix.size();
int stray_skips = 0;
while (true) {
auto [end_idx, end_len] = find_first_seq_any(ids, markers.end_msg_seqs, cursor);
if (end_idx < 0) break;
Expand All @@ -136,10 +138,24 @@ std::vector<int> find_all_boundaries(const std::vector<int32_t> & ids,
}
}
found:
if (next_match < 0) break;
if (next_match < 0) {
// Stray end-of-message marker with no following role start — chatml
// tokens embedded in message content (file dumps, terminal output,
// model-echoed markers). Skip this marker and keep scanning instead
// of truncating the whole boundary list. A lone stray previously cut
// the walk off here, hiding every real boundary after it; that pinned
// the inline-snapshot deepen target (second-to-last boundary) at the
// already-restored prefix length, so no snapshot was ever deepened and
// every turn re-prefilled the entire tail. Guard against pathological
// input (or a marker-family mismatch) by capping consecutive strays.
if (++stray_skips > 8192) break;
cursor = after_end;
continue;
}
int boundary = next_match + next_len;
out.push_back(boundary);
cursor = boundary;
stray_skips = 0;
}
return out;
}
Expand Down Expand Up @@ -171,35 +187,63 @@ static bool is_strict_prefix(const std::vector<int32_t> & a,
}

int select_inline_evict_victim(const std::vector<const std::vector<int32_t> *> & ids_lru,
const std::vector<bool> * protected_lru) {
const std::vector<bool> * protected_lru,
int skip_index) {
const int n = (int)ids_lru.size();
if (n <= 0) return 0;
auto is_protected = [&](int i) {
return protected_lru && i >= 0 && i < (int)protected_lru->size() &&
(*protected_lru)[(size_t)i];
};
// Oldest-first scan: prefer an unprotected leaf so sticky tools pins survive.
int oldest_protected_leaf = -1;
for (int i = 0; i < n; i++) {
bool is_ancestor = false;
auto is_ancestor = [&](int i) {
for (int j = 0; j < n; j++) {
if (j == i) continue;
if (is_strict_prefix(*ids_lru[i], *ids_lru[j])) { is_ancestor = true; break; }
if (is_strict_prefix(*ids_lru[i], *ids_lru[j])) return true;
}
if (is_ancestor) continue;
return false;
};
// Oldest-first scan: prefer an unprotected leaf so sticky tools pins
// survive. skip_index (the in-flight restore source) is never a victim.
int oldest_protected_leaf = -1;
for (int i = 0; i < n; i++) {
if (i == skip_index) continue;
if (is_ancestor(i)) continue;
if (!is_protected(i)) return i; // oldest unprotected leaf
if (oldest_protected_leaf < 0) oldest_protected_leaf = i;
}
if (skip_index >= 0) {
// No unprotected leaf outside the restore source — e.g. a linearly
// growing conversation whose only leaf is the restore source itself.
// Evict the shallowest non-protected ancestor (its KV is subsumed by
// every deeper entry) so the new snapshot lands in a different slot
// and the restore point can slide forward. Never the protected tools
// pin, never the restore source.
int shallowest_ancestor = -1;
for (int i = 0; i < n; i++) {
if (i == skip_index || is_protected(i)) continue;
if (!is_ancestor(i)) continue;
if (shallowest_ancestor < 0 ||
ids_lru[i]->size() < ids_lru[(size_t)shallowest_ancestor]->size()) {
shallowest_ancestor = i;
}
}
if (shallowest_ancestor >= 0) return shallowest_ancestor;
// Only the restore source and/or protected pins remain: destroying
// either would throw away the stable tools head or the in-flight
// restore, so there is no safe victim.
return -1;
}
if (oldest_protected_leaf >= 0) return oldest_protected_leaf;
return 0; // unreachable (the longest entry is always a leaf); pure-LRU fallback
}

int select_inline_evict_victim(const std::vector<std::vector<int32_t>> & ids_lru,
const std::vector<bool> * protected_lru) {
const std::vector<bool> * protected_lru,
int skip_index) {
std::vector<const std::vector<int32_t> *> ptrs;
ptrs.reserve(ids_lru.size());
for (const auto & v : ids_lru) ptrs.push_back(&v);
return select_inline_evict_victim(ptrs, protected_lru);
return select_inline_evict_victim(ptrs, protected_lru, skip_index);
}

int select_inline_snapshot_boundary(const std::vector<int> & boundaries,
Expand Down Expand Up @@ -341,7 +385,8 @@ std::pair<int, int> PrefixCache::prepare_inline_snap(
const std::vector<int32_t> & prompt_ids,
int restored_prefix_len,
bool prefer_tools_boundary,
int forced_cut) {
int forced_cut,
int restore_source_slot) {
if (disabled_) return {-1, 0};

auto candidates = find_all_boundaries(prompt_ids, markers_);
Expand All @@ -356,7 +401,22 @@ std::pair<int, int> PrefixCache::prepare_inline_snap(
target_cut = select_inline_snapshot_boundary(
candidates, restored_prefix_len, prefer_tools_boundary);
}
if (target_cut <= 0) return {-1, 0};
if (target_cut <= 0) {
// An expected no-op when the restored prefix already covers the next
// boundary (single-turn prompts, or a cache-primed conversation): log
// once only, so the diagnostic — a truncated boundary list from stray
// chatml in content — stays visible without flooding long runs. The
// HTTP layer may retry within a request, so this can otherwise fire
// twice per turn.
static std::atomic<bool> s_snap_blocked_logged{false};
if (!s_snap_blocked_logged.exchange(true)) {
std::fprintf(stderr,
"[pc] inline snap blocked: boundaries=%zu restored=%d target<=0 "
"(deepen target not past restored prefix; logged once)\n",
candidates.size(), restored_prefix_len);
}
return {-1, 0};
}

auto key = hash_prefix(prompt_ids.data(), target_cut);
if (find_entry(key) >= 0) return {-1, 0}; // already cached
Expand All @@ -372,7 +432,10 @@ std::pair<int, int> PrefixCache::prepare_inline_snap(
if ((int)entries_.size() >= cap_) {
// At capacity — reserve a slot without evicting yet. Prefix-aware: prefer
// the oldest leaf so shared ancestor prefixes (reused by later branches)
// stay resident. Skip protected tools pins when an unprotected leaf exists.
// stay resident. Skip protected tools pins when an unprotected leaf
// exists. The in-flight restore source is never a victim, so the new
// snapshot lands in a different slot and the restore point can slide
// forward past the deepest slot.
std::vector<const std::vector<int32_t> *> ids_lru;
std::vector<bool> protected_lru;
ids_lru.reserve(entries_.size());
Expand All @@ -381,7 +444,23 @@ std::pair<int, int> PrefixCache::prepare_inline_snap(
ids_lru.push_back(&e.ids);
protected_lru.push_back(e.protect);
}
int victim = select_inline_evict_victim(ids_lru, &protected_lru);
int skip_index = -1;
if (restore_source_slot >= 0) {
for (int i = 0; i < (int)entries_.size(); i++) {
if (entries_[i].slot == restore_source_slot) {
skip_index = i;
break;
}
}
}
int victim = select_inline_evict_victim(ids_lru, &protected_lru, skip_index);
if (victim < 0) {
// Nothing safe to evict (only the restore source and/or protected
// pins remain). Skip this snapshot; the restore point stays put
// rather than being destroyed.
pending_protect_ = false;
return {-1, 0};
}
pending_evict_key_ = entries_[victim].hash;
has_pending_evict_ = true;
slot = entries_[victim].slot;
Expand All @@ -393,8 +472,16 @@ std::pair<int, int> PrefixCache::prepare_inline_snap(
entries_[victim].ids.size(), entries_.front().ids.size());
}
} else {
// Skip the in-flight restore source so the new snapshot lands in a
// different slot (the http_server/agent-replay guards would cancel
// an unlucky collision, leaving the restore point pinned). With a
// vacancy and cap >= 2 there is always a non-restore slot to take;
// cap == 1 keeps the old guard behavior.
slot = next_slot_;
next_slot_ = (next_slot_ + 1) % cap_;
if (slot == restore_source_slot && cap_ > 1) {
slot = (slot + 1) % cap_;
}
next_slot_ = (slot + 1) % cap_;
has_pending_evict_ = false;
}

Expand Down
40 changes: 25 additions & 15 deletions server/src/server/prefix_cache.h
Original file line number Diff line number Diff line change
Expand Up @@ -45,23 +45,28 @@ std::vector<int> find_all_boundaries(const std::vector<int32_t> & ids,
using PrefixHash = std::array<uint8_t, 16>;
PrefixHash hash_prefix(const int32_t * ids, int count);

// Prefix-aware inline eviction policy. Given the cached prefixes in LRU order
// (index 0 = oldest), return the index of the eviction victim: the oldest entry
// whose ids are NOT a strict prefix of any other entry's ids (a "leaf"). Keeping
// shared ancestor prefixes resident avoids re-prefilling them for later branches.
// Returns 0 (pure-LRU fallback) when ids_lru is empty or, impossibly, no leaf
// is found. Pure and model-free so it can be unit-tested without a PrefixCache.
// The pointer overload is the core (the caller passes pointers into its own
// entries so no token vectors are copied); the value overload is a convenience
// wrapper for tests.
// Prefix-aware inline eviction: given cached prefixes in LRU order (0 = oldest),
// return the index of the oldest "leaf" — an entry that is not a strict prefix
// of any other — so shared ancestors stay resident. Pointer overload is the
// core (no token copies); the value overload is for tests.
//
// When `protected_lru` is non-null and same-sized, entries with
// `(*protected_lru)[i] == true` are skipped unless every leaf is protected
// (then the oldest protected leaf is chosen as a last resort).
// protected_lru (optional, same size): entries marked true are skipped.
// Without skip_index, if every leaf is protected, the oldest protected leaf
// is the last resort. With skip_index set, protected entries stay ineligible
// and the function may return -1 instead (see skip_index below).
// skip_index (default -1): the in-flight restore source, never a victim; if it
// is the only unprotected leaf, evict the shallowest non-protected ancestor
// instead so the restore point can slide. The protected pin is never evicted.
//
// Returns the victim index [0, n-1]; -1 if skip_index is set and only the
// restore source and/or protected pins remain; 0 if ids_lru is empty or,
// impossibly, no leaf exists.
int select_inline_evict_victim(const std::vector<const std::vector<int32_t> *> & ids_lru,
const std::vector<bool> * protected_lru = nullptr);
const std::vector<bool> * protected_lru = nullptr,
int skip_index = -1);
int select_inline_evict_victim(const std::vector<std::vector<int32_t>> & ids_lru,
const std::vector<bool> * protected_lru = nullptr);
const std::vector<bool> * protected_lru = nullptr,
int skip_index = -1);

// Pick the inline snapshot boundary for a request.
// Default: boundary before the current user turn (second-to-last marker),
Expand Down Expand Up @@ -121,12 +126,17 @@ class PrefixCache {
// `prefer_tools_boundary` selects the system/tools head first (see
// select_inline_snapshot_boundary). When `forced_cut` > restored, that
// cut is used instead (PPP pin_end, including mid-message LCP cuts).
// `restore_source_slot` (default -1) is the slot this request restores
// from; at capacity it is never chosen as the eviction victim, so the new
// snapshot lands in a different slot and the restore point can slide
// forward past the deepest slot.
// Returns (slot, target_cut) or (-1, 0).
std::pair<int, int> prepare_inline_snap(
const std::vector<int32_t> & prompt_ids,
int restored_prefix_len = 0,
bool prefer_tools_boundary = false,
int forced_cut = 0);
int forced_cut = 0,
int restore_source_slot = -1);

// Confirm after daemon successfully saved the snapshot.
// `protect` marks the entry non-evictable by unprotected traffic (tool pin).
Expand Down
Loading