diff --git a/mtplx/cache_bank/cold_tier.py b/mtplx/cache_bank/cold_tier.py index d11145079..eb6363bef 100644 --- a/mtplx/cache_bank/cold_tier.py +++ b/mtplx/cache_bank/cold_tier.py @@ -2,6 +2,7 @@ from __future__ import annotations +from array import array import hashlib from collections import Counter, OrderedDict, deque import json @@ -307,12 +308,37 @@ def common_prefix_len(left: tuple[int, ...], right: tuple[int, ...]) -> int: return limit +def _token_array(token_ids: list[int]) -> array | tuple[int, ...]: + """Token ids as a compact ``array('I')``; a tuple if an id does not fit.""" + try: + return array("I", token_ids) + except (OverflowError, TypeError): + return tuple(int(token) for token in token_ids) + + +def _starts_with(tokens: tuple[int, ...], prefix: array | tuple[int, ...]) -> bool: + if len(tokens) < len(prefix): + return False + head = tokens[: len(prefix)] + if isinstance(prefix, array): + try: + return array(prefix.typecode, head) == prefix + except (OverflowError, TypeError): + return False + return tuple(head) == tuple(prefix) + + def block_aligned_prefix_len(matched_tokens: int, *, block_size: int) -> int: block = max(1, int(block_size)) matched = max(0, int(matched_tokens)) return (matched // block) * block +# Upper bound on token ids held by the parsed-row memo. Rows are held as +# 4-byte arrays, so the cap is about 128 MB; a full 100 GB store of long agent +# conversations stays below it (5.7M tokens in 79 rows measured). +_ROW_TOKENS_MEMO_MAX_TOKENS = 32_000_000 + # Ranked SSD rows a prefix lookup opens before it gives up: each refusal # costs only its payload.json read. _PREFIX_BOUNDARY_ROWS_TRIED = 4 @@ -464,6 +490,12 @@ def __init__( self._block_memos: OrderedDict[str, dict[tuple, dict[str, Any]]] = OrderedDict() self._block_memos_lock = threading.Lock() self._block_memos_max = 16 + # entry_id -> (token_hash, parsed token ids) for manifest rows. A row's + # token list is immutable, so a prefix lookup parses each one once + # instead of on every request; bounded by total cached tokens. + self._row_tokens_memo: dict[str, tuple[str, array | tuple[int, ...]]] = {} + self._row_tokens_memo_tokens = 0 + self._row_tokens_memo_lock = threading.Lock() self._stop = threading.Event() self._base_lock = threading.RLock() self._disk_usage_lock = threading.Lock() @@ -1169,6 +1201,7 @@ def lookup_prefix_boundary( block_min_match = max(block, int(block_min_matched_tokens)) started = time.perf_counter() try: + floor = max(0, int(min_useful_matched_tokens)) rows = self._candidate_rows_for_prefix_boundary( model_path=model_path, mtp_enabled=mtp_enabled, @@ -1177,12 +1210,15 @@ def lookup_prefix_boundary( mtp_history_policy=mtp_history_policy, draft_head_identity=draft_head_identity, policy_fingerprint=policy_fingerprint, + min_prefix_len=floor, ) - ranked: list[tuple[tuple[int, int, int], sqlite3.Row, int, str, int]] = [] - for row in rows: - prefix = tuple(int(token) for token in json.loads(str(row["token_ids_json"]))) + + def rank_row( + row: sqlite3.Row, + ) -> tuple[tuple[int, int, int], sqlite3.Row, int, str, int] | None: + prefix = self._row_token_ids(row) if not prefix: - continue + return None matched = common_prefix_len(tokens, prefix) gap = len(prefix) - matched safe_block = min( @@ -1205,16 +1241,40 @@ def lookup_prefix_boundary( candidate_matched = int(safe_block) restore_kind = "block_prefix" else: - continue - ranked.append( - ( - (candidate_matched, int(matched), len(prefix)), - row, - candidate_matched, - restore_kind, - len(prefix), - ) + return None + return ( + (candidate_matched, int(matched), len(prefix)), + row, + candidate_matched, + restore_kind, + len(prefix), ) + + ranked: list[tuple[tuple[int, int, int], sqlite3.Row, int, str, int]] = [] + for row in rows: + item = rank_row(row) + if item is not None: + ranked.append(item) + if not ranked and floor > 0: + # Rows shorter than the caller's RAM match cannot beat it + # (candidate_matched <= prefix_len) and were left out of the + # query. Whether one of them would have matched still + # decides the miss reason, so look at them only on this + # path; their token lists come from the memo. + for row in self._candidate_rows_for_prefix_boundary( + model_path=model_path, + mtp_enabled=mtp_enabled, + hidden_variant=hidden_variant, + template_hash=template_hash, + mtp_history_policy=mtp_history_policy, + draft_head_identity=draft_head_identity, + policy_fingerprint=policy_fingerprint, + max_prefix_len=floor - 1, + ): + if rank_row(row) is not None: + self._inc("prefix_lookups_not_better_than_ram") + self._set_last_miss("ssd_prefix_not_better_than_ram") + return None if not ranked: self._inc("restore_misses") self._set_last_miss("ssd_prefix_miss") @@ -2395,6 +2455,7 @@ def _retire_entry_row(self, row: sqlite3.Row) -> set[str]: blob_hashes = self._entry_blob_hashes(entry_dir) with self._connect() as conn: conn.execute("DELETE FROM entries WHERE entry_id = ?", (entry_id,)) + self._forget_row_tokens(entry_id) try: if entry_dir.exists(): shutil.rmtree(entry_dir) @@ -2409,6 +2470,7 @@ def _archive_entry_row(self, row: sqlite3.Row) -> None: self._archive_orphan_entry_dir(entry_dir, entry_id) with self._connect() as conn: conn.execute("DELETE FROM entries WHERE entry_id = ?", (entry_id,)) + self._forget_row_tokens(entry_id) self._invalidate_disk_usage_cache() def _entry_in_manifest(self, entry_id: str) -> bool: @@ -2664,8 +2726,8 @@ def _restore_row( prefix_restore_tokens: int | None = None, ) -> ColdRestoreRecord | None: metadata = dict(row) - token_ids = tuple(int(token) for token in json.loads(str(metadata["token_ids_json"]))) - if require_exact_prefix and lookup_tokens[: len(token_ids)] != token_ids: + token_ids = self._row_token_ids(row) + if require_exact_prefix and not _starts_with(lookup_tokens, token_ids): return None if int(metadata["format_version"]) != COLD_TIER_FORMAT_VERSION: self._set_last_miss("ssd_format_mismatch") @@ -2803,7 +2865,9 @@ def read_tensor(name: str) -> bytes: ) return ColdRestoreRecord( entry_id=str(metadata["entry_id"]), - token_ids=token_ids, + # The memo holds a compact array; a restored record carries the + # tuple its callers have always received. + token_ids=tuple(token_ids), cache_snapshot=decoded.cache_snapshot, logits=decoded.logits, hidden=decoded.hidden, @@ -2861,6 +2925,37 @@ def _candidate_rows( compatible.append(row) return compatible + def _row_token_ids(self, row: sqlite3.Row) -> array | tuple[int, ...]: + """Token ids of a manifest row, parsed once per (entry_id, token_hash). + + Held as an ``array('I')`` (4 bytes per id; a tuple of ints costs about + 40). A full memo keeps its residents and hands a newcomer back + unmemoized: prefix lookups walk the rows in the same order every + time, so evicting the oldest row would drop each row just before its + next use and the memo would never hit. + """ + entry_id = str(row["entry_id"]) + token_hash = str(row["token_hash"]) + with self._row_tokens_memo_lock: + hit = self._row_tokens_memo.get(entry_id) + if hit is not None and hit[0] == token_hash: + return hit[1] + tokens = _token_array(json.loads(str(row["token_ids_json"]))) + with self._row_tokens_memo_lock: + old = self._row_tokens_memo.pop(entry_id, None) + if old is not None: + self._row_tokens_memo_tokens -= len(old[1]) + if self._row_tokens_memo_tokens + len(tokens) <= _ROW_TOKENS_MEMO_MAX_TOKENS: + self._row_tokens_memo[entry_id] = (token_hash, tokens) + self._row_tokens_memo_tokens += len(tokens) + return tokens + + def _forget_row_tokens(self, entry_id: str) -> None: + with self._row_tokens_memo_lock: + old = self._row_tokens_memo.pop(str(entry_id), None) + if old is not None: + self._row_tokens_memo_tokens -= len(old[1]) + def _candidate_rows_for_prefix_boundary( self, *, @@ -2871,12 +2966,20 @@ def _candidate_rows_for_prefix_boundary( mtp_history_policy: str | None, draft_head_identity: str | None, policy_fingerprint: str | None, + min_prefix_len: int = 0, + max_prefix_len: int | None = None, ) -> list[sqlite3.Row]: query = [ "SELECT * FROM entries WHERE model_path = ?", "AND mtp_enabled = ?", ] params: list[Any] = [str(model_path), 1 if mtp_enabled else 0] + if min_prefix_len > 0: + query.append("AND prefix_len >= ?") + params.append(int(min_prefix_len)) + if max_prefix_len is not None: + query.append("AND prefix_len <= ?") + params.append(int(max_prefix_len)) if hidden_variant is not None: query.append("AND hidden_variant = ?") params.append(str(hidden_variant)) @@ -3221,6 +3324,7 @@ def _invalidate_disk_usage_cache(self) -> None: self._disk_usage_cache["disk_usage_stale"] = True def _insert_manifest(self, metadata: dict[str, Any]) -> None: + self._forget_row_tokens(str(metadata["entry_id"])) with self._connect() as conn: conn.execute( """ diff --git a/mtplx/engine_session.py b/mtplx/engine_session.py index a263b159f..559002b32 100644 --- a/mtplx/engine_session.py +++ b/mtplx/engine_session.py @@ -34,6 +34,7 @@ block_aligned_prefix_len, ) from .runtime_options import block_prefix_restore_enabled +from .token_ids import as_token_tuple logger = logging.getLogger(__name__) @@ -1681,7 +1682,7 @@ def commit( if finish_reason not in {"stop", "length", "tool_calls"}: return EngineSessionCommit(False, f"unsafe_finish:{finish_reason}", self.prefix_len) self.note_turn_prompt_len(len(prompt_ids)) - tokens = tuple(int(token) for token in prompt_ids) + tuple(int(token) for token in generated_ids) + tokens = as_token_tuple(prompt_ids) + tuple(int(token) for token in generated_ids) self.committed_token_ids = tokens self.last_commit_s = time.time() self.last_finish_reason = finish_reason @@ -1708,7 +1709,7 @@ def commit_prompt_prefix( prefix gives the resolver a stable session anchor without claiming the assistant/tool history has landed yet. """ - tokens = tuple(int(token) for token in prompt_ids) + tokens = as_token_tuple(prompt_ids) if not tokens: return EngineSessionCommit(False, "empty_prompt_prefix", self.prefix_len) current = self.committed_token_ids @@ -1752,7 +1753,7 @@ def commit_retokenized_prefix( prevents the next OpenCode turn from resolving by a stale prompt boundary that can differ by one chat-template token. """ - tokens = tuple(int(token) for token in token_ids) + tokens = as_token_tuple(token_ids) if not tokens: return EngineSessionCommit(False, "empty_retokenized_prefix", self.prefix_len) if expected_revision is not None and self.revision != int(expected_revision): @@ -2270,7 +2271,7 @@ def generation_slot( acquired.end_generation() def longest_prefix_session(self, token_ids: list[int] | tuple[int, ...]) -> EngineSession | None: - tokens = tuple(int(token) for token in token_ids) + tokens = as_token_tuple(token_ids) best: EngineSession | None = None for session in self._sessions_snapshot(): prefix = session.committed_token_ids @@ -2297,7 +2298,7 @@ def best_common_prefix_session( is strictly better than forking. A 64-token probe rejects unrelated sessions before the O(prefix) compare. """ - tokens = tuple(int(token) for token in token_ids) + tokens = as_token_tuple(token_ids) if not tokens: return None, 0 probe = tokens[:_COMMON_PREFIX_PROBE_TOKENS] @@ -2339,7 +2340,7 @@ def pending_near_prefix_session( min_matched_tokens: int | None = None, ) -> tuple[EngineSession | None, int]: """Return a pending-postcommit session with a near-exact prompt prefix.""" - tokens = tuple(int(token) for token in token_ids) + tokens = as_token_tuple(token_ids) gap_limit = ( _near_prefix_max_token_gap() if max_token_gap is None @@ -2388,7 +2389,7 @@ def nearest_prefix_session( self, token_ids: list[int] | tuple[int, ...], ) -> tuple[EngineSession | None, int]: - tokens = tuple(int(token) for token in token_ids) + tokens = as_token_tuple(token_ids) best: EngineSession | None = None best_len = 0 for session in self._sessions_snapshot(): @@ -2408,7 +2409,7 @@ def _prefix_diagnostic( selected: EngineSession | None = None, exact: bool = False, ) -> dict[str, Any]: - tokens = tuple(int(token) for token in token_ids) + tokens = as_token_tuple(token_ids) session, matched = (selected, len(selected.committed_token_ids)) if selected is not None else self.nearest_prefix_session(tokens) if session is None: return { diff --git a/mtplx/generation.py b/mtplx/generation.py index ebbd8eca7..108f6f90a 100644 --- a/mtplx/generation.py +++ b/mtplx/generation.py @@ -143,6 +143,7 @@ sample_from_distribution, ) from .session_bank import _boundary_true_restore_enabled +from .token_ids import as_token_tuple from .runtime_options import ( block_prefix_restore_enabled, env_bool, @@ -5107,7 +5108,7 @@ def _restore_near_prefix_prompt_state( if getattr(session_bank, "SUPPORTS_NEAR_PREFIX_MIN_RESTORE", False): _candidates_kwargs["min_restore_tokens"] = int(min_restore_tokens) for entry, matched in candidates( - match_ids if image_keyed else prompt_ids, + match_ids if image_keyed else as_token_tuple(prompt_ids), max_token_gap=max_gap, min_matched_tokens=min_match, **_candidates_kwargs, @@ -5441,7 +5442,7 @@ def _near_reject(reason: str) -> None: logits=logits[:, -1, :], hidden=hidden[:, -1:, :] if hidden is not None else None, committed_mtp_cache=mtp_history_cache, - token_prefix=tuple(int(token) for token in prompt_ids), + token_prefix=as_token_tuple(prompt_ids), prompt_eval_time_s=repair_time + repage_time, cache_restore_time_s=total_cache_restore_time_s, mtp_history_policy=mtp_history_policy, @@ -5515,7 +5516,7 @@ def _near_reject(reason: str) -> None: logits=suffix_logits, hidden=suffix_hidden, committed_mtp_cache=mtp_history_cache, - token_prefix=tuple(int(token) for token in prompt_ids), + token_prefix=as_token_tuple(prompt_ids), prompt_eval_time_s=repair_time + suffix_time + mtp_history_time, prompt_mtp_history_time_s=mtp_history_time, cache_restore_time_s=total_cache_restore_time_s, @@ -6564,6 +6565,9 @@ def restore_or_prefill_prompt_state( base_hidden_variant = _resolve_runtime_base_hidden_variant(rt, base_hidden_variant) mtp_hidden_variant = _resolve_runtime_mtp_hidden_variant(rt, mtp_hidden_variant) mtp_position_mode = _resolve_runtime_mtp_position_mode(rt) + # Normalized once: the bank's per-entry compares are shared across its + # lookups when they are handed the same tuple object. + prompt_tuple = as_token_tuple(prompt_ids) os.environ["MTPLX_CURRENT_PREFILL_CONTEXT_TOKENS"] = str(len(prompt_ids)) mtp_history_policy = _resolve_mtp_history_policy( mtp_history_policy, @@ -6801,7 +6805,9 @@ def _emit_prefill_complete(state: PromptState) -> PromptState: ) # The bank only ever sees the content-keyed view of a vision prompt; # for text prompts the two views are the same list. - bank_match_ids = bank_key_ids if bank_key_ids is not None else prompt_ids + bank_match_ids = ( + as_token_tuple(bank_key_ids) if bank_key_ids is not None else prompt_tuple + ) exact_prefix_len = 0 try: longest_prefix = getattr(session_bank, "longest_prefix", None) @@ -6854,7 +6860,7 @@ def _emit_prefill_complete(state: PromptState) -> PromptState: tried_larger_near_prefix = True near_prompt_state = _restore_near_prefix_prompt_state( rt, - prompt_ids, + prompt_tuple, base_hidden_variant=base_hidden_variant, mtp_hidden_variant=mtp_hidden_variant, mtp_history_policy=mtp_history_policy, @@ -6967,7 +6973,7 @@ def _emit_prefill_complete(state: PromptState) -> PromptState: logits=restored.logits, hidden=restored.hidden, committed_mtp_cache=restored.mtp_history_cache, - token_prefix=tuple(int(token) for token in prompt_ids), + token_prefix=prompt_tuple, prompt_eval_time_s=repage_time, cache_restore_time_s=restore_elapsed_s, mtp_history_policy=mtp_history_policy, @@ -7064,7 +7070,7 @@ def _emit_prefill_complete(state: PromptState) -> PromptState: logits=suffix_logits, hidden=suffix_hidden, committed_mtp_cache=restored.mtp_history_cache, - token_prefix=tuple(int(token) for token in prompt_ids), + token_prefix=prompt_tuple, prompt_eval_time_s=suffix_time + mtp_history_time, prompt_mtp_history_time_s=mtp_history_time, cache_restore_time_s=restore_elapsed_s, @@ -7260,7 +7266,7 @@ def _emit_prefill_complete(state: PromptState) -> PromptState: logits=logits, hidden=hidden, committed_mtp_cache=mtp_history_cache, - token_prefix=tuple(int(token) for token in prompt_ids), + token_prefix=prompt_tuple, prompt_eval_time_s=prompt_eval_time, prompt_mtp_history_time_s=prompt_history_time, mtp_history_policy=mtp_history_policy, diff --git a/mtplx/server/openai.py b/mtplx/server/openai.py index 30fb37fb0..9eff3621d 100644 --- a/mtplx/server/openai.py +++ b/mtplx/server/openai.py @@ -82,6 +82,7 @@ from mtplx.adaptive import AdaptiveDepthPolicy, ExpectedValueDepthPolicy from mtplx.attention_context import attention_phase from mtplx.cache_state import snapshot_cache +from mtplx.token_ids import as_token_tuple from mtplx.one_copy import ( anchor_nbytes as one_copy_anchor_nbytes, one_copy_runtime, @@ -22888,7 +22889,7 @@ def _run_prefill_admission( prompt_tokens = len(prompt_ids) if prompt_tokens <= 0: return None - probe_ids = list(prompt_ids) + probe_ids = as_token_tuple(prompt_ids) if vision_splice is not None: # Admission must ask the same content-keyed question as restore. # Raw image pads only match the text before the first image; that @@ -22900,7 +22901,7 @@ def _run_prefill_admission( if keyed_ids is None: session_bank = None else: - probe_ids = list(keyed_ids) + probe_ids = as_token_tuple(keyed_ids) caps = getattr(state, "metal_memory_caps", None) limit = 0 if isinstance(caps, dict): diff --git a/mtplx/session_bank.py b/mtplx/session_bank.py index 458b38a55..65715663b 100644 --- a/mtplx/session_bank.py +++ b/mtplx/session_bank.py @@ -34,6 +34,7 @@ from .cache_bank.codec import ColdEncodeInterrupted from .checkpoint_anchors import AnchorPlan, default_tail_backoff, retain_checkpoints from .runtime import MTPLXRuntime +from .token_ids import TokenTuple, as_token_tuple from .runtime_options import block_prefix_restore_enabled @@ -324,6 +325,37 @@ def common_prefix_len(left: list[int] | tuple[int, ...], right: list[int] | tupl return limit +_SHARED_LEN_MEMO_MAX = 128 +_shared_len_memo: dict[tuple[int, int], tuple[Any, Any, int]] = {} +_shared_len_memo_lock = threading.Lock() + + +def shared_prefix_len(tokens: Any, prefix: Any) -> int: + """``common_prefix_len`` that remembers its answer per (prompt, entry). + + One request asks the same question of every RAM entry several times (the + dominant-prefix estimate, the exact lookup, the near-prefix lane, its + diagnostics, the reread receipt); at 80K tokens each pass costs ~0.45 ms + per entry. Both inputs are immutable ``TokenTuple`` objects, so the answer + is keyed on their identity; the memo holds both objects, which keeps the + identities from being reused while an answer is stored. Anything else + takes the plain compare. + """ + if type(tokens) is not TokenTuple or type(prefix) is not TokenTuple: + return common_prefix_len(tokens, prefix) + key = (id(tokens), id(prefix)) + with _shared_len_memo_lock: + hit = _shared_len_memo.get(key) + if hit is not None and hit[0] is tokens and hit[1] is prefix: + return hit[2] + matched = common_prefix_len(tokens, prefix) + with _shared_len_memo_lock: + if len(_shared_len_memo) >= _SHARED_LEN_MEMO_MAX: + _shared_len_memo.pop(next(iter(_shared_len_memo)), None) + _shared_len_memo[key] = (tokens, prefix, matched) + return matched + + def block_aligned_prefix_len(matched_tokens: int, *, block_size: int) -> int: block = max(1, int(block_size)) matched = max(0, int(matched_tokens)) @@ -484,7 +516,7 @@ def _near_prefix_candidate_len( prefix = entry.token_ids if not prefix: return None - matched = common_prefix_len(tokens, prefix) + matched = shared_prefix_len(tokens, prefix) # An image prompt is matched in its content-keyed view: the match is # content-true, but a restore never cuts an image (image_safe_restore_len). # A cut back to an image's start is no tiny gap; the block rules judge it. @@ -1417,7 +1449,7 @@ def put( # synchronous_serialize_elapsed_s}. Early returns leave later keys # absent. Deferred cold-tier serialization is never charged here — # only the dispatch span is; the job itself runs on the idle lane. - tokens = tuple(int(token) for token in token_ids) + tokens = as_token_tuple(token_ids) if not tokens: raise ValueError("cannot store an empty prefix") if mtp_snapshot_epoch is not None and int(mtp_snapshot_epoch) != int(snapshot_epoch): @@ -1789,7 +1821,7 @@ def put_snapshot( mtp_snapshot_epoch: int | None = None, nbytes_override: int | None = None, ) -> SessionBankEntry | None: - tokens = tuple(int(token) for token in token_ids) + tokens = as_token_tuple(token_ids) if not tokens: raise ValueError("cannot store an empty prefix") if mtp_snapshot_epoch is not None and int(mtp_snapshot_epoch) != int(snapshot_epoch): @@ -1891,18 +1923,18 @@ def longest_shared_prefix_tokens( and calls the session compacted. One compare per entry, no cold-tier scan, no lookup side effects. """ - tokens = tuple(int(token) for token in token_ids) + tokens = as_token_tuple(token_ids) best = 0 for prefix, entry in self._entries.items(): if session_id is not None and entry.session_id != session_id: continue - matched = common_prefix_len(tokens, prefix) + matched = shared_prefix_len(tokens, prefix) if matched > best: best = matched return best def longest_prefix(self, token_ids: list[int] | tuple[int, ...]) -> SessionBankEntry | None: - tokens = tuple(int(token) for token in token_ids) + tokens = as_token_tuple(token_ids) best: SessionBankEntry | None = None for prefix, entry in self._entries.items(): if len(prefix) > len(tokens): @@ -1938,7 +1970,7 @@ def near_prefix_candidates( if a new prompt shares a large stable token prefix, restore to the last full block and prefill only the changed suffix. """ - tokens = tuple(int(token) for token in token_ids) + tokens = as_token_tuple(token_ids) gap_limit = max(0, int(max_token_gap)) min_match = max(1, int(min_matched_tokens)) block = max(1, int(block_size)) @@ -1950,7 +1982,7 @@ def near_prefix_candidates( prefix = entry.token_ids if not prefix: continue - matched = common_prefix_len(tokens, prefix) + matched = shared_prefix_len(tokens, prefix) gap = len(prefix) - matched safe_block = min( block_aligned_prefix_len(matched, block_size=block), @@ -2123,7 +2155,7 @@ def near_prefix_candidates( "prompt_len": len(tokens), "session_id": entry.session_id, "stored_prefix_len": entry.prefix_len, - "common_prefix_tokens": int(common_prefix_len(tokens, entry.token_ids)), + "common_prefix_tokens": int(shared_prefix_len(tokens, entry.token_ids)), "nearest_boundary_tokens": int(matched), "new_prefill_tokens": max(0, len(tokens) - int(matched)), "miss_reason": None, @@ -2218,7 +2250,7 @@ def _cold_near_prefix_candidate( return None metadata = dict(getattr(record, "metadata", {}) or {}) entry = SessionBankEntry( - token_ids=tuple(int(token) for token in record.token_ids), + token_ids=as_token_tuple(record.token_ids), token_hash=metadata.get("token_hash") or token_prefix_hash(record.token_ids), has_recurrent=bool( getattr(record, "has_recurrent", False) @@ -2288,7 +2320,7 @@ def _oversized_skip_covering( if record is None: return None n = int(record["prefix_len"]) - tokens = tuple(int(token) for token in token_ids) + tokens = as_token_tuple(token_ids) if len(tokens) < n or token_prefix_hash(tokens[:n]) != record["token_hash"]: return None return record @@ -2524,7 +2556,7 @@ def cold_fallback() -> SessionBankRestore | None: entry.last_access_s = time.time() self.last_restore_source = "ram" self.last_ssd_restore_s = 0.0 - lookup_len = len(tuple(int(token) for token in token_ids)) + lookup_len = len(as_token_tuple(token_ids)) self.last_prefix_diagnostic = { "prompt_len": lookup_len, "session_id": entry.session_id, @@ -3316,7 +3348,7 @@ def run_live_ref_spill( views: later cache writes cannot mutate what gets encoded — the lazy-snapshot COW pin covers this), and streams to disk. """ - entry = self._entries.get(tuple(int(token) for token in token_ids)) + entry = self._entries.get(as_token_tuple(token_ids)) if entry is None or not entry.live_ref_only or entry.cache_ref is None: return False if int(entry.snapshot_epoch) != int(snapshot_epoch): @@ -3481,7 +3513,7 @@ def _restore_cold( return None metadata = dict(getattr(record, "metadata", {}) or {}) entry = SessionBankEntry( - token_ids=tuple(int(token) for token in record.token_ids), + token_ids=as_token_tuple(record.token_ids), token_hash=metadata.get("token_hash") or token_prefix_hash(record.token_ids), has_recurrent=bool( getattr(record, "has_recurrent", False) @@ -3538,7 +3570,7 @@ def _restore_cold( self.last_restore_source = "ssd" self.last_ssd_restore_s = float(getattr(record, "restore_s", 0.0) or 0.0) self.last_miss_reason = None - lookup_len = len(tuple(int(token) for token in token_ids)) + lookup_len = len(as_token_tuple(token_ids)) self.last_prefix_diagnostic = { "prompt_len": lookup_len, "session_id": entry.session_id, @@ -4091,7 +4123,7 @@ def restore_plan( } if not prompt_tokens: return plan - tokens = tuple(int(token) for token in prompt_tokens) + tokens = as_token_tuple(prompt_tokens) identity = { "model_path": model_path, "mtp_enabled": mtp_enabled, @@ -4220,7 +4252,7 @@ def session_coverage_tokens( if not session_id or not prompt_tokens: return 0 - tokens = tuple(int(token) for token in prompt_tokens) + tokens = as_token_tuple(prompt_tokens) knobs = _near_prefix_knobs() best = 0 for key, entry in list(self._entries.items()): @@ -4265,7 +4297,7 @@ def same_conversation_entries( if not session_id or not prompt_tokens: return [] - tokens = tuple(int(token) for token in prompt_tokens) + tokens = as_token_tuple(prompt_tokens) candidates: list[tuple[SessionBankEntry, bool]] = [ (entry, True) for entry in list(self._entries.values()) @@ -4283,7 +4315,7 @@ def same_conversation_entries( continue seen.add(id(entry)) length = len(entry.token_ids) - shared = common_prefix_len(tokens, entry.token_ids) + shared = shared_prefix_len(tokens, entry.token_ids) if length <= 0 or shared <= 0 or 2 * shared < length: continue rows.append( diff --git a/mtplx/token_ids.py b/mtplx/token_ids.py new file mode 100644 index 000000000..284b04bab --- /dev/null +++ b/mtplx/token_ids.py @@ -0,0 +1,24 @@ +"""Token id tuples that are converted once and recognised afterwards.""" + +from __future__ import annotations + +from typing import Iterable + + +class TokenTuple(tuple): + """A tuple whose elements are all plain ``int``. + + Equal to, hashing as, and iterating like the plain tuple. It exists so + that code receiving a request's prompt ids can tell "already converted" + in O(1) instead of re-running ``tuple(int(t) for t in ids)`` over + 80K+ tokens at every layer. + """ + + __slots__ = () + + +def as_token_tuple(token_ids: Iterable[int]) -> TokenTuple: + """Return ``token_ids`` as a tuple of ints, converting only when needed.""" + if type(token_ids) is TokenTuple: + return token_ids + return TokenTuple([int(token) for token in token_ids]) diff --git a/tests/test_cold_tier_row_token_memo.py b/tests/test_cold_tier_row_token_memo.py new file mode 100644 index 000000000..3cd64eade --- /dev/null +++ b/tests/test_cold_tier_row_token_memo.py @@ -0,0 +1,284 @@ +"""Prefix-boundary lookup: SQL length floor and parsed-row memo. + +The lookup must pick exactly what the unfiltered, unmemoized scan picks, with +the same miss reasons and counters. The reference below is that scan, written +out against the same manifest rows. +""" + +from __future__ import annotations + +import json +import random + +import pytest + +from mtplx.cache_bank import cold_tier as cold_tier_module +from mtplx.cache_bank.cold_tier import ( + ColdPrefixRestoreRecord, + SessionBankColdTier, + block_aligned_prefix_len, + common_prefix_len, +) + +BLOCK = 16 +GAP = 8 +MIN_MATCH = 8 + + +def _meta(entry_id, tokens, *, last_access=0.0): + tokens = list(tokens) + return { + "entry_id": entry_id, + "token_hash": f"hash-{entry_id}-{len(tokens)}", + "prefix_len": len(tokens), + "token_ids": tokens, + "model_path": "model", + "mtp_enabled": False, + "snapshot_epoch": len(tokens), + "capabilities": [], + "block_size": BLOCK, + "block_hashes": [], + "entry_dir": f"entries/{entry_id}", + "nbytes": 1, + "created_at_s": 0.0, + } + + +class _Stub: + def __init__(self, entry_id): + self.entry_id = entry_id + + +@pytest.fixture() +def tier(tmp_path, monkeypatch): + t = SessionBankColdTier(base_dir=tmp_path / "bank", mode="on", min_prefix_tokens=1) + calls = [] + + def fake_restore(row, tokens, **kw): + calls.append((str(row["entry_id"]), kw.get("prefix_restore_tokens"))) + return _Stub(str(row["entry_id"])) + + monkeypatch.setattr(t, "_restore_row", fake_restore) + t.calls = calls + yield t + t.close() + + +def _reference(tier, tokens, *, floor): + """The pre-change scan: every row parsed and ranked, floor applied last.""" + with tier._connect() as conn: + rows = list( + conn.execute( + "SELECT * FROM entries WHERE model_path = ? AND mtp_enabled = ? " + "ORDER BY prefix_len DESC, last_access_s DESC", + ("model", 0), + ).fetchall() + ) + block_min = max(BLOCK, BLOCK) + ranked = [] + for row in rows: + prefix = tuple(int(t) for t in json.loads(row["token_ids_json"])) + if not prefix: + continue + matched = common_prefix_len(tokens, prefix) + gap = len(prefix) - matched + safe = min( + block_aligned_prefix_len(matched, block_size=BLOCK), len(prefix), len(tokens) + ) + required = min(MIN_MATCH, max(1, len(prefix) - GAP)) + near = 0 <= gap <= GAP and matched >= required + block = safe >= block_min and safe >= 2 and safe <= matched + if near: + ranked.append(((matched, matched, len(prefix)), row["entry_id"], matched, "near_prefix")) + elif block: + ranked.append(((safe, matched, len(prefix)), row["entry_id"], safe, "block_prefix")) + if not ranked: + return None, "ssd_prefix_miss" + ranked.sort(key=lambda item: item[0], reverse=True) + if floor > 0 and ranked[0][2] < floor: + return None, "ssd_prefix_not_better_than_ram" + best = ranked[0] + return (best[1], best[2], best[3]), None + + +def _lookup(tier, tokens, floor): + tier.calls.clear() + return tier.lookup_prefix_boundary( + tokens, + model_path="model", + mtp_enabled=False, + max_token_gap=GAP, + min_matched_tokens=MIN_MATCH, + block_size=BLOCK, + block_min_matched_tokens=BLOCK, + min_useful_matched_tokens=floor, + ) + + +def _populate(tier, rng): + base = [rng.randrange(1000, 60000) for _ in range(400)] + lengths = [20, 64, 100, 128, 200, 256, 330, 400] + for i, n in enumerate(lengths): + toks = base[:n] + if i % 3 == 1: + toks = toks[: n - 5] + [7000 + i] * 5 # diverges near the end + if i % 4 == 3: + toks = toks[:40] + [8000 + j for j in range(n - 40)] # diverges early + tier._insert_manifest(_meta(f"e{i}", toks, last_access=i)) + return base + + +def test_filter_and_memo_match_reference(tier): + rng = random.Random(7) + base = _populate(tier, rng) + queries = [ + tuple(base[:n]) + tuple(rng.randrange(70000, 80000) for _ in range(k)) + for n in (30, 64, 129, 250, 331, 400) + for k in (0, 3, 40) + ] + for _round in range(2): # second round runs on a warm memo + for query in queries: + for floor in (0, 1, 64, 129, 200, 331, 400, 500): + expected, reason = _reference(tier, query, floor=floor) + got = _lookup(tier, query, floor) + if expected is None: + assert got is None + assert tier.stats()["last_miss_reason"] == reason + else: + assert isinstance(got, ColdPrefixRestoreRecord) + assert (got.record.entry_id, got.matched_tokens, got.restore_kind) == expected + assert tier.calls[0] == (expected[0], expected[1]) + + +def test_sql_floor_keeps_rows_at_the_boundary(tier): + toks = list(range(100)) + tier._insert_manifest(_meta("a", toks)) + query = tuple(toks[:100]) + (9999,) + # candidate_matched == prefix_len == floor: not below the floor, so it is kept. + got = _lookup(tier, query, 100) + assert got is not None and got.matched_tokens == 100 + assert _lookup(tier, query, 101) is None + assert tier.stats()["last_miss_reason"] == "ssd_prefix_not_better_than_ram" + + +def test_short_rows_still_decide_the_miss_reason(tier): + tier._insert_manifest(_meta("a", list(range(64)))) + query = tuple(range(64)) + (5000, 5001) + assert _lookup(tier, query, 500) is None + assert tier.stats()["last_miss_reason"] == "ssd_prefix_not_better_than_ram" + unrelated = tuple(range(9000, 9100)) + assert _lookup(tier, unrelated, 500) is None + assert tier.stats()["last_miss_reason"] == "ssd_prefix_miss" + + +def test_memo_parses_each_row_once(tier, monkeypatch): + toks = list(range(300)) + tier._insert_manifest(_meta("a", toks)) + parses = [] + real = json.loads + + def counting(text, *a, **k): + parses.append(1) + return real(text, *a, **k) + + monkeypatch.setattr(cold_tier_module.json, "loads", counting) + query = tuple(toks[:290]) + (1, 2, 3) + for _ in range(3): + assert _lookup(tier, query, 0) is not None + token_parses = len(parses) + assert token_parses == 1 + + +def test_memo_follows_delete_and_rewrite(tier): + a = list(range(200)) + tier._insert_manifest(_meta("a", a)) + query = tuple(a[:195]) + (9, 9, 9, 9, 9) + assert _lookup(tier, query, 0).matched_tokens == 195 + # Rewrite the same entry_id with different tokens (and token_hash). + b = [x + 100000 for x in a] + tier._insert_manifest(_meta("a", b)) + assert _lookup(tier, query, 0) is None + assert _lookup(tier, tuple(b[:195]) + (1, 2, 3, 4, 5), 0).matched_tokens == 195 + # Delete through the tier's own retire path. + with tier._connect() as conn: + row = conn.execute("SELECT * FROM entries WHERE entry_id = 'a'").fetchone() + tier._delete_entry_row(row) + assert "a" not in tier._row_tokens_memo + assert _lookup(tier, tuple(b[:195]) + (1, 2, 3, 4, 5), 0) is None + + +def test_memo_rewrite_with_same_id_by_other_process(tier): + """A row rewritten behind our back (other process) is caught by token_hash.""" + a = list(range(200)) + tier._insert_manifest(_meta("a", a)) + query = tuple(a[:195]) + (9, 9, 9, 9, 9) + assert _lookup(tier, query, 0) is not None + b = [x + 100000 for x in a] + with tier._connect() as conn: # bypasses _insert_manifest + conn.execute( + "UPDATE entries SET token_ids_json = ?, token_hash = ? WHERE entry_id = 'a'", + (json.dumps(b), "other-hash"), + ) + assert _lookup(tier, query, 0) is None + + +def test_memo_is_bounded(tier, monkeypatch): + monkeypatch.setattr(cold_tier_module, "_ROW_TOKENS_MEMO_MAX_TOKENS", 500) + for i in range(6): + tier._insert_manifest(_meta(f"e{i}", list(range(i, i + 200)))) + _lookup(tier, tuple(range(300)), 0) + assert tier._row_tokens_memo_tokens <= 500 + assert tier._row_tokens_memo_tokens == sum( + len(t) for _, t in tier._row_tokens_memo.values() + ) + + +def test_full_memo_keeps_residents_so_repeated_scans_still_hit(tier, monkeypatch): + """A scan over more tokens than the cap must not evict every row before its reuse.""" + monkeypatch.setattr(cold_tier_module, "_ROW_TOKENS_MEMO_MAX_TOKENS", 500) + for i in range(6): + tier._insert_manifest(_meta(f"e{i}", list(range(i, i + 200)))) + parses = [] + real = json.loads + + def counting(text, *a, **k): + parses.append(1) + return real(text, *a, **k) + + monkeypatch.setattr(cold_tier_module.json, "loads", counting) + query = tuple(range(300)) + _lookup(tier, query, 0) + first = len(parses) + _lookup(tier, query, 0) + second = len(parses) - first + assert len(tier._row_tokens_memo) == 2 + assert second == first - 2 + + +def test_memo_holds_compact_arrays(tier): + tier._insert_manifest(_meta("a", list(range(200)))) + _lookup(tier, tuple(range(195)) + (9, 9, 9, 9, 9), 0) + (_, tokens), = tier._row_tokens_memo.values() + assert isinstance(tokens, cold_tier_module.array) + assert tokens.typecode == "I" + + +def test_token_array_falls_back_to_a_tuple_for_ids_that_do_not_fit(): + assert cold_tier_module._token_array([1, 2**40]) == (1, 2**40) + assert cold_tier_module._token_array([-1]) == (-1,) + + +@pytest.mark.parametrize( + "tokens, prefix, expected", + [ + ((1, 2, 3, 4), (1, 2, 3), True), + ((1, 2, 3, 4), (1, 2, 4), False), + ((1, 2), (1, 2, 3), False), + ((1, 2, 3), (), True), + ((-1, 2), (-1,), True), + ], +) +def test_starts_with_matches_tuple_slicing_for_arrays_and_tuples(tokens, prefix, expected): + as_stored = cold_tier_module._token_array(list(prefix)) + assert cold_tier_module._starts_with(tokens, as_stored) is expected + assert (tokens[: len(prefix)] == prefix) is expected diff --git a/tests/test_session_bank_shared_prefix_memo.py b/tests/test_session_bank_shared_prefix_memo.py new file mode 100644 index 000000000..33b443927 --- /dev/null +++ b/tests/test_session_bank_shared_prefix_memo.py @@ -0,0 +1,84 @@ +"""One request, one compare per RAM entry: shared_prefix_len memo. + +Results and diagnostics of the bank lookups must be identical whether the +prompt arrives as a list (every call converts and compares on its own) or as +one normalized TokenTuple (compares shared across the lookups). +""" + +from __future__ import annotations + +import random +from pathlib import Path +from types import SimpleNamespace + +from mtplx import session_bank as sb +from mtplx.session_bank import SessionBank, common_prefix_len, shared_prefix_len +from mtplx.token_ids import as_token_tuple + + +def _bank(rng): + bank = SessionBank(max_entries=32, max_bytes=1 << 30, per_session_max_bytes=1 << 30) + runtime = SimpleNamespace(model_path=Path("models/example"), mtp_enabled=False) + base = [rng.randrange(1000, 60000) for _ in range(600)] + for i, n in enumerate((80, 150, 300, 420, 600)): + toks = base[:n] if i % 2 == 0 else base[: n - 20] + [90000 + i] * 20 + bank.put( + runtime=runtime, token_ids=toks, cache=[], logits=None, hidden=None, + session_id=f"s{i % 2}", nbytes_override=100, + ) + return bank, base + + +def test_shared_prefix_len_matches_plain(): + rng = random.Random(3) + for _ in range(50): + a = [rng.randrange(5) for _ in range(rng.randrange(0, 40))] + b = a[: rng.randrange(0, len(a) + 1)] + [rng.randrange(5) for _ in range(rng.randrange(0, 5))] + ta, tb = as_token_tuple(a), as_token_tuple(b) + want = common_prefix_len(a, b) + assert shared_prefix_len(ta, tb) == want + assert shared_prefix_len(ta, tb) == want # memo hit + assert shared_prefix_len(a, b) == want # plain inputs bypass the memo + assert shared_prefix_len(tuple(a), tb) == want + + +def test_memo_is_identity_keyed_and_bounded(): + a = as_token_tuple([1, 2, 3, 4]) + b = as_token_tuple([1, 2, 9]) + assert shared_prefix_len(a, b) == 2 + # An equal but different object is a separate key with the same answer. + assert shared_prefix_len(as_token_tuple([1, 2, 3, 4]), b) == 2 + for i in range(sb._SHARED_LEN_MEMO_MAX + 20): + shared_prefix_len(as_token_tuple([i]), b) + assert len(sb._shared_len_memo) <= sb._SHARED_LEN_MEMO_MAX + + +def test_bank_lookups_identical_for_list_and_token_tuple(): + rng = random.Random(11) + bank, base = _bank(rng) + kw = dict(max_token_gap=8, min_matched_tokens=16, block_size=16, block_min_matched_tokens=16) + prompts = [ + base[:n] + [rng.randrange(70000, 80000) for _ in range(k)] + for n in (50, 150, 299, 420, 600) + for k in (0, 5, 60) + ] + for prompt in prompts: + for session in (None, "s0", "s1"): + norm = as_token_tuple(prompt) + assert bank.longest_shared_prefix_tokens(prompt, session_id=session) == \ + bank.longest_shared_prefix_tokens(norm, session_id=session) + assert bank.longest_prefix(prompt) is bank.longest_prefix(as_token_tuple(prompt)) + got_list = bank.near_prefix_candidates(list(prompt), **kw) + diag_list = dict(bank.last_prefix_diagnostic or {}) + norm = as_token_tuple(prompt) + got_norm = bank.near_prefix_candidates(norm, **kw) + diag_norm = dict(bank.last_prefix_diagnostic or {}) + got_norm_again = bank.near_prefix_candidates(norm, **kw) + assert [(e.token_hash, m) for e, m in got_list] == [(e.token_hash, m) for e, m in got_norm] + assert [(e.token_hash, m) for e, m in got_norm] == [ + (e.token_hash, m) for e, m in got_norm_again + ] + assert diag_list == diag_norm + assert bank.same_conversation_entries("s0", prompt) == bank.same_conversation_entries( + "s0", norm + ) diff --git a/tests/test_token_ids.py b/tests/test_token_ids.py new file mode 100644 index 000000000..606c95a48 --- /dev/null +++ b/tests/test_token_ids.py @@ -0,0 +1,19 @@ +"""TokenTuple: convert a prompt's ids once, recognise them afterwards.""" + +from __future__ import annotations + +import pytest + +from mtplx.token_ids import TokenTuple, as_token_tuple + + +def test_token_tuple_fast_path(): + plain = tuple(range(1000)) + first = as_token_tuple(plain) + assert type(first) is TokenTuple and first == plain and hash(first) == hash(plain) + assert as_token_tuple(first) is first + assert as_token_tuple([1, 2, 3]) == (1, 2, 3) + assert type(as_token_tuple(iter([1, 2]))) is TokenTuple + np = pytest.importorskip("numpy") + converted = as_token_tuple(np.arange(5)) + assert all(type(t) is int for t in converted)