From: opencoti Subject: [PATCH 0010] retention-aware KV eviction — KeyDiff key-similarity, server-window When an agentic context exceeds a slot's n_ctx, the server's context-shift discards the *positional middle* [n_keep, n_keep+n_discard) regardless of what sits there. This patch makes that choice retention-aware: it discards the lowest-value *contiguous window* of the same width instead, keeping the distinctive tokens and dropping the redundant ones. Reference choice. The named technique (TRIM-KV, arXiv 2605.09649) needs per-model retention gates trained by distillation — not a drop-in runtime patch. Attention-score methods (H2O, SnapKV) need the materialized attention matrix, which Flash Attention never produces (the vendored pin defaults flash_attn_type = AUTO). So this substitutes a score-free, FA-compatible criterion — KeyDiff (arXiv 2504.15364): score each cached key by its *distinctiveness*, the mean over heads of (1 - cosine(key_head, mean_head)). Keys close to the sequence's mean are redundant (evict); keys far from it are distinctive (keep). Mechanism (the "server-window" depth — mirror-safe): - src/llama-memory.h: a new non-pure virtual seq_key_scores(seq_id, layer_hint) with a default empty body, so no backend is forced to implement it and the pure-virtual surface is unchanged. (Adding it grows the vtable — a consistent recompile of all llama.cpp TUs is required; see the build gotcha below.) - src/llama-kv-cache.cpp: the real implementation — read each cached key of one representative layer to host (amortized: only at the rare context-shift trigger, unlike per-decode-interval scoring), group into heads, return a per-position distinctiveness vector (+INF for absent positions). Empty when unsupported (no cells, or a key type with no float converter) so the caller falls back to positional discard. - src/llama-kv-cache-iswa.cpp: delegate to the non-SWA base cache. - src/llama-context.cpp + include/llama.h: a thin C shim llama_memory_seq_key_scores() — a plain virtual call (no RTTI/downcast). - tools/server/server-context.cpp: inside the EXISTING context-shift block, when rest_kv_eviction is on, slide a width-n_discard window over [n_keep, n_tokens - rest_kv_recent) and discard the minimum-retention one. It is still ONE contiguous window, so the seq_rm/seq_add + token-mirror rewrite stay 1:1 with KV positions — no rework of prefix-reuse/sampler. - common/common.h + common/arg.cpp: --rest-kv-eviction (off by default), --rest-kv-recent N (protect the recent tail), --rest-kv-layer N (-1 = auto). Regression shield: --rest-kv-eviction off (the default, and every non-opencoti client) leaves w0 = n_keep, i.e. context-shift is byte-for-byte upstream. Any failure (non-KV memory, quantized keys, no room after the recent tail) also falls back to the positional middle. Build gotcha (buglog bug-175): seq_key_scores is a new virtual on llama_memory_i, so the vtable layout changes. The llamafile Makefile tracks no header deps; force a consistent recompile of all llama.cpp src/common/tools objects (delete x86 .o, rebuild) — and because cosmocc is a fat compiler whose aarch64 sibling object is NOT restored by a ccache hit, delete only x86 .o that lack an aarch64 twin and rebuild with CCACHE_DISABLE=1, or use the full reproducible build:llamafile. Bench: perf/llamafile/rest-kv-eviction.bench.ts Milestone: F5 M1 — see docs/features/advanced_kv.md Design: docs/features/advanced_kv.md (M1) Upstreaming: deferred; candidate to propose upstream as a pluggable context-shift eviction policy once benched on opencoti workloads. --- diff --git a/llama.cpp/common/arg.cpp b/llama.cpp/common/arg.cpp --- a/llama.cpp/common/arg.cpp +++ b/llama.cpp/common/arg.cpp @@ -1398,6 +1398,29 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.ctx_shift = value; } ).set_examples({LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_IMATRIX, LLAMA_EXAMPLE_PERPLEXITY}).set_env("LLAMA_ARG_CONTEXT_SHIFT")); + // opencoti F5 M1 rest-kv-eviction — see docs/features/advanced_kv.md + add_opt(common_arg( + {"--rest-kv-eviction"}, + {"--no-rest-kv-eviction"}, + string_format("when context shift fires, evict the lowest-retention (KeyDiff key-similarity) contiguous window instead of the positional middle (default: %s)", params.rest_kv_eviction ? "enabled" : "disabled"), + [](common_params & params, bool value) { + params.rest_kv_eviction = value; + } + ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_REST_KV_EVICTION")); + add_opt(common_arg( + {"--rest-kv-recent"}, "N", + string_format("rest-kv: protect the most-recent N tokens from eviction (default: %d)", params.rest_kv_recent), + [](common_params & params, int value) { + params.rest_kv_recent = value; + } + ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_REST_KV_RECENT")); + add_opt(common_arg( + {"--rest-kv-layer"}, "N", + string_format("rest-kv: model layer whose keys score retention, -1 = auto/middle (default: %d)", params.rest_kv_layer), + [](common_params & params, int value) { + params.rest_kv_layer = value; + } + ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_REST_KV_LAYER")); add_opt(common_arg( {"--chunks"}, "N", string_format("max number of chunks to process (default: %d, -1 = all)", params.n_chunks), diff --git a/llama.cpp/common/common.h b/llama.cpp/common/common.h --- a/llama.cpp/common/common.h +++ b/llama.cpp/common/common.h @@ -554,6 +554,14 @@ struct common_params { bool swa_full = false; // use full-size SWA cache (https://github.com/ggml-org/llama.cpp/pull/13194#issuecomment-2868343055) bool kv_unified = false; // enable unified KV cache + // opencoti F5 M1 rest-kv-eviction — see docs/features/advanced_kv.md + // When a slot fills n_ctx and ctx_shift fires, discard the lowest-retention + // *contiguous* window (KeyDiff key-similarity) instead of the blind + // positional middle. Off => byte-for-byte upstream context-shift. + bool rest_kv_eviction = false; // master switch (default off) + int32_t rest_kv_recent = 256; // protect the most-recent N tokens from eviction + int32_t rest_kv_layer = -1; // representative layer for scoring; -1 = auto (middle) + bool input_prefix_bos = false; // prefix BOS to user inputs, preceding input_prefix bool use_mmap = true; // enable mmap to use filesystem cache bool use_direct_io = false; // read from disk without buffering diff --git a/llama.cpp/include/llama.h b/llama.cpp/include/llama.h --- a/llama.cpp/include/llama.h +++ b/llama.cpp/include/llama.h @@ -730,6 +730,19 @@ extern "C" { llama_pos p0, llama_pos p1); + // opencoti F5 M1 rest-kv-eviction — see docs/features/advanced_kv.md + // Fill out[0..n) with per-position KeyDiff retention scores for seq_id + // (higher = keep, +INF = absent position). layer_hint < 0 selects a default + // representative layer. Returns the number of finite scores written; 0 when + // unsupported (non-KV memory, a key type with no float converter, or an + // empty sequence) — the caller should then fall back to positional eviction. + LLAMA_API size_t llama_memory_seq_key_scores( + llama_memory_t mem, + llama_seq_id seq_id, + int32_t layer_hint, + float * out, + size_t n); + // Copy all tokens that belong to the specified sequence to another sequence // p0 < 0 : [0, p1] // p1 < 0 : [p0, inf) diff --git a/llama.cpp/src/llama-context.cpp b/llama.cpp/src/llama-context.cpp --- a/llama.cpp/src/llama-context.cpp +++ b/llama.cpp/src/llama-context.cpp @@ -3732,6 +3732,38 @@ bool llama_memory_seq_rm( return mem->seq_rm(seq_id, p0, p1); } +// opencoti F5 M1 rest-kv-eviction — see docs/features/advanced_kv.md +// Thin C shim over the (defaulted) virtual llama_memory_i::seq_key_scores — a +// plain virtual dispatch, so no RTTI/downcast and no change to the pure-virtual +// surface of the interface. Non-KV memories inherit the empty default and this +// returns 0, leaving the server on its positional fallback. +size_t llama_memory_seq_key_scores( + llama_memory_t mem, + llama_seq_id seq_id, + int32_t layer_hint, + float * out, + size_t n) { + if (!mem || !out || n == 0) { + return 0; + } + + const std::vector scores = mem->seq_key_scores(seq_id, layer_hint); + if (scores.empty()) { + return 0; + } + + size_t cnt = 0; + for (size_t p = 0; p < n; ++p) { + const float s = p < scores.size() ? scores[p] : std::numeric_limits::infinity(); + out[p] = s; + if (std::isfinite(s)) { + cnt++; + } + } + + return cnt; +} + void llama_memory_seq_cp( llama_memory_t mem, llama_seq_id seq_id_src, diff --git a/llama.cpp/src/llama-kv-cache-iswa.cpp b/llama.cpp/src/llama-kv-cache-iswa.cpp --- a/llama.cpp/src/llama-kv-cache-iswa.cpp +++ b/llama.cpp/src/llama-kv-cache-iswa.cpp @@ -117,6 +117,14 @@ llama_pos llama_kv_cache_iswa::seq_pos_max(llama_seq_id seq_id) const { return kv_swa->seq_pos_max(seq_id); } +std::vector llama_kv_cache_iswa::seq_key_scores(llama_seq_id seq_id, int32_t layer_hint) const { + // opencoti F5 M1 rest-kv-eviction — see docs/features/advanced_kv.md + // Score on the non-SWA base cache: the SWA cache only ever retains a + // sliding window, so its keys are not the right eviction target. The base + // cache is a superset of the SWA cache, so it has every position we score. + return kv_base->seq_key_scores(seq_id, layer_hint); +} + std::map llama_kv_cache_iswa::memory_breakdown() const { std::map mb = kv_base->memory_breakdown(); for (const auto & buft_size : kv_swa->memory_breakdown()) { diff --git a/llama.cpp/src/llama-kv-cache-iswa.h b/llama.cpp/src/llama-kv-cache-iswa.h --- a/llama.cpp/src/llama-kv-cache-iswa.h +++ b/llama.cpp/src/llama-kv-cache-iswa.h @@ -62,6 +62,9 @@ public: llama_pos seq_pos_min(llama_seq_id seq_id) const override; llama_pos seq_pos_max(llama_seq_id seq_id) const override; + // opencoti F5 M1 rest-kv-eviction — see docs/features/advanced_kv.md + std::vector seq_key_scores(llama_seq_id seq_id, int32_t layer_hint) const override; + std::map memory_breakdown() const override; // state write/load diff --git a/llama.cpp/src/llama-kv-cache.cpp b/llama.cpp/src/llama-kv-cache.cpp --- a/llama.cpp/src/llama-kv-cache.cpp +++ b/llama.cpp/src/llama-kv-cache.cpp @@ -711,6 +711,137 @@ bool llama_kv_cache::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) { return true; } +std::vector llama_kv_cache::seq_key_scores(llama_seq_id seq_id, int32_t layer_hint) const { + // opencoti F5 M1 rest-kv-eviction — see docs/features/advanced_kv.md + // KeyDiff-style, score-free (Flash-Attention-compatible) retention signal: + // for one representative layer, read every cached key of `seq_id`, group + // it into attention heads, and score each position by its key's + // *distinctiveness* — the mean over heads of (1 - cosine(key_head, + // mean_head)). High = far from the sequence's mean key = keep; low = + // redundant = safe to evict. The server uses this to choose the + // lowest-retention contiguous window to discard during context-shift, + // instead of the blind positional middle. Returns scores indexed by token + // position, +INF for absent positions (never selected). An empty result + // means "unsupported here" (no cells, or keys in a type with no float + // converter) and the caller falls back to positional discard. Reads happen + // only at the (infrequent) context-shift trigger, so the device->host copy + // is amortized — unlike per-decode-interval scoring. + + if (seq_id < 0 || (size_t) seq_id >= seq_to_stream.size() || layers.empty()) { + return {}; + } + + const uint32_t strm = seq_to_stream[seq_id]; + const auto & cells = v_cells[strm]; + + const llama_pos pmax = cells.seq_pos_max(seq_id); + if (pmax < 0) { + return {}; + } + + // pick a representative cache layer. A non-negative hint first tries to map + // a *model* layer id into the cache; failing that it is clamped as a direct + // cache-layer index. Default (hint < 0) is the middle layer — empirically + // the most informative for retention and cheap (one layer's keys only). + int32_t cil = (int32_t) layers.size() / 2; + if (layer_hint >= 0) { + const auto it = map_layer_ids.find(layer_hint); + cil = it != map_layer_ids.end() + ? it->second + : std::min(layer_hint, (int32_t) layers.size() - 1); + } + + const auto & layer = layers[cil]; + const uint32_t mil = layer.il; + + const uint32_t n_head_kv = hparams.n_head_kv(mil); + const uint32_t n_embd_head_k = hparams.n_embd_head_k(mil); + const uint32_t n_embd_k_gqa = n_head_kv * n_embd_head_k; + if (n_head_kv == 0 || n_embd_head_k == 0 || n_embd_k_gqa == 0) { + return {}; + } + + ggml_tensor * k = layer.k_stream[strm]; + if (!k) { + return {}; + } + + const auto * tt = ggml_get_type_traits(k->type); + if (!tt || !tt->to_float) { + return {}; // e.g. a quantized K cache with no dequantizer -> positional fallback + } + + const size_t row_bytes = ggml_row_size(k->type, n_embd_k_gqa); + + // collect (pos, cell index) for this sequence, ascending by position + std::vector> pc; + pc.reserve(cells.get_used()); + for (uint32_t i = 0; i < cells.size(); ++i) { + if (!cells.is_empty(i) && cells.seq_has(i, seq_id)) { + pc.emplace_back(cells.pos_get(i), i); + } + } + if (pc.empty()) { + return {}; + } + std::sort(pc.begin(), pc.end()); + + const size_t n = pc.size(); + + // read each key row once (amortized device->host copy), dequantized to f32 + std::vector keys((size_t) n * n_embd_k_gqa); + std::vector rowbuf(row_bytes); + for (size_t j = 0; j < n; ++j) { + const uint32_t i = pc[j].second; + ggml_backend_tensor_get(k, rowbuf.data(), (size_t) i * row_bytes, row_bytes); + tt->to_float(rowbuf.data(), keys.data() + j * n_embd_k_gqa, (int64_t) n_embd_k_gqa); + } + + // per-head mean key across the sequence, then per-head mean norm + std::vector mean((size_t) n_embd_k_gqa, 0.0f); + for (size_t j = 0; j < n; ++j) { + const float * kv = keys.data() + j * n_embd_k_gqa; + for (uint32_t e = 0; e < n_embd_k_gqa; ++e) { + mean[e] += kv[e]; + } + } + for (uint32_t e = 0; e < n_embd_k_gqa; ++e) { + mean[e] /= (float) n; + } + + std::vector mean_norm(n_head_kv, 0.0f); + for (uint32_t h = 0; h < n_head_kv; ++h) { + const float * mh = mean.data() + (size_t) h * n_embd_head_k; + double s = 0.0; + for (uint32_t d = 0; d < n_embd_head_k; ++d) { + s += (double) mh[d] * (double) mh[d]; + } + mean_norm[h] = (float) std::sqrt(s); + } + + // distinctiveness per position = mean over heads of (1 - cos(key_head, mean_head)) + std::vector out((size_t) pmax + 1, std::numeric_limits::infinity()); + for (size_t j = 0; j < n; ++j) { + const float * kv = keys.data() + j * n_embd_k_gqa; + double acc = 0.0; + for (uint32_t h = 0; h < n_head_kv; ++h) { + const float * kh = kv + (size_t) h * n_embd_head_k; + const float * mh = mean.data() + (size_t) h * n_embd_head_k; + double dot = 0.0, kn = 0.0; + for (uint32_t d = 0; d < n_embd_head_k; ++d) { + dot += (double) kh[d] * (double) mh[d]; + kn += (double) kh[d] * (double) kh[d]; + } + const double denom = std::sqrt(kn) * (double) mean_norm[h]; + const double cos = denom > 0.0 ? dot / denom : 0.0; + acc += (1.0 - cos); + } + out[(size_t) pc[j].first] = (float) (acc / (double) n_head_kv); + } + + return out; +} + void llama_kv_cache::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) { GGML_ASSERT(seq_id_src >= 0 && (size_t) seq_id_src < seq_to_stream.size()); GGML_ASSERT(seq_id_dst >= 0 && (size_t) seq_id_dst < seq_to_stream.size()); diff --git a/llama.cpp/src/llama-kv-cache.h b/llama.cpp/src/llama-kv-cache.h --- a/llama.cpp/src/llama-kv-cache.h +++ b/llama.cpp/src/llama-kv-cache.h @@ -146,6 +146,9 @@ public: llama_pos seq_pos_min(llama_seq_id seq_id) const override; llama_pos seq_pos_max(llama_seq_id seq_id) const override; + // opencoti F5 M1 rest-kv-eviction — see docs/features/advanced_kv.md + std::vector seq_key_scores(llama_seq_id seq_id, int32_t layer_hint) const override; + std::map memory_breakdown() const override; // state write/load diff --git a/llama.cpp/src/llama-memory.h b/llama.cpp/src/llama-memory.h --- a/llama.cpp/src/llama-memory.h +++ b/llama.cpp/src/llama-memory.h @@ -6,6 +6,7 @@ #include #include #include +#include struct llama_ubatch; @@ -112,6 +113,22 @@ struct llama_memory_i { virtual llama_pos seq_pos_min(llama_seq_id seq_id) const = 0; virtual llama_pos seq_pos_max(llama_seq_id seq_id) const = 0; + // opencoti F5 M1 rest-kv-eviction — see docs/features/advanced_kv.md + // Per-position KeyDiff "retention" scores for one sequence: higher = more + // distinctive (key far from the sequence's mean key) = keep; lower = more + // redundant = safe to evict. The server uses these to pick WHICH window to + // discard during context-shift, instead of the blind positional middle. + // Returns a vector indexed by token position [0, seq_pos_max]; absent + // positions are left at +INF (never evicted). The default returns empty so + // non-KV memories — and any caller that never opts in — are byte-for-byte + // unaffected (the eviction caller then falls back to positional discard). + // layer_hint < 0 selects a default representative layer. + virtual std::vector seq_key_scores(llama_seq_id seq_id, int32_t layer_hint) const { + (void) seq_id; + (void) layer_hint; + return {}; + } + virtual std::map memory_breakdown() const = 0; // diff --git a/llama.cpp/tools/server/server-context.cpp b/llama.cpp/tools/server/server-context.cpp --- a/llama.cpp/tools/server/server-context.cpp +++ b/llama.cpp/tools/server/server-context.cpp @@ -2408,14 +2408,63 @@ private: const int n_left = slot.prompt.n_tokens() - n_keep; const int n_discard = slot.task->params.n_discard ? slot.task->params.n_discard : (n_left / 2); + // opencoti F5 M1 rest-kv-eviction — see docs/features/advanced_kv.md + // Upstream discards the positional middle [n_keep, n_keep+n_discard). + // When rest-kv-eviction is on, slide that same-width window over the + // evictable region [n_keep, n_tokens - rest_kv_recent) and discard the + // window with the LOWEST retention sum (KeyDiff key-similarity) — keep the + // distinctive tokens, drop the redundant ones. It is still ONE contiguous + // window, so the seq_rm/seq_add + token-mirror rewrite below stay 1:1 with + // KV positions. Any failure (eviction off, non-KV memory, quantized keys, + // no room after protecting the recent tail) leaves w0 = n_keep, i.e. the + // byte-for-byte upstream positional discard. + // opencoti F5 M6 MTP-composition: scoring reads ctx_tgt (the authoritative + // full-model KV); the SAME window w0 is then applied to BOTH the target and + // the draft KV cache below, preserving target/draft position alignment for + // speculative decode. + int w0 = n_keep; + if (params_base.rest_kv_eviction && n_discard > 0) { + const int n_tokens = slot.prompt.n_tokens(); + const int recent = std::max(0, params_base.rest_kv_recent); + const int hi = n_tokens - recent - n_discard; // last valid window start (inclusive) + if (hi >= n_keep) { + std::vector scores(n_tokens, std::numeric_limits::infinity()); + const size_t got = llama_memory_seq_key_scores( + llama_get_memory(ctx_tgt), slot.id, params_base.rest_kv_layer, scores.data(), n_tokens); + if (got > 0) { + // prefix sums over [n_keep, n_tokens); an absent position gets a + // large finite penalty so a window covering one is never picked + // (avoids inf/inf arithmetic in the sliding sum). + const int m = n_tokens - n_keep; + std::vector pref(m + 1, 0.0); + for (int i = 0; i < m; ++i) { + const float sv = scores[n_keep + i]; + pref[i + 1] = pref[i] + (std::isfinite(sv) ? (double) sv : 1e30); + } + double best = std::numeric_limits::infinity(); + int best_w0 = n_keep; + for (int s = n_keep; s <= hi; ++s) { + const double sum = pref[s - n_keep + n_discard] - pref[s - n_keep]; + if (sum < best) { + best = sum; + best_w0 = s; + } + } + w0 = best_w0; + SLT_WRN(slot, "rest-kv eviction: discard [%d, %d) retention=%.3f (vs positional [%d, %d))\n", + w0, w0 + n_discard, best, n_keep, n_keep + n_discard); + } + } + } + SLT_WRN(slot, "slot context shift, n_keep = %d, n_left = %d, n_discard = %d\n", n_keep, n_left, n_discard); - common_context_seq_rm (ctx_tgt, slot.id, n_keep , n_keep + n_discard); - common_context_seq_add(ctx_tgt, slot.id, n_keep + n_discard, slot.prompt.n_tokens(), -n_discard); + common_context_seq_rm (ctx_tgt, slot.id, w0 , w0 + n_discard); + common_context_seq_add(ctx_tgt, slot.id, w0 + n_discard, slot.prompt.n_tokens(), -n_discard); if (ctx_dft) { - common_context_seq_rm (ctx_dft.get(), slot.id, n_keep , n_keep + n_discard); - common_context_seq_add(ctx_dft.get(), slot.id, n_keep + n_discard, slot.prompt.tokens.pos_next(), -n_discard); + common_context_seq_rm (ctx_dft.get(), slot.id, w0 , w0 + n_discard); + common_context_seq_add(ctx_dft.get(), slot.id, w0 + n_discard, slot.prompt.tokens.pos_next(), -n_discard); } // add generated tokens to cache @@ -2424,7 +2473,7 @@ private: GGML_ASSERT(!slot.prompt.tokens.has_mtmd); llama_tokens new_tokens = slot.prompt.tokens.get_tokens(); // copy - for (size_t i = n_keep + n_discard; i < new_tokens.size(); i++) { + for (size_t i = w0 + n_discard; i < new_tokens.size(); i++) { new_tokens[i - n_discard] = new_tokens[i]; }