SSD session cache telemetry: where restores spend their time, where they land, and how much the cache writes

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]
+
+### Added
+
+- **SSD session cache telemetry** (by @jvmenen). Every SSD restore reports where its time went: reading blobs, copying them into arrays and the GPU evaluation, plus the bytes read, in its log line and in the request log (`ssd_restore_read_s`, `ssd_restore_copy_s`, `ssd_restore_eval_s`, `ssd_restore_other_s`, `ssd_restore_bytes`). Every RAM and SSD restore records where it landed: the tokens matched, the recurrent checkpoint it resumed at, the tokens lost between the two, and whether the entry came from the same conversation or another one sharing the prompt head (`prefix_restore_*` in the request log, counters per source under `session_bank.restore_landings` in `/health`). `/health` and `/admin/cache/ssd` show the bytes the writer put on disk over the last hour, the bytes it did not write because the blob was already there, both split into checkpoint and KV bytes, against `MTPLX_SSD_WRITE_BUDGET_PER_HOUR` (`write_volume`). The timing costs about 0.2 µs per tensor read.
+
 ## [2.12.2] - 2026-10-03
 
 ### Fixed

docs/server.md

@@ -249,6 +249,33 @@ writes in flight; `--force` overrides that when you know the daemon is idle.
 `--dir` defaults to the `ssd_session_cache_dir` in your saved config, then
 `~/.mtplx/session-bank`.
 
+**Where restores and writes go.** Every SSD restore logs one
+`SessionBank SSD restored` line with its `restore_s` split into `read_s`
+(reading blobs from disk), `copy_s` (copying the bytes into MLX arrays),
+`eval_s` (the GPU evaluation, which joins each tensor's blocks) and the bytes
+read. The request log carries the same split as `ssd_restore_read_s`,
+`ssd_restore_copy_s`, `ssd_restore_eval_s`, `ssd_restore_other_s` (manifest
+query, payload file, rows refused before reading) and `ssd_restore_bytes`.
+For every restore from RAM or SSD (the serial lanes, the batched MTP lane and
+the Gemma 4 pair) it also records where the restore landed:
+`prefix_restore_matched_tokens` (how much of the prompt the entry matched),
+`prefix_restore_point_tokens` (where the state resumed: on a hybrid model the
+newest recurrent checkpoint at or below the match),
+`prefix_restore_checkpoint_gap_tokens` (matched tokens computed again because
+no checkpoint sat closer), `prefix_restore_partial` and
+`prefix_restore_origin` (`same_session`, or `other_session` when another
+conversation's entry served a shared prompt head); the suffix prefilled after
+the restore point is `new_prefill_tokens`. A model without recurrent layers,
+such as Gemma 4, resumes at the matched length itself: its gap is 0. `/health` sums these per source
+under `session_bank.restore_landings`. `/health` (`ssd_session_cache`) and
+`/admin/cache/ssd` report the restore split totals (`restore_read_s_total`
+and so on, `last_restore_split`) and `write_volume`: the bytes written to
+disk over the last hour, the bytes not written because the blob was already
+on disk, each split into recurrent-checkpoint (`gdn_boundary_*`) and other
+(`kv_*`: attention KV plus the final state, logits and hidden) bytes, and the
+share of `MTPLX_SSD_WRITE_BUDGET_PER_HOUR` written. `written_bytes_last_hour`
+next to it stays what the budget booked, an estimate for queued writes.
+
 ## Sharing on your network (other devices, Parallels/VM guests)
 
 The default bind is `127.0.0.1`: only this Mac can connect. To reach MTPLX

mtplx/backends/gemma4_assistant.py

@@ -477,6 +477,29 @@ class Gemma4PromptState:
     cache_hit: bool = False
     cache_miss_reason: str | None = None
     restore_mode: str = "cold"
+    # Where a session-bank restore came from and what it served, in the
+    # shape of the Qwen lanes (generation.PromptState): the request log's
+    # ssd_* and prefix_restore_* fields read these.
+    cache_source: str = "none"
+    ssd_cache_hit: bool = False
+    ssd_restore_s: float = 0.0
+    restore_served: dict[str, Any] = field(default_factory=dict)
+
+
+def _gemma4_restore_stats(prompt_state: Any) -> dict[str, Any]:
+    """GenerationStats fields for the prompt state's restore (read with
+    defaults: callers may hand in a prompt state without these fields)."""
+    ssd_hit = bool(getattr(prompt_state, "ssd_cache_hit", False))
+    return {
+        "cache_source": str(getattr(prompt_state, "cache_source", "none") or "none"),
+        "ssd_cache_hit": ssd_hit,
+        "ssd_cached_tokens": int(prompt_state.cached_tokens) if ssd_hit else 0,
+        "ssd_restore_s": (
+            float(getattr(prompt_state, "ssd_restore_s", 0.0) or 0.0) if ssd_hit else 0.0
+        ),
+        "ssd_suffix_tokens": int(prompt_state.suffix_tokens) if ssd_hit else 0,
+        "session_restore_served": dict(getattr(prompt_state, "restore_served", None) or {}),
+    }
 
 
 @dataclass(frozen=True)
@@ -3102,16 +3125,26 @@ def _restore_or_prefill_gemma4_prompt(
     prompt_ids: list[int],
     *,
     session_bank: Any | None = None,
+    session_id: str | None = None,
     session_restore_mode: str = "clone",
     session_template_hash: str | None = None,
     session_draft_head_identity: str | None = None,
     session_policy_fingerprint: str | None = None,
     require_shared_kv: bool,
     abort_check: Any | None = None,
 ) -> Gemma4PromptState:
+    """Restore the longest banked prefix, or prefill cold.
+
+    ``session_id`` is used for telemetry only (whether a restored entry came
+    from the same conversation); the restore itself does not depend on it.
+    """
     if not prompt_ids:
         raise ValueError("prompt_ids must not be empty for Gemma 4 generation")
-    from mtplx.generation import _check_postcommit_abort
+    from mtplx.generation import (
+        _check_postcommit_abort,
+        _record_restore_landing,
+        served_restore_record,
+    )
 
     def forward(tokens: list[int], cache: Any) -> tuple[Gemma4TargetOutput, float]:
         # The request's abort site (the per-chunk memory check, the
@@ -3230,7 +3263,11 @@ def cold_prefill(reason: str | None = None) -> Gemma4PromptState:
                 output, _suffix_elapsed = forward(list(prompt_ids[seed:]), cache)
                 entry.hits += 1
                 entry.last_access_s = time.time()
-                return Gemma4PromptState(
+                # Gemma 4 has no recurrent layers: the restore lands on the
+                # matched length itself, so the checkpoint gap is 0.
+                cache_source = str(getattr(entry, "cache_source", "ram") or "ram")
+                ssd_hit = cache_source == "ssd"
+                near_state = Gemma4PromptState(
                     cache=cache,
                     logits=output.logits[:, -1, :],
                     hidden=output.hidden[:, -1:, :],
@@ -3242,13 +3279,43 @@ def cold_prefill(reason: str | None = None) -> Gemma4PromptState:
                     cache_hit=True,
                     cache_miss_reason=None,
                     restore_mode=f"block_prefix_{storage_mode}",
+                    cache_source=cache_source,
+                    ssd_cache_hit=ssd_hit,
+                    ssd_restore_s=float(getattr(entry, "ssd_restore_s", 0.0) or 0.0),
+                    restore_served=served_restore_record(
+                        entry,
+                        matched=matched,
+                        restore_point=restore_point,
+                        storage_restore_mode=str(storage_mode),
+                        session_id=session_id,
+                        ssd_cache_hit=ssd_hit,
+                        ssd_restore_split=getattr(entry, "ssd_restore_split", None),
+                    ),
                 )
+                _record_restore_landing(session_bank, near_state)
+                return near_state
         return cold_prefill(getattr(session_bank, "last_miss_reason", None))
 
+    cache_source = str(getattr(restored, "cache_source", "ram") or "ram")
+    ssd_hit = bool(getattr(restored, "ssd_cache_hit", False))
+    restore_fields = {
+        "cache_source": cache_source,
+        "ssd_cache_hit": ssd_hit,
+        "ssd_restore_s": float(getattr(restored, "ssd_restore_s", 0.0) or 0.0),
+        "restore_served": served_restore_record(
+            restored.entry,
+            matched=int(restored.entry.prefix_len),
+            restore_point=int(restored.entry.prefix_len),
+            storage_restore_mode=str(restored.restore_mode or session_restore_mode),
+            session_id=session_id,
+            ssd_cache_hit=ssd_hit,
+            ssd_restore_split=getattr(restored, "ssd_restore_split", None),
+        ),
+    }
     suffix = list(prompt_ids[restored.entry.prefix_len :])
     if suffix:
         output, suffix_elapsed = forward(suffix, restored.cache)
-        return Gemma4PromptState(
+        exact_state = Gemma4PromptState(
             cache=restored.cache,
             logits=output.logits[:, -1, :],
             hidden=output.hidden[:, -1:, :],
@@ -3260,7 +3327,10 @@ def cold_prefill(reason: str | None = None) -> Gemma4PromptState:
             cache_hit=True,
             cache_miss_reason=None,
             restore_mode=str(restored.restore_mode or session_restore_mode),
+            **restore_fields,
         )
+        _record_restore_landing(session_bank, exact_state)
+        return exact_state
 
     extra_state = restored.extra_state if isinstance(restored.extra_state, dict) else {}
     shared_kv_states = extra_state.get("gemma4_shared_kv_states")
@@ -3271,7 +3341,7 @@ def cold_prefill(reason: str | None = None) -> Gemma4PromptState:
     kv_offset = extra_state.get("gemma4_kv_offset")
     if kv_offset is None:
         kv_offset = runtime.target.cache_offset(restored.cache)
-    return Gemma4PromptState(
+    exact_state = Gemma4PromptState(
         cache=restored.cache,
         logits=restored.logits,
         hidden=restored.hidden,
@@ -3283,7 +3353,10 @@ def cold_prefill(reason: str | None = None) -> Gemma4PromptState:
         cache_hit=True,
         cache_miss_reason=None,
         restore_mode=str(restored.restore_mode or session_restore_mode),
+        **restore_fields,
     )
+    _record_restore_landing(session_bank, exact_state)
+    return exact_state
 
 
 def generate_gemma4_ar(
@@ -3300,6 +3373,7 @@ def generate_gemma4_ar(
     prefill_callback: Any | None = None,
     repetition_stop: bool = False,
     session_bank: Any | None = None,
+    session_id: str | None = None,
     session_restore_mode: str = "clone",
     session_template_hash: str | None = None,
     session_draft_head_identity: str | None = None,
@@ -3358,6 +3432,7 @@ def generate_gemma4_ar(
         session_policy_fingerprint=session_policy_fingerprint,
         require_shared_kv=False,
         abort_check=abort_check,
+        session_id=session_id,
     )
     if prefill_callback is not None:
         try:
@@ -3510,6 +3585,7 @@ def generate_gemma4_ar(
         session_cache_hit=bool(prompt_state.cache_hit),
         cache_miss_reason=prompt_state.cache_miss_reason,
         session_restore_mode=prompt_state.restore_mode,
+        **_gemma4_restore_stats(prompt_state),
         target_forward_time_s=prompt_state.prompt_eval_time_s + target_decode_time,
         verify_time_s=target_decode_time,
         verify_forward_time_s=target_decode_time,
@@ -3558,6 +3634,7 @@ def generate_gemma4_assistant(
     stop_token_ids: set[int] | None = None,
     token_callback: Any | None = None,
     session_bank: Any | None = None,
+    session_id: str | None = None,
     session_restore_mode: str = "clone",
     session_template_hash: str | None = None,
     session_draft_head_identity: str | None = None,
@@ -3627,6 +3704,7 @@ def generate_gemma4_assistant(
         session_policy_fingerprint=session_policy_fingerprint,
         require_shared_kv=True,
         abort_check=abort_check,
+        session_id=session_id,
     )
     if prefill_callback is not None:
         try:
@@ -3945,6 +4023,7 @@ def generate_gemma4_assistant(
         session_cache_hit=bool(prompt_state.cache_hit),
         cache_miss_reason=prompt_state.cache_miss_reason,
         session_restore_mode=prompt_state.restore_mode,
+        **_gemma4_restore_stats(prompt_state),
         target_forward_time_s=prompt_state.prompt_eval_time_s + float(timing_totals["verify"]),
         verify_time_s=float(timing_totals["verify"]),
         verify_forward_time_s=float(timing_totals["verify"]),

mtplx/cache_bank/codec.py

@@ -12,6 +12,8 @@
 
 from mtplx.cache_state import CacheSnapshot
 
+from .telemetry import active_restore_timings
+
 
 _MLX_TO_NUMPY_DTYPE: dict[str, Any] = {
     "bool": np.bool_,
@@ -56,6 +58,8 @@ class EncodedPayload:
     reused: dict[str, dict[str, Any]] = field(default_factory=dict)
     reused_nbytes: int = 0
     fingerprints: dict[str, tuple] = field(default_factory=dict)
+    # Seconds of eval and host copy per captured name (telemetry.WriteSplit).
+    capture_s: dict[str, float] = field(default_factory=dict)
 
 
 @dataclass(frozen=True)
@@ -171,6 +175,11 @@ def __init__(
         self.reuse = reuse
         self.reused: dict[str, dict[str, Any]] = {}
         self.fingerprints: dict[str, tuple] = {}
+        # Seconds of eval and host copy behind each captured name, set just
+        # before the name's bytes land in ``tensors``, so a streaming sink can
+        # read its own blob's cost there.
+        self.capture_s: dict[str, float] = {}
+        self._unit_s = 0.0
         self.tensors: dict[str, bytes] = {}
         self.block_size = max(1, int(block_size))
         self.should_abort = should_abort
@@ -194,15 +203,23 @@ def _materialize(self, kind: str, array: Any) -> bytes:
         started = time.perf_counter()
         mx.eval(array)
         raw = bytes(memoryview(array))
+        elapsed = time.perf_counter() - started
+        self._unit_s += elapsed
         observer = self.on_unit
         if observer is not None:
             try:
-                observer(kind, time.perf_counter() - started, len(raw))
+                observer(kind, elapsed, len(raw))
             except Exception:
                 pass
         self._check_abort()
         return raw
 
+    def _store(self, name: str, raw: bytes) -> None:
+        """Hand a captured blob to ``tensors`` with the units it took."""
+        self.capture_s[name] = self._unit_s
+        self._unit_s = 0.0
+        self.tensors[name] = raw
+
     def _fingerprints(self, value: Any, *, rows: int | None) -> list[int] | None:
         """One bounded unit like _materialize: a single GPU reduction over
         ``value``, reported to on_unit, followed by the abort check."""
@@ -294,7 +311,7 @@ def _encode_tensor(self, value: Any) -> dict[str, Any]:
                         "nbytes": int(hit["nbytes"]),
                     }
         raw = self._tensor_bytes(value, shape)
-        self.tensors[name] = raw
+        self._store(name, raw)
         return {
             "kind": "tensor",
             "name": name,
@@ -375,7 +392,7 @@ def _encode_tensor_blocks(
                     )
                     continue
             raw = self._materialize("block", chunk)
-            self.tensors[name] = raw
+            self._store(name, raw)
             total += len(raw)
             blocks.append(
                 {
@@ -511,6 +528,7 @@ def encode_payload(
         reused=dict(codec.reused),
         reused_nbytes=sum(int(blob["nbytes"]) for blob in codec.reused.values()),
         fingerprints=dict(codec.fingerprints),
+        capture_s=dict(codec.capture_s),
     )
 
 
@@ -684,7 +702,18 @@ def _eval_decoded_arrays(decoded: DecodedPayload) -> None:
     ):
         _collect_arrays(part, arrays)
     if arrays:
-        mx.eval(*arrays)
+        _timed_eval(arrays)
+
+
+def _timed_eval(arrays: list[Any]) -> None:
+    """``mx.eval``, counted as restore eval time while a restore is timed:
+    the GPU side of the decode, the concatenate of each tensor's blocks
+    included."""
+    timings = active_restore_timings()
+    started = time.perf_counter() if timings is not None else 0.0
+    mx.eval(*arrays)
+    if timings is not None:
+        timings.eval_s += time.perf_counter() - started
 
 
 def decode_gdn_boundaries(
@@ -731,7 +760,7 @@ def interrupted() -> bool:
         arrays: list[Any] = []
         _collect_arrays(decoded, arrays)
         if arrays:
-            mx.eval(*arrays)
+            _timed_eval(arrays)
         boundaries.append(decoded)
     return tuple(boundaries)
 
@@ -886,10 +915,11 @@ def _decode_tree_prefix(
     return decode_tree(spec, read_tensor)
 
 
-def _decode_tensor(spec: dict[str, Any], read_tensor: Callable[[str], bytes]) -> Any:
-    dtype = str(spec["dtype"])
-    shape = tuple(int(dim) for dim in spec.get("shape") or [])
-    raw = read_tensor(str(spec["name"]))
+def _array_from_buffer(raw: Any, dtype: str) -> Any:
+    """One host copy of ``raw`` (bytes or a byte buffer) into a flat array,
+    counted as restore copy time while a restore is timed."""
+    timings = active_restore_timings()
+    started = time.perf_counter() if timings is not None else 0.0
     if dtype == "bfloat16":
         arr = mx.array(np.frombuffer(raw, dtype=np.uint16)).view(mx.bfloat16)
     else:
@@ -898,6 +928,29 @@ def _decode_tensor(spec: dict[str, Any], read_tensor: Callable[[str], bytes]) ->
         if np_dtype is None or mlx_dtype is None:
             raise ValueError(f"unsupported persisted tensor dtype: {dtype!r}")
         arr = mx.array(np.frombuffer(raw, dtype=np_dtype), dtype=mlx_dtype)
+    if timings is not None:
+        timings.copy_s += time.perf_counter() - started
+    return arr
+
+
+def _read_blob(read_tensor: Callable[[str], bytes], name: str) -> bytes:
+    """``read_tensor(name)``, counted as restore read time and bytes while a
+    restore is timed."""
+    timings = active_restore_timings()
+    if timings is None:
+        return read_tensor(name)
+    started = time.perf_counter()
+    raw = read_tensor(name)
+    timings.read_s += time.perf_counter() - started
+    timings.bytes_read += len(raw)
+    return raw
+
+
+def _decode_tensor(spec: dict[str, Any], read_tensor: Callable[[str], bytes]) -> Any:
+    dtype = str(spec["dtype"])
+    shape = tuple(int(dim) for dim in spec.get("shape") or [])
+    raw = _read_blob(read_tensor, str(spec["name"]))
+    arr = _array_from_buffer(raw, dtype)
     if shape:
         arr = arr.reshape(shape)
     return arr