serve: --prompt-cache-spread, checkpoints kept spread back from the conversation's end (opt-in)

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

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

src/program/conv_cache_test.cpp

@@ -20,6 +20,7 @@
 #include <vector>
 
 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<int64_t> lens = {root};
+    std::vector<uint64_t> 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<int64_t> lens = {100, 200, 300, 400};
+        const std::vector<uint64_t> 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<uint64_t> stamps = {1, 2, 3, 4, 5, 6, 7};   // chat A: root + six turns, one over
         check(eviction_victim(stamps.data(), stamps.size(), 6) == 1,

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<int64_t> 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;