Skip to content

fix(prefix-cache): reuse slots after aborted snapshots #563

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
dusterbloom wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(prefix-cache): reuse slots after aborted snapshots #563

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
5 changes: 4 additions & 1 deletion .github/workflows/ci.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -79,12 +79,15 @@ jobs:
-DCMAKE_BUILD_TYPE=Release
cmake --build build --target \
test_dflash test_generate test_flash_attn_sparse test_server_unit \
test_prefix_cache_state \
test_deepseek4_unit -j$(nproc)

- name: Run C++ server unit tests
run: |
cd server/build
ctest --output-on-failure -R "server_unit|deepseek4_unit" --no-tests=error
ctest --output-on-failure \
-R "server_unit|prefix_cache_state|deepseek4_unit" \
--no-tests=error

- name: Populate venv with cu128 torch + setuptools
# First pass: install the workspace's default deps. dflash declares
Expand Down
15 changes: 15 additions & 0 deletions server/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1163,6 +1163,18 @@ if(DFLASH27B_TESTS)
list(APPEND _raw_unit_test_targets test_server_unit)
endif()

# Dependency-free production state core used by native tests and ESBMC.
# Keep this target free of tokenizer, ggml, and GPU libraries so formal
# transition checks remain fast and reproducible on hosted CPU runners.
if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_prefix_cache_state.cpp")
add_executable(
test_prefix_cache_state
test/test_prefix_cache_state.cpp)
target_include_directories(test_prefix_cache_state PRIVATE
${CMAKE_CURRENT_SOURCE_DIR}/src)
add_test(NAME prefix_cache_state COMMAND test_prefix_cache_state)
endif()

# Feature/architecture gate tests. check_feature_compatibility(),
# collect_feature_warnings() and the capability table are pure functions,
# so this target deliberately compiles only feature_gate.cpp and
Expand Down Expand Up @@ -1236,6 +1248,9 @@ if(DFLASH27B_TESTS)
if(TARGET test_feature_gate)
list(APPEND _check_deps test_feature_gate)
endif()
if(TARGET test_prefix_cache_state)
list(APPEND _check_deps test_prefix_cache_state)
endif()
if(_check_deps)
add_custom_target(check
COMMAND ${CMAKE_CTEST_COMMAND} --output-on-failure
Expand Down
184 changes: 52 additions & 132 deletions server/src/server/prefix_cache.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -140,38 +140,6 @@ PrefixHash hash_prefix(const int32_t * ids, int count) {
return h;
}

// ─── Prefix-aware eviction ──────────────────────────────────────────────

static bool is_strict_prefix(const std::vector<int32_t> & a,
const std::vector<int32_t> & b) {
// True iff `a` is a strict (shorter) prefix of `b`.
if (a.size() >= b.size()) return false;
return std::equal(a.begin(), a.end(), b.begin());
}

int select_inline_evict_victim(const std::vector<const std::vector<int32_t> *> & ids_lru) {
const int n = (int)ids_lru.size();
if (n <= 0) return 0;
// Oldest-first scan: evict the first entry that is not a strict prefix of any
// other entry (a leaf). Shared ancestor prefixes are thereby kept resident.
for (int i = 0; i < n; i++) {
bool is_ancestor = false;
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_ancestor) return i; // oldest 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) {
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);
}

int select_inline_snapshot_boundary(const std::vector<int> & boundaries,
int restored_prefix_len) {
if (boundaries.empty()) return 0;
Expand All @@ -184,7 +152,8 @@ int select_inline_snapshot_boundary(const std::vector<int> & boundaries,
// ─── PrefixCache ────────────────────────────────────────────────────────

PrefixCache::PrefixCache(int cap, const Tokenizer & tokenizer)
: cap_(std::min(cap, MAX_SLOTS))
: cap_(std::min(cap, MAX_SLOTS)),
inline_state_(std::max(0, std::min(cap, MAX_SLOTS)))
{
if (cap_ <= 0) {
disabled_ = true;
Expand All @@ -203,18 +172,9 @@ PrefixCache::PrefixCache(int cap, const Tokenizer & tokenizer)

// ── LRU helpers ─────────────────────────────────────────────────────────

int PrefixCache::find_entry(const PrefixHash & h) const {
for (int i = 0; i < (int)entries_.size(); i++) {
if (entries_[i].hash == h) return i;
}
return -1;
}

void PrefixCache::move_to_end(int idx) {
if (idx < 0 || idx >= (int)entries_.size()) return;
auto e = std::move(entries_[idx]);
entries_.erase(entries_.begin() + idx);
entries_.push_back(std::move(e));
void PrefixCache::sync_inline_size() {
entries_size_count_.store(
inline_state_.size(), std::memory_order_relaxed);
}

int PrefixCache::find_full_entry(const PrefixHash & h) const {
Expand All @@ -241,26 +201,21 @@ std::pair<int, int> PrefixCache::lookup(const std::vector<int32_t> & prompt_ids)

for (int cut : boundaries) {
auto key = hash_prefix(prompt_ids.data(), cut);
int idx = find_entry(key);
if (idx >= 0) {
const int committed = (int)entries_[idx].ids.size();
if (committed != cut) {
// Slot was refreshed in-place at a deeper boundary; a shallow
// hash→slot entry would restore the wrong cur_pos.
std::fprintf(stderr,
"[pc] lookup stale slot=%d key_cut=%d committed=%d — evicting\n",
entries_[idx].slot, cut, committed);
entries_.erase(entries_.begin() + idx);
entries_size_count_.fetch_sub(1, std::memory_order_relaxed);
continue;
}
if (cut > best_len) {
best_slot = entries_[idx].slot;
best_len = cut;
}
move_to_end(idx);
const auto result = inline_state_.lookup_candidate(key, cut);
if (result.stale_removed) {
// Slot was refreshed in-place at a deeper boundary; a shallow
// hash→slot entry would restore the wrong cur_pos.
std::fprintf(stderr,
"[pc] lookup stale slot=%d key_cut=%d committed=%d — evicting\n",
result.stale_slot, cut, result.stale_committed_len);
continue;
}
if (result.slot >= 0 && cut > best_len) {
best_slot = result.slot;
best_len = result.prefix_len;
}
}
sync_inline_size();

if (best_slot >= 0) {
lifetime_hits_.fetch_add(1, std::memory_order_relaxed);
Expand All @@ -281,68 +236,44 @@ std::pair<int, int> PrefixCache::prepare_inline_snap(
if (target_cut <= 0) return {-1, 0};

auto key = hash_prefix(prompt_ids.data(), target_cut);
if (find_entry(key) >= 0) return {-1, 0}; // already cached

int slot;
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. entries_ is already in LRU order (front = oldest).
std::vector<const std::vector<int32_t> *> ids_lru;
ids_lru.reserve(entries_.size());
for (const auto & e : entries_) ids_lru.push_back(&e.ids);
int victim = select_inline_evict_victim(ids_lru);
pending_evict_key_ = entries_[victim].hash;
has_pending_evict_ = true;
slot = entries_[victim].slot;
if (victim != 0) {
std::fprintf(stderr,
"[pc] prefix-aware evict: victim idx=%d (len=%zu) kept oldest "
"ancestor (len=%zu)\n",
victim, entries_[victim].ids.size(), entries_.front().ids.size());
}
} else {
slot = next_slot_;
next_slot_ = (next_slot_ + 1) % cap_;
has_pending_evict_ = false;
const auto reservation = inline_state_.prepare(key, target_cut);
if (reservation.slot < 0) return {-1, 0};
if (reservation.victim_index > 0) {
std::fprintf(stderr,
"[pc] prefix-aware evict: victim idx=%d (len=%d) kept oldest "
"ancestor (len=%d)\n",
reservation.victim_index, reservation.victim_len,
reservation.oldest_len);
}

return {slot, target_cut};
return {reservation.slot, reservation.target_cut};
}

void PrefixCache::confirm_inline_snap(int slot, int target_cut,
const std::vector<int32_t> & prompt_ids) {
if (disabled_) return;

// Evict the reserved entry (if any).
if (has_pending_evict_) {
int idx = find_entry(pending_evict_key_);
if (idx >= 0) {
entries_.erase(entries_.begin() + idx);
entries_size_count_.fetch_sub(1, std::memory_order_relaxed);
}
has_pending_evict_ = false;
if (slot < 0 || slot >= cap_ || target_cut <= 0 ||

@cubic-dev-ai cubic-dev-ai Bot Jul 31, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2: An invalid confirmation can leave the pending eviction reservation active while the caller proceeds as if the snapshot committed, so later requests inherit stale reservation state and the cache metadata can diverge from the backend slot. Treating this rejection as an abort (or otherwise clearing the reservation) would keep the lifecycle consistent.

Prompt for AI agents 
Check if this issue is valid — if so, understand the root cause and fix it. At server/src/server/prefix_cache.cpp, line 255:

<comment>An invalid confirmation can leave the pending eviction reservation active while the caller proceeds as if the snapshot committed, so later requests inherit stale reservation state and the cache metadata can diverge from the backend slot. Treating this rejection as an abort (or otherwise clearing the reservation) would keep the lifecycle consistent.</comment>

<file context>
@@ -281,68 +236,44 @@ std::pair<int, int> PrefixCache::prepare_inline_snap(
-            entries_size_count_.fetch_sub(1, std::memory_order_relaxed);
-        }
-        has_pending_evict_ = false;
+    if (slot < 0 || slot >= cap_ || target_cut <= 0 ||
+        target_cut > (int)prompt_ids.size()) {
+        std::fprintf(stderr,
</file context>

target_cut > (int)prompt_ids.size()) {
std::fprintf(stderr,
"[pc] rejected inline-snap slot=%d prefix_len=%d prompt_len=%zu\n",
slot, target_cut, prompt_ids.size());
return;
}

// The new snapshot replaces whatever this slot previously held. Drop any
// other entries still pointing at the slot: their hashes describe a
// different (or shorter) token stream than the new snapshot, and a later
// restore through them would attach mismatched KV. Stale entries arise
// when an aborted snap burns a round-robin next_slot_ step and a later
// confirm wraps onto a slot with a live entry (PR #370 repro).
for (int i = (int)entries_.size() - 1; i >= 0; --i) {
if (entries_[(size_t)i].slot == slot) {
std::fprintf(stderr,
"[pc] dropping stale entry for reused slot=%d\n", slot);
entries_.erase(entries_.begin() + i);
entries_size_count_.fetch_sub(1, std::memory_order_relaxed);
}
const auto key = hash_prefix(prompt_ids.data(), target_cut);
const auto result =
inline_state_.confirm(slot, key, target_cut, prompt_ids);
if (!result.accepted) {
std::fprintf(stderr,
"[pc] rejected inline-snap slot=%d prefix_len=%d prompt_len=%zu\n",
slot, target_cut, prompt_ids.size());
return;
}

auto key = hash_prefix(prompt_ids.data(), target_cut);
std::vector<int32_t> ids(prompt_ids.begin(), prompt_ids.begin() + target_cut);
entries_.push_back({key, slot, std::move(ids)});
entries_size_count_.fetch_add(1, std::memory_order_relaxed);
for (int i = 0; i < result.stale_slot_entries_removed; ++i) {
std::fprintf(stderr,
"[pc] dropping stale entry for reused slot=%d\n", slot);
}
sync_inline_size();
std::fprintf(stderr, "[pc] inline-snap committed slot=%d prefix_len=%d\n",
slot, target_cut);
}
Expand All @@ -353,31 +284,20 @@ void PrefixCache::abort_inline_snap(int slot) {
// metadata still pointing at it is therefore invalid, whether the slot was
// selected through the explicit eviction path or through a round-robin
// hole left by an earlier aborted reservation.
for (int i = (int)entries_.size() - 1; i >= 0; --i) {
if (entries_[(size_t)i].slot == slot) {
entries_.erase(entries_.begin() + i);
entries_size_count_.fetch_sub(1, std::memory_order_relaxed);
}
}
has_pending_evict_ = false;
inline_state_.abort(slot);
sync_inline_size();
}

void PrefixCache::cancel_inline_snap(int slot) {
if (disabled_) return;
if (has_pending_evict_) {
const int idx = find_entry(pending_evict_key_);
if (idx >= 0 && entries_[idx].slot != slot) return;
}
has_pending_evict_ = false;
inline_state_.cancel(slot);
}

void PrefixCache::mark_all_cleared() {
if (disabled_) return;
int n = (int)entries_.size();
entries_.clear();
entries_size_count_.store(0, std::memory_order_relaxed);
next_slot_ = 0;
has_pending_evict_ = false;
const int n = inline_state_.size();
inline_state_.clear();
sync_inline_size();
std::fprintf(stderr, "[pc] all-cleared — dropped %d LRU entries\n", n);
}

Expand Down
32 changes: 9 additions & 23 deletions server/src/server/prefix_cache.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,9 +12,9 @@

#pragma once

#include "prefix_cache_state.h"
#include "tokenizer.h"

#include <array>
#include <atomic>
#include <cstdint>
#include <functional>
Expand Down Expand Up @@ -42,20 +42,15 @@ std::vector<int> find_all_boundaries(const std::vector<int32_t> & ids,
const ChatMarkers & markers);

// SHA-1 hash of a prefix (truncated to 16 bytes).
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.
int select_inline_evict_victim(const std::vector<const std::vector<int32_t> *> & ids_lru);
int select_inline_evict_victim(const std::vector<std::vector<int32_t>> & ids_lru);
// is found. The implementation lives in prefix_cache_state.h so the exact
// production policy can also be model-checked without server dependencies.

// Pick the inline snapshot boundary for a request. We cache the boundary before
// the current user turn (second-to-last marker) and only when it advances past
Expand Down Expand Up @@ -163,17 +158,9 @@ class PrefixCache {
int cap_ = 0;
ChatMarkers markers_;

// LRU for inline prefix cache: ordered map of hash → slot.
// We use a vector to maintain insertion order (front = oldest).
struct LruEntry {
PrefixHash hash;
int slot;
std::vector<int32_t> ids; // prefix tokens [0, target_cut) for prefix-aware eviction
};
std::vector<LruEntry> entries_;
int next_slot_ = 0;
PrefixHash pending_evict_key_{};
bool has_pending_evict_ = false;
// Boundary detection and hashing live in PrefixCache; all inline-cache
// transitions live in this dependency-free core shared with ESBMC.
InlinePrefixCacheState inline_state_;

// Full-cache state
bool full_disabled_ = true;
Expand All @@ -194,19 +181,18 @@ class PrefixCache {
std::atomic<int64_t> lifetime_hits_{0}; // inline cache hits
std::atomic<int64_t> full_lifetime_hits_{0}; // full-compress cache hits
std::atomic<int64_t> full_disk_bytes_{0}; // best-effort snapshot of disk usage
// Atomic mirrors of `entries_.size()` and `full_entries_.size()`.
// Atomic mirrors of `inline_state_.size()` and `full_entries_.size()`.
// The vectors themselves are mutated only on the daemon thread
// under the daemon's serialised request loop, but `/props` reads
// happen from the client thread — calling `.size()` there is a
// data race per the C++ memory model. Bump these alongside every
// push_back / erase / clear so the public introspection counters
// stay well-defined. (Codex r1 P2 follow-up.)
std::atomic<int64_t> entries_size_count_{0}; // mirrors entries_.size()
std::atomic<int64_t> entries_size_count_{0}; // mirrors inline_state_.size()
std::atomic<int64_t> full_entries_size_count_{0}; // mirrors full_entries_.size()

// Helpers
int find_entry(const PrefixHash & h) const;
void move_to_end(int idx);
void sync_inline_size();
int find_full_entry(const PrefixHash & h) const;
void move_full_to_end(int idx);
};
Expand Down
Loading
Loading