From 09b3d6e60abbc674dc5496395014dd6cd919bf53 Mon Sep 17 00:00:00 2001
From: Jeroen van Menen <joonix@MacBook-Pro-van-Jeroen.local>
Date: Tue, 13 Oct 2026 15:11:01 +0200
Subject: [PATCH 1/4] Cold tier: skip rows shorter than the RAM match and
 memoize parsed token lists in prefix lookups

lookup_prefix_boundary read and json-parsed the token list of every manifest
row on each request, then dropped the rows under min_useful_matched_tokens.
Rows with prefix_len below that floor cannot rank (matched <= prefix_len), so
the query now leaves them out; they are only inspected when nothing else
ranks, so the miss reason stays the same. Parsed token tuples are memoized
per entry (keyed on entry_id and token_hash, bounded by total tokens, dropped
on delete and rewrite) and shared with _restore_row.
---
 mtplx/cache_bank/cold_tier.py          | 105 +++++++++--
 tests/test_cold_tier_row_token_memo.py | 233 +++++++++++++++++++++++++
 2 files changed, 324 insertions(+), 14 deletions(-)
 create mode 100644 tests/test_cold_tier_row_token_memo.py

diff --git a/mtplx/cache_bank/cold_tier.py b/mtplx/cache_bank/cold_tier.py
index d11145079..ec9a5c5d9 100644
--- a/mtplx/cache_bank/cold_tier.py
+++ b/mtplx/cache_bank/cold_tier.py
@@ -313,6 +313,10 @@ def block_aligned_prefix_len(matched_tokens: int, *, block_size: int) -> int:
     return (matched // block) * block
 
 
+# Upper bound on token ids held by the parsed-row memo (a few hundred MB at
+# most for the largest rows; typical stores stay far below it).
+_ROW_TOKENS_MEMO_MAX_TOKENS = 2_500_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 +468,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: OrderedDict[str, tuple[str, tuple[int, ...]]] = OrderedDict()
+        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 +1179,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 +1188,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 +1219,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 +2433,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 +2448,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,7 +2704,7 @@ 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"])))
+        token_ids = self._row_token_ids(row)
         if require_exact_prefix and lookup_tokens[: len(token_ids)] != token_ids:
             return None
         if int(metadata["format_version"]) != COLD_TIER_FORMAT_VERSION:
@@ -2861,6 +2901,34 @@ def _candidate_rows(
                 compatible.append(row)
         return compatible
 
+    def _row_token_ids(self, row: sqlite3.Row) -> tuple[int, ...]:
+        """Token ids of a manifest row, parsed once per (entry_id, token_hash)."""
+        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:
+                self._row_tokens_memo.move_to_end(entry_id)
+                return hit[1]
+        tokens = tuple(int(token) for token in 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 len(tokens) <= _ROW_TOKENS_MEMO_MAX_TOKENS:
+                self._row_tokens_memo[entry_id] = (token_hash, tokens)
+                self._row_tokens_memo_tokens += len(tokens)
+                while self._row_tokens_memo_tokens > _ROW_TOKENS_MEMO_MAX_TOKENS:
+                    _, (_, evicted) = self._row_tokens_memo.popitem(last=False)
+                    self._row_tokens_memo_tokens -= len(evicted)
+        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 +2939,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 +3297,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/tests/test_cold_tier_row_token_memo.py b/tests/test_cold_tier_row_token_memo.py
new file mode 100644
index 000000000..e9b3caa62
--- /dev/null
+++ b/tests/test_cold_tier_row_token_memo.py
@@ -0,0 +1,233 @@
+"""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()
+    )

From fc79dce18049de3b4251211294f9f32a8ff8c105 Mon Sep 17 00:00:00 2001
From: Jeroen van Menen <joonix@MacBook-Pro-van-Jeroen.local>
Date: Tue, 13 Oct 2026 15:11:07 +0200
Subject: [PATCH 2/4] Convert a request's prompt ids to a token tuple once

The prompt was rebuilt with tuple(int(t) for t in ids) 8-12 times per
request (session bank lookups, engine session, token_prefix in generation),
about 2.5 ms each at 85K tokens. as_token_tuple returns a TokenTuple, a tuple
subclass that marks "all plain ints", and returns it unchanged when it is
handed one again. The prompt is normalized where restore_or_prefill_prompt_state
and the admission probe receive it; lists and other sequences still work
everywhere.
---
 mtplx/engine_session.py | 17 +++++++++--------
 mtplx/generation.py     | 22 ++++++++++++++--------
 mtplx/server/openai.py  |  5 +++--
 mtplx/session_bank.py   | 29 +++++++++++++++--------------
 mtplx/token_ids.py      | 24 ++++++++++++++++++++++++
 tests/test_token_ids.py | 19 +++++++++++++++++++
 6 files changed, 84 insertions(+), 32 deletions(-)
 create mode 100644 mtplx/token_ids.py
 create mode 100644 tests/test_token_ids.py

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..194c1beb3 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 as_token_tuple
 from .runtime_options import block_prefix_restore_enabled
 
 
@@ -1417,7 +1418,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 +1790,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,7 +1892,7 @@ 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:
@@ -1902,7 +1903,7 @@ def longest_shared_prefix_tokens(
         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 +1939,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))
@@ -2218,7 +2219,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 +2289,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 +2525,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 +3317,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 +3482,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 +3539,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 +4092,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 +4221,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 +4266,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())
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_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)

From d1a0542a0303e741f600879812a2ea60f85b8fa2 Mon Sep 17 00:00:00 2001
From: Jeroen van Menen <joonix@MacBook-Pro-van-Jeroen.local>
Date: Tue, 13 Oct 2026 15:11:12 +0200
Subject: [PATCH 3/4] Session bank: compute each entry's shared prefix length
 once per request

dominant-prefix estimate, exact lookup, near-prefix candidates with their
diagnostics and the reread receipt each compared the prompt with every RAM
entry (~0.45 ms per entry at 80K tokens). shared_prefix_len remembers the
answer per (prompt, entry) pair when both are TokenTuple objects, so the later
passes reuse it. Plain lists and tuples take the old compare.
---
 mtplx/session_bank.py                         | 43 ++++++++--
 tests/test_session_bank_shared_prefix_memo.py | 84 +++++++++++++++++++
 2 files changed, 121 insertions(+), 6 deletions(-)
 create mode 100644 tests/test_session_bank_shared_prefix_memo.py

diff --git a/mtplx/session_bank.py b/mtplx/session_bank.py
index 194c1beb3..65715663b 100644
--- a/mtplx/session_bank.py
+++ b/mtplx/session_bank.py
@@ -34,7 +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 as_token_tuple
+from .token_ids import TokenTuple, as_token_tuple
 from .runtime_options import block_prefix_restore_enabled
 
 
@@ -325,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))
@@ -485,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.
@@ -1897,7 +1928,7 @@ def longest_shared_prefix_tokens(
         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
@@ -1951,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),
@@ -2124,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,
@@ -4284,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/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
+        )

From 69c21d08fca1c33a2cde62ee68b408f12972fd37 Mon Sep 17 00:00:00 2001
From: Jeroen van Menen <joonix@MacBook-Pro-van-Jeroen.local>
Date: Wed, 21 Oct 2026 18:42:42 +0200
Subject: [PATCH 4/4] Cold tier: keep the parsed-row memo hitting when the
 store outgrows it

The prefix lookup walks the candidate rows in the same order on every turn.
With more tokens in those rows than the memo cap (2.5M), evicting the least
recently used row dropped each row just before its next use, so the memo
never hit and every turn parsed every row again. On a store of 79 rows and
5.7M tokens a lookup for a 56K prompt took 343 ms with 0 of 306 rows from the
memo, against 8.6 ms when the rows stay parsed.

Rows are now held as array('I') (4 bytes per id instead of about 40 for a
tuple of ints), the cap is 32M tokens (about 128 MB), and a full memo keeps
its residents and returns a newcomer unmemoized instead of evicting. The
exact-prefix check in _restore_row compares through _starts_with, which
accepts both the array and the tuple fallback for ids that do not fit.
---
 mtplx/cache_bank/cold_tier.py          | 55 +++++++++++++++++++-------
 tests/test_cold_tier_row_token_memo.py | 51 ++++++++++++++++++++++++
 2 files changed, 92 insertions(+), 14 deletions(-)

diff --git a/mtplx/cache_bank/cold_tier.py b/mtplx/cache_bank/cold_tier.py
index ec9a5c5d9..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,15 +308,36 @@ 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 (a few hundred MB at
-# most for the largest rows; typical stores stay far below it).
-_ROW_TOKENS_MEMO_MAX_TOKENS = 2_500_000
+# 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.
@@ -471,7 +493,7 @@ def __init__(
         # 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: OrderedDict[str, tuple[str, tuple[int, ...]]] = OrderedDict()
+        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()
@@ -2705,7 +2727,7 @@ def _restore_row(
     ) -> ColdRestoreRecord | None:
         metadata = dict(row)
         token_ids = self._row_token_ids(row)
-        if require_exact_prefix and lookup_tokens[: len(token_ids)] != token_ids:
+        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")
@@ -2843,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,
@@ -2901,26 +2925,29 @@ def _candidate_rows(
                 compatible.append(row)
         return compatible
 
-    def _row_token_ids(self, row: sqlite3.Row) -> tuple[int, ...]:
-        """Token ids of a manifest row, parsed once per (entry_id, token_hash)."""
+    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:
-                self._row_tokens_memo.move_to_end(entry_id)
                 return hit[1]
-        tokens = tuple(int(token) for token in json.loads(str(row["token_ids_json"])))
+        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 len(tokens) <= _ROW_TOKENS_MEMO_MAX_TOKENS:
+            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)
-                while self._row_tokens_memo_tokens > _ROW_TOKENS_MEMO_MAX_TOKENS:
-                    _, (_, evicted) = self._row_tokens_memo.popitem(last=False)
-                    self._row_tokens_memo_tokens -= len(evicted)
         return tokens
 
     def _forget_row_tokens(self, entry_id: str) -> None:
diff --git a/tests/test_cold_tier_row_token_memo.py b/tests/test_cold_tier_row_token_memo.py
index e9b3caa62..3cd64eade 100644
--- a/tests/test_cold_tier_row_token_memo.py
+++ b/tests/test_cold_tier_row_token_memo.py
@@ -231,3 +231,54 @@ def test_memo_is_bounded(tier, monkeypatch):
     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