diff --git a/docs/DETAILS.md b/docs/DETAILS.md index 3f673bff77b..f6a8bffabee 100644 --- a/docs/DETAILS.md +++ b/docs/DETAILS.md @@ -908,6 +908,20 @@ follows the reused prefix (or the root) is read in one run, and its session is n It still starts from a checkpoint it matches, and still saves the system-prompt root when that reaches `--prompt-cache-root`. Without the field (or with `true`) nothing changes. +**Checkpoints spread back from the end (opt-in, `--prompt-cache-spread`).** Least recent use keeps the newest points: +every turn adds a checkpoint at its start, so after a few short turns the 6 sit within the last 1-2K tokens. An agent +client that rewrites its history a few thousand tokens back (it trims an old tool output, drops a step) then matches +none of them and resumes from the system prompt - the whole conversation is read again. With the flag the checkpoint +that leaves is the one whose removal leaves the smallest gap for its distance from the newest, so the kept points thin +out geometrically: after 200 turns of 1,000 tokens they sit 0, 1K, 3K, 8K, 43K and 200K tokens before the end, and an +edit d tokens back reads about d again (6K back: 8K instead of 200K; 42K back: 43K). The same number of checkpoints, no +more RAM; the system-prompt root, a pinned prefix and the newest point are never dropped, and the resume is the usual +exact-prefix match. Measured on 2x CMP 50HX (layer split, IQ3_XXS) with the pi coding agent: before, a 210,528-token +prompt that was the same as the session for 204,349 tokens resumed at 9,525 and read 201,003 again (100 s). With the +flag, an 82,593-token conversation grown turn by turn with a turn ~6 steps back rewritten resumed at 73,610 tokens +(10.2 s against 40.0 s read in full), and with the expert placement fixed (`--adapt-swaps 0`, fixed `--pcie-frac`) the +greedy answer was the same text as the full read. Without the flag nothing changes. + **One long document, many questions: a pinned shared prefix (opt-in, 0.1.40.2, [RESEARCH_RUNS.md](RESEARCH_RUNS.md)).** A request can say where the document ends: `"strata_prefix": {"messages": 1}` (the first message is the document), `{"message": 0, "chars": 210000}` (the first 210,000 characters of message 0, for a client that sends the document and diff --git a/include/strata/program/conv_cache.hpp b/include/strata/program/conv_cache.hpp index 07c23adead9..d8e97f32c81 100644 --- a/include/strata/program/conv_cache.hpp +++ b/include/strata/program/conv_cache.hpp @@ -77,4 +77,26 @@ inline size_t eviction_victim(const uint64_t* stamps, size_t n, int64_t cap, con return v; } +/// --prompt-cache-spread: the item whose removal leaves the smallest gap for its distance from the newest, so the kept +/// points thin out geometrically towards the root and an edit d tokens before the end reads about d again. `lens`: +/// the chain's token counts. Never the root, the newest or a pinned item; the tail rule first, as above. +inline size_t spread_victim(const int64_t* lens, const uint64_t* stamps, size_t n, int64_t cap, const bool* tail, + const bool* pinned) { + if (n >= 3 && cap >= 2) { + auto candidate = [&](size_t i) { return !(pinned && pinned[i]); }; + if (tail != nullptr) + for (size_t i = 1; i + 1 < n; ++i) + if (tail[i] && candidate(i)) return i; + size_t v = n; + double best = 0; + for (size_t i = 1; i + 1 < n; ++i) { + if (!candidate(i)) continue; + const double gap = double(lens[i + 1] - lens[i - 1]) / double(lens[n - 1] - lens[i - 1]); + if (v == n || gap < best) { v = i; best = gap; } + } + if (v != n) return v; + } + return eviction_victim(stamps, n, cap, tail, pinned); +} + } // namespace strata::program::conv_cache diff --git a/src/program/conv_cache_test.cpp b/src/program/conv_cache_test.cpp index d76c800b4f6..40da79573e6 100644 --- a/src/program/conv_cache_test.cpp +++ b/src/program/conv_cache_test.cpp @@ -20,6 +20,7 @@ #include using strata::program::conv_cache::eviction_victim; +using strata::program::conv_cache::spread_victim; namespace { int g_fail = 0; @@ -29,8 +30,52 @@ void check(bool ok, const char* what) { } } // namespace +// a conversation of `turns` turns of `turn` tokens each after a `root`-token system prompt, a checkpoint at every turn's +// end and at most `cap` kept; returns how many tokens a request whose prompt differs `back` tokens before the end reuses +int64_t reuse_after_edit(bool spread, int64_t root, int64_t turn, int turns, int64_t cap, int64_t back) { + std::vector lens = {root}; + std::vector stamps = {1}; + for (int t = 1; t <= turns; ++t) { + lens.push_back(root + t * turn); + stamps.push_back(uint64_t(t + 1)); + while ((int64_t) lens.size() > cap) { + const size_t v = spread ? spread_victim(lens.data(), stamps.data(), lens.size(), cap, nullptr, nullptr) + : eviction_victim(stamps.data(), stamps.size(), cap); + lens.erase(lens.begin() + (std::ptrdiff_t) v); + stamps.erase(stamps.begin() + (std::ptrdiff_t) v); + } + } + const int64_t edit = lens.back() - back; + int64_t reuse = 0; + for (int64_t l : lens) + if (l <= edit) reuse = l; + return reuse; +} + int main() { std::printf("conv_cache_test\n"); + { // --prompt-cache-spread: 200 turns of 1,000 tokens after a 9,525-token system prompt, 6 checkpoints + const int64_t end = 9525 + 200 * 1000; + const int64_t lru6 = reuse_after_edit(false, 9525, 1000, 200, 6, 6000); + const int64_t spread6 = reuse_after_edit(true, 9525, 1000, 200, 6, 6000); + std::printf(" an edit 6K before the end: least recent use reuses %lld, spread %lld of %lld\n", (long long) lru6, + (long long) spread6, (long long) end); + check(lru6 == 9525, "least recent use: an edit 6K back falls through to the root"); + check(end - spread6 < 4 * 6000, "spread: an edit 6K back reads less than 24K again"); + const int64_t spread42 = reuse_after_edit(true, 9525, 1000, 200, 6, 42000); + check(end - spread42 < 4 * 42000, "spread: an edit 42K back reads less than 168K again"); + check(reuse_after_edit(true, 9525, 1000, 200, 6, 500) == end - 1000, "spread: the turn before is still kept"); + } + { + const std::vector lens = {100, 200, 300, 400}; + const std::vector stamps = {1, 2, 3, 4}; + const bool pinned[] = {false, true, true, false}; + check(spread_victim(lens.data(), stamps.data(), 4, 3, nullptr, pinned) == + eviction_victim(stamps.data(), 4, 3, nullptr, pinned), + "spread: every candidate pinned - the plain policy"); + check(spread_victim(lens.data(), stamps.data(), 2, 1, nullptr, nullptr) == 0, + "spread: two items, one slot - the plain policy"); + } { const std::vector stamps = {1, 2, 3, 4, 5, 6, 7}; // chat A: root + six turns, one over check(eviction_victim(stamps.data(), stamps.size(), 6) == 1, diff --git a/src/program/generate.cpp b/src/program/generate.cpp index ff3f2e1a659..df356270211 100644 --- a/src/program/generate.cpp +++ b/src/program/generate.cpp @@ -625,6 +625,7 @@ struct Options { /// --serve: also keep a checkpoint every N freshly read prompt tokens (0 = only at the last turn boundary) int64_t prompt_cache_every = 16384; bool prompt_cache_tail = false; // optional extra checkpoint at an existing near-tail chunk boundary + bool prompt_cache_spread = false; ///< --prompt-cache-spread: keep checkpoints spread back from the end, not the latest /// --serve: a prompt read from token 0 is also checkpointed at its first turn boundary - the end of the system /// prompt, which every chat of the same client shares - when that is at least N tokens (0 = never) int64_t prompt_cache_root = 2048; @@ -788,6 +789,9 @@ void usage() { " --prompt-cache-every N --serve: also checkpoint every N fresh prompt tokens (default 16384, 0 = off)\n" " --prompt-cache-tail --serve, single GPU: one extra checkpoint near the prompt's end, at an existing\n" " chunk boundary (default off; requires --prompt-cache and --prompt-cache-every > 0)\n" + " --prompt-cache-spread --serve: keep the checkpoints spread back from the conversation's end (dense\n" + " near it, thinner further back) instead of the most recently used: a client that\n" + " rewrites its history N tokens back reads about N again, not the whole conversation\n" " --turn-token ID --serve: the token that opens a chat turn (default 248045, <|im_start|>)\n" " --tail-role-token ID --serve: a turn of this role right before the last turn is left out of\n" " the conversation checkpoint (a trailing effort turn; default -1 = off)\n" @@ -1983,6 +1987,7 @@ int main(int argc, char** argv) { else o.conversation_cache_slots = (int) number; } else if (a == "--prompt-cache-tail") o.prompt_cache_tail = true; + else if (a == "--prompt-cache-spread") o.prompt_cache_spread = true; else if (a == "--prompt-cache-every") o.prompt_cache_every = std::max(0LL, std::atoll(next("--prompt-cache-every"))); else if (a == "--prompt-cache-root") o.prompt_cache_root = std::max(0LL, std::atoll(next("--prompt-cache-root"))); else if (a == "--turn-token") o.turn_token = std::atoll(next("--turn-token")); @@ -7948,8 +7953,13 @@ int main(int argc, char** argv) { for (size_t i = 0; i < checks.size(); ++i) pin_flags[i] = checks[i].pinned; break; } - const size_t victim = strata::program::conv_cache::eviction_victim( - stamps.data(), stamps.size(), o.prompt_cache, tail_flags.get(), pin_flags.get()); + std::vector lens; + for (const ConvCheckpoint& k : checks) lens.push_back((int64_t) k.ids.size()); + const size_t victim = o.prompt_cache_spread + ? strata::program::conv_cache::spread_victim(lens.data(), stamps.data(), stamps.size(), + o.prompt_cache, tail_flags.get(), pin_flags.get()) + : strata::program::conv_cache::eviction_victim(stamps.data(), stamps.size(), o.prompt_cache, + tail_flags.get(), pin_flags.get()); checks.erase(checks.begin() + (std::ptrdiff_t) victim); } return true;