Something went wrong. Try again.
strata but optimized for m5 (and maybe more later)
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228// Suffix cache reuse (SCR): after an edit in the middle of a prompt, the K/V of the text that survived the edit is// moved to its new positions instead of being read again. Ported from llama-cpp-bonsai's common/scr.cpp.//// The moved cells were computed over the text before the edit, so the result approximates a fresh read. The GDN// layers have no cells to move: their state continues from a conversation checkpoint at the end of the last moved// run (the plan's `adopt`), and everything after that run is read from it. Strata's indexer pools 4 cells into one// key and keeps only the pooled keys, so a run starts and ends on a block boundary of the OLD prompt and moves by a// multiple of 4: up to 3 pad tokens are read before a run to make its shift one (DeepSeek V4.1's recipe pads to an// even shift the same way). The moved blocks then keep their cells, and their pooled keys only turn by the shift.#pragma once
#include <algorithm>#include <cstdint>#include <string>#include <unordered_map>#include <vector>
namespace strata::core {
struct ScrParams { int64_t min_span = 256; // shortest run worth moving int64_t max_runs = 64; // most runs moved per request int64_t seed = 32; // n-gram length that seeds a match int64_t block = 4; // the indexer's cells per pooled key: read parts start and end on these};
// the cells old[old_start, old_start + n) reappear at [new_start, new_start + n) of the prompt, after `pad` pad tokens// that the read inserts right before the run; `at` is where the run lands in the padded (physical) promptstruct ScrRun { int64_t new_start = 0, old_start = 0, n = 0; int64_t pad = 0, at = 0; int64_t shift() const { return at - old_start; }};
struct ScrPlan { int64_t n_prefix = 0; // leading tokens equal in both prompts int64_t restore = 0; // the checkpoint the read restarts from (at or before n_prefix) int64_t adopt = -1; // old length of the checkpoint the GDN state continues from after the last run std::vector<ScrRun> runs; // in new order, at or after restore
bool empty() const { return runs.empty(); } // the prompt position the final read starts at, and the same in the padded prompt int64_t resume() const { return runs.empty() ? restore : runs.back().new_start + runs.back().n; } int64_t resume_at() const { return runs.empty() ? restore : runs.back().at + runs.back().n; } int64_t pads() const { int64_t n = 0; for (const ScrRun& r : runs) n += r.pad; return n; } int64_t reused() const { int64_t n = 0; for (const ScrRun& r : runs) n += r.n; return n; } std::string str() const { std::string s = "prefix " + std::to_string(n_prefix) + ", restore " + std::to_string(restore) + ", adopt " + std::to_string(adopt); for (const ScrRun& r : runs) s += ", R " + std::to_string(r.new_start) + " <- " + std::to_string(r.old_start) + " x " + std::to_string(r.n) + (r.pad ? " pad " + std::to_string(r.pad) : ""); return s; }};
namespace scr_detail {// greedy maximal spans of cur[from, ...) found in old[from, ...), seeded by n-gram hashes, each old token used oncetemplate <class A, class B>std::vector<ScrRun> find_spans(const std::vector<A>& old, const std::vector<B>& cur, int64_t from, const ScrParams& p) { std::vector<ScrRun> spans; const int64_t n_old = (int64_t) old.size(), n_cur = (int64_t) cur.size(), k = p.seed; if (n_old - from < p.min_span || n_cur - from < p.min_span) return spans; auto gram = [k](const auto& v, int64_t i) { uint64_t h = 1469598103934665603ULL; for (int64_t j = 0; j < k; ++j) h = (h ^ (uint32_t) v[(size_t) (i + j)]) * 1099511628211ULL; return h; }; constexpr size_t max_bucket = 16; // the first few old positions of every seed n-gram std::unordered_map<uint64_t, std::vector<int64_t>> index; index.reserve((size_t) (n_old - from)); for (int64_t o = from; o + k <= n_old; ++o) { auto& b = index[gram(old, o)]; if (b.size() < max_bucket) b.push_back(o); } std::vector<bool> used((size_t) n_old, false); for (int64_t q = from; q + k <= n_cur;) { const auto it = index.find(gram(cur, q)); int64_t best_o = -1, best_n = 0; if (it != index.end()) for (const int64_t o : it->second) { int64_t n = 0; while (q + n < n_cur && o + n < n_old && !used[(size_t) (o + n)] && (int64_t) cur[(size_t) (q + n)] == (int64_t) old[(size_t) (o + n)]) ++n; if (n > best_n) { best_n = n; best_o = o; } } if (best_n >= p.min_span) { spans.push_back({q, best_o, best_n}); std::fill(used.begin() + best_o, used.begin() + best_o + best_n, true); q += best_n; } else { ++q; } } return spans;}inline int64_t round_up(int64_t x, int64_t b) { return (x + b - 1) / b * b; }inline int64_t round_down(int64_t x, int64_t b) { return x / b * b; }} // namespace scr_detail
// Plans the reuse of a session that holds old[0, n_old) for the prompt `cur`. `checkpoints`: the old lengths whose// GDN state can be restored (the live state counts as old.size()). A run starts on a block boundary of the new// prompt; every run but the last also ends on one, and the last one ends at a checkpoint. The plan always leaves at// least one token of `cur` to read.template <class A, class B>ScrPlan scr_plan(const std::vector<A>& old, const std::vector<B>& cur, std::vector<int64_t> checkpoints, const ScrParams& p) { using namespace scr_detail; ScrPlan plan; const int64_t n_cur = (int64_t) cur.size(), bl = std::max<int64_t>(p.block, 1); std::sort(checkpoints.begin(), checkpoints.end()); int64_t a = 0; while (a < (int64_t) old.size() && a < n_cur && (int64_t) old[(size_t) a] == (int64_t) cur[(size_t) a]) ++a; plan.n_prefix = a; for (const int64_t q : checkpoints) if (q <= a) plan.restore = std::max(plan.restore, q); auto spans = find_spans(old, cur, a, p); // a run starts on a block boundary of the old prompt and (but the last) ends on one for (ScrRun& s : spans) { const int64_t cut = round_up(s.old_start, bl) - s.old_start; s.new_start += cut; s.old_start += cut; s.n -= cut; s.n = round_down(s.old_start + s.n, bl) - s.old_start; } spans.erase(std::remove_if(spans.begin(), spans.end(), [&](const ScrRun& s) { return s.n < p.min_span; }), spans.end()); // the largest runs, back in prompt order std::sort(spans.begin(), spans.end(), [](const ScrRun& x, const ScrRun& y) { return x.n > y.n; }); if ((int64_t) spans.size() > p.max_runs) spans.resize((size_t) p.max_runs); std::sort(spans.begin(), spans.end(), [](const ScrRun& x, const ScrRun& y) { return x.new_start < y.new_start; }); // the last run ends at a checkpoint, which also ends the reuse: everything after it is read from that state. // It was cut back to a block boundary, so a checkpoint up to block - 1 cells past it may still match. while (!spans.empty()) { ScrRun& last = spans.back(); int64_t end = -1; for (const int64_t q : checkpoints) { const int64_t n = q - last.old_start; if (n < p.min_span || n >= last.n + bl || last.new_start + n >= n_cur || q <= end) continue; if (n <= last.n || std::equal(old.begin() + last.old_start + last.n, old.begin() + q, cur.begin() + last.new_start + last.n, [](const A& x, const B& y) { return (int64_t) x == (int64_t) y; })) end = q; } if (end >= 0) { last.n = end - last.old_start; plan.adopt = end; break; } spans.pop_back(); } // the pads: walking the runs in prompt order, each gets what makes its padded position a multiple of the block // away from its old one int64_t pads = 0; for (ScrRun& s : spans) { s.pad = ((s.old_start - (s.new_start + pads)) % bl + bl) % bl; pads += s.pad; s.at = s.new_start + pads; } plan.runs = std::move(spans); return plan;}
// The pads an earlier splice read (`pads`: their positions in `live`, sorted) that the prompt `cur` still shares with// the session: the ones before the last token of `cur` that `live` holds at the same place once the pads are// skipped. `out` is `cur` with those pads put back, which makes it start with `live` as far as they agree.template <class A, class B>std::vector<int64_t> scr_carry(const std::vector<B>& cur, const std::vector<A>& live, const std::vector<int64_t>& pads, int64_t pad_token, std::vector<B>& out) { size_t p = 0, i = 0, k = 0, kept = 0; while (p < live.size()) { if (k < pads.size() && pads[k] == (int64_t) p) { if ((int64_t) live[p] != pad_token) break; ++k; ++p; continue; } if (i + 1 >= cur.size() || (int64_t) live[p] != (int64_t) cur[i]) break; ++i; ++p; kept = k; } std::vector<int64_t> carried(pads.begin(), pads.begin() + (long) kept); out.clear(); out.reserve(cur.size() + carried.size()); size_t c = 0; for (const int64_t at : carried) { while ((int64_t) out.size() < at) out.push_back(cur[c++]); out.push_back((B) pad_token); } out.insert(out.end(), cur.begin() + (long) c, cur.end()); return carried;}
// The positions of the pads `plan` inserts, in the padded prompt.inline std::vector<int64_t> scr_pad_positions(const ScrPlan& plan) { std::vector<int64_t> at; for (const ScrRun& r : plan.runs) for (int64_t j = r.pad; j > 0; --j) at.push_back(r.at - j); return at;}
// The prompt the engine reads under `plan`: `cur` with each run's pads inserted right before it.template <class B>std::vector<B> scr_padded(const std::vector<B>& cur, const ScrPlan& plan, B pad_token) { std::vector<B> out; out.reserve(cur.size() + (size_t) plan.pads()); int64_t from = 0; for (const ScrRun& r : plan.runs) { out.insert(out.end(), cur.begin() + from, cur.begin() + r.new_start); out.insert(out.end(), (size_t) r.pad, pad_token); from = r.new_start; } out.insert(out.end(), cur.begin() + from, cur.end()); return out;}
} // namespace strata::core