From 5ab5f86f4019aff38acbc736de1a2b230d9c699f Mon Sep 17 00:00:00 2001 From: Jeroen van Menen <235760+jvmenen@users.noreply.github.com> Date: Tue, 20 Oct 2026 13:28:14 +0200 Subject: [PATCH 1/3] SSD cold tier: the streaming spill uses the incremental-encode memo too MTPLX_SSD_INCREMENTAL_ENCODE only reached put_entry. Entries at or above the writer backlog budget, live-ref sessions and the write-out before an idle release all stream through spill_entry, which built its codec without the session's block memo: every KV block, recurrent state and checkpoint was evaluated, copied to the host and hashed again, and incremental_reused_blobs stayed 0 on a server that writes only large sessions. spill_entry now passes the same per-session memo and replaces it with the write's content keys once its row is installed, so queued and streamed writes of one session feed each other. A memo hit is referenced only after known_blob claims the digest and finds the blob on disk; a blob gone since then is captured and written in the same spill (incremental_missing_blobs) instead of failing the write. Checkpoints a turn keeps from the previous one are referenced like KV blocks. The store stays byte-identical to a full encode (tests for both writers). --- mtplx/cache_bank/codec.py | 17 +- mtplx/cache_bank/cold_tier.py | 60 ++++++- tests/test_cold_tier_incremental_encode.py | 191 +++++++++++++++++++++ 3 files changed, 257 insertions(+), 11 deletions(-) diff --git a/mtplx/cache_bank/codec.py b/mtplx/cache_bank/codec.py index 52503e5bc..2d3cf40a1 100644 --- a/mtplx/cache_bank/codec.py +++ b/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,11 @@ def _fingerprints(self, value: Any, *, rows: int | None) -> list[int] | None: self._check_abort() return fps + 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 +294,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 { @@ -359,7 +370,9 @@ def _encode_tensor_blocks( chunk_shape = [int(dim) for dim in chunk.shape] key = (dtype, tuple(chunk_shape), int(start), block_fps[block_index]) 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) diff --git a/mtplx/cache_bank/cold_tier.py b/mtplx/cache_bank/cold_tier.py index d11145079..9d293fc32 100644 --- a/mtplx/cache_bank/cold_tier.py +++ b/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) diff --git a/tests/test_cold_tier_incremental_encode.py b/tests/test_cold_tier_incremental_encode.py index 56589cf81..2c263d62b 100644 --- a/tests/test_cold_tier_incremental_encode.py +++ b/tests/test_cold_tier_incremental_encode.py @@ -449,3 +449,194 @@ def test_encode_payload_without_reuse_is_unchanged(): assert legacy.tensors == empty_memo.tensors assert legacy.reused == {} and legacy.fingerprints == {} assert empty_memo.reused == {} and empty_memo.fingerprints + + +# ---- spill_entry (the streaming writer) ---------------------------------------- + + +def _spill(tier: SessionBankColdTier, entry) -> str: + assert tier.spill_entry(entry) is True + return tier._metadata_for_entry(entry, capabilities=(), payload_nbytes=0)["entry_id"] + + +def _with_boundaries(entry, records): + entry.gdn_boundaries = tuple(records) + return entry + + +def _boundary(tokens, salt): + """A recurrent checkpoint as the bank keeps one: a 1 MiB float32 state + (a whole-tensor fingerprint) and a bfloat16 convolution state long + enough on axis 2 to be stored in blocks.""" + return ( + int(tokens), + CacheSnapshot( + states=(None, (_rand((1, 8, 128, 256), mx.float32, salt), + _rand((1, 3, 1024), mx.bfloat16, salt + 1))), + meta_states=(None, {"kind": "gdn"}), + ), + _rand((1, 1, 8), mx.bfloat16, salt + 2), + ) + + +def _spill_turns(): + """Two turns of one session as the spill path sees them: turn 2 extends + turn 1's KV and keeps its first recurrent checkpoint, adding one.""" + turn1, turn2 = _turns() + first = _boundary(512, 40) + _with_boundaries(turn1, [first]) + _with_boundaries(turn2, [first, _boundary(1024, 50)]) + return turn1, turn2 + + +@pytest.mark.parametrize("dtype", [mx.bfloat16, mx.float32]) +@pytest.mark.parametrize("seed", [1, 2]) +def test_spill_every_turn_shape_stores_what_the_full_encode_stores( + tmp_path, monkeypatch, dtype, seed +): + monkeypatch.setenv("MTPLX_SSD_INCREMENTAL_ENCODE", "0") + full = _tier(tmp_path / "full") + monkeypatch.setenv("MTPLX_SSD_INCREMENTAL_ENCODE", "1") + incremental = _tier(tmp_path / "incremental") + try: + for turn in _session_turns(dtype, seed): + entry_id = _spill(full, turn) + assert _spill(incremental, turn) == entry_id + assert _payload(incremental, entry_id) == _payload(full, entry_id) + assert _blobs(incremental) == _blobs(full) + stats = incremental.stats() + assert stats["incremental_reused_blobs"] > 0 + assert stats["incremental_missing_blobs"] == 0 + assert full.stats()["incremental_reused_blobs"] == 0 + finally: + full.close() + incremental.close() + + +def test_spill_references_kv_and_an_unchanged_checkpoint(tmp_path, monkeypatch): + monkeypatch.setenv("MTPLX_SSD_INCREMENTAL_ENCODE", "0") + full = _tier(tmp_path / "full") + monkeypatch.setenv("MTPLX_SSD_INCREMENTAL_ENCODE", "1") + incremental = _tier(tmp_path / "incremental") + try: + for turn in _spill_turns(): + entry_id = _spill(full, turn) + assert _spill(incremental, turn) == entry_id + assert _payload(incremental, entry_id) == _payload(full, entry_id) + assert _blobs(incremental) == _blobs(full) + stats = incremental.stats() + # Turn 2: 4 full blocks each of K, V and the MTP history, the 1 MiB + # final recurrent state, and the kept checkpoint's 1 MiB state plus + # the 4 blocks of its convolution state. + assert stats["incremental_encodes"] == 1 + assert stats["incremental_reused_blobs"] == 4 * 3 + 1 + 1 + 4 + # A full encode of turn 2 captures each of those again. + assert full.stats()["incremental_reused_blobs"] == 0 + finally: + full.close() + incremental.close() + + +def test_spill_restore_after_restart_is_bit_identical(tmp_path, monkeypatch): + monkeypatch.setenv("MTPLX_SSD_INCREMENTAL_ENCODE", "1") + base = tmp_path / "store" + tier = _tier(base) + turn1, turn2 = _spill_turns() + try: + _spill(tier, turn1) + _spill(tier, turn2) + assert tier.stats()["incremental_reused_blobs"] > 0 + finally: + tier.close() + monkeypatch.setenv("MTPLX_SSD_INCREMENTAL_ENCODE", "0") + reopened = _tier(base) + try: + record = reopened.lookup( + list(turn2.token_ids) + [7, 7], model_path=MODEL, mtp_enabled=True + ) + assert record is not None + assert record.token_ids == turn2.token_ids + (k, v), (gdn,) = record.cache_snapshot.states + (k2, v2), (gdn2,) = turn2.cache_snapshot.states + assert _bits_equal(k, k2) + assert _bits_equal(v, v2) + assert _bits_equal(gdn, gdn2) + assert _bits_equal( + record.mtp_history_snapshot.states[0][0], + turn2.mtp_history_snapshot.states[0][0], + ) + boundaries = record.gdn_boundary_loader() + assert len(boundaries) == 2 + for restored, source in zip(boundaries, turn2.gdn_boundaries): + assert restored[0] == source[0] + for a, b in zip(restored[1].states[1], source[1].states[1]): + assert _bits_equal(a, b) + assert _bits_equal(restored[2], source[2]) + finally: + reopened.close() + + +def test_spill_captures_a_memo_blob_that_left_the_disk(tmp_path, monkeypatch): + """A blob the memo names but the disk lost (its entries were evicted) + is captured and written again in the same spill; the write lands.""" + monkeypatch.setenv("MTPLX_SSD_INCREMENTAL_ENCODE", "0") + full = _tier(tmp_path / "full") + monkeypatch.setenv("MTPLX_SSD_INCREMENTAL_ENCODE", "1") + tier = _tier(tmp_path / "store") + try: + turn1, turn2 = _spill_turns() + _spill(full, turn1) + _spill(tier, turn1) + memo = tier._block_memo(SESSION) + victim = next(iter(memo.values()))["sha256"] + tier._blob_path(victim).unlink() + + entry_id = _spill(tier, turn2) + stats = tier.stats() + assert stats["incremental_missing_blobs"] == 1 + assert stats["incremental_reused_blobs"] > 0 + assert tier._blob_path(victim).exists() + _spill(full, turn2) + # The same entry; only the bytes written differ (the victim twice). + stored, reference = _payload(tier, entry_id), _payload(full, entry_id) + assert stored["payload_spec"] == reference["payload_spec"] + assert stored["tensor_blobs"] == reference["tensor_blobs"] + record = tier.lookup(list(turn2.token_ids), model_path=MODEL, mtp_enabled=True) + assert record is not None + assert _bits_equal( + record.cache_snapshot.states[0][0], turn2.cache_snapshot.states[0][0] + ) + finally: + full.close() + tier.close() + + +def test_spill_other_session_does_not_borrow(tmp_path, monkeypatch): + monkeypatch.setenv("MTPLX_SSD_INCREMENTAL_ENCODE", "1") + tier = _tier(tmp_path / "store") + try: + turn1, turn2 = _spill_turns() + _spill(tier, turn1) + turn2.session_id = "another-session" + _spill(tier, turn2) + stats = tier.stats() + assert stats["incremental_reused_blobs"] == 0 + finally: + tier.close() + + +def test_spill_and_queued_writes_share_the_session_memo(tmp_path, monkeypatch): + """A queued write's memo serves the next spill of the session, and the + other way round.""" + monkeypatch.setenv("MTPLX_SSD_INCREMENTAL_ENCODE", "1") + tier = _tier(tmp_path / "store") + try: + turns = _session_turns(mx.bfloat16, 4) + _write(tier, turns[0]) + _spill(tier, turns[1]) + after_spill = tier.stats()["incremental_reused_blobs"] + assert after_spill > 0 + _write(tier, turns[2]) + assert tier.stats()["incremental_reused_blobs"] > after_spill + finally: + tier.close() From 65ad77ca757f41799b95ee31f357789fc3f65508 Mon Sep 17 00:00:00 2001 From: Jeroen van Menen <235760+jvmenen@users.noreply.github.com> Date: Tue, 20 Oct 2026 13:37:48 +0200 Subject: [PATCH 2/3] SSD cold tier: incremental-encode fingerprints in chunks of one unit The fingerprint pass reduced each blocked tensor's whole full-block prefix at once: an unbounded unit on the model-owner thread (issue #505) whose uint32 temporaries are several times the tensor's size. With the memo on the streaming spill, which also runs when the memory guard writes an entry out before releasing it, that peak lands where memory is short. Blocks are now fingerprinted per chunk of at most unit_bytes, the first time a block of the chunk needs one. A block's fingerprint depends only on its own bytes, so the content keys are unchanged. --- mtplx/cache_bank/codec.py | 57 +++++++++++++++++----- tests/test_cold_tier_incremental_encode.py | 30 ++++++++++++ 2 files changed, 76 insertions(+), 11 deletions(-) diff --git a/mtplx/cache_bank/codec.py b/mtplx/cache_bank/codec.py index 2d3cf40a1..ec39ec333 100644 --- a/mtplx/cache_bank/codec.py +++ b/mtplx/cache_bank/codec.py @@ -227,6 +227,37 @@ 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 @@ -347,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. @@ -366,9 +397,13 @@ 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._confirmed( self.reuse.get(key) if self.reuse is not None else None diff --git a/tests/test_cold_tier_incremental_encode.py b/tests/test_cold_tier_incremental_encode.py index 2c263d62b..6250474ad 100644 --- a/tests/test_cold_tier_incremental_encode.py +++ b/tests/test_cold_tier_incremental_encode.py @@ -640,3 +640,33 @@ def test_spill_and_queued_writes_share_the_session_memo(tmp_path, monkeypatch): assert tier.stats()["incremental_reused_blobs"] > after_spill finally: tier.close() + + +def test_fingerprints_run_in_bounded_chunks_with_the_same_keys(): + """Block fingerprints are taken per chunk of at most one unit, with the same + content keys as one pass over the whole tensor.""" + turn1, _ = _turns() + kwargs = { + "cache_snapshot": turn1.cache_snapshot, + "logits": turn1.logits, + "hidden": turn1.hidden, + "mtp_history_snapshot": turn1.mtp_history_snapshot, + "has_recurrent": True, + "block_size": 256, + "reuse": {}, + } + units = [] + whole = encode_payload(**kwargs) + small = encode_payload( + **kwargs, + unit_bytes=2 * 2 * 256 * 64 * 2, # two K blocks + on_unit=lambda kind, seconds, nbytes: units.append((kind, nbytes)), + ) + assert small.fingerprints == whole.fingerprints + assert small.spec == whole.spec and small.tensors == whole.tensors + block_fingerprints = [n for kind, n in units if kind == "fingerprint"] + # K and V: 4 full blocks in 2 chunks each; the MTP history (2 KiB a block + # row of 32 dims) fits in one chunk; the 1 MiB recurrent state is whole. + assert sorted(block_fingerprints) == sorted( + [2 * 2 * 256 * 64 * 2] * 4 + [4 * 256 * 32 * 2] + [1024 * 1024] + ) From 02f0991aee1ea46d945106c860052867505116e9 Mon Sep 17 00:00:00 2001 From: Jeroen van Menen <235760+jvmenen@users.noreply.github.com> Date: Tue, 20 Oct 2026 14:00:04 +0200 Subject: [PATCH 3/3] Changelog: the incremental encode on the streaming SSD write --- CHANGELOG.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3a5039769..43b1cbcef 100644 --- a/CHANGELOG.md +++ b/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