fix(offload): refresh state images in LMCache's LRU on every tier hit and keep the index consistent

atom/kv_transfer/disaggregation/aggregator.py

@@ -30,6 +30,7 @@
     ReqId,
     SaveCompletionId,
     SaveOperationId,
+    StateTouchId,
 )
 
 logger = logging.getLogger("atom")
@@ -156,10 +157,12 @@ def __init__(
             terminal_tombstone_limit,
             lambda key: isinstance(key, LoadOperationId),
         )
+        # Touch ids are never reused (fresh generation), and tombstoning them would
+        # push the store / source tombstones out of the ring.
         self._connector_completions = _TPCompletionGroup[ConnectorCompletionKey](
             world_size,
             terminal_tombstone_limit,
-            lambda _key: True,
+            lambda key: not isinstance(key[1], StateTouchId),
         )
 
     @property

atom/kv_transfer/disaggregation/multi/multi_connector.py

@@ -679,6 +679,21 @@ def enqueue_state_stores(self, stores) -> bool:
             return False
         return bool(c.enqueue_state_stores(stores))
 
+    def enqueue_state_touches(self, hashes) -> bool:
+        """State touches to the sub hosting the tier (`_state_tier_sub`);
+        False if none does."""
+        c = self._state_tier_sub()
+        if c is None:
+            return False
+        return bool(c.enqueue_state_touches(hashes))
+
+    def take_state_touch_verdicts(self) -> dict:
+        """`{hash: resident}` from the tier sub; empty if none."""
+        c = self._state_tier_sub()
+        if c is None:
+            return {}
+        return c.take_state_touch_verdicts()
+
     def take_state_reports(self):
         """`(indexed, failed)` from the tier sub, else two empty sets.
 

atom/kv_transfer/disaggregation/types.py

@@ -123,12 +123,25 @@ def __post_init__(self) -> None:
             raise ValueError("state store generation must be nonnegative")
 
 
+@dataclass(frozen=True)
+class StateTouchId:
+    """One LRU refresh of a voted state image. Every TP worker reports on its
+    own shard under this id, so the failure-dominant quorum says gone when any
+    shard is. `generation` keeps repeated touches of one hash distinct
+    identities.
+    """
+
+    prefix_hash: int
+    generation: int
+
+
 ConnectorCompletionId = (
     ReqId
     | SaveOperationId
     | SaveSourceGroupId
     | LoadOperationId
     | StateStoreOperationId
+    | StateTouchId
 )
 ConnectorCompletionKey = tuple[str, ConnectorCompletionId]
 

atom/kv_transfer/offload/README.md

@@ -122,6 +122,17 @@ The engine-side counterpart of the state tier lives outside this directory:
 store/finished/failed hash sets that `hybrid/kimi_k3/state_tier.py` reports back,
 so a later prefix hit knows which checkpoints the CPU tier can actually serve.
 
+State images and paged KV share one LocalCPU pool and one LRU, and a state hit
+is answered by the index rather than an LMCache lookup. So every voted hit is
+also sent to the workers as a touch (`LMCacheOffloadMetadata.state_touches`):
+each rank refreshes the image (`StateByteCodec.touch`) and reports on
+`k3_state_touch` whether it is still there; a gone image leaves the index. An
+admission that would load an image not confirmed within `tier_verified_for_s`
+(0.25 s per GiB of `LMCACHE_MAX_LOCAL_CPU_SIZE`, 5-60 s) is held a pass, at most
+`OFFLOAD_LOOKUP_DEFER_S`, for that verdict. Touches that only confirmed presence
+(the image is held by another backend only, or the touch raised) are counted in
+the connector statistic `state_touches_unrefreshed`.
+
 ## LMCache Multiprocess (`lmcache_mp`)
 
 The `lmcache_mp` connector selects one of two paths from capabilities published
@@ -1458,6 +1469,7 @@ LMCache version, these are what to re-check.**
 | `get_or_create(…, lambda t,s: None, lambda o,s: o)` | LMCache's trailing token-processing / output-transform callbacks | We don't use LMCache's token-shaping hooks — our codec moves raw bytes | Passed as no-op / identity callables |
 | `cfg.lookup_server_worker_ids = [0]` | default: every rank answers lookup, client takes `min()` | At TP>1 a non-rank-0 shard returning 0 would zero out a real hit; rank 0 is made authoritative | `config.py` (see [TP > 1 Notes](#tp--1-notes)) |
 | `cfg.use_gds = False` | LMCache may enable cufile GDS | cufile init hangs without NVMe-GDS hardware here | `config.py` |
+| `StateByteCodec.touch` | `LocalCPUBackend.touch_cache()` after `contains(pin=True)` | A pin needs a paired unpin, and `keys_in_request` belongs to the concurrent KV lookup | `cache_policy.update_on_hit(key, hot_cache)` under `cpu_lock`, resolved once at bind; if one is missing it warns and touches only confirm presence |
 
 ### 3. Fully delegated to LMCache (we never touch the implementation)
 

atom/kv_transfer/offload/connector.py

@@ -346,6 +346,16 @@ def enqueue_state_stores(self, stores) -> bool:
         callback = getattr(self._impl, "enqueue_state_stores", None)
         return bool(callback(stores)) if callback is not None else False
 
+    def enqueue_state_touches(self, hashes) -> bool:
+        """State touches for the impl; False when it carries none."""
+        callback = getattr(self._impl, "enqueue_state_touches", None)
+        return bool(callback(hashes)) if callback is not None else False
+
+    def take_state_touch_verdicts(self) -> dict:
+        """`{hash: resident}` for touches that reached quorum."""
+        callback = getattr(self._impl, "take_state_touch_verdicts", None)
+        return callback() if callback is not None else {}
+
     def take_state_reports(self):
         # Contract is a 2-tuple (indexed, failed) -- matched by the kimi_k3 impl
         # and by the sole caller `indexed, failed = take()` in

atom/kv_transfer/offload/hybrid/kimi_k3/connector.py

@@ -16,6 +16,7 @@
 from atom.kv_transfer.disaggregation.types import (
     ConnectorCompletion,
     StateStoreOperationId,
+    StateTouchId,
 )
 from atom.kv_transfer.offload import config as offcfg
 from atom.kv_transfer.offload._offload_common import (
@@ -30,7 +31,10 @@
     DenseOffloadScheduler,
 )
 from atom.kv_transfer.offload.hybrid.kimi_k3.staging import StagedTransfer
-from atom.kv_transfer.offload.hybrid.kimi_k3.state_object import StateByteCodec
+from atom.kv_transfer.offload.hybrid.kimi_k3.state_object import (
+    StateByteCodec,
+    TouchResult,
+)
 from atom.kv_transfer.offload.hybrid.kimi_k3.state_tier import (
     StateOffloadTier,
     _JointPark,
@@ -61,6 +65,14 @@
 #: rather than leaving a partially-reported key pending forever.
 STATE_LOAD_DISPOSITION_CHANNEL = "k3_state_load_disposition"
 
+#: One verdict per `StateTouchId`: `succeeded` is whether the image is still
+#: resident. Failure-dominant across TP: one missing shard and a load misses.
+STATE_TOUCH_CHANNEL = "k3_state_touch"
+
+#: Same id: `succeeded=False` when the image was present but not refreshed
+#: (outside LocalCPU's hot cache).
+STATE_TOUCH_REFRESH_CHANNEL = "k3_state_touch_refresh"
+
 #: The connector's standalone stall clock, used only when reclamation is
 #: disabled (`LMCACHE_EC_PIN_TIMEOUT_SEC <= 0`) so there is no abandon window to
 #: stay under. A save outstanding longer than this is a backend that stopped,
@@ -115,6 +127,11 @@ def __init__(self, config) -> None:
         # the no-tier store path and took the whole step's KV loads/saves down
         # with it (super().start_load_kv never ran).
         self._store_failed_no_tier: set[StateStoreOperationId] = set()
+        # Touch verdicts since the last `get_finished`: (resident, refreshed).
+        # GONE counts as refreshed: nothing was left unrefreshed.
+        self._touch_verdicts: dict[StateTouchId, tuple[bool, bool]] = {}
+        # Only the first touch whose codec call raised is logged.
+        self._warned_touch_failure = False
 
     def close(self) -> None:
         """Drain then join the state tier before the base executors.
@@ -278,12 +295,46 @@ def start_load_kv(self, metadata) -> None:
         # KV submit to run even when arming half-completed -- an arriving KV
         # completion then settles or evicts whatever the arm did file.
         try:
+            self._touch_states(metadata)
             self._arm_joint_loads(metadata)
             self._start_state_loads(metadata)
             self._start_state_stores(metadata)
         finally:
             super().start_load_kv(metadata)
 
+    def _touch_states(self, metadata) -> None:
+        """Refresh the voted state images and record whether each is still
+        there. Inline, not on the tier's executors: a touch is one dict move
+        and must not queue behind an image-sized store. With no tier every
+        image is gone."""
+        touches = getattr(metadata, "state_touches", None)
+        if not touches:
+            return
+        tier = self._state_tier
+        for touch_id in touches:
+            if tier is None:
+                self._touch_verdicts[touch_id] = (False, True)
+                continue
+            try:
+                result = tier.codec.touch(int(touch_id.prefix_hash))
+            except Exception:
+                # Says nothing about the bytes: kept, not refreshed. One
+                # traceback, since a drifted LMCache raises on every touch.
+                if not self._warned_touch_failure:
+                    self._warned_touch_failure = True
+                    logger.warning(
+                        "kimi_k3 offload: touch of state hash %d failed; failed "
+                        "touches keep the image indexed and are counted in "
+                        "state_touches_unrefreshed",
+                        touch_id.prefix_hash,
+                        exc_info=True,
+                    )
+                result = TouchResult.PRESENT
+            self._touch_verdicts[touch_id] = (
+                bool(result),
+                result is not TouchResult.PRESENT,
+            )
+
     def _arm_joint_loads(self, metadata) -> None:
         """Hold a request owning both legs until both report.
 
@@ -463,6 +514,16 @@ def get_finished(self):
                     ConnectorCompletion(STATE_SOURCE_CHANNEL, op, True)
                 )
             self._store_failed_no_tier = set()
+        # Before the tier-None return too: `_touch_states` answers for a
+        # missing tier, and the engine may be holding an admission for it.
+        for touch_id, (resident, refreshed) in self._touch_verdicts.items():
+            out.connector_completions.add(
+                ConnectorCompletion(STATE_TOUCH_CHANNEL, touch_id, resident)
+            )
+            out.connector_completions.add(
+                ConnectorCompletion(STATE_TOUCH_REFRESH_CHANNEL, touch_id, refreshed)
+            )
+        self._touch_verdicts = {}
         if self._state_tier is None:
             # `_arm_joint_loads` armed these even with no tier and
             # `_fail_state_loads` failed their state leg; drain the park here,
@@ -605,6 +666,10 @@ def __init__(self, config) -> None:
         # block ids, no chunking -- only the park/report lifecycle.
         self._pending_state_loads: list[tuple] = []
         self._pending_state_stores: list[tuple] = []
+        self._pending_state_touches: list[StateTouchId] = []
+        self._touch_generation = 0
+        self._state_touch_verdicts: dict[int, bool] = {}
+        self._state_touches_unrefreshed = 0
         # A finished request whose save is queued keeps its blocks pinned
         # (`should_defer_free`), so the queue depth is also how much of the pool
         # a slow backend can hold. Same knob and default the DSV4 layout bounds
@@ -642,6 +707,28 @@ def enqueue_state_stores(self, stores) -> bool:
         self._pending_state_stores.extend(stores)
         return True
 
+    def enqueue_state_touches(self, hashes) -> bool:
+        """Queue touches for the next metadata, each under its own generation."""
+        for h in hashes:
+            self._pending_state_touches.append(
+                StateTouchId(int(h), self._touch_generation)
+            )
+            self._touch_generation += 1
+        return True
+
+    def take_state_touch_verdicts(self) -> dict[int, bool]:
+        """`{hash: resident}` from the touches that reached TP quorum."""
+        verdicts = self._state_touch_verdicts
+        self._state_touch_verdicts = {}
+        return verdicts
+
+    def get_statistics(self) -> dict[str, int]:
+        # Touches that only confirmed presence, without refreshing the LRU.
+        return {
+            **super().get_statistics(),
+            "state_touches_unrefreshed": self._state_touches_unrefreshed,
+        }
+
     def build_connector_meta(self) -> LMCacheOffloadMetadata:
         self._refresh_save_stall()
         meta = super().build_connector_meta()
@@ -654,6 +741,8 @@ def build_connector_meta(self) -> LMCacheOffloadMetadata:
         # already released.
         meta.state_stores = self._pending_state_stores
         self._pending_state_stores = []
+        meta.state_touches = self._pending_state_touches
+        self._pending_state_touches = []
         return meta
 
     def _may_emit_save(self) -> bool:
@@ -679,6 +768,7 @@ def has_pending_work(self) -> bool:
             super().has_pending_work()
             or bool(self._pending_state_loads)
             or bool(self._pending_state_stores)
+            or bool(self._pending_state_touches)
         )
 
     # -- save stall --------------------------------------------------------
@@ -845,7 +935,7 @@ def _claim_after_load(self, seq, hbm: int, lmc: int) -> int:
         return max(hbm, min(joint, lmc)) if joint else max(hbm, lmc)
 
     # -- connector-owned channels -----------------------------------------
-    def connector_completion(self, completion) -> bool:
+    def connector_completion(self, completion) -> bool | None:
         if completion.channel == STATE_SOURCE_CHANNEL:
             self._state_source_released.add(completion.operation_id)
             return True
@@ -857,6 +947,18 @@ def connector_completion(self, completion) -> bool:
             )
             target.add(completion.operation_id)
             return True
+        if completion.channel == STATE_TOUCH_CHANNEL:
+            h = int(completion.operation_id.prefix_hash)
+            self._state_touch_verdicts[h] = self._state_touch_verdicts.get(
+                h, True
+            ) and bool(completion.succeeded)
+            # Handled, but not a terminal save: `process_completions` would
+            # file the touch id under `finished_saving`.
+            return None
+        if completion.channel == STATE_TOUCH_REFRESH_CHANNEL:
+            if not completion.succeeded:
+                self._state_touches_unrefreshed += 1
+            return None
         if completion.channel == STATE_LOAD_DISPOSITION_CHANNEL:
             # Finding #6: a state load whose bytes survived a failed pair. Record
             # only the survivors -- their `failed_loading` report must abandon