diff --git a/CHANGELOG.md b/CHANGELOG.md index 8840402..e9a1aa7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -46,6 +46,15 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). - serve: `--dtype` and `server.dtype` set the same knob for every model the server loads. An unrecognized value is a config error rather than a silent fall back to the default. +- Prompt cache: hybrid-arch models keep an anchor checkpoint at the end + of the system prompt, so parallel requests sharing a system prompt and + tool schemas (subagent fan-out) start warm instead of re-prefilling + the shared prefix after the conversation deepens. GMLX_APC_CKPT_SYS=0 + disables. +- Prompt cache: exact-tier models (deepseek-v4-class pooling stacks) get + the same system-prompt anchor as a whole-prefix snapshot in its own + LRU, so sibling fan-out stays warm there too. GMLX_APC_ANCHOR_ENTRIES + and GMLX_APC_ANCHOR_BUDGET_MB bound it; GMLX_APC_CKPT_SYS=0 disables. ## [0.3.0] - 2026-08-09 diff --git a/docs/server-config.md b/docs/server-config.md index 3fa264c..e67888d 100644 --- a/docs/server-config.md +++ b/docs/server-config.md @@ -1110,12 +1110,21 @@ What a warm hit restores: per entry grows quadratically over a conversation. The checkpoint tier saves these models piecewise instead - near-linear memory, same warm TTFT. Along a long prefill it drops a restore point every - `GMLX_APC_CKPT_INTERVAL` tokens (default 4096), plus two targeted + `GMLX_APC_CKPT_INTERVAL` tokens (default 4096), plus three targeted ones: a replay checkpoint one token before the prompt end (recurrent state cannot rewind, so an identical resend needs a restore point - strictly below it) and a turn checkpoint at the longest prefix the + strictly below it), a turn checkpoint at the longest prefix the next turn's re-rendered history can actually replay (predicted from - the chat template, so thinking-strip divergence lands past it). What + the chat template, so thinking-strip divergence lands past it), and + an anchor checkpoint at the end of the system prompt. The anchor is + the fan-out one: requests that share a system prompt and tool schemas + but carry different user turns (parallel agents, subagent bursts) all + restore from it instead of re-prefilling the shared prefix, and it is + exempt from the pruning that otherwise keeps only the newest restore + points as a conversation deepens. Exact-tier models (deepseek-v4-class + pooling stacks) get the same anchor as a whole-prefix clone in its own + small LRU (`GMLX_APC_ANCHOR_ENTRIES`), where sibling churn through the + count-capped exact slots cannot evict it. What reuse each family gets from these: [performance.md](performance.md#the-prompt-cache). On ckpt-tier models prompt prefill runs one request at a time (batched prefill @@ -1157,6 +1166,10 @@ and store counts surface on the authed `GET /v1/metrics`. | `GMLX_APC_CKPT_REPLAY` | `0` disables the replay checkpoint (identical resends prefill cold again). | | `GMLX_APC_CKPT_REPLAY_MIN` | Minimum prompt tokens before a replay checkpoint is saved on recurrent (GDN) models (default `1024`; short prompts re-prefill cheaply and are not worth the >100 MB state snapshot). | | `GMLX_APC_CKPT_TURN` | `0` disables the turn checkpoint (next-turn reuse falls back to the interval grid). | +| `GMLX_APC_CKPT_SYS` | `0` disables the system-prompt anchor on both tiers (sibling requests sharing a system prompt prefill the shared prefix cold; on hybrid models they also fall back to the interval grid). | +| `GMLX_APC_CKPT_SYS_MIN` | Minimum tokens of shared system prefix before an anchor is saved (default `256`; raised to `GMLX_APC_CKPT_REPLAY_MIN` on recurrent models). A shorter shared prefix re-prefills in milliseconds and is not worth a record. | +| `GMLX_APC_ANCHOR_ENTRIES` | Exact-tier anchor LRU entries (default `4`). Exact-mode models (deepseek-v4-class pooling stacks) keep their system-prompt anchors here as whole-prefix clones, out of reach of the count-capped upstream exact LRU that every request's own store would churn. | +| `GMLX_APC_ANCHOR_BUDGET_MB` | Byte budget for the exact-tier anchor LRU, in MB (default `4096`). A deep shared prefix on a pooling stack clones to GBs; newest always survives. | | `GMLX_APC_CKPT_TRIPWIRE` | Requests before the dead-tier tripwires warn (default `5`; `0` silences both). | | `GMLX_APC_CKPT_RECORDS` | Checkpoint-record LRU entries (default `32`). | | `GMLX_APC_CKPT_BUDGET_MB` | Byte budget for checkpoint-record payload (recurrent states + KV tails), in MB (default `4096`). A GDN record can carry >100 MB of state and each request saves several checkpoints, so expect resident memory to grow toward this budget on hybrid models under sustained multi-turn traffic; lower it if 4 GB of cache is too much for your machine. | diff --git a/gmlx/cache_snapshot.py b/gmlx/cache_snapshot.py index b301d69..8a8c30d 100644 --- a/gmlx/cache_snapshot.py +++ b/gmlx/cache_snapshot.py @@ -355,6 +355,123 @@ def drafter_sidecar_lookup( return None +# Exact-tier anchor: one whole-prefix clone per system-prompt chain, held in +# a gmlx-owned side LRU. The upstream exact LRU is count-capped (2 slots by +# default) and every request writes its guard-column entry there, so sibling +# fan-out churns out the early shared-prefix entry the siblings need; this +# index holds nothing but anchors. +_ANCHOR_ENTRIES = max(1, env_int("GMLX_APC_ANCHOR_ENTRIES", 4)) +# Whole-prefix clones of pooling/MLA stacks run to GBs at deep prefixes, so +# the index is byte-bounded too; newest always survives. +_ANCHOR_BUDGET_BYTES = max( + 1, env_int("GMLX_APC_ANCHOR_BUDGET_MB", 4096)) << 20 + + +def _anchor_index(manager: Any) -> "OrderedDict": + with manager.lock: + idx = getattr(manager, "_kq_anchor_cache", None) + if idx is None: + idx = OrderedDict() + manager._kq_anchor_cache = idx + return idx + + +def anchor_exact_store( + manager: Any, + token_ids, + prompt_cache: list, + extra_hash: int = 0, +) -> bool: + """Store a whole-prefix clone of ``prompt_cache`` under the anchor key. + + Cloning goes through the upstream exact-clone path (the pooling arms are + installed there), so any stack the exact tier serves anchors identically. + A re-store under the same key replaces the entry in place, keeping one + anchor per chain. Best-effort; never raises. + """ + if manager is None or not prompt_cache: + return False + try: + ids = tuple(int(t) for t in token_ids) + if not ids: + return False + from mlx_vlm import apc as _apc + clones = _apc._clone_prompt_cache_for_apc(prompt_cache) + if clones is None: + _ckpt_decline(manager, "anchor_clone") + return False + nbytes = _caches_nbytes(clones) + key = (ids, int(extra_hash)) + idx = _anchor_index(manager) + with manager.lock: + idx[key] = (clones, nbytes) + idx.move_to_end(key) + while len(idx) > _ANCHOR_ENTRIES: + idx.popitem(last=False) + total = sum(n for _, n in idx.values()) + while total > _ANCHOR_BUDGET_BYTES and len(idx) > 1: + _, (_, n) = idx.popitem(last=False) + total -= n + _ckpt_bump(manager, "anchor_stores") + _log.info("APC anchor store: tokens=%d", len(ids)) + return True + except Exception: + _log.warning("APC anchor store failed; continuing", exc_info=True) + return False + + +def anchor_exact_lookup( + manager: Any, + token_ids, + extra_hash: int = 0, + min_prefix_tokens: int = 0, +) -> tuple: + """Longest anchor strictly prefixing ``token_ids`` on the same chain. + + Returns ``(warm_prompt_cache, p)`` or ``(None, 0)``. The warm list is a + fresh clone with capacity for the full query, decoupled from the stored + entry. Never raises. + """ + if manager is None or token_ids is None: + return None, 0 + try: + ids = tuple(int(t) for t in token_ids) + n = len(ids) + idx = _anchor_index(manager) + best_key = None + with manager.lock: + for key in idx: + kids, kh = key + p = len(kids) + if (kh != int(extra_hash) + or not min_prefix_tokens < p < n + or ids[:p] != kids): + continue + if best_key is None or p > len(best_key[0]): + best_key = key + if best_key is None: + return None, 0 + clones = idx[best_key][0] + idx.move_to_end(best_key) + from mlx_vlm import apc as _apc + warm = _apc._clone_prompt_cache_for_apc( + clones, min_capacity_tokens=n + 1) + if warm is None: + return None, 0 + p = len(best_key[0]) + with manager.lock: + # Anchor hits are cache-served tokens the upstream ledger never + # sees; bumping both keeps token_hit_rate honest. + manager.stats.hits += 1 + manager.stats.matched_tokens += p + _ckpt_bump(manager, "anchor_hits") + _log.info("APC anchor hit: prefix=%d", p) + return warm, p + except Exception: + _log.warning("APC anchor lookup failed; continuing", exc_info=True) + return None, 0 + + # Rotating (sliding-window) snapshot/restore inverse. A live RotatingKVCache # buffer has three regimes (contiguous concat-mode, padded in-place growth, # rotated ring); the snapshot canonicalizes all three into temporal order at @@ -504,7 +621,8 @@ def rotating_invariant(cache): _BOUNDED_SALT = 0x3A_91_C7_55_0E_D4_26 _CKPT_RECORD_ENTRIES = max(2, env_int("GMLX_APC_CKPT_RECORDS", 32)) -# Strip-on-extend: newest N restorable checkpoints per chain. +# Strip-on-extend: newest N restorable checkpoints per chain, plus the +# chain's anchor (see _record_insert). _CKPT_HEAVY_PER_CHAIN = max(1, env_int("GMLX_APC_CKPT_HEAVY", 2)) # Byte budget for record-owned payload (recurrent states + KV tails; chain # blocks are bounded by the manager pool). A GDN record can carry >100 MB @@ -603,10 +721,12 @@ def _ckpt_block_prefix(p: int, block_size: int) -> int: class _CkptRecord: - # kind in {"boundary", "replay", "retire"}: prefill-cursor boundaries, - # the N-1 identical-replay record, and retirement stores. Retention - # differs only in the strip-on-extend exemption (see _record_insert); - # the entries cap and byte budget treat all kinds alike. + # kind in {"boundary", "anchor", "replay", "retire"}: prefill-cursor + # boundaries, the chain's pinned early boundary (sibling fan-out + # reuse), the N-1 identical-replay record, and retirement stores. + # Retention differs only in _record_insert's strip-on-extend + # exemptions and eviction order; adoption gates only ever test for + # "replay". __slots__ = ("ids", "extra_hash", "p", "b_full", "layout", "main_blocks", "bounded_blocks", "rot_meta", "states", "tails", "nbytes", "kind") @@ -636,6 +756,7 @@ def _ckpt_records(manager) -> "OrderedDict": "ckpt_stores", "ckpt_hits", "ckpt_matched_tokens", "ckpt_missed_adoptions", "ckpt_skeleton_writes", "sidecar_writes", "retire_fallback_full", "ckpt_pool_evictions", + "anchor_stores", "anchor_hits", ) @@ -695,6 +816,9 @@ def ckpt_reset(manager) -> None: side = getattr(manager, "_kq_sidecar_cache", None) if side: side.clear() + anchors = getattr(manager, "_kq_anchor_cache", None) + if anchors: + anchors.clear() manager._kq_ckpt_stats = None @@ -791,7 +915,10 @@ def _evict_for_pool(manager, deficit: int) -> int: in-flight lookup pin or another request's chain do not free on release, so an unreachable deficit is detected upfront and evicts nothing rather than draining the whole index for a store that still - declines. Returns the number of records released.""" + declines. Plain LRU order on purpose: anchor records get no + protection here, since an anchor pinning window blocks on an + exhausted pool would starve every future store. Returns the number + of records released.""" if not hasattr(manager, "_free_head"): return 0 idx = _ckpt_records(manager) @@ -821,6 +948,22 @@ def _evict_for_pool(manager, deficit: int) -> int: return evicted +def _evict_lru_record(manager, idx, keep): + """Release and return the eviction victim: the least-recently-used + non-anchor record, else the least-recently-used anchor, never the + ``keep`` key (the record being inserted always survives). Lookup + hits move_to_end, so anchor order is LRU by last hit: an anchor + that never serves a sibling ages out, one that does stays hot. + Caller holds the lock and guarantees a non-keep record exists.""" + key = next((k for k, r in idx.items() + if r.kind != "anchor" and k != keep), None) + if key is None: + key = next(k for k in idx if k != keep) + victim = idx.pop(key) + _release_record(manager, victim) + return victim + + def _record_insert(manager, rec) -> None: """Insert a checkpoint record: strip-on-extend superseded records on the same chain immediately (LRU would keep exactly the wrong ones - the @@ -828,18 +971,31 @@ def _record_insert(manager, rec) -> None: then bound the index by count and payload bytes, releasing refs on everything dropped. The newest record always survives. - Replay records are exempt from the strip only: growth (a longer - boundary or retirement insert) is exactly the moment an identical - resend still needs the N-1 record, so a longer non-replay insert - never releases one; a newer replay on the same chain supersedes it. - The count and byte bounds below evict replay records normally -- the - exemption pins nothing against real memory pressure.""" + Two strip exemptions. Replay: growth (a longer boundary or + retirement insert) is exactly the moment an identical resend still + needs the N-1 record, so a longer non-replay insert never releases + one; a newer replay on the same chain supersedes it. Anchor: the + chain's early boundary is what sibling fan-out requests (shared + system prompt, disjoint user turns) can adopt, and it is exactly + the record strip-on-extend removes first as the chain deepens. The + schedule tags the system-prefix stop "anchor"; when no tagged stop + exists (no render ctx, completions API), the first restorable + boundary on a fresh chain is promoted instead. One anchor per + chain: an anchor insert supersedes tagged anchors below it. + + Against real memory pressure the exemptions pin little: the count + and byte bounds evict anchors after non-anchors (LRU by last hit), + and pool-pressure eviction (_evict_for_pool) gives anchors no + protection at all: position-salted window chains must never sit + pinned on an exhausted pool.""" idx = _ckpt_records(manager) key = (rec.ids, rec.extra_hash) rec.nbytes = _rec_nbytes(rec) with manager.lock: old = idx.pop(key, None) if old is not None: + if old.kind == "anchor" and rec.kind == "boundary": + rec.kind = "anchor" # a re-store keeps the anchor tag _release_record(manager, old) # chain = records whose ids are a strict prefix of this one chain = [k for k, r in idx.items() @@ -849,20 +1005,25 @@ def _record_insert(manager, rec) -> None: # One replay per chain: the newer one supersedes outright. for k in [k for k in chain if idx[k].kind == "replay"]: _release_record(manager, idx.pop(k)) - chain = [k for k in chain if k in idx and idx[k].kind != "replay"] + elif rec.kind == "anchor": + for k in [k for k in chain if idx[k].kind == "anchor"]: + _release_record(manager, idx.pop(k)) + elif rec.kind == "boundary" and not any( + idx[k].kind != "replay" for k in chain if k in idx): + rec.kind = "anchor" # first restorable boundary + chain = [k for k in chain + if k in idx and idx[k].kind not in ("replay", "anchor")] chain.sort(key=lambda k: idx[k].p, reverse=True) for k in chain[_CKPT_HEAVY_PER_CHAIN - 1:]: _release_record(manager, idx.pop(k)) idx[key] = rec idx.move_to_end(key) while len(idx) > _CKPT_RECORD_ENTRIES: - _, victim = idx.popitem(last=False) - _release_record(manager, victim) + _evict_lru_record(manager, idx, key) total = sum(int(getattr(r, "nbytes", 0) or 0) for r in idx.values()) while total > _CKPT_BUDGET_BYTES and len(idx) > 1: - _, victim = idx.popitem(last=False) + victim = _evict_lru_record(manager, idx, key) total -= int(getattr(victim, "nbytes", 0) or 0) - _release_record(manager, victim) def ckpt_store( @@ -1241,9 +1402,11 @@ def ckpt_lookup( """Longest checkpoint-tier warm start for ``token_ids``. Walks pinned records p-descending testing the whole conjunction, then - falls back to the disk skeleton (restart repair). ``layout`` rejects - records from a different per-layer layout. Returns - ``(warm_prompt_cache, p)`` or ``(None, 0)``. Never raises. + falls back to the disk skeleton (restart repair). When the deepest + pinned candidate is the chain's anchor, the disk tier is consulted + first for anything strictly deeper. ``layout`` rejects records from a + different per-layer layout. Returns ``(warm_prompt_cache, p)`` or + ``(None, 0)``. Never raises. """ if manager is None or token_ids is None: return None, 0 @@ -1283,6 +1446,19 @@ def ckpt_lookup( if gated: _ckpt_decline(manager, "replay_gate") cands.sort(key=lambda r: r.p, reverse=True) + if cands and cands[0].kind == "anchor": + # The anchor is a retention floor, not a depth ceiling. It + # sits early on the chain by construction, and the pinned + # walk below returns on first success, so without this the + # anchor would cap every divergent query at its own p while + # a deeper skeleton sits on disk (the position strip-on- + # extend drops from memory but the skeleton keeps). + warm, p = _ckpt_disk_lookup( + manager, ids, extra_hash=extra_hash, + min_prefix_tokens=max(min_prefix_tokens, cands[0].p), + layout=layout) + if warm is not None: + return warm, p for rec in cands: # Assembly runs unlocked (it concatenates and evals block # tensors), so the record's chains must be pinned against a diff --git a/gmlx/retire_key.py b/gmlx/retire_key.py index 9da0bd8..f79de2c 100644 --- a/gmlx/retire_key.py +++ b/gmlx/retire_key.py @@ -215,22 +215,17 @@ def build_assistant_message(ctx: dict, full_text: str) -> dict: return msg -def predict_next_ids(ctx: dict, assistant_msg: dict | None) -> list[int] | None: - """Render and tokenize the hypothetical next-turn prefix. - - ``assistant_msg=None`` renders the request's own messages with no - reply appended -- the render-stable core a next turn extends. One - predictor for both questions; a second renderer diverging silently - is how the turn-boundary bug stayed invisible.""" +def _render_ids(ctx: dict, msgs: list) -> list[int] | None: + """Render ``msgs`` through the request's own template and tokenizer + (no generation prompt) and return the token ids. The single render + path behind every prefix prediction; a second renderer diverging + silently is how the turn-boundary bug stayed invisible.""" render = ctx.get("render") preprocess = ctx.get("preprocess") if render is None or preprocess is None: return None kw = dict(ctx.get("kw") or {}) kw["add_generation_prompt"] = False - msgs = list(ctx["messages"]) - if assistant_msg is not None: - msgs.append(assistant_msg) text = render(ctx["processor"], ctx["config"], msgs, **kw) if not isinstance(text, str): return None @@ -243,6 +238,71 @@ def predict_next_ids(ctx: dict, assistant_msg: dict | None) -> list[int] | None: return [int(t) for t in ids] +def predict_next_ids(ctx: dict, assistant_msg: dict | None) -> list[int] | None: + """Render and tokenize the hypothetical next-turn prefix. + + ``assistant_msg=None`` renders the request's own messages with no + reply appended -- the render-stable core a next turn extends.""" + msgs = list(ctx["messages"]) + if assistant_msg is not None: + msgs.append(assistant_msg) + return _render_ids(ctx, msgs) + + +def _lcp_len(seq, nxt) -> int: + n = min(len(seq), len(nxt)) + lcp = 0 + while lcp < n and int(seq[lcp]) == int(nxt[lcp]): + lcp += 1 + return lcp + + +def system_prefix_lcp(ctx: dict, prompt_ids) -> int | None: + """Token length of the prefix every sibling request shares: the LCP + of two probe renders that differ only in a dummy first user turn. + + Sibling fan-out requests share the system prompt and tool schemas + but diverge at the first user message; this is the deepest position + one checkpoint can serve all of them. A system-only render cannot + measure it on every template (gemma folds the system prompt into + the first user turn and renders a lone system message to almost + nothing), so the offset comes from a divergence probe instead: + render leading-system + user "0" and leading-system + user "1" + through the same template, kwargs (tools included), and tokenizer + as the live request, and take where they split. A folding template + folds both probes identically, so the split lands exactly where + real siblings diverge. Clamped by the LCP with the live + prompt in case a template leaks user content into the header. + Memoized on the ctx; media prompts return None (expanded media ids + cannot be re-encoded from text).""" + if not ctx or ctx.get("media"): + return None + memo = ctx.get("_p_system") + if memo is not None: + return memo if memo >= 0 else None + lcp = None + try: + msgs = list(ctx.get("messages") or ()) + lead = [] + for m in msgs: + if isinstance(m, dict) and m.get("role") == "system": + lead.append(m) + else: + break + # No leading system block, or nothing after it (a system-only + # prompt is covered by the terminal checkpoint already). + if lead and len(lead) < len(msgs): + p1 = _render_ids(ctx, lead + [{"role": "user", "content": "0"}]) + p2 = _render_ids(ctx, lead + [{"role": "user", "content": "1"}]) + if p1 and p2: + lcp = min(_lcp_len(p1, p2), _lcp_len(p1, prompt_ids)) + except Exception: + _log.debug("system-prefix prediction failed", exc_info=True) + lcp = None + ctx["_p_system"] = -1 if lcp is None else lcp + return lcp + + # Model types whose prompt-stable prediction already failed once: the # first failure warns with the traceback, repeats log at debug -- one # broken template must not stack a warning onto every request. Keyed by diff --git a/gmlx/spec_engine.py b/gmlx/spec_engine.py index 7ce6f6d..732f04b 100644 --- a/gmlx/spec_engine.py +++ b/gmlx/spec_engine.py @@ -318,6 +318,22 @@ def _l1_lookup_and_arm_store(batch, manager, mode, l0_prefix) -> int: manager.release(blocks) blocks = [] warm, prefix_len, tier = cw, cp, "ckpt" + elif mode == "exact": + # Exact-tier anchor: the shared-system-prefix clone in the + # gmlx anchor LRU wins only when strictly longer than the + # stock exact pick. Media guards mirror the stock probe. + from .cache_snapshot import anchor_exact_lookup + min_p = max(prefix_len, + view._apc_safe_prefix_lookup_min(ids_list)) + aw, ap = anchor_exact_lookup( + manager, ids_list, extra_hash=extra_hash, + min_prefix_tokens=min_p) + if (aw is not None and ap > prefix_len + and view._apc_suffix_is_text_only(ids_list, ap)): + if blocks: + manager.release(blocks) + blocks = [] + warm, prefix_len, tier = aw, ap, "anchor" if warm and 0 < prefix_len < len(ids_list): batch.prompt_cache = warm # Matched blocks stay acquired until the stock post-prefill @@ -373,6 +389,9 @@ def _l1_lookup_and_arm_store(batch, manager, mode, l0_prefix) -> int: batch._kq_ckpt_armed = True from .cache_snapshot import ckpt_note_armed ckpt_note_armed(manager) + elif mode == "exact": + _exact_anchor_arm(batch, meta, guard, + max(l0_prefix, l1_prefix)) return l1_prefix @@ -496,19 +515,145 @@ def _ckpt_turn_boundaries(batch, meta, restored: int, return out -def _sched_insert(bounds: list, pos: int, kind: str) -> None: +def _ckpt_sys_boundary(batch, meta, restored: int, + block_size: int) -> int | None: + """Anchor stop at the end of the shared system prefix. + + Sibling fan-out requests share the system prompt and tool schemas + and diverge at the first user message, generally between grid + points, so the interval schedule alone wastes up to one interval of + sibling recompute, and strip-on-extend removes the early boundary + the siblings need as the chain deepens (the anchor exemption in + _record_insert keeps this one). arr layouts snap the stop down to + the chunk grid (off-grid chunking drifts GDN state) and keep the + replay byte floor (recurrent state is prompt-length-independent, so + a tiny anchor costs the same >100 MB clone as a deep one); + attention layouts snap to the block grid, which also satisfies the + rotating store's below-window grid gate. GMLX_APC_CKPT_SYS=0 + disables; GMLX_APC_CKPT_SYS_MIN floors the position (a sub-floor + shared prefix re-prefills in milliseconds and is not worth a + record). + """ + if env_int("GMLX_APC_CKPT_SYS", 1) == 0: + return None + ids = meta.get("full_input_ids") or () + tags = _ckpt_layout_for(getattr(batch, "model", None), block_size) or () + floor_min = max(block_size, env_int("GMLX_APC_CKPT_SYS_MIN", 256)) + if "arr" in tags: + floor_min = max(floor_min, + env_int("GMLX_APC_CKPT_REPLAY_MIN", 1024)) + # Below the floor no anchor can land; skip the render+tokenize + # prediction entirely (same rule as the turn boundaries). + if len(ids) - 1 < floor_min: + return None + from .retire_key import lookup_render_ctx, system_prefix_lcp + ctx = lookup_render_ctx(ids) + lcp = system_prefix_lcp(ctx, ids) if ctx else None + if not lcp: + return None + unit = _ckpt_unit(batch, block_size) if "arr" in tags else block_size + pos = (min(int(lcp), len(ids) - 1) // unit) * unit + if pos < floor_min or pos <= max(0, restored): + return None + meta["ckpt_sys_bound"] = pos + return pos + + +def _exact_anchor_boundary(batch, meta, guard: int, + restored: int) -> int | None: + """Anchor position for an exact-tier (non-ckpt) model: the sibling + divergence point, ungridded (exact clones restore at any position). + Clamped to the stock guard column, so the prefill pauses at most + twice: once for the anchor, once for the stock guard store. + GMLX_APC_CKPT_SYS=0 disables (one switch for both tiers); + GMLX_APC_CKPT_SYS_MIN floors the position (a sub-floor shared + prefix re-prefills in milliseconds and is not worth a clone). + """ + if env_int("GMLX_APC_CKPT_SYS", 1) == 0: + return None + ids = meta.get("full_input_ids") or () + floor_min = max(2, env_int("GMLX_APC_CKPT_SYS_MIN", 256)) + if len(ids) - 1 < floor_min: + return None + from .retire_key import lookup_render_ctx, system_prefix_lcp + ctx = lookup_render_ctx(ids) + lcp = system_prefix_lcp(ctx, ids) if ctx else None + if not lcp: + _log.info("APC anchor declined: no measurable system prefix " + "(render ctx %s)", "present" if ctx else "missing") + return None + pos = min(int(lcp), len(ids) - 1) + if guard > 0: + pos = min(pos, guard) + if pos < floor_min or pos <= max(0, restored): + return None + return pos + + +def _exact_anchor_arm(batch, meta, guard: int, restored: int) -> None: + """Schedule the anchor pause by mirroring its position into + ``checkpoint_len`` (the key the stock column truncation reads). + ``_exact_anchor_store`` hands the column back to the stock guard + after firing, so the stock store still runs exactly as unarmed.""" + pos = _exact_anchor_boundary(batch, meta, guard, restored) + if pos is None: + return + meta["anchor_len"] = pos + meta["anchor_guard"] = guard + if pos != guard: + meta["checkpoint_len"] = pos + batch._kq_anchor_armed = True + _log.info("APC anchor armed: pos=%d guard=%d", pos, guard) + + +def _exact_anchor_store(batch) -> None: + """Anchor store for exact-tier models: one whole-prefix clone at the + sibling divergence, into the gmlx anchor LRU. Runs from the wrapped + stock store immediately before the stock body; after firing it + restores ``checkpoint_len`` to the stock guard column without + latching ``checkpoint_done``, so the stock guard store (and its + latch) fire untouched.""" + manager = getattr(batch, "_apc_manager", None) + meta_list = getattr(batch, "_apc_meta", None) or [] + if manager is None or not meta_list or meta_list[0] is None: + return + meta = meta_list[0] + pos = int(meta.get("anchor_len") or 0) + if pos <= 0 or meta.get("anchor_done"): + return + if batch._row_real_tokens_processed(0) != pos: + return + meta["anchor_done"] = True + guard = int(meta.get("anchor_guard") or 0) + if int(meta.get("checkpoint_len") or 0) == pos and pos != guard: + meta["checkpoint_len"] = guard + cache = batch._apc_prompt_cache_for_store(0) + if cache is None: + return + from .cache_snapshot import anchor_exact_store + anchor_exact_store(manager, meta["full_input_ids"][:pos], cache, + extra_hash=int(meta.get("extra_hash", 0))) + + +def _sched_insert(bounds: list, pos: int, kind: str, *, + upgrade: bool = False) -> None: """Insert (pos, kind) keeping order. On collision the existing entry keeps its kind: a colliding position is always grid-aligned or an exact turn boundary, where a plain boundary record adopts freely -- identical resend included -- while flipping it to replay would gate turn-2 and branch adoption out on recurrent layouts (and satisfy the - p=N drop with a record turn 2 cannot use).""" + p=N drop with a record turn 2 cannot use). ``upgrade`` lets an + anchor replace a plain boundary at the same position (strictly more + retention, same free adoption), never a replay.""" import bisect pts = [b for b, _ in bounds] i = bisect.bisect_left(pts, pos) - if i >= len(pts) or pts[i] != pos: - bounds.insert(i, (pos, kind)) + if i < len(pts) and pts[i] == pos: + if upgrade and bounds[i][1] == "boundary": + bounds[i] = (pos, kind) + return + bounds.insert(i, (pos, kind)) def _ckpt_arm_schedule(batch, meta, guard: int, restored: int, @@ -525,8 +670,13 @@ def _ckpt_arm_schedule(batch, meta, guard: int, restored: int, turn = _ckpt_turn_boundaries(batch, meta, restored, block_size) for pos in turn: _sched_insert(bounds, pos, "boundary") + sysb = _ckpt_sys_boundary(batch, meta, restored, block_size) + if sysb is not None: + _sched_insert(bounds, sysb, "anchor", upgrade=True) replay = _ckpt_replay_boundary(batch, meta, restored, block_size) if replay is not None: + # Colliding with the anchor keeps the anchor (default no-upgrade): + # it adopts identical resends freely, replay semantics add nothing. _sched_insert(bounds, replay, "replay") meta["ckpt_boundaries"] = bounds meta["checkpoint_len"] = int(bounds[0][0]) if bounds else 0 @@ -597,7 +747,8 @@ def _install_ckpt_checkpoint_store() -> None: and the stock prompt_step call the stock method, so one wrap covers both paths). The cursor's advance of ``checkpoint_len`` is what suppresses the stock store -- wrapping makes that ordering - structural. Idempotent.""" + structural. Exact-tier anchor batches ride the same wrap with their + own single-stop hook. Idempotent.""" from mlx_vlm.generate.ar import PromptProcessingBatch if getattr(PromptProcessingBatch._store_apc_exact_checkpoints, _CKPT_STORE_FLAG, False): @@ -607,6 +758,8 @@ def _install_ckpt_checkpoint_store() -> None: def _store_with_ckpt_cursor(self): if getattr(self, "_kq_ckpt_armed", False): _ckpt_mid_prefill_store(self) + elif getattr(self, "_kq_anchor_armed", False): + _exact_anchor_store(self) _orig(self) _store_with_ckpt_cursor.__dict__[_CKPT_STORE_FLAG] = True @@ -716,6 +869,96 @@ def _plain_ckpt_init(batch) -> None: } +def _plain_anchor_init(batch) -> None: + """Arm the exact-tier anchor stop on a stock prompt batch (non-ckpt + exact models: DeepSeek-V4-class pooling stacks). + + Restores come from the admission pick (_install_exact_anchor_pick), + so this only schedules the store. Warm and right-padded rows are + included: a restored prefix is usually far short of the divergence + (a bare bos match off some unrelated request), and upstream's + checkpoint column and row extraction handle both shapes. Refusing + them would skip every row that rides a warm batch, which on a busy + server is nearly all of them. The restored prefix becomes the + boundary floor, so a row already past the divergence arms nothing. + """ + manager = getattr(batch, "_apc_manager", None) + mode = getattr(batch, "_apc_mode", None) + meta_list = getattr(batch, "_apc_meta", None) or [] + if (manager is None or mode != "exact" or len(meta_list) != 1 + or meta_list[0] is None or len(batch.uids) != 1 + or batch._inputs_embeds is None): + return + if _ckpt_active(batch.model, mode, int(manager.block_size)): + return # ckpt tier owns these models + meta = meta_list[0] + if len(meta.get("full_input_ids") or ()) < 2: + return + _exact_anchor_arm(batch, meta, int(meta.get("checkpoint_len") or 0), + int(meta.get("prefix_len") or 0)) + + +_ANCHOR_PICK_FLAG = "_kq_exact_anchor_pick" + + +def _install_exact_anchor_pick() -> None: + """Consult the anchor LRU inside the stock admission pick. + + The pick is where a warm prefix belongs: admission builds the batch + from it (suffix rows, right padding, warm-cache merge) and every + downstream path treats an anchor restore exactly like a stock exact + one. The anchor wins only when strictly longer than the stock pick, + so it never shortens a restore. Idempotent. + """ + from mlx_vlm.generate.ar import BatchGenerator + if getattr(BatchGenerator._apc_pick_for, _ANCHOR_PICK_FLAG, False): + return + _orig = BatchGenerator._apc_pick_for + + def _pick_with_anchor(self, sequence): + pick = _orig(self, sequence) + try: + if _SPEC_APC_DISABLED or getattr(self, "apc_mode", None) != "exact": + return pick + manager = getattr(self, "apc_manager", None) + if manager is None or _ckpt_active( + getattr(self, "model", None), "exact", + int(manager.block_size)): + return pick + _uid, ids_list, _mt, prompt_kwargs, _lps, _crit = sequence + if not ids_list or len(ids_list) < 2: + return pick + have = int((pick or {}).get("prefix_len") or 0) + extra_hash = self._apc_extra_hash(prompt_kwargs or {}) + floor = max(have, self._apc_safe_prefix_lookup_min(ids_list)) + from .cache_snapshot import anchor_exact_lookup + warm, ap = anchor_exact_lookup( + manager, ids_list, extra_hash=extra_hash, + min_prefix_tokens=floor) + if warm is None or ap <= have or ap >= len(ids_list): + return pick + if not self._apc_suffix_is_text_only(ids_list, ap): + return pick + if pick and pick.get("matched_blocks"): + manager.release(pick["matched_blocks"]) + _log.info("APC L1 hit: prefix=%d suffix=%d tier=anchor", + ap, len(ids_list) - ap) + return { + "matched_blocks": [], + "warm_cache": warm, + "prefix_len": ap, + "extra_hash": extra_hash, + "full_input_ids": list(ids_list), + } + except Exception: + _log.warning("APC anchor pick failed; continuing", + exc_info=True) + return pick + + _pick_with_anchor.__dict__[_ANCHOR_PICK_FLAG] = True + BatchGenerator._apc_pick_for = _pick_with_anchor + + _PLAIN_DECODE_FLAG = "_kq_ckpt_plain_decode" @@ -973,6 +1216,7 @@ def install_full_prompt_mtp_prefill() -> None: _install_apc_manager_stash() _install_ckpt_checkpoint_store() _install_plain_ckpt_decode() + _install_exact_anchor_pick() if getattr(PromptProcessingBatch, _FULL_PREFILL_FLAG, False): return @@ -1006,6 +1250,7 @@ def _mtp_init(self, *args, **kwargs) -> None: and not _SPEC_APC_DISABLED: try: _plain_ckpt_init(self) + _plain_anchor_init(self) except Exception: _log.warning("APC plain ckpt init failed; continuing " "stock", exc_info=True) diff --git a/tests/e2e/run_apc_depth_e2e.py b/tests/e2e/run_apc_depth_e2e.py index 2cc3a05..4d42f3a 100644 --- a/tests/e2e/run_apc_depth_e2e.py +++ b/tests/e2e/run_apc_depth_e2e.py @@ -63,6 +63,10 @@ All requests carry a system message unless --no-system (the system-render path is part of the shared prefix, so every reuse floor exercises it). +--system-words N swaps the one-liner for an agent-shaped system prompt of +N words of tool and policy blocks, the shape the system-prompt anchor +serves: the shared prefix then ends at the system boundary instead of +sitting in the user turn, and the divergent request must restore it. --speculative adds MTP to every server (spec-path ckpt arming + sidecar). `--warm-factor`/`--restart-factor` loosen the wall-clock collapse thresholds (ckpt restores clone >100 MB of GDN state; default 0.6x cold). Wall-clock @@ -136,6 +140,37 @@ "the maintenance log you are given." ) + +def agent_system(target_words: int) -> str: + """Coding-agent-shaped system prompt: identity plus numbered tool and + policy blocks, deterministic and non-repetitive (~1.3 tok/word). + + The shape is what the system-prompt anchor keys on. Agent fan-out + shares thousands of tokens of system prompt and tool schemas and + diverges at the first user message, so the reusable prefix ends at + the system boundary. A harness carrying its shared text in the user + turn arms no anchor at all. + """ + out = [SYSTEM_MSG, "\n\nTools and policies:\n"] + words = len(SYSTEM_MSG.split()) + i = 0 + while words < target_words: + t = _TOPICS[i % len(_TOPICS)] + s = ( + f"Tool {i} ({t}_probe): reads the {t} channel and returns " + f"{(i * 7) % 97} fields; call it at most {i % 4 + 1} times " + f"per turn, never before the {_TOPICS[(i + 3) % len(_TOPICS)]} " + f"check, and log the result under key K{(i * 31) % 4093}. " + f"Policy {i}: when the {t} reading drifts past " + f"{(i * 13) % 977}, escalate to the duty engineer and cite " + f"entry {max(0, i - 4)}.\n" + ) + out.append(s) + words += len(s.split()) + i += 1 + return "".join(out) + + # Session-phase system message: tool output is part of the conversation, # so the strict answer-from-the-log framing would invite refusals. SESSION_SYSTEM = ( @@ -439,6 +474,15 @@ def main() -> int: action="store_true", help="omit the system message (for templates that reject the role)", ) + ap.add_argument( + "--system-words", + type=int, + default=0, + metavar="N", + help="agent-shaped system prompt of N words of tool and policy " + "blocks (~1.3 tok/word) shared by every request, instead of the " + "one-line default; the shape the system-prompt anchor serves", + ) ap.add_argument( "--no-tripwire", action="store_true", @@ -550,8 +594,12 @@ def content_clean(content: str) -> bool: f"{f2['reading']} units by {f2['technician']}" ) - def mk_msgs(user_content: str) -> list: - msgs = [] if a.no_system else [{"role": "system", "content": SYSTEM_MSG}] + system_msg = agent_system(a.system_words) if a.system_words > 0 \ + else SYSTEM_MSG + + def mk_msgs(user_content: str, system: str | None = None) -> list: + system = system_msg if system is None else system + msgs = [] if a.no_system else [{"role": "system", "content": system}] msgs.append({"role": "user", "content": user_content}) return msgs @@ -610,8 +658,13 @@ def settle(key: str, timeout: float = 15.0) -> int: base = srv.base_url mid = model_id_of(base, srv) # absorb first-request one-time costs (kernel warmup) so wall-time - # comparisons below measure caching, not JIT - chat(base, mid, mk_msgs("Say ok."), max_tokens=4) + # comparisons below measure caching, not JIT. The warmup stays off + # the shared chain: carrying the run's system prompt behind a + # two-word user turn would put its own guard-column store at the + # system boundary, which then serves every later sibling and hides + # whether the anchor works at all. + chat(base, mid, [{"role": "user", "content": "Say ok."}], + max_tokens=4) # -- populate ------------------------------------------------------ st, text, content, ptok, cold_wall = chat( @@ -712,6 +765,42 @@ def settle(key: str, timeout: float = 15.0) -> int: f"{div_wall:.1f}s", ) + # -- anchor (only with --system-words) ------------------------------ + # A shared agent-shaped system prompt makes the divergent request a + # sibling: it must restore the system block from the anchor stored + # during the cold request, on every tier. Without --system-words the + # divergence sits a few tokens in and no anchor arms at all, which + # is why this gate is opt-in rather than always on. + if a.system_words > 0 and not a.no_system: + st_all = stats(base) + anchored = int(st_all.get("anchor_stores", 0) or 0) + int( + st_all.get("ckpt_stores", 0) or 0 + ) + # ~1.3 tok/word, less the template scaffolding and the grid or + # chunk snap the ckpt tier applies below the boundary. + anchor_floor = int(a.system_words * 0.6) + rep.check( + "anchor.armed", + anchored > 0, + f"anchor_stores={st_all.get('anchor_stores', 0)} " + f"ckpt_stores={st_all.get('ckpt_stores', 0)}", + ) + rep.check( + "anchor.divergent_adopted", + dm >= anchor_floor, + f"divergent matched +{dm} >= system-anchor floor " + f"{anchor_floor} ({a.system_words} system words)", + ) + served = int(st_all.get("anchor_hits", 0) or 0) + int( + st_all.get("ckpt_hits", 0) or 0 + ) + rep.check( + "anchor.served", + served > 0, + f"anchor_hits={st_all.get('anchor_hits', 0)} " + f"ckpt_hits={st_all.get('ckpt_hits', 0)}", + ) + # -- turns --------------------------------------------------------- # A real conversation through the real chat template: the # production render_ctx path unit tests fake with synthetic ids. @@ -1501,7 +1590,10 @@ def sess_facts(txt: str) -> bool: # fire. The prompt must tokenize under unit plus guard (no adoptable # terminal boundary) yet past the rotating window (windows over ~1300 # tokens need --no-tripwire). Disk stays off so a retire skeleton - # cannot serve the resend. + # cannot serve the resend, and the system message stays short even + # under --system-words: an agent-shaped system block arms the anchor, + # which then serves the resends exactly as designed and leaves this + # phase with nothing to refuse. if ck and not a.no_tripwire: trip_prefix, n_trip = deep_prefix(1050) trip_q = ( @@ -1527,7 +1619,9 @@ def sess_facts(txt: str) -> bool: sts = [] for _ in range(4): # populate + 3 refused resends st, _, _, _, _ = chat( - base, mid, mk_msgs(trip_prefix + trip_q), max_tokens=32 + base, mid, + mk_msgs(trip_prefix + trip_q, system=SYSTEM_MSG), + max_tokens=32, ) sts.append(st) settle("ckpt_stores") @@ -1565,6 +1659,7 @@ def sess_facts(txt: str) -> bool: "session_final_ptok": sess_final_ptok, "speculative": a.speculative, "system_message": not a.no_system, + "system_words": a.system_words or None, "template_kwargs": template_kwargs, "cold_wall_s": round(cold_wall, 1), "warm_wall_s": round(warm_wall, 1), diff --git a/tests/test_ckpt_cursor.py b/tests/test_ckpt_cursor.py index de9dbac..a89ebe9 100644 --- a/tests/test_ckpt_cursor.py +++ b/tests/test_ckpt_cursor.py @@ -208,8 +208,11 @@ def _boom(ids): raise AssertionError("render ctx consulted below the arm floor") monkeypatch.setattr(retire_key, "lookup_render_ctx", _boom) - assert _arm_meta(2048, ("kv", "arr"))["ckpt_p_stable_bounds"] == [] - assert _arm_meta(400, ("rot:512:0", "kv"))["ckpt_p_stable_bounds"] == [] + # Floors: turn needs one chunk (arr) or the window (rot-only); the + # system anchor floors at GMLX_APC_CKPT_SYS_MIN (256), raised to the + # replay byte floor (1024) on arr layouts. + assert _arm_meta(1000, ("kv", "arr"))["ckpt_p_stable_bounds"] == [] + assert _arm_meta(250, ("rot:512:0", "kv"))["ckpt_p_stable_bounds"] == [] # At the floor the prediction runs again. _stub_p_stable(monkeypatch, 550) meta = _arm_meta(600, ("rot:512:0", "kv")) @@ -242,10 +245,10 @@ def _armed_batch(man, ids, first, terminal, interval): def test_cursor_advances_and_latches(): - man = APCManager(num_blocks=64, block_size=16) - ids = list(range(500, 500 + 120)) - batch, meta = _armed_batch(man, ids, first=32, terminal=96, interval=32) - for boundary in (32, 64, 96): + man = APCManager(num_blocks=96, block_size=16) + ids = list(range(500, 500 + 160)) + batch, meta = _armed_batch(man, ids, first=32, terminal=128, interval=32) + for boundary in (32, 64, 96, 128): assert meta["checkpoint_len"] == boundary assert not meta.get("checkpoint_done") batch.prompt_cache = make_hybrid_cache(boundary, seed=boundary) @@ -254,13 +257,12 @@ def test_cursor_advances_and_latches(): se._ckpt_mid_prefill_store(batch) assert meta["ckpt_last_stored"] == boundary assert meta.get("checkpoint_done") is True - # Strip-on-extend keeps the newest two boundaries hittable; the - # superseded first boundary is stripped by design. - for boundary in (64, 96): - warm, got = ckpt_lookup(man, ids[:boundary + 8], extra_hash=7) - assert got == boundary - warm, got = ckpt_lookup(man, ids[:40], extra_hash=7) - assert warm is None and got == 0 + # Strip-on-extend keeps the newest two boundaries hittable plus the + # first one, promoted to the chain's anchor; the interior boundary + # is stripped by design. + for probe, hit in ((40, 32), (72, 32), (104, 96), (136, 128)): + warm, got = ckpt_lookup(man, ids[:probe], extra_hash=7) + assert got == hit def test_cursor_off_boundary_chunk_is_a_noop(): @@ -411,3 +413,75 @@ def kv_cache(): bg3 = ar.BatchGenerator(SimpleNamespace(), None, draft_model=object(), apc_manager=None, prefill_batch_size=8) assert bg3.prefill_batch_size == 8 + + +# -- system-prefix anchor boundary -- + +def _stub_sys(monkeypatch, lcp): + from gmlx import retire_key + monkeypatch.setattr(retire_key, "lookup_render_ctx", + lambda ids: {"stub": True}) + monkeypatch.setattr(retire_key, "prompt_stable_lcp", + lambda ctx, ids: None) + monkeypatch.setattr(retire_key, "system_prefix_lcp", + lambda ctx, ids: lcp) + + +def test_sys_boundary_block_grid_attention(monkeypatch): + # Attention layouts snap the system-prefix stop to the block grid; + # the record kind is anchor, exempt from strip-on-extend. + _stub_sys(monkeypatch, 2900) + meta = _arm_meta(5000, ("rot:512:0", "kv")) + assert meta["ckpt_boundaries"] == [ + (2896, "anchor"), (4096, "boundary"), (4999, "replay")] + assert meta["ckpt_sys_bound"] == 2896 + assert meta["checkpoint_len"] == 2896 + + +def test_sys_boundary_chunk_grid_on_arr(monkeypatch): + # arr layouts keep the chunk grid (off-grid chunking drifts GDN + # state), so the stop lands a full grid point below the offset. + _stub_sys(monkeypatch, 2900) + meta = _arm_meta(5000, ("kv", "arr")) + assert meta["ckpt_boundaries"] == [ + (2048, "anchor"), (4096, "boundary"), (4999, "replay")] + assert meta["ckpt_sys_bound"] == 2048 + + +def test_sys_boundary_floors(monkeypatch): + # Below GMLX_APC_CKPT_SYS_MIN no anchor lands; arr layouts raise the + # floor to the replay byte floor (state clones are size-independent). + _stub_sys(monkeypatch, 200) + meta = _arm_meta(5000, ("rot:512:0", "kv")) + assert "ckpt_sys_bound" not in meta + assert all(k != "anchor" for _, k in meta["ckpt_boundaries"]) + _stub_sys(monkeypatch, 1500) # grid point 0 on the 2048 grid + meta = _arm_meta(5000, ("kv", "arr")) + assert "ckpt_sys_bound" not in meta + + +def test_sys_boundary_upgrades_boundary_never_replay(monkeypatch): + # Collision with a plain boundary upgrades it to anchor; a replay + # arriving at the same position afterwards never downgrades it. + _stub_sys(monkeypatch, 4100) + meta = _arm_meta(5000, ("rot:512:0", "kv")) + assert meta["ckpt_boundaries"] == [(4096, "anchor"), (4999, "replay")] + meta = _arm_meta(4097, ("rot:512:0", "kv")) + assert meta["ckpt_boundaries"] == [(4096, "anchor")] + + +def test_sys_boundary_kill_switch(monkeypatch): + _stub_sys(monkeypatch, 2900) + monkeypatch.setenv("GMLX_APC_CKPT_SYS", "0") + meta = _arm_meta(5000, ("rot:512:0", "kv")) + assert all(k != "anchor" for _, k in meta["ckpt_boundaries"]) + assert "ckpt_sys_bound" not in meta + + +def test_sys_boundary_skipped_when_restored_past(monkeypatch): + # A warm sibling restored at or past the anchor has nothing to + # re-store; the hit already refreshed the record's LRU position. + _stub_sys(monkeypatch, 2900) + meta = _arm_meta(5000, ("rot:512:0", "kv"), restored=3000) + assert all(k != "anchor" for _, k in meta["ckpt_boundaries"]) + assert "ckpt_sys_bound" not in meta diff --git a/tests/test_ckpt_tier.py b/tests/test_ckpt_tier.py index c545876..c22a188 100644 --- a/tests/test_ckpt_tier.py +++ b/tests/test_ckpt_tier.py @@ -186,7 +186,9 @@ def test_replay_record_survives_retirement_insert(): row=0, extra_hash=0) idx = _ckpt_records(man) assert (n - 1) in [r.p for r in idx.values()] - assert {r.kind for r in idx.values()} == {"replay", "boundary", + # The terminal store is the chain's first restorable boundary (the + # replay below it is gated), so it promotes to anchor. + assert {r.kind for r in idx.values()} == {"replay", "anchor", "retire"} # The identical resend adopts at N-1. warm, got = ckpt_lookup(man, ids[:n], extra_hash=0) @@ -239,7 +241,8 @@ def test_replay_record_not_exempt_from_byte_budget(monkeypatch): assert ckpt_store(man, list(range(300, 347)), make_hybrid_cache(47, seed=2), extra_hash=1) idx = cs._ckpt_records(man) - assert [r.kind for r in idx.values()] == ["boundary"] + # The surviving boundary is its chain's first, hence anchor. + assert [r.kind for r in idx.values()] == ["anchor"] def test_salt_isolation_from_real_tiers(): @@ -529,17 +532,24 @@ def test_pinning_survives_pool_pressure(): assert_warm_matches(warm, cache, p) -def test_strip_on_extend_keeps_newest_two(): +def test_strip_on_extend_keeps_newest_two_plus_anchor(): + """A growing chain keeps the newest two records plus its anchor: + the first restorable boundary is promoted and survives the strip + (sibling fan-out adopts exactly that early prefix), while interior + boundaries strip as before.""" from gmlx.cache_snapshot import _ckpt_records - man = APCManager(num_blocks=64, block_size=16) + man = APCManager(num_blocks=96, block_size=16) ids = list(range(400, 400 + 96)) - for p in (32, 48, 64): + for p in (32, 48, 64, 80): cache = make_hybrid_cache(p, seed=p) assert ckpt_store(man, ids[:p], cache, extra_hash=0) idx = _ckpt_records(man) - assert sorted(r.p for r in idx.values()) == [48, 64] + assert sorted(r.p for r in idx.values()) == [32, 64, 80] + assert [r.p for r in idx.values() if r.kind == "anchor"] == [32] warm, got = ckpt_lookup(man, ids[:40], extra_hash=0) - assert warm is None and got == 0 # p=32 stripped + assert got == 32 # the anchor serves siblings + warm, got = ckpt_lookup(man, ids[:50], extra_hash=0) + assert got == 32 # p=48 stripped warm, got = ckpt_lookup(man, ids[:66], extra_hash=0) assert got == 64 @@ -560,7 +570,8 @@ def test_strip_on_extend_exempts_replay(): assert ckpt_store(man, ids[:80], make_hybrid_cache(80, seed=80), extra_hash=0, kind="retire") idx = _ckpt_records(man) - assert sorted(r.p for r in idx.values()) == [47, 64, 80] + # p=32 promoted to anchor (first restorable boundary), also exempt. + assert sorted(r.p for r in idx.values()) == [32, 47, 64, 80] warm, got = ckpt_lookup(man, ids[:48], extra_hash=0) assert got == 47 @@ -1109,3 +1120,143 @@ def drive(disabled): assert not hasattr(batch_off, "_apc_manager") # stock store never armed assert not hasattr(batch_off.prompt_cache[0], "_kq_apc_retire") assert se._get_spec_prefix_cache(model_off) is None # L0 off too + + +# -- anchor records: the sibling fan-out exemption -- + +def test_anchor_kind_exempt_from_strip_and_superseded(): + """A tagged anchor survives strip-on-extend as the chain deepens; a + newer tagged anchor on the same chain supersedes it (one anchor per + chain).""" + from gmlx.cache_snapshot import _ckpt_records + + man = APCManager(num_blocks=96, block_size=16) + ids = list(range(400, 400 + 96)) + assert ckpt_store(man, ids[:32], make_hybrid_cache(32, seed=32), + extra_hash=0, kind="anchor") + for p in (48, 64, 80): + assert ckpt_store(man, ids[:p], make_hybrid_cache(p, seed=p), + extra_hash=0) + idx = _ckpt_records(man) + assert sorted(r.p for r in idx.values()) == [32, 64, 80] + assert [r.p for r in idx.values() if r.kind == "anchor"] == [32] + # Deeper records exist below the new anchor position: no promotion + # happened at 48/64/80 (the chain was never fresh). + assert ckpt_store(man, ids[:48], make_hybrid_cache(48, seed=1), + extra_hash=0, kind="anchor") + idx = _ckpt_records(man) + assert [r.p for r in idx.values() if r.kind == "anchor"] == [48] + assert 32 not in [r.p for r in idx.values()] + + +def test_anchor_evicts_after_non_anchors_lru_by_hit(monkeypatch): + """Entry-cap pressure: non-anchors go first even when an anchor is + older; among anchors the least-recently-hit goes first.""" + import gmlx.cache_snapshot as cs + + monkeypatch.setattr(cs, "_CKPT_RECORD_ENTRIES", 3) + man = APCManager(num_blocks=96, block_size=16) + a = list(range(100, 148)) + b = list(range(300, 364)) + assert ckpt_store(man, a[:32], make_hybrid_cache(32, seed=1), + extra_hash=0) # anchor A + assert ckpt_store(man, b[:32], make_hybrid_cache(32, seed=2), + extra_hash=1) # anchor B + assert ckpt_store(man, b[:48], make_hybrid_cache(48, seed=3), + extra_hash=1) # plain boundary + warm, got = ckpt_lookup(man, a[:40], extra_hash=0) # hit refreshes A + assert got == 32 + assert ckpt_store(man, list(range(500, 532)), + make_hybrid_cache(32, seed=4), extra_hash=2) + idx = cs._ckpt_records(man) + # The plain boundary (B:48) went first despite being newer than both + # anchors. + assert [(r.p, r.extra_hash) for r in idx.values() if r.kind != "anchor"] \ + == [] + assert {r.extra_hash for r in idx.values()} == {0, 1, 2} + assert ckpt_store(man, list(range(700, 732)), + make_hybrid_cache(32, seed=5), extra_hash=3) + idx = cs._ckpt_records(man) + # All anchors now: the least-recently-hit one (B) went; the hit A + # record stayed. + assert {r.extra_hash for r in idx.values()} == {0, 2, 3} + + +def test_first_boundary_promotion_skips_retire_chains(): + """Promotion targets boundaries only: a chain whose first record is + a retirement store gets no anchor from it, and a later boundary + above it does not promote either (the chain is not fresh).""" + from gmlx.cache_snapshot import _ckpt_records + + man = APCManager(num_blocks=96, block_size=16) + ids = list(range(400, 400 + 96)) + assert ckpt_store(man, ids[:32], make_hybrid_cache(32, seed=1), + extra_hash=0, kind="retire") + assert ckpt_store(man, ids[:64], make_hybrid_cache(64, seed=2), + extra_hash=0) + idx = _ckpt_records(man) + assert sorted((r.p, r.kind) for r in idx.values()) == \ + [(32, "retire"), (64, "boundary")] + + +def test_anchor_never_shadows_a_deeper_disk_skeleton(tmp_path): + """Depth beats retention: with only the anchor pinned in memory and a + deeper skeleton on disk, the lookup must return the disk depth. The + pinned walk returns on first success, so an anchor left to win here + caps every divergent query at its own p (the depth e2e's divergent + and turns floors).""" + from gmlx.apc_manager import GmlxAPCManager + from gmlx.cache_snapshot import _ckpt_records, rotating_canonical_window + + ids = list(range(700, 700 + 96)) + disk = DiskBlockStore(root=tmp_path, namespace="m") + man = GmlxAPCManager(num_blocks=96, block_size=16, disk=disk) + try: + deep = make_swa_cache(64, seed=11) + # The shallow store must carry the same KV as the deep one's + # prefix: the block pool dedups the shared chain by token hash, + # so mismatched fixture content would be a fixture artifact. + shallow = [] + for c in deep: + k, v = rotating_canonical_window(c)[:2] \ + if isinstance(c, RotatingKVCache) else c.state + if isinstance(c, RotatingKVCache): + s = RotatingKVCache(max_size=ROT_W) + s.update_and_fetch(k[..., :32, :], v[..., :32, :]) + else: + s = KVCache() + s.state = (k[..., :32, :], v[..., :32, :]) + shallow.append(s) + assert ckpt_store(man, ids[:32], shallow, + extra_hash=4, kind="anchor") + assert ckpt_store(man, ids[:64], deep, extra_hash=4) + # Drop the deep record from memory, keeping its disk skeleton: + # exactly what strip-on-extend leaves behind as a chain deepens. + idx = _ckpt_records(man) + for k, r in list(idx.items()): + if r.p == 64: + idx.pop(k) + assert [r.kind for r in idx.values()] == ["anchor"] + warm, got = ckpt_lookup(man, ids + [77], extra_hash=4) + assert got == 64 + assert_swa_warm_matches(warm, deep, 64) + # Nothing deeper on disk: the anchor still serves. + warm, got = ckpt_lookup(man, ids[:48] + [77], extra_hash=4) + assert got == 32 + finally: + disk.close() + + +def test_anchor_gets_no_pool_pressure_protection(): + """_evict_for_pool stays plain LRU: an anchor's blocks reclaim like + any record's, so a pinned anchor can never starve the block pool.""" + import gmlx.cache_snapshot as cs + + man = APCManager(num_blocks=64, block_size=16) + ids = list(range(400, 448)) + assert ckpt_store(man, ids[:32], make_hybrid_cache(32, seed=1), + extra_hash=0) + idx = cs._ckpt_records(man) + assert [r.kind for r in idx.values()] == ["anchor"] + assert cs._evict_for_pool(man, 1) >= 1 + assert len(cs._ckpt_records(man)) == 0 diff --git a/tests/test_exact_anchor.py b/tests/test_exact_anchor.py new file mode 100644 index 0000000..a3b7bee --- /dev/null +++ b/tests/test_exact_anchor.py @@ -0,0 +1,328 @@ +"""Exact-tier anchor: whole-prefix clone at the sibling divergence. + +Store and lookup live in a gmlx-owned side LRU; the upstream exact LRU +is count-capped and every request writes its guard-column entry there, +so sibling fan-out churns out the early shared-prefix entry. The armed +prefill pauses twice: the anchor hook stores at the divergence and +hands the checkpoint column back to the stock guard, whose store and +checkpoint_done latch fire exactly as unarmed. +""" + +from types import SimpleNamespace + +import mlx.core as mx +from mlx_vlm.apc import APCManager +from mlx_vlm.models.cache import CacheList, KVCache + +import gmlx.cache_snapshot as cs +import gmlx.spec_engine as se +from gmlx.cache_snapshot import anchor_exact_lookup, anchor_exact_store + + +def make_kv_cache(p, layers=2, seed=0): + out = [] + for i in range(layers): + c = KVCache() + if p > 0: + k = mx.full((1, 2, p, 8), float(seed + i), dtype=mx.float16) + v = mx.full((1, 2, p, 8), float(seed + i + 1), dtype=mx.float16) + c.update_and_fetch(k, v) + out.append(c) + return out + + +IDS = list(range(500, 500 + 120)) + + +# -- store/lookup roundtrip on the side LRU -- + +def test_anchor_roundtrip_and_decoupling(): + man = APCManager(num_blocks=8, block_size=16) + src = make_kv_cache(64) + assert anchor_exact_store(man, IDS[:64], src) + # The stored entry is a clone: mutating the source afterwards must + # not leak into what siblings restore. + src[0].update_and_fetch( + mx.zeros((1, 2, 8, 8), dtype=mx.float16), + mx.zeros((1, 2, 8, 8), dtype=mx.float16)) + warm, p = anchor_exact_lookup(man, IDS) + assert p == 64 and all(int(c.offset) == 64 for c in warm) + # The warm result is a clone too: advancing it must not corrupt the + # entry the next sibling gets. + warm[0].update_and_fetch( + mx.zeros((1, 2, 8, 8), dtype=mx.float16), + mx.zeros((1, 2, 8, 8), dtype=mx.float16)) + warm2, p2 = anchor_exact_lookup(man, IDS) + assert p2 == 64 and all(int(c.offset) == 64 for c in warm2) + + +def test_anchor_lookup_gates(): + man = APCManager(num_blocks=8, block_size=16) + anchor_exact_store(man, IDS[:64], make_kv_cache(64)) + # Different chain, wrong extra_hash, equal-length query, min-prefix + # at or past the anchor: all misses. + assert anchor_exact_lookup(man, [1, 2, 3] + IDS[3:]) == (None, 0) + assert anchor_exact_lookup(man, IDS, extra_hash=9) == (None, 0) + assert anchor_exact_lookup(man, IDS[:64]) == (None, 0) + assert anchor_exact_lookup(man, IDS, min_prefix_tokens=64) == (None, 0) + + +def test_anchor_longest_prefix_wins(): + man = APCManager(num_blocks=8, block_size=16) + anchor_exact_store(man, IDS[:32], make_kv_cache(32)) + anchor_exact_store(man, IDS[:64], make_kv_cache(64)) + _, p = anchor_exact_lookup(man, IDS) + assert p == 64 + + +def test_anchor_entries_cap_lru(monkeypatch): + monkeypatch.setattr(cs, "_ANCHOR_ENTRIES", 2) + man = APCManager(num_blocks=8, block_size=16) + a, b, c = ([i] + IDS for i in (1, 2, 3)) + anchor_exact_store(man, a[:64], make_kv_cache(64)) + anchor_exact_store(man, b[:64], make_kv_cache(64)) + _, p = anchor_exact_lookup(man, a) # refresh a's LRU position + assert p == 64 + anchor_exact_store(man, c[:64], make_kv_cache(64)) + assert anchor_exact_lookup(man, a)[1] == 64 + assert anchor_exact_lookup(man, b) == (None, 0) + assert anchor_exact_lookup(man, c)[1] == 64 + + +def test_anchor_byte_budget_newest_survives(monkeypatch): + monkeypatch.setattr(cs, "_ANCHOR_BUDGET_BYTES", 1) + man = APCManager(num_blocks=8, block_size=16) + a, b = ([i] + IDS for i in (1, 2)) + anchor_exact_store(man, a[:64], make_kv_cache(64)) + anchor_exact_store(man, b[:64], make_kv_cache(64)) + assert anchor_exact_lookup(man, a) == (None, 0) + assert anchor_exact_lookup(man, b)[1] == 64 + + +def test_anchor_cleared_by_ckpt_reset(): + man = APCManager(num_blocks=8, block_size=16) + anchor_exact_store(man, IDS[:64], make_kv_cache(64)) + cs.ckpt_reset(man) + assert anchor_exact_lookup(man, IDS) == (None, 0) + + +# -- boundary: ungridded divergence, clamped to the stock guard -- + +def _stub_sys(monkeypatch, lcp): + from gmlx import retire_key + monkeypatch.setattr(retire_key, "lookup_render_ctx", + lambda ids: {"messages": ()}) + monkeypatch.setattr(retire_key, "system_prefix_lcp", + lambda ctx, ids: lcp) + + +def _bmeta(n=5000): + return SimpleNamespace(), {"full_input_ids": list(range(n))} + + +def test_anchor_boundary_ungridded(monkeypatch): + _stub_sys(monkeypatch, 2900) + batch, meta = _bmeta() + assert se._exact_anchor_boundary(batch, meta, 4000, 0) == 2900 + + +def test_anchor_boundary_clamps_to_guard(monkeypatch): + _stub_sys(monkeypatch, 4500) + batch, meta = _bmeta() + assert se._exact_anchor_boundary(batch, meta, 4000, 0) == 4000 + # Guard 0 (stock checkpoint disabled): the divergence stands alone. + assert se._exact_anchor_boundary(batch, meta, 0, 0) == 4500 + + +def test_anchor_boundary_floor_kill_restored(monkeypatch): + _stub_sys(monkeypatch, 200) + batch, meta = _bmeta() + assert se._exact_anchor_boundary(batch, meta, 4000, 0) is None + monkeypatch.setenv("GMLX_APC_CKPT_SYS_MIN", "100") + assert se._exact_anchor_boundary(batch, meta, 4000, 0) == 200 + assert se._exact_anchor_boundary(batch, meta, 4000, 200) is None + monkeypatch.setenv("GMLX_APC_CKPT_SYS", "0") + assert se._exact_anchor_boundary(batch, meta, 4000, 0) is None + + +def test_anchor_boundary_no_render_ctx(monkeypatch): + from gmlx import retire_key + monkeypatch.setattr(retire_key, "lookup_render_ctx", lambda ids: None) + batch, meta = _bmeta() + assert se._exact_anchor_boundary(batch, meta, 4000, 0) is None + + +# -- two-stop schedule: anchor store, then the untouched stock guard -- + +def _armed_exact_batch(man, guard, monkeypatch, lcp): + _stub_sys(monkeypatch, lcp) + monkeypatch.setenv("GMLX_APC_CKPT_SYS_MIN", "16") + meta = {"full_input_ids": IDS, "prefix_len": 0, "extra_hash": 7, + "checkpoint_len": guard} + batch = SimpleNamespace( + _apc_manager=man, _apc_meta=[meta], _apc_mode="exact", + prompt_cache=None) + batch._apc_prompt_cache_for_store = lambda idx: batch.prompt_cache + se._exact_anchor_arm(batch, meta, guard, 0) + return batch, meta + + +def test_anchor_two_stop_schedule(monkeypatch): + from mlx_vlm.generate import ar + + man = APCManager(num_blocks=8, block_size=16) + calls = [] + orig = man.store_exact_cache + man.store_exact_cache = ( + lambda *a, **k: (calls.append(len(a[0])), orig(*a, **k))[1]) + batch, meta = _armed_exact_batch(man, guard=96, monkeypatch=monkeypatch, + lcp=64) + assert batch._kq_anchor_armed and meta["checkpoint_len"] == 64 + stock = ar.PromptProcessingBatch._store_apc_exact_checkpoints + + # Anchor stop: the hook stores, hands the column back to the guard, + # and the stock body (running right after, as in the wrap) skips. + batch.prompt_cache = make_kv_cache(64) + batch._row_real_tokens_processed = lambda idx: 64 + se._exact_anchor_store(batch) + stock(batch) + assert meta["anchor_done"] and meta["checkpoint_len"] == 96 + assert not meta.get("checkpoint_done") and calls == [] + assert anchor_exact_lookup(man, IDS, extra_hash=7)[1] == 64 + + # Guard stop: the hook is spent; the stock store fires and latches. + batch.prompt_cache = make_kv_cache(96) + batch._row_real_tokens_processed = lambda idx: 96 + se._exact_anchor_store(batch) + stock(batch) + assert calls == [96] and meta.get("checkpoint_done") + + +def test_anchor_at_guard_single_stop(monkeypatch): + from mlx_vlm.generate import ar + + man = APCManager(num_blocks=8, block_size=16) + calls = [] + orig = man.store_exact_cache + man.store_exact_cache = ( + lambda *a, **k: (calls.append(len(a[0])), orig(*a, **k))[1]) + # Divergence past the guard clamps onto it: one pause, both stores. + batch, meta = _armed_exact_batch(man, guard=96, monkeypatch=monkeypatch, + lcp=110) + assert meta["anchor_len"] == 96 and meta["checkpoint_len"] == 96 + batch.prompt_cache = make_kv_cache(96) + batch._row_real_tokens_processed = lambda idx: 96 + se._exact_anchor_store(batch) + ar.PromptProcessingBatch._store_apc_exact_checkpoints(batch) + assert anchor_exact_lookup(man, IDS, extra_hash=7)[1] == 96 + assert calls == [96] and meta.get("checkpoint_done") + + +# -- admission pick: the anchor rides the stock warm-batch builder -- + +def _pooling_model(man): + """A pooling stack resolves exact and is not ckpt-tier (the layout + probe rejects unknown cache classes) -- the ds4 shape.""" + from gmlx.deepseek_v4_cache import PoolingCache + + return SimpleNamespace( + _kq_apc_manager=man, _kq_apc_mode="exact", + config=SimpleNamespace(), + make_cache=lambda: [KVCache(), + CacheList(KVCache(), PoolingCache(4))]) + + +def _pick_gen(man): + from mlx_vlm.generate.ar import BatchGenerator + + se._install_exact_anchor_pick() + gen = SimpleNamespace( + apc_manager=man, apc_mode="exact", model=_pooling_model(man), + _apc_media_token_ids=lambda: set()) + for name in ("_apc_pick_for", "_apc_extra_hash", + "_apc_safe_prefix_lookup_min", "_apc_suffix_is_text_only"): + setattr(gen, name, getattr(BatchGenerator, name).__get__(gen)) + return gen + + +def _seq(ids, uid="u1"): + return (uid, list(ids), 8, {}, None, None) + + +def test_anchor_pick_returns_a_warm_prefix(): + man = APCManager(num_blocks=8, block_size=16) + anchor_exact_store(man, IDS[:64], make_kv_cache(64)) + pick = _pick_gen(man)._apc_pick_for(_seq(IDS)) + assert pick["prefix_len"] == 64 and pick["matched_blocks"] == [] + assert all(int(c.offset) == 64 for c in pick["warm_cache"]) + assert pick["full_input_ids"] == IDS + + +def test_anchor_pick_never_shortens_the_stock_pick(): + man = APCManager(num_blocks=8, block_size=16) + anchor_exact_store(man, IDS[:64], make_kv_cache(64)) + # A deeper stock exact entry wins: the anchor is a floor, not a cap. + man.store_exact_cache(IDS[:100], make_kv_cache(100, seed=5)) + pick = _pick_gen(man)._apc_pick_for(_seq(IDS)) + assert pick["prefix_len"] == 100 + + +def test_anchor_pick_passes_through_without_a_hit(): + man = APCManager(num_blocks=8, block_size=16) + assert _pick_gen(man)._apc_pick_for(_seq(IDS)) is None + + +# -- stock-path init: arms the store, right padding included -- + +def _plain_batch(man, right_pad, prefix_len=0): + n = len(IDS) + meta = {"full_input_ids": list(IDS), "prefix_len": prefix_len, + "extra_hash": 0, "checkpoint_len": 96} + return SimpleNamespace( + model=_pooling_model(man), uids=["u1"], _right_pad_per_row=right_pad, + _apc_manager=man, _apc_mode="exact", _apc_meta=[meta], + _input_ids=mx.array([IDS]), _inputs_embeds=mx.zeros((1, n, 4)), + _prompt_kwargs={}, _prompt_length_aware_keys=[], + prompt_cache=None), meta + + +def test_plain_anchor_init_arms_on_a_right_padded_row(monkeypatch): + se._bind_l1_view() + man = APCManager(num_blocks=8, block_size=16) + _stub_sys(monkeypatch, 64) + monkeypatch.setenv("GMLX_APC_CKPT_SYS_MIN", "16") + # Right padding means this row rode a warm batch built for a longer + # sibling. Upstream's checkpoint column handles it, so must we: this + # is the common shape once any request is warm. + batch, meta = _plain_batch(man, right_pad=[0]) + se._plain_anchor_init(batch) + assert batch._kq_anchor_armed and meta["checkpoint_len"] == 64 + # Nothing is trimmed here: restores come from the admission pick. + assert batch._input_ids.shape[1] == len(IDS) + + +def test_plain_anchor_init_arms_above_a_shallow_warm_prefix(monkeypatch): + se._bind_l1_view() + man = APCManager(num_blocks=8, block_size=16) + _stub_sys(monkeypatch, 64) + monkeypatch.setenv("GMLX_APC_CKPT_SYS_MIN", "16") + # A warm row is not an anchored row. The stock exact tier routinely + # matches a token or two off an unrelated request; the divergence + # still needs its clone. + batch, meta = _plain_batch(man, right_pad=[0], prefix_len=1) + se._plain_anchor_init(batch) + assert batch._kq_anchor_armed and meta["checkpoint_len"] == 64 + + +def test_plain_anchor_init_skips_a_row_restored_past_the_divergence( + monkeypatch): + se._bind_l1_view() + man = APCManager(num_blocks=8, block_size=16) + _stub_sys(monkeypatch, 64) + monkeypatch.setenv("GMLX_APC_CKPT_SYS_MIN", "16") + # Restored at the divergence: the anchor it would store exists, so + # no second stop and the stock guard runs alone. + batch, meta = _plain_batch(man, right_pad=None, prefix_len=64) + se._plain_anchor_init(batch) + assert not getattr(batch, "_kq_anchor_armed", False) + assert meta["checkpoint_len"] == 96 diff --git a/tests/test_retire_key.py b/tests/test_retire_key.py index 83e007e..bcd19ed 100644 --- a/tests/test_retire_key.py +++ b/tests/test_retire_key.py @@ -364,3 +364,114 @@ def render(proc, cfg, messages, **kw): # The memoized prompt feeds the closer. assert rk._virtually_finish(ctx, "deep partial") == \ "deep partial" + + +# -- system-prefix LCP (the sibling anchor offset) -- + +def _sys_ctx(next_ids, msgs): + ctx = _fake_ctx(next_ids) + ctx["messages"] = msgs + return ctx + + +def _probe_ctx(msgs, template): + """A ctx whose render/tokenize pair actually varies with the probe + messages: ``template`` maps a message list to text; tokens are + character codes, so the probe LCP is a plain string LCP.""" + ctx = _fake_ctx([0]) + ctx["messages"] = msgs + ctx["render"] = lambda p, c, m, **kw: template(m) + ctx["preprocess"] = lambda text: [ord(ch) for ch in text] + return ctx + + +def _chatml(msgs): + return "".join(f"<{m['role']}>{m['content']}" for m in msgs) + + +def _folding(msgs): + """Gemma-style: system content folds into the first user turn; a + system-only conversation renders to almost nothing.""" + sys_txt = "".join(m["content"] for m in msgs if m["role"] == "system") + users = [m for m in msgs if m["role"] == "user"] + if not users: + return "" + return ("" + sys_txt + "\n\n" + + "".join(u["content"] for u in users) + "") + + +def test_system_prefix_lcp_probe_pair_explicit_system_block(): + msgs = [{"role": "system", "content": "policy"}, + {"role": "user", "content": "real question"}] + ctx = _probe_ctx(msgs, _chatml) + live = [ord(ch) for ch in _chatml(msgs)] + # Probes diverge right after the shared "policy" + # header; the anchor offset includes the user header, which every + # sibling also shares. + assert rk.system_prefix_lcp(ctx, live) == len("policy") + + +def test_system_prefix_lcp_probe_pair_folding_template(): + # The gemma shape: a system-only render measures nothing, but the + # probe pair folds identically and splits at the real divergence. + msgs = [{"role": "system", "content": "policy"}, + {"role": "user", "content": "real question"}] + ctx = _probe_ctx(msgs, _folding) + live = [ord(ch) for ch in _folding(msgs)] + assert rk.system_prefix_lcp(ctx, live) == len("policy\n\n") + + +def test_system_prefix_lcp_probe_renders_lead_plus_dummy(): + calls = [] + ctx = _sys_ctx([1, 2, 3], [{"role": "system", "content": "a"}, + {"role": "system", "content": "b"}, + {"role": "user", "content": "u"}]) + inner = ctx["render"] + ctx["render"] = (lambda p, c, msgs, **kw: + (calls.append(list(msgs)), inner(p, c, msgs, **kw))[1]) + assert rk.system_prefix_lcp(ctx, [1, 2, 3, 7, 8]) == 3 + lead = [{"role": "system", "content": "a"}, + {"role": "system", "content": "b"}] + assert calls == [lead + [{"role": "user", "content": "0"}], + lead + [{"role": "user", "content": "1"}]] + # Memoized: a second call does not re-render. + assert rk.system_prefix_lcp(ctx, [1, 2, 3, 7, 8]) == 3 + assert len(calls) == 2 + + +def test_system_prefix_lcp_clamped_by_live_prompt(): + # The live-prompt clamp guards against a template leaking user + # content ahead of the divergence point. + ctx = _sys_ctx([1, 2, 9, 9], [{"role": "system", "content": "s"}, + {"role": "user", "content": "u"}]) + assert rk.system_prefix_lcp(ctx, [1, 2, 3, 4, 5]) == 2 + + +def test_system_prefix_lcp_requires_system_prefix_and_a_tail(): + # No leading system block: nothing siblings share by construction. + ctx = _sys_ctx([1, 2], [{"role": "user", "content": "u"}]) + assert rk.system_prefix_lcp(ctx, [1, 2, 3]) is None + # System-only prompt: the terminal checkpoint already covers it. + ctx = _sys_ctx([1, 2], [{"role": "system", "content": "s"}]) + assert rk.system_prefix_lcp(ctx, [1, 2, 3]) is None + # System block after the first non-system message does not count. + ctx = _sys_ctx([1, 2], [{"role": "user", "content": "u"}, + {"role": "system", "content": "s"}]) + assert rk.system_prefix_lcp(ctx, [1, 2, 3]) is None + + +def test_system_prefix_lcp_media_and_failure(): + assert rk.system_prefix_lcp( + _sys_ctx([1], [{"role": "system", "content": "s"}, + {"role": "user", "content": "u"}]) | {"media": True}, + [1, 2]) is None + calls = [] + ctx = _sys_ctx([1, 2], [{"role": "system", "content": "s"}, + {"role": "user", "content": "u"}]) + ctx["render"] = (lambda p, c, msgs, **kw: + (calls.append(1), (_ for _ in ()).throw( + ValueError("template refuses")))[1]) + assert rk.system_prefix_lcp(ctx, [1, 2]) is None + # Failure memoized: no second render attempt. + assert rk.system_prefix_lcp(ctx, [1, 2]) is None + assert len(calls) == 1