SSD cold tier: the streaming spill uses the incremental-encode memo too

CHANGELOG.md

@@ -4,6 +4,12 @@ All notable user-facing changes to MTPLX. The format is based on
 [Keep a Changelog](https://keepachangelog.com/en/1.1.0/); versions follow
 [Semantic Versioning](https://semver.org/).
 
+## [Unreleased]
+
+### Changed
+
+- **`MTPLX_SSD_INCREMENTAL_ENCODE` also reaches the streaming SSD write** (by @jvmenen). Large sessions, live-reference sessions and the write before an idle release go through `spill_entry`, which ignored the switch: every block, recurrent state and checkpoint was read back and hashed on every write, and `incremental_reused_blobs` stayed 0. The spill now uses the same per-session memo as the queued write, references a block only when its blob is still on disk (otherwise it writes it again), and keeps the store byte-identical to a full encode. The fingerprint pass runs in chunks of one encode unit (32 MiB) instead of one pass per tensor. The switch stays off by default.
+
 ## [2.12.2] - 2026-10-03
 
 ### Fixed

mtplx/cache_bank/codec.py

@@ -171,6 +171,12 @@ def __init__(
         self.reuse = reuse
         self.reused: dict[str, dict[str, Any]] = {}
         self.fingerprints: dict[str, tuple] = {}
+        # Streaming writer hook (cold_tier.spill_entry): a memo hit is
+        # referenced only once known_blob confirms the blob (it claims the
+        # digest and requires it on disk, returning the blob, or None to
+        # capture the bytes after all). The staged writer checks its
+        # referenced blobs itself and leaves this unset.
+        self.known_blob: Callable[[dict[str, Any]], dict[str, Any] | None] | None = None
         self.tensors: dict[str, bytes] = {}
         self.block_size = max(1, int(block_size))
         self.should_abort = should_abort
@@ -221,6 +227,42 @@ def _fingerprints(self, value: Any, *, rows: int | None) -> list[int] | None:
         self._check_abort()
         return fps
 
+    def _fingerprint_chunk_blocks(self, value: Any, shape: list[int]) -> int:
+        """Blocks per fingerprint chunk: as many as fit in one unit."""
+        rows = max(1, int(shape[2]))
+        block_bytes = max(1, int(getattr(value, "nbytes", 0) or 0) // rows) * self.block_size
+        return max(1, self.unit_bytes // block_bytes)
+
+    def _block_fingerprint(
+        self,
+        value: Any,
+        shape: list[int],
+        block_index: int,
+        *,
+        full: int,
+        chunk_blocks: int,
+        chunk_fps: dict[int, list[int] | None] | None,
+    ) -> int | None:
+        """The content fingerprint of one full block, or None (no memo key)."""
+        if chunk_fps is None or (block_index + 1) * self.block_size > full:
+            return None
+        chunk = block_index // chunk_blocks
+        if chunk not in chunk_fps:
+            lo = chunk * chunk_blocks * self.block_size
+            slices = [slice(None)] * len(shape)
+            slices[2] = slice(lo, min(full, lo + chunk_blocks * self.block_size))
+            chunk_fps[chunk] = self._fingerprints(
+                value[tuple(slices)], rows=self.block_size
+            )
+        part = chunk_fps[chunk]
+        offset = block_index - chunk * chunk_blocks
+        return part[offset] if part is not None and offset < len(part) else None
+
+    def _confirmed(self, hit: dict[str, Any] | None) -> dict[str, Any] | None:
+        if hit is None or self.known_blob is None:
+            return hit
+        return self.known_blob(hit)
+
     def _check_abort(self) -> None:
         check = self.should_abort
         if check is None:
@@ -283,7 +325,7 @@ def _encode_tensor(self, value: Any) -> dict[str, Any]:
             if fps:
                 key = (dtype, tuple(shape), -1, fps[0])
                 self.fingerprints[name] = key
-                hit = self.reuse.get(key)
+                hit = self._confirmed(self.reuse.get(key))
                 if hit is not None:
                     self.reused[name] = dict(hit)
                     return {
@@ -336,15 +378,15 @@ def _encode_tensor_blocks(
         axis = 2
         blocks: list[dict[str, Any]] = []
         total = 0
-        block_fps: list[int] | None = None
-        if self.reuse is not None:
-            full = (shape[axis] // self.block_size) * self.block_size
-            if full:
-                slices = [slice(None)] * len(shape)
-                slices[axis] = slice(0, full)
-                block_fps = self._fingerprints(
-                    value[tuple(slices)], rows=self.block_size
-                )
+        full = (shape[axis] // self.block_size) * self.block_size
+        # Fingerprints per chunk of whole blocks of at most unit_bytes, taken the
+        # first time a block of the chunk needs one: a bounded unit like every
+        # capture, and the reduction's temporaries stay a chunk's size on the
+        # streaming writer, which also runs when memory is short.
+        chunk_fps: dict[int, list[int] | None] | None = (
+            {} if self.reuse is not None and full else None
+        )
+        chunk_blocks = self._fingerprint_chunk_blocks(value, shape)
         for start in range(0, shape[axis], self.block_size):
             # No check here: _encode_tensor checked before the first block and
             # _materialize checks after every block.
@@ -355,11 +397,17 @@ def _encode_tensor_blocks(
             name = f"tensor_{self._next_tensor_id:08d}"
             self._next_tensor_id += 1
             block_index = start // self.block_size
-            if block_fps is not None and block_index < len(block_fps):
+            fingerprint = self._block_fingerprint(
+                value, shape, block_index, full=full, chunk_blocks=chunk_blocks,
+                chunk_fps=chunk_fps,
+            )
+            if fingerprint is not None:
                 chunk_shape = [int(dim) for dim in chunk.shape]
-                key = (dtype, tuple(chunk_shape), int(start), block_fps[block_index])
+                key = (dtype, tuple(chunk_shape), int(start), fingerprint)
                 self.fingerprints[name] = key
-                hit = self.reuse.get(key) if self.reuse is not None else None
+                hit = self._confirmed(
+                    self.reuse.get(key) if self.reuse is not None else None
+                )
                 if hit is not None:
                     nbytes = int(hit["nbytes"])
                     self.reused[name] = dict(hit)

mtplx/cache_bank/cold_tier.py

@@ -449,9 +449,10 @@ def __init__(
         # agent turn re-encodes the whole session on the model-owner thread
         # (eval + host copy of every KV block) and the writer re-hashes every
         # blob, although only the new tail is new on disk. With this on,
-        # put_entry fingerprints every KV block on the GPU (one reduction per
-        # tensor) and references a block whose content key matches the
-        # session's last completed write instead of capturing it. Off by
+        # put_entry and spill_entry fingerprint every KV block on the GPU (one
+        # reduction per tensor) and reference a block whose content key
+        # matches the session's last completed write instead of capturing it.
+        # Off by
         # default: the fingerprint is not a content identity. Flipping the
         # signs of any two float32 elements leaves both of its sums unchanged,
         # so a changed tensor can borrow its predecessor's blob and restore
@@ -978,12 +979,39 @@ def __setitem__(self, name: str, raw: bytes) -> None:  # noqa: N804
                 else:
                     written_state["deduped_hits"] += 1
 
+        # Incremental encode: the same per-session memo as put_entry. A hit
+        # is referenced only when _known_blob finds the blob on disk; the
+        # memo is replaced by this write's keys once its row is installed.
+        memo_key: str | None = None
+        reuse: dict[tuple, dict[str, Any]] | None = None
+        if self._incremental_encode:
+            session_id = getattr(entry, "session_id", None)
+            memo_key = str(session_id) if session_id else None
+            reuse = dict(self._block_memo(memo_key))
         codec = TreeCodec(
             block_size=self.block_size,
             should_abort=should_abort,
             on_unit=self._observe_encode_unit,
+            reuse=reuse,
         )
         codec.tensors = _WriteThroughTensors()
+
+        def _known_blob(blob: dict[str, Any]) -> dict[str, Any] | None:
+            # Claim first (so the orphan cleanup and eviction keep it), then
+            # require it on disk; a blob gone since the memo was taken is
+            # captured and written again.
+            digest = str(blob["sha256"])
+            if digest not in claimed_digests:
+                claimed_digests.add(digest)
+                tier._claim_inflight(digests={digest})
+            if not tier._blob_path(digest).exists():
+                tier._inc("incremental_missing_blobs")
+                return None
+            written_state["logical"] += int(blob["nbytes"])
+            written_state["deduped_hits"] += 1
+            return blob
+
+        codec.known_blob = _known_blob
         boundaries = tuple(
             (int(r[0]), r[1], r[2] if len(r) > 2 else None)
             for r in (getattr(entry, "gdn_boundaries", None) or [])
@@ -1015,6 +1043,15 @@ def __setitem__(self, name: str, raw: bytes) -> None:  # noqa: N804
                 "SessionBank SSD spill skipped: %s: %s", type(exc).__name__, exc
             )
             return False
+        for name, blob in codec.reused.items():
+            tensor_blobs[name] = dict(blob)
+        if codec.reused:
+            self._inc("incremental_encodes")
+            self._inc("incremental_reused_blobs", len(codec.reused))
+            self._inc(
+                "incremental_reused_bytes",
+                sum(int(blob["nbytes"]) for blob in codec.reused.values()),
+            )
         if should_abort is not None:
             # Same fencing rationale as put_entry: the per-tensor evals
             # above schedule GPU work whose command buffers must drain in
@@ -1052,6 +1089,8 @@ def __setitem__(self, name: str, raw: bytes) -> None:  # noqa: N804
             return False
         if outcome == _INSTALL_PRESENT:
             return True
+        if codec.fingerprints:
+            self._store_block_memo(memo_key, codec.fingerprints, tensor_blobs)
         self._inc("writes_completed")
         self._inc("spill_writes_completed")
         with self._stats_lock:
@@ -1834,7 +1873,9 @@ def _write_pending(self, pending: PendingWrite) -> bool:
                 entry_hash_prefix=entry_hash_prefix,
             )
             if wrote and pending.fingerprints:
-                self._store_block_memo(pending, tensor_blobs)
+                self._store_block_memo(
+                    pending.memo_key, pending.fingerprints, tensor_blobs
+                )
             return wrote
         finally:
             self._release_inflight(entry_dirs=(entry_dir_rel,), digests=digests)
@@ -1924,24 +1965,25 @@ def _drop_block_memo(self, memo_key: str | None) -> None:
 
     def _store_block_memo(
         self,
-        pending: PendingWrite,
+        memo_key: str | None,
+        fingerprints: dict[str, tuple],
         tensor_blobs: dict[str, dict[str, Any]],
     ) -> None:
         """Newest completed write per session becomes the reuse source: its
         content keys map to the digests the manifest row now references."""
-        if pending.memo_key is None:
+        if memo_key is None:
             return
         memo = {
             key: {
                 "sha256": str(tensor_blobs[name]["sha256"]),
                 "nbytes": int(tensor_blobs[name]["nbytes"]),
             }
-            for name, key in pending.fingerprints.items()
+            for name, key in fingerprints.items()
             if name in tensor_blobs
         }
         with self._block_memos_lock:
-            self._block_memos[pending.memo_key] = memo
-            self._block_memos.move_to_end(pending.memo_key)
+            self._block_memos[memo_key] = memo
+            self._block_memos.move_to_end(memo_key)
             while len(self._block_memos) > self._block_memos_max:
                 self._block_memos.popitem(last=False)