#include "scr.h"
#include "../src/llama-ext.h"
#include <algorithm>
#include <unordered_map>
int32_t common_scr_plan::resume() const {
return runs.empty() ? restore : runs.back().new_start + runs.back().n;
}
int32_t common_scr_plan::n_reused() const {
int32_t n = 0;
for (const auto & r : runs) {
n += r.n;
}
return n;
}
std::string common_scr_plan::str() const {
std::string s = "prefix " + std::to_string(n_prefix) + ", restore " + std::to_string(restore) + ", adopt " + std::to_string(adopt);
for (const auto & r : runs) {
s += ", R " + std::to_string(r.new_start) + " <- " + std::to_string(r.old_start) + " x " + std::to_string(r.n);
}
return s;
}
// greedy maximal spans of cur[from, ...) found in old[from, ...), seeded by n-gram hashes, each old token used once
static std::vector<common_scr_run> scr_find_spans(
const std::vector<llama_token> & old,
const std::vector<llama_token> & cur,
int32_t from,
const common_scr_params & params) {
std::vector<common_scr_run> spans;
const int32_t n_old = (int32_t) old.size();
const int32_t n_cur = (int32_t) cur.size();
const int32_t k = params.seed;
if (n_old - from < params.min_span || n_cur - from < params.min_span) {
return spans;
}
auto gram = [k](const std::vector<llama_token> & v, int32_t i) {
uint64_t h = 1469598103934665603ULL;
for (int32_t j = 0; j < k; ++j) {
h = (h ^ (uint32_t) v[i + j]) * 1099511628211ULL;
}
return h;
};
// the first few old positions of every seed n-gram
const size_t max_bucket = 16;
std::unordered_map<uint64_t, std::vector<int32_t>> index;
index.reserve(n_old - from);
for (int32_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(n_old, false);
for (int32_t p = from; p + k <= n_cur; ) {
const auto it = index.find(gram(cur, p));
int32_t best_o = -1;
int32_t best_n = 0;
if (it != index.end()) {
for (const int32_t o : it->second) {
int32_t n = 0;
while (p + n < n_cur && o + n < n_old && !used[o + n] &&
cur[p + n] == old[o + n] && cur[p + n] != LLAMA_TOKEN_NULL) {
++n;
}
if (n > best_n) {
best_n = n;
best_o = o;
}
}
}
if (best_n >= params.min_span) {
spans.push_back({ p, best_o, best_n });
std::fill(used.begin() + best_o, used.begin() + best_o + best_n, true);
p += best_n;
} else {
++p;
}
}
return spans;
}
common_scr_plan common_scr_plan_build(
const std::vector<llama_token> & old,
const std::vector<llama_token> & cur,
const std::vector<int32_t> & snapshots,
bool recurrent,
const common_scr_params & params) {
common_scr_plan plan;
const int32_t n_cur = (int32_t) cur.size();
int32_t a = 0;
while (a < (int32_t) old.size() && a < n_cur && old[a] == cur[a] && cur[a] != LLAMA_TOKEN_NULL) {
++a;
}
plan.n_prefix = a;
plan.restore = a;
if (recurrent) {
plan.restore = 0;
for (const int32_t q : snapshots) {
if (q <= a) {
plan.restore = std::max(plan.restore, q);
}
}
}
auto spans = scr_find_spans(old, cur, a, params);
// the largest K spans, back in prompt order
std::sort(spans.begin(), spans.end(), [](const auto & x, const auto & y) { return x.n > y.n; });
if ((int32_t) spans.size() > params.max_runs) {
spans.resize(params.max_runs);
}
std::sort(spans.begin(), spans.end(), [](const auto & x, const auto & y) { return x.new_start < y.new_start; });
// the last span must end at a snapshot, which also ends the reuse: everything after it is prefilled from that state
while (!spans.empty()) {
auto & last = spans.back();
int32_t end = -1;
if (recurrent) {
for (const int32_t q : snapshots) {
const int32_t n = q - last.old_start;
if (n >= params.min_span && n <= last.n && last.new_start + n < n_cur) {
end = std::max(end, q);
}
}
} else {
end = last.old_start + std::min(last.n, n_cur - 1 - last.new_start);
}
if (end - last.old_start >= params.min_span) {
last.n = end - last.old_start;
plan.adopt = recurrent ? end : -1;
break;
}
spans.pop_back();
}
plan.runs = std::move(spans);
return plan;
}
int32_t common_scr_apply_kv(llama_context * ctx, llama_seq_id seq_id, const common_scr_plan & plan) {
// a run may move once its target holds no source of a run still to move
std::vector<common_scr_run> order;
{
std::vector<common_scr_run> pending = plan.runs;
while (!pending.empty()) {
auto next = std::find_if(pending.begin(), pending.end(), [&](const common_scr_run & r) {
return std::none_of(pending.begin(), pending.end(), [&](const common_scr_run & o) {
return &o != &r && o.old_start < r.new_start + r.n && r.new_start < o.old_start + o.n;
});
});
if (next == pending.end()) {
return 0;
}
order.push_back(*next);
pending.erase(next);
}
}
std::vector<common_scr_run> sources = plan.runs;
std::sort(sources.begin(), sources.end(), [](const auto & x, const auto & y) { return x.old_start < y.old_start; });
llama_pos pos = plan.restore;
for (const auto & s : sources) {
if (s.old_start > pos && !llama_scr_kv_rm(ctx, seq_id, pos, s.old_start)) {
return -1;
}
pos = std::max(pos, (llama_pos) (s.old_start + s.n));
}
if (!llama_scr_kv_rm(ctx, seq_id, pos, -1)) {
return -1;
}
for (const auto & r : order) {
if (!llama_scr_kv_relocate(ctx, seq_id, r.old_start, r.old_start + r.n, r.new_start - r.old_start)) {
return -1;
}
}
return 1;
}