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,13 @@ def __init__(
             terminal_tombstone_limit,
             lambda key: isinstance(key, LoadOperationId),
         )
+        # A state touch has a fresh round in every identity, so it is never
+        # reported twice; tombstoning touches (one per voted hash per 0.5 s)
+        # would only 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

@@ -655,6 +655,36 @@ def enqueue_state_stores(self, stores) -> bool:
             return False
         return bool(c.enqueue_state_stores(stores))
 
+    def enqueue_state_touches(self, hashes) -> bool:
+        """LRU touches for voted state images, to the one sub hosting the
+        tier; False if none does (the engine then stops sending them).
+
+        Select by `has_state_tier`, not method presence -- see
+        `_state_tier_sub`.
+        """
+        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.
+
+        Select by `has_state_tier`, not method presence -- see
+        `_state_tier_sub`.
+        """
+        c = self._state_tier_sub()
+        if c is None:
+            return {}
+        return c.take_state_touch_verdicts()
+
+    def take_state_touches_unrefreshed(self) -> int:
+        """Presence-only touches from the tier sub; 0 if none."""
+        c = self._state_tier_sub()
+        if c is None:
+            return 0
+        return int(c.take_state_touches_unrefreshed())
+
     def take_state_reports(self):
         """`(indexed, failed)` from the tier sub, else two empty sets.
 

atom/kv_transfer/disaggregation/types.py

@@ -109,10 +109,10 @@ class StateStoreOperationId:
     Not keyed by request: by the time a state store lands, the request that
     produced the checkpoint is long gone and only the prefix hash remains. But
     the hash alone is not an identity -- the same prefix is stored again after
-    an eviction or a load miss, and `KVOutputAggregator` tombstones every
-    `(channel, operation_id)` it has taken quorum on, so a second store under a
-    bare hash is dropped as a duplicate: its pin is never settled and its
-    bytes are never re-indexed. `generation` is what separates the attempts.
+    an eviction or a load miss, and `KVOutputAggregator` tombstones the
+    `(channel, operation_id)` of every store it has taken quorum on, so a
+    second store under a bare hash is dropped as a duplicate: its pin is never
+    settled and its bytes are never re-indexed. `generation` is what separates the attempts.
     """
 
     prefix_hash: int
@@ -123,12 +123,33 @@ def __post_init__(self) -> None:
             raise ValueError("state store generation must be nonnegative")
 
 
+@dataclass(frozen=True)
+class StateTouchId:
+    """One recency refresh of a state image the engine just voted a hit for.
+
+    Every TP worker touches its own shard and reports whether the image was
+    still resident under this id, so the aggregator's failure-dominant quorum
+    says "gone" as soon as any shard is gone -- which is when a joint load
+    would miss. `round` makes repeated touches of one hash distinct identities,
+    so a quorum never merges two touches and the aggregator keeps no tombstone
+    for them.
+    """
+
+    prefix_hash: int
+    round: int
+
+    def __post_init__(self) -> None:
+        if self.round < 0:
+            raise ValueError("state touch round must be nonnegative")
+
+
 ConnectorCompletionId = (
     ReqId
     | SaveOperationId
     | SaveSourceGroupId
     | LoadOperationId
     | StateStoreOperationId
+    | StateTouchId
 )
 ConnectorCompletionKey = tuple[str, ConnectorCompletionId]
 

atom/kv_transfer/offload/README.md

@@ -122,6 +122,27 @@ 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. A KV lookup
+refreshes the chunks it touches, but a state hit is answered by the index, so
+every hit the index votes for is also sent to the workers as a *touch*
+(`LMCacheOffloadMetadata.state_touches`): each rank moves the image to the young
+end of LMCache's LRU (`StateByteCodec.touch`) and reports whether it was still
+there on the `k3_state_touch` channel. A prefix that hits every turn keeps its
+image as warm as its KV, and an image LMCache evicted leaves the index when its
+verdict lands. An image only another backend (LocalDisk) still holds is
+present but outside the hot cache: the touch cannot refresh it, reports that on
+`k3_state_touch_refresh`, and it is counted in
+`state_offload_touches_unrefreshed` (0 when every touch refreshed the LRU); the
+next store of it puts it back into the CPU tier. A gone verdict from a touch
+sent before the image's latest store is ignored as stale. An admission about to
+load an image not confirmed recently is held for a pass (at most `OFFLOAD_LOOKUP_DEFER_S`) while that verdict comes
+back, so it plans a lower boundary instead of parking on a miss and
+recomputing the whole prompt. "Recently" scales with the pool the image has to
+age out of: 0.25 s per GiB of `LMCACHE_MAX_LOCAL_CPU_SIZE`, within 5-60 s
+(52 s for a 208.6 GiB pool; `tier_verified_for_s`). `OFFLOAD_STATE_CPU_SIZE` caps how much of
+the pool images may take; it reserves nothing, and the worker warns at startup
+when the cap is not below half of `LMCACHE_MAX_LOCAL_CPU_SIZE`.
+
 ## LMCache Multiprocess (`lmcache_mp`)
 
 The `lmcache_mp` connector selects one of two paths from capabilities published
@@ -1445,6 +1466,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 state hit is answered by ATOM's index, so no LMCache lookup refreshes the image; the pin would need a paired unpin, and `keys_in_request` belongs to the concurrent KV lookup | Calls `cache_policy.update_on_hit(key, hot_cache)` under `cpu_lock`. `bind_storage_manager` resolves `local_cpu_backend` and those three internals once. If one is missing, touches only confirm presence (loads still work, nothing is refreshed): it warns once at bind and every touch is counted in `state_offload_touches_unrefreshed` |
 
 ### 3. Fully delegated to LMCache (we never touch the implementation)
 

atom/kv_transfer/offload/connector.py

@@ -335,6 +335,24 @@ 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:
+        """LRU touches for voted state images. False when the impl carries
+        none, which tells the engine to stop asking (and to stop holding
+        admissions for verdicts)."""
+        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; empty when the
+        impl carries none."""
+        callback = getattr(self._impl, "take_state_touch_verdicts", None)
+        return callback() if callback is not None else {}
+
+    def take_state_touches_unrefreshed(self) -> int:
+        """Touches that only confirmed presence; 0 when the impl carries none."""
+        callback = getattr(self._impl, "take_state_touches_unrefreshed", None)
+        return int(callback()) if callback is not None else 0
+
     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